Make the model horizon a constructor argument; shorten CI

.github/workflows/main.yml

@@ -37,6 +37,6 @@ jobs:
           pip install "pylcm @
           git+https://github.com/OpenSourceEconomics/pylcm.git@7e99a19badafe20e29a517f69d9ec155bba8e0a0"
       - name: Install aca-model with test deps
-        run: pip install -e . pytest pdbp
+        run: pip install -e . pytest pytest-xdist pdbp
       - name: Run pytest
-        run: pytest tests
+        run: pytest -n 4 --dist loadscope --durations=40 tests

docs/dag-migration.md

@@ -1,6 +1,6 @@
 # ACA: one structural stage table, existing numerical programs
 
-Both factories declare the same admissible starts, `INITIAL_NODES`: ages 51–60
+Both factories admit the same starts by default, from `build_entry_initial_nodes`: ages 51–60
 as explicit `(age, regime)` pairs in the five `*_inelig_canwork`
 regimes, 50 pairs. Every later age, and `dead` at any age, is reached only through
 transitions; the engine derives the solved domain from these roots, **184** nodes.
@@ -10,7 +10,7 @@ admission.
 `simulate_with_dense_index` drops initial-condition rows older than the last
 admissible start age and raises on any kept row outside `model.initial_nodes`.
 
-`_AGE_STAGES` and the host-side clock supply both numerical next-age routing and
+`age_stages` and the host-side clock supply both numerical next-age routing and
 `Model.edges` source-age selectors, grouped by destination. Each source has one
 mapping of `StochasticTransition` probability cells; the model graph selects the
 cells allowed at each source age. Code order/global regime IDs and parameter

src/aca_model/aca/model.py

@@ -14,12 +14,12 @@
 from aca_model.aca.regimes import build_all_regimes
 from aca_model.baseline.model import _fail_if_dcegm_without_consumption_points
 from aca_model.baseline.regimes import (
-    INITIAL_NODES,
     RegimeId,
     SolverName,
+    build_entry_initial_nodes,
     build_model_slots,
 )
-from aca_model.config import MODEL_AGES, GridConfig
+from aca_model.config import MODEL_CONFIG, GridConfig, model_ages
 from aca_model.execution import execution_config_for_devices
 
 
@@ -34,7 +34,8 @@ def create_model(
     solver: SolverName = "brute_force",
     consumption_dollars_points: tuple[float, ...] | None = None,
     execution_config: ExecutionConfig | None = None,
-    initial_nodes: InitialNodes = INITIAL_NODES,
+    initial_nodes: InitialNodes | None = None,
+    final_age: int = MODEL_CONFIG.end_age,
 ) -> Model:
     """Create an ACA policy variant model.
 
@@ -62,15 +63,21 @@ 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.
+        initial_nodes: Admissible starting `(age, regime)` pairs. None admits
+            the estimation sample's entry ages and regimes that lie before
+            `final_age`; pass `build_restart_initial_nodes` to restart
+            simulations mid-life.
+        final_age: First age at which everybody is dead; the model spans
+            `MODEL_CONFIG.start_age` through `final_age`. The default is the
+            full horizon. A shorter horizon reads a prefix of the age-indexed
+            data, and `fixed_params` must hold only what it reads (see
+            `select_horizon_params`).
 
     Returns:
         pylcm Model.
 
     """
-    ages = AgeGrid(exact_values=MODEL_AGES)
+    ages = model_ages(final_age=final_age)
     _fail_if_dcegm_without_consumption_points(
         solver=solver, consumption_dollars_points=consumption_dollars_points
     )
@@ -90,13 +97,18 @@ def create_model(
         fixed_params=fixed_params,
         wage_params=wage_params,
         pref_type_grid=pref_type_grid,
+        ages=ages,
         solver=solver,
     )
 
     return Model(
         regimes=regimes,
-        initial_nodes=initial_nodes,
-        ages=ages,
+        initial_nodes=(
+            build_entry_initial_nodes(ages=ages)
+            if initial_nodes is None
+            else initial_nodes
+        ),
+        ages=AgeGrid(exact_values=ages),
         regime_id_class=RegimeId,
         description=f"Structural retirement model ({policy.name})",
         fixed_params=fixed_params,

src/aca_model/baseline/model.py

@@ -17,13 +17,13 @@
 from lcm.typing import UserParams
 
 from aca_model.baseline.regimes import (
-    INITIAL_NODES,
     RegimeId,
     SolverName,
     build_all_regimes,
+    build_entry_initial_nodes,
     build_model_slots,
 )
-from aca_model.config import MODEL_AGES, GridConfig
+from aca_model.config import MODEL_CONFIG, GridConfig, model_ages
 from aca_model.execution import execution_config_for_devices
 
 
@@ -37,7 +37,8 @@ def create_model(
     solver: SolverName = "brute_force",
     consumption_dollars_points: tuple[float, ...] | None = None,
     execution_config: ExecutionConfig | None = None,
-    initial_nodes: InitialNodes = INITIAL_NODES,
+    initial_nodes: InitialNodes | None = None,
+    final_age: int = MODEL_CONFIG.end_age,
 ) -> Model:
     """Create the baseline structural retirement model.
 
@@ -74,17 +75,24 @@ 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.
+        initial_nodes: Admissible starting `(age, regime)` pairs. None admits
+            the estimation sample's entry ages and regimes that lie before
+            `final_age`; pass `build_restart_initial_nodes` to restart
+            simulations mid-life.
+        final_age: First age at which everybody is dead; the model spans
+            `MODEL_CONFIG.start_age` through `final_age`. The default is the
+            full horizon. A shorter horizon reads a prefix of the age-indexed
+            data, and `fixed_params` must hold only what it reads (see
+            `select_horizon_params`).
 
     Returns:
         A pylcm Model with 19 regimes (18 non-terminal + dead) spanning
-        ages 51-96; 95 is the last age anyone is alive. Regime names follow
+        ages 51 through `final_age` (96 by default); `final_age - 1` is the
+        last age anyone is alive. Regime names follow
         the `<his>_<medicare>_<ss>_<work>` scheme.
 
     """
-    ages = AgeGrid(exact_values=MODEL_AGES)
+    ages = model_ages(final_age=final_age)
     _fail_if_dcegm_without_consumption_points(
         solver=solver, consumption_dollars_points=consumption_dollars_points
     )
@@ -101,13 +109,18 @@ def create_model(
         fixed_params=fixed_params,
         wage_params=wage_params,
         pref_type_grid=pref_type_grid,
+        ages=ages,
         solver=solver,
     )
 
     return Model(
         regimes=regimes,
-        initial_nodes=initial_nodes,
-        ages=ages,
+        initial_nodes=(
+            build_entry_initial_nodes(ages=ages)
+            if initial_nodes is None
+            else initial_nodes
+        ),
+        ages=AgeGrid(exact_values=ages),
         regime_id_class=RegimeId,
         description="Baseline structural retirement model (pre-ACA)",
         fixed_params=fixed_params,

src/aca_model/baseline/regimes/__init__.py

@@ -22,34 +22,38 @@
 from aca_model.baseline.regimes import _tied as tied
 from aca_model.baseline.regimes._common import (
     ENTRY_REGIMES,
-    INITIAL_NODES,
     REGIME_SPECS,
     Grids,
     RegimeId,
     SolverName,
     build_dead_regime,
+    build_entry_initial_nodes,
     build_grids,
     build_model_constraints,
     build_model_functions,
     build_model_state_transitions,
     build_model_states,
     build_restart_initial_nodes,
+    reachable_regimes,
+    select_horizon_params,
 )
 from aca_model.baseline.regimes._dcegm import build_dcegm_solver
 from aca_model.baseline.regimes._nbegm import build_nbegm_solver
 from aca_model.config import GridConfig
 
 __all__ = [
     "ENTRY_REGIMES",
-    "INITIAL_NODES",
     "REGIME_SPECS",
     "RegimeId",
     "build_all_regimes",
+    "build_entry_initial_nodes",
     "build_model_slots",
     "build_regime",
     "build_restart_initial_nodes",
     "nongroup",
+    "reachable_regimes",
     "retiree",
+    "select_horizon_params",
     "tied",
 ]
 
@@ -134,6 +138,7 @@ def build_model_slots(
     fixed_params: UserParams,
     wage_params: Mapping[str, Any],
     pref_type_grid: DiscreteGrid,
+    ages: tuple[int, ...],
     solver: SolverName = "brute_force",
 ) -> dict[str, Any]:
     """Build the shared model topology and broadcast regime slots.
@@ -144,6 +149,7 @@ def build_model_slots(
     overlay swaps only regime-level functions, so the broadcast slots are
     policy-invariant. Under DC-EGM the solver-contract functions join the
     broadcast set; the borrowing constraint is declared under every solver.
+    `ages` is the model's age grid, which the edges' age schedules cover.
     """
     grids = build_grids(
         grid_config=grid_config,
@@ -152,21 +158,21 @@ def build_model_slots(
         pref_type_grid=pref_type_grid,
     )
     return {
-        "edges": build_edges(),
+        "edges": build_edges(ages=ages),
         "functions": build_model_functions(solver=solver),
         "constraints": build_model_constraints(solver=solver),
         "states": build_model_states(grids),
         "state_transitions": build_model_state_transitions(),
     }
 
 
-def build_edges() -> dict[str, Transition]:
+def build_edges(*, ages: tuple[int, ...]) -> dict[str, Transition]:
     """Declare each living regime's outgoing edges by their law.
 
     The law's age schedule names each destination at the source ages that reach
     it. `dead` has no outgoing edges and is therefore terminal.
     """
     return {
-        name: Transition(law=_HIS_LAWS[spec["his"]](name))
+        name: Transition(law=_HIS_LAWS[spec["his"]](name, ages=ages))
         for name, spec in REGIME_SPECS.items()
     }