From dd6ac92c6c6db2a6298aa3864c3de7840f4a2bc8 Mon Sep 17 00:00:00 2001 From: Nikhil Woodruff <35577657+nikhilwoodruff@users.noreply.github.com> Date: Sun, 28 Jul 2024 20:27:51 +0100 Subject: [PATCH] Reforms improvement (#230) * Add option for loading from dataframe * Improve flexibility of reforms * Versioning --- changelog_entry.yaml | 4 + docs/python_api/reforms.ipynb | 124 ++++++++++++++++++ docs/usage/reforms.md | 42 ------ .../data/datasets/country_template_dataset.py | 1 + policyengine_core/data/dataset.py | 23 ++++ policyengine_core/parameters/parameter.py | 4 +- policyengine_core/reforms/reform.py | 7 +- policyengine_core/simulations/simulation.py | 17 +-- 8 files changed, 166 insertions(+), 56 deletions(-) create mode 100644 docs/python_api/reforms.ipynb delete mode 100644 docs/usage/reforms.md diff --git a/changelog_entry.yaml b/changelog_entry.yaml index e69de29bb..53977fef4 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -0,0 +1,4 @@ +- bump: minor + changes: + changed: + - Reform syntax to increase flexibility. diff --git a/docs/python_api/reforms.ipynb b/docs/python_api/reforms.ipynb new file mode 100644 index 000000000..06b97f14a --- /dev/null +++ b/docs/python_api/reforms.ipynb @@ -0,0 +1,124 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Reforms\n", + "\n", + "To define a reform, simply define a class inheriting from `Reform` with an `apply(self)` function. Inside it, `self` is the tax-benefit system attached to the simulation with loaded data `self.simulation: Simulation`. From this, you can run any kind of modification on the `Simulation` instance that you like- modify parameters, variable logic or even adjust simulation data." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "from policyengine_core.country_template import Microsimulation\n", + "from policyengine_core.model_api import *\n", + "\n", + "baseline = Microsimulation()\n", + "\n", + "\n", + "class reform(Reform):\n", + " def apply(self):\n", + " simulation = self.simulation\n", + "\n", + " # Modify parameters\n", + "\n", + " simulation.tax_benefit_system.parameters.taxes.housing_tax.rate.update(\n", + " 20\n", + " )\n", + "\n", + " # Modify simulation data\n", + "\n", + " salary = simulation.calculate(\"salary\", \"2022-01\")\n", + "\n", + " new_salary = salary * 1.1\n", + "\n", + " simulation.set_input(\"salary\", \"2022-01\", new_salary)\n", + "\n", + "\n", + "reformed = Microsimulation(reform=reform)" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "( value weight\n", + " 0 110.0 1000000.0\n", + " 1 0.0 1000000.0\n", + " 2 220.0 1200000.0,\n", + " value weight\n", + " 0 100.0 1000000.0\n", + " 1 0.0 1000000.0\n", + " 2 200.0 1200000.0)" + ] + }, + "execution_count": 2, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "reformed.calculate(\"salary\", \"2022-01\"), baseline.calculate(\n", + " \"salary\", \"2022-01\"\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "( value weight\n", + " 0 4000.0 1000000.0\n", + " 1 6000.0 1200000.0,\n", + " value weight\n", + " 0 2000.0 1000000.0\n", + " 1 3000.0 1200000.0)" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "reformed.calculate(\"housing_tax\", 2022), baseline.calculate(\n", + " \"housing_tax\", 2022\n", + ")" + ] + } + ], + "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.9.19" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/docs/usage/reforms.md b/docs/usage/reforms.md deleted file mode 100644 index f60dc99a5..000000000 --- a/docs/usage/reforms.md +++ /dev/null @@ -1,42 +0,0 @@ -# Writing reforms - -In PolicyEngine country packages, reforms are defined as Python objects that specify a way to modify a tax-benefit system. This can include changes to the parameter tree, or to variables. Generally, here's how to write a reform: - -```python -from policyengine_core.reforms import Reform - -class reform(Reform): # Inherit from the Reform class - def apply(self): # You must define an apply method - self.modify_parameters(modify_parameters) - self.add_variable(new_variable_class) - self.update_variable(overriding_variable_class) - self.neutralize_variable(variable_name) -``` - -## Modifying parameters - -To modify parameters, you should define `modify_parameters` method. This should be a function of a `ParameterNode` and should return the same `ParameterNode` (but modified). For example: - -```python -def modify_parameters(parameters): - parameters.benefits.child_benefit.amount.update(period="2019-01", value=100) - return parameters -``` - -For most cases, you'll be able to do everything you need with the `update` method. This takes a `period` and a `value` and updates the parameter to that value for that period. You can also optionally specify a `start` and `stop` instant to specify the period over which the change should be applied. - -## Modifying variables - -To add a variable, define it in the same way the country package does. For example: - -```python -class new_variable(Variable): - value_type = float - entity = Person - label = "New variable" - definition_period = YEAR -``` - -The `add_variable` method takes this class as its argument. It'll throw an error if the variable already exists. If you want to override an existing variable, use `update_variable` instead. - -To neutralize a variable, use the `neutralize_variable` method. This takes the name of the variable as its argument, and ensures that this variable returns zero (or the default value) for all periods (essentially, removes all formulas). diff --git a/policyengine_core/country_template/data/datasets/country_template_dataset.py b/policyengine_core/country_template/data/datasets/country_template_dataset.py index 31e093dcd..cc5c3e974 100644 --- a/policyengine_core/country_template/data/datasets/country_template_dataset.py +++ b/policyengine_core/country_template/data/datasets/country_template_dataset.py @@ -29,6 +29,7 @@ def generate(self) -> None: "person_household_id": {ETERNITY: person_household_id}, "person_household_role": {ETERNITY: person_household_role}, "salary": {salary_time_period: salary}, + "accommodation_size": {salary_time_period: [200, 300]}, "household_weight": {weight_time_period: weight}, } self.save_dataset(data) diff --git a/policyengine_core/data/dataset.py b/policyengine_core/data/dataset.py index edc89568f..3d3bc17d5 100644 --- a/policyengine_core/data/dataset.py +++ b/policyengine_core/data/dataset.py @@ -363,3 +363,26 @@ def from_file(file_path: str, time_period: str = None): )() return dataset + + @staticmethod + def from_dataframe(dataframe: pd.DataFrame, time_period: str = None): + """Creates a dataset from a DataFrame. + + Returns: + Dataset: The dataset. + """ + file_path = Path(file_path) + dataset = type( + "Dataset", + (Dataset,), + { + "name": file_path.stem, + "label": file_path.stem, + "data_format": Dataset.FLAT_FILE, + "file_path": "dataframe", + "time_period": time_period, + "load": lambda: dataframe, + }, + )() + + return dataset diff --git a/policyengine_core/parameters/parameter.py b/policyengine_core/parameters/parameter.py index da185e187..77317314c 100644 --- a/policyengine_core/parameters/parameter.py +++ b/policyengine_core/parameters/parameter.py @@ -137,7 +137,7 @@ def clone(self): ] return clone - def update(self, period=None, start=None, stop=None, value=None): + def update(self, value=None, period=None, start=None, stop=None): """ Change the value for a given period. @@ -156,7 +156,7 @@ def update(self, period=None, start=None, stop=None, value=None): start = period.start stop = period.stop if start is None: - raise ValueError("You must provide either a start or a period") + start = "0000-01-01" start_str = str(start) stop_str = str(stop.offset(1, "day")) if stop else None diff --git a/policyengine_core/reforms/reform.py b/policyengine_core/reforms/reform.py index b227dc4d0..a3747d243 100644 --- a/policyengine_core/reforms/reform.py +++ b/policyengine_core/reforms/reform.py @@ -1,10 +1,13 @@ from __future__ import annotations import copy -from typing import Callable, Union +from typing import Callable, Union, TYPE_CHECKING from policyengine_core.parameters import ParameterNode, Parameter from policyengine_core.taxbenefitsystems import TaxBenefitSystem + +if TYPE_CHECKING: + from policyengine_core.simulations import Simulation from policyengine_core.periods import ( period as period_, instant as instant_, @@ -60,6 +63,8 @@ class Reform(TaxBenefitSystem): parameter_values: dict = None """The parameter values of the reform. This is used to inform any calls to the PolicyEngine API.""" + simulation: "Simulation" = None + def __init__(self, baseline: TaxBenefitSystem): """ :param baseline: Baseline TaxBenefitSystem. diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 9f8bfbc4b..cc8d960e0 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -85,7 +85,6 @@ def __init__( reform: Reform = None, trace: bool = False, ): - reform_applied_after = False if tax_benefit_system is None: if ( self.default_tax_benefit_system_instance is not None @@ -94,20 +93,11 @@ def __init__( tax_benefit_system = self.default_tax_benefit_system_instance else: # If reform is taken as an arg, pass it - try: - tax_benefit_system = self.default_tax_benefit_system( - reform=reform - ) - except: - tax_benefit_system = self.default_tax_benefit_system() - reform_applied_after = True + tax_benefit_system = self.default_tax_benefit_system() self.tax_benefit_system = tax_benefit_system self.reform = reform self.tax_benefit_system = tax_benefit_system - - if reform_applied_after and reform is not None: - self.apply_reform(reform) self.branch_name = "default" if dataset is None: @@ -169,6 +159,11 @@ def __init__( self.dataset = dataset self.build_from_dataset() + self.tax_benefit_system.simulation = self + + if self.reform is not None: + self.tax_benefit_system.apply_reform_set(self.reform) + # Backwards compatibility methods self.calc = self.calculate self.df = self.calculate_dataframe