diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 632668d..897c37f 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -30,12 +30,12 @@ jobs: - uses: actions/setup-python@v6 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 revision aca-dev pins + # Pinned by commit so a deleted or rewritten branch cannot break the + # install; move it together with the aca-dev pylcm submodule pointer. run: >- pip install "pylcm @ - git+https://github.com/OpenSourceEconomics/pylcm.git@feat/nb-egm" + git+https://github.com/OpenSourceEconomics/pylcm.git@7e8e4c223a63a07725015ebdd5af84cd9092cee2" - name: Install aca-model with test deps run: pip install -e . pytest pdbp - name: Run pytest diff --git a/README.md b/README.md index 5f52c1a..9edf892 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,43 @@ # aca-model Core lifecycle model for the ACA structural retirement project. + +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/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/aca/model.py b/src/aca_model/aca/model.py index 0db5ff6..5cc6139 100644 --- a/src/aca_model/aca/model.py +++ b/src/aca_model/aca/model.py @@ -7,7 +7,7 @@ from collections.abc import Mapping from typing import Any -from lcm import AgeGrid, DiscreteGrid, Model +from lcm import AgeGrid, DiscreteGrid, ExecutionConfig, Model from lcm.typing import UserParams from aca_model.aca import PolicyVariant @@ -15,11 +15,11 @@ 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.execution import execution_config_for_devices def create_model( *, - n_subjects: int, policy: PolicyVariant, fixed_params: UserParams, wage_params: Mapping[str, Any], @@ -28,11 +28,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,6 +53,9 @@ 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. @@ -92,6 +95,10 @@ def create_model( 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/agent/preferences.py b/src/aca_model/agent/preferences.py index 438b2c4..3bb4b80 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( diff --git a/src/aca_model/baseline/model.py b/src/aca_model/baseline/model.py index 87acf7c..6040c45 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,7 +13,7 @@ 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 ( @@ -23,11 +23,11 @@ build_model_slots, ) from aca_model.config import MODEL_CONFIG, 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 +35,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,6 +69,9 @@ 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 @@ -106,7 +109,11 @@ def create_model( 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..dbbb5aa 100644 --- a/src/aca_model/baseline/regimes/__init__.py +++ b/src/aca_model/baseline/regimes/__init__.py @@ -36,11 +36,6 @@ 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__ = [ "REGIME_SPECS", "RegimeId", @@ -115,7 +110,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 @@ -136,7 +131,7 @@ def build_model_slots( 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, @@ -146,7 +141,7 @@ def build_model_slots( ) return { "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(), } diff --git a/src/aca_model/baseline/regimes/_common.py b/src/aca_model/baseline/regimes/_common.py index 750bce4..6fb7e5b 100644 --- a/src/aca_model/baseline/regimes/_common.py +++ b/src/aca_model/baseline/regimes/_common.py @@ -14,18 +14,24 @@ import numpy as np from lcm import ( DiscreteGrid, + GridBreakpoint, IrregSpacedGrid, LinSpacedGrid, MarkovTransition, NormalIIDProcess, Phased, - PiecewiseGridSegment, PiecewiseLinSpacedGrid, Regime, RouwenhorstAR1Process, categorical, fixed_transition, ) +from lcm.consumption_savings_regime import ( + ConsumptionSavingsRegime, + LiquidMargin, + post_decision_lower_bound, +) +from lcm.solvers import OneMarginSolver from lcm.typing import BoolND, FloatND, IntND, RegimeName, ScalarInt, UserParams from aca_model.agent import ( @@ -44,6 +50,34 @@ 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: @@ -203,10 +237,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 @@ -270,7 +301,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( @@ -289,9 +319,8 @@ def build_grids( start=assets_start, stop=500_000.0, n_points=grid_config.n_assets_gridpoints, - batch_size=grid_config.n_assets_batch_size, ), - 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) @@ -327,9 +356,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 +367,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, ) @@ -432,49 +452,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 @@ -541,8 +552,8 @@ def build_dead_regime(*, solver: SolverName = "brute_force") -> Regime: 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.) + consumption action. It is broadcast under every solver, so there is + always exactly one mask to apply. - `pension_wealth` is masked explicitly: a carried state is rejected in terminal regimes before pruning could drop it. """ @@ -551,7 +562,7 @@ 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}, @@ -667,16 +678,21 @@ 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. + Grid search evaluates the action-level predicate directly. 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. """ - 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) + ) + return {"borrowing_constraint": borrowing_constraint} def build_model_states(grids: Grids) -> dict: @@ -684,10 +700,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 +714,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, } 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..f13e3c9 100644 --- a/src/aca_model/baseline/regimes/_nongroup.py +++ b/src/aca_model/baseline/regimes/_nongroup.py @@ -4,7 +4,6 @@ Already nongroup, so no SSI/Medicaid override needed for HIS transitions. """ -import functools from collections.abc import Callable from lcm import Regime @@ -13,12 +12,12 @@ 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, @@ -78,32 +77,11 @@ def transition( 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) -> 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["ss_benefit"] = select_ss_benefit(spec) # his and crossed_oamc_threshold are fixed params (constants per regime), @@ -117,14 +95,6 @@ def _build_functions( 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.update(build_pension_functions(spec)) return functions @@ -150,24 +120,14 @@ def build_regime( 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 - ) + functions = _build_functions(spec) constraints: dict = {} - if fix_for_nbegm: + 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 +135,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( + return build_alive_regime( + egm_solver=egm_solver, transition=build_granular_regime_transition( transition_func=transition_func, target_ids=own.values() ), active=make_active_func(spec), states=states, state_transitions=build_state_transitions(spec, solver=state_solver), - actions=build_actions( - spec, - grids, - drop_buy_private=fix_for_nbegm, - drop_labor_supply=fix_labor, - ), + 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..309a71c 100644 --- a/src/aca_model/baseline/regimes/_retiree.py +++ b/src/aca_model/baseline/regimes/_retiree.py @@ -18,8 +18,10 @@ Grids, RegimeSpec, build_actions, + build_alive_regime, build_common_functions, build_granular_regime_transition, + build_nbegm_functions, build_pension_functions, build_regime_probs, build_state_transitions, @@ -127,13 +129,19 @@ def build_regime( 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( + functions = _build_functions(spec) + 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, transition=build_granular_regime_transition( transition_func=transition_func, target_ids=(*own.values(), *ng.values()) ), @@ -141,6 +149,5 @@ def build_regime( states=states, state_transitions=build_state_transitions(spec, solver=state_solver), 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..ff13664 100644 --- a/src/aca_model/baseline/regimes/_tied.py +++ b/src/aca_model/baseline/regimes/_tied.py @@ -19,8 +19,10 @@ Grids, RegimeSpec, build_actions, + build_alive_regime, build_common_functions, build_granular_regime_transition, + build_nbegm_functions, build_pension_functions, build_regime_probs, build_state_transitions, @@ -97,13 +99,19 @@ def build_regime( 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( + functions = _build_functions(spec) + 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, transition=build_granular_regime_transition( transition_func=transition_func, target_ids=(*own.values(), *ng.values()) ), @@ -111,6 +119,5 @@ def build_regime( states=states, state_transitions=build_state_transitions(spec, solver=state_solver), 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..54fe41c 100644 --- a/src/aca_model/benchmark.py +++ b/src/aca_model/benchmark.py @@ -29,7 +29,7 @@ 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 @@ -71,21 +71,19 @@ def create_benchmark_model( *, - n_subjects: int, pref_type_grid: DiscreteGrid, + 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)`. + 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( @@ -94,7 +92,7 @@ def create_benchmark_model( wage_params=wage_params, derived_categoricals=_DERIVED_CATEGORICALS, pref_type_grid=pref_type_grid, - n_subjects=n_subjects, + execution_config=execution_config, ) diff --git a/src/aca_model/config.py b/src/aca_model/config.py index 47390c2..4d65551 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 @@ -21,148 +21,27 @@ class ModelConfig: @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"`. + # 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() @@ -175,6 +54,4 @@ class GridConfig: 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/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/tests/helpers/model.py b/tests/helpers/model.py index 55571e9..4818827 100644 --- a/tests/helpers/model.py +++ b/tests/helpers/model.py @@ -26,11 +26,10 @@ } -def make_baseline_model(*, n_subjects: int) -> Model: +def make_baseline_model() -> Model: """Baseline model on `BENCHMARK_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, @@ -39,11 +38,10 @@ def make_baseline_model(*, n_subjects: int) -> Model: ) -def make_aca_model(*, n_subjects: int, policy: PolicyVariant) -> Model: +def make_aca_model(*, policy: PolicyVariant) -> Model: """ACA model on `BENCHMARK_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, wage_params=wage_params, 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_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..c3d54f6 100644 --- a/tests/test_benchmark.py +++ b/tests/test_benchmark.py @@ -2,7 +2,14 @@ import numpy as np import pytest -from lcm import DiscreteGrid +from lcm import ( + DiscreteGrid, + GridBreakpoint, + IrregSpacedGrid, + LinSpacedGrid, + PiecewiseLinSpacedGrid, + RouwenhorstAR1Process, +) from aca_model.agent.preferences import BenchmarkPrefType from aca_model.benchmark import ( @@ -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 == 45 + 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, LinSpacedGrid) + 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,7 +61,6 @@ 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", ) @@ -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,19 @@ 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]) diff --git a/tests/test_dcegm_model_creation.py b/tests/test_dcegm_model_creation.py index 2a4b8de..6670eba 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 @@ -45,7 +46,6 @@ 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, @@ -56,14 +56,15 @@ def _build_model(solver: SolverName) -> Model: 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) @@ -122,7 +123,6 @@ 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, @@ -143,27 +143,9 @@ 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, @@ -177,42 +159,6 @@ def test_dcegm_benchmark_model_builds() -> None: 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.""" 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_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..77bd52d --- /dev/null +++ b/tests/test_execution_config.py @@ -0,0 +1,122 @@ +"""Model factories expose pylcm's execution policy without changing economic grids.""" + +from functools import partial +from types import SimpleNamespace + +import jax +import pytest +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 = ( + partial(create_aca_model, policy=PolicyVariant.ACA) + if kind == "aca" + else create_model + ) + 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,)), None), + ], +) +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_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_model_creation.py b/tests/test_model_creation.py index 5b2415c..8e4398f 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,24 +51,24 @@ 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 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 def test_dead_regime_is_terminal() -> None: - model = make_baseline_model(n_subjects=1) + model = make_baseline_model() assert model.user_regimes["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 @@ -234,7 +233,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,7 +281,7 @@ 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 @@ -323,7 +322,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,53 +362,16 @@ 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`). `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"} ) @@ -419,6 +381,6 @@ def test_dead_regime_prunes_unused_broadcast_states() -> None: 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_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..f853fcc 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 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 @@ -40,6 +45,14 @@ _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( + BENCHMARK_GRID_CONFIG, nbegm_jump_read="bridged" +) + def _build_regimes(solver: SolverName) -> dict[str, Regime]: return build_all_regimes( @@ -53,7 +66,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,7 +76,7 @@ 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: @@ -76,31 +88,45 @@ def _grids() -> Grids: ) -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: @@ -120,30 +146,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 +201,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,68 +278,27 @@ 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") - ) - - bridged = dataclasses.replace(live_labor, nbegm_jump_read="bridged") - assert isinstance(_build_model_with("nbegm", bridged), Model) + _build_model_with("nbegm", BENCHMARK_GRID_CONFIG) - -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, wage_params=_WAGE_PARAMS, @@ -321,19 +324,14 @@ 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, wage_params=_WAGE_PARAMS, diff --git a/tests/test_nbegm_solve_validation.py b/tests/test_nbegm_solve_validation.py index 6d4449d..e8e6484 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. 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_social_security.py b/tests/test_social_security.py index 2ee11a0..30fc799 100644 --- a/tests/test_social_security.py +++ b/tests/test_social_security.py @@ -6,7 +6,7 @@ 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, ) 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