From 6d75a9722696e3ca2a95da8ae839cde525017c4e Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Tue, 24 Dec 2024 18:05:17 +0100 Subject: [PATCH 01/18] feat: Neutralize variables within StructuralReform builder --- policyengine_core/model_api.py | 2 +- policyengine_core/reforms/__init__.py | 1 + .../reforms/structural_reform.py | 94 +++++++++++++++++++ 3 files changed, 96 insertions(+), 1 deletion(-) create mode 100644 policyengine_core/reforms/structural_reform.py diff --git a/policyengine_core/model_api.py b/policyengine_core/model_api.py index 68f1f2dd4..aa4876d5a 100644 --- a/policyengine_core/model_api.py +++ b/policyengine_core/model_api.py @@ -20,7 +20,7 @@ ) from policyengine_core.periods import DAY, ETERNITY, MONTH, YEAR, period from policyengine_core.populations import ADD, DIVIDE -from policyengine_core.reforms import Reform +from policyengine_core.reforms import Reform, StructuralReform from policyengine_core.simulations import ( calculate_output_add, calculate_output_divide, diff --git a/policyengine_core/reforms/__init__.py b/policyengine_core/reforms/__init__.py index a7ce20a8a..85cf6ca80 100644 --- a/policyengine_core/reforms/__init__.py +++ b/policyengine_core/reforms/__init__.py @@ -1 +1,2 @@ from .reform import Reform, set_parameter +from .structural_reform import StructuralReform \ No newline at end of file diff --git a/policyengine_core/reforms/structural_reform.py b/policyengine_core/reforms/structural_reform.py new file mode 100644 index 000000000..e410c0afa --- /dev/null +++ b/policyengine_core/reforms/structural_reform.py @@ -0,0 +1,94 @@ +from typing import Annotated, Callable +from datetime import datetime +from policyengine_core.variables import Variable +from policyengine_core.taxbenefitsystems import TaxBenefitSystem +from policyengine_core.errors import VariableNotFoundError + + +class StructuralReform: # Should this inherit from Reform and/or TaxBenefitSystem? + + variables: list[Variable] = [] + + def __init__( + self, + tax_benefit_system: TaxBenefitSystem, + start_instant: Annotated[str, "YYYY-MM-DD"] | None = str(datetime.now().year) + "-01-01", + end_instant: Annotated[str, "YYYY-MM-DD"] | None = None, + ): + # TODO: Validate start_instant and end_instant + + self.tax_benefit_system = tax_benefit_system + self.start_instant = start_instant + self.end_instant = end_instant + + def neutralize_variable(self, name: str) -> None: + """ + Neutralize a variable by setting its formula to return the default value + from the StructuralReform's start_instant date onward. + + Args: + name: The name of the variable + + Raises: + VariableNotFoundError: If the variable is not found in the tax benefit system + """ + # Clone variable + variable: Variable | None = self._fetch_variable(name) + + if variable is None: + raise VariableNotFoundError (f"Unable to neutralize {name}; variable not found.") + + # Add formula to variable that returns all defaults + neutralized_formula = self._neutralized_formula(variable) + variable = self._add_evolved_formula(variable, neutralized_formula) + + def _fetch_variable(self, name: str) -> Variable | None: + """ + Fetch the variable by reference from the tax benefit system. + + Args: + name: The name of the variable + """ + return self.tax_benefit_system.get_variable(name) + + # Method to modify metadata based on new items? + + # Method to add formula based on date + def _add_evolved_formula(self, variable: Variable, formula: Callable) -> Variable: + """ + Add an evolved formula, beginning on the StructuralReform's start_instant date, + to a variable, and return said variable. + + For more on evolved formulas, consult + https://openfisca.org/doc/coding-the-legislation/40_legislation_evolutions.html + + Args: + variable: The variable to which the formula will be added + start_instant: The date on which the formula will take effect + formula: The formula to be added + + Returns: + The variable with the evolved formula + """ + + # Add formula to variable's "formulas" SortedDict + variable.formulas[self.start_instant] = formula + + return variable + # + + def _neutralized_formula(self, variable: Variable) -> Callable: + """ + Return a formula that itself returns the default value of a variable. + + Args: + variable: The variable to be neutralized + + Returns: + The neutralized formula + """ + return lambda population, period, parameters: variable.default_value + + # Validate start instant + + # Default outputs method of some sort? \ No newline at end of file From 27efed80c6e8947526e89bc4dc3bc1036d48ef50 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 26 Dec 2024 19:36:27 +0100 Subject: [PATCH 02/18] feat: update_variable; partial add_variable --- policyengine_core/reforms/__init__.py | 2 +- .../reforms/structural_reform.py | 301 ++++++++++++------ 2 files changed, 212 insertions(+), 91 deletions(-) diff --git a/policyengine_core/reforms/__init__.py b/policyengine_core/reforms/__init__.py index 85cf6ca80..ff3ff3446 100644 --- a/policyengine_core/reforms/__init__.py +++ b/policyengine_core/reforms/__init__.py @@ -1,2 +1,2 @@ from .reform import Reform, set_parameter -from .structural_reform import StructuralReform \ No newline at end of file +from .structural_reform import StructuralReform diff --git a/policyengine_core/reforms/structural_reform.py b/policyengine_core/reforms/structural_reform.py index e410c0afa..920fb0f7b 100644 --- a/policyengine_core/reforms/structural_reform.py +++ b/policyengine_core/reforms/structural_reform.py @@ -1,94 +1,215 @@ from typing import Annotated, Callable from datetime import datetime +from policyengine_core.periods import instant, Instant from policyengine_core.variables import Variable from policyengine_core.taxbenefitsystems import TaxBenefitSystem -from policyengine_core.errors import VariableNotFoundError - - -class StructuralReform: # Should this inherit from Reform and/or TaxBenefitSystem? - - variables: list[Variable] = [] - - def __init__( - self, - tax_benefit_system: TaxBenefitSystem, - start_instant: Annotated[str, "YYYY-MM-DD"] | None = str(datetime.now().year) + "-01-01", - end_instant: Annotated[str, "YYYY-MM-DD"] | None = None, - ): - # TODO: Validate start_instant and end_instant - - self.tax_benefit_system = tax_benefit_system - self.start_instant = start_instant - self.end_instant = end_instant - - def neutralize_variable(self, name: str) -> None: - """ - Neutralize a variable by setting its formula to return the default value - from the StructuralReform's start_instant date onward. - - Args: - name: The name of the variable - - Raises: - VariableNotFoundError: If the variable is not found in the tax benefit system - """ - # Clone variable - variable: Variable | None = self._fetch_variable(name) - - if variable is None: - raise VariableNotFoundError (f"Unable to neutralize {name}; variable not found.") - - # Add formula to variable that returns all defaults - neutralized_formula = self._neutralized_formula(variable) - variable = self._add_evolved_formula(variable, neutralized_formula) - - def _fetch_variable(self, name: str) -> Variable | None: - """ - Fetch the variable by reference from the tax benefit system. - - Args: - name: The name of the variable - """ - return self.tax_benefit_system.get_variable(name) - - # Method to modify metadata based on new items? - - # Method to add formula based on date - def _add_evolved_formula(self, variable: Variable, formula: Callable) -> Variable: - """ - Add an evolved formula, beginning on the StructuralReform's start_instant date, - to a variable, and return said variable. - - For more on evolved formulas, consult - https://openfisca.org/doc/coding-the-legislation/40_legislation_evolutions.html - - Args: - variable: The variable to which the formula will be added - start_instant: The date on which the formula will take effect - formula: The formula to be added - - Returns: - The variable with the evolved formula - """ - - # Add formula to variable's "formulas" SortedDict - variable.formulas[self.start_instant] = formula - - return variable - # - - def _neutralized_formula(self, variable: Variable) -> Callable: - """ - Return a formula that itself returns the default value of a variable. - - Args: - variable: The variable to be neutralized - - Returns: - The neutralized formula - """ - return lambda population, period, parameters: variable.default_value - - # Validate start instant - - # Default outputs method of some sort? \ No newline at end of file +from policyengine_core.errors import ( + VariableNotFoundError, + VariableNameConflictError, +) + + +class StructuralReform: # Should this inherit from Reform and/or TaxBenefitSystem? + + DEFAULT_START_INSTANT = "0000-01-01" + variables: list[Variable] = [] + + def __init__( + self, + tax_benefit_system: TaxBenefitSystem, + start_instant: Annotated[str, "YYYY-MM-DD"] | None = str( + datetime.now().year + ) + + "-01-01", + end_instant: Annotated[str, "YYYY-MM-DD"] | None = None, + ): + # TODO: Validate start_instant and end_instant + + self.tax_benefit_system = tax_benefit_system + self.start_instant = start_instant + self.end_instant = end_instant + + def neutralize_variable(self, name: str) -> Variable: + """ + Neutralize a variable by setting its formula to return the default value + from the StructuralReform's start_instant date onward. + + Args: + name: The name of the variable + + Raises: + VariableNotFoundError: If the variable is not found in the tax benefit system + """ + # Clone variable + fetched_variable: Variable | None = self._fetch_variable(name) + + if fetched_variable is None: + raise VariableNotFoundError( + f"Unable to neutralize {name}; variable not found." + ) + + # Add formula to variable that returns all defaults + neutralized_formula = self._neutralized_formula(fetched_variable) + self._add_formula( + fetched_variable, + neutralized_formula, + self.start_instant, + self.end_instant, + ) + return fetched_variable + + def add_variable(self, variable: Variable) -> None: + """ + Only partially implemented; Add a variable to the StructuralReform. + + Args: + variable: The variable to be added + + Raises: + VariableNameConflictError: If a variable with the same name already exists in the tax benefit system + """ + if not issubclass(variable, Variable): + raise TypeError( + "Variable must be an instance of the Variable class." + ) + + # Attempt to fetch variable + fetched_variable: Variable | None = self._fetch_variable( + variable.__name__ + ) + + if fetched_variable is not None: + raise VariableNameConflictError( + f"Unable to add {variable.__name__}; variable with the same name already exists." + ) + + # TODO: Likely need to do something to add the variable to the TBS + + # Create variable with neutralized formula until start date, then valid formula + neutralized_formula = self._neutralized_formula(variable) + self._add_formula( + variable, + neutralized_formula, + self.DEFAULT_START_INSTANT, + self.start_instant, + ) + self._add_formula( + variable, variable.formula, self.start_instant, self.end_instant + ) + + def update_variable(self, variable: Variable) -> Variable: + """ + Update a variable in the tax benefit system; if the variable does not + yet exist, it will be added. + + Args: + variable: The variable to be updated + """ + + # Ensure variable is a Variable + if not issubclass(variable, Variable): + raise TypeError( + "Variable must be an instance of the Variable class." + ) + + # Fetch variable + fetched_variable: Variable | None = self._fetch_variable( + variable.__name__ + ) + + # If variable doesn't exist, run self.add_variable + if fetched_variable is None: + self.add_variable(variable) + return + + # Otherwise, add new formula to existing variable + self._add_formula( + fetched_variable, + variable.formula, + self.start_instant, + self.end_instant, + ) + return fetched_variable + + def _fetch_variable(self, name: str) -> Variable | None: + """ + Fetch the variable by reference from the tax benefit system. + + Args: + name: The name of the variable + """ + return self.tax_benefit_system.get_variable(name) + + # Method to modify metadata based on new items? + + # Method to add formula based on date + def _add_formula( + self, + variable: Variable, + formula: Callable, + start_instant: Annotated[str, "YYYY-MM-DD"], + end_instant: Annotated[str, "YYYY-MM-DD"] | None, + ) -> Variable: + """ + Add an evolved formula, beginning at start_instant and ending at end_instant, + to a variable, and return said variable. + + For more on evolved formulas, consult + https://openfisca.org/doc/coding-the-legislation/40_legislation_evolutions.html + + Args: + variable: The variable to which the formula will be added + formula: The formula to be added + start_instant: The date on which the formula will take effect + end_instant: The final effective date of the formula; any future formula + begins the next day; if None, the formula will be applied indefinitely + + Returns: + The variable with the evolved formula + """ + + # Determine if there's already a formula at our exact start_instant + if start_instant in variable.formulas.keys(): + # If so, save it in case we need it later + old_formula = variable.formulas[start_instant] + + # Add formula to variable's formulas + variable.formulas.update({start_instant: formula}) + + # If no end_instant, remove all formulas after start_instant + if end_instant is None: + for date in variable.formulas.keys(): + if date > start_instant: + variable.formulas.pop(date) + return + + # Otherwise, only remove formulas within interval + for date in variable.formulas.keys(): + if date > start_instant and date <= end_instant: + variable.formulas.pop(date) + + # If there's no formula at end_instant + 1 day, + # add the old formula back in + typecast_end_instant: Instant = instant(end_instant) + next_formula_start: str = str(typecast_end_instant.offset(1, "day")) + + if next_formula_start not in variable.formulas.keys(): + variable.formulas[next_formula_start] = old_formula + + return variable + + def _neutralized_formula(self, variable: Variable) -> Callable: + """ + Return a formula that itself returns the default value of a variable. + + Args: + variable: The variable to be neutralized + + Returns: + The neutralized formula + """ + return lambda population, period, parameters: variable.default_value + + # Validate start instant + + # Default outputs method of some sort? From 5303513b8f5801e30970a245ba583a632260298a Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Fri, 27 Dec 2024 00:02:26 +0100 Subject: [PATCH 03/18] feat: Implement add_variable method --- .../reforms/structural_reform.py | 65 ++++++++++--------- 1 file changed, 36 insertions(+), 29 deletions(-) diff --git a/policyengine_core/reforms/structural_reform.py b/policyengine_core/reforms/structural_reform.py index 920fb0f7b..19e6289ee 100644 --- a/policyengine_core/reforms/structural_reform.py +++ b/policyengine_core/reforms/structural_reform.py @@ -58,7 +58,7 @@ def neutralize_variable(self, name: str) -> Variable: ) return fetched_variable - def add_variable(self, variable: Variable) -> None: + def add_variable(self, variable: Variable) -> Variable: """ Only partially implemented; Add a variable to the StructuralReform. @@ -83,20 +83,26 @@ def add_variable(self, variable: Variable) -> None: f"Unable to add {variable.__name__}; variable with the same name already exists." ) - # TODO: Likely need to do something to add the variable to the TBS + # Insert variable into the tax-benefit system + added_variable = self.tax_benefit_system.add_variable(variable) # Create variable with neutralized formula until start date, then valid formula neutralized_formula = self._neutralized_formula(variable) self._add_formula( - variable, + added_variable, neutralized_formula, self.DEFAULT_START_INSTANT, self.start_instant, ) self._add_formula( - variable, variable.formula, self.start_instant, self.end_instant + added_variable, + variable.formula, + self.start_instant, + self.end_instant, ) + return added_variable + def update_variable(self, variable: Variable) -> Variable: """ Update a variable in the tax benefit system; if the variable does not @@ -119,8 +125,7 @@ def update_variable(self, variable: Variable) -> Variable: # If variable doesn't exist, run self.add_variable if fetched_variable is None: - self.add_variable(variable) - return + return self.add_variable(variable) # Otherwise, add new formula to existing variable self._add_formula( @@ -151,7 +156,7 @@ def _add_formula( end_instant: Annotated[str, "YYYY-MM-DD"] | None, ) -> Variable: """ - Add an evolved formula, beginning at start_instant and ending at end_instant, + Mutatively add an evolved formula, beginning at start_instant and ending at end_instant, to a variable, and return said variable. For more on evolved formulas, consult @@ -167,34 +172,26 @@ def _add_formula( Returns: The variable with the evolved formula """ + # Prior to manipulation, get formula at end_instant + 1 day + next_formula: Callable | None = None + if end_instant is not None: + next_formula_start = self._get_next_day(end_instant) + next_formula = variable.get_formula(next_formula_start) - # Determine if there's already a formula at our exact start_instant - if start_instant in variable.formulas.keys(): - # If so, save it in case we need it later - old_formula = variable.formulas[start_instant] - - # Add formula to variable's formulas + # Insert formula into variable's formulas at start_instant variable.formulas.update({start_instant: formula}) - # If no end_instant, remove all formulas after start_instant - if end_instant is None: - for date in variable.formulas.keys(): - if date > start_instant: - variable.formulas.pop(date) - return - - # Otherwise, only remove formulas within interval + # Remove all formulas between start_instant and end_instant (or into perpetuity + # if end_instant is None) for date in variable.formulas.keys(): - if date > start_instant and date <= end_instant: + if date > start_instant and ( + date <= end_instant or end_instant is None + ): variable.formulas.pop(date) - # If there's no formula at end_instant + 1 day, - # add the old formula back in - typecast_end_instant: Instant = instant(end_instant) - next_formula_start: str = str(typecast_end_instant.offset(1, "day")) - - if next_formula_start not in variable.formulas.keys(): - variable.formulas[next_formula_start] = old_formula + # If end_instant, add back in formula at end_instant + 1 day + if end_instant is not None: + variable.formulas[next_formula_start] = next_formula return variable @@ -210,6 +207,16 @@ def _neutralized_formula(self, variable: Variable) -> Callable: """ return lambda population, period, parameters: variable.default_value + def _get_next_day(self, date: str) -> str: + """ + Return the date of the day following the input date. + + Args: + date: The date from which to calculate the next day + """ + typed_date: Instant = instant(date) + return str(typed_date.offset(1, "day")) + # Validate start instant # Default outputs method of some sort? From c8527e2505ffed26895909041cbed5bfcce94695 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Fri, 27 Dec 2024 00:11:42 +0100 Subject: [PATCH 04/18] feat: Transformation log --- .../reforms/structural_reform.py | 48 ++++++++++++++++++- 1 file changed, 47 insertions(+), 1 deletion(-) diff --git a/policyengine_core/reforms/structural_reform.py b/policyengine_core/reforms/structural_reform.py index 19e6289ee..c93e1c036 100644 --- a/policyengine_core/reforms/structural_reform.py +++ b/policyengine_core/reforms/structural_reform.py @@ -1,5 +1,6 @@ -from typing import Annotated, Callable +from typing import Annotated, Callable, Literal from datetime import datetime +from dataclasses import dataclass from policyengine_core.periods import instant, Instant from policyengine_core.variables import Variable from policyengine_core.taxbenefitsystems import TaxBenefitSystem @@ -9,10 +10,23 @@ ) +@dataclass +class TransformationLogItem: + """ + A log item for a transformation applied to a variable. + """ + + variable_name: str + transformation: Literal["neutralize", "add", "update"] + start_instant: Annotated[str, "YYYY-MM-DD"] + end_instant: Annotated[str, "YYYY-MM-DD"] | None + + class StructuralReform: # Should this inherit from Reform and/or TaxBenefitSystem? DEFAULT_START_INSTANT = "0000-01-01" variables: list[Variable] = [] + transformation_log: list[TransformationLogItem] = [] def __init__( self, @@ -56,6 +70,17 @@ def neutralize_variable(self, name: str) -> Variable: self.start_instant, self.end_instant, ) + + # Log transformation + self.transformation_log.append( + TransformationLogItem( + variable_name=name, + transformation="neutralize", + start_instant=self.start_instant, + end_instant=self.end_instant, + ) + ) + return fetched_variable def add_variable(self, variable: Variable) -> Variable: @@ -101,6 +126,16 @@ def add_variable(self, variable: Variable) -> Variable: self.end_instant, ) + # Log transformation + self.transformation_log.append( + TransformationLogItem( + variable_name=variable.__name__, + transformation="add", + start_instant=self.start_instant, + end_instant=self.end_instant, + ) + ) + return added_variable def update_variable(self, variable: Variable) -> Variable: @@ -134,6 +169,17 @@ def update_variable(self, variable: Variable) -> Variable: self.start_instant, self.end_instant, ) + + # Log transformation + self.transformation_log.append( + TransformationLogItem( + variable_name=variable.__name__, + transformation="update", + start_instant=self.start_instant, + end_instant=self.end_instant, + ) + ) + return fetched_variable def _fetch_variable(self, name: str) -> Variable | None: From 723ddf7cde03e457c737df2d3b7f647aaad4eff4 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Fri, 27 Dec 2024 00:41:05 +0100 Subject: [PATCH 05/18] fix: Fix adding of new variable --- .../reforms/structural_reform.py | 46 ++++++++++--------- 1 file changed, 24 insertions(+), 22 deletions(-) diff --git a/policyengine_core/reforms/structural_reform.py b/policyengine_core/reforms/structural_reform.py index c93e1c036..d40ebbd44 100644 --- a/policyengine_core/reforms/structural_reform.py +++ b/policyengine_core/reforms/structural_reform.py @@ -1,7 +1,7 @@ from typing import Annotated, Callable, Literal from datetime import datetime from dataclasses import dataclass -from policyengine_core.periods import instant, Instant +import inspect from policyengine_core.variables import Variable from policyengine_core.taxbenefitsystems import TaxBenefitSystem from policyengine_core.errors import ( @@ -37,6 +37,16 @@ def __init__( + "-01-01", end_instant: Annotated[str, "YYYY-MM-DD"] | None = None, ): + """ + Create a new StructuralReform. + + Args: + tax_benefit_system: The tax benefit system to which the reform will be applied + start_instant: The date on which the reform will take effect + end_instant: The date on which the reform ends, exclusive (i.e., + the reform will be applied up to but not including this date); if None, + the reform will be applied indefinitely + """ # TODO: Validate start_instant and end_instant self.tax_benefit_system = tax_benefit_system @@ -108,17 +118,19 @@ def add_variable(self, variable: Variable) -> Variable: f"Unable to add {variable.__name__}; variable with the same name already exists." ) - # Insert variable into the tax-benefit system + # Insert variable into the tax-benefit system; this will apply default formula over + # entire period, which we will modify below added_variable = self.tax_benefit_system.add_variable(variable) - # Create variable with neutralized formula until start date, then valid formula + # First, neutralize entire period neutralized_formula = self._neutralized_formula(variable) self._add_formula( added_variable, neutralized_formula, self.DEFAULT_START_INSTANT, - self.start_instant, ) + + # Then, re-add formula in order to format correctly self._add_formula( added_variable, variable.formula, @@ -199,7 +211,7 @@ def _add_formula( variable: Variable, formula: Callable, start_instant: Annotated[str, "YYYY-MM-DD"], - end_instant: Annotated[str, "YYYY-MM-DD"] | None, + end_instant: Annotated[str, "YYYY-MM-DD"] | None = None, ) -> Variable: """ Mutatively add an evolved formula, beginning at start_instant and ending at end_instant, @@ -212,8 +224,9 @@ def _add_formula( variable: The variable to which the formula will be added formula: The formula to be added start_instant: The date on which the formula will take effect - end_instant: The final effective date of the formula; any future formula - begins the next day; if None, the formula will be applied indefinitely + end_instant: The date on which the formula ends, exclusive (i.e., + the formula will be applied up to but not including this date); if None, + the formula will be applied indefinitely Returns: The variable with the evolved formula @@ -221,8 +234,7 @@ def _add_formula( # Prior to manipulation, get formula at end_instant + 1 day next_formula: Callable | None = None if end_instant is not None: - next_formula_start = self._get_next_day(end_instant) - next_formula = variable.get_formula(next_formula_start) + next_formula = variable.get_formula(end_instant) # Insert formula into variable's formulas at start_instant variable.formulas.update({start_instant: formula}) @@ -231,13 +243,13 @@ def _add_formula( # if end_instant is None) for date in variable.formulas.keys(): if date > start_instant and ( - date <= end_instant or end_instant is None + end_instant is None or date <= end_instant ): variable.formulas.pop(date) - # If end_instant, add back in formula at end_instant + 1 day + # If end_instant, add back in formula at end_instant if end_instant is not None: - variable.formulas[next_formula_start] = next_formula + variable.formulas[end_instant] = next_formula return variable @@ -253,16 +265,6 @@ def _neutralized_formula(self, variable: Variable) -> Callable: """ return lambda population, period, parameters: variable.default_value - def _get_next_day(self, date: str) -> str: - """ - Return the date of the day following the input date. - - Args: - date: The date from which to calculate the next day - """ - typed_date: Instant = instant(date) - return str(typed_date.offset(1, "day")) - # Validate start instant # Default outputs method of some sort? From 7152e761fcf0e346143f3a6184912d9b1041c752 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Fri, 27 Dec 2024 01:23:48 +0100 Subject: [PATCH 06/18] feat: First stab at testing --- .../reforms/structural_reform.py | 2 +- tests/core/{ => reforms}/test_reforms.py | 0 tests/core/reforms/test_structural_reforms.py | 135 ++++++++++++++++++ 3 files changed, 136 insertions(+), 1 deletion(-) rename tests/core/{ => reforms}/test_reforms.py (100%) create mode 100644 tests/core/reforms/test_structural_reforms.py diff --git a/policyengine_core/reforms/structural_reform.py b/policyengine_core/reforms/structural_reform.py index d40ebbd44..d0cadeacf 100644 --- a/policyengine_core/reforms/structural_reform.py +++ b/policyengine_core/reforms/structural_reform.py @@ -263,7 +263,7 @@ def _neutralized_formula(self, variable: Variable) -> Callable: Returns: The neutralized formula """ - return lambda population, period, parameters: variable.default_value + return lambda: variable.default_value # Validate start instant diff --git a/tests/core/test_reforms.py b/tests/core/reforms/test_reforms.py similarity index 100% rename from tests/core/test_reforms.py rename to tests/core/reforms/test_reforms.py diff --git a/tests/core/reforms/test_structural_reforms.py b/tests/core/reforms/test_structural_reforms.py new file mode 100644 index 000000000..32371a795 --- /dev/null +++ b/tests/core/reforms/test_structural_reforms.py @@ -0,0 +1,135 @@ +import warnings + +import pytest +from datetime import datetime +from policyengine_core.errors import VariableNotFoundError +from policyengine_core.model_api import * +from policyengine_core.country_template.entities import Person +from policyengine_core.variables import Variable +from policyengine_core.reforms import StructuralReform + + +class test_variable_to_add(Variable): + value_type = str + default_value = "Returning default value" + entity = Person + label = "Maxwell" + definition_period = YEAR + + def formula(): + return "Returning value from formula" + + +class test_existing_variable(Variable): + value_type = str + default_value = "Returning default value, existing variable" + entity = Person + label = "Dworkin" + definition_period = YEAR + + def formula(): + return "Returning value from formula, existing variable" + + +def test_structural_reform_init(tax_benefit_system): + # Given an empty tax-benefit system... + + # When a new structural reform is created with default settings... + test_reform = StructuralReform(tax_benefit_system) + + # Then the reform is created successfully for the current year + assert test_reform.tax_benefit_system == tax_benefit_system + assert test_reform.start_instant == str(datetime.now().year) + "-01-01" + assert test_reform.end_instant == None + + +def test_structural_reform_init_with_dates(tax_benefit_system): + # Given an empty tax-benefit system... + + # When a new structural reform is created with specific dates... + reform = StructuralReform(tax_benefit_system, "2020-01-01", "2021-01-01") + + # Then the reform is created successfully for the specified dates + assert reform.tax_benefit_system == tax_benefit_system + assert reform.start_instant == "2020-01-01" + assert reform.end_instant == "2021-01-01" + + +def test_empty_tbs_endless_structural_reform_add_variable(tax_benefit_system): + # Given an empty tax-benefit system with an endless structural reform... + test_reform = StructuralReform( + tax_benefit_system, + "2025-01-01", + ) + + # When a new variable is added... + test_reform.add_variable(test_variable_to_add) + + # Then add_variable(test_var) adds new variable with proper formulas + assert ( + "test_variable_to_add" + in test_reform.tax_benefit_system.variables.keys() + ) + + added_test_variable = test_reform.tax_benefit_system.get_variable( + "test_variable_to_add" + ) + assert added_test_variable.value_type == str + assert added_test_variable.label == "Maxwell" + # TODO - Figure out how to test formula additions + + +def test_empty_tbs_endless_structural_reform_update_variable( + tax_benefit_system, +): + # Given an empty tax-benefit system with an endless structural reform... + test_reform = StructuralReform( + tax_benefit_system, + "2025-01-01", + ) + + # When update_variable is called on a variable that does not exist... + test_reform.update_variable(test_variable_to_add) + + # Then update_variable(test_var) adds new variable with proper formulas + assert test_variable_to_add in test_reform.tax_benefit_system.variables + + added_test_variable = test_reform.tax_benefit_system.get_variable( + "test_variable_to_add" + ) + assert ( + added_test_variable.get_formula("2025-01-01")() + == "Returning value from formula" + ) + assert ( + added_test_variable.get_formula("2021-01-01")() + == "Returning default value" + ) + + +def test_empty_tbs_endless_structural_reform_neutralize_variable( + tax_benefit_system, +): + # Given an empty tax-benefit system with an endless structural reform... + test_reform = StructuralReform( + tax_benefit_system, + "2025-01-01", + ) + + # When neutralize_variable is called on a variable that does not exist... + + # Then neutralize_variable(test_var) raises error + with pytest.raises(VariableNotFoundError): + test_reform.neutralize_variable("test_variable_to_add") + + +# Given a TBS with a variable test_var... +# add_variable(test_var) raises error + +# update_variable(test_var) updates variable + +# neutralize_variable(test_var) neutralizes variable + + +# Given a TBS with a complex structural reform... +# The reform successfully adds a variable, updates a variable, then neutralizes a variable From 8547f8dee9a278715c8f5b78c0d19bf1cd6f8213 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Fri, 27 Dec 2024 02:05:35 +0100 Subject: [PATCH 07/18] test: Correct tests --- .../reforms/structural_reform.py | 5 ++- tests/core/reforms/test_structural_reforms.py | 44 +++++++++++++------ 2 files changed, 33 insertions(+), 16 deletions(-) diff --git a/policyengine_core/reforms/structural_reform.py b/policyengine_core/reforms/structural_reform.py index d0cadeacf..359e7e981 100644 --- a/policyengine_core/reforms/structural_reform.py +++ b/policyengine_core/reforms/structural_reform.py @@ -64,12 +64,13 @@ def neutralize_variable(self, name: str) -> Variable: Raises: VariableNotFoundError: If the variable is not found in the tax benefit system """ - # Clone variable + # Fetch variable fetched_variable: Variable | None = self._fetch_variable(name) if fetched_variable is None: raise VariableNotFoundError( - f"Unable to neutralize {name}; variable not found." + f"Unable to neutralize {name}; variable not found.", + self.tax_benefit_system, ) # Add formula to variable that returns all defaults diff --git a/tests/core/reforms/test_structural_reforms.py b/tests/core/reforms/test_structural_reforms.py index 32371a795..5bc120122 100644 --- a/tests/core/reforms/test_structural_reforms.py +++ b/tests/core/reforms/test_structural_reforms.py @@ -31,34 +31,38 @@ def formula(): return "Returning value from formula, existing variable" -def test_structural_reform_init(tax_benefit_system): +def test_structural_reform_init(isolated_tax_benefit_system): # Given an empty tax-benefit system... # When a new structural reform is created with default settings... - test_reform = StructuralReform(tax_benefit_system) + test_reform = StructuralReform(isolated_tax_benefit_system) # Then the reform is created successfully for the current year - assert test_reform.tax_benefit_system == tax_benefit_system + assert test_reform.tax_benefit_system == isolated_tax_benefit_system assert test_reform.start_instant == str(datetime.now().year) + "-01-01" assert test_reform.end_instant == None -def test_structural_reform_init_with_dates(tax_benefit_system): +def test_structural_reform_init_with_dates(isolated_tax_benefit_system): # Given an empty tax-benefit system... # When a new structural reform is created with specific dates... - reform = StructuralReform(tax_benefit_system, "2020-01-01", "2021-01-01") + reform = StructuralReform( + isolated_tax_benefit_system, "2020-01-01", "2021-01-01" + ) # Then the reform is created successfully for the specified dates - assert reform.tax_benefit_system == tax_benefit_system + assert reform.tax_benefit_system == isolated_tax_benefit_system assert reform.start_instant == "2020-01-01" assert reform.end_instant == "2021-01-01" -def test_empty_tbs_endless_structural_reform_add_variable(tax_benefit_system): +def test_empty_tbs_endless_structural_reform_add_variable( + isolated_tax_benefit_system, +): # Given an empty tax-benefit system with an endless structural reform... test_reform = StructuralReform( - tax_benefit_system, + isolated_tax_benefit_system, "2025-01-01", ) @@ -76,15 +80,22 @@ def test_empty_tbs_endless_structural_reform_add_variable(tax_benefit_system): ) assert added_test_variable.value_type == str assert added_test_variable.label == "Maxwell" - # TODO - Figure out how to test formula additions + assert ( + added_test_variable.get_formula("2025-01-01")() + == "Returning value from formula" + ) + assert ( + added_test_variable.get_formula("2021-01-01")() + == "Returning default value" + ) def test_empty_tbs_endless_structural_reform_update_variable( - tax_benefit_system, + isolated_tax_benefit_system, ): # Given an empty tax-benefit system with an endless structural reform... test_reform = StructuralReform( - tax_benefit_system, + isolated_tax_benefit_system, "2025-01-01", ) @@ -92,11 +103,16 @@ def test_empty_tbs_endless_structural_reform_update_variable( test_reform.update_variable(test_variable_to_add) # Then update_variable(test_var) adds new variable with proper formulas - assert test_variable_to_add in test_reform.tax_benefit_system.variables + assert ( + "test_variable_to_add" + in test_reform.tax_benefit_system.variables.keys() + ) added_test_variable = test_reform.tax_benefit_system.get_variable( "test_variable_to_add" ) + assert added_test_variable.value_type == str + assert added_test_variable.label == "Maxwell" assert ( added_test_variable.get_formula("2025-01-01")() == "Returning value from formula" @@ -108,11 +124,11 @@ def test_empty_tbs_endless_structural_reform_update_variable( def test_empty_tbs_endless_structural_reform_neutralize_variable( - tax_benefit_system, + isolated_tax_benefit_system, ): # Given an empty tax-benefit system with an endless structural reform... test_reform = StructuralReform( - tax_benefit_system, + isolated_tax_benefit_system, "2025-01-01", ) From f0e540c97a3826e06b26717775c4c546f5d14c4c Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Fri, 27 Dec 2024 02:29:24 +0100 Subject: [PATCH 08/18] test: Continue with testing --- tests/core/reforms/test_structural_reforms.py | 309 +++++++++++------- 1 file changed, 195 insertions(+), 114 deletions(-) diff --git a/tests/core/reforms/test_structural_reforms.py b/tests/core/reforms/test_structural_reforms.py index 5bc120122..763d7e884 100644 --- a/tests/core/reforms/test_structural_reforms.py +++ b/tests/core/reforms/test_structural_reforms.py @@ -2,11 +2,33 @@ import pytest from datetime import datetime -from policyengine_core.errors import VariableNotFoundError +from policyengine_core.errors import ( + VariableNotFoundError, + VariableNameConflictError, +) from policyengine_core.model_api import * from policyengine_core.country_template.entities import Person from policyengine_core.variables import Variable from policyengine_core.reforms import StructuralReform +from policyengine_core.country_template import CountryTaxBenefitSystem + + +@pytest.fixture(scope="module") +def prefilled_tax_benefit_system(): + + class test_variable_to_add(Variable): + value_type = str + default_value = "Returning default value [pre-loaded]" + entity = Person + label = "Dworkin" + definition_period = YEAR + + def formula(): + return "Returning value from formula [pre-loaded]" + + tax_benefit_system = CountryTaxBenefitSystem() + tax_benefit_system.add_variable(test_variable_to_add) + return tax_benefit_system class test_variable_to_add(Variable): @@ -20,131 +42,190 @@ def formula(): return "Returning value from formula" -class test_existing_variable(Variable): - value_type = str - default_value = "Returning default value, existing variable" - entity = Person - label = "Dworkin" - definition_period = YEAR - - def formula(): - return "Returning value from formula, existing variable" - - -def test_structural_reform_init(isolated_tax_benefit_system): +class TestGivenEmptyTaxBenefitSystem: # Given an empty tax-benefit system... + def test_structural_reform_init(self, isolated_tax_benefit_system): - # When a new structural reform is created with default settings... - test_reform = StructuralReform(isolated_tax_benefit_system) - - # Then the reform is created successfully for the current year - assert test_reform.tax_benefit_system == isolated_tax_benefit_system - assert test_reform.start_instant == str(datetime.now().year) + "-01-01" - assert test_reform.end_instant == None + # When a new structural reform is created with default settings... + test_reform = StructuralReform(isolated_tax_benefit_system) + # Then the reform is created successfully for the current year + assert test_reform.tax_benefit_system == isolated_tax_benefit_system + assert test_reform.start_instant == str(datetime.now().year) + "-01-01" + assert test_reform.end_instant == None -def test_structural_reform_init_with_dates(isolated_tax_benefit_system): - # Given an empty tax-benefit system... - - # When a new structural reform is created with specific dates... - reform = StructuralReform( - isolated_tax_benefit_system, "2020-01-01", "2021-01-01" - ) + def test_structural_reform_init_with_dates( + self, isolated_tax_benefit_system + ): - # Then the reform is created successfully for the specified dates - assert reform.tax_benefit_system == isolated_tax_benefit_system - assert reform.start_instant == "2020-01-01" - assert reform.end_instant == "2021-01-01" + # When a new structural reform is created with specific dates... + reform = StructuralReform( + isolated_tax_benefit_system, "2020-01-01", "2021-01-01" + ) + # Then the reform is created successfully for the specified dates + assert reform.tax_benefit_system == isolated_tax_benefit_system + assert reform.start_instant == "2020-01-01" + assert reform.end_instant == "2021-01-01" -def test_empty_tbs_endless_structural_reform_add_variable( - isolated_tax_benefit_system, -): - # Given an empty tax-benefit system with an endless structural reform... - test_reform = StructuralReform( + def test_add_variable_no_end_dates( + self, isolated_tax_benefit_system, - "2025-01-01", - ) - - # When a new variable is added... - test_reform.add_variable(test_variable_to_add) - - # Then add_variable(test_var) adds new variable with proper formulas - assert ( - "test_variable_to_add" - in test_reform.tax_benefit_system.variables.keys() - ) - - added_test_variable = test_reform.tax_benefit_system.get_variable( - "test_variable_to_add" - ) - assert added_test_variable.value_type == str - assert added_test_variable.label == "Maxwell" - assert ( - added_test_variable.get_formula("2025-01-01")() - == "Returning value from formula" - ) - assert ( - added_test_variable.get_formula("2021-01-01")() - == "Returning default value" - ) - - -def test_empty_tbs_endless_structural_reform_update_variable( - isolated_tax_benefit_system, -): - # Given an empty tax-benefit system with an endless structural reform... - test_reform = StructuralReform( + ): + # Given an empty tax-benefit system with an endless structural reform... + test_reform = StructuralReform( + isolated_tax_benefit_system, + "2025-01-01", + ) + + # When a new variable is added... + test_reform.add_variable(test_variable_to_add) + + # Then add_variable(test_var) adds new variable with proper formulas + assert ( + "test_variable_to_add" + in test_reform.tax_benefit_system.variables.keys() + ) + + added_test_variable = test_reform.tax_benefit_system.get_variable( + "test_variable_to_add" + ) + assert added_test_variable.value_type == str + assert added_test_variable.label == "Maxwell" + assert ( + added_test_variable.get_formula("2025-01-01")() + == "Returning value from formula" + ) + assert ( + added_test_variable.get_formula("2021-01-01")() + == "Returning default value" + ) + + def test_update_variable_no_end_dates( + self, isolated_tax_benefit_system, - "2025-01-01", - ) - - # When update_variable is called on a variable that does not exist... - test_reform.update_variable(test_variable_to_add) - - # Then update_variable(test_var) adds new variable with proper formulas - assert ( - "test_variable_to_add" - in test_reform.tax_benefit_system.variables.keys() - ) - - added_test_variable = test_reform.tax_benefit_system.get_variable( - "test_variable_to_add" - ) - assert added_test_variable.value_type == str - assert added_test_variable.label == "Maxwell" - assert ( - added_test_variable.get_formula("2025-01-01")() - == "Returning value from formula" - ) - assert ( - added_test_variable.get_formula("2021-01-01")() - == "Returning default value" - ) - - -def test_empty_tbs_endless_structural_reform_neutralize_variable( - isolated_tax_benefit_system, -): - # Given an empty tax-benefit system with an endless structural reform... - test_reform = StructuralReform( + ): + # Given an empty tax-benefit system with an endless structural reform... + test_reform = StructuralReform( + isolated_tax_benefit_system, + "2025-01-01", + ) + + # When update_variable is called on a variable that does not exist... + test_reform.update_variable(test_variable_to_add) + + # Then update_variable(test_var) adds new variable with proper formulas + assert ( + "test_variable_to_add" + in test_reform.tax_benefit_system.variables.keys() + ) + + added_test_variable = test_reform.tax_benefit_system.get_variable( + "test_variable_to_add" + ) + assert added_test_variable.value_type == str + assert added_test_variable.label == "Maxwell" + assert ( + added_test_variable.get_formula("2025-01-01")() + == "Returning value from formula" + ) + assert ( + added_test_variable.get_formula("2021-01-01")() + == "Returning default value" + ) + + def test_neutralize_variable_no_end_dates( + self, isolated_tax_benefit_system, - "2025-01-01", - ) - - # When neutralize_variable is called on a variable that does not exist... - - # Then neutralize_variable(test_var) raises error - with pytest.raises(VariableNotFoundError): + ): + # Given an empty tax-benefit system with an endless structural reform... + test_reform = StructuralReform( + isolated_tax_benefit_system, + "2025-01-01", + ) + + # When neutralize_variable is called on a variable that does not exist... + + # Then neutralize_variable(test_var) raises error + with pytest.raises(VariableNotFoundError): + test_reform.neutralize_variable("test_variable_to_add") + + +class TestGivenPreFilledTaxBenefitSystem: + + # Given a tax-benefit system with variable test_variable_to_add... + + def test_add_variable_no_end_dates( + self, + prefilled_tax_benefit_system, + ): + # Given a tax-benefit system with an endless structural reform... + test_reform = StructuralReform( + prefilled_tax_benefit_system, + "2025-01-01", + ) + + # When a new variable is added... + + # Then add_variable(test_var) raises error + with pytest.raises(VariableNameConflictError): + test_reform.add_variable(test_variable_to_add) + + def test_update_variable_no_end_dates( + self, + prefilled_tax_benefit_system, + ): + # Given a tax-benefit system with an endless structural reform... + test_reform = StructuralReform( + prefilled_tax_benefit_system, + "2025-01-01", + ) + + # When update_variable is called on a variable that already exists... + test_reform.update_variable(test_variable_to_add) + + # Then update_variable(test_var) updates variable with proper formulas + added_test_variable = test_reform.tax_benefit_system.get_variable( + "test_variable_to_add" + ) + assert added_test_variable.value_type == str + assert added_test_variable.label == "Maxwell" + assert ( + added_test_variable.get_formula("2025-01-01")() + == "Returning value from formula" + ) + assert ( + added_test_variable.get_formula("2021-01-01")() + == "Returning default value" + ) + + def test_neutralize_variable_no_end_dates( + self, + prefilled_tax_benefit_system, + ): + # Given a tax-benefit system with an endless structural reform... + test_reform = StructuralReform( + prefilled_tax_benefit_system, + "2025-01-01", + ) + + # When neutralize_variable is called on a variable that already exists... test_reform.neutralize_variable("test_variable_to_add") - -# Given a TBS with a variable test_var... -# add_variable(test_var) raises error - -# update_variable(test_var) updates variable - -# neutralize_variable(test_var) neutralizes variable + # Then neutralize_variable(test_var) neutralizes variable + added_test_variable = test_reform.tax_benefit_system.get_variable( + "test_variable_to_add" + ) + assert added_test_variable.value_type == str + assert added_test_variable.label == "Maxwell" + assert ( + added_test_variable.get_formula("2025-01-01")() + == "Returning default value" + ) + assert ( + added_test_variable.get_formula("2021-01-01")() + == "Returning default value" + ) # Given a TBS with a complex structural reform... From 2a41111a0a6c569fbe0812b01c1966da3c585ce6 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Fri, 27 Dec 2024 23:54:09 +0100 Subject: [PATCH 09/18] feat: Date validation for structural reforms, with tests --- .../reforms/structural_reform.py | 28 +++++++++++++++++-- tests/core/reforms/test_structural_reforms.py | 24 ++++++++++++++++ 2 files changed, 49 insertions(+), 3 deletions(-) diff --git a/policyengine_core/reforms/structural_reform.py b/policyengine_core/reforms/structural_reform.py index 359e7e981..cc509a7cd 100644 --- a/policyengine_core/reforms/structural_reform.py +++ b/policyengine_core/reforms/structural_reform.py @@ -1,8 +1,8 @@ -from typing import Annotated, Callable, Literal +from typing import Annotated, Callable, Literal, Any from datetime import datetime from dataclasses import dataclass -import inspect from policyengine_core.variables import Variable +from policyengine_core.periods import config from policyengine_core.taxbenefitsystems import TaxBenefitSystem from policyengine_core.errors import ( VariableNotFoundError, @@ -47,7 +47,9 @@ def __init__( the reform will be applied up to but not including this date); if None, the reform will be applied indefinitely """ - # TODO: Validate start_instant and end_instant + self._validate_instant(start_instant) + if end_instant is not None: + self._validate_instant(end_instant) self.tax_benefit_system = tax_benefit_system self.start_instant = start_instant @@ -267,5 +269,25 @@ def _neutralized_formula(self, variable: Variable) -> Callable: return lambda: variable.default_value # Validate start instant + def _validate_instant(self, instant: Any) -> bool: + """ + Validate an instant. + + Args: + instant: The instant to be validated + """ + if not isinstance(instant, str): + raise TypeError( + "Instant must be a string in the format 'YYYY-MM-DD'." + ) + + if not config.INSTANT_PATTERN.match(instant): + raise ValueError( + "'{}' is not a valid instant. Instants are described using the 'YYYY-MM-DD' format, for instance '2015-06-15'.".format( + instant + ) + ) + + return True # Default outputs method of some sort? diff --git a/tests/core/reforms/test_structural_reforms.py b/tests/core/reforms/test_structural_reforms.py index 763d7e884..7392297e6 100644 --- a/tests/core/reforms/test_structural_reforms.py +++ b/tests/core/reforms/test_structural_reforms.py @@ -68,6 +68,30 @@ def test_structural_reform_init_with_dates( assert reform.start_instant == "2020-01-01" assert reform.end_instant == "2021-01-01" + def test_structural_reform_init_with_invalid_date_type( + self, isolated_tax_benefit_system + ): + + # When a new structural reform is created with incorrectly typed dates... + + # Then the reform raises a TypeError + + with pytest.raises(TypeError): + StructuralReform(isolated_tax_benefit_system, "2020-01-01", 15) + + def test_structural_reform_init_with_invalid_date_format( + self, isolated_tax_benefit_system + ): + + # When a new structural reform is created with incorrectly formatted dates... + + # Then the reform raises a ValueError + + with pytest.raises(ValueError): + StructuralReform( + isolated_tax_benefit_system, "2020-01-01", "2020-13-01" + ) + def test_add_variable_no_end_dates( self, isolated_tax_benefit_system, From d5fb8772da43e16e2e2ff9c935da8c8560f7ddc4 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Sat, 28 Dec 2024 00:56:20 +0100 Subject: [PATCH 10/18] feat: Begin to migrate to definition, followed by activation, structure --- .../reforms/structural_reform.py | 174 ++++++++++++------ 1 file changed, 116 insertions(+), 58 deletions(-) diff --git a/policyengine_core/reforms/structural_reform.py b/policyengine_core/reforms/structural_reform.py index cc509a7cd..6f1195da2 100644 --- a/policyengine_core/reforms/structural_reform.py +++ b/policyengine_core/reforms/structural_reform.py @@ -17,45 +17,108 @@ class TransformationLogItem: """ variable_name: str + variable: Variable | None # None is present when using neutralize transformation: Literal["neutralize", "add", "update"] - start_instant: Annotated[str, "YYYY-MM-DD"] - end_instant: Annotated[str, "YYYY-MM-DD"] | None class StructuralReform: # Should this inherit from Reform and/or TaxBenefitSystem? DEFAULT_START_INSTANT = "0000-01-01" - variables: list[Variable] = [] transformation_log: list[TransformationLogItem] = [] - def __init__( - self, - tax_benefit_system: TaxBenefitSystem, - start_instant: Annotated[str, "YYYY-MM-DD"] | None = str( - datetime.now().year + tax_benefit_system: TaxBenefitSystem | None = None + start_instant: Annotated[str, "YYYY-MM-DD"] = DEFAULT_START_INSTANT + end_instant: Annotated[str, "YYYY-MM-DD"] | None = None + + def add_tax_benefit_system(self, tax_benefit_system: TaxBenefitSystem): + """ + Add a tax benefit system to the structural reform. + + Args: + tax_benefit_system: The tax benefit system to be added + """ + if not isinstance(tax_benefit_system, TaxBenefitSystem): + raise TypeError( + "Tax benefit system must be an instance of the TaxBenefitSystem class." + ) + self.tax_benefit_system = tax_benefit_system + + def activate(self, start_instant: Annotated[str, "YYYY-MM-DD"], end_instant: Annotated[str, "YYYY-MM-DD"] | None): + """ + Activate the structural reform. + + Args: + start_instant: The start instant to be added; must be in the format 'YYYY-MM-DD' + end_instant: The end instant to be added; must be in the format 'YYYY-MM-DD' or None + """ + if not self.tax_benefit_system or self.tax_benefit_system is None: + raise ValueError( + "Tax benefit system must be added before start instant." + ) + + self._add_start_instant(start_instant) + self._add_end_instant(end_instant) + self._activate_transformation_log() + + def neutralize_variable(self, name: str): + """ + Neutralize a variable by setting its formula to return the default value + from the StructuralReform's start_instant date to its end_instant date. + + Args: + name: The name of the variable + """ + self.transformation_log.append( + TransformationLogItem( + variable_name=name, + transformation="neutralize", + ) ) - + "-01-01", - end_instant: Annotated[str, "YYYY-MM-DD"] | None = None, - ): + + def add_variable(self, variable: Variable): """ - Create a new StructuralReform. + Add a variable to the StructuralReform. Args: - tax_benefit_system: The tax benefit system to which the reform will be applied - start_instant: The date on which the reform will take effect - end_instant: The date on which the reform ends, exclusive (i.e., - the reform will be applied up to but not including this date); if None, - the reform will be applied indefinitely + variable: The variable to be added """ - self._validate_instant(start_instant) - if end_instant is not None: - self._validate_instant(end_instant) + self.transformation_log.append( + TransformationLogItem( + variable_name=variable.__name__, + variable=variable, + transformation="add", + ) + ) - self.tax_benefit_system = tax_benefit_system - self.start_instant = start_instant - self.end_instant = end_instant + def update_variable(self, variable: Variable): + """ + Update a variable in the tax benefit system; if the variable does not + yet exist, it will be added. + + Args: + variable: The variable to be updated + """ + self.transformation_log.append( + TransformationLogItem( + variable_name=variable.__name__, + variable=variable, + transformation="update", + ) + ) - def neutralize_variable(self, name: str) -> Variable: + def _activate_transformation_log(self): + """ + Activate the transformation log by applying the transformations to the tax benefit system. + """ + for log_item in self.transformation_log: + if log_item.transformation == "neutralize": + self._neutralize_variable(log_item.variable_name) + elif log_item.transformation == "add": + self._add_variable(log_item.variable) + elif log_item.transformation == "update": + self._update_variable(log_item.variable) + + def _neutralize_variable(self, name: str) -> Variable: """ Neutralize a variable by setting its formula to return the default value from the StructuralReform's start_instant date onward. @@ -84,19 +147,9 @@ def neutralize_variable(self, name: str) -> Variable: self.end_instant, ) - # Log transformation - self.transformation_log.append( - TransformationLogItem( - variable_name=name, - transformation="neutralize", - start_instant=self.start_instant, - end_instant=self.end_instant, - ) - ) - return fetched_variable - def add_variable(self, variable: Variable) -> Variable: + def _add_variable(self, variable: Variable) -> Variable: """ Only partially implemented; Add a variable to the StructuralReform. @@ -141,19 +194,9 @@ def add_variable(self, variable: Variable) -> Variable: self.end_instant, ) - # Log transformation - self.transformation_log.append( - TransformationLogItem( - variable_name=variable.__name__, - transformation="add", - start_instant=self.start_instant, - end_instant=self.end_instant, - ) - ) - return added_variable - def update_variable(self, variable: Variable) -> Variable: + def _update_variable(self, variable: Variable) -> Variable: """ Update a variable in the tax benefit system; if the variable does not yet exist, it will be added. @@ -173,9 +216,9 @@ def update_variable(self, variable: Variable) -> Variable: variable.__name__ ) - # If variable doesn't exist, run self.add_variable + # If variable doesn't exist, run self._add_variable if fetched_variable is None: - return self.add_variable(variable) + return self._add_variable(variable) # Otherwise, add new formula to existing variable self._add_formula( @@ -185,16 +228,6 @@ def update_variable(self, variable: Variable) -> Variable: self.end_instant, ) - # Log transformation - self.transformation_log.append( - TransformationLogItem( - variable_name=variable.__name__, - transformation="update", - start_instant=self.start_instant, - end_instant=self.end_instant, - ) - ) - return fetched_variable def _fetch_variable(self, name: str) -> Variable | None: @@ -290,4 +323,29 @@ def _validate_instant(self, instant: Any) -> bool: return True + def _add_start_instant(self, start_instant: Annotated[str, "YYYY-MM-DD"]): + """ + Add a start instant to the structural reform. + + Args: + start_instant: The start instant to be added + """ + + self._validate_instant(start_instant) + self.start_instant = start_instant + + def _add_end_instant(self, end_instant: Annotated[str, "YYYY-MM-DD"] | None): + """ + Add an end instant to the structural reform. + + Args: + end_instant: The end instant to be added + """ + if not self.tax_benefit_system or self.tax_benefit_system is None: + raise ValueError( + "Tax benefit system must be added before end instant." + ) + if end_instant is not None: + self._validate_instant(end_instant) + self.end_instant = end_instant # Default outputs method of some sort? From 74d734739967858398076e20c3230a7f904f0171 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Mon, 30 Dec 2024 22:06:58 +0100 Subject: [PATCH 11/18] feat: Add trigger parameter as explicit arg --- .../reforms/structural_reform.py | 33 +++++++++++++++---- 1 file changed, 26 insertions(+), 7 deletions(-) diff --git a/policyengine_core/reforms/structural_reform.py b/policyengine_core/reforms/structural_reform.py index 6f1195da2..b7a0153db 100644 --- a/policyengine_core/reforms/structural_reform.py +++ b/policyengine_core/reforms/structural_reform.py @@ -2,6 +2,7 @@ from datetime import datetime from dataclasses import dataclass from policyengine_core.variables import Variable +from policyengine_core.parameters import Parameter from policyengine_core.periods import config from policyengine_core.taxbenefitsystems import TaxBenefitSystem from policyengine_core.errors import ( @@ -17,7 +18,7 @@ class TransformationLogItem: """ variable_name: str - variable: Variable | None # None is present when using neutralize + variable: Variable | None # None is present when using neutralize transformation: Literal["neutralize", "add", "update"] @@ -29,6 +30,17 @@ class StructuralReform: # Should this inherit from Reform and/or TaxBenefitSyst tax_benefit_system: TaxBenefitSystem | None = None start_instant: Annotated[str, "YYYY-MM-DD"] = DEFAULT_START_INSTANT end_instant: Annotated[str, "YYYY-MM-DD"] | None = None + trigger_parameter: str = "" + + def __init__(self, trigger_parameter: str): + """ + Initialize a structural reform. + + Args: + trigger_parameter: Path to the parameter that triggers the structural reform; + this parameter must be Boolean + """ + self.trigger_parameter = trigger_parameter def add_tax_benefit_system(self, tax_benefit_system: TaxBenefitSystem): """ @@ -43,7 +55,11 @@ def add_tax_benefit_system(self, tax_benefit_system: TaxBenefitSystem): ) self.tax_benefit_system = tax_benefit_system - def activate(self, start_instant: Annotated[str, "YYYY-MM-DD"], end_instant: Annotated[str, "YYYY-MM-DD"] | None): + def activate( + self, + start_instant: Annotated[str, "YYYY-MM-DD"], + end_instant: Annotated[str, "YYYY-MM-DD"] | None, + ): """ Activate the structural reform. @@ -55,7 +71,7 @@ def activate(self, start_instant: Annotated[str, "YYYY-MM-DD"], end_instant: Ann raise ValueError( "Tax benefit system must be added before start instant." ) - + self._add_start_instant(start_instant) self._add_end_instant(end_instant) self._activate_transformation_log() @@ -106,7 +122,7 @@ def update_variable(self, variable: Variable): ) ) - def _activate_transformation_log(self): + def _activate_transformation_log(self): """ Activate the transformation log by applying the transformations to the tax benefit system. """ @@ -178,7 +194,7 @@ def _add_variable(self, variable: Variable) -> Variable: # entire period, which we will modify below added_variable = self.tax_benefit_system.add_variable(variable) - # First, neutralize entire period + # First, neutralize over entire period neutralized_formula = self._neutralized_formula(variable) self._add_formula( added_variable, @@ -334,7 +350,9 @@ def _add_start_instant(self, start_instant: Annotated[str, "YYYY-MM-DD"]): self._validate_instant(start_instant) self.start_instant = start_instant - def _add_end_instant(self, end_instant: Annotated[str, "YYYY-MM-DD"] | None): + def _add_end_instant( + self, end_instant: Annotated[str, "YYYY-MM-DD"] | None + ): """ Add an end instant to the structural reform. @@ -346,6 +364,7 @@ def _add_end_instant(self, end_instant: Annotated[str, "YYYY-MM-DD"] | None): "Tax benefit system must be added before end instant." ) if end_instant is not None: - self._validate_instant(end_instant) + self._validate_instant(end_instant) self.end_instant = end_instant + # Default outputs method of some sort? From 1bcb636c8a4f5de8bed8f1622ae99b1000213ac6 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Tue, 31 Dec 2024 03:14:15 +0100 Subject: [PATCH 12/18] fix: Modify how reforms are activated --- .../reforms/structural_reform.py | 34 +++++++++---------- 1 file changed, 16 insertions(+), 18 deletions(-) diff --git a/policyengine_core/reforms/structural_reform.py b/policyengine_core/reforms/structural_reform.py index b7a0153db..8dc569d6f 100644 --- a/policyengine_core/reforms/structural_reform.py +++ b/policyengine_core/reforms/structural_reform.py @@ -42,21 +42,9 @@ def __init__(self, trigger_parameter: str): """ self.trigger_parameter = trigger_parameter - def add_tax_benefit_system(self, tax_benefit_system: TaxBenefitSystem): - """ - Add a tax benefit system to the structural reform. - - Args: - tax_benefit_system: The tax benefit system to be added - """ - if not isinstance(tax_benefit_system, TaxBenefitSystem): - raise TypeError( - "Tax benefit system must be an instance of the TaxBenefitSystem class." - ) - self.tax_benefit_system = tax_benefit_system - def activate( self, + tax_benefit_system: TaxBenefitSystem, start_instant: Annotated[str, "YYYY-MM-DD"], end_instant: Annotated[str, "YYYY-MM-DD"] | None, ): @@ -64,14 +52,11 @@ def activate( Activate the structural reform. Args: + tax_benefit_system: The tax benefit system to which the structural reform will be applied start_instant: The start instant to be added; must be in the format 'YYYY-MM-DD' end_instant: The end instant to be added; must be in the format 'YYYY-MM-DD' or None """ - if not self.tax_benefit_system or self.tax_benefit_system is None: - raise ValueError( - "Tax benefit system must be added before start instant." - ) - + self._add_tax_benefit_system(tax_benefit_system) self._add_start_instant(start_instant) self._add_end_instant(end_instant) self._activate_transformation_log() @@ -367,4 +352,17 @@ def _add_end_instant( self._validate_instant(end_instant) self.end_instant = end_instant + def _add_tax_benefit_system(self, tax_benefit_system: TaxBenefitSystem): + """ + Add a tax benefit system to the structural reform. + + Args: + tax_benefit_system: The tax benefit system to be added + """ + if not isinstance(tax_benefit_system, TaxBenefitSystem): + raise TypeError( + "Tax benefit system must be an instance of the TaxBenefitSystem class." + ) + self.tax_benefit_system = tax_benefit_system + # Default outputs method of some sort? From fa99c4c529e4d3bdf124474c058d413a888725ae Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 2 Jan 2025 19:55:59 +0100 Subject: [PATCH 13/18] feat: Add method to parse parameter from overall param structure --- .../reforms/structural_reform.py | 91 ++++++++++++++----- 1 file changed, 69 insertions(+), 22 deletions(-) diff --git a/policyengine_core/reforms/structural_reform.py b/policyengine_core/reforms/structural_reform.py index 8dc569d6f..d3ee1d403 100644 --- a/policyengine_core/reforms/structural_reform.py +++ b/policyengine_core/reforms/structural_reform.py @@ -2,7 +2,7 @@ from datetime import datetime from dataclasses import dataclass from policyengine_core.variables import Variable -from policyengine_core.parameters import Parameter +from policyengine_core.parameters import Parameter, ParameterNode from policyengine_core.periods import config from policyengine_core.taxbenefitsystems import TaxBenefitSystem from policyengine_core.errors import ( @@ -22,7 +22,7 @@ class TransformationLogItem: transformation: Literal["neutralize", "add", "update"] -class StructuralReform: # Should this inherit from Reform and/or TaxBenefitSystem? +class StructuralReform: DEFAULT_START_INSTANT = "0000-01-01" transformation_log: list[TransformationLogItem] = [] @@ -30,9 +30,9 @@ class StructuralReform: # Should this inherit from Reform and/or TaxBenefitSyst tax_benefit_system: TaxBenefitSystem | None = None start_instant: Annotated[str, "YYYY-MM-DD"] = DEFAULT_START_INSTANT end_instant: Annotated[str, "YYYY-MM-DD"] | None = None - trigger_parameter: str = "" + trigger_parameter_path: str = "" - def __init__(self, trigger_parameter: str): + def __init__(self, trigger_parameter_path: str): """ Initialize a structural reform. @@ -40,7 +40,7 @@ def __init__(self, trigger_parameter: str): trigger_parameter: Path to the parameter that triggers the structural reform; this parameter must be Boolean """ - self.trigger_parameter = trigger_parameter + self.trigger_parameter_path = trigger_parameter_path def activate( self, @@ -56,15 +56,36 @@ def activate( start_instant: The start instant to be added; must be in the format 'YYYY-MM-DD' end_instant: The end instant to be added; must be in the format 'YYYY-MM-DD' or None """ - self._add_tax_benefit_system(tax_benefit_system) - self._add_start_instant(start_instant) - self._add_end_instant(end_instant) + if tax_benefit_system is None: + raise ValueError("Tax benefit system must be provided.") + + if not isinstance(tax_benefit_system, TaxBenefitSystem): + raise TypeError( + "Tax benefit system must be an instance of the TaxBenefitSystem class." + ) + + self.tax_benefit_system = tax_benefit_system + + # Fetch the trigger parameter + trigger_parameter: Parameter = self._fetch_parameter(self.trigger_parameter_path) + + # TODO: Parse date out of trigger parameter + start_instant: Annotated[str, "YYYY-MM-DD"] | None + end_instant: Annotated[str, "YYYY-MM-DD"] | None + start_instant, end_instant = self._parse_activation_period(trigger_parameter) + + + # Set + + # self._add_start_instant(start_instant) + # self._add_end_instant(end_instant) self._activate_transformation_log() def neutralize_variable(self, name: str): """ - Neutralize a variable by setting its formula to return the default value - from the StructuralReform's start_instant date to its end_instant date. + When structural reform is activated, neutralize a variable + by setting its formula to return the default value from the + StructuralReform's start_instant date to its end_instant date. Args: name: The name of the variable @@ -78,7 +99,8 @@ def neutralize_variable(self, name: str): def add_variable(self, variable: Variable): """ - Add a variable to the StructuralReform. + When structural reform is activated, add a variable + to the StructuralReform. Args: variable: The variable to be added @@ -93,8 +115,9 @@ def add_variable(self, variable: Variable): def update_variable(self, variable: Variable): """ - Update a variable in the tax benefit system; if the variable does not - yet exist, it will be added. + When structural reform is activated, update a + variable in the tax benefit system; if the variable + does not yet exist, it will be added. Args: variable: The variable to be updated @@ -239,6 +262,33 @@ def _fetch_variable(self, name: str) -> Variable | None: name: The name of the variable """ return self.tax_benefit_system.get_variable(name) + + def _fetch_parameter(self, parameter_path: str) -> Parameter: + """ + Given a dot-notated string, fetch a parameter by + reference from the tax benefit system. + + Args: + parameter_path: The dot-notated path to the parameter + + Raises: + AttributeError: If the parameter cannot be found + """ + root: ParameterNode | Parameter = self.tax_benefit_system.parameters + current: ParameterNode | Parameter = root + full_path: str = "" + + for index, key in enumerate(parameter_path.split(".")): + full_path += f"{key}" if index == 0 else f".{key}" + try: + current = getattr(current, key) + except AttributeError: + raise AttributeError(f"Unable to find parameter at path '{full_path}'") from None + + if not isinstance(current, Parameter): + raise AttributeError(f"Parameter at path '{full_path}' is not a Parameter, but a {type(current)}") + + return current # Method to modify metadata based on new items? @@ -352,17 +402,14 @@ def _add_end_instant( self._validate_instant(end_instant) self.end_instant = end_instant - def _add_tax_benefit_system(self, tax_benefit_system: TaxBenefitSystem): + def _parse_activation_period(self, trigger_parameter: Parameter) -> tuple[Annotated[str, "YYYY-MM-DD"], Annotated[str, "YYYY-MM-DD"] | None]: """ - Add a tax benefit system to the structural reform. + Given a trigger parameter, parse the reform start and end dates and return them. - Args: - tax_benefit_system: The tax benefit system to be added + Returns: + A tuple containing the start and end dates of the reform, + or None if the reform is not triggered """ - if not isinstance(tax_benefit_system, TaxBenefitSystem): - raise TypeError( - "Tax benefit system must be an instance of the TaxBenefitSystem class." - ) - self.tax_benefit_system = tax_benefit_system + return (self.start_instant, self.end_instant) # Default outputs method of some sort? From 3adcdf81f711ed60bb55ac24d61768b8d546e9c8 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 2 Jan 2025 22:25:36 +0100 Subject: [PATCH 14/18] feat: Parse start and end dates from parameter --- .../reforms/structural_reform.py | 78 ++++++++++--------- 1 file changed, 42 insertions(+), 36 deletions(-) diff --git a/policyengine_core/reforms/structural_reform.py b/policyengine_core/reforms/structural_reform.py index d3ee1d403..cc4b8f0ef 100644 --- a/policyengine_core/reforms/structural_reform.py +++ b/policyengine_core/reforms/structural_reform.py @@ -2,7 +2,7 @@ from datetime import datetime from dataclasses import dataclass from policyengine_core.variables import Variable -from policyengine_core.parameters import Parameter, ParameterNode +from policyengine_core.parameters import Parameter, ParameterNode, ParameterAtInstant from policyengine_core.periods import config from policyengine_core.taxbenefitsystems import TaxBenefitSystem from policyengine_core.errors import ( @@ -28,7 +28,7 @@ class StructuralReform: transformation_log: list[TransformationLogItem] = [] tax_benefit_system: TaxBenefitSystem | None = None - start_instant: Annotated[str, "YYYY-MM-DD"] = DEFAULT_START_INSTANT + start_instant: Annotated[str, "YYYY-MM-DD"] | None = None end_instant: Annotated[str, "YYYY-MM-DD"] | None = None trigger_parameter_path: str = "" @@ -45,16 +45,12 @@ def __init__(self, trigger_parameter_path: str): def activate( self, tax_benefit_system: TaxBenefitSystem, - start_instant: Annotated[str, "YYYY-MM-DD"], - end_instant: Annotated[str, "YYYY-MM-DD"] | None, ): """ Activate the structural reform. Args: tax_benefit_system: The tax benefit system to which the structural reform will be applied - start_instant: The start instant to be added; must be in the format 'YYYY-MM-DD' - end_instant: The end instant to be added; must be in the format 'YYYY-MM-DD' or None """ if tax_benefit_system is None: raise ValueError("Tax benefit system must be provided.") @@ -69,16 +65,17 @@ def activate( # Fetch the trigger parameter trigger_parameter: Parameter = self._fetch_parameter(self.trigger_parameter_path) - # TODO: Parse date out of trigger parameter + # Parse date out of trigger parameter and set start_instant: Annotated[str, "YYYY-MM-DD"] | None end_instant: Annotated[str, "YYYY-MM-DD"] | None start_instant, end_instant = self._parse_activation_period(trigger_parameter) + self.start_instant = start_instant + self.end_instant = end_instant - # Set + if self.start_instant is None: + return - # self._add_start_instant(start_instant) - # self._add_end_instant(end_instant) self._activate_transformation_log() def neutralize_variable(self, name: str): @@ -374,42 +371,51 @@ def _validate_instant(self, instant: Any) -> bool: return True - def _add_start_instant(self, start_instant: Annotated[str, "YYYY-MM-DD"]): + def _parse_activation_period(self, trigger_parameter: Parameter) -> tuple[Annotated[str, "YYYY-MM-DD"] | None, Annotated[str, "YYYY-MM-DD"] | None]: """ - Add a start instant to the structural reform. + Given a trigger parameter, parse the reform start and end dates and return them. - Args: - start_instant: The start instant to be added + Returns: + A tuple containing the start and end dates of the reform, + or None if the reform is not triggered """ + # Crash if trigger param isn't Boolean; this shouldn't be used as a trigger + if (trigger_parameter.metadata is None) or (trigger_parameter.metadata["unit"] != "bool"): + raise ValueError("Trigger parameter must be a Boolean.") + + # Build custom representation of trigger parameter instants and values + values_dict: dict[Annotated[str, "YYYY-MM-DD"], int | float] = self._generate_param_values_dict(trigger_parameter.values_list) + + if list(values_dict.values()).count(True) > 1: + raise ValueError("Trigger parameter must only be activated once.") + + if list(values_dict.values()).count(True) == 0: + return (None, None) + + # Now that True only occurs once, find it + start_instant_index: int = list(values_dict.values()).index(True) + start_instant: Annotated[str, "YYYY-MM-DD"] = list(values_dict.keys())[start_instant_index] self._validate_instant(start_instant) - self.start_instant = start_instant - def _add_end_instant( - self, end_instant: Annotated[str, "YYYY-MM-DD"] | None - ): - """ - Add an end instant to the structural reform. + # If it's the last item, the reform occurs into perpetuity, else + # the reform ends at the next instant + if start_instant_index == len(values_dict) - 1: + return (start_instant, None) - Args: - end_instant: The end instant to be added - """ - if not self.tax_benefit_system or self.tax_benefit_system is None: - raise ValueError( - "Tax benefit system must be added before end instant." - ) - if end_instant is not None: - self._validate_instant(end_instant) - self.end_instant = end_instant + end_instant: Annotated[str, "YYYY-MM-DD"] = list(values_dict.keys())[start_instant_index + 1] + self._validate_instant(end_instant) + return (start_instant, end_instant) - def _parse_activation_period(self, trigger_parameter: Parameter) -> tuple[Annotated[str, "YYYY-MM-DD"], Annotated[str, "YYYY-MM-DD"] | None]: + def _generate_param_values_dict(self, values_list: list[ParameterAtInstant]) -> dict[Annotated[str, "YYYY-MM-DD"], int | float]: """ - Given a trigger parameter, parse the reform start and end dates and return them. + Given a list of ParameterAtInstant objects, generate a dictionary of the form {instant: value}. - Returns: - A tuple containing the start and end dates of the reform, - or None if the reform is not triggered + Args: + values_list: The list of ParameterAtInstant objects """ - return (self.start_instant, self.end_instant) + unsorted_dict = {value.instant_str: value.value for value in values_list} + sorted_dict = dict(sorted(unsorted_dict.items(), key=lambda item: item[0])) + return sorted_dict # Default outputs method of some sort? From 4a479594b58a3b88afc9ca3a51db9f52ca8a1733 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 2 Jan 2025 22:33:10 +0100 Subject: [PATCH 15/18] chore: Format --- .../reforms/structural_reform.py | 85 +++++++++++++------ 1 file changed, 57 insertions(+), 28 deletions(-) diff --git a/policyengine_core/reforms/structural_reform.py b/policyengine_core/reforms/structural_reform.py index cc4b8f0ef..31741d964 100644 --- a/policyengine_core/reforms/structural_reform.py +++ b/policyengine_core/reforms/structural_reform.py @@ -2,7 +2,11 @@ from datetime import datetime from dataclasses import dataclass from policyengine_core.variables import Variable -from policyengine_core.parameters import Parameter, ParameterNode, ParameterAtInstant +from policyengine_core.parameters import ( + Parameter, + ParameterNode, + ParameterAtInstant, +) from policyengine_core.periods import config from policyengine_core.taxbenefitsystems import TaxBenefitSystem from policyengine_core.errors import ( @@ -22,7 +26,7 @@ class TransformationLogItem: transformation: Literal["neutralize", "add", "update"] -class StructuralReform: +class StructuralReform: DEFAULT_START_INSTANT = "0000-01-01" transformation_log: list[TransformationLogItem] = [] @@ -53,22 +57,26 @@ def activate( tax_benefit_system: The tax benefit system to which the structural reform will be applied """ if tax_benefit_system is None: - raise ValueError("Tax benefit system must be provided.") - + raise ValueError("Tax benefit system must be provided.") + if not isinstance(tax_benefit_system, TaxBenefitSystem): raise TypeError( "Tax benefit system must be an instance of the TaxBenefitSystem class." ) - + self.tax_benefit_system = tax_benefit_system # Fetch the trigger parameter - trigger_parameter: Parameter = self._fetch_parameter(self.trigger_parameter_path) + trigger_parameter: Parameter = self._fetch_parameter( + self.trigger_parameter_path + ) # Parse date out of trigger parameter and set start_instant: Annotated[str, "YYYY-MM-DD"] | None end_instant: Annotated[str, "YYYY-MM-DD"] | None - start_instant, end_instant = self._parse_activation_period(trigger_parameter) + start_instant, end_instant = self._parse_activation_period( + trigger_parameter + ) self.start_instant = start_instant self.end_instant = end_instant @@ -112,8 +120,8 @@ def add_variable(self, variable: Variable): def update_variable(self, variable: Variable): """ - When structural reform is activated, update a - variable in the tax benefit system; if the variable + When structural reform is activated, update a + variable in the tax benefit system; if the variable does not yet exist, it will be added. Args: @@ -259,10 +267,10 @@ def _fetch_variable(self, name: str) -> Variable | None: name: The name of the variable """ return self.tax_benefit_system.get_variable(name) - + def _fetch_parameter(self, parameter_path: str) -> Parameter: """ - Given a dot-notated string, fetch a parameter by + Given a dot-notated string, fetch a parameter by reference from the tax benefit system. Args: @@ -280,11 +288,15 @@ def _fetch_parameter(self, parameter_path: str) -> Parameter: try: current = getattr(current, key) except AttributeError: - raise AttributeError(f"Unable to find parameter at path '{full_path}'") from None - + raise AttributeError( + f"Unable to find parameter at path '{full_path}'" + ) from None + if not isinstance(current, Parameter): - raise AttributeError(f"Parameter at path '{full_path}' is not a Parameter, but a {type(current)}") - + raise AttributeError( + f"Parameter at path '{full_path}' is not a Parameter, but a {type(current)}" + ) + return current # Method to modify metadata based on new items? @@ -371,51 +383,68 @@ def _validate_instant(self, instant: Any) -> bool: return True - def _parse_activation_period(self, trigger_parameter: Parameter) -> tuple[Annotated[str, "YYYY-MM-DD"] | None, Annotated[str, "YYYY-MM-DD"] | None]: + def _parse_activation_period(self, trigger_parameter: Parameter) -> tuple[ + Annotated[str, "YYYY-MM-DD"] | None, + Annotated[str, "YYYY-MM-DD"] | None, + ]: """ Given a trigger parameter, parse the reform start and end dates and return them. Returns: - A tuple containing the start and end dates of the reform, + A tuple containing the start and end dates of the reform, or None if the reform is not triggered """ # Crash if trigger param isn't Boolean; this shouldn't be used as a trigger - if (trigger_parameter.metadata is None) or (trigger_parameter.metadata["unit"] != "bool"): + if (trigger_parameter.metadata is None) or ( + trigger_parameter.metadata["unit"] != "bool" + ): raise ValueError("Trigger parameter must be a Boolean.") - + # Build custom representation of trigger parameter instants and values - values_dict: dict[Annotated[str, "YYYY-MM-DD"], int | float] = self._generate_param_values_dict(trigger_parameter.values_list) + values_dict: dict[Annotated[str, "YYYY-MM-DD"], int | float] = ( + self._generate_param_values_dict(trigger_parameter.values_list) + ) if list(values_dict.values()).count(True) > 1: raise ValueError("Trigger parameter must only be activated once.") - + if list(values_dict.values()).count(True) == 0: return (None, None) - + # Now that True only occurs once, find it start_instant_index: int = list(values_dict.values()).index(True) - start_instant: Annotated[str, "YYYY-MM-DD"] = list(values_dict.keys())[start_instant_index] + start_instant: Annotated[str, "YYYY-MM-DD"] = list(values_dict.keys())[ + start_instant_index + ] self._validate_instant(start_instant) - # If it's the last item, the reform occurs into perpetuity, else + # If it's the last item, the reform occurs into perpetuity, else # the reform ends at the next instant if start_instant_index == len(values_dict) - 1: return (start_instant, None) - end_instant: Annotated[str, "YYYY-MM-DD"] = list(values_dict.keys())[start_instant_index + 1] + end_instant: Annotated[str, "YYYY-MM-DD"] = list(values_dict.keys())[ + start_instant_index + 1 + ] self._validate_instant(end_instant) return (start_instant, end_instant) - def _generate_param_values_dict(self, values_list: list[ParameterAtInstant]) -> dict[Annotated[str, "YYYY-MM-DD"], int | float]: + def _generate_param_values_dict( + self, values_list: list[ParameterAtInstant] + ) -> dict[Annotated[str, "YYYY-MM-DD"], int | float]: """ Given a list of ParameterAtInstant objects, generate a dictionary of the form {instant: value}. Args: values_list: The list of ParameterAtInstant objects """ - unsorted_dict = {value.instant_str: value.value for value in values_list} - sorted_dict = dict(sorted(unsorted_dict.items(), key=lambda item: item[0])) + unsorted_dict = { + value.instant_str: value.value for value in values_list + } + sorted_dict = dict( + sorted(unsorted_dict.items(), key=lambda item: item[0]) + ) return sorted_dict # Default outputs method of some sort? From 6fef05f2183c44159702f06b1320ec209d636a7e Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Wed, 22 Jan 2025 19:34:32 +0100 Subject: [PATCH 16/18] fix: Load new-style structural reform data into TBS --- .../reforms/structural_reform.py | 1 + .../taxbenefitsystems/tax_benefit_system.py | 95 +++++++++++++++++++ 2 files changed, 96 insertions(+) diff --git a/policyengine_core/reforms/structural_reform.py b/policyengine_core/reforms/structural_reform.py index 31741d964..c106a30b2 100644 --- a/policyengine_core/reforms/structural_reform.py +++ b/policyengine_core/reforms/structural_reform.py @@ -44,6 +44,7 @@ def __init__(self, trigger_parameter_path: str): trigger_parameter: Path to the parameter that triggers the structural reform; this parameter must be Boolean """ + self.reform_id = hash(f"{self.__class__.__module__}_{self.__class__.__qualname__}") self.trigger_parameter_path = trigger_parameter_path def activate( diff --git a/policyengine_core/taxbenefitsystems/tax_benefit_system.py b/policyengine_core/taxbenefitsystems/tax_benefit_system.py index 0e6897296..fa345ba0a 100644 --- a/policyengine_core/taxbenefitsystems/tax_benefit_system.py +++ b/policyengine_core/taxbenefitsystems/tax_benefit_system.py @@ -50,6 +50,8 @@ log = logging.getLogger(__name__) +if TYPE_CHECKING: + from policyengine_core.reforms import StructuralReform class TaxBenefitSystem: """ @@ -84,6 +86,8 @@ class TaxBenefitSystem: """Directory containing the Python files defining the variables of the tax and benefit system.""" parameters_dir: str = None """Directory containing the YAML parameter tree.""" + structural_reforms_dir: str = None + """Directory containing the Python files defining structural reforms.""" auto_carry_over_input_variables: bool = False """Whether to automatically carry over input variables when calculating a variable for a period different from the period of the input variables.""" basic_inputs: List[str] = None @@ -102,6 +106,7 @@ def __init__(self, entities: Sequence[Entity] = None, reform=None) -> None: self.parameters: Optional[ParameterNode] = None self._parameters_at_instant_cache = {} # weakref.WeakValueDictionary() self.variables: Dict[Any, Any] = {} + self.possible_structural_reforms: list[StructuralReform] = [] # Tax benefit systems are mutable, so entities (which need to know about our variables) can't be shared among them if entities is None or len(entities) == 0: raise Exception( @@ -138,6 +143,9 @@ def __init__(self, entities: Sequence[Entity] = None, reform=None) -> None: self.parameters = propagate_parameter_metadata(self.parameters) self.add_abolition_parameters() + if self.structural_reforms_dir is not None: + self.create_structural_reform_data(self.structural_reforms_dir) + self.add_modelled_policy_metadata() def apply_reform_set(self, reform): @@ -402,6 +410,93 @@ def add_variables(self, *variables: List[Type[Variable]]): for variable in variables: self.add_variable(variable) + def create_structural_reform_data(self, directory: str): + """ + Creates a data module representing all possible structural reforms + within the current tax-benefit system + + Args: + directory: str + """ + self._create_structural_reform_data_from_dir(directory), + + # TODO: Add loading of deprecated-style structural reforms + # self._create_structural_reform_data_from_dir_deprecated(directory) + + def _create_structural_reform_data_from_dir(self, directory: str): + """ + Recursively adds data about all StructuralReform instances to a + data module within the `TaxBenefitSystem`. + + Args: + directory: str + """ + py_files = glob.glob(os.path.join(directory, "*.py")) + for py_file in py_files: + self._create_structural_reform_data_from_file(py_file) + subdirectories = glob.glob(os.path.join(directory, "*/")) + for subdirectory in subdirectories: + self._create_structural_reform_data_from_dir(subdirectory) + + def _create_structural_reform_data_from_file(self, file_path: str): + + # To avoid circular import error due to typecheck + from policyengine_core.reforms import StructuralReform + + try: + file_name = os.path.splitext(os.path.basename(file_path))[0] + + # As Python remembers loaded modules by name, in order to prevent collisions, we need to make sure that: + # - Files with the same name, but located in different directories, have a different module names. Hence the file path hash in the module name. + # - The same file, loaded by different tax and benefit systems, has distinct module names. Hence the `id(self)` in the module name. + module_name = ( + f"{id(self)}_{hash(os.path.abspath(file_path))}_{file_name}" + ) + + try: + spec = importlib.util.spec_from_file_location( + module_name, file_path + ) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + except NameError as e: + logging.error( + str(e) + ) + raise + potential_reforms = [ + getattr(module, item) + for item in module.__dict__ + if not item.startswith("__") + ] + + for pot_reform in potential_reforms: + if isinstance(pot_reform, StructuralReform): + added_reform_ids = [reform.reform_id for reform in self.possible_structural_reforms] + print("Pot reform reform ID: ", pot_reform.reform_id) + if pot_reform.reform_id not in added_reform_ids: + print("Not in seen reform ids") + self.add_structural_reform(pot_reform) + + except Exception: + log.error( + 'Unable to load structural reform(s) from file "{}"'.format( + file_path + ) + ) + raise + + def add_structural_reform(self, reform: StructuralReform): + """ + Adds a structural reform to the tax and benefit system. + + Args: + reform: The structural reform to add. Must be an instance of StructuralReform. + """ + self.possible_structural_reforms.append(reform) + + def load_extension(self, extension: str) -> None: """ Loads an extension to the tax and benefit system. From e2863166885e8a46945323c60fdae0da0d0fd33f Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Fri, 24 Jan 2025 01:15:50 +0100 Subject: [PATCH 17/18] fix: Redo how structural reform log is loaded --- .../reforms/structural_reform.py | 4 +- .../taxbenefitsystems/tax_benefit_system.py | 93 +------------------ 2 files changed, 4 insertions(+), 93 deletions(-) diff --git a/policyengine_core/reforms/structural_reform.py b/policyengine_core/reforms/structural_reform.py index c106a30b2..6f62f0202 100644 --- a/policyengine_core/reforms/structural_reform.py +++ b/policyengine_core/reforms/structural_reform.py @@ -44,7 +44,9 @@ def __init__(self, trigger_parameter_path: str): trigger_parameter: Path to the parameter that triggers the structural reform; this parameter must be Boolean """ - self.reform_id = hash(f"{self.__class__.__module__}_{self.__class__.__qualname__}") + self.reform_id = hash( + f"{self.__class__.__module__}_{self.__class__.__qualname__}" + ) self.trigger_parameter_path = trigger_parameter_path def activate( diff --git a/policyengine_core/taxbenefitsystems/tax_benefit_system.py b/policyengine_core/taxbenefitsystems/tax_benefit_system.py index fa345ba0a..e0f9b5976 100644 --- a/policyengine_core/taxbenefitsystems/tax_benefit_system.py +++ b/policyengine_core/taxbenefitsystems/tax_benefit_system.py @@ -53,6 +53,7 @@ if TYPE_CHECKING: from policyengine_core.reforms import StructuralReform + class TaxBenefitSystem: """ Represents the legislation. @@ -86,8 +87,6 @@ class TaxBenefitSystem: """Directory containing the Python files defining the variables of the tax and benefit system.""" parameters_dir: str = None """Directory containing the YAML parameter tree.""" - structural_reforms_dir: str = None - """Directory containing the Python files defining structural reforms.""" auto_carry_over_input_variables: bool = False """Whether to automatically carry over input variables when calculating a variable for a period different from the period of the input variables.""" basic_inputs: List[str] = None @@ -143,9 +142,6 @@ def __init__(self, entities: Sequence[Entity] = None, reform=None) -> None: self.parameters = propagate_parameter_metadata(self.parameters) self.add_abolition_parameters() - if self.structural_reforms_dir is not None: - self.create_structural_reform_data(self.structural_reforms_dir) - self.add_modelled_policy_metadata() def apply_reform_set(self, reform): @@ -410,93 +406,6 @@ def add_variables(self, *variables: List[Type[Variable]]): for variable in variables: self.add_variable(variable) - def create_structural_reform_data(self, directory: str): - """ - Creates a data module representing all possible structural reforms - within the current tax-benefit system - - Args: - directory: str - """ - self._create_structural_reform_data_from_dir(directory), - - # TODO: Add loading of deprecated-style structural reforms - # self._create_structural_reform_data_from_dir_deprecated(directory) - - def _create_structural_reform_data_from_dir(self, directory: str): - """ - Recursively adds data about all StructuralReform instances to a - data module within the `TaxBenefitSystem`. - - Args: - directory: str - """ - py_files = glob.glob(os.path.join(directory, "*.py")) - for py_file in py_files: - self._create_structural_reform_data_from_file(py_file) - subdirectories = glob.glob(os.path.join(directory, "*/")) - for subdirectory in subdirectories: - self._create_structural_reform_data_from_dir(subdirectory) - - def _create_structural_reform_data_from_file(self, file_path: str): - - # To avoid circular import error due to typecheck - from policyengine_core.reforms import StructuralReform - - try: - file_name = os.path.splitext(os.path.basename(file_path))[0] - - # As Python remembers loaded modules by name, in order to prevent collisions, we need to make sure that: - # - Files with the same name, but located in different directories, have a different module names. Hence the file path hash in the module name. - # - The same file, loaded by different tax and benefit systems, has distinct module names. Hence the `id(self)` in the module name. - module_name = ( - f"{id(self)}_{hash(os.path.abspath(file_path))}_{file_name}" - ) - - try: - spec = importlib.util.spec_from_file_location( - module_name, file_path - ) - module = importlib.util.module_from_spec(spec) - sys.modules[module_name] = module - spec.loader.exec_module(module) - except NameError as e: - logging.error( - str(e) - ) - raise - potential_reforms = [ - getattr(module, item) - for item in module.__dict__ - if not item.startswith("__") - ] - - for pot_reform in potential_reforms: - if isinstance(pot_reform, StructuralReform): - added_reform_ids = [reform.reform_id for reform in self.possible_structural_reforms] - print("Pot reform reform ID: ", pot_reform.reform_id) - if pot_reform.reform_id not in added_reform_ids: - print("Not in seen reform ids") - self.add_structural_reform(pot_reform) - - except Exception: - log.error( - 'Unable to load structural reform(s) from file "{}"'.format( - file_path - ) - ) - raise - - def add_structural_reform(self, reform: StructuralReform): - """ - Adds a structural reform to the tax and benefit system. - - Args: - reform: The structural reform to add. Must be an instance of StructuralReform. - """ - self.possible_structural_reforms.append(reform) - - def load_extension(self, extension: str) -> None: """ Loads an extension to the tax and benefit system. From 9a2aa4fd8279a69c23c868c730821b7e4c8ff2c0 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Fri, 24 Jan 2025 01:19:21 +0100 Subject: [PATCH 18/18] fix: Modify structural reform loading --- policyengine_core/taxbenefitsystems/tax_benefit_system.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/policyengine_core/taxbenefitsystems/tax_benefit_system.py b/policyengine_core/taxbenefitsystems/tax_benefit_system.py index e0f9b5976..5be3faebd 100644 --- a/policyengine_core/taxbenefitsystems/tax_benefit_system.py +++ b/policyengine_core/taxbenefitsystems/tax_benefit_system.py @@ -93,6 +93,8 @@ class TaxBenefitSystem: """Short list of basic inputs to get medium accuracy.""" modelled_policies: str = None """A YAML filepath containing metadata describing the modelled policies.""" + possible_structural_reforms: List[StructuralReform] = None + """List of possible structural reforms that can be applied to the tax and benefit system.""" def __init__(self, entities: Sequence[Entity] = None, reform=None) -> None: if entities is None: @@ -105,7 +107,6 @@ def __init__(self, entities: Sequence[Entity] = None, reform=None) -> None: self.parameters: Optional[ParameterNode] = None self._parameters_at_instant_cache = {} # weakref.WeakValueDictionary() self.variables: Dict[Any, Any] = {} - self.possible_structural_reforms: list[StructuralReform] = [] # Tax benefit systems are mutable, so entities (which need to know about our variables) can't be shared among them if entities is None or len(entities) == 0: raise Exception(