Add Belgium pilot: Axiom rules engine over populace-be entity tables

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.

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)"
+)

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",
+]

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),
+                )