Split the model factories into a domain spec and a thin create_model

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`,

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,