From 18ca76eab97761250797dd5f3e40a9f69e28475e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mar=C3=ADa=20Juaristi?= <127882282+juaristi22@users.noreply.github.com> Date: Sat, 5 Sep 2026 20:11:30 +0200 Subject: [PATCH 01/32] Add exact-household-count UK dataset candidates (#355) --- CLAUDE.md | 6 + changelog.d/uk-dataset-sizes-355.added.md | 1 + docs/uk-dataset-size-plan-355.md | 105 +++++++++ .../build/uk_runtime/dataset_size.py | 198 +++++++++++++++++ .../build/uk_runtime/local_rowwise.py | 62 +++++- .../tests/test_uk_local_rowwise.py | 199 +++++++++++++++++- .../tests/test_uk_rowwise_candidate.py | 80 +++++++ .../src/microcosm/calibrate/gates.py | 39 +++- .../src/microcosm/calibrate/initialization.py | 64 ++++++ .../src/microcosm/calibrate/solve.py | 41 +++- .../tests/test_informed_gates.py | 42 ++++ tools/build_uk_rowwise_candidate.py | 56 ++++- 12 files changed, 879 insertions(+), 14 deletions(-) create mode 100644 changelog.d/uk-dataset-sizes-355.added.md create mode 100644 docs/uk-dataset-size-plan-355.md create mode 100644 packages/microcosm-build/src/microcosm/build/uk_runtime/dataset_size.py create mode 100644 packages/microcosm-calibrate/src/microcosm/calibrate/initialization.py create mode 100644 packages/microcosm-calibrate/tests/test_informed_gates.py diff --git a/CLAUDE.md b/CLAUDE.md index ce1b71c3e..73a10fb72 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -106,3 +106,9 @@ with the journal pointing to them. Update this guide in the same PR whenever the workspace layout, test commands, or release flow change. If you find it contradicting the repo, trust the repo and fix this file. + +UK size experiments use `tools/build_uk_rowwise_candidate.py --dataset-households` +with the same pool inputs as the dense candidate. The flag changes exported +support, not clone K. Sizes remain candidate-only until their matched comparison +and promotion scorecard are adjudicated; see +[the size plan](docs/uk-dataset-size-plan-355.md). diff --git a/changelog.d/uk-dataset-sizes-355.added.md b/changelog.d/uk-dataset-sizes-355.added.md new file mode 100644 index 000000000..4e9bea82a --- /dev/null +++ b/changelog.d/uk-dataset-sizes-355.added.md @@ -0,0 +1 @@ +Add exact-household-count UK rowwise candidates using contribution-informed L0 initialization, protected target carriers, fixed-size sampling and refitting. Preserve the dense pool doctrine and local gates, export compact linked entities, and record selection provenance without promoting candidates to the certified dense release or changing dataset defaults. diff --git a/docs/uk-dataset-size-plan-355.md b/docs/uk-dataset-size-plan-355.md new file mode 100644 index 000000000..a098ee0ae --- /dev/null +++ b/docs/uk-dataset-size-plan-355.md @@ -0,0 +1,105 @@ +# UK dataset sizes: implementation plan and operating boundary + +This implements the candidate-building part of [#355](https://github.com/PolicyEngine/microcosm/issues/355) +on [#870](https://github.com/PolicyEngine/microcosm/pull/870)'s branch. The PR is still open as of +2026-09-05 and itself stacks on #852. Do not base the work on the old issue's +535,080-household 2023 dataset. The authoritative inputs are the current raw-FRS +2024-25 spine, OA ladder, and pinned Chronicle facts used by the joint candidate. + +## Decisions retained + +- Pool generation keeps clone count **K=15**. Requested output households is a + different parameter, applied after cloning and materialization. +- The joint surface retains national, constituency and local-authority rows, + the declared stretch bound **10**, loss cap **10**, **grain_equal** weighting, + and **1,500 epochs** per solve in the normal driver defaults. +- Existing binding adjudications, signed deferrals, measure exclusions, + census-vintage uprating, and all local gates remain in force. Size selection + does not silently drop target rows, loosen ESS floors or change registers. +- Population-normalized engine measures are frozen from the full pool. The + refit uses their selected household contributions rather than re-running + those formulas on a smaller population. +- The Frame carries complete households, benefit units and people; every + non-dry attempt retains the existing Logbook recording envelope. + +## Implemented sequence + +1. Build the usual joint dense solve from the pinned inputs. This also supplies + the same-target reference loss for the size comparison. +2. Compute each household's maximum absolute target-contribution share from + original pool weights and the compiled sparse matrix, following + [#346's correction](https://github.com/PolicyEngine/microcosm/issues/346#issuecomment-4902880142). + Protect the largest absolute weighted carrier of each nonzero target, with + first-column tie breaking. Initialize open probabilities with + `0.1 + 0.8 * score / (score + median_positive_score)`. This bounded smooth + prior is an implementation choice for candidate evaluation, not a measured + UK release ruling. Run L0 budget search; never select the largest prior scores. +3. Use the existing exact-count Sampford sampler on learned probabilities. + Probability-one gates are certainties (`pi_hi=1`), including every protected + carrier. Refuse impossible budgets or sampling designs rather than clamp. +4. Refit on the selected support through `microcosm.calibrate`, with no L0 + penalty. Reuse the existing normalized Horvitz–Thompson `w/q` baseline. + The stretch multiplier remains 10 **relative to that inclusion-adjusted + baseline**; this is the shared exact-k refit's contract, not a claim that + sparse weights remain within 10 times the unexpanded pool-row weights. + The manifest names the reference explicitly. Its empirical suitability for + UK release remains to be assessed alongside the size scorecard. +5. Restore the prepared carrier and export only selected linked entities. + Re-run the existing local gate battery on the compact frame. Each holdout + fold independently reruns selection on training targets; held targets do + not inform the prior or protected set. +6. Record requested/realized counts, pool positions, inclusion probabilities, + protected count, seed, learned penalty, dense/refit loss and target change, + the gate reports, and ordinary byte-pinned output metadata. + +## Running a candidate + +Use the inputs and environment from the existing +[UK dense assembly runbook](uk-dense-release-assembly-runbook-762.md). +Pass the same pinned source arguments to the existing driver and add: + +```bash +uv run python tools/build_uk_rowwise_candidate.py \ + --input-h5 "$UK_SPINE_H5" --input-sha256 "$UK_SPINE_SHA256" \ + --ladder "$UK_LADDER_NPZ" --ladder-sha256 "$UK_LADDER_SHA256" \ + --ledger-facts "$UK_LEDGER_FACTS" \ + --ledger-facts-sha256 "$UK_LEDGER_FACTS_SHA256" \ + --ledger-manifest-sha256 "$UK_LEDGER_MANIFEST_SHA256" \ + --dataset-households 50000 --seed 42 --out out/uk-k50000 +``` + +Repeat with another positive household count and a fresh output directory to +compare sizes. These are requested counts, not certified presets. Omit +`--dataset-households` to retain the existing dense path. Add `--dry-run` to +inspect input binding and parameters without solving or writing a candidate. +Never reduce `--n-clones` to request a smaller output. Full builds still need +the dense build's peak memory and add L0/refit work; the reduction is in the +exported dataset's storage and downstream loading/simulation footprint. + +## Certification and publication still required + +The implementation produces **candidates**, not a new certified UK default. +A size request refuses `--release-candidate`, and its manifest records +`releasable=false` even if the diagnostic gate run passes. The dense assembler +must not interpret a compact candidate as the already reviewed dense line. +The existing registry, production pointers and pe.py default are unchanged. + +Before any size can be promoted: + +1. Run the licensed full-input build, retain all gate failures and measure + local ESS and fit. A nominal 50k size is not guaranteed to clear the floors + that motivated K=15. +2. Run #355's matched sound-comparison protocol and the referenced promotion + scorecard, including reform/distributional validation and untargeted bases. + Calibration loss alone is not a certificate. The included same-target loss + comparison is a diagnostic, not a substitute for those protocols. +3. Adjudicate any new size-specific acceptance decisions, including the + inclusion-adjusted stretch reference. A national-only product would need + its own explicit scope; this implementation does not downgrade local claims. +4. Add size-specific certified release identities, assembly contracts and + downstream bundle entries against that evidence; publication remains the + repository's deliberate human step. + +Accordingly, this increment does not close #355's default-flip requirement. +Synthetic CI checks prove code behavior and artifact structure; they do not +establish licensed-data fit, storage measurements, or release eligibility. diff --git a/packages/microcosm-build/src/microcosm/build/uk_runtime/dataset_size.py b/packages/microcosm-build/src/microcosm/build/uk_runtime/dataset_size.py new file mode 100644 index 000000000..1a61ff4f8 --- /dev/null +++ b/packages/microcosm-build/src/microcosm/build/uk_runtime/dataset_size.py @@ -0,0 +1,198 @@ +"""Exact household-count UK candidates on a fixed, materialized target surface.""" + +from dataclasses import dataclass, replace + +import numpy as np +import pandas as pd + +from microcosm.calibrate import ( + CalibrationResult, + TargetSet, + assert_exact_k_support, + calibrate, + refit_l0_selection, + select_exact_k, +) +from microcosm.calibrate.initialization import contribution_initialization +from microcosm.frame import Frame + + +@dataclass(frozen=True) +class UKDatasetSize: + """A compact refit with its full-pool row identities and selection evidence.""" + + result: CalibrationResult + support: np.ndarray + receipt: dict[str, object] + + +def refit_uk_dataset_size( + frame: Frame, + dense: CalibrationResult, + *, + households: int, + epochs: int, + learning_rate: float, + seed: int, +) -> UKDatasetSize: + """Run informed L0, a fixed-size draw, and refit under the dense doctrine. + + Freeze the already compiled household contributions, including engine + measures that depend on the full population. Re-evaluating those formulas + on a subset would change the target system. Target values, order, loss + weights, cap and stretch bound remain those of the dense solve. + """ + n = frame.n("household") + if ( + isinstance(households, bool) + or not isinstance(households, int) + or not 0 < households <= n + ): + raise ValueError(f"households must be an integer in [1, {n}]; never clamped.") + if ( + dense.skipped + or len(dense.weights) != n + or dense.l0_lambda != 0 + or dense.weight_entity != "household" + or not np.array_equal( + frame.table("household")["household_id"].to_numpy(), + dense.frame.table("household")["household_id"].to_numpy(), + ) + or not np.array_equal( + frame.weights_for("household").values, dense.initial_weights + ) + ): + raise ValueError( + "size selection requires an aligned, fully compiled dense solve." + ) + if households == n: + return UKDatasetSize( + dense, + np.arange(n), + { + "method": "full_pool", + "requested_households": n, + "realized_households": n, + "pool_households": n, + "seed": seed, + }, + ) + problem = dense.problem + # Input weights are the pool design, not the concentrated dense weights. + init = contribution_initialization( + problem.matrix, dense.initial_weights, problem.target_vector + ) + if int(init.protected.sum()) > households: + raise ValueError( + f"requested {households} households cannot retain {int(init.protected.sum())} protected target carriers." + ) + common = dict( + weight_entity="household", + epochs=epochs, + learning_rate=learning_rate, + mass=dense.options["mass"], + max_weight_ratio=dense.options["max_weight_ratio"], + seed=seed, + target_loss_weights=dense.target_loss_weights, + target_loss_scales=dense.target_loss_scales, + target_loss_cap=dense.target_loss_cap, + ) + selection = calibrate( + frame, + TargetSet(problem.targets), + target_records=households, + gate_initialization=init, + mass_reason=dense.options["mass_reason"], + **common, + ) + probabilities = selection.gate_open_probabilities + if probabilities is None: + raise RuntimeError("informed L0 returned no selection probabilities.") + # Exact-one certainties are protected by the gates themselves; no top-k + # ranking or post-hoc promotion of learned boundary scores is performed. + support, sampling, q = select_exact_k( + probabilities, households, pi_hi=1.0, seed=seed + ) + support = assert_exact_k_support(support, households, pool_size=n) + if not np.isin(np.flatnonzero(init.protected), support).all(): + raise RuntimeError("exact-count selection lost a protected carrier.") + frozen = _frozen_targets(frame, dense, support) + refit = refit_l0_selection( + frame, + frozen, + selection, + support=support, + k=households, + support_inclusion_probabilities=q, + mass_reason=dense.options["mass_reason"], + **common, + ).refit + if ( + refit.skipped + or refit.frame.n("household") != households + or (refit.weights <= 0).any() + ): + raise RuntimeError("compact refit lost targets or positive household support.") + if not np.array_equal(refit.problem.target_vector, problem.target_vector): + raise RuntimeError("compact refit changed target values.") + errors_dense = np.asarray([d.final_estimate for d in dense.diagnostics]) + errors_small = np.asarray([d.final_estimate for d in refit.diagnostics]) + return UKDatasetSize( + refit, + support, + { + "method": "contribution_informed_l0_exact_count_refit", + "requested_households": households, + "realized_households": households, + "pool_households": n, + "seed": seed, + "protected_carriers": int(init.protected.sum()), + "selection_receipt": sampling, + "selection_l0_lambda": selection.l0_lambda, + "selection_epochs": epochs, + "refit_epochs": epochs, + "pool_row_indices": support.tolist(), + "inclusion_probabilities": q.tolist(), + "refit_baseline": "normalized_horvitz_thompson_w_over_q", + "stretch_reference": "normalized_horvitz_thompson_w_over_q", + "dense_loss": dense.final_loss, + "compact_loss": refit.final_loss, + "max_target_scaled_change": float( + np.max(np.abs(errors_small - errors_dense) / dense.target_loss_scales) + ), + "certification": "candidate_only_pending_matched_comparison_and_promotion_scorecard", + }, + ) + + +def _frozen_targets( + frame: Frame, dense: CalibrationResult, support: np.ndarray +) -> TargetSet: + """Bind sparse contribution rows to selected ids, with one shared id join.""" + ids = pd.Index(frame.table("household")["household_id"].iloc[support]) + matrix = dense.problem.matrix[:, support].tocsr() + cached_ids = None + cached_positions = None + + def positions(subset): + nonlocal cached_ids, cached_positions + current = subset.table("household")["household_id"] + if cached_ids is None or not current.equals(cached_ids): + cached_positions = ids.get_indexer(current) + if (cached_positions < 0).any(): + raise ValueError("frozen target requested unknown household ids.") + cached_ids = current.copy() + return cached_positions + + def measure(row): + def values(subset): + return np.asarray(matrix[[row], :].toarray()).reshape(-1)[positions(subset)] + + return values + + return TargetSet( + [ + replace(target, entity="household", measure=measure(row), filter=None) + for row, target in enumerate(dense.problem.targets) + ] + ) diff --git a/packages/microcosm-build/src/microcosm/build/uk_runtime/local_rowwise.py b/packages/microcosm-build/src/microcosm/build/uk_runtime/local_rowwise.py index a38650ef8..3a1c4cc01 100644 --- a/packages/microcosm-build/src/microcosm/build/uk_runtime/local_rowwise.py +++ b/packages/microcosm-build/src/microcosm/build/uk_runtime/local_rowwise.py @@ -151,6 +151,8 @@ class UKRowwiseDoctrineSolve: national_past_cap_census: Mapping[str, Any] all_past_cap_census: Mapping[str, Any] binding_adjudications: Mapping[str, Any] + selected_support: np.ndarray | None = None + size_receipt: Mapping[str, Any] | None = None def past_cap_census( @@ -1008,6 +1010,7 @@ def solve_uk_rowwise_weights_under_doctrine( learning_rate: float = 0.15, conserve_mass: bool = False, target_records: int | None = None, + dataset_households: int | None = None, l0_lambda: float = 0.0, budget_iters: int = 10, seed: int = 0, @@ -1142,6 +1145,26 @@ def solve_uk_rowwise_weights_under_doctrine( target_loss_weights=target_loss_weights, target_loss_cap=doctrine.target_loss_cap, ) + selected_support = None + size_receipt = None + if dataset_households is not None: + if target_records is not None or l0_lambda != 0: + raise ValueError( + "dataset_households requires an unpruned dense reference solve." + ) + from microcosm.build.uk_runtime.dataset_size import refit_uk_dataset_size + + sized = refit_uk_dataset_size( + frame, + result, + households=dataset_households, + epochs=epochs, + learning_rate=learning_rate, + seed=seed, + ) + result = sized.result + selected_support = sized.support + size_receipt = sized.receipt if result.skipped: reasons = [ f"{skipped.target.name}: {skipped.reason}" for skipped in result.skipped[:5] @@ -1280,9 +1303,32 @@ def solve_uk_rowwise_weights_under_doctrine( # ship if the kernel product held strata or metadata the rebuilt frame # lost (both trivially equal today; the guard is the boundary marker # for the day they are not). - clean_result = result.frame if restore is None else restore(result.frame) + if selected_support is None: + clean_result = result.frame if restore is None else restore(result.frame) + expected_ids = problem.household_ids + else: + # Restore the full prepared carrier before selecting: legacy restorers + # retain full original tables and cannot accept compact weight arrays. + clean_pool = frame if restore is None else restore(frame) + expected_ids = tuple(problem.household_ids[i] for i in selected_support) + clean_subset = clean_pool.select( + clean_pool.table("person")["person_household_id"] + .isin(expected_ids) + .to_numpy() + ) + clean_result = Frame( + {entity: clean_subset.table(entity) for entity in clean_subset.entities}, + clean_subset.schema, + { + entity: result.frame.weights_for(entity) + for entity in result.frame.weighted_entities + }, + clean_subset.strata, + mass_log=result.frame.mass_log, + metadata=result.frame.metadata, + ) restored_ids = tuple(clean_result.table("household")["household_id"].tolist()) - if restored_ids != problem.household_ids: + if restored_ids != expected_ids: raise ValueError( "restore returned households that do not match the problem's rows " "(same ids, same order); weights are written back by position, so a " @@ -1344,6 +1390,8 @@ def solve_uk_rowwise_weights_under_doctrine( national_past_cap_census=national_census, all_past_cap_census=all_census, binding_adjudications=binding_adjudications, + selected_support=selected_support, + size_receipt=size_receipt, ) @@ -1380,6 +1428,7 @@ def rotated_uk_local_holdout( learning_rate: float = 0.15, conserve_mass: bool = False, target_records: int | None = None, + dataset_households: int | None = None, l0_lambda: float = 0.0, budget_iters: int = 10, solve_seed: int = 0, @@ -1424,13 +1473,20 @@ def rotated_uk_local_holdout( learning_rate=learning_rate, conserve_mass=conserve_mass, target_records=target_records, + dataset_households=dataset_households, l0_lambda=l0_lambda, budget_iters=budget_iters, seed=solve_seed, ) held_targets = problem.targets[holdout_indices] held_estimates = np.asarray( - problem.matrix[holdout_indices] @ train_solve.weights, + problem.matrix[holdout_indices][ + :, + np.arange(problem.n_households) + if train_solve.selected_support is None + else train_solve.selected_support, + ] + @ train_solve.weights, dtype=np.float64, ).reshape(-1) held_weights = uk_local_target_loss_weights( diff --git a/packages/microcosm-build/tests/test_uk_local_rowwise.py b/packages/microcosm-build/tests/test_uk_local_rowwise.py index faecd4975..bf7a838bc 100644 --- a/packages/microcosm-build/tests/test_uk_local_rowwise.py +++ b/packages/microcosm-build/tests/test_uk_local_rowwise.py @@ -651,7 +651,7 @@ def test_rotated_holdout_keeps_national_rows_in_every_training_fold( def fake_solve(frame, training_problem, **kwargs): calls.append(kwargs["national_rows"]) - return type("Solve", (), {"weights": np.ones(3)})() + return type("Solve", (), {"weights": np.ones(3), "selected_support": None})() monkeypatch.setattr( local_rowwise, "solve_uk_rowwise_weights_under_doctrine", fake_solve @@ -1220,3 +1220,200 @@ def reordering_restore(frame): restore=reordering_restore, epochs=2, ) + + +def test_exact_size_preserves_targets_and_household_links(): + from microcosm.build.uk_runtime.dataset_size import refit_uk_dataset_size + from microcosm.calibrate import Target, TargetSet, calibrate + + frame = _clone_frame() + dense = calibrate( + frame, + TargetSet( + [Target("count", "household", lambda f: np.ones(f.n("household")), 3)] + ), + epochs=2, + ) + sized = refit_uk_dataset_size( + frame, dense, households=2, epochs=2, learning_rate=0.02, seed=7 + ) + assert sized.result.frame.n("household") == 2 + assert sized.result.frame.n("person") == 2 + assert sized.result.frame.n("benunit") == 2 + assert sized.result.problem.names == dense.problem.names + assert sized.result.initial_weights.sum() == pytest.approx( + dense.initial_weights.sum() + ) + assert (sized.result.weights <= 10 * sized.result.initial_weights).all() + again = refit_uk_dataset_size( + frame, dense, households=2, epochs=2, learning_rate=0.02, seed=7 + ) + np.testing.assert_array_equal(sized.support, again.support) + + +@pytest.mark.parametrize("size", [0, -1, True, 4, 1.5]) +def test_exact_size_refuses_invalid_sizes(size): + from microcosm.build.uk_runtime.dataset_size import refit_uk_dataset_size + from microcosm.calibrate import Target, TargetSet, calibrate + + frame = _clone_frame() + dense = calibrate( + frame, + TargetSet( + [Target("count", "household", lambda f: np.ones(f.n("household")), 3)] + ), + epochs=1, + ) + with pytest.raises(ValueError, match="integer"): + refit_uk_dataset_size( + frame, dense, households=size, epochs=1, learning_rate=0.02, seed=7 + ) + + +def test_size_solve_restores_full_prepared_tables_before_subsetting(): + frame = _clone_frame() + metrics = pd.DataFrame({"households": [1.0, 1.0, 1.0]}, index=[101, 102, 103]) + problem = build_uk_rowwise_local_matrix( + metrics, + _assigned(), + pd.DataFrame({"code": ["E001", "S001"], "households": [2.0, 1.0]}), + ) + restored_counts = [] + + def restore(full): + restored_counts.append(full.n("household")) + return full + + result = solve_uk_rowwise_weights_under_doctrine( + frame, + problem, + bound_families=["census_households/constituency"], + dataset_households=2, + epochs=2, + seed=7, + restore=restore, + ) + assert restored_counts == [3] + assert result.frame.n("household") == 2 + assert result.frame.n("person") == 2 + assert len(result.diagnostics) == 2 + assert "census_households/constituency" in result.frame.mass_log[-1].reason + assert result.selected_support.tolist() == [0, 2] + + +def test_size_holdout_uses_compact_support_and_reselects_each_fold(monkeypatch): + import microcosm.build.uk_runtime.local_rowwise as runtime + + metrics = pd.DataFrame({"households": [1.0, 2.0, 3.0]}, index=[101, 102, 103]) + calls = [] + + def solve(frame, training, **kwargs): + calls.append((len(training.targets), kwargs["dataset_households"])) + return type( + "Solve", + (), + {"weights": np.array([1.0, 1.0]), "selected_support": np.array([0, 2])}, + )() + + monkeypatch.setattr(runtime, "solve_uk_rowwise_weights_under_doctrine", solve) + # The holdout helper needs at least one target in each of five folds. + richer = build_uk_rowwise_local_matrix( + pd.DataFrame( + { + name: [1.0, 2.0, 3.0] + for name in ("households", "tenure/social_rent", "tenure/private_rent") + }, + index=metrics.index, + ), + _assigned(), + pd.DataFrame( + { + "code": ["E001", "S001"], + **{ + name: [3.0, 3.0] + for name in ( + "households", + "tenure/social_rent", + "tenure/private_rent", + ) + }, + } + ), + ) + monkeypatch.setattr( + runtime, "_derive_uk_local_bound_families_from_target_frame", lambda *a, **k: () + ) + receipt = rotated_uk_local_holdout( + _clone_frame(), richer, dataset_households=2, epochs=1 + ) + assert len(calls) == 5 + assert all(n < len(richer.targets) and k == 2 for n, k in calls) + assert np.isfinite(receipt["mean_holdout_loss"]) + + +def test_size_refit_freezes_population_dependent_measures(): + from microcosm.build.uk_runtime.dataset_size import refit_uk_dataset_size + from microcosm.calibrate import Target, TargetSet, calibrate + + frame = _clone_frame() + # A population-normalized measure changes if evaluated on two instead of + # three households. The refit must retain the full-pool value of 1/3. + targets = TargetSet( + [ + Target( + "normalized", + "household", + lambda f: np.full(f.n("household"), 1 / f.n("household")), + 1, + ) + ] + ) + dense = calibrate(frame, targets, epochs=2) + small = refit_uk_dataset_size( + frame, dense, households=2, epochs=2, learning_rate=0.02, seed=7 + ) + np.testing.assert_allclose(small.result.problem.matrix.toarray(), [[1 / 3, 1 / 3]]) + + +def test_full_size_returns_dense_reference_without_search(): + from microcosm.build.uk_runtime.dataset_size import refit_uk_dataset_size + from microcosm.calibrate import Target, TargetSet, calibrate + + frame = _clone_frame() + dense = calibrate( + frame, + TargetSet( + [Target("count", "household", lambda f: np.ones(f.n("household")), 3)] + ), + epochs=1, + ) + full = refit_uk_dataset_size( + frame, dense, households=3, epochs=1, learning_rate=0.02, seed=7 + ) + assert full.result is dense + assert full.support.tolist() == [0, 1, 2] + + +def test_size_refuses_budget_smaller_than_protected_carriers(): + from microcosm.build.uk_runtime.dataset_size import refit_uk_dataset_size + from microcosm.calibrate import Target, TargetSet, calibrate + + frame = _clone_frame() + targets = TargetSet( + [ + Target( + str(i), + "household", + lambda f, i=i: ( + f.table("household").household_id.to_numpy() == i + ).astype(float), + 1, + ) + for i in [101, 102, 103] + ] + ) + dense = calibrate(frame, targets, epochs=1) + with pytest.raises(ValueError, match="3 protected target carriers"): + refit_uk_dataset_size( + frame, dense, households=2, epochs=1, learning_rate=0.02, seed=7 + ) diff --git a/packages/microcosm-build/tests/test_uk_rowwise_candidate.py b/packages/microcosm-build/tests/test_uk_rowwise_candidate.py index 8bd053f90..b4623de8c 100644 --- a/packages/microcosm-build/tests/test_uk_rowwise_candidate.py +++ b/packages/microcosm-build/tests/test_uk_rowwise_candidate.py @@ -1822,3 +1822,83 @@ def test_candidate_multi_block_engine_run_is_never_releasable( "single_block_engine": False, "release_blocking_gates_passed": True, } + + +def test_size_candidate_exports_compact_links_and_cannot_claim_dense_release( + monkeypatch, tmp_path +): + builder = _load_builder_module() + input_h5 = tmp_path / "spine.h5" + ladder_path = tmp_path / "ladder.npz" + out = tmp_path / "k300" + _write_staging_h5(input_h5, households_per_region=52) + ladder = _write_ladder(ladder_path) + import microcosm.build.uk_runtime.battery_bindings as bindings + + monkeypatch.setattr( + bindings, + "_local_area_roster", + lambda _resource, levels: { + "constituency": tuple(sorted(set(ladder.constituency_code))), + "local_authority": tuple(sorted(set(ladder.local_authority_code))), + }, + ) + status = builder.main( + [ + "--input-h5", + str(input_h5), + "--ladder", + str(ladder_path), + "--out", + str(out), + "--n-clones", + "2", + "--dataset-households", + "300", + "--epochs", + "2", + "--skip-holdout", + "--seed", + "7", + ] + ) + assert status in (0, 1) # Gate failures remain reportable candidates. + manifest = json.loads((out / builder.MANIFEST_FILENAME).read_text()) + assert manifest["releasable"] is False + assert manifest["release_posture"]["size_certification_present"] is False + size = manifest["solve"]["dataset_size"] + assert size["requested_households"] == size["realized_households"] == 300 + assert size["pool_households"] == 416 + assert manifest["parameters"]["n_clones"] == 2 + assert ( + manifest["weights"]["stretch_reference"] + == "normalized_horvitz_thompson_w_over_q" + ) + path = out / builder.CANDIDATE_FILENAME_TEMPLATE.format(calibration_year=2025) + with pd.HDFStore(path, "r") as store: + households = store["household"] + persons = store["person"] + benunits = store["benunit"] + assert len(households) == 300 + assert set(persons.person_household_id) == set(households.household_id) + assert set(persons.person_benunit_id) == set(benunits.benunit_id) + assert len(_spool_rows(out)) == 1 + + +def test_size_cli_refuses_promotion_without_separate_certification(tmp_path): + builder = _load_builder_module() + args = builder._parse_args( + [ + "--input-h5", + str(tmp_path / "spine.h5"), + "--ladder", + str(tmp_path / "ladder.npz"), + "--out", + str(tmp_path / "out"), + "--dataset-households", + "50000", + "--release-candidate", + ] + ) + with pytest.raises(ValueError, match="candidate-only"): + builder._validate_cli_args(args) diff --git a/packages/microcosm-calibrate/src/microcosm/calibrate/gates.py b/packages/microcosm-calibrate/src/microcosm/calibrate/gates.py index 306d42615..5874d3803 100644 --- a/packages/microcosm-calibrate/src/microcosm/calibrate/gates.py +++ b/packages/microcosm-calibrate/src/microcosm/calibrate/gates.py @@ -19,6 +19,7 @@ import math +import numpy as np import torch from torch import nn @@ -66,6 +67,10 @@ class HardConcrete(nn.Module): penalty close them). Must lie in ``(0, 1)``. temperature: Concrete-distribution temperature; lower is closer to a hard Bernoulli. + initial_probabilities: Optional per-record actual open probabilities + in (0, 1); overrides the historical scalar initialization. + protected_mask: Optional boolean vector. These records have gates + fixed at one during training and evaluation. Raises: ValueError: If ``init_mean`` is not strictly between 0 and 1, or @@ -78,6 +83,8 @@ def __init__( *, init_mean: float = 0.999, temperature: float = 0.25, + initial_probabilities: np.ndarray | None = None, + protected_mask: np.ndarray | None = None, ) -> None: super().__init__() if not (0.0 < init_mean < 1.0): @@ -88,6 +95,22 @@ def __init__( raise ValueError( f"HardConcrete.temperature must be positive, got {temperature!r}." ) + if initial_probabilities is not None: + probabilities = np.asarray(initial_probabilities, dtype=np.float64) + if ( + probabilities.shape != (n,) + or not np.isfinite(probabilities).all() + or ((probabilities <= 0) | (probabilities >= 1)).any() + ): + raise ValueError( + "initial_probabilities must align with gates and lie in (0, 1)." + ) + if protected_mask is None: + protected_mask = np.zeros(n, dtype=bool) + protected_mask = np.asarray(protected_mask) + if protected_mask.shape != (n,) or protected_mask.dtype != np.bool_: + raise ValueError("protected_mask must be an aligned boolean vector.") + self.register_buffer("protected_mask", torch.tensor(protected_mask)) self.temperature = float(temperature) self.gamma = _GAMMA self.zeta = _ZETA @@ -95,6 +118,12 @@ def __init__( init_val = math.log(init_mean / (1.0 - init_mean)) with torch.no_grad(): self.qz_logits.fill_(init_val) + if initial_probabilities is not None: + # The new vector is an actual open probability, unlike the + # historical scalar init_mean's sigmoid-logit convention. + shift = self.temperature * math.log(-self.gamma / self.zeta) + logits = np.log(probabilities / (1 - probabilities)) + shift + self.qz_logits.copy_(torch.tensor(logits, dtype=self.qz_logits.dtype)) def forward(self) -> torch.Tensor: """Return the current gates: sampled in training, deterministic in eval.""" @@ -105,7 +134,7 @@ def forward(self) -> torch.Tensor: else: s = torch.sigmoid(self.qz_logits) stretched = s * (self.zeta - self.gamma) + self.gamma - return torch.clamp(stretched, 0.0, 1.0) + return torch.where(self.protected_mask, 1.0, torch.clamp(stretched, 0.0, 1.0)) def get_penalty(self) -> torch.Tensor: """Expected number of open gates — the differentiable L0 surrogate. @@ -116,9 +145,13 @@ def get_penalty(self) -> torch.Tensor: ``l0_lambda`` and adds to the loss. """ shift = self.temperature * math.log(-self.gamma / self.zeta) - return torch.sigmoid(self.qz_logits - shift).sum() + return torch.where( + self.protected_mask, 1.0, torch.sigmoid(self.qz_logits - shift) + ).sum() def get_active_prob(self) -> torch.Tensor: """Per-gate open-probability (the per-record version of the penalty).""" shift = self.temperature * math.log(-self.gamma / self.zeta) - return torch.sigmoid(self.qz_logits - shift) + return torch.where( + self.protected_mask, 1.0, torch.sigmoid(self.qz_logits - shift) + ) diff --git a/packages/microcosm-calibrate/src/microcosm/calibrate/initialization.py b/packages/microcosm-calibrate/src/microcosm/calibrate/initialization.py new file mode 100644 index 000000000..f8d15d2ab --- /dev/null +++ b/packages/microcosm-calibrate/src/microcosm/calibrate/initialization.py @@ -0,0 +1,64 @@ +"""Contribution-informed L0 initialization (microcosm#346, #355).""" + +from dataclasses import dataclass + +import numpy as np +from scipy import sparse + + +@dataclass(frozen=True) +class GateInitialization: + """Aligned open probabilities and records excluded from gate pruning.""" + + probabilities: np.ndarray + protected: np.ndarray + + +def contribution_initialization( + matrix: sparse.sparray, + weights: np.ndarray, + targets: np.ndarray, +) -> GateInitialization: + """Initialize search from maximum absolute target-contribution share. + + Protect the largest absolute weighted carrier of each nonzero target + (first column breaks ties deterministically). The smooth bounded prior + uses the median positive share as its scale. It only initializes L0; + neither these scores nor their ranks select the final support. Zero + targets use a unit denominator, matching the default calibration scale. + Work and storage are sparse in the target matrix, never targets x pool. + """ + matrix = sparse.csr_array(matrix, dtype=np.float64, copy=True) + matrix.sum_duplicates() + matrix.eliminate_zeros() + weights = np.asarray(weights, dtype=np.float64) + targets = np.asarray(targets, dtype=np.float64) + if weights.shape != (matrix.shape[1],) or targets.shape != (matrix.shape[0],): + raise ValueError("matrix, weights and targets must align.") + if ( + not np.isfinite(matrix.data).all() + or not np.isfinite(weights).all() + or not np.isfinite(targets).all() + or (weights <= 0).any() + ): + raise ValueError( + "contribution initialization requires finite inputs and positive weights." + ) + scores = np.zeros(len(weights)) + protected = np.zeros(len(weights), dtype=bool) + for row, target in enumerate(targets): + start, stop = matrix.indptr[row : row + 2] + indices = matrix.indices[start:stop] + contributions = np.abs(matrix.data[start:stop]) * weights[indices] + if not len(indices): + if target != 0: + raise ValueError(f"nonzero target at row {row} has no support.") + continue + shares = contributions / (abs(target) if target != 0 else 1.0) + np.maximum.at(scores, indices, shares) + if target != 0: + protected[indices[np.argmax(contributions)]] = True + positive = scores[scores > 0] + scale = float(np.median(positive)) if len(positive) else 1.0 + probabilities = 0.1 + 0.8 * scores / (scores + scale) + return GateInitialization(probabilities, protected) diff --git a/packages/microcosm-calibrate/src/microcosm/calibrate/solve.py b/packages/microcosm-calibrate/src/microcosm/calibrate/solve.py index c37ed85c9..40d0619c8 100644 --- a/packages/microcosm-calibrate/src/microcosm/calibrate/solve.py +++ b/packages/microcosm-calibrate/src/microcosm/calibrate/solve.py @@ -77,6 +77,7 @@ from microcosm.calibrate.exact_k import assert_exact_k_support from microcosm.calibrate.gates import HardConcrete +from microcosm.calibrate.initialization import GateInitialization from microcosm.calibrate.matrix import ( CalibrationProblem, SkippedTarget, @@ -757,6 +758,7 @@ def _optimize( target_loss_cap: float, initial_weights: np.ndarray, *, + gate_initialization: GateInitialization | None = None, warm_start_weights: np.ndarray | None = None, epochs: int, learning_rate: float, @@ -810,7 +812,19 @@ def _optimize( gates: HardConcrete | None = None params: list[torch.Tensor] = [log_w] if l0_lambda > 0.0 or target_records is not None: - gates = HardConcrete(len(w0), init_mean=init_mean, temperature=temperature) + gates = HardConcrete( + len(w0), + init_mean=init_mean, + temperature=temperature, + **( + {} + if gate_initialization is None + else { + "initial_probabilities": gate_initialization.probabilities, + "protected_mask": gate_initialization.protected, + } + ), + ) params = [log_w, *gates.parameters()] optimizer = torch.optim.Adam(params, lr=learning_rate) @@ -1098,6 +1112,7 @@ def _search_l0_lambda_for_budget( target_loss_cap: float, initial_weights: np.ndarray, *, + gate_initialization: GateInitialization | None = None, target_records: int, epochs: int, learning_rate: float, @@ -1189,6 +1204,8 @@ def evaluate( "l0_lambda": lam, }, } + if gate_initialization is not None: + optimize_kwargs["gate_initialization"] = gate_initialization if return_gate_open_probabilities: weights, trajectory, gate_open_probabilities = _optimize( matrix, @@ -1377,6 +1394,7 @@ def calibrate( frame: Frame, targets: TargetSet, *, + gate_initialization: GateInitialization | None = None, weight_entity: str = "household", method: str = "adam", epochs: int = 256, @@ -1478,6 +1496,9 @@ def calibrate( cannot express, e.g. the pre-selection *design* weights during a post-L0 refit. When supplied, ``l2_anchor`` becomes a free-form provenance label recorded in options (e.g. ``"design"``). + gate_initialization: Optional informed L0 prior and protected-record + mask, aligned to the weight entity. Requires an Adam L0 solve; + protected gates remain open throughout training and selection. init_mean: Initial expected open-probability of the L0 gates (only used when pruning). temperature: Hard-concrete temperature (only used when pruning). @@ -1565,6 +1586,9 @@ def calibrate( "total cannot be conserved (sum of caps < input total). Use " "max_weight_ratio >= 1, or mass='free'." ) + if gate_initialization is not None: + if method != "adam" or (target_records is None and l0_lambda <= 0): + raise ValueError("gate_initialization requires an Adam L0 solve.") if target_records is not None and ( not isinstance(target_records, int) or target_records <= 0 ): @@ -1755,6 +1779,11 @@ def calibrate( progress_callback=progress_callback, budget_iters=budget_iters, return_gate_open_probabilities=True, + **( + {} + if gate_initialization is None + else {"gate_initialization": gate_initialization} + ), ) else: effective_l0 = l0_lambda @@ -1780,6 +1809,11 @@ def calibrate( progress_callback=progress_callback, return_gate_open_probabilities=True, selection_receipt=iterate_selection_receipt, + **( + {} + if gate_initialization is None + else {"gate_initialization": gate_initialization} + ), ) n_nonzero = int((final_weights > prune_atol).sum()) if effective_l0 > 0.0 and n_nonzero == 0: @@ -1840,6 +1874,7 @@ def calibrate( target_loss_scales=target_loss_scales_np.copy(), target_loss_cap=target_loss_cap, options={ + "gate_initialization_supplied": gate_initialization is not None, "method": method, "epochs": epochs, "learning_rate": learning_rate, @@ -2012,6 +2047,7 @@ def refit_l0_selection( epochs: int = 256, learning_rate: float = 0.02, mass: str = FREE_MASS, + mass_reason: str | None = None, max_weight_ratio: float | None = None, l2_lambda: float = 0.0, l2_anchor: str = "initial", @@ -2035,6 +2071,8 @@ def refit_l0_selection( gates and L0 penalty are removed and :func:`calibrate` runs again with ``l0_lambda=0`` on only the selected records. + ``mass_reason`` carries a caller's calibration provenance into the refit. + The explicit exact-k path subsets the original ``frame`` rather than ``selection.frame``. A record may have positive open probability while its deterministic L0 gate is closed; starting from the original frame keeps @@ -2184,6 +2222,7 @@ def refit_l0_selection( epochs=epochs, learning_rate=learning_rate, mass=mass, + mass_reason=mass_reason, max_weight_ratio=max_weight_ratio, target_records=None, l0_lambda=0.0, diff --git a/packages/microcosm-calibrate/tests/test_informed_gates.py b/packages/microcosm-calibrate/tests/test_informed_gates.py new file mode 100644 index 000000000..b4a779f0c --- /dev/null +++ b/packages/microcosm-calibrate/tests/test_informed_gates.py @@ -0,0 +1,42 @@ +"""Informed L0 must preserve rare carriers without replacing the search.""" + +import numpy as np +import pytest +import torch +from scipy import sparse + +from microcosm.calibrate.gates import HardConcrete +from microcosm.calibrate.initialization import contribution_initialization + + +def test_small_weight_large_measure_is_protected(): + init = contribution_initialization( + sparse.csr_array([[1000.0, 0, 0], [0, 1, 1]]), + np.array([0.01, 100.0, 100.0]), + np.array([10.0, 200.0]), + ) + assert init.protected[0] + assert init.probabilities[0] > init.probabilities[1] + + +def test_protected_gates_stay_open_during_training_and_evaluation(): + gates = HardConcrete( + 3, + initial_probabilities=np.array([0.2, 0.4, 0.8]), + protected_mask=np.array([True, False, False]), + ) + assert np.allclose(gates.get_active_prob().detach().numpy(), [1, 0.4, 0.8]) + with torch.no_grad(): + gates.qz_logits.fill_(-100) + for training in (True, False): + gates.train(training) + assert gates()[0].item() == 1 + assert gates.get_active_prob()[0].item() == 1 + + +@pytest.mark.parametrize( + "probabilities", [[0, 0.5], [0.5, 1], [float("nan"), 0.5], [0.5]] +) +def test_invalid_initial_probabilities_refuse(probabilities): + with pytest.raises(ValueError): + HardConcrete(2, initial_probabilities=np.array(probabilities)) diff --git a/tools/build_uk_rowwise_candidate.py b/tools/build_uk_rowwise_candidate.py index 101ca0bb2..a7a0a4dd5 100644 --- a/tools/build_uk_rowwise_candidate.py +++ b/tools/build_uk_rowwise_candidate.py @@ -591,6 +591,11 @@ def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: required=True, help="Output directory for the candidate H5 and evidence sidecars.", ) + parser.add_argument( + "--dataset-households", + type=int, + help="Exact output household count after informed L0 and refit; pool clone K is unchanged. Candidate-only until size certification.", + ) parser.add_argument("--n-clones", type=int, default=UK_LOCAL_CLONE_COUNT) parser.add_argument( "--candidate-clone-counts", @@ -841,6 +846,13 @@ def _run_candidate( source_lineage_modulus=args.source_lineage_modulus, ) clone = assignment.result + if ( + args.dataset_households is not None + and args.dataset_households > clone.frame.n("household") + ): + raise ValueError( + "--dataset-households exceeds the cloned pool; selection never clamps the request." + ) if state is not None: append_phase(state, "cloned") @@ -977,6 +989,7 @@ def _run_candidate( learning_rate=args.learning_rate, conserve_mass=_CONSERVE_MASS, target_records=_TARGET_RECORDS, + dataset_households=args.dataset_households, l0_lambda=_L0_LAMBDA, budget_iters=_BUDGET_ITERS, seed=args.seed, @@ -1083,6 +1096,7 @@ def _run_candidate( learning_rate=args.learning_rate, conserve_mass=_CONSERVE_MASS, target_records=_TARGET_RECORDS, + dataset_households=args.dataset_households, l0_lambda=_L0_LAMBDA, budget_iters=_BUDGET_ITERS, solve_seed=args.seed, @@ -1587,6 +1601,7 @@ def _joint_dry_run_plan( }, "candidate_clone_counts": list(args.candidate_clone_counts or (args.n_clones,)), "candidate_clone_support": clone_support, + "parameters": _parameters(args, source_year=source_year), "releasable": False, "engine": "not_run", "ladder_target_provenance": dict(target_provenance), @@ -1995,7 +2010,9 @@ def _dry_run_plan( "fraction": args.sample_fraction, "unreachable_check": "completed", }, - "releasable": args.sample_fraction == 1.0 and args.engine_blocks == 1, + "releasable": args.sample_fraction == 1.0 + and args.engine_blocks == 1 + and args.dataset_households is None, "parameters": _parameters(args, source_year=source_year), "shapes": { "person": list(clone.frame.table("person").shape), @@ -2289,11 +2306,14 @@ def _manifest( }, "abs_delta": abs(new_total - old_total), "declared_stretch_bound": float(UK_LOCAL_MAX_WEIGHT_RATIO), + "stretch_reference": "pool_design" + if solve.size_receipt is None + else "normalized_horvitz_thompson_w_over_q", "realized_max_weight_ratio_vs_design": float( np.max( np.divide( np.asarray(solve.weights, dtype=np.float64), - np.asarray(clone.frame.weights_for("household").values), + np.asarray(solve.initial_weights), ) ) ), @@ -2305,7 +2325,11 @@ def _manifest( "ladder": ladder_rows, "national": int(len(solve.national_diagnostics)), }, - "n_households": int(problem.n_households), + "n_households": int(solve.frame.n("household")), + "pool_households": int(problem.n_households), + "dataset_size": None + if solve.size_receipt is None + else dict(solve.size_receipt), "initial_loss": float(solve.initial_loss), "final_loss": float(solve.final_loss), "max_abs_relative_error": float(abs_errors.max()), @@ -2390,8 +2414,15 @@ def _manifest( for gate_id, payload in gate_rows.items() if not isinstance(payload, Mapping) or payload.get("status") != "passed" ), - "releasable": releasable, - "release_posture": release_posture, + "releasable": releasable and args.dataset_households is None, + "release_posture": { + **release_posture, + **( + {} + if args.dataset_households is None + else {"size_certification_present": False} + ), + }, "ladder_household_uprating": dict( cross_grain.get("ladder_household_uprating") or {"applied": False, "reason": "no cross-grain receipt"} @@ -2421,6 +2452,7 @@ def _manifest( def _parameters(args: argparse.Namespace, *, source_year: int) -> dict[str, Any]: return { "n_clones": int(args.n_clones), + "dataset_households": args.dataset_households, "seed": int(args.seed), "source_year": source_year, "source_lineage_modulus": args.source_lineage_modulus, @@ -2523,7 +2555,12 @@ def _validate_solve_result( raise RuntimeError( "doctrine solve returned no past-cap census; refusing candidate." ) - if len(solve.weights) != problem.n_households: + expected_count = ( + problem.n_households + if solve.selected_support is None + else len(solve.selected_support) + ) + if len(solve.weights) != expected_count: raise RuntimeError( "doctrine solve returned a weight vector with the wrong length." ) @@ -2585,6 +2622,13 @@ def _validate_cli_args(args: argparse.Namespace) -> None: raise ValueError( "the joint registry path requires --input-sha256 and --ladder-sha256." ) + if args.dataset_households is not None: + if args.dataset_households <= 0: + raise ValueError("--dataset-households must be positive.") + if args.release_candidate: + raise ValueError( + "--dataset-households is candidate-only: size-specific matched comparison and promotion scorecard are required before release." + ) if args.release_candidate: required_release = { "--input-sha256": args.input_sha256, From a53236f892e4e9b298dc06ff4ba0b7f350399beb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mar=C3=ADa=20Juaristi?= <127882282+juaristi22@users.noreply.github.com> Date: Tue, 8 Sep 2026 13:52:59 +0200 Subject: [PATCH 02/32] Size runs keep their dense reference and selection as evidence; --selection-seed moves only the draw (#355) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A --dataset-households run replaces its calibration product with the compact refit, but the dense joint solve it was cut from is byte-identical to a standalone dense run on the same inputs, seed and epochs. Keeping it lets a size candidate be compared with its own dense reference without a second full-pool run: - local_rowwise: the evidence labelling (alignment by name, exact target values, local/national diagnostics, past-cap censuses, initial loss) moves into _doctrine_solve_evidence and runs for both results; the solve product carries `dense_reference` (weights, initial weights, diagnostics, losses, censuses). - driver: a size run writes dense_reference_diagnostics.csv (every target's dense estimate, local rows then national, with a `grain` column) and dataset_size_selection.csv (pool row index, household id, clone index, design weight, inclusion probability, certainty, Horvitz–Thompson baseline, refit weight), lists both under outputs with digests, and summarises the reference under solve.dataset_size (losses, fit by family, uk_weight_summary). Dense runs are unchanged; the two files are size-run-only in the publish order. - --selection-seed (requires --dataset-households; defaults to --seed) seeds only the informed L0 search, the exact-count draw and the refit, threaded through the holdout as well, so two selections compare on one pool and one dense reference; recorded as parameters.selection_seed. Tests: the doctrine solve keeps a reference equal to the standalone dense solve and a selection seed leaves it untouched; the CLI size run writes both sidecars, lists them, and records both seeds; --selection-seed without a size is refused. Co-Authored-By: Claude Fable 5.1 --- changelog.d/uk-dataset-sizes-355.added.md | 1 + .../build/uk_runtime/local_rowwise.py | 299 ++++++++++++------ .../tests/test_uk_local_rowwise.py | 46 +++ .../tests/test_uk_rowwise_candidate.py | 78 +++++ tools/build_uk_rowwise_candidate.py | 135 +++++++- 5 files changed, 464 insertions(+), 95 deletions(-) diff --git a/changelog.d/uk-dataset-sizes-355.added.md b/changelog.d/uk-dataset-sizes-355.added.md index 4e9bea82a..aa8e3a7c9 100644 --- a/changelog.d/uk-dataset-sizes-355.added.md +++ b/changelog.d/uk-dataset-sizes-355.added.md @@ -1 +1,2 @@ Add exact-household-count UK rowwise candidates using contribution-informed L0 initialization, protected target carriers, fixed-size sampling and refitting. Preserve the dense pool doctrine and local gates, export compact linked entities, and record selection provenance without promoting candidates to the certified dense release or changing dataset defaults. +A size run also ships the dense joint solve it was cut from (`dense_reference_diagnostics.csv`, a `dense_reference` summary under `solve.dataset_size`) and the selection itself (`dataset_size_selection.csv`: pool row, design weight, inclusion probability, Horvitz–Thompson baseline, refit weight), and `--selection-seed` re-draws the selection on one pool and one dense reference. diff --git a/packages/microcosm-build/src/microcosm/build/uk_runtime/local_rowwise.py b/packages/microcosm-build/src/microcosm/build/uk_runtime/local_rowwise.py index 3a1c4cc01..770fc304a 100644 --- a/packages/microcosm-build/src/microcosm/build/uk_runtime/local_rowwise.py +++ b/packages/microcosm-build/src/microcosm/build/uk_runtime/local_rowwise.py @@ -153,6 +153,42 @@ class UKRowwiseDoctrineSolve: binding_adjudications: Mapping[str, Any] selected_support: np.ndarray | None = None size_receipt: Mapping[str, Any] | None = None + dense_reference: UKRowwiseDenseReference | None = None + + +@dataclass(frozen=True) +class UKRowwiseDenseReference: + """The dense joint solve a size selection started from, kept as evidence. + + A ``dataset_households`` solve replaces its calibration product with the + compact refit; the full-pool solve it was cut from is byte-identical to + a standalone dense run on the same inputs, seed and epochs, so keeping + its evidence lets a size candidate be compared with its own dense + reference without a second run. Same vocabulary as the solve product: + labelled local and national diagnostics on the doctrine's scale rule, + the microcosm#492 past-cap censuses, the closing loss. + """ + + weights: np.ndarray + initial_weights: np.ndarray + diagnostics: pd.DataFrame + national_diagnostics: pd.DataFrame + initial_loss: float + final_loss: float + n_nonzero: int + past_cap_census: Mapping[str, Any] + national_past_cap_census: Mapping[str, Any] + all_past_cap_census: Mapping[str, Any] + + +@dataclass(frozen=True) +class _DoctrineSolveEvidence: + diagnostics: pd.DataFrame + national_diagnostics: pd.DataFrame + past_cap_census: Mapping[str, Any] + national_past_cap_census: Mapping[str, Any] + all_past_cap_census: Mapping[str, Any] + initial_loss: float def past_cap_census( @@ -1014,9 +1050,14 @@ def solve_uk_rowwise_weights_under_doctrine( l0_lambda: float = 0.0, budget_iters: int = 10, seed: int = 0, + selection_seed: int | None = None, ) -> UKRowwiseDoctrineSolve: """Solve rowwise household weights under the reviewed doctrine. + ``selection_seed`` (default ``seed``) seeds only the size selection — + the informed L0 search, the exact-count draw and the refit — so two + selections can be compared on one pool and one dense reference. + Structurally knob-free like before the ``calibrate()`` migration: no per-target parameters and no doctrine parameter — the bounds always come from :data:`UK_LOCAL_SOLVE_DOCTRINE` and ride into the public front door @@ -1147,6 +1188,7 @@ def solve_uk_rowwise_weights_under_doctrine( ) selected_support = None size_receipt = None + dense_result = None if dataset_households is not None: if target_records is not None or l0_lambda != 0: raise ValueError( @@ -1160,11 +1202,169 @@ def solve_uk_rowwise_weights_under_doctrine( households=dataset_households, epochs=epochs, learning_rate=learning_rate, - seed=seed, + seed=seed if selection_seed is None else selection_seed, ) + dense_result = result result = sized.result selected_support = sized.support size_receipt = sized.receipt + evidence = _doctrine_solve_evidence( + result, + target_set=target_set, + problem=problem, + national_rows=national_rows, + local_count=local_count, + target_loss_weights=target_loss_weights, + doctrine=doctrine, + ) + diagnostics = evidence.diagnostics + national_diagnostics = evidence.national_diagnostics + census = evidence.past_cap_census + national_census = evidence.national_past_cap_census + all_census = evidence.all_past_cap_census + initial_loss = evidence.initial_loss + dense_reference = None + if dense_result is not None: + dense_evidence = _doctrine_solve_evidence( + dense_result, + target_set=target_set, + problem=problem, + national_rows=national_rows, + local_count=local_count, + target_loss_weights=target_loss_weights, + doctrine=doctrine, + ) + dense_reference = UKRowwiseDenseReference( + weights=np.asarray(dense_result.weights, dtype=np.float64), + initial_weights=np.asarray(dense_result.initial_weights, dtype=np.float64), + diagnostics=dense_evidence.diagnostics, + national_diagnostics=dense_evidence.national_diagnostics, + initial_loss=dense_evidence.initial_loss, + final_loss=float(dense_result.final_loss), + n_nonzero=int(dense_result.n_nonzero), + past_cap_census=dense_evidence.past_cap_census, + national_past_cap_census=dense_evidence.national_past_cap_census, + all_past_cap_census=dense_evidence.all_past_cap_census, + ) + + # The UK carrier persists the weight column, and the kernel product is + # immutable with no table-refresh operation, so the finished frame is + # *rebuilt* through the canonical assembler from the kernel product's + # tables, mass log, and period. A rebuild can silently drop kernel + # surfaces the assembler does not carry, so the guards below refuse to + # ship if the kernel product held strata or metadata the rebuilt frame + # lost (both trivially equal today; the guard is the boundary marker + # for the day they are not). + if selected_support is None: + clean_result = result.frame if restore is None else restore(result.frame) + expected_ids = problem.household_ids + else: + # Restore the full prepared carrier before selecting: legacy restorers + # retain full original tables and cannot accept compact weight arrays. + clean_pool = frame if restore is None else restore(frame) + expected_ids = tuple(problem.household_ids[i] for i in selected_support) + clean_subset = clean_pool.select( + clean_pool.table("person")["person_household_id"] + .isin(expected_ids) + .to_numpy() + ) + clean_result = Frame( + {entity: clean_subset.table(entity) for entity in clean_subset.entities}, + clean_subset.schema, + { + entity: result.frame.weights_for(entity) + for entity in result.frame.weighted_entities + }, + clean_subset.strata, + mass_log=result.frame.mass_log, + metadata=result.frame.metadata, + ) + restored_ids = tuple(clean_result.table("household")["household_id"].tolist()) + if restored_ids != expected_ids: + raise ValueError( + "restore returned households that do not match the problem's rows " + "(same ids, same order); weights are written back by position, so a " + "reordering restore would misattribute them silently." + ) + slash_columns = { + entity: [ + str(column) + for column in clean_result.table(entity).columns + if "/" in str(column) + ] + for entity in clean_result.entities + } + slash_columns = { + entity: columns for entity, columns in slash_columns.items() if columns + } + if slash_columns: + raise ValueError( + "prepared slash-named measure columns survived the rowwise solve " + f"restore and cannot reach an H5 writer: {slash_columns}." + ) + calibrated_household = clean_result.table("household").copy() + calibrated_household["household_weight"] = np.asarray( + result.weights, dtype=np.float64 + ) + finished = uk_national_frame( + person=clean_result.table("person"), + benunit=clean_result.table("benunit"), + household=calibrated_household, + time_period=uk_time_period(clean_result), + weight_kind=WeightKind.CALIBRATED, + mass_log=clean_result.mass_log, + ) + validate_uk_national_frame(finished) + if not clean_result.strata.equals(finished.strata): + raise ValueError( + "the calibrated frame carries strata the rebuilt UK national " + "frame would drop; extend uk_national_frame to carry them " + "before shipping a strata-bearing local solve." + ) + if dict(clean_result.metadata) != dict(finished.metadata): + raise ValueError( + "the calibrated frame carries metadata beyond the UK time " + "period; the rebuild would drop it, so the solve refuses " + "instead." + ) + return UKRowwiseDoctrineSolve( + frame=finished, + calibration_result=result, + weights=np.asarray(result.weights, dtype=np.float64), + initial_weights=np.asarray(result.initial_weights, dtype=np.float64), + diagnostics=diagnostics, + national_diagnostics=national_diagnostics, + loss_trajectory=np.asarray(result.loss_trajectory, dtype=np.float64), + initial_loss=initial_loss, + final_loss=float(result.final_loss), + n_nonzero=int(result.n_nonzero), + past_cap_census=census, + national_past_cap_census=national_census, + all_past_cap_census=all_census, + binding_adjudications=binding_adjudications, + selected_support=selected_support, + size_receipt=size_receipt, + dense_reference=dense_reference, + ) + + +def _doctrine_solve_evidence( + result: CalibrationResult, + *, + target_set: TargetSet, + problem: UKRowwiseLocalMatrix, + national_rows: UKRowwiseNationalRows | None, + local_count: int, + target_loss_weights: np.ndarray, + doctrine: Any, +) -> _DoctrineSolveEvidence: + """Label one front-door result against the declared joint surface. + + Shared by the shipped solve and, on a size run, the dense reference it + was cut from: both must compile the whole surface, align by name, and + reproduce the declared target values exactly. + """ + if result.skipped: reasons = [ f"{skipped.target.name}: {skipped.reason}" for skipped in result.skipped[:5] @@ -1294,104 +1494,13 @@ def solve_uk_rowwise_weights_under_doctrine( target_loss_cap=doctrine.target_loss_cap, ) ) - - # The UK carrier persists the weight column, and the kernel product is - # immutable with no table-refresh operation, so the finished frame is - # *rebuilt* through the canonical assembler from the kernel product's - # tables, mass log, and period. A rebuild can silently drop kernel - # surfaces the assembler does not carry, so the guards below refuse to - # ship if the kernel product held strata or metadata the rebuilt frame - # lost (both trivially equal today; the guard is the boundary marker - # for the day they are not). - if selected_support is None: - clean_result = result.frame if restore is None else restore(result.frame) - expected_ids = problem.household_ids - else: - # Restore the full prepared carrier before selecting: legacy restorers - # retain full original tables and cannot accept compact weight arrays. - clean_pool = frame if restore is None else restore(frame) - expected_ids = tuple(problem.household_ids[i] for i in selected_support) - clean_subset = clean_pool.select( - clean_pool.table("person")["person_household_id"] - .isin(expected_ids) - .to_numpy() - ) - clean_result = Frame( - {entity: clean_subset.table(entity) for entity in clean_subset.entities}, - clean_subset.schema, - { - entity: result.frame.weights_for(entity) - for entity in result.frame.weighted_entities - }, - clean_subset.strata, - mass_log=result.frame.mass_log, - metadata=result.frame.metadata, - ) - restored_ids = tuple(clean_result.table("household")["household_id"].tolist()) - if restored_ids != expected_ids: - raise ValueError( - "restore returned households that do not match the problem's rows " - "(same ids, same order); weights are written back by position, so a " - "reordering restore would misattribute them silently." - ) - slash_columns = { - entity: [ - str(column) - for column in clean_result.table(entity).columns - if "/" in str(column) - ] - for entity in clean_result.entities - } - slash_columns = { - entity: columns for entity, columns in slash_columns.items() if columns - } - if slash_columns: - raise ValueError( - "prepared slash-named measure columns survived the rowwise solve " - f"restore and cannot reach an H5 writer: {slash_columns}." - ) - calibrated_household = clean_result.table("household").copy() - calibrated_household["household_weight"] = np.asarray( - result.weights, dtype=np.float64 - ) - finished = uk_national_frame( - person=clean_result.table("person"), - benunit=clean_result.table("benunit"), - household=calibrated_household, - time_period=uk_time_period(clean_result), - weight_kind=WeightKind.CALIBRATED, - mass_log=clean_result.mass_log, - ) - validate_uk_national_frame(finished) - if not clean_result.strata.equals(finished.strata): - raise ValueError( - "the calibrated frame carries strata the rebuilt UK national " - "frame would drop; extend uk_national_frame to carry them " - "before shipping a strata-bearing local solve." - ) - if dict(clean_result.metadata) != dict(finished.metadata): - raise ValueError( - "the calibrated frame carries metadata beyond the UK time " - "period; the rebuild would drop it, so the solve refuses " - "instead." - ) - return UKRowwiseDoctrineSolve( - frame=finished, - calibration_result=result, - weights=np.asarray(result.weights, dtype=np.float64), - initial_weights=np.asarray(result.initial_weights, dtype=np.float64), + return _DoctrineSolveEvidence( diagnostics=diagnostics, national_diagnostics=national_diagnostics, - loss_trajectory=np.asarray(result.loss_trajectory, dtype=np.float64), - initial_loss=initial_loss, - final_loss=float(result.final_loss), - n_nonzero=int(result.n_nonzero), past_cap_census=census, national_past_cap_census=national_census, all_past_cap_census=all_census, - binding_adjudications=binding_adjudications, - selected_support=selected_support, - size_receipt=size_receipt, + initial_loss=initial_loss, ) @@ -1432,6 +1541,7 @@ def rotated_uk_local_holdout( l0_lambda: float = 0.0, budget_iters: int = 10, solve_seed: int = 0, + selection_seed: int | None = None, ) -> dict[str, object]: """Run five local-row rotations with national rows fixed in training.""" @@ -1477,6 +1587,7 @@ def rotated_uk_local_holdout( l0_lambda=l0_lambda, budget_iters=budget_iters, seed=solve_seed, + selection_seed=selection_seed, ) held_targets = problem.targets[holdout_indices] held_estimates = np.asarray( diff --git a/packages/microcosm-build/tests/test_uk_local_rowwise.py b/packages/microcosm-build/tests/test_uk_local_rowwise.py index bf7a838bc..896b9b052 100644 --- a/packages/microcosm-build/tests/test_uk_local_rowwise.py +++ b/packages/microcosm-build/tests/test_uk_local_rowwise.py @@ -1301,6 +1301,52 @@ def restore(full): assert result.selected_support.tolist() == [0, 2] +def test_size_solve_keeps_its_dense_reference_and_selection_seed_moves_only_the_draw(): + frame = _clone_frame() + metrics = pd.DataFrame({"households": [1.0, 1.0, 1.0]}, index=[101, 102, 103]) + problem = build_uk_rowwise_local_matrix( + metrics, + _assigned(), + pd.DataFrame({"code": ["E001", "S001"], "households": [2.0, 1.0]}), + ) + common = dict( + bound_families=["census_households/constituency"], + epochs=2, + seed=7, + ) + dense = solve_uk_rowwise_weights_under_doctrine(frame, problem, **common) + sized = solve_uk_rowwise_weights_under_doctrine( + frame, problem, dataset_households=2, **common + ) + assert dense.dense_reference is None + reference = sized.dense_reference + assert reference is not None + # The reference is the standalone dense solve, byte for byte. + np.testing.assert_array_equal(reference.weights, dense.weights) + np.testing.assert_array_equal(reference.initial_weights, dense.initial_weights) + assert reference.final_loss == dense.final_loss + assert reference.final_loss == sized.size_receipt["dense_loss"] + assert reference.initial_loss == dense.initial_loss + assert reference.n_nonzero == dense.n_nonzero + pd.testing.assert_frame_equal(reference.diagnostics, dense.diagnostics) + pd.testing.assert_frame_equal( + reference.national_diagnostics, dense.national_diagnostics + ) + assert dict(reference.past_cap_census) == dict(dense.past_cap_census) + assert dict(reference.all_past_cap_census) == dict(dense.all_past_cap_census) + # The shipped product is the compact refit, not the reference. + assert sized.weights.size == 2 + assert reference.weights.size == 3 + # A selection seed moves the draw only: same pool, same dense reference. + other = solve_uk_rowwise_weights_under_doctrine( + frame, problem, dataset_households=2, selection_seed=11, **common + ) + np.testing.assert_array_equal(other.dense_reference.weights, reference.weights) + assert other.size_receipt["seed"] == 11 + assert sized.size_receipt["seed"] == 7 + assert other.size_receipt["dense_loss"] == sized.size_receipt["dense_loss"] + + def test_size_holdout_uses_compact_support_and_reselects_each_fold(monkeypatch): import microcosm.build.uk_runtime.local_rowwise as runtime diff --git a/packages/microcosm-build/tests/test_uk_rowwise_candidate.py b/packages/microcosm-build/tests/test_uk_rowwise_candidate.py index b4623de8c..15baf8f17 100644 --- a/packages/microcosm-build/tests/test_uk_rowwise_candidate.py +++ b/packages/microcosm-build/tests/test_uk_rowwise_candidate.py @@ -1855,6 +1855,8 @@ def test_size_candidate_exports_compact_links_and_cannot_claim_dense_release( "2", "--dataset-households", "300", + "--selection-seed", + "11", "--epochs", "2", "--skip-holdout", @@ -1884,6 +1886,82 @@ def test_size_candidate_exports_compact_links_and_cannot_claim_dense_release( assert set(persons.person_benunit_id) == set(benunits.benunit_id) assert len(_spool_rows(out)) == 1 + # The selection seed moves the draw only; the manifest records both seeds. + assert manifest["parameters"]["seed"] == 7 + assert manifest["parameters"]["selection_seed"] == 11 + assert size["seed"] == 11 + + # The dense solve the selection was cut from ships as evidence. + dense = size["dense_reference"] + assert dense["final_loss"] == size["dense_loss"] + assert dense["n_households"] == 416 + assert dense["weights"]["n_records"] == 416 + assert {"effective_sample_size", "max_to_median_positive_weight"} <= set( + dense["weights"] + ) + assert dense["diagnostics_file"] == builder.DENSE_REFERENCE_DIAGNOSTICS_FILENAME + outputs = manifest["outputs"] + dense_csv = out / builder.DENSE_REFERENCE_DIAGNOSTICS_FILENAME + selection_csv = out / builder.DATASET_SIZE_SELECTION_FILENAME + assert Path(outputs["dense_reference_diagnostics"]["path"]) == dense_csv.resolve() + assert Path(outputs["dataset_size_selection"]["path"]) == selection_csv.resolve() + assert ( + outputs["dataset_size_selection"]["sha256"] + == hashlib.sha256(selection_csv.read_bytes()).hexdigest() + ) + dense_rows = pd.read_csv(dense_csv) + assert len(dense_rows) == manifest["solve"]["n_targets"] + assert dense_rows.columns[0] == "grain" + assert {"target", "final_estimate", "abs_relative_error"} <= set(dense_rows.columns) + selection = pd.read_csv(selection_csv) + assert list(selection.columns) == [ + "pool_row_index", + "household_id", + "clone_index", + "design_weight", + "inclusion_probability", + "certainty", + "ht_baseline_weight", + "refit_weight", + ] + assert len(selection) == 300 + assert selection["pool_row_index"].is_unique + assert selection["pool_row_index"].max() < 416 + assert set(selection["household_id"]) == set(households.household_id) + assert (selection["refit_weight"] > 0).all() + assert (selection["design_weight"] > 0).all() + assert ( + int(selection["certainty"].sum()) + == size["selection_receipt"]["certainty_count"] + ) + assert int(selection["certainty"].sum()) == size["protected_carriers"] + + +def test_selection_seed_requires_a_dataset_size(tmp_path): + builder = _load_builder_module() + args = builder._parse_args( + [ + "--input-h5", + str(tmp_path / "spine.h5"), + "--ladder", + str(tmp_path / "ladder.npz"), + "--out", + str(tmp_path / "out"), + "--selection-seed", + "11", + ] + ) + with pytest.raises(ValueError, match="requires --dataset-households"): + builder._validate_cli_args(args) + + +def test_dense_candidate_manifest_has_no_size_sidecars(tmp_path): + builder = _load_builder_module() + paths = builder._output_paths(tmp_path, source_year=2024, calibration_year=2025) + assert paths["dense_reference"].name == builder.DENSE_REFERENCE_DIAGNOSTICS_FILENAME + assert paths["selection"].name == builder.DATASET_SIZE_SELECTION_FILENAME + assert builder._SIZE_RUN_ONLY_OUTPUTS == {"dense_reference", "selection"} + def test_size_cli_refuses_promotion_without_separate_certification(tmp_path): builder = _load_builder_module() diff --git a/tools/build_uk_rowwise_candidate.py b/tools/build_uk_rowwise_candidate.py index a7a0a4dd5..e455b7ec1 100644 --- a/tools/build_uk_rowwise_candidate.py +++ b/tools/build_uk_rowwise_candidate.py @@ -99,6 +99,7 @@ uk_local_target_surface, uk_support_limited_misses, uk_time_period, + uk_weight_summary, write_uk_calibration_diagnostics, write_uk_rowwise_dataset, ) @@ -135,6 +136,11 @@ AREA_SUPPORT_FILENAME = "area_support_summary.csv" PAST_CAP_FILENAME = "past_cap_census.json" LOCAL_REGISTRY_FILENAME = "local_target_registry.json" +DENSE_REFERENCE_DIAGNOSTICS_FILENAME = "dense_reference_diagnostics.csv" +DATASET_SIZE_SELECTION_FILENAME = "dataset_size_selection.csv" + +#: Outputs a run writes only when ``--dataset-households`` is set. +_SIZE_RUN_ONLY_OUTPUTS = frozenset({"dense_reference", "selection"}) _CONSERVE_MASS = False _TARGET_RECORDS: int | None = None @@ -603,6 +609,16 @@ def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: help="Dry-run only comma-separated candidate clone counts.", ) parser.add_argument("--seed", type=int, default=42) + parser.add_argument( + "--selection-seed", + type=int, + help=( + "Seed for the size selection only (informed L0 search, exact-count " + "draw, refit); defaults to --seed. The pool, ladder assignment and " + "dense reference stay on --seed, so two selections compare on one " + "pool. Requires --dataset-households." + ), + ) parser.add_argument( "--sample-fraction", type=float, @@ -993,6 +1009,7 @@ def _run_candidate( l0_lambda=_L0_LAMBDA, budget_iters=_BUDGET_ITERS, seed=args.seed, + selection_seed=args.selection_seed, ) _validate_solve_result(solve, problem=problem) append_phase(state, "solved") @@ -1100,6 +1117,7 @@ def _run_candidate( l0_lambda=_L0_LAMBDA, budget_iters=_BUDGET_ITERS, solve_seed=args.seed, + selection_seed=args.selection_seed, ) args._rotated_holdout = rotated_holdout @@ -2064,6 +2082,13 @@ def _write_output_bundle( ) write_uk_rowwise_dataset(candidate, staged["dataset"]) solve.diagnostics.to_csv(staged["diagnostics"], index=False) + if solve.dense_reference is not None: + _dense_reference_diagnostics_frame(solve).to_csv( + staged["dense_reference"], index=False + ) + _dataset_size_selection_frame(solve, problem=problem, clone=clone).to_csv( + staged["selection"], index=False + ) support = support.copy() support["support_below_floor"] = ( (support["assigned_households"] < 50) @@ -2128,6 +2153,15 @@ def _write_output_bundle( reported_path=output_paths["local_registry"], ), } + if solve.dense_reference is not None: + outputs["dense_reference_diagnostics"] = _artifact_info( + staged["dense_reference"], + reported_path=output_paths["dense_reference"], + ) + outputs["dataset_size_selection"] = _artifact_info( + staged["selection"], + reported_path=output_paths["selection"], + ) manifest = _manifest( args, candidate=candidate, @@ -2329,7 +2363,10 @@ def _manifest( "pool_households": int(problem.n_households), "dataset_size": None if solve.size_receipt is None - else dict(solve.size_receipt), + else { + **dict(solve.size_receipt), + "dense_reference": _dense_reference_summary(solve), + }, "initial_loss": float(solve.initial_loss), "final_loss": float(solve.final_loss), "max_abs_relative_error": float(abs_errors.max()), @@ -2454,6 +2491,9 @@ def _parameters(args: argparse.Namespace, *, source_year: int) -> dict[str, Any] "n_clones": int(args.n_clones), "dataset_households": args.dataset_households, "seed": int(args.seed), + "selection_seed": None + if args.dataset_households is None + else int(args.seed if args.selection_seed is None else args.selection_seed), "source_year": source_year, "source_lineage_modulus": args.source_lineage_modulus, "sample_fraction": float(args.sample_fraction), @@ -2498,6 +2538,91 @@ def _fit_by_family(diagnostics: pd.DataFrame) -> list[dict[str, object]]: return rows +def _dense_reference_summary(solve: UKRowwiseDoctrineSolve) -> dict[str, Any] | None: + """Manifest-sized evidence of the dense solve a size run was cut from.""" + + dense = solve.dense_reference + if dense is None: + return None + local_errors = dense.diagnostics["abs_relative_error"].to_numpy(dtype=np.float64) + national_errors = dense.national_diagnostics["abs_relative_error"].to_numpy( + dtype=np.float64 + ) + past_cap = dict(dense.past_cap_census or {}) + return { + "initial_loss": float(dense.initial_loss), + "final_loss": float(dense.final_loss), + "n_nonzero": int(dense.n_nonzero), + "n_households": int(dense.weights.size), + "max_abs_relative_error": float(local_errors.max()) + if local_errors.size + else None, + "median_abs_relative_error": float(np.median(local_errors)) + if local_errors.size + else None, + "national_max_abs_relative_error": float(national_errors.max()) + if national_errors.size + else None, + "past_cap": {key: int(past_cap[key]) for key in _PAST_CAP_COUNT_KEYS}, + "weights": uk_weight_summary(dense.weights), + "local_by_family": _fit_by_family(dense.diagnostics), + "national_by_family": _fit_by_family(dense.national_diagnostics), + "diagnostics_file": DENSE_REFERENCE_DIAGNOSTICS_FILENAME, + } + + +def _dense_reference_diagnostics_frame(solve: UKRowwiseDoctrineSolve) -> pd.DataFrame: + """Every target's dense-reference estimate, local rows then national rows.""" + + dense = solve.dense_reference + assert dense is not None + local = dense.diagnostics.copy() + local.insert(0, "grain", local["area_type"].astype(str)) + national = dense.national_diagnostics.copy() + national.insert(0, "grain", "national") + return pd.concat([local, national], ignore_index=True, sort=False) + + +def _dataset_size_selection_frame( + solve: UKRowwiseDoctrineSolve, + *, + problem: UKRowwiseLocalMatrix, + clone: UKLadderRowwiseDatasetResult, +) -> pd.DataFrame: + """One row per selected pool household: identity, design, draw and refit.""" + + dense = solve.dense_reference + receipt = solve.size_receipt + assert dense is not None and receipt is not None + support = np.asarray(solve.selected_support, dtype=np.int64) + household = clone.frame.table("household") + clone_column = ladder_clone_index_column("household") + ids = household["household_id"].to_numpy()[support] + expected = np.asarray([problem.household_ids[i] for i in support]) + if not np.array_equal(ids, expected): + raise RuntimeError( + "the cloned pool's household order does not match the solve's " + "matrix columns; the selection sidecar would misattribute rows." + ) + inclusion = np.asarray(receipt["inclusion_probabilities"], dtype=np.float64) + if inclusion.shape != support.shape: + raise RuntimeError("selection receipt inclusion probabilities are misaligned.") + return pd.DataFrame( + { + "pool_row_index": support, + "household_id": ids, + "clone_index": household[clone_column].to_numpy()[support] + if clone_column in household.columns + else np.zeros(support.size, dtype=np.int64), + "design_weight": dense.initial_weights[support], + "inclusion_probability": inclusion, + "certainty": inclusion >= 1.0, + "ht_baseline_weight": np.asarray(solve.initial_weights, dtype=np.float64), + "refit_weight": np.asarray(solve.weights, dtype=np.float64), + } + ) + + def _local_output_registry( problem: UKRowwiseLocalMatrix, *, @@ -2622,6 +2747,8 @@ def _validate_cli_args(args: argparse.Namespace) -> None: raise ValueError( "the joint registry path requires --input-sha256 and --ladder-sha256." ) + if args.selection_seed is not None and args.dataset_households is None: + raise ValueError("--selection-seed requires --dataset-households.") if args.dataset_households is not None: if args.dataset_households <= 0: raise ValueError("--dataset-households must be positive.") @@ -2709,6 +2836,8 @@ def _output_paths( "local_gates": out_dir / LOCAL_GATE_REPORT_FILENAME_TEMPLATE.format(calibration_year=calibration_year), "local_registry": out_dir / LOCAL_REGISTRY_FILENAME, + "dense_reference": out_dir / DENSE_REFERENCE_DIAGNOSTICS_FILENAME, + "selection": out_dir / DATASET_SIZE_SELECTION_FILENAME, } @@ -2749,12 +2878,16 @@ def _publish_staged_files( "past_cap", "calibration_diagnostics", "local_registry", + "dense_reference", + "selection", "manifest", ) published: list[Path] = [] succeeded = False try: for key in publish_order: + if key in _SIZE_RUN_ONLY_OUTPUTS and not staged[key].exists(): + continue destination = output_paths[key] if destination.exists(): raise FileExistsError( From 115bc9184cf23c15779b9ffa92a0febb6ef7d519 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mar=C3=ADa=20Juaristi?= <127882282+juaristi22@users.noreply.github.com> Date: Tue, 8 Sep 2026 14:33:27 +0200 Subject: [PATCH 03/32] Carry the exact-count feasibility measurement with the size refusal and receipt (#355) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The first licensed size rehearsal (55,000 of 792,690 at 100 epochs) was refused inside select_exact_k: "degenerate boundary mass … adjust pi_hi or k". With pi_hi=1.0 every gate below one stays in the boundary and its open probability is scaled to the remaining draw size m, so the design is feasible only when m * max(pi_boundary) <= sum(pi_boundary); the L0 budget search stops on the count of not-fully-closed gates, which sits above the open-probability mass while gates are only partly polarised. The refusal carried no numbers. refit_uk_dataset_size now measures selection_feasibility before the draw — certainties and boundary draw at pi_hi=1, boundary mass and largest boundary gate, feasibility, the largest household count feasible at pi_hi=1, the smallest feasible pi_hi on a fixed grid, gate quantiles and counts, the budget search's n_nonzero and lambda — and attaches it to the ValueError on refusal and to the size receipt on success. Nothing is clamped or promoted; the ruling on pi_hi stays a reviewed decision, now made from measured mass. Co-Authored-By: Claude Fable 5.1 --- .../build/uk_runtime/dataset_size.py | 118 +++++++++++++++++- .../tests/test_uk_local_rowwise.py | 68 ++++++++++ 2 files changed, 184 insertions(+), 2 deletions(-) diff --git a/packages/microcosm-build/src/microcosm/build/uk_runtime/dataset_size.py b/packages/microcosm-build/src/microcosm/build/uk_runtime/dataset_size.py index 1a61ff4f8..eb7457670 100644 --- a/packages/microcosm-build/src/microcosm/build/uk_runtime/dataset_size.py +++ b/packages/microcosm-build/src/microcosm/build/uk_runtime/dataset_size.py @@ -1,5 +1,6 @@ """Exact household-count UK candidates on a fixed, materialized target surface.""" +import json from dataclasses import dataclass, replace import numpy as np @@ -110,9 +111,24 @@ def refit_uk_dataset_size( raise RuntimeError("informed L0 returned no selection probabilities.") # Exact-one certainties are protected by the gates themselves; no top-k # ranking or post-hoc promotion of learned boundary scores is performed. - support, sampling, q = select_exact_k( - probabilities, households, pi_hi=1.0, seed=seed + feasibility = selection_feasibility( + probabilities, + households, + protected=init.protected, + n_nonzero=int(selection.n_nonzero), + l0_lambda=float(selection.l0_lambda), ) + try: + support, sampling, q = select_exact_k( + probabilities, households, pi_hi=1.0, seed=seed + ) + except ValueError as error: + # The draw refuses rather than clamps; carry the measured gate mass + # with the refusal so the ruling it needs can be made from the receipt. + raise ValueError( + f"{error} Selection feasibility at pi_hi=1.0: " + f"{json.dumps(feasibility, sort_keys=True)}" + ) from error support = assert_exact_k_support(support, households, pool_size=n) if not np.isin(np.flatnonzero(init.protected), support).all(): raise RuntimeError("exact-count selection lost a protected carrier.") @@ -148,6 +164,7 @@ def refit_uk_dataset_size( "seed": seed, "protected_carriers": int(init.protected.sum()), "selection_receipt": sampling, + "selection_feasibility": feasibility, "selection_l0_lambda": selection.l0_lambda, "selection_epochs": epochs, "refit_epochs": epochs, @@ -165,6 +182,103 @@ def refit_uk_dataset_size( ) +_FEASIBILITY_PI_HI_GRID = (0.999, 0.99, 0.98, 0.95, 0.9, 0.8, 0.7, 0.5) + + +def selection_feasibility( + probabilities: np.ndarray, + households: int, + *, + protected: np.ndarray, + n_nonzero: int, + l0_lambda: float, +) -> dict[str, object]: + """Measure whether an exact-count draw is feasible on these gate probabilities. + + ``select_exact_k`` with ``pi_hi=1.0`` keeps every gate below one in the + boundary and scales its open probabilities to the remaining draw size + ``m``; the scaled value of the largest boundary gate must not exceed one, + i.e. ``m * max(pi_boundary) <= sum(pi_boundary)``. The L0 budget search + stops on the count of not-fully-closed gates, which can sit well above the + open-probability mass when gates are only partly polarised, so the draw + can refuse. This records the measured mass and the two ways out — the + smallest certainty threshold on a fixed grid that makes the design + feasible, and the largest household count feasible at ``pi_hi=1.0`` — so + the refusal is a ruling with numbers, never a silent clamp. + """ + + pi = np.asarray(probabilities, dtype=np.float64) + protected_mask = np.asarray(protected, dtype=bool) + if pi.shape != protected_mask.shape: + raise ValueError("selection feasibility needs aligned probabilities and mask.") + certainty = pi >= 1.0 + boundary = pi[~certainty] + positive = boundary[boundary > 0.0] + m = int(households) - int(certainty.sum()) + boundary_mass = float(positive.sum()) if positive.size else 0.0 + boundary_max = float(positive.max()) if positive.size else 0.0 + feasible = m <= 0 or ( + positive.size >= m and boundary_max * m <= boundary_mass * (1.0 + 1e-12) + ) + max_feasible_k = ( + int(certainty.sum()) + int(np.floor(boundary_mass / boundary_max)) + if boundary_max > 0.0 + else int(certainty.sum()) + ) + scan: dict[str, object] = {} + smallest_feasible_pi_hi: float | None = 1.0 if feasible else None + for threshold in _FEASIBILITY_PI_HI_GRID: + certain_t = pi >= threshold + c_t = int(certain_t.sum()) + m_t = int(households) - c_t + boundary_t = pi[~certain_t] + positive_t = boundary_t[boundary_t > 0.0] + mass_t = float(positive_t.sum()) if positive_t.size else 0.0 + max_t = float(positive_t.max()) if positive_t.size else 0.0 + ok = m_t >= 0 and ( + m_t == 0 + or (positive_t.size >= m_t and max_t * m_t <= mass_t * (1.0 + 1e-12)) + ) + scan[f"{threshold:g}"] = { + "certainties": c_t, + "boundary_draw": m_t, + "boundary_mass": mass_t, + "boundary_max": max_t, + "feasible": bool(ok), + } + if ok and smallest_feasible_pi_hi is None: + smallest_feasible_pi_hi = threshold + quantiles = ( + { + f"p{q * 100:g}": float(np.quantile(pi, q)) + for q in (0.5, 0.9, 0.99, 0.999) + } + if pi.size + else {} + ) + return { + "requested_households": int(households), + "pool_households": int(pi.size), + "protected_carriers": int(protected_mask.sum()), + "certainties_at_pi_hi_1": int(certainty.sum()), + "boundary_draw": m, + "boundary_positive_gates": int(positive.size), + "boundary_mass": boundary_mass, + "boundary_max": boundary_max, + "feasible_at_pi_hi_1": bool(feasible), + "max_feasible_households_at_pi_hi_1": max_feasible_k, + "smallest_feasible_pi_hi_on_grid": smallest_feasible_pi_hi, + "pi_hi_scan": scan, + "pi_sum": float(pi.sum()), + "pi_quantiles": quantiles, + "gates_above": { + f"{t:g}": int((pi >= t).sum()) for t in (0.5, 0.9, 0.99, 0.999) + }, + "budget_search_n_nonzero": int(n_nonzero), + "selection_l0_lambda": float(l0_lambda), + } + + def _frozen_targets( frame: Frame, dense: CalibrationResult, support: np.ndarray ) -> TargetSet: diff --git a/packages/microcosm-build/tests/test_uk_local_rowwise.py b/packages/microcosm-build/tests/test_uk_local_rowwise.py index 896b9b052..35f9002aa 100644 --- a/packages/microcosm-build/tests/test_uk_local_rowwise.py +++ b/packages/microcosm-build/tests/test_uk_local_rowwise.py @@ -1347,6 +1347,74 @@ def test_size_solve_keeps_its_dense_reference_and_selection_seed_moves_only_the_ assert other.size_receipt["dense_loss"] == sized.size_receipt["dense_loss"] +def test_selection_feasibility_measures_boundary_mass_and_the_ways_out(): + from microcosm.build.uk_runtime.dataset_size import selection_feasibility + + # Two protected certainties, then a boundary whose largest gate (0.9) + # cannot be scaled to a draw of 5 out of mass 0.9 + 4 * 0.2 = 1.7. + pi = np.array([1.0, 1.0, 0.9, 0.2, 0.2, 0.2, 0.2, 0.0]) + protected = np.array([True, True, False, False, False, False, False, False]) + report = selection_feasibility( + pi, 7, protected=protected, n_nonzero=7, l0_lambda=0.5 + ) + assert report["certainties_at_pi_hi_1"] == 2 + assert report["boundary_draw"] == 5 + assert report["boundary_positive_gates"] == 5 + assert report["boundary_mass"] == pytest.approx(1.7) + assert report["boundary_max"] == pytest.approx(0.9) + assert report["feasible_at_pi_hi_1"] is False + # At pi_hi=1 the boundary supports floor(1.7 / 0.9) = 1 draw: k <= 3. + assert report["max_feasible_households_at_pi_hi_1"] == 3 + # Promoting the 0.9 gate to a certainty (pi_hi <= 0.9) leaves a draw of 4 + # over four equal 0.2 gates: feasible; 0.99 and above are not. + assert report["pi_hi_scan"]["0.99"]["feasible"] is False + assert report["pi_hi_scan"]["0.9"]["feasible"] is True + assert report["smallest_feasible_pi_hi_on_grid"] == 0.9 + assert report["gates_above"]["0.5"] == 3 + assert report["budget_search_n_nonzero"] == 7 + assert report["selection_l0_lambda"] == 0.5 + + feasible = selection_feasibility( + np.array([1.0, 0.5, 0.5, 0.5, 0.5]), + 3, + protected=np.array([True, False, False, False, False]), + n_nonzero=5, + l0_lambda=0.1, + ) + assert feasible["feasible_at_pi_hi_1"] is True + assert feasible["smallest_feasible_pi_hi_on_grid"] == 1.0 + + +def test_size_refusal_at_the_draw_carries_the_feasibility_numbers(monkeypatch): + import microcosm.build.uk_runtime.dataset_size as sizing + from microcosm.calibrate import Target, TargetSet, calibrate + + frame = _clone_frame() + dense = calibrate( + frame, + TargetSet( + [Target("count", "household", lambda f: np.ones(f.n("household")), 3)] + ), + epochs=2, + ) + + def refuse(*_args, **_kwargs): + raise ValueError("degenerate boundary mass: synthetic refusal.") + + monkeypatch.setattr(sizing, "select_exact_k", refuse) + with pytest.raises(ValueError) as caught: + sizing.refit_uk_dataset_size( + frame, dense, households=2, epochs=2, learning_rate=0.02, seed=7 + ) + message = str(caught.value) + assert message.startswith("degenerate boundary mass: synthetic refusal.") + assert "Selection feasibility at pi_hi=1.0:" in message + payload = json.loads(message.split("Selection feasibility at pi_hi=1.0: ", 1)[1]) + assert payload["requested_households"] == 2 + assert payload["pool_households"] == 3 + assert "pi_hi_scan" in payload + + def test_size_holdout_uses_compact_support_and_reselects_each_fold(monkeypatch): import microcosm.build.uk_runtime.local_rowwise as runtime From 544f957cc3878db2d836e977648ab0b92d958c88 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mar=C3=ADa=20Juaristi?= <127882282+juaristi22@users.noreply.github.com> Date: Tue, 8 Sep 2026 14:52:49 +0200 Subject: [PATCH 04/32] Evaluate a size candidate with one command: size_evaluation library, evaluate_uk_dataset_size steps 00-90, uk_fit_by_family lifted (#355) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every evaluation piece for a --dataset-households candidate now has one name and one place: - uk_runtime/size_evaluation.py — pure measurement over a run directory: load_run (tolerant of pre-#877 dense runs), run_acceptance (the size-run checklist, not_applicable on dense runs), fit_tables (by grain and family), weight_tables (Kish ESS, distinct sources, stretch vs the Horvitz–Thompson baseline and vs pool design from the selection sidecar or a spine, self-checked against the manifest), area_support_tables (per-grain floors and breaches), gate_table (six ids with criticality), paired_targets (join on name; wins/ties/losses; the reference's red rows tracked through family/area/metric), dense_reference_deltas (size-only effect from the run's own dense reference), frozen_vs_recomputed (national rows against the incumbent-surface evaluator), footprint, and summarize with PRE_REGISTERED_OUTCOMES_V1 flags. No subprocess, no engine, no licensed data in tests. - tools/evaluate_uk_dataset_size.py — the one command: steps 00-run-acceptance, 10-dense-reference, 20-vs-reference/