diff --git a/changelog.d/523.fixed.md b/changelog.d/523.fixed.md new file mode 100644 index 00000000..7cd81a4e --- /dev/null +++ b/changelog.d/523.fixed.md @@ -0,0 +1 @@ +Load certified US annual projection files for the requested calculation years, require complete annual coverage from each family's base, retain earlier annual inputs for formula lookbacks, and reject periods outside the selected coverage. Separate source and derived input caches by artifact, model, and SPM identity. diff --git a/docs/engineering/skills/data-certification.md b/docs/engineering/skills/data-certification.md index f5319517..8af0171a 100644 --- a/docs/engineering/skills/data-certification.md +++ b/docs/engineering/skills/data-certification.md @@ -104,3 +104,93 @@ before they can be updated through this path. The retired `policyengine-bundles` flow (candidates → generated bundle → archive import) is preserved read-only in that repo's history; bundles 4.15.x–4.16.x remain the historical record of earlier certifications. + +## Annual national artifact families + +A producer can advertise exact annual inputs through release metadata: + +```json +"metadata": { + "dataset_years": { + "populace_us_2024": { + "2024": "populace_us_2024", + "2025": "populace_us_2025" + } + } +} +``` + +Each value names an ordinary `artifacts` entry with an H5 path, explicit revision +and SHA256. Each family must map its earliest year to the family/base artifact +and include every year through its projection horizon. Certification rejects +gaps and omitted bases so formula lookbacks retain the complete input history. +Do not use path templates or implicit extension for missing years. Keep the source base +release and content identity in producer provenance separately. Certification +validates this mapping and copies it into `data_releases.us.dataset_years`. +It does not invent artifact pins or certify an unpublished candidate. + +Annual files use the existing single-year entity tables plus `_time_period`. +The runtime requires that stored year to match the selected manifest year. +`pe.us.ensure_datasets(years=[2025])` fetches only the requested annual file and +loads its native inputs; it does not calculate or uprate those inputs again. +Return keys retain the requested family name, such as +`populace_us_2024_2025`. Explicit annual artifact names work too, with coverage +limited to that artifact's year. Families without this metadata retain engine +extension behavior. + +`pe.us.managed_microsimulation(years=[2025, 2030])` permits external calculations +for those two years. Without `years`, the wrapper selects the years in an explicit +`default_calculation_period`, or the current calendar year. With `years` and no +explicit default period, calculations default to the first selected year. The +wrapper checks coverage before downloading any files and preserves the chosen +default after the engine's initialization. External periods outside the selected +years raise `ValueError`, including years available only for internal lookbacks. +Legacy families without annual metadata do not accept `years`; their existing +country-model period behavior remains unchanged. + +The country `USMultiYearDataset` receives the selected years and every advertised +earlier year through the last selected year, with no future inputs. This history +prefix supplies actual prior-year income for Medicare IRMAA, state tax provisions, +and recursive employment-income formulas. Selecting an individual annual artifact +also loads its parent family's earlier inputs. Internal formula lookbacks keep +their normal behavior; the wrapper supplies no history before the family's first +year, so earlier lookbacks retain the country's existing assumptions. The loader +checks schema, row counts and IDs across files, and preserves each year's source +URI, revision and digest in `policyengine_bundle["annual_datasets"]`. + +The ordinary `Simulation(dataset=ensure_datasets(...)[...])` route uses the same +annual history. `ensure_datasets` records the required source references in the +dataset's JSON metadata without fetching them; `Simulation.run()` fetches earlier +files when needed. Both the baseline and reform use the country multi-year loader, +including its income-response normalization. Regional runs retain the same current +household IDs in every earlier input year. The wrapper rejects missing or changed +history pins, and the output records those references in `annual_input_sources`. + +For a 2024–2035 family in calendar year 2026, default construction loads three +files (2024–2026). Selecting 2035 explicitly loads all 12 files. A local candidate +built during this integration measured 1,754,538,463 bytes (1.755 GB) for the +three-year prefix and 5,928,824,315 bytes (5.929 GB) for all 12 files. Its base +file occupied 826,917,837 bytes; each projected file occupied about 463.81 MB. +These measurements describe candidate artifacts; release certification remains a +separate step. `annual_selected_years`, `annual_loaded_years`, +`annual_input_bytes_by_year`, and `annual_input_bytes` record the actual selection +and file sizes. The total describes input storage and the maximum download size; +cached files do not require another download. Memory also includes decoded tables +and engine arrays, so production adoption still needs a population-scale memory +check. The wrapper performs one full entity load per annual file; its preliminary +`_time_period` check reads only that small entry. `ensure_datasets` materializes +only the requested files and does not load this history prefix. + +Source caches live beneath `.policyengine/sources/` and include the artifact's +repository, path, revision, content hash and metadata hash. Derived input caches +beneath `.policyengine/derived/` also include the installed model source/version, +core/wrapper versions, SPM selection and requested year. Legacy basename-only +caches cannot establish that identity and are not reused. Loading never rewrites +annual native files; no model-specific derived input file is necessary. +Annual output cache keys also include the selected dataset, frozen history pins, +and runtime identity, even when a caller reuses an explicit simulation ID. + +US regional analysis keeps row filters over the annual national dataset. +Positional weight replacement cannot establish annual year/ID alignment, so the +wrapper rejects it for annual inputs. A separate certified alignment contract +would be necessary before supporting such an overlay. diff --git a/src/policyengine/core/simulation.py b/src/policyengine/core/simulation.py index f337a7f2..438885fd 100644 --- a/src/policyengine/core/simulation.py +++ b/src/policyengine/core/simulation.py @@ -125,11 +125,26 @@ def spm_provenance(self) -> Optional[dict[str, Any]]: @property def storage_id(self) -> str: """Include resolved SPM settings in cache and saved-result identity.""" + identity = self.id + history = getattr(self.dataset, "metadata", {}).get("annual_input_sources") + if history: + from policyengine.tax_benefit_models.us.datasets import ( + _derived_cache_identity, + ) + + annual = { + "dataset_id": self.dataset.id, + "year": self.dataset.year, + "history": history, + "runtime": _derived_cache_identity(), + } + encoded = json.dumps(annual, sort_keys=True, separators=(",", ":")).encode() + identity += f"-annual-{hashlib.sha256(encoded).hexdigest()}" config = self.spm_config if config is None: - return self.id + return identity encoded = json.dumps(config, sort_keys=True, separators=(",", ":")).encode() - return f"{self.id}-spm-{hashlib.sha256(encoded).hexdigest()}" + return f"{identity}-spm-{hashlib.sha256(encoded).hexdigest()}" @model_validator(mode="after") def _compile_dict_reforms(self) -> "Simulation": diff --git a/src/policyengine/provenance/certification.py b/src/policyengine/provenance/certification.py index 19d08ff0..527c3e02 100644 --- a/src/policyengine/provenance/certification.py +++ b/src/policyengine/provenance/certification.py @@ -48,6 +48,7 @@ from policyengine.provenance.manifest import ( HF_REQUEST_TIMEOUT_SECONDS, + CountryReleaseManifest, DataReleaseManifest, _specifier_matches, fetch_pypi_wheel_metadata, @@ -561,7 +562,7 @@ def build_country_manifest_payload( model_build.data_build_fingerprint ) - return { + payload = { "schema_version": 1, "bundle_id": f"{country}-{policyengine_version}", "country_id": country, @@ -583,10 +584,19 @@ def build_country_manifest_payload( "default_dataset": default_dataset, "datasets": datasets, "region_datasets": region_datasets, + **( + {"dataset_years": manifest.metadata["dataset_years"]} + if "dataset_years" in manifest.metadata + else {} + ), "certified_data_artifact": certified_artifact, "certification": certification, } + # Annual metadata must survive certification only with complete artifact pins. + CountryReleaseManifest.model_validate(payload) + return payload + def build_bundle_data_release_payload( *, diff --git a/src/policyengine/provenance/dataset_materialization.py b/src/policyengine/provenance/dataset_materialization.py index 35824357..83c0778d 100644 --- a/src/policyengine/provenance/dataset_materialization.py +++ b/src/policyengine/provenance/dataset_materialization.py @@ -2,6 +2,8 @@ from __future__ import annotations +import hashlib +import json import os import tempfile from dataclasses import dataclass @@ -16,6 +18,7 @@ from .manifest import ( CountryReleaseManifest, _artifact_revision, + _dataset_for_year, build_hf_uri, dataset_logical_name, get_release_manifest, @@ -104,6 +107,18 @@ def _resolve_bundle_dataset( f"Managed dataset {dataset_name!r} is missing a certified sha256." ) + identity = { + "repo_id": reference.repo_id or country_manifest.data_package.repo_id, + "repo_type": reference.repo_type or country_manifest.data_package.repo_type, + "path": reference.path, + "revision": reference.revision + or _artifact_revision(country_manifest.data_package), + "sha256": reference.sha256, + "metadata_sha256": reference.metadata_sha256, + } + cache_key = hashlib.sha256( + json.dumps(identity, sort_keys=True).encode() + ).hexdigest() return _BundleDatasetSpec( country_id=country_id, dataset=dataset_name, @@ -116,7 +131,11 @@ def _resolve_bundle_dataset( revision=reference.revision or _artifact_revision(country_manifest.data_package), sha256=reference.sha256, - destination=data_dir / Path(reference.path).name, + destination=data_dir + / ".policyengine" + / "sources" + / cache_key + / Path(reference.path).name, metadata_sha256=reference.metadata_sha256, ) @@ -127,10 +146,12 @@ def materialize_dataset( *, allow_unmanaged: bool = False, data_dir: Path = DEFAULT_DATA_DIR, + year: Optional[int] = None, ) -> DatasetSource: """Select a dataset source and return the local file used for calculation.""" manifest = get_release_manifest(country_id) + dataset = _dataset_for_year(manifest, dataset, year) if dataset is None or dataset == manifest.default_dataset_uri: return _use_bundle_dataset( country_id, diff --git a/src/policyengine/provenance/manifest.py b/src/policyengine/provenance/manifest.py index 99ca72be..c1abc0d8 100644 --- a/src/policyengine/provenance/manifest.py +++ b/src/policyengine/provenance/manifest.py @@ -8,7 +8,7 @@ from urllib.parse import quote import requests -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, model_validator HF_REQUEST_TIMEOUT_SECONDS = 30 PYPI_REQUEST_TIMEOUT_SECONDS = 30 @@ -171,11 +171,52 @@ class CountryReleaseManifest(BaseModel): default_dataset: str datasets: dict[str, ArtifactPathReference] = Field(default_factory=dict) region_datasets: dict[str, ArtifactPathTemplate] = Field(default_factory=dict) + dataset_years: dict[str, dict[int, str]] = Field(default_factory=dict) certified_data_artifact: Optional[CertifiedDataArtifact] = None certification: Optional[DataCertification] = None source_sha256: Optional[str] = Field(default=None, exclude=True) """Byte sha256 of the bundled manifest before Pydantic normalization.""" + @model_validator(mode="after") + def validate_dataset_years(self): + if self.dataset_years and self.country_id != "us": + raise ValueError( + "Annual native dataset families currently support only the US" + ) + for family, years in self.dataset_years.items(): + if family not in self.datasets or not years: + raise ValueError( + f"Annual dataset family {family!r} must name a dataset and contain years" + ) + for year, name in years.items(): + reference = self.datasets.get(name) + if year < 1 or reference is None: + raise ValueError( + f"Unknown annual dataset {name!r} for {family!r}, year {year}" + ) + if not reference.path.endswith(".h5") or not reference.revision: + raise ValueError( + f"Annual dataset {name!r} requires an H5 path and explicit revision" + ) + digest = reference.sha256 or "" + if len(digest) != 64 or any( + c not in "0123456789abcdef" for c in digest + ): + raise ValueError( + f"Annual dataset {name!r} requires a SHA256 digest" + ) + first_year = min(years) + if len(years) != max(years) - first_year + 1: + raise ValueError( + f"Annual dataset family {family!r} must contain contiguous years" + ) + if years[first_year] != family: + raise ValueError( + f"The earliest year of annual dataset family {family!r} " + "must map to the family dataset" + ) + return self + @property def default_dataset_uri(self) -> str: if ( @@ -516,6 +557,54 @@ def certify_data_release_compatibility( ) +def _dataset_years( + manifest: CountryReleaseManifest, + dataset: Optional[str], + *, + include_history: bool = False, +) -> dict[int, str]: + """Return calculation coverage, or the parent family's input history.""" + name = dataset or manifest.default_dataset + for artifact, reference in manifest.datasets.items(): + if name == build_hf_uri( + reference.repo_id or manifest.data_package.repo_id, + reference.path, + reference.revision or _artifact_revision(manifest.data_package), + ): + name = artifact + break + if name in manifest.dataset_years: + return manifest.dataset_years[name] + families = [ + years for years in manifest.dataset_years.values() if name in years.values() + ] + if len(families) > 1: + raise ValueError(f"Annual artifact {name!r} belongs to multiple families") + if not families: + return {} + if include_history: + return families[0] + return { + year: artifact for year, artifact in families[0].items() if artifact == name + } + + +def _dataset_for_year( + manifest: CountryReleaseManifest, dataset: Optional[str], year: Optional[int] +) -> Optional[str]: + years = _dataset_years(manifest, dataset) + if not years: + return dataset + selected = min(years) if year is None else year + if selected not in years: + raise ValueError( + f"Year {selected} is outside annual dataset coverage for " + f"{dataset or manifest.default_dataset!r}: {sorted(years)}. " + "The managed family cannot fall back to engine uprating." + ) + return years[selected] + + def resolve_dataset_reference(country_id: str, dataset: str) -> str: if "://" in dataset: return dataset diff --git a/src/policyengine/tax_benefit_models/common/model_version.py b/src/policyengine/tax_benefit_models/common/model_version.py index 1d6982b9..a5282428 100644 --- a/src/policyengine/tax_benefit_models/common/model_version.py +++ b/src/policyengine/tax_benefit_models/common/model_version.py @@ -367,10 +367,25 @@ def save(self, simulation: Simulation) -> None: raise ValueError( "SPM settings changed since this output was calculated; run again before saving" ) + annual_sources = getattr(simulation.dataset, "metadata", {}).get( + "annual_input_sources" + ) + if ( + simulation.output_dataset.metadata.get("annual_input_sources") + != annual_sources + ): + raise ValueError( + "Annual input pins changed since this output was calculated; run again before saving" + ) serialized_spm = json.dumps( { "config": simulation.spm_config, "provenance": receipt.model_dump(mode="json"), + **( + {"annual_input_sources": annual_sources} + if annual_sources + else {} + ), }, sort_keys=True, ) @@ -413,6 +428,11 @@ def load(self, simulation: Simulation) -> None: recorded = json.loads(raw) if recorded["config"] != simulation.spm_config: raise ValueError("Saved US simulation uses different SPM settings") + annual_sources = getattr(simulation.dataset, "metadata", {}).get( + "annual_input_sources" + ) + if recorded.get("annual_input_sources") != annual_sources: + raise ValueError("Saved US simulation uses different annual input pins") receipt = SPMProvenance.model_validate(recorded["provenance"]) simulation.output_dataset = self._dataset_class( @@ -429,6 +449,10 @@ def load(self, simulation: Simulation) -> None: simulation.spm_receipt = receipt simulation.spm = SPMSelection.model_validate(recorded["config"]) simulation.output_dataset.metadata["spm_config"] = recorded["config"] + if recorded.get("annual_input_sources"): + simulation.output_dataset.metadata["annual_input_sources"] = recorded[ + "annual_input_sources" + ] if os.path.exists(filepath): simulation.created_at = datetime.datetime.fromtimestamp( diff --git a/src/policyengine/tax_benefit_models/us/datasets.py b/src/policyengine/tax_benefit_models/us/datasets.py index 2ff4aa47..312de7fd 100644 --- a/src/policyengine/tax_benefit_models/us/datasets.py +++ b/src/policyengine/tax_benefit_models/us/datasets.py @@ -13,13 +13,19 @@ from policyengine.core import Dataset, YearData from policyengine.provenance.dataset_materialization import ( + DatasetSource, MaterializedDataset, materialize_dataset, ) from policyengine.provenance.manifest import ( + _artifact_revision, + _dataset_for_year, + _dataset_years, + build_hf_uri, dataset_logical_name, get_release_manifest, resolve_dataset_reference, + resolve_managed_dataset_reference, ) from policyengine.tax_benefit_models.common.model_version import ( build_runtime_dataset_provenance, @@ -322,6 +328,130 @@ def _load_policyengine_core_h5(path: Path, year: int) -> USYearData: ) +def _validate_annual_file_year(path: str, year: int) -> None: + """Reject a mislabeled producer artifact before any engine interpretation.""" + with pd.HDFStore(path, mode="r") as store: + if "_time_period" not in store: + raise ValueError(f"Annual dataset {path} requires _time_period={year}") + stored = store["_time_period"] + if len(stored) != 1 or stored.iloc[0] != year: + raise ValueError( + f"Annual dataset {path} has _time_period={stored.tolist()}, expected {year}" + ) + + +def _annual_dataset( + source: DatasetSource, + family: str, + year: int, + input_sources: Optional[dict] = None, +) -> PolicyEngineUSDataset: + _validate_annual_file_year(source.path, year) + identity = hashlib.sha256( + json.dumps( + { + "source": source.source_uri, + "sha256": source.bundle_dataset.sha256 + if source.bundle_dataset is not None + else sha256_file(Path(source.path)), + "year": year, + "annual_input_sources": input_sources or {}, + }, + sort_keys=True, + ).encode() + ).hexdigest() + return PolicyEngineUSDataset( + id=f"{family}_{year}_{identity}", + name=f"{family}-year-{year}", + description=f"Certified annual US inputs for {year} from {family}", + filepath=source.path, + year=year, + metadata={ + "annual_dataset_family": family, + "annual_input_sources": input_sources or {}, + **build_runtime_dataset_provenance( + source.source_uri, source.path, source.bundle_dataset + ), + }, + ) + + +def _annual_input_sources(manifest, dataset: Optional[str], year: int) -> dict: + """Freeze the history pins without fetching any earlier inputs.""" + return { + str(input_year): { + "dataset": artifact, + "uri": build_hf_uri( + reference.repo_id or manifest.data_package.repo_id, + reference.path, + reference.revision or _artifact_revision(manifest.data_package), + ), + "sha256": reference.sha256, + } + for input_year, artifact in _dataset_years( + manifest, dataset, include_history=True + ).items() + if input_year <= year + for reference in [manifest.datasets[artifact]] + } + + +def _validate_annual_tables(reference: dict, annual: dict, year: int) -> None: + for entity, baseline_table in reference.items(): + baseline_table = pd.DataFrame(baseline_table) + annual_table = pd.DataFrame(annual[entity]) + identifiers = [name for name in baseline_table if name.endswith("_id")] + if ( + not baseline_table.columns.equals(annual_table.columns) + or len(baseline_table) != len(annual_table) + or not baseline_table[identifiers].equals(annual_table[identifiers]) + ): + raise ValueError( + f"Annual {year} {entity} schema, rows or IDs differ from the source year" + ) + + +def _derived_cache_identity() -> dict: + from .spm import resolve_spm_selection + + return { + "format": 1, + "model": _runtime_policyengine_us_metadata(), + "core": importlib_metadata.version("policyengine-core"), + "wrapper": importlib_metadata.version("policyengine"), + "spm": resolve_spm_selection(), + } + + +def _derived_dataset_path( + dataset: str, year: int, data_folder: str, runtime: dict +) -> Path: + manifest = get_release_manifest("us") + digest = None + for reference in manifest.datasets.values(): + uri = build_hf_uri( + reference.repo_id or manifest.data_package.repo_id, + reference.path, + reference.revision + or manifest.data_package.release_manifest_revision + or manifest.data_package.version, + ) + if uri == dataset: + digest = reference.sha256 + break + if "://" not in dataset and Path(dataset).is_file(): + digest = sha256_file(Path(dataset)) + identity = {"source": dataset, "sha256": digest, "year": year, **runtime} + key = hashlib.sha256(json.dumps(identity, sort_keys=True).encode()).hexdigest() + return ( + Path(data_folder) + / ".policyengine" + / "derived" + / key + / f"{dataset_logical_name(dataset)}_year_{year}.h5" + ) + + def create_datasets( datasets: Optional[list[str]] = None, years: list[int] = [2024, 2025, 2026, 2027, 2028], @@ -344,7 +474,28 @@ def create_datasets( dataset_requests: list[Optional[str]] = datasets or [None] result = {} + manifest = get_release_manifest("us") + runtime = None for dataset in dataset_requests: + annual_years = _dataset_years(manifest, dataset) + if annual_years: + family = dataset_logical_name(dataset or manifest.default_dataset) + for year in years: + _dataset_for_year(manifest, dataset, year) + source = materialize_dataset( + "us", + dataset, + year=year, + allow_unmanaged=allow_unmanaged, + data_dir=Path(data_folder), + ) + result[f"{family}_{year}"] = _annual_dataset( + source, + family, + year, + input_sources=_annual_input_sources(manifest, dataset, year), + ) + continue source = materialize_dataset( "us", dataset, @@ -355,6 +506,8 @@ def create_datasets( sim = Microsimulation(dataset=source.path, spm=resolve_spm_selection()) for year in years: + if runtime is None: + runtime = _derived_cache_identity() # Get all input variables from the simulation # We'll calculate each input variable for the specified year entity_data = { @@ -490,11 +643,14 @@ def create_datasets( elif entity_name == "tax_unit": tax_unit_df = entity_df + filepath = _derived_dataset_path( + source.source_uri, year, data_folder, runtime + ) us_dataset = PolicyEngineUSDataset( - id=f"{dataset_stem}_year_{year}", + id=f"{dataset_stem}_year_{year}_{filepath.parent.name}", 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", + filepath=str(filepath), year=int(year), data=USYearData( person=MicroDataFrame(person_df, weights="person_weight"), @@ -532,18 +688,31 @@ def load_datasets( """ datasets = datasets or [get_release_manifest("us").default_dataset] result = {} + manifest = get_release_manifest("us") + runtime = None for dataset in datasets: + if _dataset_years(manifest, dataset): + result.update( + create_datasets( + datasets=[dataset], years=years, data_folder=data_folder + ) + ) + continue + if runtime is None: + runtime = _derived_cache_identity() resolved_dataset = resolve_dataset_reference("us", dataset) dataset_stem = dataset_logical_name(resolved_dataset) for year in years: - filepath = f"{data_folder}/{dataset_stem}_year_{year}.h5" + filepath = _derived_dataset_path( + resolved_dataset, year, data_folder, runtime + ) us_dataset = PolicyEngineUSDataset( + id=f"{dataset_stem}_year_{year}_{filepath.parent.name}", name=f"{dataset_stem}-year-{year}", description=f"US Dataset for year {year} based on {dataset_stem}", - filepath=filepath, + filepath=str(filepath), year=year, ) - us_dataset.load() dataset_key = f"{dataset_stem}_{year}" result[dataset_key] = us_dataset @@ -1192,13 +1361,41 @@ def ensure_datasets( """ datasets = datasets or [get_release_manifest("us").default_dataset] - # Check if all dataset files exist + manifest = get_release_manifest("us") + if any(_dataset_years(manifest, dataset) for dataset in datasets): + result = {} + for dataset in datasets: + if _dataset_years(manifest, dataset): + result.update( + create_datasets( + datasets=[dataset], + years=years, + data_folder=data_folder, + allow_unmanaged=allow_unmanaged, + ) + ) + else: + result.update( + ensure_datasets( + datasets=[dataset], + years=years, + data_folder=data_folder, + allow_unmanaged=allow_unmanaged, + ) + ) + return result + + runtime = _derived_cache_identity() + # Check only caches matching the source revision/digest and runtime identity. all_exist = True for dataset in datasets: - resolved_dataset = resolve_dataset_reference("us", dataset) - dataset_stem = dataset_logical_name(resolved_dataset) + resolved_dataset = resolve_managed_dataset_reference( + "us", dataset, allow_unmanaged=allow_unmanaged + ) for year in years: - filepath = Path(f"{data_folder}/{dataset_stem}_year_{year}.h5") + filepath = _derived_dataset_path( + resolved_dataset, year, data_folder, runtime + ) if not filepath.exists(): all_exist = False break diff --git a/src/policyengine/tax_benefit_models/us/model.py b/src/policyengine/tax_benefit_models/us/model.py index 51338b3e..63ae89be 100644 --- a/src/policyengine/tax_benefit_models/us/model.py +++ b/src/policyengine/tax_benefit_models/us/model.py @@ -1,4 +1,7 @@ import datetime +from contextvars import ContextVar +from copy import deepcopy +from pathlib import Path from typing import TYPE_CHECKING, Optional import pandas as pd @@ -8,6 +11,7 @@ from policyengine.provenance.dataset_materialization import ( materialize_dataset, ) +from policyengine.provenance.manifest import _dataset_years, get_release_manifest from policyengine.tax_benefit_models.common import MicrosimulationModelVersion from policyengine.tax_benefit_models.common.model_version import ( build_runtime_dataset_provenance, @@ -15,8 +19,18 @@ from policyengine.tax_benefit_models.common.model_version import ( output_dataset_filepath as _output_dataset_filepath, ) - -from .datasets import PolicyEngineUSDataset, USYearData, _validate_entity_ids +from policyengine.utils.entity_utils import filter_dataset_by_household_ids +from policyengine.utils.hashing import sha256_file + +from .datasets import ( + PolicyEngineUSDataset, + USYearData, + _annual_dataset, + _annual_input_sources, + _validate_annual_file_year, + _validate_annual_tables, + _validate_entity_ids, +) from .spm import ( SPMProvenance, SPMSelection, @@ -162,6 +176,7 @@ def run(self, simulation: "Simulation") -> "Simulation": # autosave removal in datasets.py). if dataset.data is None: dataset.load() + source_dataset = dataset # Validate every entity before constructing either the reform or its # baseline, so no set_input can suppress a canonical SPM formula. @@ -171,6 +186,13 @@ def run(self, simulation: "Simulation") -> "Simulation": # Apply regional scoping if specified if simulation.scoping_strategy: + if ( + dataset.metadata.get("annual_dataset_family") + and simulation.scoping_strategy.strategy_type == "weight_replacement" + ): + raise ValueError( + "Annual national datasets require row filters; positional weight replacement lacks certified year and ID alignment" + ) scoped_data = simulation.scoping_strategy.apply( entity_data=dataset.data.entity_data, group_entities=US_GROUP_ENTITIES, @@ -186,6 +208,7 @@ def run(self, simulation: "Simulation") -> "Simulation": filepath=None, year=dataset.year, is_output_dataset=dataset.is_output_dataset, + metadata=dict(dataset.metadata), data=USYearData( person=scoped_data["person"], marital_unit=scoped_data["marital_unit"], @@ -207,27 +230,43 @@ class InputMicrosimulation(Microsimulation): # default would load another population before we build this one. default_dataset = None - microsim = InputMicrosimulation( - situation={}, - reform=reform_dict, - spm=simulation.spm_config, - default_input_period=dataset.year, - ) + if dataset.metadata.get("annual_dataset_family"): + annual_inputs = _ordinary_annual_inputs(source_dataset, dataset) + microsim = Microsimulation( + dataset=annual_inputs, + reform=reform_dict, + spm=simulation.spm_config, + default_input_period=dataset.year, + ) + microsim.default_calculation_period = dataset.year + if microsim.baseline is not None: + microsim.baseline.default_calculation_period = dataset.year + else: + microsim = InputMicrosimulation( + situation={}, + reform=reform_dict, + spm=simulation.spm_config, + default_input_period=dataset.year, + ) # Use ``microsim.tax_benefit_system``, not the module-level # ``system``: ``Microsimulation.__init__`` applies structural # reforms (e.g. ``gov.contrib.ctc.*``) to its per-sim system but # leaves the module-level one untouched. Building populations # against the module-level system would hide reform-registered # variables like ``ctc_minimum_refundable_amount`` at calc time. - if microsim.baseline is not None: + if ( + not dataset.metadata.get("annual_dataset_family") + and microsim.baseline is not None + ): self._build_simulation_from_dataset( microsim.baseline, dataset, microsim.baseline.tax_benefit_system, ) - self._build_simulation_from_dataset( - microsim, dataset, microsim.tax_benefit_system - ) + if not dataset.metadata.get("annual_dataset_family"): + self._build_simulation_from_dataset( + microsim, dataset, microsim.tax_benefit_system + ) data = { "person": pd.DataFrame(), @@ -316,7 +355,18 @@ class InputMicrosimulation(Microsimulation): filepath=str(_output_dataset_filepath(simulation)), year=simulation.dataset.year, is_output_dataset=True, - metadata={"spm_config": dict(microsim.spm_config)}, + metadata={ + "spm_config": dict(microsim.spm_config), + **( + { + "annual_input_sources": deepcopy( + dataset.metadata["annual_input_sources"] + ) + } + if dataset.metadata.get("annual_dataset_family") + else {} + ), + }, data=USYearData( person=data["person"], marital_unit=data["marital_unit"], @@ -468,9 +518,128 @@ def _build_simulation_from_dataset(self, microsim, dataset, system): microsim.set_input(column, dataset.year, df[column].values) +def _ordinary_annual_inputs(source_dataset, scoped_dataset): + """Use the country's loader and response normalization for every input year.""" + from policyengine_us.data import USMultiYearDataset, USSingleYearDataset + + references = source_dataset.metadata.get("annual_input_sources") + if not references or str(source_dataset.year) not in references: + raise ValueError("Annual simulation requires recorded source and history pins") + if references != _annual_input_sources( + get_release_manifest("us"), + references[str(source_dataset.year)]["dataset"], + source_dataset.year, + ): + raise ValueError("Annual input history differs from the certified family pins") + _validate_entity_ids(scoped_dataset.data.entity_data) + keep_ids = scoped_dataset.data.household["household_id"] + singles = [] + for year_text, reference in sorted( + references.items(), key=lambda item: int(item[0]) + ): + year = int(year_text) + if year > source_dataset.year: + raise ValueError("Annual history cannot include future input years") + if year == source_dataset.year: + frames = scoped_dataset.data.entity_data + else: + source = materialize_dataset("us", reference["uri"], year=year) + digest = ( + source.bundle_dataset.sha256 + if source.bundle_dataset is not None + else sha256_file(Path(source.path)) + ) + if source.source_uri != reference["uri"] or digest != reference["sha256"]: + raise ValueError( + f"Annual {year} input no longer matches recorded history pins" + ) + earlier = _annual_dataset(source, reference["dataset"], year) + _validate_annual_tables( + source_dataset.data.entity_data, earlier.data.entity_data, year + ) + frames = filter_dataset_by_household_ids( + earlier.data.entity_data, US_GROUP_ENTITIES, keep_ids + ) + # Country datasets require native membership names to distinguish + # person links from each group's own ID column when flattening tables. + native = { + entity: pd.DataFrame(frame).copy(deep=False) + for entity, frame in frames.items() + } + person = native["person"] + for entity in US_GROUP_ENTITIES: + bare, link = f"{entity}_id", f"person_{entity}_id" + if bare in person: + person = ( + person.drop(columns=[bare]) + if link in person + else person.rename(columns={bare: link}) + ) + native["person"] = person + singles.append(USSingleYearDataset(**native, time_period=year)) + return USMultiYearDataset(datasets=singles) + + +_annual_calculation_active = ContextVar("annual_calculation_active", default=False) + + +def _validate_annual_period(period, years): + # Use the engine's period parser so monthly and multi-year requests obey + # exactly the same syntax as the country API. This does not calculate policy. + from policyengine_core.periods import period as parse_period + + parsed = parse_period(period) + if parsed.unit == "eternity": + return + if any( + year not in years for year in range(parsed.start.year, parsed.stop.year + 1) + ): + raise ValueError( + f"Requested period {period!r} is outside annual dataset coverage: {list(years)}" + ) + + +class _AnnualDatasetCoverage: + """Guard external requests, preserving formula lookbacks and monthly calls.""" + + def _with_annual_coverage(self, method, variable, period, *args, **kwargs): + if _annual_calculation_active.get(): + return method(variable, period, *args, **kwargs) + _validate_annual_period( + self.default_calculation_period if period is None else period, + self._annual_years, + ) + token = _annual_calculation_active.set(True) + try: + return method(variable, period, *args, **kwargs) + finally: + _annual_calculation_active.reset(token) + + def calculate(self, variable_name, period=None, *args, **kwargs): + return self._with_annual_coverage( + super().calculate, variable_name, period, *args, **kwargs + ) + + def calculate_add(self, variable_name, period=None, *args, **kwargs): + return self._with_annual_coverage( + super().calculate_add, variable_name, period, *args, **kwargs + ) + + def calculate_divide(self, variable_name, period=None, *args, **kwargs): + return self._with_annual_coverage( + super().calculate_divide, variable_name, period, *args, **kwargs + ) + + def calculate_dataframe(self, variable_names, period=None, *args, **kwargs): + return self._with_annual_coverage( + super().calculate_dataframe, variable_names, period, *args, **kwargs + ) + + def managed_microsimulation( *, dataset: Optional[str] = None, + years: Optional[list[int]] = None, allow_unmanaged: bool = False, spm: Optional[SPMSelection] = None, **kwargs, @@ -479,7 +648,13 @@ def managed_microsimulation( By default this enforces the dataset selection from the bundled ``policyengine.py`` release manifest. Arbitrary dataset URIs require - ``allow_unmanaged=True``. + ``allow_unmanaged=True``. For an advertised annual family, ``years`` selects + external calculation coverage. Its default follows ``default_calculation_period`` + or the current calendar year. Load those years and every earlier advertised + input year for formula lookbacks, without extending annual inputs or loading + future years. External calculations outside the selected years raise + ``ValueError``. Downloads and memory scale with this history prefix; use + ``ensure_datasets(years=[...])`` to materialize only selected annual files. """ from policyengine_us import Microsimulation @@ -491,18 +666,103 @@ def managed_microsimulation( ) selection = resolve_spm_selection(spm) - source = materialize_dataset( - "us", - dataset, - allow_unmanaged=allow_unmanaged, - ) - microsim = Microsimulation(dataset=source.path, spm=selection, **kwargs) - microsim.policyengine_bundle = dict(us_latest.release_bundle) + manifest = get_release_manifest("us") + annual_years = _dataset_years(manifest, dataset) + if annual_years: + from policyengine_core.periods import period as parse_period + from policyengine_us.data import USMultiYearDataset, USSingleYearDataset + + default_period = kwargs.get("default_calculation_period") + if years is None: + default_period = ( + datetime.date.today().year if default_period is None else default_period + ) + parsed = parse_period(default_period) + if parsed.unit == "eternity": + raise ValueError( + "Select annual calculation years for an eternity default" + ) + selected_years = list(range(parsed.start.year, parsed.stop.year + 1)) + else: + if not years or any(type(year) is not int or year < 1 for year in years): + raise ValueError( + "years must contain positive integer calculation years" + ) + selected_years = sorted(set(years)) + if default_period is None: + default_period = min(selected_years) + for year in selected_years: + _validate_annual_period(year, annual_years) + _validate_annual_period(default_period, selected_years) + input_years = { + year: artifact + for year, artifact in _dataset_years( + manifest, dataset, include_history=True + ).items() + if year <= max(selected_years) + } + singles = [] + sources = {} + for year in sorted(input_years): + source = materialize_dataset( + "us", input_years[year], year=year, allow_unmanaged=allow_unmanaged + ) + _validate_annual_file_year(source.path, year) + single = USSingleYearDataset(file_path=source.path, time_period=year) + if singles: + _validate_annual_tables( + { + entity: getattr(singles[0], entity) + for entity in singles[0].table_names + }, + {entity: getattr(single, entity) for entity in single.table_names}, + year, + ) + singles.append(single) + sources[year] = source + + class AnnualMicrosimulation(_AnnualDatasetCoverage, Microsimulation): + _annual_years = tuple(selected_years) + + token = _annual_calculation_active.set(True) + try: + microsim = AnnualMicrosimulation( + dataset=USMultiYearDataset(datasets=singles), spm=selection, **kwargs + ) + finally: + _annual_calculation_active.reset(token) + # Core initializes this from the earliest dataset, overriding kwargs. + microsim.default_calculation_period = default_period + if getattr(microsim, "baseline", None) is not None: + microsim.baseline.default_calculation_period = default_period + microsim.policyengine_bundle = dict(us_latest.release_bundle) + microsim.policyengine_bundle["annual_datasets"] = { + str(year): build_runtime_dataset_provenance( + source.source_uri, + source.path, + source.bundle_dataset, + logical_name=input_years[year], + ) + for year, source in sources.items() + } + microsim.policyengine_bundle["annual_selected_years"] = selected_years + microsim.policyengine_bundle["annual_loaded_years"] = sorted(sources) + bytes_by_year = { + str(year): Path(source.path).stat().st_size + for year, source in sources.items() + } + microsim.policyengine_bundle["annual_input_bytes_by_year"] = bytes_by_year + microsim.policyengine_bundle["annual_input_bytes"] = sum(bytes_by_year.values()) + source = sources[min(sources)] + else: + if years is not None: + raise ValueError("years requires a certified annual dataset family") + source = materialize_dataset("us", dataset, allow_unmanaged=allow_unmanaged) + microsim = Microsimulation(dataset=source.path, spm=selection, **kwargs) + microsim.policyengine_bundle = dict(us_latest.release_bundle) microsim.policyengine_bundle.update( build_runtime_dataset_provenance( - source.source_uri, - source.path, - source.bundle_dataset, + source.source_uri, source.path, source.bundle_dataset ) ) microsim.policyengine_bundle["spm"] = dict(microsim.spm_config) diff --git a/tests/test_annual_datasets.py b/tests/test_annual_datasets.py new file mode 100644 index 00000000..1924dd7a --- /dev/null +++ b/tests/test_annual_datasets.py @@ -0,0 +1,782 @@ +"""Annual producer artifacts must replace, rather than extend, base-year inputs.""" + +from unittest.mock import Mock + +import pytest +from pydantic import ValidationError + +import policyengine.provenance.dataset_materialization as materialization +from policyengine.provenance.manifest import CountryReleaseManifest + + +def annual_manifest(): + return CountryReleaseManifest.model_validate( + { + "country_id": "us", + "policyengine_version": "6.0.0", + "model_package": {"name": "policyengine-us", "version": "2.2.1"}, + "data_package": { + "name": "microcosm-data", + "version": "0.1.0", + "repo_id": "policyengine/populace-us", + "repo_type": "dataset", + }, + "default_dataset": "populace_us_2024", + "datasets": { + f"populace_us_{year}": { + "path": f"annual/populace_us_{year}.h5", + "revision": "annual-release", + "sha256": str(year % 10) * 64, + } + for year in (2024, 2025) + }, + "dataset_years": { + "populace_us_2024": { + "2024": "populace_us_2024", + "2025": "populace_us_2025", + } + }, + } + ) + + +def test_annual_mapping_round_trips_and_selects_pinned_year(monkeypatch, tmp_path): + manifest = annual_manifest() + assert manifest.dataset_years["populace_us_2024"][2025] == "populace_us_2025" + assert ( + CountryReleaseManifest.model_validate_json(manifest.model_dump_json()) + == manifest + ) + monkeypatch.setattr( + materialization, "get_release_manifest", lambda country: manifest + ) + use = Mock() + monkeypatch.setattr(materialization, "_use_bundle_dataset", use) + + materialization.materialize_dataset("us", year=2025, data_dir=tmp_path) + + assert use.call_args.args[:2] == ("us", "populace_us_2025") + + +def test_annual_family_rejects_uncovered_year_before_download(monkeypatch): + monkeypatch.setattr( + materialization, "get_release_manifest", lambda country: annual_manifest() + ) + use = Mock() + monkeypatch.setattr(materialization, "_use_bundle_dataset", use) + with pytest.raises(ValueError, match="2026.*coverage"): + materialization.materialize_dataset("us", year=2026) + use.assert_not_called() + + +@pytest.mark.parametrize("change", ["unknown", "no_digest", "no_revision", "not_h5"]) +def test_annual_manifest_rejects_unpinned_or_unknown_artifact(change): + payload = annual_manifest().model_dump(mode="json") + if change == "unknown": + payload["dataset_years"]["populace_us_2024"]["2025"] = "absent" + else: + reference = payload["datasets"]["populace_us_2025"] + reference[ + {"no_digest": "sha256", "no_revision": "revision", "not_h5": "path"}[change] + ] = "other.csv" if change == "not_h5" else None + with pytest.raises(ValidationError): + CountryReleaseManifest.model_validate(payload) + + +@pytest.mark.parametrize( + "years, error", + [ + ( + {2024: "populace_us_2024", 2026: "populace_us_2025"}, + "contiguous", + ), + ({2025: "populace_us_2025"}, "earliest year.*family"), + ( + {2024: "populace_us_2025", 2025: "populace_us_2024"}, + "earliest year.*family", + ), + ], +) +def test_annual_manifest_rejects_incomplete_family_history(years, error): + payload = annual_manifest().model_dump(mode="json") + payload["dataset_years"]["populace_us_2024"] = years + + with pytest.raises(ValidationError, match=error): + CountryReleaseManifest.model_validate(payload) + + +def test_annual_manifest_accepts_base_only_family(): + payload = annual_manifest().model_dump(mode="json") + payload["dataset_years"]["populace_us_2024"] = {2024: "populace_us_2024"} + + manifest = CountryReleaseManifest.model_validate(payload) + + assert manifest.dataset_years == {"populace_us_2024": {2024: "populace_us_2024"}} + + +def test_managed_source_cache_separates_revisions_and_digests(tmp_path): + first = annual_manifest() + second = first.model_copy(deep=True) + second.datasets[second.default_dataset].revision = "another-release" + third = first.model_copy(deep=True) + third.datasets[third.default_dataset].sha256 = "a" * 64 + destinations = { + materialization._resolve_bundle_dataset( + "us", data_dir=tmp_path, manifest=manifest + ).destination + for manifest in (first, second, third) + } + assert len(destinations) == 3 + + +def _annual_source(tmp_path, year): + import pandas as pd + + from policyengine.provenance.dataset_materialization import DatasetSource + from tests.fixtures.filtering_fixtures import create_us_test_dataset + + original = create_us_test_dataset() + original.data.person["employment_income"] = [12345.6789012345] * 6 + original.filepath = str(tmp_path / f"annual-{year}.h5") + original.save() + with pd.HDFStore(original.filepath, "a") as store: + store["_time_period"] = pd.Series([year]) + return original, DatasetSource( + source_uri=f"hf://example/repo/annual-{year}.h5@release", path=original.filepath + ) + + +def test_annual_ensure_fetches_selected_year_and_preserves_exact_inputs( + monkeypatch, tmp_path +): + import pandas as pd + + import policyengine.tax_benefit_models.us.datasets as us_datasets + + original, source = _annual_source(tmp_path, 2025) + monkeypatch.setattr( + us_datasets, "get_release_manifest", lambda country: annual_manifest() + ) + materialize = Mock(return_value=source) + monkeypatch.setattr(us_datasets, "materialize_dataset", materialize) + result = us_datasets.ensure_datasets(years=[2025], data_folder=str(tmp_path)) + assert list(result) == ["populace_us_2024_2025"] + result = result["populace_us_2024_2025"] + assert result.filepath == original.filepath + assert result.year == 2025 + for entity, frame in original.data.entity_data.items(): + pd.testing.assert_frame_equal( + pd.DataFrame(result.data.entity_data[entity]), + pd.DataFrame(frame), + check_exact=True, + ) + assert materialize.call_count == 1 + assert materialize.call_args.kwargs["year"] == 2025 + + +@pytest.mark.parametrize("stored_year", [2024, None]) +def test_annual_ensure_rejects_wrong_or_missing_file_year( + monkeypatch, tmp_path, stored_year +): + import pandas as pd + + import policyengine.tax_benefit_models.us.datasets as us_datasets + + original, source = _annual_source(tmp_path, 2024) + if stored_year is None: + with pd.HDFStore(original.filepath, "a") as store: + del store["_time_period"] + monkeypatch.setattr( + us_datasets, "get_release_manifest", lambda country: annual_manifest() + ) + monkeypatch.setattr(us_datasets, "materialize_dataset", Mock(return_value=source)) + with pytest.raises(ValueError, match="_time_period"): + us_datasets.ensure_datasets(years=[2025], data_folder=str(tmp_path)) + + +def test_derived_cache_identity_includes_model_and_spm(monkeypatch, tmp_path): + import policyengine.tax_benefit_models.us.datasets as us_datasets + + manifest = annual_manifest() + monkeypatch.setattr(us_datasets, "get_release_manifest", lambda country: manifest) + uri = "hf://policyengine/populace-us/annual/populace_us_2024.h5@annual-release" + variants = ( + {"model": "one", "spm": "first"}, + {"model": "two", "spm": "first"}, + {"model": "one", "spm": "second"}, + ) + assert ( + len( + { + us_datasets._derived_dataset_path(uri, 2025, str(tmp_path), runtime) + for runtime in variants + } + ) + == 3 + ) + + +def test_annual_coverage_guards_external_periods_but_permits_formula_lookbacks(): + from policyengine.tax_benefit_models.us.model import _AnnualDatasetCoverage + + class Base: + default_calculation_period = 2025 + + def calculate(self, variable, period=None, **kwargs): + if variable == "lookback": + return self.calculate("previous", 2023) + return period + + def calculate_add(self, variable, period=None, **kwargs): + return self.calculate(variable, period) + + calculate_divide = calculate_add + + class Managed(_AnnualDatasetCoverage, Base): + _annual_years = (2024, 2025) + + sim = Managed() + assert sim.calculate("lookback", 2024) == 2023 + assert sim.calculate("monthly", "2025-01") == "2025-01" + for method in (sim.calculate, sim.calculate_add, sim.calculate_divide): + with pytest.raises(ValueError, match="coverage"): + method("outside", 2026) + with pytest.raises(ValueError, match="coverage"): + method("span", "year:2025:2") + # The context must reset after returning and after a rejected call. + with pytest.raises(ValueError, match="coverage"): + sim.calculate("previous", 2023) + sim.default_calculation_period = 2026 + with pytest.raises(ValueError, match="coverage"): + sim.calculate("implicit") + + +def test_certification_preserves_and_validates_producer_annual_map(): + from policyengine.provenance.certification import build_country_manifest_payload + from policyengine.provenance.manifest import DataReleaseManifest + + country = annual_manifest() + producer = DataReleaseManifest( + schema_version=1, + data_package={"name": "microcosm-data", "version": "0.1.0"}, + default_datasets={"national": country.default_dataset}, + artifacts={ + name: { + "kind": "dataset", + "repo_id": country.data_package.repo_id, + **reference.model_dump(exclude_none=True), + } + for name, reference in country.datasets.items() + }, + metadata={"dataset_years": country.model_dump(mode="json")["dataset_years"]}, + ) + kwargs = dict( + country="us", + manifest=producer, + uri_parts={ + "repo_id": country.data_package.repo_id, + "repo_type": "dataset", + "path": "releases/test/release_manifest.json", + "revision": "release", + }, + policyengine_version="6.0.0", + model_package="policyengine-us", + model_version="2.2.1", + model_wheel={}, + ) + payload = build_country_manifest_payload(**kwargs) + assert ( + CountryReleaseManifest.model_validate(payload).dataset_years + == country.dataset_years + ) + producer.metadata["dataset_years"][country.default_dataset]["2025"] = "missing" + with pytest.raises(ValidationError): + build_country_manifest_payload(**kwargs) + + +def test_managed_annual_uses_exact_multiyear_inputs_and_preserves_explicit_default( + monkeypatch, tmp_path +): + import importlib + import sys + from types import SimpleNamespace + + import pandas as pd + + us_model = importlib.import_module("policyengine.tax_benefit_models.us.model") + originals, sources = {}, {} + for year in (2024, 2025): + originals[year], sources[year] = _annual_source(tmp_path, year) + materialize = Mock(side_effect=lambda *args, **kwargs: sources[kwargs["year"]]) + monkeypatch.setattr( + us_model, "get_release_manifest", lambda country: annual_manifest() + ) + monkeypatch.setattr(us_model, "materialize_dataset", materialize) + + class FakeMicrosimulation: + def __init__(self, dataset, spm, **kwargs): + self.dataset = dataset + self.default_calculation_period = dataset.time_period + self.spm_config = dict(spm) + self.calc = self.calculate + + def calculate(self, variable, period=None, **kwargs): + return period + + monkeypatch.setitem( + sys.modules, + "policyengine_us", + SimpleNamespace(Microsimulation=FakeMicrosimulation), + ) + sim = us_model.managed_microsimulation(default_calculation_period=2025) + assert sim.default_calculation_period == 2025 + assert sim.dataset.years == [2024, 2025] + assert materialize.call_count == 2 + for year, single in sim.dataset.datasets.items(): + pd.testing.assert_frame_equal( + single.person, pd.DataFrame(originals[year].data.person), check_exact=True + ) + assert set(sim.policyengine_bundle["annual_datasets"]) == {"2024", "2025"} + assert sim.policyengine_bundle["annual_input_bytes"] > 0 + assert sim.calc(variable_name="age", period=2025) == 2025 + with pytest.raises(ValueError, match="coverage"): + sim.calc("age", 2026) + materialize.reset_mock() + with pytest.raises(ValueError, match="coverage"): + us_model.managed_microsimulation(default_calculation_period=2026) + materialize.assert_not_called() + + +@pytest.mark.parametrize("dataset", ["populace_us_2024", "populace_us_2025"]) +@pytest.mark.parametrize("regional", [False, True]) +def test_real_managed_microsimulation_uses_each_annual_input_without_extension( + monkeypatch, tmp_path, dataset, regional +): + import hashlib + import importlib + from pathlib import Path + + import pandas as pd + import policyengine_us + from policyengine_us.data import economic_assumptions + + import policyengine.tax_benefit_models.us.datasets as us_datasets + from policyengine.core import Simulation + from policyengine.core.scoping_strategy import RowFilterStrategy + from policyengine.provenance.dataset_materialization import DatasetSource + + us_model = importlib.import_module("policyengine.tax_benefit_models.us.model") + sources = {} + for year, incomes, weights in ( + (2024, [10_000.0, 20_000.0], [3.0, 7.0]), + (2025, [20_000.0, 10_000.0], [5.0, 1.0]), + ): + frames = { + "person": pd.DataFrame( + { + "person_id": [1, 2], + **{ + f"person_{entity}_id": [1, 2] + for entity in ( + "household", + "tax_unit", + "spm_unit", + "family", + "marital_unit", + ) + }, + "age": [40, 70], + "employment_income": incomes, + } + ), + "household": pd.DataFrame( + { + "household_id": [1, 2], + "household_weight": weights, + "state_fips": [39, 39], + "county_fips": ["39049", "39049"], + } + ), + **{ + entity: pd.DataFrame({f"{entity}_id": [1, 2]}) + for entity in ("tax_unit", "spm_unit", "family", "marital_unit") + }, + } + path = tmp_path / f"population-{year}.h5" + with pd.HDFStore(path, "w") as store: + for entity, frame in frames.items(): + store[entity] = frame + store["_time_period"] = pd.Series([year]) + sources[year] = DatasetSource( + source_uri=f"hf://policyengine/populace-us/annual/populace_us_{year}.h5@annual-release", + path=str(path), + ) + manifest = annual_manifest() + original_hashes = { + year: hashlib.sha256(Path(source.path).read_bytes()).hexdigest() + for year, source in sources.items() + } + for year, digest in original_hashes.items(): + manifest.datasets[f"populace_us_{year}"].sha256 = digest + monkeypatch.setattr(us_model, "get_release_manifest", lambda country: manifest) + monkeypatch.setattr( + us_model, "materialize_dataset", lambda *args, **kwargs: sources[kwargs["year"]] + ) + extension = Mock( + side_effect=AssertionError("Annual inputs must not be uprated again") + ) + monkeypatch.setattr(economic_assumptions, "extend_single_year_dataset", extension) + sim = us_model.managed_microsimulation( + dataset=dataset, default_calculation_period=2025 + ) + # Current income uses 2025 inputs and weights: 5*20,000 + 1*10,000. + assert sim.calc("employment_income_before_lsr", map_to="person").sum() == 110_000 + # Ohio uses prior-year income but current weights: 5*10,000 + 1*20,000. + # Explicit artifact selection must retain the same history as its family. + assert ( + sim.calc("oh_homestead_exemption_total_income", map_to="tax_unit").sum() + == 70_000 + ) + assert sim.calc("age", 2025, map_to="person").count() == 6 + extension.assert_not_called() + with pytest.raises(ValueError, match="coverage"): + sim.calc("employment_income_before_lsr", 2024) + with pytest.raises(ValueError, match="coverage"): + sim.calc("employment_income_before_lsr", 2026) + + monkeypatch.setattr(us_datasets, "get_release_manifest", lambda country: manifest) + materialize = Mock(side_effect=lambda *args, **kwargs: sources[kwargs["year"]]) + monkeypatch.setattr(us_datasets, "materialize_dataset", materialize) + annual = next( + iter(us_datasets.ensure_datasets(datasets=[dataset], years=[2025]).values()) + ) + assert [call.kwargs["year"] for call in materialize.call_args_list] == [2025] + monkeypatch.setattr( + us_model.PolicyEngineUSLatest, + "resolve_entity_variables", + lambda self, simulation: { + "tax_unit": ["oh_homestead_exemption_total_income"], + "person": ["employment_income", "employment_income_before_lsr"], + }, + ) + country_instances = [] + + class CaptureMicrosimulation(policyengine_us.Microsimulation): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + country_instances.append(self) + + monkeypatch.setattr(policyengine_us, "Microsimulation", CaptureMicrosimulation) + ordinary = Simulation( + dataset=annual, + tax_benefit_model_version=us_model.us_latest, + policy={"gov.irs.credits.ctc.amount.base[0].amount": 3000}, + scoping_strategy=RowFilterStrategy( + variable_name="household_weight", variable_value=5 + ) + if regional + else None, + ) + ordinary.run() + prior_income = sim.calc( + "oh_homestead_exemption_total_income", 2025, map_to="tax_unit" + ) + if regional: + prior_income = prior_income[ + sim.calc("household_id", 2025, map_to="tax_unit") == 1 + ] + assert ( + ordinary.output_dataset.data.tax_unit[ + "oh_homestead_exemption_total_income" + ].sum() + == prior_income.sum() + ) + assert ordinary.output_dataset.data.household["household_id"].tolist() == ( + [1] if regional else [1, 2] + ) + assert ordinary.output_dataset.data.person[ + "employment_income_before_lsr" + ].sum() == (100_000 if regional else 110_000) + baseline = country_instances[-1].baseline + assert baseline is not None + assert ( + baseline.calc( + "oh_homestead_exemption_total_income", 2025, map_to="tax_unit" + ).sum() + == prior_income.sum() + ) + assert original_hashes == { + year: hashlib.sha256(Path(source.path).read_bytes()).hexdigest() + for year, source in sources.items() + } + if not regional and dataset == "populace_us_2024": + ordinary.save() + restored = Simulation( + id=ordinary.id, dataset=annual, tax_benefit_model_version=us_model.us_latest + ) + restored.load() + assert ( + restored.output_dataset.metadata["annual_input_sources"] + == annual.metadata["annual_input_sources"] + ) + import json + + import h5py + + with h5py.File(ordinary.output_dataset.filepath, "a") as stream: + receipt = json.loads(stream["policyengine_spm"].asstr()[()]) + receipt["annual_input_sources"]["2024"]["sha256"] = "f" * 64 + del stream["policyengine_spm"] + stream.create_dataset( + "policyengine_spm", + data=json.dumps(receipt), + dtype=h5py.string_dtype("utf-8"), + ) + with pytest.raises(ValueError, match="different annual input pins"): + restored.load() + annual.metadata["annual_input_sources"]["2024"]["sha256"] = "b" * 64 + with pytest.raises(ValueError, match="pins changed"): + ordinary.save() + + +def test_annual_region_filter_keeps_weights_and_rejects_positional_replacement( + monkeypatch, tmp_path +): + import importlib + + from policyengine.core import Simulation + from policyengine.core.scoping_strategy import ( + RowFilterStrategy, + WeightReplacementStrategy, + ) + from policyengine.tax_benefit_models.us.datasets import _annual_dataset + + us_model = importlib.import_module("policyengine.tax_benefit_models.us.model") + _, source = _annual_source(tmp_path, 2025) + dataset = _annual_dataset(source, "populace_us_2024", 2025) + scoped = RowFilterStrategy(variable_name="state_fips", variable_value=6).apply( + entity_data=dataset.data.entity_data, + group_entities=us_model.US_GROUP_ENTITIES, + year=2025, + ) + assert scoped["household"]["household_id"].tolist() == [1, 2] + assert scoped["household"]["household_weight"].tolist() == [1000.0, 1000.0] + simulation = Simulation( + dataset=dataset, + tax_benefit_model_version=us_model.us_latest, + scoping_strategy=WeightReplacementStrategy( + weight_matrix_bucket="test", + weight_matrix_key="weights.h5", + lookup_csv_bucket="test", + lookup_csv_key="lookup.csv", + region_code="CA", + ), + ) + with pytest.raises(ValueError, match="positional weight replacement"): + us_model.us_latest.run(simulation) + + +def test_annual_schema_and_ids_must_match_before_engine_construction( + monkeypatch, tmp_path +): + import importlib + import sys + from types import SimpleNamespace + + import pandas as pd + + us_model = importlib.import_module("policyengine.tax_benefit_models.us.model") + _, first = _annual_source(tmp_path, 2024) + _, second = _annual_source(tmp_path, 2025) + with pd.HDFStore(second.path, "a") as store: + person = store["person"] + person.loc[0, "person_id"] = 999 + store["person"] = person + monkeypatch.setattr( + us_model, "get_release_manifest", lambda country: annual_manifest() + ) + monkeypatch.setattr( + us_model, + "materialize_dataset", + lambda *args, **kwargs: {2024: first, 2025: second}[kwargs["year"]], + ) + constructor = Mock() + monkeypatch.setitem( + sys.modules, "policyengine_us", SimpleNamespace(Microsimulation=constructor) + ) + with pytest.raises(ValueError, match="schema, rows or IDs"): + us_model.managed_microsimulation(years=[2025]) + constructor.assert_not_called() + + +def test_ensure_ignores_basename_only_legacy_cache(monkeypatch, tmp_path): + import policyengine.tax_benefit_models.us.datasets as us_datasets + + (tmp_path / "populace_us_2024_year_2025.h5").write_bytes(b"stale unversioned cache") + create = Mock(return_value={"fresh": object()}) + monkeypatch.setattr(us_datasets, "create_datasets", create) + result = us_datasets.ensure_datasets(years=[2025], data_folder=str(tmp_path)) + assert list(result) == ["fresh"] + create.assert_called_once() + + +def test_annual_family_default_uri_keeps_logical_return_key(monkeypatch, tmp_path): + import policyengine.tax_benefit_models.us.datasets as us_datasets + + manifest = annual_manifest() + _, source = _annual_source(tmp_path, 2025) + monkeypatch.setattr(us_datasets, "get_release_manifest", lambda country: manifest) + monkeypatch.setattr(us_datasets, "materialize_dataset", Mock(return_value=source)) + result = us_datasets.ensure_datasets( + datasets=[ + "hf://policyengine/populace-us/annual/populace_us_2024.h5@annual-release" + ], + years=[2025], + data_folder=str(tmp_path), + ) + assert list(result) == ["populace_us_2024_2025"] + + +def test_annual_metadata_does_not_certify_unsupported_country_loader(): + payload = annual_manifest().model_dump(mode="json") + payload["country_id"] = "uk" + with pytest.raises(ValidationError, match="only the US"): + CountryReleaseManifest.model_validate(payload) + + +@pytest.mark.parametrize( + "kwargs", + [ + {"years": [2025]}, + {"years": [2025], "dataset": "populace_us_2025"}, + { + "years": [2025], + "dataset": "hf://policyengine/populace-us/annual/populace_us_2025.h5@annual-release", + }, + {"default_calculation_period": "2025-06"}, + {}, + ], +) +def test_selected_years_load_history_but_never_future(monkeypatch, tmp_path, kwargs): + import importlib + import sys + from types import SimpleNamespace + + us_model = importlib.import_module("policyengine.tax_benefit_models.us.model") + + class FixedDate(us_model.datetime.date): + @classmethod + def today(cls): + return cls(2025, 9, 19) + + monkeypatch.setattr(us_model.datetime, "date", FixedDate) + manifest = annual_manifest() + for year in (2026, 2027): + manifest.datasets[f"populace_us_{year}"] = manifest.datasets[ + "populace_us_2025" + ].model_copy(update={"path": f"annual/populace_us_{year}.h5"}) + manifest.dataset_years[manifest.default_dataset][year] = f"populace_us_{year}" + sources = {year: _annual_source(tmp_path, year)[1] for year in (2024, 2025)} + materialize = Mock(side_effect=lambda *args, **kwargs: sources[kwargs["year"]]) + monkeypatch.setattr(us_model, "get_release_manifest", lambda country: manifest) + monkeypatch.setattr(us_model, "materialize_dataset", materialize) + + class FakeMicrosimulation: + def __init__(self, dataset, spm, **kwargs): + self.dataset = dataset + self.default_calculation_period = dataset.time_period + self.spm_config = dict(spm) + self.calc = self.calculate + + def calculate(self, variable, period=None, **kwargs): + return period + + monkeypatch.setitem( + sys.modules, + "policyengine_us", + SimpleNamespace(Microsimulation=FakeMicrosimulation), + ) + sim = us_model.managed_microsimulation(**kwargs) + assert [call.kwargs["year"] for call in materialize.call_args_list] == [2024, 2025] + assert sim.default_calculation_period == kwargs.get( + "default_calculation_period", 2025 + ) + assert sim.policyengine_bundle["annual_selected_years"] == [2025] + assert sim.policyengine_bundle["annual_loaded_years"] == [2024, 2025] + assert set(sim.policyengine_bundle["annual_input_bytes_by_year"]) == { + "2024", + "2025", + } + with pytest.raises(ValueError, match="coverage"): + sim.calc("age", 2024) + with pytest.raises(ValueError, match="coverage"): + sim.calc("age", 2026) + + +@pytest.mark.parametrize( + "kwargs", + [ + {"years": []}, + {"years": [True]}, + {"years": [2025.0]}, + {"years": [2026]}, + {"years": [2025], "default_calculation_period": 2024}, + {"dataset": "populace_us_2025", "years": [2024]}, + {"default_calculation_period": "eternity"}, + ], +) +def test_invalid_selected_years_fail_before_download(monkeypatch, kwargs): + import importlib + + us_model = importlib.import_module("policyengine.tax_benefit_models.us.model") + monkeypatch.setattr( + us_model, "get_release_manifest", lambda country: annual_manifest() + ) + materialize = Mock() + monkeypatch.setattr(us_model, "materialize_dataset", materialize) + with pytest.raises(ValueError): + us_model.managed_microsimulation(**kwargs) + materialize.assert_not_called() + + +def test_annual_history_pins_and_runtime_invalidate_reused_simulation_id( + monkeypatch, tmp_path +): + import policyengine.tax_benefit_models.us.datasets as us_datasets + from policyengine.core import Simulation + + _, source = _annual_source(tmp_path, 2025) + manifest = annual_manifest() + monkeypatch.setattr(us_datasets, "get_release_manifest", lambda country: manifest) + monkeypatch.setattr(us_datasets, "materialize_dataset", Mock(return_value=source)) + first = next(iter(us_datasets.ensure_datasets(years=[2025]).values())) + first_key = Simulation(id="same-run", dataset=first).storage_id + manifest.datasets["populace_us_2024"].sha256 = "a" * 64 + second = next(iter(us_datasets.ensure_datasets(years=[2025]).values())) + second_key = Simulation(id="same-run", dataset=second).storage_id + assert first.id != second.id + assert first_key != second_key + monkeypatch.setattr( + us_datasets, "_derived_cache_identity", lambda: {"model": "another-source"} + ) + assert Simulation(id="same-run", dataset=second).storage_id != second_key + + +def test_ordinary_annual_run_rejects_missing_or_changed_history(monkeypatch, tmp_path): + import importlib + + import policyengine.tax_benefit_models.us.datasets as us_datasets + + model = importlib.import_module("policyengine.tax_benefit_models.us.model") + _, source = _annual_source(tmp_path, 2025) + manifest = annual_manifest() + monkeypatch.setattr(us_datasets, "get_release_manifest", lambda country: manifest) + monkeypatch.setattr(model, "get_release_manifest", lambda country: manifest) + monkeypatch.setattr(us_datasets, "materialize_dataset", Mock(return_value=source)) + dataset = next(iter(us_datasets.ensure_datasets(years=[2025]).values())) + download = Mock() + monkeypatch.setattr(model, "materialize_dataset", download) + del dataset.metadata["annual_input_sources"]["2024"] + with pytest.raises(ValueError, match="history differs"): + model._ordinary_annual_inputs(dataset, dataset) + download.assert_not_called() diff --git a/tests/test_bundle.py b/tests/test_bundle.py index 4a3f536b..8ec64560 100644 --- a/tests/test_bundle.py +++ b/tests/test_bundle.py @@ -78,7 +78,11 @@ def test_selected_dataset_plan_uses_certified_release_metadata(tmp_path): assert plan.country_id == "uk" assert plan.data_package_name == "policyengine-uk-data" assert plan.repo_type == "model" - assert plan.destination == tmp_path / "enhanced_frs_2024_25.h5" + assert plan.destination.name == "enhanced_frs_2024_25.h5" + assert plan.destination.parent.parent == tmp_path / ".policyengine" / "sources" + cache_identity = plan.destination.parent.name + assert len(cache_identity) == 64 + assert set(cache_identity) <= set("0123456789abcdef") assert plan.sha256 == release["datasets"][plan.dataset]["sha256"] diff --git a/tests/test_dataset_materialization.py b/tests/test_dataset_materialization.py index 256e004d..df91ab71 100644 --- a/tests/test_dataset_materialization.py +++ b/tests/test_dataset_materialization.py @@ -60,7 +60,8 @@ def test_resolve_bundle_dataset_inherits_primary_package(tmp_path): assert dataset.repo_id == "policyengine/policyengine-uk-data-private" assert dataset.repo_type == "model" assert dataset.revision == "uk-release" - assert dataset.destination == tmp_path / "enhanced_frs_2024_25.h5" + assert dataset.destination.name == "enhanced_frs_2024_25.h5" + assert dataset.destination.parent.parent == tmp_path / ".policyengine" / "sources" def test_resolve_bundle_dataset_uses_cross_package_overlay(tmp_path): @@ -190,7 +191,7 @@ def test_downloads_and_verifies_metadata_sidecar(tmp_path): result = _download(manifest, tmp_path, session) - assert result.metadata_path == (tmp_path / "enhanced_frs_2024_25.h5.metadata.json") + assert result.metadata_path == result.path.with_suffix(".h5.metadata.json") assert result.metadata_path.read_bytes() == metadata_payload assert session.calls[1][0].endswith("/enhanced_frs_2024_25.h5.metadata.json") @@ -198,7 +199,10 @@ def test_downloads_and_verifies_metadata_sidecar(tmp_path): def test_reuses_destination_when_hash_matches(tmp_path): payload = b"certified" manifest = _manifest_with_hash(_sha256(payload)) - destination = tmp_path / "enhanced_frs_2024_25.h5" + destination = _resolve_bundle_dataset( + "uk", data_dir=tmp_path, manifest=manifest + ).destination + destination.parent.mkdir(parents=True, exist_ok=True) destination.write_bytes(payload) session = _Session() @@ -211,7 +215,10 @@ def test_reuses_destination_when_hash_matches(tmp_path): def test_replaces_destination_when_hash_does_not_match(tmp_path): payload = b"certified" manifest = _manifest_with_hash(_sha256(payload)) - destination = tmp_path / "enhanced_frs_2024_25.h5" + destination = _resolve_bundle_dataset( + "uk", data_dir=tmp_path, manifest=manifest + ).destination + destination.parent.mkdir(parents=True, exist_ok=True) destination.write_bytes(b"old") _download(manifest, tmp_path, _Session(_Response(payload))) @@ -221,7 +228,10 @@ def test_replaces_destination_when_hash_does_not_match(tmp_path): def test_hash_failure_preserves_existing_destination(tmp_path): manifest = _manifest_with_hash(_sha256(b"expected")) - destination = tmp_path / "enhanced_frs_2024_25.h5" + destination = _resolve_bundle_dataset( + "uk", data_dir=tmp_path, manifest=manifest + ).destination + destination.parent.mkdir(parents=True, exist_ok=True) destination.write_bytes(b"old") with pytest.raises(DatasetMaterializationError, match="sha256"):