microimpute 3.1.4__tar.gz → 3.2.0__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 (49) hide show
  1. {microimpute-3.1.4 → microimpute-3.2.0}/PKG-INFO +2 -1
  2. {microimpute-3.1.4 → microimpute-3.2.0}/microimpute.egg-info/PKG-INFO +2 -1
  3. {microimpute-3.1.4 → microimpute-3.2.0}/microimpute.egg-info/SOURCES.txt +1 -0
  4. {microimpute-3.1.4 → microimpute-3.2.0}/microimpute.egg-info/requires.txt +1 -0
  5. {microimpute-3.1.4 → microimpute-3.2.0}/pyproject.toml +2 -1
  6. microimpute-3.2.0/tests/test_moop_spm_experiment.py +477 -0
  7. {microimpute-3.1.4 → microimpute-3.2.0}/README.md +0 -0
  8. {microimpute-3.1.4 → microimpute-3.2.0}/microimpute/__init__.py +0 -0
  9. {microimpute-3.1.4 → microimpute-3.2.0}/microimpute/comparisons/__init__.py +0 -0
  10. {microimpute-3.1.4 → microimpute-3.2.0}/microimpute/comparisons/autoimpute.py +0 -0
  11. {microimpute-3.1.4 → microimpute-3.2.0}/microimpute/comparisons/autoimpute_helpers.py +0 -0
  12. {microimpute-3.1.4 → microimpute-3.2.0}/microimpute/comparisons/imputations.py +0 -0
  13. {microimpute-3.1.4 → microimpute-3.2.0}/microimpute/comparisons/metrics.py +0 -0
  14. {microimpute-3.1.4 → microimpute-3.2.0}/microimpute/comparisons/validation.py +0 -0
  15. {microimpute-3.1.4 → microimpute-3.2.0}/microimpute/config.py +0 -0
  16. {microimpute-3.1.4 → microimpute-3.2.0}/microimpute/evaluations/__init__.py +0 -0
  17. {microimpute-3.1.4 → microimpute-3.2.0}/microimpute/evaluations/cross_validation.py +0 -0
  18. {microimpute-3.1.4 → microimpute-3.2.0}/microimpute/evaluations/predictor_analysis.py +0 -0
  19. {microimpute-3.1.4 → microimpute-3.2.0}/microimpute/models/__init__.py +0 -0
  20. {microimpute-3.1.4 → microimpute-3.2.0}/microimpute/models/imputer.py +0 -0
  21. {microimpute-3.1.4 → microimpute-3.2.0}/microimpute/models/matching.py +0 -0
  22. {microimpute-3.1.4 → microimpute-3.2.0}/microimpute/models/mdn.py +0 -0
  23. {microimpute-3.1.4 → microimpute-3.2.0}/microimpute/models/ols.py +0 -0
  24. {microimpute-3.1.4 → microimpute-3.2.0}/microimpute/models/qrf.py +0 -0
  25. {microimpute-3.1.4 → microimpute-3.2.0}/microimpute/models/quantreg.py +0 -0
  26. {microimpute-3.1.4 → microimpute-3.2.0}/microimpute/models/zero_inflated.py +0 -0
  27. {microimpute-3.1.4 → microimpute-3.2.0}/microimpute/py.typed +0 -0
  28. {microimpute-3.1.4 → microimpute-3.2.0}/microimpute/utils/__init__.py +0 -0
  29. {microimpute-3.1.4 → microimpute-3.2.0}/microimpute/utils/dashboard_formatter.py +0 -0
  30. {microimpute-3.1.4 → microimpute-3.2.0}/microimpute/utils/data.py +0 -0
  31. {microimpute-3.1.4 → microimpute-3.2.0}/microimpute/utils/statmatch_hotdeck.py +0 -0
  32. {microimpute-3.1.4 → microimpute-3.2.0}/microimpute/utils/type_handling.py +0 -0
  33. {microimpute-3.1.4 → microimpute-3.2.0}/microimpute/visualizations/__init__.py +0 -0
  34. {microimpute-3.1.4 → microimpute-3.2.0}/microimpute/visualizations/comparison_plots.py +0 -0
  35. {microimpute-3.1.4 → microimpute-3.2.0}/microimpute/visualizations/performance_plots.py +0 -0
  36. {microimpute-3.1.4 → microimpute-3.2.0}/microimpute.egg-info/dependency_links.txt +0 -0
  37. {microimpute-3.1.4 → microimpute-3.2.0}/microimpute.egg-info/top_level.txt +0 -0
  38. {microimpute-3.1.4 → microimpute-3.2.0}/setup.cfg +0 -0
  39. {microimpute-3.1.4 → microimpute-3.2.0}/tests/test_autoimpute.py +0 -0
  40. {microimpute-3.1.4 → microimpute-3.2.0}/tests/test_basic.py +0 -0
  41. {microimpute-3.1.4 → microimpute-3.2.0}/tests/test_dashboard_formatter.py +0 -0
  42. {microimpute-3.1.4 → microimpute-3.2.0}/tests/test_data_preprocessing.py +0 -0
  43. {microimpute-3.1.4 → microimpute-3.2.0}/tests/test_metrics.py +0 -0
  44. {microimpute-3.1.4 → microimpute-3.2.0}/tests/test_predictor_analysis.py +0 -0
  45. {microimpute-3.1.4 → microimpute-3.2.0}/tests/test_public_api.py +0 -0
  46. {microimpute-3.1.4 → microimpute-3.2.0}/tests/test_quantile_comparison.py +0 -0
  47. {microimpute-3.1.4 → microimpute-3.2.0}/tests/test_smoke_qrf.py +0 -0
  48. {microimpute-3.1.4 → microimpute-3.2.0}/tests/test_type_handling.py +0 -0
  49. {microimpute-3.1.4 → microimpute-3.2.0}/tests/test_visualizations.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: microimpute
3
- Version: 3.1.4
3
+ Version: 3.2.0
4
4
  Summary: Benchmarking imputation methods for microdata
5
5
  Author-email: Vahid Ahmadi <vahid@policyengine.org>, Max Ghenis <max@policyengine.org>, María Juaristi <juaristi@uni.minerva.edu>, Nikhil Woodruff <nikhil.woodruff@outlook.com>
6
6
  Project-URL: Repository, https://github.com/PolicyEngine/microimpute
@@ -34,6 +34,7 @@ Requires-Dist: ruff<0.17.0,>=0.16.0; extra == "dev"
34
34
  Requires-Dist: mypy<2.0.0,>=1.2.3; extra == "dev"
35
35
  Requires-Dist: build<2.0.0,>=1.2.0; extra == "dev"
36
36
  Requires-Dist: towncrier>=24.8.0; extra == "dev"
37
+ Requires-Dist: hypothesis<7.0.0,>=6.100.0; extra == "dev"
37
38
  Provides-Extra: matching
38
39
  Requires-Dist: rpy2<4.0.0,>=3.5.0; extra == "matching"
39
40
  Provides-Extra: mdn
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: microimpute
3
- Version: 3.1.4
3
+ Version: 3.2.0
4
4
  Summary: Benchmarking imputation methods for microdata
5
5
  Author-email: Vahid Ahmadi <vahid@policyengine.org>, Max Ghenis <max@policyengine.org>, María Juaristi <juaristi@uni.minerva.edu>, Nikhil Woodruff <nikhil.woodruff@outlook.com>
6
6
  Project-URL: Repository, https://github.com/PolicyEngine/microimpute
@@ -34,6 +34,7 @@ Requires-Dist: ruff<0.17.0,>=0.16.0; extra == "dev"
34
34
  Requires-Dist: mypy<2.0.0,>=1.2.3; extra == "dev"
35
35
  Requires-Dist: build<2.0.0,>=1.2.0; extra == "dev"
36
36
  Requires-Dist: towncrier>=24.8.0; extra == "dev"
37
+ Requires-Dist: hypothesis<7.0.0,>=6.100.0; extra == "dev"
37
38
  Provides-Extra: matching
38
39
  Requires-Dist: rpy2<4.0.0,>=3.5.0; extra == "matching"
39
40
  Provides-Extra: mdn
@@ -38,6 +38,7 @@ tests/test_basic.py
38
38
  tests/test_dashboard_formatter.py
39
39
  tests/test_data_preprocessing.py
40
40
  tests/test_metrics.py
41
+ tests/test_moop_spm_experiment.py
41
42
  tests/test_predictor_analysis.py
42
43
  tests/test_public_api.py
43
44
  tests/test_quantile_comparison.py
@@ -19,6 +19,7 @@ ruff<0.17.0,>=0.16.0
19
19
  mypy<2.0.0,>=1.2.3
20
20
  build<2.0.0,>=1.2.0
21
21
  towncrier>=24.8.0
22
+ hypothesis<7.0.0,>=6.100.0
22
23
 
23
24
  [docs]
24
25
  jupyter-book
@@ -7,7 +7,7 @@ build-backend = "setuptools.build_meta"
7
7
 
8
8
  [project]
9
9
  name = "microimpute"
10
- version = "3.1.4"
10
+ version = "3.2.0"
11
11
  description = "Benchmarking imputation methods for microdata"
12
12
  readme = "README.md"
13
13
  authors = [
@@ -49,6 +49,7 @@ dev = [
49
49
  "mypy>=1.2.3,<2.0.0",
50
50
  "build>=1.2.0,<2.0.0",
51
51
  "towncrier>=24.8.0",
52
+ "hypothesis>=6.100.0,<7.0.0",
52
53
  ]
53
54
 
54
55
  matching = [
@@ -0,0 +1,477 @@
1
+ """Property tests for the MOOP-imputation SPM poverty experiment.
2
+
3
+ The experiment lives in ``experiments/moop_spm``. These tests use small
4
+ synthetic data, so they run without the CPS ASEC download. They check
5
+ the invariants the published comparison relies on:
6
+
7
+ - Fold assignment partitions households and keeps each household whole.
8
+ - Poverty is recomputed exactly from resources, MOOP and thresholds, and
9
+ it never falls when any unit's MOOP rises.
10
+ - People whose unit is poor before MOOP stay poor under any imputation.
11
+ - Weighted group rates, replicate SEs and weighted quantiles match direct
12
+ reference computations.
13
+ - Cell pools hold only same-cell donor spenders, and the nearest-donor
14
+ search agrees with a brute-force search.
15
+ - The logistic-participation variants follow their stated rules.
16
+ - Census's MOOP and Medicare premium imputations both count as allocated.
17
+ - Extracted CSVs that differ from the pinned zip are extracted again.
18
+ - The fold cache misses whenever the environment, microimpute's source
19
+ or the model inputs change.
20
+ """
21
+
22
+ import argparse
23
+ import os
24
+ import zipfile
25
+
26
+ import numpy as np
27
+ import pandas as pd
28
+ import pytest
29
+ from hypothesis import HealthCheck, given, settings
30
+ from hypothesis import strategies as st
31
+
32
+ from experiments.moop_spm import run as run_module
33
+ from experiments.moop_spm.cps_asec import allocation_flags, extract_verified
34
+ from experiments.moop_spm.evaluate import (
35
+ RATIO_BANDS,
36
+ evaluate,
37
+ group_rates,
38
+ poverty_flags,
39
+ ratio_band_index,
40
+ replicate_se,
41
+ weighted_quantile,
42
+ )
43
+ from experiments.moop_spm.methods import (
44
+ CELL_LEVELS,
45
+ PARTICIPATION_CUTOFF,
46
+ FoldData,
47
+ _key,
48
+ assign_cells,
49
+ fit_participation,
50
+ logit_family,
51
+ nearest_donors,
52
+ participation_probability,
53
+ )
54
+ from experiments.moop_spm.run import (
55
+ assign_folds,
56
+ build_config,
57
+ collect_folds,
58
+ config_hash,
59
+ frame_sha256,
60
+ )
61
+
62
+ SETTINGS = settings(
63
+ max_examples=40,
64
+ deadline=None,
65
+ suppress_health_check=[HealthCheck.too_slow],
66
+ )
67
+
68
+
69
+ def synthetic_spm(seed: int, n_units: int = 60):
70
+ """Persons and units frames with the columns ``evaluate`` reads."""
71
+ rng = np.random.default_rng(seed)
72
+ unit_ids = np.arange(1, n_units + 1) * 7
73
+ threshold = rng.uniform(20_000, 40_000, n_units)
74
+ moop = np.where(rng.random(n_units) < 0.1, 0.0, rng.gamma(1.0, 5_000, n_units))
75
+ pre_moop = threshold * rng.uniform(0.5, 3.0, n_units)
76
+ units = pd.DataFrame(
77
+ {
78
+ "moop": moop,
79
+ "unit_weight": rng.uniform(100, 3_000, n_units),
80
+ "pre_moop_ratio": pre_moop / threshold,
81
+ "any_allocated": rng.integers(0, 2, n_units),
82
+ },
83
+ index=pd.Index(unit_ids, name="SPM_ID"),
84
+ )
85
+ sizes = rng.integers(1, 4, n_units)
86
+ person_unit = np.repeat(np.arange(n_units), sizes)
87
+ n_persons = len(person_unit)
88
+ persons = pd.DataFrame(
89
+ {
90
+ "SPM_ID": unit_ids[person_unit],
91
+ "A_AGE": rng.integers(0, 90, n_persons),
92
+ "PRDISFLG": rng.choice([-1, 1, 2], n_persons),
93
+ "coverage": rng.choice(
94
+ [
95
+ "medicare",
96
+ "employer",
97
+ "direct",
98
+ "medicaid",
99
+ "military",
100
+ "uninsured",
101
+ ],
102
+ n_persons,
103
+ ),
104
+ "SPM_MEDXPNS": moop[person_unit],
105
+ "SPM_RESOURCES": (pre_moop - moop)[person_unit],
106
+ "SPM_POVTHRESHOLD": threshold[person_unit],
107
+ "weight": rng.uniform(100, 3_000, n_persons),
108
+ }
109
+ )
110
+ persons["SPM_POOR"] = (
111
+ persons["SPM_RESOURCES"] < persons["SPM_POVTHRESHOLD"]
112
+ ).astype(int)
113
+ replicates = persons["weight"].to_numpy()[:, None] * rng.uniform(
114
+ 0.5, 1.5, (n_persons, 160)
115
+ )
116
+ replicate_frame = pd.DataFrame(
117
+ replicates, columns=[f"repwgt_{i}" for i in range(1, 161)]
118
+ )
119
+ persons = pd.concat([persons, replicate_frame], axis=1)
120
+ return persons, units
121
+
122
+
123
+ @SETTINGS
124
+ @given(
125
+ households=st.lists(st.integers(1, 40), min_size=5, max_size=120),
126
+ n_folds=st.integers(2, 6),
127
+ seed=st.integers(0, 2**31),
128
+ )
129
+ def test_folds_partition_households(households, n_folds, seed):
130
+ series = pd.Series(households)
131
+ folds = assign_folds(series, n_folds, seed)
132
+ assert set(np.unique(folds)) <= set(range(n_folds))
133
+ # Every household lands in exactly one fold.
134
+ assert (series.groupby(folds).nunique().sum()) == series.nunique()
135
+ per_household = pd.Series(folds).groupby(series.to_numpy()).nunique()
136
+ assert (per_household == 1).all()
137
+ # Household counts per fold differ by at most one.
138
+ household_counts = series.groupby(folds).nunique()
139
+ assert household_counts.max() - household_counts.min() <= 1
140
+ assert (assign_folds(series, n_folds, seed) == folds).all()
141
+
142
+
143
+ @SETTINGS
144
+ @given(seed=st.integers(0, 2**31), bump_seed=st.integers(0, 2**31))
145
+ def test_poverty_recomputation_and_monotonicity(seed, bump_seed):
146
+ persons, units = synthetic_spm(seed)
147
+ reported = poverty_flags(persons, units, units["moop"].to_numpy())[0]
148
+ assert (reported == (persons["SPM_POOR"] == 1)).all()
149
+
150
+ bump = np.random.default_rng(bump_seed).uniform(0, 20_000, len(units))
151
+ higher = poverty_flags(persons, units, units["moop"].to_numpy() + bump)[0]
152
+ # More MOOP can only push people into poverty.
153
+ assert (higher >= reported).all()
154
+
155
+ zero = poverty_flags(persons, units, np.zeros(len(units)))[0]
156
+ assert (zero <= reported).all()
157
+ pre_moop_poor = persons["SPM_ID"].map(units["pre_moop_ratio"]).to_numpy() < 1.0
158
+ # Units poor before MOOP stay poor under any nonnegative MOOP.
159
+ assert zero[pre_moop_poor].all()
160
+ assert higher[pre_moop_poor].all()
161
+
162
+
163
+ @SETTINGS
164
+ @given(seed=st.integers(0, 2**31))
165
+ def test_group_rates_match_np_average(seed):
166
+ rng = np.random.default_rng(seed)
167
+ n = 50
168
+ flags = rng.random(n)
169
+ weights = rng.uniform(0.1, 10, (n, 3))
170
+ groups = rng.random((n, 4)) < 0.5
171
+ groups[0] = True # keep every group nonempty
172
+ rates = group_rates(flags, groups, weights)
173
+ for g in range(groups.shape[1]):
174
+ for w in range(weights.shape[1]):
175
+ expected = np.average(flags[groups[:, g]], weights=weights[groups[:, g], w])
176
+ assert rates[g, w] == pytest.approx(expected)
177
+ assert ((rates >= 0) & (rates <= 1)).all()
178
+
179
+
180
+ @SETTINGS
181
+ @given(
182
+ base=st.floats(-5, 5),
183
+ deviations=st.lists(st.floats(-1, 1), min_size=160, max_size=160),
184
+ scale=st.floats(-3, 3),
185
+ )
186
+ def test_replicate_se_formula(base, deviations, scale):
187
+ estimates = np.array([base] + [base + d for d in deviations])
188
+ expected = np.sqrt(4 / 160 * np.sum(np.square(deviations)))
189
+ assert replicate_se(estimates) == pytest.approx(expected)
190
+ assert replicate_se(scale * estimates) == pytest.approx(abs(scale) * expected)
191
+ assert replicate_se(np.full(161, base)) == 0
192
+
193
+
194
+ @SETTINGS
195
+ @given(
196
+ values=st.lists(st.floats(-1e6, 1e6), min_size=1, max_size=60),
197
+ q=st.floats(0.01, 0.99),
198
+ )
199
+ def test_weighted_quantile_equal_weights(values, q):
200
+ values = np.array(values)
201
+ result = weighted_quantile(values, np.ones(len(values)), q)
202
+ assert result in values
203
+ assert result == pytest.approx(np.quantile(values, q, method="inverted_cdf"))
204
+
205
+
206
+ def _cells(rng: np.random.Generator, n: int) -> pd.DataFrame:
207
+ return pd.DataFrame(
208
+ {
209
+ "group": rng.choice(["medicare", "employer", "public"], n),
210
+ "age_band": rng.choice(["u45", "45_64", "65p"], n),
211
+ "income_band": rng.choice(["lt1", "1_2", "2_4", "4p"], n),
212
+ "size_band": rng.choice(["1", "2", "3p"], n),
213
+ }
214
+ )
215
+
216
+
217
+ @SETTINGS
218
+ @given(seed=st.integers(0, 2**31), min_donors=st.integers(1, 40))
219
+ def test_cell_pools_are_same_cell_donor_spenders(seed, min_donors):
220
+ rng = np.random.default_rng(seed)
221
+ donors = _cells(rng, 400)
222
+ moop = np.where(rng.random(400) < 0.2, 0.0, rng.gamma(1.0, 3_000, 400))
223
+ recipients = _cells(rng, 50)
224
+ assignment = assign_cells(donors, moop, recipients, min_donors=min_donors)
225
+ spender = moop > 0
226
+ donor_keys = [_key(donors.loc[spender], columns) for columns in CELL_LEVELS]
227
+ recipient_keys = [_key(recipients, columns) for columns in CELL_LEVELS]
228
+ for i, (pool, level) in enumerate(zip(assignment.pools, assignment.levels)):
229
+ assert len(pool) >= min_donors
230
+ assert (pool > 0).all()
231
+ in_cell = donor_keys[level] == recipient_keys[level][i]
232
+ expected = np.sort(moop[spender][in_cell])
233
+ assert np.array_equal(np.sort(pool), expected)
234
+ # No finer level had enough donor spenders.
235
+ for finer in range(level):
236
+ assert (donor_keys[finer] == recipient_keys[finer][i]).sum() < min_donors
237
+ assert np.allclose(assignment.means(), [p.mean() for p in assignment.pools])
238
+
239
+
240
+ @SETTINGS
241
+ @given(
242
+ donor_pred=st.lists(st.floats(-50, 50, allow_nan=False), min_size=1, max_size=40),
243
+ recipient_pred=st.lists(
244
+ st.floats(-60, 60, allow_nan=False), min_size=1, max_size=20
245
+ ),
246
+ k=st.integers(1, 8),
247
+ )
248
+ def test_nearest_donors_matches_brute_force(donor_pred, recipient_pred, k):
249
+ donor_pred = np.array(donor_pred)
250
+ recipient_pred = np.array(recipient_pred)
251
+ nearest = nearest_donors(donor_pred, recipient_pred, k)
252
+ k_eff = min(k, len(donor_pred))
253
+ assert nearest.shape == (len(recipient_pred), k_eff)
254
+ for i, value in enumerate(recipient_pred):
255
+ assert len(set(nearest[i])) == k_eff
256
+ got = np.sort(np.abs(donor_pred[nearest[i]] - value))
257
+ brute = np.sort(np.abs(donor_pred - value))[:k_eff]
258
+ assert np.allclose(got, brute)
259
+
260
+
261
+ def _fold(seed: int) -> FoldData:
262
+ rng = np.random.default_rng(seed)
263
+ n_donor, n_recipient = 600, 80
264
+
265
+ def features(n):
266
+ return pd.DataFrame(
267
+ {
268
+ "x1": rng.normal(size=n),
269
+ "x2": rng.normal(size=n),
270
+ "x3": rng.integers(0, 2, n).astype(float),
271
+ }
272
+ )
273
+
274
+ X_donor = features(n_donor)
275
+ logit = 2.0 + 1.5 * X_donor["x1"].to_numpy()
276
+ spends = rng.random(n_donor) < 1 / (1 + np.exp(-logit))
277
+ moop = np.where(spends, rng.gamma(1.0, 4_000, n_donor), 0.0)
278
+ cells_donor = _cells(rng, n_donor)
279
+ cells_recipient = _cells(rng, n_recipient)
280
+ return FoldData(
281
+ X_donor=X_donor,
282
+ moop_donor=moop,
283
+ cells_donor=cells_donor,
284
+ group_donor=cells_donor["group"].to_numpy(),
285
+ X_recipient=features(n_recipient),
286
+ cells_recipient=cells_recipient,
287
+ group_recipient=cells_recipient["group"].to_numpy(),
288
+ )
289
+
290
+
291
+ @settings(max_examples=8, deadline=None)
292
+ @given(seed=st.integers(0, 2**31))
293
+ def test_logit_family_rules(seed):
294
+ fold = _fold(seed)
295
+ out = logit_family(fold, 4, np.random.default_rng(seed))
296
+ prob = participation_probability(
297
+ fit_participation(fold.X_donor, fold.moop_donor), fold.X_recipient
298
+ )
299
+ cells = assign_cells(fold.cells_donor, fold.moop_donor, fold.cells_recipient)
300
+ means = cells.means()
301
+ below = prob < PARTICIPATION_CUTOFF
302
+ for method, draws in out.items():
303
+ assert draws.shape == (4, len(prob))
304
+ assert (draws >= 0).all(), method
305
+ assert np.allclose(out["logit_cell_mean"][0], np.where(below, 0, means))
306
+ assert np.allclose(out["logit_expected_value"][0], prob * means)
307
+ assert (out["logit_reg_mean"][:, below] == 0).all()
308
+ assert (out["logit_group_mean"][:, below] == 0).all()
309
+ assert (out["logit_threshold_hotdeck"][:, below] == 0).all()
310
+ for i, pool in enumerate(cells.pools):
311
+ for method in ("logit_draw_hotdeck", "logit_threshold_hotdeck"):
312
+ positive = out[method][:, i][out[method][:, i] > 0]
313
+ assert np.isin(positive, pool).all()
314
+ drawn = out["logit_draw_cell_mean"][:, i]
315
+ assert np.all((drawn == 0) | np.isclose(drawn, means[i]))
316
+
317
+
318
+ @settings(max_examples=10, deadline=None)
319
+ @given(seed=st.integers(0, 2**31))
320
+ def test_evaluate_invariants(seed):
321
+ persons, units = synthetic_spm(seed, n_units=80)
322
+ bump = np.random.default_rng(seed + 1).uniform(0, 10_000, len(units))
323
+ imputations = {
324
+ "no_moop": np.zeros((1, len(units))),
325
+ "higher": (units["moop"].to_numpy() + bump)[None, :],
326
+ "same": np.tile(units["moop"].to_numpy(), (3, 1)),
327
+ }
328
+ results = evaluate(persons, units, imputations)
329
+ groups = results["groups"]
330
+ methods = results["methods"]
331
+ nonempty = np.array(results["group_person_counts"]) > 0
332
+ assert np.allclose(np.array(methods["same"]["diff"])[nonempty], 0)
333
+ assert methods["same"]["newly_poor_share"] == 0
334
+ assert methods["same"]["newly_nonpoor_share"] == 0
335
+ assert (np.array(methods["no_moop"]["diff"])[nonempty] <= 1e-12).all()
336
+ assert (np.array(methods["higher"]["diff"])[nonempty] >= -1e-12).all()
337
+ band = groups.index(["pre_moop_ratio", RATIO_BANDS[0][2]])
338
+ if nonempty[band]:
339
+ for method in methods.values():
340
+ assert method["rate"][band] == pytest.approx(1.0)
341
+
342
+
343
+ @SETTINGS
344
+ @given(
345
+ ratios=st.lists(
346
+ st.one_of(
347
+ st.floats(-5, 20, allow_nan=False),
348
+ st.sampled_from([1.0, 1.25, 1.5, 2.0, 3.0]),
349
+ ),
350
+ min_size=1,
351
+ max_size=50,
352
+ )
353
+ )
354
+ def test_ratio_band_index_matches_band_masks(ratios):
355
+ ratios = np.array(ratios)
356
+ index = ratio_band_index(ratios)
357
+ for b, (low, high, _) in enumerate(RATIO_BANDS):
358
+ mask = (ratios >= low) & (ratios < high)
359
+ assert np.array_equal(index == b, mask)
360
+
361
+
362
+ def test_allocation_flags_cover_moop_and_premium():
363
+ persons = pd.DataFrame(
364
+ {
365
+ "I_MOOP": [-1, 0, 1, 2, 3, 0, 0],
366
+ "I_MCPREM": [-1, -1, -1, 0, -1, 2, 0],
367
+ }
368
+ )
369
+ flags = allocation_flags(persons)
370
+ assert flags["allocated_moop"].tolist() == [0, 0, 1, 1, 1, 0, 0]
371
+ assert flags["allocated_premium"].tolist() == [0, 0, 0, 0, 0, 1, 0]
372
+
373
+
374
+ def test_extract_verified_replaces_altered_copies(tmp_path):
375
+ zip_path = tmp_path / "release.zip"
376
+ with zipfile.ZipFile(zip_path, "w") as archive:
377
+ archive.writestr("a.csv", "x,y\n1,2\n")
378
+ archive.writestr("b.csv", "z\n3\n")
379
+ out = tmp_path / "data"
380
+ out.mkdir()
381
+ extract_verified(zip_path, out, ["a.csv", "b.csv"])
382
+ assert (out / "a.csv").read_text() == "x,y\n1,2\n"
383
+
384
+ # A matching copy stays in place.
385
+ os.utime(out / "b.csv", ns=(1, 1))
386
+ extract_verified(zip_path, out, ["b.csv"])
387
+ assert (out / "b.csv").stat().st_mtime_ns == 1
388
+
389
+ # Altered copies, of the same or a different size, come back.
390
+ (out / "a.csv").write_text("x,y\n9,2\n")
391
+ (out / "b.csv").write_text("z\n3\n4\n")
392
+ extract_verified(zip_path, out, ["a.csv", "b.csv"])
393
+ assert (out / "a.csv").read_text() == "x,y\n1,2\n"
394
+ assert (out / "b.csv").read_text() == "z\n3\n"
395
+
396
+
397
+ def _cache_inputs(n_units: int = 12):
398
+ rng = np.random.default_rng(0)
399
+ index = pd.Index(np.arange(n_units) * 3 + 1, name="SPM_ID")
400
+ units = pd.DataFrame(
401
+ {
402
+ "moop": rng.gamma(1.0, 1_000, n_units),
403
+ "household": np.arange(n_units) // 2,
404
+ },
405
+ index=index,
406
+ )
407
+ X = pd.DataFrame({"x": rng.normal(size=n_units)}, index=index)
408
+ cells = pd.DataFrame({"group": ["employer"] * n_units}, index=index)
409
+ folds = assign_folds(units["household"], 3, 0)
410
+ args = argparse.Namespace(
411
+ folds=3, draws=2, seed=0, skip_forests=False, max_households=None
412
+ )
413
+ return args, units, X, cells, folds
414
+
415
+
416
+ def _stub_fit(value: float, folds: np.ndarray, calls: list):
417
+ def fit(fold: int) -> dict[str, np.ndarray]:
418
+ calls.append(fold)
419
+ n = int((folds == fold).sum())
420
+ return {
421
+ "qrf": np.full((2, n), value),
422
+ "_timings": np.array([[0.0]]),
423
+ "_timing_names": np.array(["qrf"]),
424
+ }
425
+
426
+ return fit
427
+
428
+
429
+ def test_cache_misses_when_environment_or_microimpute_changes(tmp_path, monkeypatch):
430
+ args, units, X, cells, folds = _cache_inputs()
431
+ base_env = run_module.environment()
432
+ monkeypatch.setattr(run_module, "environment", lambda: base_env)
433
+ tag = config_hash(build_config(args, units, X, cells, folds))
434
+
435
+ calls: list = []
436
+ first, _ = collect_folds(tmp_path, tag, folds, 3, 2, _stub_fit(17.0, folds, calls))
437
+ assert calls == [0, 1, 2]
438
+ assert (first["qrf"] == 17.0).all()
439
+
440
+ # Same fingerprint: every fold comes from the cache.
441
+ calls.clear()
442
+ again, _ = collect_folds(tmp_path, tag, folds, 3, 2, _stub_fit(99.0, folds, calls))
443
+ assert calls == []
444
+ assert (again["qrf"] == 17.0).all()
445
+
446
+ # A new microimpute version, or edited microimpute source, refits.
447
+ newer = {**base_env, "microimpute": base_env["microimpute"] + ".post1"}
448
+ monkeypatch.setattr(run_module, "environment", lambda: newer)
449
+ new_tag = config_hash(build_config(args, units, X, cells, folds))
450
+ assert new_tag != tag
451
+ fresh, _ = collect_folds(
452
+ tmp_path, new_tag, folds, 3, 2, _stub_fit(99.0, folds, calls)
453
+ )
454
+ assert calls == [0, 1, 2]
455
+ assert (fresh["qrf"] == 99.0).all()
456
+
457
+ monkeypatch.setattr(run_module, "environment", lambda: base_env)
458
+ monkeypatch.setattr(run_module, "package_sha256", lambda package: "edited")
459
+ assert config_hash(build_config(args, units, X, cells, folds)) != tag
460
+
461
+
462
+ def test_cache_key_tracks_model_inputs_and_row_keys():
463
+ args, units, X, cells, folds = _cache_inputs()
464
+ tag = config_hash(build_config(args, units, X, cells, folds))
465
+ assert config_hash(build_config(args, units, X, cells, folds)) == tag
466
+
467
+ changed = units.copy()
468
+ changed.iloc[0, 0] += 1.0
469
+ assert config_hash(build_config(args, changed, X, cells, folds)) != tag
470
+
471
+ rekeyed = [frame.copy() for frame in (units, X, cells)]
472
+ for frame in rekeyed:
473
+ frame.index = frame.index + 1
474
+ assert config_hash(build_config(args, *rekeyed, folds)) != tag
475
+
476
+ assert config_hash(build_config(args, units, X, cells, folds[::-1])) != tag
477
+ assert frame_sha256(units) != frame_sha256(units.rename(columns={"moop": "m"}))
File without changes
File without changes