diff --git a/src/aca_model/aca/model.py b/src/aca_model/aca/model.py index 5b17c97..971b8b0 100644 --- a/src/aca_model/aca/model.py +++ b/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})", diff --git a/src/aca_model/baseline/model.py b/src/aca_model/baseline/model.py index 3522e30..1b1efe3 100644 --- a/src/aca_model/baseline/model.py +++ b/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)", diff --git a/src/aca_model/baseline/regimes/__init__.py b/src/aca_model/baseline/regimes/__init__.py index bd83aed..65b1f19 100644 --- a/src/aca_model/baseline/regimes/__init__.py +++ b/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", diff --git a/src/aca_model/baseline/regimes/_common.py b/src/aca_model/baseline/regimes/_common.py index 2032ce1..5dd4180 100644 --- a/src/aca_model/baseline/regimes/_common.py +++ b/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. diff --git a/src/aca_model/benchmark.py b/src/aca_model/benchmark.py index 3eed460..c3ca289 100644 --- a/src/aca_model/benchmark.py +++ b/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, ) diff --git a/tests/test_restart_initial_nodes.py b/tests/test_restart_initial_nodes.py new file mode 100644 index 0000000..0fa69ea --- /dev/null +++ b/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"]