Axiom-backed RulesEngine adapter: Belgium as the first Axiom-engine country

[MaxGhenis]

Fixes #260. Part of #259 (populace-be epic).

The first non-PolicyEngine RulesEngine adapter: populace.frame.adapters.axiom.AxiomEngine wraps the Axiom rules engine's dense vectorized surface (axiom_rules_engine.CompiledDenseProgram) for a RuleSpec country module, with Belgium (rulespec-be) as the pilot. Provide-side counterpart: TheAxiomFoundation/axiom-rules-engine#63 (fixes the dense extension build, exposes execute_f64 and derived_metadata) — the engine-backed tests here need that branch.

The four design decisions #260 asked for

1. Engine-native dataset format. Entity-table HDF5 mirroring the US/UK single-year layout — one pandas table per entity plus _time_period — read/written by AxiomEntityTableDataset in the adapter module (the adapter owns the format because there is no policyengine-be package). populace.data generalizes by pointing a registry entry's engine_module/engine_class here; the registry entry itself lands with the release channel (Stand up the populace-be release channel #265).
2. Entity mapping. BE_SCHEMA = EntitySchema(group_entities=("household",)): Belgian PIT is individual with household-level elements, and BE-SILC arrives as person + household registers (Add the BE-SILC source stage with licence guardrails and a private artifact repo #262 owns populating them). RuleSpec scopes rules to engine entities (Person, Household); the adapter maps frame→engine names via entity_names (default: capitalize). Engine entities with no mapped frame entity (rulespec-be also defines Child, Vehicle, Gift, ...) are invisible to the kernel until a frame entity is mapped to them — resolving or materializing their variables raises with the mapping named. Fiscal/benefit units beyond the household enter as group entities when the encoded slice needs them.
3. Formula-owned boundary. The compiled module's derived-rule names are the formula-owned set (exposed via the new derived_metadata surface). write_dataset rejects any persisted derived column — strict from day one, matching the current US adapter posture (a stored engine output would mask reforms). Forbidden/required/defaults/closed-contract semantics mirror PolicyEngineUSEngine.
4. Reform materialization. Decision: no protocol extension now. The engine has no parameter-overlay API — a counterfactual is a different compiled module — so one adapter instance wraps one parameter world and reform runs construct a second adapter over the reform module. The BE validation oracles (Validate populace-be against EUROMOD-BE and Federal Planning Bureau scores #264) therefore sequence behind reform modules compiled upstream (rulespec-be reform variants), not behind a materialize(..., reform=...) protocol change; that extension stays on the protocol's known-future list for engines with true overlay support.

One boundary the issue didn't anticipate: RuleSpec inputs are declared by usage, not typed — the dense surface enumerates input names but no dtypes. The frame's column dtypes are therefore authoritative (bool → Bool for truthiness-context predicates, integer → Integer, float → Decimal/f64), and variable_metadata refuses input names rather than fabricating dtype/period for them. Typed input specs are named follow-up on the engine side (TheAxiomFoundation/axiom-rules-engine#62).

Acceptance, as verified

• Behavioral contract tests parameterized over adapters — new test_rules_engine_contract.py runs one shared suite over PolicyEngineUSEngine and AxiomEngine: protocol conformance, schema/bundle agreement, metadata coherence, input enumeration (sorted, unique, computed excluded), row-aligned materialization, weight authority on export, round-trip persistence. Both cases pass locally; each case skips where its engine isn't installed (CI runs the US case once policyengine-us is present; the Axiom engine is not on PyPI yet).
• Smoke test on the pilot slice — TestBelgianPilotSlice (gated on POPULACE_RULESPEC_BE pointing at a rulespec-be checkout) materializes belgium_pit_article_130_base_tax for a toy 3-person frame and reproduces the hand-computed 2025 liabilities: taxable income 10,000 / 30,000 / 60,000 → 2,500 / 9,612 / 23,620 (25% to 16,320; 40% to 28,800; 45% to 49,840; 50% above). Passing locally against rulespec-be @ 7ddb923.
• Nothing outside the adapter imports the engine — axiom_rules_engine is imported lazily inside one method; the module imports and constructs without it (tested), and the no-engine error names the install path.
• Throughput headroom for population-scale materialization (measured on the full rulespec-be federal PIT pipeline, 49 outputs): 200k persons in ~1.04s Decimal / ~0.87s f64; the adapter's arithmetic="f64" opts into the fast mode.

Notes

• populace-frame[axiom] carries only the PyPI-resolvable dependency (pytables for the HDF5 format); the engine installs from an axiom-rules-engine checkout until it publishes wheels — the ImportError documents this. tables joins the dev group so the dataset-format round-trip tests run in CI without any engine.
• Cross-entity relation batches (dense relation offsets from frame membership columns) are not wired yet: the BE pilot slice declares no relations, and materialize raises NotImplementedError naming the relations if a module declares them.
• The toy fixture module (tests/fixtures/axiom_toy_country.yaml) is a self-contained RuleSpec with person + household scopes, a truthiness predicate, and a Month-period rule, so the full adapter surface exercises without the corpus.

🤖 Generated with Claude Code

[MaxGhenis]

Adversarial re-review (re-run of the review that died at last night's session limit)

Reviewed at d2b7bb5 against the 10-probe spec. Verdict: one real contract violation, fixed in 5239028; otherwise merge-clean. Full stack verified: adapter + contract suites 52 passed/1 skipped pre-fix, whole populace-frame suite 262 passed/3 skipped, no-engine venv proves the lazy-import claim (13 passed/40 skipped, module imports without the engine), Belgian pilot hand-computed values reproduce against the live engine.

Defect → fix

write_dataset ignored contract.formula_owned_excluded (adapters/axiom.py:360-364). The gate only blocked columns that are derived rules of the compiled module and never consulted the contract list, unlike the US adapter (policyengine_us.py:293-295). Reproduced against the installed engine: with formula_owned_excluded=("legacy_output",) and a bundle carrying that column (not a module rule — the "engine reports it as an input but it must never be persisted" case), an open contract silently wrote it, and a closed contract listing it as optional silently wrote it. Neither the shared contract suite nor the Axiom tests set a non-empty formula_owned_excluded, so CI stayed green — the exact divergence class the shared tests mask.

Fixed by unioning the contract list in (one line, mirroring the US adapter) plus two regression tests: open-contract block (mirrors test_policyengine_us_adapter.py:419-444) and closed-contract-with-optional block. Post-fix: adapter + contract files 54 passed/1 skipped; full frame suite 264 passed/3 skipped; ruff clean; no-engine venv unchanged.

Clean probes (highlights)

• Protocol conformance: all 6 methods, row-aligned materialize (shape-checked), gate ordering matches the US adapter. The weight_columns handling from bundle.weighted_entities is more correct than the US adapter's hardcoded {household_weight} — person-weighted exports round-trip here that the US adapter would wrongly reject.
• _program None-caching / metadata bootstrap: household-only modules, undeclared engine entities, and schema-order permutations all behave; unmapped-entity modules raise a clear named error.
• _period_bounds: leap-year correct (2024-02-29, 2000-02-29, 1900-02-28).
• HDF5 dataset: flat keys only from our writer, so no lstrip mangling in practice; __getattr__ recursion-safe (copy/deepcopy/pickle round-trip).
• Round-trip bool: pytables format="table" preserves bool faithfully; the {"i","u","f","b"} compat set never tolerates an actual drift here (noted for symmetry with the US adapter's {"i","u","f"}).
• Fixtures: toy Household program truly needs no person inputs; US contract case verified at policyengine-us 1.754.0 (employment_income_before_lsr input; income_tax/household_net_income computed).
• PR-body claims: hand-computed BE (2500/9612/23620) and toy (500/1000/3500) values reproduce; 200k timing claims are the right order of magnitude on this machine.

Cosmetic (PR body, no code change)

• The pilot-commit reference 7ddb923 is stale — values still reproduce at rulespec-be 7f201e0.
• The f64-faster-than-decimal ordering read as noise at single-output scale; fine as a hedged session note.
• Polish item only: nullable Int64/boolean columns carrying pd.NA fail with a pandas error that doesn't name the column (object/datetime branches do). Loud, never silent — left as-is.

🤖 Generated with Claude Code