Add dataset filter and improve labour supply response logic (#1348)

nikhilwoodruff · web-flow · commit 772d13fefa8c · 2025-09-30T14:46:01.000+01:00
* Add ability to zoom in on one household * Fix potential bug in LSRs!! * Improvements to LSR logic * Add dataset filter and improve labour supply response logic This PR adds the ability to filter datasets to single households for detailed analysis and fixes several issues in the labour supply response calculations. Closes #1349 Closes #1350 * Run formatter
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -5,6 +5,22 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), 
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [2.54.0] - 2025-09-30 14:40:31
+
+### Added
+
+- Dataset filter function to extract single households for analysis.
+
+### Changed
+
+- Default target variable for labour supply calculations to HBAI household net income.
+- Default adult count for labour supply calculations to 2.
+
+### Fixed
+
+- Labour supply response calculation issues by removing inappropriate clipping of marginal rates.
+- Potential bug in labour supply response variable initialization order.
+
 ## [2.53.1] - 2025-09-29 12:26:51
 
 ### Fixed
@@ -2248,6 +2264,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 
 
+[2.54.0]: https://github.com/PolicyEngine/openfisca-uk/compare/2.53.1...2.54.0
 [2.53.1]: https://github.com/PolicyEngine/openfisca-uk/compare/2.53.0...2.53.1
 [2.53.0]: https://github.com/PolicyEngine/openfisca-uk/compare/2.52.1...2.53.0
 [2.52.1]: https://github.com/PolicyEngine/openfisca-uk/compare/2.52.0...2.52.1
diff --git a/changelog.yaml b/changelog.yaml
@@ -1904,3 +1904,16 @@
     fixed:
     - Temporarily disabled OBR participation responses.
   date: 2025-09-29 12:26:51
+- bump: minor
+  changes:
+    added:
+    - Dataset filter function to extract single households for analysis.
+    changed:
+    - Default target variable for labour supply calculations to HBAI household net
+      income.
+    - Default adult count for labour supply calculations to 2.
+    fixed:
+    - Labour supply response calculation issues by removing inappropriate clipping
+      of marginal rates.
+    - Potential bug in labour supply response variable initialization order.
+  date: 2025-09-30 14:40:31
diff --git a/policyengine_uk/data/__init__.py b/policyengine_uk/data/__init__.py
@@ -2,3 +2,4 @@
     UKMultiYearDataset,
     UKSingleYearDataset,
 )
+from .filter_dataset import filter_dataset
diff --git a/policyengine_uk/data/filter_dataset.py b/policyengine_uk/data/filter_dataset.py
@@ -0,0 +1,52 @@
+import typing
+
+if typing.TYPE_CHECKING:
+    from policyengine_uk import Microsimulation
+    from policyengine_uk.data import UKSingleYearDataset
+
+
+def filter_dataset(
+    sim: "Microsimulation", household_id: int, year: int = 2026
+) -> "UKSingleYearDataset":
+    from policyengine_uk import Microsimulation
+    from policyengine_uk.data import UKSingleYearDataset
+
+    """
+    Extract a single household from a simulation dataset.
+
+    This function creates a new dataset containing only the specified household
+    and the associated benefit units and people within that household.
+
+    Parameters
+    ----------
+    sim : Microsimulation
+        The microsimulation object containing the dataset.
+    household_id : int
+        The ID of the household to extract.
+    year : int, default 2026
+        The dataset year to filter.
+
+    Returns
+    -------
+    UKSingleYearDataset
+        A new dataset containing only data for the specified household.
+    """
+    dataset: UKSingleYearDataset = sim.dataset[year]
+    new_dataset = dataset.copy()
+    new_dataset.person = new_dataset.person[
+        new_dataset.person.person_household_id == household_id
+    ]
+    new_dataset.household = new_dataset.household[
+        new_dataset.household.household_id == household_id
+    ]
+    benunits = new_dataset.person.person_benunit_id.unique()
+    new_dataset.benunit = new_dataset.benunit[
+        new_dataset.benunit.benunit_id.isin(benunits)
+    ]
+
+    return UKSingleYearDataset(
+        person=new_dataset.person,
+        household=new_dataset.household,
+        benunit=new_dataset.benunit,
+        fiscal_year=year,
+    )
diff --git a/policyengine_uk/dynamics/labour_supply.py b/policyengine_uk/dynamics/labour_supply.py
@@ -28,7 +28,7 @@
 
 
 def calculate_excluded_from_labour_supply_responses(
-    sim: Simulation, count_adults: int = 1
+    sim: Simulation, count_adults: int = 2
 ):
     """Calculate which individuals are excluded from labour supply responses.
 
@@ -97,10 +97,10 @@ class LabourSupplyResponseData(BaseModel):
 
 def apply_labour_supply_responses(
     sim: Simulation,
-    target_variable: str = "household_net_income",
+    target_variable: str = "hbai_household_net_income",
     input_variable: str = "employment_income",
     year: int = 2025,
-    count_adults: int = 1,
+    count_adults: int = 2,
     delta: float = 1_000,
 ) -> pd.DataFrame:
     """Apply labour supply responses to simulation and return the response vector.
@@ -197,10 +197,10 @@ def apply_labour_supply_responses(
 
 def apply_progression_responses(
     sim: Simulation,
-    target_variable: str = "household_net_income",
+    target_variable: str = "hbai_household_net_income",
     input_variable: str = "employment_income",
     year: int = 2025,
-    count_adults: int = 1,
+    count_adults: int = 2,
     delta: float = 1_000,
     pre_calculated_income_rel_change: np.ndarray = None,
 ) -> pd.DataFrame:
@@ -233,6 +233,7 @@ def apply_progression_responses(
     derivative_changes = derivative_changes.rename(
         columns={col: f"deriv_{col}" for col in derivative_changes.columns}
     )
+    derivative_changes["person_id"] = sim.calculate("person_id", year).values
 
     # Add in actual implied wages
     gross_wage = sim.calculate("employment_income", year) / sim.calculate(
@@ -262,20 +263,13 @@ def apply_progression_responses(
 
     # Calculate changes in income levels (drives income effects)
     if pre_calculated_income_rel_change is not None:
-        # Use pre-calculated values
         n_people = len(sim.calculate("person_id", year))
         income_changes = pd.DataFrame(
             {
-                "baseline": np.zeros(
-                    n_people
-                ),  # Not needed for behavioral response
-                "scenario": np.zeros(
-                    n_people
-                ),  # Not needed for behavioral response
+                "baseline": np.zeros(n_people),
+                "scenario": np.zeros(n_people),
                 "rel_change": pre_calculated_income_rel_change,
-                "abs_change": np.zeros(
-                    n_people
-                ),  # Not needed for behavioral response
+                "abs_change": np.zeros(n_people),
             }
         )
     else:
diff --git a/policyengine_uk/dynamics/progression.py b/policyengine_uk/dynamics/progression.py
@@ -14,7 +14,7 @@
 
 def calculate_derivative(
     sim: Simulation,
-    target_variable: str = "household_net_income",
+    target_variable: str = "hbai_household_net_income",
     input_variable: str = "employment_income",
     year: int = 2025,
     count_adults: int = 2,
@@ -72,12 +72,12 @@ def calculate_derivative(
     sim.set_input(input_variable, year, input_variable_values)
 
     # Clip to ensure rates are between 0 and 1 (0% to 100% retention)
-    return rel_marginal_wages.clip(0, 1)
+    return rel_marginal_wages.round(4)
 
 
 def calculate_relative_income_change(
     sim: Simulation,
-    target_variable: str = "household_net_income",
+    target_variable: str = "hbai_household_net_income",
     year: int = 2025,
 ) -> pd.DataFrame:
     """Calculate relative change in income between baseline and scenario.
@@ -125,10 +125,10 @@ def calculate_relative_income_change(
 
 def calculate_derivative_change(
     sim: Simulation,
-    target_variable: str = "household_net_income",
+    target_variable: str = "hbai_household_net_income",
     input_variable: str = "employment_income",
     year: int = 2025,
-    count_adults: int = 1,
+    count_adults: int = 2,
     delta: float = 1_000,
 ) -> pd.DataFrame:
     """Calculate change in marginal rates between baseline and scenario.
diff --git a/policyengine_uk/simulation.py b/policyengine_uk/simulation.py
@@ -104,8 +104,6 @@ def __init__(
         else:
             raise ValueError(f"Unsupported dataset type: {dataset.__class__}")
 
-        self.input_variables = self.get_known_variables()
-
         # Universal Credit reform (July 2025). Needs closer integration in the baseline,
         # but adding here for ease of toggling on/off via the 'active' parameter.
         from policyengine_uk.scenarios import universal_credit_july_2025_reform
@@ -119,6 +117,8 @@ def __init__(
         self.move_values("capital_gains", "capital_gains_before_response")
         self.move_values("employment_income", "employment_income_before_lsr")
 
+        self.input_variables = self.get_known_variables()
+
         if scenario is not None:
             self.baseline = Simulation(
                 scenario=None,
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [project] 
 name = "policyengine-uk"
-version = "2.53.1"
+version = "2.54.0"
 description = "PolicyEngine tax and benefit system for the UK."
 readme = "README.md"
 authors = [

Original file line number	Diff line number	Diff line change
`@@ -2,3 +2,4 @@`
`2`	`2`	`UKMultiYearDataset,`
`3`	`3`	`UKSingleYearDataset,`
`4`	`4`	`)`
	`5`	`+from .filter_dataset import filter_dataset`