Accept shock grid columns in initial_conditions_from_dataframe (#306)

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

Author: hmgaudecker

src/lcm/pandas_utils.py

@@ -800,42 +800,38 @@ def _validate_state_columns(
     initial_regimes: list[str],
 ) -> None:
     """Validate that DataFrame columns match model states."""
-    all_states = _collect_all_state_names(
-        regimes=regimes, initial_regimes=initial_regimes
-    )
+    expected = _collect_state_names(regimes=regimes, initial_regimes=initial_regimes)
 
-    unknown = state_columns - all_states
+    unknown = state_columns - expected
     if unknown:
         msg = (
             f"Unknown columns not matching any model state: {sorted(unknown)}. "
-            f"Expected states: {sorted(all_states)}."
+            f"Expected states: {sorted(expected)}."
         )
         raise ValueError(msg)
 
-    missing = all_states - state_columns
+    missing = expected - state_columns
     if missing:
-        msg = (
-            f"Missing required state columns: {sorted(missing)}. "
-            f"All non-shock states must be provided."
-        )
+        msg = f"Missing required state columns: {sorted(missing)}."
         raise ValueError(msg)
 
 
-def _collect_all_state_names(
+def _collect_state_names(
     *,
     regimes: Mapping[str, Regime],
     initial_regimes: list[str],
 ) -> set[str]:
-    """Collect all non-shock state names from regimes present in initial_regimes."""
-    state_names: set[str] = set()
+    """Collect all state names (including shock grids) from initial regimes.
+
+    Returns:
+        Set of all state names from the initial regimes, plus `'age'`
+        (always required).
+
+    """
+    names: set[str] = {"age"}
     for regime_name in set(initial_regimes):
-        regime = regimes[regime_name]
-        for name, grid in regime.states.items():
-            if not isinstance(grid, _ShockGrid):
-                state_names.add(name)
-    # Always include age
-    state_names.add("age")
-    return state_names
+        names.update(regimes[regime_name].states.keys())
+    return names
 
 
 def _build_discrete_grid_lookup(

tests/test_pandas_utils.py

@@ -32,6 +32,7 @@
     get_model as get_basic_model,
 )
 from tests.test_models.regime_markov import get_model as get_regime_markov_model
+from tests.test_models.shock_grids import get_model as get_shock_model
 from tests.test_models.stochastic import get_model as get_stochastic_model
 
 
@@ -270,6 +271,39 @@ def test_missing_state_column_raises():
         initial_conditions_from_dataframe(df=df, model=model)
 
 
+def test_shock_state_columns_accepted():
+    """Shock grid columns are accepted as continuous float columns."""
+    model = get_shock_model(n_periods=4, distribution_type="uniform")
+    df = pd.DataFrame(
+        {
+            "regime": ["alive", "alive"],
+            "wealth": [2.0, 4.0],
+            "health": ["bad", "good"],
+            "income": [0.3, 0.7],
+            "age": [0.0, 0.0],
+        }
+    )
+    conditions = initial_conditions_from_dataframe(df=df, model=model)
+    assert jnp.allclose(conditions["income"], jnp.array([0.3, 0.7]))
+    assert jnp.allclose(conditions["wealth"], jnp.array([2.0, 4.0]))
+    assert "regime" in conditions
+
+
+def test_shock_state_columns_required():
+    """DataFrame without shock columns raises (shocks are required)."""
+    model = get_shock_model(n_periods=4, distribution_type="uniform")
+    df = pd.DataFrame(
+        {
+            "regime": ["alive", "alive"],
+            "wealth": [2.0, 4.0],
+            "health": ["bad", "good"],
+            "age": [0.0, 0.0],
+        }
+    )
+    with pytest.raises(ValueError, match=r"Missing required state columns.*income"):
+        initial_conditions_from_dataframe(df=df, model=model)
+
+
 def test_round_trip_with_discrete_model():
     """Verify DataFrame-based initial states match raw arrays."""
     from tests.test_models.deterministic.discrete import (  # noqa: PLC0415