Restart simulations mid-life on pylcm #604 and offer uncompensated-care targets

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

src/aca_model/aca/model.py

@@ -7,7 +7,7 @@
 from collections.abc import Mapping
 from typing import Any
 
-from lcm import AgeGrid, DiscreteGrid, ExecutionConfig, Model
+from lcm import AgeGrid, DiscreteGrid, ExecutionConfig, InitialNodes, Model
 from lcm.typing import UserParams
 
 from aca_model.aca import PolicyVariant
@@ -34,6 +34,7 @@ def create_model(
     solver: SolverName = "brute_force",
     consumption_dollars_points: tuple[float, ...] | None = None,
     execution_config: ExecutionConfig | None = None,
+    initial_nodes: InitialNodes = INITIAL_NODES,
 ) -> Model:
     """Create an ACA policy variant model.
 
@@ -61,6 +62,9 @@ def create_model(
         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.
+        initial_nodes: Admissible starting `(age, regime)` pairs. The default
+            admits the estimation sample's entry ages and regimes; pass
+            `build_restart_initial_nodes` to restart simulations mid-life.
 
     Returns:
         pylcm Model.
@@ -91,7 +95,7 @@ def create_model(
 
     return Model(
         regimes=regimes,
-        initial_nodes=INITIAL_NODES,
+        initial_nodes=initial_nodes,
         ages=ages,
         regime_id_class=RegimeId,
         description=f"Structural retirement model ({policy.name})",

src/aca_model/aca/regimes/_overrides.py

@@ -8,6 +8,7 @@
 
 from aca_model.aca import health_insurance as aca_hi
 from aca_model.aca.health_insurance import PolicyVariant
+from aca_model.baseline import health_insurance as baseline_hi
 from aca_model.baseline.regimes._common import RegimeSpec
 from aca_model.environment import taxes
 
@@ -55,6 +56,7 @@ def apply_aca_overrides(
     )
 
     if has_medicaid_expansion:
+        functions["aca_magi"] = baseline_hi.aca_magi
         functions["is_medicaid_eligible"] = aca_hi.is_medicaid_eligible
 
     if has_subsidies and spec["his"] == "nongroup" and spec["mc"] == "nomc":
@@ -64,8 +66,9 @@ def apply_aca_overrides(
         # dict, not a DAG function — no entry needed here.
         functions["plan_premium"] = aca_hi.community_rated_premium
         # The community-rated premium replaces the risk-rated one, so the
-        # risk-rated premium's zero-profit intercept has no consumer left.
+        # risk-rated inputs have no consumer left.
         functions.pop("private_premium_intercept", None)
+        functions.pop("predicted_hcc_insurer", None)
         functions["hic_premium_subsidy"] = aca_hi.premium_subsidy
         # Split at the bracket, which takes finitely many values, so the
         # out-of-pocket bill can be averaged per bracket.

src/aca_model/baseline/model.py

@@ -13,7 +13,7 @@
 from collections.abc import Mapping
 from typing import Any
 
-from lcm import AgeGrid, DiscreteGrid, ExecutionConfig, Model
+from lcm import AgeGrid, DiscreteGrid, ExecutionConfig, InitialNodes, Model
 from lcm.typing import UserParams
 
 from aca_model.baseline.regimes import (
@@ -37,6 +37,7 @@ def create_model(
     solver: SolverName = "brute_force",
     consumption_dollars_points: tuple[float, ...] | None = None,
     execution_config: ExecutionConfig | None = None,
+    initial_nodes: InitialNodes = INITIAL_NODES,
 ) -> Model:
     """Create the baseline structural retirement model.
 
@@ -73,6 +74,9 @@ def create_model(
         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.
+        initial_nodes: Admissible starting `(age, regime)` pairs. The default
+            admits the estimation sample's entry ages and regimes; pass
+            `build_restart_initial_nodes` to restart simulations mid-life.
 
     Returns:
         A pylcm Model with 19 regimes (18 non-terminal + dead) spanning
@@ -102,7 +106,7 @@ def create_model(
 
     return Model(
         regimes=regimes,
-        initial_nodes=INITIAL_NODES,
+        initial_nodes=initial_nodes,
         ages=ages,
         regime_id_class=RegimeId,
         description="Baseline structural retirement model (pre-ACA)",

src/aca_model/baseline/regimes/__init__.py

@@ -33,6 +33,7 @@
     build_model_functions,
     build_model_state_transitions,
     build_model_states,
+    build_restart_initial_nodes,
 )
 from aca_model.baseline.regimes._dcegm import build_dcegm_solver
 from aca_model.baseline.regimes._nbegm import build_nbegm_solver
@@ -46,6 +47,7 @@
     "build_all_regimes",
     "build_model_slots",
     "build_regime",
+    "build_restart_initial_nodes",
     "nongroup",
     "retiree",
     "tied",

src/aca_model/baseline/regimes/_common.py

@@ -342,7 +342,7 @@ def build_grids(
         consumption_dollars=(
             IrregSpacedGrid(n_points=grid_config.n_consumption_dollars_gridpoints)
             if consumption_dollars_points is None
-            else IrregSpacedGrid(points=consumption_dollars_points)
+            else IrregSpacedGrid(points=np.asarray(consumption_dollars_points))
         ),
         wage_res=wage_res,
         hcc_persistent=hcc_persistent,
@@ -378,7 +378,7 @@ def build_assets_grid(*, floor: float, n_points: int) -> IrregSpacedGrid:
     )
     negative[0] = floor
     positive[-1] = _ASSETS_TOP
-    return IrregSpacedGrid(points=np.concatenate([negative, positive]).tolist())
+    return IrregSpacedGrid(points=np.concatenate([negative, positive]))
 
 
 def get_hcc_persistent_shock(*, grid_config: GridConfig) -> RouwenhorstAR1Process:
@@ -526,6 +526,24 @@ def transition_ages(spec: RegimeSpec) -> tuple[int, ...]:
     raise ValueError(f"Unknown regime spec: {spec}")
 
 
+def build_restart_initial_nodes(*, ages: Iterable[int]) -> InitialNodes:
+    """Admit every living regime active at each of `ages` as a start.
+
+    A simulation restarted from the states another simulation reached at some
+    age starts in whichever living regime that simulation had arrived at.
+    """
+    return InitialNodes(
+        by_age={
+            age: tuple(
+                name
+                for name, spec in REGIME_SPECS.items()
+                if age in transition_ages(spec)
+            )
+            for age in ages
+        }
+    )
+
+
 def next_model_age(period: Period) -> FloatND:
     """Advance by one grid position, never by an assumed number of years.
 
@@ -845,12 +863,10 @@ def build_model_functions(*, solver: SolverName = "brute_force") -> dict:
     # assets adjustment, both of which scale by the marginal income tax rate.
     functions["marginal_tax_rate"] = taxes.marginal_rate
 
-    # HIC premium
-    functions["predicted_hcc_insurer"] = health_insurance.hcc_insurer_predicted
-
     # Transfers
     functions["consumption_dollars_floor"] = assets_and_income.consumption_dollars_floor
-    functions["transfers"] = assets_and_income.transfers
+    if solver == "brute_force":
+        functions["transfers"] = assets_and_income.transfers
     functions["consumption_equiv"] = preferences.consumption_equiv
 
     return functions
@@ -999,10 +1015,8 @@ def build_common_functions(spec: RegimeSpec, *, subtract_bill: bool = False) ->
         if spec["mc"] == "oamc"
         else health_insurance.is_disabled_from_health
     )
-    # MAGI for the ACA Medicaid expansion track; pruned in the baseline DAG.
-    functions["aca_magi"] = health_insurance.aca_magi
-
     # Swapped per policy variant by the ACA overlay, hence regime-level
+    functions["predicted_hcc_insurer"] = health_insurance.hcc_insurer_predicted
     functions["is_medicaid_eligible"] = health_insurance.is_medicaid_eligible
     functions["after_tax_income"] = taxes.after_tax_income
     if subtract_bill:

src/aca_model/baseline/regimes/_dcegm.py

@@ -11,6 +11,7 @@
 the one piece that needs `lcm.solvers`.
 """
 
+import numpy as np
 from lcm import IrregSpacedGrid
 from lcm.solvers import DCEGM
 
@@ -31,7 +32,9 @@ def build_dcegm_solver(grids: Grids) -> DCEGM:
     assets_points = grids.assets.to_jax()
     savings_stop = float(assets_points[-1]) - float(assets_points[0])
     savings_grid = IrregSpacedGrid(
-        points=tuple(savings_stop * (i / (n_points - 1)) ** 3 for i in range(n_points)),
+        points=np.array(
+            [savings_stop * (i / (n_points - 1)) ** 3 for i in range(n_points)]
+        ),
     )
     return DCEGM(savings_grid=savings_grid)
 

src/aca_model/baseline/regimes/_nbegm.py

@@ -4,6 +4,7 @@
 Device placement and compiled program widths belong to the model execution policy.
 """
 
+import numpy as np
 from lcm import IrregSpacedGrid
 from lcm.solvers import NBEGM
 
@@ -17,7 +18,9 @@ def build_nbegm_solver(grids: Grids) -> NBEGM:
     assets_points = grids.assets.to_jax()
     savings_stop = float(assets_points[-1]) - float(assets_points[0])
     savings_grid = IrregSpacedGrid(
-        points=tuple(savings_stop * (i / (n_points - 1)) ** 3 for i in range(n_points)),
+        points=np.array(
+            [savings_stop * (i / (n_points - 1)) ** 3 for i in range(n_points)]
+        ),
     )
     return NBEGM(
         savings_grid=savings_grid,

src/aca_model/benchmark.py

@@ -29,14 +29,14 @@
 import jax.numpy as jnp
 import numpy as np
 from jax import Array
-from lcm import DiscreteGrid, ExecutionConfig, Model
+from lcm import DiscreteGrid, ExecutionConfig, InitialNodes, Model
 
 from aca_model.agent.health import GoodHealth
 from aca_model.agent.labor_market import IsMarried
 from aca_model.agent.preferences import BenchmarkPrefType
 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.baseline.regimes import ENTRY_REGIMES, INITIAL_NODES
 from aca_model.config import BENCHMARK_GRID_CONFIG, GridConfig
 from aca_model.consumption_dollars_grid import (
     compute_consumption_dollars_points,
@@ -68,6 +68,7 @@ def create_benchmark_model(
     pref_type_grid: DiscreteGrid,
     grid_config: GridConfig = BENCHMARK_GRID_CONFIG,
     execution_config: ExecutionConfig | None = None,
+    initial_nodes: InitialNodes = INITIAL_NODES,
 ) -> Model:
     """Create the aca baseline with `BENCHMARK_GRID_CONFIG` and frozen fixed_params.
 
@@ -80,6 +81,8 @@ def create_benchmark_model(
         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.
+        initial_nodes: Admissible starting `(age, regime)` pairs, forwarded
+            to `create_model`.
     """
     fixed_params, wage_params, _ = get_benchmark_params(model=None)
     return create_model(
@@ -89,6 +92,7 @@ def create_benchmark_model(
         derived_categoricals=_DERIVED_CATEGORICALS,
         pref_type_grid=pref_type_grid,
         execution_config=execution_config,
+        initial_nodes=initial_nodes,
     )
 
 

tests/test_aca_plan.py

@@ -83,6 +83,21 @@ def test_aca_non_group_premium_is_community_rated(
     assert _aca_premium(policy, spousal_income) == pytest.approx(expected)
 
 
+@pytest.mark.parametrize("policy", list(PolicyVariant))
+def test_aca_non_group_risk_prediction_is_registered_only_without_reform(
+    policy: PolicyVariant,
+) -> None:
+    """The community-rated premium does not register a risk prediction node."""
+    functions = {**build_model_functions(), **_build_functions(REGIME_SPECS[REGIME])}
+    apply_aca_overrides(functions, REGIME_SPECS[REGIME], policy)
+    expected = (
+        health_insurance.hcc_insurer_predicted
+        if policy == PolicyVariant.ACA_ONLY_MEDICAID_EXPANSION
+        else None
+    )
+    assert functions.get("predicted_hcc_insurer") is expected
+
+
 RISK_RATED_PREMIUM_PARAMS = frozenset(
     {
         "premium_markup",

tests/test_benchmark.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 import pytest
+from _lcm.simulation.additional_targets import _collect_all_available_targets
 from lcm import (
     DiscreteGrid,
     GridBreakpoint,
@@ -112,6 +113,24 @@ def test_benchmark_panel_exposes_hic_premium_and_wage_targets() -> None:
     assert (defined_wage > 0).all()
 
 
+def test_benchmark_model_offers_uncompensated_care_targets() -> None:
+    """Transfers, the floor and the health-cost covariates can be added to the panel.
+
+    Uncompensated care is the part of `transfers` above
+    `consumption_dollars_floor`. Total health costs depend on the next period's
+    shocks, so they are rebuilt from `is_married`, `good_health` and the shock
+    states rather than offered as a target.
+    """
+    model = create_benchmark_model(pref_type_grid=DiscreteGrid(BenchmarkPrefType))
+    available = _collect_all_available_targets(model._regimes)  # noqa: SLF001
+    assert {
+        "transfers",
+        "consumption_dollars_floor",
+        "is_married",
+        "good_health",
+    } <= available
+
+
 @pytest.mark.long_running
 def test_benchmark_simulate_obeys_borrowing_constraint() -> None:
     """`consumption_dollars <= max(cash_on_hand, floor)` holds for every alive row.

tests/test_model_creation.py

@@ -13,11 +13,15 @@
 from aca_model.aca import health_insurance as aca_hi
 from aca_model.aca.health_insurance import PolicyVariant
 from aca_model.aca.regimes import build_all_regimes as _build_aca_regimes
+from aca_model.agent import assets_and_income
 from aca_model.agent.preferences import BenchmarkPrefType
+from aca_model.baseline import health_insurance as baseline_hi
 from aca_model.baseline.regimes import REGIME_SPECS, RegimeId
 from aca_model.baseline.regimes import build_regime as _build_regime
 from aca_model.baseline.regimes._common import (
+    SolverName,
     build_grids,
+    build_model_functions,
     build_model_state_transitions,
     build_model_states,
 )
@@ -290,6 +294,36 @@ def test_aca_model_creates_successfully() -> None:
     assert model.n_periods == 46
 
 
+@pytest.mark.parametrize("name", REGIME_SPECS)
+def test_baseline_regime_omits_expansion_magi(name: str) -> None:
+    """Baseline eligibility does not register the ACA expansion income node."""
+    assert "aca_magi" not in build_regime(name).functions
+
+
+@pytest.mark.parametrize("policy", list(PolicyVariant))
+def test_aca_regime_registers_magi_only_for_expansion(policy: PolicyVariant) -> None:
+    """The expansion eligibility consumer owns its MAGI income node."""
+    has_expansion = policy in (
+        PolicyVariant.ACA,
+        PolicyVariant.ACA_NO_MANDATE,
+        PolicyVariant.ACA_ONLY_MEDICAID_EXPANSION,
+    )
+    expected = baseline_hi.aca_magi if has_expansion else None
+    regimes = build_aca_regimes(policy)
+    assert {
+        name: regimes[name].functions.get("aca_magi") for name in REGIME_SPECS
+    } == dict.fromkeys(REGIME_SPECS, expected)
+
+
+@pytest.mark.parametrize("solver", ["brute_force", "dcegm", "nbegm"])
+def test_model_registers_transfers_only_for_brute_asset_laws(
+    solver: SolverName,
+) -> None:
+    """Savings-form asset laws already include transfers through resources."""
+    expected = assets_and_income.transfers if solver == "brute_force" else None
+    assert build_model_functions(solver=solver).get("transfers") is expected
+
+
 def test_aca_nongroup_inelig_has_real_functions() -> None:
     """Nongroup+nomc regimes get real ACA functions under ACA policy."""
     regimes = build_aca_regimes(PolicyVariant.ACA)