diff --git a/changelog.d/447.added.md b/changelog.d/447.added.md new file mode 100644 index 00000000..e63ff8e0 --- /dev/null +++ b/changelog.d/447.added.md @@ -0,0 +1 @@ +Belgium pilot: add an Axiom-backed model over Microcosm-BE entity tables with content-bound RuleSpec/engine provenance, separately labeled data and policy periods, metadata-preserving HDF5 round trips, nonuniform calibrated-weight coverage, and in-memory outputs that cannot overwrite their inputs. diff --git a/examples/belgium_axiom_pilot.py b/examples/belgium_axiom_pilot.py new file mode 100644 index 00000000..f13941dc --- /dev/null +++ b/examples/belgium_axiom_pilot.py @@ -0,0 +1,75 @@ +"""Belgium population microsimulation through the Axiom rules engine. + +Runs the calibrated populace-be pilot dataset (populace-us support records +reweighted to Statbel/SPF/ONSS/ONEM targets from PolicyEngine/ledger) +through the rulespec-be composed worker pipeline, and scores the aggregates +against Belgian administrative facts. + +This is a demonstration of the engine channel, not a certified Belgian +model: the support records are American, and coverage is the worker slice +(employee SSC with the statutory low-wage work bonus, and individual PIT +before withholding with the fiscal work-bonus credit). + +The pilot's unpublished source stack currently requires Python 3.14: install +``microcosm-frame`` from PolicyEngine/microcosm, plus the +``axiom-rules-engine`` Python wrapper and dense extension from source. Their +current checkouts must agree on the canonical RuleSpec-root loader API. The +output is intentionally in memory (``filepath=None``) so it cannot overwrite +the input artifact; set a distinct output filepath explicitly before saving. + +Usage:: + + POPULACE_BE_DATASET=.../populace_be_pilot_2026.h5 \\ + RULESPEC_BE_ROOT=.../rulespec-be \\ + uv run python examples/belgium_axiom_pilot.py +""" + +import os + +from policyengine.core.simulation import Simulation +from policyengine.tax_benefit_models.be import ( + EMPLOYEE_SSC, + PIT_BEFORE_WITHHOLDING, + AxiomBelgiumPilot, + PopulaceBelgiumDataset, +) + +DATASET = os.environ["POPULACE_BE_DATASET"] +RULESPEC = os.environ.get("RULESPEC_BE_ROOT", "~/TheAxiomFoundation/rulespec-be") + +# Ledger facts (PolicyEngine/ledger, Belgian publisher packages) +ONSS_WORKER_CONTRIBUTIONS_2024 = 20_836_582_673 +SPF_PIT_BEFORE_WITHHOLDING_2023 = 62_840_116_134 + +dataset = PopulaceBelgiumDataset( + name="populace-be-pilot", + description="populace-us support reweighted to Belgian ledger targets", + filepath=DATASET, + year=2026, +) +model_version = AxiomBelgiumPilot(rulespec_root=RULESPEC, period=2025) + +simulation = Simulation(dataset=dataset, tax_benefit_model_version=model_version) +simulation.run() + +output = simulation.output_dataset +person = output.data.person +ssc_net = person[EMPLOYEE_SSC].sum() +pit = person[PIT_BEFORE_WITHHOLDING].sum() + +print("Belgium pilot (Axiom engine over populace-be, worker slice)") +print( + f" data vintage {output.year}; policy period {output.policy_period}; " + f"model {model_version.version}" +) +print(f" employee SSC after work bonus EUR {ssc_net / 1e9:6.2f}B") +print(" (ONSS 2024 contributions: EUR 20.84B; the bonus is outsized here") +print(" because the US-support wage distribution is low-wage-heavy)") +print( + f" PIT before withholding EUR {pit / 1e9:6.2f}B " + "(SPF 2023, all PIT: EUR 62.84B)" +) +print( + f" PIT vs SPF ratio {pit / SPF_PIT_BEFORE_WITHHOLDING_2023:.3f}" + " (worker slice only)" +) diff --git a/src/policyengine/tax_benefit_models/be/__init__.py b/src/policyengine/tax_benefit_models/be/__init__.py new file mode 100644 index 00000000..74044832 --- /dev/null +++ b/src/policyengine/tax_benefit_models/be/__init__.py @@ -0,0 +1,29 @@ +"""Belgium pilot: Axiom rules engine over Microcosm entity tables. + +The first non-policyengine-core country in policyengine.py. Statutes are +encoded as RuleSpec YAML (TheAxiomFoundation/rulespec-be), compiled and +executed by axiom-rules-engine, and driven over Microcosm-BE entity tables. +See ``examples/belgium_axiom_pilot.py`` for the end-to-end population run +and ``model.py`` for scope and source-install requirements. +""" + +from .datasets import BEYearData, PopulaceBelgiumDataset +from .model import ( + EMPLOYEE_SSC, + PIT_BEFORE_WITHHOLDING, + REMUNERATION, + AxiomBelgium, + AxiomBelgiumPilot, + be_model, +) + +__all__ = [ + "AxiomBelgium", + "AxiomBelgiumPilot", + "BEYearData", + "EMPLOYEE_SSC", + "PIT_BEFORE_WITHHOLDING", + "PopulaceBelgiumDataset", + "REMUNERATION", + "be_model", +] diff --git a/src/policyengine/tax_benefit_models/be/datasets.py b/src/policyengine/tax_benefit_models/be/datasets.py new file mode 100644 index 00000000..04bb40d0 --- /dev/null +++ b/src/policyengine/tax_benefit_models/be/datasets.py @@ -0,0 +1,190 @@ +"""Belgium pilot dataset: Microcosm entity tables with calibrated weights. + +The pilot layout has two entities (person, household), mirroring the +Microcosm ``BE_SCHEMA``. Files are pandas HDF5 stores with ``person``, +``household``, and ``_time_period`` keys. The canonical Microcosm-BE layout +stores calibrated weights once, on the household table; PolicyEngine derives +effective person weights through ``person_household_id`` when loading. Legacy +files may also carry ``person_weight``, but that redundant copy must match the +household weights exactly. Dataset metadata and the optional policy period are +stored as attributes on the ``_time_period`` record so the current +``microcosm.frame`` reader can ignore them while PolicyEngine can round-trip +them. +""" + +import json +from pathlib import Path +from typing import Any, ClassVar, Optional + +import numpy as np +import pandas as pd +from microdf import MicroDataFrame +from pydantic import ConfigDict, Field + +from policyengine.core import Dataset, YearData + + +def _person_with_household_weights( + person: pd.DataFrame, + household: pd.DataFrame, +) -> pd.DataFrame: + """Return persons with effective weights derived from household weights.""" + person_membership = "person_household_id" + household_id = "household_id" + household_weight = "household_weight" + for table_name, table, required in ( + ("person", person, person_membership), + ("household", household, household_id), + ("household", household, household_weight), + ): + if required not in table.columns: + raise ValueError( + f"Belgium {table_name} table is missing required column {required!r}." + ) + duplicate_ids = household[household_id].duplicated(keep=False) + if duplicate_ids.any(): + values = household.loc[duplicate_ids, household_id].drop_duplicates().tolist() + raise ValueError(f"Belgium household_id values must be unique; found {values}.") + if household[household_weight].isna().any(): + raise ValueError("Belgium household_weight values must not be null.") + + weight_by_household = household.set_index(household_id)[household_weight] + derived = person[person_membership].map(weight_by_household) + if derived.isna().any(): + missing = ( + person.loc[derived.isna(), person_membership].drop_duplicates().tolist() + ) + raise ValueError( + "Belgium person_household_id values must resolve to household rows; " + f"missing {missing}." + ) + + if "person_weight" in person.columns: + legacy = pd.to_numeric(person["person_weight"], errors="coerce") + if legacy.isna().any() or not np.array_equal( + legacy.to_numpy(dtype=float), + derived.to_numpy(dtype=float), + ): + raise ValueError( + "Legacy Belgium person_weight values do not exactly match the " + "effective household_weight values." + ) + + result = person.copy() + result["person_weight"] = derived.to_numpy(copy=True) + return result + + +class BEYearData(YearData): + """Entity-level data for a single Belgian year.""" + + model_config = ConfigDict(arbitrary_types_allowed=True) + + person: MicroDataFrame + household: MicroDataFrame + + @property + def entity_data(self) -> dict[str, MicroDataFrame]: + return {"person": self.person, "household": self.household} + + +class PopulaceBelgiumDataset(Dataset): + """Belgium pilot dataset loaded from a Microcosm-BE HDF5 artifact. + + ``year`` identifies the input dataset vintage. ``policy_period`` is set on + derived outputs and identifies the year of law executed by Axiom. Keeping + both fields prevents a 2025-policy run over a 2026-vintage population from + being mislabeled as though those were the same period. + """ + + data: Optional[BEYearData] = None + metadata: dict[str, Any] = Field(default_factory=dict) + policy_period: Optional[int] = None + + _TIME_PERIOD_KEY: ClassVar[str] = "_time_period" + _METADATA_ATTRIBUTE: ClassVar[str] = "policyengine_metadata_json" + _POLICY_PERIOD_ATTRIBUTE: ClassVar[str] = "policyengine_policy_period" + + def load(self) -> None: + if self.filepath is None: + raise ValueError("Cannot load a Belgium pilot dataset without a filepath.") + with pd.HDFStore(self.filepath, mode="r") as store: + person = store["person"] + household = store["household"] + if f"/{self._TIME_PERIOD_KEY}" in store.keys(): + stored_year = int(store[self._TIME_PERIOD_KEY].iloc[0]) + if stored_year != self.year: + raise ValueError( + "Belgium dataset period mismatch: " + f"constructor year={self.year}, HDF5 " + f"{self._TIME_PERIOD_KEY}={stored_year}." + ) + attributes = store.get_storer(self._TIME_PERIOD_KEY).attrs + metadata_json = getattr(attributes, self._METADATA_ATTRIBUTE, None) + if metadata_json is not None: + stored_metadata = json.loads(str(metadata_json)) + if self.metadata and self.metadata != stored_metadata: + raise ValueError( + "Belgium dataset metadata differs from the metadata " + "stored in its HDF5 artifact." + ) + self.metadata = stored_metadata + stored_policy_period = getattr( + attributes, self._POLICY_PERIOD_ATTRIBUTE, None + ) + if stored_policy_period is not None: + stored_policy_period = int(stored_policy_period) + if ( + self.policy_period is not None + and self.policy_period != stored_policy_period + ): + raise ValueError( + "Belgium dataset policy-period mismatch: " + f"constructor policy_period={self.policy_period}, " + f"HDF5 policy period={stored_policy_period}." + ) + self.policy_period = stored_policy_period + person = _person_with_household_weights(person, household) + self.data = BEYearData( + person=MicroDataFrame(person, weights="person_weight"), + household=MicroDataFrame(household, weights="household_weight"), + ) + + def save(self) -> None: + if self.data is None: + raise ValueError("No data to save.") + if self.filepath is None: + raise ValueError("Cannot save a Belgium pilot dataset without a filepath.") + # Serialize before opening in mode="w": invalid metadata must not + # truncate an existing destination. + metadata_json = json.dumps( + self.metadata, + allow_nan=False, + separators=(",", ":"), + sort_keys=True, + ) + household = pd.DataFrame(self.data.household) + person = _person_with_household_weights( + pd.DataFrame(self.data.person), + household, + ) + filepath = Path(self.filepath) + filepath.parent.mkdir(parents=True, exist_ok=True) + with pd.HDFStore(filepath, mode="w") as store: + # Microcosm carries only explicit entity weights. Person weights + # are inherited from households and reconstructed on load. + store["person"] = person.drop(columns=["person_weight"]) + store["household"] = household + store.put( + self._TIME_PERIOD_KEY, + pd.Series([self.year]), + format="table", + ) + attributes = store.get_storer(self._TIME_PERIOD_KEY).attrs + setattr(attributes, self._METADATA_ATTRIBUTE, metadata_json) + if self.policy_period is not None: + setattr( + attributes, + self._POLICY_PERIOD_ATTRIBUTE, + int(self.policy_period), + ) diff --git a/src/policyengine/tax_benefit_models/be/model.py b/src/policyengine/tax_benefit_models/be/model.py new file mode 100644 index 00000000..c4fe9e85 --- /dev/null +++ b/src/policyengine/tax_benefit_models/be/model.py @@ -0,0 +1,363 @@ +"""Axiom-backed Belgium pilot model version. + +Unlike the US and UK model versions, Belgium runs on the Axiom rules +engine: statutes encoded as RuleSpec YAML in TheAxiomFoundation/rulespec-be, +compiled and executed by ``axiom-rules-engine``, and driven over populace +entity tables through the ``microcosm.frame`` Axiom adapter. There is no +policyengine-core country package and no certified release manifest, so this +version subclasses ``TaxBenefitModelVersion`` directly and stays outside the +managed-release machinery. + +Requirements (neither is on PyPI yet): + +- ``microcosm-frame`` from PolicyEngine/microcosm + (``packages/microcosm-frame``; Python 3.13+) +- ``axiom-rules-engine`` from TheAxiomFoundation/axiom-rules-engine (PyO3 + dense extension; its source Python package currently requires Python 3.14) +- a checkout of TheAxiomFoundation/rulespec-be, passed as ``rulespec_root`` + +These are source dependencies with no jointly released compatibility set. +The current Microcosm adapter and Axiom checkout must agree on the dense +``CompiledDenseProgram.from_file`` signature (including canonical RuleSpec +roots). The deterministic tests in this repository cover PolicyEngine's +integration invariants without pretending to execute statutes; the opt-in +source-stack test executes the real engine once compatible checkouts and the +dense extension are installed. + +Scope: the composed worker pipeline only — employee social security +contributions (13.07 percent ordinary worker contribution) and personal +income tax before withholding for wage earners under individual assessment. +Dependants, joint assessment, other income categories, and employment tax +reductions are not yet encoded (TheAxiomFoundation/rulespec-be#1). +""" + +import json +from copy import deepcopy +from hashlib import sha256 +from importlib import metadata as importlib_metadata +from importlib.util import find_spec +from pathlib import Path +from typing import TYPE_CHECKING, Any, ClassVar, Optional, Union + +import pandas as pd +from microdf import MicroDataFrame + +from policyengine.core import TaxBenefitModel, TaxBenefitModelVersion + +from .datasets import ( + BEYearData, + PopulaceBelgiumDataset, + _person_with_household_weights, +) + +if TYPE_CHECKING: + from policyengine.core.simulation import Simulation + +PILOT_MODULE = "be/statutes/income_tax/individual/pilot_worker_oracle_pipeline.yaml" +REMUNERATION = "belgium_pit_article_23_worker_remuneration" +EMPLOYEE_SSC = "belgium_employee_social_security_ordinary_worker_contribution" +PIT_BEFORE_WITHHOLDING = "belgium_pit_pilot_federal_and_local_tax_before_withholding" + +#: Pipeline inputs the pilot supplies as scalars when the dataset does not +#: carry them (rulespec stage boundaries are supplied inputs by convention). +#: The work-bonus reference wage bridges 0 -> worker remuneration inside the +#: pipeline, so the scalar default keeps the statutory low-wage bonus active. +SUPPLIED_DEFAULTS: dict[str, Union[float, bool]] = { + "belgium_pit_article_466_tax_share_on_nonprofessional_movable_income": 0.0, + "belgium_pit_article_466bis_hypothetical_total_tax_if_treaty_exempt_foreign_professional_income_were_belgian": 0.0, + "belgium_pit_article_466bis_treaty_exempt_foreign_professional_income_base_applies": False, + "belgium_worker_work_bonus_supplied_reference_annual_remuneration": 0.0, + "belgium_pit_communal_additional_tax_rate": 0.0, + "belgium_pit_agglomeration_additional_tax_rate": 0.0, +} + + +class AxiomBelgium(TaxBenefitModel): + id: str = "axiom-rulespec-be" + description: str = ( + "Belgium tax rules encoded as RuleSpec (TheAxiomFoundation/rulespec-be), " + "executed by the Axiom rules engine over populace entity tables." + ) + + +be_model = AxiomBelgium() + + +class AxiomBelgiumPilot(TaxBenefitModelVersion): + """Pilot Belgium model version: worker SSC and PIT via Axiom.""" + + country_code: ClassVar[str] = "be" + + rulespec_root: str + period: Optional[int] = None + output_variables: list[str] = [EMPLOYEE_SSC, PIT_BEFORE_WITHHOLDING] + communal_additional_tax_rate: float = 0.0 + runtime_provenance: dict[str, Any] + + def __init__(self, **kwargs: Any) -> None: + rulespec_root = Path(kwargs["rulespec_root"]).expanduser().resolve() + runtime_provenance = _build_runtime_provenance(rulespec_root) + version = _provenance_version(runtime_provenance) + kwargs["rulespec_root"] = str(rulespec_root) + kwargs["runtime_provenance"] = runtime_provenance + kwargs["model"] = be_model + # This is a content identity, not the unrelated 0.1.0 version shared + # by the source-only Python wrapper and dense-extension distributions. + kwargs["version"] = version + kwargs["id"] = f"{be_model.id}@{version}" + super().__init__(**kwargs) + + def run(self, simulation: "Simulation") -> "Simulation": + current_provenance = _build_runtime_provenance(Path(self.rulespec_root)) + if current_provenance != self.runtime_provenance: + raise RuntimeError( + "The RuleSpec/Axiom runtime changed after this Belgium model " + "version was constructed; construct a new AxiomBelgiumPilot " + "so its content identity matches the code that will execute." + ) + Frame, WeightKind, Weights, BE_SCHEMA, AxiomEngine = _load_axiom_runtime() + + rulespec_root = Path(self.rulespec_root) + module = rulespec_root / PILOT_MODULE + if not module.exists(): + raise FileNotFoundError( + f"rulespec-be pilot module not found at {module}; pass a " + "checkout of TheAxiomFoundation/rulespec-be as rulespec_root." + ) + + dataset = simulation.dataset + assert isinstance(dataset, PopulaceBelgiumDataset) + if dataset.data is None: + dataset.load() + assert dataset.data is not None + + household = pd.DataFrame(dataset.data.household).copy() + person = _person_with_household_weights( + pd.DataFrame(dataset.data.person), + household, + ) + for name, value in SUPPLIED_DEFAULTS.items(): + if name not in person.columns: + person[name] = value + person["belgium_pit_communal_additional_tax_rate"] = ( + self.communal_additional_tax_rate + ) + + weights = { + "household": Weights( + values=household["household_weight"].to_numpy(), + kind=WeightKind.CALIBRATED, + ) + } + # The frame kernel owns weight columns (typed Weights vectors); + # they stay on the pe.py-side MicroDataFrames only. + frame = Frame( + { + "person": person.drop(columns=["person_weight"]), + "household": household.drop(columns=["household_weight"]), + }, + BE_SCHEMA, + weights, + ) + policy_period = self.period if self.period is not None else dataset.year + engine = AxiomEngine(module, rulespec_roots=(rulespec_root,)) + outputs = engine.materialize(frame, self.output_variables, policy_period) + for name, values in outputs.items(): + person[name] = values + + output_metadata = deepcopy(dataset.metadata) + prior_runs = output_metadata.get("policyengine_axiom_runs", []) + if not isinstance(prior_runs, list): + raise ValueError( + "Belgium dataset metadata field 'policyengine_axiom_runs' " + "must be a list when present." + ) + output_metadata["policyengine_axiom_runs"] = [ + *deepcopy(prior_runs), + { + "dataset_year": dataset.year, + "policy_period": policy_period, + "model_version": self.version, + "configuration": { + "communal_additional_tax_rate": (self.communal_additional_tax_rate), + "output_variables": list(self.output_variables), + }, + "provenance": deepcopy(self.runtime_provenance), + }, + ] + + simulation.output_dataset = PopulaceBelgiumDataset( + id=simulation.id, + name=dataset.name, + description=dataset.description, + # Derived in-memory output: never alias the source HDF5 path. A + # caller may choose a new destination explicitly before saving. + filepath=None, + year=dataset.year, + policy_period=policy_period, + is_output_dataset=True, + metadata=output_metadata, + data=BEYearData( + person=MicroDataFrame(person, weights="person_weight"), + household=MicroDataFrame(household, weights="household_weight"), + ), + ) + return simulation + + def save(self, simulation: "Simulation") -> None: + """Pilot simulations are recomputed, not persisted.""" + + def load(self, simulation: "Simulation") -> None: + raise FileNotFoundError("Pilot simulations are recomputed, not persisted.") + + +def _load_axiom_runtime() -> tuple[Any, Any, Any, Any, Any]: + """Load the source-only Microcosm/Axiom runtime at the execution boundary.""" + + try: + from microcosm.frame import Frame, WeightKind, Weights + from microcosm.frame.adapters.axiom import BE_SCHEMA, AxiomEngine + except ImportError as error: + raise ImportError( + "The Belgium pilot needs microcosm-frame " + "(PolicyEngine/microcosm, packages/microcosm-frame) and " + "axiom-rules-engine (TheAxiomFoundation/axiom-rules-engine) " + "with its dense extension; install compatible source checkouts." + ) from error + return Frame, WeightKind, Weights, BE_SCHEMA, AxiomEngine + + +def _sha256_file(path: Path) -> str: + return sha256(path.read_bytes()).hexdigest() + + +def _rulespec_tree_sha256(root: Path) -> str: + """Hash every Belgian RuleSpec YAML path and byte payload deterministically.""" + + country_root = root / "be" + files = sorted( + path + for path in country_root.rglob("*") + if path.is_file() and path.suffix.lower() in {".yaml", ".yml"} + ) + if not files: + raise FileNotFoundError( + f"No Belgian RuleSpec YAML files found below {country_root}." + ) + digest = sha256() + for path in files: + relative = path.relative_to(root).as_posix().encode("utf-8") + 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 _module_sha256(module_name: str) -> str: + try: + spec = find_spec(module_name) + except (ImportError, ModuleNotFoundError) as error: + raise ImportError( + f"The Belgium pilot cannot locate runtime module {module_name!r}." + ) from error + if spec is None or spec.origin is None: + raise ImportError( + f"The Belgium pilot cannot locate runtime module {module_name!r}." + ) + path = Path(spec.origin) + if not path.is_file(): + raise ImportError( + f"Runtime module {module_name!r} has no hashable file at {path}." + ) + return _sha256_file(path) + + +def _package_tree_sha256(package_name: str) -> str: + """Hash every Python source path and payload in an import package.""" + try: + spec = find_spec(package_name) + except (ImportError, ModuleNotFoundError) as error: + raise ImportError( + f"The Belgium pilot cannot locate runtime package {package_name!r}." + ) from error + locations = None if spec is None else spec.submodule_search_locations + if not locations: + raise ImportError( + f"Runtime package {package_name!r} has no hashable source tree." + ) + + roots = sorted(Path(location).resolve() for location in locations) + files = [ + (root_index, root, path) + for root_index, root in enumerate(roots) + for path in sorted(root.rglob("*.py")) + if path.is_file() + ] + if not files: + raise ImportError( + f"Runtime package {package_name!r} has no Python source files." + ) + + digest = sha256() + for root_index, root, path in files: + relative = f"{root_index}/{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 _distribution_version(name: str) -> Optional[str]: + try: + return importlib_metadata.version(name) + except importlib_metadata.PackageNotFoundError: + return None + + +def _build_runtime_provenance(rulespec_root: Path) -> dict[str, Any]: + """Bind the exact RuleSpec corpus, adapter, wrapper, and native engine bytes.""" + + module = rulespec_root / PILOT_MODULE + if not module.is_file(): + raise FileNotFoundError( + f"rulespec-be pilot module not found at {module}; pass a " + "checkout of TheAxiomFoundation/rulespec-be as rulespec_root." + ) + return { + "rulespec": { + "repository": "TheAxiomFoundation/rulespec-be", + "module": PILOT_MODULE, + "module_sha256": _sha256_file(module), + "belgium_tree_sha256": _rulespec_tree_sha256(rulespec_root), + }, + "runtime": { + "microcosm_frame_version": _distribution_version("microcosm-frame"), + "microcosm_frame_tree_sha256": _package_tree_sha256("microcosm.frame"), + "microcosm_axiom_adapter_sha256": _module_sha256( + "microcosm.frame.adapters.axiom" + ), + "axiom_python_version": _distribution_version("axiom-rules-engine"), + "axiom_python_sha256": _module_sha256("axiom_rules_engine.dense"), + "axiom_dense_version": _distribution_version("axiom-rules-engine-dense"), + # The native binary embeds the Rust engine implementation. Its + # full-file digest remains exact even while the source-only + # packages reuse placeholder 0.1.0 distribution versions. + "axiom_dense_sha256": _module_sha256("axiom_rules_engine_dense"), + }, + } + + +def _provenance_version(provenance: dict[str, Any]) -> str: + rulespec_sha = provenance["rulespec"]["belgium_tree_sha256"] + canonical = json.dumps( + provenance, + allow_nan=False, + separators=(",", ":"), + sort_keys=True, + ).encode("utf-8") + runtime_sha = sha256(canonical).hexdigest() + return f"rulespec-be@{rulespec_sha[:12]}+runtime@{runtime_sha[:12]}" diff --git a/tests/test_be_axiom_pilot.py b/tests/test_be_axiom_pilot.py new file mode 100644 index 00000000..b75aec79 --- /dev/null +++ b/tests/test_be_axiom_pilot.py @@ -0,0 +1,613 @@ +"""Belgium pilot integration and source-stack contract tests. + +The PolicyEngine seam is deterministic and runs in ordinary CI. It uses a +strict runtime double only at the unpublished Microcosm/Axiom boundary; no +test double claims to validate Belgian law. Set ``RUN_BE_AXIOM_INTEGRATION=1`` +to additionally execute the real source-only stack once compatible checkouts +and the Axiom dense extension are installed. +""" + +import os +from copy import deepcopy +from hashlib import sha256 +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pandas as pd +import pytest +from microdf import MicroDataFrame + +from policyengine.core import Simulation +from policyengine.tax_benefit_models.be import ( + EMPLOYEE_SSC, + PIT_BEFORE_WITHHOLDING, + REMUNERATION, + AxiomBelgiumPilot, + BEYearData, + PopulaceBelgiumDataset, +) +from policyengine.tax_benefit_models.be import model as be_model_module + +RULESPEC_ROOT = Path( + os.environ.get("RULESPEC_BE_ROOT", "~/TheAxiomFoundation/rulespec-be") +).expanduser() +ORDINARY_WORKER_SSC_RATE = 0.1307 # arrete royal 28.11.1969, art. 19 +INCOMES = [0.0, 20_000.0, 30_000.0, 60_000.0] +PERSON_WEIGHTS = [2.0, 2.0, 5.0, 5.0] +HOUSEHOLD_WEIGHTS = [2.0, 5.0] +SOURCE_METADATA = { + "build_id": "microcosm-be-test-build", + "calibration": {"weight_kind": "calibrated", "target_count": 21}, +} +TEST_PROVENANCE = { + "rulespec": { + "repository": "TheAxiomFoundation/rulespec-be", + "module": be_model_module.PILOT_MODULE, + "module_sha256": "1" * 64, + "belgium_tree_sha256": "2" * 64, + }, + "runtime": { + "microcosm_frame_version": "0.1.0", + "microcosm_frame_tree_sha256": "3" * 64, + "microcosm_axiom_adapter_sha256": "4" * 64, + "axiom_python_version": "0.1.0", + "axiom_python_sha256": "5" * 64, + "axiom_dense_version": "0.1.0", + "axiom_dense_sha256": "6" * 64, + }, +} + + +class FakeWeightKind: + CALIBRATED = "calibrated" + + +class FakeWeights: + def __init__(self, *, values, kind): + self.values = np.asarray(values, dtype=float) + self.kind = kind + + +class FakeFrame: + latest = None + + def __init__(self, tables, schema, weights): + self.tables = {name: table.copy() for name, table in tables.items()} + self.schema = schema + self.weights = weights + type(self).latest = self + + +class FakeAxiomEngine: + latest = None + + def __init__(self, module, *, rulespec_roots): + self.module = Path(module) + self.rulespec_roots = tuple(Path(root) for root in rulespec_roots) + self.period = None + type(self).latest = self + + def materialize(self, frame, variables, period): + self.period = period + assert variables == [EMPLOYEE_SSC, PIT_BEFORE_WITHHOLDING] + gross = frame.tables["person"][REMUNERATION].to_numpy() + return { + EMPLOYEE_SSC: gross * 0.1, + PIT_BEFORE_WITHHOLDING: gross * 0.2, + } + + +@pytest.fixture +def stub_rulespec_root(tmp_path): + root = tmp_path / "rulespec-be" + module = root / be_model_module.PILOT_MODULE + module.parent.mkdir(parents=True) + module.write_text("format: rulespec/v1\nrules: []\n", encoding="utf-8") + return root + + +@pytest.fixture +def stub_source_runtime(monkeypatch, stub_rulespec_root): + FakeFrame.latest = None + FakeAxiomEngine.latest = None + monkeypatch.setattr( + be_model_module, + "_build_runtime_provenance", + lambda _root: deepcopy(TEST_PROVENANCE), + ) + monkeypatch.setattr( + be_model_module, + "_load_axiom_runtime", + lambda: ( + FakeFrame, + FakeWeightKind, + FakeWeights, + "current-microcosm-be-schema", + FakeAxiomEngine, + ), + ) + return stub_rulespec_root + + +@pytest.fixture +def pilot_dataset(tmp_path): + person = MicroDataFrame( + pd.DataFrame( + { + "person_id": [1, 2, 3, 4], + "person_household_id": [1, 1, 2, 2], + "age": [40.0, 38.0, 30.0, 52.0], + "is_male": [True, False, False, True], + REMUNERATION: INCOMES, + "person_weight": PERSON_WEIGHTS, + } + ), + weights="person_weight", + ) + household = MicroDataFrame( + pd.DataFrame( + { + "household_id": [1, 2], + "household_weight": HOUSEHOLD_WEIGHTS, + } + ), + weights="household_weight", + ) + path = tmp_path / "microcosm_be_test.h5" + PopulaceBelgiumDataset( + name="microcosm-be-test", + description="four-person nonuniform-weight fixture", + filepath=str(path), + year=2026, + metadata=deepcopy(SOURCE_METADATA), + data=BEYearData(person=person, household=household), + ).save() + # Exercise load(), including metadata and _time_period, during the run. + return PopulaceBelgiumDataset( + name="microcosm-be-test", + description="four-person nonuniform-weight fixture", + filepath=str(path), + year=2026, + ) + + +def _run( + pilot_dataset, + *, + rulespec_root, + period=2025, + communal_additional_tax_rate=0.0, +): + version = AxiomBelgiumPilot( + rulespec_root=str(rulespec_root), + period=period, + communal_additional_tax_rate=communal_additional_tax_rate, + ) + simulation = Simulation( + dataset=pilot_dataset, + tax_benefit_model_version=version, + ) + simulation.run() + return simulation + + +def test_run_preserves_input_periods_weights_and_metadata( + pilot_dataset, stub_source_runtime +): + input_path = Path(pilot_dataset.filepath) + before = sha256(input_path.read_bytes()).hexdigest() + + simulation = _run( + pilot_dataset, + rulespec_root=stub_source_runtime, + period=2025, + ) + output = simulation.output_dataset + + assert sha256(input_path.read_bytes()).hexdigest() == before + assert output.filepath is None + assert output.year == 2026 + assert output.policy_period == 2025 + assert FakeAxiomEngine.latest.period == 2025 + assert FakeAxiomEngine.latest.rulespec_roots == (stub_source_runtime.resolve(),) + assert pilot_dataset.metadata == SOURCE_METADATA + assert output.metadata["build_id"] == SOURCE_METADATA["build_id"] + assert output.metadata["calibration"] == SOURCE_METADATA["calibration"] + run_metadata = output.metadata["policyengine_axiom_runs"][-1] + assert run_metadata == { + "dataset_year": 2026, + "policy_period": 2025, + "model_version": simulation.tax_benefit_model_version.version, + "configuration": { + "communal_additional_tax_rate": 0.0, + "output_variables": [EMPLOYEE_SSC, PIT_BEFORE_WITHHOLDING], + }, + "provenance": TEST_PROVENANCE, + } + + frame = FakeFrame.latest + assert frame.schema == "current-microcosm-be-schema" + assert "person_weight" not in frame.tables["person"] + assert "household_weight" not in frame.tables["household"] + np.testing.assert_array_equal( + frame.weights["household"].values, + HOUSEHOLD_WEIGHTS, + ) + assert frame.weights["household"].kind == FakeWeightKind.CALIBRATED + + # MicroDataFrame keeps the nonuniform calibrated person weights on every + # output series: 0*2 + 2,000*2 + 3,000*5 + 6,000*5 = 49,000. + assert float(output.data.person[EMPLOYEE_SSC].sum()) == 49_000.0 + np.testing.assert_array_equal( + output.data.person["person_weight"], + PERSON_WEIGHTS, + ) + + # Both the model's no-op persistence hook and the derived dataset's + # missing destination leave the source bytes untouched. + simulation.save() + with pytest.raises(ValueError, match="without a filepath"): + output.save() + assert sha256(input_path.read_bytes()).hexdigest() == before + + +def test_default_policy_period_is_the_dataset_year(pilot_dataset, stub_source_runtime): + simulation = _run( + pilot_dataset, + rulespec_root=stub_source_runtime, + period=None, + ) + assert FakeAxiomEngine.latest.period == 2026 + assert simulation.output_dataset.policy_period == 2026 + + +def test_run_records_result_changing_configuration(pilot_dataset, stub_source_runtime): + simulation = _run( + pilot_dataset, + rulespec_root=stub_source_runtime, + period=2025, + communal_additional_tax_rate=0.075, + ) + run_metadata = simulation.output_dataset.metadata["policyengine_axiom_runs"][-1] + assert run_metadata["configuration"] == { + "communal_additional_tax_rate": 0.075, + "output_variables": [EMPLOYEE_SSC, PIT_BEFORE_WITHHOLDING], + } + np.testing.assert_array_equal( + FakeFrame.latest.tables["person"]["belgium_pit_communal_additional_tax_rate"], + [0.075] * 4, + ) + + +def test_output_metadata_round_trips_only_after_a_distinct_path_is_chosen( + tmp_path, pilot_dataset, stub_source_runtime +): + input_path = Path(pilot_dataset.filepath) + before = sha256(input_path.read_bytes()).hexdigest() + output = _run( + pilot_dataset, + rulespec_root=stub_source_runtime, + period=2025, + ).output_dataset + output_path = tmp_path / "belgium_output.h5" + assert output_path != input_path + + output.filepath = str(output_path) + output.save() + reloaded = PopulaceBelgiumDataset( + name=output.name, + description=output.description, + filepath=str(output_path), + year=output.year, + ) + reloaded.load() + + assert reloaded.policy_period == 2025 + assert reloaded.metadata == output.metadata + assert sha256(input_path.read_bytes()).hexdigest() == before + + +def test_dataset_rejects_a_mislabeled_hdf5_period(pilot_dataset): + mislabeled = PopulaceBelgiumDataset( + name=pilot_dataset.name, + description=pilot_dataset.description, + filepath=pilot_dataset.filepath, + year=2025, + ) + with pytest.raises(ValueError, match="period mismatch"): + mislabeled.load() + + +def test_model_version_is_derived_from_content_provenance( + monkeypatch, stub_source_runtime +): + first = AxiomBelgiumPilot( + rulespec_root=str(stub_source_runtime), + period=2025, + ) + assert first.version.startswith("rulespec-be@222222222222+runtime@") + assert first.version != "0.1.0-pilot" + assert first.id == f"{be_model_module.be_model.id}@{first.version}" + + attempted_override = AxiomBelgiumPilot( + rulespec_root=str(stub_source_runtime), + period=2025, + id="caller-supplied-id", + ) + assert attempted_override.id == f"{be_model_module.be_model.id}@{first.version}" + + for field in ( + "microcosm_frame_tree_sha256", + "microcosm_axiom_adapter_sha256", + "axiom_python_sha256", + "axiom_dense_sha256", + ): + changed = deepcopy(TEST_PROVENANCE) + changed["runtime"][field] = "a" * 64 + monkeypatch.setattr( + be_model_module, + "_build_runtime_provenance", + lambda _root, value=changed: value, + ) + changed_version = AxiomBelgiumPilot( + rulespec_root=str(stub_source_runtime), + period=2025, + ) + assert changed_version.version != first.version + + +def test_run_refuses_provenance_drift(monkeypatch, pilot_dataset, stub_source_runtime): + version = AxiomBelgiumPilot( + rulespec_root=str(stub_source_runtime), + period=2025, + ) + changed = deepcopy(TEST_PROVENANCE) + changed["rulespec"]["module_sha256"] = "f" * 64 + monkeypatch.setattr( + be_model_module, + "_build_runtime_provenance", + lambda _root: changed, + ) + simulation = Simulation( + dataset=pilot_dataset, + tax_benefit_model_version=version, + ) + with pytest.raises(RuntimeError, match="changed after"): + simulation.run() + + +def test_rulespec_tree_digest_binds_paths_and_bytes(tmp_path): + first = tmp_path / "be" / "first.yaml" + second = tmp_path / "be" / "nested" / "second.yml" + second.parent.mkdir(parents=True) + first.write_text("format: rulespec/v1\nrules: []\n", encoding="utf-8") + second.write_text("format: rulespec/v1\nrules: []\n", encoding="utf-8") + initial = be_model_module._rulespec_tree_sha256(tmp_path) + second.write_text("format: rulespec/v1\nrules: [changed]\n", encoding="utf-8") + assert be_model_module._rulespec_tree_sha256(tmp_path) != initial + + +def test_package_tree_digest_binds_execution_modules(monkeypatch, tmp_path): + package_root = tmp_path / "microcosm" / "frame" + package_root.mkdir(parents=True) + (package_root / "__init__.py").write_text("from .bundle import Frame\n") + bundle = package_root / "bundle.py" + bundle.write_text("class Frame: pass\n") + monkeypatch.setattr( + be_model_module, + "find_spec", + lambda _name: SimpleNamespace( + submodule_search_locations=[str(package_root)], + ), + ) + + initial = be_model_module._package_tree_sha256("microcosm.frame") + bundle.write_text("class Frame: changed = True\n") + assert be_model_module._package_tree_sha256("microcosm.frame") != initial + + +def test_current_microcosm_frame_accepts_the_pilot_weight_contract(): + microcosm_frame = pytest.importorskip( + "microcosm.frame", + reason="microcosm-frame is an unpublished source dependency", + ) + from microcosm.frame.adapters.axiom import BE_SCHEMA + + person = pd.DataFrame({"person_id": [1, 2], "person_household_id": [1, 2]}) + household = pd.DataFrame({"household_id": [1, 2]}) + weights = { + "household": microcosm_frame.Weights( + values=np.asarray(HOUSEHOLD_WEIGHTS), + kind=microcosm_frame.WeightKind.CALIBRATED, + ) + } + frame = microcosm_frame.Frame( + {"person": person, "household": household}, + BE_SCHEMA, + weights, + ) + np.testing.assert_array_equal( + frame.weights_for("household").values, + HOUSEHOLD_WEIGHTS, + ) + assert frame.weights_for("household").kind is microcosm_frame.WeightKind.CALIBRATED + + +def test_saved_dataset_matches_the_current_axiom_hdf5_layout(pilot_dataset): + pytest.importorskip( + "microcosm.frame", + reason="microcosm-frame is an unpublished source dependency", + ) + from microcosm.frame.adapters.axiom import AxiomEntityTableDataset + + current = AxiomEntityTableDataset(file_path=pilot_dataset.filepath) + assert current.time_period == 2026 + assert set(current.tables) == {"person", "household"} + assert "person_weight" not in current.person + assert current.household["household_weight"].tolist() == HOUSEHOLD_WEIGHTS + + pilot_dataset.load() + np.testing.assert_array_equal( + pilot_dataset.data.person["person_weight"], + PERSON_WEIGHTS, + ) + + +def test_dataset_rejects_a_mismatched_legacy_person_weight(tmp_path): + pytest.importorskip( + "microcosm.frame", + reason="microcosm-frame is an unpublished source dependency", + ) + from microcosm.frame.adapters.axiom import AxiomEntityTableDataset + + path = tmp_path / "legacy_mismatch.h5" + AxiomEntityTableDataset( + tables={ + "person": pd.DataFrame( + { + "person_id": [1, 2], + "person_household_id": [1, 2], + "person_weight": [2.0, 999.0], + } + ), + "household": pd.DataFrame( + { + "household_id": [1, 2], + "household_weight": [2.0, 5.0], + } + ), + }, + time_period=2026, + ).save(path) + + dataset = PopulaceBelgiumDataset( + name="legacy-mismatch", + description="invalid redundant person weights", + filepath=str(path), + year=2026, + ) + with pytest.raises(ValueError, match="person_weight values do not exactly match"): + dataset.load() + + +@pytest.mark.skipif( + os.environ.get("RUN_BE_AXIOM_INTEGRATION") != "1", + reason=( + "set RUN_BE_AXIOM_INTEGRATION=1 with compatible microcosm-frame, " + "axiom-rules-engine, dense extension, and rulespec-be source checkouts" + ), +) +def test_real_source_stack_computes_the_worker_slice(pilot_dataset): + simulation = _run( + pilot_dataset, + rulespec_root=RULESPEC_ROOT, + period=2025, + ) + person = pd.DataFrame(simulation.output_dataset.data.person) + gross = person[REMUNERATION].to_numpy() + ssc = person[EMPLOYEE_SSC].to_numpy() + pit = person[PIT_BEFORE_WITHHOLDING].to_numpy() + + statutory = gross * ORDINARY_WORKER_SSC_RATE + assert np.isfinite(ssc).all() + assert np.isfinite(pit).all() + assert ssc[0] == 0.0 + assert ssc[1] == 0.0 + assert 0.0 < ssc[2] < statutory[2] + np.testing.assert_allclose(ssc[3], statutory[3], rtol=1e-9) + assert pit[0] == 0.0 + # The pinned RuleSpec companion fixture establishes that this output is + # net of a refundable credit and can be negative at low remuneration. + assert pit[1] < 0.0 < pit[2] < pit[3] + + +@pytest.fixture +def in_memory_pilot_dataset(): + return PopulaceBelgiumDataset( + name="in-memory-pilot", + description="two-person household-authoritative weight fixture", + year=2026, + metadata=deepcopy(SOURCE_METADATA), + data=BEYearData( + person=MicroDataFrame( + { + "person_id": [1, 2], + "person_household_id": [1, 2], + REMUNERATION: [10.0, 20.0], + } + ), + household=MicroDataFrame( + { + # Reversed rows require an ID join, not positional copying. + "household_id": [2, 1], + "household_weight": [5.0, 2.0], + }, + weights="household_weight", + ), + ), + ) + + +def test_run_rejects_mismatched_in_memory_person_weights( + in_memory_pilot_dataset, stub_source_runtime +): + person = pd.DataFrame(in_memory_pilot_dataset.data.person).copy() + person["person_weight"] = [2.0, 999.0] + in_memory_pilot_dataset.data.person = MicroDataFrame( + person, + weights="person_weight", + ) + + with pytest.raises(ValueError, match="person_weight values do not exactly match"): + _run( + in_memory_pilot_dataset, + rulespec_root=stub_source_runtime, + ) + + assert FakeFrame.latest is None + assert FakeAxiomEngine.latest is None + np.testing.assert_array_equal( + in_memory_pilot_dataset.data.person["person_weight"], + [2.0, 999.0], + ) + + +@pytest.mark.parametrize("with_legacy_person_weights", [False, True]) +def test_run_derives_in_memory_weights_from_households( + in_memory_pilot_dataset, stub_source_runtime, with_legacy_person_weights +): + if with_legacy_person_weights: + person = pd.DataFrame(in_memory_pilot_dataset.data.person).copy() + person["person_weight"] = [2.0, 5.0] + in_memory_pilot_dataset.data.person = MicroDataFrame( + person, + weights="person_weight", + ) + original_person = pd.DataFrame(in_memory_pilot_dataset.data.person).copy() + original_household = pd.DataFrame(in_memory_pilot_dataset.data.household).copy() + + output = _run( + in_memory_pilot_dataset, + rulespec_root=stub_source_runtime, + ).output_dataset + + # The runtime double returns 1 and 2; household weights make 1*2 + 2*5 = 12. + assert float(output.data.person[EMPLOYEE_SSC].sum()) == 12.0 + np.testing.assert_array_equal(output.data.person["person_weight"], [2.0, 5.0]) + assert "person_weight" not in FakeFrame.latest.tables["person"] + np.testing.assert_array_equal( + FakeFrame.latest.weights["household"].values, + [5.0, 2.0], + ) + assert in_memory_pilot_dataset.filepath is None + assert output.filepath is None + assert in_memory_pilot_dataset.metadata == SOURCE_METADATA + pd.testing.assert_frame_equal( + pd.DataFrame(in_memory_pilot_dataset.data.person), + original_person, + ) + pd.testing.assert_frame_equal( + pd.DataFrame(in_memory_pilot_dataset.data.household), + original_household, + )