Skip unreachable target regimes in continuation value loop (#316)

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: hmgaudecker

src/lcm/regime_building/Q_and_F.py

@@ -66,15 +66,35 @@ def get_Q_and_F(
     joint_weights_from_marginals = {}
     next_V = {}
 
-    target_regime_names = tuple(transitions)
-    active_regimes_next_period = tuple(
-        target_regime_name
-        for target_regime_name in target_regime_names
-        if period + 1 in regimes_to_active_periods[target_regime_name]
+    # Enumerate all active targets, not just those in transitions — targets
+    # entirely absent from per-target dicts must also be detected.
+    all_active_next_period = tuple(
+        regime_name
+        for regime_name in regime_to_v_interpolation_info
+        if period + 1 in regimes_to_active_periods.get(regime_name, ())
     )
+
+    # Keep only targets whose stochastic state needs are all covered by
+    # `transitions`. Targets with missing stochastic transitions are dropped
+    # from the traced function; `validate_regime_transitions_all_periods`
+    # (via `_validate_no_reachable_incomplete_targets` in
+    # `lcm.utils.error_handling`) raises pre-solve if any such target has
+    # non-zero transition probability.
+    complete_targets: list[RegimeName] = []
+    for regime_name in all_active_next_period:
+        target_stochastic_needs = {
+            f"next_{s}"
+            for s in regime_to_v_interpolation_info[regime_name].state_names
+            if f"next_{s}" in stochastic_transition_names
+        }
+        if regime_name in transitions and target_stochastic_needs.issubset(
+            transitions[regime_name]
+        ):
+            complete_targets.append(regime_name)
+
     next_V_extra_param_names: dict[str, frozenset[str]] = {}
 
-    for target_regime_name in active_regimes_next_period:
+    for target_regime_name in complete_targets:
         # Transitions from the current regime to the target regime
         target_transitions = transitions[target_regime_name]
 
@@ -170,14 +190,16 @@ def Q_and_F(
             period=period,
             age=age,
         )
-        # Filter to active regimes only — inactive regimes must have 0
-        # probability (validated before solve).
+        # `complete_targets` is resolved at trace time (it is a closure over
+        # a Python list); incomplete-target validation happens outside JIT
+        # in `_validate_no_reachable_incomplete_targets` so that the traced
+        # graph contains no runtime error-raising callbacks.
         active_regime_probs = MappingProxyType(
-            {r: regime_transition_probs[r] for r in active_regimes_next_period}
+            {r: regime_transition_probs[r] for r in complete_targets}
         )
 
         E_next_V = jnp.zeros_like(U_arr)
-        for target_regime_name in active_regimes_next_period:
+        for target_regime_name in complete_targets:
             next_states = state_transitions[target_regime_name](
                 **states_actions_params,
                 period=period,

src/lcm/regime_building/processing.py

@@ -9,7 +9,7 @@
 import pandas as pd
 from dags import concatenate_functions, get_annotations, with_signature
 from dags.signature import rename_arguments
-from dags.tree import qname_from_tree_path, tree_path_from_qname
+from dags.tree import QNAME_DELIMITER, qname_from_tree_path, tree_path_from_qname
 from jax import Array
 from jax import numpy as jnp
 
@@ -85,7 +85,7 @@ def process_regimes(
         The processed regimes.
 
     """
-    states_per_regime: dict[str, set[str]] = {
+    states_per_regime: dict[RegimeName, set[str]] = {
         name: set(regime.states.keys()) for name, regime in regimes.items()
     }
 
@@ -174,7 +174,7 @@ def process_regimes(
 def _build_solve_functions(
     *,
     regime: Regime,
-    regime_name: str,
+    regime_name: RegimeName,
     nested_transitions: dict[str, dict[str, UserFunction] | UserFunction],
     all_grids: MappingProxyType[RegimeName, MappingProxyType[str, Grid]],
     regime_params_template: RegimeParamsTemplate,
@@ -262,7 +262,7 @@ def _build_solve_functions(
 def _build_simulate_functions(
     *,
     regime: Regime,
-    regime_name: str,
+    regime_name: RegimeName,
     nested_transitions: dict[str, dict[str, UserFunction] | UserFunction],
     all_grids: MappingProxyType[RegimeName, MappingProxyType[str, Grid]],
     regime_params_template: RegimeParamsTemplate,
@@ -508,13 +508,21 @@ def _process_regime_core(
         )
 
     # Shock transitions bypass the stub pipeline entirely. Build weight and
-    # next functions for ALL target regimes directly from each target's grid.
+    # next functions for reachable target regimes from each target's grid.
+    # Scope to targets already present in non-shock transitions to avoid
+    # spurious entries for unreachable regimes.
     shock_names = variable_info.query("is_shock").index.tolist()
-    target_shock_grids: dict[tuple[str, str], _ShockGrid] = {  # ty: ignore[invalid-assignment]
-        (regime, shock): grids[shock]
+    reachable_targets = {
+        tree_path_from_qname(k)[0]
+        for k in flat_nested_transitions
+        if QNAME_DELIMITER in k
+    }
+    target_shock_grids: dict[tuple[RegimeName, str], _ShockGrid] = {
+        (regime, shock): grid
         for regime, grids in all_grids.items()
+        if regime in reachable_targets
         for shock in shock_names
-        if isinstance(grids.get(shock), _ShockGrid)
+        if isinstance(grid := grids.get(shock), _ShockGrid)
     }
     functions |= {
         f"weight_{regime}__next_{shock}": _get_weights_func_for_shock(
@@ -567,7 +575,7 @@ def _process_regime_core(
 def _extract_transitions_from_regime(
     *,
     regime: Regime,
-    states_per_regime: Mapping[str, set[str]],
+    states_per_regime: Mapping[RegimeName, set[str]],
 ) -> dict[str, dict[str, UserFunction] | UserFunction]:
     """Extract transitions from `regime.state_transitions` and regime transition.
 
@@ -600,7 +608,14 @@ def _extract_transitions_from_regime(
         {"next_regime": regime.transition},
     )
 
-    for target_regime_name, target_regime_state_names in states_per_regime.items():
+    reachable_targets = _get_reachable_targets(
+        per_target_transitions=per_target_transitions,
+        simple_transitions=simple_transitions,
+        states_per_regime=states_per_regime,
+    )
+
+    for target_regime_name in reachable_targets:
+        target_regime_state_names = states_per_regime[target_regime_name]
         target_dict: dict[str, UserFunction] = {}
         for state_name in target_regime_state_names:
             next_key = f"next_{state_name}"
@@ -616,6 +631,34 @@ def _extract_transitions_from_regime(
     return nested
 
 
+def _get_reachable_targets(
+    *,
+    per_target_transitions: dict[str, dict[str, UserFunction]],
+    simple_transitions: dict[str, UserFunction],
+    states_per_regime: Mapping[RegimeName, set[str]],
+) -> set[RegimeName]:
+    """Determine which target regimes need transition entries.
+
+    When per-target transitions exist, start from the explicitly named targets
+    and add any target whose state needs are fully covered by simple
+    (non-per-target) transitions. Without per-target transitions, all regimes
+    are reachable.
+
+    """
+    if not per_target_transitions:
+        return set(states_per_regime.keys())
+
+    targets: set[RegimeName] = set()
+    for variants in per_target_transitions.values():
+        targets |= variants.keys()
+    for target_name, target_states in states_per_regime.items():
+        if target_name not in targets:
+            needed = {f"next_{s}" for s in target_states}
+            if needed and needed.issubset(simple_transitions):
+                targets.add(target_name)
+    return targets
+
+
 def _classify_transitions(
     state_transitions: dict[str, UserFunction],
 ) -> tuple[dict[str, UserFunction], dict[str, dict[str, UserFunction]]]:

src/lcm/utils/error_handling.py

@@ -61,15 +61,17 @@ def validate_V(
             "reasons:\n"
             "- The user-defined functions returned invalid values.\n"
             "- It is impossible to reach an active regime, resulting in NaN regime\n"
-            "  transition probabilities."
+            "  transition probabilities.\n"
+            "- A per-target state_transitions dict omits a reachable target\n"
+            "  (non-zero transition probability to an incomplete target)."
         )
 
 
 def validate_regime_transition_probs(
     *,
     regime_transition_probs: MappingProxyType[str, Array],
-    active_regimes_next_period: tuple[str, ...],
-    regime_name: str,
+    active_regimes_next_period: tuple[RegimeName, ...],
+    regime_name: RegimeName,
     age: ScalarInt | ScalarFloat,
     next_age: ScalarInt | ScalarFloat,
     state_action_values: MappingProxyType[str, Array] | None = None,
@@ -224,7 +226,7 @@ def validate_regime_transitions_all_periods(
                 continue
 
             _validate_regime_transition_single(
-                internal_regime=internal_regime,
+                internal_regimes=internal_regimes,
                 regime_params=internal_params[name],
                 active_regimes_next_period=active_regimes_next_period,
                 regime_name=name,
@@ -235,10 +237,10 @@ def validate_regime_transitions_all_periods(
 
 def _validate_regime_transition_single(
     *,
-    internal_regime: InternalRegime,
+    internal_regimes: MappingProxyType[RegimeName, InternalRegime],
     regime_params: FlatRegimeParams,
-    active_regimes_next_period: tuple[str, ...],
-    regime_name: str,
+    active_regimes_next_period: tuple[RegimeName, ...],
+    regime_name: RegimeName,
     period: int,
     ages: AgeGrid,
 ) -> None:
@@ -248,6 +250,7 @@ def _validate_regime_transition_single(
     variables it accepts, using `jax.vmap` for vectorised evaluation.
 
     """
+    internal_regime = internal_regimes[regime_name]
     # Non-None guaranteed: only called for non-terminal regimes
     regime_transition_func = (
         internal_regime.solve_functions.compute_regime_transition_probs
@@ -310,6 +313,64 @@ def _call(
         state_action_values=MappingProxyType(point),
     )
 
+    _validate_no_reachable_incomplete_targets(
+        internal_regimes=internal_regimes,
+        regime_transition_probs=regime_transition_probs,
+        active_regimes_next_period=active_regimes_next_period,
+        regime_name=regime_name,
+        age=ages.values[period],  # noqa: PD011
+    )
+
+
+def _validate_no_reachable_incomplete_targets(
+    *,
+    internal_regimes: MappingProxyType[RegimeName, InternalRegime],
+    regime_transition_probs: MappingProxyType[str, Array],
+    active_regimes_next_period: tuple[RegimeName, ...],
+    regime_name: RegimeName,
+    age: ScalarInt | ScalarFloat,
+) -> None:
+    """Check that targets with incomplete stochastic transitions are unreachable.
+
+    A target is "incomplete" from the source regime if the source's
+    `transitions[target_regime_name]` does not cover all of the target's
+    stochastic state needs. Such targets must have zero transition
+    probability, otherwise the continuation value cannot be computed. This
+    includes self-transitions (regime reaches itself): omitting the
+    self-entry in a per-target dict is a common user error.
+
+    """
+    solve_functions = internal_regimes[regime_name].solve_functions
+    transitions = solve_functions.transitions
+    stochastic_names = solve_functions.stochastic_transition_names
+
+    for target_regime_name in active_regimes_next_period:
+        target_regime = internal_regimes[target_regime_name]
+        target_state_names = tuple(target_regime.variable_info.query("is_state").index)
+        needs = {
+            f"next_{s}" for s in target_state_names if f"next_{s}" in stochastic_names
+        }
+        if not needs:
+            continue
+        if target_regime_name in transitions and needs.issubset(
+            transitions[target_regime_name]
+        ):
+            continue
+        if not jnp.any(regime_transition_probs[target_regime_name] > 0):
+            continue
+        missing = sorted(needs - set(transitions.get(target_regime_name, {})))
+        if target_regime_name not in transitions:
+            missing = sorted(f"next_{s}" for s in target_state_names)
+        raise InvalidRegimeTransitionProbabilitiesError(
+            f"Regime '{regime_name}' at age {age} has positive transition "
+            f"probability to '{target_regime_name}', but '{regime_name}' "
+            f"does not provide state transition(s) for: {missing}. Extend "
+            f"`state_transitions` in '{regime_name}' to cover "
+            f"'{target_regime_name}' (via a per-target dict if the "
+            f"transition differs by target), or ensure "
+            f"'{target_regime_name}' is unreachable."
+        )
+
 
 def _get_func_indexing_params(
     *,