Merge remote-tracking branch 'origin/feat/nb-egm' into feat/nested-nbegm-ez

Author: hmgaudecker

src/_lcm/egm/published_policy.py

@@ -17,18 +17,30 @@
 action-grid argmax value of the continuous action, subject to a post-read
 feasibility check (in-support, finite, positive, within the intrinsic budget).
 
-The gate exists because the stored policy is the **solve-phase** optimum of
-one conditional problem, and a linear read is faithful only where the row
-carries the coordinates and branch topology it interpolates over. Kept on the
-grid-argmax path:
+Regimes with discrete actions publish one conditional row per discrete-action
+combination — value and policy per branch, on that branch's own endogenous
+grid. Simulation then *re-decides* the branch at the subject's state: each
+branch's conditional value is interpolated at that branch's own resources
+(discrete-only constraints mask infeasible branches to `-inf`), the feasible
+branch of highest interpolated value wins, and only the winner's policy is
+read. The value read uses the cubic Hermite interpolant with the
+`marginal_utility` row as exact node slopes — the same convention the solve
+uses to publish values — so the ranking the re-decision sees is the ranking
+the solve convention implies.
+
+The gate exists because the stored rows are the **solve-phase** optimum of
+one conditional problem each, and a read is faithful only where the rows
+carry the coordinates and branch topology they are interpolated over. Kept on
+the grid-argmax path:
 - regimes with any `Phased` declaration — a phase-variant utility, budget,
   transition, or state domain (not only `H`, e.g. naive present bias) makes
   the stored policy solve the wrong simulate-phase FOC or puts the policy
   rows on the wrong coordinates;
-- regimes with discrete actions — the branch is chosen from grid-restricted
-  Q values, and a branch whose refined conditional optimum lies between
-  action-grid nodes can lose that comparison yet win after continuous
-  refinement, so the refined policy could be paired with the wrong branch;
+- regimes whose upper-envelope backend does not certify every crossing —
+  only MSS inserts the exact segment crossing by construction; FUES decides
+  segment identity by a slope-threshold heuristic (no labels from the
+  kernel), so its row can silently bridge a missed switch, and RFC/LTM leave
+  switches between retained nodes;
 - regimes with a passive continuous state — each row is the envelope policy
   conditional on one passive node, so blending rows across a passive-dimension
   branch switch would read an action from neither branch;
@@ -39,14 +51,14 @@
   the per-node solve publishes one point per exogenous asset node, not a
   crossing-complete row, so interpolating across nodes would mix branches;
 - regimes with EV1 taste shocks, whose realized draws perturb the decision.
-Publishing conditional values and re-deciding the branch at the simulated
-state lifts the discrete-action, passive, process, and asset-row exclusions
-(the tracked follow-up).
+Publishing per-passive-node / per-process-node conditional values and
+re-deciding across those axes the way the discrete-action axis already is
+lifts the passive, process, and asset-row exclusions (the tracked follow-up).
 
 Unlike the rolling `EGMCarry` (the cross-period continuation channel, overwritten
 each period), this is saved for every solved period and travels to `simulate`
-alongside the value-function arrays. Its `endog_grid` rows are shared with the
-period's carry; only the `policy` row is additional state.
+alongside the value-function arrays. Its rows are shared with the period's
+carry; only the `policy` row is additional state.
 """
 
 from collections.abc import Sequence
@@ -64,17 +76,18 @@ class EGMSimPolicy:
 
     Leading axes match the regime's combo layout (discrete states, then passive
     states, then discrete actions, as in `EGMCarry`); the trailing axis is the
-    static refined-grid length. Both rows are NaN-padded in lockstep in the tail.
+    static refined-grid length. All rows are NaN-padded in lockstep in the tail.
     """
 
     endog_grid: FloatND
     """Endogenous grid in resources space, NaN-padded in the tail.
 
     Shared with the period's `EGMCarry.endog_grid`; weakly ascending per row.
-    Under a crossing-inserting upper-envelope backend (`fues`, `mss`) the
-    envelope-switch abscissae are duplicated with one-sided policy copies —
-    the topology the off-grid read requires; RFC/LTM rows leave switches
-    between retained nodes and do not qualify for the read.
+    Under MSS the envelope-switch abscissae are duplicated with one-sided
+    policy copies — the topology the off-grid read requires; FUES rows are
+    not certified crossing-complete (segment identity is a slope-threshold
+    heuristic) and RFC/LTM rows leave switches between retained nodes, so
+    neither qualifies for the read.
     """
 
     policy: FloatND
@@ -89,6 +102,15 @@ class EGMSimPolicy:
     the branch at the subject's state.
     """
 
+    marginal_utility: FloatND
+    """Marginal utility at `endog_grid` (NaN on padding slots).
+
+    Shared with the period's `EGMCarry.marginal_utility`: the value row's
+    exact node slope by the envelope theorem. Simulation passes it as the
+    slope row of the cubic Hermite value read, matching the interpolation
+    convention the solve publishes values under.
+    """
+
     row_discrete_state_names: tuple[StateName, ...] = ()
     """Names of the leading discrete-state row axes, in axis order."""
 
@@ -100,7 +122,7 @@ class EGMSimPolicy:
     """Names of the discrete-action row axes, after the passive states."""
 
 
-_EGM_SIM_POLICY_ARRAY_FIELDS = ("endog_grid", "policy", "value")
+_EGM_SIM_POLICY_ARRAY_FIELDS = ("endog_grid", "policy", "value", "marginal_utility")
 _EGM_SIM_POLICY_STATIC_FIELDS = (
     "row_discrete_state_names",
     "row_passive_state_names",

src/_lcm/egm/step.py

@@ -617,6 +617,7 @@ def solve_one_combo_over_axes(
                 endog_grid=grid_stack,
                 policy=policy_stack,
                 value=value_stack,
+                marginal_utility=marginal_stack,
                 row_discrete_state_names=own_discrete_state_names,
                 row_passive_state_names=own_passive_state_names,
                 row_discrete_action_names=tuple(own_discrete_action_values),
@@ -630,7 +631,10 @@ def solve_one_combo_over_axes(
                 taste_shock_scale=own_taste_shock_scale,
             )
             sim_policy = EGMSimPolicy(
-                endog_grid=grid_row, policy=policy_row, value=value_row
+                endog_grid=grid_row,
+                policy=policy_row,
+                value=value_row,
+                marginal_utility=marginal_row,
             )
         return V_arr, carry, sim_policy
 

src/_lcm/regime_building/processing.py

@@ -999,12 +999,12 @@ def _build_simulation_phase(
     #   keeper and adjuster candidates but publishes only the keeper's inner
     #   consumption function, so replaying it would pair an adjuster-won
     #   value with the keeper's policy;
-    # - a crossing-inserting upper-envelope backend — RFC and LTM leave the
-    #   envelope switch between two retained nodes, so linear policy
-    #   interpolation across the switch would mix two branch policies. MSS
-    #   inserts the exact crossing by construction; FUES does so only under the
-    #   exhaustive scan, so a bounded `fues_n_points_to_scan` is excluded too
-    #   (see `_envelope_publishes_crossings`);
+    # - a crossing-certifying upper-envelope backend (MSS only) — RFC and LTM
+    #   leave the envelope switch between two retained nodes, and FUES decides
+    #   segment identity by a slope-threshold heuristic with no labels from
+    #   the kernel, so its row can bridge a missed switch; only MSS inserts
+    #   the exact crossing by construction (see
+    #   `_envelope_publishes_crossings`);
     # - the single-post-state kernel (not asset-row mode) — when a
     #   savings-stage function reads the Euler state, DC-EGM solves per
     #   exogenous asset node and publishes one optimal point per node, not a
@@ -1066,21 +1066,22 @@ def _build_simulation_phase(
 
 
 def _envelope_publishes_crossings(solver: DCEGM) -> bool:
-    """Whether the solver's upper envelope inserts every segment crossing.
+    """Whether the solver's upper envelope certifies every segment crossing.
 
     A branch-faithful policy read interpolates a row whose envelope switches sit
     at duplicated abscissae carrying both branch policies:
-    - `"mss"` ⇒ always: it inserts the exact crossing by construction.
-    - `"fues"` ⇒ only under the exhaustive scan
-      (`fues_n_points_to_scan is None`). A bounded window may miss a segment's
-      continuation past the window's worth of interleaved off-segment
-      candidates, dropping that crossing and retaining the dominated
-      interlopers.
-    - `"rfc"` / `"ltm"` ⇒ never: the switch lands between retained nodes.
+    - `"mss"` ⇒ yes: the sweep inserts the exact crossing by construction.
+    - `"fues"` ⇒ no: segment identity is decided by thresholding the
+      implied-savings slope (`fues_jump_thresh`) — a heuristic — and the
+      DC-EGM kernel supplies no segment labels. Two value branches whose
+      cross-segment slope stays below the threshold merge into one row with
+      no crossing inserted; the row then bridges the switch, and neither an
+      exhaustive scan nor a wider window repairs that, because the scan can
+      only search within the segment identity it was given. A FUES row is
+      therefore not certified crossing-complete for the read.
+    - `"rfc"` / `"ltm"` ⇒ no: the switch lands between retained nodes.
     """
-    if solver.upper_envelope == "mss":
-        return True
-    return solver.upper_envelope == "fues" and solver.fues_n_points_to_scan is None
+    return solver.upper_envelope == "mss"
 
 
 def _regime_has_process_state(v_interpolation_info: VInterpolationInfo) -> bool:

src/_lcm/simulation/simulate.py

@@ -853,19 +853,34 @@ def _interp_rows_with_support(
     The live support runs from the row's first abscissa to its last finite one
     (rows are NaN-padded in the tail). Outside it the interpolant extends an
     edge secant (below) or clamps (above) — feasible values that need not
-    approximate the published function.
+    approximate the published function. Reads by field:
+    - `"value"` ⇒ cubic Hermite with the marginal-utility row as exact node
+      slopes — the convention the solve publishes values under, so the branch
+      ranking the re-decision sees is the solve convention's ranking;
+    - `"policy"` ⇒ piecewise linear (the policy row carries no slope data).
     """
     rows_f = getattr(sim_policy, field)
     rows_x = sim_policy.endog_grid
+    rows_slope = sim_policy.marginal_utility if field == "value" else None
     if index:
         rows_x = rows_x[index]
         rows_f = rows_f[index]
+        rows_slope = rows_slope[index] if rows_slope is not None else None
     if rows_x.ndim == 1:
         rows_x = jnp.broadcast_to(rows_x, (n_subjects, *rows_x.shape))
         rows_f = jnp.broadcast_to(rows_f, (n_subjects, *rows_f.shape))
-    values = vmap(
-        lambda x_query, xp, fp: interp_on_padded_grid(x_query=x_query, xp=xp, fp=fp)
-    )(resources, rows_x, rows_f)
+        if rows_slope is not None:
+            rows_slope = jnp.broadcast_to(rows_slope, (n_subjects, *rows_slope.shape))
+    if rows_slope is None:
+        values = vmap(
+            lambda x_query, xp, fp: interp_on_padded_grid(x_query=x_query, xp=xp, fp=fp)
+        )(resources, rows_x, rows_f)
+    else:
+        values = vmap(
+            lambda x_query, xp, fp, fp_slopes: interp_on_padded_grid(
+                x_query=x_query, xp=xp, fp=fp, fp_slopes=fp_slopes
+            )
+        )(resources, rows_x, rows_f, rows_slope)
     valid_length = jnp.sum(jnp.isfinite(rows_x), axis=-1)
     last_live = jnp.take_along_axis(rows_x, (valid_length - 1)[:, None], axis=-1)[:, 0]
     in_support = (resources >= rows_x[:, 0]) & (resources <= last_live)

tests/regime_building/test_egm_policy_read_gate.py

@@ -44,24 +44,30 @@ def test_non_crossing_envelope_backends_do_not_qualify(backend: str):
     assert model._regimes["retirement"].simulation.egm_policy_read is None
 
 
-@pytest.mark.parametrize("backend", ["fues", "mss"])
-def test_crossing_inserting_envelope_backends_qualify(backend: str):
-    """FUES/MSS rows qualify: duplicated crossing abscissae carry the switch."""
-    model = _retirement_model_with_backend(backend)
+def test_the_mss_backend_qualifies_for_the_policy_read():
+    """MSS rows qualify: the sweep inserts the exact segment crossing.
+
+    Duplicated crossing abscissae carry both branch policies, so linear
+    interpolation between retained nodes never mixes branches.
+    """
+    model = _retirement_model_with_backend("mss")
     assert model._regimes["retirement"].simulation.egm_policy_read is not None
 
 
-def test_bounded_fues_scan_does_not_qualify_for_the_policy_read():
-    """A finite FUES scan window keeps the grid path: its row may miss crossings.
+@pytest.mark.parametrize("n_points_to_scan", [None, 8])
+def test_fues_does_not_qualify_for_the_policy_read(n_points_to_scan: int | None):
+    """FUES rows keep the grid path: segment identity is a heuristic.
 
-    Only the exhaustive scan (`fues_n_points_to_scan=None`) is proven to find a
-    segment's continuation when more than the window's worth of off-segment
-    candidates interleave. A bounded window may drop that crossing and retain
-    the dominated interlopers, so the row lacks the switch topology a
-    branch-faithful policy read interpolates over.
+    FUES groups candidates into segments by thresholding the implied-savings
+    slope (`fues_jump_thresh`), and the DC-EGM kernel supplies no segment
+    labels. A cross-segment slope below the threshold merges two value
+    branches into one row — no crossing is inserted and the row bridges the
+    gap between them — regardless of the scan width, so even the exhaustive
+    scan does not certify the crossing topology the off-grid read
+    interpolates over (see `test_fues_segment_detection.py`).
     """
     solver = dataclasses.replace(
-        DCEGM_SOLVER, upper_envelope="fues", fues_n_points_to_scan=8
+        DCEGM_SOLVER, upper_envelope="fues", fues_n_points_to_scan=n_points_to_scan
     )
     model = _model_from_alive(
         dcegm_retirement.replace(active=lambda age: age < 50, solver=solver)

tests/simulation/test_egm_sim_policy_value_read.py

@@ -0,0 +1,58 @@
+"""The branch re-decision value read follows the solve's Hermite convention.
+
+The solve publishes conditional values by reading the refined row with the
+cubic Hermite interpolant, using the marginal-utility row as exact node slopes
+(envelope theorem). Simulation compares conditional branch values with the
+same interpolant, so the ranking the re-decision sees is the ranking the solve
+convention implies — a linear read of the same row can rank two close branches
+differently.
+"""
+
+import jax.numpy as jnp
+import pytest
+
+from _lcm.egm.published_policy import EGMSimPolicy
+from _lcm.simulation.simulate import _interp_rows_with_support
+
+
+def test_value_read_is_cubic_hermite_with_the_marginal_slopes():
+    """On nodes `(1, 0)`, `(2, log 2)` with slopes `(1, 1/2)`, the read at `1.5`
+    is the cubic Hermite value `0.40907`, not the linear chord `0.34657`."""
+    sim_policy = EGMSimPolicy(
+        endog_grid=jnp.array([1.0, 2.0]),
+        policy=jnp.array([0.5, 1.0]),
+        value=jnp.array([0.0, jnp.log(2.0)]),
+        marginal_utility=jnp.array([1.0, 0.5]),
+    )
+    value, in_support = _interp_rows_with_support(
+        sim_policy=sim_policy,
+        field="value",
+        index=(),
+        resources=jnp.array([1.5]),
+        n_subjects=1,
+    )
+    assert bool(in_support[0])
+    assert float(value[0]) == pytest.approx(0.4090736, abs=1e-6)
+
+
+def test_policy_read_stays_piecewise_linear():
+    """The policy read is the linear chord: `0.75` midway between `0.5` and `1.0`.
+
+    Only the value row carries exact node slopes (the marginal-utility row via
+    the envelope theorem); the policy row has no slope data, so its read is
+    piecewise linear.
+    """
+    sim_policy = EGMSimPolicy(
+        endog_grid=jnp.array([1.0, 2.0]),
+        policy=jnp.array([0.5, 1.0]),
+        value=jnp.array([0.0, jnp.log(2.0)]),
+        marginal_utility=jnp.array([1.0, 0.5]),
+    )
+    policy, _ = _interp_rows_with_support(
+        sim_policy=sim_policy,
+        field="policy",
+        index=(),
+        resources=jnp.array([1.5]),
+        n_subjects=1,
+    )
+    assert float(policy[0]) == pytest.approx(0.75, abs=1e-9)

tests/simulation/test_simulate_dcegm_branch_redecision.py

@@ -8,6 +8,8 @@
 comparison, exactly as they mask the grid argmax.
 """
 
+import dataclasses
+
 import jax.numpy as jnp
 import numpy as np
 
@@ -24,6 +26,7 @@
 from lcm_examples.iskhakov_et_al_2017 import WEALTH_GRID
 from tests.test_models.deterministic import retirement_only
 from tests.test_models.deterministic.dcegm_variants import (
+    DCEGM_SOLVER,
     dcegm_retirement,
     get_retirement_only_params,
 )
@@ -65,8 +68,12 @@ def _bequest_utility(wealth: ContinuousState, age: float) -> FloatND:
 
 
 def _bonus_model(constraints: dict | None = None) -> Model:
+    # MSS is the only backend whose rows certify every envelope crossing, so
+    # it is the only one that qualifies for the off-grid read.
+    solver = dataclasses.replace(DCEGM_SOLVER, upper_envelope="mss")
     alive = dcegm_retirement.replace(
         active=lambda age: age < 50,
+        solver=solver,
         actions={
             "consumption": dcegm_retirement.actions["consumption"],
             "take_bonus": DiscreteGrid(BonusChoice),

tests/simulation/test_simulate_dcegm_offgrid.py

@@ -13,6 +13,8 @@
 hit `c*` — which a grid-restricted argmax cannot.
 """
 
+import dataclasses
+
 import jax.numpy as jnp
 import numpy as np
 
@@ -22,6 +24,7 @@
 from lcm_examples.iskhakov_et_al_2017 import WEALTH_GRID, next_wealth_from_savings
 from tests.test_models.deterministic import base, dcegm_variants, retirement_only
 from tests.test_models.deterministic.dcegm_variants import (
+    DCEGM_SOLVER,
     dcegm_retirement,
     get_retirement_only_params,
 )
@@ -39,9 +42,14 @@ def _closed_form_model() -> Model:
         states={"wealth": LogSpacedGrid(start=0.25, stop=400.0, n_points=400)},
         functions={"utility": _bequest_utility},
     )
+    # MSS is the only backend whose rows certify every envelope crossing, so
+    # it is the only one that qualifies for the off-grid read.
+    solver = dataclasses.replace(DCEGM_SOLVER, upper_envelope="mss")
     return Model(
         regimes={
-            "retirement": dcegm_retirement.replace(active=lambda age: age < 50),
+            "retirement": dcegm_retirement.replace(
+                active=lambda age: age < 50, solver=solver
+            ),
             "dead": bequest_dead,
         },
         ages=AgeGrid(start=40, stop=50, step="10Y"),
@@ -95,7 +103,7 @@ def test_discrete_action_regime_consumption_leaves_the_action_grid():
     this pins that the off-grid path engages for a discrete-action regime.)
     """
     n_periods = 3
-    model = dcegm_variants.get_full_model("dcegm", n_periods)
+    model = dcegm_variants.get_full_model("dcegm", n_periods, upper_envelope="mss")
     params = dcegm_variants.get_full_params(n_periods)
 
     wealth_nodes = np.asarray(WEALTH_GRID.to_jax())