Cap 401k contributions at wages and salaries in the CPS

[MaxGhenis]

Code to do this in cps.py

# Assume all retirement contributions are traditional 401(k) for now.
# Procedure for allocating retirement contributions:
# 1) If they report any self-employment income, allocate entirely to
#    self-employed pension contributions.
# 2) If they report any wage and salary income, allocate in this order:
#    a) Traditional 401(k) contributions up to the limit or WSAL, whichever is lower
#    b) Roth 401(k) contributions up to the limit or WSAL, whichever is lower
#    c) IRA contributions up to the limit, split according to administrative fractions
#    d) Other retirement contributions
# Disregard reported pension contributions from people who report neither wage and salary
# nor self-employment income.
# Assume no 403(b) or 457 contributions for now.
LIMIT_401K_2022 = 20_500
LIMIT_401K_CATCH_UP_2022 = 6_500
LIMIT_IRA_2022 = 6_000
LIMIT_IRA_CATCH_UP_2022 = 1_000
CATCH_UP_AGE_2022 = 50

retirement_contributions = person.RETCB_VAL

cps["self_employed_pension_contributions"] = np.where(
    person.SEMP_VAL > 0, retirement_contributions, 0
)

remaining_retirement_contributions = np.maximum(
    retirement_contributions - cps["self_employed_pension_contributions"], 0,
)

# Compute the 401(k) limit for the person's age.
catch_up_eligible = person.A_AGE >= CATCH_UP_AGE_2022
limit_401k = LIMIT_401K_2022 + catch_up_eligible * LIMIT_401K_CATCH_UP_2022
limit_ira = LIMIT_IRA_2022 + catch_up_eligible * LIMIT_IRA_CATCH_UP_2022

cps["traditional_401k_contributions"] = np.where(
    person.WSAL_VAL > 0,
    np.minimum(remaining_retirement_contributions, np.minimum(person.WSAL_VAL, limit_401k)),
    0,
)

remaining_retirement_contributions = np.maximum(
    remaining_retirement_contributions - cps["traditional_401k_contributions"], 0,
)

cps["roth_401k_contributions"] = np.where(
    person.WSAL_VAL > 0,
    np.minimum(remaining_retirement_contributions, np.minimum(person.WSAL_VAL, limit_401k)),
    0,
)

remaining_retirement_contributions = np.maximum(
    remaining_retirement_contributions - cps["roth_401k_contributions"], 0,
)

cps["traditional_ira_contributions"] = np.where(
    person.WSAL_VAL > 0,
    np.minimum(remaining_retirement_contributions, limit_ira),
    0,
)

remaining_retirement_contributions = np.maximum(
    remaining_retirement_contributions - cps["traditional_ira_contributions"], 0,
)

roth_ira_limit = limit_ira - cps["traditional_ira_contributions"]

cps["roth_ira_contributions"] = np.where(
    person.WSAL_VAL > 0,
    np.minimum(remaining_retirement_contributions, roth_ira_limit),
    0,
)