From 7938e6efa902ae9304fe644221708397bd532406 Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Fri, 23 Oct 2026 08:16:17 +0200 Subject: [PATCH 1/6] Make the model horizon a constructor argument and shorten slow tests `create_model` (baseline and ACA) and `create_benchmark_model` take `final_age`, the first age at which everybody is dead, defaulting to the full horizon. The age stages, entry initial nodes, next-age lookup, transition ages and edges toward `dead` are functions of the ages instead of import-time constants. The medical-cost timing test builds ages 51-53, the NB-EGM wiring tests are long-running, and CI runs pytest on four workers. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Hs8mfXPYyEa24KaAp6BwCW --- .github/workflows/main.yml | 4 +- docs/dag-migration.md | 4 +- src/aca_model/aca/model.py | 29 ++-- src/aca_model/baseline/model.py | 32 +++-- src/aca_model/baseline/regimes/__init__.py | 12 +- src/aca_model/baseline/regimes/_common.py | 145 ++++++++++++-------- src/aca_model/baseline/regimes/_nongroup.py | 30 ++-- src/aca_model/baseline/regimes/_retiree.py | 32 +++-- src/aca_model/baseline/regimes/_tied.py | 22 ++- src/aca_model/benchmark.py | 12 +- src/aca_model/config.py | 32 ++++- tests/test_aca_employer_coverage.py | 4 +- tests/test_benchmark.py | 8 ++ tests/test_dcegm_model_creation.py | 3 +- tests/test_di_medicare_transitions.py | 8 +- tests/test_medical_cost_timing.py | 7 +- tests/test_mortality.py | 37 ++++- tests/test_nbegm_model_creation.py | 3 + tests/test_regime_transitions.py | 24 ++-- 19 files changed, 303 insertions(+), 145 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 71f5307..0df1cb1 100644 --- a/.github/workflows/main.yml +++ b/.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 --durations=20 tests diff --git a/docs/dag-migration.md b/docs/dag-migration.md index 5e4650f..0f43357 100644 --- a/docs/dag-migration.md +++ b/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 diff --git a/src/aca_model/aca/model.py b/src/aca_model/aca/model.py index 971b8b0..1d0943e 100644 --- a/src/aca_model/aca/model.py +++ b/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,20 @@ 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. 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 +96,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, diff --git a/src/aca_model/baseline/model.py b/src/aca_model/baseline/model.py index 1b1efe3..dd4e42e 100644 --- a/src/aca_model/baseline/model.py +++ b/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,23 @@ 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. 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 `___` 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 +108,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, diff --git a/src/aca_model/baseline/regimes/__init__.py b/src/aca_model/baseline/regimes/__init__.py index 65b1f19..495d4be 100644 --- a/src/aca_model/baseline/regimes/__init__.py +++ b/src/aca_model/baseline/regimes/__init__.py @@ -22,12 +22,12 @@ 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, @@ -41,10 +41,10 @@ __all__ = [ "ENTRY_REGIMES", - "INITIAL_NODES", "REGIME_SPECS", "RegimeId", "build_all_regimes", + "build_entry_initial_nodes", "build_model_slots", "build_regime", "build_restart_initial_nodes", @@ -134,6 +134,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 +145,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,7 +154,7 @@ 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), @@ -160,13 +162,13 @@ def build_model_slots( } -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() } diff --git a/src/aca_model/baseline/regimes/_common.py b/src/aca_model/baseline/regimes/_common.py index c74812b..cf7eaa3 100644 --- a/src/aca_model/baseline/regimes/_common.py +++ b/src/aca_model/baseline/regimes/_common.py @@ -460,30 +460,28 @@ def _compute_max_annual_labor_income( ) -# Structural stage data, shared by scheduled support and the numerical router. -# Each stage covers the half-open age interval `[start, stop)`. -_AGE_STAGES = ( - ( - (config.start_age, config.ss_early_age), - ("nomc_inelig", "dimc_inelig"), - ), - ( - (config.ss_early_age, config.medicare_age), - ("nomc_choose", "dimc_choose"), - ), - ( - (config.medicare_age, config.ss_forced_age), - ("forced_choose",) * 2, - ), - ( - (config.ss_forced_age, config.work_forced_out_age), - ("forced_forced",) * 2, - ), - ( - (config.work_forced_out_age, MODEL_AGES[-1]), - ("forcedout",) * 2, - ), -) +def age_stages( + *, final_age: int +) -> tuple[tuple[tuple[int, int], tuple[str, str]], ...]: + """Return the structural stages living through ages before `final_age`. + + Shared by scheduled support and the numerical router. Each stage covers the + half-open age interval `[start, stop)`, cut off at `final_age`, the first + age at which everybody is dead. A stage that starts at or after `final_age` + covers no age. + """ + stages = ( + ((config.start_age, config.ss_early_age), ("nomc_inelig", "dimc_inelig")), + ((config.ss_early_age, config.medicare_age), ("nomc_choose", "dimc_choose")), + ((config.medicare_age, config.ss_forced_age), ("forced_choose",) * 2), + ((config.ss_forced_age, config.work_forced_out_age), ("forced_forced",) * 2), + ((config.work_forced_out_age, config.end_age), ("forcedout",) * 2), + ) + return tuple( + ((start, min(stop, final_age)), branches) for (start, stop), branches in stages + ) + + # The living regimes a subject may start in: pre-claiming, pre-Medicare. ENTRY_REGIMES = ( "retiree_nomc_inelig_canwork", @@ -493,16 +491,22 @@ def _compute_max_annual_labor_income( "nongroup_dimc_inelig_canwork", ) -# Admissible starting (age, regime) pairs, shared by every model variant: the -# baseline estimation-sample ages. Later ages, and `dead` at any age, are reached only -# through transitions. -INITIAL_NODES = InitialNodes( - by_age={ - age: ENTRY_REGIMES - for age in MODEL_AGES - if config.start_age <= age <= config.last_start_age - } -) + +def build_entry_initial_nodes(*, ages: tuple[int, ...]) -> InitialNodes: + """Admit the entry regimes at the baseline estimation-sample ages. + + Shared by every model variant. Later ages, and `dead` at any age, are reached + only through transitions. Start ages stop before the final age, at which + everybody is dead. + """ + return InitialNodes( + by_age={ + age: ENTRY_REGIMES + for age in ages[:-1] + if config.start_age <= age <= config.last_start_age + } + ) + _STAGE_KEY = { ("nomc", "inelig", "canwork"): "nomc_inelig", @@ -513,16 +517,19 @@ def _compute_max_annual_labor_income( ("oamc", "forced", "canwork"): "forced_forced", ("oamc", "forced", "forcedout"): "forcedout", } -# Dynamic JAX period lookup and construction-time schedule use the same clock. -_NEXT_AGES = np.asarray([float(a) for a in MODEL_AGES[1:]]) -def transition_ages(spec: RegimeSpec) -> tuple[int, ...]: - """Exact source coordinates at which the template declares its law.""" +def transition_ages( + spec: RegimeSpec, *, ages: tuple[int, ...] = MODEL_AGES +) -> tuple[int, ...]: + """Exact source coordinates at which the template declares its law. + + Empty when the regime's stage starts at or after the final age of `ages`. + """ key = _STAGE_KEY.get((spec["mc"], spec["ss"], spec["canwork"])) - for (start, stop), branches in _AGE_STAGES: + for (start, stop), branches in age_stages(final_age=ages[-1]): if key in branches: - return tuple(age for age in MODEL_AGES if start <= age < stop) + return tuple(age for age in ages if start <= age < stop) raise ValueError(f"Unknown regime spec: {spec}") @@ -544,21 +551,24 @@ def build_restart_initial_nodes(*, ages: Iterable[int]) -> InitialNodes: ) -def next_model_age(period: Period) -> FloatND: - """Advance by one grid position, never by an assumed number of years. +def next_model_age(period: Period, *, ages: tuple[int, ...]) -> FloatND: + """Advance by one position on `ages`, never by an assumed number of years. - Only covered nonterminal source periods call this function; no clipping - or extrapolation of invalid period indices is part of the contract. + Dynamic JAX period lookup and the construction-time schedule use the same + clock. Only covered nonterminal source periods call this function; no + clipping or extrapolation of invalid period indices is part of the contract. This does not by itself recalibrate ACA to a different period length. """ - return jnp.asarray(_NEXT_AGES)[period] + return jnp.asarray(np.asarray([float(age) for age in ages[1:]]))[period] -def _target_pair(next_age: int, group: dict[str, int]) -> tuple[int, int]: - if next_age == MODEL_AGES[-1]: +def _target_pair( + next_age: int, group: dict[str, int], *, final_age: int +) -> tuple[int, int]: + if next_age == final_age: dead = int(RegimeId.dead) return dead, dead - for (start, stop), branches in _AGE_STAGES: + for (start, stop), branches in age_stages(final_age=final_age): if start <= next_age < stop: return group[branches[0]], group[branches[1]] raise ValueError(f"Next age {next_age} is outside the declared stages") @@ -568,6 +578,7 @@ def build_regime_edges( *, spec: RegimeSpec, target_groups: tuple[dict[str, int], ...], + ages: tuple[int, ...], ) -> dict[RegimeName, tuple[int, ...]]: """Map each destination to the exact source ages that can reach it.""" # Both Medicare branches are reachable from every living regime: before 65 a @@ -576,10 +587,12 @@ def build_regime_edges( int(getattr(RegimeId, name)): name for name in (*REGIME_SPECS, "dead") } grouped: dict[RegimeName, list[int]] = {} - next_age_by_age = dict(itertools.pairwise(MODEL_AGES)) - for age in transition_ages(spec): + next_age_by_age = dict(itertools.pairwise(ages)) + for age in transition_ages(spec, ages=ages): next_age = next_age_by_age[age] - pairs = tuple(_target_pair(next_age, group) for group in target_groups) + pairs = tuple( + _target_pair(next_age, group, final_age=ages[-1]) for group in target_groups + ) supported_ids = {target for pair in pairs for target in pair} target_ids = tuple(sorted(supported_ids | {int(RegimeId.dead)})) for target_id in target_ids: @@ -587,10 +600,12 @@ def build_regime_edges( return {target: tuple(source_ages) for target, source_ages in grouped.items()} -def build_model_edges() -> dict[RegimeName, dict[RegimeName, tuple[int, ...]]]: +def build_model_edges( + *, ages: tuple[int, ...] = MODEL_AGES +) -> dict[RegimeName, dict[RegimeName, tuple[int, ...]]]: """Declare the policy-invariant topology of the living ACA regimes.""" return { - name: build_regime_edges(spec=spec, target_groups=make_targets(name)) + name: build_regime_edges(spec=spec, target_groups=make_targets(name), ages=ages) for name, spec in REGIME_SPECS.items() } @@ -600,21 +615,28 @@ def build_regime_transition( spec: RegimeSpec, transition_func: Callable[..., FloatND], target_groups: tuple[dict[str, int], ...], + ages: tuple[int, ...], ) -> ByAge: """Name each destination's probability cell at the source ages that reach it. One scalar probability cell per destination is shared by every source age; - the ages with the same destinations form one case. + the ages with the same destinations form one case. A regime whose stage + starts at or after the final age of `ages` is unreachable; its schedule + holds only the `dead` cell as a default that no period reads. """ - edges = build_regime_edges(spec=spec, target_groups=target_groups) + edges = build_regime_edges(spec=spec, target_groups=target_groups, ages=ages) cells = build_granular_regime_transition( transition_func=transition_func, target_ids=(int(getattr(RegimeId, target)) for target in edges), ) ages_by_targets: dict[tuple[RegimeName, ...], list[int]] = {} - for age in transition_ages(spec): - targets = tuple(target for target, ages in edges.items() if age in ages) + for age in transition_ages(spec, ages=ages): + targets = tuple( + target for target, source_ages in edges.items() if age in source_ages + ) ages_by_targets.setdefault(targets, []).append(age) + if not ages_by_targets: + return ByAge(cases={}, default=cells) return ByAge( cases={ tuple(ages): {target: cells[target] for target in targets} @@ -1183,10 +1205,15 @@ def select_target_for_age( next_age: int | IntND | FloatND, mc_next: bool | BoolND, tgts: dict[str, int], + *, + final_age: int = config.end_age, ) -> IntND: - """Select with the same stage table that supplies the schedule's support.""" + """Select with the same stage table that supplies the schedule's support. + + A `next_age` at or after `final_age` selects `dead`. + """ result = jnp.asarray(RegimeId.dead) - for (start, stop), branches in reversed(_AGE_STAGES): + for (start, stop), branches in reversed(age_stages(final_age=final_age)): target = jnp.where(jnp.array(mc_next), tgts[branches[1]], tgts[branches[0]]) inside = (next_age >= start) & (next_age < stop) result = jnp.where(inside, target, result) diff --git a/src/aca_model/baseline/regimes/_nongroup.py b/src/aca_model/baseline/regimes/_nongroup.py index a346e0c..4bc48e7 100644 --- a/src/aca_model/baseline/regimes/_nongroup.py +++ b/src/aca_model/baseline/regimes/_nongroup.py @@ -41,6 +41,8 @@ def _make_transition_canwork( own: dict[str, int], + *, + ages: tuple[int, ...], ) -> Callable[..., FloatND]: """Create transition for canwork nongroup regimes. @@ -58,8 +60,12 @@ def transition( ) -> FloatND: del age # Keep the legacy signature; period owns the clock lookup. return build_regime_probs_with_di_medicare( - target_dimc=select_target_for_age(next_model_age(period), True, own), - target_nomc=select_target_for_age(next_model_age(period), False, own), + target_dimc=select_target_for_age( + next_model_age(period, ages=ages), True, own, final_age=ages[-1] + ), + target_nomc=select_target_for_age( + next_model_age(period, ages=ages), False, own, final_age=ages[-1] + ), prob_dimc=prob_di_medicare_next( labor_supply, prob_disabled_next[period, health] ), @@ -72,6 +78,8 @@ def transition( def _make_transition_forcedout( gets_medicare: bool, own: dict[str, int], + *, + ages: tuple[int, ...], ) -> Callable[..., FloatND]: """Create transition for forcedout nongroup regimes. @@ -86,7 +94,9 @@ def transition( survival_probs: FloatND, ) -> FloatND: del age # Keep the legacy signature; period owns the clock lookup. - target = select_target_for_age(next_model_age(period), gets_medicare, own) + target = select_target_for_age( + next_model_age(period, ages=ages), gets_medicare, own, final_age=ages[-1] + ) return build_regime_probs(target, survival_probs[period, health]) return transition @@ -137,17 +147,21 @@ def build_constraints(spec: RegimeSpec) -> dict: return {} -def build_law(name: str) -> ByAge: - """Build the nongroup regime's per-target transition probabilities by source age.""" +def build_law(name: str, *, ages: tuple[int, ...]) -> ByAge: + """Build the nongroup regime's per-target transition probabilities by source age. + + `ages` is the model's age grid; the final age is the first at which everybody + is dead. + """ spec = REGIME_SPECS[name] gets_mc = spec["mc"] != "nomc" own, _ng = make_targets(name) if spec["canwork"] == "canwork": - transition_func = _make_transition_canwork(own) + transition_func = _make_transition_canwork(own, ages=ages) else: - transition_func = _make_transition_forcedout(gets_mc, own) + transition_func = _make_transition_forcedout(gets_mc, own, ages=ages) return build_regime_transition( - spec=spec, transition_func=transition_func, target_groups=(own,) + spec=spec, transition_func=transition_func, target_groups=(own,), ages=ages ) diff --git a/src/aca_model/baseline/regimes/_retiree.py b/src/aca_model/baseline/regimes/_retiree.py index a0f304b..0789376 100644 --- a/src/aca_model/baseline/regimes/_retiree.py +++ b/src/aca_model/baseline/regimes/_retiree.py @@ -46,6 +46,8 @@ def _make_transition_canwork( own: dict[str, int], ng: dict[str, int], + *, + ages: tuple[int, ...], ) -> Callable[..., FloatND]: """Create transition for canwork retiree regimes. @@ -65,11 +67,13 @@ def transition( prob_disabled_next: FloatND, ) -> FloatND: del age # Keep the legacy signature; period owns the clock lookup. - next_age = next_model_age(period) + next_age = next_model_age(period, ages=ages) def target(mc_next: bool) -> IntND: - own_target = select_target_for_age(next_age, mc_next, own) - ng_target = select_target_for_age(next_age, mc_next, ng) + own_target = select_target_for_age( + next_age, mc_next, own, final_age=ages[-1] + ) + ng_target = select_target_for_age(next_age, mc_next, ng, final_age=ages[-1]) return jnp.where(is_ssi_eligible, ng_target, own_target) return build_regime_probs_with_di_medicare( @@ -88,6 +92,8 @@ def _make_transition_forcedout( gets_medicare: bool, own: dict[str, int], ng: dict[str, int], + *, + ages: tuple[int, ...], ) -> Callable[..., FloatND]: """Create transition for forcedout retiree regimes. @@ -104,9 +110,9 @@ def transition( ) -> FloatND: del age # Keep the legacy signature; period owns the clock lookup. sp = survival_probs[period, health] - next_age = next_model_age(period) - target = select_target_for_age(next_age, gets_medicare, own) - ng_ssi = select_target_for_age(next_age, gets_medicare, ng) + next_age = next_model_age(period, ages=ages) + target = select_target_for_age(next_age, gets_medicare, own, final_age=ages[-1]) + ng_ssi = select_target_for_age(next_age, gets_medicare, ng, final_age=ages[-1]) target = jnp.where(is_ssi_eligible, ng_ssi, target) return build_regime_probs(target, sp) @@ -133,17 +139,21 @@ def _build_functions(spec: RegimeSpec, *, subtract_bill: bool = False) -> dict: return functions -def build_law(name: str) -> ByAge: - """Build the retiree regime's per-target transition probabilities by source age.""" +def build_law(name: str, *, ages: tuple[int, ...]) -> ByAge: + """Build the retiree regime's per-target transition probabilities by source age. + + `ages` is the model's age grid; the final age is the first at which everybody + is dead. + """ spec = REGIME_SPECS[name] gets_mc = spec["mc"] != "nomc" own, ng = make_targets(name) if spec["canwork"] == "canwork": - transition_func = _make_transition_canwork(own, ng) + transition_func = _make_transition_canwork(own, ng, ages=ages) else: - transition_func = _make_transition_forcedout(gets_mc, own, ng) + transition_func = _make_transition_forcedout(gets_mc, own, ng, ages=ages) return build_regime_transition( - spec=spec, transition_func=transition_func, target_groups=(own, ng) + spec=spec, transition_func=transition_func, target_groups=(own, ng), ages=ages ) diff --git a/src/aca_model/baseline/regimes/_tied.py b/src/aca_model/baseline/regimes/_tied.py index 9a9c4b3..e627975 100644 --- a/src/aca_model/baseline/regimes/_tied.py +++ b/src/aca_model/baseline/regimes/_tied.py @@ -47,6 +47,8 @@ def _make_transition_canwork( own: dict[str, int], ng: dict[str, int], + *, + ages: tuple[int, ...], ) -> Callable[..., FloatND]: """Create transition for canwork tied regimes. @@ -66,12 +68,14 @@ def transition( prob_disabled_next: FloatND, ) -> FloatND: del age # Keep the legacy signature; period owns the clock lookup. - next_age = next_model_age(period) + next_age = next_model_age(period, ages=ages) to_nongroup = (labor_supply == LaborSupply.do_not_work) | is_ssi_eligible def target(mc_next: bool) -> IntND: - own_target = select_target_for_age(next_age, mc_next, own) - ng_target = select_target_for_age(next_age, mc_next, ng) + own_target = select_target_for_age( + next_age, mc_next, own, final_age=ages[-1] + ) + ng_target = select_target_for_age(next_age, mc_next, ng, final_age=ages[-1]) return jnp.where(to_nongroup, ng_target, own_target) return build_regime_probs_with_di_medicare( @@ -101,13 +105,17 @@ def _build_functions(spec: RegimeSpec, *, subtract_bill: bool = False) -> dict: return functions -def build_law(name: str) -> ByAge: - """Build the tied regime's per-target transition probabilities by source age.""" +def build_law(name: str, *, ages: tuple[int, ...]) -> ByAge: + """Build the tied regime's per-target transition probabilities by source age. + + `ages` is the model's age grid; the final age is the first at which everybody + is dead. + """ spec = REGIME_SPECS[name] own, ng = make_targets(name) - transition_func = _make_transition_canwork(own, ng) + transition_func = _make_transition_canwork(own, ng, ages=ages) return build_regime_transition( - spec=spec, transition_func=transition_func, target_groups=(own, ng) + spec=spec, transition_func=transition_func, target_groups=(own, ng), ages=ages ) diff --git a/src/aca_model/benchmark.py b/src/aca_model/benchmark.py index c3ca289..836ef07 100644 --- a/src/aca_model/benchmark.py +++ b/src/aca_model/benchmark.py @@ -36,8 +36,8 @@ 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, INITIAL_NODES -from aca_model.config import BENCHMARK_GRID_CONFIG, GridConfig +from aca_model.baseline.regimes import ENTRY_REGIMES +from aca_model.config import BENCHMARK_GRID_CONFIG, MODEL_CONFIG, GridConfig from aca_model.consumption_dollars_grid import ( compute_consumption_dollars_points, inject_consumption_dollars_points, @@ -68,7 +68,8 @@ def create_benchmark_model( pref_type_grid: DiscreteGrid, grid_config: GridConfig = BENCHMARK_GRID_CONFIG, 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 aca baseline with `BENCHMARK_GRID_CONFIG` and frozen fixed_params. @@ -83,6 +84,10 @@ def create_benchmark_model( device-memory budget; CPU construction remains unbudgeted. initial_nodes: Admissible starting `(age, regime)` pairs, forwarded to `create_model`. + final_age: First age at which everybody is dead, forwarded to + `create_model`. The default is the full horizon; the frozen + age-indexed parameters cover it, and a shorter horizon reads a + prefix of them. """ fixed_params, wage_params, _ = get_benchmark_params(model=None) return create_model( @@ -93,6 +98,7 @@ def create_benchmark_model( pref_type_grid=pref_type_grid, execution_config=execution_config, initial_nodes=initial_nodes, + final_age=final_age, ) diff --git a/src/aca_model/config.py b/src/aca_model/config.py index d95569b..cb2c720 100644 --- a/src/aca_model/config.py +++ b/src/aca_model/config.py @@ -59,11 +59,37 @@ class GridConfig: MODEL_CONFIG = ModelConfig() + + +def model_ages(*, final_age: int) -> tuple[int, ...]: + """Return the model's ages, from `start_age` through `final_age`. + + `final_age` is the first age at which everybody is dead, so the model has + `final_age - start_age` nonterminal periods. The age-indexed data cover + `start_age` through `MODEL_CONFIG.end_age`; a shorter horizon reads a prefix + of them. + + Raises: + ValueError: If `final_age` leaves no nonterminal period or lies beyond + `MODEL_CONFIG.end_age`. + + """ + _fail_if_invalid_final_age(final_age) + return tuple(range(MODEL_CONFIG.start_age, final_age + 1)) + + +def _fail_if_invalid_final_age(final_age: int) -> None: + if not MODEL_CONFIG.start_age < final_age <= MODEL_CONFIG.end_age: + msg = ( + f"final_age must lie in ({MODEL_CONFIG.start_age}, " + f"{MODEL_CONFIG.end_age}], got {final_age}." + ) + raise ValueError(msg) + + # ACA is calibrated annually. Keep the exact clock on the host until Model build; # importing a task or regime builder must not initialize a JAX device allocator. -MODEL_AGES: tuple[int, ...] = tuple( - range(MODEL_CONFIG.start_age, MODEL_CONFIG.end_age + 1) -) +MODEL_AGES: tuple[int, ...] = model_ages(final_age=MODEL_CONFIG.end_age) GRID_CONFIG = GridConfig() BENCHMARK_GRID_CONFIG = GridConfig( diff --git a/tests/test_aca_employer_coverage.py b/tests/test_aca_employer_coverage.py index 53f6f6f..6b92443 100644 --- a/tests/test_aca_employer_coverage.py +++ b/tests/test_aca_employer_coverage.py @@ -21,7 +21,7 @@ _make_transition_canwork as retiree_canwork, ) from aca_model.baseline.regimes._tied import _make_transition_canwork as tied_canwork -from aca_model.config import MODEL_CONFIG +from aca_model.config import MODEL_AGES, MODEL_CONFIG N_PERIODS = MODEL_CONFIG.end_age - MODEL_CONFIG.start_age SURVIVAL = jnp.ones((N_PERIODS, 3)) * 0.99 @@ -81,7 +81,7 @@ def test_expansion_only_medicaid_keeps_employer_coverage( ) -> None: """An employer-covered household eligible only via the expansion keeps it.""" own, ng = make_targets(regime_name) - transition = make_transition(own=own, ng=ng) + transition = make_transition(own=own, ng=ng, ages=MODEL_AGES) result = _compose_with_aca_eligibility( "next_regime", transition, diff --git a/tests/test_benchmark.py b/tests/test_benchmark.py index 69b5b64..55864fb 100644 --- a/tests/test_benchmark.py +++ b/tests/test_benchmark.py @@ -48,6 +48,14 @@ def test_benchmark_model_builds_with_the_current_pylcm_grid_api() -> None: assert consumption.n_points == 5 +def test_benchmark_model_with_final_age_53_has_three_periods() -> None: + """Ages 51, 52 and 53 make three periods; everybody is dead at 53.""" + model = create_benchmark_model( + pref_type_grid=DiscreteGrid(BenchmarkPrefType), final_age=53 + ) + assert model.n_periods == 3 + + @pytest.mark.long_running def test_benchmark_model_simulates_end_to_end() -> None: n_subjects = 20 diff --git a/tests/test_dcegm_model_creation.py b/tests/test_dcegm_model_creation.py index fc7645f..8dfa78d 100644 --- a/tests/test_dcegm_model_creation.py +++ b/tests/test_dcegm_model_creation.py @@ -29,7 +29,7 @@ get_benchmark_consumption_dollars_points, get_benchmark_params, ) -from aca_model.config import BENCHMARK_GRID_CONFIG +from aca_model.config import BENCHMARK_GRID_CONFIG, MODEL_AGES _FIXED_PARAMS, _WAGE_PARAMS, _ = get_benchmark_params(model=None) @@ -103,6 +103,7 @@ def test_dcegm_model_slots_broadcast_contract_functions_and_constraint() -> None fixed_params=_FIXED_PARAMS, wage_params=_WAGE_PARAMS, pref_type_grid=DiscreteGrid(BenchmarkPrefType), + ages=MODEL_AGES, solver="dcegm", ) assert "borrowing_constraint" in slots["constraints"] diff --git a/tests/test_di_medicare_transitions.py b/tests/test_di_medicare_transitions.py index 7c88074..8b128e7 100644 --- a/tests/test_di_medicare_transitions.py +++ b/tests/test_di_medicare_transitions.py @@ -27,7 +27,7 @@ _make_transition_canwork as retiree_canwork, ) from aca_model.baseline.regimes._tied import _make_transition_canwork as tied_canwork -from aca_model.config import MODEL_CONFIG +from aca_model.config import MODEL_AGES, MODEL_CONFIG N_PERIODS = MODEL_CONFIG.end_age - MODEL_CONFIG.start_age SURVIVAL = jnp.full((N_PERIODS, 3), 0.99) @@ -40,7 +40,7 @@ def _nongroup_probs(regime: str, health_state: int, labor: int) -> jnp.ndarray: own, _ = make_targets(regime) - return nongroup_canwork(own=own)( + return nongroup_canwork(own=own, ages=MODEL_AGES)( age=jnp.int32(55), period=jnp.int32(4), health=jnp.int32(health_state), @@ -112,7 +112,7 @@ def test_nongroup_transition_di_medicare_entry_and_exit( def test_retiree_transition_newly_disabled_non_worker_keeps_retiree_cover() -> None: """A retiree who stops working and becomes disabled holds EPHI-Medicare.""" own, ng = make_targets("retiree_nomc_inelig_canwork") - probs = retiree_canwork(own=own, ng=ng)( + probs = retiree_canwork(own=own, ng=ng, ages=MODEL_AGES)( age=jnp.int32(55), period=jnp.int32(4), health=jnp.int32(HealthWithDisability.good), @@ -129,7 +129,7 @@ def test_retiree_transition_newly_disabled_non_worker_keeps_retiree_cover() -> N def test_tied_transition_newly_disabled_non_worker_moves_to_nongroup_medicare() -> None: """A tied worker who stops and becomes disabled moves to non-group Medicare.""" own, ng = make_targets("tied_nomc_inelig_canwork") - probs = tied_canwork(own=own, ng=ng)( + probs = tied_canwork(own=own, ng=ng, ages=MODEL_AGES)( age=jnp.int32(55), period=jnp.int32(4), health=jnp.int32(HealthWithDisability.bad), diff --git a/tests/test_medical_cost_timing.py b/tests/test_medical_cost_timing.py index 2ebbc8d..f29f875 100644 --- a/tests/test_medical_cost_timing.py +++ b/tests/test_medical_cost_timing.py @@ -25,6 +25,9 @@ from aca_model.config import BENCHMARK_GRID_CONFIG _SHOCKS = ("hcc_persistent", "hcc_transitory") +# Period 0's choices look one year ahead, so ages 51-53 (two living periods and +# the age at which everybody is dead) carry the whole timing question. +_FINAL_AGE = 53 @pytest.fixture(scope="module") @@ -34,7 +37,9 @@ def period_0_by_shock() -> dict[tuple[str, str], pd.Series]: BENCHMARK_GRID_CONFIG, fold_hcc_transitory=False, subtract_hcc_bill=False ) model = create_benchmark_model( - pref_type_grid=DiscreteGrid(BenchmarkPrefType), grid_config=grid_config + pref_type_grid=DiscreteGrid(BenchmarkPrefType), + grid_config=grid_config, + final_age=_FINAL_AGE, ) _, _, params = get_benchmark_params(model=model) base = get_benchmark_initial_conditions( diff --git a/tests/test_mortality.py b/tests/test_mortality.py index 85bc6f5..68df8aa 100644 --- a/tests/test_mortality.py +++ b/tests/test_mortality.py @@ -28,7 +28,7 @@ _make_transition_forcedout as retiree_forcedout, ) from aca_model.baseline.regimes._tied import _make_transition_canwork as tied_canwork -from aca_model.config import MODEL_AGES, MODEL_CONFIG +from aca_model.config import MODEL_AGES, MODEL_CONFIG, model_ages N_PERIODS = MODEL_CONFIG.end_age - MODEL_CONFIG.start_age # Survival by (period, health): distinct per health state. @@ -39,7 +39,7 @@ def _tied(health: int) -> jnp.ndarray: own, ng = make_targets("tied_nomc_inelig_canwork") - return tied_canwork(own=own, ng=ng)( + return tied_canwork(own=own, ng=ng, ages=MODEL_AGES)( age=jnp.int32(55), period=jnp.int32(4), health=jnp.int32(health), @@ -52,7 +52,7 @@ def _tied(health: int) -> jnp.ndarray: def _retiree(health: int) -> jnp.ndarray: own, ng = make_targets("retiree_nomc_inelig_canwork") - return retiree_canwork(own=own, ng=ng)( + return retiree_canwork(own=own, ng=ng, ages=MODEL_AGES)( age=jnp.int32(55), period=jnp.int32(4), health=jnp.int32(health), @@ -65,7 +65,7 @@ def _retiree(health: int) -> jnp.ndarray: def _retiree_forcedout(health: int) -> jnp.ndarray: own, ng = make_targets("retiree_oamc_forced_forcedout") - return retiree_forcedout(gets_medicare=True, own=own, ng=ng)( + return retiree_forcedout(gets_medicare=True, own=own, ng=ng, ages=MODEL_AGES)( age=jnp.int32(80), period=jnp.int32(29), health=jnp.int32(health), @@ -76,7 +76,7 @@ def _retiree_forcedout(health: int) -> jnp.ndarray: def _nongroup(health: int) -> jnp.ndarray: own, _ng = make_targets("nongroup_nomc_inelig_canwork") - return nongroup_canwork(own=own)( + return nongroup_canwork(own=own, ages=MODEL_AGES)( age=jnp.int32(55), period=jnp.int32(4), health=jnp.int32(health), @@ -88,7 +88,7 @@ def _nongroup(health: int) -> jnp.ndarray: def _nongroup_forcedout(health: int) -> jnp.ndarray: own, _ng = make_targets("nongroup_oamc_forced_forcedout") - return nongroup_forcedout(gets_medicare=True, own=own)( + return nongroup_forcedout(gets_medicare=True, own=own, ages=MODEL_AGES)( age=jnp.int32(80), period=jnp.int32(29), health=jnp.int32(health), @@ -131,6 +131,31 @@ def test_forcedout_regimes_are_active_at_95(regime: str) -> None: assert transition_ages(REGIME_SPECS[regime])[-1] == 95 +def test_everybody_is_dead_at_a_shorter_final_age() -> None: + """A survivor of age 52 in a model ending at 53 moves to `dead`.""" + own, _ng = make_targets("retiree_nomc_inelig_canwork") + target = select_target_for_age(53, True, own, final_age=53) + assert int(target) == int(RegimeId.dead) + + +def test_entry_regime_has_no_transition_at_a_shorter_final_age() -> None: + """With ages 51-53 the entry regimes transition at 51 and 52 only.""" + spec = REGIME_SPECS["retiree_nomc_inelig_canwork"] + assert transition_ages(spec, ages=model_ages(final_age=53)) == (51, 52) + + +def test_model_ages_end_at_final_age() -> None: + """A final age of 53 gives the ages 51, 52 and 53.""" + assert model_ages(final_age=53) == (51, 52, 53) + + +@pytest.mark.parametrize("final_age", [51, 97]) +def test_model_ages_rejects_final_age_outside_the_data(final_age: int) -> None: + """The final age must leave a living period and lie within the data.""" + with pytest.raises(ValueError, match=f"got {final_age}"): + model_ages(final_age=final_age) + + def test_model_ages_run_from_51_through_96() -> None: """The age grid spans 51 to 96; 96 is the age at which everybody is dead.""" assert (MODEL_AGES[0], MODEL_AGES[-1], len(MODEL_AGES)) == (51, 96, 46) diff --git a/tests/test_nbegm_model_creation.py b/tests/test_nbegm_model_creation.py index 5394c25..d384885 100644 --- a/tests/test_nbegm_model_creation.py +++ b/tests/test_nbegm_model_creation.py @@ -40,6 +40,9 @@ from aca_model.benchmark import get_benchmark_params from aca_model.config import BENCHMARK_GRID_CONFIG, GridConfig +# NB-EGM is not a production solver; its wiring checks build the full model. +pytestmark = pytest.mark.long_running + _FIXED_PARAMS, _WAGE_PARAMS, _ = get_benchmark_params(model=None) # pylcm folds a shock state only under grid search, and the EGM solvers take the diff --git a/tests/test_regime_transitions.py b/tests/test_regime_transitions.py index f7baf11..e2f3c21 100644 --- a/tests/test_regime_transitions.py +++ b/tests/test_regime_transitions.py @@ -22,7 +22,7 @@ _make_transition_forcedout as retiree_forcedout, ) from aca_model.baseline.regimes._tied import _make_transition_canwork as tied_canwork -from aca_model.config import MODEL_CONFIG +from aca_model.config import MODEL_AGES, MODEL_CONFIG N_REGIMES = 19 N_PERIODS = MODEL_CONFIG.end_age - MODEL_CONFIG.start_age @@ -47,7 +47,7 @@ def _target_from_probs(probs: jnp.ndarray) -> int: def test_tied_stop_working_becomes_nongroup() -> None: """Tied agent who stops working loses employer coverage → nongroup.""" own, ng = make_targets("tied_nomc_inelig_canwork") - transition = tied_canwork(own=own, ng=ng) + transition = tied_canwork(own=own, ng=ng, ages=MODEL_AGES) probs = transition( age=jnp.int32(55), @@ -66,7 +66,7 @@ def test_tied_stop_working_becomes_nongroup() -> None: def test_tied_keeps_working_stays_tied() -> None: """Tied agent who keeps working retains employer coverage.""" own, ng = make_targets("tied_nomc_inelig_canwork") - transition = tied_canwork(own=own, ng=ng) + transition = tied_canwork(own=own, ng=ng, ages=MODEL_AGES) probs = transition( age=jnp.int32(55), @@ -87,7 +87,7 @@ def test_tied_keeps_working_stays_tied() -> None: def test_retiree_medicaid_override_to_nongroup() -> None: """Medicaid-eligible retiree is overridden to nongroup.""" own, ng = make_targets("retiree_nomc_inelig_canwork") - transition = retiree_canwork(own=own, ng=ng) + transition = retiree_canwork(own=own, ng=ng, ages=MODEL_AGES) probs = transition( age=jnp.int32(55), @@ -105,7 +105,7 @@ def test_retiree_medicaid_override_to_nongroup() -> None: def test_retiree_not_medicaid_stays_retiree() -> None: """Non-Medicaid retiree stays retiree.""" own, ng = make_targets("retiree_nomc_inelig_canwork") - transition = retiree_canwork(own=own, ng=ng) + transition = retiree_canwork(own=own, ng=ng, ages=MODEL_AGES) probs = transition( age=jnp.int32(55), @@ -123,7 +123,7 @@ def test_retiree_not_medicaid_stays_retiree() -> None: def test_retiree_forcedout_medicaid_override() -> None: """Forcedout retiree with Medicaid → nongroup.""" own, ng = make_targets("retiree_oamc_forced_forcedout") - transition = retiree_forcedout(gets_medicare=True, own=own, ng=ng) + transition = retiree_forcedout(gets_medicare=True, own=own, ng=ng, ages=MODEL_AGES) probs = transition( age=jnp.int32(80), @@ -162,7 +162,7 @@ def test_retiree_age_bracket_transitions( # Use nomc+inelig as starting point — the transition function resolves # the target based on next_age, not current spec. own, ng = make_targets("retiree_nomc_inelig_canwork") - transition = retiree_canwork(own=own, ng=ng) + transition = retiree_canwork(own=own, ng=ng, ages=MODEL_AGES) period = jnp.int32(age - MODEL_CONFIG.start_age) probs = transition( @@ -184,7 +184,7 @@ def test_retiree_age_bracket_transitions( def test_nongroup_canwork_valid_probs() -> None: """Nongroup canwork produces valid probability vector.""" own, _ng = make_targets("nongroup_nomc_inelig_canwork") - transition = nongroup_canwork(own=own) + transition = nongroup_canwork(own=own, ages=MODEL_AGES) probs = transition( age=jnp.int32(55), @@ -201,7 +201,7 @@ def test_nongroup_canwork_valid_probs() -> None: def test_nongroup_forcedout_valid_probs() -> None: """Nongroup forcedout produces valid probability vector.""" own, _ng = make_targets("nongroup_oamc_forced_forcedout") - transition = nongroup_forcedout(gets_medicare=True, own=own) + transition = nongroup_forcedout(gets_medicare=True, own=own, ages=MODEL_AGES) probs = transition( age=jnp.int32(80), @@ -219,7 +219,7 @@ def test_nongroup_forcedout_valid_probs() -> None: def test_tied_medicaid_override_to_nongroup() -> None: """Tied + Medicaid-eligible → nongroup (SSI override).""" own, ng = make_targets("tied_nomc_inelig_canwork") - transition = tied_canwork(own=own, ng=ng) + transition = tied_canwork(own=own, ng=ng, ages=MODEL_AGES) probs = transition( age=jnp.int32(55), @@ -237,7 +237,7 @@ def test_tied_medicaid_override_to_nongroup() -> None: def test_tied_at_medicare_age_with_medicaid() -> None: """Tied at age 64→65 (Medicare onset) + Medicaid → nongroup+oamc.""" own, ng = make_targets("tied_nomc_choose_canwork") - transition = tied_canwork(own=own, ng=ng) + transition = tied_canwork(own=own, ng=ng, ages=MODEL_AGES) period = jnp.int32(64 - MODEL_CONFIG.start_age) probs = transition( @@ -257,7 +257,7 @@ def test_tied_at_medicare_age_with_medicaid() -> None: def test_survival_prob_determines_death_weight() -> None: """Dead regime gets (1 - survival) probability weight.""" own, ng = make_targets("retiree_nomc_inelig_canwork") - transition = retiree_canwork(own=own, ng=ng) + transition = retiree_canwork(own=own, ng=ng, ages=MODEL_AGES) survival = jnp.ones((N_PERIODS, 3)) * 0.85 probs = transition( From 0d258a641dbbf82e74d416c3a279ecd6752f0f56 Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Fri, 23 Oct 2026 08:30:41 +0200 Subject: [PATCH 2/6] Pass a short-horizon model only the params it reads A horizon that ends before every regime is reached prunes those regimes from the params template, so their fixed params and the model-level params only they read are dropped before the model is built. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Hs8mfXPYyEa24KaAp6BwCW --- src/aca_model/aca/model.py | 42 +++++++++------- src/aca_model/baseline/model.py | 88 ++++++++++++++++++++++++++------- src/aca_model/benchmark.py | 9 ++-- 3 files changed, 101 insertions(+), 38 deletions(-) diff --git a/src/aca_model/aca/model.py b/src/aca_model/aca/model.py index 1d0943e..8e1dc04 100644 --- a/src/aca_model/aca/model.py +++ b/src/aca_model/aca/model.py @@ -12,7 +12,10 @@ from aca_model.aca import PolicyVariant 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.model import ( + _fail_if_dcegm_without_consumption_points, + build_model_for_horizon, +) from aca_model.baseline.regimes import ( RegimeId, SolverName, @@ -100,22 +103,25 @@ def create_model( solver=solver, ) - return Model( - regimes=regimes, - 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})", + return build_model_for_horizon( fixed_params=fixed_params, - derived_categoricals=derived_categoricals, - execution_config=( - execution_config_for_devices() - if execution_config is None - else execution_config - ), - **model_slots, + final_age=final_age, + model_kwargs={ + "regimes": regimes, + "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})", + "derived_categoricals": derived_categoricals, + "execution_config": ( + execution_config_for_devices() + if execution_config is None + else execution_config + ), + **model_slots, + }, ) diff --git a/src/aca_model/baseline/model.py b/src/aca_model/baseline/model.py index dd4e42e..e9081a0 100644 --- a/src/aca_model/baseline/model.py +++ b/src/aca_model/baseline/model.py @@ -13,6 +13,9 @@ from collections.abc import Mapping from typing import Any +from _lcm.params.processing import find_param_candidates +from _lcm.utils.namespace import flatten_regime_namespace +from dags.tree import unflatten_from_qnames from lcm import AgeGrid, DiscreteGrid, ExecutionConfig, InitialNodes, Model from lcm.typing import UserParams @@ -112,24 +115,75 @@ def create_model( solver=solver, ) - return Model( - regimes=regimes, - 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)", + return build_model_for_horizon( fixed_params=fixed_params, - derived_categoricals=derived_categoricals, - execution_config=( - execution_config_for_devices() - if execution_config is None - else execution_config - ), - **model_slots, + final_age=final_age, + model_kwargs={ + "regimes": regimes, + "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)", + "derived_categoricals": derived_categoricals, + "execution_config": ( + execution_config_for_devices() + if execution_config is None + else execution_config + ), + **model_slots, + }, + ) + + +def build_model_for_horizon( + *, + fixed_params: UserParams, + final_age: int, + model_kwargs: Mapping[str, Any], +) -> Model: + """Build the pylcm Model, passing it only the fixed params its horizon reads. + + The full horizon reaches every regime and reads every fixed param. A shorter + horizon reaches only the regimes whose stages start before `final_age`, so + the params of the other regimes, and model-level params only they read, are + dropped. A first build without fixed params supplies the template that + decides which params are read. + + Args: + fixed_params: Fixed params for the full horizon. + final_age: First age at which everybody is dead. + model_kwargs: Every other keyword argument of `lcm.Model`. + + Returns: + The pylcm Model. + + """ + if final_age < MODEL_CONFIG.end_age: + fixed_params = select_params_for_model( + params=fixed_params, model=Model(**model_kwargs) + ) + return Model(fixed_params=fixed_params, **model_kwargs) + + +def select_params_for_model(*, params: UserParams, model: Model) -> UserParams: + """Keep the entries of `params` that fill a slot of `model`'s params template. + + Each entry may sit at any level pylcm resolves: the exact slot, a regime, an + edge source, or the model level. + """ + params_flat = flatten_regime_namespace(params) + template_flat = flatten_regime_namespace(model.get_params_template()) + used = { + candidate + for qname in template_flat + for candidate in find_param_candidates(qname=qname, params_flat=params_flat) + } + return unflatten_from_qnames( + {qname: value for qname, value in params_flat.items() if qname in used} ) diff --git a/src/aca_model/benchmark.py b/src/aca_model/benchmark.py index 836ef07..f00a168 100644 --- a/src/aca_model/benchmark.py +++ b/src/aca_model/benchmark.py @@ -35,9 +35,9 @@ 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.model import create_model, select_params_for_model from aca_model.baseline.regimes import ENTRY_REGIMES -from aca_model.config import BENCHMARK_GRID_CONFIG, MODEL_CONFIG, GridConfig +from aca_model.config import BENCHMARK_GRID_CONFIG, MODEL_AGES, MODEL_CONFIG, GridConfig from aca_model.consumption_dollars_grid import ( compute_consumption_dollars_points, inject_consumption_dollars_points, @@ -116,7 +116,8 @@ def get_benchmark_params( into `params` for each regime that declares `consumption_dollars` as an `IrregSpacedGrid` with runtime-supplied points. Pass `model=None` to skip injection (e.g. when constructing the model with - `fixed_params`). + `fixed_params`). A model with a shorter horizon than the snapshot's + receives only the `params` its template reads. The regime-transition law arguments (`survival_probs`, `prob_disabled_next`) are returned under `fixed_params["edges"][source]`, @@ -135,6 +136,8 @@ def get_benchmark_params( max_consumption_dollars=max_consumption_dollars, ) params = inject_consumption_floor_schedule(params=params, model=model) + if model.n_periods < len(MODEL_AGES): + params = select_params_for_model(params=params, model=model) return fixed_params, wage_params, params From 9d8ace961aabebd49cfac1e98e8879b50361991b Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Fri, 23 Oct 2026 10:18:51 +0200 Subject: [PATCH 3/6] Share the standard test models per session The benchmark model, the baseline model on unfolded and folded grids, and the ACA model per policy are built once per pytest worker and shared by every test that builds them with the same arguments. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Hs8mfXPYyEa24KaAp6BwCW --- .github/workflows/main.yml | 2 +- tests/conftest.py | 55 ++++++++++++++++++ tests/test_additive_shock_medical_bill.py | 20 +++---- tests/test_benchmark.py | 41 ++++++------- tests/test_benchmark_folding.py | 17 ++---- tests/test_dag_coverage.py | 6 +- .../test_initial_conditions_extreme_assets.py | 12 ++-- tests/test_model_creation.py | 58 ++++++++++--------- tests/test_transition_only_medical_shock.py | 4 +- 9 files changed, 130 insertions(+), 85 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 0df1cb1..34fe18d 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -39,4 +39,4 @@ jobs: - name: Install aca-model with test deps run: pip install -e . pytest pytest-xdist pdbp - name: Run pytest - run: pytest -n 4 --durations=20 tests + run: pytest -n 4 --durations=40 tests diff --git a/tests/conftest.py b/tests/conftest.py index 1da1dcf..01cd35d 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,4 +1,59 @@ +"""Session-scoped models shared by every test that builds one with the same arguments. + +Under pytest-xdist each worker builds each model at most once. Tests read these +models and may solve or simulate them; none reassigns their attributes. +""" + +import functools import sys +from collections.abc import Callable +from dataclasses import replace from pathlib import Path sys.path.insert(0, str(Path(__file__).parent)) + +import pytest +from helpers.model import make_aca_model, make_baseline_model +from lcm import DiscreteGrid, Model + +from aca_model.aca.health_insurance import PolicyVariant +from aca_model.agent.preferences import BenchmarkPrefType +from aca_model.benchmark import create_benchmark_model +from aca_model.config import BENCHMARK_GRID_CONFIG + +# Benchmark grids with the transitory shock folded and the plain asset law. +FOLDED_GRID_CONFIG = replace( + BENCHMARK_GRID_CONFIG, fold_hcc_transitory=True, subtract_hcc_bill=False +) + + +@pytest.fixture(scope="session") +def benchmark_model() -> Model: + """`create_benchmark_model` with the 2-type preference grid and its defaults. + + The same model as `make_baseline_model(grid_config=BENCHMARK_GRID_CONFIG)`. + """ + return create_benchmark_model(pref_type_grid=DiscreteGrid(BenchmarkPrefType)) + + +@pytest.fixture(scope="session") +def baseline_model() -> Model: + """`make_baseline_model()`: the benchmark snapshot on unfolded grids.""" + return make_baseline_model() + + +@pytest.fixture(scope="session") +def folded_baseline_model() -> Model: + """`make_baseline_model` on `FOLDED_GRID_CONFIG`.""" + return make_baseline_model(grid_config=FOLDED_GRID_CONFIG) + + +@pytest.fixture(scope="session") +def aca_model_for() -> Callable[[PolicyVariant], Model]: + """Return `make_aca_model` for a policy, built once per policy.""" + + @functools.cache + def build(policy: PolicyVariant) -> Model: + return make_aca_model(policy=policy) + + return build diff --git a/tests/test_additive_shock_medical_bill.py b/tests/test_additive_shock_medical_bill.py index 9fff9d0..3d055eb 100644 --- a/tests/test_additive_shock_medical_bill.py +++ b/tests/test_additive_shock_medical_bill.py @@ -10,28 +10,26 @@ """ from collections.abc import Mapping -from dataclasses import replace import jax.numpy as jnp import lcm import numpy as np import pytest -from helpers.model import make_baseline_model from aca_model.agent import assets_and_income from aca_model.baseline.health_insurance import negative_oop_costs 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) - @pytest.fixture(scope="module") -def additive_model() -> lcm.Model: - return make_baseline_model(grid_config=_ADDITIVE) +def additive_model(benchmark_model: lcm.Model) -> lcm.Model: + """The benchmark grids fold the transitory shock and subtract its bill.""" + assert ( + BENCHMARK_GRID_CONFIG.fold_hcc_transitory, + BENCHMARK_GRID_CONFIG.subtract_hcc_bill, + ) == (True, True) + return benchmark_model def test_negative_oop_costs_is_the_negated_bill() -> None: @@ -55,10 +53,10 @@ def test_every_target_adds_the_negated_bill(additive_model: lcm.Model) -> None: @pytest.mark.long_running def test_additive_shock_value_equals_the_plain_law_value( - additive_model: lcm.Model, + additive_model: lcm.Model, folded_baseline_model: lcm.Model ) -> None: """The two routes sum the same terms in a different order.""" - folded_model = make_baseline_model(grid_config=_FOLDED) + folded_model = folded_baseline_model _, _, params = get_benchmark_params(model=folded_model) folded_values = folded_model.solve(params=params, log_level="off").values additive_values = additive_model.solve(params=params, log_level="off").values diff --git a/tests/test_benchmark.py b/tests/test_benchmark.py index 25d81ba..13c552c 100644 --- a/tests/test_benchmark.py +++ b/tests/test_benchmark.py @@ -7,6 +7,7 @@ DiscreteGrid, GridBreakpoint, IrregSpacedGrid, + Model, PiecewiseLinSpacedGrid, RouwenhorstAR1Process, ) @@ -20,11 +21,11 @@ ) -def test_benchmark_model_builds_with_the_current_pylcm_grid_api() -> None: +def test_benchmark_model_builds_with_the_current_pylcm_grid_api( + benchmark_model: Model, +) -> None: """The frozen benchmark preserves its state and action grid extents.""" - model = create_benchmark_model( - pref_type_grid=DiscreteGrid(BenchmarkPrefType), - ) + model = benchmark_model aime = model.user_regimes["retiree_nomc_inelig_canwork"].states["aime"] assert isinstance(aime, PiecewiseLinSpacedGrid) @@ -57,11 +58,9 @@ def test_benchmark_model_with_final_age_53_has_three_periods() -> None: @pytest.mark.long_running -def test_benchmark_model_simulates_end_to_end() -> None: +def test_benchmark_model_simulates_end_to_end(benchmark_model: Model) -> None: n_subjects = 20 - model = create_benchmark_model( - pref_type_grid=DiscreteGrid(BenchmarkPrefType), - ) + model = benchmark_model _, _, params = get_benchmark_params(model=model) initial_conditions = get_benchmark_initial_conditions( model=model, n_subjects=n_subjects, seed=0 @@ -81,7 +80,9 @@ def test_benchmark_model_simulates_end_to_end() -> None: @pytest.mark.long_running -def test_benchmark_panel_exposes_hic_premium_and_wage_targets() -> None: +def test_benchmark_panel_exposes_hic_premium_and_wage_targets( + benchmark_model: Model, +) -> None: """`hic_premium` and `wage` are computable simulate targets. The MSM estimation matches `hic_private_premium_*` and `wage_growth_*` @@ -92,9 +93,7 @@ def test_benchmark_panel_exposes_hic_premium_and_wage_targets() -> None: target get NaN, alive/worker rows carry a finite value. """ n_subjects = 20 - model = create_benchmark_model( - pref_type_grid=DiscreteGrid(BenchmarkPrefType), - ) + model = benchmark_model _, _, params = get_benchmark_params(model=model) initial_conditions = get_benchmark_initial_conditions( model=model, n_subjects=n_subjects, seed=0 @@ -121,7 +120,9 @@ 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: +def test_benchmark_model_offers_uncompensated_care_targets( + benchmark_model: Model, +) -> None: """Transfers, the floor and the health-cost covariates can be added to the panel. Uncompensated care is the part of `transfers` above @@ -129,7 +130,7 @@ def test_benchmark_model_offers_uncompensated_care_targets() -> None: 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)) + model = benchmark_model available = _collect_all_available_targets(model._regimes) # noqa: SLF001 assert { "transfers", @@ -140,7 +141,7 @@ def test_benchmark_model_offers_uncompensated_care_targets() -> None: @pytest.mark.long_running -def test_benchmark_simulate_obeys_borrowing_constraint() -> None: +def test_benchmark_simulate_obeys_borrowing_constraint(benchmark_model: Model) -> None: """`consumption_dollars <= max(cash_on_hand, floor)` holds for every alive row. The simulator only ever picks feasible actions — the borrowing @@ -171,9 +172,7 @@ def test_benchmark_simulate_obeys_borrowing_constraint() -> None: consumption point, not the floor. """ n_subjects = 4 - model = create_benchmark_model( - pref_type_grid=DiscreteGrid(BenchmarkPrefType), - ) + model = benchmark_model _, _, params = get_benchmark_params(model=model) initial_conditions = get_benchmark_initial_conditions( model=model, n_subjects=n_subjects, seed=0 @@ -197,9 +196,11 @@ def test_benchmark_simulate_obeys_borrowing_constraint() -> None: ) -def test_initial_conditions_choose_population_size_after_model_construction() -> None: +def test_initial_conditions_choose_population_size_after_model_construction( + benchmark_model: Model, +) -> None: """One model supports reproducible initial conditions with different row counts.""" - model = create_benchmark_model(pref_type_grid=DiscreteGrid(BenchmarkPrefType)) + model = benchmark_model small = get_benchmark_initial_conditions(model=model, n_subjects=2, seed=17) repeated = get_benchmark_initial_conditions(model=model, n_subjects=2, seed=17) large = get_benchmark_initial_conditions(model=model, n_subjects=5, seed=17) diff --git a/tests/test_benchmark_folding.py b/tests/test_benchmark_folding.py index 71c096f..f320bbf 100644 --- a/tests/test_benchmark_folding.py +++ b/tests/test_benchmark_folding.py @@ -5,19 +5,11 @@ 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.benchmark import get_benchmark_initial_conditions from aca_model.config import BENCHMARK_GRID_CONFIG -@pytest.fixture(scope="module") -def benchmark_model() -> lcm.Model: - return create_benchmark_model(pref_type_grid=DiscreteGrid(BenchmarkPrefType)) - - def test_benchmark_folds_the_transitory_medical_bill( benchmark_model: lcm.Model, ) -> None: @@ -40,14 +32,13 @@ def test_benchmark_folds_the_transitory_medical_bill( def test_unfolded_control_has_identical_shared_initial_conditions( - benchmark_model: lcm.Model, + benchmark_model: lcm.Model, baseline_model: lcm.Model ) -> None: + """The unfolded control is the benchmark built on unfolded grids.""" 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 - ) + control = baseline_model assert all( "hcc_transitory" in control.state_names(regime_name=name) for name in control.user_regimes diff --git a/tests/test_dag_coverage.py b/tests/test_dag_coverage.py index 16b6472..9d01198 100644 --- a/tests/test_dag_coverage.py +++ b/tests/test_dag_coverage.py @@ -3,7 +3,7 @@ import jax.numpy as jnp import pandas as pd import pytest -from helpers.model import make_baseline_model +from lcm import Model from lcm.exceptions import InvalidInitialConditionsError from aca_model.baseline.regimes import ENTRY_REGIMES, REGIME_SPECS @@ -14,8 +14,8 @@ @pytest.fixture(scope="module") -def model(): - return make_baseline_model() +def model(baseline_model: Model) -> Model: + return baseline_model def test_initial_nodes_are_the_age_51_to_60_entry_regimes(model): diff --git a/tests/test_initial_conditions_extreme_assets.py b/tests/test_initial_conditions_extreme_assets.py index 6455bee..028928e 100644 --- a/tests/test_initial_conditions_extreme_assets.py +++ b/tests/test_initial_conditions_extreme_assets.py @@ -9,12 +9,10 @@ import jax.numpy as jnp from _lcm.simulation.initial_conditions import validate_initial_conditions -from lcm import DiscreteGrid +from lcm import Model from aca_model.agent.assets_and_income import borrowing_constraint -from aca_model.agent.preferences import BenchmarkPrefType from aca_model.benchmark import ( - create_benchmark_model, get_benchmark_initial_conditions, get_benchmark_params, ) @@ -92,7 +90,9 @@ def test_borrowing_constraint_admits_floor_at_million_dollar_negative_cash() -> assert admitted -def test_extreme_negative_assets_subject_passes_validation() -> None: +def test_extreme_negative_assets_subject_passes_validation( + benchmark_model: Model, +) -> None: """A subject placed at `assets = -1_000_000` clears initial-conditions validation. A large-but-reasonable negative value (very bad draws for both HCC shocks) @@ -100,9 +100,7 @@ def test_extreme_negative_assets_subject_passes_validation() -> None: transfer system absorbs them, with `c = c_floor` always feasible. """ n_subjects = 1 - model = create_benchmark_model( - pref_type_grid=DiscreteGrid(BenchmarkPrefType), - ) + model = benchmark_model _, _, params = get_benchmark_params(model=model) initial_conditions = get_benchmark_initial_conditions( diff --git a/tests/test_model_creation.py b/tests/test_model_creation.py index 790a98f..6b56b25 100644 --- a/tests/test_model_creation.py +++ b/tests/test_model_creation.py @@ -1,14 +1,10 @@ """Tests for baseline model creation and regime structure.""" import inspect -from collections.abc import Mapping +from collections.abc import Callable, Mapping import pytest -from helpers.model import ( - make_aca_model, - make_baseline_model, -) -from lcm import DiscreteGrid, Phased +from lcm import DiscreteGrid, Model, Phased from aca_model.aca import health_insurance as aca_hi from aca_model.aca.health_insurance import PolicyVariant @@ -54,25 +50,25 @@ def build_regime(name: str): return _build_regime(name, _GRIDS) -def test_model_creates_successfully() -> None: - model = make_baseline_model() +def test_model_creates_successfully(baseline_model: Model) -> None: + model = baseline_model assert len(model.user_regimes) == 19 assert model.n_periods == 46 -def test_model_age_range() -> None: - model = make_baseline_model() +def test_model_age_range(baseline_model: Model) -> None: + model = baseline_model assert model.ages.values[0] == 51.0 assert model.ages.values[-1] == 96.0 -def test_dead_regime_is_terminal() -> None: - model = make_baseline_model() +def test_dead_regime_is_terminal(baseline_model: Model) -> None: + model = baseline_model assert model.graph.laws["dead"].terminal -def test_non_terminal_regimes_not_terminal() -> None: - model = make_baseline_model() +def test_non_terminal_regimes_not_terminal(baseline_model: Model) -> None: + model = baseline_model for name in REGIME_SPECS: assert not model.graph.laws[name].terminal @@ -226,7 +222,9 @@ def test_regime_specs_keys_match_regime_id() -> None: assert hasattr(RegimeId, name), f"RegimeId missing field for {name}" -def test_all_non_terminal_regimes_carry_pension_wealth_as_carried_state() -> None: +def test_all_non_terminal_regimes_carry_pension_wealth_as_carried_state( + baseline_model: Model, +) -> None: """`pension_wealth` is a carried state broadcast into every living regime. Imputed from AIME during solve (never a solve grid axis) yet seeded and @@ -241,7 +239,7 @@ def test_all_non_terminal_regimes_carry_pension_wealth_as_carried_state() -> Non build_model_state_transitions()["pension_wealth"] is pensions.wealth_next_before_adjustment ) - model = make_baseline_model() + model = baseline_model for name in REGIME_SPECS: assert isinstance(model.user_regimes[name].states["pension_wealth"], Phased), ( name @@ -288,8 +286,10 @@ def test_hcc_persistent_and_transitory_are_shock_grids() -> None: assert "hcc_transitory" not in regime.state_transitions -def test_aca_model_creates_successfully() -> None: - model = make_aca_model(policy=PolicyVariant.ACA) +def test_aca_model_creates_successfully( + aca_model_for: Callable[[PolicyVariant], Model], +) -> None: + model = aca_model_for(PolicyVariant.ACA) assert len(model.user_regimes) == 19 assert model.n_periods == 46 @@ -362,9 +362,11 @@ def test_aca_other_regimes_have_no_aca_policy_keys() -> None: @pytest.mark.parametrize("policy", list(PolicyVariant)) -def test_all_policy_variants_create(policy: PolicyVariant) -> None: +def test_all_policy_variants_create( + policy: PolicyVariant, aca_model_for: Callable[[PolicyVariant], Model] +) -> None: """All policy variants create valid models.""" - model = make_aca_model(policy=policy) + model = aca_model_for(policy) assert len(model.user_regimes) == 19 @@ -402,35 +404,35 @@ def test_aca_only_medicaid_expansion() -> None: assert "cost_sharing_scale" not in nongroup.functions -def test_baseline_model_creates() -> None: +def test_baseline_model_creates(baseline_model: Model) -> None: """Baseline model creates successfully without PolicyVariant.""" - model = make_baseline_model() + model = baseline_model assert len(model.user_regimes) == 19 -def test_dead_regime_prunes_unused_broadcast_states() -> None: +def test_dead_regime_prunes_unused_broadcast_states(baseline_model: Model) -> None: """States every living regime shares are declared once at the model level; `dead` keeps what the bequest DAG reads (`assets`, `pref_type`). `pension_wealth` is masked (carried states are illegal in terminal regimes); the other unused broadcast states are pruned by reachability.""" - model = make_baseline_model() + model = baseline_model assert model.pruned_variables["dead"] == frozenset( {"aime", "spousal_income", "hcc_persistent", "hcc_transitory"} ) -def test_dead_regime_keeps_only_the_bequest_states() -> None: +def test_dead_regime_keeps_only_the_bequest_states(baseline_model: Model) -> None: """`dead` carries only assets and preference type for valuing bequests.""" - model = make_baseline_model() + model = baseline_model assert set(model.user_regimes["dead"].states) == { "assets", "pref_type", } -def test_living_regimes_keep_every_broadcast_state() -> None: +def test_living_regimes_keep_every_broadcast_state(baseline_model: Model) -> None: """Every model-level state is read by each living regime's DAG, so pruning removes nothing outside `dead`.""" - model = make_baseline_model() + model = baseline_model for name in REGIME_SPECS: assert model.pruned_variables[name] == frozenset() diff --git a/tests/test_transition_only_medical_shock.py b/tests/test_transition_only_medical_shock.py index a4a148f..8d48535 100644 --- a/tests/test_transition_only_medical_shock.py +++ b/tests/test_transition_only_medical_shock.py @@ -24,8 +24,8 @@ @pytest.fixture(scope="module") -def folded_model() -> Model: - return make_baseline_model(grid_config=_FOLDED) +def folded_model(folded_baseline_model: Model) -> Model: + return folded_baseline_model def test_living_regimes_carry_no_lagged_transitory_shock(folded_model: Model) -> None: From 5de722da8b6b643de9bbd307fa8ca650a9ce10f2 Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Fri, 23 Oct 2026 10:42:35 +0200 Subject: [PATCH 4/6] Shrink the medical-cost timing test to the smallest grids The AIME points per PIA segment become a GridConfig field with the current counts as default. The timing test keeps the two health-cost shock grids and shrinks AIME, wage residual and consumption. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Hs8mfXPYyEa24KaAp6BwCW --- src/aca_model/baseline/regimes/_common.py | 23 ++++++++++------------ src/aca_model/config.py | 5 +++++ tests/test_aime_grid.py | 24 +++++++++++++---------- tests/test_medical_cost_timing.py | 21 +++++++++++++++++--- 4 files changed, 47 insertions(+), 26 deletions(-) diff --git a/src/aca_model/baseline/regimes/_common.py b/src/aca_model/baseline/regimes/_common.py index cf7eaa3..a4df020 100644 --- a/src/aca_model/baseline/regimes/_common.py +++ b/src/aca_model/baseline/regimes/_common.py @@ -252,13 +252,6 @@ class Grids: """Grid sizes and numerical settings used to construct each solver.""" -# AIME piecewise grid: number of points per segment between the PIA -# bend points (0 → kink_0 → kink_1 → kink_2 → extension). The fourth -# segment covers the sparse delayed-retirement-credit region above the -# taxable max, so it carries fewer points than the dense lower segments. -_AIME_PIECE_N_POINTS: tuple[int, int, int, int] = (10, 11, 11, 6) - - # `pension_wealth` is a carried state (`Phased(solve=..., simulate=Grid)`): # imputed via a derived function in solve (never a solve grid axis) and # carried per subject in simulate. The grid is simulate-side metadata only — @@ -337,7 +330,10 @@ def build_grids( assets=build_assets_grid( floor=assets_start, n_points=grid_config.n_assets_gridpoints ), - aime=_build_aime_grid(fixed_params=fixed_params), + aime=_build_aime_grid( + fixed_params=fixed_params, + points_per_segment=grid_config.aime_points_per_segment, + ), pension_wealth=_PENSION_WEALTH_GRID, consumption_dollars=( IrregSpacedGrid(n_points=grid_config.n_consumption_dollars_gridpoints) @@ -402,15 +398,16 @@ def get_hcc_persistent_grid_points(*, grid_config: GridConfig) -> FloatND: return get_hcc_persistent_shock(grid_config=grid_config).to_jax() -def _build_aime_grid(*, fixed_params: UserParams) -> PiecewiseLinSpacedGrid: +def _build_aime_grid( + *, fixed_params: UserParams, points_per_segment: tuple[int, int, int, int] +) -> PiecewiseLinSpacedGrid: """Return the AIME grid. The grid is piecewise-linspaced with breakpoints at the PIA bends - in `fixed_params["pia_aime_grid"]` and `_AIME_PIECE_N_POINTS` in + in `fixed_params["pia_aime_grid"]` and `points_per_segment` in each segment. The fifth bend point is the delayed-retirement-credit extension above the taxable max, so the grid carries four segments. - `n_aime_gridpoints` from `grid_config` is ignored on this path; the - total is fixed by the PIA structure (`sum(_AIME_PIECE_N_POINTS)`). + `n_aime_gridpoints` from `grid_config` is ignored on this path. """ kinks = [float(k) for k in np.asarray(fixed_params["pia_aime_grid"])] return PiecewiseLinSpacedGrid( @@ -421,7 +418,7 @@ def _build_aime_grid(*, fixed_params: UserParams) -> PiecewiseLinSpacedGrid: GridBreakpoint(value=kinks[2]), GridBreakpoint(value=kinks[3]), ), - points_per_segment=_AIME_PIECE_N_POINTS, + points_per_segment=points_per_segment, ) diff --git a/src/aca_model/config.py b/src/aca_model/config.py index cb2c720..cdd3387 100644 --- a/src/aca_model/config.py +++ b/src/aca_model/config.py @@ -30,6 +30,11 @@ class GridConfig: n_assets_gridpoints: int = 24 n_aime_gridpoints: int = 12 + # AIME piecewise grid: number of points per segment between the PIA bend + # points (0 → kink_0 → kink_1 → kink_2 → extension), at least 2 each. The + # fourth segment covers the sparse delayed-retirement-credit region above + # the taxable max, so it carries fewer points than the dense lower segments. + aime_points_per_segment: tuple[int, int, int, int] = (10, 11, 11, 6) n_consumption_dollars_gridpoints: int = 70 n_wage_res_gridpoints: int = 5 n_hcc_persistent_gridpoints: int = 3 diff --git a/tests/test_aime_grid.py b/tests/test_aime_grid.py index d63da86..7a976bb 100644 --- a/tests/test_aime_grid.py +++ b/tests/test_aime_grid.py @@ -5,10 +5,10 @@ import jax.numpy as jnp import numpy as np -from aca_model.baseline.regimes._common import ( - _AIME_PIECE_N_POINTS, - _build_aime_grid, -) +from aca_model.baseline.regimes._common import _build_aime_grid +from aca_model.config import GridConfig + +_POINTS_PER_SEGMENT = GridConfig().aime_points_per_segment # Production SSA bend points plus the delayed-retirement-credit extension: # 0, kink_0, kink_1, taxable-max, and the extension point that carries the @@ -19,21 +19,25 @@ def test_build_aime_grid_owns_each_pia_breakpoint_on_the_right() -> None: """The AIME grid preserves four PIA pieces with right-owned interiors.""" - grid = _build_aime_grid(fixed_params=_FIXED_PARAMS) + grid = _build_aime_grid( + fixed_params=_FIXED_PARAMS, points_per_segment=_POINTS_PER_SEGMENT + ) np.testing.assert_allclose( [point.value for point in grid.breakpoints], _PIA_AIME_GRID[1:-1], ) assert tuple(point.owner for point in grid.breakpoints) == ("right",) * 3 - assert grid.points_per_segment == _AIME_PIECE_N_POINTS + assert grid.points_per_segment == _POINTS_PER_SEGMENT def test_build_aime_grid_top_point_is_extension_aime() -> None: """The grid reaches the delayed-credit extension AIME at its top.""" - grid = _build_aime_grid(fixed_params=_FIXED_PARAMS) + grid = _build_aime_grid( + fixed_params=_FIXED_PARAMS, points_per_segment=_POINTS_PER_SEGMENT + ) np.testing.assert_allclose(float(grid.to_jax().max()), 187954.752, rtol=1e-5) -def test_aime_piece_n_points_has_four_entries() -> None: - """One point-count per segment, including the sparse extension region.""" - assert len(_AIME_PIECE_N_POINTS) == 4 +def test_aime_points_per_segment_default_is_10_11_11_6() -> None: + """One point-count per segment; the extension region is the sparsest.""" + assert _POINTS_PER_SEGMENT == (10, 11, 11, 6) diff --git a/tests/test_medical_cost_timing.py b/tests/test_medical_cost_timing.py index f29f875..b72c12c 100644 --- a/tests/test_medical_cost_timing.py +++ b/tests/test_medical_cost_timing.py @@ -28,14 +28,29 @@ # Period 0's choices look one year ahead, so ages 51-53 (two living periods and # the age at which everybody is dead) carry the whole timing question. _FINAL_AGE = 53 +# The twins differ in one lagged shock at its grid ends, so only the two shock +# grids keep the benchmark resolution and the transitory shock stays a state. +# Every other grid is at the smallest size the model accepts, except that +# consumption keeps an interior point above the two floors and below the cap, +# so the consumption comparison still allows a choice that cash would move: +# - assets: borrowing limit, zero and the top holding; +# - AIME: two points per PIA segment; +# - wage residual: two points. +_GRID_CONFIG = replace( + BENCHMARK_GRID_CONFIG, + fold_hcc_transitory=False, + subtract_hcc_bill=False, + n_assets_gridpoints=3, + aime_points_per_segment=(2, 2, 2, 2), + n_consumption_dollars_gridpoints=4, + n_wage_res_gridpoints=2, +) @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 = _GRID_CONFIG model = create_benchmark_model( pref_type_grid=DiscreteGrid(BenchmarkPrefType), grid_config=grid_config, From b3fb3258f8423290cbf4d9b3c6ed418fb1360ad0 Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Fri, 23 Oct 2026 11:37:09 +0200 Subject: [PATCH 5/6] Write short-horizon params by the reachability rule The model takes fixed params as given. `select_horizon_params` keeps the regime-level entries of the regimes the age stages reach, and the claiming-only params when a claiming regime is reachable. The benchmark params follow it, so a short horizon builds once. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Hs8mfXPYyEa24KaAp6BwCW --- src/aca_model/aca/model.py | 45 +++++------ src/aca_model/baseline/model.py | 91 +++++----------------- src/aca_model/baseline/regimes/__init__.py | 4 + src/aca_model/baseline/regimes/_common.py | 39 ++++++++++ src/aca_model/benchmark.py | 33 ++++---- tests/test_medical_cost_timing.py | 2 +- tests/test_mortality.py | 40 ++++++++++ 7 files changed, 143 insertions(+), 111 deletions(-) diff --git a/src/aca_model/aca/model.py b/src/aca_model/aca/model.py index 8e1dc04..e6b2022 100644 --- a/src/aca_model/aca/model.py +++ b/src/aca_model/aca/model.py @@ -12,10 +12,7 @@ from aca_model.aca import PolicyVariant from aca_model.aca.regimes import build_all_regimes -from aca_model.baseline.model import ( - _fail_if_dcegm_without_consumption_points, - build_model_for_horizon, -) +from aca_model.baseline.model import _fail_if_dcegm_without_consumption_points from aca_model.baseline.regimes import ( RegimeId, SolverName, @@ -73,7 +70,8 @@ def create_model( 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. + data, and `fixed_params` must hold only what it reads (see + `select_horizon_params`). Returns: pylcm Model. @@ -103,25 +101,22 @@ def create_model( solver=solver, ) - return build_model_for_horizon( + return Model( + regimes=regimes, + 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, - final_age=final_age, - model_kwargs={ - "regimes": regimes, - "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})", - "derived_categoricals": derived_categoricals, - "execution_config": ( - execution_config_for_devices() - if execution_config is None - else execution_config - ), - **model_slots, - }, + derived_categoricals=derived_categoricals, + execution_config=( + execution_config_for_devices() + if execution_config is None + else execution_config + ), + **model_slots, ) diff --git a/src/aca_model/baseline/model.py b/src/aca_model/baseline/model.py index e9081a0..78909f4 100644 --- a/src/aca_model/baseline/model.py +++ b/src/aca_model/baseline/model.py @@ -13,9 +13,6 @@ from collections.abc import Mapping from typing import Any -from _lcm.params.processing import find_param_candidates -from _lcm.utils.namespace import flatten_regime_namespace -from dags.tree import unflatten_from_qnames from lcm import AgeGrid, DiscreteGrid, ExecutionConfig, InitialNodes, Model from lcm.typing import UserParams @@ -85,7 +82,8 @@ def create_model( 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. + 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 @@ -115,75 +113,24 @@ def create_model( solver=solver, ) - return build_model_for_horizon( + return Model( + regimes=regimes, + 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, - final_age=final_age, - model_kwargs={ - "regimes": regimes, - "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)", - "derived_categoricals": derived_categoricals, - "execution_config": ( - execution_config_for_devices() - if execution_config is None - else execution_config - ), - **model_slots, - }, - ) - - -def build_model_for_horizon( - *, - fixed_params: UserParams, - final_age: int, - model_kwargs: Mapping[str, Any], -) -> Model: - """Build the pylcm Model, passing it only the fixed params its horizon reads. - - The full horizon reaches every regime and reads every fixed param. A shorter - horizon reaches only the regimes whose stages start before `final_age`, so - the params of the other regimes, and model-level params only they read, are - dropped. A first build without fixed params supplies the template that - decides which params are read. - - Args: - fixed_params: Fixed params for the full horizon. - final_age: First age at which everybody is dead. - model_kwargs: Every other keyword argument of `lcm.Model`. - - Returns: - The pylcm Model. - - """ - if final_age < MODEL_CONFIG.end_age: - fixed_params = select_params_for_model( - params=fixed_params, model=Model(**model_kwargs) - ) - return Model(fixed_params=fixed_params, **model_kwargs) - - -def select_params_for_model(*, params: UserParams, model: Model) -> UserParams: - """Keep the entries of `params` that fill a slot of `model`'s params template. - - Each entry may sit at any level pylcm resolves: the exact slot, a regime, an - edge source, or the model level. - """ - params_flat = flatten_regime_namespace(params) - template_flat = flatten_regime_namespace(model.get_params_template()) - used = { - candidate - for qname in template_flat - for candidate in find_param_candidates(qname=qname, params_flat=params_flat) - } - return unflatten_from_qnames( - {qname: value for qname, value in params_flat.items() if qname in used} + derived_categoricals=derived_categoricals, + execution_config=( + execution_config_for_devices() + if execution_config is None + else execution_config + ), + **model_slots, ) diff --git a/src/aca_model/baseline/regimes/__init__.py b/src/aca_model/baseline/regimes/__init__.py index 495d4be..296058d 100644 --- a/src/aca_model/baseline/regimes/__init__.py +++ b/src/aca_model/baseline/regimes/__init__.py @@ -34,6 +34,8 @@ 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 @@ -49,7 +51,9 @@ "build_regime", "build_restart_initial_nodes", "nongroup", + "reachable_regimes", "retiree", + "select_horizon_params", "tied", ] diff --git a/src/aca_model/baseline/regimes/_common.py b/src/aca_model/baseline/regimes/_common.py index a4df020..c16e0bf 100644 --- a/src/aca_model/baseline/regimes/_common.py +++ b/src/aca_model/baseline/regimes/_common.py @@ -530,6 +530,45 @@ def transition_ages( raise ValueError(f"Unknown regime spec: {spec}") +# Model-level params that only the regimes able to claim Social Security read. +CLAIMING_PARAMS = frozenset( + { + "early_ret_adjustment", + "earnings_test_fraction", + "earnings_test_threshold", + "normal_retirement_age", + } +) + + +def reachable_regimes(*, ages: tuple[int, ...]) -> tuple[RegimeName, ...]: + """Return the living regimes a model over `ages` can reach. + + A regime is reachable when its stage covers a source age before the final + age, the same rule that gives it transition ages and edges. + """ + return tuple( + name for name, spec in REGIME_SPECS.items() if transition_ages(spec, ages=ages) + ) + + +def select_horizon_params[V]( + *, params: Mapping[str, V], ages: tuple[int, ...] +) -> dict[str, V]: + """Return the entries of `params` that a model over `ages` reads. + + `params` is written for the full horizon. The regime-level entries of + regimes `reachable_regimes` excludes are left out, and so are + `CLAIMING_PARAMS` when no reachable regime can claim Social Security. + Every entry is kept for the full horizon. + """ + reachable = reachable_regimes(ages=ages) + unreachable = set(REGIME_SPECS) - set(reachable) + if not any(REGIME_SPECS[name]["ss"] != "inelig" for name in reachable): + unreachable |= CLAIMING_PARAMS + return {key: value for key, value in params.items() if key not in unreachable} + + def build_restart_initial_nodes(*, ages: Iterable[int]) -> InitialNodes: """Admit every living regime active at each of `ages` as a start. diff --git a/src/aca_model/benchmark.py b/src/aca_model/benchmark.py index f00a168..9b32db6 100644 --- a/src/aca_model/benchmark.py +++ b/src/aca_model/benchmark.py @@ -35,9 +35,14 @@ 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, select_params_for_model -from aca_model.baseline.regimes import ENTRY_REGIMES -from aca_model.config import BENCHMARK_GRID_CONFIG, MODEL_AGES, MODEL_CONFIG, GridConfig +from aca_model.baseline.model import create_model +from aca_model.baseline.regimes import ENTRY_REGIMES, select_horizon_params +from aca_model.config import ( + BENCHMARK_GRID_CONFIG, + MODEL_CONFIG, + GridConfig, + model_ages, +) from aca_model.consumption_dollars_grid import ( compute_consumption_dollars_points, inject_consumption_dollars_points, @@ -85,11 +90,11 @@ def create_benchmark_model( initial_nodes: Admissible starting `(age, regime)` pairs, forwarded to `create_model`. final_age: First age at which everybody is dead, forwarded to - `create_model`. The default is the full horizon; the frozen - age-indexed parameters cover it, and a shorter horizon reads a - prefix of them. + `create_model` and `get_benchmark_params`. The default is the full + horizon; the frozen age-indexed parameters cover it, and a shorter + horizon reads a prefix of them. """ - fixed_params, wage_params, _ = get_benchmark_params(model=None) + fixed_params, wage_params, _ = get_benchmark_params(model=None, final_age=final_age) return create_model( grid_config=grid_config, fixed_params=fixed_params, @@ -103,7 +108,7 @@ def create_benchmark_model( def get_benchmark_params( - *, model: Model | None + *, model: Model | None, final_age: int = MODEL_CONFIG.end_age ) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]: """Load the frozen `(fixed_params, wage_params, params)` snapshot. @@ -116,8 +121,10 @@ def get_benchmark_params( into `params` for each regime that declares `consumption_dollars` as an `IrregSpacedGrid` with runtime-supplied points. Pass `model=None` to skip injection (e.g. when constructing the model with - `fixed_params`). A model with a shorter horizon than the snapshot's - receives only the `params` its template reads. + `fixed_params`). + + `fixed_params` and `params` hold only what a model ending at `final_age` + reads (`select_horizon_params`); pass the same `final_age` as the model's. The regime-transition law arguments (`survival_probs`, `prob_disabled_next`) are returned under `fixed_params["edges"][source]`, @@ -129,6 +136,8 @@ def get_benchmark_params( wage_params = data["wage_params"] params = data["params"] max_consumption_dollars = float(fixed_params.pop("max_consumption_dollars")) + ages = model_ages(final_age=final_age) + fixed_params = select_horizon_params(params=fixed_params, ages=ages) if model is not None: params = inject_consumption_dollars_points( params=params, @@ -136,9 +145,7 @@ def get_benchmark_params( max_consumption_dollars=max_consumption_dollars, ) params = inject_consumption_floor_schedule(params=params, model=model) - if model.n_periods < len(MODEL_AGES): - params = select_params_for_model(params=params, model=model) - return fixed_params, wage_params, params + return fixed_params, wage_params, select_horizon_params(params=params, ages=ages) def get_benchmark_consumption_dollars_points(*, n_points: int) -> tuple[float, ...]: diff --git a/tests/test_medical_cost_timing.py b/tests/test_medical_cost_timing.py index b72c12c..3a5945a 100644 --- a/tests/test_medical_cost_timing.py +++ b/tests/test_medical_cost_timing.py @@ -56,7 +56,7 @@ def period_0_by_shock() -> dict[tuple[str, str], pd.Series]: grid_config=grid_config, final_age=_FINAL_AGE, ) - _, _, params = get_benchmark_params(model=model) + _, _, params = get_benchmark_params(model=model, final_age=_FINAL_AGE) base = get_benchmark_initial_conditions( model=model, n_subjects=1, seed=0, grid_config=grid_config ) diff --git a/tests/test_mortality.py b/tests/test_mortality.py index 68df8aa..1a8dc20 100644 --- a/tests/test_mortality.py +++ b/tests/test_mortality.py @@ -9,9 +9,12 @@ from aca_model.agent.labor_market import LaborSupply from aca_model.baseline.regimes._common import ( + ENTRY_REGIMES, REGIME_SPECS, RegimeId, make_targets, + reachable_regimes, + select_horizon_params, select_target_for_age, transition_ages, ) @@ -144,6 +147,43 @@ def test_entry_regime_has_no_transition_at_a_shorter_final_age() -> None: assert transition_ages(spec, ages=model_ages(final_age=53)) == (51, 52) +def test_only_entry_regimes_are_reachable_at_final_age_53() -> None: + """Ages 51-53 reach only the five pre-claiming, pre-Medicare regimes.""" + assert set(reachable_regimes(ages=model_ages(final_age=53))) == set(ENTRY_REGIMES) + + +@pytest.mark.parametrize( + ("final_age", "expected"), + [ + (53, {"retiree_nomc_inelig_canwork", "medicare_age"}), + ( + 63, + { + "retiree_nomc_inelig_canwork", + "retiree_nomc_choose_canwork", + "medicare_age", + "normal_retirement_age", + }, + ), + ], +) +def test_select_horizon_params_keeps_what_the_horizon_reaches( + final_age: int, expected: set[str] +) -> None: + """Unreachable regimes' entries, and claiming params before 62, are left out.""" + params = { + "retiree_nomc_inelig_canwork": 1, + "retiree_nomc_choose_canwork": 2, + "retiree_oamc_forced_forcedout": 3, + "medicare_age": 4, + "normal_retirement_age": 5, + } + selected = select_horizon_params( + params=params, ages=model_ages(final_age=final_age) + ) + assert set(selected) == expected + + def test_model_ages_end_at_final_age() -> None: """A final age of 53 gives the ages 51, 52 and 53.""" assert model_ages(final_age=53) == (51, 52, 53) From e511da010e98491d747b9170784f9e4fb53c2299 Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Fri, 23 Oct 2026 12:05:49 +0200 Subject: [PATCH 6/6] Keep each test module on one xdist worker in CI `--dist loadscope` sends a module's tests to one worker, so a module-scoped fixture is built once. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Hs8mfXPYyEa24KaAp6BwCW --- .github/workflows/main.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 34fe18d..c8b887f 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -39,4 +39,4 @@ jobs: - name: Install aca-model with test deps run: pip install -e . pytest pytest-xdist pdbp - name: Run pytest - run: pytest -n 4 --durations=40 tests + run: pytest -n 4 --dist loadscope --durations=40 tests