Let simulations restart mid-life from any living regime

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/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

@@ -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.
 

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_restart_initial_nodes.py

@@ -0,0 +1,56 @@
+"""Starts mid-life, at the regimes a simulation can have reached by that age."""
+
+import pandas as pd
+import pytest
+from lcm import DiscreteGrid
+
+from aca_model.agent.preferences import BenchmarkPrefType
+from aca_model.baseline.regimes import build_restart_initial_nodes
+from aca_model.benchmark import create_benchmark_model
+from aca_model.simulation import select_admissible_starts
+
+
+@pytest.mark.parametrize(
+    ("age", "expected"),
+    [
+        (
+            61,
+            (
+                "retiree_nomc_inelig_canwork",
+                "tied_nomc_inelig_canwork",
+                "nongroup_nomc_inelig_canwork",
+                "retiree_dimc_inelig_canwork",
+                "nongroup_dimc_inelig_canwork",
+            ),
+        ),
+        (
+            66,
+            (
+                "retiree_oamc_choose_canwork",
+                "tied_oamc_choose_canwork",
+                "nongroup_oamc_choose_canwork",
+            ),
+        ),
+        (
+            73,
+            ("retiree_oamc_forced_forcedout", "nongroup_oamc_forced_forcedout"),
+        ),
+    ],
+)
+def test_build_restart_initial_nodes_admits_every_living_regime_of_the_age(
+    age: int, expected: tuple[str, ...]
+) -> None:
+    """At each age, the starts are exactly the living regimes active then."""
+    nodes = build_restart_initial_nodes(ages=(age,))
+    assert set(nodes.by_age[age]) == set(expected)
+
+
+def test_model_with_restart_nodes_admits_a_start_on_old_age_medicare() -> None:
+    """A model built with restart nodes accepts a start at 66 on Medicare."""
+    model = create_benchmark_model(
+        pref_type_grid=DiscreteGrid(BenchmarkPrefType),
+        initial_nodes=build_restart_initial_nodes(ages=range(51, 70)),
+    )
+    starts = pd.DataFrame({"age": [66.0], "regime_name": ["tied_oamc_choose_canwork"]})
+    admitted = select_admissible_starts(model=model, initial_conditions=starts)
+    assert admitted["regime_name"].tolist() == ["tied_oamc_choose_canwork"]