Repository navigation
Add household-chunked microsimulation calculations - #10561
Conversation
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>
|
Picking this up for the US hub. I'll post a hand-off comment here when it's ready. Generated by Claude Code |
|
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>
|
Ready for hub review Head: What was asked, and how it was resolvedMerge
Variable-centric files (
The docstring now says the factory receives a path for this layout. Docs Rankers: Guard off the recursion path ( Audit: I confirmed the four population formulas from the code:
Tests (one file per process, local 4 vCPU VM)
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 New tests:
Frame depth (deepest Python stack, Mutants (each in its own process, code restored after each, clean tree confirmed)
CIPull request run 41481260553 at Other judgment calls and limits of the new path
Proposed follow-ups
Generated by Claude Code |
|
Ready for hub review The pair stopped before computing anything, because the measuring environment could not download the default dataset. Its egress proxy refused the connection to
To rerun, use an environment whose network access allows 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.errpair_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 |
|
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.
Generated by Claude Code |
|
Ready for hub review Independent review of the round-2 head. The head hasn't moved since FindingsP1. A variable that a reform adds and that the file stores is dropped from chunks with no error
Witness: a 4-household variable-centric file storing 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)
continueSet 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 ( P2. A chunk holding exactly one marital unit pairs SSI spouses differently: a silent wrong SSI result
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 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 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
Witness: a 3-household variable-centric file, plain With the stored zero, TOB total equals the whole chapter 1 tax. Replacement: add both names to # The TOB counterfactual holders: a stored value freezes that state.
"income_tax_without_social_security",
"income_tax_with_pre_1993_social_security",
P3. A non-HDF5 file with a
|
| 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:
- Drop each chunk's last household from the person mask: killed by
test_variable_centric_chunks_keep_periods_dtypes_and_nameandtest_skips_empty_chunks_and_preserves_native_and_projected_weights. - Write
tax_unitarrays with the household mask: killed bytest_variable_centric_file_is_elementwise_exact[1,2,3,7]. The validation file passes, 74/74. - Skip deleting the chunk directory: killed by
test_variable_centric_chunks_keep_periods_dtypes_and_nameandtest_variable_centric_chunk_file_is_deleted_after_factory_error. - Remove
spm_unit_income_decile: killed bytest_guard_covers_known_population_formulas,test_core_hook_guards_readers_across_periodsandtest_population_formulas_are_guarded_or_listed. - Restore a
calculateoverride that callssuper(): killed bytest_guard_adds_no_frames_per_dependency. - Remove the
_check_period_consistencyoverride: killed by everytest_rejects_population_aggregates_in_dependency_graphcase, among others. I capped the output at 8 lines, so I didn't recount the hand-off's 22. - (new) Write
spm_unitarrays with the household mask: killed bytest_variable_centric_file_is_elementwise_exact[1,2,3,7]. The validation file passes, 74/74. - (my own) Skip the cross-period identity check (
stored_periods[1:1]): killed bytest_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.yamlat 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
|
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
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01LwQq6GfptFq3a9AdVcFg9H
|
Ready for hub review Head: Round 3 fix round on the review of Findings and what I did
Should Judgment calls
MutantsEach ran in its own process in a scratch worktree, and
Tests at this headLocal 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.
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 CIPull request run 41556562467 at ImpactThis round changes tooling ( Proposed follow-up
Generated by Claude Code |
|
Hub review in progress Round-3 independent review of head Firm so far
StatusComplete. The final review is #10561 (comment), and it supersedes this comment. Generated by Claude Code |
|
Ready for hub review Independent round-3 review of head 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. FindingsP2. The Microsimulation CI job takes 55 of its 60 minutes
The job log's 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:
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
On a population run with arbitrary # _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."
)
P3. No test shows that a file holding a single marital unit is acceptedMy own mutant, 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
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 P3. The new reform-input test builds two reformed simulations where one will do
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 ( P3. Stale statements in the body
Round-2 findings at this head
Checked and sound
What I couldn't check
Generated by Claude Code |
|
Picking this up for the US hub. I'll post a hand-off comment here when it's ready. Generated by Claude Code |
|
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.
Generated by Claude Code |
|
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.
The earlier progress comment measured a different base, 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
|
Ready for hub review Head: 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 1. Microsimulation CI time (review P2)What changed, all in
What the tests no longer cover:
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
Per changed test, local before → after, and CI at
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, 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:
Together they would save roughly 50 s, about enough for 4 minutes. 2. Marital-unit guard before any chunk (review P3)
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 I also added 3. File holding a single marital unit (review P3)
4. A temp folder per test (review P3)
Interference runs used a loop that creates and removes
5. One reformed system in the reform-input test (review P3)
6. Body (review P3)
Tests run at this headOne file per process, each with its own
I ran no full suites and no population runs. MutantsEach mutant ran in its own git worktree and pytest process, and
Hand computations
Judgment calls
What I couldn't check
Proposed follow-up issues
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>
|
Ready for hub review The brief named head While the head side ran, the PR head moved to
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 ( Runs
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"""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
"""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()
"""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 |
|
Ready for hub review The brief named head 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.
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 What this pair covers
Runs
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.pyThe scripts sit in the parent directory and run by path from
|
|
US + core hub merge audit, us#10561 at
🤖 Generated with Claude Code |
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>
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)returnsChunkedCalculation:entitiesmaps entity names to source-order DataFrames (<entity>_id,weight, requested variables), andweighted_sumsmaps 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.h5files (<variable>/<period>groups, the population projection layout thatMicrosimulation(dataset=path)loads through Core)._is_hdfstore_formattells the layouts apart. For a variable-centric file the factory receives astrpath 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 thefinallyblock, 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 unchunkedMicrosimulation(dataset=path, reform=...)would load it; a length that differs from the entity count is aValueError. 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_positionscode, on DataFrames built from their ID and membership arrays, and those arrays must be identical in every stored period;medicaid_claiming_tax_unit_idis 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_indexandmedicaid_slcsp_state_denominator, and the rankershousehold_income_decileandspm_unit_income_decile(MicroSeries.decile_rank, which would rank within a chunk). Direct, dependent (for exampleincome_decile), baseline and branch calculations of all four are rejected.The guard overrides Core's
_check_period_consistency(period, variable)instead ofcalculateand_calculate. In Core 3.33.2 and the now-locked 3.34.2 (simulation.py:1760)Simulation._calculateis 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 raiseRecursionErroronly 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.pysetspossible_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 withmin(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_securityandincome_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 inbranch.input_variables, for all periods, and keeps every period of an input-named variable.tools/tob.pycallsdrop_inherited_values, which keeps only the exact (variable, branch, period) keysset_inputstored, 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.yamlandtob_tier_split.yamlpass 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 withMicrosimulation(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 hugen_chunkswork. Each test that checks chunk folders uses its own temp folder. RealSimulationtests 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 forincome_taxon one household: 361 guarded and unguarded (a restoredcalculateoverride gives 395). A source scan fails when a variables file that mentionsMicroSeries, reads a weight variable by name or through.weights, tests.count ==, or draws random numbers throughnp.randomordefault_rngis 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.py14,test_tob_sequencing.py1, lifecycle 16, and 2 intest_social_security_preflight.py); all pass atf8ad24ebd1, one file per process (round 4 below). Ruff formatting and lint pass on the changed files.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 atf8ad24ebd1(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 is6835df5c68.tests/core/test_calculate_chunked_validation.pytests/core/test_tob_branch_lifecycle.pytests/microsimulation/test_tob_sequencing.pytests/microsimulation/test_calculate_chunked.pyRound 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.py16 / 904.8 s / 2,460 MB;test_social_security_preflight.py57 / 5.6 s. Round 3 adds validation cases for a reform-defined stored input (raises), a chunk holding one of several marital units (raises, withhousehold_chunks), and a non-HDF5.h5file (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_oasdi13,tob_revenue_total4,tob_revenue_medicare_hi6,tob_tier_split3, all passed (5.17 to 5.54 GB peak by the round-3 review); round 4 changes no formula or anything they load. CI atf8ad24ebd1: 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 atc545c47877(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:ciprofile 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.income_taxfrom the module reference. The reformed population-file case keeps its unchunked reference onMicrosimulation(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.shsetting applied and not committed; seconds are pytest time:c545c47877: passed / s / peak GBf8ad24ebd1: passed / s / peak GBtests/microsimulation/test_calculate_chunked.pytests/microsimulation/test_tob_sequencing.pytests/core/test_calculate_chunked_validation.pytests/code_health/test_social_security_preflight.pytests/core/test_tob_branch_lifecycle.pyScaling 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 atc545c47877.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, fortob_revenue_oasdi,tob_revenue_medicare_hiandtob_revenue_total: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_taxandhousehold_net_income(2026 pair, 2030 pair).axiom: n/a: microsimulation tooling, not a rule
Generated by Claude Code