diff --git a/changelog.d/560.fixed.md b/changelog.d/560.fixed.md new file mode 100644 index 00000000..b30e0325 --- /dev/null +++ b/changelog.d/560.fixed.md @@ -0,0 +1 @@ +Add opt-in verified US state/year preparation and temporarily assume WIC claiming only for missing participation decisions on identified ACS people, preserving donor decisions and country-model eligibility. diff --git a/changelog.d/us-state-preparation.added.md b/changelog.d/us-state-preparation.added.md new file mode 100644 index 00000000..b6764f7d --- /dev/null +++ b/changelog.d/us-state-preparation.added.md @@ -0,0 +1 @@ +Add opt-in, hash-verified US state input partitioning and shared state/year preparation for precompute benchmarks. Preserve native entity relationships, source periods, weights, and geography without changing dataset defaults. diff --git a/docs/engineering/runbooks/acs-local-wic-replacement.md b/docs/engineering/runbooks/acs-local-wic-replacement.md new file mode 100644 index 00000000..c64ee63b --- /dev/null +++ b/docs/engineering/runbooks/acs-local-wic-replacement.md @@ -0,0 +1,68 @@ +# Replace the temporary ACS WIC assumption + +## Current behaviour + +The certified local-area release +`populace-us-2024-buildo-acs-local-767312d60-20260923T074941Z` stores donor +participation decisions as `would_claim_wic`, with missing decisions on ACS +people. The current model consumes the monthly person input +`takes_up_wic_if_eligible`. + +The wrapper's temporary compatibility exception preserves all valid existing +decisions and sets only missing decisions on people whose +`person_support_channel` is `acs_2024_1yr` to `True`. This explicitly assumes +that every eligible affected ACS person claims WIC; it does not make them +eligible. The country model still calculates eligibility and benefits. +Missing donor decisions, missing identifying provenance, and malformed values +remain errors. Source H5 files are never modified by this exception. + +The shared loader applies this before calculation and year preparation, +including baseline/reform branches and verified state derivatives. Existing +national donor decisions are preserved. State derivatives remain opt-in +preparation artifacts, not new certified datasets or request-serving defaults. + +## Replacement ownership + +[Microcosm #1154](https://github.com/PolicyEngine/microcosm/issues/1154) tracks +the code repair, a rebuild/publication performed by a separate dataset owner, +certification in this repository, and removal of the exception. Do not close +that issue when the Microcosm code merges. Passing code tests does not certify +a replacement population. + +The Microcosm repair must generate ACS participation after the demographic +inputs are populated and before combining the population with donors. It must +reuse the existing category-specific WIC generator, preserve donor decisions, +write complete boolean `takes_up_wic_if_eligible`, and reject default-filling +missing WIC participation. The dataset owner must rebuild affected calculations +and calibration outputs and satisfy every existing release check. + +## Remove the exception + +The removal change is prepared in the separate stacked draft +[PR #563](https://github.com/PolicyEngine/policyengine.py/pull/563). It is +blocked until a new qualified local-area release is published. Do not invent +its revision or hash, weaken tests, or merge removal against the broken pin. + +When the owner supplies the release: + +1. Use the existing + [US certification runbook](build-m-us-populace-certification.md), supplying + the actual immutable replacement regional manifest. Keep the national + manifest, default dataset, and model pins unchanged unless separately + approved. A model mismatch requires a reviewed compatibility claim or a + compatible replacement build, not a bypass. +2. Confirm complete current-name participation on the replacement, with no + obsolete WIC input column, and record its immutable revision and hashes. +3. Update the already-open removal PR with certification changes and derived + bundle metadata. Rebase and retarget it to `main` after its compatibility + parent merges. +4. Delete only the ACS `True`-fill exception. Retain ordinary old-name mapping + while other certified datasets require it. Restore strict missing-value + regression tests and verify native current-name inputs. +5. Run the focused mapping/preparation tests and the certification checks on + the actual replacement. Never reuse prepared outputs derived from another + source revision/hash. + +This work does not rebuild or publish the full dataset, launch paid compute, +modify the national latest pointer, deploy services, or introduce environment +variables or database changes. diff --git a/src/policyengine/tax_benefit_models/us/__init__.py b/src/policyengine/tax_benefit_models/us/__init__.py index bfd62147..829342dd 100644 --- a/src/policyengine/tax_benefit_models/us/__init__.py +++ b/src/policyengine/tax_benefit_models/us/__init__.py @@ -55,6 +55,13 @@ us_latest, ) from .spm import SPMProvenance, SPMSelection + from .state_preparation import ( + USPartitionManifest, + USStatePartition, + USStateYearArtifact, + partition_certified_us_source, + prepare_us_state_year, + ) model = us_latest """The pinned US ``TaxBenefitModelVersion`` for this policyengine release.""" @@ -68,6 +75,11 @@ __all__ = [ "USYearData", + "USPartitionManifest", + "USStatePartition", + "USStateYearArtifact", + "partition_certified_us_source", + "prepare_us_state_year", "PolicyEngineUSDataset", "create_datasets", "load_datasets", diff --git a/src/policyengine/tax_benefit_models/us/datasets.py b/src/policyengine/tax_benefit_models/us/datasets.py index 111f4880..4ef03737 100644 --- a/src/policyengine/tax_benefit_models/us/datasets.py +++ b/src/policyengine/tax_benefit_models/us/datasets.py @@ -4,7 +4,7 @@ import warnings from importlib import metadata as importlib_metadata from pathlib import Path -from typing import Any, Optional +from typing import TYPE_CHECKING, Any, Optional import h5py import pandas as pd @@ -34,6 +34,9 @@ ) from policyengine.utils.hashing import sha256_file +if TYPE_CHECKING: + from policyengine_us import Microsimulation + class USYearData(YearData): """Entity-level data for a single year.""" @@ -399,167 +402,175 @@ def create_datasets( legacy_input_renames = apply_legacy_input_renames_to_microsimulation(sim) for year in years: - # Get all input variables from the simulation - # We'll calculate each input variable for the specified year - entity_data = { - "person": {}, - "household": {}, - "marital_unit": {}, - "family": {}, - "spm_unit": {}, - "tax_unit": {}, - } + us_dataset = _prepare_us_year( + sim, + year=year, + dataset_stem=dataset_stem, + filepath=Path(data_folder) / f"{dataset_stem}_year_{year}.h5", + legacy_input_renames=legacy_input_renames, + ) + us_dataset.save() + result[f"{dataset_stem}_{year}"] = us_dataset + return result - # First, get ID columns which are structural (not input variables) - # These define entity membership and relationships - # For person-level links to group entities, use person_X_id naming - id_variables = { - "person": [ - "person_id", - "person_household_id", - "person_marital_unit_id", - "person_family_id", - "person_spm_unit_id", - "person_tax_unit_id", - ], - "household": ["household_id"], - "marital_unit": ["marital_unit_id"], - "family": ["family_id"], - "spm_unit": ["spm_unit_id"], - "tax_unit": ["tax_unit_id"], - } - for entity_key, var_names in id_variables.items(): - for id_var in var_names: - if id_var in sim.tax_benefit_system.variables: - values = sim.calculate(id_var, period=year).values - entity_data[entity_key][id_var] = values - - # Get input variables and calculate them for this year - for variable_name in sim.input_variables: - variable = sim.tax_benefit_system.variables[variable_name] - entity_key = variable.entity.key - - # Calculate the variable for the given year - values = sim.calculate(variable_name, period=year).values - - # Store in the appropriate entity dictionary - entity_data[entity_key][variable_name] = values - - # Build entity DataFrames - person_df = pd.DataFrame(entity_data["person"]) - household_df = pd.DataFrame(entity_data["household"]) - marital_unit_df = pd.DataFrame(entity_data["marital_unit"]) - family_df = pd.DataFrame(entity_data["family"]) - spm_unit_df = pd.DataFrame(entity_data["spm_unit"]) - tax_unit_df = pd.DataFrame(entity_data["tax_unit"]) - - # Add weight columns - household weights are primary, map to all entities - # Person weights = household weights (mapped via person_household_id) - if "household_weight" in household_df.columns: - # Only add person_weight if it doesn't already exist - if "person_weight" not in person_df.columns: - person_df = person_df.merge( - household_df[["household_id", "household_weight"]], - left_on="person_household_id", - right_on="household_id", - how="left", - ) - person_df = person_df.rename( - columns={"household_weight": "person_weight"} - ) - person_df = person_df.drop( - columns=["household_id"], errors="ignore" - ) - - # Map household weights to other group entities via person table - for entity_name, entity_df, person_id_col, entity_id_col in [ - ( - "marital_unit", - marital_unit_df, - "person_marital_unit_id", - "marital_unit_id", - ), - ("family", family_df, "person_family_id", "family_id"), - ( - "spm_unit", - spm_unit_df, - "person_spm_unit_id", - "spm_unit_id", - ), - ( - "tax_unit", - tax_unit_df, - "person_tax_unit_id", - "tax_unit_id", - ), - ]: - # Only add entity weight if it doesn't already exist - if f"{entity_name}_weight" not in entity_df.columns: - # Get household_id for each entity from person table - entity_household_map = person_df[ - [person_id_col, "person_household_id"] - ].drop_duplicates() - entity_df = entity_df.merge( - entity_household_map, - left_on=entity_id_col, - right_on=person_id_col, - how="left", - ) - entity_df = entity_df.merge( - household_df[["household_id", "household_weight"]], - left_on="person_household_id", - right_on="household_id", - how="left", - ) - entity_df = entity_df.rename( - columns={"household_weight": f"{entity_name}_weight"} - ) - entity_df = entity_df.drop( - columns=[ - "household_id", - "person_household_id", - person_id_col, - ], - errors="ignore", - ) - - # Update the entity_data - if entity_name == "marital_unit": - marital_unit_df = entity_df - elif entity_name == "family": - family_df = entity_df - elif entity_name == "spm_unit": - spm_unit_df = entity_df - elif entity_name == "tax_unit": - tax_unit_df = entity_df +def _prepare_us_year( + sim: "Microsimulation", + *, + year: int, + dataset_stem: str, + filepath: Path, + legacy_input_renames: dict[str, str], +) -> PolicyEngineUSDataset: + """Extract one year identically for national inputs and verified state inputs.""" + # Get all input variables from the simulation + # We'll calculate each input variable for the specified year + entity_data = { + "person": {}, + "household": {}, + "marital_unit": {}, + "family": {}, + "spm_unit": {}, + "tax_unit": {}, + } - us_dataset = PolicyEngineUSDataset( - id=f"{dataset_stem}_year_{year}", - name=f"{dataset_stem}-year-{year}", - description=f"US Dataset for year {year} based on {dataset_stem}", - filepath=f"{data_folder}/{dataset_stem}_year_{year}.h5", - year=int(year), - metadata={ - RENAMES_RECORD_KEY: dict(sorted(legacy_input_renames.items())) - }, - data=USYearData( - person=MicroDataFrame(person_df, weights="person_weight"), - household=MicroDataFrame(household_df, weights="household_weight"), - marital_unit=MicroDataFrame( - marital_unit_df, weights="marital_unit_weight" - ), - family=MicroDataFrame(family_df, weights="family_weight"), - spm_unit=MicroDataFrame(spm_unit_df, weights="spm_unit_weight"), - tax_unit=MicroDataFrame(tax_unit_df, weights="tax_unit_weight"), - ), - ) - us_dataset.save() + # First, get ID columns which are structural (not input variables) + # These define entity membership and relationships + # For person-level links to group entities, use person_X_id naming + id_variables = { + "person": [ + "person_id", + "person_household_id", + "person_marital_unit_id", + "person_family_id", + "person_spm_unit_id", + "person_tax_unit_id", + ], + "household": ["household_id"], + "marital_unit": ["marital_unit_id"], + "family": ["family_id"], + "spm_unit": ["spm_unit_id"], + "tax_unit": ["tax_unit_id"], + } - dataset_key = f"{dataset_stem}_{year}" - result[dataset_key] = us_dataset + for entity_key, var_names in id_variables.items(): + for id_var in var_names: + if id_var in sim.tax_benefit_system.variables: + values = sim.calculate(id_var, period=year).values + entity_data[entity_key][id_var] = values + + # Get input variables and calculate them for this year + for variable_name in sim.input_variables: + variable = sim.tax_benefit_system.variables[variable_name] + entity_key = variable.entity.key + + # Calculate the variable for the given year + values = sim.calculate(variable_name, period=year).values + + # Store in the appropriate entity dictionary + entity_data[entity_key][variable_name] = values + + # Build entity DataFrames + person_df = pd.DataFrame(entity_data["person"]) + household_df = pd.DataFrame(entity_data["household"]) + marital_unit_df = pd.DataFrame(entity_data["marital_unit"]) + family_df = pd.DataFrame(entity_data["family"]) + spm_unit_df = pd.DataFrame(entity_data["spm_unit"]) + tax_unit_df = pd.DataFrame(entity_data["tax_unit"]) + + # Add weight columns - household weights are primary, map to all entities + # Person weights = household weights (mapped via person_household_id) + if "household_weight" in household_df.columns: + # Only add person_weight if it doesn't already exist + if "person_weight" not in person_df.columns: + person_df = person_df.merge( + household_df[["household_id", "household_weight"]], + left_on="person_household_id", + right_on="household_id", + how="left", + ) + person_df = person_df.rename(columns={"household_weight": "person_weight"}) + person_df = person_df.drop(columns=["household_id"], errors="ignore") + + # Map household weights to other group entities via person table + for entity_name, entity_df, person_id_col, entity_id_col in [ + ( + "marital_unit", + marital_unit_df, + "person_marital_unit_id", + "marital_unit_id", + ), + ("family", family_df, "person_family_id", "family_id"), + ( + "spm_unit", + spm_unit_df, + "person_spm_unit_id", + "spm_unit_id", + ), + ( + "tax_unit", + tax_unit_df, + "person_tax_unit_id", + "tax_unit_id", + ), + ]: + # Only add entity weight if it doesn't already exist + if f"{entity_name}_weight" not in entity_df.columns: + # Get household_id for each entity from person table + entity_household_map = person_df[ + [person_id_col, "person_household_id"] + ].drop_duplicates() + entity_df = entity_df.merge( + entity_household_map, + left_on=entity_id_col, + right_on=person_id_col, + how="left", + ) + entity_df = entity_df.merge( + household_df[["household_id", "household_weight"]], + left_on="person_household_id", + right_on="household_id", + how="left", + ) + entity_df = entity_df.rename( + columns={"household_weight": f"{entity_name}_weight"} + ) + entity_df = entity_df.drop( + columns=[ + "household_id", + "person_household_id", + person_id_col, + ], + errors="ignore", + ) - return result + # Update the entity_data + if entity_name == "marital_unit": + marital_unit_df = entity_df + elif entity_name == "family": + family_df = entity_df + elif entity_name == "spm_unit": + spm_unit_df = entity_df + elif entity_name == "tax_unit": + tax_unit_df = entity_df + + us_dataset = PolicyEngineUSDataset( + id=f"{dataset_stem}_year_{year}", + name=f"{dataset_stem}-year-{year}", + description=f"US Dataset for year {year} based on {dataset_stem}", + filepath=str(filepath), + year=int(year), + metadata={RENAMES_RECORD_KEY: dict(sorted(legacy_input_renames.items()))}, + data=USYearData( + person=MicroDataFrame(person_df, weights="person_weight"), + household=MicroDataFrame(household_df, weights="household_weight"), + marital_unit=MicroDataFrame(marital_unit_df, weights="marital_unit_weight"), + family=MicroDataFrame(family_df, weights="family_weight"), + spm_unit=MicroDataFrame(spm_unit_df, weights="spm_unit_weight"), + tax_unit=MicroDataFrame(tax_unit_df, weights="tax_unit_weight"), + ), + ) + return us_dataset def _year_file_records_renames(path: Path) -> bool: diff --git a/src/policyengine/tax_benefit_models/us/legacy_inputs.py b/src/policyengine/tax_benefit_models/us/legacy_inputs.py index 9558cd6d..9834d130 100644 --- a/src/policyengine/tax_benefit_models/us/legacy_inputs.py +++ b/src/policyengine/tax_benefit_models/us/legacy_inputs.py @@ -146,8 +146,9 @@ def apply_legacy_input_renames( Every mapped table is checked before anything is set. Its ``{entity}_id`` column must list the simulation's entity IDs in the simulation's order, and its stored values must be complete and, for a boolean live input, - boolean. Otherwise this raises ``ValueError`` rather than set misaligned - or invented values. + boolean. A temporary WIC-only exception assumes claiming for missing + decisions on explicitly identified ACS people. Missing donor decisions + and malformed stored values still raise ``ValueError``. Applying the mapping again sets the same values, so it is idempotent. @@ -165,7 +166,15 @@ def apply_legacy_input_renames( continue context = f"Cannot map stored {legacy!r} onto {live!r} for {year}" _check_order(simulation, table, entity, context) - values = _live_values(table[legacy], variable, context) + stored = table[legacy] + if ( + legacy == "would_claim_wic" + and live == "takes_up_wic_if_eligible" + and entity == "person" + and variable.value_type is bool + ): + stored = _temporary_acs_wic_values(table, context) + values = _live_values(stored, variable, context) periods = _periods_of_year(int(year), variable.definition_period, context) planned.append((legacy, live, periods, values)) @@ -347,6 +356,32 @@ def _check_order(simulation, table: pd.DataFrame, entity: str, context: str) -> ) +def _temporary_acs_wic_values(table: pd.DataFrame, context: str) -> pd.Series: + """Preserve stored decisions; assume claiming only for missing ACS cells.""" + stored = table["would_claim_wic"] + missing = stored.isna() + if not missing.any(): + return stored + channel = "person_support_channel" + if channel not in table or not ( + table.loc[missing, channel].eq("acs_2024_1yr").fillna(False).all() + ): + raise ValueError( + f"{context}: the stored column has missing values outside " + "explicitly identified ACS people." + ) + # TEMPORARY ACS WIC COMPATIBILITY — REMOVE AS SOON AS POSSIBLE. + # The currently certified ACS-local dataset did not generate WIC + # participation decisions for ACS people. For missing ACS decisions only, + # assume that every eligible person claims WIC. + # This is an explicit modelling assumption, not observed participation. + # Replace the dataset with a corrected Microcosm release and DELETE + # this exception immediately after that replacement is certified. + # Replacement tracking: https://github.com/PolicyEngine/microcosm/issues/1154 + # Prepared removal PR: https://github.com/PolicyEngine/policyengine.py/pull/563 + return stored.mask(missing, True) + + def _live_values(stored: pd.Series, variable: Any, context: str) -> np.ndarray: if stored.isna().any(): raise ValueError(f"{context}: the stored column has missing values.") @@ -356,6 +391,15 @@ def _live_values(stored: pd.Series, variable: Any, context: str) -> np.ndarray: return np.asarray(stored.to_numpy(dtype=bool)) if pd.api.types.is_numeric_dtype(stored.dtype) and stored.isin((0, 1)).all(): return np.asarray(stored.to_numpy(), dtype=bool) + # Nullable legacy H5 columns may contain genuine booleans and numeric 0/1 + # as objects. Validate each value before casting: bool("False") is True. + if stored.map( + lambda value: ( + isinstance(value, bool | np.bool_ | int | np.integer | float | np.floating) + and value in (0, 1) + ) + ).all(): + return np.asarray(stored.to_numpy(), dtype=bool) raise ValueError(f"{context}: the stored values are not boolean.") diff --git a/src/policyengine/tax_benefit_models/us/state_preparation.py b/src/policyengine/tax_benefit_models/us/state_preparation.py new file mode 100644 index 00000000..eb0ae166 --- /dev/null +++ b/src/policyengine/tax_benefit_models/us/state_preparation.py @@ -0,0 +1,350 @@ +"""Verified state derivatives of one certified source, not new datasets. + +This opt-in prototype is not selected by ``ensure_datasets`` or request serving. +The caller obtains ``MaterializedDataset`` through the installed bundle resolver. +Neither partitioning nor year preparation permits an unmanaged dataset bypass. +""" + +from __future__ import annotations + +import hashlib +from importlib.metadata import version +from pathlib import Path +from typing import Literal + +import h5py +import pandas as pd +from pydantic import BaseModel, ConfigDict, Field, model_validator + +from policyengine.countries.us.data.states import US_STATE_FIPS +from policyengine.provenance.dataset_materialization import MaterializedDataset +from policyengine.tax_benefit_models.us.datasets import ( + US_ENTITY_KEYS, + US_PERSON_ENTITY_ID_COLUMNS, + PolicyEngineUSDataset, + _prepare_us_year, + _validate_entity_ids, +) +from policyengine.tax_benefit_models.us.legacy_inputs import ( + apply_legacy_input_renames_to_microsimulation, +) +from policyengine.utils.hashing import sha256_file + + +class _StrictModel(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + +class EntityCounts(_StrictModel): + person: int = Field(ge=0) + household: int = Field(ge=0) + tax_unit: int = Field(ge=0) + spm_unit: int = Field(ge=0) + family: int = Field(ge=0) + marital_unit: int = Field(ge=0) + + +class USStatePartition(_StrictModel): + format_version: Literal[1] = 1 + parent: MaterializedDataset + state_code: str + state_fips: int + source_year: int + path: Path + sha256: str = Field(pattern=r"^[0-9a-f]{64}$") + bytes: int = Field(gt=0) + counts: EntityCounts + + @model_validator(mode="after") + def validate_state(self) -> USStatePartition: + if US_STATE_FIPS.get(self.state_code) != self.state_fips: + raise ValueError("State code and FIPS must identify the same US state/DC") + return self + + +class USPartitionManifest(_StrictModel): + format_version: Literal[1] = 1 + source: MaterializedDataset + source_year: int + counts: EntityCounts + partitions: tuple[USStatePartition, ...] + + @model_validator(mode="after") + def validate_coverage(self) -> USPartitionManifest: + if len(self.partitions) != 51 or {p.state_code for p in self.partitions} != set( + US_STATE_FIPS + ): + raise ValueError( + "Partition manifest must contain all 50 states and DC exactly once" + ) + for partition in self.partitions: + if ( + partition.parent != self.source + or partition.source_year != self.source_year + ): + raise ValueError( + "Partitions must have the same certified parent and source period" + ) + for entity in US_ENTITY_KEYS: + if sum(getattr(p.counts, entity) for p in self.partitions) != getattr( + self.counts, entity + ): + raise ValueError( + f"Partition counts do not cover the source {entity} table" + ) + return self + + +class USPreparationIdentity(_StrictModel): + package_versions: dict[str, str] + code_sha256: str = Field(pattern=r"^[0-9a-f]{64}$") + + +class USStateYearArtifact(_StrictModel): + format_version: Literal[1] = 1 + partition: USStatePartition + year: int = Field(gt=0) + path: Path + sha256: str = Field(pattern=r"^[0-9a-f]{64}$") + bytes: int = Field(gt=0) + counts: EntityCounts + identity: USPreparationIdentity + + +def _load_native(path: Path) -> tuple[dict[str, pd.DataFrame], int]: + # Pandas 3 must not re-infer an all-text subset of a stored mixed object + # column as strings. Keep stored types rather than inferring from state rows. + with ( + pd.option_context("future.infer_string", False), + pd.HDFStore(path, "r") as store, + ): + frames = {entity: store[entity] for entity in US_ENTITY_KEYS} + period = store["_time_period"] + if len(period) != 1 or pd.isna(period.iloc[0]): + raise ValueError("Native source must declare one nonmissing _time_period") + year = int(period.iloc[0]) + if year != period.iloc[0]: + raise ValueError("Native source period must be an integer year") + return frames, year + + +def _entity_states(frames: dict[str, pd.DataFrame]) -> dict[str, pd.Series]: + """Validate every membership before deriving group state assignments.""" + _validate_entity_ids(frames) + households = frames["household"] + if "state_fips" not in households: + raise ValueError("Missing source household state_fips") + state = households["state_fips"] + if state.isna().any() or not state.isin(US_STATE_FIPS.values()).all(): + raise ValueError("Source contains missing or unknown US state_fips") + lookup = households.set_index("household_id")["state_fips"] + people = frames["person"] + household_link = US_PERSON_ENTITY_ID_COLUMNS["household"] + if household_link not in people: + raise ValueError(f"Missing {household_link}") + person_states = people[household_link].map(lookup) + if person_states.isna().any(): + raise ValueError("Person has a missing or unknown household link") + if set(people[household_link]) != set(households.household_id): + raise ValueError("Source includes a household without any person") + result = {"household": state, "person": person_states} + for entity in US_ENTITY_KEYS: + if entity in result: + continue + link = US_PERSON_ENTITY_ID_COLUMNS[entity] + group = frames[entity] + if link not in people or people[link].isna().any(): + raise ValueError(f"Missing native person {link}") + if not people[link].isin(group[f"{entity}_id"]).all(): + raise ValueError(f"Unknown native person {link}") + memberships = pd.DataFrame( + {"id": people[link], "state": person_states} + ).drop_duplicates() + if memberships.id.duplicated().any(): + raise ValueError(f"{entity} spans more than one state") + group_states = group[f"{entity}_id"].map(memberships.set_index("id").state) + if group_states.isna().any(): + raise ValueError(f"Source {entity} contains a group without any person") + result[entity] = group_states + return result + + +def _write_partition( + source_path: Path, path: Path, frames: dict[str, pd.DataFrame] +) -> None: + with pd.HDFStore(path, "w") as store: + for entity, frame in frames.items(): + store.put(entity, frame, format="fixed") + # Copy non-entity HDF nodes and root attributes verbatim: periods and producer + # metadata must not be reconstructed or stripped during filtering. + with h5py.File(source_path, "r") as source, h5py.File(path, "a") as output: + for name in source: + if name not in US_ENTITY_KEYS: + source.copy(name, output) + for name, value in source.attrs.items(): + output.attrs[name] = value + with ( + pd.option_context("future.infer_string", False), + pd.HDFStore(path, "r") as store, + ): + for entity, expected in frames.items(): + pd.testing.assert_frame_equal(store[entity], expected) + + +def partition_certified_us_source( + source: MaterializedDataset, output_dir: Path +) -> USPartitionManifest: + """Read a hash-verified native source once and write 51 exact state extracts. + + Files are published only after write/read equality. An error removes only + paths created by this call, never existing cache files or the source. + """ + if sha256_file(source.path) != source.sha256: + raise ValueError( + "Certified source SHA-256 does not match the materialized file" + ) + frames, year = _load_native(source.path) + states = _entity_states(frames) + output_dir.mkdir(parents=True, exist_ok=True) + paths = [output_dir / f"state-{code.lower()}.h5" for code in US_STATE_FIPS] + if any(path.exists() or path.with_suffix(".partial").exists() for path in paths): + raise FileExistsError("Partition output already exists; refusing to overwrite") + created: list[Path] = [] + partitions: list[USStatePartition] = [] + try: + for (code, fips), path in zip(US_STATE_FIPS.items(), paths, strict=True): + subset = { + entity: frame.loc[states[entity] == fips] + for entity, frame in frames.items() + } + partial = path.with_suffix(".partial") + created.append(partial) + _write_partition(source.path, partial, subset) + digest = sha256_file(partial) + partial.rename(path) + created.append(path) + partitions.append( + USStatePartition( + parent=source, + state_code=code, + state_fips=fips, + source_year=year, + path=path, + sha256=digest, + bytes=path.stat().st_size, + counts=EntityCounts( + **{entity: len(frame) for entity, frame in subset.items()} + ), + ) + ) + return USPartitionManifest( + source=source, + source_year=year, + counts=EntityCounts( + **{entity: len(frame) for entity, frame in frames.items()} + ), + partitions=tuple(partitions), + ) + except BaseException: + for path in created: + path.unlink(missing_ok=True) + raise + + +def prepare_us_state_year( + partition: USStatePartition, year: int, output_dir: Path +) -> USStateYearArtifact: + """Run the existing country year preparation on a verified state derivative. + + Input records are passed in memory to the country API, not accepted as an + arbitrary unmanaged source. The helper also used by national preparation + retains the installed bundle's SPM selection and legacy input mapping. + """ + from policyengine_us import Microsimulation + from policyengine_us.data.dataset_schema import USSingleYearDataset + + from policyengine.tax_benefit_models.us.spm import resolve_spm_selection + + if isinstance(year, bool) or not isinstance(year, int) or year <= 0: + raise ValueError("Preparation year must be a positive integer") + if sha256_file(partition.path) != partition.sha256: + raise ValueError("State partition SHA-256 does not match the file") + frames, source_year = _load_native(partition.path) + states = _entity_states(frames) + if source_year != partition.source_year: + raise ValueError("Partition source period does not match its manifest") + if ( + EntityCounts(**{entity: len(frame) for entity, frame in frames.items()}) + != partition.counts + ): + raise ValueError("Partition entity counts do not match its manifest") + if any(not values.eq(partition.state_fips).all() for values in states.values()): + raise ValueError("Partition includes records from a different state") + output_dir.mkdir(parents=True, exist_ok=True) + path = output_dir / f"state-{partition.state_code.lower()}-year-{year}.h5" + partial = path.with_suffix(".partial") + if path.exists() or partial.exists(): + raise FileExistsError(f"Year output already exists: {path}") + inputs = USSingleYearDataset(**frames, time_period=source_year) + sim = Microsimulation(dataset=inputs, spm=resolve_spm_selection()) + renames = apply_legacy_input_renames_to_microsimulation(sim) + dataset = _prepare_us_year( + sim, + year=year, + dataset_stem=f"state-{partition.state_code.lower()}", + filepath=partial, + legacy_input_renames=renames, + ) + try: + dataset.save() + reloaded = PolicyEngineUSDataset( + name=dataset.name, + description=dataset.description, + year=year, + filepath=str(partial), + ) + if dataset.data is None or reloaded.data is None: + raise ValueError("Year preparation did not materialize entity tables") + for entity, expected in dataset.data.entity_data.items(): + pd.testing.assert_frame_equal( + pd.DataFrame(reloaded.data.entity_data[entity]), pd.DataFrame(expected) + ) + if set(expected[f"{entity}_id"]) != set(frames[entity][f"{entity}_id"]): + raise ValueError(f"Year preparation changed {entity} membership") + if reloaded.metadata != dataset.metadata: + raise ValueError("Year output did not preserve its input-rename record") + artifact = USStateYearArtifact( + partition=partition, + year=year, + path=path, + sha256=sha256_file(partial), + bytes=partial.stat().st_size, + counts=EntityCounts( + **{ + entity: len(frame) + for entity, frame in dataset.data.entity_data.items() + } + ), + identity=USPreparationIdentity( + package_versions={ + name: version(name) + for name in ( + "policyengine", + "policyengine-us", + "policyengine-core", + "spm-calculator", + ) + }, + code_sha256=hashlib.sha256( + ( + sha256_file(Path(__file__)) + + sha256_file(Path(__file__).with_name("datasets.py")) + ).encode() + ).hexdigest(), + ), + ) + partial.rename(path) + return artifact + except BaseException: + partial.unlink(missing_ok=True) + raise diff --git a/tests/test_us_legacy_inputs.py b/tests/test_us_legacy_inputs.py index b106d928..6c328a9e 100644 --- a/tests/test_us_legacy_inputs.py +++ b/tests/test_us_legacy_inputs.py @@ -322,6 +322,103 @@ def test_a_draw_that_is_not_a_complete_boolean_is_refused(draw, message): assert simulation.inputs == {} +@pytest.mark.parametrize( + "draw", + [ + [False, np.bool_(True), None, False], + [0, 1.0, np.nan, np.bool_(False)], + pd.array([False, True, pd.NA, False], dtype="boolean"), + ], +) +def test_missing_acs_wic_decisions_assume_claiming_without_changing_stored_draws(draw): + """Only a missing ACS decision uses the temporary 100% take-up assumption.""" + ids = [1, 2, 3, 4] + person = _person( + ids, + draw, + person_support_channel=[ + "asec", + "puf_tax_detail", + "acs_2024_1yr", + "acs_2024_1yr", + ], + ) + original = person.copy(deep=True) + simulation = FakeSimulation(_variables(LIVE), ids) + + assert apply_legacy_input_renames(simulation, _tables({2024: person})) == { + LEGACY: LIVE + } + for month in _months(2024): + assert simulation.inputs[(LIVE, month)].tolist() == [False, True, True, False] + assert simulation.inputs[(LIVE, month)].dtype == bool + pd.testing.assert_frame_equal(person, original) + + +@pytest.mark.parametrize("channel", ["asec", "puf_tax_detail", "acs", "unknown", None]) +def test_missing_wic_decisions_outside_identified_acs_people_are_refused(channel): + person = _person( + [1, 2], [None, None], person_support_channel=["acs_2024_1yr", channel] + ) + simulation = FakeSimulation(_variables(LIVE), [1, 2]) + with pytest.raises(ValueError, match="missing values"): + apply_legacy_input_renames(simulation, _tables({2024: person})) + assert simulation.inputs == {} + + +@pytest.mark.parametrize("invalid", ["False", "True", "0", 2, -1, np.inf]) +def test_acs_missing_value_does_not_hide_an_invalid_stored_wic_decision(invalid): + person = _person( + [1, 2], [invalid, None], person_support_channel=["asec", "acs_2024_1yr"] + ) + simulation = FakeSimulation(_variables(LIVE), [1, 2]) + with pytest.raises(ValueError, match="not boolean"): + apply_legacy_input_renames(simulation, _tables({2024: person})) + assert simulation.inputs == {} + + +def test_a_later_invalid_year_prevents_earlier_acs_assignments(): + valid = _person([1], [None], person_support_channel=["acs_2024_1yr"]) + invalid = _person([1], [None], person_support_channel=["asec"]) + simulation = FakeSimulation(_variables(LIVE), [1]) + with pytest.raises(ValueError, match="missing values"): + apply_legacy_input_renames(simulation, _tables({2024: valid, 2025: invalid})) + assert simulation.inputs == {} + + +def test_the_temporary_acs_assumption_reaches_each_existing_branch_idempotently(): + person = _person( + [1, 2], [False, None], person_support_channel=["asec", "acs_2024_1yr"] + ) + dataset = _multi_year_dataset({2024: person}) + simulation = FakeSimulation(_variables(LIVE), [1, 2], dataset=dataset) + baseline = FakeSimulation(_variables(LIVE), [1, 2], dataset=dataset) + simulation.branches["baseline"] = baseline + apply_legacy_input_renames_to_microsimulation(simulation) + before = _snapshot(simulation) + apply_legacy_input_renames_to_microsimulation(simulation) + assert _snapshot(simulation) == _snapshot(baseline) == before + for month in _months(2024): + assert simulation.inputs[(LIVE, month)].tolist() == [False, True] + + +def test_current_name_wic_decisions_are_not_overridden_by_the_acs_assumption(): + person = _person( + [1], [None], person_support_channel=["acs_2024_1yr"], **{LIVE: [False]} + ) + simulation = FakeSimulation(_variables(LIVE), [1]) + assert apply_legacy_input_renames(simulation, _tables({2024: person})) == {} + assert simulation.inputs == {} + + +def test_boolean_values_in_an_object_column_are_preserved(): + person = _person([1, 2], pd.Series([False, np.bool_(True)], dtype=object)) + person[LEGACY] = person[LEGACY].astype(object) + simulation = FakeSimulation(_variables(LIVE), [1, 2]) + apply_legacy_input_renames(simulation, _tables({2024: person})) + assert simulation.inputs[(LIVE, "2024-01")].tolist() == [False, True] + + def test_a_non_boolean_live_input_keeps_the_stored_values(): variables = {LIVE: _variable(value_type=float)} simulation = FakeSimulation(variables, [1, 2]) diff --git a/tests/test_us_legacy_inputs_integration.py b/tests/test_us_legacy_inputs_integration.py index a03e755c..68c1a65e 100644 --- a/tests/test_us_legacy_inputs_integration.py +++ b/tests/test_us_legacy_inputs_integration.py @@ -177,6 +177,28 @@ def test_run_keeps_a_stored_false_draw(mapped_run): assert mapped_run.release_bundle["legacy_input_renames"] == RENAME +def test_missing_acs_claim_decisions_do_not_make_ineligible_people_wic_eligible( + tmp_path, +): + frames = _frames() + # Keep the donor infant's False and the ACS toddler's existing True. + # Only the ACS adult is missing a participation decision. + frames["person"][LEGACY] = pd.Series([None, False, True], dtype=object) + frames["person"]["person_support_channel"] = [ + "acs_2024_1yr", + "asec", + "acs_2024_1yr", + ] + person = _person_outputs(_run(_in_memory_dataset(tmp_path, frames))) + + assert person[LIVE].tolist() == [True, False, True] + assert not person["is_wic_eligible"].iloc[0] + assert person["wic"].iloc[0] == 0 + assert person["is_wic_eligible"].iloc[INFANT] + assert person["wic"].iloc[INFANT] == 0 + assert person["wic"].iloc[TODDLER] > 0 + + def test_run_over_a_core_h5_keeps_a_stored_false_draw(tmp_path): """Path 1 over a policyengine-core ``variable/period`` file. diff --git a/tests/test_us_state_partitions.py b/tests/test_us_state_partitions.py new file mode 100644 index 00000000..a282d7bd --- /dev/null +++ b/tests/test_us_state_partitions.py @@ -0,0 +1,106 @@ +"""Partition correctness on all six native entities, without network access.""" + +import pandas as pd +import pytest + +from policyengine.tax_benefit_models.us.datasets import US_ENTITY_KEYS +from policyengine.tax_benefit_models.us.state_preparation import ( + USPartitionManifest, + _load_native, + partition_certified_us_source, +) +from policyengine.utils.hashing import sha256_file +from tests.us_partition_fixtures import write_source + + +def test_complete_disjoint_round_trip_preserves_native_inputs(tmp_path): + source = write_source(tmp_path / "source.h5", include_nulls=True) + manifest = partition_certified_us_source(source, tmp_path / "states") + assert len(manifest.partitions) == 51 + assert ( + USPartitionManifest.model_validate_json(manifest.model_dump_json()) == manifest + ) + for entity in US_ENTITY_KEYS: + combined = pd.concat( + [_load_native(part.path)[0][entity] for part in manifest.partitions] + ) + original = _load_native(source.path)[0][entity] + id_column = f"{entity}_id" + pd.testing.assert_frame_equal( + combined.sort_values(id_column), original.sort_values(id_column) + ) + assert sha256_file(source.path) == source.sha256 + for partition in manifest.partitions: + assert pd.read_hdf(partition.path, "_time_period").iloc[0] == 2024 + assert pd.read_hdf(partition.path, "_fixture_metadata").iloc[0] == "preserve-me" + assert sha256_file(partition.path) == partition.sha256 + + +@pytest.mark.parametrize( + "failure", + [ + "hash", + "unknown_state", + "missing_state", + "dangling_link", + "cross_state_group", + "missing_person", + ], +) +def test_invalid_source_does_not_publish_partial_files(tmp_path, failure): + source = write_source(tmp_path / "source.h5") + with pd.HDFStore(source.path, "a") as store: + if failure == "hash": + frame = store["household"] + frame["household_weight"] *= 2 + store["household"] = frame + if failure == "unknown_state": + frame = store["household"] + frame.loc[frame.index[0], "state_fips"] = 99 + store["household"] = frame + elif failure == "missing_state": + frame = store["household"] + frame.loc[frame.index[0], "state_fips"] = None + store["household"] = frame + elif failure in {"dangling_link", "cross_state_group", "missing_person"}: + frame = store["person"] + if failure == "missing_person": + frame = frame.iloc[1:] + else: + frame.loc[frame.index[0], "person_tax_unit_id"] = ( + 999999 + if failure == "dangling_link" + else frame.iloc[1].person_tax_unit_id + ) + store["person"] = frame + if failure != "hash": + source = source.model_copy(update={"sha256": sha256_file(source.path)}) + with pytest.raises(ValueError): + partition_certified_us_source(source, tmp_path / "states") + assert not list((tmp_path / "states").glob("*.h5")) + assert not list((tmp_path / "states").glob("*.partial")) + + +def test_write_failure_cleans_only_its_own_files(tmp_path, monkeypatch): + import policyengine.tax_benefit_models.us.state_preparation as module + + source = write_source(tmp_path / "source.h5") + output = tmp_path / "states" + output.mkdir() + sentinel = output / "unrelated.txt" + sentinel.write_text("keep") + original = module._write_partition + calls = 0 + + def fail_second(*args, **kwargs): + nonlocal calls + calls += 1 + if calls == 2: + raise OSError("fixture disk failure") + return original(*args, **kwargs) + + monkeypatch.setattr(module, "_write_partition", fail_second) + with pytest.raises(OSError, match="fixture disk failure"): + partition_certified_us_source(source, output) + assert sentinel.read_text() == "keep" + assert not list(output.glob("*.h5")) diff --git a/tests/test_us_state_year_preparation.py b/tests/test_us_state_year_preparation.py new file mode 100644 index 00000000..9917f90d --- /dev/null +++ b/tests/test_us_state_year_preparation.py @@ -0,0 +1,132 @@ +"""Real, tiny country-model preparation agrees with national-then-filter.""" + +import pandas as pd +import pytest + +from policyengine.provenance.dataset_materialization import DatasetSource +from policyengine.tax_benefit_models.us import datasets as dataset_module +from policyengine.tax_benefit_models.us.datasets import US_ENTITY_KEYS, create_datasets +from policyengine.tax_benefit_models.us.legacy_inputs import ( + RENAMES_RECORD_KEY, + read_renames_record, +) +from policyengine.tax_benefit_models.us.state_preparation import ( + USStateYearArtifact, + partition_certified_us_source, + prepare_us_state_year, +) +from tests.us_partition_fixtures import write_source + + +@pytest.fixture( + scope="module", params=[False, True], ids=["stored-wic", "missing-acs-wic"] +) +def prepared(tmp_path_factory, request): + directory = tmp_path_factory.mktemp("state_year") + source = write_source(directory / "source.h5", acs_wic_gaps=request.param) + manifest = partition_certified_us_source(source, directory / "states") + return directory, source, manifest + + +@pytest.mark.parametrize("year", [2025, 2026, 2027]) +def test_state_first_equals_national_first_for_every_entity( + prepared, monkeypatch, year +): + directory, source, manifest = prepared + monkeypatch.setattr( + dataset_module, + "materialize_dataset", + lambda *args, **kwargs: DatasetSource( + source_uri=source.source_uri, + path=str(source.path), + bundle_dataset=source, + ), + ) + national = next( + iter( + create_datasets( + years=[year], data_folder=str(directory / "national") + ).values() + ) + ) + for state in ("CA", "UT"): + partition = next(p for p in manifest.partitions if p.state_code == state) + artifact = prepare_us_state_year(partition, year, directory / "prepared") + assert ( + USStateYearArtifact.model_validate_json(artifact.model_dump_json()) + == artifact + ) + people = pd.DataFrame(national.data.person) + households = pd.DataFrame(national.data.household) + household_ids = households.loc[ + households.state_fips == partition.state_fips, "household_id" + ] + selected_people = people.loc[people.person_household_id.isin(household_ids)] + source_people = pd.read_hdf(partition.path, "person") + if source_people.would_claim_wic.isna().any(): + # Source records remain missing; both preparation paths apply the + # explicitly temporary ACS assumption before materializing outputs. + assert source_people.person_support_channel.eq("acs_2024_1yr").all() + assert selected_people.takes_up_wic_if_eligible.all() + for entity in US_ENTITY_KEYS: + frame = pd.DataFrame(national.data.entity_data[entity]) + ids = ( + selected_people.person_id + if entity == "person" + else selected_people[f"person_{entity}_id"] + ) + expected = ( + frame.loc[frame[f"{entity}_id"].isin(ids)] + .sort_values(f"{entity}_id") + .reset_index(drop=True) + ) + actual = ( + pd.read_hdf(artifact.path, entity) + .sort_values(f"{entity}_id") + .reset_index(drop=True) + ) + pd.testing.assert_frame_equal(actual, expected) + assert artifact.identity.package_versions["policyengine-us"] + assert read_renames_record(artifact.path) == { + "would_claim_wic": "takes_up_wic_if_eligible" + } + assert national.metadata[RENAMES_RECORD_KEY] == read_renames_record( + artifact.path + ) + + +def test_partition_tampering_is_rejected_before_preparation(prepared): + directory, _, manifest = prepared + bad = manifest.partitions[0].model_copy(update={"sha256": "0" * 64}) + with pytest.raises(ValueError, match="SHA-256"): + prepare_us_state_year(bad, 2025, directory / "tampered") + assert not (directory / "tampered").exists() + + +def test_existing_output_is_not_overwritten(prepared): + directory, _, manifest = prepared + partition = manifest.partitions[1] + prepare_us_state_year(partition, 2025, directory / "no_overwrite") + with pytest.raises(FileExistsError): + prepare_us_state_year(partition, 2025, directory / "no_overwrite") + + +def test_failed_year_write_removes_partial_file_and_preserves_input( + prepared, monkeypatch +): + from pathlib import Path + + from policyengine.utils.hashing import sha256_file + + directory, _, manifest = prepared + partition = manifest.partitions[0] + + def fail_save(dataset): + Path(dataset.filepath).write_bytes(b"fixture incomplete write") + raise OSError("fixture write failure") + + monkeypatch.setattr(dataset_module.PolicyEngineUSDataset, "save", fail_save) + with pytest.raises(OSError, match="fixture write failure"): + prepare_us_state_year(partition, 2025, directory / "failed_year") + assert not list((directory / "failed_year").iterdir()) + assert sha256_file(partition.path) == partition.sha256 diff --git a/tests/us_partition_fixtures.py b/tests/us_partition_fixtures.py new file mode 100644 index 00000000..5cd35224 --- /dev/null +++ b/tests/us_partition_fixtures.py @@ -0,0 +1,57 @@ +"""A synthetic native US input file; never a substitute for certified data.""" + +from pathlib import Path + +import pandas as pd + +from policyengine.countries.us.data.states import US_STATE_FIPS +from policyengine.provenance.dataset_materialization import MaterializedDataset +from policyengine.tax_benefit_models.us.datasets import US_ENTITY_KEYS +from policyengine.utils.hashing import sha256_file + + +def write_source( + path: Path, *, include_nulls: bool = False, acs_wic_gaps: bool = False +) -> MaterializedDataset: + states = list(US_STATE_FIPS.items()) + ids = list(range(100, 100 + len(states))) + frames = {entity: pd.DataFrame({f"{entity}_id": ids}) for entity in US_ENTITY_KEYS} + person = frames["person"] + for entity in US_ENTITY_KEYS: + if entity != "person": + person[f"person_{entity}_id"] = ids + person["age"] = 40 + person["employment_income"] = 40000.0 + person["is_tax_unit_head"] = True + # Certified inputs still use this legacy name. Both preparation paths must + # preserve the draw under its current model name, including later years. + person["would_claim_wic"] = [index % 2 == 0 for index in range(len(ids))] + if acs_wic_gaps: + person["person_support_channel"] = "asec" + person["would_claim_wic"] = person["would_claim_wic"].astype(object) + for index, (code, _) in enumerate(states): + if code in ("CA", "UT"): + person.loc[index, "person_support_channel"] = "acs_2024_1yr" + person.loc[index, "would_claim_wic"] = None + household = frames["household"] + household["household_weight"] = [float(index + 1) for index in range(len(ids))] + household["state_fips"] = [fips for _, fips in states] + household["congressional_district_geoid"] = [fips * 100 + 1 for _, fips in states] + if include_nulls: + household["fixture_nullable_float"] = [None] + [1.5] * (len(ids) - 1) + household["fixture_mixed_object"] = pd.Series( + [None, 7, "text"] + ["x"] * (len(ids) - 3), dtype=object + ) + with pd.HDFStore(path, "w") as store: + for entity, frame in frames.items(): + store.put(entity, frame.sample(frac=1, random_state=42), format="fixed") + store.put("_time_period", pd.Series([2024]), format="fixed") + store.put("_fixture_metadata", pd.Series(["preserve-me"]), format="fixed") + return MaterializedDataset( + data_package_name="synthetic-fixture", + repo_type="dataset", + revision="fixture", + source_uri="fixture://us-inputs", + sha256=sha256_file(path), + path=path, + )