From f8161cbb0c1652639943b812578b0e093947f607 Mon Sep 17 00:00:00 2001
From: Max Ghenis <mghenis@gmail.com>
Date: Sat, 31 Oct 2026 09:20:45 -0400
Subject: [PATCH 01/15] Plan household-chunked microsimulation tooling

---
 changelog.d/chunked-microsimulation-calculate.added.md | 1 +
 1 file changed, 1 insertion(+)
 create mode 100644 changelog.d/chunked-microsimulation-calculate.added.md

diff --git a/changelog.d/chunked-microsimulation-calculate.added.md b/changelog.d/chunked-microsimulation-calculate.added.md
new file mode 100644
index 00000000000..54395da5193
--- /dev/null
+++ b/changelog.d/chunked-microsimulation-calculate.added.md
@@ -0,0 +1 @@
+Add exact household-chunked microsimulation calculations and share sequential TOB counterfactual tax calculations.

From 9b1218929ad5ee34770aa3f3bf2e571b76652052 Mon Sep 17 00:00:00 2001
From: Max Ghenis <mghenis@gmail.com>
Date: Sat, 31 Oct 2026 10:00:56 -0400
Subject: [PATCH 02/15] Add household-chunked calculations and sequential TOB
 caches

---
 docs/_toc.yml                                 |   1 +
 docs/usage/chunked-calculations.md            |  87 ++++
 policyengine_us/system.py                     |   5 +-
 .../core/test_calculate_chunked_validation.py | 229 +++++++++
 .../tests/core/test_tob_branch_lifecycle.py   | 197 ++++++++
 .../microsimulation/test_calculate_chunked.py | 442 ++++++++++++++++++
 .../microsimulation/test_tob_sequencing.py    | 101 ++++
 policyengine_us/tools/calculate_chunked.py    | 329 +++++++++++++
 policyengine_us/tools/tob.py                  |  99 ++++
 ...ncome_tax_with_pre_1993_social_security.py |  25 +
 .../income_tax_without_social_security.py     |  18 +
 .../gov/ssa/revenue/tob_revenue_oasdi.py      |  45 +-
 .../gov/ssa/revenue/tob_revenue_total.py      |  24 +-
 13 files changed, 1535 insertions(+), 67 deletions(-)
 create mode 100644 docs/usage/chunked-calculations.md
 create mode 100644 policyengine_us/tests/core/test_calculate_chunked_validation.py
 create mode 100644 policyengine_us/tests/core/test_tob_branch_lifecycle.py
 create mode 100644 policyengine_us/tests/microsimulation/test_calculate_chunked.py
 create mode 100644 policyengine_us/tests/microsimulation/test_tob_sequencing.py
 create mode 100644 policyengine_us/tools/calculate_chunked.py
 create mode 100644 policyengine_us/tools/tob.py
 create mode 100644 policyengine_us/variables/gov/ssa/revenue/income_tax_with_pre_1993_social_security.py
 create mode 100644 policyengine_us/variables/gov/ssa/revenue/income_tax_without_social_security.py

diff --git a/docs/_toc.yml b/docs/_toc.yml
index 0cb748b3196..d7ee6a32885 100644
--- a/docs/_toc.yml
+++ b/docs/_toc.yml
@@ -64,6 +64,7 @@ parts:
   - caption: Usage
     chapters:
       - file: usage/microsimulation
+      - file: usage/chunked-calculations
       - file: usage/parameter-discovery
       - file: usage/social-security-calculator
         sections:
diff --git a/docs/usage/chunked-calculations.md b/docs/usage/chunked-calculations.md
new file mode 100644
index 00000000000..9fafc21b9bf
--- /dev/null
+++ b/docs/usage/chunked-calculations.md
@@ -0,0 +1,87 @@
+# Household-chunked calculations
+
+`policyengine_us.tools.calculate_chunked.calculate_chunked` calculates numeric
+variables over a disjoint partition of households, releasing each simulation
+before constructing the next. The source dataset and result arrays remain in
+memory; the helper bounds simulation working memory by chunk size, rather than
+promising a fixed RSS limit.
+
+```python
+from policyengine_us import Microsimulation
+from policyengine_us.tools.calculate_chunked import calculate_chunked
+
+result = calculate_chunked(
+    lambda dataset: Microsimulation(dataset=dataset, dataset_end_year=2026),
+    ["tob_revenue_oasdi", "tob_revenue_medicare_hi", "tob_revenue_total"],
+    2026,
+    4,
+    dataset=Microsimulation.default_dataset,
+)
+tax_units = result.entities["tax_unit"]
+print(result.weighted_sums)
+```
+
+The factory owns reforms and simulation options. It must return a fresh
+PolicyEngine-US Microsimulation with the supplied population and row order.
+The explicit `dataset=` keyword supplies the source without first constructing
+a full-population simulation. Supported sources are `USSingleYearDataset`,
+`USMultiYearDataset`, and local or Hugging Face entity-table HDFStore files.
+Legacy variable-array and flat-file datasets must first be converted to the US
+entity-table schema. All six entity tables and person membership links are
+required. Input data is not modified.
+
+Each output DataFrame contains `<entity>_id`, `weight`, and requested variable
+columns, with exactly the source entity order. `weight` is the weight used by
+`Microsimulation.get_weights`, projected from household weights; native input
+entity weight columns are preserved in the chunk datasets. Weighted sums are
+calculated from the assembled source-order arrays, avoiding changes in reduction
+order. A requested variable can itself be an entity weight variable.
+
+## Exact partition invariants
+
+Every household gets exactly one label. Each person inherits their household's
+label; every tax unit, SPM unit, family and marital unit must have members in
+exactly one household and inherits that label. Entity IDs must be unique,
+nonmissing, fully referenced and unmodified. Boolean selection preserves the
+source entity order, and results are written back to their original row
+positions. These facts prove disjointness, complete coverage, and whole groups.
+They also preserve each household's calculation inputs and weights. All source
+years are retained, with identical structural IDs and links across years.
+Optional `household_chunks` provides arbitrary integer labels in
+`[0, n_chunks)`; otherwise consecutive household partitions are balanced. Empty
+partitions are skipped, including when chunks outnumber households.
+
+Exact calculation additionally requires every dependency to be local to a
+household. The current model has two Medicaid state reductions,
+`medicaid_slcsp_state_average_cost_index` and `medicaid_slcsp_state_denominator`.
+The helper rejects these at calculation time, including indirect dependencies
+and branches. Known Medicaid claiming-tax-unit links crossing households are
+also rejected. The model audit found no other population ranking, quantile or
+stochastic formula reductions. Custom reform formulas remain the caller's
+responsibility: cross-household reductions, order-dependent random draws and
+other global dependencies are unsupported. These are calculation restrictions,
+not an assumption that every model variable can be chunked.
+
+The tests compare a full 32-household, five-state fixture with one, two, three
+and seven chunks, plus random household assignments. They compare every entity
+value and weighted sum, check source order and unchanged input frames, and use
+weak references to check simulation release. The fixture includes households
+with multiple tax and SPM units and Social Security recipients. Independent
+pre-change TOB branch calculations protect the statutory signed allocation.
+
+## TOB working memory
+
+Counterfactual income tax without Social Security and under pre-1993 inclusion
+is cached in parent variable holders. A direct TOB calculation prepares those
+holders in separate outer calculations: current Core retains popped branches
+in active calculation frames, so the frames must finish before collection.
+Each temporary branch clones the system, clears non-input arrays, computes and
+copies only income tax, and is removed and collected before the next direct
+counterfactual. A TOB dependency nested inside another formula can retain its
+branch identities until that outer formula returns. OASDI remains old-law tax minus no-inclusion
+tax, total remains current tax minus no-inclusion tax, and HI remains the signed
+residual. No per-unit floor is introduced.
+
+The hub will add measured full-dataset peak memory and timing after running its
+memory probe. Small-fixture exactness establishes correctness, not a measured
+full-population RSS reduction.
diff --git a/policyengine_us/system.py b/policyengine_us/system.py
index c97ca49fa4f..b402a41c78d 100644
--- a/policyengine_us/system.py
+++ b/policyengine_us/system.py
@@ -35,6 +35,7 @@
     uprate_parameters,
 )
 from .tools.default_uprating import add_default_uprating
+from .tools.tob import TobCalculationMixin
 from .tools.period_branch import delete_input_values
 from .tools.social_security_inputs import (
     install_social_security_input_compatibility,
@@ -238,7 +239,7 @@ def _backfill_state_code_from_str(simulation):
         state_code_str.delete_arrays(known_period)
 
 
-class Simulation(SPMSimulationMixin, CoreSimulation):
+class Simulation(TobCalculationMixin, SPMSimulationMixin, CoreSimulation):
     """
     A simulation of the tax-benefit system for the United States,
     defined against the base simulation class in the -core package.
@@ -441,7 +442,7 @@ def _is_hdfstore_format(file_path):
         return False
 
 
-class Microsimulation(SPMSimulationMixin, CoreMicrosimulation):
+class Microsimulation(TobCalculationMixin, SPMSimulationMixin, CoreMicrosimulation):
     """
     A microsimulation of the tax-benefit system for the United States,
     defined against the base microsimulation class in the -core package.
diff --git a/policyengine_us/tests/core/test_calculate_chunked_validation.py b/policyengine_us/tests/core/test_calculate_chunked_validation.py
new file mode 100644
index 00000000000..922b754d4b8
--- /dev/null
+++ b/policyengine_us/tests/core/test_calculate_chunked_validation.py
@@ -0,0 +1,229 @@
+"""Reject invalid household partitions before constructing a simulation."""
+
+from types import SimpleNamespace
+
+import numpy as np
+import pandas as pd
+import pytest
+
+from policyengine_us.data.dataset_schema import (
+    US_ENTITIES,
+    USMultiYearDataset,
+    USSingleYearDataset,
+)
+from policyengine_us.system import system
+from policyengine_us.tools.calculate_chunked import calculate_chunked
+
+YEAR = 2026
+
+
+@pytest.fixture
+def dataset():
+    ids = np.asarray([91, 37])
+    person = pd.DataFrame({"person_id": [503, 101]})
+    for entity in US_ENTITIES:
+        if entity != "person":
+            person[f"person_{entity}_id"] = ids
+    tables = {"person": person}
+    for entity in US_ENTITIES:
+        if entity != "person":
+            tables[entity] = pd.DataFrame({f"{entity}_id": ids})
+    tables["household"]["household_weight"] = [1.5, 2.25]
+    return USSingleYearDataset(**tables, time_period=YEAR)
+
+
+def never_build(_dataset):
+    pytest.fail("Invalid input must be rejected before constructing a simulation")
+
+
+def calculate(dataset, **kwargs):
+    return calculate_chunked(
+        never_build,
+        kwargs.pop("variables", ["income_tax"]),
+        YEAR,
+        kwargs.pop("n_chunks", 2),
+        dataset=dataset,
+        **kwargs,
+    )
+
+
+@pytest.mark.parametrize("n_chunks", [0, -1, True, 1.5])
+def test_rejects_invalid_chunk_counts(dataset, n_chunks):
+    with pytest.raises(ValueError, match="positive integer"):
+        calculate(dataset, n_chunks=n_chunks)
+
+
+@pytest.mark.parametrize("variables", [[], ["income_tax", "income_tax"]])
+def test_rejects_empty_or_duplicate_variables(dataset, variables):
+    with pytest.raises(ValueError, match="nonempty.*unique"):
+        calculate(dataset, variables=variables)
+
+
+@pytest.mark.parametrize("assignments", [[0], [[0], [1]], [0.0, 1.0], [-1, 0], [0, 2]])
+def test_rejects_invalid_household_assignments(dataset, assignments):
+    with pytest.raises(ValueError, match="household_chunks"):
+        calculate(dataset, household_chunks=assignments)
+
+
+@pytest.mark.parametrize("entity", US_ENTITIES)
+def test_rejects_duplicate_entity_ids(dataset, entity):
+    table = getattr(dataset, entity)
+    table.loc[table.index[1], f"{entity}_id"] = table[f"{entity}_id"].iloc[0]
+    with pytest.raises(ValueError, match=f"{entity}_id.*unique"):
+        calculate(dataset)
+
+
+@pytest.mark.parametrize("entity", ["tax_unit", "spm_unit", "family", "marital_unit"])
+def test_rejects_entity_spanning_households(dataset, entity):
+    dataset.person[f"person_{entity}_id"] = [91, 91]
+    with pytest.raises(ValueError, match=f"{entity} spans multiple households"):
+        calculate(dataset)
+
+
+@pytest.mark.parametrize(
+    "entity", ["household", "tax_unit", "spm_unit", "family", "marital_unit"]
+)
+def test_rejects_missing_membership_target(dataset, entity):
+    dataset.person.loc[0, f"person_{entity}_id"] = 999
+    with pytest.raises(ValueError, match="existing household|Unknown.*membership"):
+        calculate(dataset)
+
+
+@pytest.mark.parametrize("entity", ["tax_unit", "spm_unit", "family", "marital_unit"])
+def test_rejects_unreferenced_group_entity(dataset, entity):
+    table = getattr(dataset, entity)
+    table.loc[len(table)] = [999]
+    with pytest.raises(
+        ValueError, match=f"Every {entity} must have at least one person"
+    ):
+        calculate(dataset)
+
+
+def test_rejects_external_claiming_tax_unit(dataset):
+    dataset.person["medicaid_claiming_tax_unit_id"] = [37, 91]
+    with pytest.raises(ValueError, match="claiming_tax_unit_id crosses households"):
+        calculate(dataset)
+
+
+def test_rejects_changed_multi_year_memberships(dataset):
+    other = dataset.copy()
+    other.time_period = str(YEAR + 1)
+    other.person["person_tax_unit_id"] = [37, 91]
+    multi = USMultiYearDataset(datasets=[dataset, other])
+    with pytest.raises(ValueError, match="memberships must be identical"):
+        calculate(multi)
+
+
+class AggregateSimulation:
+    """Exercise aggregate guards without building a country simulation."""
+
+    def __init__(self, dataset, aggregate, route):
+        self.tax_benefit_system = system
+        self.populations = {
+            entity: SimpleNamespace(
+                ids=getattr(dataset, entity)[f"{entity}_id"].to_numpy()
+            )
+            for entity in US_ENTITIES
+        }
+        self.aggregate = aggregate
+        self.route = route
+        self.branches = {}
+        self.baseline = None
+        if route in ("branch", "baseline"):
+            target = AggregateSimulation(dataset, aggregate, "direct")
+            if route == "branch":
+                self.branches["existing"] = target
+            else:
+                self.baseline = target
+
+    def calculate(self, variable, period, use_weights=False):
+        return self._calculate(variable, period)
+
+    def _calculate(self, variable, period=None):
+        if self.route == "branch":
+            return self.branches["existing"].calculate(self.aggregate, period)
+        if self.route == "baseline":
+            return self.baseline.calculate(self.aggregate, period)
+        if variable != self.aggregate:
+            return self.calculate(self.aggregate, period)
+        pytest.fail("Population aggregate must be rejected before formula execution")
+
+
+@pytest.mark.parametrize(
+    "aggregate",
+    ["medicaid_slcsp_state_average_cost_index", "medicaid_slcsp_state_denominator"],
+)
+@pytest.mark.parametrize("route", ["direct", "dependency", "branch", "baseline"])
+def test_rejects_population_aggregates_in_dependency_graph(dataset, aggregate, route):
+    variable = aggregate if route == "direct" else "income_tax"
+    with pytest.raises(ValueError, match="aggregates across households"):
+        calculate_chunked(
+            lambda chunk: AggregateSimulation(chunk, aggregate, route),
+            [variable],
+            YEAR,
+            2,
+            dataset=dataset,
+        )
+
+
+class InputSimulation:
+    def __init__(self, dataset):
+        self.dataset = dataset
+        self.tax_benefit_system = system
+        self.branches = {}
+        self.baseline = None
+        self.populations = {
+            entity: SimpleNamespace(
+                ids=getattr(dataset, entity)[f"{entity}_id"].to_numpy()
+            )
+            for entity in US_ENTITIES
+        }
+
+    def calculate(self, variable, period, use_weights=False):
+        return self._calculate(variable, period)
+
+    def _calculate(self, variable, period=None):
+        entity = system.variables[variable].entity.key
+        return getattr(self.dataset, entity)[variable].to_numpy()
+
+    def get_weights(self, variable, period):
+        entity = system.variables[variable].entity.key
+        household = self.dataset.household.set_index("household_id")
+        if entity == "household":
+            return self.dataset.household.household_weight.to_numpy()
+        assert entity == "person"
+        return household.loc[
+            self.dataset.person.person_household_id, "household_weight"
+        ].to_numpy()
+
+
+def test_skips_empty_chunks_and_preserves_native_and_projected_weights(dataset):
+    dataset.person["employment_income"] = [10_000.0, 25_000.0]
+    # Core uses projected household weights; native input columns still survive.
+    dataset.person["person_weight"] = [999.0, 888.0]
+    snapshots = []
+
+    def factory(chunk):
+        snapshots.append(chunk.person.copy())
+        return InputSimulation(chunk)
+
+    result = calculate_chunked(
+        factory,
+        ["employment_income", "household_weight"],
+        YEAR,
+        5,
+        dataset=dataset,
+    )
+    assert len(snapshots) == 2
+    np.testing.assert_array_equal(
+        result.entities["person"].person_id, dataset.person.person_id
+    )
+    np.testing.assert_array_equal(result.entities["person"].weight, [1.5, 2.25])
+    np.testing.assert_array_equal(
+        result.entities["household"].household_weight, [1.5, 2.25]
+    )
+    assert result.weighted_sums["employment_income"] == 71_250.0
+    np.testing.assert_array_equal(
+        pd.concat(snapshots).person_weight, dataset.person.person_weight
+    )
+    np.testing.assert_array_equal(dataset.person.person_weight, [999.0, 888.0])
diff --git a/policyengine_us/tests/core/test_tob_branch_lifecycle.py b/policyengine_us/tests/core/test_tob_branch_lifecycle.py
new file mode 100644
index 00000000000..0bd8010770b
--- /dev/null
+++ b/policyengine_us/tests/core/test_tob_branch_lifecycle.py
@@ -0,0 +1,197 @@
+"""Counterfactual branches retain only their returned income-tax array."""
+
+from types import SimpleNamespace
+import weakref
+
+import numpy as np
+import pytest
+
+from policyengine_us.tools import tob
+
+
+class FakeBranch:
+    def __init__(self, parent):
+        self.parent = parent
+        self.cycle = self
+        self.tax_benefit_system = SimpleNamespace(
+            neutralize_variable=lambda name: parent.events.append(("neutralize", name))
+        )
+
+    def get_holder(self, name):
+        assert name == "tax_unit_taxable_social_security"
+        return self
+
+    def set_input(self, period, values):
+        self.parent.events.append(("input", period, tuple(values)))
+
+    def tax_unit(self, name, period):
+        self.parent.events.append(("calculate", name, period))
+        if self.parent.fail:
+            raise RuntimeError("counterfactual failed")
+        return self.parent.result
+
+
+class FakeSimulation:
+    def __init__(self, *, fail=False):
+        self.events = []
+        self.branches = {}
+        self.references = []
+        self.result = np.array([-2.0, 17.0], dtype=np.float32)
+        self.fail = fail
+
+    def get_branch(self, name, clone_system):
+        assert clone_system
+        assert not self.branches
+        # Population/holder cycles must be collected before another branch.
+        assert all(reference() is None for reference in self.references)
+        branch = FakeBranch(self)
+        self.branches[name] = branch
+        self.references.append(weakref.ref(branch))
+        return branch
+
+
+@pytest.mark.parametrize("override", [None, np.array([30.0, 40.0])])
+def test_counterfactual_releases_branch_and_copies_result(monkeypatch, override):
+    simulation = FakeSimulation()
+    monkeypatch.setattr(
+        tob,
+        "drop_inherited_values",
+        lambda branch: branch.parent.events.append(("drop",)),
+    )
+
+    result = tob.income_tax_with_social_security_override(
+        simulation, 2026, taxable_social_security=override
+    )
+
+    np.testing.assert_array_equal(result, [-2.0, 17.0])
+    assert not np.shares_memory(result, simulation.result)
+    assert not simulation.branches
+    assert simulation.references[0]() is None
+    setup = (
+        [("neutralize", "tax_unit_taxable_social_security"), ("drop",)]
+        if override is None
+        else [("drop",), ("input", 2026, tuple(override))]
+    )
+    assert simulation.events == setup + [("calculate", "income_tax", 2026)]
+
+
+def test_counterfactuals_are_released_before_next_branch(monkeypatch):
+    simulation = FakeSimulation()
+    monkeypatch.setattr(tob, "drop_inherited_values", lambda branch: None)
+
+    tob.income_tax_with_social_security_override(simulation, 2026)
+    tob.income_tax_with_social_security_override(
+        simulation, 2026, taxable_social_security=np.array([30.0, 40.0])
+    )
+
+    assert len(simulation.references) == 2
+    assert all(reference() is None for reference in simulation.references)
+
+
+def test_failed_counterfactual_is_removed(monkeypatch):
+    simulation = FakeSimulation(fail=True)
+    monkeypatch.setattr(tob, "drop_inherited_values", lambda branch: None)
+
+    with pytest.raises(RuntimeError, match="counterfactual failed"):
+        tob.income_tax_with_social_security_override(simulation, 2026)
+
+    assert not simulation.branches
+
+
+NO_SS = "income_tax_without_social_security"
+OLD_LAW = "income_tax_with_pre_1993_social_security"
+OASDI = "tob_revenue_oasdi"
+HI = "tob_revenue_medicare_hi"
+TOTAL = "tob_revenue_total"
+
+
+class FakeCalculationBase:
+    default_calculation_period = 2026
+
+    def __init__(self, *, cached=(), neutralized=(), custom=()):
+        self.cached = set(cached)
+        self.neutralized = set(neutralized)
+        self.custom = set(custom)
+        self.events = []
+
+    def get_holder(self, name):
+        return SimpleNamespace(
+            variable=SimpleNamespace(is_neutralized=name in self.neutralized)
+        )
+
+    def get_array(self, name, period):
+        return np.array([1.0]) if name in self.cached else None
+
+    def calculate(self, name, period=None, *args, **kwargs):
+        self.events.append(("calculate", name, period, args, kwargs))
+        self.cached.add(name)
+        return name
+
+
+class FakeTobCalculation(tob.TobCalculationMixin, FakeCalculationBase):
+    def _uses_standard_tob_formula(self, name, period):
+        return name not in self.custom
+
+
+def _record_collection(monkeypatch, simulation):
+    monkeypatch.setattr(
+        tob.gc, "collect", lambda: simulation.events.append(("collect",))
+    )
+
+
+@pytest.mark.parametrize(
+    "requested, states",
+    [(OASDI, [NO_SS, OLD_LAW]), (HI, [NO_SS, OLD_LAW]), (TOTAL, [NO_SS])],
+)
+def test_tob_states_finish_before_requested_calculation(monkeypatch, requested, states):
+    simulation = FakeTobCalculation()
+    _record_collection(monkeypatch, simulation)
+
+    result = simulation.calculate(requested, 2026)
+
+    expected = []
+    for state in states:
+        expected.extend([("calculate", state, 2026, (), {}), ("collect",)])
+    assert simulation.events == expected + [("calculate", requested, 2026, (), {})]
+    assert result == requested
+
+
+@pytest.mark.parametrize("available", ["cached", "neutralized", "custom"])
+def test_tob_prefetch_preserves_available_or_replaced_target(monkeypatch, available):
+    # Supplied inputs and previously computed outputs both appear as cached
+    # arrays. Neutralized and reform-provided formulas also need no warmup.
+    simulation = FakeTobCalculation(**{available: [OASDI]})
+    _record_collection(monkeypatch, simulation)
+
+    simulation.calculate(OASDI, 2026)
+
+    assert simulation.events == [("calculate", OASDI, 2026, (), {})]
+
+
+@pytest.mark.parametrize(
+    "cached, states",
+    [([TOTAL], [NO_SS, OLD_LAW]), ([OASDI], [NO_SS]), ([TOTAL, OASDI], [])],
+)
+def test_hi_only_prefetches_missing_allocation_dependencies(
+    monkeypatch, cached, states
+):
+    simulation = FakeTobCalculation(cached=cached)
+    _record_collection(monkeypatch, simulation)
+
+    simulation.calculate(HI, 2026)
+
+    calculated = [event[1] for event in simulation.events if event[0] == "calculate"]
+    assert calculated == states + [HI]
+
+
+def test_tob_reuses_states_and_forwards_original_request_options(monkeypatch):
+    simulation = FakeTobCalculation(cached=[NO_SS])
+    _record_collection(monkeypatch, simulation)
+
+    simulation.calculate(OASDI, None, "person", decode_enums=False)
+
+    assert simulation.events == [
+        ("calculate", OLD_LAW, 2026, (), {}),
+        ("collect",),
+        ("calculate", OASDI, None, ("person",), {"decode_enums": False}),
+    ]
diff --git a/policyengine_us/tests/microsimulation/test_calculate_chunked.py b/policyengine_us/tests/microsimulation/test_calculate_chunked.py
new file mode 100644
index 00000000000..38d589aef44
--- /dev/null
+++ b/policyengine_us/tests/microsimulation/test_calculate_chunked.py
@@ -0,0 +1,442 @@
+"""Household chunking must preserve real model results and dataset row order.
+
+Python is needed to exercise the dataset/factory API, arbitrary partitions and
+simulation lifetime. The 32-household reference is calculated once; only its
+arrays survive, so no reference simulation remains live during chunked runs.
+"""
+
+import gc
+import weakref
+
+import numpy as np
+import pandas as pd
+import pytest
+from hypothesis import given, settings, strategies as st
+
+from policyengine_us import Microsimulation
+from policyengine_us.data.dataset_schema import (
+    US_ENTITIES,
+    USMultiYearDataset,
+    USSingleYearDataset,
+)
+from policyengine_us.system import system
+from policyengine_us.tools.calculate_chunked import calculate_chunked
+
+YEAR = 2026
+HOUSEHOLDS = 32
+VARIABLES = (
+    "tob_revenue_oasdi",
+    "tob_revenue_medicare_hi",
+    "tob_revenue_total",
+    "household_net_income",
+    "income_tax",
+    "spm_unit_net_income",
+    # Exercise outputs for all six entity tables, including person ordering.
+    "employment_income",
+    "family_weight",
+    "marital_unit_weight",
+)
+
+
+def build_chunking_dataset() -> USSingleYearDataset:
+    """Several states, multi-unit households and taxable Social Security.
+
+    Entity IDs are sparse, and every entity table has a different unsorted row
+    order and nontrivial index. Membership must come from the person links,
+    rather than coincident entity positions, ID sorting, or pandas indices.
+    """
+    records = {entity: [] for entity in US_ENTITIES}
+    state_fips = (6, 36, 48, 20, 12)
+    for household in range(HOUSEHOLDS):
+        household_id = 90_001 + 37 * household
+        # Non-binary fractions expose accidental float32 weighted reductions.
+        weight = 1.13 + household * 0.137
+        records["household"].append(
+            {
+                "household_id": household_id,
+                "household_weight": weight,
+                "state_fips": state_fips[household % len(state_fips)],
+            }
+        )
+        # Two tax units in every household; two SPM/family units in half.
+        tax_ids = [201_003 + 41 * household + offset for offset in (0, 11)]
+        spm_ids = [301_009 + 43 * household + offset for offset in (0, 13)]
+        family_ids = [401_011 + 47 * household + offset for offset in (0, 17)]
+        marital_ids = [501_019 + 53 * household + offset for offset in (0, 19, 23)]
+        for tax_id in tax_ids:
+            records["tax_unit"].append(
+                {"tax_unit_id": tax_id, "tax_unit_weight": weight}
+            )
+        for spm_id in spm_ids[: 2 if household % 2 == 0 else 1]:
+            records["spm_unit"].append(
+                {"spm_unit_id": spm_id, "spm_unit_weight": weight}
+            )
+        for family_id in family_ids[: 2 if household % 2 == 0 else 1]:
+            records["family"].append({"family_id": family_id, "family_weight": weight})
+        for marital_id in marital_ids:
+            records["marital_unit"].append(
+                {"marital_unit_id": marital_id, "marital_unit_weight": weight}
+            )
+        for member, age in enumerate((72, 68, 42, 12)):
+            second_unit = member >= 2
+            links = {
+                "household": household_id,
+                "tax_unit": tax_ids[int(second_unit)],
+                "spm_unit": spm_ids[int(second_unit and household % 2 == 0)],
+                "family": family_ids[int(second_unit and household % 2 == 0)],
+                "marital_unit": marital_ids[max(member - 1, 0)],
+            }
+            person = {
+                "person_id": 601_031 + 59 * household + member * 7,
+                "person_weight": weight,
+                "age": age,
+                "employment_income": (
+                    35_000 + household * 1_200 if member == 2 else 0.0
+                ),
+                "taxable_interest_income": (
+                    22_000 + household * 1_500 if member == 0 else 0.0
+                ),
+                "social_security_retirement": (
+                    25_000 + household * 150 if member < 2 else 0.0
+                ),
+            }
+            for entity, entity_id in links.items():
+                person[f"person_{entity}_id"] = entity_id
+                person[f"person_{entity}_role"] = "member"
+            records["person"].append(person)
+
+    tables = {}
+    for seed, (entity, rows) in enumerate(records.items()):
+        table = pd.DataFrame(rows).sample(frac=1, random_state=100 + seed)
+        table.index = 1_001 + np.arange(len(table)) * 3
+        tables[entity] = table
+    return USSingleYearDataset(**tables, time_period=YEAR)
+
+
+def make_simulation(dataset):
+    # Bypass automatic multi-year projection for this fixed-year fixture.
+    if isinstance(dataset, USSingleYearDataset):
+        dataset = USMultiYearDataset(datasets=[dataset])
+    return Microsimulation(dataset=dataset)
+
+
+@pytest.fixture(scope="module")
+def chunking_dataset():
+    return build_chunking_dataset()
+
+
+@pytest.fixture(scope="module")
+def reference(chunking_dataset):
+    simulation = make_simulation(chunking_dataset)
+    values = {}
+    weighted_sums = {}
+    weights = {}
+    for variable in VARIABLES:
+        entity = system.variables[variable].entity.key
+        calculated = simulation.calculate(variable, YEAR)
+        values[variable] = np.asarray(calculated).copy()
+        if entity not in weights:
+            weights[entity] = np.asarray(calculated.weights).copy()
+        weighted_sums[variable] = float(calculated.sum())
+    del calculated
+    # Core stores this backreference even for the global default system.
+    if getattr(simulation.tax_benefit_system, "simulation", None) is simulation:
+        del simulation.tax_benefit_system.simulation
+    del simulation
+    gc.collect()
+    for variable in (
+        "tob_revenue_oasdi",
+        "tob_revenue_medicare_hi",
+        "tob_revenue_total",
+    ):
+        assert np.any(values[variable] != 0), f"Fixture must exercise {variable}"
+    return values, weighted_sums, weights
+
+
+def assert_matches_reference(result, dataset, reference):
+    values, weighted_sums, weights = reference
+    expected_entities = {system.variables[var].entity.key for var in VARIABLES}
+    assert set(result.entities) == expected_entities
+    assert set(result.weighted_sums) == set(VARIABLES)
+    for entity, frame in result.entities.items():
+        expected_ids = getattr(dataset, entity)[f"{entity}_id"].to_numpy()
+        np.testing.assert_array_equal(frame[f"{entity}_id"], expected_ids)
+        np.testing.assert_array_equal(frame["weight"], weights[entity])
+        assert len(frame) == len(expected_ids)
+        assert not frame[f"{entity}_id"].duplicated().any()
+    for variable in VARIABLES:
+        entity = system.variables[variable].entity.key
+        np.testing.assert_array_equal(
+            result.entities[entity][variable], values[variable]
+        )
+        assert result.weighted_sums[variable] == weighted_sums[variable]
+
+
+@pytest.mark.parametrize("n_chunks", [1, 2, 3, 7])
+def test_chunked_real_model_is_elementwise_exact(n_chunks, chunking_dataset, reference):
+    # Weak references detect retained simulations and cycles before the next
+    # chunk, including branch-to-parent references created by TOB calculations.
+    simulations = []
+    seen = {entity: [] for entity in US_ENTITIES}
+
+    def factory(dataset):
+        assert all(simulation() is None for simulation in simulations)
+        single_year = (
+            dataset[YEAR] if isinstance(dataset, USMultiYearDataset) else dataset
+        )
+        household_ids = set(single_year.household.household_id)
+        people = chunking_dataset.person[
+            chunking_dataset.person.person_household_id.isin(household_ids)
+        ]
+        pd.testing.assert_frame_equal(single_year.person, people)
+        for entity in US_ENTITIES:
+            table = getattr(single_year, entity)
+            seen[entity].extend(table[f"{entity}_id"])
+            if entity not in ("person", "household"):
+                expected = getattr(chunking_dataset, entity)
+                expected = expected[
+                    expected[f"{entity}_id"].isin(people[f"person_{entity}_id"])
+                ]
+                pd.testing.assert_frame_equal(table, expected)
+        simulation = make_simulation(dataset)
+        simulations.append(weakref.ref(simulation))
+        return simulation
+
+    result = calculate_chunked(
+        factory, VARIABLES, YEAR, n_chunks, dataset=chunking_dataset
+    )
+    assert_matches_reference(result, chunking_dataset, reference)
+    assert len(simulations) == n_chunks
+    assert all(simulation() is None for simulation in simulations)
+    for entity, ids in seen.items():
+        expected_ids = getattr(chunking_dataset, entity)[f"{entity}_id"]
+        assert len(ids) == len(expected_ids)
+        assert set(ids) == set(expected_ids)
+
+
+@st.composite
+def household_partitions(draw):
+    n_chunks = draw(st.integers(min_value=1, max_value=7))
+    # Every partition is nonempty, while all remaining household assignments
+    # vary independently. Chunk labels are intentionally not contiguous rows.
+    tail = draw(
+        st.lists(
+            st.integers(min_value=0, max_value=n_chunks - 1),
+            min_size=HOUSEHOLDS - n_chunks,
+            max_size=HOUSEHOLDS - n_chunks,
+        )
+    )
+    labels = np.asarray(list(range(n_chunks)) + tail)
+    offset = draw(st.integers(min_value=0, max_value=HOUSEHOLDS - 1))
+    return n_chunks, np.roll(labels, offset)
+
+
+@settings(max_examples=4, deadline=None)
+@given(partition=household_partitions())
+def test_random_household_partitions_are_exact(partition, chunking_dataset, reference):
+    n_chunks, household_chunks = partition
+    result = calculate_chunked(
+        make_simulation,
+        VARIABLES,
+        YEAR,
+        n_chunks,
+        dataset=chunking_dataset,
+        household_chunks=household_chunks,
+    )
+    assert_matches_reference(result, chunking_dataset, reference)
+
+
+def test_multi_year_dataset_preserves_every_year_and_original_input(
+    chunking_dataset,
+):
+    previous = chunking_dataset.copy()
+    previous.time_period = str(YEAR - 1)
+    previous.person["employment_income"] = previous.person.employment_income * 0.9
+    multi = USMultiYearDataset(datasets=[previous, chunking_dataset])
+    original = multi.copy()
+    seen = []
+
+    def factory(dataset):
+        assert isinstance(dataset, USMultiYearDataset)
+        assert dataset.years == [YEAR - 1, YEAR]
+        household_ids = set(dataset[YEAR].household.household_id)
+        for year in dataset.years:
+            single = dataset[year]
+            for entity in US_ENTITIES:
+                table = getattr(original[year], entity)
+                if entity == "household":
+                    mask = table.household_id.isin(household_ids)
+                elif entity == "person":
+                    mask = table.person_household_id.isin(household_ids)
+                else:
+                    people = original[year].person
+                    member_ids = people.loc[
+                        people.person_household_id.isin(household_ids),
+                        f"person_{entity}_id",
+                    ]
+                    mask = table[f"{entity}_id"].isin(member_ids)
+                pd.testing.assert_frame_equal(getattr(single, entity), table[mask])
+        seen.append(household_ids)
+        return make_simulation(dataset)
+
+    result = calculate_chunked(factory, ["employment_income"], YEAR, 3, dataset=multi)
+    np.testing.assert_array_equal(
+        result.entities["person"].employment_income,
+        chunking_dataset.person.employment_income,
+    )
+    assert len(seen) == 3
+    assert set.union(*seen) == set(chunking_dataset.household.household_id)
+    for year in multi.years:
+        for entity in US_ENTITIES:
+            pd.testing.assert_frame_equal(
+                getattr(multi[year], entity), getattr(original[year], entity)
+            )
+
+
+def _snapshot_calculations(simulation, variables, period):
+    expected = {}
+    for variable in variables:
+        calculated = simulation.calculate(variable, period)
+        expected[variable] = (
+            simulation.tax_benefit_system.variables[variable].entity.key,
+            np.asarray(calculated).copy(),
+            np.asarray(calculated.weights).copy(),
+            float(calculated.sum()),
+        )
+    return expected
+
+
+def _assert_snapshot_matches(result, dataset, expected):
+    assert set(result.weighted_sums) == set(expected)
+    for variable, (entity, values, weights, weighted_sum) in expected.items():
+        table = result.entities[entity]
+        np.testing.assert_array_equal(
+            table[f"{entity}_id"], getattr(dataset, entity)[f"{entity}_id"]
+        )
+        np.testing.assert_array_equal(table[variable], values)
+        np.testing.assert_array_equal(table.weight, weights)
+        assert result.weighted_sums[variable] == weighted_sum
+
+
+def _release_reference_simulation(simulation):
+    if getattr(simulation.tax_benefit_system, "simulation", None) is simulation:
+        simulation.tax_benefit_system.simulation = None
+
+
+def test_single_year_factory_automatically_projects_future_inputs(chunking_dataset):
+    variables = ("employment_income", "household_weight")
+    future_year = YEAR + 1
+    # Exercise the normal US constructor, including its automatic dataset
+    # extension, rather than the fixed-year wrapper used by the other tests.
+    simulation = Microsimulation(dataset=chunking_dataset, dataset_end_year=future_year)
+    expected = _snapshot_calculations(simulation, variables, future_year)
+    assert np.any(
+        expected["employment_income"][1]
+        != chunking_dataset.person.employment_income.to_numpy()
+    )
+    _release_reference_simulation(simulation)
+    del simulation
+    gc.collect()
+
+    result = calculate_chunked(
+        lambda dataset: Microsimulation(dataset=dataset, dataset_end_year=future_year),
+        variables,
+        future_year,
+        3,
+        dataset=chunking_dataset,
+    )
+    _assert_snapshot_matches(result, chunking_dataset, expected)
+
+
+def test_reform_factory_is_exact_and_releases_its_baseline(chunking_dataset):
+    from policyengine_core.reforms import Reform
+
+    variables = ("income_tax", "household_net_income")
+    # Share a reform specification; each factory call gets a fresh simulation
+    # and policy system. Higher standard deductions have a visible tax effect.
+    reform = Reform.from_dict(
+        {
+            "gov.irs.deductions.standard.amount.JOINT": {
+                "2026-01-01.2100-12-31": 50_000
+            },
+            "gov.irs.deductions.standard.amount.HEAD_OF_HOUSEHOLD": {
+                "2026-01-01.2100-12-31": 35_000
+            },
+        },
+        country_id="us",
+    )
+
+    def build(dataset):
+        return Microsimulation(
+            dataset=USMultiYearDataset(datasets=[dataset]), reform=reform
+        )
+
+    baseline = make_simulation(chunking_dataset)
+    baseline_tax = np.asarray(baseline.calculate("income_tax", YEAR)).copy()
+    _release_reference_simulation(baseline)
+    del baseline
+    gc.collect()
+
+    simulation = build(chunking_dataset)
+    expected = _snapshot_calculations(simulation, variables, YEAR)
+    assert np.any(expected["income_tax"][1] != baseline_tax)
+    _release_reference_simulation(simulation)
+    del simulation
+    gc.collect()
+    references = []
+
+    def factory(dataset):
+        assert all(reference() is None for reference in references)
+        simulation = build(dataset)
+        assert simulation.baseline is not None
+        references.extend((weakref.ref(simulation), weakref.ref(simulation.baseline)))
+        return simulation
+
+    result = calculate_chunked(factory, variables, YEAR, 2, dataset=chunking_dataset)
+    _assert_snapshot_matches(result, chunking_dataset, expected)
+    assert all(reference() is None for reference in references)
+
+
+@pytest.mark.parametrize("multi_year", [False, True], ids=["single-year", "multi-year"])
+def test_entity_hdfstore_paths_match_full_calculation(
+    tmp_path, chunking_dataset, multi_year
+):
+    variables = ("employment_income", "household_weight")
+    if multi_year:
+        previous = chunking_dataset.copy()
+        previous.time_period = str(YEAR - 1)
+        previous.person["employment_income"] = previous.person.employment_income * 0.9
+        source = USMultiYearDataset(datasets=[previous, chunking_dataset])
+    else:
+        source = chunking_dataset
+    path = tmp_path / "chunking.h5"
+    source.save(str(path))
+
+    # The direct constructor supports a single-year HDFStore path. Pass the
+    # explicit multi-year object for its full-population reference; the chunk
+    # API itself must detect and load both on-disk layouts.
+    simulation = Microsimulation(
+        dataset=source if multi_year else str(path), dataset_end_year=YEAR
+    )
+    expected = _snapshot_calculations(simulation, variables, YEAR)
+    _release_reference_simulation(simulation)
+    del simulation
+    gc.collect()
+
+    result = calculate_chunked(
+        lambda dataset: Microsimulation(dataset=dataset, dataset_end_year=YEAR),
+        variables,
+        YEAR,
+        2,
+        dataset=path,
+    )
+    _assert_snapshot_matches(result, chunking_dataset, expected)
+    for variable, entity in (
+        ("employment_income", "person"),
+        ("household_weight", "household"),
+    ):
+        actual = result.entities[entity][variable].to_numpy()
+        source_values = getattr(chunking_dataset, entity)[variable].to_numpy(
+            dtype=actual.dtype
+        )
+        np.testing.assert_array_equal(actual, source_values)
diff --git a/policyengine_us/tests/microsimulation/test_tob_sequencing.py b/policyengine_us/tests/microsimulation/test_tob_sequencing.py
new file mode 100644
index 00000000000..888b5c09933
--- /dev/null
+++ b/policyengine_us/tests/microsimulation/test_tob_sequencing.py
@@ -0,0 +1,101 @@
+"""Sequential TOB counterfactuals equal the original signed main formulas."""
+
+import gc
+import weakref
+
+import numpy as np
+
+from policyengine_us.tests.microsimulation.test_calculate_chunked import (
+    YEAR,
+    build_chunking_dataset,
+    make_simulation,
+)
+
+
+def _main_counterfactual(simulation, name, taxable_social_security=None):
+    """Reproduce main's branch algorithm independently of the new helper."""
+    branch = simulation.get_branch(name, clone_system=True)
+    try:
+        if taxable_social_security is None:
+            branch.tax_benefit_system.neutralize_variable(
+                "tax_unit_taxable_social_security"
+            )
+        input_variables = set(branch.input_variables)
+        for variable in list(branch.tax_benefit_system.variables):
+            if variable not in input_variables:
+                try:
+                    branch.delete_arrays(variable)
+                except Exception:
+                    pass
+        if taxable_social_security is not None:
+            branch.populations["tax_unit"].get_holder(
+                "tax_unit_taxable_social_security"
+            ).set_input(YEAR, taxable_social_security)
+        return np.asarray(branch.tax_unit("income_tax", YEAR)).copy()
+    finally:
+        simulation.branches.pop(name)
+        del branch
+        gc.collect()
+
+
+def _release_system_backreference(simulation):
+    if getattr(simulation.tax_benefit_system, "simulation", None) is simulation:
+        del simulation.tax_benefit_system.simulation
+
+
+def test_sequential_tob_matches_main_elementwise_and_reuses_states(monkeypatch):
+    dataset = build_chunking_dataset()
+    baseline = make_simulation(dataset)
+    old_law_inclusion = np.asarray(
+        baseline.calculate("tax_unit_taxable_social_security_pre_1993", YEAR)
+    ).copy()
+    # These are main's OASDI and total algorithms, including separate no-SS
+    # branches and the signed differences rather than per-tax-unit floors.
+    oasdi_no_ss = _main_counterfactual(baseline, "oasdi_no_ss")
+    old_law = _main_counterfactual(baseline, "oasdi_old_law", old_law_inclusion)
+    current = np.asarray(baseline.calculate("income_tax", YEAR)).copy()
+    total_no_ss = _main_counterfactual(baseline, "tob_calc")
+    expected = {
+        "tob_revenue_oasdi": old_law - oasdi_no_ss,
+        "tob_revenue_total": current - total_no_ss,
+    }
+    expected["tob_revenue_medicare_hi"] = (
+        expected["tob_revenue_total"] - expected["tob_revenue_oasdi"]
+    )
+    weights = np.asarray(baseline.calculate("tax_unit_weight", YEAR)).copy()
+    _release_system_backreference(baseline)
+    del baseline
+    gc.collect()
+
+    simulation = make_simulation(dataset)
+    simulation_type = type(simulation)
+    original_get_branch = simulation_type.get_branch
+    names = []
+    references = []
+
+    def get_branch(parent, name="branch", clone_system=False):
+        is_tob = parent is simulation and name.startswith("_tob_")
+        if is_tob:
+            assert all(reference() is None for reference in references)
+        branch = original_get_branch(parent, name, clone_system=clone_system)
+        if is_tob:
+            names.append(name)
+            references.append(weakref.ref(branch))
+        return branch
+
+    monkeypatch.setattr(simulation_type, "get_branch", get_branch)
+    for variable in (
+        "tob_revenue_oasdi",
+        "tob_revenue_medicare_hi",
+        "tob_revenue_total",
+    ):
+        actual = np.asarray(simulation.calculate(variable, YEAR))
+        np.testing.assert_array_equal(actual, expected[variable])
+        assert np.sum(actual * weights) == np.sum(expected[variable] * weights)
+        assert all(reference() is None for reference in references)
+
+    assert names == ["_tob_no_social_security", "_tob_pre_1993_social_security"]
+    assert not any(name.startswith("_tob_") for name in simulation.branches)
+    _release_system_backreference(simulation)
+    del simulation
+    gc.collect()
diff --git a/policyengine_us/tools/calculate_chunked.py b/policyengine_us/tools/calculate_chunked.py
new file mode 100644
index 00000000000..cd34e132c15
--- /dev/null
+++ b/policyengine_us/tools/calculate_chunked.py
@@ -0,0 +1,329 @@
+"""Exact household partitions for calculations that are local to households."""
+
+from dataclasses import dataclass
+from functools import lru_cache
+import gc
+from numbers import Integral
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+
+from policyengine_us.data.dataset_schema import (
+    US_ENTITIES,
+    USMultiYearDataset,
+    USSingleYearDataset,
+)
+
+
+# These formulas reduce across the population, rather than within a household.
+# Guard at calculation time so parameter-selected and reform dependencies are
+# covered too, including calculations made in branches.
+_POPULATION_AGGREGATES = frozenset(
+    ("medicaid_slcsp_state_average_cost_index", "medicaid_slcsp_state_denominator")
+)
+
+
+class _HouseholdLocalCalculations:
+    @staticmethod
+    def _check_household_local(variable_name):
+        if variable_name in _POPULATION_AGGREGATES:
+            raise ValueError(
+                f"Cannot calculate {variable_name} in household chunks: "
+                "it aggregates across households. Calculate it on the full dataset."
+            )
+
+    def calculate(self, variable_name, *args, **kwargs):
+        self._check_household_local(variable_name)
+        return super().calculate(variable_name, *args, **kwargs)
+
+    def _calculate(self, variable_name, period=None):
+        self._check_household_local(variable_name)
+        return super()._calculate(variable_name, period)
+
+
+@lru_cache(maxsize=None)
+def _guarded_class(simulation_class):
+    # Core clones preserve the simulation class, so new branches inherit the
+    # guard without changing global classes or other callers' simulations.
+    return type(
+        f"HouseholdLocal{simulation_class.__name__}",
+        (_HouseholdLocalCalculations, simulation_class),
+        {},
+    )
+
+
+def _guard(simulation, seen=None):
+    seen = set() if seen is None else seen
+    if id(simulation) in seen:
+        return
+    seen.add(id(simulation))
+    if hasattr(simulation, "drop_computed_arrays"):
+        # A factory may warm caches. Recompute dependencies under the guard so
+        # a cached dependent cannot conceal a population-wide aggregation.
+        simulation.drop_computed_arrays()
+        # Chunk datasets share descriptive names; they must not read or write
+        # full-population macro-cache files under those names.
+        simulation.macro_cache_read = False
+        simulation.macro_cache_write = False
+    if not isinstance(simulation, _HouseholdLocalCalculations):
+        simulation.__class__ = _guarded_class(type(simulation))
+    for branch in simulation.branches.values():
+        _guard(branch, seen)
+    baseline = getattr(simulation, "baseline", None)
+    if baseline is not None:
+        _guard(baseline, seen)
+
+
+@dataclass
+class ChunkedCalculation:
+    """Entity tables in source order, and sums using simulation weights."""
+
+    entities: dict[str, pd.DataFrame]
+    weighted_sums: dict[str, float]
+
+
+def _load_dataset(dataset):
+    if isinstance(dataset, (str, Path)):
+        from policyengine_us.system import _resolve_dataset_path
+
+        path = _resolve_dataset_path(str(dataset))
+        with pd.HDFStore(path, mode="r") as store:
+            multi_year = any(key.startswith("/person/") for key in store.keys())
+        dataset = (
+            USMultiYearDataset(file_path=path)
+            if multi_year
+            else USSingleYearDataset(file_path=path)
+        )
+    if not isinstance(dataset, (USSingleYearDataset, USMultiYearDataset)):
+        raise TypeError(
+            "dataset must be a USSingleYearDataset, USMultiYearDataset, "
+            "or entity-table HDFStore path/URL."
+        )
+    return dataset
+
+
+def _household_positions(dataset):
+    """Validate the partition and map every entity row to a household row."""
+    person = dataset.person
+    household_ids = pd.Index(dataset.household.household_id)
+    if not household_ids.is_unique or household_ids.hasnans:
+        raise ValueError("household_id must be unique and nonmissing.")
+    people_households = household_ids.get_indexer(person.person_household_id)
+    if (people_households < 0).any():
+        raise ValueError("Every person must belong to an existing household.")
+    positions = {"household": np.arange(len(household_ids))}
+    for entity in US_ENTITIES:
+        table = getattr(dataset, entity)
+        id_column = f"{entity}_id"
+        if id_column not in table:
+            raise ValueError(
+                f"Missing {id_column}; all six entity tables are required."
+            )
+        ids = pd.Index(table[id_column])
+        if not ids.is_unique or ids.hasnans:
+            raise ValueError(f"{id_column} must be unique and nonmissing.")
+        if entity == "person":
+            positions[entity] = people_households
+        elif entity != "household":
+            links = person[f"person_{entity}_id"]
+            if (ids.get_indexer(links) < 0).any():
+                raise ValueError(f"Unknown person_{entity}_id membership.")
+            membership = (
+                pd.DataFrame({"id": links.to_numpy(), "household": people_households})
+                .groupby("id", sort=False)
+                .household
+            )
+            if (membership.nunique() != 1).any():
+                raise ValueError(f"A {entity} spans multiple households.")
+            owners = membership.first().reindex(ids)
+            if owners.isna().any():
+                raise ValueError(f"Every {entity} must have at least one person.")
+            positions[entity] = owners.to_numpy(dtype=int)
+    if len(np.unique(people_households)) != len(household_ids):
+        raise ValueError("Every household must have at least one person.")
+
+    # Medicaid can refer to a claiming tax unit explicitly. Its population-wide
+    # lookup is local only when a known claimant lives in the same household.
+    if "medicaid_claiming_tax_unit_id" in person:
+        tax_ids = pd.Index(dataset.tax_unit.tax_unit_id)
+        claimant = tax_ids.get_indexer(person.medicaid_claiming_tax_unit_id)
+        known = (claimant >= 0) & (person.medicaid_claiming_tax_unit_id.to_numpy() > 0)
+        if (positions["tax_unit"][claimant[known]] != people_households[known]).any():
+            raise ValueError("medicaid_claiming_tax_unit_id crosses households.")
+    return positions
+
+
+def calculate_chunked(
+    sim_factory,
+    variables,
+    period,
+    n_chunks,
+    *,
+    dataset,
+    household_chunks=None,
+):
+    """Calculate numeric variables on disjoint, whole-household partitions.
+
+    ``sim_factory(dataset)`` must return a fresh PolicyEngine-US Microsimulation
+    and preserve its dataset's population/order. It controls reforms and other
+    simulation options. ``dataset`` accepts US entity-table dataset objects or
+    HDFStore paths/URLs. All six entity tables and membership links are required.
+    All years are retained; IDs and memberships must be constant across years.
+
+    By default households are split into balanced consecutive chunks. Optional
+    ``household_chunks`` supplies one integer in ``[0, n_chunks)`` per source
+    household, permitting arbitrary partitions. Empty chunks are skipped.
+
+    Returns ``ChunkedCalculation``: ``entities[entity]`` has source-order rows,
+    ``<entity>_id``, ``weight`` (the weights used by Microsimulation), and
+    one column per requested variable. ``weighted_sums`` reduces those tables
+    in source order, avoiding chunk-dependent floating-point summation order.
+    Input weights are never rescaled. Only one chunk simulation is live at once.
+
+    This is exact for household-local calculations. Known Medicaid state
+    aggregates and external claiming-tax-unit links are rejected. Custom reform
+    formulas must also be household-local: population reductions, ranking and
+    random draws depending on population order are not supported.
+    """
+    if isinstance(n_chunks, bool) or not isinstance(n_chunks, Integral) or n_chunks < 1:
+        raise ValueError("n_chunks must be a positive integer.")
+    variables = [variables] if isinstance(variables, str) else list(variables)
+    if not variables or len(set(variables)) != len(variables):
+        raise ValueError("variables must be a nonempty sequence of unique names.")
+    dataset = _load_dataset(dataset)
+    years = (
+        list(dataset.datasets.values())
+        if isinstance(dataset, USMultiYearDataset)
+        else [dataset]
+    )
+    source = years[0]
+    positions = _household_positions(source)
+    if not len(source.household):
+        raise ValueError("dataset must contain at least one household.")
+    for year in years[1:]:
+        _household_positions(year)
+        for entity in US_ENTITIES:
+            columns = [f"{entity}_id"]
+            if entity == "person":
+                columns += [
+                    f"person_{group}_id" for group in US_ENTITIES if group != "person"
+                ]
+            if (
+                not getattr(source, entity)[columns]
+                .reset_index(drop=True)
+                .equals(getattr(year, entity)[columns].reset_index(drop=True))
+            ):
+                raise ValueError(
+                    "Entity IDs and memberships must be identical across years."
+                )
+    if household_chunks is None:
+        assignments = np.empty(len(source.household), dtype=int)
+        for chunk, indices in enumerate(
+            np.array_split(np.arange(len(assignments)), min(n_chunks, len(assignments)))
+        ):
+            assignments[indices] = chunk
+    else:
+        assignments = np.asarray(household_chunks)
+        if (
+            assignments.shape != (len(source.household),)
+            or not np.issubdtype(assignments.dtype, np.integer)
+            or (assignments < 0).any()
+            or (assignments >= n_chunks).any()
+        ):
+            raise ValueError(
+                "household_chunks must assign each household an integer in [0, n_chunks)."
+            )
+
+    tables = {}
+    variable_entities = {}
+    for chunk in np.unique(assignments):
+        masks = {
+            entity: assignments[owners] == chunk for entity, owners in positions.items()
+        }
+        subsets = [
+            USSingleYearDataset(
+                **{
+                    entity: getattr(year, entity)
+                    .iloc[np.flatnonzero(masks[entity])]
+                    .copy()
+                    for entity in US_ENTITIES
+                },
+                time_period=int(year.time_period),
+            )
+            for year in years
+        ]
+        chunk_dataset = (
+            USMultiYearDataset(datasets=subsets)
+            if isinstance(dataset, USMultiYearDataset)
+            else subsets[0]
+        )
+        simulation = None
+        try:
+            simulation = sim_factory(chunk_dataset)
+            _guard(simulation)
+            for name in variables:
+                entity = simulation.tax_benefit_system.get_variable(
+                    name, check_existence=True
+                ).entity.key
+                if name in variable_entities and variable_entities[name] != entity:
+                    raise ValueError(
+                        "sim_factory must use the same variable entities in every chunk."
+                    )
+                variable_entities[name] = entity
+                original_ids = getattr(source, entity)[f"{entity}_id"].to_numpy()
+                if not np.array_equal(
+                    simulation.populations[entity].ids, original_ids[masks[entity]]
+                ):
+                    raise ValueError(
+                        "sim_factory must preserve dataset entity IDs and row order."
+                    )
+                values = np.asarray(
+                    simulation.calculate(name, period, use_weights=False)
+                ).copy()
+                if not np.issubdtype(values.dtype, np.number) and values.dtype != bool:
+                    raise TypeError(
+                        f"{name} is not numeric; weighted sums require numeric variables."
+                    )
+                weights = simulation.get_weights(name, period)
+                if weights is None:
+                    raise ValueError(
+                        f"No simulation weights are available for {name} in {period}."
+                    )
+                # MicroSeries converts calculation weights to float64 before
+                # multiplying; match that even for float32 dataset weights.
+                weights = np.asarray(weights, dtype=float).copy()
+                count = int(masks[entity].sum())
+                if values.shape != (count,) or weights.shape != (count,):
+                    raise ValueError(
+                        "Calculations and weights must have one value per entity."
+                    )
+                if entity not in tables:
+                    tables[entity] = pd.DataFrame(
+                        {
+                            f"{entity}_id": original_ids.copy(),
+                            "weight": np.empty(len(original_ids), dtype=weights.dtype),
+                        }
+                    )
+                table = tables[entity]
+                if name not in table:
+                    table[name] = np.empty(len(table), dtype=values.dtype)
+                table.loc[masks[entity], name] = values
+                table.loc[masks[entity], "weight"] = weights
+        finally:
+            # Populations, branches and systems can refer back to a simulation.
+            # Collection here is deliberate: the next factory must not overlap.
+            if (
+                simulation is not None
+                and getattr(simulation.tax_benefit_system, "simulation", None)
+                is simulation
+            ):
+                simulation.tax_benefit_system.simulation = None
+            del simulation
+            del chunk_dataset, subsets
+            gc.collect()
+    sums = {
+        name: float((tables[entity][name] * tables[entity]["weight"]).sum())
+        for name, entity in variable_entities.items()
+    }
+    return ChunkedCalculation(tables, sums)
diff --git a/policyengine_us/tools/tob.py b/policyengine_us/tools/tob.py
new file mode 100644
index 00000000000..e0f0aad3460
--- /dev/null
+++ b/policyengine_us/tools/tob.py
@@ -0,0 +1,99 @@
+"""Short-lived income-tax counterfactuals for Social Security revenue."""
+
+import gc
+
+from policyengine_us.tools.period_branch import drop_inherited_values
+
+
+class TobCalculationMixin:
+    """Prepare counterfactual states before entering a requested TOB formula.
+
+    Core retains popped branches in active calculation frames. Separate
+    outer calculations let those frames finish before the next state is built,
+    without modifying Core's dependency tracking. Nested TOB dependencies still
+    use the same cached states; their outer caller may retain branch identities
+    until it returns.
+    """
+
+    def _uses_standard_tob_formula(self, name, period):
+        from policyengine_us.system import system
+
+        return self.tax_benefit_system.variables[name].get_formula(
+            period
+        ) is system.variables[name].get_formula(period)
+
+    def calculate(self, variable_name, period=None, *args, **kwargs):
+        if variable_name in (
+            "tob_revenue_oasdi",
+            "tob_revenue_medicare_hi",
+            "tob_revenue_total",
+        ):
+            calculation_period = (
+                period if period is not None else self.default_calculation_period
+            )
+            holder = self.get_holder(variable_name)
+            if (
+                not holder.variable.is_neutralized
+                and self.get_array(variable_name, calculation_period) is None
+                and self._uses_standard_tob_formula(variable_name, calculation_period)
+            ):
+                needs_total = (
+                    variable_name != "tob_revenue_oasdi"
+                    and self.get_array("tob_revenue_total", calculation_period) is None
+                    and self._uses_standard_tob_formula(
+                        "tob_revenue_total", calculation_period
+                    )
+                )
+                needs_oasdi = (
+                    variable_name != "tob_revenue_total"
+                    and self.get_array("tob_revenue_oasdi", calculation_period) is None
+                    and self._uses_standard_tob_formula(
+                        "tob_revenue_oasdi", calculation_period
+                    )
+                )
+                states = []
+                if needs_total or needs_oasdi:
+                    states.append("income_tax_without_social_security")
+                if needs_oasdi:
+                    states.append("income_tax_with_pre_1993_social_security")
+                for state in states:
+                    if self.get_array(state, calculation_period) is None:
+                        super().calculate(state, calculation_period)
+                        gc.collect()
+        return super().calculate(variable_name, period, *args, **kwargs)
+
+
+def income_tax_with_social_security_override(
+    simulation, period, *, taxable_social_security=None
+):
+    """Return only counterfactual income tax, then release its simulation.
+
+    ``None`` neutralizes taxable Social Security; an array supplies an
+    alternative inclusion. The caller caches the returned array in an ordinary
+    variable holder, so different TOB variables share the counterfactual without
+    retaining its other calculated arrays or its cloned tax-benefit system.
+    """
+    name = (
+        "_tob_no_social_security"
+        if taxable_social_security is None
+        else "_tob_pre_1993_social_security"
+    )
+    branch = simulation.get_branch(name, clone_system=True)
+    try:
+        if taxable_social_security is None:
+            branch.tax_benefit_system.neutralize_variable(
+                "tax_unit_taxable_social_security"
+            )
+        drop_inherited_values(branch)
+        if taxable_social_security is not None:
+            branch.get_holder("tax_unit_taxable_social_security").set_input(
+                period, taxable_social_security
+            )
+        # The result must not keep a view into a branch's cached storage.
+        return branch.tax_unit("income_tax", period).copy()
+    finally:
+        simulation.branches.pop(name, None)
+        del branch
+        # Simulations and populations refer to each other. Collect their cycles
+        # before another counterfactual clones the parent simulation.
+        gc.collect()
diff --git a/policyengine_us/variables/gov/ssa/revenue/income_tax_with_pre_1993_social_security.py b/policyengine_us/variables/gov/ssa/revenue/income_tax_with_pre_1993_social_security.py
new file mode 100644
index 00000000000..f984c9a84ef
--- /dev/null
+++ b/policyengine_us/variables/gov/ssa/revenue/income_tax_with_pre_1993_social_security.py
@@ -0,0 +1,25 @@
+from policyengine_us.model_api import *
+from policyengine_us.tools.tob import income_tax_with_social_security_override
+
+
+class income_tax_with_pre_1993_social_security(Variable):
+    value_type = float
+    entity = TaxUnit
+    definition_period = YEAR
+    label = "Income tax with pre-1993 taxable Social Security"
+    unit = USD
+    documentation = (
+        "Cached income-tax counterfactual using the pre-1993 section 86 "
+        "includible amount under current income-tax rules. The counterfactual "
+        "branch is released after calculating income tax."
+    )
+
+    def formula(tax_unit, period, parameters):
+        taxable_social_security = tax_unit(
+            "tax_unit_taxable_social_security_pre_1993", period
+        )
+        return income_tax_with_social_security_override(
+            tax_unit.simulation,
+            period,
+            taxable_social_security=taxable_social_security,
+        )
diff --git a/policyengine_us/variables/gov/ssa/revenue/income_tax_without_social_security.py b/policyengine_us/variables/gov/ssa/revenue/income_tax_without_social_security.py
new file mode 100644
index 00000000000..f058a52ad2c
--- /dev/null
+++ b/policyengine_us/variables/gov/ssa/revenue/income_tax_without_social_security.py
@@ -0,0 +1,18 @@
+from policyengine_us.model_api import *
+from policyengine_us.tools.tob import income_tax_with_social_security_override
+
+
+class income_tax_without_social_security(Variable):
+    value_type = float
+    entity = TaxUnit
+    definition_period = YEAR
+    label = "Income tax without taxable Social Security"
+    unit = USD
+    documentation = (
+        "Cached income-tax counterfactual shared by the taxation-of-benefits "
+        "revenue variables. The counterfactual branch is released after "
+        "calculating income tax."
+    )
+
+    def formula(tax_unit, period, parameters):
+        return income_tax_with_social_security_override(tax_unit.simulation, period)
diff --git a/policyengine_us/variables/gov/ssa/revenue/tob_revenue_oasdi.py b/policyengine_us/variables/gov/ssa/revenue/tob_revenue_oasdi.py
index d4f6cd4692b..76cbd2a4cd1 100644
--- a/policyengine_us/variables/gov/ssa/revenue/tob_revenue_oasdi.py
+++ b/policyengine_us/variables/gov/ssa/revenue/tob_revenue_oasdi.py
@@ -29,47 +29,8 @@ def formula(tax_unit, period, parameters):
         2. Calculate tax with the pre-1993 section 86 includible amount
         3. Signed OASDI revenue = (2) - (1), without a per-unit floor
         """
-        sim = tax_unit.simulation
-        old_law_taxable_ss = tax_unit(
-            "tax_unit_taxable_social_security_pre_1993", period
-        )
-
-        # === Branch 1: Tax with NO taxable SS ===
-        branch_no_ss = sim.get_branch("oasdi_no_ss", clone_system=True)
-        branch_no_ss.tax_benefit_system.neutralize_variable(
-            "tax_unit_taxable_social_security"
-        )
-
-        # Clear only non-input cached calculations
-        input_vars = set(branch_no_ss.input_variables)
-        for var_name in list(branch_no_ss.tax_benefit_system.variables.keys()):
-            if var_name not in input_vars:
-                try:
-                    branch_no_ss.delete_arrays(var_name)
-                except:
-                    pass
-
-        tax_no_ss = branch_no_ss.tax_unit("income_tax", period)
-        del sim.branches["oasdi_no_ss"]
-
-        # === Branch 2: Tax with the pre-1993 includible amount ===
-        branch_old_law = sim.get_branch("oasdi_old_law", clone_system=True)
-
-        # Clear only non-input cached calculations before setting new value
-        input_vars = set(branch_old_law.input_variables)
-        for var_name in list(branch_old_law.tax_benefit_system.variables.keys()):
-            if var_name not in input_vars:
-                try:
-                    branch_old_law.delete_arrays(var_name)
-                except:
-                    pass
-
-        # Override taxable SS after clearing computed arrays, preserving inputs.
-        tax_unit_pop = branch_old_law.populations["tax_unit"]
-        holder = tax_unit_pop.get_holder("tax_unit_taxable_social_security")
-        holder.set_input(period, old_law_taxable_ss)
-
-        tax_old_law_ss = branch_old_law.tax_unit("income_tax", period)
-        del sim.branches["oasdi_old_law"]
+        # Each state is cached on the parent after its branch is released.
+        tax_no_ss = tax_unit("income_tax_without_social_security", period)
+        tax_old_law_ss = tax_unit("income_tax_with_pre_1993_social_security", period)
 
         return tax_old_law_ss - tax_no_ss
diff --git a/policyengine_us/variables/gov/ssa/revenue/tob_revenue_total.py b/policyengine_us/variables/gov/ssa/revenue/tob_revenue_total.py
index 8fa143dba94..54b2dc5b85b 100644
--- a/policyengine_us/variables/gov/ssa/revenue/tob_revenue_total.py
+++ b/policyengine_us/variables/gov/ssa/revenue/tob_revenue_total.py
@@ -22,30 +22,8 @@ def formula(tax_unit, period, parameters):
         This isolates the tax impact of including taxable SS in AGI by
         comparing tax with and without taxable SS.
         """
-        sim = tax_unit.simulation
-
         # Calculate income tax WITH taxable SS
         income_tax_with_ss = tax_unit("income_tax", period)
-
-        # Create branch and neutralize taxable SS
-        branch = sim.get_branch("tob_calc", clone_system=True)
-        branch.tax_benefit_system.neutralize_variable(
-            "tax_unit_taxable_social_security"
-        )
-
-        # Delete non-input cached variables to force recalculation
-        input_vars = set(branch.input_variables)
-        for var_name in list(branch.tax_benefit_system.variables.keys()):
-            if var_name not in input_vars:
-                try:
-                    branch.delete_arrays(var_name)
-                except:
-                    pass
-
-        # Recalculate income tax without taxable SS
-        income_tax_without_ss = branch.tax_unit("income_tax", period)
-
-        # Clean up branch
-        del sim.branches["tob_calc"]
+        income_tax_without_ss = tax_unit("income_tax_without_social_security", period)
 
         return income_tax_with_ss - income_tax_without_ss

From c2f6569a92cc51b26495ff453386042063803679 Mon Sep 17 00:00:00 2001
From: Max Ghenis <mghenis@gmail.com>
Date: Sat, 31 Oct 2026 11:12:12 -0400
Subject: [PATCH 03/15] Harden chunk lifecycle and reform sequencing coverage

---
 docs/usage/chunked-calculations.md            |  3 +-
 .../core/test_calculate_chunked_validation.py | 74 +++++++++++++++++++
 .../tests/core/test_tob_branch_lifecycle.py   | 49 +++++++++++-
 .../microsimulation/test_calculate_chunked.py | 65 +++++++++++++---
 .../microsimulation/test_tob_sequencing.py    | 11 +--
 policyengine_us/tools/calculate_chunked.py    | 39 ++++++++--
 policyengine_us/tools/tob.py                  | 20 ++++-
 7 files changed, 234 insertions(+), 27 deletions(-)

diff --git a/docs/usage/chunked-calculations.md b/docs/usage/chunked-calculations.md
index 9fafc21b9bf..e3ebcffc9b5 100644
--- a/docs/usage/chunked-calculations.md
+++ b/docs/usage/chunked-calculations.md
@@ -35,7 +35,8 @@ columns, with exactly the source entity order. `weight` is the weight used by
 `Microsimulation.get_weights`, projected from household weights; native input
 entity weight columns are preserved in the chunk datasets. Weighted sums are
 calculated from the assembled source-order arrays, avoiding changes in reduction
-order. A requested variable can itself be an entity weight variable.
+order. A requested variable can itself be an entity weight variable. The name
+`weight` is reserved for the output weight column.
 
 ## Exact partition invariants
 
diff --git a/policyengine_us/tests/core/test_calculate_chunked_validation.py b/policyengine_us/tests/core/test_calculate_chunked_validation.py
index 922b754d4b8..a84c032a29c 100644
--- a/policyengine_us/tests/core/test_calculate_chunked_validation.py
+++ b/policyengine_us/tests/core/test_calculate_chunked_validation.py
@@ -1,5 +1,6 @@
 """Reject invalid household partitions before constructing a simulation."""
 
+import weakref
 from types import SimpleNamespace
 
 import numpy as np
@@ -59,6 +60,11 @@ def test_rejects_empty_or_duplicate_variables(dataset, variables):
         calculate(dataset, variables=variables)
 
 
+def test_rejects_reserved_weight_column(dataset):
+    with pytest.raises(ValueError, match="weight is reserved"):
+        calculate(dataset, variables=["weight"])
+
+
 @pytest.mark.parametrize("assignments", [[0], [[0], [1]], [0.0, 1.0], [-1, 0], [0, 2]])
 def test_rejects_invalid_household_assignments(dataset, assignments):
     with pytest.raises(ValueError, match="household_chunks"):
@@ -227,3 +233,71 @@ def factory(chunk):
         pd.concat(snapshots).person_weight, dataset.person.person_weight
     )
     np.testing.assert_array_equal(dataset.person.person_weight, [999.0, 888.0])
+
+
+@pytest.mark.parametrize("route", ["baseline", "branch", "popped_branch"])
+def test_releases_simulations_retained_by_system_backreferences(dataset, route):
+    # A baseline can use a shared system whose backreference keeps the baseline
+    # and its parent alive. A removed branch can likewise remain reachable only
+    # through its parent's system.simulation pointer.
+    shared_system = SimpleNamespace(get_variable=system.get_variable, simulation=None)
+    references = []
+    factory_calls = []
+
+    def factory(chunk):
+        assert all(reference() is None for reference in references)
+        assert shared_system.simulation is None
+        root = InputSimulation(chunk)
+        root.parent_branch = None
+        root.tax_benefit_system = SimpleNamespace(
+            get_variable=system.get_variable, simulation=root
+        )
+        child = InputSimulation(chunk)
+        child.parent_branch = root
+        child.tax_benefit_system = shared_system
+        shared_system.simulation = child
+        if route == "baseline":
+            root.baseline = child
+        elif route == "branch":
+            root.branches["retained"] = child
+        else:
+            # Model get_branch(..., clone_system=False), followed by pop().
+            root.tax_benefit_system = shared_system
+        references.extend((weakref.ref(root), weakref.ref(child)))
+        factory_calls.append(True)
+        return root
+
+    result = calculate_chunked(factory, ["household_weight"], YEAR, 2, dataset=dataset)
+
+    assert len(factory_calls) == 2
+    assert all(reference() is None for reference in references)
+    assert shared_system.simulation is None
+    np.testing.assert_array_equal(
+        result.entities["household"].household_weight,
+        dataset.household.household_weight,
+    )
+
+
+def test_cleanup_preserves_unrelated_shared_system_simulation(dataset):
+    external_simulation = SimpleNamespace(parent_branch=None)
+    shared_system = SimpleNamespace(
+        get_variable=system.get_variable, simulation=external_simulation
+    )
+    references = []
+    factory_calls = []
+
+    def factory(chunk):
+        assert all(reference() is None for reference in references)
+        assert shared_system.simulation is external_simulation
+        simulation = InputSimulation(chunk)
+        simulation.parent_branch = None
+        simulation.tax_benefit_system = shared_system
+        references.append(weakref.ref(simulation))
+        factory_calls.append(True)
+        return simulation
+
+    calculate_chunked(factory, ["household_weight"], YEAR, 2, dataset=dataset)
+
+    assert len(factory_calls) == 2
+    assert all(reference() is None for reference in references)
+    assert shared_system.simulation is external_simulation
diff --git a/policyengine_us/tests/core/test_tob_branch_lifecycle.py b/policyengine_us/tests/core/test_tob_branch_lifecycle.py
index 0bd8010770b..a8ebef86650 100644
--- a/policyengine_us/tests/core/test_tob_branch_lifecycle.py
+++ b/policyengine_us/tests/core/test_tob_branch_lifecycle.py
@@ -1,6 +1,6 @@
 """Counterfactual branches retain only their returned income-tax array."""
 
-from types import SimpleNamespace
+from types import FunctionType, SimpleNamespace
 import weakref
 
 import numpy as np
@@ -195,3 +195,50 @@ def test_tob_reuses_states_and_forwards_original_request_options(monkeypatch):
         ("collect",),
         ("calculate", OASDI, None, ("person",), {"decode_enums": False}),
     ]
+
+
+def test_standard_formula_detection_accepts_fresh_import_but_rejects_replacement():
+    from policyengine_us.system import system
+
+    standard = system.variables[OASDI].get_formula(2026)
+    # Fresh country systems import the same source in another module, so both
+    # the function and its code object can differ in identity.
+    reimported = FunctionType(
+        standard.__code__.replace(),
+        dict(standard.__globals__),
+        argdefs=standard.__defaults__,
+        closure=standard.__closure__,
+    )
+    reimported.__kwdefaults__ = standard.__kwdefaults__
+    assert reimported is not standard
+    assert reimported.__code__ is not standard.__code__
+    variable = SimpleNamespace(get_formula=lambda period: reimported)
+    simulation = tob.TobCalculationMixin()
+    simulation.tax_benefit_system = SimpleNamespace(variables={OASDI: variable})
+
+    assert simulation._uses_standard_tob_formula(OASDI, 2026)
+
+    variable.get_formula = lambda period: lambda tax_unit, period, parameters: 0
+    assert not simulation._uses_standard_tob_formula(OASDI, 2026)
+
+
+def test_equal_formula_code_with_different_globals_is_not_standard(monkeypatch):
+    from policyengine_us.system import system
+
+    def standard(tax_unit, period, parameters):
+        return np.ones(1)
+
+    replacement = FunctionType(
+        standard.__code__, {"np": SimpleNamespace(ones=lambda size: np.zeros(size))}
+    )
+    monkeypatch.setitem(
+        system.variables, OASDI, SimpleNamespace(get_formula=lambda period: standard)
+    )
+    simulation = tob.TobCalculationMixin()
+    simulation.tax_benefit_system = SimpleNamespace(
+        variables={OASDI: SimpleNamespace(get_formula=lambda period: replacement)}
+    )
+
+    np.testing.assert_array_equal(standard(None, None, None), [1])
+    np.testing.assert_array_equal(replacement(None, None, None), [0])
+    assert not simulation._uses_standard_tob_formula(OASDI, 2026)
diff --git a/policyengine_us/tests/microsimulation/test_calculate_chunked.py b/policyengine_us/tests/microsimulation/test_calculate_chunked.py
index 38d589aef44..3be09eecd3f 100644
--- a/policyengine_us/tests/microsimulation/test_calculate_chunked.py
+++ b/policyengine_us/tests/microsimulation/test_calculate_chunked.py
@@ -3,6 +3,12 @@
 Python is needed to exercise the dataset/factory API, arbitrary partitions and
 simulation lifetime. The 32-household reference is calculated once; only its
 arrays survive, so no reference simulation remains live during chunked runs.
+
+Partition cases use the country's parameter-lending clone route: branches keep
+the complete, unchanged parameter tree and history, with copy-on-write for
+reforms, while entities, variable metadata and provider receipts remain private.
+Rebuilding the same history for every equivalent partition would dominate this
+API test's runtime. The separate TOB sequencing test uses ordinary clones.
 """
 
 import gc
@@ -38,6 +44,27 @@
 )
 
 
+@pytest.fixture(scope="module", autouse=True)
+def lend_unchanged_parameter_history():
+    import policyengine_us.spm as spm
+
+    original = spm.clone_spm_system
+
+    def clone(source, *, copy_receipts=True, parameters=None):
+        tree = source.parameters if parameters is None else parameters
+        # Lending creates another reader even when the source just acquired a
+        # private reform tree. Keep the country's write barrier armed for both
+        # readers; a parameter reform still pays for its own detached clone.
+        spm.mark_tree_shared(tree)
+        return original(source, copy_receipts=copy_receipts, parameters=tree)
+
+    # Scope the optimization to this module. In particular, importing fixture
+    # builders from test_tob_sequencing does not apply it to that test module.
+    with pytest.MonkeyPatch.context() as patch:
+        patch.setattr(spm, "clone_spm_system", clone)
+        yield
+
+
 def build_chunking_dataset() -> USSingleYearDataset:
     """Several states, multi-unit households and taxable Social Security.
 
@@ -46,7 +73,7 @@ def build_chunking_dataset() -> USSingleYearDataset:
     rather than coincident entity positions, ID sorting, or pandas indices.
     """
     records = {entity: [] for entity in US_ENTITIES}
-    state_fips = (6, 36, 48, 20, 12)
+    state_fips = (6, 48, 12, 53, 32)
     for household in range(HOUSEHOLDS):
         household_id = 90_001 + 37 * household
         # Non-binary fractions expose accidental float32 weighted reductions.
@@ -139,9 +166,7 @@ def reference(chunking_dataset):
             weights[entity] = np.asarray(calculated.weights).copy()
         weighted_sums[variable] = float(calculated.sum())
     del calculated
-    # Core stores this backreference even for the global default system.
-    if getattr(simulation.tax_benefit_system, "simulation", None) is simulation:
-        del simulation.tax_benefit_system.simulation
+    _release_reference_simulation(simulation)
     del simulation
     gc.collect()
     for variable in (
@@ -319,8 +344,24 @@ def _assert_snapshot_matches(result, dataset, expected):
 
 
 def _release_reference_simulation(simulation):
-    if getattr(simulation.tax_benefit_system, "simulation", None) is simulation:
-        simulation.tax_benefit_system.simulation = None
+    """Detach reference simulations without using production cleanup code."""
+    owned = []
+    seen = set()
+    pending = [simulation]
+    while pending:
+        current = pending.pop()
+        if id(current) in seen:
+            continue
+        seen.add(id(current))
+        owned.append(current)
+        pending.extend(current.branches.values())
+        baseline = getattr(current, "baseline", None)
+        if baseline is not None:
+            pending.append(baseline)
+    for current in owned:
+        policy = current.tax_benefit_system
+        if id(getattr(policy, "simulation", None)) in seen:
+            policy.simulation = None
 
 
 def test_single_year_factory_automatically_projects_future_inputs(chunking_dataset):
@@ -350,6 +391,7 @@ def test_single_year_factory_automatically_projects_future_inputs(chunking_datas
 
 def test_reform_factory_is_exact_and_releases_its_baseline(chunking_dataset):
     from policyengine_core.reforms import Reform
+    from policyengine_us.spm import clone_spm_system
 
     variables = ("income_tax", "household_net_income")
     # Share a reform specification; each factory call gets a fresh simulation
@@ -367,8 +409,12 @@ def test_reform_factory_is_exact_and_releases_its_baseline(chunking_dataset):
     )
 
     def build(dataset):
+        # The country clone preserves instance-level variable metadata and
+        # provider state while avoiding repeated country/parameter loading.
         return Microsimulation(
-            dataset=USMultiYearDataset(datasets=[dataset]), reform=reform
+            dataset=USMultiYearDataset(datasets=[dataset]),
+            tax_benefit_system=clone_spm_system(system),
+            reform=reform,
         )
 
     baseline = make_simulation(chunking_dataset)
@@ -415,8 +461,9 @@ def test_entity_hdfstore_paths_match_full_calculation(
     # The direct constructor supports a single-year HDFStore path. Pass the
     # explicit multi-year object for its full-population reference; the chunk
     # API itself must detect and load both on-disk layouts.
+    constructor_options = {} if multi_year else {"dataset_end_year": YEAR}
     simulation = Microsimulation(
-        dataset=source if multi_year else str(path), dataset_end_year=YEAR
+        dataset=source if multi_year else str(path), **constructor_options
     )
     expected = _snapshot_calculations(simulation, variables, YEAR)
     _release_reference_simulation(simulation)
@@ -424,7 +471,7 @@ def test_entity_hdfstore_paths_match_full_calculation(
     gc.collect()
 
     result = calculate_chunked(
-        lambda dataset: Microsimulation(dataset=dataset, dataset_end_year=YEAR),
+        lambda dataset: Microsimulation(dataset=dataset, **constructor_options),
         variables,
         YEAR,
         2,
diff --git a/policyengine_us/tests/microsimulation/test_tob_sequencing.py b/policyengine_us/tests/microsimulation/test_tob_sequencing.py
index 888b5c09933..18975a10e35 100644
--- a/policyengine_us/tests/microsimulation/test_tob_sequencing.py
+++ b/policyengine_us/tests/microsimulation/test_tob_sequencing.py
@@ -7,6 +7,7 @@
 
 from policyengine_us.tests.microsimulation.test_calculate_chunked import (
     YEAR,
+    _release_reference_simulation,
     build_chunking_dataset,
     make_simulation,
 )
@@ -34,15 +35,11 @@ def _main_counterfactual(simulation, name, taxable_social_security=None):
         return np.asarray(branch.tax_unit("income_tax", YEAR)).copy()
     finally:
         simulation.branches.pop(name)
+        _release_reference_simulation(branch)
         del branch
         gc.collect()
 
 
-def _release_system_backreference(simulation):
-    if getattr(simulation.tax_benefit_system, "simulation", None) is simulation:
-        del simulation.tax_benefit_system.simulation
-
-
 def test_sequential_tob_matches_main_elementwise_and_reuses_states(monkeypatch):
     dataset = build_chunking_dataset()
     baseline = make_simulation(dataset)
@@ -63,7 +60,7 @@ def test_sequential_tob_matches_main_elementwise_and_reuses_states(monkeypatch):
         expected["tob_revenue_total"] - expected["tob_revenue_oasdi"]
     )
     weights = np.asarray(baseline.calculate("tax_unit_weight", YEAR)).copy()
-    _release_system_backreference(baseline)
+    _release_reference_simulation(baseline)
     del baseline
     gc.collect()
 
@@ -96,6 +93,6 @@ def get_branch(parent, name="branch", clone_system=False):
 
     assert names == ["_tob_no_social_security", "_tob_pre_1993_social_security"]
     assert not any(name.startswith("_tob_") for name in simulation.branches)
-    _release_system_backreference(simulation)
+    _release_reference_simulation(simulation)
     del simulation
     gc.collect()
diff --git a/policyengine_us/tools/calculate_chunked.py b/policyengine_us/tools/calculate_chunked.py
index cd34e132c15..0156c5d373e 100644
--- a/policyengine_us/tools/calculate_chunked.py
+++ b/policyengine_us/tools/calculate_chunked.py
@@ -75,6 +75,35 @@ def _guard(simulation, seen=None):
         _guard(baseline, seen)
 
 
+def _detach_simulation_links(root):
+    """Release systems' references to this chunk and its owned branches.
+
+    A reform baseline can use the globally shared baseline system. Its
+    ``simulation`` reference otherwise retains the baseline and its parent
+    after the factory's result has been deleted.
+    """
+    pending = [root]
+    seen = set()
+    while pending:
+        simulation = pending.pop()
+        if id(simulation) in seen:
+            continue
+        seen.add(id(simulation))
+        system = simulation.tax_benefit_system
+        owner = getattr(system, "simulation", None)
+        ancestors = set()
+        while owner is not None and id(owner) not in ancestors:
+            if owner is root:
+                system.simulation = None
+                break
+            ancestors.add(id(owner))
+            owner = getattr(owner, "parent_branch", None)
+        pending.extend(getattr(simulation, "branches", {}).values())
+        baseline = getattr(simulation, "baseline", None)
+        if baseline is not None:
+            pending.append(baseline)
+
+
 @dataclass
 class ChunkedCalculation:
     """Entity tables in source order, and sums using simulation weights."""
@@ -191,6 +220,8 @@ def calculate_chunked(
     variables = [variables] if isinstance(variables, str) else list(variables)
     if not variables or len(set(variables)) != len(variables):
         raise ValueError("variables must be a nonempty sequence of unique names.")
+    if "weight" in variables:
+        raise ValueError("weight is reserved for the returned simulation weights.")
     dataset = _load_dataset(dataset)
     years = (
         list(dataset.datasets.values())
@@ -313,12 +344,8 @@ def calculate_chunked(
         finally:
             # Populations, branches and systems can refer back to a simulation.
             # Collection here is deliberate: the next factory must not overlap.
-            if (
-                simulation is not None
-                and getattr(simulation.tax_benefit_system, "simulation", None)
-                is simulation
-            ):
-                simulation.tax_benefit_system.simulation = None
+            if simulation is not None:
+                _detach_simulation_links(simulation)
             del simulation
             del chunk_dataset, subsets
             gc.collect()
diff --git a/policyengine_us/tools/tob.py b/policyengine_us/tools/tob.py
index e0f0aad3460..1f4630c0146 100644
--- a/policyengine_us/tools/tob.py
+++ b/policyengine_us/tools/tob.py
@@ -18,9 +18,23 @@ class TobCalculationMixin:
     def _uses_standard_tob_formula(self, name, period):
         from policyengine_us.system import system
 
-        return self.tax_benefit_system.variables[name].get_formula(
-            period
-        ) is system.variables[name].get_formula(period)
+        formula = self.tax_benefit_system.variables[name].get_formula(period)
+        standard = system.variables[name].get_formula(period)
+        if formula is standard:
+            return True
+        # A fresh country system reloads unchanged variable source in separate
+        # modules. These three formulas use only arguments and literals, so
+        # equal code also identifies them without accepting reform functions
+        # whose behavior depends on different module globals or closures.
+        code = getattr(formula, "__code__", None)
+        return (
+            code is not None
+            and not code.co_names
+            and not code.co_freevars
+            and code == getattr(standard, "__code__", None)
+            and formula.__defaults__ == standard.__defaults__
+            and formula.__kwdefaults__ == standard.__kwdefaults__
+        )
 
     def calculate(self, variable_name, period=None, *args, **kwargs):
         if variable_name in (

From cb32eb0824c0015a2a80fce1aefab9fd3408ab40 Mon Sep 17 00:00:00 2001
From: Max Ghenis <mghenis@gmail.com>
Date: Sat, 31 Oct 2026 11:51:44 -0400
Subject: [PATCH 04/15] Reduce repeated setup in chunking exactness tests

---
 .../microsimulation/test_calculate_chunked.py | 38 ++++++++++++++++---
 1 file changed, 33 insertions(+), 5 deletions(-)

diff --git a/policyengine_us/tests/microsimulation/test_calculate_chunked.py b/policyengine_us/tests/microsimulation/test_calculate_chunked.py
index 3be09eecd3f..591406e83a3 100644
--- a/policyengine_us/tests/microsimulation/test_calculate_chunked.py
+++ b/policyengine_us/tests/microsimulation/test_calculate_chunked.py
@@ -9,6 +9,9 @@
 reforms, while entities, variable metadata and provider receipts remain private.
 Rebuilding the same history for every equivalent partition would dominate this
 API test's runtime. The separate TOB sequencing test uses ordinary clones.
+The retained default policy graph is frozen before any simulation is built,
+avoiding repeated GC scans of it. New simulation/branch cycles remain in normal
+GC generations and must still pass the weak-reference collection checks.
 """
 
 import gc
@@ -48,6 +51,10 @@
 def lend_unchanged_parameter_history():
     import policyengine_us.spm as spm
 
+    owns_freeze = gc.get_freeze_count() == 0
+    if owns_freeze:
+        gc.collect()
+        gc.freeze()
     original = spm.clone_spm_system
 
     def clone(source, *, copy_receipts=True, parameters=None):
@@ -60,9 +67,13 @@ def clone(source, *, copy_receipts=True, parameters=None):
 
     # Scope the optimization to this module. In particular, importing fixture
     # builders from test_tob_sequencing does not apply it to that test module.
-    with pytest.MonkeyPatch.context() as patch:
-        patch.setattr(spm, "clone_spm_system", clone)
-        yield
+    try:
+        with pytest.MonkeyPatch.context() as patch:
+            patch.setattr(spm, "clone_spm_system", clone)
+            yield
+    finally:
+        if owns_freeze:
+            gc.unfreeze()
 
 
 def build_chunking_dataset() -> USSingleYearDataset:
@@ -91,12 +102,29 @@ def build_chunking_dataset() -> USSingleYearDataset:
         family_ids = [401_011 + 47 * household + offset for offset in (0, 17)]
         marital_ids = [501_019 + 53 * household + offset for offset in (0, 19, 23)]
         for tax_id in tax_ids:
+            # As in the signed TOB fixtures, these units report no itemized
+            # medical or SALT deductions. Keep optional sales-tax tables and
+            # premium attribution out of repeated income-tax counterfactuals;
+            # the standard-deduction choice and all requested formulas remain
+            # endogenous.
             records["tax_unit"].append(
-                {"tax_unit_id": tax_id, "tax_unit_weight": weight}
+                {
+                    "tax_unit_id": tax_id,
+                    "tax_unit_weight": weight,
+                    "itemized_medical_expenses": 0.0,
+                    "salt_deduction": 0.0,
+                }
             )
         for spm_id in spm_ids[: 2 if household % 2 == 0 else 1]:
+            # These positive-income units do not take up TANF. Supplying the
+            # take-up input avoids unrelated monthly eligibility calculations
+            # in each tax branch and in the net-income resource chain.
             records["spm_unit"].append(
-                {"spm_unit_id": spm_id, "spm_unit_weight": weight}
+                {
+                    "spm_unit_id": spm_id,
+                    "spm_unit_weight": weight,
+                    "takes_up_tanf_if_eligible": False,
+                }
             )
         for family_id in family_ids[: 2 if household % 2 == 0 else 1]:
             records["family"].append({"family_id": family_id, "family_weight": weight})

From babdac8d3f8dbbc58059461e1eb2864d543d666b Mon Sep 17 00:00:00 2001
From: Max Ghenis <mghenis@gmail.com>
Date: Sat, 31 Oct 2026 12:43:14 -0400
Subject: [PATCH 05/15] Use observed benefit inputs in the chunking fixture

---
 .../tests/microsimulation/test_calculate_chunked.py | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/policyengine_us/tests/microsimulation/test_calculate_chunked.py b/policyengine_us/tests/microsimulation/test_calculate_chunked.py
index 591406e83a3..c74bd8b27b3 100644
--- a/policyengine_us/tests/microsimulation/test_calculate_chunked.py
+++ b/policyengine_us/tests/microsimulation/test_calculate_chunked.py
@@ -115,7 +115,7 @@ def build_chunking_dataset() -> USSingleYearDataset:
                     "salt_deduction": 0.0,
                 }
             )
-        for spm_id in spm_ids[: 2 if household % 2 == 0 else 1]:
+        for unit_offset, spm_id in enumerate(spm_ids[: 2 if household % 2 == 0 else 1]):
             # These positive-income units do not take up TANF. Supplying the
             # take-up input avoids unrelated monthly eligibility calculations
             # in each tax branch and in the net-income resource chain.
@@ -124,6 +124,13 @@ def build_chunking_dataset() -> USSingleYearDataset:
                     "spm_unit_id": spm_id,
                     "spm_unit_weight": weight,
                     "takes_up_tanf_if_eligible": False,
+                    # Observed annual benefits and expenses keep this dataset
+                    # test independent of unrelated eligibility networks.
+                    "snap": 1_200.0 if household == 0 and unit_offset == 1 else 0.0,
+                    "housing_assistance": 0.0,
+                    "spm_unit_medical_out_of_pocket_expenses": (
+                        600.0 + 20 * household + unit_offset
+                    ),
                 }
             )
         for family_id in family_ids[: 2 if household % 2 == 0 else 1]:
@@ -154,6 +161,10 @@ def build_chunking_dataset() -> USSingleYearDataset:
                 "social_security_retirement": (
                     25_000 + household * 150 if member < 2 else 0.0
                 ),
+                # Nobody reports SSI or WIC receipt. Core distributes these
+                # annual flow inputs across their twelve monthly periods.
+                "ssi": 0.0,
+                "wic": 0.0,
             }
             for entity, entity_id in links.items():
                 person[f"person_{entity}_id"] = entity_id

From 5da5df7893f9e3c69f5a71ee2d5c41a72da5d349 Mon Sep 17 00:00:00 2001
From: Max Ghenis <mghenis@gmail.com>
Date: Sat, 31 Oct 2026 12:59:31 -0400
Subject: [PATCH 06/15] Bound childcare eligibility work in exactness tests

---
 .../tests/microsimulation/test_calculate_chunked.py           | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/policyengine_us/tests/microsimulation/test_calculate_chunked.py b/policyengine_us/tests/microsimulation/test_calculate_chunked.py
index c74bd8b27b3..f6f483f6faf 100644
--- a/policyengine_us/tests/microsimulation/test_calculate_chunked.py
+++ b/policyengine_us/tests/microsimulation/test_calculate_chunked.py
@@ -131,6 +131,10 @@ def build_chunking_dataset() -> USSingleYearDataset:
                     "spm_unit_medical_out_of_pocket_expenses": (
                         600.0 + 20 * household + unit_offset
                     ),
+                    # Nobody reports childcare expenses or subsidy receipt.
+                    "nv_child_care_subsidies": 0.0,
+                    "ca_child_care_subsidies": 0.0,
+                    "wa_child_care_subsidies": 0.0,
                 }
             )
         for family_id in family_ids[: 2 if household % 2 == 0 else 1]:

From 9d3dfa57ae32ccf03b05d9a56e6002671de0be76 Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Mon, 2 Nov 2026 15:10:57 +0000
Subject: [PATCH 07/15] Chunk variable-centric population files and guard on
 Core's period hook

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>
---
 docs/usage/chunked-calculations.md            |  60 ++-
 .../core/test_calculate_chunked_validation.py | 387 +++++++++++++++++-
 .../microsimulation/test_calculate_chunked.py | 178 +++++++-
 policyengine_us/tools/calculate_chunked.py    | 330 +++++++++++----
 4 files changed, 864 insertions(+), 91 deletions(-)

diff --git a/docs/usage/chunked-calculations.md b/docs/usage/chunked-calculations.md
index 7bcb683ad42..724682a2ed9 100644
--- a/docs/usage/chunked-calculations.md
+++ b/docs/usage/chunked-calculations.md
@@ -25,10 +25,25 @@ The factory owns reforms and simulation options. It must return a fresh
 PolicyEngine-US Microsimulation with the supplied population and row order.
 The explicit `dataset=` keyword supplies the source without first constructing
 a full-population simulation. Supported sources are `USSingleYearDataset`,
-`USMultiYearDataset`, and local or Hugging Face entity-table HDFStore files.
-Legacy variable-array and flat-file datasets must first be converted to the US
-entity-table schema. All six entity tables and person membership links are
-required. Input data is not modified.
+`USMultiYearDataset`, local or Hugging Face entity-table HDFStore files, and
+local or Hugging Face variable-centric `.h5` files: top-level groups named for
+variables, each holding one array per period (`<variable>/<period>`), as
+population projection files store them and as `Microsimulation(dataset=path)`
+loads them through Core. Flat files are not supported. All six entity ID
+arrays and person membership links are required. Input data is not modified.
+
+For entity tables the factory receives a dataset object of the source's type.
+For a variable-centric file it receives a `str` path to a temporary `.h5` file
+in the same layout, with the source file's name, every stored period and each
+array's dtype, fixed-width byte strings included, so the factory can call
+`Microsimulation(dataset=path)`. The source arrays stay on disk and are read
+again for each chunk. Each array is assigned to the entity of its variable in
+the country system; arrays whose length differs from that entity's count are
+rejected, and names that match no variable are left out, as Core's loader
+ignores them. Membership role arrays (`person_<group>_role`), which Core also
+reads, are kept with the people. Entity IDs and memberships must be identical
+in every stored period. Only one chunk file exists at once; it is deleted after
+its chunk, also when the factory or a calculation fails.
 
 Each output DataFrame contains `<entity>_id`, `weight`, and requested variable
 columns, with exactly the source entity order. `weight` is the weight used by
@@ -53,22 +68,37 @@ Optional `household_chunks` provides arbitrary integer labels in
 partitions are skipped, including when chunks outnumber households.
 
 Exact calculation additionally requires every dependency to be local to a
-household. The current model has two Medicaid state reductions,
-`medicaid_slcsp_state_average_cost_index` and `medicaid_slcsp_state_denominator`.
-The helper rejects these at calculation time, including indirect dependencies
-and branches. Known Medicaid claiming-tax-unit links crossing households are
-also rejected. The model audit found no other population ranking, quantile or
-stochastic formula reductions. Custom reform formulas remain the caller's
-responsibility: cross-household reductions, order-dependent random draws and
-other global dependencies are unsupported. These are calculation restrictions,
-not an assumption that every model variable can be chunked.
+household. The current model has four formulas that reduce across every unit
+in the 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`, which
+call `MicroSeries.decile_rank` and would rank within the chunk. The helper
+rejects these at calculation time, including indirect dependencies such as
+`income_decile`, branches and reform baselines. The check runs in Core's
+period-consistency hook, just before a formula would run, so it adds no frames
+to a dependency chain; a chunked chain reaches the recursion limit when the
+whole-file chain does. A guarded variable with a stored value returns that
+value, because no formula reduces across the chunk. Known Medicaid
+claiming-tax-unit links crossing households are also rejected. The SSI spouse
+pairing also tests whether the whole simulation holds a single marital unit, so
+a chunk with exactly one marital unit can differ from the whole file; this is
+not guarded. Custom reform formulas remain the caller's responsibility:
+cross-household reductions, order-dependent random draws and other global
+dependencies are unsupported. These are calculation restrictions, not an
+assumption that every model variable can be chunked.
 
 The tests compare a full 32-household, five-state fixture with one, two, three
 and seven chunks, plus random household assignments. They compare every entity
 value and weighted sum, check source order and unchanged input frames, and use
 weak references to check simulation release. The fixture includes households
-with multiple tax and SPM units and Social Security recipients. Independent
-pre-change TOB branch calculations protect the statutory signed allocation.
+with multiple tax and SPM units and Social Security recipients. The same
+households written as a variable-centric file are compared with an unchunked
+`Microsimulation` on that file for one, two, three and seven chunks and under a
+parameter reform, and with the entity-table route. Real simulations check the
+guard on direct, dependency, branch, baseline and cross-period routes, and a
+frame-depth test checks that guarding adds at most three frames to the deepest
+stack. Independent pre-change TOB branch calculations protect the statutory
+signed allocation.
 
 ## TOB working memory
 
diff --git a/policyengine_us/tests/core/test_calculate_chunked_validation.py b/policyengine_us/tests/core/test_calculate_chunked_validation.py
index a84c032a29c..79024a939d6 100644
--- a/policyengine_us/tests/core/test_calculate_chunked_validation.py
+++ b/policyengine_us/tests/core/test_calculate_chunked_validation.py
@@ -1,8 +1,11 @@
 """Reject invalid household partitions before constructing a simulation."""
 
+import re
 import weakref
+from pathlib import Path
 from types import SimpleNamespace
 
+import h5py
 import numpy as np
 import pandas as pd
 import pytest
@@ -13,7 +16,11 @@
     USSingleYearDataset,
 )
 from policyengine_us.system import system
-from policyengine_us.tools.calculate_chunked import calculate_chunked
+from policyengine_us.tools.calculate_chunked import (
+    _POPULATION_AGGREGATES,
+    _guard,
+    calculate_chunked,
+)
 
 YEAR = 2026
 
@@ -120,8 +127,25 @@ def test_rejects_changed_multi_year_memberships(dataset):
         calculate(multi)
 
 
+GUARDED = sorted(_POPULATION_AGGREGATES)
+
+
+def test_guard_covers_known_population_formulas():
+    assert set(GUARDED) == {
+        "medicaid_slcsp_state_average_cost_index",
+        "medicaid_slcsp_state_denominator",
+        "household_income_decile",
+        "spm_unit_income_decile",
+    }
+
+
 class AggregateSimulation:
-    """Exercise aggregate guards without building a country simulation."""
+    """Exercise aggregate guards without building a country simulation.
+
+    Like Core's ``Simulation._calculate``, ``_calculate`` checks the period
+    just before a formula would run. The real-simulation tests below cover the
+    Core hook itself.
+    """
 
     def __init__(self, dataset, aggregate, route):
         self.tax_benefit_system = system
@@ -142,10 +166,14 @@ def __init__(self, dataset, aggregate, route):
             else:
                 self.baseline = target
 
+    def _check_period_consistency(self, period, variable):
+        pass
+
     def calculate(self, variable, period, use_weights=False):
         return self._calculate(variable, period)
 
     def _calculate(self, variable, period=None):
+        self._check_period_consistency(period, system.variables[variable])
         if self.route == "branch":
             return self.branches["existing"].calculate(self.aggregate, period)
         if self.route == "baseline":
@@ -155,10 +183,7 @@ def _calculate(self, variable, period=None):
         pytest.fail("Population aggregate must be rejected before formula execution")
 
 
-@pytest.mark.parametrize(
-    "aggregate",
-    ["medicaid_slcsp_state_average_cost_index", "medicaid_slcsp_state_denominator"],
-)
+@pytest.mark.parametrize("aggregate", GUARDED)
 @pytest.mark.parametrize("route", ["direct", "dependency", "branch", "baseline"])
 def test_rejects_population_aggregates_in_dependency_graph(dataset, aggregate, route):
     variable = aggregate if route == "direct" else "income_tax"
@@ -172,6 +197,137 @@ def test_rejects_population_aggregates_in_dependency_graph(dataset, aggregate, r
         )
 
 
+def _reader_reform():
+    """Add a yearly reader of each guarded variable and one monthly reader."""
+    from policyengine_core.reforms import Reform
+    from policyengine_core.variables import Variable
+
+    readers = {}
+    for aggregate in GUARDED:
+        guarded = system.variables[aggregate]
+
+        def formula(unit, period, parameters, aggregate=aggregate):
+            return unit(aggregate, period)
+
+        readers[aggregate] = type(
+            f"reads_{aggregate}",
+            (Variable,),
+            {
+                "value_type": float,
+                "entity": guarded.entity,
+                "definition_period": "year",
+                "label": f"Reads {aggregate}",
+                "formula": formula,
+            },
+        )
+
+    def monthly_formula(spm_unit, period, parameters):
+        return spm_unit("spm_unit_income_decile", period.this_year)
+
+    monthly = type(
+        "reads_spm_unit_income_decile_monthly",
+        (Variable,),
+        {
+            "value_type": float,
+            "entity": system.variables["spm_unit_income_decile"].entity,
+            "definition_period": "month",
+            "label": "Reads spm_unit_income_decile each month",
+            "formula": monthly_formula,
+        },
+    )
+
+    class AddReaders(Reform):
+        def apply(self):
+            for reader in [*readers.values(), monthly]:
+                self.add_variable(reader)
+
+    return AddReaders
+
+
+@pytest.fixture(scope="module")
+def guarded_simulations():
+    """One reformed simulation over a tiny dataset, guarded as chunks are.
+
+    Nothing here runs a guarded formula, so one simulation serves every route.
+    """
+    from policyengine_us import Microsimulation
+
+    ids = np.asarray([91, 37])
+    person = pd.DataFrame({"person_id": [503, 101], "age": [40, 40]})
+    for entity in US_ENTITIES:
+        if entity != "person":
+            person[f"person_{entity}_id"] = ids
+    tables = {"person": person}
+    for entity in US_ENTITIES:
+        if entity != "person":
+            tables[entity] = pd.DataFrame({f"{entity}_id": ids})
+    tables["household"]["household_weight"] = [1.5, 2.25]
+    tables["household"]["state_fips"] = [6, 48]
+    single = USSingleYearDataset(**tables, time_period=YEAR)
+    simulation = Microsimulation(
+        dataset=USMultiYearDataset(datasets=[single]), reform=_reader_reform()
+    )
+    stored = Microsimulation(dataset=USMultiYearDataset(datasets=[single]))
+    stored.set_input("spm_unit_income_decile", YEAR, np.asarray([3, 8]))
+    stored.set_input("medicaid_slcsp_state_denominator", YEAR, np.asarray([2.5, 4.0]))
+    _guard(simulation)
+    _guard(stored)
+    yield simulation, stored
+    for current in (simulation, simulation.baseline, stored):
+        current.tax_benefit_system.simulation = None
+
+
+def _rejects(calculation):
+    with pytest.raises(ValueError, match="aggregates across households"):
+        calculation()
+
+
+@pytest.mark.parametrize("aggregate", GUARDED)
+def test_core_hook_guards_every_route(guarded_simulations, aggregate):
+    # These are real Core simulations: a renamed Core hook fails here rather
+    # than silently dropping the guard.
+    simulation, _ = guarded_simulations
+    assert simulation.baseline is not None
+    _rejects(lambda: simulation.calculate(aggregate, YEAR))
+    _rejects(lambda: simulation.calculate(f"reads_{aggregate}", YEAR))
+    # A yearly variable requested for a month reaches the hook for its year.
+    _rejects(lambda: simulation.calculate(aggregate, f"{YEAR}-01"))
+    branch = simulation.get_branch(f"branch_{aggregate}")
+    try:
+        _rejects(lambda: branch.calculate(aggregate, YEAR))
+        _rejects(lambda: branch.calculate(f"reads_{aggregate}", YEAR))
+    finally:
+        simulation.branches.pop(f"branch_{aggregate}")
+    _rejects(lambda: simulation.baseline.calculate(aggregate, YEAR))
+
+
+def test_core_hook_guards_readers_across_periods(guarded_simulations):
+    simulation, _ = guarded_simulations
+    # income_decile reaches household_income_decile through its formula.
+    _rejects(lambda: simulation.calculate("income_decile", YEAR))
+    _rejects(lambda: simulation.baseline.calculate("income_decile", YEAR))
+    # A monthly reader requested for a year is summed over its months.
+    _rejects(lambda: simulation.calculate("reads_spm_unit_income_decile_monthly", YEAR))
+    _rejects(
+        lambda: simulation.calculate(
+            "reads_spm_unit_income_decile_monthly", f"{YEAR}-03"
+        )
+    )
+
+
+def test_guarded_variable_with_stored_value_returns_it(guarded_simulations):
+    # No formula reduces across the chunk when the value is stored.
+    _, stored = guarded_simulations
+    np.testing.assert_array_equal(
+        stored.calculate("spm_unit_income_decile", YEAR, use_weights=False), [3, 8]
+    )
+    np.testing.assert_array_equal(
+        stored.calculate("medicaid_slcsp_state_denominator", YEAR, use_weights=False),
+        [2.5, 4.0],
+    )
+    _rejects(lambda: stored.calculate("household_income_decile", YEAR))
+
+
 class InputSimulation:
     def __init__(self, dataset):
         self.dataset = dataset
@@ -301,3 +457,222 @@ def factory(chunk):
     assert len(factory_calls) == 2
     assert all(reference() is None for reference in references)
     assert shared_system.simulation is external_simulation
+
+
+# Variable-centric files: <variable>/<period> arrays, as Core loads them.
+
+PERIODS = (str(YEAR - 1), str(YEAR))
+
+
+def write_variable_centric(path, dataset, periods=PERIODS, extra=()):
+    with h5py.File(path, "w") as file:
+        for entity in US_ENTITIES:
+            table = getattr(dataset, entity)
+            for column in table.columns:
+                values = table[column].to_numpy()
+                if values.dtype == object:
+                    # Population files store strings as fixed-width bytes.
+                    values = values.astype("S")
+                for period in periods:
+                    file.create_dataset(f"{column}/{period}", data=values)
+        for name, period, values in extra:
+            file.create_dataset(f"{name}/{period}", data=values)
+    return path
+
+
+class FileSimulation:
+    """Read a chunk file the way the guard and result checks need it."""
+
+    def __init__(self, path):
+        with h5py.File(path, "r") as file:
+            self.contents = {
+                name: {
+                    period: (array.dtype, array[()])
+                    for period, array in file[name].items()
+                }
+                for name in file
+            }
+        self.tax_benefit_system = system
+        self.branches = {}
+        self.baseline = None
+        self.populations = {
+            entity: SimpleNamespace(ids=self.contents[f"{entity}_id"][str(YEAR)][1])
+            for entity in US_ENTITIES
+        }
+
+    def calculate(self, variable, period, use_weights=False):
+        return self.contents[variable][str(YEAR)][1]
+
+    def get_weights(self, variable, period):
+        assert system.variables[variable].entity.key == "household"
+        return self.contents["household_weight"][str(YEAR)][1]
+
+
+@pytest.fixture
+def variable_centric(tmp_path, dataset):
+    dataset.person["person_household_role"] = ["member", "member"]
+    dataset.person["employment_income"] = [10_000.0, 25_000.0]
+    extra = [
+        # Core ignores names that match no variable, whatever their length.
+        ("not_a_policyengine_variable", PERIODS[1], np.arange(5)),
+        # A period stored for one variable only must survive too.
+        ("employment_income", "2024", np.asarray([1.0, 2.0], dtype=np.float32)),
+    ]
+    return write_variable_centric(tmp_path / "population_2026.h5", dataset, extra=extra)
+
+
+def _chunk_directories():
+    import tempfile
+
+    return set(Path(tempfile.gettempdir()).glob("policyengine-us-chunk-*"))
+
+
+def test_variable_centric_chunks_keep_periods_dtypes_and_name(variable_centric):
+    before = _chunk_directories()
+    paths = []
+    simulations = []
+
+    def factory(path):
+        assert isinstance(path, str)
+        # Only one chunk file exists at a time.
+        assert all(not Path(previous).exists() for previous in paths)
+        assert _chunk_directories() - before == {Path(path).parent}
+        paths.append(path)
+        simulations.append(FileSimulation(path))
+        return simulations[-1]
+
+    result = calculate_chunked(
+        factory, ["household_weight"], YEAR, 2, dataset=str(variable_centric)
+    )
+
+    np.testing.assert_array_equal(
+        result.entities["household"].household_weight, [1.5, 2.25]
+    )
+    assert [Path(path).name for path in paths] == ["population_2026.h5"] * 2
+    assert not any(Path(path).exists() for path in paths)
+    assert _chunk_directories() == before
+    with h5py.File(variable_centric, "r") as source:
+        names = set(source) - {"not_a_policyengine_variable"}
+        for row, simulation in enumerate(simulations):
+            assert set(simulation.contents) == names
+            for name in names:
+                assert set(simulation.contents[name]) == set(source[name])
+                for period, (dtype, values) in simulation.contents[name].items():
+                    stored = source[name][period]
+                    assert dtype == stored.dtype
+                    np.testing.assert_array_equal(values, stored[()][[row]])
+        assert source["person_household_role/2026"].dtype == np.dtype("S6")
+    roles = [s.contents["person_household_role"]["2026"][1] for s in simulations]
+    np.testing.assert_array_equal(np.concatenate(roles), [b"member", b"member"])
+
+
+def test_variable_centric_chunk_file_is_deleted_after_factory_error(variable_centric):
+    before = _chunk_directories()
+    paths = []
+
+    def factory(path):
+        paths.append(path)
+        raise RuntimeError("factory failed")
+
+    with pytest.raises(RuntimeError, match="factory failed"):
+        calculate_chunked(
+            factory, ["household_weight"], YEAR, 2, dataset=variable_centric
+        )
+    assert len(paths) == 1
+    assert not Path(paths[0]).exists()
+    assert _chunk_directories() == before
+
+
+def _rewrite(path, name, periods, values):
+    with h5py.File(path, "a") as file:
+        for period in periods:
+            del file[f"{name}/{period}"]
+            file.create_dataset(f"{name}/{period}", data=values)
+
+
+@pytest.mark.parametrize(
+    "name, values, match",
+    [
+        ("employment_income", [1.0, 2.0, 3.0], r"employment_income/2025.*person has 2"),
+        ("household_weight", [1.0], r"household_weight/2025.*household has 2"),
+        ("person_tax_unit_id", [91, 91], "tax_unit spans multiple households"),
+        ("household_id", [91, 91], "household_id must be unique"),
+        ("person_household_id", [91, 999], "existing household"),
+    ],
+)
+def test_variable_centric_rejects_invalid_arrays(variable_centric, name, values, match):
+    # Every stored period agrees, so only the tested rule can fail.
+    _rewrite(variable_centric, name, PERIODS, np.asarray(values))
+    with pytest.raises(ValueError, match=match):
+        calculate(str(variable_centric))
+
+
+@pytest.mark.parametrize("name", ["household_id", "person_spm_unit_id"])
+def test_variable_centric_rejects_ids_that_change_across_periods(
+    variable_centric, name
+):
+    _rewrite(variable_centric, name, PERIODS[:1], np.asarray([37, 91]))
+    with pytest.raises(ValueError, match="identical across stored periods"):
+        calculate(str(variable_centric))
+
+
+def test_variable_centric_rejects_missing_membership(variable_centric):
+    with h5py.File(variable_centric, "a") as file:
+        del file["person_family_id"]
+    with pytest.raises(ValueError, match="Missing person_family_id"):
+        calculate(str(variable_centric))
+
+
+def test_variable_centric_rejects_external_claiming_tax_unit(variable_centric):
+    with h5py.File(variable_centric, "a") as file:
+        file.create_dataset(
+            f"medicaid_claiming_tax_unit_id/{PERIODS[0]}", data=[37, 91]
+        )
+    with pytest.raises(ValueError, match="claiming_tax_unit_id crosses households"):
+        calculate(str(variable_centric))
+
+
+def test_rejects_flat_and_periodless_array_files(tmp_path, dataset):
+    flat = tmp_path / "population.csv"
+    dataset.person.to_csv(flat, index=False)
+    with pytest.raises(ValueError, match="Flat-file datasets are not supported"):
+        calculate(str(flat))
+    arrays = tmp_path / "arrays.h5"
+    with h5py.File(arrays, "w") as file:
+        for entity in US_ENTITIES:
+            file.create_dataset(
+                f"{entity}_id", data=getattr(dataset, entity)[f"{entity}_id"]
+            )
+    with pytest.raises(ValueError, match="<variable>/<period> layout"):
+        calculate(str(arrays))
+
+
+# Formulas that read a weight or build a MicroSeries can reduce across the
+# simulation. Each must be a weight variable's own definition, listed here
+# with the reason it stays household-local, or guarded in chunks.
+_HOUSEHOLD_LOCAL_WEIGHT_READERS = {}
+_POPULATION_SIGNAL = re.compile(
+    r"MicroSeries|[\"'](?:" + "|".join(US_ENTITIES) + r")_weight[\"']|\.weights\b"
+)
+
+
+def test_population_formulas_are_guarded_or_listed():
+    root = Path(__file__).resolve().parents[2] / "variables"
+    unguarded = []
+    for path in sorted(root.rglob("*.py")):
+        source = path.read_text()
+        if not _POPULATION_SIGNAL.search(source):
+            continue
+        defined = set(re.findall(r"^class (\w+)\(Variable\)", source, re.M))
+        relative = path.relative_to(root).as_posix()
+        if (
+            defined & {f"{entity}_weight" for entity in US_ENTITIES}
+            or defined & _POPULATION_AGGREGATES
+            or relative in _HOUSEHOLD_LOCAL_WEIGHT_READERS
+        ):
+            continue
+        unguarded.append(relative)
+    assert not unguarded, (
+        "Guard these in calculate_chunked._POPULATION_AGGREGATES or list them "
+        f"with a reason: {unguarded}"
+    )
diff --git a/policyengine_us/tests/microsimulation/test_calculate_chunked.py b/policyengine_us/tests/microsimulation/test_calculate_chunked.py
index f6f483f6faf..78f895359bf 100644
--- a/policyengine_us/tests/microsimulation/test_calculate_chunked.py
+++ b/policyengine_us/tests/microsimulation/test_calculate_chunked.py
@@ -15,8 +15,12 @@
 """
 
 import gc
+import sys
+import tempfile
 import weakref
+from pathlib import Path
 
+import h5py
 import numpy as np
 import pandas as pd
 import pytest
@@ -29,7 +33,7 @@
     USSingleYearDataset,
 )
 from policyengine_us.system import system
-from policyengine_us.tools.calculate_chunked import calculate_chunked
+from policyengine_us.tools.calculate_chunked import _guard, calculate_chunked
 
 YEAR = 2026
 HOUSEHOLDS = 32
@@ -530,3 +534,175 @@ def test_entity_hdfstore_paths_match_full_calculation(
             dtype=actual.dtype
         )
         np.testing.assert_array_equal(actual, source_values)
+
+
+# Population projection files store <variable>/<period> arrays, which Core's
+# variable-centric loader reads. The same 32 households are written that way.
+TOB_VARIABLES = ("tob_revenue_oasdi", "tob_revenue_medicare_hi", "tob_revenue_total")
+
+
+def write_variable_centric(dataset, path):
+    with h5py.File(path, "w") as file:
+        for entity in US_ENTITIES:
+            table = getattr(dataset, entity)
+            for column in table.columns:
+                values = table[column].to_numpy()
+                if values.dtype == object:
+                    # Roles are stored as fixed-width byte strings.
+                    values = values.astype("S")
+                file.create_dataset(f"{column}/{YEAR}", data=values)
+        # Core's loader ignores a name that matches no variable, whatever its
+        # length, so the chunked route must too.
+        file.create_dataset(f"not_a_policyengine_variable/{YEAR}", data=np.arange(5))
+    return str(path)
+
+
+@pytest.fixture(scope="module")
+def variable_centric_path(chunking_dataset, tmp_path_factory):
+    return write_variable_centric(
+        chunking_dataset, tmp_path_factory.mktemp("population") / "population.h5"
+    )
+
+
+def _parameter_reform():
+    from policyengine_core.reforms import Reform
+
+    return Reform.from_dict(
+        {
+            # A lower base threshold widens benefit inclusion.
+            "gov.irs.social_security.taxability.threshold.base.main.JOINT": {
+                "2026-01-01.2100-12-31": 20_000
+            },
+        },
+        country_id="us",
+    )
+
+
+@pytest.fixture(scope="module")
+def variable_centric_reference(variable_centric_path, reference):
+    expected = {}
+    for reform in (None, _parameter_reform()):
+        simulation = Microsimulation(dataset=variable_centric_path, reform=reform)
+        expected[reform is not None] = _snapshot_calculations(
+            simulation, TOB_VARIABLES, YEAR
+        )
+        _release_reference_simulation(simulation)
+        del simulation
+        gc.collect()
+    # Both loaders read the same households to the same results.
+    values, weighted_sums, weights = reference
+    for variable, (entity, actual, actual_weights, total) in expected[False].items():
+        np.testing.assert_array_equal(actual, values[variable])
+        np.testing.assert_array_equal(actual_weights, weights[entity])
+        assert total == weighted_sums[variable]
+    assert any(
+        np.any(expected[True][name][1] != expected[False][name][1])
+        for name in TOB_VARIABLES
+    ), "The reform must change TOB on the fixture"
+    return expected
+
+
+def _chunk_directories():
+    return set(Path(tempfile.gettempdir()).glob("policyengine-us-chunk-*"))
+
+
+@pytest.mark.parametrize("n_chunks", [1, 2, 3, 7])
+def test_variable_centric_file_is_elementwise_exact(
+    n_chunks, chunking_dataset, variable_centric_path, variable_centric_reference
+):
+    before = _chunk_directories()
+    paths = []
+    with h5py.File(variable_centric_path, "r") as file:
+        source_roles = file[f"person_household_role/{YEAR}"][()]
+
+    def factory(path):
+        assert isinstance(path, str) and Path(path).name == "population.h5"
+        assert all(not Path(previous).exists() for previous in paths)
+        assert _chunk_directories() - before == {Path(path).parent}
+        with h5py.File(path, "r") as chunk:
+            assert "not_a_policyengine_variable" not in chunk
+            roles = chunk[f"person_household_role/{YEAR}"]
+            assert roles.dtype == source_roles.dtype == np.dtype("S6")
+            people = len(chunk[f"person_id/{YEAR}"])
+            np.testing.assert_array_equal(roles[()], source_roles[:people])
+        paths.append(path)
+        return Microsimulation(dataset=path)
+
+    result = calculate_chunked(
+        factory, TOB_VARIABLES, YEAR, n_chunks, dataset=variable_centric_path
+    )
+    _assert_snapshot_matches(
+        result, chunking_dataset, variable_centric_reference[False]
+    )
+    assert len(paths) == n_chunks
+    assert not any(Path(path).exists() for path in paths)
+    assert _chunk_directories() == before
+
+
+def test_variable_centric_file_is_exact_under_reform(
+    chunking_dataset, variable_centric_path, variable_centric_reference
+):
+    reform = _parameter_reform()
+    result = calculate_chunked(
+        lambda path: Microsimulation(dataset=path, reform=reform),
+        TOB_VARIABLES,
+        YEAR,
+        3,
+        dataset=variable_centric_path,
+    )
+    _assert_snapshot_matches(result, chunking_dataset, variable_centric_reference[True])
+
+
+def _deepest_stack(calculation):
+    depth = 0
+    deepest = 0
+
+    def profile(frame, event, arg):
+        nonlocal depth, deepest
+        if event == "call":
+            depth += 1
+            deepest = max(deepest, depth)
+        elif event == "return":
+            depth -= 1
+
+    sys.setprofile(profile)
+    try:
+        calculation()
+    finally:
+        sys.setprofile(None)
+    return deepest
+
+
+def test_guard_adds_no_frames_per_dependency(chunking_dataset):
+    # One household, as in one chunk. A guard on calculate or _calculate adds
+    # frames to every variable in the chain; the period-check hook adds them
+    # only while it runs, before any dependency starts.
+    households = chunking_dataset.household.household_id.iloc[:1]
+    people = chunking_dataset.person.person_household_id.isin(households)
+    tables = {"person": chunking_dataset.person[people], "household": None}
+    tables["household"] = chunking_dataset.household[
+        chunking_dataset.household.household_id.isin(households)
+    ]
+    for entity in US_ENTITIES:
+        if entity not in tables:
+            table = getattr(chunking_dataset, entity)
+            members = tables["person"][f"person_{entity}_id"]
+            tables[entity] = table[table[f"{entity}_id"].isin(members)]
+    one_household = USSingleYearDataset(**tables, time_period=YEAR)
+    depths = {}
+    for guarded in (False, True):
+        simulation = make_simulation(one_household)
+        if guarded:
+            _guard(simulation)
+        else:
+            simulation.drop_computed_arrays()
+            simulation.macro_cache_read = simulation.macro_cache_write = False
+        depths[guarded] = _deepest_stack(
+            lambda: simulation.calculate("income_tax", YEAR)
+        )
+        _release_reference_simulation(simulation)
+        del simulation
+        gc.collect()
+    # The chain is long: each variable costs several frames.
+    assert depths[False] > 300
+    assert 0 <= depths[True] - depths[False] <= 3, depths
diff --git a/policyengine_us/tools/calculate_chunked.py b/policyengine_us/tools/calculate_chunked.py
index 0156c5d373e..a3cadbf47b0 100644
--- a/policyengine_us/tools/calculate_chunked.py
+++ b/policyengine_us/tools/calculate_chunked.py
@@ -5,6 +5,9 @@
 import gc
 from numbers import Integral
 from pathlib import Path
+import shutil
+import tempfile
+from types import SimpleNamespace
 
 import numpy as np
 import pandas as pd
@@ -16,30 +19,38 @@
 )
 
 
-# These formulas reduce across the population, rather than within a household.
-# Guard at calculation time so parameter-selected and reform dependencies are
-# covered too, including calculations made in branches.
+# These formulas reduce or rank across every unit in the simulation, rather
+# than within a household: the Medicaid state aggregates sum by state, and the
+# income deciles rank with MicroSeries.decile_rank. Guard at calculation time
+# so parameter-selected and reform dependencies are covered too, including
+# calculations made in branches, and readers such as income_decile.
 _POPULATION_AGGREGATES = frozenset(
-    ("medicaid_slcsp_state_average_cost_index", "medicaid_slcsp_state_denominator")
+    (
+        "medicaid_slcsp_state_average_cost_index",
+        "medicaid_slcsp_state_denominator",
+        "household_income_decile",
+        "spm_unit_income_decile",
+    )
 )
 
 
 class _HouseholdLocalCalculations:
-    @staticmethod
-    def _check_household_local(variable_name):
-        if variable_name in _POPULATION_AGGREGATES:
+    """Reject population aggregates just before their formula would run.
+
+    Core's ``Simulation._calculate`` calls ``_check_period_consistency`` after
+    the stored-value lookups and before the formula, once per formula
+    evaluation, and the check returns before any dependency starts. Hooking it
+    adds no frames to the dependency chain, unlike overriding ``calculate``.
+    A stored value is returned without the check: no formula reduces then.
+    """
+
+    def _check_period_consistency(self, period, variable):
+        if variable.name in _POPULATION_AGGREGATES:
             raise ValueError(
-                f"Cannot calculate {variable_name} in household chunks: "
+                f"Cannot calculate {variable.name} in household chunks: "
                 "it aggregates across households. Calculate it on the full dataset."
             )
-
-    def calculate(self, variable_name, *args, **kwargs):
-        self._check_household_local(variable_name)
-        return super().calculate(variable_name, *args, **kwargs)
-
-    def _calculate(self, variable_name, period=None):
-        self._check_household_local(variable_name)
-        return super()._calculate(variable_name, period)
+        return super()._check_period_consistency(period, variable)
 
 
 @lru_cache(maxsize=None)
@@ -113,10 +124,16 @@ class ChunkedCalculation:
 
 
 def _load_dataset(dataset):
+    """Return an entity-table dataset, or a variable-centric file source."""
     if isinstance(dataset, (str, Path)):
-        from policyengine_us.system import _resolve_dataset_path
+        from policyengine_us.system import (
+            _is_hdfstore_format,
+            _resolve_dataset_path,
+        )
 
         path = _resolve_dataset_path(str(dataset))
+        if not _is_hdfstore_format(path):
+            return _VariableCentricFile(path)
         with pd.HDFStore(path, mode="r") as store:
             multi_year = any(key.startswith("/person/") for key in store.keys())
         dataset = (
@@ -127,7 +144,7 @@ def _load_dataset(dataset):
     if not isinstance(dataset, (USSingleYearDataset, USMultiYearDataset)):
         raise TypeError(
             "dataset must be a USSingleYearDataset, USMultiYearDataset, "
-            "or entity-table HDFStore path/URL."
+            "or an entity-table HDFStore or variable-centric h5 path/URL."
         )
     return dataset
 
@@ -183,6 +200,202 @@ def _household_positions(dataset):
     return positions
 
 
+class _EntityTables:
+    """Entity-table datasets: chunks are dataset objects of the same type."""
+
+    def __init__(self, dataset):
+        self.dataset = dataset
+        self.years = (
+            list(dataset.datasets.values())
+            if isinstance(dataset, USMultiYearDataset)
+            else [dataset]
+        )
+        source = self.years[0]
+        self.positions = _household_positions(source)
+        self.household_count = len(source.household)
+        if not self.household_count:
+            raise ValueError("dataset must contain at least one household.")
+        for year in self.years[1:]:
+            _household_positions(year)
+            for entity in US_ENTITIES:
+                columns = [f"{entity}_id"]
+                if entity == "person":
+                    columns += [
+                        f"person_{group}_id"
+                        for group in US_ENTITIES
+                        if group != "person"
+                    ]
+                if (
+                    not getattr(source, entity)[columns]
+                    .reset_index(drop=True)
+                    .equals(getattr(year, entity)[columns].reset_index(drop=True))
+                ):
+                    raise ValueError(
+                        "Entity IDs and memberships must be identical across years."
+                    )
+
+    def ids(self, entity):
+        return getattr(self.years[0], entity)[f"{entity}_id"].to_numpy()
+
+    def chunk(self, masks):
+        subsets = [
+            USSingleYearDataset(
+                **{
+                    entity: getattr(year, entity)
+                    .iloc[np.flatnonzero(masks[entity])]
+                    .copy()
+                    for entity in US_ENTITIES
+                },
+                time_period=int(year.time_period),
+            )
+            for year in self.years
+        ]
+        return (
+            USMultiYearDataset(datasets=subsets)
+            if isinstance(self.dataset, USMultiYearDataset)
+            else subsets[0]
+        )
+
+    def release(self, chunk):
+        pass
+
+
+_GROUPS = tuple(entity for entity in US_ENTITIES if entity != "person")
+# Core reads these from the first stored period as the population structure.
+_STRUCTURE = tuple(f"{entity}_id" for entity in US_ENTITIES) + tuple(
+    f"person_{group}_id" for group in _GROUPS
+)
+# Core's loader also reads membership roles, which are not variables.
+_ROLES = frozenset(f"person_{group}_role" for group in _GROUPS) | {"role"}
+
+
+class _VariableCentricFile:
+    """An h5py file of ``<variable>/<period>`` arrays, as Core loads it.
+
+    Arrays stay on disk. Each chunk is written to a temporary file in the same
+    layout, under the source file's name, and is deleted after its simulation.
+    """
+
+    def __init__(self, path):
+        import h5py
+
+        from policyengine_us.system import system
+
+        self.path = Path(path)
+        if self.path.suffix != ".h5":
+            # Core loads a file without the .h5 suffix as a CSV flat file.
+            raise ValueError(
+                f"{self.path} is not an entity-table HDFStore or a "
+                "variable-centric .h5 file. Flat-file datasets are not "
+                "supported; convert them to one of those layouts."
+            )
+        self.arrays = {}
+        with h5py.File(self.path, "r") as file:
+            names = list(file.keys())
+            if not names or isinstance(file[names[0]], h5py.Dataset):
+                raise ValueError(
+                    f"{self.path} does not use the variable-centric "
+                    "<variable>/<period> layout; only that layout and entity "
+                    "tables are supported."
+                )
+            for name in names:
+                if name in _ROLES:
+                    entity = "person"
+                elif name in system.variables:
+                    entity = system.variables[name].entity.key
+                else:
+                    # Core's loader ignores names that match no variable.
+                    continue
+                if not isinstance(file[name], h5py.Group):
+                    raise ValueError(f"{name} must be a group of one array per period.")
+                self.arrays[name] = (entity, list(file[name].keys()))
+            # A group with no periods holds no array for Core to read.
+            missing = [
+                name
+                for name in _STRUCTURE
+                if name not in self.arrays or not self.arrays[name][1]
+            ]
+            if missing:
+                raise ValueError(
+                    f"Missing {', '.join(missing)}; every entity ID and person "
+                    "membership array is required."
+                )
+            self.counts = {
+                entity: file[f"{entity}_id"][self.arrays[f"{entity}_id"][1][0]].shape[0]
+                for entity in US_ENTITIES
+            }
+            for name, (entity, stored_periods) in self.arrays.items():
+                for period in stored_periods:
+                    shape = file[name][period].shape
+                    if shape != (self.counts[entity],):
+                        raise ValueError(
+                            f"{name}/{period} has shape {shape}, but {entity} "
+                            f"has {self.counts[entity]} rows."
+                        )
+            structure = {}
+            for name in _STRUCTURE:
+                stored_periods = self.arrays[name][1]
+                first = file[name][stored_periods[0]][()]
+                for period in stored_periods[1:]:
+                    if not np.array_equal(file[name][period][()], first):
+                        raise ValueError(
+                            "Entity IDs and memberships must be identical "
+                            "across stored periods."
+                        )
+                structure[name] = first
+            claimants = [None]
+            if "medicaid_claiming_tax_unit_id" in self.arrays:
+                group = file["medicaid_claiming_tax_unit_id"]
+                claimants = [
+                    group[period][()]
+                    for period in self.arrays["medicaid_claiming_tax_unit_id"][1]
+                ]
+        self.structure = structure
+        for claimant in claimants:
+            self.positions = _household_positions(self._tables(claimant))
+        self.household_count = self.counts["household"]
+        if not self.household_count:
+            raise ValueError("dataset must contain at least one household.")
+
+    def _tables(self, claimant):
+        tables = {
+            entity: pd.DataFrame({f"{entity}_id": self.structure[f"{entity}_id"]})
+            for entity in _GROUPS
+        }
+        person = {"person_id": self.structure["person_id"]}
+        for group in _GROUPS:
+            person[f"person_{group}_id"] = self.structure[f"person_{group}_id"]
+        if claimant is not None:
+            person["medicaid_claiming_tax_unit_id"] = claimant
+        tables["person"] = pd.DataFrame(person)
+        return SimpleNamespace(**tables)
+
+    def ids(self, entity):
+        return self.structure[f"{entity}_id"]
+
+    def chunk(self, masks):
+        import h5py
+
+        directory = Path(tempfile.mkdtemp(prefix="policyengine-us-chunk-"))
+        try:
+            path = directory / self.path.name
+            with h5py.File(self.path, "r") as source, h5py.File(path, "w") as chunk:
+                for name, (entity, stored_periods) in self.arrays.items():
+                    group = chunk.create_group(name)
+                    for period in stored_periods:
+                        array = source[name][period]
+                        group.create_dataset(
+                            period, data=array[()][masks[entity]], dtype=array.dtype
+                        )
+        except BaseException:
+            shutil.rmtree(directory, ignore_errors=True)
+            raise
+        return str(path)
+
+    def release(self, chunk):
+        shutil.rmtree(Path(chunk).parent, ignore_errors=True)
+
+
 def calculate_chunked(
     sim_factory,
     variables,
@@ -196,9 +409,19 @@ def calculate_chunked(
 
     ``sim_factory(dataset)`` must return a fresh PolicyEngine-US Microsimulation
     and preserve its dataset's population/order. It controls reforms and other
-    simulation options. ``dataset`` accepts US entity-table dataset objects or
-    HDFStore paths/URLs. All six entity tables and membership links are required.
-    All years are retained; IDs and memberships must be constant across years.
+    simulation options. ``dataset`` accepts US entity-table dataset objects,
+    entity-table HDFStore paths/URLs, and variable-centric h5 paths/URLs: files
+    of ``<variable>/<period>`` arrays, as population projection files store
+    them. All six entity ID arrays and membership links are required. All years
+    are retained; IDs and memberships must be constant across years.
+
+    For entity tables the factory receives a dataset object of the source's
+    type. For a variable-centric file it receives ``str`` path to a temporary
+    file in the same layout, holding every stored period with its dtype under
+    the source file's name, for ``Microsimulation(dataset=path)``. Names that
+    match no variable are left out, as Core's loader ignores them. Only one
+    chunk file exists at once; it is deleted after its chunk, on failure too.
+    Flat files are not supported.
 
     By default households are split into balanced consecutive chunks. Optional
     ``household_chunks`` supplies one integer in ``[0, n_chunks)`` per source
@@ -211,9 +434,11 @@ def calculate_chunked(
     Input weights are never rescaled. Only one chunk simulation is live at once.
 
     This is exact for household-local calculations. Known Medicaid state
-    aggregates and external claiming-tax-unit links are rejected. Custom reform
-    formulas must also be household-local: population reductions, ranking and
-    random draws depending on population order are not supported.
+    aggregates, the household and SPM unit income deciles, and external
+    claiming-tax-unit links are rejected. A rejected variable with a stored
+    value returns that value, since no formula runs. Custom reform formulas
+    must also be household-local: population reductions, ranking and random
+    draws depending on population order are not supported.
     """
     if isinstance(n_chunks, bool) or not isinstance(n_chunks, Integral) or n_chunks < 1:
         raise ValueError("n_chunks must be a positive integer.")
@@ -223,33 +448,12 @@ def calculate_chunked(
     if "weight" in variables:
         raise ValueError("weight is reserved for the returned simulation weights.")
     dataset = _load_dataset(dataset)
-    years = (
-        list(dataset.datasets.values())
-        if isinstance(dataset, USMultiYearDataset)
-        else [dataset]
+    source = (
+        dataset if isinstance(dataset, _VariableCentricFile) else _EntityTables(dataset)
     )
-    source = years[0]
-    positions = _household_positions(source)
-    if not len(source.household):
-        raise ValueError("dataset must contain at least one household.")
-    for year in years[1:]:
-        _household_positions(year)
-        for entity in US_ENTITIES:
-            columns = [f"{entity}_id"]
-            if entity == "person":
-                columns += [
-                    f"person_{group}_id" for group in US_ENTITIES if group != "person"
-                ]
-            if (
-                not getattr(source, entity)[columns]
-                .reset_index(drop=True)
-                .equals(getattr(year, entity)[columns].reset_index(drop=True))
-            ):
-                raise ValueError(
-                    "Entity IDs and memberships must be identical across years."
-                )
+    positions = source.positions
     if household_chunks is None:
-        assignments = np.empty(len(source.household), dtype=int)
+        assignments = np.empty(source.household_count, dtype=int)
         for chunk, indices in enumerate(
             np.array_split(np.arange(len(assignments)), min(n_chunks, len(assignments)))
         ):
@@ -257,7 +461,7 @@ def calculate_chunked(
     else:
         assignments = np.asarray(household_chunks)
         if (
-            assignments.shape != (len(source.household),)
+            assignments.shape != (source.household_count,)
             or not np.issubdtype(assignments.dtype, np.integer)
             or (assignments < 0).any()
             or (assignments >= n_chunks).any()
@@ -272,25 +476,10 @@ def calculate_chunked(
         masks = {
             entity: assignments[owners] == chunk for entity, owners in positions.items()
         }
-        subsets = [
-            USSingleYearDataset(
-                **{
-                    entity: getattr(year, entity)
-                    .iloc[np.flatnonzero(masks[entity])]
-                    .copy()
-                    for entity in US_ENTITIES
-                },
-                time_period=int(year.time_period),
-            )
-            for year in years
-        ]
-        chunk_dataset = (
-            USMultiYearDataset(datasets=subsets)
-            if isinstance(dataset, USMultiYearDataset)
-            else subsets[0]
-        )
+        chunk_dataset = None
         simulation = None
         try:
+            chunk_dataset = source.chunk(masks)
             simulation = sim_factory(chunk_dataset)
             _guard(simulation)
             for name in variables:
@@ -302,7 +491,7 @@ def calculate_chunked(
                         "sim_factory must use the same variable entities in every chunk."
                     )
                 variable_entities[name] = entity
-                original_ids = getattr(source, entity)[f"{entity}_id"].to_numpy()
+                original_ids = source.ids(entity)
                 if not np.array_equal(
                     simulation.populations[entity].ids, original_ids[masks[entity]]
                 ):
@@ -347,8 +536,11 @@ def calculate_chunked(
             if simulation is not None:
                 _detach_simulation_links(simulation)
             del simulation
-            del chunk_dataset, subsets
             gc.collect()
+            if chunk_dataset is not None:
+                # After collection, so no simulation still reads a chunk file.
+                source.release(chunk_dataset)
+            del chunk_dataset
     sums = {
         name: float((tables[entity][name] * tables[entity]["weight"]).sum())
         for name, entity in variable_entities.items()

From 5b2c218a411ede507d498b732bc5f9f17c06f406 Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Mon, 2 Nov 2026 15:30:55 +0000
Subject: [PATCH 08/15] Give non-HDF5 paths the flat-file error and reapply
 test readers safely

_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>
---
 .../tests/core/test_calculate_chunked_validation.py        | 3 ++-
 policyengine_us/tools/calculate_chunked.py                 | 7 ++++++-
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/policyengine_us/tests/core/test_calculate_chunked_validation.py b/policyengine_us/tests/core/test_calculate_chunked_validation.py
index 79024a939d6..3f2e19cec0b 100644
--- a/policyengine_us/tests/core/test_calculate_chunked_validation.py
+++ b/policyengine_us/tests/core/test_calculate_chunked_validation.py
@@ -239,7 +239,8 @@ def monthly_formula(spm_unit, period, parameters):
     class AddReaders(Reform):
         def apply(self):
             for reader in [*readers.values(), monthly]:
-                self.add_variable(reader)
+                # The country simulation applies a reform more than once.
+                self.update_variable(reader)
 
     return AddReaders
 
diff --git a/policyengine_us/tools/calculate_chunked.py b/policyengine_us/tools/calculate_chunked.py
index a3cadbf47b0..842e2434e30 100644
--- a/policyengine_us/tools/calculate_chunked.py
+++ b/policyengine_us/tools/calculate_chunked.py
@@ -132,7 +132,12 @@ def _load_dataset(dataset):
         )
 
         path = _resolve_dataset_path(str(dataset))
-        if not _is_hdfstore_format(path):
+        try:
+            entity_tables = _is_hdfstore_format(path)
+        except Exception:
+            # PyTables raises its own error for a file that is not HDF5.
+            entity_tables = False
+        if not entity_tables:
             return _VariableCentricFile(path)
         with pd.HDFStore(path, mode="r") as store:
             multi_year = any(key.startswith("/person/") for key in store.keys())

From bd77f4b685042950d45f4f715635e11396819713 Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Mon, 2 Nov 2026 15:44:31 +0000
Subject: [PATCH 09/15] Run the reformed population-file case in two chunks

Three reformed chunks took 166 s; chunk counts are already covered
unreformed.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 policyengine_us/tests/microsimulation/test_calculate_chunked.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/policyengine_us/tests/microsimulation/test_calculate_chunked.py b/policyengine_us/tests/microsimulation/test_calculate_chunked.py
index 78f895359bf..56a941bfc54 100644
--- a/policyengine_us/tests/microsimulation/test_calculate_chunked.py
+++ b/policyengine_us/tests/microsimulation/test_calculate_chunked.py
@@ -647,7 +647,7 @@ def test_variable_centric_file_is_exact_under_reform(
         lambda path: Microsimulation(dataset=path, reform=reform),
         TOB_VARIABLES,
         YEAR,
-        3,
+        2,
         dataset=variable_centric_path,
     )
     _assert_snapshot_matches(result, chunking_dataset, variable_centric_reference[True])

From 6835df5c68cd6541b6a4c163f8151690ec30e66c Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Mon, 2 Nov 2026 15:57:13 +0000
Subject: [PATCH 10/15] Describe population-file support in the changelog
 fragment

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 changelog.d/chunked-microsimulation-calculate.added.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/changelog.d/chunked-microsimulation-calculate.added.md b/changelog.d/chunked-microsimulation-calculate.added.md
index 54395da5193..ebc58caf6f6 100644
--- a/changelog.d/chunked-microsimulation-calculate.added.md
+++ b/changelog.d/chunked-microsimulation-calculate.added.md
@@ -1 +1 @@
-Add exact household-chunked microsimulation calculations and share sequential TOB counterfactual tax calculations.
+Add exact household-chunked microsimulation calculations, including variable-centric population files, and share sequential TOB counterfactual tax calculations.

From 21204b6a8f4b0375241b44693d88bcfb68a5768c Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Tue, 3 Nov 2026 00:33:26 +0000
Subject: [PATCH 11/15] Reject chunk inputs that would silently differ from the
 whole file

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
---
 docs/usage/chunked-calculations.md            |  15 +-
 docs/usage/social-security-pia-ratio.md       |  20 +-
 .../test_social_security_preflight.py         |   3 +
 .../core/test_calculate_chunked_validation.py | 198 +++++++++++++++---
 policyengine_us/tools/calculate_chunked.py    |  50 ++++-
 .../tools/social_security_preflight.py        |   3 +
 6 files changed, 247 insertions(+), 42 deletions(-)

diff --git a/docs/usage/chunked-calculations.md b/docs/usage/chunked-calculations.md
index 724682a2ed9..577c60a2713 100644
--- a/docs/usage/chunked-calculations.md
+++ b/docs/usage/chunked-calculations.md
@@ -39,8 +39,10 @@ array's dtype, fixed-width byte strings included, so the factory can call
 `Microsimulation(dataset=path)`. The source arrays stay on disk and are read
 again for each chunk. Each array is assigned to the entity of its variable in
 the country system; arrays whose length differs from that entity's count are
-rejected, and names that match no variable are left out, as Core's loader
-ignores them. Membership role arrays (`person_<group>_role`), which Core also
+rejected, and names that match no country-system variable are left out, as
+Core's loader ignores them. The helper raises when the factory's system (a
+reform, say) defines a name it left out, since Core would load that name in an
+unchunked simulation. Membership role arrays (`person_<group>_role`), which Core also
 reads, are kept with the people. Entity IDs and memberships must be identical
 in every stored period. Only one chunk file exists at once; it is deleted after
 its chunk, also when the factory or a calculation fails.
@@ -81,8 +83,8 @@ whole-file chain does. A guarded variable with a stored value returns that
 value, because no formula reduces across the chunk. Known Medicaid
 claiming-tax-unit links crossing households are also rejected. The SSI spouse
 pairing also tests whether the whole simulation holds a single marital unit, so
-a chunk with exactly one marital unit can differ from the whole file; this is
-not guarded. Custom reform formulas remain the caller's responsibility:
+a chunk with exactly one marital unit can differ from the whole file; such
+chunks are rejected. Custom reform formulas remain the caller's responsibility:
 cross-household reductions, order-dependent random draws and other global
 dependencies are unsupported. These are calculation restrictions, not an
 assumption that every model variable can be chunked.
@@ -120,6 +122,5 @@ no-inclusion tax, and HI remains the signed residual. No per-unit floor is
 introduced. See [OASDI modeled liabilities and benefits](social-security-oasdi-accounting.md)
 for the accounting these outputs feed.
 
-The hub will add measured full-dataset peak memory and timing after running its
-memory probe. Small-fixture exactness establishes correctness, not a measured
-full-population RSS reduction.
+The PR records measured full-dataset peak memory and timing; small-fixture
+exactness establishes correctness, not an RSS bound.
diff --git a/docs/usage/social-security-pia-ratio.md b/docs/usage/social-security-pia-ratio.md
index 50747a01651..f87f730a81a 100644
--- a/docs/usage/social-security-pia-ratio.md
+++ b/docs/usage/social-security-pia-ratio.md
@@ -310,11 +310,21 @@ from policyengine_core.reforms import Reform
 reform = Reform.from_dict(
     {
         "gov.contrib.ssa.pia.in_effect": {"2032-01-01": True},
-        "gov.contrib.ssa.pia.progressive_price_indexing.in_effect": {"2032-01-01": True},
-        "gov.contrib.ssa.pia.progressive_price_indexing.start_year": {"2032-01-01": 2032},
-        "gov.contrib.ssa.pia.progressive_price_indexing.bend_point_fraction": {"2032-01-01": 0.25},
-        "gov.contrib.ssa.pia.progressive_price_indexing.maximum_aime": {"2035-01-01": 17_000},
-        "gov.contrib.ssa.pia.progressive_price_indexing.maximum_pia_target": {"2035-01-01": 5_100},
+        "gov.contrib.ssa.pia.progressive_price_indexing.in_effect": {
+            "2032-01-01": True
+        },
+        "gov.contrib.ssa.pia.progressive_price_indexing.start_year": {
+            "2032-01-01": 2032
+        },
+        "gov.contrib.ssa.pia.progressive_price_indexing.bend_point_fraction": {
+            "2032-01-01": 0.25
+        },
+        "gov.contrib.ssa.pia.progressive_price_indexing.maximum_aime": {
+            "2035-01-01": 17_000
+        },
+        "gov.contrib.ssa.pia.progressive_price_indexing.maximum_pia_target": {
+            "2035-01-01": 5_100
+        },
     },
     country_id="us",
 )
diff --git a/policyengine_us/tests/code_health/test_social_security_preflight.py b/policyengine_us/tests/code_health/test_social_security_preflight.py
index 2cb1ccbf118..0dcb226ee70 100644
--- a/policyengine_us/tests/code_health/test_social_security_preflight.py
+++ b/policyengine_us/tests/code_health/test_social_security_preflight.py
@@ -100,6 +100,9 @@ def test_reported_columns_are_not_migrated_again(tmp_path):
         "self_employment_social_security_tax",
         "employer_total_social_security_tax",
         "income_tax_chapter_1",
+        # TOB counterfactual holders: a stored zero sets TOB to chapter 1 tax.
+        "income_tax_without_social_security",
+        "income_tax_with_pre_1993_social_security",
         "net_investment_income_tax",
         "niit_liability",
         "niit_revenue_oasdi",
diff --git a/policyengine_us/tests/core/test_calculate_chunked_validation.py b/policyengine_us/tests/core/test_calculate_chunked_validation.py
index 3f2e19cec0b..c4af02f82b5 100644
--- a/policyengine_us/tests/core/test_calculate_chunked_validation.py
+++ b/policyengine_us/tests/core/test_calculate_chunked_validation.py
@@ -23,19 +23,25 @@
 )
 
 YEAR = 2026
+MARITAL_UNIT_IDS = [91, 92, 37, 38]
 
 
 @pytest.fixture
 def dataset():
+    # Two households of two people. Each person is their own marital unit, so
+    # a one-household chunk holds two marital units and passes the SSI-spouse
+    # guard; every other entity has one row per household.
     ids = np.asarray([91, 37])
-    person = pd.DataFrame({"person_id": [503, 101]})
+    person = pd.DataFrame({"person_id": [503, 504, 101, 102]})
     for entity in US_ENTITIES:
-        if entity != "person":
-            person[f"person_{entity}_id"] = ids
+        if entity not in ("person", "marital_unit"):
+            person[f"person_{entity}_id"] = np.repeat(ids, 2)
+    person["person_marital_unit_id"] = MARITAL_UNIT_IDS
     tables = {"person": person}
     for entity in US_ENTITIES:
-        if entity != "person":
+        if entity not in ("person", "marital_unit"):
             tables[entity] = pd.DataFrame({f"{entity}_id": ids})
+    tables["marital_unit"] = pd.DataFrame({"marital_unit_id": MARITAL_UNIT_IDS})
     tables["household"]["household_weight"] = [1.5, 2.25]
     return USSingleYearDataset(**tables, time_period=YEAR)
 
@@ -88,7 +94,7 @@ def test_rejects_duplicate_entity_ids(dataset, entity):
 
 @pytest.mark.parametrize("entity", ["tax_unit", "spm_unit", "family", "marital_unit"])
 def test_rejects_entity_spanning_households(dataset, entity):
-    dataset.person[f"person_{entity}_id"] = [91, 91]
+    dataset.person[f"person_{entity}_id"] = 91
     with pytest.raises(ValueError, match=f"{entity} spans multiple households"):
         calculate(dataset)
 
@@ -113,7 +119,7 @@ def test_rejects_unreferenced_group_entity(dataset, entity):
 
 
 def test_rejects_external_claiming_tax_unit(dataset):
-    dataset.person["medicaid_claiming_tax_unit_id"] = [37, 91]
+    dataset.person["medicaid_claiming_tax_unit_id"] = [37, 37, 91, 91]
     with pytest.raises(ValueError, match="claiming_tax_unit_id crosses households"):
         calculate(dataset)
 
@@ -121,7 +127,7 @@ def test_rejects_external_claiming_tax_unit(dataset):
 def test_rejects_changed_multi_year_memberships(dataset):
     other = dataset.copy()
     other.time_period = str(YEAR + 1)
-    other.person["person_tax_unit_id"] = [37, 91]
+    other.person["person_tax_unit_id"] = [37, 37, 91, 91]
     multi = USMultiYearDataset(datasets=[dataset, other])
     with pytest.raises(ValueError, match="memberships must be identical"):
         calculate(multi)
@@ -361,9 +367,9 @@ def get_weights(self, variable, period):
 
 
 def test_skips_empty_chunks_and_preserves_native_and_projected_weights(dataset):
-    dataset.person["employment_income"] = [10_000.0, 25_000.0]
+    dataset.person["employment_income"] = [10_000.0, 0.0, 25_000.0, 0.0]
     # Core uses projected household weights; native input columns still survive.
-    dataset.person["person_weight"] = [999.0, 888.0]
+    dataset.person["person_weight"] = [999.0, 998.0, 888.0, 887.0]
     snapshots = []
 
     def factory(chunk):
@@ -381,7 +387,9 @@ def factory(chunk):
     np.testing.assert_array_equal(
         result.entities["person"].person_id, dataset.person.person_id
     )
-    np.testing.assert_array_equal(result.entities["person"].weight, [1.5, 2.25])
+    np.testing.assert_array_equal(
+        result.entities["person"].weight, [1.5, 1.5, 2.25, 2.25]
+    )
     np.testing.assert_array_equal(
         result.entities["household"].household_weight, [1.5, 2.25]
     )
@@ -389,7 +397,9 @@ def factory(chunk):
     np.testing.assert_array_equal(
         pd.concat(snapshots).person_weight, dataset.person.person_weight
     )
-    np.testing.assert_array_equal(dataset.person.person_weight, [999.0, 888.0])
+    np.testing.assert_array_equal(
+        dataset.person.person_weight, [999.0, 998.0, 888.0, 887.0]
+    )
 
 
 @pytest.mark.parametrize("route", ["baseline", "branch", "popped_branch"])
@@ -511,13 +521,17 @@ def get_weights(self, variable, period):
 
 @pytest.fixture
 def variable_centric(tmp_path, dataset):
-    dataset.person["person_household_role"] = ["member", "member"]
-    dataset.person["employment_income"] = [10_000.0, 25_000.0]
+    dataset.person["person_household_role"] = ["member"] * 4
+    dataset.person["employment_income"] = [10_000.0, 0.0, 25_000.0, 0.0]
     extra = [
         # Core ignores names that match no variable, whatever their length.
         ("not_a_policyengine_variable", PERIODS[1], np.arange(5)),
         # A period stored for one variable only must survive too.
-        ("employment_income", "2024", np.asarray([1.0, 2.0], dtype=np.float32)),
+        (
+            "employment_income",
+            "2024",
+            np.asarray([1.0, 2.0, 3.0, 4.0], dtype=np.float32),
+        ),
     ]
     return write_variable_centric(tmp_path / "population_2026.h5", dataset, extra=extra)
 
@@ -561,10 +575,19 @@ def factory(path):
                 for period, (dtype, values) in simulation.contents[name].items():
                     stored = source[name][period]
                     assert dtype == stored.dtype
-                    np.testing.assert_array_equal(values, stored[()][[row]])
+                    # People and marital units: two rows per household.
+                    entity = (
+                        "person"
+                        if name.endswith("_role")
+                        else system.variables[name].entity.key
+                    )
+                    per = 2 if entity in ("person", "marital_unit") else 1
+                    np.testing.assert_array_equal(
+                        values, stored[()][row * per : (row + 1) * per]
+                    )
         assert source["person_household_role/2026"].dtype == np.dtype("S6")
     roles = [s.contents["person_household_role"]["2026"][1] for s in simulations]
-    np.testing.assert_array_equal(np.concatenate(roles), [b"member", b"member"])
+    np.testing.assert_array_equal(np.concatenate(roles), [b"member"] * 4)
 
 
 def test_variable_centric_chunk_file_is_deleted_after_factory_error(variable_centric):
@@ -594,11 +617,11 @@ def _rewrite(path, name, periods, values):
 @pytest.mark.parametrize(
     "name, values, match",
     [
-        ("employment_income", [1.0, 2.0, 3.0], r"employment_income/2025.*person has 2"),
+        ("employment_income", [1.0, 2.0, 3.0], r"employment_income/2025.*person has 4"),
         ("household_weight", [1.0], r"household_weight/2025.*household has 2"),
-        ("person_tax_unit_id", [91, 91], "tax_unit spans multiple households"),
+        ("person_tax_unit_id", [91] * 4, "tax_unit spans multiple households"),
         ("household_id", [91, 91], "household_id must be unique"),
-        ("person_household_id", [91, 999], "existing household"),
+        ("person_household_id", [91, 91, 999, 999], "existing household"),
     ],
 )
 def test_variable_centric_rejects_invalid_arrays(variable_centric, name, values, match):
@@ -612,7 +635,8 @@ def test_variable_centric_rejects_invalid_arrays(variable_centric, name, values,
 def test_variable_centric_rejects_ids_that_change_across_periods(
     variable_centric, name
 ):
-    _rewrite(variable_centric, name, PERIODS[:1], np.asarray([37, 91]))
+    count = 2 if name == "household_id" else 4
+    _rewrite(variable_centric, name, PERIODS[:1], np.repeat([37, 91], count // 2))
     with pytest.raises(ValueError, match="identical across stored periods"):
         calculate(str(variable_centric))
 
@@ -627,7 +651,7 @@ def test_variable_centric_rejects_missing_membership(variable_centric):
 def test_variable_centric_rejects_external_claiming_tax_unit(variable_centric):
     with h5py.File(variable_centric, "a") as file:
         file.create_dataset(
-            f"medicaid_claiming_tax_unit_id/{PERIODS[0]}", data=[37, 91]
+            f"medicaid_claiming_tax_unit_id/{PERIODS[0]}", data=[37, 37, 91, 91]
         )
     with pytest.raises(ValueError, match="claiming_tax_unit_id crosses households"):
         calculate(str(variable_centric))
@@ -648,16 +672,31 @@ def test_rejects_flat_and_periodless_array_files(tmp_path, dataset):
         calculate(str(arrays))
 
 
-# Formulas that read a weight or build a MicroSeries can reduce across the
-# simulation. Each must be a weight variable's own definition, listed here
-# with the reason it stays household-local, or guarded in chunks.
+# Formulas that read a weight, build a MicroSeries, test an entity count or
+# draw random numbers can depend on the whole simulation. Each must be a weight
+# variable's own definition, listed here with the reason it is safe in chunks,
+# or guarded in chunks.
 _HOUSEHOLD_LOCAL_WEIGHT_READERS = {}
+_CHUNK_SAFE_POPULATION_READERS = {
+    # Tests whether the simulation holds one marital unit. calculate_chunked
+    # rejects a chunk holding exactly one of several marital units.
+    "gov/ssa/ssi/eligibility/resources/deemed/_ssi_spouses.py": "chunk guard",
+}
 _POPULATION_SIGNAL = re.compile(
-    r"MicroSeries|[\"'](?:" + "|".join(US_ENTITIES) + r")_weight[\"']|\.weights\b"
+    r"MicroSeries|[\"'](?:"
+    + "|".join(US_ENTITIES)
+    + r")_weight[\"']|\.weights\b|\.count\s*==|np\.random|default_rng"
 )
 
 
 def test_population_formulas_are_guarded_or_listed():
+    """Scan variable sources for population signals.
+
+    The scan sees weight reads, MicroSeries, ``.count ==`` tests of entity
+    sizes and NumPy random draws. It cannot see other whole-array reductions
+    (``np.sum``, ``np.unique``, sorting) or branches on ``.any()``/``.all()``,
+    and it does not read ``reforms/``; those still need review by hand.
+    """
     root = Path(__file__).resolve().parents[2] / "variables"
     unguarded = []
     for path in sorted(root.rglob("*.py")):
@@ -670,6 +709,7 @@ def test_population_formulas_are_guarded_or_listed():
             defined & {f"{entity}_weight" for entity in US_ENTITIES}
             or defined & _POPULATION_AGGREGATES
             or relative in _HOUSEHOLD_LOCAL_WEIGHT_READERS
+            or relative in _CHUNK_SAFE_POPULATION_READERS
         ):
             continue
         unguarded.append(relative)
@@ -677,3 +717,111 @@ def test_population_formulas_are_guarded_or_listed():
         "Guard these in calculate_chunked._POPULATION_AGGREGATES or list them "
         f"with a reason: {unguarded}"
     )
+
+
+def _spouse_and_dependent_households():
+    """Household 1: a 70-year-old and a 40-year-old dependent in one marital
+    and tax unit, the dependent holding $100,000 of countable resources.
+    Household 2: one 30-year-old.
+
+    Unchunked, ssi is [0, 0, 0] and spousal resource deeming applies to the
+    first person. Household 1 alone holds one marital unit, so _ssi_spouses
+    rules the pair out as a couple and the first person's ssi is $11,928.
+    """
+    person = pd.DataFrame(
+        {
+            "person_id": [1, 2, 3],
+            "age": [70, 40, 30],
+            "is_tax_unit_dependent": [False, True, False],
+            "ssi_countable_resources": [0.0, 100_000.0, 0.0],
+        }
+    )
+    tables = {"person": person}
+    for entity in US_ENTITIES:
+        if entity != "person":
+            person[f"person_{entity}_id"] = [1, 1, 2]
+            tables[entity] = pd.DataFrame({f"{entity}_id": [1, 2]})
+    tables["household"]["household_weight"] = [1.0, 1.0]
+    return USSingleYearDataset(**tables, time_period=YEAR)
+
+
+def test_rejects_chunk_holding_one_of_several_marital_units():
+    households = _spouse_and_dependent_households()
+    with pytest.raises(ValueError, match="exactly one marital unit"):
+        calculate(households, household_chunks=[0, 1])
+    # The whole file in one chunk holds both marital units and is accepted.
+    result = calculate_chunked(
+        InputSimulation,
+        ["household_weight"],
+        YEAR,
+        2,
+        dataset=households,
+        household_chunks=[0, 0],
+    )
+    assert result.weighted_sums == {"household_weight": 2.0}
+
+
+def test_non_hdf5_file_with_h5_suffix_gets_flat_file_error(tmp_path, dataset):
+    path = tmp_path / "population.h5"
+    dataset.person.to_csv(path, index=False)
+    with pytest.raises(ValueError, match="Flat-file datasets are not supported"):
+        calculate(str(path))
+
+
+def _reform_adding_an_input():
+    from policyengine_core.reforms import Reform
+    from policyengine_core.variables import Variable
+
+    from policyengine_us.entities import Person
+
+    class reform_only_input(Variable):
+        value_type = float
+        entity = Person
+        definition_period = "year"
+        label = "Input only a reform defines"
+
+    class reform_only_output(Variable):
+        value_type = float
+        entity = Person
+        definition_period = "year"
+        label = "Twice the reform-only input"
+
+        def formula(person, period, parameters):
+            return 2 * person("reform_only_input", period)
+
+    class AddInput(Reform):
+        def apply(self):
+            self.update_variable(reform_only_input)
+            self.update_variable(reform_only_output)
+
+    return AddInput
+
+
+def test_rejects_stored_names_only_the_factory_system_defines(tmp_path, dataset):
+    from policyengine_us import Microsimulation
+
+    stored = np.asarray([200.0, 400.0, 600.0, 800.0])
+    path = write_variable_centric(
+        tmp_path / "population.h5",
+        dataset,
+        extra=[("reform_only_input", str(YEAR), stored)],
+    )
+    reform = _reform_adding_an_input()
+    # Unchunked, Core loads the stored input through the reformed system; a
+    # chunk file, built against the country system, would leave it out.
+    whole = Microsimulation(dataset=str(path), reform=reform)
+    np.testing.assert_array_equal(
+        whole.calculate("reform_only_output", YEAR, use_weights=False), 2 * stored
+    )
+    whole.tax_benefit_system.simulation = None
+    del whole
+    before = _chunk_directories()
+    with pytest.raises(ValueError, match="reform_only_input are stored in the file"):
+        calculate_chunked(
+            lambda chunk: Microsimulation(dataset=chunk, reform=reform),
+            ["reform_only_output"],
+            YEAR,
+            2,
+            dataset=str(path),
+        )
+    assert _chunk_directories() == before
diff --git a/policyengine_us/tools/calculate_chunked.py b/policyengine_us/tools/calculate_chunked.py
index 842e2434e30..48c5c77b4fd 100644
--- a/policyengine_us/tools/calculate_chunked.py
+++ b/policyengine_us/tools/calculate_chunked.py
@@ -208,6 +208,9 @@ def _household_positions(dataset):
 class _EntityTables:
     """Entity-table datasets: chunks are dataset objects of the same type."""
 
+    # Chunks copy every column, so no stored name is left out.
+    unmatched = ()
+
     def __init__(self, dataset):
         self.dataset = dataset
         self.years = (
@@ -295,7 +298,19 @@ def __init__(self, path):
                 "supported; convert them to one of those layouts."
             )
         self.arrays = {}
-        with h5py.File(self.path, "r") as file:
+        # Stored names that match no country-system variable, left out of
+        # chunk files; calculate_chunked checks them against the factory's.
+        self.unmatched = []
+        try:
+            file = h5py.File(self.path, "r")
+        except OSError as error:
+            # h5py raises OSError for a file with an .h5 suffix that is not HDF5.
+            raise ValueError(
+                f"{self.path} is not an entity-table HDFStore or a "
+                "variable-centric .h5 file. Flat-file datasets are not "
+                "supported; convert them to one of those layouts."
+            ) from error
+        with file:
             names = list(file.keys())
             if not names or isinstance(file[names[0]], h5py.Dataset):
                 raise ValueError(
@@ -309,7 +324,9 @@ def __init__(self, path):
                 elif name in system.variables:
                     entity = system.variables[name].entity.key
                 else:
-                    # Core's loader ignores names that match no variable.
+                    # Core's loader ignores names that match no variable in
+                    # the simulation's system; a reform can add one.
+                    self.unmatched.append(name)
                     continue
                 if not isinstance(file[name], h5py.Group):
                     raise ValueError(f"{name} must be a group of one array per period.")
@@ -424,7 +441,9 @@ def calculate_chunked(
     type. For a variable-centric file it receives ``str`` path to a temporary
     file in the same layout, holding every stored period with its dtype under
     the source file's name, for ``Microsimulation(dataset=path)``. Names that
-    match no variable are left out, as Core's loader ignores them. Only one
+    match no country-system variable are left out, as Core's loader ignores
+    them; the helper raises if the factory's system (a reform, say) defines
+    one. Only one
     chunk file exists at once; it is deleted after its chunk, on failure too.
     Flat files are not supported.
 
@@ -439,8 +458,9 @@ def calculate_chunked(
     Input weights are never rescaled. Only one chunk simulation is live at once.
 
     This is exact for household-local calculations. Known Medicaid state
-    aggregates, the household and SPM unit income deciles, and external
-    claiming-tax-unit links are rejected. A rejected variable with a stored
+    aggregates, the household and SPM unit income deciles, external
+    claiming-tax-unit links, and a chunk holding exactly one of several
+    marital units (which pairs SSI spouses differently) are rejected. A rejected variable with a stored
     value returns that value, since no formula runs. Custom reform formulas
     must also be household-local: population reductions, ranking and random
     draws depending on population order are not supported.
@@ -481,12 +501,32 @@ def calculate_chunked(
         masks = {
             entity: assignments[owners] == chunk for entity, owners in positions.items()
         }
+        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."
+            )
         chunk_dataset = None
         simulation = None
         try:
             chunk_dataset = source.chunk(masks)
             simulation = sim_factory(chunk_dataset)
             _guard(simulation)
+            dropped = (
+                sorted(
+                    set(source.unmatched) & set(simulation.tax_benefit_system.variables)
+                )
+                if source.unmatched
+                else []
+            )
+            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."
+                )
             for name in variables:
                 entity = simulation.tax_benefit_system.get_variable(
                     name, check_existence=True
diff --git a/policyengine_us/tools/social_security_preflight.py b/policyengine_us/tools/social_security_preflight.py
index 77aee68b5c4..85fe2837608 100644
--- a/policyengine_us/tools/social_security_preflight.py
+++ b/policyengine_us/tools/social_security_preflight.py
@@ -53,6 +53,9 @@
         # The TOB leaves difference it, so a stored value freezes every tax
         # state they measure and zeroes TOB.
         "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",
         "net_investment_income_tax",
         "niit",
         "oasdi_payroll_tax_revenue",

From c545c47877c70b367a6fab2322037e453364997e Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Tue, 3 Nov 2026 00:33:33 +0000
Subject: [PATCH 12/15] Restore an unrelated doc that a broad formatter run
 rewrote

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01LwQq6GfptFq3a9AdVcFg9H
---
 docs/usage/social-security-pia-ratio.md | 20 +++++---------------
 1 file changed, 5 insertions(+), 15 deletions(-)

diff --git a/docs/usage/social-security-pia-ratio.md b/docs/usage/social-security-pia-ratio.md
index f87f730a81a..50747a01651 100644
--- a/docs/usage/social-security-pia-ratio.md
+++ b/docs/usage/social-security-pia-ratio.md
@@ -310,21 +310,11 @@ from policyengine_core.reforms import Reform
 reform = Reform.from_dict(
     {
         "gov.contrib.ssa.pia.in_effect": {"2032-01-01": True},
-        "gov.contrib.ssa.pia.progressive_price_indexing.in_effect": {
-            "2032-01-01": True
-        },
-        "gov.contrib.ssa.pia.progressive_price_indexing.start_year": {
-            "2032-01-01": 2032
-        },
-        "gov.contrib.ssa.pia.progressive_price_indexing.bend_point_fraction": {
-            "2032-01-01": 0.25
-        },
-        "gov.contrib.ssa.pia.progressive_price_indexing.maximum_aime": {
-            "2035-01-01": 17_000
-        },
-        "gov.contrib.ssa.pia.progressive_price_indexing.maximum_pia_target": {
-            "2035-01-01": 5_100
-        },
+        "gov.contrib.ssa.pia.progressive_price_indexing.in_effect": {"2032-01-01": True},
+        "gov.contrib.ssa.pia.progressive_price_indexing.start_year": {"2032-01-01": 2032},
+        "gov.contrib.ssa.pia.progressive_price_indexing.bend_point_fraction": {"2032-01-01": 0.25},
+        "gov.contrib.ssa.pia.progressive_price_indexing.maximum_aime": {"2035-01-01": 17_000},
+        "gov.contrib.ssa.pia.progressive_price_indexing.maximum_pia_target": {"2035-01-01": 5_100},
     },
     country_id="us",
 )

From 392c74393f9c54b5010ca691d897250e5703f5f3 Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Tue, 3 Nov 2026 09:54:00 +0000
Subject: [PATCH 13/15] Reject a one-marital-unit chunk before any chunk is
 built

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
---
 docs/usage/chunked-calculations.md            | 12 ++--
 .../core/test_calculate_chunked_validation.py | 69 +++++++++++++++++--
 policyengine_us/tools/calculate_chunked.py    | 32 +++++----
 3 files changed, 90 insertions(+), 23 deletions(-)

diff --git a/docs/usage/chunked-calculations.md b/docs/usage/chunked-calculations.md
index 577c60a2713..ec7f7b96d6f 100644
--- a/docs/usage/chunked-calculations.md
+++ b/docs/usage/chunked-calculations.md
@@ -83,11 +83,13 @@ whole-file chain does. A guarded variable with a stored value returns that
 value, because no formula reduces across the chunk. Known Medicaid
 claiming-tax-unit links crossing households are also rejected. The SSI spouse
 pairing also tests whether the whole simulation holds a single marital unit, so
-a chunk with exactly one marital unit can differ from the whole file; such
-chunks are rejected. Custom reform formulas remain the caller's responsibility:
-cross-household reductions, order-dependent random draws and other global
-dependencies are unsupported. These are calculation restrictions, not an
-assumption that every model variable can be chunked.
+a chunk holding exactly one of several marital units can pair spouses
+differently from the whole file; such partitions are rejected before any chunk
+is built. A file that is a single marital unit is accepted. Custom reform
+formulas remain the caller's responsibility: cross-household reductions,
+order-dependent random draws and other global dependencies are unsupported.
+These are calculation restrictions, not an assumption that every model variable
+can be chunked.
 
 The tests compare a full 32-household, five-state fixture with one, two, three
 and seven chunks, plus random household assignments. They compare every entity
diff --git a/policyengine_us/tests/core/test_calculate_chunked_validation.py b/policyengine_us/tests/core/test_calculate_chunked_validation.py
index c4af02f82b5..8ef24907f92 100644
--- a/policyengine_us/tests/core/test_calculate_chunked_validation.py
+++ b/policyengine_us/tests/core/test_calculate_chunked_validation.py
@@ -1,6 +1,7 @@
 """Reject invalid household partitions before constructing a simulation."""
 
 import re
+import tempfile
 import weakref
 from pathlib import Path
 from types import SimpleNamespace
@@ -26,6 +27,15 @@
 MARITAL_UNIT_IDS = [91, 92, 37, 38]
 
 
+@pytest.fixture(autouse=True)
+def private_tempdir(tmp_path, monkeypatch):
+    # calculate_chunked's mkdtemp and _chunk_directories then see only this
+    # test's chunk folders, not another process's in the shared temp folder.
+    directory = tmp_path / "tmp"
+    directory.mkdir()
+    monkeypatch.setattr(tempfile, "tempdir", str(directory))
+
+
 @pytest.fixture
 def dataset():
     # Two households of two people. Each person is their own marital unit, so
@@ -537,8 +547,6 @@ def variable_centric(tmp_path, dataset):
 
 
 def _chunk_directories():
-    import tempfile
-
     return set(Path(tempfile.gettempdir()).glob("policyengine-us-chunk-*"))
 
 
@@ -761,6 +769,49 @@ def test_rejects_chunk_holding_one_of_several_marital_units():
     assert result.weighted_sums == {"household_weight": 2.0}
 
 
+def test_rejects_one_marital_unit_chunk_before_building_any_chunk(dataset):
+    # Household 37's two people now share one marital unit. The first default
+    # chunk (household 91, two marital units) is valid, so the guard must run
+    # for every chunk before the first is built: never_build fails otherwise.
+    dataset.person["person_marital_unit_id"] = [91, 92, 37, 37]
+    dataset.marital_unit.drop(index=dataset.marital_unit.index[3], inplace=True)
+    with pytest.raises(ValueError, match="exactly one marital unit"):
+        calculate(dataset)
+
+
+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}
+
+
+def test_marital_unit_check_counts_only_chunks_that_hold_households(dataset):
+    # Labels need only be integers below n_chunks. Counting units per possible
+    # label (np.bincount with minlength=n_chunks) would try to allocate 8 TiB.
+    result = calculate_chunked(
+        InputSimulation,
+        ["household_weight"],
+        YEAR,
+        2**40 + 1,
+        dataset=dataset,
+        household_chunks=[2**40, 0],
+    )
+    np.testing.assert_array_equal(
+        result.entities["household"].household_weight, [1.5, 2.25]
+    )
+
+
 def test_non_hdf5_file_with_h5_suffix_gets_flat_file_error(tmp_path, dataset):
     path = tmp_path / "population.h5"
     dataset.person.to_csv(path, index=False)
@@ -813,13 +864,19 @@ def test_rejects_stored_names_only_the_factory_system_defines(tmp_path, dataset)
     np.testing.assert_array_equal(
         whole.calculate("reform_only_output", YEAR, use_weights=False), 2 * stored
     )
-    whole.tax_benefit_system.simulation = None
-    del whole
+    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(
-            lambda chunk: Microsimulation(dataset=chunk, reform=reform),
-            ["reform_only_output"],
+            factory,
+            ["household_weight"],
             YEAR,
             2,
             dataset=str(path),
diff --git a/policyengine_us/tools/calculate_chunked.py b/policyengine_us/tools/calculate_chunked.py
index 48c5c77b4fd..89ad1a448eb 100644
--- a/policyengine_us/tools/calculate_chunked.py
+++ b/policyengine_us/tools/calculate_chunked.py
@@ -458,12 +458,14 @@ def calculate_chunked(
     Input weights are never rescaled. Only one chunk simulation is live at once.
 
     This is exact for household-local calculations. Known Medicaid state
-    aggregates, the household and SPM unit income deciles, external
-    claiming-tax-unit links, and a chunk holding exactly one of several
-    marital units (which pairs SSI spouses differently) are rejected. A rejected variable with a stored
-    value returns that value, since no formula runs. Custom reform formulas
-    must also be household-local: population reductions, ranking and random
-    draws depending on population order are not supported.
+    aggregates and the household and SPM unit income deciles are rejected
+    when their formulas would run; a rejected variable with a stored value
+    returns that value, since no formula runs. External claiming-tax-unit
+    links, and a chunk holding exactly one of several marital units (which
+    can pair SSI spouses differently), are rejected before any chunk is built.
+    Custom reform formulas must also be household-local: population
+    reductions, ranking and random draws depending on population order are
+    not supported.
     """
     if isinstance(n_chunks, bool) or not isinstance(n_chunks, Integral) or n_chunks < 1:
         raise ValueError("n_chunks must be a positive integer.")
@@ -495,18 +497,24 @@ def calculate_chunked(
                 "household_chunks must assign each household an integer in [0, n_chunks)."
             )
 
+    # _ssi_spouses.py tests whether the simulation holds one marital unit.
+    # Count each nonempty chunk's units before any chunk is built; empty
+    # chunks are skipped, and labels can be any integers below n_chunks.
+    _, marital_units = np.unique(
+        assignments[positions["marital_unit"]], return_counts=True
+    )
+    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."
+        )
     tables = {}
     variable_entities = {}
     for chunk in np.unique(assignments):
         masks = {
             entity: assignments[owners] == chunk for entity, owners in positions.items()
         }
-        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."
-            )
         chunk_dataset = None
         simulation = None
         try:

From f8ad24ebd189d7f14b709b94d9c3c8b93e5773ee Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Tue, 3 Nov 2026 09:54:00 +0000
Subject: [PATCH 14/15] Cut this PR's time in CI's Microsimulation job

The job ran 54 min 57 s of its 60-minute limit at c545c47877 (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
---
 docs/usage/chunked-calculations.md            |  26 +--
 .../microsimulation/test_calculate_chunked.py | 173 ++++++++++++------
 .../microsimulation/test_tob_sequencing.py    |  52 ++++--
 3 files changed, 166 insertions(+), 85 deletions(-)

diff --git a/docs/usage/chunked-calculations.md b/docs/usage/chunked-calculations.md
index ec7f7b96d6f..4d3a1968e51 100644
--- a/docs/usage/chunked-calculations.md
+++ b/docs/usage/chunked-calculations.md
@@ -91,18 +91,20 @@ order-dependent random draws and other global dependencies are unsupported.
 These are calculation restrictions, not an assumption that every model variable
 can be chunked.
 
-The tests compare a full 32-household, five-state fixture with one, two, three
-and seven chunks, plus random household assignments. They compare every entity
-value and weighted sum, check source order and unchanged input frames, and use
-weak references to check simulation release. The fixture includes households
-with multiple tax and SPM units and Social Security recipients. The same
-households written as a variable-centric file are compared with an unchunked
-`Microsimulation` on that file for one, two, three and seven chunks and under a
-parameter reform, and with the entity-table route. Real simulations check the
-guard on direct, dependency, branch, baseline and cross-period routes, and a
-frame-depth test checks that guarding adds at most three frames to the deepest
-stack. Independent pre-change TOB branch calculations protect the statutory
-signed allocation.
+The tests compare a full 32-household, five-state fixture with one, two and
+three chunks, and with household assignments drawn by Hypothesis, which always
+include every household in one chunk and seven uneven, interleaved chunks of 1
+to 11 households. They compare every entity value and weighted sum, check
+source order and unchanged input frames, and use weak references to check
+simulation release. The fixture includes households with multiple tax and SPM
+units and Social Security recipients. The same households written as a
+variable-centric file are compared with an unchunked `Microsimulation` on that
+file for one, two and three chunks and under a parameter reform, and with the
+entity-table route, checking each chunk file's contents and its deletion. Real
+simulations check the guard on direct, dependency, branch, baseline and
+cross-period routes, and a frame-depth test checks that guarding adds at most
+three frames to the deepest stack. Independent pre-change TOB branch
+calculations protect the statutory signed allocation.
 
 ## TOB working memory
 
diff --git a/policyengine_us/tests/microsimulation/test_calculate_chunked.py b/policyengine_us/tests/microsimulation/test_calculate_chunked.py
index 56a941bfc54..6ffc10a886b 100644
--- a/policyengine_us/tests/microsimulation/test_calculate_chunked.py
+++ b/policyengine_us/tests/microsimulation/test_calculate_chunked.py
@@ -8,7 +8,8 @@
 the complete, unchanged parameter tree and history, with copy-on-write for
 reforms, while entities, variable metadata and provider receipts remain private.
 Rebuilding the same history for every equivalent partition would dominate this
-API test's runtime. The separate TOB sequencing test uses ordinary clones.
+API test's runtime. The separate TOB sequencing test lends the tree only to its
+reference, main's algorithm; the helper's branches there use ordinary clones.
 The retained default policy graph is frozen before any simulation is built,
 avoiding repeated GC scans of it. New simulation/branch cycles remain in normal
 GC generations and must still pass the weak-reference collection checks.
@@ -18,13 +19,14 @@
 import sys
 import tempfile
 import weakref
+from contextlib import contextmanager
 from pathlib import Path
 
 import h5py
 import numpy as np
 import pandas as pd
 import pytest
-from hypothesis import given, settings, strategies as st
+from hypothesis import example, given, settings, strategies as st
 
 from policyengine_us import Microsimulation
 from policyengine_us.data.dataset_schema import (
@@ -51,14 +53,29 @@
 )
 
 
-@pytest.fixture(scope="module", autouse=True)
-def lend_unchanged_parameter_history():
-    import policyengine_us.spm as spm
+@contextmanager
+def frozen_heap():
+    """Leave what is alive now out of later collections.
 
+    Collections inside the block then scan only objects created inside it,
+    which include every simulation and branch the weak-reference checks follow.
+    """
     owns_freeze = gc.get_freeze_count() == 0
     if owns_freeze:
         gc.collect()
         gc.freeze()
+    try:
+        yield
+    finally:
+        if owns_freeze:
+            gc.unfreeze()
+
+
+@contextmanager
+def lent_parameter_history():
+    """Give cloned systems the unchanged parameter tree instead of a copy."""
+    import policyengine_us.spm as spm
+
     original = spm.clone_spm_system
 
     def clone(source, *, copy_receipts=True, parameters=None):
@@ -69,15 +86,26 @@ def clone(source, *, copy_receipts=True, parameters=None):
         spm.mark_tree_shared(tree)
         return original(source, copy_receipts=copy_receipts, parameters=tree)
 
+    with pytest.MonkeyPatch.context() as patch:
+        patch.setattr(spm, "clone_spm_system", clone)
+        yield
+
+
+@pytest.fixture(scope="module", autouse=True)
+def lend_unchanged_parameter_history():
     # Scope the optimization to this module. In particular, importing fixture
     # builders from test_tob_sequencing does not apply it to that test module.
-    try:
-        with pytest.MonkeyPatch.context() as patch:
-            patch.setattr(spm, "clone_spm_system", clone)
-            yield
-    finally:
-        if owns_freeze:
-            gc.unfreeze()
+    with frozen_heap(), lent_parameter_history():
+        yield
+
+
+@pytest.fixture(autouse=True)
+def private_tempdir(tmp_path, monkeypatch):
+    # calculate_chunked's mkdtemp and _chunk_directories then see only this
+    # test's chunk folders, not another process's in the shared temp folder.
+    directory = tmp_path / "tmp"
+    directory.mkdir()
+    monkeypatch.setattr(tempfile, "tempdir", str(directory))
 
 
 def build_chunking_dataset() -> USSingleYearDataset:
@@ -244,10 +272,14 @@ def assert_matches_reference(result, dataset, reference):
         assert result.weighted_sums[variable] == weighted_sums[variable]
 
 
-@pytest.mark.parametrize("n_chunks", [1, 2, 3, 7])
-def test_chunked_real_model_is_elementwise_exact(n_chunks, chunking_dataset, reference):
-    # Weak references detect retained simulations and cycles before the next
-    # chunk, including branch-to-parent references created by TOB calculations.
+def calculate_with_checked_chunks(chunking_dataset, n_chunks, household_chunks=None):
+    """Calculate VARIABLES with a factory that checks every chunk it receives.
+
+    Each chunk's person and group tables must equal its households' source
+    rows, in source order. Weak references detect retained simulations and
+    cycles before the next chunk, including branch-to-parent references
+    created by TOB calculations. Returns the result and the chunks built.
+    """
     simulations = []
     seen = {entity: [] for entity in US_ENTITIES}
 
@@ -275,20 +307,42 @@ def factory(dataset):
         return simulation
 
     result = calculate_chunked(
-        factory, VARIABLES, YEAR, n_chunks, dataset=chunking_dataset
+        factory,
+        VARIABLES,
+        YEAR,
+        n_chunks,
+        dataset=chunking_dataset,
+        household_chunks=household_chunks,
     )
-    assert_matches_reference(result, chunking_dataset, reference)
-    assert len(simulations) == n_chunks
     assert all(simulation() is None for simulation in simulations)
     for entity, ids in seen.items():
         expected_ids = getattr(chunking_dataset, entity)[f"{entity}_id"]
         assert len(ids) == len(expected_ids)
         assert set(ids) == set(expected_ids)
+    return result, len(simulations)
+
+
+# Three chunks catch order errors that two cannot; the smallest of the three
+# default chunks still holds 10 households. The partition property below runs
+# the same chunk checks on smaller, uneven and interleaved chunks.
+@pytest.mark.parametrize("n_chunks", [1, 2, 3])
+def test_chunked_real_model_is_elementwise_exact(n_chunks, chunking_dataset, reference):
+    result, built = calculate_with_checked_chunks(chunking_dataset, n_chunks)
+    assert_matches_reference(result, chunking_dataset, reference)
+    assert built == n_chunks
+
+
+MAX_PARTITION_CHUNKS = 7
+# Seven chunks of 1 to 11 households. 7 is coprime with 32, so multiplying row
+# numbers by it permutes the households and interleaves the chunks.
+UNEVEN_LABELS = np.repeat(np.arange(MAX_PARTITION_CHUNKS), [1, 2, 3, 4, 5, 6, 11])[
+    np.arange(HOUSEHOLDS) * 7 % HOUSEHOLDS
+]
 
 
 @st.composite
 def household_partitions(draw):
-    n_chunks = draw(st.integers(min_value=1, max_value=7))
+    n_chunks = draw(st.integers(min_value=1, max_value=MAX_PARTITION_CHUNKS))
     # Every partition is nonempty, while all remaining household assignments
     # vary independently. Chunk labels are intentionally not contiguous rows.
     tail = draw(
@@ -303,19 +357,23 @@ def household_partitions(draw):
     return n_chunks, np.roll(labels, offset)
 
 
-@settings(max_examples=4, deadline=None)
+# Hypothesis generates the simplest partition, one chunk labelled 0, first. In
+# CI its ci profile derandomizes, so the second generated example is one fixed
+# pseudo-random partition; elsewhere it is random. The strategy uses every
+# label it draws, so it never leaves a chunk empty. The explicit examples pin
+# both ends of its range: every household in one chunk, labelled 3 so that the
+# other six labels are empty and skipped; and seven uneven, interleaved chunks.
+@example(partition=(MAX_PARTITION_CHUNKS, np.full(HOUSEHOLDS, 3)))
+@example(partition=(MAX_PARTITION_CHUNKS, UNEVEN_LABELS))
+@settings(max_examples=2, deadline=None)
 @given(partition=household_partitions())
 def test_random_household_partitions_are_exact(partition, chunking_dataset, reference):
     n_chunks, household_chunks = partition
-    result = calculate_chunked(
-        make_simulation,
-        VARIABLES,
-        YEAR,
-        n_chunks,
-        dataset=chunking_dataset,
-        household_chunks=household_chunks,
+    result, built = calculate_with_checked_chunks(
+        chunking_dataset, n_chunks, household_chunks
     )
     assert_matches_reference(result, chunking_dataset, reference)
+    assert built == len(np.unique(household_chunks))
 
 
 def test_multi_year_dataset_preserves_every_year_and_original_input(
@@ -436,7 +494,7 @@ def test_single_year_factory_automatically_projects_future_inputs(chunking_datas
     _assert_snapshot_matches(result, chunking_dataset, expected)
 
 
-def test_reform_factory_is_exact_and_releases_its_baseline(chunking_dataset):
+def test_reform_factory_is_exact_and_releases_its_baseline(chunking_dataset, reference):
     from policyengine_core.reforms import Reform
     from policyengine_us.spm import clone_spm_system
 
@@ -464,11 +522,9 @@ def build(dataset):
             reform=reform,
         )
 
-    baseline = make_simulation(chunking_dataset)
-    baseline_tax = np.asarray(baseline.calculate("income_tax", YEAR)).copy()
-    _release_reference_simulation(baseline)
-    del baseline
-    gc.collect()
+    # The module reference calculated unreformed income_tax on this dataset
+    # with make_simulation, as a separate baseline simulation would.
+    baseline_tax = reference[0]["income_tax"]
 
     simulation = build(chunking_dataset)
     expected = _snapshot_calculations(simulation, variables, YEAR)
@@ -580,25 +636,17 @@ def _parameter_reform():
 
 @pytest.fixture(scope="module")
 def variable_centric_reference(variable_centric_path, reference):
-    expected = {}
-    for reform in (None, _parameter_reform()):
-        simulation = Microsimulation(dataset=variable_centric_path, reform=reform)
-        expected[reform is not None] = _snapshot_calculations(
-            simulation, TOB_VARIABLES, YEAR
-        )
-        _release_reference_simulation(simulation)
-        del simulation
-        gc.collect()
+    simulation = Microsimulation(dataset=variable_centric_path)
+    expected = _snapshot_calculations(simulation, TOB_VARIABLES, YEAR)
+    _release_reference_simulation(simulation)
+    del simulation
+    gc.collect()
     # Both loaders read the same households to the same results.
     values, weighted_sums, weights = reference
-    for variable, (entity, actual, actual_weights, total) in expected[False].items():
+    for variable, (entity, actual, actual_weights, total) in expected.items():
         np.testing.assert_array_equal(actual, values[variable])
         np.testing.assert_array_equal(actual_weights, weights[entity])
         assert total == weighted_sums[variable]
-    assert any(
-        np.any(expected[True][name][1] != expected[False][name][1])
-        for name in TOB_VARIABLES
-    ), "The reform must change TOB on the fixture"
     return expected
 
 
@@ -606,7 +654,7 @@ def _chunk_directories():
     return set(Path(tempfile.gettempdir()).glob("policyengine-us-chunk-*"))
 
 
-@pytest.mark.parametrize("n_chunks", [1, 2, 3, 7])
+@pytest.mark.parametrize("n_chunks", [1, 2, 3])
 def test_variable_centric_file_is_elementwise_exact(
     n_chunks, chunking_dataset, variable_centric_path, variable_centric_reference
 ):
@@ -631,9 +679,7 @@ def factory(path):
     result = calculate_chunked(
         factory, TOB_VARIABLES, YEAR, n_chunks, dataset=variable_centric_path
     )
-    _assert_snapshot_matches(
-        result, chunking_dataset, variable_centric_reference[False]
-    )
+    _assert_snapshot_matches(result, chunking_dataset, variable_centric_reference)
     assert len(paths) == n_chunks
     assert not any(Path(path).exists() for path in paths)
     assert _chunk_directories() == before
@@ -642,15 +688,32 @@ def factory(path):
 def test_variable_centric_file_is_exact_under_reform(
     chunking_dataset, variable_centric_path, variable_centric_reference
 ):
-    reform = _parameter_reform()
+    simulation = Microsimulation(
+        dataset=variable_centric_path, reform=_parameter_reform()
+    )
+    expected = _snapshot_calculations(simulation, TOB_VARIABLES, YEAR)
+    assert any(
+        np.any(expected[name][1] != variable_centric_reference[name][1])
+        for name in TOB_VARIABLES
+    ), "The reform must change TOB on the fixture"
+    # The chunks share this run's reformed system instead of each loading the
+    # country system again with the reform. The reform-factory test above
+    # still builds a reformed simulation, with its baseline, for every chunk.
+    policy = simulation.tax_benefit_system
+    _release_reference_simulation(simulation)
+    del simulation
+    gc.collect()
     result = calculate_chunked(
-        lambda path: Microsimulation(dataset=path, reform=reform),
+        lambda path: Microsimulation(dataset=path, tax_benefit_system=policy),
         TOB_VARIABLES,
         YEAR,
         2,
         dataset=variable_centric_path,
     )
-    _assert_snapshot_matches(result, chunking_dataset, variable_centric_reference[True])
+    # Collect the reformed parameter tree here, not in the next test.
+    del policy
+    gc.collect()
+    _assert_snapshot_matches(result, chunking_dataset, expected)
 
 
 def _deepest_stack(calculation):
diff --git a/policyengine_us/tests/microsimulation/test_tob_sequencing.py b/policyengine_us/tests/microsimulation/test_tob_sequencing.py
index e30aab150d9..680220363e2 100644
--- a/policyengine_us/tests/microsimulation/test_tob_sequencing.py
+++ b/policyengine_us/tests/microsimulation/test_tob_sequencing.py
@@ -4,15 +4,27 @@
 import weakref
 
 import numpy as np
+import pytest
 
 from policyengine_us.tests.microsimulation.test_calculate_chunked import (
     YEAR,
     _release_reference_simulation,
     build_chunking_dataset,
+    frozen_heap,
+    lent_parameter_history,
     make_simulation,
 )
 
 
+@pytest.fixture(autouse=True)
+def frozen_unchanged_heap():
+    # CI runs this file last of the microsimulation folder's base group, in
+    # one pytest process. Collections then scan only what this test creates,
+    # including every simulation and branch the weak-reference checks follow.
+    with frozen_heap():
+        yield
+
+
 def _main_counterfactual(simulation, name, taxable_social_security=None):
     """Reproduce main's branch algorithm independently of the new helper."""
     branch = simulation.get_branch(name, clone_system=True)
@@ -42,20 +54,28 @@ def _main_counterfactual(simulation, name, taxable_social_security=None):
 
 def test_sequential_tob_matches_main_elementwise_and_reuses_states(monkeypatch):
     dataset = build_chunking_dataset()
-    baseline = make_simulation(dataset)
-    # Main caps old-law inclusion at actual inclusion, read before any branch.
-    old_law_inclusion = np.minimum(
-        np.asarray(
-            baseline.calculate("tax_unit_taxable_social_security_pre_1993", YEAR)
-        ),
-        np.asarray(baseline.calculate("tax_unit_taxable_social_security", YEAR)),
-    ).copy()
-    # These are main's OASDI and total algorithms, including separate no-SS
-    # branches and the signed differences rather than per-tax-unit floors.
-    oasdi_no_ss = _main_counterfactual(baseline, "oasdi_no_ss")
-    old_law = _main_counterfactual(baseline, "oasdi_old_law", old_law_inclusion)
-    current = np.asarray(baseline.calculate("income_tax_chapter_1", YEAR)).copy()
-    total_no_ss = _main_counterfactual(baseline, "tob_calc")
+    # Main's algorithm is the reference. Its branches change no parameter, so
+    # they lend the unchanged tree instead of copying it; the helper's branches
+    # below, which this test checks, keep ordinary clones.
+    with lent_parameter_history():
+        baseline = make_simulation(dataset)
+        # Main caps old-law inclusion at actual inclusion, read before any branch.
+        old_law_inclusion = np.minimum(
+            np.asarray(
+                baseline.calculate("tax_unit_taxable_social_security_pre_1993", YEAR)
+            ),
+            np.asarray(baseline.calculate("tax_unit_taxable_social_security", YEAR)),
+        ).copy()
+        # These are main's OASDI and total algorithms, including separate no-SS
+        # branches and the signed differences rather than per-tax-unit floors.
+        oasdi_no_ss = _main_counterfactual(baseline, "oasdi_no_ss")
+        old_law = _main_counterfactual(baseline, "oasdi_old_law", old_law_inclusion)
+        current = np.asarray(baseline.calculate("income_tax_chapter_1", YEAR)).copy()
+        total_no_ss = _main_counterfactual(baseline, "tob_calc")
+        weights = np.asarray(baseline.calculate("tax_unit_weight", YEAR)).copy()
+        _release_reference_simulation(baseline)
+        del baseline
+        gc.collect()
     expected = {
         "tob_revenue_oasdi": old_law - oasdi_no_ss,
         "tob_revenue_total": current - total_no_ss,
@@ -63,10 +83,6 @@ def test_sequential_tob_matches_main_elementwise_and_reuses_states(monkeypatch):
     expected["tob_revenue_medicare_hi"] = (
         expected["tob_revenue_total"] - expected["tob_revenue_oasdi"]
     )
-    weights = np.asarray(baseline.calculate("tax_unit_weight", YEAR)).copy()
-    _release_reference_simulation(baseline)
-    del baseline
-    gc.collect()
 
     simulation = make_simulation(dataset)
     simulation_type = type(simulation)

From 3f2700051680816df5190e84e39f2aac9c6c0c73 Mon Sep 17 00:00:00 2001
From: Max Ghenis <mghenis@gmail.com>
Date: Tue, 3 Nov 2026 06:22:04 -0500
Subject: [PATCH 15/15] Narrow the chunk-file docs claim and the CI order
 comment

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>
---
 docs/usage/chunked-calculations.md                  | 13 ++++++++-----
 .../tests/microsimulation/test_tob_sequencing.py    |  7 ++++---
 2 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/docs/usage/chunked-calculations.md b/docs/usage/chunked-calculations.md
index 4d3a1968e51..d44c5cc0f35 100644
--- a/docs/usage/chunked-calculations.md
+++ b/docs/usage/chunked-calculations.md
@@ -97,11 +97,14 @@ include every household in one chunk and seven uneven, interleaved chunks of 1
 to 11 households. They compare every entity value and weighted sum, check
 source order and unchanged input frames, and use weak references to check
 simulation release. The fixture includes households with multiple tax and SPM
-units and Social Security recipients. The same households written as a
-variable-centric file are compared with an unchunked `Microsimulation` on that
-file for one, two and three chunks and under a parameter reform, and with the
-entity-table route, checking each chunk file's contents and its deletion. Real
-simulations check the guard on direct, dependency, branch, baseline and
+units and Social Security recipients. The tests also write the same households
+as a variable-centric file and compare the chunked results with an unchunked
+`Microsimulation` on that file, for one, two and three chunks and under a
+parameter reform, and with the entity-table route. For each chunk file, they
+check its name and the dtype of its role array, and that no other chunk file
+exists beside it. They also check that the helper leaves an unknown variable
+name out of the file and deletes the file afterwards. Real simulations check
+the guard on direct, dependency, branch, baseline and
 cross-period routes, and a frame-depth test checks that guarding adds at most
 three frames to the deepest stack. Independent pre-change TOB branch
 calculations protect the statutory signed allocation.
diff --git a/policyengine_us/tests/microsimulation/test_tob_sequencing.py b/policyengine_us/tests/microsimulation/test_tob_sequencing.py
index 680220363e2..2243261c3ee 100644
--- a/policyengine_us/tests/microsimulation/test_tob_sequencing.py
+++ b/policyengine_us/tests/microsimulation/test_tob_sequencing.py
@@ -18,9 +18,10 @@
 
 @pytest.fixture(autouse=True)
 def frozen_unchanged_heap():
-    # CI runs this file last of the microsimulation folder's base group, in
-    # one pytest process. Collections then scan only what this test creates,
-    # including every simulation and branch the weak-reference checks follow.
+    # In CI this file shares the microsimulation base group's single pytest
+    # process, and it currently sorts last. With the heap frozen, collections
+    # scan only what this test creates, including every simulation and branch
+    # the weak-reference checks follow.
     with frozen_heap():
         yield