diff --git a/changelog.d/507.added.md b/changelog.d/507.added.md new file mode 100644 index 00000000..b6b0b483 --- /dev/null +++ b/changelog.d/507.added.md @@ -0,0 +1 @@ +Add a source-only New Zealand Axiom pilot for the two official-budget WFF/IWTC entitlement comparisons. The adapter reads Microcosm entity-table artifacts, inherits family weights from households through Frame, and records explicit policy-period and runtime provenance. It does not provide a certified NZ population or a Treasury fiscal-cost replication. diff --git a/docs/engineering/nz-axiom-pilot.md b/docs/engineering/nz-axiom-pilot.md new file mode 100644 index 00000000..f7a28346 --- /dev/null +++ b/docs/engineering/nz-axiom-pilot.md @@ -0,0 +1,87 @@ +# New Zealand Axiom pilot + +The source-only NZ adapter executes the two WFF/IWTC entitlement comparisons in +`rulespec-nz/nz/policies/budget/official_budget_reform_replication.yaml` through +the real Axiom dense runtime. It does not install a `policyengine-nz` model or +expose a certified `pe.nz` bundle. + +## Prerequisites + +Use Python 3.14 and compatible source installations of `microcosm-frame`, +`axiom-rules-engine`, and its compiled dense extension. Microcosm must provide +`NZ_SCHEMA`, `AxiomPeriod`, explicit `rulespec_roots`, and the shared Axiom HDF5 +reader. These dependencies remain source-only; this change does not add an +unavailable package extra or change US/UK pins. + +The rules checkout must have committed NZ source, the official-budget transport +contract, and its toolchain configuration. The adapter records the RuleSpec +commit, source and contract hashes, Microcosm/Python source hashes, and native +binary hash. A source-built extension can lack package version metadata; its +binary hash remains mandatory. Wrapper package versions do not establish the +native engine's release version. + +Restoring a serialized model configuration verifies the saved runtime identity. +It fails if source or native code has changed; create a new model explicitly to +run a different version. + +## Input and calculation contract + +Use the example in `examples/nz_axiom_pilot.py` with a supplied HDF5 artifact and +its verified SHA-256. The reader uses Microcosm's actual entity-table codec, +including Decimal and nullable-boolean preservation. + +- The artifact contains person, household, and family tables and build label + `2026`. Axiom receives the explicit tax year `2026-04-01` to `2027-03-31`. +- Only the household table stores weights. Frame resolves person and family + weights through membership; the adapter exposes those effective weights in + memory through MicroDataFrames. Family members must share one household. +- All 11 substantive family inputs must be present and satisfy the transport + contract. The 10 unrelated eager-graph padding inputs may default to their + declared zero values. Stored formula outputs and nonzero padding fail. +- Decimal inputs retain their values. The pilot rejects values outside the + declared `decimal128(18,2)` contract; the runtime also checks whether its native + numeric boundary can represent them. +- Outputs remain in memory. Running or saving a simulation never writes over + its input artifact. Receipts identify both the original artifact and the + actual in-memory inputs, including edits made after loading. +- This pilot supports neither arbitrary reforms nor dynamic or geographic + scoping controls. Use `Simulation.run()` for an explicit fresh execution. + +The two outputs are `budget_2025_wff_abatement_entitlement_change` and +`budget_2026_iwtc_entitlement_change`. Aggregate them with the returned +MicroSeries `.sum()`, not manual multiplication by person or family weights. + +NZ model configurations and direct `PopulaceNewZealandDataset` outputs support +Pydantic JSON round-trips. The NZ-only table codec retains exact Decimals, +nullable booleans, categories, indices, dtypes, and effective weight columns. +Serialize the output dataset itself with `output_dataset.model_dump_json()` and +restore it with `PopulaceNewZealandDataset.model_validate_json(...)`. + +The shared `Simulation` model still serializes through base-typed dataset/model +fields, which omit subclass details. Its generic JSON form is not a complete NZ +run archive; this pilot does not change that pre-existing cross-country behavior. + +## What the result does not establish + +The adapter does not certify the input population, repair missing family +inputs, fit weights, or calibrate to official cost estimates. A positive test +on a synthetic population is an integration check, not an NZ national result. + +These annual family-entitlement changes are not yet comparable to Treasury's +forecast operating costs. Fiscal-year payment timing, WFF debt impairment, +and the temporary IWTC petrol-trigger/payment-tail treatment require explicit +bridges. The output receipt retains the upstream `bridge_required` status. +Do not attach these totals as completed Scorecard budget-score replications. + +## Verification + +Ordinary boundary tests use non-statutory runtime doubles. The opt-in test uses +Microcosm's real writer and reader plus Axiom's compiled RuleSpec module, with +nonuniform household weights and two upstream companion cases: + +```sh +POLICYENGINE_SKIP_COUNTRY_IMPORTS=1 \ +RUN_NZ_AXIOM_INTEGRATION=1 \ +RULESPEC_NZ_ROOT=/path/to/rulespec-nz \ +uv run --no-sync pytest --noconftest -q tests/test_nz_axiom_pilot.py +``` diff --git a/examples/nz_axiom_pilot.py b/examples/nz_axiom_pilot.py new file mode 100644 index 00000000..35cd75ec --- /dev/null +++ b/examples/nz_axiom_pilot.py @@ -0,0 +1,73 @@ +"""Run the source-only NZ entitlement pilot over a supplied Microcosm artifact. + +Requires compatible Microcosm/Axiom source installs and a committed rulespec-nz +checkout. This example does not download or certify a population, convert +entitlements to Treasury fiscal costs, or publish a Scorecard result. + + POLICYENGINE_SKIP_COUNTRY_IMPORTS=1 uv run --no-sync python \ + examples/nz_axiom_pilot.py --dataset /path/to/populace_nz_2026.h5 \ + --rulespec-root /path/to/rulespec-nz --sha256 VERIFIED_ARTIFACT_SHA256 + +Use --weight-kind calibrated only for an artifact with calibrated weights. +""" + +import argparse +import json + +from policyengine.core import Simulation +from policyengine.tax_benefit_models.nz import ( + IWTC_CHANGE, + WFF_ABATEMENT_CHANGE, + AxiomNewZealandPilot, + PopulaceNewZealandDataset, +) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--dataset", required=True) + parser.add_argument("--rulespec-root", required=True) + parser.add_argument("--sha256", required=True) + parser.add_argument( + "--weight-kind", choices=("design", "calibrated"), default="design" + ) + args = parser.parse_args() + dataset = PopulaceNewZealandDataset( + name="NZ source-only transport pilot", + description="User-supplied NZ family inputs; this adapter does not certify the population", + filepath=args.dataset, + source_sha256=args.sha256, + weight_kind=args.weight_kind, + year=2026, + ) + simulation = Simulation( + dataset=dataset, + tax_benefit_model_version=AxiomNewZealandPilot( + rulespec_root=args.rulespec_root + ), + ) + simulation.run() + output = simulation.output_dataset + # The adapter resolves effective family weights through Frame; MicroSeries + # applies them here. Do not multiply by person/family weight columns. + changes = { + name: float(output.data.family[name].sum()) + for name in (WFF_ABATEMENT_CHANGE, IWTC_CHANGE) + } + print( + json.dumps( + { + "country": "nz", + "status": "source_only_entitlement_pilot", + "official_budget_score_comparable": False, + "weighted_family_entitlement_changes_nzd": changes, + "receipt": output.metadata["policyengine_axiom_runs"][-1], + }, + indent=2, + allow_nan=False, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/src/policyengine/tax_benefit_models/nz/__init__.py b/src/policyengine/tax_benefit_models/nz/__init__.py new file mode 100644 index 00000000..a566968d --- /dev/null +++ b/src/policyengine/tax_benefit_models/nz/__init__.py @@ -0,0 +1,20 @@ +"""Source-only New Zealand pilot: real Axiom over Microcosm family tables.""" + +from .datasets import NZYearData, PopulaceNewZealandDataset +from .model import ( + IWTC_CHANGE, + WFF_ABATEMENT_CHANGE, + AxiomNewZealand, + AxiomNewZealandPilot, + nz_model, +) + +__all__ = [ + "AxiomNewZealand", + "AxiomNewZealandPilot", + "IWTC_CHANGE", + "NZYearData", + "PopulaceNewZealandDataset", + "WFF_ABATEMENT_CHANGE", + "nz_model", +] diff --git a/src/policyengine/tax_benefit_models/nz/datasets.py b/src/policyengine/tax_benefit_models/nz/datasets.py new file mode 100644 index 00000000..dff9ba56 --- /dev/null +++ b/src/policyengine/tax_benefit_models/nz/datasets.py @@ -0,0 +1,241 @@ +"""Read-only NZ pilot datasets with Frame-owned household weight inheritance. + +The public transport artifact contains person, household, and family tables. +Only household weights are stored. Effective person/family weights are resolved +by Microcosm and exposed as in-memory MicroDataFrame weights for PolicyEngine. +""" + +import json +from hashlib import sha256 +from pathlib import Path +from typing import Any, Literal, Optional + +import numpy as np +import pandas as pd +from microdf import MicroDataFrame +from pydantic import Field, ValidationInfo, field_serializer, field_validator + +from policyengine.core import Dataset, YearData + +from .serialization import decode_table, encode_table + +ENTITIES = ("person", "household", "family") + + +class NZYearData(YearData): + """Entity-level NZ pilot data; effective weights exist only in memory.""" + + person: MicroDataFrame + household: MicroDataFrame + family: MicroDataFrame + + @field_serializer("person", "household", "family", when_used="json") + def serialize_entity_table(self, value: MicroDataFrame) -> dict[str, Any]: + return encode_table(value) + + @field_validator("person", "household", "family", mode="before") + @classmethod + def restore_entity_table(cls, value: Any, info: ValidationInfo) -> Any: + if isinstance(value, dict): + return decode_table(value, info.field_name) + return value + + @property + def entity_data(self) -> dict[str, MicroDataFrame]: + return {entity: getattr(self, entity) for entity in ENTITIES} + + +def _load_frame_runtime() -> tuple[Any, Any, Any, Any]: + try: + from microcosm.frame import Frame, WeightKind, Weights + from microcosm.frame.adapters.axiom import NZ_SCHEMA + except ImportError as error: + raise ImportError( + "The NZ pilot requires a compatible source checkout of " + "microcosm-frame with NZ_SCHEMA and AxiomPeriod support." + ) from error + return Frame, WeightKind, Weights, NZ_SCHEMA + + +def _load_dataset_reader() -> Any: + try: + from microcosm.frame.adapters.axiom import AxiomEntityTableDataset + except ImportError as error: + raise ImportError( + "The NZ pilot requires Microcosm's Axiom entity-table HDF5 reader." + ) from error + return AxiomEntityTableDataset + + +def _validate_family_nesting(tables: dict[str, pd.DataFrame]) -> None: + required = { + "person": ("person_id", "person_household_id", "person_family_id"), + "household": ("household_id", "household_weight"), + "family": ("family_id", "family_household_id"), + } + for entity, columns in required.items(): + for column in columns: + if column not in tables[entity]: + raise ValueError(f"NZ {entity} table is missing {column!r}.") + if tables[entity][column].isna().any(): + raise ValueError(f"NZ {entity}.{column} must not contain nulls.") + family = tables["family"] + if family["family_id"].duplicated().any(): + raise ValueError("NZ family_id values must be unique.") + membership = tables["person"][ + ["person_family_id", "person_household_id"] + ].drop_duplicates() + if membership["person_family_id"].duplicated().any(): + raise ValueError("Each NZ family must belong to exactly one household.") + expected = family["family_id"].map( + membership.set_index("person_family_id")["person_household_id"] + ) + if expected.isna().any() or not np.array_equal( + expected.to_numpy(), family["family_household_id"].to_numpy() + ): + raise ValueError( + "NZ family_household_id must match the household of every family member." + ) + + +def frame_from_tables( + tables: dict[str, pd.DataFrame], + *, + allow_effective_weights: bool = False, + weight_kind: Literal["design", "calibrated"] = "design", + metadata: Optional[dict[str, Any]] = None, +) -> Any: + """Validate NZ structure and construct a household-weighted Frame.""" + if set(tables) != set(ENTITIES): + raise ValueError(f"NZ requires exactly the entity tables {ENTITIES}.") + copies = {entity: pd.DataFrame(table).copy() for entity, table in tables.items()} + effective: dict[str, np.ndarray] = {} + for entity, table in copies.items(): + for column in list(table.columns): + if not column.endswith("_weight"): + continue + owner = column.removesuffix("_weight") + if owner != entity or ( + owner != "household" and not allow_effective_weights + ): + raise ValueError( + "NZ household_weight must be the sole stored weight vector." + ) + if owner != "household": + effective[owner] = table.pop(column).to_numpy() + _validate_family_nesting(copies) + Frame, WeightKind, Weights, schema = _load_frame_runtime() + values = copies["household"].pop("household_weight").to_numpy() + person = copies["person"] + strata = person["support_stratum"] if "support_stratum" in person else None + frame = Frame( + copies, + schema, + { + "household": Weights( + values=values, + kind=WeightKind.CALIBRATED + if weight_kind == "calibrated" + else WeightKind.DESIGN, + ) + }, + strata, + metadata=metadata, + ) + for entity, prior in effective.items(): + resolved = frame.resolve_weights(entity).values + if not np.array_equal(prior, resolved): + raise ValueError( + f"NZ effective {entity}_weight differs from Frame-resolved household weights." + ) + return frame + + +def year_data_from_frame(frame: Any) -> NZYearData: + """Expose Frame-resolved weights through PolicyEngine MicroDataFrames.""" + tables = {} + for entity in ENTITIES: + table = frame.table(entity).copy() + column = f"{entity}_weight" + table[column] = frame.resolve_weights(entity).values + tables[entity] = MicroDataFrame(table, weights=column) + return NZYearData(**tables) + + +class PopulaceNewZealandDataset(Dataset): + """A source-only NZ transport input, never a certified country bundle. + + ``year`` is the Microcosm build-period label (2026 for tax year 2026–27). + ``policy_period`` on outputs carries the explicit Axiom start/end dates. + This pilot deliberately has no file-writing method; simulation outputs are + in memory and cannot overwrite the input artifact. + """ + + data: Optional[NZYearData] = None + metadata: dict[str, Any] = Field(default_factory=dict) + policy_period: Optional[dict[str, str]] = None + source_sha256: Optional[str] = None + weight_kind: Literal["design", "calibrated"] = "design" + + def load(self) -> None: + if self.filepath is None: + raise ValueError("Cannot load an NZ pilot dataset without a filepath.") + path = Path(self.filepath) + before = path.stat() + digest = sha256() + with path.open("rb") as source: + for block in iter(lambda: source.read(1024 * 1024), b""): + digest.update(block) + source_sha256 = digest.hexdigest() + if self.source_sha256 is not None and self.source_sha256 != source_sha256: + raise ValueError("NZ input artifact does not match its expected SHA-256.") + with pd.HDFStore(path, mode="r") as store: + keys = {key.lstrip("/") for key in store.keys()} + if keys != {*ENTITIES, "_time_period"}: + raise ValueError( + "NZ HDF5 must contain person, household, family, and _time_period only." + ) + period = store["_time_period"] + if ( + len(period) != 1 + or not pd.api.types.is_integer_dtype(period.dtype) + or period.iloc[0] != self.year + ): + raise ValueError( + f"NZ dataset period mismatch: expected build label {self.year}." + ) + attrs = store.get_storer("_time_period").attrs + encoded = getattr(attrs, "policyengine_metadata_json", None) + metadata = ( + json.loads(str(encoded)) if encoded is not None else dict(self.metadata) + ) + if not isinstance(metadata, dict): + raise ValueError("NZ dataset metadata must be a JSON object.") + if self.metadata and self.metadata != metadata: + raise ValueError("NZ dataset metadata differs from its HDF5 artifact.") + reader = _load_dataset_reader()(file_path=path) + if reader.time_period != self.year: + raise ValueError( + f"NZ dataset period mismatch: expected build label {self.year}." + ) + after = path.stat() + if (before.st_size, before.st_mtime_ns) != (after.st_size, after.st_mtime_ns): + raise RuntimeError("NZ input artifact changed while it was being loaded.") + frame = frame_from_tables( + reader.tables, weight_kind=self.weight_kind, metadata=metadata + ) + self.data = year_data_from_frame(frame) + self.metadata = metadata + self.source_sha256 = source_sha256 + + def to_frame(self) -> Any: + """Return a fresh validated Frame without redundant effective weights.""" + if self.data is None: + self.load() + assert self.data is not None + return frame_from_tables( + self.data.entity_data, + allow_effective_weights=True, + weight_kind=self.weight_kind, + metadata=self.metadata, + ) diff --git a/src/policyengine/tax_benefit_models/nz/model.py b/src/policyengine/tax_benefit_models/nz/model.py new file mode 100644 index 00000000..76eed806 --- /dev/null +++ b/src/policyengine/tax_benefit_models/nz/model.py @@ -0,0 +1,442 @@ +"""Execute the NZ official-reform entitlement module through real Axiom. + +This is an opt-in, source-only PolicyEngine integration, not policyengine-nz +or a certified NZ bundle. The module computes legal family entitlement deltas; +Treasury cash/accrual timing, debt impairment, and petrol-trigger scenarios are +separate unimplemented bridges. Ordinary Simulation reform/dynamic/scoping +controls are rejected because this pilot does not implement them. +""" + +import json +import subprocess +from copy import deepcopy +from decimal import Decimal, InvalidOperation +from hashlib import sha256 +from importlib import metadata as importlib_metadata +from importlib.util import find_spec +from numbers import Real +from pathlib import Path +from typing import TYPE_CHECKING, Any, ClassVar + +import numpy as np +import pandas as pd +from pandas.api.types import is_bool_dtype, is_integer_dtype, is_numeric_dtype +from pydantic import Field + +from policyengine.core import TaxBenefitModel, TaxBenefitModelVersion + +from . import datasets as nz_datasets +from .datasets import PopulaceNewZealandDataset, year_data_from_frame + +if TYPE_CHECKING: + from policyengine.core import Simulation + +PILOT_MODULE = "nz/policies/budget/official_budget_reform_replication.yaml" +TRANSPORT_CONTRACT = "data/microsimulation/official-budget-reform-transport.json" +WFF_ABATEMENT_CHANGE = "budget_2025_wff_abatement_entitlement_change" +IWTC_CHANGE = "budget_2026_iwtc_entitlement_change" +OUTPUTS = [WFF_ABATEMENT_CHANGE, IWTC_CHANGE] +POLICY_PERIOD = {"start": "2026-04-01", "end": "2027-03-31", "kind": "tax_year"} + + +class AxiomNewZealand(TaxBenefitModel): + id: str = "axiom-rulespec-nz" + description: str = ( + "New Zealand RuleSpec executed by Axiom over Microcosm family tables." + ) + + +nz_model = AxiomNewZealand() + + +class AxiomNewZealandPilot(TaxBenefitModelVersion): + """Source-only 2026–27 WFF/IWTC entitlement comparison pilot.""" + + country_code: ClassVar[str] = "nz" + rulespec_root: str + runtime_provenance: dict[str, Any] + transport_contract: dict[str, Any] = Field(exclude=True) + + def __init__(self, **kwargs: Any) -> None: + root = Path(kwargs["rulespec_root"]).expanduser().resolve() + contract = _load_transport_contract(root) + provenance = _build_runtime_provenance(root) + version = sha256( + json.dumps(provenance, sort_keys=True, separators=(",", ":")).encode() + ).hexdigest() + identity = { + "runtime_provenance": provenance, + "version": version, + "id": f"{nz_model.id}@{version}", + } + for field, current in identity.items(): + if field in kwargs and kwargs[field] != current: + raise ValueError( + f"The saved NZ model {field} has changed relative to this runtime; " + "construct a new model explicitly instead of rebinding saved configuration." + ) + kwargs.update( + rulespec_root=str(root), + transport_contract=contract, + runtime_provenance=provenance, + model=nz_model, + version=version, + id=f"{nz_model.id}@{version}", + ) + super().__init__(**kwargs) + + def run(self, simulation: "Simulation") -> "Simulation": + for name in ("policy", "dynamic", "scoping_strategy", "extra_variables"): + if getattr(simulation, name): + raise ValueError( + f"The NZ pilot does not support Simulation.{name}; it executes the pinned official-reform module only." + ) + root = Path(self.rulespec_root) + if _build_runtime_provenance(root) != self.runtime_provenance: + raise RuntimeError( + "The NZ RuleSpec/Axiom runtime changed; construct a new model version before running." + ) + dataset = simulation.dataset + if not isinstance(dataset, PopulaceNewZealandDataset): + raise TypeError("AxiomNewZealandPilot requires PopulaceNewZealandDataset.") + if dataset.is_output_dataset: + raise ValueError("An NZ output dataset cannot be reused as a rules input.") + if dataset.year != 2026: + raise ValueError( + "The NZ pilot only supports build label 2026 and tax year 2026–27." + ) + frame = dataset.to_frame() + input_frame_sha256 = _frame_sha256(frame) + contract = self.transport_contract + forbidden = set( + contract["output_contract"]["formula_owned_excluded_from_dataset"] + ) + stored = { + name for entity in frame.entities for name in frame.table(entity).columns + } + if forbidden.intersection(stored): + raise ValueError( + f"NZ input contains formula-owned outputs: {sorted(forbidden.intersection(stored))}." + ) + family = frame.table("family").copy() + for item in contract["input_contract"]["required_target_inputs"]: + _validate_input(family, item) + padding_applied = [] + for item in contract["input_contract"]["adapter_padding_defaults"]: + if item["name"] not in family: + family.loc[:, item["name"]] = item["value"] + padding_applied.append(item["name"]) + values = family[item["name"]] + if values.isna().any() or not all( + isinstance(value, (Real, Decimal)) + and not isinstance(value, (bool, np.bool_)) + and value == 0 + for value in values + ): + raise ValueError( + f"NZ adapter padding {item['name']!r} must contain numeric zeros only." + ) + + Frame, _, _, schema = nz_datasets._load_frame_runtime() + tables = {entity: frame.table(entity).copy() for entity in frame.entities} + tables["family"] = family + frame = Frame( + tables, + schema, + {"household": frame.weights_for("household")}, + frame.strata, + metadata=frame.metadata, + ) + AxiomEngine, AxiomPeriod = _load_axiom_runtime() + source_period = contract["period"] + period = { + "start": source_period["start"], + "end": source_period["end"], + "kind": source_period["period_kind"], + } + engine = AxiomEngine(root / PILOT_MODULE, schema=schema, rulespec_roots=(root,)) + outputs = engine.materialize(frame, OUTPUTS, AxiomPeriod(**period)) + for name in OUTPUTS: + values = np.asarray(outputs[name]) + if ( + values.shape != (frame.n("family"),) + or not np.isfinite(values.astype(float)).all() + ): + raise ValueError(f"Axiom returned invalid NZ output {name!r}.") + tables["family"][name] = values + output_frame = Frame( + tables, + schema, + {"household": frame.weights_for("household")}, + frame.strata, + metadata=frame.metadata, + ) + metadata = deepcopy(dataset.metadata) + previous = metadata.get("policyengine_axiom_runs", []) + if not isinstance(previous, list): + raise ValueError("NZ policyengine_axiom_runs metadata must be a list.") + metadata["policyengine_axiom_runs"] = [ + *previous, + { + "dataset_year": dataset.year, + "input_artifact_sha256": dataset.source_sha256, + "input_frame_sha256": input_frame_sha256, + "weight_kind": dataset.weight_kind, + "certified_population": False, + "policy_period": period, + "model_version": self.version, + "output_variables": list(OUTPUTS), + "padding_applied": padding_applied, + "provenance": deepcopy(self.runtime_provenance), + "official_score_bridge": deepcopy(contract["official_score_bridge"]), + }, + ] + simulation.output_dataset = PopulaceNewZealandDataset( + id=simulation.id, + name=dataset.name, + description=dataset.description, + filepath=None, + year=dataset.year, + policy_period=period, + source_sha256=dataset.source_sha256, + weight_kind=dataset.weight_kind, + is_output_dataset=True, + metadata=metadata, + data=year_data_from_frame(output_frame), + ) + return simulation + + def save(self, simulation: "Simulation") -> None: + """Pilot outputs are in memory and recomputed, never saved over inputs.""" + + def load(self, simulation: "Simulation") -> None: + raise FileNotFoundError("NZ pilot simulations are recomputed, not persisted.") + + +def _load_axiom_runtime() -> tuple[Any, Any]: + try: + from microcosm.frame.adapters.axiom import AxiomEngine, AxiomPeriod + except ImportError as error: + raise ImportError( + "The NZ pilot needs compatible microcosm-frame and Axiom source installs with AxiomPeriod support." + ) from error + return AxiomEngine, AxiomPeriod + + +def _validate_input(family: Any, item: dict[str, Any]) -> None: + name, dtype = item["name"], item["dtype"] + if name not in family: + raise ValueError( + f"Missing required NZ Family input {name!r}; no default is permitted." + ) + values = family[name] + if values.isna().any(): + raise ValueError(f"NZ Family input {name!r} contains nulls.") + if dtype == "bool": + valid = is_bool_dtype(values.dtype) + elif dtype == "int16": + valid = is_integer_dtype(values.dtype) and not is_bool_dtype(values.dtype) + else: + valid = ( + is_numeric_dtype(values.dtype) and not is_bool_dtype(values.dtype) + ) or all(isinstance(value, Decimal) for value in values) + if valid: + try: + decimals = [ + value if isinstance(value, Decimal) else Decimal(str(value)) + for value in values + ] + valid = all( + value.is_finite() + and abs(value) < Decimal("1e16") + and value == value.quantize(Decimal("0.01")) + for value in decimals + ) + except (InvalidOperation, ValueError): + valid = False + if not valid: + raise ValueError(f"NZ Family input {name!r} must satisfy dtype {dtype}.") + if ( + dtype == "int16" + and ( + (values < np.iinfo(np.int16).min) | (values > np.iinfo(np.int16).max) + ).any() + ): + raise ValueError(f"NZ Family input {name!r} is outside int16 bounds.") + + +def _load_transport_contract(root: Path) -> dict[str, Any]: + if not (root / PILOT_MODULE).is_file(): + raise FileNotFoundError(f"NZ Axiom module not found below {root}.") + contract = json.loads((root / TRANSPORT_CONTRACT).read_text()) + if ( + contract.get("schema") != "axiom/nz-official-budget-reform-transport/1" + or contract.get("jurisdiction") != "nz" + ): + raise ValueError("Unsupported NZ transport contract.") + if ( + contract.get("runtime", {}).get("rulespec_module") != PILOT_MODULE + or contract.get("runtime", {}).get("root_entity") != "Family" + ): + raise ValueError( + "NZ transport contract names an unsupported module or root entity." + ) + period = contract.get("period", {}) + if { + "start": period.get("start"), + "end": period.get("end"), + "kind": period.get("period_kind"), + } != POLICY_PERIOD: + raise ValueError("The NZ pilot requires the explicit 2026–27 tax-year period.") + inputs = contract["input_contract"] + required = inputs["required_target_inputs"] + padding = inputs["adapter_padding_defaults"] + names = [item["name"] for item in required + padding] + if len(names) != len(set(names)) or len(names) != inputs["engine_root_input_count"]: + raise ValueError( + "NZ transport input names must be unique and match the declared count." + ) + if any( + item.get("entity") != "family" + or item.get("missing") != "fail_closed" + or item.get("dtype") not in {"bool", "int16", "decimal128(18,2)"} + for item in required + ): + raise ValueError( + "NZ transport required inputs must be typed, fail-closed Family fields." + ) + if any( + item.get("value") != 0 or isinstance(item.get("value"), bool) + for item in padding + ): + raise ValueError("NZ transport padding must be explicit numeric zero values.") + outputs = contract["output_contract"]["requested"] + if [item["name"] for item in outputs] != OUTPUTS or any( + item.get("entity") != "family" or item.get("unit") != "NZD" for item in outputs + ): + raise ValueError( + "NZ transport outputs must be the two Family entitlement deltas in NZD." + ) + return contract + + +def _file_sha256(path: Path) -> str: + digest = sha256() + with path.open("rb") as source: + for block in iter(lambda: source.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def _frame_sha256(frame: Any) -> str: + """Identify actual in-memory inputs, including edits since artifact load. + + This is a fingerprint, not a weight aggregation. The receipt records the + pandas version defining its row-hash algorithm alongside the runtime. + """ + digest = sha256() + for entity in nz_datasets.ENTITIES: + table = frame.table(entity) + header = [entity, list(table.columns), [str(dtype) for dtype in table.dtypes]] + digest.update(json.dumps(header, separators=(",", ":")).encode()) + digest.update( + pd.util.hash_pandas_object(table, index=True) + .to_numpy(dtype=" Path: + spec = find_spec(name) + if spec is None or spec.origin is None or not Path(spec.origin).is_file(): + raise ImportError(f"Cannot locate the source-only NZ runtime module {name!r}.") + return Path(spec.origin) + + +def _package_sha256(name: str) -> str: + root = _module_path(name).parent + digest = sha256() + for path in sorted(root.rglob("*.py")): + relative = path.relative_to(root).as_posix().encode() + payload = path.read_bytes() + digest.update(len(relative).to_bytes(8, "big")) + digest.update(relative) + digest.update(len(payload).to_bytes(8, "big")) + digest.update(payload) + return digest.hexdigest() + + +def _git(root: Path, *args: str) -> str: + result = subprocess.run( + ["git", "-C", str(root), *args], check=True, capture_output=True, text=True + ) + return result.stdout.strip() + + +def _optional_version(distribution: str) -> str | None: + try: + return importlib_metadata.version(distribution) + except importlib_metadata.PackageNotFoundError: + # Source-built native extensions can have no distribution metadata. + # Their file hash remains mandatory; do not invent an engine version. + return None + + +def _build_runtime_provenance(root: Path) -> dict[str, Any]: + """Identify all executed source trees without scanning unrelated files.""" + try: + if Path(_git(root, "rev-parse", "--show-toplevel")).resolve() != root: + raise ValueError("rulespec_root must be the rulespec-nz repository root.") + dirty = _git( + root, + "status", + "--porcelain", + "--untracked-files=all", + "--", + "nz", + TRANSPORT_CONTRACT, + ".axiom/toolchain.toml", + ) + if dirty: + raise RuntimeError( + "Commit the NZ RuleSpec source/transport changes before constructing a content-identified pilot." + ) + return { + "rulespec": { + "repository": "TheAxiomFoundation/rulespec-nz", + "commit": _git(root, "rev-parse", "HEAD"), + "nz_tree": _git(root, "rev-parse", "HEAD:nz"), + "module": PILOT_MODULE, + "module_sha256": _file_sha256(root / PILOT_MODULE), + "transport_contract_sha256": _file_sha256(root / TRANSPORT_CONTRACT), + }, + "runtime": { + "policyengine_nz_sha256": _package_sha256( + "policyengine.tax_benefit_models.nz" + ), + "pandas_version": importlib_metadata.version("pandas"), + "numpy_version": importlib_metadata.version("numpy"), + "microcosm_frame_version": importlib_metadata.version( + "microcosm-frame" + ), + "microcosm_frame_sha256": _package_sha256("microcosm.frame"), + "axiom_python_version": importlib_metadata.version( + "axiom-rules-engine" + ), + "axiom_python_sha256": _package_sha256("axiom_rules_engine"), + "axiom_dense_version": _optional_version("axiom-rules-engine-dense"), + "axiom_dense_sha256": _file_sha256( + _module_path("axiom_rules_engine_dense") + ), + }, + } + except ( + subprocess.CalledProcessError, + importlib_metadata.PackageNotFoundError, + ModuleNotFoundError, + ) as error: + raise ImportError( + "The NZ pilot requires committed rulespec-nz plus compatible microcosm-frame, Axiom Python, and dense-extension source installs." + ) from error diff --git a/src/policyengine/tax_benefit_models/nz/serialization.py b/src/policyengine/tax_benefit_models/nz/serialization.py new file mode 100644 index 00000000..a9c5737e --- /dev/null +++ b/src/policyengine/tax_benefit_models/nz/serialization.py @@ -0,0 +1,196 @@ +"""JSON transport for the NZ pilot's primitive entity-table values. + +This is a serialization codec, not a policy evaluator or dataset writer. +Decimal values use explicit tags so JSON cannot turn exact money into floats. +""" + +import math +from datetime import date, datetime +from decimal import Decimal +from typing import Any + +import numpy as np +import pandas as pd +from microdf import MicroDataFrame + +TABLE_SCHEMA = "policyengine/nz-entity-table/1" + + +def _encode_cell(value: Any) -> Any: + if value is pd.NA: + return {"_type": "na"} + if value is pd.NaT: + return {"_type": "nat"} + if isinstance(value, Decimal): + return {"_type": "decimal", "value": str(value)} + if isinstance(value, pd.Timestamp): + return {"_type": "timestamp", "value": value.isoformat()} + if isinstance(value, (datetime, date)): + return {"_type": type(value).__name__, "value": value.isoformat()} + if isinstance(value, np.generic): + return _encode_cell(value.item()) + if isinstance(value, float) and not math.isfinite(value): + return {"_type": "float", "value": str(value)} + if value is None or isinstance(value, (str, bool, int, float)): + return value + raise ValueError(f"Unsupported NZ entity-table JSON value: {type(value).__name__}.") + + +def _decode_cell(value: Any) -> Any: + if not isinstance(value, dict): + return value + kind = value.get("_type") + if kind == "na": + return pd.NA + if kind == "nat": + return pd.NaT + decoders = { + "decimal": Decimal, + "timestamp": pd.Timestamp, + "datetime": datetime.fromisoformat, + "date": date.fromisoformat, + "float": float, + } + if kind not in decoders: + raise ValueError(f"Unsupported NZ JSON scalar tag {kind!r}.") + return decoders[kind](value["value"]) + + +def _encode_dtype(dtype: Any) -> dict[str, Any]: + if isinstance(dtype, pd.CategoricalDtype): + return { + "kind": "category", + "categories": _encode_index(dtype.categories), + "ordered": dtype.ordered, + } + if isinstance(dtype, pd.StringDtype): + return { + "kind": "string", + "storage": dtype.storage, + "na_value": _encode_cell(dtype.na_value), + } + return {"kind": "pandas", "name": str(dtype)} + + +def _decode_dtype(payload: dict[str, Any]) -> Any: + if payload["kind"] == "category": + return pd.CategoricalDtype( + categories=_decode_index(payload["categories"]), ordered=payload["ordered"] + ) + if payload["kind"] == "string": + na_value = _decode_cell(payload.get("na_value", {"_type": "na"})) + if na_value is pd.NA: + # The omitted keyword keeps older pandas versions compatible. + return pd.StringDtype(storage=payload["storage"]) + if not isinstance(na_value, float) or not math.isnan(na_value): + raise ValueError("Invalid NZ JSON string missing-value descriptor.") + try: + return pd.StringDtype(storage=payload["storage"], na_value=na_value) + except TypeError as exc: + # pandas 2.2's pyarrow_numpy storage already has NaN semantics, + # but its constructor does not accept the na_value keyword. + legacy_dtype = pd.StringDtype(storage=payload["storage"]) + if _encode_cell(legacy_dtype.na_value) == _encode_cell(na_value): + return legacy_dtype + raise ValueError( + "This pandas version cannot restore NaN-semantics string data." + ) from exc + if payload["kind"] != "pandas": + raise ValueError("Unsupported NZ JSON dtype descriptor.") + return pd.api.types.pandas_dtype(payload["name"]) + + +def _encode_index(index: pd.Index) -> dict[str, Any]: + if isinstance(index, pd.MultiIndex): + return { + "kind": "multi", + "levels": [_encode_index(level) for level in index.levels], + "codes": [code.tolist() for code in index.codes], + "names": [_encode_cell(name) for name in index.names], + } + name = _encode_cell(index.name) + if isinstance(index, pd.RangeIndex): + return { + "kind": "range", + "start": index.start, + "stop": index.stop, + "step": index.step, + "name": name, + } + return { + "kind": "index", + "values": [_encode_cell(value) for value in index], + "dtype": _encode_dtype(index.dtype), + "name": name, + } + + +def _decode_index(payload: dict[str, Any]) -> pd.Index: + if payload["kind"] == "multi": + return pd.MultiIndex( + levels=[_decode_index(level) for level in payload["levels"]], + codes=payload["codes"], + names=[_decode_cell(name) for name in payload["names"]], + verify_integrity=True, + ) + name = _decode_cell(payload["name"]) + if payload["kind"] == "range": + return pd.RangeIndex( + payload["start"], payload["stop"], payload["step"], name=name + ) + if payload["kind"] != "index": + raise ValueError("Unsupported NZ JSON index descriptor.") + return pd.Index( + [_decode_cell(value) for value in payload["values"]], + dtype=_decode_dtype(payload["dtype"]), + name=name, + ) + + +def encode_table(value: MicroDataFrame) -> dict[str, Any]: + table = pd.DataFrame(value) + if not table.columns.is_unique or not all( + isinstance(name, str) for name in table.columns + ): + raise ValueError("NZ JSON entity tables need unique string column names.") + return { + "schema": TABLE_SCHEMA, + "columns": list(table.columns), + "columns_name": _encode_cell(table.columns.name), + "dtypes": [_encode_dtype(dtype) for dtype in table.dtypes], + "index": _encode_index(table.index), + "data": [ + [_encode_cell(value) for value in row] + for row in table.itertuples(index=False, name=None) + ], + } + + +def decode_table(payload: dict[str, Any], entity: str) -> MicroDataFrame: + if payload.get("schema") != TABLE_SCHEMA: + raise ValueError("Unsupported NZ entity-table JSON schema.") + columns = payload["columns"] + if ( + not all(isinstance(name, str) for name in columns) + or len(set(columns)) != len(columns) + or len(columns) != len(payload["dtypes"]) + ): + raise ValueError("Invalid NZ entity-table JSON columns/dtypes.") + table = pd.DataFrame( + [[_decode_cell(value) for value in row] for row in payload["data"]], + index=_decode_index(payload["index"]), + columns=columns, + # Avoid pandas inferring float for integer+None object columns before + # the recorded dtype is restored: that can round integers above 2**53. + dtype=object, + ).astype( + { + name: _decode_dtype(dtype) + for name, dtype in zip(columns, payload["dtypes"], strict=True) + } + ) + table.columns.name = _decode_cell(payload["columns_name"]) + weight = f"{entity}_weight" + if weight not in table: + raise ValueError(f"NZ JSON {entity} table must contain effective {weight}.") + return MicroDataFrame(table, weights=weight) diff --git a/tests/test_nz_axiom_pilot.py b/tests/test_nz_axiom_pilot.py new file mode 100644 index 00000000..b33a3037 --- /dev/null +++ b/tests/test_nz_axiom_pilot.py @@ -0,0 +1,664 @@ +"""PolicyEngine boundary tests for the source-only New Zealand Axiom pilot. + +The ordinary tests stub only the unpublished runtime boundary. The opt-in +integration test executes the real Axiom module; no stub validates NZ law. +""" + +import json +import os +from copy import deepcopy +from dataclasses import dataclass +from decimal import Decimal +from hashlib import sha256 +from pathlib import Path + +import numpy as np +import pandas as pd +import pytest + +from policyengine.core import Simulation +from policyengine.tax_benefit_models.nz import ( + IWTC_CHANGE, + WFF_ABATEMENT_CHANGE, + AxiomNewZealandPilot, + PopulaceNewZealandDataset, +) +from policyengine.tax_benefit_models.nz import datasets as nz_datasets +from policyengine.tax_benefit_models.nz import model as nz_model +from policyengine.tax_benefit_models.nz import serialization as nz_serialization + +REQUIRED_INPUTS = { + "family_tax_credit_eldest_dependent_child_care_units": "decimal128(18,2)", + "family_tax_credit_subsequent_dependent_child_care_units": "decimal128(18,2)", + "family_tax_credit_entitlement_days": "int16", + "wff_family_scheme_income_for_relationship_period": "decimal128(18,2)", + "wff_family_credit_abatement_days": "int16", + "entitled_to_in_work_tax_credit": "bool", + "in_work_tax_credit_allowed_children_count": "int16", + "in_work_tax_credit_weekly_periods": "int16", + "child_tax_credit_for_entitlement_period": "decimal128(18,2)", + "parental_tax_credit_for_entitlement_period": "decimal128(18,2)", + "parental_tax_credit_additional_abatement": "decimal128(18,2)", +} +PADDING_INPUTS = [ + "best_start_abatement_days", + "best_start_child_care_fraction", + "best_start_entitlement_days", + "best_start_family_scheme_income_for_relationship_period", + "minimum_family_adjusted_income_tax_liability", + "minimum_family_amount_paid", + "minimum_family_amount_received", + "minimum_family_full_time_earner_weeks", + "minimum_family_scheme_income_attributable_to_full_time_weeks", + "minimum_family_tax_credit_weekly_periods", +] +PERIOD = {"start": "2026-04-01", "end": "2027-03-31", "kind": "tax_year"} +SOURCE_METADATA = {"data_build_id": "nz-test-build", "donor_country": "US"} +TEST_PROVENANCE = {"rulespec": {"commit": "a" * 40}, "runtime": {"hash": "b" * 64}} + + +class FakeWeightKind: + CALIBRATED = "calibrated" + DESIGN = "design" + + +class FakeWeights: + def __init__(self, *, values, kind): + self.values = np.asarray(values, dtype=float) + self.kind = kind + + +class FakeFrame: + """Strictly an integration seam double, not a rules evaluator.""" + + resolved_entities = [] + + def __init__(self, tables, schema, weights, strata=None, *, metadata=None): + self.tables = {name: table.copy() for name, table in tables.items()} + self.schema = schema + self.weights = weights + self.metadata = metadata or {} + self.strata = strata + self.entities = ("person", "household", "family") + self.weighted_entities = ("household",) + + def table(self, entity): + return self.tables[entity] + + def n(self, entity): + return len(self.tables[entity]) + + def weights_for(self, entity): + return self.weights[entity] + + def resolve_weights(self, entity): + type(self).resolved_entities.append(entity) + if entity == "household": + return self.weights[entity] + households = self.tables["household"]["household_id"] + lookup = dict(zip(households, self.weights["household"].values)) + membership = ( + self.tables["person"]["person_household_id"] + if entity == "person" + else self.tables["family"]["family_household_id"] + ) + return FakeWeights( + values=membership.map(lookup), kind=FakeWeightKind.CALIBRATED + ) + + +@dataclass +class FakeAxiomPeriod: + start: str + end: str + kind: str + + +class FakeAxiomEngine: + latest = None + + def __init__(self, module, *, schema, rulespec_roots): + self.module = Path(module) + self.schema = schema + self.rulespec_roots = tuple(rulespec_roots) + self.frame = None + self.period = None + type(self).latest = self + + def materialize(self, frame, variables, period): + self.frame = frame + self.period = period + assert variables == [WFF_ABATEMENT_CHANGE, IWTC_CHANGE] + # Deliberately non-statutory values: this tests weighted plumbing only. + return { + WFF_ABATEMENT_CHANGE: np.array([10.0, 20.0]), + IWTC_CHANGE: np.array([30.0, 40.0]), + } + + +class FakeEntityTableDataset: + """Reader seam double; real-codec coverage lives in the opt-in test.""" + + def __init__(self, *, file_path): + with pd.HDFStore(file_path, mode="r") as store: + self.tables = {entity: store[entity] for entity in nz_datasets.ENTITIES} + self.time_period = int(store["_time_period"].iloc[0]) + + +@pytest.fixture +def stub_runtime(monkeypatch): + FakeFrame.resolved_entities = [] + FakeAxiomEngine.latest = None + monkeypatch.setattr( + nz_datasets, + "_load_frame_runtime", + lambda: (FakeFrame, FakeWeightKind, FakeWeights, "nz-schema"), + ) + monkeypatch.setattr( + nz_model, + "_load_axiom_runtime", + lambda: (FakeAxiomEngine, FakeAxiomPeriod), + ) + monkeypatch.setattr( + nz_model, + "_build_runtime_provenance", + lambda _root: deepcopy(TEST_PROVENANCE), + ) + monkeypatch.setattr( + nz_datasets, "_load_dataset_reader", lambda: FakeEntityTableDataset + ) + + +@pytest.fixture +def rulespec_root(tmp_path): + root = tmp_path / "rulespec-nz" + module = root / nz_model.PILOT_MODULE + module.parent.mkdir(parents=True) + module.write_text("# runtime boundary fixture; not executable RuleSpec\n") + contract = { + "schema": "axiom/nz-official-budget-reform-transport/1", + "jurisdiction": "nz", + "period": { + "period_kind": "tax_year", + **{k: PERIOD[k] for k in ("start", "end")}, + }, + "runtime": {"rulespec_module": nz_model.PILOT_MODULE, "root_entity": "Family"}, + "input_contract": { + "engine_root_input_count": 21, + "required_target_inputs": [ + { + "name": name, + "entity": "family", + "dtype": dtype, + "missing": "fail_closed", + } + for name, dtype in REQUIRED_INPUTS.items() + ], + "adapter_padding_defaults": [ + {"name": name, "value": 0} for name in PADDING_INPUTS + ], + }, + "output_contract": { + "requested": [ + {"name": name, "entity": "family", "unit": "NZD"} + for name in (WFF_ABATEMENT_CHANGE, IWTC_CHANGE) + ], + "formula_owned_excluded_from_dataset": [WFF_ABATEMENT_CHANGE, IWTC_CHANGE], + }, + "official_score_bridge": { + "model_measure": "annual family entitlement change", + "official_measure": "forecast operating cost change", + "like_for_like_status": "bridge_required", + }, + } + contract_path = root / nz_model.TRANSPORT_CONTRACT + contract_path.parent.mkdir(parents=True) + contract_path.write_text(json.dumps(contract)) + return root + + +def _tables(): + person = pd.DataFrame( + { + "person_id": [1, 2, 3, 4, 5], + "person_household_id": [10, 10, 10, 20, 20], + "person_family_id": [100, 100, 100, 200, 200], + } + ) + household = pd.DataFrame({"household_id": [10, 20], "household_weight": [2.0, 5.0]}) + family = pd.DataFrame( + { + "family_id": [100, 200], + "family_household_id": [10, 20], + "family_tax_credit_eldest_dependent_child_care_units": [1.0, 1.0], + "family_tax_credit_subsequent_dependent_child_care_units": [1.0, 0.0], + "family_tax_credit_entitlement_days": [365, 365], + "wff_family_scheme_income_for_relationship_period": [50000.0, 100000.0], + "wff_family_credit_abatement_days": [365, 365], + "entitled_to_in_work_tax_credit": [False, True], + "in_work_tax_credit_allowed_children_count": [0, 1], + "in_work_tax_credit_weekly_periods": [0, 52], + "child_tax_credit_for_entitlement_period": [0.0, 0.0], + "parental_tax_credit_for_entitlement_period": [0.0, 0.0], + "parental_tax_credit_additional_abatement": [0.0, 0.0], + } + ) + return {"person": person, "household": household, "family": family} + + +def _write_dataset(path, tables, *, year=2026): + with pd.HDFStore(path, mode="w") as store: + for entity, table in tables.items(): + store[entity] = table + store.put("_time_period", pd.Series([year]), format="table") + store.get_storer("_time_period").attrs.policyengine_metadata_json = json.dumps( + SOURCE_METADATA + ) + return PopulaceNewZealandDataset( + name="nz-nonuniform-fixture", + description="Synthetic boundary fixture, not calibrated NZ data", + filepath=str(path), + year=year, + ) + + +@pytest.fixture +def dataset(tmp_path): + return _write_dataset(tmp_path / "nz.h5", _tables()) + + +def _run(dataset, rulespec_root): + model = AxiomNewZealandPilot(rulespec_root=str(rulespec_root)) + simulation = Simulation(dataset=dataset, tax_benefit_model_version=model) + simulation.run() + return simulation + + +def test_run_preserves_source_and_uses_family_weights( + dataset, rulespec_root, stub_runtime +): + before = sha256(Path(dataset.filepath).read_bytes()).hexdigest() + simulation = _run(dataset, rulespec_root) + output = simulation.output_dataset + + assert sha256(Path(dataset.filepath).read_bytes()).hexdigest() == before + assert output.filepath is None + assert output.is_output_dataset + assert output.year == 2026 + assert output.policy_period == PERIOD + assert output.metadata["data_build_id"] == SOURCE_METADATA["data_build_id"] + assert dataset.metadata == SOURCE_METADATA + assert "family" in FakeFrame.resolved_entities + assert output.data.family[WFF_ABATEMENT_CHANGE].sum() == 120.0 + assert output.data.family[IWTC_CHANGE].sum() == 260.0 + np.testing.assert_array_equal(output.data.person["person_weight"], [2, 2, 2, 5, 5]) + np.testing.assert_array_equal(output.data.family["family_weight"], [2, 5]) + + engine = FakeAxiomEngine.latest + assert engine.rulespec_roots == (rulespec_root.resolve(),) + assert engine.period == FakeAxiomPeriod(**PERIOD) + assert engine.frame.weighted_entities == ("household",) + assert "family_weight" not in engine.frame.table("family") + assert "person_weight" not in engine.frame.table("person") + assert all(name in engine.frame.table("family") for name in PADDING_INPUTS) + assert not any(name in dataset.data.family for name in PADDING_INPUTS) + receipt = output.metadata["policyengine_axiom_runs"][-1] + assert receipt["provenance"] == TEST_PROVENANCE + assert receipt["policy_period"] == PERIOD + assert receipt["official_score_bridge"]["like_for_like_status"] == "bridge_required" + simulation.save() + assert sha256(Path(dataset.filepath).read_bytes()).hexdigest() == before + + +def test_dataset_rejects_mismatched_period(dataset, stub_runtime): + dataset.year = 2027 + with pytest.raises(ValueError, match="period mismatch"): + dataset.load() + + +@pytest.mark.parametrize("entity", ["person", "family"]) +def test_source_dataset_rejects_extra_weight_vectors(tmp_path, stub_runtime, entity): + tables = _tables() + tables[entity][f"{entity}_weight"] = 1.0 + dataset = _write_dataset(tmp_path / "bad_weights.h5", tables) + with pytest.raises(ValueError, match="sole stored weight"): + dataset.load() + + +def test_dataset_rejects_family_crossing_households(tmp_path, stub_runtime): + tables = _tables() + tables["person"].loc[4, "person_family_id"] = 100 + dataset = _write_dataset(tmp_path / "cross_household.h5", tables) + with pytest.raises(ValueError, match="exactly one household"): + dataset.load() + + +@pytest.mark.parametrize("problem", ["missing", "null", "wrong_dtype"]) +def test_substantive_family_inputs_never_default( + dataset, rulespec_root, stub_runtime, problem +): + dataset.load() + name = "entitled_to_in_work_tax_credit" + if problem == "missing": + dataset.data.family = dataset.data.family.drop(columns=[name]) + elif problem == "null": + dataset.data.family[name] = dataset.data.family[name].astype("boolean") + dataset.data.family.loc[0, name] = None + else: + dataset.data.family[name] = [0, 1] + with pytest.raises(ValueError, match=name): + _run(dataset, rulespec_root) + + +def test_formula_owned_inputs_are_rejected(dataset, rulespec_root, stub_runtime): + dataset.load() + dataset.data.family[WFF_ABATEMENT_CHANGE] = 999.0 + with pytest.raises(ValueError, match="formula-owned"): + _run(dataset, rulespec_root) + + +def test_changed_runtime_requires_new_model( + dataset, rulespec_root, stub_runtime, monkeypatch +): + model = AxiomNewZealandPilot(rulespec_root=str(rulespec_root)) + monkeypatch.setattr( + nz_model, "_build_runtime_provenance", lambda _root: {"changed": True} + ) + simulation = Simulation(dataset=dataset, tax_benefit_model_version=model) + with pytest.raises(RuntimeError, match="changed"): + simulation.run() + + +def test_extra_simulation_controls_fail_closed(dataset, rulespec_root, stub_runtime): + model = AxiomNewZealandPilot(rulespec_root=str(rulespec_root)) + simulation = Simulation( + dataset=dataset, + tax_benefit_model_version=model, + extra_variables={"family": ["unknown"]}, + ) + with pytest.raises(ValueError, match="extra_variables"): + simulation.run() + + +def test_model_configuration_json_roundtrip(rulespec_root, stub_runtime): + model = AxiomNewZealandPilot(rulespec_root=str(rulespec_root)) + restored = AxiomNewZealandPilot.model_validate_json(model.model_dump_json()) + assert restored.id == model.id + assert restored.runtime_provenance == model.runtime_provenance + assert restored.transport_contract == model.transport_contract + + +def test_saved_model_refuses_runtime_drift(rulespec_root, stub_runtime, monkeypatch): + saved = AxiomNewZealandPilot(rulespec_root=str(rulespec_root)).model_dump_json() + monkeypatch.setattr( + nz_model, "_build_runtime_provenance", lambda _root: {"changed": True} + ) + with pytest.raises(ValueError, match="changed"): + AxiomNewZealandPilot.model_validate_json(saved) + + +def test_tampered_effective_weights_fail(dataset, rulespec_root, stub_runtime): + dataset.load() + dataset.data.family["family_weight"] = [1.0, 1.0] + with pytest.raises(ValueError, match="Frame-resolved"): + _run(dataset, rulespec_root) + + +def test_output_datasets_cannot_be_reused(dataset, rulespec_root, stub_runtime): + output = _run(dataset, rulespec_root).output_dataset + with pytest.raises(ValueError, match="reused"): + _run(output, rulespec_root) + + +@pytest.mark.parametrize("values", [[1.0], [float("nan"), 2.0]]) +def test_invalid_engine_outputs_fail( + dataset, rulespec_root, stub_runtime, monkeypatch, values +): + monkeypatch.setattr( + FakeAxiomEngine, + "materialize", + lambda *_: { + WFF_ABATEMENT_CHANGE: np.array(values), + IWTC_CHANGE: np.array([0.0, 0.0]), + }, + ) + with pytest.raises(ValueError, match="invalid NZ output"): + _run(dataset, rulespec_root) + + +def test_wrong_policy_build_year_fails_closed(tmp_path, rulespec_root, stub_runtime): + dataset = _write_dataset(tmp_path / "wrong_year.h5", _tables(), year=2025) + with pytest.raises(ValueError, match="2026"): + _run(dataset, rulespec_root) + + +def test_expected_source_hash_is_verified(dataset, stub_runtime): + dataset.source_sha256 = "0" * 64 + with pytest.raises(ValueError, match="SHA-256"): + dataset.load() + + +def test_decimal_inputs_are_preserved(dataset, rulespec_root, stub_runtime): + dataset.load() + name = "wff_family_scheme_income_for_relationship_period" + dataset.data.family[name] = [Decimal("50000.00"), Decimal("100000.00")] + simulation = _run(dataset, rulespec_root) + assert list(FakeAxiomEngine.latest.frame.table("family")[name]) == [ + Decimal("50000.00"), + Decimal("100000.00"), + ] + assert simulation.output_dataset.data.family[WFF_ABATEMENT_CHANGE].sum() == 120.0 + + +@pytest.mark.parametrize( + "value", [Decimal("1.001"), Decimal("10000000000000000"), "50000.00"] +) +def test_out_of_contract_decimal_inputs_fail( + dataset, rulespec_root, stub_runtime, value +): + dataset.load() + name = "wff_family_scheme_income_for_relationship_period" + dataset.data.family[name] = [value, value] + with pytest.raises(ValueError, match=name): + _run(dataset, rulespec_root) + + +def test_nonzero_adapter_padding_fails(dataset, rulespec_root, stub_runtime): + dataset.load() + dataset.data.family[PADDING_INPUTS[0]] = 1 + with pytest.raises(ValueError, match="padding"): + _run(dataset, rulespec_root) + + +def test_decimal_adapter_padding_is_valid(dataset, rulespec_root, stub_runtime): + dataset.load() + dataset.data.family[PADDING_INPUTS[0]] = [Decimal("0.00"), Decimal("0.00")] + simulation = _run(dataset, rulespec_root) + assert simulation.output_dataset.data.family[WFF_ABATEMENT_CHANGE].sum() == 120.0 + + +def test_mutated_input_has_distinct_execution_fingerprint( + dataset, rulespec_root, stub_runtime +): + first = _run(dataset, rulespec_root).output_dataset.metadata[ + "policyengine_axiom_runs" + ][-1] + dataset.data.family["wff_family_scheme_income_for_relationship_period"] = [ + 50001.0, + 100000.0, + ] + second = _run(dataset, rulespec_root).output_dataset.metadata[ + "policyengine_axiom_runs" + ][-1] + assert first["input_artifact_sha256"] == second["input_artifact_sha256"] + assert first["input_frame_sha256"] != second["input_frame_sha256"] + json.dumps(second, allow_nan=False) + + +def test_output_dataset_json_roundtrip_preserves_tables_and_receipt( + dataset, rulespec_root, stub_runtime +): + dataset.load() + name = "wff_family_scheme_income_for_relationship_period" + dataset.data.family[name] = [Decimal("50000.00"), Decimal("100000.00")] + output = _run(dataset, rulespec_root).output_dataset + # These extra columns exercise JSON transport, not policy inputs. + output.data.family["nullable_marker"] = pd.Series([True, pd.NA], dtype="boolean") + output.data.family["source_category"] = pd.Categorical( + ["b", "a"], categories=["a", "b", "unused"], ordered=True + ) + encoded = output.model_dump_json() + restored = PopulaceNewZealandDataset.model_validate_json(encoded) + assert restored.filepath is None + assert restored.is_output_dataset + assert restored.policy_period == output.policy_period + assert restored.metadata == output.metadata + for entity in nz_datasets.ENTITIES: + pd.testing.assert_frame_equal( + pd.DataFrame(getattr(restored.data, entity)), + pd.DataFrame(getattr(output.data, entity)), + ) + assert list(restored.data.family[name]) == [ + Decimal("50000.00"), + Decimal("100000.00"), + ] + assert restored.data.family[WFF_ABATEMENT_CHANGE].sum() == 120.0 + + +def test_year_data_json_roundtrip_preserves_named_nonrange_indices(stub_runtime): + tables = _tables() + tables["family"].index = pd.Index(["second", "first"], name="family_row") + tables["person"].index = pd.MultiIndex.from_tuples( + [("a", 1), ("a", 2), ("a", 3), ("b", 1), ("b", 2)], + names=["group", "row"], + ) + data = nz_datasets.year_data_from_frame(nz_datasets.frame_from_tables(tables)) + restored = nz_datasets.NZYearData.model_validate_json(data.model_dump_json()) + for entity in nz_datasets.ENTITIES: + pd.testing.assert_frame_equal( + pd.DataFrame(getattr(restored, entity)), + pd.DataFrame(getattr(data, entity)), + ) + + +def test_json_codec_does_not_coerce_large_object_integers( + dataset, rulespec_root, stub_runtime +): + output = _run(dataset, rulespec_root).output_dataset + large_id = 9007199254740993 + output.data.family["optional_id"] = pd.Series([large_id, None], dtype=object) + restored = PopulaceNewZealandDataset.model_validate_json(output.model_dump_json()) + values = restored.data.family["optional_id"] + assert values.iloc[0] == large_id + assert type(values.iloc[0]) is int + assert values.iloc[1] is None + + +@pytest.mark.parametrize("missing_value", [pd.NA, np.nan], ids=["NA", "NaN"]) +def test_json_codec_preserves_string_missing_value_semantics( + stub_runtime, missing_value +): + try: + dtype = pd.StringDtype(storage="python", na_value=missing_value) + except TypeError: + if missing_value is not pd.NA: + pytest.skip("Installed pandas does not support NaN-semantics strings") + dtype = pd.StringDtype(storage="python") + tables = _tables() + family = tables["family"] + family["optional_text"] = pd.Series(["a", None], dtype=dtype) + family["source_category"] = pd.Categorical( + ["a", None], categories=pd.Index(["a", "unused"], dtype=dtype), ordered=True + ) + family.index = pd.Index(["second", None], dtype=dtype, name="family_row") + data = nz_datasets.year_data_from_frame(nz_datasets.frame_from_tables(tables)) + restored = nz_datasets.NZYearData.model_validate_json(data.model_dump_json()) + pd.testing.assert_frame_equal( + pd.DataFrame(restored.family), pd.DataFrame(data.family) + ) + + +def test_json_codec_reads_original_string_descriptor(): + dtype = nz_serialization._decode_dtype({"kind": "string", "storage": "python"}) + assert dtype == pd.StringDtype(storage="python") + assert dtype.na_value is pd.NA + + +def test_json_codec_rejects_invalid_string_missing_value_descriptor(): + with pytest.raises(ValueError, match="missing-value descriptor"): + nz_serialization._decode_dtype( + {"kind": "string", "storage": "python", "na_value": "missing"} + ) + + +@pytest.mark.parametrize("supports_nan", [True, False]) +def test_json_codec_legacy_string_constructor_preserves_semantics( + monkeypatch, supports_nan +): + # Exercise the old constructor signature without requiring an Arrow + # installation: only dtype metadata is involved in this fallback. + try: + legacy_dtype = pd.StringDtype( + storage="python", na_value=np.nan if supports_nan else pd.NA + ) + except TypeError: + if supports_nan: + pytest.skip("Installed pandas does not support NaN-semantics strings") + legacy_dtype = pd.StringDtype(storage="python") + + def legacy_constructor(*, storage): + assert storage == "pyarrow_numpy" + return legacy_dtype + + monkeypatch.setattr(nz_serialization.pd, "StringDtype", legacy_constructor) + payload = { + "kind": "string", + "storage": "pyarrow_numpy", + "na_value": {"_type": "float", "value": "nan"}, + } + if supports_nan: + assert nz_serialization._decode_dtype(payload) is legacy_dtype + else: + with pytest.raises(ValueError, match="cannot restore NaN-semantics"): + nz_serialization._decode_dtype(payload) + + +@pytest.mark.skipif( + os.environ.get("RUN_NZ_AXIOM_INTEGRATION") != "1", + reason="requires compatible source-only Microcosm/Axiom runtime", +) +def test_real_axiom_source_stack_executes_both_reforms(tmp_path): + from microcosm.frame.adapters.axiom import AxiomEntityTableDataset + + root = Path(os.environ["RULESPEC_NZ_ROOT"]).resolve() + tables = _tables() + family_inputs = tables["family"] + for name, dtype in REQUIRED_INPUTS.items(): + if dtype == "decimal128(18,2)": + family_inputs = family_inputs.assign( + **{name: family_inputs[name].map(lambda value: Decimal(str(value)))} + ) + family_inputs = family_inputs.assign( + entitled_to_in_work_tax_credit=family_inputs[ + "entitled_to_in_work_tax_credit" + ].astype("boolean") + ) + tables["family"] = family_inputs + path = tmp_path / "real_nz_fixture.h5" + # Use the producer's actual codec, including exact Decimal persistence + # and nullable booleans, rather than a parallel PolicyEngine writer. + AxiomEntityTableDataset(tables=tables, time_period=2026).save(path) + dataset = PopulaceNewZealandDataset( + name="nz-real-runtime-fixture", + description="Synthetic integration fixture, not calibrated NZ data", + filepath=str(path), + year=2026, + ) + simulation = _run(dataset, root) + family = simulation.output_dataset.data.family + # These are the two upstream RuleSpec companion cases, with nonuniform + # household weights 2 and 5. No official population score is asserted. + np.testing.assert_allclose(np.asarray(family[WFF_ABATEMENT_CHANGE]), [568.5, 0.0]) + np.testing.assert_allclose(np.asarray(family[IWTC_CHANGE]), [0.0, 438.5]) + assert family[WFF_ABATEMENT_CHANGE].sum() == 1137.0 + assert family[IWTC_CHANGE].sum() == 2192.5