From 9058caa9685d13c4944d1a59f51459aa5fd64084 Mon Sep 17 00:00:00 2001 From: Nikhil Woodruff Date: Thu, 13 Feb 2025 11:08:01 +0000 Subject: [PATCH 1/2] Add budget window function --- changelog_entry.yaml | 4 ++ .../outputs/calculate_budget_window_impact.md | 7 +++ .../calculate_budget_window_comparison.py | 58 +++++++++++++++++++ tests/country/test_budget_window.py | 19 ++++++ 4 files changed, 88 insertions(+) create mode 100644 docs/outputs/calculate_budget_window_impact.md create mode 100644 policyengine/outputs/macro/comparison/calculate_budget_window_comparison.py create mode 100644 tests/country/test_budget_window.py diff --git a/changelog_entry.yaml b/changelog_entry.yaml index e69de29b..685549c8 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -0,0 +1,4 @@ +- bump: minor + changes: + added: + - Budget window impact function. diff --git a/docs/outputs/calculate_budget_window_impact.md b/docs/outputs/calculate_budget_window_impact.md new file mode 100644 index 00000000..9bb71b0d --- /dev/null +++ b/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 +``` \ No newline at end of file diff --git a/policyengine/outputs/macro/comparison/calculate_budget_window_comparison.py b/policyengine/outputs/macro/comparison/calculate_budget_window_comparison.py new file mode 100644 index 00000000..f1c9bfe4 --- /dev/null +++ b/policyengine/outputs/macro/comparison/calculate_budget_window_comparison.py @@ -0,0 +1,58 @@ +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 \ No newline at end of file diff --git a/tests/country/test_budget_window.py b/tests/country/test_budget_window.py new file mode 100644 index 00000000..e528aa97 --- /dev/null +++ b/tests/country/test_budget_window.py @@ -0,0 +1,19 @@ + + +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 From 5c98d0a6149aa387a6c940f5ebe25dca6545e6d1 Mon Sep 17 00:00:00 2001 From: Nikhil Woodruff Date: Thu, 13 Feb 2025 11:47:07 +0000 Subject: [PATCH 2/2] Format --- .../calculate_budget_window_comparison.py | 36 ++++++++++++------- tests/country/test_budget_window.py | 2 -- 2 files changed, 24 insertions(+), 14 deletions(-) diff --git a/policyengine/outputs/macro/comparison/calculate_budget_window_comparison.py b/policyengine/outputs/macro/comparison/calculate_budget_window_comparison.py index f1c9bfe4..bf316dae 100644 --- a/policyengine/outputs/macro/comparison/calculate_budget_window_comparison.py +++ b/policyengine/outputs/macro/comparison/calculate_budget_window_comparison.py @@ -2,15 +2,20 @@ from policyengine import Simulation from policyengine_core.simulations import Microsimulation -from policyengine.outputs.macro.single.budget import _calculate_government_balance +from policyengine.outputs.macro.single.budget import ( + _calculate_government_balance, +) -def calculate_budget_window_comparison(simulation: Simulation, count_years: int = 10) -> pd.DataFrame: + +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. """ @@ -23,17 +28,24 @@ def calculate_budget_window_comparison(simulation: Simulation, count_years: int 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) + 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 - }) + + return pd.DataFrame( + { + "year": years, + "federal_budget_impact": federal_budget_impacts, + "state_budget_impact": state_budget_impacts, + } + ) + def _get_balance_for_year( sim: Microsimulation, @@ -55,4 +67,4 @@ def _get_balance_for_year( national_balance = national_tax - total_spending state_balance = total_state_tax - return national_balance, state_balance \ No newline at end of file + return national_balance, state_balance diff --git a/tests/country/test_budget_window.py b/tests/country/test_budget_window.py index e528aa97..0f61efb5 100644 --- a/tests/country/test_budget_window.py +++ b/tests/country/test_budget_window.py @@ -1,5 +1,3 @@ - - def test_budget_window(): from policyengine import Simulation