Skip to content

Commit b45b728

Browse files
thisiscamclaude
andcommitted
README: replace quickstart with verified-working code
The previous quickstart was paraphrased from the paper's Listing 1 and contained several bugs: undefined Weibull(), non-existent module-level `model` reference, and jax.grad(ll) where ll was a scalar value rather than a function. None of it ran. Replace with two runnable blocks: (1) Two-level nested Frank with Uniform leaves. Demonstrates @copula, @defmodel, marginal, log_likelihood, and parameter gradients via jax.grad over a closure. Verified end-to-end with the actual installed package. (2) Same model with Weibull marginals and a per-leaf censoring flag, showing the survival-data pattern. Includes the float64 hint needed because tfp distributions default to float32 but acopula enables jax_enable_x64 at import time. Drop the "Conditional sampling" feature claim — sampling is currently broken in the installed code (oryx Talbot evaluation feeds complex inputs through real-valued grad). To be re-added once the underlying path is fixed. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent abce343 commit b45b728

1 file changed

Lines changed: 48 additions & 28 deletions

File tree

‎README.md‎

Lines changed: 48 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -27,46 +27,68 @@ inversion. At import time we register a missing ILDJ rule for the
2727

2828
## Quickstart
2929

30+
A two-level nested Archimedean copula with Frank inner and outer, four
31+
Uniform leaves, and gradient-based MLE in twenty lines:
32+
3033
```python
3134
import jax
3235
import jax.numpy as jnp
36+
from tensorflow_probability.substrates import jax as tfp
3337
from acopula import copula, defmodel, marginal
3438

3539
@copula
36-
class Clayton:
40+
class Frank:
3741
theta: float
3842
def generator(self, u):
39-
return (1 + u)**(-1.0 / self.theta)
43+
return -jnp.log1p(jnp.expm1(-self.theta) * jnp.exp(-u)) / self.theta
4044

41-
@copula
42-
class AMH:
43-
theta: float
44-
def generator(self, u):
45-
return (1 - self.theta) / (jnp.exp(u) - self.theta)
45+
@defmodel
46+
def model(params, obs):
47+
outer = Frank(params[0])
48+
inner = Frank(params[1])
49+
return outer(
50+
inner(marginal(tfp.distributions.Uniform(0.0, 1.0), obs=obs[i, j])
51+
for j in range(2))
52+
for i in range(2))
53+
54+
obs = jnp.array([[0.3, 0.7],
55+
[0.4, 0.8]])
4656

57+
model.set_params(jnp.array([2.0, 5.0]))
58+
print(model.log_likelihood(obs)) # scalar log-likelihood
59+
60+
def neg_ll(params, obs):
61+
model.set_params(params)
62+
return -model.log_likelihood(obs)
63+
64+
grad = jax.grad(neg_ll)(jnp.array([2.0, 5.0]), obs)
65+
print(grad) # ∂(-ll)/∂params, exact
66+
```
67+
68+
### With Weibull marginals and right-censoring
69+
70+
For survival data, swap the leaf distribution and mark censored observations.
71+
Use `float64` for the distribution parameters since `acopula` enables
72+
`jax_enable_x64` at import time.
73+
74+
```python
4775
@defmodel
48-
def nested_amh_clayton(params, obs):
49-
outer = AMH(params[0])
50-
inner = Clayton(params[1])
76+
def survival_model(params, obs):
77+
outer = Frank(params[0])
78+
inner = Frank(params[1])
79+
weib = tfp.distributions.Weibull(
80+
concentration=jnp.float64(1.5), scale=jnp.float64(1.0))
5181
return outer(
52-
inner(marginal(Weibull(),
53-
obs=obs[i, j],
54-
censored=((i, j) == (1, 3)))
55-
for j in range(5))
56-
for i in range(4))
57-
58-
model = nested_amh_clayton
59-
model.set_params(jnp.array([0.5, 2.0]))
60-
ll = model.log_likelihood(observations)
61-
gradient = jax.grad(ll)(model.params)
82+
inner(marginal(weib, obs=obs[i, j], censored=((i, j) == (1, 1)))
83+
for j in range(2))
84+
for i in range(2))
6285
```
6386

64-
The `@copula` decorator registers parameters and derives the generator inverse
65-
symbolically (falling back to bisection with implicit-function-theorem
66-
gradients when symbolic inversion fails). The `@defmodel` decorator traces the
67-
function into a copula tree, flattens parameters into a single array, and
68-
exposes `log_likelihood`, `sample`, and `cdf` — all `jit`/`grad`-compatible.
69-
The `marginal` primitive pairs each leaf with a distribution and an optional
87+
`@copula` registers the parameters declared as type-annotated class
88+
attributes and derives the inverse generator symbolically via `oryx`.
89+
`@defmodel` traces the function into a copula tree, flattens parameters
90+
into a single array, and exposes `log_likelihood`, all `jit`/`grad`-compatible.
91+
`marginal` pairs each leaf with a distribution and an optional
7092
per-observation censoring flag.
7193

7294
## Features
@@ -77,8 +99,6 @@ per-observation censoring flag.
7799
- **Density via Bell polynomials**, computed from a Taylor expansion of the
78100
generator rather than nested first-order AD.
79101
- **Symbolic generator inversion** via `oryx`, with bisection + IFT fallback.
80-
- **Conditional sampling** of nested Archimedean models, made practical by
81-
the same higher-order derivatives.
82102
- **Validity diagnostic** — per-edge `d_c`-monotonicity check for
83103
cross-family nesting.
84104

0 commit comments

Comments
 (0)