consumption_grid: read upper bound from max_consumption fixed param

The grid floor already tracks the per-iteration `consumption_floor`
parameter; the ceiling was a hardcoded 300k constant. Surface it as a
fixed param via a marker function (`consumption_grid_upper_bound`) so
callers can declare the bracket per model creation, and read it back
at inject time from each regime's `resolved_fixed_params`.

The marker function's output is intentionally unused — its only job
is to put `max_consumption` in the regime params template so pylcm's
fixed-param machinery captures it. dags.tree pruning drops the call
at solve / simulate.

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

Author: hmgaudecker

src/aca_model/baseline/regimes/_common.py

@@ -37,6 +37,7 @@
 from aca_model.baseline import health_insurance
 from aca_model.baseline.health_insurance import BuyPrivate
 from aca_model.config import GRID_CONFIG, MODEL_CONFIG, GridConfig
+from aca_model.consumption_grid import consumption_grid_upper_bound
 from aca_model.environment import social_security, taxes
 from aca_model.environment.social_security import ClaimedSS
 
@@ -537,6 +538,12 @@ def build_common_functions(spec: dict[str, str]) -> dict:
     functions["cash_on_hand"] = assets_and_income.cash_on_hand
     functions["transfers"] = assets_and_income.transfers
 
+    # Marker: surfaces `max_consumption` in the params template so it
+    # can be supplied via fixed_params and read back at inject time
+    # by `inject_consumption_points`. Output unused; pruned at
+    # solve / simulate.
+    functions["consumption_grid_upper_bound"] = consumption_grid_upper_bound
+
     return functions
 
 

src/aca_model/consumption_grid.py

@@ -1,10 +1,12 @@
 """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()`.
+`baseline.regimes._common.build_grids` so the bounds can track
+runtime parameters: the lower bound from the per-iteration
+`consumption_floor` parameter, the upper bound from the per-creation-time
+`max_consumption` fixed param. Callers must inject the actual
+gridpoints into `params` via `inject_consumption_points` before
+calling `model.solve()` / `model.simulate()`.
 """
 
 from collections.abc import Mapping
@@ -14,59 +16,72 @@
 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>}`.
 
+    Lower bound: `params["consumption_floor"]` (varies per iteration).
+    Upper bound: `max_consumption` from the regime's resolved
+    fixed-params (set once at model creation).
+
     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"])
+    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
+        # Runtime-points grids always have `n_points` set (the constructor
+        # rejects the (points=None, n_points=None) combo); narrow for ty.
+        assert grid.n_points is not None
+        max_consumption = float(
+            model.internal_regimes[regime_name].resolved_fixed_params["max_consumption"]
+        )
+        points = _compute_consumption_points(
+            consumption_floor=consumption_floor,
+            max_consumption=max_consumption,
+            n_points=grid.n_points,
         )
         regime_entry = dict(out.get(regime_name, {}))
         regime_entry["consumption"] = {"points": points}
         out[regime_name] = regime_entry
     return out
+
+
+def consumption_grid_upper_bound(max_consumption: float) -> float:
+    """Surface `max_consumption` in the regime params template.
+
+    pylcm builds the params template from each regime function's
+    signature. `max_consumption` is read at runtime by
+    `inject_consumption_points` from `resolved_fixed_params`; for
+    that to work via pylcm's fixed-params machinery, the key must
+    appear in some function's signature. This marker function is
+    the entry point — its output is intentionally unused, and
+    dags.tree pruning drops the call at solve / simulate time.
+    """
+    return max_consumption
+
+
+def _compute_consumption_points(
+    *,
+    consumption_floor: float,
+    max_consumption: float,
+    n_points: int,
+) -> Array:
+    """Return log-spaced consumption gridpoints from floor to max."""
+    return jnp.geomspace(consumption_floor, max_consumption, num=n_points)