Stricter param validation: reject numpy arrays, validate Series labels (#302)

Co-authored-by: timmens <mensingertim@gmail.com>

Author: hmgaudecker

AGENTS.md

@@ -220,6 +220,13 @@ Model(
   \- Simulate forward given solution. `period_to_regime_to_V_arr` is optional; when
   `None`, the model is solved automatically before simulating.
 
+### Derived Categoricals
+
+When `solve()` / `simulate()` parameters are indexed by a DAG function output (not a
+model state/action), pass `derived_categoricals={"name": DiscreteGrid(...)}`. Functions
+used as derived categoricals must return **integer** types, not booleans — JAX cannot
+use booleans as array indices inside JIT. Use `jnp.int32(...)` to cast.
+
 ### SimulationResult
 
 `simulate()` returns a `SimulationResult` object:

docs/user_guide/pandas_interop.md

@@ -143,6 +143,23 @@ derived_categoricals = {
 }
 ```
 
+### Integer return types required
+
+Functions used as derived categoricals must return **integer** values, not booleans. JAX
+cannot use boolean values as array indices inside JIT-compiled code
+(`NonConcreteBooleanIndexError`). If your derived categorical compares states:
+
+```python
+# Wrong — returns bool, fails inside JIT
+def is_good_health(health: DiscreteState) -> BoolND:
+    return health == Health.good
+
+
+# Correct — returns int32
+def is_good_health(health: DiscreteState) -> IntND:
+    return jnp.int32(health == Health.good)
+```
+
 ## Validating Transition Probabilities
 
 Check that a transition probability array has the correct shape, values in $[0, 1]$, and

src/lcm/model.py

@@ -4,14 +4,22 @@
 from pathlib import Path
 from types import MappingProxyType
 
+import pandas as pd
 from jax import Array
 
 from lcm.ages import AgeGrid
+from lcm.exceptions import InvalidParamsError
 from lcm.grids import DiscreteGrid
 from lcm.model_processing import (
     build_regimes_and_template,
     validate_model_inputs,
 )
+from lcm.pandas_utils import (
+    convert_series_in_params,
+    has_series,
+    initial_conditions_from_dataframe,
+)
+from lcm.params import MappingLeaf, SequenceLeaf
 from lcm.params.processing import (
     process_params,
 )
@@ -195,6 +203,7 @@ def solve(
         internal_params = _maybe_convert_series(
             internal_params, model=self, derived_categoricals=derived_categoricals
         )
+        _validate_param_types(internal_params)
         validate_regime_transitions_all_periods(
             internal_regimes=self.internal_regimes,
             internal_params=internal_params,
@@ -283,6 +292,7 @@ def simulate(
         internal_params = _maybe_convert_series(
             internal_params, model=self, derived_categoricals=derived_categoricals
         )
+        _validate_param_types(internal_params)
         if check_initial_conditions:
             validate_initial_conditions(
                 initial_conditions=initial_conditions,
@@ -336,8 +346,6 @@ def _maybe_convert_series(
     | None,
 ) -> InternalParams:
     """Convert pd.Series leaves in params to JAX arrays if any are present."""
-    from lcm.pandas_utils import convert_series_in_params, has_series  # noqa: PLC0415
-
     if derived_categoricals is not None or has_series(internal_params):
         return convert_series_in_params(
             internal_params=internal_params,
@@ -347,17 +355,51 @@ def _maybe_convert_series(
     return internal_params
 
 
+def _validate_param_types(internal_params: InternalParams) -> None:
+    """Raise if any param leaf is not a Python scalar or JAX array.
+
+    After processing, every leaf value (including inside MappingLeaf /
+    SequenceLeaf containers) must be a Python scalar (float, int, bool) or a
+    JAX array. Notably, numpy arrays and pandas Series are not accepted.
+    """
+    for regime_name, regime_params in internal_params.items():
+        for key, value in regime_params.items():
+            _check_leaf(value, f"{regime_name}__{key}")
+
+
+def _check_leaf(value: object, path: str) -> None:
+    """Check a single leaf value, recursing into MappingLeaf/SequenceLeaf."""
+    if isinstance(value, MappingLeaf):
+        for k, v in value.data.items():
+            _check_leaf(v, f"{path}.{k}")
+        return
+    if isinstance(value, SequenceLeaf):
+        for i, v in enumerate(value.data):
+            _check_leaf(v, f"{path}[{i}]")
+        return
+    if isinstance(value, (float, int, bool)):
+        return
+    if hasattr(value, "dtype") and hasattr(value, "shape"):
+        if isinstance(value, Array):
+            return
+        type_name = type(value).__module__ + "." + type(value).__name__
+        msg = (
+            f"Parameter '{path}' is a {type_name} (shape {value.shape}). "
+            f"Use jax.numpy.array() or pass a pd.Series with a named index."
+        )
+        raise InvalidParamsError(msg)
+    type_name = type(value).__module__ + "." + type(value).__name__
+    msg = f"Parameter '{path}' has unexpected type {type_name}."
+    raise InvalidParamsError(msg)
+
+
 def _maybe_convert_dataframe(
     initial_conditions: Mapping[str, Array],
     *,
     model: Model,
 ) -> Mapping[str, Array]:
     """Convert a DataFrame to initial_conditions dict if needed."""
-    import pandas as pd  # noqa: PLC0415
-
     if isinstance(initial_conditions, pd.DataFrame):
-        from lcm.pandas_utils import initial_conditions_from_dataframe  # noqa: PLC0415
-
         return initial_conditions_from_dataframe(df=initial_conditions, model=model)
     return initial_conditions
 

src/lcm/pandas_utils.py

@@ -3,7 +3,7 @@
 from collections.abc import Callable, Mapping
 from dataclasses import dataclass
 from types import MappingProxyType
-from typing import cast
+from typing import TYPE_CHECKING, cast
 
 import jax.numpy as jnp
 import numpy as np
@@ -13,8 +13,11 @@
 
 from lcm.ages import AgeGrid
 from lcm.grids import DiscreteGrid, IrregSpacedGrid
-from lcm.model import Model
 from lcm.params import MappingLeaf
+
+if TYPE_CHECKING:
+    from lcm.model import Model  # avoid circular import: pandas_utils ↔ model
+
 from lcm.params.sequence_leaf import SequenceLeaf
 from lcm.regime import Regime
 from lcm.shocks import _ShockGrid
@@ -726,6 +729,20 @@ def _map_level(*, mapping: _LevelMapping, level_values: pd.Index) -> np.ndarray:
         ValueError: If any label is not valid for the mapping.
 
     """
+    # Categorical levels must use string labels matching grid category names.
+    # Reject integer labels early with a clear message instead of a cryptic KeyError.
+    if mapping.valid_labels and any(not isinstance(v, str) for v in level_values):
+        non_str_types = sorted(
+            {type(v).__name__ for v in level_values if not isinstance(v, str)}
+        )
+        msg = (
+            f"Series index level '{mapping.name}' uses non-string labels "
+            f"(types: {non_str_types}) but the DiscreteGrid expects string "
+            f"category names. Use string labels matching: "
+            f"{sorted(mapping.valid_labels)}."
+        )
+        raise ValueError(msg)
+
     try:
         return np.array([mapping.label_to_index(v) for v in level_values])
     except ValueError:

tests/test_pandas_utils.py

@@ -1097,6 +1097,27 @@ def test_array_from_series_wrong_level_names_raises() -> None:
         )
 
 
+def test_array_from_series_integer_labels_rejected() -> None:
+    """Integer labels on a categorical level raise ValueError."""
+    model = get_stochastic_model(3)
+    # Use integer labels (0, 1) instead of string category names
+    index = pd.MultiIndex.from_tuples(
+        [(40.0, 0, "single", "single"), (40.0, 1, "single", "single")],
+        names=["age", "labor_supply", "partner", "next_partner"],
+    )
+    series = pd.Series([0.5, 0.5], index=index)
+    func = model.regimes["working_life"].get_all_functions()["next_partner"]
+    with pytest.raises(ValueError, match="non-string labels"):
+        array_from_series(
+            sr=series,
+            func=func,
+            param_name="probs_array",
+            func_name="next_partner",
+            model=model,
+            regime_name="working_life",
+        )
+
+
 def test_array_from_series_scalar_param_explicit_lookup() -> None:
     """Scalar parameter with explicit func lookup returns 1D array."""
     model = get_stochastic_model(3)

tests/test_validate_param_types.py

@@ -0,0 +1,68 @@
+"""Test that numpy arrays in params are rejected after processing."""
+
+import jax.numpy as jnp
+import numpy as np
+import pytest
+
+from lcm import AgeGrid, DiscreteGrid, LinSpacedGrid, Model, Regime, categorical
+from lcm.exceptions import InvalidParamsError
+
+
+@categorical(ordered=True)
+class Health:
+    bad: int
+    good: int
+
+
+@categorical(ordered=False)
+class RegimeId:
+    working: int
+    dead: int
+
+
+def _next_regime() -> int:
+    return RegimeId.dead
+
+
+working = Regime(
+    transition=_next_regime,
+    active=lambda age: age < 30,
+    states={
+        "health": DiscreteGrid(Health),
+        "wealth": LinSpacedGrid(start=0, stop=100, n_points=5),
+    },
+    state_transitions={"health": None, "wealth": lambda wealth: wealth},
+    functions={"utility": lambda wealth, health, bonus: wealth + health + bonus},
+)
+
+dead = Regime(
+    transition=None,
+    functions={"utility": lambda: 0.0},
+)
+
+
+def _make_model() -> Model:
+    return Model(
+        regimes={"working": working, "dead": dead},
+        ages=AgeGrid(start=25, stop=30, step="Y"),
+        regime_id_class=RegimeId,
+    )
+
+
+def test_numpy_array_param_rejected() -> None:
+    """Passing a numpy array as a param should raise InvalidParamsError."""
+    model = _make_model()
+    with pytest.raises(InvalidParamsError, match=r"numpy\.ndarray"):
+        model.solve(params={"bonus": np.array(1.0), "discount_factor": 0.95})  # ty: ignore[invalid-argument-type]
+
+
+def test_jax_array_param_accepted() -> None:
+    """JAX arrays should be accepted."""
+    model = _make_model()
+    model.solve(params={"bonus": jnp.array(1.0), "discount_factor": 0.95})
+
+
+def test_python_scalar_param_accepted() -> None:
+    """Python scalars should be accepted."""
+    model = _make_model()
+    model.solve(params={"bonus": 1.0, "discount_factor": 0.95})