Wire pension imputation correction (FJ 2011 Appendix A.5)

Two new DAG functions in canwork & ss != "forced" regimes:

- target_his(his, labor_supply, is_medicaid_eligible): HIS class of the
  surviving target regime. Mirrors the cross-HIS branches inside
  _make_transition_canwork (tied → nongroup when stopping work,
  Medicaid override → nongroup).
- imputed_pension_wealth_next_period(next_aime, target_his, period, ...):
  computes pw_next_imputed = benefit_imputed(next_pia, next_period,
  target_his) · epdv_constant_pension[next_period] using bare-name
  parameters into 1-period-shifted views of the imputation arrays
  (`*_next_period`). Inlining is required because pylcm's AST shape
  inference doesn't trace nested calls into pensions.benefit.

next_assets continues to consume pension_assets_adjustment, which now
sees a real imputed_pension_wealth_next_period via the DAG (previously
fixed to 0.0 in aca-estimation). The chained dependency
next_aime → imputed_pension_wealth_next_period → pension_assets_adjustment
is unblocked by pylcm exempting next_<state> names from
fixed_param extraction (PR pylcm#342).

Also drops pension_assets_adjustment from borrowing_constraint: a
negative correction at a cross-HIS transition can leave no feasible
action and inject `-inf` into V via `argmax_and_max(initial=-inf,
where=F_arr)`, which then cancels with `0 * -inf = NaN`. The
correction is a post-decision shift on next-period assets and must
not gate the current consumption choice.

Author: hmgaudecker

src/aca_model/agent/assets_and_income.py

@@ -69,7 +69,14 @@ def borrowing_constraint(
     consumption: ContinuousAction,
     cash_on_hand: FloatND,
     transfers: FloatND,
-    pension_assets_adjustment: FloatND,
 ) -> BoolND:
-    """Consumption cannot exceed available resources (no borrowing)."""
-    return consumption <= cash_on_hand + transfers + pension_assets_adjustment
+    """Consumption cannot exceed available resources (no borrowing).
+
+    `pension_assets_adjustment` is excluded: it can be negative (e.g.,
+    when the imputation overstates next-period pension wealth at a
+    cross-HIS transition), and including it here can leave no feasible
+    action at low-asset / mid-AIME corners. The correction enters
+    `next_assets` instead — a post-decision shift that does not gate
+    the current consumption choice.
+    """
+    return consumption <= cash_on_hand + transfers

src/aca_model/baseline/health_insurance.py

@@ -246,6 +246,29 @@ def is_medicaid_eligible(is_ssi_eligible: BoolND) -> BoolND:
     return is_ssi_eligible
 
 
+def target_his(
+    his: IntND,
+    labor_supply: DiscreteAction,
+    is_medicaid_eligible: BoolND,
+) -> IntND:
+    """Return the HIS class of the surviving target regime.
+
+    Mirrors the cross-HIS branches inside `_make_transition_canwork` (retiree,
+    tied, nongroup): tied agents who stop working become nongroup, and
+    Medicaid-eligible agents are overridden to nongroup. Used by
+    `imputed_pension_wealth_next_period` to look up next-period imputation
+    coefficients at the target's HIS.
+    """
+    tied_to_ng = (his == HealthInsuranceState.tied) & (
+        labor_supply == LaborSupply.do_not_work
+    )
+    return jnp.where(
+        tied_to_ng | is_medicaid_eligible,
+        HealthInsuranceState.nongroup,
+        his,
+    ).astype(jnp.int32)
+
+
 def oop_with_medicaid(
     primary_oop: FloatND,
     is_medicaid_eligible: BoolND,

src/aca_model/baseline/regimes/_nongroup.py

@@ -99,6 +99,10 @@ def _build_functions(spec: dict[str, str]) -> dict:
         functions["pension_wealth_next_before_adjustment"] = (
             pensions.wealth_next_before_adjustment
         )
+        functions["target_his"] = health_insurance.target_his
+        functions["imputed_pension_wealth_next_period"] = (
+            pensions.imputed_pension_wealth_next_period
+        )
         functions["pension_assets_adjustment"] = pensions.assets_adjustment
         functions["total_to_pia"] = pensions.total_to_pia
 

src/aca_model/baseline/regimes/_retiree.py

@@ -109,6 +109,10 @@ def _build_functions(spec: dict[str, str]) -> dict:
         functions["pension_wealth_next_before_adjustment"] = (
             pensions.wealth_next_before_adjustment
         )
+        functions["target_his"] = health_insurance.target_his
+        functions["imputed_pension_wealth_next_period"] = (
+            pensions.imputed_pension_wealth_next_period
+        )
         functions["pension_assets_adjustment"] = pensions.assets_adjustment
         functions["total_to_pia"] = pensions.total_to_pia
 

src/aca_model/baseline/regimes/_tied.py

@@ -83,6 +83,10 @@ def _build_functions(spec: dict[str, str]) -> dict:
         functions["pension_wealth_next_before_adjustment"] = (
             pensions.wealth_next_before_adjustment
         )
+        functions["target_his"] = health_insurance.target_his
+        functions["imputed_pension_wealth_next_period"] = (
+            pensions.imputed_pension_wealth_next_period
+        )
         functions["pension_assets_adjustment"] = pensions.assets_adjustment
         functions["total_to_pia"] = pensions.total_to_pia
 

src/aca_model/environment/pensions.py

@@ -4,7 +4,7 @@
 """
 
 import jax.numpy as jnp
-from lcm.typing import FloatND, IntND, Period
+from lcm.typing import ContinuousState, FloatND, IntND, Period
 
 
 def benefit(
@@ -164,3 +164,42 @@ def assets_adjustment(
         * unconditional_survival_prob[period]
         * (pension_wealth_next_before_adjustment - imputed_pension_wealth_next_period)
     )
+
+
+def imputed_pension_wealth_next_period(
+    next_aime: ContinuousState,
+    target_his: IntND,
+    period: Period,
+    pia_table: FloatND,
+    pia_aime_grid: FloatND,
+    imp_intercept_next_period: FloatND,
+    imp_pia_coeff_next_period: FloatND,
+    imp_pia_kink_0_coeff_next_period: FloatND,
+    imp_pia_kink_1_coeff_next_period: FloatND,
+    imp_kink_0_next_period: FloatND,
+    imp_kink_1_next_period: FloatND,
+    imp_fraction_receiving_next_period: FloatND,
+    epdv_constant_pension_next_period: FloatND,
+) -> FloatND:
+    """Imputed pension wealth at next period using the target regime's HIS.
+
+    Mirrors `benefit` and `wealth` but indexes into 1-period-shifted views
+    of the imputation arrays so all subscripts use bare-name parameters
+    (`period`, `target_his`). Inlining is required: pylcm's AST shape
+    inference inspects the registered function's body and does not trace
+    through nested calls into `benefit`.
+    """
+    next_pia = jnp.interp(next_aime, pia_aime_grid, pia_table)
+
+    intercept = imp_intercept_next_period[period, target_his]
+    pia_pred = imp_pia_coeff_next_period[period, target_his] * next_pia
+    kink_0_adj = imp_pia_kink_0_coeff_next_period[period, target_his] * jnp.maximum(
+        0.0, next_pia - imp_kink_0_next_period[period]
+    )
+    kink_1_adj = imp_pia_kink_1_coeff_next_period[period, target_his] * jnp.maximum(
+        0.0, next_pia - imp_kink_1_next_period[period]
+    )
+
+    full_benefit = jnp.maximum(0.0, intercept + pia_pred + kink_0_adj + kink_1_adj)
+    benefit_next = full_benefit * imp_fraction_receiving_next_period[period]
+    return benefit_next * epdv_constant_pension_next_period[period]