From 105c0625dc51b4aabcc6a65578b81686dec48c6f Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 7 Oct 2026 03:31:44 -0400 Subject: [PATCH 1/2] Add Axiom graph adapter surfaces and simulate.rules_by_ref@1 (NZ G1) Salvage the live parts of #821's adapter onto main's axiom.py and add what graph nodes need to run the real Axiom engine (NZ graph plan, package G1): - NZ_SCHEMA (persons in households and families) and NZ_NESTING. With AxiomEngine(nesting=...), a family split across two households is refused from the person memberships alone, whatever the weights. #821's explicit {group}_{parent}_id relation-column check is ported; a contract-required relation column is enforced at export only. - AxiomPeriod(start, end, kind) and periods={label: AxiomPeriod}: "2026-27" resolves to explicit tax_year bounds, and an unmapped label fails closed. - output_dtypes="graph": judgments -> int64 codes, integers -> int64, decimals -> float64, bool -> bool, all lossless; text and date outputs are refused before the engine runs. graph_dtype() names the Owned dtype token. The default ("native") is unchanged, so Belgian behaviour is unchanged. - axiom_engine_ref(): canonical JSON of the engine commit, wheel SHA-256, RuleSpec commit, module path and SHA-256, the whole-root digest (rulespec_tree_digest, the graph's directory-identity formula), arithmetic, output dtypes, schema, entity names, period map and nesting. The reference pins the adapter: it refuses an adapter that already compiled, and every later compile or reference refuses a root whose digest moved. A git checkout root must hold exactly the declared commit (no modified, untracked or ignored files); a root with symlinks is refused. - assert_no_relations(engine, entity) names the module path when a dense program declares relations. - New microcosm.frame.rules_kernels.SimulateRulesByRefKernel (simulate.rules_by_ref@1): {engine_ref: RulesEngine}, one SimulateRulesKernel per binding by composition, so each node keeps simulate.rules@1's params contract, outputs check, all-rows rule and receipt. Its implementation hash binds rules_kernels.py, kernels.py, the frame bundle/rules/schema modules and each distinct adapter class, not the set of refs. kernels.py is untouched (parity pin). Invariants (Hypothesis property tests): - A family nests in exactly one household (refused iff some family spans households). - Only households carry explicit weights in an NZ export (any explicit person or family weight blocks the write under the NZ contract; families resolve their household's weight). - An unmapped period label fails closed (resolves iff mapped). - Graph-typed judgments are lossless (codes -1/0/1 round-trip bit-for-bit). - engine_ref moves iff a pinned byte moves (relocation-invariant; any one-byte edit moves it). Differential tests: rulespec_tree_digest vs microcosm.graph.keys (_directory_identity and source_content_key); simulate.rules_by_ref@1 vs simulate.rules@1 on kernel contexts and in two-engine run_graph runs. Not ported (salvage note): #821's Decimal-object input path (_checked_decimal_array) and its Decimal HDF write path (fixed-format tables for Decimal columns). Graph columns cannot hold object dtype, so they stay on branch codex/nz-transport-microcosm (9f7deadd) for a later lane. Coordination: drafts #826 (feature/axiom-relation-bindings, 952952a1) and #822 (feat/axiom-concept-coverage-be, e6d3c92b) also edit axiom.py. Both already conflict with main there (2 hunks each). With this change, merge-tree shows 5 and 3 hunks, all additive (imports, __all__, docstring, the period line in materialize, appended helpers); no new conflicted files. Co-Authored-By: Claude Opus 5.5 --- changelog.d/axiom-graph-adapter.added.md | 21 + .../src/microcosm/frame/adapters/axiom.py | 739 ++++++- .../src/microcosm/frame/rules_kernels.py | 160 ++ .../shared/test_axiom_graph_adapter.py | 1824 +++++++++++++++++ .../engine_free/shared/test_rules_kernels.py | 665 ++++++ .../zz/policies/tests/axiom_toy_family.yaml | 91 + 6 files changed, 3491 insertions(+), 9 deletions(-) create mode 100644 changelog.d/axiom-graph-adapter.added.md create mode 100644 packages/microcosm-frame/src/microcosm/frame/rules_kernels.py create mode 100644 packages/microcosm-frame/tests/engine_free/shared/test_axiom_graph_adapter.py create mode 100644 packages/microcosm-frame/tests/engine_free/shared/test_rules_kernels.py create mode 100644 packages/microcosm-frame/tests/fixtures/rulespec-zz/zz/policies/tests/axiom_toy_family.yaml diff --git a/changelog.d/axiom-graph-adapter.added.md b/changelog.d/axiom-graph-adapter.added.md new file mode 100644 index 0000000000..52a2de0504 --- /dev/null +++ b/changelog.d/axiom-graph-adapter.added.md @@ -0,0 +1,21 @@ +Prepare the Axiom adapter for graph nodes. `NZ_SCHEMA` declares persons in +households and families. `AxiomEngine(nesting=NZ_NESTING)` refuses a family +whose members sit in two households, from the person memberships alone, and an +explicit `{group}_{parent}_id` column must agree with those memberships. A +`periods={label: AxiomPeriod(...)}` mapping resolves labels such as `"2026-27"` +to explicit `tax_year` bounds and refuses any label it does not map. + +`output_dtypes="graph"` casts outputs to graph-ownable dtypes without loss: +judgments become int64 codes, integers int64, decimals float64. It refuses text +and date outputs. The default output stays as before. + +`axiom_engine_ref` builds an `engine_ref` from the engine commit and wheel +digest, the RuleSpec commit, the module and whole-root content digests, and the +adapter configuration, so a RuleSpec edit re-keys every node that names it. It +also pins the adapter, which then refuses to compile from a root whose bytes +have moved. A git-checkout root must hold exactly the declared commit. +`assert_no_relations` refuses relation-bearing modules. + +The new `simulate.rules_by_ref@1` kernel (`microcosm.frame.rules_kernels`) runs +several rules engines in one graph run. It routes each node to the adapter its +`engine_ref` names, with the `simulate.rules@1` contract unchanged. diff --git a/packages/microcosm-frame/src/microcosm/frame/adapters/axiom.py b/packages/microcosm-frame/src/microcosm/frame/adapters/axiom.py index ad5d996c1e..17e6bb823d 100644 --- a/packages/microcosm-frame/src/microcosm/frame/adapters/axiom.py +++ b/packages/microcosm-frame/src/microcosm/frame/adapters/axiom.py @@ -57,14 +57,49 @@ ``materialize(..., reform=...)`` protocol extension is needed for the BE validation oracles (microcosm#264), which sequence behind reform modules compiled upstream, not behind a protocol change. + +Explicit periods +---------------- +A year label maps to the calendar year and ``YYYY-MM`` to the month unless +the adapter is given a ``periods`` mapping. With one, every period label must +appear in it: ``"2026-27"`` resolves to the :class:`AxiomPeriod` bounds the +caller declared (for example a New Zealand ``tax_year`` from 1 April), and an +unmapped label fails instead of falling back to a calendar year. + +Graph output types +------------------ +By default :meth:`AxiomEngine.materialize` returns the arrays the dense +surface returns: judgments as int8 codes, text and dates as Python string +lists. A graph node may own only a closed set of dtypes, so +``output_dtypes="graph"`` casts each output to one of them without loss — +judgments to int64 codes, integers to int64, decimals to float64, booleans +to bool — and refuses text and date outputs before the engine runs. + +Engine references +----------------- +``simulate.rules@1`` binds a node to its adapter through the node's +``engine_ref`` parameter, and its implementation hash covers adapter source, +not RuleSpec bytes or the native engine build. :func:`axiom_engine_ref` builds +that reference from the pins that decide the adapter's outputs: the engine +commit and wheel digest, the RuleSpec commit, the module path and digest, the +content digest of the whole RuleSpec root, and the adapter's configuration. +Any edit to a RuleSpec byte therefore moves the reference, and with it every +node key that names it. The reference also pins the adapter: from then on it +compiles only from a root whose digest still matches, so a node keyed by the +reference is never computed from other bytes. """ +import hashlib import json import re +import subprocess from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from datetime import date from functools import cache from importlib import resources from pathlib import Path +from types import MappingProxyType from typing import Any import numpy as np @@ -75,12 +110,19 @@ from microcosm.frame.materialize import engine_tables, put_frame_table, read_frame_table from microcosm.frame.rules import ExportContract from microcosm.frame.schema import EntitySchema, VariableMetadata +from microcosm.graph.canonical import canonical_json __all__ = [ "AxiomEngine", "AxiomEntityTableDataset", + "AxiomPeriod", "BE_SCHEMA", + "NZ_NESTING", + "NZ_SCHEMA", + "assert_no_relations", "axiom_concept_mapping", + "axiom_engine_ref", + "rulespec_tree_digest", ] #: The Belgian frame schema for the populace-be pilot: persons in households. @@ -90,6 +132,16 @@ #: needs them, mapped via ``entity_names``. BE_SCHEMA = EntitySchema(group_entities=("household",)) +#: The New Zealand transport schema: persons in households and in benefit +#: units (``family``, the engine's ``Family`` entity). Only households carry +#: explicit weights; a family inherits its household's weight through +#: membership (:meth:`Frame.resolve_weights`). +NZ_SCHEMA = EntitySchema(group_entities=("household", "family")) + +#: Group nesting the New Zealand schema requires: every family lies inside +#: exactly one household. Pass it as ``AxiomEngine(nesting=NZ_NESTING)``. +NZ_NESTING: Mapping[str, str] = MappingProxyType({"family": "household"}) + #: Engine dtype vocabulary -> kernel dtype kind. ``judgment`` is tri-state #: (holds / not holds / undetermined) and materializes as int8 codes #: ``1 / -1 / 0``, so it reports as ``int``, not ``bool``. @@ -113,6 +165,78 @@ _COUNTRY_TREE = re.compile(r"^([a-z]{2})(?:-[a-z0-9]+)*$") _MAPPINGS_DIRECTORY = "axiom_concept_mappings" +#: ``native`` returns the dense surface's arrays unchanged; ``graph`` casts +#: them to graph-ownable dtypes (see "Graph output types" above). +_OUTPUT_DTYPE_MODES: tuple[str, ...] = ("native", "graph") + +#: Engine dtype -> the owned-column dtype token of its graph-typed output. +#: Text and date outputs reach Python as string lists and have no entry. +_GRAPH_DTYPE_BY_ENGINE: dict[str, str] = { + "bool": "bool", + "integer": "int64", + "decimal": "float64", + "judgment": "int64", +} +_GRAPH_OUTPUT_DTYPES: frozenset[str] = frozenset(_GRAPH_DTYPE_BY_ENGINE) + +#: The tri-state judgment codes the dense surface emits. +_JUDGMENT_CODES = np.asarray((-1, 0, 1), dtype=np.int64) + +#: Versioned schema of :func:`axiom_engine_ref` documents. +_ENGINE_REF_SCHEMA = "microcosm.frame.axiom-engine-ref/1" + +# A full git object name (SHA-1 or SHA-256 repositories), lowercase. +_GIT_COMMIT = re.compile(r"^(?:[0-9a-f]{40}|[0-9a-f]{64})$") +_SHA256 = re.compile(r"^[0-9a-f]{64}$") + + +@dataclass(frozen=True) +class AxiomPeriod: + """Explicit dense-execution bounds for one policy period. + + Attributes: + start: First day, ISO ``YYYY-MM-DD``. + end: Last day, ISO ``YYYY-MM-DD``, not before ``start``. + kind: The engine's period identifier (for example ``tax_year``), not + a display label. No fiscal-year convention is inferred from a + year; the caller states the bounds. + + Raises: + ValueError: If ``kind`` is empty or padded with whitespace, a date is + not a valid ISO ``YYYY-MM-DD`` string, or ``start`` follows + ``end``. + """ + + start: str + end: str + kind: str + + def __post_init__(self) -> None: + if ( + not isinstance(self.kind, str) + or not self.kind.strip() + or self.kind != self.kind.strip() + ): + raise ValueError( + "Axiom period kind must be a non-empty string without " + "surrounding whitespace." + ) + try: + start = date.fromisoformat(self.start) + end = date.fromisoformat(self.end) + except (TypeError, ValueError) as exc: + raise ValueError( + "Axiom period dates must be valid ISO YYYY-MM-DD dates." + ) from exc + if start.isoformat() != self.start or end.isoformat() != self.end: + raise ValueError("Axiom period dates must use ISO YYYY-MM-DD format.") + if start > end: + raise ValueError("Axiom period start must not follow end.") + + def bounds(self) -> tuple[str, str, str]: + """Return the dense executor's ``(start, end, period_kind)`` tuple.""" + return self.start, self.end, self.kind + @cache def axiom_concept_mapping(country: str) -> ConceptMapping: @@ -163,6 +287,19 @@ class AxiomEngine: arithmetic: ``"decimal"`` (exact, canonical) or ``"f64"`` (faster, floating-point rounding) — which dense execution mode :meth:`materialize` uses. + periods: Optional period label -> :class:`AxiomPeriod` mapping. When + given, :meth:`materialize` resolves every label through it and an + unmapped label fails, rather than becoming a calendar year or + month. Labels compare as strings, so ``2026`` and ``"2026"`` are + the same label and may not both appear. :meth:`materialize` also + accepts an :class:`AxiomPeriod` directly. + nesting: Optional child group -> parent group mapping (for example + :data:`NZ_NESTING`). Each declared child group must lie inside + exactly one parent group, which :meth:`materialize` and + :meth:`write_dataset` check from the person memberships alone. + output_dtypes: ``"native"`` (default) returns the dense surface's + arrays unchanged; ``"graph"`` casts them to graph-ownable dtypes + and refuses text and date outputs (see the module docstring). The compiled dense programs (one per engine entity) and the module's variable metadata are loaded lazily and cached; constructing the adapter @@ -179,11 +316,19 @@ def __init__( defaults: Mapping[str, object] | None = None, entity_names: Mapping[str, str] | None = None, arithmetic: str = "decimal", + periods: Mapping[int | str, AxiomPeriod] | None = None, + nesting: Mapping[str, str] | None = None, + output_dtypes: str = "native", ) -> None: if arithmetic not in ("decimal", "f64"): raise ValueError( f"arithmetic must be 'decimal' or 'f64', got {arithmetic!r}." ) + if output_dtypes not in _OUTPUT_DTYPE_MODES: + raise ValueError( + f"output_dtypes must be one of {list(_OUTPUT_DTYPE_MODES)}, " + f"got {output_dtypes!r}." + ) if isinstance(rulespec_roots, (str, Path)): raise TypeError( "rulespec_roots must be a non-empty sequence of explicit " @@ -215,11 +360,17 @@ def __init__( f"schema declares {list(schema.entities)}." ) self._arithmetic = arithmetic + self._output_dtypes = output_dtypes + self._periods = _validated_periods(periods) + self._nesting = _validated_nesting(nesting, schema) self._frame_entity_by_engine = { engine: frame for frame, engine in self._entity_names.items() } self._programs: dict[str, Any] = {} self._metadata: dict[str, Any] | None = None + # Set by axiom_engine_ref: the RuleSpec tree digest this adapter's + # reference names. Every later compile re-hashes the root against it. + self._pinned_tree_sha256: str | None = None # ------------------------------------------------------------------ # Variable metadata @@ -265,6 +416,36 @@ def variable_metadata(self, name: str) -> VariableMetadata: period=_PERIOD_BY_ENGINE.get(period, "point"), ) + def graph_dtype(self, name: str) -> str: + """The owned-column dtype token of ``name``'s graph-typed output. + + What a graph node declares in ``Owned(entity, name, dtype)`` when it + materializes ``name`` through an ``output_dtypes="graph"`` adapter: + ``bool``, ``int64`` (integers and judgment codes), or ``float64`` + (decimals). + + Raises: + ImportError: If ``axiom_rules_engine`` is not installed. + ValueError: If the adapter is not ``output_dtypes="graph"`` (its + native arrays, int8 judgment codes among them, are not graph + dtypes), or the variable is unknown, an input, on an unmapped + engine entity, or a text or date output. + """ + if self._output_dtypes != "graph": + raise ValueError( + "graph_dtype describes output_dtypes='graph' adapters only; this " + "adapter returns native engine arrays." + ) + self.variable_metadata(name) + engine_dtype = self._derived_metadata()[name].dtype + token = _GRAPH_DTYPE_BY_ENGINE.get(engine_dtype) + if token is None: + raise ValueError( + f"{name!r} is a {engine_dtype!r} output; no graph column holds " + "text or date values." + ) + return token + def variables(self) -> list[str]: """Return the input variables the engine accepts on a dataset. @@ -332,7 +513,7 @@ def materialize( self, bundle: Frame, variables: Sequence[str], - period: int | str, + period: int | str | AxiomPeriod, ) -> Mapping[str, np.ndarray]: """Compute ``variables`` for ``period`` over the bundle's tables. @@ -344,29 +525,50 @@ def materialize( bundle: A bundle whose entities match the adapter's schema. variables: Computed (derived) variable names. period: ``2025`` / ``"2025"`` for a calendar year, ``"2025-01"`` - for a month. + for a month, or explicit :class:`AxiomPeriod` bounds. When + the adapter has a ``periods`` mapping, a label resolves only + through it. Returns: One array per variable, row-aligned to the variable's entity - table. Judgment variables come back as int8 codes (``1`` holds, - ``-1`` not holds, ``0`` undetermined). + table. By default judgment variables come back as int8 codes + (``1`` holds, ``-1`` not holds, ``0`` undetermined); with + ``output_dtypes="graph"`` they come back as int64 codes and every + array holds a graph-ownable dtype. Raises: ImportError: If ``axiom_rules_engine`` is not installed. ValueError: If the bundle's entities do not match the schema, a - requested variable is unknown or an input, or a computed - array's length does not match its entity table. + declared nesting or relation column is violated, the period + label is unmapped, a requested variable is unknown or an + input, a graph-typed output is text, date, or not losslessly + castable, or a computed array's length does not match its + entity table. NotImplementedError: If the compiled module declares relations (cross-entity aggregation batches; not yet wired — the BE pilot slice declares none). """ self._require_schema(bundle) - start, end, period_kind = _period_bounds(period) + start, end, period_kind = self._materialization_period(period) by_entity: dict[str, list[str]] = {} for name in variables: metadata = self.variable_metadata(name) by_entity.setdefault(metadata.entity, []).append(name) + engine_dtypes: dict[str, str] = {} + if self._output_dtypes == "graph": + derived = self._derived_metadata() + engine_dtypes = {name: derived[name].dtype for name in variables} + unsupported = sorted( + f"{name} ({dtype})" + for name, dtype in engine_dtypes.items() + if dtype not in _GRAPH_OUTPUT_DTYPES + ) + if unsupported: + raise ValueError( + "output_dtypes='graph' cannot type text or date outputs as " + f"graph columns; refused: {unsupported}." + ) results: dict[str, np.ndarray] = {} for frame_entity, names in by_entity.items(): @@ -398,9 +600,33 @@ def materialize( f"{values.shape} but entity {frame_entity!r} has " f"{expected} row(s)." ) + if self._output_dtypes == "graph": + values = _graph_output(name, engine_dtypes[name], values) results[name] = values return results + def _materialization_period( + self, period: int | str | AxiomPeriod + ) -> tuple[str, str, str]: + """Resolve a period to dense ``(start, end, period_kind)`` bounds. + + Raises: + ValueError: If the adapter has a ``periods`` mapping and the label + is not in it, or (without one) the label is neither a year + nor a month. + """ + if isinstance(period, AxiomPeriod): + return period.bounds() + if self._periods is not None: + bounds = self._periods.get(str(period)) + if bounds is None: + raise ValueError( + f"No explicit Axiom period bounds for {period!r}; the " + f"adapter maps only {sorted(self._periods)}." + ) + return bounds.bounds() + return _period_bounds(period) + # ------------------------------------------------------------------ # Export # ------------------------------------------------------------------ @@ -500,6 +726,19 @@ def write_dataset( # Lazy engine plumbing # ------------------------------------------------------------------ + def _module_label(self) -> str: + """The module's path under its RuleSpec root, for error messages. + + Many modules share a file name (``core.yaml``), so a message names the + relative path; a module under no root falls back to its given path. + """ + module = self._module.resolve() + for root in self._rulespec_roots: + resolved = root.resolve() + if module.is_relative_to(resolved): + return module.relative_to(resolved).as_posix() + return str(self._module) + def _import_engine(self) -> Any: try: import axiom_rules_engine @@ -535,6 +774,8 @@ def _program(self, frame_entity: str, *, missing_ok: bool = False) -> Any: f"engine entity {engine_entity!r}." ) return program + if self._pinned_tree_sha256 is not None: + self._require_pinned_tree() engine = self._import_engine() try: program = engine.CompiledDenseProgram.from_file( @@ -566,6 +807,25 @@ def _program(self, frame_entity: str, *, missing_ok: bool = False) -> Any: self._metadata = {item.name: item for item in program.derived_metadata} return program + def _require_pinned_tree(self) -> None: + """Refuse to compile when the root no longer holds the referenced bytes. + + Once :func:`axiom_engine_ref` has named this adapter's RuleSpec tree, + every compile re-hashes the root, so a node keyed by that reference is + never computed from other bytes. + + Raises: + ValueError: If the root's digest moved since the reference. + """ + current, _ = rulespec_tree_digest(self._rulespec_roots[0].resolve()) + if current != self._pinned_tree_sha256: + raise ValueError( + f"RuleSpec root {self._rulespec_roots[0]} changed after this " + "adapter's engine_ref named it (tree digest " + f"{self._pinned_tree_sha256} -> {current}); construct a new " + "adapter and reference." + ) + def _derived_metadata(self) -> dict[str, Any]: """Name -> authoring metadata for every derived rule in the module.""" if self._metadata is None: @@ -580,13 +840,76 @@ def _derived_metadata(self) -> dict[str, Any]: ) return self._metadata - def _require_schema(self, bundle: Frame) -> None: + def _require_schema(self, bundle: Frame, *, export: bool = False) -> None: + """Refuse a bundle whose entities or group nesting the adapter cannot run. + + Beyond matching the schema's entities, two nesting checks run, both + against the person memberships. Weight agreement alone never proves + nesting: a family split across two equally weighted households still + resolves a family weight. + + * Every child -> parent pair in the adapter's ``nesting`` must nest: + all members of one child group share one parent group. + * An explicit group-to-group id column ``{group}_{parent}_id`` (for + example ``family_household_id``) must sit on ``group``'s table and + agree with the members' ``parent`` membership. When ``export`` is + true (the :meth:`write_dataset` path), a relation column the export + contract requires must be present; a graph node materializing + through the adapter need not slice it. + + Raises: + ValueError: On an entity mismatch or a violated nesting check. + """ if set(bundle.entities) != set(self._schema.entities): raise ValueError( f"Axiom adapter requires the schema entities " f"{list(self._schema.entities)}; bundle has " f"{list(bundle.entities)}." ) + schema = self._schema + person = bundle.table(schema.person_entity) + for child, parent in self._nesting.items(): + pairs = pd.DataFrame( + { + "child": person[schema.membership_column(child)].to_numpy(), + "parent": person[schema.membership_column(parent)].to_numpy(), + } + ).drop_duplicates() + spanning = pairs.loc[pairs["child"].duplicated(keep=False), "child"] + if not spanning.empty: + ids = sorted(spanning.unique().tolist()) + raise ValueError( + f"Every {child!r} must nest in exactly one {parent!r}; " + f"{child} id(s) {ids[:5]} have members in more than one " + f"{parent}." + ) + for group in schema.group_entities: + for parent in schema.group_entities: + if group == parent: + continue + column = f"{group}_{parent}_id" + try: + owner = bundle.column_entity(column) + except ValueError: + if export and column in self._contract.required: + raise ValueError( + f"Required relation column {column!r} is missing " + f"from entity {group!r}." + ) from None + continue + if owner != group: + raise ValueError( + f"{column!r} must be on entity {group!r}, not {owner!r}." + ) + broadcast = bundle.place( + column, schema.person_entity, how="broadcast" + ).table(schema.person_entity)[column] + membership = person[schema.membership_column(parent)] + if not np.array_equal(broadcast.to_numpy(), membership.to_numpy()): + raise ValueError( + f"{column!r} disagrees with person membership; every " + f"{group!r} must be nested in its declared {parent!r}." + ) def _engine_tables(self, bundle: Frame) -> dict[str, pd.DataFrame]: """Copy the bundle's tables and materialize typed weights as columns. @@ -595,7 +918,7 @@ def _engine_tables(self, bundle: Frame) -> dict[str, pd.DataFrame]: (typed weights authoritative, any existing ``{entity}_weight`` column overwritten, never trusted), keyed to this adapter's schema order. """ - self._require_schema(bundle) + self._require_schema(bundle, export=True) tables = engine_tables(bundle) return {name: tables[name] for name in self._schema.entities} @@ -764,6 +1087,215 @@ def _read(cls, path: Path) -> tuple[dict[str, pd.DataFrame], int]: return tables, time_period +def rulespec_tree_digest(root: str | Path) -> tuple[str, int]: + """Content digest and payload size of a RuleSpec root directory. + + The formula is the graph's directory source identity + (``microcosm.graph.keys``): SHA-256 over a fixed prefix and then, for every + regular file in sorted path order, its length-prefixed relative POSIX path + and its length-prefixed bytes. The digest therefore equals the content hash + :func:`microcosm.graph.keys.source_content_key` binds when the same + directory is declared as a graph source. Renaming the root is inert; + adding, removing, renaming, or editing any file inside it changes the + digest. A ``.git`` directory, if present, is hashed like any other. + + Args: + root: The RuleSpec root directory (for example a ``git archive`` + export of rulespec-nz). + + Returns: + ``(sha256 hex digest, total bytes of the hashed files)``. + + Raises: + ValueError: If ``root`` is not a directory. + """ + path = Path(root) + if not path.is_dir(): + raise ValueError(f"RuleSpec root {path} is not a directory.") + digest = hashlib.sha256(b"microcosm-graph/source-directory/1\0") + size = 0 + files = sorted(candidate for candidate in path.rglob("*") if candidate.is_file()) + for candidate in files: + relative = candidate.relative_to(path).as_posix().encode("utf-8") + content = candidate.read_bytes() + digest.update(len(relative).to_bytes(8, "little")) + digest.update(relative) + digest.update(len(content).to_bytes(8, "little")) + digest.update(content) + size += len(content) + return digest.hexdigest(), size + + +def axiom_engine_ref( + engine: AxiomEngine, + *, + engine_commit: str, + wheel_sha256: str, + rulespec_root: str | Path, + rulespec_commit: str, +) -> str: + """The ``engine_ref`` that names one Axiom adapter and the bytes behind it. + + A graph node binds its rules engine through ``params["engine_ref"]``, and + that parameter enters the node key. The reference this function returns is + canonical JSON of everything that decides the adapter's outputs: + + * the engine (``axiom-rules-engine``), its commit, and the SHA-256 of the + wheel built from it; + * the RuleSpec commit, the module's path relative to the root, the + module's SHA-256, and :func:`rulespec_tree_digest` of the whole root + (imports resolve anywhere under it); + * the adapter's arithmetic, output dtypes, entity schema, entity names, + period map, and declared nesting. + + For a ``git archive`` export (the normal root, with no ``.git``) no + absolute path enters the reference: the same tree exported to two places + gives the same reference, and editing any byte under the root gives a + different one. The commits are declarations taken from the pin file, and + the directory digest is the authority. + + A git checkout root is accepted only when clean, with no untracked or + ignored files, and at ``rulespec_commit``. Its ``.git`` is then hashed + like any other file, as the graph's source key hashes it. That makes the + reference specific to one clone: two clones or worktrees of one commit + (a worktree's ``.git`` file names an absolute path) and an export of that + commit all give different references, and any later git operation in the + checkout moves it. Prefer the export. + + The first call pins the adapter to the tree digest it names: it refuses + an adapter that has already compiled (those programs were read from bytes + no reference names), and from then on every compile, and every later + reference to the same adapter, refuses a root whose digest has moved. + + Args: + engine: The adapter the reference names. + engine_commit: Full lowercase git commit of axiom-rules-engine. + wheel_sha256: Lowercase SHA-256 of the installed engine wheel. + rulespec_root: The adapter's only RuleSpec root. + rulespec_commit: Full lowercase git commit the root was exported from. + + Returns: + The reference, a canonical JSON string. + + Raises: + TypeError: If ``engine`` is not an :class:`AxiomEngine`. + ValueError: If a pin is malformed; ``rulespec_root`` is not the + adapter's only root; the module is not a file under it; the root + contains a symbolic link (the digest would not cover what it + points at); a git checkout root is dirty or at another commit; or + the adapter compiled before its first reference, or its root moved + after it. + """ + if not isinstance(engine, AxiomEngine): + raise TypeError(f"{engine!r} is not an AxiomEngine.") + pins = ( + ("engine_commit", engine_commit, _GIT_COMMIT, "a full lowercase git commit"), + ( + "rulespec_commit", + rulespec_commit, + _GIT_COMMIT, + "a full lowercase git commit", + ), + ("wheel_sha256", wheel_sha256, _SHA256, "a lowercase SHA-256 hex digest"), + ) + for label, value, pattern, shape in pins: + if not isinstance(value, str) or pattern.fullmatch(value) is None: + raise ValueError(f"{label} must be {shape}, got {value!r}.") + root = Path(rulespec_root) + if not root.is_dir(): + raise ValueError(f"rulespec_root {root} is not a directory.") + resolved_root = root.resolve() + if tuple(item.resolve() for item in engine._rulespec_roots) != (resolved_root,): + raise ValueError( + "axiom_engine_ref digests one RuleSpec root, which must be the " + "adapter's only root; adapter roots " + f"{[str(item) for item in engine._rulespec_roots]}, rulespec_root " + f"{root}." + ) + module = engine._module.resolve() + if not module.is_file() or not module.is_relative_to(resolved_root): + raise ValueError( + f"Module {engine._module} is not a file under rulespec_root {root}." + ) + _refuse_symlinks(resolved_root) + if (resolved_root / ".git").exists(): + _require_clean_checkout(resolved_root, rulespec_commit) + tree_sha256, tree_bytes = rulespec_tree_digest(resolved_root) + if engine._pinned_tree_sha256 is None: + if engine._programs: + raise ValueError( + "axiom_engine_ref must name the adapter's RuleSpec tree before " + "the adapter compiles: it has already compiled programs from " + "bytes no reference names. Compute the reference on a fresh " + "adapter." + ) + engine._pinned_tree_sha256 = tree_sha256 + elif engine._pinned_tree_sha256 != tree_sha256: + raise ValueError( + f"RuleSpec root {root} changed after this adapter's engine_ref named " + f"it (tree digest {engine._pinned_tree_sha256} -> {tree_sha256}); " + "construct a new adapter and reference." + ) + schema = engine._schema + periods = ( + None + if engine._periods is None + else { + label: {"start": item.start, "end": item.end, "kind": item.kind} + for label, item in sorted(engine._periods.items()) + } + ) + document = { + "format": _ENGINE_REF_SCHEMA, + "engine": "axiom-rules-engine", + "engine_commit": engine_commit, + "engine_wheel_sha256": wheel_sha256, + "rulespec_commit": rulespec_commit, + "rulespec_tree_sha256": tree_sha256, + "rulespec_tree_bytes": tree_bytes, + "module": module.relative_to(resolved_root).as_posix(), + "module_sha256": hashlib.sha256(module.read_bytes()).hexdigest(), + "arithmetic": engine._arithmetic, + "output_dtypes": engine._output_dtypes, + "person_entity": schema.person_entity, + "group_entities": list(schema.group_entities), + "entity_names": dict(sorted(engine._entity_names.items())), + "periods": periods, + "nesting": dict(sorted(engine._nesting.items())), + } + return canonical_json(document).decode("utf-8") + + +def assert_no_relations(engine: AxiomEngine, entity: str) -> None: + """Refuse a module whose dense program on ``entity`` declares relations. + + The adapter builds no relation batches from frame membership, so a module + with dense relations cannot run on the graph. Callers check this when + they compose a node, before any data reaches the engine. + + Args: + engine: The adapter to check. + entity: The frame entity the node materializes (``"family"``). + + Raises: + TypeError: If ``engine`` is not an :class:`AxiomEngine`. + ImportError: If ``axiom_rules_engine`` is not installed. + ValueError: If the module has no derived rules on ``entity``. + NotImplementedError: If the program declares relations; the message + names the module and the relations. + """ + if not isinstance(engine, AxiomEngine): + raise TypeError(f"{engine!r} is not an AxiomEngine.") + relations = list(engine._program(entity).relations) + if relations: + raise NotImplementedError( + f"Module {engine._module_label()} declares dense relations " + f"{[item.name for item in relations]} on entity {entity!r}; " + "relation batches from frame membership are not wired, so it " + "cannot run on the graph." + ) + + def _period_bounds(period: int | str) -> tuple[str, str, str]: """Map a kernel period to dense-execution (start, end, period_kind). @@ -820,3 +1352,192 @@ def _batch_from_table( "be bool, integer, or float columns." ) return batch + + +def _validated_periods( + periods: Mapping[int | str, AxiomPeriod] | None, +) -> dict[str, AxiomPeriod] | None: + """Normalize a ``periods`` mapping to string labels, refusing ambiguity. + + Raises: + TypeError: If ``periods`` is not a mapping, a label is not an int or + a non-empty string, or a value is not an :class:`AxiomPeriod`. + ValueError: If the mapping is empty, or two labels share a string + form (``2026`` and ``"2026"``). + """ + if periods is None: + return None + if not isinstance(periods, Mapping): + raise TypeError("periods must map period labels to AxiomPeriod bounds.") + if not periods: + raise ValueError( + "periods must map at least one label; omit it to use calendar years " + "and months." + ) + resolved: dict[str, AxiomPeriod] = {} + for label, bounds in periods.items(): + if ( + isinstance(label, bool) + or not isinstance(label, int | str) + or (isinstance(label, str) and not label) + ): + raise TypeError( + f"periods labels must be ints or non-empty strings, got {label!r}." + ) + if not isinstance(bounds, AxiomPeriod): + raise TypeError("periods values must be AxiomPeriod instances.") + key = str(label) + if key in resolved: + raise ValueError(f"Duplicate explicit Axiom period label {key!r}.") + resolved[key] = bounds + return resolved + + +def _validated_nesting( + nesting: Mapping[str, str] | None, schema: EntitySchema +) -> dict[str, str]: + """Validate a child group -> parent group mapping against ``schema``. + + Raises: + TypeError: If ``nesting`` is not a mapping. + ValueError: If a side is not a group entity of ``schema``, a group + maps to itself, or the declared parents form a cycle. + """ + if nesting is None: + return {} + if not isinstance(nesting, Mapping): + raise TypeError("nesting must map child group entities to parent groups.") + groups = set(schema.group_entities) + resolved: dict[str, str] = {} + for child, parent in nesting.items(): + for side in (child, parent): + if not isinstance(side, str) or side not in groups: + raise ValueError( + f"nesting names {side!r}, which is not a group entity of " + f"the schema {list(schema.group_entities)}." + ) + if child == parent: + raise ValueError(f"nesting maps group {child!r} to itself.") + resolved[child] = parent + for start in resolved: + seen = {start} + current = resolved[start] + while current in resolved: + if current in seen: + raise ValueError(f"nesting is cyclic through group {current!r}.") + seen.add(current) + current = resolved[current] + return resolved + + +def _graph_output(name: str, engine_dtype: str, values: np.ndarray) -> np.ndarray: + """Cast one dense output array to its graph-ownable dtype without loss. + + Bool stays bool; integers and judgment codes become int64; decimals + become float64. Judgment codes must be ``-1``, ``0`` or ``1``. The input + must already hold the engine's numpy kind for that dtype: nothing is + parsed, rounded, or truncated. + + Raises: + ValueError: If ``engine_dtype`` has no graph type, the array's numpy + dtype does not match it, or a value would not survive the cast. + """ + if engine_dtype not in _GRAPH_DTYPE_BY_ENGINE: + raise ValueError( + f"Output {name!r} has engine dtype {engine_dtype!r}, which no graph " + "column holds." + ) + kind = values.dtype.kind + if engine_dtype == "bool" and kind == "b": + return values.astype(np.bool_) + if engine_dtype == "decimal" and kind == "f": + return values.astype(np.float64) + if engine_dtype in ("integer", "judgment") and kind in ("i", "u"): + if kind == "u" and values.size and values.max() > np.iinfo(np.int64).max: + raise ValueError( + f"Output {name!r} holds unsigned values beyond the int64 range." + ) + cast = values.astype(np.int64) + if engine_dtype == "judgment" and not np.isin(cast, _JUDGMENT_CODES).all(): + invalid = sorted(set(np.unique(cast).tolist()) - {-1, 0, 1}) + raise ValueError( + f"Judgment output {name!r} holds codes outside -1/0/1: {invalid[:5]}." + ) + return cast + raise ValueError( + f"Output {name!r} ({engine_dtype}) arrived as numpy dtype {values.dtype}, " + "which does not cast to a graph column without loss." + ) + + +def _refuse_symlinks(root: Path) -> None: + """Refuse a RuleSpec root holding symbolic links. + + The tree digest walks the root without following directory links, so a + linked directory's files would reach the engine without entering the + reference; a linked file would bind its target's bytes under the link's + name. A ``git archive`` export of a RuleSpec repository has none. + + Raises: + ValueError: Naming the first link found. + """ + for directory, directories, files in root.walk(follow_symlinks=False): + for name in sorted((*directories, *files)): + candidate = directory / name + if candidate.is_symlink(): + raise ValueError( + f"RuleSpec root {root} contains the symbolic link " + f"{candidate.relative_to(root).as_posix()}; export the tree " + "without links." + ) + + +def _require_clean_checkout(root: Path, commit: str) -> None: + """Require a git-checkout RuleSpec root to hold exactly ``commit``. + + Ignored files count as changes: the digest hashes them, but the commit + does not hold them. ``--no-optional-locks`` stops ``git status`` + refreshing the index, which would otherwise rewrite ``.git/index`` and + move the digest taken next. + + Raises: + ValueError: If git fails or times out, the checkout has modified, + untracked, or ignored files, or ``HEAD`` is not ``commit``. + """ + + def git(*args: str) -> str: + try: + completed = subprocess.run( + ["git", "--no-optional-locks", "-C", str(root), *args], + capture_output=True, + text=True, + check=False, + timeout=120, + ) + except FileNotFoundError as exc: + raise ValueError( + f"git is required to verify the RuleSpec checkout {root}." + ) from exc + except subprocess.TimeoutExpired as exc: + raise ValueError( + f"git {' '.join(args)} timed out in RuleSpec checkout {root}." + ) from exc + if completed.returncode != 0: + raise ValueError( + f"git {' '.join(args)} failed in RuleSpec checkout {root}: " + f"{completed.stderr.strip()}" + ) + return completed.stdout + + if git("status", "--porcelain=v1", "--untracked-files=all", "--ignored").strip(): + raise ValueError( + f"RuleSpec checkout {root} has modified, untracked, or ignored " + f"files that commit {commit} does not hold; reference a git " + "archive export instead." + ) + head = git("rev-parse", "--verify", "HEAD^{commit}").strip() + if head != commit: + raise ValueError( + f"RuleSpec checkout {root} is at {head}, not the declared " + f"rulespec_commit {commit}." + ) diff --git a/packages/microcosm-frame/src/microcosm/frame/rules_kernels.py b/packages/microcosm-frame/src/microcosm/frame/rules_kernels.py new file mode 100644 index 0000000000..54ccb9345a --- /dev/null +++ b/packages/microcosm-frame/src/microcosm/frame/rules_kernels.py @@ -0,0 +1,160 @@ +"""Graph kernel routing each node to one of several bound rules engines. + +``simulate.rules@1`` (:class:`~microcosm.frame.kernels.SimulateRulesKernel`) +binds one adapter per kernel instance and refuses any other ``engine_ref``, +and a :class:`~microcosm.graph.KernelRegistry` holds one kernel per ref. A run +can therefore reach only one rules engine through ``simulate.rules@1``. An +Axiom adapter wraps one RuleSpec module, so a graph that runs several modules +needs ``simulate.rules_by_ref@1``: it holds an ``engine_ref -> adapter`` +mapping and runs each node exactly as ``simulate.rules@1`` does, through the +adapter the node's ``params["engine_ref"]`` names. + +Node identity is unchanged by the routing: the node's ``engine_ref`` parameter +already enters its key, and the kernel's implementation hash binds code, not +the set of bound references, so adding a binding of an adapter class already +bound re-keys no existing node. +""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Mapping +from types import MappingProxyType + +import microcosm.frame.bundle as frame_bundle_module +import microcosm.frame.rules as frame_rules_module +import microcosm.frame.schema as frame_schema_module +from microcosm.frame.kernels import SimulateRulesKernel +from microcosm.frame.rules import RulesEngine +from microcosm.graph import KernelBase, KernelContext, KernelResult, source_hash + +__all__ = ["SimulateRulesByRefKernel"] + +_PARAMS = frozenset({"engine_ref", "period", "variables"}) + + +class SimulateRulesByRefKernel(KernelBase): + """Materialize declared variables through the adapter a node's ref names. + + Args: + engines: Non-empty mapping from ``engine_ref`` to the adapter it names. + Each reference must uniquely identify its adapter and that + adapter's configuration (for an Axiom adapter, use + :func:`~microcosm.frame.adapters.axiom.axiom_engine_ref`). + dependencies: Distribution names whose versions affect every bound + engine's behavior, as for ``simulate.rules@1``. + + Each binding is a :class:`~microcosm.frame.kernels.SimulateRulesKernel` + held by composition, so a node gets that kernel's parameter contract + (exactly ``engine_ref``, ``period`` and ``variables``), its declared- + outputs check, its all-rows rule, and its receipt, unchanged. + + The implementation hash binds this module, ``microcosm.frame.kernels``, + the frame modules ``simulate.rules@1`` binds, and the source of each + distinct adapter *class*. It does not bind the references: a second + binding of a class already bound leaves every node key where it was. A + binding of a new adapter class adds that class's source to the hash and so + re-keys every node of this kernel, as replacing an adapter class does for + ``simulate.rules@1``. + + Raises: + TypeError: If ``engines`` is not a mapping, or a binding is not a + :class:`~microcosm.frame.rules.RulesEngine`. + ValueError: If ``engines`` is empty, or a reference is empty. + """ + + ref = "simulate.rules_by_ref@1" + + def __init__( + self, + engines: Mapping[str, RulesEngine], + dependencies: tuple[str, ...] = (), + ) -> None: + if not isinstance(engines, Mapping): + raise TypeError( + "engines must be a mapping from engine_ref to a RulesEngine." + ) + if not engines: + raise ValueError("engines must bind at least one engine_ref.") + delegates = { + engine_ref: SimulateRulesKernel(engine_ref, engine, dependencies) + for engine_ref, engine in engines.items() + } + self._delegates = MappingProxyType(dict(sorted(delegates.items()))) + self._adapter_classes = tuple( + sorted( + {type(engine) for engine in engines.values()}, + key=lambda cls: (cls.__module__, cls.__qualname__), + ) + ) + # Every delegate declares the same contract; take it from one so the + # two kernels cannot drift apart. + self.capabilities = next(iter(self._delegates.values())).capabilities + + def engine_refs(self) -> tuple[str, ...]: + """The bound references, sorted.""" + + return tuple(self._delegates) + + def implementation_hash(self) -> str: + """Hash the routing and delegate kernels plus every bound adapter class.""" + + return source_hash( + type(self), + SimulateRulesKernel, + *self._adapter_classes, + frame_bundle_module, + frame_rules_module, + frame_schema_module, + dependencies=self.capabilities.dependencies, + ) + + def run(self, context: KernelContext) -> KernelResult: + """Run the node through the delegate its ``engine_ref`` names. + + The parameter set and the reference's type are checked before + routing, in the order and with the exception types + ``simulate.rules@1`` uses, so a malformed node fails the same way + under either kernel. + """ + + actual = set(context.params) + if actual != _PARAMS: + raise ValueError( + "simulate.rules_by_ref parameters must be exactly " + f"{sorted(_PARAMS)!r}; got {sorted(actual)!r}." + ) + engine_ref = context.params["engine_ref"] + if not isinstance(engine_ref, str) or not engine_ref: + raise TypeError( + "simulate.rules_by_ref engine_ref must be a non-empty string." + ) + delegate = self._delegates.get(engine_ref) + if delegate is None: + bound = ", ".join(_describe(ref) for ref in self._delegates) + raise ValueError( + "simulate.rules_by_ref has no engine bound to engine_ref " + f"{_describe(engine_ref)}; bound: {bound}." + ) + return delegate.run(context) + + +def _describe(engine_ref: str) -> str: + """Name a reference briefly: its SHA-256 prefix, and its module if any. + + References are often long canonical JSON whose leading keys coincide + (every Axiom reference starts with the same arithmetic and engine pins), so + a truncated prefix would not tell two apart. + """ + + digest = hashlib.sha256(engine_ref.encode("utf-8")).hexdigest()[:16] + try: + document = json.loads(engine_ref) + except ValueError: + document = None + module = document.get("module") if isinstance(document, dict) else None + if isinstance(module, str): + return f"sha256:{digest} (module {module!r})" + preview = engine_ref if len(engine_ref) <= 60 else engine_ref[:60] + "..." + return f"sha256:{digest} ({preview!r})" diff --git a/packages/microcosm-frame/tests/engine_free/shared/test_axiom_graph_adapter.py b/packages/microcosm-frame/tests/engine_free/shared/test_axiom_graph_adapter.py new file mode 100644 index 0000000000..30803e9041 --- /dev/null +++ b/packages/microcosm-frame/tests/engine_free/shared/test_axiom_graph_adapter.py @@ -0,0 +1,1824 @@ +"""Axiom adapter surfaces for graph nodes: periods, nesting, dtypes, engine refs. + +Everything above the ``needs_engine`` classes runs without +``axiom_rules_engine``: period and nesting checks run before the engine is +touched, the graph dtype cast is a pure function, ``axiom_engine_ref`` reads +only bytes on disk, and materialization is exercised through a recording +stand-in for the dense surface that returns exactly the array types the native +extension returns (int8 judgment codes, float64 decimals, int64 integers, bool, +and Python string lists for text and dates). + +The ``needs_engine`` classes run the real engine on the ``rulespec-zz`` +fixture (the existing skip pattern): a graph-typed ``simulate.rules@1`` node +in a FILTER-free graph, the ``tax_year`` period label, family-level judgment +codes, and ``simulate.rules_by_ref@1`` with two Axiom bindings in one run. +""" + +import hashlib +import importlib.util +import json +import os +import re +import shutil +import subprocess +import tarfile +from collections.abc import Mapping +from pathlib import Path + +import numpy as np +import pandas as pd +import pytest +from hypothesis import HealthCheck, given, settings +from hypothesis import strategies as st + +from microcosm.frame import ExportContract, Frame, WeightKind, Weights +from microcosm.frame.adapters.axiom import ( + BE_SCHEMA, + NZ_NESTING, + NZ_SCHEMA, + AxiomEngine, + AxiomPeriod, + _graph_output, + assert_no_relations, + axiom_engine_ref, + rulespec_tree_digest, +) +from microcosm.frame.kernels import SimulateRulesKernel +from microcosm.frame.rules_kernels import SimulateRulesByRefKernel +from microcosm.graph import ( + Capabilities, + ContentStore, + Determinism, + Graph, + KernelBase, + KernelContext, + KernelRegistry, + KernelResult, + Node, + Numeric, + Owned, + SeedSource, + Slice, + SourceRef, + StructuralDelta, + compile_graph, + run_graph, +) +from microcosm.graph.canonical import canonical_json, sha256_domain +from microcosm.graph.keys import _directory_identity, source_content_key +from test_support.paths import paths_for + +_TEST_PATHS = paths_for("microcosm-frame") + +_ENGINE_INSTALLED = importlib.util.find_spec("axiom_rules_engine") is not None +if _ENGINE_INSTALLED: + from axiom_rules_engine.dense import NativeCompiledDenseProgram + + _DENSE_AVAILABLE = NativeCompiledDenseProgram is not None +else: + _DENSE_AVAILABLE = False + +needs_engine = pytest.mark.skipif( + not _DENSE_AVAILABLE, + reason="axiom_rules_engine (with the dense native extension) is not installed", +) +needs_tables = pytest.mark.skipif( + importlib.util.find_spec("tables") is None, + reason="pytables (microcosm-frame[axiom]) is not installed", +) + +FIXTURE_RULESPEC_ROOT = _TEST_PATHS.tests / "fixtures" / "rulespec-zz" +FIXTURE_MODULE = FIXTURE_RULESPEC_ROOT / "zz/policies/tests/axiom_toy_country.yaml" +FAMILY_FIXTURE_MODULE = ( + FIXTURE_RULESPEC_ROOT / "zz/policies/tests/axiom_toy_family.yaml" +) +FIXTURE_RULESPEC_ROOTS = (FIXTURE_RULESPEC_ROOT,) + +TAX_YEAR = AxiomPeriod(start="2026-04-01", end="2027-03-31", kind="tax_year") +ENGINE_COMMIT = "a" * 40 +WHEEL_SHA256 = "b" * 64 +RULESPEC_COMMIT = "c" * 40 + +# A full override of every git setting a test repository depends on, so the +# developer's global configuration (hooks, signing) never reaches it. +_GIT_ENV = { + **os.environ, + "GIT_CONFIG_GLOBAL": os.devnull, + "GIT_CONFIG_NOSYSTEM": "1", +} + + +def _git(root: Path, *args: str) -> str: + completed = subprocess.run( + [ + "git", + "-c", + "user.name=microcosm-test", + "-c", + "user.email=microcosm-test@example.invalid", + "-c", + "commit.gpgsign=false", + "-c", + f"core.hooksPath={os.devnull}", + "-C", + str(root), + *args, + ], + capture_output=True, + text=True, + check=True, + env=_GIT_ENV, + ) + return completed.stdout.strip() + + +# ---------------------------------------------------------------------- +# Frames +# ---------------------------------------------------------------------- + + +def _nz_frame( + person_households: list[int], + person_families: list[int], + household_weights: list[float], + *, + family_columns: Mapping[str, list[object]] | None = None, + person_columns: Mapping[str, list[object]] | None = None, +) -> Frame: + """An NZ_SCHEMA frame with household-only design weights.""" + + count = len(person_households) + person = pd.DataFrame( + { + "person_id": np.arange(1, count + 1, dtype=np.int64), + "person_household_id": np.asarray(person_households, dtype=np.int64), + "person_family_id": np.asarray(person_families, dtype=np.int64), + **(person_columns or {}), + } + ) + household = pd.DataFrame( + {"household_id": np.unique(np.asarray(person_households, dtype=np.int64))} + ) + family = pd.DataFrame( + { + "family_id": np.unique(np.asarray(person_families, dtype=np.int64)), + **(family_columns or {}), + } + ) + return Frame( + {"person": person, "household": household, "family": family}, + NZ_SCHEMA, + { + "household": Weights( + values=np.asarray(household_weights, dtype=np.float64), + kind=WeightKind.DESIGN, + ) + }, + ) + + +@st.composite +def nz_populations(draw: st.DrawFn) -> dict[str, object]: + """A nested NZ population, optionally with one family split across households. + + Households hold one to three families and families one to three persons, + so every family nests in exactly one household. When ``split`` is drawn, + one person of a family with at least two members moves to another + household while keeping the family, and that family then spans two + households. Weights are drawn equal or distinct so a test can show weight + agreement never hides the split. + """ + + household_count = draw(st.integers(min_value=2, max_value=5)) + person_households: list[int] = [] + person_families: list[int] = [] + family_household: dict[int, int] = {} + next_family = 10 + for household in range(1, household_count + 1): + for _ in range(draw(st.integers(min_value=1, max_value=3))): + family = next_family + next_family += 1 + family_household[family] = household + for _ in range(draw(st.integers(min_value=1, max_value=3))): + person_households.append(household) + person_families.append(family) + splittable = [ + family for family in family_household if person_families.count(family) >= 2 + ] + split = bool(splittable) and draw(st.booleans()) + if split: + family = draw(st.sampled_from(splittable)) + position = person_families.index(family) + others = [ + household + for household in range(1, household_count + 1) + if household != family_household[family] + ] + person_households[position] = draw(st.sampled_from(others)) + equal_weights = draw(st.booleans()) + weights = [ + 250.0 if equal_weights else float(100 + 37 * household) + for household in range(1, household_count + 1) + ] + return { + "person_households": person_households, + "person_families": person_families, + "household_weights": weights, + "split": split, + } + + +# ---------------------------------------------------------------------- +# A recording stand-in for the dense surface +# ---------------------------------------------------------------------- + + +class _Metadata: + """The authoring metadata fields the adapter reads.""" + + def __init__(self, name: str, entity: str, dtype: str, period: str = "Year"): + self.name = name + self.entity = entity + self.dtype = dtype + self.period = period + + +class _Relation: + def __init__(self, name: str) -> None: + self.name = name + + +_FAKE_METADATA = ( + _Metadata("person_benefit", "Person", "decimal"), + _Metadata("family_assets_ok", "Family", "judgment"), + _Metadata("family_assistance", "Family", "decimal"), + _Metadata("family_band", "Family", "integer"), + _Metadata("family_is_large", "Family", "bool"), + _Metadata("family_category", "Family", "text"), + _Metadata("family_review_date", "Family", "date"), +) + + +class _FakeProgram: + """One entity's dense program; outputs mirror the native return types.""" + + def __init__(self, entity: str, relations: tuple[_Relation, ...] = ()) -> None: + self.entity = entity + self.root_inputs = ["family_rent", "person_income"] + self.relations = list(relations) + self.derived_metadata = list(_FAKE_METADATA) + self.calls: list[dict[str, object]] = [] + + def execute(self, *, period_kind, start, end, inputs, outputs): + self.calls.append( + { + "period_kind": period_kind, + "start": start, + "end": end, + "inputs": sorted(inputs), + "outputs": list(outputs), + } + ) + rows = len(next(iter(inputs.values()))) + codes = np.resize(np.asarray([1, -1, 0], dtype=np.int8), rows) + available = { + "person_benefit": np.full(rows, 12.5), + "family_assets_ok": codes, + "family_assistance": np.linspace(0.0, 70.25, rows), + "family_band": np.arange(rows, dtype=np.int64), + "family_is_large": codes > 0, + "family_category": ["small"] * rows, + "family_review_date": ["2026-04-01"] * rows, + } + return { + "row_count": rows, + "outputs": {name: available[name] for name in outputs}, + } + + execute_f64 = execute + + +def _fake_engine( + adapter: AxiomEngine, + monkeypatch: pytest.MonkeyPatch, + *, + relations: tuple[_Relation, ...] = (), +) -> dict[str, _FakeProgram]: + """Swap the adapter's engine import for the recording stand-in.""" + + programs: dict[str, _FakeProgram] = {} + + class _Compiled: + @classmethod + def from_file(cls, path, *, rulespec_roots, entity): + if entity not in {"Person", "Family"}: + raise ValueError( + "dense compilation could not find derived outputs for entity " + f"`{entity}`" + ) + program = _FakeProgram(entity, relations) + programs[entity] = program + return program + + class _Engine: + CompiledDenseProgram = _Compiled + + monkeypatch.setattr(adapter, "_import_engine", lambda: _Engine) + return programs + + +def _fake_frame() -> Frame: + return _nz_frame( + [1, 1, 2, 2, 2], + [10, 10, 20, 21, 21], + [300.0, 500.0], + family_columns={"family_rent": [320.0, 0.0, 180.0]}, + person_columns={"person_income": [1.0, 2.0, 3.0, 4.0, 5.0]}, + ) + + +def _nz_adapter(**kwargs) -> AxiomEngine: + return AxiomEngine( + FIXTURE_MODULE, + schema=NZ_SCHEMA, + rulespec_roots=FIXTURE_RULESPEC_ROOTS, + **kwargs, + ) + + +# ---------------------------------------------------------------------- +# NZ schema +# ---------------------------------------------------------------------- + + +class TestNzSchema: + def test_nz_schema_has_households_and_families(self) -> None: + assert NZ_SCHEMA.entities == ("person", "household", "family") + assert dict(NZ_NESTING) == {"family": "household"} + + def test_nz_nesting_is_read_only(self) -> None: + with pytest.raises(TypeError): + NZ_NESTING["family"] = "person" # type: ignore[index] + + @settings(max_examples=60, deadline=None) + @given(population=nz_populations()) + def test_household_weights_resolve_onto_families(self, population) -> None: + """Invariant (layout half): families inherit their household's weight. + + With household-only weights, every family's effective weight is its + one household's weight, and the adapter's export tables persist + ``household_weight`` alone. The enforcing half, an export refusing any + other explicit weight, is + ``test_property_only_household_weights_reach_an_export``. + """ + + if population["split"]: + return + frame = _nz_frame( + population["person_households"], + population["person_families"], + population["household_weights"], + ) + assert frame.weighted_entities == ("household",) + household_weight = dict( + zip( + frame.table("household")["household_id"], + population["household_weights"], + strict=True, + ) + ) + person = frame.table("person") + family_household = person.groupby("person_family_id")[ + "person_household_id" + ].first() + expected = [ + household_weight[family_household[family]] + for family in frame.table("family")["family_id"] + ] + np.testing.assert_array_equal( + frame.resolve_weights("family").values, np.asarray(expected) + ) + tables = _nz_adapter(nesting=NZ_NESTING)._engine_tables(frame) + assert "household_weight" in tables["household"].columns + assert "family_weight" not in tables["family"].columns + assert "person_weight" not in tables["person"].columns + + +#: The New Zealand export rule #821 carried: weights persist on households +#: only, so the person and family weight columns are forbidden. +_NZ_WEIGHT_CONTRACT = ExportContract( + required=(), + forbidden=("person_weight", "family_weight"), + optional=(), + formula_owned_excluded=(), +) + + +class TestExportWeights: + @settings(max_examples=60, deadline=None) + @given( + population=nz_populations(), + extra=st.sets(st.sampled_from(["person", "family"]), min_size=1), + ) + def test_property_only_household_weights_reach_an_export( + self, tmp_path_factory, population, extra + ) -> None: + """Invariant: only households carry explicit weights in an NZ export. + + The executor hands a kernel resolved weights for every entity it + projects (``G/executor.py`` builds ``KernelContext.weights`` for each + one), so ``materialize`` cannot refuse explicit family weights. The + export can: with the NZ contract, any explicit person or family weight + blocks the write, and nothing is written. + """ + + if population["split"]: + return + base = _nz_frame( + population["person_households"], + population["person_families"], + population["household_weights"], + ) + weights = {"household": base.weights_for("household")} + for entity in extra: + weights[entity] = Weights( + values=np.full(base.n(entity), 7.0), kind=WeightKind.DESIGN + ) + frame = Frame( + {entity: base.table(entity) for entity in NZ_SCHEMA.entities}, + NZ_SCHEMA, + weights, + ) + adapter = _nz_adapter(nesting=NZ_NESTING, contract=_NZ_WEIGHT_CONTRACT) + forbidden = sorted(f"{entity}_weight" for entity in extra) + with pytest.MonkeyPatch.context() as monkeypatch: + _fake_engine(adapter, monkeypatch) + path = tmp_path_factory.mktemp("export") / "weighted.h5" + with pytest.raises( + ValueError, + match=f"forbidden column\\(s\\) present: {re.escape(repr(forbidden))}", + ): + adapter.write_dataset(frame, path, 2026) + assert not path.exists() + + @needs_tables + def test_a_household_weighted_frame_exports_household_weight_alone( + self, tmp_path, monkeypatch + ) -> None: + from microcosm.frame.adapters.axiom import AxiomEntityTableDataset + + adapter = _nz_adapter(nesting=NZ_NESTING, contract=_NZ_WEIGHT_CONTRACT) + _fake_engine(adapter, monkeypatch) + path = tmp_path / "household_weighted.h5" + adapter.write_dataset(_nz_frame([1, 1, 2], [1, 1, 2], [2.0, 3.0]), path, 2026) + reloaded = AxiomEntityTableDataset(file_path=path) + assert reloaded.household["household_weight"].tolist() == [2.0, 3.0] + assert "family_weight" not in reloaded.family.columns + assert "person_weight" not in reloaded.person.columns + + +# ---------------------------------------------------------------------- +# Periods +# ---------------------------------------------------------------------- + + +class TestAxiomPeriod: + def test_bounds_are_the_dense_executor_tuple(self) -> None: + assert TAX_YEAR.bounds() == ("2026-04-01", "2027-03-31", "tax_year") + + def test_refuses_reversed_dates(self) -> None: + with pytest.raises(ValueError, match="start.*end"): + AxiomPeriod(start="2027-03-31", end="2026-04-01", kind="tax_year") + + @pytest.mark.parametrize( + "start", ["2026-02-30", "20260401", "2026-4-1", "", None, 20260401] + ) + def test_refuses_dates_that_are_not_canonical_iso(self, start) -> None: + with pytest.raises(ValueError, match="ISO"): + AxiomPeriod(start=start, end="2027-03-31", kind="tax_year") + + @pytest.mark.parametrize("kind", ["", " ", " tax_year", "tax_year\n", None]) + def test_refuses_an_empty_or_padded_kind(self, kind) -> None: + with pytest.raises(ValueError, match="kind"): + AxiomPeriod(start="2026-04-01", end="2027-03-31", kind=kind) + + def test_a_single_day_period_is_valid(self) -> None: + assert AxiomPeriod("2026-04-01", "2026-04-01", "day").bounds()[0] == ( + "2026-04-01" + ) + + +class TestPeriodLabels: + def test_a_mapped_label_resolves_to_its_explicit_bounds(self) -> None: + adapter = _nz_adapter(periods={"2026-27": TAX_YEAR}) + assert adapter._materialization_period("2026-27") == TAX_YEAR.bounds() + + def test_int_and_string_labels_are_one_label(self) -> None: + adapter = _nz_adapter(periods={2026: TAX_YEAR}) + assert adapter._materialization_period(2026) == TAX_YEAR.bounds() + assert adapter._materialization_period("2026") == TAX_YEAR.bounds() + + @pytest.mark.parametrize("label", ["2026", 2026, "2026-04", "2027-28", "2026-27 "]) + def test_an_unmapped_label_fails_closed(self, label) -> None: + adapter = _nz_adapter(periods={"2026-27": TAX_YEAR}) + with pytest.raises(ValueError, match="No explicit Axiom period bounds"): + adapter._materialization_period(label) + + def test_explicit_bounds_pass_through_a_mapping(self) -> None: + other = AxiomPeriod("2025-04-01", "2026-03-31", "tax_year") + adapter = _nz_adapter(periods={"2026-27": TAX_YEAR}) + assert adapter._materialization_period(other) == other.bounds() + + def test_without_a_mapping_calendar_labels_are_unchanged(self) -> None: + adapter = _nz_adapter() + assert adapter._materialization_period(2025) == ( + "2025-01-01", + "2025-12-31", + "calendar_year", + ) + # "2026-27" reads as month 27 of 2026 and fails; it never becomes + # the April-to-March tax year without an explicit mapping. + with pytest.raises(ValueError, match="Invalid month"): + adapter._materialization_period("2026-27") + + def test_duplicate_string_forms_are_refused(self) -> None: + with pytest.raises(ValueError, match="Duplicate"): + _nz_adapter(periods={2026: TAX_YEAR, "2026": TAX_YEAR}) + + @pytest.mark.parametrize("label", [True, "", 2026.0, None]) + def test_labels_must_be_ints_or_non_empty_strings(self, label) -> None: + with pytest.raises(TypeError, match="labels"): + _nz_adapter(periods={label: TAX_YEAR}) + + def test_values_must_be_axiom_periods(self) -> None: + with pytest.raises(TypeError, match="AxiomPeriod"): + _nz_adapter(periods={"2026-27": ("2026-04-01", "2027-03-31", "x")}) + + def test_an_empty_mapping_is_refused(self) -> None: + with pytest.raises(ValueError, match="at least one label"): + _nz_adapter(periods={}) + + @settings(max_examples=100, deadline=None) + @given( + labels=st.sets( + st.text(alphabet="0123456789-ab", min_size=1, max_size=9), + min_size=1, + max_size=6, + ), + query=st.text(alphabet="0123456789-ab", min_size=1, max_size=9), + ) + def test_property_a_label_resolves_iff_it_is_mapped(self, labels, query) -> None: + """Invariant: an unmapped period label fails closed, never defaults.""" + + periods = { + label: AxiomPeriod("2026-04-01", f"2027-03-{10 + index:02d}", "tax_year") + for index, label in enumerate(sorted(labels)) + } + adapter = _nz_adapter(periods=periods) + if query in periods: + assert adapter._materialization_period(query) == periods[query].bounds() + else: + with pytest.raises(ValueError, match="No explicit Axiom period bounds"): + adapter._materialization_period(query) + + def test_materialize_hands_the_mapped_bounds_to_the_engine( + self, monkeypatch + ) -> None: + adapter = _nz_adapter(periods={"2026-27": TAX_YEAR}) + programs = _fake_engine(adapter, monkeypatch) + adapter.materialize(_fake_frame(), ["family_assistance"], "2026-27") + assert programs["Family"].calls == [ + { + "period_kind": "tax_year", + "start": "2026-04-01", + "end": "2027-03-31", + "inputs": ["family_rent"], + "outputs": ["family_assistance"], + } + ] + + def test_materialize_refuses_an_unmapped_label_before_the_engine( + self, monkeypatch + ) -> None: + adapter = _nz_adapter(periods={"2026-27": TAX_YEAR}) + programs = _fake_engine(adapter, monkeypatch) + with pytest.raises(ValueError, match="No explicit Axiom period bounds"): + adapter.materialize(_fake_frame(), ["family_assistance"], 2026) + assert programs == {} + + +# ---------------------------------------------------------------------- +# Nesting +# ---------------------------------------------------------------------- + + +class TestNesting: + @pytest.mark.parametrize( + ("nesting", "message"), + [ + ({"family": "tax_unit"}, "not a group entity"), + ({"person": "household"}, "not a group entity"), + ({"family": "family"}, "to itself"), + ({"family": "household", "household": "family"}, "cyclic"), + ], + ) + def test_declarations_are_validated_against_the_schema( + self, nesting, message + ) -> None: + with pytest.raises(ValueError, match=message): + _nz_adapter(nesting=nesting) + + def test_a_non_mapping_is_refused(self) -> None: + with pytest.raises(TypeError, match="nesting"): + _nz_adapter(nesting=[("family", "household")]) + + def test_cross_household_family_is_refused_from_memberships_alone( + self, + ) -> None: + # Equal household weights: the family weight still resolves, so weight + # agreement cannot be what detects the split. + frame = _nz_frame([1, 2], [1, 1], [2.0, 2.0]) + assert frame.resolve_weights("family").values.tolist() == [2.0] + with pytest.raises(ValueError, match="exactly one 'household'.*\\[1\\]"): + _nz_adapter(nesting=NZ_NESTING).materialize(frame, [], 2026) + + def test_without_a_declaration_the_split_is_not_checked(self) -> None: + frame = _nz_frame([1, 2], [1, 1], [2.0, 2.0]) + assert _nz_adapter().materialize(frame, [], 2026) == {} + + def test_a_nested_population_passes(self) -> None: + frame = _nz_frame([1, 1, 2, 2], [1, 1, 2, 3], [2.0, 3.0]) + assert _nz_adapter(nesting=NZ_NESTING).materialize(frame, [], 2026) == {} + + def test_the_split_blocks_an_export_before_anything_is_written( + self, tmp_path + ) -> None: + frame = _nz_frame([1, 2], [1, 1], [2.0, 2.0]) + path = tmp_path / "split.h5" + with pytest.raises(ValueError, match="exactly one 'household'"): + _nz_adapter(nesting=NZ_NESTING).write_dataset(frame, path, 2026) + assert not path.exists() + + def test_explicit_relation_column_disagreeing_with_membership_is_refused( + self, + ) -> None: + frame = _nz_frame( + [1, 2], [1, 1], [2.0, 2.0], family_columns={"family_household_id": [1]} + ) + with pytest.raises(ValueError, match="family_household_id.*membership"): + _nz_adapter().materialize(frame, [], 2026) + + def test_explicit_relation_column_agreeing_with_membership_passes( + self, + ) -> None: + frame = _nz_frame( + [1, 1, 2], + [1, 1, 2], + [2.0, 3.0], + family_columns={"family_household_id": [1, 2]}, + ) + assert _nz_adapter().materialize(frame, [], 2026) == {} + + def test_relation_column_on_the_wrong_entity_blocks_materialize_and_export( + self, tmp_path + ) -> None: + frame = _nz_frame( + [1, 2], [1, 1], [2.0, 2.0], person_columns={"family_household_id": [1, 2]} + ) + adapter = _nz_adapter() + with pytest.raises(ValueError, match="family_household_id.*family.*person"): + adapter.materialize(frame, [], 2026) + path = tmp_path / "wrong_owner.h5" + with pytest.raises(ValueError, match="family_household_id.*family.*person"): + adapter.write_dataset(frame, path, 2026) + assert not path.exists() + + def test_a_required_relation_column_is_an_export_rule(self, tmp_path) -> None: + # A graph node need not slice an export-only column to materialize; + # the export itself refuses the frame and writes nothing. + contract = ExportContract( + required=("family_household_id",), + forbidden=(), + optional=(), + formula_owned_excluded=(), + ) + frame = _nz_frame([1, 1, 2], [1, 1, 2], [2.0, 3.0]) + adapter = _nz_adapter(contract=contract) + assert adapter.materialize(frame, [], 2026) == {} + path = tmp_path / "missing_relation.h5" + with pytest.raises(ValueError, match="Required relation column"): + adapter.write_dataset(frame, path, 2026) + assert not path.exists() + + @settings( + max_examples=80, + deadline=None, + suppress_health_check=[HealthCheck.too_slow], + ) + @given(population=nz_populations()) + def test_property_refused_iff_a_family_spans_households(self, population) -> None: + """Invariant: a family nests in exactly one household. + + The declared nesting refuses a population exactly when some family + has members in two households, whatever the weights. + """ + + frame = _nz_frame( + population["person_households"], + population["person_families"], + population["household_weights"], + ) + person = frame.table("person") + spans = ( + person.groupby("person_family_id")["person_household_id"].nunique() > 1 + ).any() + assert bool(spans) == population["split"] + adapter = _nz_adapter(nesting=NZ_NESTING) + if spans: + with pytest.raises(ValueError, match="exactly one 'household'"): + adapter.materialize(frame, [], 2026) + else: + assert adapter.materialize(frame, [], 2026) == {} + + +# ---------------------------------------------------------------------- +# Graph output dtypes +# ---------------------------------------------------------------------- + + +class TestGraphOutputCast: + def test_judgment_int8_codes_become_int64(self) -> None: + codes = np.asarray([1, -1, 0, 1], dtype=np.int8) + cast = _graph_output("j", "judgment", codes) + assert cast.dtype == np.dtype(np.int64) + assert cast.tolist() == [1, -1, 0, 1] + + @settings(max_examples=200, deadline=None) + @given( + codes=st.lists(st.sampled_from([-1, 0, 1]), min_size=0, max_size=64), + source=st.sampled_from([np.int8, np.int16, np.int32, np.int64]), + ) + def test_property_judgment_codes_survive_the_cast_losslessly( + self, codes, source + ) -> None: + """Invariant: graph-typed judgments are lossless (codes -1/0/1).""" + + native = np.asarray(codes, dtype=source) + cast = _graph_output("j", "judgment", native) + assert cast.dtype == np.dtype(np.int64) + assert cast.astype(source).tobytes() == native.tobytes() + + @pytest.mark.parametrize("bad", [2, -2, 127]) + def test_judgment_codes_outside_the_tri_state_are_refused(self, bad) -> None: + with pytest.raises(ValueError, match="outside -1/0/1"): + _graph_output("j", "judgment", np.asarray([0, bad], dtype=np.int8)) + + def test_unsigned_judgment_codes_cast_when_in_range(self) -> None: + cast = _graph_output("j", "judgment", np.asarray([0, 1], dtype=np.uint8)) + assert cast.dtype == np.dtype(np.int64) and cast.tolist() == [0, 1] + + @pytest.mark.parametrize( + ("engine_dtype", "values", "expected"), + [ + ("integer", np.asarray([3, -4], dtype=np.int32), np.int64), + ("integer", np.asarray([3, 4], dtype=np.uint16), np.int64), + ("decimal", np.asarray([1.25, -0.5], dtype=np.float64), np.float64), + ("decimal", np.asarray([1.25, -0.5], dtype=np.float32), np.float64), + ("bool", np.asarray([True, False]), np.bool_), + ], + ) + def test_each_supported_dtype_casts_to_its_graph_dtype( + self, engine_dtype, values, expected + ) -> None: + cast = _graph_output("v", engine_dtype, values) + assert cast.dtype == np.dtype(expected) + np.testing.assert_array_equal(cast, values) + + @pytest.mark.parametrize( + ("engine_dtype", "values"), + [ + ("integer", np.asarray([1.5, 2.0])), + ("integer", np.asarray([True, False])), + ("decimal", np.asarray([1, 2], dtype=np.int64)), + ("decimal", np.asarray(["1.5"])), + ("bool", np.asarray([0, 1], dtype=np.int8)), + ("judgment", np.asarray([0.0, 1.0])), + ("integer", np.asarray([2**64 - 1], dtype=np.uint64)), + ], + ) + def test_a_lossy_or_mismatched_array_is_refused(self, engine_dtype, values) -> None: + with pytest.raises(ValueError, match="'v'"): + _graph_output("v", engine_dtype, values) + + @pytest.mark.parametrize("engine_dtype", ["text", "date", "unknown"]) + def test_text_and_date_have_no_graph_dtype(self, engine_dtype) -> None: + with pytest.raises(ValueError, match="no graph"): + _graph_output("v", engine_dtype, np.asarray(["a"])) + + +class TestGraphTypedMaterialize: + def test_output_dtypes_must_be_native_or_graph(self) -> None: + with pytest.raises(ValueError, match="output_dtypes"): + _nz_adapter(output_dtypes="pandas") + + def test_native_default_returns_the_engine_arrays_unchanged( + self, monkeypatch + ) -> None: + adapter = _nz_adapter() + _fake_engine(adapter, monkeypatch) + results = adapter.materialize( + _fake_frame(), ["family_assets_ok", "family_category"], 2026 + ) + assert results["family_assets_ok"].dtype == np.dtype(np.int8) + assert results["family_category"].tolist() == ["small"] * 3 + + def test_graph_mode_returns_graph_ownable_dtypes(self, monkeypatch) -> None: + adapter = _nz_adapter(output_dtypes="graph", periods={"2026-27": TAX_YEAR}) + _fake_engine(adapter, monkeypatch) + variables = [ + "person_benefit", + "family_assets_ok", + "family_assistance", + "family_band", + "family_is_large", + ] + results = adapter.materialize(_fake_frame(), variables, "2026-27") + assert {name: results[name].dtype.str for name in variables} == { + "person_benefit": np.dtype(np.float64).str, + "family_assets_ok": np.dtype(np.int64).str, + "family_assistance": np.dtype(np.float64).str, + "family_band": np.dtype(np.int64).str, + "family_is_large": np.dtype(np.bool_).str, + } + assert results["family_assets_ok"].tolist() == [1, -1, 0] + assert results["person_benefit"].shape == (5,) + + @pytest.mark.parametrize("variable", ["family_category", "family_review_date"]) + def test_graph_mode_refuses_text_and_date_before_the_engine_runs( + self, monkeypatch, variable + ) -> None: + adapter = _nz_adapter(output_dtypes="graph") + programs = _fake_engine(adapter, monkeypatch) + with pytest.raises( + ValueError, match=f"{variable}.*refused|refused.*{variable}" + ): + adapter.materialize(_fake_frame(), ["family_assistance", variable], 2026) + assert all(not program.calls for program in programs.values()) + + def test_graph_dtype_refuses_a_native_adapter(self, monkeypatch) -> None: + adapter = _nz_adapter() + _fake_engine(adapter, monkeypatch) + with pytest.raises(ValueError, match="output_dtypes='graph' adapters only"): + adapter.graph_dtype("family_assets_ok") + + def test_graph_dtype_names_the_owned_column_token(self, monkeypatch) -> None: + adapter = _nz_adapter(output_dtypes="graph") + _fake_engine(adapter, monkeypatch) + assert { + name: adapter.graph_dtype(name) + for name in ( + "person_benefit", + "family_assets_ok", + "family_band", + "family_is_large", + ) + } == { + "person_benefit": "float64", + "family_assets_ok": "int64", + "family_band": "int64", + "family_is_large": "bool", + } + with pytest.raises(ValueError, match="no graph column"): + adapter.graph_dtype("family_category") + with pytest.raises(ValueError, match="Unknown Axiom variable"): + adapter.graph_dtype("not_a_variable") + + +# ---------------------------------------------------------------------- +# Relations +# ---------------------------------------------------------------------- + + +class TestAssertNoRelations: + def test_a_relation_free_program_passes(self, monkeypatch) -> None: + adapter = _nz_adapter() + _fake_engine(adapter, monkeypatch) + assert assert_no_relations(adapter, "family") is None + + def test_relations_are_refused_naming_the_module_path(self, monkeypatch) -> None: + adapter = _nz_adapter() + _fake_engine( + adapter, monkeypatch, relations=(_Relation("member_of_household"),) + ) + with pytest.raises( + NotImplementedError, + match="zz/policies/tests/axiom_toy_country.yaml.*member_of_household", + ): + assert_no_relations(adapter, "family") + + def test_an_entity_without_derived_rules_is_named(self, monkeypatch) -> None: + adapter = _nz_adapter() + _fake_engine(adapter, monkeypatch) + with pytest.raises(ValueError, match="no derived rules.*Household"): + assert_no_relations(adapter, "household") + + def test_only_axiom_engines_are_accepted(self) -> None: + with pytest.raises(TypeError, match="AxiomEngine"): + assert_no_relations(object(), "family") # type: ignore[arg-type] + + +# ---------------------------------------------------------------------- +# Engine references +# ---------------------------------------------------------------------- + +_MODULE = "zz/policies/tests/toy.yaml" +_BASE_TREE = { + _MODULE: b"format: rulespec/v1\nrules: []\n", + "zz/policies/shared/rates.yaml": b"format: rulespec/v1\nrules:\n - name: r\n", + ".axiom/toolchain.toml": b'engine = "pinned"\n', + "README.md": b"RuleSpec fixture tree.\n", +} + + +def _write_tree(root: Path, files: Mapping[str, bytes]) -> Path: + for name, content in files.items(): + path = root / name + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(content) + return root + + +def _ref(root: Path, **engine_kwargs) -> str: + engine = AxiomEngine( + root / _MODULE, schema=NZ_SCHEMA, rulespec_roots=(root,), **engine_kwargs + ) + return axiom_engine_ref( + engine, + engine_commit=ENGINE_COMMIT, + wheel_sha256=WHEEL_SHA256, + rulespec_root=root, + rulespec_commit=RULESPEC_COMMIT, + ) + + +_NAME = st.text( + alphabet="abcdefghijklmnopqrstuvwxyz0123456789_-", min_size=1, max_size=6 +) + + +@st.composite +def rulespec_trees(draw: st.DrawFn) -> dict[str, bytes]: + """A random RuleSpec-like tree that always holds the module. + + Directory components carry a ``d_`` prefix and files a ``.yaml`` suffix, + so no path is both a file and a directory; names are lower case, so a + case-insensitive file system cannot fold two of them together. + """ + + paths = draw( + st.sets( + st.builds( + lambda directories, name: "/".join( + [*(f"d_{part}" for part in directories), f"{name}.yaml"] + ), + st.lists(_NAME, max_size=3), + _NAME, + ), + max_size=6, + ) + ) + files = {path: draw(st.binary(min_size=1, max_size=48)) for path in sorted(paths)} + files[_MODULE] = draw(st.binary(min_size=1, max_size=48)) + return files + + +class TestRulespecTreeDigest: + @settings(max_examples=60, deadline=None) + @given(files=rulespec_trees()) + def test_differential_digest_equals_the_graph_directory_identity( + self, tmp_path_factory, files + ) -> None: + """Two implementations of one formula: adapter and graph agree.""" + + root = _write_tree(tmp_path_factory.mktemp("tree"), files) + digest, size = rulespec_tree_digest(root) + assert (digest, size) == _directory_identity(root) + assert size == sum(len(content) for content in files.values()) + # The public source key wraps exactly this digest and size. + assert source_content_key("rulespec_nz", root) == sha256_domain( + "source", canonical_json(("rulespec_nz", digest, size)) + ) + + def test_a_file_is_not_a_root(self, tmp_path) -> None: + (tmp_path / "file.yaml").write_bytes(b"x") + with pytest.raises(ValueError, match="not a directory"): + rulespec_tree_digest(tmp_path / "file.yaml") + + +class TestAxiomEngineRef: + def test_the_reference_is_canonical_json_of_every_pin(self, tmp_path) -> None: + root = _write_tree(tmp_path / "rulespec-zz", _BASE_TREE) + reference = _ref( + root, periods={"2026-27": TAX_YEAR}, nesting=NZ_NESTING, arithmetic="f64" + ) + document = json.loads(reference) + assert reference == canonical_json(document).decode("utf-8") + digest, size = rulespec_tree_digest(root) + assert document == { + "format": "microcosm.frame.axiom-engine-ref/1", + "engine": "axiom-rules-engine", + "engine_commit": ENGINE_COMMIT, + "engine_wheel_sha256": WHEEL_SHA256, + "rulespec_commit": RULESPEC_COMMIT, + "rulespec_tree_sha256": digest, + "rulespec_tree_bytes": size, + "module": _MODULE, + "module_sha256": hashlib.sha256(_BASE_TREE[_MODULE]).hexdigest(), + "arithmetic": "f64", + "output_dtypes": "native", + "person_entity": "person", + "group_entities": ["household", "family"], + "entity_names": { + "family": "Family", + "household": "Household", + "person": "Person", + }, + "periods": { + "2026-27": { + "start": "2026-04-01", + "end": "2027-03-31", + "kind": "tax_year", + } + }, + "nesting": {"family": "household"}, + } + assert str(tmp_path) not in reference + + def test_determinism_relocation_and_timestamps_leave_it_unchanged( + self, tmp_path + ) -> None: + root = _write_tree(tmp_path / "one", _BASE_TREE) + first = _ref(root) + assert _ref(root) == first + relocated = tmp_path / "elsewhere" / "two" + shutil.copytree(root, relocated) + os.utime(relocated / _MODULE, (0, 0)) + assert _ref(relocated) == first + + @pytest.mark.parametrize( + "change", + [ + {"arithmetic": "f64"}, + {"output_dtypes": "graph"}, + {"periods": {"2026-27": TAX_YEAR}}, + {"nesting": NZ_NESTING}, + { + "entity_names": { + "person": "Person", + "household": "Household", + "family": "BenefitUnit", + } + }, + ], + ) + def test_every_configuration_field_moves_it(self, tmp_path, change) -> None: + root = _write_tree(tmp_path / "rulespec", _BASE_TREE) + assert _ref(root, **change) != _ref(root) + + @pytest.mark.parametrize( + ("field", "value"), + [ + ("engine_commit", "d" * 40), + ("wheel_sha256", "e" * 64), + ("rulespec_commit", "f" * 40), + ], + ) + def test_every_declared_pin_moves_it(self, tmp_path, field, value) -> None: + root = _write_tree(tmp_path / "rulespec", _BASE_TREE) + engine = AxiomEngine(root / _MODULE, rulespec_roots=(root,)) + pins = { + "engine_commit": ENGINE_COMMIT, + "wheel_sha256": WHEEL_SHA256, + "rulespec_commit": RULESPEC_COMMIT, + } + base = axiom_engine_ref(engine, rulespec_root=root, **pins) + moved = axiom_engine_ref(engine, rulespec_root=root, **{**pins, field: value}) + assert moved != base + + @settings(max_examples=60, deadline=None) + @given(files=rulespec_trees(), data=st.data()) + def test_property_it_moves_iff_a_pinned_byte_moves( + self, tmp_path_factory, files, data + ) -> None: + """Invariant: ``engine_ref`` changes iff a pinned byte changes. + + The same bytes in another place give the same reference; editing any + one byte of any file under the root gives a different one. + """ + + root = _write_tree(tmp_path_factory.mktemp("base"), files) + reference = _ref(root) + copy = tmp_path_factory.mktemp("copy") + shutil.copytree(root, copy, dirs_exist_ok=True) + assert _ref(copy) == reference + name = data.draw(st.sampled_from(sorted(files))) + content = files[name] + offset = data.draw(st.integers(min_value=0, max_value=len(content) - 1)) + replacement = data.draw( + st.integers(min_value=0, max_value=255).filter( + lambda value: value != content[offset] + ) + ) + edited = bytearray(content) + edited[offset] = replacement + (copy / name).write_bytes(bytes(edited)) + assert _ref(copy) != reference + + def test_adding_or_renaming_a_file_moves_it(self, tmp_path) -> None: + root = _write_tree(tmp_path / "rulespec", _BASE_TREE) + reference = _ref(root) + (root / "zz/policies/shared/new.yaml").write_bytes(b"x") + added = _ref(root) + assert added != reference + (root / "zz/policies/shared/new.yaml").rename( + root / "zz/policies/shared/renamed.yaml" + ) + assert _ref(root) not in {reference, added} + + @pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("engine_commit", "a" * 39, "engine_commit"), + ("engine_commit", "A" * 40, "engine_commit"), + ("engine_commit", "a" * 12, "engine_commit"), + ("rulespec_commit", "main", "rulespec_commit"), + ("rulespec_commit", None, "rulespec_commit"), + ("wheel_sha256", "b" * 63, "wheel_sha256"), + ("wheel_sha256", "B" * 64, "wheel_sha256"), + ], + ) + def test_malformed_pins_are_refused(self, tmp_path, field, value, message) -> None: + root = _write_tree(tmp_path / "rulespec", _BASE_TREE) + engine = AxiomEngine(root / _MODULE, rulespec_roots=(root,)) + pins = { + "engine_commit": ENGINE_COMMIT, + "wheel_sha256": WHEEL_SHA256, + "rulespec_commit": RULESPEC_COMMIT, + field: value, + } + with pytest.raises(ValueError, match=message): + axiom_engine_ref(engine, rulespec_root=root, **pins) + + def test_a_64_hex_commit_is_accepted(self, tmp_path) -> None: + root = _write_tree(tmp_path / "rulespec", _BASE_TREE) + engine = AxiomEngine(root / _MODULE, rulespec_roots=(root,)) + assert axiom_engine_ref( + engine, + engine_commit="a" * 64, + wheel_sha256=WHEEL_SHA256, + rulespec_root=root, + rulespec_commit="c" * 64, + ) + + def test_the_root_must_be_the_adapters_only_root(self, tmp_path) -> None: + root = _write_tree(tmp_path / "rulespec", _BASE_TREE) + other = _write_tree(tmp_path / "other", {"x.yaml": b"x"}) + pins = { + "engine_commit": ENGINE_COMMIT, + "wheel_sha256": WHEEL_SHA256, + "rulespec_commit": RULESPEC_COMMIT, + } + two_roots = AxiomEngine(root / _MODULE, rulespec_roots=(root, other)) + with pytest.raises(ValueError, match="only root"): + axiom_engine_ref(two_roots, rulespec_root=root, **pins) + one_root = AxiomEngine(root / _MODULE, rulespec_roots=(root,)) + with pytest.raises(ValueError, match="only root"): + axiom_engine_ref(one_root, rulespec_root=other, **pins) + + def test_the_module_must_be_a_file_under_the_root(self, tmp_path) -> None: + root = _write_tree(tmp_path / "rulespec", _BASE_TREE) + outside = _write_tree(tmp_path / "outside", {"m.yaml": b"x"}) / "m.yaml" + pins = { + "engine_commit": ENGINE_COMMIT, + "wheel_sha256": WHEEL_SHA256, + "rulespec_commit": RULESPEC_COMMIT, + } + for module in (outside, root / "zz/missing.yaml"): + engine = AxiomEngine(module, rulespec_roots=(root,)) + with pytest.raises(ValueError, match="not a file under"): + axiom_engine_ref(engine, rulespec_root=root, **pins) + + def test_a_symbolic_link_under_the_root_is_refused(self, tmp_path) -> None: + root = _write_tree(tmp_path / "rulespec", _BASE_TREE) + target = _write_tree(tmp_path / "linked", {"hidden.yaml": b"x"}) + (root / "zz/policies/linked").symlink_to(target, target_is_directory=True) + with pytest.raises(ValueError, match="symbolic link zz/policies/linked"): + _ref(root) + + def test_only_axiom_engines_are_accepted(self, tmp_path) -> None: + with pytest.raises(TypeError, match="AxiomEngine"): + axiom_engine_ref( + object(), # type: ignore[arg-type] + engine_commit=ENGINE_COMMIT, + wheel_sha256=WHEEL_SHA256, + rulespec_root=tmp_path, + rulespec_commit=RULESPEC_COMMIT, + ) + + +class TestEngineRefPin: + """The adapter computes only from the bytes its reference names.""" + + def _adapter(self, root: Path) -> AxiomEngine: + return AxiomEngine( + root / _MODULE, + schema=NZ_SCHEMA, + rulespec_roots=(root,), + output_dtypes="graph", + ) + + def _reference(self, adapter: AxiomEngine, root: Path) -> str: + return axiom_engine_ref( + adapter, + engine_commit=ENGINE_COMMIT, + wheel_sha256=WHEEL_SHA256, + rulespec_root=root, + rulespec_commit=RULESPEC_COMMIT, + ) + + def test_a_reference_then_unchanged_compiles_succeed( + self, tmp_path, monkeypatch + ) -> None: + root = _write_tree(tmp_path / "rulespec", _BASE_TREE) + adapter = self._adapter(root) + _fake_engine(adapter, monkeypatch) + reference = self._reference(adapter, root) + adapter.materialize(_fake_frame(), ["family_assistance"], 2026) + adapter.materialize(_fake_frame(), ["person_benefit"], 2026) + assert self._reference(adapter, root) == reference + + def test_an_edit_after_the_reference_blocks_the_next_compile( + self, tmp_path, monkeypatch + ) -> None: + root = _write_tree(tmp_path / "rulespec", _BASE_TREE) + adapter = self._adapter(root) + programs = _fake_engine(adapter, monkeypatch) + self._reference(adapter, root) + (root / "zz/policies/shared/rates.yaml").write_bytes(b"edited") + with pytest.raises(ValueError, match="changed after this adapter"): + adapter.materialize(_fake_frame(), ["family_assistance"], 2026) + assert programs == {} + + def test_an_edit_after_the_reference_blocks_a_second_reference( + self, tmp_path, monkeypatch + ) -> None: + root = _write_tree(tmp_path / "rulespec", _BASE_TREE) + adapter = self._adapter(root) + _fake_engine(adapter, monkeypatch) + self._reference(adapter, root) + adapter.materialize(_fake_frame(), ["family_assistance"], 2026) + (root / _MODULE).write_bytes(b"edited") + with pytest.raises(ValueError, match="changed after this adapter"): + self._reference(adapter, root) + + def test_a_reference_after_compiling_is_refused( + self, tmp_path, monkeypatch + ) -> None: + root = _write_tree(tmp_path / "rulespec", _BASE_TREE) + adapter = self._adapter(root) + _fake_engine(adapter, monkeypatch) + adapter.materialize(_fake_frame(), ["family_assistance"], 2026) + with pytest.raises(ValueError, match="before the adapter compiles"): + self._reference(adapter, root) + + def test_an_adapter_without_a_reference_never_hashes_its_root( + self, monkeypatch + ) -> None: + # The Belgian path is unchanged: no reference, no digest on compile. + adapter = _nz_adapter() + _fake_engine(adapter, monkeypatch) + + def refuse(root): + raise AssertionError("an unreferenced adapter hashed its root") + + monkeypatch.setattr( + "microcosm.frame.adapters.axiom.rulespec_tree_digest", refuse + ) + adapter.materialize(_fake_frame(), ["family_assistance"], 2026) + + +class TestAxiomEngineRefGitCheckout: + @pytest.fixture + def checkout(self, tmp_path) -> tuple[Path, str]: + root = _write_tree( + tmp_path / "rulespec-git", + {**_BASE_TREE, ".gitignore": b"*.local.yaml\n"}, + ) + _git(root, "init", "-q") + _git(root, "add", "-A") + _git(root, "commit", "-q", "-m", "fixture") + return root, _git(root, "rev-parse", "HEAD") + + def _checkout_ref(self, root: Path, commit: str) -> str: + engine = AxiomEngine(root / _MODULE, rulespec_roots=(root,)) + return axiom_engine_ref( + engine, + engine_commit=ENGINE_COMMIT, + wheel_sha256=WHEEL_SHA256, + rulespec_root=root, + rulespec_commit=commit, + ) + + def test_a_clean_checkout_at_the_declared_commit_is_accepted( + self, checkout + ) -> None: + root, head = checkout + first = self._checkout_ref(root, head) + # The status check takes no optional locks, so it cannot rewrite the + # index and move the digest between two calls. + assert self._checkout_ref(root, head) == first + assert json.loads(first)["rulespec_commit"] == head + + def test_a_modified_file_is_refused(self, checkout) -> None: + root, head = checkout + (root / _MODULE).write_bytes(b"format: rulespec/v1\nrules: [edited]\n") + with pytest.raises(ValueError, match="modified, untracked, or ignored"): + self._checkout_ref(root, head) + + def test_an_untracked_file_is_refused(self, checkout) -> None: + root, head = checkout + (root / "zz/untracked.yaml").write_bytes(b"x") + with pytest.raises(ValueError, match="modified, untracked, or ignored"): + self._checkout_ref(root, head) + + def test_an_ignored_file_is_refused(self, checkout) -> None: + # The digest would hash it, but the declared commit does not hold it. + root, head = checkout + (root / "zz/override.local.yaml").write_bytes(b"x") + assert _git(root, "status", "--porcelain") == "" + with pytest.raises(ValueError, match="ignored"): + self._checkout_ref(root, head) + + def test_a_checkout_reference_is_clone_specific_and_an_export_is_not( + self, checkout, tmp_path + ) -> None: + # Documented: a checkout's .git is hashed, so its reference names that + # clone; exports of the same commit agree wherever they land. + root, head = checkout + archive = tmp_path / "export.tar" + _git(root, "archive", "--format=tar", "-o", str(archive), "HEAD") + exports = [] + for name in ("export-one", "export-two"): + destination = tmp_path / name + with tarfile.open(archive) as bundle: + bundle.extractall(destination, filter="data") + exports.append(self._checkout_ref(destination, head)) + assert exports[0] == exports[1] + assert self._checkout_ref(root, head) != exports[0] + + def test_a_checkout_at_another_commit_is_refused(self, checkout) -> None: + root, head = checkout + other = "0" * 40 if head != "0" * 40 else "1" * 40 + with pytest.raises(ValueError, match="not the declared rulespec_commit"): + self._checkout_ref(root, other) + + +# ---------------------------------------------------------------------- +# Real engine (skipped without axiom_rules_engine) +# ---------------------------------------------------------------------- + + +_FIXTURE_COMMIT = "0" * 40 + + +def _fixture_ref(engine: AxiomEngine) -> str: + """A reference for the in-repo fixture; the commit is a test placeholder.""" + + return axiom_engine_ref( + engine, + engine_commit=_FIXTURE_COMMIT, + wheel_sha256=WHEEL_SHA256, + rulespec_root=FIXTURE_RULESPEC_ROOT, + rulespec_commit=_FIXTURE_COMMIT, + ) + + +class _ToyPopulation(KernelBase): + """CREATE kernel loading the toy population from a JSON source.""" + + ref = "test.toy_population@1" + capabilities = Capabilities( + determinism=Determinism.DETERMINISTIC, + numeric=Numeric.BITWISE, + seed_source=SeedSource.NONE, + structural=StructuralDelta.CREATE, + ) + + def run(self, context: KernelContext) -> KernelResult: + document = json.loads(context.sources["population"].read_text()) + schema = NZ_SCHEMA if "family" in document else BE_SCHEMA + tables = { + entity: pd.DataFrame( + { + column: np.asarray(values, dtype=document["dtypes"][column]) + for column, values in document[entity].items() + } + ) + for entity in schema.entities + } + frame = Frame( + tables, + schema, + { + "household": Weights( + np.asarray(document["household_weights"], dtype=np.float64), + WeightKind.DESIGN, + ) + }, + ) + return KernelResult(frame=frame, receipt={"persons": len(tables["person"])}) + + +_TOY_COUNTRY_DOCUMENT = { + "person": { + "person_id": [1, 2, 3], + "person_household_id": [1, 1, 2], + "toy_taxable_income": [5_000.0, 10_000.0, 20_000.0], + "toy_is_exempt": [False, False, False], + "toy_child_count": [0, 2, 1], + }, + "household": {"household_id": [1, 2], "toy_household_rent": [7_200.0, 4_800.0]}, + "household_weights": [1500.0, 900.0], + "dtypes": { + "person_id": "int64", + "person_household_id": "int64", + "toy_taxable_income": "float64", + "toy_is_exempt": "bool", + "toy_child_count": "int64", + "household_id": "int64", + "toy_household_rent": "float64", + }, +} + +_TOY_FAMILY_DOCUMENT = { + "person": { + "person_id": [1, 2, 3, 4, 5], + "person_household_id": [1, 1, 2, 2, 2], + "person_family_id": [10, 10, 20, 21, 21], + "toy_taxable_income": [5_000.0, 10_000.0, 20_000.0, 0.0, 0.0], + "toy_is_exempt": [False, False, False, False, False], + "toy_child_count": [0, 0, 0, 1, 0], + }, + "household": {"household_id": [1, 2]}, + "family": { + "family_id": [10, 20, 21], + "toy_family_weekly_rent": [300.0, 80.0, 180.0], + "toy_family_cash_assets": [9_000.0, 1_000.0, 8_000.0], + "toy_family_size": [2, 1, 3], + }, + "household_weights": [1200.0, 800.0], + "dtypes": { + "person_id": "int64", + "person_household_id": "int64", + "person_family_id": "int64", + "toy_taxable_income": "float64", + "toy_is_exempt": "bool", + "toy_child_count": "int64", + "household_id": "int64", + "family_id": "int64", + "toy_family_weekly_rent": "float64", + "toy_family_cash_assets": "float64", + "toy_family_size": "int64", + }, +} + + +def _source_node(document: Mapping[str, object]) -> Node: + outputs = tuple( + Owned(entity, column, document["dtypes"][column]) + for entity in ("person", "household", "family") + if entity in document + for column in document[entity] + if not column.endswith("_id") + ) + return Node( + "population", + _ToyPopulation.ref, + outputs=outputs, + structural=StructuralDelta.CREATE, + sources=("population",), + ) + + +def _run( + tmp_path: Path, + document: Mapping[str, object], + nodes: tuple[Node, ...], + kernels: tuple[object, ...], +): + tmp_path.mkdir(parents=True, exist_ok=True) + source = tmp_path / "population.json" + source.write_text(json.dumps(document)) + registry = KernelRegistry() + registry.register(_ToyPopulation()) + for kernel in kernels: + registry.register(kernel) + graph = Graph( + "axiom-graph-adapter", + (SourceRef("population", "raw-bytes-v1"),), + (_source_node(document), *nodes), + ) + store = ContentStore(tmp_path / "store") + manifest = run_graph( + compile_graph(graph), + sources={"population": source}, + store=store, + kernels=registry, + resume="forbid", + decisions=(), + ) + return manifest, store + + +def _column(manifest, store, node_id: str, entity: str, column: str) -> pd.Series: + return store.load_column(manifest.nodes[node_id].artifacts[(entity, column)]) + + +@needs_engine +class TestRealEngineOnTheGraph: + def test_graph_typed_simulate_rules_runs_in_a_filter_free_graph( + self, tmp_path + ) -> None: + engine = AxiomEngine( + FIXTURE_MODULE, rulespec_roots=FIXTURE_RULESPEC_ROOTS, output_dtypes="graph" + ) + engine_ref = _fixture_ref(engine) + variables = ("toy_income_tax", "toy_housing_allowance") + node = Node( + "toy_rules", + SimulateRulesKernel.ref, + inputs=( + Slice( + "person", + ("toy_taxable_income", "toy_is_exempt", "toy_child_count"), + ), + Slice("household", ("toy_household_rent",)), + ), + outputs=tuple( + Owned( + engine.variable_metadata(name).entity, + name, + engine.graph_dtype(name), + ) + for name in variables + ), + params={"engine_ref": engine_ref, "variables": variables, "period": 2025}, + ) + manifest, store = _run( + tmp_path, + _TOY_COUNTRY_DOCUMENT, + (node,), + (SimulateRulesKernel(engine_ref, engine),), + ) + tax = _column(manifest, store, "toy_rules", "person", "toy_income_tax") + allowance = _column( + manifest, store, "toy_rules", "household", "toy_housing_allowance" + ) + assert tax.dtype == np.dtype(np.float64) + np.testing.assert_allclose(tax.to_numpy(), [500.0, 1_000.0, 3_500.0]) + np.testing.assert_allclose(allowance.to_numpy(), [1_200.0, 0.0]) + + def test_the_tax_year_label_reaches_the_engine_as_explicit_bounds( + self, monkeypatch + ) -> None: + adapter = AxiomEngine( + FIXTURE_MODULE, + rulespec_roots=FIXTURE_RULESPEC_ROOTS, + periods={"2026-27": TAX_YEAR}, + output_dtypes="graph", + ) + program = adapter._program("person") + calls: list[tuple[str, str, str]] = [] + original = program.execute + + def recording_execute(*, period_kind, start, end, **kwargs): + calls.append((start, end, period_kind)) + return original(period_kind=period_kind, start=start, end=end, **kwargs) + + monkeypatch.setattr(program, "execute", recording_execute) + frame = Frame( + { + "person": pd.DataFrame( + { + "person_id": [1, 2], + "person_household_id": [1, 2], + "toy_taxable_income": [5_000.0, 20_000.0], + "toy_is_exempt": [False, False], + "toy_child_count": [0, 1], + } + ), + "household": pd.DataFrame({"household_id": [1, 2]}), + }, + BE_SCHEMA, + {"household": Weights(np.asarray([1.0, 1.0]), WeightKind.DESIGN)}, + ) + results = adapter.materialize(frame, ["toy_income_tax"], "2026-27") + assert calls == [("2026-04-01", "2027-03-31", "tax_year")] + np.testing.assert_allclose(results["toy_income_tax"], [500.0, 3_500.0]) + with pytest.raises(ValueError, match="No explicit Axiom period bounds"): + adapter.materialize(frame, ["toy_income_tax"], 2026) + + def test_family_judgments_are_lossless_graph_codes(self) -> None: + adapter = AxiomEngine( + FAMILY_FIXTURE_MODULE, + schema=NZ_SCHEMA, + rulespec_roots=FIXTURE_RULESPEC_ROOTS, + nesting=NZ_NESTING, + ) + typed = AxiomEngine( + FAMILY_FIXTURE_MODULE, + schema=NZ_SCHEMA, + rulespec_roots=FIXTURE_RULESPEC_ROOTS, + nesting=NZ_NESTING, + output_dtypes="graph", + ) + assert_no_relations(adapter, "family") + frame = _nz_frame( + [1, 1, 2, 2, 2], + [10, 10, 20, 21, 21], + [1200.0, 800.0], + family_columns={ + "toy_family_weekly_rent": [300.0, 80.0, 180.0], + "toy_family_cash_assets": [9_000.0, 1_000.0, 8_000.0], + "toy_family_size": np.asarray([2, 1, 3], dtype=np.int64), + }, + ) + variables = [ + "toy_family_assets_within_limit", + "toy_family_rent_assistance", + "toy_family_size_band", + "toy_family_is_large", + ] + native = adapter.materialize(frame, variables, 2026) + graph = typed.materialize(frame, variables, 2026) + judgment = "toy_family_assets_within_limit" + assert native[judgment].dtype == np.dtype(np.int8) + assert graph[judgment].dtype == np.dtype(np.int64) + assert graph[judgment].tolist() == [-1, 1, 1] + assert graph[judgment].astype(np.int8).tobytes() == native[judgment].tobytes() + np.testing.assert_allclose( + graph["toy_family_rent_assistance"], [100.0, 0.0, 40.0] + ) + assert graph["toy_family_size_band"].dtype == np.dtype(np.int64) + assert graph["toy_family_size_band"].tolist() == [1, 1, 2] + assert graph["toy_family_is_large"].dtype == np.dtype(np.bool_) + assert graph["toy_family_is_large"].tolist() == [False, False, True] + assert {name: typed.graph_dtype(name) for name in variables} == { + "toy_family_assets_within_limit": "int64", + "toy_family_rent_assistance": "float64", + "toy_family_size_band": "int64", + "toy_family_is_large": "bool", + } + + def test_rules_by_ref_runs_two_axiom_engines_in_one_graph(self, tmp_path) -> None: + person_engine = AxiomEngine( + FIXTURE_MODULE, + schema=NZ_SCHEMA, + rulespec_roots=FIXTURE_RULESPEC_ROOTS, + nesting=NZ_NESTING, + output_dtypes="graph", + periods={"2026-27": TAX_YEAR}, + ) + family_engine = AxiomEngine( + FAMILY_FIXTURE_MODULE, + schema=NZ_SCHEMA, + rulespec_roots=FIXTURE_RULESPEC_ROOTS, + nesting=NZ_NESTING, + output_dtypes="graph", + periods={"2026-27": TAX_YEAR}, + ) + person_ref = _fixture_ref(person_engine) + family_ref = _fixture_ref(family_engine) + assert person_ref != family_ref + person_variables = ("toy_income_tax",) + family_variables = ( + "toy_family_assets_within_limit", + "toy_family_rent_assistance", + ) + + def node(node_id, kernel_ref, engine, engine_ref, variables, inputs): + return Node( + node_id, + kernel_ref, + inputs=inputs, + outputs=tuple( + Owned( + engine.variable_metadata(name).entity, + name, + engine.graph_dtype(name), + ) + for name in variables + ), + params={ + "engine_ref": engine_ref, + "variables": variables, + "period": "2026-27", + }, + ) + + person_inputs = ( + Slice( + "person", + ("toy_taxable_income", "toy_is_exempt", "toy_child_count"), + ), + ) + # A family-only node still slices one person data column: the + # kernel rebuilds group tables from the person table it is given. + family_inputs = ( + Slice( + "family", + ( + "toy_family_weekly_rent", + "toy_family_cash_assets", + "toy_family_size", + ), + ), + Slice("person", ("toy_taxable_income",)), + ) + by_ref = SimulateRulesByRefKernel( + {person_ref: person_engine, family_ref: family_engine} + ) + manifest, store = _run( + tmp_path / "by-ref", + _TOY_FAMILY_DOCUMENT, + ( + node( + "rules.person", + by_ref.ref, + person_engine, + person_ref, + person_variables, + person_inputs, + ), + node( + "rules.family", + by_ref.ref, + family_engine, + family_ref, + family_variables, + family_inputs, + ), + ), + (by_ref,), + ) + np.testing.assert_allclose( + _column( + manifest, store, "rules.person", "person", "toy_income_tax" + ).to_numpy(), + [500.0, 1_000.0, 3_500.0, 0.0, 0.0], + ) + codes = _column( + manifest, store, "rules.family", "family", "toy_family_assets_within_limit" + ) + assert codes.dtype == np.dtype(np.int64) + assert codes.tolist() == [-1, 1, 1] + + # Differential: each node equals a single-engine simulate.rules@1 run. + for node_id, engine, engine_ref, variables, inputs in ( + ( + "rules.person", + person_engine, + person_ref, + person_variables, + person_inputs, + ), + ( + "rules.family", + family_engine, + family_ref, + family_variables, + family_inputs, + ), + ): + single, single_store = _run( + tmp_path / f"single-{node_id}", + _TOY_FAMILY_DOCUMENT, + ( + node( + node_id, + SimulateRulesKernel.ref, + engine, + engine_ref, + variables, + inputs, + ), + ), + (SimulateRulesKernel(engine_ref, engine),), + ) + assert single.nodes[node_id].receipt == manifest.nodes[node_id].receipt + for name in variables: + entity = engine.variable_metadata(name).entity + expected = _column(single, single_store, node_id, entity, name) + actual = _column(manifest, store, node_id, entity, name) + pd.testing.assert_series_equal(actual, expected) + assert actual.to_numpy().tobytes() == expected.to_numpy().tobytes() diff --git a/packages/microcosm-frame/tests/engine_free/shared/test_rules_kernels.py b/packages/microcosm-frame/tests/engine_free/shared/test_rules_kernels.py new file mode 100644 index 0000000000..3aea3324e5 --- /dev/null +++ b/packages/microcosm-frame/tests/engine_free/shared/test_rules_kernels.py @@ -0,0 +1,665 @@ +"""``simulate.rules_by_ref@1``: one kernel, several bound rules engines. + +The kernel must behave, node for node, exactly like ``simulate.rules@1`` +bound to the engine the node names. These tests use two pure-Python engines +on the New Zealand schema, one with person outputs and one with family +outputs, and check that property three ways: directly on kernel contexts +(Hypothesis over random populations), through one ``run_graph`` that reaches +both engines, and against two single-engine ``simulate.rules@1`` graph runs. +""" + +import json +from collections.abc import Mapping, Sequence +from pathlib import Path + +import numpy as np +import pandas as pd +import pytest +from hypothesis import given, settings +from hypothesis import strategies as st + +import microcosm.frame.bundle as frame_bundle_module +import microcosm.frame.rules as frame_rules_module +import microcosm.frame.schema as frame_schema_module +from microcosm.frame import ( + EntitySchema, + ExportContract, + Frame, + VariableMetadata, + WeightKind, + Weights, +) +from microcosm.frame.adapters.axiom import NZ_SCHEMA +from microcosm.frame.kernels import SimulateRulesKernel +from microcosm.frame.rules_kernels import SimulateRulesByRefKernel +from microcosm.graph import ( + Capabilities, + ContentStore, + Determinism, + Graph, + Kernel, + KernelBase, + KernelContext, + KernelRegistry, + KernelResult, + Node, + Numeric, + Owned, + SeedSource, + Slice, + SourceRef, + StructuralDelta, + compile_graph, + run_graph, + source_hash, +) +from test_support.microcosm_frame.kernels import _context, _StubRulesEngine + +PERSON_REF = "toy-person-engine" +FAMILY_REF = "toy-family-engine" +PERSON_VARIABLES = ("person_net_income", "person_is_adult") +FAMILY_VARIABLES = ("family_support", "family_code") + + +class _PersonEngine: + """Pure-Python engine with two person outputs.""" + + def variable_metadata(self, name: str) -> VariableMetadata: + metadata = { + "person_net_income": VariableMetadata( + name="person_net_income", entity="person", dtype="float", period="year" + ), + "person_is_adult": VariableMetadata( + name="person_is_adult", entity="person", dtype="bool", period="year" + ), + } + try: + return metadata[name] + except KeyError as error: + raise ValueError(f"Unknown person variable {name!r}.") from error + + def variables(self) -> Sequence[str]: + return ("earnings", "age") + + def entity_schema(self) -> EntitySchema: + return NZ_SCHEMA + + def materialize( + self, bundle: Frame, variables: Sequence[str], period: int | str + ) -> Mapping[str, np.ndarray]: + year = int(str(period)[:4]) + person = bundle.table("person") + earnings = person["earnings"].to_numpy(dtype=np.float64) + age = person["age"].to_numpy(dtype=np.int64) + available = { + "person_net_income": earnings * np.float64(0.8) + np.float64(year - 2025), + "person_is_adult": age >= 18, + } + return {name: available[name] for name in variables} + + def export_contract(self) -> ExportContract: + return ExportContract.empty() + + def write_dataset(self, bundle: Frame, path: str | Path, period: int | str): + return None + + +class _FamilyEngine: + """Pure-Python engine with two family outputs, one tri-state code.""" + + def variable_metadata(self, name: str) -> VariableMetadata: + metadata = { + "family_support": VariableMetadata( + name="family_support", entity="family", dtype="float", period="year" + ), + "family_code": VariableMetadata( + name="family_code", entity="family", dtype="int", period="year" + ), + } + try: + return metadata[name] + except KeyError as error: + raise ValueError(f"Unknown family variable {name!r}.") from error + + def variables(self) -> Sequence[str]: + return ("family_rent",) + + def entity_schema(self) -> EntitySchema: + return NZ_SCHEMA + + def materialize( + self, bundle: Frame, variables: Sequence[str], period: int | str + ) -> Mapping[str, np.ndarray]: + rent = bundle.table("family")["family_rent"].to_numpy(dtype=np.float64) + available = { + "family_support": np.maximum( + (rent - np.float64(100.0)) * np.float64(0.7), np.float64(0.0) + ), + "family_code": np.sign(rent - np.float64(250.0)).astype(np.int64), + } + return {name: available[name] for name in variables} + + def export_contract(self) -> ExportContract: + return ExportContract.empty() + + def write_dataset(self, bundle: Frame, path: str | Path, period: int | str): + return None + + +def _person_node(kernel_ref: str, engine_ref: str = PERSON_REF) -> Node: + return Node( + "rules.person", + kernel_ref, + inputs=(Slice("person", ("earnings", "age")),), + outputs=( + Owned("person", "person_net_income", "float64"), + Owned("person", "person_is_adult", "bool"), + ), + params={ + "engine_ref": engine_ref, + "variables": PERSON_VARIABLES, + "period": 2026, + }, + ) + + +def _family_node(kernel_ref: str, engine_ref: str = FAMILY_REF) -> Node: + # A family-only engine node also slices one person data column: the + # kernel context holds only declared entities, and the kernel rebuilds + # group tables from the person memberships. + return Node( + "rules.family", + kernel_ref, + inputs=(Slice("family", ("family_rent",)), Slice("person", ("age",))), + outputs=( + Owned("family", "family_support", "float64"), + Owned("family", "family_code", "int64"), + ), + params={ + "engine_ref": engine_ref, + "variables": FAMILY_VARIABLES, + "period": "2026-27", + }, + ) + + +def _by_ref(**extra) -> SimulateRulesByRefKernel: + return SimulateRulesByRefKernel( + {PERSON_REF: _PersonEngine(), FAMILY_REF: _FamilyEngine(), **extra} + ) + + +@st.composite +def nz_frames(draw: st.DrawFn) -> Frame: + """A nested NZ frame with person and family inputs and household weights.""" + + person_households: list[int] = [] + person_families: list[int] = [] + rents: list[float] = [] + family = 100 + household_count = draw(st.integers(min_value=1, max_value=5)) + for household in range(1, household_count + 1): + for _ in range(draw(st.integers(min_value=1, max_value=3))): + family += 1 + rents.append( + draw(st.floats(min_value=0.0, max_value=900.0, allow_nan=False)) + ) + for _ in range(draw(st.integers(min_value=1, max_value=3))): + person_households.append(household) + person_families.append(family) + count = len(person_households) + person = pd.DataFrame( + { + "person_id": np.arange(1, count + 1, dtype=np.int64) * 7, + "person_household_id": np.asarray(person_households, dtype=np.int64), + "person_family_id": np.asarray(person_families, dtype=np.int64), + "earnings": np.asarray( + draw( + st.lists( + st.floats(min_value=0.0, max_value=2e5, allow_nan=False), + min_size=count, + max_size=count, + ) + ), + dtype=np.float64, + ), + "age": np.asarray( + draw( + st.lists( + st.integers(min_value=0, max_value=99), + min_size=count, + max_size=count, + ) + ), + dtype=np.int64, + ), + } + ) + household = pd.DataFrame( + {"household_id": np.arange(1, household_count + 1, dtype=np.int64)} + ) + families = pd.DataFrame( + { + "family_id": np.arange(101, family + 1, dtype=np.int64), + "family_rent": np.asarray(rents, dtype=np.float64), + } + ) + weights = draw( + st.lists( + st.floats(min_value=1.0, max_value=5e3, allow_nan=False), + min_size=household_count, + max_size=household_count, + ) + ) + return Frame( + {"person": person, "household": household, "family": families}, + NZ_SCHEMA, + {"household": Weights(np.asarray(weights), WeightKind.DESIGN)}, + ) + + +def _assert_same_result(actual: KernelResult, expected: KernelResult) -> None: + assert set(actual.columns) == set(expected.columns) + for coordinate, series in expected.columns.items(): + observed = actual.columns[coordinate] + assert observed.dtype == series.dtype + assert observed.index.equals(series.index) + assert observed.index.name == series.index.name + assert observed.to_numpy().tobytes() == series.to_numpy().tobytes() + assert dict(actual.receipt) == dict(expected.receipt) + + +# ---------------------------------------------------------------------- +# Construction and identity +# ---------------------------------------------------------------------- + + +class TestConstruction: + def test_satisfies_the_kernel_protocol_with_simulate_rules_capabilities( + self, + ) -> None: + kernel = _by_ref() + assert isinstance(kernel, Kernel) + assert kernel.ref == "simulate.rules_by_ref@1" + assert ( + kernel.capabilities + == SimulateRulesKernel(PERSON_REF, _PersonEngine()).capabilities + ) + assert kernel.capabilities.determinism is Determinism.DETERMINISTIC + assert kernel.capabilities.numeric is Numeric.BITWISE + assert kernel.capabilities.structural is StructuralDelta.NONE + assert kernel.engine_refs() == (FAMILY_REF, PERSON_REF) + + def test_dependencies_reach_the_capabilities(self) -> None: + kernel = SimulateRulesByRefKernel( + {PERSON_REF: _PersonEngine()}, dependencies=("numpy",) + ) + assert kernel.capabilities.dependencies == ("numpy",) + + @pytest.mark.parametrize( + ("engines", "error", "message"), + [ + ([(PERSON_REF, _PersonEngine())], TypeError, "mapping"), + ({}, ValueError, "at least one"), + ({"": _PersonEngine()}, ValueError, "engine_ref"), + ({PERSON_REF: object()}, TypeError, "RulesEngine"), + ], + ) + def test_invalid_bindings_are_refused(self, engines, error, message) -> None: + with pytest.raises(error, match=message): + SimulateRulesByRefKernel(engines) + + def test_implementation_hash_binds_both_kernels_and_each_adapter_class( + self, + ) -> None: + kernel = _by_ref() + assert kernel.implementation_hash() == source_hash( + SimulateRulesByRefKernel, + SimulateRulesKernel, + _FamilyEngine, + _PersonEngine, + frame_bundle_module, + frame_rules_module, + frame_schema_module, + ) + assert ( + kernel.implementation_hash() + != SimulateRulesKernel(PERSON_REF, _PersonEngine()).implementation_hash() + ) + + def test_adding_a_binding_of_a_bound_class_moves_no_hash(self) -> None: + base = _by_ref() + extended = _by_ref(**{"another-person-engine": _PersonEngine()}) + reordered = SimulateRulesByRefKernel( + {FAMILY_REF: _FamilyEngine(), PERSON_REF: _PersonEngine()} + ) + assert extended.implementation_hash() == base.implementation_hash() + assert reordered.implementation_hash() == base.implementation_hash() + assert extended.capabilities == base.capabilities + + def test_binding_a_new_adapter_class_moves_the_hash(self) -> None: + # Documented: a new class's source joins the hash. + assert ( + _by_ref(stub=_StubRulesEngine()).implementation_hash() + != _by_ref().implementation_hash() + ) + + +# ---------------------------------------------------------------------- +# Routing on kernel contexts +# ---------------------------------------------------------------------- + + +class TestRouting: + def test_an_unknown_engine_ref_is_refused(self) -> None: + frame = _population_frame() + node = _person_node(SimulateRulesByRefKernel.ref, engine_ref="unbound") + context = _context(frame, node, weighted_entities=("household",)) + with pytest.raises( + ValueError, + match=r"no engine bound to engine_ref sha256:[0-9a-f]{16} \('unbound'\); " + r"bound: sha256:", + ): + _by_ref().run(context) + + def test_an_unknown_axiom_style_ref_is_named_by_digest_and_module(self) -> None: + # Canonical JSON references share their leading keys, so the message + # names a digest and the module rather than a truncated prefix. + bound = json.dumps({"engine_commit": "a" * 40, "module": "nz/as/core.yaml"}) + unbound = json.dumps({"engine_commit": "a" * 40, "module": "nz/mb/rates.yaml"}) + frame = _population_frame() + node = _person_node(SimulateRulesByRefKernel.ref, engine_ref=unbound) + context = _context(frame, node, weighted_entities=("household",)) + with pytest.raises( + ValueError, + match=r"engine_ref sha256:[0-9a-f]{16} \(module 'nz/mb/rates\.yaml'\); " + r"bound: sha256:[0-9a-f]{16} \(module 'nz/as/core\.yaml'\)", + ): + SimulateRulesByRefKernel({bound: _PersonEngine()}).run(context) + + @pytest.mark.parametrize( + ("change", "error"), + [ + ({"drop": "engine_ref"}, ValueError), + ({"drop": "period"}, ValueError), + ({"set": ("extra", 1)}, ValueError), + ({"set": ("engine_ref", "")}, TypeError), + ({"set": ("engine_ref", 7)}, TypeError), + ({"set": ("variables", ["person_net_income"])}, TypeError), + ({"set": ("period", True)}, TypeError), + ], + ) + def test_differential_malformed_params_fail_as_under_simulate_rules( + self, change, error + ) -> None: + """A malformed node fails with the same exception type either way.""" + + frame = _population_frame() + context = _context( + frame, + _person_node(SimulateRulesByRefKernel.ref), + weighted_entities=("household",), + ) + params = dict(context.params) + if "drop" in change: + params.pop(change["drop"]) + else: + key, value = change["set"] + params[key] = value + context = KernelContext( + node=context.node, + tables=context.tables, + weights=context.weights, + strata=context.strata, + params=params, + rng=context.rng, + ) + with pytest.raises(error): + SimulateRulesKernel(PERSON_REF, _PersonEngine()).run(context) + with pytest.raises(error): + _by_ref().run(context) + + def test_the_delegate_parameter_contract_still_applies(self) -> None: + frame = _population_frame() + node = _person_node(SimulateRulesByRefKernel.ref) + context = _context(frame, node, weighted_entities=("household",)) + context = KernelContext( + node=context.node, + tables=context.tables, + weights=context.weights, + strata=context.strata, + params={**context.params, "extra": 1}, + rng=context.rng, + ) + with pytest.raises(ValueError, match="parameters must be exactly"): + _by_ref().run(context) + + @settings(max_examples=60, deadline=None) + @given(frame=nz_frames()) + def test_differential_each_node_equals_its_single_engine_kernel( + self, frame + ) -> None: + """Two implementations of one semantics: routed and direct agree.""" + + kernel = _by_ref() + for node, single in ( + ( + _person_node(SimulateRulesByRefKernel.ref), + SimulateRulesKernel(PERSON_REF, _PersonEngine()), + ), + ( + _family_node(SimulateRulesByRefKernel.ref), + SimulateRulesKernel(FAMILY_REF, _FamilyEngine()), + ), + ): + context = _context( + frame, node, weighted_entities=("person", "household", "family") + ) + _assert_same_result(kernel.run(context), single.run(context)) + + +# ---------------------------------------------------------------------- +# The graph +# ---------------------------------------------------------------------- + + +class _JsonPopulation(KernelBase): + """CREATE kernel loading an NZ population from a JSON source.""" + + ref = "test.json_population@1" + capabilities = Capabilities( + determinism=Determinism.DETERMINISTIC, + numeric=Numeric.BITWISE, + seed_source=SeedSource.NONE, + structural=StructuralDelta.CREATE, + ) + + def run(self, context: KernelContext) -> KernelResult: + document = json.loads(context.sources["population"].read_text()) + frame = _frame_from_document(document) + return KernelResult(frame=frame, receipt={"persons": frame.n("person")}) + + +def _frame_from_document(document: Mapping[str, object]) -> Frame: + person = pd.DataFrame(document["person"]).astype( + { + "person_id": "int64", + "person_household_id": "int64", + "person_family_id": "int64", + "earnings": "float64", + "age": "int64", + } + ) + household = pd.DataFrame(document["household"]).astype({"household_id": "int64"}) + family = pd.DataFrame(document["family"]).astype( + {"family_id": "int64", "family_rent": "float64"} + ) + return Frame( + {"person": person, "household": household, "family": family}, + NZ_SCHEMA, + { + "household": Weights( + np.asarray(document["household_weights"], dtype=np.float64), + WeightKind.DESIGN, + ) + }, + ) + + +def _population_frame() -> Frame: + return _frame_from_document(_POPULATION) + + +_POPULATION = { + "person": { + "person_id": [1, 2, 3, 4, 5, 6], + "person_household_id": [1, 1, 2, 2, 2, 3], + "person_family_id": [10, 10, 20, 21, 21, 30], + "earnings": [52_000.0, 0.0, 31_500.0, 18_250.0, 0.0, 74_000.0], + "age": [41, 9, 67, 35, 4, 23], + }, + "household": {"household_id": [1, 2, 3]}, + "family": { + "family_id": [10, 20, 21, 30], + "family_rent": [420.0, 0.0, 250.0, 180.0], + }, + "household_weights": [310.0, 145.5, 890.25], +} + +_POPULATION_NODE = Node( + "population", + _JsonPopulation.ref, + outputs=( + Owned("person", "earnings", "float64"), + Owned("person", "age", "int64"), + Owned("family", "family_rent", "float64"), + ), + structural=StructuralDelta.CREATE, + sources=("population",), +) + + +def _run(root: Path, nodes: tuple[Node, ...], kernels: tuple[object, ...]): + source = root / "population.json" + root.mkdir(parents=True, exist_ok=True) + source.write_text(json.dumps(_POPULATION)) + registry = KernelRegistry() + registry.register(_JsonPopulation()) + for kernel in kernels: + registry.register(kernel) + graph = Graph( + "rules-by-ref", + (SourceRef("population", "raw-bytes-v1"),), + (_POPULATION_NODE, *nodes), + ) + store = ContentStore(root / "store") + manifest = run_graph( + compile_graph(graph), + sources={"population": source}, + store=store, + kernels=registry, + resume="forbid", + decisions=(), + ) + return manifest, store + + +def _load(manifest, store, node_id: str, entity: str, column: str) -> pd.Series: + return store.load_column(manifest.nodes[node_id].artifacts[(entity, column)]) + + +class TestTwoEngineGraph: + def test_two_engines_compile_and_run_in_one_run_graph(self, tmp_path) -> None: + kernel = _by_ref() + manifest, store = _run( + tmp_path / "by-ref", + ( + _person_node(kernel.ref), + _family_node(kernel.ref), + ), + (kernel,), + ) + np.testing.assert_allclose( + _load(manifest, store, "rules.person", "person", "person_net_income"), + [41_601.0, 1.0, 25_201.0, 14_601.0, 1.0, 59_201.0], + ) + assert _load( + manifest, store, "rules.person", "person", "person_is_adult" + ).tolist() == [True, False, True, True, False, True] + np.testing.assert_allclose( + _load(manifest, store, "rules.family", "family", "family_support"), + [224.0, 0.0, 105.0, 56.0], + ) + codes = _load(manifest, store, "rules.family", "family", "family_code") + assert codes.dtype == np.dtype(np.int64) + assert codes.tolist() == [1, -1, 0, -1] + + def test_differential_outputs_equal_two_single_engine_runs(self, tmp_path) -> None: + kernel = _by_ref() + manifest, store = _run( + tmp_path / "by-ref", + (_person_node(kernel.ref), _family_node(kernel.ref)), + (kernel,), + ) + for node_id, node, single, variables in ( + ( + "rules.person", + _person_node(SimulateRulesKernel.ref), + SimulateRulesKernel(PERSON_REF, _PersonEngine()), + (("person", name) for name in PERSON_VARIABLES), + ), + ( + "rules.family", + _family_node(SimulateRulesKernel.ref), + SimulateRulesKernel(FAMILY_REF, _FamilyEngine()), + (("family", name) for name in FAMILY_VARIABLES), + ), + ): + single_manifest, single_store = _run( + tmp_path / f"single-{node_id}", (node,), (single,) + ) + assert ( + single_manifest.nodes[node_id].receipt + == manifest.nodes[node_id].receipt + ) + for entity, column in variables: + expected = _load(single_manifest, single_store, node_id, entity, column) + actual = _load(manifest, store, node_id, entity, column) + pd.testing.assert_series_equal(actual, expected) + assert actual.to_numpy().tobytes() == expected.to_numpy().tobytes() + + def test_adding_a_binding_re_keys_no_existing_node(self, tmp_path) -> None: + kernel = _by_ref() + manifest, _ = _run( + tmp_path / "two", + (_person_node(kernel.ref), _family_node(kernel.ref)), + (kernel,), + ) + extended = _by_ref(**{"third-engine": _PersonEngine()}) + extended_manifest, _ = _run( + tmp_path / "three", + (_person_node(extended.ref), _family_node(extended.ref)), + (extended,), + ) + for node_id in ("population", "rules.person", "rules.family"): + assert extended_manifest.nodes[node_id].key == manifest.nodes[node_id].key + + def test_each_node_key_binds_its_own_engine_ref(self, tmp_path) -> None: + kernel = _by_ref(**{"renamed-person-engine": _PersonEngine()}) + manifest, _ = _run(tmp_path / "base", (_person_node(kernel.ref),), (kernel,)) + renamed, _ = _run( + tmp_path / "renamed", + (_person_node(kernel.ref, engine_ref="renamed-person-engine"),), + (kernel,), + ) + assert renamed.nodes["rules.person"].key != manifest.nodes["rules.person"].key + assert renamed.nodes["population"].key == manifest.nodes["population"].key + + def test_a_node_naming_an_unbound_engine_rejects_the_run(self, tmp_path) -> None: + kernel = _by_ref() + with pytest.raises(Exception, match="no engine bound"): + _run( + tmp_path / "unbound", + (_person_node(kernel.ref, engine_ref="unbound"),), + (kernel,), + ) diff --git a/packages/microcosm-frame/tests/fixtures/rulespec-zz/zz/policies/tests/axiom_toy_family.yaml b/packages/microcosm-frame/tests/fixtures/rulespec-zz/zz/policies/tests/axiom_toy_family.yaml new file mode 100644 index 0000000000..5a1b619d78 --- /dev/null +++ b/packages/microcosm-frame/tests/fixtures/rulespec-zz/zz/policies/tests/axiom_toy_family.yaml @@ -0,0 +1,91 @@ +# Self-contained canonical-root RuleSpec module for the Axiom adapter's graph +# tests: family-scope outputs of every graph-typed engine dtype — a decimal +# rent assistance, a judgment asset test, an integer size band, and a boolean +# flag — shaped like the New Zealand family benefit unit (no relations). +format: rulespec/v1 +module: + summary: |- + microcosm-frame test fixture: family-scope rent assistance above a + weekly threshold, a cash-asset limit judgment, a family size band, and a + large-family flag. Values are round numbers chosen for hand verification. +units: + - name: EUR + kind: currency + minor_units: 2 +rules: + - name: toy_family_rent_threshold + kind: parameter + dtype: Money + unit: EUR + source: populace-frame fixture + versions: + - effective_from: '2025-01-01' + formula: |- + 100 + - name: toy_family_assistance_rate + kind: parameter + dtype: Rate + source: populace-frame fixture + versions: + - effective_from: '2025-01-01' + formula: |- + 0.5 + - name: toy_family_asset_limit + kind: parameter + dtype: Money + unit: EUR + source: populace-frame fixture + versions: + - effective_from: '2025-01-01' + formula: |- + 8000 + - name: toy_family_rent_assistance + kind: derived + entity: Family + dtype: Money + period: Year + unit: EUR + source: populace-frame fixture + versions: + - effective_from: '2025-01-01' + formula: |- + if toy_family_weekly_rent > toy_family_rent_threshold: + (toy_family_weekly_rent - toy_family_rent_threshold) * toy_family_assistance_rate + else: + 0 + - name: toy_family_assets_within_limit + kind: derived + entity: Family + dtype: Judgment + period: Year + source: populace-frame fixture + versions: + - effective_from: '2025-01-01' + formula: |- + toy_family_cash_assets <= toy_family_asset_limit + - name: toy_family_size_band + kind: derived + entity: Family + dtype: Integer + period: Year + source: populace-frame fixture + versions: + - effective_from: '2025-01-01' + formula: |- + if toy_family_size >= 3: + 2 + else: + 1 + - name: toy_family_is_large + kind: derived + entity: Family + dtype: Boolean + period: Year + source: populace-frame fixture + versions: + - effective_from: '2025-01-01' + formula: |- + if toy_family_size >= 3: + true + else: + false From 065d8eb396195d342ea14c60dc75cdf333dbd750 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 7 Oct 2026 05:31:23 -0400 Subject: [PATCH 2/2] Close review nits on the Axiom engine reference (#1118) Apply findings 1, 2, 3 and 5 of the independent review of 105c0625: - _require_clean_checkout now refuses any tracked file that `git ls-files -v` tags S (skip-worktree) or lowercase (assume-unchanged). git status reports no edit to such a file, so an edited, flagged file used to pass as a clean checkout at the declared commit. The check runs before the HEAD comparison and refuses a flagged file even when it is unedited. - AxiomEngine._require_pinned_tree refuses a symbolic link under the root before re-hashing it. The tree digest does not descend into a linked directory, so a directory link added after axiom_engine_ref left the pinned digest unchanged, and later compiles went ahead. - The axiom_engine_ref docstring now says the engine commit and wheel SHA-256 are pins the caller declares. They are recorded as given and not read from the installed engine, and the directory digest is the authority. Binding the installed engine is left for a later lane. - A new test checks that SimulateRulesByRefKernel's implementation hash is the same for all six binding orders of adapter classes from three modules. The hash must equal source_hash over the classes in (module, qualname) order. The test kills the review's surviving -id(cls) sort-key mutant and an unsorted insertion-order mutant; every earlier test let both through. New tests: a hidden-edit refusal for each flag, an acceptance once each flag is cleared, and a refusal of a symlink added after the reference. All five fail on 105c0625 with "DID NOT RAISE". Findings 4, 6 and 7 are outside this change: 4 lands with G2's export contract, 6 is a separate docstring nit (native text and date outputs are arrays, not lists), and 7 is the missing CI engine job. This narrows the PR description's checkout claim to: "A git checkout root must be clean and at the declared commit: no modified, untracked or ignored files, and no tracked file flagged skip-worktree or assume-unchanged (git status hides edits to those). A root holding a symbolic link is refused when the reference is made and on every later compile." The engine pins are described as "the declared engine commit and wheel SHA-256 (recorded as given, not read from the installed engine)". Co-Authored-By: Claude Opus 5.5 --- .../src/microcosm/frame/adapters/axiom.py | 66 ++++++++++++++----- .../shared/test_axiom_graph_adapter.py | 45 +++++++++++++ .../engine_free/shared/test_rules_kernels.py | 44 ++++++++++++- 3 files changed, 138 insertions(+), 17 deletions(-) diff --git a/packages/microcosm-frame/src/microcosm/frame/adapters/axiom.py b/packages/microcosm-frame/src/microcosm/frame/adapters/axiom.py index 17e6bb823d..1689fbc383 100644 --- a/packages/microcosm-frame/src/microcosm/frame/adapters/axiom.py +++ b/packages/microcosm-frame/src/microcosm/frame/adapters/axiom.py @@ -85,8 +85,9 @@ content digest of the whole RuleSpec root, and the adapter's configuration. Any edit to a RuleSpec byte therefore moves the reference, and with it every node key that names it. The reference also pins the adapter: from then on it -compiles only from a root whose digest still matches, so a node keyed by the -reference is never computed from other bytes. +compiles only from a root whose digest still matches and that holds no +symbolic link, so a node keyed by the reference is never computed from other +bytes. """ import hashlib @@ -811,12 +812,16 @@ def _require_pinned_tree(self) -> None: """Refuse to compile when the root no longer holds the referenced bytes. Once :func:`axiom_engine_ref` has named this adapter's RuleSpec tree, - every compile re-hashes the root, so a node keyed by that reference is - never computed from other bytes. + every compile refuses a symbolic link under the root (the digest does + not descend into a linked directory, so a directory or dangling link + added after the reference would not move it) and re-hashes the root, + so a node keyed by that reference is never computed from other bytes. Raises: - ValueError: If the root's digest moved since the reference. + ValueError: If the root holds a symbolic link, or its digest moved + since the reference. """ + _refuse_symlinks(self._rulespec_roots[0].resolve()) current, _ = rulespec_tree_digest(self._rulespec_roots[0].resolve()) if current != self._pinned_tree_sha256: raise ValueError( @@ -1140,8 +1145,8 @@ def axiom_engine_ref( that parameter enters the node key. The reference this function returns is canonical JSON of everything that decides the adapter's outputs: - * the engine (``axiom-rules-engine``), its commit, and the SHA-256 of the - wheel built from it; + * the engine (``axiom-rules-engine``) and the commit and wheel SHA-256 + the caller declares for it; * the RuleSpec commit, the module's path relative to the root, the module's SHA-256, and :func:`rulespec_tree_digest` of the whole root (imports resolve anywhere under it); @@ -1151,11 +1156,17 @@ def axiom_engine_ref( For a ``git archive`` export (the normal root, with no ``.git``) no absolute path enters the reference: the same tree exported to two places gives the same reference, and editing any byte under the root gives a - different one. The commits are declarations taken from the pin file, and - the directory digest is the authority. + different one. The engine commit and wheel SHA-256 are pins the caller + declares, not values read from the installed engine: nothing here + inspects ``axiom_rules_engine``, so a different engine installed under the + same pins gives the same reference. The RuleSpec commit is also a + declaration, checked only for a git checkout root (below). The directory + digest is the authority. A git checkout root is accepted only when clean, with no untracked or - ignored files, and at ``rulespec_commit``. Its ``.git`` is then hashed + ignored files and no tracked file flagged skip-worktree or + assume-unchanged (``git status`` does not report edits to those), and at + ``rulespec_commit``. Its ``.git`` is then hashed like any other file, as the graph's source key hashes it. That makes the reference specific to one clone: two clones or worktrees of one commit (a worktree's ``.git`` file names an absolute path) and an export of that @@ -1165,12 +1176,17 @@ def axiom_engine_ref( The first call pins the adapter to the tree digest it names: it refuses an adapter that has already compiled (those programs were read from bytes no reference names), and from then on every compile, and every later - reference to the same adapter, refuses a root whose digest has moved. + reference to the same adapter, refuses a root whose digest has moved or + that holds a symbolic link. Args: engine: The adapter the reference names. - engine_commit: Full lowercase git commit of axiom-rules-engine. - wheel_sha256: Lowercase SHA-256 of the installed engine wheel. + engine_commit: Full lowercase git commit of axiom-rules-engine, as + the caller declares it. It is recorded as given; the installed + engine is not checked against it. + wheel_sha256: Lowercase SHA-256 of the engine wheel, as the caller + declares it. It is recorded as given; no installed wheel or + native library is hashed to confirm it. rulespec_root: The adapter's only RuleSpec root. rulespec_commit: Full lowercase git commit the root was exported from. @@ -1182,7 +1198,8 @@ def axiom_engine_ref( ValueError: If a pin is malformed; ``rulespec_root`` is not the adapter's only root; the module is not a file under it; the root contains a symbolic link (the digest would not cover what it - points at); a git checkout root is dirty or at another commit; or + points at); a git checkout root is dirty, has a file flagged + skip-worktree or assume-unchanged, or is at another commit; or the adapter compiled before its first reference, or its root moved after it. """ @@ -1496,13 +1513,16 @@ def _require_clean_checkout(root: Path, commit: str) -> None: """Require a git-checkout RuleSpec root to hold exactly ``commit``. Ignored files count as changes: the digest hashes them, but the commit - does not hold them. ``--no-optional-locks`` stops ``git status`` + does not hold them. A tracked file flagged skip-worktree or + assume-unchanged is refused even when unedited, because ``git status`` + does not report edits to it. ``--no-optional-locks`` stops ``git status`` refreshing the index, which would otherwise rewrite ``.git/index`` and move the digest taken next. Raises: ValueError: If git fails or times out, the checkout has modified, - untracked, or ignored files, or ``HEAD`` is not ``commit``. + untracked, or ignored files or a file flagged skip-worktree or + assume-unchanged, or ``HEAD`` is not ``commit``. """ def git(*args: str) -> str: @@ -1535,6 +1555,20 @@ def git(*args: str) -> str: f"files that commit {commit} does not hold; reference a git " "archive export instead." ) + # ``ls-files -v`` tags a skip-worktree file ``S`` and an assume-unchanged + # file in lower case; ``git status`` reports no edit to either. + flagged = [ + entry[2:] + for entry in git("ls-files", "-v", "-z").split("\0") + if entry and (entry[0] == "S" or entry[0].islower()) + ] + if flagged: + raise ValueError( + f"RuleSpec checkout {root}: {flagged[0]} is flagged skip-worktree " + f"or assume-unchanged ({len(flagged)} flagged file(s)), so git " + "status does not report edits to it; clear the flags or reference " + "a git archive export instead." + ) head = git("rev-parse", "--verify", "HEAD^{commit}").strip() if head != commit: raise ValueError( diff --git a/packages/microcosm-frame/tests/engine_free/shared/test_axiom_graph_adapter.py b/packages/microcosm-frame/tests/engine_free/shared/test_axiom_graph_adapter.py index 30803e9041..609e971cb1 100644 --- a/packages/microcosm-frame/tests/engine_free/shared/test_axiom_graph_adapter.py +++ b/packages/microcosm-frame/tests/engine_free/shared/test_axiom_graph_adapter.py @@ -1281,6 +1281,23 @@ def test_an_edit_after_the_reference_blocks_a_second_reference( with pytest.raises(ValueError, match="changed after this adapter"): self._reference(adapter, root) + def test_a_symbolic_link_added_after_the_reference_blocks_the_next_compile( + self, tmp_path, monkeypatch + ) -> None: + # The digest does not descend into a linked directory, so it cannot + # see the link; the compile-time check must refuse it directly. + root = _write_tree(tmp_path / "rulespec", _BASE_TREE) + adapter = self._adapter(root) + programs = _fake_engine(adapter, monkeypatch) + self._reference(adapter, root) + before = rulespec_tree_digest(root) + target = _write_tree(tmp_path / "linked", {"hidden.yaml": b"x"}) + (root / "zz/policies/linked").symlink_to(target, target_is_directory=True) + assert rulespec_tree_digest(root) == before + with pytest.raises(ValueError, match="symbolic link zz/policies/linked"): + adapter.materialize(_fake_frame(), ["family_assistance"], 2026) + assert programs == {} + def test_a_reference_after_compiling_is_refused( self, tmp_path, monkeypatch ) -> None: @@ -1359,6 +1376,34 @@ def test_an_ignored_file_is_refused(self, checkout) -> None: with pytest.raises(ValueError, match="ignored"): self._checkout_ref(root, head) + @pytest.mark.parametrize("flag", ["--skip-worktree", "--assume-unchanged"]) + def test_an_edit_hidden_by_an_index_flag_is_refused(self, checkout, flag) -> None: + # git status trusts the flag and reports nothing, so the checkout would + # pass as clean while the module no longer holds the commit's bytes. + root, head = checkout + _git(root, "update-index", flag, _MODULE) + (root / _MODULE).write_bytes(b"format: rulespec/v1\nrules: [hidden]\n") + assert _git(root, "status", "--porcelain", "--ignored") == "" + with pytest.raises(ValueError, match=f"{_MODULE} is flagged"): + self._checkout_ref(root, head) + + @pytest.mark.parametrize( + ("flag", "unflag"), + [ + ("--skip-worktree", "--no-skip-worktree"), + ("--assume-unchanged", "--no-assume-unchanged"), + ], + ) + def test_a_clean_checkout_with_its_index_flags_cleared_is_accepted( + self, checkout, flag, unflag + ) -> None: + root, head = checkout + _git(root, "update-index", flag, _MODULE) + with pytest.raises(ValueError, match="flagged"): + self._checkout_ref(root, head) + _git(root, "update-index", unflag, _MODULE) + assert json.loads(self._checkout_ref(root, head))["rulespec_commit"] == head + def test_a_checkout_reference_is_clone_specific_and_an_export_is_not( self, checkout, tmp_path ) -> None: diff --git a/packages/microcosm-frame/tests/engine_free/shared/test_rules_kernels.py b/packages/microcosm-frame/tests/engine_free/shared/test_rules_kernels.py index 3aea3324e5..7a75d2ec89 100644 --- a/packages/microcosm-frame/tests/engine_free/shared/test_rules_kernels.py +++ b/packages/microcosm-frame/tests/engine_free/shared/test_rules_kernels.py @@ -8,6 +8,7 @@ both engines, and against two single-engine ``simulate.rules@1`` graph runs. """ +import itertools import json from collections.abc import Mapping, Sequence from pathlib import Path @@ -29,7 +30,7 @@ WeightKind, Weights, ) -from microcosm.frame.adapters.axiom import NZ_SCHEMA +from microcosm.frame.adapters.axiom import NZ_SCHEMA, AxiomEngine from microcosm.frame.kernels import SimulateRulesKernel from microcosm.frame.rules_kernels import SimulateRulesByRefKernel from microcosm.graph import ( @@ -54,6 +55,9 @@ source_hash, ) from test_support.microcosm_frame.kernels import _context, _StubRulesEngine +from test_support.paths import paths_for + +_RULESPEC_ROOT = paths_for("microcosm-frame").tests / "fixtures" / "rulespec-zz" PERSON_REF = "toy-person-engine" FAMILY_REF = "toy-family-engine" @@ -337,6 +341,44 @@ def test_adding_a_binding_of_a_bound_class_moves_no_hash(self) -> None: assert reordered.implementation_hash() == base.implementation_hash() assert extended.capabilities == base.capabilities + def test_the_hash_is_independent_of_binding_order(self) -> None: + """Invariant: the hash never depends on the order bindings arrive in. + + Three adapter classes from three modules make the source order + observable: all six binding orders must give the hash of the classes + in ``(module, qualname)`` order. These three classes sort the same way + by qualname alone, so the test pins order independence, not the + choice of sort key. + """ + + bindings = { + "axiom": AxiomEngine( + _RULESPEC_ROOT / "zz/policies/tests/axiom_toy_family.yaml", + schema=NZ_SCHEMA, + rulespec_roots=(_RULESPEC_ROOT,), + ), + "stub": _StubRulesEngine(), + "person": _PersonEngine(), + } + classes = sorted( + {type(engine) for engine in bindings.values()}, + key=lambda cls: (cls.__module__, cls.__qualname__), + ) + assert len({cls.__module__ for cls in classes}) == 3 + expected = source_hash( + SimulateRulesByRefKernel, + SimulateRulesKernel, + *classes, + frame_bundle_module, + frame_rules_module, + frame_schema_module, + ) + orders = list(itertools.permutations(bindings)) + assert len(orders) == 6 + for order in orders: + kernel = SimulateRulesByRefKernel({ref: bindings[ref] for ref in order}) + assert kernel.implementation_hash() == expected, order + def test_binding_a_new_adapter_class_moves_the_hash(self) -> None: # Documented: a new class's source joins the hash. assert (