Standardise save() and load() functions for simulations

changelog_entry.yaml

@@ -0,0 +1,4 @@
+- bump: patch
+  changes:
+    fixed:
+    - Standardised saving and loading of simulations.

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

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

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

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)

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.
 

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()