emu-gmm 0.5.1

#+TITLE: emu-gmm
#+SUBTITLE: Measure-theoretic GMM; estimation via $\mathbb{E}_\mu$
#+AUTHOR: Ethan Ligon
#+OPTIONS: toc:nil num:nil

[[https://doi.org/10.5281/zenodo.21629016][https://img.shields.io/badge/DOI-10.5281%2Fzenodo.21629016-blue.svg]]

=emu-gmm= is a JAX-native framework for Generalized Method of Moments estimation. The framework is named for the operator at its centre: =emu= reads as $\mathbb{E}_\mu$, the expectation under a measure $\mu$.  One operator interface --- implemented against empirical, analytical, or synthetic measures --- drives sample estimation, identification analysis, and simulation-based inference through a single computational pipeline.

The most important architectural commitment is that variance construction is orthogonal to integration: $\mathbb{E}_\mu[\psi]$ is a property of the measure; $V_\mu(\theta)$ is a property of a separate =CovarianceStrategy= that can be swapped (iid vs cluster-robust vs replicate-weight) without changing what is being integrated.

* Status

v2: the v1 pipeline plus Riemannian-manifold parameter geometry and the stratified / design-aware covariance ladder. The implemented menu --- everything importable and composable today --- is tabulated in =docs/design.org= under "Contract surface (implemented)", the single source of truth for status claims; the green gate is =make check= (ruff + black + mypy + the full pytest suite). The three measure paths (synthetic, analytical, empirical) all run end-to-end against the bundled multi-asset Euler example, and the empirical path additionally has a real-data acceptance test (=tests/test_estimator_realdata.py=). The architecture and theoretical scope are pinned in =docs/=.

* Quickstart

Bootstrap and try the worked example:

#+begin_src shell
make setup                                 # poetry install + .venv/
poetry run python examples/run_euler.py    # multi-asset Hansen-Singleton demo
#+end_src

The minimal estimation surface in code (synthetic-data variant):

#+begin_src python
import jax
from emu_gmm import estimate, SyntheticMeasure, SyntheticCovariance
from emu_gmm.examples.euler import (
    EulerParams, euler_residual, euler_sampler_factory,
)

measure = SyntheticMeasure(
    key=jax.random.PRNGKey(0),
    n_sim=5000,
    sampler=euler_sampler_factory(5000),
)

result = estimate(
    model=euler_residual,
    measure=measure,
    covariance=SyntheticCovariance(),
    theta_init=EulerParams(beta=0.9, gamma=1.0),
)

print(f"beta = {result.theta_hat.beta:.4f}  (truth 0.96)")
print(f"gamma = {result.theta_hat.gamma:.4f}  (truth 2.00)")
law = result.asymptotic()
print(f"J-stat = {result.objective_value:.3f}  dof={result.n_overid}  p={law.J_pvalue:.3f}")
print(law.cov())  # asymptotic parameter covariance
#+end_src

The same =euler_residual= drives the analytical and empirical variants of the demo; only the =Measure= / =CovarianceStrategy= change. See =docs/api-sketch.org= Section 5 for the three side-by-side, and =examples/run_euler.py= for the runnable script.

* Sample reductions

On =EmpiricalMeasure=, a model may supply =aux(measure, theta)= and =psi(row, theta, aux)= instead of a row callable.
The package recomputes aux at each parameter value inside differentiation and shares the resulting residuals across moments and covariance.
The model declares the reduction's population and handles excluded inputs and empty denominators.
Use one =build_estimator= factory for fresh same-structure datasets; model configuration and covariance layouts stay fixed.
Reducing synthetic and analytical models require a future explicit adapter.
See [[file:docs/design.org::*Empirical sample reductions (#113)][the sample-reduction contract]] for derivative, labeling and inference limits.

* Inference

Estimation gives you =theta_hat=; inference gives you the uncertainty around it. =emu-gmm= hands you a menu, not a single number, and every tool reads off the same pipeline --- you never recompute a variance by hand (see "Using emu-gmm correctly" below).

| You want                              | Use                                       | Kind                         |
|---------------------------------------+-------------------------------------------+------------------------------|
| parameter SEs, fast                   | =result.asymptotic().se()= / =.cov()=     | analytic sandwich            |
| parameter SEs / CIs, robust           | =cluster_bootstrap= \(\to\) =theta_boot=  | refit: resample + re-solve   |
| the overidentification test \(J\)     | =moment_wild_bootstrap=                    | refit-free: perturb moments  |
| identification-robust inference       | =k_statistic= / =k_confidence_set=        | weak-id-robust               |
| gauge-invariant functionals of \(AA'\) | =result.asymptotic().functional_se(f)=   | analytic                     |
| =n_boot= chosen to a precision target | =adaptive_bootstrap=                      | wrapper over the above       |

*Two bootstraps, deliberately distinct.* =moment_wild_bootstrap= perturbs the moments and never re-solves --- cheap, and the right tool for the \(J\) distribution. =cluster_bootstrap= is the /refit/ (re-estimation) parameter bootstrap people usually mean: it resamples whole clusters with replacement, re-runs =estimate()= on each resample, and returns the bootstrap distribution =theta_boot= of =theta_hat= (Cameron, Gelbach & Miller 2008). Reach for it when the analytic sandwich is suspect --- few clusters, a boundary optimum, or a functional whose nonlinearity the first-order delta method can't capture.

The capability is easy to miss, so it's worth naming why it exists: a refit bootstrap is just =estimate()= on a resampled =Measure=. Because variance construction is orthogonal to integration (the commitment at the top of this README), resampling is a measure-construction concern, and the bootstrap is a thin re-run of the ordinary pipeline rather than a bespoke variance routine. That's also why there's no monolithic =bootstrap_se=: the capability falls out of the architecture.

This example groups ten synthetic observations per cluster to show the interface; use the actual sampling-cluster IDs for real data.

#+begin_src python
import numpy as np
import jax.numpy as jnp
from emu_gmm import ClusteredCovariance, EmpiricalMeasure
from emu_gmm.examples.euler import N_ASSETS, euler_data
from emu_gmm.inference import cluster_bootstrap

n_clusters, rows_per_cluster = 30, 10
n = n_clusters * rows_per_cluster
empirical = EmpiricalMeasure(
    x=euler_data(seed=0, n=n),
    mask=jnp.ones((n, N_ASSETS)),
    weights=jnp.ones(n),
)
clustered = ClusteredCovariance(
    cluster_ids=jnp.repeat(jnp.arange(n_clusters), rows_per_cluster),
    n_clusters=n_clusters,
)
empirical_result = estimate(
    model=euler_residual, measure=empirical, covariance=clustered,
    theta_init=result.theta_hat,
)
boot = cluster_bootstrap(
    model=euler_residual, theta_init=empirical_result.theta_hat, measure=empirical,
    covariance=clustered, n_boot=99, key=jax.random.PRNGKey(1),
)
ok = np.asarray(boot.convergence)        # non-converged refits are NaN rows (#209)
se = np.asarray(boot.theta_boot.array)[ok].std(axis=0, ddof=1) if ok.sum() > 1 else None
#+end_src

A refit that does not converge --- a resampled world can extinguish a rare moment's support entirely --- is stored as a =NaN= row, flagged in =boot.convergence= and explained by =boot.zero_support= / =boot.tau_saturated=; =boot.n_nonconverged= counts them. Reduce over the converged rows as above; =se= is =None= with fewer than two valid draws. Alternatively, hand raw columns to =adaptive_bootstrap=, which excludes and counts =NaN= draws.

=docs/howto.org= has the expanded menu and the refit-vs-refit-free decision in full.
For =StratifiedCovariance(fpc=False)=, =stratified_bootstrap= resamples PSUs within each stratum-by-arm cell.
For mixed =DesignAwareCovariance(fixed_strata=True)=, =mixed_bootstrap= resamples complete balanced paired sampling bundles within strata, requiring =design.fpc=False= and =sampling.dof_correction=False=.
For default =fixed_strata=False= with sampling clusters and strata defining the same partition, it draws whole strata with replacement, preserving variable sizes, cell counts, rows, masks, weights and the sampling centering/dof flags.
This targets unconditional iid-stratum variation under the usual finite-moment, positive-support and regular-estimator assumptions; it need not match the conditional design =V_TT= when random common stratum profiles are omitted (#145), and does not repair analytic covariance, J chi-square calibration or general studentization.
Copied covariance IDs are fresh at every level; IDs inside arbitrary model-data columns are not rewritten, so model grouping reductions must respect their intended relabeling/population contract.
Both modes reject FPC: without-replacement resampling requires an explicit population/assignment law.
Finer or crossing sampling partitions and broader calibration remain tracked in issue #251.

* Documents

- =docs/howto.org= --- *start here.* Integration HOWTO: the architecture
  at the level you need to wire an application to =estimate()= --- what
  code you write (the moment function, the parameter container, the
  measure/covariance choice) and what the framework hands back.
- =docs/design.org= --- architecture specification (four review rounds; stable). Its "Contract surface (implemented)" section is the single source of truth for the implemented menu.
- =docs/api-sketch.org= --- the v1 API surface, retained as the v1 design record; superseded on status by =design.org='s contract-surface section.
- =docs/implementation-plan.org= --- phased task list (Phases 1-7 complete, Phase 8 polish underway; the v2 roadmap is Section 13).
- =docs/mcar-asymptotics.org= --- companion theoretical note; consistency, asymptotic normality, PD properties of the pairwise-overlap estimator under MCAR.
- =docs/refs.bib= --- project-local bibliography (entries not in =~/bibtex/main.bib=).

* Setup

Requires Python >= 3.11.

Install the published package from PyPI with pip:

#+begin_src shell
pip install emu-gmm
#+end_src

With uv, explicitly allow prereleases:

#+begin_src shell
uv pip install --prerelease=allow emu-gmm
#+end_src

The transitive =haliax>=1.4.dev452,<1.5= dependency deliberately uses a development release for compatibility with current JAX.
The flag explicitly permits that prerelease under [[https://docs.astral.sh/uv/concepts/resolution/#pre-release-handling][uv's resolution policy]].

** Development checkout

For a source checkout, install [[https://python-poetry.org/][Poetry]] and run:

#+begin_src shell
make setup
#+end_src

This runs =poetry install= and creates =.venv/= in the project root.
Activate with =direnv allow= or =poetry shell=.

* Development

#+begin_src shell
make check         # ruff + black + mypy + full pytest
make quick-check   # same, skipping slow tests
make test          # pytest only
#+end_src

To install pre-commit hooks (ruff + black on every commit):

#+begin_src shell
poetry run pre-commit install
#+end_src

** Releasing

Releases are cut locally and published from GitHub; no PyPI token is
stored anywhere (PyPI trusted publishing).

#+begin_src shell
make release BUMP=patch   # full check gate, clean-tree guard, version bump,
                          # commit, annotated tag vX.Y.Z, poetry build
make release-publish      # push commit + tag, create the GitHub Release
#+end_src

Publishing the GitHub Release triggers =.github/workflows/publish.yml=,
which rebuilds the sdist and wheel, checks the tag against =pyproject.toml=,
attaches them to the Release, and uploads to PyPI through the =pypi=
environment; the same Release event mints the Zenodo version DOI.  One-time
setup lives in the workflow's header comment (a trusted publisher on
pypi.org and a =pypi= environment on GitHub).  =make publish= (=poetry
publish=, token-based) remains the manual fallback.

* Layout

#+begin_example
docs/                      design specs and theoretical notes
src/emu_gmm/
    __init__.py            public API re-exports
    types.py               protocols + OptimizationResult / Diagnostics
    law.py                 AsymptoticLaw and inference readouts
    estimator.py           estimate() / build_estimator() entry points
    measures/              SyntheticMeasure, AnalyticalMeasure, EmpiricalMeasure
    covariance/            the CovarianceStrategy ladder (iid, clustered,
                           stratified / design-aware, sum, analytical, synthetic)
    weighting.py           Identity, Fixed, IteratedWeighting, ContinuouslyUpdated
    regularization.py      DiagonalTikhonov
    penalty.py             TikhonovPenalty
    optimizer.py           optimistix_lm, scipy_lm, linear_solver
    manifolds/             parameter geometry (Euclidean, Positive, PSDFixedRank,
                           Product, ManifoldLeaf) + riemannian_lm
    parameter_space.py     ParameterSpace / on field-to-manifold declarations
    inference/             j_test, k_statistic, k_confidence_set, bootstraps
    numerics/              ridge_inverse
    studies/               Monte Carlo driver (subpackage-only; not re-exported)
    diagnostics.py         build_diagnostics, log_to_stdout
    examples/              shared example models (euler.py)
    _internal/             axes, params, cholesky, labels (private)
tests/                     test suite mirroring src/emu_gmm/
examples/                  runnable demo scripts
Makefile                   build automation
pyproject.toml             Poetry configuration
poetry.lock                pinned dependencies (committed)
#+end_example

* Public API

Everything user-facing is re-exported at the package top level
(=from emu_gmm import ...=), with one deliberate exception: the Monte
Carlo =studies= subpackage is subpackage-only (import from
=emu_gmm.studies= directly). The implemented menu --- entry points,
measures, covariance strategies, weighting, regularization, penalty,
optimisers, manifold geometry, inference helpers, and result types ---
is tabulated, one line of semantics and a module pointer per row, in
=docs/design.org= under "Contract surface (implemented)". That table is
the single source of truth this README defers to;
=sorted(emu_gmm.__all__)= is the export list it tracks.

* Using emu-gmm correctly (and reporting problems)

=emu-gmm= is intended as a /spare set of correct interfaces/. Your only
modelling input is a per-observation residual =psi(x_i, theta) -> R^M=;
the package owns the moment expectation, the design-aware covariance
=V_X=, the criterion, the \(J\)-statistic, standard errors, and
\(p\)-values. Two rules keep estimation correct:

1. *Read results off the package, never recompute them by hand.* Take the
   criterion and \(J\)-statistic from =result.objective_value=, then
   standard errors and \(p\)-values from =result.asymptotic()= and the
   inference helpers. A
   hand-rolled criterion is how subtle scaling bugs get reintroduced ---
   notably the textbook habit of scaling moments by \(\sqrt{N}\) and
   weighting by \(V_X\), which is correct only when every moment has the
   /same/ number of observations. =emu-gmm= keeps moments as
   per-coordinate means and folds all the per-moment \(N_j\) bookkeeping
   into \(V_X\) (the \(1/(N_j N_k)\) normalisation); see =docs/design.org=
   "Scaling convention" and =docs/mcar-asymptotics.org=.

2. *Express per-moment observability only through the =mask=.* Different
   moments may be observed for different numbers of units; the \((N, M)\)
   =mask= on =EmpiricalMeasure= is the single place that lives. Do not
   work around it by pre-aggregating or by assuming a common \(N\).

If an interface is missing a knob, or you believe it gives a wrong
answer: *file an issue against =emu-gmm=, do not reimplement the
statistic in your own project.* A local reimplementation forfeits the
"single correct implementation" guarantee for everyone downstream. The
fix belongs in the shared package.

* License

[[file:LICENSE.org][CC-BY-NC-SA-4.0]].