diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 71f5307..c8b887f 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 --dist loadscope --durations=40 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..e6b2022 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,21 @@ def create_model( execution_config: Explicit hardware-local policy forwarded unchanged. None uses the smallest selected accelerator allocator limit as the device-memory budget; CPU construction remains unbudgeted. - initial_nodes: Admissible starting `(age, regime)` pairs. The default - admits the estimation sample's entry ages and regimes; pass - `build_restart_initial_nodes` to restart simulations mid-life. + initial_nodes: Admissible starting `(age, regime)` pairs. None admits + the estimation sample's entry ages and regimes that lie before + `final_age`; pass `build_restart_initial_nodes` to restart + simulations mid-life. + final_age: First age at which everybody is dead; the model spans + `MODEL_CONFIG.start_age` through `final_age`. The default is the + full horizon. A shorter horizon reads a prefix of the age-indexed + data, and `fixed_params` must hold only what it reads (see + `select_horizon_params`). Returns: pylcm Model. """ - ages = AgeGrid(exact_values=MODEL_AGES) + ages = model_ages(final_age=final_age) _fail_if_dcegm_without_consumption_points( solver=solver, consumption_dollars_points=consumption_dollars_points ) @@ -90,13 +97,18 @@ def create_model( fixed_params=fixed_params, wage_params=wage_params, pref_type_grid=pref_type_grid, + ages=ages, solver=solver, ) return Model( regimes=regimes, - initial_nodes=initial_nodes, - ages=ages, + initial_nodes=( + build_entry_initial_nodes(ages=ages) + if initial_nodes is None + else initial_nodes + ), + ages=AgeGrid(exact_values=ages), regime_id_class=RegimeId, description=f"Structural retirement model ({policy.name})", fixed_params=fixed_params, diff --git a/src/aca_model/baseline/model.py b/src/aca_model/baseline/model.py index 1b1efe3..78909f4 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,24 @@ def create_model( execution_config: Explicit hardware-local policy forwarded unchanged. None uses the smallest selected accelerator allocator limit as the device-memory budget; CPU construction remains unbudgeted. - initial_nodes: Admissible starting `(age, regime)` pairs. The default - admits the estimation sample's entry ages and regimes; pass - `build_restart_initial_nodes` to restart simulations mid-life. + initial_nodes: Admissible starting `(age, regime)` pairs. None admits + the estimation sample's entry ages and regimes that lie before + `final_age`; pass `build_restart_initial_nodes` to restart + simulations mid-life. + final_age: First age at which everybody is dead; the model spans + `MODEL_CONFIG.start_age` through `final_age`. The default is the + full horizon. A shorter horizon reads a prefix of the age-indexed + data, and `fixed_params` must hold only what it reads (see + `select_horizon_params`). Returns: A pylcm Model with 19 regimes (18 non-terminal + dead) spanning - ages 51-96; 95 is the last age anyone is alive. Regime names follow + ages 51 through `final_age` (96 by default); `final_age - 1` is the + last age anyone is alive. Regime names follow the `___` scheme. """ - ages = AgeGrid(exact_values=MODEL_AGES) + ages = model_ages(final_age=final_age) _fail_if_dcegm_without_consumption_points( solver=solver, consumption_dollars_points=consumption_dollars_points ) @@ -101,13 +109,18 @@ def create_model( fixed_params=fixed_params, wage_params=wage_params, pref_type_grid=pref_type_grid, + ages=ages, solver=solver, ) return Model( regimes=regimes, - initial_nodes=initial_nodes, - ages=ages, + initial_nodes=( + build_entry_initial_nodes(ages=ages) + if initial_nodes is None + else initial_nodes + ), + ages=AgeGrid(exact_values=ages), regime_id_class=RegimeId, description="Baseline structural retirement model (pre-ACA)", fixed_params=fixed_params, diff --git a/src/aca_model/baseline/regimes/__init__.py b/src/aca_model/baseline/regimes/__init__.py index 65b1f19..296058d 100644 --- a/src/aca_model/baseline/regimes/__init__.py +++ b/src/aca_model/baseline/regimes/__init__.py @@ -22,18 +22,20 @@ from aca_model.baseline.regimes import _tied as tied from aca_model.baseline.regimes._common import ( ENTRY_REGIMES, - INITIAL_NODES, REGIME_SPECS, Grids, RegimeId, SolverName, build_dead_regime, + build_entry_initial_nodes, build_grids, build_model_constraints, build_model_functions, build_model_state_transitions, build_model_states, build_restart_initial_nodes, + reachable_regimes, + select_horizon_params, ) from aca_model.baseline.regimes._dcegm import build_dcegm_solver from aca_model.baseline.regimes._nbegm import build_nbegm_solver @@ -41,15 +43,17 @@ __all__ = [ "ENTRY_REGIMES", - "INITIAL_NODES", "REGIME_SPECS", "RegimeId", "build_all_regimes", + "build_entry_initial_nodes", "build_model_slots", "build_regime", "build_restart_initial_nodes", "nongroup", + "reachable_regimes", "retiree", + "select_horizon_params", "tied", ] @@ -134,6 +138,7 @@ def build_model_slots( fixed_params: UserParams, wage_params: Mapping[str, Any], pref_type_grid: DiscreteGrid, + ages: tuple[int, ...], solver: SolverName = "brute_force", ) -> dict[str, Any]: """Build the shared model topology and broadcast regime slots. @@ -144,6 +149,7 @@ def build_model_slots( overlay swaps only regime-level functions, so the broadcast slots are policy-invariant. Under DC-EGM the solver-contract functions join the broadcast set; the borrowing constraint is declared under every solver. + `ages` is the model's age grid, which the edges' age schedules cover. """ grids = build_grids( grid_config=grid_config, @@ -152,7 +158,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 +166,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..c16e0bf 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, ) @@ -460,30 +457,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 +488,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,19 +514,61 @@ 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}") +# 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. @@ -544,21 +587,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 +614,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 +623,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 +636,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 +651,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 +1241,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..9b32db6 100644 --- a/src/aca_model/benchmark.py +++ b/src/aca_model/benchmark.py @@ -36,8 +36,13 @@ 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, 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, @@ -68,7 +73,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,8 +89,12 @@ 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` 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, @@ -93,11 +103,12 @@ def create_benchmark_model( pref_type_grid=pref_type_grid, execution_config=execution_config, initial_nodes=initial_nodes, + final_age=final_age, ) 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. @@ -112,6 +123,9 @@ def get_benchmark_params( to skip injection (e.g. when constructing the model with `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]`, where pylcm looks for the parameters of edge-declared callables. @@ -122,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, @@ -129,7 +145,7 @@ def get_benchmark_params( max_consumption_dollars=max_consumption_dollars, ) params = inject_consumption_floor_schedule(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/src/aca_model/config.py b/src/aca_model/config.py index d95569b..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 @@ -59,11 +64,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/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_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_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_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_benchmark.py b/tests/test_benchmark.py index 6e3a970..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) @@ -48,12 +49,18 @@ def test_benchmark_model_builds_with_the_current_pylcm_grid_api() -> None: assert consumption.n_points == 5 -@pytest.mark.long_running -def test_benchmark_model_simulates_end_to_end() -> None: - n_subjects = 20 +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), + 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(benchmark_model: Model) -> None: + n_subjects = 20 + model = benchmark_model _, _, params = get_benchmark_params(model=model) initial_conditions = get_benchmark_initial_conditions( model=model, n_subjects=n_subjects, seed=0 @@ -73,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_*` @@ -84,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 @@ -113,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 @@ -121,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", @@ -132,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 @@ -163,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 @@ -189,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_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_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_medical_cost_timing.py b/tests/test_medical_cost_timing.py index 2ebbc8d..3a5945a 100644 --- a/tests/test_medical_cost_timing.py +++ b/tests/test_medical_cost_timing.py @@ -25,18 +25,38 @@ 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 +# 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 + pref_type_grid=DiscreteGrid(BenchmarkPrefType), + 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_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_mortality.py b/tests/test_mortality.py index 85bc6f5..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, ) @@ -28,7 +31,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 +42,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 +55,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 +68,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 +79,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 +91,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 +134,68 @@ 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_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) + + +@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( 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: