diff --git a/changelog.d/transport-compose-cli.added.md b/changelog.d/transport-compose-cli.added.md index 3f78eb3878..1b57572732 100644 --- a/changelog.d/transport-compose-cli.added.md +++ b/changelog.d/transport-compose-cli.added.md @@ -2,16 +2,18 @@ Add the transport skeleton composer, its local-only build driver and `tools/build_transport.py`. A country's `transport_graph.json` spec resource declares the skeleton graph; each node binds only the spec resources it selects, every rules node runs through the one `simulate.rules_by_ref@1` -router, and later packages append branches through a typed extension point -that cannot re-key the skeleton. The driver runs ancestor-closed checkpoints +router, and later packages append branches through a typed extension point. +Explicit artifact edges into the terminal package receipt may re-key that +receipt; every other skeleton predecessor set stays fixed. The driver runs ancestor-closed checkpoints on `/.graph-store` and writes every output under `--out`; it refuses an existing store tree that contains a link, which could redirect store writes. Activation checks run before registry preparation, CREATE or any graph source -read. Declaration selections must name present resources; value and JSON +read. Resource selections must name present resources; value and JSON selections must resolve existing paths to non-null selected values. Parameter -resolution shares these checks: JSON selections require objects, while value -selections reject objects. Optional null fields within selected JSON remain -allowed. Preflight also checks the dependent-child age limit, mandatory and +resolution shares these checks: skeleton JSON selections require objects; +entitlement JSON selections also accept lists, and value selections reject +objects. Optional null fields within selected JSON remain allowed. +Preflight also checks the dependent-child age limit, mandatory and selected reference activation rows, required engine commit and wheel pin presence, and every null scenario knob. The receipt kernel's shared validator checks contract fields, required text, program and exclusion structure, and diff --git a/changelog.d/transport-entitlement-gap.added.md b/changelog.d/transport-entitlement-gap.added.md new file mode 100644 index 0000000000..388472e5d5 --- /dev/null +++ b/changelog.d/transport-entitlement-gap.added.md @@ -0,0 +1,7 @@ +Add country-neutral counterfactual and row-independent fixed-bisection rules +kernels with final residual verification, weighted entitlement gaps and labelled +scenario bands, and a spec-declared transport extension for entitlement +scenarios, calibration variants and held-out checks. Preflight validates the +selected scenario rows, and the local differential driver refuses ambiguous +calls and invalid oracle pointers. Real-country activation and the approved +golden-08 executable bridge composition remain deferred. diff --git a/docs/agent-guide.md b/docs/agent-guide.md index 065af1f864..0860d1e94a 100644 --- a/docs/agent-guide.md +++ b/docs/agent-guide.md @@ -363,11 +363,14 @@ upload or staging path. Before it opens the store it refuses any link in an existing store tree (the root, `objects`, `tmp` or a shard), because the store creates those directories by following links; it does not guard against links planted while a build is running. Activation checks run before -registry preparation, CREATE or any graph source read. Every declaration -selection must name a present resource; value and JSON selections must have an -existing path and a non-null selected value. JSON selections require objects, -and value selections reject objects, using the same checks as parameter -resolution. Selected JSON may retain optional null fields. The preflight also +registry preparation, CREATE or any graph source read. Resource selections +must name a present resource; value and JSON selections must have an +existing path and a non-null selected value. Skeleton JSON selections require +objects; entitlement JSON selections also accept lists. Value selections +reject objects, using the same checks as parameter resolution. Selected JSON +may retain optional null fields. Entitlement resource checks use the factory's +common and selected scenario rows; explicit scenario `nodes` replaces the tier +template, so unused templates contribute no selections. The preflight also checks the dependent-child age limit, mandatory and selected reference activation rows, required engine commit and wheel pin presence, and every null scenario knob. The receipt kernel's shared validator checks the executable receipt @@ -386,12 +389,24 @@ build runs CREATE twice (inventory probe and graph execution); warm inventory lookups hash the inputs without running the probe. Cold `--resume require` refuses before preparation. Output directories must also be disjoint from a local `--spec-dir`. -G7 extensions can wrap a factory in `TransportExtension` to declare additional +Extensions can wrap a factory in `TransportExtension` to declare additional sources and ordered checkpoints. Source names must be new; skeleton predecessor -sets must stay identical. Checkpoints run after calibration and before export, +sets must stay identical except for explicit `package_inputs` artifact edges +from extension nodes to a terminal `transport.package@1` receipt. Other skeleton +nodes cannot gain predecessors. Checkpoints run after calibration and before export, and must not depend on the exported dataset. `transport_rules_node` supplies the prepared engine reference and period from the binding, so extensions do not repeat period values. +An optional `entitlement_graph.json` legacy resource installs the entitlement +extension in both composition and the local driver. Its scenario templates, +bridge parameters, comparison references and receipt connections are described +in [the entitlement declaration contract](transport-entitlement.md). +The zero solver requires independent, non-increasing row residuals and verifies +the returned answers after its fixed iterations. Its receipt counts the final +verification. The local differential driver validates call and output identities +before execution and uses canonical RFC 6901 oracle pointers. The packaged NZ +spec remains inactive; the approved golden-08 evaluation assembly, real WFF +inputs and law certification remain deferred. UK size experiments use `microcosm-build-uk --release-role dense --dataset-households` (`tools/build_uk_full.py`; `tools/build_uk_rowwise_candidate.py` is a stub over it) diff --git a/docs/transport-entitlement.md b/docs/transport-entitlement.md new file mode 100644 index 0000000000..bee2aaef38 --- /dev/null +++ b/docs/transport-entitlement.md @@ -0,0 +1,148 @@ +# Entitlement declaration contract + +This is an internal spec and hub hand-off contract. The executable example is +`test_support/microcosm_build/transport_entitlement.py`: its donor support, +facts, rates, references, pins and engine programs are invented test data. + +The optional `entitlement_graph.json` resource has `schema_version: 1` and is +listed as a `legacy_json` row in the country package. Composition and the local +driver install it automatically. Its fields are: + +| Field | Declaration | +| --- | --- | +| `nodes` | Ordered common bridge and validation node declarations, using the transport skeleton grammar. | +| `scenario_nodes`, `variant_nodes` | Node templates for the entitlement and calibration tiers of `scenarios.json`. A scenario row's `nodes` replaces its tier's template. Strings may contain `{scenario}` or `{variant}`. | +| `scenario_gap`, `variant_gap` | Gap producer id templates for the two tiers. | +| `bands` | `id`, `population`, `baseline_scenario`; all declared scenario gaps become typed inputs. | +| `package` | Final skeleton receipt `id`, bands `artifact_name`, optional additional `inputs` in the ordinary artifact-edge grammar. | +| `sources`, `checkpoints` | Additional sources with new names and ordered pre-export endpoints. The bands node is the default checkpoint. | + +Activation checks resource selections in the common nodes and each scenario's +selected nodes, using the entitlement factory's row selection. A scenario's +explicit `nodes` replaces its tier template; an unused template contributes no +resource selections. Selected resources, paths and reference activation remain +preflight checks before registry preparation or graph-source reads. + +Only explicit extension artifact edges may change a terminal skeleton package +receipt. +Every other skeleton predecessor set stays fixed. Scenarios branch below the +shared AS FILTER; calibration variants branch from the pre-calibration FILTER +and declare their own target/problem/calibration chain. AS inputs, AS engine +references and hold-out sources remain outside every calibration ancestor set. + +Parameters accept the ordinary resource selectors and two additional selectors: + +```json +{"binding": "rates_binding", "field": "engine_ref"} +{"binding": "rates_binding", "field": "period"} +{"scenario": "knobs", "path": [], "encoding": "json"} +{"scenario": "knobs", "path": [], "encoding": "sha256"} +``` + +Scenario selectors read the current row (`row` selects the entire row). +Their digest covers the selected subtree, so a sibling scenario edit or addition +does not move existing keys. Scenario overrides bind their selected knobs; +the keep-all scenario FILTER retains its key. Binding selectors inside selected +bridge JSON are resolved before canonical encoding. Structured literal records, +record lists and empty +lists are also canonical JSON; scalar lists remain graph tuples. Production +bridge declarations should select their coefficients and overrides from a +dedicated resource and bind its `resource_sha256`, rather than put policy data +in Python or bind the whole country fingerprint. + +`simulate.counterfactual@1` uses a prepared engine-reference mapping. Primary +parameters are `engine_ref`, `period`, `variables`, `input_overrides` and +`output_coefficients`; `components` supplies independent additional engines. +Overrides name an entity/input column and a scalar value or source column. +Terms name a requested engine output, coefficient and target column. A +person-to-family term declares `sum`, `max` or `first` aggregation; masks select +receipt branches. Components evaluate on independent copies of the input frame. +All bridge outputs are produced float64 columns. Each family engine node and +keep-all FILTER declares a person data-column slice. + +`simulate.solve_zero@1` adds `solve_input`, `bracket`, `tolerance` and +`iterations`. Each output row's residual must be independent of other rows' +trial inputs and non-increasing in its own trial input over the bracket. The +kernel does not support coupled systems. Fixed-count vector bisection locates +the boundary of positive residuals, including a benefit's zero plateau; +an already nonpositive lower endpoint returns that endpoint. A residual of +exactly zero at a positive trial input retains the nonzero boundary answer. +The iteration count must achieve the declared input-interval tolerance in +float64. After the fixed iterations the kernel evaluates the returned answers +and refuses if any row's positive residual exceeds the same declared tolerance. +A successful run reports `iterations + 3` residual evaluations: two bracket +endpoints, the fixed iterations and final verification. Each component engine +runs once per residual evaluation. The caller supplies the reference-case +overrides. + +`takeup.gap@1` parameters are `country`, `references`, `references_sha256`, +`entities`, `weight_entity`, `annualisation_factors` and `scenario` (optional +`facts_sha256` and `annualisation_sha256`). References select measured columns, +filters, comparator facts and table-cell metadata. Metadata labels are flat +string-valued fields, following the target compiler contract. Annualisation +factors map each reference name to its positive factor; count measures can +declare a factor of one. The artifact +records weighted entitlement `E`, published actual `A`, exactly `E - A`, the +scenario and reference/fact trace. Groupings, including decile by family type +and receipt strata, are reference data. `takeup.bands@1` requires identical +table definitions across scenarios, records low/baseline/high for E, A and gap, +and names each bound's scenario; ties use the first label in lexical order. + +The NZ package remains inactive. Its activation owner must supply the graph +declaration, prepared engine and RuleSpec pins, executable bridge overrides, +non-null knobs, active comparison references and annualisation/grouping data. +V1/V2 need explicit aged target-reference resources: the current target compiler +does not accept an `aging_index` parameter. If pre-calibration quantile maps are +aged too, those variants must branch from CREATE. V3 declares its receipt +override before its independent problem and calibration. + +Engine components and masks can declare an NZS/VP alternative; an explicit +exclusion input and counted excluded-unit reference can declare the fallback. +The toy tests exercise these declaration and grouping shapes with invented +rules. Actual reg 17 treatment requires approved bridge data, legal-source +review and real-engine differential evidence. + +Real WFF remains deferred pending approved input encoding, RuleSpec bindings +and Child-rule treatment. The supplied graph plan describes three chained +modules; this checkout does not certify their law semantics or external +approval. A tripwire declaration must calculate the WFF measure before +comparing it with the hold-out reference. The toy graph +does this with its invented family-credit engine. It supplies no real WFF +inputs or statutory claim. + +The committed `build/nz/as_rate_bridge.json` sole-parent declaration rules out +a `solve_zero` recipe for golden-08: no engine input combination reproduces +the harness's single-with-children gross rate paired with the Income Test 3 +slope. It instead declares three evaluations, `gross_rate`, +`reduction_at_anchor` and `reduction_after_step`, whose outputs recover the +slope and lower threshold before assembling the cutout. The resource declares +the overrides, anchor, step, composition and refusal conditions. The driver +executes declared kernel calls but does not assemble this nonlinear +composition. The approved executable golden-08 case and evaluation assembly +therefore remain unresolved; no golden-08 result is claimed. + +The hub can compare approved small cases and sampled-unit outputs without donor +H5 I/O using the declared case format in +`tools/transport_bridge_differential.py`: + +```bash +PYTHONDONTWRITEBYTECODE=1 \ +PYTHONPATH="$(ls -d $PWD/packages/*/src | tr '\n' ':')$PWD:$PWD/test_support" \ +UV_CACHE_DIR=$PWD/.hub-scratch/uv-cache \ +uv run --no-project .venv/bin/python tools/transport_bridge_differential.py \ + --case \ + --oracle \ + --rulespec-root \ + --out .hub-scratch/bridge-receipt.json +``` + +The case supplies unit-rule data, bindings, small entity tables and weights, +ordered kernel calls, and oracle JSON pointers with explicit absolute +tolerances. Call IDs must be distinct and nonempty, each output entity/column +coordinate must have one producer, and comparisons must name a declared call +output. These checks precede execution. Oracle pointers use RFC 6901 escapes; +array indices must be canonical nonnegative integers within bounds. A case can +include sampled units and matching oracle arrays. With the engine available, +the receipt records each comparison and the case/oracle hashes. If the engine +is absent the command writes a skipped receipt without reading case data; +no real-engine result is claimed. diff --git a/packages/microcosm-build/src/microcosm/build/transport/compose.py b/packages/microcosm-build/src/microcosm/build/transport/compose.py index 97296c6c69..30b0f2b000 100644 --- a/packages/microcosm-build/src/microcosm/build/transport/compose.py +++ b/packages/microcosm-build/src/microcosm/build/transport/compose.py @@ -39,7 +39,8 @@ from its binding, including for extension nodes. Composition refuses an extension whose nodes change any skeleton node's compiled predecessors (for example a new member of a calibration base -version), so an extension can never re-key the skeleton. +version). Explicit ``package_inputs`` may add extension artifacts to the final +receipt; every other skeleton predecessor set remains fixed. Composition reads no donor, no facts and no engine, and writes nothing. Activation checks selected resource paths, non-null selected roots and encoding @@ -208,12 +209,15 @@ class TransportExtension: Sources must have new names, so they cannot replace a skeleton input. Checkpoints name nodes in the extended graph and run after calibration, before the skeleton H5 is written; they cannot read the exported dataset. + ``package_inputs`` explicitly connect typed extension evidence to a + skeleton package receipt. No other skeleton node can gain predecessors. The factory receives a deep copy of the caller's spec. """ factory: TransportGraphExtension sources: tuple[SourceRef, ...] = () checkpoints: tuple[str, ...] = () + package_inputs: tuple[tuple[str, ArtifactInput], ...] = () def __post_init__(self) -> None: if not callable(self.factory): @@ -226,6 +230,14 @@ def __post_init__(self) -> None: raise TypeError("Extension checkpoints must be a tuple of node ids.") for checkpoint in self.checkpoints: _safe_node_id(checkpoint) + if not isinstance(self.package_inputs, tuple): + raise TypeError("Extension package inputs must be a tuple.") + for target, edge in self.package_inputs: + _safe_node_id(target) + if not isinstance(edge, ArtifactInput): + raise TypeError( + "Extension package inputs require typed artifact edges." + ) def __call__( self, @@ -416,6 +428,8 @@ def validate_transport_activation(spec: Mapping[str, object]) -> None: Also check dependent-child age and required engine pin presence, mandatory and selected reference activation rows, and every null scenario knob. Selections inside parameter lists count; selected resource data is literal. + Entitlement selections use the factory's instantiated rows, so explicit + scenario nodes replace their tier templates during preflight too. """ missing = [] @@ -424,14 +438,25 @@ def validate_transport_activation(spec: Mapping[str, object]) -> None: if _has_resource(spec, GRAPH_RESOURCE) else () ) + selections = [(selection, False) for selection in selections] + if _has_resource(spec, "entitlement_graph") and _has_resource(spec, "scenarios"): + from .entitlement import _rows + + rows, _ = _rows(spec) + selections.extend( + (selection, True) + for selection in _declared_resource_selections( + {"nodes": [row for row, _ in rows]} + ) + ) selected = { selection["resource"].removesuffix(".json") - for selection in selections + for selection, _ in selections if isinstance(selection["resource"], str) and selection["resource"] } - for selection in selections: + for selection, allow_json_sequences in selections: try: - _resource_param(spec, selection) + _resource_param(spec, selection, allow_json_sequences=allow_json_sequences) except (ValueError, TypeError) as error: missing.append( f"Resource {selection['resource']!r} path {selection.get('path', [])!r}: " @@ -764,6 +789,9 @@ def transport_endpoints( ) -> dict[str, object]: """Geography, calibration, extension checkpoints, then export and terminals.""" + from .entitlement import transport_extensions + + extensions = transport_extensions(spec, extensions) values = _mapping( transport_graph_document(spec).get("endpoints"), "transport endpoints" ) @@ -836,11 +864,14 @@ def _param( return _resource_param(spec, value) -def _resource_param(spec: Mapping[str, object], value: Mapping) -> object: +def _resource_param( + spec: Mapping[str, object], value: Mapping, *, allow_json_sequences: bool = False +) -> object: """Resolve and validate one literal resource selection for preflight/params. - Selected roots cannot be null; JSON roots must be objects, and value roots - must not be objects. Optional JSON fields remain unchanged. + Selected roots cannot be null; skeleton JSON roots must be objects. + Entitlement JSON may also select lists. Value roots must not be objects. + Optional JSON fields remain unchanged. """ if set(value) - {"resource", "path", "encoding"} or "resource" not in value: @@ -868,7 +899,9 @@ def _resource_param(spec: Mapping[str, object], value: Mapping) -> object: if document is None: raise ValueError(f"Resource {name!r} path {list(path)!r} is not activated.") if encoding == "json": - if not isinstance(document, Mapping): + if not isinstance(document, Mapping) and not ( + allow_json_sequences and isinstance(document, list | tuple) + ): raise ValueError( f"Resource {name!r} path {list(path)!r} json selection must be an object." ) @@ -1012,6 +1045,9 @@ def _node( row: Mapping, config: TransportGraphConfig, types: Mapping[tuple[str, str], ArtifactType], + *, + allowed_kernels: frozenset[str] = SKELETON_KERNELS, + resolved_params: Mapping[str, object] | None = None, ) -> Node: _safe_node_id(row.get("id")) unknown = set(row) - _NODE_FIELDS @@ -1020,14 +1056,18 @@ def _node( f"Transport node {row.get('id')!r} has unknown fields {sorted(unknown)}." ) kernel = row.get("kernel") - if kernel not in SKELETON_KERNELS: + if kernel not in allowed_kernels: raise ValueError( f"Kernel {kernel!r} is not a skeleton kernel; add it through an extension." ) - params = { - name: _param(spec, value, config.engine_refs) - for name, value in _mapping(row.get("params", {}), "node params").items() - } + params = ( + dict(resolved_params) + if resolved_params is not None + else { + name: _param(spec, value, config.engine_refs) + for name, value in _mapping(row.get("params", {}), "node params").items() + } + ) if "rules_binding" in row: params = _rules_params(spec, row, params, config) elif kernel == RULES_ROUTER: @@ -1136,6 +1176,7 @@ def compose_transport_graph( ) -> Graph: """The skeleton declared by the spec, plus any extension branches. + An optional entitlement graph resource installs its extension automatically. The result is a pure function of the spec's resources and ``config``: composing twice gives equal ``graph_to_json``. Activation is checked first, so pending country evidence fails with its items named. @@ -1144,6 +1185,9 @@ def compose_transport_graph( if not isinstance(config, TransportGraphConfig): raise TypeError("config must be a TransportGraphConfig.") validate_transport_activation(spec) + from .entitlement import transport_extensions + + config = replace(config, extensions=transport_extensions(spec, config.extensions)) document = transport_graph_document(spec) validate_transport_calibration_ancestry(spec, config=config) rows = tuple( @@ -1175,9 +1219,26 @@ def compose_transport_graph( transport_pending_outputs(spec) additional: list[Node] = [] additional_sources: list[SourceRef] = [] + package_inputs: dict[str, list[ArtifactInput]] = {} source_names = {source.name for source in skeleton.sources} for factory in config.extensions: if isinstance(factory, TransportExtension): + for target, edge in factory.package_inputs: + if ( + target not in compiled.order + or skeleton.node(target).kernel != "transport.package@1" + ): + raise ValueError( + "Extension artifact edges may target only a skeleton package node." + ) + if any( + target in predecessors + for predecessors in compiled.predecessors.values() + ): + raise ValueError( + "Extension package inputs require a terminal skeleton package node." + ) + package_inputs.setdefault(target, []).append(edge) for source in factory.sources: if source.name in source_names: raise ValueError( @@ -1193,20 +1254,36 @@ def compose_transport_graph( for node in nodes: _safe_node_id(node.id) additional.extend(nodes) - if not additional and not additional_sources: + if not additional and not additional_sources and not package_inputs: _check_extension_checkpoints(spec, config, skeleton, compiled) return skeleton + skeleton_nodes = [] + for node in skeleton.nodes: + edges = tuple(package_inputs.get(node.id, ())) + aliases = [edge.name for edge in (*node.artifact_inputs, *edges)] + if len(aliases) != len(set(aliases)): + raise ValueError( + "Extension package artifact aliases must be new and distinct." + ) + if any(edge.producer not in {item.id for item in additional} for edge in edges): + raise ValueError("Extension package inputs must come from extension nodes.") + skeleton_nodes.append( + replace(node, artifact_inputs=(*node.artifact_inputs, *edges)) + ) graph = Graph( skeleton.country, (*skeleton.sources, *additional_sources), - (*skeleton.nodes, *additional), + (*skeleton_nodes, *additional), skeleton.mass_partition, ) extended = compile_graph(graph) _check_extension_checkpoints(spec, config, graph, extended) validate_transport_calibration_ancestry(spec, graph=graph, config=config) for node in skeleton.nodes: - if extended.predecessors[node.id] != compiled.predecessors[node.id]: + expected = set(compiled.predecessors[node.id]) | { + edge.producer for edge in package_inputs.get(node.id, ()) + } + if set(extended.predecessors[node.id]) != expected: raise ValueError( f"An extension would re-key skeleton node {node.id!r}; " "branch a new FILTER population instead." diff --git a/packages/microcosm-build/src/microcosm/build/transport/entitlement.py b/packages/microcosm-build/src/microcosm/build/transport/entitlement.py new file mode 100644 index 0000000000..1c644d05ce --- /dev/null +++ b/packages/microcosm-build/src/microcosm/build/transport/entitlement.py @@ -0,0 +1,306 @@ +"""Declare entitlement bridges, scenario branches and their held-out bands. + +``entitlement_graph.json`` is an optional legacy-JSON resource. Its ``nodes`` +use the skeleton node declaration grammar. ``scenario_nodes`` and +``variant_nodes`` are templates, expanded for the entitlement and calibration +tiers of ``scenarios.json`` respectively. A row's ``nodes`` can replace its +tier's template, for example to select an engine-derived alternative bridge. +Node ids and data strings can contain ``{scenario}`` or ``{variant}``. + +Extension parameters additionally accept ``{binding, field}`` selections +(field is engine_ref or period), and ``{scenario, path?, encoding?}`` selections +of the current scenario row. The latter's SHA-256 binds only the selected +subtree, so another scenario can be added without re-keying existing branches. +Binding selections inside selected JSON bridge data are resolved explicitly. +Every policy coefficient, override and solver setting remains spec data. + +``bands`` declares id, population and baseline_scenario; ``scenario_gap`` and +``variant_gap`` name each tier's gap producer template. ``package`` names the +skeleton receipt node and the bands artifact alias; its optional ``inputs`` +adds other typed extension evidence, such as a held-out validation tripwire. +The only skeleton change is these explicit artifact edges on the final receipt. +""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Mapping +from dataclasses import dataclass + +from microcosm.graph import ArtifactInput, ArtifactOutput, Graph, Node, SourceRef +from microcosm.graph.canonical import canonical_json + +from .compose import ( + SKELETON_KERNELS, + TransportExtension, + TransportGraphConfig, + _artifact_outputs, + _literal_data, + _mapping, + _node, + _param, + _resource_param, + _safe_node_id, + _sequence, + transport_resource, +) +from .takeup_kernels import BANDS_TYPE, GAP_TYPE + +__all__ = ["ENTITLEMENT_RESOURCE", "entitlement_extension", "transport_extensions"] + +ENTITLEMENT_RESOURCE = "entitlement_graph" +_KERNELS = SKELETON_KERNELS | frozenset( + { + "simulate.counterfactual@1", + "simulate.solve_zero@1", + "takeup.gap@1", + "takeup.bands@1", + "takeup.compare@1", + "transport.scenario_override@1", + } +) +_FIELDS = frozenset( + { + "schema_version", + "nodes", + "scenario_nodes", + "variant_nodes", + "scenario_gap", + "variant_gap", + "sources", + "checkpoints", + "bands", + "package", + } +) + + +def _document(spec): + document = transport_resource(spec, ENTITLEMENT_RESOURCE) + if document.get("schema_version") != 1 or set(document) - _FIELDS: + raise ValueError( + "An entitlement graph requires schema_version 1 and known fields." + ) + return document + + +def _format(value, substitutions): + if isinstance(value, str): + for key, replacement in substitutions.items(): + value = value.replace("{" + key + "}", replacement) + return value + if isinstance(value, Mapping): + return {key: _format(item, substitutions) for key, item in value.items()} + if isinstance(value, list | tuple): + return [_format(item, substitutions) for item in value] + return value + + +def _binding(spec, config, selection): + if set(selection) != {"binding", "field"}: + raise ValueError("A binding selection names exactly binding and field.") + name, field = selection["binding"], selection["field"] + rows = transport_resource(spec, "axiom_rules_bindings")["bindings"] + matches = [row for row in rows if row["id"] == name] + if len(matches) != 1 or name not in config.engine_refs: + raise ValueError(f"Entitlement binding {name!r} needs one prepared reference.") + if field == "engine_ref": + return config.engine_refs[name] + if field == "period": + return matches[0]["period"] + raise ValueError("A binding selection field must be engine_ref or period.") + + +def _bind_data(spec, config, value): + if isinstance(value, Mapping): + if "binding" in value and "field" in value: + return _binding(spec, config, value) + return {key: _bind_data(spec, config, item) for key, item in value.items()} + if isinstance(value, list | tuple): + return [_bind_data(spec, config, item) for item in value] + return value + + +def _parameter(spec, config, value, scenario): + if isinstance(value, list | tuple): + if not value or any(isinstance(item, Mapping) for item in value): + return canonical_json(_bind_data(spec, config, value)).decode("utf-8") + return tuple(_parameter(spec, config, item, scenario) for item in value) + if not isinstance(value, Mapping): + return value + if "binding" in value: + return _binding(spec, config, value) + if "scenario" in value: + if scenario is None or set(value) - {"scenario", "path", "encoding"}: + raise ValueError( + "A scenario selection requires a current row and known fields." + ) + key = value["scenario"] + selected = scenario if key == "row" else scenario[key] + for component in value.get("path", ()): + selected = _mapping(selected, "scenario selection")[component] + encoding = value.get("encoding", "value") + if encoding == "sha256": + return hashlib.sha256(canonical_json(selected)).hexdigest() + if encoding == "json": + return canonical_json(_bind_data(spec, config, selected)).decode("utf-8") + if encoding == "value" and not isinstance(selected, Mapping): + return _literal_data(selected) + raise ValueError( + "Scenario objects require json encoding; encodings are value, json, sha256." + ) + if value.get("encoding") == "json" and "resource" in value: + selected = json.loads(_resource_param(spec, value, allow_json_sequences=True)) + return canonical_json(_bind_data(spec, config, selected)).decode("utf-8") + if not ({"resource", "prepared"} & set(value)): + return canonical_json(_bind_data(spec, config, value)).decode("utf-8") + return _param(spec, value, config.engine_refs) + + +def _rows(spec): + """Rows the factory instantiates, shared with activation preflight.""" + document = _document(spec) + result = [(row, None) for row in document.get("nodes", ())] + gaps = {} + rows = transport_resource(spec, "scenarios")["scenarios"] + for scenario in _sequence(rows, "scenarios.scenarios"): + scenario = _mapping(scenario, "scenario") + name = _safe_node_id(scenario.get("id")) + tier = scenario.get("tier") + if tier not in {"entitlement", "calibration"} or name in gaps: + raise ValueError( + "Scenario ids must be distinct and tiers entitlement or calibration." + ) + token = "scenario" if tier == "entitlement" else "variant" + template = "scenario_nodes" if tier == "entitlement" else "variant_nodes" + declared = scenario.get("nodes", document.get(template, ())) + if not declared: + raise ValueError(f"Scenario {name!r} requires a node template.") + result.extend( + (_format(_mapping(row, "entitlement node"), {token: name}), scenario) + for row in _sequence(declared, template) + ) + gap = document["scenario_gap" if tier == "entitlement" else "variant_gap"] + gaps[name] = _format(gap, {token: name}) + if not gaps: + raise ValueError("An entitlement graph needs at least its baseline scenario.") + return result, gaps + + +def _types(rows, skeleton=None): + types = ( + {} + if skeleton is None + else { + (node.id, output.name): output.type + for node in skeleton.nodes + for output in node.artifact_outputs + } + ) + for row, _ in rows: + outputs = {"takeup.gap@1": (ArtifactOutput("gap", GAP_TYPE),)}.get( + row["kernel"], _artifact_outputs(row) + ) + types.update({(row["id"], output.name): output.type for output in outputs}) + return types + + +@dataclass(frozen=True) +class _EntitlementFactory: + def __call__(self, spec: Mapping, config: TransportGraphConfig, skeleton: Graph): + document = _document(spec) + rows, gaps = _rows(spec) + types = _types(rows, skeleton) + nodes: list[Node] = [] + for row, scenario in rows: + params = { + key: _parameter(spec, config, value, scenario) + for key, value in row.get("params", {}).items() + } + if row["kernel"] == "takeup.gap@1" and "artifact_outputs" not in row: + row = { + **row, + "artifact_outputs": [ + { + "name": "gap", + "type": { + "name": GAP_TYPE.name, + "schema_version": GAP_TYPE.schema_version, + }, + } + ], + } + nodes.append( + _node( + spec, + row, + config, + types, + allowed_kernels=_KERNELS, + resolved_params=params, + ) + ) + bands = _mapping(document["bands"], "entitlement bands") + aliases = {f"scenario_{index}": name for index, name in enumerate(sorted(gaps))} + nodes.append( + Node( + id=bands["id"], + kernel="takeup.bands@1", + population=bands["population"], + params={ + "baseline_scenario": bands["baseline_scenario"], + "artifact_scenarios": canonical_json(aliases).decode("utf-8"), + }, + artifact_inputs=tuple( + ArtifactInput(alias, gaps[name], "gap", GAP_TYPE) + for alias, name in aliases.items() + ), + artifact_outputs=(ArtifactOutput("bands", BANDS_TYPE),), + ) + ) + return tuple(nodes) + + +def entitlement_extension(spec: Mapping) -> TransportExtension: + """The optional entitlement declarations as an ordered transport extension.""" + document = _document(spec) + rows, _ = _rows(spec) + bands = _mapping(document["bands"], "entitlement bands") + package = _mapping(document["package"], "entitlement package") + types = _types(rows) + types[bands["id"], "bands"] = BANDS_TYPE + edges = [ + ArtifactInput( + package.get("artifact_name", "bands"), bands["id"], "bands", BANDS_TYPE + ) + ] + for row in package.get("inputs", ()): + artifact = row.get("artifact", row["name"]) + edges.append( + ArtifactInput( + row["name"], row["producer"], artifact, types[row["producer"], artifact] + ) + ) + return TransportExtension( + factory=_EntitlementFactory(), + sources=tuple(SourceRef(**row) for row in document.get("sources", ())), + checkpoints=tuple(document.get("checkpoints", (bands["id"],))), + package_inputs=tuple((package["id"], edge) for edge in edges), + ) + + +def transport_extensions(spec: Mapping, extensions=()) -> tuple: + """Add the spec-declared entitlement extension once, preserving caller probes.""" + resources = spec.get("resources", {}) + if not any( + name in resources + for name in (ENTITLEMENT_RESOURCE, ENTITLEMENT_RESOURCE + ".json") + ): + return extensions + if any( + isinstance(getattr(item, "factory", None), _EntitlementFactory) + for item in extensions + ): + return extensions + return (*extensions, entitlement_extension(spec)) diff --git a/packages/microcosm-build/src/microcosm/build/transport/registry.py b/packages/microcosm-build/src/microcosm/build/transport/registry.py index ea25ced9fa..5a65a1a806 100644 --- a/packages/microcosm-build/src/microcosm/build/transport/registry.py +++ b/packages/microcosm-build/src/microcosm/build/transport/registry.py @@ -7,9 +7,8 @@ bytes and tree contents are hashed; installed engine provenance is not checked. Kernel implementation identity binds adapter classes, never spec data, so adding a binding of an existing class leaves it unchanged. -The AS bridge and gap kernels are registered by their owning package when -that package is available; this module -does not substitute computations for kernels that have not been implemented. +Counterfactual and zero-cutout bridges share the same prepared engine mapping. +Gap and scenario-band kernels are registered with the other transport kernels. """ from __future__ import annotations @@ -32,7 +31,11 @@ rulespec_tree_digest, ) from microcosm.frame.rules import RulesEngine -from microcosm.frame.rules_kernels import SimulateRulesByRefKernel +from microcosm.frame.rules_kernels import ( + SimulateCounterfactualKernel, + SimulateRulesByRefKernel, + SimulateSolveZeroKernel, +) from microcosm.frame.unit_construction import BenefitUnitRule from microcosm.graph import KernelRegistry from microcosm.graph.canonical import canonical_json @@ -44,6 +47,7 @@ from .gate_kernels import register_gate_kernel from .geography_kernels import register_geography_kernels from .population_kernels import register_population_kernels +from .takeup_kernels import register_takeup_kernels from .target_kernels import register_target_kernels from .terminal_kernels import register_terminal_kernels @@ -373,6 +377,9 @@ def build_transport_registry( register_gate_kernel(kernels, TRANSPORT_GATE_REGISTRY) register_terminal_kernels(kernels, codecs=codecs) kernels.register(SimulateRulesByRefKernel(engine_mapping, dependencies)) + kernels.register(SimulateCounterfactualKernel(engine_mapping, dependencies)) + kernels.register(SimulateSolveZeroKernel(engine_mapping, dependencies)) + register_takeup_kernels(kernels) return TransportRegistry( kernels, codecs, engine_mapping, MappingProxyType(refs), schema, nesting ) diff --git a/packages/microcosm-build/src/microcosm/build/transport/takeup_kernels.py b/packages/microcosm-build/src/microcosm/build/transport/takeup_kernels.py new file mode 100644 index 0000000000..89a0f7a8a8 --- /dev/null +++ b/packages/microcosm-build/src/microcosm/build/transport/takeup_kernels.py @@ -0,0 +1,496 @@ +"""Weighted entitlement gaps and labelled scenario envelopes. + +``takeup.gap@1`` reads published comparators through the same executable +target-reference compiler as calibration. Its references select the measured +columns and filters; each reference has an explicit annualisation factor. +The estimate is the sum of effective weight times annualised measure and the +gap is exactly that estimate minus the selected comparator. Reference +metadata can label any declared table cells, including receipt strata. + +``takeup.bands@1`` consumes the gap artifacts from declared scenarios. For +each table cell and each of its estimate, actual and gap values it records +the minimum, central scenario and maximum, naming the scenario at every +bound. Equal-valued ties use the lexicographically first scenario label. +No scenario list, rate or comparator is held in the kernel implementation. +""" + +from __future__ import annotations + +import json +import math +from collections.abc import Mapping + +import numpy as np + +from microcosm.graph import ( + ArtifactType, + Capabilities, + Determinism, + KernelBase, + KernelContext, + KernelRegistry, + KernelResult, + Numeric, + SeedSource, + source_hash, +) +from microcosm.graph.canonical import canonical_json + +from . import graph_inputs, target_kernels +from .graph_inputs import ( + canonical_document_param, + context_frame, + entities_param, + require_params, + sha256_param, + sha256_text, + single_source, + string_param, +) +from .target_kernels import ( + compile_target_surface, + decode_target_surface, + parse_reference_document, +) + +__all__ = [ + "BANDS_TYPE", + "GAP_TYPE", + "TAKEUP_BANDS", + "TAKEUP_GAP", + "TakeupBandsKernel", + "TakeupGapKernel", + "decode_bands", + "decode_gap", + "register_takeup_kernels", +] + +GAP_TYPE = ArtifactType("microcosm.takeup.gap", 1) +BANDS_TYPE = ArtifactType("microcosm.takeup.bands", 1) +_GAP_KIND = "transport_entitlement_gap" +_BANDS_KIND = "transport_entitlement_bands" +_METRICS = ("E", "A", "gap") +_ROW_SEMANTICS = ("name", "period", "entity", "measure", "filter", "metadata") +_SCENARIO_SEMANTICS = ( + *_ROW_SEMANTICS, + "annualisation_factor", + "reference", + "ledger_fact_key", +) +_ROW_KEYS = frozenset( + {*_ROW_SEMANTICS, "annualisation_factor", "reference", "ledger_fact_key", *_METRICS} +) + + +def _number(value: object, label: str, *, positive: bool = False) -> float: + if isinstance(value, bool) or not isinstance(value, int | float): + raise TypeError(f"{label} must be a finite number.") + result = float(value) + if not math.isfinite(result) or (positive and result <= 0): + raise ValueError( + f"{label} must be finite" + (" and positive." if positive else ".") + ) + return result + + +def _text(value: object, label: str) -> str: + if not isinstance(value, str) or not value: + raise ValueError(f"{label} must be a non-empty string.") + return value + + +def _document(payload: bytes, kind: str, keys: set[str]) -> dict: + if type(payload) is not bytes: + raise TypeError("Transport artifacts are immutable bytes.") + document = json.loads(payload) + if ( + not isinstance(document, dict) + or canonical_json(document) != payload + or document.get("schema_version") != 1 + or document.get("kind") != kind + or set(document) != keys | {"schema_version", "kind"} + ): + raise ValueError(f"Unsupported or noncanonical {kind} artifact.") + return document + + +def _cell(row: Mapping) -> tuple[str, int | str]: + name = _text(row.get("name"), "gap row.name") + period = row.get("period") + if isinstance(period, bool) or not isinstance(period, int | str): + raise ValueError("gap row.period must be an integer or string.") + return name, period + + +def decode_gap(payload: bytes) -> dict: + """Validate canonical gap bytes, including exact subtraction.""" + + document = _document( + payload, + _GAP_KIND, + { + "country", + "scenario", + "weight_entity", + "weight_kind", + "facts", + "references_sha256", + "gaps", + }, + ) + for name in ("country", "scenario", "weight_entity", "weight_kind"): + _text(document[name], f"gap.{name}") + if not isinstance(document["facts"], dict): + raise ValueError("gap.facts must carry the comparator source identity.") + rows = document["gaps"] + if not isinstance(rows, list) or not rows: + raise ValueError("A gap artifact needs a non-empty list of table cells.") + seen = set() + for row in rows: + if not isinstance(row, dict) or set(row) != _ROW_KEYS: + raise ValueError("Unsupported gap row shape.") + cell = _cell(row) + if cell in seen: + raise ValueError("A gap artifact repeats a table cell.") + seen.add(cell) + for name in ("entity", "measure", "reference"): + _text(row[name], f"gap row.{name}") + if row["reference"] != row["name"]: + raise ValueError("A gap row must trace to its named reference.") + if row["filter"] is not None: + _text(row["filter"], "gap row.filter") + if not isinstance(row["metadata"], dict): + raise ValueError("gap row.metadata must be an object.") + _number(row["annualisation_factor"], "annualisation_factor", positive=True) + estimate = _number(row["E"], "E") + actual = _number(row["A"], "A") + gap = _number(row["gap"], "gap") + if gap != estimate - actual: + raise ValueError("The gap must equal E minus A exactly.") + return document + + +def decode_bands(payload: bytes) -> dict: + """Validate a labelled envelope, its ordering and its scenario coverage.""" + + document = _document( + payload, + _BANDS_KIND, + {"country", "baseline_scenario", "scenarios", "bands"}, + ) + _text(document["country"], "bands.country") + baseline = _text(document["baseline_scenario"], "bands.baseline_scenario") + scenarios = document["scenarios"] + if not isinstance(scenarios, list) or not scenarios: + raise ValueError("Bands need at least the central scenario.") + labels = [] + for scenario in scenarios: + if not isinstance(scenario, dict) or set(scenario) != { + "scenario", + "gap_sha256", + }: + raise ValueError("Bands scenario rows must name their gap artifact.") + labels.append(_text(scenario["scenario"], "bands scenario")) + digest = scenario["gap_sha256"] + if ( + not isinstance(digest, str) + or len(digest) != 64 + or not set(digest) <= set("0123456789abcdef") + ): + raise ValueError("Bands scenario gap_sha256 must be a SHA-256 digest.") + if labels != sorted(set(labels)) or baseline not in labels: + raise ValueError( + "Bands scenarios must be distinct, ordered and include the central scenario." + ) + rows = document["bands"] + if not isinstance(rows, list) or not rows: + raise ValueError("Bands need a non-empty list of table cells.") + seen = set() + for row in rows: + if not isinstance(row, dict) or set(row) != { + *_ROW_SEMANTICS, + "metric", + "low", + "baseline", + "high", + }: + raise ValueError("Unsupported bands row shape.") + cell = (*_cell(row), row["metric"]) + if cell in seen or row["metric"] not in _METRICS: + raise ValueError("Bands repeat a cell or declare an unknown metric.") + seen.add(cell) + values = [] + for bound in ("low", "baseline", "high"): + value = row[bound] + if not isinstance(value, dict) or set(value) != {"value", "scenario"}: + raise ValueError("Every band bound must name its scenario.") + if value["scenario"] not in labels: + raise ValueError("A band bound names an undeclared scenario.") + values.append(_number(value["value"], f"bands.{bound}.value")) + if row["baseline"]["scenario"] != baseline: + raise ValueError("The central band must name the central scenario.") + if not values[0] <= values[1] <= values[2]: + raise ValueError("Band bounds must be ordered low <= central <= high.") + return document + + +def _require_output( + context: KernelContext, ref: str, name: str, kind: ArtifactType +) -> None: + declared = {item.name: item.type for item in context.node.artifact_outputs} + if declared != {name: kind}: + raise ValueError( + f"{ref} must declare exactly the {name!r} typed artifact output." + ) + + +class TakeupGapKernel(KernelBase): + """``takeup.gap@1``: reference-driven annualised weighted entitlement.""" + + ref = "takeup.gap@1" + capabilities = Capabilities( + Determinism.DETERMINISTIC, + numeric=Numeric.BITWISE, + seed_source=SeedSource.NONE, + ) + _required = frozenset( + { + "country", + "references", + "references_sha256", + "entities", + "weight_entity", + "annualisation_factors", + "scenario", + } + ) + _optional = frozenset({"facts_sha256", "annualisation_sha256"}) + + def implementation_hash(self) -> str: + return source_hash( + type(self), + graph_inputs, + target_kernels, + *target_kernels._COMPILE_MODULES, + *target_kernels._MATRIX_MODULES, + dependencies=self.capabilities.dependencies, + ) + + def run(self, context: KernelContext) -> KernelResult: + require_params( + context, self.ref, required=self._required, optional=self._optional + ) + _require_output(context, self.ref, "gap", GAP_TYPE) + country = string_param(context, self.ref, "country") + scenario = string_param(context, self.ref, "scenario") + document = canonical_document_param(context, self.ref, "references") + if "annualisation_sha256" in context.params: + sha256_param(context, self.ref, "annualisation_sha256") + references = parse_reference_document(document, country=country) + factors = canonical_document_param(context, self.ref, "annualisation_factors") + names = {reference.name for reference in references} + if set(factors) != names: + raise ValueError( + "Annualisation factors must cover exactly the declared references." + ) + factors = { + name: _number(value, f"annualisation factor {name!r}", positive=True) + for name, value in factors.items() + } + weight_entity = string_param(context, self.ref, "weight_entity") + frame = context_frame( + context, + self.ref, + entities=entities_param(context, self.ref, "entities"), + weight_entity=weight_entity, + ) + source = single_source(context, self.ref) + surface_payload, registry = compile_target_surface( + context.sources[source], + document, + country=country, + references_sha256=sha256_param(context, self.ref, "references_sha256"), + expected_facts_sha256=( + sha256_param(context, self.ref, "facts_sha256") + if "facts_sha256" in context.params + else None + ), + ) + surface = decode_target_surface(surface_payload) + problem = target_kernels._compiled_problem( + frame, registry, weight_entity, self.ref + ) + weights = frame.weights_for(weight_entity) + rows = [] + for index, (spec, reference, trace) in enumerate( + zip(registry.specs, references, surface.trace, strict=True) + ): + factor = factors[spec.name] + with np.errstate(over="raise", invalid="raise"): + annualised = ( + problem.matrix[index : index + 1].toarray().ravel() * factor + ) + estimate = float(np.sum(weights.values * annualised)) + actual = float(spec.value) + rows.append( + { + "name": spec.name, + "period": spec.period, + "entity": spec.entity, + "measure": spec.measure, + "filter": spec.filter, + "metadata": dict(reference.metadata), + "annualisation_factor": factor, + "reference": trace["reference"], + "ledger_fact_key": trace["ledger_fact_key"], + "E": estimate, + "A": actual, + "gap": estimate - actual, + } + ) + payload = canonical_json( + { + "schema_version": 1, + "kind": _GAP_KIND, + "country": country, + "scenario": scenario, + "weight_entity": weight_entity, + "weight_kind": weights.kind.value, + "facts": dict(surface.facts), + "references_sha256": surface.references_sha256, + "gaps": rows, + } + ) + decode_gap(payload) + return KernelResult( + artifacts={"gap": payload}, + receipt={ + "scenario": scenario, + "n_cells": len(rows), + "gap_sha256": sha256_text(payload), + }, + ) + + +class TakeupBandsKernel(KernelBase): + """``takeup.bands@1``: labelled envelopes of matching scenario tables.""" + + ref = "takeup.bands@1" + capabilities = Capabilities( + Determinism.DETERMINISTIC, + numeric=Numeric.BITWISE, + seed_source=SeedSource.NONE, + ) + _required = frozenset({"baseline_scenario", "artifact_scenarios"}) + _optional = frozenset({"scenarios_sha256"}) + + def implementation_hash(self) -> str: + return source_hash( + type(self), graph_inputs, dependencies=self.capabilities.dependencies + ) + + def run(self, context: KernelContext) -> KernelResult: + require_params( + context, self.ref, required=self._required, optional=self._optional + ) + _require_output(context, self.ref, "bands", BANDS_TYPE) + baseline = string_param(context, self.ref, "baseline_scenario") + aliases = canonical_document_param(context, self.ref, "artifact_scenarios") + if "scenarios_sha256" in context.params: + sha256_param(context, self.ref, "scenarios_sha256") + labels = [_text(label, "artifact scenario") for label in aliases.values()] + if not labels or len(set(labels)) != len(labels) or baseline not in labels: + raise ValueError( + "Declared scenarios must be distinct and include the central scenario." + ) + if set(aliases) != set(context.artifacts): + raise ValueError( + "Scenario declarations must cover exactly the consumed artifacts." + ) + documents = {} + identities = {} + for alias, label in aliases.items(): + artifact = context.artifacts[alias] + if artifact.type != GAP_TYPE: + raise ValueError(f"{self.ref} needs gap artifacts for every scenario.") + document = decode_gap(artifact.payload) + if document["scenario"] != label: + raise ValueError( + "A consumed gap artifact disagrees with its declared scenario." + ) + documents[label] = document + identities[label] = sha256_text(artifact.payload) + central = documents[baseline] + cells = { + label: {_cell(row): row for row in document["gaps"]} + for label, document in documents.items() + } + for label, document in documents.items(): + if document["country"] != central["country"] or set(cells[label]) != set( + cells[baseline] + ): + raise ValueError( + "Scenario gap artifacts must cover the same country and table cells." + ) + for cell, row in cells[label].items(): + if any( + row[name] != cells[baseline][cell][name] + for name in _SCENARIO_SEMANTICS + ): + raise ValueError( + "Scenario gap artifacts disagree on a table cell's semantics." + ) + rows = [] + ordered_labels = sorted(documents) + for cell in sorted(cells[baseline], key=lambda value: canonical_json(value)): + central_row = cells[baseline][cell] + for metric in _METRICS: + low = min(ordered_labels, key=lambda label: cells[label][cell][metric]) + high = min( + ordered_labels, key=lambda label: -cells[label][cell][metric] + ) + rows.append( + { + **{name: central_row[name] for name in _ROW_SEMANTICS}, + "metric": metric, + "low": {"value": cells[low][cell][metric], "scenario": low}, + "baseline": { + "value": central_row[metric], + "scenario": baseline, + }, + "high": {"value": cells[high][cell][metric], "scenario": high}, + } + ) + payload = canonical_json( + { + "schema_version": 1, + "kind": _BANDS_KIND, + "country": central["country"], + "baseline_scenario": baseline, + "scenarios": [ + {"scenario": label, "gap_sha256": identities[label]} + for label in ordered_labels + ], + "bands": rows, + } + ) + decode_bands(payload) + return KernelResult( + artifacts={"bands": payload}, + receipt={ + "n_scenarios": len(labels), + "n_cells": len(rows), + "bands_sha256": sha256_text(payload), + }, + ) + + +TAKEUP_GAP = TakeupGapKernel() +TAKEUP_BANDS = TakeupBandsKernel() + + +def register_takeup_kernels(registry: KernelRegistry) -> None: + """Register the two kernels without modifying a global registry.""" + + for kernel in (TAKEUP_GAP, TAKEUP_BANDS): + registry.register(kernel) diff --git a/packages/microcosm-build/tests/engine_free/shared/test_transport_bridge_differential.py b/packages/microcosm-build/tests/engine_free/shared/test_transport_bridge_differential.py new file mode 100644 index 0000000000..8fc1859699 --- /dev/null +++ b/packages/microcosm-build/tests/engine_free/shared/test_transport_bridge_differential.py @@ -0,0 +1,289 @@ +"""Refusals for declared differentials, using invented engine outputs only.""" + +from copy import deepcopy +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest + +from microcosm.frame import EntitySchema, ExportContract, VariableMetadata +from microcosm.frame.rules_kernels import SimulateCounterfactualKernel +from microcosm.graph import KernelRegistry +from tools import transport_bridge_differential as driver + + +class InventedRatesEngine: + """A tiny protocol engine; these rates are invented test data.""" + + def __init__(self): + self.calls = 0 + + def variable_metadata(self, name): + return VariableMetadata(name, "person", "float", "year") + + def variables(self): + return ("income", "rate") + + def entity_schema(self): + return EntitySchema(group_entities=("family",)) + + def materialize(self, bundle, variables, period): + self.calls += 1 + person = bundle.table("person") + values = np.maximum(person["rate"] - person["income"], 0).to_numpy() + return {variable: values for variable in variables} + + def export_contract(self): + return ExportContract.empty() + + def write_dataset(self, bundle, path, period): + return None + + +@pytest.fixture +def differential_case(monkeypatch): + engine = InventedRatesEngine() + kernel = SimulateCounterfactualKernel({"prepared-rates": engine}) + registry = KernelRegistry() + registry.register(kernel) + prepared = SimpleNamespace( + schema=engine.entity_schema(), + kernels=registry, + engine_refs={"rates": "prepared-rates"}, + ) + preparation = [] + + def prepare(*args, **kwargs): + preparation.append("registry") + return prepared + + def unit_rule(data): + preparation.append("unit_rule") + return object() + + monkeypatch.setattr(driver, "build_transport_registry", prepare) + monkeypatch.setattr(driver.BenefitUnitRule, "from_dict", unit_rule) + case = { + "unit_rule": {}, + "rules_bindings": {}, + "tables": { + "person": [ + { + "person_id": 1, + "person_family_id": 100, + "income": 10, + "rate": 50, + }, + { + "person_id": 2, + "person_family_id": 100, + "income": 11, + "rate": 51, + }, + ], + "family": [{"family_id": 100}], + }, + "weights": {"family": {"kind": "design", "values": [1]}}, + "calls": [ + { + "id": "bridge", + "kernel": kernel.ref, + "inputs": {"person": ["income", "rate"]}, + "outputs": [ + {"entity": "family", "column": "bridge", "dtype": "float64"} + ], + "params": { + "engine_ref": "rates", + "period": "2026-27", + "variables": ["weekly_benefit"], + "input_overrides": [ + {"entity": "person", "column": "income", "value": 0} + ], + "output_coefficients": [ + { + "variable": "weekly_benefit", + "coefficient": 1, + "entity": "family", + "column": "bridge", + "aggregation": "sum", + } + ], + "resource_sha256": "a" * 64, + }, + } + ], + "comparisons": [ + { + "call": "bridge", + "entity": "family", + "column": "bridge", + "oracle_pointer": "/weekly", + "tolerance": 0, + } + ], + } + return case, engine, preparation + + +@pytest.mark.parametrize("call_id", ["", None], ids=["empty", "non-string"]) +def test_differential_refuses_empty_call_id_before_preparation( + differential_case, call_id +): + case, engine, preparation = differential_case + case["calls"][0]["id"] = call_id + with pytest.raises(ValueError, match="non-empty call id"): + driver.differential(case, {"weekly": 101}, Path("unused")) + assert preparation == [] + assert engine.calls == 0 + + +def test_differential_refuses_duplicate_bridge_id_with_different_coefficients( + differential_case, +): + case, engine, preparation = differential_case + replacement = deepcopy(case["calls"][0]) + replacement["params"]["output_coefficients"][0]["coefficient"] = 2 + case["calls"].append(replacement) + # The old driver overwrote the coefficient-1 result and falsely passed 202. + with pytest.raises(ValueError, match="duplicate call id.*bridge"): + driver.differential(case, {"weekly": 202}, Path("unused")) + assert preparation == [] + assert engine.calls == 0 + + +@pytest.mark.parametrize("within_call", [False, True], ids=["two-calls", "one-call"]) +def test_differential_refuses_multiple_output_owners_before_preparation( + differential_case, within_call +): + case, engine, preparation = differential_case + if within_call: + case["calls"][0]["outputs"] *= 2 + else: + replacement = deepcopy(case["calls"][0]) + replacement["id"] = "other-bridge" + case["calls"].append(replacement) + with pytest.raises(ValueError, match="output coordinate.*family.*bridge"): + driver.differential(case, {"weekly": 101}, Path("unused")) + assert preparation == [] + assert engine.calls == 0 + + +@pytest.mark.parametrize( + "reference,match", + [ + ({"call": "absent"}, "comparison names missing call"), + ({"column": "absent"}, "comparison names missing output"), + ({"entity": "person"}, "comparison names missing output"), + ], + ids=["missing-call", "missing-column", "wrong-entity"], +) +def test_differential_refuses_missing_comparison_references_before_preparation( + differential_case, reference, match +): + case, engine, preparation = differential_case + case["comparisons"][0].update(reference) + with pytest.raises(ValueError, match=match): + driver.differential(case, {"weekly": 101}, Path("unused")) + assert preparation == [] + assert engine.calls == 0 + + +def test_differential_refuses_comparison_output_owned_by_another_call( + differential_case, +): + case, engine, preparation = differential_case + other = deepcopy(case["calls"][0]) + other["id"] = "other-bridge" + other["outputs"][0]["column"] = "other-output" + other["params"]["output_coefficients"][0]["column"] = "other-output" + case["calls"].append(other) + case["comparisons"][0]["column"] = "other-output" + with pytest.raises(ValueError, match="comparison names missing output"): + driver.differential(case, {"weekly": 101}, Path("unused")) + assert preparation == [] + assert engine.calls == 0 + + +def test_differential_requires_comparisons_before_preparation(differential_case): + case, engine, preparation = differential_case + case["comparisons"] = [] + with pytest.raises(ValueError, match="must declare comparisons"): + driver.differential(case, {}, Path("unused")) + assert preparation == [] + assert engine.calls == 0 + + +@pytest.mark.parametrize("token", ["-1", "01", "+1", " 1", "1 ", "-", "", "\u0661"]) +def test_oracle_pointer_refuses_noncanonical_array_indices(token): + with pytest.raises(ValueError, match="canonical non-negative array index"): + driver._pointer({"items": [10, 3]}, f"/items/{token}") + + +@pytest.mark.parametrize("token", ["2", "123456789012345678901234567890"]) +def test_oracle_pointer_refuses_out_of_bounds_array_indices(token): + with pytest.raises(ValueError, match="array index.*out of bounds"): + driver._pointer({"items": [10, 3]}, f"/items/{token}") + + +@pytest.mark.parametrize("token", ["~", "~2", "name~x", "~01~"]) +def test_oracle_pointer_refuses_invalid_escape_sequences(token): + with pytest.raises(ValueError, match="invalid JSON pointer escape"): + driver._pointer({token: 7}, f"/{token}") + + +@pytest.mark.parametrize( + "pointer,expected", + [ + ("", {"items": [10, 3], "a/b": {"~value": [7]}, "~1": 9, "01": 8, "": 5}), + ("/", 5), + ("/items/0", 10), + ("/items/1", 3), + ("/a~1b/~0value/0", 7), + ("/~01", 9), + ("/01", 8), + ], + ids=[ + "document", + "empty-key", + "zero", + "one", + "escaped", + "decode-once", + "object-key", + ], +) +def test_oracle_pointer_accepts_canonical_indices_and_rfc6901_escapes( + pointer, expected +): + document = {"items": [10, 3], "a/b": {"~value": [7]}, "~1": 9, "01": 8, "": 5} + assert driver._pointer(document, pointer) == expected + + +@pytest.mark.parametrize( + "document,pointer", [({}, "/missing"), ({"item": 7}, "/item/x")] +) +def test_oracle_pointer_refuses_missing_or_scalar_targets(document, pointer): + with pytest.raises(ValueError, match="oracle_pointer does not resolve"): + driver._pointer(document, pointer) + + +@pytest.mark.parametrize( + "comparison,oracle,match", + [ + ({"oracle_pointer": "/items/-1"}, {"items": [101]}, "canonical"), + ({"oracle_pointer": "/absent"}, {"weekly": 101}, "does not resolve"), + ({"tolerance": -1}, {"weekly": 101}, "finite nonnegative tolerance"), + ({}, {"weekly": float("nan")}, "values must be finite"), + ], + ids=["invalid-index", "missing-pointer", "negative-tolerance", "nonfinite-oracle"], +) +def test_differential_refuses_invalid_oracle_comparison_before_preparation( + differential_case, comparison, oracle, match +): + case, engine, preparation = differential_case + case["comparisons"][0].update(comparison) + with pytest.raises(ValueError, match=match): + driver.differential(case, oracle, Path("unused")) + assert preparation == [] + assert engine.calls == 0 diff --git a/packages/microcosm-build/tests/engine_free/shared/test_transport_entitlement_graph.py b/packages/microcosm-build/tests/engine_free/shared/test_transport_entitlement_graph.py new file mode 100644 index 0000000000..0cc5c07c74 --- /dev/null +++ b/packages/microcosm-build/tests/engine_free/shared/test_transport_entitlement_graph.py @@ -0,0 +1,542 @@ +"""Executed entitlement invariants and cache locality on a synthetic graph. + +The composed graph uses the public engine protocol, all four new kernels, +six entitlement scenarios, three independently calibrated variants and a +computed toy WFF tripwire. Every amount belongs to synthetic fixture data. +These tests make no claim about statutory correctness or real donor outputs. +""" + +from __future__ import annotations + +import copy +import json +from dataclasses import replace + +import numpy as np +import pytest +from hypothesis import HealthCheck, example, given, settings +from hypothesis import strategies as st + +from microcosm.build.transport.compose import ( + compose_transport_graph, + validate_transport_activation, +) +from microcosm.build.transport.takeup_kernels import decode_bands, decode_gap +from microcosm.graph import compile_graph, graph_to_json +from microcosm.graph.keys import node_key, source_content_key +from test_support.microcosm_build.transport_entitlement import ( + AS_INPUTS, + AS_OUTPUTS, + TOY_DECLARATIONS, + make_entitlement_fixture, + run_entitlement_graph, +) + +EXECUTED = settings( + max_examples=2, + deadline=None, + database=None, + suppress_health_check=[HealthCheck.function_scoped_fixture], +) +NONNEGATIVE = st.floats( + min_value=0, max_value=1500, allow_nan=False, allow_infinity=False +) + + +def _case_root(tmp_path): + """Give every Hypothesis execution an independent cold input/store tree.""" + root = tmp_path / f"case-{len(tuple(tmp_path.iterdir()))}" + root.mkdir() + return root + + +def _artifact(manifest, store, node, name): + return store.load_bytes(manifest.nodes[node].opaque_artifacts[name]) + + +def _frame(manifest, population="nz.scn.S0"): + return manifest.populations[population] + + +def _misses(manifest): + return {name for name, receipt in manifest.nodes.items() if not receipt.hit} + + +def _keys(manifest): + return {name: receipt.key for name, receipt in manifest.nodes.items()} + + +def _seeded(fixture, overrides): + spec = copy.deepcopy(fixture.spec) + spec["resources"]["entitlement_seed"]["input_overrides"] = [ + {"entity": "family", "column": name, "value": value} + for name, value in overrides.items() + ] + return fixture.reprepare(spec) + + +def _declared_keys(fixture, graph, exported): + compiled = compile_graph(graph) + sources = { + name: source_content_key(name, path) + for name, path in fixture.source_mapping(exported).items() + } + keys = {} + for name in compiled.order: + kernel = fixture.registry.kernels.get(graph.node(name).kernel) + keys[name] = node_key( + compiled, + name, + keys, + kernel.implementation_hash(), + sources, + kernel_capabilities=kernel.capabilities, + ) + return keys + + +@EXECUTED +@example(income=0, costs=1000, assets=0, beneficiary=False) +@given( + income=NONNEGATIVE, costs=NONNEGATIVE, assets=NONNEGATIVE, beneficiary=st.booleans() +) +def test_as_bounds_weighted_total_exact_gap_labelled_bands_and_determinism( + tmp_path, income, costs, assets, beneficiary +): + """Bounds, weighted E, exact E-A, labelled order and deterministic bytes.""" + tmp_path = _case_root(tmp_path) + fixture = make_entitlement_fixture( + tmp_path / "fixture", + seed_overrides={ + "as_income": income, + "as_costs": costs, + "as_assets": assets, + "as_beneficiary": int(beneficiary), + }, + ) + graph = fixture.graph() + assert graph_to_json(graph) == graph_to_json(fixture.graph()) + compiled = compile_graph(graph) + for root in ( + node.id for node in graph.nodes if node.kernel.startswith("calibrate.") + ): + pending, ancestors = [root], set() + while pending: + name = pending.pop() + if name in ancestors: + continue + ancestors.add(name) + pending.extend(compiled.predecessors[name]) + for name in ancestors: + node = graph.node(name) + assert "nz_holdout_facts" not in node.sources + assert ( + node.params.get("engine_ref") + != fixture.registry.engine_refs["entitlement"] + ) + assert not { + column + for item in node.inputs + if item.entity == "family" + for column in item.columns + }.intersection(AS_INPUTS) + first, store, exported = run_entitlement_graph( + tmp_path / "first", fixture, graph=graph + ) + assert _keys(first) == _declared_keys(fixture, graph, exported) + for scenario in fixture.spec["resources"]["scenarios"]["scenarios"]: + name = scenario["id"] + population = ( + f"nz.scn.{name}" + if scenario["tier"] == "entitlement" + else f"nz.v.{name}.housing" + ) + frame = _frame(first, population) + family = frame.table("family") + amounts = family["toy_as_weekly"].to_numpy() + assert np.all(amounts >= 0) + assert np.all(amounts <= family["as_maximum"].to_numpy()) + prefix = ( + f"nz.scn.{name}" if scenario["tier"] == "entitlement" else f"nz.v.{name}" + ) + gap = decode_gap(_artifact(first, store, f"{prefix}.gap", "gap")) + rows = {row["name"]: row for row in gap["gaps"]} + expenditure = rows["toy_as_expenditure"] + factor = fixture.spec["resources"]["entitlement_annualisation"][ + "toy_as_expenditure" + ] + expected = np.sum(frame.resolve_weights("family").values * amounts * factor) + # The kernel aggregates each household's families before weighting; + # float64 reassociation is allowed only on this independent total. + assert expenditure["E"] == pytest.approx(float(expected), rel=1e-12, abs=1e-12) + references = { + row["name"]: row + for row in fixture.spec["resources"]["entitlement_gap"]["target_references"] + } + for row in gap["gaps"]: + assert row["gap"] == row["E"] - row["A"] + assert row["metadata"] == references[row["name"]]["metadata"] + values = ( + amounts if row["measure"] == "toy_as_weekly" else np.ones(len(family)) + ) + if row["filter"] is not None: + values = values * (family[row["filter"]].to_numpy() != 0) + expected_cell = np.sum( + frame.resolve_weights("family").values + * values + * row["annualisation_factor"] + ) + assert row["E"] == pytest.approx(float(expected_cell), rel=1e-12, abs=1e-12) + if row["measure"] == "toy_as_unit_count": + assert row["annualisation_factor"] == 1 + bands = decode_bands(_artifact(first, store, "nz.as.bands", "bands")) + labels = {row["scenario"] for row in bands["scenarios"]} + assert labels == { + row["id"] for row in fixture.spec["resources"]["scenarios"]["scenarios"] + } + for row in bands["bands"]: + assert row["low"]["value"] <= row["baseline"]["value"] <= row["high"]["value"] + assert row["baseline"]["scenario"] == "S0" + for bound in ("low", "baseline", "high"): + assert row[bound]["scenario"] in labels + # Independent cold store: equality checks actual execution, not memoization. + second, second_store, _ = run_entitlement_graph( + tmp_path / "second", fixture, graph=graph, exported=exported + ) + assert _keys(first) == _keys(second) + for name, receipt in first.nodes.items(): + assert receipt.frame_key == second.nodes[name].frame_key + for artifact in receipt.opaque_artifacts: + assert _artifact(first, store, name, artifact) == _artifact( + second, second_store, name, artifact + ) + warm, _, _ = run_entitlement_graph( + tmp_path / "warm", fixture, store=store, exported=exported + ) + assert _misses(warm) == set() + + +@EXECUTED +@given( + excess=st.floats(min_value=1, max_value=1000, allow_nan=False, allow_infinity=False) +) +def test_as_is_zero_for_excess_assets_and_each_statutory_exclusion(tmp_path, excess): + """An encoded asset limit or either independent exclusion zeros AS.""" + tmp_path = _case_root(tmp_path) + fixture = make_entitlement_fixture(tmp_path / "fixture", variants=False) + _, store, exported = run_entitlement_graph(tmp_path / "initial", fixture) + policy = fixture.engines_by_binding["entitlement"].program()["policy"] + for case in ( + {"as_assets": policy["asset_limit"] + excess}, + {"as_assets": 0, "as_s66": 1}, + {"as_assets": 0, "as_s67": 1}, + ): + changed = _seeded(fixture, case) + result, _, _ = run_entitlement_graph( + tmp_path / "excluded", changed, store=store, exported=exported + ) + assert np.all(_frame(result).table("family")["toy_as_weekly"].to_numpy() == 0) + assert result.nodes["nz.calibrate"].hit is True + + +@EXECUTED +@given(income=NONNEGATIVE, increment=NONNEGATIVE) +def test_nonbeneficiary_as_is_nonincreasing_in_income(tmp_path, income, increment): + tmp_path = _case_root(tmp_path) + fixture = make_entitlement_fixture( + tmp_path / "fixture", + variants=False, + seed_overrides={"as_income": income, "as_beneficiary": 0, "as_assets": 0}, + ) + low, store, exported = run_entitlement_graph(tmp_path / "low", fixture) + changed = _seeded( + fixture, {"as_income": income + increment, "as_beneficiary": 0, "as_assets": 0} + ) + high, _, _ = run_entitlement_graph( + tmp_path / "high", changed, store=store, exported=exported + ) + assert np.all( + _frame(high).table("family")["toy_as_weekly"].to_numpy() + <= _frame(low).table("family")["toy_as_weekly"].to_numpy() + ) + + +@EXECUTED +@given(costs=NONNEGATIVE, increment=NONNEGATIVE) +def test_as_is_nondecreasing_in_qualifying_costs_until_the_cap( + tmp_path, costs, increment +): + tmp_path = _case_root(tmp_path) + fixture = make_entitlement_fixture( + tmp_path / "fixture", + variants=False, + seed_overrides={"as_costs": costs, "as_income": 0, "as_assets": 0}, + ) + low, store, exported = run_entitlement_graph(tmp_path / "low", fixture) + changed = _seeded( + fixture, {"as_costs": costs + increment, "as_income": 0, "as_assets": 0} + ) + high, _, _ = run_entitlement_graph( + tmp_path / "high", changed, store=store, exported=exported + ) + before = _frame(low).table("family")["toy_as_weekly"].to_numpy() + after = _frame(high).table("family")["toy_as_weekly"].to_numpy() + assert np.all(after >= before) + assert np.all(after <= _frame(high).table("family")["as_maximum"].to_numpy()) + + +@EXECUTED +@given( + asset_knob=st.floats( + min_value=1, max_value=1000, allow_nan=False, allow_infinity=False + ) +) +def test_a3_s1_knob_misses_only_s1_bands_and_package(tmp_path, asset_knob): + tmp_path = _case_root(tmp_path) + fixture = make_entitlement_fixture(tmp_path / "fixture") + initial, store, exported = run_entitlement_graph(tmp_path / "initial", fixture) + spec = copy.deepcopy(fixture.spec) + spec["resources"]["scenarios"]["scenarios"][1]["knobs"]["input_overrides"][0][ + "value" + ] = asset_knob + changed = fixture.reprepare(spec) + result, _, _ = run_entitlement_graph( + tmp_path / "changed", changed, store=store, exported=exported + ) + expected = {name for name in initial.nodes if name.startswith("nz.scn.S1.")} | { + "nz.as.bands", + "nz.package", + } + assert _misses(result) == expected + assert result.nodes["nz.scn.S1"].hit is True + assert { + name + for name in initial.nodes + if initial.nodes[name].key != result.nodes[name].key + } == expected + + +@EXECUTED +@given(newlines=st.integers(min_value=1, max_value=8)) +def test_a6_holdout_bytes_miss_only_gaps_validation_bands_and_package( + tmp_path, newlines +): + tmp_path = _case_root(tmp_path) + fixture = make_entitlement_fixture(tmp_path / "fixture") + initial, store, exported = run_entitlement_graph(tmp_path / "initial", fixture) + source = fixture.sources["nz_holdout_facts"] + source.write_bytes(source.read_bytes() + b"\n" * newlines) + result, _, _ = run_entitlement_graph( + tmp_path / "changed", fixture, store=store, exported=exported + ) + expected = {name for name in initial.nodes if name.endswith(".gap")} | { + "nz.validate.wff_tripwire", + "nz.as.bands", + "nz.package", + } + assert _misses(result) == expected + assert { + name + for name in initial.nodes + if initial.nodes[name].key != result.nodes[name].key + } == expected + assert result.nodes["nz.calibrate"].hit is True + + +@EXECUTED +@given(assets=NONNEGATIVE) +def test_c2_adding_scenario_rekeys_only_bands_and_package(tmp_path, assets): + tmp_path = _case_root(tmp_path) + fixture = make_entitlement_fixture(tmp_path / "fixture") + initial, store, exported = run_entitlement_graph(tmp_path / "initial", fixture) + spec = copy.deepcopy(fixture.spec) + spec["resources"]["scenarios"]["scenarios"].append( + { + "id": "SX", + "tier": "entitlement", + "knobs": { + "input_overrides": [ + {"entity": "family", "column": "as_assets", "value": assets} + ] + }, + } + ) + changed = fixture.reprepare(spec) + result, _, _ = run_entitlement_graph( + tmp_path / "changed", changed, store=store, exported=exported + ) + assert { + name + for name in initial.nodes + if initial.nodes[name].key != result.nodes[name].key + } == {"nz.as.bands", "nz.package"} + new = { + name + for name in result.nodes + if name == "nz.scn.SX" or name.startswith("nz.scn.SX.") + } + assert _misses(result) == new | {"nz.as.bands", "nz.package"} + assert result.nodes["nz.calibrate"].hit is True + + +def test_bridges_and_as_graph_match_direct_toy_engine_calls(tmp_path): + """A differential independently calls the adapter on the executed frame.""" + fixture = make_entitlement_fixture(tmp_path / "fixture", variants=False) + manifest, store, _ = run_entitlement_graph(tmp_path / "run", fixture) + (tmp_path / "node-durations.json").write_text( + json.dumps( + {name: receipt.wall_time for name, receipt in manifest.nodes.items()}, + sort_keys=True, + ) + ) + frame = _frame(manifest) + table = frame.table("family") + rate_policy = fixture.engines_by_binding["bridge_rates"].program()["policy"] + credit_policy = fixture.engines_by_binding["bridge_credit"].program()["policy"] + expected_base = ( + rate_policy["rate"] + + credit_policy["credit"] * TOY_DECLARATIONS["base_credit_coefficient"] + ) + assert np.all(table["as_base_rate"].to_numpy() == expected_base) + assert table["as_cutout"].to_numpy() == pytest.approx( + rate_policy["rate"] / rate_policy["abatement"], + abs=TOY_DECLARATIONS["solve_tolerance"], + ) + engine = fixture.engines_by_binding["entitlement"] + direct = engine.materialize(frame, tuple(AS_OUTPUTS), "2026-27") + for name in AS_OUTPUTS: + np.testing.assert_array_equal(table[name].to_numpy(), direct[name]) + alternative = ( + _frame(manifest, "nz.scn.S5").table("family")["as_base_rate"].to_numpy() + ) + assert np.all( + alternative + == rate_policy["alternative_rate"] + + credit_policy["credit"] * TOY_DECLARATIONS["base_credit_coefficient"] + ) + comparison = json.loads( + _artifact(manifest, store, "nz.validate.wff_tripwire", "comparison") + ) + validation_frame = _frame(manifest, "nz.validate") + expected_credit = np.sum( + validation_frame.resolve_weights("family").values + * validation_frame.table("family")["toy_wff_credit"].to_numpy() + ) + assert comparison["comparisons"][0]["model"] == pytest.approx( + expected_credit, rel=1e-12 + ) + receipt = json.loads(_artifact(manifest, store, "nz.package", "receipt")) + assert ( + receipt["artifacts"]["bands"]["key"] + == manifest.nodes["nz.as.bands"].opaque_artifacts["bands"] + ) + assert ( + receipt["artifacts"]["wff_tripwire"]["key"] + == manifest.nodes["nz.validate.wff_tripwire"].opaque_artifacts["comparison"] + ) + + +def test_automatic_extension_installation_and_package_edge_restrictions(tmp_path): + """Automatic composition equals explicit installation; receipt edges are narrow.""" + fixture = make_entitlement_fixture(tmp_path / "fixture", variants=False) + explicit = fixture.graph() + automatic = compose_transport_graph( + fixture.spec, replace(fixture.config, extensions=()) + ) + assert graph_to_json(explicit) == graph_to_json(automatic) + extension = fixture.config.extensions[0] + target, edge = extension.package_inputs[0] + skeleton_output = explicit.node("nz.targets.compile").artifact_outputs[0] + invalid = ( + ( + replace(extension, package_inputs=(("nz.export.prepare", edge),)), + "only a skeleton package node", + ), + ( + replace( + extension, package_inputs=((target, replace(edge, name="diagnostics")),) + ), + "new and distinct", + ), + ( + replace( + extension, + package_inputs=( + ( + target, + replace( + edge, + producer="nz.targets.compile", + artifact="surface", + type=skeleton_output.type, + ), + ), + ), + ), + "from extension nodes", + ), + ) + for changed, message in invalid: + with pytest.raises(ValueError, match=message): + compose_transport_graph( + fixture.spec, replace(fixture.config, extensions=(changed,)) + ) + spec = copy.deepcopy(fixture.spec) + rows = spec["resources"]["transport_graph"]["nodes"] + consumer = copy.deepcopy(next(row for row in rows if row["id"] == "nz.package")) + consumer["id"] = "nz.package.consumer" + consumer["artifact_inputs"].append( + {"name": "prior_receipt", "producer": "nz.package", "artifact": "receipt"} + ) + rows.append(consumer) + with pytest.raises(ValueError, match="terminal skeleton package node"): + compose_transport_graph(spec, fixture.config) + + +@pytest.mark.parametrize("template", ["scenario_nodes", "variant_nodes"]) +def test_activation_ignores_unused_absent_tier_template(tmp_path, template): + """Explicit scenario nodes replace their tier's entire unused template.""" + fixture = make_entitlement_fixture(tmp_path / "fixture") + spec = copy.deepcopy(fixture.spec) + document = spec["resources"]["entitlement_graph"] + unused = copy.deepcopy(document["scenario_nodes"][0]) + unused["params"]["unused"] = {"resource": "unused_absent"} + document[template] = [unused] + validate_transport_activation(spec) + assert graph_to_json( + compose_transport_graph(spec, fixture.config) + ) == graph_to_json(fixture.graph()) + + +@pytest.mark.parametrize("location", ["common", "scenario", "variant", "template"]) +@pytest.mark.parametrize("gap", ["null", "path", "name"]) +def test_activation_checks_every_instantiated_entitlement_selection( + tmp_path, location, gap +): + """G6's resource preflight also covers the G7 rows that will be used.""" + fixture = make_entitlement_fixture(tmp_path / "fixture") + spec = copy.deepcopy(fixture.spec) + document = spec["resources"]["entitlement_graph"] + scenarios = spec["resources"]["scenarios"]["scenarios"] + if location == "common": + node = document["nodes"][0] + elif location == "template": + scenarios[0].pop("nodes") + node = document["scenario_nodes"][0] + else: + tier = "entitlement" if location == "scenario" else "calibration" + node = next(row for row in scenarios if row["tier"] == tier)["nodes"][0] + spec["resources"]["selected_probe"] = {"value": None} + selection = {"resource": "selected_probe", "path": ["value"]} + if gap == "path": + selection["path"] = ["absent"] + elif gap == "name": + selection["resource"] = None + node["params"]["probe"] = selection + match = { + "null": "selected_probe.*not activated", + "path": "selected_probe.*no selected path", + "name": "must name a nonempty string", + }[gap] + with pytest.raises(ValueError, match=match): + validate_transport_activation(spec) diff --git a/packages/microcosm-build/tests/engine_free/shared/test_transport_takeup_kernels.py b/packages/microcosm-build/tests/engine_free/shared/test_transport_takeup_kernels.py new file mode 100644 index 0000000000..bd7010a2c5 --- /dev/null +++ b/packages/microcosm-build/tests/engine_free/shared/test_transport_takeup_kernels.py @@ -0,0 +1,541 @@ +"""Gap and band invariants over synthetic populations and published toy facts. + +The weighted-gap path is compared with both a plain family-level sum and the +existing target compiler/matrix estimator. Hypothesis checks exact subtraction, +determinism, ordered labelled bounds, and deterministic tie breaking. Values +and annualisation factors in this file are invented test inputs. +""" + +from __future__ import annotations + +import copy +import json +from dataclasses import replace + +import numpy as np +import pandas as pd +import pytest +from hypothesis import HealthCheck, Phase, example, given, settings +from hypothesis import strategies as st + +from microcosm.build.transport.takeup_kernels import ( + BANDS_TYPE, + GAP_TYPE, + TAKEUP_BANDS, + TAKEUP_GAP, + decode_bands, + decode_gap, + register_takeup_kernels, +) +from microcosm.build.transport.target_kernels import compile_target_surface +from microcosm.calibrate.matrix import build_constraint_matrix +from microcosm.frame import Frame +from microcosm.graph import ( + ArtifactOutput, + ArtifactType, + ArtifactValue, + Kernel, + KernelContext, + KernelRegistry, + Node, + NumericScope, + Slice, +) +from microcosm.graph.canonical import canonical_json +from test_support.microcosm_build.transport_graph import ( + COUNTRY, + ENTITIES, + canonical_text, + reference_document, + reference_row, + toy_fact, + toy_frame, + write_facts, +) + +_SHA = "3" * 64 +_KEY = "a" * 64 +_FACT_NAMES = ("annual_support", "recipient_units") +_FACTOR = 52 +_PHASES = tuple(phase for phase in Phase if phase != Phase.explain) + + +def _frame(values=None, weights=None) -> Frame: + original = toy_frame(weights=weights) + tables = {entity: original.table(entity).copy(deep=True) for entity in ENTITIES} + if values is not None: + tables["family"]["family_support"] = np.asarray(values, dtype=np.float64) + tables["family"]["support_recipient"] = tables["family"]["family_support"] > 0 + return Frame( + tables, + original.schema, + {"household": original.weights_for("household")}, + original.strata.copy(), + metadata=dict(original.metadata), + ) + + +def _references() -> dict: + amount = reference_row( + _FACT_NAMES[0], + entity="family", + measure="family_support", + filter_="support_recipient", + ) + amount["metadata"].update( + { + "decile": "toy_decile", + "family_type": "toy_type", + "receipt_status": "toy_receipt", + } + ) + count = reference_row(_FACT_NAMES[1], entity="family", measure="support_recipient") + return reference_document([amount, count]) + + +def _facts(tmp_path, values=(600, 5)): + return write_facts( + tmp_path / "facts.jsonl", + [ + toy_fact(name, value, entity="family") + for name, value in zip(_FACT_NAMES, values, strict=True) + ], + ) + + +def _gap_node(*, scenario="S0", factors=None, references=None, **params) -> Node: + return Node( + "gap", + TAKEUP_GAP.ref, + population="open", + inputs=( + Slice("person", ("age",)), + Slice("household", ("rent",)), + Slice("family", ("family_support", "support_recipient")), + ), + sources=("facts",), + params={ + "country": COUNTRY, + "references": canonical_text( + _references() if references is None else references + ), + "references_sha256": _SHA, + "entities": ENTITIES, + "weight_entity": "family", + "annualisation_factors": canonical_text( + {_FACT_NAMES[0]: _FACTOR, _FACT_NAMES[1]: 1} + if factors is None + else factors + ), + "annualisation_sha256": _SHA, + "scenario": scenario, + **params, + }, + artifact_outputs=(ArtifactOutput("gap", GAP_TYPE),), + ) + + +def _context(node, *, frame=None, facts=None, artifacts=None) -> KernelContext: + return KernelContext( + node=node, + tables={} + if frame is None + else {entity: frame.table(entity).copy() for entity in ENTITIES}, + weights={} if frame is None else {"family": frame.resolve_weights("family")}, + strata=pd.Series([], dtype=object, name="stratum") + if frame is None + else frame.strata.copy(), + params=node.params, + rng=np.random.default_rng(0), + sources={} if facts is None else {"facts": facts}, + artifacts={} if artifacts is None else artifacts, + ) + + +def _gap( + tmp_path, *, values=None, weights=None, actual=(600, 5), scenario="S0" +) -> bytes: + return TAKEUP_GAP.run( + _context( + _gap_node(scenario=scenario), + frame=_frame(values, weights), + facts=_facts(tmp_path, actual), + ) + ).artifacts["gap"] + + +def _value(payload, kind=GAP_TYPE) -> ArtifactValue: + return ArtifactValue( + payload=payload, type=kind, key=_KEY, producer_key=_KEY, numerics=NumericScope() + ) + + +def _bands_node(labels, *, baseline="S0", **params) -> Node: + return Node( + "bands", + TAKEUP_BANDS.ref, + population="open", + params={ + "baseline_scenario": baseline, + "artifact_scenarios": canonical_text( + {f"gap_{index}": label for index, label in enumerate(labels)} + ), + **params, + }, + artifact_outputs=(ArtifactOutput("bands", BANDS_TYPE),), + ) + + +def _bands(payloads, *, baseline="S0") -> bytes: + labels = [decode_gap(payload)["scenario"] for payload in payloads] + node = _bands_node(labels, baseline=baseline) + return TAKEUP_BANDS.run( + _context( + node, + artifacts={ + f"gap_{index}": _value(payload) + for index, payload in enumerate(payloads) + }, + ) + ).artifacts["bands"] + + +def test_neutral_kernel_contracts_and_idempotent_registration() -> None: + registry = KernelRegistry() + for _ in range(2): + register_takeup_kernels(registry) + for kernel in (TAKEUP_GAP, TAKEUP_BANDS): + assert isinstance(kernel, Kernel) + assert "nz" not in kernel.ref + assert kernel.implementation_hash() == kernel.implementation_hash() + assert kernel.capabilities.consumes_se is False + + +@settings( + max_examples=24, + deadline=None, + phases=_PHASES, + suppress_health_check=[HealthCheck.function_scoped_fixture], +) +@given( + values=st.lists( + st.floats(min_value=0, max_value=2000, allow_nan=False, allow_infinity=False), + min_size=8, + max_size=8, + ), + weights=st.lists( + st.floats(min_value=1, max_value=1000, allow_nan=False, allow_infinity=False), + min_size=6, + max_size=6, + ), + actual=st.floats(min_value=0, max_value=1e8, allow_nan=False, allow_infinity=False), +) +@example(values=[0] * 8, weights=[1] * 6, actual=0) +@example(values=[1] * 8, weights=[1] * 6, actual=1) +def test_weighted_total_is_sum_family_weight_times_annualised_support_and_gap_exact( + tmp_path, values, weights, actual +) -> None: + frame = _frame(values, weights) + payload = _gap(tmp_path, values=values, weights=weights, actual=(actual, 5)) + rows = {row["name"]: row for row in decode_gap(payload)["gaps"]} + expected = float( + np.sum(frame.resolve_weights("family").values * (np.asarray(values) * _FACTOR)) + ) + assert rows["annual_support"]["E"] == expected + assert rows["annual_support"]["A"] == actual + assert rows["annual_support"]["gap"] == expected - actual + count = float( + np.sum(frame.resolve_weights("family").values * (np.asarray(values) > 0)) + ) + assert rows["recipient_units"]["E"] == count + assert all(row["gap"] == row["E"] - row["A"] for row in rows.values()) + + +@settings( + max_examples=12, + deadline=None, + phases=_PHASES, + suppress_health_check=[HealthCheck.function_scoped_fixture], +) +@given( + values=st.lists(st.integers(min_value=0, max_value=2000), min_size=8, max_size=8) +) +def test_gap_differential_matches_existing_target_compiler_and_matrix( + tmp_path, values +) -> None: + frame = _frame(values) + facts = _facts(tmp_path) + _, registry = compile_target_surface( + facts, _references(), country=COUNTRY, references_sha256=_SHA + ) + problem = build_constraint_matrix(frame, registry.to_target_set(), "family") + estimates = problem.estimates(frame.resolve_weights("family").values) + rows = decode_gap(_gap(tmp_path, values=values))["gaps"] + for row, spec, estimate in zip(rows, registry.specs, estimates, strict=True): + assert row["A"] == spec.value + assert row["E"] == pytest.approx( + estimate * row["annualisation_factor"], rel=1e-12 + ) + assert row["reference"] == spec.name + assert row["ledger_fact_key"] == spec.metadata["ledger_fact_key"] + + +@settings( + max_examples=12, + deadline=None, + phases=_PHASES, + suppress_health_check=[HealthCheck.function_scoped_fixture], +) +@given( + values=st.lists(st.integers(min_value=0, max_value=2000), min_size=8, max_size=8) +) +def test_gap_and_bands_deterministic_and_path_independent(tmp_path, values) -> None: + payload = _gap(tmp_path, values=values) + assert payload == _gap(tmp_path, values=values) + assert _bands([payload]) == _bands([payload]) + node = _gap_node() + moved = tmp_path / "moved.jsonl" + moved.write_bytes(_facts(tmp_path).read_bytes()) + result = TAKEUP_GAP.run(_context(node, frame=_frame(values), facts=moved)) + assert result.artifacts["gap"] == payload + assert str(tmp_path) not in payload.decode() + + +def test_reference_metadata_carries_declared_group_labels_and_annualisation( + tmp_path, +) -> None: + document = decode_gap(_gap(tmp_path)) + assert document["scenario"] == "S0" + assert document["weight_entity"] == "family" + row = document["gaps"][0] + assert row["annualisation_factor"] == _FACTOR + assert row["metadata"] == { + "activation_status": "active", + "decile": "toy_decile", + "family_type": "toy_type", + "receipt_status": "toy_receipt", + } + + +@settings( + max_examples=12, + deadline=None, + phases=_PHASES, + suppress_health_check=[HealthCheck.function_scoped_fixture], +) +@given( + central=st.integers(min_value=0, max_value=2000), + alternate=st.integers(min_value=0, max_value=2000), + other=st.integers(min_value=0, max_value=2000), +) +def test_bands_ordered_with_each_bound_naming_its_scenario( + tmp_path, central, alternate, other +) -> None: + payloads = [ + _gap(tmp_path, values=[value] * 8, scenario=label) + for label, value in (("S0", central), ("S1", alternate), ("V1", other)) + ] + bands = decode_bands(_bands(payloads)) + source_rows = { + decode_gap(payload)["scenario"]: { + (row["name"], row["period"]): row for row in decode_gap(payload)["gaps"] + } + for payload in payloads + } + for row in bands["bands"]: + assert row["low"]["value"] <= row["baseline"]["value"] <= row["high"]["value"] + assert row["baseline"]["scenario"] == "S0" + for bound in ("low", "baseline", "high"): + selected = row[bound] + assert ( + selected["value"] + == source_rows[selected["scenario"]][row["name"], row["period"]][ + row["metric"] + ] + ) + values = [ + source_rows[label][row["name"], row["period"]][row["metric"]] + for label in source_rows + ] + assert row["low"]["value"] == min(values) + assert row["high"]["value"] == max(values) + assert _bands(payloads) == _bands(list(reversed(payloads))) + + +@settings( + max_examples=8, + deadline=None, + phases=_PHASES, + suppress_health_check=[HealthCheck.function_scoped_fixture], +) +@given(value=st.integers(min_value=0, max_value=2000)) +def test_equal_bounds_choose_first_scenario_label_deterministically( + tmp_path, value +) -> None: + payloads = [ + _gap(tmp_path, values=[value] * 8, scenario=label) for label in ("Z", "S0", "A") + ] + for row in decode_bands(_bands(payloads))["bands"]: + assert row["low"]["scenario"] == "A" + assert row["high"]["scenario"] == "A" + + +@pytest.mark.parametrize( + "factors", + [ + {}, + {"annual_support": 52}, + {"annual_support": 52, "recipient_units": 1, "extra": 1}, + ], +) +def test_gap_requires_factor_for_exactly_every_reference(tmp_path, factors) -> None: + with pytest.raises(ValueError, match="cover exactly"): + TAKEUP_GAP.run( + _context(_gap_node(factors=factors), frame=_frame(), facts=_facts(tmp_path)) + ) + + +@pytest.mark.parametrize("factor", [0, -1, True, "52"]) +def test_gap_rejects_invalid_annualisation(tmp_path, factor) -> None: + with pytest.raises((ValueError, TypeError), match="number|positive"): + TAKEUP_GAP.run( + _context( + _gap_node(factors={"annual_support": factor, "recipient_units": 1}), + frame=_frame(), + facts=_facts(tmp_path), + ) + ) + + +def test_gap_refuses_uncompiled_placeholder_or_missing_column(tmp_path) -> None: + references = _references() + references["target_references"][0]["metadata"]["activation_status"] = ( + "awaiting_evidence" + ) + with pytest.raises(ValueError, match="placeholder"): + TAKEUP_GAP.run( + _context( + _gap_node(references=references), frame=_frame(), facts=_facts(tmp_path) + ) + ) + references = _references() + references["target_references"][0]["measure"] = "missing_support" + with pytest.raises(ValueError, match="uncompilable|No targets"): + TAKEUP_GAP.run( + _context( + _gap_node(references=references), frame=_frame(), facts=_facts(tmp_path) + ) + ) + + +def test_gap_requires_honest_params_typed_outputs_and_canonical_documents( + tmp_path, +) -> None: + node = _gap_node(extra=True) + with pytest.raises(ValueError, match="unknown parameter"): + TAKEUP_GAP.run(_context(node, frame=_frame(), facts=_facts(tmp_path))) + node = replace(_gap_node(), artifact_outputs=(ArtifactOutput("gap", BANDS_TYPE),)) + with pytest.raises(ValueError, match="typed artifact output"): + TAKEUP_GAP.run(_context(node, frame=_frame(), facts=_facts(tmp_path))) + node = _gap_node() + node = replace( + node, params={**node.params, "references": json.dumps(_references())} + ) + with pytest.raises(ValueError, match="canonical JSON"): + TAKEUP_GAP.run(_context(node, frame=_frame(), facts=_facts(tmp_path))) + + +def test_gap_decoder_refuses_wrong_subtraction_and_noncanonical_bytes(tmp_path) -> None: + payload = _gap(tmp_path) + document = json.loads(payload) + document["gaps"][0]["gap"] += 1 + with pytest.raises(ValueError, match="minus A exactly"): + decode_gap(canonical_json(document)) + with pytest.raises(ValueError, match="noncanonical"): + decode_gap(json.dumps(json.loads(payload)).encode()) + + +def test_bands_refuses_missing_central_scenario_duplicates_and_wrong_artifact_type( + tmp_path, +) -> None: + payload = _gap(tmp_path) + node = _bands_node(["S1"]) + with pytest.raises(ValueError, match="central scenario"): + TAKEUP_BANDS.run(_context(node, artifacts={"gap_0": _value(payload)})) + node = _bands_node(["S0", "S0"]) + with pytest.raises(ValueError, match="distinct"): + TAKEUP_BANDS.run( + _context( + node, artifacts={"gap_0": _value(payload), "gap_1": _value(payload)} + ) + ) + node = _bands_node(["S0"]) + with pytest.raises(ValueError, match="needs gap artifacts"): + TAKEUP_BANDS.run( + _context( + node, artifacts={"gap_0": _value(payload, ArtifactType("wrong", 1))} + ) + ) + + +def test_bands_refuses_mismatched_scenario_country_cells_and_semantics( + tmp_path, +) -> None: + central = _gap(tmp_path) + alternate = decode_gap(_gap(tmp_path, scenario="S1")) + node = _bands_node(["S0", "S1"]) + for field, value, message in ( + ("scenario", "S2", "declared scenario"), + ("country", "yy", "same country"), + ("gaps", alternate["gaps"][:1], "table cells"), + ): + changed = {**alternate, field: value} + with pytest.raises(ValueError, match=message): + TAKEUP_BANDS.run( + _context( + node, + artifacts={ + "gap_0": _value(central), + "gap_1": _value(canonical_json(changed)), + }, + ) + ) + changed = copy.deepcopy(alternate) + changed["gaps"][0]["metadata"]["decile"] = "different" + with pytest.raises(ValueError, match="semantics"): + _bands([central, canonical_json(changed)]) + + +def test_bands_decoder_refuses_unlabelled_reversed_or_unmatched_bounds( + tmp_path, +) -> None: + payload = _bands([_gap(tmp_path)]) + for mutation, message in ( + (lambda row: row["low"].pop("scenario"), "name its scenario"), + (lambda row: row["low"].update(value=row["high"]["value"] + 1), "ordered"), + (lambda row: row["high"].update(scenario="absent"), "undeclared scenario"), + ): + changed = json.loads(payload) + mutation(changed["bands"][0]) + with pytest.raises(ValueError, match=message): + decode_bands(canonical_json(changed)) + + +def test_pinned_facts_hash_is_checked(tmp_path) -> None: + with pytest.raises(ValueError, match="pinned hash"): + TAKEUP_GAP.run( + _context( + _gap_node(facts_sha256="0" * 64), frame=_frame(), facts=_facts(tmp_path) + ) + ) + + +def test_annualisation_resource_digest_must_be_sha256(tmp_path) -> None: + with pytest.raises(ValueError, match="SHA-256"): + TAKEUP_GAP.run( + _context( + _gap_node(annualisation_sha256="unbound"), + frame=_frame(), + facts=_facts(tmp_path), + ) + ) diff --git a/packages/microcosm-frame/src/microcosm/frame/rules_kernels.py b/packages/microcosm-frame/src/microcosm/frame/rules_kernels.py index 54ccb9345a..9c7ad185c3 100644 --- a/packages/microcosm-frame/src/microcosm/frame/rules_kernels.py +++ b/packages/microcosm-frame/src/microcosm/frame/rules_kernels.py @@ -19,17 +19,33 @@ import hashlib import json +import math from collections.abc import Mapping from types import MappingProxyType +import numpy as np +import pandas as pd + 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.bundle import Frame from microcosm.frame.kernels import SimulateRulesKernel from microcosm.frame.rules import RulesEngine -from microcosm.graph import KernelBase, KernelContext, KernelResult, source_hash +from microcosm.graph import ( + ROWS_ALL, + KernelBase, + KernelContext, + KernelResult, + Ownership, + source_hash, +) -__all__ = ["SimulateRulesByRefKernel"] +__all__ = [ + "SimulateRulesByRefKernel", + "SimulateCounterfactualKernel", + "SimulateSolveZeroKernel", +] _PARAMS = frozenset({"engine_ref", "period", "variables"}) @@ -158,3 +174,502 @@ def _describe(engine_ref: str) -> 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})" + + +class SimulateCounterfactualKernel(SimulateRulesByRefKernel): + """Combine declared engine outputs after declared input overrides. + + The primary component uses ``engine_ref``, ``period`` and ``variables``. + ``input_overrides`` is a JSON list of records with ``entity``, ``column`` + and either ``value`` or ``source_entity``/``source_column``. A source group + column broadcasts to persons through structural membership. An optional + ``mask`` record (``entity``, ``column``, ``equals``) limits an override. + Override columns must be inputs of that component's engine. + + ``components`` is an optional JSON list of additional components, each + declaring those four primary parameters (variables and overrides as JSON + arrays). Components evaluate independently on copies of the original + frame, allowing modules with distinct input worlds to share a bridge. + + ``output_coefficients`` is a JSON list of terms. Each declares an engine + ``variable``, numeric ``coefficient``, and target ``entity``/``column``. + ``engine_ref`` defaults to the primary reference. A person output mapped + to a group must declare ``aggregation`` as ``sum``, ``max`` or ``first``; + group outputs may broadcast to persons. ``first`` selects the member with + the lowest structural person id. Terms may declare the same mask shape, + evaluated on the original frame, to select receipt or eligibility flags. + Terms targeting the same column add in their declared order. + + Every owned output is a produced float64 column on all rows. Parameters + contain all policy coefficients and overrides; ``resource_sha256`` may + bind the declaring spec resource. Neither engine bindings nor spec + fingerprints enter the implementation hash. + """ + + ref = "simulate.counterfactual@1" + + def run(self, context: KernelContext) -> KernelResult: + components, terms, frame = self._prepare(context) + values = self._evaluate(frame, components, terms) + return self._result(context, frame, values) + + def _prepare(self, context: KernelContext, *, solver: bool = False): + required = { + "engine_ref", + "period", + "variables", + "input_overrides", + "output_coefficients", + } + if solver: + required |= {"solve_input", "bracket", "tolerance", "iterations"} + allowed = required | {"components", "resource_sha256"} + if not required <= set(context.params) or set(context.params) - allowed: + raise ValueError( + f"{self.ref} parameters require {sorted(required)!r} and allow " + f"only {sorted(allowed)!r}." + ) + if "resource_sha256" in context.params: + digest = context.params["resource_sha256"] + if ( + not isinstance(digest, str) + or len(digest) != 64 + or any(char not in "0123456789abcdef" for char in digest) + ): + raise ValueError("resource_sha256 must be a lowercase SHA-256.") + primary = { + key: context.params[key] for key in ("engine_ref", "period", "variables") + } + primary["input_overrides"] = _json_list( + context.params["input_overrides"], "input_overrides" + ) + additional = _json_list(context.params.get("components", "[]"), "components") + components = [primary, *additional] + references = set() + schema = None + for component in components: + if set(component) != { + "engine_ref", + "period", + "variables", + "input_overrides", + }: + raise ValueError( + "Each component requires engine_ref, period, variables and input_overrides." + ) + ref = component["engine_ref"] + if not isinstance(ref, str) or ref not in self._delegates: + raise ValueError( + f"{self.ref} has no engine bound to engine_ref {ref!r}." + ) + if ref in references: + raise ValueError("Components must not repeat engine_ref.") + references.add(ref) + period = component["period"] + if ( + isinstance(period, bool) + or not isinstance(period, int | str) + or period == "" + ): + raise TypeError("Component period must be an int or non-empty string.") + variables = component["variables"] + if ( + not isinstance(variables, list | tuple) + or not variables + or any(not isinstance(item, str) or not item for item in variables) + or len(set(variables)) != len(variables) + ): + raise ValueError( + "Component variables must be distinct non-empty names." + ) + if not isinstance(component["input_overrides"], list): + raise TypeError("Component input_overrides must be a list.") + engine_schema = self._delegates[ref]._engine.entity_schema() + if schema is not None and schema != engine_schema: + raise ValueError("Bridge components must share an entity schema.") + schema = engine_schema + tables = SimulateRulesKernel._tables(context, schema) + frame = Frame( + tables=tables, + schema=schema, + weights=dict(context.weights), + strata=context.strata, + ) + terms = _json_list(context.params["output_coefficients"], "output_coefficients") + if not terms: + raise ValueError("output_coefficients must not be empty.") + outputs = set() + for term in terms: + _fields( + term, + {"variable", "coefficient", "entity", "column"}, + {"engine_ref", "aggregation", "mask"}, + "output coefficient", + ) + ref = term.get("engine_ref", primary["engine_ref"]) + component = next( + (item for item in components if item["engine_ref"] == ref), None + ) + if component is None or term["variable"] not in component["variables"]: + raise ValueError( + "Output coefficient must name a requested component variable." + ) + _number(term["coefficient"], "coefficient") + frame.schema.entity_id_column(term["entity"]) + if not isinstance(term["column"], str) or not term["column"]: + raise ValueError("Output column must be a non-empty string.") + outputs.add((term["entity"], term["column"])) + declared = {(owned.entity, owned.column) for owned in context.node.outputs} + if declared != outputs: + raise ValueError("Declared outputs do not match output_coefficients.") + if any( + owned.rows != ROWS_ALL + or owned.ownership is not Ownership.PRODUCED + or owned.dtype != "float64" + for owned in context.node.outputs + ): + raise ValueError( + "Bridge outputs must be produced float64 columns on all rows." + ) + return components, terms, frame + + def _evaluate(self, frame, components, terms, *, trial=None, solve_input=None): + materialized = {} + for component in components: + ref = component["engine_ref"] + engine = self._delegates[ref]._engine + tables = {entity: frame.table(entity).copy() for entity in frame.entities} + inputs = set(engine.variables()) + overridden = set() + for override in component["input_overrides"]: + _fields( + override, + {"entity", "column"}, + {"value", "source_entity", "source_column", "mask"}, + "input override", + ) + entity, column = override["entity"], override["column"] + frame.schema.entity_id_column(entity) + if column not in inputs or (entity, column) in overridden: + raise ValueError( + "Overrides must name distinct engine input columns." + ) + overridden.add((entity, column)) + if column in _structural_columns(frame.schema, entity): + raise ValueError( + "Input overrides cannot change structural columns." + ) + if "value" in override: + if "source_entity" in override or "source_column" in override: + raise ValueError( + "An override must declare either value or source column." + ) + value = override["value"] + if not isinstance(value, bool | int | float | str) or ( + isinstance(value, float) and not math.isfinite(value) + ): + raise ValueError("Override values must be finite scalar data.") + values = np.full(len(tables[entity]), value) + else: + if not {"source_entity", "source_column"} <= set(override): + raise ValueError( + "A column override needs source_entity and source_column." + ) + source_entity = override["source_entity"] + source = frame.table(source_entity)[ + override["source_column"] + ].to_numpy() + values = _map_values( + frame, source, source_entity, entity, aggregation=None + ) + if "mask" in override: + if column not in tables[entity]: + raise ValueError( + "A masked override requires the existing input column." + ) + mask = _mask(frame, override["mask"], entity) + values = np.where(mask, values, tables[entity][column].to_numpy()) + tables[entity][column] = values + if trial is not None and solve_input["engine_ref"] == ref: + entity, column = solve_input["entity"], solve_input["column"] + if column not in inputs: + raise ValueError("solve_input must name an engine input column.") + tables[entity][column] = _map_values( + frame, trial[1], trial[0], entity, aggregation=None + ) + counterfactual = Frame( + tables, + frame.schema, + { + entity: frame.weights_for(entity) + for entity in frame.weighted_entities + }, + strata=frame.strata, + ) + variables = tuple(component["variables"]) + raw = engine.materialize(counterfactual, variables, component["period"]) + for variable in variables: + metadata = engine.variable_metadata(variable) + if variable not in raw: + raise ValueError( + f"Engine did not return requested variable {variable!r}." + ) + values = np.asarray(raw[variable], dtype=np.float64) + if values.shape != (len(frame.table(metadata.entity)),): + raise ValueError(f"Engine output {variable!r} has the wrong shape.") + if not np.all(np.isfinite(values)): + raise ValueError(f"Engine output {variable!r} must be finite.") + materialized[(ref, variable)] = metadata.entity, values + output = {} + for term in terms: + ref = term.get("engine_ref", components[0]["engine_ref"]) + entity, values = materialized[(ref, term["variable"])] + if "mask" in term: + values = np.where(_mask(frame, term["mask"], entity), values, 0) + mapped = _map_values( + frame, + values, + entity, + term["entity"], + aggregation=term.get("aggregation"), + ) + key = term["entity"], term["column"] + contribution = mapped * _number(term["coefficient"], "coefficient") + output[key] = ( + output.get(key, np.zeros(len(mapped), dtype=np.float64)) + contribution + ) + if not np.all(np.isfinite(output[key])): + raise ValueError("Bridge coefficient sum must be finite.") + return output + + def _result(self, context, frame, values, **receipt): + columns = {} + for (entity, column), data in values.items(): + id_column = frame.schema.entity_id_column(entity) + columns[(entity, column)] = pd.Series( + data, + index=pd.Index(frame.table(entity)[id_column], name=id_column), + name=column, + copy=False, + ) + return KernelResult( + columns=columns, + receipt={ + "engine_ref": context.params["engine_ref"], + "period": context.params["period"], + "variables": context.params["variables"], + "components": context.params.get("components", "[]"), + **receipt, + }, + ) + + +class SimulateSolveZeroKernel(SimulateCounterfactualKernel): + """Find the first nonpositive engine residual by fixed vector bisection. + + The counterfactual contract describes exactly one residual/output column. + ``solve_input`` is a JSON record naming ``entity`` and input ``column``; + its optional ``engine_ref`` defaults to the primary component. Trial values + on the output entity broadcast to that input entity. ``bracket`` is a + finite increasing pair, ``tolerance`` bounds the final income interval + and positive residual at the returned answer, and ``iterations`` is a + fixed positive integer sufficient for the interval bound. Each output + row's residual must be independent of other rows' trial inputs and must + be non-increasing in its own trial input on the bracket (the caller's rules + contract). Coupled rows are unsupported. A positive upper endpoint is + refused, and an already nonpositive lower endpoint returns that endpoint. + All rows perform every iteration, without value-dependent early stopping. + Returned answers are + evaluated once more; a positive residual beyond ``tolerance`` is refused. + Successful runs perform exactly ``iterations + 3`` residual evaluations, + each evaluating every component once. Outputs are the upper endpoints, + so a benefit that floors at zero yields its cutout. An exact-zero residual + does not replace a nonzero answer with zero. + """ + + ref = "simulate.solve_zero@1" + + def run(self, context: KernelContext) -> KernelResult: + components, terms, frame = self._prepare(context, solver=True) + output_keys = {(term["entity"], term["column"]) for term in terms} + if len(output_keys) != 1: + raise ValueError("solve_zero requires exactly one output column.") + key = next(iter(output_keys)) + solve_input = _json_record(context.params["solve_input"], "solve_input") + _fields(solve_input, {"entity", "column"}, {"engine_ref"}, "solve_input") + solve_input = {"engine_ref": components[0]["engine_ref"], **solve_input} + if solve_input["engine_ref"] not in {item["engine_ref"] for item in components}: + raise ValueError("solve_input must name a declared component.") + frame.schema.entity_id_column(solve_input["entity"]) + if solve_input["column"] in _structural_columns( + frame.schema, solve_input["entity"] + ): + raise ValueError("solve_input cannot change structural columns.") + bracket = context.params["bracket"] + if not isinstance(bracket, tuple) or len(bracket) != 2: + raise ValueError("bracket must be an increasing finite pair.") + lower, upper = (_number(value, "bracket endpoint") for value in bracket) + tolerance = _number(context.params["tolerance"], "tolerance") + iterations = context.params["iterations"] + if ( + lower >= upper + or tolerance <= 0 + or not math.isfinite(upper - lower) + or isinstance(iterations, bool) + or not isinstance(iterations, int) + or iterations <= 0 + ): + raise ValueError( + "solve_zero requires an increasing bracket, positive tolerance and iterations." + ) + if math.ldexp(upper - lower, -iterations) > tolerance: + raise ValueError( + "iterations do not achieve the declared bracket tolerance." + ) + lo = np.full(len(frame.table(key[0])), lower, dtype=np.float64) + hi = np.full(len(lo), upper, dtype=np.float64) + + def residual(trial): + return self._evaluate( + frame, components, terms, trial=(key[0], trial), solve_input=solve_input + )[key] + + lower_values = residual(lo) + if np.any(residual(hi) > 0): + raise ValueError( + "solve_zero bracket upper endpoint has a positive residual." + ) + already_zero = lower_values <= 0 + for _ in range(iterations): + mid = lo + (hi - lo) / 2 + positive = residual(mid) > 0 + lo = np.where(positive, mid, lo) + hi = np.where(positive, hi, mid) + if np.any((hi - lo) > tolerance): + raise ValueError("Float64 bisection cannot meet the declared tolerance.") + answer = np.where(already_zero, lower, hi) + if np.any(residual(answer) > tolerance): + raise ValueError( + "solve_zero final answer has a positive residual beyond the " + "declared tolerance; output rows must be independent." + ) + return self._result( + context, + frame, + {key: answer}, + bracket=bracket, + tolerance=tolerance, + iterations=iterations, + evaluations=iterations + 3, + ) + + +def _json_list(value, name): + payload = _json_value(value, name) + if not isinstance(payload, list) or any( + not isinstance(item, dict) for item in payload + ): + raise TypeError(f"{name} must be a JSON list of records.") + return payload + + +def _json_record(value, name): + payload = _json_value(value, name) + if not isinstance(payload, dict): + raise TypeError(f"{name} must be a JSON record.") + return payload + + +def _json_value(value, name): + if not isinstance(value, str): + raise TypeError(f"{name} must be a JSON string.") + try: + return json.loads(value, parse_constant=lambda item: _invalid_constant(item)) + except ValueError as error: + raise ValueError(f"{name} is not valid finite JSON.") from error + + +def _invalid_constant(value): + raise ValueError(f"Nonfinite JSON constant {value!r}.") + + +def _fields(record, required, optional, name): + if ( + not isinstance(record, dict) + or not required <= set(record) + or set(record) - required - optional + ): + raise ValueError( + f"Invalid {name} fields; require {sorted(required)!r}, allow {sorted(optional)!r}." + ) + + +def _number(value, name): + if isinstance(value, bool) or not isinstance(value, int | float): + raise ValueError(f"{name} must be a finite number.") + try: + number = float(value) + except OverflowError as error: + raise ValueError(f"{name} must be a finite number.") from error + if not math.isfinite(number): + raise ValueError(f"{name} must be a finite number.") + return number + + +def _structural_columns(schema, entity): + columns = {schema.entity_id_column(entity)} + if entity == schema.person_entity: + columns.update( + schema.membership_column(group) for group in schema.group_entities + ) + return columns + + +def _mask(frame, record, entity): + _fields(record, {"entity", "column", "equals"}, set(), "mask") + match = record["equals"] + if not isinstance(match, bool | int | float | str) or ( + isinstance(match, float) and not math.isfinite(match) + ): + raise ValueError("Mask equals must be finite scalar data.") + source_entity = record["entity"] + values = frame.table(source_entity)[record["column"]].to_numpy() == record["equals"] + return _map_values(frame, values, source_entity, entity, aggregation=None) + + +def _map_values(frame, values, source, target, *, aggregation): + """Map arrays through declared partitions, retaining the target id order.""" + schema = frame.schema + schema.entity_id_column(source) + schema.entity_id_column(target) + if source == target: + if aggregation is not None: + raise ValueError("Same-entity terms cannot declare aggregation.") + return values + person = schema.person_entity + if source in schema.group_entities and target == person: + if aggregation is not None: + raise ValueError("Broadcast terms cannot declare aggregation.") + index = pd.Index(frame.table(source)[schema.id_column(source)]) + positions = index.get_indexer( + frame.table(person)[schema.membership_column(source)] + ) + return values[positions] + if source == person and target in schema.group_entities: + if aggregation not in {"sum", "max", "first"}: + raise ValueError( + "Person-to-group terms require sum, max or first aggregation." + ) + persons = frame.table(person) + membership = schema.membership_column(target) + series = pd.Series(values, index=pd.Index(persons[membership])) + if aggregation == "first": + order = np.argsort( + persons[schema.person_id_column].to_numpy(), kind="stable" + ) + series = series.iloc[order] + grouped = series.groupby(level=0, sort=False).agg(aggregation) + return grouped.reindex(frame.table(target)[schema.id_column(target)]).to_numpy() + raise ValueError( + "Bridge mapping requires the same entity or a person/group partition." + ) diff --git a/packages/microcosm-frame/tests/engine_free/shared/test_counterfactual_kernels.py b/packages/microcosm-frame/tests/engine_free/shared/test_counterfactual_kernels.py new file mode 100644 index 0000000000..581e80ae7c --- /dev/null +++ b/packages/microcosm-frame/tests/engine_free/shared/test_counterfactual_kernels.py @@ -0,0 +1,617 @@ +"""Engine-free bridge properties and differentials against direct engines.""" + +import json +from collections.abc import Sequence +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pandas as pd +import pytest +from hypothesis import example, given, settings +from hypothesis import strategies as st + +from microcosm.frame import ( + EntitySchema, + ExportContract, + Frame, + VariableMetadata, + WeightKind, + Weights, +) +from microcosm.frame.rules_kernels import ( + SimulateCounterfactualKernel, + SimulateSolveZeroKernel, +) +from microcosm.graph import KernelRegistry, Node, Owned, Slice +from test_support.microcosm_frame.kernels import _context +from tools import transport_bridge_differential as differential_tool + +SCHEMA = EntitySchema(group_entities=("family",)) +RATE_REF = "toy-rates" +CREDIT_REF = "toy-credit" +POPULATIONS = st.lists( + st.tuples( + st.integers(0, 100), + st.integers(1, 100), + st.integers(1, 100), + st.booleans(), + ), + min_size=1, + max_size=5, +) + + +class RatesEngine: + """Rates live in this toy engine; bridge code has no policy formula.""" + + def __init__(self): + self.calls = 0 + + def variable_metadata(self, name): + if name != "weekly_benefit": + raise ValueError(name) + return VariableMetadata(name, "person", "float", "year") + + def variables(self): + return ("income", "rate", "reduction", "receives") + + def entity_schema(self): + return SCHEMA + + def materialize(self, bundle, variables: Sequence[str], period): + self.calls += 1 + person = bundle.table("person") + values = np.maximum( + person["rate"].to_numpy() + - person["income"].to_numpy() * person["reduction"].to_numpy(), + 0, + ).astype(np.float64) + return {variable: values for variable in variables} + + def export_contract(self): + return ExportContract.empty() + + def write_dataset(self, bundle, path: str | Path, period): + return None + + +class CreditEngine(RatesEngine): + def variable_metadata(self, name): + if name != "eldest_credit": + raise ValueError(name) + return VariableMetadata(name, "family", "float", "year") + + def variables(self): + return ("credit",) + + def materialize(self, bundle, variables, period): + self.calls += 1 + return { + variable: bundle.table("family")["credit"].to_numpy(dtype=np.float64) + for variable in variables + } + + +def _json(value): + return json.dumps(value, sort_keys=True, separators=(",", ":")) + + +def _frame(population): + persons, families = [], [] + for offset, (income, rate, reduction, receives) in enumerate(population): + group = 100 + offset + # Reverse the person ids within a group: "first" must use ids, not + # incidental input row order. Group tables retain Frame's sorted ids. + for member in (1, 0): + persons.append( + { + "person_id": 10 * group + member, + "person_family_id": group, + "income": float(income + member), + "rate": float(rate + member), + "reduction": float(reduction), + "receives": receives if member == 0 else False, + } + ) + families.append({"family_id": group, "credit": float(rate * 2)}) + return Frame( + {"person": pd.DataFrame(persons), "family": pd.DataFrame(families)}, + SCHEMA, + {"family": Weights(np.ones(len(families)), WeightKind.DESIGN)}, + ) + + +def _node(kernel, **changes): + params = { + "engine_ref": RATE_REF, + "period": "2026-27", + "variables": ("weekly_benefit",), + "input_overrides": _json( + [{"entity": "person", "column": "income", "value": 0}] + ), + "output_coefficients": _json( + [ + { + "variable": "weekly_benefit", + "coefficient": 1, + "entity": "family", + "column": "bridge", + "aggregation": "sum", + } + ] + ), + "resource_sha256": "a" * 64, + } + if isinstance(kernel, SimulateSolveZeroKernel): + params.update( + input_overrides="[]", + solve_input=_json({"entity": "person", "column": "income"}), + bracket=(0, 256), + tolerance=1e-7, + iterations=34, + ) + params.update(changes) + return Node( + "bridge", + kernel.ref, + inputs=( + Slice("person", ("income", "rate", "reduction", "receives")), + Slice("family", ("credit",)), + ), + outputs=(Owned("family", "bridge", "float64"),), + params=params, + ) + + +def _run(kernel, frame, node): + return kernel.run(_context(frame, node, weighted_entities=("family",))) + + +@settings(max_examples=25, deadline=None) +@given(POPULATIONS, st.sampled_from(("sum", "max", "first"))) +def test_counterfactual_matches_direct_engine_and_declared_reduction( + population, aggregation +): + frame = _frame(population) + original = frame.table("person").copy(deep=True) + engine = RatesEngine() + kernel = SimulateCounterfactualKernel({RATE_REF: engine}) + term = { + "variable": "weekly_benefit", + "coefficient": 2, + "entity": "family", + "column": "bridge", + "aggregation": aggregation, + } + node = _node(kernel, output_coefficients=_json([term])) + result = _run(kernel, frame, node) + tables = {entity: frame.table(entity).copy() for entity in frame.entities} + tables["person"]["income"] = 0 + direct = engine.materialize( + Frame(tables, SCHEMA, {"family": frame.weights_for("family")}), + ("weekly_benefit",), + "2026-27", + )["weekly_benefit"] + expected = [] + for group in frame.table("family")["family_id"]: + positions = np.flatnonzero(original["person_family_id"].to_numpy() == group) + if aggregation == "first": + position = min(positions, key=lambda item: original["person_id"].iloc[item]) + expected.append(direct[position] * 2) + else: + expected.append(getattr(direct[positions], aggregation)() * 2) + series = result.columns[("family", "bridge")] + np.testing.assert_array_equal(series.to_numpy(), expected) + assert series.index.tolist() == frame.table("family")["family_id"].tolist() + assert series.dtype == np.dtype("float64") + pd.testing.assert_frame_equal(frame.table("person"), original) + repeated = _run(kernel, frame, node) + assert ( + repeated.columns[("family", "bridge")].to_numpy().tobytes() + == series.to_numpy().tobytes() + ) + + +@settings(max_examples=20, deadline=None) +@given(POPULATIONS) +@example([(0, 50, 1, True)]) +def test_multiple_engines_masks_and_group_input_broadcast_equal_direct_calls( + population, +): + frame = _frame(population) + rates, credit = RatesEngine(), CreditEngine() + kernel = SimulateCounterfactualKernel({RATE_REF: rates, CREDIT_REF: credit}) + primary_overrides = [ + { + "entity": "person", + "column": "rate", + "source_entity": "family", + "source_column": "credit", + "mask": {"entity": "person", "column": "receives", "equals": True}, + } + ] + terms = [ + { + "variable": "weekly_benefit", + "coefficient": 1, + "entity": "family", + "column": "bridge", + "aggregation": "sum", + "mask": {"entity": "person", "column": "receives", "equals": True}, + }, + { + "engine_ref": CREDIT_REF, + "variable": "eldest_credit", + "coefficient": 0.5, + "entity": "family", + "column": "bridge", + }, + ] + node = _node( + kernel, + input_overrides=_json(primary_overrides), + output_coefficients=_json(terms), + components=_json( + [ + { + "engine_ref": CREDIT_REF, + "period": "2026-27", + "variables": ["eldest_credit"], + "input_overrides": [ + {"entity": "family", "column": "credit", "value": 4} + ], + } + ] + ), + ) + result = _run(kernel, frame, node) + person = frame.table("person").copy() + family = frame.table("family").copy() + lookup = family.set_index("family_id")["credit"] + person.loc[person["receives"], "rate"] = ( + person.loc[person["receives"], "person_family_id"].map(lookup).to_numpy() + ) + direct_frame = Frame( + {"person": person, "family": family}, + SCHEMA, + {"family": frame.weights_for("family")}, + ) + direct_rate = rates.materialize(direct_frame, ("weekly_benefit",), "2026-27")[ + "weekly_benefit" + ] + family["credit"] = 4 + direct_credit = credit.materialize( + Frame( + {"person": person, "family": family}, + SCHEMA, + {"family": frame.weights_for("family")}, + ), + ("eldest_credit",), + "2026-27", + )["eldest_credit"] + expected = [] + for position, group in enumerate(family["family_id"]): + mask = (person["person_family_id"].to_numpy() == group) & person[ + "receives" + ].to_numpy() + expected.append(direct_rate[mask].sum() + direct_credit[position] * 0.5) + np.testing.assert_array_equal(result.columns[("family", "bridge")], expected) + assert frame.table("family")["credit"].tolist() == [ + item[1] * 2 for item in population + ] + + +@settings(max_examples=15, deadline=None) +@given(POPULATIONS) +def test_fixed_vector_bisection_matches_analytic_cutout_and_is_deterministic( + population, +): + frame = _frame(population) + engine = RatesEngine() + kernel = SimulateSolveZeroKernel({RATE_REF: engine}) + node = _node(kernel) + result = _run(kernel, frame, node) + cutouts = result.columns[("family", "bridge")].to_numpy() + person = frame.table("person") + direct = person["rate"].to_numpy() / person["reduction"].to_numpy() + expected = np.asarray( + [ + direct[person["person_family_id"].to_numpy() == group].max() + for group in frame.table("family")["family_id"] + ] + ) + assert np.all(cutouts >= expected - np.spacing(expected)) + assert np.all(cutouts - expected <= node.params["tolerance"]) + assert engine.calls == node.params["iterations"] + 3 + assert result.receipt["evaluations"] == engine.calls + repeated = _run(kernel, frame, node) + assert ( + cutouts.tobytes() == repeated.columns[("family", "bridge")].to_numpy().tobytes() + ) + + +def test_cutout_override_uses_declared_reference_case_instead_of_current_benefit(): + frame = _frame([(10, 5, 1, True)]) + kernel = SimulateSolveZeroKernel({RATE_REF: RatesEngine()}) + node = _node( + kernel, + input_overrides=_json( + [ + {"entity": "person", "column": "rate", "value": 100}, + {"entity": "person", "column": "reduction", "value": 2}, + ] + ), + ) + result = _run(kernel, frame, node) + assert result.columns[("family", "bridge")].iloc[0] == 50 + + +def test_nonpositive_lower_endpoint_is_returned_with_fixed_evaluation_count(): + frame = _frame([(0, 1, 1, False)]) + engine = RatesEngine() + kernel = SimulateSolveZeroKernel({RATE_REF: engine}) + result = _run(kernel, frame, _node(kernel, bracket=(10, 256))) + assert result.columns[("family", "bridge")].iloc[0] == 10 + assert engine.calls == 37 + assert result.receipt["evaluations"] == engine.calls + + +def test_solver_evaluates_final_nonzero_answer_on_zero_plateau(): + class RecordingEngine(RatesEngine): + def materialize(self, bundle, variables, period): + self.last_inputs = bundle.table("person")["income"].to_numpy().copy() + self.last_residuals = super().materialize(bundle, variables, period) + return self.last_residuals + + frame = _frame([(10, 5, 1, True)]) + engine = RecordingEngine() + kernel = SimulateSolveZeroKernel({RATE_REF: engine}) + node = _node( + kernel, + input_overrides=_json( + [ + {"entity": "person", "column": "rate", "value": 100}, + {"entity": "person", "column": "reduction", "value": 2}, + ] + ), + ) + result = _run(kernel, frame, node) + answer = result.columns[("family", "bridge")].iloc[0] + assert answer == 50 + np.testing.assert_array_equal(engine.last_inputs, np.full(2, answer)) + np.testing.assert_array_equal(engine.last_residuals["weekly_benefit"], [0, 0]) + assert engine.calls == node.params["iterations"] + 3 + assert result.receipt["evaluations"] == engine.calls + + +def test_solver_refuses_coupled_rows_with_positive_final_residual(): + class CoupledEngine(RatesEngine): + def materialize(self, bundle, variables, period): + self.calls += 1 + person = bundle.table("person") + family = bundle.table("family") + first_incomes = person.groupby("person_family_id")["income"].first() + total_income = first_incomes.sum() + residuals = pd.Series( + np.maximum(np.asarray([100, 50]) - total_income, 0), + index=family["family_id"], + ) + values = person["person_family_id"].map(residuals).to_numpy() + return {variable: values for variable in variables} + + frame = _frame([(0, 100, 1, True), (0, 50, 1, True)]) + engine = CoupledEngine() + kernel = SimulateSolveZeroKernel({RATE_REF: engine}) + node = _node(kernel) + with pytest.raises(ValueError, match="final answer.*positive residual"): + _run(kernel, frame, node) + assert engine.calls == node.params["iterations"] + 3 + + +@pytest.mark.parametrize( + "changes,match", + [ + ({"bracket": (0, 1)}, "positive residual"), + ({"iterations": 1}, "tolerance"), + ({"tolerance": 0}, "positive tolerance"), + ( + {"solve_input": _json({"entity": "person", "column": "not_an_input"})}, + "engine input", + ), + ], +) +def test_solver_refuses_unbracketed_or_inadequately_declared_roots(changes, match): + frame = _frame([(10, 50, 1, True)]) + kernel = SimulateSolveZeroKernel({RATE_REF: RatesEngine()}) + with pytest.raises(ValueError, match=match): + _run(kernel, frame, _node(kernel, **changes)) + + +@pytest.mark.parametrize( + "changes,match", + [ + ({"engine_ref": "missing"}, "no engine bound"), + ( + { + "input_overrides": _json( + [{"entity": "person", "column": "not_an_input", "value": 0}] + ) + }, + "engine input", + ), + ( + { + "output_coefficients": _json( + [ + { + "variable": "weekly_benefit", + "coefficient": 1, + "entity": "family", + "column": "bridge", + } + ] + ) + }, + "aggregation", + ), + ({"resource_sha256": "no digest"}, "SHA-256"), + ({"components": "[{}]"}, "Each component"), + ({"output_coefficients": "NaN"}, "finite JSON"), + ], +) +def test_bridge_refuses_unbound_or_undeclared_inputs(changes, match): + frame = _frame([(10, 50, 1, True)]) + kernel = SimulateCounterfactualKernel({RATE_REF: RatesEngine()}) + with pytest.raises(ValueError, match=match): + _run(kernel, frame, _node(kernel, **changes)) + + +def test_adding_same_adapter_binding_does_not_change_bridge_implementation_hash(): + for cls in (SimulateCounterfactualKernel, SimulateSolveZeroKernel): + first = cls({RATE_REF: RatesEngine()}) + second = cls({RATE_REF: RatesEngine(), "unused": RatesEngine()}) + assert first.implementation_hash() == second.implementation_hash() + + +@pytest.mark.parametrize( + "returned,match", [([1], "wrong shape"), ([np.nan, 1], "finite")] +) +def test_bridge_refuses_malformed_engine_outputs(returned, match): + class MalformedEngine(RatesEngine): + def materialize(self, bundle, variables, period): + return {variable: np.asarray(returned) for variable in variables} + + frame = _frame([(10, 50, 1, True)]) + kernel = SimulateCounterfactualKernel({RATE_REF: MalformedEngine()}) + with pytest.raises(ValueError, match=match): + _run(kernel, frame, _node(kernel)) + + +@pytest.mark.parametrize( + "entity,column", + [ + ("person", "person_id"), + ("person", "person_family_id"), + ("family", "family_id"), + ], +) +def test_solver_rejects_structural_trial_inputs_before_calling_engine(entity, column): + class StructuralInputEngine(RatesEngine): + def variables(self): + return (*super().variables(), column) + + frame = _frame([(10, 50, 1, True)]) + original = {name: frame.table(name).copy() for name in frame.entities} + engine = StructuralInputEngine() + kernel = SimulateSolveZeroKernel({RATE_REF: engine}) + node = _node(kernel, solve_input=_json({"entity": entity, "column": column})) + with pytest.raises(ValueError, match="solve_input cannot change structural"): + _run(kernel, frame, node) + assert engine.calls == 0 + for name in frame.entities: + pd.testing.assert_frame_equal(frame.table(name), original[name]) + + +@pytest.mark.parametrize("expected,status", [(101, "passed"), (100, "failed")]) +def test_hub_differential_receipt_runs_kernel_and_compares_external_oracle( + monkeypatch, expected, status +): + frame = _frame([(10, 50, 1, True)]) + kernel = SimulateCounterfactualKernel({"prepared-reference": RatesEngine()}) + registry = KernelRegistry() + registry.register(kernel) + prepared = SimpleNamespace( + schema=SCHEMA, kernels=registry, engine_refs={RATE_REF: "prepared-reference"} + ) + monkeypatch.setattr( + differential_tool.BenefitUnitRule, "from_dict", lambda data: object() + ) + monkeypatch.setattr( + differential_tool, "build_transport_registry", lambda *args, **kwargs: prepared + ) + node = _node(kernel) + case = { + "unit_rule": {}, + "rules_bindings": {}, + "tables": { + entity: frame.table(entity).to_dict("records") for entity in frame.entities + }, + "weights": {"family": {"kind": "design", "values": [1]}}, + "calls": [ + { + "id": node.id, + "kernel": node.kernel, + "inputs": {item.entity: list(item.columns) for item in node.inputs}, + "outputs": [ + {"entity": "family", "column": "bridge", "dtype": "float64"} + ], + "params": dict(node.params), + } + ], + "comparisons": [ + { + "call": node.id, + "entity": "family", + "column": "bridge", + "oracle_pointer": "/requests/0/expected/weekly", + "tolerance": 0, + } + ], + } + receipt = differential_tool.differential( + case, {"requests": [{"expected": {"weekly": expected}}]}, Path("unused") + ) + assert receipt["status"] == status + assert receipt["comparisons"][0]["actual"] == [101] + assert receipt["comparisons"][0]["expected"] == [expected] + + +def test_hub_differential_skips_without_engine_before_reading_case( + monkeypatch, tmp_path +): + receipt_path = tmp_path / "receipt.json" + monkeypatch.setattr( + differential_tool.importlib.util, "find_spec", lambda name: None + ) + monkeypatch.setattr( + "sys.argv", + [ + "transport_bridge_differential.py", + "--case", + "missing-case.json", + "--oracle", + "missing-oracle.json", + "--rulespec-root", + "missing-root", + "--out", + str(receipt_path), + ], + ) + assert differential_tool.main() == 0 + assert json.loads(receipt_path.read_text())["status"] == "skipped" + + +def test_hub_binding_resolution_preserves_case_for_repeated_runs(): + params = { + "engine_ref": "rates", + "variables": ["benefit"], + "components": [ + { + "engine_ref": "credit", + "variables": ["credit_amount"], + "period": "2026-27", + "input_overrides": [], + } + ], + } + original = _json(params) + first = differential_tool._resolve_params( + params, {"rates": "ref-a", "credit": "ref-b"} + ) + second = differential_tool._resolve_params( + params, {"rates": "ref-a", "credit": "ref-b"} + ) + assert _json(params) == original + assert first == second + assert json.loads(first["components"])[0]["engine_ref"] == "ref-b" diff --git a/test_support/microcosm_build/transport_entitlement.py b/test_support/microcosm_build/transport_entitlement.py new file mode 100644 index 0000000000..c95e4b3587 --- /dev/null +++ b/test_support/microcosm_build/transport_entitlement.py @@ -0,0 +1,833 @@ +"""Synthetic entitlement graph over invented donor support and rules. + +The numbers in ``TOY_PROGRAMS`` and ``TOY_DECLARATIONS`` are test data, never +statutory rates or destination microdata. The adapters read their complete +programs from the pinned fixture tree on every call. Scenario knobs are +separate selected subtrees, so a scenario edit cannot change sibling keys. +""" + +from __future__ import annotations + +import copy +import json +import shutil +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, replace +from hashlib import sha256 +from pathlib import Path + +import numpy as np + +from microcosm.build.transport.cli import EXPORT_FILENAME, through +from microcosm.build.transport.compose import compose_transport_graph +from microcosm.build.transport.registry import build_transport_registry +from microcosm.build.transport.terminal_kernels import materialize_export +from microcosm.frame import Frame, VariableMetadata +from microcosm.graph import ContentStore, compile_graph, run_graph +from microcosm.graph.canonical import canonical_json +from test_support.microcosm_build.transport_composed import ( + COUNTRY, + PERIOD, + ComposedFixture, + ToyRulesEngine, + _edge, + _fact, + _node, + _owned, + _references, + _select, + _slice, + make_composed_fixture, +) +from test_support.microcosm_build.transport_graph import reference_row, write_facts +from test_support.microcosm_build.transport_population import UNIT_RULE + +INPUT_FIELDS = ( + "as_assets", + "as_costs", + "as_income", + "as_beneficiary", + "as_super", + "as_s66", + "as_s67", + "as_maximum", + "as_credit_units", +) +AS_INPUTS = (*INPUT_FIELDS, "as_base_rate", "as_cutout") +AS_OUTPUTS = { + "toy_as_weekly": "float64", + "toy_as_unit_count": "int64", + "toy_as_eligible": "int64", + "toy_as_single": "int64", + "toy_as_other": "int64", + "toy_as_excluded": "int64", +} + +# Invented fixture programs. Every output-affecting amount is serialized into +# the RuleSpec source tree, so these examples exercise content-bound execution. +TOY_PROGRAMS = { + "entitlement": { + "path": "toy/as.yaml", + "entity": "family", + "kind": "as", + "policy": { + "asset_limit": 350.0, + "cost_share": 0.25, + "payment_share": 0.7, + "income_floor": 10.0, + "income_abatement": 0.2, + "family_type": "single", + }, + "defaults": { + "as_assets": {"source": "input_assets", "value": 200.0}, + "as_costs": {"source": "input_rent", "value": 300.0}, + "as_income": {"value": 80.0}, + "as_beneficiary": {"value": 0.0}, + "as_super": {"value": 0.0}, + "as_s66": {"value": 0.0}, + "as_s67": {"value": 0.0}, + "as_maximum": {"value": 90.0}, + "as_credit_units": {"value": 1.0}, + }, + "outputs": { + **{f"seed_{name}": {"dtype": "float64"} for name in INPUT_FIELDS}, + **{name: {"dtype": dtype} for name, dtype in AS_OUTPUTS.items()}, + }, + }, + "bridge_rates": { + "path": "toy/rates.yaml", + "entity": "person", + "kind": "rates", + "policy": {"rate": 40.0, "alternative_rate": 60.0, "abatement": 0.25}, + "outputs": {"toy_bridge_benefit": {"dtype": "float64"}}, + }, + "bridge_credit": { + "path": "toy/credits.yaml", + "entity": "family", + "kind": "credit", + "policy": {"credit": 16.0}, + "outputs": {"toy_bridge_credit": {"dtype": "float64"}}, + }, +} +TOY_DECLARATIONS = { + "base_credit_coefficient": 0.25, + # The invented cutout 160 is exact on this dyadic bisection grid. Eight + # fixed iterations preserve the graph semantics; core solver tests cover + # general brackets and tighter accuracy separately. + "solve_bracket": [0.0, 256.0], + "solve_tolerance": 1.0, + "solve_iterations": 8, + "annualisation": 4.0, + "target_factors": {"V1": 0.95, "V2": 1.05}, + "alternative_receipt_rate": 0.2, +} + + +class EntitlementToyEngine(ToyRulesEngine): + """AS-shaped toy arithmetic using only the public ``RulesEngine`` protocol.""" + + def cache_identity(self) -> Mapping: + return {"formula": "synthetic-as-and-bridge-v1"} + + def variable_metadata(self, name: str) -> VariableMetadata: + program = self.program() + dtype = program["outputs"][name]["dtype"] + return VariableMetadata( + name, + program["entity"], + "int" if dtype == "int64" else "float", + "year", + ) + + def variables(self) -> Sequence[str]: + kind = self.program()["kind"] + if kind == "rates": + return ("input_income", "bridge_alternative") + if kind == "credit": + return ("as_credit_units",) + return (*AS_INPUTS, "input_assets", "input_rent") + + def materialize(self, bundle: Frame, variables, period): + self.calls.append((tuple(variables), period)) + program = self.program() + table = bundle.table(program["entity"]) + policy = program["policy"] + + def column(name, fallback=0): + if name in table: + return table[name].to_numpy(dtype=np.float64) + return np.full(len(table), fallback, dtype=np.float64) + + if program["kind"] == "rates": + rate = np.where( + column("bridge_alternative") != 0, + policy["alternative_rate"], + policy["rate"], + ) + result = { + "toy_bridge_benefit": np.maximum( + rate - column("input_income") * policy["abatement"], 0 + ) + } + elif program["kind"] == "credit": + result = { + "toy_bridge_credit": column("as_credit_units", 1) * policy["credit"] + } + else: + result = {} + for name, row in program["defaults"].items(): + result[f"seed_{name}"] = ( + column(name, row["value"]) + if name in table + else column(row.get("source", name), row["value"]) + ) + # Seed-only counterfactuals need no derived AS inputs. + if set(variables).intersection(AS_OUTPUTS): + eligible = ( + (column("as_assets") <= policy["asset_limit"]) + & (column("as_s66") == 0) + & (column("as_s67") == 0) + ) + costs = np.maximum( + column("as_costs") - column("as_base_rate") * policy["cost_share"], + 0, + ) + amount = np.minimum( + column("as_maximum"), costs * policy["payment_share"] + ) + reduction = ( + np.maximum(column("as_income") - policy["income_floor"], 0) + * policy["income_abatement"] + ) + amount = np.maximum( + amount - np.where(column("as_beneficiary") != 0, 0, reduction), 0 + ) + eligible &= (column("as_beneficiary") != 0) | ( + column("as_income") < column("as_cutout") + ) + result["toy_as_weekly"] = np.where(eligible, amount, 0) + result["toy_as_unit_count"] = np.ones(len(table), dtype=np.int64) + result["toy_as_eligible"] = np.where(eligible, 1, -1) + result["toy_as_single"] = ( + table["family_type"].to_numpy() == policy["family_type"] + ).astype(np.int64) + result["toy_as_other"] = ( + (column("as_beneficiary") == 0) & (column("as_super") == 0) + ).astype(np.int64) + result["toy_as_excluded"] = ( + (column("as_s66") != 0) | (column("as_s67") != 0) + ).astype(np.int64) + return { + name: np.asarray(result[name], dtype=program["outputs"][name]["dtype"]) + for name in variables + } + + +def _binding(name, field="engine_ref"): + return {"binding": name, "field": field} + + +def _terms(variables, columns): + return [ + {"variable": variable, "coefficient": 1, "entity": "family", "column": column} + for variable, column in zip(variables, columns, strict=True) + ] + + +def _seed(population, id, *, scenario=False, fields=INPUT_FIELDS): + return _node( + id, + "simulate.counterfactual@1", + population=population, + sources=["rulespec_nz"], + inputs=[ + _slice("person", ["age"]), + _slice("family", fields if scenario else ["input_assets", "input_rent"]), + ], + outputs=[ + _owned("family", name, "float64", rewrite=scenario) for name in fields + ], + params={ + "engine_ref": _binding("entitlement"), + "period": _binding("entitlement", "period"), + "variables": [f"seed_{name}" for name in fields], + "input_overrides": { + "scenario": "knobs", + "encoding": "json", + "path": ["input_overrides"], + } + if scenario + else _select("entitlement_seed", "json", "input_overrides"), + "output_coefficients": _terms([f"seed_{name}" for name in fields], fields), + "resource_sha256": {"scenario": "knobs", "encoding": "sha256"} + if scenario + else _select("entitlement_seed", "sha256"), + }, + ) + + +def _bridge(population, prefix): + slices = [ + _slice("person", ["age", "input_income"]), + _slice("family", ["as_credit_units"]), + ] + base = _node( + f"{prefix}.bridge.base_rate", + "simulate.counterfactual@1", + population=population, + sources=["rulespec_nz"], + inputs=slices, + outputs=[_owned("family", "as_base_rate", "float64")], + params={ + "engine_ref": _binding("bridge_rates"), + "period": _binding("bridge_rates", "period"), + "variables": ["toy_bridge_benefit"], + "resource_sha256": _select("entitlement_bridge", "sha256"), + "input_overrides": [ + {"entity": "person", "column": "input_income", "value": 0} + ], + "components": [ + { + "engine_ref": _binding("bridge_credit"), + "period": _binding("bridge_credit", "period"), + "variables": ["toy_bridge_credit"], + "input_overrides": [ + {"entity": "family", "column": "as_credit_units", "value": 1} + ], + } + ], + "output_coefficients": [ + { + "variable": "toy_bridge_benefit", + "coefficient": 1, + "entity": "family", + "column": "as_base_rate", + "aggregation": "max", + }, + { + "engine_ref": _binding("bridge_credit"), + "variable": "toy_bridge_credit", + "coefficient": TOY_DECLARATIONS["base_credit_coefficient"], + "entity": "family", + "column": "as_base_rate", + }, + ], + }, + ) + cutout = _node( + f"{prefix}.bridge.cutout", + "simulate.solve_zero@1", + population=population, + sources=["rulespec_nz"], + inputs=slices, + outputs=[_owned("family", "as_cutout", "float64")], + params={ + "engine_ref": _binding("bridge_rates"), + "period": _binding("bridge_rates", "period"), + "variables": ["toy_bridge_benefit"], + "resource_sha256": _select("entitlement_bridge", "sha256"), + "input_overrides": [ + {"entity": "person", "column": "bridge_alternative", "value": 0} + ], + "output_coefficients": [ + { + "variable": "toy_bridge_benefit", + "coefficient": 1, + "entity": "family", + "column": "as_cutout", + "aggregation": "max", + } + ], + "solve_input": {"entity": "person", "column": "input_income"}, + "bracket": TOY_DECLARATIONS["solve_bracket"], + "tolerance": TOY_DECLARATIONS["solve_tolerance"], + "iterations": TOY_DECLARATIONS["solve_iterations"], + }, + ) + return [base, cutout] + + +def _as_and_gap(population, prefix, scenario): + return [ + _node( + f"{prefix}.as", + "simulate.rules_by_ref@1", + population=population, + sources=["rulespec_nz"], + inputs=[ + _slice("person", ["age"]), + _slice("family", [*AS_INPUTS, "family_type"]), + ], + outputs=[ + _owned("family", name, dtype) for name, dtype in AS_OUTPUTS.items() + ], + rules_binding="entitlement", + params={"variables": list(AS_OUTPUTS)}, + ), + _node( + f"{prefix}.gap", + "takeup.gap@1", + population=population, + sources=["nz_holdout_facts"], + inputs=[ + _slice("person", ["age"]), + _slice("household", ["rent"]), + _slice("family", [*AS_OUTPUTS, "as_beneficiary", "as_super"]), + ], + params={ + "country": COUNTRY, + "references": _select("entitlement_gap"), + "references_sha256": _select("entitlement_gap", "sha256"), + "entities": ["person", "household", "family"], + "weight_entity": "household", + "annualisation_factors": _select("entitlement_annualisation"), + "annualisation_sha256": _select("entitlement_annualisation", "sha256"), + "scenario": scenario, + }, + ), + ] + + +def _boundary(id, base): + return _node( + id, + "transport.boundary@1", + structural="filter", + base=base, + inputs=[_slice("person", ["age"])], + params={}, + ) + + +def _scenario_nodes(fields): + nodes = [_boundary("nz.scn.{scenario}", "nz.as")] + if fields: + nodes.append( + _seed( + "nz.scn.{scenario}", + "nz.scn.{scenario}.override", + scenario=True, + fields=fields, + ) + ) + return [ + *nodes, + *_as_and_gap("nz.scn.{scenario}", "nz.scn.{scenario}", "{scenario}"), + ] + + +def _entitlement_document(skeleton): + shared = [_seed("nz.as", "nz.as.seed"), *_bridge("nz.as", "nz.as")] + shared += [ + _boundary("nz.validate", "nz.calibrate"), + _node( + "nz.validate.wff", + "simulate.counterfactual@1", + population="nz.validate", + sources=["rulespec_nz"], + inputs=[_slice("person", ["age"])], + outputs=[_owned("family", "toy_wff_credit", "float64")], + params={ + "engine_ref": _binding("bridge_credit"), + "period": _binding("bridge_credit", "period"), + "variables": ["toy_bridge_credit"], + "input_overrides": [ + {"entity": "family", "column": "as_credit_units", "value": 1} + ], + "output_coefficients": _terms( + ["toy_bridge_credit"], ["toy_wff_credit"] + ), + }, + ), + _node( + "nz.validate.wff_tripwire", + "takeup.compare@1", + population="nz.validate", + sources=["nz_holdout_facts"], + inputs=[ + _slice("person", ["age"]), + _slice("household", ["rent"]), + _slice("family", ["toy_wff_credit"]), + ], + params={ + "country": COUNTRY, + "references": _select("wff_holdout_references"), + "references_sha256": _select("wff_holdout_references", "sha256"), + "entities": ["person", "household", "family"], + "weight_entity": "household", + }, + ), + ] + return { + "schema_version": 1, + "nodes": shared, + # A newly added scenario in the C2 fixture declares an asset override. + "scenario_nodes": _scenario_nodes(("as_assets",)), + "scenario_gap": "nz.scn.{scenario}.gap", + "variant_gap": "nz.v.{variant}.gap", + "bands": { + "id": "nz.as.bands", + "population": "nz.terminal", + "baseline_scenario": "S0", + }, + "package": { + "id": "nz.package", + "artifact_name": "bands", + "inputs": [_edge("wff_tripwire", "nz.validate.wff_tripwire", "comparison")], + }, + "checkpoints": ["nz.as.bands", "nz.validate.wff_tripwire"], + } + + +def _variant_nodes(skeleton, variant): + """An actual independent target/problem/calibration branch, then AS.""" + by_id = {node["id"]: node for node in skeleton["nodes"]} + prefix = f"nz.v.{variant}" + nodes = [_boundary(prefix, "nz.open")] + if variant == "V3": + receipt = copy.deepcopy(by_id["nz.takeup.assign"]) + receipt.update(id=f"{prefix}.receipt", population=prefix) + for output in receipt["outputs"]: + output["rewrite"] = True + receipt["params"] = { + "receipt_contract": _select("variant_receipt", "json", "receipts"), + "receipt_contract_sha256": _select("variant_receipt", "sha256"), + } + nodes.append(receipt) + targets = copy.deepcopy(by_id["nz.targets.compile"]) + targets.update(id=f"{prefix}.targets", population=prefix) + reference = "target_references" if variant == "V3" else f"variant_{variant}_targets" + targets["params"] = { + "country": COUNTRY, + "references": _select(reference), + "references_sha256": _select(reference, "sha256"), + } + problem = copy.deepcopy(by_id["nz.targets.problem"]) + problem.update( + id=f"{prefix}.problem", + population=prefix, + artifact_inputs=[_edge("surface", targets["id"])], + ) + calibrated = copy.deepcopy(by_id["nz.calibrate"]) + calibrated.update( + id=f"{prefix}.calibrate", + base=prefix, + artifact_inputs=[_edge("problem", problem["id"])], + ) + as_population = f"{prefix}.housing" + group = copy.deepcopy(by_id["nz.as.inputs"]) + group.update(id=f"{prefix}.inputs", population=as_population) + nodes += [ + targets, + problem, + calibrated, + _boundary(as_population, calibrated["id"]), + group, + _seed(as_population, f"{prefix}.seed"), + *_bridge(as_population, prefix), + *_as_and_gap(as_population, prefix, variant), + ] + return nodes + + +@dataclass(frozen=True) +class EntitlementFixture(ComposedFixture): + """Prepared composed toy graph with real entitlement extension kernels.""" + + def graph(self, *, extended=True): + return compose_transport_graph(self.spec, self.config) + + def reprepare(self, spec=None): + from microcosm.build.transport.entitlement import entitlement_extension + + spec = self.spec if spec is None else spec + bindings = spec["resources"]["axiom_rules_bindings"] + root = self.sources["rulespec_nz"] + engines = { + row["id"]: ( + EntitlementToyEngine if row["id"] in TOY_PROGRAMS else ToyRulesEngine + )(root, row) + for row in bindings["bindings"] + } + registry = build_transport_registry( + bindings, root, unit_rule=UNIT_RULE, engines_by_binding=engines + ) + config = replace( + self.config, + engine_refs=registry.engine_refs, + extensions=(entitlement_extension(spec),), + ) + return replace( + self, + spec=spec, + config=config, + registry=registry, + engines_by_binding=engines, + ) + + def relocated(self, root): + shutil.copytree(self.root, root) + sources = { + name: root / path.relative_to(self.root) + for name, path in self.sources.items() + } + return replace(self, root=root, sources=sources).reprepare() + + +def make_entitlement_fixture(root: Path, *, seed_overrides=None, variants=True): + base = make_composed_fixture(root) + spec = copy.deepcopy(base.spec) + resources = spec["resources"] + tree = base.sources["rulespec_nz"] + for name, raw in TOY_PROGRAMS.items(): + program = {key: value for key, value in raw.items() if key != "path"} + program["relations"] = [] + payload = canonical_json(program) + path = tree / raw["path"] + path.write_bytes(payload) + resources["axiom_rules_bindings"]["bindings"].append( + { + "id": name, + "rulespec_path": raw["path"], + "sha256": sha256(payload).hexdigest(), + "entity": raw["entity"], + "engine_entity": raw["entity"].capitalize(), + "period": PERIOD, + "variables": list(raw["outputs"]), + } + ) + resources["entitlement_seed"] = { + "input_overrides": [ + {"entity": "family", "column": name, "value": value} + for name, value in (seed_overrides or {}).items() + ] + } + resources["entitlement_bridge"] = copy.deepcopy(TOY_DECLARATIONS) + gap_rows = [ + reference_row("toy_as_expenditure", entity="family", measure="toy_as_weekly"), + reference_row( + "toy_as_recipients", + entity="family", + measure="toy_as_unit_count", + filter_="toy_as_weekly", + ), + reference_row( + "toy_as_single_expenditure", + entity="family", + measure="toy_as_weekly", + filter_="toy_as_single", + ), + reference_row( + "toy_as_beneficiary_expenditure", + entity="family", + measure="toy_as_weekly", + filter_="as_beneficiary", + ), + reference_row( + "toy_as_super_expenditure", + entity="family", + measure="toy_as_weekly", + filter_="as_super", + ), + reference_row( + "toy_as_other_expenditure", + entity="family", + measure="toy_as_weekly", + filter_="toy_as_other", + ), + reference_row( + "toy_as_excluded_units", + entity="family", + measure="toy_as_unit_count", + filter_="toy_as_excluded", + ), + ] + for row, cell in zip( + gap_rows[2:], + ( + {"income_decile": "toy_decile_1", "family_type": "single"}, + {"receipt_stratum": "beneficiary"}, + {"receipt_stratum": "superannuitant"}, + {"receipt_stratum": "other"}, + {"eligibility_stratum": "excluded_s66_or_s67"}, + ), + strict=True, + ): + # LedgerTargetReference metadata is a flat string-to-string mapping. + row["metadata"].update(cell) + resources["entitlement_gap"] = _references(gap_rows) + # Recipient target counts positivity, rather than summing -1/1 judgments. + resources["entitlement_annualisation"] = { + row["name"]: TOY_DECLARATIONS["annualisation"] + if row["measure"] == "toy_as_weekly" + else 1 + for row in gap_rows + } + resources["wff_holdout_references"] = _references( + [ + reference_row( + "toy_wff_expenditure", entity="family", measure="toy_wff_credit" + ) + ] + ) + old_facts = [ + json.loads(row) + for row in base.sources["nz_holdout_facts"].read_text().splitlines() + ] + write_facts( + base.sources["nz_holdout_facts"], + [ + *old_facts, + _fact("toy_as_expenditure", 100.0, "family"), + _fact("toy_as_recipients", 15.0, "family"), + *[_fact(row["name"], 0.0, "family") for row in gap_rows[2:]], + _fact("toy_wff_expenditure", 200.0, "family"), + ], + ) + resources["scenarios"] = { + "scenarios": [ + {"id": "S0", "tier": "entitlement", "knobs": {"input_overrides": []}}, + { + "id": "S1", + "tier": "entitlement", + "knobs": { + "input_overrides": [ + {"entity": "family", "column": "as_assets", "value": 0.0} + ] + }, + }, + { + "id": "S2", + "tier": "entitlement", + "knobs": { + "input_overrides": [ + {"entity": "family", "column": "as_costs", "value": 150.0} + ] + }, + }, + { + "id": "S3", + "tier": "entitlement", + "knobs": { + "input_overrides": [ + {"entity": "family", "column": "as_maximum", "value": 60.0} + ] + }, + }, + { + "id": "S4", + "tier": "entitlement", + "knobs": { + "input_overrides": [ + {"entity": "family", "column": "as_costs", "value": 450.0} + ] + }, + }, + { + "id": "S5", + "tier": "entitlement", + "knobs": { + "input_overrides": [ + {"entity": "family", "column": "as_super", "value": 1}, + {"entity": "family", "column": "as_beneficiary", "value": 0}, + ] + }, + }, + ] + } + document = _entitlement_document(resources["transport_graph"]) + for row in resources["scenarios"]["scenarios"]: + row["nodes"] = _scenario_nodes( + tuple(item["column"] for item in row["knobs"]["input_overrides"]) + ) + # The alternate base-rate scenario runs the engine again; policy-derived + # bridge columns are never overwritten by a data-only transform. + s5 = resources["scenarios"]["scenarios"][-1] + alternate = copy.deepcopy(document["nodes"][1]) + alternate.update(id="nz.scn.{scenario}.base_rate", population="nz.scn.{scenario}") + alternate["outputs"][0]["rewrite"] = True + alternate["params"]["input_overrides"].append( + {"entity": "person", "column": "bridge_alternative", "value": 1} + ) + s5["nodes"].insert(2, alternate) + if variants: + facts = [ + json.loads(row) + for row in base.sources["nz_calibration_facts"].read_text().splitlines() + ] + fact_by_name = { + row["semantic_fact_key"].rsplit(":", 1)[-1]: row for row in facts + } + for variant, factor in TOY_DECLARATIONS["target_factors"].items(): + refs = copy.deepcopy(resources["target_references"]) + for row in refs["target_references"]: + original = row["name"] + name = f"{variant}_{original}" + row.update(name=name, ledger_fact_key=f"toy.aggregate_fact.v1:{name}") + facts.append( + _fact(name, fact_by_name[original]["value"] * factor, row["entity"]) + ) + resources[f"variant_{variant}_targets"] = refs + write_facts(base.sources["nz_calibration_facts"], facts) + receipt = copy.deepcopy(resources["receipt_contract"]) + receipt["receipts"]["programs"][0]["rate"] = TOY_DECLARATIONS[ + "alternative_receipt_rate" + ] + resources["variant_receipt"] = receipt + for variant in ("V1", "V2", "V3"): + resources["scenarios"]["scenarios"].append( + { + "id": variant, + "tier": "calibration", + "knobs": {}, + "nodes": _variant_nodes(resources["transport_graph"], variant), + } + ) + resources["entitlement_graph"] = document + fixture = EntitlementFixture( + base.root, + spec, + base.config, + base.registry, + base.sources, + base.engines_by_binding, + ) + return fixture.reprepare() + + +def run_entitlement_graph( + root, fixture, *, store=None, exported=None, graph=None, resume="auto" +): + """Execute the complete graph with one pinned H5 outer-service input.""" + root.mkdir(parents=True, exist_ok=True) + graph = fixture.graph() if graph is None else graph + store = ( + ContentStore(root / "store", codecs=fixture.registry.codecs) + if store is None + else store + ) + if exported is None: + prepared = run_graph( + compile_graph(through(graph, "nz.export.prepare")), + sources=fixture.sources, + store=store, + kernels=fixture.registry.kernels, + ) + descriptor = json.loads( + store.load_bytes( + prepared.nodes["nz.export.prepare"].opaque_artifacts[ + "export_descriptor" + ] + ) + ) + version = compile_graph(graph).versions["nz.export.prepare"] + exported = root / EXPORT_FILENAME + materialize_export(prepared.populations[version], descriptor, exported) + manifest = run_graph( + compile_graph(graph), + sources=fixture.source_mapping(exported), + store=store, + kernels=fixture.registry.kernels, + resume=resume, + ) + return manifest, store, exported diff --git a/tools/transport_bridge_differential.py b/tools/transport_bridge_differential.py new file mode 100644 index 0000000000..13e64628e3 --- /dev/null +++ b/tools/transport_bridge_differential.py @@ -0,0 +1,272 @@ +"""Run a small declared bridge/AS differential and write a local JSON receipt. + +No country engine is installed by this command. If ``axiom_rules_engine`` is +absent it writes a skipped receipt and exits successfully. A hub case JSON +declares ``unit_rule`` (BenefitUnitRule data), ``rules_bindings`` (the transport +registry contract), ``tables`` (entity -> records), ``weights`` (entity -> +{values, kind}), and ordered ``calls``. Each call has ``id``, ``kernel``, +``inputs`` (entity -> data-column names), ``outputs`` (Owned records), and +``params`` (the kernel contract). ``engine_ref`` values name binding ids and +are replaced with prepared references, including embedded component/term refs. + +``comparisons`` contains records with ``call``, ``entity``, ``column``, +``oracle_pointer`` (RFC 6901 pointer into the oracle JSON) and explicit +``tolerance``. Call ids must be nonempty and distinct, each entity/column must +have one output owner, and comparisons must name an output of their call. +These declarations and oracle comparisons are checked before registry +preparation or engine execution. All numbers come from case data or the +oracle; the script has no policy constants. + +The golden-08 executable case is unresolved pending an approved declaration. +``build/nz/as_rate_bridge.json`` declares three sole-parent cutout evaluations: +``gross_rate``, ``reduction_at_anchor`` and ``reduction_after_step``, followed by +slope, threshold and cutout composition. Its ``why_not_solve_zero`` explains +why a zero solve on the unit's own Jobseeker schedule cannot reproduce that +harness pairing. This driver has no arithmetic-composition contract for those +evaluations and cannot yet run that recipe. Approved sampled-unit cases may +declare additional small entity-table rows and corresponding oracle arrays. +This command does not read donor HDF5. +""" + +from __future__ import annotations + +import argparse +import hashlib +import importlib.util +import json +import re +from pathlib import Path + +import numpy as np +import pandas as pd + +from microcosm.build.transport.registry import build_transport_registry +from microcosm.frame import Frame, WeightKind, Weights +from microcosm.frame.unit_construction import BenefitUnitRule +from microcosm.graph import KernelContext, Node, Owned, Slice + + +def _pointer(document, pointer): + if pointer == "": + return document + if not isinstance(pointer, str) or not pointer.startswith("/"): + raise ValueError("oracle_pointer must be an RFC 6901 JSON pointer.") + value = document + for token in pointer[1:].split("/"): + if re.search(r"~(?:[^01]|$)", token): + raise ValueError("oracle_pointer contains an invalid JSON pointer escape.") + token = token.replace("~1", "/").replace("~0", "~") + if isinstance(value, list): + if re.fullmatch(r"0|[1-9][0-9]*", token) is None: + raise ValueError( + "oracle_pointer requires a canonical non-negative array index." + ) + index = int(token) + if index >= len(value): + raise ValueError("oracle_pointer array index is out of bounds.") + value = value[index] + elif isinstance(value, dict) and token in value: + value = value[token] + else: + raise ValueError("oracle_pointer does not resolve to an oracle value.") + return value + + +def _validate_case(case, oracle): + """Refuse ambiguous calls and comparisons before preparing any engine.""" + declared_outputs = {} + for call in case["calls"]: + call_id = call.get("id") + if not isinstance(call_id, str) or not call_id: + raise ValueError("Every differential call needs a non-empty call id.") + if call_id in declared_outputs: + raise ValueError( + f"The differential case has duplicate call id {call_id!r}." + ) + declared_outputs[call_id] = set() + owners = {} + for call in case["calls"]: + call_id = call["id"] + for output in call["outputs"]: + coordinate = (output["entity"], output["column"]) + if coordinate in owners: + raise ValueError( + f"The differential output coordinate {coordinate!r} has " + f"multiple owners: {owners[coordinate]!r} and {call_id!r}." + ) + owners[coordinate] = call_id + declared_outputs[call_id].add(coordinate) + if not case["comparisons"]: + raise ValueError("The differential case must declare comparisons.") + expected_values = [] + for row in case["comparisons"]: + call_id = row["call"] + if call_id not in declared_outputs: + raise ValueError(f"A comparison names missing call {call_id!r}.") + coordinate = (row["entity"], row["column"]) + if coordinate not in declared_outputs[call_id]: + raise ValueError( + f"A comparison names missing output {coordinate!r} " + f"for call {call_id!r}." + ) + tolerance = row["tolerance"] + if ( + isinstance(tolerance, bool) + or not isinstance(tolerance, int | float) + or not np.isfinite(tolerance) + or tolerance < 0 + ): + raise ValueError( + "Every comparison needs an explicit finite nonnegative tolerance." + ) + expected = np.atleast_1d( + np.asarray(_pointer(oracle, row["oracle_pointer"]), dtype=np.float64) + ) + if not np.all(np.isfinite(expected)): + raise ValueError("Oracle comparison values must be finite.") + expected_values.append(expected) + return expected_values + + +def _resolve_params(params, refs): + result = dict(params) + result["engine_ref"] = refs[result["engine_ref"]] + if "variables" in result: + result["variables"] = tuple(result["variables"]) + if "bracket" in result: + result["bracket"] = tuple(result["bracket"]) + for name in ("input_overrides", "output_coefficients", "components", "solve_input"): + if name not in result: + continue + data = result[name] + if isinstance(data, str): + data = json.loads(data) + else: + data = json.loads(json.dumps(data)) + records = data if isinstance(data, list) else [data] + for record in records: + if "engine_ref" in record: + record["engine_ref"] = refs[record["engine_ref"]] + result[name] = json.dumps(data, sort_keys=True, separators=(",", ":")) + return result + + +def _call_tables(frame, node): + """Expose only declared columns and structural ids, as graph execution does.""" + columns = {output.entity: [] for output in node.outputs} + for item in node.inputs: + columns.setdefault(item.entity, []).extend(item.columns) + tables = {} + for entity, declared in columns.items(): + structural = [frame.schema.entity_id_column(entity)] + if entity == frame.schema.person_entity: + structural.extend( + frame.schema.membership_column(group) + for group in frame.schema.group_entities + ) + tables[entity] = frame.table(entity).loc[ + :, list(dict.fromkeys(structural + declared)) + ] + return tables + + +def differential(case, oracle, rulespec_root): + """Validate declarations, execute unique calls, and compare oracle data.""" + expected_values = _validate_case(case, oracle) + rule = BenefitUnitRule.from_dict(case["unit_rule"]) + prepared = build_transport_registry( + case["rules_bindings"], rulespec_root, unit_rule=rule + ) + tables = { + entity: pd.DataFrame(records) for entity, records in case["tables"].items() + } + weights = { + entity: Weights( + np.asarray(row["values"], dtype=np.float64), WeightKind(row["kind"]) + ) + for entity, row in case["weights"].items() + } + frame = Frame(tables, prepared.schema, weights) + outputs = {} + receipts = [] + for call in case["calls"]: + node = Node( + call["id"], + call["kernel"], + inputs=tuple( + Slice(entity, tuple(columns)) + for entity, columns in call["inputs"].items() + ), + outputs=tuple(Owned(**output) for output in call["outputs"]), + params=_resolve_params(call["params"], prepared.engine_refs), + ) + result = prepared.kernels.get(node.kernel).run( + KernelContext( + node=node, + tables=_call_tables(frame, node), + weights=weights, + strata=frame.strata, + params=node.params, + rng=np.random.default_rng(0), + ) + ) + receipts.append({"call": node.id, "receipt": dict(result.receipt)}) + for (entity, column), values in result.columns.items(): + outputs[(node.id, entity, column)] = values.to_numpy() + ids = frame.schema.entity_id_column(entity) + tables[entity][column] = tables[entity][ids].map(values).to_numpy() + frame = Frame(tables, prepared.schema, weights) + comparisons = [] + for row, expected in zip(case["comparisons"], expected_values, strict=True): + actual = outputs[(row["call"], row["entity"], row["column"])] + tolerance = row["tolerance"] + same_shape = actual.shape == expected.shape + differences = np.abs(actual - expected) if same_shape else None + comparisons.append( + { + **row, + "actual": actual.tolist(), + "expected": expected.tolist(), + "maximum_absolute_difference": float(differences.max()) + if same_shape and len(actual) + else None, + "passed": bool( + same_shape and len(actual) and np.all(differences <= tolerance) + ), + } + ) + return { + "status": "passed" if all(row["passed"] for row in comparisons) else "failed", + "comparisons": comparisons, + "calls": receipts, + } + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--case", type=Path, required=True) + parser.add_argument("--oracle", type=Path, required=True) + parser.add_argument("--rulespec-root", type=Path, required=True) + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + if importlib.util.find_spec("axiom_rules_engine") is None: + receipt = {"status": "skipped", "reason": "axiom_rules_engine is unavailable"} + else: + case_bytes, oracle_bytes = args.case.read_bytes(), args.oracle.read_bytes() + receipt = differential( + json.loads(case_bytes), json.loads(oracle_bytes), args.rulespec_root + ) + receipt.update( + case_sha256=hashlib.sha256(case_bytes).hexdigest(), + oracle_sha256=hashlib.sha256(oracle_bytes).hexdigest(), + ) + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text( + json.dumps(receipt, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + print(json.dumps({"status": receipt["status"], "receipt": str(args.out)})) + return int(receipt["status"] == "failed") + + +if __name__ == "__main__": + raise SystemExit(main())