policyengine-uk 2.40.1__py3-none-any.whl → 2.65.6__py3-none-any.whl

policyengine_uk/__init__.py

@@ -1,6 +1,10 @@
 # -*- coding: utf-8 -*-
 
 import os
+from pathlib import Path
+
+from policyengine_core.taxbenefitsystems import TaxBenefitSystem
+
 from policyengine_uk import entities
 from policyengine_uk.system import (
     CountryTaxBenefitSystem,
@@ -10,8 +14,6 @@ from policyengine_uk.system import (
     parameters,
     variables,
 )
-from pathlib import Path
-import os
-from policyengine_core.taxbenefitsystems import TaxBenefitSystem
+from .model_api import *
 
 REPO = Path(__file__).parent

policyengine_uk/data/__init__.py

@@ -2,3 +2,4 @@ from policyengine_uk.data.dataset_schema import (
     UKMultiYearDataset,
     UKSingleYearDataset,
 )
+from .filter_dataset import filter_dataset

policyengine_uk/data/dataset_schema.py

@@ -14,11 +14,17 @@ class UKSingleYearDataset:
     household: pd.DataFrame
 
     @staticmethod
-    def validate_file_path(file_path: str):
+    def validate_file_path(file_path: str, raise_exception: bool = True):
         if not file_path.endswith(".h5"):
-            raise ValueError("File path must end with '.h5' for UKDataset.")
+            if raise_exception:
+                raise ValueError(
+                    "File path must end with '.h5' for UKDataset."
+                )
+            return False
         if not Path(file_path).exists():
-            raise FileNotFoundError(f"File not found: {file_path}")
+            if raise_exception:
+                raise FileNotFoundError(f"File not found: {file_path}")
+            return False
 
         # Check if the file contains time_period, person, benunit, and household datasets
         with h5py.File(file_path, "r") as f:
@@ -30,9 +36,14 @@ class UKSingleYearDataset:
             ]
             for dataset in required_datasets:
                 if dataset not in f:
-                    raise ValueError(
-                        f"Dataset '{dataset}' not found in the file: {file_path}"
-                    )
+                    if raise_exception:
+                        raise ValueError(
+                            f"Dataset '{dataset}' not found in the file: {file_path}"
+                        )
+                    else:
+                        return False
+
+        return True
 
     def __init__(
         self,
@@ -42,6 +53,7 @@ class UKSingleYearDataset:
         household: pd.DataFrame = None,
         fiscal_year: int = 2025,
     ):
+        file_path = str(file_path) if file_path else None
         if file_path is not None:
             self.validate_file_path(file_path)
             with pd.HDFStore(file_path) as f:
@@ -108,9 +120,12 @@ class UKSingleYearDataset:
                 if simulation.tax_benefit_system.variables[variable].entity.key
                 == entity
             ]
-            entity_dfs[entity] = simulation.calculate_dataframe(
-                input_variables, period=fiscal_year
-            )
+            if len(input_variables) == 0:
+                entity_dfs[entity] = pd.DataFrame()
+            else:
+                entity_dfs[entity] = simulation.calculate_dataframe(
+                    input_variables, period=fiscal_year
+                )
 
         return UKSingleYearDataset(
             person=entity_dfs["person"],
@@ -137,7 +152,7 @@ class UKMultiYearDataset:
                 self.datasets[year] = dataset
 
         if file_path is not None:
-            UKSingleYearDataset.validate_file_path(file_path)
+            UKMultiYearDataset.validate_file_path(file_path)
             with pd.HDFStore(file_path) as f:
                 self.datasets = {}
                 for year in f.keys():
@@ -162,6 +177,10 @@ class UKMultiYearDataset:
         else:
             raise ValueError(f"No dataset found for year {fiscal_year}.")
 
+    @property
+    def years(self):
+        return list(self.datasets.keys())
+
     def __getitem__(self, fiscal_year: int):
         return self.get_year(fiscal_year)
 
@@ -203,21 +222,49 @@ class UKMultiYearDataset:
         return UKMultiYearDataset(datasets=list(new_datasets.values()))
 
     @staticmethod
-    def validate_file_path(file_path: str):
+    def validate_file_path(file_path: str, raise_exception: bool = False):
         if not file_path.endswith(".h5"):
-            raise ValueError(
-                "File path must end with '.h5' for UKMultiYearDataset."
-            )
+            if raise_exception:
+                raise ValueError(
+                    "File path must end with '.h5' for UKMultiYearDataset."
+                )
+            else:
+                return False
         if not Path(file_path).exists():
-            raise FileNotFoundError(f"File not found: {file_path}")
+            if raise_exception:
+                raise FileNotFoundError(f"File not found: {file_path}")
+            else:
+                return False
 
         # Check if the file contains datasets for multiple years
         with h5py.File(file_path, "r") as f:
             for required_dataset in ["person", "benunit", "household"]:
                 if not any(f"{required_dataset}" in key for key in f.keys()):
-                    raise ValueError(
-                        f"Dataset '{required_dataset}' not found in the file: {file_path}"
-                    )
+                    if raise_exception:
+                        raise ValueError(
+                            f"Dataset '{required_dataset}' not found in the file: {file_path}"
+                        )
+                    else:
+                        return False
+
+                # Check that there is at least one dataset year in the folder (keys include e.g. person/2025)
+
+                # Check that there is at least one dataset year in the folder
+                years_found = False
+                for key in f.keys():
+                    parts = key.split("/")
+                    if len(parts) >= 2 and required_dataset == parts[0]:
+                        years_found = True
+                        break
+
+                if not years_found:
+                    if raise_exception:
+                        raise ValueError(
+                            f"No yearly data found for '{required_dataset}' in file: {file_path}"
+                        )
+                    else:
+                        return False
+        return True
 
     def load(self):
         data = {}
@@ -228,3 +275,8 @@ class UKMultiYearDataset:
                         data[col] = {}
                     data[col][year] = df[col].values
         return data
+
+    def reset_uprating(self):
+        from policyengine_uk.data.economic_assumptions import reset_uprating
+
+        reset_uprating(self)

policyengine_uk/data/economic_assumptions.py

@@ -1,14 +1,37 @@
 from policyengine_uk.data import UKMultiYearDataset, UKSingleYearDataset
-from policyengine_uk.system import system
+
 import yaml
 from policyengine_core.parameters import ParameterNode
 from pathlib import Path
 import numpy as np
+import logging
+
+
+def extend_single_year_dataset(
+    dataset: UKSingleYearDataset,
+    tax_benefit_system_parameters: ParameterNode,
+    end_year: int = 2030,
+) -> UKMultiYearDataset:
+    # Extend years and uprate
+    start_year = int(dataset.time_period)
+    datasets = [dataset]
+    for year in range(start_year, end_year + 1):
+        next_year = dataset.copy()
+        next_year.time_period = str(year)
+        datasets.append(next_year)
+    multi_year_dataset = UKMultiYearDataset(datasets=datasets)
+    return apply_uprating(
+        multi_year_dataset,
+        tax_benefit_system_parameters=tax_benefit_system_parameters,
+    )
 
 
 def apply_uprating(
     dataset: UKMultiYearDataset,
+    tax_benefit_system_parameters: ParameterNode = None,
 ):
+    from policyengine_uk.system import system
+
     # Apply uprating to the dataset.
     dataset = dataset.copy()
 
@@ -20,7 +43,9 @@ def apply_uprating(
             continue  # Don't uprate the first year
         current_year = dataset.datasets[year]
         prev_year = dataset.datasets[year - 1]
-        apply_single_year_uprating(current_year, prev_year, system.parameters)
+        apply_single_year_uprating(
+            current_year, prev_year, tax_benefit_system_parameters
+        )
 
     return dataset
 
@@ -48,7 +73,9 @@ def apply_single_year_uprating(
                     prev_year_value = getattr(previous_year, table_name)[
                         variable
                     ]
-                    current_year_value = prev_year_value * index_rel_change
+                    current_year_value = prev_year_value * (
+                        1 + index_rel_change
+                    )
                     getattr(current_year, table_name)[
                         variable
                     ] = current_year_value
@@ -128,9 +155,10 @@ def uprate_rent(
     social_rent_growth = growth.obr.social_rent(year)
 
     if year < 2022:
-        raise ValueError(
-            "Rent uprating is not supported for years before 2022."
+        logging.warning(
+            "Rent uprating is not supported for years before 2022. Not applying uprating."
         )
+        pass
     elif year < 2025:
         # We have regional growth rates for private rent.
         regional_growth_rate = growth.ons.private_rental_prices(year)[
@@ -158,11 +186,6 @@ def uprate_rent(
         private_rent_growth = (
             aggregate_growth - social_weight * social_rent_growth
         ) / private_weight
-        print(
-            f"Backed out private rent growth: {private_rent_growth:.1%} in {year}"
-        )
-        print(f"OBR aggregate rent growth: {aggregate_growth:.1%} in {year}")
-        print(f"Social rent growth: {social_rent_growth:.1%} in {year}")
 
         current_year.household["rent"] = np.where(
             is_private_rented,
@@ -182,3 +205,6 @@ def reset_uprating(
     for year in dataset.datasets:
         if year != first_year:
             dataset.datasets[year] = dataset.datasets[first_year].copy()
+            dataset.datasets[year].time_period = str(year)
+
+    return dataset

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,
+    )

policyengine_uk/dynamics/labour_supply.py

@@ -0,0 +1,343 @@
+"""Labour supply dynamics module for PolicyEngine UK.
+
+This module coordinates labour supply responses by combining progression (intensive margin)
+and participation (extensive margin) models. It implements the methodology from the OBR's
+labour supply elasticity framework to estimate how changes in tax and benefit policies
+affect employment and working hours.
+
+Reference: https://obr.uk/docs/dlm_uploads/NICS-Cut-Impact-on-Labour-Supply-Note.pdf
+"""
+
+import numpy as np
+import pandas as pd
+from policyengine_uk import Simulation
+from microdf import MicroDataFrame
+from pydantic import BaseModel
+from typing import Optional
+
+from .progression import (
+    calculate_derivative_change,
+    calculate_relative_income_change,
+    calculate_labour_substitution_elasticities,
+    calculate_labour_net_income_elasticities,
+    calculate_employment_income_change,
+)
+from .participation import (
+    apply_participation_responses,
+)
+
+
+def calculate_excluded_from_labour_supply_responses(
+    sim: Simulation, count_adults: int = 2
+):
+    """Calculate which individuals are excluded from labour supply responses.
+
+    Excludes self-employed, full-time students, aged 60+, and individuals
+    outside the specified adult range.
+
+    Args:
+        sim: PolicyEngine simulation object
+        count_adults: Number of adults to include in calculations
+
+    Returns:
+        Boolean array indicating which individuals are excluded
+    """
+    # Exclude self-employed, full-time students, aged 60+, and adult_index == (0, >= count_adults + 1)
+    employment_status = sim.calculate("employment_status")
+    self_employed = np.isin(
+        employment_status, ["FT_SELF_EMPLOYED", "PT_SELF_EMPLOYED"]
+    )
+    student = employment_status == "STUDENT"
+    age = sim.calculate("age")
+    age_60_plus = age >= 60
+    adult_index = sim.calculate("adult_index")
+    excluded = (
+        self_employed
+        | student
+        | age_60_plus
+        | (adult_index == 0)
+        | (adult_index >= count_adults + 1)
+    )
+    return excluded.values
+
+
+class FTEImpacts(BaseModel):
+    """Data model for FTE impacts of labour supply responses."""
+
+    substitution_response_ftes: float
+    """FTE impact from substitution effects."""
+    income_response_ftes: float
+    """FTE impact from income effects."""
+    total_response_ftes: float
+    """Total FTE impact from both substitution and income effects."""
+
+    participation_response_employment: Optional[float] = None
+    """FTE impact from participation responses, if available."""
+    participation_response_ftes: Optional[float] = None
+    """FTE impact from participation responses, if available."""
+
+    ftes: Optional[float] = None
+    """Total FTE impact across all responses."""
+
+
+class LabourSupplyResponseData(BaseModel):
+    progression: pd.DataFrame
+    """DataFrame containing intensive margin (progression) responses."""
+    participation: Optional[pd.DataFrame] = None
+    """DataFrame containing extensive margin (participation) responses, if available."""
+
+    # Specific outputs for comparison with OBR outputs
+    fte_impacts: FTEImpacts
+    """FTE impacts of labour supply responses, including both progression and participation."""
+
+    model_config = {
+        "arbitrary_types_allowed": True,
+    }
+
+
+def apply_labour_supply_responses(
+    sim: Simulation,
+    target_variable: str = "hbai_household_net_income",
+    input_variable: str = "employment_income",
+    year: int = 2025,
+    count_adults: int = 2,
+    delta: float = 1_000,
+) -> pd.DataFrame:
+    """Apply labour supply responses to simulation and return the response vector.
+
+    This is the main function for applying dynamic labour supply responses to
+    a PolicyEngine simulation. It coordinates both intensive margin (progression)
+    and extensive margin (participation) responses to policy changes.
+
+    Args:
+        sim: PolicyEngine simulation object (should have baseline attribute)
+        target_variable: Variable that drives labour supply decisions
+        input_variable: Variable representing labour supply (typically employment_income)
+        year: Year for calculation
+        count_adults: Number of adults to calculate responses for
+        delta: Size of change for marginal rate calculation (£)
+
+    Returns:
+        DataFrame with labour supply response information
+    """
+    follow_obr = sim.tax_benefit_system.parameters.gov.dynamic.obr_labour_supply_assumptions(
+        year
+    )
+    if (not follow_obr) or (sim.baseline is None):
+        return
+
+    # Calculate income changes using household_net_income
+    baseline_income = sim.baseline.calculate(
+        target_variable, year, map_to="person"
+    )
+    reform_income = sim.calculate(target_variable, year, map_to="person")
+
+    baseline_income = baseline_income
+    reform_income = reform_income
+    baseline_income = pd.Series(baseline_income, index=reform_income.index)
+
+    # Calculate relative changes
+    income_rel_change = np.where(
+        baseline_income != 0,
+        (reform_income - baseline_income) / baseline_income,
+        0,
+    )
+
+    # Apply intensive margin responses (progression model)
+    progression_responses = apply_progression_responses(
+        sim=sim,
+        target_variable=target_variable,
+        input_variable=input_variable,
+        year=year,
+        count_adults=count_adults,
+        delta=delta,
+        pre_calculated_income_rel_change=income_rel_change,
+    )
+
+    # Apply extensive margin responses (participation model)
+    participation_responses = (
+        None  # = apply_participation_responses(sim=sim, year=year)
+    )
+
+    # Add FTE impacts to the response data
+    fte_impacts = FTEImpacts(
+        substitution_response_ftes=progression_responses[
+            "substitution_response_ftes"
+        ].sum(),
+        income_response_ftes=progression_responses[
+            "income_response_ftes"
+        ].sum(),
+        total_response_ftes=progression_responses["total_response_ftes"].sum(),
+        participation_response_employment=(
+            participation_responses["participation_change"].sum()
+            if participation_responses is not None
+            else None
+        ),
+        participation_response_ftes=(
+            participation_responses["participation_change_ftes"].sum()
+            if participation_responses is not None
+            else None
+        ),
+    )
+
+    fte_impacts.ftes = fte_impacts.total_response_ftes + (
+        fte_impacts.participation_response_ftes
+        if fte_impacts.participation_response_ftes is not None
+        else 0
+    )
+
+    # For now, return only progression responses since participation is placeholder
+    # TODO: Combine progression and participation responses when participation model is implemented
+    return LabourSupplyResponseData(
+        progression=progression_responses,
+        participation=participation_responses,
+        fte_impacts=fte_impacts,
+    )
+
+
+def apply_progression_responses(
+    sim: Simulation,
+    target_variable: str = "hbai_household_net_income",
+    input_variable: str = "employment_income",
+    year: int = 2025,
+    count_adults: int = 2,
+    delta: float = 1_000,
+    pre_calculated_income_rel_change: np.ndarray = None,
+) -> pd.DataFrame:
+    """Apply progression (intensive margin) labour supply responses.
+
+    This function handles the intensive margin of labour supply by calculating
+    how individuals adjust their working hours in response to policy changes.
+
+    Args:
+        sim: PolicyEngine simulation object (should have baseline attribute)
+        target_variable: Variable that drives labour supply decisions
+        input_variable: Variable representing labour supply (typically employment_income)
+        year: Year for calculation
+        count_adults: Number of adults to calculate responses for
+        delta: Size of change for marginal rate calculation (£)
+
+    Returns:
+        DataFrame with progression response information
+    """
+    # Calculate changes in marginal rates (drives substitution effects)
+    derivative_changes = calculate_derivative_change(
+        sim=sim,
+        target_variable=target_variable,
+        input_variable=input_variable,
+        year=year,
+        count_adults=count_adults,
+        delta=delta,
+    )
+
+    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(
+        "hours_worked", year
+    )
+    gross_wage = gross_wage.fillna(0).replace([np.inf, -np.inf], 0)
+    derivative_changes["wage_gross"] = gross_wage
+    derivative_changes["wage_baseline"] = (
+        gross_wage * derivative_changes["deriv_baseline"]
+    )
+    derivative_changes["wage_scenario"] = (
+        gross_wage * derivative_changes["deriv_scenario"]
+    )
+    derivative_changes["wage_rel_change"] = (
+        (
+            derivative_changes["wage_scenario"]
+            / derivative_changes["wage_baseline"]
+            - 1
+        )
+        .replace([np.inf, -np.inf, np.nan], 0)
+        .fillna(0)
+    )
+    derivative_changes["wage_abs_change"] = (
+        derivative_changes["wage_scenario"]
+        - derivative_changes["wage_baseline"]
+    )
+
+    # Calculate changes in income levels (drives income effects)
+    if pre_calculated_income_rel_change is not None:
+        n_people = len(sim.calculate("person_id", year))
+        income_changes = pd.DataFrame(
+            {
+                "baseline": np.zeros(n_people),
+                "scenario": np.zeros(n_people),
+                "rel_change": pre_calculated_income_rel_change,
+                "abs_change": np.zeros(n_people),
+            }
+        )
+    else:
+        income_changes = calculate_relative_income_change(
+            sim, target_variable, year
+        )
+
+    income_changes = income_changes.rename(
+        columns={col: f"income_{col}" for col in income_changes.columns}
+    )
+
+    df = pd.concat([derivative_changes, income_changes], axis=1).fillna(0)
+
+    # Get elasticity parameters by demographic group
+    substitution_elasticities = calculate_labour_substitution_elasticities(sim)
+    income_elasticities = calculate_labour_net_income_elasticities(sim)
+
+    df["income_elasticity"] = income_elasticities
+    df["substitution_elasticity"] = substitution_elasticities
+
+    # Get baseline employment income levels
+    employment_income = sim.calculate(input_variable, year)
+
+    df["employment_income"] = employment_income
+    df["hours_per_week"] = sim.calculate("hours_worked", year) / 52
+
+    # Calculate total labour supply response
+    response_df = calculate_employment_income_change(
+        employment_income=employment_income,
+        derivative_changes=derivative_changes,
+        income_changes=income_changes,
+        substitution_elasticities=substitution_elasticities,
+        income_elasticities=income_elasticities,
+    )
+
+    df = pd.concat([df, response_df], axis=1)
+
+    # Apply relative {substitution, income, total} changes to hours as well
+    # Apply relative changes to hours using the same factor for all response types
+    for response_type in [
+        "substitution_response",
+        "income_response",
+        "total_response",
+    ]:
+        df[f"{response_type}_ftes"] = (
+            df[response_type]
+            / df["employment_income"]
+            * df["hours_per_week"]
+            / 37.5
+        )
+
+    excluded = calculate_excluded_from_labour_supply_responses(
+        sim, count_adults=count_adults
+    )
+
+    for col in df.columns:
+        df.loc[excluded, col] = 0
+
+    df["excluded"] = excluded
+
+    response = response_df["total_response"].values
+
+    # Apply the labour supply response to the simulation
+    sim.reset_calculations()
+    sim.set_input(input_variable, year, employment_income + response)
+
+    weight = sim.calculate("household_weight", year, map_to="person").values
+
+    result = MicroDataFrame(df, weights=weight)
+
+    return result[~result.excluded].drop(columns=["excluded"])