Merge pull request #2 from thisiscam/examples-and-docs

Curate examples + MkDocs docs site (with llms.txt)

Author: thisiscam

.github/workflows/docs.yml

@@ -0,0 +1,43 @@
+# Build the MkDocs site and deploy to GitHub Pages.
+# Enable Pages once in repo Settings → Pages → Source: "GitHub Actions".
+name: docs
+
+on:
+  push:
+    branches: [main]
+  workflow_dispatch:
+
+permissions:
+  contents: read
+  pages: write
+  id-token: write
+
+concurrency:
+  group: pages
+  cancel-in-progress: false
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v5
+        with:
+          python-version: "3.12"
+          enable-cache: true
+      # mkdocstrings imports acopula to introspect, so install the full env.
+      - name: Build site
+        run: uv run --extra docs mkdocs build --strict
+      - uses: actions/upload-pages-artifact@v3
+        with:
+          path: site
+
+  deploy:
+    needs: build
+    runs-on: ubuntu-latest
+    environment:
+      name: github-pages
+      url: ${{ steps.deployment.outputs.page_url }}
+    steps:
+      - id: deployment
+        uses: actions/deploy-pages@v4

.gitignore

@@ -18,3 +18,6 @@ env/
 .idea/
 .vscode/
 *.swp
+
+# MkDocs build output
+site/

README.md

@@ -12,6 +12,9 @@ forward pass via Taylor-mode AD, scaling polynomially in the dimension.
 For background, see the paper *Archimedean Copula Inference via Taylor-Mode
 AD* (arXiv:TBD).
 
+**Documentation:** <https://thisiscam.github.io/acopula/> ·
+**Examples:** [`examples/`](examples/)
+
 ## Install
 
 ```bash

acopula/compile.py

@@ -323,14 +323,13 @@ def visualize(
         annotations=None,
         ax=None,
     ):
-        from .visualize import draw
-        return draw(
-            self.graph,
-            include_leaves=include_leaves,
+        from .visualize import draw_networkx
+        G = self.as_networkx(include_leaves=include_leaves, annotations=annotations)
+        return draw_networkx(
+            G,
             layout=layout,
             with_labels=with_labels,
             graphviz_prog=graphviz_prog,
-            annotations=annotations,
             ax=ax,
         )
 

acopula/compose.py

@@ -28,7 +28,7 @@
 _FALLBACK = "implicit"  # "implicit" (default, stable, O(d_c^3)) or "direct_jet" (fast, may overflow)
 
 
-def register_composition(outer_cls, inner_cls, fn):
+def register_composition(outer_cls: type, inner_cls: type, fn: Callable):
     """Register a stable closed-form composition h(t) = psi_outer_inv(psi_inner(t)).
 
     Args:

acopula/core.py

@@ -1847,16 +1847,11 @@ def _full_mixed_partials(
     #     = F^{(k)}(S) * prod(|g'|) = Likelihood!
     T_k = _kth_derivative(outer_fun, S, k, scale=scaling_factor)
 
-    # 5. Assemble log-likelihood (mixed partial)
-    log_derivative_term = jnp.log(T_k)
-    jax.debug.print(
-        "log_derivative_term: {} T_k: {} S: {} scaling_factor: {} g_primes: {}",
-        log_derivative_term,
-        T_k,
-        S,
-        scaling_factor,
-        g_primes,
-    )
+    # 5. Assemble log-likelihood (mixed partial).  The mixed partial of an
+    # Archimedean copula is (-1)^k psi^{(k)}; take the magnitude so the log is
+    # well-defined for odd k as well (e.g. a single observed dimension under
+    # censoring), where the raw k-th derivative is negative.
+    log_derivative_term = jnp.log(jnp.abs(T_k))
 
     return log_derivative_term
 

docs/api.md

@@ -0,0 +1,29 @@
+# API reference
+
+The public surface of `acopula`.
+
+## Defining copulas
+
+::: acopula.copula
+
+::: acopula.Copula
+
+::: acopula.marginal
+
+## Compiling and evaluating
+
+::: acopula.compile_model
+
+::: acopula.CompiledModel
+
+## Composition registry
+
+::: acopula.register_composition
+
+::: acopula.set_composition_fallback
+
+## Configuration
+
+::: acopula.set_stable_log
+
+::: acopula.set_compile_cache_dir

docs/assets/acopula_pairs.png

@@ -0,0 +1,59 @@
+# Examples
+
+The [`examples/`](https://github.com/thisiscam/acopula/tree/main/examples)
+directory has runnable, self-contained scripts. Each uses small synthetic /
+copula-scale data — no external datasets required.
+
+Run them from the repo root:
+
+```bash
+uv run python examples/01_quickstart.py
+uv run python examples/02_fit_mle.py
+uv run python examples/03_censored_survival.py
+uv run --extra examples python examples/04_sampling_and_visualize.py [output_dir]
+```
+
+## 01 — Quickstart
+
+Build a two-level nested copula, evaluate the log-likelihood and its exact
+`jax.grad`. See [Quickstart](quickstart.md).
+
+## 02 — Maximum-likelihood fitting
+
+Simulate 1000 observations from a known `d=10` Clayton copula (two sectors of
+five), then recover the parameters by gradient-descent MLE with `optax`. The
+whole likelihood is differentiable, so the fit is plain
+`jax.value_and_grad` + Adam:
+
+```
+parameter         true     fitted
+  theta_root      2.00      1.958
+  theta_sector    5.00      4.840
+```
+
+## 03 — Censored survival
+
+A bivariate Frank copula over two Weibull event times, where the second is
+right-censored. A censored leaf still enters the copula argument (through its
+survival function) but is not differentiated in the density — `acopula`
+assembles the correct mixed partial over only the observed dimensions.
+
+## 04 — Sampling and visualization
+
+Draw samples from a Clayton-over-Frank copula (one root, three sectors of two
+leaves) with the Marshall-Olkin sampler, and render the copula tree.
+
+### Copula tree
+
+`CompiledModel.visualize()` draws the structural graph — root and sector
+copulas in blue, leaves in orange:
+
+![Copula tree](assets/acopula_tree.png)
+
+### Sampled margins
+
+Pairwise scatter of the six sampled margins. Within-sector pairs
+(`u[i,0]`–`u[i,1]`) are visibly more dependent than across-sector pairs,
+reflecting the nesting:
+
+![Pairwise margins](assets/acopula_pairs.png)