From 061e03c1d43c7785ef812ee9193d489801a68d09 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 3 Jul 2026 00:39:53 +0200 Subject: [PATCH 1/6] Add Belgium pilot: Axiom rules engine over populace-be entity tables The first non-policyengine-core country channel: an Axiom-backed TaxBenefitModelVersion outside the certified-release machinery, running the rulespec-be composed worker pipeline (employee SSC and PIT before withholding) over populace-be person/household tables with calibrated weights. Includes a populace-style two-entity dataset class, a population example scored against ONSS/SPF facts, and tests that skip cleanly when the source-only dependencies are absent. Fixes #447 Co-Authored-By: Claude Fable 5 --- changelog.d/447.added.md | 1 + examples/belgium_axiom_pilot.py | 59 +++++++ .../tax_benefit_models/be/__init__.py | 29 ++++ .../tax_benefit_models/be/datasets.py | 49 ++++++ .../tax_benefit_models/be/model.py | 158 ++++++++++++++++++ tests/test_be_axiom_pilot.py | 96 +++++++++++ 6 files changed, 392 insertions(+) create mode 100644 changelog.d/447.added.md create mode 100644 examples/belgium_axiom_pilot.py create mode 100644 src/policyengine/tax_benefit_models/be/__init__.py create mode 100644 src/policyengine/tax_benefit_models/be/datasets.py create mode 100644 src/policyengine/tax_benefit_models/be/model.py create mode 100644 tests/test_be_axiom_pilot.py diff --git a/changelog.d/447.added.md b/changelog.d/447.added.md new file mode 100644 index 00000000..a95863ab --- /dev/null +++ b/changelog.d/447.added.md @@ -0,0 +1 @@ +Belgium pilot: an Axiom-rules-engine-backed model version (tax_benefit_models/be) running the rulespec-be worker pipeline over populace-be entity tables, with a two-entity dataset class, example, and source-dependency-gated tests. diff --git a/examples/belgium_axiom_pilot.py b/examples/belgium_axiom_pilot.py new file mode 100644 index 00000000..918b32be --- /dev/null +++ b/examples/belgium_axiom_pilot.py @@ -0,0 +1,59 @@ +"""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 only worker SSC and individual +PIT are encoded so far. + +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() + +person = simulation.output_dataset.data.person +ssc = person[EMPLOYEE_SSC].sum() +pit = person[PIT_BEFORE_WITHHOLDING].sum() + +print("Belgium pilot (Axiom engine over populace-be, worker slice)") +print(f" employee SSC EUR {ssc / 1e9:6.2f}B (ONSS 2024: EUR 20.84B)") +print( + f" PIT before withholding EUR {pit / 1e9:6.2f}B (SPF 2023, all PIT: EUR 62.84B)" +) +print(f" SSC vs ONSS ratio {ssc / ONSS_WORKER_CONTRIBUTIONS_2024:.3f}") +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..d05692ae --- /dev/null +++ b/src/policyengine/tax_benefit_models/be/__init__.py @@ -0,0 +1,29 @@ +"""Belgium pilot: Axiom rules engine over populace 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 populace-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..0a28b8d8 --- /dev/null +++ b/src/policyengine/tax_benefit_models/be/datasets.py @@ -0,0 +1,49 @@ +"""Belgium pilot dataset: populace entity tables with calibrated weights. + +The pilot layout has two entities (person, household), mirroring the +populace ``BE_SCHEMA``. Files are plain pandas HDF5 stores with ``person`` +and ``household`` keys; weights live in ``person_weight`` / +``household_weight`` columns, as in the US/UK single-year layouts. +""" + +from typing import Any, Optional + +import pandas as pd +from microdf import MicroDataFrame +from pydantic import ConfigDict, Field + +from policyengine.core.dataset import Dataset, YearData + + +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 populace-be HDF5 artifact.""" + + data: Optional[BEYearData] = None + metadata: dict[str, Any] = Field(default_factory=dict) + + def load(self) -> None: + person = pd.read_hdf(self.filepath, key="person") + household = pd.read_hdf(self.filepath, key="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.") + pd.DataFrame(self.data.person).to_hdf(self.filepath, key="person", mode="w") + pd.DataFrame(self.data.household).to_hdf(self.filepath, key="household") 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..036a5472 --- /dev/null +++ b/src/policyengine/tax_benefit_models/be/model.py @@ -0,0 +1,158 @@ +"""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 ``populace-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): + +- ``populace-frame`` from PolicyEngine/populace (``packages/populace-frame``) +- ``axiom-rules-engine`` from TheAxiomFoundation/axiom-rules-engine (PyO3 + dense extension) +- a checkout of TheAxiomFoundation/rulespec-be, passed as ``rulespec_root`` + +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). +""" + +from pathlib import Path +from typing import TYPE_CHECKING, ClassVar, Optional, Union + +import pandas as pd +from microdf import MicroDataFrame + +from policyengine.core import TaxBenefitModel +from policyengine.core.tax_benefit_model_version import TaxBenefitModelVersion + +from .datasets import BEYearData, PopulaceBelgiumDataset + +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). +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_pit_pilot_article_289ter1_work_bonus_a_amount": 0.0, + "belgium_pit_pilot_article_289ter1_work_bonus_b_amount": 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 + + def __init__(self, **kwargs) -> None: + kwargs.setdefault("model", be_model) + kwargs.setdefault("version", "0.1.0-pilot") + super().__init__(**kwargs) + + def run(self, simulation: "Simulation") -> "Simulation": + try: + from populace.frame import Frame, WeightKind, Weights + from populace.frame.adapters.axiom import BE_SCHEMA, AxiomEngine + except ImportError as error: + raise ImportError( + "The Belgium pilot needs populace-frame (PolicyEngine/populace, " + "packages/populace-frame) and axiom-rules-engine " + "(TheAxiomFoundation/axiom-rules-engine); neither is on PyPI " + "yet, install both from source." + ) from error + + module = Path(self.rulespec_root).expanduser() / 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 + + person = pd.DataFrame(dataset.data.person).copy() + household = pd.DataFrame(dataset.data.household).copy() + 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, + ) + engine = AxiomEngine(str(module)) + outputs = engine.materialize( + frame, self.output_variables, self.period or dataset.year + ) + for name, values in outputs.items(): + person[name] = values + + simulation.output_dataset = PopulaceBelgiumDataset( + id=simulation.id, + name=dataset.name, + description=dataset.description, + filepath=dataset.filepath, + year=dataset.year, + is_output_dataset=True, + 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.") diff --git a/tests/test_be_axiom_pilot.py b/tests/test_be_axiom_pilot.py new file mode 100644 index 00000000..364bff98 --- /dev/null +++ b/tests/test_be_axiom_pilot.py @@ -0,0 +1,96 @@ +"""Belgium pilot: Axiom engine over a tiny populace-style dataset. + +Skips cleanly unless the source-only dependencies (populace-frame, +axiom-rules-engine) are importable and a rulespec-be checkout is available +via ``RULESPEC_BE_ROOT`` (or the default sibling path). +""" + +import os +from importlib.util import find_spec +from pathlib import Path + +import numpy as np +import pandas as pd +import pytest + +RULESPEC_ROOT = Path( + os.environ.get("RULESPEC_BE_ROOT", "~/TheAxiomFoundation/rulespec-be") +).expanduser() +PILOT_MODULE = ( + RULESPEC_ROOT + / "be/statutes/income_tax/individual/pilot_worker_oracle_pipeline.yaml" +) + +requires_axiom = pytest.mark.skipif( + find_spec("populace") is None + or find_spec("axiom_rules_engine") is None + or not PILOT_MODULE.exists(), + reason="needs populace-frame, axiom-rules-engine, and a rulespec-be checkout", +) + +ORDINARY_WORKER_SSC_RATE = 0.1307 # arrete royal 28.11.1969, art. 19 + + +@pytest.fixture +def pilot_dataset(tmp_path): + from policyengine.tax_benefit_models.be import PopulaceBelgiumDataset + + person = pd.DataFrame( + { + "person_id": [1, 2, 3], + "person_household_id": [1, 1, 2], + "age": [40.0, 38.0, 30.0], + "is_male": [True, False, False], + "belgium_pit_article_23_worker_remuneration": [0.0, 30_000.0, 60_000.0], + "person_weight": [1.0, 1.0, 1.0], + } + ) + household = pd.DataFrame({"household_id": [1, 2], "household_weight": [1.0, 1.0]}) + path = tmp_path / "populace_be_test.h5" + person.to_hdf(path, key="person", mode="w") + household.to_hdf(path, key="household") + return PopulaceBelgiumDataset( + name="populace-be-test", + description="three-person fixture", + filepath=str(path), + year=2026, + ) + + +@requires_axiom +def test_pilot_run_computes_statutory_ssc_and_progressive_pit(pilot_dataset): + from policyengine.core.simulation import Simulation + from policyengine.tax_benefit_models.be import ( + EMPLOYEE_SSC, + PIT_BEFORE_WITHHOLDING, + REMUNERATION, + AxiomBelgiumPilot, + ) + + version = AxiomBelgiumPilot(rulespec_root=str(RULESPEC_ROOT), period=2025) + simulation = Simulation(dataset=pilot_dataset, tax_benefit_model_version=version) + simulation.run() + + 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() + + np.testing.assert_allclose(ssc, gross * ORDINARY_WORKER_SSC_RATE, rtol=1e-9) + assert pit[0] == 0.0 + assert 0.0 < pit[1] < pit[2] + assert simulation.output_dataset.is_output_dataset + + +@requires_axiom +def test_pilot_weighted_aggregates_use_calibrated_weights(pilot_dataset): + from policyengine.core.simulation import Simulation + from policyengine.tax_benefit_models.be import EMPLOYEE_SSC, AxiomBelgiumPilot + + version = AxiomBelgiumPilot(rulespec_root=str(RULESPEC_ROOT), period=2025) + simulation = Simulation(dataset=pilot_dataset, tax_benefit_model_version=version) + simulation.run() + + person = simulation.output_dataset.data.person + expected = 90_000.0 * ORDINARY_WORKER_SSC_RATE + assert float(person[EMPLOYEE_SSC].sum()) == pytest.approx(expected, rel=1e-9) From 7cbf45cd0c4fce483c4af0d6c0ca9c7bf3faf7b5 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 3 Jul 2026 07:47:17 +0200 Subject: [PATCH 2/6] Track rulespec-be work-bonus encoding in the Belgium pilot The composed pipeline now computes the ONSS work-bonus amounts from law (the a/b supplied inputs are gone; the reference-wage input bridges 0 to worker remuneration), and the ordinary-contribution concept nets the bonus. Update the supplied defaults, exercise the phase-out in the test fixture (wiped at 20k, partial at 30k, statutory 13.07 percent by 60k), and label the example aggregates accordingly. Co-Authored-By: Claude Fable 5 --- examples/belgium_axiom_pilot.py | 18 +++--- .../tax_benefit_models/be/model.py | 5 +- tests/test_be_axiom_pilot.py | 63 ++++++++++++------- 3 files changed, 53 insertions(+), 33 deletions(-) diff --git a/examples/belgium_axiom_pilot.py b/examples/belgium_axiom_pilot.py index 918b32be..4d26bd48 100644 --- a/examples/belgium_axiom_pilot.py +++ b/examples/belgium_axiom_pilot.py @@ -6,8 +6,9 @@ against Belgian administrative facts. This is a demonstration of the engine channel, not a certified Belgian -model: the support records are American, and only worker SSC and individual -PIT are encoded so far. +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). Usage:: @@ -45,15 +46,18 @@ simulation.run() person = simulation.output_dataset.data.person -ssc = person[EMPLOYEE_SSC].sum() +ssc_net = person[EMPLOYEE_SSC].sum() pit = person[PIT_BEFORE_WITHHOLDING].sum() print("Belgium pilot (Axiom engine over populace-be, worker slice)") -print(f" employee SSC EUR {ssc / 1e9:6.2f}B (ONSS 2024: EUR 20.84B)") +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)" + f" PIT before withholding EUR {pit / 1e9:6.2f}B " + "(SPF 2023, all PIT: EUR 62.84B)" ) -print(f" SSC vs ONSS ratio {ssc / ONSS_WORKER_CONTRIBUTIONS_2024:.3f}") print( - f" PIT vs SPF ratio {pit / SPF_PIT_BEFORE_WITHHOLDING_2023:.3f} (worker slice only)" + f" PIT vs SPF ratio {pit / SPF_PIT_BEFORE_WITHHOLDING_2023:.3f}" + " (worker slice only)" ) diff --git a/src/policyengine/tax_benefit_models/be/model.py b/src/policyengine/tax_benefit_models/be/model.py index 036a5472..a19f4ba7 100644 --- a/src/policyengine/tax_benefit_models/be/model.py +++ b/src/policyengine/tax_benefit_models/be/model.py @@ -43,12 +43,13 @@ #: 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_pit_pilot_article_289ter1_work_bonus_a_amount": 0.0, - "belgium_pit_pilot_article_289ter1_work_bonus_b_amount": 0.0, + "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, } diff --git a/tests/test_be_axiom_pilot.py b/tests/test_be_axiom_pilot.py index 364bff98..d95ba26f 100644 --- a/tests/test_be_axiom_pilot.py +++ b/tests/test_be_axiom_pilot.py @@ -29,6 +29,10 @@ ) ORDINARY_WORKER_SSC_RATE = 0.1307 # arrete royal 28.11.1969, art. 19 +# incomes chosen around the 2025 work-bonus phase-out: the low-wage bonus +# wipes the employee contribution at 20k, partially reduces it at 30k, and +# is exhausted well before 60k (ONSS DMFA 2025 tables, as encoded). +INCOMES = [0.0, 20_000.0, 30_000.0, 60_000.0] @pytest.fixture @@ -37,12 +41,12 @@ def pilot_dataset(tmp_path): person = pd.DataFrame( { - "person_id": [1, 2, 3], - "person_household_id": [1, 1, 2], - "age": [40.0, 38.0, 30.0], - "is_male": [True, False, False], - "belgium_pit_article_23_worker_remuneration": [0.0, 30_000.0, 60_000.0], - "person_weight": [1.0, 1.0, 1.0], + "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], + "belgium_pit_article_23_worker_remuneration": INCOMES, + "person_weight": [1.0, 1.0, 1.0, 1.0], } ) household = pd.DataFrame({"household_id": [1, 2], "household_weight": [1.0, 1.0]}) @@ -51,46 +55,57 @@ def pilot_dataset(tmp_path): household.to_hdf(path, key="household") return PopulaceBelgiumDataset( name="populace-be-test", - description="three-person fixture", + description="four-person fixture around the work-bonus phase-out", filepath=str(path), year=2026, ) -@requires_axiom -def test_pilot_run_computes_statutory_ssc_and_progressive_pit(pilot_dataset): +def _run(pilot_dataset): from policyengine.core.simulation import Simulation + from policyengine.tax_benefit_models.be import AxiomBelgiumPilot + + version = AxiomBelgiumPilot(rulespec_root=str(RULESPEC_ROOT), period=2025) + simulation = Simulation(dataset=pilot_dataset, tax_benefit_model_version=version) + simulation.run() + return simulation + + +@requires_axiom +def test_pilot_run_computes_ssc_with_work_bonus_and_progressive_pit(pilot_dataset): from policyengine.tax_benefit_models.be import ( EMPLOYEE_SSC, PIT_BEFORE_WITHHOLDING, REMUNERATION, - AxiomBelgiumPilot, ) - version = AxiomBelgiumPilot(rulespec_root=str(RULESPEC_ROOT), period=2025) - simulation = Simulation(dataset=pilot_dataset, tax_benefit_model_version=version) - simulation.run() - + simulation = _run(pilot_dataset) 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() - np.testing.assert_allclose(ssc, gross * ORDINARY_WORKER_SSC_RATE, rtol=1e-9) + statutory = gross * ORDINARY_WORKER_SSC_RATE + assert ssc[0] == 0.0 + assert ssc[1] == 0.0 # work bonus wipes the contribution at 20k + assert 0.0 < ssc[2] < statutory[2] # partial bonus at 30k + np.testing.assert_allclose(ssc[3], statutory[3], rtol=1e-9) # exhausted + assert pit[0] == 0.0 - assert 0.0 < pit[1] < pit[2] + assert 0.0 <= pit[1] <= pit[2] < pit[3] assert simulation.output_dataset.is_output_dataset @requires_axiom def test_pilot_weighted_aggregates_use_calibrated_weights(pilot_dataset): - from policyengine.core.simulation import Simulation - from policyengine.tax_benefit_models.be import EMPLOYEE_SSC, AxiomBelgiumPilot - - version = AxiomBelgiumPilot(rulespec_root=str(RULESPEC_ROOT), period=2025) - simulation = Simulation(dataset=pilot_dataset, tax_benefit_model_version=version) - simulation.run() + from policyengine.tax_benefit_models.be import EMPLOYEE_SSC + simulation = _run(pilot_dataset) person = simulation.output_dataset.data.person - expected = 90_000.0 * ORDINARY_WORKER_SSC_RATE - assert float(person[EMPLOYEE_SSC].sum()) == pytest.approx(expected, rel=1e-9) + total = float(person[EMPLOYEE_SSC].sum()) + # 0 + 0 (bonus-wiped) + partial at 30k + full 13.07% at 60k + assert ( + 60_000.0 * ORDINARY_WORKER_SSC_RATE + < total + < 90_000.0 * ORDINARY_WORKER_SSC_RATE + ) From aea7d435f226cd191f6954e2fb92e3b768c4b538 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 18 Aug 2026 21:44:36 -0400 Subject: [PATCH 3/6] Refresh Belgium pilot for current core APIs --- src/policyengine/tax_benefit_models/be/datasets.py | 6 +++++- src/policyengine/tax_benefit_models/be/model.py | 7 +++---- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/src/policyengine/tax_benefit_models/be/datasets.py b/src/policyengine/tax_benefit_models/be/datasets.py index 0a28b8d8..ca5d72c5 100644 --- a/src/policyengine/tax_benefit_models/be/datasets.py +++ b/src/policyengine/tax_benefit_models/be/datasets.py @@ -12,7 +12,7 @@ from microdf import MicroDataFrame from pydantic import ConfigDict, Field -from policyengine.core.dataset import Dataset, YearData +from policyengine.core import Dataset, YearData class BEYearData(YearData): @@ -35,6 +35,8 @@ class PopulaceBelgiumDataset(Dataset): metadata: dict[str, Any] = Field(default_factory=dict) def load(self) -> None: + if self.filepath is None: + raise ValueError("Cannot load a Belgium pilot dataset without a filepath.") person = pd.read_hdf(self.filepath, key="person") household = pd.read_hdf(self.filepath, key="household") self.data = BEYearData( @@ -45,5 +47,7 @@ def load(self) -> None: 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.") pd.DataFrame(self.data.person).to_hdf(self.filepath, key="person", mode="w") pd.DataFrame(self.data.household).to_hdf(self.filepath, key="household") diff --git a/src/policyengine/tax_benefit_models/be/model.py b/src/policyengine/tax_benefit_models/be/model.py index a19f4ba7..64ce1b23 100644 --- a/src/policyengine/tax_benefit_models/be/model.py +++ b/src/policyengine/tax_benefit_models/be/model.py @@ -23,13 +23,12 @@ """ from pathlib import Path -from typing import TYPE_CHECKING, ClassVar, Optional, Union +from typing import TYPE_CHECKING, Any, ClassVar, Optional, Union import pandas as pd from microdf import MicroDataFrame -from policyengine.core import TaxBenefitModel -from policyengine.core.tax_benefit_model_version import TaxBenefitModelVersion +from policyengine.core import TaxBenefitModel, TaxBenefitModelVersion from .datasets import BEYearData, PopulaceBelgiumDataset @@ -76,7 +75,7 @@ class AxiomBelgiumPilot(TaxBenefitModelVersion): output_variables: list[str] = [EMPLOYEE_SSC, PIT_BEFORE_WITHHOLDING] communal_additional_tax_rate: float = 0.0 - def __init__(self, **kwargs) -> None: + def __init__(self, **kwargs: Any) -> None: kwargs.setdefault("model", be_model) kwargs.setdefault("version", "0.1.0-pilot") super().__init__(**kwargs) From 78896567da56007b78eee93ee41acb00e06423e3 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sat, 29 Aug 2026 18:38:21 -0400 Subject: [PATCH 4/6] Harden Belgium Axiom pilot integration --- changelog.d/447.added.md | 2 +- examples/belgium_axiom_pilot.py | 14 +- .../tax_benefit_models/be/__init__.py | 4 +- .../tax_benefit_models/be/datasets.py | 157 +++++- .../tax_benefit_models/be/model.py | 242 ++++++++- tests/test_be_axiom_pilot.py | 508 +++++++++++++++--- 6 files changed, 821 insertions(+), 106 deletions(-) diff --git a/changelog.d/447.added.md b/changelog.d/447.added.md index a95863ab..e63ff8e0 100644 --- a/changelog.d/447.added.md +++ b/changelog.d/447.added.md @@ -1 +1 @@ -Belgium pilot: an Axiom-rules-engine-backed model version (tax_benefit_models/be) running the rulespec-be worker pipeline over populace-be entity tables, with a two-entity dataset class, example, and source-dependency-gated tests. +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 index 4d26bd48..f13941dc 100644 --- a/examples/belgium_axiom_pilot.py +++ b/examples/belgium_axiom_pilot.py @@ -10,6 +10,13 @@ (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 \\ @@ -45,11 +52,16 @@ simulation = Simulation(dataset=dataset, tax_benefit_model_version=model_version) simulation.run() -person = simulation.output_dataset.data.person +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)") diff --git a/src/policyengine/tax_benefit_models/be/__init__.py b/src/policyengine/tax_benefit_models/be/__init__.py index d05692ae..74044832 100644 --- a/src/policyengine/tax_benefit_models/be/__init__.py +++ b/src/policyengine/tax_benefit_models/be/__init__.py @@ -1,8 +1,8 @@ -"""Belgium pilot: Axiom rules engine over populace entity tables. +"""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 populace-be entity tables. +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. """ diff --git a/src/policyengine/tax_benefit_models/be/datasets.py b/src/policyengine/tax_benefit_models/be/datasets.py index ca5d72c5..04bb40d0 100644 --- a/src/policyengine/tax_benefit_models/be/datasets.py +++ b/src/policyengine/tax_benefit_models/be/datasets.py @@ -1,13 +1,22 @@ -"""Belgium pilot dataset: populace entity tables with calibrated weights. +"""Belgium pilot dataset: Microcosm entity tables with calibrated weights. The pilot layout has two entities (person, household), mirroring the -populace ``BE_SCHEMA``. Files are plain pandas HDF5 stores with ``person`` -and ``household`` keys; weights live in ``person_weight`` / -``household_weight`` columns, as in the US/UK single-year layouts. +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. """ -from typing import Any, Optional +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 @@ -15,6 +24,57 @@ 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.""" @@ -29,16 +89,62 @@ def entity_data(self) -> dict[str, MicroDataFrame]: class PopulaceBelgiumDataset(Dataset): - """Belgium pilot dataset loaded from a populace-be HDF5 artifact.""" + """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.") - person = pd.read_hdf(self.filepath, key="person") - household = pd.read_hdf(self.filepath, key="household") + 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"), @@ -49,5 +155,36 @@ def save(self) -> None: raise ValueError("No data to save.") if self.filepath is None: raise ValueError("Cannot save a Belgium pilot dataset without a filepath.") - pd.DataFrame(self.data.person).to_hdf(self.filepath, key="person", mode="w") - pd.DataFrame(self.data.household).to_hdf(self.filepath, key="household") + # 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 index 64ce1b23..5317fcc6 100644 --- a/src/policyengine/tax_benefit_models/be/model.py +++ b/src/policyengine/tax_benefit_models/be/model.py @@ -3,18 +3,27 @@ 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 ``populace-frame`` Axiom adapter. There is no +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): -- ``populace-frame`` from PolicyEngine/populace (``packages/populace-frame``) +- ``microcosm-frame`` from PolicyEngine/microcosm + (``packages/microcosm-frame``; Python 3.13+) - ``axiom-rules-engine`` from TheAxiomFoundation/axiom-rules-engine (PyO3 - dense extension) + 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. @@ -22,6 +31,11 @@ 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 @@ -74,25 +88,33 @@ class AxiomBelgiumPilot(TaxBenefitModelVersion): 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: - kwargs.setdefault("model", be_model) - kwargs.setdefault("version", "0.1.0-pilot") + 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": - try: - from populace.frame import Frame, WeightKind, Weights - from populace.frame.adapters.axiom import BE_SCHEMA, AxiomEngine - except ImportError as error: - raise ImportError( - "The Belgium pilot needs populace-frame (PolicyEngine/populace, " - "packages/populace-frame) and axiom-rules-engine " - "(TheAxiomFoundation/axiom-rules-engine); neither is on PyPI " - "yet, install both from source." - ) from error - - module = Path(self.rulespec_root).expanduser() / PILOT_MODULE + 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 " @@ -130,20 +152,44 @@ def run(self, simulation: "Simulation") -> "Simulation": BE_SCHEMA, weights, ) - engine = AxiomEngine(str(module)) - outputs = engine.materialize( - frame, self.output_variables, self.period or dataset.year - ) + 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, - filepath=dataset.filepath, + # 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"), @@ -156,3 +202,155 @@ def save(self, simulation: "Simulation") -> None: 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 index d95ba26f..ffd572ae 100644 --- a/tests/test_be_axiom_pilot.py +++ b/tests/test_be_axiom_pilot.py @@ -1,85 +1,470 @@ -"""Belgium pilot: Axiom engine over a tiny populace-style dataset. +"""Belgium pilot integration and source-stack contract tests. -Skips cleanly unless the source-only dependencies (populace-frame, -axiom-rules-engine) are importable and a rulespec-be checkout is available -via ``RULESPEC_BE_ROOT`` (or the default sibling path). +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 importlib.util import find_spec +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() -PILOT_MODULE = ( - RULESPEC_ROOT - / "be/statutes/income_tax/individual/pilot_worker_oracle_pipeline.yaml" -) - -requires_axiom = pytest.mark.skipif( - find_spec("populace") is None - or find_spec("axiom_rules_engine") is None - or not PILOT_MODULE.exists(), - reason="needs populace-frame, axiom-rules-engine, and a rulespec-be checkout", -) - ORDINARY_WORKER_SSC_RATE = 0.1307 # arrete royal 28.11.1969, art. 19 -# incomes chosen around the 2025 work-bonus phase-out: the low-wage bonus -# wipes the employee contribution at 20k, partially reduces it at 30k, and -# is exhausted well before 60k (ONSS DMFA 2025 tables, as encoded). 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_source_runtime(monkeypatch): + 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, + ), + ) @pytest.fixture def pilot_dataset(tmp_path): - from policyengine.tax_benefit_models.be import PopulaceBelgiumDataset - - person = 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], - "belgium_pit_article_23_worker_remuneration": INCOMES, - "person_weight": [1.0, 1.0, 1.0, 1.0], - } + 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 = pd.DataFrame({"household_id": [1, 2], "household_weight": [1.0, 1.0]}) - path = tmp_path / "populace_be_test.h5" - person.to_hdf(path, key="person", mode="w") - household.to_hdf(path, key="household") + 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="populace-be-test", - description="four-person fixture around the work-bonus phase-out", + name="microcosm-be-test", + description="four-person nonuniform-weight fixture", filepath=str(path), year=2026, ) -def _run(pilot_dataset): - from policyengine.core.simulation import Simulation - from policyengine.tax_benefit_models.be import AxiomBelgiumPilot - - version = AxiomBelgiumPilot(rulespec_root=str(RULESPEC_ROOT), period=2025) - simulation = Simulation(dataset=pilot_dataset, tax_benefit_model_version=version) +def _run(pilot_dataset, *, 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 -@requires_axiom -def test_pilot_run_computes_ssc_with_work_bonus_and_progressive_pit(pilot_dataset): - from policyengine.tax_benefit_models.be import ( - EMPLOYEE_SSC, - PIT_BEFORE_WITHHOLDING, - REMUNERATION, +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, 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 == (RULESPEC_ROOT.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, 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, + 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, 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() - simulation = _run(pilot_dataset) + 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(RULESPEC_ROOT), 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(RULESPEC_ROOT), + 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(RULESPEC_ROOT), + 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(RULESPEC_ROOT), 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, period=2025) person = pd.DataFrame(simulation.output_dataset.data.person) gross = person[REMUNERATION].to_numpy() ssc = person[EMPLOYEE_SSC].to_numpy() @@ -87,25 +472,8 @@ def test_pilot_run_computes_ssc_with_work_bonus_and_progressive_pit(pilot_datase statutory = gross * ORDINARY_WORKER_SSC_RATE assert ssc[0] == 0.0 - assert ssc[1] == 0.0 # work bonus wipes the contribution at 20k - assert 0.0 < ssc[2] < statutory[2] # partial bonus at 30k - np.testing.assert_allclose(ssc[3], statutory[3], rtol=1e-9) # exhausted - + 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 assert 0.0 <= pit[1] <= pit[2] < pit[3] - assert simulation.output_dataset.is_output_dataset - - -@requires_axiom -def test_pilot_weighted_aggregates_use_calibrated_weights(pilot_dataset): - from policyengine.tax_benefit_models.be import EMPLOYEE_SSC - - simulation = _run(pilot_dataset) - person = simulation.output_dataset.data.person - total = float(person[EMPLOYEE_SSC].sum()) - # 0 + 0 (bonus-wiped) + partial at 30k + full 13.07% at 60k - assert ( - 60_000.0 * ORDINARY_WORKER_SSC_RATE - < total - < 90_000.0 * ORDINARY_WORKER_SSC_RATE - ) From 521a66b60b1935b5cefb93c790c2a39e7ec7dc95 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sat, 29 Aug 2026 21:33:40 -0400 Subject: [PATCH 5/6] Fix issues from review: validate Belgian in-memory weights Keep the shared household-weight validator authoritative on the populated-data path. Refs #447. --- .../tax_benefit_models/be/model.py | 11 ++- tests/test_be_axiom_pilot.py | 85 +++++++++++++++++++ 2 files changed, 94 insertions(+), 2 deletions(-) diff --git a/src/policyengine/tax_benefit_models/be/model.py b/src/policyengine/tax_benefit_models/be/model.py index 5317fcc6..c4fe9e85 100644 --- a/src/policyengine/tax_benefit_models/be/model.py +++ b/src/policyengine/tax_benefit_models/be/model.py @@ -44,7 +44,11 @@ from policyengine.core import TaxBenefitModel, TaxBenefitModelVersion -from .datasets import BEYearData, PopulaceBelgiumDataset +from .datasets import ( + BEYearData, + PopulaceBelgiumDataset, + _person_with_household_weights, +) if TYPE_CHECKING: from policyengine.core.simulation import Simulation @@ -127,8 +131,11 @@ def run(self, simulation: "Simulation") -> "Simulation": dataset.load() assert dataset.data is not None - person = pd.DataFrame(dataset.data.person).copy() 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 diff --git a/tests/test_be_axiom_pilot.py b/tests/test_be_axiom_pilot.py index ffd572ae..4d1d03a4 100644 --- a/tests/test_be_axiom_pilot.py +++ b/tests/test_be_axiom_pilot.py @@ -477,3 +477,88 @@ def test_real_source_stack_computes_the_worker_slice(pilot_dataset): np.testing.assert_allclose(ssc[3], statutory[3], rtol=1e-9) assert pit[0] == 0.0 assert 0.0 <= pit[1] <= 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) + + 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).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, + ) From 910bd210e41267e02b6bc611727b04c1d72faf15 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 30 Aug 2026 07:10:03 -0400 Subject: [PATCH 6/6] Fix issues from review: harden Belgian pilot tests --- tests/test_be_axiom_pilot.py | 79 +++++++++++++++++++++++++++++------- 1 file changed, 64 insertions(+), 15 deletions(-) diff --git a/tests/test_be_axiom_pilot.py b/tests/test_be_axiom_pilot.py index 4d1d03a4..b75aec79 100644 --- a/tests/test_be_axiom_pilot.py +++ b/tests/test_be_axiom_pilot.py @@ -99,7 +99,16 @@ def materialize(self, frame, variables, period): @pytest.fixture -def stub_source_runtime(monkeypatch): +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( @@ -118,6 +127,7 @@ def stub_source_runtime(monkeypatch): FakeAxiomEngine, ), ) + return stub_rulespec_root @pytest.fixture @@ -162,9 +172,15 @@ def pilot_dataset(tmp_path): ) -def _run(pilot_dataset, *, period=2025, communal_additional_tax_rate=0.0): +def _run( + pilot_dataset, + *, + rulespec_root, + period=2025, + communal_additional_tax_rate=0.0, +): version = AxiomBelgiumPilot( - rulespec_root=str(RULESPEC_ROOT), + rulespec_root=str(rulespec_root), period=period, communal_additional_tax_rate=communal_additional_tax_rate, ) @@ -182,7 +198,11 @@ def test_run_preserves_input_periods_weights_and_metadata( input_path = Path(pilot_dataset.filepath) before = sha256(input_path.read_bytes()).hexdigest() - simulation = _run(pilot_dataset, period=2025) + simulation = _run( + pilot_dataset, + rulespec_root=stub_source_runtime, + period=2025, + ) output = simulation.output_dataset assert sha256(input_path.read_bytes()).hexdigest() == before @@ -190,7 +210,7 @@ def test_run_preserves_input_periods_weights_and_metadata( assert output.year == 2026 assert output.policy_period == 2025 assert FakeAxiomEngine.latest.period == 2025 - assert FakeAxiomEngine.latest.rulespec_roots == (RULESPEC_ROOT.resolve(),) + 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"] @@ -233,7 +253,11 @@ def test_run_preserves_input_periods_weights_and_metadata( def test_default_policy_period_is_the_dataset_year(pilot_dataset, stub_source_runtime): - simulation = _run(pilot_dataset, period=None) + simulation = _run( + pilot_dataset, + rulespec_root=stub_source_runtime, + period=None, + ) assert FakeAxiomEngine.latest.period == 2026 assert simulation.output_dataset.policy_period == 2026 @@ -241,6 +265,7 @@ def test_default_policy_period_is_the_dataset_year(pilot_dataset, stub_source_ru 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, ) @@ -260,7 +285,11 @@ def test_output_metadata_round_trips_only_after_a_distinct_path_is_chosen( ): input_path = Path(pilot_dataset.filepath) before = sha256(input_path.read_bytes()).hexdigest() - output = _run(pilot_dataset, period=2025).output_dataset + 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 @@ -293,13 +322,16 @@ def test_dataset_rejects_a_mislabeled_hdf5_period(pilot_dataset): def test_model_version_is_derived_from_content_provenance( monkeypatch, stub_source_runtime ): - first = AxiomBelgiumPilot(rulespec_root=str(RULESPEC_ROOT), period=2025) + 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(RULESPEC_ROOT), + rulespec_root=str(stub_source_runtime), period=2025, id="caller-supplied-id", ) @@ -319,14 +351,17 @@ def test_model_version_is_derived_from_content_provenance( lambda _root, value=changed: value, ) changed_version = AxiomBelgiumPilot( - rulespec_root=str(RULESPEC_ROOT), + 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(RULESPEC_ROOT), period=2025) + version = AxiomBelgiumPilot( + rulespec_root=str(stub_source_runtime), + period=2025, + ) changed = deepcopy(TEST_PROVENANCE) changed["rulespec"]["module_sha256"] = "f" * 64 monkeypatch.setattr( @@ -464,19 +499,27 @@ def test_dataset_rejects_a_mismatched_legacy_person_weight(tmp_path): ), ) def test_real_source_stack_computes_the_worker_slice(pilot_dataset): - simulation = _run(pilot_dataset, period=2025) + 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 - assert 0.0 <= pit[1] <= pit[2] < pit[3] + # 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 @@ -517,7 +560,10 @@ def test_run_rejects_mismatched_in_memory_person_weights( ) with pytest.raises(ValueError, match="person_weight values do not exactly match"): - _run(in_memory_pilot_dataset) + _run( + in_memory_pilot_dataset, + rulespec_root=stub_source_runtime, + ) assert FakeFrame.latest is None assert FakeAxiomEngine.latest is None @@ -541,7 +587,10 @@ def test_run_derives_in_memory_weights_from_households( 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).output_dataset + 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