Repair state-conditioned shocks per PR #405 rounds 2-3 review

Round-2 findings:
- F1: broadcast pruner and usage validator missed the metadata-declared
  `state_conditioned.on` dependency, rejecting legitimate models. Both passes
  now register the conditioner explicitly (pruner keeps it only when a retained
  or reachable-target process reads it; validator credits it model-wide).
- F2: `_validate_conditioned_grid_is_fixed` now rejects singleton and
  non-strictly-increasing node axes.
- F4: added AR(1) simulate-side regression (innovation std tracks conditioned
  sigma) and a through-time timing test.

Round-3 finding:
- F1 (cross-regime): a conditioned process in a target regime creates a
  dependency in every source regime that can reach it (per-target transitions),
  evaluated at the source's own `on` state. Pruner keeps the conditioner for
  reachable targets; validator credits it via the model-wide set. Added three
  cross-regime regression tests (build + solve + semantic).

F3 accepted as a documented v1 limitation (cross-regime code agreement).

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

Author: hmgaudecker

src/_lcm/model_processing.py

@@ -26,6 +26,7 @@
     materialize_granular_transition_params,
 )
 from _lcm.params.sequence_leaf import SequenceLeaf
+from _lcm.processes import _ContinuousStochasticProcess
 from _lcm.regime_building.finalize import FinalizedUserRegime
 from _lcm.regime_building.h_dag import get_dag_targets_consumed_by_H
 from _lcm.regime_building.max_Q_over_a import TASTE_SHOCK_SCALE_PARAM
@@ -257,6 +258,30 @@ def _fail_if_invalid_n_subjects(*, n_subjects: int | None) -> None:
         raise ValueError(msg)
 
 
+def _model_wide_conditioning_names(
+    user_regimes: Mapping[RegimeName, UserRegime],
+) -> set[str]:
+    """Every conditioning state read by any state-conditioned process in the model.
+
+    `state_conditioned.on` is a real dependency of the generated weights/draw functions,
+    but it lives in grid metadata rather than in a user function, so a callable-DAG scan
+    misses it. A conditioned process's transition weight is also built into the Q of
+    every *source* regime that can reach the process's regime, evaluated at that
+    source's current `on` state — so the dependency is not local to the process's regime
+    (round-3 review F1). We collect the conditioners model-wide and, at the call site,
+    credit each regime for those it actually carries: a conservative over-approximation
+    of reachability whose only cost is not flagging a genuinely unused state, never a
+    wrong policy.
+    """
+    return {
+        grid.state_conditioned.on
+        for user_regime in user_regimes.values()
+        for grid in user_regime.states.values()
+        if isinstance(grid, _ContinuousStochasticProcess)
+        and grid.state_conditioned is not None
+    }
+
+
 def _validate_all_variables_used(
     user_regimes: Mapping[RegimeName, UserRegime],
     *,
@@ -285,6 +310,7 @@ def _validate_all_variables_used(
 
     """
     error_messages = []
+    conditioning_names = _model_wide_conditioning_names(user_regimes)
 
     for regime_name, user_regime in user_regimes.items():
         variable_names = set(user_regime.states) | set(user_regime.actions)
@@ -306,6 +332,14 @@ def _validate_all_variables_used(
         reachable = get_ancestors(
             user_functions, targets=targets, include_targets=False
         )
+        # A state-conditioned process reads `state_conditioned.on`: the generated solve
+        # weights and the simulation draw both take it as an argument. But it is
+        # declared as grid *metadata* rather than in a user function, so the ancestry
+        # above cannot see it. A process conditioned in one regime is also drawn into a
+        # *source* regime's Q when that source can reach it, evaluated at the source's
+        # own `on` state — so a conditioner is credited to every regime that carries it,
+        # not only the process's own regime (round-2 review F1 / round-3 review F1).
+        reachable = set(reachable) | (conditioning_names & variable_names)
         unused_variables = sorted(variable_names - reachable)
 
         if unused_variables:

src/_lcm/regime_building/broadcast.py

@@ -21,6 +21,7 @@
 from dags import get_ancestors
 
 from _lcm.grids import Grid
+from _lcm.processes import _ContinuousStochasticProcess
 from _lcm.regime_building.phases import (
     PhasedRegimeSpec,
     RegimePhaseSpec,
@@ -264,6 +265,13 @@ def _phase_fixed_point(
                 reachable_targets=reachable[regime_name],
                 kept=grown,
             )
+            needed |= _state_conditioned_names(
+                user_regime=user_regime,
+                reachable_targets=reachable[regime_name],
+                user_regimes=user_regimes,
+                candidates=broadcast_variables[regime_name],
+                grown_here=grown[regime_name],
+            )
             newly_kept = candidates & needed
             if newly_kept:
                 grown[regime_name] = grown[regime_name] | newly_kept
@@ -272,6 +280,57 @@ def _phase_fixed_point(
             return grown
 
 
+def _state_conditioned_names(
+    *,
+    user_regime: UserRegime,
+    reachable_targets: frozenset[RegimeName],
+    user_regimes: Mapping[RegimeName, UserRegime],
+    candidates: frozenset[StateOrActionName],
+    grown_here: frozenset[StateOrActionName],
+) -> set[str]:
+    """Collect the conditioning states this regime's Q reads through a process.
+
+    A state-conditioned process reads ``state_conditioned.on``: the generated solve
+    weights and simulation draw both take it as an argument. But it is declared as
+    *metadata* on the grid rather than in a user function, so `_needed_names`'s
+    callable-DAG ancestry cannot see it, and a conditioning state declared at model
+    level (a broadcast candidate, reaching nothing else) was pruned here before the
+    feature could ask for it — leaving a misleading "must name a DiscreteGrid state in
+    the same regime" error at build time (code-review round 2, F1).
+
+    Two sources contribute:
+
+    - **local processes** — declared in this regime. Only a *retained* process forces
+      its conditioner (a process that is itself a pruned broadcast candidate reads
+      nothing), so this stays monotone in ``grown_here`` and the caller's fixed point
+      terminates.
+    - **reachable-target processes** — a conditioned process in a regime this one can
+      transition into has its transition weight built into *this* regime's Q, evaluated
+      at *this* regime's ``on`` state. So the conditioner is needed in the source too,
+      not only the process's own regime (round-3 review F1).
+
+    Keeping the state also keeps its law of motion, which the caller filters by the
+    same pruned-set.
+    """
+    names: set[str] = set()
+    for name, grid in user_regime.states.items():
+        if not isinstance(grid, _ContinuousStochasticProcess):
+            continue
+        if grid.state_conditioned is None:
+            continue
+        if name in candidates and name not in grown_here:
+            continue  # the process itself is not (yet) retained
+        names.add(grid.state_conditioned.on)
+    for target in reachable_targets:
+        for grid in user_regimes[target].states.values():
+            if (
+                isinstance(grid, _ContinuousStochasticProcess)
+                and grid.state_conditioned is not None
+            ):
+                names.add(grid.state_conditioned.on)
+    return names
+
+
 def _needed_names(
     *,
     phase_slice: RegimePhaseSpec,

src/_lcm/regime_building/processing.py

@@ -1680,6 +1680,13 @@ def _validate_conditioning_codes_agree_across_regimes(
     do; a regime that relabels `{low: 0, high: 1}` to `{high: 0, low: 1}` would silently
     swap the two volatilities (code-review F4). Rather than thread the source grid
     through every seam, v1 requires one shared map and rejects otherwise.
+
+    Round-2 review F3 noted this global scan is not edge-precise (it fails a
+    state-less reachable source late rather than at the guard, and over-rejects
+    disconnected components reusing the name), but established **no silent policy
+    error**. Kept deliberately for v1: the conservative rule is safe, and the
+    per-edge alternative would need source-reachability plumbing that could itself
+    introduce the sigma-swap this prevents. Accepted limitation, not an oversight.
     """
     maps = {
         regime: dict(zip(grid.categories, grid.codes, strict=True))
@@ -1712,6 +1719,10 @@ def _validate_conditioned_sigmas(by: Mapping[str, float]) -> None:
         raise ModelInitializationError(msg)
 
 
+#: A CDF row bins on node midpoints, so it needs at least one interior edge.
+_MIN_CONDITIONED_NODES = 2
+
+
 def _validate_conditioned_grid_is_fixed(
     *, name: str, grid: _ContinuousStochasticProcess
 ) -> None:
@@ -1735,6 +1746,30 @@ def _validate_conditioned_grid_is_fixed(
             f"state-conditioned process '{name}' resolved to non-finite nodes: {nodes}"
         )
         raise ModelInitializationError(msg)
+    # Finiteness is not enough. The direct-CDF row bins on the MIDPOINTS of these nodes,
+    # so a strictly increasing axis of >= 2 points is what makes the CDF differences
+    # probabilities at all (code-review round 2, F2). All of this is reachable from the
+    # public API: `sigma=-0.3` or `n_std=-3.0` yields a DESCENDING axis (pylcm does not
+    # require a positive sigma), `sigma=0.0` collapses every node, and a large `mu` in
+    # float32 rounds them together. The descending case is the dangerous one — the row
+    # still sums to 1.0, so a row-sum check passes while masses go negative and a
+    # continuation decision flips sign.
+    if nodes.ndim != 1 or nodes.shape[0] < _MIN_CONDITIONED_NODES:
+        msg = (
+            f"state-conditioned process '{name}' needs a 1-D axis of at least "
+            f"{_MIN_CONDITIONED_NODES} nodes; got shape {nodes.shape}. A single node "
+            f"leaves the CDF row no bin edges."
+        )
+        raise ModelInitializationError(msg)
+    if not bool(jnp.all(jnp.diff(nodes) > 0)):
+        msg = (
+            f"state-conditioned process '{name}' resolved to a node axis that is not "
+            f"strictly increasing: {nodes}. Midpoint-CDF binning would return negative "
+            f"transition masses (which still sum to 1.0, so this fails silently). "
+            f"Check for a negative or zero sigma / n_std, or nodes collapsed by "
+            f"float32 precision at a large mu."
+        )
+        raise ModelInitializationError(msg)
 
 
 def _get_conditioned_weights_func(