diff --git a/changelog.d/1273-uk-brma-clone-spread.changed.md b/changelog.d/1273-uk-brma-clone-spread.changed.md new file mode 100644 index 0000000000..4bdc339012 --- /dev/null +++ b/changelog.d/1273-uk-brma-clone-spread.changed.md @@ -0,0 +1 @@ +The UK local pool spreads each spine household's BRMA across its K clones: `uk.full.brma_table` computes each household's exact single-draw BRMA distribution, an LHA-rate order over its region's BRMAs and a lineage-keyed uniform before expansion, and `uk.full.brma` gives clone c the BRMA at quantile `(u + c/K) mod 1`, so each clone keeps the single-draw distribution while the clones average its noise. A terminal `uk_local_brma_enum_domain` gate rechecks the rewritten values; the spine's `frs_brma` stage and its pins are unchanged. diff --git a/docs/uk-full-build-graph.md b/docs/uk-full-build-graph.md index bc11e51252..1890e5d4de 100644 --- a/docs/uk-full-build-graph.md +++ b/docs/uk-full-build-graph.md @@ -24,6 +24,12 @@ uv run --no-sync python tools/build_uk_full.py --release-role dense \ The three atomic-area supports are the artifacts pinned in `uk/uk_atomic_area_supports.provenance.json` and `uk/spec/sources.yaml` (built by `tools/build_uk_atomic_area_supports.py` from the published ONS, NRS and NISRA lookups; publisher registration PolicyEngine/chronicle#269). `--geography-assignment atomic` is the default and requires all three; `--atomic-support-sha256-{ew,scotland,ni}` pin them like `--ladder-sha256`, and a release candidate requires all four pins. `--geography-assignment legacy` keeps the previous sequential ladder draw for measurement builds, takes no supports, and is refused by `--release-candidate`. +The local pool spreads each spine household's BRMA draw across its K clones (K = 15 in the local doctrine). Before expansion, `uk.full.brma_table` reads the normalized population and engine: it recomputes benefit-unit `LHA_category`, averages those units' census private-rented-household BRMA distributions within each household, and orders the household's regional BRMAs by the engine's `uncapped_BRMA_LHA_rate` (falling back to `BRMA_LHA_rate` when the pinned engine lacks the uncapped variable). A stable uniform `u_h` is keyed by the spine household's lineage before geographic cloning, with seed 0 and salt `brma:clone_spread`; SPI and CGT copies therefore have their own uniforms. After expansion, the engine-free `uk.full.brma` owns the household `brma` rewrite: clone c takes the BRMA at quantile `(u_h + c / K) mod 1` in that rate order. Each clone retains exactly the household's single-draw marginal distribution, while the clone mean spreads the draw over K evenly spaced quantiles. This step changes only `brma`: it adds no rows and changes no weights or other inputs. + +The assembled spine enum check stays in place, and the local terminal battery separately runs `uk_local_brma_enum_domain` on the rewritten values before export. Invalid BRMAs or missing domain evidence stop development and release-candidate builds. The spine's `frs_brma` stage, source-stage and coverage pins, and `uk_spine.json` contract remain unchanged; the national line continues to use the spine's original BRMA draws. `uc_reporter_redraw` screened reporters using those original spine BRMAs. Local clone BRMAs are drawn within the household's FRS region and are not conditioned on the clone's assigned local authority. + +Clone weights are still calibrated independently. The solve can put different weights on a household's BRMA alternatives, so the variance reduction at fixed weights does not establish equivalent outcomes after calibration. In uk-data, independently calibrating split records shifted outcomes compared with calibrating them as one household. A follow-up should add per-household clone weight tying to `microcosm-calibrate`, which currently has no tie-groups feature, and then measure the solved outcomes. PR tests check the draw and graph contracts; calibrated outcome comparisons and native certification require separate measured build evidence. + The checkpoint must bind the exact frame content, the current spine stage roster and the gate-report bytes. A spine build records its H5's file digest in the sidecar (`output {filename, sha256, size_bytes}`), in a `
.sha256` file beside it, in its `spine_h5_creation` stage event and in its Logbook `pipeline` verdict (`artifact_sha256`); when the sidecar carries `output.sha256`, the full build refuses an `--input-h5` whose measured digest differs (older sidecars without the key bind by content identity alone). The bound-spine node compares the checkpoint's gate report digests with the branch's own gate declarations and refuses a spine whose gate manifest differs from them, so `--input-h5` needs a spine built by a branch with the same declarations; every acceptance spine on disk when this registration landed predates them and is not admitted. Historical candidate H5 files and reviewed-bypass sidecars are not alternate build sources. Chronicle facts and manifest must match the independently reviewed national and local feed declarations; filtering targets does not relax source validation. The target registry binds Census household and demographic rows from the reviewed Chronicle feed, including the approved Northern Ireland constituency geography. Geography is assigned after expansion by the shared atomic-geography operators (microcosm#931): `uk.full.identity` keys every household with `household_draw_key` from the spine's explicit lineage (source household id, SPI support channel and clone index, CGT clone and donor flags) and the pool clone index; `uk.full.geography.assign` draws one atomic area per household (E&W 2021 Output Area, Scotland 2022 Output Area, NI 2021 Data Zone) by census household count within the household's FRS region, with a keyed `sha256-u53-v1` stream so a household's draw never depends on row order, K or any other household; `uk.full.geography.derive` reads every larger geography off the support's versioned mappings; `uk.full.geography.local_authority` resolves the engine's `local_authority` key from the derived authority code (the same resolver the ladder path uses); the shared `uk.full.geography.gate` and the UK distribution gate `uk.full.geography_gate` follow, and `uk.full.pool` refuses unless the shared gate passed. The OA ladder still supplies the constituency and local-authority rosters, household dispersion and lookup support for target compilation; its household counts are not a second source of calibration targets. Source receipts retain the Chronicle identity, the paired ladder digest and the geography binding (assignment mode, definition sha256, support pins, identity column, stream), so target values and the geography used to assign households can be audited separately. @@ -68,7 +74,9 @@ The spine graph is the 40-stage source roster declared by the UK spec, built by | Assembled and transferred spine gates | `spine.gates.assembled`, `spine.gates.transferred` | | Bound checkpoint admission | `uk.full.spine_checkpoint`, when resuming a saved spine | | Pool sample and mass normalization | `uk.full.sample`, `uk.full.normalize` | +| Household BRMA distributions, engine LHA-rate order and lineage uniforms | `uk.full.brma_table` | | Linked entity expansion and ancestry | `uk.full.expand`, `uk.full.expand.owned` | +| BRMAs at evenly spaced clone quantiles, with an owned household rewrite | `uk.full.brma` | | Post-clone household identity | `uk.full.identity` | | Atomic-area support import, identity-keyed assignment, derivation and shared integrity gate | `uk.full.geography.support.{0,1,2}`, `uk.full.geography.assign`, `uk.full.geography.derive`, `uk.full.geography.gate` (shared `geography.*@1` kernels) | | Engine local-authority key from the derived authority code | `uk.full.geography.local_authority` | @@ -89,7 +97,7 @@ The spine graph is the 40-stage source roster declared by the UK spec, built by The modules under `uk_runtime` divide the roster: `graph_build` composes the graph over the bound spine, `graph_population` owns sample through the geography gate, `graph_targets` owns target compilation through the problem, `graph_calibration` owns the dense solve, the size nodes and the calibrated population, `graph_terminal` owns the gate nodes, the holdout, the export and the package, `full_certification` owns the terminal certification node, and `full_build_cli` resolves the request, executes graph endpoints and atomically materialises their stored artifacts. -`operations.json` is generated from the compiled graph, including actual dependencies and artifact owners. It is the execution inventory, rather than a second manually maintained pipeline roster. Composite source/model stages preserve their existing numerical boundary; for example, WAS retains its joint donor/recipient encoding dependency. This registration does not change imputation order, RNG consumption or clone IDs; the geography draw is the identity-keyed atomic-area assignment of microcosm#931 (measured in `experiments/931-uk-atomic-assignment-receipts.md`), with the sequential ladder draw retained behind `--geography-assignment legacy` for measurement builds. +`operations.json` is generated from the compiled graph, including actual dependencies and artifact owners. It is the execution inventory, rather than a second manually maintained pipeline roster. Composite source/model stages preserve their existing numerical boundary; for example, WAS retains its joint donor/recipient encoding dependency. Spine imputation order and clone IDs are unchanged; local BRMA spreading has its own lineage-keyed stream. The geography draw is the identity-keyed atomic-area assignment of microcosm#931 (measured in `experiments/931-uk-atomic-assignment-receipts.md`), with the sequential ladder draw retained behind `--geography-assignment legacy` for measurement builds. Shared machinery includes graph execution/storage/replay, typed artifacts, explicit same-kind weight updates, target selection receipts, ordered sparse calibration problem/solution/result codecs, exact-count selection, atomic artifact materialization and bundle publication. UK adapters retain source interpretation, entity relationships, geography mappings, measure bindings and gate prescriptions. diff --git a/packages/microcosm-build/src/microcosm/build/uk/gates.json b/packages/microcosm-build/src/microcosm/build/uk/gates.json index 87b31fc261..ee01184c37 100644 --- a/packages/microcosm-build/src/microcosm/build/uk/gates.json +++ b/packages/microcosm-build/src/microcosm/build/uk/gates.json @@ -1261,6 +1261,19 @@ "parameters": {}, "notes": "The existing UK geography-ladder evaluator is rerun on the locally calibrated candidate weights and adapted through UKGateBinding; no second ladder verdict implementation is introduced." }, + { + "id": "uk_local_brma_enum_domain", + "gate": "enum_domain", + "phase": "terminal", + "criticality": "release_blocking", + "evidence_absent_blocks": true, + "parameters": { + "columns": [ + "brma" + ] + }, + "notes": "The local clone-spread node rewrites household brma after the spine's assembled enum check. Recheck the actual local candidate against the PolicyEngine-UK brma enum domain before export; an invalid value or missing domain evidence is structural and stops even a development build. The separate local entry preserves the assembled spine gate declaration and its checkpoint identity." + }, { "id": "uk_local_area_support", "gate": "area_support", diff --git a/packages/microcosm-build/src/microcosm/build/uk_runtime/brma_spread.py b/packages/microcosm-build/src/microcosm/build/uk_runtime/brma_spread.py new file mode 100644 index 0000000000..c2bfa59f26 --- /dev/null +++ b/packages/microcosm-build/src/microcosm/build/uk_runtime/brma_spread.py @@ -0,0 +1,383 @@ +"""Spread each normalized UK household's BRMA draw across its local clones. + +The household distribution is the mean of its benefit units' census cells, +the exact marginal of ``frs_brma``'s unit draws and uniform household pick. +For clone c of K, invert that distribution at ``(u_h + c / K) mod 1`` in +increasing household-mean LHA-rate order. The uniform is keyed to lineage +before geographic expansion, including the SPI and CGT branches. + +Invariants: +1. Exact marginals: over uniform u_h, every clone has the household's BRMA + distribution, for any rate order. +2. Support: only positive-probability BRMAs in that household's region occur. +3. Spacing: K clones have K evenly spaced quantiles. For an outcome monotone + in the walk order, their mean differs from its expectation by at most + the outcome's range divided by K. +4. Determinism: the same lineage keys give the same draws regardless of row + order; source-support branches have distinct keys. +5. Identity: K=1 retains the household distribution and supported clone 0. +6. Conservation: assignment returns only BRMA values; it changes no other + column, weight, or row count and never mutates its input frame or table. + +This ports the quantile logic from policyengine-uk-data without importing it. +The preparation uses one engine materialization per BRMA slot, avoiding a +frame stacked across all slots; the probability tables still use O(H * W) +memory. Applying the prepared distribution +after geographic expansion requires no country engine. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from numbers import Integral +from typing import Any + +import numpy as np +import pandas as pd + +from microcosm.build.stochastic_assignment import stable_identity_uniforms +from microcosm.build.uk_runtime.frs_brma import ( + _benunit_household_map, + _enum_name, + load_brma_count_resource, +) +from microcosm.build.uk_runtime.national_frame import uk_time_period +from microcosm.frame import Frame +from microcosm.frame.rules import assert_rules_engine_country + + +def _benunit_household_indices(frame: Frame) -> np.ndarray: + household = frame.table("household") + placements = _benunit_household_map(frame.table("person")) + positions = pd.Series(np.arange(len(household)), index=household["household_id"]) + indices = frame.table("benunit")["benunit_id"].map(placements).map(positions) + if indices.isna().any(): + raise ValueError("Every benefit unit needs a household for BRMA spreading.") + return indices.to_numpy(dtype=np.int64) + + +def household_brma_probabilities( + frame: Frame, + *, + lha_category: Sequence[object], + count_resource: Mapping[str, Any], +) -> tuple[dict[str, list[str]], np.ndarray]: + """Return regional alphabetical BRMAs and household-by-BRMA probabilities. + + Cells are the committed ``frs_brma`` census resource. Each unit's cell is + normalized separately before averaging over its household's units. Rows + are aligned to the household table, columns to its region's BRMA list; + columns beyond the region's list have probability zero. + """ + categories = np.asarray([_enum_name(value) for value in lha_category]) + if categories.shape != (frame.n("benunit"),): + raise ValueError("LHA_category materialization must align to benunit rows.") + cells = count_resource["cells"] + brmas: dict[str, list[str]] = {} + for region, region_cells in cells.items(): + names = set() + for counts in region_cells.values(): + for name, count in counts.items(): + count = float(count) + if not np.isfinite(count) or count < 0: + raise ValueError("BRMA counts must be finite and non-negative.") + if count > 0: + names.add(str(name)) + if names: + brmas[str(region)] = sorted(names) + if not brmas: + raise ValueError("BRMA count resource has no positive counts.") + width = max(map(len, brmas.values())) + household = frame.table("household") + regions = np.asarray([_enum_name(value) for value in household["region"]]) + indices = _benunit_household_indices(frame) + units = np.bincount(indices, minlength=len(household)) + if (units == 0).any(): + raise ValueError("Every household needs at least one benefit unit.") + probabilities = np.zeros((len(household), width), dtype=float) + unit_regions = regions[indices] + for region, category in sorted(set(zip(unit_regions, categories, strict=True))): + try: + counts = cells[region][category] + except KeyError as exc: + raise KeyError( + f"missing BRMA count-table cell for region={region!r}, " + f"LHA_category={category!r}." + ) from exc + total = sum(float(count) for count in counts.values()) + if not np.isfinite(total) or total <= 0: + raise ValueError("BRMA count-table cell has no positive counts.") + p = np.zeros(width) + names = brmas[region] + p[: len(names)] = [float(counts.get(name, 0)) / total for name in names] + mask = (unit_regions == region) & (categories == category) + np.add.at(probabilities, indices[mask], p) + return brmas, probabilities / units[:, None] + + +def lha_rate_keys( + frame: Frame, + *, + engine: object, + brmas: Mapping[str, Sequence[str]], + benunit_household: np.ndarray | None = None, +) -> np.ndarray: + """Materialize each household's mean benefit-unit LHA rate per BRMA. + + Use ``uncapped_BRMA_LHA_rate`` where the engine defines it, otherwise + ``BRMA_LHA_rate``. Engine ``variables()`` lists inputs only, so inspect + variable metadata to detect this computed variable. Padded slots get inf. + """ + assert_rules_engine_country(engine, "uk") + variable = "uncapped_BRMA_LHA_rate" + try: + metadata = engine.variable_metadata(variable) + except (KeyError, ValueError): + variable = "BRMA_LHA_rate" + metadata = engine.variable_metadata(variable) + if metadata.entity != "benunit": + raise ValueError(f"{variable} must be a benefit-unit variable.") + household = frame.table("household") + regions = np.asarray([_enum_name(value) for value in household["region"]]) + indices = ( + _benunit_household_indices(frame) + if benunit_household is None + else np.asarray(benunit_household) + ) + if ( + indices.shape != (frame.n("benunit"),) + or indices.dtype.kind not in "iu" + or (indices < 0).any() + or (indices >= len(household)).any() + ): + raise ValueError("Benefit-unit household indices must align to the frame.") + units = np.bincount(indices, minlength=len(household)) + if (units == 0).any(): + raise ValueError("Every household needs at least one benefit unit.") + if not brmas or any(not names for names in brmas.values()): + raise ValueError("Every BRMA region needs a non-empty BRMA list.") + try: + count = np.asarray([len(brmas[region]) for region in regions]) + except KeyError as exc: + raise ValueError("A household's region has no BRMA support.") from exc + width = max(map(len, brmas.values())) + keys = np.full((len(household), width), np.inf) + for slot in range(width): + moved = household.copy() + moved["brma"] = [ + brmas[region][min(slot, size - 1)] + for region, size in zip(regions, count, strict=True) + ] + # A fresh carrier preserves every stored input, weight and receipt. + # Each materialize call starts its own simulation, so cached LHA + # calculations from another BRMA cannot contaminate this slot. + tables = {entity: frame.table(entity) for entity in frame.entities} + tables.update({link: frame.link(link) for link in frame.links}) + tables["household"] = moved + candidate = Frame( + tables=tables, + schema=frame.schema, + weights={ + entity: frame.weights_for(entity) for entity in frame.weighted_entities + }, + strata=frame.strata, + mass_log=frame.mass_log, + metadata=frame.metadata, + ) + rate = np.asarray( + engine.materialize(candidate, [variable], uk_time_period(frame))[variable], + dtype=float, + ) + if rate.shape != (frame.n("benunit"),) or not np.isfinite(rate).all(): + raise ValueError( + f"{variable} must give finite rates aligned to benefit units." + ) + mean = np.bincount(indices, weights=rate, minlength=len(household)) / units + keys[:, slot] = np.where(slot < count, mean, np.inf) + del candidate + return keys + + +def _ordered_cdf( + probabilities: np.ndarray, keys: np.ndarray | None +) -> tuple[np.ndarray, np.ndarray]: + probabilities = np.asarray(probabilities, dtype=float) + if ( + probabilities.ndim != 2 + or probabilities.shape[1] == 0 + or not np.isfinite(probabilities).all() + or (probabilities < 0).any() + or not np.allclose(probabilities.sum(axis=1), 1, rtol=0, atol=1e-12) + ): + raise ValueError("Probabilities must be finite non-negative rows summing to 1.") + rows, width = probabilities.shape + if keys is None: + order = np.broadcast_to(np.arange(width), probabilities.shape) + else: + keys = np.asarray(keys, dtype=float) + if ( + keys.shape != probabilities.shape + or not np.isfinite(keys[probabilities > 0]).all() + ): + raise ValueError("Keys must align and be finite on positive BRMA support.") + order = np.argsort( + np.where(probabilities > 0, keys, np.inf), axis=1, kind="stable" + ) + p = np.take_along_axis(probabilities, order, axis=1) + cumulative = np.cumsum(p, axis=1) + last = width - 1 - np.argmax(p[:, ::-1] > 0, axis=1) + cumulative[np.arange(width)[None, :] >= last[:, None]] = 1.0 + return order, cumulative + + +def _validate_quantiles(quantiles: np.ndarray) -> np.ndarray: + quantiles = np.asarray(quantiles, dtype=float) + if ( + quantiles.ndim != 1 + or not np.isfinite(quantiles).all() + or np.any((quantiles < 0) | (quantiles >= 1)) + ): + raise ValueError("Quantiles must be a vector in [0, 1).") + return quantiles + + +def brmas_at_quantiles( + probabilities: np.ndarray, keys: np.ndarray | None, quantiles: np.ndarray +) -> np.ndarray: + """Return each row's positive-probability column at its quantile. + + Walk ascending keys, breaking ties by the alphabetical input column. + The CDF is exactly 1 from the last positive column onwards, preventing + floating-point rounding from selecting a zero-probability padded column. + """ + quantiles = _validate_quantiles(quantiles) + order, cumulative = _ordered_cdf(probabilities, keys) + if quantiles.shape != (len(order),): + raise ValueError("Quantiles must align to probability rows.") + position = (cumulative <= quantiles[:, None]).sum(axis=1) + return order[np.arange(len(order)), position] + + +def spaced_quantiles(uniforms: np.ndarray, k: int) -> np.ndarray: + """Return ``(u_h + c / k) mod 1`` for each household and clone c.""" + if not isinstance(k, Integral) or isinstance(k, bool) or k < 1: + raise ValueError("k must be a positive integer.") + uniforms = _validate_quantiles(uniforms) + return (uniforms[:, None] + np.arange(k)[None, :] / k) % 1.0 + + +def prepare_brma_spread( + frame: Frame, + *, + engine: object, + lineage_keys: Sequence[str], + count_resource: Mapping[str, Any] | None = None, +) -> dict[str, Any]: + """Prepare JSON-safe distributions before geographic expansion. + + The host supplies pre-geographic lineage keys from ``household_draw_key``; + distinct SPI/CGT rows must have distinct keys, even with a shared source + household. Names in zero-probability padding repeat a regional BRMA and + cannot be selected. No engine work remains when applying this payload. + """ + assert_rules_engine_country(engine, "uk") + if len(lineage_keys) != frame.n("household") or any( + not isinstance(key, str) or not key for key in lineage_keys + ): + raise ValueError( + "Lineage keys must be non-empty strings aligned to households." + ) + if len(set(lineage_keys)) != len(lineage_keys): + raise ValueError("Every normalized household needs a distinct lineage key.") + resource = load_brma_count_resource() if count_resource is None else count_resource + category = engine.materialize(frame, ["LHA_category"], uk_time_period(frame))[ + "LHA_category" + ] + brmas, probabilities = household_brma_probabilities( + frame, lha_category=category, count_resource=resource + ) + keys = lha_rate_keys(frame, engine=engine, brmas=brmas) + order, cumulative = _ordered_cdf(probabilities, keys) + household = frame.table("household") + width = probabilities.shape[1] + names = np.asarray( + [ + brmas[_enum_name(region)] + + [brmas[_enum_name(region)][-1]] * (width - len(brmas[_enum_name(region)])) + for region in household["region"] + ], + dtype=object, + ) + return { + "household_ids": household["household_id"].to_list(), + "ordered_brmas": np.take_along_axis(names, order, axis=1).tolist(), + "cumulative_probabilities": cumulative.tolist(), + "uniforms": stable_identity_uniforms( + lineage_keys, seed=0, salt="brma:clone_spread" + ).tolist(), + } + + +def assign_clone_brmas( + household: pd.DataFrame, + payload: Mapping[str, Any], + *, + n_clones: int, + id_multiplier: int, + clone_index_column: str, +) -> np.ndarray: + """Return clone-aligned BRMAs, preserving the table and every other value. + + Geographic expansion owns ``household_id = pre_id + c * id_multiplier`` + and the clone index column. Recover its parent by that exact integer + relationship, without assuming any row ordering. Process one clone at a + time to avoid allocating an expanded-household-by-BRMA matrix. + """ + if ( + not isinstance(id_multiplier, Integral) + or isinstance(id_multiplier, bool) + or id_multiplier <= 0 + ): + raise ValueError("id_multiplier must be a positive integer.") + uniforms = _validate_quantiles(payload["uniforms"]) + # Validate K without allocating the full household-by-clone grid. + spaced_quantiles(np.empty(0), n_clones) + ids = pd.Index(payload["household_ids"]) + if len(ids) != len(uniforms) or ids.has_duplicates or ids.dtype.kind not in "iu": + raise ValueError( + "Prepared household IDs must be unique integers aligned to draws." + ) + names = np.asarray(payload["ordered_brmas"], dtype=object) + cumulative = np.asarray(payload["cumulative_probabilities"], dtype=float) + if ( + cumulative.ndim != 2 + or cumulative.shape[0] != len(ids) + or cumulative.shape[1] == 0 + or names.shape != cumulative.shape + or not np.isfinite(cumulative).all() + or (cumulative < 0).any() + or (cumulative > 1).any() + or (np.diff(cumulative, axis=1) < 0).any() + or (cumulative[:, -1] != 1).any() + ): + raise ValueError("Prepared BRMA names and cumulative probabilities must align.") + clones = household[clone_index_column].to_numpy() + household_ids = household["household_id"].to_numpy() + if ( + clones.dtype.kind not in "iu" + or household_ids.dtype.kind not in "iu" + or (clones < 0).any() + or (clones >= n_clones).any() + ): + raise ValueError("Geographic clone IDs and indices must be valid integers.") + assigned = np.empty(len(household), dtype=object) + for clone in np.unique(clones): + rows = np.flatnonzero(clones == clone) + parents = household_ids[rows] - int(clone) * int(id_multiplier) + positions = ids.get_indexer(parents) + if (positions < 0).any(): + raise ValueError("A geographic clone has no prepared parent household.") + quantiles = (uniforms[positions] + int(clone) / n_clones) % 1.0 + columns = (cumulative[positions] <= quantiles[:, None]).sum(axis=1) + assigned[rows] = names[positions, columns] + return assigned diff --git a/packages/microcosm-build/src/microcosm/build/uk_runtime/calibration_run.py b/packages/microcosm-build/src/microcosm/build/uk_runtime/calibration_run.py index b9f38bd871..0a66cbf3be 100644 --- a/packages/microcosm-build/src/microcosm/build/uk_runtime/calibration_run.py +++ b/packages/microcosm-build/src/microcosm/build/uk_runtime/calibration_run.py @@ -80,6 +80,7 @@ UK_LOCAL_GATE_SCOPE = ( "uk_local_geography_ladder_post_calibration", + "uk_local_brma_enum_domain", "uk_local_area_support", "uk_local_target_fit", "uk_local_per_family_fit", @@ -130,7 +131,8 @@ # declared by the stage module (microcosm#725). "uk_capital_gains_asset_type_enum_domain", # Weight-independent, and its column exists from frs_brma onward, so the - # spine checks it at the assembled boundary instead of the release end. + # spine checks it at the assembled boundary. Local cloning owns a separate + # enum check after its BRMA rewrite; the spine declaration stays unchanged. "uk_brma_enum_domain", ) diff --git a/packages/microcosm-build/src/microcosm/build/uk_runtime/full_gates.py b/packages/microcosm-build/src/microcosm/build/uk_runtime/full_gates.py index 1bc886834b..9a4a4e7009 100644 --- a/packages/microcosm-build/src/microcosm/build/uk_runtime/full_gates.py +++ b/packages/microcosm-build/src/microcosm/build/uk_runtime/full_gates.py @@ -311,7 +311,8 @@ def classify_full_gate_outcomes( The combined driver exports a diagnostic candidate after failures of its five local fit/support/weight checks, but stops on a non-passing geography - ladder. Below f100 those five checks are recorded without enforcement. + ladder or local BRMA domain check. Below f100 those five statistical checks + are recorded without enforcement. All newly included national/source checks keep the shared battery's own BLOCKS_ARTIFACT policy, including its declared missing-evidence rules. """ @@ -322,8 +323,11 @@ def classify_full_gate_outcomes( raise TypeError("Classify a validated GatePhaseReport, not unbound JSON.") if not 0.0 < sample_fraction <= 1.0: raise ValueError("sample_fraction must be in (0, 1].") - geography = "uk_local_geography_ladder_post_calibration" - local_export_exception = set(UK_LOCAL_GATE_SCOPE) - {geography} + structural_local = { + "uk_local_geography_ladder_post_calibration", + "uk_local_brma_enum_domain", + } + local_export_exception = set(UK_LOCAL_GATE_SCOPE) - structural_local blocking = { outcome.entry.id for outcome in report.blocking_outcomes(release_candidate=release_candidate) @@ -331,7 +335,7 @@ def classify_full_gate_outcomes( stop, exported_failures, unenforced, diagnostic = [], [], [], [] for outcome in report.outcomes: gate_id = outcome.entry.id - if gate_id == geography and outcome.status is not GateStatus.PASSED: + if gate_id in structural_local and outcome.status is not GateStatus.PASSED: stop.append(gate_id) continue if outcome.status not in {GateStatus.FAILED, GateStatus.EVIDENCE_ABSENT}: diff --git a/packages/microcosm-build/src/microcosm/build/uk_runtime/graph_population.py b/packages/microcosm-build/src/microcosm/build/uk_runtime/graph_population.py index a0a56481d9..2df301c072 100644 --- a/packages/microcosm-build/src/microcosm/build/uk_runtime/graph_population.py +++ b/packages/microcosm-build/src/microcosm/build/uk_runtime/graph_population.py @@ -7,6 +7,7 @@ from __future__ import annotations +import hashlib import importlib.util import json from collections.abc import Mapping @@ -15,7 +16,9 @@ import numpy as np import pandas as pd +from microcosm.build import stochastic_assignment from microcosm.frame import Frame, Weights +from microcosm.frame.adapters import policyengine_uk as uk_engine_adapter from microcosm.graph import ( ArtifactInput, ArtifactOutput, @@ -45,6 +48,8 @@ from . import ( atomic_area_support, atomic_household_identity, + brma_spread, + frs_brma, geography_ladder, local_authority_input, national_sampling, @@ -68,12 +73,14 @@ POPULATION_RECEIPT_TYPE = ArtifactType("microcosm.uk.population-receipt", 1) LOCATION_DRAW_TYPE = ArtifactType("microcosm.uk.ladder-location-draw", 1) +BRMA_TABLE_TYPE = ArtifactType("microcosm.uk.brma-spread-table", 1) GEOGRAPHY_GATE_TYPE = ArtifactType("microcosm.uk.geography-gate", 1) GEOGRAPHY_ASSIGNMENTS = ("atomic", "legacy") #: Lineage columns the post-clone household identity is keyed on. The support -#: channel, the CGT support split's copy flag and index (microcosm#1045) and -#: the clone flag come from the spine stages; the geographic clone index from -#: ``uk.full.expand``. Calibration K, weights and remapped ids are not inputs. +#: channel, the CGT support split's copy flag and index (microcosm#1045), and +#: the incidence and residential-arm lineage come from the spine stages; +#: the geographic clone index comes from ``uk.full.expand``. Calibration K, +#: weights and remapped ids are not inputs. UK_GEOGRAPHY_IDENTITY_INPUTS = ( "source_household_id", "household_support_channel", @@ -309,6 +316,88 @@ def run(self, context: KernelContext) -> KernelResult: ) +class UKBRMATableKernel(_PopulationKernel): + """Compute rate-ordered household mixtures once, before K-way expansion.""" + + ref = "uk.full.brma_table@1" + capabilities = Capabilities( + Determinism.DETERMINISTIC, + dependencies=("policyengine-uk",), + seed_source=SeedSource.PARAM, + ) + + def __init__(self, *, engine=None, count_resource=None): + self.engine = engine + self.count_resource = count_resource + + def implementation_hash(self) -> str: + implementation = source_hash( + type(self), + brma_spread, + frs_brma, + stochastic_assignment, + uk_engine_adapter, + atomic_household_identity, + dependencies=self.capabilities.dependencies, + ) + resource = self.count_resource + if resource is None: + resource = frs_brma.load_brma_count_resource() + return hashlib.sha256( + canonical_json( + {"implementation": implementation, "count_resource": resource} + ) + ).hexdigest() + + def run(self, context: KernelContext) -> KernelResult: + frame = context_frame(context) + keys = household_lineage_keys( + frame.table("household"), + source=str(context.params["source"]), + source_vintage=str(context.params["source_vintage"]), + geographic=False, + ) + payload = brma_spread.prepare_brma_spread( + frame, + engine=self.engine + if self.engine is not None + else uk_engine_adapter.PolicyEngineUKEngine(), + lineage_keys=keys, + count_resource=self.count_resource, + ) + return KernelResult(artifacts={"brma_table": canonical_json(payload)}) + + +class UKBRMACloneSpreadKernel(_PopulationKernel): + """Rewrite only household BRMA at the clone's evenly spaced quantile.""" + + ref = "uk.full.brma@1" + capabilities = Capabilities(Determinism.DETERMINISTIC) + + def implementation_hash(self) -> str: + return source_hash(type(self), brma_spread) + + def run(self, context: KernelContext) -> KernelResult: + payload = json.loads(context.artifacts["brma_table"].payload) + expansion = json.loads(context.artifacts["expansion"].payload) + household = context.tables["household"] + names = brma_spread.assign_clone_brmas( + household, + payload, + n_clones=int(expansion["n_clones"]), + id_multiplier=int(expansion["id_multiplier"]), + clone_index_column=ladder_clone_index_column("household"), + ) + index = pd.Index(household["household_id"], name="household_id") + return KernelResult( + columns={ + ("household", "brma"): pd.Series( + names, index=index, name="brma", dtype=dtype_for_token("string") + ) + } + ) + + class UKLocationDrawKernel(_PopulationKernel): ref = "uk.full.locations@1" capabilities = Capabilities(Determinism.DETERMINISTIC, seed_source=SeedSource.PARAM) @@ -365,6 +454,111 @@ def run(self, context: KernelContext) -> KernelResult: ) +def household_lineage_keys( + household: pd.DataFrame, *, source: str, source_vintage: str, geographic: bool +) -> list[str]: + """Project spine ancestry, optionally including the geographic clone step. + + The BRMA offset uses the same source/SPI/CGT path as geography, before + geographic support is appended. It never depends on remapped IDs or K. + """ + + inputs = ( + UK_GEOGRAPHY_IDENTITY_INPUTS + if geographic + else UK_GEOGRAPHY_IDENTITY_INPUTS[:-1] + ) + missing = [c for c in inputs if c not in household] + if missing: + raise ValueError(f"Geography identity lacks lineage column(s) {missing}.") + lineage = household.loc[:, list(inputs)] + if lineage.isna().any().any(): + raise ValueError("Geography identity refuses null lineage values.") + if lineage["source_household_id"].dtype != np.dtype("int64"): + raise ValueError("Geography identity requires int64 source household ids.") + if lineage["cgt_support_copy_index"].dtype != np.dtype("int64"): + raise ValueError("Geography identity requires an int64 support copy index.") + if lineage["cgt_residential_clone_index"].dtype != np.dtype("int64"): + raise ValueError("Geography identity requires an int64 residential arm index.") + for column in ( + "household_is_cgt_support_copy", + "household_is_capital_gains_clone", + "household_is_cgt_residential_clone", + ): + dtype = lineage[column].dtype + if dtype != np.dtype("bool") and not isinstance(dtype, pd.BooleanDtype): + raise ValueError(f"Geography identity requires a boolean {column}.") + keys = [] + for ( + source_id, + channel, + support_index, + support_copy, + copy_index, + cgt, + residential_arm, + residential_index, + clone_index, + ) in zip( + lineage["source_household_id"].to_numpy(), + lineage["household_support_channel"].to_numpy(), + lineage["household_support_clone_index"].to_numpy(), + lineage["household_is_cgt_support_copy"].to_numpy(), + lineage["cgt_support_copy_index"].to_numpy(), + lineage["household_is_capital_gains_clone"].to_numpy(), + lineage["household_is_cgt_residential_clone"].to_numpy(), + lineage["cgt_residential_clone_index"].to_numpy(), + ( + lineage[ladder_clone_index_column("household")].to_numpy() + if geographic + else np.zeros(len(lineage), dtype=np.int64) + ), + strict=True, + ): + channel = str(channel) + support_index, clone_index = int(support_index), int(clone_index) + copy_index = int(copy_index) + residential_index = int(residential_index) + if residential_index < 0: + raise ValueError("Geography identity refuses negative clone indices.") + if bool(residential_arm) != (residential_index != 0): + raise ValueError("Household residential-arm flag and arm index disagree.") + if channel not in {_BASE_SUPPORT_CHANNEL, _SPI_SUPPORT_CHANNEL}: + raise ValueError(f"Unknown household support channel {channel!r}.") + if support_index < 0 or clone_index < 0 or copy_index < 0: + raise ValueError("Geography identity refuses negative clone indices.") + if (channel == _SPI_SUPPORT_CHANNEL) != (support_index != 0): + raise ValueError( + "Household support channel and support clone index disagree." + ) + if bool(support_copy) != (copy_index != 0): + raise ValueError( + "Household support-copy flag and support copy index disagree." + ) + path = [] + if support_index: + path.append(("spi_support_channel", support_index)) + if copy_index: + path.append(("cgt_support_split", copy_index)) + if bool(cgt): + path.append(("cgt_incidence_clone", 1)) + if residential_index: + path.append(("cgt_residential_split", residential_index)) + if clone_index: + path.append(("geographic_support", clone_index)) + keys.append( + atomic_household_identity.household_draw_key( + source=source, + source_vintage=source_vintage, + source_household_id=int(source_id), + clone_path=tuple(path), + ) + ) + if len(set(keys)) != len(keys): + raise ValueError("Geography identity keys are not unique.") + return keys + + class UKGeographyIdentityKernel(_PopulationKernel): """Key every post-clone household by its source id and ordered clone path. @@ -373,6 +567,7 @@ class UKGeographyIdentityKernel(_PopulationKernel): support clone index when the household is an SPI support copy, ``cgt_support_split`` with the copy index when it is a support copy (microcosm#1045), ``cgt_incidence_clone`` (ordinal 1) from its flag, + ``cgt_residential_split`` with the residential arm index, ``geographic_support`` with the pool clone index when it is not the original copy. Growing K never changes an existing household's key. """ @@ -389,94 +584,9 @@ def run(self, context: KernelContext) -> KernelResult: household = context.tables["household"] source = str(context.params["source"]) vintage = str(context.params["source_vintage"]) - missing = [c for c in UK_GEOGRAPHY_IDENTITY_INPUTS if c not in household] - if missing: - raise ValueError(f"Geography identity lacks lineage column(s) {missing}.") - lineage = household.loc[:, list(UK_GEOGRAPHY_IDENTITY_INPUTS)] - if lineage.isna().any().any(): - raise ValueError("Geography identity refuses null lineage values.") - if lineage["source_household_id"].dtype != np.dtype("int64"): - raise ValueError("Geography identity requires int64 source household ids.") - if lineage["cgt_support_copy_index"].dtype != np.dtype("int64"): - raise ValueError("Geography identity requires an int64 support copy index.") - if lineage["cgt_residential_clone_index"].dtype != np.dtype("int64"): - raise ValueError( - "Geography identity requires an int64 residential arm index." - ) - for column in ( - "household_is_cgt_support_copy", - "household_is_capital_gains_clone", - "household_is_cgt_residential_clone", - ): - dtype = lineage[column].dtype - if dtype != np.dtype("bool") and not isinstance(dtype, pd.BooleanDtype): - raise ValueError(f"Geography identity requires a boolean {column}.") - keys = [] - for ( - source_id, - channel, - support_index, - support_copy, - copy_index, - cgt, - residential_arm, - residential_index, - clone_index, - ) in zip( - lineage["source_household_id"].to_numpy(), - lineage["household_support_channel"].to_numpy(), - lineage["household_support_clone_index"].to_numpy(), - lineage["household_is_cgt_support_copy"].to_numpy(), - lineage["cgt_support_copy_index"].to_numpy(), - lineage["household_is_capital_gains_clone"].to_numpy(), - lineage["household_is_cgt_residential_clone"].to_numpy(), - lineage["cgt_residential_clone_index"].to_numpy(), - lineage[ladder_clone_index_column("household")].to_numpy(), - strict=True, - ): - channel = str(channel) - support_index, clone_index = int(support_index), int(clone_index) - copy_index = int(copy_index) - residential_index = int(residential_index) - if residential_index < 0: - raise ValueError("Geography identity refuses negative clone indices.") - if bool(residential_arm) != (residential_index != 0): - raise ValueError( - "Household residential-arm flag and arm index disagree." - ) - if channel not in {_BASE_SUPPORT_CHANNEL, _SPI_SUPPORT_CHANNEL}: - raise ValueError(f"Unknown household support channel {channel!r}.") - if support_index < 0 or clone_index < 0 or copy_index < 0: - raise ValueError("Geography identity refuses negative clone indices.") - if (channel == _SPI_SUPPORT_CHANNEL) != (support_index != 0): - raise ValueError( - "Household support channel and support clone index disagree." - ) - if bool(support_copy) != (copy_index != 0): - raise ValueError( - "Household support-copy flag and support copy index disagree." - ) - path = [] - if support_index: - path.append(("spi_support_channel", support_index)) - if copy_index: - path.append(("cgt_support_split", copy_index)) - if bool(cgt): - path.append(("cgt_incidence_clone", 1)) - if residential_index: - path.append(("cgt_residential_split", residential_index)) - if clone_index: - path.append(("geographic_support", clone_index)) - keys.append( - atomic_household_identity.household_draw_key( - source=source, - source_vintage=vintage, - source_household_id=int(source_id), - clone_path=tuple(path), - ) - ) - if len(set(keys)) != len(keys): - raise ValueError("Geography identity keys are not unique.") + keys = household_lineage_keys( + household, source=source, source_vintage=vintage, geographic=True + ) index = pd.Index(household["household_id"], name="household_id") column = atomic_area_support.IDENTITY_COLUMN return KernelResult( @@ -712,6 +822,23 @@ def append_uk_population_nodes( description="Install normalized weights without changing their kind.", ) ) + nodes.append( + Node( + id="uk.full.brma_table", + kernel=UKBRMATableKernel.ref, + population="uk.full.normalize", + inputs=population_slices(cells), + params={ + **common, + "source": identity_source, + "source_vintage": source_vintage or str(source_year), + "seed": 0, + "salt": "brma:clone_spread", + }, + artifact_outputs=(ArtifactOutput("brma_table", BRMA_TABLE_TYPE),), + description="Compute household BRMA mixtures and LHA-rate order on normalized spine rows.", + ) + ) expansion_cells = { ("household", "source_household_id"): "int64", ("household", "source_year"): "int64", @@ -770,6 +897,30 @@ def append_uk_population_nodes( description="Declare the lineage and replicate columns materialized by expansion.", ) ) + nodes.append( + Node( + id="uk.full.brma", + kernel=UKBRMACloneSpreadKernel.ref, + population="uk.full.expand", + # Primary keys are implicit in every context table. + inputs=( + Slice( + "household", + (ladder_clone_index_column("household"),), + ), + ), + outputs=(Owned("household", "brma", "string", rewrite=True),), + artifact_inputs=( + ArtifactInput( + "brma_table", "uk.full.brma_table", "brma_table", BRMA_TABLE_TYPE + ), + ArtifactInput( + "expansion", "uk.full.expand", "expansion", POPULATION_RECEIPT_TYPE + ), + ), + description="Spread each household's K BRMA draws at (u + c/K) mod 1 in LHA-rate order.", + ) + ) cells.update(expansion_cells) if geography_assignment == "legacy": geography_nodes, geography_sources = _legacy_geography_nodes( @@ -932,7 +1083,12 @@ def _atomic_geography_nodes( return (identity, *shared[:-1], local_authority, shared[-1], gate), sources -def register_uk_population_kernels(registry: KernelRegistry) -> None: +def register_uk_population_kernels( + registry: KernelRegistry, + *, + engine=None, + brma_count_resource=None, +) -> None: from microcosm.build.graph_atomic_geography import ( register_atomic_geography_kernels, ) @@ -941,6 +1097,8 @@ def register_uk_population_kernels(registry: KernelRegistry) -> None: UKFamilySampleKernel(), UKSampleNormalizationKernel(), UKGeographicExpansionKernel(), + UKBRMATableKernel(engine=engine, count_resource=brma_count_resource), + UKBRMACloneSpreadKernel(), UKGeographyIdentityKernel(), UKLocalAuthorityResolveKernel(), UKLocationDrawKernel(), diff --git a/packages/microcosm-build/src/microcosm/build/uk_runtime/orrery_contract.py b/packages/microcosm-build/src/microcosm/build/uk_runtime/orrery_contract.py index 66787056db..2709ee541a 100644 --- a/packages/microcosm-build/src/microcosm/build/uk_runtime/orrery_contract.py +++ b/packages/microcosm-build/src/microcosm/build/uk_runtime/orrery_contract.py @@ -286,7 +286,8 @@ def checkpoint_references(sidecar_path: Path | None) -> tuple[list[dict], str | "population", "Population pool", "Checkpoint admission, sampling, normalisation, expansion into K " - "geographic copies and the pooled problem checkpoint.", + "geographic copies, BRMA spreading across clones and the pooled problem " + "checkpoint.", ), ( "geography", @@ -321,7 +322,9 @@ def checkpoint_references(sidecar_path: Path | None) -> tuple[list[dict], str | "uk.full.spine_checkpoint", "uk.full.sample", "uk.full.normalize", + "uk.full.brma_table", "uk.full.expand", + "uk.full.brma", "uk.full.pool", ), "geography": ( diff --git a/packages/microcosm-build/src/microcosm/build/uk_runtime/rowwise_posture.py b/packages/microcosm-build/src/microcosm/build/uk_runtime/rowwise_posture.py index 59c6b02160..ec49eb9ff0 100644 --- a/packages/microcosm-build/src/microcosm/build/uk_runtime/rowwise_posture.py +++ b/packages/microcosm-build/src/microcosm/build/uk_runtime/rowwise_posture.py @@ -14,7 +14,7 @@ in-process gates; the seam-shaped build record it emits is what ``tools/certify_uk_release_cut.py`` certifies. ``dense`` is the K-clone joint national + local surface under the local doctrine (:mod:`local_doctrine`) -and the six local gates; its dense pool is the intermediate the exact-K +and the local gate battery; its dense pool is the intermediate the exact-K local-area cuts are drawn from. Every value here is a reviewed constant. Changing one is a doctrine change: diff --git a/packages/microcosm-build/tests/conftest.py b/packages/microcosm-build/tests/conftest.py index 8de2e3dc35..fb13efc55e 100644 --- a/packages/microcosm-build/tests/conftest.py +++ b/packages/microcosm-build/tests/conftest.py @@ -1,7 +1,6 @@ """Shared fixtures: a small person+household frame, an ODS writer, and the autouse guard that keeps every test in this shard away from a live Logbook.""" -import sys from pathlib import Path from zipfile import ZIP_DEFLATED, ZIP_STORED, ZipFile @@ -10,6 +9,9 @@ import pytest from microcosm.frame import EntitySchema, Frame, WeightKind, Weights +from test_support.microcosm_build.logbook_isolation import ( + logbook_urlopen_holders as _logbook_urlopen_holders, +) @pytest.fixture @@ -162,22 +164,6 @@ def _refuse_logbook_network(*_args: object, **_kwargs: object) -> None: ) -def _logbook_urlopen_holders(real_urlopen: object) -> list[object]: - """Every imported module whose ``urlopen`` is the Logbook one. - - ``tools/logbook.py`` binds the function at import - (``from microcosm.build.logbook import urlopen``), so patching the - defining module alone leaves the CLI's copy live. Rather than name the - importers — the next one would be missed — this finds every module - currently holding the same object. - """ - return [ - module - for module in list(sys.modules.values()) - if getattr(module, "urlopen", None) is real_urlopen - ] - - @pytest.fixture(autouse=True) def _isolate_logbook(monkeypatch: pytest.MonkeyPatch) -> dict[str, object]: """Detach every test from any live Logbook the shell may point at. diff --git a/packages/microcosm-build/tests/engine/uk/test_uk_brma_spread.py b/packages/microcosm-build/tests/engine/uk/test_uk_brma_spread.py new file mode 100644 index 0000000000..1d96fa8f23 --- /dev/null +++ b/packages/microcosm-build/tests/engine/uk/test_uk_brma_spread.py @@ -0,0 +1,131 @@ +"""Small UK-engine contract for BRMA-spread mixtures and rate ordering.""" + +import numpy as np +import pandas as pd + +from microcosm.build.stochastic_assignment import stable_identity_uniforms +from microcosm.build.uk_runtime.atomic_household_identity import household_draw_key +from microcosm.build.uk_runtime.brma_spread import lha_rate_keys, prepare_brma_spread +from microcosm.build.uk_runtime.national_frame import uk_national_frame +from microcosm.frame.adapters.policyengine_uk import PolicyEngineUKEngine + + +def test_engine_materializes_household_mean_lha_rates_for_each_regional_brma(): + # Two differently sized units in the London household make aggregation + # across benefit units observable; Scotland has only one BRMA slot. + person = pd.DataFrame( + { + "person_id": [1, 2, 3, 4], + "person_benunit_id": [10, 20, 20, 30], + "person_household_id": [1, 1, 1, 2], + "age": [40, 40, 39, 40], + "is_benunit_head": [True, True, False, True], + "is_parent": True, + } + ) + benunit = pd.DataFrame( + {"benunit_id": [10, 20, 30], "LHA_category": ["B", "C", "B"]} + ) + household = pd.DataFrame( + { + "household_id": [1, 2], + "region": ["LONDON", "SCOTLAND"], + "brma": ["CENTRAL_LONDON", "LOTHIAN"], + "tenure_type": "RENT_PRIVATELY", + "rent": 12_000.0, + # The single-year loader extends and uprates the dataset even + # for a same-year LHA calculation, requiring council tax input. + "council_tax": 0.0, + } + ) + + def frame_for(table): + return uk_national_frame( + person=person, + benunit=benunit, + household=table, + household_weights=[2.0, 3.0], + time_period=2025, + ) + + engine = PolicyEngineUKEngine() + frame = frame_for(household) + original = {entity: frame.table(entity).copy() for entity in frame.entities} + original_weights = frame.weights_for("household").values.copy() + brmas = { + "LONDON": ["CENTRAL_LONDON", "INNER_EAST_LONDON"], + "SCOTLAND": ["LOTHIAN"], + } + keys = lha_rate_keys(frame, engine=engine, brmas=brmas) + variable = "uncapped_BRMA_LHA_rate" + try: + engine.variable_metadata(variable) + except ValueError: + variable = "BRMA_LHA_rate" + assert variable not in engine.variables() + for slot, london_brma in enumerate(brmas["LONDON"]): + moved = household.assign(brma=[london_brma, "LOTHIAN"]) + rate = engine.materialize(frame_for(moved), [variable], 2025)[variable] + assert np.isfinite(rate).all() + assert rate[0] != rate[1] + np.testing.assert_allclose(keys[0, slot], rate[:2].mean()) + if slot == 0: + np.testing.assert_allclose(keys[1, slot], rate[2]) + assert keys[1, 1] == np.inf + + # The engine's actual category materialization must select B/C/B census + # cells. Different cell totals distinguish averaging normalized unit + # distributions from pooling raw counts or averaging across people. + resource = { + "cells": { + "LONDON": { + "B": {"CENTRAL_LONDON": 1, "INNER_EAST_LONDON": 3}, + "C": {"CENTRAL_LONDON": 9, "INNER_EAST_LONDON": 1}, + }, + "SCOTLAND": {"B": {"LOTHIAN": 7}}, + } + } + lineage_keys = [ + household_draw_key( + source="frs", + source_vintage="2023_24", + source_household_id=household_id, + clone_path=(), + ) + for household_id in (1, 2) + ] + payload = prepare_brma_spread( + frame, engine=engine, lineage_keys=lineage_keys, count_resource=resource + ) + assert payload["household_ids"] == [1, 2] + assert payload["ordered_brmas"][0] == [ + brmas["LONDON"][index] for index in np.argsort(keys[0], kind="stable") + ] + known_probabilities = { + "CENTRAL_LONDON": (1 / 4 + 9 / 10) / 2, + "INNER_EAST_LONDON": (3 / 4 + 1 / 10) / 2, + } + probabilities = np.diff( + np.asarray(payload["cumulative_probabilities"]), axis=1, prepend=0.0 + ) + np.testing.assert_allclose( + probabilities[0], + [known_probabilities[name] for name in payload["ordered_brmas"][0]], + rtol=0, + atol=1e-12, + ) + # Scotland's sole supported BRMA precedes a safely padded zero column. + assert payload["ordered_brmas"][1] == ["LOTHIAN", "LOTHIAN"] + np.testing.assert_array_equal(probabilities[1], [1.0, 0.0]) + np.testing.assert_array_equal( + np.asarray(payload["cumulative_probabilities"])[:, -1], [1.0, 1.0] + ) + np.testing.assert_array_equal( + payload["uniforms"], + stable_identity_uniforms(lineage_keys, seed=0, salt="brma:clone_spread"), + ) + for entity in frame.entities: + pd.testing.assert_frame_equal(frame.table(entity), original[entity]) + np.testing.assert_array_equal( + frame.weights_for("household").values, original_weights + ) diff --git a/packages/microcosm-build/tests/engine_free/shared/test_country_spec.py b/packages/microcosm-build/tests/engine_free/shared/test_country_spec.py index 33235ec5a6..eaaf6c64c2 100644 --- a/packages/microcosm-build/tests/engine_free/shared/test_country_spec.py +++ b/packages/microcosm-build/tests/engine_free/shared/test_country_spec.py @@ -1459,6 +1459,7 @@ def test_declares_the_full_june_battery(self, manifest) -> None: "uk_input_mass_parity", "uk_qrf_tail_concentration", "uk_local_geography_ladder_post_calibration", + "uk_local_brma_enum_domain", "uk_local_area_support", "uk_local_target_fit", "uk_local_per_family_fit", @@ -1477,6 +1478,18 @@ def test_declares_the_full_june_battery(self, manifest) -> None: if gate.id not in diagnostic ) + def test_brma_domain_is_checked_on_the_spine_and_local_candidate( + self, manifest + ) -> None: + entries = {gate.id: gate for gate in manifest.gates} + assembled = entries["uk_brma_enum_domain"] + local = entries["uk_local_brma_enum_domain"] + assert assembled.phase == "assembled" + assert local.phase == "terminal" + assert assembled.gate == local.gate == "enum_domain" + assert assembled.parameters == local.parameters == {"columns": ("brma",)} + assert local.evidence_absent_blocks is True + def test_ledger_compile_parity_gates_pin_their_fixture_periods( self, manifest ) -> None: @@ -1552,6 +1565,7 @@ def test_strict_absent_evidence_entries_are_declared(self, manifest) -> None: # The #970 projection fence: a seam that cannot project cannot # certify the candidate's sub-exempt gainers. "uk_cgt_projection_entrants", + "uk_local_brma_enum_domain", ] assert all(g.not_applicable is None for g in manifest.gates) diff --git a/packages/microcosm-build/tests/engine_free/shared/test_logbook_isolation.py b/packages/microcosm-build/tests/engine_free/shared/test_logbook_isolation.py index cb38af7bfb..4820858bec 100644 --- a/packages/microcosm-build/tests/engine_free/shared/test_logbook_isolation.py +++ b/packages/microcosm-build/tests/engine_free/shared/test_logbook_isolation.py @@ -14,11 +14,14 @@ from __future__ import annotations import os +import sys +from types import ModuleType, SimpleNamespace import pytest from microcosm.build import logbook from microcosm.build.logbook_env import LOGBOOK_ENV_LEGACY_NAMES, logbook_env +from test_support.microcosm_build.logbook_isolation import logbook_urlopen_holders _ALL_LOGBOOK_ENV_NAMES = tuple( sorted( @@ -33,6 +36,28 @@ ) +def test_alias_discovery_reads_only_module_globals(monkeypatch) -> None: + """Find imported aliases without invoking native modules' lazy hooks.""" + real_urlopen = object() + lazy = ModuleType("test_logbook_lazy_module") + lookups = [] + + def lazy_getattr(name): + lookups.append(name) + raise AssertionError(f"Unexpected lazy module lookup: {name}") + + lazy.__getattr__ = lazy_getattr + alias = ModuleType("test_logbook_alias_module") + alias.urlopen = real_urlopen + foreign = SimpleNamespace(urlopen=real_urlopen) + monkeypatch.setitem(sys.modules, lazy.__name__, lazy) + monkeypatch.setitem(sys.modules, alias.__name__, alias) + monkeypatch.setitem(sys.modules, "test_logbook_foreign_holder", foreign) + + assert logbook_urlopen_holders(real_urlopen) == [alias] + assert lookups == [] + + @pytest.mark.parametrize("name", _ALL_LOGBOOK_ENV_NAMES) def test_no_logbook_variable_survives_into_a_test(name: str) -> None: """Both generations are cleared, not just the ledger-era spelling. diff --git a/packages/microcosm-build/tests/engine_free/uk/test_uk_atomic_identity_graph.py b/packages/microcosm-build/tests/engine_free/uk/test_uk_atomic_identity_graph.py index 568bd92451..64312dddb5 100644 --- a/packages/microcosm-build/tests/engine_free/uk/test_uk_atomic_identity_graph.py +++ b/packages/microcosm-build/tests/engine_free/uk/test_uk_atomic_identity_graph.py @@ -11,12 +11,15 @@ from __future__ import annotations import itertools +import json from dataclasses import replace +from types import SimpleNamespace import numpy as np import pandas as pd import pytest +from microcosm.build.stochastic_assignment import stable_identity_uniforms from microcosm.build.uk_runtime import uk_national_frame from microcosm.build.uk_runtime.atomic_area_support import ( IDENTITY_COLUMN, @@ -31,11 +34,15 @@ from microcosm.build.uk_runtime.graph_kernels import UKClaimKernel from microcosm.build.uk_runtime.graph_population import ( UK_GEOGRAPHY_IDENTITY_INPUTS, + UKGeographyIdentityKernel, append_uk_population_nodes, - register_uk_population_kernels, + household_lineage_keys, ) -from microcosm.build.uk_runtime.rowwise_dataset import ladder_clone_index_column -from microcosm.frame import MassChangeRecord, WeightKind +from microcosm.build.uk_runtime.rowwise_dataset import ( + expand_uk_geographic_pool, + ladder_clone_index_column, +) +from microcosm.frame import MassChangeRecord, WeightKind, engine_tables from microcosm.graph import ( ContentStore, Graph, @@ -53,8 +60,12 @@ from test_support.microcosm_build.uk_full_population_graph import ( SOURCE_VINTAGE, Source, + expected_clone_brmas, + pregeographic_key, + register_toy_population_kernels, source_node, ) +from test_support.paths import paths_for CLONE = ladder_clone_index_column("household") @@ -82,6 +93,7 @@ def lineage_frame(rows=LINEAGE, mutate=None): "household_id": ids, "household_weight": np.ones(len(rows)), "region": pd.array([r[6] for r in rows], dtype="string"), + "brma": pd.array([f"{r[6]}_A" for r in rows], dtype="string"), "source_household_id": np.asarray([r[1] for r in rows], dtype=np.int64), "household_support_channel": pd.array([r[2] for r in rows], dtype="string"), "household_support_clone_index": np.asarray( @@ -153,7 +165,7 @@ def graph_for(frame, *, k, definition, seed=7): registry = KernelRegistry() registry.register(LineageSource(frame)) registry.register(UKClaimKernel()) - register_uk_population_kernels(registry) + register_toy_population_kernels(registry) return graph, registry @@ -264,14 +276,102 @@ def test_identity_node_matches_the_receipt_based_projector(tmp_path): assert len(actual) == 3 * len(LINEAGE) -@pytest.mark.parametrize("k", [1, 2]) -def test_every_lineage_combination_encodes_its_declared_branch_path(tmp_path, k): - rows = [ +@pytest.mark.parametrize("k", [1, 3, 15]) +@pytest.mark.parametrize("fixture", ["lineage", "residential_arms"]) +def test_brma_spread_uses_spine_lineage_and_preserves_expanded_population( + tmp_path, k, fixture +): + rows = LINEAGE if fixture == "lineage" else RESIDENTIAL_ROWS + source = lineage_frame(rows) + manifest, graph = run_identity(source, tmp_path, k=k, endpoint="uk.full.brma") + compiled = compile_graph(graph) + table_node = graph.node("uk.full.brma_table") + rewrite_node = graph.node("uk.full.brma") + assert table_node.population == "uk.full.normalize" + assert {table_node.id, "uk.full.expand"} <= set( + compiled.predecessors[rewrite_node.id] + ) + assert [ + (output.entity, output.column, output.rewrite) + for output in rewrite_node.outputs + ] == [("household", "brma", True)] + actual = manifest.population("uk.full.expand") + household = source.table("household").copy() + household["household_weight"] = source.weights_for("household").values + expected = expand_uk_geographic_pool( + person=source.table("person"), + benunit=source.table("benunit"), + household=household, + n_clones=k, + source_year=2023, + time_period="2023", + household_weight_kind=source.weights_for("household").kind, + mass_log=source.mass_log, + ).frame + for entity in expected.entities: + columns = expected.table(entity).columns + if entity == "household": + columns = columns.drop("brma") + pd.testing.assert_frame_equal( + actual.table(entity)[columns], + expected.table(entity)[columns], + check_dtype=False, + ) + np.testing.assert_array_equal( + actual.weights_for("household").values, + expected.weights_for("household").values, + ) + assert actual.mass_log == expected.mass_log + expanded = actual.table("household") + np.testing.assert_array_equal(expanded["brma"], expected_clone_brmas(expanded, k)) + # Each SPI, CGT or residential branch has its own pre-geographic offset. + keys = [pregeographic_key(row) for row in source.table("household").itertuples()] + assert len(set(keys)) == len(rows) + table_key = manifest.nodes[table_node.id].opaque_artifacts["brma_table"] + table = json.loads(ContentStore(tmp_path / "store").load_bytes(table_key)) + assert table["household_ids"] == source.table("household")["household_id"].tolist() + np.testing.assert_array_equal( + table["uniforms"], + stable_identity_uniforms(keys, seed=0, salt="brma:clone_spread"), + ) + assert all( + row.brma.startswith(f"{row.region}_") + for row in expanded.itertuples(index=False) + ) + + +@pytest.mark.parametrize("k", [1, 3, 15]) +def test_brma_lineage_draws_do_not_depend_on_population_row_order(tmp_path, k): + by_lineage = [] + for name, rows in (("forward", LINEAGE), ("reverse", list(reversed(LINEAGE)))): + # Frame requires sorted structural IDs. Remap them for each row order, + # preserving source lineage: draws must also ignore concrete local IDs. + remapped = [(index, *row[1:]) for index, row in enumerate(rows, start=1)] + manifest, _ = run_identity( + lineage_frame(remapped), tmp_path / name, k=k, endpoint="uk.full.brma" + ) + household = manifest.population("uk.full.expand").table("household") + by_lineage.append( + { + (pregeographic_key(row), getattr(row, CLONE)): row.brma + for row in household.itertuples(index=False) + } + ) + assert by_lineage[0] == by_lineage[1] + + +def lineage_combinations(): + return [ (100 + index, 1, channel, int(channel == "spi"), cgt, copy_index, "LONDON") for index, (channel, cgt, copy_index) in enumerate( itertools.product(("frs", "spi"), (False, True), (0, 2)) ) ] + + +@pytest.mark.parametrize("k", [1, 2]) +def test_every_lineage_combination_encodes_its_declared_branch_path(tmp_path, k): + rows = lineage_combinations() manifest, _ = run_identity(lineage_frame(rows), tmp_path, k=k) household = manifest.population("uk.full.expand").table("household") keys = set() @@ -423,3 +523,132 @@ def clear_flag(household): NodeRejectedError, match="residential-arm flag and arm index disagree" ): run_identity(frame, tmp_path, k=1) + + +def main_identity_reference(): + # Generated by executing main's unchanged identity kernel on its own + # fixtures at a91c7536f7e318bdfdd8ee6214a7411257f0033a. The golden records + # both source-file hashes; it does not reconstruct the expected key loop. + path = ( + paths_for("microcosm-build").tests + / "golden" + / "uk_geography_identity_main.json" + ) + return json.loads(path.read_text()) + + +@pytest.mark.parametrize("k", [1, 2]) +@pytest.mark.parametrize("fixture", ["lineage", "residential_arms", "combinations"]) +def test_geographic_identity_keys_are_byte_identical_to_main(fixture, k): + rows = { + "lineage": LINEAGE, + "residential_arms": RESIDENTIAL_ROWS, + "combinations": lineage_combinations(), + }[fixture] + frame = lineage_frame(rows) + pool = expand_uk_geographic_pool( + **engine_tables(frame, weighted_entities=("household",)), + n_clones=k, + time_period="2023", + source_year=2023, + household_weight_kind=WeightKind.IMPORTANCE, + mass_log=frame.mass_log, + ) + household = pool.frame.table("household") + reference = main_identity_reference() + result = UKGeographyIdentityKernel().run( + SimpleNamespace( + tables={"household": household}, + params={ + "source": reference["source"], + "source_vintage": reference["source_vintage"], + }, + ) + ) + expected = reference["cases"][fixture][str(k)] + actual = result.columns[("household", IDENTITY_COLUMN)] + assert household["household_id"].tolist() == expected["household_ids"] + assert [key.encode("utf-8") for key in actual] == [ + key.encode("utf-8") for key in expected["keys"] + ] + + +def test_pregeographic_residential_arms_have_distinct_main_lineage_keys(): + household = lineage_frame(RESIDENTIAL_ROWS).table("household") + assert CLONE not in household + keys = household_lineage_keys( + household, + source="frs", + source_vintage=SOURCE_VINTAGE, + geographic=False, + ) + expected = main_identity_reference()["cases"]["residential_arms"]["1"]["keys"] + assert [key.encode("utf-8") for key in keys] == [ + key.encode("utf-8") for key in expected + ] + assert len(set(keys)) == len(RESIDENTIAL_ROWS) + arms = household["cgt_residential_clone_index"].to_numpy() != 0 + assert sum(arms) == 3 + assert all( + "cgt_residential_split" in key + for key, arm in zip(keys, arms, strict=True) + if arm + ) + + +@pytest.mark.parametrize( + ("defect", "message"), + [ + ("flag_disagrees", "residential-arm flag and arm index disagree"), + ("negative_index", "negative clone indices"), + ("index_dtype", "int64 residential arm index"), + ("flag_dtype", "boolean household_is_cgt_residential_clone"), + ("null_index", "null lineage"), + ("missing_index", "lacks lineage column"), + ("missing_flag", "lacks lineage column"), + ], +) +def test_pregeographic_lineage_validates_residential_arms(defect, message): + household = lineage_frame(RESIDENTIAL_ROWS).table("household") + arm = household["household_id"] == 201 + if defect == "flag_disagrees": + household.loc[arm, "household_is_cgt_residential_clone"] = False + elif defect == "negative_index": + household.loc[arm, "cgt_residential_clone_index"] = -1 + elif defect == "index_dtype": + household["cgt_residential_clone_index"] = household[ + "cgt_residential_clone_index" + ].astype("float64") + elif defect == "flag_dtype": + household["household_is_cgt_residential_clone"] = household[ + "household_is_cgt_residential_clone" + ].astype("int64") + elif defect == "null_index": + household["cgt_residential_clone_index"] = household[ + "cgt_residential_clone_index" + ].astype("Int64") + household.loc[arm, "cgt_residential_clone_index"] = pd.NA + elif defect == "missing_index": + household.drop(columns="cgt_residential_clone_index", inplace=True) + else: + household.drop(columns="household_is_cgt_residential_clone", inplace=True) + with pytest.raises(ValueError, match=message): + household_lineage_keys( + household, + source="frs", + source_vintage=SOURCE_VINTAGE, + geographic=False, + ) + + +def test_brma_table_reads_residential_lineage_before_geographic_expansion(tmp_path): + payloads, _ = write_toy_supports(tmp_path / "supports") + definition = uk_atomic_assignment_definition(payloads, seed=7) + graph, _ = graph_for(lineage_frame(RESIDENTIAL_ROWS), k=2, definition=definition) + node = graph.node("uk.full.brma_table") + household_inputs = next( + item.columns for item in node.inputs if item.entity == "household" + ) + assert set(UK_GEOGRAPHY_IDENTITY_INPUTS[:-1]) <= set(household_inputs) + assert CLONE not in household_inputs + assert node.population == "uk.full.normalize" diff --git a/packages/microcosm-build/tests/engine_free/uk/test_uk_battery_bindings.py b/packages/microcosm-build/tests/engine_free/uk/test_uk_battery_bindings.py index 499711977f..d997f4e28c 100644 --- a/packages/microcosm-build/tests/engine_free/uk/test_uk_battery_bindings.py +++ b/packages/microcosm-build/tests/engine_free/uk/test_uk_battery_bindings.py @@ -974,7 +974,10 @@ def test_fully_armed_battery_evaluates_gate_for_gate(self) -> None: # Plus the #970 projection fence, armed here with the pinned growth # path over distinct sub-exempt gains that never cross the exempt # amount. - assert len(passed) == 20 + # Plus uk_local_brma_enum_domain, the terminal recheck of the local + # clone-spread BRMA rewrite (its enum_domain binding is stubbed here). + assert len(passed) == 21 + assert "uk_local_brma_enum_domain" in passed qrf = by_id["uk_qrf_tail_concentration"] assert qrf.status is GateStatus.FAILED assert "declared QRF output is absent" in qrf.result.failures[0] diff --git a/packages/microcosm-build/tests/engine_free/uk/test_uk_brma_spread.py b/packages/microcosm-build/tests/engine_free/uk/test_uk_brma_spread.py new file mode 100644 index 0000000000..d3bedefddc --- /dev/null +++ b/packages/microcosm-build/tests/engine_free/uk/test_uk_brma_spread.py @@ -0,0 +1,414 @@ +"""BRMA clone-spreading invariants, without importing a country engine.""" + +import copy +import json +from types import SimpleNamespace + +import numpy as np +import pandas as pd +import pytest + +from microcosm.build.stochastic_assignment import stable_identity_uniforms +from microcosm.build.uk_runtime.atomic_household_identity import household_draw_key +from microcosm.build.uk_runtime.brma_spread import ( + assign_clone_brmas, + brmas_at_quantiles, + household_brma_probabilities, + lha_rate_keys, + prepare_brma_spread, + spaced_quantiles, +) +from microcosm.build.uk_runtime.national_frame import uk_national_frame +from microcosm.build.uk_runtime.rowwise_dataset import ( + expand_uk_geographic_pool, + ladder_clone_index_column, +) +from microcosm.frame import engine_tables + +_DISTRIBUTIONS = ( + (1.0,), + (0.1, 0.6, 0.3, 0.0), + (0.0, 0.2, 0.0, 0.8, 0.0), + (0.7, 0.0, 0.1, 0.0, 0.2), + (0.0, 1e-13, 1.0 - 1e-13, 0.0), + (0.03, 0.08, 0.89, 0.0), +) + + +def _frame(): + return uk_national_frame( + person=pd.DataFrame( + { + "person_id": [1, 2, 3, 4], + "person_benunit_id": [10, 10, 20, 30], + "person_household_id": [5, 5, 5, 9], + "age": [40, 39, 60, 50], + } + ), + benunit=pd.DataFrame({"benunit_id": [10, 20, 30], "would_claim_uc": True}), + household=pd.DataFrame( + { + "household_id": [5, 9], + "region": ["LONDON", "SCOTLAND"], + "brma": ["LON_A", "SCOT_A"], + "rent": [100.0, 200.0], + } + ), + household_weights=[7.0, 13.0], + time_period=2025, + ) + + +def _resource(): + return { + "cells": { + "LONDON": { + "A": {"LON_C": 3, "LON_A": 1, "LON_ZERO": 0}, + "E": {"LON_B": 1, "LON_A": 0}, + }, + "SCOTLAND": {"B": {"SCOT_B": 2, "SCOT_A": 2}}, + } + } + + +class _Engine: + country = "uk" + + def __init__(self, *, uncapped=True): + self.uncapped = uncapped + self.calls = [] + + def variables(self): + # Formula-owned rates are deliberately absent from this adapter API. + return ["brma"] + + def variable_metadata(self, name): + if name == "uncapped_BRMA_LHA_rate" and not self.uncapped: + raise ValueError(f"Unknown PolicyEngine-UK variable {name!r}.") + return SimpleNamespace(entity="benunit") + + def materialize(self, frame, variables, period): + names = frame.table("household")["brma"].to_list() + self.calls.append((list(variables), names, period)) + if variables == ["LHA_category"]: + return {"LHA_category": np.asarray(["A", "E", "B"])} + rates = {"LON_A": 60, "LON_B": 30, "LON_C": 10, "SCOT_A": 100, "SCOT_B": 50} + return { + variables[0]: np.asarray( + [rates[names[0]], rates[names[0]] + 20, rates[names[1]] + 7], + dtype=float, + ) + } + + +def _lineage_keys(): + return [ + household_draw_key( + source="frs", + source_vintage="2023_24", + source_household_id=5, + clone_path=(("spi_support_channel", spi), ("cgt_support_split", 0)), + ) + for spi in (0, 1) + ] + + +@pytest.mark.parametrize("probabilities", _DISTRIBUTIONS) +@pytest.mark.parametrize("k", [1, 2, 4, 15]) +@pytest.mark.parametrize("ordered", [False, True]) +def test_each_clone_has_exact_household_marginals_and_positive_support( + probabilities, k, ordered +): + grid = 6_000 + p = np.asarray(probabilities)[None, :] + keys = np.arange(p.shape[1], dtype=float)[None, ::-1] if ordered else None + uniforms = (np.arange(grid) + 0.5) / grid + quantiles = spaced_quantiles(uniforms, k) + p_grid = np.repeat(p, grid, axis=0) + keys_grid = None if keys is None else np.repeat(keys, grid, axis=0) + for clone in range(k): + columns = brmas_at_quantiles(p_grid, keys_grid, quantiles[:, clone]) + frequencies = np.bincount(columns, minlength=p.shape[1]) / grid + np.testing.assert_allclose(frequencies, p[0], rtol=0, atol=2 / grid) + assert (p[0, columns] > 0).all() + assert (frequencies[p[0] == 0] == 0).all() + + +@pytest.mark.parametrize("probabilities", _DISTRIBUTIONS) +@pytest.mark.parametrize("quantile", [0.0, 0.37, np.nextafter(1.0, 0.0)]) +@pytest.mark.parametrize("ordered", [False, True]) +def test_boundary_quantiles_never_choose_zero_probability( + probabilities, quantile, ordered +): + p = np.asarray(probabilities)[None, :] + keys = np.arange(p.shape[1], dtype=float)[None, ::-1] if ordered else None + columns = brmas_at_quantiles(p, keys, np.asarray([quantile])) + assert p[0, columns[0]] > 0 + + +@pytest.mark.parametrize("probabilities", _DISTRIBUTIONS) +@pytest.mark.parametrize("k", [1, 2, 4, 15]) +@pytest.mark.parametrize("uniform", [0.0, 0.113, 0.491, 0.997]) +def test_spaced_clones_bound_monotone_outcome_error(probabilities, k, uniform): + p = np.asarray(probabilities)[None, :] + keys = np.arange(p.shape[1], dtype=float)[None, ::-1] + quantiles = spaced_quantiles(np.asarray([uniform]), k) + gaps = np.diff(np.concatenate([np.sort(quantiles[0]), [quantiles.min() + 1]])) + np.testing.assert_allclose(gaps, 1 / k, rtol=0, atol=1e-15) + assert sorted(np.floor(quantiles[0] * k).astype(int)) == list(range(k)) + columns = brmas_at_quantiles( + np.repeat(p, k, axis=0), np.repeat(keys, k, axis=0), quantiles.ravel() + ) + # A non-linear monotone transform also obeys the range/K bound. + outcome = keys[0] ** 3 + expected = p[0] @ outcome + support = outcome[p[0] > 0] + assert abs(outcome[columns].mean() - expected) <= np.ptp(support) / k + 1e-12 + + +def test_walk_uses_increasing_rate_and_breaks_ties_by_column(): + p = np.asarray([[0.25, 0.25, 0.25, 0.0, 0.25]]) + keys = np.asarray([[3, 1, 1, -100, 2]]) + quantiles = np.asarray([0, 0.25, 0.5, 0.75]) + columns = brmas_at_quantiles( + np.repeat(p, 4, axis=0), np.repeat(keys, 4, axis=0), quantiles + ) + np.testing.assert_array_equal(columns, [1, 2, 4, 0]) + + +def test_household_distribution_averages_units_not_people_or_cell_counts(): + frame = _frame() + resource = _resource() + original = copy.deepcopy(resource) + brmas, probabilities = household_brma_probabilities( + frame, lha_category=["A", "E", "B"], count_resource=resource + ) + assert brmas == { + "LONDON": ["LON_A", "LON_B", "LON_C"], + "SCOTLAND": ["SCOT_A", "SCOT_B"], + } + np.testing.assert_allclose(probabilities, [[0.125, 0.5, 0.375], [0.5, 0.5, 0]]) + assert resource == original + + +@pytest.mark.parametrize("uncapped", [True, False]) +def test_lha_rate_order_uses_household_mean_and_uncapped_fallback(uncapped): + frame = _frame() + original = frame.table("household").copy() + engine = _Engine(uncapped=uncapped) + brmas, _ = household_brma_probabilities( + frame, lha_category=["A", "E", "B"], count_resource=_resource() + ) + keys = lha_rate_keys(frame, engine=engine, brmas=brmas) + np.testing.assert_array_equal(keys, [[70, 40, 20], [107, 57, np.inf]]) + variable = "uncapped_BRMA_LHA_rate" if uncapped else "BRMA_LHA_rate" + assert [call[0] for call in engine.calls] == [[variable]] * 3 + assert [call[1] for call in engine.calls] == [ + ["LON_A", "SCOT_A"], + ["LON_B", "SCOT_B"], + ["LON_C", "SCOT_B"], + ] + pd.testing.assert_frame_equal(frame.table("household"), original) + + +def test_preparation_uses_pregeographic_lineage_keys_and_serializes(): + frame = _frame() + lineage = _lineage_keys() + payload = prepare_brma_spread( + frame, engine=_Engine(), lineage_keys=lineage, count_resource=_resource() + ) + assert payload["household_ids"] == [5, 9] + assert payload["ordered_brmas"] == [ + ["LON_C", "LON_B", "LON_A"], + ["SCOT_B", "SCOT_A", "SCOT_B"], + ] + np.testing.assert_allclose( + payload["cumulative_probabilities"], [[0.375, 0.875, 1], [0.5, 1, 1]] + ) + np.testing.assert_array_equal( + payload["uniforms"], + stable_identity_uniforms(lineage, seed=0, salt="brma:clone_spread"), + ) + assert payload["uniforms"][0] != payload["uniforms"][1] + assert json.loads(json.dumps(payload, allow_nan=False)) == payload + + +@pytest.mark.parametrize("k", [1, 4, 15]) +def test_clone_assignment_has_regional_support_conserves_and_ignores_row_order(k): + frame = _frame() + payload = prepare_brma_spread( + frame, + engine=_Engine(), + lineage_keys=_lineage_keys(), + count_resource=_resource(), + ) + tables = engine_tables(frame, weighted_entities=("household",)) + pool = expand_uk_geographic_pool( + **tables, n_clones=k, time_period=2025, id_multiplier=1_000 + ) + household = pool.frame.table("household") + original = household.copy() + weights = pool.frame.weights_for("household").values.copy() + kwargs = { + "n_clones": k, + "id_multiplier": pool.id_multiplier, + "clone_index_column": ladder_clone_index_column("household"), + } + assigned = assign_clone_brmas(household, payload, **kwargs) + for region, supported in { + "LONDON": {"LON_A", "LON_B", "LON_C"}, + "SCOTLAND": {"SCOT_A", "SCOT_B"}, + }.items(): + assert set(assigned[household.region == region]) <= supported + assert len(assigned) == len(household) == 2 * k + pd.testing.assert_frame_equal(household, original) + np.testing.assert_array_equal(pool.frame.weights_for("household").values, weights) + rewritten = household.assign(brma=assigned) + pd.testing.assert_frame_equal( + rewritten.drop(columns="brma"), original.drop(columns="brma") + ) + reordered = household.iloc[::-1] + np.testing.assert_array_equal( + assign_clone_brmas(reordered, payload, **kwargs), assigned[::-1] + ) + # Reorder the prepared rows as well: parent lookup is by exact ID. + reverse_payload = {key: value[::-1] for key, value in payload.items()} + np.testing.assert_array_equal( + assign_clone_brmas(household, reverse_payload, **kwargs), assigned + ) + brmas, probabilities = household_brma_probabilities( + frame, lha_category=["A", "E", "B"], count_resource=_resource() + ) + keys = lha_rate_keys(frame, engine=_Engine(), brmas=brmas) + quantiles = spaced_quantiles(np.asarray(payload["uniforms"]), k) + for clone in range(k): + columns = brmas_at_quantiles(probabilities, keys, quantiles[:, clone]) + expected = [ + brmas[region][column] + for region, column in zip( + frame.table("household").region, columns, strict=True + ) + ] + np.testing.assert_array_equal(assigned[2 * clone : 2 * (clone + 1)], expected) + + +@pytest.mark.parametrize( + ("names", "cumulative", "uniform", "k", "expected"), + [ + # A quantile exactly on a cumulative boundary takes the next BRMA. + (["A", "B", "B"], [0.5, 1.0, 1.0], 0.5, 2, ["B", "A"]), + (["A", "B", "B"], [0.5, 1.0, 1.0], 0.0, 4, ["A", "A", "B", "B"]), + # A zero-probability column is never chosen, even at quantile zero. + (["Z", "A", "B"], [0.0, 0.5, 1.0], 0.0, 2, ["A", "B"]), + (["A", "Z", "B"], [0.5, 0.5, 1.0], 0.5, 1, ["B"]), + # Just below one stays inside the support. + (["A", "B", "B"], [0.5, 1.0, 1.0], float(np.nextafter(1.0, 0.0)), 1, ["B"]), + ], +) +def test_clone_assignment_cdf_boundaries(names, cumulative, uniform, k, expected): + payload = { + "household_ids": [7], + "ordered_brmas": [names], + "cumulative_probabilities": [cumulative], + "uniforms": [uniform], + } + household = pd.DataFrame( + {"household_id": [7 + clone * 100 for clone in range(k)], "clone": range(k)} + ) + assigned = assign_clone_brmas( + household, payload, n_clones=k, id_multiplier=100, clone_index_column="clone" + ) + assert assigned.tolist() == expected + + +def test_lineage_draws_and_quantile_assignments_are_independent_of_row_order(): + lineage = _lineage_keys() + p = np.asarray([[0.125, 0.5, 0.375], [0.5, 0.5, 0]]) + keys = np.asarray([[70, 40, 20], [107, 57, np.inf]]) + u = stable_identity_uniforms(lineage, seed=0, salt="brma:clone_spread") + reversed_u = stable_identity_uniforms( + lineage[::-1], seed=0, salt="brma:clone_spread" + ) + np.testing.assert_array_equal(reversed_u, u[::-1]) + np.testing.assert_array_equal( + brmas_at_quantiles(p[::-1], keys[::-1], reversed_u), + brmas_at_quantiles(p, keys, u)[::-1], + ) + + +@pytest.mark.parametrize("quantile", [-1, 1, np.nan, np.inf]) +def test_quantiles_reject_invalid_values(quantile): + with pytest.raises(ValueError, match="Quantiles"): + brmas_at_quantiles(np.asarray([[1.0]]), None, np.asarray([quantile])) + + +@pytest.mark.parametrize("k", [0, -1, 1.5, True]) +def test_spaced_quantiles_require_a_positive_integer(k): + with pytest.raises(ValueError, match="positive integer"): + spaced_quantiles(np.asarray([0.5]), k) + + +@pytest.mark.parametrize( + "probabilities", [[0, 0], [-0.1, 1.1], [0.1, 0.2], [np.nan, 1]] +) +def test_inverse_cdf_rejects_invalid_probabilities(probabilities): + with pytest.raises(ValueError, match="Probabilities"): + brmas_at_quantiles(np.asarray([probabilities]), None, np.asarray([0.5])) + + +@pytest.mark.parametrize("bad_count", [-1, np.nan, np.inf]) +def test_household_probabilities_reject_invalid_counts(bad_count): + resource = _resource() + resource["cells"]["LONDON"]["A"]["LON_A"] = bad_count + with pytest.raises(ValueError, match="finite and non-negative"): + household_brma_probabilities( + _frame(), lha_category=["A", "E", "B"], count_resource=resource + ) + + +def test_household_probabilities_refuse_missing_and_zero_only_cells(): + with pytest.raises(KeyError, match="missing BRMA"): + household_brma_probabilities( + _frame(), lha_category=["A", "C", "B"], count_resource=_resource() + ) + resource = _resource() + resource["cells"]["LONDON"]["E"] = {"LON_B": 0} + with pytest.raises(ValueError, match="no positive counts"): + household_brma_probabilities( + _frame(), lha_category=["A", "E", "B"], count_resource=resource + ) + + +def test_preparation_refuses_duplicate_lineage_keys(): + with pytest.raises(ValueError, match="distinct lineage key"): + prepare_brma_spread( + _frame(), + engine=_Engine(), + lineage_keys=["same", "same"], + count_resource=_resource(), + ) + + +@pytest.mark.parametrize("error", ["unknown_parent", "invalid_clone", "invalid_cdf"]) +def test_clone_assignment_refuses_malformed_inputs(error): + payload = prepare_brma_spread( + _frame(), + engine=_Engine(), + lineage_keys=_lineage_keys(), + count_resource=_resource(), + ) + household = pd.DataFrame({"household_id": [5, 9], "household_clone_index": [0, 0]}) + if error == "unknown_parent": + household.loc[0, "household_id"] = 7 + elif error == "invalid_clone": + household.loc[0, "household_clone_index"] = 2 + else: + payload["cumulative_probabilities"][0][-1] = 0.99 + with pytest.raises(ValueError): + assign_clone_brmas( + household, + payload, + n_clones=2, + id_multiplier=1_000, + clone_index_column="household_clone_index", + ) diff --git a/packages/microcosm-build/tests/engine_free/uk/test_uk_full_gates.py b/packages/microcosm-build/tests/engine_free/uk/test_uk_full_gates.py index 5a4ec666c9..5ca31fbb43 100644 --- a/packages/microcosm-build/tests/engine_free/uk/test_uk_full_gates.py +++ b/packages/microcosm-build/tests/engine_free/uk/test_uk_full_gates.py @@ -50,6 +50,7 @@ def test_country_filter_retains_integrity_and_registry_checks_without_local_fit_ } assert { "uk_local_geography_ladder_post_calibration", + "uk_local_brma_enum_domain", "uk_weight_ess", "uk_release_family_build_stages", "uk_ledger_compile_parity_local_incumbent_2025", @@ -284,6 +285,7 @@ def test_existing_local_statistical_failure_can_export_but_blocks_full_release() "gate_id", [ "uk_local_geography_ladder_post_calibration", + "uk_local_brma_enum_domain", "uk_nonnegative_columns", "uk_input_mass_parity", "uk_target_fit", @@ -299,6 +301,18 @@ def test_geography_and_migrated_national_gates_keep_artifact_enforcement(gate_id assert result["structural_failures"] == [gate_id] +@pytest.mark.parametrize("fraction", [0.1, 1.0]) +@pytest.mark.parametrize("release_candidate", [False, True]) +def test_local_brma_missing_domain_evidence_stops_export(fraction, release_candidate): + result = runtime.classify_full_gate_outcomes( + _phase_failure("uk_local_brma_enum_domain", absent=True), + sample_fraction=fraction, + release_candidate=release_candidate, + ) + assert result["artifact_permitted"] is False + assert result["structural_failures"] == ["uk_local_brma_enum_domain"] + + def test_missing_evidence_uses_declared_development_policy(): ordinary = _phase_failure("uk_input_mass_parity", absent=True) assert ( diff --git a/packages/microcosm-build/tests/engine_free/uk/test_uk_full_population_graph.py b/packages/microcosm-build/tests/engine_free/uk/test_uk_full_population_graph.py index 6abbf82f5a..18d2347d48 100644 --- a/packages/microcosm-build/tests/engine_free/uk/test_uk_full_population_graph.py +++ b/packages/microcosm-build/tests/engine_free/uk/test_uk_full_population_graph.py @@ -27,6 +27,7 @@ ) from test_support.microcosm_build.uk_full_population_graph import ( SOURCE_VINTAGE, + expected_clone_brmas, graph_and_registry, source_frame, ) @@ -51,6 +52,7 @@ def test_legacy_population_graph_preserves_rows_weights_geography_and_replays( graph, registry = graph_and_registry(k, geography_assignment="legacy") store = ContentStore(tmp_path / "store") compiled = compile_graph(graph) + _assert_brma_nodes(graph, compiled) assert "uk.full.locations" in compiled.predecessors["uk.full.geography_mapping"] assert "uk.full.identity" not in compiled.order first = run_graph( @@ -61,11 +63,18 @@ def test_legacy_population_graph_preserves_rows_weights_geography_and_replays( ) actual = first.population("uk.full.expand") for entity in expected.entities: + columns = expected.table(entity).columns + if entity == "household": + columns = columns.drop("brma") pd.testing.assert_frame_equal( - actual.table(entity)[expected.table(entity).columns], - expected.table(entity), + actual.table(entity)[columns], + expected.table(entity)[columns], check_dtype=False, ) + np.testing.assert_array_equal( + actual.table("household")["brma"], + expected_clone_brmas(actual.table("household"), k), + ) np.testing.assert_array_equal( actual.weights_for("household").values, expected.weights_for("household").values ) @@ -80,6 +89,23 @@ def test_legacy_population_graph_preserves_rows_weights_geography_and_replays( resume="require", ) assert replay.population("uk.full.expand").mass_log == actual.mass_log + assert all(receipt.hit for receipt in replay.nodes.values()) + pd.testing.assert_frame_equal( + replay.population("uk.full.expand").table("household"), + actual.table("household"), + ) + + +def _assert_brma_nodes(graph, compiled): + table = graph.node("uk.full.brma_table") + rewrite = graph.node("uk.full.brma") + assert table.population == "uk.full.normalize" + assert "uk.full.normalize" in compiled.predecessors[table.id] + assert rewrite.population == "uk.full.expand" + assert {table.id, "uk.full.expand"} <= set(compiled.predecessors[rewrite.id]) + assert [(o.entity, o.column, o.rewrite) for o in rewrite.outputs] == [ + ("household", "brma", True) + ] def test_k_changes_expansion_but_never_sampling(): @@ -87,16 +113,16 @@ def test_k_changes_expansion_but_never_sampling(): second, _ = graph_and_registry(3) assert first.node("uk.full.sample") == second.node("uk.full.sample") assert first.node("uk.full.normalize") == second.node("uk.full.normalize") + assert first.node("uk.full.brma_table") == second.node("uk.full.brma_table") assert first.node("uk.full.expand").params["n_clones"] == 1 assert second.node("uk.full.expand").params["n_clones"] == 3 -def _atomic_oracle(k, payloads, definition): - """The same pool, keyed and assigned in process through the shared operators.""" +def _expanded_oracle(k): frame = source_frame() household = frame.table("household").copy() household["household_weight"] = frame.weights_for("household").values - pool = expand_uk_geographic_pool( + return expand_uk_geographic_pool( person=frame.table("person"), benunit=frame.table("benunit"), household=household, @@ -105,7 +131,12 @@ def _atomic_oracle(k, payloads, definition): time_period="2023", household_weight_kind=frame.weights_for("household").kind, mass_log=frame.mass_log, - ).frame.table("household") + ).frame + + +def _atomic_oracle(k, payloads, definition): + """The same pool, keyed and assigned in process through the shared operators.""" + pool = _expanded_oracle(k).table("household") clone = ladder_clone_index_column("household") keys = [ household_draw_key( @@ -143,6 +174,7 @@ def test_atomic_population_graph_keys_assigns_derives_gates_and_replays( k, geography_assignment="atomic", definition=definition ) compiled = compile_graph(graph) + _assert_brma_nodes(graph, compiled) for absent in ("uk.full.locations", "uk.full.geography_mapping"): assert absent not in compiled.order gate = compiled.predecessors["uk.full.geography_gate"] @@ -164,9 +196,26 @@ def test_atomic_population_graph_keys_assigns_derives_gates_and_replays( } store = ContentStore(tmp_path / "store") first = run_graph(compiled, sources=sources, store=store, kernels=registry) - household = first.population("uk.full.expand").table("household") + frame = first.population("uk.full.expand") + household = frame.table("household") expected = _atomic_oracle(k, payloads, definition) assert len(household) == 4 * k + expanded = _expanded_oracle(k) + for entity in expanded.entities: + columns = expanded.table(entity).columns + if entity == "household": + columns = columns.drop("brma") + pd.testing.assert_frame_equal( + frame.table(entity)[columns], + expanded.table(entity)[columns], + check_dtype=False, + ) + np.testing.assert_array_equal( + frame.weights_for("household").values, + expanded.weights_for("household").values, + ) + assert frame.mass_log == expanded.mass_log + np.testing.assert_array_equal(household["brma"], expected_clone_brmas(household, k)) actual = household.set_index("household_id") for column in ( IDENTITY_COLUMN, @@ -201,6 +250,9 @@ def test_atomic_population_graph_keys_assigns_derives_gates_and_replays( compiled, sources=sources, store=store, kernels=fresh, resume="require" ) assert all(receipt.hit for receipt in replay.nodes.values()) + pd.testing.assert_frame_equal( + replay.population("uk.full.expand").table("household"), household + ) def test_atomic_keys_and_draws_are_stable_when_k_grows(toy_ladder, tmp_path): diff --git a/packages/microcosm-build/tests/engine_free/uk/test_uk_graph_terminal.py b/packages/microcosm-build/tests/engine_free/uk/test_uk_graph_terminal.py index c576ffe05c..f5a3db81ff 100644 --- a/packages/microcosm-build/tests/engine_free/uk/test_uk_graph_terminal.py +++ b/packages/microcosm-build/tests/engine_free/uk/test_uk_graph_terminal.py @@ -280,7 +280,10 @@ def test_full_gate_nodes_precede_dense_and_bind_final_problem_axis(households): from microcosm.build.uk_runtime.graph import uk_spine_endpoint, uk_spine_graph from microcosm.build.uk_runtime.graph_build import UKFullBuildConfig, uk_full_graph from microcosm.build.uk_runtime.graph_calibration import UKGraphCalibrationConfig - from microcosm.build.uk_runtime.graph_terminal import append_uk_full_gate_nodes + from microcosm.build.uk_runtime.graph_terminal import ( + add_uk_export_preparation, + append_uk_full_gate_nodes, + ) from microcosm.graph import compile_graph from test_support.microcosm_build.uk_atomic_support_fixtures import ( toy_support_payloads, @@ -305,7 +308,48 @@ def test_full_gate_nodes_precede_dense_and_bind_final_problem_axis(households): engine_identity="fixture-engine", review_date="2026-09-10", ) + graph = add_uk_export_preparation( + graph, population=full.calibration.population, bindings={} + ) compiled = compile_graph(graph) + assert compiled.order.index("uk.full.expand") < compiled.order.index("uk.full.brma") + + def reads_local_population(node): + population = node.population or node.base + while population is not None: + if population == "uk.full.expand": + return True + population = graph.node(population).base + return False + + def ancestors(node_id): + pending = list(compiled.predecessors[node_id]) + found = set() + while pending: + predecessor = pending.pop() + if predecessor not in found: + found.add(predecessor) + pending.extend(compiled.predecessors[predecessor]) + return found + + # Scan the current graph so new scoring, selection or export consumers + # cannot accidentally read the spine BRMA before the local rewrite. + brma_consumers = { + node.id + for node in graph.nodes + if reads_local_population(node) + and any( + item.entity == "household" and "brma" in item.columns + for item in node.inputs + ) + } + assert { + "uk.full.measures", + "uk.full.gates.calibrated", + "uk.full.export.prepare", + } <= (brma_consumers) + for consumer in brma_consumers: + assert "uk.full.brma" in ancestors(consumer), consumer assert ( "uk.full.gates.preflight" in compiled.predecessors[full.calibration.dense_producer] @@ -324,6 +368,34 @@ def test_full_gate_nodes_precede_dense_and_bind_final_problem_axis(households): ) +@pytest.mark.parametrize("candidate_valid", [False, True]) +def test_local_brma_terminal_gate_reads_candidate_instead_of_spine(candidate_valid): + from dataclasses import replace + + from microcosm.build.country_spec import load_country_spec + from microcosm.build.gate_battery import EvidenceContext, GateStatus, evaluate_phase + from microcosm.build.uk_runtime.battery_bindings import UK_GATE_REGISTRY + + candidate, spine = _frame(), _frame() + valid, invalid = "CENTRAL_LONDON", "INVALID_BRMA" + candidate.table("household")["brma"] = valid if candidate_valid else invalid + spine.table("household")["brma"] = invalid if candidate_valid else valid + manifest = load_country_spec("uk").gates + entry = next(g for g in manifest.gates if g.id == "uk_local_brma_enum_domain") + report = evaluate_phase( + replace(manifest, gates=(entry,)), + "terminal", + EvidenceContext( + frame=candidate, + artifacts={"brma_enum_domain": [valid], "spine_frame": spine}, + ), + registry=UK_GATE_REGISTRY, + ) + assert report.outcomes[0].status is ( + GateStatus.PASSED if candidate_valid else GateStatus.FAILED + ) + + def test_full_preflight_persists_real_source_failures_without_matrix_diagnostics( monkeypatch, ): diff --git a/packages/microcosm-build/tests/engine_free/uk/test_uk_local_gate_battery.py b/packages/microcosm-build/tests/engine_free/uk/test_uk_local_gate_battery.py index 0b18a5432e..10cc07eef5 100644 --- a/packages/microcosm-build/tests/engine_free/uk/test_uk_local_gate_battery.py +++ b/packages/microcosm-build/tests/engine_free/uk/test_uk_local_gate_battery.py @@ -8,9 +8,10 @@ ) -def test_local_scope_is_the_declared_six_gate_terminal_battery() -> None: +def test_local_scope_is_the_declared_seven_gate_terminal_battery() -> None: assert UK_LOCAL_GATE_SCOPE == ( "uk_local_geography_ladder_post_calibration", + "uk_local_brma_enum_domain", "uk_local_area_support", "uk_local_target_fit", "uk_local_per_family_fit", diff --git a/packages/microcosm-build/tests/engine_free/uk/test_uk_orrery_contract.py b/packages/microcosm-build/tests/engine_free/uk/test_uk_orrery_contract.py index b645213162..282fd4bea2 100644 --- a/packages/microcosm-build/tests/engine_free/uk/test_uk_orrery_contract.py +++ b/packages/microcosm-build/tests/engine_free/uk/test_uk_orrery_contract.py @@ -116,9 +116,13 @@ def test_complete_dense_and_size_declarations_fit_existing_export_limits(househo assert set(groups["enrichment"]["operations"]) == { node.id for node in spine.nodes if not node.id.startswith("spine.gates.") } - assert {"uk.full.pool", "uk.full.expand", "uk.full.expand.owned"} <= set( - groups["population"]["operations"] - ) + assert { + "uk.full.pool", + "uk.full.brma_table", + "uk.full.expand", + "uk.full.expand.owned", + "uk.full.brma", + } <= set(groups["population"]["operations"]) assert set(groups["checks"]["operations"]) >= { "spine.gates.assembled", "uk.full.gates.preflight", diff --git a/packages/microcosm-build/tests/engine_free/uk/test_uk_rowwise_candidate.py b/packages/microcosm-build/tests/engine_free/uk/test_uk_rowwise_candidate.py index ed35194466..0b420c3125 100644 --- a/packages/microcosm-build/tests/engine_free/uk/test_uk_rowwise_candidate.py +++ b/packages/microcosm-build/tests/engine_free/uk/test_uk_rowwise_candidate.py @@ -100,6 +100,8 @@ def test_graph_driver_dry_run_prints_the_operation_inventory( assert plan["configuration"]["calibration"]["epochs"] == 1500 assert plan["configuration"]["target_families"] is None assert {node["id"] for node in plan["nodes"]} >= { + "uk.full.brma_table", + "uk.full.brma", "uk.full.dense", "uk.full.target_selection", "uk.full.gates.calibrated", diff --git a/packages/microcosm-build/tests/engine_free/uk/test_uk_take_up_gate.py b/packages/microcosm-build/tests/engine_free/uk/test_uk_take_up_gate.py index 26a0170830..0f58abab6a 100644 --- a/packages/microcosm-build/tests/engine_free/uk/test_uk_take_up_gate.py +++ b/packages/microcosm-build/tests/engine_free/uk/test_uk_take_up_gate.py @@ -3,8 +3,13 @@ import numpy as np import pandas as pd -from microcosm.build.gate_battery import EvidenceContext +from microcosm.build.gate_battery import ( + EvidenceContext, + GateStatus, + evaluate_phase, +) from microcosm.build.uk_runtime.battery_bindings import UK_GATE_REGISTRY +from microcosm.build.uk_runtime.calibration_run import uk_scoped_gate_manifest from microcosm.build.uk_runtime.frs_brma import FRS_BRMA_OUTPUT_COLUMNS from microcosm.build.uk_runtime.frs_household_draws import ( FRS_HOUSEHOLD_DRAW_OUTPUT_COLUMNS, @@ -135,6 +140,30 @@ def test_brma_enum_domain_binding_fails_off_domain() -> None: assert "OFF_DOMAIN" in result.failures[0] +def test_local_phase_rechecks_brma_after_a_passing_assembled_check() -> None: + gates = uk_scoped_gate_manifest( + ("uk_brma_enum_domain", "uk_local_brma_enum_domain"), + phases=("assembled", "terminal"), + policy_suffix="brma_rewrite_test", + ) + frame = _frame() + context = EvidenceContext( + frame=frame, + artifacts={"brma_enum_domain": ("LONDON_A", "LONDON_B")}, + ) + assembled = evaluate_phase(gates, "assembled", context, registry=UK_GATE_REGISTRY) + assert len(assembled.outcomes) == 1 + assert assembled.outcomes[0].entry.id == "uk_brma_enum_domain" + assert assembled.outcomes[0].status is GateStatus.PASSED + + frame.table("household").loc[0, "brma"] = "OFF_DOMAIN" + local = evaluate_phase(gates, "terminal", context, registry=UK_GATE_REGISTRY) + assert len(local.outcomes) == 1 + assert local.outcomes[0].entry.id == "uk_local_brma_enum_domain" + assert local.outcomes[0].status is GateStatus.FAILED + assert "OFF_DOMAIN" in local.outcomes[0].result.failures[0] + + def test_student_loan_enum_domain_binding_resolves_person_column() -> None: frame = _frame() frame.table("person")["student_loan_plan"] = ["NONE"] * 9 + ["PLAN_4"] diff --git a/packages/microcosm-build/tests/golden/uk_geography_identity_main.json b/packages/microcosm-build/tests/golden/uk_geography_identity_main.json new file mode 100644 index 0000000000..af1e18d450 --- /dev/null +++ b/packages/microcosm-build/tests/golden/uk_geography_identity_main.json @@ -0,0 +1,208 @@ +{ + "schema_version": 1, + "main_commit": "a91c7536f7e318bdfdd8ee6214a7411257f0033a", + "main_kernel_source_sha256": "24668350de487949bbdc8df93ec4bfca460a4f69afe18da3d0811366fcb49c0a", + "main_fixture_source_sha256": "827430a03f011418222519dfbdcdb557e8d974c41f8f148e428ff1d5f168ac17", + "source": "frs", + "source_vintage": "toy_2023_24", + "cases": { + "lineage": { + "1": { + "household_ids": [ + 1, + 2, + 11, + 21, + 101, + 111, + 1001 + ], + "keys": [ + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"2\",[]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"spi_support_channel\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"2\",[[\"spi_support_channel\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"cgt_incidence_clone\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"spi_support_channel\",\"1\"],[\"cgt_incidence_clone\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"2\",[[\"cgt_support_split\",\"1\"]]]" + ] + }, + "2": { + "household_ids": [ + 1, + 2, + 11, + 21, + 101, + 111, + 1001, + 1000001, + 1000002, + 1000011, + 1000021, + 1000101, + 1000111, + 1001001 + ], + "keys": [ + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"2\",[]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"spi_support_channel\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"2\",[[\"spi_support_channel\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"cgt_incidence_clone\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"spi_support_channel\",\"1\"],[\"cgt_incidence_clone\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"2\",[[\"cgt_support_split\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"geographic_support\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"2\",[[\"geographic_support\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"spi_support_channel\",\"1\"],[\"geographic_support\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"2\",[[\"spi_support_channel\",\"1\"],[\"geographic_support\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"cgt_incidence_clone\",\"1\"],[\"geographic_support\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"spi_support_channel\",\"1\"],[\"cgt_incidence_clone\",\"1\"],[\"geographic_support\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"2\",[[\"cgt_support_split\",\"1\"],[\"geographic_support\",\"1\"]]]" + ] + } + }, + "residential_arms": { + "1": { + "household_ids": [ + 1, + 2, + 11, + 21, + 101, + 111, + 201, + 211, + 1001, + 1101, + 1201 + ], + "keys": [ + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"2\",[]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"spi_support_channel\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"2\",[[\"spi_support_channel\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"cgt_incidence_clone\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"spi_support_channel\",\"1\"],[\"cgt_incidence_clone\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"cgt_incidence_clone\",\"1\"],[\"cgt_residential_split\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"spi_support_channel\",\"1\"],[\"cgt_incidence_clone\",\"1\"],[\"cgt_residential_split\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"2\",[[\"cgt_support_split\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"2\",[[\"cgt_support_split\",\"1\"],[\"cgt_incidence_clone\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"2\",[[\"cgt_support_split\",\"1\"],[\"cgt_incidence_clone\",\"1\"],[\"cgt_residential_split\",\"1\"]]]" + ] + }, + "2": { + "household_ids": [ + 1, + 2, + 11, + 21, + 101, + 111, + 201, + 211, + 1001, + 1101, + 1201, + 1000001, + 1000002, + 1000011, + 1000021, + 1000101, + 1000111, + 1000201, + 1000211, + 1001001, + 1001101, + 1001201 + ], + "keys": [ + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"2\",[]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"spi_support_channel\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"2\",[[\"spi_support_channel\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"cgt_incidence_clone\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"spi_support_channel\",\"1\"],[\"cgt_incidence_clone\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"cgt_incidence_clone\",\"1\"],[\"cgt_residential_split\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"spi_support_channel\",\"1\"],[\"cgt_incidence_clone\",\"1\"],[\"cgt_residential_split\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"2\",[[\"cgt_support_split\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"2\",[[\"cgt_support_split\",\"1\"],[\"cgt_incidence_clone\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"2\",[[\"cgt_support_split\",\"1\"],[\"cgt_incidence_clone\",\"1\"],[\"cgt_residential_split\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"geographic_support\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"2\",[[\"geographic_support\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"spi_support_channel\",\"1\"],[\"geographic_support\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"2\",[[\"spi_support_channel\",\"1\"],[\"geographic_support\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"cgt_incidence_clone\",\"1\"],[\"geographic_support\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"spi_support_channel\",\"1\"],[\"cgt_incidence_clone\",\"1\"],[\"geographic_support\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"cgt_incidence_clone\",\"1\"],[\"cgt_residential_split\",\"1\"],[\"geographic_support\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"spi_support_channel\",\"1\"],[\"cgt_incidence_clone\",\"1\"],[\"cgt_residential_split\",\"1\"],[\"geographic_support\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"2\",[[\"cgt_support_split\",\"1\"],[\"geographic_support\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"2\",[[\"cgt_support_split\",\"1\"],[\"cgt_incidence_clone\",\"1\"],[\"geographic_support\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"2\",[[\"cgt_support_split\",\"1\"],[\"cgt_incidence_clone\",\"1\"],[\"cgt_residential_split\",\"1\"],[\"geographic_support\",\"1\"]]]" + ] + } + }, + "combinations": { + "1": { + "household_ids": [ + 100, + 101, + 102, + 103, + 104, + 105, + 106, + 107 + ], + "keys": [ + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"cgt_support_split\",\"2\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"cgt_incidence_clone\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"cgt_support_split\",\"2\"],[\"cgt_incidence_clone\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"spi_support_channel\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"spi_support_channel\",\"1\"],[\"cgt_support_split\",\"2\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"spi_support_channel\",\"1\"],[\"cgt_incidence_clone\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"spi_support_channel\",\"1\"],[\"cgt_support_split\",\"2\"],[\"cgt_incidence_clone\",\"1\"]]]" + ] + }, + "2": { + "household_ids": [ + 100, + 101, + 102, + 103, + 104, + 105, + 106, + 107, + 100100, + 100101, + 100102, + 100103, + 100104, + 100105, + 100106, + 100107 + ], + "keys": [ + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"cgt_support_split\",\"2\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"cgt_incidence_clone\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"cgt_support_split\",\"2\"],[\"cgt_incidence_clone\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"spi_support_channel\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"spi_support_channel\",\"1\"],[\"cgt_support_split\",\"2\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"spi_support_channel\",\"1\"],[\"cgt_incidence_clone\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"spi_support_channel\",\"1\"],[\"cgt_support_split\",\"2\"],[\"cgt_incidence_clone\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"geographic_support\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"cgt_support_split\",\"2\"],[\"geographic_support\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"cgt_incidence_clone\",\"1\"],[\"geographic_support\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"cgt_support_split\",\"2\"],[\"cgt_incidence_clone\",\"1\"],[\"geographic_support\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"spi_support_channel\",\"1\"],[\"geographic_support\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"spi_support_channel\",\"1\"],[\"cgt_support_split\",\"2\"],[\"geographic_support\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"spi_support_channel\",\"1\"],[\"cgt_incidence_clone\",\"1\"],[\"geographic_support\",\"1\"]]]", + "uk-household-v1:[\"frs\",\"toy_2023_24\",\"1\",[[\"spi_support_channel\",\"1\"],[\"cgt_support_split\",\"2\"],[\"cgt_incidence_clone\",\"1\"],[\"geographic_support\",\"1\"]]]" + ] + } + } + } +} diff --git a/packages/microcosm-data/src/microcosm/data/contract.py b/packages/microcosm-data/src/microcosm/data/contract.py index 38ba3e05af..040fcba818 100644 --- a/packages/microcosm-data/src/microcosm/data/contract.py +++ b/packages/microcosm-data/src/microcosm/data/contract.py @@ -442,13 +442,13 @@ def compatibility_claim_declarer_error(declared_by: object) -> str | None: # fingerprint derives from the manifest digest. Editing the spec moves all # three here in the same reviewed change. _UK_GATE_BATTERY_POLICY_SHA256 = ( - "bb727d8c8af98b9d3376d5e8eabb2eaeb708a2dd26379cbc3599af38f29a03c4" + "e6298f9017081efe4b813fa467a7cc9aca5a22e553f26f434a50d65cd0604a17" ) _UK_GATE_BATTERY_GATES_MANIFEST_SHA256 = ( - "a34f6782a2f2f2da7126d2e3a59881745130ff4d73733b8bd4c38eb425549061" + "c8b980c5367d937ac42247fcd14a3023cc1ea9559e798246107af2b7aee2dd8f" ) _UK_GATE_BATTERY_SPEC_FINGERPRINT = ( - "7fffef7f4f2aaeb25a42c06f94388416858ac8d6305d64f0c59d8bc840333872" + "2018dc382accdc470d280c266e4a37615c61b26639e3f7faa894f3da7cecc47b" ) #: Spec entry id -> the legacy gate name whose observable detail checks #: apply unchanged (the battery re-keys the report by entry id; the gate @@ -467,6 +467,7 @@ def compatibility_claim_declarer_error(declared_by: object) -> str | None: "uk_export_surface": "export_surface", "uk_take_up_signal": "take_up_signal", "uk_brma_enum_domain": "enum_domain", + "uk_local_brma_enum_domain": "enum_domain", "uk_ons_household_type_enum_domain": "enum_domain", "uk_capital_gains_asset_type_enum_domain": "enum_domain", "uk_uc_deduction_combination_enum_domain": "enum_domain", @@ -578,6 +579,7 @@ def compatibility_claim_declarer_error(declared_by: object) -> str | None: "spine_agreement", "terminal", ), + "uk_local_brma_enum_domain": ("enum_domain", "terminal"), "uk_local_area_support": ("area_support", "terminal"), "uk_local_target_fit": ("target_fit", "terminal"), "uk_local_per_family_fit": ("per_family_fit", "terminal"), @@ -737,6 +739,7 @@ def dataset_role_for_line(line: str) -> str: _UK_DENSE_GATE_ENTRY_IDS = frozenset( { "uk_local_area_support", + "uk_local_brma_enum_domain", "uk_local_geography_ladder_post_calibration", "uk_local_per_family_fit", "uk_local_target_fit", @@ -747,9 +750,9 @@ def dataset_role_for_line(line: str) -> str: _UK_DENSE_RELEASE_BLOCKING_IDS = _UK_DENSE_GATE_ENTRY_IDS _UK_DENSE_GATE_PHASES = ("terminal",) _UK_DENSE_GATE_DIGESTS = { - "gates_manifest_sha256": "3f0d6cf8ad29f960d3bb1006528a288a360747d203e7ca090a36dbabd781d301", - "policy_sha256": "ad63a0aa05ac127d45a42d82c5a0e376f23cf0f708c7b6d25958479b3483bbeb", - "spec_fingerprint": "c722d3741e806269d70003fb536947950dabf44bf5b5b1c27e835bdc6d8cdb00", + "gates_manifest_sha256": "aaab15a7e2d0d8b9ba6873fe1a8a41910902cb02bbc4d028a9016d717ad2c527", + "policy_sha256": "b3601174408fd9035f3415c85f688e40d14ef8be37b5e1d6f95e224787bf056a", + "spec_fingerprint": "337f37f42dbf54c212fd4f170dd61db71f61b9568fe92b18ded650b75044be5f", } _UK_DENSE_SOURCE_COVERAGE_KEYS = ( "spine", @@ -794,6 +797,7 @@ def _uk_line_cut_tag_re(release_id: str) -> re.Pattern[str] | None: _UK_CERTIFICATION_EXCLUDED_GATE_IDS = frozenset( { "uk_local_geography_ladder_post_calibration", + "uk_local_brma_enum_domain", "uk_local_area_support", "uk_local_target_fit", "uk_local_per_family_fit", diff --git a/packages/microcosm-data/tests/engine_free/shared/test_contract.py b/packages/microcosm-data/tests/engine_free/shared/test_contract.py index f92f4d3edb..de30f67ec1 100644 --- a/packages/microcosm-data/tests/engine_free/shared/test_contract.py +++ b/packages/microcosm-data/tests/engine_free/shared/test_contract.py @@ -214,13 +214,13 @@ def _trusted_terminal_gate_signing_key(monkeypatch) -> None: UK_GATE_BATTERY_PRODUCER = "microcosm.build.gate_battery" UK_GATE_BATTERY_SIGNING_KEY_ENV = "MICROCOSM_UK_TERMINAL_GATE_SIGNING_KEY" UK_GATE_BATTERY_POLICY_SHA256 = ( - "bb727d8c8af98b9d3376d5e8eabb2eaeb708a2dd26379cbc3599af38f29a03c4" + "e6298f9017081efe4b813fa467a7cc9aca5a22e553f26f434a50d65cd0604a17" ) UK_GATE_BATTERY_GATES_MANIFEST_SHA256 = ( - "a34f6782a2f2f2da7126d2e3a59881745130ff4d73733b8bd4c38eb425549061" + "c8b980c5367d937ac42247fcd14a3023cc1ea9559e798246107af2b7aee2dd8f" ) UK_GATE_BATTERY_SPEC_FINGERPRINT = ( - "7fffef7f4f2aaeb25a42c06f94388416858ac8d6305d64f0c59d8bc840333872" + "2018dc382accdc470d280c266e4a37615c61b26639e3f7faa894f3da7cecc47b" ) UK_GATE_BATTERY_DEGENERATE_EVIDENCE_SHA256 = ( "6f0243bcda09dad26945376230c44ec3cf55d4e417c3a25e29bae8c59bc1a69d" @@ -409,6 +409,7 @@ def _trusted_terminal_gate_signing_key(monkeypatch) -> None: "terminal", None, ), + "uk_local_brma_enum_domain": ("enum_domain", "terminal", "enum_domain"), "uk_local_area_support": ("area_support", "terminal", None), "uk_local_target_fit": ("target_fit", "terminal", None), "uk_local_per_family_fit": ("per_family_fit", "terminal", None), diff --git a/test_support/microcosm_build/logbook_isolation.py b/test_support/microcosm_build/logbook_isolation.py new file mode 100644 index 0000000000..43680c7d8d --- /dev/null +++ b/test_support/microcosm_build/logbook_isolation.py @@ -0,0 +1,20 @@ +"""Discover literal Logbook aliases without triggering foreign lazy imports.""" + +import sys +from types import ModuleType + + +def logbook_urlopen_holders(real_urlopen: object) -> list[ModuleType]: + """Return imported modules binding the Logbook HTTP function directly. + + The CLI imports ``urlopen`` by value, so the isolation fixture must patch + aliases as well as the defining module. Inspect module globals rather than + looking up attributes: unrelated modules may resolve missing attributes + through ``__getattr__`` hooks that import or initialize native libraries. + """ + return [ + module + for module in list(sys.modules.values()) + if isinstance(module, ModuleType) + and vars(module).get("urlopen") is real_urlopen + ] diff --git a/test_support/microcosm_build/uk_full_population_graph.py b/test_support/microcosm_build/uk_full_population_graph.py index 91d985d95a..b3f588d636 100644 --- a/test_support/microcosm_build/uk_full_population_graph.py +++ b/test_support/microcosm_build/uk_full_population_graph.py @@ -5,20 +5,26 @@ # ruff: noqa: F401 import hashlib +from dataclasses import replace import numpy as np import pandas as pd import pytest +from microcosm.build.stochastic_assignment import stable_identity_uniforms +from microcosm.build.uk_runtime.atomic_household_identity import household_draw_key from microcosm.build.uk_runtime.graph_kernels import UKClaimKernel from microcosm.build.uk_runtime.graph_population import ( + UKBRMATableKernel, append_uk_population_nodes, register_uk_population_kernels, ) from microcosm.build.uk_runtime.rowwise_dataset import ( clone_uk_dataset_with_ladder_geography, + ladder_clone_index_column, ) from microcosm.frame import Frame +from microcosm.frame.schema import VariableMetadata from microcosm.graph import ( Capabilities, ContentStore, @@ -40,6 +46,125 @@ ) SOURCE_VINTAGE = "toy_2023_24" +BRMA_REGIONS = ("LONDON", "WALES", "SCOTLAND", "NORTHERN_IRELAND") + + +def toy_brma_count_resource(): + """Invented cells: alphabetical order differs from the engine's rate order.""" + return { + "cells": { + region: { + category: { + f"{region}_A": 1, + f"{region}_B": 2, + f"{region}_C": 1, + } + for category in "ABCDE" + } + for region in BRMA_REGIONS + } + } + + +class FakeBRMAEngine: + """A UK protocol double calculating per-unit category and BRMA rates.""" + + country = "uk" + + def __init__(self, *, uncapped=True): + self.uncapped = uncapped + self.calls = [] + + def variable_metadata(self, name): + if name == "uncapped_BRMA_LHA_rate" and not self.uncapped: + raise ValueError("Unknown variable 'uncapped_BRMA_LHA_rate'.") + if name not in { + "LHA_category", + "uncapped_BRMA_LHA_rate", + "BRMA_LHA_rate", + }: + raise ValueError(f"Unknown variable {name!r}.") + return VariableMetadata( + name, "benunit", "str" if name == "LHA_category" else "float", "year" + ) + + def materialize(self, frame, variables, period): + assert str(period) == "2023" + self.calls.append(tuple(variables)) + benunit = frame.table("benunit") + placements = frame.table("person").drop_duplicates("person_benunit_id") + household_ids = benunit["benunit_id"].map( + placements.set_index("person_benunit_id")["person_household_id"] + ) + values = {} + for variable in variables: + self.variable_metadata(variable) + if variable == "LHA_category": + values[variable] = np.full(len(benunit), "B", dtype=object) + else: + names = household_ids.map( + frame.table("household").set_index("household_id")["brma"] + ) + values[variable] = np.asarray( + [{"A": 300.0, "B": 100.0, "C": 200.0}[n[-1]] for n in names] + ) + return values + + +class FakeBRMATableKernel(UKBRMATableKernel): + """Exercise the real calculation without importing a country distribution.""" + + capabilities = replace(UKBRMATableKernel.capabilities, dependencies=()) + + +def register_toy_population_kernels(registry, *, engine=None): + engine = FakeBRMAEngine() if engine is None else engine + resource = toy_brma_count_resource() + production = KernelRegistry() + register_uk_population_kernels( + production, engine=engine, brma_count_resource=resource + ) + for kernel in production.as_mapping().values(): + registry.register( + FakeBRMATableKernel(engine=engine, count_resource=resource) + if kernel.ref == UKBRMATableKernel.ref + else kernel + ) + + +def pregeographic_key(row, *, vintage=SOURCE_VINTAGE): + """Independent lineage oracle, deliberately excluding local clone identity.""" + path = [] + if row.household_support_channel == "spi": + path.append(("spi_support_channel", int(row.household_support_clone_index))) + if row.cgt_support_copy_index: + path.append(("cgt_support_split", int(row.cgt_support_copy_index))) + if row.household_is_capital_gains_clone: + path.append(("cgt_incidence_clone", 1)) + if row.cgt_residential_clone_index: + path.append(("cgt_residential_split", int(row.cgt_residential_clone_index))) + return household_draw_key( + source="frs", + source_vintage=vintage, + source_household_id=int(row.source_household_id), + clone_path=tuple(path), + ) + + +def expected_clone_brmas(household, k, *, vintage=SOURCE_VINTAGE): + """Inverse the invented B,C,A walk directly, without BRMA spread helpers.""" + rows = list(household.itertuples(index=False)) + uniforms = stable_identity_uniforms( + [pregeographic_key(row, vintage=vintage) for row in rows], + seed=0, + salt="brma:clone_spread", + ) + clone_column = ladder_clone_index_column("household") + quantiles = (uniforms + household[clone_column].to_numpy() / k) % 1.0 + suffixes = np.where(quantiles < 0.5, "B", np.where(quantiles < 0.75, "C", "A")) + return np.asarray( + [f"{row.region}_{suffix}" for row, suffix in zip(rows, suffixes, strict=True)] + ) def source_frame(): @@ -49,6 +174,9 @@ def source_frame(): tables["benunit"]["would_claim_uc"] = True household = tables["household"] household["region"] = household["region"].astype("string") + household["brma"] = pd.array( + [f"{region}_A" for region in household["region"]], dtype="string" + ) # Explicit spine lineage: every toy household is its own FRS original. household["source_household_id"] = household["household_id"].astype("int64") household["household_support_channel"] = pd.array( @@ -126,12 +254,12 @@ def graph_and_registry(k, *, geography_assignment="legacy", definition=None, see source_year=2023, geography_assignment=geography_assignment, atomic_geography_definition=definition, - source_vintage=SOURCE_VINTAGE if geography_assignment == "atomic" else None, + source_vintage=SOURCE_VINTAGE, ) registry = KernelRegistry() registry.register(Source()) registry.register(UKClaimKernel()) - register_uk_population_kernels(registry) + register_toy_population_kernels(registry) return graph, registry diff --git a/test_support/microcosm_build/uk_rowwise_candidate.py b/test_support/microcosm_build/uk_rowwise_candidate.py index bd79c09cf0..b31a123c22 100644 --- a/test_support/microcosm_build/uk_rowwise_candidate.py +++ b/test_support/microcosm_build/uk_rowwise_candidate.py @@ -24,6 +24,7 @@ rowwise_staging, write_uk_national_frame, ) +from microcosm.build.uk_runtime.frs_brma import load_brma_count_resource from microcosm.build.uk_runtime.national_frame import ( uk_national_frame, ) @@ -243,6 +244,7 @@ def _write_staging_h5( "SCOTLAND", "NORTHERN_IRELAND", ) + brma_cells = load_brma_count_resource()["cells"] household_ids = list(range(1, 4 * households_per_region + 1)) source_household_ids: list[int] = [] support_clone_indices: list[int] = [] @@ -265,6 +267,16 @@ def _write_staging_h5( "region": [ region for region in region_names for _ in range(households_per_region) ], + "brma": [ + next( + name + for counts in brma_cells[region].values() + for name, count in counts.items() + if count > 0 + ) + for region in region_names + for _ in range(households_per_region) + ], "source_household_id": source_household_ids, "source_household_key": [ f"2023:{source_id}" for source_id in source_household_ids