Log-space Taylor override hook + API-docs reorg

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

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)

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]] = []

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

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
 

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:

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}"