diff --git a/changelog.d/chunked-microsimulation-calculate.added.md b/changelog.d/chunked-microsimulation-calculate.added.md new file mode 100644 index 00000000000..ebc58caf6f6 --- /dev/null +++ b/changelog.d/chunked-microsimulation-calculate.added.md @@ -0,0 +1 @@ +Add exact household-chunked microsimulation calculations, including variable-centric population files, and share sequential TOB counterfactual tax calculations. diff --git a/docs/_toc.yml b/docs/_toc.yml index 32f129bcb7f..2e6f404dc49 100644 --- a/docs/_toc.yml +++ b/docs/_toc.yml @@ -65,6 +65,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..d44c5cc0f35 --- /dev/null +++ b/docs/usage/chunked-calculations.md @@ -0,0 +1,133 @@ +# 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`, 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 (`/`), 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 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__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 `_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. The name +`weight` is reserved for the output weight column. + +## 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 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 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 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 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. + +## TOB working memory + +The TOB outputs difference `income_tax_chapter_1`, chapter 1 liability after +nonrefundable and before refundable credits, which leaves out the chapter 2A +net investment income tax. Counterfactual chapter 1 tax without Social +Security, and under bounded pre-1993 inclusion, is cached in parent variable +holders. The bound is the lesser of the pre-1993 section 86 amount and actual +inclusion, read in the outer simulation before either counterfactual branch +exists. A direct TOB calculation prepares those holders in separate outer +calculations, old-law first: 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 +chapter 1 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 bounded +old-law tax minus no-inclusion tax, total remains current chapter 1 tax minus +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 PR records measured full-dataset peak memory and timing; small-fixture +exactness establishes correctness, not an RSS bound. diff --git a/policyengine_us/system.py b/policyengine_us/system.py index 65a99683f52..10639e64381 100644 --- a/policyengine_us/system.py +++ b/policyengine_us/system.py @@ -38,6 +38,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, @@ -313,7 +314,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. @@ -480,7 +481,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/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 new file mode 100644 index 00000000000..8ef24907f92 --- /dev/null +++ b/policyengine_us/tests/core/test_calculate_chunked_validation.py @@ -0,0 +1,884 @@ +"""Reject invalid household partitions before constructing a simulation.""" + +import re +import tempfile +import weakref +from pathlib import Path +from types import SimpleNamespace + +import h5py +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 ( + _POPULATION_AGGREGATES, + _guard, + calculate_chunked, +) + +YEAR = 2026 +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 + # 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, 504, 101, 102]}) + for entity in US_ENTITIES: + 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 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) + + +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) + + +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"): + 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 + 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, 37, 91, 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, 37, 91, 91] + multi = USMultiYearDataset(datasets=[dataset, other]) + with pytest.raises(ValueError, match="memberships must be identical"): + 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. + + 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 + 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 _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": + 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", 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" + with pytest.raises(ValueError, match="aggregates across households"): + calculate_chunked( + lambda chunk: AggregateSimulation(chunk, aggregate, route), + [variable], + YEAR, + 2, + dataset=dataset, + ) + + +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]: + # The country simulation applies a reform more than once. + self.update_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 + 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, 0.0, 25_000.0, 0.0] + # Core uses projected household weights; native input columns still survive. + dataset.person["person_weight"] = [999.0, 998.0, 888.0, 887.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, 1.5, 2.25, 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, 998.0, 888.0, 887.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 + + +# Variable-centric files: / 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"] * 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, 3.0, 4.0], dtype=np.float32), + ), + ] + return write_variable_centric(tmp_path / "population_2026.h5", dataset, extra=extra) + + +def _chunk_directories(): + 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 + # 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"] * 4) + + +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 4"), + ("household_weight", [1.0], r"household_weight/2025.*household has 2"), + ("person_tax_unit_id", [91] * 4, "tax_unit spans multiple households"), + ("household_id", [91, 91], "household_id must be unique"), + ("person_household_id", [91, 91, 999, 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 +): + 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)) + + +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, 37, 91, 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="/ layout"): + calculate(str(arrays)) + + +# 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|\.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")): + 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 + or relative in _CHUNK_SAFE_POPULATION_READERS + ): + continue + unguarded.append(relative) + assert not unguarded, ( + "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_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) + 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 + ) + reformed = whole.tax_benefit_system + + def factory(chunk): + # Only the chunk system matters; reuse the unchunked reformed one. + simulation = FileSimulation(chunk) + simulation.tax_benefit_system = reformed + return simulation + + before = _chunk_directories() + with pytest.raises(ValueError, match="reform_only_input are stored in the file"): + calculate_chunked( + factory, + ["household_weight"], + YEAR, + 2, + dataset=str(path), + ) + assert _chunk_directories() == before 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..f798d518ba6 --- /dev/null +++ b/policyengine_us/tests/core/test_tob_branch_lifecycle.py @@ -0,0 +1,244 @@ +"""Counterfactual branches retain only their returned chapter 1 tax array.""" + +from types import FunctionType, 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_chapter_1", 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, [OLD_LAW, NO_SS]), (HI, [OLD_LAW, NO_SS]), (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], [OLD_LAW, NO_SS]), ([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}), + ] + + +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 new file mode 100644 index 00000000000..6ffc10a886b --- /dev/null +++ b/policyengine_us/tests/microsimulation/test_calculate_chunked.py @@ -0,0 +1,771 @@ +"""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. + +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 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. +""" + +import gc +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 example, 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 _guard, 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", +) + + +@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): + 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) + + 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. + 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: + """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, 48, 12, 53, 32) + 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: + # 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, + "itemized_medical_expenses": 0.0, + "salt_deduction": 0.0, + } + ) + 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. + records["spm_unit"].append( + { + "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 + ), + # 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]: + 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 + ), + # 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 + 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 + _release_reference_simulation(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] + + +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} + + 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, + household_chunks=household_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=MAX_PARTITION_CHUNKS)) + # 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) + + +# 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, 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( + 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): + """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): + 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, reference): + 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 + # 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): + # The country clone preserves instance-level variable metadata and + # provider state while avoiding repeated country/parameter loading. + return Microsimulation( + dataset=USMultiYearDataset(datasets=[dataset]), + tax_benefit_system=clone_spm_system(system), + reform=reform, + ) + + # 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) + 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. + constructor_options = {} if multi_year else {"dataset_end_year": YEAR} + simulation = Microsimulation( + dataset=source if multi_year else str(path), **constructor_options + ) + expected = _snapshot_calculations(simulation, variables, YEAR) + _release_reference_simulation(simulation) + del simulation + gc.collect() + + result = calculate_chunked( + lambda dataset: Microsimulation(dataset=dataset, **constructor_options), + 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) + + +# Population projection files store / 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): + 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.items(): + np.testing.assert_array_equal(actual, values[variable]) + np.testing.assert_array_equal(actual_weights, weights[entity]) + assert total == weighted_sums[variable] + return expected + + +def _chunk_directories(): + return set(Path(tempfile.gettempdir()).glob("policyengine-us-chunk-*")) + + +@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 +): + 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) + 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 +): + 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, tax_benefit_system=policy), + TOB_VARIABLES, + YEAR, + 2, + dataset=variable_centric_path, + ) + # 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): + 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/tests/microsimulation/test_tob_sequencing.py b/policyengine_us/tests/microsimulation/test_tob_sequencing.py new file mode 100644 index 00000000000..2243261c3ee --- /dev/null +++ b/policyengine_us/tests/microsimulation/test_tob_sequencing.py @@ -0,0 +1,119 @@ +"""Sequential TOB counterfactuals equal the original signed main formulas.""" + +import gc +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(): + # 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 + + +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_chapter_1", YEAR)).copy() + finally: + simulation.branches.pop(name) + _release_reference_simulation(branch) + del branch + gc.collect() + + +def test_sequential_tob_matches_main_elementwise_and_reuses_states(monkeypatch): + dataset = build_chunking_dataset() + # 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, + } + expected["tob_revenue_medicare_hi"] = ( + expected["tob_revenue_total"] - expected["tob_revenue_oasdi"] + ) + + 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_pre_1993_social_security", "_tob_no_social_security"] + assert not any(name.startswith("_tob_") for name in simulation.branches) + _release_reference_simulation(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..89ad1a448eb --- /dev/null +++ b/policyengine_us/tools/calculate_chunked.py @@ -0,0 +1,601 @@ +"""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 shutil +import tempfile +from types import SimpleNamespace + +import numpy as np +import pandas as pd + +from policyengine_us.data.dataset_schema import ( + US_ENTITIES, + USMultiYearDataset, + USSingleYearDataset, +) + + +# 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", + "household_income_decile", + "spm_unit_income_decile", + ) +) + + +class _HouseholdLocalCalculations: + """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: " + "it aggregates across households. Calculate it on the full dataset." + ) + return super()._check_period_consistency(period, variable) + + +@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) + + +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.""" + + entities: dict[str, pd.DataFrame] + weighted_sums: dict[str, float] + + +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 ( + _is_hdfstore_format, + _resolve_dataset_path, + ) + + path = _resolve_dataset_path(str(dataset)) + 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()) + 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 an entity-table HDFStore or variable-centric h5 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 + + +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 = ( + 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 ``/`` 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 = {} + # 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( + f"{self.path} does not use the variable-centric " + "/ 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 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.") + 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, + 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, + entity-table HDFStore paths/URLs, and variable-centric h5 paths/URLs: files + of ``/`` 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 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. + + 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, + ``_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 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.") + 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) + source = ( + dataset if isinstance(dataset, _VariableCentricFile) else _EntityTables(dataset) + ) + positions = source.positions + if household_chunks is None: + 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))) + ): + assignments[indices] = chunk + else: + assignments = np.asarray(household_chunks) + if ( + assignments.shape != (source.household_count,) + 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)." + ) + + # _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() + } + 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 + ).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 = source.ids(entity) + 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: + _detach_simulation_links(simulation) + del simulation + 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() + } + return ChunkedCalculation(tables, sums) 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", diff --git a/policyengine_us/tools/tob.py b/policyengine_us/tools/tob.py new file mode 100644 index 00000000000..614339ee124 --- /dev/null +++ b/policyengine_us/tools/tob.py @@ -0,0 +1,118 @@ +"""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 + + 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 ( + "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 = [] + # Prepare old-law first: its bound reads actual inclusion in + # this simulation before either counterfactual branch exists. + if needs_oasdi: + states.append("income_tax_with_pre_1993_social_security") + if needs_total or needs_oasdi: + states.append("income_tax_without_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 chapter 1 tax, then release its simulation. + + The measure is ``income_tax_chapter_1``: liability after nonrefundable and + before refundable credits, without the chapter 2A net investment income + tax. ``None`` neutralizes taxable Social Security; an array supplies an + alternative inclusion, which the caller has already bounded. 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_chapter_1", 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..88ea48bdf2b --- /dev/null +++ b/policyengine_us/variables/gov/ssa/revenue/income_tax_with_pre_1993_social_security.py @@ -0,0 +1,33 @@ +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 = "Chapter 1 income tax with bounded pre-1993 taxable Social Security" + unit = USD + documentation = ( + "Cached chapter 1 income-tax counterfactual (after nonrefundable and " + "before refundable credits, excluding chapter 2A net investment " + "income tax) using the lesser of the pre-1993 section 86 includible " + "amount and actual inclusion, under current income-tax rules. The " + "counterfactual branch is released after calculating the tax." + ) + reference = ( + "https://www.law.cornell.edu/uscode/text/42/401", + "https://www.ssa.gov/OP_Home/comp2/F098-021.html", + ) + + def formula(tax_unit, period, parameters): + # Read actual inclusion here, before any branch neutralizes it. + taxable_social_security = min_( + tax_unit("tax_unit_taxable_social_security_pre_1993", period), + tax_unit("tax_unit_taxable_social_security", 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..cbdc6b17b58 --- /dev/null +++ b/policyengine_us/variables/gov/ssa/revenue/income_tax_without_social_security.py @@ -0,0 +1,20 @@ +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 = "Chapter 1 income tax without taxable Social Security" + unit = USD + documentation = ( + "Cached chapter 1 income-tax counterfactual (after nonrefundable and " + "before refundable credits, excluding chapter 2A net investment " + "income tax) shared by the taxation-of-benefits revenue variables. " + "The counterfactual branch is released after calculating the tax." + ) + reference = "https://www.law.cornell.edu/uscode/text/42/401" + + 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 db8df732912..4798a41038c 100644 --- a/policyengine_us/variables/gov/ssa/revenue/tob_revenue_oasdi.py +++ b/policyengine_us/variables/gov/ssa/revenue/tob_revenue_oasdi.py @@ -38,55 +38,14 @@ def formula(tax_unit, period, parameters): nothing, so OASDI takes the whole increase and HI none. Method: - 1. Calculate tax with NO taxable SS - 2. Calculate tax with the lesser of the pre-1993 section 86 includible - amount and actual inclusion + 1. Calculate chapter 1 tax with NO taxable SS + 2. Calculate chapter 1 tax with the lesser of the pre-1993 section 86 + includible amount and actual inclusion 3. Signed OASDI revenue = (2) - (1), without a per-unit floor """ - sim = tax_unit.simulation - # Read actual inclusion here, before any branch neutralizes it. - taxable_ss = tax_unit("tax_unit_taxable_social_security", period) - old_law_taxable_ss = min_( - tax_unit("tax_unit_taxable_social_security_pre_1993", period), - taxable_ss, - ) - - # === 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_chapter_1", period) - del sim.branches["oasdi_no_ss"] - - # === Branch 2: Tax with the capped 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_chapter_1", period) - del sim.branches["oasdi_old_law"] + # Each state is cached on the parent after its branch is released. + # The capped old-law state reads actual inclusion in this simulation. + tax_old_law_ss = tax_unit("income_tax_with_pre_1993_social_security", period) + tax_no_ss = tax_unit("income_tax_without_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 cc85d335016..784a510d5f8 100644 --- a/policyengine_us/variables/gov/ssa/revenue/tob_revenue_total.py +++ b/policyengine_us/variables/gov/ssa/revenue/tob_revenue_total.py @@ -24,33 +24,12 @@ def formula(tax_unit, period, parameters): """ Calculate total trust fund revenue using branching + neutralization. - This isolates the tax impact of including taxable SS in AGI by - comparing tax with and without taxable SS. + This isolates the chapter 1 tax impact of including taxable SS in AGI + by comparing chapter 1 liability with and without taxable SS. """ - sim = tax_unit.simulation - - # Calculate chapter 1 liability WITH taxable SS + # Chapter 1 liability WITH taxable SS, in the current state income_tax_with_ss = tax_unit("income_tax_chapter_1", 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 chapter 1 liability without taxable SS - income_tax_without_ss = branch.tax_unit("income_tax_chapter_1", period) - - # Clean up branch - del sim.branches["tob_calc"] + # Cached on the parent after its branch is released + income_tax_without_ss = tax_unit("income_tax_without_social_security", period) return income_tax_with_ss - income_tax_without_ss