Anchor consumption grid lower bound to consumption_floor parameter

Consumption is now declared as `IrregSpacedGrid(n_points=N)` (no
fixed points). Callers inject log-spaced gridpoints from
`consumption_floor` to $300k via `aca_model.consumption_grid.
inject_consumption_points(params=..., model=...)` before solving.
This means the lowest consumption choice equals the per-iteration
floor, removing a degree of freedom from the grid and eliminating
the previous mismatch where c < floor was a legal grid choice.

Requires pylcm support for runtime-supplied points on continuous
action grids (PR OpenSourceEconomics/pylcm#338). aca-model CI now
installs pylcm from the matching `feature/runtime-action-grids`
branch.

Other changes:
- `consumption_grid.py`: new module with `compute_consumption_points`
  and `inject_consumption_points` helpers.
- `benchmark.get_benchmark_params(*, model=None)`: when `model` is
  given, returns params with consumption points injected.
- `benchmark.get_benchmark_initial_conditions`: switch from `.start`
  / `.stop` to `to_jax().min()` / `.max()` so it works on both
  `LinSpacedGrid` and `PiecewiseLinSpacedGrid` (the AIME grid is now
  piecewise; this was a pre-existing bug surfacing as
  `AttributeError`).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: hmgaudecker

.github/workflows/main.yml

@@ -26,10 +26,10 @@ jobs:
       - uses: actions/setup-python@v6
         with:
           python-version: ${{ matrix.python-version }}
-      - name: Install pylcm
+      - name: Install pylcm (unreleased feature branch required)
         run: >-
           pip install "pylcm @
-          git+https://github.com/OpenSourceEconomics/pylcm.git@main"
+          git+https://github.com/OpenSourceEconomics/pylcm.git@feature/runtime-action-grids"
       - name: Install aca-model with test deps
         run: pip install -e . pytest pdbp
       - name: Run pytest

src/aca_model/baseline/regimes/_common.py

@@ -194,14 +194,6 @@ class Grids:
 # bend points (0 → kink_0 → kink_1 → kink_2). Total = 32.
 _AIME_PIECE_N_POINTS: tuple[int, int, int] = (10, 11, 11)
 
-# Consumption grid: log-spaced from the lower bound of the
-# `consumption_floor` parameter (BOUNDS in task_estimate_parameters)
-# up to a high value that brackets the unconstrained optimum for the
-# richest agents in the state space. Mirrors the struct-ret design
-# (concentrate gridpoints where CRRA curvature is highest).
-_CONSUMPTION_GRID_START: float = 100.0
-_CONSUMPTION_GRID_STOP: float = 300_000.0
-
 
 def build_grids(
     grid_config: GridConfig = GRID_CONFIG,
@@ -273,14 +265,7 @@ def build_grids(
         ),
         aime=_build_aime_grid(grid_config=grid_config, fixed_params=fixed_params),
         consumption=IrregSpacedGrid(
-            points=tuple(
-                float(c)
-                for c in np.geomspace(
-                    _CONSUMPTION_GRID_START,
-                    _CONSUMPTION_GRID_STOP,
-                    num=grid_config.n_consumption_gridpoints,
-                )
-            ),
+            n_points=grid_config.n_consumption_gridpoints,
         ),
         wage_res=wage_res,
         hcc_persistent=hcc_persistent,

src/aca_model/benchmark.py

@@ -44,6 +44,7 @@
 from aca_model.baseline.health_insurance import HealthInsuranceState
 from aca_model.baseline.model import create_model
 from aca_model.config import BENCHMARK_GRID_CONFIG
+from aca_model.consumption_grid import inject_consumption_points
 
 _PARAMS_FILE = (
     Path(__file__).resolve().parent / "_benchmark_data" / "benchmark_params.pkl"
@@ -96,17 +97,26 @@ def create_benchmark_model(*, pref_type_grid: DiscreteGrid | None = None) -> Mod
     )
 
 
-def get_benchmark_params() -> tuple[dict[str, Any], dict[str, Any]]:
+def get_benchmark_params(
+    *, model: Model | None = None
+) -> tuple[dict[str, Any], dict[str, Any]]:
     """Load the frozen `(fixed_params, params)` snapshot.
 
     Pref-type-indexed `pd.Series` in `params` are truncated to
     `_N_BENCHMARK_PREF_TYPES` rows so they line up with
     `BenchmarkPrefType`'s categories.
+
+    When `model` is provided, consumption gridpoints are injected into
+    `params` for each regime that declares `consumption` as an
+    `IrregSpacedGrid` with runtime-supplied points. The lower bound is
+    read from `params["consumption_floor"]`.
     """
     with _PARAMS_FILE.open("rb") as fh:
         data = cloudpickle.load(fh)
     fixed_params = data["fixed_params"]
     params = _truncate_pref_type_indexed(data["params"])
+    if model is not None:
+        params = inject_consumption_points(params=params, model=model)
     return fixed_params, params
 
 
@@ -143,10 +153,14 @@ def get_benchmark_initial_conditions(
     regime = rng.choice(regime_ids, size=n_subjects).astype(np.int32)
 
     # Grid ranges come from any of the five regimes (shared structure).
+    # Use to_jax() so the helper handles both LinSpacedGrid and
+    # PiecewiseLinSpacedGrid (the latter has no `.start` / `.stop`).
     ref_regime = model.regimes[_INITIAL_REGIMES[0]]
     grids = ref_regime.states
-    assets_lo, assets_hi = grids["assets"].start, grids["assets"].stop
-    aime_lo, aime_hi = grids["aime"].start, grids["aime"].stop
+    assets_pts = np.asarray(grids["assets"].to_jax())
+    aime_pts = np.asarray(grids["aime"].to_jax())
+    assets_lo, assets_hi = float(assets_pts.min()), float(assets_pts.max())
+    aime_lo, aime_hi = float(aime_pts.min()), float(aime_pts.max())
     hcc_p_pts = np.asarray(grids["hcc_persistent"].to_jax())
     hcc_t_pts = np.asarray(grids["hcc_transitory"].to_jax())
     wage_res_pts = np.asarray(grids["log_ft_wage_res"].to_jax())

src/aca_model/consumption_grid.py

@@ -0,0 +1,76 @@
+"""Runtime-supplied gridpoints for the consumption action.
+
+Consumption is declared as `IrregSpacedGrid(n_points=N)` in
+`baseline.regimes._common.build_grids` so the lower bound can track
+the per-iteration `consumption_floor` parameter. Callers must inject
+the actual gridpoints into `params` via `inject_consumption_points`
+before calling `model.solve()` / `model.simulate()`.
+"""
+
+from collections.abc import Mapping
+from typing import Any
+
+import jax.numpy as jnp
+from jax import Array
+from lcm import IrregSpacedGrid, Model
+
+MAX_CONSUMPTION: float = 300_000.0
+"""Upper bound of the consumption grid in $/year. Brackets the unconstrained
+CRRA optimum for the highest-asset, highest-income agents in the state space."""
+
+
+def compute_consumption_points(
+    *, consumption_floor: float, n_points: int
+) -> Array:
+    """Return log-spaced consumption gridpoints from the floor to `MAX_CONSUMPTION`.
+
+    Args:
+        consumption_floor: Lowest gridpoint, equal to the `consumption_floor`
+            parameter so the agent cannot pick `c < floor` even when saving
+            from a transfer top-up.
+        n_points: Total number of gridpoints.
+
+    Returns:
+        1-D float array of length `n_points`.
+    """
+    return jnp.geomspace(consumption_floor, MAX_CONSUMPTION, num=n_points)
+
+
+def inject_consumption_points(
+    *,
+    params: Mapping[str, Any],
+    model: Model,
+    consumption_floor: float | None = None,
+) -> dict[str, Any]:
+    """Inject consumption gridpoints into per-regime params.
+
+    Walks `model.regimes`, finds those with `consumption` declared as
+    `IrregSpacedGrid` with runtime-supplied points, and writes
+    `params[regime_name]["consumption"] = {"points": <pts>}`.
+
+    Args:
+        params: Existing params mapping. Returned as a new dict; the input is
+            not mutated.
+        model: Model whose regime specs determine which regimes need points.
+        consumption_floor: Lowest gridpoint. When `None`, taken from
+            `params["consumption_floor"]`.
+
+    Returns:
+        New params dict with consumption points injected.
+    """
+    if consumption_floor is None:
+        consumption_floor = float(params["consumption_floor"])
+    out: dict[str, Any] = dict(params)
+    for regime_name, regime in model.regimes.items():
+        grid = regime.actions.get("consumption")
+        if not (
+            isinstance(grid, IrregSpacedGrid) and grid.pass_points_at_runtime
+        ):
+            continue
+        points = compute_consumption_points(
+            consumption_floor=consumption_floor, n_points=grid.n_points
+        )
+        regime_entry = dict(out.get(regime_name, {}))
+        regime_entry["consumption"] = {"points": points}
+        out[regime_name] = regime_entry
+    return out

tests/test_benchmark.py

@@ -13,7 +13,7 @@
 def test_benchmark_model_simulates_end_to_end() -> None:
     n_subjects = 20
     model = create_benchmark_model()
-    _, params = get_benchmark_params()
+    _, params = get_benchmark_params(model=model)
     initial_conditions = get_benchmark_initial_conditions(
         model=model, n_subjects=n_subjects, seed=0
     )