From 76a137c9e64c5f10ac76832667464e1444369c58 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 1 Nov 2026 10:11:54 -0500 Subject: [PATCH 1/3] Add declared tax-unit partitions and binary Axiom relation bindings Add a versioned MultiUnitRule list with registered benefit and tax membership kinds, so a country can declare several unit partitions beside its households. transport.create@1 lifts the unchanged NZ rule and transport.create@2 reads the versioned list. Tax membership requires an explicit non-null boolean per person. Both eligible reciprocal partners pair; everyone else stays a singleton. Add frozen binary relation manifests to the Axiom adapter, with source and compiled tuple checks, related-input enumeration and identity-ordered native batches. A binding without relations keeps its identity and its legacy native call. Co-Authored-By: Claude Opus 5.5 --- changelog.d/declared-unit-partitions.added.md | 5 + docs/transport-declared-units.md | 104 ++++ .../microcosm/build/calibration_columns.py | 59 +- .../src/microcosm/build/transport/cli.py | 20 +- .../build/transport/column_kernels.py | 84 ++- .../src/microcosm/build/transport/compose.py | 92 ++- .../build/transport/population_kernels.py | 61 +- .../src/microcosm/build/transport/registry.py | 162 +++++- .../build/transport/target_kernels.py | 11 +- .../build/transport/unit_declarations.py | 30 + .../shared/test_calibration_columns.py | 77 +++ .../shared/test_nz_spec_package.py | 48 +- .../shared/test_transport_column_kernels.py | 66 +++ .../shared/test_transport_declared_units.py | 518 +++++++++++++++++ .../test_transport_population_kernels.py | 9 +- .../shared/test_transport_registry.py | 296 +++++++++- .../shared/test_transport_target_kernels.py | 14 + .../calibration_columns_09a2fce.py.txt | 223 ++++++++ .../src/microcosm/frame/adapters/axiom.py | 419 +++++++++++++- .../src/microcosm/frame/rules_kernels.py | 5 + .../src/microcosm/frame/unit_construction.py | 472 +++++++++++++-- .../shared/test_axiom_relations.py | 539 ++++++++++++++++++ .../engine_free/shared/test_declared_units.py | 414 ++++++++++++++ .../policies/tests/axiom_toy_relations.yaml | 33 ++ 24 files changed, 3557 insertions(+), 204 deletions(-) create mode 100644 changelog.d/declared-unit-partitions.added.md create mode 100644 docs/transport-declared-units.md create mode 100644 packages/microcosm-build/src/microcosm/build/transport/unit_declarations.py create mode 100644 packages/microcosm-build/tests/engine_free/shared/test_transport_declared_units.py create mode 100644 packages/microcosm-build/tests/fixtures/calibration_columns_09a2fce.py.txt create mode 100644 packages/microcosm-frame/tests/engine_free/shared/test_axiom_relations.py create mode 100644 packages/microcosm-frame/tests/engine_free/shared/test_declared_units.py create mode 100644 packages/microcosm-frame/tests/fixtures/rulespec-zz/zz/policies/tests/axiom_toy_relations.yaml diff --git a/changelog.d/declared-unit-partitions.added.md b/changelog.d/declared-unit-partitions.added.md new file mode 100644 index 0000000000..c025be6716 --- /dev/null +++ b/changelog.d/declared-unit-partitions.added.md @@ -0,0 +1,5 @@ +Add opt-in versioned unit declarations with shared benefit-unit and tax-unit +partitions, preserving the legacy NZ CREATE interface. Validate declared binary +Axiom relations, gather related inputs by identity and pass native dense batches +without changing references for bindings without relations. Move calibration +roster preparation into the shared unit framework. diff --git a/docs/transport-declared-units.md b/docs/transport-declared-units.md new file mode 100644 index 0000000000..87ddca0e01 --- /dev/null +++ b/docs/transport-declared-units.md @@ -0,0 +1,104 @@ +# Declared transport units and binary relations + +`transport.create@2` accepts a versioned `unit_rule` document, using the same +donor, fact, seed and digest parameters as `transport.create@1`: + +```json +{ + "format": "microcosm.multi_unit_rule.v1", + "units": [ + { + "kind": "tax_unit", + "rule": { + "format": "microcosm.tax_unit_rule.v1", + "entity": "tax_unit", + "pairing_eligibility_column": "joint_assessment_eligible", + "provenance": {"eligibility": "The declared source of eligibility"} + } + } + ] +} +``` + +A declaration requires at least one entry and distinct physical entity names; +unknown kinds are refused. Each entry dispatches to a registered membership constructor. `benefit_unit` +uses the existing `BenefitUnitRule.to_dict()` format. The legacy CREATE reads +the unchanged NZ country rule and lifts it to a one-entry declaration. CREATE, +CLI inventory preparation, unit attributes, group encoding and the rules +registry share the declarations and their derived entity schema and nesting. + +Tax units consume an explicit, non-null boolean eligibility column. Both +members must be eligible reciprocal co-resident partners to share an assessment; +every other person in a valid concept frame has a singleton assessment. Invalid +self, non-reciprocal or dangling pointers are refused by the shared concept +validation. Minors' assessment membership +does not allocate dependant relief. The smaller physical person ID is the +head (role `head`, spouse A); the other is `partner` (spouse B). This rule reads +no income and implements no Belgian legal eligibility test. Its eligibility +input must already exist in the supplied concept tables; CREATE's donor +boundary continues to carry only declared transport concepts. A country must +provide and document that input before activation. + +All partitions use the same validation, household nesting and source-household +lineage. Units store no weights. `Frame.resolve_weights(unit)` inherits each +unit's household weight once. Household totals use the household table; +unit totals count the units explicitly. Changed membership changes CREATE's +population identity. Unchanged membership keeps IDs under row permutations and +income changes. Physical support clones keep distinct IDs and shared original +source-household lineage. + +`transport.unit_attributes@1` accepts either declaration format. It can emit +`person._role` for each partition when the node owns that output. Legacy +nodes keep their declared output inventory. Group encoding uses the sole +declared partition for legacy and one-entry declarations. With several +partitions, its node must name the physical `unit_entity`; `engine_entity` +separately names the logical engine entity. The person/household concept +mapping does not choose group partitions. + +Binary Axiom bindings opt in with `relations` on the binding row. For example: + +```json +{ + "predicate": "be:statutes/example#relation.spouse", + "key": "be:statutes/example#relation.spouse:0:1", + "current_slot": 0, + "related_slot": 1, + "arguments": ["TaxUnit", "Person"], + "root_entity": "tax_unit", + "related_entity": "person", + "membership_column": "person_tax_unit_id", + "role_column": "tax_unit_role", + "roles": ["head", "partner"], + "cardinality": "many", + "period": "2026", + "related_inputs": {"professional_income": "input_professional_income"} +} +``` + +The predicate and key use the compiler's canonical scoped names. Explicit tuple +coordinates are checked against compiled current/related slots and the source +relation declaration. The expected related input map must match the union of +compiled schemas sharing that key. `one` means exactly one selected person per +root; `many` permits zero. A role selector requires both its column and distinct +role values. This interface supports person members of genuine group partitions; +binding-local views and arbitrary links remain separate work. + +Rules nodes declare whole-table slices for related physical IDs, partition +membership, selector columns and every mapped related input; composition refuses +a missing dependency. The adapter gathers per-root offsets and input arrays, +ordering people by physical identity and retaining distinct equal-valued people. +Data-derived batches are built at execution and never enter prepared adapter +configuration. Relation-enabled execution requires the native dense API with +`DenseRelationBatch.related_row_count`, including for relations with no related +input columns. Undeclared compiled relations continue to refuse. + +Counterfactual and solve-zero bridges currently reject overrides of aliased +related inputs, where the engine input name differs from its physical column. +Either spelling of the alias refuses before native execution. Resolving these +overrides is deferred; ordinary relation materialization uses the declared map. + +Relation declarations enter the binding and engine identity. An absent +declaration remains absent; an explicitly supplied empty declaration is distinct. +Bindings without relations retain their existing engine-reference serialization +and native execution call shape. Calibration roster counts use the same frame +unit framework, with frozen-kernel differential tests preserving NZ results. diff --git a/packages/microcosm-build/src/microcosm/build/calibration_columns.py b/packages/microcosm-build/src/microcosm/build/calibration_columns.py index c41ee074d1..d1f5614921 100644 --- a/packages/microcosm-build/src/microcosm/build/calibration_columns.py +++ b/packages/microcosm-build/src/microcosm/build/calibration_columns.py @@ -16,7 +16,6 @@ from microcosm.build.calibration_predicates import ( _finite_number, - _identity, predicate_mask, validate_calibration_inputs, ) @@ -24,6 +23,7 @@ interval_indicators, validate_interval_document, ) +from microcosm.frame.unit_construction import prepare_roster_counts def prepare_calibration_columns( @@ -44,7 +44,7 @@ def prepare_calibration_columns( required, columns = validate_calibration_inputs(tables, bindings) interval_columns = _interval_columns(tables, columns, required, intervals) out = {entity: table.copy(deep=True) for entity, table in tables.items()} - derived = _prepare_roster_counts(out, bindings.get("roster_counts", ()), required) + derived = prepare_roster_counts(out, bindings.get("roster_counts", ()), required) prepared: dict[str, dict[str, object]] = {} seen: set[tuple[str, str]] = set() for entry in columns: @@ -164,58 +164,3 @@ def _interval_columns( for entity, name, _ in outputs: result[entity, name] = computed[table_id][name] return result - - -def _prepare_roster_counts( - tables: dict[str, pd.DataFrame], entries: object, required: Mapping[str, Any] -) -> dict[str, set[str]]: - if not isinstance(entries, (list, tuple)): - raise ValueError("Calibration roster_counts must be a list.") - derived: dict[str, set[str]] = {} - for entry in entries: - if not isinstance(entry, Mapping): - raise ValueError("Calibration roster count must be an object.") - entity, members, name = ( - entry.get("entity"), - entry.get("member_entity"), - entry.get("name"), - ) - group_id, member_id, membership = ( - entry.get("entity_id"), - entry.get("member_id"), - entry.get("membership"), - ) - if ( - entity not in required - or members not in required - or not isinstance(name, str) - or not name - or name.endswith("_id") - or name in required[entity] - or name in derived.get(entity, ()) - or group_id not in required[entity] - or member_id not in required[members] - or membership not in required[members] - or not isinstance(entry.get("representation"), str) - or not entry["representation"] - ): - raise ValueError( - "Calibration roster count needs declared identities and representation." - ) - groups = [_identity(value) for value in tables[entity][group_id].tolist()] - identities = [_identity(value) for value in tables[members][member_id].tolist()] - if len(set(groups)) != len(groups) or len(set(identities)) != len(identities): - raise ValueError( - "Calibration roster contains duplicate membership identities." - ) - counts = dict.fromkeys(groups, 0) - for value in tables[members][membership].tolist(): - key = _identity(value) - if key not in counts: - raise ValueError("Calibration roster contains orphan membership.") - counts[key] += 1 - if any(count == 0 for count in counts.values()): - raise ValueError("Calibration roster contains empty households or groups.") - tables[entity][name] = [counts[key] for key in groups] - derived.setdefault(entity, set()).add(name) - return derived diff --git a/packages/microcosm-build/src/microcosm/build/transport/cli.py b/packages/microcosm-build/src/microcosm/build/transport/cli.py index 844d32ccc2..8294c5037c 100644 --- a/packages/microcosm-build/src/microcosm/build/transport/cli.py +++ b/packages/microcosm-build/src/microcosm/build/transport/cli.py @@ -42,7 +42,6 @@ from microcosm.frame.adapters.axiom import AxiomEntityTableDataset from microcosm.frame.materialize import put_frame_table from microcosm.frame.rules import RulesEngine -from microcosm.frame.unit_construction import BenefitUnitRule from microcosm.graph import ( CompiledGraph, ContentStore, @@ -76,9 +75,13 @@ transport_resource, validate_transport_activation, ) -from .population_kernels import TRANSPORT_CREATE +from .population_kernels import ( + TRANSPORT_CREATE, + TRANSPORT_CREATE_DECLARED, +) from .registry import TransportRegistry, build_transport_registry from .terminal_kernels import materialize_export, validate_export +from .unit_declarations import CREATE_V1_REF, create_unit_declaration __all__ = [ "EXPORT_FILENAME", @@ -138,6 +141,11 @@ def prepare_create_outputs( if resume == "require" and inventory_store is None: raise StoreMissError("--resume require needs a CREATE inventory store.") declaration = skeleton_create_declaration(spec) + kernel = ( + TRANSPORT_CREATE + if declaration.kernel == CREATE_V1_REF + else TRANSPORT_CREATE_DECLARED + ) missing = sorted(set(declaration.sources) - set(sources)) if missing: raise ValueError(f"CREATE needs local paths for sources {missing}.") @@ -148,7 +156,7 @@ def prepare_create_outputs( canonical_json( { "kernel": declaration.kernel, - "implementation_hash": TRANSPORT_CREATE.implementation_hash(), + "implementation_hash": kernel.implementation_hash(), "params": declaration.params, "sources": { name: source_content_key(name, sources[name]) @@ -181,7 +189,7 @@ def prepare_create_outputs( rng=np.random.default_rng(0), sources={name: Path(sources[name]).resolve() for name in probe.sources}, ) - frame = TRANSPORT_CREATE.run(context).frame + frame = kernel.run(context).frame schema = frame.schema structural = { schema.person_entity: { @@ -279,7 +287,9 @@ def prepare_transport_build( if rules_source not in paths: raise ValueError(f"Source {rules_source!r} (the RuleSpec tree) needs a path.") create = skeleton_create_declaration(spec) - rule = BenefitUnitRule.from_spec(json.loads(create.params["unit_rule"])) + rule = create_unit_declaration( + create.kernel, json.loads(create.params["unit_rule"]) + ) registry = build_transport_registry( transport_resource(spec, "axiom_rules_bindings"), paths[rules_source], diff --git a/packages/microcosm-build/src/microcosm/build/transport/column_kernels.py b/packages/microcosm-build/src/microcosm/build/transport/column_kernels.py index fd062678db..a1a482479b 100644 --- a/packages/microcosm-build/src/microcosm/build/transport/column_kernels.py +++ b/packages/microcosm-build/src/microcosm/build/transport/column_kernels.py @@ -54,9 +54,11 @@ resolve_judgments, ) from microcosm.frame.unit_construction import ( + MULTI_UNIT_RULE_FORMAT, BenefitUnitRule, - benefit_unit_attributes, + MultiUnitRule, benefit_unit_membership, + declared_unit_attributes, units_per_household, ) from microcosm.graph import ( @@ -138,8 +140,11 @@ def _document(context: KernelContext, ref: str, name: str): return canonical_document_param(context, ref, name) -def _rule(context: KernelContext, ref: str) -> BenefitUnitRule: - return BenefitUnitRule.from_spec(_document(context, ref, "unit_rule")) +def _unit_declaration(context: KernelContext, ref: str) -> MultiUnitRule: + document = _document(context, ref, "unit_rule") + if document.get("format") == MULTI_UNIT_RULE_FORMAT: + return MultiUnitRule.from_dict(document) + return MultiUnitRule.from_benefit_unit(BenefitUnitRule.from_spec(document)) def _paths(context: KernelContext, ref: str) -> tuple[str, ...] | None: @@ -314,31 +319,35 @@ def run(self, context: KernelContext) -> KernelResult: require_params( context, self.ref, required=frozenset({"unit_rule", "unit_rule_sha256"}) ) - rule = _rule(context, self.ref) - tables = _whole_tables(context, self.ref, ("person", "household", rule.entity)) - person, household, units = ( - tables["person"], - tables["household"], - tables[rule.entity], - ) - attributes = benefit_unit_attributes( - person, units, person[rule.membership_column], rule - ) - # These are aliases used by input-bridge declarations, with the same - # adult composition as the wrapped function's is_couple attribute. - attributes["is_partnered"] = attributes["is_couple"] - attributes["is_single"] = ~attributes["is_couple"] + declaration = _unit_declaration(context, self.ref) + entities = tuple(item.entity for item in declaration.units) + tables = _whole_tables(context, self.ref, ("person", "household", *entities)) + person, household = tables["person"], tables["household"] household_attributes = household.loc[:, ["household_id"]].copy() - household_attributes[f"n_{rule.entity}_units"] = units_per_household( - household, units, rule - ) + person_attributes = person.loc[:, ["person_id"]].copy() + outputs = {"household": household_attributes, "person": person_attributes} + counts = {} + for item in declaration.units: + rule, units = item.rule, tables[item.entity] + attributes = declared_unit_attributes( + person, units, person[rule.membership_column], item + ) + attributes["is_partnered"] = attributes["is_couple"] + attributes["is_single"] = ~attributes["is_couple"] + household_attributes[f"n_{rule.entity}_units"] = units_per_household( + household, units, rule + ) + outputs[item.entity] = attributes + membership = benefit_unit_membership( + person, units, person[rule.membership_column], rule + ) + person_attributes[f"{item.entity}_role"] = membership.person_role + counts[item.entity] = len(units) return KernelResult( - columns=_columns( - context, - self.ref, - {rule.entity: attributes, "household": household_attributes}, - ), - receipt={"entity": rule.entity, "n_units": len(units)}, + columns=_columns(context, self.ref, outputs), + receipt={"entity": entities[0], "n_units": counts[entities[0]]} + if len(entities) == 1 + else {"units": counts}, ) @@ -411,15 +420,31 @@ def run(self, context: KernelContext) -> KernelResult: _GROUP_REQUIRED = _ENCODE_REQUIRED | frozenset( {"unit_rule", "unit_rule_sha256", "engine_entity"} ) -_GROUP_OPTIONAL = _ENCODE_OPTIONAL | frozenset({"knobs", "scenario_sha256"}) +_GROUP_OPTIONAL = _ENCODE_OPTIONAL | frozenset( + {"knobs", "scenario_sha256", "unit_entity"} +) def _group_inputs(context: KernelContext, ref: str): - rule = _rule(context, ref) - tables = _whole_tables(context, ref, ("person", "household", rule.entity)) # The whole mapping: encode_groups selects modules itself and checks # state bindings against every concept binding's input name. mapping = _full_mapping(context, ref) + entity = string_param(context, ref, "engine_entity") + declaration = _unit_declaration(context, ref) + physical = context.params.get("unit_entity") + if physical is None: + if len(declaration.units) != 1: + raise ValueError( + f"{ref} requires unit_entity when several units are declared." + ) + matches = list(declaration.units) + else: + physical = string_param(context, ref, "unit_entity") + matches = [item for item in declaration.units if item.entity == physical] + if len(matches) != 1: + raise ValueError(f"{ref} needs exactly one declared unit for {physical!r}.") + rule = matches[0].rule + tables = _whole_tables(context, ref, ("person", "household", rule.entity)) closure = _closure(context, ref) knobs = _optional_document(context, ref, "knobs") if knobs and closure is None: @@ -428,7 +453,6 @@ def _group_inputs(context: KernelContext, ref: str): sha256_param(context, ref, "scenario_sha256") elif "scenario_sha256" in context.params: sha256_param(context, ref, "scenario_sha256") - entity = string_param(context, ref, "engine_entity") membership = benefit_unit_membership( tables["person"], tables[rule.entity], diff --git a/packages/microcosm-build/src/microcosm/build/transport/compose.py b/packages/microcosm-build/src/microcosm/build/transport/compose.py index 6f7b9a2d25..dd2a6b2e40 100644 --- a/packages/microcosm-build/src/microcosm/build/transport/compose.py +++ b/packages/microcosm-build/src/microcosm/build/transport/compose.py @@ -117,7 +117,14 @@ MONEY_NORMALIZE_REF, ) from .reference_selection import selection_resources, validate_exemptions +from .registry import binding_relation_dependencies from .target_kernels import TARGET_COMPILATION_KERNELS +from .unit_declarations import ( + CREATE_REFS, + CREATE_V1_REF, + CREATE_V2_REF, + create_unit_declaration, +) __all__ = [ "GRAPH_RESOURCE", @@ -146,9 +153,8 @@ RULES_ROUTER = "simulate.rules_by_ref@1" #: The kernels a skeleton declaration may use. A kernel outside this set #: (for example an entitlement bridge) enters only through an extension. -SKELETON_KERNELS = frozenset( +SKELETON_KERNELS = CREATE_REFS | frozenset( { - "transport.create@1", "transport.boundary@1", "transport.currency@1", "transport.quantile_map@1", @@ -1073,11 +1079,12 @@ def validate_transport_activation(spec: Mapping[str, object]) -> None: if _has_resource(spec, "reference_selection") else {} ) - selections = ( - tuple(_declared_resource_selections(transport_resource(spec, GRAPH_RESOURCE))) + graph_document = ( + transport_resource(spec, GRAPH_RESOURCE) if _has_resource(spec, GRAPH_RESOURCE) - else () + else {} ) + selections = tuple(_declared_resource_selections(graph_document)) selections = [(selection, False, None) for selection in selections] if _has_resource(spec, "entitlement_graph") and _has_resource(spec, "scenarios"): from .entitlement import _rows @@ -1116,10 +1123,22 @@ def validate_transport_activation(spec: Mapping[str, object]) -> None: f"Resource {selection['resource']!r} path {selection.get('path', [])!r}: " f"{error}" ) - for name in ("benefit_unit_rule", "axiom_rules_bindings", "scenarios"): + versioned_create = any( + isinstance(row, Mapping) and row.get("kernel") == CREATE_V2_REF + for row in graph_document.get("nodes", ()) + ) + required_resources = ("axiom_rules_bindings", "scenarios") + if not versioned_create: + required_resources += ("benefit_unit_rule",) + for name in required_resources: if not _has_resource(spec, name): missing.append(f"{name}.json") - if _has_resource(spec, "benefit_unit_rule"): + if versioned_create: + try: + _versioned_create_units(spec) + except (ValueError, TypeError, KeyError) as error: + missing.append(f"unit declarations: {error}") + elif _has_resource(spec, "benefit_unit_rule"): unit = transport_resource(spec, "benefit_unit_rule") child = _mapping( unit.get("dependent_child", {}), "benefit_unit_rule.dependent_child" @@ -1337,10 +1356,31 @@ def _embeds_engine_ref( def _unit_entities(spec: Mapping[str, object]) -> tuple[object, ...]: - """The declared benefit-unit entities (S2-d PR 1 adds every unit kind).""" + """Every declared partition, including the legacy benefit-unit alias.""" + declaration = _versioned_create_units(spec) + if declaration is not None: + return tuple(item.entity for item in declaration.units) return (transport_resource(spec, "benefit_unit_rule").get("entity"),) +def _versioned_create_units(spec: Mapping[str, object]): + if not _has_resource(spec, GRAPH_RESOURCE): + return None + rows = [ + row + for row in transport_graph_document(spec).get("nodes", ()) + if isinstance(row, Mapping) and row.get("kernel") == CREATE_V2_REF + ] + if not rows: + return None + if len(rows) != 1: + raise ValueError("A transport skeleton declares exactly one CREATE.") + document = _param(spec, rows[0].get("params", {}).get("unit_rule"), None) + if not isinstance(document, str): + raise ValueError("Declared CREATE requires a canonical JSON unit_rule.") + return create_unit_declaration(CREATE_V2_REF, json.loads(document)) + + def _is_unit_entity(entity: object, unit_entities: tuple[object, ...]) -> bool: # Compare, never hash, so a malformed entity behaves as it always did. return any(entity == unit for unit in unit_entities) @@ -2761,6 +2801,7 @@ def transport_rules_node( if not isinstance(node, Node): raise TypeError("A transport rules node must be a Node.") _safe_node_id(node.id) + _require_relation_inputs(spec, binding_id, node.inputs) declaration = { "id": node.id, "kernel": node.kernel, @@ -2773,6 +2814,35 @@ def transport_rules_node( ) +def _require_relation_inputs(spec, binding_id, inputs) -> None: + document = transport_resource(spec, "axiom_rules_bindings") + bindings = [ + row for row in document.get("bindings", ()) if row.get("id") == binding_id + ] + if len(bindings) != 1 or "relations" not in bindings[0]: + return + declaration = _versioned_create_units(spec) + if declaration is None: + declaration = create_unit_declaration( + CREATE_V1_REF, transport_resource(spec, "benefit_unit_rule") + ) + schema = declaration.entity_schema()[0] + dependencies = binding_relation_dependencies( + bindings[0], schema=schema, entity_names=document["entity_names"] + ) + declared = { + (item.entity, column) + for item in inputs + if item.rows == ROWS_ALL + for column in item.columns + } + missing = set(dependencies) - declared + if missing: + raise ValueError( + f"Rules binding {binding_id!r} needs whole-table relation inputs {sorted(missing)}." + ) + + def _check_extension_checkpoints(spec, config, graph, compiled) -> None: checkpoints = transport_endpoints(spec, extensions=config.extensions)["checkpoints"] if not set(checkpoints) <= set(compiled.order): @@ -2915,6 +2985,8 @@ def _node( Slice(**{**value, "columns": tuple(value["columns"])}) for value in _sequence(row.get("inputs", ()), "input slices") ) + if "rules_binding" in row: + _require_relation_inputs(spec, row["rules_binding"], inputs) artifacts = [] for value in _sequence(row.get("artifact_inputs", ()), "artifact inputs"): edge = _mapping(value, "artifact input") @@ -2977,9 +3049,9 @@ def skeleton_create_declaration(spec: Mapping[str, object]) -> CreateDeclaration ) if isinstance(row, Mapping) and row.get("structural") == "create" ] - if len(rows) != 1 or rows[0].get("kernel") != "transport.create@1": + if len(rows) != 1 or rows[0].get("kernel") not in CREATE_REFS: raise ValueError( - "A transport skeleton declares exactly one transport.create@1." + "A transport skeleton declares exactly one transport.create@1 or @2." ) row = rows[0] return CreateDeclaration( diff --git a/packages/microcosm-build/src/microcosm/build/transport/population_kernels.py b/packages/microcosm-build/src/microcosm/build/transport/population_kernels.py index ed6ba2368f..f8d24121fd 100644 --- a/packages/microcosm-build/src/microcosm/build/transport/population_kernels.py +++ b/packages/microcosm-build/src/microcosm/build/transport/population_kernels.py @@ -32,7 +32,7 @@ import microcosm.frame.transport as transport_module import microcosm.frame.unit_construction as units_module import microcosm.frame.weights as weights_module -from microcosm.frame import EntitySchema, Frame, WeightKind, Weights +from microcosm.frame import Frame, WeightKind, Weights from microcosm.frame.concepts import ( concept_for_column, concept_schema_sha256, @@ -45,9 +45,8 @@ read_populace_us_donor, ) from microcosm.frame.unit_construction import ( - BenefitUnitRule, benefit_unit_diagnostics, - build_benefit_units, + build_declared_units, ) from microcosm.graph import ( ROWS_ALL, @@ -64,7 +63,7 @@ ) from microcosm.graph.population import dtype_for_token -from . import graph_inputs, target_kernels +from . import graph_inputs, target_kernels, unit_declarations from .artifact_types import TARGET_SURFACE_TYPE from .graph_inputs import ( canonical_document_param, @@ -73,17 +72,21 @@ string_param, ) from .target_kernels import compile_target_surface, decode_target_surface +from .unit_declarations import CREATE_V1_REF, CREATE_V2_REF, create_unit_declaration __all__ = [ "TRANSPORT_BOUNDARY", "TRANSPORT_CREATE", + "TRANSPORT_CREATE_DECLARED", "TRANSPORT_CURRENCY", "TRANSPORT_QUANTILE_MAP", "TransportBoundaryKernel", "TransportCreateKernel", + "TransportCreateDeclaredKernel", "TransportCurrencyKernel", "TransportQuantileMapKernel", "register_population_kernels", + "create_unit_declaration", ] @@ -114,7 +117,7 @@ def _column_result( class TransportCreateKernel(KernelBase): """``transport.create@1``: verified support records at destination mass.""" - ref = "transport.create@1" + ref = CREATE_V1_REF capabilities = Capabilities( Determinism.DETERMINISTIC, numeric=Numeric.BITWISE, @@ -144,6 +147,7 @@ def implementation_hash(self) -> str: type(self), graph_inputs, target_kernels, + unit_declarations, transport_module, units_module, concepts_module, @@ -227,17 +231,17 @@ def run(self, context: KernelContext) -> KernelResult: item = concept_for_column(entity, column) if item is not None and item.dtype == "str": table[column] = table[column].astype(text) - rule = BenefitUnitRule.from_spec( - canonical_document_param(context, self.ref, "unit_rule") + declaration = create_unit_declaration( + self.ref, canonical_document_param(context, self.ref, "unit_rule") ) - if rule.entity != "family": - raise ValueError("CREATE's declared schema requires family benefit units.") - family, membership = build_benefit_units( - tables["person"], tables["household"], rule + partitions = build_declared_units( + tables["person"], tables["household"], declaration ) - tables["family"] = family person, household = tables["person"], tables["household"] - person[rule.membership_column] = membership + for unit in declaration.units: + membership = partitions[unit.rule.entity] + tables[unit.rule.entity] = membership.units + person[unit.rule.membership_column] = membership.person_unit person["take_up_seed"] = derive_transport_seed( donor.source_person_ids, string_param(context, self.ref, "seed_stream") ) @@ -261,11 +265,9 @@ def run(self, context: KernelContext) -> KernelResult: dtype=object, ) surface = decode_target_surface(surface_bytes) - # The persons, their households and the rule's benefit units: the - # entity check above makes this the destination's declared schema. frame = Frame( tables, - EntitySchema(group_entities=("household", rule.entity)), + declaration.entity_schema()[0], {"household": weights}, strata, metadata={ @@ -276,6 +278,17 @@ def run(self, context: KernelContext) -> KernelResult: "donor_time_period": donor.time_period, }, ) + diagnostics = { + unit.rule.entity: benefit_unit_diagnostics( + person, + household, + partitions[unit.rule.entity].units, + partitions[unit.rule.entity].person_unit, + unit.rule, + ) + for unit in declaration.units + if unit.kind == "benefit_unit" + } return KernelResult( frame=frame, receipt={ @@ -292,16 +305,24 @@ def run(self, context: KernelContext) -> KernelResult: "design_scale": scale, "n_persons": len(person), "n_households": len(household), - "n_units": len(family), - "unit_construction": benefit_unit_diagnostics( - person, household, family, membership, rule + "n_units": sum( + len(partition.units) for partition in partitions.values() ), + "unit_construction": next(iter(diagnostics.values())) + if len(declaration.units) == 1 and diagnostics + else diagnostics, "mass_reference": registry.specs[0].name, "facts": dict(surface.facts), }, ) +class TransportCreateDeclaredKernel(TransportCreateKernel): + """Versioned, typed declarations; the legacy CREATE lifts one benefit unit.""" + + ref = CREATE_V2_REF + + class TransportBoundaryKernel(KernelBase): """``transport.boundary@1``: keep every person and conserve mass.""" @@ -586,6 +607,7 @@ def run(self, context: KernelContext) -> KernelResult: TRANSPORT_CREATE = TransportCreateKernel() +TRANSPORT_CREATE_DECLARED = TransportCreateDeclaredKernel() TRANSPORT_BOUNDARY = TransportBoundaryKernel() TRANSPORT_CURRENCY = TransportCurrencyKernel() TRANSPORT_QUANTILE_MAP = TransportQuantileMapKernel() @@ -595,6 +617,7 @@ def register_population_kernels(registry: KernelRegistry) -> KernelRegistry: """Register the population kernels explicitly and idempotently.""" for kernel in ( TRANSPORT_CREATE, + TRANSPORT_CREATE_DECLARED, TRANSPORT_BOUNDARY, TRANSPORT_CURRENCY, TRANSPORT_QUANTILE_MAP, diff --git a/packages/microcosm-build/src/microcosm/build/transport/registry.py b/packages/microcosm-build/src/microcosm/build/transport/registry.py index e13f2af096..82a5cec524 100644 --- a/packages/microcosm-build/src/microcosm/build/transport/registry.py +++ b/packages/microcosm-build/src/microcosm/build/transport/registry.py @@ -26,6 +26,7 @@ from microcosm.frame.adapters.axiom import ( AxiomEngine, AxiomPeriod, + AxiomRelation, assert_no_relations, axiom_engine_ref, rulespec_tree_digest, @@ -36,7 +37,7 @@ SimulateRulesByRefKernel, SimulateSolveZeroKernel, ) -from microcosm.frame.unit_construction import BenefitUnitRule +from microcosm.frame.unit_construction import BenefitUnitRule, MultiUnitRule from microcosm.graph import KernelRegistry from microcosm.graph.canonical import canonical_json from microcosm.graph.codecs import SourceCodecRegistry @@ -56,6 +57,7 @@ __all__ = [ "PYTHON_ENGINE_REF_FORMAT", "TransportRegistry", + "binding_relation_dependencies", "build_transport_registry", "register_transport_population_kernels", "transport_entity_schema", @@ -86,18 +88,21 @@ class TransportRegistry: def transport_entity_schema( - rule: BenefitUnitRule, + rule: BenefitUnitRule | MultiUnitRule, ) -> tuple[EntitySchema, Mapping[str, str]]: - """The persons, their households and the rule's benefit units. + """The persons, households and declared partitions CREATE builds. - This is the schema ``transport.create@1`` gives its frame. Each unit is - built inside one household (``build_benefit_units`` returns - ``_household_id``), so the unit nests in the household. + Both CREATE versions derive their physical schema from the shared unit + declarations. Every constructed unit nests inside one household; the + legacy benefit-unit interface lifts to a one-element declaration list. """ - if not isinstance(rule, BenefitUnitRule): - raise TypeError("transport_entity_schema requires a BenefitUnitRule.") - schema = EntitySchema(group_entities=("household", rule.entity)) - return schema, MappingProxyType({rule.entity: "household"}) + if isinstance(rule, BenefitUnitRule): + rule = MultiUnitRule.from_benefit_unit(rule) + if not isinstance(rule, MultiUnitRule): + raise TypeError( + "transport_entity_schema requires a BenefitUnitRule or MultiUnitRule." + ) + return rule.entity_schema() def _mapping(value: object, label: str) -> Mapping: @@ -165,7 +170,7 @@ def _binding_semantics( if not variables or len(variables) != len(set(variables)): raise ValueError(f"{label}.variables must be non-empty and distinct.") # Metadata prose and declarations do not decide engine outputs. - return { + semantics = { "id": row["id"], "rulespec_path": relative, "sha256": digest, @@ -174,6 +179,120 @@ def _binding_semantics( "period": period, "variables": variables, } + if "views" in row: + raise NotImplementedError(f"{label}.views are not supported.") + if "relations" in row: + semantics["relations"] = tuple( + relation.to_dict() + for relation in _binding_relations(row, entity_names, schema) + ) + return semantics + + +def _binding_relations( + row: Mapping, + entity_names: Mapping[str, str], + schema: EntitySchema, +) -> tuple[AxiomRelation, ...]: + """Parse explicit partition relations without inventing absent declarations. + + These are person members of a declared group partition. Arbitrary links + and binding-local logical views require a separate declaration contract. + Compiled key, slot orientation and input surfaces are checked after the + adapter is pinned to its RuleSpec bytes. + """ + label = f"rules binding {row['id']!r}.relations" + values = row["relations"] + if isinstance(values, str | bytes) or not isinstance(values, Sequence): + raise TypeError(f"{label} must be a list of relation manifests.") + required = { + "predicate", + "key", + "current_slot", + "related_slot", + "arguments", + "root_entity", + "related_entity", + "membership_column", + "related_inputs", + "cardinality", + "period", + } + optional = {"role_column", "roles"} + relations = [] + for value in values: + manifest = _mapping(value, f"{label} manifest") + if not required <= set(manifest) or set(manifest) - required - optional: + raise ValueError( + f"{label} manifests require exactly {sorted(required)} and " + "optional role_column and roles." + ) + data = dict(manifest) + for name in ("arguments", "roles"): + if name not in data: + continue + items = data[name] + if isinstance(items, str | bytes) or not isinstance(items, Sequence): + raise TypeError(f"{label}.{name} must be a list of strings.") + data[name] = tuple(items) + data["related_inputs"] = dict( + _mapping(data["related_inputs"], f"{label}.related_inputs") + ) + relation = AxiomRelation(**data) + if relation.root_entity != row["entity"]: + raise ValueError(f"{label}.root_entity must equal the binding entity.") + if ( + relation.root_entity not in schema.group_entities + or relation.related_entity != schema.person_entity + ): + raise ValueError( + f"{label} supports a declared group partition with related persons." + ) + if relation.membership_column != schema.membership_column(relation.root_entity): + raise ValueError( + f"{label}.membership_column must name the declared partition membership." + ) + if relation.period != row["period"]: + raise ValueError(f"{label}.period must equal the binding period.") + if ( + relation.arguments[relation.current_slot] + != entity_names[relation.root_entity] + or relation.arguments[relation.related_slot] + != entity_names[relation.related_entity] + ): + raise ValueError( + f"{label}.arguments disagree with the declared entity names and slots." + ) + relations.append(relation) + keys = [relation.key for relation in relations] + if len(keys) != len(set(keys)): + raise ValueError(f"{label} repeats a compiled relation key.") + return tuple(relations) + + +def binding_relation_dependencies( + binding: Mapping, + *, + entity_names: Mapping[str, str], + schema: EntitySchema, +) -> tuple[tuple[str, str], ...]: + """The physical graph slices an explicit relation declaration reads. + + Composition uses the same parser and physical dependency implementation + as preparation and native batching. No declaration contributes no extra + columns; identifiers, memberships and role selectors enter when present. + """ + if "relations" not in binding: + return () + return tuple( + sorted( + { + column + for relation in _binding_relations(binding, entity_names, schema) + for column in relation.input_dependencies(schema) + } + ) + ) def _check_engine( @@ -207,6 +326,13 @@ def _check_engine( raise ValueError( f"{label} engine must use declared periods, unit nesting and graph dtypes." ) + expected_relations = ( + None + if "relations" not in semantics + else _binding_relations(semantics, entity_names, schema) + ) + if engine._relations != expected_relations: + raise ValueError(f"{label} engine relations do not match the binding.") else: declared_root = getattr(engine, "rulespec_root", None) if ( @@ -222,13 +348,15 @@ def _check_engine( ) if not callable(getattr(engine, "cache_identity", None)): raise TypeError(f"{label} injected engine must expose cache_identity().") + if "relations" in semantics: + raise TypeError(f"{label} relation manifests require an AxiomEngine.") def build_transport_registry( rules_bindings: Mapping, rulespec_root: str | Path, *, - unit_rule: BenefitUnitRule, + unit_rule: BenefitUnitRule | MultiUnitRule, engines_by_binding: Mapping[str, RulesEngine] | None = None, dependencies: tuple[str, ...] = (), ) -> TransportRegistry: @@ -325,6 +453,11 @@ def build_transport_registry( entity_names=entity_names, nesting=nesting, output_dtypes="graph", + **( + {"relations": _binding_relations(semantics, entity_names, schema)} + if "relations" in semantics + else {} + ), ) if engines_by_binding is None else engines_by_binding[row["id"]] @@ -334,7 +467,10 @@ def build_transport_registry( ) if isinstance(engine, AxiomEngine): reference = axiom_engine_ref(engine, rulespec_root=root, **pins) - assert_no_relations(engine, semantics["entity"]) + if "relations" in semantics: + engine.assert_declared_relations(semantics["entity"]) + else: + assert_no_relations(engine, semantics["entity"]) else: identity = _mapping( engine.cache_identity(), diff --git a/packages/microcosm-build/src/microcosm/build/transport/target_kernels.py b/packages/microcosm-build/src/microcosm/build/transport/target_kernels.py index 6e29496e0e..7b3d926540 100644 --- a/packages/microcosm-build/src/microcosm/build/transport/target_kernels.py +++ b/packages/microcosm-build/src/microcosm/build/transport/target_kernels.py @@ -65,6 +65,7 @@ import microcosm.calibrate.target as target_module import microcosm.frame.bundle as frame_bundle_module import microcosm.frame.schema as frame_schema_module +import microcosm.frame.unit_construction as units_module import microcosm.frame.weights as frame_weights_module from microcosm.build.calibration_columns import prepare_calibration_columns from microcosm.build.country_spec import _validate_target_references @@ -116,6 +117,7 @@ string_param, ) from .reference_selection import exemption_rows, validate_exemptions +from .unit_declarations import CREATE_REFS __all__ = [ "TARGET_COMPILATION_KERNELS", @@ -138,8 +140,12 @@ ] #: Every transport kernel that resolves references against consumer facts. -TARGET_COMPILATION_KERNELS = frozenset( - {"transport.create@1", "targets.compile@1", "takeup.compare@1", "takeup.gap@1"} +TARGET_COMPILATION_KERNELS = CREATE_REFS | frozenset( + { + "targets.compile@1", + "takeup.compare@1", + "takeup.gap@1", + } ) _SURFACE_KIND = "transport_target_surface" @@ -603,6 +609,7 @@ def implementation_hash(self) -> str: graph_inputs, artifact_types, calibration_columns_module, + units_module, calibration_predicates_module, reference_intervals_module, reference_factors_module, diff --git a/packages/microcosm-build/src/microcosm/build/transport/unit_declarations.py b/packages/microcosm-build/src/microcosm/build/transport/unit_declarations.py new file mode 100644 index 0000000000..1d6325354c --- /dev/null +++ b/packages/microcosm-build/src/microcosm/build/transport/unit_declarations.py @@ -0,0 +1,30 @@ +"""CREATE version identities and their shared unit-declaration parser. + +This module depends only on the Frame declaration framework, so population +kernels, target compilation, composition and CLI preparation share one contract +without importing each other's execution or registration modules. +""" + +from collections.abc import Mapping + +from microcosm.frame.unit_construction import BenefitUnitRule, MultiUnitRule + +__all__ = [ + "CREATE_V1_REF", + "CREATE_V2_REF", + "CREATE_REFS", + "create_unit_declaration", +] + +CREATE_V1_REF = "transport.create@1" +CREATE_V2_REF = "transport.create@2" +CREATE_REFS = frozenset({CREATE_V1_REF, CREATE_V2_REF}) + + +def create_unit_declaration(kernel: str, document: Mapping) -> MultiUnitRule: + """Lift legacy benefit declarations or read the versioned declaration list.""" + if kernel == CREATE_V1_REF: + return MultiUnitRule.from_benefit_unit(BenefitUnitRule.from_spec(document)) + if kernel == CREATE_V2_REF: + return MultiUnitRule.from_spec(document) + raise ValueError(f"Unknown transport CREATE kernel {kernel!r}.") diff --git a/packages/microcosm-build/tests/engine_free/shared/test_calibration_columns.py b/packages/microcosm-build/tests/engine_free/shared/test_calibration_columns.py index 943a7b7ada..b8fda11757 100644 --- a/packages/microcosm-build/tests/engine_free/shared/test_calibration_columns.py +++ b/packages/microcosm-build/tests/engine_free/shared/test_calibration_columns.py @@ -836,6 +836,83 @@ def test_prepare_matches_the_frozen_legacy_on_nz_bindings(): assert _prepared_document(live) == _prepared_document(legacy) +@settings( + max_examples=8, + deadline=None, + database=None, + suppress_health_check=[HealthCheck.function_scoped_fixture], +) +@given(ages=st.lists(st.integers(0, 200), min_size=1, max_size=20)) +def test_calibration_kernel_roster_framework_matches_previous_kernel(monkeypatch, ages): + import microcosm.build.transport.target_kernels as kernels + from microcosm.build.transport.artifact_types import CALIBRATION_PREPARATION_TYPE + from microcosm.graph import ArtifactOutput, KernelContext, Node, Owned, Slice + from microcosm.graph.canonical import canonical_json + from test_support.microcosm_build.transport_graph import ( + canonical_text, + reference_document, + ) + from test_support.microcosm_build.transport_population import digest + + tables, bindings = _tables(), _bindings() + _replace_people(tables, ages) + path = ( + paths_for("microcosm-build").tests + / "fixtures/calibration_columns_09a2fce.py.txt" + ) + scope = {} + exec(compile(path.read_text(), str(path), "exec"), scope) + previous_prepare = scope["prepare_calibration_columns"] + expected = previous_prepare(tables, bindings) + inventory = [*bindings["columns"], *bindings.get("roster_counts", ())] + references = reference_document([]) + node = Node( + "prepare", + "calibration.prepare@1", + population="open", + params={ + "bindings": canonical_text(bindings), + "references": canonical_text(references), + "references_sha256": digest(references), + }, + inputs=tuple( + Slice(entity, tuple(names)) + for entity, names in bindings["required_inputs"].items() + ), + outputs=tuple( + Owned( + row["entity"], + row["name"], + str(expected[row["entity"]][row["name"]].dtype), + ) + for row in inventory + ), + artifact_outputs=(ArtifactOutput("preparation", CALIBRATION_PREPARATION_TYPE),), + ) + + def run(): + return kernels.CALIBRATION_PREPARE.run( + KernelContext( + node=node, + tables=copy.deepcopy(tables), + weights={}, + strata=pd.Series(dtype=object), + params=node.params, + rng=np.random.default_rng(0), + ) + ) + + live = run() + with monkeypatch.context() as patch: + patch.setattr(kernels, "prepare_calibration_columns", previous_prepare) + previous = run() + assert live.artifacts["preparation"] == previous.artifacts["preparation"] + assert canonical_json(live.receipt) == canonical_json(previous.receipt) + assert live.columns.keys() == previous.columns.keys() + for key, values in live.columns.items(): + pd.testing.assert_series_equal(values, previous.columns[key], check_exact=True) + + _NZ_DEFECTS = ( "negative_age", "fractional_age", diff --git a/packages/microcosm-build/tests/engine_free/shared/test_nz_spec_package.py b/packages/microcosm-build/tests/engine_free/shared/test_nz_spec_package.py index 2e2389b246..6334f759d5 100644 --- a/packages/microcosm-build/tests/engine_free/shared/test_nz_spec_package.py +++ b/packages/microcosm-build/tests/engine_free/shared/test_nz_spec_package.py @@ -3225,7 +3225,9 @@ def test_json_pointer_keys_use_rfc6901_slash_escaping(self) -> None: assert selected["item"] == "A1" def test_create_and_donor_reader_citations_still_describe_unit_inputs(self) -> None: + from microcosm.build.transport import population_kernels, unit_declarations from microcosm.build.transport.population_kernels import TransportCreateKernel + from microcosm.frame import unit_construction from microcosm.frame.transport import read_populace_us_donor from microcosm.graph.decl import StructuralDelta @@ -3243,14 +3245,50 @@ def test_create_and_donor_reader_citations_still_describe_unit_inputs(self) -> N ) assert TransportCreateKernel.capabilities.structural is StructuralDelta.CREATE # Differential source evidence: CREATE reads the verified concept donor, - # then builds family units using the declared benefit-unit rule. + # then delegates its unchanged NZ rule through the canonical parser and + # the registered benefit-unit membership constructor. create_source = inspect.getsource(TransportCreateKernel) - assert create_source.index("read_populace_us_donor(") < create_source.index( - "build_benefit_units(" + assert ( + create_source.index("read_populace_us_donor(") + < create_source.index("create_unit_declaration(") + < create_source.index("build_declared_units(") ) assert "tables, dropped = split_for_transport(donor.tables)" in create_source - assert "BenefitUnitRule.from_spec(" in create_source - assert 'tables["person"], tables["household"], rule' in create_source + assert 'tables["person"], tables["household"], declaration' in create_source + assert ( + population_kernels.create_unit_declaration + is unit_declarations.create_unit_declaration + ) + assert TransportCreateKernel.ref == unit_declarations.CREATE_V1_REF + parser_source = inspect.getsource(unit_declarations.create_unit_declaration) + assert "if kernel == CREATE_V1_REF:" in parser_source + assert ( + "MultiUnitRule.from_benefit_unit(BenefitUnitRule.from_spec(document))" + in parser_source + ) + declaration = unit_declarations.create_unit_declaration( + TransportCreateKernel.ref, rule + ) + (unit,) = declaration.units + assert unit.kind == "benefit_unit" + assert ( + unit.rule.to_dict() + == unit_construction.BenefitUnitRule.from_spec(rule).to_dict() + ) + rule_type, constructor = unit_construction._UNIT_KINDS[unit.kind] + assert rule_type is unit_construction.BenefitUnitRule + constructor_source = inspect.getsource(constructor) + assert "_heads(rows, _dependents(rows, rule), rule)" in constructor_source + framework_source = inspect.getsource(unit_construction.build_declared_units) + assert "_UNIT_KINDS[item.kind][1](rows, person, item.rule)" in framework_source + assert ( + "_build_partition(person, household, item.rule, rows, head)" + in framework_source + ) + legacy_source = inspect.getsource(unit_construction.build_benefit_units) + assert "_dependents(rows, rule)" in legacy_source + assert "_heads(rows, dependent, rule)" in legacy_source + assert "_build_partition(person, household, rule, rows, head)" in legacy_source donor_source = inspect.getsource(read_populace_us_donor) assert "tables = POLICYENGINE_US_CONCEPT_MAPPING.decode(raw)" in donor_source assert 'tables["person"][name] = values' in donor_source diff --git a/packages/microcosm-build/tests/engine_free/shared/test_transport_column_kernels.py b/packages/microcosm-build/tests/engine_free/shared/test_transport_column_kernels.py index 52ce03b610..8848d99782 100644 --- a/packages/microcosm-build/tests/engine_free/shared/test_transport_column_kernels.py +++ b/packages/microcosm-build/tests/engine_free/shared/test_transport_column_kernels.py @@ -57,6 +57,9 @@ resolve_judgments, ) from microcosm.frame.unit_construction import ( + MultiUnitRule, + TaxUnitRule, + UnitDeclaration, benefit_unit_attributes, benefit_unit_membership, build_benefit_units, @@ -400,6 +403,69 @@ def _context(kernel, frame, params, outputs, *, artifacts=None): ) +@PROPERTY +@given(tables=concept_frames(max_households=2, max_members=3)) +def test_group_encoding_requires_one_selected_declared_unit(tables): + frame = _generated_frame(tables) + params = {**_encode_params(groups=True), "unit_entity": "tax_unit"} + with pytest.raises(ValueError, match="exactly one declared unit"): + CONCEPTS_ENCODE_GROUPS.run( + _context( + CONCEPTS_ENCODE_GROUPS, + frame, + params, + (Owned("family", "input_assets", "float64"),), + ) + ) + declaration = MultiUnitRule( + units=( + UnitDeclaration(kind="benefit_unit", rule=UNIT_RULE), + UnitDeclaration( + kind="tax_unit", + rule=TaxUnitRule( + entity="tax_unit", + pairing_eligibility_column="explicit_eligibility", + provenance={"eligibility": "Synthetic unused declaration."}, + ), + ), + ) + ) + document = declaration.to_dict() + base_params = _encode_params( + groups=True, + knobs={ + "rent_allocation": "reference_unit", + "asset_test": True, + "rent_factor": 1.0, + }, + ) + versioned = { + **base_params, + "unit_rule": canonical_text(document), + "unit_rule_sha256": digest(document), + } + outputs = (Owned("family", "input_assets", "float64"),) + with pytest.raises(ValueError, match="requires unit_entity"): + CONCEPTS_ENCODE_GROUPS.run( + _context(CONCEPTS_ENCODE_GROUPS, frame, versioned, outputs) + ) + expected = CONCEPTS_ENCODE_GROUPS.run( + _context(CONCEPTS_ENCODE_GROUPS, frame, base_params, outputs) + ) + actual = CONCEPTS_ENCODE_GROUPS.run( + _context( + CONCEPTS_ENCODE_GROUPS, + frame, + {**versioned, "unit_entity": "family"}, + outputs, + ) + ) + pd.testing.assert_series_equal( + actual.columns["family", "input_assets"], + expected.columns["family", "input_assets"], + ) + + def _generated_frame(tables): tables, _ = split_for_transport(tables) person = tables["person"].copy() diff --git a/packages/microcosm-build/tests/engine_free/shared/test_transport_declared_units.py b/packages/microcosm-build/tests/engine_free/shared/test_transport_declared_units.py new file mode 100644 index 0000000000..453b53e7e7 --- /dev/null +++ b/packages/microcosm-build/tests/engine_free/shared/test_transport_declared_units.py @@ -0,0 +1,518 @@ +"""CREATE aliases, inventories and population identity for declared partitions.""" + +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, given, settings +from hypothesis import strategies as st + +from microcosm.build.transport.cli import prepare_create_outputs +from microcosm.build.transport.column_kernels import TRANSPORT_UNIT_ATTRIBUTES +from microcosm.build.transport.compose import ( + skeleton_create_declaration, + transport_rules_node, + validate_transport_activation, +) +from microcosm.build.transport.population_kernels import ( + TRANSPORT_CREATE, + TRANSPORT_CREATE_DECLARED, + create_unit_declaration, +) +from microcosm.frame.unit_construction import ( + BenefitUnitRule, + MultiUnitRule, + TaxUnitRule, + UnitDeclaration, + benefit_unit_roles, +) +from microcosm.graph import KernelContext, Node, Owned, Slice +from test_support.microcosm_build.transport_composed import make_composed_fixture +from test_support.microcosm_build.transport_graph import canonical_text +from test_support.microcosm_build.transport_population import ( + UNIT_RULE, + create_node, + digest, + population_graph, + run_population, + write_population_sources, +) +from test_support.paths import paths_for + + +def _run_create(sources, node, kernel): + return kernel.run( + KernelContext( + node=node, + tables={}, + weights={}, + strata=pd.Series(dtype=object), + params=node.params, + rng=np.random.default_rng(0), + sources={"donor": sources.donor, "facts": sources.facts}, + ) + ) + + +def _inventory(frame): + schema = frame.schema + structural = { + "person": { + "person_id", + *(schema.membership_column(entity) for entity in schema.group_entities), + }, + **{entity: {schema.id_column(entity)} for entity in schema.group_entities}, + } + return tuple( + Owned(entity, column, str(frame.table(entity)[column].dtype)) + for entity in schema.entities + for column in frame.table(entity) + if column not in structural[entity] + ) + + +def _versioned(node, declaration): + document = declaration.to_dict() + return replace( + node, + kernel=TRANSPORT_CREATE_DECLARED.ref, + params={ + **node.params, + "unit_rule": canonical_text(document), + "unit_rule_sha256": digest(document), + }, + ) + + +@settings( + max_examples=6, + deadline=None, + database=None, + suppress_health_check=[HealthCheck.function_scoped_fixture, HealthCheck.too_slow], +) +@given(hours=st.integers(0, 60), mass=st.integers(1, 100_000)) +def test_nz_create_alias_is_byte_identical_to_declared_create(tmp_path, hours, mass): + def edit(tables): + tables["person"]["weekly_hours_worked_before_lsr"] = 0 + tables["person"].loc[2, "age"] = 17 + tables["person"].loc[2, "weekly_hours_worked_before_lsr"] = hours + + sources = write_population_sources( + tmp_path / "sources", mass=mass, edit_tables=edit + ) + path = ( + paths_for("microcosm-build").package + / "src/microcosm/build/nz/benefit_unit_rule.json" + ) + document = json.loads(path.read_text()) + rule = BenefitUnitRule.from_spec(document) + legacy = create_node(sources) + legacy = replace( + legacy, + params={ + **legacy.params, + "unit_rule": canonical_text(document), + "unit_rule_sha256": digest(document), + }, + ) + declared = _versioned(legacy, MultiUnitRule.from_benefit_unit(rule)) + first = _run_create(sources, legacy, TRANSPORT_CREATE) + second = _run_create(sources, declared, TRANSPORT_CREATE_DECLARED) + assert first.frame.schema == second.frame.schema + assert first.receipt == second.receipt + assert _inventory(first.frame) == _inventory(second.frame) + for entity in first.frame.entities: + before, after = first.frame.table(entity), second.frame.table(entity) + pd.testing.assert_frame_equal(before, after, check_exact=True) + assert ( + before.to_json(orient="table").encode() + == after.to_json(orient="table").encode() + ) + roles = [] + for frame in (first.frame, second.frame): + roles.append( + benefit_unit_roles( + frame.person, + frame.table("family"), + frame.person[rule.membership_column], + rule, + ) + .to_json() + .encode() + ) + assert roles[0] == roles[1] + assert ( + first.frame.resolve_weights("family").values.tobytes() + == second.frame.resolve_weights("family").values.tobytes() + ) + + +def test_cli_create_versions_have_identical_cold_and_warm_inventories(tmp_path): + from microcosm.graph import ContentStore + + fixture = make_composed_fixture(tmp_path / "fixture") + spec = copy.deepcopy(fixture.spec) + create = next( + row + for row in spec["resources"]["transport_graph"]["nodes"] + if row["kernel"] == TRANSPORT_CREATE.ref + ) + create["kernel"] = TRANSPORT_CREATE_DECLARED.ref + spec["resources"]["benefit_unit_rule"] = MultiUnitRule.from_benefit_unit( + UNIT_RULE + ).to_dict() + assert skeleton_create_declaration(spec).kernel == TRANSPORT_CREATE_DECLARED.ref + store = ContentStore(tmp_path / "inventory") + legacy = prepare_create_outputs(fixture.spec, fixture.sources) + cold = prepare_create_outputs(spec, fixture.sources, inventory_store=store) + warm = prepare_create_outputs( + spec, fixture.sources, inventory_store=store, resume="require" + ) + assert legacy == cold == warm + + +@settings( + max_examples=4, + deadline=None, + database=None, + suppress_health_check=[HealthCheck.function_scoped_fixture, HealthCheck.too_slow], +) +@given(eligible=st.booleans()) +def test_tax_partition_membership_changes_population_identity(tmp_path, eligible): + # A synthetic boolean uses a carried concept to exercise CREATE plumbing; + # this fixture makes no legal claim about disability or Belgian eligibility. + def edit(tables): + tables["person"]["is_disabled"] = eligible + + sources = write_population_sources(tmp_path / "sources", edit_tables=edit) + declaration = MultiUnitRule( + units=( + UnitDeclaration(kind="benefit_unit", rule=UNIT_RULE), + UnitDeclaration( + kind="tax_unit", + rule=TaxUnitRule( + entity="tax_unit", + pairing_eligibility_column="has_disability", + provenance={ + "eligibility": "Synthetic explicit boolean for plumbing only." + }, + ), + ), + ) + ) + graph = population_graph(sources, through="boundary") + node = _versioned(graph.nodes[0], declaration) + frame = _run_create(sources, node, TRANSPORT_CREATE_DECLARED).frame + node = replace(node, outputs=_inventory(frame)) + graph = replace(graph, nodes=(node,)) + first = run_population(tmp_path / "first", sources, graph=graph) + old_ids = frame.person["person_tax_unit_id"].tolist() + + def opposite(tables): + tables["person"]["is_disabled"] = not eligible + + changed_sources = write_population_sources( + tmp_path / "changed", edit_tables=opposite + ) + changed = _versioned(create_node(changed_sources), declaration) + changed_frame = _run_create( + changed_sources, changed, TRANSPORT_CREATE_DECLARED + ).frame + changed = replace(changed, outputs=_inventory(changed_frame)) + second = run_population( + tmp_path / "second", changed_sources, graph=replace(graph, nodes=(changed,)) + ) + assert old_ids != changed_frame.person["person_tax_unit_id"].tolist() + assert ( + first.manifest.nodes[node.id].frame_key + != second.manifest.nodes[changed.id].frame_key + ) + assert frame.weighted_entities == ("household",) + assert "tax_unit" in frame.schema.group_entities + + +def test_create_parser_rejects_an_unregistered_version(): + with pytest.raises(ValueError, match="Unknown transport CREATE"): + create_unit_declaration("transport.create@99", {}) + + +def test_multiple_unit_attributes_produce_explicit_person_role_columns(tmp_path): + def edit(tables): + tables["person"]["is_disabled"] = True + + sources = write_population_sources(tmp_path / "sources", edit_tables=edit) + declaration = MultiUnitRule( + units=( + UnitDeclaration(kind="benefit_unit", rule=UNIT_RULE), + UnitDeclaration( + kind="tax_unit", + rule=TaxUnitRule( + entity="tax_unit", + pairing_eligibility_column="has_disability", + provenance={ + "eligibility": "Synthetic explicit boolean, no legal claim." + }, + ), + ), + ) + ) + frame = _run_create( + sources, + _versioned(create_node(sources), declaration), + TRANSPORT_CREATE_DECLARED, + ).frame + document = declaration.to_dict() + node = Node( + "attributes", + TRANSPORT_UNIT_ATTRIBUTES.ref, + population="open", + params={ + "unit_rule": canonical_text(document), + "unit_rule_sha256": digest(document), + }, + outputs=( + Owned("person", "family_role", "string"), + Owned("person", "tax_unit_role", "string"), + Owned("family", "n_members", "int64"), + Owned("tax_unit", "n_members", "int64"), + Owned("household", "n_family_units", "int64"), + Owned("household", "n_tax_unit_units", "int64"), + ), + ) + result = TRANSPORT_UNIT_ATTRIBUTES.run( + KernelContext( + node=node, + tables={entity: frame.table(entity) for entity in frame.entities}, + weights={}, + strata=frame.strata, + params=node.params, + rng=np.random.default_rng(0), + ) + ) + assert set(result.columns) == {(item.entity, item.column) for item in node.outputs} + for item in declaration.units: + expected = benefit_unit_roles( + frame.person, + frame.table(item.entity), + frame.person[item.rule.membership_column], + item.rule, + ) + assert ( + result.columns["person", f"{item.entity}_role"].tolist() + == expected.tolist() + ) + assert result.columns[item.entity, "n_members"].sum() == frame.n("person") + + +@settings( + max_examples=5, + deadline=None, + database=None, + suppress_health_check=[HealthCheck.function_scoped_fixture, HealthCheck.too_slow], +) +@given( + omitted=st.sampled_from( + [ + ("family", "family_id"), + ("person", "person_id"), + ("person", "person_family_id"), + ("person", "age"), + ("person", "family_role"), + ] + ) +) +def test_relation_inputs_must_be_explicit_graph_dependencies(tmp_path, omitted): + fixture = make_composed_fixture(tmp_path / "fixture") + spec = copy.deepcopy(fixture.spec) + binding = next( + row + for row in spec["resources"]["axiom_rules_bindings"]["bindings"] + if row["id"] == "family_housing" + ) + binding["relations"] = [ + { + "predicate": "toy_members", + "key": "toy_members:0:1", + "current_slot": 0, + "related_slot": 1, + "arguments": ["Family", "Person"], + "root_entity": "family", + "related_entity": "person", + "membership_column": "person_family_id", + "related_inputs": {"toy_age": "age"}, + "cardinality": "many", + "period": binding["period"], + "role_column": "family_role", + "roles": ["head", "partner"], + } + ] + dependencies = { + ("family", "family_id"), + ("person", "person_id"), + ("person", "person_family_id"), + ("person", "age"), + ("person", "family_role"), + } + node = Node( + "relations", + "simulate.rules_by_ref@1", + population="open", + inputs=tuple( + Slice(entity, (column,)) for entity, column in sorted(dependencies) + ), + outputs=tuple( + Owned("family", name, "float64") for name in binding["variables"] + ), + ) + assert ( + transport_rules_node(spec, fixture.config, binding["id"], node).params["period"] + == binding["period"] + ) + restricted = replace( + node, + inputs=( + replace( + node.inputs[0], + columns=(*node.inputs[0].columns, "selected"), + rows="selected", + ), + *node.inputs[1:], + ), + ) + with pytest.raises(ValueError, match="whole-table relation inputs"): + transport_rules_node(spec, fixture.config, binding["id"], restricted) + missing = replace( + node, + inputs=tuple( + Slice(entity, (column,)) + for entity, column in sorted(dependencies - {omitted}) + ), + ) + with pytest.raises(ValueError, match="whole-table relation inputs"): + transport_rules_node(spec, fixture.config, binding["id"], missing) + + +@pytest.mark.parametrize("document", [None, {}, "{}"]) +def test_versioned_create_requires_a_versioned_canonical_document(tmp_path, document): + fixture = make_composed_fixture(tmp_path / "fixture") + spec = copy.deepcopy(fixture.spec) + create = next( + row + for row in spec["resources"]["transport_graph"]["nodes"] + if row["kernel"] == TRANSPORT_CREATE.ref + ) + create["kernel"] = TRANSPORT_CREATE_DECLARED.ref + create["params"]["unit_rule"] = document + with pytest.raises((ValueError, TypeError), match="unit|param|selection|format"): + skeleton_create_declaration(spec) + + +def test_versioned_create_parser_requires_one_create_and_json_text(tmp_path): + from microcosm.build.transport.compose import _versioned_create_units + + fixture = make_composed_fixture(tmp_path / "fixture") + spec = copy.deepcopy(fixture.spec) + rows = spec["resources"]["transport_graph"]["nodes"] + create = next(row for row in rows if row["kernel"] == TRANSPORT_CREATE.ref) + create["kernel"] = TRANSPORT_CREATE_DECLARED.ref + create["params"]["unit_rule"] = None + with pytest.raises(ValueError, match="canonical JSON unit_rule"): + _versioned_create_units(spec) + rows.append(copy.deepcopy(create)) + with pytest.raises(ValueError, match="exactly one CREATE"): + _versioned_create_units(spec) + + +def test_legacy_unit_entities_and_ancestry_accept_an_absent_graph(): + from microcosm.build.transport.compose import ( + _unit_entities, + _versioned_create_units, + load_transport_spec, + validate_transport_calibration_ancestry, + ) + + spec = load_transport_spec("nz") + resources = spec["resources"] + resources.pop("transport_graph", None) + resources.pop("transport_graph.json", None) + assert resources["precalibration_stages"]["stages"] + assert _versioned_create_units(spec) is None + assert _unit_entities(spec) == (resources["benefit_unit_rule"]["entity"],) + validate_transport_calibration_ancestry(spec) + + +@pytest.mark.parametrize( + "document", + [ + {}, + {"schema_version": 2, "country": "toy", "nodes": []}, + {"schema_version": 1, "country": "other", "nodes": []}, + ], +) +def test_legacy_unit_entities_validate_a_present_graph(document): + from microcosm.build.transport.compose import _unit_entities + + spec = { + "country": "toy", + "resources": { + "benefit_unit_rule": {"entity": "family"}, + "transport_graph": document, + }, + } + with pytest.raises(ValueError, match="schema_version 1 and the spec country"): + _unit_entities(spec) + + +def test_skeleton_refuses_unregistered_create_version(tmp_path): + fixture = make_composed_fixture(tmp_path / "fixture") + spec = copy.deepcopy(fixture.spec) + create = next( + row + for row in spec["resources"]["transport_graph"]["nodes"] + if row["kernel"] == TRANSPORT_CREATE.ref + ) + create["kernel"] = "transport.create@99" + with pytest.raises(ValueError, match="exactly one transport.create"): + skeleton_create_declaration(spec) + + +@pytest.mark.parametrize("include_graph", [False, True], ids=["absent", "partial"]) +def test_activation_aggregates_gaps_without_a_complete_graph(include_graph): + resources = { + "target_references": {"target_references": []}, + "precal_references": {"target_references": []}, + "holdout_references": { + "target_references": [ + { + "name": "parked", + "metadata": {"activation_status": "requires_review"}, + } + ] + }, + } + if include_graph: + resources["transport_graph"] = { + "nodes": [ + { + "params": { + "reference_document": { + "resource": "holdout_references", + "encoding": "json", + } + } + } + ] + } + with pytest.raises(ValueError, match="Transport spec is not activated:") as error: + validate_transport_activation({"country": "toy", "resources": resources}) + message = str(error.value) + assert "placeholder reference holdout_references.parked" in message + for resource in ("axiom_rules_bindings", "scenarios", "benefit_unit_rule"): + assert f"{resource}.json" in message + if not include_graph: + assert "transport_graph.json (the skeleton declaration)" in message diff --git a/packages/microcosm-build/tests/engine_free/shared/test_transport_population_kernels.py b/packages/microcosm-build/tests/engine_free/shared/test_transport_population_kernels.py index dda02072d1..a91b6a9f02 100644 --- a/packages/microcosm-build/tests/engine_free/shared/test_transport_population_kernels.py +++ b/packages/microcosm-build/tests/engine_free/shared/test_transport_population_kernels.py @@ -25,6 +25,7 @@ from microcosm.build.transport.population_kernels import ( TRANSPORT_BOUNDARY, TRANSPORT_CREATE, + TRANSPORT_CREATE_DECLARED, TRANSPORT_CURRENCY, TRANSPORT_QUANTILE_MAP, register_population_kernels, @@ -121,6 +122,7 @@ def test_kernel_registration_is_explicit_idempotent_and_country_neutral(): assert register_population_kernels(registry) is registry kernels = ( TRANSPORT_CREATE, + TRANSPORT_CREATE_DECLARED, TRANSPORT_BOUNDARY, TRANSPORT_CURRENCY, TRANSPORT_QUANTILE_MAP, @@ -132,11 +134,13 @@ def test_kernel_registration_is_explicit_idempotent_and_country_neutral(): assert kernel.capabilities.consumes_se is False assert kernel.implementation_hash() == kernel.implementation_hash() assert kernel.capabilities.seed_source is ( - SeedSource.KEYED if kernel is TRANSPORT_CREATE else SeedSource.NONE + SeedSource.KEYED + if kernel in (TRANSPORT_CREATE, TRANSPORT_CREATE_DECLARED) + else SeedSource.NONE ) -def test_aggregate_population_registration_installs_exactly_thirteen_kernels(): +def test_aggregate_population_registration_installs_exactly_fourteen_kernels(): registry = KernelRegistry() assert register_transport_population_kernels(registry) is registry assert register_transport_population_kernels(registry) is registry @@ -145,6 +149,7 @@ def test_aggregate_population_registration_installs_exactly_thirteen_kernels(): "money.minor_units@1", "money.annual_transfers@1", "transport.create@1", + "transport.create@2", "transport.boundary@1", "transport.currency@1", "transport.quantile_map@1", diff --git a/packages/microcosm-build/tests/engine_free/shared/test_transport_registry.py b/packages/microcosm-build/tests/engine_free/shared/test_transport_registry.py index 7547bce961..eb02d36193 100644 --- a/packages/microcosm-build/tests/engine_free/shared/test_transport_registry.py +++ b/packages/microcosm-build/tests/engine_free/shared/test_transport_registry.py @@ -21,7 +21,11 @@ RULESPEC_SOURCE_CODEC, ) from microcosm.build.transport.gate_bindings import TRANSPORT_GATE_REGISTRY -from microcosm.build.transport.registry import build_transport_registry +from microcosm.build.transport.registry import ( + binding_relation_dependencies, + build_transport_registry, + transport_entity_schema, +) from microcosm.build.transport.terminal_kernels import EXPORT_SOURCE_CODEC from microcosm.frame import ( EntitySchema, @@ -34,6 +38,11 @@ from microcosm.frame.adapters.axiom import NZ_SCHEMA, AxiomEngine, AxiomPeriod from microcosm.frame.kernels import SimulateRulesKernel from microcosm.frame.rules_kernels import SimulateRulesByRefKernel +from microcosm.frame.unit_construction import ( + MultiUnitRule, + TaxUnitRule, + UnitDeclaration, +) from microcosm.graph import ( Capabilities, ContentStore, @@ -336,6 +345,291 @@ def metadata(engine, name): assert ref["periods"] == {row["period"]: document["periods"][row["period"]]} +def _relation_manifest(): + return { + "predicate": "family_member", + "key": "family_member:0:1", + "current_slot": 0, + "related_slot": 1, + "arguments": ["Family", "Person"], + "root_entity": "family", + "related_entity": "person", + "membership_column": "person_family_id", + "related_inputs": {"member_income": "income"}, + "cardinality": "many", + "period": "2026-27", + "role_column": "person_family_role", + "roles": ["head", "partner"], + } + + +def _native_relation_fixture(tmp_path, monkeypatch): + root, document, _ = _fixture(tmp_path) + document["bindings"] = [document["bindings"][-1]] + row = document["bindings"][0] + row["relations"] = [_relation_manifest()] + checks = [] + + def check(engine, entity): + assert engine._pinned_tree_sha256 is not None + checks.append((entity, tuple(item.to_dict() for item in engine._relations))) + + monkeypatch.setattr(AxiomEngine, "assert_declared_relations", check) + monkeypatch.setattr(registry_module, "assert_no_relations", lambda *args: None) + monkeypatch.setattr( + AxiomEngine, + "variable_metadata", + lambda engine, name: VariableMetadata(name, "family", "float", "year"), + ) + monkeypatch.setattr(AxiomEngine, "graph_dtype", lambda engine, name: "float64") + return root, document, checks + + +def test_explicit_relation_is_pinned_and_checked_before_preparation_returns( + tmp_path, monkeypatch +): + root, document, checks = _native_relation_fixture(tmp_path, monkeypatch) + prepared = build_transport_registry(document, root, unit_rule=UNIT_RULE) + reference = json.loads(prepared.engine_refs["housing"]) + assert reference["relations"] == document["bindings"][0]["relations"] + assert checks == [("family", tuple(reference["relations"]))] + engine = prepared.engines[prepared.engine_refs["housing"]] + # Preparation keeps declarations, never data-derived offsets or arrays. + assert len(engine._relations) == 1 + assert engine._relations[0].related_inputs == {"member_income": "income"} + + +@settings( + max_examples=12, + deadline=None, + suppress_health_check=[HealthCheck.function_scoped_fixture], +) +@given(st.integers(min_value=1, max_value=30)) +def test_relation_input_and_selector_semantics_change_only_own_reference( + tmp_path, monkeypatch, suffix +): + root, document, _ = _native_relation_fixture(tmp_path, monkeypatch) + document["bindings"].append(_fixture(tmp_path)[1]["bindings"][0]) + monkeypatch.setattr( + AxiomEngine, + "variable_metadata", + lambda engine, name: VariableMetadata( + name, + next( + row["entity"] + for row in document["bindings"] + if name in row["variables"] + ), + "float", + "year", + ), + ) + first = build_transport_registry(document, root, unit_rule=UNIT_RULE) + row = document["bindings"][0] + row["relations"][0]["related_inputs"]["member_income"] = f"income_{suffix}" + second = build_transport_registry(document, root, unit_rule=UNIT_RULE) + assert first.engine_refs["housing"] != second.engine_refs["housing"] + row["relations"][0]["roles"] = [f"role_{suffix}"] + third = build_transport_registry(document, root, unit_rule=UNIT_RULE) + assert second.engine_refs["housing"] != third.engine_refs["housing"] + assert ( + first.engine_refs["main_benefit"] + == second.engine_refs["main_benefit"] + == third.engine_refs["main_benefit"] + ) + assert first.kernels.implementation_hash( + "simulate.rules_by_ref@1" + ) == third.kernels.implementation_hash("simulate.rules_by_ref@1") + + +def test_absent_relation_declaration_stays_absent(tmp_path, monkeypatch): + root, document, _ = _native_relation_fixture(tmp_path, monkeypatch) + document["bindings"][0].pop("relations") + prepared = build_transport_registry(document, root, unit_rule=UNIT_RULE) + reference = json.loads(prepared.engine_refs["housing"]) + assert "relations" not in reference + assert "views" not in reference + engine = prepared.engines[prepared.engine_refs["housing"]] + assert engine._relations is None + + +def test_relation_dependencies_share_physical_native_inputs(tmp_path, monkeypatch): + root, document, _ = _native_relation_fixture(tmp_path, monkeypatch) + binding = document["bindings"][0] + prepared = build_transport_registry(document, root, unit_rule=UNIT_RULE) + engine = prepared.engines[prepared.engine_refs["housing"]] + expected = ( + ("family", "family_id"), + ("person", "income"), + ("person", "person_family_id"), + ("person", "person_family_role"), + ("person", "person_id"), + ) + assert ( + binding_relation_dependencies( + binding, entity_names=document["entity_names"], schema=prepared.schema + ) + == expected + ) + assert engine.relation_dependencies("family") == expected + binding.pop("relations") + assert ( + binding_relation_dependencies( + binding, entity_names=document["entity_names"], schema=prepared.schema + ) + == () + ) + + +def test_relation_dependencies_refuse_a_binding_period_mismatch(tmp_path, monkeypatch): + _, document, _ = _native_relation_fixture(tmp_path, monkeypatch) + binding = document["bindings"][0] + binding["relations"][0]["period"] = "2027-28" + with pytest.raises(ValueError, match="period must equal the binding period"): + binding_relation_dependencies( + binding, entity_names=document["entity_names"], schema=NZ_SCHEMA + ) + + +def test_declared_units_derive_the_same_registry_schema_as_create(): + declaration = MultiUnitRule( + units=( + UnitDeclaration(kind="benefit_unit", rule=UNIT_RULE), + UnitDeclaration( + kind="tax_unit", + rule=TaxUnitRule( + entity="tax_unit", + pairing_eligibility_column="joint_assessment_eligible", + provenance={"method": "pinned classification input"}, + ), + ), + ) + ) + schema, nesting = transport_entity_schema(declaration) + assert (schema, nesting) == declaration.entity_schema() + assert schema.group_entities == ("household", "family", "tax_unit") + assert nesting == {"family": "household", "tax_unit": "household"} + assert transport_entity_schema(UNIT_RULE) == transport_entity_schema( + MultiUnitRule.from_benefit_unit(UNIT_RULE) + ) + with pytest.raises(TypeError, match="BenefitUnitRule or MultiUnitRule"): + transport_entity_schema(object()) + + +def test_declared_tax_unit_relation_uses_the_shared_partition_schema( + tmp_path, monkeypatch +): + root, document, _ = _native_relation_fixture(tmp_path, monkeypatch) + declaration = MultiUnitRule( + units=( + UnitDeclaration(kind="benefit_unit", rule=UNIT_RULE), + UnitDeclaration( + kind="tax_unit", + rule=TaxUnitRule( + entity="tax_unit", + pairing_eligibility_column="joint_assessment_eligible", + provenance={"method": "pinned classification input"}, + ), + ), + ) + ) + document["entity_names"]["tax_unit"] = "TaxUnit" + binding = document["bindings"][0] + binding["entity"] = "tax_unit" + binding["engine_entity"] = "TaxUnit" + relation = binding["relations"][0] + relation["root_entity"] = "tax_unit" + relation["arguments"] = ["TaxUnit", "Person"] + relation["membership_column"] = "person_tax_unit_id" + relation["role_column"] = "person_tax_unit_role" + monkeypatch.setattr( + AxiomEngine, + "variable_metadata", + lambda engine, name: VariableMetadata(name, "tax_unit", "float", "year"), + ) + prepared = build_transport_registry(document, root, unit_rule=declaration) + assert (prepared.schema, prepared.nesting) == declaration.entity_schema() + engine = prepared.engines[prepared.engine_refs["housing"]] + assert engine._nesting == {"family": "household", "tax_unit": "household"} + assert ("person", "person_tax_unit_id") in engine.relation_dependencies("tax_unit") + + +@pytest.mark.parametrize( + ("field", "value", "match"), + [ + ("root_entity", "household", "root_entity"), + ("related_entity", "household", "related persons"), + ("membership_column", "person_household_id", "partition membership"), + ("period", "2027-28", "period"), + ("arguments", ["Person", "Family"], "arguments"), + ("cardinality", "maybe", "cardinality"), + ("current_slot", True, "slot"), + ("current_slot", 1, "slot"), + ("key", "family_member:1:0", "key"), + ("arguments", "Family,Person", "list"), + ("related_inputs", [], "mapping"), + ("roles", "head", "list"), + ("roles", ["head", "head"], "roles"), + ("role_column", None, "role"), + ("undocumented", "value", "exactly"), + ], +) +def test_relation_manifest_validation(tmp_path, monkeypatch, field, value, match): + root, document, _ = _native_relation_fixture(tmp_path, monkeypatch) + document["bindings"][0]["relations"][0][field] = value + with pytest.raises((TypeError, ValueError), match=match): + build_transport_registry(document, root, unit_rule=UNIT_RULE) + + +@pytest.mark.parametrize( + ("problem", "error", "match"), + [ + ("scalar", TypeError, "list of relation manifests"), + ("duplicate", ValueError, "repeats a compiled relation key"), + ("missing", ValueError, "manifests require exactly"), + ("views", NotImplementedError, "views are not supported"), + ], +) +def test_relation_declaration_refuses_ambiguous_or_future_shapes( + tmp_path, monkeypatch, problem, error, match +): + root, document, _ = _native_relation_fixture(tmp_path, monkeypatch) + row = document["bindings"][0] + if problem == "scalar": + row["relations"] = None + elif problem == "duplicate": + row["relations"].append(copy.deepcopy(row["relations"][0])) + elif problem == "missing": + row["relations"][0].pop("related_inputs") + else: + row["views"] = [] + with pytest.raises(error, match=match): + build_transport_registry(document, root, unit_rule=UNIT_RULE) + + +def test_injected_axiom_relation_configuration_must_match(tmp_path, monkeypatch): + root, document, _ = _native_relation_fixture(tmp_path, monkeypatch) + prepared = build_transport_registry(document, root, unit_rule=UNIT_RULE) + engine = prepared.engines[prepared.engine_refs["housing"]] + engine._relations = None + with pytest.raises(ValueError, match="relations do not match"): + build_transport_registry( + document, + root, + unit_rule=UNIT_RULE, + engines_by_binding={"housing": engine}, + ) + + +def test_python_adapter_does_not_claim_native_relation_execution(tmp_path): + root, document, engines = _fixture(tmp_path) + document["bindings"][-1]["relations"] = [_relation_manifest()] + with pytest.raises(TypeError, match="relation manifests require an AxiomEngine"): + build_transport_registry( + document, root, unit_rule=UNIT_RULE, engines_by_binding=engines + ) + + @pytest.mark.parametrize( "problem", ["root", "module", "schema", "periods", "names", "nesting", "dtypes"] ) diff --git a/packages/microcosm-build/tests/engine_free/shared/test_transport_target_kernels.py b/packages/microcosm-build/tests/engine_free/shared/test_transport_target_kernels.py index 10c2d3c7df..7c74b41200 100644 --- a/packages/microcosm-build/tests/engine_free/shared/test_transport_target_kernels.py +++ b/packages/microcosm-build/tests/engine_free/shared/test_transport_target_kernels.py @@ -501,6 +501,20 @@ def test_compile_emits_a_traced_surface_bound_to_facts_not_paths(tmp_path) -> No assert again.artifacts["surface"] == payload +def test_calibration_preparation_attests_the_framework_roster_source( + tmp_path, monkeypatch +): + import microcosm.frame.unit_construction as units + + before = CALIBRATION_PREPARE.implementation_hash() + changed = tmp_path / "unit_construction.py" + changed.write_bytes( + Path(units.__file__).read_bytes() + b"\n# changed roster implementation\n" + ) + monkeypatch.setattr(units, "__file__", str(changed)) + assert CALIBRATION_PREPARE.implementation_hash() != before + + def test_compile_values_are_the_fact_values(tmp_path) -> None: surface = decode_target_surface(_surface(tmp_path)) facts = {fact["label"]: fact["value"] for fact in toy_facts(CALIBRATION_TARGETS)} diff --git a/packages/microcosm-build/tests/fixtures/calibration_columns_09a2fce.py.txt b/packages/microcosm-build/tests/fixtures/calibration_columns_09a2fce.py.txt new file mode 100644 index 0000000000..a194567ac1 --- /dev/null +++ b/packages/microcosm-build/tests/fixtures/calibration_columns_09a2fce.py.txt @@ -0,0 +1,223 @@ +# Frozen current-kernel helper from 09a2fce53968 before S1+S2a. +"""Prepare simple calibration columns from an explicit, engine-free contract. + +Geography and canonical concepts must already be present. This helper creates +constants, validated roster counts, conjunctions of declared predicates and +interval indicators classified by :mod:`microcosm.build.reference_intervals`; +receipt assignment and statistical-unit construction belong to their own producers. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +import numpy as np +import pandas as pd + +from microcosm.build.calibration_predicates import ( + _finite_number, + _identity, + predicate_mask, + validate_calibration_inputs, +) +from microcosm.build.reference_intervals import ( + interval_indicators, + validate_interval_document, +) + + +def prepare_calibration_columns( + tables: Mapping[str, pd.DataFrame], + bindings: Mapping[str, Any], + *, + intervals: Mapping[str, Any] | None = None, +) -> dict[str, pd.DataFrame]: + """Return independent tables with the contract's prepared columns. + + All required inputs are validated before any column is produced. Existing + prepared columns are recomputed, so a changed geography or age cannot retain + stale indicators. Missing/unknown inputs never become all-false categories. + An interval output names its entity, an effective table of ``intervals`` + and one of its bands; every record of each named table is validated, and + any refusal raises, before an output exists. + """ + required, columns = validate_calibration_inputs(tables, bindings) + interval_columns = _interval_columns(tables, columns, required, intervals) + out = {entity: table.copy(deep=True) for entity, table in tables.items()} + derived = _prepare_roster_counts(out, bindings.get("roster_counts", ()), required) + prepared: dict[str, dict[str, object]] = {} + seen: set[tuple[str, str]] = set() + for entry in columns: + if not isinstance(entry, Mapping): + raise ValueError("Calibration column entries must be objects.") + entity, name = entry.get("entity"), entry.get("name") + if entity not in required or not isinstance(name, str) or not name: + raise ValueError("Calibration column needs a declared entity and name.") + key = (entity, name) + if ( + key in seen + or name in required[entity] + or name in derived.get(entity, ()) + or name.endswith("_id") + ): + raise ValueError( + f"Calibration output {entity}.{name} is repeated or structural." + ) + seen.add(key) + if "interval" in entry: + prepared.setdefault(entity, {})[name] = interval_columns[key] + continue + if "constant" in entry: + value = entry["constant"] + if "predicates" in entry or not _finite_number(value): + raise ValueError(f"Calibration constant {entity}.{name} is invalid.") + prepared.setdefault(entity, {})[name] = value + continue + prepared.setdefault(entity, {})[name] = predicate_mask( + out[entity], + entry.get("predicates"), + entity=entity, + name=name, + validated_inputs={*required[entity], *derived.get(entity, ())}, + ) + for entity, values in prepared.items(): + original_order = list(out[entity].columns) + computed = pd.DataFrame(values, index=out[entity].index) + out[entity] = pd.concat( + [out[entity].drop(columns=list(values), errors="ignore"), computed], axis=1 + ) + order = original_order + [name for name in values if name not in original_order] + out[entity] = out[entity].loc[:, order] + return out + + +def _interval_columns( + tables: Mapping[str, pd.DataFrame], + columns: object, + required: Mapping[str, Any], + intervals: Mapping[str, Any] | None, +) -> dict[tuple[str, str], np.ndarray]: + """Classify every named table once; refuse before any output is produced.""" + entries = [ + entry for entry in columns if isinstance(entry, Mapping) and "interval" in entry + ] + if not entries: + return {} + if intervals is None: + raise ValueError("Interval outputs need the effective interval document.") + document = validate_interval_document(intervals) + requested: dict[str, list[tuple[str, str, str]]] = {} + for entry in entries: + entity, name, binding = ( + entry.get("entity"), + entry.get("name"), + entry["interval"], + ) + if set(entry) != {"entity", "name", "interval"}: + raise ValueError( + f"Interval output {entity}.{name} declares only its entity, name " + "and interval." + ) + if not isinstance(binding, Mapping) or set(binding) != {"table", "band"}: + raise ValueError("An interval output names exactly its table and band.") + table = document.get(binding["table"]) + if table is None or binding["band"] not in table.band_ids: + raise ValueError( + f"Interval output {entity}.{name} names an undeclared table or band." + ) + if entity not in required or entity != table.entity: + raise ValueError( + f"Interval output {entity}.{name} must name its table's entity." + ) + if name != table.indicators[table.band_ids.index(binding["band"])]: + raise ValueError( + f"Interval output {entity}.{name} must be its band's declared indicator." + ) + if ( + table.declaration["source_period"] != table.declaration["target_period"] + and not table.scaled + ): + raise ValueError( + f"Interval table {table.id!r} has source-period cuts; scale it to its " + "target period before preparing indicators." + ) + requested.setdefault(table.id, []).append((entity, name, binding["band"])) + result: dict[tuple[str, str], np.ndarray] = {} + computed = {} + for table_id in requested: + table = document[table_id] + data = tables.get(table.entity) + if data is None or any( + column not in data + for column in (table.income_column, table.universe_column) + ): + raise ValueError( + f"Interval table {table_id!r} needs {table.entity}." + f"{table.income_column} and {table.entity}.{table.universe_column}." + ) + computed[table_id] = interval_indicators( + table, + data[table.income_column].tolist(), + data[table.universe_column].tolist(), + ) + for table_id, outputs in requested.items(): + for entity, name, _ in outputs: + result[entity, name] = computed[table_id][name] + return result + + +def _prepare_roster_counts( + tables: dict[str, pd.DataFrame], entries: object, required: Mapping[str, Any] +) -> dict[str, set[str]]: + if not isinstance(entries, (list, tuple)): + raise ValueError("Calibration roster_counts must be a list.") + derived: dict[str, set[str]] = {} + for entry in entries: + if not isinstance(entry, Mapping): + raise ValueError("Calibration roster count must be an object.") + entity, members, name = ( + entry.get("entity"), + entry.get("member_entity"), + entry.get("name"), + ) + group_id, member_id, membership = ( + entry.get("entity_id"), + entry.get("member_id"), + entry.get("membership"), + ) + if ( + entity not in required + or members not in required + or not isinstance(name, str) + or not name + or name.endswith("_id") + or name in required[entity] + or name in derived.get(entity, ()) + or group_id not in required[entity] + or member_id not in required[members] + or membership not in required[members] + or not isinstance(entry.get("representation"), str) + or not entry["representation"] + ): + raise ValueError( + "Calibration roster count needs declared identities and representation." + ) + groups = [_identity(value) for value in tables[entity][group_id].tolist()] + identities = [_identity(value) for value in tables[members][member_id].tolist()] + if len(set(groups)) != len(groups) or len(set(identities)) != len(identities): + raise ValueError( + "Calibration roster contains duplicate membership identities." + ) + counts = dict.fromkeys(groups, 0) + for value in tables[members][membership].tolist(): + key = _identity(value) + if key not in counts: + raise ValueError("Calibration roster contains orphan membership.") + counts[key] += 1 + if any(count == 0 for count in counts.values()): + raise ValueError("Calibration roster contains empty households or groups.") + tables[entity][name] = [counts[key] for key in groups] + derived.setdefault(entity, set()).add(name) + return derived + diff --git a/packages/microcosm-frame/src/microcosm/frame/adapters/axiom.py b/packages/microcosm-frame/src/microcosm/frame/adapters/axiom.py index a5c3f5d9d0..aa5b83f032 100644 --- a/packages/microcosm-frame/src/microcosm/frame/adapters/axiom.py +++ b/packages/microcosm-frame/src/microcosm/frame/adapters/axiom.py @@ -122,6 +122,7 @@ "AxiomEngine", "AxiomEntityTableDataset", "AxiomPeriod", + "AxiomRelation", "BE_SCHEMA", "NZ_NESTING", "NZ_SCHEMA", @@ -196,6 +197,135 @@ _SHA256 = re.compile(r"^[0-9a-f]{64}$") +@dataclass(frozen=True) +class AxiomRelation: + """An explicit binary relation from a partition to its person members. + + ``arguments`` records the source declaration's engine entity tuple; slots + and ``key`` must agree with the compiled dense program. Related columns + are named explicitly so they also form graph input dependencies. No rows + or data-derived batches are retained in this declaration. + """ + + predicate: str + key: str + current_slot: int + related_slot: int + arguments: tuple[str, str] + root_entity: str + related_entity: str + membership_column: str + related_inputs: Mapping[str, str] + cardinality: str + period: str + role_column: str | None = None + roles: tuple[str, ...] | None = None + + def __post_init__(self) -> None: + for name in ( + "predicate", + "key", + "root_entity", + "related_entity", + "membership_column", + "period", + ): + value = getattr(self, name) + if not isinstance(value, str) or not value: + raise ValueError(f"Relation {name} must be a non-empty string.") + if ( + type(self.current_slot) is not int + or type(self.related_slot) is not int + or {self.current_slot, self.related_slot} != {0, 1} + ): + raise ValueError( + "Relation slots must be distinct binary coordinates 0 and 1." + ) + if isinstance(self.arguments, str | bytes) or not isinstance( + self.arguments, Sequence + ): + raise ValueError( + "Relation arguments must be a sequence of two engine entity names." + ) + arguments = tuple(self.arguments) + if len(arguments) != 2 or any( + not isinstance(x, str) or not x for x in arguments + ): + raise ValueError("Relation arguments must contain two engine entity names.") + object.__setattr__(self, "arguments", arguments) + expected = f"{self.predicate}:{self.current_slot}:{self.related_slot}" + if self.key != expected: + raise ValueError( + f"Relation key must agree with predicate and slots: {expected!r}." + ) + if self.cardinality not in ("one", "many"): + raise ValueError("Relation cardinality must be 'one' or 'many'.") + if not isinstance(self.related_inputs, Mapping) or any( + not isinstance(k, str) or not k or not isinstance(v, str) or not v + for k, v in self.related_inputs.items() + ): + raise ValueError( + "Relation related_inputs must map engine names to physical columns." + ) + object.__setattr__( + self, "related_inputs", MappingProxyType(dict(self.related_inputs)) + ) + if (self.role_column is None) != (self.roles is None): + raise ValueError( + "Relation role_column and roles must be supplied together." + ) + if self.roles is not None: + if isinstance(self.roles, str | bytes) or not isinstance( + self.roles, Sequence + ): + raise ValueError("Relation roles must be a sequence of role names.") + roles = tuple(self.roles) + if ( + not isinstance(self.role_column, str) + or not self.role_column + or not roles + or any(not isinstance(x, str) or not x for x in roles) + or len(set(roles)) != len(roles) + ): + raise ValueError( + "Relation roles must be unique non-empty strings and name a role_column." + ) + object.__setattr__(self, "roles", roles) + + def to_dict(self) -> dict[str, object]: + """Return the JSON declaration, preserving absent role selection.""" + result: dict[str, object] = { + "predicate": self.predicate, + "key": self.key, + "current_slot": self.current_slot, + "related_slot": self.related_slot, + "arguments": list(self.arguments), + "root_entity": self.root_entity, + "related_entity": self.related_entity, + "membership_column": self.membership_column, + "related_inputs": dict(sorted(self.related_inputs.items())), + "cardinality": self.cardinality, + "period": self.period, + } + if self.role_column is not None: + result["role_column"] = self.role_column + result["roles"] = list(self.roles) + return result + + def input_dependencies(self, schema: EntitySchema) -> tuple[tuple[str, str], ...]: + """The physical identity, membership, selector and input columns.""" + related = self.related_entity + dependencies = { + (self.root_entity, schema.entity_id_column(self.root_entity)), + (related, schema.entity_id_column(related)), + (related, self.membership_column), + *((related, column) for column in self.related_inputs.values()), + } + if self.role_column is not None: + dependencies.add((related, self.role_column)) + return tuple(sorted(dependencies)) + + @dataclass(frozen=True) class AxiomPeriod: """Explicit dense-execution bounds for one policy period. @@ -325,6 +455,7 @@ def __init__( periods: Mapping[int | str, AxiomPeriod] | None = None, nesting: Mapping[str, str] | None = None, output_dtypes: str = "native", + relations: Sequence[AxiomRelation] | None = None, ) -> None: if arithmetic not in ("decimal", "f64"): raise ValueError( @@ -372,6 +503,39 @@ def __init__( self._frame_entity_by_engine = { engine: frame for frame, engine in self._entity_names.items() } + self._relations = None if relations is None else tuple(relations) + if self._relations is not None: + keys: set[tuple[str, str]] = set() + for relation in self._relations: + if not isinstance(relation, AxiomRelation): + raise TypeError( + "relations must contain AxiomRelation declarations." + ) + key = (relation.root_entity, relation.key) + if key in keys: + raise ValueError(f"Duplicate relation declaration {key!r}.") + keys.add(key) + if ( + relation.root_entity not in schema.group_entities + or relation.related_entity != schema.person_entity + or relation.membership_column + != schema.membership_column(relation.root_entity) + ): + raise ValueError( + "Relations require a declared partition's person membership column." + ) + if relation.arguments[relation.current_slot] != self._entity_names.get( + relation.root_entity + ) or relation.arguments[ + relation.related_slot + ] != self._entity_names.get(relation.related_entity): + raise ValueError( + "Relation argument slots disagree with the entity mapping." + ) + if self._periods is not None and relation.period not in self._periods: + raise ValueError( + f"Relation period {relation.period!r} is not declared." + ) self._programs: dict[str, Any] = {} self._metadata: dict[str, Any] | None = None # Set by axiom_engine_ref: the RuleSpec tree digest this adapter's @@ -457,8 +621,8 @@ def graph_dtype(self, name: str) -> str: def variables(self) -> list[str]: """Return the input variables the engine accepts on a dataset. - The union of the dense root inputs across every mapped engine - entity, sorted. Computed (derived) outputs are not included. + The union of root and declared related inputs across every mapped + engine entity, sorted. Computed outputs are not included. Raises: ImportError: If ``axiom_rules_engine`` is not installed. @@ -468,8 +632,165 @@ def variables(self) -> list[str]: program = self._program(frame_entity, missing_ok=True) if program is not None: names.update(program.root_inputs) + if self._relations is not None: + self.assert_declared_relations(frame_entity) + for relation in self._relations: + if relation.root_entity == frame_entity: + names.update(relation.related_inputs) return sorted(names) + def relation_dependencies(self, entity: str) -> tuple[tuple[str, str], ...]: + """Physical columns required to construct this root's relation batches.""" + dependencies: set[tuple[str, str]] = set() + for relation in self._relations or (): + if relation.root_entity != entity: + continue + dependencies.update(relation.input_dependencies(self._schema)) + return tuple(sorted(dependencies)) + + def validate_input_override(self, entity: str, column: str) -> None: + """Refuse aliased related overrides until bridges can resolve them. + + Bridges write physical table columns after checking logical engine + inputs. An aliased relation input needs both names resolved together; + accepting only the engine name would silently leave its source intact. + """ + for relation in self._relations or (): + for name, physical in relation.related_inputs.items(): + if ( + entity == relation.related_entity + and name != physical + and column in (name, physical) + ): + raise NotImplementedError( + "Input overrides of aliased related Axiom inputs are not supported." + ) + + def assert_declared_relations(self, entity: str) -> None: + """Validate manifests against compiled slots, input union and source tuple.""" + program = self._program(entity) + manifests = { + item.key: item + for item in self._relations or () + if item.root_entity == entity + } + compiled: dict[str, set[str]] = {} + for item in program.relations: + manifest = manifests.get(item.key) + if manifest is None: + raise NotImplementedError( + f"Module {self._module_label()} declares dense relations " + f"{[relation.name for relation in program.relations]} on entity {entity!r}; " + "an explicit relation declaration is required." + ) + if ( + item.name != manifest.predicate + or item.current_slot != manifest.current_slot + or item.related_slot != manifest.related_slot + ): + raise ValueError( + f"Compiled relation slots disagree with manifest {manifest.key!r}." + ) + compiled.setdefault(item.key, set()).update(item.related_inputs) + if set(manifests) != set(compiled): + raise ValueError( + f"Relation declarations do not match compiled keys on {entity!r}." + ) + for key, inputs in compiled.items(): + manifest = manifests[key] + if set(manifest.related_inputs) != inputs: + raise ValueError( + f"Relation {key!r} related inputs disagree with compiled input union." + ) + arguments = _source_relation_arguments( + self._module, self._rulespec_roots, manifest.predicate + ) + if arguments != manifest.arguments: + raise ValueError( + f"Relation {key!r} source arguments {arguments!r} disagree with manifest." + ) + + def _relation_batches( + self, bundle: Frame, entity: str, period: object + ) -> dict[str, Any]: + """Gather identity-ordered related rows following the native dense contract.""" + from axiom_rules_engine.dense import DenseRelationBatch + + batches: dict[str, Any] = {} + root = bundle.table(entity) + root_ids = root[self._schema.entity_id_column(entity)].to_numpy() + for relation in self._relations or (): + if relation.root_entity != entity: + continue + if self._materialization_period( + relation.period + ) != self._materialization_period(period): + raise ValueError( + f"Relation {relation.key!r} does not cover the requested period." + ) + table = bundle.table(relation.related_entity) + required = { + relation.membership_column, + self._schema.entity_id_column(relation.related_entity), + *relation.related_inputs.values(), + } + if relation.role_column is not None: + required.add(relation.role_column) + missing = sorted(required - set(table.columns)) + if missing: + raise ValueError( + f"Relation {relation.key!r} is missing related columns {missing}." + ) + selected = table + if relation.role_column is not None: + selected = selected.loc[ + selected[relation.role_column].isin(relation.roles) + ] + if not selected[relation.membership_column].isin(root_ids).all(): + raise ValueError( + f"Relation {relation.key!r} contains dangling membership identities." + ) + identity = self._schema.entity_id_column(relation.related_entity) + if selected[identity].isna().any() or selected[identity].duplicated().any(): + raise ValueError( + f"Relation {relation.key!r} requires unique non-null related identities." + ) + # A physical person identity is the deduplication key, never its + # input values; equal-valued people must remain separate targets. + selected = selected.drop_duplicates(identity).sort_values( + identity, kind="stable" + ) + positions_by_root: dict[object, list[int]] = {} + for position, root_id in enumerate( + selected[relation.membership_column].to_numpy() + ): + positions_by_root.setdefault(root_id, []).append(position) + rows: list[int] = [] + offsets = [0] + for root_id in root_ids: + positions = positions_by_root.get(root_id, []) + if relation.cardinality == "one" and len(positions) != 1: + raise ValueError( + f"Relation {relation.key!r} requires exactly one related identity per root." + ) + rows.extend(positions) + offsets.append(len(rows)) + # Mapping two engine inputs to one physical column is valid. + gathered = _batch_from_table( + selected.iloc[rows], + tuple(dict.fromkeys(relation.related_inputs.values())), + ) + inputs = { + name: gathered[column] + for name, column in relation.related_inputs.items() + } + batches[relation.key] = DenseRelationBatch( + np.asarray(offsets, dtype=np.int64), + inputs, + related_row_count=len(rows), + ) + return batches + def entity_schema(self) -> EntitySchema: """Return the frame entity schema (no engine import required).""" return self._schema @@ -552,9 +873,8 @@ def materialize( input, a graph-typed output is text, date, or not losslessly castable, or a computed array's length does not match its entity table. - NotImplementedError: If the compiled module declares relations - (cross-entity aggregation batches; not yet wired — the BE - pilot slice declares none). + NotImplementedError: If a compiled relation has no explicit + validated binding declaration. """ self._require_schema(bundle) start, end, period_kind = self._materialization_period(period) @@ -581,24 +901,27 @@ def materialize( results: dict[str, np.ndarray] = {} for frame_entity, names in by_entity.items(): program = self._program(frame_entity) - if program.relations: - raise NotImplementedError( - f"Module {self._module.name!r} declares dense relations " - f"{[item.name for item in program.relations]}; relation " - "batches from frame membership are not wired yet." - ) + if program.relations or self._relations: + self.assert_declared_relations(frame_entity) table = bundle.table(frame_entity) inputs = _batch_from_table(table, program.root_inputs) execute = ( program.execute_f64 if self._arithmetic == "f64" else program.execute ) - outputs = execute( + arguments: dict[str, Any] = dict( period_kind=period_kind, start=start, end=end, inputs=inputs, outputs=list(names), - )["outputs"] + ) + # Omit the keyword entirely for legacy bindings, preserving the + # existing execution call as well as its engine reference. + if program.relations: + arguments["relations"] = self._relation_batches( + bundle, frame_entity, period + ) + outputs = execute(**arguments)["outputs"] expected = bundle.n(frame_entity) for name in names: values = np.asarray(outputs[name]) @@ -1341,9 +1664,79 @@ def axiom_engine_ref( "periods": periods, "nesting": dict(sorted(engine._nesting.items())), } + if engine._relations is not None: + document["relations"] = [item.to_dict() for item in engine._relations] return canonical_json(document).decode("utf-8") +def _source_relation_arguments( + module: Path, roots: Sequence[Path], predicate: str +) -> tuple[str, ...] | None: + """Read the authored data-relation tuple along the module's import graph. + + PyYAML is a dependency of Axiom itself and is imported only for opt-in + relation bindings. Compiled slot coordinates alone cannot establish what + entity the author's tuple places at each coordinate. + """ + import yaml + + seen: set[Path] = set() + pending = [module] + found: set[tuple[str, ...]] = set() + canonical_roots = tuple(root.resolve() for root in roots) + while pending: + path = pending.pop().resolve() + if path in seen: + continue + seen.add(path) + if not any(path.is_relative_to(root) for root in canonical_roots): + raise ValueError("Relation declaration import escapes the RuleSpec roots.") + document = yaml.safe_load(path.read_text()) or {} + for rule in document.get("rules", ()): + declaration = rule.get("data_relation") + if rule.get("kind") == "data_relation" and isinstance(declaration, Mapping): + name = declaration.get("predicate", rule.get("name")) + qualified = { + f"{relative.parts[0]}:{relative.with_suffix('').relative_to(relative.parts[0]).as_posix()}#relation.{rule['name']}" + for root in canonical_roots + if path.is_relative_to(root) + for relative in (path.relative_to(root),) + if len(relative.parts) > 1 + } + if name == predicate or predicate in qualified: + if declaration.get("arity") != 2: + raise ValueError( + f"Relation {predicate!r} source must be binary." + ) + found.add(tuple(declaration.get("arguments", ()))) + for reference in document.get("imports", ()): + if not isinstance(reference, str): + raise ValueError( + "Relation declaration imports must name RuleSpec modules." + ) + target = reference.split("#", 1)[0] + if ":" in target: + country, relative = target.split(":", 1) + candidates = [ + root / country / f"{relative}.yaml" for root in canonical_roots + ] + else: + candidates = [path.parent / target] + if not Path(target).suffix: + candidates = [ + candidate.with_suffix(".yaml") for candidate in candidates + ] + existing = [candidate for candidate in candidates if candidate.is_file()] + if len(existing) != 1: + raise ValueError( + f"Relation declaration import {reference!r} is missing or ambiguous." + ) + pending.extend(existing) + if len(found) > 1: + raise ValueError(f"Relation {predicate!r} has conflicting source declarations.") + return next(iter(found), None) + + def assert_no_relations(engine: AxiomEngine, entity: str) -> None: """Refuse a module whose dense program on ``entity`` declares relations. diff --git a/packages/microcosm-frame/src/microcosm/frame/rules_kernels.py b/packages/microcosm-frame/src/microcosm/frame/rules_kernels.py index 2acebab145..861d24c266 100644 --- a/packages/microcosm-frame/src/microcosm/frame/rules_kernels.py +++ b/packages/microcosm-frame/src/microcosm/frame/rules_kernels.py @@ -448,6 +448,7 @@ def _evaluate( engine = self._delegates[ref]._engine tables = {entity: frame.table(entity).copy() for entity in frame.entities} inputs = set(engine.variables()) + validate_override = getattr(engine, "validate_input_override", None) overridden = set() for override in component["input_overrides"]: _fields( @@ -458,6 +459,8 @@ def _evaluate( ) entity, column = override["entity"], override["column"] frame.schema.entity_id_column(entity) + if callable(validate_override): + validate_override(entity, column) if column not in inputs or (entity, column) in overridden: raise ValueError( "Overrides must name distinct engine input columns." @@ -502,6 +505,8 @@ def _evaluate( tables[entity][column] = values if trial is not None and solve_input["engine_ref"] == ref: entity, column = solve_input["entity"], solve_input["column"] + if callable(validate_override): + validate_override(entity, column) if column not in inputs: raise ValueError("solve_input must name an engine input column.") tables[entity][column] = _map_values( diff --git a/packages/microcosm-frame/src/microcosm/frame/unit_construction.py b/packages/microcosm-frame/src/microcosm/frame/unit_construction.py index e8303f9452..1e25a19310 100644 --- a/packages/microcosm-frame/src/microcosm/frame/unit_construction.py +++ b/packages/microcosm-frame/src/microcosm/frame/unit_construction.py @@ -1,4 +1,10 @@ -"""Benefit units built from concept relationship pointers. +"""Declared unit partitions built from concept relationship pointers. + +``MultiUnitRule`` dispatches each declaration to its registered membership +constructor. NZ benefit units and tax assessment units share partition +validation, household nesting, stable physical ids, lineage and inherited +household weights. Tax units consume an explicit eligibility column; they +never infer legal status or place children with their parents for tax relief. A transported concept frame (:mod:`microcosm.frame.concepts`) has persons and households only. A target country's rules assess many programs on a smaller @@ -69,9 +75,11 @@ import json import math import re -from collections.abc import Mapping +from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass from enum import StrEnum +from numbers import Integral +from types import MappingProxyType import numpy as np import pandas as pd @@ -90,7 +98,16 @@ ) __all__ = [ + "prepare_roster_counts", "BENEFIT_UNIT_RULE_FORMAT", + "MULTI_UNIT_RULE_FORMAT", + "TAX_UNIT_RULE_FORMAT", + "MultiUnitRule", + "TaxUnitRule", + "UnitDeclaration", + "build_declared_units", + "declared_unit_attributes", + "reciprocal_partner_rows", "BenefitUnitRule", "DependentChildRule", "FinancialIndependenceTest", @@ -106,6 +123,8 @@ #: The serialized rule format :meth:`BenefitUnitRule.to_dict` writes. BENEFIT_UNIT_RULE_FORMAT = "microcosm.benefit_unit_rule.v1" +MULTI_UNIT_RULE_FORMAT = "microcosm.multi_unit_rule.v1" +TAX_UNIT_RULE_FORMAT = "microcosm.tax_unit_rule.v1" _PERSON_ID = PERSON_ID_COLUMN _PERSON_HOUSEHOLD_ID = PERSON_HOUSEHOLD_ID_COLUMN @@ -214,8 +233,39 @@ def __post_init__(self) -> None: raise ValueError("A parent-independence age must lie on [0, max_age].") +class _UnitColumns: + """Shared physical columns for every declared partition.""" + + entity: str + + @property + def id_column(self) -> str: + """The unit table's id column.""" + return f"{self.entity}_id" + + @property + def household_column(self) -> str: + """The unit table's column naming the household the unit nests in.""" + return f"{self.entity}_household_id" + + @property + def source_household_column(self) -> str: + """Original donor lineage, distinct from every other entity column.""" + return f"{self.entity}_{SOURCE_HOUSEHOLD_ID_COLUMN}" + + @property + def head_column(self) -> str: + """The unit table's column naming the unit's head.""" + return f"{self.entity}_head_person_id" + + @property + def membership_column(self) -> str: + """The person table's column naming each person's unit.""" + return f"person_{self.entity}_id" + + @dataclass(frozen=True, kw_only=True) -class BenefitUnitRule: +class BenefitUnitRule(_UnitColumns): """How a country groups persons into benefit units. Attributes: @@ -276,31 +326,6 @@ def __hash__(self) -> int: ) ) - @property - def id_column(self) -> str: - """The unit table's id column.""" - return f"{self.entity}_id" - - @property - def household_column(self) -> str: - """The unit table's column naming the household the unit nests in.""" - return f"{self.entity}_household_id" - - @property - def source_household_column(self) -> str: - """Original donor lineage, distinct from every other entity column.""" - return f"{self.entity}_{SOURCE_HOUSEHOLD_ID_COLUMN}" - - @property - def head_column(self) -> str: - """The unit table's column naming the unit's head.""" - return f"{self.entity}_head_person_id" - - @property - def membership_column(self) -> str: - """The person table's column naming each person's unit.""" - return f"person_{self.entity}_id" - def to_dict(self) -> dict[str, object]: """A JSON-ready form that :meth:`from_dict` reads back unchanged.""" test = self.dependent_child.financial_independence @@ -468,26 +493,295 @@ def build_benefit_units( rows = _Pointers.read(person, household, rule) dependent = _dependents(rows, rule) head = _heads(rows, dependent, rule) - ids = rows.ids + return _build_partition(person, household, rule, rows, head) + + +@dataclass(frozen=True, kw_only=True) +class TaxUnitRule(_UnitColumns): + """Assessment membership from an explicitly supplied eligibility result. + + Eligibility is a non-null boolean person column. Two reciprocal partners + form a pair only when both are eligible. Every other person, children + included, gets an individual assessment unit; dependant relief is a + separate association. A is the smaller person id (the ``head`` role), + B is their ``partner``. No income or implicit legal-status test is read. + """ + + entity: str + pairing_eligibility_column: str + provenance: Mapping[str, str] + + def __post_init__(self) -> None: + if not isinstance(self.entity, str) or not _ENTITY_NAME.fullmatch(self.entity): + raise ValueError("A tax-unit entity must be snake_case.") + if self.entity in ("person", "household"): + raise ValueError(f"Units cannot be named {self.entity!r}.") + column = self.pairing_eligibility_column + if not isinstance(column, str) or not _ENTITY_NAME.fullmatch(column): + raise ValueError( + "A tax-unit pairing eligibility column must be snake_case." + ) + if not isinstance(self.provenance, Mapping) or not all( + isinstance(key, str) and isinstance(value, str) and key and value + for key, value in self.provenance.items() + ): + raise ValueError( + "A tax-unit rule's provenance maps text to non-empty text." + ) + object.__setattr__(self, "provenance", MappingProxyType(dict(self.provenance))) + + def to_dict(self) -> dict[str, object]: + return { + "format": TAX_UNIT_RULE_FORMAT, + "entity": self.entity, + "pairing_eligibility_column": self.pairing_eligibility_column, + "provenance": dict(self.provenance), + } + + @classmethod + def from_dict(cls, data: Mapping[str, object]) -> TaxUnitRule: + with _parsing("tax unit rule"): + fields = _record_fields( + data, + "tax unit rule", + required=( + "format", + "entity", + "pairing_eligibility_column", + "provenance", + ), + ) + if fields.pop("format") != TAX_UNIT_RULE_FORMAT: + raise ValueError( + f"A tax unit rule must declare format {TAX_UNIT_RULE_FORMAT!r}." + ) + return cls(**fields) + + +@dataclass(frozen=True, kw_only=True) +class UnitDeclaration: + """A registered membership kind and its cited construction rule.""" + + kind: str + rule: BenefitUnitRule | TaxUnitRule + + def __post_init__(self) -> None: + registration = ( + _UNIT_KINDS.get(self.kind) if isinstance(self.kind, str) else None + ) + if registration is None: + raise ValueError(f"Unknown unit kind {self.kind!r}.") + if not isinstance(self.rule, registration[0]): + raise ValueError( + f"Unit kind {self.kind!r} requires {registration[0].__name__}." + ) + + @property + def entity(self) -> str: + return self.rule.entity + + def to_dict(self) -> dict[str, object]: + return {"kind": self.kind, "rule": self.rule.to_dict()} + + @classmethod + def from_dict(cls, data: Mapping[str, object]) -> UnitDeclaration: + fields = _record_fields(data, "unit declaration", required=("kind", "rule")) + kind = fields["kind"] + registration = _UNIT_KINDS.get(kind) if isinstance(kind, str) else None + if registration is None: + raise ValueError(f"Unknown unit kind {kind!r}.") + return cls(kind=kind, rule=registration[0].from_dict(fields["rule"])) + + +@dataclass(frozen=True, kw_only=True) +class MultiUnitRule: + """Versioned ordered declarations shared by CREATE, CLI and registry.""" + + units: tuple[UnitDeclaration, ...] + + def __post_init__(self) -> None: + if isinstance(self.units, str | bytes) or not isinstance(self.units, Sequence): + raise ValueError("Declared units must be a non-empty list.") + units = tuple(self.units) + if not units or not all(isinstance(item, UnitDeclaration) for item in units): + raise ValueError( + "Declared units must be a non-empty list of UnitDeclaration values." + ) + entities = [item.entity for item in units] + if len(entities) != len(set(entities)): + raise ValueError("Declared units repeat an entity.") + object.__setattr__(self, "units", units) + + @classmethod + def from_benefit_unit(cls, rule: BenefitUnitRule) -> MultiUnitRule: + return cls(units=(UnitDeclaration(kind="benefit_unit", rule=rule),)) + + def to_dict(self) -> dict[str, object]: + return { + "format": MULTI_UNIT_RULE_FORMAT, + "units": [item.to_dict() for item in self.units], + } + + @classmethod + def from_dict(cls, data: Mapping[str, object]) -> MultiUnitRule: + with _parsing("multi unit rule"): + fields = _record_fields( + data, "multi unit rule", required=("format", "units") + ) + if fields["format"] != MULTI_UNIT_RULE_FORMAT: + raise ValueError( + f"Declared units must declare format {MULTI_UNIT_RULE_FORMAT!r}." + ) + units = fields["units"] + if not isinstance(units, list): + raise ValueError("Declared units must be a list.") + return cls(units=tuple(UnitDeclaration.from_dict(item) for item in units)) + + @classmethod + def from_spec(cls, data: Mapping[str, object]) -> MultiUnitRule: + """Read the versioned country unit-declaration resource.""" + return cls.from_dict(data) + + def entity_schema(self): + """The physical partition schema and shared household nesting.""" + from microcosm.frame.schema import EntitySchema + + entities = tuple(item.entity for item in self.units) + return ( + EntitySchema(group_entities=("household", *entities)), + MappingProxyType(dict.fromkeys(entities, "household")), + ) + + +def _benefit_unit_heads( + rows: _Pointers, person: pd.DataFrame, rule: BenefitUnitRule +) -> np.ndarray: + return _heads(rows, _dependents(rows, rule), rule) + + +def _tax_unit_heads( + rows: _Pointers, person: pd.DataFrame, rule: TaxUnitRule +) -> np.ndarray: + column = rule.pairing_eligibility_column + if column not in person: + raise ValueError(f"Building tax units needs eligibility column {column!r}.") + eligibility = person[column] + if not pd.api.types.is_bool_dtype(eligibility.dtype) or eligibility.isna().any(): + raise ValueError("Tax-unit pairing eligibility must hold non-null booleans.") + eligible = eligibility.to_numpy(dtype=bool) + own = np.arange(len(rows.ids)) + other = np.maximum(rows.partner, 0) + paired = (rows.partner >= 0) & eligible & eligible[other] + return np.where(paired & (rows.ids[other] < rows.ids), other, own) + + +# Registration owns parsing and membership dispatch. All physical construction, +# validation, roles, lineage and weight semantics are shared below. +_UNIT_KINDS: Mapping[str, tuple[type, Callable]] = MappingProxyType( + { + "benefit_unit": (BenefitUnitRule, _benefit_unit_heads), + "tax_unit": (TaxUnitRule, _tax_unit_heads), + } +) + + +def build_declared_units( + person: pd.DataFrame, + household: pd.DataFrame, + declaration: MultiUnitRule, +) -> dict[str, GroupMembership]: + """Construct each declared partition without adding unit weight columns.""" + if not isinstance(declaration, MultiUnitRule): + raise TypeError("build_declared_units requires a MultiUnitRule.") + result = {} + for item in declaration.units: + rows = _Pointers.read(person, household, item.rule) + head = _UNIT_KINDS[item.kind][1](rows, person, item.rule) + units, membership = _build_partition(person, household, item.rule, rows, head) + result[item.entity] = benefit_unit_membership( + person, units, membership, item.rule + ) + return result + + +def declared_unit_attributes( + person: pd.DataFrame, + units: pd.DataFrame, + membership: pd.Series, + declaration: UnitDeclaration, +) -> pd.DataFrame: + """Composition attributes from the shared head/partner/child role contract.""" + return benefit_unit_attributes(person, units, membership, declaration.rule) + + +def _build_partition( + person: pd.DataFrame, + household: pd.DataFrame, + rule: BenefitUnitRule | TaxUnitRule, + rows: _Pointers, + head: np.ndarray, +) -> tuple[pd.DataFrame, pd.Series]: + """Materialize and validate the one physical partition contract.""" + if ( + head.shape != (len(person),) + or head.dtype.kind not in "iu" + or ((head < 0).any() or (head >= len(person)).any()) + ): + raise ValueError( + "A declared partition must assign exactly one head to every person." + ) + if not np.array_equal(rows.household_ids[head], rows.household_ids): + raise ValueError("A declared unit must nest inside one household.") head_rows = np.unique(head) - order = np.argsort(ids[head_rows], kind="stable") - head_rows = head_rows[order] - family = pd.DataFrame( + if not np.array_equal(head[head_rows], head_rows): + raise ValueError("A declared unit head must belong to its own unit.") + head_rows = head_rows[np.argsort(rows.ids[head_rows], kind="stable")] + units = pd.DataFrame( { - rule.id_column: ids[head_rows], + rule.id_column: rows.ids[head_rows], rule.household_column: rows.household_ids[head_rows], - rule.head_column: ids[head_rows], + rule.head_column: rows.ids[head_rows], } ) if SOURCE_HOUSEHOLD_ID_COLUMN in household: lineage = household.set_index(_HOUSEHOLD_ID)[SOURCE_HOUSEHOLD_ID_COLUMN] - family[rule.source_household_column] = family[rule.household_column].map( - lineage - ) + units[rule.source_household_column] = units[rule.household_column].map(lineage) membership = pd.Series( - ids[head], index=person.index, name=rule.membership_column, dtype=np.int64 + rows.ids[head], index=person.index, name=rule.membership_column, dtype=np.int64 ) - return family, membership + return units, membership + + +def reciprocal_partner_rows( + person: pd.DataFrame, household: pd.DataFrame +) -> np.ndarray: + """Resolve reciprocal co-resident partner pointers, refusing invalid ones. + + Row positions use -1 for absence, leaving every integer person id valid. + Both NZ and tax-unit construction use this one pointer implementation. + """ + columns = [_PERSON_ID, _PERSON_HOUSEHOLD_ID, "partner_person_id"] + missing = [name for name in columns if name not in person] + if missing: + raise ValueError(f"Partner pointers need columns {missing}.") + subset = { + "person": person.loc[:, columns], + "household": household.loc[:, [_HOUSEHOLD_ID, "reference_person_id"]], + } + violations = validate_concept_tables(subset) + if violations: + detail = "; ".join( + f"{item.column} {item.code}: {item.message}" for item in violations + ) + raise ValueError(f"Partner pointers refuse an invalid concept frame: {detail}") + ids = _int64_ids(person[_PERSON_ID], f"person.{_PERSON_ID}") + pointers = person["partner_person_id"] + present = pointers.notna().to_numpy() + rows = np.full(len(person), -1, dtype=np.int64) + rows[present] = pd.Index(ids).get_indexer( + _int64_ids(pointers[present], "person.partner_person_id") + ) + return rows def benefit_unit_roles( @@ -580,10 +874,11 @@ def benefit_unit_attributes( n_members = np.bincount(unit_rows, minlength=n_units).astype(np.int64) n_adults = np.bincount(unit_rows[is_adult], minlength=n_units).astype(np.int64) n_children = n_members - n_adults - ages = person["age"].to_numpy(dtype=np.float64) youngest = np.full(n_units, np.inf) child_rows = np.flatnonzero(~is_adult) - np.minimum.at(youngest, unit_rows[child_rows], ages[child_rows]) + if len(child_rows): + ages = person["age"].to_numpy(dtype=np.float64) + np.minimum.at(youngest, unit_rows[child_rows], ages[child_rows]) youngest_age = pd.array( [None if math.isinf(value) else int(value) for value in youngest], dtype="Int64", @@ -694,10 +989,16 @@ def read( cls, person: pd.DataFrame, household: pd.DataFrame, - rule: BenefitUnitRule, + rule: BenefitUnitRule | TaxUnitRule, ) -> _Pointers: - test = rule.dependent_child.financial_independence - person_columns = [_PERSON_ID, _PERSON_HOUSEHOLD_ID, "age", *_POINTERS] + benefit = isinstance(rule, BenefitUnitRule) + test = rule.dependent_child.financial_independence if benefit else None + person_columns = [_PERSON_ID, _PERSON_HOUSEHOLD_ID, "partner_person_id"] + person_columns += [ + name + for name in ("age", "parent_1_person_id", "parent_2_person_id") + if benefit or name in person + ] if test is not None: person_columns.append("usual_weekly_hours") missing = [name for name in person_columns if name not in person.columns] @@ -718,7 +1019,7 @@ def read( f"{entity}.{column}" for entity, columns in _ID_COLUMNS.items() for column in columns - if _exceeds_int64(subset[entity][column]) + if column in subset[entity] and _exceeds_int64(subset[entity][column]) ] if wide: raise ValueError( @@ -737,6 +1038,8 @@ def read( index = pd.Index(ids) def rows(table: pd.DataFrame, entity: str, column: str) -> np.ndarray: + if column not in table: + return np.full(len(table), -1, dtype=np.int64) values = table[column] present = values.notna().to_numpy() found = np.full(len(values), -1, dtype=np.int64) @@ -756,11 +1059,13 @@ def rows(table: pd.DataFrame, entity: str, column: str) -> np.ndarray: return cls( ids=ids, household_ids=household_ids, - age=person["age"].to_numpy(dtype=np.float64), + age=person["age"].to_numpy(dtype=np.float64) + if "age" in person + else np.zeros(len(person)), hours=None if test is None else person["usual_weekly_hours"].to_numpy(dtype=np.float64), - partner=rows(person, "person", "partner_person_id"), + partner=reciprocal_partner_rows(person, household), parent_1=rows(person, "person", "parent_1_person_id"), parent_2=rows(person, "person", "parent_2_person_id"), reference=reference_by_household[household_rows], @@ -907,3 +1212,76 @@ def _int64_ids(values: pd.Series, label: str) -> np.ndarray: "int64, which cannot represent them." ) return values.to_numpy(dtype=np.int64) + + +def _roster_identity(value: object) -> tuple[type, object]: + if ( + isinstance(value, bool) + or not isinstance(value, (str, Integral)) + or isinstance(value, str) + and not value + ): + raise ValueError( + "Calibration membership identities must be strings or integers." + ) + return (str, value) if isinstance(value, str) else (int, int(value)) + + +def prepare_roster_counts( + tables: dict[str, pd.DataFrame], entries: object, required: Mapping[str, object] +) -> dict[str, set[str]]: + """Validate and prepare declared roster counts on independent caller tables.""" + if not isinstance(entries, (list, tuple)): + raise ValueError("Calibration roster_counts must be a list.") + derived: dict[str, set[str]] = {} + for entry in entries: + if not isinstance(entry, Mapping): + raise ValueError("Calibration roster count must be an object.") + entity, members, name = ( + entry.get("entity"), + entry.get("member_entity"), + entry.get("name"), + ) + group_id, member_id, membership = ( + entry.get("entity_id"), + entry.get("member_id"), + entry.get("membership"), + ) + if ( + entity not in required + or members not in required + or not isinstance(name, str) + or not name + or name.endswith("_id") + or name in required[entity] + or name in derived.get(entity, ()) + or group_id not in required[entity] + or member_id not in required[members] + or membership not in required[members] + or not isinstance(entry.get("representation"), str) + or not entry["representation"] + ): + raise ValueError( + "Calibration roster count needs declared identities and representation." + ) + groups = [ + _roster_identity(value) for value in tables[entity][group_id].tolist() + ] + identities = [ + _roster_identity(value) for value in tables[members][member_id].tolist() + ] + if len(set(groups)) != len(groups) or len(set(identities)) != len(identities): + raise ValueError( + "Calibration roster contains duplicate membership identities." + ) + counts = dict.fromkeys(groups, 0) + for value in tables[members][membership].tolist(): + key = _roster_identity(value) + if key not in counts: + raise ValueError("Calibration roster contains orphan membership.") + counts[key] += 1 + if any(count == 0 for count in counts.values()): + raise ValueError("Calibration roster contains empty households or groups.") + tables[entity][name] = [counts[key] for key in groups] + derived.setdefault(entity, set()).add(name) + return derived diff --git a/packages/microcosm-frame/tests/engine_free/shared/test_axiom_relations.py b/packages/microcosm-frame/tests/engine_free/shared/test_axiom_relations.py new file mode 100644 index 0000000000..43eb6a486b --- /dev/null +++ b/packages/microcosm-frame/tests/engine_free/shared/test_axiom_relations.py @@ -0,0 +1,539 @@ +"""Relation declaration and native API contract tests; no substitute evaluator. + +The recording program only returns preselected output arrays and captures the +arguments supplied to the documented dense API. It does not implement Axiom +formulas. Real dense execution remains conditional on the installed extension. +""" + +import importlib.util +import json +import sys +from dataclasses import dataclass, replace +from types import ModuleType, SimpleNamespace + +import numpy as np +import pandas as pd +import pytest +from hypothesis import given, settings +from hypothesis import strategies as st + +from microcosm.frame import Frame, WeightKind, Weights +from microcosm.frame.adapters.axiom import ( + AxiomEngine, + AxiomRelation, + _source_relation_arguments, + axiom_engine_ref, +) +from microcosm.frame.rules_kernels import SimulateCounterfactualKernel +from microcosm.frame.schema import EntitySchema +from test_support.paths import paths_for + +ROOT = paths_for("microcosm-frame").tests / "fixtures/rulespec-zz" +MODULE = ROOT / "zz/policies/tests/axiom_toy_relations.yaml" +SCHEMA = EntitySchema(group_entities=("household", "tax_unit")) + + +def _relation(**updates): + values = dict( + predicate="toy_member", + key="toy_member:0:1", + current_slot=0, + related_slot=1, + arguments=("TaxUnit", "Person"), + root_entity="tax_unit", + related_entity="person", + membership_column="person_tax_unit_id", + related_inputs={"toy_income": "income"}, + cardinality="many", + period="2025", + ) + values.update(updates) + return AxiomRelation(**values) + + +def _adapter(*relations, **updates): + values = dict( + module=MODULE, + rulespec_roots=(ROOT,), + schema=SCHEMA, + entity_names={ + "person": "Person", + "household": "Household", + "tax_unit": "TaxUnit", + }, + nesting={"tax_unit": "household"}, + relations=relations or (_relation(),), + ) + values.update(updates) + return AxiomEngine(**values) + + +def _frame(ids=(3, 1, 2), groups=(11, 10, 10), income=None, roles=None): + group_ids = sorted(set(groups)) + people = pd.DataFrame( + { + "person_id": ids, + "person_household_id": np.ones(len(ids), dtype=np.int64), + "person_tax_unit_id": groups, + "income": np.ones(len(ids)) if income is None else income, + "role": ["a"] * len(ids) if roles is None else roles, + } + ) + return Frame( + tables={ + "person": people, + "household": pd.DataFrame({"household_id": [1]}), + "tax_unit": pd.DataFrame( + {"tax_unit_id": group_ids, "scale": np.ones(len(group_ids))} + ), + }, + schema=SCHEMA, + weights={"household": Weights(np.ones(1), WeightKind.DESIGN)}, + ) + + +@dataclass +class _ContractBatch: + offsets: np.ndarray + inputs: dict[str, np.ndarray] + related_row_count: int | None = None + + +class _RecordingProgram: + """Records the native signature; never evaluates a RuleSpec expression.""" + + root_entity = "TaxUnit" + root_inputs = ["scale"] + + def __init__(self, relations=None): + self.relations = ( + relations + if relations is not None + else [ + SimpleNamespace( + name="toy_member", + key="toy_member:0:1", + current_slot=0, + related_slot=1, + related_inputs=("toy_income",), + ) + ] + ) + self.calls = [] + + def execute(self, **arguments): + self.calls.append(arguments) + return {"outputs": {"toy_total": np.zeros(len(arguments["inputs"]["scale"]))}} + + execute_f64 = execute + + +def _wire(adapter, monkeypatch, program=None): + program = program or _RecordingProgram() + monkeypatch.setattr( + adapter, + "_program", + lambda entity, **kw: program if entity == "tax_unit" else None, + ) + monkeypatch.setattr( + adapter, + "_derived_metadata", + lambda: { + "toy_total": SimpleNamespace( + entity="TaxUnit", dtype="decimal", period="Year" + ) + }, + ) + dense = ModuleType("axiom_rules_engine.dense") + dense.DenseRelationBatch = _ContractBatch + monkeypatch.setitem(sys.modules, "axiom_rules_engine.dense", dense) + return program + + +@pytest.mark.parametrize( + "updates,match", + [ + ({"predicate": ""}, "non-empty"), + ({"current_slot": True}, "binary coordinates"), + ({"related_slot": 0}, "binary coordinates"), + ({"arguments": ("TaxUnit",)}, "two engine"), + ({"arguments": "ab"}, "arguments must be a sequence"), + ({"arguments": {"TaxUnit": 0, "Person": 1}}, "arguments must be a sequence"), + ({"key": "toy_member:1:0"}, "key must agree"), + ({"cardinality": "optional"}, "cardinality"), + ({"related_inputs": {"toy_income": ""}}, "related_inputs"), + ({"role_column": "role"}, "supplied together"), + ({"role_column": "role", "roles": ()}, "unique non-empty"), + ({"role_column": "role", "roles": "head"}, "roles must be a sequence"), + ({"role_column": "role", "roles": {"head": 1}}, "roles must be a sequence"), + ({"role_column": "role", "roles": ("a", "a")}, "unique non-empty"), + ], +) +def test_manifest_shape_refuses_invalid_declarations(updates, match): + with pytest.raises(ValueError, match=match): + _relation(**updates) + + +@pytest.mark.parametrize( + "relation,match", + [ + (_relation(root_entity="household"), "partition"), + (_relation(arguments=("Person", "TaxUnit")), "argument slots"), + ], +) +def test_manifest_matches_physical_partition_and_entity_slots(relation, match): + with pytest.raises(ValueError, match=match): + _adapter(relation) + + +def test_duplicate_binding_and_undeclared_period_refused(): + from microcosm.frame.adapters.axiom import AxiomPeriod + + with pytest.raises(ValueError, match="Duplicate relation"): + _adapter(_relation(), _relation()) + with pytest.raises(ValueError, match="period.*not declared"): + _adapter(periods={"2026": AxiomPeriod("2026-01-01", "2026-12-31", "year")}) + with pytest.raises(TypeError, match="AxiomRelation"): + _adapter(relations=({},)) + + +def test_source_tuple_and_compiled_coordinates_are_verified(monkeypatch, tmp_path): + adapter = _adapter() + program = _wire(adapter, monkeypatch) + adapter.assert_declared_relations("tax_unit") + assert _source_relation_arguments(MODULE, (ROOT,), "toy_member") == ( + "TaxUnit", + "Person", + ) + assert _source_relation_arguments( + MODULE, (ROOT,), "zz:policies/tests/axiom_toy_relations#relation.toy_member" + ) == ("TaxUnit", "Person") + for field, value, match in [ + ("current_slot", 1, "slots disagree"), + ("related_slot", 0, "slots disagree"), + ("name", "other", "slots disagree"), + ("key", "other:0:1", "explicit relation declaration"), + ("related_inputs", ("other",), "input union"), + ]: + previous = getattr(program.relations[0], field) + setattr(program.relations[0], field, value) + with pytest.raises((ValueError, NotImplementedError), match=match): + adapter.assert_declared_relations("tax_unit") + setattr(program.relations[0], field, previous) + bad = tmp_path / "rulespec-zz" + bad.mkdir() + source = bad / "x.yaml" + source.write_text( + MODULE.read_text().replace("[TaxUnit, Person]", "[Person, TaxUnit]") + ) + scoped = bad / "zz/policies/tests/alias.yaml" + scoped.parent.mkdir(parents=True) + scoped.write_text( + MODULE.read_text().replace("predicate: toy_member", "predicate: authored_alias") + ) + assert _source_relation_arguments( + scoped, (bad,), "zz:policies/tests/alias#relation.toy_member" + ) == ("TaxUnit", "Person") + adapter._module, adapter._rulespec_roots = source, (bad,) + with pytest.raises(ValueError, match="source arguments"): + adapter.assert_declared_relations("tax_unit") + + +@pytest.mark.parametrize( + "problem,match", + [ + ("arity", "source must be binary"), + ("outside", "escapes"), + ("import_type", "imports must name"), + ("missing", "missing or ambiguous"), + ("ambiguous", "missing or ambiguous"), + ("conflicting", "conflicting source"), + ], +) +def test_source_relation_declaration_integrity(tmp_path, problem, match): + root = tmp_path / "rulespec-zz" + root.mkdir() + source = root / "x.yaml" + source.write_text(MODULE.read_text()) + roots = (root,) + if problem == "arity": + source.write_text(source.read_text().replace("arity: 2", "arity: 3")) + elif problem == "outside": + roots = (tmp_path / "other",) + elif problem == "import_type": + source.write_text(source.read_text() + "\nimports: [123]\n") + elif problem == "missing": + source.write_text(source.read_text() + "\nimports: [missing.yaml]\n") + elif problem == "ambiguous": + other = tmp_path / "other" + (other / "zz").mkdir(parents=True) + (root / "zz").mkdir() + (root / "zz/import.yaml").write_text("{}") + (other / "zz/import.yaml").write_text("{}") + roots = (root, other) + source.write_text(source.read_text() + "\nimports: [zz:import]\n") + elif problem == "conflicting": + (root / "other.yaml").write_text( + source.read_text().replace("[TaxUnit, Person]", "[Person, TaxUnit]") + ) + source.write_text(source.read_text() + "\nimports: [other.yaml]\n") + with pytest.raises(ValueError, match=match): + _source_relation_arguments(source, roots, "toy_member") + + +def test_extra_manifest_and_undeclared_relation_refuse(monkeypatch): + adapter = _adapter() + _wire(adapter, monkeypatch, _RecordingProgram(relations=[])) + with pytest.raises(ValueError, match="compiled keys"): + adapter.assert_declared_relations("tax_unit") + adapter = _adapter(relations=None) + _wire(adapter, monkeypatch) + with pytest.raises(NotImplementedError, match="explicit relation declaration"): + adapter.materialize(_frame(), ["toy_total"], "2025") + + +def test_related_input_union_and_graph_dependencies(monkeypatch): + relation = _relation( + related_inputs={"toy_income": "income", "toy_eligible": "eligible"}, + role_column="role", + roles=("a",), + ) + adapter = _adapter(relation) + program = _wire(adapter, monkeypatch) + program.relations.append( + SimpleNamespace( + name="toy_member", + key="toy_member:0:1", + current_slot=0, + related_slot=1, + related_inputs=("toy_eligible",), + ) + ) + assert adapter.variables() == ["scale", "toy_eligible", "toy_income"] + assert adapter.relation_dependencies("tax_unit") == ( + ("person", "eligible"), + ("person", "income"), + ("person", "person_id"), + ("person", "person_tax_unit_id"), + ("person", "role"), + ("tax_unit", "tax_unit_id"), + ) + frame = _frame() + frame._tables["person"]["eligible"] = True + adapter.materialize(frame, ["toy_total"], "2025") + batch = program.calls[-1]["relations"]["toy_member:0:1"] + assert set(batch.inputs) == {"toy_income", "toy_eligible"} + assert batch.offsets.tolist() == [0, 2, 3] + assert batch.inputs["toy_eligible"].tolist() == [True, True, True] + + +def test_relation_without_related_inputs_states_native_row_count(monkeypatch): + adapter = _adapter(_relation(related_inputs={})) + program = _RecordingProgram() + program.relations[0].related_inputs = () + _wire(adapter, monkeypatch, program) + adapter.materialize(_frame(), ["toy_total"], "2025") + batch = program.calls[-1]["relations"]["toy_member:0:1"] + assert batch.inputs == {} + assert batch.offsets.tolist() == [0, 2, 3] + assert batch.related_row_count == 3 + + +@pytest.mark.parametrize("mode", ("override", "solve")) +@pytest.mark.parametrize("column", ("toy_income", "income")) +def test_bridge_aliased_related_input_override_refused(monkeypatch, mode, column): + adapter = _adapter() + program = _wire(adapter, monkeypatch) + kernel = SimulateCounterfactualKernel({"fixture": adapter}) + component = { + "engine_ref": "fixture", + "variables": ["toy_total"], + "period": "2025", + "input_overrides": [], + } + options = {} + if mode == "override": + component["input_overrides"] = [ + {"entity": "person", "column": column, "value": 100} + ] + else: + options = { + "trial": ("tax_unit", np.asarray([100.0, 200.0])), + "solve_input": { + "engine_ref": "fixture", + "entity": "person", + "column": column, + }, + } + with pytest.raises(NotImplementedError, match="aliased related Axiom inputs"): + kernel._evaluate(_frame(), [component], [], **options) + assert program.calls == [] + + +@pytest.mark.parametrize("mode", ("override", "solve")) +def test_bridge_unaliased_related_input_override_reaches_native_batch( + monkeypatch, mode +): + adapter = _adapter(_relation(related_inputs={"toy_income": "toy_income"})) + program = _wire(adapter, monkeypatch) + frame = _frame() + frame._tables["person"]["toy_income"] = frame.table("person")["income"] + kernel = SimulateCounterfactualKernel({"fixture": adapter}) + component = { + "engine_ref": "fixture", + "variables": ["toy_total"], + "period": "2025", + "input_overrides": [], + } + options = {} + if mode == "override": + component["input_overrides"] = [ + {"entity": "person", "column": "toy_income", "value": 100} + ] + expected = [100.0, 100.0, 100.0] + else: + options = { + "trial": ("tax_unit", np.asarray([100.0, 200.0])), + "solve_input": { + "engine_ref": "fixture", + "entity": "person", + "column": "toy_income", + }, + } + expected = [100.0, 100.0, 200.0] + kernel._evaluate(frame, [component], [], **options) + batch = program.calls[-1]["relations"]["toy_member:0:1"] + assert batch.inputs["toy_income"].tolist() == expected + + +@settings(max_examples=40, deadline=None, database=None) +@given( + st.lists( + st.integers(min_value=1, max_value=500), min_size=1, max_size=20, unique=True + ), + st.data(), +) +def test_related_rows_order_by_identity_and_preserve_equal_values(ids, data): + permutation = data.draw(st.permutations(ids)) + groups = [10 + value % 3 for value in permutation] + adapter = _adapter() + # Capture the native API batch without implementing the engine evaluator. + with pytest.MonkeyPatch.context() as patch: + program = _wire(adapter, patch) + frame = _frame(permutation, groups, income=np.asarray(permutation, dtype=float)) + adapter.materialize(frame, ["toy_total"], "2025") + batch = program.calls[-1]["relations"]["toy_member:0:1"] + expected = [ + identity + for group in sorted(set(groups)) + for identity in sorted(ids) + if 10 + identity % 3 == group + ] + assert batch.inputs["toy_income"].tolist() == expected + assert batch.offsets.tolist() == [ + 0, + *np.cumsum([groups.count(group) for group in sorted(set(groups))]).tolist(), + ] + assert len(batch.inputs["toy_income"]) == len(set(ids)) + equal = _frame(permutation, groups) + adapter.materialize(equal, ["toy_total"], "2025") + batch = program.calls[-1]["relations"]["toy_member:0:1"] + assert batch.inputs["toy_income"].tolist() == [1.0] * len(ids) + + +def test_role_selection_empty_sets_cardinality_and_arithmetic(monkeypatch): + adapter = _adapter(_relation(role_column="role", roles=("b",)), arithmetic="f64") + program = _wire(adapter, monkeypatch) + adapter.materialize(_frame(roles=["b", "a", "a"]), ["toy_total"], "2025") + batch = program.calls[-1]["relations"]["toy_member:0:1"] + assert batch.offsets.tolist() == [0, 0, 1] + adapter = _adapter(_relation(cardinality="one")) + _wire(adapter, monkeypatch) + with pytest.raises(ValueError, match="exactly one"): + adapter.materialize(_frame(), ["toy_total"], "2025") + + +@pytest.mark.parametrize( + "problem,match", + [ + ("missing", "missing related columns"), + ("dangling", "dangling membership"), + ("duplicate", "unique non-null"), + ("period", "requested period"), + ("cross_household", "nest in exactly one"), + ], +) +def test_relation_data_integrity_refused(monkeypatch, problem, match): + adapter = _adapter() + _wire(adapter, monkeypatch) + frame = _frame() + if problem == "missing": + frame._tables["person"].drop(columns="income", inplace=True) + elif problem == "dangling": + frame._tables["person"].loc[0, "person_tax_unit_id"] = 999 + elif problem == "duplicate": + frame._tables["person"].loc[0, "person_id"] = 1 + elif problem == "cross_household": + frame._tables["person"].loc[1, "person_household_id"] = 2 + with pytest.raises(ValueError, match=match): + adapter.materialize( + frame, ["toy_total"], "2026" if problem == "period" else "2025" + ) + + +def test_absent_relation_identity_and_call_are_unchanged(monkeypatch): + absent = _adapter(relations=None) + explicit = _adapter(relations=()) + pins = dict( + engine_commit="a" * 40, + wheel_sha256="b" * 64, + rulespec_root=ROOT, + rulespec_commit="c" * 40, + ) + reference = json.loads(axiom_engine_ref(absent, **pins)) + assert "relations" not in reference + assert json.loads(axiom_engine_ref(explicit, **pins))["relations"] == [] + program = _wire(absent, monkeypatch, _RecordingProgram(relations=[])) + absent.materialize(_frame(), ["toy_total"], "2025") + assert "relations" not in program.calls[-1] + declared = json.loads(axiom_engine_ref(_adapter(), **pins)) + assert declared["relations"] == [_relation().to_dict()] + assert "roles" not in declared["relations"][0] + + +def test_relation_manifest_mapping_is_immutable(): + source = {"toy_income": "income"} + relation = _relation(related_inputs=source) + source["toy_income"] = "changed" + assert relation.related_inputs["toy_income"] == "income" + with pytest.raises(TypeError): + relation.related_inputs["toy_income"] = "changed" + assert replace(relation, role_column="role", roles=("a",)).to_dict()["roles"] == [ + "a" + ] + + +_ENGINE = importlib.util.find_spec("axiom_rules_engine") +if _ENGINE is not None: + from axiom_rules_engine.dense import NativeCompiledDenseProgram +else: + NativeCompiledDenseProgram = None + + +@pytest.mark.skipif( + NativeCompiledDenseProgram is None, + reason="Axiom dense native extension is not installed", +) +def test_real_dense_relation_fixture(): + predicate = "zz:policies/tests/axiom_toy_relations#relation.toy_member" + adapter = _adapter( + _relation( + predicate=predicate, + key=f"{predicate}:0:1", + related_inputs={"toy_income": "toy_income"}, + ) + ) + frame = _frame() + frame._tables["person"].rename(columns={"income": "toy_income"}, inplace=True) + result = adapter.materialize(frame, ["toy_total"], "2025") + assert result["toy_total"].tolist() == [2.0, 1.0] diff --git a/packages/microcosm-frame/tests/engine_free/shared/test_declared_units.py b/packages/microcosm-frame/tests/engine_free/shared/test_declared_units.py new file mode 100644 index 0000000000..ed8c11352e --- /dev/null +++ b/packages/microcosm-frame/tests/engine_free/shared/test_declared_units.py @@ -0,0 +1,414 @@ +"""Shared declared partitions: compatibility, identity, lineage and refusals.""" + +from types import MappingProxyType + +import numpy as np +import pandas as pd +import pytest +from hypothesis import given, settings +from hypothesis import strategies as st + +from microcosm.frame import Frame, WeightKind, Weights +from microcosm.frame.concepts import SOURCE_HOUSEHOLD_ID_COLUMN +from microcosm.frame.unit_construction import ( + MULTI_UNIT_RULE_FORMAT, + BenefitUnitRule, + DependentChildRule, + MultiUnitRule, + TaxUnitRule, + UnitDeclaration, + _build_partition, + _Pointers, + benefit_unit_membership, + build_benefit_units, + build_declared_units, + declared_unit_attributes, + prepare_roster_counts, + reciprocal_partner_rows, + units_per_household, +) +from test_support.microcosm_frame.concept_frames import concept_frames +from test_support.microcosm_frame.unit_rules import unit_rules + +PROPERTY = settings(max_examples=70, deadline=None) +FRAMES = concept_frames(max_households=4, max_members=5, min_id=-(10**8)) +TAX = TaxUnitRule( + entity="tax_unit", + pairing_eligibility_column="joint_assessment_eligible", + provenance={ + "eligibility": "Explicit upstream eligibility result, no implicit legal test." + }, +) +BENEFIT = BenefitUnitRule( + entity="family", + dependent_child=DependentChildRule(max_age=17, financial_independence=None), + unparented_child="own_unit", + split_parents="first_parent_unit", + provenance={"test": "fixture"}, +) +DECLARATION = MultiUnitRule( + units=( + UnitDeclaration(kind="benefit_unit", rule=BENEFIT), + UnitDeclaration(kind="tax_unit", rule=TAX), + ) +) + + +def _tables(): + person = pd.DataFrame( + { + "person_id": [10, 20, 30], + "person_household_id": [1, 1, 1], + "age": [45, 40, 10], + "partner_person_id": pd.array([20, 10, None], dtype="Int64"), + "parent_1_person_id": pd.array([None, None, 10], dtype="Int64"), + "parent_2_person_id": pd.array([None, None, 20], dtype="Int64"), + "joint_assessment_eligible": [True, True, False], + } + ) + household = pd.DataFrame( + { + "household_id": [1], + "reference_person_id": pd.array([20], dtype="Int64"), + SOURCE_HOUSEHOLD_ID_COLUMN: [901], + } + ) + return person, household + + +@PROPERTY +@given(tables=FRAMES, rule=unit_rules()) +def test_benefit_kind_is_byte_identical_to_legacy(tables, rule): + person, household = tables["person"], tables["household"] + legacy_units, legacy_membership = build_benefit_units(person, household, rule) + legacy = benefit_unit_membership(person, legacy_units, legacy_membership, rule) + built = build_declared_units( + person, household, MultiUnitRule.from_benefit_unit(rule) + )[rule.entity] + for left, right in ( + (legacy.units, built.units), + (legacy.person_unit, built.person_unit), + (legacy.person_role, built.person_role), + ): + assert left.to_json() == right.to_json() + if isinstance(left, pd.DataFrame): + pd.testing.assert_frame_equal(left, right) + else: + pd.testing.assert_series_equal(left, right) + + +@PROPERTY +@given(tables=FRAMES, data=st.data()) +def test_every_declared_partition_nests_and_preserves_identity(tables, data): + person, household = tables["person"].copy(), tables["household"].copy() + person[TAX.pairing_eligibility_column] = data.draw( + st.lists(st.booleans(), min_size=len(person), max_size=len(person)) + ) + household[SOURCE_HOUSEHOLD_ID_COLUMN] = household["household_id"] + 10**9 + built = build_declared_units(person, household, DECLARATION) + shuffled = person.iloc[list(data.draw(st.permutations(range(len(person)))))].copy() + shuffled["employment_income"] = data.draw(st.integers(0, 10**8)) + again = build_declared_units(shuffled, household, DECLARATION) + for item in DECLARATION.units: + group, other = built[item.entity], again[item.entity] + rule = item.rule + ids = group.units[rule.id_column] + assert group.person_unit.index.equals(person.index) + assert group.person_unit.notna().all() + assert ids.is_unique + assert set(group.person_unit) == set(ids) + nested = group.units.set_index(rule.id_column)[rule.household_column] + assert ( + nested.loc[group.person_unit.to_numpy()].tolist() + == person["person_household_id"].tolist() + ) + assert rule.source_household_column in group.units + assert all("weight" not in column for column in group.units) + pd.testing.assert_frame_equal(group.units, other.units) + pd.testing.assert_series_equal( + group.person_unit.sort_index(), other.person_unit.sort_index() + ) + pd.testing.assert_series_equal( + group.person_role.sort_index(), other.person_role.sort_index() + ) + tax = built[TAX.entity] + assert tax.person_role.isin(["head", "partner"]).all() + assert tax.person_unit.value_counts().max() <= 2 + + +@PROPERTY +@given(tables=FRAMES, data=st.data()) +def test_unit_weights_inherit_household_once(tables, data): + person, household = tables["person"].copy(), tables["household"].copy() + person[TAX.pairing_eligibility_column] = True + weights = np.asarray( + data.draw( + st.lists( + st.integers(1, 10000), min_size=len(household), max_size=len(household) + ) + ), + dtype=float, + ) + groups = build_declared_units(person, household, DECLARATION) + tables = {"person": person, "household": household} + for entity, group in groups.items(): + person[group.person_unit.name] = group.person_unit + tables[entity] = group.units + schema, _ = DECLARATION.entity_schema() + order = np.argsort(household["household_id"].to_numpy()) + tables["household"] = household.iloc[order].reset_index(drop=True) + frame = Frame( + tables=tables, + schema=schema, + weights={"household": Weights(values=weights[order], kind=WeightKind.DESIGN)}, + ) + for item in DECLARATION.units: + resolved = frame.resolve_weights(item.entity).values + counts = units_per_household(household, groups[item.entity].units, item.rule) + assert float(resolved.sum()) == float(np.dot(weights, counts)) + hh_weight = dict(zip(household["household_id"], weights, strict=True)) + assert resolved.tolist() == [ + hh_weight[value] + for value in groups[item.entity].units[item.rule.household_column] + ] + + +@PROPERTY +@given(eligible=st.tuples(st.booleans(), st.booleans())) +def test_pairing_requires_explicit_eligibility_on_both_members(eligible): + person, household = _tables() + person[TAX.pairing_eligibility_column] = [*eligible, False] + group = build_declared_units(person, household, DECLARATION)[TAX.entity] + paired = all(eligible) + assert group.person_unit.tolist() == [10, 10 if paired else 20, 30] + assert group.person_role.tolist() == [ + "head", + "partner" if paired else "head", + "head", + ] + + +@PROPERTY +@given(pointer=st.sampled_from([None, 10, 30])) +def test_a_nonreciprocal_or_self_pointer_never_forms_a_pair(pointer): + person, household = _tables() + person.loc[1, "partner_person_id"] = pointer + if pointer == 10: + # Make the first pointer self-referential instead. + person.loc[0, "partner_person_id"] = 10 + with pytest.raises(ValueError, match="asymmetric|self"): + reciprocal_partner_rows(person, household) + with pytest.raises(ValueError, match="asymmetric|self"): + build_declared_units(person, household, DECLARATION) + + +def test_serialized_declaration_is_strict_and_roundtrips(): + encoded = DECLARATION.to_dict() + assert encoded["format"] == MULTI_UNIT_RULE_FORMAT + assert MultiUnitRule.from_dict(encoded).to_dict() == encoded + assert MultiUnitRule.from_spec(encoded).to_dict() == encoded + schema, nesting = DECLARATION.entity_schema() + assert schema.group_entities == ("household", "family", "tax_unit") + assert nesting == {"family": "household", "tax_unit": "household"} + assert isinstance(nesting, MappingProxyType) + assert MultiUnitRule.from_benefit_unit(BENEFIT).units[0].rule is BENEFIT + + +@pytest.mark.parametrize( + "mutate, message", + [ + (lambda row: row.update(extra=True), "unexpected"), + (lambda row: row["units"][0].update(extra=True), "unexpected"), + (lambda row: row.update(format="unknown"), "format"), + (lambda row: row.update(units=[]), "non-empty"), + (lambda row: row.update(units={}), "Declared units must be a list\\."), + (lambda row: row["units"].append(row["units"][0]), "repeat"), + (lambda row: row["units"][0].update(kind="unknown"), "Unknown unit kind"), + (lambda row: row["units"][1]["rule"].update(extra=True), "unexpected"), + (lambda row: row["units"][1]["rule"].update(format="unknown"), "format"), + ( + lambda row: row["units"][1]["rule"].update(entity="person"), + "cannot be named", + ), + (lambda row: row["units"][1]["rule"].update(entity="TaxUnit"), "snake_case"), + ( + lambda row: row["units"][1]["rule"].update( + pairing_eligibility_column="bad column" + ), + "snake_case", + ), + ( + lambda row: row["units"][1]["rule"].update(provenance={"missing": ""}), + "provenance", + ), + ], +) +def test_bad_serialized_declarations_refuse(mutate, message): + encoded = DECLARATION.to_dict() + mutate(encoded) + with pytest.raises(ValueError, match=message): + MultiUnitRule.from_dict(encoded) + + +@pytest.mark.parametrize("units", [None, "tax_unit", [], [TAX]]) +def test_bad_declaration_lists_refuse(units): + with pytest.raises(ValueError, match="non-empty"): + MultiUnitRule(units=units) + + +def test_kind_and_rule_must_agree(): + with pytest.raises(ValueError, match="requires BenefitUnitRule"): + UnitDeclaration(kind="benefit_unit", rule=TAX) + with pytest.raises(ValueError, match="Unknown unit kind"): + UnitDeclaration(kind="unknown", rule=TAX) + with pytest.raises(TypeError, match="MultiUnitRule"): + person, household = _tables() + build_declared_units(person, household, BENEFIT) + + +@pytest.mark.parametrize( + "bad", + [ + None, + [True, None, False], + pd.array([True, None, False], dtype="boolean"), + [1, 1, 0], + ["true"] * 3, + ], +) +def test_missing_or_malformed_eligibility_refuses(bad): + person, household = _tables() + if bad is None: + person = person.drop(columns=TAX.pairing_eligibility_column) + else: + person[TAX.pairing_eligibility_column] = bad + with pytest.raises(ValueError, match="eligibility"): + build_declared_units(person, household, DECLARATION) + + +@pytest.mark.parametrize( + "heads, message", + [ + (np.array([0, 0]), "exactly one"), + (np.array([-1, 0, 0]), "exactly one"), + (np.array([0, 0, 3]), "exactly one"), + (np.array([0.0, 0.0, 0.0]), "exactly one"), + (np.array([1, 0, 2]), "own unit"), + ], +) +def test_registered_constructor_must_return_a_valid_partition(heads, message): + person, household = _tables() + rows = _Pointers.read(person, household, TAX) + with pytest.raises(ValueError, match=message): + _build_partition(person, household, TAX, rows, heads) + + +def test_registered_constructor_cannot_cross_households(): + person, household = _tables() + person.loc[2, "person_household_id"] = 2 + person.loc[2, ["parent_1_person_id", "parent_2_person_id"]] = None + household.loc[1] = [2, 30, 902] + rows = _Pointers.read(person, household, TAX) + with pytest.raises(ValueError, match="one household"): + _build_partition(person, household, TAX, rows, np.array([0, 0, 0])) + + +def test_changed_eligibility_changes_partition_identity(): + person, household = _tables() + before = build_declared_units(person, household, DECLARATION)[TAX.entity] + person[TAX.pairing_eligibility_column] = False + after = build_declared_units(person, household, DECLARATION)[TAX.entity] + # Structural membership is part of population identity: no legal proxy is + # hidden in prepared state, even when the head id remains the same. + assert before.person_unit.to_json() != after.person_unit.to_json() + + +def test_support_clones_keep_physical_ids_and_share_original_lineage(): + person, household = _tables() + clone = person.copy() + clone["person_id"] += 100 + clone["person_household_id"] += 1 + for column in ("partner_person_id", "parent_1_person_id", "parent_2_person_id"): + clone[column] += 100 + cloned_household = household.copy() + cloned_household["household_id"] += 1 + cloned_household["reference_person_id"] += 100 + both = build_declared_units( + pd.concat([person, clone], ignore_index=True), + pd.concat([household, cloned_household], ignore_index=True), + DECLARATION, + ) + for item in DECLARATION.units: + group, rule = both[item.entity], item.rule + assert group.units[rule.id_column].is_unique + assert set(group.units[rule.household_column]) == {1, 2} + assert set(group.units[rule.source_household_column]) == {901} + assert set(group.person_unit[:3]).isdisjoint(set(group.person_unit[3:])) + + +def test_partner_pointer_input_surface_is_required(): + person, household = _tables() + with pytest.raises(ValueError, match="need columns"): + reciprocal_partner_rows(person.drop(columns="partner_person_id"), household) + + +def test_tax_membership_does_not_need_dependency_facts(): + person, household = _tables() + person = person.drop(columns=["age", "parent_1_person_id", "parent_2_person_id"]) + declaration = UnitDeclaration(kind="tax_unit", rule=TAX) + group = build_declared_units( + person, household, MultiUnitRule(units=(declaration,)) + )[TAX.entity] + attributes = declared_unit_attributes( + person, group.units, group.person_unit, declaration + ) + assert attributes["n_members"].tolist() == [2, 1] + assert attributes["n_dependent_children"].tolist() == [0, 0] + assert attributes["youngest_dependent_child_age"].isna().all() + + +@pytest.mark.parametrize( + "defect,message", + [ + ("entries", "must be a list"), + ("entry", "must be an object"), + ("declaration", "declared identities"), + ("identity", "strings or integers"), + ("duplicate", "duplicate membership identities"), + ("orphan", "orphan membership"), + ("empty", "empty households"), + ], +) +def test_roster_preparation_is_validated_by_the_unit_framework(defect, message): + tables = { + "household": pd.DataFrame({"household_id": [10, 20]}), + "person": pd.DataFrame({"person_id": [1, 2], "person_household_id": [10, 20]}), + } + required = {entity: list(table.columns) for entity, table in tables.items()} + entries = [ + { + "entity": "household", + "member_entity": "person", + "name": "size", + "entity_id": "household_id", + "member_id": "person_id", + "membership": "person_household_id", + "representation": "one row per physical person", + } + ] + if defect == "entries": + entries = {} + elif defect == "entry": + entries = [None] + elif defect == "declaration": + entries[0]["representation"] = "" + elif defect == "identity": + tables["person"]["person_id"] = [True, False] + elif defect == "duplicate": + tables["person"]["person_id"] = [1, 1] + elif defect == "orphan": + tables["person"]["person_household_id"] = [10, 30] + else: + tables["person"]["person_household_id"] = [10, 10] + with pytest.raises(ValueError, match=message): + prepare_roster_counts(tables, entries, required) diff --git a/packages/microcosm-frame/tests/fixtures/rulespec-zz/zz/policies/tests/axiom_toy_relations.yaml b/packages/microcosm-frame/tests/fixtures/rulespec-zz/zz/policies/tests/axiom_toy_relations.yaml new file mode 100644 index 0000000000..f6d769d2e4 --- /dev/null +++ b/packages/microcosm-frame/tests/fixtures/rulespec-zz/zz/policies/tests/axiom_toy_relations.yaml @@ -0,0 +1,33 @@ +format: rulespec/v1 +module: + summary: Binary relation batching fixture for declared partition members. +inputs: + - name: toy_income + entity: Person + dtype: Decimal +rules: + - name: toy_person_income + kind: derived + entity: Person + dtype: Decimal + period: Year + source: microcosm test fixture + versions: + - effective_from: '2025-01-01' + formula: toy_income + - name: toy_member + kind: data_relation + data_relation: + predicate: toy_member + arity: 2 + arguments: [TaxUnit, Person] + source: microcosm test fixture + - name: toy_total + kind: derived + entity: TaxUnit + dtype: Decimal + period: Year + source: microcosm test fixture + versions: + - effective_from: '2025-01-01' + formula: sum(toy_member.toy_person_income) From a395c1e374940789100824c0621d00cabea4a034 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 1 Nov 2026 11:01:07 -0500 Subject: [PATCH 2/3] Refuse partial relations and wrong-entity overrides Sol's review of 76a137c9e found three defects, and an end-to-end test found a fourth: - A row mask could drop a relation member yet leave every unit populated, so no partition check saw it and the engine got partial data. Every rules kernel now refuses a node whose slices don't deliver its relations whole. - A related-input override on another entity wrote a column no batch reads, so it changed nothing. It now refuses, and a real root input of that name stays overridable. - transport.create@1 accepted non-family units. It keeps its old refusal, and transport.create@2 carries the wider domain. - Composition required id and membership columns in slices, which compile_graph refuses, so no relation node could compile. Tables always carry those columns, so only data columns need a slice. Composition and the rules kernels now share missing_relation_inputs. frame/kernels.py moves simulate.rules@1's hash, so the simulate parity pin is re-recorded with tools/graph_parity_repin.py. Co-Authored-By: Claude Opus 5.5 --- docs/transport-declared-units.md | 8 +- .../src/microcosm/build/transport/compose.py | 11 +- .../build/transport/unit_declarations.py | 5 +- .../shared/test_transport_declared_units.py | 35 ++- .../src/microcosm/frame/adapters/axiom.py | 35 ++- .../src/microcosm/frame/kernels.py | 61 +++- .../src/microcosm/frame/rules_kernels.py | 3 + .../shared/test_axiom_relations.py | 277 +++++++++++++++++- .../parity/kernels/simulate/pins.json | 2 +- 9 files changed, 398 insertions(+), 39 deletions(-) diff --git a/docs/transport-declared-units.md b/docs/transport-declared-units.md index 87ddca0e01..07b455a7af 100644 --- a/docs/transport-declared-units.md +++ b/docs/transport-declared-units.md @@ -83,9 +83,11 @@ root; `many` permits zero. A role selector requires both its column and distinct role values. This interface supports person members of genuine group partitions; binding-local views and arbitrary links remain separate work. -Rules nodes declare whole-table slices for related physical IDs, partition -membership, selector columns and every mapped related input; composition refuses -a missing dependency. The adapter gathers per-root offsets and input arrays, +Rules nodes declare whole-table slices for selector columns and every mapped +related input. Tables always carry their IDs and memberships, and no slice can +declare those, but a row mask on any dependency's entity still fails, because it +can drop a member and leave every unit populated. Composition and every rules +kernel apply this one check. The adapter gathers per-root offsets and input arrays, ordering people by physical identity and retaining distinct equal-valued people. Data-derived batches are built at execution and never enter prepared adapter configuration. Relation-enabled execution requires the native dense API with diff --git a/packages/microcosm-build/src/microcosm/build/transport/compose.py b/packages/microcosm-build/src/microcosm/build/transport/compose.py index dd2a6b2e40..3cf0c437b6 100644 --- a/packages/microcosm-build/src/microcosm/build/transport/compose.py +++ b/packages/microcosm-build/src/microcosm/build/transport/compose.py @@ -86,6 +86,7 @@ ) from microcosm.build.reference_profiles import resolve_reference_profiles from microcosm.calibrate.ordered_kernels import OUTPUT_TYPES +from microcosm.frame.kernels import missing_relation_inputs from microcosm.graph import ( ROWS_ALL, ArtifactInput, @@ -2830,16 +2831,10 @@ def _require_relation_inputs(spec, binding_id, inputs) -> None: dependencies = binding_relation_dependencies( bindings[0], schema=schema, entity_names=document["entity_names"] ) - declared = { - (item.entity, column) - for item in inputs - if item.rows == ROWS_ALL - for column in item.columns - } - missing = set(dependencies) - declared + missing = missing_relation_inputs(dependencies, inputs, schema) if missing: raise ValueError( - f"Rules binding {binding_id!r} needs whole-table relation inputs {sorted(missing)}." + f"Rules binding {binding_id!r} needs whole-table relation inputs {missing}." ) diff --git a/packages/microcosm-build/src/microcosm/build/transport/unit_declarations.py b/packages/microcosm-build/src/microcosm/build/transport/unit_declarations.py index 1d6325354c..1bd334a292 100644 --- a/packages/microcosm-build/src/microcosm/build/transport/unit_declarations.py +++ b/packages/microcosm-build/src/microcosm/build/transport/unit_declarations.py @@ -24,7 +24,10 @@ def create_unit_declaration(kernel: str, document: Mapping) -> MultiUnitRule: """Lift legacy benefit declarations or read the versioned declaration list.""" if kernel == CREATE_V1_REF: - return MultiUnitRule.from_benefit_unit(BenefitUnitRule.from_spec(document)) + rule = BenefitUnitRule.from_spec(document) + if rule.entity != "family": + raise ValueError("CREATE's declared schema requires family benefit units.") + return MultiUnitRule.from_benefit_unit(rule) if kernel == CREATE_V2_REF: return MultiUnitRule.from_spec(document) raise ValueError(f"Unknown transport CREATE kernel {kernel!r}.") diff --git a/packages/microcosm-build/tests/engine_free/shared/test_transport_declared_units.py b/packages/microcosm-build/tests/engine_free/shared/test_transport_declared_units.py index 453b53e7e7..1e67cb92e3 100644 --- a/packages/microcosm-build/tests/engine_free/shared/test_transport_declared_units.py +++ b/packages/microcosm-build/tests/engine_free/shared/test_transport_declared_units.py @@ -241,6 +241,30 @@ def test_create_parser_rejects_an_unregistered_version(): create_unit_declaration("transport.create@99", {}) +def test_legacy_create_keeps_its_family_only_domain(tmp_path): + """create@1 refuses what it always refused; create@2 carries the wider domain.""" + sources = write_population_sources(tmp_path / "sources", mass=1000) + document = {**UNIT_RULE.to_dict(), "entity": "assessment"} + legacy = create_node(sources) + legacy = replace( + legacy, + params={ + **legacy.params, + "unit_rule": canonical_text(document), + "unit_rule_sha256": digest(document), + }, + ) + with pytest.raises(ValueError, match="requires family benefit units"): + create_unit_declaration(TRANSPORT_CREATE.ref, document) + with pytest.raises(ValueError, match="requires family benefit units"): + _run_create(sources, legacy, TRANSPORT_CREATE) + declared = _versioned( + legacy, MultiUnitRule.from_benefit_unit(BenefitUnitRule.from_spec(document)) + ) + result = _run_create(sources, declared, TRANSPORT_CREATE_DECLARED) + assert "assessment" in result.frame.schema.group_entities + + def test_multiple_unit_attributes_produce_explicit_person_role_columns(tmp_path): def edit(tables): tables["person"]["is_disabled"] = True @@ -318,9 +342,6 @@ def edit(tables): @given( omitted=st.sampled_from( [ - ("family", "family_id"), - ("person", "person_id"), - ("person", "person_family_id"), ("person", "age"), ("person", "family_role"), ] @@ -351,13 +372,7 @@ def test_relation_inputs_must_be_explicit_graph_dependencies(tmp_path, omitted): "roles": ["head", "partner"], } ] - dependencies = { - ("family", "family_id"), - ("person", "person_id"), - ("person", "person_family_id"), - ("person", "age"), - ("person", "family_role"), - } + dependencies = {("person", "age"), ("person", "family_role")} node = Node( "relations", "simulate.rules_by_ref@1", diff --git a/packages/microcosm-frame/src/microcosm/frame/adapters/axiom.py b/packages/microcosm-frame/src/microcosm/frame/adapters/axiom.py index aa5b83f032..b11c0f7070 100644 --- a/packages/microcosm-frame/src/microcosm/frame/adapters/axiom.py +++ b/packages/microcosm-frame/src/microcosm/frame/adapters/axiom.py @@ -649,19 +649,28 @@ def relation_dependencies(self, entity: str) -> tuple[tuple[str, str], ...]: return tuple(sorted(dependencies)) def validate_input_override(self, entity: str, column: str) -> None: - """Refuse aliased related overrides until bridges can resolve them. + """Refuse related overrides a bridge cannot apply to their source. Bridges write physical table columns after checking logical engine - inputs. An aliased relation input needs both names resolved together; - accepting only the engine name would silently leave its source intact. + inputs. An aliased relation input needs both names resolved together, + and only the related entity's table feeds a batch. Accepting either + override would silently leave the source intact. """ - for relation in self._relations or (): + if not self._relations: + return + program = self._program(entity, missing_ok=True) + root_inputs = () if program is None else program.root_inputs + for relation in self._relations: for name, physical in relation.related_inputs.items(): - if ( - entity == relation.related_entity - and name != physical - and column in (name, physical) - ): + if column not in (name, physical): + continue + if entity != relation.related_entity: + if column not in root_inputs: + raise ValueError( + f"Related Axiom input {column!r} belongs to " + f"{relation.related_entity!r}, not {entity!r}." + ) + elif name != physical: raise NotImplementedError( "Input overrides of aliased related Axiom inputs are not supported." ) @@ -755,11 +764,9 @@ def _relation_batches( raise ValueError( f"Relation {relation.key!r} requires unique non-null related identities." ) - # A physical person identity is the deduplication key, never its - # input values; equal-valued people must remain separate targets. - selected = selected.drop_duplicates(identity).sort_values( - identity, kind="stable" - ) + # Physical identity orders the rows, never input values, so + # equal-valued people remain separate targets. + selected = selected.sort_values(identity, kind="stable") positions_by_root: dict[object, list[int]] = {} for position, root_id in enumerate( selected[relation.membership_column].to_numpy() diff --git a/packages/microcosm-frame/src/microcosm/frame/kernels.py b/packages/microcosm-frame/src/microcosm/frame/kernels.py index d453ca7b54..e93b66ca6a 100644 --- a/packages/microcosm-frame/src/microcosm/frame/kernels.py +++ b/packages/microcosm-frame/src/microcosm/frame/kernels.py @@ -8,6 +8,8 @@ from __future__ import annotations +from collections.abc import Iterable + import numpy as np import pandas as pd @@ -24,14 +26,48 @@ KernelBase, KernelContext, KernelResult, + Node, Numeric, Ownership, SeedSource, + Slice, StructuralDelta, source_hash, ) -__all__ = ["RulesKernel", "SimulateRulesKernel"] +__all__ = ["RulesKernel", "SimulateRulesKernel", "missing_relation_inputs"] + + +def missing_relation_inputs( + dependencies: Iterable[tuple[str, str]], + inputs: Iterable[Slice], + schema: EntitySchema, +) -> list[tuple[str, str]]: + """The relation dependencies a node's slices do not deliver whole. + + A table always carries its id and memberships, and a slice cannot declare + them, so only data columns need a slice. Related rows reach an engine only + through these slices, and a row mask can drop a member yet leave every unit + populated. No partition check sees that loss, so a mask on any dependency's + entity fails. + """ + + inputs = tuple(inputs) + masked = {item.entity for item in inputs if item.rows != ROWS_ALL} + declared = {(item.entity, column) for item in inputs for column in item.columns} + memberships = {schema.membership_column(group) for group in schema.group_entities} + + def carried(entity: str, column: str) -> bool: + return column == schema.entity_id_column(entity) or ( + entity == schema.person_entity and column in memberships + ) + + return sorted( + (entity, column) + for entity, column in set(dependencies) + if entity in masked + or ((entity, column) not in declared and not carried(entity, column)) + ) class SimulateRulesKernel(KernelBase): @@ -94,6 +130,28 @@ def implementation_hash(self) -> str: dependencies=self.capabilities.dependencies, ) + @staticmethod + def _require_relation_inputs(node: Node, engine: RulesEngine) -> None: + """Refuse a node whose slices cannot hold the engine's whole relations.""" + + dependencies = getattr(engine, "relation_dependencies", None) + if not callable(dependencies): + return + schema = engine.entity_schema() + missing = missing_relation_inputs( + { + dependency + for entity in schema.entities + for dependency in dependencies(entity) + }, + node.inputs, + schema, + ) + if missing: + raise ValueError( + f"Node {node.id!r} needs whole-table relation inputs {missing}." + ) + @staticmethod def _tables( context: KernelContext, schema: EntitySchema @@ -134,6 +192,7 @@ def run(self, context: KernelContext) -> KernelResult: """Call ``RulesEngine.materialize`` and index its arrays by entity id.""" engine_ref, variables, period = self._validated_params(context) + self._require_relation_inputs(context.node, self._engine) schema = self._engine.entity_schema() tables = self._tables(context, schema) diff --git a/packages/microcosm-frame/src/microcosm/frame/rules_kernels.py b/packages/microcosm-frame/src/microcosm/frame/rules_kernels.py index 861d24c266..7d3b9b4030 100644 --- a/packages/microcosm-frame/src/microcosm/frame/rules_kernels.py +++ b/packages/microcosm-frame/src/microcosm/frame/rules_kernels.py @@ -330,6 +330,9 @@ def _prepare(self, context: KernelContext, *, solver: bool = False): if schema is not None and schema != engine_schema: raise ValueError("Bridge components must share an entity schema.") schema = engine_schema + SimulateRulesKernel._require_relation_inputs( + context.node, self._delegates[ref]._engine + ) tables = SimulateRulesKernel._tables(context, schema) frame = Frame( tables=tables, diff --git a/packages/microcosm-frame/tests/engine_free/shared/test_axiom_relations.py b/packages/microcosm-frame/tests/engine_free/shared/test_axiom_relations.py index 43eb6a486b..1517d5a6e0 100644 --- a/packages/microcosm-frame/tests/engine_free/shared/test_axiom_relations.py +++ b/packages/microcosm-frame/tests/engine_free/shared/test_axiom_relations.py @@ -24,8 +24,34 @@ _source_relation_arguments, axiom_engine_ref, ) -from microcosm.frame.rules_kernels import SimulateCounterfactualKernel +from microcosm.frame.kernels import SimulateRulesKernel, missing_relation_inputs +from microcosm.frame.rules_kernels import ( + SimulateCounterfactualKernel, + SimulateRulesByRefKernel, + SimulateSolveZeroKernel, +) from microcosm.frame.schema import EntitySchema +from microcosm.graph import ( + Capabilities, + ContentStore, + Determinism, + Graph, + GraphError, + KernelBase, + KernelContext, + KernelRegistry, + KernelResult, + Node, + NodeRejected, + Numeric, + Owned, + SeedSource, + Slice, + SourceRef, + StructuralDelta, + compile_graph, + run_graph, +) from test_support.paths import paths_for ROOT = paths_for("microcosm-frame").tests / "fixtures/rulespec-zz" @@ -406,6 +432,255 @@ def test_bridge_unaliased_related_input_override_reaches_native_batch( assert batch.inputs["toy_income"].tolist() == expected +def _bridge_options(mode, entity, column): + component = { + "engine_ref": "fixture", + "variables": ["toy_total"], + "period": "2025", + "input_overrides": [], + } + if mode == "override": + component["input_overrides"] = [ + {"entity": entity, "column": column, "value": 100} + ] + return component, {} + return component, { + "trial": ("tax_unit", np.asarray([100.0, 200.0])), + "solve_input": {"engine_ref": "fixture", "entity": entity, "column": column}, + } + + +@pytest.mark.parametrize("mode", ("override", "solve")) +@pytest.mark.parametrize( + ("related_inputs", "entity", "column"), + ( + ({"toy_income": "income"}, "tax_unit", "toy_income"), + ({"toy_income": "income"}, "household", "toy_income"), + ({"toy_income": "toy_income"}, "tax_unit", "toy_income"), + ), +) +def test_bridge_override_on_another_entity_refused( + monkeypatch, mode, related_inputs, entity, column +): + """An override on a table no batch reads would change nothing, silently.""" + adapter = _adapter(_relation(related_inputs=related_inputs)) + program = _wire(adapter, monkeypatch) + component, options = _bridge_options(mode, entity, column) + kernel = SimulateCounterfactualKernel({"fixture": adapter}) + with pytest.raises(ValueError, match="belongs to 'person'"): + kernel._evaluate(_frame(), [component], [], **options) + assert program.calls == [] + + +@pytest.mark.parametrize("mode", ("override", "solve")) +def test_bridge_root_input_sharing_a_related_name_stays_overridable(monkeypatch, mode): + adapter = _adapter() + program = _wire(adapter, monkeypatch) + program.root_inputs = ["scale", "toy_income"] + frame = _frame() + frame._tables["tax_unit"]["toy_income"] = 1.0 + component, options = _bridge_options(mode, "tax_unit", "toy_income") + kernel = SimulateCounterfactualKernel({"fixture": adapter}) + kernel._evaluate(frame, [component], [], **options) + assert set(program.calls[-1]["inputs"]["toy_income"].tolist()) != {1.0} + batch = program.calls[-1]["relations"]["toy_member:0:1"] + assert batch.inputs["toy_income"].tolist() == [1.0, 1.0, 1.0] + + +class _RelationPopulation(KernelBase): + """Three people, two units, and a mask that keeps one member of the couple.""" + + ref = "test.relation_population@1" + capabilities = Capabilities( + determinism=Determinism.DETERMINISTIC, + numeric=Numeric.BITWISE, + seed_source=SeedSource.NONE, + structural=StructuralDelta.CREATE, + consumes_se=False, + ) + + def run(self, context): + frame = _frame(ids=(1, 2, 3), groups=(10, 10, 11), income=(10.0, 20.0, 30.0)) + frame.table("person").drop(columns="role", inplace=True) + frame.table("person")["selected"] = np.array([True, False, True]) + return KernelResult(frame=frame) + + +_STRUCTURAL = ( + ("person", "person_id"), + ("person", "person_tax_unit_id"), + ("tax_unit", "tax_unit_id"), +) + + +def _relation_node(*, rows="all", omitted=(), kernel=SimulateRulesByRefKernel.ref): + person = ["selected"] if ("person", "income") in omitted else ["income", "selected"] + return Node( + "relations", + kernel, + population="create", + inputs=( + Slice("person", tuple(person), rows=rows), + Slice("tax_unit", ("scale",)), + ), + outputs=(Owned("tax_unit", "toy_total", "float64"),), + params={"engine_ref": "fixture", "variables": ("toy_total",), "period": "2025"}, + ) + + +def _relation_graph(node): + create = Node( + "create", + _RelationPopulation.ref, + sources=("raw",), + structural=StructuralDelta.CREATE, + outputs=( + Owned("person", "income", "float64"), + Owned("person", "selected", "bool"), + Owned("tax_unit", "scale", "float64"), + ), + ) + return Graph("relations", (SourceRef("raw", "raw-bytes-v1"),), (create, node)) + + +def _graph_with_identity_slice(): + node = _relation_node() + return _relation_graph( + replace(node, inputs=(Slice("person", ("person_id",)), *node.inputs[1:])) + ) + + +def _run_relation_graph(tmp_path, adapter, node): + registry = KernelRegistry() + registry.register(_RelationPopulation()) + registry.register(SimulateRulesByRefKernel({"fixture": adapter})) + raw = tmp_path / "raw.txt" + raw.write_text("source") + return run_graph( + compile_graph(_relation_graph(node)), + sources={"raw": raw}, + store=ContentStore(tmp_path / "store"), + kernels=registry, + resume="forbid", + decisions=(), + ) + + +def test_graph_refuses_a_row_mask_that_drops_a_relation_member(monkeypatch, tmp_path): + """Both units keep a member here, so only the declared mask shows the loss.""" + adapter = _adapter() + program = _wire(adapter, monkeypatch) + with pytest.raises((NodeRejected, ValueError), match="whole-table relation inputs"): + _run_relation_graph(tmp_path, adapter, _relation_node(rows="selected")) + assert program.calls == [] + + +def test_graph_gives_the_engine_every_member_of_a_whole_slice(monkeypatch, tmp_path): + adapter = _adapter() + program = _wire(adapter, monkeypatch) + _run_relation_graph(tmp_path, adapter, _relation_node()) + batch = program.calls[-1]["relations"]["toy_member:0:1"] + assert batch.offsets.tolist() == [0, 2, 3] + assert batch.inputs["toy_income"].tolist() == [10.0, 20.0, 30.0] + + +def _entry_points(adapter): + bridge = { + "input_overrides": "[]", + "output_coefficients": "[]", + } + solve = { + **bridge, + "solve_input": "{}", + "bracket": (0.0, 1.0), + "tolerance": 1.0, + "iterations": 1, + } + return { + "rules": (SimulateRulesKernel("fixture", adapter), {}), + "by_ref": (SimulateRulesByRefKernel({"fixture": adapter}), {}), + "counterfactual": (SimulateCounterfactualKernel({"fixture": adapter}), bridge), + "solve": (SimulateSolveZeroKernel({"fixture": adapter}), solve), + } + + +@pytest.mark.parametrize("entry", ("rules", "by_ref", "counterfactual", "solve")) +@pytest.mark.parametrize( + "problem", + ({"rows": "selected"}, {"omitted": (("person", "income"),)}), + ids=("masked", "undeclared"), +) +def test_every_rules_entry_point_refuses_a_partial_relation( + monkeypatch, entry, problem +): + adapter = _adapter() + program = _wire(adapter, monkeypatch) + kernel, params = _entry_points(adapter)[entry] + node = _relation_node(**problem) + node = replace(node, params={**node.params, **params}) + frame = _frame(ids=(1, 3), groups=(10, 11)) + context = KernelContext( + node=node, + tables={entity: frame.table(entity) for entity in frame.entities}, + weights={}, + strata=pd.Series(dtype=object), + params=node.params, + rng=np.random.default_rng(0), + sources={}, + ) + with pytest.raises(ValueError, match="whole-table relation inputs"): + kernel.run(context) + assert program.calls == [] + + +@settings(max_examples=40, deadline=None, database=None) +@given( + entity=st.sampled_from(("person", "tax_unit", "household")), + masked=st.booleans(), + declared=st.booleans(), +) +def test_a_node_passes_only_with_every_relation_dependency_whole( + entity, masked, declared +): + """Any mask on a dependency's entity or an undeclared data column refuses.""" + adapter = _adapter() + inputs = [ + Slice("person", ("income",) if declared else ("selected",)), + Slice("tax_unit", ("scale",)), + Slice("household", ("region",)), + ] + if masked: + inputs = [ + replace(item, columns=(*item.columns, "keep"), rows="keep") + if item.entity == entity + else item + for item in inputs + ] + node = replace(_relation_node(), inputs=tuple(inputs)) + if not declared or (masked and entity != "household"): + with pytest.raises(ValueError, match="whole-table relation inputs"): + SimulateRulesKernel._require_relation_inputs(node, adapter) + else: + SimulateRulesKernel._require_relation_inputs(node, adapter) + + +def test_structural_columns_need_no_slice_because_no_graph_can_declare_them(): + """A table carries its id and memberships; compile refuses a slice of them.""" + assert missing_relation_inputs(_STRUCTURAL, (), SCHEMA) == [] + assert missing_relation_inputs( + _STRUCTURAL, (Slice("person", ("income",), rows="selected"),), SCHEMA + ) == [("person", "person_id"), ("person", "person_tax_unit_id")] + graph = _graph_with_identity_slice() + with pytest.raises(GraphError, match="which no node owns"): + compile_graph(graph) + + +def test_absent_relation_declaration_needs_no_extra_slice(): + adapter = _adapter(relations=None) + node = _relation_node(rows="selected", omitted=(("person", "income"),)) + SimulateRulesKernel._require_relation_inputs(node, adapter) + + @settings(max_examples=40, deadline=None, database=None) @given( st.lists( diff --git a/packages/microcosm-graph/tests/fixtures/parity/kernels/simulate/pins.json b/packages/microcosm-graph/tests/fixtures/parity/kernels/simulate/pins.json index 15f5d879b8..1a3cecc086 100644 --- a/packages/microcosm-graph/tests/fixtures/parity/kernels/simulate/pins.json +++ b/packages/microcosm-graph/tests/fixtures/parity/kernels/simulate/pins.json @@ -1 +1 @@ -{"dependencies":{},"implementation_hash":"eed54eaf27b53faf446068aa833ec16ab398840faf870445dbdcd6a4604c566c","kernel":"simulate.rules@1","node":"simulate","node_key":"a643736e821e841fa74996f0768d61110fe4d6e590d8db4286fd241f7d75b123","numeric":"bitwise","platform":"arm64/darwin/py3.14","platforms":{"arm64/darwin/py3.14":{"direct":"direct.csv","node_key":"a643736e821e841fa74996f0768d61110fe4d6e590d8db4286fd241f7d75b123"}},"seed":null} +{"dependencies":{},"implementation_hash":"60664a4f28b84a693dd5f3b916ad04be6b9d160c090e549659a674b0c2d81e42","kernel":"simulate.rules@1","node":"simulate","node_key":"5968c72d2bb1f0c909946e3a75d6757ba3bd5d945c4dd95e2c9c680a993f3536","numeric":"bitwise","platform":"arm64/darwin/py3.14","platforms":{"arm64/darwin/py3.14":{"direct":"direct.csv","node_key":"5968c72d2bb1f0c909946e3a75d6757ba3bd5d945c4dd95e2c9c680a993f3536"}},"seed":null} From 9ecfa8742c6491412e18fc6682fbd968ec7fcbfe Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 1 Nov 2026 11:18:47 -0500 Subject: [PATCH 3/3] Follow the create@1 parser in its citation test The NZ source-citation test pinned the old one-line lift. The parser now binds the rule first, to keep the family-only refusal. Co-Authored-By: Claude Opus 5.5 --- .../tests/engine_free/shared/test_nz_spec_package.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/packages/microcosm-build/tests/engine_free/shared/test_nz_spec_package.py b/packages/microcosm-build/tests/engine_free/shared/test_nz_spec_package.py index 6334f759d5..2eb64d21a8 100644 --- a/packages/microcosm-build/tests/engine_free/shared/test_nz_spec_package.py +++ b/packages/microcosm-build/tests/engine_free/shared/test_nz_spec_package.py @@ -3262,10 +3262,8 @@ def test_create_and_donor_reader_citations_still_describe_unit_inputs(self) -> N assert TransportCreateKernel.ref == unit_declarations.CREATE_V1_REF parser_source = inspect.getsource(unit_declarations.create_unit_declaration) assert "if kernel == CREATE_V1_REF:" in parser_source - assert ( - "MultiUnitRule.from_benefit_unit(BenefitUnitRule.from_spec(document))" - in parser_source - ) + assert "rule = BenefitUnitRule.from_spec(document)" in parser_source + assert "return MultiUnitRule.from_benefit_unit(rule)" in parser_source declaration = unit_declarations.create_unit_declaration( TransportCreateKernel.ref, rule )