Fix outcome axis for cross-grid stochastic transitions (#321)

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

Author: hmgaudecker

src/lcm/pandas_utils.py

@@ -379,6 +379,7 @@ def array_from_series(
             outcome_mapping = _build_outcome_mapping(
                 func_name=func_name,
                 grids=grids,
+                regimes=regimes,
                 regime_names_to_ids=regime_names_to_ids,
             )
             level_mappings = (*level_mappings, outcome_mapping)
@@ -607,16 +608,20 @@ def _build_outcome_mapping(
     *,
     func_name: str,
     grids: dict[str, DiscreteGrid],
+    regimes: Mapping[str, Regime],
     regime_names_to_ids: RegimeNamesToIds,
 ) -> _LevelMapping:
     """Build a `_LevelMapping` for the outcome axis of a `next_*` function.
 
     For state transitions (e.g. `"next_partner"`), look up the state grid.
+    For per-target transitions (e.g. `"next_health__post65"`), use the target
+    regime's grid for the outcome axis.
     For regime transitions (`"next_regime"`), use `regime_names_to_ids`.
 
     Args:
         func_name: Function name starting with `"next_"`.
         grids: Categorical grid lookup.
+        regimes: Mapping of regime names to user Regime instances.
         regime_names_to_ids: Immutable mapping from regime names to integer
             indices.
 
@@ -634,6 +639,17 @@ def _build_outcome_mapping(
 
     path = tree_path_from_qname(func_name)
     state_name = path[0].removeprefix("next_")
+
+    # Per-target transitions (e.g. "next_health__post65") must use the TARGET
+    # regime's grid for the outcome axis, not the source regime's grid.
+    if len(path) > 1:
+        target_regime_name = path[1]
+        target_regime = regimes.get(target_regime_name)
+        if target_regime is not None and state_name in target_regime.states:
+            target_grid = target_regime.states[state_name]
+            if isinstance(target_grid, DiscreteGrid):
+                return _grid_level_mapping(name=f"next_{state_name}", grid=target_grid)
+
     return _grid_level_mapping(name=f"next_{state_name}", grid=grids[state_name])
 
 

tests/test_pandas_utils.py

@@ -1614,6 +1614,7 @@ def test_build_outcome_mapping_qualified_func_name() -> None:
     result = _build_outcome_mapping(
         func_name="next_health__working",
         grids=grids,
+        regimes=model.regimes,
         regime_names_to_ids=model.regime_names_to_ids,
     )
     assert result.size == 2
@@ -1787,6 +1788,107 @@ class WrongPartner:
         )
 
 
+def test_convert_series_cross_grid_transition() -> None:
+    """Outcome axis must use the TARGET regime's grid, not the source's.
+
+    When a per-target MarkovTransition crosses grid sizes (e.g. 3-state
+    source → 2-state target), the converted array's last dimension must
+    match the target's grid size (2), not the source's (3).
+    """
+    from lcm import MarkovTransition  # noqa: PLC0415
+    from lcm.typing import DiscreteState, FloatND, Period  # noqa: PLC0415
+
+    @categorical(ordered=True)
+    class _HealthPre:
+        disabled: int
+        bad: int
+        good: int
+
+    @categorical(ordered=True)
+    class _HealthPost:
+        bad: int
+        good: int
+
+    @categorical(ordered=False)
+    class _RId:
+        pre65: int
+        post65: int
+
+    def _health_probs_same(
+        period: Period, health: DiscreteState, health_trans_probs: FloatND
+    ) -> FloatND:
+        return health_trans_probs[period, health]
+
+    def _health_probs_cross(
+        period: Period, health: DiscreteState, health_trans_probs_cross: FloatND
+    ) -> FloatND:
+        return health_trans_probs_cross[period, health]
+
+    pre65 = Regime(
+        states={
+            "health": DiscreteGrid(_HealthPre),
+            "wealth": LinSpacedGrid(start=0, stop=10, n_points=5),
+        },
+        state_transitions={
+            "health": {
+                "pre65": MarkovTransition(_health_probs_same),
+                "post65": MarkovTransition(_health_probs_cross),
+            },
+            "wealth": lambda wealth: wealth,
+        },
+        functions={"utility": lambda health, wealth: wealth + health},
+        transition=lambda age: jnp.where(age >= 1, _RId.post65, _RId.pre65),
+        active=lambda age: age < 1,
+    )
+    post65 = Regime(
+        transition=None,
+        states={
+            "health": DiscreteGrid(_HealthPost),
+            "wealth": LinSpacedGrid(start=0, stop=10, n_points=5),
+        },
+        functions={"utility": lambda health, wealth: wealth + health},
+    )
+    model = Model(
+        regimes={"pre65": pre65, "post65": post65},
+        ages=AgeGrid(start=0, stop=1, step="Y"),
+        regime_id_class=_RId,
+    )
+
+    # Cross-grid transition probs: 3 source states → 2 target states
+    index_cross = pd.MultiIndex.from_tuples(
+        [
+            (0.0, "disabled", "bad"),
+            (0.0, "disabled", "good"),
+            (0.0, "bad", "bad"),
+            (0.0, "bad", "good"),
+            (0.0, "good", "bad"),
+            (0.0, "good", "good"),
+        ],
+        names=["age", "health", "next_health"],
+    )
+    sr_cross = pd.Series([0.65, 0.35, 0.81, 0.19, 0.06, 0.94], index=index_cross)
+
+    params = {
+        "pre65": {
+            "to_post65_next_health": {"health_trans_probs_cross": sr_cross},
+        },
+    }
+    internal = broadcast_to_template(
+        params=params, template=model.get_params_template(), required=False
+    )
+    result = convert_series_in_params(
+        internal_params=internal,
+        regimes=model.regimes,
+        ages=model.ages,
+        regime_names_to_ids=model.regime_names_to_ids,
+    )
+
+    arr = result["pre65"]["to_post65_next_health__health_trans_probs_cross"]
+    # Shape: (n_ages=2, n_source_health=3, n_target_health=2)
+    # n_ages=2 because AgeGrid has ages [0, 1]; missing age 1 is NaN-filled.
+    assert arr.shape == (2, 3, 2)  # ty: ignore[unresolved-attribute]
+
+
 def test_resolve_categoricals_includes_derived_when_no_regime_name() -> None:
     """derived_categoricals are included even when regime_name is None."""
     from lcm.pandas_utils import _resolve_categoricals  # noqa: PLC0415