diff --git a/docs/user_guide/grids.md b/docs/user_guide/grids.md index a15ae8c83..6d84e0844 100644 --- a/docs/user_guide/grids.md +++ b/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, diff --git a/src/lcm/grids/continuous.py b/src/lcm/grids/continuous.py index 839d4cefb..25503a00e 100644 --- a/src/lcm/grids/continuous.py +++ b/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: diff --git a/src/lcm/params/regime_template.py b/src/lcm/params/regime_template.py index 67ddfe75b..97c6c74e4 100644 --- a/src/lcm/params/regime_template.py +++ b/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( diff --git a/tests/test_runtime_params.py b/tests/test_runtime_params.py index 58fdfc27c..5064ab0e4 100644 --- a/tests/test_runtime_params.py +++ b/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"}