Repair state-conditioned shocks per the PR #405 code review

pro-math-code-review (5.5-Pro, 2026-07-17) returned serious_gap on the built
primitive. All six findings reproduced against the real source and repaired.
The design was sound; the implementation had two independent law mismatches
that the entire 37-test suite passed over.

F1 (serious) — solve gathered the conditioned sigma, simulate did not. The
continuous next-state wrappers copied dict(grid.params) and called draw_shock,
so every category simulated at the scalar common-grid sigma. Measured on the
pre-fix code: the low regime drew std=0.3009 where its law says 0.05, a
variance 36x too large. This moves simulated states, moments and every later
decision, not just value levels. Both wrappers now take the conditioning state
and override sigma via gather_sigma, reusing sigma_array_by_code so the solve
and simulate orderings cannot drift apart.

F2 (serious) — both row builders dropped a fixed non-zero mu. Stock pylcm
builds Tauchen rows in demeaned coordinates, where the intercept vanishes; the
conditioned branch is handed physical nodes (centred on mu/(1-rho)) and a
physical from-value, so the conditional mean is mu + rho*y and the intercept
must appear. Threaded mu through conditioned_row and tauchen_row.

F3 — the conditioned branch is selected before the runtime-parameter
mechanism, so a grid parameter left for runtime is never bound: get_gridpoints
returns all-NaN and the closure captures it forever. Now require
is_fully_specified, finite nodes, finite positive sigmas and exact category
keys.

F4 — sigma is ordered by the target regime's grid but indexed by the source
state's code. Require one shared category-to-code map across every regime
carrying the conditioning state.

F5 — the tests guarded nothing: every reduction test used mu=0 and every model
test solved but never simulated. Adds 13 regressions, each verified to fail on
the pre-fix code.

F6 — the GH rejection was documented as blanket but implemented for IID only.
Now rejects Tauchen GH too, matching the stated v1 scope.

50 tests pass (37 before). Lint clean; ty skipped locally (the worktree has no
type-checking env) and left to CI.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01DW9D1WhLN8JhhiEPCScbaw

Author: hmgaudecker

CHANGES.md

@@ -13,10 +13,12 @@ chronological order. We follow [semantic versioning](https://semver.org/).
   transition row is evaluated directly at the from-value with that regime's `sigma`
   (no precomputed-row interpolation). This expresses regime-switching income risk /
   stochastic volatility. Supported for CDF-binned `NormalIIDProcess`
-  (`gauss_hermite=False`) and `TauchenAR1Process`; Gauss-Hermite IID and Rouwenhorst
-  are rejected at construction (their fixed-node kernels cannot carry a
-  state-conditioned `sigma`). Current-regime conditioning only. See
-  `lcm_examples/stochastic_volatility.py`.
+  (`gauss_hermite=False`) and `TauchenAR1Process`; Gauss-Hermite node placement and
+  Rouwenhorst are rejected at construction (their fixed-node kernels cannot carry a
+  state-conditioned `sigma`). Solving and simulating use the same conditioned law.
+  Every grid parameter must be fixed at construction, and the conditioning state must
+  map its categories to the same integer codes in every regime that carries it.
+  Current-regime conditioning only. See `lcm_examples/stochastic_volatility.py`.
 
 - Adds the DC-EGM solver (Iskhakov, Jørgensen, Rust & Schjerning 2017) as a
   per-regime alternative to grid search: `Regime(solver=lcm.DCEGM(...))`.

src/_lcm/processes/state_conditioned.py

@@ -48,13 +48,21 @@ def sigma_array_by_code(cond_grid: DiscreteGrid, by: Mapping[str, float]) -> Flo
 
     Stacking by `Mapping` insertion order silently permutes regimes when the code order
     differs (audit F5/RT5); indexing the returned array by the conditioning state's code
-    is therefore correct by construction. Every category must appear in `by`.
+    is therefore correct by construction. `by` must name exactly the categories of
+    `cond_grid` — an extra key is a typo or a stale category, never a no-op.
     """
     cats, codes = cond_grid.categories, cond_grid.codes
     missing = set(cats) - set(by)
     if missing:
         msg = f"StateConditioned.by is missing categories {sorted(missing)}"
         raise ValueError(msg)
+    extra = set(by) - set(cats)
+    if extra:
+        msg = (
+            f"StateConditioned.by has categories {sorted(extra)} that are not in the "
+            f"conditioning grid {sorted(cats)}"
+        )
+        raise ValueError(msg)
     # Runtime indexing (`gather_sigma`) uses the conditioning state's code directly,
     # which is safe because `@categorical` assigns contiguous 0..n-1 codes and
     # `DiscreteGrid` accepts only such classes — so position == code here.
@@ -68,20 +76,23 @@ def conditioned_row(
     nodes: Float1D,
     sigma: _Scalar,
     from_value: _Scalar,
+    mu: _Scalar,
     rho: _Scalar | None = None,
 ) -> Float1D:
     """Dispatch to the direct-CDF row builder for the given process `family`.
 
     `nodes` are the FIXED common nodes (from `grid_sigma`); `sigma` is the current
-    regime's innovation std (already gathered by code). `rho` is required for Tauchen.
+    regime's innovation std (already gathered by code). `mu` is the process's fixed
+    location (IID mean / AR(1) intercept) — dropping it misplaces the entire row
+    (code-review F2). `rho` is required for Tauchen.
     """
     if family == "iid_normal":
-        return iid_normal_row(nodes, mu=0.0, sigma=sigma)
+        return iid_normal_row(nodes, mu=mu, sigma=sigma)
     if family == "tauchen":
         if rho is None:
             msg = "conditioned_row(family='tauchen') requires rho"
             raise ValueError(msg)
-        return tauchen_row(nodes, rho=rho, sigma=sigma, from_value=from_value)
+        return tauchen_row(nodes, rho=rho, sigma=sigma, from_value=from_value, mu=mu)
     msg = f"unsupported family {family!r} (v1: 'iid_normal' | 'tauchen')"
     raise ValueError(msg)
 
@@ -118,12 +129,19 @@ def tauchen_row(
     rho: _Scalar,
     sigma: _Scalar,
     from_value: _Scalar,
+    mu: _Scalar = 0.0,
 ) -> Float1D:
-    """Conditional AR(1) row on FIXED ``nodes``, innovation ``~ N(0, sigma**2)``.
+    """Conditional AR(1) row for ``y' = mu + rho*y + eps``, ``eps ~ N(0, sigma**2)``.
 
     Evaluated DIRECTLY at ``from_value`` (audit F1) rather than by interpolating node
     rows. The denominator is the innovation ``sigma`` (conditional std of ``y' | y``),
     matching ``TauchenAR1Process.compute_transition_probs``.
+
+    ``nodes`` and ``from_value`` are in PHYSICAL units — i.e. the axis
+    ``TauchenAR1Process.compute_gridpoints`` returns, centred on ``mu/(1-rho)`` —
+    so the conditional mean is ``mu + rho*from_value``. Stock pylcm builds the same
+    row in demeaned coordinates, where the intercept vanishes; here it does not, and
+    omitting it misplaces every row unless ``mu == 0`` (code-review F2).
     """
     edges = (nodes[:-1] + nodes[1:]) / 2.0
-    return _row_from_edge_cdf(cdf((edges - rho * from_value) / sigma))
+    return _row_from_edge_cdf(cdf((edges - mu - rho * from_value) / sigma))

src/_lcm/regime_building/next_state.py

@@ -1,17 +1,23 @@
 """Generate function that compute the next states for solution and simulation."""
 
-from collections.abc import Callable
+from collections.abc import Callable, Mapping
 from types import MappingProxyType
+from typing import Any
 
 import jax
 from dags import concatenate_functions, with_signature
 from dags.tree import qname_from_tree_path
 
 from _lcm.engine import Variables
-from _lcm.grids import Grid
+from _lcm.grids import DiscreteGrid, Grid
 from _lcm.processes import _ContinuousStochasticProcess
 from _lcm.processes.ar1 import _AR1Process
 from _lcm.processes.iid import _IIDProcess
+from _lcm.processes.state_conditioned import (
+    StateConditioned,
+    gather_sigma,
+    sigma_array_by_code,
+)
 from _lcm.typing import (
     EconFunctionsMapping,
     NextStateSimulationFunction,
@@ -24,7 +30,8 @@
     TransitionFunctionName,
     TransitionFunctionsMapping,
 )
-from lcm.typing import ContinuousState, DiscreteState, FloatND, IntND
+from lcm.exceptions import ModelInitializationError
+from lcm.typing import ContinuousState, DiscreteState, Float1D, FloatND, IntND
 
 
 def get_next_state_function_for_solution(
@@ -273,17 +280,55 @@ def _create_continuous_stochastic_next_func(
     grid: _ContinuousStochasticProcess = all_grids[target_regime_name][state_name]  # ty: ignore [invalid-assignment]
     qname = qname_from_tree_path((target_regime_name, next_state_name))
 
+    # A state-conditioned process must DRAW with the current regime's sigma, not the
+    # scalar common-grid sigma — otherwise solve and simulate run different laws
+    # (code-review F1).
+    conditioned = _resolve_conditioned_sigma(
+        grid=grid, grids=all_grids[target_regime_name]
+    )
+
     if isinstance(grid, _AR1Process):
-        return _create_ar1_next_func(qname=qname, state_name=state_name, grid=grid)
+        return _create_ar1_next_func(
+            qname=qname, state_name=state_name, grid=grid, conditioned=conditioned
+        )
     if isinstance(grid, _IIDProcess):
-        return _create_iid_next_func(qname=qname, state_name=state_name, grid=grid)
+        return _create_iid_next_func(
+            qname=qname, state_name=state_name, grid=grid, conditioned=conditioned
+        )
 
     msg = f"Expected _IIDProcess or _AR1Process, got {type(grid)}"
     raise TypeError(msg)
 
 
+def _resolve_conditioned_sigma(
+    *,
+    grid: _ContinuousStochasticProcess,
+    grids: Mapping[StateOrActionName, Grid],
+) -> tuple[StateConditioned, Float1D] | None:
+    """Resolve the code-ordered sigma array of a state-conditioned process, else None.
+
+    Reuses `sigma_array_by_code` so simulation gathers sigma under exactly the
+    category-code ordering the solve-side transition rows use (code-review F1).
+    """
+    sc = grid.state_conditioned
+    if sc is None:
+        return None
+    conditioning_grid = grids.get(sc.on)
+    if not isinstance(conditioning_grid, DiscreteGrid):
+        msg = (
+            f"state_conditioned.on='{sc.on}' must name a DiscreteGrid state in the "
+            f"same regime as the process."
+        )
+        raise ModelInitializationError(msg)
+    return sc, sigma_array_by_code(conditioning_grid, sc.by)
+
+
 def _create_ar1_next_func(
-    *, qname: str, state_name: StateName, grid: _AR1Process
+    *,
+    qname: str,
+    state_name: StateName,
+    grid: _AR1Process,
+    conditioned: tuple[StateConditioned, Float1D] | None = None,
 ) -> StochasticNextFunction:
     fixed_params = dict(grid.params)
     runtime_param_names = {
@@ -294,6 +339,8 @@ def _create_ar1_next_func(
         state_name: "ContinuousState",
         **dict.fromkeys(runtime_param_names, "FloatND"),
     }
+    if conditioned is not None:
+        args[conditioned[0].on] = "DiscreteState"
     _draw_shock = grid.draw_shock
 
     @with_signature(args=args, return_annotation="ContinuousState")
@@ -302,6 +349,7 @@ def next_stochastic_state(**kwargs: FloatND) -> ContinuousState:
             {
                 **fixed_params,
                 **{raw: kwargs[qn] for qn, raw in runtime_param_names.items()},
+                **_conditioned_sigma(conditioned, kwargs),
             }
         )
         return _draw_shock(
@@ -314,7 +362,11 @@ def next_stochastic_state(**kwargs: FloatND) -> ContinuousState:
 
 
 def _create_iid_next_func(
-    *, qname: str, state_name: StateName, grid: _IIDProcess
+    *,
+    qname: str,
+    state_name: StateName,
+    grid: _IIDProcess,
+    conditioned: tuple[StateConditioned, Float1D] | None = None,
 ) -> StochasticNextFunction:
     fixed_params = dict(grid.params)
     runtime_param_names = {
@@ -324,6 +376,8 @@ def _create_iid_next_func(
         f"key_{qname}": "PRNGKeyND",
         **dict.fromkeys(runtime_param_names, "FloatND"),
     }
+    if conditioned is not None:
+        args[conditioned[0].on] = "DiscreteState"
     _draw_shock = grid.draw_shock
 
     @with_signature(args=args, return_annotation="ContinuousState")
@@ -332,6 +386,7 @@ def next_stochastic_state(**kwargs: FloatND) -> ContinuousState:
             {
                 **fixed_params,
                 **{raw: kwargs[qn] for qn, raw in runtime_param_names.items()},
+                **_conditioned_sigma(conditioned, kwargs),
             }
         )
         return _draw_shock(
@@ -340,3 +395,18 @@ def next_stochastic_state(**kwargs: FloatND) -> ContinuousState:
         )
 
     return next_stochastic_state
+
+
+def _conditioned_sigma(
+    conditioned: tuple[StateConditioned, Float1D] | None,
+    kwargs: Mapping[str, Any],
+) -> Mapping[str, Any]:
+    """`{"sigma": <current regime's sigma>}` for a conditioned process, else `{}`.
+
+    Overrides the scalar common-grid sigma the draw would otherwise use, which is what
+    makes the simulated law match the solved one (code-review F1).
+    """
+    if conditioned is None:
+        return {}
+    sc, sigma_by_code = conditioned
+    return {"sigma": gather_sigma(sigma_by_code, kwargs[sc.on])}