From 79e2d31abf931ce4c9f50d3e8839c0487d8cb61e Mon Sep 17 00:00:00 2001 From: Nikhil Woodruff Date: Tue, 14 Jan 2025 22:03:33 +1100 Subject: [PATCH 1/3] Use actual python functions --- docs/_toc.yml | 2 +- docs/reference/balance.md | 13 + docs/reference/calculate.md | 608 ------------------ .../outputs/macro/single/gov/balance.py | 29 +- policyengine/simulation.py | 169 ----- 5 files changed, 41 insertions(+), 780 deletions(-) create mode 100644 docs/reference/balance.md delete mode 100644 docs/reference/calculate.md diff --git a/docs/_toc.yml b/docs/_toc.yml index 65bcb57..3897247 100644 --- a/docs/_toc.yml +++ b/docs/_toc.yml @@ -7,4 +7,4 @@ parts: - caption: Reference chapters: - file: reference/simulation - - file: reference/calculate \ No newline at end of file + - file: reference/balance \ No newline at end of file diff --git a/docs/reference/balance.md b/docs/reference/balance.md new file mode 100644 index 0000000..650e5ef --- /dev/null +++ b/docs/reference/balance.md @@ -0,0 +1,13 @@ +# macro.single.gov.balance + + +```{eval-rst} +.. autofunction:: policyengine.outputs.macro.single.gov.balance.balance +``` + +```{eval-rst} +.. autoclass:: policyengine.outputs.macro.single.gov.balance.BalanceOutput + :members: + :undoc-members: + :show-inheritance: +``` \ No newline at end of file diff --git a/docs/reference/calculate.md b/docs/reference/calculate.md deleted file mode 100644 index bdb6c0c..0000000 --- a/docs/reference/calculate.md +++ /dev/null @@ -1,608 +0,0 @@ -# Calculate - -The `Simulation.calculate` function is the most important function in the `Simulation` class. Think of initialising the `Simulation` class as setting up the world model (with assumptions and data specified), and `calculate` as asking it a question (we refer to this as an `endpoint`). - -Generally, you call `calculate('folder/structure/question')`, with the questions being organised in a folder structure that makes sense. For example, we might call `Simulation.calculate('macro/comparison/budget/general/tax_revenue_impact')` to ask the model to calculate the impact of a tax revenue change on the budget. - -The set of possible questions depends on the parameters of the `Simulation` you've defined. For example, if you set up a simulation with household data describing a person who earns £30,000 a year, you can't ask it for the budgetary impact of a reform, but you can ask it for the change to that person's net income. If you set up a simulation from PolicyEngine's survey data, you can ask it for the impact of a policy on the distribution of income, but you can't ask "how does this reform affect this person's net income?". - -This page contains a list of all the questions you can ask the model, organised by folder structure. Under each one, we add the description, output type, and conditions under which you can ask that question. - -To reduce duplication, we've organised them into four categories, by the type of simulation you're running: - -* Single macro: simulations of a single policy on a large survey dataset. e.g. "What is the poverty rate?" -* Comparison macro: simulations of two policies (and the effect of comparing them) on a large survey dataset. e.g. "How much revenue would this policy raise?" -* Single household: simulations of a single policy on a single household. e.g. "What is my net income?" -* Comparison household: simulations of two policies (and the effect of comparing them) on a single household. e.g. "How much better off would I be under this policy?" - - -## Single macro - - - -### `macro/gov/balance/total_spending` - -The total spending of the government. - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/gov/balance/total_state_tax` - -The total tax revenue collected at the state level. - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/gov/balance/total_tax` - -The total tax revenue collected across all levels of government. - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/gov/budget_window/total_budget` - -The total government budget over the specified budget window. - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/gov/budget_window/total_federal_budget` - -The total federal government budget over the specified budget window. - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/gov/budget_window/total_spending` - -The total government spending over the specified budget window. - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/gov/budget_window/total_state_tax` - -The total state tax revenue over the specified budget window. - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/gov/budget_window/total_tax` - -The total tax revenue over the specified budget window. - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/gov/programs/{program}` - -Total expenditure or revenue from a given program. Available programs in the UK are: child_benefit, council_tax, fuel_duty, income_tax, national_insurance, ni_employer, pension_credit, state_pension, tax_credits, universal_credit, vat. - -*Output type*: `float` - -*Conditions*: `country='uk'` - - - - -### `macro/household/demographics/total_households` - -Total number of households in the simulation. - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/household/finance/deep_poverty_gap` - -The aggregate gap between household incomes and the deep poverty line. - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/household/finance/deep_poverty_rate` - -The proportion of households in deep poverty. - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/household/finance/poverty_gap` - -The aggregate gap between household incomes and the poverty line. - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/household/finance/poverty_rate` - -The proportion of households in poverty. - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/household/finance/total_benefits` - -Total benefits received by all households. - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/household/finance/total_market_income` - -Total market income (before taxes and transfers) for all households. - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/household/finance/total_net_income` - -Total net income (after taxes and transfers) for all households. - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/household/finance/total_tax` - -Total tax paid by all households. - -*Output type*: `float` - -*Conditions*: None - - - -### `macro/household/income_distribution/{lower_bound}` - -Total households with income above the lower bound. Request `household/income_distribution` for the full dictionary (this will be more helpful). - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/household/inequality/gini` - -Gini coefficient measuring income inequality (0 = perfect equality, 1 = perfect inequality). - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/household/inequality/top_10_percent_share` - -Share of total income held by the top 10% of households. - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/household/inequality/top_1_percent_share` - -Share of total income held by the top 1% of households. - -*Output type*: `float` - -*Conditions*: None - -## Comparison macro - - - -### `macro/comparison/budget/general/baseline_net_income` - -Total household net income under the baseline policy. - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/comparison/budget/general/benefit_spending_impact` - -Change in total government benefit spending between reform and baseline. - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/comparison/budget/general/budgetary_impact` - -Overall budgetary impact of the reform (positive means the reform raises revenue). - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/comparison/budget/general/households` - -Number of households affected by the reform. - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/comparison/budget/general/state_tax_revenue_impact` - -Change in state tax revenue between reform and baseline. - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/comparison/budget/general/tax_revenue_impact` - -Change in total tax revenue between reform and baseline. - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/comparison/budget/programs/{program}/baseline` - -Program spending/revenue under baseline policy, where program is one of: child_benefit, council_tax, fuel_duty, income_tax, national_insurance, ni_employer, pension_credit, state_pension, tax_credits, universal_credit, vat. - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/comparison/budget/programs/{program}/reform` - -Program spending/revenue under reform policy. See baseline endpoint for available programs. - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/comparison/budget/programs/{program}/difference` - -Change in program spending/revenue (reform - baseline). See baseline endpoint for available programs. - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/comparison/budget/window/federal_budget` - -Impact on federal budget over the budget window. - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/comparison/budget/window/total_budget` - -Impact on total budget over the budget window. - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/comparison/decile/income/average/{n}` - -Average change in household net income for income decile n (1-10). - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/comparison/decile/income/relative/{n}` - -Percentage change in household net income for income decile n (1-10). - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/comparison/decile/wealth/average/{n}` - -Average change in household net income for wealth decile n (1-10). - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/comparison/decile/wealth/relative/{n}` - -Percentage change in household net income for wealth decile n (1-10). - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/comparison/inequality/gini/baseline` - -See baseline endpoint for `household/inequality/gini`. - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/comparison/inequality/gini/reform` - -Gini coefficient under reform policy. - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/comparison/inequality/top_10_pct_share/baseline` - -See baseline endpoint for `household/inequality/top_10_percent_share`. - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/comparison/inequality/top_10_pct_share/reform` - -Share of income held by top 10% under reform policy. - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/comparison/inequality/top_1_pct_share/baseline` - -See baseline endpoint for `household/inequality/top_1_percent_share`. - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/comparison/inequality/top_1_pct_share/reform` - -Share of income held by top 1% under reform policy. - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/comparison/poverty/age/{poverty_type}/{group}/{metric}` - -Poverty metrics by age group, where: -- poverty_type is 'poverty' or 'deep_poverty' -- group is 'all', 'child', 'adult', or 'senior' -- metric is 'baseline', 'reform', or 'change_count' (only available for poverty) - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/comparison/poverty/gender/{poverty_type}/{gender}/{metric}` - -Poverty metrics by gender, where: -- poverty_type is 'poverty' or 'deep_poverty' -- gender is 'male' or 'female' -- metric is 'baseline' or 'reform' - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/comparison/winners/income_decile/all/{outcome}` - -Count of households experiencing each outcome across all income deciles, where outcome is: -'Gain more than 5%', 'Gain less than 5%', 'No change', 'Lose less than 5%', 'Lose more than 5%' - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/comparison/winners/income_decile/deciles/{outcome}` - -Count of households experiencing each outcome by income decile. See all endpoint for outcome options. - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/comparison/winners/wealth_decile/all/{outcome}` - -Count of households experiencing each outcome across all wealth deciles. See income_decile/all endpoint for outcome options. - -*Output type*: `float` - -*Conditions*: None - - - - -### `macro/comparison/winners/wealth_decile/deciles/{outcome}` - -Count of households experiencing each outcome by wealth decile. See income_decile/all endpoint for outcome options. - -*Output type*: `float` - -*Conditions*: None - -## Single household - - - -### `household/net_income` - -Net income of the household after taxes and benefits. - -*Output type*: `float` - -*Conditions*: None - -## Comparison household - - - -### `household/baseline/net_income` - -Net income of the household under the baseline policy. - -*Output type*: `float` - -*Conditions*: None - - - -### `household/comparison/net_income_change` - -Change in net income of the household between reform and baseline. - -*Output type*: `float` - -*Conditions*: None - - - -### `household/reform/net_income` - -Net income of the household under the reform policy. - -*Output type*: `float` - -*Conditions*: None diff --git a/policyengine/outputs/macro/single/gov/balance.py b/policyengine/outputs/macro/single/gov/balance.py index 3bfeb5c..6696356 100644 --- a/policyengine/outputs/macro/single/gov/balance.py +++ b/policyengine/outputs/macro/single/gov/balance.py @@ -1,8 +1,33 @@ +from typing import TypedDict, Literal from policyengine import Simulation -def balance(simulation: Simulation) -> dict: - sim = simulation.selected_sim +class BalanceOutput(TypedDict): + total_tax: float + total_spending: float + total_state_tax: float + + +def balance(simulation: Simulation) -> BalanceOutput: + """Calculate the total tax and spending for the selected simulation. + + Args: + simulation (Simulation): The selected simulation. + Must have attributes: + - country: Either "uk" or "us" + - selected_sim: Simulation object with calculate() method + + Returns: + BalanceOutput: A dictionary containing: + - total_tax (float): Total tax collected + - total_spending (float): Total government spending + - total_state_tax (float): Total state-level tax (US only, 0 for UK) + + Raises: + AttributeError: If simulation lacks required attributes + ValueError: If country is neither "uk" nor "us" + """ + sim = simulation.baseline_sim if simulation.country == "uk": total_tax = sim.calculate("gov_tax").sum() total_spending = sim.calculate("gov_spending").sum() diff --git a/policyengine/simulation.py b/policyengine/simulation.py index 9cdfcb1..b613b51 100644 --- a/policyengine/simulation.py +++ b/policyengine/simulation.py @@ -73,8 +73,6 @@ def __init__( self.comparison = reform is not None - self.output_functions, self.outputs = self._get_outputs() - self._initialise_simulations() def _set_dataset(self, dataset: str | dict | None): @@ -105,168 +103,6 @@ def _set_dataset(self, dataset: str | dict | None): ) self.data = Dataset.from_file(self.data, "2023") - def calculate(self, output: str, force: bool = False, **kwargs) -> Any: - """Calculate the given output (path). - - Args: - output (str): The output to calculate. Must be a valid path in the output tree. - force (bool): Whether to force recalculation of the output, even if it has already been calculated. - **kwargs: Any additional arguments to pass to the output function. - - Returns: - Any: The output of the calculation (using the cache if possible). - """ - if self.verbose: - print(f"Calculating {output}...") - if output.endswith("/"): - output = output[:-1] - - if output == "": - output = self.scope - - node = self.outputs - - for child_key in output.split("/")[:-1]: - if child_key not in node: - raise KeyError( - f"Output '{child_key}' not found in '{node}'. Available keys are: {list(node.keys())}" - ) - node = node[child_key] - - parent = node - child_key = output.split("/")[-1] - if parent is None: - parent = self.calculate("/".join(output.split("/")[:-1])) - if child_key not in parent: - try: - is_numeric_key = int(child_key) in parent - except ValueError: - is_numeric_key = False - if is_numeric_key: - child_key = int(child_key) - else: - # Maybe you've requested a key that is only available as a dictionary - # item in one of the output functions. Let's try to calculate the full output, - # then check the result to see if it's there. - self.calculate("/".join(output.split("/")[:-1])) - if child_key not in parent: - raise KeyError( - f"Output '{child_key}' not found in '{output}'. Available keys are: {list(parent.keys())}" - ) - node = parent[child_key] - - # Check if any descendants are None - - if ( - force or parent[child_key] is None or len(kwargs) > 0 - ) and output in self.output_functions: - output_function = self.output_functions[output] - node = output_function(self, **kwargs) - if len(kwargs) == 0: - # Only save as part of the larger tree if no non-standard args are passed - parent[child_key] = node - - if isinstance(node, dict) and len(kwargs) == 0: - for child_key in node.keys(): - self.calculate(output + "/" + str(child_key)) - - return node - - def _get_outputs(self) -> Tuple[dict, dict]: - """Get all the output functions and construct the output tree. - - Returns: - Tuple[dict, dict]: A tuple containing the output functions and the output tree. - """ - from pathlib import Path - import importlib.util - - output_functions = {} - for output in Path(__file__).parent.glob("outputs/**/*.py"): - module_name = output.stem - spec = importlib.util.spec_from_file_location(module_name, output) - if spec is None: - raise RuntimeError( - f"Expected to load a spec from file '{output.absolute}'" - ) - module = importlib.util.module_from_spec(spec) - relative_path = str( - output.relative_to(Path(__file__).parent / "outputs") - ).replace(".py", "") - if not self.comparison and "/comparison/" in relative_path: - # If we're just analysing one scenario, skip loading the comparison modules. - continue - if f"{self.scope}/" not in relative_path: - # Don't load household modules for macro comparisons, etc. - continue - - if spec.loader is None: - raise RuntimeError( - f"Expected module from '{output.absolute}' to have a loader, but it does not" - ) - spec.loader.exec_module(module) - - # Only import the function with the same name as the module, enforcing one function per file - try: - output_functions[str(relative_path)] = getattr( - module, module_name - ) - except AttributeError: - raise AttributeError( - f"Each module must contain a function with the same name as the module. Module '{str(relative_path)}.py' does not." - ) - - # If we are just calculating for a single scenario, put all 'macro/single/' children under 'macro/'. - # If not, duplicate them into 'macro/baseline/' and 'single/reform'. - - single_keys = [key for key in output_functions if "single/" in key] - for key in single_keys: - root = key.split("/")[0] - rest = "/".join(key.split("/")[2:]) - func = output_functions[key] - - def passed_reform_simulation(func, is_reform): - def adjusted_func(simulation: Simulation, **kwargs): - if is_reform: - simulation.selected_sim = simulation.reformed_sim - else: - simulation.selected_sim = simulation.baseline_sim - return func(simulation, **kwargs) - - adjusted_func.__name__ = func.__name__ - adjusted_func.__doc__ = func.__doc__ - - return adjusted_func - - if self.comparison: - output_functions[f"{root}/baseline/{rest}"] = ( - passed_reform_simulation(func, False) - ) - output_functions[f"{root}/reform/{rest}"] = ( - passed_reform_simulation(func, True) - ) - else: - output_functions[f"{root}/{rest}"] = passed_reform_simulation( - func, False - ) - - del output_functions[key] - - # Construct the output tree, fill with Nones for now - - outputs = {} - - for output_path in output_functions.keys(): - parts = output_path.split("/") - current = outputs - for part in parts[:-1]: - if part not in current: - current[part] = {} - current = current[part] - current[parts[-1]] = None - - return output_functions, outputs - def _to_reform(self, value: int | dict): if isinstance(value, dict): return Reform.from_dict(value, country_id=self.country) @@ -422,11 +258,6 @@ def _apply_region_to_simulation( with h5py.File(weights_file_path, "r") as f: weights = f[str(self.time_period)][...] - print( - weights[constituency_id], - simulation.default_calculation_period, - ) - simulation.calculate("household_net_income") simulation.set_input( From e0b48eb26c66b630b46031ddd5caaad06ae67fd5 Mon Sep 17 00:00:00 2001 From: Nikhil Woodruff Date: Thu, 16 Jan 2025 09:28:28 +1100 Subject: [PATCH 2/3] Add example use case notebook --- .gitignore | 1 + .../outputs/household/single/net_income.py | 2 +- .../macro/comparison/budget/general.py | 6 ++ test.ipynb | 102 ++++++++++++++++++ 4 files changed, 110 insertions(+), 1 deletion(-) create mode 100644 test.ipynb diff --git a/.gitignore b/.gitignore index b149da0..036062f 100644 --- a/.gitignore +++ b/.gitignore @@ -160,3 +160,4 @@ cython_debug/ #.idea/ *.ipynb +!test.ipynb diff --git a/policyengine/outputs/household/single/net_income.py b/policyengine/outputs/household/single/net_income.py index 01a97f0..bb0d69d 100644 --- a/policyengine/outputs/household/single/net_income.py +++ b/policyengine/outputs/household/single/net_income.py @@ -1,2 +1,2 @@ -def net_income(simulation): +def calculate_net_income(simulation) -> float: return simulation.selected_sim.calculate("household_net_income").sum() diff --git a/policyengine/outputs/macro/comparison/budget/general.py b/policyengine/outputs/macro/comparison/budget/general.py index c998661..cca2c72 100644 --- a/policyengine/outputs/macro/comparison/budget/general.py +++ b/policyengine/outputs/macro/comparison/budget/general.py @@ -5,6 +5,12 @@ from policyengine.utils.charts import * +def calculate_budgetary_impact( + simulation: Simulation, federal_only: bool = True +) -> float: + return 3e9 + + def general(simulation: Simulation, chart: bool = False): """Calculate the budgetary impact of the given simulation. diff --git a/test.ipynb b/test.ipynb new file mode 100644 index 0000000..f8cf4c9 --- /dev/null +++ b/test.ipynb @@ -0,0 +1,102 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "from policyengine import Simulation\n", + "\n", + "sim = Simulation(\n", + " country=\"uk\",\n", + " scope=\"macro\",\n", + " time_period=2025,\n", + " reform={\n", + " \"basic_rate\": 0.2,\n", + " }\n", + ")\n", + "\n", + "from policyengine.outputs.macro.comparison.budget.general import calculate_budgetary_impact\n", + "\n", + "budgetary_impact = calculate_budgetary_impact(sim, federal_only=False)" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "303552701743.09863" + ] + }, + "execution_count": 3, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from policyengine import Simulation\n", + "\n", + "sim = Simulation(\n", + " country=\"uk\",\n", + " scope=\"household\",\n", + " data={\n", + " \"employment_income\": 20_000,\n", + " }\n", + " time_period=2025,\n", + ")\n", + "\n", + "from policyengine.outputs.household.single.net_income import calculate_net_income\n", + "\n", + "net_income = calculate_net_income(sim)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from policyengine import Simulation\n", + "\n", + "sim = Simulation(\n", + " country=\"uk\",\n", + " scope=\"macro\",\n", + " time_period=2025,\n", + " reform={\n", + " \"basic_rate\": 0.2,\n", + " }\n", + ")\n", + "\n", + "from policyengine.outputs.macro.comparison.local_areas.parliamentary_constituencies import calculate_impact_by_constituency\n", + "\n", + "impact = calculate_impact_by_constituency(sim, metric=\"household_tax\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "base", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.14" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} From 85b90937d7ccfbd7fadc1d7f3c70841c01a4872f Mon Sep 17 00:00:00 2001 From: Nikhil Woodruff Date: Thu, 16 Jan 2025 10:50:11 +1100 Subject: [PATCH 3/3] Add changes --- .../outputs/macro/comparison/budget/general.py | 7 +++---- policyengine/simulation.py | 11 ++++++++++- 2 files changed, 13 insertions(+), 5 deletions(-) diff --git a/policyengine/outputs/macro/comparison/budget/general.py b/policyengine/outputs/macro/comparison/budget/general.py index cca2c72..0711afb 100644 --- a/policyengine/outputs/macro/comparison/budget/general.py +++ b/policyengine/outputs/macro/comparison/budget/general.py @@ -1,4 +1,3 @@ -from policyengine import Simulation import plotly.graph_objects as go import plotly.express as px import pandas as pd @@ -6,12 +5,12 @@ def calculate_budgetary_impact( - simulation: Simulation, federal_only: bool = True + self, federal_only: bool = True ) -> float: return 3e9 -def general(simulation: Simulation, chart: bool = False): +def general(simulation, chart: bool = False): """Calculate the budgetary impact of the given simulation. Args: @@ -58,7 +57,7 @@ def general(simulation: Simulation, chart: bool = False): return result -def budget_chart(simulation: Simulation, data: dict) -> go.Figure: +def budget_chart(simulation, data: dict) -> go.Figure: if simulation.country == "uk": x = ["Tax revenues", "Benefit spending", "Budgetary impact"] y = [ diff --git a/policyengine/simulation.py b/policyengine/simulation.py index b613b51..a1cc153 100644 --- a/policyengine/simulation.py +++ b/policyengine/simulation.py @@ -8,7 +8,13 @@ import pandas as pd import h5py from pathlib import Path -from typing import Literal +from typing import Literal, Callable + +from policyengine.outputs.macro.comparison.budget.general import calculate_budgetary_impact + +FUNCTIONS = ( + calculate_budgetary_impact, +) class Simulation: @@ -304,3 +310,6 @@ def _apply_region_to_simulation( simulation.default_calculation_period = self.time_period return simulation + + def calculate_budgetary_impact(self, federal_only: bool = True) -> float: + return calculate_budgetary_impact(self, federal_only=federal_only) \ No newline at end of file