Activate transitory medical-shock folding in ACA benchmarks

.github/workflows/main.yml

@@ -35,7 +35,7 @@ jobs:
         # install; keep it compatible with the model's graph declarations.
         run: >-
           pip install "pylcm @
-          git+https://github.com/OpenSourceEconomics/pylcm.git@43a4bf2b0ad003b7308dd3478cd008458b528102"
+          git+https://github.com/OpenSourceEconomics/pylcm.git@eea73a8f0e301e88f02278df3595c797fdba24df"
       - name: Install aca-model with test deps
         run: pip install -e . pytest pdbp
       - name: Run pytest

README.md

@@ -2,6 +2,14 @@
 
 Core lifecycle model for the ACA structural retirement project.
 
+The benchmark folds the IID transitory medical-cost shock and declares its asset
+transition as resources plus the negated medical bill. Production grid defaults remain
+unfolded. To measure the unfolded benchmark control, pass
+`grid_config=replace(BENCHMARK_GRID_CONFIG, fold_hcc_transitory=False, subtract_hcc_bill=False)`
+to `create_benchmark_model`. Pass the same grid configuration to
+`get_benchmark_initial_conditions`; both routes retain identical seeded draws for their
+shared initial states.
+
 Model factories accept `execution_config=ExecutionConfig(...)`. An omitted policy uses
 the smallest JAX allocator `bytes_limit` among the selected accelerators as an explicit
 per-device budget. This enables budget-aware width selection; CPU construction remains

src/aca_model/aca/health_insurance.py

@@ -129,6 +129,39 @@ def cost_sharing(
     return jnp.where(is_insured & ~is_medicaid_eligible, scale, 1.0)
 
 
+def cost_sharing_bracket(
+    gross_income: FloatND,
+    spousal_income: DiscreteState,
+    buy_private: DiscreteAction,
+    is_medicaid_eligible: BoolND,
+    cost_sharing_schedule: MappingLeaf,
+) -> IntND:
+    """Return the household's cost-sharing bracket, `n_kinks + 1` when neutral.
+
+    The bracket is the FPL bracket of `cost_sharing`; a household without
+    private cover, or Medicaid-eligible, falls in the neutral bracket one past
+    the last, whose scale is 1.0. The bracket takes finitely many values, so
+    the out-of-pocket bill can be averaged per bracket.
+    """
+    sched = cast("Mapping[str, Any]", cost_sharing_schedule.data)
+    kinks = sched["kinks"]  # [n_kinks, 3]
+    bracket = jnp.searchsorted(kinks[:, spousal_income], gross_income, side="right")
+    is_insured = buy_private == BuyPrivate.yes
+    return jnp.where(
+        is_insured & ~is_medicaid_eligible, bracket, kinks.shape[0] + 1
+    ).astype(jnp.int32)
+
+
+def cost_sharing_scale_of_bracket(
+    cost_sharing_bracket: IntND,
+    cost_sharing_schedule: MappingLeaf,
+) -> FloatND:
+    """Return the cost-sharing scale of a bracket from `cost_sharing_bracket`."""
+    sched = cast("Mapping[str, Any]", cost_sharing_schedule.data)
+    factors = sched["factors"]  # [n_kinks + 1]
+    return jnp.append(factors, jnp.ones((), dtype=factors.dtype))[cost_sharing_bracket]
+
+
 def is_medicaid_eligible(
     is_ssi_eligible: BoolND,
     aca_magi: FloatND,

src/aca_model/aca/regimes/_overrides.py

@@ -67,7 +67,10 @@ def apply_aca_overrides(
         # risk-rated premium's zero-profit intercept has no consumer left.
         functions.pop("private_premium_intercept", None)
         functions["hic_premium_subsidy"] = aca_hi.premium_subsidy
-        functions["cost_sharing_scale"] = aca_hi.cost_sharing
+        # Split at the bracket, which takes finitely many values, so the
+        # out-of-pocket bill can be averaged per bracket.
+        functions["cost_sharing_bracket"] = aca_hi.cost_sharing_bracket
+        functions["cost_sharing_scale"] = aca_hi.cost_sharing_scale_of_bracket
         functions["premium_default"] = aca_hi.premium_default
         functions["cash_on_hand"] = aca_hi.cash_on_hand
         functions["primary_oop"] = aca_hi.primary_oop

src/aca_model/agent/assets_and_income.py

@@ -109,6 +109,31 @@ def next_assets(
     )
 
 
+def assets_before_bill(
+    cash_on_hand: FloatND,
+    transfers: FloatND,
+    pension_assets_adjustment: FloatND,
+    consumption_dollars: ContinuousAction,
+) -> FloatND:
+    """Compute next-period assets for non-terminal targets before the medical bill.
+
+    `next_assets` is this minus `oop_costs`.
+    """
+    return cash_on_hand + transfers + pension_assets_adjustment - consumption_dollars
+
+
+def assets_before_bill_when_dead(
+    cash_on_hand: FloatND,
+    transfers: FloatND,
+    consumption_dollars: ContinuousAction,
+) -> FloatND:
+    """Compute assets for the dead/terminal target before the medical bill.
+
+    `next_assets_when_dead` is this minus `oop_costs`.
+    """
+    return cash_on_hand + transfers - consumption_dollars
+
+
 def next_assets_when_dead(
     cash_on_hand: FloatND,
     transfers: FloatND,

src/aca_model/baseline/health_insurance.py

@@ -357,6 +357,11 @@ def oop_costs(
     return jnp.minimum(oop, oop_max)
 
 
+def negative_oop_costs(oop_costs: FloatND) -> FloatND:
+    """Return the out-of-pocket bill as the shock added to assets."""
+    return -oop_costs
+
+
 def insured_oop(
     total_health_costs: FloatND,
     is_married: IntND,

src/aca_model/baseline/regimes/_common.py

@@ -14,6 +14,7 @@
 import jax.numpy as jnp
 import numpy as np
 from lcm import (
+    AdditiveShockTransition,
     ByAge,
     DiscreteGrid,
     GridBreakpoint,
@@ -325,6 +326,7 @@ def build_grids(
         gauss_hermite=True,
         mu=0.0,
         sigma=1.0,
+        fold=grid_config.fold_hcc_transitory,
     )
 
     assets_start = -_compute_max_annual_labor_income(
@@ -948,7 +950,7 @@ def build_model_state_transitions() -> dict:
     }
 
 
-def build_common_functions(spec: RegimeSpec) -> dict:
+def build_common_functions(spec: RegimeSpec, *, subtract_bill: bool = False) -> dict:
     """Build the regime-level functions dict for a non-dead regime.
 
     Contains the spec-dependent selections and the overlay-swapped names;
@@ -1003,6 +1005,12 @@ def build_common_functions(spec: RegimeSpec) -> dict:
     # Swapped per policy variant by the ACA overlay, hence regime-level
     functions["is_medicaid_eligible"] = health_insurance.is_medicaid_eligible
     functions["after_tax_income"] = taxes.after_tax_income
+    if subtract_bill:
+        functions["assets_before_bill"] = assets_and_income.assets_before_bill
+        functions["negative_oop_costs"] = health_insurance.negative_oop_costs
+        functions["assets_before_bill_when_dead"] = (
+            assets_and_income.assets_before_bill_when_dead
+        )
     functions["premium_default"] = assets_and_income.premium_default
     functions["cash_on_hand"] = assets_and_income.cash_on_hand
 
@@ -1172,7 +1180,7 @@ def select_target_for_age(
 
 
 def build_state_transitions(
-    spec: RegimeSpec, *, solver: SolverName = "brute_force"
+    spec: RegimeSpec, *, solver: SolverName = "brute_force", subtract_bill: bool = False
 ) -> dict:
     """Build the regime-level state transitions dict for a non-dead regime.
 
@@ -1183,7 +1191,9 @@ def build_state_transitions(
     post-decision (savings) form.
     """
     transitions: dict = {}
-    transitions["assets"] = _build_per_target_regime_assets(spec, solver=solver)
+    transitions["assets"] = _build_per_target_regime_assets(
+        spec, solver=solver, subtract_bill=subtract_bill
+    )
     transitions["health"] = _build_per_target_regime_health(spec)
     claimed_ss_transition = _build_per_target_regime_claimed_ss(spec)
     if claimed_ss_transition:
@@ -1221,8 +1231,24 @@ def _select_aime_law(spec: RegimeSpec) -> Callable[..., FloatND]:
     )
 
 
+# Every input of the out-of-pocket bill that varies across a source's states and
+# actions, with its support. The bill reads nothing else besides next period's
+# medical-cost draws, the period and parameters.
+_BILL_CONDITIONERS: Mapping[str, tuple[int | bool, ...]] = MappingProxyType(
+    {
+        "is_married": (0, 1),
+        "good_health": (0, 1),
+        "is_medicaid_eligible": (False, True),
+        "buy_private": (int(BuyPrivate.no), int(BuyPrivate.yes)),
+        # Under the ACA policy: the four FPL kinks of the cost-sharing schedule
+        # give brackets 0-4, and 5 is the neutral bracket.
+        "cost_sharing_bracket": tuple(range(6)),
+    }
+)
+
+
 def _build_per_target_regime_assets(
-    spec: RegimeSpec, *, solver: SolverName = "brute_force"
+    spec: RegimeSpec, *, solver: SolverName = "brute_force", subtract_bill: bool = False
 ) -> dict[RegimeName, Callable[..., FloatND]]:
     """Build per-target assets transitions.
 
@@ -1231,11 +1257,30 @@ def _build_per_target_regime_assets(
     pull in the `next_aime`-dependent imputation chain — `dead` has no
     `aime` state and pylcm cannot resolve `next_aime` there. Non-dead
     targets use the full `next_assets` with the pension correction.
-    Under DC-EGM both laws take their post-decision (savings) form.
+    Under DC-EGM both laws take their post-decision (savings) form. With
+    `subtract_bill`, each law is declared as its resources before the bill
+    (`assets_before_bill`, `assets_before_bill_when_dead`) plus the shock
+    `negative_oop_costs`, conditioned on every source-side input of the bill.
     """
+    if subtract_bill and solver != "brute_force":
+        msg = f"subtract_hcc_bill needs the brute-force asset law, not {solver!r}."
+        raise ValueError(msg)
+    living_law: Callable[..., FloatND]
+    dead_law: Callable[..., FloatND]
     if solver in ("dcegm", "nbegm"):
         living_law = assets_and_income.next_assets_from_savings
         dead_law = assets_and_income.next_assets_when_dead_from_savings
+    elif subtract_bill:
+        living_law = AdditiveShockTransition(
+            base="assets_before_bill",
+            shock="negative_oop_costs",
+            conditioners=_BILL_CONDITIONERS,
+        )
+        dead_law = AdditiveShockTransition(
+            base="assets_before_bill_when_dead",
+            shock="negative_oop_costs",
+            conditioners=_BILL_CONDITIONERS,
+        )
     else:
         living_law = assets_and_income.next_assets
         dead_law = assets_and_income.next_assets_when_dead

src/aca_model/baseline/regimes/_nongroup.py

@@ -92,10 +92,10 @@ def transition(
     return transition
 
 
-def _build_functions(spec: RegimeSpec) -> dict:
+def _build_functions(spec: RegimeSpec, *, subtract_bill: bool = False) -> dict:
     """Build functions dict for a nongroup regime."""
     can_work = spec["canwork"] == "canwork"
-    functions = build_common_functions(spec)
+    functions = build_common_functions(spec, subtract_bill=subtract_bill)
 
     functions["ss_benefit"] = select_ss_benefit(spec)
 
@@ -168,7 +168,9 @@ def build_regime(
         if egm_solver is None
         else ("nbegm" if nbegm_solver is not None else "dcegm")
     )
-    functions = _build_functions(spec)
+    functions = _build_functions(
+        spec, subtract_bill=grids.grid_config.subtract_hcc_bill
+    )
     # The EGM solvers reject constraints that read the continuous state, which
     # the private-cover constraint does through `premium_default`.
     constraints = build_constraints(spec) if egm_solver is None else {}
@@ -183,7 +185,11 @@ def build_regime(
     return build_alive_regime(
         egm_solver=egm_solver,
         states=states,
-        state_transitions=build_state_transitions(spec, solver=state_solver),
+        state_transitions=build_state_transitions(
+            spec,
+            solver=state_solver,
+            subtract_bill=grids.grid_config.subtract_hcc_bill,
+        ),
         actions=build_actions(spec, grids),
         functions=functions,
         constraints=constraints,

src/aca_model/baseline/regimes/_retiree.py

@@ -113,10 +113,10 @@ def transition(
     return transition
 
 
-def _build_functions(spec: RegimeSpec) -> dict:
+def _build_functions(spec: RegimeSpec, *, subtract_bill: bool = False) -> dict:
     """Build functions dict for a retiree regime."""
     can_work = spec["canwork"] == "canwork"
-    functions = build_common_functions(spec)
+    functions = build_common_functions(spec, subtract_bill=subtract_bill)
 
     functions["ss_benefit"] = select_ss_benefit(spec)
 
@@ -164,7 +164,9 @@ def build_regime(
         if egm_solver is None
         else ("nbegm" if nbegm_solver is not None else "dcegm")
     )
-    functions = _build_functions(spec)
+    functions = _build_functions(
+        spec, subtract_bill=grids.grid_config.subtract_hcc_bill
+    )
     if nbegm_solver is not None:
         # NBEGM's solver contract is stated per regime: it reads the budget in
         # savings form off `resources` and the post-decision node off
@@ -173,7 +175,11 @@ def build_regime(
     return build_alive_regime(
         egm_solver=egm_solver,
         states=states,
-        state_transitions=build_state_transitions(spec, solver=state_solver),
+        state_transitions=build_state_transitions(
+            spec,
+            solver=state_solver,
+            subtract_bill=grids.grid_config.subtract_hcc_bill,
+        ),
         actions=build_actions(spec, grids),
         functions=functions,
     )

src/aca_model/baseline/regimes/_tied.py

@@ -86,9 +86,9 @@ def target(mc_next: bool) -> IntND:
     return transition
 
 
-def _build_functions(spec: RegimeSpec) -> dict:
+def _build_functions(spec: RegimeSpec, *, subtract_bill: bool = False) -> dict:
     """Build functions dict for a tied regime."""
-    functions = build_common_functions(spec)
+    functions = build_common_functions(spec, subtract_bill=subtract_bill)
 
     functions["ss_benefit"] = select_ss_benefit(spec)
 
@@ -127,7 +127,9 @@ def build_regime(
         if egm_solver is None
         else ("nbegm" if nbegm_solver is not None else "dcegm")
     )
-    functions = _build_functions(spec)
+    functions = _build_functions(
+        spec, subtract_bill=grids.grid_config.subtract_hcc_bill
+    )
     if nbegm_solver is not None:
         # NBEGM's solver contract is stated per regime: it reads the budget in
         # savings form off `resources` and the post-decision node off
@@ -136,7 +138,11 @@ def build_regime(
     return build_alive_regime(
         egm_solver=egm_solver,
         states=states,
-        state_transitions=build_state_transitions(spec, solver=state_solver),
+        state_transitions=build_state_transitions(
+            spec,
+            solver=state_solver,
+            subtract_bill=grids.grid_config.subtract_hcc_bill,
+        ),
         actions=build_actions(spec, grids),
         functions=functions,
     )