Skip to content

Commit 7b4de93

Browse files
thisiscamthisiscam
authored
Merge pull request #7 from thisiscam/example-notebooks
Examples as rendered notebooks + neural ACNet example
2 parents 394b70a + 778ec9c commit 7b4de93

20 files changed

Lines changed: 2110 additions & 422 deletions

‎.gitignore‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,3 +21,6 @@ env/
2121

2222
# MkDocs build output
2323
site/
24+
25+
# mkdocs-jupyter build cache
26+
.cache/

‎README.md‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ For background, see the paper *Archimedean Copula Inference via Taylor-Mode
1313
AD* (arXiv:TBD).
1414

1515
**Documentation:** <https://thisiscam.github.io/acopula/> ·
16-
**Examples:** [`examples/`](examples/)
16+
**Examples:** [rendered notebooks](https://thisiscam.github.io/acopula/examples/01_quickstart/) ([source](docs/examples/))
1717

1818
## Install
1919

@@ -28,8 +28,8 @@ with jax 0.8. For development, clone and use [uv](https://docs.astral.sh/uv/):
2828
```bash
2929
git clone https://github.com/thisiscam/acopula
3030
cd acopula
31-
uv sync # add --extra examples for the plotting examples
32-
uv run python examples/01_quickstart.py
31+
uv sync --extra examples
32+
uv run jupyter lab docs/examples/ # the example notebooks
3333
```
3434

3535
A PyPI `pip install acopula` is **not available yet**: the git-pinned `oryx`

‎docs/assets/acopula_pairs.png‎

-774 KB
Binary file not shown.

‎docs/assets/acopula_tree.png‎

-41.8 KB
Binary file not shown.

‎docs/examples.md‎

Lines changed: 0 additions & 59 deletions
This file was deleted.

‎docs/examples/01_quickstart.ipynb‎

Lines changed: 234 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,234 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "markdown",
5+
"id": "ec609f4a",
6+
"metadata": {},
7+
"source": [
8+
"# Quickstart: a two-level nested copula\n",
9+
"\n",
10+
"This notebook builds a **nested Archimedean copula**, evaluates its\n",
11+
"log-likelihood, and takes an exact gradient with `jax.grad` — the whole\n",
12+
"pipeline in a few cells.\n",
13+
"\n",
14+
"The model is a Frank copula over two Frank \"sectors\", each with two uniform\n",
15+
"leaves (4 dimensions total):\n",
16+
"\n",
17+
"```\n",
18+
" Frank (outer)\n",
19+
" / \\\n",
20+
" Frank Frank (sectors)\n",
21+
" / \\ / \\\n",
22+
"u00 u01 u10 u11 (uniform leaves)\n",
23+
"```"
24+
]
25+
},
26+
{
27+
"cell_type": "markdown",
28+
"id": "eca84c8d",
29+
"metadata": {},
30+
"source": [
31+
"## 1. Define a copula family\n",
32+
"\n",
33+
"A family subclasses `Copula` via the `@copula` decorator and implements\n",
34+
"`generator` — the Archimedean generator $\\varphi$. Everything else (density,\n",
35+
"inverse, sampling) is derived automatically."
36+
]
37+
},
38+
{
39+
"cell_type": "code",
40+
"execution_count": 1,
41+
"id": "904083af",
42+
"metadata": {
43+
"execution": {
44+
"iopub.execute_input": "2026-05-24T21:36:28.908968Z",
45+
"iopub.status.busy": "2026-05-24T21:36:28.908835Z",
46+
"iopub.status.idle": "2026-05-24T21:36:30.033947Z",
47+
"shell.execute_reply": "2026-05-24T21:36:30.033144Z"
48+
}
49+
},
50+
"outputs": [],
51+
"source": [
52+
"import jax\n",
53+
"import jax.numpy as jnp\n",
54+
"from tensorflow_probability.substrates import jax as tfp\n",
55+
"\n",
56+
"from acopula import compile_model, copula, marginal\n",
57+
"\n",
58+
"\n",
59+
"@copula\n",
60+
"class Frank:\n",
61+
" theta: float\n",
62+
"\n",
63+
" def generator(self, t):\n",
64+
" return -jnp.log1p(jnp.expm1(-self.theta) * jnp.exp(-t)) / self.theta"
65+
]
66+
},
67+
{
68+
"cell_type": "markdown",
69+
"id": "7079393b",
70+
"metadata": {},
71+
"source": [
72+
"## 2. Describe the model\n",
73+
"\n",
74+
"A model is a plain function `(params, obs) -> Node`. Calling a copula\n",
75+
"instance on its children builds the tree; `marginal` attaches a distribution\n",
76+
"(here a `Uniform`, since the data is already on the copula scale) and the\n",
77+
"observation index it reads."
78+
]
79+
},
80+
{
81+
"cell_type": "code",
82+
"execution_count": 2,
83+
"id": "f250721b",
84+
"metadata": {
85+
"execution": {
86+
"iopub.execute_input": "2026-05-24T21:36:30.035492Z",
87+
"iopub.status.busy": "2026-05-24T21:36:30.035290Z",
88+
"iopub.status.idle": "2026-05-24T21:36:30.038167Z",
89+
"shell.execute_reply": "2026-05-24T21:36:30.037623Z"
90+
}
91+
},
92+
"outputs": [],
93+
"source": [
94+
"def model(params, obs):\n",
95+
" outer = Frank(params[0])\n",
96+
" inner = Frank(params[1])\n",
97+
" return outer(\n",
98+
" inner(marginal(tfp.distributions.Uniform(0.0, 1.0), obs=obs[i, j])\n",
99+
" for j in range(2))\n",
100+
" for i in range(2)\n",
101+
" )"
102+
]
103+
},
104+
{
105+
"cell_type": "markdown",
106+
"id": "faf17f24",
107+
"metadata": {},
108+
"source": [
109+
"## 3. Compile it\n",
110+
"\n",
111+
"`compile_model` traces the function once into a copula tree, builds a\n",
112+
"parameter flattener, and returns a jit-compiled, `grad`/`vmap`-friendly\n",
113+
"evaluator. The `template` only conveys the parameter *structure* — its values\n",
114+
"are placeholders."
115+
]
116+
},
117+
{
118+
"cell_type": "code",
119+
"execution_count": 3,
120+
"id": "14f8ba3f",
121+
"metadata": {
122+
"execution": {
123+
"iopub.execute_input": "2026-05-24T21:36:30.039277Z",
124+
"iopub.status.busy": "2026-05-24T21:36:30.039160Z",
125+
"iopub.status.idle": "2026-05-24T21:36:30.120701Z",
126+
"shell.execute_reply": "2026-05-24T21:36:30.120007Z"
127+
}
128+
},
129+
"outputs": [],
130+
"source": [
131+
"cm = compile_model(model, template=jnp.array([1.0, 1.0]))"
132+
]
133+
},
134+
{
135+
"cell_type": "markdown",
136+
"id": "33c6032d",
137+
"metadata": {},
138+
"source": [
139+
"## 4. Evaluate the log-likelihood"
140+
]
141+
},
142+
{
143+
"cell_type": "code",
144+
"execution_count": 4,
145+
"id": "a612d689",
146+
"metadata": {
147+
"execution": {
148+
"iopub.execute_input": "2026-05-24T21:36:30.122666Z",
149+
"iopub.status.busy": "2026-05-24T21:36:30.122536Z",
150+
"iopub.status.idle": "2026-05-24T21:36:31.088230Z",
151+
"shell.execute_reply": "2026-05-24T21:36:31.087537Z"
152+
}
153+
},
154+
"outputs": [
155+
{
156+
"name": "stdout",
157+
"output_type": "stream",
158+
"text": [
159+
"log-likelihood: -0.958893\n"
160+
]
161+
}
162+
],
163+
"source": [
164+
"obs = jnp.array([[0.3, 0.7],\n",
165+
" [0.4, 0.8]])\n",
166+
"params = jnp.array([2.0, 5.0]) # theta_outer, theta_inner\n",
167+
"\n",
168+
"ll = cm.eval(obs, params)\n",
169+
"print(f\"log-likelihood: {ll:.6f}\")"
170+
]
171+
},
172+
{
173+
"cell_type": "markdown",
174+
"id": "61dd0233",
175+
"metadata": {},
176+
"source": [
177+
"## 5. Differentiate it\n",
178+
"\n",
179+
"The log-likelihood is an ordinary JAX function, so its exact gradient w.r.t.\n",
180+
"the parameters is one `jax.grad` away — this is what makes gradient-based MLE\n",
181+
"(the [next notebook](02_fit_mle.ipynb)) trivial."
182+
]
183+
},
184+
{
185+
"cell_type": "code",
186+
"execution_count": 5,
187+
"id": "7d7a095b",
188+
"metadata": {
189+
"execution": {
190+
"iopub.execute_input": "2026-05-24T21:36:31.091797Z",
191+
"iopub.status.busy": "2026-05-24T21:36:31.091592Z",
192+
"iopub.status.idle": "2026-05-24T21:36:33.198063Z",
193+
"shell.execute_reply": "2026-05-24T21:36:33.197304Z"
194+
}
195+
},
196+
"outputs": [
197+
{
198+
"name": "stdout",
199+
"output_type": "stream",
200+
"text": [
201+
"d(ll)/d(theta_outer): 0.109353\n",
202+
"d(ll)/d(theta_inner): -0.356627\n"
203+
]
204+
}
205+
],
206+
"source": [
207+
"grad = jax.grad(cm.eval, argnums=1)(obs, params)\n",
208+
"print(f\"d(ll)/d(theta_outer): {grad[0]:.6f}\")\n",
209+
"print(f\"d(ll)/d(theta_inner): {grad[1]:.6f}\")"
210+
]
211+
}
212+
],
213+
"metadata": {
214+
"jupytext": {
215+
"cell_metadata_filter": "-all",
216+
"main_language": "python",
217+
"notebook_metadata_filter": "-all"
218+
},
219+
"language_info": {
220+
"codemirror_mode": {
221+
"name": "ipython",
222+
"version": 3
223+
},
224+
"file_extension": ".py",
225+
"mimetype": "text/x-python",
226+
"name": "python",
227+
"nbconvert_exporter": "python",
228+
"pygments_lexer": "ipython3",
229+
"version": "3.12.4"
230+
}
231+
},
232+
"nbformat": 4,
233+
"nbformat_minor": 5
234+
}

0 commit comments

Comments
 (0)