From 8f3c5f25853ebc9a860483ae115e6ff938a1417c Mon Sep 17 00:00:00 2001 From: Hans-Martin von Gaudecker Date: Fri, 23 Oct 2026 11:42:16 +0200 Subject: [PATCH] Split the model factories into a domain spec and a thin create_model `build_model_spec` returns the domain keyword arguments of `lcm.Model` for the baseline (`policy=None`) or an ACA policy, including the entry initial nodes and derived categoricals, and reads only `pia_aime_grid` from the parameters. `create_model(spec=, fixed_params=, **model_options)` forwards generic options unchanged; without `execution_config` pylcm's device default applies. The ACA factory module is gone and the benchmark forwards `**model_options`. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Hs8mfXPYyEa24KaAp6BwCW --- README.md | 30 ++-- src/aca_model/aca/model.py | 122 -------------- src/aca_model/aca/regimes/__init__.py | 6 +- src/aca_model/baseline/model.py | 186 +++++++++++++-------- src/aca_model/baseline/regimes/__init__.py | 17 +- src/aca_model/baseline/regimes/_common.py | 14 +- src/aca_model/benchmark.py | 40 ++--- tests/helpers/model.py | 40 ++--- tests/test_aime_grid.py | 6 +- tests/test_dcegm_model_creation.py | 53 +++--- tests/test_dcegm_parity.py | 27 +-- tests/test_execution_config.py | 149 +++++++++-------- tests/test_model_creation.py | 4 +- tests/test_nbegm_labor_live_validation.py | 15 +- tests/test_nbegm_model_creation.py | 57 ++++--- tests/test_nbegm_solve_validation.py | 15 +- 16 files changed, 351 insertions(+), 430 deletions(-) delete mode 100644 src/aca_model/aca/model.py diff --git a/README.md b/README.md index f22ff8f..271e353 100644 --- a/README.md +++ b/README.md @@ -10,11 +10,16 @@ to `create_benchmark_model`. Pass the same grid configuration to `get_benchmark_initial_conditions`; both routes retain identical seeded draws for their shared initial states. -Model factories accept `execution_config=ExecutionConfig(...)`. An omitted policy uses -the smallest JAX allocator `bytes_limit` among the selected accelerators as an explicit -per-device budget. This enables budget-aware width selection; CPU construction remains -unbudgeted. Missing accelerator limits require an explicit policy instead of silently -falling back to unbudgeted execution. +`build_model_spec(...)` returns the model's domain: regimes, ages, edges, broadcast +slots and the default entry `initial_nodes`. `policy=None` builds the baseline; a +`PolicyVariant` overlays its ACA schedules. +`create_model(spec=..., fixed_params=..., **model_options)` forwards every other +`lcm.Model` option unchanged, and an option the spec also holds, such as +`initial_nodes`, overrides it. `create_benchmark_model` takes the same options. + +Without `execution_config`, pylcm's default applies: the budget is the smallest selected +device pool limit less headroom, which requires a preallocated allocator pool +(`XLA_PYTHON_CLIENT_PREALLOCATE` not `false`); CPU runs unbudgeted. To combine a measured budget with device selection and sharding: @@ -29,14 +34,15 @@ execution_config = replace( ) ``` -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. +Pass this policy as `execution_config` to `create_model` or `create_benchmark_model`. +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. +Explicit policies pass through unchanged. In particular, +`ExecutionConfig(device_memory_bytes=None)` 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`, diff --git a/src/aca_model/aca/model.py b/src/aca_model/aca/model.py deleted file mode 100644 index e6b2022..0000000 --- a/src/aca_model/aca/model.py +++ /dev/null @@ -1,122 +0,0 @@ -"""ACA structural retirement model variants. - -Creates model variants for counterfactual ACA policy analysis by applying -function overrides on top of baseline regimes. -""" - -from collections.abc import Mapping -from typing import Any - -from lcm import AgeGrid, DiscreteGrid, ExecutionConfig, InitialNodes, Model -from lcm.typing import UserParams - -from aca_model.aca import PolicyVariant -from aca_model.aca.regimes import build_all_regimes -from aca_model.baseline.model import _fail_if_dcegm_without_consumption_points -from aca_model.baseline.regimes import ( - RegimeId, - SolverName, - build_entry_initial_nodes, - build_model_slots, -) -from aca_model.config import MODEL_CONFIG, GridConfig, model_ages -from aca_model.execution import execution_config_for_devices - - -def create_model( - *, - policy: PolicyVariant, - fixed_params: UserParams, - wage_params: Mapping[str, Any], - derived_categoricals: Mapping[str, DiscreteGrid], - grid_config: GridConfig, - pref_type_grid: DiscreteGrid, - solver: SolverName = "brute_force", - consumption_dollars_points: tuple[float, ...] | None = None, - execution_config: ExecutionConfig | None = None, - initial_nodes: InitialNodes | None = None, - final_age: int = MODEL_CONFIG.end_age, -) -> Model: - """Create an ACA policy variant model. - - Args: - policy: Which ACA policy combination to apply (e.g. - `PolicyVariant.ACA`). - fixed_params: Parameters to fix at model creation time. Pass - data-derived constants here; only estimation parameters - should go through `model.simulate(params=...)`. - wage_params: Data-derived wage profile dict (`log_ft_wage_mean`, - `log_ft_wage_std`, `adj_wage_hours_*`) used only at grid-build - time to size the assets-floor to `-max_annual_labor_income`. - Not routed to the pylcm Model. - derived_categoricals: Categorical mappings for `pd.Series` - fixed_params index levels that aren't model state/action - grids — `target_his`, `his`, `good_health`, `is_married`, - `pref_type`. - grid_config: Continuous-grid point counts. - pref_type_grid: Pref-type `DiscreteGrid`. - solver: `"brute_force"` (the default) or `"dcegm"`; see - `aca_model.baseline.model.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. - initial_nodes: Admissible starting `(age, regime)` pairs. None admits - the estimation sample's entry ages and regimes that lie before - `final_age`; pass `build_restart_initial_nodes` to restart - simulations mid-life. - final_age: First age at which everybody is dead; the model spans - `MODEL_CONFIG.start_age` through `final_age`. The default is the - full horizon. A shorter horizon reads a prefix of the age-indexed - data, and `fixed_params` must hold only what it reads (see - `select_horizon_params`). - - Returns: - pylcm Model. - - """ - ages = model_ages(final_age=final_age) - _fail_if_dcegm_without_consumption_points( - solver=solver, consumption_dollars_points=consumption_dollars_points - ) - regimes = build_all_regimes( - policy=policy, - grid_config=grid_config, - fixed_params=fixed_params, - wage_params=wage_params, - pref_type_grid=pref_type_grid, - solver=solver, - consumption_dollars_points=consumption_dollars_points, - ) - # The overlay swaps only regime-level functions; the broadcast slots - # are policy-invariant. - model_slots = build_model_slots( - grid_config=grid_config, - fixed_params=fixed_params, - wage_params=wage_params, - pref_type_grid=pref_type_grid, - ages=ages, - solver=solver, - ) - - return Model( - regimes=regimes, - initial_nodes=( - build_entry_initial_nodes(ages=ages) - if initial_nodes is None - else initial_nodes - ), - ages=AgeGrid(exact_values=ages), - regime_id_class=RegimeId, - description=f"Structural retirement model ({policy.name})", - fixed_params=fixed_params, - derived_categoricals=derived_categoricals, - execution_config=( - execution_config_for_devices() - if execution_config is None - else execution_config - ), - **model_slots, - ) diff --git a/src/aca_model/aca/regimes/__init__.py b/src/aca_model/aca/regimes/__init__.py index bd9d237..747ce1e 100644 --- a/src/aca_model/aca/regimes/__init__.py +++ b/src/aca_model/aca/regimes/__init__.py @@ -6,8 +6,8 @@ from collections.abc import Mapping from typing import Any +from jax.typing import ArrayLike from lcm import DiscreteGrid, Regime -from lcm.typing import UserParams from aca_model.aca.health_insurance import PolicyVariant from aca_model.aca.regimes._overrides import apply_aca_overrides @@ -22,7 +22,7 @@ def build_all_regimes( *, policy: PolicyVariant, grid_config: GridConfig, - fixed_params: UserParams, + pia_aime_grid: ArrayLike, wage_params: Mapping[str, Any], pref_type_grid: DiscreteGrid, solver: SolverName = "brute_force", @@ -31,7 +31,7 @@ def build_all_regimes( """Build all 19 regimes with ACA policy overrides.""" regimes = baseline_build_all_regimes( grid_config=grid_config, - fixed_params=fixed_params, + pia_aime_grid=pia_aime_grid, wage_params=wage_params, pref_type_grid=pref_type_grid, solver=solver, diff --git a/src/aca_model/baseline/model.py b/src/aca_model/baseline/model.py index 78909f4..70d82af 100644 --- a/src/aca_model/baseline/model.py +++ b/src/aca_model/baseline/model.py @@ -1,21 +1,27 @@ -"""Baseline structural retirement model using pylcm. +"""Structural retirement model specification and construction using pylcm. -This is the main model specification, ported from struct-ret/. The baseline -represents pre-ACA rules (no individual mandate, no ACA subsidies). +The baseline represents pre-ACA rules (no individual mandate, no ACA +subsidies); an ACA `PolicyVariant` overlays its schedules on the baseline +regimes. Usage: - from aca_model.baseline.model import create_model - model = create_model(fixed_params=..., wage_params=..., ...) - params = get_default_params() - V = model.solve(params) + from aca_model.baseline.model import build_model_spec, create_model + spec = build_model_spec(grid_config=..., pia_aime_grid=..., ...) + model = create_model(spec=spec, fixed_params=...) """ from collections.abc import Mapping from typing import Any -from lcm import AgeGrid, DiscreteGrid, ExecutionConfig, InitialNodes, Model +from jax.typing import ArrayLike +from lcm import AgeGrid, DiscreteGrid, Model from lcm.typing import UserParams +from aca_model.aca.health_insurance import PolicyVariant +from aca_model.aca.regimes import build_all_regimes as build_all_aca_regimes +from aca_model.agent.health import GoodHealth +from aca_model.agent.labor_market import IsMarried +from aca_model.baseline.health_insurance import HealthInsuranceState from aca_model.baseline.regimes import ( RegimeId, SolverName, @@ -24,114 +30,150 @@ build_model_slots, ) from aca_model.config import MODEL_CONFIG, GridConfig, model_ages -from aca_model.execution import execution_config_for_devices -def create_model( +def build_model_spec( *, - fixed_params: UserParams, - wage_params: Mapping[str, Any], - derived_categoricals: Mapping[str, DiscreteGrid], + policy: PolicyVariant | None = None, grid_config: GridConfig, + pia_aime_grid: ArrayLike, + wage_params: Mapping[str, Any], pref_type_grid: DiscreteGrid, solver: SolverName = "brute_force", consumption_dollars_points: tuple[float, ...] | None = None, - execution_config: ExecutionConfig | None = None, - initial_nodes: InitialNodes | None = None, final_age: int = MODEL_CONFIG.end_age, -) -> Model: - """Create the baseline structural retirement model. +) -> dict[str, Any]: + """Return the domain keyword arguments of `lcm.Model` for the retirement model. + + The spec holds `regimes`, `ages`, `regime_id_class`, `description`, + `derived_categoricals`, the entry `initial_nodes`, `edges`, `functions`, + `constraints`, `states` and `state_transitions`. It holds no generic + `Model` option; pass those to `create_model`. Merge an override with `|`, + e.g. `spec | {"initial_nodes": build_restart_initial_nodes(ages=...)}`. Args: - 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; - only estimation parameters should go through - `model.simulate(params=...)`. - wage_params: Data-derived wage profile dict (`log_ft_wage_mean`, - `log_ft_wage_std`, `adj_wage_hours_*`) used only at grid-build - time to size the assets-floor to `-max_annual_labor_income`. - Not routed to the pylcm Model. - derived_categoricals: Categorical mappings for `pd.Series` - fixed_params index levels that aren't model state/action - grids — `target_his`, `his`, `good_health`, `is_married`, - `pref_type`. + policy: ACA policy whose schedules overlay the baseline regimes. None + builds the baseline (pre-ACA) model. grid_config: Continuous-grid point counts. Pass `GRID_CONFIG` for production values or `BENCHMARK_GRID_CONFIG` for the fast-but-structurally-faithful benchmark. + pia_aime_grid: PIA bend points in AIME, which break the AIME grid + into its segments. + wage_params: Data-derived wage profile dict (`log_ft_wage_mean`, + `log_ft_wage_std`, `adj_wage_hours_*`) used only at grid-build + time to size the assets-floor to `-max_annual_labor_income`. pref_type_grid: Pref-type `DiscreteGrid`. Pass `DiscreteGrid(PrefType)` for the production 3-type layout, or a compact variant (e.g. `DiscreteGrid(BenchmarkPrefType)`). - solver: `"brute_force"` (the default) or `"dcegm"`. DC-EGM swaps - the spec to its post-decision form: savings-form assets laws, + solver: `"brute_force"` (the default), `"dcegm"` or `"nbegm"`. DC-EGM + swaps the spec to its post-decision form: savings-form assets laws, broadcast `resources`/`savings`/`inverse_marginal_utility` - functions, no borrowing constraint, and a `DCEGM` solver - config on every living regime. + functions, and a `DCEGM` solver config on every living regime. consumption_dollars_points: Construction-time consumption action gridpoints. Required under DC-EGM (the kernel needs the 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. - initial_nodes: Admissible starting `(age, regime)` pairs. None admits - the estimation sample's entry ages and regimes that lie before - `final_age`; pass `build_restart_initial_nodes` to restart - simulations mid-life. final_age: First age at which everybody is dead; the model spans `MODEL_CONFIG.start_age` through `final_age`. The default is the full horizon. A shorter horizon reads a prefix of the age-indexed - data, and `fixed_params` must hold only what it reads (see + data, and its fixed params must hold only what it reads (see `select_horizon_params`). Returns: - A pylcm Model with 19 regimes (18 non-terminal + dead) spanning - ages 51 through `final_age` (96 by default); `final_age - 1` is the - last age anyone is alive. Regime names follow - the `___` scheme. + Keyword arguments for `lcm.Model` describing 19 regimes (18 + non-terminal + dead) over ages 51 through `final_age`; `final_age - 1` + is the last age anyone is alive. Regime names follow the + `___` scheme. """ - ages = model_ages(final_age=final_age) _fail_if_dcegm_without_consumption_points( solver=solver, consumption_dollars_points=consumption_dollars_points ) - regimes = build_all_regimes( - grid_config=grid_config, - fixed_params=fixed_params, - wage_params=wage_params, - pref_type_grid=pref_type_grid, - solver=solver, - consumption_dollars_points=consumption_dollars_points, + ages = model_ages(final_age=final_age) + regimes = ( + build_all_regimes( + grid_config=grid_config, + pia_aime_grid=pia_aime_grid, + wage_params=wage_params, + pref_type_grid=pref_type_grid, + solver=solver, + consumption_dollars_points=consumption_dollars_points, + ) + if policy is None + else build_all_aca_regimes( + policy=policy, + grid_config=grid_config, + pia_aime_grid=pia_aime_grid, + wage_params=wage_params, + pref_type_grid=pref_type_grid, + solver=solver, + consumption_dollars_points=consumption_dollars_points, + ) ) + # The ACA overlay swaps only regime-level functions; the broadcast slots + # are policy-invariant. model_slots = build_model_slots( grid_config=grid_config, - fixed_params=fixed_params, + pia_aime_grid=pia_aime_grid, wage_params=wage_params, pref_type_grid=pref_type_grid, ages=ages, solver=solver, ) - - return Model( - regimes=regimes, - initial_nodes=( - build_entry_initial_nodes(ages=ages) - if initial_nodes is None - else initial_nodes + return { + "regimes": regimes, + "ages": AgeGrid(exact_values=ages), + "regime_id_class": RegimeId, + "description": ( + "Baseline structural retirement model (pre-ACA)" + if policy is None + else f"Structural retirement model ({policy.name})" ), - ages=AgeGrid(exact_values=ages), - regime_id_class=RegimeId, - description="Baseline structural retirement model (pre-ACA)", - fixed_params=fixed_params, - derived_categoricals=derived_categoricals, - execution_config=( - execution_config_for_devices() - if execution_config is None - else execution_config + "derived_categoricals": build_derived_categoricals( + pref_type_grid=pref_type_grid ), + "initial_nodes": build_entry_initial_nodes(ages=ages), **model_slots, - ) + } + + +def create_model( + *, spec: Mapping[str, Any], fixed_params: UserParams, **model_options: Any +) -> Model: + """Build the pylcm Model from `spec`, `fixed_params` and generic options. + + Args: + spec: Domain keyword arguments from `build_model_spec`. + fixed_params: Parameters to fix at model creation time. Fixed params + are partialled into compiled functions and removed from the params + template. They must hold exactly what the spec's horizon reads. + **model_options: Any other `lcm.Model` keyword argument, forwarded + unchanged, e.g. `execution_config`. A key the spec also holds, + such as `initial_nodes`, overrides the spec's value. + + Returns: + The pylcm Model. + + """ + return Model(**{**spec, **model_options}, fixed_params=fixed_params) + + +def build_derived_categoricals( + *, pref_type_grid: DiscreteGrid +) -> dict[str, DiscreteGrid]: + """Return the labels of params index levels that are not grids of every regime. + + `pd.Series` params are indexed by `target_his`, `his`, `good_health`, + `is_married` and `pref_type`. + """ + return { + "good_health": DiscreteGrid(GoodHealth), + "is_married": DiscreteGrid(IsMarried), + "his": DiscreteGrid(HealthInsuranceState), + "target_his": DiscreteGrid(HealthInsuranceState), + "pref_type": pref_type_grid, + } def _fail_if_dcegm_without_consumption_points( diff --git a/src/aca_model/baseline/regimes/__init__.py b/src/aca_model/baseline/regimes/__init__.py index 296058d..7034033 100644 --- a/src/aca_model/baseline/regimes/__init__.py +++ b/src/aca_model/baseline/regimes/__init__.py @@ -13,9 +13,9 @@ from collections.abc import Mapping from typing import Any +from jax.typing import ArrayLike from lcm import DiscreteGrid, Regime, Transition from lcm.solvers import DCEGM, NBEGM -from lcm.typing import UserParams from aca_model.baseline.regimes import _nongroup as nongroup from aca_model.baseline.regimes import _retiree as retiree @@ -94,7 +94,7 @@ def build_regime( def build_all_regimes( *, grid_config: GridConfig, - fixed_params: UserParams, + pia_aime_grid: ArrayLike, wage_params: Mapping[str, Any], pref_type_grid: DiscreteGrid, solver: SolverName = "brute_force", @@ -102,7 +102,7 @@ def build_all_regimes( ) -> dict[str, Regime]: """Build all 19 baseline regimes (18 non-terminal + dead). - `fixed_params` carries the PIA bends for the AIME piecewise grid; + `pia_aime_grid` holds the PIA bends for the AIME piecewise grid; `wage_params` sizes the assets-floor to `-max_annual_labor_income`; `pref_type_grid` selects the pref-type cardinality (production `DiscreteGrid(PrefType)` or the benchmark's 2-type variant); @@ -113,7 +113,7 @@ def build_all_regimes( """ grids = build_grids( grid_config=grid_config, - fixed_params=fixed_params, + pia_aime_grid=pia_aime_grid, wage_params=wage_params, pref_type_grid=pref_type_grid, consumption_dollars_points=consumption_dollars_points, @@ -135,7 +135,7 @@ def build_all_regimes( def build_model_slots( *, grid_config: GridConfig, - fixed_params: UserParams, + pia_aime_grid: ArrayLike, wage_params: Mapping[str, Any], pref_type_grid: DiscreteGrid, ages: tuple[int, ...], @@ -145,15 +145,14 @@ def build_model_slots( Returns keyword arguments for `lcm.Model(...)`: the edges, functions, constraint, states, and laws of motion shared by all living regimes. - Both the baseline and the ACA `create_model` consume this — the ACA - overlay swaps only regime-level functions, so the broadcast slots are - policy-invariant. Under DC-EGM the solver-contract functions join the + 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; the borrowing constraint is declared under every solver. `ages` is the model's age grid, which the edges' age schedules cover. """ grids = build_grids( grid_config=grid_config, - fixed_params=fixed_params, + pia_aime_grid=pia_aime_grid, wage_params=wage_params, pref_type_grid=pref_type_grid, ) diff --git a/src/aca_model/baseline/regimes/_common.py b/src/aca_model/baseline/regimes/_common.py index c16e0bf..52faade 100644 --- a/src/aca_model/baseline/regimes/_common.py +++ b/src/aca_model/baseline/regimes/_common.py @@ -13,6 +13,7 @@ import jax.numpy as jnp import numpy as np +from jax.typing import ArrayLike from lcm import ( AdditiveShockTransition, ByAge, @@ -43,7 +44,6 @@ Period, RegimeName, ScalarInt, - UserParams, ) from aca_model.agent import ( @@ -284,7 +284,7 @@ class Grids: def build_grids( *, grid_config: GridConfig, - fixed_params: UserParams, + pia_aime_grid: ArrayLike, wage_params: Mapping[str, Any], pref_type_grid: DiscreteGrid, consumption_dollars_points: tuple[float, ...] | None = None, @@ -292,7 +292,7 @@ def build_grids( """Build continuous-state/action grids from a `GridConfig`. The AIME grid is `PiecewiseLinSpacedGrid` breakpointed at the PIA - bends from `fixed_params["pia_aime_grid"]`. The assets grid + bends in `pia_aime_grid`. The assets grid (`build_assets_grid`) runs from `-max_annual_labor_income`, computed from `wage_params` (`log_ft_wage_mean`, `log_ft_wage_std`, `adj_wage_hours_*`), to the largest initial asset holding. @@ -331,7 +331,7 @@ def build_grids( floor=assets_start, n_points=grid_config.n_assets_gridpoints ), aime=_build_aime_grid( - fixed_params=fixed_params, + pia_aime_grid=pia_aime_grid, points_per_segment=grid_config.aime_points_per_segment, ), pension_wealth=_PENSION_WEALTH_GRID, @@ -399,17 +399,17 @@ def get_hcc_persistent_grid_points(*, grid_config: GridConfig) -> FloatND: def _build_aime_grid( - *, fixed_params: UserParams, points_per_segment: tuple[int, int, int, int] + *, pia_aime_grid: ArrayLike, points_per_segment: tuple[int, int, int, int] ) -> PiecewiseLinSpacedGrid: """Return the AIME grid. The grid is piecewise-linspaced with breakpoints at the PIA bends - in `fixed_params["pia_aime_grid"]` and `points_per_segment` in + in `pia_aime_grid` and `points_per_segment` in each segment. The fifth bend point is the delayed-retirement-credit extension above the taxable max, so the grid carries four segments. `n_aime_gridpoints` from `grid_config` is ignored on this path. """ - kinks = [float(k) for k in np.asarray(fixed_params["pia_aime_grid"])] + kinks = [float(k) for k in np.asarray(pia_aime_grid)] return PiecewiseLinSpacedGrid( start=kinks[0], stop=kinks[4], diff --git a/src/aca_model/benchmark.py b/src/aca_model/benchmark.py index 9b32db6..7b599bf 100644 --- a/src/aca_model/benchmark.py +++ b/src/aca_model/benchmark.py @@ -29,13 +29,10 @@ import jax.numpy as jnp import numpy as np from jax import Array -from lcm import DiscreteGrid, ExecutionConfig, InitialNodes, Model +from lcm import DiscreteGrid, Model -from aca_model.agent.health import GoodHealth -from aca_model.agent.labor_market import IsMarried from aca_model.agent.preferences import BenchmarkPrefType -from aca_model.baseline.health_insurance import HealthInsuranceState -from aca_model.baseline.model import create_model +from aca_model.baseline.model import build_model_spec, create_model from aca_model.baseline.regimes import ENTRY_REGIMES, select_horizon_params from aca_model.config import ( BENCHMARK_GRID_CONFIG, @@ -55,14 +52,6 @@ _N_BENCHMARK_PREF_TYPES = len(DiscreteGrid(BenchmarkPrefType).categories) -_DERIVED_CATEGORICALS = { - "good_health": DiscreteGrid(GoodHealth), - "is_married": DiscreteGrid(IsMarried), - "his": DiscreteGrid(HealthInsuranceState), - "target_his": DiscreteGrid(HealthInsuranceState), - "pref_type": DiscreteGrid(BenchmarkPrefType), -} - # The five entry regimes (inelig + canwork). All have # HealthWithDisability (3-state); four have lagged_labor_supply (tied # does not — it's implied by the regime). @@ -72,9 +61,8 @@ def create_benchmark_model( *, pref_type_grid: DiscreteGrid, grid_config: GridConfig = BENCHMARK_GRID_CONFIG, - execution_config: ExecutionConfig | None = None, - initial_nodes: InitialNodes | None = None, final_age: int = MODEL_CONFIG.end_age, + **model_options: Any, ) -> Model: """Create the aca baseline with `BENCHMARK_GRID_CONFIG` and frozen fixed_params. @@ -84,27 +72,23 @@ def create_benchmark_model( Args: pref_type_grid: Pref-type grid; pass `DiscreteGrid(BenchmarkPrefType)`. grid_config: Economic resolution and medical-shock folding declarations. - execution_config: Explicit hardware-local policy forwarded unchanged. - None uses the smallest selected accelerator allocator limit as the - device-memory budget; CPU construction remains unbudgeted. - initial_nodes: Admissible starting `(age, regime)` pairs, forwarded - to `create_model`. final_age: First age at which everybody is dead, forwarded to - `create_model` and `get_benchmark_params`. The default is the full - horizon; the frozen age-indexed parameters cover it, and a shorter - horizon reads a prefix of them. + `build_model_spec` and `get_benchmark_params`. The default is the + full horizon; the frozen age-indexed parameters cover it, and a + shorter horizon reads a prefix of them. + **model_options: Generic `lcm.Model` options forwarded to + `create_model`, e.g. `execution_config` or an `initial_nodes` + override. """ fixed_params, wage_params, _ = get_benchmark_params(model=None, final_age=final_age) - return create_model( + spec = build_model_spec( grid_config=grid_config, - fixed_params=fixed_params, + pia_aime_grid=fixed_params["pia_aime_grid"], wage_params=wage_params, - derived_categoricals=_DERIVED_CATEGORICALS, pref_type_grid=pref_type_grid, - execution_config=execution_config, - initial_nodes=initial_nodes, final_age=final_age, ) + return create_model(spec=spec, fixed_params=fixed_params, **model_options) def get_benchmark_params( diff --git a/tests/helpers/model.py b/tests/helpers/model.py index 3645ab7..d54c551 100644 --- a/tests/helpers/model.py +++ b/tests/helpers/model.py @@ -10,12 +10,8 @@ from lcm import DiscreteGrid, Model 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.health_insurance import HealthInsuranceState -from aca_model.baseline.model import create_model as _create_baseline_model +from aca_model.baseline.model import build_model_spec, create_model from aca_model.baseline.regimes import REGIME_SPECS from aca_model.benchmark import get_benchmark_params from aca_model.config import BENCHMARK_GRID_CONFIG, GridConfig @@ -26,14 +22,6 @@ if spec["his"] == "nongroup" and spec["mc"] == "nomc" ) -_DERIVED_CATEGORICALS = { - "good_health": DiscreteGrid(GoodHealth), - "is_married": DiscreteGrid(IsMarried), - "his": DiscreteGrid(HealthInsuranceState), - "target_his": DiscreteGrid(HealthInsuranceState), - "pref_type": DiscreteGrid(BenchmarkPrefType), -} - _UNFOLDED_GRID_CONFIG = replace( BENCHMARK_GRID_CONFIG, fold_hcc_transitory=False, subtract_hcc_bill=False ) @@ -42,12 +30,14 @@ def make_baseline_model(*, grid_config: GridConfig = _UNFOLDED_GRID_CONFIG) -> Model: """Baseline model on `grid_config` with the benchmark snapshot params.""" fixed_params, wage_params, _ = get_benchmark_params(model=None) - return _create_baseline_model( + return create_model( + spec=build_model_spec( + pia_aime_grid=fixed_params["pia_aime_grid"], + wage_params=wage_params, + grid_config=grid_config, + pref_type_grid=DiscreteGrid(BenchmarkPrefType), + ), fixed_params=fixed_params, - wage_params=wage_params, - derived_categoricals=_DERIVED_CATEGORICALS, - grid_config=grid_config, - pref_type_grid=DiscreteGrid(BenchmarkPrefType), ) @@ -76,11 +66,13 @@ def make_aca_model( ) -> Model: """ACA model on `grid_config` with the benchmark snapshot params.""" fixed_params, wage_params, _ = get_benchmark_params(model=None) - return _create_aca_model( - policy=policy, + return create_model( + spec=build_model_spec( + pia_aime_grid=fixed_params["pia_aime_grid"], + policy=policy, + wage_params=wage_params, + grid_config=grid_config, + pref_type_grid=DiscreteGrid(BenchmarkPrefType), + ), fixed_params=aca_fixed_params(fixed_params=fixed_params, policy=policy), - wage_params=wage_params, - derived_categoricals=_DERIVED_CATEGORICALS, - grid_config=grid_config, - pref_type_grid=DiscreteGrid(BenchmarkPrefType), ) diff --git a/tests/test_aime_grid.py b/tests/test_aime_grid.py index 7a976bb..170d156 100644 --- a/tests/test_aime_grid.py +++ b/tests/test_aime_grid.py @@ -20,7 +20,8 @@ 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, points_per_segment=_POINTS_PER_SEGMENT + pia_aime_grid=_FIXED_PARAMS["pia_aime_grid"], + points_per_segment=_POINTS_PER_SEGMENT, ) np.testing.assert_allclose( [point.value for point in grid.breakpoints], @@ -33,7 +34,8 @@ def test_build_aime_grid_owns_each_pia_breakpoint_on_the_right() -> None: def test_build_aime_grid_top_point_is_extension_aime() -> None: """The grid reaches the delayed-credit extension AIME at its top.""" grid = _build_aime_grid( - fixed_params=_FIXED_PARAMS, points_per_segment=_POINTS_PER_SEGMENT + pia_aime_grid=_FIXED_PARAMS["pia_aime_grid"], + points_per_segment=_POINTS_PER_SEGMENT, ) np.testing.assert_allclose(float(grid.to_jax().max()), 187954.752, rtol=1e-5) diff --git a/tests/test_dcegm_model_creation.py b/tests/test_dcegm_model_creation.py index 8dfa78d..6065f9a 100644 --- a/tests/test_dcegm_model_creation.py +++ b/tests/test_dcegm_model_creation.py @@ -11,14 +11,13 @@ import numpy as np import pytest -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 from aca_model.agent.preferences import BenchmarkPrefType -from aca_model.baseline.model import create_model +from aca_model.baseline.model import build_model_spec, create_model from aca_model.baseline.regimes import ( REGIME_SPECS, SolverName, @@ -44,7 +43,7 @@ def _build_regimes(solver: SolverName) -> dict[str, Regime]: return build_all_regimes( grid_config=_EGM_GRID_CONFIG, - fixed_params=_FIXED_PARAMS, + pia_aime_grid=_FIXED_PARAMS["pia_aime_grid"], wage_params=_WAGE_PARAMS, pref_type_grid=DiscreteGrid(BenchmarkPrefType), solver=solver, @@ -53,12 +52,14 @@ def _build_regimes(solver: SolverName) -> dict[str, Regime]: def _build_model(solver: SolverName) -> Model: return create_model( + spec=build_model_spec( + pia_aime_grid=_FIXED_PARAMS["pia_aime_grid"], + wage_params=_WAGE_PARAMS, + grid_config=_EGM_GRID_CONFIG, + pref_type_grid=DiscreteGrid(BenchmarkPrefType), + solver=solver, + ), fixed_params=_FIXED_PARAMS, - wage_params=_WAGE_PARAMS, - derived_categoricals=_DERIVED_CATEGORICALS, - grid_config=_EGM_GRID_CONFIG, - pref_type_grid=DiscreteGrid(BenchmarkPrefType), - solver=solver, ) @@ -100,7 +101,7 @@ def test_dcegm_model_slots_broadcast_contract_functions_and_constraint() -> None consumption grid and needs the explicit feasibility mask.""" slots = build_model_slots( grid_config=_EGM_GRID_CONFIG, - fixed_params=_FIXED_PARAMS, + pia_aime_grid=_FIXED_PARAMS["pia_aime_grid"], wage_params=_WAGE_PARAMS, pref_type_grid=DiscreteGrid(BenchmarkPrefType), ages=MODEL_AGES, @@ -131,12 +132,14 @@ 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( + spec=build_model_spec( + pia_aime_grid=_FIXED_PARAMS["pia_aime_grid"], + wage_params=_WAGE_PARAMS, + grid_config=_EGM_GRID_CONFIG, + pref_type_grid=DiscreteGrid(BenchmarkPrefType), + solver="dcegm", + ), fixed_params=_FIXED_PARAMS, - wage_params=_WAGE_PARAMS, - derived_categoricals=_DERIVED_CATEGORICALS, - grid_config=_EGM_GRID_CONFIG, - pref_type_grid=DiscreteGrid(BenchmarkPrefType), - solver="dcegm", ) @@ -154,15 +157,17 @@ def test_benchmark_consumption_points_pin_both_floors() -> None: def test_dcegm_benchmark_model_builds() -> None: """The benchmark model accepts `solver="dcegm"` end to end.""" model = create_model( - fixed_params=_FIXED_PARAMS, - wage_params=_WAGE_PARAMS, - derived_categoricals=_DERIVED_CATEGORICALS, - grid_config=_EGM_GRID_CONFIG, - pref_type_grid=DiscreteGrid(BenchmarkPrefType), - solver="dcegm", - consumption_dollars_points=get_benchmark_consumption_dollars_points( - n_points=_EGM_GRID_CONFIG.n_consumption_dollars_gridpoints + spec=build_model_spec( + pia_aime_grid=_FIXED_PARAMS["pia_aime_grid"], + wage_params=_WAGE_PARAMS, + grid_config=_EGM_GRID_CONFIG, + pref_type_grid=DiscreteGrid(BenchmarkPrefType), + solver="dcegm", + consumption_dollars_points=get_benchmark_consumption_dollars_points( + n_points=_EGM_GRID_CONFIG.n_consumption_dollars_gridpoints + ), ), + fixed_params=_FIXED_PARAMS, ) assert isinstance(model.user_regimes["retiree_nomc_inelig_canwork"].solver, DCEGM) @@ -173,7 +178,7 @@ def test_savings_grid_length_follows_grid_config() -> None: grid_config = dataclasses.replace(_EGM_GRID_CONFIG, n_savings_gridpoints=70) regimes = build_all_regimes( grid_config=grid_config, - fixed_params=_FIXED_PARAMS, + pia_aime_grid=_FIXED_PARAMS["pia_aime_grid"], wage_params=_WAGE_PARAMS, pref_type_grid=DiscreteGrid(BenchmarkPrefType), solver="dcegm", @@ -190,7 +195,7 @@ def test_dcegm_savings_grid_rejects_too_few_points() -> None: with pytest.raises(ValueError, match="n_savings_gridpoints"): build_all_regimes( grid_config=grid_config, - fixed_params=_FIXED_PARAMS, + pia_aime_grid=_FIXED_PARAMS["pia_aime_grid"], wage_params=_WAGE_PARAMS, pref_type_grid=DiscreteGrid(BenchmarkPrefType), solver="dcegm", diff --git a/tests/test_dcegm_parity.py b/tests/test_dcegm_parity.py index 70bfe05..3bd86af 100644 --- a/tests/test_dcegm_parity.py +++ b/tests/test_dcegm_parity.py @@ -18,11 +18,8 @@ import pytest from lcm import DiscreteGrid, Model -from aca_model.agent.health import GoodHealth -from aca_model.agent.labor_market import IsMarried from aca_model.agent.preferences import BenchmarkPrefType -from aca_model.baseline.health_insurance import HealthInsuranceState -from aca_model.baseline.model import SolverName, create_model +from aca_model.baseline.model import SolverName, build_model_spec, create_model from aca_model.benchmark import ( get_benchmark_consumption_dollars_points, get_benchmark_initial_conditions, @@ -78,14 +75,6 @@ def _recorded_noise_floor() -> tuple[float, float, float]: ) -_DERIVED_CATEGORICALS = { - "good_health": DiscreteGrid(GoodHealth), - "is_married": DiscreteGrid(IsMarried), - "his": DiscreteGrid(HealthInsuranceState), - "target_his": DiscreteGrid(HealthInsuranceState), - "pref_type": DiscreteGrid(BenchmarkPrefType), -} - _DISCRETE_COLUMNS = ("regime_name", "claim_ss", "labor_supply", "buy_private") _XFAIL_UNTIL_DCEGM_ADMITS_THE_ACA_BUDGET = pytest.mark.xfail( @@ -114,13 +103,15 @@ def _make_model(*, solver: SolverName, grid_config: GridConfig) -> Model: else None ) return create_model( + spec=build_model_spec( + pia_aime_grid=fixed_params["pia_aime_grid"], + wage_params=wage_params, + grid_config=grid_config, + pref_type_grid=DiscreteGrid(BenchmarkPrefType), + solver=solver, + consumption_dollars_points=consumption_dollars_points, + ), fixed_params=fixed_params, - wage_params=wage_params, - derived_categoricals=_DERIVED_CATEGORICALS, - grid_config=grid_config, - pref_type_grid=DiscreteGrid(BenchmarkPrefType), - solver=solver, - consumption_dollars_points=consumption_dollars_points, ) diff --git a/tests/test_execution_config.py b/tests/test_execution_config.py index 79a47d6..743dfbc 100644 --- a/tests/test_execution_config.py +++ b/tests/test_execution_config.py @@ -1,26 +1,47 @@ -"""Model factories expose pylcm's execution policy without changing economic grids.""" +"""Model factories forward pylcm's generic options without changing economic grids.""" from functools import partial -from types import SimpleNamespace import jax import pytest from helpers.model import aca_fixed_params -from lcm import DiscreteGrid, ExecutionConfig +from lcm import DiscreteGrid, ExecutionConfig, InitialNodes 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.baseline.model import build_model_spec, create_model from aca_model.benchmark import create_benchmark_model, get_benchmark_params from aca_model.config import BENCHMARK_GRID_CONFIG +_DOMAIN_KEYS = frozenset( + { + "regimes", + "ages", + "regime_id_class", + "description", + "derived_categoricals", + "initial_nodes", + "edges", + "functions", + "constraints", + "states", + "state_transitions", + } +) + + +def _spec(*, policy: PolicyVariant | None = None) -> dict: + fixed_params, wage_params, _ = get_benchmark_params(model=None) + return build_model_spec( + policy=policy, + grid_config=BENCHMARK_GRID_CONFIG, + pia_aime_grid=fixed_params["pia_aime_grid"], + wage_params=wage_params, + pref_type_grid=DiscreteGrid(BenchmarkPrefType), + ) + def _factory(kind): if kind == "benchmark": @@ -28,27 +49,19 @@ def _factory(kind): create_benchmark_model, pref_type_grid=DiscreteGrid(BenchmarkPrefType), ) - fixed_params, wage_params, _ = get_benchmark_params(model=None) - factory = create_model - if kind == "aca": - factory = partial(create_aca_model, policy=PolicyVariant.ACA) - fixed_params = aca_fixed_params( - fixed_params=fixed_params, policy=PolicyVariant.ACA - ) - return partial( - factory, - fixed_params=fixed_params, - wage_params=wage_params, - derived_categoricals={ - "good_health": DiscreteGrid(GoodHealth), - "is_married": DiscreteGrid(IsMarried), - "his": DiscreteGrid(HealthInsuranceState), - "target_his": DiscreteGrid(HealthInsuranceState), - "pref_type": DiscreteGrid(BenchmarkPrefType), - }, - grid_config=BENCHMARK_GRID_CONFIG, - pref_type_grid=DiscreteGrid(BenchmarkPrefType), - ) + fixed_params, _, _ = get_benchmark_params(model=None) + policy = PolicyVariant.ACA if kind == "aca" else None + if policy is not None: + fixed_params = aca_fixed_params(fixed_params=fixed_params, policy=policy) + return partial(create_model, spec=_spec(policy=policy), fixed_params=fixed_params) + + +@pytest.mark.parametrize("policy", [None, PolicyVariant.ACA]) +def test_build_model_spec_holds_exactly_the_domain_arguments( + policy: PolicyVariant | None, +) -> None: + """The spec carries the model's domain and no generic `Model` option.""" + assert set(_spec(policy=policy)) == _DOMAIN_KEYS @pytest.mark.parametrize("kind", ["baseline", "aca", "benchmark"]) @@ -73,48 +86,50 @@ def test_factory_preserves_axis_width_validation(kind): _factory(kind)(execution_config=config) -@pytest.mark.parametrize("kind", ["baseline", "aca", "benchmark"]) -@pytest.mark.parametrize( - ("requested_policy", "expected_budget"), - [ - (None, 800), - ( - ExecutionConfig( - devices=(0,), - device_memory_bytes=400, - sharded_states=("pref_type",), - axis_widths={"cell": 32}, - ), - 400, - ), - (ExecutionConfig(devices=(0,)), ExecutionConfig().device_memory_bytes), - ], -) -def test_factory_supplies_the_requested_or_measured_budget( - monkeypatch, kind, requested_policy, expected_budget -): - """Defaults use measured limits; explicit budgets and unbudgeted controls survive.""" - factory = _factory(kind) - device = SimpleNamespace( - id=0, platform="gpu", memory_stats=lambda: {"bytes_limit": 800} - ) - monkeypatch.setattr(jax, "devices", lambda: [device]) +class _ConstructionObservedError(Exception): + def __init__(self, kwargs: dict) -> None: + super().__init__() + self.kwargs = kwargs - class ConstructionObservedError(Exception): - pass +def _observe_model_kwargs(monkeypatch, factory, **options) -> dict: 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 + raise _ConstructionObservedError(kwargs) 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) + with pytest.raises(_ConstructionObservedError) as observed: + factory(**options) + return observed.value.kwargs + + +@pytest.mark.parametrize("kind", ["baseline", "aca", "benchmark"]) +def test_factory_without_execution_config_leaves_pylcm_default(monkeypatch, kind): + """No execution policy reaches `Model`, so pylcm's device default applies.""" + kwargs = _observe_model_kwargs(monkeypatch, _factory(kind)) + assert "execution_config" not in kwargs + + +@pytest.mark.parametrize("kind", ["baseline", "aca", "benchmark"]) +def test_factory_forwards_an_explicit_execution_config_unchanged(monkeypatch, kind): + """An explicit policy reaches `Model` as the same object.""" + config = ExecutionConfig(devices=(0,), device_memory_bytes=400) + kwargs = _observe_model_kwargs(monkeypatch, _factory(kind), execution_config=config) + assert kwargs["execution_config"] is config + + +@pytest.mark.parametrize("kind", ["baseline", "benchmark"]) +def test_factory_forwards_any_model_option(monkeypatch, kind): + """A generic `Model` option the spec does not hold is passed through.""" + kwargs = _observe_model_kwargs(monkeypatch, _factory(kind), durable_identity=False) + assert kwargs["durable_identity"] is False + + +@pytest.mark.parametrize("kind", ["baseline", "benchmark"]) +def test_factory_initial_nodes_override_the_entry_default(monkeypatch, kind): + """`initial_nodes` passed as an option replaces the spec's entry starts.""" + nodes = InitialNodes(by_age={52: ("retiree_nomc_inelig_canwork",)}) + kwargs = _observe_model_kwargs(monkeypatch, _factory(kind), initial_nodes=nodes) + assert kwargs["initial_nodes"] is nodes @pytest.mark.parametrize("kind", ["baseline", "aca", "benchmark"]) diff --git a/tests/test_model_creation.py b/tests/test_model_creation.py index 6b56b25..608db83 100644 --- a/tests/test_model_creation.py +++ b/tests/test_model_creation.py @@ -32,7 +32,7 @@ def build_aca_regimes(policy: PolicyVariant) -> dict: return _build_aca_regimes( policy=policy, grid_config=BENCHMARK_GRID_CONFIG, - fixed_params=_FIXED_PARAMS, + pia_aime_grid=_FIXED_PARAMS["pia_aime_grid"], wage_params=_WAGE_PARAMS, pref_type_grid=DiscreteGrid(BenchmarkPrefType), ) @@ -40,7 +40,7 @@ def build_aca_regimes(policy: PolicyVariant) -> dict: _GRIDS = build_grids( grid_config=BENCHMARK_GRID_CONFIG, - fixed_params=_FIXED_PARAMS, + pia_aime_grid=_FIXED_PARAMS["pia_aime_grid"], wage_params=_WAGE_PARAMS, pref_type_grid=DiscreteGrid(BenchmarkPrefType), ) diff --git a/tests/test_nbegm_labor_live_validation.py b/tests/test_nbegm_labor_live_validation.py index 62b4f9e..cbd07c9 100644 --- a/tests/test_nbegm_labor_live_validation.py +++ b/tests/test_nbegm_labor_live_validation.py @@ -16,11 +16,10 @@ import numpy as np import pytest -from helpers.model import _DERIVED_CATEGORICALS from lcm import DiscreteGrid, ExecutionConfig from aca_model.agent.preferences import BenchmarkPrefType -from aca_model.baseline.model import create_model +from aca_model.baseline.model import build_model_spec, create_model from aca_model.baseline.regimes import SolverName from aca_model.benchmark import get_benchmark_params from aca_model.config import BENCHMARK_GRID_CONFIG @@ -38,12 +37,14 @@ def _solve_m1(solver: SolverName) -> tuple[dict[int, np.ndarray], int]: ) fixed_params, wage_params, _ = get_benchmark_params(model=None) model = create_model( + spec=build_model_spec( + pia_aime_grid=fixed_params["pia_aime_grid"], + wage_params=wage_params, + grid_config=grid_config, + pref_type_grid=DiscreteGrid(BenchmarkPrefType), + solver=solver, + ), 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 {} ), diff --git a/tests/test_nbegm_model_creation.py b/tests/test_nbegm_model_creation.py index d384885..ef4b452 100644 --- a/tests/test_nbegm_model_creation.py +++ b/tests/test_nbegm_model_creation.py @@ -17,19 +17,18 @@ from typing import cast import pytest -from helpers.model import _DERIVED_CATEGORICALS, aca_fixed_params +from helpers.model import aca_fixed_params from lcm import DiscreteGrid, Model, Regime from lcm.consumption_savings_regime import ConsumptionSavingsRegime, LiquidMargin from lcm.exceptions import RegimeInitializationError from lcm.solvers import NBEGM from aca_model.aca import PolicyVariant -from aca_model.aca.model import create_model as create_aca_model from aca_model.aca.regimes import build_all_regimes as build_all_aca_regimes from aca_model.agent import assets_and_income from aca_model.agent.preferences import BenchmarkPrefType from aca_model.baseline import health_insurance -from aca_model.baseline.model import create_model +from aca_model.baseline.model import build_model_spec, create_model from aca_model.baseline.regimes import ( REGIME_SPECS, SolverName, @@ -65,7 +64,7 @@ def _build_regimes(solver: SolverName) -> dict[str, Regime]: return build_all_regimes( grid_config=_EGM_GRID_CONFIG, - fixed_params=_FIXED_PARAMS, + pia_aime_grid=_FIXED_PARAMS["pia_aime_grid"], wage_params=_WAGE_PARAMS, pref_type_grid=DiscreteGrid(BenchmarkPrefType), solver=solver, @@ -74,12 +73,14 @@ def _build_regimes(solver: SolverName) -> dict[str, Regime]: def _build_model_with(solver: SolverName, grid_config: GridConfig) -> Model: return create_model( + spec=build_model_spec( + pia_aime_grid=_FIXED_PARAMS["pia_aime_grid"], + wage_params=_WAGE_PARAMS, + grid_config=grid_config, + pref_type_grid=DiscreteGrid(BenchmarkPrefType), + solver=solver, + ), fixed_params=_FIXED_PARAMS, - wage_params=_WAGE_PARAMS, - derived_categoricals=_DERIVED_CATEGORICALS, - grid_config=grid_config, - pref_type_grid=DiscreteGrid(BenchmarkPrefType), - solver=solver, ) @@ -90,7 +91,7 @@ def _build_model(solver: SolverName) -> Model: def _grids() -> Grids: return build_grids( grid_config=_EGM_GRID_CONFIG, - fixed_params=_FIXED_PARAMS, + pia_aime_grid=_FIXED_PARAMS["pia_aime_grid"], wage_params=_WAGE_PARAMS, pref_type_grid=DiscreteGrid(BenchmarkPrefType), ) @@ -146,7 +147,7 @@ def test_build_nbegm_solver_forwards_the_jump_read_mode() -> None: grid_config = dataclasses.replace(_EGM_GRID_CONFIG, nbegm_jump_read="bridged") grids = build_grids( grid_config=grid_config, - fixed_params=_FIXED_PARAMS, + pia_aime_grid=_FIXED_PARAMS["pia_aime_grid"], wage_params=_WAGE_PARAMS, pref_type_grid=DiscreteGrid(BenchmarkPrefType), ) @@ -306,20 +307,22 @@ def test_nbegm_builds_every_aca_policy_variant(policy: PolicyVariant) -> None: 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( - policy=policy, + model = create_model( + spec=build_model_spec( + pia_aime_grid=_FIXED_PARAMS["pia_aime_grid"], + policy=policy, + wage_params=_WAGE_PARAMS, + grid_config=grid_config, + pref_type_grid=DiscreteGrid(BenchmarkPrefType), + solver="nbegm", + ), fixed_params=aca_fixed_params(fixed_params=_FIXED_PARAMS, policy=policy), - wage_params=_WAGE_PARAMS, - derived_categoricals=_DERIVED_CATEGORICALS, - grid_config=grid_config, - pref_type_grid=DiscreteGrid(BenchmarkPrefType), - solver="nbegm", ) assert isinstance(model, Model) regimes = build_all_aca_regimes( policy=policy, grid_config=grid_config, - fixed_params=_FIXED_PARAMS, + pia_aime_grid=_FIXED_PARAMS["pia_aime_grid"], wage_params=_WAGE_PARAMS, pref_type_grid=DiscreteGrid(BenchmarkPrefType), solver="nbegm", @@ -339,14 +342,16 @@ def test_nbegm_aca_variants_leave_no_free_buy_private_params( `*__buy_private` entry the pipeline cannot supply. """ grid_config = _BRIDGED_GRID_CONFIG - model = create_aca_model( - policy=policy, + model = create_model( + spec=build_model_spec( + pia_aime_grid=_FIXED_PARAMS["pia_aime_grid"], + policy=policy, + wage_params=_WAGE_PARAMS, + grid_config=grid_config, + pref_type_grid=DiscreteGrid(BenchmarkPrefType), + solver="nbegm", + ), fixed_params=aca_fixed_params(fixed_params=_FIXED_PARAMS, policy=policy), - wage_params=_WAGE_PARAMS, - derived_categoricals=_DERIVED_CATEGORICALS, - grid_config=grid_config, - pref_type_grid=DiscreteGrid(BenchmarkPrefType), - solver="nbegm", ) template = model.get_params_template() offenders = [ diff --git a/tests/test_nbegm_solve_validation.py b/tests/test_nbegm_solve_validation.py index bbe2265..d37b624 100644 --- a/tests/test_nbegm_solve_validation.py +++ b/tests/test_nbegm_solve_validation.py @@ -19,11 +19,10 @@ import numpy as np import pytest -from helpers.model import _DERIVED_CATEGORICALS from lcm import DiscreteGrid from aca_model.agent.preferences import BenchmarkPrefType -from aca_model.baseline.model import create_model +from aca_model.baseline.model import build_model_spec, create_model from aca_model.baseline.regimes import SolverName from aca_model.benchmark import get_benchmark_params from aca_model.config import BENCHMARK_GRID_CONFIG @@ -35,12 +34,14 @@ 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( + spec=build_model_spec( + pia_aime_grid=fixed_params["pia_aime_grid"], + wage_params=wage_params, + grid_config=grid_config, + pref_type_grid=DiscreteGrid(BenchmarkPrefType), + solver=solver, + ), fixed_params=fixed_params, - wage_params=wage_params, - derived_categoricals=_DERIVED_CATEGORICALS, - grid_config=grid_config, - pref_type_grid=DiscreteGrid(BenchmarkPrefType), - solver=solver, ) _, _, params = get_benchmark_params(model=model) solution = model.solve(params=params, log_level="off")