emu-gmm 0.3.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (58) hide show
  1. emu_gmm-0.3.1/LICENSE.org +19 -0
  2. emu_gmm-0.3.1/PKG-INFO +191 -0
  3. emu_gmm-0.3.1/README.org +165 -0
  4. emu_gmm-0.3.1/pyproject.toml +96 -0
  5. emu_gmm-0.3.1/src/emu_gmm/__init__.py +213 -0
  6. emu_gmm-0.3.1/src/emu_gmm/_internal/__init__.py +6 -0
  7. emu_gmm-0.3.1/src/emu_gmm/_internal/axes.py +93 -0
  8. emu_gmm-0.3.1/src/emu_gmm/_internal/cholesky.py +128 -0
  9. emu_gmm-0.3.1/src/emu_gmm/_internal/fn_cache.py +120 -0
  10. emu_gmm-0.3.1/src/emu_gmm/_internal/labels.py +399 -0
  11. emu_gmm-0.3.1/src/emu_gmm/_internal/nan_safety.py +118 -0
  12. emu_gmm-0.3.1/src/emu_gmm/_internal/params.py +579 -0
  13. emu_gmm-0.3.1/src/emu_gmm/_internal/pinv_eigvalrule.py +133 -0
  14. emu_gmm-0.3.1/src/emu_gmm/covariance/__init__.py +21 -0
  15. emu_gmm-0.3.1/src/emu_gmm/covariance/analytical.py +76 -0
  16. emu_gmm-0.3.1/src/emu_gmm/covariance/clustered.py +239 -0
  17. emu_gmm-0.3.1/src/emu_gmm/covariance/iid.py +174 -0
  18. emu_gmm-0.3.1/src/emu_gmm/covariance/stratified.py +664 -0
  19. emu_gmm-0.3.1/src/emu_gmm/covariance/sum.py +167 -0
  20. emu_gmm-0.3.1/src/emu_gmm/covariance/synthetic.py +108 -0
  21. emu_gmm-0.3.1/src/emu_gmm/diagnostics.py +543 -0
  22. emu_gmm-0.3.1/src/emu_gmm/estimator.py +1347 -0
  23. emu_gmm-0.3.1/src/emu_gmm/examples/__init__.py +5 -0
  24. emu_gmm-0.3.1/src/emu_gmm/examples/euler.py +176 -0
  25. emu_gmm-0.3.1/src/emu_gmm/inference/__init__.py +93 -0
  26. emu_gmm-0.3.1/src/emu_gmm/inference/adaptive.py +433 -0
  27. emu_gmm-0.3.1/src/emu_gmm/inference/cluster_bootstrap.py +464 -0
  28. emu_gmm-0.3.1/src/emu_gmm/inference/confidence_set.py +382 -0
  29. emu_gmm-0.3.1/src/emu_gmm/inference/functional_se.py +364 -0
  30. emu_gmm-0.3.1/src/emu_gmm/inference/j_test.py +197 -0
  31. emu_gmm-0.3.1/src/emu_gmm/inference/k_statistic.py +781 -0
  32. emu_gmm-0.3.1/src/emu_gmm/inference/wild_bootstrap.py +379 -0
  33. emu_gmm-0.3.1/src/emu_gmm/manifolds/__init__.py +57 -0
  34. emu_gmm-0.3.1/src/emu_gmm/manifolds/base.py +117 -0
  35. emu_gmm-0.3.1/src/emu_gmm/manifolds/euclidean.py +131 -0
  36. emu_gmm-0.3.1/src/emu_gmm/manifolds/manifold_leaf.py +179 -0
  37. emu_gmm-0.3.1/src/emu_gmm/manifolds/optimizer.py +43 -0
  38. emu_gmm-0.3.1/src/emu_gmm/manifolds/positive.py +159 -0
  39. emu_gmm-0.3.1/src/emu_gmm/manifolds/product.py +165 -0
  40. emu_gmm-0.3.1/src/emu_gmm/manifolds/psd_fixed_rank.py +230 -0
  41. emu_gmm-0.3.1/src/emu_gmm/manifolds/riemannian_lm.py +462 -0
  42. emu_gmm-0.3.1/src/emu_gmm/manifolds/spec.py +79 -0
  43. emu_gmm-0.3.1/src/emu_gmm/measures/__init__.py +7 -0
  44. emu_gmm-0.3.1/src/emu_gmm/measures/analytical.py +108 -0
  45. emu_gmm-0.3.1/src/emu_gmm/measures/empirical.py +838 -0
  46. emu_gmm-0.3.1/src/emu_gmm/measures/synthetic.py +227 -0
  47. emu_gmm-0.3.1/src/emu_gmm/numerics/__init__.py +17 -0
  48. emu_gmm-0.3.1/src/emu_gmm/numerics/ridge_inverse.py +152 -0
  49. emu_gmm-0.3.1/src/emu_gmm/optimizer.py +654 -0
  50. emu_gmm-0.3.1/src/emu_gmm/parameter_space.py +260 -0
  51. emu_gmm-0.3.1/src/emu_gmm/penalty.py +104 -0
  52. emu_gmm-0.3.1/src/emu_gmm/regularization.py +336 -0
  53. emu_gmm-0.3.1/src/emu_gmm/studies/__init__.py +66 -0
  54. emu_gmm-0.3.1/src/emu_gmm/studies/driver.py +266 -0
  55. emu_gmm-0.3.1/src/emu_gmm/studies/study.py +122 -0
  56. emu_gmm-0.3.1/src/emu_gmm/studies/summaries.py +345 -0
  57. emu_gmm-0.3.1/src/emu_gmm/types.py +1185 -0
  58. emu_gmm-0.3.1/src/emu_gmm/weighting.py +616 -0
@@ -0,0 +1,19 @@
1
+ #+TITLE: License
2
+ #+OPTIONS: toc:nil num:nil
3
+
4
+ * Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International (CC BY-NC-SA 4.0)
5
+
6
+ This work is licensed under the
7
+ [[http://creativecommons.org/licenses/by-nc-sa/4.0/][Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International License]].
8
+
9
+ To view a copy of this license, visit:
10
+ http://creativecommons.org/licenses/by-nc-sa/4.0/
11
+ or send a letter to Creative Commons, PO Box 1866, Mountain View, CA 94042, USA.
12
+
13
+ ** Summary (non-legal)
14
+ - *Attribution*: credit "Ethan Ligon et al., emu-gmm project"
15
+ - *Non-Commercial*: may not be used for commercial purposes
16
+ - *ShareAlike*: derivative works must carry the same license
17
+ - *No additional restrictions*: you may not impose legal or technological barriers that prevent others from exercising the license terms.
18
+
19
+ Full legal code available at the above URL.
emu_gmm-0.3.1/PKG-INFO ADDED
@@ -0,0 +1,191 @@
1
+ Metadata-Version: 2.4
2
+ Name: emu-gmm
3
+ Version: 0.3.1
4
+ Summary: Measure-theoretic Generalized Method of Moments estimation; estimation via E_mu.
5
+ License: CC-BY-NC-SA-4.0
6
+ License-File: LICENSE.org
7
+ Author: Ethan Ligon
8
+ Author-email: ligon@berkeley.edu
9
+ Requires-Python: >=3.11,<4.0
10
+ Classifier: Programming Language :: Python :: 3
11
+ Classifier: Programming Language :: Python :: 3.11
12
+ Classifier: License :: Other/Proprietary License
13
+ Classifier: Operating System :: OS Independent
14
+ Classifier: Intended Audience :: Science/Research
15
+ Classifier: Topic :: Scientific/Engineering
16
+ Requires-Dist: haliax (>=1.4.dev452,<1.5)
17
+ Requires-Dist: jax (>=0.4.25)
18
+ Requires-Dist: jax_dataclasses (>=1.6.1)
19
+ Requires-Dist: jaxtyping (>=0.2.30)
20
+ Requires-Dist: numpy (>=1.26)
21
+ Requires-Dist: optimistix (>=0.0.7)
22
+ Requires-Dist: pandas (>=2.0)
23
+ Requires-Dist: scipy (>=1.11)
24
+ Description-Content-Type: text/plain
25
+
26
+ #+TITLE: emu-gmm
27
+ #+SUBTITLE: Measure-theoretic GMM; estimation via $\mathbb{E}_\mu$
28
+ #+AUTHOR: Ethan Ligon
29
+ #+OPTIONS: toc:nil num:nil
30
+
31
+ =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.
32
+
33
+ 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.
34
+
35
+ * Status
36
+
37
+ 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/=.
38
+
39
+ * Quickstart
40
+
41
+ Bootstrap and try the worked example:
42
+
43
+ #+begin_src shell
44
+ make setup # poetry install + .venv/
45
+ poetry run python examples/run_euler.py # multi-asset Hansen-Singleton demo
46
+ #+end_src
47
+
48
+ The minimal estimation surface in code (synthetic-data variant):
49
+
50
+ #+begin_src python
51
+ import jax
52
+ from emu_gmm import estimate, SyntheticMeasure, SyntheticCovariance
53
+ from emu_gmm.examples.euler import (
54
+ EulerParams, euler_residual, euler_sampler_factory,
55
+ )
56
+
57
+ measure = SyntheticMeasure(
58
+ key=jax.random.PRNGKey(0),
59
+ n_sim=5000,
60
+ sampler=euler_sampler_factory(5000),
61
+ )
62
+
63
+ result = estimate(
64
+ model=euler_residual,
65
+ measure=measure,
66
+ covariance=SyntheticCovariance(),
67
+ theta_init=EulerParams(beta=0.9, gamma=1.0),
68
+ )
69
+
70
+ print(f"beta = {result.theta_hat.beta:.4f} (truth 0.96)")
71
+ print(f"gamma = {result.theta_hat.gamma:.4f} (truth 2.00)")
72
+ print(f"J-stat = {result.J_stat:.3f} dof={result.J_dof} p={result.J_pvalue:.3f}")
73
+ print(result.to_pandas()["Sigma_theta"])
74
+ #+end_src
75
+
76
+ 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.
77
+
78
+ * Documents
79
+
80
+ - =docs/howto.org= --- *start here.* Integration HOWTO: the architecture
81
+ at the level you need to wire an application to =estimate()= --- what
82
+ code you write (the moment function, the parameter container, the
83
+ measure/covariance choice) and what the framework hands back.
84
+ - =docs/design.org= --- architecture specification (four review rounds; stable). Its "Contract surface (implemented)" section is the single source of truth for the implemented menu.
85
+ - =docs/api-sketch.org= --- the v1 API surface, retained as the v1 design record; superseded on status by =design.org='s contract-surface section.
86
+ - =docs/implementation-plan.org= --- phased task list (Phases 1-7 complete, Phase 8 polish underway; the v2 roadmap is Section 13).
87
+ - =docs/mcar-asymptotics.org= --- companion theoretical note; consistency, asymptotic normality, PD properties of the pairwise-overlap estimator under MCAR.
88
+ - =docs/refs.bib= --- project-local bibliography (entries not in =~/bibtex/main.bib=).
89
+
90
+ * Setup
91
+
92
+ Requires Python >= 3.11 and [[https://python-poetry.org/][Poetry]].
93
+
94
+ #+begin_src shell
95
+ make setup
96
+ #+end_src
97
+
98
+ This runs =poetry install= and creates =.venv/= in the project root. Activate with =direnv allow= or =poetry shell=.
99
+
100
+ * Development
101
+
102
+ #+begin_src shell
103
+ make check # ruff + black + mypy + full pytest
104
+ make quick-check # same, skipping slow tests
105
+ make test # pytest only
106
+ #+end_src
107
+
108
+ To install pre-commit hooks (ruff + black on every commit):
109
+
110
+ #+begin_src shell
111
+ poetry run pre-commit install
112
+ #+end_src
113
+
114
+ * Layout
115
+
116
+ #+begin_example
117
+ docs/ design specs and theoretical notes
118
+ src/emu_gmm/
119
+ __init__.py public API re-exports
120
+ types.py protocols + EstimationResult / Diagnostics
121
+ estimator.py estimate() / build_estimator() entry points
122
+ measures/ SyntheticMeasure, AnalyticalMeasure, EmpiricalMeasure
123
+ covariance/ the CovarianceStrategy ladder (iid, clustered,
124
+ stratified / design-aware, sum, analytical, synthetic)
125
+ weighting.py Identity, Fixed, IteratedWeighting, ContinuouslyUpdated
126
+ regularization.py DiagonalTikhonov
127
+ penalty.py TikhonovPenalty
128
+ optimizer.py optimistix_lm, scipy_lm, linear_solver
129
+ manifolds/ parameter geometry (Euclidean, Positive, PSDFixedRank,
130
+ Product, ManifoldLeaf) + riemannian_lm
131
+ parameter_space.py ParameterSpace / on field-to-manifold declarations
132
+ inference/ j_test, k_statistic, k_confidence_set, bootstraps
133
+ numerics/ ridge_inverse
134
+ studies/ Monte Carlo driver (subpackage-only; not re-exported)
135
+ diagnostics.py build_diagnostics, log_to_stdout
136
+ examples/ shared example models (euler.py)
137
+ _internal/ axes, params, cholesky, labels (private)
138
+ tests/ test suite mirroring src/emu_gmm/
139
+ examples/ runnable demo scripts
140
+ Makefile build automation
141
+ pyproject.toml Poetry configuration
142
+ poetry.lock pinned dependencies (committed)
143
+ #+end_example
144
+
145
+ * Public API
146
+
147
+ Everything user-facing is re-exported at the package top level
148
+ (=from emu_gmm import ...=), with one deliberate exception: the Monte
149
+ Carlo =studies= subpackage is subpackage-only (import from
150
+ =emu_gmm.studies= directly). The implemented menu --- entry points,
151
+ measures, covariance strategies, weighting, regularization, penalty,
152
+ optimisers, manifold geometry, inference helpers, and result types ---
153
+ is tabulated, one line of semantics and a module pointer per row, in
154
+ =docs/design.org= under "Contract surface (implemented)". That table is
155
+ the single source of truth this README defers to;
156
+ =sorted(emu_gmm.__all__)= is the export list it tracks.
157
+
158
+ * Using emu-gmm correctly (and reporting problems)
159
+
160
+ =emu-gmm= is intended as a /spare set of correct interfaces/. Your only
161
+ modelling input is a per-observation residual =psi(x_i, theta) -> R^M=;
162
+ the package owns the moment expectation, the design-aware covariance
163
+ =V_X=, the criterion, the \(J\)-statistic, standard errors, and
164
+ \(p\)-values. Two rules keep estimation correct:
165
+
166
+ 1. *Read results off the package, never recompute them by hand.* Take the
167
+ criterion, =J_stat=, =standard_errors=, and \(p\)-values from
168
+ =EstimationResult= / =k_statistic= / the inference helpers. A
169
+ hand-rolled criterion is how subtle scaling bugs get reintroduced ---
170
+ notably the textbook habit of scaling moments by \(\sqrt{N}\) and
171
+ weighting by \(V_X\), which is correct only when every moment has the
172
+ /same/ number of observations. =emu-gmm= keeps moments as
173
+ per-coordinate means and folds all the per-moment \(N_j\) bookkeeping
174
+ into \(V_X\) (the \(1/(N_j N_k)\) normalisation); see =docs/design.org=
175
+ "Scaling convention" and =docs/mcar-asymptotics.org=.
176
+
177
+ 2. *Express per-moment observability only through the =mask=.* Different
178
+ moments may be observed for different numbers of units; the \((N, M)\)
179
+ =mask= on =EmpiricalMeasure= is the single place that lives. Do not
180
+ work around it by pre-aggregating or by assuming a common \(N\).
181
+
182
+ If an interface is missing a knob, or you believe it gives a wrong
183
+ answer: *file an issue against =emu-gmm=, do not reimplement the
184
+ statistic in your own project.* A local reimplementation forfeits the
185
+ "single correct implementation" guarantee for everyone downstream. The
186
+ fix belongs in the shared package.
187
+
188
+ * License
189
+
190
+ [[file:LICENSE.org][CC-BY-NC-SA-4.0]].
191
+
@@ -0,0 +1,165 @@
1
+ #+TITLE: emu-gmm
2
+ #+SUBTITLE: Measure-theoretic GMM; estimation via $\mathbb{E}_\mu$
3
+ #+AUTHOR: Ethan Ligon
4
+ #+OPTIONS: toc:nil num:nil
5
+
6
+ =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.
7
+
8
+ 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.
9
+
10
+ * Status
11
+
12
+ 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/=.
13
+
14
+ * Quickstart
15
+
16
+ Bootstrap and try the worked example:
17
+
18
+ #+begin_src shell
19
+ make setup # poetry install + .venv/
20
+ poetry run python examples/run_euler.py # multi-asset Hansen-Singleton demo
21
+ #+end_src
22
+
23
+ The minimal estimation surface in code (synthetic-data variant):
24
+
25
+ #+begin_src python
26
+ import jax
27
+ from emu_gmm import estimate, SyntheticMeasure, SyntheticCovariance
28
+ from emu_gmm.examples.euler import (
29
+ EulerParams, euler_residual, euler_sampler_factory,
30
+ )
31
+
32
+ measure = SyntheticMeasure(
33
+ key=jax.random.PRNGKey(0),
34
+ n_sim=5000,
35
+ sampler=euler_sampler_factory(5000),
36
+ )
37
+
38
+ result = estimate(
39
+ model=euler_residual,
40
+ measure=measure,
41
+ covariance=SyntheticCovariance(),
42
+ theta_init=EulerParams(beta=0.9, gamma=1.0),
43
+ )
44
+
45
+ print(f"beta = {result.theta_hat.beta:.4f} (truth 0.96)")
46
+ print(f"gamma = {result.theta_hat.gamma:.4f} (truth 2.00)")
47
+ print(f"J-stat = {result.J_stat:.3f} dof={result.J_dof} p={result.J_pvalue:.3f}")
48
+ print(result.to_pandas()["Sigma_theta"])
49
+ #+end_src
50
+
51
+ 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.
52
+
53
+ * Documents
54
+
55
+ - =docs/howto.org= --- *start here.* Integration HOWTO: the architecture
56
+ at the level you need to wire an application to =estimate()= --- what
57
+ code you write (the moment function, the parameter container, the
58
+ measure/covariance choice) and what the framework hands back.
59
+ - =docs/design.org= --- architecture specification (four review rounds; stable). Its "Contract surface (implemented)" section is the single source of truth for the implemented menu.
60
+ - =docs/api-sketch.org= --- the v1 API surface, retained as the v1 design record; superseded on status by =design.org='s contract-surface section.
61
+ - =docs/implementation-plan.org= --- phased task list (Phases 1-7 complete, Phase 8 polish underway; the v2 roadmap is Section 13).
62
+ - =docs/mcar-asymptotics.org= --- companion theoretical note; consistency, asymptotic normality, PD properties of the pairwise-overlap estimator under MCAR.
63
+ - =docs/refs.bib= --- project-local bibliography (entries not in =~/bibtex/main.bib=).
64
+
65
+ * Setup
66
+
67
+ Requires Python >= 3.11 and [[https://python-poetry.org/][Poetry]].
68
+
69
+ #+begin_src shell
70
+ make setup
71
+ #+end_src
72
+
73
+ This runs =poetry install= and creates =.venv/= in the project root. Activate with =direnv allow= or =poetry shell=.
74
+
75
+ * Development
76
+
77
+ #+begin_src shell
78
+ make check # ruff + black + mypy + full pytest
79
+ make quick-check # same, skipping slow tests
80
+ make test # pytest only
81
+ #+end_src
82
+
83
+ To install pre-commit hooks (ruff + black on every commit):
84
+
85
+ #+begin_src shell
86
+ poetry run pre-commit install
87
+ #+end_src
88
+
89
+ * Layout
90
+
91
+ #+begin_example
92
+ docs/ design specs and theoretical notes
93
+ src/emu_gmm/
94
+ __init__.py public API re-exports
95
+ types.py protocols + EstimationResult / Diagnostics
96
+ estimator.py estimate() / build_estimator() entry points
97
+ measures/ SyntheticMeasure, AnalyticalMeasure, EmpiricalMeasure
98
+ covariance/ the CovarianceStrategy ladder (iid, clustered,
99
+ stratified / design-aware, sum, analytical, synthetic)
100
+ weighting.py Identity, Fixed, IteratedWeighting, ContinuouslyUpdated
101
+ regularization.py DiagonalTikhonov
102
+ penalty.py TikhonovPenalty
103
+ optimizer.py optimistix_lm, scipy_lm, linear_solver
104
+ manifolds/ parameter geometry (Euclidean, Positive, PSDFixedRank,
105
+ Product, ManifoldLeaf) + riemannian_lm
106
+ parameter_space.py ParameterSpace / on field-to-manifold declarations
107
+ inference/ j_test, k_statistic, k_confidence_set, bootstraps
108
+ numerics/ ridge_inverse
109
+ studies/ Monte Carlo driver (subpackage-only; not re-exported)
110
+ diagnostics.py build_diagnostics, log_to_stdout
111
+ examples/ shared example models (euler.py)
112
+ _internal/ axes, params, cholesky, labels (private)
113
+ tests/ test suite mirroring src/emu_gmm/
114
+ examples/ runnable demo scripts
115
+ Makefile build automation
116
+ pyproject.toml Poetry configuration
117
+ poetry.lock pinned dependencies (committed)
118
+ #+end_example
119
+
120
+ * Public API
121
+
122
+ Everything user-facing is re-exported at the package top level
123
+ (=from emu_gmm import ...=), with one deliberate exception: the Monte
124
+ Carlo =studies= subpackage is subpackage-only (import from
125
+ =emu_gmm.studies= directly). The implemented menu --- entry points,
126
+ measures, covariance strategies, weighting, regularization, penalty,
127
+ optimisers, manifold geometry, inference helpers, and result types ---
128
+ is tabulated, one line of semantics and a module pointer per row, in
129
+ =docs/design.org= under "Contract surface (implemented)". That table is
130
+ the single source of truth this README defers to;
131
+ =sorted(emu_gmm.__all__)= is the export list it tracks.
132
+
133
+ * Using emu-gmm correctly (and reporting problems)
134
+
135
+ =emu-gmm= is intended as a /spare set of correct interfaces/. Your only
136
+ modelling input is a per-observation residual =psi(x_i, theta) -> R^M=;
137
+ the package owns the moment expectation, the design-aware covariance
138
+ =V_X=, the criterion, the \(J\)-statistic, standard errors, and
139
+ \(p\)-values. Two rules keep estimation correct:
140
+
141
+ 1. *Read results off the package, never recompute them by hand.* Take the
142
+ criterion, =J_stat=, =standard_errors=, and \(p\)-values from
143
+ =EstimationResult= / =k_statistic= / the inference helpers. A
144
+ hand-rolled criterion is how subtle scaling bugs get reintroduced ---
145
+ notably the textbook habit of scaling moments by \(\sqrt{N}\) and
146
+ weighting by \(V_X\), which is correct only when every moment has the
147
+ /same/ number of observations. =emu-gmm= keeps moments as
148
+ per-coordinate means and folds all the per-moment \(N_j\) bookkeeping
149
+ into \(V_X\) (the \(1/(N_j N_k)\) normalisation); see =docs/design.org=
150
+ "Scaling convention" and =docs/mcar-asymptotics.org=.
151
+
152
+ 2. *Express per-moment observability only through the =mask=.* Different
153
+ moments may be observed for different numbers of units; the \((N, M)\)
154
+ =mask= on =EmpiricalMeasure= is the single place that lives. Do not
155
+ work around it by pre-aggregating or by assuming a common \(N\).
156
+
157
+ If an interface is missing a knob, or you believe it gives a wrong
158
+ answer: *file an issue against =emu-gmm=, do not reimplement the
159
+ statistic in your own project.* A local reimplementation forfeits the
160
+ "single correct implementation" guarantee for everyone downstream. The
161
+ fix belongs in the shared package.
162
+
163
+ * License
164
+
165
+ [[file:LICENSE.org][CC-BY-NC-SA-4.0]].
@@ -0,0 +1,96 @@
1
+ [project]
2
+ name = "emu-gmm"
3
+ version = "0.3.1"
4
+ description = "Measure-theoretic Generalized Method of Moments estimation; estimation via E_mu."
5
+ readme = "README.org"
6
+ requires-python = ">=3.11,<4.0"
7
+ license = { text = "CC-BY-NC-SA-4.0" }
8
+ authors = [{ name = "Ethan Ligon", email = "ligon@berkeley.edu" }]
9
+ classifiers = [
10
+ "Programming Language :: Python :: 3",
11
+ "Programming Language :: Python :: 3.11",
12
+ "License :: Other/Proprietary License",
13
+ "Operating System :: OS Independent",
14
+ "Intended Audience :: Science/Research",
15
+ "Topic :: Scientific/Engineering"
16
+ ]
17
+ dependencies = [
18
+ "jax>=0.4.25",
19
+ "jaxtyping>=0.2.30",
20
+ "haliax>=1.4.dev452,<1.5",
21
+ "jax_dataclasses>=1.6.1",
22
+ "optimistix>=0.0.7",
23
+ "numpy>=1.26",
24
+ "scipy>=1.11",
25
+ "pandas>=2.0"
26
+ ]
27
+
28
+ [tool.poetry]
29
+ package-mode = true
30
+ packages = [{ include = "emu_gmm", from = "src" }]
31
+
32
+ [tool.poetry.dependencies]
33
+ python = ">=3.11,<4.0"
34
+ jax = ">=0.4.25"
35
+ jaxtyping = ">=0.2.30"
36
+ # PyPI's stable haliax 1.3 imports jax._src.tree_util.BuiltInKeyEntry, which
37
+ # current JAX no longer exposes; upstream main is compatible but PyPI forbids
38
+ # direct (git URL) dependencies in published metadata (the v0.3.0 publish
39
+ # rejection). haliax auto-publishes dev wheels from main: 1.4.dev452 is built
40
+ # from 280f49fb (2025-11-08) -- the EXACT revision the previous git pin had
41
+ # locked and the whole v2 validation ran against. The pre-release specifier
42
+ # is deliberate and PEP 440-explicit (it enables pre-release resolution for
43
+ # this requirement); the <1.5 ceiling guards API drift, and a future stable
44
+ # 1.4 satisfies the range. Suite re-validated against the PyPI wheel before
45
+ # this change landed.
46
+ haliax = { version = ">=1.4.dev452,<1.5", allow-prereleases = true }
47
+ jax_dataclasses = ">=1.6.1"
48
+ optimistix = ">=0.0.7"
49
+ numpy = ">=1.26"
50
+ scipy = ">=1.11"
51
+ pandas = ">=2.0"
52
+
53
+ [tool.poetry.group.dev.dependencies]
54
+ pytest = "^8.2.0"
55
+ pytest-xdist = "^3.0"
56
+ ipython = "^8.26.0"
57
+ ruff = "^0.5.4"
58
+ black = "^24.4.2"
59
+ mypy = "^1.10.0"
60
+
61
+ [tool.pytest.ini_options]
62
+ testpaths = ["tests"]
63
+ markers = [
64
+ "slow: tests that perform heavier computations"
65
+ ]
66
+
67
+ [tool.black]
68
+ target-version = ["py311"]
69
+ line-length = 88
70
+
71
+ [tool.ruff]
72
+ target-version = "py311"
73
+ line-length = 88
74
+ extend-exclude = [".claude", ".coder-session", "examples"]
75
+
76
+ [tool.ruff.lint]
77
+ select = ["E", "F", "I", "B", "UP"]
78
+ # E501 (line length): handled by black.
79
+ # F722 (syntax error in forward annotation), F821 (undefined name in
80
+ # forward annotation), UP037 (remove quotes from type annotation):
81
+ # all false positives on jaxtyping shape strings like
82
+ # Float[Array, "M K"]. Suppressed project-wide.
83
+ ignore = ["E501", "F722", "F821", "UP037"]
84
+
85
+ [tool.ruff.lint.isort]
86
+ known-third-party = ["jax", "jaxtyping", "haliax", "jax_dataclasses", "optimistix", "numpy", "scipy"]
87
+
88
+ [tool.mypy]
89
+ python_version = "3.11"
90
+ warn_unused_ignores = true
91
+ ignore_missing_imports = true
92
+ strict_optional = true
93
+
94
+ [build-system]
95
+ requires = ["poetry-core>=1.7.0"]
96
+ build-backend = "poetry.core.masonry.api"
@@ -0,0 +1,213 @@
1
+ """emu-gmm: Measure-theoretic GMM. Estimation via E_mu.
2
+
3
+ The public API is re-exported here for convenience::
4
+
5
+ from emu_gmm import (
6
+ estimate,
7
+ SyntheticMeasure, SyntheticCovariance,
8
+ AnalyticalMeasure, AnalyticalCovariance,
9
+ EmpiricalMeasure, IIDCovariance, ClusteredCovariance,
10
+ Identity, Fixed, ContinuouslyUpdated, IteratedWeighting,
11
+ DiagonalTikhonov,
12
+ TikhonovPenalty,
13
+ optimistix_lm, scipy_lm,
14
+ EstimationResult, Diagnostics, OptimizerInfo,
15
+ )
16
+
17
+ See docs/design.org for the architectural specification, docs/api-sketch.org
18
+ for the v1 API surface, and docs/mcar-asymptotics.org for the asymptotic
19
+ theory under MCAR.
20
+ """
21
+
22
+ # ruff: noqa: E402
23
+ # Imports below the float64 config update intentionally come after a
24
+ # non-import statement.
25
+ from importlib.metadata import PackageNotFoundError, version
26
+
27
+ # Enable float64 in JAX *before* any module-under-this-package import
28
+ # touches jax.numpy. JAX defaults to float32 (a deep-learning convention);
29
+ # for a numerical-statistics framework where Cholesky pivots, condition
30
+ # numbers, and gradient norms cross many orders of magnitude, float32
31
+ # precision is the wrong baseline. Notable symptoms with float32:
32
+ # - optimistix LM cannot certify convergence at rtol=1e-8 because the
33
+ # float32 noise floor on a O(0.1)-magnitude whitened residual is
34
+ # around 2e-8;
35
+ # - theta_hat itself drifts at the third significant digit relative to
36
+ # the float64 solution.
37
+ # Users who explicitly want float32 can override after import with:
38
+ # jax.config.update("jax_enable_x64", False)
39
+ import jax as _jax
40
+
41
+ _jax.config.update("jax_enable_x64", True)
42
+
43
+ # Public re-exports.
44
+ from emu_gmm.covariance import (
45
+ AnalyticalCovariance,
46
+ ClusteredCovariance,
47
+ DesignAwareCovariance,
48
+ IIDCovariance,
49
+ StratifiedCovariance,
50
+ SumCovariance,
51
+ SyntheticCovariance,
52
+ )
53
+ from emu_gmm.diagnostics import (
54
+ build_diagnostics,
55
+ build_optimizer_health,
56
+ compute_cond_info,
57
+ )
58
+ from emu_gmm.estimator import build_estimator, estimate
59
+ from emu_gmm.inference import (
60
+ AdaptiveBootstrapResult,
61
+ BootstrapMean,
62
+ BootstrapPValue,
63
+ BootstrapQuantile,
64
+ BootstrapSE,
65
+ ClusterBootstrapResult,
66
+ JTestResult,
67
+ KConfidenceSet,
68
+ KStatisticResult,
69
+ WildBootstrapResult,
70
+ adaptive_bootstrap,
71
+ cluster_bootstrap,
72
+ j_test,
73
+ k_confidence_set,
74
+ k_statistic,
75
+ moment_wild_bootstrap,
76
+ )
77
+ from emu_gmm.manifolds import (
78
+ Euclidean,
79
+ ManifoldLeaf,
80
+ ManifoldParam,
81
+ Positive,
82
+ Product,
83
+ PSDFixedRank,
84
+ riemannian_lm,
85
+ )
86
+ from emu_gmm.measures import (
87
+ AnalyticalMeasure,
88
+ EmpiricalMeasure,
89
+ SyntheticMeasure,
90
+ )
91
+ from emu_gmm.numerics import ridge_inverse
92
+ from emu_gmm.optimizer import linear_solver, optimistix_lm, scipy_lm
93
+ from emu_gmm.parameter_space import ParameterSpace, on
94
+ from emu_gmm.penalty import PenaltyStrategy, TikhonovPenalty
95
+ from emu_gmm.regularization import DiagonalTikhonov
96
+ from emu_gmm.types import (
97
+ CovarianceStrategy,
98
+ Diagnostics,
99
+ Emu_GMM_DimensionError,
100
+ EstimationResult,
101
+ FitRecord,
102
+ LabelContext,
103
+ Measure,
104
+ Optimizer,
105
+ OptimizerInfo,
106
+ RegularizationStrategy,
107
+ StructuralModel,
108
+ WeightingStrategy,
109
+ )
110
+ from emu_gmm.weighting import (
111
+ CUE,
112
+ ContinuouslyUpdated,
113
+ Fixed,
114
+ Identity,
115
+ IteratedWeighting,
116
+ )
117
+
118
+ try:
119
+ __version__ = version("emu-gmm")
120
+ except PackageNotFoundError:
121
+ __version__ = "0.0.0"
122
+
123
+ __all__ = [
124
+ "__version__",
125
+ # Entry point
126
+ "estimate",
127
+ "build_estimator",
128
+ # Inference
129
+ "k_confidence_set",
130
+ "k_statistic",
131
+ "KConfidenceSet",
132
+ "KStatisticResult",
133
+ # Measures
134
+ "SyntheticMeasure",
135
+ "AnalyticalMeasure",
136
+ "EmpiricalMeasure",
137
+ # Covariance strategies
138
+ "SyntheticCovariance",
139
+ "AnalyticalCovariance",
140
+ "IIDCovariance",
141
+ "ClusteredCovariance",
142
+ "StratifiedCovariance",
143
+ "DesignAwareCovariance",
144
+ "SumCovariance",
145
+ # Weighting strategies
146
+ "Identity",
147
+ "Fixed",
148
+ "ContinuouslyUpdated",
149
+ "CUE",
150
+ "IteratedWeighting",
151
+ # Regularization
152
+ "DiagonalTikhonov",
153
+ # Numerics helpers
154
+ "ridge_inverse",
155
+ # Penalty (in-objective parameter penalty)
156
+ "TikhonovPenalty",
157
+ # Optimisers
158
+ "optimistix_lm",
159
+ "scipy_lm",
160
+ "linear_solver",
161
+ "riemannian_lm",
162
+ # Inference helpers
163
+ "cluster_bootstrap",
164
+ "ClusterBootstrapResult",
165
+ "j_test",
166
+ "JTestResult",
167
+ # Result / diagnostics types
168
+ "EstimationResult",
169
+ "FitRecord",
170
+ "Diagnostics",
171
+ "OptimizerInfo",
172
+ "LabelContext",
173
+ # Errors
174
+ "Emu_GMM_DimensionError",
175
+ # Diagnostics builders
176
+ "build_diagnostics",
177
+ "build_optimizer_health",
178
+ "compute_cond_info",
179
+ # Inference
180
+ "moment_wild_bootstrap",
181
+ "WildBootstrapResult",
182
+ # Adaptive (precision-targeted) bootstrap stopping (#91)
183
+ "adaptive_bootstrap",
184
+ "AdaptiveBootstrapResult",
185
+ "BootstrapMean",
186
+ "BootstrapSE",
187
+ "BootstrapQuantile",
188
+ "BootstrapPValue",
189
+ # Manifolds (parameter geometry — a first-class problem-tuple menu, peer
190
+ # to Measures / Covariance strategies; wrap a non-scalar parameter leaf in
191
+ # ManifoldLeaf(array, manifold) and estimate via the ordinary entry point.
192
+ # ManifoldSpec / LeafSpec stay submodule-internal — they are machinery,
193
+ # not part of the user-composed tuple.)
194
+ "Euclidean",
195
+ "PSDFixedRank",
196
+ "Product",
197
+ "Positive",
198
+ "ManifoldLeaf",
199
+ # Parameter-space declaration layer (#107): declare field -> manifold
200
+ # geometry once in a class, then ParameterSpace.point([seed]) lowers to a
201
+ # ManifoldLeaf pytree (a valid theta_init) consumed by estimate(parameters=).
202
+ "ParameterSpace",
203
+ "on",
204
+ # Protocols (for type-checking user code)
205
+ "Measure",
206
+ "CovarianceStrategy",
207
+ "WeightingStrategy",
208
+ "RegularizationStrategy",
209
+ "PenaltyStrategy",
210
+ "Optimizer",
211
+ "StructuralModel",
212
+ "ManifoldParam",
213
+ ]