Add budget window extension

changelog_entry.yaml

@@ -0,0 +1,4 @@
+- bump: minor
+  changes:
+    added:
+    - Budget window impact function.

docs/outputs/calculate_budget_window_impact.md

@@ -0,0 +1,7 @@
+# Budget windows
+
+For a given policy reform, this function calculates the budget window impact at the federal (national) and State level.
+
+```{eval-rst}
+.. autofunction:: policyengine.outputs.macro.comparison.calculate_budget_window_comparison.calculate_budget_window_comparison
+```

policyengine/outputs/macro/comparison/calculate_budget_window_comparison.py

@@ -0,0 +1,70 @@
+import pandas as pd
+from policyengine import Simulation
+from policyengine_core.simulations import Microsimulation
+
+from policyengine.outputs.macro.single.budget import (
+    _calculate_government_balance,
+)
+
+
+def calculate_budget_window_comparison(
+    simulation: Simulation, count_years: int = 10
+) -> pd.DataFrame:
+    """Calculate how a reform affects the budget over a specified window.
+
+    Args:
+        simulation (Simulation): The simulation object containing baseline and reform simulations.
+        count_years (int, optional): The number of years over which to calculate the budget impact. Defaults to 10.
+
+    Returns:
+        pd.DataFrame: A DataFrame containing the years, federal budget impacts, and state budget impacts.
+    """
+
+    start_year = simulation.options.time_period
+    end_year = start_year + count_years
+
+    years = []
+    federal_budget_impacts = []
+    state_budget_impacts = []
+
+    for year in range(start_year, end_year):
+        baseline_f, baseline_s = _get_balance_for_year(
+            simulation.baseline_simulation, year, simulation.options.country
+        )
+        reform_f, reform_s = _get_balance_for_year(
+            simulation.reform_simulation, year, simulation.options.country
+        )
+        years.append(year)
+        federal_budget_impacts.append(reform_f - baseline_f)
+        state_budget_impacts.append(reform_s - baseline_s)
+
+    return pd.DataFrame(
+        {
+            "year": years,
+            "federal_budget_impact": federal_budget_impacts,
+            "state_budget_impact": state_budget_impacts,
+        }
+    )
+
+
+def _get_balance_for_year(
+    sim: Microsimulation,
+    year: int,
+    country: str,
+):
+    if country == "uk":
+        total_tax = sim.calculate("gov_tax", period=year).sum()
+        total_spending = sim.calculate("gov_spending", period=year).sum()
+        total_state_tax = 0
+    else:
+        total_tax = sim.calculate("household_tax", period=year).sum()
+        total_spending = sim.calculate("household_benefits", period=year).sum()
+        total_state_tax = sim.calculate(
+            "household_state_income_tax", period=year
+        ).sum()
+
+    national_tax = total_tax - total_state_tax
+
+    national_balance = national_tax - total_spending
+    state_balance = total_state_tax
+    return national_balance, state_balance

tests/country/test_budget_window.py

@@ -0,0 +1,17 @@
+def test_budget_window():
+    from policyengine import Simulation
+
+    sim = Simulation(
+        country="uk",
+        scope="macro",
+        reform={
+            "gov.hmrc.income_tax.allowances.personal_allowance.amount": 10_000,
+        },
+    )
+
+    window = sim.calculate_budget_window_comparison(count_years=3)
+
+    assert len(window) == 3
+
+    assert (window.federal_budget_impact > 0).all()
+    assert window.federal_budget_impact.mean() > 15e9