diff --git a/changelog.d/transport-population-graph-kernels.added.md b/changelog.d/transport-population-graph-kernels.added.md new file mode 100644 index 000000000..e9fcb7517 --- /dev/null +++ b/changelog.d/transport-population-graph-kernels.added.md @@ -0,0 +1 @@ +Add explicit donor/facts/RuleSpec codecs and ten country-neutral transport population graph kernels, with destination-scale design weights, benefit units, concept encoding, fact-backed quantile maps and geography support, deterministic receipt assignment, and scenario input overrides. Synthetic engine-free graph properties check identity, memoization, provenance, mass, and eligibility. diff --git a/docs/concept-schema-transport-adr.md b/docs/concept-schema-transport-adr.md index 6367f711b..c3480a629 100644 --- a/docs/concept-schema-transport-adr.md +++ b/docs/concept-schema-transport-adr.md @@ -11,7 +11,7 @@ bank from its pinned donor: 166,321 persons, 57,240 households, 57,240 heads, 74,572 partner pointers, 62,096 `parent_1_person_id` pointers and 43,731 `parent_2_person_id` pointers, after dropping two parent edges. Concept validation reports no violations. Benefit-unit construction from the concept -pointers, which step 6 below encodes, is built for New Zealand. The composed +pointers, which step 6 below relies on, is built for New Zealand. The composed transport pipeline remains a design, and no build has migrated. **Date:** 2026-09-28. **Source:** Max's 27 September rulings on international populations and law-anchored concepts. **Amended 2026-10-07:** diff --git a/packages/microcosm-build/src/microcosm/build/transport/codecs.py b/packages/microcosm-build/src/microcosm/build/transport/codecs.py new file mode 100644 index 000000000..58b3beba9 --- /dev/null +++ b/packages/microcosm-build/src/microcosm/build/transport/codecs.py @@ -0,0 +1,129 @@ +"""Explicit codecs for donor populations, Chronicle facts and RuleSpec trees. + +The graph registry supports Frame and immutable-byte loaders. The donor codec +therefore returns a donor-scale concept Frame; CREATE installs the destination +schema and mass itself. Facts are validated by the shared Chronicle loader and +returned as their original JSONL bytes. RuleSpec trees have a canonical, +root-path-independent byte envelope containing every relative file and byte. +Importing this module does not mutate the shared codec registry. +""" + +from __future__ import annotations + +import base64 +from hashlib import file_digest +from pathlib import Path + +import pandas as pd + +from microcosm.build.ledger_artifact import load_ledger_consumer_artifact +from microcosm.frame import EntitySchema, Frame, WeightKind, Weights +from microcosm.frame.transport import read_populace_us_donor +from microcosm.graph.canonical import canonical_json +from microcosm.graph.codecs import SOURCE_CODECS, SourceCodecRegistry +from microcosm.graph.store import ContentStore + +__all__ = [ + "DONOR_SOURCE_CODEC", + "FACTS_SOURCE_CODEC", + "RULESPEC_SOURCE_CODEC", + "load_ledger_consumer_bytes", + "load_populace_us_h5", + "load_rulespec_tree_bytes", + "register_transport_codecs", +] + +DONOR_SOURCE_CODEC = "populace-us-h5-v1" +FACTS_SOURCE_CODEC = "ledger-consumer-artifact-v1" +RULESPEC_SOURCE_CODEC = "rulespec-tree-v1" + + +def load_populace_us_h5(path: Path, *, store: ContentStore | None = None) -> Frame: + """Read a donor at its original scale and currency. + + SourceRef identity checks belong to the executor. This loader computes the + wrapped reader's mandatory size/hash pins from the file itself, so it + authenticates nothing; the CREATE kernel passes its declared pins to the + reader directly, and the reader refuses a file that differs from them. + Source lineage stays donor provenance, never destination observations. + """ + + path = Path(path) + with path.open("rb") as source: + digest = file_digest(source, "sha256").hexdigest() + donor = read_populace_us_donor(path, sha256=digest, size=path.stat().st_size) + person = donor.tables["person"] + households = donor.tables["household"] + rows = pd.Index(households["household_id"]).get_indexer( + person["person_household_id"] + ) + return Frame( + donor.tables, + EntitySchema(group_entities=("household",)), + {"household": Weights(donor.weights, WeightKind.DESIGN)}, + pd.Series( + donor.support_strata[rows], + index=person.index, + dtype=object, + name="stratum", + ), + metadata={ + "content_basis": donor.content_basis.value, + "donor_country": donor.donor_country, + "currency": donor.currency, + "source_person_ids": donor.source_person_ids.tolist(), + "dropped_parent_cycle_edges": donor.dropped_parent_cycle_edges, + }, + ) + + +def load_ledger_consumer_bytes( + path: Path, *, store: ContentStore | None = None +) -> bytes: + """Validate a Chronicle artifact and return its unchanged fact-file bytes.""" + + path = Path(path) + load_ledger_consumer_artifact(path) + facts_path = path / "consumer_facts.jsonl" if path.is_dir() else path + return facts_path.read_bytes() + + +def load_rulespec_tree_bytes(path: Path, *, store: ContentStore | None = None) -> bytes: + """Pack a nonempty tree deterministically without its absolute root path. + + Every file is included, including auxiliary files. Symlinks are refused: + an input tree must contain its own bytes rather than borrow undeclared + resources outside the source root. + """ + + path = Path(path) + if not path.is_dir() or path.is_symlink(): + raise ValueError("A RuleSpec tree source must be a regular directory.") + entries = sorted( + path.rglob("*"), key=lambda entry: entry.relative_to(path).as_posix() + ) + if any(entry.is_symlink() for entry in entries): + raise ValueError("A RuleSpec tree source must not contain symlinks.") + files = [ + { + "path": entry.relative_to(path).as_posix(), + "bytes_base64": base64.b64encode(entry.read_bytes()).decode("ascii"), + } + for entry in entries + if entry.is_file() + ] + if not files: + raise ValueError("A RuleSpec tree source must contain at least one file.") + return canonical_json({"format": RULESPEC_SOURCE_CODEC, "files": files}) + + +def register_transport_codecs(registry: SourceCodecRegistry = SOURCE_CODECS) -> None: + """Register the three loaders explicitly and idempotently. + + The registry refuses an incumbent registered under a different loader or + mode; registering these same functions again changes nothing. + """ + + registry.register(DONOR_SOURCE_CODEC, load_populace_us_h5) + registry.register_bytes(FACTS_SOURCE_CODEC, load_ledger_consumer_bytes) + registry.register_bytes(RULESPEC_SOURCE_CODEC, load_rulespec_tree_bytes) diff --git a/packages/microcosm-build/src/microcosm/build/transport/column_kernels.py b/packages/microcosm-build/src/microcosm/build/transport/column_kernels.py new file mode 100644 index 000000000..088b829ae --- /dev/null +++ b/packages/microcosm-build/src/microcosm/build/transport/column_kernels.py @@ -0,0 +1,739 @@ +"""Column kernels for transported concepts, units, receipts and scenarios. + +Structured parameters are canonical JSON, each beside a ``*_sha256`` naming +the resource it came from. The kernels check only that digest's format, so +it moves the node key and authenticates nothing. The kernels never load a +country package. ``unit_rule`` is the resolved :class:`BenefitUnitRule` +representation, rather than an unresolved country proposal. ``mapping`` is a +:class:`ConceptMapping` representation; ``rulespec_paths`` selects its +bindings and the closure's defaults. + +The encode and scenario kernels derive no take-up draws: they refuse +``take_up_rates``, so the encoder refuses any selected take-up-threshold +binding it would execute, for want of a rate. ``takeup.assign@1`` owns the +keyed draws. + +Scenario overrides re-run the same group encoder with the closure's knobs. +This preserves the encoder's allocation and state-column semantics, including +household rent conservation, instead of implementing another set of transforms. +Only columns declared as encoded inputs or closure defaults may be owned. + +Receipt assignment reads a ``receipt_contract`` with ``seed_column``, ordered +``programs`` and ``exclusion_groups``. Each program names a ``program`` draw +key, ``output``, ``judgment_column`` and either ``rate`` or ``target`` (a +compiled person-count target in the declared ``surface`` artifact). A program +may require ``payment_column`` to be positive and may name existing boolean +``exclude_columns``. Program order is an explicit conflict priority. No law, +priority or take-up rate is inferred from names. A target sets the expected +take-up rate, target / eligible weighted mass; the realized weighted count is +recorded separately. All judgments must be determined. +""" + +from __future__ import annotations + +import math +from collections.abc import Mapping +from dataclasses import replace + +import numpy as np +import pandas as pd + +from microcosm.calibrate import hierarchy as hierarchy_module +from microcosm.calibrate import registry as registry_module +from microcosm.frame import concept_mapping as mapping_module +from microcosm.frame import concepts as concepts_module +from microcosm.frame import input_closure as closure_module +from microcosm.frame import unit_construction as units_module +from microcosm.frame.concept_mapping import ConceptMapping +from microcosm.frame.concepts import derive_take_up_draws +from microcosm.frame.input_closure import ( + InputClosure, + UndeterminedAction, + apply_defaults, + resolve_judgments, +) +from microcosm.frame.unit_construction import ( + BenefitUnitRule, + benefit_unit_attributes, + benefit_unit_membership, + units_per_household, +) +from microcosm.graph import ( + ROWS_ALL, + Capabilities, + Determinism, + KernelBase, + KernelContext, + KernelRegistry, + KernelResult, + Numeric, + SeedSource, + source_hash, +) +from microcosm.graph.population import dtype_for_token + +from . import graph_inputs +from .artifact_types import TARGET_SURFACE_TYPE +from .graph_inputs import ( + canonical_document_param, + require_params, + sha256_param, + string_param, +) +from .target_kernels import decode_target_surface + +__all__ = [ + "CONCEPTS_ENCODE", + "CONCEPTS_ENCODE_GROUPS", + "TAKEUP_ASSIGN", + "UNIT_ATTRIBUTES", + "SCENARIO_OVERRIDE", + "TRANSPORT_SCENARIO_OVERRIDE", + "TRANSPORT_UNIT_ATTRIBUTES", + "ConceptsEncodeGroupsKernel", + "ConceptsEncodeKernel", + "ScenarioOverrideKernel", + "TakeupAssignKernel", + "UnitAttributesKernel", + "assign_receipts", + "register_column_kernels", +] + + +def _whole_tables( + context: KernelContext, + ref: str, + entities: tuple[str, ...], + *, + artifact_aliases: frozenset[str] = frozenset(), +): + if any( + item.rows != ROWS_ALL for item in (*context.node.inputs, *context.node.outputs) + ): + raise ValueError(f"{ref} requires all-row slices.") + missing = set(entities) - set(context.tables) + if missing: + raise ValueError( + f"{ref} needs {sorted(missing)} tables; declare a data-column slice " + "on each entity." + ) + if context.node.sources or context.node.artifact_outputs: + raise ValueError(f"{ref} declares no sources or artifact outputs.") + if set(context.artifacts) - artifact_aliases: + raise ValueError(f"{ref} received undeclared artifact aliases.") + return {entity: context.tables[entity].copy(deep=True) for entity in entities} + + +def _document(context: KernelContext, ref: str, name: str): + sha256_param(context, ref, f"{name}_sha256") + return canonical_document_param(context, ref, name) + + +def _rule(context: KernelContext, ref: str) -> BenefitUnitRule: + return BenefitUnitRule.from_dict(_document(context, ref, "unit_rule")) + + +def _paths(context: KernelContext, ref: str) -> tuple[str, ...] | None: + value = context.params.get("rulespec_paths") + if value is None: + return None + if ( + not isinstance(value, tuple) + or not value + or any(not isinstance(path, str) or not path for path in value) + or len(set(value)) != len(value) + ): + raise ValueError(f"{ref} rulespec_paths must be distinct non-empty paths.") + return value + + +def _full_mapping(context: KernelContext, ref: str) -> ConceptMapping: + """The whole declared mapping, after checking its selected paths exist.""" + mapping = ConceptMapping.from_dict(_document(context, ref, "mapping")) + paths = _paths(context, ref) + if paths is not None: + unknown = set(paths) - {binding.module for binding in mapping.bindings} + if unknown: + raise ValueError(f"{ref} has unknown mapping paths {sorted(unknown)}.") + return mapping + + +def _mapping(context: KernelContext, ref: str) -> ConceptMapping: + """The mapping restricted to the selected paths, for person encoding.""" + mapping = _full_mapping(context, ref) + paths = _paths(context, ref) + if paths is None: + return mapping + bindings = tuple(binding for binding in mapping.bindings if binding.module in paths) + bound = {concept for binding in bindings for concept in binding.reads} + unmapped = { + item.id: mapping.unmapped.get(item.id, "Outside the selected RuleSpec paths.") + for item in concepts_module.CONCEPTS + if item.id not in bound + } + return replace(mapping, bindings=bindings, unmapped=unmapped) + + +def _closure(context: KernelContext, ref: str) -> InputClosure | None: + if "closure" not in context.params: + if "closure_sha256" in context.params: + raise ValueError(f"{ref} closure_sha256 needs a closure document.") + return None + closure = InputClosure.from_dict(_document(context, ref, "closure")) + mapping = ConceptMapping.from_dict(_document(context, ref, "mapping")) + if (closure.mapping_engine, closure.mapping_engine_version) != ( + mapping.engine, + mapping.engine_version, + ): + raise ValueError(f"{ref} closure and mapping identities differ.") + paths = _paths(context, ref) + if paths is None: + raise ValueError(f"{ref} needs rulespec_paths when a closure is supplied.") + unknown = set(paths) - set(closure.modules) + if unknown: + raise ValueError(f"{ref} closure has no paths {sorted(unknown)}.") + return closure + + +def _defaults( + context: KernelContext, + ref: str, + tables: Mapping[str, pd.DataFrame], + closure: InputClosure | None, + engine_entities: Mapping[str, str], +) -> dict[str, pd.DataFrame]: + out = dict(tables) + if closure is None: + return out + for frame_entity, engine_entity in engine_entities.items(): + defaults: dict[str, bool | int | float] = {} + for path in _paths(context, ref): + for name, value in closure.defaults(path, engine_entity).items(): + if name in defaults and ( + type(defaults[name]) is not type(value) or defaults[name] != value + ): + raise ValueError(f"{ref} has conflicting defaults for {name!r}.") + defaults[name] = value + out[frame_entity] = apply_defaults(out[frame_entity], defaults) + return out + + +def _equal_scalars(left: pd.Series, right: pd.Series) -> bool: + # Python scalar int/float equality preserves the integer's precision; + # numpy array comparison would promote an int64 to an inexact float64. + for before, after in zip(left.tolist(), right.tolist(), strict=True): + before_null, after_null = pd.isna(before), pd.isna(after) + if before_null or after_null: + if not (before_null and after_null): + return False + elif before != after: + return False + return True + + +def _columns(context: KernelContext, ref: str, tables: Mapping[str, pd.DataFrame]): + """Project computed values by id, casting only explicit graph dtypes.""" + columns = {} + if not context.node.outputs: + raise ValueError(f"{ref} must own at least one computed input column.") + for owned in context.node.outputs: + table = tables.get(owned.entity) + if table is None or owned.column not in table: + raise ValueError( + f"{ref} did not encode declared column {owned.entity}.{owned.column}." + ) + id_column = f"{owned.entity}_id" + if ( + owned.column == id_column + or owned.column.startswith("person_") + and owned.column.endswith("_id") + ): + raise ValueError(f"{ref} cannot return structural columns.") + values = table[owned.column].copy(deep=True) + # The graph's dtype for the token; its "string" pins Python storage, + # where a bare "string" would store pyarrow strings when installed. + dtype = None if owned.dtype is None else dtype_for_token(owned.dtype) + if dtype is not None and values.dtype != dtype: + # The group encoder uses numpy object for labels; the graph has + # an explicit string dtype. Integer casts must never truncate. + if ( + owned.dtype == "string" + and not values.dropna().map(lambda value: isinstance(value, str)).all() + ): + raise ValueError(f"{ref} refuses non-text values in {owned.column}.") + if ( + owned.dtype in {"bool", "boolean"} + and not values.dropna() + .map(lambda value: isinstance(value, bool | np.bool_)) + .all() + ): + raise ValueError(f"{ref} refuses non-boolean values in {owned.column}.") + cast = values.astype(dtype) + if owned.dtype != "string" and not _equal_scalars(values, cast): + raise ValueError(f"{ref} would lose values casting {owned.column}.") + values = cast + values.index = pd.Index(table[id_column].to_numpy(copy=True)) + columns[owned.entity, owned.column] = values + return columns + + +class _ColumnKernel(KernelBase): + capabilities = Capabilities( + Determinism.DETERMINISTIC, + numeric=Numeric.BITWISE, + seed_source=SeedSource.NONE, + dependencies=("numpy", "pandas"), + ) + _modules = () + + def implementation_hash(self) -> str: + return source_hash( + type(self), + graph_inputs, + *self._modules, + dependencies=self.capabilities.dependencies, + ) + + +class UnitAttributesKernel(_ColumnKernel): + """Attributes of the units CREATE already constructed; never new ids.""" + + ref = "transport.unit_attributes@1" + _modules = (units_module, mapping_module, concepts_module) + + 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"] + household_attributes = household.loc[:, ["household_id"]].copy() + household_attributes[f"n_{rule.entity}_units"] = units_per_household( + household, units, rule + ) + return KernelResult( + columns=_columns( + context, + self.ref, + {rule.entity: attributes, "household": household_attributes}, + ), + receipt={"entity": rule.entity, "n_units": len(units)}, + ) + + +_ENCODE_REQUIRED = frozenset({"mapping", "mapping_sha256"}) +_ENCODE_OPTIONAL = frozenset( + { + "rulespec_paths", + "shares", + "shares_sha256", + "closure", + "closure_sha256", + } +) +_TAKE_UP_PARAMS = frozenset({"take_up_rates", "take_up_rates_sha256"}) + + +def _refuse_take_up(context: KernelContext, ref: str) -> None: + """Keep keyed take-up draws in the receipt layer (``SeedSource.NONE``).""" + if _TAKE_UP_PARAMS & set(context.params): + raise ValueError( + f"{ref} derives no take-up draws; takeup.assign@1 owns them, so " + "take_up_rates are refused." + ) + + +def _optional_document(context: KernelContext, ref: str, name: str): + return ( + canonical_document_param(context, ref, name) if name in context.params else {} + ) + + +def _optional_resource_document(context: KernelContext, ref: str, name: str): + if name not in context.params: + if f"{name}_sha256" in context.params: + raise ValueError(f"{ref} {name}_sha256 needs a {name} document.") + return {} + return _document(context, ref, name) + + +class ConceptsEncodeKernel(_ColumnKernel): + """Encode person and household bindings plus declared closure defaults.""" + + ref = "concepts.encode@1" + _modules = (mapping_module, concepts_module, closure_module) + + def run(self, context: KernelContext) -> KernelResult: + _refuse_take_up(context, self.ref) + require_params( + context, self.ref, required=_ENCODE_REQUIRED, optional=_ENCODE_OPTIONAL + ) + tables = _whole_tables(context, self.ref, ("person", "household")) + mapping = _mapping(context, self.ref) + encoded = mapping.encode( + tables, + shares=_optional_resource_document(context, self.ref, "shares"), + ) + outputs = _defaults( + context, + self.ref, + encoded.tables, + _closure(context, self.ref), + mapping.entity_correspondence, + ) + return KernelResult( + columns=_columns(context, self.ref, outputs), + receipt={"engine": mapping.engine, "n_deferred": len(encoded.deferred)}, + ) + + +_GROUP_REQUIRED = _ENCODE_REQUIRED | frozenset( + {"unit_rule", "unit_rule_sha256", "engine_entity"} +) +_GROUP_OPTIONAL = _ENCODE_OPTIONAL | frozenset({"knobs", "scenario_sha256"}) + + +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) + closure = _closure(context, ref) + knobs = _optional_document(context, ref, "knobs") + if knobs and closure is None: + raise ValueError(f"{ref} requires a closure for scenario knobs.") + if knobs: + 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], + tables["person"][rule.membership_column], + rule, + ) + encoded = mapping.encode_groups( + tables, + {entity: membership}, + modules=_paths(context, ref), + knobs=None if closure is None else closure.group_knobs(knobs), + state_bindings=() if closure is None else closure.state_bindings(), + shares=_optional_resource_document(context, ref, "shares"), + ) + outputs = _defaults(context, ref, encoded.tables, closure, {rule.entity: entity}) + return mapping, encoded, outputs + + +class ConceptsEncodeGroupsKernel(_ColumnKernel): + """Group bindings execute through G4's encoder and built memberships.""" + + ref = "concepts.encode_groups@1" + _modules = (mapping_module, units_module, concepts_module, closure_module) + + def run(self, context: KernelContext) -> KernelResult: + _refuse_take_up(context, self.ref) + require_params( + context, self.ref, required=_GROUP_REQUIRED, optional=_GROUP_OPTIONAL + ) + mapping, encoded, outputs = _group_inputs(context, self.ref) + return KernelResult( + columns=_columns(context, self.ref, outputs), + receipt={"engine": mapping.engine, "n_deferred": len(encoded.deferred)}, + ) + + +class ScenarioOverrideKernel(ConceptsEncodeGroupsKernel): + """Re-encode declared group inputs under the closure's scenario knobs.""" + + ref = "transport.scenario_override@1" + + def run(self, context: KernelContext) -> KernelResult: + _refuse_take_up(context, self.ref) + require_params( + context, + self.ref, + required=_GROUP_REQUIRED + | frozenset({"closure", "closure_sha256", "knobs", "scenario_sha256"}), + optional=_GROUP_OPTIONAL, + ) + if any(not owned.rewrite for owned in context.node.outputs): + raise ValueError(f"{self.ref} owns rewrites of encoded inputs only.") + mapping, encoded, outputs = _group_inputs(context, self.ref) + return KernelResult( + columns=_columns(context, self.ref, outputs), + receipt={"engine": mapping.engine, "n_deferred": len(encoded.deferred)}, + ) + + +def _text(value: object, name: str) -> str: + if not isinstance(value, str) or not value: + raise ValueError(f"Receipt {name} must be non-empty text.") + return value + + +def _finite_nonnegative(value: object, name: str) -> float: + if isinstance(value, bool) or not isinstance(value, int | float): + raise ValueError(f"Receipt {name} must be a number.") + value = float(value) + if not math.isfinite(value) or value < 0: + raise ValueError(f"Receipt {name} must be finite and non-negative.") + return value + + +def assign_receipts( + person: pd.DataFrame, + contract: Mapping[str, object], + *, + weights: np.ndarray | None = None, + targets: Mapping[str, float] | None = None, +) -> tuple[pd.DataFrame, dict[str, object]]: + """Apply explicit receipt priorities to determined eligible program draws. + + This shared pure operation contains receipt orchestration only. G4's + judgment decoder and the content layer's stable program draws supply its + semantics. Existing exclusions are read as booleans and never rewritten. + """ + if set(contract) != {"seed_column", "programs", "exclusion_groups"}: + raise ValueError( + "Receipt contract declares seed_column, programs and exclusion_groups." + ) + seed_column = _text(contract["seed_column"], "seed_column") + programs = contract["programs"] + groups = contract["exclusion_groups"] + if not isinstance(programs, list | tuple) or not programs: + raise ValueError("Receipt programs must be a non-empty ordered list.") + if not isinstance(groups, list | tuple): + raise ValueError("Receipt exclusion_groups must be a list.") + parsed = [] + for row in programs: + if not isinstance(row, Mapping): + raise ValueError("A receipt program must be an object.") + required = {"program", "output", "judgment_column"} + optional = { + "rate", + "target", + "target_measure", + "payment_column", + "exclude_columns", + } + if not required <= set(row) or set(row) - required - optional: + raise ValueError("A receipt program has missing or unknown fields.") + if ("rate" in row) == ("target" in row): + raise ValueError("A receipt program declares exactly one rate or target.") + if "target_measure" in row and "target" not in row: + raise ValueError("A target_measure requires a target reference.") + parsed.append({**row, **{key: _text(row[key], key) for key in required}}) + outputs = [row["output"] for row in parsed] + keys = [row["program"] for row in parsed] + if len(set(outputs)) != len(outputs) or len(set(keys)) != len(keys): + raise ValueError("Receipt program keys and output names must be distinct.") + for group in groups: + if ( + not isinstance(group, list | tuple) + or len(group) < 2 + or len(set(group)) != len(group) + or not set(group) <= set(outputs) + ): + raise ValueError("Each receipt exclusion group names distinct outputs.") + out = person.loc[:, ["person_id"]].copy() + audit = {} + if weights is not None: + weights = np.asarray(weights, dtype=np.float64) + if ( + weights.shape != (len(person),) + or not np.isfinite(weights).all() + or (weights < 0).any() + ): + raise ValueError( + "Receipt weights must be aligned, finite and non-negative." + ) + for row in parsed: + output = row["output"] + eligible = resolve_judgments( + person[row["judgment_column"]].to_numpy(), UndeterminedAction.REFUSE + ) + payment_column = row.get("payment_column") + if payment_column is not None: + payment = person[_text(payment_column, "payment_column")].to_numpy( + dtype=np.float64 + ) + if not np.isfinite(payment).all(): + raise ValueError("Receipt payment values must be finite.") + eligible &= payment > 0 + excludes = row.get("exclude_columns", []) + if not isinstance(excludes, list | tuple) or len(set(excludes)) != len( + excludes + ): + raise ValueError("Receipt exclude_columns must be distinct columns.") + for column in excludes: + values = person[_text(column, "exclude column")] + if str(values.dtype) not in {"bool", "boolean"} or values.isna().any(): + raise ValueError( + "Existing receipt exclusions must be non-null booleans." + ) + eligible &= ~values.to_numpy(dtype=bool) + for group in groups: + if output in group: + for previous in group: + if previous in out: + eligible &= ~out[previous].to_numpy() + target = None + mass = None if weights is None else math.fsum(weights[eligible].tolist()) + if "target" in row: + name = _text(row["target"], "target") + if targets is None or name not in targets or weights is None: + raise ValueError( + f"Receipt target {name!r} needs a surface and weights." + ) + target = _finite_nonnegative(targets[name], "target") + if target > mass: + raise ValueError( + f"Receipt target {name!r} exceeds eligible weighted mass." + ) + # target <= mass, so a zero eligible mass has a zero target. + rate = target / mass if mass else target + else: + rate = _finite_nonnegative(row["rate"], "rate") + if rate > 1: + raise ValueError("Receipt rate must lie on [0, 1].") + draws = derive_take_up_draws(person[seed_column].to_numpy(), row["program"]) + assigned = eligible & (draws < rate) + out[output] = assigned + audit[output] = { + "program": row["program"], + "rate": rate, + "target": target, + "eligible_mass": mass, + "n_receipts": int(assigned.sum()), + "realized_mass": None + if weights is None + else math.fsum(weights[assigned].tolist()), + } + return out, audit + + +class TakeupAssignKernel(_ColumnKernel): + """Receipt assignment over deterministic stable program-key draws.""" + + ref = "takeup.assign@1" + _modules = (concepts_module, closure_module) + capabilities = Capabilities( + Determinism.DETERMINISTIC, + numeric=Numeric.BITWISE, + seed_source=SeedSource.KEYED, + dependencies=("numpy", "pandas"), + ) + + def implementation_hash(self) -> str: + from . import target_kernels + + return source_hash( + type(self), + graph_inputs, + concepts_module, + closure_module, + target_kernels, + registry_module, + hierarchy_module, + dependencies=self.capabilities.dependencies, + ) + + def run(self, context: KernelContext) -> KernelResult: + require_params( + context, + self.ref, + required=frozenset({"receipt_contract", "receipt_contract_sha256"}), + ) + tables = _whole_tables( + context, self.ref, ("person",), artifact_aliases=frozenset({"surface"}) + ) + contract = _document(context, self.ref, "receipt_contract") + targets = None + if context.artifacts: + if ( + set(context.artifacts) != {"surface"} + or context.artifacts["surface"].type != TARGET_SURFACE_TYPE + ): + raise ValueError(f"{self.ref} reads only the typed surface artifact.") + surface = decode_target_surface(context.artifacts["surface"].payload) + targets = {} + programs = contract.get("programs", ()) + for row in programs: + if not isinstance(row, Mapping) or "target" not in row: + continue + name = _text(row["target"], "target") + expected_measure = _text( + row.get("target_measure", row.get("output")), "target_measure" + ) + matches = [ + target for target in surface.registry.specs if target.name == name + ] + if len(matches) != 1: + raise ValueError( + f"Receipt target {name!r} must identify one surface row." + ) + (target,) = matches + if ( + target.entity != "person" + or target.filter is not None + or target.measure != expected_measure + or target.metadata.get("ledger_measure_unit") != "count" + ): + raise ValueError( + f"Receipt target {name!r} must be an unfiltered person count of {expected_measure!r}." + ) + targets[name] = target.value + table, audit = assign_receipts( + tables["person"], + contract, + weights=None + if "person" not in context.weights + else context.weights["person"].values, + targets=targets, + ) + expected = set(table.columns) - {"person_id"} + owned = {(item.entity, item.column) for item in context.node.outputs} + if owned != {("person", column) for column in expected}: + raise ValueError( + f"{self.ref} must own exactly the contract's receipt flags." + ) + return KernelResult( + columns=_columns(context, self.ref, {"person": table}), + receipt={"programs": audit}, + ) + + +TRANSPORT_UNIT_ATTRIBUTES = UnitAttributesKernel() +CONCEPTS_ENCODE = ConceptsEncodeKernel() +CONCEPTS_ENCODE_GROUPS = ConceptsEncodeGroupsKernel() +TAKEUP_ASSIGN = TakeupAssignKernel() +TRANSPORT_SCENARIO_OVERRIDE = ScenarioOverrideKernel() +UNIT_ATTRIBUTES = TRANSPORT_UNIT_ATTRIBUTES +SCENARIO_OVERRIDE = TRANSPORT_SCENARIO_OVERRIDE + + +def register_column_kernels(registry: KernelRegistry) -> KernelRegistry: + """Register one instance of each column kernel in the run's registry.""" + for kernel in ( + TRANSPORT_UNIT_ATTRIBUTES, + CONCEPTS_ENCODE, + CONCEPTS_ENCODE_GROUPS, + TAKEUP_ASSIGN, + TRANSPORT_SCENARIO_OVERRIDE, + ): + registry.register(kernel) + return registry diff --git a/packages/microcosm-build/src/microcosm/build/transport/geography_kernels.py b/packages/microcosm-build/src/microcosm/build/transport/geography_kernels.py new file mode 100644 index 000000000..759957a82 --- /dev/null +++ b/packages/microcosm-build/src/microcosm/build/transport/geography_kernels.py @@ -0,0 +1,336 @@ +"""Compile atomic geography support from a precalibration target surface. + +The resolved definition is spec data: atomic-area codes, each row's compiled +population target, population shares, and the existing atomic-support column +metadata. Integer population is apportioned within each target with exact +largest remainders, conserving its total. No geography mapping or population +value is inferred from country names. Empty, incomplete or ambiguous support +refuses execution, including an unharvested crosswalk. +""" + +from __future__ import annotations + +import math +import sys +from collections.abc import Mapping +from fractions import Fraction + +import numpy as np + +import microcosm.build.atomic_geography as atomic_geography_module +import microcosm.calibrate.hierarchy as hierarchy_module +import microcosm.calibrate.registry as registry_module +from microcosm.build.atomic_geography import _U53, encode_atomic_support +from microcosm.build.graph_atomic_geography import ATOMIC_SUPPORT_TYPE +from microcosm.graph import ( + Capabilities, + Determinism, + KernelBase, + KernelContext, + KernelRegistry, + KernelResult, + Numeric, + SeedSource, + source_hash, +) +from microcosm.graph.canonical import canonical_json + +from . import graph_inputs, target_kernels +from .artifact_types import TARGET_SURFACE_TYPE +from .graph_inputs import ( + canonical_document_param, + require_params, + sha256_param, + sha256_text, + string_param, +) +from .target_kernels import TargetSurface, decode_target_surface + +__all__ = [ + "ATOMIC_SUPPORT_TYPE", + "GEOGRAPHY_SUPPORT_FROM_FACTS", + "GeographySupportFromFactsKernel", + "support_from_surface", + "register_geography_kernels", +] + +_DEFINITION_KEYS = frozenset( + { + "version", + "level", + "code_system", + "vintage", + "columns", + "rows", + "share_tolerance", + "population_geography_level", + "population_code_column", + } +) + + +def support_from_surface( + surface: TargetSurface, definition: Mapping[str, object], *, system: str +) -> bytes: + """Resolve declared support rows through already compiled population facts. + + Each row is ``{target, codes, population_share}``; ``codes`` contains + exactly the definition's code columns, including a unique ``area``. + The population column is the sole weight column and follows the existing + atomic-support contract. A target's hierarchy geographic id must equal + the declared population code. Population targets carry + ``precal_use=atomic_geography_support`` and ``ledger_measure_unit=count`` + and have integral values. + + Largest-remainder apportionment rounds shares to integer sampling counts, + preserving each published target exactly; equal remainders use ascending + atomic-area code. Sorting those codes also makes row permutation inert. + """ + + if ( + set(definition) != _DEFINITION_KEYS + or type(definition["version"]) is not int + or definition["version"] != 1 + ): + raise ValueError("Atomic support definition has an unsupported shape.") + columns = definition["columns"] + if not isinstance(columns, Mapping): + raise ValueError("Atomic support definition requires column metadata.") + code_columns = { + name + for name, description in columns.items() + if isinstance(description, Mapping) and description.get("kind") == "code" + } + weight_columns = { + name + for name, description in columns.items() + if isinstance(description, Mapping) and description.get("kind") == "weight" + } + if ( + "area" not in code_columns + or weight_columns != {"population"} + or set(columns) != code_columns | weight_columns + ): + raise ValueError( + "Atomic support declares code columns and one population weight column." + ) + population_column = definition["population_code_column"] + if population_column not in code_columns: + raise ValueError("Atomic support population code column must be a code column.") + tolerance = definition["share_tolerance"] + if ( + isinstance(tolerance, bool) + or not isinstance(tolerance, (int, float)) + or not math.isfinite(tolerance) + or tolerance < 0 + ): + raise ValueError( + "Atomic support share tolerance must be finite and nonnegative." + ) + rows = definition["rows"] + if not isinstance(rows, list) or not rows: + raise ValueError("Atomic support requires harvested nonempty rows.") + groups: dict[str, list[tuple[Mapping[str, str], Fraction]]] = {} + seen = set() + for row in rows: + if not isinstance(row, Mapping) or set(row) != { + "target", + "codes", + "population_share", + }: + raise ValueError( + "Atomic support row requires target, codes and population_share." + ) + target, codes, share = row["target"], row["codes"], row["population_share"] + if ( + not isinstance(target, str) + or not target + or not isinstance(codes, Mapping) + or set(codes) != code_columns + ): + raise ValueError("Atomic support row target or code declaration differs.") + if any( + not isinstance(code, str) or not code or code.strip() != code + for code in codes.values() + ): + raise ValueError( + "Atomic support codes must be nonempty strings without whitespace." + ) + if codes["area"] in seen: + raise ValueError("Atomic support area codes must be unique.") + seen.add(codes["area"]) + if ( + isinstance(share, bool) + or not isinstance(share, (int, float)) + or not math.isfinite(share) + or not 0 <= share <= 1 + ): + raise ValueError( + "Atomic support shares must be finite fractions in [0, 1]." + ) + groups.setdefault(target, []).append((codes, Fraction(str(share)))) + specs = {spec.name: spec for spec in surface.registry.specs} + if len(specs) != len(surface.registry.specs): + raise ValueError( + "Atomic support requires unambiguous target names across periods." + ) + eligible = { + spec.name + for spec in surface.registry.specs + if spec.metadata.get("precal_use") == "atomic_geography_support" + } + if set(groups) != eligible: + raise ValueError( + "Atomic support rows must cover exactly the surface's geography population targets." + ) + apportioned = [] + for target, members in sorted(groups.items()): + spec = specs[target] + if spec.metadata.get("ledger_measure_unit") != "count": + raise ValueError( + "Atomic support population facts must declare count units." + ) + if spec.value < 0 or not spec.value.is_integer(): + raise ValueError( + "Atomic support population facts must be nonnegative integers." + ) + if ( + spec.hierarchy is None + or spec.hierarchy.geography.level + != definition["population_geography_level"] + ): + raise ValueError( + "Atomic support population fact has the wrong geographic level." + ) + if any( + codes[population_column] != spec.hierarchy.geography.id + for codes, _ in members + ): + raise ValueError( + "Atomic support row codes differ from the population fact geography." + ) + members.sort(key=lambda member: member[0]["area"]) + total_share = sum((share for _, share in members), Fraction()) + if total_share <= 0 or abs(total_share - 1) > Fraction(str(tolerance)): + raise ValueError( + "Atomic support shares for each population target must sum to one." + ) + population = int(spec.value) + quotas = [population * share / total_share for _, share in members] + counts = [quota.numerator // quota.denominator for quota in quotas] + remainder = population - sum(counts) + ranked = sorted( + range(len(members)), + key=lambda i: (-(quotas[i] - counts[i]), members[i][0]["area"]), + ) + for index in ranked[:remainder]: + counts[index] += 1 + apportioned.extend( + (codes, count) for (codes, _), count in zip(members, counts, strict=True) + ) + apportioned.sort(key=lambda item: item[0]["area"]) + # Validate the exact sampling bound before converting Python integers to + # int64, avoiding a wrap or a numpy conversion error on excessive counts. + if not 0 < sum(count for _, count in apportioned) <= _U53: + raise ValueError( + "Atomic support population total is outside the exact sampling range." + ) + arrays = { + name: np.asarray([codes[name] for codes, _ in apportioned], dtype=str) + for name in sorted(code_columns) + } + arrays["population"] = np.asarray( + [count for _, count in apportioned], dtype=np.int64 + ) + return encode_atomic_support( + { + "version": 1, + "system": system, + "level": definition["level"], + "code_system": definition["code_system"], + "vintage": definition["vintage"], + "columns": dict(columns), + }, + arrays, + ) + + +class GeographySupportFromFactsKernel(KernelBase): + """``geography.support_from_facts@1`` over a declared target surface. + + The support payload is a deflated NPZ archive. CPython's zip writer + records the creating platform in each entry (``ZipInfo.create_system`` + is 0 on Windows and 3 elsewhere), so the bytes can differ between + platforms and the kernel declares the atomic geography kernels' + ``platform_bitwise`` numeric class. ``definition_sha256`` is + format-checked and recorded: it moves the node key and authenticates + nothing. + """ + + ref = "geography.support_from_facts@1" + capabilities = Capabilities( + Determinism.DETERMINISTIC, + numeric=Numeric.PLATFORM_BITWISE, + seed_source=SeedSource.NONE, + dependencies=("numpy",), + consumes_se=False, + ) + + def implementation_hash(self) -> str: + return source_hash( + sys.modules[__name__], + atomic_geography_module, + target_kernels, + graph_inputs, + registry_module, + hierarchy_module, + canonical_json, + dependencies=self.capabilities.dependencies, + ) + + def run(self, context: KernelContext) -> KernelResult: + require_params( + context, + self.ref, + required=frozenset({"definition", "definition_sha256", "system"}), + ) + digest = sha256_param(context, self.ref, "definition_sha256") + definition = canonical_document_param(context, self.ref, "definition") + system = string_param(context, self.ref, "system") + declared = { + output.name: output.type for output in context.node.artifact_outputs + } + if declared != {"support": ATOMIC_SUPPORT_TYPE}: + raise ValueError(f"{self.ref} declares exactly its typed support output.") + value = context.artifacts.get("surface") + if ( + set(context.artifacts) != {"surface"} + or value is None + or value.type != TARGET_SURFACE_TYPE + ): + raise ValueError( + f"{self.ref} reads one typed target surface under alias 'surface'." + ) + if context.node.sources: + raise ValueError(f"{self.ref} reads an artifact and declares no sources.") + surface = decode_target_surface(value.payload) + payload = support_from_surface(surface, definition, system=system) + return KernelResult( + artifacts={"support": payload}, + receipt={ + "system": system, + "surface_sha256": surface.sha256, + "definition_sha256": digest, + "support_sha256": sha256_text(payload), + "areas": len(definition["rows"]), + }, + ) + + +GEOGRAPHY_SUPPORT_FROM_FACTS = GeographySupportFromFactsKernel() + + +def register_geography_kernels(registry: KernelRegistry) -> None: + """Register the population-facts support adapter explicitly.""" + + registry.register(GEOGRAPHY_SUPPORT_FROM_FACTS) diff --git a/packages/microcosm-build/src/microcosm/build/transport/population_kernels.py b/packages/microcosm-build/src/microcosm/build/transport/population_kernels.py new file mode 100644 index 000000000..0c32524e7 --- /dev/null +++ b/packages/microcosm-build/src/microcosm/build/transport/population_kernels.py @@ -0,0 +1,589 @@ +"""Authenticated donor CREATE and declared transport rewrites. + +Country data enters only through node parameters and declared facts. The +CREATE installs destination-scale design weights before the graph captures +its design anchor; the records retain their donor support provenance. + +Four CREATE parameters are checked against content: the donor reader +refuses a file whose byte size or SHA-256 differs from ``donor_size`` and +``donor_sha256``, the Chronicle loader refuses facts whose digest differs from +an optional ``facts_sha256``, and ``concept_schema_sha256`` must equal the +reader's concept schema. Every other ``*_sha256`` parameter here +(``donor_pin_sha256``, ``unit_rule_sha256``, ``mass_reference_sha256``, +``currency_sha256`` and ``bands_sha256``) is checked for format only: it +moves the node key and authenticates nothing. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from math import fsum + +import numpy as np +import pandas as pd + +import microcosm.frame.adapters.policyengine_us_concepts as donor_mapping_module +import microcosm.frame.bundle as bundle_module +import microcosm.frame.concept_mapping as mapping_module +import microcosm.frame.concepts as concepts_module +import microcosm.frame.materialize as materialize_module +import microcosm.frame.scaling as scaling_module +import microcosm.frame.schema as schema_module +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.concepts import ( + concept_for_column, + concept_schema_sha256, + split_for_transport, +) +from microcosm.frame.transport import ( + currency_bridge, + derive_transport_seed, + quantile_map, + read_populace_us_donor, +) +from microcosm.frame.unit_construction import BenefitUnitRule, build_benefit_units +from microcosm.graph import ( + ROWS_ALL, + Capabilities, + Determinism, + KernelBase, + KernelContext, + KernelRegistry, + KernelResult, + Numeric, + SeedSource, + StructuralDelta, + source_hash, +) +from microcosm.graph.population import dtype_for_token + +from . import graph_inputs, target_kernels +from .artifact_types import TARGET_SURFACE_TYPE +from .graph_inputs import ( + canonical_document_param, + require_params, + sha256_param, + string_param, +) +from .target_kernels import compile_target_surface, decode_target_surface + +__all__ = [ + "TRANSPORT_BOUNDARY", + "TRANSPORT_CREATE", + "TRANSPORT_CURRENCY", + "TRANSPORT_QUANTILE_MAP", + "TransportBoundaryKernel", + "TransportCreateKernel", + "TransportCurrencyKernel", + "TransportQuantileMapKernel", + "register_population_kernels", +] + + +def _whole_rows(context: KernelContext, ref: str) -> None: + if any(item.rows != ROWS_ALL for item in context.node.inputs): + raise ValueError(f"{ref} requires whole-row slices.") + + +def _column_result( + context: KernelContext, ref: str, values: Mapping[tuple[str, str], pd.Series] +) -> KernelResult: + declared = {(item.entity, item.column) for item in context.node.outputs} + if declared != set(values): + raise ValueError(f"{ref} outputs must equal its declared rewritten columns.") + columns = {} + for (entity, name), series in values.items(): + table = context.tables[entity] + ids = table[f"{entity}_id"] + columns[entity, name] = pd.Series( + series.to_numpy(copy=True), + index=pd.Index(ids.to_numpy(copy=True), name=f"{entity}_id"), + name=name, + dtype=series.dtype, + ) + return KernelResult(columns=columns) + + +class TransportCreateKernel(KernelBase): + """``transport.create@1``: verified support records at destination mass.""" + + ref = "transport.create@1" + capabilities = Capabilities( + Determinism.DETERMINISTIC, + numeric=Numeric.BITWISE, + seed_source=SeedSource.KEYED, + structural=StructuralDelta.CREATE, + dependencies=("numpy", "pandas", "tables"), + ) + _required = frozenset( + { + "donor_source", + "facts_source", + "donor_sha256", + "donor_size", + "donor_pin_sha256", + "unit_rule", + "unit_rule_sha256", + "mass_reference", + "mass_reference_sha256", + "country", + "concept_schema_sha256", + "seed_stream", + } + ) + + def implementation_hash(self) -> str: + return source_hash( + type(self), + graph_inputs, + target_kernels, + transport_module, + units_module, + concepts_module, + mapping_module, + donor_mapping_module, + materialize_module, + scaling_module, + bundle_module, + schema_module, + weights_module, + *target_kernels._COMPILE_MODULES, + dependencies=self.capabilities.dependencies, + ) + + def run(self, context: KernelContext) -> KernelResult: + require_params( + context, + self.ref, + required=self._required, + optional=frozenset({"facts_sha256"}), + ) + donor_source = string_param(context, self.ref, "donor_source") + facts_source = string_param(context, self.ref, "facts_source") + if ( + donor_source == facts_source + or set(context.node.sources) != {donor_source, facts_source} + or any(name not in context.sources for name in context.node.sources) + ): + raise ValueError(f"{self.ref} reads exactly its donor and facts sources.") + for name in ("donor_pin_sha256", "unit_rule_sha256", "concept_schema_sha256"): + sha256_param(context, self.ref, name) + if context.params["concept_schema_sha256"] != concept_schema_sha256(): + raise ValueError("CREATE concept schema SHA-256 differs from this reader.") + country = string_param(context, self.ref, "country") + surface_bytes, registry = compile_target_surface( + context.sources[facts_source], + canonical_document_param(context, self.ref, "mass_reference"), + country=country, + references_sha256=sha256_param(context, self.ref, "mass_reference_sha256"), + expected_facts_sha256=( + sha256_param(context, self.ref, "facts_sha256") + if "facts_sha256" in context.params + else None + ), + ) + if len(registry.specs) != 1: + raise ValueError( + "CREATE requires exactly one destination population reference." + ) + mass_spec = registry.specs[0] + if ( + mass_spec.entity != "person" + or mass_spec.filter + or mass_spec.metadata.get("ledger_measure_unit") != "count" + ): + raise ValueError( + "CREATE destination population reference must be an unfiltered person count fact." + ) + mass = float(mass_spec.value) + if not np.isfinite(mass) or mass <= 0: + raise ValueError("Destination population mass must be finite and positive.") + donor = read_populace_us_donor( + context.sources[donor_source], + sha256=sha256_param(context, self.ref, "donor_sha256"), + size=context.params["donor_size"], + ) + tables, dropped = split_for_transport(donor.tables) + # pandas 3's inferred text dtype spells itself "str"; graph columns + # have an explicit "string" token. Keep the concept values intact + # while fixing the portable dtype at this boundary, with the graph's + # pinned Python storage: a bare "string" stores pyarrow strings + # whenever pyarrow is installed. + text = dtype_for_token("string") + for entity, table in tables.items(): + for column in table: + item = concept_for_column(entity, column) + if item is not None and item.dtype == "str": + table[column] = table[column].astype(text) + rule = BenefitUnitRule.from_dict( + 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 + ) + tables["family"] = family + person, household = tables["person"], tables["household"] + person[rule.membership_column] = membership + person["take_up_seed"] = derive_transport_seed( + donor.source_person_ids, string_param(context, self.ref, "seed_stream") + ) + household["donor_support_stratum"] = pd.array(donor.support_strata, dtype=text) + household_ids = household["household_id"] + donor_weights = pd.Series(donor.weights, index=household_ids) + person_weights = person["person_household_id"].map(donor_weights) + donor_mass = fsum(person_weights.to_numpy(dtype=np.float64)) + if not np.isfinite(donor_mass) or donor_mass <= 0: + raise ValueError("Donor person mass must be finite and positive.") + scale = mass / donor_mass + weights = Weights(donor.weights * scale, WeightKind.DESIGN) + strata = pd.Series( + person["person_household_id"] + .map(pd.Series(donor.support_strata, index=household_ids)) + .to_numpy(copy=True), + index=person.index, + name="stratum", + # Object storage, as the donor codec stores strata: an inferred + # pandas string dtype would follow whether pyarrow is installed. + 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)), + {"household": weights}, + strata, + metadata={ + "content_basis": donor.content_basis.value, + "donor_country": donor.donor_country, + "country": country, + "donor_currency": donor.currency, + }, + ) + return KernelResult( + frame=frame, + receipt={ + "content_basis": donor.content_basis.value, + "donor_country": donor.donor_country, + "country": country, + "dropped_columns": { + entity: list(names) for entity, names in dropped.items() + }, + "dropped_parent_cycle_edges": donor.dropped_parent_cycle_edges, + "donor_mass": donor_mass, + "destination_mass": mass, + "design_scale": scale, + "n_persons": len(person), + "n_households": len(household), + "n_units": len(family), + "mass_reference": registry.specs[0].name, + "facts": dict(surface.facts), + }, + ) + + +class TransportBoundaryKernel(KernelBase): + """``transport.boundary@1``: keep every person and conserve mass.""" + + ref = "transport.boundary@1" + capabilities = Capabilities( + Determinism.DETERMINISTIC, + structural=StructuralDelta.FILTER, + dependencies=("numpy", "pandas"), + ) + + def implementation_hash(self) -> str: + return source_hash( + type(self), + _whole_rows, + graph_inputs, + dependencies=self.capabilities.dependencies, + ) + + def run(self, context: KernelContext) -> KernelResult: + require_params(context, self.ref, required=frozenset()) + _whole_rows(context, self.ref) + if not any(item.entity == "person" for item in context.node.inputs): + raise ValueError( + "Boundary FILTER requires a declared person data-column slice." + ) + if context.node.mass != "conserve": + raise ValueError("Boundary FILTER requires mass='conserve'.") + ids = context.tables["person"]["person_id"] + return KernelResult( + keep=pd.Series(True, index=pd.Index(ids.to_numpy(), name="person_id")), + receipt={"kept_persons": len(ids)}, + ) + + +class TransportCurrencyKernel(KernelBase): + """``transport.currency@1``: one declared currency multiplication.""" + + ref = "transport.currency@1" + capabilities = Capabilities( + Determinism.DETERMINISTIC, + dependencies=("numpy", "pandas"), + ) + + def implementation_hash(self) -> str: + return source_hash( + type(self), + _whole_rows, + _column_result, + graph_inputs, + transport_module, + concepts_module, + scaling_module, + dependencies=self.capabilities.dependencies, + ) + + def run(self, context: KernelContext) -> KernelResult: + require_params( + context, + self.ref, + required=frozenset({"columns", "rate", "currency_sha256"}), + ) + _whole_rows(context, self.ref) + sha256_param(context, self.ref, "currency_sha256") + selected = canonical_document_param(context, self.ref, "columns") + if not selected or any( + not isinstance(names, list) + or not names + or any(not isinstance(name, str) or not name for name in names) + for names in selected.values() + ): + raise ValueError( + "Currency columns must name nonempty lists of money concepts." + ) + tables = currency_bridge(context.tables, selected, context.params["rate"]) + return _column_result( + context, + self.ref, + { + (entity, name): tables[entity][name] + for entity, names in selected.items() + for name in names + }, + ) + + +def _surface_bands(document: Mapping[str, object], surface, *, unit: str) -> object: + """Resolve band shares from the exact declared population facts. + + ``unit`` is the entity whose values are ranked. Each band reference must + be an unfiltered count fact on that entity, so the shares divide the + same population the ranks do. + """ + specs = {spec.name: spec for spec in surface.registry.specs} + + def rows(raw): + if not isinstance(raw, list) or not raw: + raise ValueError("Quantile bands require a nonempty list.") + if any( + not isinstance(row, dict) or set(row) != {"lower", "upper", "reference"} + for row in raw + ): + raise ValueError("Each quantile band declares lower, upper and reference.") + names = [row["reference"] for row in raw] + if any(not isinstance(name, str) or name not in specs for name in names): + raise ValueError("A quantile band's reference is absent from its surface.") + if len(set(names)) != len(names): + raise ValueError("Quantile band references must be distinct.") + if any( + specs[name].entity != unit + or specs[name].filter + or specs[name].metadata.get("ledger_measure_unit") != "count" + for name in names + ): + raise ValueError( + f"Quantile band references must be unfiltered {unit} count facts." + ) + counts = [float(specs[name].value) for name in names] + if any(not np.isfinite(value) or value < 0 for value in counts): + raise ValueError( + "Quantile band population counts must be finite and nonnegative." + ) + total = fsum(counts) + if total <= 0 or not np.isfinite(total): + raise ValueError( + "Quantile band population total must be positive and finite." + ) + return [ + {"lower": row["lower"], "upper": row["upper"], "share": count / total} + for row, count in zip(raw, counts, strict=True) + ] + + if set(document) == {"bands"}: + return rows(document["bands"]) + if set(document) == {"positive", "negative"}: + return {sign: rows(raw) for sign, raw in document.items()} + if set(document) == {"components"}: + components = document["components"] + if not isinstance(components, list) or not components: + raise ValueError("Component bands require a nonempty list.") + result = {} + for component in components: + if not isinstance(component, dict) or set(component) != {"value", "bands"}: + raise ValueError("Each component declares value and bands.") + value = component["value"] + if ( + not isinstance(value, (str, int)) + or isinstance(value, bool) + or value in result + ): + raise ValueError( + "Band component values must be distinct string or integer codes." + ) + result[value] = rows(component["bands"]) + return {"components": result} + raise ValueError("Quantile bands declare bands or positive and negative rows.") + + +class TransportQuantileMapKernel(KernelBase): + """``transport.quantile_map@1``: ranks mapped to compiled fact bands. + + ``aggregate_entity`` requests household totals and pro rata member + allocation. ``component_column`` requests independent distributions + within a sliced column (for example a household region). + """ + + ref = "transport.quantile_map@1" + capabilities = Capabilities( + Determinism.DETERMINISTIC, + dependencies=("numpy", "pandas"), + ) + + def implementation_hash(self) -> str: + return source_hash( + type(self), + _surface_bands, + _whole_rows, + _column_result, + graph_inputs, + target_kernels, + transport_module, + *target_kernels._COMPILE_MODULES, + dependencies=self.capabilities.dependencies, + ) + + def run(self, context: KernelContext) -> KernelResult: + require_params( + context, + self.ref, + required=frozenset( + {"entity", "column", "bands", "interpolation", "bands_sha256"} + ), + optional=frozenset({"component_column", "aggregate_entity"}), + ) + _whole_rows(context, self.ref) + sha256_param(context, self.ref, "bands_sha256") + artifact = context.artifacts.get("surface") + if artifact is None or artifact.type != TARGET_SURFACE_TYPE: + raise ValueError( + "Quantile map requires a target surface under alias 'surface'." + ) + entity = string_param(context, self.ref, "entity") + column = string_param(context, self.ref, "column") + aggregate = ( + string_param(context, self.ref, "aggregate_entity") + if "aggregate_entity" in context.params + else None + ) + # Aggregate maps rank each aggregate entity once at its own weight. + bands = _surface_bands( + canonical_document_param(context, self.ref, "bands"), + decode_target_surface(artifact.payload), + unit=entity if aggregate is None else aggregate, + ) + interpolation = canonical_document_param(context, self.ref, "interpolation") + if interpolation == {"method": "uniform"}: + interpolation = "uniform" + table = context.tables[entity] + component = ( + table[string_param(context, self.ref, "component_column")].to_numpy() + if "component_column" in context.params + else None + ) + group = component + if aggregate is not None: + if entity != "person": + raise ValueError( + "Aggregate quantile maps require person membership rows." + ) + group = {"entity_ids": table[f"person_{aggregate}_id"].to_numpy()} + if component is not None: + group["components"] = component + membership_components = ( + pd.DataFrame( + {"entity": group["entity_ids"], "component": component} + ) + .groupby("entity")["component"] + .nunique(dropna=False) + ) + if (membership_components > 1).any(): + raise ValueError( + "Aggregate members must repeat one component code." + ) + values = table[column].to_numpy() + weights = context.weights[entity].values + if isinstance(bands, dict) and set(bands) == {"components"}: + if component is None or pd.isna(component).any(): + raise ValueError( + "Component-specific bands require a nonmissing component_column." + ) + if set(component) != set(bands["components"]): + raise ValueError( + "Component bands must exactly cover the sliced component codes." + ) + mapped = np.empty(len(values), dtype=np.float64) + for value, component_bands in bands["components"].items(): + selected = component == value + component_group = ( + {"entity_ids": group["entity_ids"][selected]} + if isinstance(group, dict) + else None + ) + mapped[selected] = quantile_map( + values[selected], + weights[selected], + component_bands, + interpolation=interpolation, + group=component_group, + ) + else: + mapped = quantile_map( + values, + weights, + bands, + interpolation=interpolation, + group=group, + ) + return _column_result( + context, + self.ref, + {(entity, column): pd.Series(mapped, index=table.index, name=column)}, + ) + + +TRANSPORT_CREATE = TransportCreateKernel() +TRANSPORT_BOUNDARY = TransportBoundaryKernel() +TRANSPORT_CURRENCY = TransportCurrencyKernel() +TRANSPORT_QUANTILE_MAP = TransportQuantileMapKernel() + + +def register_population_kernels(registry: KernelRegistry) -> KernelRegistry: + """Register the population kernels explicitly and idempotently.""" + for kernel in ( + TRANSPORT_CREATE, + TRANSPORT_BOUNDARY, + TRANSPORT_CURRENCY, + TRANSPORT_QUANTILE_MAP, + ): + registry.register(kernel) + return registry diff --git a/packages/microcosm-build/src/microcosm/build/transport/registry.py b/packages/microcosm-build/src/microcosm/build/transport/registry.py new file mode 100644 index 000000000..3c969e2a7 --- /dev/null +++ b/packages/microcosm-build/src/microcosm/build/transport/registry.py @@ -0,0 +1,22 @@ +"""Explicit registration for the population side of a transport graph.""" + +from microcosm.graph import KernelRegistry + +from .column_kernels import register_column_kernels +from .geography_kernels import register_geography_kernels +from .population_kernels import register_population_kernels + +__all__ = ["register_transport_population_kernels"] + + +def register_transport_population_kernels(registry: KernelRegistry) -> KernelRegistry: + """Install the ten population kernels without import-time side effects. + + Source codecs are separately installed with ``register_transport_codecs``; + target, solver and terminal kernels keep their own registration functions. + Repeating registration with the same kernel objects is idempotent. + """ + register_population_kernels(registry) + register_column_kernels(registry) + register_geography_kernels(registry) + return registry diff --git a/packages/microcosm-build/tests/engine_free/shared/test_transport_codecs_geography.py b/packages/microcosm-build/tests/engine_free/shared/test_transport_codecs_geography.py new file mode 100644 index 000000000..ac5087d2d --- /dev/null +++ b/packages/microcosm-build/tests/engine_free/shared/test_transport_codecs_geography.py @@ -0,0 +1,574 @@ +"""Synthetic donor codecs and facts-derived atomic support. + +Properties: codec registration is idempotent; tree bytes ignore the root path +and file creation order; donor columns, original design mass, row order and +support channels match the shared reader; atomic counts conserve each published +population, differ from exact fractional quotas by less than one, and ignore +row permutation. Differential tests compare codec/kernel outputs with their +wrapped reader and support builder. No real donor or destination data is used. +""" + +from __future__ import annotations + +import base64 +import json +from copy import deepcopy +from dataclasses import replace +from fractions import Fraction +from hashlib import sha256 +from pathlib import Path +from unittest import mock + +import numpy as np +import pandas as pd +import pytest +from hypothesis import HealthCheck, example, given, settings +from hypothesis import strategies as st + +from microcosm.build.atomic_geography import decode_atomic_support +from microcosm.build.graph_atomic_geography import ( + AtomicAssignKernel, + AtomicSupportImportKernel, +) +from microcosm.build.ledger_artifact import ( + CONSUMER_ARTIFACT_SCHEMA_VERSION, + load_ledger_consumer_artifact, +) +from microcosm.build.ledger_targets import LedgerTargetReference +from microcosm.build.transport.artifact_types import TARGET_SURFACE_TYPE +from microcosm.build.transport.codecs import ( + DONOR_SOURCE_CODEC, + FACTS_SOURCE_CODEC, + RULESPEC_SOURCE_CODEC, + load_ledger_consumer_bytes, + load_populace_us_h5, + load_rulespec_tree_bytes, + register_transport_codecs, +) +from microcosm.build.transport.geography_kernels import ( + ATOMIC_SUPPORT_TYPE, + GEOGRAPHY_SUPPORT_FROM_FACTS, + register_geography_kernels, + support_from_surface, +) +from microcosm.build.transport.target_kernels import ( + decode_target_surface, + encode_target_surface, +) +from microcosm.calibrate import TargetRegistry, TargetSpec +from microcosm.calibrate.hierarchy import ( + CalibrationHierarchy, + HierarchyCategory, + HierarchyGeography, + HierarchyNode, +) +from microcosm.frame import WeightKind +from microcosm.frame.transport import read_populace_us_donor +from microcosm.graph import ( + ArtifactInput, + ArtifactOutput, + ArtifactValue, + Capabilities, + ContentStore, + Determinism, + Graph, + KernelBase, + KernelContext, + KernelResult, + Node, + Numeric, + NumericScope, + SourceRef, + compile_graph, + run_graph, +) +from microcosm.graph.canonical import canonical_json +from microcosm.graph.codecs import SourceCodecRegistry +from microcosm.graph.keys import source_content_key +from test_support.microcosm_build.transport_graph import toy_fact, write_facts +from test_support.paths import paths_for + +PROPERTY = settings(max_examples=40, deadline=None, database=None) +FILE_PROPERTY = settings( + max_examples=8, + deadline=None, + database=None, + suppress_health_check=[HealthCheck.function_scoped_fixture], +) +SHA = "4" * 64 + + +def _surface(populations=(100, 61), *, unit="count") -> bytes: + specs, references = [], [] + for index, population in enumerate(populations): + code = f"0{index + 1}" + name = f"population_{code}" + specs.append( + TargetSpec( + name=name, + entity="person", + measure="people", + value=population, + source="Synthetic population publication", + family="demography", + metadata={ + "precal_use": "atomic_geography_support", + "ledger_measure_unit": unit, + }, + hierarchy=CalibrationHierarchy( + provider=HierarchyNode("toy", "Toy"), + category=HierarchyCategory("population", "Population", "toy"), + geography=HierarchyGeography( + code, f"Toy area {code}", "territorial_authority" + ), + dimensions=(), + target=HierarchyNode(name, name), + ), + ) + ) + references.append( + LedgerTargetReference( + name=name, + ledger_fact_key=f"toy:{name}", + entity="person", + measure="people", + metadata={"activation_status": "active"}, + ) + ) + return encode_target_surface( + TargetRegistry(specs, country="xx"), + references, + references_sha256=SHA, + facts={"facts_sha256": SHA}, + ) + + +def _definition(shares=(0.3, 0.7)) -> dict: + codes = ("area", "ta", "region", "as_area", "as_area_alt") + columns = { + column: { + "kind": "code", + "source": "synthetic-crosswalk", + "vintage": "toy-2025", + "relation": "exact", + } + for column in codes + } + columns["population"] = { + "kind": "weight", + "source": "synthetic-facts", + "basis": "persons", + } + return { + "version": 1, + "level": "atomic", + "code_system": "toy", + "vintage": "toy-2025", + "columns": columns, + "share_tolerance": 1e-9, + "population_geography_level": "territorial_authority", + "population_code_column": "ta", + "rows": [ + { + "target": f"population_0{index + 1}", + "codes": { + "area": f"0{index + 1}:{area}", + "ta": f"0{index + 1}", + "region": "R" if index == 0 else "S", + "as_area": str(area), + "as_area_alt": str(1 + shares.index(max(shares))), + }, + "population_share": share, + } + for index in range(2) + for area, share in enumerate(shares, 1) + ], + } + + +def _node(definition=None, **params) -> Node: + return Node( + "geo.support", + GEOGRAPHY_SUPPORT_FROM_FACTS.ref, + population="nz.open", + artifact_inputs=( + ArtifactInput("surface", "precal", "surface", TARGET_SURFACE_TYPE), + ), + artifact_outputs=(ArtifactOutput("support", ATOMIC_SUPPORT_TYPE),), + params={ + "definition": canonical_json( + _definition() if definition is None else definition + ).decode(), + "definition_sha256": SHA, + "system": "toy-atomic", + **params, + }, + ) + + +def _context(node=None, payload=None) -> KernelContext: + node = _node() if node is None else node + payload = _surface() if payload is None else payload + return KernelContext( + node=node, + params=node.params, + tables={}, + weights={}, + strata=pd.Series([], dtype=object), + rng=np.random.default_rng(0), + sources={}, + artifacts={ + "surface": ArtifactValue( + payload, TARGET_SURFACE_TYPE, SHA, SHA, NumericScope() + ) + }, + ) + + +def test_codec_registration_is_explicit_idempotent_and_mode_correct() -> None: + registry = SourceCodecRegistry() + assert registry.names() == registry.bytes_names() == () + register_transport_codecs(registry) + before = registry.as_mapping(), registry.as_bytes_mapping() + register_transport_codecs(registry) + assert before == (registry.as_mapping(), registry.as_bytes_mapping()) + assert registry.names() == (DONOR_SOURCE_CODEC,) + assert registry.bytes_names() == tuple( + sorted((FACTS_SOURCE_CODEC, RULESPEC_SOURCE_CODEC)) + ) + with pytest.raises(TypeError, match="Frame codec"): + registry.load_bytes(DONOR_SOURCE_CODEC, "unused") + with pytest.raises(TypeError, match="raw-bytes codec"): + registry.load(FACTS_SOURCE_CODEC, "unused") + conflicting = SourceCodecRegistry() + conflicting.register_bytes(FACTS_SOURCE_CODEC, lambda path, **kwargs: b"wrong") + with pytest.raises(ValueError, match="already registered"): + register_transport_codecs(conflicting) + + +@FILE_PROPERTY +@given(permutation=st.permutations(tuple(range(6)))) +def test_donor_codec_matches_reader_columns_weights_channels_and_row_order( + tmp_path, permutation +) -> None: + fixture = ( + paths_for("microcosm-frame").tests / "fixtures/transport/synthetic_donor.json" + ) + raw = { + entity: pd.DataFrame(rows) + for entity, rows in json.loads(fixture.read_text())["tables"].items() + } + raw["person"] = raw["person"].iloc[list(permutation)].reset_index(drop=True) + path = tmp_path / "donor.h5" + with pd.HDFStore(path, mode="w") as store: + for entity, table in raw.items(): + store.put(entity, table, format="table") + direct = read_populace_us_donor( + path, sha256=sha256(path.read_bytes()).hexdigest(), size=path.stat().st_size + ) + frame = load_populace_us_h5(path) + for entity, table in direct.tables.items(): + pd.testing.assert_frame_equal(frame.table(entity), table) + assert frame.weights_for("household").kind == WeightKind.DESIGN + np.testing.assert_array_equal(frame.weights_for("household").values, direct.weights) + household_rows = pd.Index(direct.tables["household"]["household_id"]).get_indexer( + frame.person["person_household_id"] + ) + assert frame.strata.tolist() == direct.support_strata[household_rows].tolist() + assert frame.metadata["donor_country"] == "us" + assert frame.metadata["currency"] == "USD" + assert frame.metadata["source_person_ids"] == tuple(direct.source_person_ids) + + +@FILE_PROPERTY +@given(value=st.integers(min_value=0, max_value=1_000_000)) +def test_facts_codec_matches_validated_feed_bytes_and_refuses_tampering( + tmp_path, value +) -> None: + path = write_facts( + tmp_path / "facts.jsonl", [toy_fact("population", value, entity="person")] + ) + payload = path.read_bytes() + direct = load_ledger_consumer_artifact(path) + assert direct.facts[0]["value"] == value + assert load_ledger_consumer_bytes(path) == payload + directory = tmp_path / "artifact" + directory.mkdir(exist_ok=True) + (directory / "consumer_facts.jsonl").write_bytes(payload) + (directory / "manifest.json").write_bytes( + canonical_json( + { + "schema_version": CONSUMER_ARTIFACT_SCHEMA_VERSION, + "facts_sha256": sha256(payload).hexdigest(), + "fact_row_count": 1, + } + ) + ) + assert load_ledger_consumer_bytes(directory) == payload + (directory / "consumer_facts.jsonl").write_bytes(payload + b"\n") + with pytest.raises(ValueError, match="manifest hash"): + load_ledger_consumer_bytes(directory) + + +@FILE_PROPERTY +@given(first=st.binary(max_size=64), second=st.binary(max_size=64)) +def test_rulespec_codec_is_root_path_and_creation_order_inert_and_byte_sensitive( + tmp_path, first, second +) -> None: + trees = [tmp_path / "first", tmp_path / "renamed"] + for tree, files in zip( + trees, (("b.yaml", "a.yaml"), ("a.yaml", "b.yaml")), strict=True + ): + tree.mkdir(exist_ok=True) + for name in files: + (tree / name).write_bytes(first if name == "a.yaml" else second) + payload = load_rulespec_tree_bytes(trees[0]) + assert load_rulespec_tree_bytes(trees[1]) == payload + assert source_content_key("rules", trees[0]) == source_content_key( + "rules", trees[1] + ) + document = json.loads(payload) + assert [file["path"] for file in document["files"]] == ["a.yaml", "b.yaml"] + assert base64.b64decode(document["files"][0]["bytes_base64"]) == first + (trees[1] / "a.yaml").write_bytes(first + b"x") + assert load_rulespec_tree_bytes(trees[1]) != payload + assert source_content_key("rules", trees[0]) != source_content_key( + "rules", trees[1] + ) + + +@FILE_PROPERTY +@example(order=(3, 2, 1, 0)) +@given( + order=st.permutations(tuple(range(4))).filter(lambda order: order != (0, 1, 2, 3)) +) +def test_rulespec_codec_sorts_whatever_order_the_directory_lists( + tmp_path, order +) -> None: + """The envelope sorts entries itself; listing order never reaches it.""" + + tree = tmp_path / "tree" + names = ("a.yaml", "b/c.yaml", "b/d.yaml", "e.yaml") + for name in names: + (tree / name).parent.mkdir(parents=True, exist_ok=True) + (tree / name).write_bytes(name.encode()) + payload = load_rulespec_tree_bytes(tree) + listing = Path.rglob + + def reordered(self, pattern, *args, **kwargs): + entries = sorted(listing(self, pattern, *args, **kwargs)) + files = [entry for entry in entries if entry.is_file()] + # Directories first, then the files in a non-sorted order. + return iter( + [ + *(entry for entry in entries if not entry.is_file()), + *(files[index] for index in order), + ] + ) + + with mock.patch.object(Path, "rglob", reordered): + assert load_rulespec_tree_bytes(tree) == payload + document = json.loads(payload) + assert [file["path"] for file in document["files"]] == sorted(names) + + +def test_rulespec_codec_refuses_empty_tree_file_and_symlink(tmp_path) -> None: + with pytest.raises(ValueError, match="at least one"): + load_rulespec_tree_bytes(tmp_path) + file = tmp_path / "file" + file.write_bytes(b"x") + with pytest.raises(ValueError, match="regular directory"): + load_rulespec_tree_bytes(file) + (tmp_path / "link").symlink_to(file) + with pytest.raises(ValueError, match="symlinks"): + load_rulespec_tree_bytes(tmp_path) + + +@PROPERTY +@given( + populations=st.tuples(st.integers(1, 1_000_000), st.integers(1, 1_000_000)), + numerator=st.integers(0, 100), + permutation=st.permutations(tuple(range(4))), +) +def test_support_conserves_facts_apportions_within_one_and_ignores_row_order( + populations, numerator, permutation +) -> None: + shares = (numerator / 100, (100 - numerator) / 100) + definition = _definition(shares) + surface = decode_target_surface(_surface(populations)) + payload = support_from_surface(surface, definition, system="toy-atomic") + support = decode_atomic_support(payload) + for index, population in enumerate(populations): + selected = support.arrays["ta"] == f"0{index + 1}" + actual = support.arrays["population"][selected] + assert sum(int(count) for count in actual) == population + denominator = sum((Fraction(str(share)) for share in shares), Fraction()) + for count, share in zip(actual, shares, strict=True): + assert ( + abs( + Fraction(int(count)) + - population * Fraction(str(share)) / denominator + ) + < 1 + ) + permuted = deepcopy(definition) + permuted["rows"] = [definition["rows"][index] for index in permutation] + assert support_from_surface(surface, permuted, system="toy-atomic") == payload + for array in support.arrays.values(): + assert not array.flags.writeable + + +def test_geography_kernel_declares_its_siblings_platform_bitwise_numerics() -> None: + """Its deflated NPZ payload records the creating platform (zip header).""" + + numeric = GEOGRAPHY_SUPPORT_FROM_FACTS.capabilities.numeric + assert numeric is Numeric.PLATFORM_BITWISE + assert numeric is AtomicSupportImportKernel.capabilities.numeric + assert numeric is AtomicAssignKernel.capabilities.numeric + + +def test_geography_kernel_is_differential_to_shared_support_builder() -> None: + result = GEOGRAPHY_SUPPORT_FROM_FACTS.run(_context()) + assert result.artifacts["support"] == support_from_surface( + decode_target_surface(_surface()), _definition(), system="toy-atomic" + ) + assert result.receipt["definition_sha256"] == SHA + assert result.receipt["areas"] == 4 + assert GEOGRAPHY_SUPPORT_FROM_FACTS.capabilities.consumes_se is False + assert ( + GEOGRAPHY_SUPPORT_FROM_FACTS.implementation_hash() + == GEOGRAPHY_SUPPORT_FROM_FACTS.implementation_hash() + ) + + +@pytest.mark.parametrize( + "defect", + ( + "empty", + "missing_target", + "duplicate_area", + "bad_share", + "bad_sum", + "wrong_code", + "fractional_population", + "wrong_unit", + "wrong_metadata", + ), +) +def test_support_refuses_unharvested_incomplete_and_ambiguous_data(defect) -> None: + definition = _definition() + payload = _surface() + if defect == "empty": + definition["rows"] = [] + elif defect == "missing_target": + definition["rows"] = definition["rows"][:2] + elif defect == "duplicate_area": + definition["rows"][1]["codes"]["area"] = definition["rows"][0]["codes"]["area"] + elif defect == "bad_share": + definition["rows"][0]["population_share"] = -1 + elif defect == "bad_sum": + definition["rows"][0]["population_share"] = 0.1 + elif defect == "wrong_code": + definition["rows"][0]["codes"]["ta"] = "unknown" + elif defect == "fractional_population": + payload = _surface((100.5, 61)) + elif defect == "wrong_unit": + payload = _surface(unit="currency") + else: + definition["columns"]["as_area"]["relation"] = "guess" + with pytest.raises(ValueError): + GEOGRAPHY_SUPPORT_FROM_FACTS.run(_context(_node(definition), payload)) + + +@pytest.mark.parametrize( + "params", + ({"definition_sha256": "bad"}, {"unused": "x"}, {"definition": '{"version": 1}'}), +) +def test_geography_kernel_refuses_unknown_unbound_or_noncanonical_params( + params, +) -> None: + with pytest.raises(ValueError): + GEOGRAPHY_SUPPORT_FROM_FACTS.run(_context(_node(**params))) + + +class _SurfaceSource(KernelBase): + ref = "test.surface_source@1" + capabilities = Capabilities(Determinism.DETERMINISTIC) + + def run(self, context): + return KernelResult( + artifacts={"surface": context.sources["geo_surface"].read_bytes()} + ) + + +def test_geography_kernel_runs_after_synthetic_h5_create_and_memoizes(tmp_path) -> None: + from test_support.microcosm_build.transport_population import ( + population_graph, + population_registry, + write_population_sources, + ) + + sources = write_population_sources(tmp_path / "inputs") + base = population_graph(sources, through="boundary") + surface_path = tmp_path / "surface.json" + surface_path.write_bytes(_surface()) + surface_node = Node( + "precal", + _SurfaceSource.ref, + population="nz.open", + sources=("geo_surface",), + artifact_outputs=(ArtifactOutput("surface", TARGET_SURFACE_TYPE),), + ) + graph = compile_graph( + Graph( + base.country, + (*base.sources, SourceRef("geo_surface", "raw-bytes-v1")), + (*base.nodes, surface_node, _node()), + ) + ) + registry = population_registry() + registry.register(_SurfaceSource()) + register_geography_kernels(registry) + mapping = {**sources.mapping(), "geo_surface": surface_path} + store = ContentStore(tmp_path / "store") + cold = run_graph(graph, sources=mapping, store=store, kernels=registry) + warm = run_graph( + graph, sources=mapping, store=store, kernels=registry, resume="require" + ) + output = cold.nodes["geo.support"].opaque_artifacts["support"] + expected = GEOGRAPHY_SUPPORT_FROM_FACTS.run(_context()).artifacts["support"] + assert store.load_bytes(output) == expected + assert all(node.hit for node in warm.nodes.values()) + fresh = ContentStore(tmp_path / "fresh") + again = run_graph(graph, sources=mapping, store=fresh, kernels=registry) + assert cold.nodes["geo.support"].key == again.nodes["geo.support"].key + assert ( + fresh.load_bytes(again.nodes["geo.support"].opaque_artifacts["support"]) + == expected + ) + inert = compile_graph( + replace( + graph.graph, + nodes=tuple( + replace(node, description="Synthetic citation-only description") + for node in graph.graph.nodes + ), + ) + ) + edited = run_graph( + inert, sources=mapping, store=store, kernels=registry, resume="require" + ) + assert all(node.hit for node in edited.nodes.values()) + renamed = tmp_path / "renamed-surface.json" + renamed.write_bytes(surface_path.read_bytes()) + renamed_mapping = {**mapping, "geo_surface": renamed} + moved = run_graph( + graph, sources=renamed_mapping, store=store, kernels=registry, resume="require" + ) + assert all(node.hit for node in moved.nodes.values()) + renamed.write_bytes(_surface((101, 61))) + changed = run_graph(graph, sources=renamed_mapping, store=store, kernels=registry) + assert {name for name, node in changed.nodes.items() if not node.hit} == { + "precal", + "geo.support", + } 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 new file mode 100644 index 000000000..aacd6a368 --- /dev/null +++ b/packages/microcosm-build/tests/engine_free/shared/test_transport_column_kernels.py @@ -0,0 +1,1191 @@ +"""Column wrappers on synthetic donor graphs and generated concept tables. + +Hypothesis properties hold ids, row order, G4 differential equality, stable +take-up draws, eligibility implication, exclusivity and household rent mass. +The graph fixture uses production transport CREATE with a tiny synthetic H5; +every receipt, rule, rate and target here is a toy modelling choice. +""" + +from __future__ import annotations + +from dataclasses import replace +from unittest import mock + +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.artifact_types import TARGET_SURFACE_TYPE +from microcosm.build.transport.column_kernels import ( + CONCEPTS_ENCODE, + CONCEPTS_ENCODE_GROUPS, + TAKEUP_ASSIGN, + TRANSPORT_SCENARIO_OVERRIDE, + TRANSPORT_UNIT_ATTRIBUTES, + assign_receipts, + register_column_kernels, +) +from microcosm.build.transport.target_kernels import compile_target_surface +from microcosm.frame import Frame, WeightKind, Weights +from microcosm.frame import concept_mapping as mapping_module +from microcosm.frame.adapters.axiom import NZ_SCHEMA +from microcosm.frame.concept_mapping import ( + ConceptMapping, + GroupRule, + Identity, + InputBinding, + InputDeclaration, + StateBinding, + TakeUpThreshold, +) +from microcosm.frame.concepts import ( + CONCEPTS, + ContentBasis, + derive_take_up_draws, + split_for_transport, +) +from microcosm.frame.input_closure import ( + ClosureEntry, + InputClosure, + Knob, + UndeterminedPolicy, + apply_defaults, + resolve_judgments, +) +from microcosm.frame.unit_construction import ( + benefit_unit_attributes, + benefit_unit_membership, + build_benefit_units, + units_per_household, +) +from microcosm.graph import ( + ArtifactValue, + KernelContext, + Node, + NumericScope, + Owned, + SeedSource, + Slice, + StructuralDelta, +) +from microcosm.graph.population import dtype_for_token +from test_support.microcosm_build.transport_graph import ( + canonical_text, + reference_document, + reference_row, + toy_fact, + write_facts, +) +from test_support.microcosm_build.transport_population import ( + UNIT_RULE as DONOR_UNIT_RULE, +) +from test_support.microcosm_build.transport_population import ( + digest, + direct_population, + population_graph, + population_registry, + run_population, + write_population_sources, +) +from test_support.microcosm_frame.concept_frames import concept_frames + +PROPERTY = settings(max_examples=35, deadline=None, database=None) +GRAPH_PROPERTY = settings( + max_examples=6, + deadline=None, + database=None, + suppress_health_check=[HealthCheck.function_scoped_fixture], +) +PERSON_PATH = "xx/person.yaml" +GROUP_PATH = "xx/group.yaml" +OTHER_PATH = "xx/other.yaml" +RESOURCE_SHA = "b" * 64 +# Generated frames include unrelated co-resident parents. The toy rule makes +# their child's placement explicit instead of refusing those valid frames. +UNIT_RULE = replace(DONOR_UNIT_RULE, split_parents="first_parent_unit") + + +def _binding(name, entity, concept, *, group_rule=None, module=None, transform=None): + return InputBinding( + engine_input=name, + engine_entity=entity, + concepts=(concept,), + transform=Identity() if transform is None else transform, + relation="exact", + note="Synthetic exact input, not a statutory definition.", + group_rule=group_rule, + module=(PERSON_PATH if entity == "Person" else GROUP_PATH) + if module is None + else module, + canonical_input=f"xx:{entity}#input.{name}", + ) + + +def _mapping_with(bindings): + return ConceptMapping( + engine="axiom:xx", + engine_version="synthetic", + entity_correspondence={"person": "Person", "household": "Household"}, + input_declaration=InputDeclaration.USAGE_INFERRED, + bindings=tuple(bindings), + unmapped={ + item.id: "Outside this synthetic mapping." + for item in CONCEPTS + if item.id + not in {concept for binding in bindings for concept in binding.reads} + }, + ) + + +BINDINGS = ( + _binding("input_age", "Person", "fact:person.age"), + _binding("input_income", "Person", "fact:person.employment_income"), + _binding( + "input_assets", + "Family", + "fact:person.liquid_financial_assets", + group_rule="sum_over_members", + ), + _binding( + "input_rent", + "Family", + "fact:household.rent", + group_rule="allocate_to_reference_unit", + ), +) +MAPPING = _mapping_with(BINDINGS) +STATES = ( + StateBinding( + engine_input="input_non_beneficiary", + engine_entity="Family", + columns=( + "person.receives_alpha", + "person.receives_beta", + "person.receives_super", + ), + test="any_true", + group_rule=GroupRule.ANY_MEMBER, + negate=True, + note="Synthetic receipt state.", + module=GROUP_PATH, + canonical_input="xx:Family#input.input_non_beneficiary", + ), + StateBinding( + engine_input="input_support_channel", + engine_entity="Family", + columns=("household.donor_support_stratum",), + test="equals", + value="us:asec", + group_rule=GroupRule.HOUSEHOLD_VALUE, + note="Synthetic donor-channel indicator.", + module=GROUP_PATH, + canonical_input="xx:Family#input.input_support_channel", + ), +) +CLOSURE = InputClosure( + country="xx", + content_basis=ContentBasis.TRANSPORT, + mapping_engine=MAPPING.engine, + mapping_engine_version=MAPPING.engine_version, + surface_rulespec_commit="synthetic", + surface_engine_repository="toy/engine", + surface_engine_commit="synthetic", + modules={PERSON_PATH: RESOURCE_SHA, GROUP_PATH: RESOURCE_SHA}, + engine_optional_evidence="Synthetic engine declares no optional inputs.", + undetermined=UndeterminedPolicy( + action="refuse", + reason="Toy refusal.", + alternatives=("not_eligible", "eligible"), + new_default=False, + ), + knobs=( + Knob( + name="asset_test", + kind="switch_zeroes", + targets=("input_assets",), + central=True, + values=(), + method_card_row="toy", + note="Toy zero-assets alternative.", + ), + Knob( + name="rent_allocation", + kind="allocation", + targets=("input_rent",), + central="reference_unit", + values=("reference_unit", "per_adult_share"), + method_card_row="toy", + note="Toy rent allocation.", + ), + Knob( + name="rent_factor", + kind="factor", + targets=("input_rent",), + central=1.0, + values=(), + method_card_row="toy", + note="Toy scaling.", + ), + Knob( + name="channel_column", + kind="state_column", + targets=("household.donor_support_stratum",), + central="donor_support_stratum", + values=("donor_support_stratum", "tenure"), + method_card_row="toy", + note="Toy state replacement.", + ), + ), + entries=tuple( + ClosureEntry( + module=binding.module, + entity=binding.engine_entity, + input=binding.engine_input, + closure_class="encoded", + encoded_by="concept_binding", + ) + for binding in BINDINGS + ) + + tuple( + ClosureEntry( + module=GROUP_PATH, + entity="Family", + input=state.engine_input, + closure_class="encoded", + encoded_by="state_binding", + state_binding=state, + reason="Toy state.", + ) + for state in STATES + ) + + ( + ClosureEntry( + module=PERSON_PATH, + entity="Person", + input="eligibility_code", + closure_class="defaulted", + value=1, + reason="Synthetic determined eligibility.", + citation="Toy fixture.", + ), + ClosureEntry( + module=GROUP_PATH, + entity="Family", + input="input_default", + closure_class="defaulted", + value=False, + reason="Synthetic default.", + citation="Toy fixture.", + ), + ), +) + + +def _rule_params(rule=UNIT_RULE): + return { + "unit_rule": canonical_text(rule.to_dict()), + "unit_rule_sha256": digest(rule.to_dict()), + } + + +def _encode_params(*, groups=False, knobs=None, mapping=MAPPING): + params = { + "mapping": canonical_text(mapping.to_dict()), + "mapping_sha256": digest(mapping.to_dict()), + "closure": canonical_text(CLOSURE.to_dict()), + "closure_sha256": digest(CLOSURE.to_dict()), + "rulespec_paths": (GROUP_PATH if groups else PERSON_PATH,), + } + if groups: + params.update(_rule_params()) + params["engine_entity"] = "Family" + if knobs is not None: + params.update(knobs=canonical_text(knobs), scenario_sha256=digest(knobs)) + return params + + +def _contract(rate=0.6): + return { + "seed_column": "take_up_seed", + "programs": [ + { + "program": "xx.alpha", + "output": "receives_alpha", + "judgment_column": "eligibility_code", + "payment_column": "input_income", + "rate": rate, + }, + { + "program": "xx.beta", + "output": "receives_beta", + "judgment_column": "eligibility_code", + "payment_column": "input_income", + "rate": rate, + }, + { + "program": "xx.super", + "output": "receives_super", + "judgment_column": "eligibility_code", + "rate": rate, + }, + ], + "exclusion_groups": [["receives_alpha", "receives_beta", "receives_super"]], + } + + +def _context(kernel, frame, params, outputs, *, artifacts=None): + node = Node("toy", kernel.ref, population="open", params=params, outputs=outputs) + return KernelContext( + node=node, + tables={ + entity: frame.table(entity).copy(deep=True) + for entity in frame.schema.entities + }, + weights={ + entity: frame.resolve_weights(entity) for entity in frame.schema.entities + }, + strata=frame.strata.copy(deep=True), + params=node.params, + rng=np.random.default_rng(1), + artifacts={} if artifacts is None else artifacts, + ) + + +def _generated_frame(tables): + tables, _ = split_for_transport(tables) + person = tables["person"].copy() + household = tables["household"].sort_values("household_id").reset_index(drop=True) + units, members = build_benefit_units(person, household, UNIT_RULE) + person[UNIT_RULE.membership_column] = members + person["receives_alpha"] = False + person["receives_beta"] = False + person["receives_super"] = False + household["donor_support_stratum"] = pd.array( + ["us:asec"] * len(household), dtype="string" + ) + return Frame( + {"person": person, "household": household, "family": units}, + NZ_SCHEMA, + {"household": Weights(np.ones(len(household)), WeightKind.DESIGN)}, + pd.Series(["toy"] * len(person), name="stratum"), + ) + + +def _assert_column(result, frame, entity, column, expected): + actual = result.columns[entity, column] + assert actual.index.tolist() == frame.table(entity)[f"{entity}_id"].tolist() + assert actual.tolist() == expected.tolist() + + +@PROPERTY +@given(tables=concept_frames(max_households=3, max_members=4)) +def test_unit_attributes_equal_g4_and_keep_ids_and_order(tables): + frame = _generated_frame(tables) + family = frame.table("family") + person = frame.person + expected = benefit_unit_attributes( + person, family, person[UNIT_RULE.membership_column], UNIT_RULE + ) + # The input-bridge aliases share the wrapped function's adult composition. + expected["is_partnered"] = expected["is_couple"] + expected["is_single"] = ~expected["is_couple"] + outputs = tuple( + Owned( + "family", + column, + "string" if column == "family_type" else str(expected[column].dtype), + ) + for column in expected + if column != "family_id" + ) + (Owned("household", "n_family_units", "int64"),) + result = TRANSPORT_UNIT_ATTRIBUTES.run( + _context(TRANSPORT_UNIT_ATTRIBUTES, frame, _rule_params(), outputs) + ) + for output in outputs[:-1]: + _assert_column(result, frame, "family", output.column, expected[output.column]) + # Text takes the graph's pinned Python string storage, not pyarrow's. + assert result.columns["family", "family_type"].dtype == dtype_for_token("string") + _assert_column( + result, + frame, + "household", + "n_family_units", + units_per_household(frame.table("household"), family, UNIT_RULE), + ) + assert np.array_equal(frame.person["person_id"], tables["person"]["person_id"]) + + +@PROPERTY +@given(tables=concept_frames(max_households=3, max_members=4)) +def test_person_encoding_equals_direct_mapping_and_defaults(tables): + frame = _generated_frame(tables) + outputs = ( + Owned("person", "input_age", "int64"), + Owned("person", "input_income", "float64"), + Owned("person", "eligibility_code", "int64"), + ) + expected = apply_defaults( + MAPPING.encode( + {entity: frame.table(entity) for entity in ("person", "household")} + ).tables["person"], + CLOSURE.defaults(PERSON_PATH, "Person"), + ) + result = CONCEPTS_ENCODE.run( + _context(CONCEPTS_ENCODE, frame, _encode_params(), outputs) + ) + for output in outputs: + _assert_column(result, frame, "person", output.column, expected[output.column]) + + +@PROPERTY +@given( + tables=concept_frames(max_households=3, max_members=4), + allocation=st.sampled_from(["reference_unit", "per_adult_share"]), + asset_test=st.booleans(), + factor=st.floats(min_value=0.5, max_value=2, allow_nan=False), +) +def test_group_and_scenario_encoding_equal_g4_and_conserve_rent( + tables, allocation, asset_test, factor +): + frame = _generated_frame(tables) + knobs = { + "rent_allocation": allocation, + "asset_test": asset_test, + "rent_factor": factor, + } + membership = benefit_unit_membership( + frame.person, + frame.table("family"), + frame.person[UNIT_RULE.membership_column], + UNIT_RULE, + ) + expected = MAPPING.encode_groups( + {entity: frame.table(entity) for entity in ("person", "household")}, + {"Family": membership}, + modules=(GROUP_PATH,), + knobs=CLOSURE.group_knobs(knobs), + state_bindings=CLOSURE.state_bindings(), + ).tables["family"] + outputs = ( + Owned("family", "input_assets", "float64", rewrite=True), + Owned("family", "input_rent", "float64", rewrite=True), + Owned("family", "input_non_beneficiary", "bool", rewrite=True), + Owned("family", "input_support_channel", "bool", rewrite=True), + ) + params = _encode_params(groups=True, knobs=knobs) + for kernel in (CONCEPTS_ENCODE_GROUPS, TRANSPORT_SCENARIO_OVERRIDE): + result = kernel.run(_context(kernel, frame, params, outputs)) + for output in outputs: + _assert_column( + result, frame, "family", output.column, expected[output.column] + ) + rents = expected.groupby(frame.table("family")[UNIT_RULE.household_column])[ + "input_rent" + ].sum() + household_rents = ( + frame.table("household").set_index("household_id")["rent"] * factor + ) + np.testing.assert_allclose( + rents.reindex(household_rents.index), household_rents, rtol=1e-12 + ) + + +@st.composite +def receipt_rows(draw): + count = draw(st.integers(1, 15)) + codes = draw(st.lists(st.sampled_from([-1, 1]), min_size=count, max_size=count)) + seeds = draw( + st.lists( + st.floats(min_value=0, max_value=np.nextafter(1.0, 0), allow_nan=False), + min_size=count, + max_size=count, + ) + ) + income = draw( + st.lists( + st.floats(min_value=-100, max_value=100, allow_nan=False), + min_size=count, + max_size=count, + ) + ) + return pd.DataFrame( + { + "person_id": np.arange(count, dtype=np.int64) * 7 + 1, + "eligibility_code": codes, + "take_up_seed": seeds, + "input_income": income, + } + ) + + +@PROPERTY +@given( + person=receipt_rows(), + rate=st.floats(min_value=0, max_value=1, allow_nan=False), + data=st.data(), +) +def test_receipts_imply_eligibility_positive_payment_and_exclusivity_and_stable_draws( + person, rate, data +): + contract = _contract(rate) + out, audit = assign_receipts(person, contract) + eligible = resolve_judgments(person["eligibility_code"].to_numpy(), "refuse") + chosen = np.zeros(len(person), dtype=bool) + for program in contract["programs"]: + draws = derive_take_up_draws( + person["take_up_seed"].to_numpy(), program["program"] + ) + assert ((draws >= 0) & (draws < 1)).all() + flags = out[program["output"]].to_numpy() + permitted = eligible.copy() + if "payment_column" in program: + permitted &= person[program["payment_column"]].to_numpy() > 0 + expected = permitted & ~chosen & (draws < rate) + np.testing.assert_array_equal(flags, expected) + assert not (flags & ~permitted).any() + chosen |= flags + assert (out.iloc[:, 1:].sum(axis=1) <= 1).all() + permutation = data.draw(st.permutations(range(len(person)))) + permuted, permuted_audit = assign_receipts(person.iloc[list(permutation)], contract) + pd.testing.assert_frame_equal( + out.set_index("person_id").sort_index(), + permuted.set_index("person_id").sort_index(), + ) + assert audit == permuted_audit + + +def _column_graph(sources, *, rate=0.6, knobs=None): + graph = population_graph(sources, through="boundary") + frame, _ = direct_population(sources) + all_slices = tuple( + Slice( + entity, + tuple( + column + for column in frame.table(entity) + if column != f"{entity}_id" + and not ( + entity == "person" + and column.startswith("person_") + and column.endswith("_id") + ) + ), + ) + for entity in frame.schema.entities + ) + attrs = Node( + "nz.attributes", + TRANSPORT_UNIT_ATTRIBUTES.ref, + population="nz.open", + inputs=all_slices, + params=_rule_params(), + outputs=( + Owned("family", "family_type", "string"), + Owned("family", "n_dependent_children", "int64"), + Owned("family", "is_sole_parent", "bool"), + Owned("family", "youngest_dependent_child_age", "Int64"), + Owned("household", "n_family_units", "int64"), + ), + ) + encode = Node( + "nz.encode", + CONCEPTS_ENCODE.ref, + population="nz.open", + inputs=all_slices[:2], + params=_encode_params(), + outputs=( + Owned("person", "input_age", "int64"), + Owned("person", "input_income", "float64"), + Owned("person", "eligibility_code", "int64"), + ), + ) + contract = _contract(rate) + receipt = Node( + "nz.receipts", + TAKEUP_ASSIGN.ref, + population="nz.open", + inputs=(Slice("person", ("eligibility_code", "take_up_seed", "input_income")),), + params={ + "receipt_contract": canonical_text(contract), + "receipt_contract_sha256": digest(contract), + }, + outputs=tuple( + Owned("person", row["output"], "bool") for row in contract["programs"] + ), + ) + group_slices = ( + Slice( + "person", + all_slices[0].columns + + tuple(row["output"] for row in contract["programs"]), + ), + all_slices[1], + all_slices[2], + ) + group_outputs = ( + Owned("family", "input_assets", "float64"), + Owned("family", "input_rent", "float64"), + Owned("family", "input_non_beneficiary", "bool"), + Owned("family", "input_support_channel", "bool"), + Owned("family", "input_default", "bool"), + ) + group = Node( + "nz.groups", + CONCEPTS_ENCODE_GROUPS.ref, + population="nz.open", + inputs=group_slices, + params=_encode_params(groups=True), + outputs=group_outputs, + ) + branch = Node( + "nz.scenario", + "transport.boundary@1", + structural=StructuralDelta.FILTER, + base="nz.open", + inputs=(Slice("person", ("age",)),), + ) + knobs = ( + { + "asset_test": False, + "rent_allocation": "per_adult_share", + "channel_column": "tenure", + } + if knobs is None + else knobs + ) + scenario = Node( + "nz.override", + TRANSPORT_SCENARIO_OVERRIDE.ref, + population=branch.id, + inputs=group_slices, + params=_encode_params(groups=True, knobs=knobs), + outputs=tuple(replace(output, rewrite=True) for output in group_outputs[:4]), + ) + return replace( + graph, nodes=graph.nodes + (attrs, encode, receipt, group, branch, scenario) + ) + + +@GRAPH_PROPERTY +@given( + rate=st.floats(min_value=0, max_value=1, allow_nan=False), + factor=st.floats(min_value=0.5, max_value=2, allow_nan=False), +) +def test_each_column_kernel_runs_in_real_synthetic_donor_graph_and_rewrites_keep_order( + tmp_path, rate, factor +): + sources = write_population_sources(tmp_path / "sources") + graph = _column_graph( + sources, rate=rate, knobs={"rent_factor": factor, "asset_test": False} + ) + registry = population_registry() + register_column_kernels(registry) + run = run_population(tmp_path / "run", sources, graph=graph, registry=registry) + before = run.manifest.population("nz.open") + after = run.manifest.population("nz.scenario") + for entity in before.schema.entities: + np.testing.assert_array_equal( + before.table(entity)[f"{entity}_id"], after.table(entity)[f"{entity}_id"] + ) + np.testing.assert_array_equal(after.table("family")["input_assets"], 0.0) + np.testing.assert_allclose( + after.table("family")["input_rent"], + before.table("family")["input_rent"] * factor, + ) + flags = before.person[["receives_alpha", "receives_beta", "receives_super"]] + assert (flags.sum(axis=1) <= 1).all() + assert not (flags["receives_alpha"] & (before.person["input_income"] <= 0)).any() + assert not before.table("family")["input_default"].any() + assert run.manifest.known_failures == () + + +def test_take_up_refuses_undetermined_even_when_rate_zero_and_keeps_existing_exclusions(): + person = pd.DataFrame( + { + "person_id": [7], + "eligibility_code": [0], + "input_income": [10.0], + "take_up_seed": [0.5], + } + ) + with pytest.raises(ValueError, match="undetermined"): + assign_receipts(person, _contract(0)) + person["eligibility_code"] = 1 + person["existing_receipt"] = True + contract = _contract(1) + for row in contract["programs"]: + row["exclude_columns"] = ["existing_receipt"] + out, _ = assign_receipts(person, contract) + assert not out.iloc[:, 1:].to_numpy().any() + assert person["existing_receipt"].tolist() == [True] + + +def test_target_driven_receipts_bind_surface_and_report_expected_and_realized_mass( + tmp_path, +): + sources = write_population_sources(tmp_path / "sources") + frame, _ = direct_population(sources) + person = frame.person.copy() + person["eligibility_code"] = 1 + person["input_income"] = person["employment_income"] + frame = Frame( + { + "person": person, + "household": frame.table("household"), + "family": frame.table("family"), + }, + frame.schema, + {"household": frame.weights_for("household")}, + frame.strata, + ) + document = reference_document( + [reference_row("toy_receipts", entity="person", measure="receives_alpha")] + ) + facts = write_facts( + tmp_path / "receipt-facts.jsonl", + [toy_fact("toy_receipts", 60, entity="person")], + ) + surface, _ = compile_target_surface( + facts, document, country="xx", references_sha256=digest(document) + ) + artifact = ArtifactValue( + payload=surface, + type=TARGET_SURFACE_TYPE, + key=RESOURCE_SHA, + producer_key=RESOURCE_SHA, + numerics=NumericScope(), + ) + contract = _contract(1) + contract["programs"] = [ + { + "program": "xx.alpha", + "output": "receives_alpha", + "judgment_column": "eligibility_code", + "target": "toy_receipts", + } + ] + contract["exclusion_groups"] = [] + params = { + "receipt_contract": canonical_text(contract), + "receipt_contract_sha256": digest(contract), + } + result = TAKEUP_ASSIGN.run( + _context( + TAKEUP_ASSIGN, + frame, + params, + (Owned("person", "receives_alpha", "bool"),), + artifacts={"surface": artifact}, + ) + ) + audit = result.receipt["programs"]["receives_alpha"] + assert audit["target"] == 60 + assert audit["eligible_mass"] == 120 + assert audit["rate"] == 0.5 + direct, _ = assign_receipts( + person, + contract, + weights=frame.resolve_weights("person").values, + targets={"toy_receipts": 60}, + ) + _assert_column(result, frame, "person", "receives_alpha", direct["receives_alpha"]) + assert audit["realized_mass"] == sum( + frame.resolve_weights("person").values[direct["receives_alpha"]] + ) + with pytest.raises(ValueError, match="exceeds eligible"): + assign_receipts( + person, + contract, + weights=frame.resolve_weights("person").values, + targets={"toy_receipts": 121}, + ) + money_document = reference_document( + [reference_row("toy_receipts", entity="person", measure="input_income")] + ) + money_surface, _ = compile_target_surface( + facts, money_document, country="xx", references_sha256=digest(money_document) + ) + with pytest.raises(ValueError, match="person count"): + TAKEUP_ASSIGN.run( + _context( + TAKEUP_ASSIGN, + frame, + params, + (Owned("person", "receives_alpha", "bool"),), + artifacts={"surface": replace(artifact, payload=money_surface)}, + ) + ) + money_fact = toy_fact("toy_receipts", 60, entity="person") + money_fact["observed_measure"]["unit"] = "NZD" + wrong_unit_facts = write_facts(tmp_path / "monetary-facts.jsonl", [money_fact]) + wrong_unit_surface, _ = compile_target_surface( + wrong_unit_facts, document, country="xx", references_sha256=digest(document) + ) + with pytest.raises(ValueError, match="person count"): + TAKEUP_ASSIGN.run( + _context( + TAKEUP_ASSIGN, + frame, + params, + (Owned("person", "receives_alpha", "bool"),), + artifacts={"surface": replace(artifact, payload=wrong_unit_surface)}, + ) + ) + + +def test_wrappers_refuse_missing_resource_ids_missing_columns_and_non_input_overrides( + tmp_path, +): + sources = write_population_sources(tmp_path / "sources") + frame, _ = direct_population(sources) + params = _encode_params() + params.pop("mapping_sha256") + with pytest.raises(ValueError, match="missing required"): + CONCEPTS_ENCODE.run( + _context( + CONCEPTS_ENCODE, frame, params, (Owned("person", "input_age", "int64"),) + ) + ) + with pytest.raises(ValueError, match="did not encode"): + CONCEPTS_ENCODE.run( + _context( + CONCEPTS_ENCODE, + frame, + _encode_params(), + (Owned("person", "unmapped_input", "float64"),), + ) + ) + with pytest.raises(ValueError, match="sha256"): + CONCEPTS_ENCODE.run( + _context( + CONCEPTS_ENCODE, + frame, + {**_encode_params(), "shares": "{}"}, + (Owned("person", "input_age", "int64"),), + ) + ) + with pytest.raises(ValueError, match="non-text"): + CONCEPTS_ENCODE.run( + _context( + CONCEPTS_ENCODE, + frame, + _encode_params(), + (Owned("person", "input_age", "string"),), + ) + ) + with pytest.raises(ValueError, match="non-boolean"): + CONCEPTS_ENCODE.run( + _context( + CONCEPTS_ENCODE, + frame, + _encode_params(), + (Owned("person", "input_age", "bool"),), + ) + ) + huge_person = frame.person.copy() + huge_person["age"] = np.int64(2**53 + 1) + huge_frame = Frame( + { + "person": huge_person, + "household": frame.table("household"), + "family": frame.table("family"), + }, + frame.schema, + {"household": frame.weights_for("household")}, + frame.strata, + ) + with pytest.raises(ValueError, match="lose values"): + CONCEPTS_ENCODE.run( + _context( + CONCEPTS_ENCODE, + huge_frame, + _encode_params(), + (Owned("person", "input_age", "float64"),), + ) + ) + complete = _context( + CONCEPTS_ENCODE, + frame, + _encode_params(), + (Owned("person", "input_age", "int64"),), + ) + masked_person = complete.tables["person"].copy() + masked_person["is_in_scope"] = True + mask_context = replace( + complete, + tables={**complete.tables, "person": masked_person}, + ) + with pytest.raises(ValueError, match="all-row"): + CONCEPTS_ENCODE.run( + replace( + mask_context, + node=replace( + complete.node, + inputs=( + Slice("person", ("age", "is_in_scope"), rows="is_in_scope"), + ), + ), + ) + ) + with pytest.raises(ValueError, match="all-row"): + CONCEPTS_ENCODE.run( + replace( + mask_context, + node=replace( + complete.node, + inputs=(Slice("person", ("age", "is_in_scope")),), + outputs=( + Owned("person", "input_age", "int64", rows="is_in_scope"), + ), + ), + ) + ) + with pytest.raises(ValueError, match="declare a data-column slice"): + CONCEPTS_ENCODE.run( + replace(complete, tables={"person": complete.tables["person"]}) + ) + with pytest.raises(ValueError, match="rewrites"): + TRANSPORT_SCENARIO_OVERRIDE.run( + _context( + TRANSPORT_SCENARIO_OVERRIDE, + frame, + _encode_params(groups=True, knobs={"asset_test": False}), + (Owned("family", "input_assets", "float64"),), + ) + ) + + +def test_capabilities_are_deterministic_country_neutral_and_hash_code_only(): + for kernel in ( + TRANSPORT_UNIT_ATTRIBUTES, + CONCEPTS_ENCODE, + CONCEPTS_ENCODE_GROUPS, + TRANSPORT_SCENARIO_OVERRIDE, + TAKEUP_ASSIGN, + ): + assert "nz" not in kernel.ref + assert kernel.capabilities.determinism.value == "deterministic" + 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 TAKEUP_ASSIGN else SeedSource.NONE + ) + + +def _donor_frame(sources, **person_columns): + frame, _ = direct_population(sources) + person = frame.person.copy() + for name, value in person_columns.items(): + person[name] = value + return Frame( + { + "person": person, + "household": frame.table("household"), + "family": frame.table("family"), + }, + frame.schema, + {"household": frame.weights_for("household")}, + frame.strata, + ) + + +TAKE_UP_RATES = {"xx.claim": 1} + + +def _with_take_up_rates(params): + return { + **params, + "take_up_rates": canonical_text(TAKE_UP_RATES), + "take_up_rates_sha256": digest(TAKE_UP_RATES), + } + + +@pytest.mark.parametrize( + "kernel", [CONCEPTS_ENCODE, CONCEPTS_ENCODE_GROUPS, TRANSPORT_SCENARIO_OVERRIDE] +) +def test_encode_kernels_refuse_take_up_rates_so_they_derive_no_draws(tmp_path, kernel): + """SeedSource.NONE holds: takeup.assign@1 owns the keyed take-up draws.""" + sources = write_population_sources(tmp_path / "sources") + frame = _donor_frame(sources) + if kernel is CONCEPTS_ENCODE: + params = _encode_params() + outputs = (Owned("person", "input_age", "int64"),) + else: + params = _encode_params(groups=True, knobs={"asset_test": False}) + outputs = (Owned("family", "input_assets", "float64", rewrite=True),) + assert kernel.capabilities.seed_source is SeedSource.NONE + with pytest.raises(ValueError, match="derives no take-up draws"): + kernel.run(_context(kernel, frame, _with_take_up_rates(params), outputs)) + + +def test_selected_take_up_threshold_binding_is_refused_without_a_draw(tmp_path): + mapping = _mapping_with( + ( + *BINDINGS, + _binding( + "input_claims", + "Person", + "fact:person.take_up_seed", + transform=TakeUpThreshold(program="xx.claim"), + ), + ) + ) + sources = write_population_sources(tmp_path / "sources") + frame = _donor_frame(sources) + params = _encode_params(mapping=mapping) + outputs = (Owned("person", "input_claims", "bool"),) + with mock.patch.object( + mapping_module, + "derive_take_up_draws", + side_effect=AssertionError("An encode kernel derived a take-up draw."), + ): + with pytest.raises(ValueError, match="derives no take-up draws"): + CONCEPTS_ENCODE.run( + _context(CONCEPTS_ENCODE, frame, _with_take_up_rates(params), outputs) + ) + with pytest.raises(ValueError, match="take-up rate 'xx.claim'"): + CONCEPTS_ENCODE.run(_context(CONCEPTS_ENCODE, frame, params, outputs)) + + +def test_person_encoding_encodes_only_the_selected_paths(tmp_path): + mapping = _mapping_with( + ( + *BINDINGS, + _binding("input_other", "Person", "fact:person.age", module=OTHER_PATH), + ) + ) + sources = write_population_sources(tmp_path / "sources") + frame = _donor_frame(sources) + params = _encode_params(mapping=mapping) + result = CONCEPTS_ENCODE.run( + _context( + CONCEPTS_ENCODE, frame, params, (Owned("person", "input_age", "int64"),) + ) + ) + # The group path's two bindings are outside the selection, so none defer. + assert result.receipt["n_deferred"] == 0 + with pytest.raises(ValueError, match="did not encode declared column"): + CONCEPTS_ENCODE.run( + _context( + CONCEPTS_ENCODE, + frame, + params, + (Owned("person", "input_other", "int64"),), + ) + ) + selected = { + key: value + for key, value in params.items() + if key not in {"closure", "closure_sha256"} + } + selected["rulespec_paths"] = (OTHER_PATH,) + other = CONCEPTS_ENCODE.run( + _context( + CONCEPTS_ENCODE, frame, selected, (Owned("person", "input_other", "int64"),) + ) + ) + _assert_column(other, frame, "person", "input_other", frame.person["age"]) + + +def test_group_encoding_keeps_the_state_collision_guard_outside_selected_paths( + tmp_path, +): + """encode_groups checks state bindings against every concept binding.""" + mapping = _mapping_with( + ( + *BINDINGS, + _binding( + "input_non_beneficiary", + "Family", + "fact:person.age", + group_rule="sum_over_members", + module=OTHER_PATH, + ), + ) + ) + sources = write_population_sources(tmp_path / "sources") + frame = _donor_frame( + sources, receives_alpha=False, receives_beta=False, receives_super=False + ) + membership = benefit_unit_membership( + frame.person, + frame.table("family"), + frame.person[UNIT_RULE.membership_column], + UNIT_RULE, + ) + with pytest.raises(ValueError, match="bound by a concept and by state"): + mapping.encode_groups( + {entity: frame.table(entity) for entity in ("person", "household")}, + {"Family": membership}, + modules=(GROUP_PATH,), + state_bindings=CLOSURE.state_bindings(), + ) + outputs = (Owned("family", "input_non_beneficiary", "bool"),) + with pytest.raises(ValueError, match="bound by a concept and by state"): + CONCEPTS_ENCODE_GROUPS.run( + _context( + CONCEPTS_ENCODE_GROUPS, + frame, + _encode_params(groups=True, mapping=mapping), + outputs, + ) + ) + + +@pytest.mark.parametrize("defect", ["filter", "entity"]) +def test_target_receipts_refuse_a_filtered_or_non_person_count_target(tmp_path, defect): + sources = write_population_sources(tmp_path / "sources") + frame = _donor_frame(sources, eligibility_code=1) + entity = "household" if defect == "entity" else "person" + document = reference_document( + [ + reference_row( + "toy_receipts", + entity=entity, + measure="receives_alpha", + filter_="is_female" if defect == "filter" else None, + ) + ] + ) + facts = write_facts( + tmp_path / "receipt-facts.jsonl", + [toy_fact("toy_receipts", 60, entity=entity)], + ) + surface, _ = compile_target_surface( + facts, document, country="xx", references_sha256=digest(document) + ) + contract = { + "seed_column": "take_up_seed", + "programs": [ + { + "program": "xx.alpha", + "output": "receives_alpha", + "judgment_column": "eligibility_code", + "target": "toy_receipts", + } + ], + "exclusion_groups": [], + } + params = { + "receipt_contract": canonical_text(contract), + "receipt_contract_sha256": digest(contract), + } + artifact = ArtifactValue( + payload=surface, + type=TARGET_SURFACE_TYPE, + key=RESOURCE_SHA, + producer_key=RESOURCE_SHA, + numerics=NumericScope(), + ) + with pytest.raises(ValueError, match="unfiltered person count"): + TAKEUP_ASSIGN.run( + _context( + TAKEUP_ASSIGN, + frame, + params, + (Owned("person", "receives_alpha", "bool"),), + artifacts={"surface": artifact}, + ) + ) diff --git a/packages/microcosm-build/tests/engine_free/shared/test_transport_gate_bindings.py b/packages/microcosm-build/tests/engine_free/shared/test_transport_gate_bindings.py index 9beeb1a00..5b24952c1 100644 --- a/packages/microcosm-build/tests/engine_free/shared/test_transport_gate_bindings.py +++ b/packages/microcosm-build/tests/engine_free/shared/test_transport_gate_bindings.py @@ -16,9 +16,9 @@ - for every finite non-negative weight vector, subnormals included, the reported positive-weight median is the float64 nearest the exact median of the original positive weights (ties to even), and the summary is refused - exactly when the exact max/median ratio exceeds the float64 maximum; within - a relative 1e-12 of that maximum the max-normalized ratio may round either - way, an intended boundary pinned by a deterministic test (Hypothesis); + when the max-normalized ratio is unrepresentable; near the float64 maximum, + normalization can round a finite exact ratio to an overflowing reciprocal, + a refusal-only boundary pinned by a deterministic test (Hypothesis); - the declared surfaces in NZ ``gates.json`` equal the package resources they restate (differential checks against ``source_stages.json`` and ``export_contract.json``); @@ -1318,7 +1318,7 @@ def _exact_median(positive: list[Fraction]) -> Fraction: def _near_the_float64_limit(ratio: Fraction) -> bool: - """Whether max normalization may round ``ratio`` either side of overflow.""" + """Whether normalization can cause refusal near the float64 overflow limit.""" limit = Fraction(_FLOAT64_MAX) return abs(ratio - limit) <= limit * Fraction(1, 10**12) @@ -1459,6 +1459,7 @@ def test_a_ratio_at_the_float64_limit_fails_closed(self) -> None: @example(weights=[_FLOAT64_MAX, _FLOAT64_MAX]) @example(weights=[_FLOAT64_MAX / 2, _FLOAT64_MAX]) @example(weights=[0.0, 5e-324, _FLOAT64_MAX]) + @example(weights=[5e-324, 1e300]) @given(weights=_FULL_RANGE_WEIGHTS) def test_property_summary_is_exact_over_the_full_float64_range( self, weights 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 new file mode 100644 index 000000000..5f3e9f5e1 --- /dev/null +++ b/packages/microcosm-build/tests/engine_free/shared/test_transport_population_kernels.py @@ -0,0 +1,776 @@ +"""Synthetic H5 graphs exercise the population transport contracts. + +Hypothesis invariants: exact CREATE declarations, destination design mass, +boundary mass conservation, deterministic keys/artifacts, zero memoized +executions, inert prose, source content identity, stable seeds and rewrite +identity. The wrapped-function differentials use the same authenticated +synthetic donor, never destination microdata. +""" + +from __future__ import annotations + +import json +import shutil +from dataclasses import replace +from pathlib import Path +from tempfile import TemporaryDirectory + +import numpy as np +import pandas as pd +import pytest +from hypothesis import HealthCheck, example, given, settings +from hypothesis import strategies as st + +from microcosm.build.transport.population_kernels import ( + TRANSPORT_BOUNDARY, + TRANSPORT_CREATE, + TRANSPORT_CURRENCY, + TRANSPORT_QUANTILE_MAP, + register_population_kernels, +) +from microcosm.build.transport.registry import register_transport_population_kernels +from microcosm.frame import WeightKind +from microcosm.frame.adapters.axiom import NZ_SCHEMA +from microcosm.frame.transport import currency_bridge, quantile_map +from microcosm.graph import ( + Capabilities, + Determinism, + KernelBase, + KernelContext, + KernelRegistry, + Node, + NodeRejectedError, + Owned, + SeedSource, + Slice, + compile_graph, +) +from microcosm.graph.population import dtype_for_token +from test_support.microcosm_build.transport_graph import ( + canonical_text, + descendants, + reference_document, + reference_row, + toy_fact, + write_facts, +) +from test_support.microcosm_build.transport_population import ( + BAND_FACTS, + ENTITIES, + PopulationSources, + band_document, + create_node, + digest, + direct_population, + expected_bands, + population_graph, + population_registry, + run_population, + write_population_sources, +) + +PROPERTY = settings( + max_examples=8, + deadline=None, + database=None, + suppress_health_check=[HealthCheck.function_scoped_fixture], +) + + +def _keys(run): + return {name: receipt.key for name, receipt in run.manifest.nodes.items()} + + +def _artifact_bytes(run): + return { + node: { + name: run.store.load_bytes(key) + for name, key in receipt.opaque_artifacts.items() + } + for node, receipt in run.manifest.nodes.items() + } + + +def _structural_columns(entity): + if entity == "person": + return {"person_id", "person_household_id", "person_family_id"} + return {f"{entity}_id"} + + +class CountingKernel(KernelBase): + def __init__(self, kernel, calls): + self.kernel = kernel + self.calls = calls + self.ref = kernel.ref + self.capabilities = kernel.capabilities + + def implementation_hash(self): + return self.kernel.implementation_hash() + + def run(self, context): + self.calls.append(self.ref) + return self.kernel.run(context) + + +def test_kernel_registration_is_explicit_idempotent_and_country_neutral(): + registry = KernelRegistry() + assert register_population_kernels(registry) is registry + assert register_population_kernels(registry) is registry + kernels = ( + TRANSPORT_CREATE, + TRANSPORT_BOUNDARY, + TRANSPORT_CURRENCY, + TRANSPORT_QUANTILE_MAP, + ) + assert set(registry.refs()) == {kernel.ref for kernel in kernels} + for kernel in kernels: + assert "nz" not in kernel.ref + assert kernel.capabilities.determinism is Determinism.DETERMINISTIC + 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 + ) + + +def test_aggregate_population_registration_installs_exactly_ten_kernels(): + registry = KernelRegistry() + assert register_transport_population_kernels(registry) is registry + assert register_transport_population_kernels(registry) is registry + assert set(registry.refs()) == { + "transport.create@1", + "transport.boundary@1", + "transport.currency@1", + "transport.quantile_map@1", + "transport.unit_attributes@1", + "concepts.encode@1", + "concepts.encode_groups@1", + "takeup.assign@1", + "geography.support_from_facts@1", + "transport.scenario_override@1", + } + + +@PROPERTY +@example(mass=1) +@given(mass=st.integers(min_value=1, max_value=100_000)) +def test_a1_a2_d2_create_contracts_on_a_synthetic_graph(tmp_path, mass): + with TemporaryDirectory(dir=tmp_path) as temporary: + root = Path(temporary) + sources = write_population_sources(root / "sources", mass=float(mass)) + graph = population_graph(sources) + calls = [] + counted = KernelRegistry() + registry = population_registry() + for ref in registry.refs(): + counted.register(CountingKernel(registry.get(ref), calls)) + first = run_population(root / "one", sources, graph=graph, registry=counted) + assert len(calls) == len(graph.nodes) + calls.clear() + second = run_population(root / "two", sources, graph=graph) + assert _keys(first) == _keys(second) + assert _artifact_bytes(first) == _artifact_bytes(second) + assert first.manifest.key == second.manifest.key + again = run_population( + root / "one", sources, graph=graph, registry=counted, store=first.store + ) + assert calls == [] + assert all(receipt.hit for receipt in again.manifest.nodes.values()) + assert _keys(again) == _keys(first) + + create = first.manifest.population("nz.create") + boundary = first.store.load_frame(first.manifest.nodes["nz.open"].frame_key) + assert create.schema == NZ_SCHEMA + assert create.weights_for("household").kind is WeightKind.DESIGN + for frame in (create, boundary): + assert frame.resolve_weights("person").values.sum() == pytest.approx( + mass, rel=1e-6 + ) + declared = { + (item.entity, item.column) for item in graph.node("nz.create").outputs + } + actual = { + (entity, column) + for entity in ENTITIES + for column in create.table(entity) + if column not in _structural_columns(entity) + } + assert actual == declared + receipt = first.manifest.nodes["nz.create"].receipt + expected, dropped = direct_population(sources) + assert dict(receipt["dropped_columns"]) == dropped + assert receipt["destination_mass"] == mass + for entity in ENTITIES: + pd.testing.assert_frame_equal(create.table(entity), expected.table(entity)) + np.testing.assert_array_equal( + create.weights_for("household").values, + expected.weights_for("household").values, + ) + pd.testing.assert_series_equal(create.strata, expected.strata) + assert create.strata.tolist() == ["us:asec"] * 3 + ["us:puf_tax_detail"] * 3 + # Strata keep object storage whether or not pyarrow is installed. + assert create.strata.dtype == object + + +@PROPERTY +@example(prose="Edited prose and citations.") +@given(prose=st.text(alphabet="abcdefghijklmnopqrstuvwxyz .", max_size=35)) +def test_a4_description_and_citation_change_no_key(tmp_path, prose): + with TemporaryDirectory(dir=tmp_path) as temporary: + root = Path(temporary) + sources = write_population_sources(root / "sources") + first = run_population(root / "one", sources) + second = run_population( + root / "two", sources, graph=population_graph(sources, description=prose) + ) + assert _keys(first) == _keys(second) + assert _artifact_bytes(first) == _artifact_bytes(second) + + +def _relocated(sources, root): + root.mkdir(parents=True) + paths = {} + for name, original in sources.mapping().items(): + paths[name] = root / f"renamed-{name}.bin" + shutil.copyfile(original, paths[name]) + return PopulationSources(**paths) + + +@PROPERTY +@example(source="donor", whitespace=b" ") +@example(source="facts", whitespace=b" ") +@example(source="band_facts", whitespace=b" ") +@given( + source=st.sampled_from(["donor", "facts", "band_facts"]), + whitespace=st.sampled_from([b" ", b"\n", b"\t"]), +) +def test_a6_source_paths_are_inert_and_bytes_rekey_exact_consumers( + tmp_path, source, whitespace +): + with TemporaryDirectory(dir=tmp_path) as temporary: + root = Path(temporary) + sources = write_population_sources(root / "sources") + first = run_population(root / "run", sources) + moved = _relocated(sources, root / "moved") + renamed = run_population(root / "run", moved, store=first.store) + assert _keys(renamed) == _keys(first) + assert _artifact_bytes(renamed) == _artifact_bytes(first) + assert all(receipt.hit for receipt in renamed.manifest.nodes.values()) + changed_path = moved.mapping()[source] + changed_path.write_bytes(changed_path.read_bytes() + whitespace) + # The HDF decoder and facts decoder accept trailing whitespace. For + # the donor, run_population re-pins CREATE (population_graph reads + # the new size and SHA-256), so CREATE's key moves through its params + # as well as its source content; stale pins are refused below. + changed = run_population(root / "run", moved, store=first.store) + consumer = "nz.bands" if source == "band_facts" else "nz.create" + expected = descendants(changed.graph, {consumer}) + rekeyed = { + node for node in _keys(first) if _keys(first)[node] != _keys(changed)[node] + } + assert rekeyed == expected + assert { + node for node, receipt in changed.manifest.nodes.items() if not receipt.hit + } == expected + + +@pytest.mark.parametrize( + ("edit", "message"), + [("append", "Donor size mismatch"), ("same_size", "Donor SHA-256 mismatch")], +) +def test_a6_donor_byte_change_with_stale_pins_is_refused(tmp_path, edit, message): + sources = write_population_sources(tmp_path / "sources") + stale = population_graph(sources) + first = run_population(tmp_path / "run", sources, graph=stale) + payload = sources.donor.read_bytes() + if edit == "append": + changed = payload + b" " + else: + changed = payload[:-1] + bytes([payload[-1] ^ 1]) + assert changed != payload + sources.donor.write_bytes(changed) + with pytest.raises(NodeRejectedError, match=message): + run_population(tmp_path / "run", sources, graph=stale, store=first.store) + + +@PROPERTY +@example(permutation=tuple(range(6)), rate=1.6) +@given( + permutation=st.permutations(tuple(range(6))), + rate=st.floats(min_value=0.1, max_value=10, allow_nan=False, allow_infinity=False), +) +def test_seeds_and_rewrites_preserve_identity_with_wrapped_function_differentials( + tmp_path, permutation, rate +): + with TemporaryDirectory(dir=tmp_path) as temporary: + root = Path(temporary) + sources = write_population_sources(root / "sources", permutation=permutation) + graph = population_graph(sources, rate=rate) + run = run_population(root / "run", sources, graph=graph) + created = run.manifest.population("nz.create") + rewritten = run.manifest.population("nz.open") + direct, _ = direct_population(sources) + baseline = write_population_sources(root / "baseline") + stable, _ = direct_population(baseline) + seeds = created.table("person").set_index("person_id")["take_up_seed"] + expected_seeds = stable.table("person").set_index("person_id")["take_up_seed"] + assert ((seeds >= 0) & (seeds < 1)).all() + pd.testing.assert_series_equal(seeds.sort_index(), expected_seeds.sort_index()) + for entity in ENTITIES: + for column in _structural_columns(entity): + np.testing.assert_array_equal( + rewritten.table(entity)[column], created.table(entity)[column] + ) + np.testing.assert_array_equal( + rewritten.table(entity)[f"{entity}_id"], + direct.table(entity)[f"{entity}_id"], + ) + bridged = currency_bridge( + {entity: direct.table(entity) for entity in ENTITIES}, + {"person": ["interest_income"], "household": ["rent"]}, + rate, + ) + for entity, column in (("person", "interest_income"), ("household", "rent")): + np.testing.assert_array_equal( + rewritten.table(entity)[column], bridged[entity][column] + ) + bands = expected_bands(sources) + expected = quantile_map( + direct.table("person")["employment_income"], + direct.resolve_weights("person").values, + bands, + interpolation="uniform", + ) + np.testing.assert_array_equal( + rewritten.table("person")["employment_income"], expected + ) + # The unequal facts must matter: equal or reversed shares map + # differently, so the differential checks the facts' shares. + for shares in ( + [1 / len(bands)] * len(bands), + [band["share"] for band in reversed(bands)], + ): + assert not np.array_equal( + expected, + quantile_map( + direct.table("person")["employment_income"], + direct.resolve_weights("person").values, + [ + {**band, "share": share} + for band, share in zip(bands, shares, strict=True) + ], + interpolation="uniform", + ), + ) + + +def _two_holders_in_one_household(tables): + """Give the first household's second member assets beside its head's.""" + person = tables["person"] + second = person.index[person["person_household_id"] == 101][1] + person.loc[second, "bank_account_assets"] = 40 + + +@PROPERTY +@example(mass=1) +@given(mass=st.integers(min_value=1, max_value=100_000)) +def test_household_aggregate_quantile_map_equals_direct_function(tmp_path, mass): + with TemporaryDirectory(dir=tmp_path) as temporary: + root = Path(temporary) + sources = write_population_sources( + root / "sources", + mass=float(mass), + band_entity="household", + edit_tables=_two_holders_in_one_household, + ) + graph = population_graph( + sources, qmap_column="liquid_financial_assets", aggregate_entity="household" + ) + run = run_population(root / "run", sources, graph=graph) + direct, _ = direct_population(sources) + person = direct.table("person") + values = person["liquid_financial_assets"] + weights = direct.resolve_weights("person").values + holders = (values != 0).groupby(person["person_household_id"]).sum() + assert (holders >= 2).any(), "A household needs two nonzero members." + expected = quantile_map( + values, + weights, + expected_bands(sources), + interpolation="uniform", + group={"entity_ids": person["person_household_id"]}, + ) + # Aggregation must matter: mapping persons one by one differs. + assert not np.array_equal( + expected, + quantile_map( + values, weights, expected_bands(sources), interpolation="uniform" + ), + ) + rewritten = run.manifest.population("nz.open").table("person") + np.testing.assert_array_equal(rewritten["liquid_financial_assets"], expected) + np.testing.assert_array_equal(rewritten["person_id"], person["person_id"]) + + +@PROPERTY +@example(multipliers=(1, 2)) +@given(multipliers=st.tuples(st.integers(1, 4), st.integers(1, 4))) +def test_component_specific_bands_equal_separate_direct_maps(tmp_path, multipliers): + with TemporaryDirectory(dir=tmp_path) as temporary: + root = Path(temporary) + sources = write_population_sources(root / "sources") + direct, _ = direct_population(sources) + person = direct.table("person") + components = sorted(person["sex"].unique().tolist()) + assert len(components) == len(multipliers) + fixture_bands = expected_bands(sources) + document = { + "components": [ + { + "value": component, + "bands": [ + { + "lower": band["lower"] * multiplier, + "upper": band["upper"] * multiplier, + "reference": reference, + } + for band, (reference, _) in zip( + fixture_bands, BAND_FACTS, strict=True + ) + ], + } + for component, multiplier in zip(components, multipliers, strict=True) + ] + } + graph = population_graph(sources) + quantile = replace( + graph.node("nz.quantile"), + inputs=(Slice("person", ("employment_income", "sex")),), + params={ + **graph.node("nz.quantile").params, + "component_column": "sex", + "bands": canonical_text(document), + "bands_sha256": digest(document), + }, + ) + graph = replace( + graph, + nodes=tuple( + quantile if node.id == quantile.id else node for node in graph.nodes + ), + ) + run = run_population(root / "run", sources, graph=graph) + expected = np.empty(len(person), dtype=np.float64) + weights = direct.resolve_weights("person").values + for component, multiplier in zip(components, multipliers, strict=True): + selected = person["sex"].to_numpy() == component + expected[selected] = quantile_map( + person["employment_income"].to_numpy()[selected], + weights[selected], + [ + { + "lower": band["lower"] * multiplier, + "upper": band["upper"] * multiplier, + "share": band["share"], + } + for band in fixture_bands + ], + interpolation="uniform", + ) + rewritten = run.manifest.population("nz.open").table("person") + np.testing.assert_array_equal(rewritten["employment_income"], expected) + np.testing.assert_array_equal(rewritten["person_id"], person["person_id"]) + + +@PROPERTY +@example(resource=("nz.currency", "currency_sha256"), digit="0") +@given( + resource=st.sampled_from( + [ + ("nz.create", "unit_rule_sha256"), + ("nz.currency", "currency_sha256"), + ("nz.quantile", "bands_sha256"), + ] + ), + digit=st.sampled_from("0123456789abcdef"), +) +def test_resource_identity_rekeys_only_its_node_and_consumers( + tmp_path, resource, digit +): + """These digests are format-checked only: any hex moves keys, not bytes.""" + with TemporaryDirectory(dir=tmp_path) as temporary: + root = Path(temporary) + sources = write_population_sources(root / "sources") + before = run_population(root / "run", sources) + node_id, parameter = resource + graph = before.graph + node = replace( + graph.node(node_id), + params={**graph.node(node_id).params, parameter: digit * 64}, + ) + edited = replace( + graph, + nodes=tuple(node if item.id == node_id else item for item in graph.nodes), + ) + after = run_population(root / "run", sources, graph=edited, store=before.store) + assert { + name for name in _keys(before) if _keys(before)[name] != _keys(after)[name] + } == descendants(edited, {node_id}) + assert _artifact_bytes(before) == _artifact_bytes(after) + + +@pytest.mark.parametrize( + "problem", ["absent_reference", "missing_component", "extra_component"] +) +def test_quantile_map_refuses_absent_facts_and_mismatched_component_coverage( + tmp_path, problem +): + sources = write_population_sources(tmp_path / "sources") + graph = population_graph(sources) + quantile = graph.node("nz.quantile") + document = json.loads(quantile.params["bands"]) + if problem == "absent_reference": + document["bands"][0]["reference"] = "absent_fact" + inputs = quantile.inputs + extra_params = {} + message = "reference is absent" + else: + direct, _ = direct_population(sources) + components = sorted(direct.table("person")["sex"].unique().tolist()) + codes = ( + components[:-1] + if problem == "missing_component" + else [*components, "absent_code"] + ) + document = { + "components": [ + {"value": code, "bands": document["bands"]} for code in codes + ] + } + inputs = (Slice("person", ("employment_income", "sex")),) + extra_params = {"component_column": "sex"} + message = "exactly cover" + quantile = replace( + quantile, + inputs=inputs, + params={ + **quantile.params, + "bands": canonical_text(document), + "bands_sha256": digest(document), + **extra_params, + }, + ) + graph = replace( + graph, + nodes=tuple( + quantile if node.id == quantile.id else node for node in graph.nodes + ), + ) + with pytest.raises(NodeRejectedError, match=message): + run_population(tmp_path / "run", sources, graph=graph) + + +def test_create_runtime_rejects_a_missing_declared_column(tmp_path): + sources = write_population_sources(tmp_path / "sources") + graph = population_graph(sources, through="boundary") + create = replace( + graph.node("nz.create"), outputs=graph.node("nz.create").outputs[:-1] + ) + graph = replace(graph, nodes=(create, graph.node("nz.open"))) + compiled = compile_graph(graph) + assert compiled is not None + with pytest.raises(NodeRejectedError, match="exactly equal its declaration"): + run_population(tmp_path / "run", sources, graph=graph) + + +class StructuralMutation(KernelBase): + ref = "test.structural_mutation@1" + capabilities = Capabilities(Determinism.DETERMINISTIC) + + def run(self, context): + raise AssertionError( + "The runtime must refuse structural ownership before execution." + ) + + +@PROPERTY +@example(column="person_household_id") +@given(column=st.sampled_from(["person_id", "person_household_id", "person_family_id"])) +def test_variant_structural_ownership_is_rejected_at_run_time(tmp_path, column): + with TemporaryDirectory(dir=tmp_path) as temporary: + root = Path(temporary) + sources = write_population_sources(root / "sources") + graph = population_graph(sources, through="boundary") + node = Node( + "nz.variant.structural", + StructuralMutation.ref, + population="nz.open", + inputs=(Slice("person", ("age",)),), + outputs=(Owned("person", column, "int64"),), + ) + graph = replace(graph, nodes=(*graph.nodes, node)) + compile_graph(graph) + registry = population_registry() + registry.register(StructuralMutation()) + with pytest.raises(NodeRejectedError, match="cannot own structural column"): + run_population(root / "run", sources, graph=graph, registry=registry) + + +@pytest.mark.parametrize( + ("entity", "measure", "filter_"), + [ + ("household", "person_count", None), + ("person", "employment_income", None), + ("person", "person_count", "adults"), + ], +) +def test_create_refuses_a_destination_mass_that_is_not_total_person_population( + tmp_path, entity, measure, filter_ +): + sources = write_population_sources(tmp_path / "sources") + document = reference_document( + [ + reference_row( + "toy_population", entity=entity, measure=measure, filter_=filter_ + ) + ] + ) + fact = toy_fact("toy_population", 120.0, entity=entity) + if measure == "employment_income": + fact["observed_measure"]["unit"] = "currency" + write_facts(sources.facts, [fact]) + node = create_node(sources) + node = replace( + node, + params={ + **node.params, + "mass_reference": canonical_text(document), + "mass_reference_sha256": digest(document), + }, + ) + context = 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}, + ) + with pytest.raises(ValueError, match="population|person count"): + TRANSPORT_CREATE.run(context) + + +def test_create_accepts_a_person_count_fact_with_a_country_prepared_measure_name( + tmp_path, +): + sources = write_population_sources(tmp_path / "sources") + document = reference_document( + [reference_row("toy_population", entity="person", measure="people")] + ) + node = create_node(sources) + node = replace( + node, + params={ + **node.params, + "mass_reference": canonical_text(document), + "mass_reference_sha256": digest(document), + }, + ) + context = 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}, + ) + result = TRANSPORT_CREATE.run(context) + assert result.frame.resolve_weights("person").values.sum() == pytest.approx(120.0) + assert result.receipt["mass_reference"] == "toy_population" + + +def _with_band_references(graph, references): + bands = graph.node("nz.bands") + bands = replace( + bands, + params={ + **bands.params, + "references": canonical_text(references), + "references_sha256": digest(references), + }, + ) + return replace( + graph, + nodes=tuple(bands if node.id == bands.id else node for node in graph.nodes), + ) + + +@pytest.mark.parametrize( + ("aggregate_entity", "fact_entity", "filter_", "unit"), + [ + (None, "household", None, "count"), + ("household", "person", None, "count"), + (None, "person", "is_female", "count"), + (None, "person", None, "currency"), + ], +) +def test_quantile_bands_refuse_facts_that_do_not_count_the_ranked_unit( + tmp_path, aggregate_entity, fact_entity, filter_, unit +): + sources = write_population_sources(tmp_path / "sources", band_entity=fact_entity) + facts = [toy_fact(name, value, entity=fact_entity) for name, value in BAND_FACTS] + for fact in facts: + fact["observed_measure"]["unit"] = unit + write_facts(sources.band_facts, facts) + graph = population_graph( + sources, + qmap_column="employment_income" + if aggregate_entity is None + else "liquid_financial_assets", + aggregate_entity=aggregate_entity, + ) + references = band_document(fact_entity) + if filter_ is not None: + for row in references["target_references"]: + row["filter"] = filter_ + graph = _with_band_references(graph, references) + ranked = "person" if aggregate_entity is None else aggregate_entity + with pytest.raises(NodeRejectedError, match=f"unfiltered {ranked} count facts"): + run_population(tmp_path / "run", sources, graph=graph) + + +@pytest.mark.parametrize("problem", ["household_slice", "free_mass"]) +def test_boundary_refuses_without_a_person_slice_or_conserved_mass(tmp_path, problem): + sources = write_population_sources(tmp_path / "sources") + graph = population_graph(sources, through="boundary") + boundary = graph.node("nz.open") + if problem == "household_slice": + boundary = replace(boundary, inputs=(Slice("household", ("rent",)),)) + message = "person data-column slice" + else: + boundary = replace(boundary, mass="free") + message = "mass='conserve'" + graph = replace(graph, nodes=(graph.node("nz.create"), boundary)) + with pytest.raises(NodeRejectedError, match=message): + run_population(tmp_path / "run", sources, graph=graph) + + +def test_create_artifacts_pin_python_string_storage_with_pyarrow_installed(tmp_path): + pinned = dtype_for_token("string") + # microcosm-build depends on pyarrow; with it installed, a bare "string" + # dtype stores pyarrow strings, which the graph's token does not. + assert pd.Series(["x"]).astype("string").dtype != pinned + sources = write_population_sources(tmp_path / "sources") + run = run_population(tmp_path / "run", sources) + text = [ + (item.entity, item.column) + for item in run.graph.node("nz.create").outputs + if item.dtype == "string" + ] + assert ("household", "donor_support_stratum") in text + for node in ("nz.create", "nz.open"): + frame = run.manifest.population(node) + for entity, column in text: + assert frame.table(entity)[column].dtype == pinned, (node, column) diff --git a/test_support/microcosm_build/transport_population.py b/test_support/microcosm_build/transport_population.py new file mode 100644 index 000000000..60a48bd02 --- /dev/null +++ b/test_support/microcosm_build/transport_population.py @@ -0,0 +1,382 @@ +"""Small synthetic donor graphs for the population transport kernels. + +The six made-up US support records come from the frame package's committed +JSON fixture. They are donor support, never New Zealand microdata. All facts, +rates and benefit-unit rules here describe a toy country. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from hashlib import sha256 +from math import fsum +from pathlib import Path + +import numpy as np +import pandas as pd + +from microcosm.build.transport.artifact_types import TARGET_SURFACE_TYPE +from microcosm.build.transport.target_kernels import register_target_kernels +from microcosm.frame import Frame, WeightKind, Weights +from microcosm.frame.adapters.axiom import NZ_SCHEMA +from microcosm.frame.concepts import ( + concept_for_column, + concept_schema_sha256, + split_for_transport, +) +from microcosm.frame.transport import derive_transport_seed, read_populace_us_donor +from microcosm.frame.unit_construction import ( + BenefitUnitRule, + DependentChildRule, + build_benefit_units, +) +from microcosm.graph import ( + ArtifactInput, + ArtifactOutput, + ContentStore, + Graph, + KernelRegistry, + Node, + Owned, + RunManifest, + Slice, + SourceRef, + StructuralDelta, + compile_graph, + run_graph, +) +from microcosm.graph.canonical import canonical_json +from microcosm.graph.population import dtype_for_token +from test_support.microcosm_build.transport_graph import ( + COUNTRY, + canonical_text, + reference_document, + reference_row, + toy_fact, + write_facts, +) +from test_support.paths import paths_for + +SYNTHETIC_DONOR = ( + paths_for("microcosm-frame").tests / "fixtures/transport/synthetic_donor.json" +) +ENTITIES = ("person", "household", "family") +SEED_STREAM = "synthetic:takeup" +UNIT_RULE = BenefitUnitRule( + entity="family", + dependent_child=DependentChildRule(max_age=17, financial_independence=None), + unparented_child="reference_person_unit", + split_parents="refuse", + provenance={"max_age": "Synthetic test rule, not a country's legal rule."}, +) +MASS_DOCUMENT = reference_document( + [reference_row("toy_population", entity="person", measure="person_count")] +) +#: Unequal toy band counts, so a share that ignores or reorders the facts +#: differs from the facts' own shares. Both are exact binary fractions. +BAND_FACTS = (("toy_lower_band", 1.0), ("toy_upper_band", 3.0)) + + +def band_document(entity: str = "person") -> dict: + """References to the band count facts on the ranked ``entity``.""" + return reference_document( + [ + reference_row(name, entity=entity, measure=f"{entity}_count") + for name, _ in BAND_FACTS + ] + ) + + +def digest(document) -> str: + return sha256(canonical_json(document)).hexdigest() + + +def donor_pin(path: Path) -> dict: + return { + "sha256": sha256(path.read_bytes()).hexdigest(), + "size": path.stat().st_size, + } + + +@dataclass(frozen=True) +class PopulationSources: + donor: Path + facts: Path + band_facts: Path + + def mapping(self) -> dict[str, Path]: + return {"donor": self.donor, "facts": self.facts, "band_facts": self.band_facts} + + +def write_population_sources( + root: Path, + *, + mass: float = 120.0, + permutation=None, + band_entity: str = "person", + edit_tables=None, +) -> PopulationSources: + """Write the synthetic donor H5 and toy facts. + + ``edit_tables`` may change the raw donor tables (a dict of DataFrames) in + place before they are written, for a test-local donor. + """ + root.mkdir(parents=True, exist_ok=True) + document = json.loads(SYNTHETIC_DONOR.read_text()) + tables = {entity: pd.DataFrame(rows) for entity, rows in document["tables"].items()} + if edit_tables is not None: + edit_tables(tables) + if permutation is not None: + tables["person"] = ( + tables["person"].iloc[list(permutation)].reset_index(drop=True) + ) + donor = root / "synthetic-donor.h5" + with pd.HDFStore(donor, mode="w") as store: + for entity, table in tables.items(): + store.put(entity, table, format="table") + return PopulationSources( + donor=donor, + facts=write_facts( + root / "mass-facts.jsonl", + [toy_fact("toy_population", mass, entity="person")], + ), + band_facts=write_facts( + root / "band-facts.jsonl", + [toy_fact(name, value, entity=band_entity) for name, value in BAND_FACTS], + ), + ) + + +def expected_bands(sources: PopulationSources, bands=None) -> list[dict]: + """The fixture's band bounds with shares derived from the written facts. + + This reads the facts file directly, independently of the kernels' target + surface: each share is its band's count over the bands' total count. + """ + counts = { + row["semantic_fact_key"].rsplit(":", 1)[-1]: row["value"] + for row in map(json.loads, sources.band_facts.read_text().splitlines()) + } + values = [counts[name] for name, _ in BAND_FACTS] + total = fsum(values) + if bands is None: + bands = json.loads(SYNTHETIC_DONOR.read_text())["target_bands"] + return [ + {"lower": band["lower"], "upper": band["upper"], "share": value / total} + for band, value in zip(bands, values, strict=True) + ] + + +def create_params(sources: PopulationSources, *, seed_stream=SEED_STREAM) -> dict: + pin = donor_pin(sources.donor) + return { + "donor_source": "donor", + "facts_source": "facts", + "donor_sha256": pin["sha256"], + "donor_size": pin["size"], + "donor_pin_sha256": digest(pin), + "unit_rule": canonical_text(UNIT_RULE.to_dict()), + "unit_rule_sha256": digest(UNIT_RULE.to_dict()), + "mass_reference": canonical_text(MASS_DOCUMENT), + "mass_reference_sha256": digest(MASS_DOCUMENT), + "country": COUNTRY, + "concept_schema_sha256": concept_schema_sha256(), + "seed_stream": seed_stream, + } + + +def direct_population( + sources: PopulationSources, *, seed_stream=SEED_STREAM +) -> tuple[Frame, dict]: + """Independent composition of the functions the CREATE wrapper calls.""" + donor = read_populace_us_donor(sources.donor, **donor_pin(sources.donor)) + tables, dropped = split_for_transport(donor.tables) + text = dtype_for_token("string") + for entity, table in tables.items(): + for column in table: + concept = concept_for_column(entity, column) + if concept is not None and concept.dtype == "str": + table[column] = pd.array(table[column], dtype=text) + person, household = tables["person"], tables["household"] + family, membership = build_benefit_units(person, household, UNIT_RULE) + person[UNIT_RULE.membership_column] = membership + person["take_up_seed"] = derive_transport_seed(donor.source_person_ids, seed_stream) + household["donor_support_stratum"] = pd.array(donor.support_strata, dtype=text) + tables["family"] = family + ids = household["household_id"] + mass = json.loads(sources.facts.read_text().splitlines()[0])["value"] + person_weights = person["person_household_id"].map( + pd.Series(donor.weights, index=ids) + ) + scale = mass / fsum(person_weights.to_numpy(dtype=np.float64)) + strata = pd.Series( + person["person_household_id"].map(pd.Series(donor.support_strata, index=ids)), + name="stratum", + dtype=object, + ) + return ( + Frame( + tables, + NZ_SCHEMA, + {"household": Weights(donor.weights * scale, WeightKind.DESIGN)}, + strata, + ), + dropped, + ) + + +def create_node(sources: PopulationSources, *, description="") -> Node: + frame, _ = direct_population(sources) + structural = { + "person": {"person_id", "person_household_id", "person_family_id"}, + "household": {"household_id"}, + "family": {"family_id"}, + } + return Node( + "nz.create", + "transport.create@1", + sources=("donor", "facts"), + structural=StructuralDelta.CREATE, + params=create_params(sources), + outputs=tuple( + Owned(entity, column, str(table[column].dtype)) + for entity in ENTITIES + for table in (frame.table(entity),) + for column in table + if column not in structural[entity] + ), + description=description, + ) + + +def population_graph( + sources: PopulationSources, + *, + rate=1.6, + description="", + through="quantile", + qmap_column="employment_income", + aggregate_entity=None, +) -> Graph: + create = create_node(sources, description=description) + boundary = Node( + "nz.open", + "transport.boundary@1", + structural=StructuralDelta.FILTER, + base=create.id, + inputs=(Slice("person", ("age",)),), + description=description, + ) + nodes = [create, boundary] + if through != "boundary": + # Aggregate maps rank aggregate entities: their band facts count them. + references = band_document( + "person" if aggregate_entity is None else aggregate_entity + ) + bands = Node( + "nz.bands", + "targets.compile@1", + population=boundary.id, + sources=("band_facts",), + params={ + "country": COUNTRY, + "references": canonical_text(references), + "references_sha256": digest(references), + }, + artifact_outputs=(ArtifactOutput("surface", TARGET_SURFACE_TYPE),), + citation=description, + ) + columns = {"person": ["interest_income"], "household": ["rent"]} + currency = Node( + "nz.currency", + "transport.currency@1", + population=boundary.id, + inputs=( + Slice("person", ("interest_income",)), + Slice("household", ("rent",)), + ), + outputs=( + Owned("person", "interest_income", "float64", rewrite=True), + Owned("household", "rent", "float64", rewrite=True), + ), + params={ + "columns": canonical_text(columns), + "rate": rate, + "currency_sha256": digest({"columns": columns, "rate": rate}), + }, + ) + band_rows = [ + {"lower": row["lower"], "upper": row["upper"], "reference": name} + for row, (name, _) in zip( + json.loads(SYNTHETIC_DONOR.read_text())["target_bands"], + BAND_FACTS, + strict=True, + ) + ] + qparams = { + "entity": "person", + "column": qmap_column, + "bands": canonical_text({"bands": band_rows}), + "interpolation": canonical_text({"method": "uniform"}), + "bands_sha256": digest({"bands": band_rows}), + } + if aggregate_entity is not None: + qparams["aggregate_entity"] = aggregate_entity + quantile = Node( + "nz.quantile", + "transport.quantile_map@1", + population=boundary.id, + inputs=(Slice("person", (qmap_column,)),), + outputs=(Owned("person", qmap_column, "float64", rewrite=True),), + params=qparams, + artifact_inputs=( + ArtifactInput("surface", bands.id, "surface", TARGET_SURFACE_TYPE), + ), + ) + nodes.extend([bands, currency, quantile]) + return Graph( + "synthetic-population", + ( + SourceRef("donor", "populace-us-h5-v1"), + SourceRef("facts", "ledger-consumer-artifact-v1"), + SourceRef("band_facts", "ledger-consumer-artifact-v1"), + ), + tuple(nodes), + ) + + +def population_registry() -> KernelRegistry: + from microcosm.build.transport.codecs import register_transport_codecs + from microcosm.build.transport.population_kernels import register_population_kernels + from microcosm.graph import SOURCE_CODECS + + register_transport_codecs(SOURCE_CODECS) + registry = KernelRegistry() + register_population_kernels(registry) + register_target_kernels(registry) + return registry + + +@dataclass(frozen=True) +class PopulationRun: + manifest: RunManifest + store: ContentStore + graph: Graph + + +def run_population( + root: Path, + sources: PopulationSources, + *, + graph=None, + store=None, + registry=None, +) -> PopulationRun: + graph = population_graph(sources) if graph is None else graph + store = ContentStore(root / "store") if store is None else store + registry = population_registry() if registry is None else registry + manifest = run_graph( + compile_graph(graph), sources=sources.mapping(), store=store, kernels=registry + ) + return PopulationRun(manifest, store, graph)