Skip to content

Move entity variables out to constant #187

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Nov 18, 2025
Merged

Move entity variables out to constant #187

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
26a968d
Move entity variables out to constant
nwoodruff-co Nov 18, 2025
22209de
Merge branch 'main' of https://github.com/PolicyEngine/policyengine.p…
nwoodruff-co Nov 18, 2025
152891d
Versioning
nwoodruff-co Nov 18, 2025
e41b243
Annotate
nwoodruff-co Nov 18, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions changelog_entry.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
- bump: patch
changes:
fixed:
- Entity variables moved out to an editable constant.
136 changes: 68 additions & 68 deletions src/policyengine/tax_benefit_models/uk/model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,73 @@ class PolicyEngineUKLatest(TaxBenefitModelVersion):
upload_time
)

entity_variables: dict[str, list[str]] = {
"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",
],
}

def __init__(self, **kwargs: dict):
super().__init__(**kwargs)
from policyengine_core.enums import Enum
Expand Down Expand Up @@ -153,80 +220,13 @@ def run(self, simulation: "Simulation") -> "Simulation":
)
modifier(microsim)

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(),
"benunit": pd.DataFrame(),
"household": pd.DataFrame(),
}

for entity, variables in entity_variables.items():
for entity, variables in self.entity_variables.items():
for var in variables:
data[entity][var] = microsim.calculate(
var, period=simulation.dataset.year, map_to=entity
Expand Down
108 changes: 54 additions & 54 deletions src/policyengine/tax_benefit_models/us/model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,59 @@ class PolicyEngineUSLatest(TaxBenefitModelVersion):
version: str = None
created_at: datetime.datetime = None

entity_variables: dict[str, list[str]] = {
"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",
],
}

def __init__(self, **kwargs: dict):
# Lazy-load package metadata if not provided
if "version" not in kwargs or kwargs.get("version") is None:
Expand Down Expand Up @@ -156,59 +209,6 @@ def run(self, simulation: "Simulation") -> "Simulation":
)
modifier(microsim)

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(),
"marital_unit": pd.DataFrame(),
Expand Down Expand Up @@ -259,7 +259,7 @@ def run(self, simulation: "Simulation") -> "Simulation":
data["person"][target_col] = person_input_df[col].values

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