Give UK simulations a working subsample

.github/test_outcome_allowlist.toml

@@ -19,6 +19,18 @@
 # Keep entries narrow. An entry that matches nothing in a run is reported as a
 # warning; remove it once its bug is fixed. Failures can never be accepted.
 
+[[allow]]
+nodeid = "policyengine_uk/tests/test_subsample.py::test_neutralized_inputs_and_model_values_survive_subsampling[[]*]"
+outcome = "xfailed"
+reason = "^needs policyengine-core holder resync$"
+why = "Issue #2254 / PR #2265 waits for core to resync restored holders after subsampling. Both ownership variants retain every assertion under strict xfail."
+
+[[allow]]
+nodeid = "policyengine_uk/tests/test_input_records.py::test_restoring_reform_can_replace_a_neutralized_input[[]subsampled-*]"
+outcome = "xfailed"
+reason = "^needs policyengine-core holder resync$"
+why = "Issue #2254 / PR #2265 waits for core to accept writes immediately after restoring a subsampled holder. Direct supplied/model controls remain unmarked."
+
 [[allow]]
 nodeid = "policyengine_uk/tests/test_uc_state_pension_properties.py::test_state_pension_does_not_change_partners_marriage_allowance"
 outcome = "xfailed"

changelog.d/fix-subsample-2254.fixed.md

@@ -0,0 +1,2 @@
+- Fixed `Simulation.subsample()` and `Microsimulation.subsample()`, which raised `TypeError: Simulation.build_from_dataset() missing 1 required positional argument: 'dataset'`. They now keep a weighted sample of households with every supplied weight total unchanged, and replace the baseline and `dataset` with copies holding the same households.
+- Preserved stored inputs and model-owned values for neutralised variables when subsampling. Restoring their definitions after subsampling awaits policyengine-core holder resync.

docs/book/usage/simulations.md

@@ -479,6 +479,21 @@ weighted_total = (unweighted_values * modified_weights).sum()
 print(f"Total with 10% higher weights: £{weighted_total / 1e9:.1f}bn")
 ```
 
+### Subsampling a microsimulation
+
+`subsample` cuts a simulation built from a dataset down to a weighted random sample of its households:
+
+```python
+sim = Microsimulation(
+    dataset="hf://policyengine/policyengine-uk-data/enhanced_frs_2022_23.h5"
+)
+sim.subsample(n=10_000, seed=0, time_period=2026)
+```
+
+It draws `n` households with replacement, each with probability proportional to its weight in `time_period` (with the default `quantize_weights=True`), and keeps every household drawn at least once, with its benefit units and people. So the sample holds at most `n` households. A household's new weight is proportional to the number of times it was drawn, and the total household weight of every year is unchanged. The values the simulation was given are kept for the sampled records; results are calculated again on the sample, so population estimates then carry sampling error.
+
+The simulation's `baseline` and `dataset` are replaced with copies cut down to the same households; the objects they held before are left as they were. The `dataset` holds the rows of the dataset the simulation was built from, so it holds whole years only and does not include values set with `set_input` after the simulation was built. The same `seed` always draws the same households.
+
 ## Debugging simulations
 
 When simulations don't behave as expected:

policyengine_uk/simulation.py

@@ -52,6 +52,7 @@
 )
 from policyengine_uk.utils.dependencies import get_variable_dependencies
 from policyengine_uk.utils.input_dataframe_export import input_dataframe
+from policyengine_uk.utils.subsample import subsample_simulation
 from policyengine_uk.utils.supplied_inputs import (
     SUPPLIED_INPUT_VARIABLES,
     copy_model_values,
@@ -1273,6 +1274,74 @@ def build_from_ids(
 
         self.build_from_populations(builder.populations)
 
+    def subsample(
+        self,
+        n: Optional[int] = None,
+        frac: Optional[float] = None,
+        seed=None,
+        time_period=None,
+        quantize_weights: bool = True,
+    ) -> "Simulation":
+        """Cut the simulation down to a weighted random sample of households.
+
+        Draws ``n`` households with replacement and keeps each household
+        drawn at least once, with its benefit units and people. So the sample
+        holds at most ``n`` households. Only households with a positive
+        weight in ``time_period`` can be drawn: with ``quantize_weights``
+        (the default), each with probability proportional to that weight;
+        without it, each with equal probability. A period before the first
+        one the simulation was given household weights for has none, so the
+        first given weights are used instead.
+
+        Every value the simulation was given (from its dataset or through
+        ``set_input``) is kept for the sampled records, for every period, and
+        so is every value the model stored as its own (``set_model_value``),
+        which stays the model's. So does the record of where each capital
+        gains input came from. Calculated values are dropped and recalculated
+        on the sample, and the tracer starts again.
+
+        Each supplied weight is replaced so that its total is unchanged in
+        every period: a household's new weight is proportional to the number
+        of times it was drawn, and also to its old weight without
+        ``quantize_weights``. A zero weight stays zero. With
+        ``quantize_weights``, a household's new weight in ``time_period``,
+        averaged over draws, is its old one; in another period that holds
+        only if the household's share of the total weight is the same there.
+
+        ``baseline`` is replaced with a copy subsampled to the same
+        households, which keeps its own policy and inputs. The baseline the
+        simulation held before is left as it was, since another simulation
+        (a clone of this one) can hold it too. ``dataset`` is replaced with
+        the rows of the dataset the simulation was built from for the
+        sampled records, with its weights rescaled in the same way. Like the
+        dataset before, it holds whole years only, and not values set with
+        ``set_input`` after the simulation was built.
+
+        Args:
+            n: The number of households to draw.
+            frac: The number to draw as a share of all households, used when
+                ``n`` is not given.
+            seed: A number or text that fixes the draw. The same seed always
+                draws the same households; without one, a fixed seed is used.
+            time_period: Draw households by their weight in this period.
+                Defaults to the default calculation period.
+            quantize_weights: If false, draw households with a positive
+                weight with equal probability, and make a household's new
+                weight proportional to its old weight times the number of
+                times it was drawn.
+
+        Returns:
+            This simulation, subsampled.
+        """
+        return subsample_simulation(
+            self,
+            n=n,
+            frac=frac,
+            seed=seed,
+            time_period=time_period,
+            quantize_weights=quantize_weights,
+        )
+
     def move_values(self, variable_donor: str, variable_target: str) -> None:
         """Move values from one variable to another across all branches.
 

policyengine_uk/tests/test_input_records.py

@@ -8,9 +8,20 @@
 input.
 """
 
+import numpy as np
+import pandas as pd
 import pytest
+from policyengine_core.periods import period as period_
+from policyengine_core.reforms import Reform
 
-from policyengine_uk import Simulation
+from policyengine_uk import Microsimulation, Simulation
+from policyengine_uk.data import UKSingleYearDataset
+from policyengine_uk.utils.supplied_inputs import (
+    is_model_value,
+    set_model_value,
+    stored_supplied_value,
+    supplied_input,
+)
 
 pytestmark = pytest.mark.usefixtures("cloned_uk_tax_benefit_system")
 
@@ -30,3 +41,95 @@ def test_deleting_after_a_rebuild_does_not_export_a_carried_over_input():
     # The engine carries 2025's hours into 2026; they are not an input.
     assert sim.calculate("hours_worked", 2026).tolist() == [1_000]
     assert sim.to_input_dict()["hours_worked"] == {"2025": [1_000]}
+
+
+@pytest.mark.parametrize("model", [False, True], ids=["supplied", "model"])
+@pytest.mark.parametrize(
+    "subsampled",
+    [
+        pytest.param(False, id="direct"),
+        pytest.param(
+            True,
+            id="subsampled",
+            marks=pytest.mark.xfail(
+                strict=True,
+                raises=AssertionError,
+                reason="needs policyengine-core holder resync",
+            ),
+        ),
+    ],
+)
+def test_restoring_reform_can_replace_a_neutralized_input(model, subsampled):
+    """A write after restoration replaces the hidden value and its ownership."""
+    if subsampled:
+        dataset = UKSingleYearDataset(
+            person=pd.DataFrame(
+                {
+                    "person_id": [1],
+                    "person_benunit_id": [1],
+                    "person_household_id": [1],
+                    "age": [40],
+                    "employment_income": [30_000.0],
+                }
+            ),
+            benunit=pd.DataFrame({"benunit_id": [1]}),
+            household=pd.DataFrame(
+                {
+                    "household_id": [1],
+                    "household_weight": [1.0],
+                    "region": ["LONDON"],
+                    "tenure_type": ["OWNED_OUTRIGHT"],
+                    "rent": [0.0],
+                    "council_tax": [0.0],
+                }
+            ),
+            fiscal_year=2025,
+        )
+        simulation = Microsimulation(dataset=dataset)
+    else:
+        simulation = Simulation(situation={"age": 40, "employment_income": 30_000})
+    name = "adjusted_net_income"
+    period = period_(2025)
+    original_variable = simulation.tax_benefit_system.get_variable(name).__class__
+    if model:
+        set_model_value(simulation, name, period, [10_000])
+    else:
+        simulation.set_input(name, period, [10_000])
+
+    class NeutralizeIncome(Reform):
+        def apply(self):
+            self.neutralize_variable(name)
+
+    class RestoreAndReplaceIncome(Reform):
+        def apply(self):
+            self.replace_variable(original_variable)
+            assert not self.get_variable(name).is_neutralized
+            assert self.get_variable(name).__class__ is original_variable
+            self.simulation.set_input(name, period, [20_000])
+            stored = stored_supplied_value(
+                self.simulation.get_holder(name), period, self.simulation.branch_name
+            )
+            np.testing.assert_array_equal(stored, [20_000])
+            np.testing.assert_array_equal(
+                self.simulation.calculate(name, period), [20_000]
+            )
+            assert not is_model_value(self.simulation, name, period, stored)
+            np.testing.assert_array_equal(
+                supplied_input(self.simulation.populations["person"], name, period),
+                [20_000],
+            )
+
+    simulation.apply_reform(NeutralizeIncome)
+    if subsampled:
+        simulation.subsample(n=1, seed=0, time_period=2025)
+    simulation.apply_reform(RestoreAndReplaceIncome)
+
+    stored = stored_supplied_value(
+        simulation.get_holder(name), period, simulation.branch_name
+    )
+    np.testing.assert_array_equal(stored, [20_000])
+    np.testing.assert_array_equal(simulation.calculate(name, period), [20_000])
+    assert not is_model_value(simulation, name, period, stored)
+    np.testing.assert_array_equal(
+        supplied_input(simulation.populations["person"], name, period), [20_000]
+    )