From a5e695080ef528652e5c73521aeee0b99c5c5eed Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Fri, 9 Oct 2026 15:00:05 +0200 Subject: [PATCH 1/9] ENH Add a grid flag that folds the transitory medical-cost shock `GridConfig.fold_hcc_transitory` declares the transitory shock `fold=True`. Nothing a period computes reads last year's transitory shock, so pylcm then draws it on the transition only: living regimes keep the lagged persistent shock, which prices the private premium, and lose the transitory axis. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Hs8mfXPYyEa24KaAp6BwCW --- src/aca_model/baseline/regimes/_common.py | 1 + src/aca_model/config.py | 4 ++ tests/helpers/model.py | 8 +-- tests/test_transition_only_medical_shock.py | 62 +++++++++++++++++++++ 4 files changed, 71 insertions(+), 4 deletions(-) create mode 100644 tests/test_transition_only_medical_shock.py diff --git a/src/aca_model/baseline/regimes/_common.py b/src/aca_model/baseline/regimes/_common.py index 81ce991..d5909f0 100644 --- a/src/aca_model/baseline/regimes/_common.py +++ b/src/aca_model/baseline/regimes/_common.py @@ -325,6 +325,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( diff --git a/src/aca_model/config.py b/src/aca_model/config.py index 8b8f6fd..b33e8e2 100644 --- a/src/aca_model/config.py +++ b/src/aca_model/config.py @@ -34,6 +34,10 @@ class GridConfig: n_wage_res_gridpoints: int = 5 n_hcc_persistent_gridpoints: int = 3 n_hcc_transitory_gridpoints: int = 5 + # Fold the transitory medical-cost shock out of the state space. Nothing a + # period computes reads last year's transitory shock, so with the fold it is + # drawn on the transition only and no value carries its axis. + fold_hcc_transitory: bool = False # The post-decision savings grid is cubically clustered toward the borrowing # constraint. Its node count controls resolution under DC-EGM and NB-EGM. n_savings_gridpoints: int = 200 diff --git a/tests/helpers/model.py b/tests/helpers/model.py index 28cbb83..df958e8 100644 --- a/tests/helpers/model.py +++ b/tests/helpers/model.py @@ -16,7 +16,7 @@ from aca_model.baseline.model import create_model as _create_baseline_model from aca_model.baseline.regimes import REGIME_SPECS from aca_model.benchmark import get_benchmark_params -from aca_model.config import BENCHMARK_GRID_CONFIG +from aca_model.config import BENCHMARK_GRID_CONFIG, GridConfig _NONGROUP_BEFORE_MEDICARE = frozenset( name @@ -33,14 +33,14 @@ } -def make_baseline_model() -> Model: - """Baseline model on `BENCHMARK_GRID_CONFIG` with the benchmark snapshot params.""" +def make_baseline_model(*, grid_config: GridConfig = BENCHMARK_GRID_CONFIG) -> Model: + """Baseline model on `grid_config` with the benchmark snapshot params.""" fixed_params, wage_params, _ = get_benchmark_params(model=None) return _create_baseline_model( fixed_params=fixed_params, wage_params=wage_params, derived_categoricals=_DERIVED_CATEGORICALS, - grid_config=BENCHMARK_GRID_CONFIG, + grid_config=grid_config, pref_type_grid=DiscreteGrid(BenchmarkPrefType), ) diff --git a/tests/test_transition_only_medical_shock.py b/tests/test_transition_only_medical_shock.py new file mode 100644 index 0000000..4cb07f1 --- /dev/null +++ b/tests/test_transition_only_medical_shock.py @@ -0,0 +1,62 @@ +"""Folding the transitory medical-cost shock takes it out of every state space. + +Nothing a period computes reads last year's transitory shock: this year's draw +enters the out-of-pocket bill inside the end-of-period asset law, and the draw is +IID. With `fold_hcc_transitory`, the shock is therefore drawn on the transition +only. Living regimes keep the lagged persistent shock, which prices the private +premium, and lose the transitory axis; values agree with the unfolded model on +every transitory slice. +""" + +from dataclasses import replace + +import numpy as np +import pytest +from helpers.model import make_baseline_model + +from aca_model.benchmark import get_benchmark_params +from aca_model.config import BENCHMARK_GRID_CONFIG + +_FOLDED = replace(BENCHMARK_GRID_CONFIG, fold_hcc_transitory=True) +_REGIME = "retiree_nomc_inelig_canwork" + + +@pytest.fixture(scope="module") +def folded_model(): + return make_baseline_model(grid_config=_FOLDED) + + +def test_living_regimes_carry_no_lagged_transitory_shock(folded_model): + assert not any( + "hcc_transitory" in folded_model.state_names(regime_name=name) + for name in folded_model.user_regimes + ) + + +def test_living_regimes_keep_the_lagged_persistent_shock(folded_model): + assert all( + "hcc_persistent" in folded_model.state_names(regime_name=name) + for name in folded_model.user_regimes + if name != "dead" + ) + + +@pytest.mark.long_running +def test_folded_value_equals_every_transitory_slice_of_the_unfolded_value( + folded_model, +): + """The two models' sums differ only in order, so they agree to rounding.""" + unfolded_model = make_baseline_model() + _, _, params = get_benchmark_params(model=unfolded_model) + unfolded = np.asarray( + unfolded_model.solve(params=params, log_level="off").values[0][_REGIME] + ) + folded = np.asarray( + folded_model.solve(params=params, log_level="off").values[0][_REGIME] + ) + axis = unfolded_model.state_names(regime_name=_REGIME).index("hcc_transitory") + by_slice = np.moveaxis(unfolded, axis, 0) + + np.testing.assert_allclose( + by_slice, np.broadcast_to(folded, by_slice.shape), rtol=1e-6, atol=0 + ) From 649bed97cf05e662a79115f342437023872b50fd Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Fri, 9 Oct 2026 15:54:43 +0200 Subject: [PATCH 2/9] ENH Declare the asset law as resources minus the medical bill `GridConfig.subtract_hcc_bill` declares the asset law toward every target as `lcm.SubtractedBill`: assets before the out-of-pocket bill minus `oop_costs`, conditioned on marital status, health, Medicaid eligibility, the private-cover choice and the ACA cost-sharing bracket. With `fold_hcc_transitory`, pylcm then averages each target's value over the transitory medical-cost shock once per period instead of at every source point. The ACA cost-sharing scale is split at its income bracket, which takes finitely many values, so the bill can be conditioned on it; the scale is unchanged. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Hs8mfXPYyEa24KaAp6BwCW --- src/aca_model/aca/health_insurance.py | 33 ++++++++++++ src/aca_model/aca/regimes/_overrides.py | 5 +- src/aca_model/agent/assets_and_income.py | 25 +++++++++ src/aca_model/baseline/regimes/_common.py | 53 +++++++++++++++++-- src/aca_model/baseline/regimes/_nongroup.py | 14 ++++-- src/aca_model/baseline/regimes/_retiree.py | 14 ++++-- src/aca_model/baseline/regimes/_tied.py | 14 ++++-- src/aca_model/config.py | 5 ++ tests/test_aca_policies.py | 31 ++++++++++++ tests/test_subtracted_medical_bill.py | 56 +++++++++++++++++++++ 10 files changed, 232 insertions(+), 18 deletions(-) create mode 100644 tests/test_subtracted_medical_bill.py diff --git a/src/aca_model/aca/health_insurance.py b/src/aca_model/aca/health_insurance.py index c05e8ef..7e0b5bb 100644 --- a/src/aca_model/aca/health_insurance.py +++ b/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, diff --git a/src/aca_model/aca/regimes/_overrides.py b/src/aca_model/aca/regimes/_overrides.py index cdc904d..eb92ffe 100644 --- a/src/aca_model/aca/regimes/_overrides.py +++ b/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 diff --git a/src/aca_model/agent/assets_and_income.py b/src/aca_model/agent/assets_and_income.py index 2dad637..251c7f8 100644 --- a/src/aca_model/agent/assets_and_income.py +++ b/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, diff --git a/src/aca_model/baseline/regimes/_common.py b/src/aca_model/baseline/regimes/_common.py index d5909f0..4f63eab 100644 --- a/src/aca_model/baseline/regimes/_common.py +++ b/src/aca_model/baseline/regimes/_common.py @@ -26,6 +26,7 @@ Regime, RouwenhorstAR1Process, StochasticTransition, + SubtractedBill, categorical, fixed_transition, ) @@ -949,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; @@ -1004,6 +1005,11 @@ 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["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 @@ -1173,7 +1179,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. @@ -1184,7 +1190,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: @@ -1222,8 +1230,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": (BuyPrivate.no, 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. @@ -1232,11 +1256,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`) minus `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 = SubtractedBill( + resources="assets_before_bill", + bill="oop_costs", + conditioners=_BILL_CONDITIONERS, + ) + dead_law = SubtractedBill( + resources="assets_before_bill_when_dead", + bill="oop_costs", + conditioners=_BILL_CONDITIONERS, + ) else: living_law = assets_and_income.next_assets dead_law = assets_and_income.next_assets_when_dead diff --git a/src/aca_model/baseline/regimes/_nongroup.py b/src/aca_model/baseline/regimes/_nongroup.py index b6a2eb6..a346e0c 100644 --- a/src/aca_model/baseline/regimes/_nongroup.py +++ b/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, diff --git a/src/aca_model/baseline/regimes/_retiree.py b/src/aca_model/baseline/regimes/_retiree.py index f4010e0..a0f304b 100644 --- a/src/aca_model/baseline/regimes/_retiree.py +++ b/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, ) diff --git a/src/aca_model/baseline/regimes/_tied.py b/src/aca_model/baseline/regimes/_tied.py index 96bcd15..9a9c4b3 100644 --- a/src/aca_model/baseline/regimes/_tied.py +++ b/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, ) diff --git a/src/aca_model/config.py b/src/aca_model/config.py index b33e8e2..2031952 100644 --- a/src/aca_model/config.py +++ b/src/aca_model/config.py @@ -38,6 +38,11 @@ class GridConfig: # period computes reads last year's transitory shock, so with the fold it is # drawn on the transition only and no value carries its axis. fold_hcc_transitory: bool = False + # Declare the asset law toward living targets as resources minus the + # out-of-pocket bill, so the solve averages each target's value over the + # transitory shock once per period instead of at every source point. Needs + # `fold_hcc_transitory` and the brute-force asset law. + subtract_hcc_bill: bool = False # The post-decision savings grid is cubically clustered toward the borrowing # constraint. Its node count controls resolution under DC-EGM and NB-EGM. n_savings_gridpoints: int = 200 diff --git a/tests/test_aca_policies.py b/tests/test_aca_policies.py index 1c327d6..f2c9dee 100644 --- a/tests/test_aca_policies.py +++ b/tests/test_aca_policies.py @@ -372,3 +372,34 @@ def test_premium_subsidy_just_below_400_fpl() -> None: premium_credit_schedule=PREMIUM_CREDIT_SCHEDULE, ) assert result > 0.0 + + +@pytest.mark.parametrize( + ("income_fpl_frac", "buy_private", "is_medicaid_eligible"), + [ + (0.5, BuyPrivate.yes, False), + (1.2, BuyPrivate.yes, False), + (1.7, BuyPrivate.yes, False), + (2.3, BuyPrivate.yes, False), + (3.0, BuyPrivate.yes, False), + (1.2, BuyPrivate.no, False), + (1.2, BuyPrivate.yes, True), + ], +) +def test_cost_sharing_scale_of_its_bracket_equals_the_cost_sharing_scale( + income_fpl_frac: float, buy_private: int, is_medicaid_eligible: bool +) -> None: + """The scale read off the bracket is the scale the bracket lookup gives.""" + kwargs = { + "gross_income": jnp.array(income_fpl_frac * FPL_SINGLE), + "spousal_income": jnp.int32(0), + "buy_private": jnp.array(buy_private), + "is_medicaid_eligible": jnp.array(is_medicaid_eligible), + "cost_sharing_schedule": COST_SHARING_SCHEDULE, + } + bracket = aca_hi.cost_sharing_bracket(**kwargs) + scale = aca_hi.cost_sharing_scale_of_bracket( + cost_sharing_bracket=bracket, cost_sharing_schedule=COST_SHARING_SCHEDULE + ) + + assert scale == aca_hi.cost_sharing(**kwargs) diff --git a/tests/test_subtracted_medical_bill.py b/tests/test_subtracted_medical_bill.py new file mode 100644 index 0000000..9a97705 --- /dev/null +++ b/tests/test_subtracted_medical_bill.py @@ -0,0 +1,56 @@ +"""Declaring the asset law as resources minus the medical bill is exact. + +With `subtract_hcc_bill`, the asset law toward every target is +`lcm.SubtractedBill`: end-of-period resources before the out-of-pocket bill, +minus the bill. pylcm then averages each target's value over the transitory +medical-cost shock once per period, on the merged points `{a_j + bill_k}`, +rather than interpolating it at every shock node of every source point. With +linear interpolation in assets that average is exact, so the values agree with +the plain asset law to rounding. +""" + +from dataclasses import replace + +import lcm +import numpy as np +import pytest +from helpers.model import make_baseline_model + +from aca_model.benchmark import get_benchmark_params +from aca_model.config import BENCHMARK_GRID_CONFIG + +_FOLDED = replace(BENCHMARK_GRID_CONFIG, fold_hcc_transitory=True) +_SUBTRACTED = replace(_FOLDED, subtract_hcc_bill=True) +_REGIME = "retiree_nomc_inelig_canwork" + + +@pytest.fixture(scope="module") +def subtracted_model(): + return make_baseline_model(grid_config=_SUBTRACTED) + + +def test_every_target_takes_the_subtracted_bill(subtracted_model): + laws = [ + law + for name, regime in subtracted_model.user_regimes.items() + if name != "dead" + for law in regime.state_transitions["assets"].values() + ] + + assert laws + assert all(isinstance(law, lcm.SubtractedBill) for law in laws) + + +@pytest.mark.long_running +def test_subtracted_bill_value_equals_the_plain_law_value(subtracted_model): + """The two routes sum the same terms in a different order.""" + folded_model = make_baseline_model(grid_config=_FOLDED) + _, _, params = get_benchmark_params(model=folded_model) + folded = np.asarray( + folded_model.solve(params=params, log_level="off").values[0][_REGIME] + ) + subtracted = np.asarray( + subtracted_model.solve(params=params, log_level="off").values[0][_REGIME] + ) + + np.testing.assert_allclose(subtracted, folded, rtol=1e-9, atol=0) From 622ba0038639d1c82b83af21ed9f78dfbf23f51a Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Sat, 10 Oct 2026 16:19:28 +0200 Subject: [PATCH 3/9] ENH Declare the asset law as lcm.AdditiveShockTransition The shock is `negative_oop_costs`, the out-of-pocket bill negated, added to the assets before the bill. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Hs8mfXPYyEa24KaAp6BwCW --- src/aca_model/baseline/health_insurance.py | 5 ++ src/aca_model/baseline/regimes/_common.py | 19 ++++--- tests/test_additive_shock_medical_bill.py | 65 ++++++++++++++++++++++ tests/test_subtracted_medical_bill.py | 56 ------------------- 4 files changed, 80 insertions(+), 65 deletions(-) create mode 100644 tests/test_additive_shock_medical_bill.py delete mode 100644 tests/test_subtracted_medical_bill.py diff --git a/src/aca_model/baseline/health_insurance.py b/src/aca_model/baseline/health_insurance.py index 190e7d2..d073287 100644 --- a/src/aca_model/baseline/health_insurance.py +++ b/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, diff --git a/src/aca_model/baseline/regimes/_common.py b/src/aca_model/baseline/regimes/_common.py index 4f63eab..7a45538 100644 --- a/src/aca_model/baseline/regimes/_common.py +++ b/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, @@ -26,7 +27,6 @@ Regime, RouwenhorstAR1Process, StochasticTransition, - SubtractedBill, categorical, fixed_transition, ) @@ -1007,6 +1007,7 @@ def build_common_functions(spec: RegimeSpec, *, subtract_bill: bool = False) -> 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 ) @@ -1258,8 +1259,8 @@ def _build_per_target_regime_assets( targets use the full `next_assets` with the pension correction. 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`) minus `oop_costs`, - conditioned on every source-side input of 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}." @@ -1270,14 +1271,14 @@ def _build_per_target_regime_assets( 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 = SubtractedBill( - resources="assets_before_bill", - bill="oop_costs", + living_law = AdditiveShockTransition( + base="assets_before_bill", + shock="negative_oop_costs", conditioners=_BILL_CONDITIONERS, ) - dead_law = SubtractedBill( - resources="assets_before_bill_when_dead", - bill="oop_costs", + dead_law = AdditiveShockTransition( + base="assets_before_bill_when_dead", + shock="negative_oop_costs", conditioners=_BILL_CONDITIONERS, ) else: diff --git a/tests/test_additive_shock_medical_bill.py b/tests/test_additive_shock_medical_bill.py new file mode 100644 index 0000000..c7932af --- /dev/null +++ b/tests/test_additive_shock_medical_bill.py @@ -0,0 +1,65 @@ +"""Declaring the asset law as resources plus the negated medical bill is exact. + +With `subtract_hcc_bill`, the asset law toward every target is +`lcm.AdditiveShockTransition`: end-of-period resources before the out-of-pocket +bill, plus the shock `negative_oop_costs`. pylcm then averages each target's value +over the transitory medical-cost shock once per period, on the merged points +`{a_j - shock_k}`, rather than interpolating it at every shock node of every source +point. With linear interpolation in assets that average is exact, so the values +agree with the plain asset law to rounding. +""" + +from dataclasses import replace + +import jax.numpy as jnp +import lcm +import numpy as np +import pytest +from helpers.model import make_baseline_model + +from aca_model.baseline.health_insurance import negative_oop_costs +from aca_model.benchmark import get_benchmark_params +from aca_model.config import BENCHMARK_GRID_CONFIG + +_FOLDED = replace(BENCHMARK_GRID_CONFIG, fold_hcc_transitory=True) +_ADDITIVE = replace(_FOLDED, subtract_hcc_bill=True) +_REGIME = "retiree_nomc_inelig_canwork" + + +@pytest.fixture(scope="module") +def additive_model(): + return make_baseline_model(grid_config=_ADDITIVE) + + +def test_negative_oop_costs_is_the_negated_bill(): + np.testing.assert_array_equal( + negative_oop_costs(oop_costs=jnp.array([0.0, 1.5, 250.0])), + np.array([0.0, -1.5, -250.0]), + ) + + +def test_every_target_adds_the_negated_bill(additive_model): + laws = [ + (type(law), law.shock) + for name, regime in additive_model.user_regimes.items() + if name != "dead" + for law in regime.state_transitions["assets"].values() + ] + + assert laws + assert set(laws) == {(lcm.AdditiveShockTransition, "negative_oop_costs")} + + +@pytest.mark.long_running +def test_additive_shock_value_equals_the_plain_law_value(additive_model): + """The two routes sum the same terms in a different order.""" + folded_model = make_baseline_model(grid_config=_FOLDED) + _, _, params = get_benchmark_params(model=folded_model) + folded = np.asarray( + folded_model.solve(params=params, log_level="off").values[0][_REGIME] + ) + additive = np.asarray( + additive_model.solve(params=params, log_level="off").values[0][_REGIME] + ) + + np.testing.assert_allclose(additive, folded, rtol=1e-9, atol=0) diff --git a/tests/test_subtracted_medical_bill.py b/tests/test_subtracted_medical_bill.py deleted file mode 100644 index 9a97705..0000000 --- a/tests/test_subtracted_medical_bill.py +++ /dev/null @@ -1,56 +0,0 @@ -"""Declaring the asset law as resources minus the medical bill is exact. - -With `subtract_hcc_bill`, the asset law toward every target is -`lcm.SubtractedBill`: end-of-period resources before the out-of-pocket bill, -minus the bill. pylcm then averages each target's value over the transitory -medical-cost shock once per period, on the merged points `{a_j + bill_k}`, -rather than interpolating it at every shock node of every source point. With -linear interpolation in assets that average is exact, so the values agree with -the plain asset law to rounding. -""" - -from dataclasses import replace - -import lcm -import numpy as np -import pytest -from helpers.model import make_baseline_model - -from aca_model.benchmark import get_benchmark_params -from aca_model.config import BENCHMARK_GRID_CONFIG - -_FOLDED = replace(BENCHMARK_GRID_CONFIG, fold_hcc_transitory=True) -_SUBTRACTED = replace(_FOLDED, subtract_hcc_bill=True) -_REGIME = "retiree_nomc_inelig_canwork" - - -@pytest.fixture(scope="module") -def subtracted_model(): - return make_baseline_model(grid_config=_SUBTRACTED) - - -def test_every_target_takes_the_subtracted_bill(subtracted_model): - laws = [ - law - for name, regime in subtracted_model.user_regimes.items() - if name != "dead" - for law in regime.state_transitions["assets"].values() - ] - - assert laws - assert all(isinstance(law, lcm.SubtractedBill) for law in laws) - - -@pytest.mark.long_running -def test_subtracted_bill_value_equals_the_plain_law_value(subtracted_model): - """The two routes sum the same terms in a different order.""" - folded_model = make_baseline_model(grid_config=_FOLDED) - _, _, params = get_benchmark_params(model=folded_model) - folded = np.asarray( - folded_model.solve(params=params, log_level="off").values[0][_REGIME] - ) - subtracted = np.asarray( - subtracted_model.solve(params=params, log_level="off").values[0][_REGIME] - ) - - np.testing.assert_allclose(subtracted, folded, rtol=1e-9, atol=0) From 5b75bee408cf5dca40da5ba5b087057f9e15524e Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Sun, 11 Oct 2026 06:53:08 +0200 Subject: [PATCH 4/9] Enable medical-shock folding in the ACA benchmark --- .github/workflows/main.yml | 2 +- README.md | 8 +++ src/aca_model/baseline/regimes/_common.py | 2 +- src/aca_model/benchmark.py | 25 ++++++-- src/aca_model/config.py | 2 + tests/test_additive_shock_medical_bill.py | 58 ++++++++++++++----- tests/test_benchmark_folding.py | 64 +++++++++++++++++++++ tests/test_transition_only_medical_shock.py | 17 ++++-- 8 files changed, 151 insertions(+), 27 deletions(-) create mode 100644 tests/test_benchmark_folding.py diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 1624b74..3d2f21d 100644 --- a/.github/workflows/main.yml +++ b/.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 diff --git a/README.md b/README.md index 9edf892..f22ff8f 100644 --- a/README.md +++ b/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 diff --git a/src/aca_model/baseline/regimes/_common.py b/src/aca_model/baseline/regimes/_common.py index 7a45538..2032ce1 100644 --- a/src/aca_model/baseline/regimes/_common.py +++ b/src/aca_model/baseline/regimes/_common.py @@ -1239,7 +1239,7 @@ def _select_aime_law(spec: RegimeSpec) -> Callable[..., FloatND]: "is_married": (0, 1), "good_health": (0, 1), "is_medicaid_eligible": (False, True), - "buy_private": (BuyPrivate.no, BuyPrivate.yes), + "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)), diff --git a/src/aca_model/benchmark.py b/src/aca_model/benchmark.py index 093e65e..3eed460 100644 --- a/src/aca_model/benchmark.py +++ b/src/aca_model/benchmark.py @@ -37,7 +37,7 @@ from aca_model.baseline.health_insurance import HealthInsuranceState from aca_model.baseline.model import create_model from aca_model.baseline.regimes import ENTRY_REGIMES -from aca_model.config import BENCHMARK_GRID_CONFIG +from aca_model.config import BENCHMARK_GRID_CONFIG, GridConfig from aca_model.consumption_dollars_grid import ( compute_consumption_dollars_points, inject_consumption_dollars_points, @@ -66,6 +66,7 @@ def create_benchmark_model( *, pref_type_grid: DiscreteGrid, + grid_config: GridConfig = BENCHMARK_GRID_CONFIG, execution_config: ExecutionConfig | None = None, ) -> Model: """Create the aca baseline with `BENCHMARK_GRID_CONFIG` and frozen fixed_params. @@ -75,13 +76,14 @@ def create_benchmark_model( Args: pref_type_grid: Pref-type grid; pass `DiscreteGrid(BenchmarkPrefType)`. + grid_config: Economic resolution and medical-shock folding declarations. execution_config: Explicit hardware-local policy forwarded unchanged. None uses the smallest selected accelerator allocator limit as the device-memory budget; CPU construction remains unbudgeted. """ fixed_params, wage_params, _ = get_benchmark_params(model=None) return create_model( - grid_config=BENCHMARK_GRID_CONFIG, + grid_config=grid_config, fixed_params=fixed_params, wage_params=wage_params, derived_categoricals=_DERIVED_CATEGORICALS, @@ -151,7 +153,11 @@ def get_benchmark_consumption_dollars_points(*, n_points: int) -> tuple[float, . def get_benchmark_initial_conditions( - *, model: Model, n_subjects: int, seed: int + *, + model: Model, + n_subjects: int, + seed: int, + grid_config: GridConfig = BENCHMARK_GRID_CONFIG, ) -> dict[str, Array]: """Draw random feasible initial conditions across five age-51 regimes. @@ -177,7 +183,13 @@ def get_benchmark_initial_conditions( assets_lo, assets_hi = float(assets_pts.min()), float(assets_pts.max()) aime_lo, aime_hi = float(aime_pts.min()), float(aime_pts.max()) hcc_p_pts = np.asarray(grids["hcc_persistent"].to_jax()) # ty: ignore[unresolved-attribute] - hcc_t_pts = np.asarray(grids["hcc_transitory"].to_jax()) # ty: ignore[unresolved-attribute] + hcc_t_grid = grids.get("hcc_transitory") + # Preserve the random stream for shared states when the folded state is absent. + hcc_t_pts = ( + np.asarray(hcc_t_grid.to_jax()) # ty: ignore[unresolved-attribute] + if hcc_t_grid is not None + else np.arange(grid_config.n_hcc_transitory_gridpoints) + ) wage_res_pts = np.asarray(grids["log_ft_wage_res"].to_jax()) # ty: ignore[unresolved-attribute] # For lagged_labor_supply: 0 for tied rows (regime lacks the state), @@ -187,7 +199,7 @@ def get_benchmark_initial_conditions( np.int32 ) - return { + initial = { "regime_id": jnp.asarray(regime), "age": jnp.full(n_subjects, 51.0), "assets": jnp.asarray(rng.uniform(assets_lo, assets_hi, n_subjects)), @@ -208,3 +220,6 @@ def get_benchmark_initial_conditions( # everyone starts as not-yet-claimed. "claimed_ss": jnp.zeros(n_subjects, dtype=jnp.int32), } + if "hcc_transitory" not in model.state_names(regime_name=ENTRY_REGIMES[0]): + initial.pop("hcc_transitory") + return initial diff --git a/src/aca_model/config.py b/src/aca_model/config.py index 2031952..d95569b 100644 --- a/src/aca_model/config.py +++ b/src/aca_model/config.py @@ -67,6 +67,8 @@ class GridConfig: GRID_CONFIG = GridConfig() BENCHMARK_GRID_CONFIG = GridConfig( + fold_hcc_transitory=True, + subtract_hcc_bill=True, n_assets_gridpoints=3, n_aime_gridpoints=3, n_consumption_dollars_gridpoints=5, diff --git a/tests/test_additive_shock_medical_bill.py b/tests/test_additive_shock_medical_bill.py index c7932af..9017edb 100644 --- a/tests/test_additive_shock_medical_bill.py +++ b/tests/test_additive_shock_medical_bill.py @@ -9,6 +9,7 @@ agree with the plain asset law to rounding. """ +from collections.abc import Mapping from dataclasses import replace import jax.numpy as jnp @@ -17,41 +18,46 @@ import pytest from helpers.model import make_baseline_model +from aca_model.agent import assets_and_income from aca_model.baseline.health_insurance import negative_oop_costs from aca_model.benchmark import get_benchmark_params from aca_model.config import BENCHMARK_GRID_CONFIG -_FOLDED = replace(BENCHMARK_GRID_CONFIG, fold_hcc_transitory=True) +_FOLDED = replace( + BENCHMARK_GRID_CONFIG, fold_hcc_transitory=True, subtract_hcc_bill=False +) _ADDITIVE = replace(_FOLDED, subtract_hcc_bill=True) _REGIME = "retiree_nomc_inelig_canwork" @pytest.fixture(scope="module") -def additive_model(): +def additive_model() -> lcm.Model: return make_baseline_model(grid_config=_ADDITIVE) -def test_negative_oop_costs_is_the_negated_bill(): +def test_negative_oop_costs_is_the_negated_bill() -> None: np.testing.assert_array_equal( negative_oop_costs(oop_costs=jnp.array([0.0, 1.5, 250.0])), np.array([0.0, -1.5, -250.0]), ) -def test_every_target_adds_the_negated_bill(additive_model): - laws = [ - (type(law), law.shock) - for name, regime in additive_model.user_regimes.items() - if name != "dead" - for law in regime.state_transitions["assets"].values() - ] - - assert laws - assert set(laws) == {(lcm.AdditiveShockTransition, "negative_oop_costs")} +def test_every_target_adds_the_negated_bill(additive_model: lcm.Model) -> None: + for name, regime in additive_model.user_regimes.items(): + if name == "dead": + continue + laws = regime.state_transitions["assets"] + assert isinstance(laws, Mapping) + assert laws + for law in laws.values(): + assert isinstance(law, lcm.AdditiveShockTransition) + assert law.shock == "negative_oop_costs" @pytest.mark.long_running -def test_additive_shock_value_equals_the_plain_law_value(additive_model): +def test_additive_shock_value_equals_the_plain_law_value( + additive_model: lcm.Model, +) -> None: """The two routes sum the same terms in a different order.""" folded_model = make_baseline_model(grid_config=_FOLDED) _, _, params = get_benchmark_params(model=folded_model) @@ -63,3 +69,27 @@ def test_additive_shock_value_equals_the_plain_law_value(additive_model): ) np.testing.assert_allclose(additive, folded, rtol=1e-9, atol=0) + + +@pytest.mark.parametrize( + ("bill", "expected_living", "expected_dead"), + [(0.0, 85.0, 75.0), (20.0, 65.0, 55.0), (250.0, -165.0, -175.0)], +) +def test_asset_law_preserves_resources_minus_medical_costs( + bill: float, expected_living: float, expected_dead: float +) -> None: + common = { + "cash_on_hand": jnp.array(100.0), + "transfers": jnp.array(5.0), + "consumption_dollars": jnp.array(30.0), + } + shock = negative_oop_costs(oop_costs=jnp.array(bill)) + living = ( + assets_and_income.assets_before_bill( + **common, pension_assets_adjustment=jnp.array(10.0) + ) + + shock + ) + dead = assets_and_income.assets_before_bill_when_dead(**common) + shock + np.testing.assert_allclose(living, expected_living, rtol=0, atol=0) + np.testing.assert_allclose(dead, expected_dead, rtol=0, atol=0) diff --git a/tests/test_benchmark_folding.py b/tests/test_benchmark_folding.py new file mode 100644 index 0000000..71c096f --- /dev/null +++ b/tests/test_benchmark_folding.py @@ -0,0 +1,64 @@ +"""The ACA benchmark activates transitory-shock folding explicitly.""" + +from collections.abc import Mapping +from dataclasses import replace + +import lcm +import numpy as np +import pytest +from lcm import DiscreteGrid + +from aca_model.agent.preferences import BenchmarkPrefType +from aca_model.benchmark import create_benchmark_model, get_benchmark_initial_conditions +from aca_model.config import BENCHMARK_GRID_CONFIG + + +@pytest.fixture(scope="module") +def benchmark_model() -> lcm.Model: + return create_benchmark_model(pref_type_grid=DiscreteGrid(BenchmarkPrefType)) + + +def test_benchmark_folds_the_transitory_medical_bill( + benchmark_model: lcm.Model, +) -> None: + model = benchmark_model + living = { + name: regime for name, regime in model.user_regimes.items() if name != "dead" + } + assert living + assert all( + "hcc_transitory" not in model.state_names(regime_name=name) for name in living + ) + for regime in living.values(): + laws = regime.state_transitions["assets"] + assert isinstance(laws, Mapping) + assert all( + isinstance(law, lcm.AdditiveShockTransition) for law in laws.values() + ) + initial = get_benchmark_initial_conditions(model=model, n_subjects=8, seed=42) + assert "hcc_transitory" not in initial + + +def test_unfolded_control_has_identical_shared_initial_conditions( + benchmark_model: lcm.Model, +) -> None: + control_grid = replace( + BENCHMARK_GRID_CONFIG, fold_hcc_transitory=False, subtract_hcc_bill=False + ) + control = create_benchmark_model( + pref_type_grid=DiscreteGrid(BenchmarkPrefType), grid_config=control_grid + ) + assert all( + "hcc_transitory" in control.state_names(regime_name=name) + for name in control.user_regimes + if name != "dead" + ) + folded_initial = get_benchmark_initial_conditions( + model=benchmark_model, n_subjects=8, seed=42 + ) + control_initial = get_benchmark_initial_conditions( + model=control, n_subjects=8, seed=42, grid_config=control_grid + ) + assert set(control_initial) == set(folded_initial) | {"hcc_transitory"} + for name, values in folded_initial.items(): + np.testing.assert_array_equal(values, control_initial[name]) diff --git a/tests/test_transition_only_medical_shock.py b/tests/test_transition_only_medical_shock.py index 4cb07f1..b2a1f42 100644 --- a/tests/test_transition_only_medical_shock.py +++ b/tests/test_transition_only_medical_shock.py @@ -13,6 +13,7 @@ import numpy as np import pytest from helpers.model import make_baseline_model +from lcm import Model from aca_model.benchmark import get_benchmark_params from aca_model.config import BENCHMARK_GRID_CONFIG @@ -22,18 +23,18 @@ @pytest.fixture(scope="module") -def folded_model(): +def folded_model() -> Model: return make_baseline_model(grid_config=_FOLDED) -def test_living_regimes_carry_no_lagged_transitory_shock(folded_model): +def test_living_regimes_carry_no_lagged_transitory_shock(folded_model: Model) -> None: assert not any( "hcc_transitory" in folded_model.state_names(regime_name=name) for name in folded_model.user_regimes ) -def test_living_regimes_keep_the_lagged_persistent_shock(folded_model): +def test_living_regimes_keep_the_lagged_persistent_shock(folded_model: Model) -> None: assert all( "hcc_persistent" in folded_model.state_names(regime_name=name) for name in folded_model.user_regimes @@ -43,10 +44,14 @@ def test_living_regimes_keep_the_lagged_persistent_shock(folded_model): @pytest.mark.long_running def test_folded_value_equals_every_transitory_slice_of_the_unfolded_value( - folded_model, -): + folded_model: Model, +) -> None: """The two models' sums differ only in order, so they agree to rounding.""" - unfolded_model = make_baseline_model() + unfolded_model = make_baseline_model( + grid_config=replace( + BENCHMARK_GRID_CONFIG, fold_hcc_transitory=False, subtract_hcc_bill=False + ) + ) _, _, params = get_benchmark_params(model=unfolded_model) unfolded = np.asarray( unfolded_model.solve(params=params, log_level="off").values[0][_REGIME] From ee48d9b5ed6739d054291a9a2e8734c4a1f9a2fd Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Sun, 11 Oct 2026 06:54:29 +0200 Subject: [PATCH 5/9] Preserve unfolded controls in legacy model test fixtures --- tests/helpers/model.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/tests/helpers/model.py b/tests/helpers/model.py index df958e8..3645ab7 100644 --- a/tests/helpers/model.py +++ b/tests/helpers/model.py @@ -5,6 +5,8 @@ Production callers (aca-slurm, scripts) assemble these explicitly. """ +from dataclasses import replace + from lcm import DiscreteGrid, Model from aca_model.aca.health_insurance import PolicyVariant @@ -32,8 +34,12 @@ "pref_type": DiscreteGrid(BenchmarkPrefType), } +_UNFOLDED_GRID_CONFIG = replace( + BENCHMARK_GRID_CONFIG, fold_hcc_transitory=False, subtract_hcc_bill=False +) + -def make_baseline_model(*, grid_config: GridConfig = BENCHMARK_GRID_CONFIG) -> Model: +def make_baseline_model(*, grid_config: GridConfig = _UNFOLDED_GRID_CONFIG) -> Model: """Baseline model on `grid_config` with the benchmark snapshot params.""" fixed_params, wage_params, _ = get_benchmark_params(model=None) return _create_baseline_model( @@ -65,14 +71,16 @@ def aca_fixed_params(*, fixed_params: dict, policy: PolicyVariant) -> dict: } -def make_aca_model(*, policy: PolicyVariant) -> Model: - """ACA model on `BENCHMARK_GRID_CONFIG` with the benchmark snapshot params.""" +def make_aca_model( + *, policy: PolicyVariant, grid_config: GridConfig = _UNFOLDED_GRID_CONFIG +) -> Model: + """ACA model on `grid_config` with the benchmark snapshot params.""" fixed_params, wage_params, _ = get_benchmark_params(model=None) return _create_aca_model( policy=policy, fixed_params=aca_fixed_params(fixed_params=fixed_params, policy=policy), wage_params=wage_params, derived_categoricals=_DERIVED_CATEGORICALS, - grid_config=BENCHMARK_GRID_CONFIG, + grid_config=grid_config, pref_type_grid=DiscreteGrid(BenchmarkPrefType), ) From b962ebd27f01f03e7ec6fe5888fdca4feae0dd50 Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Sun, 11 Oct 2026 06:56:47 +0200 Subject: [PATCH 6/9] Check folded values across every solved model cell --- tests/test_additive_shock_medical_bill.py | 26 ++++++++++----- tests/test_transition_only_medical_shock.py | 37 +++++++++++++-------- 2 files changed, 41 insertions(+), 22 deletions(-) diff --git a/tests/test_additive_shock_medical_bill.py b/tests/test_additive_shock_medical_bill.py index 9017edb..9fff9d0 100644 --- a/tests/test_additive_shock_medical_bill.py +++ b/tests/test_additive_shock_medical_bill.py @@ -27,7 +27,6 @@ BENCHMARK_GRID_CONFIG, fold_hcc_transitory=True, subtract_hcc_bill=False ) _ADDITIVE = replace(_FOLDED, subtract_hcc_bill=True) -_REGIME = "retiree_nomc_inelig_canwork" @pytest.fixture(scope="module") @@ -61,14 +60,23 @@ def test_additive_shock_value_equals_the_plain_law_value( """The two routes sum the same terms in a different order.""" folded_model = make_baseline_model(grid_config=_FOLDED) _, _, params = get_benchmark_params(model=folded_model) - folded = np.asarray( - folded_model.solve(params=params, log_level="off").values[0][_REGIME] - ) - additive = np.asarray( - additive_model.solve(params=params, log_level="off").values[0][_REGIME] - ) - - np.testing.assert_allclose(additive, folded, rtol=1e-9, atol=0) + folded_values = folded_model.solve(params=params, log_level="off").values + additive_values = additive_model.solve(params=params, log_level="off").values + assert set(additive_values) == set(folded_values) + for period, regimes in folded_values.items(): + assert set(additive_values[period]) == set(regimes) + for name, values in regimes.items(): + np.testing.assert_array_equal( + np.isfinite(additive_values[period][name]), np.isfinite(values) + ) + np.testing.assert_allclose( + additive_values[period][name], + values, + rtol=1e-9, + atol=0, + equal_nan=False, + err_msg=f"period={period}, regime={name}", + ) @pytest.mark.parametrize( diff --git a/tests/test_transition_only_medical_shock.py b/tests/test_transition_only_medical_shock.py index b2a1f42..a112d16 100644 --- a/tests/test_transition_only_medical_shock.py +++ b/tests/test_transition_only_medical_shock.py @@ -19,7 +19,6 @@ from aca_model.config import BENCHMARK_GRID_CONFIG _FOLDED = replace(BENCHMARK_GRID_CONFIG, fold_hcc_transitory=True) -_REGIME = "retiree_nomc_inelig_canwork" @pytest.fixture(scope="module") @@ -53,15 +52,27 @@ def test_folded_value_equals_every_transitory_slice_of_the_unfolded_value( ) ) _, _, params = get_benchmark_params(model=unfolded_model) - unfolded = np.asarray( - unfolded_model.solve(params=params, log_level="off").values[0][_REGIME] - ) - folded = np.asarray( - folded_model.solve(params=params, log_level="off").values[0][_REGIME] - ) - axis = unfolded_model.state_names(regime_name=_REGIME).index("hcc_transitory") - by_slice = np.moveaxis(unfolded, axis, 0) - - np.testing.assert_allclose( - by_slice, np.broadcast_to(folded, by_slice.shape), rtol=1e-6, atol=0 - ) + unfolded_values = unfolded_model.solve(params=params, log_level="off").values + folded_values = folded_model.solve(params=params, log_level="off").values + assert set(folded_values) == set(unfolded_values) + for period, regimes in unfolded_values.items(): + assert set(folded_values[period]) == set(regimes) + for name, values in regimes.items(): + unfolded = np.asarray(values) + folded = np.asarray(folded_values[period][name]) + states = unfolded_model.state_names(regime_name=name) + by_slice = ( + np.moveaxis(unfolded, states.index("hcc_transitory"), 0) + if "hcc_transitory" in states + else unfolded + ) + expected = np.broadcast_to(folded, by_slice.shape) + np.testing.assert_array_equal(np.isfinite(by_slice), np.isfinite(expected)) + np.testing.assert_allclose( + by_slice, + expected, + rtol=1e-6, + atol=0, + equal_nan=False, + err_msg=f"period={period}, regime={name}", + ) From 89d26dda693c92cfda02e2352fa914c174175fec Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Sun, 11 Oct 2026 06:58:42 +0200 Subject: [PATCH 7/9] State exact ACA cost-sharing parity with a numeric assertion --- tests/test_aca_policies.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_aca_policies.py b/tests/test_aca_policies.py index f2c9dee..b16fda4 100644 --- a/tests/test_aca_policies.py +++ b/tests/test_aca_policies.py @@ -3,6 +3,7 @@ import inspect import jax.numpy as jnp +import numpy as np import pytest from lcm.params import MappingLeaf @@ -402,4 +403,4 @@ def test_cost_sharing_scale_of_its_bracket_equals_the_cost_sharing_scale( cost_sharing_bracket=bracket, cost_sharing_schedule=COST_SHARING_SCHEDULE ) - assert scale == aca_hi.cost_sharing(**kwargs) + np.testing.assert_allclose(scale, aca_hi.cost_sharing(**kwargs), rtol=0, atol=0) From 54397e8dc6df4c3de53c4470e92f4754ce9595dc Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Sun, 11 Oct 2026 07:06:38 +0200 Subject: [PATCH 8/9] Use the unfolded control for lagged medical-shock twins --- tests/test_medical_cost_timing.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/tests/test_medical_cost_timing.py b/tests/test_medical_cost_timing.py index 0705808..2ebbc8d 100644 --- a/tests/test_medical_cost_timing.py +++ b/tests/test_medical_cost_timing.py @@ -8,6 +8,8 @@ year's costs, and not at all on the lagged transitory shock. """ +from dataclasses import replace + import jax.numpy as jnp import numpy as np import pandas as pd @@ -20,6 +22,7 @@ get_benchmark_initial_conditions, get_benchmark_params, ) +from aca_model.config import BENCHMARK_GRID_CONFIG _SHOCKS = ("hcc_persistent", "hcc_transitory") @@ -27,9 +30,16 @@ @pytest.fixture(scope="module") def period_0_by_shock() -> dict[tuple[str, str], pd.Series]: """Period-0 rows of twins that differ in one lagged shock, at its grid ends.""" - model = create_benchmark_model(pref_type_grid=DiscreteGrid(BenchmarkPrefType)) + grid_config = replace( + BENCHMARK_GRID_CONFIG, fold_hcc_transitory=False, subtract_hcc_bill=False + ) + model = create_benchmark_model( + pref_type_grid=DiscreteGrid(BenchmarkPrefType), grid_config=grid_config + ) _, _, params = get_benchmark_params(model=model) - base = get_benchmark_initial_conditions(model=model, n_subjects=1, seed=0) + base = get_benchmark_initial_conditions( + model=model, n_subjects=1, seed=0, grid_config=grid_config + ) grids = model.user_regimes["retiree_nomc_inelig_canwork"].states initial_conditions = {name: jnp.tile(values, 4) for name, values in base.items()} for row, shock in enumerate(_SHOCKS): From 5e69b66e20a76bd715151697b6c1e07b34c461d3 Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Sun, 11 Oct 2026 07:06:54 +0200 Subject: [PATCH 9/9] Keep the transition-only fold witness separate from additive folding --- tests/test_transition_only_medical_shock.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/test_transition_only_medical_shock.py b/tests/test_transition_only_medical_shock.py index a112d16..a4a148f 100644 --- a/tests/test_transition_only_medical_shock.py +++ b/tests/test_transition_only_medical_shock.py @@ -18,7 +18,9 @@ from aca_model.benchmark import get_benchmark_params from aca_model.config import BENCHMARK_GRID_CONFIG -_FOLDED = replace(BENCHMARK_GRID_CONFIG, fold_hcc_transitory=True) +_FOLDED = replace( + BENCHMARK_GRID_CONFIG, fold_hcc_transitory=True, subtract_hcc_bill=False +) @pytest.fixture(scope="module")