diff --git a/changelog.d/uk-keyed-qrf-draws.changed.md b/changelog.d/uk-keyed-qrf-draws.changed.md new file mode 100644 index 0000000000..9b69b84ba6 --- /dev/null +++ b/changelog.d/uk-keyed-qrf-draws.changed.md @@ -0,0 +1 @@ +UK QRF stages draw from uniforms keyed on each record's own id, so a household's imputed wealth, consumption, services, UC report, SPI fill and salary sacrifice move only when its own inputs do: an engine change that moved a few households' net income no longer redraws the whole spine through a shifted positional stream. The switch itself redraws these stages once for every household, benefit unit and person, so a comparison of this release with the previous one shows that one-off shift. Each stage refuses ids that are not unique integers, since a float id would key a different draw. diff --git a/changelog.d/uk-pension-contribution-coherence.fixed.md b/changelog.d/uk-pension-contribution-coherence.fixed.md new file mode 100644 index 0000000000..6d46c3bcbd --- /dev/null +++ b/changelog.d/uk-pension-contribution-coherence.fixed.md @@ -0,0 +1 @@ +UK employee pension contributions are capped at the pay they come out of, on FRS rows and on SPI-channel rows. On an SPI-channel row the forest still decides, from the row's own drawn pay, whether it has an employer contribution; a member whose FRS twin has an ASHE rate takes that rate times its own pay, so the amount follows the redrawn pay, and any other member keeps the forest's amount, bounded by ASHE's top rate. diff --git a/docs/evidence/spec-engine/us-f0-coverage.json b/docs/evidence/spec-engine/us-f0-coverage.json index b6153a67b0..52fe0adb4f 100644 --- a/docs/evidence/spec-engine/us-f0-coverage.json +++ b/docs/evidence/spec-engine/us-f0-coverage.json @@ -1656,13 +1656,13 @@ "compiler_ir.node_slices" ], "expected": { - "map_sha256": "b3d3f769647212a81d796a28f2b5ff759f513572bd3c70e253816c88ff496091", - "protocol_sha256": "d756be07ab66c179f8612eb5c614527586004506d88be3e1ab62b934de0b61ce" + "map_sha256": "0db599df419dd16779097361dc72afdefb8df2abf8b134078da9c52d31d8e017", + "protocol_sha256": "2b7bdc1d6b0da1bdaf54b4cac29f10f2d3c8bed2e177729c2c89e6b33a099cc5" }, "failures": [], "observed": { - "map_sha256": "b3d3f769647212a81d796a28f2b5ff759f513572bd3c70e253816c88ff496091", - "protocol_sha256": "d756be07ab66c179f8612eb5c614527586004506d88be3e1ab62b934de0b61ce" + "map_sha256": "0db599df419dd16779097361dc72afdefb8df2abf8b134078da9c52d31d8e017", + "protocol_sha256": "2b7bdc1d6b0da1bdaf54b4cac29f10f2d3c8bed2e177729c2c89e6b33a099cc5" }, "status": "covered" }, @@ -1677,7 +1677,7 @@ "compiler_ir.seed_stream_map" ], "expected": { - "implementation_sha256": "d756be07ab66c179f8612eb5c614527586004506d88be3e1ab62b934de0b61ce", + "implementation_sha256": "2b7bdc1d6b0da1bdaf54b4cac29f10f2d3c8bed2e177729c2c89e6b33a099cc5", "protocol": "legacy-v1", "streams": [ "build_model", @@ -1698,7 +1698,7 @@ }, "failures": [], "observed": { - "implementation_sha256": "d756be07ab66c179f8612eb5c614527586004506d88be3e1ab62b934de0b61ce", + "implementation_sha256": "2b7bdc1d6b0da1bdaf54b4cac29f10f2d3c8bed2e177729c2c89e6b33a099cc5", "protocol": "legacy-v1", "streams": [ "build_model", @@ -2599,7 +2599,7 @@ "country": "us", "schema_id": "country_spec", "schema_version": 1, - "spec_sha256": "740d1ff5b2e9fa6da3154c44f8676f80109d404e6aa7a16c48ed8db83055b15d" + "spec_sha256": "19460fc168a6f9ed44e580d0e36777cce78bb33cb52b75630ee86203fa795684" } }, "report_schema_version": 3, @@ -2609,7 +2609,7 @@ "country": "us", "schema_id": "country_spec", "schema_version": 1, - "spec_sha256": "740d1ff5b2e9fa6da3154c44f8676f80109d404e6aa7a16c48ed8db83055b15d" + "spec_sha256": "19460fc168a6f9ed44e580d0e36777cce78bb33cb52b75630ee86203fa795684" }, "status": "pass" } diff --git a/packages/microcosm-build/src/microcosm/build/spec_engine/inventory_coverage.py b/packages/microcosm-build/src/microcosm/build/spec_engine/inventory_coverage.py index 01287d21a3..4a20e69687 100644 --- a/packages/microcosm-build/src/microcosm/build/spec_engine/inventory_coverage.py +++ b/packages/microcosm-build/src/microcosm/build/spec_engine/inventory_coverage.py @@ -367,8 +367,8 @@ # moves both again: calibrate gains budget_search_bracket, whose default # keeps the US bracket, and the budget search, which the US solve does # not run, changes its probes, so again no US solve changes. - "seed_map": "b3d3f769647212a81d796a28f2b5ff759f513572bd3c70e253816c88ff496091", - "seed_protocol": "d756be07ab66c179f8612eb5c614527586004506d88be3e1ab62b934de0b61ce", + "seed_map": "0db599df419dd16779097361dc72afdefb8df2abf8b134078da9c52d31d8e017", + "seed_protocol": "2b7bdc1d6b0da1bdaf54b4cac29f10f2d3c8bed2e177729c2c89e6b33a099cc5", "source_manifest": "f4f8bfeb79268a4043a843b47cec1d1c0e58e9f5dd6bc23417ada95446dd7b8e", "take_up": "9522ce40f7dea569312dd7a7beb474e5a5afadd2f1e10cb5876534c3ef623d35", "tail": "ac92829c88a1a4fb6460d61190918d5d99c6c377fc8dd8f62f02b332d09bf59c", diff --git a/packages/microcosm-build/src/microcosm/build/stochastic_assignment.py b/packages/microcosm-build/src/microcosm/build/stochastic_assignment.py index c7441b557d..ee35a0cc55 100644 --- a/packages/microcosm-build/src/microcosm/build/stochastic_assignment.py +++ b/packages/microcosm-build/src/microcosm/build/stochastic_assignment.py @@ -13,6 +13,7 @@ "assign_binary_from_rate", "assign_binary_with_anchored_residual", "clipped_normal_from_uniforms", + "integer_record_keys", "sample_categorical_from_counts", "stable_identity_uniforms", ] @@ -44,6 +45,25 @@ def stable_identity_uniforms( ) +def integer_record_keys( + ids: Sequence[object] | np.ndarray, *, label: str +) -> np.ndarray: + """Record ids to key ``stable_identity_uniforms`` on: integers, each once. + + The key string formats each id, so the integer ``1001`` and the float + ``1001.0`` draw differently, and a merge that turned an id column into + floats would silently redraw every record. A repeated id would give two + records the same uniforms. + """ + + keys = np.asarray(ids) + if keys.ndim != 1 or not np.issubdtype(keys.dtype, np.integer): + raise TypeError(f"{label} must be a one-dimensional array of integer ids.") + if np.unique(keys).size != keys.size: + raise ValueError(f"{label} must name each record once.") + return keys + + def assign_binary_from_rate( draws: Sequence[float] | np.ndarray, rate: float, diff --git a/packages/microcosm-build/src/microcosm/build/uk_runtime/etb_services.py b/packages/microcosm-build/src/microcosm/build/uk_runtime/etb_services.py index ccafd36e2b..ead29566b9 100644 --- a/packages/microcosm-build/src/microcosm/build/uk_runtime/etb_services.py +++ b/packages/microcosm-build/src/microcosm/build/uk_runtime/etb_services.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +from collections.abc import Sequence from dataclasses import dataclass, field from importlib.resources import files from pathlib import Path @@ -12,6 +13,10 @@ from microcosm.build.gates import FitWeightRecord from microcosm.build.source_manifest import SourceStageSpec +from microcosm.build.stochastic_assignment import ( + integer_record_keys, + stable_identity_uniforms, +) from microcosm.build.uk_runtime.bus_support_pricing import ( PRICE_BUS_SUPPORT_KIND, bus_support_pricing_operation, @@ -186,7 +191,10 @@ def __call__(self, frame: Frame) -> Frame: ) predictors = recipient_predictors(frame, self.engine) draws, records = impute_etb_services( - donor, predictors, seed=_qrf_seed(self.stage) + donor, + predictors, + seed=_qrf_seed(self.stage), + record_ids=frame.table("household")["household_id"].to_numpy(), ) clip_result = support_clip_to_donor( draws, donor, exempt=support_clip_exempt(self.stage) @@ -436,11 +444,28 @@ def person_sum(values: np.ndarray) -> np.ndarray: def impute_etb_services( - donor: pd.DataFrame, recipient: pd.DataFrame, *, seed: int, n_estimators: int = 100 + donor: pd.DataFrame, + recipient: pd.DataFrame, + *, + seed: int, + record_ids: Sequence[object] | np.ndarray, + n_estimators: int = 100, ) -> tuple[pd.DataFrame, tuple[FitWeightRecord, ...]]: + """Chain-draw the ETB services targets from household-keyed uniforms. + + Each draw reads uniforms keyed on the household's own id (``record_ids``, + in recipient row order) and the target, so a household's draws move only + when its own predictors do, whatever happens to the other recipients. + """ + from microcosm.fit import RegimeGatedQRF targets = UK_ETB_SERVICES_HOUSEHOLD_OUTPUT_COLUMNS[:3] + ids = integer_record_keys(record_ids, label="ETB record_ids") + if ids.shape != (len(recipient),) or pd.Series(ids).duplicated().any(): + raise ValueError( + "ETB record_ids must name every recipient once, in recipient row order." + ) model = RegimeGatedQRF(n_estimators=n_estimators, seed=seed) state = model.start_chain( donor, @@ -451,12 +476,15 @@ def impute_etb_services( raw = pd.DataFrame(index=recipient.index) records: list[FitWeightRecord] = [] for target in targets: + salt = f"{UK_ETB_SERVICES_FIT_NAME}:{target}" result = model.fit_draw_next( donor, recipient.loc[:, list(UK_ETB_SERVICES_PREDICTORS)], raw, state=state, weights="weight", + quantiles=stable_identity_uniforms(ids, seed=seed, salt=salt), + sign_uniforms=stable_identity_uniforms(ids, seed=seed, salt=f"{salt}:sign"), ) raw[target] = result.raw_draw records.append( diff --git a/packages/microcosm-build/src/microcosm/build/uk_runtime/frs_spine.py b/packages/microcosm-build/src/microcosm/build/uk_runtime/frs_spine.py index b7100e275a..9ce3ffb083 100644 --- a/packages/microcosm-build/src/microcosm/build/uk_runtime/frs_spine.py +++ b/packages/microcosm-build/src/microcosm/build/uk_runtime/frs_spine.py @@ -24,6 +24,10 @@ uk_national_frame, validate_uk_national_frame, ) +from microcosm.build.uk_runtime.pension_coherence import ( + cap_net_pay_contributions, + coherence_receipt, +) from microcosm.build.uk_runtime.uc_relationships import frs_uc_claimant_mask from microcosm.frame import Frame, WeightKind @@ -319,6 +323,7 @@ def __init__(self, raw_dir: str | Path, *, stage: SourceStageSpec) -> None: self.stage = stage self._sentinel_mapped_rows: int | None = None self._access_fund: dict[str, object] | None = None + self._net_pay: dict[str, object] | None = None def __call__(self, frame: Frame) -> Frame: evidence: dict[str, object] = {} @@ -328,6 +333,7 @@ def __call__(self, frame: Frame) -> Frame: capital = result.table("benunit")["frs_benunit_capital"] self._sentinel_mapped_rows = int((capital == UC_CAPITAL_UNAVAILABLE).sum()) self._access_fund = dict(evidence["access_fund"]) + self._net_pay = dict(evidence["net_pay_contributions"]) return result @staticmethod @@ -338,7 +344,11 @@ def checkpoint_metadata(self) -> dict[str, object]: """Report how loudly the capital availability rule and the access-fund repair fired.""" - if self._sentinel_mapped_rows is None or self._access_fund is None: + if ( + self._sentinel_mapped_rows is None + or self._access_fund is None + or self._net_pay is None + ): raise RuntimeError("checkpoint metadata requires a completed stage run.") return { "evidence": { @@ -348,6 +358,7 @@ def checkpoint_metadata(self) -> dict[str, object]: "mapped_rows": self._sentinel_mapped_rows, }, "access_fund": dict(self._access_fund), + "net_pay_contributions": dict(self._net_pay), } } @@ -585,7 +596,7 @@ def _assemble_frame( pe_person["would_claim_carers_allowance"] = ( pe_person["carers_allowance_reported"] > 0 ) - _add_person_expenses(pe_person, person, frs) + _add_person_expenses(pe_person, person, frs, evidence=evidence) pe_benunit["is_married"] = _number(benunit_raw, "famtypb2").isin([5, 7]) pe_benunit["dependent_children"] = dependent_children @@ -1004,7 +1015,11 @@ def _add_benefits( def _add_person_expenses( - pe_person: pd.DataFrame, person: pd.DataFrame, frs: Mapping[str, pd.DataFrame] + pe_person: pd.DataFrame, + person: pd.DataFrame, + frs: Mapping[str, pd.DataFrame], + *, + evidence: dict[str, object] | None = None, ) -> None: household = frs["househol"].set_index("household_id").sort_index() ctrebamt = pd.Series(_number(household, "ctrebamt").values, index=household.index) @@ -1055,7 +1070,7 @@ def _add_person_expenses( * WEEKS_IN_YEAR, ) job = frs["job"] - pe_person["employee_pension_contributions"] = np.maximum( + deducted = np.maximum( 0, _sum_to_entity( _number(job, "deduc1").fillna(0), @@ -1064,6 +1079,36 @@ def _add_person_expenses( ) * WEEKS_IN_YEAR, ) + # A net-pay deduction cannot exceed the pay it comes out of; the FRS + # records each job's deduction apart from the adult's usual earnings, so + # the deduction is capped at them (uk_runtime.pension_coherence). + capped, changed = cap_net_pay_contributions( + np.asarray(deducted, dtype=np.float64), + pe_person["employment_income"].to_numpy(dtype=np.float64), + ) + pe_person["employee_pension_contributions"] = capped + if evidence is not None: + # A deduction beside no recorded pay may stand for pay the adult tab + # left blank, not for a contribution above pay, so the two cases are + # counted apart. + recorded_pay = pe_person["employment_income"].to_numpy(dtype=np.float64) > 0.0 + receipt = { + "before": np.asarray(deducted, dtype=np.float64), + "after": capped, + "weights": person["household_id"] + .map(_raw_number(household, "gross4")) + .fillna(0.0) + .to_numpy(dtype=np.float64), + "source_ids": person["household_id"].to_numpy(), + } + evidence["net_pay_contributions"] = { + "above_recorded_pay": coherence_receipt( + changed=changed & recorded_pay, **receipt + ), + "beside_no_recorded_pay": coherence_receipt( + changed=changed & ~recorded_pay, **receipt + ), + } pe_person["pension_contributions_via_salary_sacrifice"] = np.maximum( 0, _sum_to_entity( diff --git a/packages/microcosm-build/src/microcosm/build/uk_runtime/lcfs_consumption.py b/packages/microcosm-build/src/microcosm/build/uk_runtime/lcfs_consumption.py index 9749b97ee3..c46484068d 100644 --- a/packages/microcosm-build/src/microcosm/build/uk_runtime/lcfs_consumption.py +++ b/packages/microcosm-build/src/microcosm/build/uk_runtime/lcfs_consumption.py @@ -13,7 +13,10 @@ from microcosm.build.gates import EXPLICIT_KIND, FitWeightRecord from microcosm.build.source_manifest import SourceStageSpec -from microcosm.build.stochastic_assignment import stable_identity_uniforms +from microcosm.build.stochastic_assignment import ( + integer_record_keys, + stable_identity_uniforms, +) from microcosm.build.uk_runtime.bus_fare_pricing import ( PRICE_BUS_JOURNEYS_KIND, UK_DFT_BUS_JOURNEYS_RESOURCE, @@ -575,6 +578,7 @@ def __call__(self, frame: Frame) -> Frame: recipient, seed=_operation_seed(self.stage, "fit_weighted_qrf_chain"), n_estimators=_qrf_n_estimators(self.stage), + record_ids=frame.table("household")["household_id"].to_numpy(), ) clip_result = support_clip_to_donor( imputation.draws, donor, exempt=support_clip_exempt(self.stage) @@ -3108,12 +3112,17 @@ def impute_lcfs_consumption( *, seed: int, n_estimators: int, + record_ids: Sequence[object] | np.ndarray, ) -> UKLCFSConsumptionImputationResult: """Chain-draw the targets in the declared order. ``bus_fare_spending`` stays a chain target so later targets condition on the same raw draw; the emitted fares are then replaced by the priced value where an area prices (microcosm#930). + + Each draw reads uniforms keyed on the household's own id (``record_ids``, + in recipient row order) and the target, never a positional stream, so a + household's consumption moves only when its own predictors do. """ from microcosm.fit import RegimeGatedQRF @@ -3121,6 +3130,11 @@ def impute_lcfs_consumption( donor_encoded, recipient_encoded, predictors = _encode_consumption_predictors( donor, recipient_predictor_frame ) + ids = integer_record_keys(record_ids, label="LCFS record_ids") + if ids.shape != (len(recipient_encoded),) or pd.Series(ids).duplicated().any(): + raise ValueError( + "LCFS record_ids must name every recipient once, in recipient row order." + ) model = RegimeGatedQRF(n_estimators=n_estimators, seed=seed) state = model.start_chain( donor_encoded, @@ -3133,12 +3147,15 @@ def impute_lcfs_consumption( fit_records: list[FitWeightRecord] = [] for target in UK_LCFS_CONSUMPTION_TARGET_COLUMNS: features = recipient_encoded.loc[:, list(predictors)] + salt = f"{UK_LCFS_CONSUMPTION_FIT_NAME}:{target}" result = model.fit_draw_next( donor_encoded, features, raw, state=state, weights="household_weight", + quantiles=stable_identity_uniforms(ids, seed=seed, salt=salt), + sign_uniforms=stable_identity_uniforms(ids, seed=seed, salt=f"{salt}:sign"), ) raw[target] = result.raw_draw draws[target] = result.raw_draw diff --git a/packages/microcosm-build/src/microcosm/build/uk_runtime/pension_coherence.py b/packages/microcosm-build/src/microcosm/build/uk_runtime/pension_coherence.py new file mode 100644 index 0000000000..e5e3ad90ca --- /dev/null +++ b/packages/microcosm-build/src/microcosm/build/uk_runtime/pension_coherence.py @@ -0,0 +1,118 @@ +"""Pension contributions that cohere with the pay they come from. + +policyengine-uk relieves an employee's pension contributions through the net +pay arrangement, which cannot take more than the pay the contributions are +deducted from, and from 2.134.5 (pe-uk#2240) it counts every contribution, the +employer's included, as pension input for the annual allowance. Two readings +of the spine broke that link: + +- the FRS records each job's pension deduction (JOB.DEDUC1) apart from the + adult's usual earnings (ADULT.INEARNS), so a deduction can exceed the pay + recorded beside it; +- an SPI-channel row redraws its pay from the tape, while a forest trained on + FRS rows draws its contributions. The forest reads the drawn pay, so it + draws whether the row has an employer contribution in line with that pay, + but it can only draw amounts that FRS rows in the same leaf received, so it + cannot scale them to pay above the FRS's range. + +So employee contributions are capped at pay. An SPI-channel row keeps the +forest's decision on whether it has an employer contribution, and a row that +has one takes its FRS twin's rate, which the ASHE draw of microcosm#1069 set, +times its own pay. The twin's scheme record is not carried over: it describes +the twin's job at the twin's pay, and the row's pay is redrawn, so carrying it +put members at pay where few people are members (María's ruling of 2026-10-21, +after Vahid's review of microcosm#1332). Where the twin has no rate to read, +the forest's amount stays, bounded by ASHE's top employer rate times the row's +pay. Personal (relief-at-source) contributions are left as recorded: the +engine itself limits their relief. +""" + +from __future__ import annotations + +from collections.abc import Mapping + +import numpy as np + +from microcosm.build.uk_runtime.employer_pension_contributions import ( + ASHE_EMPLOYER_RATE_BANDS, +) +from microcosm.build.uk_runtime.receipt_counts import distinct_count, support_band + +#: ASHE's highest employer contribution rate, the top of its open band as the +#: employer draw closes it (``employer_pension_contributions.py``). +ASHE_EMPLOYER_RATE_CEILING = ASHE_EMPLOYER_RATE_BANDS[-1][2] + + +def cap_net_pay_contributions( + contributions: np.ndarray, pay: np.ndarray +) -> tuple[np.ndarray, np.ndarray]: + """Employee contributions capped at pay, and the rows the cap changed.""" + + recorded = np.asarray(contributions, dtype=np.float64) + ceiling = np.maximum(np.asarray(pay, dtype=np.float64), 0.0) + if recorded.shape != ceiling.shape: + raise ValueError("contributions and pay must have one value per person.") + capped = np.minimum(recorded, ceiling) + return capped, capped < recorded + + +def twin_rate_employer_contributions( + *, + drawn: np.ndarray, + pay: np.ndarray, + twin_pay: np.ndarray, + twin_employer: np.ndarray, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Employer contributions at the FRS twin's rate times the row's own pay. + + The drawn contribution decides membership: a row the forest gave none + keeps none. A member whose twin has pay and an employer contribution + takes the twin's rate times its own pay; any other member keeps its drawn + contribution, bounded by ASHE's top rate times its pay. + + Returns the contributions, the rows that took their twin's rate, and the + rows whose drawn contribution the ASHE ceiling bounded. + """ + + drawn = np.maximum(np.asarray(drawn, dtype=np.float64), 0.0) + pay = np.maximum(np.asarray(pay, dtype=np.float64), 0.0) + twin_pay = np.asarray(twin_pay, dtype=np.float64) + twin_employer = np.maximum(np.asarray(twin_employer, dtype=np.float64), 0.0) + if not drawn.shape == pay.shape == twin_pay.shape == twin_employer.shape: + raise ValueError("employer inputs must have one value per person.") + has_rate = (twin_pay > 0.0) & (twin_employer > 0.0) + rate = np.divide( + twin_employer, twin_pay, out=np.zeros_like(twin_pay), where=has_rate + ) + from_twin = (drawn > 0.0) & has_rate + bounded = np.minimum(drawn, ASHE_EMPLOYER_RATE_CEILING * pay) + employer = np.where(from_twin, rate * pay, bounded) + return employer, from_twin, ~from_twin & (bounded < drawn) + + +def coherence_receipt( + *, + changed: np.ndarray, + before: np.ndarray, + after: np.ndarray, + weights: np.ndarray, + source_ids: np.ndarray, +) -> Mapping[str, object]: + """Banded support and weighted amounts for one coherence step. + + Support is counted on distinct source households (``support_band``), and + the weighted amounts are withheld as "<10" where one to nine of them carry + the change, so the receipt gives back no small cell. + """ + + mask = np.asarray(changed, dtype=bool) + weight = np.asarray(weights, dtype=np.float64) + sources = distinct_count(source_ids, mask, weight) + small = 0 < sources < 10 + before_mass = float(np.dot(weight[mask], np.asarray(before, dtype=float)[mask])) + after_mass = float(np.dot(weight[mask], np.asarray(after, dtype=float)[mask])) + return { + "support": support_band(sources), + "weighted_before": "<10" if small else before_mass, + "weighted_after": "<10" if small else after_mass, + } diff --git a/packages/microcosm-build/src/microcosm/build/uk_runtime/salary_sacrifice.py b/packages/microcosm-build/src/microcosm/build/uk_runtime/salary_sacrifice.py index 016aa9922a..32cc952bcd 100644 --- a/packages/microcosm-build/src/microcosm/build/uk_runtime/salary_sacrifice.py +++ b/packages/microcosm-build/src/microcosm/build/uk_runtime/salary_sacrifice.py @@ -13,7 +13,10 @@ import pandas as pd from microcosm.build.source_manifest import SourceStageSpec -from microcosm.build.stochastic_assignment import stable_identity_uniforms +from microcosm.build.stochastic_assignment import ( + integer_record_keys, + stable_identity_uniforms, +) from microcosm.build.uk_runtime.cgt_structure import ( HOUSEHOLD_IS_CGT_CLONE, _assert_closed_world_operations, @@ -278,7 +281,26 @@ def impute_salary_sacrifice(frame: Frame) -> UKSalarySacrificeResult: weights="_fit_weight", ) if predict_mask.any(): - predictions = fitted.predict(numeric.loc[predict_mask, list(SALSAC_PREDICTORS)]) + # Uniforms keyed on each person's own id, never a positional stream: + # a person's predicted sacrifice moves only when their own predictors do. + keys = integer_record_keys( + person.loc[predict_mask, "person_id"].to_numpy(), + label="salary sacrifice person_id", + ) + salt = f"salary_sacrifice:{SALSAC_OUTPUT}" + predictions = fitted.predict_from_uniforms( + numeric.loc[predict_mask, list(SALSAC_PREDICTORS)], + quantiles={ + SALSAC_OUTPUT: stable_identity_uniforms( + keys, seed=SALSAC_QRF_SEED, salt=salt + ) + }, + sign_uniforms={ + SALSAC_OUTPUT: stable_identity_uniforms( + keys, seed=SALSAC_QRF_SEED, salt=f"{salt}:sign" + ) + }, + ) predicted = pd.to_numeric(predictions[SALSAC_OUTPUT], errors="coerce").to_numpy( dtype=float ) @@ -346,6 +368,16 @@ def impute_salary_sacrifice(frame: Frame) -> UKSalarySacrificeResult: raise ValueError( "Salary-sacrifice amounts must not exceed the pay that remains." ) + # Net-pay contributions come out of pay: the spine caps them at pay + # (uk_runtime.pension_coherence), and the conversion takes the same amount + # off both, so a contribution above the pay that remains is a defect. The + # caps hold here only while no stage between them and this one lowers + # pay: today only hmrc_spi_income_spine and salary_sacrifice write + # employment_income, so a new pay rewrite must cap contributions again. + if (employee > np.maximum(employment_income, 0.0) + 1e-6).any(): + raise ValueError( + "Employee pension contributions must not exceed the pay that remains." + ) person[SALSAC_OUTPUT] = final_ss person["employee_pension_contributions"] = employee person["employment_income"] = employment_income diff --git a/packages/microcosm-build/src/microcosm/build/uk_runtime/spi_income.py b/packages/microcosm-build/src/microcosm/build/uk_runtime/spi_income.py index 2450b4d587..c784a7bba0 100644 --- a/packages/microcosm-build/src/microcosm/build/uk_runtime/spi_income.py +++ b/packages/microcosm-build/src/microcosm/build/uk_runtime/spi_income.py @@ -15,7 +15,10 @@ import pandas as pd from microcosm.build.gates import FitWeightRecord -from microcosm.build.stochastic_assignment import stable_identity_uniforms +from microcosm.build.stochastic_assignment import ( + integer_record_keys, + stable_identity_uniforms, +) from microcosm.build.uk_runtime.frs_disability import ( UKDWPDisabilityCategoryRates, UKDWPDisabilityFlagRates, @@ -35,6 +38,11 @@ LOSS_RELIEF_COLUMN, self_employment_income_after_loss_relief, ) +from microcosm.build.uk_runtime.pension_coherence import ( + cap_net_pay_contributions, + coherence_receipt, + twin_rate_employer_contributions, +) from microcosm.build.uk_runtime.spi_rank_transport import ( SPI_STAGE1_ANALOGUE_COLUMNS, SPI_TRANSPORT_FIT_NAME, @@ -674,6 +682,7 @@ class UKSPIIncomeImputationResult: reference_training_rows: int = 0 state_pension_age: int | None = None gainful_self_employment_refresh: Mapping[str, object] | None = None + pension_coherence: Mapping[str, object] | None = None @dataclass(frozen=True) @@ -912,6 +921,14 @@ def impute_uk_spi_income_support( undrawn = role_domain & ~spi_people if not spi_people.any(): raise ValueError("SPI support has no recipients with taxable income.") + # The FRS twin's pay and employer contribution, before the tape redraws + # the pay: the twin's employer rate is carried onto the drawn pay below. + twin_pay = person.loc[spi_people, "employment_income"].to_numpy( + dtype=np.float64, copy=True + ) + twin_employer = person.loc[spi_people, "employer_pension_contributions"].to_numpy( + dtype=np.float64, copy=True + ) recipient_domain = { "minimum_age": SPI_MINIMUM_RECIPIENT_AGE, "role_column": SPI_RECIPIENT_ROLE_COLUMN, @@ -1151,8 +1168,12 @@ def fit_reference(encoded: pd.DataFrame, observed: pd.DataFrame): income_predictors=income_predictors, ) # Queried over the age domain like stage 1; a dependant's row conditions - # on its twin's values and is discarded below, and the draws consume the - # RNG stream row by row, so no recipient's draw depends on it. + # on its twin's values and is discarded below, and each row draws from + # uniforms keyed on its own person_id, so no recipient's draw depends on + # it. Stage 1 keys on the twin's person_source_id on purpose and counts + # the shared identities; stage 2 does not, because the SPI channel and + # the band donors copy one FRS person more than once, and each copy + # draws its stage-2 values apart, as under the positional stream. target_predictors = _person_predictors( person.loc[spi_age_domain], household, @@ -1204,7 +1225,25 @@ def fit_reference(encoded: pd.DataFrame, observed: pd.DataFrame): list(stage2_outputs), weights="importance", ) - stage2_draws = stage2.predict(encoded_target) + stage2_keys = integer_record_keys( + person.loc[spi_age_domain, "person_id"].to_numpy(), + label="SPI stage-2 person_id", + ) + stage2_draws = stage2.predict_from_uniforms( + encoded_target, + quantiles={ + target: stable_identity_uniforms( + stage2_keys, seed=seed + 1, salt=f"{FRS_ONLY_FIT_NAME}:{target}" + ) + for target in stage2_outputs + }, + sign_uniforms={ + target: stable_identity_uniforms( + stage2_keys, seed=seed + 1, salt=f"{FRS_ONLY_FIT_NAME}:{target}:sign" + ) + for target in stage2_outputs + }, + ) _validate_predictions( stage2_draws, expected=stage2_outputs, @@ -1215,6 +1254,13 @@ def fit_reference(encoded: pd.DataFrame, observed: pd.DataFrame): raise ValueError("FRS-only stage-2 produced negative non-negative outputs.") for column in stage2_outputs: person.loc[spi_people, column] = stage2_draws[column].to_numpy() + person, pension_coherence = _spi_pension_coherence( + person, + household, + spi_people=spi_people, + twin_pay=twin_pay, + twin_employer=twin_employer, + ) person = _owner_manager_flags_from_draws(person, spi_people=spi_people) if ( state_pension_age_guard is not None @@ -1284,6 +1330,7 @@ def fit_reference(encoded: pd.DataFrame, observed: pd.DataFrame): reference_training_rows=int(reference_people.sum()), state_pension_age=int(state_pension_age), gainful_self_employment_refresh=gainful_self_employment, + pension_coherence=pension_coherence, ) @@ -2459,6 +2506,73 @@ def _stage2_outputs(person: pd.DataFrame) -> tuple[tuple[str, ...], dict[str, st return tuple(outputs), reviewed +def _spi_pension_coherence( + person: pd.DataFrame, + household: pd.DataFrame, + *, + spi_people: pd.Series, + twin_pay: np.ndarray, + twin_employer: np.ndarray, +) -> tuple[pd.DataFrame, Mapping[str, object]]: + """Tie an SPI row's contributions to the pay the tape drew for it. + + Stage 2 draws the contributions from a forest trained on FRS rows. It reads + the drawn pay, so whether the row has an employer contribution follows that + pay, but it cannot scale amounts beyond the FRS's pay. So the employee + (net-pay) contribution is capped at the row's pay, and a row the forest + gave an employer contribution takes the FRS twin's rate times its own pay + where the twin has one (``uk_runtime.pension_coherence``). + """ + + rows = spi_people.to_numpy(dtype=bool) + pay = person.loc[spi_people, "employment_income"].to_numpy(dtype=np.float64) + weights = _person_household_weights(person, household).to_numpy(dtype=float)[rows] + sources = ( + person["person_household_id"] + .map(household.set_index("household_id")["household_source_id"]) + .to_numpy()[rows] + ) + employee_before = person.loc[spi_people, "employee_pension_contributions"].to_numpy( + dtype=np.float64 + ) + employee, capped = cap_net_pay_contributions(employee_before, pay) + employer_before = person.loc[spi_people, "employer_pension_contributions"].to_numpy( + dtype=np.float64 + ) + employer, from_twin, bounded = twin_rate_employer_contributions( + drawn=employer_before, + pay=pay, + twin_pay=twin_pay, + twin_employer=twin_employer, + ) + person.loc[spi_people, "employee_pension_contributions"] = employee + person.loc[spi_people, "employer_pension_contributions"] = employer + receipt = { + "employee_capped_at_pay": coherence_receipt( + changed=capped, + before=employee_before, + after=employee, + weights=weights, + source_ids=sources, + ), + "employer_at_twin_rate": coherence_receipt( + changed=from_twin, + before=employer_before, + after=employer, + weights=weights, + source_ids=sources, + ), + "employer_bounded_by_ashe_ceiling": coherence_receipt( + changed=bounded, + before=employer_before, + after=employer, + weights=weights, + source_ids=sources, + ), + } + return person, receipt + + def _person_household_weights( person: pd.DataFrame, household: pd.DataFrame, diff --git a/packages/microcosm-build/src/microcosm/build/uk_runtime/spi_spine.py b/packages/microcosm-build/src/microcosm/build/uk_runtime/spi_spine.py index befaa16588..606e307672 100644 --- a/packages/microcosm-build/src/microcosm/build/uk_runtime/spi_spine.py +++ b/packages/microcosm-build/src/microcosm/build/uk_runtime/spi_spine.py @@ -389,6 +389,7 @@ def evidence(self) -> dict[str, object]: "gainful_self_employment_refresh": ( self.imputation.gainful_self_employment_refresh ), + "pension_coherence": self.imputation.pension_coherence, "state_pension_age_guard": list(self.imputation.state_pension_age_guard), "targets": { "count": len(self.source_targets.targets), @@ -1032,6 +1033,7 @@ def _build_spine_replay_report( "bridge_validation": imputation.bridge_validation, "recipient_domain": imputation.recipient_domain, "gainful_self_employment_refresh": imputation.gainful_self_employment_refresh, + "pension_coherence": imputation.pension_coherence, "post_draw_identity": { "formula": SPI_SOURCE_TI_FORMULA, "rows_checked": identity_rows, diff --git a/packages/microcosm-build/src/microcosm/build/uk_runtime/uc_reporter_redraw.py b/packages/microcosm-build/src/microcosm/build/uk_runtime/uc_reporter_redraw.py index 4cd18ec7e9..50bad156e6 100644 --- a/packages/microcosm-build/src/microcosm/build/uk_runtime/uc_reporter_redraw.py +++ b/packages/microcosm-build/src/microcosm/build/uk_runtime/uc_reporter_redraw.py @@ -17,6 +17,10 @@ import pandas as pd from microcosm.build.source_manifest import SourceStageSpec +from microcosm.build.stochastic_assignment import ( + integer_record_keys, + stable_identity_uniforms, +) from microcosm.build.uk_runtime.national_frame import ( uk_household_weight_kind, uk_national_frame, @@ -314,7 +318,25 @@ def redraw_spi_reported_uc( weights=weights[training], ) draws = np.zeros(len(benunit), dtype=np.float64) - predicted = fitted.predict(encoded_target) + # Uniforms keyed on each benefit unit's own id, never a positional + # stream: a unit's redraw moves only when its own predictors do. + keys = integer_record_keys( + benunit["benunit_id"].to_numpy()[target], label="UC reporter benunit_id" + ) + salt = f"uc_reporter_redraw:{UC_REPORTER_TARGET}" + predicted = fitted.predict_from_uniforms( + encoded_target, + quantiles={ + UC_REPORTER_TARGET: stable_identity_uniforms( + keys, seed=UC_REPORTER_REDRAW_SEED, salt=salt + ) + }, + sign_uniforms={ + UC_REPORTER_TARGET: stable_identity_uniforms( + keys, seed=UC_REPORTER_REDRAW_SEED, salt=f"{salt}:sign" + ) + }, + ) if list(predicted.columns) != [UC_REPORTER_TARGET]: raise ValueError( "UC reporter redraw model must return exactly the declared target." diff --git a/packages/microcosm-build/src/microcosm/build/uk_runtime/was_wealth.py b/packages/microcosm-build/src/microcosm/build/uk_runtime/was_wealth.py index 7fe9dbf68d..d1318749f0 100644 --- a/packages/microcosm-build/src/microcosm/build/uk_runtime/was_wealth.py +++ b/packages/microcosm-build/src/microcosm/build/uk_runtime/was_wealth.py @@ -13,6 +13,10 @@ from microcosm.build.gates import FitWeightRecord from microcosm.build.source_manifest import SourceStageSpec +from microcosm.build.stochastic_assignment import ( + integer_record_keys, + stable_identity_uniforms, +) from microcosm.build.uk_runtime.frs_brma import _benunit_household_map from microcosm.build.uk_runtime.frs_spine import read_pinned_tab from microcosm.build.uk_runtime.national_frame import ( @@ -369,6 +373,7 @@ def __call__(self, frame: Frame) -> Frame: household_predictors, seed=UK_WAS_WEALTH_DECLARED_SEEDS["was_wealth"], n_estimators=_qrf_n_estimators(self.stage), + record_ids=frame.table("household")["household_id"].to_numpy(), ) self.last_fit_weight_records = imputation.fit_weight_records capped, capped_rows = cap_derived_totals_to_donor_range(imputation.draws, donor) @@ -735,6 +740,7 @@ def impute_was_wealth( *, seed: int, n_estimators: int, + record_ids: Sequence[object] | np.ndarray, segments: int = UK_WAS_CHAIN_SEGMENTS, ) -> UKWASWealthImputationResult: """Fit segmented checkpointed QRF chains and draw WAS wealth outputs. @@ -746,6 +752,14 @@ def impute_was_wealth( each is computed from its drawn components, so the accounting identities the donor satisfies hold on every recipient. + Every draw reads uniforms keyed on the recipient's own identity + (``record_ids``, the household ids in recipient row order) and the target, + never a positional stream: a household's draws move only when its own + predictors do, whatever happens to the other recipients. A positional + stream shifted every later household's draws whenever an engine change + moved one household's net income and, through the high-value copies, the + frame's length. + ``segments`` runs only the first *k* chain segments (a test seam: each segment draws from its own child seed, so the earlier segments' draws are identical whether or not the later ones run). @@ -762,6 +776,14 @@ def impute_was_wealth( raise KeyError( f"WAS {side} table is missing {UK_WAS_TENURE_CATEGORY_COLUMN!r}." ) + ids = pd.Series( + integer_record_keys(record_ids, label="WAS record_ids"), + index=recipient_predictor_frame.index, + ) + if len(ids) != len(recipient_predictor_frame) or ids.duplicated().any(): + raise ValueError( + "WAS record_ids must name every recipient once, in recipient row order." + ) donor_category = donor[UK_WAS_TENURE_CATEGORY_COLUMN].astype(str).to_numpy() recipient_category = ( recipient_predictor_frame[UK_WAS_TENURE_CATEGORY_COLUMN].astype(str).to_numpy() @@ -820,12 +842,18 @@ def run_segment( ) ) continue + keys = ids.loc[recipient_rows.index].to_numpy() + salt = f"{UK_WAS_WEALTH_FIT_NAME}:{target}" result = model.fit_draw_next( donor_rows, recipient_base, segment_raw, state=state, weights="weight", + quantiles=stable_identity_uniforms(keys, seed=seed, salt=salt), + sign_uniforms=stable_identity_uniforms( + keys, seed=seed, salt=f"{salt}:sign" + ), ) fit_records.append( FitWeightRecord( diff --git a/packages/microcosm-build/tests/engine/uk/test_uk_spi_income.py b/packages/microcosm-build/tests/engine/uk/test_uk_spi_income.py index 8344260548..9008564d2c 100644 --- a/packages/microcosm-build/tests/engine/uk/test_uk_spi_income.py +++ b/packages/microcosm-build/tests/engine/uk/test_uk_spi_income.py @@ -105,12 +105,13 @@ def test_spi_rebases_income_before_frs_fill_and_preserves_child_inputs( def test_child_exclusion_keeps_the_other_recipients_draws( monkeypatch, tmp_path ) -> None: - """Dropping a recipient leaves every other SPI row's stage-1 draw unchanged. + """Dropping a recipient leaves every other SPI row's draws unchanged. - The forest is queried on the drawn recipients alone and every draw is - keyed on the recipient's identity (its twin's id), so the draw a recipient - receives does not depend on who else is drawn or in what order - (microcosm#840). FRS rows are never redrawn. + The stage-1 forest is queried on the drawn recipients alone and keys each + draw on the recipient's identity (its twin's id); the stage-2 forest keys + each draw on the row's own person_id. So the draws a recipient receives do + not depend on who else is drawn or in what order (microcosm#840, + microcosm#1095). FRS rows are never redrawn. """ support = _dead_support() @@ -122,9 +123,11 @@ def predict_from_uniforms(self, predictors, *, quantiles, sign_uniforms): result = original_draw( self, predictors, quantiles=quantiles, sign_uniforms=sign_uniforms ) - if "dividend_income" in result: - # A draw that reads its uniform: the identity-keyed quantile. - result["dividend_income"] = 100.0 * quantiles["dividend_income"] + # Draws that read their uniforms: the stage-1 dividend and the + # stage-2 UC report take their keyed quantiles. + for target in ("dividend_income", "universal_credit_reported"): + if target in result: + result[target] = 100.0 * quantiles[target] return result monkeypatch.setattr(_FakeFittedQRF, "predict_from_uniforms", predict_from_uniforms) @@ -147,11 +150,12 @@ def predict_from_uniforms(self, predictors, *, quantiles, sign_uniforms): replace(support, person=person), donor_path, **options ) others = spi_rows.drop(child) - np.testing.assert_array_equal( - reference.person.loc[others, "dividend_income"], - candidate.person.loc[others, "dividend_income"], - ) - assert reference.person.loc[others, "dividend_income"].nunique() == len(others) + for column in ("dividend_income", "universal_credit_reported"): + np.testing.assert_array_equal( + reference.person.loc[others, column], + candidate.person.loc[others, column], + ) + assert reference.person.loc[others, column].nunique() == len(others) base = person[channel].eq("frs") np.testing.assert_array_equal( candidate.person.loc[base, "dividend_income"], @@ -160,10 +164,11 @@ def predict_from_uniforms(self, predictors, *, quantiles, sign_uniforms): # Reversing the recipients' order changes no draw either. reversed_support = replace(support, person=support.person.iloc[::-1].copy()) reordered = impute_uk_spi_income_support(reversed_support, donor_path, **options) - pd.testing.assert_series_equal( - reordered.person.loc[spi_rows, "dividend_income"], - reference.person.loc[spi_rows, "dividend_income"], - ) + for column in ("dividend_income", "universal_credit_reported"): + pd.testing.assert_series_equal( + reordered.person.loc[spi_rows, column], + reference.person.loc[spi_rows, column], + ) def test_dependants_keep_their_twin_values_and_recipients_their_draws( @@ -892,13 +897,19 @@ def test_stage1_queries_only_drawn_recipients_with_the_bridge( support = replace(support, person=person) donor_path = tmp_path / SPI_DONOR_FILENAME _write_donor(donor_path) - captured: dict[str, object] = {} + captured: dict[str, dict[str, object]] = {} original_draw = _FakeFittedQRF.predict_from_uniforms def predict_from_uniforms(self, predictors, *, quantiles, sign_uniforms): - captured["columns"] = list(predictors.columns) - captured["rows"] = len(predictors) - captured["sign"] = sign_uniforms["gift_aid"].copy() + # Both stages draw from keyed uniforms; stage 1 draws the tape's + # targets, gift aid among them. + stage = "stage1" if "gift_aid" in sign_uniforms else "stage2" + captured[stage] = { + "columns": list(predictors.columns), + "rows": len(predictors), + "quantiles": {k: np.array(v, dtype=float) for k, v in quantiles.items()}, + "sign": {k: np.array(v, dtype=float) for k, v in sign_uniforms.items()}, + } return original_draw( self, predictors, quantiles=quantiles, sign_uniforms=sign_uniforms ) @@ -915,24 +926,41 @@ def predict_from_uniforms(self, predictors, *, quantiles, sign_uniforms): state_pension_age=66, allow_in_sample_reference=True, ) - assert captured["rows"] == result.spi_prediction_rows == 3 + stage1 = captured["stage1"] + assert stage1["rows"] == result.spi_prediction_rows == 3 assert {"has_pay", "has_trade", "at_or_over_spa", "main_source_pay", "age"} <= set( - captured["columns"] + stage1["columns"] ) # The no-analogue leaf's sign uniform is the identity-keyed one for the # drawn recipients, in their order. + age_domain = (result.person[channel] == "spi") & result.person["age"].ge(16) drawn_ids = result.person.loc[ - (result.person[channel] == "spi") - & result.person["age"].ge(16) - & ~result.person.index.isin(copy), - "person_source_id", + age_domain & ~result.person.index.isin(copy), "person_source_id" ].to_numpy() + from microcosm.build.stochastic_assignment import stable_identity_uniforms + from microcosm.build.uk_runtime.spi_income import FRS_ONLY_FIT_NAME from microcosm.build.uk_runtime.spi_rank_transport import identity_uniforms np.testing.assert_array_equal( - captured["sign"], + stage1["sign"]["gift_aid"], identity_uniforms(drawn_ids, target="gift_aid", part="sign", seed=9), ) + # Stage 2 is queried over the whole age domain, and each row draws from + # uniforms keyed on its own person_id under the stage's own seed, never + # from a stream its row position decides, and never shared with another + # copy of the same FRS person (microcosm#1095). + domain_ids = result.person.loc[age_domain, "person_id"].to_numpy() + stage2 = captured["stage2"] + assert stage2["rows"] == len(domain_ids) + for target in stage2["quantiles"]: + for part, salt in ( + ("quantiles", f"{FRS_ONLY_FIT_NAME}:{target}"), + ("sign", f"{FRS_ONLY_FIT_NAME}:{target}:sign"), + ): + np.testing.assert_array_equal( + stage2[part][target], + stable_identity_uniforms(domain_ids, seed=10, salt=salt), + ) assert result.stage1_transport["identity_key"] == "person_source_id" assert result.stage1_transport["recipients"] == 3 assert result.fit_weight_records[1].fit_name == "uk_frs_reference_conditional" diff --git a/packages/microcosm-build/tests/engine_free/shared/test_spec_engine_loader.py b/packages/microcosm-build/tests/engine_free/shared/test_spec_engine_loader.py index c43fbd6544..2d4a494760 100644 --- a/packages/microcosm-build/tests/engine_free/shared/test_spec_engine_loader.py +++ b/packages/microcosm-build/tests/engine_free/shared/test_spec_engine_loader.py @@ -13,9 +13,10 @@ def test_semantic_hash_has_golden_vector_and_surface_separation(tmp_path) -> Non # exact normative projection as one reviewable golden vector. The seed # protocol inside it attests microcosm.calibrate.solve's bytes, so # microcosm#1124's opt-in retain_best_objective and microcosm#1301's - # budget_search_bracket move it. + # budget_search_bracket move it, and the fit kernels' bytes, so the QRF + # chain step's caller-owned uniforms move it too. assert first.spec_sha256 == ( - "58a60d16e514df91eb7b5eacb812906594e95f9ca2f76fd5ce4db7ffc008e160" + "4f71334334c67cc17d1aab5b02a838def57bcc54ac4413b89b70b8ae3086d9ee" ) second_root = _rich_minimal(tmp_path / "xy", note="second", store="local:b") diff --git a/packages/microcosm-build/tests/engine_free/shared/test_stochastic_assignment.py b/packages/microcosm-build/tests/engine_free/shared/test_stochastic_assignment.py index ede0c0a063..30f41616a7 100644 --- a/packages/microcosm-build/tests/engine_free/shared/test_stochastic_assignment.py +++ b/packages/microcosm-build/tests/engine_free/shared/test_stochastic_assignment.py @@ -7,6 +7,7 @@ assign_binary_from_rate, assign_binary_with_anchored_residual, clipped_normal_from_uniforms, + integer_record_keys, sample_categorical_from_counts, stable_identity_uniforms, ) @@ -160,3 +161,27 @@ def test_anchored_residual_population_without_anchor_and_shape_refusal() -> None assign_binary_with_anchored_residual( draws, 0.4, np.zeros(10, dtype=bool), population=np.ones(3, dtype=bool) ) + + +def test_integer_record_keys_pass_unique_integer_ids_through() -> None: + ids = np.array([3, 1, 2], dtype="int64") + + np.testing.assert_array_equal(integer_record_keys(ids, label="ids"), ids) + + +def test_a_float_id_keys_a_different_draw_so_it_is_refused() -> None: + # The key string formats each id, so 1001 and 1001.0 draw differently: + # a merge that turned an id column into floats would redraw every record. + assert not np.array_equal( + stable_identity_uniforms([1001], seed=0, salt="s"), + stable_identity_uniforms([1001.0], seed=0, salt="s"), + ) + with pytest.raises(TypeError, match="integer ids"): + integer_record_keys(np.array([1001.0, 1002.0]), label="ids") + with pytest.raises(TypeError, match="integer ids"): + integer_record_keys(np.array(["1001"], dtype=object), label="ids") + + +def test_a_repeated_id_is_refused() -> None: + with pytest.raises(ValueError, match="each record once"): + integer_record_keys(np.array([1, 2, 1]), label="ids") diff --git a/packages/microcosm-build/tests/engine_free/uk/test_uk_etb_services.py b/packages/microcosm-build/tests/engine_free/uk/test_uk_etb_services.py index 7f72136d28..d62ce38521 100644 --- a/packages/microcosm-build/tests/engine_free/uk/test_uk_etb_services.py +++ b/packages/microcosm-build/tests/engine_free/uk/test_uk_etb_services.py @@ -103,7 +103,20 @@ def start_chain(self, donor, predictors, targets, *, weights): ] return {"targets": targets} - def fit_draw_next(self, donor, recipient_base, raw, *, state, weights): + def fit_draw_next( + self, + donor, + recipient_base, + raw, + *, + state, + weights, + quantiles, + sign_uniforms, + ): + # Each step draws from uniforms keyed on the recipients' own ids. + for uniforms in (quantiles, sign_uniforms): + assert uniforms.shape == (len(recipient_base),) return type( "Result", (), @@ -130,7 +143,9 @@ def fit_draw_next(self, donor, recipient_base, raw, *, state, weights): ] ) - draws, records = impute_etb_services(donor, recipient, seed=0) + draws, records = impute_etb_services( + donor, recipient, seed=0, record_ids=np.arange(1, len(recipient) + 1) + ) assert draws.columns.tolist() == [ "dfe_education_spending", @@ -415,3 +430,35 @@ def test_stage_transform_records_its_mass_conservation_receipt( assert receipt.reason == module.UK_ETB_SERVICES_MASS_CONSERVATION_REASON assert receipt.old_total == receipt.new_total == 40.0 assert receipt.declared_factor == 1.0 + + +def test_etb_draws_move_only_for_the_households_whose_inputs_move() -> None: + """Keyed on household ids: reordering, one changed and one added move no other.""" + + from test_support.microcosm_build.uk_keyed_draws import ( + assert_draws_are_keyed_on_ids, + ) + + donor = clean_etb_services_table(_raw_etb()) + recipient = donor.iloc[:8].drop( + columns=[ + "weight", + "dfe_education_spending", + "rail_subsidy_spending", + "bus_subsidy_spending", + ] + ) + + def draw(rows: pd.DataFrame, ids: np.ndarray) -> pd.DataFrame: + draws, _ = impute_etb_services( + donor, rows, seed=0, record_ids=ids, n_estimators=5 + ) + return draws.reset_index(drop=True) + + assert_draws_are_keyed_on_ids( + draw, + recipient, + np.arange(1, len(recipient) + 1), + perturbed_column="hbai_household_net_income", + perturbed_value=54_321.0, + ) diff --git a/packages/microcosm-build/tests/engine_free/uk/test_uk_frs_spine.py b/packages/microcosm-build/tests/engine_free/uk/test_uk_frs_spine.py index e017b349e7..c2b1d1db56 100644 --- a/packages/microcosm-build/tests/engine_free/uk/test_uk_frs_spine.py +++ b/packages/microcosm-build/tests/engine_free/uk/test_uk_frs_spine.py @@ -1298,6 +1298,51 @@ def test_no_frs_field_feeds_two_income_columns(tmp_path: Path) -> None: assert moved[("pension", "PENPAY")] == {"private_pension_income"} +@pytest.mark.parametrize( + ("inearns", "deduction", "capped_at", "case"), + [ + (10.0, 15.0, 10.0, "above_recorded_pay"), + (0.0, 2.0, 0.0, "beside_no_recorded_pay"), + ], + ids=["above_pay", "no_pay"], +) +def test_a_net_pay_deduction_above_pay_is_capped_at_it( + tmp_path: Path, inearns: float, deduction: float, capped_at: float, case: str +) -> None: + """The FRS records each job's pension deduction apart from the adult's + usual earnings, so a deduction can exceed them; the spine caps it at them + (uk_runtime.pension_coherence). A deduction beside no recorded pay may + stand for pay the adult tab left blank, so the receipt counts it apart + (Vahid's review of microcosm#1332).""" + + from microcosm.build.uk_runtime.receipt_counts import support_band + + tables = _fixture_tables() + adult = next( + row for row in tables["adult"] if row["SERNUM"] == 1 and row["PERSON"] == 1 + ) + adult["INEARNS"] = inearns + tables["job"][0]["DEDUC1"] = deduction + stage = _write_fixture(tmp_path, tables) + evidence: dict[str, object] = {} + person = ( + build_uk_frs_spine_frame(tmp_path, stage=stage, evidence=evidence) + .table("person") + .set_index("person_id") + ) + + assert person.loc[1001, "employment_income"] == pytest.approx( + inearns * WEEKS_IN_YEAR + ) + assert person.loc[1001, "employee_pension_contributions"] == pytest.approx( + capped_at * WEEKS_IN_YEAR + ) + receipt = evidence["net_pay_contributions"] + other = ({"above_recorded_pay", "beside_no_recorded_pay"} - {case}).pop() + assert receipt[case]["support"] == support_band(1) + assert receipt[other]["support"] == 0 + + @pytest.mark.parametrize("seed", range(8)) def test_property_income_matches_a_merge_of_the_raw_tabs( tmp_path: Path, seed: int diff --git a/packages/microcosm-build/tests/engine_free/uk/test_uk_identity_stability_receipts.py b/packages/microcosm-build/tests/engine_free/uk/test_uk_identity_stability_receipts.py index 340b03d596..9b876c8b67 100644 --- a/packages/microcosm-build/tests/engine_free/uk/test_uk_identity_stability_receipts.py +++ b/packages/microcosm-build/tests/engine_free/uk/test_uk_identity_stability_receipts.py @@ -184,8 +184,18 @@ def start_chain(self, donor, predictors, targets, *, weights): del donor, predictors, weights return {"targets": list(targets)} - def fit_draw_next(self, donor, recipient_base, raw, *, state, weights): - del recipient_base, raw, weights + def fit_draw_next( + self, + donor, + recipient_base, + raw, + *, + state, + weights, + quantiles, + sign_uniforms, + ): + del recipient_base, raw, weights, quantiles, sign_uniforms target = state["targets"][0] state = {"targets": state["targets"][1:]} return SimpleNamespace( diff --git a/packages/microcosm-build/tests/engine_free/uk/test_uk_lcfs_consumption.py b/packages/microcosm-build/tests/engine_free/uk/test_uk_lcfs_consumption.py index 447450a830..640ba6d250 100644 --- a/packages/microcosm-build/tests/engine_free/uk/test_uk_lcfs_consumption.py +++ b/packages/microcosm-build/tests/engine_free/uk/test_uk_lcfs_consumption.py @@ -16,6 +16,7 @@ assign_recipient_has_fuel, clean_lcfs_consumption_table, derive_energy_from_lcfs, + impute_lcfs_consumption, recipient_predictors, support_clip_to_donor, ) @@ -2923,3 +2924,27 @@ def table(rows: int) -> pd.DataFrame: household_ids=np.arange(60) + 1, energy=energy, ) + + +def test_lcfs_draws_move_only_for_the_households_whose_inputs_move() -> None: + """Keyed on household ids: reordering, one changed and one added move no other.""" + + from test_support.microcosm_build.uk_keyed_draws import ( + assert_draws_are_keyed_on_ids, + ) + + donor = clean_lcfs_consumption_table(*_synthetic_lcfs_donor()) + recipient = donor.iloc[:10].loc[:, list(UK_LCFS_CONSUMPTION_PREDICTORS)] + + def draw(rows: pd.DataFrame, ids: np.ndarray) -> pd.DataFrame: + return impute_lcfs_consumption( + donor, rows, seed=0, n_estimators=5, record_ids=ids + ).draws.reset_index(drop=True) + + assert_draws_are_keyed_on_ids( + draw, + recipient, + np.arange(1, len(recipient) + 1), + perturbed_column="hbai_household_net_income", + perturbed_value=777.0, + ) diff --git a/packages/microcosm-build/tests/engine_free/uk/test_uk_pension_coherence.py b/packages/microcosm-build/tests/engine_free/uk/test_uk_pension_coherence.py new file mode 100644 index 0000000000..29d2d8c33e --- /dev/null +++ b/packages/microcosm-build/tests/engine_free/uk/test_uk_pension_coherence.py @@ -0,0 +1,137 @@ +"""Pension contributions that cohere with the pay they come from (microcosm#1095).""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import pytest + +from microcosm.build.uk_runtime.pension_coherence import ( + ASHE_EMPLOYER_RATE_CEILING, + cap_net_pay_contributions, + coherence_receipt, + twin_rate_employer_contributions, +) +from microcosm.build.uk_runtime.receipt_counts import support_band +from microcosm.build.uk_runtime.spi_income import _spi_pension_coherence + + +def test_net_pay_contributions_are_capped_at_pay() -> None: + capped, changed = cap_net_pay_contributions( + np.array([500.0, 3_000.0, 200.0]), np.array([1_000.0, 2_000.0, -50.0]) + ) + + np.testing.assert_array_equal(capped, [500.0, 2_000.0, 0.0]) + np.testing.assert_array_equal(changed, [False, True, True]) + + +def test_the_net_pay_cap_refuses_misaligned_inputs() -> None: + with pytest.raises(ValueError, match="one value per person"): + cap_net_pay_contributions(np.zeros(2), np.zeros(3)) + + +def test_the_forest_decides_membership_and_the_twin_rate_scales_the_amount() -> None: + """A row the forest gave no employer contribution keeps none, whatever its + twin's record; a member whose twin has a rate takes that rate times its + own pay; any other member keeps its drawn amount, bounded by ASHE's top + rate (María's ruling of 2026-10-21, after Vahid's review of + microcosm#1332).""" + + employer, from_twin, bounded = twin_rate_employer_contributions( + drawn=np.array([0.0, 4_000.0, 9_000.0, 50_000.0, 1_000.0, 3_000.0]), + pay=np.array([100_000.0, 60_000.0, 80_000.0, 100_000.0, 0.0, 50_000.0]), + twin_pay=np.array([20_000.0, 30_000.0, 0.0, 0.0, 25_000.0, 40_000.0]), + twin_employer=np.array([1_000.0, 0.0, 0.0, 0.0, 2_500.0, 2_000.0]), + ) + + np.testing.assert_allclose( + employer, + [0.0, 4_000.0, 9_000.0, ASHE_EMPLOYER_RATE_CEILING * 100_000.0, 0.0, 2_500.0], + ) + np.testing.assert_array_equal(from_twin, [False, False, False, False, True, True]) + np.testing.assert_array_equal(bounded, [False, False, False, True, False, False]) + + +def test_a_receipt_withholds_amounts_resting_on_fewer_than_ten_households() -> None: + receipt = coherence_receipt( + changed=np.array([True, True, False, True]), + before=np.array([10.0, 20.0, 30.0, 40.0]), + after=np.array([5.0, 10.0, 30.0, 0.0]), + weights=np.ones(4), + source_ids=np.array([1, 1, 2, 3]), + ) + + assert receipt == { + "support": support_band(2), + "weighted_before": "<10", + "weighted_after": "<10", + } + + +def test_a_receipt_gives_weighted_amounts_on_ten_or_more_households() -> None: + rows = 12 + receipt = coherence_receipt( + changed=np.ones(rows, dtype=bool), + before=np.full(rows, 10.0), + after=np.full(rows, 4.0), + weights=np.full(rows, 2.0), + source_ids=np.arange(rows), + ) + + assert receipt["support"] == support_band(rows) + assert receipt["weighted_before"] == pytest.approx(240.0) + assert receipt["weighted_after"] == pytest.approx(96.0) + + +def test_the_spi_stage_applies_the_rule_to_its_recipients() -> None: + """The SPI stage caps each recipient's net-pay contribution at its drawn + pay, scales a member's employer contribution by its twin's rate, leaves a + row the forest gave none at none, and counts each rule apart.""" + + person = pd.DataFrame( + { + "person_id": [1, 2, 3, 4, 5], + "person_household_id": [10, 20, 30, 40, 50], + "employment_income": [100_000.0, 60_000.0, 50_000.0, 40_000.0, 30_000.0], + "employee_pension_contributions": [5_000.0, 70_000.0, 0.0, 1_000.0, 0.0], + "employer_pension_contributions": [ + 2_000.0, + 4_000.0, + 20_000.0, + 3_000.0, + 0.0, + ], + } + ) + household = pd.DataFrame( + { + "household_id": [10, 20, 30, 40, 50], + "household_weight": [1.0] * 5, + "household_source_id": [10, 20, 30, 40, 50], + } + ) + recipients = pd.Series([True, True, True, False, True]) + + out, receipt = _spi_pension_coherence( + person.copy(), + household, + spi_people=recipients, + twin_pay=np.array([20_000.0, 30_000.0, 0.0, 25_000.0]), + twin_employer=np.array([1_000.0, 0.0, 0.0, 2_500.0]), + ) + + np.testing.assert_allclose( + out["employer_pension_contributions"], + [5_000.0, 4_000.0, 15_000.0, 3_000.0, 0.0], + ) + np.testing.assert_allclose( + out["employee_pension_contributions"], [5_000.0, 60_000.0, 0.0, 1_000.0, 0.0] + ) + assert set(receipt) == { + "employee_capped_at_pay", + "employer_at_twin_rate", + "employer_bounded_by_ashe_ceiling", + } + for rule in receipt.values(): + assert rule["support"] == support_band(1) + assert rule["weighted_before"] == rule["weighted_after"] == "<10" diff --git a/packages/microcosm-build/tests/engine_free/uk/test_uk_salary_sacrifice.py b/packages/microcosm-build/tests/engine_free/uk/test_uk_salary_sacrifice.py index 0aa9556eff..fd784edbd5 100644 --- a/packages/microcosm-build/tests/engine_free/uk/test_uk_salary_sacrifice.py +++ b/packages/microcosm-build/tests/engine_free/uk/test_uk_salary_sacrifice.py @@ -31,7 +31,12 @@ class _Fitted: value = 0.0 - def predict(self, predictors: pd.DataFrame) -> pd.DataFrame: + def predict_from_uniforms( + self, predictors: pd.DataFrame, *, quantiles, sign_uniforms + ) -> pd.DataFrame: + # The stage keys each person's uniforms on their own id. + for uniforms in (quantiles[SALSAC_OUTPUT], sign_uniforms[SALSAC_OUTPUT]): + assert uniforms.shape == (len(predictors),) return pd.DataFrame({SALSAC_OUTPUT: self.value}, index=predictors.index) @@ -322,9 +327,10 @@ def test_conversion_moves_at_most_half_the_pay_and_keeps_the_rest( monkeypatch.setattr(salary_sacrifice, "QRF", _FakeQRF) n_donors = 100 pay = np.asarray([30_000.0, *([1_000.0] * n_donors)]) - # Contributions from 100 to 1,500: above half the pay from 500 up, and - # above the whole pay from 1,000 up. - contributions = np.asarray([0.0, *np.linspace(100.0, 1_500.0, n_donors)]) + # Contributions from 100 to 1,000: above half the pay from 500 up, and at + # most the whole pay, where the spine caps a net-pay contribution + # (uk_runtime.pension_coherence). + contributions = np.asarray([0.0, *np.linspace(100.0, 1_000.0, n_donors)]) frame = _frame( asked=[1, *([0] * n_donors)], salary_sacrifice_values=[0.0] * (n_donors + 1), @@ -398,8 +404,10 @@ def test_no_amount_exceeds_pay_and_nothing_is_lost(rows, predicted) -> None: rows = [(True, 0.0, 0.0, 30_000.0, "frs"), *rows] asked = [1 if row[0] else 0 for row in rows] reported = np.asarray([row[1] if row[0] else 0.0 for row in rows]) - contributions = np.asarray([row[2] for row in rows]) pay = np.asarray([row[3] for row in rows]) + # The spine caps net-pay contributions at pay before this stage + # (uk_runtime.pension_coherence), so the stage only ever sees them there. + contributions = np.minimum(np.asarray([row[2] for row in rows]), pay) frame = _frame( asked=asked, salary_sacrifice_values=reported, @@ -435,6 +443,95 @@ def test_no_amount_exceeds_pay_and_nothing_is_lost(rows, predicted) -> None: assert (sacrifice[converted] <= SALSAC_MOVE_PAY_SHARE_CAP * pay[converted]).all() +def test_a_contribution_above_the_pay_that_remains_is_refused( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The spine caps net-pay contributions at pay; the stage refuses one above it.""" + + monkeypatch.setattr(salary_sacrifice, "QRF", _FakeQRF) + frame = _frame( + asked=[1, 1], + salary_sacrifice_values=[0.0, 0.0], + employee_pension=[0.0, 3_000.0], + employment_income=[30_000.0, 1_000.0], + ) + + with pytest.raises(ValueError, match="must not exceed the pay that remains"): + impute_salary_sacrifice(frame) + + +def test_draws_move_only_for_the_people_whose_inputs_move( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Each prediction reads uniforms keyed on the person's own id: reversing + the people, changing one and adding another moves no other person's draw + (Vahid's review of microcosm#1332).""" + + from test_support.microcosm_build.uk_keyed_draws import ( + assert_draws_are_keyed_on_ids, + ) + + drawn: dict[str, pd.Series] = {} + + class _Keyed(_Fitted): + def predict_from_uniforms(self, predictors, *, quantiles, sign_uniforms): + drawn["quantile"] = pd.Series( + quantiles[SALSAC_OUTPUT], index=predictors.index + ) + return pd.DataFrame({SALSAC_OUTPUT: 0.0}, index=predictors.index) + + class _KeyedQRF(_FakeQRF): + def fit(self, frame, predictors, targets, *, weights): + super().fit(frame, predictors, targets, weights=weights) + return _Keyed() + + monkeypatch.setattr(salary_sacrifice, "QRF", _KeyedQRF) + + def draw(recipient: pd.DataFrame, ids: np.ndarray) -> pd.DataFrame: + # The person rows keep the recipients' order; the group tables are + # sorted by id, as a frame requires. + groups = np.sort(ids) + frame = uk_national_frame( + person=recipient.assign( + person_id=ids, person_benunit_id=ids, person_household_id=ids + ), + benunit=pd.DataFrame({"benunit_id": groups}), + household=pd.DataFrame( + { + "household_id": groups, + "household_support_channel": "frs", + HOUSEHOLD_IS_CGT_CLONE: False, + HOUSEHOLD_IS_CGT_SUPPORT_COPY: False, + } + ), + household_weights=np.ones(len(ids)), + time_period="2024", + ) + drawn.clear() + impute_salary_sacrifice(frame) + quantile = pd.Series(np.nan, index=frame.table("person").index) + quantile.loc[drawn["quantile"].index] = drawn["quantile"].to_numpy() + return pd.DataFrame({"quantile": quantile.to_numpy()}) + + rows = 6 + recipient = pd.DataFrame( + { + "age": np.linspace(25, 55, rows), + "employment_income": np.full(rows, 30_000.0), + "salary_sacrifice_asked": [1, 1, 0, 0, 0, 0], + SALSAC_OUTPUT: [10.0, 20.0, 0.0, 0.0, 0.0, 0.0], + "employee_pension_contributions": np.zeros(rows), + } + ) + assert_draws_are_keyed_on_ids( + draw, + recipient, + np.arange(101, 101 + rows, dtype="int64"), + perturbed_column="age", + perturbed_value=61.0, + ) + + def test_anchor_is_self_consistent() -> None: anchor = load_salary_sacrifice_anchor() assert anchor["hmrc_anchor"]["total_users"] * anchor["derived"][ diff --git a/packages/microcosm-build/tests/engine_free/uk/test_uk_uc_reporter_redraw.py b/packages/microcosm-build/tests/engine_free/uk/test_uk_uc_reporter_redraw.py index 66d7aa29de..73c3e66b3a 100644 --- a/packages/microcosm-build/tests/engine_free/uk/test_uk_uc_reporter_redraw.py +++ b/packages/microcosm-build/tests/engine_free/uk/test_uk_uc_reporter_redraw.py @@ -173,6 +173,32 @@ def test_twin_calls_are_byte_identical_and_emit_transition_receipt() -> None: } +def test_each_spi_unit_draws_from_uniforms_keyed_on_its_own_id() -> None: + """Each redrawn unit reads uniforms keyed on its own benefit unit id, never + a stream its row position decides, so a unit's draw cannot move when other + units are added, dropped or reordered (Vahid's review of microcosm#1332).""" + + from microcosm.build.stochastic_assignment import stable_identity_uniforms + from microcosm.build.uk_runtime.uc_reporter_redraw import UC_REPORTER_TARGET + + frame = _frame() + factory = _StubQRFFactory() + redraw_spi_reported_uc(frame, engine=_StubEngine(), qrf_factory=factory) + quantiles, signs = factory.record["predict_uniforms"] + benunit = frame.table("benunit") + units = benunit.loc[factory.record["predict_index"], "benunit_id"].to_numpy() + + assert set(units) == {201, 203} + salt = f"uc_reporter_redraw:{UC_REPORTER_TARGET}" + for drawn, part_salt in ((quantiles, salt), (signs, f"{salt}:sign")): + np.testing.assert_array_equal( + drawn[UC_REPORTER_TARGET], + stable_identity_uniforms( + units, seed=UC_REPORTER_REDRAW_SEED, salt=part_salt + ), + ) + + def test_seeded_regime_gated_qrf_twin_calls_are_byte_identical() -> None: first = redraw_spi_reported_uc(_frame(), engine=_StubEngine()) twin = redraw_spi_reported_uc(_frame(), engine=_StubEngine()) diff --git a/packages/microcosm-build/tests/engine_free/uk/test_uk_was_wealth.py b/packages/microcosm-build/tests/engine_free/uk/test_uk_was_wealth.py index 2f9d69f748..76b3686808 100644 --- a/packages/microcosm-build/tests/engine_free/uk/test_uk_was_wealth.py +++ b/packages/microcosm-build/tests/engine_free/uk/test_uk_was_wealth.py @@ -208,6 +208,12 @@ def _frame() -> object: ) +def _ids(recipient: pd.DataFrame) -> np.ndarray: + """Stable household ids for a synthetic recipient table, in row order.""" + + return np.arange(1, len(recipient) + 1) + + def test_was_donor_cleaning_arithmetic_and_exact_case_insensitive_columns() -> None: donor = clean_was_household_table(_raw_was()) @@ -611,8 +617,14 @@ def fit_draw_next( *, state, weights, + quantiles, + sign_uniforms, ): assert weights == "weight" + # Every step draws from uniforms keyed on the recipients' own ids. + for uniforms in (quantiles, sign_uniforms): + assert uniforms.shape == (len(recipient_predictors),) + assert ((uniforms >= 0) & (uniforms < 1)).all() target = state.targets[state.position] if state.position: previous = state.targets[state.position - 1] @@ -633,7 +645,9 @@ def fit_draw_next( donor = clean_was_household_table(_raw_was()) recipient = recipient_predictors(_frame(), _FakeEngine()) - result = module.impute_was_wealth(donor, recipient, seed=0, n_estimators=7) + result = module.impute_was_wealth( + donor, recipient, seed=0, n_estimators=7, record_ids=_ids(recipient) + ) assert result.draws.columns.tolist() == _DRAW_COLUMNS assert [targets for _, targets in calls] == [ @@ -709,7 +723,9 @@ def test_stratified_target_refuses_a_donor_without_the_stratum() -> None: recipient = recipient_predictors(_frame(), _FakeEngine()) with pytest.raises(ValueError, match="no household in the tenure stratum"): - impute_was_wealth(donor, recipient, seed=0, n_estimators=2) + impute_was_wealth( + donor, recipient, seed=0, n_estimators=2, record_ids=_ids(recipient) + ) def test_chain_declaration_drift_is_refused() -> None: @@ -865,11 +881,18 @@ def test_later_segments_leave_the_earlier_draws_byte_equal() -> None: donor = _fixture_donor() recipient = _synthetic_recipients(48, seed=685) - full = impute_was_wealth(donor, recipient, seed=0, n_estimators=2) + full = impute_was_wealth( + donor, recipient, seed=0, n_estimators=2, record_ids=_ids(recipient) + ) assert list(full.draws.columns) == _DRAW_COLUMNS for segments in range(1, 8): partial = impute_was_wealth( - donor, recipient, seed=0, n_estimators=2, segments=segments + donor, + recipient, + seed=0, + n_estimators=2, + record_ids=_ids(recipient), + segments=segments, ) pd.testing.assert_frame_equal( partial.draws, full.draws.loc[:, list(partial.draws.columns)] @@ -885,7 +908,9 @@ def test_real_chain_holds_the_tenure_rules_and_the_identities() -> None: recipient = _synthetic_recipients() category = recipient[UK_WAS_TENURE_CATEGORY_COLUMN].to_numpy() - draws = impute_was_wealth(donor, recipient, seed=0, n_estimators=3).draws + draws = impute_was_wealth( + donor, recipient, seed=0, n_estimators=3, record_ids=_ids(recipient) + ).draws owner = np.isin(category, ["owned_outright", "owned_with_mortgage"]) mortgaged = category == "owned_with_mortgage" @@ -968,3 +993,27 @@ def test_total_cap_takes_the_excess_from_the_remainder_first() -> None: assert capped.loc[0, "non_residential_property_value"] == 3_090.0 assert capped.loc[0, "main_residence_value"] == 100_000.0 assert set(wealth_identity_violations(capped).values()) == {0} + + +def test_was_draws_move_only_for_the_households_whose_inputs_move() -> None: + """Keyed on household ids: reordering, one changed and one added move no other.""" + + from test_support.microcosm_build.uk_keyed_draws import ( + assert_draws_are_keyed_on_ids, + ) + + donor = clean_was_household_table(_raw_was()) + recipient = recipient_predictors(_frame(), _FakeEngine()) + + def draw(rows: pd.DataFrame, ids: np.ndarray) -> pd.DataFrame: + return impute_was_wealth( + donor, rows, seed=0, n_estimators=3, record_ids=ids + ).draws.reset_index(drop=True) + + assert_draws_are_keyed_on_ids( + draw, + recipient, + _ids(recipient), + perturbed_column="household_net_income", + perturbed_value=12_345.0, + ) diff --git a/packages/microcosm-build/tests/engine_workflow/us/test_us_multispine_pool_tool.py b/packages/microcosm-build/tests/engine_workflow/us/test_us_multispine_pool_tool.py index bbeac4342f..ed52466467 100644 --- a/packages/microcosm-build/tests/engine_workflow/us/test_us_multispine_pool_tool.py +++ b/packages/microcosm-build/tests/engine_workflow/us/test_us_multispine_pool_tool.py @@ -341,8 +341,9 @@ def capture_equality(expected: object, actual: object) -> None: # The US bundle's semantic hash covers microcosm.calibrate.solve's # bytes (its seed protocol); microcosm#1124's opt-in # retain_best_objective and microcosm#1301's budget_search_bracket - # move it. - "spec_sha256": "740d1ff5b2e9fa6da3154c44f8676f80109d404e6aa7a16c48ed8db83055b15d", + # move it. It covers microcosm.fit.qrf's attested bytes too, so + # fit_draw_next's caller-owned uniforms (microcosm#1332) move it. + "spec_sha256": "19460fc168a6f9ed44e580d0e36777cce78bb33cb52b75630ee86203fa795684", }, } diff --git a/packages/microcosm-fit/src/microcosm/fit/qrf.py b/packages/microcosm-fit/src/microcosm/fit/qrf.py index f49d866ccd..01446085e6 100644 --- a/packages/microcosm-fit/src/microcosm/fit/qrf.py +++ b/packages/microcosm-fit/src/microcosm/fit/qrf.py @@ -1424,6 +1424,8 @@ def fit_draw_next( *, state: QRFChainState, weights: WeightSpec = DESIGN_WEIGHTS, + quantiles: np.ndarray | None = None, + sign_uniforms: np.ndarray | None = None, ) -> QRFChainStepResult: """Fit and draw exactly the next target in a checkpointed chain. @@ -1436,6 +1438,14 @@ def fit_draw_next( model-config drift, target-prefix drift, donor/recipient row reordering, weight drift, an entity/front-door change, or non-float64 raw priors. + ``quantiles`` and ``sign_uniforms`` (both or neither) are caller-owned + uniforms, one in ``[0, 1)`` per recipient row. With them the step draws + as :meth:`FittedRegimeGatedQRF.predict_from_uniforms` does and leaves + the chain's draw stream where it was, so a row's draw depends on its + own uniforms and features only: keyed on stable record identities, + adding, dropping or reordering other recipients moves no row's draw. + Without them the step consumes the chain's positional draw stream. + Returns: A result whose ``state`` is the state *after* this target and must be checkpointed for the next subprocess. @@ -1478,10 +1488,27 @@ def fit_draw_next( augmented[prior] = raw_prior_draws[prior].to_numpy( dtype=np.float64, copy=False ) - raw_draw = np.asarray( - _draw_target_with_rng(augmented, target_model, draw_rng), - dtype=np.float64, - ) + if (quantiles is None) != (sign_uniforms is None): + raise ValueError("Supply both quantiles and sign_uniforms, or neither.") + if quantiles is None: + drawn = _draw_target_with_rng(augmented, target_model, draw_rng) + else: + supplied = { + name: _validate_uniform_mapping( + name, {target: values}, [target], len(augmented) + )[target] + for name, values in ( + ("quantiles", quantiles), + ("sign_uniforms", sign_uniforms), + ) + } + drawn = _draw_target_from_uniforms( + augmented, + target_model, + supplied["quantiles"], + supplied["sign_uniforms"], + ) + raw_draw = np.asarray(drawn, dtype=np.float64) raw_draw.setflags(write=False) advanced = replace( state, diff --git a/packages/microcosm-fit/tests/engine_free/shared/test_qrf_stateless.py b/packages/microcosm-fit/tests/engine_free/shared/test_qrf_stateless.py index c436a67496..132d4ac99c 100644 --- a/packages/microcosm-fit/tests/engine_free/shared/test_qrf_stateless.py +++ b/packages/microcosm-fit/tests/engine_free/shared/test_qrf_stateless.py @@ -303,6 +303,64 @@ def test_chain_step_exposes_a_fitted_one_target_view(): np.testing.assert_array_equal(step_two.raw_draw, step_two_again.raw_draw) +def test_chain_step_with_caller_uniforms_draws_each_row_on_its_own(): + """Keyed uniforms make a chain row's draw independent of the other rows. + + Adding a recipient, changing another recipient's predictor and reordering + the rest leave every untouched row's draws where they were, and a step + with caller uniforms leaves the chain's positional draw stream unconsumed. + """ + from microcosm.fit.qrf import RegimeGatedQRF + + x = np.tile(np.arange(30, dtype=float), 6) + donor = pd.DataFrame( + { + "x": x, + "first": x + 1 + np.tile([0.0, 0.5, 1.0], len(x) // 3), + "inflated": np.tile([0.0, 4.0, 5.0], len(x) // 3), + } + ) + qrf = RegimeGatedQRF(n_estimators=4, seed=7) + + def draw(recipient: pd.DataFrame, ids: np.ndarray) -> pd.DataFrame: + state = qrf.start_chain(donor, ["x"], ["first", "inflated"], weights="none") + raw = pd.DataFrame(index=recipient.index) + for salt in (0.137, 0.611): + step = qrf.fit_draw_next( + donor, + recipient, + raw, + state=state, + weights="none", + quantiles=(ids * salt) % 1.0, + sign_uniforms=(ids * (salt + 0.29)) % 1.0, + ) + assert step.state.draw_rng_state_json == state.draw_rng_state_json + raw[step.target] = step.raw_draw + state = step.state + return raw.set_axis(ids) + + ids = np.arange(1.0, 13.0) + base = draw(pd.DataFrame({"x": np.arange(12, dtype=float)}), ids) + + changed_ids = np.concatenate([[99.0], ids[::-1]]) + changed_x = np.concatenate([[3.0], np.arange(12, dtype=float)[::-1]]) + changed_x[1] = 29.0 # the row with id 12 gets a new predictor + changed = draw(pd.DataFrame({"x": changed_x}), changed_ids) + + untouched = ids[:-1] + pd.testing.assert_frame_equal(changed.loc[untouched], base.loc[untouched]) + with pytest.raises(ValueError, match="both quantiles and sign_uniforms"): + qrf.fit_draw_next( + donor, + pd.DataFrame({"x": [1.0]}), + pd.DataFrame(index=[0]), + state=qrf.start_chain(donor, ["x"], ["first"], weights="none"), + weights="none", + quantiles=np.array([0.5]), + ) + + def _placements(model, n, seed): rng = np.random.default_rng(seed) return { diff --git a/packages/microcosm-graph/tests/fixtures/parity/kernels/fit.qrf/pins.json b/packages/microcosm-graph/tests/fixtures/parity/kernels/fit.qrf/pins.json index 043e8f4660..43f777a08a 100644 --- a/packages/microcosm-graph/tests/fixtures/parity/kernels/fit.qrf/pins.json +++ b/packages/microcosm-graph/tests/fixtures/parity/kernels/fit.qrf/pins.json @@ -1 +1 @@ -{"dependencies":{"numpy":"2.4.6","pandas":"3.0.3","quantile-forest":"1.4.2","scikit-learn":"1.8.0"},"implementation_hash":"fc75df4cc7b44f9277714661ffbc79728a30fb89957da52cfd6d82b1dcd0d3fa","kernel":"fit.qrf@1","node":"fit_qrf","node_key":"d76005a7a0767f62cd55f2d9a941d8ace68fa53448650e0f4ee97040960a6b22","numeric":"platform_bitwise","platform":"arm64/darwin/py3.14","platforms":{"arm64/darwin/py3.14":{"direct":"direct.csv","node_key":"d76005a7a0767f62cd55f2d9a941d8ace68fa53448650e0f4ee97040960a6b22"},"x86_64/linux/py3.13":{"direct":"platforms/x86_64-linux-py3_13/direct.csv","node_key":"eea71eb4437c9a1db975f22b7fa878c06746c5746c8497f2217ed5132a52646c"},"x86_64/linux/py3.14":{"direct":"platforms/x86_64-linux-py3_14/direct.csv","node_key":"2be10b1d3ed415552985c1581eb673070bccffc340778119bef70f9192039516"}},"seed":947} +{"dependencies":{"numpy":"2.4.6","pandas":"3.0.3","quantile-forest":"1.4.2","scikit-learn":"1.8.0"},"implementation_hash":"5eb07d7c48a24c7b11426d67879b6dcca0af8e4de2835228a4eab9b5f9ba2f72","kernel":"fit.qrf@1","node":"fit_qrf","node_key":"e5383f06ebd1ee0573233c4573ab4fd463030b012674e6ddc8b46bc759b5df00","numeric":"platform_bitwise","platform":"arm64/darwin/py3.14","platforms":{"arm64/darwin/py3.14":{"direct":"direct.csv","node_key":"e5383f06ebd1ee0573233c4573ab4fd463030b012674e6ddc8b46bc759b5df00"},"x86_64/linux/py3.13":{"direct":"platforms/x86_64-linux-py3_13/direct.csv","node_key":"c56ada7dc523c155af5ae3cbc8b9619a79356b27632ab0d338e07023877c771f"},"x86_64/linux/py3.14":{"direct":"platforms/x86_64-linux-py3_14/direct.csv","node_key":"2c78582881c8a45d9112cc72899c41d7d1254a46e834f4f1edc2a67c8e1d086f"}},"seed":947} diff --git a/packages/microcosm-graph/tests/fixtures/parity/uk_spine/sources/cgt_asset_type_facts.json b/packages/microcosm-graph/tests/fixtures/parity/uk_spine/sources/cgt_asset_type_facts.json index 79100c4e9a..6ec576e485 100644 --- a/packages/microcosm-graph/tests/fixtures/parity/uk_spine/sources/cgt_asset_type_facts.json +++ b/packages/microcosm-graph/tests/fixtures/parity/uk_spine/sources/cgt_asset_type_facts.json @@ -2,70 +2,70 @@ "resource": "synthetic-cgt-asset-type.json", "resource_sha256": "synthetic", "source_commit": "synthetic", - "table4_all_gains": 6341728.91006761, + "table4_all_gains": 6509588.154029989, "table4_all_tax": 0.0, - "table4_all_taxpayers": 15.64025028989186, + "table4_all_taxpayers": 15.64584768489452, "table4_bands": [ { - "gains": 5614.755050771031, + "gains": 6044.921017395652, "lower_bound": 0, "tax": null, - "taxpayers": 1.4498610400415817, + "taxpayers": 1.3781098872785849, "upper_bound": 10000 }, { - "gains": 27083.345909709733, + "gains": 24922.967696084444, "lower_bound": 10000, "tax": null, - "taxpayers": 1.5144705014116808, + "taxpayers": 1.5208346263072319, "upper_bound": 25000 }, { - "gains": 76331.17897703486, + "gains": 79282.68292974937, "lower_bound": 25000, "tax": null, - "taxpayers": 2.1412992916855758, + "taxpayers": 2.1786247372384744, "upper_bound": 50000 }, { - "gains": 157905.0946424909, + "gains": 137743.4527659259, "lower_bound": 50000, "tax": null, - "taxpayers": 2.178039286606179, + "taxpayers": 1.97394515547429, "upper_bound": 100000 }, { - "gains": 155213.69790998782, + "gains": 151515.32471723968, "lower_bound": 100000, "tax": null, - "taxpayers": 0.9562896654462721, + "taxpayers": 0.9626944897911075, "upper_bound": 250000 }, { - "gains": 604360.2752446379, + "gains": 648737.6777786638, "lower_bound": 250000, "tax": null, - "taxpayers": 1.7593341347179783, + "taxpayers": 1.8012777557759854, "upper_bound": 500000 }, { - "gains": 754797.8046933388, + "gains": 840435.5717438484, "lower_bound": 500000, "tax": null, - "taxpayers": 1.0805336123429539, + "taxpayers": 1.2094554776477642, "upper_bound": 1000000 }, { - "gains": 4560422.757639639, + "gains": 4620905.555381081, "lower_bound": 1000000, "tax": null, - "taxpayers": 4.560422757639639, + "taxpayers": 4.620905555381081, "upper_bound": null } ], - "table4_individuals_gains": 6341728.91006761, + "table4_individuals_gains": 6509588.154029989, "table4_individuals_tax": 0.0, - "table4_individuals_taxpayers": 15.64025028989186, + "table4_individuals_taxpayers": 15.64584768489452, "table4_trusts_gains": 0.0, "table4_trusts_tax": 0.0, "table7_total_disposals": 2743095.2380952383, @@ -114,12 +114,12 @@ "proceeds": 9000000000.0 } ], - "table8a_disposals_total": 29.493043403796086, - "table8a_gains_total": 36452873.707876876, - "table8a_tax_total": 7290574.741575375, - "table8a_taxpayers_total": 26.811857639814622, - "table8b_all_gains": 36452873.707876876, - "table8b_all_taxpayers": 26.811857639814622, - "table8b_individuals_gains": 36452873.707876876, - "table8b_individuals_taxpayers": 26.811857639814622 + "table8a_disposals_total": 29.50359849151538, + "table8a_gains_total": 36565227.645255364, + "table8a_tax_total": 7313045.529051073, + "table8a_taxpayers_total": 26.82145317410489, + "table8b_all_gains": 36565227.645255364, + "table8b_all_taxpayers": 26.82145317410489, + "table8b_individuals_gains": 36565227.645255364, + "table8b_individuals_taxpayers": 26.82145317410489 } diff --git a/packages/microcosm-graph/tests/fixtures/parity/uk_spine/sources/property_purchase_targets.json b/packages/microcosm-graph/tests/fixtures/parity/uk_spine/sources/property_purchase_targets.json index d98cd250c7..626cc85c32 100644 --- a/packages/microcosm-graph/tests/fixtures/parity/uk_spine/sources/property_purchase_targets.json +++ b/packages/microcosm-graph/tests/fixtures/parity/uk_spine/sources/property_purchase_targets.json @@ -26,8 +26,8 @@ ], "minimum_price": 1000.0, "property_type": "additional_residential", - "transactions": 26.999194565555904, - "value": 204508.05790306703 + "transactions": 27.097631578947365, + "value": 179110.4255908138 }, "derivation": {}, "non_residential": { @@ -57,7 +57,7 @@ ], "minimum_price": 1000.0, "property_type": "non_residential", - "transactions": 9.653686298356586, - "value": 123175.93019207372 + "transactions": 9.064159936519806, + "value": 111830.66958141114 } } diff --git a/test_support/microcosm_build/uk_keyed_draws.py b/test_support/microcosm_build/uk_keyed_draws.py new file mode 100644 index 0000000000..f4c08c7082 --- /dev/null +++ b/test_support/microcosm_build/uk_keyed_draws.py @@ -0,0 +1,50 @@ +"""A stage's QRF draws move only for the records whose own inputs move. + +Every UK stage that imputes with ``RegimeGatedQRF`` draws from uniforms keyed +on each record's id (``stable_identity_uniforms``), never from a positional +stream. A positional stream let one engine change re-randomise the whole spine: +a few households' net income moved, the high-value copies changed the frame's +length, and every later household got different random numbers. This helper +checks the property the keying gives, on a stage's own imputation function. +""" + +from __future__ import annotations + +from collections.abc import Callable + +import numpy as np +import pandas as pd + +#: An id far from any fixture's, for the recipient the check adds. +_ADDED_ID_OFFSET = 1_000_000 + + +def assert_draws_are_keyed_on_ids( + draw: Callable[[pd.DataFrame, np.ndarray], pd.DataFrame], + recipient: pd.DataFrame, + ids: np.ndarray, + *, + perturbed_column: str, + perturbed_value: object, +) -> None: + """Reorder, change one recipient and add another: no other draw moves. + + ``draw(recipient, ids)`` returns the draws in recipient row order. The + check reverses the recipients, gives the last one (now first) a new + ``perturbed_column`` value, appends a copy of the first recipient under a + new id, and asserts that every recipient it did not touch draws exactly + what it drew before. + """ + + ids = np.asarray(ids) + base = draw(recipient.reset_index(drop=True), ids).set_axis(ids) + changed = recipient.iloc[::-1].reset_index(drop=True) + changed_ids = ids[::-1].copy() + changed.loc[0, perturbed_column] = perturbed_value + changed = pd.concat([changed, recipient.iloc[[0]]], ignore_index=True) + changed_ids = np.append(changed_ids, ids.max() + _ADDED_ID_OFFSET) + after = draw(changed, changed_ids).set_axis(changed_ids) + untouched = ids[:-1] + pd.testing.assert_frame_equal( + after.loc[untouched], base.loc[untouched], check_dtype=False + ) diff --git a/test_support/microcosm_build/uk_uc_reporter_redraw.py b/test_support/microcosm_build/uk_uc_reporter_redraw.py index 9254e8235d..3cd5b8caba 100644 --- a/test_support/microcosm_build/uk_uc_reporter_redraw.py +++ b/test_support/microcosm_build/uk_uc_reporter_redraw.py @@ -75,8 +75,11 @@ class _StubFittedQRF: def __init__(self, record: dict[str, object]) -> None: self.record = record - def predict(self, table: pd.DataFrame) -> pd.DataFrame: + def predict_from_uniforms( + self, table: pd.DataFrame, *, quantiles, sign_uniforms + ) -> pd.DataFrame: self.record["predict_index"] = table.index.tolist() + self.record["predict_uniforms"] = (quantiles, sign_uniforms) draws = np.zeros(len(table), dtype=float) if len(draws): draws[0] = 250.0