Write 401(k) deferrals to the new *_desired inputs

PolicyEngine-US PR #8391 makes traditional_401k_contributions and
roth_401k_contributions formula variables that proportionally cap the
combined elective deferrals at the IRC §402(g) limit. The formulas read
the new pre-cap inputs traditional_401k_contributions_desired and
roth_401k_contributions_desired.

This change updates the CPS imputation and PUF clone pipelines to write
the desired inputs instead of the canonical names, so the §402(g) cap
takes effect in microsimulation. Calibration targets continue to point
at the canonical names (post-cap formula output), so the loss matrix
and database targets are unchanged.

IRA contribution variables are not renamed yet — those will move when
PolicyEngine/policyengine-us#8388 lands.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: MaxGhenis

changelog.d/rename-401k-contributions-to-desired.changed.md

@@ -0,0 +1 @@
+Write 401(k) elective deferrals to `traditional_401k_contributions_desired` and `roth_401k_contributions_desired` so PolicyEngine-US can apply the §402(g) cap to the canonical variable.

policyengine_us_data/calibration/puf_impute.py

@@ -160,8 +160,8 @@
 ]
 
 CPS_RETIREMENT_VARIABLES = [
-    "traditional_401k_contributions",
-    "roth_401k_contributions",
+    "traditional_401k_contributions_desired",
+    "roth_401k_contributions_desired",
     "traditional_ira_contributions",
     "roth_ira_contributions",
     "self_employed_pension_contributions",
@@ -886,8 +886,8 @@ def _impute_retirement_contributions(
     logger.info(
         "Imputed retirement contributions for PUF: "
         "401k mean=$%.0f, IRA mean=$%.0f, SE pension mean=$%.0f",
-        result["traditional_401k_contributions"].mean()
-        + result["roth_401k_contributions"].mean(),
+        result["traditional_401k_contributions_desired"].mean()
+        + result["roth_401k_contributions_desired"].mean(),
         result["traditional_ira_contributions"].mean()
         + result["roth_ira_contributions"].mean(),
         result["self_employed_pension_contributions"].mean(),

policyengine_us_data/datasets/cps/cps.py

@@ -1410,8 +1410,8 @@ def add_personal_income_variables(cps: h5py.File, person: DataFrame, year: int):
     # DC pool: split into traditional/Roth 401(k), cap at combined
     # 401(k) limit.
     dc_capped = np.minimum(dc_pool, limit_401k)
-    cps["traditional_401k_contributions"] = dc_capped * (1 - roth_dc_share)
-    cps["roth_401k_contributions"] = dc_capped * roth_dc_share
+    cps["traditional_401k_contributions_desired"] = dc_capped * (1 - roth_dc_share)
+    cps["roth_401k_contributions_desired"] = dc_capped * roth_dc_share
 
     # IRA pool: split into traditional/Roth IRA, cap at combined
     # IRA limit.

policyengine_us_data/datasets/cps/extended_cps.py

@@ -159,8 +159,8 @@ def _supports_structural_mortgage_inputs() -> bool:
     "taxable_sep_distributions",
     "tax_exempt_sep_distributions",
     # Retirement contributions
-    "traditional_401k_contributions",
-    "roth_401k_contributions",
+    "traditional_401k_contributions_desired",
+    "roth_401k_contributions_desired",
     "traditional_ira_contributions",
     "roth_ira_contributions",
     "self_employed_pension_contributions",
@@ -587,8 +587,8 @@ def apply_retirement_constraints(predictions, X_test, time_period):
 
     # Explicit mapping: variable -> (cap array, zero_mask or None).
     _CONSTRAINT_MAP = {
-        "traditional_401k_contributions": (limit_401k, emp_income == 0),
-        "roth_401k_contributions": (limit_401k, emp_income == 0),
+        "traditional_401k_contributions_desired": (limit_401k, emp_income == 0),
+        "roth_401k_contributions_desired": (limit_401k, emp_income == 0),
         "traditional_ira_contributions": (limit_ira, None),
         "roth_ira_contributions": (limit_ira, None),
         "self_employed_pension_contributions": (
@@ -641,8 +641,8 @@ def reconcile_ss_subcomponents(predictions, total_ss):
 
 
 _RETIREMENT_VARS = {
-    "traditional_401k_contributions",
-    "roth_401k_contributions",
+    "traditional_401k_contributions_desired",
+    "roth_401k_contributions_desired",
     "traditional_ira_contributions",
     "roth_ira_contributions",
     "self_employed_pension_contributions",

tests/unit/calibration/test_calibration_puf_impute.py

@@ -406,8 +406,8 @@ def calculate_dataframe(self, columns):
                         "qualified_dividend_income": [0.0, 0.0],
                         "taxable_pension_income": [0.0, 0.0],
                         "social_security": [0.0, 0.0],
-                        "traditional_401k_contributions": [0.0, 0.0],
-                        "roth_401k_contributions": [0.0, 0.0],
+                        "traditional_401k_contributions_desired": [0.0, 0.0],
+                        "roth_401k_contributions_desired": [0.0, 0.0],
                         "traditional_ira_contributions": [0.0, 0.0],
                         "roth_ira_contributions": [0.0, 0.0],
                         "self_employed_pension_contributions": [0.0, 0.0],
@@ -446,8 +446,8 @@ def fit_predict(
             )
             return pd.DataFrame(
                 {
-                    "traditional_401k_contributions": [0.0, 0.0],
-                    "roth_401k_contributions": [0.0, 0.0],
+                    "traditional_401k_contributions_desired": [0.0, 0.0],
+                    "roth_401k_contributions_desired": [0.0, 0.0],
                     "traditional_ira_contributions": [0.0, 0.0],
                     "roth_ira_contributions": [0.0, 0.0],
                     "self_employed_pension_contributions": [50_000.0, 50_000.0],

tests/unit/calibration/test_retirement_imputation.py

@@ -89,8 +89,8 @@ def _make_cps_df(n, rng):
             "taxable_pension_income": rng.uniform(0, 20_000, n),
             "social_security": rng.uniform(0, 15_000, n),
             # Targets
-            "traditional_401k_contributions": rng.uniform(0, 5000, n),
-            "roth_401k_contributions": rng.uniform(0, 3000, n),
+            "traditional_401k_contributions_desired": rng.uniform(0, 5000, n),
+            "roth_401k_contributions_desired": rng.uniform(0, 3000, n),
             "traditional_ira_contributions": rng.uniform(0, 2000, n),
             "roth_ira_contributions": rng.uniform(0, 2000, n),
             "self_employed_pension_contributions": rng.uniform(0, 10_000, n),
@@ -142,8 +142,8 @@ def test_five_retirement_variables(self):
 
     def test_retirement_variable_names(self):
         expected = {
-            "traditional_401k_contributions",
-            "roth_401k_contributions",
+            "traditional_401k_contributions_desired",
+            "roth_401k_contributions_desired",
             "traditional_ira_contributions",
             "roth_ira_contributions",
             "self_employed_pension_contributions",
@@ -321,8 +321,8 @@ def test_401k_capped(self):
         max_401k = lim["401k"] + lim["401k_catch_up"]
 
         for var in (
-            "traditional_401k_contributions",
-            "roth_401k_contributions",
+            "traditional_401k_contributions_desired",
+            "roth_401k_contributions_desired",
         ):
             assert np.all(result[var] <= max_401k), f"{var} exceeds 401k limit"
 
@@ -343,8 +343,8 @@ def test_401k_zero_when_no_wages(self):
         assert zero_wage.sum() == 10
 
         for var in (
-            "traditional_401k_contributions",
-            "roth_401k_contributions",
+            "traditional_401k_contributions_desired",
+            "roth_401k_contributions_desired",
         ):
             assert np.all(result[var][zero_wage] == 0), (
                 f"{var} should be 0 when employment_income is 0"
@@ -369,8 +369,8 @@ def test_catch_up_age_threshold(self):
 
         result = self._call_with_mocks(self._uniform_preds(val))
 
-        young_401k = result["traditional_401k_contributions"][:25]
-        old_401k = result["traditional_401k_contributions"][25:]
+        young_401k = result["traditional_401k_contributions_desired"][:25]
+        old_401k = result["traditional_401k_contributions_desired"][25:]
 
         # Young capped at base limit
         assert np.all(young_401k == lim["401k"])
@@ -397,8 +397,8 @@ def test_401k_nonzero_for_positive_wages(self):
         result = self._call_with_mocks(self._uniform_preds(5_000.0))
         pos_wage = self.puf_imputations["employment_income"] > 0
         for var in (
-            "traditional_401k_contributions",
-            "roth_401k_contributions",
+            "traditional_401k_contributions_desired",
+            "roth_401k_contributions_desired",
         ):
             assert np.all(result[var][pos_wage] > 0)
 

tests/unit/test_extended_cps.py

@@ -143,8 +143,8 @@ def test_cps_only_vars_mostly_exist_in_tbs(self):
     def test_retirement_contributions_in_cps_only(self):
         """All 5 retirement contribution vars should be in CPS_ONLY."""
         expected = {
-            "traditional_401k_contributions",
-            "roth_401k_contributions",
+            "traditional_401k_contributions_desired",
+            "roth_401k_contributions_desired",
             "traditional_ira_contributions",
             "roth_ira_contributions",
             "self_employed_pension_contributions",
@@ -912,8 +912,14 @@ class TestRetirementConstraints:
     def sample_predictions(self):
         return pd.DataFrame(
             {
-                "traditional_401k_contributions": [25000, -500, 5000, 10000, 3000],
-                "roth_401k_contributions": [30000, 2000, 0, 50000, 1000],
+                "traditional_401k_contributions_desired": [
+                    25000,
+                    -500,
+                    5000,
+                    10000,
+                    3000,
+                ],
+                "roth_401k_contributions_desired": [30000, 2000, 0, 50000, 1000],
                 "traditional_ira_contributions": [8000, -100, 3000, 15000, 500],
                 "roth_ira_contributions": [10000, 1000, 0, 20000, 200],
                 "self_employed_pension_contributions": [80000, -200, 5000, 0, 100000],
@@ -943,7 +949,10 @@ def test_401k_capped_at_limit(self, sample_predictions, sample_features):
         age = sample_features["age"].values
         catch_up = age >= 50
         cap = limits["401k"] + catch_up * limits["401k_catch_up"]
-        for var in ["traditional_401k_contributions", "roth_401k_contributions"]:
+        for var in [
+            "traditional_401k_contributions_desired",
+            "roth_401k_contributions_desired",
+        ]:
             assert (result[var].values <= cap).all(), f"{var} exceeds 401k cap"
 
     def test_ira_capped_at_limit(self, sample_predictions, sample_features):
@@ -962,7 +971,10 @@ def test_401k_zeroed_without_employment_income(
     ):
         result = apply_retirement_constraints(sample_predictions, sample_features, 2024)
         no_emp = sample_features["employment_income"] == 0
-        for var in ["traditional_401k_contributions", "roth_401k_contributions"]:
+        for var in [
+            "traditional_401k_contributions_desired",
+            "roth_401k_contributions_desired",
+        ]:
             assert (result[var].values[no_emp] == 0).all(), (
                 f"{var} should be zero without employment income"
             )