Initial commit
This commit is contained in:
Zhongwei Li
@@ -0,0 +1,257 @@
|
||||
"""
|
||||
Helper functions for creating PolicyEngine-US situations.
|
||||
|
||||
These utilities simplify the creation of situation dictionaries
|
||||
for common household configurations.
|
||||
"""
|
||||
|
||||
CURRENT_YEAR = 2024
|
||||
|
||||
|
||||
def create_single_filer(income, state="CA", age=35, **kwargs):
|
||||
"""
|
||||
Create a situation for a single tax filer.
|
||||
|
||||
Args:
|
||||
income (float): Employment income
|
||||
state (str): Two-letter state code (e.g., "CA", "NY")
|
||||
age (int): Person's age
|
||||
**kwargs: Additional person attributes (e.g., self_employment_income)
|
||||
|
||||
Returns:
|
||||
dict: PolicyEngine situation dictionary
|
||||
"""
|
||||
person_attrs = {
|
||||
"age": {CURRENT_YEAR: age},
|
||||
"employment_income": {CURRENT_YEAR: income},
|
||||
}
|
||||
person_attrs.update({k: {CURRENT_YEAR: v} for k, v in kwargs.items()})
|
||||
|
||||
return {
|
||||
"people": {"person": person_attrs},
|
||||
"families": {"family": {"members": ["person"]}},
|
||||
"marital_units": {"marital_unit": {"members": ["person"]}},
|
||||
"tax_units": {"tax_unit": {"members": ["person"]}},
|
||||
"spm_units": {"spm_unit": {"members": ["person"]}},
|
||||
"households": {
|
||||
"household": {
|
||||
"members": ["person"],
|
||||
"state_name": {CURRENT_YEAR: state}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def create_married_couple(
|
||||
income_1, income_2=0, state="CA", age_1=35, age_2=35, **kwargs
|
||||
):
|
||||
"""
|
||||
Create a situation for a married couple filing jointly.
|
||||
|
||||
Args:
|
||||
income_1 (float): First spouse's employment income
|
||||
income_2 (float): Second spouse's employment income
|
||||
state (str): Two-letter state code
|
||||
age_1 (int): First spouse's age
|
||||
age_2 (int): Second spouse's age
|
||||
**kwargs: Additional household attributes
|
||||
|
||||
Returns:
|
||||
dict: PolicyEngine situation dictionary
|
||||
"""
|
||||
members = ["spouse_1", "spouse_2"]
|
||||
|
||||
household_attrs = {
|
||||
"members": members,
|
||||
"state_name": {CURRENT_YEAR: state}
|
||||
}
|
||||
household_attrs.update({k: {CURRENT_YEAR: v} for k, v in kwargs.items()})
|
||||
|
||||
return {
|
||||
"people": {
|
||||
"spouse_1": {
|
||||
"age": {CURRENT_YEAR: age_1},
|
||||
"employment_income": {CURRENT_YEAR: income_1}
|
||||
},
|
||||
"spouse_2": {
|
||||
"age": {CURRENT_YEAR: age_2},
|
||||
"employment_income": {CURRENT_YEAR: income_2}
|
||||
}
|
||||
},
|
||||
"families": {"family": {"members": members}},
|
||||
"marital_units": {"marital_unit": {"members": members}},
|
||||
"tax_units": {"tax_unit": {"members": members}},
|
||||
"spm_units": {"spm_unit": {"members": members}},
|
||||
"households": {"household": household_attrs}
|
||||
}
|
||||
|
||||
|
||||
def create_family_with_children(
|
||||
parent_income,
|
||||
num_children=1,
|
||||
child_ages=None,
|
||||
state="CA",
|
||||
parent_age=35,
|
||||
married=False,
|
||||
spouse_income=0,
|
||||
**kwargs
|
||||
):
|
||||
"""
|
||||
Create a situation for a family with children.
|
||||
|
||||
Args:
|
||||
parent_income (float): Primary parent's employment income
|
||||
num_children (int): Number of children
|
||||
child_ages (list): List of child ages (defaults to [5, 8, 12, ...])
|
||||
state (str): Two-letter state code
|
||||
parent_age (int): Parent's age
|
||||
married (bool): Whether parents are married
|
||||
spouse_income (float): Spouse's income if married
|
||||
**kwargs: Additional household attributes
|
||||
|
||||
Returns:
|
||||
dict: PolicyEngine situation dictionary
|
||||
"""
|
||||
if child_ages is None:
|
||||
child_ages = [5 + i * 3 for i in range(num_children)]
|
||||
elif len(child_ages) != num_children:
|
||||
raise ValueError("Length of child_ages must match num_children")
|
||||
|
||||
people = {
|
||||
"parent": {
|
||||
"age": {CURRENT_YEAR: parent_age},
|
||||
"employment_income": {CURRENT_YEAR: parent_income}
|
||||
}
|
||||
}
|
||||
|
||||
members = ["parent"]
|
||||
|
||||
if married:
|
||||
people["spouse"] = {
|
||||
"age": {CURRENT_YEAR: parent_age},
|
||||
"employment_income": {CURRENT_YEAR: spouse_income}
|
||||
}
|
||||
members.append("spouse")
|
||||
|
||||
for i, age in enumerate(child_ages):
|
||||
child_id = f"child_{i+1}"
|
||||
people[child_id] = {"age": {CURRENT_YEAR: age}}
|
||||
members.append(child_id)
|
||||
|
||||
household_attrs = {
|
||||
"members": members,
|
||||
"state_name": {CURRENT_YEAR: state}
|
||||
}
|
||||
household_attrs.update({k: {CURRENT_YEAR: v} for k, v in kwargs.items()})
|
||||
|
||||
return {
|
||||
"people": people,
|
||||
"families": {"family": {"members": members}},
|
||||
"marital_units": {
|
||||
"marital_unit": {
|
||||
"members": members if married else ["parent"]
|
||||
}
|
||||
},
|
||||
"tax_units": {"tax_unit": {"members": members}},
|
||||
"spm_units": {"spm_unit": {"members": members}},
|
||||
"households": {"household": household_attrs}
|
||||
}
|
||||
|
||||
|
||||
def add_itemized_deductions(
|
||||
situation,
|
||||
charitable_donations=0,
|
||||
mortgage_interest=0,
|
||||
real_estate_taxes=0,
|
||||
medical_expenses=0,
|
||||
casualty_losses=0
|
||||
):
|
||||
"""
|
||||
Add itemized deductions to an existing situation.
|
||||
|
||||
Adds deductions to the first person in the situation.
|
||||
|
||||
Args:
|
||||
situation (dict): Existing PolicyEngine situation
|
||||
charitable_donations (float): Cash charitable contributions
|
||||
mortgage_interest (float): Mortgage interest paid
|
||||
real_estate_taxes (float): State and local property taxes
|
||||
medical_expenses (float): Medical and dental expenses
|
||||
casualty_losses (float): Casualty and theft losses
|
||||
|
||||
Returns:
|
||||
dict: Updated situation with deductions
|
||||
"""
|
||||
# Get first person ID
|
||||
first_person = list(situation["people"].keys())[0]
|
||||
|
||||
# Add deductions
|
||||
if charitable_donations > 0:
|
||||
situation["people"][first_person]["charitable_cash_donations"] = {
|
||||
CURRENT_YEAR: charitable_donations
|
||||
}
|
||||
|
||||
if mortgage_interest > 0:
|
||||
situation["people"][first_person]["mortgage_interest"] = {
|
||||
CURRENT_YEAR: mortgage_interest
|
||||
}
|
||||
|
||||
if real_estate_taxes > 0:
|
||||
situation["people"][first_person]["real_estate_taxes"] = {
|
||||
CURRENT_YEAR: real_estate_taxes
|
||||
}
|
||||
|
||||
if medical_expenses > 0:
|
||||
situation["people"][first_person]["medical_expense"] = {
|
||||
CURRENT_YEAR: medical_expenses
|
||||
}
|
||||
|
||||
if casualty_losses > 0:
|
||||
situation["people"][first_person]["casualty_loss"] = {
|
||||
CURRENT_YEAR: casualty_losses
|
||||
}
|
||||
|
||||
return situation
|
||||
|
||||
|
||||
def add_axes(situation, variable_name, min_val, max_val, count=1001):
|
||||
"""
|
||||
Add axes to a situation for parameter sweeps.
|
||||
|
||||
Args:
|
||||
situation (dict): Existing PolicyEngine situation
|
||||
variable_name (str): Variable to vary (e.g., "employment_income")
|
||||
min_val (float): Minimum value
|
||||
max_val (float): Maximum value
|
||||
count (int): Number of points (default: 1001)
|
||||
|
||||
Returns:
|
||||
dict: Updated situation with axes
|
||||
"""
|
||||
situation["axes"] = [[{
|
||||
"name": variable_name,
|
||||
"count": count,
|
||||
"min": min_val,
|
||||
"max": max_val,
|
||||
"period": CURRENT_YEAR
|
||||
}]]
|
||||
|
||||
return situation
|
||||
|
||||
|
||||
def set_state_nyc(situation, in_nyc=True):
|
||||
"""
|
||||
Set state to NY and configure NYC residence.
|
||||
|
||||
Args:
|
||||
situation (dict): Existing PolicyEngine situation
|
||||
in_nyc (bool): Whether household is in NYC
|
||||
|
||||
Returns:
|
||||
dict: Updated situation
|
||||
"""
|
||||
household_id = list(situation["households"].keys())[0]
|
||||
situation["households"][household_id]["state_name"] = {CURRENT_YEAR: "NY"}
|
||||
situation["households"][household_id]["in_nyc"] = {CURRENT_YEAR: in_nyc}
|
||||
|
||||
return situation
|
||||
Reference in New Issue
Block a user