Bind a subsampled reform's rebuilt baseline arm to baseline policy

[MaxGhenis]

Summary

Microsimulation(reform=...).subsample(...) (and Simulation(dataset=..., reform=...).subsample(...)) rebuilt the baseline arm still bound to the reform. This PR binds the rebuilt arm to its own baseline policy, the way construction binds a reform simulation's baseline arm.

Found in the round-3 review of #9698 (finding P3-1). It does not depend on #9698: _rebind_holders and _isolate_parameter_tracing are already on main (policyengine_us/spm.py).

Root cause

Core's Simulation.subsample rebuilds the baseline arm with self.get_branch("baseline"), a branch of the reform simulation, and then assigns it the replaced arm's tax-benefit system. Read in the installed core packages:

• Core 3.32.8, the version uv.lock pins. subsample assigns the system and nothing else.
  • The branch's holders keep the reform's variable objects, so a variable the reform neutralized or replaced reads as the reform in the baseline.
  • Core's clone copies every instance attribute except branches, trace, tracer, debug and _fast_cache, so the new arm's own baseline is the arm it replaced, on the population from before subsampling.
  • labor_supply_behavioral_response and capital_gains_responses compute responses whenever simulation.baseline is not None, so the baseline arm would measure responses against that stale arm.
• Core 3.32.21 (core Edit SNAP normal allotment and ACP labels #587), what CI resolves. CI runs uv lock --upgrade-package policyengine-core before testing. subsample now calls baseline._bind_to_tax_benefit_system() and sets baseline.baseline = None, which fixes the holder bindings. That binder also drops holders the reform built in another entity's population, or with storage for (or not for) an ETERNITY variable; core 3.32.8 keeps them, at construction as well as in subsample. It cannot do the US-only part:
  • point the policy's simulation back-reference at the new arm;
  • prime the policy's private parameter-tracing root, which _isolate_parameter_tracing skips unless policy.simulation is self.

Observed on main (2 earners at $50,000, UserReform from test_spm_system.py, which neutralizes income_tax and replaces household_market_income, with tracing turned on after construction, then subsample(2)):

	core 3.32.8	core 3.32.21
baseline income_tax	$0 (should be $4,016)	$4,016
baseline.tax_benefit_system.simulation is baseline	False	False
baseline.baseline is None	False (replaced arm)	True
baseline's parameter root	the shared tree	the shared tree
shared system.parameters.trace after a baseline calculation	True: the request's FullTracer is written onto the process-wide tree	False

Fix

In SPMSimulationMixin.subsample, after core rebuilds the arm:

baseline.baseline = None
baseline._drop_holders_built_for_another_policy()
baseline._rebind_holders()
baseline._isolate_parameter_tracing()

This is what __init__ gives a reform's baseline arm, and what #9698's _rebuild_baseline_branch does.

New _drop_holders_built_for_another_policy(). It applies core 3.32.21's compatibility rule, so results no longer depend on the core version:

• it drops a holder in another entity's population, or with the wrong ETERNITY storage, together with its recorded inputs;
• __init__ now runs it on the baseline arm too, where core 3.32.8 had the same gap;
• it does not run in apply_reform, which (like core's) keeps a simulation's own inputs.

All the calls are idempotent after core 3.32.21's own binding.

Does core own this? Core's generic part is already done: core #587, released in 3.32.21, rebinds holders and clears the arm's baseline. What remains is this package's own state (the policy's simulation back-reference and the private tracing root), so the fix belongs in the US wrapper. No core PR is needed for it.

Tests

All in policyengine_us/tests/core/test_spm_system.py, next to the existing subsample test (SPM CI group):

1. test_subsampled_reform_baseline_and_clones_calculate_original_tax, run for both Simulation and Microsimulation. After subsample:

   • baseline income_tax is $4,016 and the reform's is $0;
   • household_market_income is $50,000 in the baseline and $123 in the reform;
   • each arm keeps its policy object, and every population, holder and recorded input belongs to its arm's policy, with each holder in its variable's entity's population and with storage that matches whether the variable is ETERNITY;
   • the arm has no baseline of its own;
   • the arm is traced through its own primed root, and the shared tree stays untraced;
   • clones of both arms recompute the right tax, also after delete_arrays, and are bound to their own policy.

   The $4,016 matches what the existing test_reform_baseline_and_clones_calculate_original_tax asserts for an ordinary, unreformed simulation of the same households.

2. test_subsampling_never_changes_which_policy_a_holder_belongs_to, a Hypothesis property test:

   • Inputs: three reform kinds, a set of variables calculated on either arm before subsampling, and one or two subsamples with random n, seed and quantize_weights. The reform kinds are:
     • replaces variables;
     • adds a reform-only variable and redefines the ETERNITY is_household_head as yearly;
     • changes a parameter.
   • Invariant after every subsample: each arm keeps its policy, and every holder, population and recorded input belongs to it and suits it. The baseline arm has no baseline and holds nothing reform-only.
   • Budget: max_examples=2 plus one explicit example, with derandomize=True, matching the repo's other simulation-backed Hypothesis tests. The explicit example drives both deletion paths: the reform-only variable, and the redefined variable's yearly holder. A guard asserts that both holders came through the rebuild.

Fails on main, passes here, run locally and labelled by core version:

	core 3.32.8 (lock)	core 3.32.21 (CI)
main	3 new tests fail (baseline income_tax $0, $4,016 expected; the property test's explicit example finds the stale baseline.baseline). Existing subsample test passes.	3 new tests fail (policy.simulation is the replaced arm). Existing subsample test passes.
this PR	4 passed	4 passed
this PR without the new drop	the property test fails: is_household_head keeps yearly storage	(core drops it first)

Invariants

• Subsampling never changes which policy an arm, or any holder in it, belongs to.
• A reform's baseline arm computes baseline policy: the original tax, no reform-only state, and no baseline of its own, so no behavioural responses.
• A traced baseline arm records against its own primed root and never marks the shared parameter tree.

The property test covers the first two; the regression test covers all three.

Impact

No formula, parameter or variable changes. The new lines run only for a reform simulation's baseline arm, in subsample and at construction.

• Default microsimulations (no reform) have no baseline arm and run none of this code.
• Reform simulations that are never subsampled run only the construction-time drop. It removes a holder only when the reform changes that variable's entity or ETERNITY-ness, and core ≥ 3.32.21 already drops it first.
• No shipped reform does that.
  • The 140 registry creators that accept bypass override 191 baseline variables when applied, and none changes entity or ETERNITY-ness (script: scan_reforms.py, run 2026-10-07).
  • The scan misses creators outside that registry: the taxsim class, NYC STC, Rhode Island, and Trustees (no bypass). Review r2 checked those statically. Their overrides keep entity and period, and Trustees changes parameters only.

So nothing they compute changes, and no microsim run is needed.

• GitHub code search (2026-10-07, default branches) finds no subsample calls in policyengine.py, policyengine-api, policyengine-api-v2 or policyengine-household-api. In this repo the only callers are tests, and the only ones that subsample a reform simulation are the RUN_HEAVY_TESTS=1-gated test_behavioral_structural.py and test_lsr_cg_interaction.py.
• Changes for subsampled reform workflows. On core < 3.32.21, a structural reform's baseline arm now computes baseline policy. The arm also no longer has a baseline of its own, so the simulation.baseline is None guards in the labour-supply and capital-gains responses now skip responses there. On core ≥ 3.32.21 values were already right, and the change is the policy back-reference and tracing isolation.
• Dynamic scores of a subsampled reform also need Measure a branch's behavioural responses against baseline policy, with its own cache #9913 (draft). It drops the _behavioral_response_measurements cache in _invalidate_all_caches, which subsample calls; without it, measurements taken before subsampling survive it.

CI cost

The new tests run in the SPM group (make test-other-python-spm, the Microsimulation runner). No job, matrix entry or group was added.

• Baseline (Linux, CI run 37619343624): 14:21 wall clock, 5.47 GB peak RSS for the group; test_spm_system.py takes 302 s.
• Local macOS on a shared machine, so these may overstate CI:
  • the regression test reforms the supplied system and takes 16–21 s per wrapper, below the existing test_reform_baseline_and_clones_calculate_original_tax (36–42 s per wrapper on CI);
  • the property test takes 35–42 s for two generated examples plus the explicit one; an earlier version with three generated examples took 86–136 s;
  • the 4 subsample tests run in 94 s and peak at 3.5 GB RSS in their own process.
• This PR's CI timing and memory will be in the SPM job's simulation-test-timings artifact.

Interactions

• Make branch-measured marginal tax rates match fresh re-simulation #9698 adds _rebuild_baseline_branch (the same pattern, at construction) and edits test_spm_system.py at a different spot. There's no textual overlap. Whichever lands second should add rebuilt._drop_holders_built_for_another_policy() before rebuilt._rebind_holders() there.
• Use Core cache ownership APIs #9980 (draft) removes _isolate_parameter_tracing in favour of core's per-call tracing. Whichever lands second should drop the baseline._isolate_parameter_tracing() line added here; the regression test's tracing block will flag it.

Not fixed here (follow-up filed)

Core's subsample exports to_input_dataframe(include_computed_variables=True), which core's own docstring calls "the historical unsafe behavior". Anything the reform arm calculated before subsampling becomes an input of the rebuilt population, in both arms. For example, with a standard-deduction-to-zero reform, calculating taxable_income and then subsampling leaves the baseline reading the reform's $48,755 instead of $35,400. That is a separate core defect, filed as a follow-up. The property test therefore checks bindings, not values, after pre-subsample calculations.

axiom: n/a: simulation infrastructure (subsample wiring); no policy provision changes.

🤖 Generated with Claude Code

[MaxGhenis]

Hand-off to the US + core hub

• Head: 22183cb3b9f7597d07edf5384b1d3b6d8a4e1892. Not a draft, MERGEABLE.
• Independent review: r2 (Subfleet --tier hard) APPROVE at this head, with no P1/P2/P3 findings left. r1 had requested changes (P2: incompatible holders kept on core 3.32.8; P3: test setup cost); both are fixed in 22183cb3b9. Files are in ~/reviews/us-subsample-baseline-rebind/: review-r1.md, review-r1-response.md, review-r2.md.
• Passed locally:
  • Main vs this head. Under core 3.32.8 (lock) and 3.32.21 (what CI resolves), main fails all 3 new tests and this head passes all 4 subsample tests (matrix-r2.log).
  • Mutation check. Removing only the new drop makes the property test fail on core 3.32.8.
  • SPM files on core 3.32.8. All three pass: 76 + 28 + 11 tests (spm-files-core3.32.8-r2.log).
• Pending: CI run 37674828003 is queued. The first push's 5 YAML-shard "failures" were runner-acquisition failures ("failed to be acquired (5 attempts)"); no test ran in them.
• Impact: no formula, parameter or variable changes, and no shipped reform reaches the new drop. No microsim run is needed (see the body's Impact section).
• Depends on: nothing.
• Hold or decision: none.
• Merge order:
  • If Make branch-measured marginal tax rates match fresh re-simulation #9698 lands second, it should add rebuilt._drop_holders_built_for_another_policy() before rebuilt._rebind_holders() in _rebuild_baseline_branch.
  • If Use Core cache ownership APIs #9980 lands second, it removes the new _isolate_parameter_tracing() call and turns the private-root assertions into attribution checks.
  • Measure a branch's behavioural responses against baseline policy, with its own cache #9913 covers the behavioural-measurement cache across subsample.
• Follow-ups filed as separate tasks:
  • core subsample exports pre-subsample computed values as inputs (core fix);
  • the hours-held MTR-branch test (Make branch-measured marginal tax rates match fresh re-simulation #9698 review P3-2), whose patch is ready, after Make branch-measured marginal tax rates match fresh re-simulation #9698 merges.

[MaxGhenis]

US + core hub merge audit, us#9994 at 22183cb3b9f7597d07edf5384b1d3b6d8a4e1892 (squash)

• Origin: P3-1 of the round-3 review of Make branch-measured marginal tax rates match fresh re-simulation #9698. Microsimulation(reform=...).subsample(...) rebuilt the baseline arm while it was still bound to the reform. This PR binds the rebuilt arm to its own baseline policy, using _rebind_holders and _isolate_parameter_tracing, which are already on main. It doesn't depend on Make branch-measured marginal tax rates match fresh re-simulation #9698.
• Review: independent review r2 APPROVE at this exact head (~/reviews/us-subsample-baseline-rebind/review-r2.md): no P1, P2 or P3 findings remain, and both r1 requests are resolved.
• Hand-off: from local_47e2f077. No decision is needed.
• Gates checked live at merge:
  • gh pr checks exits 0;
  • the latest Pull request run at the head concluded success;
  • MERGEABLE;
  • not a draft;
  • no CHANGES_REQUESTED review;
  • the semantic-overlap guard;
  • head pinned.
• Impact: none on default microsimulations. A run without a reform has no baseline arm and doesn't reach this code. Only subsampled reform simulations change, and they now compute baseline values from the baseline policy.
• --admin: used only because a required approving review is the sole block.