diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 632668d..48163c1 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -27,15 +27,15 @@ jobs: - '3.14' steps: - uses: actions/checkout@v7 - - uses: actions/setup-python@v6 + - uses: actions/setup-python@v7 with: python-version: ${{ matrix.python-version }} - - name: Install pylcm from feat/nb-egm - # This slice depends on the NB-EGM solver, which lives on pylcm's - # feat/nb-egm branch; revert to @main once that branch merges. + - name: Install pylcm at the supported API revision + # Pinned by commit so a deleted or rewritten branch cannot break the + # install; keep it compatible with the model's graph declarations. run: >- pip install "pylcm @ - git+https://github.com/OpenSourceEconomics/pylcm.git@feat/nb-egm" + git+https://github.com/OpenSourceEconomics/pylcm.git@eea73a8f0e301e88f02278df3595c797fdba24df" - name: Install aca-model with test deps run: pip install -e . pytest pdbp - name: Run pytest diff --git a/AGENTS.md b/AGENTS.md index a11a970..c17aa67 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -72,27 +72,41 @@ ACA variants don't create new regimes — they swap functions on baseline regime - ACA replaces stubs with real policy functions from `aca/policies.py` - `PolicyVariant` enum controls which policies are active (full ACA, no mandate, Medicaid only) +- Every variant with the reformed non-group market (all but Medicaid-only) swaps the + non-group-nomc `plan_premium` for the community-rated + `aca_premium_single + aca_premium_married_extra * is_married`; aca-slurm puts the ACA + plan's copay parameters into those regimes and drops the risk-rated premium's params ### Key State Variables -- `assets`: Savings grid `[≈ −221k, 500k]`, 24 points (lower bound = minus one year of - maximum full-time earnings, so it shifts with the wage parameters; finer at low - levels) +- `assets`: Savings grid from minus one year of maximum full-time earnings (so the lower + bound shifts with the wage parameters) to 12M (the largest initial holding), 24 + points, sinh-spaced around zero (5 nodes below zero, a node at zero, 5 nodes in + `[0, 30k]`, geometric above; `build_assets_grid`) - `aime`: Average Indexed Monthly Earnings — piecewise grid at the PIA bend points (32 points total; `n_aime_gridpoints` is ignored on this path) - `health`: `HealthWithDisability` (disabled/bad/good) pre-65, `Health` (bad/good) post-65 - `log_ft_wage_res`: AR(1) wage residual shock (5-point Rouwenhorst) -- `hcc_persistent` / `hcc_transitory`: Health cost shocks (`_ShockGrid` — integrated - over, policy does not condition on them) -- Regime transitions determined by `select_target_for_age()` based on age and actions +- `hcc_persistent` / `hcc_transitory`: last period's health-cost shocks (Rouwenhorst + AR(1) and Gauss-Hermite IID). Policy conditions on the lagged persistent shock, which + predicts this period's costs and prices the private premium; nothing reads the lagged + transitory shock. The draw dependencies retain both states in living regimes; `dead` + carries only assets and preference type. +- Regime transitions determined by `select_target_for_age()` based on age and actions; + death has probability `1 − survival_probs[age, health]`, a per-regime param on the + regime's own health grid. Ages run 51–96: 95 is the last age anyone is alive (survival + 1e-4 at 94, 0 at 95), and only `dead` is active at 96. ### Key Design Decisions - **AgeGrid with int ages**: Model uses `AgeGrid` with integer start/stop — all `age` parameters are `int`, not `float` -- **OOP timing**: OOP health costs are in `next_assets` (post-consumption), not - `cash_on_hand`. Matches struct-ret: agent doesn't condition consumption on OOP. +- **OOP timing**: medical costs are realised after the period's choices, and the state + carries the lagged persistent shock ζ\_{t−1}. `total_health_costs` reads the + transition's draws `next_hcc_persistent` / `next_hcc_transitory`, so OOP enters + `next_assets` and the continuation integrates over it; `cash_on_hand`, transfers and + the consumption floor see no OOP. - **DAG key ≠ function name**: `functions["marginal_tax_rate"] = taxes.marginal_rate` — the dict key must match consuming functions' parameter names, not the definition name. Don't stutter module name in function name (`taxes.marginal_rate` not @@ -107,6 +121,13 @@ ACA variants don't create new regimes — they swap functions on baseline regime an action. Those use `premium()` and `primary_oop()` which condition on it. All other regimes use `premium_insured()` / `premium_retired()` and `oop_costs` directly — no `buy_private` parameter at all. +- **Medicaid in the non-group regimes**: their `hic_premium` is + `medicaid_adjusted_premium(plan_premium)`, zero when Medicaid-eligible (Medicaid + replaces private cover before Medicare and pays the Medicare premium after). Before + Medicare, `oop_costs` is Medicaid's cost-sharing on total costs when eligible, + whatever `buy_private` says, so a Medicaid-eligible household can neither hold private + cover nor stay uninsured; both `buy_private` values give the same outcome. Elsewhere + Medicaid pays on top of the primary OOP (`oop_with_medicaid`). - **`reference_age` parameter**: Fixed cost of work uses `age - reference_age` (not a hardcoded constant). Same parameter appears in `leisure()`, `tied()`, `with_hours()`, and `utility_scale_factor()`. @@ -116,36 +137,57 @@ ACA variants don't create new regimes — they swap functions on baseline regime `premium_default = max(0, hic_premium − max(0, resources − consumption_dollars_floor))` is a tracked DAG node; `cash_on_hand` subtracts only the affordable part. The floor is therefore on pre-premium resources; OOP stays the post-decision shock in - `next_assets`. Coverage is unchanged on default (the defaulted premium is - uncompensated care). + `next_assets`. Employer and Medicare coverage is unchanged on default (the defaulted + premium is uncompensated care). Private non-group cover is not: under the brute-force + solver the `private_cover_paid_in_full` constraint makes `buy_private = yes` + infeasible whenever `premium_default > 0`, so no household holds private cover + financed by the floor. +- **Private non-group premium**: + `max(premium_minimum[is_married], p0 + b (1 + markup) E[insurer cost])`, with `b` + estimated. The intercept `p0` is the DAG node `private_premium_intercept`, solved by + bisection from the insurer's zero-profit condition on the data sample of private + buyers under 65 (`premium_sample_*`), so it moves with `b`. +- **Couples' cost sharing**: `deductible` and `oop_max` are the single household's; + `insured_oop` doubles both for married households. - **Two-track Medicaid eligibility**: `is_medicaid_eligible` is the union of a *categorical* track — `(crossed_oamc_threshold OR is_disabled)` AND the SSI asset and income tests (on SSI countable income) — and, under the ACA Medicaid-expansion variant, an *income-only* track — `aca_magi < 138% FPL`, scoped to the under-65 - non-disabled population. The expansion uses MAGI (full income via `aca_magi`), - distinct from the half-counted SSI countable income of the categorical track. + population whatever their health (the disabled under 65 included, without the asset + test). The expansion uses MAGI (full income via `aca_magi`), distinct from the + half-counted SSI countable income of the categorical track. Regime (HIS) transitions + and the pension `target_his` read the categorical track only (`is_ssi_eligible`): the + ACA leaves employer coverage untouched, so a retiree/tied household eligible only + through the expansion keeps its employer coverage, with Medicaid paying on top of its + OOP. - **`crossed_oamc_threshold`**: per-regime constant fixed param (`= spec["mc"] == "oamc"`, i.e. age ≥ 65), the *aged* indicator in eligibility. It replaced the `gets_medicare` gate there; `is_disabled` (= `health == disabled`, a DAG - function in `nomc`/`dimc` regimes) supplies the disabled arm. The Medicare - *transition* still uses its own build-time `gets_medicare` constant (`mc != nomc`) — - distinct from this. + function in `nomc`/`dimc` regimes) supplies the disabled arm. +- **Disability Medicare before 65**: a household is on Medicare (`dimc`) next period + exactly when it does not work this period and is disabled next period. The canwork + regime transitions split survivors between the `dimc` and `nomc` targets by + `P(disabled next | health)`, and the per-target health laws condition on the outcome + (`next_health_into_dimc`, `next_health_into_nomc`). - **SS claim-age adjustment baked into AIME**: at the voluntary claim the - early-retirement reduction / delayed-retirement credit is applied to PIA and converted - back to AIME via `find_aime` (the exact inverse of `pia`), so the permanent adjustment - rides in the carried `aime` with no extra state and persists into the forced regimes. + early-retirement reduction / delayed-retirement credit is applied to PIA + (`pia_adjusted_next_period`) and converted back to AIME via the exact inverse of + `pia`, so the permanent adjustment rides in the carried `aime` with no extra state. The `pia_aime_grid` / `pia_table` carry a fifth bend point above the taxable max (`max_delayed_factor * max_pia`, round-tripped to AIME), so a top earner who delays keeps the delayed-retirement credit in the carried AIME instead of clamping at the - taxable max — matching the reference (struct-ret's - `convert_pia_to_aime(..., impose_upper_bound_on_aime=False)`). Post-claim labor - accrual runs on this adjusted AIME; `_accrue_aime` still caps the labor-earnings base - at the taxable max, and the solve-phase `pension_assets_adjustment` corrects the - pension imputation gap, with pension imputation reading an *unadjusted* PIA - (`pia_unadjusted_next_period`) so the claim-age credit never feeds the pension node — - faithful to the reference. The DI path (`ssdi_pia`) reads the un-baked AIME. - `inelig`/`forced` regimes use a plain `next_aime` with no claim inputs - (`_select_aime_law` routes on `spec["ss"]`). + taxable max. Which PIA the next AIME encodes is the phased `carried_pia` node: + - solve: `pensions.total_to_pia`, the PIA `PIA*` with + `PIA* + (1 − τ) pbmax_{t+1}(PIA*) = PIA_adj + (1 − τ) pbmax_{t+1}(PIA_unadj)`, + because solve re-imputes pension wealth from next period's PIA (French & Jones 2011, + app. D; struct-ret `convert_total_ben_to_pia`); + - simulate: the adjusted PIA itself, since simulate carries true pension wealth. + Pension imputation reads the unadjusted PIA (`pia_unadjusted_next_period`), and the + solve-phase `pension_assets_adjustment` corrects the accrual-evolved pension against + it. The earnings-test credit covers only benefits the earnings test withheld. In the + `forced` regimes (70+) AIME is frozen (`next_aime_forced`); `inelig` regimes use + plain accrual with no claim inputs (`_select_aime_law` routes on `spec["ss"]`). The + DI path reads the un-baked AIME after indexing and labor accrual. - **ACA subsidies/mandate respect Medicaid**: `premium_subsidy`, `cost_sharing`, and `mandate_penalty` take `is_medicaid_eligible` and return the neutral value when the household is Medicaid-eligible (Medicaid is minimum-essential coverage). diff --git a/README.md b/README.md index 5f52c1a..f22ff8f 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,51 @@ # aca-model Core lifecycle model for the ACA structural retirement project. + +The benchmark folds the IID transitory medical-cost shock and declares its asset +transition as resources plus the negated medical bill. Production grid defaults remain +unfolded. To measure the unfolded benchmark control, pass +`grid_config=replace(BENCHMARK_GRID_CONFIG, fold_hcc_transitory=False, subtract_hcc_bill=False)` +to `create_benchmark_model`. Pass the same grid configuration to +`get_benchmark_initial_conditions`; both routes retain identical seeded draws for their +shared initial states. + +Model factories accept `execution_config=ExecutionConfig(...)`. An omitted policy uses +the smallest JAX allocator `bytes_limit` among the selected accelerators as an explicit +per-device budget. This enables budget-aware width selection; CPU construction remains +unbudgeted. Missing accelerator limits require an explicit policy instead of silently +falling back to unbudgeted execution. + +To combine a measured budget with device selection and sharding: + +```python +from dataclasses import replace + +from aca_model.execution import execution_config_for_devices + +execution_config = replace( + execution_config_for_devices(devices=(0, 1, 2)), + sharded_states=("pref_type",), +) +``` + +Pass this policy to the baseline, ACA or benchmark factory. Record its exact +`device_memory_bytes`, the selected devices' allocator observations, executed precision +and selected program widths with performance evidence. The budget is an allocator +ceiling; it does not establish that all workload phases fit. + +Explicit policies pass through unchanged. In particular, `ExecutionConfig()` is an +unbudgeted control and retains pylcm's bootstrap width selection. Economic grids, +numerical solver settings and result retention are separate choices. + +`GridConfig` accepts economic grid sizes and numerical solver choices only. Execution +controls belong to `ExecutionConfig`: sharded state names go in `sharded_states`, +declared program widths in `axis_widths`, and the per-device byte ceiling in +`device_memory_bytes`. Obsolete grid execution keywords are rejected. A state-specific +batch size has no automatic conversion to a flattened program width; callers must choose +and record an explicit execution policy. + +For comparisons with an older pylcm version, retain its model/configuration pin and use +a separately reviewed API adapter. Preserve grids, transition declarations, numerical +solver choices, retention and seeds across the comparison. An unspecified width in the +new policy means planner selection, not a translation of an old zero-valued batch size. diff --git a/docs/dag-migration.md b/docs/dag-migration.md new file mode 100644 index 0000000..5e4650f --- /dev/null +++ b/docs/dag-migration.md @@ -0,0 +1,31 @@ +# ACA: one structural stage table, existing numerical programs + +Both factories declare the same admissible starts, `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. +`(51, "dead")` is not solved. Starts at age ≥ 61 or in `dead` are rejected at +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 +`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 +paths remain; different target schemas can still need different continuation +programs, so no backend compile-time improvement is claimed. + +The dead regime has no outgoing edges in `Model.edges` and is terminal at every age. Age 95 retains its +living consumption/saving problem and handoff into the separate age-96 bequest. The +three-to-two health grid, claim/lagged-work entry and exit, carried pension state +and phase semantics are not simplified away. + +Use `model.graph.nodes`, `model.graph.initial_nodes`, and `model.graph.edges`; +`model.graph.solution` and `model.graph.simulation` expose the effective phase +graphs. The declared edge selectors are identical across policy variants and +phases; carried pension wealth retains its distinct solve/simulate laws. Keep the phase +target queries and period-keyed solutions. GridSearch is the first numerical +acceptance target. Annual ACA calibration is not automatically valid on a +differently spaced clock. diff --git a/src/aca_model/__init__.py b/src/aca_model/__init__.py index 1fc43f7..d26e6e4 100644 --- a/src/aca_model/__init__.py +++ b/src/aca_model/__init__.py @@ -1,6 +1,13 @@ +import os + import jax -jax.config.update("jax_enable_x64", True) +_x64_requested = os.environ.get("ACA_JAX_ENABLE_X64", "1").lower() not in { + "0", + "false", + "no", +} +jax.config.update("jax_enable_x64", _x64_requested) # Import lcm before installing the claw so its `_jaxtyping_patch` (picklable # jaxtyping sentinel) and `MappingProxyType` pytree registration are in place. diff --git a/src/aca_model/_benchmark_data/benchmark_params.pkl b/src/aca_model/_benchmark_data/benchmark_params.pkl index 0447289..c5844b5 100644 Binary files a/src/aca_model/_benchmark_data/benchmark_params.pkl and b/src/aca_model/_benchmark_data/benchmark_params.pkl differ diff --git a/src/aca_model/aca/health_insurance.py b/src/aca_model/aca/health_insurance.py index 9e23c71..7e0b5bb 100644 --- a/src/aca_model/aca/health_insurance.py +++ b/src/aca_model/aca/health_insurance.py @@ -18,6 +18,7 @@ DiscreteAction, DiscreteState, FloatND, + IntND, ScalarBool, ScalarFloat, ) @@ -35,6 +36,23 @@ class PolicyVariant(Enum): ACA_ONLY_MEDICAID_EXPANSION = auto() +def community_rated_premium( + is_married: IntND, + buy_private: DiscreteAction, + aca_premium_single: ScalarFloat, + aca_premium_married_extra: ScalarFloat, +) -> FloatND: + """Compute the ACA private non-group premium before premium credits. + + With community rating the premium varies only by household size, not + with expected medical costs: `aca_premium_single` for a single + household plus `aca_premium_married_extra` for a married one. Return 0 + when not buying (`buy_private=no`). + """ + premium = aca_premium_single + aca_premium_married_extra * is_married + return jnp.where(buy_private == BuyPrivate.yes, premium, 0.0) + + def mandate_penalty( gross_income: FloatND, spousal_income: DiscreteState, @@ -111,12 +129,44 @@ def cost_sharing( return jnp.where(is_insured & ~is_medicaid_eligible, scale, 1.0) +def cost_sharing_bracket( + gross_income: FloatND, + spousal_income: DiscreteState, + buy_private: DiscreteAction, + is_medicaid_eligible: BoolND, + cost_sharing_schedule: MappingLeaf, +) -> IntND: + """Return the household's cost-sharing bracket, `n_kinks + 1` when neutral. + + The bracket is the FPL bracket of `cost_sharing`; a household without + private cover, or Medicaid-eligible, falls in the neutral bracket one past + the last, whose scale is 1.0. The bracket takes finitely many values, so + the out-of-pocket bill can be averaged per bracket. + """ + sched = cast("Mapping[str, Any]", cost_sharing_schedule.data) + kinks = sched["kinks"] # [n_kinks, 3] + bracket = jnp.searchsorted(kinks[:, spousal_income], gross_income, side="right") + is_insured = buy_private == BuyPrivate.yes + return jnp.where( + is_insured & ~is_medicaid_eligible, bracket, kinks.shape[0] + 1 + ).astype(jnp.int32) + + +def cost_sharing_scale_of_bracket( + cost_sharing_bracket: IntND, + cost_sharing_schedule: MappingLeaf, +) -> FloatND: + """Return the cost-sharing scale of a bracket from `cost_sharing_bracket`.""" + sched = cast("Mapping[str, Any]", cost_sharing_schedule.data) + factors = sched["factors"] # [n_kinks + 1] + return jnp.append(factors, jnp.ones((), dtype=factors.dtype))[cost_sharing_bracket] + + def is_medicaid_eligible( is_ssi_eligible: BoolND, aca_magi: FloatND, spousal_income: DiscreteState, crossed_oamc_threshold: ScalarBool, - is_disabled: BoolND, medicaid_schedule: MappingLeaf, ) -> BoolND: """Determine Medicaid eligibility on two tracks under ACA expansion. @@ -125,22 +175,18 @@ def is_medicaid_eligible( - **Categorical** (SSI-linked): `is_ssi_eligible`, which applies the aged-or-disabled gate plus the SSI asset and income tests. - - **Expansion** (ACA): the under-65, non-disabled population with MAGI - below the expansion threshold (138% FPL encoded in - `medicaid_schedule["income_threshold"]`). Expansion never reaches the - aged or the disabled; they stay on the categorical track with its - asset test. + - **Expansion** (ACA): every under-65 household with MAGI below the + expansion threshold (138% FPL encoded in + `medicaid_schedule["income_threshold"]`), regardless of assets or + health, so the disabled under 65 qualify without the asset test. The + aged stay on the categorical track with its asset test. Expansion uses full-count MAGI, not the half-counted SSI `countable_income`. """ sched = cast("Mapping[str, Any]", medicaid_schedule.data) threshold = sched["income_threshold"] - expansion = ( - (~crossed_oamc_threshold) - & (~is_disabled) - & (aca_magi < threshold[spousal_income]) - ) + expansion = (~crossed_oamc_threshold) & (aca_magi < threshold[spousal_income]) return is_ssi_eligible | expansion @@ -203,6 +249,7 @@ def primary_oop( total_health_costs: FloatND, cost_sharing_scale: FloatND, buy_private: DiscreteAction, + is_married: IntND, deductible: ScalarFloat, coinsurance_rate: ScalarFloat, oop_max: ScalarFloat, @@ -210,13 +257,15 @@ def primary_oop( """Compute primary OOP costs with ACA cost-sharing reductions. Scale deductible, coinsurance rate, and OOP max by the cost-sharing - factor before applying the standard OOP calculation. When uninsured + factor before applying the standard OOP calculation; couples' + deductible and OOP max are twice the single ones. When uninsured (`buy_private=no`), OOP equals total health costs (no coverage). """ + couple_scale = 1.0 + is_married insured_oop = oop_costs( total_health_costs, - deductible * cost_sharing_scale, + deductible * cost_sharing_scale * couple_scale, coinsurance_rate * cost_sharing_scale, - oop_max * cost_sharing_scale, + oop_max * cost_sharing_scale * couple_scale, ) return jnp.where(buy_private == BuyPrivate.yes, insured_oop, total_health_costs) diff --git a/src/aca_model/aca/model.py b/src/aca_model/aca/model.py index 0db5ff6..5b17c97 100644 --- a/src/aca_model/aca/model.py +++ b/src/aca_model/aca/model.py @@ -7,19 +7,24 @@ from collections.abc import Mapping from typing import Any -from lcm import AgeGrid, DiscreteGrid, Model +from lcm import AgeGrid, DiscreteGrid, ExecutionConfig, Model from lcm.typing import UserParams 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.regimes import RegimeId, SolverName, build_model_slots -from aca_model.config import MODEL_CONFIG, GridConfig +from aca_model.baseline.regimes import ( + INITIAL_NODES, + RegimeId, + SolverName, + build_model_slots, +) +from aca_model.config import MODEL_AGES, GridConfig +from aca_model.execution import execution_config_for_devices def create_model( *, - n_subjects: int, policy: PolicyVariant, fixed_params: UserParams, wage_params: Mapping[str, Any], @@ -28,11 +33,11 @@ def create_model( pref_type_grid: DiscreteGrid, solver: SolverName = "brute_force", consumption_dollars_points: tuple[float, ...] | None = None, + execution_config: ExecutionConfig | None = None, ) -> Model: """Create an ACA policy variant model. Args: - n_subjects: Forwarded to `lcm.Model(n_subjects=...)`. policy: Which ACA policy combination to apply (e.g. `PolicyVariant.ACA`). fixed_params: Parameters to fix at model creation time. Pass @@ -53,16 +58,15 @@ def create_model( consumption_dollars_points: Construction-time consumption action gridpoints; required under DC-EGM. See `aca_model.baseline.model.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. Returns: pylcm Model. """ - ages = AgeGrid( - start=MODEL_CONFIG.start_age, - stop=MODEL_CONFIG.end_age - 1, - step="Y", - ) + ages = AgeGrid(exact_values=MODEL_AGES) _fail_if_dcegm_without_consumption_points( solver=solver, consumption_dollars_points=consumption_dollars_points ) @@ -87,11 +91,16 @@ def create_model( return Model( regimes=regimes, + initial_nodes=INITIAL_NODES, ages=ages, regime_id_class=RegimeId, description=f"Structural retirement model ({policy.name})", fixed_params=fixed_params, derived_categoricals=derived_categoricals, - n_subjects=n_subjects, + execution_config=( + execution_config_for_devices() + if execution_config is None + else execution_config + ), **model_slots, ) diff --git a/src/aca_model/aca/regimes/_overrides.py b/src/aca_model/aca/regimes/_overrides.py index ff78245..eb92ffe 100644 --- a/src/aca_model/aca/regimes/_overrides.py +++ b/src/aca_model/aca/regimes/_overrides.py @@ -9,6 +9,7 @@ from aca_model.aca import health_insurance as aca_hi from aca_model.aca.health_insurance import PolicyVariant from aca_model.baseline.regimes._common import RegimeSpec +from aca_model.environment import taxes def apply_aca_overrides( @@ -18,16 +19,21 @@ def apply_aca_overrides( ) -> None: """Override baseline functions with ACA versions in-place. + Every variant charges the ACA surtaxes in every regime: the 3.8% surtax + on unearned income above 200,000, plus, under the full ACA only, the + 0.9% surtax on earnings above 200,000 (struct-ret's composition). + Three orthogonal feature flags derived from the policy variant: - **Medicaid expansion**: two-track eligibility (categorical SSI plus the - under-65 non-disabled MAGI expansion) installed on all regimes. The - expansion arm is internally scoped to the under-65 non-disabled - population, so post-65 and disabled households keep the categorical - track with its asset test. - - **Subsidies**: premium credits, cost-sharing reductions, and their - consuming functions (nongroup+nomc only). All mask to their neutral - value when Medicaid-eligible (minimum-essential coverage). + under-65 MAGI expansion) installed on all regimes. The expansion arm is + internally scoped to the under-65 population, disabled or not, so only + post-65 households keep the categorical track with its asset test. + - **Subsidies**: the reformed non-group market (nongroup+nomc only): the + community-rated plan premium, premium credits, cost-sharing + reductions, and their consuming functions. Credits and cost-sharing + mask to their neutral value when Medicaid-eligible (minimum-essential + coverage). - **Mandate**: individual mandate penalty (nongroup+nomc only, requires subsidies), waived when Medicaid-eligible. """ @@ -41,6 +47,13 @@ def apply_aca_overrides( PolicyVariant.ACA_NO_MEDICAID_EXPANSION, ) + functions["after_tax_income_before_aca_surtaxes"] = taxes.after_tax_income + functions["after_tax_income"] = ( + taxes.after_tax_income_with_aca_surtaxes + if policy == PolicyVariant.ACA + else taxes.after_tax_income_with_aca_investment_surtax + ) + if has_medicaid_expansion: functions["is_medicaid_eligible"] = aca_hi.is_medicaid_eligible @@ -49,8 +62,15 @@ def apply_aca_overrides( functions["mandate_penalty"] = aca_hi.mandate_penalty # No mandate: mandate_penalty is a fixed param (0.0) in the params # 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. + functions.pop("private_premium_intercept", None) functions["hic_premium_subsidy"] = aca_hi.premium_subsidy - functions["cost_sharing_scale"] = aca_hi.cost_sharing + # Split at the bracket, which takes finitely many values, so the + # out-of-pocket bill can be averaged per bracket. + functions["cost_sharing_bracket"] = aca_hi.cost_sharing_bracket + functions["cost_sharing_scale"] = aca_hi.cost_sharing_scale_of_bracket functions["premium_default"] = aca_hi.premium_default functions["cash_on_hand"] = aca_hi.cash_on_hand functions["primary_oop"] = aca_hi.primary_oop diff --git a/src/aca_model/agent/assets_and_income.py b/src/aca_model/agent/assets_and_income.py index 147adf5..251c7f8 100644 --- a/src/aca_model/agent/assets_and_income.py +++ b/src/aca_model/agent/assets_and_income.py @@ -63,9 +63,8 @@ def cash_on_hand( part (`premium_default`) is never paid, so the effective premium is `hic_premium - premium_default`. - OOP health costs are NOT deducted here — they are deducted from - next-period assets instead, matching the timing where HCC shocks are - integrated over (agent does not condition consumption on OOP). + Out-of-pocket health costs are not deducted here: they are realised after + the period's choices and come off end-of-period assets (`next_assets`). """ effective_premium = hic_premium - premium_default return assets + after_tax_income + ssi_benefit - effective_premium @@ -96,8 +95,10 @@ def next_assets( ) -> ContinuousState: """Compute beginning-of-next-period assets for non-terminal targets. - OOP health costs are deducted here (not from cash_on_hand) so that the - consumption choice does not condition on the HCC shock realization. + `oop_costs` is the bill for this period's medical-cost shocks, which are + drawn after the period's choices; the continuation integrates over them. + End-of-period assets can therefore fall below the borrowing limit that + binds consumption (medical debt). """ return ( cash_on_hand @@ -108,6 +109,31 @@ def next_assets( ) +def assets_before_bill( + cash_on_hand: FloatND, + transfers: FloatND, + pension_assets_adjustment: FloatND, + consumption_dollars: ContinuousAction, +) -> FloatND: + """Compute next-period assets for non-terminal targets before the medical bill. + + `next_assets` is this minus `oop_costs`. + """ + return cash_on_hand + transfers + pension_assets_adjustment - consumption_dollars + + +def assets_before_bill_when_dead( + cash_on_hand: FloatND, + transfers: FloatND, + consumption_dollars: ContinuousAction, +) -> FloatND: + """Compute assets for the dead/terminal target before the medical bill. + + `next_assets_when_dead` is this minus `oop_costs`. + """ + return cash_on_hand + transfers - consumption_dollars + + def next_assets_when_dead( cash_on_hand: FloatND, transfers: FloatND, @@ -211,3 +237,17 @@ def borrowing_constraint( The `max` form returns `floor` exactly. """ return consumption_dollars <= jnp.maximum(cash_on_hand, consumption_dollars_floor) + + +def consumption_above_floor( + consumption_dollars: ContinuousAction, + consumption_dollars_floor: FloatND, +) -> BoolND: + """Consumption cannot fall below the household's own $-floor. + + The consumption grid starts at the single floor; a married household's + floor is higher, and consuming below it would let the household save part + of a floor-topping transfer. With cash on hand below the floor, this and + the borrowing constraint together pin consumption at the floor. + """ + return consumption_dollars >= consumption_dollars_floor diff --git a/src/aca_model/agent/health.py b/src/aca_model/agent/health.py index 66f04fe..4246bcb 100644 --- a/src/aca_model/agent/health.py +++ b/src/aca_model/agent/health.py @@ -6,7 +6,16 @@ import jax.numpy as jnp from lcm import categorical -from lcm.typing import DiscreteState, FloatND, IntND, Period, ScalarInt +from lcm.typing import ( + DiscreteAction, + DiscreteState, + FloatND, + IntND, + Period, + ScalarInt, +) + +from aca_model.agent.labor_market import LaborSupply @categorical(ordered=True) @@ -60,3 +69,37 @@ def next_health_cross( regimes (Health) at the age-65 boundary. """ return health_trans_probs_cross[period, health] + + +def next_health_into_dimc( + health: DiscreteState, + period: Period, + health_trans_probs: FloatND, +) -> FloatND: + """Next health of a household that lands on disability Medicare: disabled. + + Pre-65 Medicare goes to households disabled next period, so landing in a + `dimc` regime reveals the next health state. + """ + probs = health_trans_probs[period, health] + return jnp.zeros_like(probs).at[HealthWithDisability.disabled].set(1.0) + + +def next_health_into_nomc( + health: DiscreteState, + period: Period, + labor_supply: DiscreteAction, + health_trans_probs: FloatND, +) -> FloatND: + """Next health of a pre-65 household that lands without Medicare. + + A household that does not work this period gets disability Medicare + whenever it is disabled next period, so landing without Medicare means it + is not disabled: its next health is the transition row with the disabled + state removed and renormalised. A worker keeps the unconditional row. + """ + probs = health_trans_probs[period, health] + off_disability = probs.at[HealthWithDisability.disabled].set(0.0) + mass = off_disability.sum() + conditional = off_disability / jnp.where(mass > 0.0, mass, 1.0) + return jnp.where(labor_supply == LaborSupply.do_not_work, conditional, probs) diff --git a/src/aca_model/agent/labor_market.py b/src/aca_model/agent/labor_market.py index 5665d92..24a012a 100644 --- a/src/aca_model/agent/labor_market.py +++ b/src/aca_model/agent/labor_market.py @@ -121,3 +121,18 @@ def next_spousal_income( ) -> FloatND: """Stochastic spousal income transition.""" return spousal_income_trans_probs[period, spousal_income] + + +def spousal_income_amount( + spousal_income: DiscreteState, + good_health: IntND, + period: Period, + spousal_income_amounts: FloatND, +) -> FloatND: + """Return the household's spousal income this period. + + `spousal_income_amounts` is indexed by period, the respondent's own health + (good or not) and the `SpousalIncome` state; only a working spouse + (`married_has_inc`) brings income. + """ + return spousal_income_amounts[period, good_health, spousal_income] diff --git a/src/aca_model/agent/preferences.py b/src/aca_model/agent/preferences.py index 438b2c4..9abdefe 100644 --- a/src/aca_model/agent/preferences.py +++ b/src/aca_model/agent/preferences.py @@ -23,6 +23,13 @@ # endowment; it only bends the map near and beyond the endowment. _LEISURE_SMOOTHING_FRACTION = 0.01 +# Positive floor that leisure approaches as work costs exceed the endowment, as a +# fraction of the time endowment. It keeps felicity and its consumption derivative +# inside float32's range when leisure_available is far below zero, down to the +# NB-EGM inverse bracket's lowest consumption; away from the floor its effect decays +# exponentially. +_LEISURE_FLOOR_FRACTION = 1e-5 + @categorical(ordered=False) class PrefType: @@ -71,17 +78,25 @@ def fixed_cost_of_work( def _smooth_leisure_floor( leisure_available: FloatND, time_endowment: ScalarFloat ) -> FloatND: - """Bend leisure to a strictly positive floor as work costs approach the endowment. - - `softplus(x) = log(1 + e^x)` via `jnp.logaddexp(0, x)`, scaled by a small fraction - of the endowment. Where `leisure_available` is large relative to the smoothing width - the map reduces to `leisure_available` (bulk unchanged); as it falls to zero leisure - bends to `0⁺` — never negative, never a kinked clamp — so the CRRA aggregator never - receives a non-positive base. The smoothing width scales with the endowment, so the - map is scale-invariant. + """Bend leisure smoothly onto a positive floor as work costs approach the endowment. + + Leisure is a smooth maximum of `leisure_available` and the floor + `F = _LEISURE_FLOOR_FRACTION * time_endowment`: + `s * logaddexp(F / s, leisure_available / s)`, with smoothing width + `s = _LEISURE_SMOOTHING_FRACTION * time_endowment`. + + - `leisure_available` large relative to `s`: leisure equals `leisure_available` + up to a term that decays like `e^(-leisure_available / s)`. + - `leisure_available` at or below zero: leisure bends to `F⁺` — never below the + floor, never a kinked clamp — so every work choice stays feasible and the CRRA + aggregator's base stays large enough for float32 felicity and marginal + felicity. + + Both widths scale with the endowment, so the map is scale-invariant. """ smoothing = _LEISURE_SMOOTHING_FRACTION * time_endowment - return smoothing * jnp.logaddexp(0.0, leisure_available / smoothing) + floor = _LEISURE_FLOOR_FRACTION * time_endowment + return smoothing * jnp.logaddexp(floor / smoothing, leisure_available / smoothing) def leisure_canwork_retiree_or_nongroup( @@ -269,16 +284,20 @@ def utility_scale_factor( def scaled_bequest_weight( bequest_weight: ScalarFloat, - consumption_weight: ScalarFloat, - coefficient_rra: ScalarFloat, + consumption_weight: FloatND, + coefficient_rra: FloatND, time_endowment: ScalarFloat, - time_discount_factor: ScalarFloat, + discount_factor: FloatND, rate_of_return: ScalarFloat, ) -> FloatND: - """Transform raw bequest weight into the form used in the bequest function. + """Per-type bequest weight `θ_B` from the estimated bequest MPC. result = T^ξ * (bw / (1+r-bw))^ξ₂ / β where ξ = (1-α)(1-γ) and ξ₂ = α(1-γ) - 1. + + `bw` is the marginal propensity to consume out of final-period wealth, one + scalar for all types; α, γ and β are the type's consumption weight, CRRA + coefficient and discount factor, so `θ_B` differs by preference type. """ xi = (1.0 - consumption_weight) * (1.0 - coefficient_rra) xi2 = consumption_weight * (1.0 - coefficient_rra) - 1.0 @@ -287,7 +306,7 @@ def scaled_bequest_weight( bequest_weight > 0.0, time_endowment**xi * (safe_bw / (1.0 + rate_of_return - safe_bw)) ** xi2 - / time_discount_factor, + / discount_factor, 0.0, ) @@ -295,7 +314,7 @@ def scaled_bequest_weight( def bequest( assets: ContinuousState, bequest_shifter: ScalarFloat, - scaled_bequest_weight: ScalarFloat, + scaled_bequest_weight: FloatND, consumption_weight: FloatND, coefficient_rra: FloatND, utility_scale_factor: FloatND, diff --git a/src/aca_model/baseline/health_insurance.py b/src/aca_model/baseline/health_insurance.py index 91f2503..d073287 100644 --- a/src/aca_model/baseline/health_insurance.py +++ b/src/aca_model/baseline/health_insurance.py @@ -12,6 +12,7 @@ - SSI benefit computation """ +import jax import jax.numpy as jnp import lcm from lcm import categorical @@ -49,8 +50,7 @@ class HealthInsuranceState: def countable_income( labor_income: FloatND, capital_income: FloatND, - spousal_income: DiscreteState, - spousal_income_amounts: FloatND, + spousal_income_amount: FloatND, ss_benefit: FloatND, pension_benefit: FloatND, ssi_ignored_overall: ScalarFloat, @@ -64,7 +64,7 @@ def countable_income( return ( earned + capital_income - + spousal_income_amounts[spousal_income] + + spousal_income_amount + ss_benefit + pension_benefit - ssi_ignored_overall @@ -131,8 +131,7 @@ def is_disabled_never() -> BoolND: def aca_magi( labor_income: FloatND, capital_income: FloatND, - spousal_income: DiscreteState, - spousal_income_amounts: FloatND, + spousal_income_amount: FloatND, ss_benefit: FloatND, pension_benefit: FloatND, ) -> FloatND: @@ -145,7 +144,7 @@ def aca_magi( return ( labor_income + capital_income - + spousal_income_amounts[spousal_income] + + spousal_income_amount + ss_benefit + pension_benefit ) @@ -183,47 +182,96 @@ def ssi_benefit( def premium( - age: Age, - good_health: IntND, - is_married: IntND, - labor_supply: DiscreteAction, buy_private: DiscreteAction, - premium_intercept: ScalarFloat, - premium_age: ScalarFloat, - premium_age_sq: ScalarFloat, - premium_age_cub: ScalarFloat, - premium_predicted_hcc: ScalarFloat, - premium_good_health: ScalarFloat, - premium_married: ScalarFloat, - premium_works: ScalarFloat, - premium_married_works: ScalarFloat, - premium_minimum: ScalarFloat, + is_married: IntND, predicted_hcc_insurer: FloatND, + private_premium_intercept: FloatND, + premium_predicted_hcc: ScalarFloat, + premium_markup: ScalarFloat, + premium_minimum: FloatND, ) -> FloatND: - """Compute health insurance premium for canwork regimes. + """Compute the private non-group premium. - Premium coefficients are regime-specific (different for each HIC type). - Return 0 when uninsured (`buy_private=no`). + `max(minimum[married], p0 + b (1 + markup) E[insurer cost])` when buying + private cover, 0 when uninsured (`buy_private=no`). The insurer cannot + price health or marital status beyond the predicted cost and the + married minimum. """ - works = labor_supply != LaborSupply.do_not_work raw = ( - premium_intercept - + premium_age * age - + premium_age_sq * age**2 - + premium_age_cub * age**3 - + premium_predicted_hcc * predicted_hcc_insurer - + premium_good_health * good_health - + premium_married * is_married - + premium_works * works - + premium_married_works * is_married * works + private_premium_intercept + + premium_predicted_hcc * (1.0 + premium_markup) * predicted_hcc_insurer ) return jnp.where( buy_private == BuyPrivate.yes, - jnp.maximum(premium_minimum, raw), + jnp.maximum(premium_minimum[is_married], raw), 0.0, ) +# Bracket and step count of the intercept bisection. The bracket matches the +# reference implementation; 100 halvings shrink it below float64 resolution. +_INTERCEPT_BRACKET = (-50_000.0, 10_000.0) +_N_INTERCEPT_BISECTIONS = 100 + + +def private_premium_intercept( + premium_predicted_hcc: ScalarFloat, + premium_markup: ScalarFloat, + premium_minimum: FloatND, + premium_sample_insurer_cost: FloatND, + premium_sample_is_married: IntND, +) -> FloatND: + """Solve the private premium intercept from the insurer's zero-profit condition. + + On the sample of private buyers under 65 in the data, the intercept `p0` + makes premium revenue equal expected payouts plus the markup: + + ``` + mean_i max(minimum[married_i], p0 + b (1 + markup) cost_i) + = (1 + markup) mean_i cost_i + ``` + + The left side is continuous and non-decreasing in `p0`; bisection on + `[-50000, 10000]` finds the root. Without a sign change on the bracket + the result is NaN, so an unpriceable slope `b` fails loudly downstream. + """ + slope = premium_predicted_hcc * (1.0 + premium_markup) + minimum = premium_minimum[premium_sample_is_married] + target = (1.0 + premium_markup) * jnp.mean(premium_sample_insurer_cost) + + def excess_revenue(intercept: FloatND) -> FloatND: + premiums = jnp.maximum(minimum, intercept + slope * premium_sample_insurer_cost) + return jnp.mean(premiums) - target + + def halve( + _: jax.Array, bracket: tuple[FloatND, FloatND] + ) -> tuple[FloatND, FloatND]: + low, high = bracket + mid = 0.5 * (low + high) + too_high = excess_revenue(mid) > 0.0 + return jnp.where(too_high, low, mid), jnp.where(too_high, mid, high) + + low0 = jnp.asarray(_INTERCEPT_BRACKET[0]) + high0 = jnp.asarray(_INTERCEPT_BRACKET[1]) + low, high = jax.lax.fori_loop(0, _N_INTERCEPT_BISECTIONS, halve, (low0, high0)) + bracketed = (excess_revenue(low0) <= 0.0) & (excess_revenue(high0) >= 0.0) + return jnp.where(bracketed, 0.5 * (low + high), jnp.nan) + + +def private_cover_paid_in_full( + buy_private: DiscreteAction, + premium_default: FloatND, +) -> BoolND: + """Private cover is a feasible choice only when its premium is paid in full. + + A household whose resources above the consumption floor fall short of the + premium cannot buy private cover: the government does not finance private + insurance through floor transfers, and an unpaid premium buys no private + cover. Not buying is always feasible. + """ + return (buy_private == BuyPrivate.no) | (premium_default <= 0.0) + + def premium_insured( age: Age, good_health: IntND, @@ -309,20 +357,49 @@ def oop_costs( return jnp.minimum(oop, oop_max) +def negative_oop_costs(oop_costs: FloatND) -> FloatND: + """Return the out-of-pocket bill as the shock added to assets.""" + return -oop_costs + + +def insured_oop( + total_health_costs: FloatND, + is_married: IntND, + deductible: ScalarFloat, + coinsurance_rate: ScalarFloat, + oop_max: ScalarFloat, +) -> FloatND: + """Compute OOP costs under the regime's plan. + + `deductible` and `oop_max` are the single household's; couples' + deductible and OOP maximum are twice as large. + """ + couple_scale = 1.0 + is_married + return oop_costs( + total_health_costs, + deductible * couple_scale, + coinsurance_rate, + oop_max * couple_scale, + ) + + def primary_oop( total_health_costs: FloatND, buy_private: DiscreteAction, + is_married: IntND, deductible: ScalarFloat, coinsurance_rate: ScalarFloat, oop_max: ScalarFloat, ) -> FloatND: - """Compute primary OOP costs. + """Compute primary OOP costs for regimes with the private-cover choice. - When uninsured (`buy_private=no`), OOP equals total health costs - (no coverage). + Private buyers pay `insured_oop`; when uninsured (`buy_private=no`), OOP + equals total health costs (no coverage). """ - insured_oop = oop_costs(total_health_costs, deductible, coinsurance_rate, oop_max) - return jnp.where(buy_private == BuyPrivate.yes, insured_oop, total_health_costs) + insured = insured_oop( + total_health_costs, is_married, deductible, coinsurance_rate, oop_max + ) + return jnp.where(buy_private == BuyPrivate.yes, insured, total_health_costs) def is_medicaid_eligible(is_ssi_eligible: BoolND) -> BoolND: @@ -333,13 +410,14 @@ def is_medicaid_eligible(is_ssi_eligible: BoolND) -> BoolND: def target_his( his: IntND, labor_supply: DiscreteAction, - is_medicaid_eligible: BoolND, + is_ssi_eligible: BoolND, ) -> IntND: """Return the HIS class of the surviving target regime. Mirrors the cross-HIS branches inside `_make_transition_canwork` (retiree, tied, nongroup): tied agents who stop working become nongroup, and - Medicaid-eligible agents are overridden to nongroup. Used by + categorically (SSI-) Medicaid-eligible agents are overridden to nongroup. + The ACA Medicaid expansion leaves employer coverage untouched. Used by `imputed_pension_wealth_next_period` to look up next-period imputation coefficients at the target's HIS. """ @@ -347,7 +425,7 @@ def target_his( labor_supply == LaborSupply.do_not_work ) return jnp.where( - tied_to_ng | is_medicaid_eligible, + tied_to_ng | is_ssi_eligible, HealthInsuranceState.nongroup, his, ).astype(jnp.int32) @@ -355,18 +433,18 @@ def target_his( def target_his_forcedout( his: IntND, - is_medicaid_eligible: BoolND, + is_ssi_eligible: BoolND, ) -> IntND: """Return the HIS class of the surviving target regime in forced-out regimes. Forced-out regimes have no labor-supply choice, and tied agents have already moved to nongroup before the forced-out age, so the only HIS - override is Medicaid eligibility → nongroup. Used by + override is categorical (SSI) Medicaid eligibility → nongroup. Used by `imputed_pension_wealth_next_period` to look up next-period imputation coefficients at the target's HIS. """ return jnp.where( - is_medicaid_eligible, + is_ssi_eligible, HealthInsuranceState.nongroup, his, ).astype(jnp.int32) @@ -393,18 +471,65 @@ def oop_with_medicaid( return jnp.where(is_medicaid_eligible, medicaid_oop, primary_oop) +def oop_with_medicaid_replacing_private( + total_health_costs: FloatND, + primary_oop: FloatND, + is_medicaid_eligible: BoolND, + deductible_medicaid: ScalarFloat, + coinsurance_rate_medicaid: ScalarFloat, + oop_max_medicaid: ScalarFloat, +) -> FloatND: + """Compute OOP costs in the non-group regimes before Medicare. + + A Medicaid-eligible household is covered by Medicaid alone: Medicaid is + the payer of last resort and does not cover private premia, so it can + neither hold private cover nor stay uninsured. Its OOP is Medicaid's + cost-sharing on total costs, whatever its `buy_private` choice. Other + households pay the primary (private or uninsured) OOP. + """ + medicaid_oop = oop_costs( + total_health_costs=total_health_costs, + deductible=deductible_medicaid, + coinsurance_rate=coinsurance_rate_medicaid, + oop_max=oop_max_medicaid, + ) + return jnp.where(is_medicaid_eligible, medicaid_oop, primary_oop) + + +def medicaid_adjusted_premium( + plan_premium: FloatND, + is_medicaid_eligible: BoolND, +) -> FloatND: + """Compute the premium a non-group household pays. + + A Medicaid-eligible household pays no premium: before Medicare its only + cover is Medicaid, whose premium is zero; with Medicare, Medicaid pays + the Medicare premium. + """ + return jnp.where(is_medicaid_eligible, 0.0, plan_premium) + + def hcc_insurer_predicted( + period: Period, + is_married: IntND, + good_health: IntND, hcc_persistent: ContinuousState, predicted_hcc_insurer_table: FloatND, hcc_persistent_grid: FloatND, ) -> FloatND: - """Interpolate pre-computed expected insurer cost for the current HCC state. + """Interpolate the private insurer's expected cost for the household. - The table contains E[total_costs - oop_costs | hcc_persistent] at each - persistent grid point. Linear interpolation handles off-grid values - during simulation (where draw_shock returns continuous AR1 values). + `hcc_persistent` is last period's persistent shock, so the table holds + this period's E[total_costs - oop_costs | age, married, good health, + lagged persistent node] on `[period, is_married, good_health, node]`. Linear + interpolation handles off-grid values during simulation (where + draw_shock returns continuous AR1 values). """ - return jnp.interp(hcc_persistent, hcc_persistent_grid, predicted_hcc_insurer_table) + return jnp.interp( + hcc_persistent, + hcc_persistent_grid, + predicted_hcc_insurer_table[period, is_married, good_health], + ) def compute_hcc_insurer_table( @@ -421,8 +546,9 @@ def compute_hcc_insurer_table( ) -> FloatND: """Compute predicted insurer costs table for all persistent grid points. - For each source persistent state i, integrate over (target persistent j, - transitory k) weighted by transition probs and quadrature weights. + For each lagged persistent state i, integrate over this period's + (persistent j, transitory k) weighted by transition probs and quadrature + weights. """ std_trans = jnp.sqrt(1.0 - std_xsect_persistent**2) # total_costs[j, k] for all (persistent, transitory) combinations @@ -453,11 +579,17 @@ def total_costs( good_health: IntND, log_mean: FloatND, log_std: FloatND, - hcc_persistent: ContinuousState, - hcc_transitory: ContinuousState, + next_hcc_persistent: ContinuousState, + next_hcc_transitory: ContinuousState, std_xsect_persistent: ScalarFloat, ) -> FloatND: - """Compute total health care costs from log-normal model. + """Compute this period's total health care costs from the log-normal model. + + Costs are realised after the period's choices. The persistent and + transitory shocks that drive them are the draws the transition makes, + `next_hcc_persistent` and `next_hcc_transitory`: the next period carries + them as its lagged shocks. Age, marital status and health are this + period's. ``log_mean`` and ``log_std`` are ``pd.Series`` with ``(age, is_married, good_health)`` MultiIndex, resolved by pylcm via ``derived_categoricals``. @@ -466,5 +598,5 @@ def total_costs( return jnp.exp( log_mean[period, is_married, good_health] + log_std[period, is_married, good_health] - * (hcc_persistent * std_xsect_persistent + hcc_transitory * std_trans) + * (next_hcc_persistent * std_xsect_persistent + next_hcc_transitory * std_trans) ) diff --git a/src/aca_model/baseline/model.py b/src/aca_model/baseline/model.py index 87acf7c..3522e30 100644 --- a/src/aca_model/baseline/model.py +++ b/src/aca_model/baseline/model.py @@ -5,7 +5,7 @@ Usage: from aca_model.baseline.model import create_model - model = create_model(n_subjects=..., fixed_params=..., wage_params=..., ...) + model = create_model(fixed_params=..., wage_params=..., ...) params = get_default_params() V = model.solve(params) """ @@ -13,21 +13,22 @@ from collections.abc import Mapping from typing import Any -from lcm import AgeGrid, DiscreteGrid, Model +from lcm import AgeGrid, DiscreteGrid, ExecutionConfig, Model from lcm.typing import UserParams from aca_model.baseline.regimes import ( + INITIAL_NODES, RegimeId, SolverName, build_all_regimes, build_model_slots, ) -from aca_model.config import MODEL_CONFIG, GridConfig +from aca_model.config import MODEL_AGES, GridConfig +from aca_model.execution import execution_config_for_devices def create_model( *, - n_subjects: int, fixed_params: UserParams, wage_params: Mapping[str, Any], derived_categoricals: Mapping[str, DiscreteGrid], @@ -35,11 +36,11 @@ def create_model( pref_type_grid: DiscreteGrid, solver: SolverName = "brute_force", consumption_dollars_points: tuple[float, ...] | None = None, + execution_config: ExecutionConfig | None = None, ) -> Model: """Create the baseline structural retirement model. Args: - n_subjects: Forwarded to `lcm.Model(n_subjects=...)`. fixed_params: Parameters to fix at model creation time. Fixed params are partialled into compiled functions and removed from the params template. Pass data-derived constants here; @@ -69,17 +70,17 @@ def create_model( continuous-action grid at model construction); `None` keeps the runtime-points grid completed per iteration via `inject_consumption_dollars_points`. + 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. Returns: A pylcm Model with 19 regimes (18 non-terminal + dead) spanning - ages 51-95. Regime names follow the `___` scheme. + ages 51-96; 95 is the last age anyone is alive. Regime names follow + the `___` scheme. """ - ages = AgeGrid( - start=MODEL_CONFIG.start_age, - stop=MODEL_CONFIG.end_age - 1, - step="Y", - ) + ages = AgeGrid(exact_values=MODEL_AGES) _fail_if_dcegm_without_consumption_points( solver=solver, consumption_dollars_points=consumption_dollars_points ) @@ -101,12 +102,17 @@ def create_model( return Model( regimes=regimes, + initial_nodes=INITIAL_NODES, ages=ages, regime_id_class=RegimeId, description="Baseline structural retirement model (pre-ACA)", fixed_params=fixed_params, derived_categoricals=derived_categoricals, - n_subjects=n_subjects, + 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 a0c5543..bd83aed 100644 --- a/src/aca_model/baseline/regimes/__init__.py +++ b/src/aca_model/baseline/regimes/__init__.py @@ -13,7 +13,7 @@ from collections.abc import Mapping from typing import Any -from lcm import DiscreteGrid, Regime +from lcm import DiscreteGrid, Regime, Transition from lcm.solvers import DCEGM, NBEGM from lcm.typing import UserParams @@ -21,6 +21,8 @@ from aca_model.baseline.regimes import _retiree as retiree from aca_model.baseline.regimes import _tied as tied from aca_model.baseline.regimes._common import ( + ENTRY_REGIMES, + INITIAL_NODES, REGIME_SPECS, Grids, RegimeId, @@ -36,12 +38,9 @@ from aca_model.baseline.regimes._nbegm import build_nbegm_solver from aca_model.config import GridConfig -# NBEGM is a per-regime (not global) solver: it solves one 1-D -# consumption/savings regime with at most one discrete action, so it attaches -# only to the M1 vertical-slice regime. -_NBEGM_REGIME = "nongroup_nomc_inelig_canwork" - __all__ = [ + "ENTRY_REGIMES", + "INITIAL_NODES", "REGIME_SPECS", "RegimeId", "build_all_regimes", @@ -58,6 +57,12 @@ "nongroup": nongroup.build_regime, } +_HIS_LAWS = { + "retiree": retiree.build_law, + "tied": tied.build_law, + "nongroup": nongroup.build_law, +} + def build_regime( name: str, @@ -115,7 +120,7 @@ def build_all_regimes( name, grids, dcegm_solver=dcegm_solver, - nbegm_solver=nbegm_solver if name == _NBEGM_REGIME else None, + nbegm_solver=nbegm_solver, ) regimes["dead"] = build_dead_regime(solver=solver) return regimes @@ -129,14 +134,14 @@ def build_model_slots( pref_type_grid: DiscreteGrid, solver: SolverName = "brute_force", ) -> dict[str, Any]: - """Build the model-level regime slots broadcast into every regime. + """Build the shared model topology and broadcast regime slots. - Returns keyword arguments for `lcm.Model(...)`: the functions, + Returns keyword arguments for `lcm.Model(...)`: the edges, functions, constraint, states, and laws of motion shared by all living regimes. Both the baseline and the ACA `create_model` consume this — the ACA overlay swaps only regime-level functions, so the broadcast slots are policy-invariant. Under DC-EGM the solver-contract functions join the - broadcast set and no borrowing constraint is declared. + broadcast set; the borrowing constraint is declared under every solver. """ grids = build_grids( grid_config=grid_config, @@ -145,8 +150,21 @@ def build_model_slots( pref_type_grid=pref_type_grid, ) return { + "edges": build_edges(), "functions": build_model_functions(solver=solver), - "constraints": build_model_constraints(), + "constraints": build_model_constraints(solver=solver), "states": build_model_states(grids), "state_transitions": build_model_state_transitions(), } + + +def build_edges() -> 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)) + 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 750bce4..2032ce1 100644 --- a/src/aca_model/baseline/regimes/_common.py +++ b/src/aca_model/baseline/regimes/_common.py @@ -5,6 +5,7 @@ """ import functools +import itertools from collections.abc import Callable, Iterable, Mapping from dataclasses import dataclass from types import MappingProxyType @@ -13,20 +14,37 @@ import jax.numpy as jnp import numpy as np from lcm import ( + AdditiveShockTransition, + ByAge, DiscreteGrid, + GridBreakpoint, + InitialNodes, IrregSpacedGrid, LinSpacedGrid, - MarkovTransition, NormalIIDProcess, Phased, - PiecewiseGridSegment, PiecewiseLinSpacedGrid, Regime, RouwenhorstAR1Process, + StochasticTransition, categorical, fixed_transition, ) -from lcm.typing import BoolND, FloatND, IntND, RegimeName, ScalarInt, UserParams +from lcm.consumption_savings_regime import ( + ConsumptionSavingsRegime, + LiquidMargin, + post_decision_lower_bound, +) +from lcm.solvers import OneMarginSolver +from lcm.typing import ( + BoolND, + FloatND, + IntND, + Period, + RegimeName, + ScalarInt, + UserParams, +) from aca_model.agent import ( assets_and_income, @@ -38,12 +56,40 @@ from aca_model.agent.labor_market import LaborSupply, LaggedLaborSupply, SpousalIncome from aca_model.baseline import health_insurance from aca_model.baseline.health_insurance import BuyPrivate -from aca_model.config import MODEL_CONFIG, GridConfig +from aca_model.config import MODEL_AGES, MODEL_CONFIG, GridConfig from aca_model.environment import pensions, social_security, taxes from aca_model.environment.social_security import ClaimedSS SolverName = Literal["brute_force", "dcegm", "nbegm"] +# The liquid Euler margin every EGM-solved ACA regime shares. `resources` is +# post-transfer cash-on-hand (`max(cash_on_hand, floor)`) and `savings` is the +# post-decision assets node; both are supplied by the savings-form rewiring in +# `build_dcegm_functions` / `build_nbegm_functions`, so the margin is +# declarable only on a regime that carries them. +ACA_LIQUID_MARGIN = LiquidMargin( + state="assets", + action="consumption_dollars", + resources="resources", + post_decision_state="savings", +) + + +def build_alive_regime( + *, egm_solver: OneMarginSolver | None, **regime_kwargs: Any +) -> Regime: + """Build one alive regime, declaring the liquid margin when EGM solves it. + + A brute-force regime has no `resources` or `savings` node to name, so it + stays a plain `Regime`. An EGM-solved regime owns the four margin names + and hands them to the solver, which takes numerical configuration only. + """ + if egm_solver is None: + return Regime(**regime_kwargs) + return ConsumptionSavingsRegime( + solver=egm_solver, liquid=ACA_LIQUID_MARGIN, **regime_kwargs + ) + @categorical(ordered=False) class RegimeId: @@ -194,7 +240,7 @@ class RegimeSpec(TypedDict): @dataclass(frozen=True) class Grids: - assets: LinSpacedGrid + assets: IrregSpacedGrid aime: PiecewiseLinSpacedGrid pension_wealth: LinSpacedGrid consumption_dollars: IrregSpacedGrid @@ -203,10 +249,7 @@ class Grids: hcc_transitory: Any pref_type: DiscreteGrid grid_config: GridConfig - """The originating `GridConfig`. Exposed on `Grids` so `build_states` - can read per-axis `batch_size` settings for the discrete states it - constructs inline (health, spousal_income, lagged_labor_supply, - claimed_ss) without changing the `build_states`/`build_regime` API.""" + """Grid sizes and numerical settings used to construct each solver.""" # AIME piecewise grid: number of points per segment between the PIA @@ -226,6 +269,17 @@ class Grids: _PENSION_WEALTH_GRID = LinSpacedGrid(start=0.0, stop=2_000_000.0, n_points=2) +# Asset grid: `_ASSETS_SINH_SCALE * sinh(u)` with `u` evenly spaced on each side +# of zero. The spacing is about `_ASSETS_SINH_SCALE * du` near zero, where the +# consumption floor and the SSI asset tests bite, and geometric (ratio `e^du`) for +# large holdings. The top node is the largest initial asset holding in the data. +_ASSETS_TOP = 12_000_000.0 +_ASSETS_SINH_SCALE = 10_000.0 +# Share of nodes below zero; the negative branch only has to reach the borrowing +# limit. +_ASSETS_NEGATIVE_SHARE = 5 / 24 + + # AR(1) persistence of the Rouwenhorst shocks. Calibrated once; not # routed through fixed_params because they shape the grid topology # rather than feed any DAG function. The Rouwenhorst innovation std is @@ -245,15 +299,10 @@ def build_grids( """Build continuous-state/action grids from a `GridConfig`. The AIME grid is `PiecewiseLinSpacedGrid` breakpointed at the PIA - bends from `fixed_params["pia_aime_grid"]` (total 32 points). The - assets grid's lower bound is `-max_annual_labor_income` computed + bends from `fixed_params["pia_aime_grid"]`. The assets grid + (`build_assets_grid`) runs from `-max_annual_labor_income`, computed from `wage_params` (`log_ft_wage_mean`, `log_ft_wage_std`, - `adj_wage_hours_*`). - - `wage_params` is passed separately rather than via `fixed_params` - because `log_ft_wage_mean` is a per-iteration param at estimation - time (reconstructed from `wage_bias_coeffs_*`), not a fixed one; - the grid floor must still be known at build time. + `adj_wage_hours_*`), to the largest initial asset holding. `consumption_dollars_points` fixes the consumption action grid at construction (the DC-EGM kernel needs it then); `None` keeps the @@ -270,7 +319,6 @@ def build_grids( rho=_WAGE_RHO, sigma=(1.0 - _WAGE_RHO**2) ** 0.5, mu=0.0, - batch_size=grid_config.n_wage_res_batch_size, ) hcc_persistent = get_hcc_persistent_shock(grid_config=grid_config) hcc_transitory = NormalIIDProcess( @@ -278,6 +326,7 @@ def build_grids( gauss_hermite=True, mu=0.0, sigma=1.0, + fold=grid_config.fold_hcc_transitory, ) assets_start = -_compute_max_annual_labor_income( @@ -285,13 +334,10 @@ def build_grids( ) return Grids( - assets=LinSpacedGrid( - start=assets_start, - stop=500_000.0, - n_points=grid_config.n_assets_gridpoints, - batch_size=grid_config.n_assets_batch_size, + assets=build_assets_grid( + floor=assets_start, n_points=grid_config.n_assets_gridpoints ), - aime=_build_aime_grid(grid_config=grid_config, fixed_params=fixed_params), + aime=_build_aime_grid(fixed_params=fixed_params), pension_wealth=_PENSION_WEALTH_GRID, consumption_dollars=( IrregSpacedGrid(n_points=grid_config.n_consumption_dollars_gridpoints) @@ -306,6 +352,35 @@ def build_grids( ) +def build_assets_grid(*, floor: float, n_points: int) -> IrregSpacedGrid: + """Return the sinh-spaced asset grid from `floor` to `_ASSETS_TOP`. + + `round(n_points * _ASSETS_NEGATIVE_SHARE)` nodes (at least one) lie below + zero, starting at `floor`; the rest run from exactly zero to `_ASSETS_TOP`. + On each side the nodes are `_ASSETS_SINH_SCALE * sinh(u)` with `u` evenly + spaced, so the grid is fine near zero and geometric for large holdings. + + Args: + floor: Borrowing limit, the bottom node. Must be negative. + n_points: Total number of nodes. + + Returns: + The asset grid. + + """ + n_negative = max(1, round(n_points * _ASSETS_NEGATIVE_SHARE)) + scale = _ASSETS_SINH_SCALE + negative = scale * np.sinh( + np.linspace(np.arcsinh(floor / scale), 0.0, n_negative + 1)[:-1] + ) + positive = scale * np.sinh( + np.linspace(0.0, np.arcsinh(_ASSETS_TOP / scale), n_points - n_negative) + ) + negative[0] = floor + positive[-1] = _ASSETS_TOP + return IrregSpacedGrid(points=np.concatenate([negative, positive]).tolist()) + + def get_hcc_persistent_shock(*, grid_config: GridConfig) -> RouwenhorstAR1Process: """Return the persistent-HCC AR(1) shock grid for a given `grid_config`. @@ -327,9 +402,7 @@ 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( - *, grid_config: GridConfig, fixed_params: UserParams -) -> PiecewiseLinSpacedGrid: +def _build_aime_grid(*, fixed_params: UserParams) -> PiecewiseLinSpacedGrid: """Return the AIME grid. The grid is piecewise-linspaced with breakpoints at the PIA bends @@ -340,22 +413,15 @@ def _build_aime_grid( total is fixed by the PIA structure (`sum(_AIME_PIECE_N_POINTS)`). """ kinks = [float(k) for k in np.asarray(fixed_params["pia_aime_grid"])] - segments = ( - PiecewiseGridSegment( - interval=f"[{kinks[0]}, {kinks[1]})", n_points=_AIME_PIECE_N_POINTS[0] - ), - PiecewiseGridSegment( - interval=f"[{kinks[1]}, {kinks[2]})", n_points=_AIME_PIECE_N_POINTS[1] - ), - PiecewiseGridSegment( - interval=f"[{kinks[2]}, {kinks[3]})", n_points=_AIME_PIECE_N_POINTS[2] - ), - PiecewiseGridSegment( - interval=f"[{kinks[3]}, {kinks[4]}]", n_points=_AIME_PIECE_N_POINTS[3] - ), - ) return PiecewiseLinSpacedGrid( - segments=segments, batch_size=grid_config.n_aime_batch_size + start=kinks[0], + stop=kinks[4], + breakpoints=( + GridBreakpoint(value=kinks[1]), + GridBreakpoint(value=kinks[2]), + GridBreakpoint(value=kinks[3]), + ), + points_per_segment=_AIME_PIECE_N_POINTS, ) @@ -394,35 +460,149 @@ def _compute_max_annual_labor_income( ) -_ACTIVE_PREDICATES: dict[tuple[str, str, str], Callable[..., Any]] = { - ("nomc", "inelig", "canwork"): lambda age: age < config.ss_early_age, - ("dimc", "inelig", "canwork"): lambda age: age < config.ss_early_age, - ("nomc", "choose", "canwork"): lambda age: ( - (age >= config.ss_early_age) & (age < config.medicare_age) +# 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"), ), - ("dimc", "choose", "canwork"): lambda age: ( - (age >= config.ss_early_age) & (age < config.medicare_age) + ( + (config.ss_early_age, config.medicare_age), + ("nomc_choose", "dimc_choose"), ), - ("oamc", "choose", "canwork"): lambda age: ( - (age >= config.medicare_age) & (age < config.ss_forced_age) + ( + (config.medicare_age, config.ss_forced_age), + ("forced_choose",) * 2, ), - ("oamc", "forced", "canwork"): lambda age: ( - (age >= config.ss_forced_age) & (age < config.work_forced_out_age) + ( + (config.ss_forced_age, config.work_forced_out_age), + ("forced_forced",) * 2, ), - ("oamc", "forced", "forcedout"): lambda age: ( - (age >= config.work_forced_out_age) & (age < config.end_age - 1) + ( + (config.work_forced_out_age, MODEL_AGES[-1]), + ("forcedout",) * 2, ), +) +# The living regimes a subject may start in: pre-claiming, pre-Medicare. +ENTRY_REGIMES = ( + "retiree_nomc_inelig_canwork", + "tied_nomc_inelig_canwork", + "nongroup_nomc_inelig_canwork", + "retiree_dimc_inelig_canwork", + "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 + } +) + +_STAGE_KEY = { + ("nomc", "inelig", "canwork"): "nomc_inelig", + ("dimc", "inelig", "canwork"): "dimc_inelig", + ("nomc", "choose", "canwork"): "nomc_choose", + ("dimc", "choose", "canwork"): "dimc_choose", + ("oamc", "choose", "canwork"): "forced_choose", + ("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 make_active_func(spec: RegimeSpec) -> Callable[..., Any]: - """Return the age predicate for a regime spec.""" - key = (spec["mc"], spec["ss"], spec["canwork"]) - predicate = _ACTIVE_PREDICATES.get(key) - if predicate is None: - msg = f"Unknown regime spec: {spec}" - raise ValueError(msg) - return predicate +def transition_ages(spec: RegimeSpec) -> tuple[int, ...]: + """Exact source coordinates at which the template declares its law.""" + key = _STAGE_KEY.get((spec["mc"], spec["ss"], spec["canwork"])) + for (start, stop), branches in _AGE_STAGES: + if key in branches: + return tuple(age for age in MODEL_AGES if start <= age < stop) + raise ValueError(f"Unknown regime spec: {spec}") + + +def next_model_age(period: Period) -> FloatND: + """Advance by one grid position, 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. + This does not by itself recalibrate ACA to a different period length. + """ + return jnp.asarray(_NEXT_AGES)[period] + + +def _target_pair(next_age: int, group: dict[str, int]) -> tuple[int, int]: + if next_age == MODEL_AGES[-1]: + dead = int(RegimeId.dead) + return dead, dead + for (start, stop), branches in _AGE_STAGES: + if start <= next_age < stop: + return group[branches[0]], group[branches[1]] + raise ValueError(f"Next age {next_age} is outside the declared stages") + + +def build_regime_edges( + *, + spec: RegimeSpec, + target_groups: tuple[dict[str, 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 + # non-worker who becomes disabled moves onto disability Medicare. + id_to_name = { + 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 = next_age_by_age[age] + pairs = tuple(_target_pair(next_age, group) 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: + grouped.setdefault(id_to_name[target_id], []).append(age) + return {target: tuple(source_ages) for target, source_ages in grouped.items()} + + +def build_model_edges() -> 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)) + for name, spec in REGIME_SPECS.items() + } + + +def build_regime_transition( + *, + spec: RegimeSpec, + transition_func: Callable[..., FloatND], + target_groups: tuple[dict[str, 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. + """ + edges = build_regime_edges(spec=spec, target_groups=target_groups) + 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) + ages_by_targets.setdefault(targets, []).append(age) + return ByAge( + cases={ + tuple(ages): {target: cells[target] for target in targets} + for targets, ages in ages_by_targets.items() + } + ) def build_states(spec: RegimeSpec, grids: Grids) -> dict: @@ -432,49 +612,40 @@ def build_states(spec: RegimeSpec, grids: Grids) -> dict: living regime are broadcast from the model level (`build_model_states`). """ can_work = spec["canwork"] == "canwork" - gc = grids.grid_config states: dict = {} states["health"] = DiscreteGrid( Health if spec["mc"] == "oamc" else HealthWithDisability, - batch_size=gc.n_health_batch_size, ) if can_work: states["log_ft_wage_res"] = grids.wage_res if can_work and spec["his"] != "tied": states["lagged_labor_supply"] = DiscreteGrid( LaggedLaborSupply, - batch_size=gc.n_lagged_labor_supply_batch_size, ) if spec["ss"] == "choose": states["claimed_ss"] = DiscreteGrid( ClaimedSS, - batch_size=gc.n_claimed_ss_batch_size, ) return states -def build_actions( - spec: RegimeSpec, - grids: Grids, - *, - drop_buy_private: bool = False, - drop_labor_supply: bool = False, -) -> dict: +def build_actions(spec: RegimeSpec, grids: Grids) -> dict: """Build the action dict for a non-dead regime. - The `drop_*` flags fix a discrete action to a single level for the NBEGM - M1 vertical slice (its case-piece envelope handles at most one discrete - action). The dropped action's former consumers are rebound to the fixed - level at the regime builder, so removing it here is the action side of the - dags remove-and-fix. + Every choice the regime's structure affords is a live action: whether to + claim Social Security where claiming is neither impossible nor already + forced, how many hours to work where work is still available, and whether + to buy non-group coverage before Medicare. Which solver runs the regime + does not enter — a solver that cannot carry a choice refuses the regime + rather than being handed a narrower one. """ actions: dict = {} if spec["ss"] == "choose": actions["claim_ss"] = DiscreteGrid(ClaimedSS) - if spec["canwork"] == "canwork" and not drop_labor_supply: + if spec["canwork"] == "canwork": actions["labor_supply"] = DiscreteGrid(LaborSupply) - if spec["his"] == "nongroup" and spec["mc"] == "nomc" and not drop_buy_private: + if spec["his"] == "nongroup" and spec["mc"] == "nomc": actions["buy_private"] = DiscreteGrid(BuyPrivate) actions["consumption_dollars"] = grids.consumption_dollars return actions @@ -487,24 +658,61 @@ def build_regime_probs(target: IntND, survival: FloatND) -> FloatND: return probs.at[target].add(survival) +def build_regime_probs_with_di_medicare( + *, + target_dimc: IntND, + target_nomc: IntND, + prob_dimc: FloatND, + survival: FloatND, +) -> FloatND: + """Build the regime probability vector with the disability-Medicare split. + + Survivors land in `target_dimc` with probability `prob_dimc` and in + `target_nomc` otherwise. From age 65 on both targets are the same `oamc` + regime, so the split is immaterial there. + """ + probs = jnp.zeros(19) + probs = probs.at[RegimeId.dead].set(1.0 - survival) + probs = probs.at[target_dimc].add(survival * prob_dimc) + return probs.at[target_nomc].add(survival * (1.0 - prob_dimc)) + + +def prob_di_medicare_next( + labor_supply: IntND, + prob_disabled: FloatND, +) -> FloatND: + """Probability of holding disability Medicare next period. + + Before 65, Medicare covers households that are disabled and did not work + in the previous period: the household must not work this period and be + disabled next period, which happens with probability `prob_disabled` + (`prob_disabled_next[period, health]`, 0 from 65 on, where the split does + not matter; see `build_regime_probs_with_di_medicare`). + + Callers index `prob_disabled_next` in their own body: pylcm reads the + indexing expression from the consuming function's source to lay out the + `(age, health)` Series. + """ + return jnp.where(labor_supply == LaborSupply.do_not_work, prob_disabled, 0.0) + + def build_granular_regime_transition( *, transition_func: Callable[..., FloatND], target_ids: Iterable[int], -) -> dict[RegimeName, MarkovTransition]: - """Declare the regime's reachable targets via per-target probability cells. +) -> dict[RegimeName, StochasticTransition]: + """Build the regime's per-target probability cells. Each cell evaluates the regime's probability vector and selects its - target's entry — identical arithmetic to the coarse vector form, with the - key set making every other regime structurally unreachable. + target's entry. `build_regime_transition` schedules the cells by source age. """ id_to_name = { int(getattr(RegimeId, name)): name for name in (*REGIME_SPECS, "dead") } declared = sorted({*(int(i) for i in target_ids), int(RegimeId.dead)}) return { - id_to_name[target_id]: MarkovTransition( - _prob_of_target(transition_func=transition_func, target_id=target_id) + id_to_name[target_id]: StochasticTransition( + func=_prob_of_target(transition_func=transition_func, target_id=target_id) ) for target_id in declared } @@ -525,14 +733,21 @@ def cell(*args: Any, **kwargs: Any) -> FloatND: # Broadcast functions the bequest DAG reads: the pref-type-indexed scalars # resolve their per-cell value from the broadcast `pref_type` state. _DEAD_KEEPS = frozenset( - {"consumption_weight", "coefficient_rra", "utility_scale_factor"} + { + "consumption_weight", + "coefficient_rra", + "discount_factor", + "utility_scale_factor", + } ) def build_dead_regime(*, solver: SolverName = "brute_force") -> Regime: """Build the terminal dead regime. - Everything `dead` carries arrives via the model-level broadcast: + `dead` owns two functions, the bequest `utility` and the per-type bequest + weight `scaled_bequest_weight`; everything else arrives via the + model-level broadcast: - states: `assets` and `pref_type` survive DAG pruning (the bequest utility reads them); the remaining broadcast states are pruned. @@ -540,9 +755,8 @@ def build_dead_regime(*, solver: SolverName = "brute_force") -> Regime: other broadcast function is masked with `None` so its unresolvable inputs (e.g. `pension_benefit`) don't surface as params in the dead template. - - constraints: the borrowing constraint is masked — `dead` has no - consumption action. (Under DC-EGM no constraint is broadcast, so - there is nothing to mask.) + - constraints: every broadcast constraint is masked because `dead` has no + consumption action. - `pension_wealth` is masked explicitly: a carried state is rejected in terminal regimes before pruning could drop it. """ @@ -551,13 +765,15 @@ def build_dead_regime(*, solver: SolverName = "brute_force") -> Regime: for name in build_model_functions(solver=solver) if name not in _DEAD_KEEPS } - constraint_masks = dict.fromkeys(build_model_constraints()) + constraint_masks = dict.fromkeys(build_model_constraints(solver=solver)) return Regime( - transition=None, - functions={"utility": preferences.bequest, **function_masks}, + functions={ + "utility": preferences.bequest, + "scaled_bequest_weight": preferences.scaled_bequest_weight, + **function_masks, + }, constraints=constraint_masks, states={"pension_wealth": None}, - active=lambda _age: True, ) @@ -588,9 +804,11 @@ def build_model_functions(*, solver: SolverName = "brute_force") -> dict: Contains exactly the functions that are identical across all 18 living regimes AND are never swapped by the ACA policy overlay. Spec-dependent - selections (`good_health`, `leisure`, …) and overlay-swapped names - (`is_medicaid_eligible`, `cash_on_hand`, `primary_oop`) stay regime-level - in `build_common_functions`. The `dead` regime masks every entry the + selections (`good_health`, `leisure`, `oop_costs`, …) and overlay-swapped + names (`is_medicaid_eligible`, `cash_on_hand`, `primary_oop`, + `after_tax_income`) stay regime-level in `build_common_functions`: pylcm + refuses a name defined at both levels, so a regime can only swap a + regime-level entry. The `dead` regime masks every entry the bequest DAG does not read (see `build_dead_regime`). Under DC-EGM the solver-contract functions join the broadcast set. """ @@ -603,9 +821,8 @@ def build_model_functions(*, solver: SolverName = "brute_force") -> dict: # `marginal_continuation`. functions |= build_dcegm_functions() functions["total_health_costs"] = health_insurance.total_costs - functions["oop_costs"] = health_insurance.oop_with_medicaid functions["capital_income"] = assets_and_income.capital_income - # spousal_income_amounts is a lookup table param, not a DAG function + functions["spousal_income_amount"] = labor_market.spousal_income_amount functions["is_married"] = labor_market.is_married functions["equivalence_scale"] = preferences.equivalence_scale functions["utility_scale_factor"] = preferences.utility_scale_factor @@ -624,7 +841,6 @@ def build_model_functions(*, solver: SolverName = "brute_force") -> dict: # Taxes functions["taxable_ss_benefit"] = taxes.taxable_ss_benefit functions["gross_income"] = taxes.gross_income - functions["after_tax_income"] = taxes.after_tax_income # Every living regime carries pension wealth and the solve-phase pension # assets adjustment, both of which scale by the marginal income tax rate. functions["marginal_tax_rate"] = taxes.marginal_rate @@ -667,16 +883,30 @@ def build_nbegm_functions() -> dict: } -def build_model_constraints() -> dict: +def build_model_constraints(*, solver: SolverName) -> dict: """Build the model-level constraints broadcast into every regime. - `dead` masks the borrowing constraint — it has no consumption action. - The constraint is broadcast under every solver: an EGM solve (DC-EGM or - NBEGM) enforces the borrowing limit through the savings grid's lower - bound, but forward simulation re-decides consumption by an argmax over - the consumption grid and needs the explicit feasibility mask. + `dead` masks every constraint because it has no consumption action. + + Grid search evaluates the action-level predicates directly: the borrowing + constraint and the household's own consumption floor. An EGM-family + solve proves the equivalent post-decision lower bound from its savings + grid, while forward simulation receives the solver's intrinsic budget + mask over the consumption grid; it does not carry the own-floor bound. """ - return {"borrowing_constraint": assets_and_income.borrowing_constraint} + borrowing_constraint = ( + assets_and_income.borrowing_constraint + if solver == "brute_force" + else post_decision_lower_bound(margin=ACA_LIQUID_MARGIN, lower=0.0) + ) + constraints = { + "borrowing_constraint": borrowing_constraint, + } + if solver == "brute_force": + constraints["consumption_floor_constraint"] = ( + assets_and_income.consumption_above_floor + ) + return constraints def build_model_states(grids: Grids) -> dict: @@ -684,10 +914,9 @@ def build_model_states(grids: Grids) -> dict: These are the states every living regime carries with an identical grid. pylcm prunes them per regime by DAG reachability, so `dead` keeps only - `assets` and `pref_type` (the bequest DAG). `spousal_income` carries the - `distributed` flag — sharding is legal only on model-level states. + `assets` and `pref_type` (the bequest DAG). Placement is configured + separately through the model execution policy. """ - gc = grids.grid_config return { "assets": grids.assets, "aime": grids.aime, @@ -699,8 +928,6 @@ def build_model_states(grids: Grids) -> dict: "hcc_transitory": grids.hcc_transitory, "spousal_income": DiscreteGrid( SpousalIncome, - batch_size=gc.n_spousal_income_batch_size, - distributed=gc.spousal_income_distributed, ), "pref_type": grids.pref_type, } @@ -716,14 +943,14 @@ def build_model_state_transitions() -> dict: """ return { "pref_type": fixed_transition("pref_type"), - "spousal_income": MarkovTransition(labor_market.next_spousal_income), + "spousal_income": StochasticTransition(func=labor_market.next_spousal_income), # Carried state: evolved only in simulate (in solve, `pension_wealth` # is re-imputed from AIME each period and has no transition). "pension_wealth": pensions.wealth_next_before_adjustment, } -def build_common_functions(spec: RegimeSpec) -> dict: +def build_common_functions(spec: RegimeSpec, *, subtract_bill: bool = False) -> dict: """Build the regime-level functions dict for a non-dead regime. Contains the spec-dependent selections and the overlay-swapped names; @@ -739,7 +966,16 @@ def build_common_functions(spec: RegimeSpec) -> dict: ) has_buy_private = spec["his"] == "nongroup" and spec["mc"] == "nomc" functions["primary_oop"] = ( - health_insurance.primary_oop if has_buy_private else health_insurance.oop_costs + health_insurance.primary_oop + if has_buy_private + else health_insurance.insured_oop + ) + # Medicaid pays on top of the primary cover, except in the non-group + # regimes before Medicare, where it replaces private or no cover. + functions["oop_costs"] = ( + health_insurance.oop_with_medicaid_replacing_private + if has_buy_private + else health_insurance.oop_with_medicaid ) if can_work: @@ -768,6 +1004,13 @@ def build_common_functions(spec: RegimeSpec) -> dict: # Swapped per policy variant by the ACA overlay, hence regime-level functions["is_medicaid_eligible"] = health_insurance.is_medicaid_eligible + functions["after_tax_income"] = taxes.after_tax_income + if subtract_bill: + functions["assets_before_bill"] = assets_and_income.assets_before_bill + functions["negative_oop_costs"] = health_insurance.negative_oop_costs + functions["assets_before_bill_when_dead"] = ( + assets_and_income.assets_before_bill_when_dead + ) functions["premium_default"] = assets_and_income.premium_default functions["cash_on_hand"] = assets_and_income.cash_on_hand @@ -819,6 +1062,10 @@ def build_pension_functions(spec: RegimeSpec) -> dict: `pension_assets_adjustment`, which reconciles the accrual-evolved pension with next period's AIME imputation; in simulate the adjustment is zero because the true wealth is carried directly. + - in `ss=choose` regimes, `carried_pia` is the PIA the next-period AIME + encodes: in solve the claim-adjusted PIA corrected so that SS plus + imputed pension benefits keep their adjusted total + (`pensions.total_to_pia`); in simulate the claim-adjusted PIA itself. """ can_work = spec["canwork"] == "canwork" @@ -836,10 +1083,21 @@ def build_pension_functions(spec: RegimeSpec) -> dict: if can_work else health_insurance.target_his_forcedout ) - functions["pia_unadjusted_next_period"] = social_security.pia_unadjusted_next_period + functions["pia_unadjusted_next_period"] = ( + social_security.pia_unadjusted_next_period_forced + if spec["ss"] == "forced" + else social_security.pia_unadjusted_next_period + ) + functions["full_benefit_next_period"] = pensions.full_benefit_next_period functions["imputed_pension_wealth_next_period"] = ( pensions.imputed_pension_wealth_next_period ) + if spec["ss"] == "choose": + functions["pia_adjusted_next_period"] = social_security.pia_adjusted_next_period + functions["carried_pia"] = Phased( + solve=pensions.total_to_pia, + simulate=social_security.carried_pia_simulate, + ) functions["pension_assets_adjustment"] = Phased( solve=pensions.assets_adjustment, simulate=_zero_pension_assets_adjustment, @@ -912,38 +1170,17 @@ def select_target_for_age( mc_next: bool | BoolND, tgts: dict[str, int], ) -> IntND: - """Select target regime ID based on next-period age bracket.""" - ss_choose = jnp.where( - jnp.array(mc_next), - tgts["dimc_choose"], - tgts["nomc_choose"], - ) - ss_inelig = jnp.where( - jnp.array(mc_next), - tgts["dimc_inelig"], - tgts["nomc_inelig"], - ) - return jnp.where( - next_age >= config.end_age - 1, - RegimeId.dead, - jnp.where( - next_age >= config.work_forced_out_age, - tgts["forcedout"], - jnp.where( - next_age >= config.ss_forced_age, - tgts["forced_forced"], - jnp.where( - next_age >= config.medicare_age, - tgts["forced_choose"], - jnp.where(next_age >= config.ss_early_age, ss_choose, ss_inelig), - ), - ), - ), - ) + """Select with the same stage table that supplies the schedule's support.""" + result = jnp.asarray(RegimeId.dead) + for (start, stop), branches in reversed(_AGE_STAGES): + 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) + return result def build_state_transitions( - spec: RegimeSpec, *, solver: SolverName = "brute_force" + spec: RegimeSpec, *, solver: SolverName = "brute_force", subtract_bill: bool = False ) -> dict: """Build the regime-level state transitions dict for a non-dead regime. @@ -954,7 +1191,9 @@ def build_state_transitions( post-decision (savings) form. """ transitions: dict = {} - transitions["assets"] = _build_per_target_regime_assets(spec, solver=solver) + transitions["assets"] = _build_per_target_regime_assets( + spec, solver=solver, subtract_bill=subtract_bill + ) transitions["health"] = _build_per_target_regime_health(spec) claimed_ss_transition = _build_per_target_regime_claimed_ss(spec) if claimed_ss_transition: @@ -971,20 +1210,19 @@ def _select_aime_law(spec: RegimeSpec) -> Callable[..., FloatND]: The claim-age actuarial bake applies only where the agent chooses when to claim (`ss=choose`), so only those regimes carry the `claim_ss`/`claimed_ss` - inputs. `ss=inelig` (cannot claim) and `ss=forced` (claims by rule) use the - plain-accrual variant: no claim adjustment, no claim inputs. A forced - claimant who claimed early carries the reduction in from the choose regime; - plain accrual preserves it. + inputs. `ss=inelig` (cannot claim) uses plain accrual with no claim inputs; + in `ss=forced` (claims by rule) AIME stays as carried in, with any early + reduction or delayed credit from the choose regime. - post-65 (`oamc`), `choose` → `next_aime` (claim-adjusted) - - post-65 (`oamc`), `forced` → `next_aime_plain` + - post-65 (`oamc`), `forced` → `next_aime_forced` (AIME frozen) - pre-65 (`nomc`/`dimc`), `choose` → `next_aime_disabled` (claim-adjusted) - pre-65 (`nomc`/`dimc`), `inelig` → `next_aime_disabled_plain` """ is_choose = spec["ss"] == "choose" if spec["mc"] == "oamc": return ( - social_security.next_aime if is_choose else social_security.next_aime_plain + social_security.next_aime if is_choose else social_security.next_aime_forced ) return ( social_security.next_aime_disabled @@ -993,8 +1231,24 @@ def _select_aime_law(spec: RegimeSpec) -> Callable[..., FloatND]: ) +# Every input of the out-of-pocket bill that varies across a source's states and +# actions, with its support. The bill reads nothing else besides next period's +# medical-cost draws, the period and parameters. +_BILL_CONDITIONERS: Mapping[str, tuple[int | bool, ...]] = MappingProxyType( + { + "is_married": (0, 1), + "good_health": (0, 1), + "is_medicaid_eligible": (False, True), + "buy_private": (int(BuyPrivate.no), int(BuyPrivate.yes)), + # Under the ACA policy: the four FPL kinks of the cost-sharing schedule + # give brackets 0-4, and 5 is the neutral bracket. + "cost_sharing_bracket": tuple(range(6)), + } +) + + def _build_per_target_regime_assets( - spec: RegimeSpec, *, solver: SolverName = "brute_force" + spec: RegimeSpec, *, solver: SolverName = "brute_force", subtract_bill: bool = False ) -> dict[RegimeName, Callable[..., FloatND]]: """Build per-target assets transitions. @@ -1003,11 +1257,30 @@ def _build_per_target_regime_assets( pull in the `next_aime`-dependent imputation chain — `dead` has no `aime` state and pylcm cannot resolve `next_aime` there. Non-dead targets use the full `next_assets` with the pension correction. - Under DC-EGM both laws take their post-decision (savings) form. + Under DC-EGM both laws take their post-decision (savings) form. With + `subtract_bill`, each law is declared as its resources before the bill + (`assets_before_bill`, `assets_before_bill_when_dead`) plus the shock + `negative_oop_costs`, conditioned on every source-side input of the bill. """ + if subtract_bill and solver != "brute_force": + msg = f"subtract_hcc_bill needs the brute-force asset law, not {solver!r}." + raise ValueError(msg) + living_law: Callable[..., FloatND] + dead_law: Callable[..., FloatND] if solver in ("dcegm", "nbegm"): living_law = assets_and_income.next_assets_from_savings dead_law = assets_and_income.next_assets_when_dead_from_savings + elif subtract_bill: + living_law = AdditiveShockTransition( + base="assets_before_bill", + shock="negative_oop_costs", + conditioners=_BILL_CONDITIONERS, + ) + dead_law = AdditiveShockTransition( + base="assets_before_bill_when_dead", + shock="negative_oop_costs", + conditioners=_BILL_CONDITIONERS, + ) else: living_law = assets_and_income.next_assets dead_law = assets_and_income.next_assets_when_dead @@ -1033,16 +1306,21 @@ def _build_per_target_regime_assets( def _build_per_target_regime_health( spec: RegimeSpec, -) -> dict[RegimeName, MarkovTransition]: +) -> dict[RegimeName, StochasticTransition]: """Build per-target health transitions. Pre-65 regimes use HealthWithDisability (3-state), post-65 use Health (2-state). - Cross-grid transitions (3->2) happen at the age-65 boundary. + Cross-grid transitions (3->2) happen at the age-65 boundary. Between pre-65 + regimes the Medicare outcome reveals next-period disability for + non-workers, so each target carries the health law conditional on it: + + - `dimc` target: disabled (`next_health_into_dimc`); + - `nomc` target: not disabled if not working (`next_health_into_nomc`). """ target_regimes = precompute_target_regimes(spec) id_to_name = {int(getattr(RegimeId, name)): name for name in REGIME_SPECS} - result: dict[RegimeName, MarkovTransition] = {} + result: dict[RegimeName, StochasticTransition] = {} seen_ids: set[int] = set() for target_id in target_regimes.values(): @@ -1052,13 +1330,19 @@ def _build_per_target_regime_health( target_name = id_to_name.get(target_id) if target_name is None: continue - target_spec = REGIME_SPECS[target_name] - target_is_post65 = target_spec["mc"] == "oamc" - - if spec["mc"] != "oamc" and target_is_post65: - result[target_name] = MarkovTransition(health.next_health_cross) + target_mc = REGIME_SPECS[target_name]["mc"] + if spec["mc"] == "oamc": + result[target_name] = StochasticTransition(func=health.next_health) + elif target_mc == "oamc": + result[target_name] = StochasticTransition(func=health.next_health_cross) + elif target_mc == "dimc": + result[target_name] = StochasticTransition( + func=health.next_health_into_dimc + ) else: - result[target_name] = MarkovTransition(health.next_health) + result[target_name] = StochasticTransition( + func=health.next_health_into_nomc + ) return result diff --git a/src/aca_model/baseline/regimes/_dcegm.py b/src/aca_model/baseline/regimes/_dcegm.py index b0ee129..54e3ab2 100644 --- a/src/aca_model/baseline/regimes/_dcegm.py +++ b/src/aca_model/baseline/regimes/_dcegm.py @@ -32,16 +32,8 @@ def build_dcegm_solver(grids: Grids) -> DCEGM: 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)), - batch_size=grids.grid_config.n_savings_batch_size, - ) - return DCEGM( - continuous_state="assets", - continuous_action="consumption_dollars", - resources="resources", - post_decision_function="savings", - savings_grid=savings_grid, - stochastic_node_batch_size=grids.grid_config.n_stochastic_node_batch_size, ) + return DCEGM(savings_grid=savings_grid) def _fail_if_too_few_savings_gridpoints(n_savings_gridpoints: int) -> None: diff --git a/src/aca_model/baseline/regimes/_nbegm.py b/src/aca_model/baseline/regimes/_nbegm.py index d1cda4f..23fb96a 100644 --- a/src/aca_model/baseline/regimes/_nbegm.py +++ b/src/aca_model/baseline/regimes/_nbegm.py @@ -1,15 +1,7 @@ -"""NBEGM solver configuration for the ACA M1 vertical-slice regime. +"""NB-EGM numerical configuration for the structural retirement model. -NBEGM is the case-piece endogenous-grid solver for a single 1-D -consumption/savings regime whose budget is split by institutional breakpoints -on a derived monotone income quantity. It shares DC-EGM's post-decision -(savings) spec — consumption is recovered from `resources = max(cash_on_hand, -floor)`, the assets laws are in savings form, the borrowing constraint is the -savings grid's lower bound — but solves only one regime with at most one -discrete action, so it attaches per regime rather than globally. The -function-level rewiring is shared with DC-EGM (`build_dcegm_functions`, the -savings-form assets laws in `_common`); this module holds only the solver -config. +Savings gridpoints, cliff reads, and envelope arithmetic belong to the solver. +Device placement and compiled program widths belong to the model execution policy. """ from lcm import IrregSpacedGrid @@ -19,53 +11,17 @@ def build_nbegm_solver(grids: Grids) -> NBEGM: - """Build the per-regime NBEGM configuration. - - The savings grid mirrors DC-EGM's: lower bound 0 (the borrowing constraint - in post-decision form), upper bound the assets span, cubically clustered - toward the constraint. The budget node is `resources` (post-floor - cash-on-hand) and the post-decision function is `savings`, matching the - shared savings-form spec. - """ + """Build NB-EGM with the model's savings grid and numerical read rules.""" n_points = grids.grid_config.n_savings_gridpoints _fail_if_too_few_savings_gridpoints(n_points) 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)), - batch_size=grids.grid_config.n_savings_batch_size, ) return NBEGM( savings_grid=savings_grid, - continuous_state="assets", - budget_target="resources", - post_decision_function="savings", - # The budget DAG mixes scalar- and array-annotated parameters (tax - # tables, threshold schedules) that the build-time derivative probes - # cannot synthesize, so affinity/constancy is asserted here and - # validated by the full-model brute-agreement gates. - probe_failure="assume_declared", - # Splay the child stochastic-node expectation per the grid config: `0` (the - # default) reads the whole node mesh in one pass on a memory-rich device; a - # positive value loops it in blocks to fit a tighter budget (a CPU run). - stochastic_node_batch_size=grids.grid_config.n_nbegm_stochastic_node_batch_size, - # Stream the per-interval upper envelope over candidate-segment blocks per - # the grid config; `0` keeps the one-shot dense envelope. - envelope_segment_block_size=( - grids.grid_config.n_nbegm_envelope_segment_block_size - ), - # Which arithmetic decides envelope ownership per the grid config; - # "certified" is exact and can abstain, "ordinary" reads in the working - # format at a fraction of the cost. envelope_arithmetic=grids.grid_config.nbegm_envelope_arithmetic, - # Stream both ride-along cores over ride-cell blocks per the grid config; - # `0` vmaps the whole flattened mesh at once. - cell_block_size=grids.grid_config.n_nbegm_cell_block_size, - # Stream the discrete-action branch axis in blocks per the grid config; - # `0` runs the whole axis in one vectorized pass. - branch_batch_size=grids.grid_config.n_nbegm_branch_batch_size, - # Cliff-read mode: exact one-sided limits (default) or the fast bridged - # read for inner estimation loops (see `GridConfig.nbegm_jump_read`). jump_read=grids.grid_config.nbegm_jump_read, ) diff --git a/src/aca_model/baseline/regimes/_nongroup.py b/src/aca_model/baseline/regimes/_nongroup.py index 8d329ea..a346e0c 100644 --- a/src/aca_model/baseline/regimes/_nongroup.py +++ b/src/aca_model/baseline/regimes/_nongroup.py @@ -4,55 +4,67 @@ Already nongroup, so no SSI/Medicaid override needed for HIS transitions. """ -import functools from collections.abc import Callable -from lcm import Regime +from lcm import ByAge, Regime from lcm.solvers import DCEGM, NBEGM -from lcm.typing import Age, DiscreteAction, FloatND, Period +from lcm.typing import ( + Age, + DiscreteAction, + DiscreteState, + FloatND, + Period, +) -from aca_model.agent.labor_market import LaborSupply from aca_model.baseline import health_insurance -from aca_model.baseline.health_insurance import BuyPrivate from aca_model.baseline.regimes._common import ( REGIME_SPECS, Grids, RegimeSpec, build_actions, + build_alive_regime, build_common_functions, - build_granular_regime_transition, build_nbegm_functions, build_pension_functions, build_regime_probs, + build_regime_probs_with_di_medicare, + build_regime_transition, build_state_transitions, build_states, - make_active_func, make_targets, + next_model_age, + prob_di_medicare_next, select_ss_benefit, select_target_for_age, ) def _make_transition_canwork( - gets_medicare: bool, own: dict[str, int], ) -> Callable[..., FloatND]: """Create transition for canwork nongroup regimes. - Already nongroup — no SSI override needed. Gets Medicare if stops working - (when gets_medicare is True). + Already nongroup — no SSI override needed. Before 65, a household that + does not work and is disabled next period holds disability Medicare. """ def transition( age: Age, period: Period, + health: DiscreteState, labor_supply: DiscreteAction, survival_probs: FloatND, + prob_disabled_next: FloatND, ) -> FloatND: - sp = survival_probs[period] - mc_next = gets_medicare & (labor_supply == LaborSupply.do_not_work) - target = select_target_for_age(age + 1, mc_next, own) - return build_regime_probs(target, sp) + 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), + prob_dimc=prob_di_medicare_next( + labor_supply, prob_disabled_next[period, health] + ), + survival=survival_probs[period, health], + ) return transition @@ -70,66 +82,75 @@ def _make_transition_forcedout( def transition( age: Age, period: Period, + health: DiscreteState, survival_probs: FloatND, ) -> FloatND: - target = select_target_for_age(age + 1, gets_medicare, own) - return build_regime_probs(target, survival_probs[period]) + del age # Keep the legacy signature; period owns the clock lookup. + target = select_target_for_age(next_model_age(period), gets_medicare, own) + return build_regime_probs(target, survival_probs[period, health]) return transition -def _fixed_full_time_labor_supply() -> DiscreteAction: - """Labor supply fixed to full-time work for the NBEGM M1 slice.""" - return LaborSupply.h2000 - - -def _build_functions( - spec: RegimeSpec, *, fix_buy_private: bool = False, fix_labor_supply: bool = False -) -> dict: - """Build functions dict for a nongroup regime. - - The NBEGM M1 slice fixes both discrete actions to a single level so the - only choice is continuous consumption: - - - `fix_buy_private` binds `buy_private` to `BuyPrivate.yes` in its consumers - (premium, OOP) — the `buy_private == BuyPrivate.yes` arm — leaving the - remaining budget structure untouched. - - `fix_labor_supply` supplies `labor_supply` as a fixed full-time node read - by labor income, AIME accrual, and the lagged-supply transition (which - stays a state, so the cross-regime continuation space is unchanged). - """ +def _build_functions(spec: RegimeSpec, *, subtract_bill: bool = False) -> dict: + """Build functions dict for a nongroup regime.""" can_work = spec["canwork"] == "canwork" - functions = build_common_functions(spec) - - if can_work and fix_labor_supply: - functions["labor_supply"] = _fixed_full_time_labor_supply + functions = build_common_functions(spec, subtract_bill=subtract_bill) functions["ss_benefit"] = select_ss_benefit(spec) # his and crossed_oamc_threshold are fixed params (constants per regime), # not DAG functions. pylcm resolves them from the params dict. - has_buy_private = spec["his"] == "nongroup" and spec["mc"] == "nomc" - if has_buy_private: - functions["hic_premium"] = health_insurance.premium + if _has_buy_private(spec): + functions["plan_premium"] = health_insurance.premium + functions["private_premium_intercept"] = ( + health_insurance.private_premium_intercept + ) elif can_work: - functions["hic_premium"] = health_insurance.premium_insured + functions["plan_premium"] = health_insurance.premium_insured else: - functions["hic_premium"] = health_insurance.premium_retired - - if has_buy_private and fix_buy_private: - functions["hic_premium"] = functools.partial( - health_insurance.premium, buy_private=BuyPrivate.yes - ) - functions["primary_oop"] = functools.partial( - health_insurance.primary_oop, buy_private=BuyPrivate.yes - ) + functions["plan_premium"] = health_insurance.premium_retired + # Medicaid-eligible households pay no premium: Medicaid replaces private + # cover and pays the Medicare premium. + functions["hic_premium"] = health_insurance.medicaid_adjusted_premium functions.update(build_pension_functions(spec)) return functions +def _has_buy_private(spec: RegimeSpec) -> bool: + return spec["his"] == "nongroup" and spec["mc"] == "nomc" + + +def build_constraints(spec: RegimeSpec) -> dict: + """Build the regime's choice constraints. + + Where private cover is a choice, it is feasible only if its premium is + paid in full (`private_cover_paid_in_full`). + """ + if _has_buy_private(spec): + return { + "private_cover_paid_in_full": health_insurance.private_cover_paid_in_full + } + return {} + + +def build_law(name: str) -> ByAge: + """Build the nongroup regime's per-target transition probabilities by source age.""" + spec = REGIME_SPECS[name] + gets_mc = spec["mc"] != "nomc" + own, _ng = make_targets(name) + if spec["canwork"] == "canwork": + transition_func = _make_transition_canwork(own) + else: + transition_func = _make_transition_forcedout(gets_mc, own) + return build_regime_transition( + spec=spec, transition_func=transition_func, target_groups=(own,) + ) + + def build_regime( name: str, grids: Grids, @@ -139,35 +160,21 @@ def build_regime( ) -> Regime: """Build a nongroup regime.""" spec = REGIME_SPECS[name] - gets_mc = spec["mc"] != "nomc" - own, _ng = make_targets(name) - - if spec["canwork"] == "canwork": - transition_func = _make_transition_canwork(gets_mc, own) - else: - transition_func = _make_transition_forcedout(gets_mc, own) - states = build_states(spec, grids) egm_solver = dcegm_solver if dcegm_solver is not None else nbegm_solver - solver_kwargs: dict = {} if egm_solver is None else {"solver": egm_solver} state_solver = ( "brute_force" if egm_solver is None else ("nbegm" if nbegm_solver is not None else "dcegm") ) - # Under NBEGM the M1 slice fixes `buy_private` (a second discrete action the - # branch compiler does not yet solve) to a single level. `labor_supply` is fixed - # too by default, leaving only continuous consumption; with - # `nbegm_live_labor_supply` it stays a live action and the branch compiler solves - # each labor level against the cliffed budget. - fix_for_nbegm = nbegm_solver is not None - fix_labor = fix_for_nbegm and not grids.grid_config.nbegm_live_labor_supply functions = _build_functions( - spec, fix_buy_private=fix_for_nbegm, fix_labor_supply=fix_labor + spec, subtract_bill=grids.grid_config.subtract_hcc_bill ) - constraints: dict = {} - if fix_for_nbegm: + # The EGM solvers reject constraints that read the continuous state, which + # the private-cover constraint does through `premium_default`. + constraints = build_constraints(spec) if egm_solver is None else {} + if nbegm_solver is not None: # NBEGM solves only this regime, so its solver-contract functions are # regime-level here rather than broadcast model-wide. The broadcast # borrowing constraint stays: the EGM solve enforces the limit through @@ -175,20 +182,15 @@ def build_regime( # consumption by an argmax over the consumption grid and needs the # explicit feasibility mask. functions = {**functions, **build_nbegm_functions()} - return Regime( - transition=build_granular_regime_transition( - transition_func=transition_func, target_ids=own.values() - ), - active=make_active_func(spec), + return build_alive_regime( + egm_solver=egm_solver, states=states, - state_transitions=build_state_transitions(spec, solver=state_solver), - actions=build_actions( + state_transitions=build_state_transitions( spec, - grids, - drop_buy_private=fix_for_nbegm, - drop_labor_supply=fix_labor, + solver=state_solver, + subtract_bill=grids.grid_config.subtract_hcc_bill, ), + actions=build_actions(spec, grids), functions=functions, constraints=constraints, - **solver_kwargs, ) diff --git a/src/aca_model/baseline/regimes/_retiree.py b/src/aca_model/baseline/regimes/_retiree.py index 3ac87d2..a0f304b 100644 --- a/src/aca_model/baseline/regimes/_retiree.py +++ b/src/aca_model/baseline/regimes/_retiree.py @@ -1,62 +1,85 @@ """Regime transitions and builder for retiree HIS regimes. Retiree regimes: agents with employer-sponsored retiree health insurance. -Medicaid-eligible agents are overridden to nongroup. +Categorically (SSI-) Medicaid-eligible agents are overridden to nongroup; the +ACA Medicaid expansion leaves employer coverage untouched. """ from collections.abc import Callable import jax.numpy as jnp -from lcm import Regime +from lcm import ByAge, Regime from lcm.solvers import DCEGM, NBEGM -from lcm.typing import Age, BoolND, DiscreteAction, FloatND, Period +from lcm.typing import ( + Age, + BoolND, + DiscreteAction, + DiscreteState, + FloatND, + IntND, + Period, +) -from aca_model.agent.labor_market import LaborSupply from aca_model.baseline import health_insurance from aca_model.baseline.regimes._common import ( REGIME_SPECS, Grids, RegimeSpec, build_actions, + build_alive_regime, build_common_functions, - build_granular_regime_transition, + build_nbegm_functions, build_pension_functions, build_regime_probs, + build_regime_probs_with_di_medicare, + build_regime_transition, build_state_transitions, build_states, - make_active_func, make_targets, + next_model_age, + prob_di_medicare_next, select_ss_benefit, select_target_for_age, ) def _make_transition_canwork( - gets_medicare: bool, own: dict[str, int], ng: dict[str, int], ) -> Callable[..., FloatND]: """Create transition for canwork retiree regimes. - Retirees who stop working get Medicare (if gets_medicare). - Medicaid-eligible agents are overridden to nongroup targets. + Before 65, a household that does not work and is disabled next period + holds disability Medicare. Categorically (SSI-) Medicaid-eligible + agents are overridden to + nongroup targets. """ def transition( age: Age, period: Period, + health: DiscreteState, labor_supply: DiscreteAction, - is_medicaid_eligible: BoolND, + is_ssi_eligible: BoolND, survival_probs: FloatND, + prob_disabled_next: FloatND, ) -> FloatND: - sp = survival_probs[period] - next_age = age + 1 - mc_next = gets_medicare & (labor_supply == LaborSupply.do_not_work) - target = select_target_for_age(next_age, mc_next, own) - # Medicaid eligibility overrides to nongroup - ng_ssi = select_target_for_age(next_age, mc_next, ng) - target = jnp.where(is_medicaid_eligible, ng_ssi, target) - return build_regime_probs(target, sp) + del age # Keep the legacy signature; period owns the clock lookup. + next_age = next_model_age(period) + + 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) + return jnp.where(is_ssi_eligible, ng_target, own_target) + + return build_regime_probs_with_di_medicare( + target_dimc=target(True), + target_nomc=target(False), + prob_dimc=prob_di_medicare_next( + labor_supply, prob_disabled_next[period, health] + ), + survival=survival_probs[period, health], + ) return transition @@ -68,29 +91,32 @@ def _make_transition_forcedout( ) -> Callable[..., FloatND]: """Create transition for forcedout retiree regimes. - No labor supply action. Medicaid-eligible agents are overridden to nongroup. + No labor supply action. Categorically (SSI-) Medicaid-eligible agents are + overridden to nongroup. """ def transition( age: Age, period: Period, - is_medicaid_eligible: BoolND, + health: DiscreteState, + is_ssi_eligible: BoolND, survival_probs: FloatND, ) -> FloatND: - sp = survival_probs[period] - next_age = age + 1 + 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) - target = jnp.where(is_medicaid_eligible, ng_ssi, target) + target = jnp.where(is_ssi_eligible, ng_ssi, target) return build_regime_probs(target, sp) return transition -def _build_functions(spec: RegimeSpec) -> dict: +def _build_functions(spec: RegimeSpec, *, subtract_bill: bool = False) -> dict: """Build functions dict for a retiree regime.""" can_work = spec["canwork"] == "canwork" - functions = build_common_functions(spec) + functions = build_common_functions(spec, subtract_bill=subtract_bill) functions["ss_benefit"] = select_ss_benefit(spec) @@ -107,6 +133,20 @@ def _build_functions(spec: RegimeSpec) -> dict: return functions +def build_law(name: str) -> ByAge: + """Build the retiree regime's per-target transition probabilities by source age.""" + 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) + else: + transition_func = _make_transition_forcedout(gets_mc, own, ng) + return build_regime_transition( + spec=spec, transition_func=transition_func, target_groups=(own, ng) + ) + + def build_regime( name: str, grids: Grids, @@ -116,31 +156,30 @@ def build_regime( ) -> Regime: """Build a retiree regime.""" spec = REGIME_SPECS[name] - gets_mc = spec["mc"] != "nomc" - own, ng = make_targets(name) - - if spec["canwork"] == "canwork": - transition_func = _make_transition_canwork(gets_mc, own, ng) - else: - transition_func = _make_transition_forcedout(gets_mc, own, ng) - states = build_states(spec, grids) egm_solver = dcegm_solver if dcegm_solver is not None else nbegm_solver - solver_kwargs: dict = {} if egm_solver is None else {"solver": egm_solver} state_solver = ( "brute_force" if egm_solver is None else ("nbegm" if nbegm_solver is not None else "dcegm") ) - return Regime( - transition=build_granular_regime_transition( - transition_func=transition_func, target_ids=(*own.values(), *ng.values()) - ), - active=make_active_func(spec), + functions = _build_functions( + spec, subtract_bill=grids.grid_config.subtract_hcc_bill + ) + if nbegm_solver is not None: + # NBEGM's solver contract is stated per regime: it reads the budget in + # savings form off `resources` and the post-decision node off + # `savings`, neither of which the brute-force build needs. + functions = {**functions, **build_nbegm_functions()} + return build_alive_regime( + egm_solver=egm_solver, states=states, - state_transitions=build_state_transitions(spec, solver=state_solver), + state_transitions=build_state_transitions( + spec, + solver=state_solver, + subtract_bill=grids.grid_config.subtract_hcc_bill, + ), actions=build_actions(spec, grids), - functions=_build_functions(spec), - **solver_kwargs, + functions=functions, ) diff --git a/src/aca_model/baseline/regimes/_tied.py b/src/aca_model/baseline/regimes/_tied.py index 08a2196..9a9c4b3 100644 --- a/src/aca_model/baseline/regimes/_tied.py +++ b/src/aca_model/baseline/regimes/_tied.py @@ -2,15 +2,24 @@ Tied regimes: agents with employer-tied health insurance. Tied agents who stop working become nongroup. -Medicaid-eligible agents are also overridden to nongroup. +Categorically (SSI-) Medicaid-eligible agents are also overridden to nongroup; +the ACA Medicaid expansion leaves employer coverage untouched. """ from collections.abc import Callable import jax.numpy as jnp -from lcm import Regime +from lcm import ByAge, Regime from lcm.solvers import DCEGM, NBEGM -from lcm.typing import Age, BoolND, DiscreteAction, FloatND, Period +from lcm.typing import ( + Age, + BoolND, + DiscreteAction, + DiscreteState, + FloatND, + IntND, + Period, +) from aca_model.agent.labor_market import LaborSupply from aca_model.baseline import health_insurance @@ -19,56 +28,67 @@ Grids, RegimeSpec, build_actions, + build_alive_regime, build_common_functions, - build_granular_regime_transition, + build_nbegm_functions, build_pension_functions, - build_regime_probs, + build_regime_probs_with_di_medicare, + build_regime_transition, build_state_transitions, build_states, - make_active_func, make_targets, + next_model_age, + prob_di_medicare_next, select_ss_benefit, select_target_for_age, ) def _make_transition_canwork( - gets_medicare: bool, own: dict[str, int], ng: dict[str, int], ) -> Callable[..., FloatND]: """Create transition for canwork tied regimes. - Tied agents who stop working become nongroup (lose employer coverage). - Medicaid-eligible agents are also overridden to nongroup targets. + Tied agents who stop working become nongroup (lose employer coverage); + before 65, those disabled next period hold disability Medicare there. + Categorically (SSI-) Medicaid-eligible agents are also overridden to + nongroup targets. """ def transition( age: Age, period: Period, + health: DiscreteState, labor_supply: DiscreteAction, - is_medicaid_eligible: BoolND, + is_ssi_eligible: BoolND, survival_probs: FloatND, + prob_disabled_next: FloatND, ) -> FloatND: - sp = survival_probs[period] - next_age = age + 1 - mc_next = gets_medicare & (labor_supply == LaborSupply.do_not_work) - target = select_target_for_age(next_age, mc_next, own) - # Tied agents who stop working become nongroup - stopped = labor_supply == LaborSupply.do_not_work - ng_target = select_target_for_age(next_age, mc_next, ng) - target = jnp.where(stopped, ng_target, target) - # Medicaid eligibility overrides to nongroup - ng_ssi = select_target_for_age(next_age, mc_next, ng) - target = jnp.where(is_medicaid_eligible, ng_ssi, target) - return build_regime_probs(target, sp) + del age # Keep the legacy signature; period owns the clock lookup. + next_age = next_model_age(period) + 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) + return jnp.where(to_nongroup, ng_target, own_target) + + return build_regime_probs_with_di_medicare( + target_dimc=target(True), + target_nomc=target(False), + prob_dimc=prob_di_medicare_next( + labor_supply, prob_disabled_next[period, health] + ), + survival=survival_probs[period, health], + ) return transition -def _build_functions(spec: RegimeSpec) -> dict: +def _build_functions(spec: RegimeSpec, *, subtract_bill: bool = False) -> dict: """Build functions dict for a tied regime.""" - functions = build_common_functions(spec) + functions = build_common_functions(spec, subtract_bill=subtract_bill) functions["ss_benefit"] = select_ss_benefit(spec) @@ -81,6 +101,16 @@ def _build_functions(spec: RegimeSpec) -> dict: return functions +def build_law(name: str) -> ByAge: + """Build the tied regime's per-target transition probabilities by source age.""" + spec = REGIME_SPECS[name] + own, ng = make_targets(name) + transition_func = _make_transition_canwork(own, ng) + return build_regime_transition( + spec=spec, transition_func=transition_func, target_groups=(own, ng) + ) + + def build_regime( name: str, grids: Grids, @@ -90,27 +120,29 @@ def build_regime( ) -> Regime: """Build a tied regime (all tied regimes are canwork).""" spec = REGIME_SPECS[name] - gets_mc = spec["mc"] != "nomc" - own, ng = make_targets(name) - - transition_func = _make_transition_canwork(gets_mc, own, ng) - states = build_states(spec, grids) egm_solver = dcegm_solver if dcegm_solver is not None else nbegm_solver - solver_kwargs: dict = {} if egm_solver is None else {"solver": egm_solver} state_solver = ( "brute_force" if egm_solver is None else ("nbegm" if nbegm_solver is not None else "dcegm") ) - return Regime( - transition=build_granular_regime_transition( - transition_func=transition_func, target_ids=(*own.values(), *ng.values()) - ), - active=make_active_func(spec), + functions = _build_functions( + spec, subtract_bill=grids.grid_config.subtract_hcc_bill + ) + if nbegm_solver is not None: + # NBEGM's solver contract is stated per regime: it reads the budget in + # savings form off `resources` and the post-decision node off + # `savings`, neither of which the brute-force build needs. + functions = {**functions, **build_nbegm_functions()} + return build_alive_regime( + egm_solver=egm_solver, states=states, - state_transitions=build_state_transitions(spec, solver=state_solver), + state_transitions=build_state_transitions( + spec, + solver=state_solver, + subtract_bill=grids.grid_config.subtract_hcc_bill, + ), actions=build_actions(spec, grids), - functions=_build_functions(spec), - **solver_kwargs, + functions=functions, ) diff --git a/src/aca_model/benchmark.py b/src/aca_model/benchmark.py index 52b8691..3eed460 100644 --- a/src/aca_model/benchmark.py +++ b/src/aca_model/benchmark.py @@ -29,14 +29,15 @@ import jax.numpy as jnp import numpy as np from jax import Array -from lcm import DiscreteGrid, Model +from lcm import DiscreteGrid, ExecutionConfig, 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.config import BENCHMARK_GRID_CONFIG +from aca_model.baseline.regimes import ENTRY_REGIMES +from aca_model.config import BENCHMARK_GRID_CONFIG, GridConfig from aca_model.consumption_dollars_grid import ( compute_consumption_dollars_points, inject_consumption_dollars_points, @@ -57,44 +58,37 @@ "pref_type": DiscreteGrid(BenchmarkPrefType), } -# Five regimes active at start_age=51 (inelig + canwork). All have +# The five entry regimes (inelig + canwork). All have # HealthWithDisability (3-state); four have lagged_labor_supply (tied # does not — it's implied by the regime). -_INITIAL_REGIMES = ( - "retiree_nomc_inelig_canwork", - "tied_nomc_inelig_canwork", - "nongroup_nomc_inelig_canwork", - "retiree_dimc_inelig_canwork", - "nongroup_dimc_inelig_canwork", -) def create_benchmark_model( *, - n_subjects: int, pref_type_grid: DiscreteGrid, + grid_config: GridConfig = BENCHMARK_GRID_CONFIG, + execution_config: ExecutionConfig | None = None, ) -> Model: """Create the aca baseline with `BENCHMARK_GRID_CONFIG` and frozen fixed_params. - The benchmark uses a 2-type `BenchmarkPrefType`. No `batch_size != 0` - on any grid (continuous grids inherit - `BENCHMARK_GRID_CONFIG.n_assets_batch_size = 0` and - `n_aime_batch_size = 0`). + The benchmark uses a 2-type `BenchmarkPrefType`. Grids describe economic + outcomes; the execution policy selects devices and program widths. Args: - n_subjects: Forwarded to `lcm.Model(n_subjects=...)`. When set, the - first matching `simulate(...)` call AOT-compiles all simulate - functions for that batch shape. pref_type_grid: Pref-type grid; pass `DiscreteGrid(BenchmarkPrefType)`. + grid_config: Economic resolution and medical-shock folding declarations. + 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. """ fixed_params, wage_params, _ = get_benchmark_params(model=None) return create_model( - grid_config=BENCHMARK_GRID_CONFIG, + grid_config=grid_config, fixed_params=fixed_params, wage_params=wage_params, derived_categoricals=_DERIVED_CATEGORICALS, pref_type_grid=pref_type_grid, - n_subjects=n_subjects, + execution_config=execution_config, ) @@ -113,6 +107,10 @@ def get_benchmark_params( an `IrregSpacedGrid` with runtime-supplied points. Pass `model=None` to skip injection (e.g. when constructing the model with `fixed_params`). + + 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. """ with _PARAMS_FILE.open("rb") as fh: data = cloudpickle.load(fh) @@ -155,23 +153,27 @@ def get_benchmark_consumption_dollars_points(*, n_points: int) -> tuple[float, . def get_benchmark_initial_conditions( - *, model: Model, n_subjects: int, seed: int + *, + model: Model, + n_subjects: int, + seed: int, + grid_config: GridConfig = BENCHMARK_GRID_CONFIG, ) -> dict[str, Array]: """Draw random feasible initial conditions across five age-51 regimes. - Every subject gets a random regime from `_INITIAL_REGIMES`; continuous + Every subject gets a random regime from `ENTRY_REGIMES`; continuous states are drawn uniformly over the regime's grid range, discrete states uniformly over categories. States absent from a subject's regime are filled with 0 (pylcm ignores them for that regime). """ rng = np.random.default_rng(seed) - regime_ids = tuple(model.regime_names_to_ids[n] for n in _INITIAL_REGIMES) + regime_ids = tuple(model.regime_names_to_ids[n] for n in ENTRY_REGIMES) regime = rng.choice(regime_ids, size=n_subjects).astype(np.int32) # Grid ranges come from any of the five regimes (shared structure). # Use to_jax() so the helper handles both LinSpacedGrid and # PiecewiseLinSpacedGrid (the latter has no `.start` / `.stop`). - ref_regime = model.user_regimes[_INITIAL_REGIMES[0]] + ref_regime = model.user_regimes[ENTRY_REGIMES[0]] grids = ref_regime.states # Every state read here is a plain `Grid`; `pension_wealth` (a # carried state) is never indexed, so the union widening @@ -181,7 +183,13 @@ def get_benchmark_initial_conditions( assets_lo, assets_hi = float(assets_pts.min()), float(assets_pts.max()) aime_lo, aime_hi = float(aime_pts.min()), float(aime_pts.max()) hcc_p_pts = np.asarray(grids["hcc_persistent"].to_jax()) # ty: ignore[unresolved-attribute] - hcc_t_pts = np.asarray(grids["hcc_transitory"].to_jax()) # ty: ignore[unresolved-attribute] + hcc_t_grid = grids.get("hcc_transitory") + # Preserve the random stream for shared states when the folded state is absent. + hcc_t_pts = ( + np.asarray(hcc_t_grid.to_jax()) # ty: ignore[unresolved-attribute] + if hcc_t_grid is not None + else np.arange(grid_config.n_hcc_transitory_gridpoints) + ) wage_res_pts = np.asarray(grids["log_ft_wage_res"].to_jax()) # ty: ignore[unresolved-attribute] # For lagged_labor_supply: 0 for tied rows (regime lacks the state), @@ -191,7 +199,7 @@ def get_benchmark_initial_conditions( np.int32 ) - return { + initial = { "regime_id": jnp.asarray(regime), "age": jnp.full(n_subjects, 51.0), "assets": jnp.asarray(rng.uniform(assets_lo, assets_hi, n_subjects)), @@ -212,3 +220,6 @@ def get_benchmark_initial_conditions( # everyone starts as not-yet-claimed. "claimed_ss": jnp.zeros(n_subjects, dtype=jnp.int32), } + if "hcc_transitory" not in model.state_names(regime_name=ENTRY_REGIMES[0]): + initial.pop("hcc_transitory") + return initial diff --git a/src/aca_model/config.py b/src/aca_model/config.py index 47390c2..d95569b 100644 --- a/src/aca_model/config.py +++ b/src/aca_model/config.py @@ -1,4 +1,4 @@ -"""Configuration for the aca_model package.""" +"""Economic grids and numerical solver choices for the aca_model package.""" from dataclasses import dataclass from pathlib import Path @@ -12,169 +12,67 @@ @dataclass(frozen=True) class ModelConfig: start_age: int = 51 + last_start_age: int = 60 end_age: int = 96 + """First age at which everybody is dead; the model's last period.""" ss_early_age: int = 62 ss_forced_age: int = 70 work_forced_out_age: int = 72 medicare_age: int = 65 + # Pension benefits must be drawn from this age on; pension wealth no longer + # accrues from earnings. + pension_must_receive_age: int = 71 @dataclass(frozen=True) class GridConfig: + """Economic resolution and numerical choices, independent of device execution.""" + n_assets_gridpoints: int = 24 n_aime_gridpoints: int = 12 n_consumption_dollars_gridpoints: int = 70 n_wage_res_gridpoints: int = 5 n_hcc_persistent_gridpoints: int = 3 n_hcc_transitory_gridpoints: int = 5 - # `batch_size` on the assets / AIME grids: chunked vmap stride for the - # outer state loop. `1` shrinks the per-period Q intermediate by that - # axis's cardinality on hosts where the unsplayed kernel doesn't fit; - # `0` lets a single kernel span the axis. - n_assets_batch_size: int = 0 - n_aime_batch_size: int = 0 - # Sharding flags for discrete state grids. pylcm distributes the - # grid across available devices when the flag is `True`. Sharding - # is only supported on discrete state grids; continuous axes - # (`assets`, `aime`, `wage_res`, `hcc_*`) compile to an all-gather - # of the full V-array per device and are rejected at grid - # construction. Mutually exclusive with `batch_size>0` on the same - # axis (pylcm rejects the combination). `spousal_income_distributed` - # routes through `baseline/regimes/_common.py:build_states` to its - # inline-built `DiscreteGrid(...)` call. - pref_type_distributed: bool = False - spousal_income_distributed: bool = False - # `batch_size` on the inline-constructed discrete state grids — - # health, spousal_income, lagged_labor_supply, claimed_ss. These - # are read in `build_states` via `grids.grid_config`. Setting any - # of them to `1` puts that axis in a Python-level outer loop within - # the discrete-states block of the productmap - # (`_ordered_state_action_names`), shrinking the per-call Q - # intermediate by that axis's cardinality at the cost of one extra - # lax.scan layer. Defaults to `0`; production overrides set to `1` - # to compound the splay across the unsharded discretes. - n_health_batch_size: int = 0 - n_spousal_income_batch_size: int = 0 - n_lagged_labor_supply_batch_size: int = 0 - n_claimed_ss_batch_size: int = 0 - # `batch_size` on the `pref_type` discrete grid: chunked vmap stride - # for the pref-type axis during solve. `1` (one pref-type per Python - # dispatch) shrinks the per-period Q intermediate by `n_pref_types` - # at the cost of an outer Python loop; `0` lets a single kernel span - # all pref-types. Defaults to `0` — the production overrides set it - # to `1` on hardware where the unsplayed kernel doesn't fit. - n_pref_type_batch_size: int = 0 - # `batch_size` on the `wage_res` stochastic shock process: chunked - # productmap stride along the wage-residual stoch axis inside Q_and_F. - # `1` shrinks the per-target Q intermediate by `n_wage_res_gridpoints` - # at the cost of an inner Python loop; `0` lets the productmap span - # the full axis. Defaults to `0` — production overrides set it to `1` - # on hardware where the ACA-overlay per-cell DAG blows the kernel's - # compile-time working set past device HBM. - n_wage_res_batch_size: int = 0 - # Number of nodes on the DC-EGM savings grid (the post-decision endogenous - # grid), cubically clustered toward the borrowing constraint. Drives the - # padded-grid dimension the egm_step kernel carries, so it scales both the - # rolling carry and the gather mesh roughly linearly — the dominant lever on - # DC-EGM device memory. Only consulted under `solver="dcegm"`. + # Fold the transitory medical-cost shock out of the state space. Nothing a + # period computes reads last year's transitory shock, so with the fold it is + # drawn on the transition only and no value carries its axis. + fold_hcc_transitory: bool = False + # Declare the asset law toward living targets as resources minus the + # out-of-pocket bill, so the solve averages each target's value over the + # transitory shock once per period instead of at every source point. Needs + # `fold_hcc_transitory` and the brute-force asset law. + subtract_hcc_bill: bool = False + # The post-decision savings grid is cubically clustered toward the borrowing + # constraint. Its node count controls resolution under DC-EGM and NB-EGM. n_savings_gridpoints: int = 200 - # `batch_size` on the DC-EGM savings grid (the post-decision endogenous - # grid). pylcm splays the per-savings-node continuation into `lax.map` - # blocks of this size, shrinking the binding `egm_step` working buffer by - # roughly the block factor while the upper envelope still runs on the full - # gathered grid (value function unchanged). `0` keeps the whole grid in one - # kernel. Only consulted under `solver="dcegm"`. - n_savings_batch_size: int = 0 - # `batch_size` on the DC-EGM child stochastic-node expectation (the - # process-state mesh the egm_step kernel sums over). pylcm splays that - # expectation into `lax.map` blocks of this size, collapsing the gather - # buffer's node axis by roughly the block factor (value function - # unchanged) — the grid-independent lever for the residual egm_step - # transient that savings-batching leaves behind. `0` evaluates the whole - # mesh in one kernel. Only consulted under `solver="dcegm"`. - n_stochastic_node_batch_size: int = 0 - # Block size for splaying the NBEGM continuation's child stochastic-node - # expectation (health, health-cost shocks, the wage residual). `0` reads the - # whole joint node mesh in one pass — fast, but its peak intermediate scales - # with the full ride-along × node × child-grid product. A positive value loops - # the mesh in blocks of that size, trading runtime for a much smaller peak; `1` - # (one node at a time) is the memory-minimal setting for a CPU validation grid. - # Only consulted under `solver="nbegm"`. - n_nbegm_stochastic_node_batch_size: int = 0 - # Streams the per-interval upper envelope over candidate-segment blocks of this - # size instead of materialising the full (query x candidate) bracket matrix per - # ride cell. `0` keeps the one-shot dense envelope; the result is identical - # either way — the knob trades peak device memory against a sequential scan. - # Only consulted under `solver="nbegm"`. - n_nbegm_envelope_segment_block_size: int = 0 - # Which arithmetic decides ownership in the merged upper envelope: - # - "certified" (default) — candidates are compared in double-double precision - # and no winner is published where none is separated, so the reported owner is - # one the arithmetic could prove. Ordering survives the cancellation a nearly - # tied crossing produces, at roughly an order of magnitude more arithmetic per - # read — and the envelope read is the dominant per-cell cost of a case-piece - # solve, so this is the setting that decides the solver's runtime. - # - "ordinary" — each candidate is read in the working format and the largest - # owns the query. Adequate wherever candidate values are separated by much - # more than the format's resolution at their own magnitude. - # Incompatible with a positive `n_nbegm_envelope_segment_block_size`, which - # selects a blocked scan that carries the certified arithmetic only. - # Only consulted under `solver="nbegm"`. + # Arithmetic used to compare upper-envelope candidates: + # - "certified": double-double comparisons publish only separated winners. + # - "ordinary": comparisons use the working floating-point format. nbegm_envelope_arithmetic: Literal["certified", "ordinary"] = "certified" - # Streams both NBEGM ride-along cores (continuation fan-out and envelope - # solve) over ride-cell blocks of this size instead of vmapping the whole - # flattened ride mesh at once — the dominant peak-memory term at production - # mesh sizes. `0` keeps the whole-mesh vmap; the result is identical either - # way. Only consulted under `solver="nbegm"`. - # - # Backend-dependent tuning at production ride-mesh sizes (under both cliff-read - # modes). On GPU the whole-mesh vmap stays within a few GiB, so `0` is fine. The - # CPU XLA backend does not fuse the fan-out and materialises the whole flattened - # ride mesh at once — a production-grid solve then needs hundreds of GiB even at - # a small asset grid (the blow-up rides the aime/shock/health mesh, not assets). - # Set this to a positive block (e.g. 64) for a CPU solve; it bounds the peak to - # the GPU's few-GiB footprint at the cost of serialising the mesh into a - # `lax.map` scan. - n_nbegm_cell_block_size: int = 0 - # How NBEGM parents read the child value's institutional cliffs: - # - "one_sided" (default) — carry rows hold each cliff preimage as a duplicated - # abscissa with exact one-sided limits; reads never average across a cliff, - # but publishing the topology gates the stochastic-dim fold off (slower). - # - "bridged" — plain carry rows; interpolation may bridge a cliff like any - # finite-grid solver, and the fold stays available. The fast setting for - # inner estimation loops, polished afterwards under "one_sided". - # Only consulted under `solver="nbegm"`. + # How NB-EGM parents read institutional cliffs in the child value: + # - "one_sided": duplicated abscissae preserve one-sided limits. + # - "bridged": interpolation may bridge finite-grid discontinuities. + # Live labor choices require "bridged" because branch-dependent income + # thresholds cannot share one query grid of one-sided cliff limits. nbegm_jump_read: Literal["one_sided", "bridged"] = "one_sided" - # Block size for streaming the discrete-action branch axis in both NBEGM - # ride-along cores (one continuation row / one continuous subproblem per live - # labor level). `0` runs the whole branch axis in one vectorized pass; a positive - # value scans it in blocks of that many branches (identical result, per-branch - # intermediates bounded by one block). Only consulted under `solver="nbegm"`. - n_nbegm_branch_batch_size: int = 0 - # Keep `labor_supply` a live discrete action on the M1 regime under NBEGM (the - # branch compiler solves each labor level against its own continuation, utility, - # and breakpoint partition); `buy_private` stays fixed. `False` fixes both actions - # to a single level so the only choice is continuous consumption. Only consulted - # under `solver="nbegm"`. - # - # Requires `nbegm_jump_read="bridged"`. `labor_supply` enters `countable_income`, - # which carries the SSI income test, so each labor level puts that breakpoint at a - # different liquid level; the one-sided read publishes its cliff limits on a single - # query grid shared across branches, which the two cannot both satisfy. Building a - # model with live labor under `"one_sided"` raises `RegimeInitializationError`. - nbegm_live_labor_supply: bool = False MODEL_CONFIG = ModelConfig() +# 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) +) GRID_CONFIG = GridConfig() BENCHMARK_GRID_CONFIG = GridConfig( + fold_hcc_transitory=True, + subtract_hcc_bill=True, n_assets_gridpoints=3, n_aime_gridpoints=3, n_consumption_dollars_gridpoints=5, n_wage_res_gridpoints=3, n_hcc_persistent_gridpoints=3, n_hcc_transitory_gridpoints=3, - n_assets_batch_size=0, - n_aime_batch_size=0, ) diff --git a/src/aca_model/consumption_dollars_grid.py b/src/aca_model/consumption_dollars_grid.py index 3cb6179..abb2dfb 100644 --- a/src/aca_model/consumption_dollars_grid.py +++ b/src/aca_model/consumption_dollars_grid.py @@ -68,7 +68,7 @@ def inject_consumption_dollars_points( max_consumption_dollars_arr = jnp.asarray(max_consumption_dollars) out: dict[str, Any] = dict(params) for regime_name, regime in model.user_regimes.items(): - if regime.terminal: + if model.graph.laws[regime_name].terminal: continue grid = regime.actions.get("consumption_dollars") if grid is None: @@ -131,7 +131,10 @@ def inject_consumption_floor_schedule( ) out: dict[str, Any] = dict(params) for regime_name, regime in model.user_regimes.items(): - if regime.terminal or "resources" not in regime.functions: + if ( + model.graph.laws[regime_name].terminal + or "resources" not in regime.functions + ): continue regime_entry = dict(out.get(regime_name, {})) resources_entry = dict(regime_entry.get("resources", {})) diff --git a/src/aca_model/environment/pensions.py b/src/aca_model/environment/pensions.py index 5318b9d..e9a55d8 100644 --- a/src/aca_model/environment/pensions.py +++ b/src/aca_model/environment/pensions.py @@ -4,7 +4,9 @@ """ import jax.numpy as jnp -from lcm.typing import FloatND, IntND, Period, ScalarFloat +from lcm.typing import Age, FloatND, IntND, Period, ScalarFloat + +from aca_model.config import MODEL_CONFIG def full_benefit( @@ -59,37 +61,49 @@ def benefit( def total_to_pia( - pension_benefit: FloatND, - pia: FloatND, + pia_adjusted_next_period: FloatND, + pia_unadjusted_next_period: FloatND, + full_benefit_next_period: FloatND, + target_his: IntND, period: Period, - his: IntND, marginal_tax_rate: FloatND, - imp_intercept: FloatND, - imp_pia_coeff: FloatND, - imp_pia_kink_0_coeff: FloatND, - imp_pia_kink_1_coeff: FloatND, - imp_kink_0: FloatND, - imp_kink_1: FloatND, + imp_intercept_next_period: FloatND, + imp_pia_coeff_next_period: FloatND, + imp_pia_kink_0_coeff_next_period: FloatND, + imp_pia_kink_1_coeff_next_period: FloatND, + imp_kink_0_next_period: FloatND, + imp_kink_1_next_period: FloatND, ) -> FloatND: - """Invert pension imputation to recover PIA from total after-tax benefits. + """PIA the next-period AIME encodes in solve (French & Jones 2011, app. D). + + Solve re-imputes pension wealth from next period's PIA, so a claim-age + adjustment or earnings-test credit baked into that PIA would also move the + imputed pension. The carried PIA `PIA*` instead solves + + ``` + PIA* + (1 - τ) pbmax_{t+1}(PIA*) = PIA_adj + (1 - τ) pbmax_{t+1}(PIA_unadj) + ``` - Piecewise-linear inverse of the pension_benefit mapping, adjusted for - marginal tax rates. + which keeps the after-tax total of Social Security and imputed pension + benefits at its claim-adjusted value. `pbmax_{t+1}` uses next period's + coefficients for the target HIS and `τ` this period's marginal tax rate. + Without an adjustment (`PIA_adj == PIA_unadj`) the carried PIA is the + accrued PIA. """ after_tax = 1.0 - marginal_tax_rate - total_ben = after_tax * pension_benefit + pia + total_ben = after_tax * full_benefit_next_period + pia_adjusted_next_period - at_intercept = after_tax * imp_intercept[period, his] - at_pia = after_tax * imp_pia_coeff[period, his] - at_kink_0 = after_tax * imp_pia_kink_0_coeff[period, his] - at_kink_1 = after_tax * imp_pia_kink_1_coeff[period, his] + at_intercept = after_tax * imp_intercept_next_period[period, target_his] + at_pia = after_tax * imp_pia_coeff_next_period[period, target_his] + at_kink_0 = after_tax * imp_pia_kink_0_coeff_next_period[period, target_his] + at_kink_1 = after_tax * imp_pia_kink_1_coeff_next_period[period, target_his] - k0 = imp_kink_0[period] - k1 = imp_kink_1[period] + k0 = imp_kink_0_next_period[period] + k1 = imp_kink_1_next_period[period] kink_0_tb = at_intercept + k0 * (1.0 + at_pia) kink_1_tb = kink_0_tb + (k1 - k0) * (1.0 + at_pia + at_kink_0) - return jnp.where( + inverted = jnp.where( total_ben < at_intercept, 0.0, jnp.where( @@ -102,10 +116,16 @@ def total_to_pia( ), ), ) + return jnp.where( + pia_adjusted_next_period == pia_unadjusted_next_period, + pia_unadjusted_next_period, + inverted, + ) def accrual( labor_income: FloatND, + age: Age, period: Period, his: IntND, accrual_intercept: FloatND, @@ -119,6 +139,9 @@ def accrual( Accrual has two components: - Accrual rate among holders (linear in log earnings) - Probability of accrual (logistic in log earnings) + + Accrual stops at `MODEL_CONFIG.pension_must_receive_age`, from which + pension benefits must be drawn. """ lli = jnp.log(jnp.maximum(1.0, labor_income)) @@ -134,7 +157,8 @@ def accrual( ) prob = jnp.exp(logit) / (1.0 + jnp.exp(logit)) - return jnp.where(labor_income > 0.0, rate * prob * labor_income, 0.0) + accrues = (labor_income > 0.0) & (age < MODEL_CONFIG.pension_must_receive_age) + return jnp.where(accrues, rate * prob * labor_income, 0.0) def wealth( @@ -164,11 +188,18 @@ def wealth_next_before_adjustment( next = ((1 + r) * wealth + accrual - benefit) / survival_prob The (1 + r) factor compounds pension wealth at the market return rate. - The division by survival prob accounts for the annuity pricing. + The division by survival prob accounts for the annuity pricing. At an age + where nobody survives, the wealth is left undivided. """ - return ( + survival_prob = unconditional_survival_prob[period] + compounded = ( (1.0 + rate_of_return) * pension_wealth + pension_accrual - pension_benefit - ) / unconditional_survival_prob[period] + ) + return jnp.where( + survival_prob > 0.0, + compounded / jnp.where(survival_prob > 0.0, survival_prob, 1.0), + compounded, + ) def assets_adjustment( @@ -191,7 +222,7 @@ def assets_adjustment( ) -def imputed_pension_wealth_next_period( +def full_benefit_next_period( pia_unadjusted_next_period: FloatND, target_his: IntND, period: Period, @@ -201,21 +232,17 @@ def imputed_pension_wealth_next_period( imp_pia_kink_1_coeff_next_period: FloatND, imp_kink_0_next_period: FloatND, imp_kink_1_next_period: FloatND, - epdv_constant_pension_next_period: FloatND, ) -> FloatND: - """Imputed pension wealth at next period using the target regime's HIS. - - `pw_{t+1} = Γ_{t+1} · pbmax_{t+1}` (French & Jones 2011, eq. D.3): the - *full* next-period benefit times the next-period annuity factor, with no - current-period fraction (the fraction lives inside `Γ`). Mirrors - `full_benefit` and `wealth` but indexes 1-period-shifted views so all - subscripts use bare-name parameters (`period`, `target_his`). Inlining is - required: pylcm's AST shape inference inspects the registered function's - body and does not trace through nested calls. - - The PIA input is `pia_unadjusted_next_period` — the next-period PIA from pure labor - accrual. French & Jones impute pension wealth from the unadjusted PIA, so the - claim-age reduction or credit baked into the carried AIME never enters here. + """Next period's full pension benefit `pbmax_{t+1}` at the target regime's HIS. + + Mirrors `full_benefit` but indexes 1-period-shifted views so all subscripts + use bare-name parameters (`period`, `target_his`). Inlining is required: + pylcm's AST shape inference inspects the registered function's body and + does not trace through nested calls. + + The PIA input is `pia_unadjusted_next_period` — the next-period PIA from pure + labor accrual. French & Jones impute pension wealth from the unadjusted PIA, + so the claim-age reduction or credit never enters here. """ intercept = imp_intercept_next_period[period, target_his] pia_pred = ( @@ -227,6 +254,18 @@ def imputed_pension_wealth_next_period( kink_1_adj = imp_pia_kink_1_coeff_next_period[period, target_his] * jnp.maximum( 0.0, pia_unadjusted_next_period - imp_kink_1_next_period[period] ) + return jnp.maximum(0.0, intercept + pia_pred + kink_0_adj + kink_1_adj) + + +def imputed_pension_wealth_next_period( + full_benefit_next_period: FloatND, + period: Period, + epdv_constant_pension_next_period: FloatND, +) -> FloatND: + """Imputed pension wealth at next period (French & Jones 2011, eq. D.3). - pbmax_next = jnp.maximum(0.0, intercept + pia_pred + kink_0_adj + kink_1_adj) - return pbmax_next * epdv_constant_pension_next_period[period] + `pw_{t+1} = Γ_{t+1} · pbmax_{t+1}`: the *full* next-period benefit times the + next-period annuity factor, with no current-period fraction (the fraction + lives inside `Γ`). + """ + return full_benefit_next_period * epdv_constant_pension_next_period[period] diff --git a/src/aca_model/environment/social_security.py b/src/aca_model/environment/social_security.py index 412e44b..129c168 100644 --- a/src/aca_model/environment/social_security.py +++ b/src/aca_model/environment/social_security.py @@ -202,18 +202,30 @@ def benefit_inelig_pre65( def benefit_withheld_fraction( pia: FloatND, ss_benefit: FloatND, + age: Age, + period: Period, claim_ss: DiscreteAction, claimed_ss: DiscreteState, + early_ret_adjustment: FloatND, + normal_retirement_age: ScalarInt, ) -> FloatND: - """Fraction of raw PIA withheld (early retirement + earnings test). + """Fraction of the claimed benefit the earnings test withholds. - Zero when not claiming regular SS. Used by `next_aime` to credit back - future AIME for benefit withholding during the earnings test. + The claimed benefit is the PIA times the early-claim factor below the + normal retirement age. Only earnings-test withholding is credited back to + future AIME, so the fraction is zero when not claiming regular SS, when not + working, and when earnings stay below the threshold; the early-claim + reduction itself is never credited. """ is_claiming = jnp.maximum(claim_ss, claimed_ss) > 0 + claimed_benefit = jnp.where( + age < normal_retirement_age, pia * early_ret_adjustment[period], pia + ) + has_benefit = claimed_benefit > 0 + safe_benefit = jnp.where(has_benefit, claimed_benefit, 1.0) return jnp.where( - is_claiming & (pia > 0), - jnp.maximum(0.0, 1.0 - ss_benefit / pia), + is_claiming & has_benefit, + jnp.maximum(0.0, 1.0 - ss_benefit / safe_benefit), 0.0, ) @@ -266,7 +278,7 @@ def _apply_benefit_rules( # --- AIME transition functions (state transitions, no aime_to_pia calls) --- -def next_aime( +def pia_adjusted_next_period( aime: ContinuousState, labor_income: FloatND, period: Period, @@ -285,16 +297,15 @@ def next_aime( aime_last_age_with_indexing: ScalarInt, aime_kink_2: ScalarFloat, ratio_lowest_earnings: FloatND, -) -> ContinuousState: - """Compute next period's AIME given labor earnings and the claim decision. +) -> FloatND: + """Next period's PIA after accrual, earnings-test credit and claim adjustment. Steps: 1. Accrue AIME from labor income (indexing, taxable cap, lowest-year drop). - 2. Credit back for earnings-test withholding (PIA round-trip). - 3. Bake the claim-age actuarial factor into AIME (`_apply_claim_adjustment`), + 2. Credit back for earnings-test withholding. + 3. Apply the per-year claim-age actuarial factor (`_apply_claim_adjustment`), so the carried AIME permanently encodes the early-retirement reduction or - the delayed-retirement credit. The flat-PIA benefit read off this AIME is - then correct at every later age, including the forced regimes. + the delayed-retirement credit. """ credited_pia = _accrue_and_credit_back_pia( aime=aime, @@ -312,8 +323,7 @@ def next_aime( aime_kink_2=aime_kink_2, ratio_lowest_earnings=ratio_lowest_earnings, ) - - adjusted_pia = _apply_claim_adjustment( + return _apply_claim_adjustment( pia=credited_pia, period=period, age=age, @@ -322,57 +332,41 @@ def next_aime( normal_retirement_age=normal_retirement_age, early_ret_adjustment=early_ret_adjustment, ) - # The extended `pia_table`/`pia_aime_grid` reach above the taxable max so a - # delayed-retirement credit on a top earner's PIA round-trips to an AIME - # above `aime_kink_2` rather than clamping there. `_accrue_aime` already - # capped the labor-earnings base at the taxable max; the actuarial credit is - # the only thing carried beyond it. - return jnp.interp(adjusted_pia, pia_table, pia_aime_grid) -def next_aime_plain( - aime: ContinuousState, - labor_income: FloatND, - period: Period, - age: Age, - benefit_withheld_fraction: FloatND, - earnings_test_credited_back: FloatND, - earnings_test_repealed_age: ScalarInt, +def carried_pia_simulate(pia_adjusted_next_period: FloatND) -> FloatND: + """PIA the next-period AIME encodes in simulate: the adjusted PIA itself. + + Simulate carries the true pension wealth as a state, so no pension + imputation has to be offset. The solve-phase counterpart is + `pensions.total_to_pia`. + """ + return pia_adjusted_next_period + + +def next_aime( + carried_pia: FloatND, pia_table: FloatND, pia_aime_grid: FloatND, - aime_accrual_factor: ScalarFloat, - aggregate_wage_growth: ScalarFloat, - aime_last_age_with_indexing: ScalarInt, - aime_kink_2: ScalarFloat, - ratio_lowest_earnings: FloatND, ) -> ContinuousState: - """Compute next period's AIME without any claim-age adjustment. - - Used by post-65 `ss=forced` regimes, where the agent cannot choose when to - claim and so carries no `claim_ss` action / `claimed_ss` state. A forced - claimant who claimed early already has the actuarial reduction baked into - the AIME carried in from the `ss=choose` regime; plain accrual preserves it. - A forced claimant who never claimed early keeps a pristine AIME, so the - flat-PIA benefit equals the full PIA, which is correct. + """AIME law of the post-65 `ss=choose` regimes: the AIME of the carried PIA. + + The extended `pia_table`/`pia_aime_grid` reach above the taxable max, so a + delayed-retirement credit on a top earner's PIA round-trips to an AIME + above `aime_kink_2` instead of clamping there. """ - credited_pia = _accrue_and_credit_back_pia( - aime=aime, - labor_income=labor_income, - period=period, - age=age, - benefit_withheld_fraction=benefit_withheld_fraction, - earnings_test_credited_back=earnings_test_credited_back, - earnings_test_repealed_age=earnings_test_repealed_age, - pia_table=pia_table, - pia_aime_grid=pia_aime_grid, - aime_accrual_factor=aime_accrual_factor, - aggregate_wage_growth=aggregate_wage_growth, - aime_last_age_with_indexing=aime_last_age_with_indexing, - aime_kink_2=aime_kink_2, - ratio_lowest_earnings=ratio_lowest_earnings, - ) - accrued_aime = jnp.interp(credited_pia, pia_table, pia_aime_grid) - return jnp.minimum(accrued_aime, aime_kink_2) + return jnp.interp(carried_pia, pia_table, pia_aime_grid) + + +def next_aime_forced(aime: ContinuousState) -> ContinuousState: + """AIME law of the post-65 `ss=forced` regimes: AIME stays as it is. + + Once claiming is forced, earnings no longer accrue to AIME and nothing is + capped. A forced claimant who claimed early carries the reduction in from + the `ss=choose` regime; one who delayed carries the delayed-retirement + credit, including the part above the taxable maximum. + """ + return aime def pia_unadjusted_next_period( @@ -408,6 +402,11 @@ def pia_unadjusted_next_period( return jnp.interp(accrued_aime, pia_aime_grid, pia_table) +def pia_unadjusted_next_period_forced(pia: FloatND) -> FloatND: + """Next-period PIA in the `ss=forced` regimes, where AIME stays as it is.""" + return pia + + def _accrue_aime( *, aime: ContinuousState, @@ -420,7 +419,10 @@ def _accrue_aime( aime_kink_2: ScalarFloat, ratio_lowest_earnings: FloatND, ) -> FloatND: - """Accrue AIME from labor income (indexing, taxable cap, lowest-year drop).""" + """Accrue AIME from labor income (indexing, taxable cap, lowest-year drop). + + The accrued AIME is capped at the taxable maximum. + """ indexed_aime = jnp.where( age <= aime_last_age_with_indexing, aime * (1.0 + aggregate_wage_growth), @@ -432,7 +434,7 @@ def _accrue_aime( accrual = ( jnp.maximum(0.0, capped_labor - lowest_year_earnings) * aime_accrual_factor ) - return capped_aime + accrual + return jnp.minimum(capped_aime + accrual, aime_kink_2) def _accrue_and_credit_back_pia( @@ -521,17 +523,11 @@ def _apply_claim_adjustment( def next_aime_disabled( aime: ContinuousState, + carried_pia: FloatND, labor_income: FloatND, period: Period, age: Age, health: DiscreteState, - claim_ss: DiscreteAction, - claimed_ss: DiscreteState, - normal_retirement_age: ScalarInt, - early_ret_adjustment: FloatND, - benefit_withheld_fraction: FloatND, - earnings_test_credited_back: FloatND, - earnings_test_repealed_age: ScalarInt, pia_table: FloatND, pia_aime_grid: FloatND, aime_accrual_factor: ScalarFloat, @@ -543,50 +539,33 @@ def next_aime_disabled( di_dropout_scale: FloatND, di_dropout_next_period_ratio: FloatND, ) -> ContinuousState: - """AIME transition for pre-65 regimes handling both disabled and non-disabled. + """AIME transition for pre-65 `ss=choose` regimes, disabled and non-disabled. - Non-disabled: standard AIME accrual from labor income, earnings-test - credit-back, and the claim-age actuarial bake (early reduction for claimants - below NRA — pre-65 never reaches the delayed-credit window). + Non-disabled: the AIME of the carried PIA (accrual, earnings-test credit and + claim-age adjustment; pre-65 never reaches the delayed-credit window). Disabled: maintain PIA continuity across ages by adjusting AIME so the DI dropout-year scale factor change doesn't alter the benefit. The disabled path reads the un-baked AIME — DI benefits are routed only when never claimed, so the claim adjustment never touches them. At Medicare transition, stores the dropout-adjusted AIME (switching to OA). """ - credited_pia = _accrue_and_credit_back_pia( + regular = jnp.minimum( + jnp.interp(carried_pia, pia_table, pia_aime_grid), + aime_kink_2, + ) + accrued_aime = _accrue_aime( aime=aime, labor_income=labor_income, period=period, age=age, - benefit_withheld_fraction=benefit_withheld_fraction, - earnings_test_credited_back=earnings_test_credited_back, - earnings_test_repealed_age=earnings_test_repealed_age, - pia_table=pia_table, - pia_aime_grid=pia_aime_grid, aime_accrual_factor=aime_accrual_factor, aggregate_wage_growth=aggregate_wage_growth, aime_last_age_with_indexing=aime_last_age_with_indexing, aime_kink_2=aime_kink_2, ratio_lowest_earnings=ratio_lowest_earnings, ) - - adjusted_pia = _apply_claim_adjustment( - pia=credited_pia, - period=period, - age=age, - claim_ss=claim_ss, - claimed_ss=claimed_ss, - normal_retirement_age=normal_retirement_age, - early_ret_adjustment=early_ret_adjustment, - ) - regular = jnp.minimum( - jnp.interp(adjusted_pia, pia_table, pia_aime_grid), - aime_kink_2, - ) - return _select_disabled_or_regular( - aime=aime, + accrued_aime=accrued_aime, regular=regular, period=period, age=age, @@ -647,8 +626,19 @@ def next_aime_disabled_plain( aime_kink_2, ) - return _select_disabled_or_regular( + accrued_aime = _accrue_aime( aime=aime, + labor_income=labor_income, + period=period, + age=age, + aime_accrual_factor=aime_accrual_factor, + aggregate_wage_growth=aggregate_wage_growth, + aime_last_age_with_indexing=aime_last_age_with_indexing, + aime_kink_2=aime_kink_2, + ratio_lowest_earnings=ratio_lowest_earnings, + ) + return _select_disabled_or_regular( + accrued_aime=accrued_aime, regular=regular, period=period, age=age, @@ -661,7 +651,7 @@ def next_aime_disabled_plain( def _select_disabled_or_regular( *, - aime: ContinuousState, + accrued_aime: FloatND, regular: FloatND, period: Period, age: Age, @@ -672,14 +662,16 @@ def _select_disabled_or_regular( ) -> ContinuousState: """Route the disabled DI-continuity AIME against the non-disabled AIME. - The disabled path reads the un-baked `aime`, scaling it so the DI dropout-year - factor change leaves the benefit unchanged. At the Medicare transition it - switches to the dropout-adjusted AIME (OA from then on). + The disabled path starts from the un-baked AIME after indexing and labor + accrual, the same accrual as for anyone else, and scales it so the DI + dropout-year factor change alone leaves the benefit unchanged. At the + Medicare transition it switches to the dropout-adjusted AIME (OA from then + on). """ disabled_next = jnp.where( age + 1 < medicare_age, - aime * di_dropout_next_period_ratio[period], - aime * di_dropout_scale[period], + accrued_aime * di_dropout_next_period_ratio[period], + accrued_aime * di_dropout_scale[period], ) is_disabled = health == 0 return jnp.where(is_disabled, disabled_next, regular) diff --git a/src/aca_model/environment/taxes.py b/src/aca_model/environment/taxes.py index dced975..c61d014 100644 --- a/src/aca_model/environment/taxes.py +++ b/src/aca_model/environment/taxes.py @@ -25,8 +25,7 @@ def gross_income( capital_income: FloatND, labor_income: FloatND, - spousal_income: DiscreteState, - spousal_income_amounts: FloatND, + spousal_income_amount: FloatND, taxable_ss_benefit: FloatND, pension_benefit: FloatND, ) -> FloatND: @@ -34,7 +33,7 @@ def gross_income( return ( capital_income + labor_income - + spousal_income_amounts[spousal_income] + + spousal_income_amount + taxable_ss_benefit + pension_benefit ) @@ -44,7 +43,7 @@ def taxable_ss_benefit( capital_income: FloatND, labor_income: FloatND, spousal_income: DiscreteState, - spousal_income_amounts: FloatND, + spousal_income_amount: FloatND, ss_benefit: FloatND, pension_benefit: FloatND, ss_tax_schedule: MappingLeaf, @@ -58,7 +57,7 @@ def taxable_ss_benefit( prov_income = ( capital_income + labor_income - + spousal_income_amounts[spousal_income] + + spousal_income_amount + sched["ben_fraction_prov_income"] * ss_benefit + pension_benefit ) @@ -179,3 +178,43 @@ def _payroll_tax( taxes_at_lower[bracket_id] + (labor_income - bracket_lower) * marginal_rates[bracket_id] ) + + +# ACA surtaxes (the 2013 net investment income tax and additional Medicare tax). +ACA_INVESTMENT_SURTAX_RATE = 0.038 +ACA_INVESTMENT_SURTAX_THRESHOLD = 200_000.0 +ACA_PAYROLL_SURTAX_RATE = 0.009 +ACA_PAYROLL_SURTAX_THRESHOLD = 200_000.0 + + +def after_tax_income_with_aca_surtaxes( + after_tax_income_before_aca_surtaxes: FloatND, + gross_income: FloatND, + labor_income: FloatND, +) -> FloatND: + """Subtract both ACA surtaxes from after-tax income. + + - 3.8% on unearned income (`gross_income - labor_income`: asset income, + spousal income, taxable SS benefits and pension benefits) above 200,000; + - 0.9% on earnings above 200,000. + """ + unearned_income = gross_income - labor_income + return ( + after_tax_income_before_aca_surtaxes + - ACA_INVESTMENT_SURTAX_RATE + * jnp.maximum(0.0, unearned_income - ACA_INVESTMENT_SURTAX_THRESHOLD) + - ACA_PAYROLL_SURTAX_RATE + * jnp.maximum(0.0, labor_income - ACA_PAYROLL_SURTAX_THRESHOLD) + ) + + +def after_tax_income_with_aca_investment_surtax( + after_tax_income_before_aca_surtaxes: FloatND, + gross_income: FloatND, + labor_income: FloatND, +) -> FloatND: + """Subtract the 3.8% ACA surtax on unearned income above 200,000.""" + unearned_income = gross_income - labor_income + return after_tax_income_before_aca_surtaxes - ACA_INVESTMENT_SURTAX_RATE * ( + jnp.maximum(0.0, unearned_income - ACA_INVESTMENT_SURTAX_THRESHOLD) + ) diff --git a/src/aca_model/execution.py b/src/aca_model/execution.py new file mode 100644 index 0000000..b940674 --- /dev/null +++ b/src/aca_model/execution.py @@ -0,0 +1,53 @@ +"""Hardware-local execution policies derived from allocator observations.""" + +import jax +from lcm import ExecutionConfig + + +def execution_config_for_devices( + *, devices: tuple[int, ...] | None = None +) -> ExecutionConfig: + """Create an execution policy bounded by the selected accelerator pools. + + The smallest allocator limit is a total per-device ceiling. Pylcm accounts + for its resident arrays within that ceiling, so current allocator usage is + not subtracted here. CPU devices do not supply an accelerator-pool budget. + + Args: + devices: Visible device IDs to select, in order; None selects all devices + on JAX's default backend. + + Returns: + ExecutionConfig with explicit device IDs and the measured accelerator + budget. Axis widths and sharded states remain unspecified. + + Raises: + ValueError: A selected device is not visible or an accelerator has no + positive integer allocator limit. Missing limits require an explicit + policy. + """ + policy = ExecutionConfig(devices=devices) + visible = {device.id: device for device in jax.devices()} + selected = tuple(visible) if policy.devices is None else policy.devices + missing = tuple(device_id for device_id in selected if device_id not in visible) + if missing: + msg = f"Selected devices are not visible: {missing}." + raise ValueError(msg) + limits = [] + for device_id in selected: + device = visible[device_id] + if device.platform == "cpu": + continue + stats = device.memory_stats() or {} + limit = stats.get("bytes_limit") + if type(limit) is not int or limit <= 0: + msg = ( + "No positive integer allocator byte limit reported for " + f"device {device_id}." + ) + raise ValueError(msg) + limits.append(limit) + return ExecutionConfig( + devices=selected, + device_memory_bytes=min(limits) if limits else None, + ) diff --git a/src/aca_model/simulation.py b/src/aca_model/simulation.py index 6760bab..743394a 100644 --- a/src/aca_model/simulation.py +++ b/src/aca_model/simulation.py @@ -40,12 +40,50 @@ def simulate_with_dense_index( on the result panel. """ - original_ids = np.asarray(initial_conditions.index) - dense = initial_conditions.reset_index(drop=True) + admissible = select_admissible_starts( + model=model, initial_conditions=initial_conditions + ) + original_ids = np.asarray(admissible.index) + dense = admissible.reset_index(drop=True) result = model.simulate(initial_conditions=dense, **simulate_kwargs) return result, original_ids +def select_admissible_starts( + *, model: Any, initial_conditions: pd.DataFrame +) -> pd.DataFrame: + """Keep the rows that start at or before the model's last admissible start age. + + Rows older than every age in `model.graph.initial_nodes` are outside the sample + the model is meant to simulate and are dropped. Every kept row must start at + an admissible `(age, regime_name)` pair; any other kept row raises. + + Args: + model: A pylcm `Model` (anything exposing `graph.initial_nodes`, the + admissible `(age, regime_name)` pairs). + initial_conditions: Seed DataFrame with `age` and `regime_name` columns. + + Returns: + The rows of `initial_conditions` inside the model's entry ages. + + Raises: + ValueError: If a kept row starts at a pair outside `model.graph.initial_nodes`. + + """ + admitted = {(float(age), regime) for age, regime in model.graph.initial_nodes} + last_entry_age = max(age for age, _ in admitted) + kept = initial_conditions.loc[initial_conditions["age"] <= last_entry_age] + pairs = set(zip(kept["age"].astype(float), kept["regime_name"], strict=True)) + _fail_if_starts_not_admitted(refused=pairs - admitted) + return kept + + +def _fail_if_starts_not_admitted(*, refused: set[tuple[float, str]]) -> None: + if refused: + msg = f"Initial conditions start at inadmissible pairs: {sorted(refused)}" + raise ValueError(msg) + + def restore_subject_ids( panel: pd.DataFrame, original_ids: NDArray[Any], diff --git a/tests/helpers/model.py b/tests/helpers/model.py index 55571e9..3645ab7 100644 --- a/tests/helpers/model.py +++ b/tests/helpers/model.py @@ -5,6 +5,8 @@ Production callers (aca-slurm, scripts) assemble these explicitly. """ +from dataclasses import replace + from lcm import DiscreteGrid, Model from aca_model.aca.health_insurance import PolicyVariant @@ -14,8 +16,15 @@ from aca_model.agent.preferences import BenchmarkPrefType from aca_model.baseline.health_insurance import HealthInsuranceState from aca_model.baseline.model import create_model as _create_baseline_model +from aca_model.baseline.regimes import REGIME_SPECS from aca_model.benchmark import get_benchmark_params -from aca_model.config import BENCHMARK_GRID_CONFIG +from aca_model.config import BENCHMARK_GRID_CONFIG, GridConfig + +_NONGROUP_BEFORE_MEDICARE = frozenset( + name + for name, spec in REGIME_SPECS.items() + if spec["his"] == "nongroup" and spec["mc"] == "nomc" +) _DERIVED_CATEGORICALS = { "good_health": DiscreteGrid(GoodHealth), @@ -25,29 +34,53 @@ "pref_type": DiscreteGrid(BenchmarkPrefType), } +_UNFOLDED_GRID_CONFIG = replace( + BENCHMARK_GRID_CONFIG, fold_hcc_transitory=False, subtract_hcc_bill=False +) -def make_baseline_model(*, n_subjects: int) -> Model: - """Baseline model on `BENCHMARK_GRID_CONFIG` with the benchmark snapshot params.""" + +def make_baseline_model(*, grid_config: GridConfig = _UNFOLDED_GRID_CONFIG) -> Model: + """Baseline model on `grid_config` with the benchmark snapshot params.""" fixed_params, wage_params, _ = get_benchmark_params(model=None) return _create_baseline_model( - n_subjects=n_subjects, fixed_params=fixed_params, wage_params=wage_params, derived_categoricals=_DERIVED_CATEGORICALS, - grid_config=BENCHMARK_GRID_CONFIG, + grid_config=grid_config, pref_type_grid=DiscreteGrid(BenchmarkPrefType), ) -def make_aca_model(*, n_subjects: int, policy: PolicyVariant) -> Model: - """ACA model on `BENCHMARK_GRID_CONFIG` with the benchmark snapshot params.""" +def aca_fixed_params(*, fixed_params: dict, policy: PolicyVariant) -> dict: + """Return baseline `fixed_params` fit for the ACA `policy` model. + + Under a variant with the reformed non-group market, the non-group regimes + before Medicare price cover at the community-rated ACA premium, so the + risk-rated premium parameters (prefix `premium_`) are dropped from those + regimes, as aca-slurm does. + """ + if policy == PolicyVariant.ACA_ONLY_MEDICAID_EXPANSION: + return fixed_params + return { + name: ( + {k: v for k, v in value.items() if not k.startswith("premium_")} + if name in _NONGROUP_BEFORE_MEDICARE + else value + ) + for name, value in fixed_params.items() + } + + +def make_aca_model( + *, policy: PolicyVariant, grid_config: GridConfig = _UNFOLDED_GRID_CONFIG +) -> Model: + """ACA model on `grid_config` with the benchmark snapshot params.""" fixed_params, wage_params, _ = get_benchmark_params(model=None) return _create_aca_model( - n_subjects=n_subjects, policy=policy, - fixed_params=fixed_params, + fixed_params=aca_fixed_params(fixed_params=fixed_params, policy=policy), wage_params=wage_params, derived_categoricals=_DERIVED_CATEGORICALS, - grid_config=BENCHMARK_GRID_CONFIG, + grid_config=grid_config, pref_type_grid=DiscreteGrid(BenchmarkPrefType), ) diff --git a/tests/helpers/production_rules.py b/tests/helpers/production_rules.py new file mode 100644 index 0000000..83ac3a0 --- /dev/null +++ b/tests/helpers/production_rules.py @@ -0,0 +1,731 @@ +"""Social Security and pension inputs at their production values (2014 rules). + +Built from the raw aca-data inputs (`social_security_rules.json`, `AIMEfrac_012207.txt`, +the pension coefficient CSVs). Every per-age tuple is period-indexed: entry `i` holds +age `51 + i`, ages 51-95. +""" + +START_AGE = 51 + +PIA_AIME_GRID = (0.0, 9792.0, 59004.0, 117000.0, 187954.75) +PIA_TABLE = (0.0, 8812.7998046875, 24560.640625, 33260.0390625, 43903.25390625) +NORMAL_RETIREMENT_AGE = 66 +EARNINGS_TEST_REPEALED_AGE = 67 +AIME_KINK_2 = 117000.0 +AIME_ACCRUAL_FACTOR = 0.02857142857142857 +AGGREGATE_WAGE_GROWTH = 0.008 +SSDI_SUBSTANTIAL_GAINFUL_ACTIVITY = 12840.0 +AIME_LAST_AGE_WITH_INDEXING = 60 +MEDICARE_AGE = 65 +EARLY_RET_ADJUSTMENT = ( + 0.75, + 0.75, + 0.75, + 0.75, + 0.75, + 0.75, + 0.75, + 0.75, + 0.75, + 0.75, + 0.75, + 0.75, + 0.8, + 0.866666666667, + 0.933333333333, + 1.0, + 1.08, + 1.16, + 1.24, + 1.32, + 1.32, + 1.32, + 1.32, + 1.32, + 1.32, + 1.32, + 1.32, + 1.32, + 1.32, + 1.32, + 1.32, + 1.32, + 1.32, + 1.32, + 1.32, + 1.32, + 1.32, + 1.32, + 1.32, + 1.32, + 1.32, + 1.32, + 1.32, + 1.32, + 1.32, +) +EARNINGS_TEST_THRESHOLD = ( + 15480.0, + 15480.0, + 15480.0, + 15480.0, + 15480.0, + 15480.0, + 15480.0, + 15480.0, + 15480.0, + 15480.0, + 15480.0, + 15480.0, + 15480.0, + 15480.0, + 15480.0, + 41400.0, + 1000000000000.0, + 1000000000000.0, + 1000000000000.0, + 1000000000000.0, + 1000000000000.0, + 1000000000000.0, + 1000000000000.0, + 1000000000000.0, + 1000000000000.0, + 1000000000000.0, + 1000000000000.0, + 1000000000000.0, + 1000000000000.0, + 1000000000000.0, + 1000000000000.0, + 1000000000000.0, + 1000000000000.0, + 1000000000000.0, + 1000000000000.0, + 1000000000000.0, + 1000000000000.0, + 1000000000000.0, + 1000000000000.0, + 1000000000000.0, + 1000000000000.0, + 1000000000000.0, + 1000000000000.0, + 1000000000000.0, + 1000000000000.0, +) +EARNINGS_TEST_FRACTION = ( + 0.5, + 0.5, + 0.5, + 0.5, + 0.5, + 0.5, + 0.5, + 0.5, + 0.5, + 0.5, + 0.5, + 0.5, + 0.5, + 0.5, + 0.5, + 0.33333333333333, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, +) +EARNINGS_TEST_CREDITED_BACK = ( + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.06666666666666665, + 0.08333333333374981, + 0.07692307692227796, + 0.0714285714289542, + 0.08000000000000007, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, +) +RATIO_LOWEST_EARNINGS = ( + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.10672, + 0.21343, + 0.32015, + 0.42687, + 0.53358, + 0.56986, + 0.58915, + 0.60016, + 0.60782, + 0.61424, + 0.61571, + 0.61694, + 0.61769, + 0.61821, + 0.61873, + 0.61925, + 0.61925, + 0.61925, + 0.61925, + 0.61925, + 0.61925, + 0.61925, + 0.61925, + 0.61925, + 0.61925, + 0.61925, + 0.61925, + 0.61925, + 0.61925, + 0.61925, + 0.61925, + 0.61925, + 0.61925, + 0.61925, + 0.61925, + 0.61925, + 0.61925, + 0.61925, + 0.61925, + 0.61925, +) +DI_DROPOUT_SCALE = ( + 1.2083333333333333, + 1.2, + 1.1923076923076923, + 1.1851851851851851, + 1.1785714285714286, + 1.1724137931034482, + 1.1666666666666667, + 1.1631093333333333, + 1.1560166926186666, + 1.1454424427860417, + 1.1314722994845539, + 1.1142228426400118, + 1.099433715124906, + 1.087559039482935, + 1.0787779344134025, + 1.0731317905082343, + 1.0706203950623945, + 1.0691946690817968, + 1.0683311937577702, + 1.0677867771027343, + 1.0674641949164902, + 1.0673248083147346, + 1.0672149229623997, + 1.0671432229312787, + 1.0670948043279977, + 1.0670625262355393, + 1.0670463874334053, + 1.0670463874334053, + 1.0670463874334053, + 1.0670463874334053, + 1.0670463874334053, + 1.0670463874334053, + 1.0670463874334053, + 1.0670463874334053, + 1.0670463874334053, + 1.0670463874334053, + 1.0670463874334053, + 1.0670463874334053, + 1.0670463874334053, + 1.0670463874334053, + 1.0670463874334053, + 1.0670463874334053, + 1.0670463874334053, + 1.0670463874334053, + 1.0670463874334053, +) +DI_DROPOUT_NEXT_PERIOD_RATIO = ( + 1.0069444444444444, + 1.0064516129032257, + 1.0060096153846154, + 1.0056116722783388, + 1.0052521008403363, + 1.0049261083743841, + 1.0030584685647208, + 1.0061354137530663, + 1.0092315854884033, + 1.0123468716890833, + 1.0154811552809953, + 1.0134515863136193, + 1.0109186492051196, + 1.0081398634412255, + 1.0052613704627036, + 1.0023457384684824, + 1.0013334578087842, + 1.0008082468517927, + 1.000509855213335, + 1.0003021948537292, + 1.000130594361407, + 1.0001029645950132, + 1.0000671887611525, + 1.0000453742283109, + 1.0000302494855406, + 1.0000151247427704, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, +) +IMP_INTERCEPT = ( + (-12910.289654931415, -11444.233693705028, -13263.514860994796), + (-12109.012655612423, -11277.312062238974, -12910.546833118284), + (-11279.981383814986, -11120.553833266282, -12536.714764203674), + (-10423.195839539105, -10973.959006786954, -12142.01865425097), + (-9538.656022784777, -10837.527582800987, -11726.458503260168), + (-8626.361933552005, -10711.25956130839, -11290.034311231271), + (-7686.313571840789, -10595.154942309153, -10832.74607816428), + (-6718.510937651128, -10489.21372580328, -10354.593804059192), + (-5722.954030983024, -10393.43591179077, -9855.577488916007), + (-4699.642851836475, -10307.821500271624, -9335.697132734727), + (-3648.57740021148, -10232.370491245842, -8794.952735515351), + (-2569.757676108042, -10167.082884713423, -8233.34429725788), + (-1463.1836795261593, -10111.958680674372, -7650.871817962313), + (-328.85541046583194, -10066.99787912868, -7047.5352976286495), + (833.2271310729416, -10032.200480076352, -6423.334736256889), + (2023.0639450901576, -10007.566483517388, -5778.270133847034), + (3240.6550315858185, -9993.09588945179, -5112.341490399083), + (4486.000390559924, -9988.788697879554, -4425.548805913037), + (5759.100022012473, -9994.64490880068, -3717.892080388894), + (7059.95392594347, -10010.664522215173, -2989.371313826653), + (8388.56210235291, -10036.847538123026, -2239.986506226318), + (9744.92455124079, -10073.193956524248, -1469.7376575878916), + (9744.92455124079, -10073.193956524248, -1469.7376575878916), + (9744.92455124079, -10073.193956524248, -1469.7376575878916), + (9744.92455124079, -10073.193956524248, -1469.7376575878916), + (9744.92455124079, -10073.193956524248, -1469.7376575878916), + (9744.92455124079, -10073.193956524248, -1469.7376575878916), + (9744.92455124079, -10073.193956524248, -1469.7376575878916), + (9744.92455124079, -10073.193956524248, -1469.7376575878916), + (9744.92455124079, -10073.193956524248, -1469.7376575878916), + (9744.92455124079, -10073.193956524248, -1469.7376575878916), + (9744.92455124079, -10073.193956524248, -1469.7376575878916), + (9744.92455124079, -10073.193956524248, -1469.7376575878916), + (9744.92455124079, -10073.193956524248, -1469.7376575878916), + (9744.92455124079, -10073.193956524248, -1469.7376575878916), + (9744.92455124079, -10073.193956524248, -1469.7376575878916), + (9744.92455124079, -10073.193956524248, -1469.7376575878916), + (9744.92455124079, -10073.193956524248, -1469.7376575878916), + (9744.92455124079, -10073.193956524248, -1469.7376575878916), + (9744.92455124079, -10073.193956524248, -1469.7376575878916), + (9744.92455124079, -10073.193956524248, -1469.7376575878916), + (9744.92455124079, -10073.193956524248, -1469.7376575878916), + (9744.92455124079, -10073.193956524248, -1469.7376575878916), + (9744.92455124079, -10073.193956524248, -1469.7376575878916), + (9744.92455124079, -10073.193956524248, -1469.7376575878916), +) +IMP_PIA_COEFF = ( + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), + (1.877047452602557, 1.877047452602557, 1.877047452602557), +) +IMP_PIA_KINK_0_COEFF = ( + (-2.6823413559119853, -2.6823413559119853, -2.6823413559119853), + (-2.5062022601604563, -2.5062022601604563, -2.5062022601604563), + (-2.321110595527509, -2.321110595527509, -2.321110595527509), + (-2.127066362013144, -2.127066362013144, -2.127066362013144), + (-1.92406955961736, -1.92406955961736, -1.92406955961736), + (-1.7121201883401571, -1.7121201883401571, -1.7121201883401571), + (-1.4912182481815368, -1.4912182481815368, -1.4912182481815368), + (-1.261363739141498, -1.261363739141498, -1.261363739141498), + (-1.0225566612200407, -1.0225566612200407, -1.0225566612200407), + (-0.7747970144171654, -0.7747970144171654, -0.7747970144171654), + (-0.5180847987328717, -0.5180847987328717, -0.5180847987328717), + (-0.2524200141671596, -0.2524200141671596, -0.2524200141671596), + (0.0221973392799706, 0.0221973392799706, 0.0221973392799706), + (0.3057672616085193, 0.3057672616085193, 0.3057672616085193), + (0.5982897528184863, 0.5982897528184863, 0.5982897528184863), + (0.8997648129098714, 0.8997648129098714, 0.8997648129098714), + (1.210192441882675, 1.210192441882675, 1.210192441882675), + (1.529572639736897, 1.529572639736897, 1.529572639736897), + (1.857905406472537, 1.857905406472537, 1.857905406472537), + (2.1951907420895957, 2.1951907420895957, 2.1951907420895957), + (2.541428646588072, 2.541428646588072, 2.541428646588072), + (2.8966191199679674, 2.8966191199679674, 2.8966191199679674), + (2.8966191199679674, 2.8966191199679674, 2.8966191199679674), + (2.8966191199679674, 2.8966191199679674, 2.8966191199679674), + (2.8966191199679674, 2.8966191199679674, 2.8966191199679674), + (2.8966191199679674, 2.8966191199679674, 2.8966191199679674), + (2.8966191199679674, 2.8966191199679674, 2.8966191199679674), + (2.8966191199679674, 2.8966191199679674, 2.8966191199679674), + (2.8966191199679674, 2.8966191199679674, 2.8966191199679674), + (2.8966191199679674, 2.8966191199679674, 2.8966191199679674), + (2.8966191199679674, 2.8966191199679674, 2.8966191199679674), + (2.8966191199679674, 2.8966191199679674, 2.8966191199679674), + (2.8966191199679674, 2.8966191199679674, 2.8966191199679674), + (2.8966191199679674, 2.8966191199679674, 2.8966191199679674), + (2.8966191199679674, 2.8966191199679674, 2.8966191199679674), + (2.8966191199679674, 2.8966191199679674, 2.8966191199679674), + (2.8966191199679674, 2.8966191199679674, 2.8966191199679674), + (2.8966191199679674, 2.8966191199679674, 2.8966191199679674), + (2.8966191199679674, 2.8966191199679674, 2.8966191199679674), + (2.8966191199679674, 2.8966191199679674, 2.8966191199679674), + (2.8966191199679674, 2.8966191199679674, 2.8966191199679674), + (2.8966191199679674, 2.8966191199679674, 2.8966191199679674), + (2.8966191199679674, 2.8966191199679674, 2.8966191199679674), + (2.8966191199679674, 2.8966191199679674, 2.8966191199679674), + (2.8966191199679674, 2.8966191199679674, 2.8966191199679674), +) +IMP_PIA_KINK_1_COEFF = ( + (-3.1582574731568576, -3.1582574731568576, -3.1582574731568576), + (-2.3263985293592504, -2.3263985293592504, -2.3263985293592504), + (-1.586173944257751, -1.586173944257751, -1.586173944257751), + (-0.9375837178523592, -0.9375837178523592, -0.9375837178523592), + (-0.3806278501430754, -0.3806278501430754, -0.3806278501430754), + (0.084693658870101, 0.084693658870101, 0.084693658870101), + (0.4583808091871695, 0.4583808091871695, 0.4583808091871695), + (0.7404336008081298, 0.7404336008081298, 0.7404336008081298), + (0.9308520337329828, 0.9308520337329828, 0.9308520337329828), + (1.029636107961728, 1.029636107961728, 1.029636107961728), + (1.036785823494366, 1.036785823494366, 1.036785823494366), + (0.9523011803308954, 0.9523011803308954, 0.9523011803308954), + (0.7761821784713181, 0.7761821784713181, 0.7761821784713181), + (0.5084288179156298, 0.5084288179156298, 0.5084288179156298), + (0.1490410986638366, 0.1490410986638366, 0.1490410986638366), + (-0.3019809792840622, -0.3019809792840622, -0.3019809792840622), + (-0.8446374159280712, -0.8446374159280712, -0.8446374159280712), + (-1.4789282112681867, -1.4789282112681867, -1.4789282112681867), + (-2.2048533653044142, -2.2048533653044142, -2.2048533653044142), + (-3.0224128780367465, -3.0224128780367465, -3.0224128780367465), + (-3.9316067494651854, -3.9316067494651854, -3.9316067494651854), + (-4.932434979589736, -4.932434979589736, -4.932434979589736), + (-4.932434979589736, -4.932434979589736, -4.932434979589736), + (-4.932434979589736, -4.932434979589736, -4.932434979589736), + (-4.932434979589736, -4.932434979589736, -4.932434979589736), + (-4.932434979589736, -4.932434979589736, -4.932434979589736), + (-4.932434979589736, -4.932434979589736, -4.932434979589736), + (-4.932434979589736, -4.932434979589736, -4.932434979589736), + (-4.932434979589736, -4.932434979589736, -4.932434979589736), + (-4.932434979589736, -4.932434979589736, -4.932434979589736), + (-4.932434979589736, -4.932434979589736, -4.932434979589736), + (-4.932434979589736, -4.932434979589736, -4.932434979589736), + (-4.932434979589736, -4.932434979589736, -4.932434979589736), + (-4.932434979589736, -4.932434979589736, -4.932434979589736), + (-4.932434979589736, -4.932434979589736, -4.932434979589736), + (-4.932434979589736, -4.932434979589736, -4.932434979589736), + (-4.932434979589736, -4.932434979589736, -4.932434979589736), + (-4.932434979589736, -4.932434979589736, -4.932434979589736), + (-4.932434979589736, -4.932434979589736, -4.932434979589736), + (-4.932434979589736, -4.932434979589736, -4.932434979589736), + (-4.932434979589736, -4.932434979589736, -4.932434979589736), + (-4.932434979589736, -4.932434979589736, -4.932434979589736), + (-4.932434979589736, -4.932434979589736, -4.932434979589736), + (-4.932434979589736, -4.932434979589736, -4.932434979589736), + (-4.932434979589736, -4.932434979589736, -4.932434979589736), +) +IMP_KINK_0 = ( + 8575.974479744078, + 8731.448831139582, + 8886.923182535085, + 9042.397533930589, + 9256.587112476413, + 9402.182994638755, + 9319.930704810571, + 9524.355960657978, + 9617.03283358649, + 9799.654678462908, + 10423.8659145107, + 10691.70033560747, + 11394.924408430852, + 12025.115878719977, + 12173.105198328823, + 12722.358154802065, + 12565.048054151552, + 12565.048054151552, + 12565.048054151552, + 12565.048054151552, + 12565.048054151552, + 12565.048054151552, + 12565.048054151552, + 12565.048054151552, + 12565.048054151552, + 12565.048054151552, + 12565.048054151552, + 12565.048054151552, + 12565.048054151552, + 12565.048054151552, + 12565.048054151552, + 12565.048054151552, + 12565.048054151552, + 12565.048054151552, + 12565.048054151552, + 12565.048054151552, + 12565.048054151552, + 12565.048054151552, + 12565.048054151552, + 12565.048054151552, + 12565.048054151552, + 12565.048054151552, + 12565.048054151552, + 12565.048054151552, + 12565.048054151552, +) +IMP_KINK_1 = ( + 11572.45542008772, + 11713.599948182897, + 11854.744476278069, + 11995.889004373244, + 12238.228113539872, + 12363.52639198924, + 12385.693839410424, + 12563.353470856087, + 12714.252790399893, + 12996.920909081768, + 13858.934848147588, + 14514.609937658775, + 15516.08194652413, + 16305.915470536642, + 16454.00202972665, + 17087.67191139636, + 17329.483780991442, + 17329.483780991442, + 17329.483780991442, + 17329.483780991442, + 17329.483780991442, + 17329.483780991442, + 17329.483780991442, + 17329.483780991442, + 17329.483780991442, + 17329.483780991442, + 17329.483780991442, + 17329.483780991442, + 17329.483780991442, + 17329.483780991442, + 17329.483780991442, + 17329.483780991442, + 17329.483780991442, + 17329.483780991442, + 17329.483780991442, + 17329.483780991442, + 17329.483780991442, + 17329.483780991442, + 17329.483780991442, + 17329.483780991442, + 17329.483780991442, + 17329.483780991442, + 17329.483780991442, + 17329.483780991442, + 17329.483780991442, +) +ACCRUAL_INTERCEPT = ( + (-0.2810633901901035, -0.3583711901901036, -0.3345061901901036), + (-0.2810633901901035, -0.3583711901901036, -0.3345061901901036), + (-0.1727824073816185, -0.2373314073816186, -0.2020886073816186), + (-0.0725590702150753, -0.1238664702150753, -0.0813164702150753), + (0.0203416095257619, -0.017375690474238, 0.0289008095257619), + (0.1065212124433541, 0.0826084124433541, 0.1295206124433542), + (0.1864475575493455, 0.1764196575493454, 0.2213666575493454), + (0.2604555946507882, 0.2642586946507883, 0.3051294946507882), + (0.328746373554735, 0.346192673554735, 0.3813660735547349), + (0.3913877132728298, 0.4221552132728298, 0.4504999132728298), + (0.4483137944075337, 0.4919463944075336, 0.5128209944075337), + (0.3206231352723865, 0.3696765352723865, 0.3405695352723865), + (0.3653881562752405, 0.4259920562752404, 0.3896019562752404), + (0.4034406373204525, 0.4747365373204524, 0.4318915373204525), + (0.074266228649094, 0.118775128649094, 0.144751328649094), + (0.096955948404771, 0.150037648404771, 0.172685148404771), + (0.1107455160331468, 0.1711392160331468, 0.1927560160331468), + (0.1146131121366821, 0.1809239121366821, 0.2042977121366821), + (0.1074476557270211, 0.1781465557270211, 0.2065551557270211), + (0.0879604966112162, 0.1613843966112162, 0.1985952966112162), + (0.0879604966112162, 0.1613843966112162, 0.1985952966112162), + (0.0879604966112162, 0.1613843966112162, 0.1985952966112162), + (0.0879604966112162, 0.1613843966112162, 0.1985952966112162), + (0.0879604966112162, 0.1613843966112162, 0.1985952966112162), + (0.0879604966112162, 0.1613843966112162, 0.1985952966112162), + (0.0879604966112162, 0.1613843966112162, 0.1985952966112162), + (0.0879604966112162, 0.1613843966112162, 0.1985952966112162), + (0.0879604966112162, 0.1613843966112162, 0.1985952966112162), + (0.0879604966112162, 0.1613843966112162, 0.1985952966112162), + (0.0879604966112162, 0.1613843966112162, 0.1985952966112162), + (0.0879604966112162, 0.1613843966112162, 0.1985952966112162), + (0.0879604966112162, 0.1613843966112162, 0.1985952966112162), + (0.0879604966112162, 0.1613843966112162, 0.1985952966112162), + (0.0879604966112162, 0.1613843966112162, 0.1985952966112162), + (0.0879604966112162, 0.1613843966112162, 0.1985952966112162), + (0.0879604966112162, 0.1613843966112162, 0.1985952966112162), + (0.0879604966112162, 0.1613843966112162, 0.1985952966112162), + (0.0879604966112162, 0.1613843966112162, 0.1985952966112162), + (0.0879604966112162, 0.1613843966112162, 0.1985952966112162), + (0.0879604966112162, 0.1613843966112162, 0.1985952966112162), + (0.0879604966112162, 0.1613843966112162, 0.1985952966112162), + (0.0879604966112162, 0.1613843966112162, 0.1985952966112162), + (0.0879604966112162, 0.1613843966112162, 0.1985952966112162), + (0.0879604966112162, 0.1613843966112162, 0.1985952966112162), + (0.0879604966112162, 0.1613843966112162, 0.1985952966112162), +) +ACCRUAL_LOG_EARNINGS = ( + (0.0438052, 0.0438052, 0.0438052), + (0.0438052, 0.0438052, 0.0438052), + (0.0339093, 0.0339093, 0.0339093), + (0.0240996, 0.0240996, 0.0240996), + (0.0144505, 0.0144505, 0.0144505), + (0.0050361, 0.0050361, 0.0050361), + (-0.0040693, -0.0040693, -0.0040693), + (-0.0127915, -0.0127915, -0.0127915), + (-0.0210563, -0.0210563, -0.0210563), + (-0.0287894, -0.0287894, -0.0287894), + (-0.0359167, -0.0359167, -0.0359167), + (-0.0271691, -0.0271691, -0.0271691), + (-0.0328619, -0.0328619, -0.0328619), + (-0.0377262, -0.0377262, -0.0377262), + (-0.009307, -0.009307, -0.009307), + (-0.0122916, -0.0122916, -0.0122916), + (-0.0142249, -0.0142249, -0.0142249), + (-0.0150328, -0.0150328, -0.0150328), + (-0.014641, -0.014641, -0.014641), + (-0.0129753, -0.0129753, -0.0129753), + (-0.0129753, -0.0129753, -0.0129753), + (-0.0129753, -0.0129753, -0.0129753), + (-0.0129753, -0.0129753, -0.0129753), + (-0.0129753, -0.0129753, -0.0129753), + (-0.0129753, -0.0129753, -0.0129753), + (-0.0129753, -0.0129753, -0.0129753), + (-0.0129753, -0.0129753, -0.0129753), + (-0.0129753, -0.0129753, -0.0129753), + (-0.0129753, -0.0129753, -0.0129753), + (-0.0129753, -0.0129753, -0.0129753), + (-0.0129753, -0.0129753, -0.0129753), + (-0.0129753, -0.0129753, -0.0129753), + (-0.0129753, -0.0129753, -0.0129753), + (-0.0129753, -0.0129753, -0.0129753), + (-0.0129753, -0.0129753, -0.0129753), + (-0.0129753, -0.0129753, -0.0129753), + (-0.0129753, -0.0129753, -0.0129753), + (-0.0129753, -0.0129753, -0.0129753), + (-0.0129753, -0.0129753, -0.0129753), + (-0.0129753, -0.0129753, -0.0129753), + (-0.0129753, -0.0129753, -0.0129753), + (-0.0129753, -0.0129753, -0.0129753), + (-0.0129753, -0.0129753, -0.0129753), + (-0.0129753, -0.0129753, -0.0129753), + (-0.0129753, -0.0129753, -0.0129753), +) +ACCRUAL_PROB_INTERCEPT = (-21.141658072405324, -24.804149517460925, -17.228156917864453) +ACCRUAL_PROB_LOG_EARNINGS = (2.9740129249502, 3.6243613249502, 2.8730428249502) +ACCRUAL_PROB_LOG_EARNINGS_SQ = ( + -0.1123136999999999, + -0.1123136999999999, + -0.1123136999999999, +) diff --git a/tests/helpers/social_security.py b/tests/helpers/social_security.py index 0c4ff2d..bbb911b 100644 --- a/tests/helpers/social_security.py +++ b/tests/helpers/social_security.py @@ -4,10 +4,14 @@ runtime dependency on the data-prep package. """ +from typing import Any + import numpy as np import pandas as pd from numpy.typing import NDArray +from aca_model.environment import social_security + def compute_pia_table( aime_kink_0: float, @@ -71,3 +75,17 @@ def compute_di_dropout_scale( result[period] = product * elapsed / (elapsed - years_to_drop) return result + + +def next_aime_from_state(**kwargs: Any) -> Any: + """Advance AIME one period through `pia_adjusted_next_period` and `next_aime`. + + Composes the two DAG steps the way simulate wires them (the carried PIA is + the claim-adjusted PIA). + """ + carried_pia = social_security.pia_adjusted_next_period(**kwargs) + return social_security.next_aime( + carried_pia=carried_pia, + pia_table=kwargs["pia_table"], + pia_aime_grid=kwargs["pia_aime_grid"], + ) diff --git a/tests/test_aca_employer_coverage.py b/tests/test_aca_employer_coverage.py new file mode 100644 index 0000000..53f6f6f --- /dev/null +++ b/tests/test_aca_employer_coverage.py @@ -0,0 +1,112 @@ +"""The ACA Medicaid expansion leaves employer-provided coverage untouched. + +A retiree or tied household whose only Medicaid route is the ACA MAGI +expansion keeps its employer coverage next period. Only categorical (SSI) +eligibility moves an employer-covered household to non-group. +""" + +from collections.abc import Callable + +import jax.numpy as jnp +import pytest +from dags import concatenate_functions +from lcm.params import MappingLeaf + +from aca_model.aca import health_insurance as aca_hi +from aca_model.agent.labor_market import LaborSupply +from aca_model.baseline import health_insurance +from aca_model.baseline.health_insurance import HealthInsuranceState +from aca_model.baseline.regimes._common import RegimeId, make_targets +from aca_model.baseline.regimes._retiree import ( + _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 + +N_PERIODS = MODEL_CONFIG.end_age - MODEL_CONFIG.start_age +SURVIVAL = jnp.ones((N_PERIODS, 3)) * 0.99 +MEDICAID_SCHEDULE = MappingLeaf( + {"income_threshold": jnp.array([16104.6, 21707.4, 21707.4])} +) + +# Under 65, not categorically eligible, MAGI below 138% FPL: eligible only +# through the expansion. +_EXPANSION_ONLY_INPUTS = { + "is_ssi_eligible": jnp.array(False), + "aca_magi": jnp.array(10000.0), + "spousal_income": jnp.int32(0), + "crossed_oamc_threshold": jnp.asarray(False), + "medicaid_schedule": MEDICAID_SCHEDULE, +} + + +def _live_target(probs: jnp.ndarray) -> int: + return int(jnp.argmax(probs.at[RegimeId.dead].set(0.0))) + + +def _compose_with_aca_eligibility( + target_name: str, target_func: Callable, inputs: dict +) -> dict: + """Evaluate `target_func` in a DAG that also carries ACA Medicaid eligibility.""" + combined = concatenate_functions( + {"is_medicaid_eligible": aca_hi.is_medicaid_eligible, target_name: target_func}, + targets=[target_name, "is_medicaid_eligible"], + return_type="dict", + ) + return combined(**inputs) + + +@pytest.mark.parametrize( + ("regime_name", "make_transition", "extra_inputs", "expected"), + [ + ( + "retiree_nomc_inelig_canwork", + retiree_canwork, + {"labor_supply": jnp.array(LaborSupply.h2000)}, + RegimeId.retiree_nomc_inelig_canwork, + ), + ( + "tied_nomc_inelig_canwork", + tied_canwork, + {"labor_supply": jnp.array(LaborSupply.h2000)}, + RegimeId.tied_nomc_inelig_canwork, + ), + ], +) +def test_expansion_only_medicaid_keeps_employer_coverage( + regime_name: str, + make_transition: Callable, + extra_inputs: dict, + expected: int, +) -> 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) + result = _compose_with_aca_eligibility( + "next_regime", + transition, + { + **_EXPANSION_ONLY_INPUTS, + **extra_inputs, + "age": jnp.int32(55), + "period": jnp.int32(55 - MODEL_CONFIG.start_age), + "health": jnp.int32(2), + "survival_probs": SURVIVAL, + "prob_disabled_next": jnp.zeros((N_PERIODS, 3)), + }, + ) + assert _live_target(result["next_regime"]) == expected + + +def test_expansion_only_medicaid_keeps_retiree_his_for_pension_imputation() -> None: + """The pension imputation's target HIS stays retiree for an expansion-only case.""" + result = _compose_with_aca_eligibility( + "target_his", + health_insurance.target_his, + { + **_EXPANSION_ONLY_INPUTS, + "his": jnp.int32(HealthInsuranceState.retiree), + "labor_supply": jnp.array(LaborSupply.h2000), + }, + ) + assert int(result["target_his"]) == HealthInsuranceState.retiree diff --git a/tests/test_aca_plan.py b/tests/test_aca_plan.py new file mode 100644 index 0000000..c8b0e3c --- /dev/null +++ b/tests/test_aca_plan.py @@ -0,0 +1,132 @@ +"""The ACA private non-group plan is community-rated. + +Under every ACA variant with the reformed non-group market, the premium of +a private non-group policy before premium credits no longer depends on the +household's expected medical costs: it is a flat amount per household +size. The tests evaluate the real non-group regime DAG after the ACA +overlay, with the baseline risk-rated premium pinned to a constant, so they +check which premium the overlay wires in. +""" + +import inspect + +import jax.numpy as jnp +import pytest +from dags import concatenate_functions + +from aca_model.aca.health_insurance import PolicyVariant +from aca_model.aca.regimes._overrides import apply_aca_overrides +from aca_model.agent.labor_market import LaborSupply +from aca_model.baseline import health_insurance +from aca_model.baseline.health_insurance import BuyPrivate +from aca_model.baseline.regimes._common import REGIME_SPECS, build_model_functions +from aca_model.baseline.regimes._nongroup import _build_functions + +RISK_RATED_PREMIUM = 4620.0 +ACA_PREMIUM_SINGLE = 6000.0 +ACA_PREMIUM_MARRIED_EXTRA = 3000.0 +REGIME = "nongroup_nomc_inelig_canwork" + + +def _is_base_premium(func: object) -> bool: + """Whether `func` is one of the baseline premium formulas.""" + return getattr(func, "__module__", "") == health_insurance.__name__ and getattr( + func, "__name__", "" + ).startswith("premium") + + +def _aca_premium(policy: PolicyVariant, spousal_income: int) -> float: + """Evaluate `hic_premium` of a non-Medicaid household buying private cover.""" + spec = REGIME_SPECS[REGIME] + functions = {**build_model_functions(), **_build_functions(spec)} + functions = { + name: (lambda: jnp.asarray(RISK_RATED_PREMIUM)) + if _is_base_premium(func) + else func + for name, func in functions.items() + if callable(func) + } + apply_aca_overrides(functions, spec, policy) + functions["is_medicaid_eligible"] = lambda: jnp.asarray(False) + combined = concatenate_functions( + functions, targets=["hic_premium"], return_type="tuple" + ) + inputs = { + "buy_private": jnp.asarray(BuyPrivate.yes), + "labor_supply": jnp.asarray(LaborSupply.do_not_work), + "spousal_income": jnp.int32(spousal_income), + "aca_premium_single": jnp.asarray(ACA_PREMIUM_SINGLE), + "aca_premium_married_extra": jnp.asarray(ACA_PREMIUM_MARRIED_EXTRA), + } + needed = inspect.signature(combined).parameters + return float(combined(**{k: v for k, v in inputs.items() if k in needed})[0]) + + +@pytest.mark.parametrize( + ("policy", "spousal_income", "expected"), + [ + (PolicyVariant.ACA, 0, ACA_PREMIUM_SINGLE), + (PolicyVariant.ACA, 2, ACA_PREMIUM_SINGLE + ACA_PREMIUM_MARRIED_EXTRA), + (PolicyVariant.ACA_NO_MANDATE, 1, 9000.0), + (PolicyVariant.ACA_NO_MEDICAID_EXPANSION, 0, 6000.0), + (PolicyVariant.ACA_NO_MEDICAID_EXPANSION_NO_MANDATE, 2, 9000.0), + (PolicyVariant.ACA_ONLY_MEDICAID_EXPANSION, 0, RISK_RATED_PREMIUM), + ], +) +def test_aca_non_group_premium_is_community_rated( + policy: PolicyVariant, spousal_income: int, expected: float +) -> None: + """The ACA plan costs 6000 for a single and 9000 for a married household. + + The variant without the non-group reform keeps the risk-rated premium. + """ + assert _aca_premium(policy, spousal_income) == pytest.approx(expected) + + +RISK_RATED_PREMIUM_PARAMS = frozenset( + { + "premium_markup", + "premium_minimum", + "premium_predicted_hcc", + "premium_sample_insurer_cost", + "premium_sample_is_married", + } +) + + +def _free_arguments(policy: PolicyVariant, regime: str) -> frozenset[str]: + """Arguments of a non-group regime's DAG that no function of it produces.""" + spec = REGIME_SPECS[regime] + functions = {**build_model_functions(), **_build_functions(spec)} + functions = {name: func for name, func in functions.items() if callable(func)} + apply_aca_overrides(functions, spec, policy) + arguments = { + name + for func in functions.values() + for name in inspect.signature(func).parameters + } + return frozenset(arguments - functions.keys()) + + +@pytest.mark.parametrize( + "regime", ["nongroup_nomc_inelig_canwork", "nongroup_nomc_choose_canwork"] +) +@pytest.mark.parametrize( + ("policy", "expected"), + [ + (PolicyVariant.ACA, frozenset()), + (PolicyVariant.ACA_NO_MANDATE, frozenset()), + (PolicyVariant.ACA_NO_MEDICAID_EXPANSION, frozenset()), + (PolicyVariant.ACA_NO_MEDICAID_EXPANSION_NO_MANDATE, frozenset()), + (PolicyVariant.ACA_ONLY_MEDICAID_EXPANSION, RISK_RATED_PREMIUM_PARAMS), + ], +) +def test_aca_non_group_regime_reads_risk_rated_premium_params_only_without_reform( + policy: PolicyVariant, regime: str, expected: frozenset[str] +) -> None: + """A community-rated non-group market reads none of the risk-rated inputs. + + The variant without the non-group reform still prices private cover by + risk and reads all of them. + """ + assert _free_arguments(policy, regime) & RISK_RATED_PREMIUM_PARAMS == expected diff --git a/tests/test_aca_policies.py b/tests/test_aca_policies.py index 8392e90..b16fda4 100644 --- a/tests/test_aca_policies.py +++ b/tests/test_aca_policies.py @@ -3,6 +3,7 @@ import inspect import jax.numpy as jnp +import numpy as np import pytest from lcm.params import MappingLeaf @@ -254,6 +255,7 @@ def _aca_medicaid_kwargs(**overrides: object) -> dict: Defaults describe an under-65, non-disabled household that is not categorically eligible; tests override the expansion-relevant fields. + Only the arguments `is_medicaid_eligible` reads are passed on. """ base = { "is_ssi_eligible": jnp.array(False), @@ -264,7 +266,8 @@ def _aca_medicaid_kwargs(**overrides: object) -> dict: "medicaid_schedule": MEDICAID_SCHEDULE, } base.update(overrides) - return base + read = inspect.signature(aca_hi.is_medicaid_eligible).parameters + return {name: value for name, value in base.items() if name in read} def test_medicaid_eligible_aca_expansion_below_threshold() -> None: @@ -303,14 +306,18 @@ def test_medicaid_eligible_aca_expansion_excludes_aged() -> None: assert not result -def test_medicaid_eligible_aca_expansion_excludes_disabled() -> None: - """Disabled households are not reached by the expansion track.""" +def test_medicaid_eligible_aca_expansion_includes_disabled_under_65() -> None: + """A disabled under-65 household below the MAGI threshold gets expansion cover. + + The expansion has no asset test, so the household qualifies even though + it fails the categorical track (`is_ssi_eligible` is False). + """ result = aca_hi.is_medicaid_eligible( **_aca_medicaid_kwargs( is_disabled=jnp.asarray(True), aca_magi=jnp.array(10000.0) ) ) - assert not result + assert result def test_medicaid_eligible_aca_ignores_assets() -> None: @@ -366,3 +373,34 @@ def test_premium_subsidy_just_below_400_fpl() -> None: premium_credit_schedule=PREMIUM_CREDIT_SCHEDULE, ) assert result > 0.0 + + +@pytest.mark.parametrize( + ("income_fpl_frac", "buy_private", "is_medicaid_eligible"), + [ + (0.5, BuyPrivate.yes, False), + (1.2, BuyPrivate.yes, False), + (1.7, BuyPrivate.yes, False), + (2.3, BuyPrivate.yes, False), + (3.0, BuyPrivate.yes, False), + (1.2, BuyPrivate.no, False), + (1.2, BuyPrivate.yes, True), + ], +) +def test_cost_sharing_scale_of_its_bracket_equals_the_cost_sharing_scale( + income_fpl_frac: float, buy_private: int, is_medicaid_eligible: bool +) -> None: + """The scale read off the bracket is the scale the bracket lookup gives.""" + kwargs = { + "gross_income": jnp.array(income_fpl_frac * FPL_SINGLE), + "spousal_income": jnp.int32(0), + "buy_private": jnp.array(buy_private), + "is_medicaid_eligible": jnp.array(is_medicaid_eligible), + "cost_sharing_schedule": COST_SHARING_SCHEDULE, + } + bracket = aca_hi.cost_sharing_bracket(**kwargs) + scale = aca_hi.cost_sharing_scale_of_bracket( + cost_sharing_bracket=bracket, cost_sharing_schedule=COST_SHARING_SCHEDULE + ) + + np.testing.assert_allclose(scale, aca_hi.cost_sharing(**kwargs), rtol=0, atol=0) diff --git a/tests/test_aca_surtaxes_wiring.py b/tests/test_aca_surtaxes_wiring.py new file mode 100644 index 0000000..7378ad2 --- /dev/null +++ b/tests/test_aca_surtaxes_wiring.py @@ -0,0 +1,105 @@ +"""The ACA overlay charges the ACA surtaxes in every regime's after-tax income. + +- Every ACA variant charges the 3.8% surtax on unearned income above 200,000. +- The full ACA also charges the 0.9% surtax on earnings above 200,000. + +The tests evaluate the real regime DAG after the ACA overlay, with income +before the surtaxes pinned to a constant, so they check the wiring rather +than the tax schedule. +""" + +import inspect + +import jax.numpy as jnp +import pytest +from dags import concatenate_functions + +from aca_model.aca.health_insurance import PolicyVariant +from aca_model.aca.regimes._overrides import apply_aca_overrides +from aca_model.baseline.regimes import REGIME_SPECS +from aca_model.baseline.regimes._common import build_model_functions +from aca_model.baseline.regimes._nongroup import _build_functions as nongroup_functions +from aca_model.baseline.regimes._retiree import _build_functions as retiree_functions +from aca_model.baseline.regimes._tied import _build_functions as tied_functions +from aca_model.environment import taxes + +INCOME_BEFORE_SURTAXES = 100_000.0 + + +def _after_tax_income( + *, + regime_name: str, + policy: PolicyVariant, + labor_income: float, + unearned_income: float, +) -> float: + functions = { + **build_model_functions(), + **_regime_functions(regime_name), + } + apply_aca_overrides(functions, REGIME_SPECS[regime_name], policy) + functions = { + name: (lambda: jnp.asarray(INCOME_BEFORE_SURTAXES)) + if func is taxes.after_tax_income + else func + for name, func in functions.items() + if callable(func) + } + functions["labor_income"] = lambda: jnp.asarray(labor_income) + functions["gross_income"] = lambda: jnp.asarray(labor_income + unearned_income) + combined = concatenate_functions( + functions, targets=["after_tax_income"], return_type="tuple" + ) + assert not inspect.signature(combined).parameters + return float(combined()[0]) + + +def _regime_functions(regime_name: str) -> dict: + """Regime-level functions of the baseline regime, as its builder sets them.""" + build = { + "retiree": retiree_functions, + "tied": tied_functions, + "nongroup": nongroup_functions, + }[REGIME_SPECS[regime_name]["his"]] + return build(REGIME_SPECS[regime_name]) + + +@pytest.mark.parametrize( + ("regime_name", "policy", "labor_income", "unearned_income", "expected"), + [ + # 3.8% of 50,000 unearned income above 200,000. + ("retiree_nomc_inelig_canwork", PolicyVariant.ACA, 0.0, 250_000.0, 98_100.0), + # 0.9% of 50,000 earnings above 200,000. + ("nongroup_nomc_inelig_canwork", PolicyVariant.ACA, 250_000.0, 0.0, 99_550.0), + # No earnings surtax outside the full ACA. + ( + "tied_nomc_inelig_canwork", + PolicyVariant.ACA_NO_MANDATE, + 250_000.0, + 0.0, + 100_000.0, + ), + ( + "nongroup_oamc_forced_forcedout", + PolicyVariant.ACA_ONLY_MEDICAID_EXPANSION, + 0.0, + 250_000.0, + 98_100.0, + ), + ], +) +def test_aca_after_tax_income_charges_the_aca_surtaxes( + regime_name: str, + policy: PolicyVariant, + labor_income: float, + unearned_income: float, + expected: float, +) -> None: + """After-tax income under the ACA is income before the surtaxes less them.""" + result = _after_tax_income( + regime_name=regime_name, + policy=policy, + labor_income=labor_income, + unearned_income=unearned_income, + ) + assert result == pytest.approx(expected) diff --git a/tests/test_additive_shock_medical_bill.py b/tests/test_additive_shock_medical_bill.py new file mode 100644 index 0000000..9fff9d0 --- /dev/null +++ b/tests/test_additive_shock_medical_bill.py @@ -0,0 +1,103 @@ +"""Declaring the asset law as resources plus the negated medical bill is exact. + +With `subtract_hcc_bill`, the asset law toward every target is +`lcm.AdditiveShockTransition`: end-of-period resources before the out-of-pocket +bill, plus the shock `negative_oop_costs`. pylcm then averages each target's value +over the transitory medical-cost shock once per period, on the merged points +`{a_j - shock_k}`, rather than interpolating it at every shock node of every source +point. With linear interpolation in assets that average is exact, so the values +agree with the plain asset law to rounding. +""" + +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 test_negative_oop_costs_is_the_negated_bill() -> None: + np.testing.assert_array_equal( + negative_oop_costs(oop_costs=jnp.array([0.0, 1.5, 250.0])), + np.array([0.0, -1.5, -250.0]), + ) + + +def test_every_target_adds_the_negated_bill(additive_model: lcm.Model) -> None: + for name, regime in additive_model.user_regimes.items(): + if name == "dead": + continue + laws = regime.state_transitions["assets"] + assert isinstance(laws, Mapping) + assert laws + for law in laws.values(): + assert isinstance(law, lcm.AdditiveShockTransition) + assert law.shock == "negative_oop_costs" + + +@pytest.mark.long_running +def test_additive_shock_value_equals_the_plain_law_value( + additive_model: lcm.Model, +) -> None: + """The two routes sum the same terms in a different order.""" + folded_model = make_baseline_model(grid_config=_FOLDED) + _, _, 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 + assert set(additive_values) == set(folded_values) + for period, regimes in folded_values.items(): + assert set(additive_values[period]) == set(regimes) + for name, values in regimes.items(): + np.testing.assert_array_equal( + np.isfinite(additive_values[period][name]), np.isfinite(values) + ) + np.testing.assert_allclose( + additive_values[period][name], + values, + rtol=1e-9, + atol=0, + equal_nan=False, + err_msg=f"period={period}, regime={name}", + ) + + +@pytest.mark.parametrize( + ("bill", "expected_living", "expected_dead"), + [(0.0, 85.0, 75.0), (20.0, 65.0, 55.0), (250.0, -165.0, -175.0)], +) +def test_asset_law_preserves_resources_minus_medical_costs( + bill: float, expected_living: float, expected_dead: float +) -> None: + common = { + "cash_on_hand": jnp.array(100.0), + "transfers": jnp.array(5.0), + "consumption_dollars": jnp.array(30.0), + } + shock = negative_oop_costs(oop_costs=jnp.array(bill)) + living = ( + assets_and_income.assets_before_bill( + **common, pension_assets_adjustment=jnp.array(10.0) + ) + + shock + ) + dead = assets_and_income.assets_before_bill_when_dead(**common) + shock + np.testing.assert_allclose(living, expected_living, rtol=0, atol=0) + np.testing.assert_allclose(dead, expected_dead, rtol=0, atol=0) diff --git a/tests/test_aime_grid.py b/tests/test_aime_grid.py index a249ed3..d63da86 100644 --- a/tests/test_aime_grid.py +++ b/tests/test_aime_grid.py @@ -9,7 +9,6 @@ _AIME_PIECE_N_POINTS, _build_aime_grid, ) -from aca_model.config import BENCHMARK_GRID_CONFIG # Production SSA bend points plus the delayed-retirement-credit extension: # 0, kink_0, kink_1, taxable-max, and the extension point that carries the @@ -18,19 +17,20 @@ _FIXED_PARAMS = MappingProxyType({"pia_aime_grid": _PIA_AIME_GRID}) -def test_build_aime_grid_has_four_segments() -> None: - """The AIME grid spans all four PIA intervals, including the extension.""" - grid = _build_aime_grid( - grid_config=BENCHMARK_GRID_CONFIG, fixed_params=_FIXED_PARAMS +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) + np.testing.assert_allclose( + [point.value for point in grid.breakpoints], + _PIA_AIME_GRID[1:-1], ) - assert len(grid.segments) == 4 + assert tuple(point.owner for point in grid.breakpoints) == ("right",) * 3 + assert grid.points_per_segment == _AIME_PIECE_N_POINTS 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( - grid_config=BENCHMARK_GRID_CONFIG, fixed_params=_FIXED_PARAMS - ) + grid = _build_aime_grid(fixed_params=_FIXED_PARAMS) np.testing.assert_allclose(float(grid.to_jax().max()), 187954.752, rtol=1e-5) diff --git a/tests/test_assets_grid.py b/tests/test_assets_grid.py new file mode 100644 index 0000000..6d61223 --- /dev/null +++ b/tests/test_assets_grid.py @@ -0,0 +1,49 @@ +"""Tests for the asset state grid.""" + +import numpy as np +import pytest + +from aca_model.baseline.regimes._common import build_assets_grid + +FLOOR = -221_270.85 + + +def _points(n_points: int = 24) -> np.ndarray: + return np.asarray(build_assets_grid(floor=FLOOR, n_points=n_points).to_jax()) + + +def test_assets_grid_top_covers_largest_initial_assets() -> None: + """The top node is 12 million, the largest initial asset holding.""" + assert _points()[-1] == 12_000_000.0 + + +def test_assets_grid_bottom_is_the_borrowing_limit() -> None: + """The bottom node is the borrowing limit passed in.""" + assert _points()[0] == FLOOR + + +def test_assets_grid_has_a_node_at_zero() -> None: + """Zero assets is a node: the asset tests and the floor bite next to it.""" + assert np.any(_points() == 0.0) + + +def test_assets_grid_is_dense_near_zero() -> None: + """At least five nodes lie in [0, 30k], as in struct-ret's grid.""" + points = _points() + assert int(((points >= 0.0) & (points <= 30_000.0)).sum()) >= 5 + + +@pytest.mark.parametrize("n_points", [3, 8, 24]) +def test_assets_grid_has_requested_number_of_points(n_points: int) -> None: + """The node count is the configured `n_assets_gridpoints`.""" + assert len(_points(n_points)) == n_points + + +@pytest.mark.parametrize( + ("index", "expected"), + [(5, 0.0), (6, 4_460.02), (9, 27_305.37), (15, 377_388.77), (19, 2_128_241.24)], +) +def test_assets_grid_nodes_are_sinh_spaced(index: int, expected: float) -> None: + """Positive nodes are `10k * sinh(u)` with `u` evenly spaced from 0 to + `asinh(12M / 10k)` over 19 nodes.""" + np.testing.assert_allclose(_points()[index], expected, rtol=1e-6) diff --git a/tests/test_baseline_equivalence.py b/tests/test_baseline_equivalence.py index e7da710..d1e30c9 100644 --- a/tests/test_baseline_equivalence.py +++ b/tests/test_baseline_equivalence.py @@ -92,6 +92,7 @@ def test_baseline_primary_oop_no_cost_sharing_scale() -> None: result = health_insurance.primary_oop( total_health_costs=costs, buy_private=jnp.array(BuyPrivate.yes), + is_married=jnp.int32(0), deductible=deductible, coinsurance_rate=coinsurance, oop_max=oop_max_val, @@ -115,6 +116,7 @@ def test_aca_primary_oop_scaled_reduces_costs() -> None: total_health_costs=costs, cost_sharing_scale=jnp.array(1.0), buy_private=jnp.array(BuyPrivate.yes), + is_married=jnp.int32(0), deductible=deductible, coinsurance_rate=coinsurance, oop_max=oop_max_val, @@ -123,6 +125,7 @@ def test_aca_primary_oop_scaled_reduces_costs() -> None: total_health_costs=costs, cost_sharing_scale=jnp.array(0.3), buy_private=jnp.array(BuyPrivate.yes), + is_married=jnp.int32(0), deductible=deductible, coinsurance_rate=coinsurance, oop_max=oop_max_val, diff --git a/tests/test_beartype_claw.py b/tests/test_beartype_claw.py index e4130be..b5824d1 100644 --- a/tests/test_beartype_claw.py +++ b/tests/test_beartype_claw.py @@ -12,14 +12,15 @@ import pytest from beartype.roar import BeartypeCallHintViolation -from helpers.model import make_baseline_model # ty: ignore[unresolved-import] + +from aca_model.benchmark import create_benchmark_model def test_claw_checks_aca_model() -> None: """An ill-typed argument to an `aca_model` function is rejected by beartype. - `create_model` annotates `n_subjects` as `int`; passing a string is caught - by the claw before the value reaches pylcm's own `Model` perimeter. + The benchmark factory requires a DiscreteGrid for preference types. + Invalid types are rejected before model construction. """ with pytest.raises(BeartypeCallHintViolation): - make_baseline_model(n_subjects="not an int") + create_benchmark_model(pref_type_grid="not a grid") # ty: ignore[invalid-argument-type] diff --git a/tests/test_benchmark.py b/tests/test_benchmark.py index fb1a9d6..74b3af8 100644 --- a/tests/test_benchmark.py +++ b/tests/test_benchmark.py @@ -2,9 +2,16 @@ import numpy as np import pytest -from lcm import DiscreteGrid +from lcm import ( + DiscreteGrid, + GridBreakpoint, + IrregSpacedGrid, + PiecewiseLinSpacedGrid, + RouwenhorstAR1Process, +) from aca_model.agent.preferences import BenchmarkPrefType +from aca_model.baseline.regimes import REGIME_SPECS from aca_model.benchmark import ( create_benchmark_model, get_benchmark_initial_conditions, @@ -12,11 +19,38 @@ ) +def test_benchmark_model_builds_with_the_current_pylcm_grid_api() -> None: + """The frozen benchmark preserves its state and action grid extents.""" + model = create_benchmark_model( + pref_type_grid=DiscreteGrid(BenchmarkPrefType), + ) + + aime = model.user_regimes["retiree_nomc_inelig_canwork"].states["aime"] + assert isinstance(aime, PiecewiseLinSpacedGrid) + assert all(isinstance(point, GridBreakpoint) for point in aime.breakpoints) + assert all(point.owner == "right" for point in aime.breakpoints) + assert aime.n_points == 38 + regime = model.user_regimes["retiree_nomc_inelig_canwork"] + assert len(model.user_regimes) == 19 + assert model.n_periods == 46 + assets = regime.states["assets"] + wage_res = regime.states["log_ft_wage_res"] + pref_type = regime.states["pref_type"] + consumption = regime.actions["consumption_dollars"] + assert isinstance(assets, IrregSpacedGrid) + assert isinstance(wage_res, RouwenhorstAR1Process) + assert isinstance(pref_type, DiscreteGrid) + assert isinstance(consumption, IrregSpacedGrid) + assert assets.n_points == 3 + assert wage_res.n_points == 3 + assert len(pref_type.categories) == 2 + assert consumption.n_points == 5 + + @pytest.mark.long_running def test_benchmark_model_simulates_end_to_end() -> None: n_subjects = 20 model = create_benchmark_model( - n_subjects=n_subjects, pref_type_grid=DiscreteGrid(BenchmarkPrefType), ) _, _, params = get_benchmark_params(model=model) @@ -27,11 +61,10 @@ def test_benchmark_model_simulates_end_to_end() -> None: result = model.simulate( params=params, initial_conditions=initial_conditions, - period_to_regime_to_V_arr=None, log_level="off", ) - df = result.to_dataframe() + df = result.to_dataframe(terminal_rows="all") assert len(df) == n_subjects * model.n_periods # Period 0 rows reflect initial conditions — no NaN in continuous states. period_0 = df.loc[df["period"] == 0] @@ -51,7 +84,6 @@ def test_benchmark_panel_exposes_hic_premium_and_wage_targets() -> None: """ n_subjects = 20 model = create_benchmark_model( - n_subjects=n_subjects, pref_type_grid=DiscreteGrid(BenchmarkPrefType), ) _, _, params = get_benchmark_params(model=model) @@ -62,7 +94,6 @@ def test_benchmark_panel_exposes_hic_premium_and_wage_targets() -> None: result = model.simulate( params=params, initial_conditions=initial_conditions, - period_to_regime_to_V_arr=None, log_level="off", ) @@ -114,7 +145,6 @@ def test_benchmark_simulate_obeys_borrowing_constraint() -> None: """ n_subjects = 4 model = create_benchmark_model( - n_subjects=n_subjects, pref_type_grid=DiscreteGrid(BenchmarkPrefType), ) _, _, params = get_benchmark_params(model=model) @@ -125,7 +155,6 @@ def test_benchmark_simulate_obeys_borrowing_constraint() -> None: result = model.simulate( params=params, initial_conditions=initial_conditions, - period_to_regime_to_V_arr=None, log_level="off", ) @@ -139,3 +168,54 @@ def test_benchmark_simulate_obeys_borrowing_constraint() -> None: f"borrowing_constraint violated on {int((slack < 0).sum())} row(s); " f"min slack = {slack.min():.6g}" ) + + +def test_initial_conditions_choose_population_size_after_model_construction() -> None: + """One model supports reproducible initial conditions with different row counts.""" + model = create_benchmark_model(pref_type_grid=DiscreteGrid(BenchmarkPrefType)) + 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) + + assert small.keys() == repeated.keys() == large.keys() + for name, values in small.items(): + assert values.shape == (2,) + assert large[name].shape == (5,) + np.testing.assert_array_equal(values, repeated[name]) + np.testing.assert_array_equal(small["age"], [51.0, 51.0]) + np.testing.assert_array_equal(small["claimed_ss"], [0, 0]) + + +_LAW_ARGS = ("survival_probs", "prob_disabled_next") + + +def test_benchmark_law_args_live_under_their_source_edges() -> None: + """Each source's law arguments sit at `fixed_params["edges"][source]`.""" + fixed_params, _, _ = get_benchmark_params(model=None) + sources_with_both = { + source + for source, args in fixed_params["edges"].items() + if set(args) == set(_LAW_ARGS) + } + sources_with_survival_only = { + source + for source, args in fixed_params["edges"].items() + if set(args) == {"survival_probs"} + } + forcedout = {name for name in REGIME_SPECS if name.endswith("_forcedout")} + assert (sources_with_both, sources_with_survival_only) == ( + set(REGIME_SPECS) - forcedout, + forcedout, + ) + + +def test_benchmark_regimes_hold_no_law_args() -> None: + """No law argument stays under a regime, where it would not reach the law.""" + fixed_params, _, _ = get_benchmark_params(model=None) + leftovers = [ + (name, arg) + for name in REGIME_SPECS + for arg in _LAW_ARGS + if arg in fixed_params[name] + ] + assert leftovers == [] diff --git a/tests/test_benchmark_folding.py b/tests/test_benchmark_folding.py new file mode 100644 index 0000000..71c096f --- /dev/null +++ b/tests/test_benchmark_folding.py @@ -0,0 +1,64 @@ +"""The ACA benchmark activates transitory-shock folding explicitly.""" + +from collections.abc import Mapping +from dataclasses import replace + +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.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: + model = benchmark_model + living = { + name: regime for name, regime in model.user_regimes.items() if name != "dead" + } + assert living + assert all( + "hcc_transitory" not in model.state_names(regime_name=name) for name in living + ) + for regime in living.values(): + laws = regime.state_transitions["assets"] + assert isinstance(laws, Mapping) + assert all( + isinstance(law, lcm.AdditiveShockTransition) for law in laws.values() + ) + initial = get_benchmark_initial_conditions(model=model, n_subjects=8, seed=42) + assert "hcc_transitory" not in initial + + +def test_unfolded_control_has_identical_shared_initial_conditions( + benchmark_model: lcm.Model, +) -> None: + 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 + ) + assert all( + "hcc_transitory" in control.state_names(regime_name=name) + for name in control.user_regimes + if name != "dead" + ) + folded_initial = get_benchmark_initial_conditions( + model=benchmark_model, n_subjects=8, seed=42 + ) + control_initial = get_benchmark_initial_conditions( + model=control, n_subjects=8, seed=42, grid_config=control_grid + ) + assert set(control_initial) == set(folded_initial) | {"hcc_transitory"} + for name, values in folded_initial.items(): + np.testing.assert_array_equal(values, control_initial[name]) diff --git a/tests/test_budget_chain_integration.py b/tests/test_budget_chain_integration.py index 9d38f86..8388061 100644 --- a/tests/test_budget_chain_integration.py +++ b/tests/test_budget_chain_integration.py @@ -61,7 +61,7 @@ def test_working_agent_cash_on_hand() -> None: assets=jnp.array(50000.0), rate_of_return=jnp.asarray(0.03), labor_income=jnp.array(40000.0), - spousal_income_amounts=jnp.array([0.0, 0.0, 20000.0]), + spousal_income_amount=jnp.array(0.0), ss_benefit=jnp.array(0.0), pension_benefit=jnp.array(0.0), spousal_income=jnp.int32(0), @@ -97,7 +97,7 @@ def test_retired_agent_with_pension() -> None: assets=jnp.array(200000.0), rate_of_return=jnp.asarray(0.03), labor_income=jnp.array(0.0), - spousal_income_amounts=jnp.array([0.0, 0.0, 20000.0]), + spousal_income_amount=jnp.array(0.0), ss_benefit=jnp.array(15000.0), pension_benefit=jnp.array(10000.0), spousal_income=jnp.int32(0), diff --git a/tests/test_budget_rules.py b/tests/test_budget_rules.py new file mode 100644 index 0000000..7797ce0 --- /dev/null +++ b/tests/test_budget_rules.py @@ -0,0 +1,389 @@ +"""Budget, Social Security and pension rules at their production values. + +Expected values come from a literal scalar transcription of struct-ret's +formulas (or of the draft's formula where struct-ret departs from it), +evaluated at the same inputs. +""" + +import inspect +from collections.abc import Callable +from typing import Any + +import jax.numpy as jnp +import numpy as np +import pytest +from helpers import production_rules as pr + +from aca_model.agent import labor_market +from aca_model.agent.labor_market import LaborSupply, SpousalIncome +from aca_model.baseline.regimes._common import ( + REGIME_SPECS, + _select_aime_law, + build_model_constraints, +) +from aca_model.environment import pensions, social_security, taxes +from aca_model.environment.social_security import ClaimedSS + +PIA_AIME_GRID = jnp.asarray(pr.PIA_AIME_GRID) +PIA_TABLE = jnp.asarray(pr.PIA_TABLE) +EARLY_RET_ADJUSTMENT = jnp.asarray(pr.EARLY_RET_ADJUSTMENT) +EARNINGS_TEST_CREDITED_BACK = jnp.asarray(pr.EARNINGS_TEST_CREDITED_BACK) +RATIO_LOWEST_EARNINGS = jnp.asarray(pr.RATIO_LOWEST_EARNINGS) +DI_DROPOUT_SCALE = jnp.asarray(pr.DI_DROPOUT_SCALE) +DI_DROPOUT_NEXT_PERIOD_RATIO = jnp.asarray(pr.DI_DROPOUT_NEXT_PERIOD_RATIO) + + +def _shift_one_period_forward(values: tuple) -> jnp.ndarray: + """Return the period-indexed array whose row `t` carries the value at `t + 1`.""" + arr = jnp.asarray(values) + return jnp.concatenate([arr[1:], arr[-1:]], axis=0) + + +IMP_NEXT_PERIOD = { + "imp_intercept_next_period": _shift_one_period_forward(pr.IMP_INTERCEPT), + "imp_pia_coeff_next_period": _shift_one_period_forward(pr.IMP_PIA_COEFF), + "imp_pia_kink_0_coeff_next_period": _shift_one_period_forward( + pr.IMP_PIA_KINK_0_COEFF + ), + "imp_pia_kink_1_coeff_next_period": _shift_one_period_forward( + pr.IMP_PIA_KINK_1_COEFF + ), + "imp_kink_0_next_period": _shift_one_period_forward(pr.IMP_KINK_0), + "imp_kink_1_next_period": _shift_one_period_forward(pr.IMP_KINK_1), +} + +SS_RULES = { + "pia_table": PIA_TABLE, + "pia_aime_grid": PIA_AIME_GRID, + "normal_retirement_age": jnp.int32(pr.NORMAL_RETIREMENT_AGE), + "early_ret_adjustment": EARLY_RET_ADJUSTMENT, + "earnings_test_credited_back": EARNINGS_TEST_CREDITED_BACK, + "earnings_test_repealed_age": jnp.int32(pr.EARNINGS_TEST_REPEALED_AGE), + "aime_accrual_factor": jnp.asarray(pr.AIME_ACCRUAL_FACTOR), + "aggregate_wage_growth": jnp.asarray(pr.AGGREGATE_WAGE_GROWTH), + "aime_last_age_with_indexing": jnp.int32(pr.AIME_LAST_AGE_WITH_INDEXING), + "aime_kink_2": jnp.asarray(pr.AIME_KINK_2), + "ratio_lowest_earnings": RATIO_LOWEST_EARNINGS, + "medicare_age": jnp.int32(pr.MEDICARE_AGE), + "di_dropout_scale": DI_DROPOUT_SCALE, + "di_dropout_next_period_ratio": DI_DROPOUT_NEXT_PERIOD_RATIO, +} + + +def _period(age: int) -> jnp.ndarray: + return jnp.int32(age - pr.START_AGE) + + +def _call(func: Callable[..., Any], **kwargs: Any) -> Any: + """Call `func` with the subset of `kwargs` its signature names.""" + names = inspect.signature(func).parameters + return func(**{k: v for k, v in kwargs.items() if k in names}) + + +@pytest.mark.parametrize( + ("age", "pia_unadjusted", "pia_adjusted", "full_benefit_next", "expected"), + [ + (62, 18479.36, 17324.4, 35680.75116959574, 18117.163555158222), + (63, 18479.36, 17057.87076922421, 37436.31858689668, 18035.38135778553), + (66, 18479.36, 19957.7088, 44113.51778285309, 19000.883559951508), + (68, 24710.04, 26414.180689655175, 58432.219638177565, 25466.839034926863), + ], +) +def test_total_to_pia_keeps_social_security_plus_imputed_pension_total( + age: int, + pia_unadjusted: float, + pia_adjusted: float, + full_benefit_next: float, + expected: float, +) -> None: + """The PIA carried into the next period reproduces the claim-adjusted total. + + For a retiree-HIS household at a 18.19% marginal tax rate, + `PIA* + (1 - τ) pbmax_{t+1}(PIA*) = PIA_adj + (1 - τ) pbmax_{t+1}(PIA_unadj)`. + """ + result = pensions.total_to_pia( + pia_adjusted_next_period=jnp.asarray(pia_adjusted), + pia_unadjusted_next_period=jnp.asarray(pia_unadjusted), + full_benefit_next_period=jnp.asarray(full_benefit_next), + target_his=jnp.int32(0), + period=_period(age), + marginal_tax_rate=jnp.asarray(0.1819), + **IMP_NEXT_PERIOD, + ) + np.testing.assert_allclose(result, expected, rtol=1e-10) + + +def test_total_to_pia_without_claim_adjustment_returns_unadjusted_pia() -> None: + """Without a claim adjustment or credit the carried PIA is the accrued PIA.""" + result = pensions.total_to_pia( + pia_adjusted_next_period=jnp.asarray(5000.0), + pia_unadjusted_next_period=jnp.asarray(5000.0), + full_benefit_next_period=jnp.asarray(0.0), + target_his=jnp.int32(0), + period=_period(62), + marginal_tax_rate=jnp.asarray(0.1819), + **IMP_NEXT_PERIOD, + ) + np.testing.assert_allclose(result, 5000.0, rtol=1e-12) + + +@pytest.mark.parametrize( + ("labor_supply", "ss_benefit", "expected"), + [ + # Claims at 62 and does not work: nothing is withheld. + (LaborSupply.do_not_work, 13859.52, 0.0), + # Claims at 62 and earns 10,000, below the 15,480 threshold. + (LaborSupply.h1000, 13859.52, 0.0), + # Claims at 62 and earns 30,000: 7,260 of the 13,859.52 benefit withheld. + (LaborSupply.h1000, 6599.52, 0.5238276650273602), + ], +) +def test_benefit_withheld_fraction_counts_only_earnings_test_withholding( + labor_supply: int, ss_benefit: float, expected: float +) -> None: + """The withheld fraction is the earnings-test share of the claimed benefit. + + AIME is credited back only for benefits the earnings test withheld, not + for the early-claiming reduction. AIME 40,000 (PIA 18,479.36), age 62. + """ + result = _call( + social_security.benefit_withheld_fraction, + pia=jnp.asarray(18479.36), + ss_benefit=jnp.asarray(ss_benefit), + age=jnp.int32(62), + period=_period(62), + claim_ss=jnp.int32(ClaimedSS.yes), + claimed_ss=jnp.int32(ClaimedSS.no), + labor_supply=jnp.int32(labor_supply), + normal_retirement_age=jnp.int32(pr.NORMAL_RETIREMENT_AGE), + early_ret_adjustment=EARLY_RET_ADJUSTMENT, + ) + np.testing.assert_allclose(result, expected, atol=1e-6) + + +SPOUSAL_INCOME_TABLE = jnp.zeros((45, 2, 3)) +SPOUSAL_INCOME_TABLE = SPOUSAL_INCOME_TABLE.at[4, :, SpousalIncome.married_has_inc].set( + jnp.asarray([20299.103515625, 21513.8984375]) +) +SPOUSAL_INCOME_TABLE = SPOUSAL_INCOME_TABLE.at[ + 11, :, SpousalIncome.married_has_inc +].set(jnp.asarray([22837.513671875, 23848.978515625])) + + +@pytest.mark.parametrize( + ("age", "good_health", "spousal_income", "expected"), + [ + (55, 0, SpousalIncome.married_has_inc, 20299.103515625), + (55, 1, SpousalIncome.married_has_inc, 21513.8984375), + (62, 1, SpousalIncome.married_has_inc, 23848.978515625), + (62, 1, SpousalIncome.married_no_inc, 0.0), + (62, 0, SpousalIncome.single, 0.0), + ], +) +def test_spousal_income_amount_varies_with_age_and_own_health( + age: int, good_health: int, spousal_income: int, expected: float +) -> None: + """Spousal income is read from the age × own-health × spousal-income table.""" + result = labor_market.spousal_income_amount( + period=_period(age), + good_health=jnp.int32(good_health), + spousal_income=jnp.int32(spousal_income), + spousal_income_amounts=SPOUSAL_INCOME_TABLE, + ) + np.testing.assert_allclose(result, expected, rtol=1e-12) + + +@pytest.mark.parametrize( + ("labor_income", "expected"), + [ + # Indexed by 0.8%, no earnings: 40,320 times the DI continuity ratio. + (0.0, 40531.764705882364), + # Indexed and accrues 12,000/35 of earnings, below the SGA threshold. + (12000.0, 40876.42256902762), + ], +) +def test_disabled_aime_accrues_indexing_and_earnings( + labor_income: float, expected: float +) -> None: + """A disabled agent's AIME is indexed and accrues earnings like anyone's. + + The DI continuity ratio then scales the accrued AIME. Age 55, AIME 40,000. + """ + result = social_security.next_aime_disabled_plain( + aime=jnp.asarray(40000.0), + labor_income=jnp.asarray(labor_income), + period=_period(55), + age=jnp.int32(55), + health=jnp.int32(0), + benefit_withheld_fraction=jnp.asarray(0.0), + earnings_test_credited_back=EARNINGS_TEST_CREDITED_BACK, + earnings_test_repealed_age=jnp.int32(pr.EARNINGS_TEST_REPEALED_AGE), + pia_table=PIA_TABLE, + pia_aime_grid=PIA_AIME_GRID, + aime_accrual_factor=jnp.asarray(pr.AIME_ACCRUAL_FACTOR), + aggregate_wage_growth=jnp.asarray(pr.AGGREGATE_WAGE_GROWTH), + aime_last_age_with_indexing=jnp.int32(pr.AIME_LAST_AGE_WITH_INDEXING), + aime_kink_2=jnp.asarray(pr.AIME_KINK_2), + ratio_lowest_earnings=RATIO_LOWEST_EARNINGS, + medicare_age=jnp.int32(pr.MEDICARE_AGE), + di_dropout_scale=DI_DROPOUT_SCALE, + di_dropout_next_period_ratio=DI_DROPOUT_NEXT_PERIOD_RATIO, + ) + np.testing.assert_allclose(result, expected, rtol=1e-10) + + +@pytest.mark.parametrize( + ("age", "aime", "labor_income"), + [ + (70, 40000.0, 60000.0), + (71, 131305.0, 0.0), + ], +) +def test_aime_is_frozen_once_claiming_is_forced( + age: int, aime: float, labor_income: float +) -> None: + """From the forced-claim age on, AIME neither accrues nor is capped.""" + law = _select_aime_law(REGIME_SPECS["retiree_oamc_forced_canwork"]) + result = _call( + law, + aime=jnp.asarray(aime), + labor_income=jnp.asarray(labor_income), + period=_period(age), + age=jnp.int32(age), + benefit_withheld_fraction=jnp.asarray(0.0), + **SS_RULES, + ) + np.testing.assert_allclose(result, aime, rtol=1e-12) + + +def test_accrued_aime_is_capped_at_the_taxable_maximum() -> None: + """Accrual from earnings cannot lift AIME above the taxable maximum. + + Age 66, AIME 116,000, earnings 117,000: AIME accrues to the 117,000 cap, + whose PIA is 33,260.04. + """ + result = social_security.pia_unadjusted_next_period( + aime=jnp.asarray(116000.0), + labor_income=jnp.asarray(117000.0), + period=_period(66), + age=jnp.int32(66), + pia_table=PIA_TABLE, + pia_aime_grid=PIA_AIME_GRID, + aime_accrual_factor=jnp.asarray(pr.AIME_ACCRUAL_FACTOR), + aggregate_wage_growth=jnp.asarray(pr.AGGREGATE_WAGE_GROWTH), + aime_last_age_with_indexing=jnp.int32(pr.AIME_LAST_AGE_WITH_INDEXING), + aime_kink_2=jnp.asarray(pr.AIME_KINK_2), + ratio_lowest_earnings=RATIO_LOWEST_EARNINGS, + ) + np.testing.assert_allclose(result, 33260.04, atol=0.01) + + +@pytest.mark.parametrize( + ("age", "his", "expected"), + [ + (70, 0, -190.86435823930114), + (70, 1, 696.7748524892276), + (71, 0, 0.0), + (71, 1, 0.0), + ], +) +def test_pension_accrual_stops_at_the_pension_claiming_age( + age: int, his: int, expected: float +) -> None: + """Pension wealth accrues from earnings through age 70 and not at 71. + + Earnings 40,000. + """ + result = _call( + pensions.accrual, + labor_income=jnp.asarray(40000.0), + period=_period(age), + age=jnp.int32(age), + his=jnp.int32(his), + accrual_intercept=jnp.asarray(pr.ACCRUAL_INTERCEPT), + accrual_log_earnings=jnp.asarray(pr.ACCRUAL_LOG_EARNINGS), + accrual_prob_intercept=jnp.asarray(pr.ACCRUAL_PROB_INTERCEPT), + accrual_prob_log_earnings=jnp.asarray(pr.ACCRUAL_PROB_LOG_EARNINGS), + accrual_prob_log_earnings_sq=jnp.asarray(pr.ACCRUAL_PROB_LOG_EARNINGS_SQ), + ) + np.testing.assert_allclose(result, expected, rtol=1e-8, atol=1e-8) + + +SINGLE_FLOOR = 1597.09 +MARRIED_FLOOR = SINGLE_FLOOR * 2.0**0.7 + + +@pytest.mark.parametrize( + ("consumption", "expected"), + [ + (SINGLE_FLOOR, False), + (MARRIED_FLOOR, True), + ], +) +def test_married_household_cannot_consume_below_its_own_floor( + consumption: float, expected: bool +) -> None: + """A married household at the transfer floor consumes exactly its own floor. + + Cash on hand 500 is topped up to the married floor; consuming the lower + single floor is infeasible. + """ + constraints = build_model_constraints(solver="brute_force") + feasible = all( + bool( + _call( + constraint, + consumption_dollars=jnp.asarray(consumption), + cash_on_hand=jnp.asarray(500.0), + consumption_dollars_floor=jnp.asarray(MARRIED_FLOOR), + hcc_persistent=jnp.asarray(0.0), + hcc_transitory=jnp.asarray(0.0), + ) + ) + for constraint in constraints.values() + ) + assert feasible is expected + + +@pytest.mark.parametrize( + ("gross_income", "labor_income", "expected_surtax"), + [ + # Unearned income 250,000: 3.8% on the 50,000 above 200,000. + (250000.0, 0.0, 1900.0), + # Earnings 250,000: 0.9% on the 50,000 above 200,000. + (250000.0, 250000.0, 450.0), + # Earnings 250,000 and unearned income 250,000. + (500000.0, 250000.0, 2350.0), + # Unearned income 150,000 and earnings 150,000: both below the thresholds. + (300000.0, 150000.0, 0.0), + ], +) +def test_aca_surtaxes_reduce_after_tax_income( + gross_income: float, labor_income: float, expected_surtax: float +) -> None: + """The ACA adds a 3.8% tax on unearned and 0.9% on earned income above 200,000.""" + result = taxes.after_tax_income_with_aca_surtaxes( + after_tax_income_before_aca_surtaxes=jnp.asarray(100000.0), + gross_income=jnp.asarray(gross_income), + labor_income=jnp.asarray(labor_income), + ) + np.testing.assert_allclose(result, 100000.0 - expected_surtax, rtol=1e-12) + + +@pytest.mark.parametrize( + ("gross_income", "labor_income", "expected_surtax"), + [ + (250000.0, 0.0, 1900.0), + (250000.0, 250000.0, 0.0), + ], +) +def test_aca_investment_surtax_alone_spares_earnings( + gross_income: float, labor_income: float, expected_surtax: float +) -> None: + """Without the payroll surtax only unearned income above 200,000 is taxed.""" + result = taxes.after_tax_income_with_aca_investment_surtax( + after_tax_income_before_aca_surtaxes=jnp.asarray(100000.0), + gross_income=jnp.asarray(gross_income), + labor_income=jnp.asarray(labor_income), + ) + np.testing.assert_allclose(result, 100000.0 - expected_surtax, rtol=1e-12) diff --git a/tests/test_dag_coverage.py b/tests/test_dag_coverage.py new file mode 100644 index 0000000..16b6472 --- /dev/null +++ b/tests/test_dag_coverage.py @@ -0,0 +1,143 @@ +"""The baseline's solved domain derives from its declared starting pairs.""" + +import jax.numpy as jnp +import pandas as pd +import pytest +from helpers.model import make_baseline_model +from lcm.exceptions import InvalidInitialConditionsError + +from aca_model.baseline.regimes import ENTRY_REGIMES, REGIME_SPECS +from aca_model.baseline.regimes._common import build_model_edges +from aca_model.benchmark import get_benchmark_initial_conditions, get_benchmark_params +from aca_model.config import MODEL_CONFIG +from aca_model.simulation import select_admissible_starts + + +@pytest.fixture(scope="module") +def model(): + return make_baseline_model() + + +def test_initial_nodes_are_the_age_51_to_60_entry_regimes(model): + expected = { + (age, regime) + for age in range(MODEL_CONFIG.start_age, MODEL_CONFIG.last_start_age + 1) + for regime in ENTRY_REGIMES + } + assert {(float(a), r) for a, r in model.graph.initial_nodes} == { + (float(a), r) for a, r in expected + } + + +def test_solved_domain_has_184_nodes(model): + assert len(model.graph.nodes) == 184 + + +def test_dead_at_the_first_age_is_not_solved(model): + assert (51, "dead") not in model.graph.nodes + + +def test_later_living_nodes_are_solved_but_not_admissible(model): + node = (64, "retiree_nomc_choose_canwork") + assert (node in model.graph.nodes, node in model.graph.initial_nodes) == ( + True, + False, + ) + + +def test_regimes_carry_no_active_attribute(model): + assert not any(hasattr(r, "active") for r in model.user_regimes.values()) + + +@pytest.mark.parametrize( + ("age", "regime"), + [ + (61.0, "retiree_nomc_inelig_canwork"), + (62.0, "retiree_nomc_choose_canwork"), + (69.0, "retiree_oamc_choose_canwork"), + (51.0, "dead"), + (60.0, "dead"), + ], +) +def test_simulate_rejects_starts_outside_the_entry_contract(model, age, regime): + _, _, params = get_benchmark_params(model=model) + ic = get_benchmark_initial_conditions(model=model, n_subjects=2, seed=0) + ic = { + **ic, + "age": jnp.full(2, age), + "regime_id": jnp.full(2, model.regime_names_to_ids[regime], dtype=jnp.int32), + } + with pytest.raises(InvalidInitialConditionsError): + model.simulate(params=params, initial_conditions=ic, log_level="off") + + +def test_boundary_targets_change_with_source_age(model): + assert set( + model.graph.solution.targets(period=9, source="retiree_nomc_inelig_canwork") + ) == { + "retiree_nomc_inelig_canwork", + "retiree_dimc_inelig_canwork", + "nongroup_nomc_inelig_canwork", + "nongroup_dimc_inelig_canwork", + "dead", + } + assert set( + model.graph.solution.targets(period=10, source="retiree_nomc_inelig_canwork") + ) == { + "retiree_nomc_choose_canwork", + "retiree_dimc_choose_canwork", + "nongroup_nomc_choose_canwork", + "nongroup_dimc_choose_canwork", + "dead", + } + edges = model.graph.edges.solve["retiree_nomc_inelig_canwork"] + assert ( + 60 in edges["retiree_nomc_inelig_canwork"], + 61 in edges["retiree_nomc_inelig_canwork"], + ) == (True, False) + + +@pytest.mark.parametrize("regime", list(REGIME_SPECS)) +def test_each_regime_law_reaches_exactly_its_declared_targets(model, regime): + """The law's age schedule names each target at the ages that can reach it.""" + assert dict(model.graph.edges.solve[regime]) == { + target: frozenset(ages) for target, ages in build_model_edges()[regime].items() + } + + +def test_bequest_continuation_is_not_removed(model): + regime = model.user_regimes["retiree_oamc_forced_forcedout"] + assert model.graph.solution.targets( + period=44, source="retiree_oamc_forced_forcedout" + ) == ("dead",) + assert "consumption_dollars" in regime.actions + assert (95, "retiree_oamc_forced_forcedout") in model.graph.nodes + assert (96, "dead") in model.graph.nodes + + +def test_dead_regime_is_terminal_at_every_age(model): + assert model.graph.laws["dead"].terminal + assert {(age, "dead") for age in model.ages.exact_values[1:]} <= model.graph.nodes + + +def test_select_admissible_starts_drops_starts_at_61_and_above(model): + ic = pd.DataFrame( + { + "age": [60.0, 61.0, 62.0], + "regime_name": ["retiree_nomc_inelig_canwork"] * 2 + + ["retiree_nomc_choose_canwork"], + }, + index=pd.Index([1, 2, 3], name="id"), + ) + kept = select_admissible_starts(model=model, initial_conditions=ic) + assert list(kept.index) == [1] + + +def test_medical_cost_draws_keep_living_states_without_choice_constraints(model): + """Living regimes carry medical shocks; bequests depend only on assets and type.""" + for name, regime in model.user_regimes.items(): + assert "medical_cost_shocks_carried" not in regime.constraints + if name == "dead": + assert set(regime.states) == {"assets", "pref_type"} + else: + assert {"hcc_persistent", "hcc_transitory"} <= set(regime.states) diff --git a/tests/test_dcegm_model_creation.py b/tests/test_dcegm_model_creation.py index 2a4b8de..fc7645f 100644 --- a/tests/test_dcegm_model_creation.py +++ b/tests/test_dcegm_model_creation.py @@ -11,8 +11,9 @@ import numpy as np import pytest -from helpers.model import _DERIVED_CATEGORICALS # ty: ignore[unresolved-import] +from helpers.model import _DERIVED_CATEGORICALS from lcm import DiscreteGrid, IrregSpacedGrid, Model, Regime +from lcm.consumption_savings_regime import ConsumptionSavingsRegime from lcm.solvers import DCEGM from aca_model.agent import assets_and_income @@ -32,10 +33,17 @@ _FIXED_PARAMS, _WAGE_PARAMS, _ = get_benchmark_params(model=None) +# pylcm folds a shock state only under grid search, and the EGM solvers take the +# savings-form asset law, so the transitory medical-cost shock stays a state and +# the bill stays inside `next_assets`. +_EGM_GRID_CONFIG = dataclasses.replace( + BENCHMARK_GRID_CONFIG, fold_hcc_transitory=False, subtract_hcc_bill=False +) + def _build_regimes(solver: SolverName) -> dict[str, Regime]: return build_all_regimes( - grid_config=BENCHMARK_GRID_CONFIG, + grid_config=_EGM_GRID_CONFIG, fixed_params=_FIXED_PARAMS, wage_params=_WAGE_PARAMS, pref_type_grid=DiscreteGrid(BenchmarkPrefType), @@ -45,25 +53,25 @@ def _build_regimes(solver: SolverName) -> dict[str, Regime]: def _build_model(solver: SolverName) -> Model: return create_model( - n_subjects=1, fixed_params=_FIXED_PARAMS, wage_params=_WAGE_PARAMS, derived_categoricals=_DERIVED_CATEGORICALS, - grid_config=BENCHMARK_GRID_CONFIG, + grid_config=_EGM_GRID_CONFIG, pref_type_grid=DiscreteGrid(BenchmarkPrefType), solver=solver, ) def test_every_living_regime_gets_the_dcegm_solver() -> None: - """`solver="dcegm"` attaches a `DCEGM` config with assets as the Euler - state to every living regime; the terminal regime keeps the default.""" + """`solver="dcegm"` attaches a `DCEGM` config to every living regime, each + declaring `assets` as its liquid margin; the terminal regime keeps the + default.""" regimes = _build_regimes("dcegm") for name in REGIME_SPECS: - solver = regimes[name].solver - assert isinstance(solver, DCEGM), name - assert solver.continuous_state == "assets" - assert solver.continuous_action == "consumption_dollars" + regime = cast("ConsumptionSavingsRegime", regimes[name]) + assert isinstance(regime.solver, DCEGM), name + assert regime.liquid.state == "assets", name + assert regime.liquid.action == "consumption_dollars", name assert not isinstance(regimes["dead"].solver, DCEGM) @@ -91,7 +99,7 @@ def test_dcegm_model_slots_broadcast_contract_functions_and_constraint() -> None forward simulation re-decides consumption by an argmax over the consumption grid and needs the explicit feasibility mask.""" slots = build_model_slots( - grid_config=BENCHMARK_GRID_CONFIG, + grid_config=_EGM_GRID_CONFIG, fixed_params=_FIXED_PARAMS, wage_params=_WAGE_PARAMS, pref_type_grid=DiscreteGrid(BenchmarkPrefType), @@ -122,11 +130,10 @@ def test_dcegm_requires_construction_time_consumption_points() -> None: at model construction, so the runtime-injection path cannot be used.""" with pytest.raises(ValueError, match="consumption_dollars_points"): create_model( - n_subjects=1, fixed_params=_FIXED_PARAMS, wage_params=_WAGE_PARAMS, derived_categoricals=_DERIVED_CATEGORICALS, - grid_config=BENCHMARK_GRID_CONFIG, + grid_config=_EGM_GRID_CONFIG, pref_type_grid=DiscreteGrid(BenchmarkPrefType), solver="dcegm", ) @@ -143,80 +150,26 @@ def test_benchmark_consumption_points_pin_both_floors() -> None: np.testing.assert_allclose(points[:2], [floor, floor * 2.0**exponent], rtol=1e-12) -@pytest.mark.xfail( - strict=False, - reason=( - "pylcm's DC-EGM contract does not yet admit the ACA budget: the " - "assets law reaches `assets` outside the post-decision function — " - "through `oop_costs` (Medicaid eligibility → `countable_income` → " - "`capital_income`) and `pension_assets_adjustment` " - "(`marginal_tax_rate` → `gross_income` → `capital_income`). " - "Fixes land upstream in pylcm, not here." - ), -) def test_dcegm_benchmark_model_builds() -> None: - """The benchmark model accepts `solver="dcegm"` end to end. - - The acceptance criterion for the upstream DC-EGM stack: once pylcm's - contract admits the ACA budget chains, this builds without error. The - construction-time consumption points are supplied so the build reaches - the upstream limitation rather than the missing-points guard. - """ + """The benchmark model accepts `solver="dcegm"` end to end.""" model = create_model( - n_subjects=1, fixed_params=_FIXED_PARAMS, wage_params=_WAGE_PARAMS, derived_categoricals=_DERIVED_CATEGORICALS, - grid_config=BENCHMARK_GRID_CONFIG, + grid_config=_EGM_GRID_CONFIG, pref_type_grid=DiscreteGrid(BenchmarkPrefType), solver="dcegm", consumption_dollars_points=get_benchmark_consumption_dollars_points( - n_points=BENCHMARK_GRID_CONFIG.n_consumption_dollars_gridpoints + n_points=_EGM_GRID_CONFIG.n_consumption_dollars_gridpoints ), ) assert isinstance(model.user_regimes["retiree_nomc_inelig_canwork"].solver, DCEGM) -def test_savings_grid_batch_size_follows_grid_config() -> None: - """`GridConfig.n_savings_batch_size` sets the `batch_size` on every - living regime's DC-EGM savings grid, so the post-decision continuation - splays into `lax.map` blocks of that width.""" - grid_config = dataclasses.replace(BENCHMARK_GRID_CONFIG, n_savings_batch_size=50) - regimes = build_all_regimes( - grid_config=grid_config, - fixed_params=_FIXED_PARAMS, - wage_params=_WAGE_PARAMS, - pref_type_grid=DiscreteGrid(BenchmarkPrefType), - solver="dcegm", - ) - for name in REGIME_SPECS: - solver = cast("DCEGM", regimes[name].solver) - assert solver.savings_grid.batch_size == 50, name - - -def test_stochastic_node_batch_size_follows_grid_config() -> None: - """`GridConfig.n_stochastic_node_batch_size` sets `stochastic_node_batch_size` - on every living regime's DC-EGM solver, so the child stochastic-node - expectation splays into `lax.map` blocks of that width.""" - grid_config = dataclasses.replace( - BENCHMARK_GRID_CONFIG, n_stochastic_node_batch_size=7 - ) - regimes = build_all_regimes( - grid_config=grid_config, - fixed_params=_FIXED_PARAMS, - wage_params=_WAGE_PARAMS, - pref_type_grid=DiscreteGrid(BenchmarkPrefType), - solver="dcegm", - ) - for name in REGIME_SPECS: - solver = cast("DCEGM", regimes[name].solver) - assert solver.stochastic_node_batch_size == 7, name - - def test_savings_grid_length_follows_grid_config() -> None: """`GridConfig.n_savings_gridpoints` sets the number of nodes on every living regime's DC-EGM savings grid.""" - grid_config = dataclasses.replace(BENCHMARK_GRID_CONFIG, n_savings_gridpoints=70) + grid_config = dataclasses.replace(_EGM_GRID_CONFIG, n_savings_gridpoints=70) regimes = build_all_regimes( grid_config=grid_config, fixed_params=_FIXED_PARAMS, @@ -232,7 +185,7 @@ def test_savings_grid_length_follows_grid_config() -> None: def test_dcegm_savings_grid_rejects_too_few_points() -> None: """`n_savings_gridpoints < 2` cannot form the cubically clustered DC-EGM savings grid, so building the dcegm regimes raises a clear `ValueError`.""" - grid_config = dataclasses.replace(BENCHMARK_GRID_CONFIG, n_savings_gridpoints=1) + grid_config = dataclasses.replace(_EGM_GRID_CONFIG, n_savings_gridpoints=1) with pytest.raises(ValueError, match="n_savings_gridpoints"): build_all_regimes( grid_config=grid_config, diff --git a/tests/test_dcegm_parity.py b/tests/test_dcegm_parity.py index 49e1695..70bfe05 100644 --- a/tests/test_dcegm_parity.py +++ b/tests/test_dcegm_parity.py @@ -114,7 +114,6 @@ def _make_model(*, solver: SolverName, grid_config: GridConfig) -> Model: else None ) return create_model( - n_subjects=N_SUBJECTS, fixed_params=fixed_params, wage_params=wage_params, derived_categoricals=_DERIVED_CATEGORICALS, @@ -134,7 +133,6 @@ def _seeded_panel(model: Model) -> pd.DataFrame: result = model.simulate( params=params, initial_conditions=initial_conditions, - period_to_regime_to_V_arr=None, log_level="off", ) df = result.to_dataframe() @@ -235,7 +233,7 @@ def test_dcegm_solves_the_parity_model() -> None: (regime, period) cell carries a finite value-function array.""" model = _make_model(solver="dcegm", grid_config=PARITY_GRID_CONFIG) _, _, params = get_benchmark_params(model=model) - period_to_regime_to_v = model.solve(params=params, log_level="off") - for regime_to_v in period_to_regime_to_v.values(): + solution = model.solve(params=params, log_level="off") + for regime_to_v in solution.values.values(): for v_arr in regime_to_v.values(): assert bool(np.isfinite(np.asarray(v_arr)).any()) diff --git a/tests/test_di_medicare_transitions.py b/tests/test_di_medicare_transitions.py new file mode 100644 index 0000000..7c88074 --- /dev/null +++ b/tests/test_di_medicare_transitions.py @@ -0,0 +1,230 @@ +"""Disability Medicare before 65: disabled and not working last year. + +A pre-65 household is on disability Medicare next period exactly when it does +not work this period and is disabled next period. The regime transition splits +the surviving mass by next-period disability; each target's health law is the +health distribution conditional on landing there. +""" + +import jax.numpy as jnp +import numpy as np +import pytest + +from aca_model.agent import health +from aca_model.agent.health import HealthWithDisability +from aca_model.agent.labor_market import LaborSupply +from aca_model.baseline.regimes._common import ( + REGIME_SPECS, + RegimeId, + _build_per_target_regime_health, + build_model_edges, + make_targets, +) +from aca_model.baseline.regimes._nongroup import ( + _make_transition_canwork as nongroup_canwork, +) +from aca_model.baseline.regimes._retiree import ( + _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 + +N_PERIODS = MODEL_CONFIG.end_age - MODEL_CONFIG.start_age +SURVIVAL = jnp.full((N_PERIODS, 3), 0.99) +# Rows: current health (disabled, bad, good); columns: next health. +_HEALTH_ROWS = jnp.array([[0.98, 0.01, 0.01], [0.05, 0.75, 0.20], [0.002, 0.058, 0.94]]) +HEALTH_TRANS = jnp.broadcast_to(_HEALTH_ROWS, (N_PERIODS, 3, 3)) +# P(disabled next period | health), the first column of the rows above. +PROB_DISABLED = HEALTH_TRANS[:, :, 0] + + +def _nongroup_probs(regime: str, health_state: int, labor: int) -> jnp.ndarray: + own, _ = make_targets(regime) + return nongroup_canwork(own=own)( + age=jnp.int32(55), + period=jnp.int32(4), + health=jnp.int32(health_state), + labor_supply=jnp.array(labor), + survival_probs=SURVIVAL, + prob_disabled_next=PROB_DISABLED, + ) + + +@pytest.mark.parametrize( + ("regime", "health_state", "labor", "target", "expected"), + [ + # Bad health, not working: disabled next year with prob 0.05. + ( + "nongroup_nomc_inelig_canwork", + HealthWithDisability.bad, + LaborSupply.do_not_work, + "nongroup_dimc_inelig_canwork", + 0.99 * 0.05, + ), + ( + "nongroup_nomc_inelig_canwork", + HealthWithDisability.bad, + LaborSupply.do_not_work, + "nongroup_nomc_inelig_canwork", + 0.99 * 0.95, + ), + # Working: no disability Medicare next year. + ( + "nongroup_nomc_inelig_canwork", + HealthWithDisability.bad, + LaborSupply.h2000, + "nongroup_dimc_inelig_canwork", + 0.0, + ), + # On disability Medicare and recovering: leaves it. + ( + "nongroup_dimc_inelig_canwork", + HealthWithDisability.disabled, + LaborSupply.do_not_work, + "nongroup_nomc_inelig_canwork", + 0.99 * 0.02, + ), + ( + "nongroup_dimc_inelig_canwork", + HealthWithDisability.disabled, + LaborSupply.do_not_work, + "nongroup_dimc_inelig_canwork", + 0.99 * 0.98, + ), + # On disability Medicare and working: leaves it. + ( + "nongroup_dimc_inelig_canwork", + HealthWithDisability.disabled, + LaborSupply.h2000, + "nongroup_nomc_inelig_canwork", + 0.99, + ), + ], +) +def test_nongroup_transition_di_medicare_entry_and_exit( + regime: str, health_state: int, labor: int, target: str, expected: float +) -> None: + """Probability of each target splits survival by next-period disability.""" + probs = _nongroup_probs(regime, health_state, labor) + np.testing.assert_allclose(probs[getattr(RegimeId, target)], expected, atol=1e-12) + + +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)( + age=jnp.int32(55), + period=jnp.int32(4), + health=jnp.int32(HealthWithDisability.good), + labor_supply=jnp.array(LaborSupply.do_not_work), + is_ssi_eligible=jnp.array(False), + survival_probs=SURVIVAL, + prob_disabled_next=PROB_DISABLED, + ) + np.testing.assert_allclose( + probs[RegimeId.retiree_dimc_inelig_canwork], 0.99 * 0.002, atol=1e-12 + ) + + +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)( + age=jnp.int32(55), + period=jnp.int32(4), + health=jnp.int32(HealthWithDisability.bad), + labor_supply=jnp.array(LaborSupply.do_not_work), + is_ssi_eligible=jnp.array(False), + survival_probs=SURVIVAL, + prob_disabled_next=PROB_DISABLED, + ) + np.testing.assert_allclose( + probs[RegimeId.nongroup_dimc_inelig_canwork], 0.99 * 0.05, atol=1e-12 + ) + + +def test_next_health_into_dimc_is_disabled() -> None: + """Landing on disability Medicare means being disabled.""" + result = health.next_health_into_dimc( + health=jnp.int32(HealthWithDisability.bad), + period=jnp.int32(4), + health_trans_probs=HEALTH_TRANS, + ) + np.testing.assert_allclose(result, [1.0, 0.0, 0.0]) + + +@pytest.mark.parametrize( + ("labor", "expected"), + [ + (LaborSupply.do_not_work, [0.0, 0.75 / 0.95, 0.20 / 0.95]), + (LaborSupply.h2000, [0.05, 0.75, 0.20]), + ], +) +def test_next_health_into_nomc_excludes_disability_for_non_workers( + labor: int, expected: list[float] +) -> None: + """A non-worker without disability Medicare next year is not disabled then.""" + result = health.next_health_into_nomc( + health=jnp.int32(HealthWithDisability.bad), + period=jnp.int32(4), + labor_supply=jnp.array(labor), + health_trans_probs=HEALTH_TRANS, + ) + np.testing.assert_allclose(result, expected, atol=1e-12) + + +@pytest.mark.parametrize( + ("source", "target", "law"), + [ + ( + "nongroup_nomc_inelig_canwork", + "nongroup_dimc_inelig_canwork", + "next_health_into_dimc", + ), + ( + "nongroup_nomc_inelig_canwork", + "nongroup_nomc_inelig_canwork", + "next_health_into_nomc", + ), + ( + "retiree_dimc_choose_canwork", + "retiree_nomc_choose_canwork", + "next_health_into_nomc", + ), + ( + "retiree_dimc_choose_canwork", + "retiree_oamc_choose_canwork", + "next_health_cross", + ), + ( + "nongroup_oamc_choose_canwork", + "nongroup_oamc_forced_canwork", + "next_health", + ), + ], +) +def test_per_target_health_law(source: str, target: str, law: str) -> None: + """Pre-65 targets condition next health on the Medicare outcome.""" + laws = _build_per_target_regime_health(REGIME_SPECS[source]) + assert laws[target].func is getattr(health, law) + + +@pytest.mark.parametrize( + ("regime", "dimc_target"), + [ + ( + "nongroup_nomc_inelig_canwork", + "nongroup_dimc_inelig_canwork", + ), + ( + "retiree_nomc_inelig_canwork", + "retiree_dimc_inelig_canwork", + ), + ("tied_nomc_inelig_canwork", "nongroup_dimc_inelig_canwork"), + ], +) +def test_graph_targets_include_disability_medicare_entry( + regime: str, dimc_target: str +) -> None: + """A pre-65 regime without Medicare declares its disability-Medicare target.""" + assert 55 in build_model_edges()[regime][dimc_target] diff --git a/tests/test_execution_budget.py b/tests/test_execution_budget.py new file mode 100644 index 0000000..0dafef2 --- /dev/null +++ b/tests/test_execution_budget.py @@ -0,0 +1,62 @@ +"""Allocator observations determine the execution ceiling on selected devices.""" + +from dataclasses import dataclass + +import jax +import pytest + +from aca_model import execution + + +@dataclass +class _Device: + id: int + platform: str + stats: dict[str, int] | None + + def memory_stats(self): + return self.stats + + +def test_budget_uses_smallest_selected_allocator_limit(monkeypatch): + """Only selected GPU pools bound execution; current usage is not subtracted.""" + devices = [ + _Device(0, "gpu", {"bytes_limit": 1000, "bytes_in_use": 200}), + _Device(1, "gpu", {"bytes_limit": 800, "bytes_in_use": 100}), + _Device(2, "gpu", None), + ] + monkeypatch.setattr(jax, "devices", lambda: devices) + + config = execution.execution_config_for_devices(devices=(0, 1)) + + assert config.devices == (0, 1) + assert config.device_memory_bytes == 800 + assert config.axis_widths == {} + assert config.sharded_states == () + + +@pytest.mark.parametrize("stats", [None, {}, {"bytes_limit": 0}]) +def test_missing_gpu_limit_is_refused(monkeypatch, stats): + """An unknown selected GPU ceiling must not silently select bootstrap widths.""" + monkeypatch.setattr(jax, "devices", lambda: [_Device(0, "gpu", stats)]) + + with pytest.raises(ValueError, match=r"allocator.*limit.*device 0"): + execution.execution_config_for_devices() + + +def test_cpu_construction_does_not_invent_an_allocator_budget(monkeypatch): + """CPU devices retain an unspecified byte budget when no pool limit exists.""" + monkeypatch.setattr(jax, "devices", lambda: [_Device(0, "cpu", None)]) + + config = execution.execution_config_for_devices() + + assert config.devices == (0,) + assert config.device_memory_bytes is None + + +def test_selected_device_must_be_visible(monkeypatch): + """An unavailable selected device is rejected before reading allocator stats.""" + monkeypatch.setattr(jax, "devices", lambda: [_Device(0, "gpu", None)]) + + with pytest.raises(ValueError, match="Selected devices are not visible"): + execution.execution_config_for_devices(devices=(1,)) diff --git a/tests/test_execution_config.py b/tests/test_execution_config.py new file mode 100644 index 0000000..79a47d6 --- /dev/null +++ b/tests/test_execution_config.py @@ -0,0 +1,124 @@ +"""Model factories expose pylcm's execution policy without changing economic grids.""" + +from functools import partial +from types import SimpleNamespace + +import jax +import pytest +from helpers.model import aca_fixed_params +from lcm import DiscreteGrid, ExecutionConfig +from lcm.exceptions import ExecutionPlanningError + +from aca_model.aca import model as aca_model_module +from aca_model.aca.health_insurance import PolicyVariant +from aca_model.aca.model import create_model as create_aca_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 import model as baseline_model_module +from aca_model.baseline.health_insurance import HealthInsuranceState +from aca_model.baseline.model import create_model +from aca_model.benchmark import create_benchmark_model, get_benchmark_params +from aca_model.config import BENCHMARK_GRID_CONFIG + + +def _factory(kind): + if kind == "benchmark": + return partial( + create_benchmark_model, + pref_type_grid=DiscreteGrid(BenchmarkPrefType), + ) + fixed_params, wage_params, _ = get_benchmark_params(model=None) + factory = create_model + if kind == "aca": + factory = partial(create_aca_model, policy=PolicyVariant.ACA) + fixed_params = aca_fixed_params( + fixed_params=fixed_params, policy=PolicyVariant.ACA + ) + return partial( + factory, + fixed_params=fixed_params, + wage_params=wage_params, + derived_categoricals={ + "good_health": DiscreteGrid(GoodHealth), + "is_married": DiscreteGrid(IsMarried), + "his": DiscreteGrid(HealthInsuranceState), + "target_his": DiscreteGrid(HealthInsuranceState), + "pref_type": DiscreteGrid(BenchmarkPrefType), + }, + grid_config=BENCHMARK_GRID_CONFIG, + pref_type_grid=DiscreteGrid(BenchmarkPrefType), + ) + + +@pytest.mark.parametrize("kind", ["baseline", "aca", "benchmark"]) +def test_factory_preserves_explicit_device_selection(kind): + """A factory uses exactly the selected device and preserves the model grids.""" + devices = (jax.devices()[-1].id,) + model = _factory(kind)(execution_config=ExecutionConfig(devices=devices)) + + assert model.execution_devices == devices + assert len(model.user_regimes) == 19 + regime = model.user_regimes["retiree_nomc_inelig_canwork"] + assert regime.states["assets"].n_points == 3 + assert regime.states["aime"].n_points == 38 + assert regime.actions["consumption_dollars"].n_points == 5 + + +@pytest.mark.parametrize("kind", ["baseline", "aca", "benchmark"]) +def test_factory_preserves_axis_width_validation(kind): + """A width for an undeclared program axis is refused by the model.""" + config = ExecutionConfig(axis_widths={"aca_undeclared_axis": 1}) + with pytest.raises(ExecutionPlanningError, match="aca_undeclared_axis"): + _factory(kind)(execution_config=config) + + +@pytest.mark.parametrize("kind", ["baseline", "aca", "benchmark"]) +@pytest.mark.parametrize( + ("requested_policy", "expected_budget"), + [ + (None, 800), + ( + ExecutionConfig( + devices=(0,), + device_memory_bytes=400, + sharded_states=("pref_type",), + axis_widths={"cell": 32}, + ), + 400, + ), + (ExecutionConfig(devices=(0,)), ExecutionConfig().device_memory_bytes), + ], +) +def test_factory_supplies_the_requested_or_measured_budget( + monkeypatch, kind, requested_policy, expected_budget +): + """Defaults use measured limits; explicit budgets and unbudgeted controls survive.""" + factory = _factory(kind) + device = SimpleNamespace( + id=0, platform="gpu", memory_stats=lambda: {"bytes_limit": 800} + ) + monkeypatch.setattr(jax, "devices", lambda: [device]) + + class ConstructionObservedError(Exception): + pass + + def observe_pylcm_constructor(**kwargs): + policy = kwargs["execution_config"] + assert policy.device_memory_bytes == expected_budget + assert policy.devices == (0,) + if requested_policy is not None: + assert policy == requested_policy + raise ConstructionObservedError + + monkeypatch.setattr(baseline_model_module, "Model", observe_pylcm_constructor) + monkeypatch.setattr(aca_model_module, "Model", observe_pylcm_constructor) + with pytest.raises(ConstructionObservedError): + factory(execution_config=requested_policy) + + +@pytest.mark.parametrize("kind", ["baseline", "aca", "benchmark"]) +def test_factory_rejects_construction_subject_count(kind): + """Population size is supplied by initial conditions, not model construction.""" + with pytest.raises(TypeError, match="n_subjects"): + _factory(kind)(n_subjects=1) diff --git a/tests/test_grid_config.py b/tests/test_grid_config.py new file mode 100644 index 0000000..7eb6564 --- /dev/null +++ b/tests/test_grid_config.py @@ -0,0 +1,35 @@ +"""Economic-grid configuration rejects hardware execution controls.""" + +import pytest + +from aca_model.config import GridConfig + + +@pytest.mark.parametrize( + "field", + [ + "n_assets_batch_size", + "n_aime_batch_size", + "pref_type_distributed", + "spousal_income_distributed", + "n_health_batch_size", + "n_spousal_income_batch_size", + "n_lagged_labor_supply_batch_size", + "n_claimed_ss_batch_size", + "n_pref_type_batch_size", + "n_wage_res_batch_size", + "n_savings_batch_size", + "n_stochastic_node_batch_size", + "n_nbegm_stochastic_node_batch_size", + "n_nbegm_envelope_segment_block_size", + "n_nbegm_interval_batch_size", + "n_nbegm_max_device_workspace_bytes", + "n_nbegm_cell_block_size", + "n_nbegm_branch_batch_size", + ], +) +def test_grid_config_refuses_execution_controls(field): + """A misplaced execution control fails visibly instead of being ignored.""" + value = True if field.endswith("distributed") else 1 + with pytest.raises(TypeError, match=field): + GridConfig(**{field: value}) # ty: ignore[invalid-argument-type] diff --git a/tests/test_health_insurance.py b/tests/test_health_insurance.py index a3546de..15d379d 100644 --- a/tests/test_health_insurance.py +++ b/tests/test_health_insurance.py @@ -3,7 +3,6 @@ import jax.numpy as jnp import pytest -from aca_model.agent.labor_market import LaborSupply from aca_model.baseline import health_insurance from aca_model.baseline.health_insurance import BuyPrivate @@ -103,11 +102,14 @@ def test_ssi_benefit_not_eligible() -> None: # --- predicted_hcc_insurer (DAG interpolation function) --- GRID = jnp.array([-1.0, 0.0, 1.0]) -TABLE = jnp.array([100.0, 200.0, 400.0]) +TABLE = jnp.zeros((1, 2, 2, 3)).at[0, 0, 1].set(jnp.array([100.0, 200.0, 400.0])) def test_predicted_hcc_insurer_on_grid() -> None: result = health_insurance.hcc_insurer_predicted( + period=jnp.int32(0), + is_married=jnp.int32(0), + good_health=jnp.int32(1), hcc_persistent=jnp.array(0.0), predicted_hcc_insurer_table=TABLE, hcc_persistent_grid=GRID, @@ -117,6 +119,9 @@ def test_predicted_hcc_insurer_on_grid() -> None: def test_predicted_hcc_insurer_off_grid() -> None: result = health_insurance.hcc_insurer_predicted( + period=jnp.int32(0), + is_married=jnp.int32(0), + good_health=jnp.int32(1), hcc_persistent=jnp.array(-0.5), predicted_hcc_insurer_table=TABLE, hcc_persistent_grid=GRID, @@ -163,42 +168,6 @@ def test_compute_table_uniform_transition(table_inputs: dict) -> None: assert jnp.allclose(result, result[0], atol=ATOL) -# --- premium: buy_private conditioning --- - - -_PREMIUM_KWARGS: dict = { - "age": jnp.int32(60), - "good_health": jnp.int32(1), - "is_married": jnp.int32(0), - "labor_supply": jnp.array(LaborSupply.h2000), - "premium_intercept": jnp.asarray(1000.0), - "premium_age": jnp.asarray(0.0), - "premium_age_sq": jnp.asarray(0.0), - "premium_age_cub": jnp.asarray(0.0), - "premium_predicted_hcc": jnp.asarray(0.0), - "premium_good_health": jnp.asarray(0.0), - "premium_married": jnp.asarray(0.0), - "premium_works": jnp.asarray(0.0), - "premium_married_works": jnp.asarray(0.0), - "premium_minimum": jnp.asarray(500.0), - "predicted_hcc_insurer": jnp.array(0.0), -} - - -def test_premium_insured_positive() -> None: - result = health_insurance.premium( - buy_private=jnp.array(BuyPrivate.yes), **_PREMIUM_KWARGS - ) - assert result > 0.0 - - -def test_premium_uninsured_zero() -> None: - result = health_insurance.premium( - buy_private=jnp.array(BuyPrivate.no), **_PREMIUM_KWARGS - ) - assert jnp.isclose(result, 0.0) - - # --- primary_oop: buy_private conditioning --- @@ -206,6 +175,7 @@ def test_primary_oop_insured_applies_deductible_coinsurance() -> None: result = health_insurance.primary_oop( total_health_costs=jnp.array(10000.0), buy_private=jnp.array(BuyPrivate.yes), + is_married=jnp.int32(0), deductible=jnp.asarray(500.0), coinsurance_rate=jnp.asarray(0.2), oop_max=jnp.asarray(5000.0), @@ -219,6 +189,7 @@ def test_primary_oop_uninsured_equals_total_costs() -> None: result = health_insurance.primary_oop( total_health_costs=total, buy_private=jnp.array(BuyPrivate.no), + is_married=jnp.int32(0), deductible=jnp.asarray(500.0), coinsurance_rate=jnp.asarray(0.2), oop_max=jnp.asarray(5000.0), diff --git a/tests/test_initial_conditions_extreme_assets.py b/tests/test_initial_conditions_extreme_assets.py index 407b1cd..6455bee 100644 --- a/tests/test_initial_conditions_extreme_assets.py +++ b/tests/test_initial_conditions_extreme_assets.py @@ -8,8 +8,8 @@ """ import jax.numpy as jnp +from _lcm.simulation.initial_conditions import validate_initial_conditions from lcm import DiscreteGrid -from lcm.model import validate_initial_conditions from aca_model.agent.assets_and_income import borrowing_constraint from aca_model.agent.preferences import BenchmarkPrefType @@ -101,7 +101,6 @@ def test_extreme_negative_assets_subject_passes_validation() -> None: """ n_subjects = 1 model = create_benchmark_model( - n_subjects=n_subjects, pref_type_grid=DiscreteGrid(BenchmarkPrefType), ) _, _, params = get_benchmark_params(model=model) diff --git a/tests/test_leisure_floor.py b/tests/test_leisure_floor.py new file mode 100644 index 0000000..2c1e470 --- /dev/null +++ b/tests/test_leisure_floor.py @@ -0,0 +1,190 @@ +"""Float32 felicity at the leisure floor, against a 60-digit reference. + +The class covers the canwork cells at production preference values: every pref +type, health, lagged labor supply, hours choice and canwork age, at consumption +points spanning the NB-EGM numeric-inverse bracket (`1e-8` to the top of the +action range). Marginal utility is `jax.grad` of the consumption-dollar +felicity, the route pylcm's NB-EGM uses. +""" + +import itertools +from decimal import Decimal, getcontext + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +from aca_model.agent import preferences + +getcontext().prec = 60 + +CONSUMPTION_WEIGHTS = (0.6776541629520845, 0.8805772328591686, 0.0718086283445225) +COEFFICIENTS_RRA = (3.841252231680976, 0.9990771146810682, 3.8328505891095137) +TIME_ENDOWMENT = 3926.9478390365557 +BAD_HEALTH_COST = 408.9190313043897 +REENTRY_COST = 119.61095526191421 +FIXED_COST_INTERCEPT = 337.5223349543413 +FIXED_COST_AGE_TREND = 84.09242960636178 +REFERENCE_AGE = 50 +REFERENCE_HOURS = 1000.0 +AVERAGE_CONSUMPTION = 20000.0 +HOURS = (0.0, 1000.0, 1500.0, 2000.0, 2500.0) +CANWORK_AGES = tuple(range(51, 72)) +CONSUMPTION = (1e-8, 1597.0921419521899, 20000.0, 300_000.0, 721271853.846792) +FLOAT32_REL_TOL = 1e-4 +# Leisure-map widths as fractions of the endowment: smoothing width and floor. +SMOOTHING_FRACTION = 0.01 +FLOOR_FRACTION = 1e-5 + + +def _f32(x: float) -> jnp.ndarray: + return jnp.asarray(x, dtype=jnp.float32) + + +def _felicity_of_dollars(consumption, leisure, weight, rra, scale): + return preferences.u_alive( + consumption_equiv=preferences.consumption_equiv( + consumption_dollars=consumption, + equivalence_scale=jnp.ones_like(consumption), + ), + leisure=leisure, + consumption_weight=weight, + coefficient_rra=rra, + utility_scale_factor=scale, + ) + + +def _leisure_available(good: int, lagged: int, hours: float, age: int) -> float: + fixed_cost = FIXED_COST_INTERCEPT + FIXED_COST_AGE_TREND * (age - REFERENCE_AGE) + reentry = REENTRY_COST if lagged == 0 else 0.0 + work = hours + fixed_cost + reentry if hours > 0 else 0.0 + return TIME_ENDOWMENT - (0.0 if good else BAD_HEALTH_COST) - work + + +def _cells() -> list[tuple[float, float]]: + """All `(leisure_available, consumption)` pairs of the class.""" + return [ + (_leisure_available(good, lagged, hours, age), c) + for good, lagged, hours, age, c in itertools.product( + (0, 1), (0, 1), HOURS, CANWORK_AGES, CONSUMPTION + ) + ] + + +def _float32_u_and_marginal( + pref_type: int, cells: list[tuple[float, float]] +) -> tuple[np.ndarray, np.ndarray]: + available = jnp.asarray([a for a, _ in cells], dtype=jnp.float32) + consumption = jnp.asarray([c for _, c in cells], dtype=jnp.float32) + weight = _f32(CONSUMPTION_WEIGHTS[pref_type]) + rra = _f32(COEFFICIENTS_RRA[pref_type]) + leisure = _floored_leisure(available, _f32(TIME_ENDOWMENT)) + scale = preferences.utility_scale_factor( + average_consumption_equiv=_f32(AVERAGE_CONSUMPTION), + consumption_weight=weight, + coefficient_rra=rra, + time_endowment=_f32(TIME_ENDOWMENT), + fixed_cost_of_work_intercept=_f32(FIXED_COST_INTERCEPT), + reference_hours=_f32(REFERENCE_HOURS), + ) + args = (weight, rra, scale) + u = jax.vmap(_felicity_of_dollars, in_axes=(0, 0, None, None, None))( + consumption, leisure, *args + ) + marginal = jax.vmap( + jax.grad(_felicity_of_dollars), in_axes=(0, 0, None, None, None) + )(consumption, leisure, *args) + assert u.dtype == marginal.dtype == jnp.float32 + return np.asarray(u, dtype=np.float64), np.asarray(marginal, dtype=np.float64) + + +def _floored_leisure(available, endowment): + """Production leisure at a given `leisure_available`, via the tied-regime map. + + Good health and no fixed cost make the work loss the only deduction, so hours of + `endowment - available` leave exactly `available`. + """ + return preferences.leisure_canwork_tied( + working_hours_value=endowment - available, + good_health=jnp.ones(available.shape, dtype=jnp.int32), + time_endowment=endowment, + leisure_cost_of_bad_health=jnp.zeros_like(endowment), + fixed_cost_of_work=jnp.zeros_like(endowment), + ) + + +def _d(x: float) -> Decimal: + return Decimal(repr(float(x))) + + +def _reference_u_and_marginal( + pref_type: int, cells: list[tuple[float, float]] +) -> tuple[np.ndarray, np.ndarray]: + """Closed forms in 60-digit log space, independent of the JAX code path.""" + weight, rra = _d(CONSUMPTION_WEIGHTS[pref_type]), _d(COEFFICIENTS_RRA[pref_type]) + endowment = _d(TIME_ENDOWMENT) + smoothing = _d(SMOOTHING_FRACTION) * endowment + floor = _d(FLOOR_FRACTION) * endowment + average_leisure = endowment - _d(REFERENCE_HOURS) - _d(FIXED_COST_INTERCEPT) + log_average = weight * _d(AVERAGE_CONSUMPTION).ln() + (1 - weight) * ( + average_leisure.ln() + ) + scale = abs((1 - rra) / ((1 - rra) * log_average).exp()) + u, marginal = [], [] + for available, c in cells: + leisure = ( + smoothing + * ((floor / smoothing).exp() + (_d(available) / smoothing).exp()).ln() + ) + log_composite = weight * _d(c).ln() + (1 - weight) * leisure.ln() + power = ((1 - rra) * log_composite).exp() + u.append(float(scale * power / (1 - rra))) + marginal.append(float(scale * weight / _d(c) * power)) + return np.array(u), np.array(marginal) + + +OUTPUTS = {"u": 0, "marginal": 1} +WITNESS = [(_leisure_available(0, 0, 2500.0, 71), c) for c in CONSUMPTION] + + +@pytest.mark.parametrize("output", ["u", "marginal"]) +def test_float32_felicity_at_witness_cell_matches_reference(output: str) -> None: + """Pref type 2, bad health, did not work, 2500 h, age 71: finite and accurate.""" + got = _float32_u_and_marginal(2, WITNESS)[OUTPUTS[output]] + ref = _reference_u_and_marginal(2, WITNESS)[OUTPUTS[output]] + np.testing.assert_allclose(got, ref, rtol=FLOAT32_REL_TOL) + + +@pytest.mark.parametrize("output", ["u", "marginal"]) +@pytest.mark.parametrize("pref_type", [0, 1, 2]) +def test_float32_felicity_over_canwork_class_matches_reference( + pref_type: int, output: str +) -> None: + """Every canwork cell and consumption point: float32 equals the reference.""" + cells = _cells() + got = _float32_u_and_marginal(pref_type, cells)[OUTPUTS[output]] + ref = _reference_u_and_marginal(pref_type, cells)[OUTPUTS[output]] + np.testing.assert_allclose(got, ref, rtol=FLOAT32_REL_TOL) + + +@pytest.mark.parametrize("pref_type", [0, 1, 2]) +def test_leisure_floor_leaves_utility_unchanged_away_from_the_floor( + pref_type: int, +) -> None: + """With leisure_available above 10% of the endowment, u moves by at most 1e-7. + + Compared against the unfloored softplus `s * log(1 + e^(x/s))`. + """ + available = np.array([a for a, _ in _cells() if a > 0.1 * TIME_ENDOWMENT]) + assert available.size > 0 + smoothing = SMOOTHING_FRACTION * TIME_ENDOWMENT + unfloored = smoothing * np.logaddexp(0.0, available / smoothing) + floored = np.asarray( + _floored_leisure(jnp.asarray(available), jnp.asarray(TIME_ENDOWMENT)) + ) + exponent = (1.0 - CONSUMPTION_WEIGHTS[pref_type]) * ( + 1.0 - COEFFICIENTS_RRA[pref_type] + ) + relative_change = np.abs((floored / unfloored) ** exponent - 1.0) + np.testing.assert_array_less(relative_change, 1e-7) diff --git a/tests/test_medicaid_nongroup.py b/tests/test_medicaid_nongroup.py new file mode 100644 index 0000000..b8aa0af --- /dev/null +++ b/tests/test_medicaid_nongroup.py @@ -0,0 +1,98 @@ +"""Medicaid in the non-group regimes: last-resort payer and Medicare premia. + +The tests evaluate the real non-group regime DAG with the upstream cost and +premium nodes pinned to constants, so they check how Medicaid eligibility +routes those quantities rather than how the quantities themselves are built. + +- A Medicaid-eligible non-group household before Medicare is covered by + Medicaid alone: it can neither hold private cover nor stay uninsured, so + its OOP is Medicaid's cost-sharing on total costs and it pays no premium + whatever its `buy_private` choice. +- Medicaid pays the Medicare premium of a Medicaid-eligible household. +""" + +import inspect + +import jax.numpy as jnp +import pytest +from dags import concatenate_functions + +from aca_model.agent.labor_market import LaborSupply +from aca_model.baseline import health_insurance +from aca_model.baseline.health_insurance import BuyPrivate +from aca_model.baseline.regimes._common import REGIME_SPECS, build_model_functions +from aca_model.baseline.regimes._nongroup import _build_functions + +TOTAL_COSTS = 20000.0 +PRIVATE_OOP = 5285.0 +PREMIUM_BEFORE_MEDICAID = 4620.0 +MEDICAID_COINSURANCE = 0.0352 + +_MEDICAID_INPUTS = { + "deductible_medicaid": jnp.asarray(0.0), + "coinsurance_rate_medicaid": jnp.asarray(MEDICAID_COINSURANCE), + "oop_max_medicaid": jnp.asarray(jnp.inf), +} + + +def _is_base_premium(func: object) -> bool: + """Whether `func` is one of the baseline premium formulas.""" + return getattr(func, "__module__", "") == health_insurance.__name__ and getattr( + func, "__name__", "" + ).startswith("premium") + + +def _evaluate_nongroup(regime_name: str, buy_private: int) -> dict: + """Evaluate `oop_costs` and `hic_premium` of a Medicaid-eligible household.""" + functions = { + **build_model_functions(), + **_build_functions(REGIME_SPECS[regime_name]), + } + functions = { + name: (lambda: jnp.asarray(PREMIUM_BEFORE_MEDICAID)) + if _is_base_premium(func) + else func + for name, func in functions.items() + if callable(func) + } + functions["total_health_costs"] = lambda: jnp.asarray(TOTAL_COSTS) + functions["primary_oop"] = lambda: jnp.asarray(PRIVATE_OOP) + functions["is_medicaid_eligible"] = lambda: jnp.asarray(True) + combined = concatenate_functions( + functions, + targets=["oop_costs", "hic_premium"], + return_type="dict", + ) + inputs = { + **_MEDICAID_INPUTS, + "buy_private": jnp.asarray(buy_private), + "labor_supply": jnp.asarray(LaborSupply.do_not_work), + } + needed = inspect.signature(combined).parameters + return combined(**{k: v for k, v in inputs.items() if k in needed}) + + +def test_medicaid_eligible_buyer_pays_medicaid_cost_sharing_on_total_costs() -> None: + """With `buy_private=yes`, OOP is Medicaid's 3.52% of total costs: 704.""" + result = _evaluate_nongroup("nongroup_nomc_inelig_canwork", BuyPrivate.yes) + assert jnp.isclose(result["oop_costs"], MEDICAID_COINSURANCE * TOTAL_COSTS) + + +def test_medicaid_eligible_buyer_pays_no_private_premium() -> None: + """With `buy_private=yes`, a Medicaid-eligible household pays no premium.""" + result = _evaluate_nongroup("nongroup_nomc_inelig_canwork", BuyPrivate.yes) + assert jnp.isclose(result["hic_premium"], 0.0) + + +@pytest.mark.parametrize( + "regime_name", + [ + "nongroup_dimc_inelig_canwork", + "nongroup_oamc_choose_canwork", + "nongroup_oamc_forced_forcedout", + ], +) +def test_medicaid_pays_the_medicare_premium(regime_name: str) -> None: + """A Medicaid-eligible household with Medicare pays no Medicare premium.""" + result = _evaluate_nongroup(regime_name, BuyPrivate.no) + assert jnp.isclose(result["hic_premium"], 0.0) diff --git a/tests/test_medical_cost_timing.py b/tests/test_medical_cost_timing.py new file mode 100644 index 0000000..2ebbc8d --- /dev/null +++ b/tests/test_medical_cost_timing.py @@ -0,0 +1,86 @@ +"""Medical costs are realised after the period's choices. + +The state a household carries into period t holds last year's persistent +shock and last year's transitory shock. This year's shocks are drawn after +consumption, work, claiming and insurance are chosen; the resulting +out-of-pocket bill is taken off end-of-period assets. So the value and every +choice in period t depend on the lagged persistent shock, which predicts this +year's costs, and not at all on the lagged transitory shock. +""" + +from dataclasses import replace + +import jax.numpy as jnp +import numpy as np +import pandas as pd +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, + get_benchmark_params, +) +from aca_model.config import BENCHMARK_GRID_CONFIG + +_SHOCKS = ("hcc_persistent", "hcc_transitory") + + +@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 + ) + model = create_benchmark_model( + pref_type_grid=DiscreteGrid(BenchmarkPrefType), grid_config=grid_config + ) + _, _, params = get_benchmark_params(model=model) + base = get_benchmark_initial_conditions( + model=model, n_subjects=1, seed=0, grid_config=grid_config + ) + grids = model.user_regimes["retiree_nomc_inelig_canwork"].states + initial_conditions = {name: jnp.tile(values, 4) for name, values in base.items()} + for row, shock in enumerate(_SHOCKS): + points = np.asarray(grids[shock].to_jax()) # ty: ignore[unresolved-attribute] + values = np.asarray(initial_conditions[shock]).copy() + values[2 * row : 2 * row + 2] = (points.min(), points.max()) + initial_conditions[shock] = jnp.asarray(values) + result = model.simulate( + params=params, initial_conditions=initial_conditions, log_level="off" + ) + df = result.to_dataframe(terminal_rows="all") + period_0 = df.loc[df["period"] == 0].sort_values("subject_id") + return { + (shock, end): period_0.iloc[2 * row + i] + for row, shock in enumerate(_SHOCKS) + for i, end in enumerate(("low", "high")) + } + + +def test_value_does_not_depend_on_lagged_transitory_shock( + period_0_by_shock: dict[tuple[str, str], pd.Series], +) -> None: + """Twins differing only in last year's transitory shock have equal value.""" + low = period_0_by_shock["hcc_transitory", "low"]["value"] + high = period_0_by_shock["hcc_transitory", "high"]["value"] + assert low == high + + +def test_consumption_does_not_depend_on_lagged_transitory_shock( + period_0_by_shock: dict[tuple[str, str], pd.Series], +) -> None: + """Twins differing only in last year's transitory shock consume the same.""" + low = period_0_by_shock["hcc_transitory", "low"]["consumption_dollars"] + high = period_0_by_shock["hcc_transitory", "high"]["consumption_dollars"] + assert low == high + + +def test_value_falls_with_lagged_persistent_shock( + period_0_by_shock: dict[tuple[str, str], pd.Series], +) -> None: + """A higher lagged persistent shock predicts higher costs and lowers value.""" + low = period_0_by_shock["hcc_persistent", "low"]["value"] + high = period_0_by_shock["hcc_persistent", "high"]["value"] + assert high < low diff --git a/tests/test_model_components.py b/tests/test_model_components.py index 5caa981..6f18dc5 100644 --- a/tests/test_model_components.py +++ b/tests/test_model_components.py @@ -1,6 +1,7 @@ """Tests for individual model component functions.""" import jax.numpy as jnp +from helpers.social_security import next_aime_from_state from aca_model.agent import preferences from aca_model.environment import social_security @@ -191,7 +192,7 @@ def test_aime_to_pia_three_brackets() -> None: def test_next_aime_accrual() -> None: ratio = jnp.zeros(100) ratio = ratio.at[55].set(0.5) - result = social_security.next_aime( + result = next_aime_from_state( aime=jnp.array(1000.0), labor_income=jnp.array(50000.0), period=jnp.int32(55), diff --git a/tests/test_model_creation.py b/tests/test_model_creation.py index 5b2415c..9940ffd 100644 --- a/tests/test_model_creation.py +++ b/tests/test_model_creation.py @@ -2,10 +2,9 @@ import inspect from collections.abc import Mapping -from dataclasses import replace import pytest -from helpers.model import ( # ty: ignore[unresolved-import] +from helpers.model import ( make_aca_model, make_baseline_model, ) @@ -52,26 +51,26 @@ def build_regime(name: str): def test_model_creates_successfully() -> None: - model = make_baseline_model(n_subjects=1) + model = make_baseline_model() assert len(model.user_regimes) == 19 - assert model.n_periods == 45 + assert model.n_periods == 46 def test_model_age_range() -> None: - model = make_baseline_model(n_subjects=1) + model = make_baseline_model() assert model.ages.values[0] == 51.0 - assert model.ages.values[-1] == 95.0 + assert model.ages.values[-1] == 96.0 def test_dead_regime_is_terminal() -> None: - model = make_baseline_model(n_subjects=1) - assert model.user_regimes["dead"].terminal + model = make_baseline_model() + assert model.graph.laws["dead"].terminal def test_non_terminal_regimes_not_terminal() -> None: - model = make_baseline_model(n_subjects=1) + model = make_baseline_model() for name in REGIME_SPECS: - assert not model.user_regimes[name].terminal + assert not model.graph.laws[name].terminal def test_regime_id_dead_is_last() -> None: @@ -204,13 +203,17 @@ def test_inelig_and_forced_regimes_aime_law_takes_no_claim_inputs(name: str) -> [n for n, s in REGIME_SPECS.items() if s["ss"] == "choose"], ) def test_choose_regimes_aime_law_takes_claim_inputs(name: str) -> None: - """`ss=choose` regimes bake the claim-age adjustment, so their AIME law - reads the `claim_ss` action and `claimed_ss` state.""" + """`ss=choose` regimes bake the claim-age adjustment into the carried PIA. + + The AIME law reads `carried_pia`, whose claim-adjusted input + `pia_adjusted_next_period` reads the `claim_ss` action and `claimed_ss` state. + """ regime = build_regime(name) - aime_law = regime.state_transitions["aime"] - params = set(inspect.signature(aime_law).parameters) - assert "claim_ss" in params - assert "claimed_ss" in params + aime_params = set(inspect.signature(regime.state_transitions["aime"]).parameters) + adjusted = regime.functions["pia_adjusted_next_period"] + adjusted_params = set(inspect.signature(adjusted).parameters) + assert "carried_pia" in aime_params + assert {"claim_ss", "claimed_ss"} <= adjusted_params def test_regime_specs_keys_match_regime_id() -> None: @@ -234,7 +237,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(n_subjects=1) + model = make_baseline_model() for name in REGIME_SPECS: assert isinstance(model.user_regimes[name].states["pension_wealth"], Phased), ( name @@ -282,9 +285,9 @@ def test_hcc_persistent_and_transitory_are_shock_grids() -> None: def test_aca_model_creates_successfully() -> None: - model = make_aca_model(n_subjects=1, policy=PolicyVariant.ACA) + model = make_aca_model(policy=PolicyVariant.ACA) assert len(model.user_regimes) == 19 - assert model.n_periods == 45 + assert model.n_periods == 46 def test_aca_nongroup_inelig_has_real_functions() -> None: @@ -293,7 +296,9 @@ def test_aca_nongroup_inelig_has_real_functions() -> None: regime = regimes["nongroup_nomc_inelig_canwork"] assert regime.functions["mandate_penalty"] is aca_hi.mandate_penalty assert regime.functions["hic_premium_subsidy"] is aca_hi.premium_subsidy - assert regime.functions["cost_sharing_scale"] is aca_hi.cost_sharing + assert ( + regime.functions["cost_sharing_scale"] is aca_hi.cost_sharing_scale_of_bracket + ) assert regime.functions["cash_on_hand"] is aca_hi.cash_on_hand assert regime.functions["primary_oop"] is aca_hi.primary_oop assert regime.functions["is_medicaid_eligible"] is aca_hi.is_medicaid_eligible @@ -305,7 +310,9 @@ def test_aca_no_mandate_has_no_mandate_function() -> None: regime = regimes["nongroup_nomc_inelig_canwork"] assert "mandate_penalty" not in regime.functions assert regime.functions["hic_premium_subsidy"] is aca_hi.premium_subsidy - assert regime.functions["cost_sharing_scale"] is aca_hi.cost_sharing + assert ( + regime.functions["cost_sharing_scale"] is aca_hi.cost_sharing_scale_of_bracket + ) assert regime.functions["cash_on_hand"] is aca_hi.cash_on_hand @@ -323,7 +330,7 @@ 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: """All policy variants create valid models.""" - model = make_aca_model(n_subjects=1, policy=policy) + model = make_aca_model(policy=policy) assert len(model.user_regimes) == 19 @@ -363,62 +370,33 @@ def test_aca_only_medicaid_expansion() -> None: def test_baseline_model_creates() -> None: """Baseline model creates successfully without PolicyVariant.""" - model = make_baseline_model(n_subjects=1) + model = make_baseline_model() assert len(model.user_regimes) == 19 -@pytest.mark.parametrize( - ("config_field", "state_name"), - [ - ("spousal_income_distributed", "spousal_income"), - ], -) -def test_discrete_state_distributed_flag_propagates_to_model_states( - config_field: str, state_name: str -) -> None: - """`GridConfig._distributed=True` sets `distributed=True` on the - model-level `DiscreteGrid` for that axis (sharding is legal only on - model-level states).""" - gc = replace(BENCHMARK_GRID_CONFIG, **{config_field: True}) - grids = build_grids( - grid_config=gc, - fixed_params=_FIXED_PARAMS, - wage_params=_WAGE_PARAMS, - pref_type_grid=DiscreteGrid(BenchmarkPrefType), - ) - model_states = build_model_states(grids) - assert model_states[state_name].distributed is True - - -@pytest.mark.parametrize( - "state_name", - ["lagged_labor_supply", "claimed_ss", "spousal_income"], -) -def test_discrete_state_distributed_flag_defaults_to_false(state_name: str) -> None: - """`distributed` on inline-built discrete states defaults to `False` so - configurations that do not opt in see no behaviour change.""" - if state_name == "spousal_income": - grid = build_model_states(_GRIDS)[state_name] - else: - grid = build_regime("retiree_dimc_choose_canwork").states[state_name] - assert grid.distributed is False - - def test_dead_regime_prunes_unused_broadcast_states() -> None: """States every living regime shares are declared once at the model level; - `dead` keeps only what the bequest DAG reads (`assets`, `pref_type`). + `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(n_subjects=1) + model = make_baseline_model() assert model.pruned_variables["dead"] == frozenset( {"aime", "spousal_income", "hcc_persistent", "hcc_transitory"} ) - assert set(model.user_regimes["dead"].states) == {"assets", "pref_type"} + + +def test_dead_regime_keeps_only_the_bequest_states() -> None: + """`dead` carries only assets and preference type for valuing bequests.""" + model = make_baseline_model() + assert set(model.user_regimes["dead"].states) == { + "assets", + "pref_type", + } def test_living_regimes_keep_every_broadcast_state() -> None: """Every model-level state is read by each living regime's DAG, so pruning removes nothing outside `dead`.""" - model = make_baseline_model(n_subjects=1) + model = make_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 new file mode 100644 index 0000000..85bc6f5 --- /dev/null +++ b/tests/test_mortality.py @@ -0,0 +1,136 @@ +"""Tests for health-dependent mortality and the terminal age.""" + +from collections.abc import Callable +from typing import Any + +import jax.numpy as jnp +import numpy as np +import pytest + +from aca_model.agent.labor_market import LaborSupply +from aca_model.baseline.regimes._common import ( + REGIME_SPECS, + RegimeId, + make_targets, + select_target_for_age, + transition_ages, +) +from aca_model.baseline.regimes._nongroup import ( + _make_transition_canwork as nongroup_canwork, +) +from aca_model.baseline.regimes._nongroup import ( + _make_transition_forcedout as nongroup_forcedout, +) +from aca_model.baseline.regimes._retiree import ( + _make_transition_canwork as retiree_canwork, +) +from aca_model.baseline.regimes._retiree import ( + _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 + +N_PERIODS = MODEL_CONFIG.end_age - MODEL_CONFIG.start_age +# Survival by (period, health): distinct per health state. +SURVIVAL = jnp.tile(jnp.array([0.90, 0.95, 0.99]), (N_PERIODS, 1)) +# P(disabled next period | health); inert for the working households here. +PROB_DISABLED = jnp.zeros((N_PERIODS, 3)) + + +def _tied(health: int) -> jnp.ndarray: + own, ng = make_targets("tied_nomc_inelig_canwork") + return tied_canwork(own=own, ng=ng)( + age=jnp.int32(55), + period=jnp.int32(4), + health=jnp.int32(health), + labor_supply=jnp.array(LaborSupply.h2000), + prob_disabled_next=PROB_DISABLED, + is_ssi_eligible=jnp.array(False), + survival_probs=SURVIVAL, + ) + + +def _retiree(health: int) -> jnp.ndarray: + own, ng = make_targets("retiree_nomc_inelig_canwork") + return retiree_canwork(own=own, ng=ng)( + age=jnp.int32(55), + period=jnp.int32(4), + health=jnp.int32(health), + labor_supply=jnp.array(LaborSupply.h2000), + prob_disabled_next=PROB_DISABLED, + is_ssi_eligible=jnp.array(False), + survival_probs=SURVIVAL, + ) + + +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)( + age=jnp.int32(80), + period=jnp.int32(29), + health=jnp.int32(health), + is_ssi_eligible=jnp.array(False), + survival_probs=SURVIVAL, + ) + + +def _nongroup(health: int) -> jnp.ndarray: + own, _ng = make_targets("nongroup_nomc_inelig_canwork") + return nongroup_canwork(own=own)( + age=jnp.int32(55), + period=jnp.int32(4), + health=jnp.int32(health), + labor_supply=jnp.array(LaborSupply.h2000), + prob_disabled_next=PROB_DISABLED, + survival_probs=SURVIVAL, + ) + + +def _nongroup_forcedout(health: int) -> jnp.ndarray: + own, _ng = make_targets("nongroup_oamc_forced_forcedout") + return nongroup_forcedout(gets_medicare=True, own=own)( + age=jnp.int32(80), + period=jnp.int32(29), + health=jnp.int32(health), + survival_probs=SURVIVAL, + ) + + +@pytest.mark.parametrize( + "transition", + [_tied, _retiree, _retiree_forcedout, _nongroup, _nongroup_forcedout], +) +@pytest.mark.parametrize(("health", "expected_death"), [(0, 0.10), (1, 0.05)]) +def test_death_probability_is_one_minus_survival_at_current_health( + transition: Callable[[int], Any], health: int, expected_death: float +) -> None: + """The dead regime gets `1 - survival_probs[period, health]`.""" + probs = transition(health) + np.testing.assert_allclose(probs[RegimeId.dead], expected_death, atol=1e-12) + + +def test_survivor_at_94_moves_to_forcedout_regime_at_95() -> None: + """95 is an age at which people can be alive.""" + own, _ng = make_targets("nongroup_oamc_forced_forcedout") + target = select_target_for_age(95, True, own) + assert int(target) == int(RegimeId.nongroup_oamc_forced_forcedout) + + +def test_everybody_is_dead_at_96() -> None: + """Nobody is alive at `end_age` (96).""" + own, _ng = make_targets("nongroup_oamc_forced_forcedout") + target = select_target_for_age(MODEL_CONFIG.end_age, True, own) + assert int(target) == int(RegimeId.dead) + + +@pytest.mark.parametrize( + "regime", ["nongroup_oamc_forced_forcedout", "retiree_oamc_forced_forcedout"] +) +def test_forcedout_regimes_are_active_at_95(regime: str) -> None: + """The work-forced-out regimes are active through the last alive age, 95.""" + assert transition_ages(REGIME_SPECS[regime])[-1] == 95 + + +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_labor_live_validation.py b/tests/test_nbegm_labor_live_validation.py index 81eccf4..62b4f9e 100644 --- a/tests/test_nbegm_labor_live_validation.py +++ b/tests/test_nbegm_labor_live_validation.py @@ -1,93 +1,52 @@ -"""NBEGM solves the M1 regime with a live labor-supply choice, matching brute. - -The M1 regime `nongroup_nomc_inelig_canwork` carries `labor_supply` (5 levels) as a -genuine discrete action while `buy_private` is fixed. Labor supply feeds the `aime` -co-state (earnings accrual), the `lagged_labor_supply` co-state, and the leisure term -in period utility — every branch-dependent channel at once. NBEGM's ride-along -discrete envelope solves each labor branch against its own continuation and utility; -the value function must match a brute-force solve on the same state-action space. - -Both solvers run the full 18-regime model at the benchmark grid with `buy_private` -fixed and `labor_supply` live, in NBEGM's `"bridged"` cliff-read mode so the -comparison isolates the solver machinery from the asset-grid read convention. +"""NBEGM solves the M1 regime with its discrete choices live, matching brute. + +The M1 regime `nongroup_nomc_inelig_canwork` carries `labor_supply` (5 levels) and +`buy_private` as genuine discrete actions. Labor supply feeds the `aime` co-state +(earnings accrual), the `lagged_labor_supply` co-state, and the leisure term in period +utility — every branch-dependent channel at once. NBEGM's ride-along discrete envelope +solves each branch against its own continuation and utility; the value function must +match a brute-force solve on the same state-action space. + +Both solvers run the full 18-regime model at the benchmark grid in NBEGM's `"bridged"` +cliff-read mode, so the comparison isolates the solver machinery from the asset-grid +read convention. """ import dataclasses import numpy as np import pytest -from helpers.model import _DERIVED_CATEGORICALS # ty: ignore[unresolved-import] -from lcm import DiscreteGrid +from helpers.model import _DERIVED_CATEGORICALS +from lcm import DiscreteGrid, ExecutionConfig -import aca_model.baseline.regimes._nongroup as nongroup_mod from aca_model.agent.preferences import BenchmarkPrefType from aca_model.baseline.model import create_model from aca_model.baseline.regimes import SolverName -from aca_model.baseline.regimes._common import Grids, RegimeSpec from aca_model.benchmark import get_benchmark_params from aca_model.config import BENCHMARK_GRID_CONFIG _M1_REGIME = "nongroup_nomc_inelig_canwork" -def _is_m1(spec: RegimeSpec) -> bool: - return spec["ss"] == "inelig" and spec["mc"] == "nomc" - - -@pytest.fixture -def m1_labor_live(monkeypatch: pytest.MonkeyPatch) -> None: - """Keep `labor_supply` live and fix only `buy_private` on M1, under every solver. - - The NBEGM wiring fixes both discrete actions on its own; this override keeps - labor supply as a genuine action (so the branch compiler solves it) and fixes - `buy_private`, applying the same choice to the brute build so both solvers share - one state-action space. - """ - original_build_functions = nongroup_mod._build_functions # noqa: SLF001 - original_build_actions = nongroup_mod.build_actions - - def build_functions_labor_live(spec: RegimeSpec, **kwargs: bool) -> dict: - if _is_m1(spec): - return original_build_functions( - spec, fix_buy_private=True, fix_labor_supply=False - ) - return original_build_functions(spec, **kwargs) - - def build_actions_labor_live( - spec: RegimeSpec, grids: Grids, **kwargs: bool - ) -> dict: - if _is_m1(spec): - return original_build_actions( - spec, grids, drop_buy_private=True, drop_labor_supply=False - ) - return original_build_actions(spec, grids, **kwargs) - - monkeypatch.setattr(nongroup_mod, "_build_functions", build_functions_labor_live) - monkeypatch.setattr(nongroup_mod, "build_actions", build_actions_labor_live) - - def _solve_m1(solver: SolverName) -> tuple[dict[int, np.ndarray], int]: - # The CPU XLA backend does not fuse the ride-cell fan-out and materialises the - # whole flattened ride mesh at once, so a full-model solve needs hundreds of GiB on - # host even at a tiny asset grid. `n_nbegm_cell_block_size` streams the mesh in - # blocks (identical result) to bound the peak to the GPU's few-GiB footprint; the - # live-labor branch axis makes this essential on CPU. A coarser savings grid keeps - # the check quick. On GPU the whole-mesh vmap stays small, so production sets 0. + # A bounded ride-cell width limits continuation fan-out on CPU. + # Savings resolution is shared by both numerical comparison arms. grid_config = dataclasses.replace( BENCHMARK_GRID_CONFIG, nbegm_jump_read="bridged", - n_nbegm_cell_block_size=32, n_savings_gridpoints=50, ) fixed_params, wage_params, _ = get_benchmark_params(model=None) model = create_model( - n_subjects=1, fixed_params=fixed_params, wage_params=wage_params, derived_categoricals=_DERIVED_CATEGORICALS, grid_config=grid_config, pref_type_grid=DiscreteGrid(BenchmarkPrefType), solver=solver, + execution_config=ExecutionConfig( + axis_widths={"cell": 32} if solver == "nbegm" else {} + ), ) _, _, params = get_benchmark_params(model=model) solution = model.solve(params=params, log_level="off") @@ -96,7 +55,7 @@ def _solve_m1(solver: SolverName) -> tuple[dict[int, np.ndarray], int]: ).index("assets") return { period: np.asarray(regimes[_M1_REGIME]) - for period, regimes in solution.items() + for period, regimes in solution.values.items() if _M1_REGIME in regimes }, assets_axis @@ -127,7 +86,6 @@ def _cliff_band_mask(reference: np.ndarray, assets_axis: int) -> np.ndarray: @pytest.mark.long_running -@pytest.mark.usefixtures("m1_labor_live") def test_nbegm_m1_labor_live_agrees_with_brute_split_by_cliff_band() -> None: """The M1 value functions agree cell-wise with labor live, gated per region. diff --git a/tests/test_nbegm_model_creation.py b/tests/test_nbegm_model_creation.py index d98bebb..5394c25 100644 --- a/tests/test_nbegm_model_creation.py +++ b/tests/test_nbegm_model_creation.py @@ -1,11 +1,15 @@ -"""NBEGM solver wiring: `solver="nbegm"` is a per-regime option. - -Unlike DC-EGM (a global Euler solver on every living regime), NBEGM solves a -single 1-D consumption/savings regime with at most one discrete action, so it -attaches only to the M1 vertical-slice regime `nongroup_nomc_inelig_canwork`; -every other living regime keeps brute force. The savings-form spec is shared -with DC-EGM (NBEGM's budget node is `resources`, the post-decision function is -`savings`). +"""NBEGM solver wiring. + +`solver="nbegm"` solves every living regime with NB-EGM, so each of them carries +the savings-form budget the solver reads: the budget node is `resources` and the +post-decision function is `savings`, the same spec DC-EGM uses. The `dead` regime +is terminal and keeps its own solver. + +The regime declares every choice its structure affords — whether to buy +non-group coverage, and how many hours to work — and the discrete envelope +branches over the Cartesian product of those grids. The regime is never narrowed +to fit the solver: a solver that cannot carry a choice refuses the regime rather +than being handed a model that omits it. """ import dataclasses @@ -13,10 +17,11 @@ from typing import cast import pytest -from helpers.model import _DERIVED_CATEGORICALS # ty: ignore[unresolved-import] +from helpers.model import _DERIVED_CATEGORICALS, aca_fixed_params from lcm import DiscreteGrid, Model, Regime +from lcm.consumption_savings_regime import ConsumptionSavingsRegime, LiquidMargin from lcm.exceptions import RegimeInitializationError -from lcm.solvers import NBEGM, GridSearch +from lcm.solvers import NBEGM from aca_model.aca import PolicyVariant from aca_model.aca.model import create_model as create_aca_model @@ -37,13 +42,26 @@ _FIXED_PARAMS, _WAGE_PARAMS, _ = get_benchmark_params(model=None) +# pylcm folds a shock state only under grid search, and the EGM solvers take the +# savings-form asset law, so the transitory medical-cost shock stays a state and +# the bill stays inside `next_assets`. +_EGM_GRID_CONFIG = dataclasses.replace( + BENCHMARK_GRID_CONFIG, fold_hcc_transitory=False, subtract_hcc_bill=False +) + _M1_REGIME = "nongroup_nomc_inelig_canwork" _BRUTE_REGIME = "retiree_nomc_inelig_canwork" +# `labor_supply` enters `countable_income`, which carries the SSI income test, so each +# labor level puts that breakpoint at a different liquid level. The one-sided read +# publishes its cliff limits on a single query grid shared across branches, which the +# two cannot both satisfy, so a regime carrying `labor_supply` needs the bridged read. +_BRIDGED_GRID_CONFIG = dataclasses.replace(_EGM_GRID_CONFIG, nbegm_jump_read="bridged") + def _build_regimes(solver: SolverName) -> dict[str, Regime]: return build_all_regimes( - grid_config=BENCHMARK_GRID_CONFIG, + grid_config=_EGM_GRID_CONFIG, fixed_params=_FIXED_PARAMS, wage_params=_WAGE_PARAMS, pref_type_grid=DiscreteGrid(BenchmarkPrefType), @@ -53,7 +71,6 @@ def _build_regimes(solver: SolverName) -> dict[str, Regime]: def _build_model_with(solver: SolverName, grid_config: GridConfig) -> Model: return create_model( - n_subjects=1, fixed_params=_FIXED_PARAMS, wage_params=_WAGE_PARAMS, derived_categoricals=_DERIVED_CATEGORICALS, @@ -64,43 +81,57 @@ def _build_model_with(solver: SolverName, grid_config: GridConfig) -> Model: def _build_model(solver: SolverName) -> Model: - return _build_model_with(solver, BENCHMARK_GRID_CONFIG) + return _build_model_with(solver, _BRIDGED_GRID_CONFIG) def _grids() -> Grids: return build_grids( - grid_config=BENCHMARK_GRID_CONFIG, + grid_config=_EGM_GRID_CONFIG, fixed_params=_FIXED_PARAMS, wage_params=_WAGE_PARAMS, pref_type_grid=DiscreteGrid(BenchmarkPrefType), ) -def test_nbegm_attaches_only_to_the_m1_regime() -> None: - """`solver="nbegm"` gives the M1 slice regime a `NBEGM` config and leaves - every other living regime on brute force.""" +def _liquid_margins(regimes: Mapping[str, Regime]) -> set[LiquidMargin]: + """Return the liquid margin every living regime declares.""" + return { + cast("ConsumptionSavingsRegime", regimes[name]).liquid for name in REGIME_SPECS + } + + +def test_nbegm_attaches_to_every_living_regime() -> None: + """`solver="nbegm"` solves every living regime with NB-EGM. + + A solver that reached only some regimes would leave the rest on brute + force while still reporting itself as the model's solver, so the choice of + solver would not be visible in the result it produced. + """ regimes = _build_regimes("nbegm") - assert isinstance(regimes[_M1_REGIME].solver, NBEGM) - for name in REGIME_SPECS: - if name == _M1_REGIME: - continue - assert isinstance(regimes[name].solver, GridSearch), name + on_brute_force = [ + name for name in REGIME_SPECS if not isinstance(regimes[name].solver, NBEGM) + ] + assert on_brute_force == [] -def test_build_nbegm_solver_uses_the_savings_form_resources_budget() -> None: - """The NBEGM config inverts against `resources` in post-decision savings +def test_nbegm_regimes_declare_the_savings_form_resources_budget() -> None: + """Every NB-EGM regime inverts against `resources` in post-decision savings form, matching the DC-EGM contract the regime is rewired into.""" - solver = build_nbegm_solver(_grids()) - assert isinstance(solver, NBEGM) - assert solver.budget_target == "resources" - assert solver.post_decision_function == "savings" + declared = { + (margin.resources, margin.post_decision_state) + for margin in _liquid_margins(_build_regimes("nbegm")) + } + assert declared == {("resources", "savings")} -def test_build_nbegm_solver_names_assets_as_the_euler_axis() -> None: - """`assets` is the liquid (Euler) axis; `aime` and the stochastic shock grids - ride along, so the solver names the Euler axis explicitly.""" - solver = build_nbegm_solver(_grids()) - assert solver.continuous_state == "assets" +def test_nbegm_regimes_declare_assets_as_the_liquid_euler_axis() -> None: + """`assets` paid down by `consumption_dollars` is the liquid margin of every + NB-EGM regime; `aime` and the stochastic shock grids ride along.""" + declared = { + (margin.state, margin.action) + for margin in _liquid_margins(_build_regimes("nbegm")) + } + assert declared == {("assets", "consumption_dollars")} def test_build_nbegm_solver_forwards_the_jump_read_mode() -> None: @@ -109,7 +140,7 @@ def test_build_nbegm_solver_forwards_the_jump_read_mode() -> None: `"bridged"` is the fast estimation setting (plain carry rows, fold kept); `"one_sided"` (the default) publishes exact one-sided cliff limits. """ - grid_config = dataclasses.replace(BENCHMARK_GRID_CONFIG, nbegm_jump_read="bridged") + grid_config = dataclasses.replace(_EGM_GRID_CONFIG, nbegm_jump_read="bridged") grids = build_grids( grid_config=grid_config, fixed_params=_FIXED_PARAMS, @@ -120,30 +151,43 @@ def test_build_nbegm_solver_forwards_the_jump_read_mode() -> None: assert solver.jump_read == "bridged" -def test_nbegm_m1_regime_fixes_buy_private() -> None: - """The NBEGM M1 slice drops `buy_private` as an action (fixed to purchase), - so the only choice is continuous consumption; the brute M1 regime keeps it.""" +def test_nbegm_m1_regime_declares_the_same_actions_as_brute_force() -> None: + """Which solver runs a regime does not change the choices the household has. + + The M1 regime affords a coverage choice and an hours choice, so it declares + both under either solver. + """ nbegm_m1 = _build_regimes("nbegm")[_M1_REGIME] brute_m1 = _build_regimes("brute_force")[_M1_REGIME] - assert "buy_private" not in nbegm_m1.actions - assert "buy_private" in brute_m1.actions + assert set(nbegm_m1.actions) == set(brute_m1.actions) -def test_nbegm_m1_regime_fixes_labor_supply() -> None: - """The NBEGM M1 slice drops `labor_supply` as an action (fixed to full-time - work), so no discrete action remains and the only choice is continuous - consumption; the brute M1 regime keeps `labor_supply`.""" +def test_nbegm_m1_regime_declares_both_discrete_actions() -> None: + """The M1 regime's discrete choices are whether to buy non-group coverage and + how many hours to work.""" nbegm_m1 = _build_regimes("nbegm")[_M1_REGIME] - brute_m1 = _build_regimes("brute_force")[_M1_REGIME] - assert "labor_supply" not in nbegm_m1.actions - assert "labor_supply" in brute_m1.actions + discrete = { + name + for name, grid in nbegm_m1.actions.items() + if isinstance(grid, DiscreteGrid) + } + assert discrete == {"buy_private", "labor_supply"} -def test_nbegm_m1_regime_has_no_discrete_action() -> None: - """With both discrete actions fixed, the NBEGM M1 slice leaves only the - continuous consumption choice — no `DiscreteGrid` action remains.""" - nbegm_m1 = _build_regimes("nbegm")[_M1_REGIME] - assert not any(isinstance(grid, DiscreteGrid) for grid in nbegm_m1.actions.values()) +def test_nbegm_model_builds_a_regime_declaring_several_discrete_actions() -> None: + """The M1 regime builds under NBEGM with both of its discrete actions live. + + The discrete envelope branches over the Cartesian product of the regime's + discrete action grids, so a regime is never narrowed to fit the solver. + """ + model = _build_model("nbegm") + nbegm_m1 = model.user_regimes[_M1_REGIME] + discrete = { + name + for name, grid in nbegm_m1.actions.items() + if isinstance(grid, DiscreteGrid) + } + assert discrete == {"buy_private", "labor_supply"} def test_nbegm_m1_regime_takes_the_savings_form_assets_laws() -> None: @@ -162,15 +206,20 @@ def test_nbegm_m1_regime_takes_the_savings_form_assets_laws() -> None: assert law is expected, target_name -def test_nbegm_savings_form_functions_are_scoped_to_the_m1_regime() -> None: - """Under NBEGM only the M1 regime carries the savings-form budget functions - (`resources`, `savings`); brute regimes keep the cash-on-hand form and carry - neither.""" +def test_nbegm_gives_every_living_regime_the_savings_form_budget() -> None: + """Under NBEGM every living regime carries `resources` and `savings`. + + They are the solver's budget contract, so a regime NB-EGM solves without + them cannot be built at all; a regime it does not solve has no use for + them, which is why the brute-force build omits them. + """ model = _build_model("nbegm") - m1_functions = model.user_regimes[_M1_REGIME].functions - assert "resources" in m1_functions - assert "savings" in m1_functions - assert "resources" not in model.user_regimes[_BRUTE_REGIME].functions + missing = [ + name + for name in REGIME_SPECS + if not {"resources", "savings"} <= set(model.user_regimes[name].functions) + ] + assert missing == [] def test_nbegm_m1_regime_does_not_carry_inverse_marginal_utility() -> None: @@ -234,70 +283,29 @@ def test_ssi_benefit_declares_the_income_test_kink() -> None: assert income_test.indexed_by == "spousal_income" -def test_nbegm_keeps_labor_supply_live_when_configured() -> None: - """With `nbegm_live_labor_supply=True`, the M1 regime carries `labor_supply` - as a genuine discrete action under NBEGM while `buy_private` stays fixed, so - the branch compiler solves the labor choice against the cliffed budget.""" - grid_config = dataclasses.replace( - BENCHMARK_GRID_CONFIG, nbegm_live_labor_supply=True - ) - regimes = build_all_regimes( - grid_config=grid_config, - fixed_params=_FIXED_PARAMS, - wage_params=_WAGE_PARAMS, - pref_type_grid=DiscreteGrid(BenchmarkPrefType), - solver="nbegm", - ) - actions = regimes[_M1_REGIME].actions - assert "labor_supply" in actions - assert "buy_private" not in actions - - -def test_nbegm_live_labor_supply_requires_the_bridged_cliff_read() -> None: - """A live `labor_supply` action builds only under `nbegm_jump_read="bridged"`. +def test_nbegm_labor_supply_requires_the_bridged_cliff_read() -> None: + """The M1 regime builds under NBEGM only with `nbegm_jump_read="bridged"`. `labor_supply` enters `countable_income`, which carries the SSI income test, so each labor level puts that breakpoint at a different liquid level. The one-sided read publishes its cliff limits on one query grid shared across branches, so the two cannot both hold and the build is refused. """ - live_labor = dataclasses.replace( - BENCHMARK_GRID_CONFIG, nbegm_live_labor_supply=True - ) with pytest.raises(RegimeInitializationError, match="must not enter any schedule"): - _build_model_with( - "nbegm", dataclasses.replace(live_labor, nbegm_jump_read="one_sided") - ) + _build_model_with("nbegm", _EGM_GRID_CONFIG) - bridged = dataclasses.replace(live_labor, nbegm_jump_read="bridged") - assert isinstance(_build_model_with("nbegm", bridged), Model) - - -def test_nbegm_fixes_labor_supply_by_default() -> None: - """By default NBEGM fixes both discrete actions on the M1 regime, so the - only remaining choice is continuous consumption against the cliffed budget.""" - regimes = _build_regimes("nbegm") - actions = regimes[_M1_REGIME].actions - assert "labor_supply" not in actions - assert "buy_private" not in actions + assert isinstance(_build_model_with("nbegm", _BRIDGED_GRID_CONFIG), Model) @pytest.mark.parametrize("policy", list(PolicyVariant)) def test_nbegm_builds_every_aca_policy_variant(policy: PolicyVariant) -> None: """Every ACA policy variant builds a model under NBEGM with the M1 regime on - the solver and labor live — the overlay's function swaps compose with the - branch compiler's per-regime wiring.""" - # Live labor requires the bridged cliff read — see - # `test_nbegm_live_labor_supply_requires_the_bridged_cliff_read`. - grid_config = dataclasses.replace( - BENCHMARK_GRID_CONFIG, - nbegm_live_labor_supply=True, - nbegm_jump_read="bridged", - ) + the solver and both discrete choices live — the overlay's function swaps + compose with the branch compiler's per-regime wiring.""" + grid_config = _BRIDGED_GRID_CONFIG model = create_aca_model( - n_subjects=1, policy=policy, - fixed_params=_FIXED_PARAMS, + fixed_params=aca_fixed_params(fixed_params=_FIXED_PARAMS, policy=policy), wage_params=_WAGE_PARAMS, derived_categoricals=_DERIVED_CATEGORICALS, grid_config=grid_config, @@ -321,21 +329,16 @@ def test_nbegm_builds_every_aca_policy_variant(policy: PolicyVariant) -> None: def test_nbegm_aca_variants_leave_no_free_buy_private_params( policy: PolicyVariant, ) -> None: - """With `buy_private` fixed under the NBEGM M1 slice, no ACA-swapped - function may leave `buy_private` as a free parameter — the params template - holds no `buy_private` leaves, so solve/simulate never demand a - `*__buy_private` entry the pipeline cannot supply.""" - # Live labor requires the bridged cliff read — see - # `test_nbegm_live_labor_supply_requires_the_bridged_cliff_read`. - grid_config = dataclasses.replace( - BENCHMARK_GRID_CONFIG, - nbegm_live_labor_supply=True, - nbegm_jump_read="bridged", - ) + """`buy_private` is a choice, never a parameter. + + No ACA-swapped function may leave it as a free parameter — the params + template holds no `buy_private` leaves, so solve/simulate never demand a + `*__buy_private` entry the pipeline cannot supply. + """ + grid_config = _BRIDGED_GRID_CONFIG model = create_aca_model( - n_subjects=1, policy=policy, - fixed_params=_FIXED_PARAMS, + fixed_params=aca_fixed_params(fixed_params=_FIXED_PARAMS, policy=policy), wage_params=_WAGE_PARAMS, derived_categoricals=_DERIVED_CATEGORICALS, grid_config=grid_config, @@ -346,7 +349,7 @@ def test_nbegm_aca_variants_leave_no_free_buy_private_params( offenders = [ (regime_name, function_name) for regime_name, functions in template.items() - if isinstance(functions, dict) + if isinstance(functions, dict) and regime_name != "edges" for function_name, params in functions.items() if isinstance(params, dict) and "buy_private" in params ] diff --git a/tests/test_nbegm_solve_validation.py b/tests/test_nbegm_solve_validation.py index 6d4449d..bbe2265 100644 --- a/tests/test_nbegm_solve_validation.py +++ b/tests/test_nbegm_solve_validation.py @@ -1,9 +1,9 @@ """NBEGM vs brute-force agreement on the M1 slice value function. -Both solvers run the full 18-regime model at the benchmark grid with the M1 -regime's discrete actions fixed (labor supply to full-time, `buy_private` to -purchase) so the value functions are defined on identical state spaces and -differ only through the continuous-consumption solver. +Both solvers run the full 18-regime model at the benchmark grid on the same +state-action space — the M1 regime declares its `labor_supply` and `buy_private` +choices under either solver — so the value functions differ only through the +continuous-consumption solver. NBEGM runs its `"bridged"` cliff-read mode so both solvers share the same read convention (a finite brute reads child values by linear interpolation @@ -19,57 +19,22 @@ import numpy as np import pytest -from helpers.model import _DERIVED_CATEGORICALS # ty: ignore[unresolved-import] +from helpers.model import _DERIVED_CATEGORICALS from lcm import DiscreteGrid -import aca_model.baseline.regimes._nongroup as nongroup_mod from aca_model.agent.preferences import BenchmarkPrefType from aca_model.baseline.model import create_model from aca_model.baseline.regimes import SolverName -from aca_model.baseline.regimes._common import Grids, RegimeSpec from aca_model.benchmark import get_benchmark_params from aca_model.config import BENCHMARK_GRID_CONFIG _M1_REGIME = "nongroup_nomc_inelig_canwork" -@pytest.fixture -def m1_actions_fixed_for_brute(monkeypatch: pytest.MonkeyPatch) -> None: - """Fix the M1 regime's discrete actions under every solver. - - The NBEGM wiring fixes `labor_supply` and `buy_private` on its own; this - fixture applies the same fixing to the brute build so both solvers produce - an M1 value function on the same state-action space. - """ - original_build_functions = nongroup_mod._build_functions # noqa: SLF001 - original_build_actions = nongroup_mod.build_actions - - def build_functions_fixed(spec: RegimeSpec, **kwargs: bool) -> dict: - if spec["ss"] == "inelig" and spec["mc"] == "nomc": - return original_build_functions( - spec, fix_buy_private=True, fix_labor_supply=True - ) - return original_build_functions(spec, **kwargs) - - def build_actions_fixed(spec: RegimeSpec, grids: Grids, **kwargs: bool) -> dict: - if spec["ss"] == "inelig" and spec["mc"] == "nomc": - return original_build_actions( - spec, - grids, - drop_buy_private=True, - drop_labor_supply=True, - ) - return original_build_actions(spec, grids, **kwargs) - - monkeypatch.setattr(nongroup_mod, "_build_functions", build_functions_fixed) - monkeypatch.setattr(nongroup_mod, "build_actions", build_actions_fixed) - - def _solve_m1(solver: SolverName) -> dict[int, np.ndarray]: grid_config = dataclasses.replace(BENCHMARK_GRID_CONFIG, nbegm_jump_read="bridged") fixed_params, wage_params, _ = get_benchmark_params(model=None) model = create_model( - n_subjects=1, fixed_params=fixed_params, wage_params=wage_params, derived_categoricals=_DERIVED_CATEGORICALS, @@ -81,13 +46,12 @@ def _solve_m1(solver: SolverName) -> dict[int, np.ndarray]: solution = model.solve(params=params, log_level="off") return { period: np.asarray(regimes[_M1_REGIME]) - for period, regimes in solution.items() + for period, regimes in solution.values.items() if _M1_REGIME in regimes } @pytest.mark.long_running -@pytest.mark.usefixtures("m1_actions_fixed_for_brute") def test_nbegm_m1_value_function_agrees_with_brute_in_the_bulk() -> None: """The M1 value functions agree cell-wise away from the cliff tail. @@ -102,7 +66,7 @@ def test_nbegm_m1_value_function_agrees_with_brute_in_the_bulk() -> None: bq = _solve_m1("nbegm") brute = _solve_m1("brute_force") - assert bq.keys() == brute.keys() + assert bq.keys() == brute.keys() == set(range(11)) rel_diffs = [] for period in bq: finite_bq = np.isfinite(bq[period]) diff --git a/tests/test_pension_integration.py b/tests/test_pension_integration.py index b6af1c3..0d1c6df 100644 --- a/tests/test_pension_integration.py +++ b/tests/test_pension_integration.py @@ -80,21 +80,39 @@ def test_imputation_chain_full_benefit_to_wealth() -> None: assert jnp.isclose(result, 500.0, atol=ATOL) -def test_total_to_pia_inverts_imputed_benefit_via_dag() -> None: - """`pbmax → total_to_pia` round-trip via dags recovers original PIA.""" +def test_total_to_pia_keeps_adjusted_total_via_dag() -> None: + """The carried PIA keeps SS plus after-tax imputed pension at its adjusted total. + + Accrued PIA 8,000 is adjusted to 7,000; the carried PIA `p` satisfies + `p + (1 - τ) pbmax(p) = 7,000 + (1 - τ) pbmax(8,000)`. + """ + next_period_kwargs = {f"{k}_next_period": v for k, v in PBMAX_KWARGS.items()} functions = { - "pension_benefit": pensions.full_benefit, + "full_benefit_next_period": pensions.full_benefit_next_period, "total_to_pia": pensions.total_to_pia, } combined = concatenate_functions(functions, targets="total_to_pia") - recovered = combined( - pia=jnp.array(8000.0), + mtr = jnp.array(0.2) + carried = combined( + pia_adjusted_next_period=jnp.array(7000.0), + pia_unadjusted_next_period=jnp.array(8000.0), + target_his=jnp.int32(0), period=PERIOD, - his=jnp.int32(0), - marginal_tax_rate=jnp.array(0.2), - **PBMAX_KWARGS, + marginal_tax_rate=mtr, + **next_period_kwargs, ) - assert jnp.isclose(recovered, 8000.0, atol=ATOL) + + def pbmax(pia: jnp.ndarray) -> jnp.ndarray: + return pensions.full_benefit_next_period( + pia_unadjusted_next_period=pia, + target_his=jnp.int32(0), + period=PERIOD, + **next_period_kwargs, + ) + + total_carried = carried + (1.0 - mtr) * pbmax(carried) + target = 7000.0 + (1.0 - mtr) * pbmax(jnp.array(8000.0)) + assert jnp.isclose(total_carried, target, atol=ATOL) def test_next_assets_includes_pension_adjustment() -> None: @@ -126,7 +144,11 @@ def test_zero_adjustment_when_his_unchanged() -> None: period=PERIOD, ) accrual_val = pensions.accrual( - labor_income=labor_income, period=PERIOD, his=his, **ACCRUAL_KWARGS + labor_income=labor_income, + age=jnp.int32(60), + period=PERIOD, + his=his, + **ACCRUAL_KWARGS, ) next_exact = pensions.wealth_next_before_adjustment( @@ -173,7 +195,11 @@ def test_rebalancing_preserves_total_wealth_across_his_change() -> None: period=PERIOD, ) accrual_val = pensions.accrual( - labor_income=labor_income, period=PERIOD, his=old_his, **ACCRUAL_KWARGS + labor_income=labor_income, + age=jnp.int32(60), + period=PERIOD, + his=old_his, + **ACCRUAL_KWARGS, ) next_exact = pensions.wealth_next_before_adjustment( @@ -224,7 +250,11 @@ def _solve_phase_adjustment_across_his_change() -> jnp.ndarray: period=PERIOD, ) accrual_val = pensions.accrual( - labor_income=labor_income, period=PERIOD, his=old_his, **ACCRUAL_KWARGS + labor_income=labor_income, + age=jnp.int32(60), + period=PERIOD, + his=old_his, + **ACCRUAL_KWARGS, ) next_exact = pensions.wealth_next_before_adjustment( pension_wealth=pw_old, diff --git a/tests/test_pensions.py b/tests/test_pensions.py index 3da46ff..28c59b4 100644 --- a/tests/test_pensions.py +++ b/tests/test_pensions.py @@ -125,6 +125,7 @@ def test_imputed_benefit_round_trips_through_wealth() -> None: def test_pension_accrual_no_income() -> None: result = pensions.accrual( labor_income=jnp.array(-1000.0), + age=jnp.int32(60), period=jnp.int32(20), his=jnp.int32(0), accrual_intercept=ACCRUAL_INTERCEPT, @@ -139,6 +140,7 @@ def test_pension_accrual_no_income() -> None: def test_pension_accrual_positive() -> None: result = pensions.accrual( labor_income=jnp.array(10000.0), + age=jnp.int32(60), period=jnp.int32(20), his=jnp.int32(0), accrual_intercept=ACCRUAL_INTERCEPT, @@ -186,6 +188,20 @@ def test_pension_wealth_next_with_benefit() -> None: assert jnp.isclose(result, expected, atol=ATOL) +def test_pension_wealth_next_is_undivided_where_nobody_survives() -> None: + """At an age with zero survival probability, next pension wealth is the + compounded wealth net of accrual and benefit, not divided by survival.""" + result = pensions.wealth_next_before_adjustment( + pension_wealth=jnp.array(3000.0), + pension_benefit=jnp.array(2000.0), + pension_accrual=jnp.array(500.0), + rate_of_return=jnp.asarray(0.03), + unconditional_survival_prob=jnp.array([0.99, 0.0]), + period=jnp.int32(1), + ) + assert jnp.isclose(result, 1.03 * 3000.0 + 500.0 - 2000.0, atol=ATOL) + + def test_convert_total_ben_to_pia_below_kink_0() -> None: """Round-trip: PIA below first kink recovers original PIA.""" pia_input = jnp.array(500.0) @@ -195,12 +211,13 @@ def test_convert_total_ben_to_pia_below_kink_0() -> None: pia=pia_input, period=jnp.int32(29), his=jnp.int32(0), **_PBMAX_KWARGS ) recovered = pensions.total_to_pia( - pension_benefit=pbmax, - pia=pia_input, + pia_adjusted_next_period=pia_input, + pia_unadjusted_next_period=pia_input + 1.0, + full_benefit_next_period=pbmax, + target_his=jnp.int32(0), period=jnp.int32(29), - his=jnp.int32(0), marginal_tax_rate=mtr, - **_PBMAX_KWARGS, + **{f"{k}_next_period": v for k, v in _PBMAX_KWARGS.items()}, ) assert jnp.isclose(recovered, pia_input, atol=ATOL) @@ -214,12 +231,13 @@ def test_convert_total_ben_to_pia_between_kinks() -> None: pia=pia_input, period=jnp.int32(29), his=jnp.int32(0), **_PBMAX_KWARGS ) recovered = pensions.total_to_pia( - pension_benefit=pbmax, - pia=pia_input, + pia_adjusted_next_period=pia_input, + pia_unadjusted_next_period=pia_input + 1.0, + full_benefit_next_period=pbmax, + target_his=jnp.int32(0), period=jnp.int32(29), - his=jnp.int32(0), marginal_tax_rate=mtr, - **_PBMAX_KWARGS, + **{f"{k}_next_period": v for k, v in _PBMAX_KWARGS.items()}, ) assert jnp.isclose(recovered, pia_input, atol=ATOL) @@ -233,12 +251,13 @@ def test_convert_total_ben_to_pia_above_kink_1() -> None: pia=pia_input, period=jnp.int32(29), his=jnp.int32(0), **_PBMAX_KWARGS ) recovered = pensions.total_to_pia( - pension_benefit=pbmax, - pia=pia_input, + pia_adjusted_next_period=pia_input, + pia_unadjusted_next_period=pia_input + 1.0, + full_benefit_next_period=pbmax, + target_his=jnp.int32(0), period=jnp.int32(29), - his=jnp.int32(0), marginal_tax_rate=mtr, - **_PBMAX_KWARGS, + **{f"{k}_next_period": v for k, v in _PBMAX_KWARGS.items()}, ) assert jnp.isclose(recovered, pia_input, atol=ATOL) @@ -252,12 +271,13 @@ def test_convert_total_ben_to_pia_zero_mtr() -> None: pia=pia_input, period=jnp.int32(29), his=jnp.int32(0), **_PBMAX_KWARGS ) recovered = pensions.total_to_pia( - pension_benefit=pbmax, - pia=pia_input, + pia_adjusted_next_period=pia_input, + pia_unadjusted_next_period=pia_input + 1.0, + full_benefit_next_period=pbmax, + target_his=jnp.int32(0), period=jnp.int32(29), - his=jnp.int32(0), marginal_tax_rate=mtr, - **_PBMAX_KWARGS, + **{f"{k}_next_period": v for k, v in _PBMAX_KWARGS.items()}, ) assert jnp.isclose(recovered, pia_input, atol=ATOL) @@ -270,7 +290,7 @@ def test_imputed_pension_wealth_uses_unadjusted_pia() -> None: PIA yields the unreduced imputed wealth, independent of any claim-age reduction baked into the carried AIME. """ - result = pensions.imputed_pension_wealth_next_period( + full_benefit_next = pensions.full_benefit_next_period( pia_unadjusted_next_period=jnp.array(500.0), target_his=jnp.int32(0), period=jnp.int32(28), @@ -280,6 +300,10 @@ def test_imputed_pension_wealth_uses_unadjusted_pia() -> None: imp_pia_kink_1_coeff_next_period=jnp.zeros((30, 1)), imp_kink_0_next_period=jnp.full(30, 99999.0), imp_kink_1_next_period=jnp.full(30, 99999.0), + ) + result = pensions.imputed_pension_wealth_next_period( + full_benefit_next_period=full_benefit_next, + period=jnp.int32(28), epdv_constant_pension_next_period=EPDV, ) # pbmax = max(0, -50 + 500·0.2) = 50, pw = 50 · Γ(=10) = 500. diff --git a/tests/test_precision.py b/tests/test_precision.py new file mode 100644 index 0000000..e807bd3 --- /dev/null +++ b/tests/test_precision.py @@ -0,0 +1,28 @@ +"""Process-level floating-point precision selection.""" + +import os +import subprocess +import sys + + +def test_aca_precision_environment_selects_fp32() -> None: + """`ACA_JAX_ENABLE_X64=0` makes ACA computations use 32-bit floats.""" + env = { + **os.environ, + "JAX_ENABLE_X64": "1", + "ACA_JAX_ENABLE_X64": "0", + } + + result = subprocess.run( + [ + sys.executable, + "-c", + "import aca_model, jax; print(jax.config.jax_enable_x64)", + ], + check=True, + capture_output=True, + env=env, + text=True, + ) + + assert result.stdout.strip() == "False" diff --git a/tests/test_preferences.py b/tests/test_preferences.py index 26aa941..b83760b 100644 --- a/tests/test_preferences.py +++ b/tests/test_preferences.py @@ -3,10 +3,16 @@ Parameter values from struct-ret PreferenceParameters fixture. """ +from collections.abc import Callable +from typing import Any + import jax.numpy as jnp import numpy as np +import pytest +from dags import concatenate_functions from aca_model.agent import preferences +from aca_model.baseline.regimes._common import build_dead_regime, build_model_functions # Struct-ret preference parameters. Tests call DAG functions directly, so # every scalar fixed_param is supplied as a 0-d jax array (the type pylcm @@ -58,7 +64,7 @@ def test_scaled_bequest_weight_positive() -> None: consumption_weight=CONSUMPTION_WEIGHT, coefficient_rra=jnp.asarray(5.0), time_endowment=TIME_ENDOWMENT, - time_discount_factor=TIME_DISCOUNT_FACTOR, + discount_factor=TIME_DISCOUNT_FACTOR, rate_of_return=RATE_OF_RETURN, ) assert jnp.isclose(result, 0.820_137_639_127_977_3, rtol=1e-6) @@ -70,7 +76,7 @@ def test_scaled_bequest_weight_log() -> None: consumption_weight=CONSUMPTION_WEIGHT, coefficient_rra=jnp.asarray(1.0), time_endowment=TIME_ENDOWMENT, - time_discount_factor=TIME_DISCOUNT_FACTOR, + discount_factor=TIME_DISCOUNT_FACTOR, rate_of_return=RATE_OF_RETURN, ) assert jnp.isclose(result, 58.235_294_117_647_05, rtol=1e-6) @@ -82,7 +88,7 @@ def test_scaled_bequest_weight_zero() -> None: consumption_weight=CONSUMPTION_WEIGHT, coefficient_rra=jnp.asarray(5.0), time_endowment=TIME_ENDOWMENT, - time_discount_factor=TIME_DISCOUNT_FACTOR, + discount_factor=TIME_DISCOUNT_FACTOR, rate_of_return=RATE_OF_RETURN, ) assert result == 0.0 @@ -173,7 +179,7 @@ def test_bequest_log_regression() -> None: consumption_weight=CONSUMPTION_WEIGHT, coefficient_rra=jnp.asarray(1.0), time_endowment=TIME_ENDOWMENT, - time_discount_factor=TIME_DISCOUNT_FACTOR, + discount_factor=TIME_DISCOUNT_FACTOR, rate_of_return=RATE_OF_RETURN, ) result = preferences.bequest( @@ -201,7 +207,7 @@ def test_bequest_crra_regression() -> None: consumption_weight=CONSUMPTION_WEIGHT, coefficient_rra=jnp.asarray(5.0), time_endowment=TIME_ENDOWMENT, - time_discount_factor=TIME_DISCOUNT_FACTOR, + discount_factor=TIME_DISCOUNT_FACTOR, rate_of_return=RATE_OF_RETURN, ) result = preferences.bequest( @@ -238,7 +244,7 @@ def test_bequest_uses_signed_assets_for_indebted_decedent() -> None: consumption_weight=consumption_weight, coefficient_rra=coefficient_rra, time_endowment=TIME_ENDOWMENT, - time_discount_factor=TIME_DISCOUNT_FACTOR, + discount_factor=TIME_DISCOUNT_FACTOR, rate_of_return=RATE_OF_RETURN, ) result = preferences.bequest( @@ -282,7 +288,7 @@ def test_bequest_floors_estate_base_when_debt_exceeds_the_shifter() -> None: consumption_weight=consumption_weight, coefficient_rra=coefficient_rra, time_endowment=TIME_ENDOWMENT, - time_discount_factor=TIME_DISCOUNT_FACTOR, + discount_factor=TIME_DISCOUNT_FACTOR, rate_of_return=RATE_OF_RETURN, ) result = preferences.bequest( @@ -296,3 +302,64 @@ def test_bequest_floors_estate_base_when_debt_exceeds_the_shifter() -> None: expected = scale * bwt * 1.0 / (1.0 - coefficient_rra) np.testing.assert_allclose(result, expected, atol=1e-9) + + +# --- dead-regime bequest utility per preference type --- + +# Draft estimates (struct-ret original_data/baseline/final/prefs.json). +_DRAFT_CONSUMPTION_WEIGHTS = jnp.asarray( + [0.6776541629520845, 0.8805772328591686, 0.0718086283445225] +) +_DRAFT_COEFFICIENTS_RRA = jnp.asarray( + [3.841252231680976, 0.9990771146810682, 3.8328505891095137] +) +_DRAFT_DISCOUNT_FACTORS = jnp.asarray( + [0.839255583514979, 0.9121549199895648, 1.0595836603747046] +) + + +def _dead_regime_utility() -> Callable[..., Any]: + """Compose the dead regime's utility DAG as pylcm merges it. + + Model-level functions come first; the dead regime's own entries override + them, and `None` entries mask a model-level function out of the regime. + """ + functions = dict(build_model_functions()) + for name, func in build_dead_regime().functions.items(): + if func is None: + functions.pop(name, None) + else: + functions[name] = func + return concatenate_functions(functions, targets="utility") + + +@pytest.mark.parametrize( + ("pref_type", "expected"), + [ + (0, -75.59884789663738), + (1, 39.12342314306149), + (2, -12.47026859098475), + ], +) +def test_dead_regime_bequest_uses_per_type_transformed_weight( + pref_type: int, expected: float +) -> None: + """At the draft estimates and r = 0.05, the bequest utility of 100k assets + equals `θ_B(type) · (A + κ)^((1-ν)γ) / (1-ν) · scale(type)` with the + per-type `θ_B = T^ξ (bw / (1 + r - bw))^ξ₂ / β(type)`.""" + utility = _dead_regime_utility() + result = utility( + assets=jnp.asarray(100_000.0), + pref_type=jnp.asarray(pref_type, dtype=jnp.int32), + bequest_shifter=jnp.asarray(334460.20744719007), + bequest_weight=jnp.asarray(0.02861082652580584), + consumption_weights=_DRAFT_CONSUMPTION_WEIGHTS, + coefficients_rra=_DRAFT_COEFFICIENTS_RRA, + discount_factor_by_type=_DRAFT_DISCOUNT_FACTORS, + time_endowment=jnp.asarray(3926.9478390365557), + rate_of_return=jnp.asarray(0.05), + average_consumption_equiv=jnp.asarray(20_000.0), + fixed_cost_of_work_intercept=jnp.asarray(337.5223349543413), + reference_hours=jnp.asarray(1000.0), + ) + np.testing.assert_allclose(result, expected, rtol=1e-9) diff --git a/tests/test_private_premium.py b/tests/test_private_premium.py new file mode 100644 index 0000000..40e8891 --- /dev/null +++ b/tests/test_private_premium.py @@ -0,0 +1,236 @@ +"""Private non-group premium: zero-profit intercept, pricing rule, and floor rule.""" + +import jax.numpy as jnp +import numpy as np +import pytest + +from aca_model.aca import health_insurance as aca_health_insurance +from aca_model.baseline import health_insurance +from aca_model.baseline.health_insurance import BuyPrivate +from aca_model.baseline.regimes import _nongroup +from aca_model.baseline.regimes._common import REGIME_SPECS + +_MARKUP = jnp.asarray(0.1696) +_MINIMUM = jnp.array([1000.0, 2000.0]) +_SAMPLE_COST = jnp.array([1000.0, 3000.0, 8000.0]) +_SAMPLE_MARRIED = jnp.array([0, 1, 0], dtype=jnp.int32) + + +@pytest.mark.parametrize( + ("coeff", "expected"), + [ + # 3 p0 + 0.5 * 1.1696 * 12000 = 1.1696 * 12000; no minimum binds. + (0.5, 2339.2), + # The single 1000-cost buyer sits at the 1000 minimum: + # 1000 + 2 p0 + 1.5 * 1.1696 * 11000 = 1.1696 * 12000. + (1.5, -3131.6), + ], +) +def test_private_premium_intercept_concrete_values( + coeff: float, expected: float +) -> None: + """The intercept solves mean premium = (1 + markup) * mean insurer cost.""" + intercept = health_insurance.private_premium_intercept( + premium_predicted_hcc=jnp.asarray(coeff), + premium_markup=_MARKUP, + premium_minimum=_MINIMUM, + premium_sample_insurer_cost=_SAMPLE_COST, + premium_sample_is_married=_SAMPLE_MARRIED, + ) + np.testing.assert_allclose(intercept, expected, atol=1e-6) + + +def test_private_premium_intercept_zero_profit_on_sample() -> None: + """At the solved intercept, premium revenue covers insurer cost plus markup. + + The sample mean of `max(minimum[married], p0 + b (1 + markup) cost)` + equals `(1 + markup)` times the sample mean of the insurer cost. + """ + rng = np.random.default_rng(seed=3) + cost = jnp.asarray(np.exp(rng.normal(loc=7.5, scale=1.2, size=2000))) + married = jnp.asarray(rng.integers(0, 2, size=2000), dtype=jnp.int32) + coeff = jnp.asarray(1.161245437600243) + intercept = health_insurance.private_premium_intercept( + premium_predicted_hcc=coeff, + premium_markup=_MARKUP, + premium_minimum=_MINIMUM, + premium_sample_insurer_cost=cost, + premium_sample_is_married=married, + ) + premiums = jnp.maximum(_MINIMUM[married], intercept + coeff * (1 + _MARKUP) * cost) + assert bool((premiums == _MINIMUM[married]).any()) + np.testing.assert_allclose(premiums.mean(), (1 + _MARKUP) * cost.mean(), rtol=1e-10) + + +def test_private_premium_intercept_without_root_in_bracket_is_nan() -> None: + """No intercept in [-50000, 10000] balances the books: the result is NaN.""" + intercept = health_insurance.private_premium_intercept( + premium_predicted_hcc=jnp.asarray(100.0), + premium_markup=_MARKUP, + premium_minimum=_MINIMUM, + premium_sample_insurer_cost=_SAMPLE_COST, + premium_sample_is_married=_SAMPLE_MARRIED, + ) + assert bool(jnp.isnan(intercept)) + + +@pytest.mark.parametrize( + ("is_married", "insurer_cost", "expected"), + [ + (0, 1913.49, 1529.488080), + (1, 1913.49, 2000.0), + (0, 371.13, 1000.0), + (1, 8319.48, 10230.056703), + ], +) +def test_premium_private_buyer( + is_married: int, insurer_cost: float, expected: float +) -> None: + """Premium is `max(minimum[married], p0 + b (1 + markup) E[insurer cost])`. + + No marital-status shifter: couples differ only through the minimum and + their insurer cost. + """ + result = health_insurance.premium( + buy_private=jnp.array(BuyPrivate.yes), + is_married=jnp.int32(is_married), + predicted_hcc_insurer=jnp.asarray(insurer_cost), + private_premium_intercept=jnp.asarray(-1069.4), + premium_predicted_hcc=jnp.asarray(1.161245437600243), + premium_markup=_MARKUP, + premium_minimum=_MINIMUM, + ) + np.testing.assert_allclose(result, expected, atol=1e-6) + + +def test_premium_uninsured_pays_nothing() -> None: + """A household that does not buy private cover pays no premium.""" + result = health_insurance.premium( + buy_private=jnp.array(BuyPrivate.no), + is_married=jnp.int32(0), + predicted_hcc_insurer=jnp.asarray(1913.49), + private_premium_intercept=jnp.asarray(-1069.4), + premium_predicted_hcc=jnp.asarray(1.161245437600243), + premium_markup=_MARKUP, + premium_minimum=_MINIMUM, + ) + assert float(result) == 0.0 + + +@pytest.mark.parametrize( + ("buy_private", "premium_default", "expected"), + [ + (BuyPrivate.yes, 0.0, True), + (BuyPrivate.yes, 0.01, False), + (BuyPrivate.no, 500.0, True), + (BuyPrivate.no, 0.0, True), + ], +) +def test_private_cover_paid_in_full( + buy_private: int, premium_default: float, expected: bool +) -> None: + """Private cover is a valid choice only if the household pays its premium. + + A household at the consumption floor cannot hold private cover financed by + transfers or by not paying the premium. + """ + result = health_insurance.private_cover_paid_in_full( + buy_private=jnp.array(buy_private), + premium_default=jnp.asarray(premium_default), + ) + assert bool(result) is expected + + +def test_predicted_hcc_insurer_reads_age_marriage_health_cell() -> None: + """Predicted insurer cost interpolates the (period, married, health) row.""" + table = jnp.zeros((2, 2, 2, 3)).at[1, 1, 0].set(jnp.array([100.0, 200.0, 400.0])) + result = health_insurance.hcc_insurer_predicted( + period=jnp.int32(1), + is_married=jnp.int32(1), + good_health=jnp.int32(0), + hcc_persistent=jnp.asarray(0.5), + predicted_hcc_insurer_table=table, + hcc_persistent_grid=jnp.array([-1.0, 0.0, 1.0]), + ) + np.testing.assert_allclose(result, 300.0, atol=1e-9) + + +@pytest.mark.parametrize( + ("is_married", "total", "expected"), + [ + (0, 60000.0, 7472.3), + (1, 60000.0, 2652.2 + (60000.0 - 2652.2) * 0.2120), + (1, 2000.0, 2000.0), + (1, 200000.0, 14944.6), + ], +) +def test_insured_oop_doubles_couples_deductible_and_oop_max( + is_married: int, total: float, expected: float +) -> None: + """Couples' deductible and OOP maximum are twice the single ones.""" + result = health_insurance.insured_oop( + total_health_costs=jnp.asarray(total), + is_married=jnp.int32(is_married), + deductible=jnp.asarray(1326.1), + coinsurance_rate=jnp.asarray(0.2120), + oop_max=jnp.asarray(7472.3), + ) + np.testing.assert_allclose(result, expected, atol=1e-6) + + +@pytest.mark.parametrize( + ("buy_private", "expected"), + [(BuyPrivate.yes, 2652.2 + (20000.0 - 2652.2) * 0.2120), (BuyPrivate.no, 20000.0)], +) +def test_primary_oop_married_private_buyer(buy_private: int, expected: float) -> None: + """A married private buyer faces the doubled deductible; uninsured pay all.""" + result = health_insurance.primary_oop( + total_health_costs=jnp.asarray(20000.0), + buy_private=jnp.array(buy_private), + is_married=jnp.int32(1), + deductible=jnp.asarray(1326.1), + coinsurance_rate=jnp.asarray(0.2120), + oop_max=jnp.asarray(7472.3), + ) + np.testing.assert_allclose(result, expected, atol=1e-6) + + +@pytest.mark.parametrize( + ("regime", "expected"), + [ + ("nongroup_nomc_choose_canwork", ("private_cover_paid_in_full",)), + ("nongroup_dimc_choose_canwork", ()), + ("nongroup_oamc_choose_canwork", ()), + ], +) +def test_nongroup_regimes_constrain_private_cover_to_paid_premiums( + regime: str, expected: tuple[str, ...] +) -> None: + """Regimes offering private cover require its premium to be paid in full.""" + constraints = _nongroup.build_constraints(REGIME_SPECS[regime]) + assert {name: func.__name__ for name, func in constraints.items()} == { + name: name for name in expected + } + + +@pytest.mark.parametrize( + ("is_married", "expected"), + [ + (0, 1109.843 + (10000.0 - 1109.843) * 0.1774133), + (1, 2219.686 + (10000.0 - 2219.686) * 0.1774133), + ], +) +def test_aca_primary_oop_doubles_couples_deductible_and_oop_max( + is_married: int, expected: float +) -> None: + """Under the ACA plan too, couples' deductible and OOP maximum double.""" + result = aca_health_insurance.primary_oop( + total_health_costs=jnp.asarray(10000.0), + cost_sharing_scale=jnp.asarray(1.0), + buy_private=jnp.array(BuyPrivate.yes), + is_married=jnp.int32(is_married), + deductible=jnp.asarray(1109.843), + coinsurance_rate=jnp.asarray(0.1774133), + oop_max=jnp.asarray(6500.0), + ) + np.testing.assert_allclose(result, expected, atol=1e-6) diff --git a/tests/test_regime_transitions.py b/tests/test_regime_transitions.py index 9040fdc..f7baf11 100644 --- a/tests/test_regime_transitions.py +++ b/tests/test_regime_transitions.py @@ -26,7 +26,12 @@ N_REGIMES = 19 N_PERIODS = MODEL_CONFIG.end_age - MODEL_CONFIG.start_age -SURVIVAL = jnp.ones(N_PERIODS) * 0.99 +SURVIVAL = jnp.ones((N_PERIODS, 3)) * 0.99 +# P(disabled next period | disabled, bad, good health). +PROB_DISABLED = jnp.broadcast_to( + jnp.array([0.98, 0.05, 0.002]), + (N_PERIODS, 3), +) def _target_from_probs(probs: jnp.ndarray) -> int: @@ -42,13 +47,15 @@ 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(gets_medicare=False, own=own, ng=ng) + transition = tied_canwork(own=own, ng=ng) probs = transition( age=jnp.int32(55), period=jnp.int32(4), + health=jnp.int32(2), + prob_disabled_next=PROB_DISABLED, labor_supply=jnp.array(LaborSupply.do_not_work), - is_medicaid_eligible=jnp.array(False), + is_ssi_eligible=jnp.array(False), survival_probs=SURVIVAL, ) target = _target_from_probs(probs) @@ -59,13 +66,15 @@ 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(gets_medicare=False, own=own, ng=ng) + transition = tied_canwork(own=own, ng=ng) probs = transition( age=jnp.int32(55), period=jnp.int32(4), + health=jnp.int32(2), + prob_disabled_next=PROB_DISABLED, labor_supply=jnp.array(LaborSupply.h2000), - is_medicaid_eligible=jnp.array(False), + is_ssi_eligible=jnp.array(False), survival_probs=SURVIVAL, ) target = _target_from_probs(probs) @@ -78,13 +87,15 @@ 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(gets_medicare=False, own=own, ng=ng) + transition = retiree_canwork(own=own, ng=ng) probs = transition( age=jnp.int32(55), period=jnp.int32(4), + health=jnp.int32(2), + prob_disabled_next=PROB_DISABLED, labor_supply=jnp.array(LaborSupply.h2000), - is_medicaid_eligible=jnp.array(True), + is_ssi_eligible=jnp.array(True), survival_probs=SURVIVAL, ) target = _target_from_probs(probs) @@ -94,13 +105,15 @@ 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(gets_medicare=False, own=own, ng=ng) + transition = retiree_canwork(own=own, ng=ng) probs = transition( age=jnp.int32(55), period=jnp.int32(4), + health=jnp.int32(2), + prob_disabled_next=PROB_DISABLED, labor_supply=jnp.array(LaborSupply.h2000), - is_medicaid_eligible=jnp.array(False), + is_ssi_eligible=jnp.array(False), survival_probs=SURVIVAL, ) target = _target_from_probs(probs) @@ -115,7 +128,8 @@ def test_retiree_forcedout_medicaid_override() -> None: probs = transition( age=jnp.int32(80), period=jnp.int32(29), - is_medicaid_eligible=jnp.array(True), + health=jnp.int32(2), + is_ssi_eligible=jnp.array(True), survival_probs=SURVIVAL, ) target = _target_from_probs(probs) @@ -148,14 +162,16 @@ 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(gets_medicare=False, own=own, ng=ng) + transition = retiree_canwork(own=own, ng=ng) period = jnp.int32(age - MODEL_CONFIG.start_age) probs = transition( age=jnp.asarray(age), period=period, + health=jnp.int32(2), + prob_disabled_next=PROB_DISABLED, labor_supply=jnp.array(LaborSupply.h2000), - is_medicaid_eligible=jnp.array(False), + is_ssi_eligible=jnp.array(False), survival_probs=SURVIVAL, ) target = _target_from_probs(probs) @@ -168,11 +184,13 @@ 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(gets_medicare=False, own=own) + transition = nongroup_canwork(own=own) probs = transition( age=jnp.int32(55), period=jnp.int32(4), + health=jnp.int32(2), + prob_disabled_next=PROB_DISABLED, labor_supply=jnp.array(LaborSupply.h2000), survival_probs=SURVIVAL, ) @@ -188,6 +206,7 @@ def test_nongroup_forcedout_valid_probs() -> None: probs = transition( age=jnp.int32(80), period=jnp.int32(29), + health=jnp.int32(2), survival_probs=SURVIVAL, ) assert jnp.isclose(jnp.sum(probs), 1.0, atol=1e-6) @@ -200,13 +219,15 @@ 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(gets_medicare=False, own=own, ng=ng) + transition = tied_canwork(own=own, ng=ng) probs = transition( age=jnp.int32(55), period=jnp.int32(4), + health=jnp.int32(2), + prob_disabled_next=PROB_DISABLED, labor_supply=jnp.array(LaborSupply.h2000), - is_medicaid_eligible=jnp.array(True), + is_ssi_eligible=jnp.array(True), survival_probs=SURVIVAL, ) target = _target_from_probs(probs) @@ -216,14 +237,16 @@ 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(gets_medicare=False, own=own, ng=ng) + transition = tied_canwork(own=own, ng=ng) period = jnp.int32(64 - MODEL_CONFIG.start_age) probs = transition( age=jnp.int32(64), period=period, + health=jnp.int32(2), + prob_disabled_next=PROB_DISABLED, labor_supply=jnp.array(LaborSupply.h2000), - is_medicaid_eligible=jnp.array(True), + is_ssi_eligible=jnp.array(True), survival_probs=SURVIVAL, ) target = _target_from_probs(probs) @@ -234,14 +257,16 @@ 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(gets_medicare=False, own=own, ng=ng) + transition = retiree_canwork(own=own, ng=ng) - survival = jnp.ones(N_PERIODS) * 0.85 + survival = jnp.ones((N_PERIODS, 3)) * 0.85 probs = transition( age=jnp.int32(55), period=jnp.int32(4), + health=jnp.int32(2), + prob_disabled_next=PROB_DISABLED, labor_supply=jnp.array(LaborSupply.h2000), - is_medicaid_eligible=jnp.array(False), + is_ssi_eligible=jnp.array(False), survival_probs=survival, ) assert jnp.isclose(probs[RegimeId.dead], 0.15, atol=1e-6) diff --git a/tests/test_simulation.py b/tests/test_simulation.py index fd50755..5abf531 100644 --- a/tests/test_simulation.py +++ b/tests/test_simulation.py @@ -1,7 +1,11 @@ """Tests for the pylcm simulate-interaction helpers.""" +from types import SimpleNamespace + import numpy as np import pandas as pd +import pytest +from lcm import InitialNodes from aca_model.simulation import restore_subject_ids, simulate_with_dense_index @@ -11,6 +15,10 @@ class _StubModel: def __init__(self) -> None: self.seen_index: list[int] = [] + self.initial_nodes = InitialNodes(by_age={51: ("a", "b", "c"), 52: "a"}) + self.graph = SimpleNamespace( + initial_nodes=frozenset({(51, "a"), (51, "b"), (51, "c"), (52, "a")}) + ) def simulate(self, *, initial_conditions: pd.DataFrame, **_kwargs: object) -> str: self.seen_index = list(initial_conditions.index) @@ -24,7 +32,7 @@ def test_simulate_with_dense_index_hands_pylcm_a_dense_range() -> None: `subject_id` positionally, so a sparse id index would index out of bounds. """ ic = pd.DataFrame( - {"regime_name": ["a", "b", "c"]}, + {"regime_name": ["a", "b", "c"], "age": [51, 51, 51]}, index=pd.Index([3010, 500_000_000, 959_738_010], name="id"), ) model = _StubModel() @@ -35,7 +43,7 @@ def test_simulate_with_dense_index_hands_pylcm_a_dense_range() -> None: def test_simulate_with_dense_index_returns_ids_aligned_to_subject() -> None: """Returned ids align to subject_id: position i is the id of subject i.""" ic = pd.DataFrame( - {"regime_name": ["a", "b"]}, + {"regime_name": ["a", "b"], "age": [51, 51]}, index=pd.Index([3010, 959_738_010], name="id"), ) _result, ids = simulate_with_dense_index( @@ -52,3 +60,30 @@ def test_restore_subject_ids_maps_positional_subject_back_to_original() -> None: np.testing.assert_array_equal( out["id"].to_numpy(), [3010, 959_738_010, 3010, 959_738_010] ) + + +def test_simulate_with_dense_index_drops_starts_after_the_last_entry_age() -> None: + """Rows older than every admissible start age never reach the simulator.""" + ic = pd.DataFrame( + {"regime_name": ["a", "b", "a"], "age": [51, 62, 52]}, + index=pd.Index([10, 20, 30], name="id"), + ) + _result, ids = simulate_with_dense_index( + model=_StubModel(), initial_conditions=ic, params={} + ) + np.testing.assert_array_equal(ids, [10, 30]) + + +@pytest.mark.parametrize( + ("age", "regime"), [(51, "dead"), (52, "b")], ids=["dead", "off-contract"] +) +def test_simulate_with_dense_index_rejects_inadmissible_pairs_within_entry_ages( + age: int, regime: str +) -> None: + """A row inside the entry ages but outside the entry pairs raises.""" + ic = pd.DataFrame( + {"regime_name": ["a", regime], "age": [51, age]}, + index=pd.Index([10, 20], name="id"), + ) + with pytest.raises(ValueError, match=regime): + simulate_with_dense_index(model=_StubModel(), initial_conditions=ic, params={}) diff --git a/tests/test_social_security.py b/tests/test_social_security.py index 2ee11a0..68e193b 100644 --- a/tests/test_social_security.py +++ b/tests/test_social_security.py @@ -6,9 +6,10 @@ import jax.numpy as jnp import numpy as np import pandas as pd -from helpers.social_security import ( # ty: ignore[unresolved-import] +from helpers.social_security import ( compute_di_dropout_scale, compute_pia_table, + next_aime_from_state, ) from lcm.typing import ScalarInt @@ -154,7 +155,7 @@ def test_pia_to_aime_above_kink_1() -> None: def test_next_aime_indexing_high_income() -> None: - result = social_security.next_aime( + result = next_aime_from_state( claim_ss=jnp.array(ClaimedSS.no), claimed_ss=jnp.array(ClaimedSS.no), normal_retirement_age=NORMAL_RETIREMENT_AGE, @@ -179,7 +180,7 @@ def test_next_aime_indexing_high_income() -> None: def test_next_aime_indexing_low_income() -> None: - result = social_security.next_aime( + result = next_aime_from_state( claim_ss=jnp.array(ClaimedSS.no), claimed_ss=jnp.array(ClaimedSS.no), normal_retirement_age=NORMAL_RETIREMENT_AGE, @@ -203,7 +204,7 @@ def test_next_aime_indexing_low_income() -> None: def test_next_aime_no_indexing_high_income() -> None: - result = social_security.next_aime( + result = next_aime_from_state( claim_ss=jnp.array(ClaimedSS.no), claimed_ss=jnp.array(ClaimedSS.no), normal_retirement_age=NORMAL_RETIREMENT_AGE, @@ -228,7 +229,7 @@ def test_next_aime_no_indexing_high_income() -> None: def test_next_aime_no_indexing_low_income() -> None: - result = social_security.next_aime( + result = next_aime_from_state( claim_ss=jnp.array(ClaimedSS.no), claimed_ss=jnp.array(ClaimedSS.no), normal_retirement_age=NORMAL_RETIREMENT_AGE, @@ -251,15 +252,9 @@ def test_next_aime_no_indexing_low_income() -> None: assert jnp.isclose(result, 1000, atol=ATOL) -def test_next_aime_high_aime_high_income_accrues_above_taxable_max() -> None: - """Within-period labor accrual above the taxable-max-indexed base is preserved. - - `_accrue_aime` caps the *indexed* base at the taxable max; the small extra - accrual from current labor earnings then rides on top. The carried AIME is - the round-trip of that accrued PIA — not re-clamped at the taxable max — so - it lands just above `aime_kink_2`. - """ - result = social_security.next_aime( +def test_next_aime_high_aime_high_income_is_capped_at_taxable_max() -> None: + """Labor accrual cannot lift an AIME at the taxable max above it.""" + result = next_aime_from_state( claim_ss=jnp.array(ClaimedSS.no), claimed_ss=jnp.array(ClaimedSS.no), normal_retirement_age=NORMAL_RETIREMENT_AGE, @@ -279,14 +274,11 @@ def test_next_aime_high_aime_high_income_accrues_above_taxable_max() -> None: aime_kink_2=AIME_KINK_2_SCALAR, ratio_lowest_earnings=RATIO, ) - capped_base = AIME_KINK_2 - lowest_year = _RATIO_NP[62] * capped_base - expected = capped_base + (20000.0 - lowest_year) * float(AIME_ACCRUAL_FACTOR) - assert jnp.isclose(result, expected, atol=ATOL) + assert jnp.isclose(result, AIME_KINK_2, atol=ATOL) def test_next_aime_cap_high_aime_low_income() -> None: - result = social_security.next_aime( + result = next_aime_from_state( claim_ss=jnp.array(ClaimedSS.no), claimed_ss=jnp.array(ClaimedSS.no), normal_retirement_age=NORMAL_RETIREMENT_AGE, @@ -503,7 +495,7 @@ def _next_aime_claiming( labor_income: float = 0.0, ) -> jnp.ndarray: """Advance AIME one period under a given claim state, no labor accrual.""" - return social_security.next_aime( + return next_aime_from_state( aime=aime, labor_income=jnp.array(labor_income), period=jnp.int32(age), @@ -580,7 +572,7 @@ def test_delayed_claim_at_taxable_max_carries_credit_above_max() -> None: taxable max instead of clamping there, so the credit is carried forward. """ aime_at_max = AIME_KINK_2_SCALAR - result = social_security.next_aime( + result = next_aime_from_state( claim_ss=jnp.array(ClaimedSS.no), claimed_ss=jnp.array(ClaimedSS.no), normal_retirement_age=NORMAL_RETIREMENT_AGE, diff --git a/tests/test_ss_benefit_integration.py b/tests/test_ss_benefit_integration.py index 0b70990..e111b96 100644 --- a/tests/test_ss_benefit_integration.py +++ b/tests/test_ss_benefit_integration.py @@ -5,7 +5,7 @@ """ import jax.numpy as jnp -from helpers.social_security import compute_pia_table # ty: ignore[unresolved-import] +from helpers.social_security import compute_pia_table from aca_model.agent.labor_market import LaborSupply from aca_model.environment import social_security diff --git a/tests/test_ssi_medicaid_integration.py b/tests/test_ssi_medicaid_integration.py index 20572f9..75a41f4 100644 --- a/tests/test_ssi_medicaid_integration.py +++ b/tests/test_ssi_medicaid_integration.py @@ -37,7 +37,7 @@ def test_low_income_qualifies_for_ssi_and_medicaid() -> None: result = combined( labor_income=jnp.array(0.0), capital_income=jnp.array(0.0), - spousal_income_amounts=jnp.array([0.0, 0.0, 20000.0]), + spousal_income_amount=jnp.array(0.0), ss_benefit=jnp.array(500.0), pension_benefit=jnp.array(0.0), ssi_ignored_overall=jnp.asarray(20.0), @@ -69,7 +69,7 @@ def test_high_income_ineligible_for_ssi() -> None: result = combined( labor_income=jnp.array(50000.0), capital_income=jnp.array(5000.0), - spousal_income_amounts=jnp.array([0.0, 0.0, 20000.0]), + spousal_income_amount=jnp.array(0.0), ss_benefit=jnp.array(2000.0), pension_benefit=jnp.array(0.0), ssi_ignored_overall=jnp.asarray(20.0), @@ -105,7 +105,7 @@ def test_disabled_under_65_qualifies_for_ssi_and_medicaid() -> None: result = combined( labor_income=jnp.array(0.0), capital_income=jnp.array(0.0), - spousal_income_amounts=jnp.array([0.0, 0.0, 20000.0]), + spousal_income_amount=jnp.array(0.0), ss_benefit=jnp.array(0.0), pension_benefit=jnp.array(0.0), ssi_ignored_overall=jnp.asarray(20.0), @@ -144,7 +144,7 @@ def test_aca_expansion_uses_magi_not_countable_income() -> None: # but full MAGI 28000 > 15000 threshold. labor_income=jnp.array(28000.0), capital_income=jnp.array(0.0), - spousal_income_amounts=jnp.array([0.0, 0.0, 20000.0]), + spousal_income_amount=jnp.array(0.0), ss_benefit=jnp.array(0.0), pension_benefit=jnp.array(0.0), ssi_ignored_overall=jnp.asarray(20.0), @@ -175,6 +175,7 @@ def test_medicaid_reduces_oop() -> None: oop_medicaid = combined( total_health_costs=jnp.array(10000.0), buy_private=jnp.array(BuyPrivate.yes), + is_married=jnp.int32(0), deductible=jnp.asarray(500.0), coinsurance_rate=jnp.asarray(0.2), oop_max=jnp.asarray(5000.0), @@ -188,6 +189,7 @@ def test_medicaid_reduces_oop() -> None: oop_no_medicaid = combined( total_health_costs=jnp.array(10000.0), buy_private=jnp.array(BuyPrivate.yes), + is_married=jnp.int32(0), deductible=jnp.asarray(500.0), coinsurance_rate=jnp.asarray(0.2), oop_max=jnp.asarray(5000.0), diff --git a/tests/test_taxes.py b/tests/test_taxes.py index 862c024..d9cab16 100644 --- a/tests/test_taxes.py +++ b/tests/test_taxes.py @@ -221,7 +221,7 @@ def test_taxable_ss_benefit_below_threshold() -> None: result = taxes.taxable_ss_benefit( capital_income=jnp.array(0.0), labor_income=jnp.array(10000.0), - spousal_income_amounts=jnp.array([0.0, 0.0, 20000.0]), + spousal_income_amount=jnp.array(0.0), ss_benefit=jnp.array(5000.0), pension_benefit=jnp.array(0.0), spousal_income=jnp.int32(0), @@ -235,8 +235,7 @@ def test_gross_income_basic() -> None: result = taxes.gross_income( capital_income=jnp.array(1000.0), labor_income=jnp.array(5000.0), - spousal_income=jnp.int32(1), - spousal_income_amounts=jnp.array([0.0, 2000.0, 20000.0]), + spousal_income_amount=jnp.array(2000.0), taxable_ss_benefit=jnp.array(500.0), pension_benefit=jnp.array(300.0), ) @@ -247,8 +246,7 @@ def test_after_tax_income_zero() -> None: gi = taxes.gross_income( capital_income=jnp.array(0.0), labor_income=jnp.array(0.0), - spousal_income=jnp.int32(0), - spousal_income_amounts=jnp.array([0.0, 0.0, 20000.0]), + spousal_income_amount=jnp.array(0.0), taxable_ss_benefit=jnp.array(0.0), pension_benefit=jnp.array(0.0), ) @@ -269,8 +267,7 @@ def test_after_tax_income_low_bracket() -> None: gi = taxes.gross_income( capital_income=jnp.array(0.0), labor_income=jnp.array(gross), - spousal_income=jnp.int32(0), - spousal_income_amounts=jnp.array([0.0, 0.0, 20000.0]), + spousal_income_amount=jnp.array(0.0), taxable_ss_benefit=jnp.array(0.0), pension_benefit=jnp.array(0.0), ) diff --git a/tests/test_transition_only_medical_shock.py b/tests/test_transition_only_medical_shock.py new file mode 100644 index 0000000..a4a148f --- /dev/null +++ b/tests/test_transition_only_medical_shock.py @@ -0,0 +1,80 @@ +"""Folding the transitory medical-cost shock takes it out of every state space. + +Nothing a period computes reads last year's transitory shock: this year's draw +enters the out-of-pocket bill inside the end-of-period asset law, and the draw is +IID. With `fold_hcc_transitory`, the shock is therefore drawn on the transition +only. Living regimes keep the lagged persistent shock, which prices the private +premium, and lose the transitory axis; values agree with the unfolded model on +every transitory slice. +""" + +from dataclasses import replace + +import numpy as np +import pytest +from helpers.model import make_baseline_model +from lcm import Model + +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 +) + + +@pytest.fixture(scope="module") +def folded_model() -> Model: + return make_baseline_model(grid_config=_FOLDED) + + +def test_living_regimes_carry_no_lagged_transitory_shock(folded_model: Model) -> None: + assert not any( + "hcc_transitory" in folded_model.state_names(regime_name=name) + for name in folded_model.user_regimes + ) + + +def test_living_regimes_keep_the_lagged_persistent_shock(folded_model: Model) -> None: + assert all( + "hcc_persistent" in folded_model.state_names(regime_name=name) + for name in folded_model.user_regimes + if name != "dead" + ) + + +@pytest.mark.long_running +def test_folded_value_equals_every_transitory_slice_of_the_unfolded_value( + folded_model: Model, +) -> None: + """The two models' sums differ only in order, so they agree to rounding.""" + unfolded_model = make_baseline_model( + grid_config=replace( + BENCHMARK_GRID_CONFIG, fold_hcc_transitory=False, subtract_hcc_bill=False + ) + ) + _, _, params = get_benchmark_params(model=unfolded_model) + unfolded_values = unfolded_model.solve(params=params, log_level="off").values + folded_values = folded_model.solve(params=params, log_level="off").values + assert set(folded_values) == set(unfolded_values) + for period, regimes in unfolded_values.items(): + assert set(folded_values[period]) == set(regimes) + for name, values in regimes.items(): + unfolded = np.asarray(values) + folded = np.asarray(folded_values[period][name]) + states = unfolded_model.state_names(regime_name=name) + by_slice = ( + np.moveaxis(unfolded, states.index("hcc_transitory"), 0) + if "hcc_transitory" in states + else unfolded + ) + expected = np.broadcast_to(folded, by_slice.shape) + np.testing.assert_array_equal(np.isfinite(by_slice), np.isfinite(expected)) + np.testing.assert_allclose( + by_slice, + expected, + rtol=1e-6, + atol=0, + equal_nan=False, + err_msg=f"period={period}, regime={name}", + )