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..1689fbc383 100644 --- a/packages/microcosm-frame/src/microcosm/frame/adapters/axiom.py +++ b/packages/microcosm-frame/src/microcosm/frame/adapters/axiom.py @@ -57,14 +57,50 @@ ``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 and that holds no +symbolic link, 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 +111,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 +133,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 +166,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 +288,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 +317,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 +361,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 +417,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 +514,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 +526,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 +601,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 +727,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 +775,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 +808,29 @@ 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 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 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( + 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 +845,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 +923,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 +1092,227 @@ 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``) 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); + * 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 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 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 + 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 or + that holds a symbolic link. + + Args: + engine: The adapter the reference names. + 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. + + 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, 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. + """ + 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 +1369,209 @@ 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. 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 a file flagged skip-worktree or + assume-unchanged, 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." + ) + # ``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( + 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..609e971cb1 --- /dev/null +++ b/packages/microcosm-frame/tests/engine_free/shared/test_axiom_graph_adapter.py @@ -0,0 +1,1869 @@ +"""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_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: + 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) + + @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: + # 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..7a75d2ec89 --- /dev/null +++ b/packages/microcosm-frame/tests/engine_free/shared/test_rules_kernels.py @@ -0,0 +1,707 @@ +"""``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 itertools +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, AxiomEngine +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 +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" +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_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 ( + _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