subsample reloads calculated values as inputs, so a reform's baseline reads the reform's results #633

Problem

Simulation.subsample reloads every value the simulation holds as an input, calculated values included. After it:

1. A reform simulation's baseline reads values calculated under the reform. The rebuilt baseline branch starts from the rebuilt simulation's stored values, so any variable calculated before subsample is a baseline input holding the reform's result. Reform-vs-baseline comparisons of those variables, and of everything computed from them, show no change.
2. Input presence is lost. Holder.get_input_periods counts the reloaded values as inputs, so a calculated default looks like a supplied column. policyengine-us PR #10447 decides behaviour from whether a dataset supplies unemployment_compensation_start_month, and has to drop that variable's calculated defaults before calling subsample.

Where (3.32.27)

• simulations/simulation.py:2330: df = self.to_input_dataframe(include_computed_variables=True). With that flag, _get_exportable_input_periods returns holder.get_known_periods() (:2213-2214), which includes calculated periods, and to_input_dataframe exports each (:2254-2268).
• :2408-2410: the subsampled frame becomes the dataset, _user_input_keys is reset, and build_from_dataset() sets every exported column as an input.
• :2430-2444: the baseline branch is rebuilt with get_branch("baseline"), which starts from the values the rebuilt simulation holds (get_branch docstring, :1961-1966), then is rebound to the baseline system.

The comment at :2328-2329 gives the reason for the flag: keeping formula-backed structural variables such as ids (#498).

Reproduction (policyengine-core 3.32.27, policyengine-us 2.63.3)

import numpy as np
import pandas as pd
from policyengine_core.reforms import Reform
from policyengine_us import Microsimulation
from policyengine_us.data.dataset_schema import USSingleYearDataset

ids = np.arange(1, 41)
groups = ("household", "tax_unit", "spm_unit", "family", "marital_unit")
dataset = USSingleYearDataset(
    person=pd.DataFrame(
        {"person_id": ids, "age": np.full(ids.size, 40),
         "employment_income": np.linspace(0, 200_000, ids.size)}
        | {f"person_{group}_id": ids for group in groups}
    ),
    household=pd.DataFrame(
        {"household_id": ids, "state_fips": np.full(ids.size, 48),
         "household_weight": np.ones(ids.size)}
    ),
    **{group: pd.DataFrame({f"{group}_id": ids}) for group in groups[1:]},
    time_period=2024,
)
reform = Reform.from_dict(
    {f"gov.irs.deductions.standard.amount.{status}": {"2024-01-01.2100-12-31": 99_000}
     for status in ("SINGLE", "JOINT", "HEAD_OF_HOUSEHOLD", "SEPARATE", "SURVIVING_SPOUSE")},
    country_id="us",
)
for calculate_first in (False, True):
    sim = Microsimulation(dataset=dataset, reform=reform)
    if calculate_first:
        sim.calculate("standard_deduction", 2024)
    sim.subsample(n=20, seed="repro", quantize_weights=False)
    print(
        f"calculated first: {calculate_first}; "
        f"reform {sim.calculate('standard_deduction', 2024).values.mean():,.0f}; "
        f"baseline {sim.baseline.calculate('standard_deduction', 2024).values.mean():,.0f}"
    )

Output:

calculated first: False; reform 99,000; baseline 14,600
calculated first: True; reform 99,000; baseline 99,000

Suggested fix

Export inputs, not every calculated value: the set_input record (include_computed_variables=False) plus the structural variables a rebuild needs (ids, entity links and weights, formula-backed or not, per #498). Calculated values are then recalculated on the subsample. Regression tests:

• the reproduction above: the baseline result after subsample is the same whether or not the reform result was calculated first;
• get_input_periods of a variable that was only calculated is empty after subsample, and a variable set with set_input keeps its periods.

Callers that calculate before subsampling are affected; those that subsample straight after building are not. I have not checked which callers do which.

Related: #605 (the input record after build_from_populations), #632 (a neutralized input lost from the baseline), #522 (the baseline branch's system after subsample).

🤖 Generated with Claude Code