policyengine 0.1.0__py3-none-any.whl

policyengine/__init__.py

@@ -0,0 +1 @@
+from .simulation import Simulation, SimulationOptions

policyengine/constants.py

@@ -0,0 +1,14 @@
+"""Mainly simulation options and parameters."""
+
+# Datasets
+
+ENHANCED_FRS = "hf://policyengine/policyengine-uk-data/enhanced_frs_2022_23.h5"
+FRS = "hf://policyengine/policyengine-uk-data/frs_2022_23.h5"
+ENHANCED_CPS = "hf://policyengine/policyengine-us-data/enhanced_cps_2024.h5"
+CPS = "hf://policyengine/policyengine-us-data/cps_2023.h5"
+POOLED_CPS = "hf://policyengine/policyengine-us-data/pooled_3_year_cps_2023.h5"
+
+DEFAULT_DATASETS_BY_COUNTRY = {
+    "uk": ENHANCED_FRS,
+    "us": CPS,
+}

policyengine/outputs/household/comparison/calculate_household_comparison.py

@@ -0,0 +1,53 @@
+"""Calculate comparison statistics between two economic scenarios."""
+
+import typing
+
+from policyengine import Simulation
+
+from pydantic import BaseModel
+from policyengine.utils.calculations import get_change
+from policyengine_core.simulations import Simulation as CountrySimulation
+from policyengine.outputs.household.single.calculate_single_household import (
+    SingleHousehold,
+    fill_and_calculate,
+    FullHouseholdSpecification,
+)
+from typing import Literal, List
+
+
+class HouseholdComparison(BaseModel):
+    full_household_baseline: FullHouseholdSpecification
+    """The full completion of the household under the baseline scenario."""
+
+    full_household_reform: FullHouseholdSpecification
+    """The full completion of the household under the reform scenario."""
+
+    change: FullHouseholdSpecification
+    """The change in the household from the baseline to the reform scenario."""
+
+
+def calculate_household_comparison(
+    simulation: Simulation,
+) -> HouseholdComparison:
+    """Calculate comparison statistics between two household scenarios."""
+    if not simulation.is_comparison:
+        raise ValueError("Simulation must be a comparison simulation.")
+
+    baseline_household = fill_and_calculate(
+        simulation.options.data, simulation.baseline_simulation
+    )
+    reform_household = fill_and_calculate(
+        simulation.options.data, simulation.reform_simulation
+    )
+    change = get_change(
+        baseline_household,
+        reform_household,
+        relative=False,
+        skip_mismatch=True,
+    )
+
+    return HouseholdComparison(
+        full_household_baseline=baseline_household,
+        full_household_reform=reform_household,
+        change=change,
+    )

policyengine/outputs/household/single/calculate_single_household.py

@@ -0,0 +1,203 @@
+"""Calculate comparison statistics between two economic scenarios."""
+
+import typing
+
+from policyengine import Simulation
+
+from pydantic import BaseModel
+from policyengine_core.simulations import Simulation as CountrySimulation
+from typing import List, Dict
+from datetime import date
+from policyengine_core.variables import Variable
+from policyengine_core.entities import Entity
+from policyengine_core.model_api import YEAR, MONTH, ETERNITY, Enum
+import dpath.util
+import math
+import json
+
+Value = float | str | bool | list | None
+Axes = List[List[Dict[str, str | int]]]
+TimePeriodValues = Dict[str, Value]
+EntityValues = Dict[str, TimePeriodValues]
+EntityGroupValues = Dict[str, EntityValues]
+FullHouseholdSpecification = Dict[
+    str, EntityGroupValues | Axes
+]  # {people: {person: {variable: {time_period: value}}}}
+
+
+class SingleHousehold(BaseModel):
+    """Statistics for a single household scenario."""
+
+    full_household: FullHouseholdSpecification
+    """Full variable calculations for the household."""
+
+
+def calculate_single_household(
+    simulation: Simulation,
+) -> SingleHousehold:
+    """Calculate household statistics for a single household scenario."""
+    if simulation.is_comparison:
+        raise ValueError(
+            "This function is for single economy simulations only."
+        )
+
+    return SingleHousehold(
+        full_household=fill_and_calculate(
+            simulation.options.data, simulation.baseline_simulation
+        )
+    )
+
+
+def fill_and_calculate(
+    household: FullHouseholdSpecification, simulation: CountrySimulation
+):
+    """Fill in missing variables and calculate all variables for a household"""
+    # Copy the household to avoid modifying the original
+    household = json.loads(json.dumps(household))
+    household = add_yearly_variables(household, simulation)
+    household = calculate_all_variables(household, simulation)
+    household.pop("axes", None)
+    return household
+
+
+def get_requested_computations(
+    household: FullHouseholdSpecification,
+) -> List[tuple[str, str, str, str]]:
+    requested_computations = dpath.util.search(
+        {k: v for k, v in household.items() if k != "axes"},
+        "*/*/*/*",
+        # afilter=lambda t: t is None,
+        yielded=True,
+    )
+    requested_computation_data = []
+
+    for computation in requested_computations:
+        path = computation[0]
+        entity_plural, entity_id, variable_name, period = path.split("/")
+        requested_computation_data.append(
+            (entity_plural, entity_id, variable_name, period)
+        )
+
+    return requested_computation_data
+
+
+def calculate_all_variables(
+    household: FullHouseholdSpecification, simulation: CountrySimulation
+) -> FullHouseholdSpecification:
+    requested_computations = get_requested_computations(household)
+
+    for (
+        entity_plural,
+        entity_id,
+        variable_name,
+        period,
+    ) in requested_computations:
+        variable = simulation.tax_benefit_system.get_variable(variable_name)
+        result = simulation.calculate(variable_name, period)
+        population = simulation.get_population(entity_plural)
+
+        if "axes" in household:
+            count_entities = len(household[entity_plural])
+            entity_index = 0
+            for _entity_id in household[entity_plural].keys():
+                if _entity_id == entity_id:
+                    break
+                entity_index += 1
+            try:
+                result = result.astype(float)
+            except:
+                pass
+            result = (
+                result.reshape((-1, count_entities)).T[entity_index].tolist()
+            )
+            # If the result contains infinities, throw an error
+            if any(
+                [
+                    not isinstance(value, str) and math.isinf(value)
+                    for value in result
+                ]
+            ):
+                raise ValueError("Infinite value")
+            else:
+                household[entity_plural][entity_id][variable_name][
+                    period
+                ] = result
+        else:
+            entity_index = population.get_index(entity_id)
+            if variable.value_type == Enum:
+                entity_result = result.decode()[entity_index].name
+            elif variable.value_type == float:
+                entity_result = float(str(result[entity_index]))
+                # Convert infinities to JSON infinities
+                if entity_result == float("inf"):
+                    entity_result = "Infinity"
+                elif entity_result == float("-inf"):
+                    entity_result = "-Infinity"
+            elif variable.value_type == str:
+                entity_result = str(result[entity_index])
+            else:
+                entity_result = result.tolist()[entity_index]
+
+            household[entity_plural][entity_id][variable_name][
+                period
+            ] = entity_result
+
+    return household
+
+
+def get_household_year(household: FullHouseholdSpecification) -> str:
+    """Given a household dict, get the household's year
+
+    Args:
+        household (dict): The household itself
+    """
+
+    # Set household_year based on current year
+    household_year = date.today().year
+
+    # Determine if "age" variable present within household and return list of values at it
+    household_age_list = list(
+        household.get("people", {}).get("you", {}).get("age", {}).keys()
+    )
+    # If it is, overwrite household_year with the value present
+    if len(household_age_list) > 0:
+        household_year = household_age_list[0]
+
+    return str(household_year)
+
+
+def add_yearly_variables(
+    household: FullHouseholdSpecification, simulation: CountrySimulation
+) -> FullHouseholdSpecification:
+    """
+    Add yearly variables to a household dict before enqueueing calculation
+    """
+
+    variables: Dict[str, Variable] = simulation.tax_benefit_system.variables
+    entities: Dict[str, Entity] = (
+        simulation.tax_benefit_system.entities_by_singular()
+    )
+    household_year = get_household_year(household)
+
+    for variable in variables:
+        if variables[variable].definition_period in (YEAR, MONTH, ETERNITY):
+            entity_plural = entities[variables[variable].entity.key].plural
+            if entity_plural in household:
+                possible_entities = household[entity_plural].keys()
+                for entity in possible_entities:
+                    if (
+                        not variables[variable].name
+                        in household[entity_plural][entity]
+                    ):
+                        if variables[variable].is_input_variable():
+                            value = variables[variable].default_value
+                            if isinstance(value, Enum):
+                                value = value.name
+                            household[entity_plural][entity][
+                                variables[variable].name
+                            ] = {household_year: value}
+                        else:
+                            household[entity_plural][entity][
+                                variables[variable].name
+                            ] = {household_year: None}
+    return household

policyengine/outputs/macro/comparison/calculate_economy_comparison.py

@@ -0,0 +1,110 @@
+"""Calculate comparison statistics between two economic scenarios."""
+
+import typing
+
+from policyengine import Simulation
+
+from pydantic import BaseModel
+from policyengine.utils.calculations import get_change
+
+from policyengine.outputs.macro.single import (
+    _calculate_government_balance,
+    FiscalSummary,
+    _calculate_inequality,
+    InequalitySummary,
+)
+
+from .decile import calculate_decile_impacts, DecileImpacts
+
+from typing import Literal, List
+
+
+class FiscalComparison(BaseModel):
+    baseline: FiscalSummary
+    reform: FiscalSummary
+    change: FiscalSummary
+    relative_change: FiscalSummary
+
+
+class InequalityComparison(BaseModel):
+    baseline: InequalitySummary
+    reform: InequalitySummary
+    change: InequalitySummary
+    relative_change: InequalitySummary
+
+
+class Headlines(BaseModel):
+    budgetary_impact: float
+    """The change in the (federal) government budget balance."""
+    winner_share: float
+    """The share of people that are better off in the reform scenario."""
+
+
+class EconomyComparison(BaseModel):
+    headlines: Headlines
+    """Headline statistics for the comparison."""
+    fiscal: FiscalComparison
+    """Government budgets and other top-level fiscal statistics."""
+    inequality: InequalityComparison
+    """Inequality statistics for the household sector."""
+    distributional: DecileImpacts
+    """Distributional impacts of the reform."""
+
+
+def calculate_economy_comparison(
+    simulation: Simulation,
+) -> EconomyComparison:
+    """Calculate comparison statistics between two economic scenarios."""
+    if not simulation.is_comparison:
+        raise ValueError("Simulation must be a comparison simulation.")
+
+    baseline = simulation.baseline_simulation
+    reform = simulation.reform_simulation
+    options = simulation.options
+
+    baseline_balance = _calculate_government_balance(baseline, options)
+    reform_balance = _calculate_government_balance(reform, options)
+    balance_change = get_change(
+        baseline_balance, reform_balance, relative=False
+    )
+    balance_rel_change = get_change(
+        baseline_balance, reform_balance, relative=True
+    )
+    fiscal_comparison = FiscalComparison(
+        baseline=baseline_balance,
+        reform=reform_balance,
+        change=balance_change,
+        relative_change=balance_rel_change,
+    )
+
+    baseline_inequality = _calculate_inequality(baseline)
+    reform_inequality = _calculate_inequality(reform)
+    inequality_change = get_change(
+        baseline_inequality, reform_inequality, relative=False
+    )
+    inequality_rel_change = get_change(
+        baseline_inequality, reform_inequality, relative=True
+    )
+    inequality_comparison = InequalityComparison(
+        baseline=baseline_inequality,
+        reform=reform_inequality,
+        change=inequality_change,
+        relative_change=inequality_rel_change,
+    )
+
+    decile_impacts = calculate_decile_impacts(baseline, reform, options)
+
+    # Headlines
+    budgetary_impact = fiscal_comparison.change.federal_balance
+    winner_share = decile_impacts.income.winners_and_losers.all.gain_share
+    headlines = Headlines(
+        budgetary_impact=budgetary_impact,
+        winner_share=winner_share,
+    )
+
+    return EconomyComparison(
+        headlines=headlines,
+        fiscal=fiscal_comparison,
+        inequality=inequality_comparison,
+        distributional=decile_impacts,
+    )

policyengine/outputs/macro/comparison/charts/__init__.py

@@ -0,0 +1,5 @@
+from .budget import create_budget_comparison_chart
+from .budget_by_program import create_budget_program_comparison_chart
+from .decile import create_decile_chart
+from .winners_losers import create_winners_losers_chart
+from .inequality import create_inequality_chart

policyengine/outputs/macro/comparison/charts/budget.py

@@ -0,0 +1,90 @@
+import plotly.express as px
+import plotly.graph_objects as go
+import typing
+
+from policyengine import Simulation
+
+from pydantic import BaseModel
+from policyengine.utils.charts import *
+
+
+def create_budget_comparison_chart(
+    simulation: Simulation,
+) -> go.Figure:
+    """Create a budget comparison chart."""
+    if not simulation.is_comparison:
+        raise ValueError("Simulation must be a comparison simulation.")
+
+    economy = simulation.calculate_economy_comparison()
+
+    if simulation.options.country == "uk":
+        x_values = [
+            "Tax revenues",
+            "Government spending",
+            "Public sector net worth",
+        ]
+        y_values = [
+            economy.fiscal.change.federal_tax,
+            -economy.fiscal.change.government_spending,
+            economy.fiscal.change.federal_balance,
+        ]
+    else:
+        x_values = [
+            "Federal tax revenues",
+            "State tax revenues",
+            "Federal government spending",
+        ]
+        y_values = [
+            economy.fiscal.change.federal_tax,
+            economy.fiscal.change.state_tax,
+            -economy.fiscal.change.government_spending,
+        ]
+
+    y_values = [value / 1e9 for value in y_values]
+
+    net_change = round(economy.fiscal.change.federal_balance / 1e9, 1)
+
+    if net_change > 0:
+        description = f"raise ${net_change}bn"
+    elif net_change < 0:
+        description = f"cost ${-net_change}bn"
+    else:
+        description = "have no effect on government finances"
+
+    chart = go.Figure(
+        data=[
+            go.Waterfall(
+                x=x_values,
+                y=y_values,
+                measure=["relative"] * (len(x_values) - 1) + ["total"],
+                textposition="inside",
+                text=[f"${value:.1f}bn" for value in y_values],
+                increasing=dict(
+                    marker=dict(
+                        color=BLUE,
+                    )
+                ),
+                decreasing=dict(
+                    marker=dict(
+                        color=DARK_GRAY,
+                    )
+                ),
+                totals=dict(
+                    marker=dict(
+                        color=BLUE if net_change > 0 else DARK_GRAY,
+                    )
+                ),
+            ),
+        ]
+    ).update_layout(
+        title=f"{simulation.options.title} would {description}",
+        yaxis_title="Budgetary impact ($bn)",
+        uniformtext=dict(
+            mode="hide",
+            minsize=12,
+        ),
+    )
+
+    return format_fig(
+        chart, country=simulation.options.country, add_zero_line=True
+    )

policyengine/outputs/macro/comparison/charts/budget_by_program.py

@@ -0,0 +1,96 @@
+import plotly.express as px
+import plotly.graph_objects as go
+import typing
+
+from policyengine import Simulation
+
+from pydantic import BaseModel
+from policyengine.utils.charts import *
+
+
+def create_budget_program_comparison_chart(
+    simulation: Simulation,
+) -> go.Figure:
+    """Create a budget comparison chart."""
+    if not simulation.is_comparison:
+        raise ValueError("Simulation must be a comparison simulation.")
+
+    if not simulation.options.country == "uk":
+        raise ValueError("This chart is only available for the UK.")
+
+    economy = simulation.calculate_economy_comparison()
+
+    change_programs = economy.fiscal.change.tax_benefit_programs
+
+    change_programs = {
+        program: change_programs[program]
+        for program in change_programs
+        if round(change_programs[program] / 1e9, 1) != 0
+    }
+
+    labels = [
+        simulation.baseline_simulation.tax_benefit_system.variables.get(
+            program
+        ).label
+        for program in change_programs
+    ]
+
+    x_values = labels
+    y_values = [
+        round(change_programs[program] / 1e9, 1) for program in change_programs
+    ]
+
+    total_from_programs = round(sum(y_values))
+    total_overall = round(economy.fiscal.change.federal_balance / 1e9)
+
+    if total_from_programs != total_overall:
+        x_values.append("Other")
+        y_values.append(total_overall - total_from_programs)
+
+    x_values.append("Combined")
+    y_values.append(total_overall)
+
+    if total_overall > 0:
+        description = f"raise ${total_overall}bn"
+    elif total_overall < 0:
+        description = f"cost ${-total_overall}bn"
+    else:
+        description = "have no effect on government finances"
+
+    chart = go.Figure(
+        data=[
+            go.Waterfall(
+                x=x_values,
+                y=y_values,
+                measure=["relative"] * (len(x_values) - 1) + ["total"],
+                textposition="inside",
+                text=[f"${value:.1f}bn" for value in y_values],
+                increasing=dict(
+                    marker=dict(
+                        color=BLUE,
+                    )
+                ),
+                decreasing=dict(
+                    marker=dict(
+                        color=DARK_GRAY,
+                    )
+                ),
+                totals=dict(
+                    marker=dict(
+                        color=BLUE if total_overall > 0 else DARK_GRAY,
+                    )
+                ),
+            ),
+        ]
+    ).update_layout(
+        title=f"{simulation.options.title} would {description}",
+        yaxis_title="Budgetary impact (bn)",
+        uniformtext=dict(
+            mode="hide",
+            minsize=12,
+        ),
+    )
+
+    return format_fig(
+        chart, country=simulation.options.country, add_zero_line=True
+    )

policyengine/outputs/macro/comparison/charts/decile.py

@@ -0,0 +1,89 @@
+import plotly.express as px
+import plotly.graph_objects as go
+import typing
+
+from policyengine import Simulation
+
+from pydantic import BaseModel
+from policyengine.utils.charts import *
+from typing import Literal
+
+
+def create_decile_chart(
+    simulation: Simulation,
+    decile_variable: Literal["income", "wealth"],
+    relative: bool,
+) -> go.Figure:
+    """Create a budget comparison chart."""
+    if not simulation.is_comparison:
+        raise ValueError("Simulation must be a comparison simulation.")
+
+    economy = simulation.calculate_economy_comparison()
+
+    if decile_variable == "income":
+        data = economy.distributional.income.income_change
+    else:
+        data = economy.distributional.wealth.income_change
+
+    if relative:
+        data = data.relative
+    else:
+        data = data.average
+
+    avg_change = sum(data.values()) / len(data)
+
+    if relative:
+        text = [f"{value:.1%}" for value in data.values()]
+        avg_change = round(avg_change, 3)
+    else:
+        text = [f"${value:,.0f}" for value in data.values()]
+        avg_change = round(avg_change)
+
+    avg_change_str = (
+        f"${abs(avg_change):,.0f}"
+        if not relative
+        else f"{abs(avg_change):.1%}"
+    )
+
+    if avg_change > 0:
+        description = (
+            f"increase the net income of households by {avg_change_str}"
+        )
+    elif avg_change < 0:
+        description = (
+            f"decrease the net income of households by {avg_change_str}"
+        )
+    else:
+        description = "have no effect on household net income"
+
+    chart = go.Figure(
+        data=[
+            go.Bar(
+                x=list(data.keys()),
+                y=list(data.values()),
+                text=text,
+                marker=dict(
+                    color=[
+                        BLUE if value > 0 else DARK_GRAY
+                        for value in data.values()
+                    ]
+                ),
+            )
+        ]
+    ).update_layout(
+        title=f"{simulation.options.title} would {description}",
+        yaxis_title=f"Average change to net income ({'%' if relative else '$'})",
+        yaxis_tickformat="$,.0f" if not relative else ".0%",
+        xaxis_title=(
+            "Income decile" if decile_variable == "income" else "Wealth decile"
+        ),
+        uniformtext=dict(
+            mode="hide",
+            minsize=12,
+        ),
+        xaxis_tickvals=list(data.keys()),
+    )
+
+    return format_fig(
+        chart, country=simulation.options.country, add_zero_line=True
+    )