diff --git a/packages/microcosm-build/src/microcosm/build/tail_estimators.py b/packages/microcosm-build/src/microcosm/build/tail_estimators.py new file mode 100644 index 0000000000..ac1ee6cd11 --- /dev/null +++ b/packages/microcosm-build/src/microcosm/build/tail_estimators.py @@ -0,0 +1,375 @@ +"""Declared estimators for an open top band's tail shape. + +``grouped_pareto_mle@1`` fits a Pareto alpha by maximum likelihood to grouped +counts: contiguous anchor bands ``(a_k, b_k]`` with counts ``n_k`` and the +open top above ``c`` with count ``n_T``. With ``L0 = a_1`` and +``S(x) = (x / L0) ** -alpha``, band ``k`` has probability +``p_k = S(a_k) - S(b_k)`` and the top ``p_T = S(c)``, so the log likelihood +is ``l(alpha) = sum n_k ln p_k + n_T ln p_T``. + +Writing ``u_k = ln(a_k / L0)``, ``d_k = ln(b_k / a_k)`` and ``t = ln(c / L0)``, +``ln p_k = -alpha u_k + ln(1 - exp(-alpha d_k))`` and ``ln p_T = -alpha t``. +Each term is concave in alpha, because ``ln(1 - exp(-alpha d))`` has second +derivative ``-d**2 exp(alpha d) / (exp(alpha d) - 1)**2 < 0``. So + +``l'(alpha) = sum n_k (d_k / expm1(alpha d_k) - u_k) - n_T t`` + +is strictly decreasing and has at most one root. The estimator bisects a +declared bracket for a declared number of iterations, rounds the midpoint +once to six decimals (half-even) and carries alpha as ``p / 10**6``. It +refuses, by name, a bracket on which ``l'`` does not change sign and any +alpha at or below one, where the top band's mean ``alpha c / (alpha - 1)`` +is infinite. Under the r2c design's bracket ``[1 + 2**-20, 64]`` the +rounded alpha is at least 1.000001, so ``alpha_not_above_one`` cannot fire +there: heavy-tail data whose root lies at or below one is refused as +``bracket_without_sign_change`` instead, and the refusal names that reason. + +The estimator reads no country data: the caller passes the anchors, their +counts and the declared bracket and iteration count. +""" + +from __future__ import annotations + +import math +import re +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from decimal import Decimal +from fractions import Fraction + +__all__ = [ + "ALPHA_DECIMALS", + "GROUPED_PARETO_MLE_REF", + "MAX_BISECTION_ITERATIONS", + "ParetoAnchor", + "ParetoFit", + "TailEstimatorRefusalError", + "grouped_pareto_mle", + "grouped_pareto_score", + "parse_grouped_pareto_declaration", +] + +GROUPED_PARETO_MLE_REF = "grouped_pareto_mle@1" +#: Alpha is rounded once to this many decimals and carried as p / 10**6. +ALPHA_DECIMALS = 6 +#: Float bisection stops narrowing after about 64 halvings; more is refused +#: so a declaration cannot claim precision the arithmetic cannot give. +MAX_BISECTION_ITERATIONS = 200 + +_DECIMAL_TEXT = re.compile(r"^(0|[1-9][0-9]*)(\.[0-9]+)?$") + + +class TailEstimatorRefusalError(ValueError): + """A named refusal from a tail estimator; ``code`` names the reason.""" + + def __init__(self, code: str, message: str) -> None: + super().__init__(f"{GROUPED_PARETO_MLE_REF} refuses {code}: {message}") + self.code = code + + +@dataclass(frozen=True) +class ParetoAnchor: + """One anchor band: ``(lower, upper]`` or the open top when ``upper`` is None.""" + + reference: str + lower: Fraction + upper: Fraction | None + count: float + + +@dataclass(frozen=True) +class ParetoFit: + """The fitted alpha and everything its receipt records.""" + + numerator: int + denominator: int + anchors: tuple[ParetoAnchor, ...] + bracket: tuple[Fraction, Fraction] + iterations: int + final_bracket: tuple[float, float] + score_at_final_bracket: tuple[float, float] + midpoint: float + + @property + def alpha(self) -> Fraction: + return Fraction(self.numerator, self.denominator) + + @property + def alpha_float(self) -> float: + """The binary64 value nearest ``p / 10**6``, as ``float(text)`` gives.""" + return self.numerator / self.denominator + + @property + def top_lower(self) -> Fraction: + return self.anchors[-1].lower + + @property + def top_band_mean(self) -> Fraction: + """``alpha c / (alpha - 1)``, exact.""" + alpha = self.alpha + return alpha * self.top_lower / (alpha - 1) + + def receipt(self) -> dict[str, object]: + """The deterministic receipt the design lists (r2c design section 2).""" + + def text(value: Fraction) -> str: + return _fraction_text(value) + + return { + "estimator": GROUPED_PARETO_MLE_REF, + "estimator_id": GROUPED_PARETO_MLE_REF.split("@")[0], + "estimator_version": int(GROUPED_PARETO_MLE_REF.split("@")[1]), + "anchor": text(self.anchors[0].lower), + "anchors": [ + { + "reference": anchor.reference, + "lower": text(anchor.lower), + "upper": None if anchor.upper is None else text(anchor.upper), + "count": anchor.count, + } + for anchor in self.anchors + ], + "bracket": [text(self.bracket[0]), text(self.bracket[1])], + "iterations": self.iterations, + "final_bracket": list(self.final_bracket), + "score_at_final_bracket": list(self.score_at_final_bracket), + "alpha": {"numerator": self.numerator, "denominator": self.denominator}, + "alpha_decimal": text(self.alpha), + "top_band_lower": text(self.top_lower), + "top_band_mean": { + "numerator": self.top_band_mean.numerator, + "denominator": self.top_band_mean.denominator, + }, + "top_band_mean_decimal": str( + _round_half_even(self.top_band_mean, ALPHA_DECIMALS) + ), + } + + +def _fraction_text(value: Fraction) -> str: + """Exact decimal text when one exists, else ``p/q``.""" + denominator = value.denominator + for prime in (2, 5): + while denominator % prime == 0: + denominator //= prime + if denominator != 1: + return f"{value.numerator}/{value.denominator}" + decimal = Decimal(value.numerator) / Decimal(value.denominator) + # Exact: the denominator divides a power of ten. + text = format(decimal.normalize(), "f") + return text + + +def _round_half_even(value: Fraction, places: int) -> Decimal: + quantum = Decimal(1).scaleb(-places) + scaled = value * 10**places + floor = scaled.numerator // scaled.denominator + remainder = scaled - floor + half = Fraction(1, 2) + if remainder > half or (remainder == half and floor % 2 == 1): + floor += 1 + return (Decimal(floor) * quantum).quantize(quantum) + + +def _decimal_fraction(value: object, name: str) -> Fraction: + if not isinstance(value, str) or not _DECIMAL_TEXT.match(value): + raise ValueError( + f"{GROUPED_PARETO_MLE_REF} {name} must be nonnegative decimal text." + ) + return Fraction(value) + + +def _edge(value: object, name: str) -> Fraction: + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise ValueError(f"{GROUPED_PARETO_MLE_REF} {name} must be a number.") + if not math.isfinite(value) or value <= 0: + raise ValueError( + f"{GROUPED_PARETO_MLE_REF} {name} must be finite and positive." + ) + return Fraction(value) + + +def parse_grouped_pareto_declaration( + estimator: Mapping[str, object], +) -> tuple[tuple[str, ...], tuple[Fraction, Fraction], int]: + """Validate ``{id, anchor_references, bracket, iterations}``. + + Returns the anchor references, the bracket as exact fractions and the + iteration count. The bracket is two decimal strings, so a declaration + states an exact rational rather than a binary float. + """ + + if not isinstance(estimator, Mapping) or set(estimator) != { + "id", + "anchor_references", + "bracket", + "iterations", + }: + raise ValueError( + "A grouped Pareto estimator declares exactly id, anchor_references, " + "bracket and iterations." + ) + if estimator["id"] != GROUPED_PARETO_MLE_REF: + raise ValueError( + f"The estimator id must be {GROUPED_PARETO_MLE_REF!r}; " + f"{estimator['id']!r} has no implementation." + ) + references = estimator["anchor_references"] + if ( + not isinstance(references, list) + or len(references) < 2 + or any(not isinstance(item, str) or not item for item in references) + or len(set(references)) != len(references) + ): + raise ValueError( + f"{GROUPED_PARETO_MLE_REF} anchor_references must list at least one " + "finite band and the open top, each once." + ) + bracket = estimator["bracket"] + if not isinstance(bracket, list) or len(bracket) != 2: + raise ValueError(f"{GROUPED_PARETO_MLE_REF} bracket must be [low, high].") + low = _decimal_fraction(bracket[0], "bracket low") + high = _decimal_fraction(bracket[1], "bracket high") + if not 0 < low < high: + raise ValueError( + f"{GROUPED_PARETO_MLE_REF} bracket must satisfy 0 < low < high." + ) + iterations = estimator["iterations"] + if ( + isinstance(iterations, bool) + or not isinstance(iterations, int) + or not 1 <= iterations <= MAX_BISECTION_ITERATIONS + ): + raise ValueError( + f"{GROUPED_PARETO_MLE_REF} iterations must be an integer from 1 to " + f"{MAX_BISECTION_ITERATIONS}." + ) + return tuple(references), (low, high), iterations + + +def _anchors( + bands: Sequence[tuple[str, object, object, object]], +) -> tuple[ParetoAnchor, ...]: + if len(bands) < 2: + raise ValueError( + f"{GROUPED_PARETO_MLE_REF} needs at least one finite band and the open top." + ) + anchors = [] + for index, (reference, lower, upper, count) in enumerate(bands): + last = index == len(bands) - 1 + low = _edge(lower, "band lower") + if last: + if upper is not None: + raise ValueError( + f"{GROUPED_PARETO_MLE_REF} fits an open top band; " + f"{reference!r} declares an upper bound." + ) + high = None + else: + if upper is None: + raise ValueError( + f"{GROUPED_PARETO_MLE_REF}: only the top anchor may be open." + ) + high = _edge(upper, "band upper") + if high <= low: + raise ValueError( + f"{GROUPED_PARETO_MLE_REF} band {reference!r} must have " + "upper above lower." + ) + if anchors and anchors[-1].upper != low: + raise ValueError( + f"{GROUPED_PARETO_MLE_REF} anchor bands must be contiguous and " + "increasing." + ) + if ( + isinstance(count, bool) + or not isinstance(count, (int, float)) + or not math.isfinite(count) + or count < 0 + ): + raise ValueError( + f"{GROUPED_PARETO_MLE_REF} counts must be finite and nonnegative." + ) + anchors.append(ParetoAnchor(reference, low, high, float(count))) + return tuple(anchors) + + +def _terms( + anchors: Sequence[ParetoAnchor], +) -> tuple[list[tuple[float, float, float]], float, float]: + """``(n_k, u_k, d_k)`` per finite band, then ``n_T`` and ``t``.""" + origin = anchors[0].lower + finite = [ + ( + anchor.count, + math.log(anchor.lower / origin), + math.log(anchor.upper / anchor.lower), + ) + for anchor in anchors[:-1] + ] + top = anchors[-1] + return finite, top.count, math.log(top.lower / origin) + + +def grouped_pareto_score(anchors: Sequence[ParetoAnchor], alpha: float) -> float: + """``l'(alpha)`` for the anchors, as the module docstring writes it.""" + finite, top_count, top_offset = _terms(anchors) + parts = [ + count * (width / math.expm1(alpha * width) - offset) + for count, offset, width in finite + ] + parts.append(-top_count * top_offset) + return math.fsum(parts) + + +def grouped_pareto_mle( + bands: Sequence[tuple[str, object, object, object]], + *, + bracket: tuple[Fraction, Fraction], + iterations: int, +) -> ParetoFit: + """Fit alpha to ``(reference, lower, upper, count)`` anchor rows. + + The rows are contiguous and increasing; the last is the open top + (``upper`` is None). ``bracket`` and ``iterations`` come from the node's + declaration. + """ + + anchors = _anchors(bands) + low, high = (float(bracket[0]), float(bracket[1])) + score_low = grouped_pareto_score(anchors, low) + score_high = grouped_pareto_score(anchors, high) + if not (score_low > 0 > score_high): + raise TailEstimatorRefusalError( + "bracket_without_sign_change", + f"l'(alpha) is {score_low!r} at {_fraction_text(bracket[0])} and " + f"{score_high!r} at {_fraction_text(bracket[1])}; the root is not " + "inside the declared bracket.", + ) + for _ in range(iterations): + middle = (low + high) / 2 + score = grouped_pareto_score(anchors, middle) + if score > 0: + low, score_low = middle, score + else: + high, score_high = middle, score + midpoint = (low + high) / 2 + rounded = _round_half_even(Fraction(midpoint), ALPHA_DECIMALS) + denominator = 10**ALPHA_DECIMALS + numerator = int(rounded.scaleb(ALPHA_DECIMALS)) + if numerator <= denominator: + raise TailEstimatorRefusalError( + "alpha_not_above_one", + f"the fitted alpha {rounded} is at most one, so the open top band " + "has no finite mean.", + ) + return ParetoFit( + numerator=numerator, + denominator=denominator, + anchors=anchors, + bracket=bracket, + iterations=iterations, + final_bracket=(low, high), + score_at_final_bracket=(score_low, score_high), + midpoint=midpoint, + ) diff --git a/packages/microcosm-build/src/microcosm/build/transport/compose.py b/packages/microcosm-build/src/microcosm/build/transport/compose.py index 0fef19eef3..7bc8799593 100644 --- a/packages/microcosm-build/src/microcosm/build/transport/compose.py +++ b/packages/microcosm-build/src/microcosm/build/transport/compose.py @@ -181,6 +181,7 @@ "transport.boundary@1", "transport.currency@1", "transport.quantile_map@1", + "transport.quantile_map@2", MONEY_NORMALIZE_REF, "transport.unit_attributes@1", "concepts.encode@1", @@ -215,7 +216,8 @@ MONEY_NORMALIZATION_RESOURCE = "money_normalization" #: The only kernels that may produce donor money after CREATE once a spec #: declares money normalization. -_MONEY_PRODUCERS = frozenset({MONEY_NORMALIZE_REF, "transport.quantile_map@1"}) +_QUANTILE_MAPS = frozenset({"transport.quantile_map@1", "transport.quantile_map@2"}) +_MONEY_PRODUCERS = frozenset({MONEY_NORMALIZE_REF}) | _QUANTILE_MAPS #: Money bridges own new integer columns in their whole population lineage. _MONEY_BRIDGES = frozenset({MONEY_MINOR_UNITS_REF, MONEY_ANNUAL_TRANSFERS_REF}) #: Conversion changes the basis of an existing donor column; derived money @@ -842,10 +844,10 @@ def _money_producer_gaps(spec: Mapping[str, object]) -> list[str]: and reads nothing. Otherwise every donor money concept a node reads after CREATE needs exactly one producer in its population lineage (its own version and each base version): ``money.normalize@1`` for a - normalized column, ``transport.quantile_map@1`` for a quantile-mapped - one. The concept's default transport rule is not consulted. A consumed - pending or quantile-mapped column without its producer names what it - awaits. New destination money outputs are classified separately. + normalized column, ``transport.quantile_map@1`` or ``@2`` for a + quantile-mapped one. The concept's default transport rule is not + consulted. A consumed pending or quantile-mapped column without its + producer names what it awaits. New destination money outputs are classified separately. """ if not _has_resource(spec, MONEY_NORMALIZATION_RESOURCE): return [] @@ -941,16 +943,17 @@ def lineage(version: object) -> list[str]: ): continue expected = { - "normalized": MONEY_NORMALIZE_REF, - "quantile_mapped": "transport.quantile_map@1", - }.get(route) + "normalized": frozenset({MONEY_NORMALIZE_REF}), + "quantile_mapped": _QUANTILE_MAPS, + }.get(route, frozenset()) if kernel not in _MONEY_PRODUCERS: message = ( f"{name} is produced by {kernel} node " - f"{producer.get('id')!r}; only {MONEY_NORMALIZE_REF} or " - "transport.quantile_map@1 may produce it after CREATE" + f"{producer.get('id')!r}; only {MONEY_NORMALIZE_REF}, " + "transport.quantile_map@1 or transport.quantile_map@2 " + "may produce it after CREATE" ) - elif kernel != expected: + elif kernel not in expected: message = ( f"{name} is declared {route} but produced by {kernel} " f"node {producer.get('id')!r}" @@ -1013,6 +1016,7 @@ def _check_single_money_normalization( expected = { MONEY_NORMALIZE_REF: "normalized", "transport.quantile_map@1": "quantile_mapped", + "transport.quantile_map@2": "quantile_mapped", } by_id = {node.id: node for node in graph.nodes} for (_, entity, column), owner in compiled.owners.items(): @@ -1629,6 +1633,7 @@ def _check_weights_audit_entity(spec: Mapping[str, object], graph: Graph) -> Non { MONEY_NORMALIZE_REF: "normalized", "transport.quantile_map@1": "quantile_mapped", + "transport.quantile_map@2": "quantile_mapped", "transport.currency@1": "currency", } ) diff --git a/packages/microcosm-build/src/microcosm/build/transport/quantile_kernels.py b/packages/microcosm-build/src/microcosm/build/transport/quantile_kernels.py new file mode 100644 index 0000000000..828fb1e651 --- /dev/null +++ b/packages/microcosm-build/src/microcosm/build/transport/quantile_kernels.py @@ -0,0 +1,158 @@ +"""``transport.quantile_map@2``: a quantile map whose open top fits its own alpha. + +The kernel estimates the Pareto alpha of the open top band from the declared +anchor facts on the same compiled band surface the map's shares come from, +then maps exactly as ``transport.quantile_map@1`` does with that alpha as a +literal. It delegates the map to ``@1``'s own ``run``, so no second +interpolation path exists, and it lives in its own module because +``source_hash`` digests a kernel's whole defining module: a class added to +``population_kernels`` would re-key ``@1``. + +The interpolation parameter is ``{"method": "pareto", "estimator": {"id": +"grouped_pareto_mle@1", "anchor_references": [...], "bracket": [low, high], +"iterations": n}}``. The anchor references are the top contiguous suffix of +the band list, ending with the open top band, and their counts are the +surface facts the band rows reference. The receipt records the fit. +""" + +from __future__ import annotations + +from dataclasses import replace + +import microcosm.build.tail_estimators as tail_estimators_module +import microcosm.frame.transport as transport_module +from microcosm.build.tail_estimators import ( + grouped_pareto_mle, + parse_grouped_pareto_declaration, +) +from microcosm.graph import ( + KernelBase, + KernelContext, + KernelRegistry, + KernelResult, + source_hash, +) +from microcosm.graph.canonical import canonical_json + +from . import graph_inputs, population_kernels, target_kernels +from .artifact_types import TARGET_SURFACE_TYPE +from .graph_inputs import canonical_document_param, require_params, string_param +from .population_kernels import TRANSPORT_QUANTILE_MAP, _surface_bands +from .target_kernels import decode_target_surface + +__all__ = [ + "QUANTILE_MAP_V2_REF", + "TRANSPORT_QUANTILE_MAP_V2", + "TransportQuantileMapV2Kernel", + "estimated_alpha", + "register_quantile_kernels", +] + +QUANTILE_MAP_V2_REF = "transport.quantile_map@2" + + +def estimated_alpha(context: KernelContext, ref: str = QUANTILE_MAP_V2_REF): + """Fit the declared estimator to the anchor counts on the node's surface. + + Returns the fit and the literal interpolation document ``@1`` reads. + """ + + interpolation = canonical_document_param(context, ref, "interpolation") + if set(interpolation) != {"method", "estimator"} or ( + interpolation["method"] != "pareto" + ): + raise ValueError( + f"{ref} interpolation is {{method: pareto, estimator: ...}}; a literal " + "alpha belongs to transport.quantile_map@1." + ) + references, bracket, iterations = parse_grouped_pareto_declaration( + interpolation["estimator"] + ) + document = canonical_document_param(context, ref, "bands") + if set(document) != {"bands"}: + raise ValueError( + f"{ref} fits one alpha to a bare band list; signed or component " + "bands are refused." + ) + artifact = context.artifacts.get("surface") + if artifact is None or artifact.type != TARGET_SURFACE_TYPE: + raise ValueError(f"{ref} requires a target surface under alias 'surface'.") + surface = decode_target_surface(artifact.payload) + entity = string_param(context, ref, "entity") + unit = ( + string_param(context, ref, "aggregate_entity") + if "aggregate_entity" in context.params + else entity + ) + # The same validation and the same facts the map's shares use. + _surface_bands(document, surface, unit=unit) + rows = document["bands"] + if [row["reference"] for row in rows[-len(references) :]] != list(references): + raise ValueError( + f"{ref} anchor_references must be the top contiguous bands of the " + "band list, in order, ending with the open top." + ) + specs = {spec.name: spec for spec in surface.registry.specs} + fit = grouped_pareto_mle( + [ + ( + row["reference"], + row["lower"], + row["upper"], + specs[row["reference"]].value, + ) + for row in rows[-len(references) :] + ], + bracket=bracket, + iterations=iterations, + ) + literal = canonical_json({"method": "pareto", "alpha": fit.alpha_float}).decode( + "utf-8" + ) + return fit, literal + + +class TransportQuantileMapV2Kernel(KernelBase): + """``transport.quantile_map@2``: ``@1`` with a fitted open-top alpha.""" + + ref = QUANTILE_MAP_V2_REF + capabilities = TRANSPORT_QUANTILE_MAP.capabilities + + def implementation_hash(self) -> str: + return source_hash( + type(self), + tail_estimators_module, + population_kernels, + 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"}), + ) + fit, literal = estimated_alpha(context, self.ref) + delegated = replace( + context, params={**context.params, "interpolation": literal} + ) + result = TRANSPORT_QUANTILE_MAP.run(delegated) + return replace( + result, receipt={**result.receipt, "tail_estimator": fit.receipt()} + ) + + +TRANSPORT_QUANTILE_MAP_V2 = TransportQuantileMapV2Kernel() + + +def register_quantile_kernels(registry: KernelRegistry) -> KernelRegistry: + """Register ``transport.quantile_map@2`` explicitly and idempotently.""" + registry.register(TRANSPORT_QUANTILE_MAP_V2) + return registry diff --git a/packages/microcosm-build/src/microcosm/build/transport/registry.py b/packages/microcosm-build/src/microcosm/build/transport/registry.py index 30a15db350..ab434ae026 100644 --- a/packages/microcosm-build/src/microcosm/build/transport/registry.py +++ b/packages/microcosm-build/src/microcosm/build/transport/registry.py @@ -53,6 +53,7 @@ from .money_kernels import register_money_kernels from .person_location_kernels import register_person_location_kernels from .population_kernels import register_population_kernels +from .quantile_kernels import register_quantile_kernels from .reporting_kernels import register_reporting_kernels from .takeup_kernels import register_takeup_kernels from .target_kernels import register_target_kernels @@ -531,13 +532,14 @@ def build_transport_registry( def register_transport_population_kernels(registry: KernelRegistry) -> KernelRegistry: - """Install the thirteen population kernels without import-time side effects. + """Install the fifteen 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_quantile_kernels(registry) register_column_kernels(registry) register_geography_kernels(registry) register_money_kernels(registry) diff --git a/packages/microcosm-build/tests/engine_free/shared/test_transport_population_kernels.py b/packages/microcosm-build/tests/engine_free/shared/test_transport_population_kernels.py index a91b6a9f02..3cf21a60b8 100644 --- a/packages/microcosm-build/tests/engine_free/shared/test_transport_population_kernels.py +++ b/packages/microcosm-build/tests/engine_free/shared/test_transport_population_kernels.py @@ -140,7 +140,7 @@ def test_kernel_registration_is_explicit_idempotent_and_country_neutral(): ) -def test_aggregate_population_registration_installs_exactly_fourteen_kernels(): +def test_aggregate_population_registration_installs_exactly_fifteen_kernels(): registry = KernelRegistry() assert register_transport_population_kernels(registry) is registry assert register_transport_population_kernels(registry) is registry @@ -153,6 +153,7 @@ def test_aggregate_population_registration_installs_exactly_fourteen_kernels(): "transport.boundary@1", "transport.currency@1", "transport.quantile_map@1", + "transport.quantile_map@2", "transport.unit_attributes@1", "concepts.encode@1", "concepts.encode_groups@1", diff --git a/packages/microcosm-build/tests/engine_free/shared/test_transport_quantile_map_v2.py b/packages/microcosm-build/tests/engine_free/shared/test_transport_quantile_map_v2.py new file mode 100644 index 0000000000..d60fbecc55 --- /dev/null +++ b/packages/microcosm-build/tests/engine_free/shared/test_transport_quantile_map_v2.py @@ -0,0 +1,806 @@ +"""``grouped_pareto_mle@1`` and ``transport.quantile_map@2`` (G8-r2c hunk B). + +The estimator fits a Pareto alpha to grouped counts; ``@2`` fits it to the +anchor facts on its own band surface and then maps exactly as ``@1`` does +with that alpha as a literal. Every count here is synthetic except +Belgium's two RSZ 2026-Q1 counts, which the hub's ruling R2C-9 quotes. +""" + +from __future__ import annotations + +import copy +import json +import math +from decimal import Decimal +from fractions import Fraction +from pathlib import Path +from tempfile import TemporaryDirectory + +import numpy as np +import pandas as pd +import pytest +from hypothesis import HealthCheck, assume, given, settings +from hypothesis import strategies as st + +from microcosm.build.tail_estimators import ( + GROUPED_PARETO_MLE_REF, + TailEstimatorRefusalError, + _anchors, + _round_half_even, + grouped_pareto_mle, + grouped_pareto_score, + parse_grouped_pareto_declaration, +) +from microcosm.build.transport import compose +from microcosm.build.transport.artifact_types import TARGET_SURFACE_TYPE +from microcosm.build.transport.evidence_contract import ( + DONOR_RANGES, + DONOR_RANGES_REF, + DONOR_RANGES_TYPE, + DONOR_SUPPORT_BOUNDS, + SUPPORT_BOUNDS_REF, + SUPPORT_BOUNDS_TYPE, + decode_support_bounds, + quantile_bounds, + quantile_support, +) +from microcosm.build.transport.evidence_kernels import ( + TRANSPORT_DONOR_RANGES, + TRANSPORT_SUPPORT_BOUNDS, +) +from microcosm.build.transport.population_kernels import TRANSPORT_QUANTILE_MAP +from microcosm.build.transport.quantile_kernels import ( + QUANTILE_MAP_V2_REF, + TRANSPORT_QUANTILE_MAP_V2, +) +from microcosm.build.transport.registry import register_transport_population_kernels +from microcosm.build.transport.target_kernels import compile_target_surface +from microcosm.frame import EntitySchema, Frame, WeightKind, Weights +from microcosm.graph import ( + ArtifactInput, + ArtifactOutput, + KernelRegistry, + Node, + Owned, + Slice, +) +from microcosm.graph.canonical import canonical_json +from microcosm.graph.executor import Population, _project_context +from microcosm.graph.kernel import ArtifactValue, NumericScope +from test_support.microcosm_build.transport_graph import ( + COUNTRY, + reference_document, + reference_row, + toy_fact, + write_facts, +) +from test_support.microcosm_build.transport_population import digest + +#: The bracket the r2c design declares, [1 + 2**-20, 64], as exact text. +BRACKET_TEXT = ["1.00000095367431640625", "64"] +BRACKET = (Fraction(BRACKET_TEXT[0]), Fraction(BRACKET_TEXT[1])) +ITERATIONS = 64 +#: Recovery tolerance on exact expected counts. Rounding to six decimals +#: contributes at most 5e-7; 64 halvings of [1 + 2**-20, 64] leave a bracket +#: of 63 / 2**64 (about 3e-18); and the binary64 score's root moves by its +#: summation error over l'' (about 1e-13 relative here). 1e-6 covers the +#: three with room, and no closer tolerance survives the half-unit rounding. +TOLERANCE = 1e-6 + +PROPERTY = settings( + max_examples=60, + deadline=None, + suppress_health_check=(HealthCheck.function_scoped_fixture,), +) + + +def _fit(rows, *, bracket=BRACKET, iterations=ITERATIONS): + return grouped_pareto_mle(rows, bracket=bracket, iterations=iterations) + + +def _expected_rows(alpha, edges, total=1e6): + """Exact Pareto expected counts on contiguous ``edges`` plus the open top.""" + origin = edges[0] + + def survival(x): + return (x / origin) ** -alpha + + rows = [ + (f"b{index}", lower, upper, total * (survival(lower) - survival(upper))) + for index, (lower, upper) in enumerate(zip(edges, edges[1:], strict=False)) + ] + rows.append(("top", edges[-1], None, total * survival(edges[-1]))) + return rows + + +@st.composite +def anchors(draw): + origin = draw(st.floats(min_value=1, max_value=1e6)) + ratios = draw( + st.lists(st.floats(min_value=1.01, max_value=3), min_size=1, max_size=10) + ) + edges = [origin] + for ratio in ratios: + edges.append(edges[-1] * ratio) + return edges + + +# --------------------------------------------------------------------------- +# The estimator: examples +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("alpha", [1.2, 1.5, 2.093, 3.348955, 5.0]) +def test_recovers_a_known_alpha_from_exact_pareto_counts(alpha) -> None: + edges = [500.0, 550.0, 600.0, 650.0, 700.0, 750.0, 800.0, 1000.0] + fit = _fit(_expected_rows(alpha, edges)) + assert abs(fit.alpha_float - alpha) <= TOLERANCE + assert fit.denominator == 10**6 + assert fit.alpha == Fraction(fit.numerator, 10**6) + + +def test_refuses_a_bracket_without_a_sign_change_by_name() -> None: + rows = _expected_rows(1.5, [100.0, 200.0, 400.0]) + with pytest.raises(TailEstimatorRefusalError, match="bracket_without_sign_change"): + _fit(rows, bracket=(Fraction(2), Fraction(64))) + with pytest.raises(TailEstimatorRefusalError) as caught: + _fit(rows, bracket=(Fraction(1, 10), Fraction(1))) + assert caught.value.code == "bracket_without_sign_change" + + +def test_refuses_an_alpha_at_or_below_one_by_name() -> None: + rows = _expected_rows(0.8, [100.0, 200.0, 400.0]) + with pytest.raises(TailEstimatorRefusalError) as caught: + _fit(rows, bracket=(Fraction(1, 2), Fraction(64))) + assert caught.value.code == "alpha_not_above_one" + # Alpha exactly one after rounding is refused too: its mean is infinite. + rows = _expected_rows(1.0000001, [100.0, 200.0, 400.0]) + with pytest.raises(TailEstimatorRefusalError, match="alpha_not_above_one"): + _fit(rows, bracket=(Fraction(1, 2), Fraction(64))) + + +def test_degenerate_counts_fall_out_as_no_sign_change() -> None: + no_top = [("a", 100.0, 200.0, 10.0), ("t", 200.0, None, 0.0)] + no_body = [("a", 100.0, 200.0, 0.0), ("t", 200.0, None, 10.0)] + for rows in (no_top, no_body): + with pytest.raises(TailEstimatorRefusalError, match="no_sign_change|sign"): + _fit(rows) + + +@pytest.mark.parametrize( + ("rows", "message"), + [ + ([("t", 100.0, None, 1.0)], "at least one finite band"), + ([("a", 100.0, 200.0, 1.0), ("t", 200.0, 300.0, 1.0)], "open top"), + ([("a", 100.0, None, 1.0), ("t", 200.0, None, 1.0)], "only the top"), + ([("a", 100.0, 200.0, 1.0), ("t", 250.0, None, 1.0)], "contiguous"), + ([("a", 0.0, 200.0, 1.0), ("t", 200.0, None, 1.0)], "positive"), + ([("a", 100.0, 200.0, -1.0), ("t", 200.0, None, 1.0)], "nonnegative"), + ], +) +def test_refuses_malformed_anchors(rows, message) -> None: + with pytest.raises(ValueError, match=message): + _fit(rows) + + +def test_the_declaration_is_parsed_exactly() -> None: + declared = { + "id": GROUPED_PARETO_MLE_REF, + "anchor_references": ["a", "t"], + "bracket": BRACKET_TEXT, + "iterations": ITERATIONS, + } + references, bracket, iterations = parse_grouped_pareto_declaration(declared) + assert references == ("a", "t") + assert bracket == (1 + Fraction(1, 2**20), Fraction(64)) + assert iterations == ITERATIONS + for change, message in ( + ({"id": "grouped_pareto_mle@2"}, "no implementation"), + ({"bracket": [1.5, 64]}, "decimal text"), + ({"bracket": ["64", "2"]}, "low < high"), + ({"iterations": 0}, "iterations"), + ({"iterations": 201}, "iterations"), + ({"anchor_references": ["t"]}, "at least one"), + ): + with pytest.raises(ValueError, match=message): + parse_grouped_pareto_declaration({**declared, **change}) + with pytest.raises(ValueError, match="exactly id"): + parse_grouped_pareto_declaration({**declared, "alpha": 2}) + + +def test_the_receipt_records_the_fit() -> None: + fit = _fit([("a", 200, 300, 1_066_243), ("t", 300, None, 369_203)]) + receipt = fit.receipt() + assert receipt["estimator"] == GROUPED_PARETO_MLE_REF + assert receipt["estimator_id"] == "grouped_pareto_mle" + assert receipt["estimator_version"] == 1 + assert receipt["anchor"] == "200" + assert receipt["anchors"] == [ + {"reference": "a", "lower": "200", "upper": "300", "count": 1_066_243.0}, + {"reference": "t", "lower": "300", "upper": None, "count": 369_203.0}, + ] + assert receipt["bracket"] == BRACKET_TEXT + assert receipt["iterations"] == ITERATIONS + low, high = receipt["final_bracket"] + assert low <= fit.midpoint <= high + score_low, score_high = receipt["score_at_final_bracket"] + assert score_low > 0 >= score_high + assert receipt["alpha"] == {"numerator": 3_348_955, "denominator": 10**6} + assert receipt["alpha_decimal"] == "3.348955" + mean = Fraction(3_348_955, 10**6) * 300 / (Fraction(3_348_955, 10**6) - 1) + assert ( + Fraction( + receipt["top_band_mean"]["numerator"], + receipt["top_band_mean"]["denominator"], + ) + == mean + ) + assert receipt["top_band_mean_decimal"] == "427.716367" + json.dumps(receipt, allow_nan=False) + + +# --------------------------------------------------------------------------- +# Belgium's minimal anchor set: the closed form +# --------------------------------------------------------------------------- + + +def _closed_form(n_band, n_top, lower, upper): + return math.log((n_band + n_top) / n_top) / math.log(upper / lower) + + +def test_belgium_rsz_counts_give_the_closed_form_alpha() -> None: + n_band, n_top = 1_066_243, 369_203 + fit = _fit([("rsz_200_300", 200, 300, n_band), ("rsz_over_300", 300, None, n_top)]) + assert fit.receipt()["alpha_decimal"] == "3.348955" + assert f"{_closed_form(n_band, n_top, 200, 300):.6f}" == "3.348955" + + +def test_six_decimal_rounding_is_half_even_at_an_exact_tie() -> None: + """A tie goes to the even sixth decimal; half-up would differ on the first.""" + + # 1.0078125 is a binary64 value lying exactly halfway between 1.007812 + # and 1.007813: the even neighbour is 1.007812. + assert _round_half_even(Fraction("1.0078125"), 6) == Decimal("1.007812") + # Halfway between 1.007813 and 1.007814: the even neighbour is 1.007814. + assert _round_half_even(Fraction("1.0078135"), 6) == Decimal("1.007814") + # Off a tie, the nearer neighbour wins either way. + assert _round_half_even(Fraction("1.00781249"), 6) == Decimal("1.007812") + assert _round_half_even(Fraction("1.00781251"), 6) == Decimal("1.007813") + + +@PROPERTY +@given( + n_band=st.integers(min_value=1, max_value=10**8), + n_top=st.integers(min_value=1, max_value=10**8), + lower=st.floats(min_value=1, max_value=1e6), + ratio=st.floats(min_value=1.05, max_value=10), +) +def test_property_one_band_and_the_top_match_the_closed_form( + n_band, n_top, lower, ratio +) -> None: + upper = lower * ratio + expected = _closed_form(n_band, n_top, lower, upper) + # Hypothesis replaces a draw whose root lies outside the bracket, so every + # counted example reaches the assertion. + assume(1.01 < expected < 63) + fit = _fit([("a", lower, upper, n_band), ("t", upper, None, n_top)]) + assert abs(fit.alpha_float - expected) <= TOLERANCE + + +# --------------------------------------------------------------------------- +# The estimator: properties +# --------------------------------------------------------------------------- + + +@PROPERTY +@given(alpha=st.floats(min_value=1.2, max_value=5), edges=anchors()) +def test_property_recovers_alpha_within_the_declared_tolerance(alpha, edges) -> None: + fit = _fit(_expected_rows(alpha, edges)) + assert abs(fit.alpha_float - alpha) <= TOLERANCE + + +@PROPERTY +@given( + alpha=st.floats(min_value=1.2, max_value=5), + edges=anchors(), + k=st.integers(min_value=2, max_value=10**6), + power=st.integers(min_value=-20, max_value=20), +) +def test_property_scaling_every_count_leaves_alpha_unchanged( + alpha, edges, k, power +) -> None: + rows = _expected_rows(alpha, edges) + base = _fit(rows) + # A power of two scales every float exactly, so the fit is bitwise equal. + exact = _fit([(r, lo, up, n * 2.0**power) for r, lo, up, n in rows]) + assert exact.numerator == base.numerator + # Any other factor rounds each count once; alpha moves by at most a unit + # in the sixth decimal, and only at a rounding tie. + scaled = _fit([(r, lo, up, n * k) for r, lo, up, n in rows]) + assert abs(scaled.numerator - base.numerator) <= 1 + + +@PROPERTY +@given( + alpha=st.floats(min_value=1.2, max_value=5), + edges=anchors(), + data=st.data(), +) +def test_property_moving_mass_upward_lowers_alpha(alpha, edges, data) -> None: + rows = _expected_rows(alpha, edges) + source = data.draw(st.integers(min_value=0, max_value=len(rows) - 2)) + target = data.draw(st.integers(min_value=source + 1, max_value=len(rows) - 1)) + share = data.draw(st.floats(min_value=0.01, max_value=0.99)) + moved = rows[source][3] * share + counts = [row[3] for row in rows] + counts[source] -= moved + counts[target] += moved + shifted = [(r, lo, up, n) for (r, lo, up, _), n in zip(rows, counts, strict=True)] + before = _fit(rows) + shifted_anchors = _anchors(shifted) + # Before rounding, the move lowers the score at every alpha, so the root + # moves down. + assert grouped_pareto_score(shifted_anchors, before.midpoint) < 0 + try: + after = _fit(shifted) + except TailEstimatorRefusalError as refusal: + # The root fell below the bracket: alpha is lower still. + assert refusal.code in {"bracket_without_sign_change", "alpha_not_above_one"} + assert grouped_pareto_score(shifted_anchors, float(BRACKET[0])) <= 0 + return + assert after.numerator <= before.numerator + + +@PROPERTY +@given( + counts=st.lists( + st.floats(min_value=0, max_value=1e7), min_size=2, max_size=8 + ).filter(lambda c: c[-1] > 0 and sum(c[:-1]) > 0), + edges=anchors(), +) +def test_property_the_score_changes_sign_once_on_the_bracket(counts, edges) -> None: + size = min(len(counts), len(edges)) + counts, edges = counts[-size:], edges[:size] + rows = [ + (f"b{index}", lower, upper, count) + for index, (lower, upper, count) in enumerate( + zip(edges, edges[1:], counts, strict=False) + ) + ] + rows.append(("top", edges[-1], None, counts[-1])) + assume(len(rows) >= 2 and sum(row[3] for row in rows[:-1]) > 0) + fit_anchors = _anchors(rows) + grid = np.linspace(float(BRACKET[0]), float(BRACKET[1]), 257) + scores = [grouped_pareto_score(fit_anchors, float(a)) for a in grid] + # Nonincreasing (strictly, until expm1 overflows to infinity), hence + # at most one sign change. + assert all(b <= a for a, b in zip(scores, scores[1:], strict=False)) + signs = [s > 0 for s in scores] + assert sum(1 for a, b in zip(signs, signs[1:], strict=False) if a != b) <= 1 + + +# --------------------------------------------------------------------------- +# The kernel +# --------------------------------------------------------------------------- + +SCHEMA = EntitySchema(group_entities=("household",)) +#: Synthetic band rows: a uniform body, two anchor bands and the open top. +BAND_ROWS = ( + ("toy_wage_0", 1.0, 1000.0), + ("toy_wage_1", 1000.0, 2000.0), + ("toy_wage_2", 2000.0, 3000.0), + ("toy_wage_top", 3000.0, None), +) +COUNTS = {"toy_wage_0": 500.0, "toy_wage_1": 300.0, "toy_wage_2": 120.0} +ANCHORS = ["toy_wage_1", "toy_wage_2", "toy_wage_top"] + + +def _frame(values) -> Frame: + count = len(values) + person = pd.DataFrame( + { + "person_id": np.arange(1, count + 1, dtype=np.int64), + "person_household_id": np.arange(1, count + 1, dtype=np.int64), + "employment_income": np.asarray(values, dtype=np.float64), + } + ) + household = pd.DataFrame({"household_id": np.arange(1, count + 1, dtype=np.int64)}) + weights = np.linspace(1, 3, count) + return Frame( + {"person": person, "household": household}, + SCHEMA, + {"household": Weights(weights, WeightKind.DESIGN)}, + ) + + +def _surface(root: Path, counts: dict[str, float]) -> ArtifactValue: + document = reference_document( + [ + reference_row(name, entity="person", measure="person_count") + for name, _, _ in BAND_ROWS + ] + ) + facts = write_facts( + root / "band-facts.jsonl", + [toy_fact(name, counts[name], entity="person") for name, _, _ in BAND_ROWS], + ) + payload, _ = compile_target_surface( + facts, document, country=COUNTRY, references_sha256=digest(document) + ) + return ArtifactValue( + payload=payload, + type=TARGET_SURFACE_TYPE, + key="a" * 64, + producer_key="b" * 64, + numerics=NumericScope(), + ) + + +def _bands(): + return { + "bands": [ + {"lower": lower, "upper": upper, "reference": name} + for name, lower, upper in BAND_ROWS + ] + } + + +def _estimator(**change): + return { + "id": GROUPED_PARETO_MLE_REF, + "anchor_references": ANCHORS, + "bracket": BRACKET_TEXT, + "iterations": ITERATIONS, + **change, + } + + +def _node(kernel, interpolation, bands=None) -> Node: + bands = _bands() if bands is None else bands + return Node( + "toy.wage.map", + kernel, + inputs=(Slice("person", ("employment_income",)),), + outputs=(Owned("person", "employment_income", "float64", rewrite=True),), + params={ + "entity": "person", + "column": "employment_income", + "bands": canonical_json(bands).decode(), + "interpolation": canonical_json(interpolation).decode(), + "bands_sha256": digest(bands), + }, + artifact_inputs=( + ArtifactInput("surface", "toy.bands", "surface", TARGET_SURFACE_TYPE), + ), + ) + + +def _run(kernel, node, frame, surface): + context = _project_context( + node, + Population.from_frame(frame, "toy.population"), + key="0" * 64, + sources={}, + tolerances={}, + numerics={}, + artifacts={"surface": surface}, + ) + return kernel.run(context) + + +def _counts(top=60.0, **change): + return {**COUNTS, "toy_wage_top": top, **change} + + +def _values(seed=7, size=40): + rng = np.random.default_rng(seed) + values = np.round(rng.lognormal(7, 1, size), 2) + values[::9] = 0 + return values + + +def _pareto(estimator=None): + return {"method": "pareto", "estimator": estimator or _estimator()} + + +def test_v2_equals_v1_with_its_estimated_alpha_as_a_literal(tmp_path) -> None: + frame, surface = _frame(_values()), _surface(tmp_path, _counts()) + v2 = _run( + TRANSPORT_QUANTILE_MAP_V2, + _node(QUANTILE_MAP_V2_REF, _pareto()), + frame, + surface, + ) + decimal = v2.receipt["tail_estimator"]["alpha_decimal"] + literal = {"method": "pareto", "alpha": float(decimal)} + v1 = _run( + TRANSPORT_QUANTILE_MAP, + _node("transport.quantile_map@1", literal), + frame, + surface, + ) + ((key, mapped),) = v2.columns.items() + assert key == ("person", "employment_income") + assert mapped.to_numpy().tobytes() == v1.columns[key].to_numpy().tobytes() + pd.testing.assert_index_equal(mapped.index, v1.columns[key].index) + + +@PROPERTY +@given( + seed=st.integers(min_value=0, max_value=2**32 - 1), + top=st.floats(min_value=5, max_value=400), + middle=st.floats(min_value=20, max_value=800), +) +def test_property_v2_is_bitwise_v1_at_the_fitted_alpha(tmp_path, seed, top, middle): + with TemporaryDirectory(dir=tmp_path) as temporary: + frame = _frame(_values(seed)) + surface = _surface(Path(temporary), _counts(top, toy_wage_2=middle)) + try: + v2 = _run( + TRANSPORT_QUANTILE_MAP_V2, + _node(QUANTILE_MAP_V2_REF, _pareto()), + frame, + surface, + ) + except TailEstimatorRefusalError: + # Hypothesis replaces a draw the estimator refuses, so every + # counted example reaches the differential. + assume(False) + alpha = v2.receipt["tail_estimator"]["alpha"] + literal = { + "method": "pareto", + "alpha": alpha["numerator"] / alpha["denominator"], + } + v1 = _run( + TRANSPORT_QUANTILE_MAP, + _node("transport.quantile_map@1", literal), + frame, + surface, + ) + key = ("person", "employment_income") + assert v2.columns[key].to_numpy().tobytes() == ( + v1.columns[key].to_numpy().tobytes() + ) + + +def test_anchor_counts_are_the_band_rows_counts(tmp_path) -> None: + """One source: the anchors read the same surface facts as the shares.""" + + counts = _counts() + frame, surface = _frame(_values()), _surface(tmp_path / "a", counts) + receipt = _run( + TRANSPORT_QUANTILE_MAP_V2, + _node(QUANTILE_MAP_V2_REF, _pareto()), + frame, + surface, + ).receipt["tail_estimator"] + assert [(row["reference"], row["count"]) for row in receipt["anchors"]] == [ + (name, counts[name]) for name in ANCHORS + ] + expected = grouped_pareto_mle( + [(name, lower, upper, counts[name]) for name, lower, upper in BAND_ROWS[1:]], + bracket=BRACKET, + iterations=ITERATIONS, + ) + assert receipt["alpha"]["numerator"] == expected.numerator + # Changing one band row's count changes alpha. + changed = _run( + TRANSPORT_QUANTILE_MAP_V2, + _node(QUANTILE_MAP_V2_REF, _pareto()), + frame, + _surface(tmp_path / "b", _counts(toy_wage_2=counts["toy_wage_2"] * 2)), + ).receipt["tail_estimator"] + assert changed["alpha"]["numerator"] != receipt["alpha"]["numerator"] + # A count below the anchors changes the shares but not alpha. + body = _run( + TRANSPORT_QUANTILE_MAP_V2, + _node(QUANTILE_MAP_V2_REF, _pareto()), + frame, + _surface(tmp_path / "c", _counts(toy_wage_0=counts["toy_wage_0"] * 2)), + ).receipt["tail_estimator"] + assert body["alpha"] == receipt["alpha"] + + +@pytest.mark.parametrize( + ("interpolation", "bands", "message"), + [ + ({"method": "pareto", "alpha": 2.0}, None, "literal alpha belongs"), + ({"method": "uniform"}, None, "literal alpha belongs"), + ( + _pareto(_estimator(anchor_references=["toy_wage_1", "toy_wage_2"])), + None, + "top contiguous bands", + ), + ( + _pareto(_estimator(anchor_references=["toy_wage_top", "toy_wage_2"])), + None, + "top contiguous bands", + ), + ( + _pareto(), + {"positive": _bands()["bands"], "negative": _bands()["bands"]}, + "bare band list", + ), + (_pareto(_estimator(bracket=["4", "64"])), None, "bracket_without_sign_change"), + ], +) +def test_v2_refusals(tmp_path, interpolation, bands, message) -> None: + with pytest.raises(ValueError, match=message): + _run( + TRANSPORT_QUANTILE_MAP_V2, + _node(QUANTILE_MAP_V2_REF, interpolation, bands), + _frame(_values()), + _surface(tmp_path, _counts()), + ) + + +def test_v2_refuses_a_count_that_is_not_on_the_ranked_unit(tmp_path) -> None: + node = _node(QUANTILE_MAP_V2_REF, _pareto()) + node = Node( + node.id, + node.kernel, + inputs=node.inputs, + outputs=node.outputs, + params={**node.params, "aggregate_entity": "household"}, + artifact_inputs=node.artifact_inputs, + ) + with pytest.raises(ValueError, match="unfiltered household count facts"): + _run( + TRANSPORT_QUANTILE_MAP_V2, + node, + _frame(_values()), + _surface(tmp_path, _counts()), + ) + + +# --------------------------------------------------------------------------- +# Registration and support +# --------------------------------------------------------------------------- + + +def test_v2_is_registered_beside_v1() -> None: + registry = register_transport_population_kernels(KernelRegistry()) + assert registry.get(QUANTILE_MAP_V2_REF) is TRANSPORT_QUANTILE_MAP_V2 + assert registry.get("transport.quantile_map@1") is TRANSPORT_QUANTILE_MAP + assert {"transport.quantile_map@1", QUANTILE_MAP_V2_REF} <= compose.SKELETON_KERNELS + assert compose._SUPPORT_PRODUCERS[QUANTILE_MAP_V2_REF] == "quantile_mapped" + assert QUANTILE_MAP_V2_REF in compose._MONEY_PRODUCERS + assert TRANSPORT_QUANTILE_MAP_V2.implementation_hash() != ( + TRANSPORT_QUANTILE_MAP.implementation_hash() + ) + + +def _spec_with_map(kernel): + spec = copy.deepcopy(compose.load_transport_spec("nz")) + spec["resources"]["transport_graph"]["nodes"].append( + { + "id": "toy.wage.map", + "kernel": kernel, + "population": "nz.central", + "inputs": [{"entity": "person", "columns": ["employment_income"]}], + "outputs": [ + { + "entity": "person", + "column": "employment_income", + "dtype": "float64", + "rewrite": True, + } + ], + } + ) + return spec + + +def test_composition_accepts_v2_as_the_quantile_mapped_money_producer() -> None: + """A map in ``nz.central`` satisfies the central encode's wage read.""" + + def wage_gaps(spec): + return [ + gap + for gap in compose._money_producer_gaps(spec) + if "person.employment_income" in gap + ] + + def central_readers(gaps): + return [gap for gap in gaps if "'nz.encode.person'" in gap] + + assert central_readers(wage_gaps(compose.load_transport_spec("nz"))) + for kernel in (QUANTILE_MAP_V2_REF, "transport.quantile_map@1"): + gaps = wage_gaps(_spec_with_map(kernel)) + assert central_readers(gaps) == [] + assert not any("produced by" in gap for gap in gaps) + (gap,) = central_readers(wage_gaps(_spec_with_map("transport.currency@1"))) + assert "transport.quantile_map@2 may produce it" in gap + + +def _support_params(): + bands = _bands() + return {"bands": bands, "interpolation": _pareto()} + + +def test_the_w4_support_producer_reads_v2_params_with_an_open_top() -> None: + params = _support_params() + support = quantile_support(params["bands"], params["interpolation"]) + assert support.positive == ( + (1.0, 1000.0), + (1000.0, 2000.0), + (2000.0, 3000.0), + (3000.0, None), + ) + # A bare band list feeds both signs, as quantile_map applies it. + assert support.negative == support.positive + low, high = quantile_bounds(support.positive, support.negative) + assert high is None # expected_max: null + + +def _support_row(params, values): + """Run w4's donor-ranges and support-bounds producers on one roster row.""" + roster = { + "columns": [ + { + "entity": "person", + "column": "employment_income", + "class": "quantile_mapped", + "owner": "toy.wage.map", + "producer_params": params, + } + ] + } + frame = _frame(values) + text = canonical_json(roster).decode() + + def node(ref, alias, artifact_type, inputs=(), artifact_inputs=()): + return Node( + f"toy.evidence.{alias}", + ref, + inputs=inputs, + params={"support_roster": text}, + artifact_inputs=artifact_inputs, + artifact_outputs=(ArtifactOutput(alias, artifact_type),), + ) + + def context(item, artifacts): + return _project_context( + item, + Population.from_frame(frame, "toy.terminal"), + key="0" * 64, + sources={}, + tolerances={}, + numerics={}, + artifacts=artifacts, + ) + + ranges_node = node(DONOR_RANGES_REF, DONOR_RANGES, DONOR_RANGES_TYPE) + ranges = TRANSPORT_DONOR_RANGES.run(context(ranges_node, {})).artifacts[ + DONOR_RANGES + ] + bounds_node = node( + SUPPORT_BOUNDS_REF, + DONOR_SUPPORT_BOUNDS, + SUPPORT_BOUNDS_TYPE, + inputs=(Slice("person", ("employment_income",)),), + artifact_inputs=( + ArtifactInput( + DONOR_RANGES, ranges_node.id, DONOR_RANGES, DONOR_RANGES_TYPE + ), + ), + ) + value = ArtifactValue(ranges, DONOR_RANGES_TYPE, "c" * 64, "d" * 64, NumericScope()) + bounds = TRANSPORT_SUPPORT_BOUNDS.run( + context(bounds_node, {DONOR_RANGES: value}) + ).artifacts[DONOR_SUPPORT_BOUNDS] + (row,) = decode_support_bounds(bounds).document["columns"] + return row + + +def test_the_support_bounds_kernel_reports_a_null_expected_max_for_v2() -> None: + row = _support_row(_support_params(), [0.0, 5.0, 4e6]) + assert row["class"] == "quantile_mapped" + assert row["expected_max"] is None + assert row["records_outside"] == 0 + + +def test_v2_support_bounds_both_signs_for_a_bare_band_list() -> None: + row = _support_row(_support_params(), [0.0, 5.0, 4e6]) + assert row["expected_min"] is None + assert row["expected_max"] is None