Support heterogeneous state sets in initial conditions (#315)

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

Author: hmgaudecker

src/lcm/ages.py

@@ -20,6 +20,10 @@
     }
 )
 
+# Names that behave like states in initial conditions but are not declared on
+# any `Regime.states`. `age` is required for every subject regardless of regime.
+PSEUDO_STATE_NAMES: frozenset[str] = frozenset({"age"})
+
 
 class AgeGrid:
     """Age grid for life-cycle models.

src/lcm/pandas_utils.py

@@ -11,12 +11,13 @@
 from dags.tree import qname_from_tree_path, tree_path_from_qname
 from jax import Array
 
-from lcm.ages import AgeGrid
+from lcm.ages import PSEUDO_STATE_NAMES, AgeGrid
 from lcm.grids import DiscreteGrid, IrregSpacedGrid
 from lcm.params import MappingLeaf
 from lcm.params.sequence_leaf import SequenceLeaf
 from lcm.regime import Regime
 from lcm.shocks import _ShockGrid
+from lcm.simulation.initial_conditions import MISSING_CAT_CODE
 from lcm.typing import InternalParams, RegimeNamesToIds
 from lcm.utils.error_handling import (
     _get_func_indexing_params,
@@ -39,7 +40,7 @@ def has_series(params: Mapping) -> bool:
     return False
 
 
-def initial_conditions_from_dataframe(
+def initial_conditions_from_dataframe(  # noqa: C901
     *,
     df: pd.DataFrame,
     regimes: Mapping[str, Regime],
@@ -92,9 +93,9 @@ def initial_conditions_from_dataframe(
     n_subjects = len(df)
     state_cols = [col for col in df.columns if col != "regime"]
 
-    # Pre-allocate result arrays
+    # Pre-allocate result arrays (NaN default surfaces bugs for missing states)
     result_arrays: dict[str, np.ndarray] = {
-        col: np.empty(n_subjects, dtype=float) for col in state_cols
+        col: np.full(n_subjects, np.nan) for col in state_cols
     }
     discrete_state_names: set[str] = set()
 
@@ -109,7 +110,12 @@ def initial_conditions_from_dataframe(
         }
         discrete_state_names |= discrete_grids.keys()
 
+        regime_state_names = set(regime.states.keys()) | PSEUDO_STATE_NAMES
+
         for col in state_cols:
+            if col not in regime_state_names:
+                continue
+
             values = group[col]
             if hasattr(values, "cat"):
                 values = values.astype(str)
@@ -126,6 +132,14 @@ def initial_conditions_from_dataframe(
             else:
                 result_arrays[col][idx] = values.to_numpy(dtype=float)
 
+    # Replace remaining NaN in discrete columns with an explicit int sentinel
+    # before casting to int32. This avoids platform-undefined NaN→int behavior
+    # and the associated RuntimeWarning.
+    for col in discrete_state_names:
+        if col in result_arrays:
+            nan_mask = np.isnan(result_arrays[col])
+            result_arrays[col][nan_mask] = MISSING_CAT_CODE
+
     initial_conditions: dict[str, Array] = {
         col: jnp.array(arr, dtype=jnp.int32)
         if col in discrete_state_names
@@ -786,17 +800,35 @@ def _validate_state_columns(
     unknown = state_columns - expected
     if unknown:
         msg = (
-            f"Unknown columns not matching any model state: {sorted(unknown)}. "
+            f"Unknown columns not matching any state of an initial regime: "
+            f"{sorted(unknown)}. "
             f"Expected states: {sorted(expected)}."
         )
         raise ValueError(msg)
 
     missing = expected - state_columns
     if missing:
-        msg = f"Missing required state columns: {sorted(missing)}."
+        required_by: dict[str, list[str]] = {name: [] for name in missing}
+        for regime_name in set(initial_regimes):
+            for name in regimes[regime_name].states:
+                if name in required_by:
+                    required_by[name].append(regime_name)
+        details = ", ".join(
+            _format_missing_state_detail(name=name, required_by=required_by[name])
+            for name in sorted(missing)
+        )
+        msg = f"Missing required state columns: {details}."
         raise ValueError(msg)
 
 
+def _format_missing_state_detail(*, name: str, required_by: list[str]) -> str:
+    if name in PSEUDO_STATE_NAMES:
+        return f"'{name}' (required for every subject)"
+    if required_by:
+        return f"'{name}' (required by {sorted(required_by)})"
+    return f"'{name}' (required by an initial regime)"
+
+
 def _collect_state_names(
     *,
     regimes: Mapping[str, Regime],
@@ -805,11 +837,11 @@ def _collect_state_names(
     """Collect all state names (including shock grids) from initial regimes.
 
     Returns:
-        Set of all state names from the initial regimes, plus `'age'`
-        (always required).
+        Set of all state names from the initial regimes, plus the pseudo-state
+        names from `PSEUDO_STATE_NAMES` (always required).
 
     """
-    names: set[str] = {"age"}
+    names: set[str] = set(PSEUDO_STATE_NAMES)
     for regime_name in set(initial_regimes):
         names.update(regimes[regime_name].states.keys())
     return names

src/lcm/simulation/initial_conditions.py

@@ -7,14 +7,15 @@
 
 from collections.abc import Callable, Mapping, Sequence
 from types import MappingProxyType
+from typing import Never
 
 import jax
 import numpy as np
 import pandas as pd
 from jax import Array
 from jax import numpy as jnp
 
-from lcm.ages import AgeGrid
+from lcm.ages import PSEUDO_STATE_NAMES, AgeGrid
 from lcm.exceptions import (
     InvalidInitialConditionsError,
     format_messages,
@@ -143,7 +144,10 @@ def validate_initial_conditions(
 
     # Validate discrete state values
     _validate_discrete_state_values(
-        initial_states=initial_states, internal_regimes=internal_regimes
+        initial_states=initial_states,
+        internal_regimes=internal_regimes,
+        regime_id_arr=regime_arr,
+        regime_names_to_ids=regime_names_to_ids,
     )
 
     # Validate feasibility
@@ -194,7 +198,7 @@ def _format_missing_states_message(missing: set[str], required: set[str]) -> str
             "knows each subject's starting age. Example: "
             "initial_states={'age': jnp.array([25.0, 25.0]), ...}"
         )
-    missing_model_states = sorted(missing - {"age"})
+    missing_model_states = sorted(missing - PSEUDO_STATE_NAMES)
     if missing_model_states:
         parts.append(f"Missing model states: {missing_model_states}.")
     parts.append(f"Required initial states are: {sorted(required)}")
@@ -230,12 +234,12 @@ def _collect_state_name_errors(
     errors: list[str] = []
 
     # All known states (union across all regimes) — used for the "extra" check
-    all_known_states: set[str] = {"age"}
+    all_known_states: set[str] = set(PSEUDO_STATE_NAMES)
     for internal_regime in internal_regimes.values():
         all_known_states.update(_get_regime_state_names(internal_regime))
 
     # Required states — only from regimes subjects actually start in
-    required_states: set[str] = {"age"}
+    required_states: set[str] = set(PSEUDO_STATE_NAMES)
     used_ids = jnp.unique(regime_id_arr)
     used_regime_names = {
         ids_to_regime_names[int(i)] for i in used_ids if int(i) in ids_to_regime_names
@@ -414,35 +418,53 @@ def _validate_discrete_state_values(
     *,
     initial_states: Mapping[str, Array],
     internal_regimes: MappingProxyType[RegimeName, InternalRegime],
+    regime_id_arr: Array,
+    regime_names_to_ids: Mapping[str, int],
 ) -> None:
     """Validate that discrete state values are valid codes.
 
+    Only check subjects in regimes that actually have the state.
+
     Args:
         initial_states: Mapping of state names to arrays.
         internal_regimes: Immutable mapping of regime names to internal regime
             instances.
+        regime_id_arr: Array of regime IDs for each subject.
+        regime_names_to_ids: Mapping from regime names to integer IDs.
 
     Raises:
         InvalidInitialConditionsError: If any discrete state contains invalid codes.
 
     """
-    discrete_valid_codes: dict[str, set[int]] = {}
-    for internal_regime in internal_regimes.values():
+    # Build per-state: valid codes + regime IDs that have this state
+    discrete_info: dict[str, tuple[set[int], set[int]]] = {}
+    for regime_name, internal_regime in internal_regimes.items():
+        regime_id = regime_names_to_ids[regime_name]
         for state_name in internal_regime.variable_info.query(
             "is_state and is_discrete"
         ).index:
             grid = internal_regime.grids[state_name]
             if isinstance(grid, DiscreteGrid):
-                existing = discrete_valid_codes.get(state_name, set())
-                discrete_valid_codes[state_name] = existing | set(grid.codes)
+                codes, regime_ids = discrete_info.get(state_name, (set(), set()))
+                discrete_info[state_name] = (
+                    codes | set(grid.codes),
+                    regime_ids | {regime_id},
+                )
 
-    for state_name, valid_codes in discrete_valid_codes.items():
+    for state_name, (valid_codes, regime_ids) in discrete_info.items():
         if state_name not in initial_states:
             continue
         values = initial_states[state_name]
-        invalid_mask = jnp.isin(values, jnp.array(sorted(valid_codes)), invert=True)
+        # Only validate subjects in regimes that have this state
+        in_relevant_regime = jnp.isin(regime_id_arr, jnp.array(sorted(regime_ids)))
+        relevant_values = values[in_relevant_regime]
+        if relevant_values.size == 0:
+            continue
+        invalid_mask = jnp.isin(
+            relevant_values, jnp.array(sorted(valid_codes)), invert=True
+        )
         if jnp.any(invalid_mask):
-            invalid_vals = sorted({int(v) for v in values[invalid_mask]})
+            invalid_vals = sorted({int(v) for v in relevant_values[invalid_mask]})
             raise InvalidInitialConditionsError(
                 f"Invalid values {invalid_vals} for discrete state "
                 f"'{state_name}'. Valid codes are: {sorted(valid_codes)}"
@@ -523,7 +545,7 @@ def _is_any_action_feasible(per_subject_kwargs: dict[str, Array]) -> Array:
     return jnp.concatenate(results)
 
 
-def _check_regime_feasibility(
+def _check_regime_feasibility(  # noqa: C901
     *,
     internal_regime: InternalRegime,
     regime_name: str,
@@ -587,13 +609,21 @@ def _check_regime_feasibility(
     }
 
     if subject_states:
-        any_feasible = _batched_feasibility_check(
-            feasibility_func=feasibility_func,
-            subject_states=subject_states,
-            action_kwargs=action_kwargs,
-            filtered_params=filtered_params,
-            flat_actions=flat_actions,
-        )
+        try:
+            any_feasible = _batched_feasibility_check(
+                feasibility_func=feasibility_func,
+                subject_states=subject_states,
+                action_kwargs=action_kwargs,
+                filtered_params=filtered_params,
+                flat_actions=flat_actions,
+            )
+        except TypeError as exc:
+            _raise_feasibility_type_error(
+                exc=exc,
+                regime_name=regime_name,
+                internal_regime=internal_regime,
+                subject_states=subject_states,
+            )
         infeasible_mask = np.asarray(~any_feasible)
         infeasible_indices = np.asarray(idx_arr)[infeasible_mask].tolist()
     else:
@@ -620,6 +650,52 @@ def _check_combo(action_kw: dict[str, Array]) -> Array:
     )
 
 
+def _raise_feasibility_type_error(
+    *,
+    exc: TypeError,
+    regime_name: str,
+    internal_regime: InternalRegime,
+    subject_states: dict[str, Array],
+) -> Never:
+    """Re-raise a TypeError from feasibility checking with diagnostic context.
+
+    Args:
+        exc: The original TypeError from the feasibility check.
+        regime_name: Name of the regime being checked.
+        internal_regime: The internal regime containing variable info.
+        subject_states: Mapping of state names to arrays for subjects in
+            this regime.
+
+    Raises:
+        InvalidInitialConditionsError: Always — wraps `exc` with a dtype hint
+            when any discrete state has a non-integer dtype.
+
+    """
+    discrete_names = {
+        name
+        for name, grid in internal_regime.grids.items()
+        if isinstance(grid, DiscreteGrid)
+    }
+
+    bad_dtypes: list[str] = []
+    for name, arr in subject_states.items():
+        if name in discrete_names and not jnp.issubdtype(arr.dtype, jnp.integer):
+            bad_dtypes.append(f"  {name!r}: dtype={arr.dtype} (expected integer)")
+
+    hint = ""
+    if bad_dtypes:
+        hint = (
+            "\n\nDiscrete states with wrong dtype:\n"
+            + "\n".join(bad_dtypes)
+            + "\n\nDiscrete states are used as array indices and must have integer "
+            "dtype. Check that initial conditions encode categorical states as int "
+            "codes, not floats."
+        )
+
+    msg = f"TypeError in feasibility check for regime {regime_name!r}: {exc}{hint}"
+    raise InvalidInitialConditionsError(msg) from exc
+
+
 def _format_infeasibility_message(
     *,
     infeasible_indices: Sequence[int],