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

policyengine/outputs/aggregate.py

@@ -0,0 +1,124 @@
+from enum import Enum
+from typing import Any
+
+from policyengine.core import Output, Simulation
+
+
+class AggregateType(str, Enum):
+    SUM = "sum"
+    MEAN = "mean"
+    COUNT = "count"
+
+
+class Aggregate(Output):
+    simulation: Simulation
+    variable: str
+    aggregate_type: AggregateType
+    entity: str | None = None
+
+    filter_variable: str | None = None
+    filter_variable_eq: Any | None = None
+    filter_variable_leq: Any | None = None
+    filter_variable_geq: Any | None = None
+    filter_variable_describes_quantiles: bool = False
+
+    # Convenient quantile specification (alternative to describes_quantiles)
+    quantile: int | None = (
+        None  # Number of quantiles (e.g., 10 for deciles, 5 for quintiles)
+    )
+    quantile_eq: int | None = None  # Exact quantile (e.g., 3 for 3rd decile)
+    quantile_leq: int | None = (
+        None  # Maximum quantile (e.g., 5 for bottom 5 deciles)
+    )
+    quantile_geq: int | None = (
+        None  # Minimum quantile (e.g., 9 for top 2 deciles)
+    )
+
+    result: Any | None = None
+
+    def run(self):
+        # Convert quantile specification to describes_quantiles format
+        if self.quantile is not None:
+            self.filter_variable_describes_quantiles = True
+            if self.quantile_eq is not None:
+                # For a specific quantile, filter between (quantile-1)/n and quantile/n
+                self.filter_variable_geq = (
+                    self.quantile_eq - 1
+                ) / self.quantile
+                self.filter_variable_leq = self.quantile_eq / self.quantile
+            elif self.quantile_leq is not None:
+                self.filter_variable_leq = self.quantile_leq / self.quantile
+            elif self.quantile_geq is not None:
+                self.filter_variable_geq = (
+                    self.quantile_geq - 1
+                ) / self.quantile
+
+        # Get variable object
+        var_obj = next(
+            v
+            for v in self.simulation.tax_benefit_model_version.variables
+            if v.name == self.variable
+        )
+
+        # Get the target entity data
+        target_entity = self.entity or var_obj.entity
+        data = getattr(self.simulation.output_dataset.data, target_entity)
+
+        # Map variable to target entity if needed
+        if var_obj.entity != target_entity:
+            mapped = self.simulation.output_dataset.data.map_to_entity(
+                var_obj.entity, target_entity, columns=[self.variable]
+            )
+            series = mapped[self.variable]
+        else:
+            series = data[self.variable]
+
+        # Apply filters
+        if self.filter_variable is not None:
+            filter_var_obj = next(
+                v
+                for v in self.simulation.tax_benefit_model_version.variables
+                if v.name == self.filter_variable
+            )
+
+            if filter_var_obj.entity != target_entity:
+                filter_mapped = (
+                    self.simulation.output_dataset.data.map_to_entity(
+                        filter_var_obj.entity,
+                        target_entity,
+                        columns=[self.filter_variable],
+                    )
+                )
+                filter_series = filter_mapped[self.filter_variable]
+            else:
+                filter_series = data[self.filter_variable]
+
+            if self.filter_variable_describes_quantiles:
+                if self.filter_variable_eq is not None:
+                    threshold = filter_series.quantile(self.filter_variable_eq)
+                    series = series[filter_series <= threshold]
+                if self.filter_variable_leq is not None:
+                    threshold = filter_series.quantile(
+                        self.filter_variable_leq
+                    )
+                    series = series[filter_series <= threshold]
+                if self.filter_variable_geq is not None:
+                    threshold = filter_series.quantile(
+                        self.filter_variable_geq
+                    )
+                    series = series[filter_series >= threshold]
+            else:
+                if self.filter_variable_eq is not None:
+                    series = series[filter_series == self.filter_variable_eq]
+                if self.filter_variable_leq is not None:
+                    series = series[filter_series <= self.filter_variable_leq]
+                if self.filter_variable_geq is not None:
+                    series = series[filter_series >= self.filter_variable_geq]
+
+        # Aggregate - MicroSeries will automatically apply weights
+        if self.aggregate_type == AggregateType.SUM:
+            self.result = series.sum()
+        elif self.aggregate_type == AggregateType.MEAN:
+            self.result = series.mean()
+        elif self.aggregate_type == AggregateType.COUNT:
+            self.result = series.count()

policyengine/outputs/change_aggregate.py

@@ -0,0 +1,184 @@
+from enum import Enum
+from typing import Any
+
+from policyengine.core import Output, Simulation
+
+
+class ChangeAggregateType(str, Enum):
+    COUNT = "count"
+    SUM = "sum"
+    MEAN = "mean"
+
+
+class ChangeAggregate(Output):
+    baseline_simulation: Simulation
+    reform_simulation: Simulation
+    variable: str
+    aggregate_type: ChangeAggregateType
+    entity: str | None = None
+
+    # Filter by absolute change
+    change_geq: float | None = None  # Change >= value (e.g., gain >= 500)
+    change_leq: float | None = None  # Change <= value (e.g., loss <= -500)
+    change_eq: float | None = None  # Change == value
+
+    # Filter by relative change (as decimal, e.g., 0.05 = 5%)
+    relative_change_geq: float | None = None  # Relative change >= value
+    relative_change_leq: float | None = None  # Relative change <= value
+    relative_change_eq: float | None = None  # Relative change == value
+
+    # Filter by another variable (e.g., only count people with age >= 30)
+    filter_variable: str | None = None
+    filter_variable_eq: Any | None = None
+    filter_variable_leq: Any | None = None
+    filter_variable_geq: Any | None = None
+    filter_variable_describes_quantiles: bool = False
+
+    # Convenient quantile specification (alternative to describes_quantiles)
+    quantile: int | None = (
+        None  # Number of quantiles (e.g., 10 for deciles, 5 for quintiles)
+    )
+    quantile_eq: int | None = None  # Exact quantile (e.g., 3 for 3rd decile)
+    quantile_leq: int | None = (
+        None  # Maximum quantile (e.g., 5 for bottom 5 deciles)
+    )
+    quantile_geq: int | None = (
+        None  # Minimum quantile (e.g., 9 for top 2 deciles)
+    )
+
+    result: Any | None = None
+
+    def run(self):
+        # Convert quantile specification to describes_quantiles format
+        if self.quantile is not None:
+            self.filter_variable_describes_quantiles = True
+            if self.quantile_eq is not None:
+                # For a specific quantile, filter between (quantile-1)/n and quantile/n
+                self.filter_variable_geq = (
+                    self.quantile_eq - 1
+                ) / self.quantile
+                self.filter_variable_leq = self.quantile_eq / self.quantile
+            elif self.quantile_leq is not None:
+                self.filter_variable_leq = self.quantile_leq / self.quantile
+            elif self.quantile_geq is not None:
+                self.filter_variable_geq = (
+                    self.quantile_geq - 1
+                ) / self.quantile
+
+        # Get variable object
+        var_obj = next(
+            v
+            for v in self.baseline_simulation.tax_benefit_model_version.variables
+            if v.name == self.variable
+        )
+
+        # Get the target entity data
+        target_entity = self.entity or var_obj.entity
+        baseline_data = getattr(
+            self.baseline_simulation.output_dataset.data, target_entity
+        )
+        reform_data = getattr(
+            self.reform_simulation.output_dataset.data, target_entity
+        )
+
+        # Map variable to target entity if needed
+        if var_obj.entity != target_entity:
+            baseline_mapped = (
+                self.baseline_simulation.output_dataset.data.map_to_entity(
+                    var_obj.entity, target_entity
+                )
+            )
+            baseline_series = baseline_mapped[self.variable]
+
+            reform_mapped = (
+                self.reform_simulation.output_dataset.data.map_to_entity(
+                    var_obj.entity, target_entity
+                )
+            )
+            reform_series = reform_mapped[self.variable]
+        else:
+            baseline_series = baseline_data[self.variable]
+            reform_series = reform_data[self.variable]
+
+        # Calculate change (reform - baseline)
+        change_series = reform_series - baseline_series
+
+        # Calculate relative change (handling division by zero)
+        # Where baseline is 0, relative change is undefined; we'll mask these out if relative filters are used
+        import numpy as np
+
+        with np.errstate(divide="ignore", invalid="ignore"):
+            relative_change_series = change_series / baseline_series
+            relative_change_series = relative_change_series.replace(
+                [np.inf, -np.inf], np.nan
+            )
+
+        # Start with all rows
+        mask = baseline_series.notna()
+
+        # Apply absolute change filters
+        if self.change_eq is not None:
+            mask &= change_series == self.change_eq
+        if self.change_leq is not None:
+            mask &= change_series <= self.change_leq
+        if self.change_geq is not None:
+            mask &= change_series >= self.change_geq
+
+        # Apply relative change filters
+        if self.relative_change_eq is not None:
+            mask &= relative_change_series == self.relative_change_eq
+        if self.relative_change_leq is not None:
+            mask &= relative_change_series <= self.relative_change_leq
+        if self.relative_change_geq is not None:
+            mask &= relative_change_series >= self.relative_change_geq
+
+        # Apply filter_variable filters
+        if self.filter_variable is not None:
+            filter_var_obj = next(
+                v
+                for v in self.baseline_simulation.tax_benefit_model_version.variables
+                if v.name == self.filter_variable
+            )
+
+            if filter_var_obj.entity != target_entity:
+                filter_mapped = (
+                    self.baseline_simulation.output_dataset.data.map_to_entity(
+                        filter_var_obj.entity, target_entity
+                    )
+                )
+                filter_series = filter_mapped[self.filter_variable]
+            else:
+                filter_series = baseline_data[self.filter_variable]
+
+            if self.filter_variable_describes_quantiles:
+                if self.filter_variable_eq is not None:
+                    threshold = filter_series.quantile(self.filter_variable_eq)
+                    mask &= filter_series <= threshold
+                if self.filter_variable_leq is not None:
+                    threshold = filter_series.quantile(
+                        self.filter_variable_leq
+                    )
+                    mask &= filter_series <= threshold
+                if self.filter_variable_geq is not None:
+                    threshold = filter_series.quantile(
+                        self.filter_variable_geq
+                    )
+                    mask &= filter_series >= threshold
+            else:
+                if self.filter_variable_eq is not None:
+                    mask &= filter_series == self.filter_variable_eq
+                if self.filter_variable_leq is not None:
+                    mask &= filter_series <= self.filter_variable_leq
+                if self.filter_variable_geq is not None:
+                    mask &= filter_series >= self.filter_variable_geq
+
+        # Apply mask to get filtered data
+        filtered_change = change_series[mask]
+
+        # Aggregate
+        if self.aggregate_type == ChangeAggregateType.COUNT:
+            self.result = filtered_change.count()
+        elif self.aggregate_type == ChangeAggregateType.SUM:
+            self.result = filtered_change.sum()
+        elif self.aggregate_type == ChangeAggregateType.MEAN:
+            self.result = filtered_change.mean()

policyengine/outputs/decile_impact.py

@@ -0,0 +1,140 @@
+import pandas as pd
+from pydantic import ConfigDict
+
+from policyengine.core import Output, OutputCollection, Simulation
+
+
+class DecileImpact(Output):
+    """Single decile's impact from a policy reform - represents one database row."""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    baseline_simulation: Simulation
+    reform_simulation: Simulation
+    income_variable: str = "equiv_hbai_household_net_income"
+    entity: str | None = None
+    decile: int
+    quantiles: int = 10
+
+    # Results populated by run()
+    baseline_mean: float | None = None
+    reform_mean: float | None = None
+    absolute_change: float | None = None
+    relative_change: float | None = None
+    count_better_off: float | None = None
+    count_worse_off: float | None = None
+    count_no_change: float | None = None
+
+    def run(self):
+        """Calculate impact for this specific decile."""
+        # Get variable object to determine entity
+        var_obj = next(
+            v
+            for v in self.baseline_simulation.tax_benefit_model_version.variables
+            if v.name == self.income_variable
+        )
+
+        # Get target entity
+        target_entity = self.entity or var_obj.entity
+
+        # Get data from both simulations
+        baseline_data = getattr(
+            self.baseline_simulation.output_dataset.data, target_entity
+        )
+        reform_data = getattr(
+            self.reform_simulation.output_dataset.data, target_entity
+        )
+
+        # Map income variable to target entity if needed
+        if var_obj.entity != target_entity:
+            baseline_mapped = (
+                self.baseline_simulation.output_dataset.data.map_to_entity(
+                    var_obj.entity, target_entity
+                )
+            )
+            baseline_income = baseline_mapped[self.income_variable]
+
+            reform_mapped = (
+                self.reform_simulation.output_dataset.data.map_to_entity(
+                    var_obj.entity, target_entity
+                )
+            )
+            reform_income = reform_mapped[self.income_variable]
+        else:
+            baseline_income = baseline_data[self.income_variable]
+            reform_income = reform_data[self.income_variable]
+
+        # Calculate deciles based on baseline income
+        decile_series = (
+            pd.qcut(
+                baseline_income,
+                self.quantiles,
+                labels=False,
+                duplicates="drop",
+            )
+            + 1
+        )
+
+        # Calculate changes
+        absolute_change = reform_income - baseline_income
+        relative_change = (absolute_change / baseline_income) * 100
+
+        # Filter to this decile
+        mask = decile_series == self.decile
+
+        # Populate results
+        self.baseline_mean = float(baseline_income[mask].mean())
+        self.reform_mean = float(reform_income[mask].mean())
+        self.absolute_change = float(absolute_change[mask].mean())
+        self.relative_change = float(relative_change[mask].mean())
+        self.count_better_off = float((absolute_change[mask] > 0).sum())
+        self.count_worse_off = float((absolute_change[mask] < 0).sum())
+        self.count_no_change = float((absolute_change[mask] == 0).sum())
+
+
+def calculate_decile_impacts(
+    baseline_simulation: Simulation,
+    reform_simulation: Simulation,
+    income_variable: str = "equiv_hbai_household_net_income",
+    entity: str | None = None,
+    quantiles: int = 10,
+) -> OutputCollection[DecileImpact]:
+    """Calculate decile-by-decile impact of a reform.
+
+    Returns:
+        OutputCollection containing list of DecileImpact objects and DataFrame
+    """
+    results = []
+    for decile in range(1, quantiles + 1):
+        impact = DecileImpact(
+            baseline_simulation=baseline_simulation,
+            reform_simulation=reform_simulation,
+            income_variable=income_variable,
+            entity=entity,
+            decile=decile,
+            quantiles=quantiles,
+        )
+        impact.run()
+        results.append(impact)
+
+    # Create DataFrame
+    df = pd.DataFrame(
+        [
+            {
+                "baseline_simulation_id": r.baseline_simulation.id,
+                "reform_simulation_id": r.reform_simulation.id,
+                "income_variable": r.income_variable,
+                "decile": r.decile,
+                "baseline_mean": r.baseline_mean,
+                "reform_mean": r.reform_mean,
+                "absolute_change": r.absolute_change,
+                "relative_change": r.relative_change,
+                "count_better_off": r.count_better_off,
+                "count_worse_off": r.count_worse_off,
+                "count_no_change": r.count_no_change,
+            }
+            for r in results
+        ]
+    )
+
+    return OutputCollection(outputs=results, dataframe=df)

policyengine/tax_benefit_models/uk/__init__.py

@@ -0,0 +1,26 @@
+"""PolicyEngine UK tax-benefit model."""
+
+from .analysis import general_policy_reform_analysis
+from .datasets import PolicyEngineUKDataset, UKYearData, create_datasets
+from .model import PolicyEngineUK, PolicyEngineUKLatest, uk_latest, uk_model
+from .outputs import ProgrammeStatistics
+
+__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/uk/analysis.py

@@ -0,0 +1,97 @@
+"""General utility functions for UK policy reform analysis."""
+
+import pandas as pd
+from pydantic import BaseModel
+
+from policyengine.core import OutputCollection, Simulation
+from policyengine.outputs.decile_impact import (
+    DecileImpact,
+    calculate_decile_impacts,
+)
+
+from .outputs import ProgrammeStatistics
+
+
+class PolicyReformAnalysis(BaseModel):
+    """Complete policy reform analysis result."""
+
+    decile_impacts: OutputCollection[DecileImpact]
+    programme_statistics: OutputCollection[ProgrammeStatistics]
+
+
+def general_policy_reform_analysis(
+    baseline_simulation: Simulation,
+    reform_simulation: Simulation,
+) -> PolicyReformAnalysis:
+    """Perform comprehensive analysis of a policy reform.
+
+    Returns:
+        PolicyReformAnalysis containing decile impacts and programme statistics
+    """
+    # Decile impact
+    decile_impacts = calculate_decile_impacts(
+        baseline_simulation=baseline_simulation,
+        reform_simulation=reform_simulation,
+    )
+
+    # Major programmes to analyse
+    programmes = {
+        # Tax
+        "income_tax": {"entity": "person", "is_tax": True},
+        "national_insurance": {"entity": "person", "is_tax": True},
+        "vat": {"entity": "household", "is_tax": True},
+        "council_tax": {"entity": "household", "is_tax": True},
+        # Benefits
+        "universal_credit": {"entity": "person", "is_tax": False},
+        "child_benefit": {"entity": "person", "is_tax": False},
+        "pension_credit": {"entity": "person", "is_tax": False},
+        "income_support": {"entity": "person", "is_tax": False},
+        "working_tax_credit": {"entity": "person", "is_tax": False},
+        "child_tax_credit": {"entity": "person", "is_tax": False},
+    }
+
+    programme_statistics = []
+
+    for programme_name, programme_info in programmes.items():
+        entity = programme_info["entity"]
+        is_tax = programme_info["is_tax"]
+
+        stats = ProgrammeStatistics(
+            baseline_simulation=baseline_simulation,
+            reform_simulation=reform_simulation,
+            programme_name=programme_name,
+            entity=entity,
+            is_tax=is_tax,
+        )
+        stats.run()
+        programme_statistics.append(stats)
+
+    # Create DataFrame
+    programme_df = pd.DataFrame(
+        [
+            {
+                "baseline_simulation_id": p.baseline_simulation.id,
+                "reform_simulation_id": p.reform_simulation.id,
+                "programme_name": p.programme_name,
+                "entity": p.entity,
+                "is_tax": p.is_tax,
+                "baseline_total": p.baseline_total,
+                "reform_total": p.reform_total,
+                "change": p.change,
+                "baseline_count": p.baseline_count,
+                "reform_count": p.reform_count,
+                "winners": p.winners,
+                "losers": p.losers,
+            }
+            for p in programme_statistics
+        ]
+    )
+
+    programme_collection = OutputCollection(
+        outputs=programme_statistics, dataframe=programme_df
+    )
+
+    return PolicyReformAnalysis(
+        decile_impacts=decile_impacts,
+        programme_statistics=programme_collection,
+    )