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..96e8ac6d6d 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", @@ -189,6 +190,7 @@ "simulate.solve_zero@1", # [BE] Non-structural pre-calibration solves. "money.minor_units@1", "money.annual_transfers@1", + "measure.combine@1", "takeup.assign@1", "targets.compile@1", "calibration.prepare@1", @@ -215,7 +217,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 +845,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 +944,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 +1017,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 +1634,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/measure_kernels.py b/packages/microcosm-build/src/microcosm/build/transport/measure_kernels.py new file mode 100644 index 0000000000..f93aac85ee --- /dev/null +++ b/packages/microcosm-build/src/microcosm/build/transport/measure_kernels.py @@ -0,0 +1,393 @@ +"""``measure.combine@1``: declared signed integer-cent components summed into a measure. + +A declaration document, never this module, names the entity, every component +column and its sign, the add-backs, the floor, the scale and the three output +columns. For each record the kernel computes, in Python integers, + +``x = sum(sign_k * c_k) + sum(max(0, -(sign_j * c_j)) for each add-back j)`` + +so an add-back component counts only its gains (its losses are not deducted). +It writes ``x`` in cents, a unit column equal to ``max(floor, x) / scale`` +(correctly rounded binary64), and a universe status: ``included`` where +``x > 0``, ``excluded`` where ``x <= 0`` and ``unresolved`` where any +component is null. An unresolved record's measure is null and its unit is +NaN; it is never read as zero. A measure outside int64 is refused by name. + +The node may declare ``measurement_id``; the kernel carries it into its +receipt for the composition proof that reads it, and computes nothing from it. + +Declaration (``measure`` parameter, canonical JSON, with ``measure_sha256``):: + + {"entity": "person", + "components": [{"column": "...", "sign": 1}, {"column": "...", "sign": -1}], + "add_backs": [""], + "floor": , "scale": , + "outputs": {"measure": "...", "unit": "...", "universe_status": "..."}} +""" + +from __future__ import annotations + +from collections.abc import Mapping +from math import fsum + +import numpy as np +import pandas as pd + +import microcosm.build.reference_intervals as reference_intervals_module +from microcosm.build.reference_intervals import UNIVERSE_STATES +from microcosm.graph import ( + ROWS_ALL, + Capabilities, + Determinism, + KernelBase, + KernelContext, + KernelRegistry, + KernelResult, + Numeric, + source_hash, +) +from microcosm.graph.population import dtype_for_token + +from . import graph_inputs +from .graph_inputs import ( + canonical_document_param, + require_params, + sha256_param, + sha256_text, + string_param, +) + +__all__ = [ + "MEASURE_COMBINE", + "MEASURE_COMBINE_REF", + "MeasureCombineKernel", + "combine_measure", + "register_measure_kernels", + "validate_measure_declaration", +] + +MEASURE_COMBINE_REF = "measure.combine@1" +INCLUDED, EXCLUDED, UNRESOLVED = UNIVERSE_STATES +_INT64_MIN = int(np.iinfo(np.int64).min) +_INT64_MAX = int(np.iinfo(np.int64).max) +_OUTPUT_DTYPES = { + "measure": frozenset({"Int64", "int64"}), + "unit": frozenset({"float64"}), + "universe_status": frozenset({"string"}), +} + + +def _integer(value: object) -> bool: + return isinstance(value, int) and not isinstance(value, bool) + + +def _name(value: object) -> bool: + return isinstance(value, str) and bool(value) + + +def validate_measure_declaration(document: Mapping[str, object]) -> dict: + """Check a declaration's shape; every value comes from the document.""" + + ref = MEASURE_COMBINE_REF + if not isinstance(document, Mapping) or set(document) != { + "entity", + "components", + "add_backs", + "floor", + "scale", + "outputs", + }: + raise ValueError( + f"{ref} declares exactly entity, components, add_backs, floor, scale " + "and outputs." + ) + if not _name(document["entity"]): + raise ValueError(f"{ref} entity must be a nonempty name.") + components = document["components"] + if ( + not isinstance(components, list) + or not components + or any( + not isinstance(row, Mapping) + or set(row) != {"column", "sign"} + or not _name(row["column"]) + or not _integer(row["sign"]) + or row["sign"] not in {-1, 1} + for row in components + ) + ): + raise ValueError( + f"{ref} components are a nonempty list of {{column, sign}} with sign " + "1 or -1." + ) + columns = [row["column"] for row in components] + if len(set(columns)) != len(columns): + raise ValueError(f"{ref} component columns must be distinct.") + add_backs = document["add_backs"] + if ( + not isinstance(add_backs, list) + or any(not _name(name) for name in add_backs) + or len(set(add_backs)) != len(add_backs) + or not set(add_backs) <= set(columns) + ): + raise ValueError( + f"{ref} add_backs must name distinct declared component columns." + ) + if not _integer(document["floor"]): + raise ValueError(f"{ref} floor must be an integer number of minor units.") + if not _integer(document["scale"]) or document["scale"] <= 0: + raise ValueError(f"{ref} scale must be a positive integer.") + if not _INT64_MIN <= document["floor"] <= _INT64_MAX: + raise ValueError(f"{ref} floor must fit int64.") + outputs = document["outputs"] + if ( + not isinstance(outputs, Mapping) + or set(outputs) != set(_OUTPUT_DTYPES) + or any(not _name(name) for name in outputs.values()) + or len(set(outputs.values())) != len(outputs) + or set(outputs.values()) & set(columns) + ): + raise ValueError( + f"{ref} outputs name distinct measure, unit and universe_status " + "columns, none of them a component." + ) + return { + "entity": document["entity"], + "components": [dict(row) for row in components], + "add_backs": list(add_backs), + "floor": document["floor"], + "scale": document["scale"], + "outputs": dict(outputs), + } + + +def _component_values(series: pd.Series, column: str) -> list[int | None]: + """Exact Python integers, ``None`` where the component is null.""" + dtype = series.dtype + if isinstance(dtype, pd.Int64Dtype) or ( + isinstance(dtype, np.dtype) and np.issubdtype(dtype, np.integer) + ): + return [None if pd.isna(value) else int(value) for value in series.tolist()] + raise ValueError( + f"{MEASURE_COMBINE_REF} component {column!r} must be an integer minor-unit " + f"column, not {dtype}." + ) + + +def combine_measure( + declaration: Mapping[str, object], + columns: Mapping[str, pd.Series], + weights: np.ndarray, +) -> tuple[list[int | None], list[str], dict]: + """Combine one entity's components; returns measures, statuses and a receipt. + + ``columns`` holds each declared component; ``weights`` are the entity's + weights in the same row order. All arithmetic is in Python integers. + """ + + components = declaration["components"] + add_backs = set(declaration["add_backs"]) + values = { + row["column"]: _component_values(columns[row["column"]], row["column"]) + for row in components + } + size = len(weights) + if any(len(column) != size for column in values.values()): + raise ValueError( + f"{MEASURE_COMBINE_REF} components and weights differ in length." + ) + measures: list[int | None] = [] + statuses: list[str] = [] + back_totals = {name: 0 for name in declaration["add_backs"]} + for index in range(size): + signed = [ + (row["column"], row["sign"] * values[row["column"]][index]) + if values[row["column"]][index] is not None + else (row["column"], None) + for row in components + ] + if any(value is None for _, value in signed): + measures.append(None) + statuses.append(UNRESOLVED) + continue + backs = { + column: max(0, -value) for column, value in signed if column in add_backs + } + total = sum(value for _, value in signed) + sum(backs.values()) + if not _INT64_MIN <= total <= _INT64_MAX: + raise ValueError( + f"{MEASURE_COMBINE_REF} refuses int64_overflow: record {index} sums " + f"to {total}, outside int64." + ) + for column, value in backs.items(): + back_totals[column] += value + measures.append(total) + statuses.append(INCLUDED if total > 0 else EXCLUDED) + + resolved = [index for index, status in enumerate(statuses) if status != UNRESOLVED] + weights = np.asarray(weights, dtype=np.float64) + + def weighted(indices) -> float: + return fsum(float(weights[index]) for index in indices) + + receipt_components = [] + for row in components: + column = row["column"] + signed = [row["sign"] * values[column][index] for index in resolved] + receipt_components.append( + { + "column": column, + "sign": row["sign"], + "add_back": column in add_backs, + "n_null": sum(value is None for value in values[column]), + "unweighted_signed_total": sum(signed), + "weighted_signed_total": fsum( + float(weights[index]) * value + for index, value in zip(resolved, signed, strict=True) + ), + } + ) + negative = [index for index in resolved if measures[index] < 0] + zero = [index for index in resolved if measures[index] == 0] + receipt = { + "entity": declaration["entity"], + "n": size, + "components": receipt_components, + "add_backs": [ + {"column": name, "total": back_totals[name]} + for name in declaration["add_backs"] + ], + "sum_measure": sum(measures[index] for index in resolved), + "universe_statuses": { + state: statuses.count(state) for state in UNIVERSE_STATES + }, + "below_zero": {"n": len(negative), "weighted": weighted(negative)}, + "at_zero": {"n": len(zero), "weighted": weighted(zero)}, + "floor": declaration["floor"], + "scale": declaration["scale"], + } + # The receipt's identities, checked exactly before anything is written. + if receipt["sum_measure"] != sum( + row["unweighted_signed_total"] for row in receipt_components + ) + sum(back_totals.values()): + raise AssertionError("components and add-backs must sum to the measure.") + return measures, statuses, receipt + + +class MeasureCombineKernel(KernelBase): + """``measure.combine@1``: one measure from declared signed components.""" + + ref = MEASURE_COMBINE_REF + capabilities = Capabilities( + Determinism.DETERMINISTIC, + numeric=Numeric.BITWISE, + dependencies=("numpy", "pandas"), + ) + + def implementation_hash(self) -> str: + return source_hash( + type(self), + graph_inputs, + reference_intervals_module, + dependencies=self.capabilities.dependencies, + ) + + def run(self, context: KernelContext) -> KernelResult: + ref = self.ref + require_params( + context, + ref, + required=frozenset({"measure", "measure_sha256"}), + optional=frozenset({"measurement_id"}), + ) + document = canonical_document_param(context, ref, "measure") + if sha256_param(context, ref, "measure_sha256") != sha256_text( + context.params["measure"].encode("utf-8") + ): + raise ValueError(f"{ref} measure_sha256 differs from its declaration.") + declaration = validate_measure_declaration(document) + measurement_id = ( + string_param(context, ref, "measurement_id") + if "measurement_id" in context.params + else None + ) + entity = declaration["entity"] + if any(item.rows != ROWS_ALL for item in context.node.inputs): + raise ValueError(f"{ref} requires whole-row slices.") + declared_inputs = { + (item.entity, column) + for item in context.node.inputs + for column in item.columns + } + expected_inputs = {(entity, row["column"]) for row in declaration["components"]} + if declared_inputs != expected_inputs: + raise ValueError( + f"{ref} inputs must slice exactly the declared components on {entity!r}." + ) + outputs = declaration["outputs"] + declared_outputs = { + (item.entity, item.column): item for item in context.node.outputs + } + if set(declared_outputs) != {(entity, name) for name in outputs.values()}: + raise ValueError( + f"{ref} outputs must equal the declaration's three columns on " + f"{entity!r}." + ) + dtypes = {} + for role, name in outputs.items(): + item = declared_outputs[entity, name] + if item.rewrite or item.dtype not in _OUTPUT_DTYPES[role]: + raise ValueError( + f"{ref} output {name!r} ({role}) is a new column of dtype " + f"{sorted(_OUTPUT_DTYPES[role])}." + ) + dtypes[role] = item.dtype + table = context.tables[entity] + weights = context.weights[entity].values + measures, statuses, receipt = combine_measure( + declaration, + {row["column"]: table[row["column"]] for row in declaration["components"]}, + weights, + ) + if dtypes["measure"] == "int64" and any(value is None for value in measures): + raise ValueError( + f"{ref} has unresolved records; declare the measure Int64 so they " + "stay null rather than zero." + ) + floor, scale = declaration["floor"], declaration["scale"] + # Python's int / int is correctly rounded for any magnitude. + units = [ + float("nan") if value is None else max(floor, value) / scale + for value in measures + ] + index = pd.Index(table[f"{entity}_id"].to_numpy(copy=True), name=f"{entity}_id") + columns = { + (entity, outputs["measure"]): pd.Series( + pd.array(measures, dtype=dtype_for_token(dtypes["measure"])), + index=index, + name=outputs["measure"], + ), + (entity, outputs["unit"]): pd.Series( + np.asarray(units, dtype=np.float64), index=index, name=outputs["unit"] + ), + (entity, outputs["universe_status"]): pd.Series( + pd.array(statuses, dtype=dtype_for_token("string")), + index=index, + name=outputs["universe_status"], + ), + } + receipt = { + **receipt, + "outputs": dict(outputs), + "measurement_id": measurement_id, + } + return KernelResult(columns=columns, receipt=receipt) + + +MEASURE_COMBINE = MeasureCombineKernel() + + +def register_measure_kernels(registry: KernelRegistry) -> KernelRegistry: + """Register ``measure.combine@1`` explicitly and idempotently.""" + registry.register(MEASURE_COMBINE) + return registry 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..77c1cf3fc6 100644 --- a/packages/microcosm-build/src/microcosm/build/transport/registry.py +++ b/packages/microcosm-build/src/microcosm/build/transport/registry.py @@ -50,9 +50,11 @@ from .gate_bindings import TRANSPORT_GATE_REGISTRY from .gate_kernels import register_gate_kernel from .geography_kernels import register_geography_kernels +from .measure_kernels import register_measure_kernels 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,14 +533,16 @@ def build_transport_registry( def register_transport_population_kernels(registry: KernelRegistry) -> KernelRegistry: - """Install the thirteen population kernels without import-time side effects. + """Install the sixteen 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) + register_measure_kernels(registry) return registry diff --git a/packages/microcosm-build/tests/engine_free/shared/test_transport_measure_kernels.py b/packages/microcosm-build/tests/engine_free/shared/test_transport_measure_kernels.py new file mode 100644 index 0000000000..7434324543 --- /dev/null +++ b/packages/microcosm-build/tests/engine_free/shared/test_transport_measure_kernels.py @@ -0,0 +1,575 @@ +"""``measure.combine@1`` (G8-r2c hunk C). + +The kernel sums declared signed integer-cent components on one entity. Every +column name, the floor and the scale here are synthetic declarations; the +kernel reads them from its declaration and holds none of them. +""" + +from __future__ import annotations + +import random +from fractions import Fraction +from math import fsum + +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.reference_intervals import UNIVERSE_STATES +from microcosm.build.transport import compose +from microcosm.build.transport.measure_kernels import ( + MEASURE_COMBINE, + MEASURE_COMBINE_REF, + validate_measure_declaration, +) +from microcosm.build.transport.registry import register_transport_population_kernels +from microcosm.frame import EntitySchema, Frame, WeightKind, Weights +from microcosm.graph import KernelRegistry, Node, Owned, Slice +from microcosm.graph.canonical import canonical_json +from microcosm.graph.executor import Population, _project_context +from test_support.microcosm_build.transport_population import digest + +SCHEMA = EntitySchema(group_entities=("household",)) +INT64_MAX = int(np.iinfo(np.int64).max) +INT64_MIN = int(np.iinfo(np.int64).min) +OUTPUTS = { + "measure": "toy_measure_cents", + "unit": "toy_measure_units", + "universe_status": "toy_measure_universe_status", +} +PROPERTY = settings( + max_examples=80, + deadline=None, + suppress_health_check=(HealthCheck.function_scoped_fixture,), +) + + +def _declaration(components, add_backs=(), *, floor=0, scale=100, outputs=None): + return { + "entity": "person", + "components": [{"column": c, "sign": s} for c, s in components], + "add_backs": list(add_backs), + "floor": floor, + "scale": scale, + "outputs": dict(outputs or OUTPUTS), + } + + +def _frame(columns: dict[str, list], dtype: str, extra=None) -> Frame: + size = len(next(iter(columns.values()))) + person = pd.DataFrame( + { + "person_id": np.arange(1, size + 1, dtype=np.int64), + "person_household_id": np.arange(1, size + 1, dtype=np.int64), + } + ) + for name, values in columns.items(): + person[name] = pd.array(values, dtype=dtype) + for name, values in (extra or {}).items(): + person[name] = values + household = pd.DataFrame({"household_id": np.arange(1, size + 1, dtype=np.int64)}) + return Frame( + {"person": person, "household": household}, + SCHEMA, + {"household": Weights(np.linspace(1, 2, size), WeightKind.DESIGN)}, + ) + + +def _node(declaration, *, measure_dtype="Int64", measurement_id=None, params=None): + text = canonical_json(declaration).decode() + outputs = declaration["outputs"] + node_params = {"measure": text, "measure_sha256": digest(declaration)} + if measurement_id is not None: + node_params["measurement_id"] = measurement_id + node_params.update(params or {}) + return Node( + "toy.measure", + MEASURE_COMBINE_REF, + inputs=( + Slice( + declaration["entity"], + tuple(row["column"] for row in declaration["components"]), + ), + ), + outputs=( + Owned("person", outputs["measure"], measure_dtype), + Owned("person", outputs["unit"], "float64"), + Owned("person", outputs["universe_status"], "string"), + ), + params=node_params, + ) + + +def _run(columns, declaration, *, dtype="Int64", **node): + frame = _frame(columns, dtype) + context = _project_context( + _node(declaration, **node), + Population.from_frame(frame, "toy.population"), + key="0" * 64, + sources={}, + tolerances={}, + numerics={}, + ) + result = MEASURE_COMBINE.run(context) + out = {role: result.columns["person", name] for role, name in OUTPUTS.items()} + return out, result.receipt + + +def _measure(out): + return [None if pd.isna(v) else int(v) for v in out["measure"].tolist()] + + +# --------------------------------------------------------------------------- +# Examples +# --------------------------------------------------------------------------- + +WAGES = "toy_wages_cents" +BUSINESS = "toy_business_cents" +RENT = "toy_rent_cents" +DEDUCTION = "toy_deduction_cents" + + +def test_losses_net_against_gains() -> None: + out, receipt = _run( + {WAGES: [10_000, 0, 500], BUSINESS: [-5_000, -200, 0]}, + _declaration([(WAGES, 1), (BUSINESS, 1)]), + ) + assert _measure(out) == [5_000, -200, 500] + assert out["universe_status"].tolist() == ["included", "excluded", "included"] + assert receipt["components"][1]["unweighted_signed_total"] == -5_200 + + +def test_a_negative_sign_subtracts() -> None: + out, _ = _run( + {WAGES: [10_000, 100], DEDUCTION: [2_500, 300]}, + _declaration([(WAGES, 1), (DEDUCTION, -1)]), + ) + assert _measure(out) == [7_500, -200] + + +def test_the_add_back_ignores_an_add_back_components_losses() -> None: + out, receipt = _run( + {WAGES: [10_000, 10_000, 0], RENT: [-3_000, 3_000, -1]}, + _declaration([(WAGES, 1), (RENT, 1)], [RENT]), + ) + assert _measure(out) == [10_000, 13_000, 0] + assert receipt["add_backs"] == [{"column": RENT, "total": 3_001}] + assert out["universe_status"].tolist() == ["included", "included", "excluded"] + + +def test_records_at_or_below_zero_sit_outside_the_positive_universe() -> None: + out, receipt = _run( + {WAGES: [0, -1, 1, -250]}, + _declaration([(WAGES, 1)]), + ) + assert out["universe_status"].tolist() == [ + "excluded", + "excluded", + "included", + "excluded", + ] + # The unit column is max(floor, x) / scale: excluded records sit at the floor. + assert out["unit"].tolist() == [0.0, 0.0, 0.01, 0.0] + assert receipt["below_zero"]["n"] == 2 + assert receipt["at_zero"]["n"] == 1 + + +def test_the_receipt_weights_each_record_by_its_frame_weight() -> None: + """The weighted totals and masses use each record's own weight.""" + + out, receipt = _run( + { + WAGES: [10_000, 0, 500, 300, 0], + BUSINESS: [-5_000, -200, 0, -300, 0], + }, + _declaration([(WAGES, 1), (BUSINESS, 1)]), + ) + assert _measure(out) == [5_000, -200, 500, 0, 0] + # _frame weights its five households (one person each) 1, 1.25, 1.5, 1.75 + # and 2; every record here is resolved. + w = [1.0, 1.25, 1.5, 1.75, 2.0] + assert np.linspace(1, 2, 5).tolist() == w + wages, business = receipt["components"] + assert wages["weighted_signed_total"] == fsum( + [w[0] * 10_000, w[1] * 0, w[2] * 500, w[3] * 300, w[4] * 0] + ) + assert wages["weighted_signed_total"] == 11_275.0 + assert business["weighted_signed_total"] == fsum( + [w[0] * -5_000, w[1] * -200, w[2] * 0, w[3] * -300, w[4] * 0] + ) + assert business["weighted_signed_total"] == -5_775.0 + # x < 0 is record 1 alone; x = 0 is records 3 and 4. + assert receipt["below_zero"] == {"n": 1, "weighted": w[1]} + assert receipt["at_zero"] == {"n": 2, "weighted": fsum([w[3], w[4]])} + assert receipt["at_zero"]["weighted"] == 3.75 + + +def test_the_receipt_masses_count_resolved_records_only() -> None: + """An unresolved record joins neither ``below_zero`` nor ``at_zero``.""" + + out, receipt = _run({WAGES: [0, None, -5, None, 7]}, _declaration([(WAGES, 1)])) + assert _measure(out) == [0, None, -5, None, 7] + # _frame weights the five records 1, 1.25, 1.5, 1.75 and 2. Records 1 and 3 + # are unresolved, so x < 0 is record 2 alone and x = 0 is record 0 alone. + assert receipt["below_zero"] == {"n": 1, "weighted": 1.5} + assert receipt["at_zero"] == {"n": 1, "weighted": 1.0} + + +def test_an_int64_overflow_is_refused_by_name() -> None: + with pytest.raises(ValueError, match="int64_overflow"): + _run( + {WAGES: [INT64_MAX, 1], BUSINESS: [1, 1]}, + _declaration([(WAGES, 1), (BUSINESS, 1)]), + ) + with pytest.raises(ValueError, match="int64_overflow"): + _run( + {WAGES: [INT64_MIN, 1], BUSINESS: [1, 1]}, + _declaration([(WAGES, 1), (BUSINESS, -1)]), + ) + + +def test_a_null_component_leaves_the_record_unresolved_not_zero() -> None: + out, receipt = _run( + {WAGES: [100, None, 0], BUSINESS: [0, 50, None]}, + _declaration([(WAGES, 1), (BUSINESS, 1)]), + ) + assert _measure(out) == [100, None, None] + assert out["universe_status"].tolist() == ["included", "unresolved", "unresolved"] + assert np.isnan(out["unit"].to_numpy()[1:]).all() + assert receipt["universe_statuses"] == { + "included": 1, + "excluded": 0, + "unresolved": 2, + } + assert [row["n_null"] for row in receipt["components"]] == [1, 1] + with pytest.raises(ValueError, match="declare the measure Int64"): + _run( + {WAGES: [100, None]}, + _declaration([(WAGES, 1)]), + measure_dtype="int64", + ) + + +def test_a_non_null_int64_measure_is_written_as_int64() -> None: + out, _ = _run( + {WAGES: [3, -4]}, + _declaration([(WAGES, 1)]), + dtype="int64", + measure_dtype="int64", + ) + assert out["measure"].dtype == np.dtype("int64") + assert out["measure"].tolist() == [3, -4] + + +def test_the_floor_and_scale_come_from_the_declaration() -> None: + out, receipt = _run( + {WAGES: [0, 1_000, 5_000, -7]}, + _declaration([(WAGES, 1)], floor=2_000, scale=1_000), + ) + assert out["unit"].tolist() == [2.0, 2.0, 5.0, 2.0] + assert (receipt["floor"], receipt["scale"]) == (2_000, 1_000) + + +def test_the_measurement_id_is_carried_not_used() -> None: + columns, declaration = {WAGES: [5, -5]}, _declaration([(WAGES, 1)]) + out, receipt = _run(columns, declaration, measurement_id="toy_measurement_v1") + bare, bare_receipt = _run(columns, declaration) + assert receipt["measurement_id"] == "toy_measurement_v1" + assert bare_receipt["measurement_id"] is None + for role in OUTPUTS: + pd.testing.assert_series_equal(out[role], bare[role]) + + +@pytest.mark.parametrize( + "node", + [ + {"measurement_id": 123}, + {"measurement_id": ""}, + {"measurement_id": True}, + {"params": {"measurement_id": None}}, + ], + ids=["integer", "empty", "boolean", "explicit_null"], +) +def test_a_measurement_id_that_is_not_a_non_empty_string_is_refused(node) -> None: + with pytest.raises(TypeError) as refusal: + _run({WAGES: [5, -5]}, _declaration([(WAGES, 1)]), **node) + assert str(refusal.value) == ( + f"{MEASURE_COMBINE_REF} parameter 'measurement_id' must be a non-empty string." + ) + + +def test_the_receipt_names_the_declared_outputs() -> None: + _, receipt = _run({WAGES: [1]}, _declaration([(WAGES, 1)])) + assert receipt["outputs"] == OUTPUTS + assert receipt["entity"] == "person" + + +@pytest.mark.parametrize( + ("change", "message"), + [ + ({"floor": 0.5}, "floor must be an integer"), + ({"floor": True}, "floor must be an integer"), + ({"floor": INT64_MAX + 1}, "floor must fit int64"), + ({"floor": INT64_MIN - 1}, "floor must fit int64"), + ({"scale": 0}, "scale must be a positive integer"), + ({"scale": 1.0}, "scale must be a positive integer"), + ({"add_backs": ["toy_other_cents"]}, "add_backs must name"), + ({"components": []}, "components are a nonempty list"), + ({"components": [{"column": WAGES, "sign": 2}]}, "sign"), + ( + {"components": [{"column": WAGES, "sign": 1}] * 2}, + "must be distinct", + ), + ({"outputs": {**OUTPUTS, "unit": OUTPUTS["measure"]}}, "distinct measure"), + ({"outputs": {**OUTPUTS, "unit": WAGES}}, "none of them a component"), + ({"outputs": {"measure": "a", "unit": "b"}}, "distinct measure"), + ], +) +def test_malformed_declarations_are_refused(change, message) -> None: + declaration = {**_declaration([(WAGES, 1)]), **change} + with pytest.raises(ValueError, match=message): + validate_measure_declaration(declaration) + missing = _declaration([(WAGES, 1)]) + del missing["floor"] + with pytest.raises(ValueError, match="declares exactly"): + validate_measure_declaration(missing) + + +def _run_node(node, frame) -> None: + context = _project_context( + node, + Population.from_frame(frame, "toy.population"), + key="0" * 64, + sources={}, + tolerances={}, + numerics={}, + ) + MEASURE_COMBINE.run(context) + + +def _reshaped(node, **change) -> Node: + fields = { + "inputs": node.inputs, + "outputs": node.outputs, + "params": node.params, + **change, + } + return Node(node.id, node.kernel, **fields) + + +def test_node_shape_refusals() -> None: + declaration = _declaration([(WAGES, 1)]) + with pytest.raises(ValueError, match="measure_sha256 differs"): + _run({WAGES: [1]}, declaration, params={"measure_sha256": "0" * 64}) + with pytest.raises(ValueError, match="unknown parameter"): + _run({WAGES: [1]}, declaration, params={"floor": 0}) + with pytest.raises(ValueError, match="is a new column of dtype"): + _run({WAGES: [1]}, declaration, measure_dtype="float64") + frame = _frame({WAGES: [1], BUSINESS: [2]}, "Int64") + node = _node(declaration) + wider = Node( + node.id, + node.kernel, + inputs=(Slice("person", (WAGES, BUSINESS)),), + outputs=node.outputs, + params=node.params, + ) + context = _project_context( + wider, + Population.from_frame(frame, "toy.population"), + key="0" * 64, + sources={}, + tolerances={}, + numerics={}, + ) + with pytest.raises(ValueError, match="slice exactly the declared components"): + MEASURE_COMBINE.run(context) + + # A row mask would combine a subset of the entity: whole rows only. + frame = _frame({WAGES: [1, 2]}, "Int64", {"toy_mask": np.array([True, False])}) + masked = _reshaped( + node, inputs=(Slice("person", (WAGES, "toy_mask"), rows="toy_mask"),) + ) + with pytest.raises(ValueError, match="requires whole-row slices"): + _run_node(masked, frame) + + # The outputs are exactly the declaration's three columns. + frame = _frame({WAGES: [1, 2]}, "Int64") + extra = _reshaped( + node, outputs=(*node.outputs, Owned("person", "toy_extra_cents", "Int64")) + ) + renamed = _reshaped( + node, + outputs=( + node.outputs[0], + Owned("person", "toy_other_units", "float64"), + node.outputs[2], + ), + ) + for item in (extra, renamed): + with pytest.raises(ValueError, match="outputs must equal the declaration"): + _run_node(item, frame) + + # Each output is a new column; the stand-in never rewrites one. + frame = _frame({WAGES: [1, 2]}, "Int64", {OUTPUTS["unit"]: np.zeros(2)}) + rewrite = _reshaped( + node, + outputs=( + node.outputs[0], + Owned("person", OUTPUTS["unit"], "float64", rewrite=True), + node.outputs[2], + ), + ) + with pytest.raises(ValueError, match="is a new column of dtype"): + _run_node(rewrite, frame) + + +def test_a_float_component_is_refused() -> None: + with pytest.raises(ValueError, match="integer minor-unit column"): + _run({WAGES: [1.0, 2.0]}, _declaration([(WAGES, 1)]), dtype="float64") + + +def test_measure_combine_is_registered_in_the_skeleton() -> None: + registry = register_transport_population_kernels(KernelRegistry()) + assert registry.get(MEASURE_COMBINE_REF) is MEASURE_COMBINE + assert MEASURE_COMBINE_REF in compose.SKELETON_KERNELS + + +# --------------------------------------------------------------------------- +# Properties and the differential +# --------------------------------------------------------------------------- + +_CENTS = st.one_of( + st.integers(min_value=-(10**12), max_value=10**12), + st.integers(min_value=-(2**60), max_value=2**60), + st.none(), +) + + +@st.composite +def measures(draw): + width = draw(st.integers(min_value=1, max_value=6)) + size = draw(st.integers(min_value=1, max_value=12)) + names = [f"toy_component_{index}_cents" for index in range(width)] + signs = [draw(st.sampled_from((1, -1))) for _ in names] + backs = [name for name in names if draw(st.booleans())] + columns = {name: [draw(_CENTS) for _ in range(size)] for name in names} + floor = draw(st.integers(min_value=-(10**6), max_value=10**6)) + scale = draw(st.integers(min_value=1, max_value=10**4)) + return columns, list(zip(names, signs, strict=True)), backs, floor, scale + + +def _independent(columns, components, backs): + """The measure by plain Python integers, independently of the kernel.""" + size = len(next(iter(columns.values()))) + result = [] + for index in range(size): + cells = {name: columns[name][index] for name, _ in components} + if any(value is None for value in cells.values()): + result.append(None) + continue + total = 0 + for name, sign in components: + total += sign * cells[name] + for name, sign in components: + if name in backs and sign * cells[name] < 0: + total += -(sign * cells[name]) + result.append(total) + return result + + +def _run_or_overflow(columns, components, backs, floor=0, scale=100): + try: + return _run(columns, _declaration(components, backs, floor=floor, scale=scale)) + except ValueError as error: + assert "int64_overflow" in str(error) + expected = _independent(columns, components, backs) + assert any(v is not None and not INT64_MIN <= v <= INT64_MAX for v in expected) + return None + + +@PROPERTY +@given(case=measures()) +def test_property_differential_against_an_independent_integer_sum(case) -> None: + columns, components, backs, _, _ = case + result = _run_or_overflow(columns, components, backs) + if result is None: + return + out, _ = result + assert _measure(out) == _independent(columns, components, backs) + + +@PROPERTY +@given(case=measures()) +def test_property_components_plus_add_back_equal_the_measure(case) -> None: + columns, components, backs, _, _ = case + result = _run_or_overflow(columns, components, backs) + if result is None: + return + out, receipt = result + measure = _measure(out) + resolved = [m for m in measure if m is not None] + assert receipt["sum_measure"] == sum(resolved) + assert receipt["sum_measure"] == sum( + row["unweighted_signed_total"] for row in receipt["components"] + ) + sum(row["total"] for row in receipt["add_backs"]) + for index, value in enumerate(measure): + if value is None: + continue + signed = [sign * columns[name][index] for name, sign in components] + back = sum( + max(0, -(sign * columns[name][index])) + for name, sign in components + if name in backs + ) + assert sum(signed) + back == value + + +@PROPERTY +@given(case=measures()) +def test_property_the_universe_statuses_partition_the_records(case) -> None: + columns, components, backs, _, _ = case + result = _run_or_overflow(columns, components, backs) + if result is None: + return + out, receipt = result + statuses = out["universe_status"].tolist() + assert set(statuses) <= set(UNIVERSE_STATES) + assert sum(receipt["universe_statuses"].values()) == len(statuses) + for value, status in zip(_measure(out), statuses, strict=True): + expected = ( + "unresolved" if value is None else "included" if value > 0 else "excluded" + ) + assert status == expected + + +@PROPERTY +@given(case=measures(), seed=st.integers(min_value=0, max_value=2**32 - 1)) +def test_property_the_output_does_not_depend_on_component_order(case, seed) -> None: + columns, components, backs, _, _ = case + result = _run_or_overflow(columns, components, backs) + if result is None: + return + shuffled = list(components) + random.Random(seed).shuffle(shuffled) + other, _ = _run(columns, _declaration(shuffled, list(reversed(backs)))) + out, _ = result + for role in OUTPUTS: + pd.testing.assert_series_equal(out[role], other[role]) + + +@PROPERTY +@given(case=measures()) +def test_property_the_unit_is_max_floor_x_over_scale(case) -> None: + columns, components, backs, floor, scale = case + result = _run_or_overflow(columns, components, backs, floor, scale) + if result is None: + return + out, _ = result + for value, unit in zip(_measure(out), out["unit"].tolist(), strict=True): + if value is None: + assert np.isnan(unit) + else: + assert unit == float(Fraction(max(floor, value), scale)) 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..e2249440bb 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_sixteen_kernels(): registry = KernelRegistry() assert register_transport_population_kernels(registry) is registry assert register_transport_population_kernels(registry) is registry @@ -148,11 +148,13 @@ def test_aggregate_population_registration_installs_exactly_fourteen_kernels(): "money.normalize@1", "money.minor_units@1", "money.annual_transfers@1", + "measure.combine@1", "transport.create@1", "transport.create@2", "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