IrregSpacedGrid: extra_param_names for user-side runtime scalars

docs/user_guide/grids.md

@@ -98,6 +98,56 @@ are then supplied via the params dict:
 IrregSpacedGrid(n_points=4)
 ```
 
+#### Runtime points driven by `extra_param_names`
+
+When the gridpoints depend on a value that varies across solver iterations — a borrowing
+limit, a per-iteration upper bound, a calibrated cap — declare the dependency explicitly
+with `extra_param_names`:
+
+```python
+IrregSpacedGrid(n_points=64, extra_param_names=("max_consumption",))
+```
+
+`extra_param_names` is only valid together with runtime points (`points=None`). Each
+name is added to the params template alongside the grid's `points` slot, and
+`broadcast_to_template` carries it through even though no DAG function references it.
+The names exist so user-side code that *constructs* the gridpoints can read them without
+tripping the template validator's `Unknown keys` check.
+
+The typical workflow looks like this. Declare the grid with `extra_param_names`,
+populate the extras (e.g. as `model.fixed_params["max_consumption"]`), then inject the
+points before solving or simulating:
+
+```python
+import jax.numpy as jnp
+from lcm import IrregSpacedGrid, Model
+
+
+def inject_consumption_points(*, params: dict, model: Model) -> dict:
+    """Fill the runtime `consumption` gridpoints on every non-terminal regime."""
+    max_consumption = jnp.asarray(model.fixed_params["max_consumption"])
+    out: dict = dict(params)
+    for regime_name, regime in model.regimes.items():
+        if regime.terminal:
+            continue
+        grid = regime.actions["consumption"]
+        assert isinstance(grid, IrregSpacedGrid) and grid.pass_points_at_runtime
+        points = jnp.geomspace(1e-3, max_consumption, num=grid.n_points)
+        regime_entry = dict(out.get(regime_name, {}))
+        regime_entry["consumption"] = {"points": points}
+        out[regime_name] = regime_entry
+    return out
+
+
+params = inject_consumption_points(params=model.get_params_template(), model=model)
+model.solve(params=params)
+```
+
+`extra_param_names` carries any number of scalars (`("max_consumption", "floor")`,
+etc.). The construction logic in the user-side helper is free to combine them however
+the model needs — pylcm only validates that each declared name is present in the
+resolved params.
+
 ### PiecewiseLinSpacedGrid
 
 Multiple linearly spaced segments joined at breakpoints. Dense where you need precision,

src/lcm/grids/continuous.py

@@ -227,10 +227,21 @@ class IrregSpacedGrid(ContinuousGrid):
     omitted and only `n_points` is given, the points must be supplied at
     runtime via the params.
 
+    `extra_param_names` declares additional scalar parameters consumed by
+    user-side code that *constructs* the runtime points (e.g., a grid
+    upper bound that changes across optimizer iterations). These names
+    enter the params template alongside `points` and pass through
+    `broadcast_to_template` even though no DAG function references them
+    — the user's injection code reads them from the resolved params /
+    `model.fixed_params` to derive the points before calling
+    `solve` / `simulate`.
+
     Example:
     --------
     Fixed grid: `IrregSpacedGrid(points=[-1.73, -0.58, 0.58, 1.73])` Grid that is only
     completed at runtime via params: `IrregSpacedGrid(n_points=4)`
+    Grid that pairs runtime `points` with an extra scalar bound:
+    `IrregSpacedGrid(n_points=4, extra_param_names=("max_consumption",))`
 
     """
 
@@ -240,12 +251,23 @@ class IrregSpacedGrid(ContinuousGrid):
     n_points: int
     """Number of points. Derived from `len(points)` when points are given."""
 
+    extra_param_names: tuple[str, ...]
+    """Names of additional scalar params surfaced in the template.
+
+    Only meaningful when points are supplied at runtime
+    (`pass_points_at_runtime=True`); pylcm itself never reads these
+    values — they're carried through the params machinery so user-side
+    injection code can pick them up without fighting the `Unknown keys`
+    validator.
+    """
+
     def __init__(
         self,
         *,
         points: Sequence[float] | Float1D | None = None,
         n_points: int | None = None,
         batch_size: int = 0,
+        extra_param_names: tuple[str, ...] = (),
     ) -> None:
         if points is not None:
             _validate_irreg_spaced_grid(points)
@@ -269,9 +291,16 @@ def __init__(
             )
         else:
             stored_points = None
+        if extra_param_names and stored_points is not None:
+            raise GridInitializationError(
+                "`extra_param_names` is only valid when points are supplied at "
+                "runtime (i.e. `points=None`); a fixed-points grid has no "
+                "user-side params to thread through."
+            )
         object.__setattr__(self, "points", stored_points)
         object.__setattr__(self, "n_points", n_points)
         object.__setattr__(self, "batch_size", batch_size)
+        object.__setattr__(self, "extra_param_names", tuple(extra_param_names))
 
     @property
     def pass_points_at_runtime(self) -> bool:

src/lcm/params/regime_template.py

@@ -92,7 +92,7 @@ def _add_runtime_grid_params(
             _fail_if_runtime_grid_shadows_function(
                 function_params=function_params, name=state_name, kind="state"
             )
-            function_params[state_name] = {"points": "Float1D"}
+            function_params[state_name] = _irreg_grid_template_entry(grid)
         elif isinstance(grid, _ShockGrid) and grid.params_to_pass_at_runtime:
             _fail_if_runtime_grid_shadows_function(
                 function_params=function_params,
@@ -108,7 +108,22 @@ def _add_runtime_grid_params(
             _fail_if_runtime_grid_shadows_function(
                 function_params=function_params, name=action_name, kind="action"
             )
-            function_params[action_name] = {"points": "Float1D"}
+            function_params[action_name] = _irreg_grid_template_entry(grid)
+
+
+def _irreg_grid_template_entry(grid: IrregSpacedGrid) -> dict[str, str]:
+    """Template slots for a runtime-points `IrregSpacedGrid`.
+
+    Always exposes `points: Float1D`. Each entry in
+    `grid.extra_param_names` adds a `ScalarFloat` slot so user-side
+    injection code can thread its scalar bounds through `fixed_params`
+    or per-iteration params without tripping `broadcast_to_template`'s
+    unknown-keys check.
+    """
+    entry: dict[str, str] = {"points": "Float1D"}
+    for name in grid.extra_param_names:
+        entry[name] = "ScalarFloat"
+    return entry
 
 
 def _fail_if_runtime_grid_shadows_function(

tests/test_runtime_params.py

@@ -313,3 +313,58 @@ def test_simulate_with_runtime_action_grid_no_nan() -> None:
     )
     df = result.to_dataframe()
     assert not df["value"].isna().any()
+
+
+# --- extra_param_names: scalar params consumed by user-side injection code ---
+
+
+def test_extra_param_names_added_to_template():
+    """`extra_param_names` show up in the action's template alongside `points`."""
+    model = _make_action_grid_model(
+        consumption_grid=IrregSpacedGrid(
+            n_points=5,
+            extra_param_names=("max_consumption",),
+        ),
+    )
+    alive_template = model._params_template["alive"]
+    assert alive_template["consumption"] == {
+        "points": "Float1D",
+        "max_consumption": "ScalarFloat",
+    }
+
+
+def test_extra_param_names_accepted_via_fixed_params():
+    """Model-level `fixed_params` with extra-grid-param keys broadcast cleanly."""
+    grid = IrregSpacedGrid(n_points=5, extra_param_names=("max_consumption",))
+    model_fixed = _make_action_grid_model(consumption_grid=grid)
+    # Build via fresh Model to inject `fixed_params`.
+    alive = model_fixed.regimes["alive"]
+    dead = model_fixed.regimes["dead"]
+    model = Model(
+        regimes={"alive": alive, "dead": dead},
+        ages=AgeGrid(start=0, stop=2, step="Y"),
+        regime_id_class=RegimeId,
+        fixed_params={"max_consumption": 5.0},
+    )
+    params = {
+        "discount_factor": 0.95,
+        "interest_rate": 0.05,
+        "alive": {"consumption": {"points": jnp.linspace(0.1, 5.0, 5)}},
+    }
+    period_to_regime_to_V_arr = model.solve(params=params, log_level="off")
+    assert len(period_to_regime_to_V_arr) > 0
+
+
+def test_extra_param_names_rejected_on_fixed_points_grid():
+    """`extra_param_names` is meaningless when points are baked in at construction."""
+    with pytest.raises(Exception, match="only valid when points are supplied"):
+        IrregSpacedGrid(points=[1.0, 2.0, 3.0], extra_param_names=("foo",))
+
+
+def test_extra_param_names_empty_by_default():
+    """No `extra_param_names` keeps the template's grid entry to just `points`."""
+    model = _make_action_grid_model(
+        consumption_grid=IrregSpacedGrid(n_points=5),
+    )
+    alive_template = model._params_template["alive"]
+    assert alive_template["consumption"] == {"points": "Float1D"}