Skip to content

Commit d067f2a

Browse files
nwoodruff-conwoodruff-co
authored
Move entity variables out to constant (#187)
* Move entity variables out to constant * Versioning * Annotate
1 parent ccac3e3 commit d067f2a

File tree

3 files changed

+126
-122
lines changed

3 files changed

+126
-122
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+
- Entity variables moved out to an editable constant.

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

Lines changed: 68 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,73 @@ class PolicyEngineUKLatest(TaxBenefitModelVersion):
4444
upload_time
4545
)
4646

47+
entity_variables: dict[str, list[str]] = {
48+
"person": [
49+
# IDs and weights
50+
"person_id",
51+
"benunit_id",
52+
"household_id",
53+
"person_weight",
54+
# Demographics
55+
"age",
56+
"gender",
57+
"is_adult",
58+
"is_SP_age",
59+
"is_child",
60+
# Income
61+
"employment_income",
62+
"self_employment_income",
63+
"pension_income",
64+
"private_pension_income",
65+
"savings_interest_income",
66+
"dividend_income",
67+
"property_income",
68+
"total_income",
69+
"earned_income",
70+
# Benefits
71+
"universal_credit",
72+
"child_benefit",
73+
"pension_credit",
74+
"income_support",
75+
"working_tax_credit",
76+
"child_tax_credit",
77+
# Tax
78+
"income_tax",
79+
"national_insurance",
80+
],
81+
"benunit": [
82+
# IDs and weights
83+
"benunit_id",
84+
"benunit_weight",
85+
# Structure
86+
"family_type",
87+
# Income and benefits
88+
"universal_credit",
89+
"child_benefit",
90+
"working_tax_credit",
91+
"child_tax_credit",
92+
],
93+
"household": [
94+
# IDs and weights
95+
"household_id",
96+
"household_weight",
97+
# Income measures
98+
"household_net_income",
99+
"hbai_household_net_income",
100+
"equiv_hbai_household_net_income",
101+
"household_market_income",
102+
"household_gross_income",
103+
# Benefits and tax
104+
"household_benefits",
105+
"household_tax",
106+
"vat",
107+
# Housing
108+
"rent",
109+
"council_tax",
110+
"tenure_type",
111+
],
112+
}
113+
47114
def __init__(self, **kwargs: dict):
48115
super().__init__(**kwargs)
49116
from policyengine_core.enums import Enum
@@ -153,80 +220,13 @@ def run(self, simulation: "Simulation") -> "Simulation":
153220
)
154221
modifier(microsim)
155222

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-
}
222-
223223
data = {
224224
"person": pd.DataFrame(),
225225
"benunit": pd.DataFrame(),
226226
"household": pd.DataFrame(),
227227
}
228228

229-
for entity, variables in entity_variables.items():
229+
for entity, variables in self.entity_variables.items():
230230
for var in variables:
231231
data[entity][var] = microsim.calculate(
232232
var, period=simulation.dataset.year, map_to=entity

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

Lines changed: 54 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,59 @@ class PolicyEngineUSLatest(TaxBenefitModelVersion):
4545
version: str = None
4646
created_at: datetime.datetime = None
4747

48+
entity_variables: dict[str, list[str]] = {
49+
"person": [
50+
# IDs and weights
51+
"person_id",
52+
"marital_unit_id",
53+
"family_id",
54+
"spm_unit_id",
55+
"tax_unit_id",
56+
"household_id",
57+
"person_weight",
58+
# Demographics
59+
"age",
60+
# Income
61+
"employment_income",
62+
# Benefits
63+
"ssi",
64+
"social_security",
65+
"medicaid",
66+
"unemployment_compensation",
67+
],
68+
"marital_unit": [
69+
"marital_unit_id",
70+
"marital_unit_weight",
71+
],
72+
"family": [
73+
"family_id",
74+
"family_weight",
75+
],
76+
"spm_unit": [
77+
"spm_unit_id",
78+
"spm_unit_weight",
79+
"snap",
80+
"tanf",
81+
"spm_unit_net_income",
82+
],
83+
"tax_unit": [
84+
"tax_unit_id",
85+
"tax_unit_weight",
86+
"income_tax",
87+
"employee_payroll_tax",
88+
"eitc",
89+
"ctc",
90+
],
91+
"household": [
92+
"household_id",
93+
"household_weight",
94+
"household_net_income",
95+
"household_benefits",
96+
"household_tax",
97+
"household_market_income",
98+
],
99+
}
100+
48101
def __init__(self, **kwargs: dict):
49102
# Lazy-load package metadata if not provided
50103
if "version" not in kwargs or kwargs.get("version") is None:
@@ -156,59 +209,6 @@ def run(self, simulation: "Simulation") -> "Simulation":
156209
)
157210
modifier(microsim)
158211

159-
entity_variables = {
160-
"person": [
161-
# IDs and weights
162-
"person_id",
163-
"marital_unit_id",
164-
"family_id",
165-
"spm_unit_id",
166-
"tax_unit_id",
167-
"household_id",
168-
"person_weight",
169-
# Demographics
170-
"age",
171-
# Income
172-
"employment_income",
173-
# Benefits
174-
"ssi",
175-
"social_security",
176-
"medicaid",
177-
"unemployment_compensation",
178-
],
179-
"marital_unit": [
180-
"marital_unit_id",
181-
"marital_unit_weight",
182-
],
183-
"family": [
184-
"family_id",
185-
"family_weight",
186-
],
187-
"spm_unit": [
188-
"spm_unit_id",
189-
"spm_unit_weight",
190-
"snap",
191-
"tanf",
192-
"spm_unit_net_income",
193-
],
194-
"tax_unit": [
195-
"tax_unit_id",
196-
"tax_unit_weight",
197-
"income_tax",
198-
"employee_payroll_tax",
199-
"eitc",
200-
"ctc",
201-
],
202-
"household": [
203-
"household_id",
204-
"household_weight",
205-
"household_net_income",
206-
"household_benefits",
207-
"household_tax",
208-
"household_market_income",
209-
],
210-
}
211-
212212
data = {
213213
"person": pd.DataFrame(),
214214
"marital_unit": pd.DataFrame(),
@@ -259,7 +259,7 @@ def run(self, simulation: "Simulation") -> "Simulation":
259259
data["person"][target_col] = person_input_df[col].values
260260

261261
# Then calculate non-ID, non-weight variables from simulation
262-
for entity, variables in entity_variables.items():
262+
for entity, variables in self.entity_variables.items():
263263
for var in variables:
264264
if var not in id_columns and var not in weight_columns:
265265
data[entity][var] = microsim.calculate(

0 commit comments

Comments
 (0)