diff --git a/changelog_entry.yaml b/changelog_entry.yaml index e69de29b..1311929d 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -0,0 +1,4 @@ +- bump: patch + changes: + fixed: + - Standardised saving and loading of simulations. diff --git a/examples/employment_income_variation_uk.py b/examples/employment_income_variation_uk.py index 173c78ff..22bcb93c 100644 --- a/examples/employment_income_variation_uk.py +++ b/examples/employment_income_variation_uk.py @@ -162,42 +162,10 @@ def create_dataset_with_varied_employment_income( def run_simulation(dataset: PolicyEngineUKDataset) -> Simulation: """Run a single simulation for all employment income variations.""" - # Specify additional variables to calculate beyond defaults - variables = { - "household": [ - # Default variables - "household_id", - "household_weight", - "household_net_income", - "hbai_household_net_income", - "household_benefits", - "household_tax", - ], - "person": [ - "person_id", - "benunit_id", - "household_id", - "person_weight", - "employment_income", - "age", - ], - "benunit": [ - "benunit_id", - "benunit_weight", - # Individual benefits (at benunit level) - "universal_credit", - "child_benefit", - "working_tax_credit", - "child_tax_credit", - "pension_credit", - "income_support", - ], - } simulation = Simulation( dataset=dataset, tax_benefit_model_version=uk_latest, - variables=variables, ) simulation.run() return simulation diff --git a/examples/employment_income_variation_us.py b/examples/employment_income_variation_us.py index 863d8018..f4ceb80e 100644 --- a/examples/employment_income_variation_us.py +++ b/examples/employment_income_variation_us.py @@ -171,56 +171,10 @@ def create_dataset_with_varied_employment_income( def run_simulation(dataset: PolicyEngineUSDataset) -> Simulation: """Run a single simulation for all employment income variations.""" - # Specify variables to calculate - variables = { - "household": [ - "household_id", - "household_weight", - "household_net_income", - "household_benefits", - "household_tax", - "household_market_income", - ], - "person": [ - "person_id", - "household_id", - "marital_unit_id", - "family_id", - "spm_unit_id", - "tax_unit_id", - "person_weight", - "employment_income", - "age", - ], - "spm_unit": [ - "spm_unit_id", - "spm_unit_weight", - "snap", - "tanf", - "spm_unit_net_income", - ], - "tax_unit": [ - "tax_unit_id", - "tax_unit_weight", - "income_tax", - "employee_payroll_tax", - "eitc", - "ctc", - ], - "marital_unit": [ - "marital_unit_id", - "marital_unit_weight", - ], - "family": [ - "family_id", - "family_weight", - ], - } simulation = Simulation( dataset=dataset, tax_benefit_model_version=us_latest, - variables=variables, ) simulation.run() return simulation diff --git a/src/policyengine/core/parameter.py b/src/policyengine/core/parameter.py index 54e3e116..79c4f175 100644 --- a/src/policyengine/core/parameter.py +++ b/src/policyengine/core/parameter.py @@ -8,6 +8,7 @@ class Parameter(BaseModel): id: str = Field(default_factory=lambda: str(uuid4())) name: str + label: str | None = None description: str | None = None data_type: type | None = None tax_benefit_model_version: TaxBenefitModelVersion diff --git a/src/policyengine/core/simulation.py b/src/policyengine/core/simulation.py index 1e493b9a..f7c214e4 100644 --- a/src/policyengine/core/simulation.py +++ b/src/policyengine/core/simulation.py @@ -21,10 +21,13 @@ class Simulation(BaseModel): tax_benefit_model_version: TaxBenefitModelVersion = None output_dataset: Dataset | None = None - variables: dict[str, list[str]] | None = Field( - default=None, - description="Optional dictionary mapping entity names to lists of variable names to calculate. If None, uses model defaults.", - ) - def run(self): self.tax_benefit_model_version.run(self) + + def save(self): + """Save the simulation's output dataset.""" + self.tax_benefit_model_version.save(self) + + def load(self): + """Load the simulation's output dataset.""" + self.tax_benefit_model_version.load(self) diff --git a/src/policyengine/core/tax_benefit_model_version.py b/src/policyengine/core/tax_benefit_model_version.py index b03c73eb..53b9936e 100644 --- a/src/policyengine/core/tax_benefit_model_version.py +++ b/src/policyengine/core/tax_benefit_model_version.py @@ -29,6 +29,16 @@ def run(self, simulation: "Simulation") -> "Simulation": "The TaxBenefitModel class must define a method to execute simulations." ) + def save(self, simulation: "Simulation"): + raise NotImplementedError( + "The TaxBenefitModel class must define a method to save simulations." + ) + + def load(self, simulation: "Simulation"): + raise NotImplementedError( + "The TaxBenefitModel class must define a method to load simulations." + ) + def get_parameter(self, name: str) -> "Parameter": """Get a parameter by name. diff --git a/src/policyengine/tax_benefit_models/uk/model.py b/src/policyengine/tax_benefit_models/uk/model.py index 18f1ef25..6b8c5c76 100644 --- a/src/policyengine/tax_benefit_models/uk/model.py +++ b/src/policyengine/tax_benefit_models/uk/model.py @@ -83,6 +83,7 @@ def __init__(self, **kwargs: dict): parameter = Parameter( id=self.id + "-" + param_node.name, name=param_node.name, + label=param_node.metadata.get("label", param_node.name), tax_benefit_model_version=self, description=param_node.description, data_type=type( @@ -152,77 +153,72 @@ def run(self, simulation: "Simulation") -> "Simulation": ) modifier(microsim) - # Allow custom variable selection, or use defaults - if simulation.variables is not None: - entity_variables = simulation.variables - else: - # Default comprehensive variable set - entity_variables = { - "person": [ - # IDs and weights - "person_id", - "benunit_id", - "household_id", - "person_weight", - # Demographics - "age", - "gender", - "is_adult", - "is_SP_age", - "is_child", - # Income - "employment_income", - "self_employment_income", - "pension_income", - "private_pension_income", - "savings_interest_income", - "dividend_income", - "property_income", - "total_income", - "earned_income", - # Benefits - "universal_credit", - "child_benefit", - "pension_credit", - "income_support", - "working_tax_credit", - "child_tax_credit", - # Tax - "income_tax", - "national_insurance", - ], - "benunit": [ - # IDs and weights - "benunit_id", - "benunit_weight", - # Structure - "family_type", - # Income and benefits - "universal_credit", - "child_benefit", - "working_tax_credit", - "child_tax_credit", - ], - "household": [ - # IDs and weights - "household_id", - "household_weight", - # Income measures - "household_net_income", - "hbai_household_net_income", - "equiv_hbai_household_net_income", - "household_market_income", - "household_gross_income", - # Benefits and tax - "household_benefits", - "household_tax", - "vat", - # Housing - "rent", - "council_tax", - "tenure_type", - ], - } + entity_variables = { + "person": [ + # IDs and weights + "person_id", + "benunit_id", + "household_id", + "person_weight", + # Demographics + "age", + "gender", + "is_adult", + "is_SP_age", + "is_child", + # Income + "employment_income", + "self_employment_income", + "pension_income", + "private_pension_income", + "savings_interest_income", + "dividend_income", + "property_income", + "total_income", + "earned_income", + # Benefits + "universal_credit", + "child_benefit", + "pension_credit", + "income_support", + "working_tax_credit", + "child_tax_credit", + # Tax + "income_tax", + "national_insurance", + ], + "benunit": [ + # IDs and weights + "benunit_id", + "benunit_weight", + # Structure + "family_type", + # Income and benefits + "universal_credit", + "child_benefit", + "working_tax_credit", + "child_tax_credit", + ], + "household": [ + # IDs and weights + "household_id", + "household_weight", + # Income measures + "household_net_income", + "hbai_household_net_income", + "equiv_hbai_household_net_income", + "household_market_income", + "household_gross_income", + # Benefits and tax + "household_benefits", + "household_tax", + "vat", + # Housing + "rent", + "council_tax", + "tenure_type", + ], + } data = { "person": pd.DataFrame(), @@ -247,6 +243,7 @@ def run(self, simulation: "Simulation") -> "Simulation": ) simulation.output_dataset = PolicyEngineUKDataset( + id=simulation.id, name=dataset.name, description=dataset.description, filepath=str( @@ -262,7 +259,23 @@ def run(self, simulation: "Simulation") -> "Simulation": ), ) + def save(self, simulation: "Simulation"): + """Save the simulation's output dataset.""" simulation.output_dataset.save() + def load(self, simulation: "Simulation"): + """Load the simulation's output dataset.""" + simulation.output_dataset = PolicyEngineUKDataset( + id=simulation.id, + name=simulation.dataset.name, + description=simulation.dataset.description, + filepath=str( + Path(simulation.dataset.filepath).parent + / (simulation.id + ".h5") + ), + year=simulation.dataset.year, + is_output_dataset=True, + ) + uk_latest = PolicyEngineUKLatest() diff --git a/src/policyengine/tax_benefit_models/us/model.py b/src/policyengine/tax_benefit_models/us/model.py index 5e2068c5..a5a267a2 100644 --- a/src/policyengine/tax_benefit_models/us/model.py +++ b/src/policyengine/tax_benefit_models/us/model.py @@ -156,63 +156,58 @@ def run(self, simulation: "Simulation") -> "Simulation": ) modifier(microsim) - # Allow custom variable selection, or use defaults - if simulation.variables is not None: - entity_variables = simulation.variables - else: - # Default comprehensive variable set - entity_variables = { - "person": [ - # IDs and weights - "person_id", - "marital_unit_id", - "family_id", - "spm_unit_id", - "tax_unit_id", - "household_id", - "person_weight", - # Demographics - "age", - # Income - "employment_income", - # Benefits - "ssi", - "social_security", - "medicaid", - "unemployment_compensation", - ], - "marital_unit": [ - "marital_unit_id", - "marital_unit_weight", - ], - "family": [ - "family_id", - "family_weight", - ], - "spm_unit": [ - "spm_unit_id", - "spm_unit_weight", - "snap", - "tanf", - "spm_unit_net_income", - ], - "tax_unit": [ - "tax_unit_id", - "tax_unit_weight", - "income_tax", - "employee_payroll_tax", - "eitc", - "ctc", - ], - "household": [ - "household_id", - "household_weight", - "household_net_income", - "household_benefits", - "household_tax", - "household_market_income", - ], - } + entity_variables = { + "person": [ + # IDs and weights + "person_id", + "marital_unit_id", + "family_id", + "spm_unit_id", + "tax_unit_id", + "household_id", + "person_weight", + # Demographics + "age", + # Income + "employment_income", + # Benefits + "ssi", + "social_security", + "medicaid", + "unemployment_compensation", + ], + "marital_unit": [ + "marital_unit_id", + "marital_unit_weight", + ], + "family": [ + "family_id", + "family_weight", + ], + "spm_unit": [ + "spm_unit_id", + "spm_unit_weight", + "snap", + "tanf", + "spm_unit_net_income", + ], + "tax_unit": [ + "tax_unit_id", + "tax_unit_weight", + "income_tax", + "employee_payroll_tax", + "eitc", + "ctc", + ], + "household": [ + "household_id", + "household_weight", + "household_net_income", + "household_benefits", + "household_tax", + "household_market_income", + ], + } data = { "person": pd.DataFrame(), @@ -291,6 +286,7 @@ def run(self, simulation: "Simulation") -> "Simulation": ) simulation.output_dataset = PolicyEngineUSDataset( + id=simulation.id, name=dataset.name, description=dataset.description, filepath=str( @@ -309,8 +305,24 @@ def run(self, simulation: "Simulation") -> "Simulation": ), ) + def save(self, simulation: "Simulation"): + """Save the simulation's output dataset.""" simulation.output_dataset.save() + def load(self, simulation: "Simulation"): + """Load the simulation's output dataset.""" + simulation.output_dataset = PolicyEngineUSDataset( + id=simulation.id, + name=simulation.dataset.name, + description=simulation.dataset.description, + filepath=str( + Path(simulation.dataset.filepath).parent + / (simulation.id + ".h5") + ), + year=simulation.dataset.year, + is_output_dataset=True, + ) + def _build_simulation_from_dataset(self, microsim, dataset, system): """Build a PolicyEngine Core simulation from dataset entity IDs. diff --git a/tests/test_us_simulation.py b/tests/test_us_simulation.py index 4de79691..aad5f9bb 100644 --- a/tests/test_us_simulation.py +++ b/tests/test_us_simulation.py @@ -227,19 +227,6 @@ def test_us_simulation_from_dataset(): simulation = Simulation( dataset=dataset, tax_benefit_model_version=us_latest, - variables={ - "person": [ - "person_id", - "person_weight", - "age", - "employment_income", - ], - "household": ["household_id", "household_weight"], - "marital_unit": ["marital_unit_id", "marital_unit_weight"], - "family": ["family_id", "family_weight"], - "spm_unit": ["spm_unit_id", "spm_unit_weight"], - "tax_unit": ["tax_unit_id", "tax_unit_weight"], - }, ) simulation.run()