Runtime-supplied points on action grids; tighten grid + regime validators

[hmgaudecker]

Summary

Three threads bundled into one PR:

1. Feature. IrregSpacedGrid runtime-supplied points (previously state-only) extended to continuous action grids. Action grids declared as IrregSpacedGrid(n_points=N) now contribute {action_name: {"points": "Float1D"}} to the regime params template; InternalRegime.state_action_space(regime_params=...) substitutes runtime points into continuous_actions at solve and simulate time. Motivation: lets a model tie its consumption (or other continuous-action) grid bounds to a per-iteration parameter without rebuilding the model each estimation step.

2. Two silent-failure fixes that surfaced while completing the feature.

   • Simulate was using the NaN placeholder for runtime action grids. create_regime_state_action_space (forward simulation) bypassed the substitution that solve uses, so state_action_space.continuous_actions carried the compile-time NaN placeholder. That fed into argmax_and_max_Q_over_a and _lookup_values_from_indices, optimal actions came back NaN, the source regime's next_state propagated NaN to every target regime's namespaced state, and validate_V raised on the first downstream regime whose utility depended on those states. Now routes through the same substitution path solve uses; restores _validate_all_states_present on the per-subject states overlay.
   • IrregSpacedGrid.to_jax() silently returned a placeholder. Previously runtime grids returned jnp.full(N, NaN) from to_jax(), so any caller forgetting to call state_action_space(regime_params=...) got an undetectably broken grid. Now raises with a message pointing the caller at the substituted-grid path.

3. Adjacent validator hardening (also surfaced while completing the feature).

   • Function-output / discrete-grid-indexed-input clash detector. A regime function f that takes a discrete grid g (state, action, or derived categorical) as input has a per-cell scalar output, and a consumer that does f[g] raises IndexError: Too many indices: array is 0-dimensional, but 1 were indexed at trace time — well after the user could see what they wrote. New AST validator catches the pattern at Regime.__post_init__. The clashing function in the original spec produced this footgun and survived all existing validators.
   • The validator only fires when the producing function actually takes the discrete grid as input — verified by exercising both shapes with the validator disabled (func_output[grid] is correct code when the producer doesn't take grid; the validator now correctly exempts that shape).
   • Log-spaced grids rejected non-positive start. LogSpacedGrid(start=0, …) (or negative) silently produced NaN/-inf via to_jax() because the base validator only checked start < stop. Now refused at construction. Tightened two related silent-failure modes for free: _validate_continuous_grid rejects non-finite start/stop (the start >= stop check is False for NaN, so a NaN bound previously slipped through), and _validate_irreg_spaced_grid rejects non-finite points (the ascending-order test uses >=, also False for NaN, so an all-NaN runtime-supplied points array previously passed the order check silently).

Changes

• src/lcm/grids/continuous.py — LogSpacedGrid.__post_init__ enforces start > 0; _validate_continuous_grid rejects non-finite bounds via jnp.isfinite; _validate_irreg_spaced_grid rejects non-finite points; IrregSpacedGrid.to_jax() raises with a message pointing at state_action_space(regime_params=...).states[name] / .continuous_actions[name].
• src/lcm/params/regime_template.py — walk regime.actions alongside regime.states for runtime-grid params; shared _fail_if_runtime_grid_shadows_function helper.
• src/lcm/interfaces.py — InternalRegime.state_action_space(regime_params=...) builds both state and continuous-action replacements in one pass and returns via _base_state_action_space.replace(...).
• src/lcm/state_action_space.py — _grid_to_jax_or_placeholder private helper centralises the NaN-placeholder convention; deep-module ordering (public function on top, helpers below).
• src/lcm/regime_building/validation.py — new _validate_function_output_grid_indexing AST validator covering states, actions, and derived categoricals; only fires when the producing function takes the discrete grid as input.
• src/lcm/simulation/transitions.py + src/lcm/simulation/simulate.py — create_regime_state_action_space routes through internal_regime.state_action_space(regime_params=...) and overlays per-subject states with _validate_all_states_present.
• src/lcm/pandas_utils.py — _is_runtime_grid_param also recognises action grids.
• tests/test_runtime_params.py — TDD tests for runtime action grids: template entry, solve works, runtime ≡ fixed equivalence, varying runtime points changes V, simulate path returns finite V.
• tests/test_function_output_grid_indexing.py — TDD tests for the validator across all three discrete-grid sources, plus a regression test that the array-valued-producer + state-indexed-consumer shape (correct code) is not flagged.
• tests/test_grids.py — TDD tests for the new log-grid / non-finite guards.
• tests/test_single_feasible_action.py — regression tests pinning down map_coordinates NaN-on-non-finite-coordinate, get_irreg_coordinate divide-by-zero on duplicate points, runtime-state-grid validate_initial_conditions placeholder feasibility, and CRRA-bequest-with-discrete-state-indexing under jnp.where.

Test plan

• pixi run -e tests-cpu tests (880 passed, 5 skipped)
• pixi run ty
• prek run --all-files
• aca-model long-running end-to-end test (test_benchmark_model_simulates_end_to_end) — full 18-regime baseline solve+simulate with runtime-supplied consumption points
• benchmark CI on a clean GPU: aca-baseline +7% peak GPU mem, +10% exec time, +1% compile — all within run-to-run noise

🤖 Generated with Claude Code

[read-the-docs-community]

Documentation build overview

📚 pylcm | 🛠️ Build #32523614 | 📁 Comparing f18d27c against latest (8f2a4cf)

  🔍 Preview build  

32 files changed · ± 32 modified

± Modified

• index.html
• approximating-continuous-shocks/index.html
• benchmarking/index.html
• benchmarking-1/index.html
• beta-delta/index.html
• conventions/index.html
• debugging/index.html
• defining-models/index.html
• dispatchers/index.html
• function-representation/index.html
• and 22 more...

[github-actions]

Benchmark comparison (main → HEAD)

Comparing 8f2a4cfc (main) → f18d27c7 (HEAD)

Benchmark	Statistic	before	after	Ratio	Alert
aca-baseline	execution time	41.879 s	47.643 s	1.14	❌
	peak GPU mem	539 MB	579 MB	1.07
	compilation time	385.10 s	399.41 s	1.04
	peak CPU mem	8.08 GB	8.16 GB	1.01
Mahler-Yum	execution time	5.261 s	4.742 s	0.90
	peak GPU mem	522 MB	522 MB	1.00
	compilation time	17.01 s	15.99 s	0.94
	peak CPU mem	1.69 GB	1.69 GB	1.00
Precautionary Savings - Solve	execution time	52.7 ms	45.0 ms	0.85
	peak GPU mem	101 MB	101 MB	1.00
	compilation time	2.61 s	2.60 s	1.00
	peak CPU mem	1.12 GB	1.13 GB	1.01
Precautionary Savings - Simulate	execution time	146.7 ms	119.1 ms	0.81
	peak GPU mem	340 MB	340 MB	1.00
	compilation time	6.08 s	6.00 s	0.99
	peak CPU mem	1.28 GB	1.28 GB	1.00
Precautionary Savings - Solve & Simulate	execution time	166.1 ms	148.1 ms	0.89
	peak GPU mem	577 MB	577 MB	1.00
	compilation time	8.30 s	8.17 s	0.98
	peak CPU mem	1.28 GB	1.27 GB	1.00
Precautionary Savings - Solve & Simulate (irreg)	execution time	299.6 ms	282.8 ms	0.94
	peak GPU mem	2.19 GB	2.19 GB	1.00
	compilation time	8.72 s	8.70 s	1.00
	peak CPU mem	1.33 GB	1.33 GB	1.00

[hmgaudecker]

Code review

Found 6 issues (reporting all, no score filter applied per request).

1. Decorative section-separator comment blocks in a new test file (AGENTS.md says "Never add decorative section-separator comments like # ---...---. Code structure should be self-evident from function names and ordering.") Three blocks of dashed banners were added.

pylcm/tests/test_single_feasible_action.py

Lines 246 to 251 in db6214f

	# ---------------------------------------------------------------------------
	# Replicas of the aca-baseline failure path: dead regime with a CRRA bequest
	# whose `gamma` is per-pref_type, evaluated through `jnp.where`.
	# ---------------------------------------------------------------------------

pylcm/tests/test_single_feasible_action.py

Lines 397 to 408 in db6214f

	# ---------------------------------------------------------------------------
	# Direct probe: `map_coordinates` produces NaN at ±inf / NaN coordinates.
	# This is the concrete NaN source — `lower_weight = 1 - inf = -inf` and
	# `upper_weight = inf` combined with positive grid values gives `inf - inf =
	# NaN`. The aca-baseline NaN-in-V at age 51 is most plausibly traced back to
	# *some* upstream computation (next_assets / next_aime, or a state coordinate
	# from `get_irreg_coordinate` / `get_*_coordinate` that divides by zero on a
	# degenerate grid segment) producing inf, which then poisons the value
	# function via this interpolation path.
	# ---------------------------------------------------------------------------

pylcm/tests/test_single_feasible_action.py

Lines 448 to 457 in db6214f

	# ---------------------------------------------------------------------------
	# `validate_initial_conditions` uses `_base_state_action_space` directly,
	# which still holds the placeholder zeros for runtime-supplied
	# `IrregSpacedGrid`. With a feasibility constraint that the all-zero
	# placeholder fails, every subject is reported infeasible — even though the
	# real (post-substitution) grid would pass. This affects runtime grids
	# regardless of whether they are state or action grids.
	# ---------------------------------------------------------------------------

2. Test helpers added in this PR have parameters with no type annotations (AGENTS.md says "Type Hints — Always use type hints in all function signatures. This is mandatory."). _make_action_grid_model(*, consumption_grid) and _make_action_grid_model_with_stateful_dead(*, consumption_grid) lack annotations on consumption_grid; _crra_bequest and _alive_utility similarly omit annotations on pref_type, consumption_weight, coefficient_rra.

pylcm/tests/test_runtime_params.py

Lines 153 to 175 in db6214f

	def _make_action_grid_model(*, consumption_grid):
	"""Create a 2-regime model where consumption is the runtime-points action grid."""
	alive = Regime(
	functions={"utility": _utility},
	states={"wealth": LinSpacedGrid(start=1, stop=10, n_points=5)},
	state_transitions={"wealth": _next_wealth},
	actions={"consumption": consumption_grid},
	constraints={"borrowing_constraint": _borrowing_constraint},
	transition=_next_regime,
	active=lambda age: age < 2,
	)
	dead = Regime(
	transition=None,
	functions={"utility": lambda: 0.0},
	active=lambda age: age >= 2,
	)
	return Model(
	regimes={"alive": alive, "dead": dead},
	ages=AgeGrid(start=0, stop=2, step="Y"),
	regime_id_class=RegimeId,
	)

pylcm/tests/test_runtime_params.py

Lines 243 to 286 in db6214f

	def _make_action_grid_model_with_stateful_dead(*, consumption_grid):
	"""Variant where `dead` has a `wealth` state so its utility depends on it.
	Mirrors the aca-model dead regime (carries assets / pref_type so the
	bequest function can read them). Used to surface NaN propagation
	when the simulate path forgets to substitute runtime-supplied action
	gridpoints.
	"""
	def _alive_utility(
	consumption: ContinuousAction, wealth: ContinuousState
	) -> FloatND:
	return jnp.log(consumption + 1) + 0.01 * wealth
	def _dead_utility(wealth: ContinuousState) -> FloatND:
	return jnp.log(wealth + 1)
	alive = Regime(
	functions={"utility": _alive_utility},
	states={"wealth": LinSpacedGrid(start=1, stop=10, n_points=5)},
	state_transitions={
	"wealth": {
	"alive": _next_wealth,
	"dead": _next_wealth,
	},
	},
	actions={"consumption": consumption_grid},
	constraints={"borrowing_constraint": _borrowing_constraint},
	transition=_next_regime,
	active=lambda age: age < 2,
	)
	dead = Regime(
	transition=None,
	functions={"utility": _dead_utility},
	states={"wealth": LinSpacedGrid(start=1, stop=10, n_points=5)},
	active=lambda _age: True,
	)
	return Model(
	regimes={"alive": alive, "dead": dead},
	ages=AgeGrid(start=0, stop=2, step="Y"),
	regime_id_class=RegimeId,
	)

pylcm/tests/test_single_feasible_action.py

Lines 264 to 294 in db6214f

	def _crra_bequest(
	assets: ContinuousState,
	pref_type,
	bequest_shifter: float,
	consumption_weight,
	coefficient_rra,
	) -> FloatND:
	"""Replica of aca_model.agent.preferences.bequest, simplified.
	`consumption_weight` and `coefficient_rra` are FloatND indexed by
	`pref_type`. Both branches of the `jnp.where` are traced.
	"""
	alpha = consumption_weight[pref_type]
	gamma = coefficient_rra[pref_type]
	assets_shifted = jnp.maximum(0.0, assets) + bequest_shifter
	one_minus_gamma = jnp.where(jnp.isclose(gamma, 1.0), 1.0, 1.0 - gamma)
	return jnp.where(
	jnp.isclose(gamma, 1.0),
	jnp.log(assets_shifted),
	assets_shifted ** (one_minus_gamma * alpha) / one_minus_gamma,
	)
	def _alive_utility(
	consumption: ContinuousAction, pref_type, consumption_weight
	) -> FloatND:
	"""Make pref_type matter in alive's utility too (otherwise pylcm complains)."""
	alpha = consumption_weight[pref_type]
	return alpha * jnp.log(consumption)

3. Simulate's new create_regime_state_action_space no longer runs _validate_all_states_present (bug due to bypass: the previous path through create_state_action_space(states=...) validated that the per-subject states mapping covered every required state name; the new path calls base.replace(states=...) directly and skips that check). If a future caller forgets a state, the failure deferred to a less helpful site.

pylcm/src/lcm/simulation/transitions.py

Lines 52 to 61 in db6214f

	"""
	base = internal_regime.state_action_space(regime_params=regime_params)
	relevant_state_names = internal_regime.variable_info.query("is_state").index
	states_for_state_action_space = MappingProxyType(
	{sn: states[f"{internal_regime.name}__{sn}"] for sn in relevant_state_names}
	)
	return base.replace(states=states_for_state_action_space)

vs the dropped validation in:

pylcm/src/lcm/state_action_space.py

Lines 55 to 62 in db6214f

	else:
	_validate_all_states_present(
	provided_states=states,
	required_state_names=set(variable_info.query("is_state").index),
	)
	_states = states
	discrete_actions = {

4. New private helper has no docstring (AGENTS.md says "Docstrings — Use Google convention. Imperative mood in summary lines."). Sibling helpers in the same file (_add_runtime_grid_params, _collect_all_functions_for_template, _validate_no_shadowing) all carry docstrings; this one was missed.

pylcm/src/lcm/params/regime_template.py

Lines 107 to 121 in db6214f

	def _fail_if_runtime_grid_shadows_function(
	*,
	function_params: dict[FunctionName, dict[str, str]],
	name: str,
	kind: str,
	) -> None:
	if name in function_params:
	raise InvalidNameError(
	f"IrregSpacedGrid {kind} '{name}' (with runtime-supplied "
	f"points/params) conflicts with a function of the same name in the regime."
	)

5. Public create_regime_params_template docstring still says template entries match "the state name", but the PR extends the template to runtime-supplied action grids (the private helper _add_runtime_grid_params has the updated docstring). The public-facing description is now stale.

pylcm/src/lcm/params/regime_template.py

Lines 31 to 43 in db6214f

	dict that satisfies both phases.
	Grids with runtime-supplied values (IrregSpacedGrid without points,
	`_ShockGrid` without full shock_params) add entries to the template under
	pseudo-function keys matching the state name.
	Args:
	regime: The regime as provided by the user.
	Returns:
	The regime parameter template with type annotations as values.
	"""

6. _ShockGrid runtime-param substitution is gated by in_states and isinstance(spec, _ShockGrid) — silently a no-op if a _ShockGrid ever appears as a continuous action (bug due to asymmetry between the new IrregSpacedGrid branch which handles both in_states and in_continuous_actions). In practice _ShockGrid is state-only by design (intrinsic transitions, AGENTS.md forbids them as actions), so the path is unreachable today; flagging for completeness.

pylcm/src/lcm/interfaces.py

Lines 266 to 296 in db6214f

	)
	if not (in_states or in_continuous_actions):
	continue
	if isinstance(spec, IrregSpacedGrid) and spec.pass_points_at_runtime:
	points_key = f"{name}__points"
	if points_key not in all_params:
	continue
	if in_states:
	state_replacements[name] = cast(
	"ContinuousState", all_params[points_key]
	)
	else:
	action_replacements[name] = cast(
	"ContinuousAction", all_params[points_key]
	)
	elif (
	in_states
	and isinstance(spec, _ShockGrid)
	and spec.params_to_pass_at_runtime
	):
	all_present = all(
	f"{name}__{p}" in all_params for p in spec.params_to_pass_at_runtime
	)
	if not all_present:
	continue
	shock_kw: dict[str, float] = dict(spec.params)
	for p in spec.params_to_pass_at_runtime:
	shock_kw[p] = cast("float", all_params[f"{name}__{p}"])
	state_replacements[name] = spec.compute_gridpoints(**shock_kw)
	if not state_replacements and not action_replacements:

🤖 Generated with Claude Code

_{- If this code review was useful, please react with 👍. Otherwise, react with 👎.}

[hmgaudecker]

Autoreview.

@timmens @mj023: I am sorry this became larger once more -- I needed runtime-supplied points on an action grid, and this surfaced lots of other things which seemed in scope initially. Hope it's still reasonable, but if you want me to tear the PR apart, lmk!

[hmgaudecker]

Note: aca-baseline performance regression is because the model switched to using an IrregSpacedGrid with runtime-supplied points so XLA can't propagate the action values through algebraic simplifications

[mj023]

Everything else looks correct.

[mj023]

I would prefer to not do this check. I think it's quite intuitive that indexing like this is wrong and checking all user functions for mistakes like this is too much. I understand that it can be hard to debug just from the error message, but I think the examples are enough to show how to correctly use additional functions.

[hmgaudecker]

It's always a fine line between trying to error early and expecting a certain level of understanding from the user. In this case, it actually happened to me when I got a little lost during refactoring and forgot an instance.

It would be great if you could spell out what precisely you are worried about? E.g., runtime, fragility of the approach, ...?

[mj023]

I think checking for this is quite complicated and there are many ways to hide the pattern from the check and still get essentially the same error. Users are also supposed to mostly treat their functions like they were using scalars, so it should be clear that this won't work. I just want to avoid that we start checking the user supplied functions for all sorts of coding errors, as the same justification can be made for many other cases, like for example a user having a function utility( consumption, ..) and thinking they can index into the consumption variable consumption[0] because they provided consumption as a grid.

If you both want this check then I'm fine with it, but we should really try not to add too many of these, as they add complexity and are not very reliable.

[hmgaudecker]

I agree on the general sentiment. At the same time, I don't think this will hurt much and do at least some good (it certainly would have saved me a couple of hours had it been around last week).

So let's keep it in but remove it should it ever cause trouble!