Add opt-in declared tax-unit partitions and validated binary Axiom relation bindings, with NZ benefit units as a declared unit kind

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.

docs/transport-declared-units.md

@@ -0,0 +1,106 @@
+# 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.<entity>_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 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
+`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.

packages/microcosm-build/src/microcosm/build/calibration_columns.py

@@ -16,14 +16,14 @@
 
 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,
 )
+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

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,
@@ -77,9 +76,13 @@
     validate_transport_activation,
 )
 from .entitlement import money_node_factor_entries
-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",
@@ -139,6 +142,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}.")
@@ -149,7 +157,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])
@@ -182,7 +190,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: {
@@ -282,7 +290,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],

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],