IrregSpacedGrid: extra_param_names for user-side runtime scalars (#348)

`broadcast_to_template` rejects any `fixed_params` key that no DAG
function lists in its signature. That assumption breaks for
runtime-supplied grids whose points are computed by user-side code
from per-iteration scalars (e.g., a grid upper bound) — pylcm itself
never reads those scalars, but they still need to flow through the
params machinery instead of being funneled through
`simulate(params=...)` where they compete with estimation parameters.

This adds an `extra_param_names: tuple[str, ...]` keyword to
`IrregSpacedGrid`. When `pass_points_at_runtime=True`, each name in
the tuple surfaces as a `ScalarFloat` slot in the action/state's
pseudo-function template entry alongside `points`. The user's
injection code reads the resolved values from `model.fixed_params`
(or per-iteration params) before computing the points; pylcm just
threads them through `broadcast_to_template` unchanged.

Rejected on fixed-points grids — a baked-in grid has no user-side
runtime computation to feed.

Closes #348.

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

Author: hmgaudecker

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"}