policyengine 3.0.0__py3-none-any.whl → 3.1.1__py3-none-any.whl

policyengine/tax_benefit_models/uk/datasets.py

@@ -0,0 +1,176 @@
+from pathlib import Path
+
+import pandas as pd
+from microdf import MicroDataFrame
+from pydantic import BaseModel, ConfigDict
+
+from policyengine.core import Dataset, map_to_entity
+
+
+class UKYearData(BaseModel):
+    """Entity-level data for a single year."""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    person: MicroDataFrame
+    benunit: MicroDataFrame
+    household: MicroDataFrame
+
+    def map_to_entity(
+        self, source_entity: str, target_entity: str, columns: list[str] = None
+    ) -> MicroDataFrame:
+        """Map data from source entity to target entity using join keys.
+
+        Args:
+            source_entity (str): The source entity name ('person', 'benunit', 'household').
+            target_entity (str): The target entity name ('person', 'benunit', 'household').
+            columns (list[str], optional): List of column names to map. If None, maps all columns.
+
+        Returns:
+            MicroDataFrame: The mapped data at the target entity level.
+
+        Raises:
+            ValueError: If source or target entity is invalid.
+        """
+        entity_data = {
+            "person": self.person,
+            "benunit": self.benunit,
+            "household": self.household,
+        }
+        return map_to_entity(
+            entity_data=entity_data,
+            source_entity=source_entity,
+            target_entity=target_entity,
+            person_entity="person",
+            columns=columns,
+        )
+
+
+class PolicyEngineUKDataset(Dataset):
+    """UK dataset with multi-year entity-level data."""
+
+    data: UKYearData | None = None
+
+    def model_post_init(self, __context):
+        """Called after Pydantic initialization."""
+        # Make sure we are synchronised between in-memory and storage, at least on initialisation
+        if self.data is not None:
+            self.save()
+        elif self.filepath and not self.data:
+            try:
+                self.load()
+            except FileNotFoundError:
+                # File doesn't exist yet, that's OK
+                pass
+
+    def save(self) -> None:
+        """Save dataset to HDF5 file."""
+        filepath = Path(self.filepath)
+        if not filepath.parent.exists():
+            filepath.parent.mkdir(parents=True, exist_ok=True)
+        with pd.HDFStore(filepath, mode="w") as store:
+            store["person"] = pd.DataFrame(self.data.person)
+            store["benunit"] = pd.DataFrame(self.data.benunit)
+            store["household"] = pd.DataFrame(self.data.household)
+
+    def load(self) -> None:
+        """Load dataset from HDF5 file into this instance."""
+        filepath = self.filepath
+        with pd.HDFStore(filepath, mode="r") as store:
+            self.data = UKYearData(
+                person=MicroDataFrame(
+                    store["person"], weights="person_weight"
+                ),
+                benunit=MicroDataFrame(
+                    store["benunit"], weights="benunit_weight"
+                ),
+                household=MicroDataFrame(
+                    store["household"], weights="household_weight"
+                ),
+            )
+
+    def __repr__(self) -> str:
+        if self.data is None:
+            return f"<PolicyEngineUKDataset id={self.id} year={self.year} filepath={self.filepath} (not loaded)>"
+        else:
+            n_people = len(self.data.person)
+            n_benunits = len(self.data.benunit)
+            n_households = len(self.data.household)
+            return f"<PolicyEngineUKDataset id={self.id} year={self.year} filepath={self.filepath} people={n_people} benunits={n_benunits} households={n_households}>"
+
+
+def create_datasets(
+    datasets: list[str] = [
+        "hf://policyengine/policyengine-uk-data/frs_2023_24.h5",
+        "hf://policyengine/policyengine-uk-data/enhanced_frs_2023_24.h5",
+    ],
+    years: list[int] = [2026, 2027, 2028, 2029, 2030],
+) -> None:
+    for dataset in datasets:
+        from policyengine_uk import Microsimulation
+
+        sim = Microsimulation(dataset=dataset)
+        for year in years:
+            year_dataset = sim.dataset[year]
+
+            # Convert to pandas DataFrames and add weight columns
+            person_df = pd.DataFrame(year_dataset.person)
+            benunit_df = pd.DataFrame(year_dataset.benunit)
+            household_df = pd.DataFrame(year_dataset.household)
+
+            # Map household weights to person and benunit levels
+            person_df = person_df.merge(
+                household_df[["household_id", "household_weight"]],
+                left_on="person_household_id",
+                right_on="household_id",
+                how="left",
+            )
+            person_df = person_df.rename(
+                columns={"household_weight": "person_weight"}
+            )
+            person_df = person_df.drop(columns=["household_id"])
+
+            # Get household_id for each benunit from person table
+            benunit_household_map = person_df[
+                ["person_benunit_id", "person_household_id"]
+            ].drop_duplicates()
+            benunit_df = benunit_df.merge(
+                benunit_household_map,
+                left_on="benunit_id",
+                right_on="person_benunit_id",
+                how="left",
+            )
+            benunit_df = benunit_df.merge(
+                household_df[["household_id", "household_weight"]],
+                left_on="person_household_id",
+                right_on="household_id",
+                how="left",
+            )
+            benunit_df = benunit_df.rename(
+                columns={"household_weight": "benunit_weight"}
+            )
+            benunit_df = benunit_df.drop(
+                columns=[
+                    "person_benunit_id",
+                    "person_household_id",
+                    "household_id",
+                ],
+                errors="ignore",
+            )
+
+            uk_dataset = PolicyEngineUKDataset(
+                name=f"{dataset}-year-{year}",
+                description=f"UK Dataset for year {year} based on {dataset}",
+                filepath=f"./data/{Path(dataset).stem}_year_{year}.h5",
+                year=year,
+                data=UKYearData(
+                    person=MicroDataFrame(person_df, weights="person_weight"),
+                    benunit=MicroDataFrame(
+                        benunit_df, weights="benunit_weight"
+                    ),
+                    household=MicroDataFrame(
+                        household_df, weights="household_weight"
+                    ),
+                ),
+            )
+            uk_dataset.save()

policyengine/tax_benefit_models/uk/model.py

@@ -0,0 +1,268 @@
+import datetime
+from importlib.metadata import version
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import pandas as pd
+import requests
+from microdf import MicroDataFrame
+
+from policyengine.core import (
+    Parameter,
+    ParameterValue,
+    TaxBenefitModel,
+    TaxBenefitModelVersion,
+    Variable,
+)
+from policyengine.utils import parse_safe_date
+
+from .datasets import PolicyEngineUKDataset, UKYearData
+
+if TYPE_CHECKING:
+    from policyengine.core.simulation import Simulation
+
+
+class PolicyEngineUK(TaxBenefitModel):
+    id: str = "policyengine-uk"
+    description: str = "The UK's open-source dynamic tax and benefit microsimulation model maintained by PolicyEngine."
+
+
+uk_model = PolicyEngineUK()
+
+pkg_version = version("policyengine-uk")
+
+# Get published time from PyPI
+response = requests.get("https://pypi.org/pypi/policyengine-uk/json")
+data = response.json()
+upload_time = data["releases"][pkg_version][0]["upload_time_iso_8601"]
+
+
+class PolicyEngineUKLatest(TaxBenefitModelVersion):
+    model: TaxBenefitModel = uk_model
+    version: str = pkg_version
+    created_at: datetime.datetime = datetime.datetime.fromisoformat(
+        upload_time
+    )
+
+    def __init__(self, **kwargs: dict):
+        super().__init__(**kwargs)
+        from policyengine_core.enums import Enum
+        from policyengine_uk.system import system
+
+        self.id = f"{self.model.id}@{self.version}"
+
+        self.variables = []
+        for var_obj in system.variables.values():
+            variable = Variable(
+                id=self.id + "-" + var_obj.name,
+                name=var_obj.name,
+                tax_benefit_model_version=self,
+                entity=var_obj.entity.key,
+                description=var_obj.documentation,
+                data_type=var_obj.value_type
+                if var_obj.value_type is not Enum
+                else str,
+            )
+            if (
+                hasattr(var_obj, "possible_values")
+                and var_obj.possible_values is not None
+            ):
+                variable.possible_values = list(
+                    map(
+                        lambda x: x.name,
+                        var_obj.possible_values._value2member_map_.values(),
+                    )
+                )
+            self.variables.append(variable)
+
+        self.parameters = []
+        from policyengine_core.parameters import Parameter as CoreParameter
+
+        for param_node in system.parameters.get_descendants():
+            if isinstance(param_node, CoreParameter):
+                parameter = Parameter(
+                    id=self.id + "-" + param_node.name,
+                    name=param_node.name,
+                    tax_benefit_model_version=self,
+                    description=param_node.description,
+                    data_type=type(
+                        param_node(2025)
+                    ),  # Example year to infer type
+                    unit=param_node.metadata.get("unit"),
+                )
+                self.parameters.append(parameter)
+
+                for i in range(len(param_node.values_list)):
+                    param_at_instant = param_node.values_list[i]
+                    if i + 1 < len(param_node.values_list):
+                        next_instant = param_node.values_list[i + 1]
+                    else:
+                        next_instant = None
+                    parameter_value = ParameterValue(
+                        parameter=parameter,
+                        start_date=parse_safe_date(
+                            param_at_instant.instant_str
+                        ),
+                        end_date=parse_safe_date(next_instant.instant_str)
+                        if next_instant
+                        else None,
+                        value=param_at_instant.value,
+                    )
+                    self.parameter_values.append(parameter_value)
+
+    def run(self, simulation: "Simulation") -> "Simulation":
+        from policyengine_uk import Microsimulation
+        from policyengine_uk.data import UKSingleYearDataset
+
+        from policyengine.utils.parametric_reforms import (
+            simulation_modifier_from_parameter_values,
+        )
+
+        assert isinstance(simulation.dataset, PolicyEngineUKDataset)
+
+        dataset = simulation.dataset
+        dataset.load()
+        input_data = UKSingleYearDataset(
+            person=dataset.data.person,
+            benunit=dataset.data.benunit,
+            household=dataset.data.household,
+            fiscal_year=dataset.year,
+        )
+        microsim = Microsimulation(dataset=input_data)
+
+        if (
+            simulation.policy
+            and simulation.policy.simulation_modifier is not None
+        ):
+            simulation.policy.simulation_modifier(microsim)
+        elif simulation.policy:
+            modifier = simulation_modifier_from_parameter_values(
+                simulation.policy.parameter_values
+            )
+            modifier(microsim)
+
+        if (
+            simulation.dynamic
+            and simulation.dynamic.simulation_modifier is not None
+        ):
+            simulation.dynamic.simulation_modifier(microsim)
+        elif simulation.dynamic:
+            modifier = simulation_modifier_from_parameter_values(
+                simulation.dynamic.parameter_values
+            )
+            modifier(microsim)
+
+        # Allow custom variable selection, or use defaults
+        if simulation.variables is not None:
+            entity_variables = simulation.variables
+        else:
+            # Default comprehensive variable set
+            entity_variables = {
+                "person": [
+                    # IDs and weights
+                    "person_id",
+                    "benunit_id",
+                    "household_id",
+                    "person_weight",
+                    # Demographics
+                    "age",
+                    "gender",
+                    "is_adult",
+                    "is_SP_age",
+                    "is_child",
+                    # Income
+                    "employment_income",
+                    "self_employment_income",
+                    "pension_income",
+                    "private_pension_income",
+                    "savings_interest_income",
+                    "dividend_income",
+                    "property_income",
+                    "total_income",
+                    "earned_income",
+                    # Benefits
+                    "universal_credit",
+                    "child_benefit",
+                    "pension_credit",
+                    "income_support",
+                    "working_tax_credit",
+                    "child_tax_credit",
+                    # Tax
+                    "income_tax",
+                    "national_insurance",
+                ],
+                "benunit": [
+                    # IDs and weights
+                    "benunit_id",
+                    "benunit_weight",
+                    # Structure
+                    "family_type",
+                    # Income and benefits
+                    "universal_credit",
+                    "child_benefit",
+                    "working_tax_credit",
+                    "child_tax_credit",
+                ],
+                "household": [
+                    # IDs and weights
+                    "household_id",
+                    "household_weight",
+                    # Income measures
+                    "household_net_income",
+                    "hbai_household_net_income",
+                    "equiv_hbai_household_net_income",
+                    "household_market_income",
+                    "household_gross_income",
+                    # Benefits and tax
+                    "household_benefits",
+                    "household_tax",
+                    "vat",
+                    # Housing
+                    "rent",
+                    "council_tax",
+                    "tenure_type",
+                ],
+            }
+
+        data = {
+            "person": pd.DataFrame(),
+            "benunit": pd.DataFrame(),
+            "household": pd.DataFrame(),
+        }
+
+        for entity, variables in entity_variables.items():
+            for var in variables:
+                data[entity][var] = microsim.calculate(
+                    var, period=simulation.dataset.year, map_to=entity
+                ).values
+
+        data["person"] = MicroDataFrame(
+            data["person"], weights="person_weight"
+        )
+        data["benunit"] = MicroDataFrame(
+            data["benunit"], weights="benunit_weight"
+        )
+        data["household"] = MicroDataFrame(
+            data["household"], weights="household_weight"
+        )
+
+        simulation.output_dataset = PolicyEngineUKDataset(
+            name=dataset.name,
+            description=dataset.description,
+            filepath=str(
+                Path(simulation.dataset.filepath).parent
+                / (simulation.id + ".h5")
+            ),
+            year=simulation.dataset.year,
+            is_output_dataset=True,
+            data=UKYearData(
+                person=data["person"],
+                benunit=data["benunit"],
+                household=data["household"],
+            ),
+        )
+
+        simulation.output_dataset.save()
+
+
+uk_latest = PolicyEngineUKLatest()

policyengine/tax_benefit_models/uk/outputs.py

@@ -0,0 +1,108 @@
+"""UK-specific output templates."""
+
+from typing import TYPE_CHECKING
+
+from pydantic import ConfigDict
+
+from policyengine.core import Output
+from policyengine.outputs.aggregate import Aggregate, AggregateType
+from policyengine.outputs.change_aggregate import (
+    ChangeAggregate,
+    ChangeAggregateType,
+)
+
+if TYPE_CHECKING:
+    from policyengine.core.simulation import Simulation
+
+
+class ProgrammeStatistics(Output):
+    """Single programme's statistics from a policy reform - represents one database row."""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    baseline_simulation: "Simulation"
+    reform_simulation: "Simulation"
+    programme_name: str
+    entity: str
+    is_tax: bool = False
+
+    # Results populated by run()
+    baseline_total: float | None = None
+    reform_total: float | None = None
+    change: float | None = None
+    baseline_count: float | None = None
+    reform_count: float | None = None
+    winners: float | None = None
+    losers: float | None = None
+
+    def run(self):
+        """Calculate statistics for this programme."""
+        # Baseline totals
+        baseline_total = Aggregate(
+            simulation=self.baseline_simulation,
+            variable=self.programme_name,
+            aggregate_type=AggregateType.SUM,
+            entity=self.entity,
+        )
+        baseline_total.run()
+
+        # Reform totals
+        reform_total = Aggregate(
+            simulation=self.reform_simulation,
+            variable=self.programme_name,
+            aggregate_type=AggregateType.SUM,
+            entity=self.entity,
+        )
+        reform_total.run()
+
+        # Count of recipients/payers (baseline)
+        baseline_count = Aggregate(
+            simulation=self.baseline_simulation,
+            variable=self.programme_name,
+            aggregate_type=AggregateType.COUNT,
+            entity=self.entity,
+            filter_variable=self.programme_name,
+            filter_variable_geq=0.01,
+        )
+        baseline_count.run()
+
+        # Count of recipients/payers (reform)
+        reform_count = Aggregate(
+            simulation=self.reform_simulation,
+            variable=self.programme_name,
+            aggregate_type=AggregateType.COUNT,
+            entity=self.entity,
+            filter_variable=self.programme_name,
+            filter_variable_geq=0.01,
+        )
+        reform_count.run()
+
+        # Winners and losers
+        winners = ChangeAggregate(
+            baseline_simulation=self.baseline_simulation,
+            reform_simulation=self.reform_simulation,
+            variable=self.programme_name,
+            aggregate_type=ChangeAggregateType.COUNT,
+            entity=self.entity,
+            change_geq=0.01 if not self.is_tax else -0.01,
+        )
+        winners.run()
+
+        losers = ChangeAggregate(
+            baseline_simulation=self.baseline_simulation,
+            reform_simulation=self.reform_simulation,
+            variable=self.programme_name,
+            aggregate_type=ChangeAggregateType.COUNT,
+            entity=self.entity,
+            change_leq=-0.01 if not self.is_tax else 0.01,
+        )
+        losers.run()
+
+        # Populate results
+        self.baseline_total = float(baseline_total.result)
+        self.reform_total = float(reform_total.result)
+        self.change = float(reform_total.result - baseline_total.result)
+        self.baseline_count = float(baseline_count.result)
+        self.reform_count = float(reform_count.result)
+        self.winners = float(winners.result)
+        self.losers = float(losers.result)

policyengine/tax_benefit_models/uk.py

@@ -0,0 +1,33 @@
+"""PolicyEngine UK tax-benefit model - imports from uk/ module."""
+
+from .uk import (
+    PolicyEngineUK,
+    PolicyEngineUKDataset,
+    PolicyEngineUKLatest,
+    ProgrammeStatistics,
+    UKYearData,
+    create_datasets,
+    general_policy_reform_analysis,
+    uk_latest,
+    uk_model,
+)
+
+__all__ = [
+    "UKYearData",
+    "PolicyEngineUKDataset",
+    "create_datasets",
+    "PolicyEngineUK",
+    "PolicyEngineUKLatest",
+    "uk_model",
+    "uk_latest",
+    "general_policy_reform_analysis",
+    "ProgrammeStatistics",
+]
+
+# Rebuild models to resolve forward references
+from policyengine.core import Dataset
+
+Dataset.model_rebuild()
+UKYearData.model_rebuild()
+PolicyEngineUKDataset.model_rebuild()
+PolicyEngineUKLatest.model_rebuild()

policyengine/tax_benefit_models/us/__init__.py

@@ -0,0 +1,36 @@
+"""PolicyEngine US tax-benefit model."""
+
+from importlib.util import find_spec
+
+if find_spec("policyengine_us") is not None:
+    from policyengine.core import Dataset
+
+    from .analysis import general_policy_reform_analysis
+    from .datasets import PolicyEngineUSDataset, USYearData, create_datasets
+    from .model import (
+        PolicyEngineUS,
+        PolicyEngineUSLatest,
+        us_latest,
+        us_model,
+    )
+    from .outputs import ProgramStatistics
+
+    # Rebuild Pydantic models to resolve forward references
+    Dataset.model_rebuild()
+    USYearData.model_rebuild()
+    PolicyEngineUSDataset.model_rebuild()
+    PolicyEngineUSLatest.model_rebuild()
+
+    __all__ = [
+        "USYearData",
+        "PolicyEngineUSDataset",
+        "create_datasets",
+        "PolicyEngineUS",
+        "PolicyEngineUSLatest",
+        "us_model",
+        "us_latest",
+        "general_policy_reform_analysis",
+        "ProgramStatistics",
+    ]
+else:
+    __all__ = []