Test against pylcm #630

.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@7e99a19badafe20e29a517f69d9ec155bba8e0a0"
+          git+https://github.com/OpenSourceEconomics/pylcm.git@1e1a36621e7d82a4ce0530491797118ae03effa5"
       - name: Install aca-model with test deps
         run: pip install -e . pytest pdbp
       - name: Run pytest

README.md

@@ -2,13 +2,12 @@
 
 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.
+The IID transitory medical-cost shock is read only as its draw `next_hcc_transitory`, so
+pylcm integrates it out and no regime carries it as a state. The benchmark declares the
+asset transition as resources plus the negated medical bill; to measure the plain asset
+law instead, pass `grid_config=replace(BENCHMARK_GRID_CONFIG, subtract_hcc_bill=False)`
+to `create_benchmark_model`, and the same grid configuration to
+`get_benchmark_initial_conditions`.
 
 Model factories accept `execution_config=ExecutionConfig(...)`. An omitted policy uses
 the smallest JAX allocator `bytes_limit` among the selected accelerators as an explicit

src/aca_model/baseline/regimes/_common.py

@@ -326,7 +326,6 @@ 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(

src/aca_model/benchmark.py

@@ -77,7 +77,7 @@ def create_benchmark_model(
 
     Args:
         pref_type_grid: Pref-type grid; pass `DiscreteGrid(BenchmarkPrefType)`.
-        grid_config: Economic resolution and medical-shock folding declarations.
+        grid_config: Economic resolution and numerical choices.
         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.
@@ -188,7 +188,7 @@ def get_benchmark_initial_conditions(
     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_grid = grids.get("hcc_transitory")
-    # Preserve the random stream for shared states when the folded state is absent.
+    # Preserve the random stream for shared states; the transitory shock is no state.
     hcc_t_pts = (
         np.asarray(hcc_t_grid.to_jax())  # ty: ignore[unresolved-attribute]
         if hcc_t_grid is not None

src/aca_model/config.py

@@ -34,14 +34,12 @@ 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
     # 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.
+    # the brute-force asset law. The transitory shock is read only as its draw
+    # `next_hcc_transitory`, so pylcm integrates it out and no value carries its
+    # axis, with or without this option.
     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.
@@ -67,7 +65,6 @@ 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,

tests/helpers/model.py

@@ -34,12 +34,10 @@
     "pref_type": DiscreteGrid(BenchmarkPrefType),
 }
 
-_UNFOLDED_GRID_CONFIG = replace(
-    BENCHMARK_GRID_CONFIG, fold_hcc_transitory=False, subtract_hcc_bill=False
-)
+_PLAIN_LAW_GRID_CONFIG = replace(BENCHMARK_GRID_CONFIG, subtract_hcc_bill=False)
 
 
-def make_baseline_model(*, grid_config: GridConfig = _UNFOLDED_GRID_CONFIG) -> Model:
+def make_baseline_model(*, grid_config: GridConfig = _PLAIN_LAW_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(
@@ -72,7 +70,7 @@ def aca_fixed_params(*, fixed_params: dict, policy: PolicyVariant) -> dict:
 
 
 def make_aca_model(
-    *, policy: PolicyVariant, grid_config: GridConfig = _UNFOLDED_GRID_CONFIG
+    *, policy: PolicyVariant, grid_config: GridConfig = _PLAIN_LAW_GRID_CONFIG
 ) -> Model:
     """ACA model on `grid_config` with the benchmark snapshot params."""
     fixed_params, wage_params, _ = get_benchmark_params(model=None)

tests/test_additive_shock_medical_bill.py

@@ -23,10 +23,8 @@
 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, subtract_hcc_bill=False
-)
-_ADDITIVE = replace(_FOLDED, subtract_hcc_bill=True)
+_PLAIN_LAW = replace(BENCHMARK_GRID_CONFIG, subtract_hcc_bill=False)
+_ADDITIVE = replace(_PLAIN_LAW, subtract_hcc_bill=True)
 
 
 @pytest.fixture(scope="module")
@@ -58,12 +56,12 @@ 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)
-    folded_values = folded_model.solve(params=params, log_level="off").values
+    plain_model = make_baseline_model(grid_config=_PLAIN_LAW)
+    _, _, params = get_benchmark_params(model=plain_model)
+    plain_values = plain_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) == set(plain_values)
+    for period, regimes in plain_values.items():
         assert set(additive_values[period]) == set(regimes)
         for name, values in regimes.items():
             np.testing.assert_array_equal(

tests/test_benchmark_folding.py

@@ -1,16 +1,13 @@
-"""The ACA benchmark activates transitory-shock folding explicitly."""
+"""The ACA benchmark integrates the transitory medical bill out of every state."""
 
 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")
@@ -37,28 +34,3 @@ def test_benchmark_folds_the_transitory_medical_bill(
         )
     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])

tests/test_dag_coverage.py

@@ -134,10 +134,13 @@ def test_select_admissible_starts_drops_starts_at_61_and_above(model):
 
 
 def test_medical_cost_draws_keep_living_states_without_choice_constraints(model):
-    """Living regimes carry medical shocks; bequests depend only on assets and type."""
+    """Living regimes carry the persistent medical shock as a state; the transitory
+    one is read only as its next-period draw. Bequests depend only on assets and type.
+    """
     for name, regime in model.user_regimes.items():
         assert "medical_cost_shocks_carried" not in regime.constraints
         if name == "dead":
             assert set(regime.states) == {"assets", "pref_type"}
         else:
-            assert {"hcc_persistent", "hcc_transitory"} <= set(regime.states)
+            assert "hcc_persistent" in regime.states
+            assert "hcc_transitory" not in regime.states

tests/test_dcegm_model_creation.py

@@ -33,12 +33,9 @@
 
 _FIXED_PARAMS, _WAGE_PARAMS, _ = get_benchmark_params(model=None)
 
-# pylcm folds a shock state only under grid search, and the EGM solvers take the
-# savings-form asset law, so the transitory medical-cost shock stays a state and
-# the bill stays inside `next_assets`.
-_EGM_GRID_CONFIG = dataclasses.replace(
-    BENCHMARK_GRID_CONFIG, fold_hcc_transitory=False, subtract_hcc_bill=False
-)
+# The EGM solvers take the savings-form asset law, so the bill stays inside
+# `next_assets`; the transitory medical-cost shock is read only as its draw.
+_EGM_GRID_CONFIG = dataclasses.replace(BENCHMARK_GRID_CONFIG, subtract_hcc_bill=False)
 
 
 def _build_regimes(solver: SolverName) -> dict[str, Regime]:

tests/test_medical_cost_timing.py

@@ -24,15 +24,13 @@
 )
 from aca_model.config import BENCHMARK_GRID_CONFIG
 
-_SHOCKS = ("hcc_persistent", "hcc_transitory")
+_SHOCKS = ("hcc_persistent",)
 
 
 @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."""
-    grid_config = replace(
-        BENCHMARK_GRID_CONFIG, fold_hcc_transitory=False, subtract_hcc_bill=False
-    )
+    grid_config = replace(BENCHMARK_GRID_CONFIG, subtract_hcc_bill=False)
     model = create_benchmark_model(
         pref_type_grid=DiscreteGrid(BenchmarkPrefType), grid_config=grid_config
     )
@@ -59,22 +57,13 @@ def period_0_by_shock() -> dict[tuple[str, str], pd.Series]:
     }
 
 
-def test_value_does_not_depend_on_lagged_transitory_shock(
-    period_0_by_shock: dict[tuple[str, str], pd.Series],
-) -> None:
-    """Twins differing only in last year's transitory shock have equal value."""
-    low = period_0_by_shock["hcc_transitory", "low"]["value"]
-    high = period_0_by_shock["hcc_transitory", "high"]["value"]
-    assert low == high
-
-
-def test_consumption_does_not_depend_on_lagged_transitory_shock(
-    period_0_by_shock: dict[tuple[str, str], pd.Series],
-) -> None:
-    """Twins differing only in last year's transitory shock consume the same."""
-    low = period_0_by_shock["hcc_transitory", "low"]["consumption_dollars"]
-    high = period_0_by_shock["hcc_transitory", "high"]["consumption_dollars"]
-    assert low == high
+def test_no_regime_carries_a_lagged_transitory_shock() -> None:
+    """Last year's transitory shock is read by nothing, so it is no state."""
+    model = create_benchmark_model(pref_type_grid=DiscreteGrid(BenchmarkPrefType))
+    assert not any(
+        "hcc_transitory" in model.state_names(regime_name=name)
+        for name in model.user_regimes
+    )
 
 
 def test_value_falls_with_lagged_persistent_shock(

tests/test_nbegm_model_creation.py

@@ -42,12 +42,9 @@
 
 _FIXED_PARAMS, _WAGE_PARAMS, _ = get_benchmark_params(model=None)
 
-# pylcm folds a shock state only under grid search, and the EGM solvers take the
-# savings-form asset law, so the transitory medical-cost shock stays a state and
-# the bill stays inside `next_assets`.
-_EGM_GRID_CONFIG = dataclasses.replace(
-    BENCHMARK_GRID_CONFIG, fold_hcc_transitory=False, subtract_hcc_bill=False
-)
+# The EGM solvers take the savings-form asset law, so the bill stays inside
+# `next_assets`; the transitory medical-cost shock is read only as its draw.
+_EGM_GRID_CONFIG = dataclasses.replace(BENCHMARK_GRID_CONFIG, subtract_hcc_bill=False)
 
 _M1_REGIME = "nongroup_nomc_inelig_canwork"
 _BRUTE_REGIME = "retiree_nomc_inelig_canwork"

tests/test_transition_only_medical_shock.py

@@ -1,80 +1,38 @@
-"""Folding the transitory medical-cost shock takes it out of every state space.
+"""The transitory medical-cost shock is in no 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.
+enters the out-of-pocket bill inside the end-of-period asset law as
+`next_hcc_transitory`, and the draw is IID. pylcm therefore draws it on the
+transition only and integrates it in the continuation. Living regimes keep the
+lagged persistent shock, which prices the private premium.
 """
 
 from dataclasses import replace
 
-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
 
-_FOLDED = replace(
-    BENCHMARK_GRID_CONFIG, fold_hcc_transitory=True, subtract_hcc_bill=False
-)
+_GRID_CONFIG = replace(BENCHMARK_GRID_CONFIG, subtract_hcc_bill=False)
 
 
 @pytest.fixture(scope="module")
-def folded_model() -> Model:
-    return make_baseline_model(grid_config=_FOLDED)
+def model() -> Model:
+    return make_baseline_model(grid_config=_GRID_CONFIG)
 
 
-def test_living_regimes_carry_no_lagged_transitory_shock(folded_model: Model) -> None:
+def test_living_regimes_carry_no_lagged_transitory_shock(model: Model) -> None:
     assert not any(
-        "hcc_transitory" in folded_model.state_names(regime_name=name)
-        for name in folded_model.user_regimes
+        "hcc_transitory" in model.state_names(regime_name=name)
+        for name in model.user_regimes
     )
 
 
-def test_living_regimes_keep_the_lagged_persistent_shock(folded_model: Model) -> None:
+def test_living_regimes_keep_the_lagged_persistent_shock(model: Model) -> None:
     assert all(
-        "hcc_persistent" in folded_model.state_names(regime_name=name)
-        for name in folded_model.user_regimes
+        "hcc_persistent" in model.state_names(regime_name=name)
+        for name in model.user_regimes
         if name != "dead"
     )
-
-
-@pytest.mark.long_running
-def test_folded_value_equals_every_transitory_slice_of_the_unfolded_value(
-    folded_model: Model,
-) -> None:
-    """The two models' sums differ only in order, so they agree to rounding."""
-    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_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}",
-            )