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.
- emu_gmm-0.3.1/LICENSE.org +19 -0
- emu_gmm-0.3.1/PKG-INFO +191 -0
- emu_gmm-0.3.1/README.org +165 -0
- emu_gmm-0.3.1/pyproject.toml +96 -0
- emu_gmm-0.3.1/src/emu_gmm/__init__.py +213 -0
- emu_gmm-0.3.1/src/emu_gmm/_internal/__init__.py +6 -0
- emu_gmm-0.3.1/src/emu_gmm/_internal/axes.py +93 -0
- emu_gmm-0.3.1/src/emu_gmm/_internal/cholesky.py +128 -0
- emu_gmm-0.3.1/src/emu_gmm/_internal/fn_cache.py +120 -0
- emu_gmm-0.3.1/src/emu_gmm/_internal/labels.py +399 -0
- emu_gmm-0.3.1/src/emu_gmm/_internal/nan_safety.py +118 -0
- emu_gmm-0.3.1/src/emu_gmm/_internal/params.py +579 -0
- emu_gmm-0.3.1/src/emu_gmm/_internal/pinv_eigvalrule.py +133 -0
- emu_gmm-0.3.1/src/emu_gmm/covariance/__init__.py +21 -0
- emu_gmm-0.3.1/src/emu_gmm/covariance/analytical.py +76 -0
- emu_gmm-0.3.1/src/emu_gmm/covariance/clustered.py +239 -0
- emu_gmm-0.3.1/src/emu_gmm/covariance/iid.py +174 -0
- emu_gmm-0.3.1/src/emu_gmm/covariance/stratified.py +664 -0
- emu_gmm-0.3.1/src/emu_gmm/covariance/sum.py +167 -0
- emu_gmm-0.3.1/src/emu_gmm/covariance/synthetic.py +108 -0
- emu_gmm-0.3.1/src/emu_gmm/diagnostics.py +543 -0
- emu_gmm-0.3.1/src/emu_gmm/estimator.py +1347 -0
- emu_gmm-0.3.1/src/emu_gmm/examples/__init__.py +5 -0
- emu_gmm-0.3.1/src/emu_gmm/examples/euler.py +176 -0
- emu_gmm-0.3.1/src/emu_gmm/inference/__init__.py +93 -0
- emu_gmm-0.3.1/src/emu_gmm/inference/adaptive.py +433 -0
- emu_gmm-0.3.1/src/emu_gmm/inference/cluster_bootstrap.py +464 -0
- emu_gmm-0.3.1/src/emu_gmm/inference/confidence_set.py +382 -0
- emu_gmm-0.3.1/src/emu_gmm/inference/functional_se.py +364 -0
- emu_gmm-0.3.1/src/emu_gmm/inference/j_test.py +197 -0
- emu_gmm-0.3.1/src/emu_gmm/inference/k_statistic.py +781 -0
- emu_gmm-0.3.1/src/emu_gmm/inference/wild_bootstrap.py +379 -0
- emu_gmm-0.3.1/src/emu_gmm/manifolds/__init__.py +57 -0
- emu_gmm-0.3.1/src/emu_gmm/manifolds/base.py +117 -0
- emu_gmm-0.3.1/src/emu_gmm/manifolds/euclidean.py +131 -0
- emu_gmm-0.3.1/src/emu_gmm/manifolds/manifold_leaf.py +179 -0
- emu_gmm-0.3.1/src/emu_gmm/manifolds/optimizer.py +43 -0
- emu_gmm-0.3.1/src/emu_gmm/manifolds/positive.py +159 -0
- emu_gmm-0.3.1/src/emu_gmm/manifolds/product.py +165 -0
- emu_gmm-0.3.1/src/emu_gmm/manifolds/psd_fixed_rank.py +230 -0
- emu_gmm-0.3.1/src/emu_gmm/manifolds/riemannian_lm.py +462 -0
- emu_gmm-0.3.1/src/emu_gmm/manifolds/spec.py +79 -0
- emu_gmm-0.3.1/src/emu_gmm/measures/__init__.py +7 -0
- emu_gmm-0.3.1/src/emu_gmm/measures/analytical.py +108 -0
- emu_gmm-0.3.1/src/emu_gmm/measures/empirical.py +838 -0
- emu_gmm-0.3.1/src/emu_gmm/measures/synthetic.py +227 -0
- emu_gmm-0.3.1/src/emu_gmm/numerics/__init__.py +17 -0
- emu_gmm-0.3.1/src/emu_gmm/numerics/ridge_inverse.py +152 -0
- emu_gmm-0.3.1/src/emu_gmm/optimizer.py +654 -0
- emu_gmm-0.3.1/src/emu_gmm/parameter_space.py +260 -0
- emu_gmm-0.3.1/src/emu_gmm/penalty.py +104 -0
- emu_gmm-0.3.1/src/emu_gmm/regularization.py +336 -0
- emu_gmm-0.3.1/src/emu_gmm/studies/__init__.py +66 -0
- emu_gmm-0.3.1/src/emu_gmm/studies/driver.py +266 -0
- emu_gmm-0.3.1/src/emu_gmm/studies/study.py +122 -0
- emu_gmm-0.3.1/src/emu_gmm/studies/summaries.py +345 -0
- emu_gmm-0.3.1/src/emu_gmm/types.py +1185 -0
- 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
|
+
|
emu_gmm-0.3.1/README.org
ADDED
|
@@ -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
|
+
]
|