Skip to content

Add household-chunked microsimulation calculations - #10561

Merged
MaxGhenis merged 18 commits into
mainfrom
chunked-microsimulation-calculate
Nov 3, 2026
Merged

MaxGhenis merged 18 commits into
mainfrom
chunked-microsimulation-calculate

Conversation

@MaxGhenis

@MaxGhenis MaxGhenis commented Oct 31, 2026 •

Copy link
Copy Markdown
Contributor

Problem

Full-population taxation-of-benefits calculations can exceed a 20 GB working-memory budget because counterfactual income-tax calculations clone simulations. Callers need the three TOB variables and exact per-entity results across the whole population.

Interface and behavior

policyengine_us.tools.calculate_chunked.calculate_chunked(sim_factory, variables, period, n_chunks, *, dataset, household_chunks=None) returns ChunkedCalculation: entities maps entity names to source-order DataFrames (<entity>_id, weight, requested variables), and weighted_sums maps variable names to weighted totals.

The explicit dataset= keyword supplies the source without constructing a full-population simulation. The factory controls reforms, projection bounds, SPM configuration and other constructor options. Supported sources are US single/multi-year entity-table datasets, local/Hugging Face entity-table HDFStore files, and local/Hugging Face variable-centric .h5 files (<variable>/<period> groups, the population projection layout that Microsimulation(dataset=path) loads through Core). _is_hdfstore_format tells the layouts apart. For a variable-centric file the factory receives a str path to a temporary chunk file in the same layout, with every stored period, each array's dtype (fixed-width byte strings included) and the source file's name; one chunk file exists at a time and is deleted in the finally block, on failure too. Each array is assigned to its variable's entity in the country system; names that match no country-system variable are dropped, as Core's loader ignores them, except the membership role arrays Core also reads; the helper raises when the factory's system (for example a reform) defines a dropped name, since an unchunked Microsimulation(dataset=path, reform=...) would load it; a length that differs from the entity count is a ValueError. Flat files and period-less array files remain unsupported, with a clear error. All source periods and native IDs/weights remain unchanged; output weights match Microsimulation's projected household weights, promoted to float64 as MicroSeries does.

The inspected Core APIs had no raw entity-table partition helper. Simulation.subsample() starts with a full simulation and rescales weights, so it cannot supply these disjoint, unchanged-weight chunks.

Exactness proof and invariants

Every household receives one label. People inherit that household label, and each tax unit, SPM unit, family and marital unit must have members in exactly one household and receives the same label. Unique IDs, complete memberships and complete entity coverage are validated. Variable-centric files go through the same _household_positions code, on DataFrames built from their ID and membership arrays, and those arrays must be identical in every stored period; medicaid_claiming_tax_unit_id is checked in each stored period. Boolean selection preserves entity order; results are assigned to original row positions. Therefore chunks are disjoint, cover the population exactly, preserve whole groups, and present identical household-local inputs. Weighted totals reduce assembled source-order arrays, preserving MicroSeries precision and summation order. A chunk simulation is collected before the next factory call.

Four formulas reduce across the whole simulation: the Medicaid state reductions medicaid_slcsp_state_average_cost_index and medicaid_slcsp_state_denominator, and the rankers household_income_decile and spm_unit_income_decile (MicroSeries.decile_rank, which would rank within a chunk). Direct, dependent (for example income_decile), baseline and branch calculations of all four are rejected.

The guard overrides Core's _check_period_consistency(period, variable) instead of calculate and _calculate. In Core 3.33.2 and the now-locked 3.34.2 (simulation.py:1760) Simulation._calculate is its only caller, after the stored, frame and macro-cache lookups and before the formula, and it returns before any dependency starts; a monthly variable requested for a year and a yearly one requested for a month both reach it at their own period. The old overrides added two frames per variable in a chain, so a chain that fit a whole-file run could raise RecursionError only when chunked. Behaviour change, on purpose: a guarded variable with a stored value now returns that value, since no formula reduces across the chunk. The recursion limit is unchanged. An audit found no other formula that reduces across the population. The source scan below sees weight reads, MicroSeries, .count == tests and NumPy random draws; other whole-array reductions and .any()/.all() branches still need review by hand. One formula reads the simulation's size: _ssi_spouses.py sets possible_default = person.marital_unit.count == 1, so a chunk holding exactly one marital unit can pair SSI spouses differently from the whole file; such chunks are rejected when the file holds more than one marital unit, and every chunk is checked before the first is built. Known claiming-tax-unit links crossing households are also rejected. Custom reform formulas must remain household-local. Computed factory caches are cleared and macro-cache reads/writes are disabled for chunk datasets.

TOB sequencing

This branch carries #10566: TOB measures income_tax_chapter_1 (chapter 1 liability after nonrefundable and before refundable credits, excluding chapter 2A NIIT), and old-law inclusion is capped at actual inclusion. Two ordinary parent variable holders cache chapter 1 tax without Social Security and with min(pre-1993 inclusion, actual inclusion), the bound formed in the outer simulation before any branch exists. Old-law is prepared before no-benefits. Standard direct TOB requests prepare each state in a separate outer calculation, allowing Core's calculation frames to finish and branches to be collected before the next counterfactual. This preserves Core's dependency tracking. Supplied/cached/neutralized or reform-replaced TOB outputs skip unnecessary preparation. Nested TOB dependencies may retain branch identities until their outer formula returns.

The signed allocation matches main: OASDI is bounded old-law minus no-inclusion chapter 1 tax; total is current minus no-inclusion chapter 1 tax; HI is total minus OASDI. There is no per-unit floor.

What differs from main in the TOB formulas: (1) the two counterfactual states are computed once into the parent holders of income_tax_without_social_security and income_tax_with_pre_1993_social_security, instead of in branches inside each TOB formula; (2) branch clearing. Main deletes every variable whose name is not in branch.input_variables, for all periods, and keeps every period of an input-named variable. tools/tob.py calls drop_inherited_values, which keeps only the exact (variable, branch, period) keys set_input stored, so a formula-computed value of an input-named variable at another period is recomputed in the branch rather than inherited. Projected entity-table datasets store every year and don't reach this case; a single-period variable-centric file can. On the 32-household fixture, every result equals main's, and so does every record in the 2026 population pair (impact below). Both new holder names are rejected by the Social Security preflight when stored, since a stored value freezes that counterfactual state.

Validation and CI cost

Python coverage is required for factory/dataset APIs, original row ordering, arbitrary partitions, cache isolation and simulation lifetime. The new 32-household fixture has five states, multiple tax and SPM units, independently shuffled entity tables and Social Security recipients. Tests compare exact values and actual unchunked MicroSeries sums for 1, 2 and 3 chunks plus Hypothesis partitions: two explicit examples (every household in one chunk labelled 3, leaving six labels empty, and seven uneven, interleaved chunks of 1 to 11 households) and two generated ones (Hypothesis tries the one-chunk partition first, then one more, a fixed draw in CI). Every partition also checks each chunk's tables and the release of earlier chunk simulations. Additional cases cover reforms, automatic projection, HDFStore layouts, malformed partitions and branch collection. An independent calculation reproduces main's TOB branch algorithms (chapter 1 measure, bounded old-law inclusion) elementwise. The main witnesses in tob_revenue_oasdi.yaml, tob_revenue_total.yaml, tob_revenue_medicare_hi.yaml and tob_tier_split.yaml pass on the helper path; no cheaper separate agreement test exists, since each witness needs its own simulation and those files already run them.

Population-file tests write the same 32 households as a variable-centric file and compare 1, 2 and 3 chunks with an unchunked Microsimulation(dataset=path), and 2 chunks under a parameter reform with Microsimulation(dataset=path, reform=...), whose reformed system the two chunks share, elementwise and in weighted sums for the three TOB variables; the unchunked file run also equals the entity-table route. Validation tests cover an ignored unknown name, preserved periods, dtypes, byte strings and file name, one chunk file at a time, deletion after success and after a factory exception, and clear errors for a wrong length, a group spanning households, duplicate household IDs, IDs changing across periods, a missing membership, a crossing claiming tax unit, and flat or period-less files. A chunk holding exactly one of several marital units is rejected before any chunk is built, a file holding a single marital unit is accepted, and sparse labels below a huge n_chunks work. Each test that checks chunk folders uses its own temp folder. Real Simulation tests cover the guard on direct, dependency, branch, baseline and month/year routes for all four variables, and a stored guarded value. A frame-depth test measures the deepest Python stack for income_tax on one household: 361 guarded and unguarded (a restored calculate override gives 395). A source scan fails when a variables file that mentions MicroSeries, reads a weight variable by name or through .weights, tests .count ==, or draws random numbers through np.random or default_rng is neither a weight's own file, listed, nor guarded.

Repeated partition cases use the existing country parameter-lending clone path, preserving the complete unchanged policy history and copy-on-write protection for reforms. The retained default-policy object graph is frozen before creating any simulation, so repeated collections scan new simulations and branches rather than the unchanged global policy; weak-reference collection assertions remain active. The independent TOB comparison lends the tree only to its reference, main's algorithm, whose branches change no parameter; the helper's branches, which it checks, keep default policy cloning. It freezes the heap alive before it starts, as the chunking module does, so in CI's single process its collections no longer rescan the objects earlier tests left. The fixture supplies observed benefits, medical expenses and take-up inputs, including one positive SNAP amount, avoiding unrelated repeated eligibility calculations while the requested output variables retain their normal formulas.

The round-1 table below passed 73 new cases and 18 existing dataset/projection cases across capped selections. The PR now adds 113 new cases (validation 80, test_calculate_chunked.py 14, test_tob_sequencing.py 1, lifecycle 16, and 2 in test_social_security_preflight.py); all pass at f8ad24ebd1, one file per process (round 4 below). Ruff formatting and lint pass on the changed files.

Selection Passed Seconds including startup Peak RSS GiB
Chunk validation 46 1623.58 0.812
Deterministic chunking and dataset inputs 8 1109.27 1.231
Hypothesis partitions 1 (4 examples) See below Not recorded
Reform factory 1 995.07 2.225
Independent TOB sequencing 1 1417.90 2.167
TOB lifecycle 16 352.88 0.825
Existing per-capita dataset regressions 6 371.49 0.827
Existing role-input dataset regressions 12 432.18 0.868

Seconds run through pytest completion; interpreter shutdown adds time. Hypothesis passed before a combined property/reform process reached its 1800-second limit during reform; reform then passed alone. The deterministic, reform, independent sequencing, lifecycle, existing projection and role-input processes exited within their caps at 1179, 1062, 1450, 389, 433 and 446 seconds respectively.

Simulation/data tests route to the existing Microsimulation group; validation and lifecycle tests use the existing Core group. No CI runners or concurrent heavy processes were added. The round-1 selections ran one foreground file per process with a 4 GB RSS cap and four numeric threads; that cap applied to those local runs only. The round-3 review measured the four TOB YAML files cited below at 5.17 to 5.54 GB peak. CI's Microsimulation job took 54 min 57 s at c545c47877 (run 41556562467), where this PR's tests took 552.9 s by --durations, and 49 min 8 s at f8ad24ebd1 (run 41614839743), where they took 282.7 s.

Population files and guard move: before and after

Local Linux VM (4 vCPU, 16 GB), one file per process; seconds are pytest time, peak RSS is the largest child process. Before is 5da5df7893, after is 6835df5c68.

File Before: passed / s / peak MB After: passed / s / peak MB
tests/core/test_calculate_chunked_validation.py 46 / 20.9 / 839 74 / 66.1 / 1,529
tests/core/test_tob_branch_lifecycle.py 16 / 5.9 / 838 16 / 6.9 / 836
tests/microsimulation/test_tob_sequencing.py 1 / 180.4 / 2,198 1 / 161.0 / 2,198
tests/microsimulation/test_calculate_chunked.py 10 / 399.7 / 2,280 16 / 728.8 / 2,492

Round 3 head c545c47877, same VM and method, Core 3.34.2 after merging main: validation 77 passed / 152.7 s / 2,620 MB (the new reform-added-input case builds two small reformed simulations); lifecycle 16 / 6.6 s / 860 MB; test_calculate_chunked.py 16 / 904.8 s / 2,460 MB; test_social_security_preflight.py 57 / 5.6 s. Round 3 adds validation cases for a reform-defined stored input (raises), a chunk holding one of several marital units (raises, with household_chunks), and a non-HDF5 .h5 file (flat-file error); the validation fixture now has two people and two marital units per household.

Round 2's validation increase (839 to 1,529 MB) is one module-scoped reformed simulation for the real-guard tests. Round 3's, to 2,620 MB, came from the reform-input case, which built two reformed simulations; it now builds one and reuses its system in the file stub (round 4 below). The round-3 microsim increase is the population-file reference (unchunked and reformed) plus 1, 2, 3, 7 and reformed 2-chunk runs; round 4 below trims it. TOB YAML at c545c47877: tob_revenue_oasdi 13, tob_revenue_total 4, tob_revenue_medicare_hi 6, tob_tier_split 3, all passed (5.17 to 5.54 GB peak by the round-3 review); round 4 changes no formula or anything they load. CI at f8ad24ebd1: all 35 jobs of Pull request run 41614839743 succeeded, and so did the Documentation run 41614839518.

Round 4: Microsimulation CI time

CI's Microsimulation job (timeout-minutes: 60) took 54 min 57 s at c545c47877 (run 41556562467), and this PR's tests took 552.9 s of it by --durations. In round 4, the tests keep every kind of check, run some of them on fewer partitions (listed below), and build less:

  • The 7-chunk cases of the real-model and population-file tests are gone; three chunks catch order errors that two cannot. The real-model test's per-chunk checks (each chunk's source rows and order, release of earlier simulations, id coverage) now live in a helper that the Hypothesis property uses too, so they still run on seven chunks.
  • The Hypothesis property runs two generated partitions instead of four, plus two explicit ones: every household in one chunk labelled 3 (six empty labels skipped), and seven uneven, interleaved chunks of 1 to 11 households. About 13 chunk simulations are expected either way, so this pins the boundaries rather than saving time. In CI, Hypothesis's ci profile derandomizes, so the non-trivial generated partition is one fixed draw.
  • test_tob_sequencing.py: main's reference algorithm lends the unchanged parameter tree to its three branches, each of which paid an 8.2 s copy and a 3.1 s collection locally; the helper's branches keep ordinary clones. The heap alive before the test is frozen.
  • The reform-factory test reads unreformed income_tax from the module reference. The reformed population-file case keeps its unchunked reference on Microsimulation(dataset=path, reform=...), and its two chunks share that run's reformed system instead of loading the country system again (21.0 s each locally).

What the tests no longer cover: the default balanced 7-chunk split (5, 5, 5, 5, 4, 4, 4 households); the population file's chunk-file checks above three chunks (the validation file still writes one-household chunk files); two of the three non-trivial generated partitions; main's reference with full parameter copies; and, for the population file under a reform, a reformed system and reform baseline built per chunk (the entity-table reform-factory test still builds both and checks the baseline's release).

Local, same VM and method as above, with CI's update_itemization.sh setting applied and not committed; seconds are pytest time:

File Before c545c47877: passed / s / peak GB After f8ad24ebd1: passed / s / peak GB
tests/microsimulation/test_calculate_chunked.py 16 / 399.6 and 425.4 / 2.39 14 / 296.9 / 2.28
tests/microsimulation/test_tob_sequencing.py 1 / 99.7 and 103.8 / 2.07 1 / 43.4 and 42.9 / 2.06
tests/core/test_calculate_chunked_validation.py 77 / 104.1 / 2.50 80 / 93.1 / 2.10
tests/code_health/test_social_security_preflight.py 57 / 3.5 / 0.84
tests/core/test_tob_branch_lifecycle.py 16 / 5.3 / 0.83

Scaling each test's duration in run 41556562467 by its local after/before ratio projects this PR's tests at 343 s and the job at 51 min 27 s. That keeps the TOB test's old CI slowdown (166.7 s in CI, 103.7 s locally); without it, 318 s and 51 min 2 s. Measured in CI at f8ad24ebd1 (run 41614839743): the job took 49 min 8 s and this PR's tests 282.7 s. By test, CI before → after: TOB sequencing 166.7 → 39.9 s (32.8 s plus 7.1 s to collect and freeze the heap); population file under reform 81.8 → 60.7; Hypothesis property 73.1 → 54.0; reform factory 49.1 → 45.2; population-file module setup 42.0 → 4.7; the 7-chunk cases 63.4 → 0; everything else 76.8 → 78.1. Runner variance is large: the SPM step, which round 4 doesn't touch, took 1,085.8 s against 1,202.6 s at c545c47877.

That is 43 s over the 4-minute aim, and I found no further cut that keeps every kind of check. Narrowing options, for the hub to decide: max_examples=1 (no non-trivial generated partition, about 4 chunk simulations fewer on average); a 3-chunk ceiling for the property's strategy and explicit example (4 chunk simulations fewer, a narrower range); and the population-file reform reference on a cloned system instead of the default route (about 10 s locally). Together they would save roughly 50 s.

Memory

Measured before the #10566 carry-over and the population-file path, on the 2026 default dataset at 5da5df7893, in one process on a shared 18-CPU Mac, for tob_revenue_oasdi, tob_revenue_medicare_hi and tob_revenue_total:

Chunks Result Wall seconds Peak RSS (GB) Peak footprint (GB)
8 finished 3,128 7.16 4.75
4 finished 3,053 9.78 6.55
None stopped at the 20 GB memory ceiling, no result about 1,630 at the stop not recorded 21.5 at the stop

The 8- and 4-chunk runs agree to the cent: OASDI $64,926,259,711.73, HI $47,006,113,215.90, total $111,932,372,935.85. No unchunked full-population sum exists under that ceiling, so the small-fixture tests carry the comparison with an unchunked run. A run of the base formulas on the same eight chunks is queued, to compare base and head at full population.

impact: no output moves in 2026 or 2030. In each year, every record matches main's in all nine variables we compared, including the three TOB revenue variables, income_tax and household_net_income (2026 pair, 2030 pair).

axiom: n/a: microsimulation tooling, not a rule


Generated by Claude Code

@MaxGhenis MaxGhenis changed the title WIP: Add household-chunked microsimulation calculations Add household-chunked microsimulation calculations Oct 31, 2026
@MaxGhenis
MaxGhenis marked this pull request as ready for review October 31, 2026 19:02
MaxGhenis added a commit to PolicyEngine/crfb-tob-impacts that referenced this pull request Nov 1, 2026
The TOB variables run counterfactual income tax, which is too large for one
whole-population simulation on the full base. Target vectors now come from a
written year file through calculate_chunked, and each gamma probe writes a
temporary year file. This needs a release holding PolicyEngine/policyengine-us#10561.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@MaxGhenis

Copy link
Copy Markdown
Contributor Author

Picking this up for the US hub. I'll post a hand-off comment here when it's ready.


Generated by Claude Code

@MaxGhenis

Copy link
Copy Markdown
Contributor Author

Picking this up for the US hub, taking over from the cloud session that posted the pick-up comment on 2026-11-01 (it pushed nothing). I'll post a hand-off comment here when it's ready.


Generated by Claude Code

#10566 changed main's TOB formulas to difference income_tax_chapter_1 and
to cap old-law inclusion at actual inclusion. This branch had moved those
formulas onto cached counterfactual helpers, so taking either side of the
conflict would have dropped one change.

The shared branch helper now measures income_tax_chapter_1, the total
measures current-state chapter 1 liability, and the old-law helper forms
min(pre-1993 inclusion, actual inclusion) in the outer simulation before
any branch exists. The mixin prepares old-law before no-benefits so the
bound is read before either counterfactual. Main's TOB documentation and
references are kept, and the sequencing and lifecycle oracles follow the
new measure and order.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
calculate_chunked accepted only entity-table datasets, so a population
projection file, stored as <variable>/<period> arrays, failed with
"Dataset 'person' not found". Those files need chunking most: unchunked
TOB on one passed a 20 GB ceiling.

A path that _is_hdfstore_format rejects is now read as a variable-centric
file. Its ID and membership arrays go through the same _household_positions
rules, and must match in every stored period. Each array takes its
variable's entity; names that match no variable are dropped, as Core's
loader ignores them, and a wrong length is an error. Each chunk is written
to a temporary file in the same layout, under the source's name, and
deleted in the finally block.

The household and SPM unit income deciles rank over the whole simulation,
so in a chunk they ranked within it with no error. Both are now guarded.

The guard overrode calculate and _calculate, adding two frames per
variable in a chain, so a chain that fit a whole-file run could hit the
recursion limit only when chunked. It now overrides
_check_period_consistency, which Core calls after the cached-value lookups
and before the formula, and which returns before any dependency starts. A
guarded variable with a stored value now returns it.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
_is_hdfstore_format lets PyTables' own error escape for a file that is
not HDF5, so a CSV path failed with an HDF5 trace instead of the
flat-file message. The country simulation applies a reform more than
once, so the test readers use update_variable.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Three reformed chunks took 166 s; chunk counts are already covered
unreformed.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@MaxGhenis

Copy link
Copy Markdown
Contributor Author

Ready for hub review

Head: 6835df5c68cd6541b6a4c163f8151690ec30e66c (started from 5da5df7893f9e3c69f5a71ee2d5c41a72da5d349; nobody else pushed in between). Takes over from the earlier cloud session, which pushed nothing.

What was asked, and how it was resolved

Merge origin/main (normal merge commit, 3df8fb10), carrying #10566. The conflicts in tob_revenue_oasdi.py and tob_revenue_total.py are resolved with main's documentation, labels and references, on this branch's shared helpers. #10565 comes in with the rest of main.

  • policyengine_us/tools/tob.py:112: the shared branch helper measures income_tax_chapter_1; its docstring (:87) says what that measure is.
  • variables/gov/ssa/revenue/tob_revenue_total.py:31: the current state measures chapter 1.
  • variables/gov/ssa/revenue/income_tax_with_pre_1993_social_security.py:25: min_(pre-1993, actual) is formed in the outer simulation before the helper is called.
  • tools/tob.py:70-74: old-law is prepared before no-benefits.
  • tests/microsimulation/test_tob_sequencing.py: the oracle uses chapter 1 and bounded old-law inclusion, and the expected branch order is old-law then no-SS.
  • tests/core/test_tob_branch_lifecycle.py: the measured variable is now chapter 1, with the new preparation order.
  • Both helper variables' labels and docs, and docs/usage/chunked-calculations.md "TOB working memory", now describe chapter 1 and bounded inclusion, and link main's OASDI accounting page.
  • Body: impact: pending (hub).
  • Agreement with main on the tob_revenue_oasdi.yaml witnesses: there's no cheap separate test, since each witness needs its own simulation. Instead, main's TOB YAML files pass on the helper path: tob_revenue_oasdi 13/13 (528 s, 5.7 GB), tob_revenue_total 4/4, tob_revenue_medicare_hi 6/6, tob_tier_split 3/3. The sequencing test also re-implements main's algorithm independently and matches it elementwise on 32 households.

Variable-centric files (policyengine_us/tools/calculate_chunked.py):

  1. Layout detection, :126-150: any path that _is_hdfstore_format rejects goes to _VariableCentricFile. Judgment call: PyTables' own error for a non-HDF5 file (for example a CSV) escapes _is_hdfstore_format, so I catch it at :135-139 and give the flat-file error instead. system.py is unchanged.
  2. Validation (_VariableCentricFile.__init__, :277-366): reads every stored period of the six <entity>_id arrays and the five person_<group>_id arrays, requires them to be identical in each period, then calls the unchanged _household_positions on DataFrames built from them (:360, _tables). medicaid_claiming_tax_unit_id is checked in every stored period.
  3. Each array takes the entity of its variable in system.variables. Unknown names are dropped. A wrong length raises ValueError: <name>/<period> has shape ..., but <entity> has N rows. Judgment call: person_<group>_role and role are not variables, but Core's build_from_dataset reads them as roles, so they are kept as person arrays.
  4. chunk(), :381-398: writes <tmp>/policyengine-us-chunk-*/<source name> with every period and dtype=source.dtype, so fixed-width and variable-length byte strings are kept. The factory receives str(path). release() at :400 runs in the loop's finally (:547), after the simulation is collected, on failure too. The source stays on disk and is re-read once per chunk, so memory is bounded by one array rather than the file.
  5. The rest of the loop is unchanged: the ID/order check, _guard, source-order tables and weighted_sums. The entity-table path (_EntityTables, :208-265) is a move of the old code; its tests pass unchanged.
  6. Flat files (no .h5 suffix, which Core reads as CSV) and period-less array files (ARRAYS) raise clear errors.

The docstring now says the factory receives a path for this layout. Docs :27-50 and the body's supported-sources sentence are corrected.

Rankers: _POPULATION_AGGREGATES (:27) now holds all four variables. The comment above it is rewritten, and the variables are named in the docs (:71) and the body. No variables changed.

Guard off the recursion path (:37-53): both overrides are removed, and _check_period_consistency(self, period, variable) checks variable.name and then calls super(). I confirmed in Core 3.33.2 (the locked version) that simulation.py:1708 is its only caller, after the cached-array, frame-recall and macro-cache returns and before defined_for and the formula, and that it returns before any dependency starts. A monthly variable asked for a year goes through calculate_add, and a yearly one asked for a month goes through calculate_divide or this_year; both re-enter _calculate at the variable's own period and reach the hook. The tests cover both. Error text and _guard are unchanged, and the recursion limit is not raised. The behaviour change (a stored value returns) is stated in the docs and the body.

Audit: I confirmed the four population formulas from the code: state_aggregate_helpers.sum_by_state, and decile_rank in both deciles. I found two more things and added neither to the guard:

  • variables/gov/ssa/ssi/eligibility/resources/deemed/_ssi_spouses.py:35, possible_default = person.marital_unit.count == 1, reads the size of the whole simulation. A chunk with exactly one marital unit can pair SSI spouses differently from the whole file. The docs and body now say so.
  • tax_unit_roles_supplied.py:19,28 uses .any()/.sum() only in error messages. In a chunk only the count in the message differs.

Tests (one file per process, local 4 vCPU VM)

File Before (5da5df78): passed / s / peak MB After: passed / s / peak MB
tests/core/test_calculate_chunked_validation.py 46 / 20.9 / 839 74 / 66.1 / 1,529
tests/core/test_tob_branch_lifecycle.py 16 / 5.9 / 838 16 / 6.9 / 836
tests/microsimulation/test_tob_sequencing.py 1 / 180.4 / 2,198 1 / 161.0 / 2,198
tests/microsimulation/test_calculate_chunked.py 10 / 399.7 / 2,280 16 / 728.8 / 2,492

Most of the microsim increase is the variable-centric reference: unchunked and reformed, 64 s in setup. Then come the exactness runs for 1, 2, 3 and 7 chunks (11, 23, 35 and 74 s) and the reformed 2-chunk run; 3 reformed chunks took 166 s, so I cut the reform case to 2 chunks. Most of the validation increase is one module-scoped reformed Microsimulation on two households for the real-guard tests. ruff format --check and ruff check are clean. No full suites and no microsimulation on population data were run.

New tests:

  • validation :133, :188, :287, :305, :319, :531, :570, :604, :612, :620, :627, :636, :660;
  • microsim :610, :642, :676.
    They cover:
  • exact 1/2/3/7 chunks and the reform, against an unchunked Microsimulation(dataset=path) for the three TOB variables, and agreement with the entity-table reference;
  • an ignored unknown name, written at the wrong length on purpose;
  • S6 role byte strings, periods, dtypes and the file name kept;
  • one chunk file at a time, with no leftover after success or after a factory exception;
  • errors for a wrong length, a group spanning households, duplicate household IDs, a person in an unknown household, IDs changing across periods, a missing membership, a crossing claiming unit, and flat or period-less files;
  • real Simulation guard routes for all four variables: direct, a reform-added reader, month-for-year, branch, branch reader and baseline, plus income_decile and a monthly reader asked for a year and a month;
  • a stored guarded value is returned;
  • the frame-depth test;
  • the source scan;
  • the stub tests, which now call the hook as Core does, parametrized over all four variables.

Frame depth (deepest Python stack, income_tax, one household): 361 unguarded, 361 guarded. With the old-style calculate override restored it is 395 (+34), and the test allows at most +3.

Mutants (each in its own process, code restored after each, clean tree confirmed)

  1. Last household dropped from the person mask: test_variable_centric_chunks_keep_periods_dtypes_and_name fails.
  2. tax_unit arrays written with the household mask: test_variable_centric_file_is_elementwise_exact[2] fails. The 2-household validation fixture can't catch this, because its masks coincide.
  3. Temporary file not deleted: test_variable_centric_chunks_keep_periods_dtypes_and_name and test_variable_centric_chunk_file_is_deleted_after_factory_error fail.
  4. spm_unit_income_decile removed from the set: test_guard_covers_known_population_formulas, test_core_hook_guards_readers_across_periods and test_population_formulas_are_guarded_or_listed fail.
  5. A calculate override calling super() restored: test_guard_adds_no_frames_per_dependency fails (361 vs 395).
  6. _check_period_consistency override removed: 22 tests fail, including every test_rejects_population_aggregates_in_dependency_graph case, every test_core_hook_guards_every_route case, test_core_hook_guards_readers_across_periods and test_guarded_variable_with_stored_value_returns_it.

CI

Pull request run 41481260553 at 6835df5c68 was still in progress 60 minutes after the push (17:24Z), so it has no conclusion yet. Every job that finished passed, including Quick Feedback, PR validation (lock, guard tests, lint, changelog), Package and compatibility, Household API Partners, all Rest Python shards, Full Suite - Microsimulation, and every Baseline shard. Two jobs were still running: Contrib other-shard-1 and other-shard-3, both started at 16:56Z. Documentation run 41481259902 started alongside. No partner test file changed.

Other judgment calls and limits of the new path

  • Entity assignment uses the country system's variables, as the brief asks. A variable that only a reform adds and that is stored in the file is dropped from chunks.
  • The source file is re-read for every chunk, which costs I/O but keeps memory bounded.
  • Chunk files go to the system temp directory, which needs room for about one chunk.
  • .h5 suffix only: Core treats any other suffix as CSV.
  • medicaid_claiming_tax_unit_id IDs are validated with the same rule as entity tables, per period.
  • The changelog fragment now mentions variable-centric population files.
  • The body's Memory section is kept and marked as measured before the carry-over and this change.

Proposed follow-ups

  • Make SSI spouse defaults independent of simulation size. _ssi_spouses.py:35 tests person.marital_unit.count == 1 across the whole simulation, so household chunks, or any subsample with one marital unit, can pair spouses differently. Decide whether this should be a per-situation flag instead.
  • Reject stored TOB counterfactual helpers in the Social Security preflight. The preflight rejects a stored income_tax_chapter_1. A stored income_tax_without_social_security or income_tax_with_pre_1993_social_security would freeze a counterfactual the same way.

Generated by Claude Code

@MaxGhenis

Copy link
Copy Markdown
Contributor Author

Ready for hub review
IMPACT PAIR: could not run
head 6835df5
base 0ea8894

The pair stopped before computing anything, because the measuring environment could not download the default dataset. Its egress proxy refused the connection to huggingface.co:443 (HTTP 403 on CONNECT, a network-policy denial). Neither side produced outputs, so this comment gives no numbers and no verdict on the diff.

  • Dataset: hf://datasets/policyengine/populace-us/populace_us_2024.h5@populace-us-2024-spm-20260909. DEFAULT_DATASET and DEFAULT_DATASET_SHA256 (6496cc4393d4d3c6574f76eca231de5898c803b9067645591fd5c4d3e65aee84) are identical in both trees. No file was downloaded, so no local SHA-256 exists to check.
  • Year 2026, 8 chunks.
  • git diff <base> <head> -- pyproject.toml uv.lock is empty, so both sides used the head's uv sync environment (Python 3.14), with PYTHONPATH choosing the tree.
  • Base side: exited with status 1 after 53 s wall time, at a maximum resident set size of 0.86 GB. It failed in huggingface_hub.hf_hub_download with httpx.ProxyError: 403 Forbidden. The head side was not run, because it needs the same download.
  • No variables were dropped.

To rerun, use an environment whose network access allows huggingface.co (and its file CDN), or a machine that already has the file cached.

Command:

export OMP_NUM_THREADS=4 OPENBLAS_NUM_THREADS=4 MKL_NUM_THREADS=4
PYTHONPATH=$PWD/../base /usr/bin/time -v ../head/.venv/bin/python ../pair_side.py ../head/policyengine_us/tools/calculate_chunked.py 8 2026 ../out/base > ../out/base.log 2> ../out/base.err
pair_side.py (as specified in the brief, unchanged)
"""One side of #10561's pair: python pair_side.py <helper.py> <chunks> <year> <out_dir>."""
import importlib.util
import json
import sys
import time
from pathlib import Path

helper, chunks, year, out = sys.argv[1], int(sys.argv[2]), int(sys.argv[3]), Path(sys.argv[4])
VARIABLES = (
    "tob_revenue_oasdi",
    "tob_revenue_medicare_hi",
    "tob_revenue_total",
    "income_tax_chapter_1",
    "income_tax",
    "state_income_tax",
    "social_security",
    "household_net_income",
    "household_benefits",
)
spec = importlib.util.spec_from_file_location("hub_calculate_chunked", helper)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)

import policyengine_us  # noqa: E402
from policyengine_us import Microsimulation  # noqa: E402

start = time.monotonic()
result = module.calculate_chunked(
    lambda dataset: Microsimulation(dataset=dataset, dataset_end_year=year),
    VARIABLES,
    year,
    chunks,
    dataset=Microsimulation.default_dataset,
)
out.mkdir(parents=True, exist_ok=True)
for entity, table in result.entities.items():
    table.to_pickle(out / f"{entity}.pkl")
summary = {
    "package": str(Path(policyengine_us.__file__).parent),
    "chunks": chunks,
    "year": year,
    "weighted_sums": result.weighted_sums,
    "wall_seconds": time.monotonic() - start,
}
(out / "summary.json").write_text(json.dumps(summary, indent=1, sort_keys=True))
print(json.dumps(summary, sort_keys=True), flush=True)
compare.py (written, not run)
"""Compare #10561 pair outputs: python compare.py <base_dir> <head_dir>."""
import json
import sys
from pathlib import Path

import numpy as np
import pandas as pd

base, head = Path(sys.argv[1]), Path(sys.argv[2])
bs = json.loads((base / "summary.json").read_text())
hs = json.loads((head / "summary.json").read_text())
print("base package", bs["package"]); print("head package", hs["package"])
rows, samples = [], []
entities = sorted(p.stem for p in base.glob("*.pkl"))
assert entities == sorted(p.stem for p in head.glob("*.pkl")), "entity sets differ"
for entity in entities:
    b = pd.read_pickle(base / f"{entity}.pkl")
    h = pd.read_pickle(head / f"{entity}.pkl")
    idc = f"{entity}_id"
    assert idc in b.columns and idc in h.columns, f"{idc} missing"
    assert len(b) == len(h) and np.array_equal(b[idc].values, h[idc].values), f"{idc} order differs"
    print(f"{entity}: {len(b)} rows, {idc} identical and same order; columns {list(b.columns)}")
    for var in b.columns:
        if var == idc or var.endswith("_weight"):
            continue
        bv = b[var].to_numpy(dtype=float); hv = h[var].to_numpy(dtype=float)
        diff = np.abs(hv - bv)
        nan_mismatch = np.isnan(bv) != np.isnan(hv)
        diff = np.where(np.isnan(diff), 0.0, diff)
        differ = (diff > 1e-6) | nan_mismatch
        rows.append(dict(variable=var, entity=entity,
                         base_sum=bs["weighted_sums"][var], head_sum=hs["weighted_sums"][var],
                         difference=hs["weighted_sums"][var] - bs["weighted_sums"][var],
                         rows_differ=int(differ.sum()), max_abs_diff=float(diff.max())))
        for i in np.flatnonzero(differ)[:20]:
            samples.append((var, entity, int(b[idc].iloc[i]), bv[i], hv[i]))
order = [v for v in bs["weighted_sums"]]
print("| variable | entity | base sum | head sum | difference | rows that differ | largest abs diff |")
print("|---|---|---:|---:|---:|---:|---:|")
for r in sorted(rows, key=lambda r: order.index(r["variable"]) if r["variable"] in order else 99):
    print(f"| `{r['variable']}` | {r['entity']} | {r['base_sum']:,.2f} | {r['head_sum']:,.2f} | {r['difference']:,.6g} | {r['rows_differ']} | {r['max_abs_diff']:.6g} |")
print("weighted sums identical:", bs["weighted_sums"] == hs["weighted_sums"])
print("samples:", samples[:20])

Generated by Claude Code

@MaxGhenis

Copy link
Copy Markdown
Contributor Author

Hub impact pair, progress

The base side finished. The head side runs next, on the same chunks and in the same environment, and a result comment follows.

  • Base 0ea88946c6584d59057d1dac23bdd8a85cb6ea6a, the merge base of head 6835df5c68cd6541b6a4c163f8151690ec30e66c with main. The head's calculate_chunked ran the base formulas on 8 household chunks for 2026, in the head's uv sync environment (Python 3.14.6, policyengine-core 3.33.2). pyproject.toml and uv.lock don't differ between the two commits.
  • Dataset: hf://datasets/policyengine/populace-us/populace_us_2024.h5@populace-us-2024-spm-20260909, 57,240 households. The downloaded file's SHA-256 equals DEFAULT_DATASET_SHA256 (6496cc4393d4d3c6574f76eca231de5898c803b9067645591fd5c4d3e65aee84).
  • Wall time: 4,140 s inside calculate_chunked, and 1:09:37 for the whole process.
  • Maximum resident set size: 10,226,828 kB (9.75 GiB), under the 14 GB limit that would call for a 16-chunk rerun.
  • The helper accepted all nine variables, and Microsimulation accepted dataset_end_year.
Variable Base weighted sum, 2026
tob_revenue_oasdi 63,221,571,702.76
tob_revenue_medicare_hi 46,005,803,316.51
tob_revenue_total 109,227,375,027.48
income_tax_chapter_1 2,633,056,572,090.40
income_tax 2,595,105,852,713.51
state_income_tax 578,044,599,346.72
social_security 1,677,873,892,455.98
household_net_income 17,161,633,654,134.80
household_benefits 2,059,999,418,114.87

Generated by Claude Code

@MaxGhenis

Copy link
Copy Markdown
Contributor Author

Ready for hub review
VERDICT: REQUEST_CHANGES
6835df5
METHODOLOGY_CHOICE: no

Independent review of the round-2 head. The head hasn't moved since 6835df5c68. I found no earlier "Hub review in progress" comment, so this review starts from scratch. Every number below comes from a run in this session: a 4 vCPU / 16 GB Linux VM, Python 3.14.6, the locked policyengine-core 3.33.2, one process at a time.

Findings

P1. A variable that a reform adds and that the file stores is dropped from chunks with no error

policyengine_us/tools/calculate_chunked.py:306-313. _VariableCentricFile.__init__ assigns each stored array an entity by looking it up in the baseline system.variables, and it skips any name it doesn't find there. Core's loader checks the simulation's own system instead. Simulation.__init__ builds that system with the reform (default_tax_benefit_system(reform=reform)) before build_from_dataset runs (simulation.py:1178, if variable in self.tax_benefit_system.variables). So for Microsimulation(dataset=path, reform=R), unchunked loads a reform-added input from the file, and every chunk file leaves it out. The entity-table route copies every column, so it keeps such inputs. The two layouts therefore disagree as well.

Witness: a 4-household variable-centric file storing reform_only_input/2026. The reform adds reform_only_input and reform_only_output = 2 * reform_only_input.

unchunked [ 200.  400.  600. ... 3200.]
chunked   [0. 0. 0. ... 0.]
weighted sums: unchunked 117813.19994926453 chunked {'reform_only_output': 0.0}
DIFFER

The hand-off calls this a limit ("A variable that only a reform adds and that is stored in the file is dropped from chunks"). The docs page and the body only say that names matching no variable "are left out, as Core's loader ignores them". That statement is false for reformed factories. The helper should raise. Only the factory's simulation knows the reformed system, so check after the factory returns:

# _VariableCentricFile.__init__: replace the `else: continue` branch
                else:
                    # Core's loader ignores names that match no variable in
                    # the simulation's system; a reform can add one.
                    self.unmatched.append(name)
                    continue

Set self.unmatched = [] next to self.arrays = {}, add unmatched = () as a class attribute on _EntityTables, and after _guard(simulation) in calculate_chunked:

            dropped = sorted(
                set(source.unmatched) & set(simulation.tax_benefit_system.variables)
            )
            if dropped:
                raise ValueError(
                    f"{', '.join(dropped)} are stored in the file and defined by "
                    "the factory's system but not the country system, so chunk "
                    "files cannot carry them. Store them in an entity-table "
                    "dataset or add the variables to the country system."
                )

Then change the docs sentence (docs/usage/chunked-calculations.md, "names that match no variable are left out, as Core's loader ignores them") and the matching body sentence to say the helper raises when the factory's system defines a dropped name. A test with the reform above would kill the old behaviour.

P2. A chunk holding exactly one marital unit pairs SSI spouses differently: a silent wrong SSI result

policyengine_us/variables/gov/ssa/ssi/eligibility/resources/deemed/_ssi_spouses.py:35, possible_default = person.marital_unit.count == 1, measures the whole simulation. When it is true, a two-person marital unit in one tax unit where exactly one member is a dependent is ruled out as a couple (claims, line 50). Readers: is_ssi_spousal_resource_deeming_applies, ssi_resources_deemed_from_ineligible_spouse and ssi_resources_deemed_from_ineligible_parent, and through them ssi; is_ssi_resource_deeming_child; NH SSP (_nh_ssp_couple.py), IA (ia_ssa_has_spouse.py), NY (ny_ssp_is_minor.py). Through ssi it also reaches household_benefits, household_net_income, SPM resources and SSI-linked Medicaid categories. TOB reaches it only through benefit-dependent income inputs. TOB's taxable Social Security doesn't read SSI, so the fixture's TOB values don't move.

Witness: household 1 holds a 70-year-old and a 40-year-old in one marital unit and one tax unit. The 40-year-old has is_tax_unit_dependent = True and $100,000 of ssi_countable_resources. Household 2 holds one 30-year-old. household_chunks=[0, 1]:

unchunked ssi [0. 0. 0.]
unchunked is_ssi_spousal_resource_deeming_applies [ True False False]
chunked   ssi [11928.     0.     0.] {'ssi': 11928.0}
chunked   is_ssi_spousal_resource_deeming_applies [False False False]

The docs and the body disclose this. But the helper accepts it silently, and a cheap exact guard exists. The formula changes only when a chunk has exactly one marital unit while the file has more than one. With the default balanced split on a population file, that can't happen unless n_chunks approaches the household count. Arbitrary household_chunks can trigger it, and so can a small file. I'd add the guard here and leave fixing the formula to a follow-up. In the loop, after masks is built:

        if masks["marital_unit"].sum() == 1 < len(masks["marital_unit"]):
            raise ValueError(
                "A chunk holding exactly one marital unit can pair SSI spouses "
                "differently from the whole file (_ssi_spouses.py); use fewer "
                "chunks or different household_chunks."
            )

Then replace "this is not guarded" in the docs and the body with "such chunks are rejected". As for the fixture: its smallest chunk (7 chunks of 32 households) holds 4 or 5 households, so no existing test reaches this case.

P2. The new counterfactual holders can be stored and freeze TOB; the Social Security preflight doesn't catch them

policyengine_us/tools/social_security_preflight.py:47-63 (CALCULATED_TAX_LEAVES). This PR adds two ordinary variables, income_tax_without_social_security and income_tax_with_pre_1993_social_security. Core loads a stored value for either one as an input, and TOB then differences against it. On main these names match no variable, so Core ignores them in a file. That makes this a defect the PR makes reachable, so it belongs here rather than in a follow-up. The preflight already rejects a stored income_tax_chapter_1 for exactly this reason ("a stored value freezes every tax state they measure and zeroes TOB"). The new names pass it.

Witness: a 3-household variable-centric file, plain Microsimulation(dataset=path):

preflight rejects stored income_tax_chapter_1 True
preflight rejects stored income_tax_without_social_security False
preflight rejects stored income_tax_with_pre_1993_social_security False
clean tob_revenue_total [   0.   5572.4     0.   5770.7     0.    336.75]
stored zero helper -> tob_revenue_total [2924.   7222.4  2780.   7570.7     0.    336.75]

With the stored zero, TOB total equals the whole chapter 1 tax. Replacement: add both names to CALCULATED_TAX_LEAVES, below "income_tax_chapter_1",:

        # The TOB counterfactual holders: a stored value freezes that state.
        "income_tax_without_social_security",
        "income_tax_with_pre_1993_social_security",

calculate_chunked could also refuse a source that stores either name. The preflight line is the minimum.

P3. A non-HDF5 file with a .h5 suffix gives a raw h5py error

calculate_chunked.py:135-141, 298. A .csv path gets the flat-file ValueError. A CSV saved as pop2.h5 raises OSError: Unable to synchronously open file (file signature not found). That isn't a wrong answer, but it isn't the "clear error" the brief asks for. Wrap the h5py.File(self.path, "r") open in try/except OSError and raise the flat-file message.

P3. Stale or overstated text

  • docs/usage/chunked-calculations.md, last paragraph: "The hub will add measured full-dataset peak memory and timing after running its memory probe." The body now carries those measurements. Replace it with "The PR records measured full-dataset peak memory and timing; small-fixture exactness establishes correctness, not an RSS bound."
  • Body, "An audit found no other formula that reduces across the population. One formula reads the simulation's size". I agree with the audit (see below). But the source scan test_population_formulas_are_guarded_or_listed looks only for MicroSeries, *_weight and .weights. It would not catch a second .count == test like _ssi_spouses.py:35, or a seeded draw. Say so in the test docstring, or extend _POPULATION_SIGNAL with \.count\s*== and np\.random|default_rng, listing _ssi_spouses.py until it's fixed.
  • "No formula changes", as a reading of the TOB restructure: the branch clearing isn't main's. Main deletes every variable not in branch.input_variables, a list of names, all periods. tools/tob.py:106 calls drop_inherited_values, which keeps only the exact (variable, branch, period) keys set_input stored. The two differ when a stored-input variable also holds a formula-computed value for another period. Main keeps that value in the branch and the head recomputes it. Projected entity-table datasets store every year, so they don't hit this, and the fixture comparison below is exact. A single-period variable-centric file could. The body's impact: pending (hub) is the right line. The hub pair (base formulas against head, same chunks) is what settles it at population.

Checked and sound

  • No regression against main on the fixture. I ran the 32-household fixture with the merge base 0ea88946's formulas and with the head's, in the same environment. Both used the head's calculate_chunked.py, loaded by path. Each side ran the baseline and a reform (threshold.base.main.JOINT = $20,000), unchunked and in 3 chunks. All 24 arrays and sums are identical:

    OASDI HI Total
    baseline, unchunked = 3 chunks, base = head 223,525.76135742664 147,697.79816761817 371,223.5595250448
    reform, unchunked = 3 chunks, base = head 223,525.76135742664 198,148.5367371195 421,674.29809454613

    The reform moves only total and HI, as the bound predicts: old-law inclusion is capped at actual inclusion, and a lower base threshold raises actual inclusion only above the pre-1993 amount. Peak RSS was 4.55 GB for base and 3.08 GB for head. The base side passed my 4 GB budget; I report that rather than hide it.

  • Merge with Measure benefit-taxation revenue as chapter 1 liability #10566 (git show --remerge-diff 3df8fb10). tools/tob.py:112 measures income_tax_chapter_1. tob_revenue_total.py:31 measures chapter 1 in the current state. income_tax_with_pre_1993_social_security.py:25-28 forms min_(pre-1993, actual) in the outer simulation, inside the helper variable's formula, before income_tax_with_social_security_override creates a branch. The no-benefits helper never reads the bound, and every branch is popped before the next one starts. So whatever the order, each state reads the same outer arrays: the old-law-then-no-benefits order can't change a value. The fixture comparison above confirms it. test_tob_sequencing.py (an independent copy of main's algorithm) passes: 1 passed, 126 s, 2.20 GB.

  • Variable-centric edge cases (a 4-household file, apart from P1). IDs that change across stored periods raise "Entity IDs and memberships must be identical across stored periods". So does a permuted person_tax_unit_id in one period. A person array one row short raises employment_income_before_lsr/2026 has shape (15,), but person has 16 rows. A role array stored as S6 in 2026 and S16 in 2027 keeps both dtypes in the chunk, and a variable-length UTF-8 string array keeps string_info(encoding='utf-8'). A factory exception leaves no policyengine-us-chunk-* directory, and neither does any of the errors above. A .csv path gets the flat-file error. Rows out of household order: boolean masks keep source order, the ID check at :500 compares chunk IDs with original_ids[mask], and the fixture's entity tables are shuffled independently. Equal IDs stored as float in a later period pass, which is harmless: Core reads only the first period (get_eternity_array).

  • Guard hook. In Core 3.33.2, simulation.py:1708 is the only call to _check_period_consistency. It comes after the stored-array, frame-recall and macro-cache returns and after the month/year alternate handling, and before defined_for and _run_formula. _run_formula is called only at :1738. The month-for-year and year-for-month routes re-enter _calculate at :1692/:1699 (via calculate_add/calculate_divide), at the variable's own period. Main has since locked Core 3.34.2 (Lock policyengine-core 3.34.2 #10589). I fetched that wheel from PyPI: the hook is still the only caller, at simulation.py:1760, in the same place, so the guard survives the next merge with main. A stored value of a guarded variable returning is sound. The chunk file holds the source's own stored rows for that chunk, which is exactly what the unchunked run returns for those rows, so no chunk uses a value it couldn't have read whole. A guarded variable at a period with no stored value still reaches the formula and raises.

  • Frame-depth test. The guarded and unguarded depths are measured in the same interpreter, so their difference is structural. The hook frame exists only while the check runs, at a leaf. CI runs the microsimulation group on Python 3.14 only, and the test passed there and here. The depths[False] > 300 floor is the only version-sensitive part. A failure there would show as a test failure, not a wrong guard.

  • Population audit (round 1's check 3). I grepped policyengine_us/variables, reforms and tools for MicroSeries, decile_rank, quantiles, ranks, np.random, default_rng, weight reads, .count ==, whole-array np.sum/np.unique/np.sort, and conditionals on .any()/.all(). Population reductions: the two Medicaid state aggregates (state_aggregate_helpers.sum_by_state) and the two decile_rank formulas. These four are exactly _POPULATION_AGGREGATES. One simulation-size read: _ssi_spouses.py:35 (P2 above). The .any() branches I read are short cuts whose result doesn't depend on the population: county.py, spm_unit_allocated_tenant_payment, weekly_hours_worked LSR, substitution_elasticity, snap_utility_region. tax_unit_roles_supplied uses .sum() only in error text. ssi_resources_deemed_from_ineligible_parent.py:91 pools rows by sorted keys within the array, which is order-independent. LSR earnings deciles use fixed markers (substitution_elasticity.py:42), not ranks. Neither parameters/ nor reforms/ turned up anything.

  • Tests at this head, one file per process: test_calculate_chunked_validation.py 74 passed (59 s, 1.53 GB); test_tob_branch_lifecycle.py 16 passed (5 s, 0.84 GB); test_tob_sequencing.py 1 passed (126 s, 2.20 GB); test_calculate_chunked.py 16 passed (592 s, 2.44 GB). CI: Pull request run 41481260553 concluded success at this head, with all 35 jobs successful and none skipped or cancelled. The microsimulation job log shows every new case collected and passing. That job took 46 min, against 42-46 min for four other PRs' microsimulation jobs the same evening.

  • Mutants, each in its own process on a scratch worktree, which git status showed clean afterwards:

    1. Drop each chunk's last household from the person mask: killed by test_variable_centric_chunks_keep_periods_dtypes_and_name and test_skips_empty_chunks_and_preserves_native_and_projected_weights.
    2. Write tax_unit arrays with the household mask: killed by test_variable_centric_file_is_elementwise_exact[1,2,3,7]. The validation file passes, 74/74.
    3. Skip deleting the chunk directory: killed by test_variable_centric_chunks_keep_periods_dtypes_and_name and test_variable_centric_chunk_file_is_deleted_after_factory_error.
    4. Remove spm_unit_income_decile: killed by test_guard_covers_known_population_formulas, test_core_hook_guards_readers_across_periods and test_population_formulas_are_guarded_or_listed.
    5. Restore a calculate override that calls super(): killed by test_guard_adds_no_frames_per_dependency.
    6. Remove the _check_period_consistency override: killed by every test_rejects_population_aggregates_in_dependency_graph case, among others. I capped the output at 8 lines, so I didn't recount the hand-off's 22.
    7. (new) Write spm_unit arrays with the household mask: killed by test_variable_centric_file_is_elementwise_exact[1,2,3,7]. The validation file passes, 74/74.
    8. (my own) Skip the cross-period identity check (stored_periods[1:1]): killed by test_variable_centric_rejects_ids_that_change_across_periods[household_id] and [person_spm_unit_id].

    The 2-household validation fixture has one row per entity per household, so every mask is identical and it can't catch a mask-substitution error (mutants 2 and 7). The 32-household microsim fixture catches both: its even households hold two SPM units, and every household holds two tax units and three marital units. The reformed case's 2 chunks are enough. The code path is reform-independent, so the baseline 3- and 7-chunk cases cover order errors (a 3-chunk cyclic shift, say) that 2 chunks can't tell from a swap.

  • Readers of changed variables: tob_revenue_medicare_hi, oasdi_benefit_taxation_revenue (adds = ["tob_revenue_oasdi"]), oasdi_income, and the S.770 revenue variables. In these the TOB call is nested, so the mixin's preparation doesn't run. The formulas then run old-law then no-benefits in formula order, with the same values as above. CI's YAML for them passed.

  • Population totals. The pair's base side (OASDI $63,221,571,702.76, HI $46,005,803,316.51, total $109,227,375,027.48) already measures chapter 1. The earlier $64.926bn / $47.006bn at 5da5df78 measured income_tax (before-refundable tax, NIIT included per Measure benefit-taxation revenue as chapter 1 liability #10566, less refundable credits), so those totals were expected to move and don't compare. The head side must match base to the cent on all nine variables, the three TOB outputs included. Both sides use the same chunks and main's measure and bound, and nothing else in the PR touches income_tax_chapter_1, income_tax, state_income_tax, social_security, household_net_income or household_benefits. Any TOB difference would come from the branch-clearing difference noted in P3. The pair chunks both sides identically, so it can't see chunk-only effects like P1 and the SSI case. Only fixture tests cover those.

  • Scope. Nothing under .github/ or policyengine_us/tests/policy/baseline/partners/. Changelog fragment changelog.d/chunked-microsimulation-calculate.added.md is a top-level file. The impact: and axiom: lines are both present.

  • Open overlap (REST search on the changed file names, plus earlier candidates). Make branch-measured marginal tax rates match fresh re-simulation #9698 edits system.py (input-variable sync) and spm.py, but not the mixin lines. It changes input_variables, which head TOB no longer reads, but main's formula and test_tob_sequencing.py's oracle do. Land Make branch-measured marginal tax rates match fresh re-simulation #9698 after this PR and rerun the sequencing test. Keep inputs when NY and AL recompute federal credits under pinned rules #10587 adds drop_recomputed_values to tools/period_branch.py, next to the drop_inherited_values that tob.py uses, and doesn't change it. Apply 416.1163(e)(2) to deemed SSI and identify SSI couples by marriage, not tax roles #10141 edits _ssi_spouses.py but keeps line 35, so the P2 case persists after it lands. Add annual Social Security earnings histories and per-year reforms #10560 touches only tools/social_security/. Add modeled OASDI liabilities, benefits and balance #10563 and Add configurable Social Security COLA means test #10486 have merged. No textual conflict is expected with any of these.

What I couldn't check

  • Legal sources. The egress proxy refused uscode.house.gov, govinfo.gov, ecfr.gov, ssa.gov and law.cornell.edu. This PR carries Measure benefit-taxation revenue as chapter 1 liability #10566's chapter 1 measure unchanged and adds no legal claim of its own. I confirmed the measure against main's code, not against the text of PL 98-21 section 121(e).
  • tob_revenue_oasdi.yaml at this head: I skipped it, as the brief asks (5.7 GB). The other TOB YAML files ran in CI.
  • Full-population runs. The head side of the hub pair hasn't posted yet.
  • Python 3.11 to 3.13 for the frame-depth test: CI runs that group on 3.14 only, and so did I.
  • The hand-off's mutant 6 count of 22: my output was capped at 8 lines.

Generated by Claude Code

@MaxGhenis

Copy link
Copy Markdown
Contributor Author

Picking this up for the US hub. I'll post a hand-off comment here when it's ready.


Generated by Claude Code

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
A reform can define a name the file stores but the country system lacks;
Core loads it unchunked, chunk files dropped it. A chunk holding exactly one
of several marital units pairs SSI spouses differently (_ssi_spouses.py).
Both now raise. A non-HDF5 .h5 file gets the flat-file error, the TOB
counterfactual holders join the Social Security preflight's calculated tax
leaves, and the population scan also flags .count == tests and random draws.
The validation fixture gives each household two marital units so its
one-household chunks still pass the new guard.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01LwQq6GfptFq3a9AdVcFg9H
@MaxGhenis

Copy link
Copy Markdown
Contributor Author

Ready for hub review

Head: c545c47877c70b367a6fab2322037e453364997e

Round 3 fix round on the review of 6835df5c68. The head had not moved when I started, and its last push was more than 6 hours before. I merged origin/main as a merge commit (041bdd36), which brings in the Core 3.34.2 lock (#10589). I added two commits on top of it and did not rebase or force-push.

Findings and what I did

# Finding What I did Where
1 P1: reform-added stored inputs are silently dropped Added the review's fix as written. _VariableCentricFile.unmatched lists names that match no country-system variable, _EntityTables.unmatched = (), and after _guard(simulation) the helper raises when the factory's system defines one of those names. I changed the docstring, the docs and the body to say so. New test uses the review's reform (reform_only_input, reform_only_output = 2 × input) and first shows that an unchunked run loads the input. tools/calculate_chunked.py:212, 303, 329, 516-530; docs/usage/chunked-calculations.md:41-45; test_calculate_chunked_validation.py:800
2 P2: a chunk holding one of several marital units changes SSI pairing Added the review's guard after masks is built. Docs and body now say "such chunks are rejected". New test uses the review's witness with household_chunks=[0, 1], which raises. [0, 0], the whole file in one chunk, is accepted. _ssi_spouses.py is unchanged. calculate_chunked.py:504-509; docs :85-86; test :748
3 P2: the new counterfactual holders can be stored Added both names to CALCULATED_TAX_LEAVES, worded as the review gives them. Each is now a case in test_stored_outputs_rejected_even_when_zero. tools/social_security_preflight.py:56-58; tests/code_health/test_social_security_preflight.py:103-105
4 P3: an .h5 file that isn't HDF5 gives a raw h5py error Wrapped the h5py.File open in try/except OSError; it now raises the flat-file ValueError. Test: a CSV saved as population.h5. calculate_chunked.py:304-313; test :764
5 P3: stale text Replaced the docs' last paragraph with the review's wording. Extended _POPULATION_SIGNAL with \.count\s*== and np\.random|default_rng (the first option). Across variables/, these match only _ssi_spouses.py:35, so there are no false positives. That file is listed in _CHUNK_SAFE_POPULATION_READERS with the reason. The test docstring also says what the scan can't see: other whole-array reductions, .any()/.all() branches, and reforms/. docs :125-126; test :672-713
6 P3: the "No formula changes" claim The body (TOB sequencing section) now has a "What differs from main in the TOB formulas" paragraph. It says the two states are cached once in parent holders, and that branch clearing differs. Main deletes non-input names for all periods and keeps every period of input-named variables. drop_inherited_values keeps only exact (variable, branch, period) set_input keys, so it recomputes formula values of input-named variables at other periods. I did not restructure tools/tob.py: the review's witness is the fixture comparison, which is exact, and I found no wrong value. What I ran is the 32-household fixture file test_calculate_chunked.py at this head (16 passed, which includes the TOB variable-centric 1/2/3/7-chunk and reform cases), plus test_tob_branch_lifecycle.py. I did not rerun the base-vs-head comparison with main's formulas. PR body

Should calculate_chunked also refuse a source that stores either holder name? No, and I didn't add it. A stored holder freezes TOB the same way in an unchunked run and in every chunk, so chunked results still equal the unchunked ones. The helper's job is to match the unchunked run, and the preflight is the scoring check for stored tax leaves. A refusal in the helper would make it disagree with Microsimulation(dataset=path) on the same file. It would be one line if the hub wants it anyway.

Judgment calls

  • Validation fixture changed. The review's marital-unit guard rejects any one-household chunk of the old 2-household validation fixture, since each household had one marital unit. The first run gave 4 failures among 73 passes in the first file run, plus the cases stopped by -x. The dataset fixture now holds two people per household, each their own marital unit (IDs 91, 92, 37, 38). Every other entity still has one row per household. I updated the person-length arrays in that file's tests to match: employment_income, person_weight, the claiming-tax-unit and membership rewrites, the cross-period case, and the per-chunk row comparison. The weighted sum is unchanged at 71,250.
  • The reform-name check only reads simulation.tax_benefit_system.variables when the file actually dropped a name. Without that, the stub systems in the release and cleanup tests (which have no variables) failed, and an entity-table source never needs the check.
  • The witness for finding 2 reproduced here. Unchunked, ssi is [0, 0, 0] and deeming applies to [True, False, False]. Household 1 alone gives ssi [11928, 0] and deeming [False, False]. The test doesn't build that simulation (about 50 s) because the guard raises before the factory runs. The numbers are in the fixture's docstring.
  • Core hook at 3.34.2. In the locked wheel, simulation.py:1760 is still the only caller of _check_period_consistency, and _run_formula is called at :1790. The body now cites 3.34.2.
  • One commit (c545c478) only restores docs/usage/social-security-pia-ratio.md. A broad ruff format docs call had rewrapped a code block in it; the net diff to that file is zero.

Mutants

Each ran in its own process in a scratch worktree, and git status showed it clean after each one.

Mutant Killed by
Remove the raise for reform-added names (if dropped: → if False:) test_rejects_stored_names_only_the_factory_system_defines
Remove the marital-unit guard test_rejects_chunk_holding_one_of_several_marital_units
Drop income_tax_without_social_security from CALCULATED_TAX_LEAVES test_stored_outputs_rejected_even_when_zero[income_tax_without_social_security]
Drop income_tax_with_pre_1993_social_security test_stored_outputs_rejected_even_when_zero[income_tax_with_pre_1993_social_security]
Remove the OSError wrap (except OSError → except ImportError) test_non_hdf5_file_with_h5_suffix_gets_flat_file_error

Tests at this head

Local Linux VM: 4 vCPU, 16 GB, Python 3.14, Core 3.34.2. One file per process; times are pytest time, and peak RSS is the child process's maximum.

File Result Seconds Peak RSS
tests/core/test_calculate_chunked_validation.py 77 passed 152.7 2.62 GB
tests/core/test_tob_branch_lifecycle.py 16 passed 6.6 0.86 GB
tests/microsimulation/test_calculate_chunked.py 16 passed 904.8 2.46 GB
tests/code_health/test_social_security_preflight.py 57 passed 5.6 not measured

The validation file went from 74 cases to 77 and from 1.53 GB to 2.62 GB peak RSS. The rise comes from the new reform-added-input case, which builds two small reformed Microsimulations. I didn't rerun test_tob_sequencing.py or the TOB YAML, because no TOB code changed in this round.

CI

Pull request run 41556562467 at c545c478 concluded success: all 34 jobs succeeded, with none skipped or cancelled. That includes PR validation (lock, guard tests, lint, changelog), Household API Partners, every Baseline shard, all seven Rest Python shards and Microsimulation, which took 55 min (00:58 to 01:53 UTC). The Documentation run also succeeded.

Impact

This round changes tooling (tools/calculate_chunked.py, tools/social_security_preflight.py), tests and docs only. It touches no formula in policyengine_us/variables/, nothing under .github/, and nothing under the partner tests. The merge from main brings main's own changes. The new raises fire only for inputs the old code would have silently miscomputed (reform-only stored names, one-of-several marital-unit chunks) or failed on (a non-HDF5 .h5). The pair's chunks cover a population file, where each chunk holds thousands of marital units. So the pair's result at the earlier head still applies to the PR's own changes. A pair against the new merge base would also pick up main's changes since 0ea88946. impact: pending (hub) and the axiom: line are unchanged in the body.

Proposed follow-up

  • Make _ssi_spouses independent of the simulation's marital-unit count. possible_default = person.marital_unit.count == 1 depends on the whole simulation's size, so a household's SSI couple pairing changes with what else is in the simulation. Once that test is per household or per input, calculate_chunked's marital-unit guard and the scan's listing can be removed.

Generated by Claude Code

@MaxGhenis

MaxGhenis commented Nov 3, 2026 •

Copy link
Copy Markdown
Contributor Author

Hub review in progress

Round-3 independent review of head c545c47877c70b367a6fab2322037e453364997e. The head hasn't moved. This comment holds results that are already firm and will be updated; the final review follows as a separate comment.

Firm so far

  • CI at the head. Pull request run 41556562467 concluded success: 34 of 34 jobs succeeded, with none skipped or cancelled. Documentation run 41556562302 also succeeded. Round 2's run had 35 jobs; this one plans 7 Rest Python shards instead of 8, and "Check Rest Python shards" passed.

  • Changed test files at the head, one file per process, on a 4 vCPU / 16 GB Linux VM with Python 3.14.6 and Core 3.34.2:

    • test_social_security_preflight.py: 57 passed, peak RSS 0.88 GB.
    • test_tob_branch_lifecycle.py: 16 passed, 0.86 GB.
    • test_calculate_chunked_validation.py: 77 passed, 108.6 s pytest time, peak RSS 2.62 GB. This matches the hand-off's figure.
    • TOB YAML: tob_revenue_total 4 passed (5.33 GB), tob_revenue_medicare_hi 6 passed (5.17 GB), tob_tier_split 3 passed (4.03 GB), tob_revenue_oasdi 13 passed (5.54 GB). Main's tob_revenue_total peaks at 5.23 GB on the same VM.
    • test_tob_sequencing.py: 1 passed (102 s, 2.17 GB).
    • test_calculate_chunked.py: 16 passed (544.5 s, 2.51 GB) with a private TMPDIR. In the shared /tmp, while other processes on the VM ran calculate_chunked on variable-centric files, two test_variable_centric_file_is_elementwise_exact cases failed. Both failures were the chunk-directory assertions at :621 and :639, not exactness (see below).
  • Merged with main. A scratch merge of main 6fa4eaf6d3 (which includes Make branch-measured marginal tax rates match fresh re-simulation #9698) into the head has no conflicts. With it: preflight 57 passed, lifecycle 16, validation 77 (2.61 GB), test_tob_sequencing.py 1, test_calculate_chunked.py 16 (private TMPDIR), and the four TOB YAML files 26. The YAML peaks equal the head's.

  • TOB values, head against the merge. On the 32-household fixture the TOB arrays are elementwise identical between the head and the merge in every route. That covers: baseline unchunked; 3 chunks; the threshold reform unchunked and in 2 chunks; and the variable-centric file. Main equals the head on the unchunked baseline too.

  • Test isolation (P3). _chunk_directories() globs tempfile.gettempdir() for policyengine-us-chunk-* machine-wide, at test_calculate_chunked.py:605-639 and test_calculate_chunked_validation.py:539-607, 818-827. So these tests fail whenever another process on the machine runs a variable-centric calculate_chunked, such as a population run on the hub machine. An autouse fixture that points tempfile.tempdir at a per-test directory fixes this.

  • Legal sources. govinfo.gov, ecfr.gov and law.cornell.edu were reachable from this session. ssa.gov and congress.gov returned 403, and web.archive.org reset the connection. From govinfo's 2023 edition of 42 U.S.C. 401, the note for Pub. L. 98-21 §121(e)(1)(A) appropriates to each payor fund "the aggregate increase in tax liabilities under chapter 1 of the Internal Revenue Code of 1986 which is attributable to the application of sections 86". It subtracts the increase "attributable to the amendments to section 86 ... made by section 13215" of the 1993 Act, and (B) sends that amount to the HI trust fund. 26 U.S.C. 86 sits in chapter 1, and §1411 (NIIT) sits in chapter 2A. So the PR's allocation is consistent with the statute: OASDI is old-law minus no inclusion, total is current minus no inclusion, and HI is the difference, with NIIT excluded. The PR carries main's chapter 1 measure unchanged and adds no legal reading of its own.

  • CI cost (likely P2). The Microsimulation job has timeout-minutes: 60, in both python-tests.yaml and push.yaml. At this head it ran 54 min 57 s; round 2's run took 46.4 min, and main's push run at 6fa4eaf6 took 45.5 min. Six other PRs' runs the same night took 38.3 to 46.5 min. The simulation-test-timings artifacts give the reason. CI runs all of tests/microsimulation in one pytest process (make test-microsimulation-base). There, this PR's test_calculate_chunked.py and test_tob_sequencing.py take 386 s and 167 s (553 s) at this head, and 481 s at round 2. That step's peak RSS was 14.77 GB at this head and 15.02 GB at round 2, against 14.10 to 14.49 GB for the six other PRs, on 16 GB runners. The body still says "Linux CI measurements are not yet available". It also says "every test runs as a separate foreground file with a 4 GB RSS cap", which describes the local runs, not CI.

  • Hand values. In _spouse_and_dependent_households, the 70-year-old alone gets $11,928. That is 12 × $994, the 2026 individual amount in parameters/gov/ssa/ssi/amount/individual.yaml. Its cited source, ssa.gov, returns 403 here. For test_skips_empty_chunks..., 10,000 × 1.5 + 25,000 × 2.25 = 71,250, and the two people the rewrite added have zero earnings, so the sum is unchanged.

  • Main moved again. Main is now 8604dc52. git merge-tree against it reports no conflict, and a scratch merge builds. The only change under tools/ is Keep inputs when NY and AL recompute federal credits under pinned rules #10587's drop_recomputed_values in period_branch.py; drop_inherited_values, which tools/tob.py calls, is unchanged.

  • Stored-holder refusal claim: confirmed. On the fixture's last four households, written as a variable-centric file with a stored zero for each holder, every TOB array is elementwise equal, unchunked against 2 chunks (3 runs × 3 variables, all equal=True). The stored zero moves TOB as the arithmetic predicts. For tax unit 2, a stored zero no-benefit tax gives total 8,072.6 (chapter 1 tax), against 5,972.6 clean, so no-benefit tax is 2,100. A stored zero old-law tax gives OASDI −2,100 and HI 8,072.6. So the helper matches Microsimulation(dataset=path), and the preflight is the right place to catch it.

  • Make branch-measured marginal tax rates match fresh re-simulation #9698's rebuilt baseline at the head. A reformed Microsimulation (base threshold 0,000) on the same four households: its baseline branch has the mixin, owns separate holders, and gives TOB equal to a fresh baseline whether the reform or the baseline computes first. The reform's own TOB doesn't depend on the order either.

  • Validation memory. Run alone, the new reform-added-input test peaks at 1.77 GB, and the other 76 cases at 1.57 GB. Together they reach 2.62 GB.

  • CI per-file time, from the Microsimulation job log's --durations: test_tob_sequencing 166.7 s, and the listed test_calculate_chunked entries sum to about 386 s.

Status

Complete. The final review is #10561 (comment), and it supersedes this comment.


Generated by Claude Code

@MaxGhenis

Copy link
Copy Markdown
Contributor Author

Ready for hub review
VERDICT: REQUEST_CHANGES
c545c47
METHODOLOGY_CHOICE: no

Independent round-3 review of head c545c47877. The head hasn't moved. This review continues the interim comment above, which an earlier session started; I reran or rechecked each of its results that I rely on here. Every number below comes from a run in this session: a 4 vCPU / 16 GB Linux VM, Python 3.14.6, the locked policyengine-core 3.34.2, one test file per process, each with a private TMPDIR.

All five round-2 findings are closed at this head, and their fixes are sound. I found no wrong result. I'm asking for changes for one reason: CI's Microsimulation job now runs within five minutes of its timeout (P2). The rest are P3s.

Findings

P2. The Microsimulation CI job takes 55 of its 60 minutes

.github/workflows/python-tests.yaml:155 sets timeout-minutes: 60 for "Full Suite - Microsimulation". At this head (run 41556562467) that job ran from 00:58:44 to 01:53:41, which is 54 min 57 s. The same job took:

  • 47 min 19 s on main's push run at 8604dc52 (run 41588249384);
  • 43.4, 41.1 and 34.9 min on three other PRs' runs today (runs 41588134995, 41588082387, 41587172431).

The job log's --durations (the base group runs every tests/microsimulation file in one pytest process) lists 166.66 s for test_tob_sequencing.py::test_sequential_tob_matches_main_elementwise_and_reuses_states. The listed test_calculate_chunked.py entries add up to about 386 s, the biggest being test_variable_centric_file_is_exact_under_reform (81.8 s), test_random_household_partitions_are_exact (73.1 s), test_reform_factory_is_exact_and_releases_its_baseline (49.1 s) and the module fixture's setup (42.0 s). Together that is about 9 minutes, which matches the gap to main. With a 5-minute margin, runner variance or the next test added to that group can time the job out. That would block every PR, not just this one.

The project guide ("Keep the existing runner budget") rules out adding a job or matrix entry without authorization. So the fix has to come from this PR's own tests, or the hub authorizes the budget. Options, cheapest first:

  1. Drop the 7-chunk cases from test_chunked_real_model_is_elementwise_exact and test_variable_centric_file_is_elementwise_exact (63 s together). Round 2 found that 3 chunks already cover order errors that 2 can't, and the 32-household fixture's smallest 3-chunk chunk holds 10 households.
  2. Cap the Hypothesis property test (73 s) at fewer examples, or move it to the existing Rest Python group, which builds a Core model anyway.
  3. Ask the hub to rule on raising timeout-minutes. That edit is under .github/, so this PR can't make it.

The body also says "Linux CI measurements are not yet available" and "Linux CI results pending". Replace both with the measured job time above.

P3. The marital-unit guard fires only after earlier chunks have run

policyengine_us/tools/calculate_chunked.py:504. The guard checks each chunk inside the loop, so every chunk before the bad one is built and calculated first. Witness: a file with household 1 holding two people in two marital units and household 2 one person, default 2 chunks, a counting stub factory:

raised: A chunk holding exactly one marital unit can pair SSI spouse
simulations built before the raise: 1

On a population run with arbitrary household_chunks, the bad chunk can be the last of eight, and the run fails after hours of work. The check needs only assignments and positions, so it can run once before the loop. Replace the in-loop block with this, placed just before tables = {}:

    # _ssi_spouses.py tests whether the simulation holds one marital unit.
    marital_units = np.bincount(
        assignments[positions["marital_unit"]], minlength=n_chunks
    )
    if len(positions["marital_unit"]) > 1 and (marital_units == 1).any():
        raise ValueError(
            "A chunk holding exactly one marital unit can pair SSI spouses "
            "differently from the whole file (_ssi_spouses.py); use fewer "
            "chunks or different household_chunks."
        )

assignments is already validated as non-negative integers below n_chunks. An empty chunk counts 0 and is still skipped.

P3. No test shows that a file holding a single marital unit is accepted

My own mutant, sum() == 1 < len(...) changed to sum() == 1 <= len(...), rejects every file whose only marital unit lands in one chunk. It survives the whole validation file (77 passed) and, by construction, the 32-household microsimulation file, where every file holds 96 marital units. This test kills it (passes at the head; with the mutant it fails with the guard's ValueError at calculate_chunked.py:505):

def test_accepts_file_holding_a_single_marital_unit():
    # The whole file is one marital unit, so the chunk pairs SSI spouses
    # exactly as the unchunked simulation does.
    person = pd.DataFrame({"person_id": [1, 2]})
    tables = {"person": person}
    for entity in US_ENTITIES:
        if entity != "person":
            person[f"person_{entity}_id"] = [1, 1]
            tables[entity] = pd.DataFrame({f"{entity}_id": [1]})
    tables["household"]["household_weight"] = [1.0]
    one = USSingleYearDataset(**tables, time_period=YEAR)
    result = calculate_chunked(InputSimulation, ["household_weight"], YEAR, 1, dataset=one)
    assert result.weighted_sums == {"household_weight": 1.0}

P3. The chunk-directory assertions read the machine-wide temp folder

test_calculate_chunked_validation.py:539-542 and test_calculate_chunked.py:605 glob tempfile.gettempdir() for policyengine-us-chunk-*, and the tests assert set equality at validation:554, 568, 607, 827 and test_calculate_chunked.py:621, 639. Any other process on the machine running a variable-centric calculate_chunked breaks them. A hub machine running a population pair while it tests is that case. Witness: a shell loop creating and removing policyengine-us-chunk-other<i> in the same TMPDIR while -k variable_centric_chunk runs:

>       assert _chunk_directories() - before == {Path(path).parent}
E       AssertionError: assert {PosixPath('/...nk-other508')} == {PosixPath('/...nk-4yk0mx0h')}

Give each test its own temp folder. Add this to both files:

@pytest.fixture(autouse=True)
def private_tempdir(tmp_path, monkeypatch):
    # calculate_chunked's mkdtemp and _chunk_directories then see only this test's chunks.
    directory = tmp_path / "tmp"
    directory.mkdir()
    monkeypatch.setattr(tempfile, "tempdir", str(directory))

Add import tempfile at the top of both files, and drop the local import in the validation file's _chunk_directories. I applied this to the validation file on a scratch copy and reran -k "variable_centric_chunk or stored_names_only" under the same interference loop: 3 passed.

P3. The new reform-input test builds two reformed simulations where one will do

test_calculate_chunked_validation.py:800-827. Run alone, this test peaks at 1.77 GB, and the other 76 cases peak at 1.57 GB. The whole file peaks at 2.62 GB (2.61 GB at the merge with main). The chunk side needs only a simulation whose tax_benefit_system defines reform_only_input, because the helper raises before it calculates anything. So it can reuse the unchunked simulation's reformed system through the file's own FileSimulation stub. Replace lines 816-822 (from whole.tax_benefit_system.simulation = None through ["reform_only_output"],) with:

    reformed = whole.tax_benefit_system

    def factory(chunk):
        # Only the chunk system matters; reuse the unchunked reformed one.
        simulation = FileSimulation(chunk)
        simulation.tax_benefit_system = reformed
        return simulation

    before = _chunk_directories()
    with pytest.raises(ValueError, match="reform_only_input are stored in the file"):
        calculate_chunked(
            factory,
            ["household_weight"],

With that change, all 77 pass, in 108.8 s, and the file peaks at 2.20 GB, down from 2.62. Mutant 1 (if dropped: → if False:) still fails the test, because the stub then returns household_weight without raising. The unchunked half, which shows Core loading the reform-added input, still uses a real reformed Microsimulation.

P3. Stale statements in the body

  • "Linux CI measurements are not yet available" and "Linux CI results pending": see P2.
  • "Every test runs as a separate foreground file with a 4 GB RSS cap": that describes local runs. The four TOB YAML files the body cites peak at 5.17 to 5.54 GB (measured again here at the merge with main).
  • "All 73 new cases": the validation file alone now has 77.
  • "A source scan fails when a variables file that builds a MicroSeries or reads a weight variable": the scan now also matches .count ==, np.random and default_rng (the earlier paragraph says so). Make this sentence match it.
  • "The validation increase is one module-scoped reformed simulation for the real-guard tests" explains round 2's rise. Round 3's rise, from 1.53 to 2.62 GB, comes from the reform-input test (see above).

Round-2 findings at this head

Round-2 finding Status How I checked
P1, reform-added stored inputs dropped Closed calculate_chunked.py:301-330, 516-529; the new test runs the round-2 reform and asserts the unchunked load. Mutant 1 is killed by test_rejects_stored_names_only_the_factory_system_defines. The check runs after the first factory call, before any calculation, and the chunk directory is still removed (asserted at :827).
P2, one marital unit in a chunk Closed (see the P3 on timing) The guard condition sum() == 1 < len(...) is exactly the case where person.marital_unit.count == 1 (_ssi_spouses.py:35) differs between chunk and file. Mutant 2 is killed by test_rejects_chunk_holding_one_of_several_marital_units.
P2, stored counterfactual holders Closed Both names are in CALCULATED_TAX_LEAVES. Mutants 3 and 4 are each killed by their own test_stored_outputs_rejected_even_when_zero[...] case.
P3, .h5 that isn't HDF5 Closed calculate_chunked.py:304-313. Mutant 5 is killed by test_non_hdf5_file_with_h5_suffix_gets_flat_file_error.
P3, stale text Mostly closed The docs paragraph is replaced as asked. The scan now matches .count == and random draws, _ssi_spouses.py is listed, and the docstring names what the scan can't see. The body's "No formula changes" is now a correct "What differs from main" paragraph. The body sentences listed above remain.

Checked and sound

  • The holder-name refusal claim holds. I took the 32-household fixture's last four households (the first four have no taxable benefits, so every TOB value there is 0), wrote them as a variable-centric file, and stored a zero for each holder in turn. Unchunked Microsimulation(dataset=path) against calculate_chunked(..., 2, ...): all 9 arrays are elementwise equal (equal=True throughout, peak 2.13 GB). The stored zero moves TOB as the formulas predict. For tax unit 2, the clean values are total 5,972.6, OASDI 3,482 and HI 2,490.6. A stored zero no-benefit tax gives total 8,072.6, which is the whole chapter 1 tax, so the no-benefit tax is 8,072.6 − 5,972.6 = 2,100. OASDI is then 5,582 (old-law tax), and 5,582 − 2,100 = 3,482 matches the clean OASDI. A stored zero old-law tax gives OASDI 0 − 2,100 = −2,100 and HI 5,972.6 + 2,100 = 8,072.6. Both match the output. The helper agrees with an unchunked run, so a refusal there would make it disagree with Microsimulation(dataset=path). The preflight is the right check, and it now catches both names.
  • Make branch-measured marginal tax rates match fresh re-simulation #9698 and the cached holders. At this head, and at a scratch merge with main 8604dc52 (which includes Make branch-measured marginal tax rates match fresh re-simulation #9698 and has no conflict), I built a reformed Microsimulation (base threshold JOINT = $20,000) on the same four households. Its baseline branch is a TobCalculationMixin, and each holder in it is a different object from the parent's. The baseline's TOB equals a fresh baseline simulation whether the reform or the baseline calculates first. The reform's TOB doesn't depend on the order, and it differs from baseline (HI 2,452.80 against 2,234.40 for tax unit 4). _rebuild_baseline_branch runs in __init__, before any calculation, so the branch clones empty holders, and each simulation caches its own states afterwards. All 15 arrays are identical between the head and the merge (peak 3.11 GB each side).
  • Tests at the merge with 8604dc52, one file per process: lifecycle 16 passed (0.87 GB); preflight 57 passed (0.89 GB); validation 77 passed (141 s, 2.61 GB); test_tob_sequencing.py 1 passed (136 s, 2.17 GB); TOB YAML: tob_revenue_medicare_hi 6 passed (5.17 GB), tob_revenue_oasdi 13 (5.54 GB), tob_revenue_total 4 (5.33 GB), tob_tier_split 3 (4.03 GB). The earlier session ran the same set, plus test_calculate_chunked.py, at the merge with 6fa4eaf6, with the same counts, and found the fixture's TOB arrays identical between head and merge in every route. Since 6fa4eaf6, main changed tools/period_branch.py only by adding drop_recomputed_values (Keep inputs when NY and AL recompute federal credits under pinned rules #10587). drop_inherited_values, which tools/tob.py calls, is unchanged.
  • Fixture rewrite. Each rewritten array still asserts what the old one did. [91, 91] becomes all 91 (every person in one unit, so one group spans both households). [37, 91] becomes [37, 37, 91, 91] for the external claiming unit and the multi-year swap (each household still points to the other's unit). person_household_id [91, 999] becomes [91, 91, 999, 999] (the unknown household is still referenced). The short array has 3 values for 4 people, so it is still one row short. The ID-change case still swaps the two households' IDs in one period. In the weights test, 10,000 × 1.5 + 25,000 × 2.25 = 71,250, and the two added people earn 0, so the sum is unchanged. Person weights [1.5, 1.5, 2.25, 2.25] are the projected household weights. The per-chunk row comparison now slices two rows for person and marital-unit arrays. That also means the fixture now catches a marital-unit array written with the household mask (shape 1 against 2), which the old one-row fixture couldn't.
  • Mutants, each in its own process on a scratch worktree, with git status clean after each:
    1. if dropped: → if False:: killed by test_rejects_stored_names_only_the_factory_system_defines (76 others pass).
    2. Marital-unit guard → if False:: killed by test_rejects_chunk_holding_one_of_several_marital_units.
    3. Drop income_tax_without_social_security from CALCULATED_TAX_LEAVES: killed by test_stored_outputs_rejected_even_when_zero[income_tax_without_social_security].
    4. Drop income_tax_with_pre_1993_social_security: killed by [income_tax_with_pre_1993_social_security].
    5. except OSError → except ImportError: killed by test_non_hdf5_file_with_h5_suffix_gets_flat_file_error.
    6. (mine) == 1 < len → == 1 <= len: survives all 77; see the P3 above for a test that kills it.
    7. (mine) Never record unmatched names (self.unmatched.append(name) → pass): killed by test_rejects_stored_names_only_the_factory_system_defines.
  • Readers of the changed variables: tob_revenue_medicare_hi, then oasdi_benefit_taxation_revenue (adds = ["tob_revenue_oasdi"]), then oasdi_income, then oasdi_modeled_balance (adds = ["oasdi_income"]). The S.770 files mention tob_revenue_oasdi only in documentation strings. This round changes no formula, so the readers are as round 2 found them.
  • Baseline outputs. This round touches only tools/, tests and docs, plus the merge from main. The new raises fire only for a reform-defined stored name, a one-of-several marital-unit chunk, or a non-HDF5 .h5. A population run in 8 balanced chunks hits none of them, so the hub pair's base side still stands, and the PR still needs its head side (impact: pending (hub)).
  • CI at the head: Pull request run 41556562467, 34 of 34 jobs succeeded, none skipped or cancelled. The Documentation run also succeeded.
  • Scope: nothing under .github/ or policyengine_us/tests/policy/baseline/partners/. The PR's 16 files include the top-level changelog fragment. The impact: and axiom: lines are present.
  • Open overlap: REST searches on calculate_chunked, tob.py, tob_revenue_oasdi, tob_revenue_total, income_tax_without_social_security, social_security_preflight and TobCalculationMixin return no other PR that changes these files. Keep ABLE contributions in the saver's credit and apply 6012(a)(1) to every 2026 filer #10590, Count Michigan household resources' benefits and premiums for the claimant and spouse only #10325 and Apply the 2026 section 68 reduction in Maryland, Maine and New Mexico itemized deductions #10594 match only on text and share no file with this PR. Apply 416.1163(e)(2) to deemed SSI and identify SSI couples by marriage, not tax roles #10141 (head 0327519344) edits _ssi_spouses.py but keeps possible_default = person.marital_unit.count == 1 at line 35, so the guard and the scan's listing are still needed after it lands.

What I couldn't check

  • Legal sources. This session's proxy refused govinfo.gov, uscode.house.gov, ecfr.gov, law.cornell.edu and ssa.gov (CONNECT 403). The interim comment's quote of the Pub. L. 98-21 §121(e) note came from govinfo in the earlier session; I couldn't re-fetch it. This PR adds no legal reading: it carries main's chapter 1 measure.
  • CI timing artifact. The artifact's storage host was refused. The per-test times above come from the job log's --durations instead. The step's peak RSS (14.77 GB in the interim comment) comes from that artifact, and I couldn't recheck it.
  • test_calculate_chunked.py at the merge with 8604dc52. I didn't rerun it. The earlier session ran it at the merge with 6fa4eaf6 (16 passed), and nothing it reaches has changed since, apart from the unused drop_recomputed_values.
  • Full-population runs, including the hub pair's head side.
  • Python 3.11 to 3.13.

Generated by Claude Code

@MaxGhenis

Copy link
Copy Markdown
Contributor Author

Picking this up for the US hub. I'll post a hand-off comment here when it's ready.


Generated by Claude Code

@MaxGhenis

Copy link
Copy Markdown
Contributor Author

Hub impact pair (2030), progress

The base side finished. The head side runs next, on the same household chunks and in the same environment, and a result comment follows.

  • Base 0d933918a82827b92243c0a43d0d8d03571604fd, the merge base of head c545c47877c70b367a6fab2322037e453364997e with main. The brief named head 6835df5c and base 0ea88946. The PR's head has since moved to c545c478, and its merge base with main moved with it, so this pair measures c545c478 against 0d933918.
  • The head's calculate_chunked, loaded by path, ran the base formulas for 2030 on 8 household chunks. The side ran in the head's uv sync environment (Python 3.14.6, policyengine-core 3.34.2) and imported policyengine_us from the base tree. pyproject.toml and uv.lock don't differ between the two commits.
  • Dataset: hf://datasets/policyengine/populace-us/populace_us_2024.h5@populace-us-2024-spm-20260909, the default in both trees. The downloaded file's SHA-256 equals DEFAULT_DATASET_SHA256 (6496cc4393d4d3c6574f76eca231de5898c803b9067645591fd5c4d3e65aee84). I downloaded it before the base side, so the side's time doesn't include the download.
  • Microsimulation accepted dataset_end_year=2030, and the helper accepted all nine variables.
  • Wall time: 4,396 s (73.3 min) inside calculate_chunked, and 1:13:58 for the whole process.
  • Maximum resident set size: 10,357,244 KiB (9.88 GiB, 10.61 GB), under the 14 GB limit that would call for a 16-chunk rerun. Exit status 0.
Variable Base weighted sum, 2030
tob_revenue_oasdi 85,000,478,688.04
tob_revenue_medicare_hi 60,733,320,670.50
tob_revenue_total 145,733,799,367.91
income_tax_chapter_1 3,153,596,038,642.86
income_tax 3,112,583,162,904.35
state_income_tax 676,000,734,185.94
social_security 2,097,686,306,929.41
household_net_income 19,488,922,620,773.88
household_benefits 2,512,646,035,391.43

Generated by Claude Code

@MaxGhenis

Copy link
Copy Markdown
Contributor Author

Hub impact pair, progress

The base side finished. The head side runs next, on the same chunks and in the same environment, and a result comment follows.

  • Head c545c47877c70b367a6fab2322037e453364997e. Base 0d933918a82827b92243c0a43d0d8d03571604fd, the head's merge base with main and the main commit that 041bdd36 merged in. The earlier progress comment measured base 0ea88946 for head 6835df5c. The head has since moved and merged main, so this pair measures c545c478 against its own merge base.
  • The head's calculate_chunked, loaded by path, ran the base formulas on 8 household chunks for 2026. It ran in the head's uv sync environment (Python 3.14.6, policyengine-core 3.34.2). pyproject.toml and uv.lock don't differ between the two commits, and summary.json names the base tree as the imported package.
  • Dataset: hf://datasets/policyengine/populace-us/populace_us_2024.h5@populace-us-2024-spm-20260909, 57,240 households. It is Microsimulation.default_dataset in both trees, and DEFAULT_DATASET and DEFAULT_DATASET_SHA256 are identical in both. The downloaded file's SHA-256 equals DEFAULT_DATASET_SHA256 (6496cc4393d4d3c6574f76eca231de5898c803b9067645591fd5c4d3e65aee84). I downloaded it before this side, so the times below exclude the download.
  • Wall time: 4,804 s inside calculate_chunked, and 1:20:47 for the whole process.
  • Maximum resident set size: 9,808,564 kB (9.35 GiB), under the 14 GB limit that would call for a 16-chunk rerun. Exit status 0.
  • The helper accepted all nine variables, and Microsimulation accepted dataset_end_year.
Variable Base weighted sum, 2026
tob_revenue_oasdi 63,221,571,702.76
tob_revenue_medicare_hi 46,005,803,316.51
tob_revenue_total 109,227,375,027.48
income_tax_chapter_1 2,633,056,694,752.26
income_tax 2,595,105,975,379.62
state_income_tax 578,044,599,361.57
social_security 1,677,873,892,455.98
household_net_income 17,161,428,039,502.38
household_benefits 2,059,793,926,050.13

The earlier progress comment measured a different base, 0ea88946, under policyengine-core 3.33.2, so its sums don't carry over. Its three TOB sums and social_security equal these to the cent. The other five differ.


Generated by Claude Code

The marital-unit guard ran inside the chunk loop, so every chunk before
the bad one was built and calculated first; on a population run that can
be hours of work. It now counts each nonempty chunk's marital units from
the household assignments before the loop. np.unique counts only chunks
that hold households: np.bincount with minlength=n_chunks would allocate
one integer per possible label, 8 TiB for n_chunks = 2**40.

New validation cases: a bad second chunk raises before the first is
built; a file holding a single marital unit is accepted (the review's
mutant 6 rejected it unnoticed); sparse labels below a huge n_chunks
work. Each validation test now gets its own temp folder, so another
process's chunk folders in the shared one cannot break the directory
assertions. The reform-input test reuses the unchunked run's reformed
system in the file stub instead of building a second reformed
simulation. The docs now say a file that is one marital unit is
accepted.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01WR8LHf7LXw1C77y8n8AqCm
The job ran 54 min 57 s of its 60-minute limit at c545c47 (run
41556562467), and this PR's tests took 552.9 s of it. Every assertion is
kept:

- Drop the 7-chunk cases of the real-model and population-file tests;
  three chunks catch order errors that two cannot. The real-model test's
  per-chunk checks (source rows and order, release of earlier
  simulations, id coverage) move into a helper that the partition
  property also uses, so they still run on seven chunks, now uneven ones
  of 1 to 11 households. The population file's chunk-file checks now run
  on up to three chunks; the validation file still writes one-household
  chunk files.
- Hypothesis property: two generated partitions instead of four, plus
  two explicit ones that pin the strategy's ends: every household in one
  chunk (labelled 3, so six labels are empty and skipped) and seven
  uneven, interleaved chunks. This rebalances rather than cuts: about 13
  chunk simulations are expected either way. In CI, Hypothesis's ci
  profile derandomizes, so the one non-trivial generated partition is a
  fixed draw.
- TOB sequencing: main's reference algorithm lends the unchanged
  parameter tree to its three branches instead of copying it (locally an
  8 s copy and a 3 s collection each); the helper's own branches keep
  ordinary clones. The heap alive before the test is frozen, as the
  chunking module does, so its collections no longer rescan what earlier
  tests in the same CI process left.
- The reform-factory test reads unreformed income_tax from the module
  reference instead of building another baseline simulation.
- The reformed population-file case keeps its unchunked reference on
  Microsimulation(dataset=path, reform=...); its two chunks share that
  run's reformed system instead of each loading the country system again
  (about 20 s each locally). The module fixture keeps only the unreformed
  reference.
- Each test gets its own temp folder for chunk files.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01WR8LHf7LXw1C77y8n8AqCm
@MaxGhenis

Copy link
Copy Markdown
Contributor Author

Ready for hub review

Head: f8ad24ebd189d7f14b709b94d9c3c8b93e5773ee, from c545c47877 (nobody else pushed in between). Two commits: 392c74393f (items 2 to 5) and f8ad24ebd1 (item 1). No formula, system.py, parameter file, .github/ or partner test changes: the diff from c545c47877 touches tools/calculate_chunked.py, the three chunking/TOB test files and docs/usage/chunked-calculations.md.

CI at this head: Pull request run https://github.com/PolicyEngine/policyengine-us/actions/runs/41614839743. At hand-off (10:51 UTC), 33 of 35 jobs had succeeded, including the Microsimulation job: 49 min 8 s, against 54 min 57 s at c545c47877. This PR's tests took 282.7 s of that job, against 552.9 s before. Contrib other-shard-1 and other-shard-3 started at 10:24 UTC and were still running; they run the contrib reform suites, which this round doesn't touch. The Documentation run 41614839518 succeeded. The PR body is updated (item 6).

1. Microsimulation CI time (review P2)

What changed, all in policyengine_us/tests/microsimulation/:

  • 7-chunk cases dropped from test_chunked_real_model_is_elementwise_exact (test_calculate_chunked.py:325-333) and test_variable_centric_file_is_elementwise_exact (:657-658). Three chunks catch order errors that two can't, and the smallest of the three default chunks holds 10 households (11, 11, 10). The real-model test's per-chunk checks (each chunk's tables equal its households' source rows in order, earlier chunk simulations already collected, every id seen once) moved into calculate_with_checked_chunks (:275-322), and the Hypothesis property uses it too, so those checks still run on seven chunks.
  • Hypothesis property (:335-376): max_examples 4 → 2, plus two explicit examples. (7, every household labelled 3) is one chunk with six empty labels skipped; the strategy itself can't draw an empty chunk, since it uses every label in range(n_chunks). (7, UNEVEN_LABELS) is seven interleaved chunks of 1, 2, 3, 4, 5, 6 and 11 households. This pins the boundaries but doesn't save time: about 13 chunk simulations are expected either way (hand computation below).
  • test_tob_sequencing.py (:19-25, 57-78): main's reference algorithm runs under lent_parameter_history(). Its three branches change no parameter, so they get the unchanged tree instead of an 8.2 s copy plus a 3.1 s collection each (local profile). The helper's two branches, the code under test, keep ordinary clones. An autouse fixture freezes the heap alive before the test. In CI this file runs last in the base group's single pytest process, where it took 166.7 s against 103.7 s locally.
  • Reform-factory test (test_calculate_chunked.py:525-527): unreformed income_tax now comes from the module reference fixture, which makes the same make_simulation calculation, instead of from another baseline simulation.
  • Population file under a reform (:637-650, 688-716): the module fixture keeps only the unreformed reference. The reform test builds its reference with Microsimulation(dataset=path, reform=...), and its two chunks share that run's reformed system (tax_benefit_system=policy) instead of each loading the country system again (21.0 s each, local profile). The system is collected inside the test (:713-715), not in the next one.
  • The module fixture is split into frozen_heap() and lent_parameter_history() (:56-99), same order and behaviour as before.

What the tests no longer cover:

  • the default balanced 7-chunk split (5, 5, 5, 5, 4, 4, 4);
  • the population file's chunk-file checks above three chunks (the validation file still writes and checks one-household chunk files);
  • two of the three non-trivial generated partitions, each a fixed draw in CI (see judgment call 5);
  • main's reference algorithm with full parameter copies;
  • for the population file under a reform, a reformed system and reform baseline built per chunk. The entity-table reform-factory test still builds both per chunk and checks that the baseline is released.

Local measurements. The VM has 4 vCPU and 16 GB, with Python 3.14.6 and Core 3.34.2. Each file ran in its own process with a private TMPDIR. CI's update_itemization.sh setting was applied to the working tree and never committed. Seconds are pytest time; peak is the largest child RSS. Before runs are at c545c47877.

File Before: passed / s / peak GB After: passed / s / peak GB
test_calculate_chunked.py 16 / 399.6 and 425.4 / 2.34 and 2.39 14 / 296.9 / 2.28
test_tob_sequencing.py 1 / 99.7 and 103.8 / 2.06 and 2.07 1 / 43.4 and 42.9 / 2.06
test_calculate_chunked_validation.py 77 / 104.1 / 2.50 80 / 98.1 and 93.1 / 2.10 and 2.10

Per changed test, local before → after, and CI at c545c47877 → projected (each CI duration scaled by its local after/before ratio):

Test Local before → after (s) CI → projected (s)
TOB sequencing (call) 103.7 → 41.6 (+1.1 setup) 166.7 → 66.8
population file under reform 86.0 → 74.8 81.8 → 71.1
Hypothesis property 64.0 → 61.2 73.1 → 69.9
reform factory 56.1 → 56.5 49.1 → 49.4
population-file module setup 44.9 → 6.5 42.0 → 6.1
population file, 7 chunks 42.6 → removed 34.5 → 0
real model, 7 chunks 35.4 → removed 28.9 → 0
the other 12 entries (1-3 chunks, setup, small tests) 93.0 → 94.7 76.8 → 78.4

Projection: this PR's tests 552.9 → 343 s, and the job 54:57 → 51:27. That keeps the TOB test's old CI slowdown. At its local time, which freezing should approach, the projection is 318 s and 51:02. Measured at this head (run 41614839743, --durations): this PR's tests took 282.7 s and the job 49 min 8 s. CI before → after by test: TOB sequencing 166.7 → 39.9 s (32.8 s call plus 7.1 s to collect and freeze the heap, below its local time); population file under reform 81.8 → 60.7; Hypothesis property 73.1 → 54.0; reform factory 49.1 → 45.2; population-file module setup 42.0 → 4.7; 7-chunk cases 63.4 → 0; the other entries 76.8 → 78.1. Runner variance is large. The SPM step, untouched here, took 1,085.8 s against 1,202.6 s at c545c47877, and the base group's other tests took 938.5 s against 863.7 s.

Four minutes isn't reached: 282.7 s is 43 s over, and I found no further cut that keeps every assertion. These would narrow a test, for the hub to decide:

  • max_examples=1 leaves no non-trivial generated partition, about 4 chunk simulations (~18 s) fewer on average.
  • A 3-chunk ceiling for the strategy and the explicit example saves 4 chunk simulations (~18 s) and narrows the range.
  • A cloned system for the population-file reform reference, instead of the default route, saves about 10 s locally (21.0 s vs about 11 s).

Together they would save roughly 50 s, about enough for 4 minutes.

2. Marital-unit guard before any chunk (review P3)

policyengine_us/tools/calculate_chunked.py:500-511, placed just before tables = {}, with the review's condition and message, and its comment plus two lines. One deviation (judgment call 1): it counts with np.unique(assignments[positions["marital_unit"]], return_counts=True), not np.bincount(..., minlength=n_chunks). Docstring :460-468 and docs updated.

Review's witness: household 1 holds two people in two marital units, household 2 one person; default 2 chunks; a counting stub factory. The head's helper was loaded from git show c545c47877:... in the same process:

head c545c47877: raised: A chunk holding exactly one marital unit can pair SSI spouses diff
head c545c47877: simulations built before the raise: 1
this branch: raised: A chunk holding exactly one marital unit can pair SSI spouses diff
this branch: simulations built before the raise: 0

I also added test_rejects_one_marital_unit_chunk_before_building_any_chunk (tests/core/test_calculate_chunked_validation.py:772-779). It builds the same situation from the fixture and uses the file's never_build factory, and the head's in-loop guard fails it (mutant 8).

3. File holding a single marital unit (review P3)

test_accepts_file_holding_a_single_marital_unit (test_calculate_chunked_validation.py:782-796) is the review's code. Mutant 6 is killed in both forms (table below): the review's literal == 1 <= len(...) on the head helper, and this head's len(positions["marital_unit"]) >= 1.

4. A temp folder per test (review P3)

private_tempdir is at test_calculate_chunked_validation.py:30-36 and test_calculate_chunked.py:102-108. import tempfile is at the top of both (validation :4; chunked already had it), and the local import in _chunk_directories is gone (validation :549-550).

Interference runs used a loop that creates and removes policyengine-us-chunk-other<i> in the same TMPDIR:

  • this head, validation -k "variable_centric_chunk or stored_names_only": 3 passed, with 5,639 folders created during the run;
  • this head, test_calculate_chunked.py -k variable_centric_file_is_elementwise_exact: 3 passed (7,268 folders);
  • negative controls at c545c47877, same loop: the validation selection gives 2 failed and 1 passed (AssertionError: assert {PosixPath('/...k-other3112')} == {PosixPath('/...k-other3010')}), and the chunked [2] case 1 failed.

5. One reformed system in the reform-input test (review P3)

test_calculate_chunked_validation.py:867-884 is the review's replacement. Mutant 1 still fails this test, and the other 79 pass. Locally, the file's peak went from 2.50 GB (77 passed) to 2.10 GB (80 passed, two runs), and this test from 43.4 s to 21.6 s. In CI the file's durations summed to 85.0 s at this head (Rest Python shard 3, 80 passed), against 180.7 s at c545c47877 (shard 1, 77 passed); this test took 19.9 s against 52.0 s. The two shards run different files in the same process, so these aren't strictly comparable.

6. Body (review P3)

  • "Linux CI measurements are not yet available" and "Linux CI results pending" are replaced with the CI job times.
  • The 4 GB RSS cap is now described as applying to the round-1 local selections only, citing the round-3 review's 5.17 to 5.54 GB peaks for the four TOB YAML files.
  • "All 73 new cases" is now 113 new cases: validation 80, test_calculate_chunked.py 14, test_tob_sequencing.py 1, lifecycle 16, preflight 2.
  • The source-scan sentence names MicroSeries, weight names, .weights, .count ==, np.random and default_rng, as the regex at test_calculate_chunked_validation.py:693-697 does.
  • Round 3's validation memory rise is attributed to the reform-input case.
  • A round-4 section adds these measurements.
  • impact: pending (hub) and the axiom: line are kept.

Tests run at this head

One file per process, each with its own TMPDIR; numbers are from the table above.

  • tests/core/test_calculate_chunked_validation.py: 80 passed, 93.1 s, 2.10 GB.
  • tests/microsimulation/test_calculate_chunked.py: 14 passed, 296.9 s, 2.28 GB. Hypothesis reported 2 generated examples, stopped at max_examples=2.
  • tests/microsimulation/test_tob_sequencing.py: 1 passed, 42.9 s, 2.06 GB.
  • tests/code_health/test_social_security_preflight.py: 57 passed, 3.5 s, 0.84 GB.
  • tests/core/test_tob_branch_lifecycle.py: 16 passed, 5.3 s, 0.83 GB.
  • Ruff format and lint pass on the changed files.

I ran no full suites and no population runs.

Mutants

Each mutant ran in its own git worktree and pytest process, and git status was clean after each restore. They ran on an unpushed intermediate commit whose helper, validation, preflight and TOB files match this head, apart from one docstring word.

# Mutant Killed by Others
1 if dropped: → if False: test_rejects_stored_names_only_the_factory_system_defines 79 pass
2 marital-unit guard → if False: test_rejects_chunk_holding_one_of_several_marital_units and test_rejects_one_marital_unit_chunk_before_building_any_chunk 78 pass
3 drop income_tax_without_social_security from CALCULATED_TAX_LEAVES test_stored_outputs_rejected_even_when_zero[income_tax_without_social_security] 56 pass
4 drop income_tax_with_pre_1993_social_security test_stored_outputs_rejected_even_when_zero[income_tax_with_pre_1993_social_security] 56 pass
5 except OSError → except ImportError test_non_hdf5_file_with_h5_suffix_gets_flat_file_error 79 pass
6 len(positions["marital_unit"]) > 1 → >= 1 (this head's form) test_accepts_file_holding_a_single_marital_unit 79 pass
6 literal sum() == 1 < len(...) → sum() == 1 <= len(...) on the head helper test_accepts_file_holding_a_single_marital_unit only that test selected
7 self.unmatched.append(name) → pass test_rejects_stored_names_only_the_factory_system_defines 79 pass
8 guard back inside the loop (the head helper) test_rejects_one_marital_unit_chunk_before_building_any_chunk 79 pass
9 the review's np.bincount(..., minlength=n_chunks) test_marital_unit_check_counts_only_chunks_that_hold_households (Unable to allocate 8.00 TiB) 79 pass
F1 tools/tob.py keeps every TOB branch on its parent test_sequential_tob_matches_main_elementwise_and_reuses_states, on the weak-reference check at the second branch, with the heap frozen n/a
F2 tools/tob.py skips the bounded old-law set_input same test: 26 of 64 tax units differ from main's lent-tree reference n/a

Hand computations

  • 8 TiB: (2**40 + 1) int64 counts × 8 bytes; numpy reported "Unable to allocate 8.00 TiB".
  • Default splits of 32 households: 3 chunks give 11, 11, 10; 7 chunks give 5, 5, 5, 5, 4, 4, 4.
  • UNEVEN_LABELS: 1 + 2 + 3 + 4 + 5 + 6 + 11 = 32. Since gcd(7, 32) = 1, i ↦ 7i mod 32 permutes the households.
  • Expected chunk simulations in the property: before, 1 (Hypothesis's simplest example) + 3 × 4 = 13; after, 1 + 7 (explicit) + 1 + 4 = 13. n_chunks is drawn from 1..7, mean 4.
  • Projection example: the population-file module setup, 42.04 s in CI × 6.52 / 44.88 = 6.11 s.

Judgment calls

  1. np.unique instead of the review's np.bincount(..., minlength=n_chunks). minlength allocates one integer per possible label, so n_chunks = 2**40, which is legal input and works at the head, raises MemoryError. np.unique counts only chunks that hold households; empty chunks still count as zero. test_marital_unit_check_counts_only_chunks_that_hold_households (:799-812) pins it (mutant 9).

  2. Two regression tests added beyond the brief: the hoisting witness (item 2) and the large-label case above.

  3. "The most chunks the fixture allows" I read as the strategy's ceiling, 7. One household per chunk would be 32 chunk simulations, about 150 s on its own.

  4. The explicit one-chunk example is labelled 3 of 7. Hypothesis's first generated example is already the 0-labelled one-chunk partition, so an explicit (1, zeros) would duplicate it. Label 3 also covers empty labels with the real model at no extra simulation.

  5. max_examples=2, not 1, so one non-trivial generated partition remains. Under GitHub Actions, Hypothesis loads its built-in ci profile (derandomize=True), so that partition is one fixed draw, as each of the three was before. The test comment says so.

  6. TOB test. Lending applies only to the reference phase; the code under test still clones normally. Freezing excludes only objects alive before the test from its collections. F1 shows a retained branch is still caught, and F2 shows wrong values still differ from the lent reference.

  7. Population file under a reform. The chunks share the reference's reformed system, so the default Microsimulation(dataset=path, reform=...) route remains only in the unchunked reference.

  8. Independent review before pushing. I ran a read-only review of the diff before pushing. It found two things, both fixed:

    • the property's 7-chunk example lacked the per-chunk checks, so they now come from the shared helper;
    • CI derandomizes Hypothesis, so the 7-chunk example is now uneven and the comment says so.

    It also asked for precise wording: "can pair" in the docstring, and a docs sentence that says a single-marital-unit file is accepted.

  9. Docs (docs/usage/chunked-calculations.md): updated the chunk counts, the explicit partitions, the before-any-chunk rejection and the single-unit acceptance.

  10. Commit trailers. Both commits end with a Claude-Session trailer after Co-Authored-By, as this PR's earlier commits do. The brief asks for the Co-Authored-By line only, so drop the trailer in the squash message. I rebuilt my two local commits once before pushing; no pushed history changed.

What I couldn't check

  • CI's timing artifact (per-step peak RSS).
  • Full-population runs, including the hub pair's head side. This round changes no formula and no system.py, and only moves a rejection earlier, so the pair at c545c47877 should stand.
  • Python 3.11 to 3.13 locally; CI smoke-imports those.

Proposed follow-up issues

  • Give CI's Microsimulation job headroom under its 60-minute timeout. Main's run took 47 min 19 s: SPM tests 20 min, the base group 16 min and Social Security current law 9 min, run in sequence. Splitting a group into its own job or raising the timeout needs a .github/ change the hub authorizes.
  • Lend the parameter tree to TOB counterfactual branches in production. Each of the helper's two branches copies the full parameter tree (8.2 s plus a 3.1 s collection locally) per simulation, although neither branch changes a parameter. Lending, as these tests now do for main's reference, could cut that per chunk if the copy-on-write barrier makes it safe.

Generated by Claude Code

The population-file tests check each chunk file's name, its role dtype, that no other chunk file exists beside it, that the helper leaves an unknown variable name out, and that the helper deletes it; only the validation file compares array contents. Say that in docs/usage/chunked-calculations.md. In test_tob_sequencing.py, say that the file shares the microsimulation base group's pytest process and currently sorts last, since nothing fixes the order.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@MaxGhenis

Copy link
Copy Markdown
Contributor Author

Ready for hub review
IMPACT PAIR: zero-diff
head c545c47
base 0d93391

The brief named head 6835df5c and base 0ea88946. When this run started, the PR head was c545c478, whose history merges main at 0d933918 (merge commit 041bdd36), so this pair measures c545c478 against its own merge base, 0d933918. This run computed both sides and reuses no earlier output. Its base side is in this progress comment, and the base column below equals that comment's nine sums to the cent. Both come from this run, so that match isn't an independent check. No other progress comment gives a 2026 base side at 0d933918. A 2030 pair's progress comment uses the same base, but its sums are for 2030. The earlier 2026 progress comment measured 0ea88946 under policyengine-core 3.33.2, so its sums don't apply here. Its three TOB sums and social_security equal this run's to the cent, and the other five differ.

While the head side ran, the PR head moved to f8ad24eb. After the pair finished it moved again, to 3f2700051680816df5190e84e39f2aac9c6c0c73. The three commits since c545c478 change docs, tests and calculate_chunked.py, where the one-marital-unit check moves ahead of the chunk loop. They change no file under policyengine_us/variables/ or policyengine_us/parameters/, and they leave system.py, tools/tob.py, pyproject.toml and uv.lock unchanged. The latest head's merge base with main is still 0d933918. This pair measured c545c478, and both sides loaded its calculate_chunked.py.

  • Dataset: hf://datasets/policyengine/populace-us/populace_us_2024.h5@populace-us-2024-spm-20260909, which is Microsimulation.default_dataset in both trees. DEFAULT_DATASET and DEFAULT_DATASET_SHA256 are identical in both trees. The downloaded file (Hugging Face snapshot 9a814a3b) has SHA-256 6496cc4393d4d3c6574f76eca231de5898c803b9067645591fd5c4d3e65aee84, which matches DEFAULT_DATASET_SHA256. I downloaded it before the base side, so neither side's time includes the download. I hashed it again after the head side and got the same digest.
  • Year 2026; 8 chunks on both sides. Both sides load calculate_chunked.py from c545c478 by path.
  • One environment, the head's uv sync (Python 3.14.6, policyengine-core 3.34.2), serves both sides, since pyproject.toml and uv.lock don't differ between base and head. PYTHONPATH picks the tree, and each side's summary.json names the tree it imported. Before the pair, each tree's variable registry listed each of the nine variables on the same entity in both trees (six on tax_unit, one on person, two on household), each a float. The head's Microsimulation includes TobCalculationMixin and the base's doesn't. Microsimulation accepted dataset_end_year on both sides.
Variable Entity Base sum Head sum Difference Rows that differ Largest abs. difference
tob_revenue_oasdi tax_unit 63,221,571,702.76 63,221,571,702.76 0 0 0
tob_revenue_medicare_hi tax_unit 46,005,803,316.51 46,005,803,316.51 0 0 0
tob_revenue_total tax_unit 109,227,375,027.48 109,227,375,027.48 0 0 0
income_tax_chapter_1 tax_unit 2,633,056,694,752.26 2,633,056,694,752.26 0 0 0
income_tax tax_unit 2,595,105,975,379.62 2,595,105,975,379.62 0 0 0
state_income_tax tax_unit 578,044,599,361.57 578,044,599,361.57 0 0 0
social_security person 1,677,873,892,455.98 1,677,873,892,455.98 0 0 0
household_net_income household 17,161,428,039,502.38 17,161,428,039,502.38 0 0 0
household_benefits household 2,059,793,926,050.13 2,059,793,926,050.13 0 0 0

The sums use the simulation weights. "Rows that differ" counts entity records (tax units, people or households) whose values differ by more than 1e-6. Every value is equal between the sides (np.array_equal) in all nine variables, across 79,729 tax units, 166,321 people and 57,240 households, and the saved base and head tables are byte-identical files (cmp). Before comparing values, the script checked that each entity table's <entity>_id column is identical and in the same order on both sides and in the dataset's own entity table, and that the weight columns are equal. It also recomputed each weighted sum from the saved tables and found it equal to that side's summary.json. The TOB variables are nonzero for 13,727 (OASDI), 13,536 (HI) and 13,991 (total) tax units on each side.

Runs

  • Base: 1:20:47 for the whole process (GNU time), and 4,804 s (80.1 min) from the helper call through saving the tables (pair_side.py's wall_seconds). Maximum resident set size 9,808,564 KiB (9.35 GiB, 10.04 GB); exit status 0.
  • Head: 1:09:50 for the whole process (GNU time), and 4,149 s (69.1 min) from the helper call through saving the tables (pair_side.py's wall_seconds). Maximum resident set size 6,305,684 KiB (6.01 GiB, 6.46 GB); exit status 0.
  • The sides ran one at a time. During each side the VM ran only that side, a shell loop that read its resident set size every 20 s, and brief git commands and small file reads and writes. Neither side passed 14 GB (the larger peak was the base side's 10.04 GB), and both exited with status 0, so I didn't rerun with 16 chunks.

Dropped variables: none. The helper's guard rejected none of the nine on either side.

Commands

# Once on this VM
curl -sI https://huggingface.co/api/datasets/policyengine/populace-us | head -1   # HTTP/1.1 200 Connection Established (then HTTP/2 200)
apt-get install -y time   # the VM had no GNU time

# In the policyengine-us checkout; ../head, ../base, ../out and the scripts sit beside it
git fetch origin main pull/10561/head:pr-10561
git fetch --deepen=400 origin main pull/10561/head:pr-10561   # the checkout was shallow, so I deepened it before trusting merge-base
HEAD_SHA=$(git rev-parse pr-10561); BASE_SHA=$(git merge-base origin/main $HEAD_SHA); echo $HEAD_SHA $BASE_SHA
# c545c47877c70b367a6fab2322037e453364997e 0d933918a82827b92243c0a43d0d8d03571604fd
git diff --stat $BASE_SHA $HEAD_SHA -- pyproject.toml uv.lock   # empty
git worktree add --detach ../head $HEAD_SHA
git worktree add --detach ../base $BASE_SHA
mkdir -p ../out
(cd ../head && uv sync)

# Checks before the pair: the dataset download and hash, then each tree's registry
PY=$PWD/../head/.venv/bin/python
DS=$(PYTHONPATH=$PWD/../head $PY -c 'from policyengine_us.system import DEFAULT_DATASET, _resolve_dataset_path; print(_resolve_dataset_path(DEFAULT_DATASET))')
sha256sum -b "$(readlink -f "$DS")"
for side in base head; do PYTHONPATH=$PWD/../$side $PY ../registry_check.py > ../out/registry_$side.json 2> ../out/registry_$side.err; done

# The pair, one side at a time
H=$PWD/../head/policyengine_us/tools/calculate_chunked.py; PY=$PWD/../head/.venv/bin/python
export OMP_NUM_THREADS=4 OPENBLAS_NUM_THREADS=4 MKL_NUM_THREADS=4
date -u +%Y-%m-%dT%H:%M:%SZ > ../out/base.start
PYTHONPATH=$PWD/../base /usr/bin/time -v $PY ../pair_side.py $H 8 2026 ../out/base > ../out/base.log 2> ../out/base.err
date -u +%Y-%m-%dT%H:%M:%SZ > ../out/base.end
# Between the sides: the progress comment, and a self-test of the comparison on base against base (zero-diff)
$PY -I ../compare.py ../out/base ../out/base "$(readlink -f "$DS")" "$SCRATCH/selftest.json"   # SCRATCH: a scratch directory
date -u +%Y-%m-%dT%H:%M:%SZ > ../out/head.start
PYTHONPATH=$PWD/../head /usr/bin/time -v $PY ../pair_side.py $H 8 2026 ../out/head > ../out/head.log 2> ../out/head.err
date -u +%Y-%m-%dT%H:%M:%SZ > ../out/head.end

# Alongside each side, a sampler on the side's Python process (the child of /usr/bin/time);
# it printed each new whole-GiB peak
PID=$(pgrep -P "$(pgrep -f '^/usr/bin/time -v .*pair_side.py')")
peak=0; while kill -0 $PID 2>/dev/null; do rss=$(ps -o rss= -p $PID | tr -d ' '); [ "${rss:-0}" -gt $peak ] && peak=$rss; sleep 20; done
# During the base side I also ran `git fetch -q --deepen=50 origin main`, to place earlier pairs' bases on main;
# it moves no ref the pair uses.

# After both sides
sha256sum -b "$(readlink -f "$DS")"   # same digest as before
for side in base head; do $PY -I ../time_stats.py ../out/$side.err > ../out/$side.time.json; done
$PY -I ../compare.py ../out/base ../out/head "$(readlink -f "$DS")" ../out/compare.json
for e in household person tax_unit; do cmp ../out/base/$e.pkl ../out/head/$e.pkl && echo "$e.pkl byte-identical"; done

# The PR head moved: ls-remote gave f8ad24eb right after the pair and 3f270005 before this comment
git ls-remote origin refs/pull/10561/head
git fetch origin pull/10561/head:pr-10561-new     # then f8ad24ebd189d7f14b709b94d9c3c8b93e5773ee
git fetch origin pull/10561/head:pr-10561-newer   # later 3f2700051680816df5190e84e39f2aac9c6c0c73
for ref in pr-10561-new pr-10561-newer; do
  git log --oneline $HEAD_SHA..$ref; git diff --stat $HEAD_SHA $ref
  git diff --name-only $HEAD_SHA $ref | grep -E '^policyengine_us/(variables|parameters)/|^policyengine_us/system.py$|^policyengine_us/tools/tob.py$|^pyproject.toml$|^uv.lock$'   # no output
  git merge-base origin/main $ref   # 0d933918a82827b92243c0a43d0d8d03571604fd
done

pair_side.py is the brief's script, unchanged.

pair_side.py
"""One side of #10561's pair: python pair_side.py <helper.py> <chunks> <year> <out_dir>."""
import importlib.util
import json
import sys
import time
from pathlib import Path

helper, chunks, year, out = sys.argv[1], int(sys.argv[2]), int(sys.argv[3]), Path(sys.argv[4])
VARIABLES = (
    "tob_revenue_oasdi",
    "tob_revenue_medicare_hi",
    "tob_revenue_total",
    "income_tax_chapter_1",
    "income_tax",
    "state_income_tax",
    "social_security",
    "household_net_income",
    "household_benefits",
)
spec = importlib.util.spec_from_file_location("hub_calculate_chunked", helper)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)

import policyengine_us  # noqa: E402
from policyengine_us import Microsimulation  # noqa: E402

start = time.monotonic()
result = module.calculate_chunked(
    lambda dataset: Microsimulation(dataset=dataset, dataset_end_year=year),
    VARIABLES,
    year,
    chunks,
    dataset=Microsimulation.default_dataset,
)
out.mkdir(parents=True, exist_ok=True)
for entity, table in result.entities.items():
    table.to_pickle(out / f"{entity}.pkl")
summary = {
    "package": str(Path(policyengine_us.__file__).parent),
    "chunks": chunks,
    "year": year,
    "weighted_sums": result.weighted_sums,
    "wall_seconds": time.monotonic() - start,
}
(out / "summary.json").write_text(json.dumps(summary, indent=1, sort_keys=True))
print(json.dumps(summary, sort_keys=True), flush=True)
compare.py
"""Compare the two sides of #10561's pair.

python compare.py <base_dir> <head_dir> <dataset.h5> <out.json>

Checks that every entity table has the same <entity>_id column, in the same
order, on both sides and in the dataset's own entity tables; then, for each
variable, compares weighted sums and per-row values.
"""

import json
import sys
from pathlib import Path

import numpy as np
import pandas as pd

TOLERANCE = 1e-6
VARIABLES = (
    "tob_revenue_oasdi",
    "tob_revenue_medicare_hi",
    "tob_revenue_total",
    "income_tax_chapter_1",
    "income_tax",
    "state_income_tax",
    "social_security",
    "household_net_income",
    "household_benefits",
)

base_dir, head_dir, dataset, out_path = (
    Path(sys.argv[1]),
    Path(sys.argv[2]),
    sys.argv[3],
    Path(sys.argv[4]),
)
sides = {"base": base_dir, "head": head_dir}
summaries = {
    side: json.loads((path / "summary.json").read_text())
    for side, path in sides.items()
}
for side, summary in summaries.items():
    print(f"{side}: package {summary['package']}, chunks {summary['chunks']}, "
          f"year {summary['year']}, wall {summary['wall_seconds']:.1f} s")
assert summaries["base"]["chunks"] == summaries["head"]["chunks"], "chunk counts differ"
assert summaries["base"]["year"] == summaries["head"]["year"], "years differ"
for side, summary in summaries.items():
    assert sorted(summary["weighted_sums"]) == sorted(VARIABLES), f"{side} variables"

entities = sorted(path.stem for path in base_dir.glob("*.pkl"))
assert entities == sorted(path.stem for path in head_dir.glob("*.pkl")), "entity sets differ"
tables = {
    side: {entity: pd.read_pickle(path / f"{entity}.pkl") for entity in entities}
    for side, path in sides.items()
}

# 1. Identical ID columns, in the same order, on both sides and in the dataset.
report = {"entities": {}, "variables": {}, "differing_rows": []}
with pd.HDFStore(dataset, mode="r") as store:
    source_ids = {entity: store[f"/{entity}"][f"{entity}_id"].to_numpy() for entity in entities}
for entity in entities:
    base, head = tables["base"][entity], tables["head"][entity]
    id_column = f"{entity}_id"
    assert list(base.columns) == list(head.columns), f"{entity} columns differ"
    assert base[id_column].dtype == head[id_column].dtype, f"{id_column} dtype differs"
    assert np.array_equal(base[id_column].to_numpy(), head[id_column].to_numpy()), (
        f"{id_column} differs or is in a different order"
    )
    assert np.array_equal(base[id_column].to_numpy(), source_ids[entity]), (
        f"{id_column} does not follow the dataset's order"
    )
    weights_equal = np.array_equal(base["weight"].to_numpy(), head["weight"].to_numpy())
    report["entities"][entity] = {
        "rows": len(base),
        "columns": list(base.columns),
        "ids_identical_and_in_dataset_order": True,
        "weights_bitwise_identical": bool(weights_equal),
    }
    print(f"{entity}: {len(base)} rows; {id_column} identical, same order, matches "
          f"the dataset; weights identical: {weights_equal}; columns {list(base.columns)}")

# 2. Per-variable sums and row comparisons.
for name in VARIABLES:
    entity = next(e for e in entities if name in tables["base"][e].columns)
    assert name in tables["head"][entity].columns, f"{name} missing on head"
    base, head = tables["base"][entity], tables["head"][entity]
    b = base[name].to_numpy(dtype=float)
    h = head[name].to_numpy(dtype=float)
    nan_mismatch = np.isnan(b) != np.isnan(h)
    gap = np.abs(h - b)
    gap[np.isnan(gap)] = 0.0  # both NaN; one-sided NaN is counted by nan_mismatch
    differ = (gap > TOLERANCE) | nan_mismatch
    # The helper's own reduction, recomputed from the saved tables.
    recomputed = {
        side: float((tables[side][entity][name] * tables[side][entity]["weight"]).sum())
        for side in sides
    }
    for side in sides:
        assert recomputed[side] == summaries[side]["weighted_sums"][name], (
            f"{side} {name}: summary sum does not match the saved table"
        )
    base_sum = summaries["base"]["weighted_sums"][name]
    head_sum = summaries["head"]["weighted_sums"][name]
    report["variables"][name] = {
        "entity": entity,
        "base_sum": base_sum,
        "head_sum": head_sum,
        "difference": head_sum - base_sum,
        "rows_that_differ": int(differ.sum()),
        "largest_abs_difference": float(gap.max()) if len(gap) else 0.0,
        "bitwise_identical": bool(np.array_equal(b, h, equal_nan=True)),
        "nonzero_rows_base": int(np.count_nonzero(b)),
        "nonzero_rows_head": int(np.count_nonzero(h)),
    }
    id_column = f"{entity}_id"
    for i in np.flatnonzero(differ):
        report["differing_rows"].append({
            "variable": name,
            "entity": entity,
            id_column: int(base[id_column].iloc[i]),
            "base": float(b[i]),
            "head": float(h[i]),
        })

out_path.write_text(json.dumps(report, indent=1))


def money(x):
    return f"{x:,.2f}"


def gap_text(x):
    return "0" if x == 0 else f"{x:+.6g}"


print()
print("| Variable | Entity | Base sum | Head sum | Difference | Rows that differ | Largest abs. difference |")
print("|---|---|---:|---:|---:|---:|---:|")
for name, row in report["variables"].items():
    print(f"| `{name}` | {row['entity']} | {money(row['base_sum'])} | {money(row['head_sum'])} | "
          f"{gap_text(row['difference'])} | {row['rows_that_differ']:,} | "
          f"{gap_text(row['largest_abs_difference']).lstrip('+')} |")
print()
moved = [name for name, row in report["variables"].items() if row["rows_that_differ"]]
print("bitwise identical:", {n: r["bitwise_identical"] for n, r in report["variables"].items()})
print("variables that move:", moved or "none")
print("differing rows (up to 20):", report["differing_rows"][:20])
print("VERDICT:", "differs" if moved or any(
    r["difference"] != 0 for r in report["variables"].values()) else "zero-diff")
Checks: registry_check.py and time_stats.py

registry_check.py

"""Print each pair variable's entity and value type in the tree on PYTHONPATH."""
import json
import sys

import policyengine_us
from policyengine_us import CountryTaxBenefitSystem, Microsimulation
from policyengine_us.system import DEFAULT_DATASET, DEFAULT_DATASET_SHA256

VARIABLES = (
    "tob_revenue_oasdi",
    "tob_revenue_medicare_hi",
    "tob_revenue_total",
    "income_tax_chapter_1",
    "income_tax",
    "state_income_tax",
    "social_security",
    "household_net_income",
    "household_benefits",
)
system = CountryTaxBenefitSystem()
out = {
    "package": policyengine_us.__file__,
    "DEFAULT_DATASET": DEFAULT_DATASET,
    "DEFAULT_DATASET_SHA256": DEFAULT_DATASET_SHA256,
    "Microsimulation.default_dataset": Microsimulation.default_dataset,
    "microsimulation_mro": [cls.__name__ for cls in Microsimulation.__mro__],
    "variables": {},
}
for name in VARIABLES:
    variable = system.variables.get(name)
    out["variables"][name] = (
        None
        if variable is None
        else [variable.entity.key, variable.value_type.__name__]
    )
json.dump(out, sys.stdout, indent=1)
print()

time_stats.py

"""Print GNU time -v figures from a side's stderr: python time_stats.py <side.err>."""
import json
import re
import sys

text = open(sys.argv[1]).read()


def field(label):
    match = re.search(rf"^\s*{re.escape(label)}: (.*)$", text, re.M)
    return match.group(1).strip() if match else None


rss_kb = int(field("Maximum resident set size (kbytes)"))
print(json.dumps({
    "elapsed": field("Elapsed (wall clock) time (h:mm:ss or m:ss)"),
    "max_rss_kb": rss_kb,
    "max_rss_gib": round(rss_kb / 1024**2, 2),
    "max_rss_gb": round(rss_kb * 1024 / 1e9, 2),
    "exit_status": int(field("Exit status")),
    "user_seconds": float(field("User time (seconds)")),
    "system_seconds": float(field("System time (seconds)")),
}, indent=1))

Generated by Claude Code

@MaxGhenis

Copy link
Copy Markdown
Contributor Author

Ready for hub review
IMPACT PAIR 2030: zero-diff
head c545c47
base 0d93391

The brief named head 6835df5c and base 0ea88946. Before the pair started, the PR's head had moved to c545c478, whose merge base with main is 0d933918, so this pair measures c545c478 against 0d933918. The base side's numbers are also in the progress comment; its 4,396 s figure also covers saving the tables, as the Runs section below says.

The PR's head moved three times after the pair started. As the brief asked, this pair still measures the head it started with, and it ran none of the three.

  • At 09:54Z, while the head side ran, the head moved to f8ad24eb, and at 11:22Z, after the pair, to 3f270005. Both keep merge base 0d933918 and change no formula, system.py, parameter, data or lock file. Apart from docs and tests, their only change is to calculate_chunked.py: the check that rejects a chunk holding exactly one of several marital units now runs before the helper builds any chunk. Every chunk here holds 7,155 households, and so at least 7,155 marital units, so that check can't fire under either version.

  • At 11:32Z the head moved to 205f1b92, which merges main at c0689058 and so moves the PR's merge base to c0689058bba98631634f0d6aa1389fd8a50c0f4b. Main's 45 commits from 0d933918 to c0689058 change 294 files, among them system.py and 187 files under policyengine_us/variables/ and policyengine_us/parameters/. In pyproject.toml and uv.lock they change only the package's own version, from 2.91.11 to 2.92.17. The PR's own changes against c0689058 touch the same files as its changes against 0d933918, and add and remove the same lines in every file except calculate_chunked.py (the check above). This pair didn't run those changes on main's newer code, so it doesn't measure 205f1b92 against its merge base.

  • Dataset: hf://datasets/policyengine/populace-us/populace_us_2024.h5@populace-us-2024-spm-20260909, which is Microsimulation.default_dataset in both trees. DEFAULT_DATASET and DEFAULT_DATASET_SHA256 are identical in both trees. The downloaded file (Hugging Face snapshot 9a814a3b, 826,917,837 bytes) has SHA-256 6496cc4393d4d3c6574f76eca231de5898c803b9067645591fd5c4d3e65aee84, which matches DEFAULT_DATASET_SHA256. The helper resolves the URI itself and passes dataset objects to Microsimulation, so the model's own hash check of the default file doesn't run on this path. The sha256sum below is the check: I ran it before the base side and again after the head side. The Hugging Face cache holds one tag reference, one snapshot and one blob, and the head's registry check wrote all three before the base side started, so both sides read these bytes. I downloaded the file before the base side, so neither side's time includes the download.

  • Year 2030; 8 chunks on both sides, each of 7,155 households. Both sides load calculate_chunked.py from the head tree (c545c478) by path.

  • One environment, the head's uv sync (Python 3.14.6, policyengine-core 3.34.2), serves both sides, since pyproject.toml and uv.lock don't differ between base and head. PYTHONPATH picks the tree. Each side's summary.json names the tree it imported, and the registry check finds TobCalculationMixin among the bases of the head's Microsimulation only.

  • dataset_end_year: Microsimulation accepted it on both sides. The file stores _time_period 2024, and the helper builds each chunk as a USSingleYearDataset for that year. Microsimulation.__init__ passes the chunk to extend_single_year_dataset(dataset, end_year=2030), which copies it into one frame per year from 2024 through 2030 and uprates the copies with growth factors from the parameter tree. Microsimulation.__init__ is identical in both trees (system.py differs only in adding TobCalculationMixin to the bases of Simulation and Microsimulation), and nothing under policyengine_us/data/ or policyengine_us/parameters/ differs, so both sides take the dataset to 2030 the same way.

Variable Entity Base sum Head sum Difference Rows that differ Largest abs. difference
tob_revenue_oasdi tax_unit 85,000,478,688.04 85,000,478,688.04 0.00 0 0
tob_revenue_medicare_hi tax_unit 60,733,320,670.50 60,733,320,670.50 0.00 0 0
tob_revenue_total tax_unit 145,733,799,367.91 145,733,799,367.91 0.00 0 0
income_tax_chapter_1 tax_unit 3,153,596,038,642.86 3,153,596,038,642.86 0.00 0 0
income_tax tax_unit 3,112,583,162,904.35 3,112,583,162,904.35 0.00 0 0
state_income_tax tax_unit 676,000,734,185.94 676,000,734,185.94 0.00 0 0
social_security person 2,097,686,306,929.41 2,097,686,306,929.41 0.00 0 0
household_net_income household 19,488,922,620,773.88 19,488,922,620,773.88 0.00 0 0
household_benefits household 2,512,646,035,391.43 2,512,646,035,391.43 0.00 0 0

The sums use the simulation weights. "Rows that differ" counts the records (tax units, people or households) whose value differs between the sides by more than 1e-6, or is missing on one side only. No record differs at all. Every saved value is equal on the two sides, and the three saved tables are byte-identical files, so every weighted sum matches exactly. Before comparing values, the script checked that each entity table has the same <entity>_id column in the same order on both sides: 79,729 tax units, 166,321 people and 57,240 households. The weight columns are identical too.

What this pair covers

  • Both sides run the head's helper on the same chunks, so the pair compares the base and head formulas under that chunking. It doesn't compare chunked results with an unchunked run.
  • The pair checks values, not the head's memory sequencing. If TobCalculationMixin didn't prepare the counterfactual states first, the head's TOB formulas would compute the same cached variables on demand, so equal values don't show that the sequencing ran.
  • It measures c545c478 against 0d933918. It doesn't measure the PR's current head, 205f1b92, against its merge base, c0689058.

Runs

  • Base: 1:13:58 for the whole process (GNU time), and 4,396 s (73.3 min) inside pair_side.py's timed block (the helper call and saving its tables). Maximum resident set size 10,357,244 KiB (9.88 GiB, 10.61 GB); exit status 0.
  • Head: 1:07:06 for the whole process (GNU time), and 3,986 s (66.4 min) inside pair_side.py's timed block (the helper call and saving its tables). Maximum resident set size 6,800,420 KiB (6.49 GiB, 6.96 GB); exit status 0.
  • The sides ran one at a time, and no other workload ran during either side. A watcher read the side's peak resident set size from /proc every 30 s, so that I'd see early if a side passed 14 GB. During the base side I also wrote compare.py and read a few files. Neither side passed 14 GB (the base side peaked at 10.61 GB, the head side at 6.96 GB), and the VM killed neither, so I didn't rerun with 16 chunks.

Dropped variables: none. Before the pair, both trees' variable registries listed all nine variables on the same entities, each a float. Both runs computed every variable, and the helper's guard rejected none.

Commands

# Once on this VM
curl -sI https://huggingface.co/api/datasets/policyengine/populace-us | grep -i '^HTTP'
#   HTTP/1.1 200 Connection Established   (the proxy's reply to CONNECT)
#   HTTP/2 200
apt-get install -y time   # the VM had no GNU time

# In a fresh parent directory
git clone -q --filter=blob:none https://github.com/PolicyEngine/policyengine-us.git pe && cd pe
git fetch -q origin pull/10561/head:pr-10561
HEAD_SHA=$(git rev-parse pr-10561); BASE_SHA=$(git merge-base origin/main $HEAD_SHA); echo $HEAD_SHA $BASE_SHA
# at run time: c545c47877c70b367a6fab2322037e453364997e 0d933918a82827b92243c0a43d0d8d03571604fd
git diff --stat $BASE_SHA $HEAD_SHA -- pyproject.toml uv.lock   # empty
git diff --stat $BASE_SHA $HEAD_SHA -- policyengine_us/data/ policyengine_us/parameters/   # empty
git worktree add -q --detach ../head $HEAD_SHA
git worktree add -q --detach ../base $BASE_SHA
cd ../head && (time uv sync) > ../uv_sync.log 2>&1   # Python 3.14.6, policyengine-core 3.34.2
cd ../pe

# Checks before the pair, from pe: each tree's registry and dataset settings.
# The head's registry check downloads the dataset.
PY=$PWD/../head/.venv/bin/python; mkdir -p ../out ../checks
PYTHONPATH=$PWD/../head $PY ../registry_check.py --resolve > ../checks/registry_head.json 2> ../checks/registry_head.err
PYTHONPATH=$PWD/../base $PY ../registry_check.py --resolve > ../checks/registry_base.json 2> ../checks/registry_base.err
DATA="$(python3 -c 'import json; print(json.load(open("../checks/registry_head.json"))["resolved_dataset"])')"
sha256sum "$DATA"   # again after the head side
$PY -I -c '
import sys, pandas as pd
with pd.HDFStore(sys.argv[1], mode="r") as s:
    keys = s.keys(); print(sorted(keys))
    print("_time_period", s["_time_period"].tolist() if "/_time_period" in keys else "absent")
' "$DATA"   # _time_period [2024]
find "${DATA%/snapshots/*}" \( -path '*/refs/*' -o -path '*/snapshots/*' -o -path '*/blobs/*' \) -printf '%y %TF %TT %s %P\n'
#   one ref, one snapshot (a directory and its symlink) and one 826,917,837-byte blob, all 08:25:45-08:25:54Z
find "${DATA%/snapshots/*}" -newer ../out/base.start   # nothing

# The pair, from pe, one side at a time
H=$PWD/../head/policyengine_us/tools/calculate_chunked.py; PY=$PWD/../head/.venv/bin/python
export OMP_NUM_THREADS=4 OPENBLAS_NUM_THREADS=4 MKL_NUM_THREADS=4
PYTHONPATH=$PWD/../base /usr/bin/time -v $PY ../pair_side.py $H 8 2030 ../out/base > ../out/base.log 2> ../out/base.err
PYTHONPATH=$PWD/../head /usr/bin/time -v $PY ../pair_side.py $H 8 2030 ../out/head > ../out/head.log 2> ../out/head.err

# How each side ran: detached, from the parent directory, through run_side.sh (below),
# which changes to pe and runs that side's line above. watch_side.sh (below) sampled it.
setsid nohup bash ./run_side.sh base 8 > /dev/null 2>&1 < /dev/null &
./watch_side.sh base
setsid nohup bash ./run_side.sh head 8 > /dev/null 2>&1 < /dev/null &   # after the base side and the progress comment
./watch_side.sh head

# The comparison, from pe. Its second run, identical to the first, saved the table above.
$PY ../compare.py ../out/base ../out/head ../out/compare.json > ../out/compare.txt 2> ../out/compare.err
for t in tax_unit person household; do cmp ../out/base/$t.pkl ../out/head/$t.pkl && echo "$t identical"; done

# The PR's newer heads, after the pair. This fetch ran at 10:50Z, when the head was f8ad24eb.
git fetch -q origin pull/10561/head:pr-10561-new && N=$(git rev-parse pr-10561-new)   # f8ad24ebd189d7f14b709b94d9c3c8b93e5773ee
git log --format='%h %cI %s' $HEAD_SHA..$N
git merge-base origin/main $N   # 0d933918a82827b92243c0a43d0d8d03571604fd
git diff --stat $HEAD_SHA $N
git diff --name-only $HEAD_SHA $N -- policyengine_us/variables policyengine_us/parameters policyengine_us/system.py policyengine_us/data policyengine_us/tools/tob.py pyproject.toml uv.lock   # empty
git diff $HEAD_SHA $N -- policyengine_us/tools/calculate_chunked.py
# This fetch ran at 11:42Z.
git fetch -q origin main pull/10561/head:pr-10561-latest && L=$(git rev-parse pr-10561-latest)   # 205f1b927f34ddbad339b2c683067250fffc7f62
git log --format='%h %cI %s' --first-parent $N..$L
git diff --stat $N $L^1   # 3f270005: docs and one test
MB=$(git merge-base origin/main $L); echo $MB   # c0689058bba98631634f0d6aa1389fd8a50c0f4b
git rev-list --count $BASE_SHA..$MB; git diff --name-only $BASE_SHA $MB | wc -l; git diff --name-only $BASE_SHA $MB -- policyengine_us/variables policyengine_us/parameters | wc -l   # 45, 294, 187
git diff --stat $BASE_SHA $MB -- policyengine_us/system.py pyproject.toml uv.lock policyengine_us/data
git diff $BASE_SHA $MB -- pyproject.toml uv.lock | grep -E '^[-+][^-+]'   # only the version, 2.91.11 to 2.92.17
for f in $(git diff --name-only $BASE_SHA $HEAD_SHA | grep -v -E '^(docs/|changelog.d/|policyengine_us/tests/)'); do
  a=$(git diff $BASE_SHA $HEAD_SHA -- $f | grep -E '^[-+]' | grep -v -E '^(\+\+\+|---) ' | sha256sum); b=$(git diff $MB $L -- $f | grep -E '^[-+]' | grep -v -E '^(\+\+\+|---) ' | sha256sum)
  [ "$a" = "$b" ] && echo "same $f" || echo "different $f"
done   # "different" only for calculate_chunked.py

The scripts sit in the parent directory and run by path from pe, so Python puts the parent directory first on sys.path and leaves out pe, a main checkout. PYTHONPATH then decides which policyengine_us they import.

pair_side.py is the brief's script, unchanged.

pair_side.py
"""One side of #10561's pair: python pair_side.py <helper.py> <chunks> <year> <out_dir>."""
import importlib.util
import json
import sys
import time
from pathlib import Path

helper, chunks, year, out = sys.argv[1], int(sys.argv[2]), int(sys.argv[3]), Path(sys.argv[4])
VARIABLES = (
    "tob_revenue_oasdi",
    "tob_revenue_medicare_hi",
    "tob_revenue_total",
    "income_tax_chapter_1",
    "income_tax",
    "state_income_tax",
    "social_security",
    "household_net_income",
    "household_benefits",
)
spec = importlib.util.spec_from_file_location("hub_calculate_chunked", helper)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)

import policyengine_us  # noqa: E402
from policyengine_us import Microsimulation  # noqa: E402

start = time.monotonic()
result = module.calculate_chunked(
    lambda dataset: Microsimulation(dataset=dataset, dataset_end_year=year),
    VARIABLES,
    year,
    chunks,
    dataset=Microsimulation.default_dataset,
)
out.mkdir(parents=True, exist_ok=True)
for entity, table in result.entities.items():
    table.to_pickle(out / f"{entity}.pkl")
summary = {
    "package": str(Path(policyengine_us.__file__).parent),
    "chunks": chunks,
    "year": year,
    "weighted_sums": result.weighted_sums,
    "wall_seconds": time.monotonic() - start,
}
(out / "summary.json").write_text(json.dumps(summary, indent=1, sort_keys=True))
print(json.dumps(summary, sort_keys=True), flush=True)
compare.py (the comparison)
"""Compare the pair's two sides: python compare.py <base_out> <head_out> <report.json>.

For each variable: the weighted sum on each side (from summary.json, checked
against a recomputation from the saved tables), their difference (head minus
base), the number of rows whose values differ by more than 1e-6, and the
largest absolute difference. Before comparing values it checks that both sides
saved the same entity tables with the same columns, and that each table's
<entity>_id column holds the same IDs in the same order.
"""

import json
import sys
from pathlib import Path

import numpy as np
import pandas as pd

base_dir, head_dir, report_path = (Path(arg) for arg in sys.argv[1:4])
TOLERANCE = 1e-6
MAX_LISTED = 20
ORDER = (
    "tob_revenue_oasdi",
    "tob_revenue_medicare_hi",
    "tob_revenue_total",
    "income_tax_chapter_1",
    "income_tax",
    "state_income_tax",
    "social_security",
    "household_net_income",
    "household_benefits",
)

summaries = {
    side: json.loads((directory / "summary.json").read_text())
    for side, directory in (("base", base_dir), ("head", head_dir))
}
for key in ("year", "chunks"):
    if summaries["base"][key] != summaries["head"][key]:
        raise SystemExit(f"The sides differ in {key}.")
names = [n for n in ORDER if n in summaries["base"]["weighted_sums"]]
for side in ("base", "head"):
    if set(summaries[side]["weighted_sums"]) != set(names):
        raise SystemExit(f"The {side} side computed other variables.")

entities = sorted(p.stem for p in base_dir.glob("*.pkl"))
if entities != sorted(p.stem for p in head_dir.glob("*.pkl")):
    raise SystemExit("The sides saved different entity tables.")

report = {
    "year": summaries["base"]["year"],
    "chunks": summaries["base"]["chunks"],
    "tolerance": TOLERANCE,
    "entities": {},
    "variables": {},
    "listed_rows": [],
}
candidates = {}
for entity in entities:
    base = pd.read_pickle(base_dir / f"{entity}.pkl")
    head = pd.read_pickle(head_dir / f"{entity}.pkl")
    id_column = f"{entity}_id"
    if list(base.columns) != list(head.columns):
        raise SystemExit(f"{entity}: columns differ: {list(base.columns)} vs {list(head.columns)}")
    if id_column not in base.columns:
        raise SystemExit(f"{entity}: no {id_column} column.")
    base_ids, head_ids = base[id_column].to_numpy(), head[id_column].to_numpy()
    if base_ids.shape != head_ids.shape or not np.array_equal(base_ids, head_ids):
        raise SystemExit(f"{entity}: {id_column} differs in content or order.")
    if not pd.Index(base_ids).is_unique:
        raise SystemExit(f"{entity}: {id_column} is not unique.")
    base_weight = base["weight"].to_numpy(dtype=float)
    head_weight = head["weight"].to_numpy(dtype=float)
    report["entities"][entity] = {
        "rows": int(len(base)),
        "ids_identical_in_order": True,
        "weights_identical": bool(np.array_equal(base_weight, head_weight)),
        "weight_sum_base": float(base_weight.sum()),
        "weight_sum_head": float(head_weight.sum()),
    }
    for name in base.columns:
        if name in (id_column, "weight"):
            continue
        base_values = base[name].to_numpy(dtype=float)
        head_values = head[name].to_numpy(dtype=float)
        base_nan, head_nan = np.isnan(base_values), np.isnan(head_values)
        absolute = np.abs(head_values - base_values)
        differs = (absolute > TOLERANCE) | (base_nan != head_nan)
        identical = (base_values == head_values) | (base_nan & head_nan)
        base_sum = summaries["base"]["weighted_sums"][name]
        head_sum = summaries["head"]["weighted_sums"][name]
        # The helper's own reduction, in source order, on each saved table.
        recomputed = {
            "base": float((base[name] * base["weight"]).sum()),
            "head": float((head[name] * head["weight"]).sum()),
        }
        finite = absolute[~np.isnan(absolute)]
        report["variables"][name] = {
            "entity": entity,
            "base_sum": base_sum,
            "head_sum": head_sum,
            "difference": head_sum - base_sum,
            "rows_differ": int(differs.sum()),
            "rows_not_bitwise_identical": int((~identical).sum()),
            "largest_abs_difference": float(finite.max()) if finite.size else 0.0,
            "nan_rows_base": int(base_nan.sum()),
            "nan_rows_head": int(head_nan.sum()),
            "sums_match_tables": recomputed["base"] == base_sum
            and recomputed["head"] == head_sum,
        }
        order = np.flatnonzero(differs)
        order = order[np.argsort(-np.nan_to_num(absolute[order], nan=np.inf), kind="stable")]
        candidates[name] = [
            {
                "entity": entity,
                "id": base_ids[i].item(),
                "variable": name,
                "base": float(base_values[i]),
                "head": float(head_values[i]),
                "difference": float(head_values[i] - base_values[i]),
                "weight": float(base_weight[i]),
            }
            for i in order[:MAX_LISTED]
        ]

# Up to MAX_LISTED differing rows, largest differences first, taken in turn
# from each variable that moves so that every one of them is represented.
moving = [n for n in names if candidates.get(n)]
rank = 0
while len(report["listed_rows"]) < MAX_LISTED and any(rank < len(candidates[n]) for n in moving):
    for name in moving:
        if rank < len(candidates[name]) and len(report["listed_rows"]) < MAX_LISTED:
            report["listed_rows"].append(candidates[name][rank])
    rank += 1
report["variables_that_move"] = moving
report["zero_diff"] = not moving and all(
    v["difference"] == 0 for v in report["variables"].values()
)
report_path.write_text(json.dumps(report, indent=1))


def money(x):
    return f"{x:,.2f}" if abs(x) >= 0.005 or x == 0 else f"{x:.3e}"


print("| Variable | Entity | Base sum | Head sum | Difference | Rows that differ | Largest abs. difference |")
print("|---|---|---:|---:|---:|---:|---:|")
for name in names:
    v = report["variables"][name]
    print(
        f"| `{name}` | {v['entity']} | {money(v['base_sum'])} | {money(v['head_sum'])} | "
        f"{money(v['difference'])} | {v['rows_differ']:,} | {v['largest_abs_difference']:.6g} |"
    )
print()
print(json.dumps({k: report[k] for k in ("entities", "variables_that_move", "zero_diff")}, indent=1))
for name in names:
    v = report["variables"][name]
    if not v["sums_match_tables"] or v["nan_rows_base"] or v["nan_rows_head"] or v["rows_not_bitwise_identical"]:
        print(name, {k: v[k] for k in ("sums_match_tables", "nan_rows_base", "nan_rows_head", "rows_not_bitwise_identical")})
registry_check.py (the checks before the pair)
"""Print the tree on PYTHONPATH, its dataset settings and the pair's variables.

python registry_check.py [--resolve]

With --resolve, also resolve (and if needed download) the default dataset and
print its local path. Run it like pair_side.py, from the same directory, so
sys.path matches the pair's runs.
"""
import json
import sys

import policyengine_us
from policyengine_us import CountryTaxBenefitSystem, Microsimulation
from policyengine_us.system import (
    DEFAULT_DATASET,
    DEFAULT_DATASET_SHA256,
    _resolve_dataset_path,
)

VARIABLES = (
    "tob_revenue_oasdi",
    "tob_revenue_medicare_hi",
    "tob_revenue_total",
    "income_tax_chapter_1",
    "income_tax",
    "state_income_tax",
    "social_security",
    "household_net_income",
    "household_benefits",
)
system = CountryTaxBenefitSystem()
out = {
    "package": policyengine_us.__file__,
    "microsimulation_mro": [c.__name__ for c in Microsimulation.__mro__],
    "DEFAULT_DATASET": DEFAULT_DATASET,
    "DEFAULT_DATASET_SHA256": DEFAULT_DATASET_SHA256,
    "Microsimulation.default_dataset": Microsimulation.default_dataset,
    "variables": {},
}
for name in VARIABLES:
    variable = system.variables.get(name)
    out["variables"][name] = (
        None
        if variable is None
        else [variable.entity.key, variable.value_type.__name__]
    )
if "--resolve" in sys.argv:
    out["resolved_dataset"] = _resolve_dataset_path(DEFAULT_DATASET)
json.dump(out, sys.stdout, indent=1)
print()
run_side.sh and watch_side.sh (with <parent> in place of the parent directory's absolute path)
#!/bin/bash
# run_side.sh <side> <chunks>: one side of the pair, using the brief's command.
set -u
cd <parent>/pe
side=$1; chunks=$2
H=$PWD/../head/policyengine_us/tools/calculate_chunked.py; PY=$PWD/../head/.venv/bin/python
export OMP_NUM_THREADS=4 OPENBLAS_NUM_THREADS=4 MKL_NUM_THREADS=4
date -u +%FT%TZ > ../out/$side.start
PYTHONPATH=$PWD/../$side /usr/bin/time -v $PY ../pair_side.py $H $chunks 2030 ../out/$side > ../out/$side.log 2> ../out/$side.err
echo $? > ../out/$side.exit
date -u +%FT%TZ > ../out/$side.end
#!/bin/bash
# watch_side.sh <side>: sample the side's Python process every 30 s into
# ../out/<side>.rss.tsv; exit when the side ends or its peak RSS passes 14 GB.
side=$1; out=<parent>/out; limit_kib=13671875   # 14e9 bytes
pat="^<parent>/pe/\.\./head/\.venv/bin/python \.\./pair_side\.py .* \.\./out/$side\$"
while [ ! -f $out/$side.exit ]; do
  pid=$(pgrep -f "$pat" | head -1)
  if [ -n "$pid" ] && [ -r /proc/$pid/status ]; then
    rss=$(awk '/^VmRSS/{print $2}' /proc/$pid/status); hwm=$(awk '/^VmHWM/{print $2}' /proc/$pid/status)
    echo -e "$(date -u +%FT%TZ)\t$rss\t$hwm" >> $out/$side.rss.tsv
    if [ -n "$hwm" ] && [ "$hwm" -gt $limit_kib ]; then echo "PEAK RSS PASSED 14 GB: VmHWM=${hwm} KiB"; exit 0; fi
  fi
  sleep 30
done
echo "SIDE $side ENDED exit=$(cat $out/$side.exit) $(grep -E 'Maximum resident|Elapsed|Exit status|signal' $out/$side.err | tr -s ' \t' ' ' | tr '\n' ';')"

Generated by Claude Code

@MaxGhenis
MaxGhenis merged commit 8c005b9 into main Nov 3, 2026
38 checks passed
@MaxGhenis

Copy link
Copy Markdown
Contributor Author

US + core hub merge audit, us#10561 at 205f1b927f34ddbad339b2c683067250fffc7f62 (squash): Add household-chunked microsimulation calculations.

  • Review: independent cloud reviewers asked for changes at 6835df5c68 and c545c47877 (round 2, round 3). Round 3 found no wrong result and asked the PR to cut its tests' CI time. An Opus 5.5 agent reviewed the fix round at f8ad24ebd1 and approved with four wording notes; it measured the Microsimulation job at 49 min 8 s of its 60-minute limit. The hub applied the notes at 3f27000516, which corrects a docs sentence on the chunk-file tests and a test comment on file order, and in the PR body. No formula, parameter or test assertion changed after that review. The hub merged main c0689058bb into 3f27000516 with GitHub's branch update, giving 205f1b927f, and each file's patch matches the one at 3f27000516.
  • Gates the hub checked live at merge:
    • gh pr checks exits 0;
    • the latest Pull request run at the head concluded success;
    • MERGEABLE;
    • not a draft;
    • no reviewer's latest review requests changes;
    • the semantic-overlap guard;
    • head pinned.
  • Impact (2026 pair and 2030 pair, both at c545c47877 against main 0d933918a8): no output moves in 2026 or 2030. In each year, every record matches in all nine variables we compared: the three TOB revenue variables, income_tax_chapter_1, income_tax, state_income_tax, social_security, household_net_income and household_benefits. Apart from the main merge, the commits after c545c47877 change docs, tests and where calculate_chunked.py checks for a single marital unit, and no file under variables/ or parameters/.
  • Axiom: n/a: microsimulation tooling, not a rule.

🤖 Generated with Claude Code

@MaxGhenis
MaxGhenis deleted the chunked-microsimulation-calculate branch November 3, 2026 12:39
MaxGhenis pushed a commit that referenced this pull request Nov 4, 2026
Main's #10561 scans variables for population signals. gov/states/_spouses.py
tests marital_unit.count == 1, the same one-marital-unit test as
_ssi_spouses.py, and calculate_chunked already rejects a chunk holding
exactly one of several marital units, so both give the whole file's
answer in every chunk it allows.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants