diff --git a/changelog.d/fix-subsample-inputs-only.fixed.md b/changelog.d/fix-subsample-inputs-only.fixed.md new file mode 100644 index 000000000..61ade9b4b --- /dev/null +++ b/changelog.d/fix-subsample-inputs-only.fixed.md @@ -0,0 +1 @@ +Rebuild subsamples from current source inputs and loaded memberships, preserve each input weight column's entity total, select with available requested-period or dataset-period weights, and discard pre-sample calculation branches, including caller-created branches. diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index e4628e244..ddda2ab94 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -1150,12 +1150,12 @@ def get_eternity_array(name): person_membership_ids = get_eternity_array(person_membership_id_field) person_role_field = f"{person_entity.key}_{group_entity.key}_role" - if person_role_field in data: + if any(column.split("__")[0] == person_role_field for column in data): person_roles = get_eternity_array(person_role_field) elif "role" in data: person_roles = get_eternity_array("role") elif self.default_role is not None: - person_roles = np.full(len(entity_ids), self.default_role) + person_roles = np.full(self.persons.count, self.default_role) else: raise ValueError( f"Missing {person_role_field} column in the dataset. Each group entity must have a person role array defined for ETERNITY." @@ -3161,6 +3161,15 @@ def _get_exportable_input_periods( if not self._is_exportable_input_variable(variable_name): return [] + return self._get_set_input_periods(variable_name) + + def _get_set_input_periods(self, variable_name: str) -> List[Period]: + """Periods of ``variable_name`` whose stored value the simulation was given. + + Use the holder's per-tier input record, including for variables with + formulas. Storage's carry-over marks also include cache writes, which + must not become source inputs when a simulation is rebuilt. + """ user_input_periods = { period for input_variable_name, branch_name, period in getattr( @@ -3168,15 +3177,29 @@ def _get_exportable_input_periods( ) if input_variable_name == variable_name and branch_name in self._get_visible_branch_names() + and self.get_holder(variable_name)._stores_user_input(period, branch_name) } - if not user_input_periods: - return [] - variable = self.tax_benefit_system.get_variable(variable_name) + return sorted(user_input_periods, key=str) + + def _get_recorded_input_array(self, variable_name: str, period: Period): + """Read a surviving supplied value from the tier that still records it. + + Match input replay's branch and tier precedence. A cache in memory + can shadow a recorded disk input, so ``calculate`` is not an input + export operation even when the period remains in the input record. + """ holder = self.get_holder(variable_name) - if variable.definition_period == ETERNITY: - return holder.get_known_periods() - known_periods = set(holder.get_known_periods()) - return sorted(user_input_periods & known_periods, key=str) + input_keys = getattr(self, "_user_input_keys", set()) + for branch_name in self._get_visible_branch_names(): + if (variable_name, branch_name, period) not in input_keys: + continue + for storage_name, storage in ( + ("memory", holder._memory_storage), + ("disk", holder._disk_storage), + ): + if holder._stores_user_input(period, branch_name, storage_name): + return storage.get(period, branch_name) + return None def to_input_dataframe( self, @@ -3185,8 +3208,9 @@ def to_input_dataframe( """Exports a DataFrame that can be loaded back into a new Simulation. By default, only structurally input variables populated through - ``set_input`` are exported. This avoids serializing pseudo-inputs and - stale calculated values that would override formulas when reloaded. + ``set_input`` are exported from their recorded storage tiers. This + avoids serializing pseudo-inputs and stale calculated values that + would override formulas when reloaded. Args: include_computed_variables: If ``True``, export every variable with @@ -3196,17 +3220,33 @@ def to_input_dataframe( pd.DataFrame: The DataFrame containing the input values. """ + return self._to_person_dataframe( + lambda variable: self._get_exportable_input_periods( + variable, include_computed_variables + ), + include_computed_variables=include_computed_variables, + ) + + def _to_person_dataframe( + self, get_periods, include_computed_variables: bool = False + ) -> pd.DataFrame: + """Person-level DataFrame of each variable at ``get_periods(variable)``.""" df = pd.DataFrame() for variable in self.tax_benefit_system.variables: variable_meta = self.tax_benefit_system.variables[variable] - for period in self._get_exportable_input_periods( - variable, include_computed_variables - ): + for period in get_periods(variable): # Test if period matches entity definition period if variable_meta.definition_period != period.unit: continue - values = self.calculate(variable, period, map_to="person") + if include_computed_variables: + values = self.calculate(variable, period, map_to="person") + else: + values = self._get_recorded_input_array(variable, period) + if values is not None: + values = self.map_result( + values, variable_meta.entity.key, self.persons.entity.key + ) if values is not None: df[f"{variable}__{period}"] = values @@ -3216,8 +3256,9 @@ def to_input_dict(self, include_computed_variables: bool = False) -> dict: """Exports a dictionary that can be loaded back into a new Simulation. By default, only structurally input variables populated through - ``set_input`` are exported. This avoids serializing pseudo-inputs and - stale calculated values that would override formulas when reloaded. + ``set_input`` are exported from their recorded storage tiers. This + avoids serializing pseudo-inputs and stale calculated values that + would override formulas when reloaded. Args: include_computed_variables: If ``True``, export every variable with @@ -3233,7 +3274,15 @@ def to_input_dict(self, include_computed_variables: bool = False) -> dict: for period in self._get_exportable_input_periods( variable, include_computed_variables ): - values = self.calculate(variable, period, map_to="person") + if include_computed_variables: + values = self.calculate(variable, period, map_to="person") + else: + values = self._get_recorded_input_array(variable, period) + if values is not None: + variable_meta = self.tax_benefit_system.get_variable(variable) + values = self.map_result( + values, variable_meta.entity.key, self.persons.entity.key + ) if values is not None: data[variable][str(period)] = values.tolist() @@ -3242,6 +3291,68 @@ def to_input_dict(self, include_computed_variables: bool = False) -> dict: return data + def _subsample_weight_period(self, time_period) -> Period: + """Use dataset input weights if the requested period has no source. + + A calculated default is not evidence of available sampling weights. + Check the input/formula paths, rather than comparing values with the + default, so a formula intentionally returning uniform weights is kept. + """ + requested = periods.period(time_period) + variable = self.tax_benefit_system.get_variable("household_weight") + calculation_period = ( + requested.this_year + if variable.definition_period == YEAR and requested.unit == MONTH + else requested + ) + calculation_periods = ( + requested.get_subperiods(MONTH) + if variable.definition_period == MONTH and requested.unit == YEAR + else [calculation_period] + ) + input_period = ( + periods.period(ETERNITY) + if variable.definition_period == ETERNITY + else calculation_period + ) + inputs = self._get_set_input_periods("household_weight") + has_formula = ( + any( + variable.get_formula(value) is not None for value in calculation_periods + ) + or variable.adds + or variable.subtracts + or variable.is_neutralized + ) + can_uprate = variable.uprating is not None and any( + source.unit == variable.definition_period + and source.start < calculation_period.start + for source in inputs + ) + can_carry = ( + self.tax_benefit_system.auto_carry_over_input_variables + and variable.calculate_output is None + and any(source.start <= calculation_period.start for source in inputs) + ) + if ( + requested in inputs + or input_period in inputs + or any(value in inputs for value in calculation_periods) + or has_formula + or can_uprate + or can_carry + ): + return requested + dataset_period = periods.period(self.dataset.time_period).start.period( + variable.definition_period + ) + if dataset_period in inputs: + return dataset_period + raise ValueError( + f"No household_weight source is available for {requested}, " + f"and no input weights exist for dataset period {dataset_period}." + ) + def subsample( self, n=None, @@ -3252,6 +3363,19 @@ def subsample( ) -> "Simulation": """Quantize the simulation to a smaller size by sampling households. + Rebuild from the holder's recorded input tiers, including inputs of + formula-backed variables. Calculate household weights for + ``time_period`` to choose the sample, falling back to dataset-period + input weights when that period has no input, formula, uprating or + carry-over source. Rescale each input weight column to preserve its + own entity's total using the memberships exported for the rebuild. + Calculated weights and other formula results are recomputed on the + sample. Calculation branches are recreated on its populations; + baseline policy is kept. + If structural inputs are absent, loaded population IDs, memberships + and roles supply the leaves needed to rebuild the same partition. + This also supports structural variables whose IDs have formulas. + Args: n (int, optional): The number of households to sample. Defaults to 10_000. frac (float, optional): The fraction of households to sample. Defaults to None. @@ -3270,33 +3394,82 @@ def subsample( if time_period is None: time_period = self.default_calculation_period - # Subsampling rebuilds the complete dataset, so preserve computed - # structural variables such as formula-backed IDs. - df = self.to_input_dataframe(include_computed_variables=True) + # Subsampling rebuilds the simulation from what it was given: every + # value loaded from the dataset or set with ``set_input``, including + # for variables that have a formula (structural IDs, for instance). + # Calculated values are left out. Reloaded as inputs they would + # replace their formulas for good, so later results (after a reform, + # or carried over to another year) would depend on what happened to + # be calculated before subsampling. + df = self._to_person_dataframe(self._get_set_input_periods) + + def retain_structure(name, values): + variable = self.tax_benefit_system.get_variable(name) + if variable is None or any(column.split("__")[0] == name for column in df): + return + structural_period = ( + ETERNITY + if variable.definition_period == ETERNITY + else periods.period(self.dataset.time_period).start.period( + variable.definition_period + ) + ) + df[f"{name}__{structural_period}"] = values - # Extract time period from DataFrame columns - df_time_period = ( - df.columns[df.columns.str.contains("household_id__")] - .values[0] - .split("__")[1] + # These are loaded population leaves, not calculated ID formulas. + # A simulation built from populations may have no stored ID inputs. + person_key = self.persons.entity.key + retain_structure(f"{person_key}_id", self.persons.ids) + for entity in self.tax_benefit_system.group_entities: + population = self.populations[entity.key] + retain_structure( + f"{person_key}_{entity.key}_id", + population.ids[population.members_entity_id], + ) + role_name = f"{person_key}_{entity.key}_role" + role_variable = self.tax_benefit_system.get_variable(role_name) + roles = np.broadcast_to(population.members_role, self.persons.count) + if role_variable is not None and role_variable.value_type == int: + role_values = [ + entity.flattened_roles.index(role) + if role in entity.flattened_roles + else 0 + for role in roles + ] + else: + role_values = [getattr(role, "key", "") for role in roles] + retain_structure(role_name, role_values) + + # Flat-file loading can infer households from membership IDs even + # when the household ID variable has a formula. Use the same leaves + # for selection without materializing a calculated ID input. + household_id_columns = [ + column for column in df if column.startswith("household_id__") + ] + if not household_id_columns: + household_id_columns = [ + column for column in df if column.startswith("person_household_id__") + ] + df_household_id_column = household_id_columns[0] + + # Sampling needs the requested year's weight even when it was + # calculated rather than supplied as an input. Use it for selection + # without adding it to the data that will be reloaded as inputs. + sampling_df = df[[df_household_id_column]].copy() + sampling_df["household_weight"] = np.asarray( + self.calculate( + "household_weight", + self._subsample_weight_period(time_period), + map_to="person", + ) ) - df_household_id_column = f"household_id__{df_time_period}" - df_person_id_column = f"person_id__{df_time_period}" - - # Determine the appropriate household weight column - if f"household_weight__{time_period}" in df.columns: - household_weight_column = f"household_weight__{time_period}" - else: - household_weight_column = f"household_weight__{df_time_period}" - - # Group by household ID and get the first entry for each group - h_df = df.groupby(df_household_id_column).first() + h_df = sampling_df.groupby(df_household_id_column).first() h_ids = pd.Series(h_df.index) if n is None and frac is None: raise ValueError("Either n or frac must be provided.") if n is None: n = int(len(h_ids) * frac) - h_weights = pd.Series(h_df[household_weight_column].values) + h_weights = pd.Series(h_df["household_weight"].values) frac = n / len(h_ids) @@ -3334,16 +3507,46 @@ def subsample( lambda x: household_id_to_count.get(x, 0) ) - # Adjust household weights to maintain the total weight + # Conserve each source column's native entity total. Group weights + # repeat on person rows, so summing those rows counts larger groups + # more often. The loader reads each retained group's first member; + # that member can differ from its original first member when another + # group partition crosses the sampled households. Group by the same + # exported membership the rebuild reads, including later overrides. for col in subset_df.columns: if "weight__" in col: - target_total_weight = df[col].values.sum() + variable_name, weight_period = col.split("__") + entity = self.tax_benefit_system.get_variable(variable_name).entity + target_total_weight = np.asarray( + self._get_recorded_input_array( + variable_name, periods.period(weight_period) + ) + ).sum() + if target_total_weight == 0: + subset_df[col] = 0.0 + continue if not quantize_weights: subset_df[col] *= household_counts.values else: subset_df[col] = household_counts.values - subset_df[col] *= target_total_weight / subset_df[col].values.sum() + if entity.is_person: + selected_total = subset_df[col].values.sum() + else: + membership_column = next( + column + for column in subset_df + if column.split("__")[0] == f"{person_key}_{entity.key}_id" + ) + memberships = subset_df[membership_column].values + _, first_retained = np.unique(memberships, return_index=True) + selected_total = subset_df[col].values[first_retained].sum() + if selected_total == 0: + raise ValueError( + f"Cannot preserve {col} total: the retained entities have " + "zero weight. Use quantize_weights=True or a different sample." + ) + subset_df[col] *= target_total_weight / selected_total df = subset_df @@ -3364,6 +3567,13 @@ def subsample( # as "size X != Y = count" projection errors. self._invalidate_all_caches() + # Other branches still contain the pre-subsample populations. Forget + # them so a formula's next ``get_branch`` starts from the new sample. + # Keep the baseline long enough to preserve its policy below. + self.branches = { + name: branch for name, branch in self.branches.items() if name == "baseline" + } + # Ensure the baseline branch has the new data: rebuild it from the # subsampled simulation through ``get_branch`` (the same wiring # ``__init__`` uses), then restore the saved baseline tax-benefit diff --git a/policyengine_core/tools/simulation_dumper.py b/policyengine_core/tools/simulation_dumper.py index daa5045e8..d351bb5d0 100644 --- a/policyengine_core/tools/simulation_dumper.py +++ b/policyengine_core/tools/simulation_dumper.py @@ -230,6 +230,7 @@ def _restore_holder( # every value counts as an input, as it did then. derived_periods_path = os.path.join(storage_dir, DERIVED_PERIODS_FILE) derived_periods = set() + # Legacy dumps have no provenance record; their calculated values may restore as inputs. if os.path.exists(derived_periods_path): with open(derived_periods_path) as file: derived_periods = set(file.read().split()) diff --git a/tests/core/test_subsample_deleted_inputs.py b/tests/core/test_subsample_deleted_inputs.py new file mode 100644 index 000000000..e6b276673 --- /dev/null +++ b/tests/core/test_subsample_deleted_inputs.py @@ -0,0 +1,159 @@ +"""Deleted source inputs must not turn later calculations into source data. + +Subsampling after deleting an override must match subsampling without reading +the deleted variable first. This holds for branches and both storage backends; +new overrides are still preserved. +""" + +import numpy as np +import pytest + +from policyengine_core.country_template.entities import Person +from policyengine_core.experimental import MemoryConfig +from policyengine_core.periods import ETERNITY, YEAR, period +from policyengine_core.variables import Variable +from tests.fixtures.subsample_inputs import DATASET_YEAR, build, stored + + +class deleted_leaf(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "An input-only variable whose deleted value becomes a default" + default_value = 17.0 + + +class deleted_eternal(Variable): + value_type = float + entity = Person + definition_period = ETERNITY + label = "An eternal formula variable with a deleted override" + + def formula(person, period): + return person("person_id", period) * 2 + + +def prepare(variable, on_disk=False, branch=False): + simulation = build() + simulation.tax_benefit_system.add_variables(deleted_leaf, deleted_eternal) + holder = simulation.get_holder(variable) + if on_disk: + # Exercise this variable's provenance on disk, including after the + # rebuild, without spilling unrelated object-valued role arrays. + simulation.memory_config = MemoryConfig( + max_memory_occupation=0, + priority_variables=[ + name + for name in simulation.tax_benefit_system.variables + if name != variable + ], + ) + holder._disk_storage = holder.create_disk_storage() + holder._on_disk_storable = True + simulation.set_input(variable, DATASET_YEAR, np.full(simulation.persons.count, 99)) + if on_disk: + assert holder._disk_storage.has(period(DATASET_YEAR)) + if branch: + simulation = simulation.get_branch("child") + return simulation + + +def delete_override(simulation, variable, through_holder): + if through_holder: + holder = simulation.get_holder(variable) + for branch_name in simulation._get_visible_branch_names(): + holder.delete_arrays(DATASET_YEAR, branch_name) + else: + simulation.delete_arrays(variable, DATASET_YEAR) + + +@pytest.mark.parametrize("variable", ["doubled", "deleted_leaf", "deleted_eternal"]) +@pytest.mark.parametrize("on_disk", [False, True], ids=["memory", "disk"]) +@pytest.mark.parametrize("through_holder", [False, True], ids=["simulation", "holder"]) +def test_deleted_input_recalculated_as_derived_is_not_exportable( + variable, on_disk, through_holder +): + simulation = prepare(variable, on_disk) + delete_override(simulation, variable, through_holder) + simulation.calculate(variable, DATASET_YEAR) + + assert simulation.get_holder(variable).is_derived( + period(DATASET_YEAR), simulation.branch_name + ) + assert simulation._get_set_input_periods(variable) == [] + + +@pytest.mark.parametrize("on_disk", [False, True], ids=["memory", "disk"]) +@pytest.mark.parametrize("branch", [False, True], ids=["default", "child"]) +def test_deleted_formula_override_does_not_freeze_after_subsample(on_disk, branch): + simulation = prepare("doubled", on_disk, branch) + simulation.delete_arrays("doubled", DATASET_YEAR) + simulation.calculate("doubled", DATASET_YEAR) + + simulation.subsample(n=3, seed="deleted-input", time_period=DATASET_YEAR) + simulation.set_input("base", DATASET_YEAR, np.full(simulation.persons.count, 5)) + + np.testing.assert_array_equal( + simulation.calculate("doubled", DATASET_YEAR), + np.full(simulation.persons.count, 10), + ) + + +@pytest.mark.parametrize("on_disk", [False, True], ids=["memory", "disk"]) +@pytest.mark.parametrize("branch", [False, True], ids=["default", "child"]) +def test_new_override_after_deletion_is_still_exportable(on_disk, branch): + simulation = prepare("doubled", on_disk, branch) + simulation.delete_arrays("doubled", DATASET_YEAR) + simulation.calculate("doubled", DATASET_YEAR) + simulation.set_input("doubled", DATASET_YEAR, np.full(simulation.persons.count, 7)) + + assert simulation._get_set_input_periods("doubled") == [period(DATASET_YEAR)] + simulation.subsample(n=3, seed="deleted-input", time_period=DATASET_YEAR) + np.testing.assert_array_equal( + simulation.calculate("doubled", DATASET_YEAR), + np.full(simulation.persons.count, 7), + ) + + +def test_reading_deleted_overrides_does_not_change_subsample(): + """Property: storage and all later formulas ignore pre-sample reads. + + Vary the deleted override, storage, branch, deletion API, sample size, + and subsequent leaf input. A never-read twin is the differential oracle. + """ + hypothesis = pytest.importorskip("hypothesis") + st = hypothesis.strategies + + @hypothesis.settings(max_examples=40, deadline=None) + @hypothesis.given( + override=st.integers(-100, 100), + on_disk=st.booleans(), + branch=st.booleans(), + through_holder=st.booleans(), + n=st.integers(1, 6), + later_input=st.integers(-100, 100), + ) + def check(override, on_disk, branch, through_holder, n, later_input): + fresh = prepare("doubled", on_disk, branch) + used = prepare("doubled", on_disk, branch) + for simulation in (fresh, used): + simulation.set_input( + "doubled", DATASET_YEAR, np.full(simulation.persons.count, override) + ) + delete_override(simulation, "doubled", through_holder) + used.calculate("doubled", DATASET_YEAR) + for simulation in (fresh, used): + simulation.subsample(n=n, seed="deleted-input", time_period=DATASET_YEAR) + + assert stored(used) == stored(fresh) + for simulation in (fresh, used): + simulation.set_input( + "base", DATASET_YEAR, np.full(simulation.persons.count, later_input) + ) + np.testing.assert_array_equal( + used.calculate("doubled", DATASET_YEAR), + fresh.calculate("doubled", DATASET_YEAR), + ) + assert stored(used) == stored(fresh) + + check() diff --git a/tests/core/test_subsample_entity_weight_totals.py b/tests/core/test_subsample_entity_weight_totals.py new file mode 100644 index 000000000..1f5d8f098 --- /dev/null +++ b/tests/core/test_subsample_entity_weight_totals.py @@ -0,0 +1,337 @@ +"""Every recorded weight column conserves its own entity's total. + +The flat-file representation repeats group weights per person. Those repeated +values must not determine the total conserved by subsampling. The independent +reference below groups by the first surviving member, matching the loader even +when a second group partition crosses households. +""" + +from collections import Counter +from unittest.mock import patch + +import numpy as np +import pandas as pd +import pytest + +from policyengine_core.data import Dataset +from policyengine_core.entities import build_entity +from policyengine_core.periods import ETERNITY, YEAR, period +from policyengine_core.simulations import Simulation +from policyengine_core.taxbenefitsystems import TaxBenefitSystem +from policyengine_core.variables import Variable +from tests.core.test_user_input_keys import _store_on_disk + +YEAR_INPUT = "2022" +Person = build_entity(key="person", plural="persons", label="Person", is_person=True) +Household = build_entity( + key="household", + plural="households", + label="Household", + roles=[{"key": "member", "plural": "members", "label": "Members"}], +) +Family = build_entity( + key="family", + plural="families", + label="Family", + roles=[{"key": "member", "plural": "members", "label": "Members"}], +) +ENTITIES = {"person": Person, "household": Household, "family": Family} +WEIGHT_YEARS = (YEAR_INPUT, "2023") + + +def weight_simulation(data): + """Load only leaf inputs into a small three-entity system.""" + system = TaxBenefitSystem(list(ENTITIES.values())) + for name, entity, period in ( + ("person_id", Person, ETERNITY), + ("household_id", Household, ETERNITY), + ("person_household_id", Person, ETERNITY), + ("family_id", Family, ETERNITY), + ("person_family_id", Person, ETERNITY), + ("person_weight", Person, YEAR), + ("household_weight", Household, YEAR), + ("family_weight", Family, YEAR), + ): + system.add_variable( + type( + name, + (Variable,), + { + "value_type": float if period == YEAR else int, + "entity": entity, + "definition_period": period, + "label": name.replace("_", " "), + }, + ) + ) + return Simulation( + tax_benefit_system=system, + dataset=Dataset.from_dataframe(data, YEAR_INPUT), + ) + + +def weight_data( + households, families, household_weights, family_weights, person_weights +): + data = pd.DataFrame( + { + "person_id__ETERNITY": np.arange(len(households)) + 101, + "household_id__ETERNITY": households, + "person_household_id__ETERNITY": households, + "family_id__ETERNITY": families, + "person_family_id__ETERNITY": families, + } + ) + for year_index, year in enumerate(WEIGHT_YEARS): + # Distinct period distributions expose accidental reuse of one column's + # total or scale factor when rescaling all recorded weight inputs. + data[f"household_weight__{year}"] = data["household_id__ETERNITY"].map( + { + group: weight * (1 + year_index * (index % 3 + 1)) + for index, (group, weight) in enumerate( + sorted(household_weights.items()) + ) + } + ) + data[f"family_weight__{year}"] = data["family_id__ETERNITY"].map( + { + group: weight * (1 + year_index * (index % 2 + 1)) + for index, (group, weight) in enumerate(sorted(family_weights.items())) + } + ) + data[f"person_weight__{year}"] = np.asarray(person_weights) * ( + 1 + year_index * (np.arange(len(households)) % 4 + 1) + ) + return data + + +def assert_weight_totals_and_reference(data, chosen, quantize_weights): + simulation = weight_simulation(data) + originals = { + (entity, year): simulation.calculate(f"{entity}_weight", year).sum() + for entity in ENTITIES + for year in WEIGHT_YEARS + } + with patch("numpy.random.choice", return_value=np.asarray(chosen)): + simulation.subsample( + n=len(chosen), + seed="entity-weight-totals", + time_period=YEAR_INPUT, + quantize_weights=quantize_weights, + ) + + retained = data[data["household_id__ETERNITY"].isin(chosen)].copy() + counts = retained["household_id__ETERNITY"].map(Counter(chosen)) + for entity in ENTITIES: + # Native entity order is independent of the flat-file person order. + entity_ids = simulation.calculate(f"{entity}_id", ETERNITY) + for year in WEIGHT_YEARS: + column = f"{entity}_weight__{year}" + actual = simulation.calculate(f"{entity}_weight", year) + target = originals[(entity, year)] + assert actual.sum() == pytest.approx(target, rel=2e-6) + + # This reference uses pandas group labels, rather than simulation + # population indices or member counts. A family can lose its first + # source member and retain members with unequal household counts. + retained["candidate_weight"] = ( + counts if quantize_weights else retained[column] * counts + ) + reference = retained.groupby(f"{entity}_id__ETERNITY", sort=True)[ + "candidate_weight" + ].first() + reference *= target / reference.sum() + np.testing.assert_allclose( + actual, reference.loc[entity_ids].to_numpy(), rtol=2e-6 + ) + + +@pytest.mark.parametrize("quantize_weights", [False, True]) +@pytest.mark.parametrize("chosen", [[20, 40, 20], [30, 20, 20]]) +def test_subsample_preserves_native_entity_weight_totals(chosen, quantize_weights): + data = weight_data( + households=[30, 10, 20, 20, 30, 40, 40, 20, 40], + families=[110, 110, 220, 110, 330, 330, 220, 220, 440], + household_weights={10: 100.0, 20: 200.0, 30: 300.0, 40: 400.0}, + family_weights={110: 5.0, 220: 11.0, 330: 17.0, 440: 23.0}, + person_weights=np.arange(2.0, 11.0), + ) + assert_weight_totals_and_reference(data, chosen, quantize_weights) + + +@pytest.mark.parametrize("quantize_weights", [False, True]) +def test_subsample_normalizes_using_exported_membership(quantize_weights): + data = weight_data( + households=[30, 10, 20, 20, 30, 40, 40, 20, 40], + families=[110, 110, 220, 110, 330, 330, 220, 220, 440], + household_weights={10: 100.0, 20: 200.0, 30: 300.0, 40: 400.0}, + family_weights={110: 5.0, 220: 11.0, 330: 17.0, 440: 23.0}, + person_weights=np.arange(2.0, 11.0), + ) + simulation = weight_simulation(data) + exported_memberships = np.array([110, 110, 220, 220, 330, 330, 330, 220, 440]) + simulation.set_input("person_family_id", ETERNITY, exported_memberships) + + # set_input changes the recorded column without changing loaded populations. + # The retained rows contain three exported families, but four old families. + chosen = [20, 40, 20] + retained = data[data["household_id__ETERNITY"].isin(chosen)].copy() + retained["person_family_id__ETERNITY"] = exported_memberships[retained.index] + counts = retained["household_id__ETERNITY"].map(Counter(chosen)) + np.testing.assert_array_equal( + simulation.to_input_dataframe()["person_family_id__ETERNITY"], + exported_memberships, + ) + with patch("numpy.random.choice", return_value=np.asarray(chosen)): + simulation.subsample( + n=len(chosen), + seed="exported-membership", + time_period=YEAR_INPUT, + quantize_weights=quantize_weights, + ) + + for entity in ENTITIES: + entity_ids = simulation.calculate(f"{entity}_id", ETERNITY) + membership_column = ( + "person_id__ETERNITY" + if entity == "person" + else f"person_{entity}_id__ETERNITY" + ) + for year in WEIGHT_YEARS: + column = f"{entity}_weight__{year}" + # Each period's target comes independently from the source's native + # entity groups, before the membership override or subsampling. + target = data.groupby(f"{entity}_id__ETERNITY")[column].first().sum() + candidate = ( + counts.astype(float) if quantize_weights else retained[column] * counts + ) + reference = candidate.groupby(retained[membership_column]).first() + reference *= target / reference.sum() + actual = simulation.calculate(f"{entity}_weight", year) + assert actual.sum() == pytest.approx(target, rel=2e-6) + np.testing.assert_allclose( + actual, reference.loc[entity_ids].to_numpy(), rtol=2e-6 + ) + + +@pytest.mark.parametrize("quantize_weights", [False, True]) +def test_subsample_conserves_recorded_disk_weight_total_under_memory_cache( + quantize_weights, +): + data = weight_data( + households=[10, 10, 20, 30, 30], + families=[110, 110, 220, 330, 110], + household_weights={10: 100.0, 20: 200.0, 30: 300.0}, + family_weights={110: 7.0, 220: 11.0, 330: 13.0}, + person_weights=np.arange(2.0, 7.0), + ) + simulation = weight_simulation(data) + simulation.delete_arrays("family_weight") + _store_on_disk(simulation, "family_weight") + source_weights = data.groupby("family_id__ETERNITY").first() + holder = simulation.get_holder("family_weight") + for year in WEIGHT_YEARS: + values = source_weights[f"family_weight__{year}"].to_numpy() + simulation.set_input("family_weight", year, values) + np.testing.assert_array_equal(holder._disk_storage.get(period(year)), values) + + simulation.memory_config.max_memory_occupation_pc = 101 + for year in WEIGHT_YEARS: + # An unmarked memory cache can shadow a surviving disk input in reads. + # Normalization must preserve the recorded source total from disk. + shadow = np.full(len(source_weights), 999.0) + holder._set(period(year), shadow, derived=False, is_input=False) + assert not holder._stores_user_input(period(year), "default", "memory") + assert holder._stores_user_input(period(year), "default", "disk") + np.testing.assert_array_equal( + simulation.calculate("family_weight", year), shadow + ) + + chosen = [20, 30, 20] + with patch("numpy.random.choice", return_value=np.asarray(chosen)): + simulation.subsample( + n=len(chosen), + seed="recorded-disk-weight-total", + time_period=YEAR_INPUT, + quantize_weights=quantize_weights, + ) + + retained = data[data["household_id__ETERNITY"].isin(chosen)] + counts = retained["household_id__ETERNITY"].map(Counter(chosen)) + entity_ids = simulation.calculate("family_id", ETERNITY) + for year in WEIGHT_YEARS: + column = f"family_weight__{year}" + target = source_weights[column].sum() + candidate = ( + counts.astype(float) if quantize_weights else retained[column] * counts + ) + reference = candidate.groupby(retained["person_family_id__ETERNITY"]).first() + reference *= target / reference.sum() + actual = simulation.calculate("family_weight", year) + assert actual.sum() == pytest.approx(target, rel=2e-6) + np.testing.assert_allclose( + actual, reference.loc[entity_ids].to_numpy(), rtol=2e-6 + ) + + +@pytest.mark.parametrize("quantize_weights", [False, True]) +@pytest.mark.parametrize("seed", range(8)) +def test_entity_total_invariant_across_deterministic_random_partitions( + seed, quantize_weights +): + rng = np.random.default_rng(seed) + households = np.repeat([10, 20, 30, 40], rng.integers(1, 5, size=4)) + rng.shuffle(households) + families = rng.choice([110, 220, 330], size=len(households)) + family_ids = np.unique(families) + data = weight_data( + households=households, + families=families, + household_weights=dict(zip([10, 20, 30, 40], rng.integers(1, 100, 4))), + family_weights=dict(zip(family_ids, rng.integers(1, 100, len(family_ids)))), + person_weights=rng.integers(1, 100, len(households)), + ) + chosen = rng.choice([10, 20, 30, 40], size=5, replace=True) + assert_weight_totals_and_reference(data, chosen, quantize_weights) + + +@pytest.mark.parametrize("quantize_weights", [False, True]) +def test_zero_total_input_weight_column_remains_zero(quantize_weights): + data = weight_data( + households=[10, 10, 20], + families=[110, 110, 220], + household_weights={10: 100, 20: 200}, + family_weights={110: 0, 220: 0}, + person_weights=[0, 0, 0], + ) + simulation = weight_simulation(data) + simulation.subsample(n=2, seed="zero-column", quantize_weights=quantize_weights) + for entity in ("person", "family"): + for year in WEIGHT_YEARS: + np.testing.assert_array_equal( + simulation.calculate(f"{entity}_weight", year), 0 + ) + + +@pytest.mark.parametrize("quantize_weights", [False, True]) +def test_positive_total_with_only_zero_retained_weights(quantize_weights): + data = weight_data( + households=[10, 20], + families=[110, 220], + household_weights={10: 1, 20: 1}, + family_weights={110: 1, 220: 1}, + person_weights=[0, 5], + ) + simulation = weight_simulation(data) + original_dataset = simulation.dataset + with patch("numpy.random.choice", return_value=np.array([10])): + if quantize_weights: + simulation.subsample(n=1, seed="zero-retained", quantize_weights=True) + assert simulation.calculate("person_weight", YEAR_INPUT).sum() == 5 + else: + with pytest.raises(ValueError, match="person_weight__2022.*zero weight"): + simulation.subsample(n=1, seed="zero-retained", quantize_weights=False) + assert simulation.dataset is original_dataset + np.testing.assert_array_equal( + simulation.calculate("person_weight", YEAR_INPUT), [0, 5] + ) diff --git a/tests/core/test_subsample_inputs_only.py b/tests/core/test_subsample_inputs_only.py new file mode 100644 index 000000000..dd7a51b81 --- /dev/null +++ b/tests/core/test_subsample_inputs_only.py @@ -0,0 +1,220 @@ +"""``subsample`` rebuilds a simulation from its inputs, not its calculations. + +``subsample`` used to export every stored value, calculated ones included, +and load them all back as inputs. A formula result calculated before +subsampling then replaced its formula for good: it was carried over past the +formula's end, and survived ``apply_reform`` and later ``set_input`` calls, +so results after subsampling depended on what had been calculated before. +""" + +import numpy as np +import pytest + +from policyengine_core.periods import period +from policyengine_core.reforms import Reform +from policyengine_core.variables import Variable +from tests.core.test_user_input_keys import _store_on_disk +from tests.fixtures.subsample_inputs import ( + DATASET_YEAR, + build, + doubled, + stored, +) + +SEED = "inputs-only" + + +def subsampled(prior=(), n=4): + """A simulation subsampled after calculating ``prior`` (variable, period).""" + simulation = build() + for variable, period in prior: + simulation.calculate(variable, period) + simulation.subsample(n=n, seed=SEED, time_period=DATASET_YEAR) + return simulation + + +def test_formula_result_is_not_carried_past_its_end_after_subsample(): + fresh = subsampled() + used = subsampled(prior=[("ended", DATASET_YEAR)]) + assert used.calculate("ended", "2023").tolist() == [0.0] * used.persons.count + assert ( + used.calculate("ended", "2023").tolist() + == fresh.calculate("ended", "2023").tolist() + ) + + +def test_subsample_stores_the_same_values_whatever_was_calculated_before(): + fresh = subsampled() + used = subsampled( + prior=[ + ("doubled", DATASET_YEAR), + ("household_total", DATASET_YEAR), + ("person_weight", DATASET_YEAR), + ("ended", DATASET_YEAR), + ("base", "2024"), + ] + ) + assert stored(used) == stored(fresh) + assert not any(name == "doubled" for name, _, _ in stored(used)) + + +def test_calculated_value_follows_a_new_input_after_subsample(): + used = subsampled(prior=[("doubled", DATASET_YEAR)]) + count = used.persons.count + used.set_input("base", DATASET_YEAR, np.full(count, 5.0)) + assert used.calculate("doubled", DATASET_YEAR).tolist() == [10.0] * count + + +def test_calculated_value_follows_a_reform_after_subsample(): + class tripled(Variable): + value_type = float + entity = doubled.entity + definition_period = doubled.definition_period + label = "Three times the input amount" + + def formula(person, period): + return 3 * person("base", period) + + tripled.__name__ = "doubled" + + class reform(Reform): + def apply(self): + self.update_variable(tripled) + + used = subsampled(prior=[("doubled", DATASET_YEAR)]) + used.apply_reform(reform) + base = used.calculate("base", DATASET_YEAR) + assert used.calculate("doubled", DATASET_YEAR).tolist() == (3 * base).tolist() + + +def test_dataset_values_for_a_formula_variable_survive_subsample(): + # ``overridden`` has a formula returning -1, and dataset values that + # replace it: person ``i`` (1-based ID) has ``i - 1``. + simulation = subsampled() + ids = simulation.calculate("person_id", DATASET_YEAR) + assert ( + simulation.calculate("overridden", DATASET_YEAR).tolist() + == (ids - 1).astype(float).tolist() + ) + assert ("overridden", "default", DATASET_YEAR) in stored(simulation) + + +@pytest.mark.parametrize("variable", ["base", "overridden"]) +@pytest.mark.parametrize("derived", [False, True], ids=["unmarked", "derived"]) +def test_subsample_uses_recorded_disk_input_under_a_memory_cache(variable, derived): + simulation = build() + simulation.delete_arrays(variable) + _store_on_disk(simulation, variable) + source_values = np.arange(simulation.persons.count, dtype=float) + 10 + simulation.set_input(variable, DATASET_YEAR, source_values) + holder = simulation.get_holder(variable) + input_period = period(DATASET_YEAR) + np.testing.assert_array_equal(holder._disk_storage.get(input_period), source_values) + + simulation.memory_config.max_memory_occupation_pc = 101 + # A low-level cache write can coexist with the recorded disk input. + # The holder's tier record distinguishes both marked and unmarked caches + # from the actual input even when ordinary reads see the memory value. + holder._set( + input_period, + np.full(simulation.persons.count, 999.0), + derived=derived, + is_input=False, + ) + assert not holder._stores_user_input(input_period, "default", "memory") + assert holder._stores_user_input(input_period, "default", "disk") + np.testing.assert_array_equal( + simulation.calculate(variable, DATASET_YEAR), + np.full(simulation.persons.count, 999.0), + ) + column = f"{variable}__{DATASET_YEAR}" + if variable == "base": + assert ( + simulation.to_input_dataframe()[column].tolist() == source_values.tolist() + ) + assert ( + simulation.to_input_dict()[variable][DATASET_YEAR] == source_values.tolist() + ) + else: + assert column not in simulation.to_input_dataframe() + assert variable not in simulation.to_input_dict() + cached_values = [999.0] * simulation.persons.count + assert ( + simulation.to_input_dataframe(include_computed_variables=True)[column].tolist() + == cached_values + ) + assert ( + simulation.to_input_dict(include_computed_variables=True)[variable][ + DATASET_YEAR + ] + == cached_values + ) + + simulation.subsample(n=4, seed=SEED, time_period=DATASET_YEAR) + + retained_ids = simulation.calculate("person_id", DATASET_YEAR) + np.testing.assert_array_equal( + simulation.calculate(variable, DATASET_YEAR), + source_values[retained_ids - 1], + ) + + +@pytest.mark.parametrize("variable,expected", [("base", 0.0), ("overridden", -1.0)]) +def test_subsample_discards_a_cache_that_replaced_the_only_input(variable, expected): + simulation = build() + holder = simulation.get_holder(variable) + holder.put_in_cache( + np.full(simulation.persons.count, 999.0), + period(DATASET_YEAR), + derived=False, + ) + if variable == "base": + assert f"{variable}__{DATASET_YEAR}" not in simulation.to_input_dataframe() + assert variable not in simulation.to_input_dict() + + simulation.subsample(n=4, seed=SEED, time_period=DATASET_YEAR) + + np.testing.assert_array_equal( + simulation.calculate(variable, DATASET_YEAR), + np.full(simulation.persons.count, expected), + ) + + +def test_formulas_run_on_the_subsample_after_subsample(): + simulation = subsampled(prior=[("household_total", DATASET_YEAR)]) + assert ( + simulation.calculate("doubled", DATASET_YEAR).tolist() + == (2 * simulation.calculate("base", DATASET_YEAR)).tolist() + ) + household_total = simulation.calculate("household_total", DATASET_YEAR) + assert household_total.sum() == simulation.calculate("doubled", DATASET_YEAR).sum() + + +@pytest.mark.parametrize("with_baseline", [False, True]) +def test_subsample_recreates_formula_branches_on_the_new_population(with_baseline): + fresh = build() + used = build() + if with_baseline: + for simulation in (fresh, used): + simulation.baseline = simulation.get_branch("baseline", clone_system=True) + baseline_system = used.baseline.tax_benefit_system + + used.calculate("via_branch", DATASET_YEAR) + stale_branch = used.get_branch("probe") + for simulation in (fresh, used): + simulation.subsample(n=4, seed=SEED, time_period=DATASET_YEAR) + + assert "probe" not in used.branches + if with_baseline: + assert used.baseline is used.branches["baseline"] + assert used.baseline.tax_benefit_system is baseline_system + assert used.baseline.persons.count == used.persons.count + + for year in (DATASET_YEAR, "2023"): + np.testing.assert_array_equal( + used.calculate("via_branch", year), fresh.calculate("via_branch", year) + ) + branch = used.get_branch("probe") + assert branch is not stale_branch + assert branch.persons.count == used.persons.count + assert branch.household.count == used.household.count diff --git a/tests/core/test_subsample_inputs_only_property.py b/tests/core/test_subsample_inputs_only_property.py new file mode 100644 index 000000000..e26dfddb1 --- /dev/null +++ b/tests/core/test_subsample_inputs_only_property.py @@ -0,0 +1,60 @@ +"""Property: what was calculated before ``subsample`` leaves no trace. + +For any calculations made before subsampling (formula variables, weights and +inputs, at the dataset year and other years, with auto-carry-over on or off), +the subsampled simulation stores exactly the values a simulation subsampled +straight after loading does, and every later result agrees. +""" + +import pytest + +from tests.fixtures.subsample_inputs import ( + DATASET_YEAR, + FORMULA_VARIABLES, + build, + stored, +) + +hypothesis = pytest.importorskip("hypothesis") +st = hypothesis.strategies + +VARIABLES = FORMULA_VARIABLES + ["base", "overridden", "household_weight"] +YEARS = ["2021", DATASET_YEAR, "2023", "2024"] +REQUESTS = st.tuples(st.sampled_from(VARIABLES), st.sampled_from(YEARS)) + + +def subsampled(carry_over, prior, n, seed): + simulation = build(carry_over) + for variable, period in prior: + simulation.calculate(variable, period) + simulation.subsample(n=n, seed=seed, time_period=DATASET_YEAR) + return simulation + + +@hypothesis.settings(max_examples=150, deadline=None) +@hypothesis.example( + carry_over=True, + prior=[("via_branch", DATASET_YEAR)], + later=[("via_branch", DATASET_YEAR), ("via_branch", "2023")], + n=4, + seed="a", +) +@hypothesis.given( + carry_over=st.booleans(), + prior=st.lists(REQUESTS, min_size=1, max_size=6), + later=st.lists(REQUESTS, min_size=1, max_size=4), + n=st.integers(1, 8), + seed=st.sampled_from(["a", "b", "c"]), +) +def test_calculations_before_subsample_leave_no_trace( + carry_over, prior, later, n, seed +): + fresh = subsampled(carry_over, [], n, seed) + used = subsampled(carry_over, prior, n, seed) + assert stored(used) == stored(fresh) + for variable, period in later: + assert ( + used.calculate(variable, period).tolist() + == fresh.calculate(variable, period).tolist() + ) + assert stored(used) == stored(fresh) diff --git a/tests/core/test_subsample_membership_ids.py b/tests/core/test_subsample_membership_ids.py new file mode 100644 index 000000000..9c434f456 --- /dev/null +++ b/tests/core/test_subsample_membership_ids.py @@ -0,0 +1,95 @@ +"""Subsampling preserves a flat-file partition recorded by membership IDs.""" + +from unittest.mock import patch + +import numpy as np +import pytest + +from policyengine_core.country_template import Simulation +from policyengine_core.data import Dataset +from tests.fixtures.subsample_inputs import DATASET_YEAR, build, dataframe, stored + + +def partition(simulation): + """Person IDs grouped by the population's actual household membership.""" + members = simulation.household.members_entity_id + return { + frozenset(simulation.persons.ids[members == index].tolist()) + for index in range(simulation.household.count) + } + + +@pytest.mark.parametrize("nonconsecutive_ids", [False, True]) +@pytest.mark.parametrize("quantize_weights", [False, True]) +@pytest.mark.parametrize("n", [1, 4]) +def test_missing_input_household_id_uses_recorded_membership( + nonconsecutive_ids, quantize_weights, n +): + data = dataframe().drop(columns=[f"household_id__{DATASET_YEAR}"]) + membership_column = f"person_household_id__{DATASET_YEAR}" + if nonconsecutive_ids: + data[membership_column] = data[membership_column].map( + {1: 901, 2: -20, 3: 33, 4: 88, 5: 205, 6: 411} + ) + data = data.iloc[[9, 0, 4, 6, 2, 8, 1, 5, 3, 7]].reset_index(drop=True) + + def make_simulation(): + return Simulation( + tax_benefit_system=build().tax_benefit_system, + dataset=Dataset.from_dataframe(data, DATASET_YEAR), + ) + + fresh, used = make_simulation(), make_simulation() + source_partition = partition(fresh) + assert fresh.household.count == 6 + used.calculate("household_id", DATASET_YEAR) + + choice = np.random.choice + for simulation in (fresh, used): + selected = {} + + def choose_households(*args, **kwargs): + selected["ids"] = choice(*args, **kwargs) + return selected["ids"] + + with patch("numpy.random.choice", side_effect=choose_households): + simulation.subsample( + n=n, + seed="recorded-memberships", + time_period=DATASET_YEAR, + quantize_weights=quantize_weights, + ) + # The resulting inputs select whole recorded groups, without adding + # the calculated default household ID. + sampled = simulation.dataset.load() + selected_ids = set(selected["ids"]) + assert set(sampled["person_household_id__ETERNITY"]) == selected_ids + expected_people = set( + data.loc[ + data[membership_column].isin(selected_ids), f"person_id__{DATASET_YEAR}" + ] + ) + assert set(simulation.persons.ids) == expected_people + assert simulation.household.count == len(selected_ids) + assert 1 <= len(selected_ids) <= n + assert partition(simulation) <= source_partition + assert not any(column.startswith("household_id__") for column in sampled) + assert not any(name == "household_id" for name, _, _ in stored(simulation)) + assert simulation.default_calculation_period == DATASET_YEAR + expected_totals = 2 * data.groupby(membership_column)[ + f"base__{DATASET_YEAR}" + ].transform("sum") + totals_by_person = dict( + zip(data[f"person_id__{DATASET_YEAR}"], expected_totals) + ) + np.testing.assert_array_equal( + simulation.calculate("household_total", DATASET_YEAR, map_to="person"), + [totals_by_person[person_id] for person_id in simulation.persons.ids], + ) + + assert stored(used) == stored(fresh) + assert partition(used) == partition(fresh) + np.testing.assert_array_equal( + used.calculate("household_weight", DATASET_YEAR), + fresh.calculate("household_weight", DATASET_YEAR), + ) diff --git a/tests/core/test_subsample_restore_inputs.py b/tests/core/test_subsample_restore_inputs.py new file mode 100644 index 000000000..77fec9f8e --- /dev/null +++ b/tests/core/test_subsample_restore_inputs.py @@ -0,0 +1,106 @@ +"""Dump/restore preserves the source inputs used to rebuild a subsample. + +For the same inputs and seed, a restored simulation and one never dumped +produce the same subsample regardless of prior calculations. Cached derived +values stay derived after restore and do not become inputs to the subsample; +explicit values of formula-backed variables remain inputs. +""" + +import numpy as np +import pytest + +from policyengine_core.periods import period +from policyengine_core.tools.simulation_dumper import ( + dump_simulation, + restore_simulation, +) +from tests.fixtures.subsample_inputs import DATASET_YEAR, build, stored + +SEED = "restored-inputs" + + +def build_with_later_input(carry_over): + simulation = build(carry_over) + simulation.set_input("overridden", "2023", np.full(simulation.persons.count, 7.0)) + return simulation + + +@pytest.mark.parametrize("carry_over", [False, True]) +@pytest.mark.parametrize("calculate_before", [False, True]) +@pytest.mark.parametrize("n", [1, 4]) +def test_restored_subsample_matches_never_dumped_inputs( + tmp_path, carry_over, calculate_before, n +): + fresh = build_with_later_input(carry_over) + source = build_with_later_input(carry_over) + if calculate_before: + for variable, requested_period in [ + ("doubled", DATASET_YEAR), + ("household_total", DATASET_YEAR), + ("person_weight", DATASET_YEAR), + ("ended", DATASET_YEAR), + ("base", "2024"), + ("overridden", "2024"), + ]: + source.calculate(variable, requested_period) + + directory = str(tmp_path / "dump") + dump_simulation(source, directory) + restored = restore_simulation(directory, source.tax_benefit_system) + # The dump stores populations and holder values, not dataset metadata. + restored.dataset = source.dataset + restored.default_calculation_period = source.default_calculation_period + + for variable in ["base", "overridden", "household_weight"]: + assert (variable, "default", period(DATASET_YEAR)) in restored._user_input_keys + assert ("overridden", "default", period("2023")) in restored._user_input_keys + if calculate_before: + for variable in ["doubled", "household_total", "person_weight", "ended"]: + holder = restored.get_holder(variable) + assert holder.get_array(period(DATASET_YEAR)) is not None + assert holder.is_derived(period(DATASET_YEAR)) + assert ( + variable, + "default", + period(DATASET_YEAR), + ) not in restored._user_input_keys + assert ("base", "default", period("2024")) not in restored._user_input_keys + assert ( + "overridden", + "default", + period("2024"), + ) not in restored._user_input_keys + + fresh.subsample(n=n, seed=SEED, time_period=DATASET_YEAR) + restored.subsample(n=n, seed=SEED, time_period=DATASET_YEAR) + assert stored(restored) == stored(fresh) + assert not any( + variable in ["doubled", "household_total", "person_weight", "ended"] + for variable, _, _ in stored(restored) + ) + for variable, requested_period in [ + ("base", DATASET_YEAR), + ("person_weight", DATASET_YEAR), + ("overridden", DATASET_YEAR), + ("overridden", "2023"), + ("ended", "2023"), + ]: + np.testing.assert_array_equal( + restored.calculate(variable, requested_period), + fresh.calculate(variable, requested_period), + ) + + # The old calculation cannot override the formula after new source data. + for simulation in [fresh, restored]: + simulation.set_input( + "base", DATASET_YEAR, np.full(simulation.persons.count, 5.0) + ) + np.testing.assert_array_equal( + simulation.calculate("doubled", DATASET_YEAR), + np.full(simulation.persons.count, 10.0), + ) + np.testing.assert_array_equal( + restored.calculate("household_total", DATASET_YEAR), + fresh.calculate("household_total", DATASET_YEAR), + ) + assert stored(restored) == stored(fresh) diff --git a/tests/core/test_subsample_sampling_period.py b/tests/core/test_subsample_sampling_period.py new file mode 100644 index 000000000..ed71a4605 --- /dev/null +++ b/tests/core/test_subsample_sampling_period.py @@ -0,0 +1,108 @@ +"""Requested-period weights select the sample without becoming source inputs.""" + +from unittest.mock import patch + +import numpy as np +import pytest + +from policyengine_core.country_template import Microsimulation, Simulation +from policyengine_core.country_template.entities import Household +from policyengine_core.data import Dataset +from policyengine_core.periods import YEAR +from policyengine_core.variables import Variable +from tests.fixtures.subsample_inputs import DATASET_YEAR, build, dataframe, stored + + +class household_weight(Variable): + value_type = float + entity = Household + definition_period = YEAR + label = "Weights growing at a different rate in each household" + + def formula(household, period): + growth = household("household_id", period) % 3 + 1 + return household("household_weight", period.last_year) * growth + + +def make_simulation(simulation_class): + system = build().tax_benefit_system + system.update_variable(household_weight) + data = dataframe() + # Microsimulation's export calculates weights as a side effect of + # reading this input, reproducing the s17 variable-order case. + data["base__2023"] = data[f"base__{DATASET_YEAR}"] + 1 + return simulation_class( + tax_benefit_system=system, + dataset=Dataset.from_dataframe(data, DATASET_YEAR), + ) + + +def check_sampling_period(simulation_class, year, n, seed, quantize_weights): + reference = make_simulation(simulation_class) + expected_ids = np.asarray(reference.calculate("household_id", DATASET_YEAR)) + expected_weights = np.asarray(reference.calculate("household_weight", year)) + expected_probabilities = expected_weights / expected_weights.sum() + fresh = make_simulation(simulation_class) + used = make_simulation(simulation_class) + used.calculate("household_weight", year) + used.calculate("person_weight", year) + + for simulation in (fresh, used): + original_weight_total = np.asarray( + simulation.calculate("household_weight", DATASET_YEAR) + ).sum() + with patch("numpy.random.choice", wraps=np.random.choice) as choose: + simulation.subsample( + n=n, seed=seed, time_period=year, quantize_weights=quantize_weights + ) + args, kwargs = choose.call_args + np.testing.assert_array_equal(args[0], expected_ids) + if quantize_weights: + np.testing.assert_allclose(kwargs["p"], expected_probabilities) + else: + assert kwargs["p"] is None + data = simulation.dataset.load() + assert np.asarray( + simulation.calculate("household_weight", DATASET_YEAR) + ).sum() == pytest.approx(original_weight_total) + if year != DATASET_YEAR: + assert f"household_weight__{year}" not in data + assert not any(name == "person_weight" for name, _, _ in stored(simulation)) + assert simulation.default_calculation_period == DATASET_YEAR + + assert stored(used) == stored(fresh) + for variable in ("household_weight", "person_weight"): + np.testing.assert_array_equal( + used.calculate(variable, year), fresh.calculate(variable, year) + ) + if year != DATASET_YEAR: + assert used.get_holder(variable).is_derived(year) + + +@pytest.mark.parametrize("simulation_class", [Simulation, Microsimulation]) +@pytest.mark.parametrize("year", [DATASET_YEAR, "2023", "2024", "2022-07", "2023-07"]) +@pytest.mark.parametrize("quantize_weights", [False, True]) +def test_sampling_uses_requested_period_weights_independently_of_calculations( + simulation_class, year, quantize_weights +): + check_sampling_period(simulation_class, year, 6, "probe", quantize_weights) + + +def test_sampling_period_is_preserved_after_repeated_subsamples(): + simulation = make_simulation(Simulation) + for n in (6, 3): + simulation.subsample(n=n, seed="again", time_period="2024") + assert simulation.default_calculation_period == DATASET_YEAR + assert "household_weight__2024" not in simulation.dataset.load() + assert not any( + name == "household_weight" and period == "2024" + for name, _, period in stored(simulation) + ) + + +def test_monthly_sampling_uses_existing_annual_flow_conversion(): + simulation = make_simulation(Simulation) + annual = np.asarray(simulation.calculate("household_weight", "2023")) + monthly = np.asarray(simulation.calculate("household_weight", "2023-07")) + np.testing.assert_allclose(monthly, annual / 12) + np.testing.assert_allclose(monthly / monthly.sum(), annual / annual.sum()) diff --git a/tests/core/test_subsample_sampling_period_property.py b/tests/core/test_subsample_sampling_period_property.py new file mode 100644 index 000000000..21489b91d --- /dev/null +++ b/tests/core/test_subsample_sampling_period_property.py @@ -0,0 +1,24 @@ +"""Differential sampling coverage across requested periods, seeds and sizes.""" + +import pytest + +from policyengine_core.country_template import Microsimulation, Simulation +from tests.core.test_subsample_sampling_period import check_sampling_period +from tests.fixtures.subsample_inputs import DATASET_YEAR + +hypothesis = pytest.importorskip("hypothesis") +st = hypothesis.strategies + + +@hypothesis.settings(max_examples=35, deadline=None) +@hypothesis.given( + simulation_class=st.sampled_from([Simulation, Microsimulation]), + year=st.sampled_from([DATASET_YEAR, "2023", "2024", "2022-07", "2023-07"]), + n=st.integers(1, 8), + seed=st.sampled_from(["a", "b", "c"]), + quantize_weights=st.booleans(), +) +def test_requested_period_sampling_matches_direct_weights( + simulation_class, year, n, seed, quantize_weights +): + check_sampling_period(simulation_class, year, n, seed, quantize_weights) diff --git a/tests/core/test_subsample_structural_ids.py b/tests/core/test_subsample_structural_ids.py new file mode 100644 index 000000000..9216904d8 --- /dev/null +++ b/tests/core/test_subsample_structural_ids.py @@ -0,0 +1,222 @@ +"""Subsampling uses the loaded population when structural IDs have formulas. + +The source partition and roles are leaves: a structural formula's cache must +neither replace them nor decide which people constitute a whole household. +""" + +import numpy as np +import pytest + +from policyengine_core.country_template import CountryTaxBenefitSystem, Simulation +from policyengine_core.country_template.entities import Household, Person +from policyengine_core.data import Dataset +from policyengine_core.periods import ETERNITY +from policyengine_core.simulations.simulation_builder import SimulationBuilder +from policyengine_core.variables import Variable +from tests.fixtures.subsample_inputs import ( + DATASET_YEAR, + HOUSEHOLD_OF_PERSON, + base, + dataframe, + doubled, + household_total, + stored, +) + + +class household_id(Variable): + value_type = int + entity = Household + definition_period = ETERNITY + label = "Formula household identifier" + + def formula(household, period): + # Deliberately unsuitable for deciding population membership. + return household.filled_array(-999) + + +class person_id(Variable): + value_type = int + entity = Person + definition_period = ETERNITY + label = "Formula person identifier" + + def formula(person, period): + return np.arange(person.count) + 700 + + +class person_household_id(Variable): + value_type = int + entity = Person + definition_period = ETERNITY + label = "Formula person household identifier" + + def formula(person, period): + return person.filled_array(-999) + + +class person_household_role(Variable): + value_type = str + entity = Person + definition_period = ETERNITY + label = "Formula person household role" + + def formula(person, period): + return person.filled_array("child") + + +STRUCTURAL_VARIABLES = ( + household_id, + person_id, + person_household_id, + person_household_role, +) + + +def partition(simulation): + return { + frozenset( + simulation.persons.ids[ + simulation.household.members_entity_id == index + ].tolist() + ) + for index in range(simulation.household.count) + } + + +def roles_by_person(simulation): + return { + person: getattr(role, "key", role) + for person, role in zip( + simulation.persons.ids, simulation.household.members_role + ) + } + + +def make_population_simulation(recorded_structural_inputs, shuffled): + """A population with structural leaves and optional formula overrides.""" + system = CountryTaxBenefitSystem() + system.add_variables(base, doubled, household_total) + for variable in STRUCTURAL_VARIABLES: + system.replace_variable(variable) + + data = dataframe() + labels = {1: 901, 2: -20, 3: 33, 4: 88, 5: 205, 6: 411} + membership = np.array([labels[group] for group in HOUSEHOLD_OF_PERSON]) + # Explicit first/second parent and child roles make role loss observable. + roles = np.array([0, 2, 0, 1, 2, 0, 0, 2, 0, 0]) + order = np.array([9, 0, 4, 6, 2, 8, 1, 5, 3, 7]) if shuffled else np.arange(10) + data = data.iloc[order].reset_index(drop=True) + membership, roles = membership[order], roles[order] + + builder = SimulationBuilder() + builder.create_entities(system) + builder.declare_person_entity("person", data[f"person_id__{DATASET_YEAR}"]) + household = builder.declare_entity("household", list(labels.values())) + builder.join_with_persons(household, membership, roles) + simulation = Simulation( + tax_benefit_system=system, + populations=builder.populations, + default_calculation_period=DATASET_YEAR, + ) + # Attach period metadata without reloading the explicitly built population. + simulation.dataset = Dataset.from_dataframe(data, DATASET_YEAR) + simulation.is_over_dataset = True + simulation.set_input("base", DATASET_YEAR, data[f"base__{DATASET_YEAR}"].values) + simulation.set_input( + "household_weight", DATASET_YEAR, [100, 200, 300, 400, 500, 600] + ) + if recorded_structural_inputs: + simulation.set_input("person_id", ETERNITY, simulation.persons.ids) + simulation.set_input("person_household_id", ETERNITY, membership) + simulation.set_input( + "person_household_role", + ETERNITY, + [role.key for role in simulation.household.members_role], + ) + return simulation + + +def test_formula_only_household_id_uses_recorded_membership(): + """The reviewed failure: no household ID source gives an IndexError.""" + data = dataframe().drop(columns=[f"household_id__{DATASET_YEAR}"]) + system = CountryTaxBenefitSystem() + system.replace_variable(household_id) + system.add_variables(base, doubled, household_total) + simulation = Simulation( + tax_benefit_system=system, + dataset=Dataset.from_dataframe(data, DATASET_YEAR), + ) + source_partition = partition(simulation) + simulation.calculate("household_id", ETERNITY) + + simulation.subsample(n=4, seed="formula-household-id", time_period=DATASET_YEAR) + + assert partition(simulation) <= source_partition + assert 1 <= simulation.household.count <= 4 + assert not any(name == "household_id" for name, _, _ in stored(simulation)) + + +@pytest.mark.parametrize("recorded_structural_inputs", [False, True]) +@pytest.mark.parametrize("shuffled", [False, True]) +@pytest.mark.parametrize("quantize_weights", [False, True]) +@pytest.mark.parametrize("n,seed", [(1, "one"), (4, "four"), (8, "repeats")]) +def test_loaded_structural_leaves_survive_independent_of_formula_caches( + recorded_structural_inputs, shuffled, quantize_weights, n, seed +): + """Differential: formula caches cannot alter IDs, groups, roles or sums.""" + fresh = make_population_simulation(recorded_structural_inputs, shuffled) + used = make_population_simulation(recorded_structural_inputs, shuffled) + source_partition = partition(fresh) + source_roles = roles_by_person(fresh) + source_totals = dict( + zip( + fresh.persons.ids, + fresh.calculate("household_total", DATASET_YEAR, map_to="person"), + ) + ) + for variable in STRUCTURAL_VARIABLES: + used.calculate(variable.__name__, ETERNITY) + + for simulation in (fresh, used): + simulation.subsample( + n=n, + seed=seed, + time_period=DATASET_YEAR, + quantize_weights=quantize_weights, + ) + assert partition(simulation) <= source_partition + assert 1 <= simulation.household.count <= min(n, 6) + assert roles_by_person(simulation) == { + person: source_roles[person] for person in simulation.persons.ids + } + np.testing.assert_array_equal( + simulation.calculate("household_total", DATASET_YEAR, map_to="person"), + [source_totals[person] for person in simulation.persons.ids], + ) + assert not any(name == "household_id" for name, _, _ in stored(simulation)) + + assert partition(used) == partition(fresh) + assert stored(used) == stored(fresh) + + +@pytest.mark.parametrize("shuffled", [False, True]) +def test_formula_only_integer_roles_preserve_loaded_subroles(shuffled): + simulation = make_population_simulation(False, shuffled) + source_roles = roles_by_person(simulation) + + class integer_role(Variable): + value_type = int + entity = Person + definition_period = ETERNITY + label = "Formula integer role" + + def formula(person, period): + return person.filled_array(2) + + integer_role.__name__ = "person_household_role" + simulation.tax_benefit_system.replace_variable(integer_role) + simulation.subsample(n=4, seed="integer-roles", time_period=DATASET_YEAR) + assert roles_by_person(simulation) == { + person: source_roles[person] for person in simulation.persons.ids + } diff --git a/tests/core/test_subsample_weight_fallback.py b/tests/core/test_subsample_weight_fallback.py new file mode 100644 index 000000000..8acc27152 --- /dev/null +++ b/tests/core/test_subsample_weight_fallback.py @@ -0,0 +1,143 @@ +"""Unavailable sampling periods use recorded weights rather than defaults.""" + +from unittest.mock import patch + +import numpy as np +import pytest + +hypothesis = pytest.importorskip("hypothesis") +st = hypothesis.strategies + +from policyengine_core.country_template import Microsimulation, Simulation +from policyengine_core.country_template.entities import Household +from policyengine_core.data import Dataset +from policyengine_core.parameters import Parameter +from policyengine_core.periods import MONTH, YEAR +from policyengine_core.variables import Variable +from tests.fixtures.subsample_inputs import DATASET_YEAR, build, dataframe, stored + + +class household_weight(Variable): + value_type = float + entity = Household + definition_period = YEAR + default_value = 1 + uprating = "weight_index" + + +def make_simulation(simulation_class, uprated, carry_over): + system = build(carry_over=carry_over).tax_benefit_system + if uprated: + system.update_variable(household_weight) + system.parameters.add_child( + "weight_index", + Parameter( + "weight_index", + data={"values": {"2021-01-01": 1, "2022-01-01": 2, "2023-01-01": 3}}, + ), + ) + return simulation_class( + tax_benefit_system=system, + dataset=Dataset.from_dataframe(dataframe(), DATASET_YEAR), + ) + + +def check_fallback(simulation_class, uprated, carry_over, year, n, seed): + fresh = make_simulation(simulation_class, uprated, carry_over) + used = make_simulation(simulation_class, uprated, carry_over) + expected_weights = np.asarray(fresh.calculate("household_weight", DATASET_YEAR)) + expected_ids = fresh.calculate("household_id", DATASET_YEAR) + used.calculate("household_weight", year) + for simulation in (fresh, used): + with patch("numpy.random.choice", wraps=np.random.choice) as choose: + simulation.subsample(n=n, seed=seed, time_period=year) + args, kwargs = choose.call_args + np.testing.assert_array_equal(args[0], expected_ids) + np.testing.assert_allclose( + kwargs["p"], expected_weights / expected_weights.sum() + ) + assert f"household_weight__{year}" not in simulation.dataset.load() + assert simulation.default_calculation_period == DATASET_YEAR + assert stored(fresh) == stored(used) + + +@pytest.mark.parametrize("simulation_class", [Simulation, Microsimulation]) +@pytest.mark.parametrize( + "uprated,carry_over,year", + [ + (False, False, "2023"), + (False, False, "2023-07"), + (False, True, "2021"), + (True, False, "2021"), + (True, True, "2021-07"), + ], +) +def test_missing_period_weights_use_dataset_input( + simulation_class, uprated, carry_over, year +): + check_fallback(simulation_class, uprated, carry_over, year, 3, "fallback") + + +@hypothesis.settings(max_examples=25, deadline=None) +@hypothesis.given( + simulation_class=st.sampled_from([Simulation, Microsimulation]), + case=st.sampled_from( + [(False, False, "2023"), (False, True, "2021"), (True, True, "2021-07")] + ), + n=st.integers(1, 8), + seed=st.integers(0, 100), +) +def test_weight_fallback_is_independent_of_prior_default_calculation( + simulation_class, case, n, seed +): + check_fallback(simulation_class, *case, n, seed) + + +def test_available_requested_input_takes_precedence_over_dataset_weights(): + simulation = make_simulation(Simulation, False, False) + weights = np.arange(1, simulation.household.count + 1, dtype=float)[::-1] + simulation.set_input("household_weight", "2023", weights) + with patch("numpy.random.choice", wraps=np.random.choice) as choose: + simulation.subsample(n=3, seed="explicit", time_period="2023") + np.testing.assert_allclose(choose.call_args.kwargs["p"], weights / weights.sum()) + + +def test_absent_weight_source_has_clear_error(): + simulation = make_simulation(Simulation, False, False) + simulation.delete_arrays("household_weight") + with pytest.raises(ValueError, match="household_weight.*2023.*2022"): + simulation.subsample(n=1, seed="missing", time_period="2023") + + +def test_annual_weight_fallback_normalizes_dataset_period(): + simulation = make_simulation(Simulation, False, False) + simulation.dataset.time_period = "2022-01" + weights = np.asarray(simulation.calculate("household_weight", DATASET_YEAR)) + with patch("numpy.random.choice", wraps=np.random.choice) as choose: + simulation.subsample(n=3, seed="dataset-month", time_period="2023") + np.testing.assert_allclose(choose.call_args.kwargs["p"], weights / weights.sum()) + + +def test_annual_request_uses_monthly_formula_starting_midyear(): + class monthly_weight(Variable): + value_type = float + entity = Household + definition_period = MONTH + label = "Monthly household weight" + + def formula_2023_07(household, period): + return household.filled_array(2.0) + + monthly_weight.__name__ = "household_weight" + system = build(carry_over=False).tax_benefit_system + system.replace_variable(monthly_weight) + data = dataframe().rename( + columns={"household_weight__2022": "household_weight__2022-01"} + ) + simulation = Simulation( + tax_benefit_system=system, + dataset=Dataset.from_dataframe(data, DATASET_YEAR), + ) + with patch("numpy.random.choice", wraps=np.random.choice) as choose: + simulation.subsample(n=3, seed="monthly-formula", time_period="2023") + np.testing.assert_allclose(choose.call_args.kwargs["p"], np.full(6, 1 / 6)) diff --git a/tests/fixtures/subsample_inputs.py b/tests/fixtures/subsample_inputs.py new file mode 100644 index 000000000..9d6f1ed41 --- /dev/null +++ b/tests/fixtures/subsample_inputs.py @@ -0,0 +1,127 @@ +"""A small dataset-backed simulation for testing what ``subsample`` keeps. + +Shared by ``tests/core/test_subsample_inputs_only.py`` and its Hypothesis +property module, which pytest must be able to skip on its own when +Hypothesis is not installed. +""" + +import numpy as np +import pandas as pd + +from policyengine_core.country_template import CountryTaxBenefitSystem, Simulation +from policyengine_core.country_template.entities import Household, Person +from policyengine_core.data import Dataset +from policyengine_core.periods import YEAR +from policyengine_core.variables import Variable + +DATASET_YEAR = "2022" +HOUSEHOLD_OF_PERSON = [1, 1, 2, 2, 2, 3, 4, 4, 5, 6] +BASE = [10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0, 90.0, 100.0] +WEIGHTS = {1: 100.0, 2: 200.0, 3: 300.0, 4: 400.0, 5: 500.0, 6: 600.0} + + +class base(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Input amount" + + +class doubled(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Twice the input amount" + + def formula(person, period): + return 2 * person("base", period) + + +class ended(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Seven, until its formula ends with the dataset year" + end = f"{DATASET_YEAR}-12-31" + + def formula(person, period): + return person.filled_array(7.0) + + +class household_total(Variable): + value_type = float + entity = Household + definition_period = YEAR + label = "Household sum of the doubled amount" + + def formula(household, period): + return household.sum(household.members("doubled", period)) + + +class overridden(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "A formula variable the dataset gives values for" + + def formula(person, period): + return person.filled_array(-1.0) + + +class via_branch(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "A formula that changes an input in its own branch" + + def formula(person, period): + branch = person.simulation.get_branch("probe") + branch.set_input("base", period, np.full(person.count, 1.0)) + return branch.calculate("doubled", period) + + +FORMULA_VARIABLES = [ + "doubled", + "ended", + "household_total", + "person_weight", + "via_branch", +] +OVERRIDDEN = [float(index) for index in range(len(BASE))] + + +def dataframe() -> pd.DataFrame: + return pd.DataFrame( + { + f"person_id__{DATASET_YEAR}": list(range(1, len(BASE) + 1)), + f"household_id__{DATASET_YEAR}": HOUSEHOLD_OF_PERSON, + f"person_household_id__{DATASET_YEAR}": HOUSEHOLD_OF_PERSON, + f"household_weight__{DATASET_YEAR}": [ + WEIGHTS[household] for household in HOUSEHOLD_OF_PERSON + ], + f"base__{DATASET_YEAR}": BASE, + f"overridden__{DATASET_YEAR}": OVERRIDDEN, + } + ) + + +def build(carry_over: bool = True) -> Simulation: + """A ten-person, six-household simulation loaded from ``dataframe()``.""" + system = CountryTaxBenefitSystem() + system.auto_carry_over_input_variables = carry_over + system.add_variables(base, doubled, ended, household_total, overridden, via_branch) + return Simulation( + tax_benefit_system=system, + dataset=Dataset.from_dataframe(dataframe(), DATASET_YEAR), + ) + + +def stored(simulation) -> dict: + """Every stored value: ``{(variable, branch, period): values}``.""" + values = {} + for population in simulation.populations.values(): + for name, holder in population._holders.items(): + for branch, period in holder.get_known_branch_periods(): + array = holder.get_array(period, branch) + values[(name, branch, str(period))] = np.asarray(array).tolist() + return values