Skip to content

Commit ceaee0a

Browse files
nwoodruff-conwoodruff-co
committed
Final updates
1 parent 1ab55a4 commit ceaee0a

File tree

8 files changed

+174
-224
lines changed

8 files changed

+174
-224
lines changed

‎changelog_entry.yaml‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
- bump: patch
2+
changes:
3+
fixed:
4+
- Standardised saving and loading of simulations.

‎examples/employment_income_variation_uk.py‎

Lines changed: 0 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -162,42 +162,10 @@ def create_dataset_with_varied_employment_income(
162162

163163
def run_simulation(dataset: PolicyEngineUKDataset) -> Simulation:
164164
"""Run a single simulation for all employment income variations."""
165-
# Specify additional variables to calculate beyond defaults
166-
variables = {
167-
"household": [
168-
# Default variables
169-
"household_id",
170-
"household_weight",
171-
"household_net_income",
172-
"hbai_household_net_income",
173-
"household_benefits",
174-
"household_tax",
175-
],
176-
"person": [
177-
"person_id",
178-
"benunit_id",
179-
"household_id",
180-
"person_weight",
181-
"employment_income",
182-
"age",
183-
],
184-
"benunit": [
185-
"benunit_id",
186-
"benunit_weight",
187-
# Individual benefits (at benunit level)
188-
"universal_credit",
189-
"child_benefit",
190-
"working_tax_credit",
191-
"child_tax_credit",
192-
"pension_credit",
193-
"income_support",
194-
],
195-
}
196165

197166
simulation = Simulation(
198167
dataset=dataset,
199168
tax_benefit_model_version=uk_latest,
200-
variables=variables,
201169
)
202170
simulation.run()
203171
return simulation

‎examples/employment_income_variation_us.py‎

Lines changed: 0 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -171,56 +171,10 @@ def create_dataset_with_varied_employment_income(
171171

172172
def run_simulation(dataset: PolicyEngineUSDataset) -> Simulation:
173173
"""Run a single simulation for all employment income variations."""
174-
# Specify variables to calculate
175-
variables = {
176-
"household": [
177-
"household_id",
178-
"household_weight",
179-
"household_net_income",
180-
"household_benefits",
181-
"household_tax",
182-
"household_market_income",
183-
],
184-
"person": [
185-
"person_id",
186-
"household_id",
187-
"marital_unit_id",
188-
"family_id",
189-
"spm_unit_id",
190-
"tax_unit_id",
191-
"person_weight",
192-
"employment_income",
193-
"age",
194-
],
195-
"spm_unit": [
196-
"spm_unit_id",
197-
"spm_unit_weight",
198-
"snap",
199-
"tanf",
200-
"spm_unit_net_income",
201-
],
202-
"tax_unit": [
203-
"tax_unit_id",
204-
"tax_unit_weight",
205-
"income_tax",
206-
"employee_payroll_tax",
207-
"eitc",
208-
"ctc",
209-
],
210-
"marital_unit": [
211-
"marital_unit_id",
212-
"marital_unit_weight",
213-
],
214-
"family": [
215-
"family_id",
216-
"family_weight",
217-
],
218-
}
219174

220175
simulation = Simulation(
221176
dataset=dataset,
222177
tax_benefit_model_version=us_latest,
223-
variables=variables,
224178
)
225179
simulation.run()
226180
return simulation

‎src/policyengine/core/simulation.py‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -21,10 +21,13 @@ class Simulation(BaseModel):
2121
tax_benefit_model_version: TaxBenefitModelVersion = None
2222
output_dataset: Dataset | None = None
2323

24-
variables: dict[str, list[str]] | None = Field(
25-
default=None,
26-
description="Optional dictionary mapping entity names to lists of variable names to calculate. If None, uses model defaults.",
27-
)
28-
2924
def run(self):
3025
self.tax_benefit_model_version.run(self)
26+
27+
def save(self):
28+
"""Save the simulation's output dataset."""
29+
self.tax_benefit_model_version.save(self)
30+
31+
def load(self):
32+
"""Load the simulation's output dataset."""
33+
self.tax_benefit_model_version.load(self)

‎src/policyengine/core/tax_benefit_model_version.py‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,16 @@ def run(self, simulation: "Simulation") -> "Simulation":
2929
"The TaxBenefitModel class must define a method to execute simulations."
3030
)
3131

32+
def save(self, simulation: "Simulation"):
33+
raise NotImplementedError(
34+
"The TaxBenefitModel class must define a method to save simulations."
35+
)
36+
37+
def load(self, simulation: "Simulation"):
38+
raise NotImplementedError(
39+
"The TaxBenefitModel class must define a method to load simulations."
40+
)
41+
3242
def get_parameter(self, name: str) -> "Parameter":
3343
"""Get a parameter by name.
3444

‎src/policyengine/tax_benefit_models/uk/model.py‎

Lines changed: 83 additions & 71 deletions
Original file line numberDiff line numberDiff line change
@@ -153,77 +153,72 @@ def run(self, simulation: "Simulation") -> "Simulation":
153153
)
154154
modifier(microsim)
155155

156-
# Allow custom variable selection, or use defaults
157-
if simulation.variables is not None:
158-
entity_variables = simulation.variables
159-
else:
160-
# Default comprehensive variable set
161-
entity_variables = {
162-
"person": [
163-
# IDs and weights
164-
"person_id",
165-
"benunit_id",
166-
"household_id",
167-
"person_weight",
168-
# Demographics
169-
"age",
170-
"gender",
171-
"is_adult",
172-
"is_SP_age",
173-
"is_child",
174-
# Income
175-
"employment_income",
176-
"self_employment_income",
177-
"pension_income",
178-
"private_pension_income",
179-
"savings_interest_income",
180-
"dividend_income",
181-
"property_income",
182-
"total_income",
183-
"earned_income",
184-
# Benefits
185-
"universal_credit",
186-
"child_benefit",
187-
"pension_credit",
188-
"income_support",
189-
"working_tax_credit",
190-
"child_tax_credit",
191-
# Tax
192-
"income_tax",
193-
"national_insurance",
194-
],
195-
"benunit": [
196-
# IDs and weights
197-
"benunit_id",
198-
"benunit_weight",
199-
# Structure
200-
"family_type",
201-
# Income and benefits
202-
"universal_credit",
203-
"child_benefit",
204-
"working_tax_credit",
205-
"child_tax_credit",
206-
],
207-
"household": [
208-
# IDs and weights
209-
"household_id",
210-
"household_weight",
211-
# Income measures
212-
"household_net_income",
213-
"hbai_household_net_income",
214-
"equiv_hbai_household_net_income",
215-
"household_market_income",
216-
"household_gross_income",
217-
# Benefits and tax
218-
"household_benefits",
219-
"household_tax",
220-
"vat",
221-
# Housing
222-
"rent",
223-
"council_tax",
224-
"tenure_type",
225-
],
226-
}
156+
entity_variables = {
157+
"person": [
158+
# IDs and weights
159+
"person_id",
160+
"benunit_id",
161+
"household_id",
162+
"person_weight",
163+
# Demographics
164+
"age",
165+
"gender",
166+
"is_adult",
167+
"is_SP_age",
168+
"is_child",
169+
# Income
170+
"employment_income",
171+
"self_employment_income",
172+
"pension_income",
173+
"private_pension_income",
174+
"savings_interest_income",
175+
"dividend_income",
176+
"property_income",
177+
"total_income",
178+
"earned_income",
179+
# Benefits
180+
"universal_credit",
181+
"child_benefit",
182+
"pension_credit",
183+
"income_support",
184+
"working_tax_credit",
185+
"child_tax_credit",
186+
# Tax
187+
"income_tax",
188+
"national_insurance",
189+
],
190+
"benunit": [
191+
# IDs and weights
192+
"benunit_id",
193+
"benunit_weight",
194+
# Structure
195+
"family_type",
196+
# Income and benefits
197+
"universal_credit",
198+
"child_benefit",
199+
"working_tax_credit",
200+
"child_tax_credit",
201+
],
202+
"household": [
203+
# IDs and weights
204+
"household_id",
205+
"household_weight",
206+
# Income measures
207+
"household_net_income",
208+
"hbai_household_net_income",
209+
"equiv_hbai_household_net_income",
210+
"household_market_income",
211+
"household_gross_income",
212+
# Benefits and tax
213+
"household_benefits",
214+
"household_tax",
215+
"vat",
216+
# Housing
217+
"rent",
218+
"council_tax",
219+
"tenure_type",
220+
],
221+
}
227222

228223
data = {
229224
"person": pd.DataFrame(),
@@ -248,6 +243,7 @@ def run(self, simulation: "Simulation") -> "Simulation":
248243
)
249244

250245
simulation.output_dataset = PolicyEngineUKDataset(
246+
id=simulation.id,
251247
name=dataset.name,
252248
description=dataset.description,
253249
filepath=str(
@@ -263,7 +259,23 @@ def run(self, simulation: "Simulation") -> "Simulation":
263259
),
264260
)
265261

262+
def save(self, simulation: "Simulation"):
263+
"""Save the simulation's output dataset."""
266264
simulation.output_dataset.save()
267265

266+
def load(self, simulation: "Simulation"):
267+
"""Load the simulation's output dataset."""
268+
simulation.output_dataset = PolicyEngineUKDataset(
269+
id=simulation.id,
270+
name=simulation.dataset.name,
271+
description=simulation.dataset.description,
272+
filepath=str(
273+
Path(simulation.dataset.filepath).parent
274+
/ (simulation.id + ".h5")
275+
),
276+
year=simulation.dataset.year,
277+
is_output_dataset=True,
278+
)
279+
268280

269281
uk_latest = PolicyEngineUKLatest()

0 commit comments

Comments
 (0)