emu-gmm 0.5.0

#+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 (v0.2.0): 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)")
print(f"J-stat = {result.J_stat:.3f}  dof={result.J_dof}  p={result.J_pvalue:.3f}")
print(result.to_pandas()["Sigma_theta"])
#+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.

* 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.standard_errors= / =Sigma_theta= | 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'\) | =functional_se= / =gamma_se=             | 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.

#+begin_src python
import numpy as np
from emu_gmm.inference import cluster_bootstrap

boot = cluster_bootstrap(
    model=euler_residual, theta_init=result.theta_hat, measure=measure,
    covariance=cov, n_boot=999, key=jax.random.PRNGKey(1),
)
ok = np.asarray(boot.convergence)        # non-converged refits are NaN rows (#209)
se = boot.theta_boot.array[ok].std(axis=0)   # bootstrap SE, per parameter
#+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, or 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. Design-aware note: =cluster_bootstrap= resamples flat clusters; the /stratified/ refit bootstrap (resample PSUs within strata) is tracked in issue #155.

* 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 and [[https://python-poetry.org/][Poetry]].

#+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 + EstimationResult / Diagnostics
    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, =J_stat=, =standard_errors=, and \(p\)-values from
   =EstimationResult= / =k_statistic= / 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]].