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

policyengine_uk/utils/dependencies.py

@@ -0,0 +1,259 @@
+import re
+from policyengine_uk.system import system
+import plotly.express as px
+import pandas as pd
+from plotly import graph_objects as go
+
+
+def find_variable_in_trees(variable_name, tracer):
+    for tree in tracer.trees:
+        try:
+            node = find_variable_in_tree_recursive(variable_name, tree)
+            if node:
+                return node
+        except ValueError:
+            continue  # Variable not in this tree, try next
+
+    raise ValueError(
+        f"Variable '{variable_name}' not found in any simulation tree. "
+        f"Make sure you've calculated this variable or a parent variable that depends on it."
+    )
+
+
+def find_variable_in_tree_recursive(variable_name, node):
+    if node.name == variable_name:
+        return node
+
+    for child in node.children:
+        result = find_variable_in_tree_recursive(variable_name, child)
+        if result:
+            return result
+
+    return None
+
+
+def get_variable_dependencies(variable_name, sim):
+    node = find_variable_in_trees(variable_name, sim.tracer)
+    if not node:
+        return []
+    return [child.name for child in node.children]
+
+
+def extract_variables_regex(formula_source):
+    # Find all strings in double quotes that look like variable names
+    pattern = r'benunit\(\s*["\']([^"\']+)["\']\s*,'
+    matches = re.findall(pattern, formula_source)
+    return matches
+
+
+def calculate_dependency_contributions(
+    sim, variable_name, year, top_n=None, filter=None, map_to=None
+):
+    original_values = sim.calculate(variable_name, year)
+
+    if map_to is not None:
+        source_entity = sim.tax_benefit_system.get_variable(
+            variable_name
+        ).entity.key
+        original_values_mapped = sim.map_result(
+            original_values,
+            source_entity,
+            map_to,
+        )
+    else:
+        original_values_mapped = original_values
+
+    dependency_contributions = {}
+    first_level_dependencies = get_variable_dependencies(variable_name, sim)
+    for variable in first_level_dependencies:
+        if "weight" in variable:
+            continue
+        sim.get_holder(variable_name).delete_arrays(year)
+        value_type = sim.tax_benefit_system.get_variable(variable).value_type
+        current_values = sim.calculate(variable, year)
+        if value_type == float:
+            sim.set_input(variable, year, (current_values * 0).astype(float))
+
+        new_values_mapped = sim.calculate(variable_name, year, map_to=map_to)
+        if filter is not None:
+            contribution = (
+                original_values_mapped[filter] - new_values_mapped[filter]
+            ).mean()
+        else:
+            contribution = (original_values_mapped - new_values_mapped).mean()
+        dependency_contributions[variable] = contribution
+        sim.set_input(variable_name, year, original_values)
+        sim.set_input(variable, year, current_values)
+
+    result = pd.Series(dependency_contributions)
+
+    if top_n is not None:
+        # Keep the top N variables by absolute contribution
+        result = result.reindex(result.abs().nlargest(top_n).index)
+
+    return result.sort_values(ascending=False)
+
+
+def calculate_dependency_contribution_change(
+    baseline_sim, reform_sim, variable_name, year, reform_year=None, top_n=5
+):
+    baseline_dependency = calculate_dependency_contributions(
+        baseline_sim, variable_name, year
+    )
+    reform_dependency = calculate_dependency_contributions(
+        reform_sim, variable_name, reform_year or year
+    )
+
+    df = pd.DataFrame(
+        {
+            "baseline": pd.Series(baseline_dependency),
+            "reform": pd.Series(reform_dependency),
+        }
+    ).fillna(0)
+
+    df["change"] = df["reform"] - df["baseline"]
+    df["relative_change"] = df["change"] / df["baseline"].abs().replace(0, 1)
+
+    # Keep the top N variables by absolute change
+    if top_n is not None:
+        df = df.reindex(df["change"].abs().nlargest(top_n).index)
+
+    return df.sort_values(by="change", ascending=False)
+
+
+def create_waterfall_chart(sim, variable_name, year, top_n=5):
+    if not sim.trace:
+        raise ValueError(
+            "Simulation must have trace enabled to create a waterfall chart."
+        )
+
+    df = calculate_dependency_contributions(
+        sim, variable_name, year, top_n=top_n
+    )
+
+    # make a waterfall chart
+
+    fig = go.Figure(
+        go.Waterfall(
+            name="Waterfall",
+            orientation="v",
+            measure=["relative"] * len(df),
+            x=df.index,
+            y=df.values,
+            connector={"line": {"color": "rgb(63, 63, 63)", "width": 2}},
+            increasing={"marker": {"color": "green"}},
+            decreasing={"marker": {"color": "red"}},
+            totals={"marker": {"color": "blue"}},
+        )
+    )
+    return format_fig(fig).update_layout(
+        title=f"Dependency contributions for {variable_name} in {year}",
+    )
+
+
+def create_waterfall_change_chart(
+    sim_1, sim_2, variable_name, year, sim_2_year=None, top_n=5
+):
+    df = calculate_dependency_contribution_change(
+        sim_1, sim_2, variable_name, year, reform_year=sim_2_year, top_n=top_n
+    )
+
+    # make a waterfall chart
+
+    fig = go.Figure(
+        go.Waterfall(
+            name="Waterfall",
+            orientation="v",
+            measure=["relative"] * len(df),
+            x=df.index,
+            y=df.change,
+            connector={"line": {"color": "rgb(63, 63, 63)", "width": 2}},
+            increasing={"marker": {"color": "green"}},
+            decreasing={"marker": {"color": "red"}},
+            totals={"marker": {"color": "blue"}},
+        )
+    )
+    return format_fig(fig).update_layout(
+        title=f"Change in {variable_name} contributions from reform",
+    )
+
+
+def add_fonts():
+    fonts = HTML(
+        """
+    <link rel="preconnect" href="https://fonts.googleapis.com">
+    <link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
+    <link href="https://fonts.googleapis.com/css2?family=Roboto+Serif:ital,opsz,wght@0,8..144,100..900;1,8..144,100..900&display=swap" rel="stylesheet">
+    <link href="https://fonts.googleapis.com/css2?family=Roboto+Mono:ital,wght@0,100..700;1,100..700&display=swap" rel="stylesheet">
+    """
+    )
+    return display_html(fonts)
+
+
+from IPython.core.display import HTML, display_html
+
+
+def format_fig(fig):
+    # PolicyEngine style (roboto mono for numbers, roboto serif for text), dark grey for negative, blue for positive, spacing, etc.
+
+    # Set layout properties for a clean, professional look
+    fig.update_layout(
+        plot_bgcolor="white",
+        paper_bgcolor="white",
+        font_family="Roboto Serif",
+        font_size=12,
+        margin=dict(l=60, r=60, t=80, b=60),
+        title_font_size=16,
+        title_font_family="Roboto Serif",
+        title_x=0.5,
+        title_xanchor="center",
+    )
+
+    # Update axes
+    fig.update_xaxes(
+        title_font_family="Roboto Serif",
+        tickfont_family="Roboto Mono",
+        showline=True,
+        linewidth=1,
+        linecolor="lightgray",
+    )
+
+    fig.update_yaxes(
+        title_font_family="Roboto Serif",
+        tickfont_family="Roboto Mono",
+        tickformat=",.0f",
+        tickprefix="£",
+        gridcolor="lightgray",
+        showline=True,
+        linewidth=1,
+        linecolor="lightgray",
+    )
+
+    # Update waterfall specific styles
+    fig.update_traces(
+        increasing=dict(marker=dict(color="#2C6496")),
+        decreasing=dict(marker=dict(color="#555555")),
+        connector=dict(line=dict(color="rgb(63, 63, 63)", width=1)),
+    )
+
+    # Height/width adjustments for better visibility
+    fig.update_layout(
+        height=600,
+        width=800,
+        # set font to black
+        font_color="black",
+    )
+
+    # Add rounded numbers on bars
+    for trace in fig.data:
+        if isinstance(trace, go.Waterfall):
+            trace.texttemplate = "£%{y:,.0f}"
+            trace.textposition = "outside"
+            trace.hovertemplate = "%{x}: £%{y:,.0f}<extra></extra>"
+            # mono font for numbers
+            trace.textfont = dict(family="Roboto Mono", size=12, color="black")
+
+    # Add padding round the edges
+    fig.update_layout(margin=dict(l=100, r=100, t=100, b=100))
+
+    return fig

policyengine_uk/utils/parameters.py

@@ -26,7 +26,18 @@ def backdate_parameters(
 
 
 def convert_to_fiscal_year_parameters(parameters):
-    YEARS = list(range(2015, 2026))
+    """
+    Convert parameters to use UK fiscal year values.
+
+    The UK fiscal year runs April 6 to April 5. When querying a parameter
+    for a year (e.g., param("2026")), we want the value at April 30 of
+    that year (which represents the fiscal year starting April 6).
+
+    This function samples each parameter at April 30 of each year and
+    sets that as the value for the entire year period.
+    """
+    # Cover years from 2015 through 2040 for long-term projections
+    YEARS = list(range(2015, 2041))
     for param in parameters.get_descendants():
         if isinstance(param, Parameter):
             for year in YEARS:

policyengine_uk/utils/scenario.py

@@ -0,0 +1,225 @@
+from pydantic import BaseModel
+from typing import Optional, Callable, Dict, Type, Union
+from policyengine_core.simulations import Simulation
+from policyengine_core.reforms import Reform
+from policyengine_core.periods import period, instant
+
+
+class Scenario(BaseModel):
+    """Represents a scenario configuration for policy simulations.
+
+    A scenario can include parameter changes and/or simulation modifications
+    that are applied before running a simulation. Scenarios can be combined
+    using the + operator.
+    """
+
+    applied_before_data_load: bool = False
+
+    parameter_changes: Optional[
+        Dict[
+            str,
+            Union[
+                int,
+                float,
+                bool,
+                Dict[Union[str, int], Union[int, float, bool]],
+            ],
+        ]
+    ] = None
+    """A dictionary of parameter changes to apply to the simulation. These are applied *before* parameter operations."""
+
+    simulation_modifier: Optional[Callable[["Simulation"], None]] = None
+    """A function that modifies the simulation before running it."""
+
+    class Config:
+        """Pydantic configuration."""
+
+        arbitrary_types_allowed = True  # Allow Callable types
+
+    def __add__(self, other: "Scenario") -> "Scenario":
+        """Combine two scenarios by merging parameter changes and chaining modifiers.
+
+        Args:
+            other: Another Scenario to combine with this one
+
+        Returns:
+            A new Scenario with merged parameter changes and combined modifiers
+        """
+        # Merge parameter changes (other's changes take precedence in conflicts)
+        merged_params = {}
+
+        if self.parameter_changes:
+            merged_params.update(self.parameter_changes)
+
+        if other.parameter_changes:
+            for key, value in other.parameter_changes.items():
+                if (
+                    key in merged_params
+                    and isinstance(merged_params[key], dict)
+                    and isinstance(value, dict)
+                ):
+                    # Deep merge nested dictionaries
+                    merged_params[key] = {**merged_params[key], **value}
+                else:
+                    # Simple override
+                    merged_params[key] = value
+
+        # Chain simulation modifiers
+        combined_modifier = None
+
+        if self.simulation_modifier and other.simulation_modifier:
+            # Both have modifiers - chain them
+            def combined_modifier(simulation: Simulation) -> None:
+                self.simulation_modifier(simulation)
+                other.simulation_modifier(simulation)
+
+        elif self.simulation_modifier:
+            combined_modifier = self.simulation_modifier
+        elif other.simulation_modifier:
+            combined_modifier = other.simulation_modifier
+
+        # Return new scenario with merged configuration
+        return Scenario(
+            parameter_changes=merged_params if merged_params else None,
+            simulation_modifier=combined_modifier,
+        )
+
+    @classmethod
+    def from_reform(
+        cls, reform: Union[tuple, dict, Type[Reform]]
+    ) -> "Scenario":
+        """Create a Scenario from various reform representations.
+
+        Args:
+            reform: Can be:
+                - A Reform class type (will be applied via simulation modifier)
+                - A dict of parameter changes
+                - A tuple (treated as a Reform for backward compatibility)
+
+        Returns:
+            A new Scenario configured with the reform
+
+        Raises:
+            ValueError: If reform type is not supported
+        """
+        if isinstance(reform, type) and issubclass(reform, Reform):
+            # Reform class - create modifier function
+            def modifier(simulation: Simulation) -> None:
+                reform_instance = reform()
+                reform_instance.apply(simulation.tax_benefit_system)
+
+            return cls(
+                simulation_modifier=modifier,
+            )
+
+        elif isinstance(reform, dict):
+            # Dictionary of parameter changes
+            # Make sure to capture YYYY-MM-DD.YYYY-MM-DD.
+
+            def modifier(sim: Simulation):
+                for parameter in reform:
+                    if isinstance(reform[parameter], dict):
+                        for period_str, value in reform[parameter].items():
+                            if "." in period_str:
+                                start = instant(period_str.split(".")[0])
+                                stop = instant(period_str.split(".")[1])
+                                period_ = None
+                            else:
+                                start = None
+                                stop = None
+                                period_ = period(period_str)
+                            sim.tax_benefit_system.parameters.get_child(
+                                parameter
+                            ).update(
+                                start=start,
+                                stop=stop,
+                                period=period_,
+                                value=value,
+                            )
+                    else:
+                        start = instant("2023-01-01")
+                        stop = None
+                        period_ = None
+
+                        sim.tax_benefit_system.parameters.get_child(
+                            parameter
+                        ).update(
+                            start=start,
+                            stop=stop,
+                            period=period_,
+                            value=reform[parameter],
+                        )
+
+            return Scenario(
+                simulation_modifier=modifier,
+            )
+
+        elif isinstance(reform, tuple):
+            # Tuple format (legacy support) - treat as a Reform class
+            # Assuming the tuple contains (reform_class, *args)
+            if (
+                len(reform) > 0
+                and isinstance(reform[0], type)
+                and issubclass(reform[0], Reform)
+            ):
+                reform_class = reform[0]
+                reform_args = reform[1:] if len(reform) > 1 else ()
+
+                def modifier(simulation: Simulation) -> None:
+                    reform_instance = reform_class(*reform_args)
+                    reform_instance.apply(simulation.tax_benefit_system)
+
+                return cls(
+                    simulation_modifier=modifier,
+                )
+            else:
+                raise ValueError(f"Invalid tuple format for reform: {reform}")
+
+        else:
+            raise ValueError(
+                f"Unsupported reform type: {type(reform)}. "
+                "Expected Reform class, dict, or tuple."
+            )
+
+    def apply(self, simulation: Simulation) -> None:
+        """Apply this scenario to a simulation.
+
+        First applies parameter changes, then runs the simulation modifier if present.
+
+        Args:
+            simulation: The simulation to modify
+        """
+        # Apply parameter changes first
+        if self.parameter_changes:
+            for path, value in self.parameter_changes.items():
+                if isinstance(value, dict):
+                    # Handle nested parameter changes
+                    for sub_path, sub_value in value.items():
+                        full_path = f"{path}.{sub_path}"
+                        simulation.tax_benefit_system.parameters.update(
+                            full_path,
+                            period=None,  # Apply to all periods
+                            value=sub_value,
+                        )
+                else:
+                    # Simple parameter change
+                    simulation.tax_benefit_system.parameters.update(
+                        path,
+                        period=None,
+                        value=value,  # Apply to all periods
+                    )
+
+        # Then apply simulation modifier
+        if self.simulation_modifier:
+            self.simulation_modifier(simulation)
+
+    def __repr__(self) -> str:
+        """String representation of the Scenario."""
+        parts = []
+        if self.parameter_changes:
+            parts.append(
+                f"parameter_changes={len(self.parameter_changes)} items"
+            )
+        if self.simulation_modifier:
+            parts.append("simulation_modifier=<function>")
+        return f"Scenario({', '.join(parts)})"

policyengine_uk/utils/solve_private_school_attendance_factor.py

@@ -1,4 +1,4 @@
-from policyengine import Simulation
+from policyengine_uk import Microsimulation
 from policyengine_core.reforms import Reform
 from tqdm import tqdm
 
@@ -19,11 +19,9 @@ for factor in tqdm([round(x * 0.01, 2) for x in range(70, 91)]):
     }
 
     # Run the reformed microsimulation
-    reformed = Simulation(
-        scope="macro",
-        baseline=reform,
-        country="uk",
-    ).baseline_simulation
+    reformed = Microsimulation(
+        reform=reform,
+    )
     reformed.baseline_simulation.get_holder(
         "attends_private_school"
     ).delete_arrays()

policyengine_uk/variables/contrib/policyengine/education_budget_change.py

@@ -8,7 +8,6 @@ class education_budget_change(Variable):
     value_type = float
 
     def formula(household, period, parameters):
-
         budget_increase = (
             parameters(period).gov.contrib.policyengine.budget.education * 1e9
         )

policyengine_uk/variables/contrib/policyengine/employer_ni/baseline_employer_cost.py

@@ -35,7 +35,11 @@ class baseline_employer_cost(Variable):
         )
 
         # Calculate baseline employer cost
-        baseline_parameters = parameters(period).baseline
+        if person.simulation.baseline is None:
+            return 0
+        baseline_parameters = (
+            person.simulation.baseline.tax_benefit_system.parameters(period)
+        )
         baseline_class_1 = (
             baseline_parameters.gov.hmrc.national_insurance.class_1
         )

policyengine_uk/variables/contrib/policyengine/employer_ni/employer_ni_fixed_employer_cost_change.py

@@ -9,14 +9,13 @@ class employer_ni_fixed_employer_cost_change(Variable):
     unit = GBP
 
     def formula(person, period, parameters):
-        return 0
         employee_incidence = parameters(
             period
         ).gov.contrib.policyengine.employer_ni.employee_incidence
         if employee_incidence == 0:
             return 0
-        # First, calculate baseline and reformed employer NI contributions.
 
+        # First, calculate baseline and reformed employer NI contributions.
         prior_employment_income = person(
             "employment_income_before_lsr", period
         )
@@ -43,7 +42,12 @@ class employer_ni_fixed_employer_cost_change(Variable):
         )
 
         # Calculate baseline employer cost
-        baseline_parameters = parameters(period).baseline
+        if person.simulation.baseline is None:
+            return 0
+        baseline_parameters = (
+            person.simulation.baseline.tax_benefit_system.parameters(period)
+        )
+
         baseline_class_1 = (
             baseline_parameters.gov.hmrc.national_insurance.class_1
         )
@@ -52,19 +56,14 @@ class employer_ni_fixed_employer_cost_change(Variable):
         p_b = (
             baseline_parameters.gov.contrib.policyengine.employer_ni.exempt_employer_pension_contributions
         )
-        pen_con_subtracted_b = employer_pension_contributions
-        if p_b:
-            pen_con_subtracted_b = employer_pension_contributions
-        else:
-            pen_con_subtracted_b = 0
+        pen_con_subtracted_b = employer_pension_contributions if p_b else 0
 
-        baseline_employer_ni = r_b * (
-            ni_class_1_income - pen_con_subtracted_b - t_b
+        baseline_employer_ni = r_b * max_(
+            ni_class_1_income - pen_con_subtracted_b - t_b, 0
         )
         c_b = ni_class_1_income + baseline_employer_ni
 
-        # Calculate new employment income keeping employer cost constant
-
+        # Calculate reform employer cost
         reform_parameters = parameters(period)
         reform_class_1 = reform_parameters.gov.hmrc.national_insurance.class_1
         r_r = reform_class_1.rates.employer
@@ -72,32 +71,33 @@ class employer_ni_fixed_employer_cost_change(Variable):
         p_r = (
             reform_parameters.gov.contrib.policyengine.employer_ni.exempt_employer_pension_contributions
         )
-        if p_r:
-            pen_con_subtracted_r = employer_pension_contributions
-        else:
-            pen_con_subtracted_r = 0
+        pen_con_subtracted_r = employer_pension_contributions if p_r else 0
 
+        # Early return if no change in parameters
+        if r_b == r_r and t_b == t_r and p_b == p_r:
+            return 0
+
+        # Calculate new employment income keeping employer cost constant
+        # Solve: c_b = new_ni_class_1_income + r_r * max(new_ni_class_1_income - pen_con_subtracted_r - t_r, 0)
+        # Rearranging: new_ni_class_1_income = (c_b + r_r * (pen_con_subtracted_r + t_r)) / (1 + r_r)
         new_ni_class_1_income = (c_b + r_r * (pen_con_subtracted_r + t_r)) / (
             1 + r_r
         )
 
         # Find difference in employment income
-
         previous_employment_income = (
-            ni_class_1_income - benefits
-        ) - employer_pension_contributions
+            ni_class_1_income - benefits - employer_pension_contributions
+        )
         new_employment_income = (
-            new_ni_class_1_income - benefits
-        ) - employer_pension_contributions
+            new_ni_class_1_income - benefits - employer_pension_contributions
+        )
 
         pay_change = new_employment_income - previous_employment_income
 
         # Apply incidence percentage
-
         interpolated_pay_change = pay_change * employee_incidence
 
         # Where a person's prior employment income was below the secondary threshold, the formula doesn't hold, so assume no change.
-
         below_threshold = previous_employment_income < t_b
 
         return where(below_threshold, 0, interpolated_pay_change)

policyengine_uk/variables/contrib/policyengine/employer_ni/employer_ni_response_capital_incidence.py

@@ -36,7 +36,7 @@ class employer_ni_response_capital_incidence(Variable):
         value = (
             amount_paid_by_employers
             * share_of_total_wealth
-            * capital_incidence
+            * emp_ni.capital_incidence
         )
 
         if total_wealth == 0:

policyengine_uk/variables/contrib/policyengine/employer_ni/employer_ni_response_consumer_incidence.py

@@ -36,7 +36,7 @@ class employer_ni_response_consumer_incidence(Variable):
         value = (
             amount_paid_by_employers
             * share_of_total_consumption
-            * consumer_incidence
+            * emp_ni.consumer_incidence
         )
 
         if total_consumption == 0:

policyengine_uk/variables/contrib/policyengine/other_public_spending_budget_change.py

@@ -8,7 +8,6 @@ class other_public_spending_budget_change(Variable):
     value_type = float
 
     def formula(household, period, parameters):
-
         budget_increase = (
             parameters(
                 period

policyengine_uk/variables/gov/dfe/targeted_childcare_entitlement/targeted_childcare_entitlement_eligible.py

@@ -9,7 +9,6 @@ class targeted_childcare_entitlement_eligible(Variable):
     defined_for = "would_claim_targeted_childcare"
 
     def formula(benunit, period, parameters):
-
         # Check if household is in England
         country = benunit.household("country", period)
         in_england = country == country.possible_values.ENGLAND