Stream the asset-row refine-to-query into the FUES scan

In asset-row mode the per-node solve refines a full NaN-padded upper
envelope of static length n_pad and interpolates it at exactly one
query (resources_at_node) to publish a scalar (V_node, policy_node),
then discards it. Batched over combos x Euler-nodes, that n_pad scratch
axis is the binder that OOMs large solves.

Fold the single-query interpolation into the upper-envelope scan so the
n_pad rows never materialize:

- Extract `_interp_between_nodes` in interp.py — the pure two-bracket-node
  arithmetic plus the Hermite correction, shared by `interp_on_prepared_grid`
  and the streamed path. The existing interp unit tests pin it (behavior-
  preserving). This guarantees the streamed value cannot diverge from
  row-then-interp: only which two nodes differs.
- Add `refine_to_bracket` + `QueryBracket` in fues.py — a new scan driver
  reusing `_inspect_candidate` verbatim, changing only the emission sink: each
  step's up-to-3 emitted points fold into an O(1) bracket-capture carry
  (first/second, rolling prev/last, lo = latest emitted with grid <= q, hi =
  first with grid > q, running counts), no [n_input,3] stack and no [n_pad]
  rows. Post-scan, the bracket reproduces
  clip(searchsorted(side="right"), 1, max(n_kept-1, 1)) node-for-node: the
  side="right" tie-break at a duplicated kink (right copy wins the lower slot),
  the below-first clamp (first, second), and the at-or-above-last clamp
  (second-last, last; single-live falls back to the NaN-padded slot like the
  row path). Geometry only — no utility, borrowing limit, or floor.
- Add `publish_node_from_bracket` in asset_row.py — the scalar-bracket
  counterpart of `_publish_node_V_and_policy`: the shared `_interp_between_nodes`
  with the value Hermite slope = grad(utility) at the two bracket policies only,
  plus the constrained floor and the n_kept > n_pad overflow poison, preserved
  identically. Wire the asset-row node solve to refine_to_bracket ->
  publish_node_from_bracket and drop the n_pad envelope from that path.

Scope: asset-row mode only. Single-post-state mode publishes the refined
envelope AS its inter-period carry (queried later at many parent points), so
step_core's single-post-state path keeps calling `refine` unchanged.

The new equivalence test pins the streamed pair against refine +
_publish_node_V_and_policy for the same candidates + query, to fp tolerance:
smooth, kinked, multi-crossing, all-dead, single-live, overflow, and query
below-first / above-last / exactly on a duplicated kink abscissa.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: hmgaudecker

src/_lcm/egm/asset_row.py

@@ -24,6 +24,7 @@
     bind_continuation,
 )
 from _lcm.egm.interp import (
+    _interp_between_nodes,
     interp_on_padded_grid,
 )
 from _lcm.egm.step_core import (
@@ -32,6 +33,9 @@
     _EgmKernelPieces,
     _get_compute_node,
 )
+from _lcm.egm.upper_envelope.fues import (
+    QueryBracket,
+)
 from _lcm.typing import (
     RegimeName,
 )
@@ -171,17 +175,19 @@ def utility_of_action(action_value: ScalarFloat) -> ScalarFloat:
             # Same `-inf` masking as the default per-combo computation: dead
             # candidates become the envelope scan's absent form (NaN).
             candidate_dead = jnp.isneginf(candidate_value)
-            refined_grid, refined_policy, refined_value, n_kept = pieces.refine(
+            # The node reads its refined envelope at exactly one query
+            # (`resources_at_node`), so the scan streams the bracketing pair
+            # instead of materializing the NaN-padded `n_pad` envelope rows —
+            # the per-(combo, node) envelope working set is O(1), not O(n_pad).
+            bracket = pieces.refine_to_bracket(
                 endog_grid=jnp.where(candidate_dead, jnp.nan, candidate_grid),
                 policy=jnp.where(candidate_dead, jnp.nan, candidate_policy),
                 value=jnp.where(candidate_dead, jnp.nan, candidate_value),
+                x_query=resources_at_node,
             )
 
-            V_node, policy_node = _publish_node_V_and_policy(
-                refined_grid=refined_grid,
-                refined_policy=refined_policy,
-                refined_value=refined_value,
-                n_kept=n_kept,
+            V_node, policy_node = publish_node_from_bracket(
+                bracket=bracket,
                 n_pad=pieces.n_pad,
                 resources_at_node=resources_at_node,
                 borrowing_limit=pieces.borrowing_limit,
@@ -372,3 +378,127 @@ def _publish_node_V_and_policy(
     V_node = jnp.where(overflowed, jnp.nan, V_node).astype(dtype)
     policy_node = jnp.where(overflowed, jnp.nan, policy_node).astype(dtype)
     return V_node, policy_node
+
+
+def publish_node_from_bracket(
+    *,
+    bracket: QueryBracket,
+    n_pad: int,
+    resources_at_node: ScalarFloat,
+    borrowing_limit: ScalarFloat,
+    utility_of_action: Callable[[ScalarFloat], ScalarFloat],
+    discounted_expected_value_at_limit: ScalarFloat,
+) -> tuple[ScalarFloat, ScalarFloat]:
+    """Publish one asset node's value and optimal action from its query bracket.
+
+    The streamed counterpart of `_publish_node_V_and_policy`: it consumes the
+    two envelope nodes that `refine_to_bracket` captured around
+    `resources_at_node` instead of the full NaN-padded refined row, so the
+    `n_pad` envelope is never materialized. The published economics are
+    identical:
+
+    - The value is the cubic-Hermite read of the envelope between the two
+      bracket nodes (the value slope at each node is `grad(utility_of_action)`
+      at that node's policy, the envelope-theorem marginal), floored at the
+      closed-form constrained value, which is a feasible-policy lower bound.
+    - Below the lowest envelope node the closed-form constrained value is
+      published outright; the winning branch sets the action (the closed-form
+      `R - borrowing_limit` where constrained wins, the interpolated policy
+      otherwise).
+    - Envelope overflow (`n_kept > n_pad`) NaN-poisons both outputs, identical
+      to the row path, so the solve loop's NaN diagnostics surface the offending
+      (regime, period).
+
+    Because the value and policy arithmetic is the shared `_interp_between_nodes`
+    primitive — the same one the row path reaches through
+    `interp_on_padded_grid` — the streamed publish cannot diverge from
+    row-then-interpolate: only the bracket-finding differs.
+
+    Args:
+        bracket: The query bracket from `refine_to_bracket`.
+        n_pad: Static length of the envelope-refinement workspace (the overflow
+            threshold).
+        resources_at_node: Resources at this exogenous Euler node (the row's
+            single publish query).
+        borrowing_limit: Lower bound of the savings grid.
+        utility_of_action: Utility with everything but the continuous action
+            bound.
+        discounted_expected_value_at_limit: Discounted expected continuation
+            value at the lowest savings node.
+
+    Returns:
+        Tuple of the node's published value and published optimal action.
+
+    """
+    dtype = resources_at_node.dtype
+    overflowed = bracket.n_kept > n_pad
+
+    # The value Hermite slope is the envelope-theorem marginal `u'(c*)`, masked
+    # exactly as the row path: NaN where the node's policy is NaN (a padded
+    # slot), 0.0 where the node's value is `-inf` (an infeasible endpoint), so
+    # `_interp_between_nodes` falls back to the linear rule on those brackets.
+    slope_lower = _node_value_slope(
+        policy=bracket.lower_policy,
+        value=bracket.lower_value,
+        utility_of_action=utility_of_action,
+    )
+    slope_upper = _node_value_slope(
+        policy=bracket.upper_policy,
+        value=bracket.upper_value,
+        utility_of_action=utility_of_action,
+    )
+
+    value_interpolated = _interp_between_nodes(
+        x_query=resources_at_node,
+        xp_lower=bracket.lower_grid,
+        xp_upper=bracket.upper_grid,
+        fp_lower=bracket.lower_value,
+        fp_upper=bracket.upper_value,
+        slope_lower=slope_lower,
+        slope_upper=slope_upper,
+    )
+    policy_interpolated = _interp_between_nodes(
+        x_query=resources_at_node,
+        xp_lower=bracket.lower_grid,
+        xp_upper=bracket.upper_grid,
+        fp_lower=bracket.lower_policy,
+        fp_upper=bracket.upper_policy,
+    )
+
+    closed_form_action = resources_at_node - borrowing_limit
+    value_constrained = jnp.where(
+        closed_form_action > 0.0,
+        utility_of_action(jnp.maximum(closed_form_action, jnp.finfo(dtype).tiny))
+        + discounted_expected_value_at_limit,
+        -jnp.inf,
+    )
+    below_refined = (closed_form_action > 0.0) & (
+        resources_at_node <= bracket.first_grid
+    )
+    constrained_wins = below_refined | (value_constrained >= value_interpolated)
+    V_node = jnp.where(
+        below_refined,
+        value_constrained,
+        jnp.maximum(value_interpolated, value_constrained),
+    )
+    policy_node = jnp.where(constrained_wins, closed_form_action, policy_interpolated)
+    V_node = jnp.where(overflowed, jnp.nan, V_node).astype(dtype)
+    policy_node = jnp.where(overflowed, jnp.nan, policy_node).astype(dtype)
+    return V_node, policy_node
+
+
+def _node_value_slope(
+    *,
+    policy: ScalarFloat,
+    value: ScalarFloat,
+    utility_of_action: Callable[[ScalarFloat], ScalarFloat],
+) -> ScalarFloat:
+    """Envelope-theorem value slope at one bracket node, masked like the row.
+
+    The slope is `grad(utility_of_action)` at the node's policy — NaN where the
+    policy is a padded NaN slot, 0.0 where the node value is `-inf` — matching
+    the per-node masking the row path applies before interpolating.
+    """
+    slope = jax.grad(utility_of_action)(jnp.where(jnp.isnan(policy), 1.0, policy))
+    slope = jnp.where(jnp.isnan(policy), jnp.nan, slope)
+    return jnp.where(jnp.isneginf(value), 0.0, slope)

src/_lcm/egm/interp.py

@@ -151,10 +151,63 @@ def interp_on_prepared_grid(
         jnp.maximum(valid_length - 1, 1),
     ).astype(jnp.int32)
     lower = upper - 1
-    xp_lower = xp[lower]
-    fp_lower = fp[lower]
-    fp_upper = fp[upper]
-    bracket_width = xp[upper] - xp_lower
+    return _interp_between_nodes(
+        x_query=x_query,
+        xp_lower=xp[lower],
+        xp_upper=xp[upper],
+        fp_lower=fp[lower],
+        fp_upper=fp[upper],
+        slope_lower=None if fp_slopes is None else fp_slopes[lower],
+        slope_upper=None if fp_slopes is None else fp_slopes[upper],
+    )
+
+
+def _interp_between_nodes(
+    *,
+    x_query: FloatND,
+    xp_lower: FloatND,
+    xp_upper: FloatND,
+    fp_lower: FloatND,
+    fp_upper: FloatND,
+    slope_lower: FloatND | None = None,
+    slope_upper: FloatND | None = None,
+) -> FloatND:
+    """Interpolate a query between its two bracketing grid nodes.
+
+    The pure two-node arithmetic of the padded-grid interpolant, shared by
+    `interp_on_prepared_grid` (which gathers the bracket from a full row) and
+    the streamed asset-row publish (which captures the bracket directly during
+    the upper-envelope scan). Having both paths reduce to this one function
+    guarantees the streamed value cannot diverge from the row-then-interpolate
+    value: only *which two nodes* differs, not the arithmetic on them.
+
+    The bracket must already be edge-clamped to a real pair of nodes (queries
+    below the first node bracket the first pair, queries at or above the last
+    bracket the last pair). The interpolant is then edge-safe by construction:
+
+    - At a zero-width bracket (a duplicated kink abscissa) the relative position
+      is forced to `1.0`, so the right node's value applies and the zero width
+      is never used as a divisor.
+    - A `-inf` endpoint yields `-inf` wherever it carries positive weight
+      (instead of the NaN of `fp_lower + rel * (fp_upper - fp_lower)`) and
+      contributes exactly nothing at weight zero.
+
+    Args:
+        x_query: Point(s) at which to evaluate the interpolant.
+        xp_lower: Lower bracket node abscissa.
+        xp_upper: Upper bracket node abscissa.
+        fp_lower: Function value at the lower node.
+        fp_upper: Function value at the upper node.
+        slope_lower: Node derivative at the lower node; `None` selects linear
+            interpolation.
+        slope_upper: Node derivative at the upper node; `None` selects linear
+            interpolation.
+
+    Returns:
+        Interpolated value(s) with the shape of `x_query`.
+
+    """
+    bracket_width = xp_upper - xp_lower
     safe_width = jnp.where(bracket_width == 0.0, 1.0, bracket_width)
     # Zero-width brackets arise only when a duplicated abscissa sits at the end
     # of the non-NaN prefix; queries there are at or above the duplicate, so
@@ -172,16 +225,16 @@ def interp_on_prepared_grid(
     linear = jnp.where(weight_lower > 0.0, weight_lower * fp_lower, 0.0) + jnp.where(
         relative_position > 0.0, relative_position * fp_upper, 0.0
     )
-    if fp_slopes is None:
+    if slope_lower is None or slope_upper is None:
         return linear
     return linear + _hermite_correction(
         relative_position=relative_position,
         bracket_width=bracket_width,
         safe_width=safe_width,
         fp_lower=fp_lower,
         fp_upper=fp_upper,
-        slope_lower=fp_slopes[lower],
-        slope_upper=fp_slopes[upper],
+        slope_lower=slope_lower,
+        slope_upper=slope_upper,
     )
 
 

src/_lcm/egm/step.py

@@ -161,7 +161,7 @@
     _EgmKernelPieces,
     _get_solve_one_combo,
 )
-from _lcm.egm.upper_envelope import get_upper_envelope
+from _lcm.egm.upper_envelope import get_bracket_finder, get_upper_envelope
 from _lcm.egm.validation import _reachable_target_names, savings_stage_reads_euler_state
 from _lcm.engine import StateActionSpace
 from _lcm.grids import ContinuousGrid, Grid
@@ -754,6 +754,7 @@ def _build_kernel_pieces(
         ),
         build_H_kwargs=_get_build_H_kwargs(functions),
         refine=get_upper_envelope(solver=solver, n_refined=n_pad),
+        refine_to_bracket=get_bracket_finder(solver=solver),
         continuation_plan=continuation_plan,
     )
 

src/_lcm/egm/step_core.py

@@ -27,6 +27,7 @@
 from _lcm.egm.interp import (
     interp_on_padded_grid,
 )
+from _lcm.egm.upper_envelope.fues import QueryBracket
 from _lcm.typing import (
     ActionName,
     RegimeName,
@@ -101,7 +102,10 @@ class _EgmKernelPieces:
     """Closure assembling the Bellman aggregator's keyword arguments."""
 
     refine: Callable[..., tuple[Float1D, Float1D, Float1D, ScalarInt]]
-    """The configured upper-envelope backend."""
+    """The configured upper-envelope backend (single-post-state carry)."""
+
+    refine_to_bracket: Callable[..., QueryBracket]
+    """The streaming single-query bracket finder (asset-row publish)."""
 
     continuation_plan: ContinuationPlan
     """Build-time statics of the per-savings-node continuation aggregation."""

src/_lcm/egm/upper_envelope/__init__.py

@@ -8,11 +8,16 @@
 (`_lcm.egm.upper_envelope.fues`).
 """
 
+from collections.abc import Callable
 from typing import Protocol, runtime_checkable
 
-from _lcm.egm.upper_envelope.fues import refine_envelope
+from _lcm.egm.upper_envelope.fues import (
+    QueryBracket,
+    refine_envelope,
+    refine_to_bracket,
+)
 from lcm.solvers import DCEGM
-from lcm.typing import Float1D, ScalarInt
+from lcm.typing import Float1D, ScalarFloat, ScalarInt
 
 
 @runtime_checkable
@@ -69,3 +74,47 @@ def fues_backend(
 
     msg = f"Unknown upper-envelope backend: {solver.upper_envelope!r}."
     raise ValueError(msg)
+
+
+def get_bracket_finder(*, solver: DCEGM) -> Callable[..., QueryBracket]:
+    """Build the streaming single-query bracket finder for the asset-row solve.
+
+    The geometry-only counterpart of `get_upper_envelope` for asset-row mode,
+    where the refined envelope is read at exactly one query per node: it returns
+    the two bracketing envelope nodes (plus the first node and the kept count)
+    without materializing the NaN-padded `n_pad` rows. It is FUES-local and
+    deliberately not on the `UpperEnvelopeBackend` Protocol — the backend
+    returns envelope geometry; the asset-row module owns the EGM economics
+    (utility gradients, the borrowing limit, the constrained floor).
+
+    Args:
+        solver: The regime's DC-EGM solver configuration; the `fues_*` fields
+            parametrize the scan.
+
+    Returns:
+        The configured bracket finder.
+
+    """
+    if solver.upper_envelope == "fues":
+
+        def fues_bracket_finder(
+            *,
+            endog_grid: Float1D,
+            policy: Float1D,
+            value: Float1D,
+            x_query: ScalarFloat,
+        ) -> QueryBracket:
+            """Run the streaming FUES scan with the solver's thresholds."""
+            return refine_to_bracket(
+                endog_grid=endog_grid,
+                policy=policy,
+                value=value,
+                x_query=x_query,
+                jump_thresh=solver.fues_jump_thresh,
+                n_points_to_scan=solver.fues_n_points_to_scan,
+            )
+
+        return fues_bracket_finder
+
+    msg = f"Unknown upper-envelope backend: {solver.upper_envelope!r}."
+    raise ValueError(msg)