diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 48163c1..71f5307 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -35,7 +35,7 @@ jobs: # install; keep it compatible with the model's graph declarations. run: >- pip install "pylcm @ - git+https://github.com/OpenSourceEconomics/pylcm.git@eea73a8f0e301e88f02278df3595c797fdba24df" + git+https://github.com/OpenSourceEconomics/pylcm.git@7e99a19badafe20e29a517f69d9ec155bba8e0a0" - name: Install aca-model with test deps run: pip install -e . pytest pdbp - name: Run pytest diff --git a/src/aca_model/aca/model.py b/src/aca_model/aca/model.py index 5b17c97..971b8b0 100644 --- a/src/aca_model/aca/model.py +++ b/src/aca_model/aca/model.py @@ -7,7 +7,7 @@ from collections.abc import Mapping from typing import Any -from lcm import AgeGrid, DiscreteGrid, ExecutionConfig, Model +from lcm import AgeGrid, DiscreteGrid, ExecutionConfig, InitialNodes, Model from lcm.typing import UserParams from aca_model.aca import PolicyVariant @@ -34,6 +34,7 @@ def create_model( solver: SolverName = "brute_force", consumption_dollars_points: tuple[float, ...] | None = None, execution_config: ExecutionConfig | None = None, + initial_nodes: InitialNodes = INITIAL_NODES, ) -> Model: """Create an ACA policy variant model. @@ -61,6 +62,9 @@ def create_model( execution_config: Explicit hardware-local policy forwarded unchanged. None uses the smallest selected accelerator allocator limit as the device-memory budget; CPU construction remains unbudgeted. + initial_nodes: Admissible starting `(age, regime)` pairs. The default + admits the estimation sample's entry ages and regimes; pass + `build_restart_initial_nodes` to restart simulations mid-life. Returns: pylcm Model. @@ -91,7 +95,7 @@ def create_model( return Model( regimes=regimes, - initial_nodes=INITIAL_NODES, + initial_nodes=initial_nodes, ages=ages, regime_id_class=RegimeId, description=f"Structural retirement model ({policy.name})", diff --git a/src/aca_model/aca/regimes/_overrides.py b/src/aca_model/aca/regimes/_overrides.py index eb92ffe..23cbd77 100644 --- a/src/aca_model/aca/regimes/_overrides.py +++ b/src/aca_model/aca/regimes/_overrides.py @@ -8,6 +8,7 @@ from aca_model.aca import health_insurance as aca_hi from aca_model.aca.health_insurance import PolicyVariant +from aca_model.baseline import health_insurance as baseline_hi from aca_model.baseline.regimes._common import RegimeSpec from aca_model.environment import taxes @@ -55,6 +56,7 @@ def apply_aca_overrides( ) if has_medicaid_expansion: + functions["aca_magi"] = baseline_hi.aca_magi functions["is_medicaid_eligible"] = aca_hi.is_medicaid_eligible if has_subsidies and spec["his"] == "nongroup" and spec["mc"] == "nomc": @@ -64,8 +66,9 @@ def apply_aca_overrides( # dict, not a DAG function — no entry needed here. functions["plan_premium"] = aca_hi.community_rated_premium # The community-rated premium replaces the risk-rated one, so the - # risk-rated premium's zero-profit intercept has no consumer left. + # risk-rated inputs have no consumer left. functions.pop("private_premium_intercept", None) + functions.pop("predicted_hcc_insurer", None) functions["hic_premium_subsidy"] = aca_hi.premium_subsidy # Split at the bracket, which takes finitely many values, so the # out-of-pocket bill can be averaged per bracket. diff --git a/src/aca_model/baseline/model.py b/src/aca_model/baseline/model.py index 3522e30..1b1efe3 100644 --- a/src/aca_model/baseline/model.py +++ b/src/aca_model/baseline/model.py @@ -13,7 +13,7 @@ from collections.abc import Mapping from typing import Any -from lcm import AgeGrid, DiscreteGrid, ExecutionConfig, Model +from lcm import AgeGrid, DiscreteGrid, ExecutionConfig, InitialNodes, Model from lcm.typing import UserParams from aca_model.baseline.regimes import ( @@ -37,6 +37,7 @@ def create_model( solver: SolverName = "brute_force", consumption_dollars_points: tuple[float, ...] | None = None, execution_config: ExecutionConfig | None = None, + initial_nodes: InitialNodes = INITIAL_NODES, ) -> Model: """Create the baseline structural retirement model. @@ -73,6 +74,9 @@ def create_model( execution_config: Explicit hardware-local policy forwarded unchanged. None uses the smallest selected accelerator allocator limit as the device-memory budget; CPU construction remains unbudgeted. + initial_nodes: Admissible starting `(age, regime)` pairs. The default + admits the estimation sample's entry ages and regimes; pass + `build_restart_initial_nodes` to restart simulations mid-life. Returns: A pylcm Model with 19 regimes (18 non-terminal + dead) spanning @@ -102,7 +106,7 @@ def create_model( return Model( regimes=regimes, - initial_nodes=INITIAL_NODES, + initial_nodes=initial_nodes, ages=ages, regime_id_class=RegimeId, description="Baseline structural retirement model (pre-ACA)", diff --git a/src/aca_model/baseline/regimes/__init__.py b/src/aca_model/baseline/regimes/__init__.py index bd83aed..65b1f19 100644 --- a/src/aca_model/baseline/regimes/__init__.py +++ b/src/aca_model/baseline/regimes/__init__.py @@ -33,6 +33,7 @@ build_model_functions, build_model_state_transitions, build_model_states, + build_restart_initial_nodes, ) from aca_model.baseline.regimes._dcegm import build_dcegm_solver from aca_model.baseline.regimes._nbegm import build_nbegm_solver @@ -46,6 +47,7 @@ "build_all_regimes", "build_model_slots", "build_regime", + "build_restart_initial_nodes", "nongroup", "retiree", "tied", diff --git a/src/aca_model/baseline/regimes/_common.py b/src/aca_model/baseline/regimes/_common.py index 2032ce1..c74812b 100644 --- a/src/aca_model/baseline/regimes/_common.py +++ b/src/aca_model/baseline/regimes/_common.py @@ -342,7 +342,7 @@ def build_grids( consumption_dollars=( IrregSpacedGrid(n_points=grid_config.n_consumption_dollars_gridpoints) if consumption_dollars_points is None - else IrregSpacedGrid(points=consumption_dollars_points) + else IrregSpacedGrid(points=np.asarray(consumption_dollars_points)) ), wage_res=wage_res, hcc_persistent=hcc_persistent, @@ -378,7 +378,7 @@ def build_assets_grid(*, floor: float, n_points: int) -> IrregSpacedGrid: ) negative[0] = floor positive[-1] = _ASSETS_TOP - return IrregSpacedGrid(points=np.concatenate([negative, positive]).tolist()) + return IrregSpacedGrid(points=np.concatenate([negative, positive])) def get_hcc_persistent_shock(*, grid_config: GridConfig) -> RouwenhorstAR1Process: @@ -526,6 +526,24 @@ def transition_ages(spec: RegimeSpec) -> tuple[int, ...]: raise ValueError(f"Unknown regime spec: {spec}") +def build_restart_initial_nodes(*, ages: Iterable[int]) -> InitialNodes: + """Admit every living regime active at each of `ages` as a start. + + A simulation restarted from the states another simulation reached at some + age starts in whichever living regime that simulation had arrived at. + """ + return InitialNodes( + by_age={ + age: tuple( + name + for name, spec in REGIME_SPECS.items() + if age in transition_ages(spec) + ) + for age in ages + } + ) + + def next_model_age(period: Period) -> FloatND: """Advance by one grid position, never by an assumed number of years. @@ -845,12 +863,10 @@ def build_model_functions(*, solver: SolverName = "brute_force") -> dict: # assets adjustment, both of which scale by the marginal income tax rate. functions["marginal_tax_rate"] = taxes.marginal_rate - # HIC premium - functions["predicted_hcc_insurer"] = health_insurance.hcc_insurer_predicted - # Transfers functions["consumption_dollars_floor"] = assets_and_income.consumption_dollars_floor - functions["transfers"] = assets_and_income.transfers + if solver == "brute_force": + functions["transfers"] = assets_and_income.transfers functions["consumption_equiv"] = preferences.consumption_equiv return functions @@ -999,10 +1015,8 @@ def build_common_functions(spec: RegimeSpec, *, subtract_bill: bool = False) -> if spec["mc"] == "oamc" else health_insurance.is_disabled_from_health ) - # MAGI for the ACA Medicaid expansion track; pruned in the baseline DAG. - functions["aca_magi"] = health_insurance.aca_magi - # Swapped per policy variant by the ACA overlay, hence regime-level + functions["predicted_hcc_insurer"] = health_insurance.hcc_insurer_predicted functions["is_medicaid_eligible"] = health_insurance.is_medicaid_eligible functions["after_tax_income"] = taxes.after_tax_income if subtract_bill: diff --git a/src/aca_model/baseline/regimes/_dcegm.py b/src/aca_model/baseline/regimes/_dcegm.py index 54e3ab2..588cbae 100644 --- a/src/aca_model/baseline/regimes/_dcegm.py +++ b/src/aca_model/baseline/regimes/_dcegm.py @@ -11,6 +11,7 @@ the one piece that needs `lcm.solvers`. """ +import numpy as np from lcm import IrregSpacedGrid from lcm.solvers import DCEGM @@ -31,7 +32,9 @@ def build_dcegm_solver(grids: Grids) -> DCEGM: assets_points = grids.assets.to_jax() savings_stop = float(assets_points[-1]) - float(assets_points[0]) savings_grid = IrregSpacedGrid( - points=tuple(savings_stop * (i / (n_points - 1)) ** 3 for i in range(n_points)), + points=np.array( + [savings_stop * (i / (n_points - 1)) ** 3 for i in range(n_points)] + ), ) return DCEGM(savings_grid=savings_grid) diff --git a/src/aca_model/baseline/regimes/_nbegm.py b/src/aca_model/baseline/regimes/_nbegm.py index 23fb96a..1007ed0 100644 --- a/src/aca_model/baseline/regimes/_nbegm.py +++ b/src/aca_model/baseline/regimes/_nbegm.py @@ -4,6 +4,7 @@ Device placement and compiled program widths belong to the model execution policy. """ +import numpy as np from lcm import IrregSpacedGrid from lcm.solvers import NBEGM @@ -17,7 +18,9 @@ def build_nbegm_solver(grids: Grids) -> NBEGM: assets_points = grids.assets.to_jax() savings_stop = float(assets_points[-1]) - float(assets_points[0]) savings_grid = IrregSpacedGrid( - points=tuple(savings_stop * (i / (n_points - 1)) ** 3 for i in range(n_points)), + points=np.array( + [savings_stop * (i / (n_points - 1)) ** 3 for i in range(n_points)] + ), ) return NBEGM( savings_grid=savings_grid, diff --git a/src/aca_model/benchmark.py b/src/aca_model/benchmark.py index 3eed460..c3ca289 100644 --- a/src/aca_model/benchmark.py +++ b/src/aca_model/benchmark.py @@ -29,14 +29,14 @@ import jax.numpy as jnp import numpy as np from jax import Array -from lcm import DiscreteGrid, ExecutionConfig, Model +from lcm import DiscreteGrid, ExecutionConfig, InitialNodes, Model from aca_model.agent.health import GoodHealth from aca_model.agent.labor_market import IsMarried from aca_model.agent.preferences import BenchmarkPrefType from aca_model.baseline.health_insurance import HealthInsuranceState from aca_model.baseline.model import create_model -from aca_model.baseline.regimes import ENTRY_REGIMES +from aca_model.baseline.regimes import ENTRY_REGIMES, INITIAL_NODES from aca_model.config import BENCHMARK_GRID_CONFIG, GridConfig from aca_model.consumption_dollars_grid import ( compute_consumption_dollars_points, @@ -68,6 +68,7 @@ def create_benchmark_model( pref_type_grid: DiscreteGrid, grid_config: GridConfig = BENCHMARK_GRID_CONFIG, execution_config: ExecutionConfig | None = None, + initial_nodes: InitialNodes = INITIAL_NODES, ) -> Model: """Create the aca baseline with `BENCHMARK_GRID_CONFIG` and frozen fixed_params. @@ -80,6 +81,8 @@ def create_benchmark_model( execution_config: Explicit hardware-local policy forwarded unchanged. None uses the smallest selected accelerator allocator limit as the device-memory budget; CPU construction remains unbudgeted. + initial_nodes: Admissible starting `(age, regime)` pairs, forwarded + to `create_model`. """ fixed_params, wage_params, _ = get_benchmark_params(model=None) return create_model( @@ -89,6 +92,7 @@ def create_benchmark_model( derived_categoricals=_DERIVED_CATEGORICALS, pref_type_grid=pref_type_grid, execution_config=execution_config, + initial_nodes=initial_nodes, ) diff --git a/tests/test_aca_plan.py b/tests/test_aca_plan.py index c8b0e3c..8686faa 100644 --- a/tests/test_aca_plan.py +++ b/tests/test_aca_plan.py @@ -83,6 +83,21 @@ def test_aca_non_group_premium_is_community_rated( assert _aca_premium(policy, spousal_income) == pytest.approx(expected) +@pytest.mark.parametrize("policy", list(PolicyVariant)) +def test_aca_non_group_risk_prediction_is_registered_only_without_reform( + policy: PolicyVariant, +) -> None: + """The community-rated premium does not register a risk prediction node.""" + functions = {**build_model_functions(), **_build_functions(REGIME_SPECS[REGIME])} + apply_aca_overrides(functions, REGIME_SPECS[REGIME], policy) + expected = ( + health_insurance.hcc_insurer_predicted + if policy == PolicyVariant.ACA_ONLY_MEDICAID_EXPANSION + else None + ) + assert functions.get("predicted_hcc_insurer") is expected + + RISK_RATED_PREMIUM_PARAMS = frozenset( { "premium_markup", diff --git a/tests/test_benchmark.py b/tests/test_benchmark.py index 74b3af8..6e3a970 100644 --- a/tests/test_benchmark.py +++ b/tests/test_benchmark.py @@ -2,6 +2,7 @@ import numpy as np import pytest +from _lcm.simulation.additional_targets import _collect_all_available_targets from lcm import ( DiscreteGrid, GridBreakpoint, @@ -112,6 +113,24 @@ 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: + """Transfers, the floor and the health-cost covariates can be added to the panel. + + Uncompensated care is the part of `transfers` above + `consumption_dollars_floor`. Total health costs depend on the next period's + 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)) + available = _collect_all_available_targets(model._regimes) # noqa: SLF001 + assert { + "transfers", + "consumption_dollars_floor", + "is_married", + "good_health", + } <= available + + @pytest.mark.long_running def test_benchmark_simulate_obeys_borrowing_constraint() -> None: """`consumption_dollars <= max(cash_on_hand, floor)` holds for every alive row. diff --git a/tests/test_model_creation.py b/tests/test_model_creation.py index 9940ffd..790a98f 100644 --- a/tests/test_model_creation.py +++ b/tests/test_model_creation.py @@ -13,11 +13,15 @@ from aca_model.aca import health_insurance as aca_hi from aca_model.aca.health_insurance import PolicyVariant from aca_model.aca.regimes import build_all_regimes as _build_aca_regimes +from aca_model.agent import assets_and_income from aca_model.agent.preferences import BenchmarkPrefType +from aca_model.baseline import health_insurance as baseline_hi from aca_model.baseline.regimes import REGIME_SPECS, RegimeId from aca_model.baseline.regimes import build_regime as _build_regime from aca_model.baseline.regimes._common import ( + SolverName, build_grids, + build_model_functions, build_model_state_transitions, build_model_states, ) @@ -290,6 +294,36 @@ def test_aca_model_creates_successfully() -> None: assert model.n_periods == 46 +@pytest.mark.parametrize("name", REGIME_SPECS) +def test_baseline_regime_omits_expansion_magi(name: str) -> None: + """Baseline eligibility does not register the ACA expansion income node.""" + assert "aca_magi" not in build_regime(name).functions + + +@pytest.mark.parametrize("policy", list(PolicyVariant)) +def test_aca_regime_registers_magi_only_for_expansion(policy: PolicyVariant) -> None: + """The expansion eligibility consumer owns its MAGI income node.""" + has_expansion = policy in ( + PolicyVariant.ACA, + PolicyVariant.ACA_NO_MANDATE, + PolicyVariant.ACA_ONLY_MEDICAID_EXPANSION, + ) + expected = baseline_hi.aca_magi if has_expansion else None + regimes = build_aca_regimes(policy) + assert { + name: regimes[name].functions.get("aca_magi") for name in REGIME_SPECS + } == dict.fromkeys(REGIME_SPECS, expected) + + +@pytest.mark.parametrize("solver", ["brute_force", "dcegm", "nbegm"]) +def test_model_registers_transfers_only_for_brute_asset_laws( + solver: SolverName, +) -> None: + """Savings-form asset laws already include transfers through resources.""" + expected = assets_and_income.transfers if solver == "brute_force" else None + assert build_model_functions(solver=solver).get("transfers") is expected + + def test_aca_nongroup_inelig_has_real_functions() -> None: """Nongroup+nomc regimes get real ACA functions under ACA policy.""" regimes = build_aca_regimes(PolicyVariant.ACA) diff --git a/tests/test_restart_initial_nodes.py b/tests/test_restart_initial_nodes.py new file mode 100644 index 0000000..0fa69ea --- /dev/null +++ b/tests/test_restart_initial_nodes.py @@ -0,0 +1,56 @@ +"""Starts mid-life, at the regimes a simulation can have reached by that age.""" + +import pandas as pd +import pytest +from lcm import DiscreteGrid + +from aca_model.agent.preferences import BenchmarkPrefType +from aca_model.baseline.regimes import build_restart_initial_nodes +from aca_model.benchmark import create_benchmark_model +from aca_model.simulation import select_admissible_starts + + +@pytest.mark.parametrize( + ("age", "expected"), + [ + ( + 61, + ( + "retiree_nomc_inelig_canwork", + "tied_nomc_inelig_canwork", + "nongroup_nomc_inelig_canwork", + "retiree_dimc_inelig_canwork", + "nongroup_dimc_inelig_canwork", + ), + ), + ( + 66, + ( + "retiree_oamc_choose_canwork", + "tied_oamc_choose_canwork", + "nongroup_oamc_choose_canwork", + ), + ), + ( + 73, + ("retiree_oamc_forced_forcedout", "nongroup_oamc_forced_forcedout"), + ), + ], +) +def test_build_restart_initial_nodes_admits_every_living_regime_of_the_age( + age: int, expected: tuple[str, ...] +) -> None: + """At each age, the starts are exactly the living regimes active then.""" + nodes = build_restart_initial_nodes(ages=(age,)) + assert set(nodes.by_age[age]) == set(expected) + + +def test_model_with_restart_nodes_admits_a_start_on_old_age_medicare() -> None: + """A model built with restart nodes accepts a start at 66 on Medicare.""" + model = create_benchmark_model( + pref_type_grid=DiscreteGrid(BenchmarkPrefType), + initial_nodes=build_restart_initial_nodes(ages=range(51, 70)), + ) + starts = pd.DataFrame({"age": [66.0], "regime_name": ["tied_oamc_choose_canwork"]}) + admitted = select_admissible_starts(model=model, initial_conditions=starts) + assert admitted["regime_name"].tolist() == ["tied_oamc_choose_canwork"]