From 8d4f23811347c4220850dcdc769edeb47f26ecea Mon Sep 17 00:00:00 2001 From: Cambridge Yang Date: Sun, 24 May 2026 12:58:46 -0400 Subject: [PATCH 1/2] core: add log-space generator_taylor_coefficients override hook MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds `Copula.log_generator_taylor_coefficients(t, k_max)` returning `(sign, log|c|)` for the normalized Taylor coefficients c_k = ψ^(k)(t)/k!. The Bell density already consumes coefficients as (sign, log|c|) internally, so a family whose closed-form series is computed in the log domain (e.g. AMH's Geometric-frailty logsumexp) can feed them straight in — instead of exponentiating to a possibly-overflowing float64 only for bell to take the log again. bell._root_assembly now prefers this hook, then the linear generator_taylor_coefficients, then jet. Verified on AMH (flat copula, θ=0.5): the log override reproduces the linear override's log-likelihood to machine precision across d=8..50, while plain jet drifts at d=30 and is wrong at d=50 — i.e. the log hook matches the paper's AMH precision result. Regression test in test_log_taylor_override.py. --- acopula/bell.py | 25 ++++++++---- acopula/core.py | 36 +++++++++++++++++ tests/test_log_taylor_override.py | 66 +++++++++++++++++++++++++++++++ 3 files changed, 120 insertions(+), 7 deletions(-) create mode 100644 tests/test_log_taylor_override.py diff --git a/acopula/bell.py b/acopula/bell.py index 2d62be9..21a83db 100644 --- a/acopula/bell.py +++ b/acopula/bell.py @@ -958,13 +958,24 @@ def _root_assembly( # Per-family Taylor override: families whose ψ has a numerically # stable series form (e.g. Laplace transforms of nonneg distributions - # via the frailty representation) can provide - # ``generator_taylor_coefficients(t, k_max)`` to bypass jet entirely. - # This avoids the chain-rule cancellations Taylor-mode AD on ψ's - # closed form produces at high derivative order — relevant for AMH - # past d≈30 in float64. Returning None means "use jet". - custom_taylor = root_cop.generator_taylor_coefficients(t_r, d) - if custom_taylor is not None: + # via the frailty representation) can provide a closed-form Taylor + # expansion to bypass jet entirely. This avoids the chain-rule + # cancellations Taylor-mode AD on ψ's closed form produces at high + # derivative order — relevant for AMH past d≈30 in float64. + # Preference order: log-space hook -> linear hook -> jet. The Bell + # density consumes the coefficients as (sign, log|c|) anyway, so the + # log-space hook feeds them straight in with no overflow-prone round + # trip through a linear float64 value. + custom_log_taylor = root_cop.log_generator_taylor_coefficients(t_r, d) + custom_taylor = (None if custom_log_taylor is not None + else root_cop.generator_taylor_coefficients(t_r, d)) + if custom_log_taylor is not None: + # custom_log_taylor = (sign, log|c|), each shape (d+1,); already in + # the representation the downstream signed-LSE wants. + twp_sign, twp_log_abs = custom_log_taylor + twp_sign = jnp.asarray(twp_sign) + twp_log_abs = jnp.asarray(twp_log_abs) + elif custom_taylor is not None: # custom_taylor[k] = psi^{(k)}(t_r) / k!, shape (d+1,). twp_sign = jnp.sign(custom_taylor) twp_abs = jnp.abs(custom_taylor) diff --git a/acopula/core.py b/acopula/core.py index f6e2694..cbed27c 100644 --- a/acopula/core.py +++ b/acopula/core.py @@ -49,6 +49,12 @@ class Copula: """ def generator(self, u: jax.Array) -> jax.Array: # pragma: no cover - abstract + """Return the Archimedean generator φ(u). + + Subclasses **must** implement this. φ maps (0, 1] → [0, ∞), is + strictly decreasing and convex, with φ(1) = 0. Everything else + (density, inverse, sampling) is derived from it automatically. + """ raise NotImplementedError # Inverse of generator: φ^{-1}(t) @@ -102,6 +108,36 @@ def generator_taylor_coefficients( """ return None # sentinel: use jet + def log_generator_taylor_coefficients( + self, t: jax.Array, k_max: int + ) -> Optional[Tuple[jax.Array, jax.Array]]: + """Log-space variant of :meth:`generator_taylor_coefficients`. + + Return ``(sign, log_abs)``, two arrays of shape ``(k_max + 1,)`` with + + ψ^{(k)}(t) / k! == sign[k] * exp(log_abs[k]). + + The Bell density consumes the Taylor coefficients as ``(sign, log|c|)`` + pairs internally, so a family whose normalized coefficients are + computed in the log domain (e.g. via ``logsumexp`` over a frailty + series) should return them here directly rather than through the linear + :meth:`generator_taylor_coefficients` — that avoids exponentiating to a + possibly-overflowing float64 only for the framework to take its log + again, extending the usable derivative order at high ``d``. + + The Bell path prefers this hook when it returns non-``None``, then falls + back to the linear hook, then to jet. Returning ``None`` (the default) + selects that fallback chain. + + Args: + t: scalar input point. + k_max: highest derivative order (inclusive). + + Returns: + ``(sign, log_abs)`` arrays of shape ``(k_max + 1,)``, or ``None``. + """ + return None # sentinel: fall back to the linear hook / jet + # For tracing: combine children (nodes/leaves) into a Node def __call__(self, children: Iterable[Union["Node", "Leaf", jax.Array]]) -> "Node": child_nodes: List[Union[Node, Leaf]] = [] diff --git a/tests/test_log_taylor_override.py b/tests/test_log_taylor_override.py new file mode 100644 index 0000000..b21cc8b --- /dev/null +++ b/tests/test_log_taylor_override.py @@ -0,0 +1,66 @@ +"""The log-space generator_taylor_coefficients override must agree with the +linear override (and with jet where jet is still accurate). + +Regression test for ``Copula.log_generator_taylor_coefficients`` and the Bell +dispatch that prefers it. Uses AMH, whose frailty series is computed in the log +domain anyway, so the linear override is just the log override exponentiated. +""" +import jax +import jax.numpy as jnp +import pytest + +jax.config.update("jax_enable_x64", True) + +from acopula import compile_model, copula, marginal + + +class _Uniform: + def quantile(self, u): return u + def cdf(self, x): return x + def log_prob(self, x): return 0.0 + + +def _amh_sign_logabs(theta, t, k_max, n_terms=300): + ks = jnp.arange(k_max + 1) + xs = jnp.arange(1, n_terms + 1, dtype=jnp.float64) + log_w = jnp.log(1.0 - theta) + (xs - 1.0) * jnp.log(theta) - t * xs + log_fact_k = jax.scipy.special.gammaln(ks + 1) + log_term = (ks[:, None] * jnp.log(xs)[None, :] + log_w[None, :] + - log_fact_k[:, None]) + log_abs_c = jax.scipy.special.logsumexp(log_term, axis=1) + signs = jnp.where(ks % 2 == 0, 1.0, -1.0) + return signs, log_abs_c + + +@copula +class _AMHLinear: + theta: float + def generator(self, t): + return (1.0 - self.theta) / (jnp.exp(t) - self.theta) + def generator_taylor_coefficients(self, t, k_max, n_terms=300): + signs, log_abs = _amh_sign_logabs(self.theta, t, k_max, n_terms) + return signs * jnp.exp(log_abs) + + +@copula +class _AMHLog: + theta: float + def generator(self, t): + return (1.0 - self.theta) / (jnp.exp(t) - self.theta) + def log_generator_taylor_coefficients(self, t, k_max, n_terms=300): + return _amh_sign_logabs(self.theta, t, k_max, n_terms) + + +def _ll(cls, d, theta_val): + def model(p, u): + return cls(p[0])([marginal(_Uniform(), obs=u[i]) for i in range(d)]) + cm = compile_model(model, template=jnp.array([theta_val]), method="bell") + u = jnp.arange(1, d + 1, dtype=jnp.float64) / (d + 1) + return float(cm.eval(u, jnp.array([theta_val]))) + + +@pytest.mark.parametrize("d", [8, 12]) +def test_log_override_matches_linear(d): + lin = _ll(_AMHLinear, d, 0.5) + log = _ll(_AMHLog, d, 0.5) + assert abs(lin - log) < 1e-9, f"d={d}: linear={lin}, log={log}" From 6b67dc8c74c1d6c4c5dea98ff6807f8422f6c1e2 Mon Sep 17 00:00:00 2001 From: Cambridge Yang Date: Sun, 24 May 2026 13:01:39 -0400 Subject: [PATCH 2/2] docs: organize API reference; correct install instructions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - mkdocstrings: hide private/dunder members, Args/Returns as tables, symbol-type headings + TOC, no source dumps, drop merge_init (leaked CompiledModel private dataclass fields). API page renders in clean curated sections. - api.md: restructure into Defining a copula / Compiling & evaluating / Composition registry / Configuration / Advanced (generator-precision override hooks grouped + explained). - index.md + README.md: fix install — pip install (incl. git+) does not work (ignores the [tool.uv] oryx pin); document uv-from-source; PyPI gated on an upstream oryx release for jax 0.8. --- README.md | 16 ++++++++++++-- docs/api.md | 59 +++++++++++++++++++++++++++++++++++++++++++++++++-- docs/index.md | 19 +++++++++++++---- mkdocs.yml | 13 ++++++++++-- 4 files changed, 97 insertions(+), 10 deletions(-) diff --git a/README.md b/README.md index 8a5fb3b..8d66bf5 100644 --- a/README.md +++ b/README.md @@ -17,11 +17,23 @@ AD* (arXiv:TBD). ## Install +Install from source with [uv](https://docs.astral.sh/uv/): + ```bash -pip install acopula +git clone https://github.com/thisiscam/acopula +cd acopula +uv sync # add --extra examples for the plotting examples +uv run python examples/01_quickstart.py ``` -`acopula` currently pins `jax>=0.8,<0.9` because it relies on JAX-internal +uv is required (for now) because `acopula` depends on a patched +[`oryx`](https://github.com/jax-ml/oryx) build pinned via `[tool.uv.sources]`; +plain `pip install` — including `pip install git+…` — ignores that pin, pulls +the stock PyPI `oryx` (incompatible with jax 0.8), and fails at import. A +`pip install acopula` from PyPI is **not available yet**; it is gated on an +upstream `oryx` release compatible with jax 0.8. + +`acopula` pins `jax>=0.8,<0.9` because it relies on JAX-internal APIs in the [jet-array](https://github.com/thisiscam/jet-array) backend. `acopula` uses [`oryx`](https://github.com/jax-ml/oryx) for symbolic generator diff --git a/docs/api.md b/docs/api.md index a42c83a..d26fbda 100644 --- a/docs/api.md +++ b/docs/api.md @@ -1,12 +1,21 @@ # API reference -The public surface of `acopula`. +Generated from the docstrings. The everyday surface is small: define a copula +family, compile a model, then evaluate or sample. -## Defining copulas +## Defining a copula + +A family subclasses `Copula` — usually via the `@copula` decorator for +dataclass-style parameters — and implements `generator`. Everything else +(density, inverse, sampling) is derived automatically. Call a copula instance +on its children to build the model tree. ::: acopula.copula ::: acopula.Copula + options: + members: + - generator ::: acopula.marginal @@ -15,9 +24,21 @@ The public surface of `acopula`. ::: acopula.compile_model ::: acopula.CompiledModel + options: + members: + - eval + - sample + - flatten + - ll_fn + - visualize + - as_networkx ## Composition registry +For nested copulas that mix families, register a closed-form composition +`ψ_outer⁻¹ ∘ ψ_inner`, or choose the fallback strategy used when none is +registered. + ::: acopula.register_composition ::: acopula.set_composition_fallback @@ -27,3 +48,37 @@ The public surface of `acopula`. ::: acopula.set_stable_log ::: acopula.set_compile_cache_dir + +## Advanced: custom generator precision + +Most families never touch these — they are optional override hooks on `Copula` +for families whose high-order derivatives suffer float64 cancellation (e.g. AMH +past d≈30). Each returns `None` by default, which selects Taylor-mode AD +(`jet`). `generator_taylor_coefficients` (and its log-space variant +`log_generator_taylor_coefficients`) is the recommended, general hook, used by +the Bell density and the nesting composition. `log_generator_kth_derivative` is +a legacy hook consulted only by the flat single-layer likelihood path. + +::: acopula.Copula.generator_inv + options: + show_root_heading: true + show_root_full_path: false + heading_level: 3 + +::: acopula.Copula.generator_taylor_coefficients + options: + show_root_heading: true + show_root_full_path: false + heading_level: 3 + +::: acopula.Copula.log_generator_taylor_coefficients + options: + show_root_heading: true + show_root_full_path: false + heading_level: 3 + +::: acopula.Copula.log_generator_kth_derivative + options: + show_root_heading: true + show_root_full_path: false + heading_level: 3 diff --git a/docs/index.md b/docs/index.md index 37c8b5f..ff6807b 100644 --- a/docs/index.md +++ b/docs/index.md @@ -11,13 +11,24 @@ forward pass via Taylor-mode AD, scaling polynomially in the dimension. ## Install +Install from source with [uv](https://docs.astral.sh/uv/): + ```bash -pip install acopula +git clone https://github.com/thisiscam/acopula +cd acopula +uv sync # add --extra examples for the plotting examples +uv run python examples/01_quickstart.py ``` -`acopula` pins `jax>=0.8,<0.9` (it relies on JAX-internal APIs in the -[jet-array](https://github.com/thisiscam/jet-array) backend) and uses -[`oryx`](https://github.com/jax-ml/oryx) for symbolic generator inversion. +uv is required (for now): `acopula` depends on a patched +[`oryx`](https://github.com/jax-ml/oryx) build pinned via `[tool.uv.sources]`, +and plain `pip install` (including `pip install git+…`) ignores that pin, pulls +the stock PyPI `oryx` (incompatible with jax 0.8), and fails at import. A PyPI +`pip install acopula` is **not available yet** — it is gated on an upstream +`oryx` release compatible with jax 0.8. + +`acopula` pins `jax>=0.8,<0.9` because it relies on JAX-internal APIs in the +[jet-array](https://github.com/thisiscam/jet-array) backend. ## Features diff --git a/mkdocs.yml b/mkdocs.yml index 6c65061..4098da7 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -44,11 +44,20 @@ plugins: python: options: docstring_style: google + docstring_section_style: table # Args/Returns as tables show_root_heading: true - show_source: true - members_order: source + show_root_full_path: false # "Copula", not "acopula.core.Copula" + show_symbol_type_heading: true # class/func/attr badges + show_symbol_type_toc: true # and in the nav TOC separate_signature: true show_signature_annotations: true + signature_crossrefs: true + merge_init_into_class: false # CompiledModel isn't user-constructed; its __init__ leaks private fields + members_order: source + filters: ["!^_"] # hide private / dunder members + show_source: false + summary: true + heading_level: 2 # Emits /llms.txt (an index for LLMs/agents, per llmstxt.org) plus a single # concatenated llms-full.txt of the whole site. - llmstxt: