Fix the ty failures the state-conditioned field introduced

CI has been failing `ty` since 3dd6ef9 and I never saw it, because I skip the
hook locally (this worktree has no type-checking env). It is reachable with
`uvx ty@0.0.56 check --python <shared pylcm tree>/.pixi/envs/type-checking`.
36 diagnostics -> 0.

Two real annotation bugs of mine:

- `_get_conditioned_weights_func` read `sc.on` / `sc.by` off a
  `StateConditioned | None` without narrowing, and the call site passed
  `Grid | None` into a `DiscreteGrid | None` parameter. Both go away by passing
  the regime's grid mapping plus the already-narrowed `StateConditioned`, and
  resolving the conditioning grid inside — the same shape `next_state.py`
  already uses, so the two seams now read alike.
- `_process_family` returned `str` where `conditioned_row` wants the `Family`
  literal.

The rest are fallout in four test files that splat a kwargs table into a
process constructor with `**`. Those tables were annotated `dict[str, bool]` /
inferred as `dict[str, float]`, and only ever type-checked by luck: every
constructor parameter happened to accept a number (`bool` is an `int`, and
`mu`/`sigma` take `int | float | None`). `state_conditioned: StateConditioned |
None` is the first parameter that does not, so the checker started reporting a
possible landing. `Any` is the honest value type for a splat table; the
alternative was a dozen ignore comments.

Also adds the missing `log_level` to seven `Model.solve` calls in the
state-conditioned tests, which its overloads require.

66 targeted tests pass; the full suite runs on CI.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01DW9D1WhLN8JhhiEPCScbaw

Author: hmgaudecker

src/_lcm/processes/state_conditioned.py

@@ -27,6 +27,7 @@
 from lcm.typing import Float1D, ScalarFloat, ScalarInt
 
 __all__ = [
+    "Family",
     "StateConditioned",
     "conditioned_row",
     "gather_sigma",

src/_lcm/regime_building/processing.py

@@ -46,6 +46,8 @@
 from _lcm.processes.ar1 import TauchenAR1Process
 from _lcm.processes.iid import NormalIIDProcess
 from _lcm.processes.state_conditioned import (
+    Family,
+    StateConditioned,
     conditioned_row,
     gather_sigma,
     sigma_array_by_code,
@@ -1235,13 +1237,7 @@ def _process_regime_core(
             )
     processed_functions |= {
         f"weight_{user_regime}__next_{process}": _get_weights_func_for_process(
-            name=process,
-            grid=grid,
-            conditioning_grid=(
-                all_grids[user_regime].get(grid.state_conditioned.on)
-                if grid.state_conditioned is not None
-                else None
-            ),
+            name=process, grid=grid, grids=all_grids[user_regime]
         )
         for (user_regime, process), grid in target_process_grids.items()
     } | {
@@ -1570,7 +1566,7 @@ def next_func(**kwargs: Any) -> Int1D:  # noqa: ARG001, ANN401
     return next_func
 
 
-def _process_family(grid: _ContinuousStochasticProcess) -> str:
+def _process_family(grid: _ContinuousStochasticProcess) -> Family:
     """Map a supported process to its state-conditioned family (audit F2).
 
     Only families whose transition CDF carries ``sigma`` can express a fixed-node,
@@ -1672,15 +1668,16 @@ def _get_conditioned_weights_func(
     *,
     name: str,
     grid: _ContinuousStochasticProcess,
-    conditioning_grid: DiscreteGrid | None,
+    sc: StateConditioned,
+    grids: Mapping[StateOrActionName, Grid],
 ) -> UserFunction:
     """Weights function for a state-conditioned process (direct-CDF, audit F1/F5/F6).
 
     The transition row is computed DIRECTLY at the from-value on the FIXED common nodes
     (from the scalar ``sigma``), with the per-regime ``sigma`` gathered by the
     conditioning state's integer code. No precomputed-row interpolation (F1).
     """
-    sc = grid.state_conditioned
+    conditioning_grid = grids.get(sc.on)
     if not isinstance(conditioning_grid, DiscreteGrid):
         msg = (
             f"state_conditioned.on='{sc.on}' must name a DiscreteGrid state in the "
@@ -1716,22 +1713,22 @@ def _get_weights_func_for_process(
     *,
     name: str,
     grid: _ContinuousStochasticProcess,
-    conditioning_grid: DiscreteGrid | None = None,
+    grids: Mapping[StateOrActionName, Grid] = MappingProxyType({}),
 ) -> UserFunction:
     """Get function that uses linear interpolation to calculate the process weights.
 
     For processes whose params are supplied at runtime, the grid points and
     transition probabilities are computed inside JIT from those runtime params. For a
     state-conditioned process the row is instead computed directly at the from-value
-    (``_get_conditioned_weights_func``).
+    (``_get_conditioned_weights_func``), and `grids` is that regime's grid mapping, from
+    which the conditioning state is resolved.
 
     """
-    if grid.state_conditioned is not None:
-        # `conditioning_grid` may be None here (unknown `on`); the callee raises a
-        # clear message naming the offending state.
-        return _get_conditioned_weights_func(
-            name=name, grid=grid, conditioning_grid=conditioning_grid
-        )
+    sc = grid.state_conditioned
+    if sc is not None:
+        # `sc.on` may be absent from `grids`; the callee raises a clear message naming
+        # the offending state.
+        return _get_conditioned_weights_func(name=name, grid=grid, sc=sc, grids=grids)
 
     if grid.params_to_pass_at_runtime:
         n_points = grid.n_points

tests/test_models/processes.py

@@ -1,5 +1,5 @@
 import functools
-from typing import Literal
+from typing import Any, Literal
 
 from jax import numpy as jnp
 
@@ -32,7 +32,11 @@
     "rouwenhorst": RouwenhorstAR1Process,
 }
 
-_SHOCK_GRID_KWARGS: dict[str, dict[str, bool]] = {
+# Heterogeneous per-class constructor kwargs, splatted with `**`, so the value type
+# is genuinely `Any`: a checker must assume any key of the target class could receive
+# one. (It typed as `bool` only by luck — every param used to accept a bool, since
+# `bool` is an `int`.)
+_SHOCK_GRID_KWARGS: dict[str, dict[str, Any]] = {
     "uniform": {},
     "normal": {"gauss_hermite": True},
     "lognormal": {"gauss_hermite": True},

tests/test_processes.py

@@ -1,3 +1,5 @@
+from typing import Any
+
 import pandas as pd
 import pytest
 from jax import numpy as jnp
@@ -203,7 +205,7 @@ def test_process_grid_fully_specified_with_all_params(grid_cls, kwargs):
 def test_ar1_grid_centers_on_unconditional_mean(grid_cls):
     """Midpoint of AR(1) gridpoints is approximately mu / (1 - rho)."""
     mu, rho = 2.0, 0.8
-    kwargs = {"rho": rho, "sigma": 0.5, "mu": mu}
+    kwargs: dict[str, Any] = {"rho": rho, "sigma": 0.5, "mu": mu}
     if grid_cls is TauchenAR1Process:
         kwargs["gauss_hermite"] = True
     grid = grid_cls(n_points=11, **kwargs)
@@ -216,7 +218,7 @@ def test_ar1_grid_centers_on_unconditional_mean(grid_cls):
 @pytest.mark.parametrize("grid_cls", _AR1_GRID_CLASSES)
 def test_ar1_transition_probs_rows_sum_to_one(grid_cls):
     """Each row of the transition matrix sums to 1."""
-    kwargs = {"rho": 0.9, "sigma": 0.5, "mu": 1.0}
+    kwargs: dict[str, Any] = {"rho": 0.9, "sigma": 0.5, "mu": 1.0}
     if grid_cls is TauchenAR1Process:
         kwargs["gauss_hermite"] = True
     grid = grid_cls(n_points=7, **kwargs)
@@ -398,7 +400,7 @@ def test_lognormal_gauss_hermite_weights_sum_to_one():
     aaae(P[0].sum(), 1.0, decimal=DECIMAL_PRECISION)
 
 
-_NORMAL_MIXTURE_KWARGS = {
+_NORMAL_MIXTURE_KWARGS: dict[str, Any] = {
     "n_std": 3.0,
     "p1": 0.9,
     "mu1": 0.0,
@@ -438,7 +440,7 @@ def test_iid_normal_mixture_stationary_moments():
     aaae(got_std, float(jnp.sqrt(expected_var)), decimal=1)
 
 
-_TAUCHEN_NORMAL_MIXTURE_KWARGS = {
+_TAUCHEN_NORMAL_MIXTURE_KWARGS: dict[str, Any] = {
     "rho": 0.8,
     "mu": 1.0,
     "n_std": 3.0,

tests/test_runtime_process_params.py

@@ -1,5 +1,7 @@
 """Tests for runtime-supplied stochastic-process parameters."""
 
+from typing import Any
+
 import jax.numpy as jnp
 import pytest
 
@@ -41,7 +43,7 @@ def _constraint(consumption: ContinuousAction, wealth: ContinuousState) -> Float
     return consumption <= wealth
 
 
-_TAUCHEN_PARAMS = {"rho": 0.9, "sigma": 1.0, "mu": 0.0, "n_std": 2}
+_TAUCHEN_PARAMS: dict[str, Any] = {"rho": 0.9, "sigma": 1.0, "mu": 0.0, "n_std": 2}
 
 
 def _make_model(*, fixed_params=None):

tests/test_shock_draw.py

@@ -1,6 +1,7 @@
 """Tests for draw_shock sampling correctness across all shock grid types."""
 
 from types import MappingProxyType
+from typing import Any
 
 import jax
 import pytest
@@ -19,7 +20,7 @@
 
 _N_DRAWS = 10_000
 
-_NORMAL_MIXTURE_KWARGS = {
+_NORMAL_MIXTURE_KWARGS: dict[str, Any] = {
     "n_std": 3.0,
     "p1": 0.9,
     "mu1": 0.0,
@@ -28,7 +29,7 @@
     "sigma2": 0.3,
 }
 
-_TAUCHEN_NORMAL_MIXTURE_KWARGS = {
+_TAUCHEN_NORMAL_MIXTURE_KWARGS: dict[str, Any] = {
     "rho": 0.8,
     "mu": 1.0,
     "n_std": 3.0,
@@ -54,7 +55,7 @@ def _draw_many(grid, params, key_seed=0, current_value=None):
 @pytest.mark.parametrize("params_at_init", [True, False])
 def test_draw_shock_uniform(params_at_init):
     """Uniform.draw_shock uses start/stop params."""
-    kwargs = {"start": 2.0, "stop": 4.0}
+    kwargs: dict[str, Any] = {"start": 2.0, "stop": 4.0}
     if params_at_init:
         grid = UniformIIDProcess(n_points=5, batch_size=0, distributed=False, **kwargs)
         params = grid.params
@@ -70,7 +71,7 @@ def test_draw_shock_uniform(params_at_init):
 @pytest.mark.parametrize("params_at_init", [True, False])
 def test_draw_shock_normal(params_at_init):
     """Normal.draw_shock uses mu/sigma params."""
-    kwargs = {"mu": 5.0, "sigma": 0.1}
+    kwargs: dict[str, Any] = {"mu": 5.0, "sigma": 0.1}
     if params_at_init:
         grid = NormalIIDProcess(
             n_points=5, batch_size=0, distributed=False, gauss_hermite=True, **kwargs
@@ -89,7 +90,7 @@ def test_draw_shock_normal(params_at_init):
 @pytest.mark.parametrize("params_at_init", [True, False])
 def test_draw_shock_lognormal(params_at_init):
     """LogNormal.draw_shock produces positive samples with correct log-moments."""
-    kwargs = {"mu": 1.0, "sigma": 0.1}
+    kwargs: dict[str, Any] = {"mu": 1.0, "sigma": 0.1}
     if params_at_init:
         grid = LogNormalIIDProcess(
             n_points=5, batch_size=0, distributed=False, gauss_hermite=True, **kwargs
@@ -107,7 +108,7 @@ def test_draw_shock_lognormal(params_at_init):
 @pytest.mark.parametrize("params_at_init", [True, False])
 def test_draw_shock_tauchen(params_at_init):
     """Tauchen.draw_shock uses mu/sigma/rho params."""
-    kwargs = {"rho": 0.5, "sigma": 0.1, "mu": 2.0}
+    kwargs: dict[str, Any] = {"rho": 0.5, "sigma": 0.1, "mu": 2.0}
     if params_at_init:
         grid = TauchenAR1Process(
             n_points=5, batch_size=0, distributed=False, gauss_hermite=True, **kwargs
@@ -124,7 +125,7 @@ def test_draw_shock_tauchen(params_at_init):
 @pytest.mark.parametrize("params_at_init", [True, False])
 def test_draw_shock_rouwenhorst(params_at_init):
     """Rouwenhorst.draw_shock uses mu/sigma/rho params."""
-    kwargs = {"rho": 0.5, "sigma": 0.1, "mu": 2.0}
+    kwargs: dict[str, Any] = {"rho": 0.5, "sigma": 0.1, "mu": 2.0}
     if params_at_init:
         grid = RouwenhorstAR1Process(
             n_points=5, batch_size=0, distributed=False, **kwargs
@@ -186,7 +187,7 @@ def test_draw_shock_tauchen_normal_mixture(params_at_init):
 def test_ar1_draw_shock_unconditional_moments(grid_cls):
     """Long-run simulated moments match AR(1) unconditional moments."""
     mu, rho, sigma = 0.5, 0.7, 0.3
-    kwargs = {"rho": rho, "sigma": sigma, "mu": mu}
+    kwargs: dict[str, Any] = {"rho": rho, "sigma": sigma, "mu": mu}
     if grid_cls is TauchenAR1Process:
         kwargs["gauss_hermite"] = True
     grid = grid_cls(n_points=11, **kwargs)

tests/test_state_conditioned_model.py

@@ -269,7 +269,7 @@ def test_gauss_hermite_state_conditioned_rejected():
         ),
     )
     with pytest.raises(ModelInitializationError, match="Gauss-Hermite"):
-        _model_with_income(income).solve(params=_params())
+        _model_with_income(income).solve(params=_params(), log_level="warning")
 
 
 def test_rouwenhorst_state_conditioned_rejected():
@@ -284,7 +284,7 @@ def test_rouwenhorst_state_conditioned_rejected():
         ),
     )
     with pytest.raises(ModelInitializationError, match="only supported for CDF-binned"):
-        _model_with_income(income).solve(params=_params())
+        _model_with_income(income).solve(params=_params(), log_level="warning")
 
 
 def test_unknown_conditioning_state_rejected():
@@ -298,7 +298,7 @@ def test_unknown_conditioning_state_rejected():
         state_conditioned=StateConditioned(on="nope", by={"low": 0.1, "high": 0.3}),
     )
     with pytest.raises(ModelInitializationError, match="must name a DiscreteGrid"):
-        _model_with_income(income).solve(params=_params())
+        _model_with_income(income).solve(params=_params(), log_level="warning")
 
 
 @categorical(ordered=True)
@@ -355,7 +355,7 @@ def test_gauss_hermite_tauchen_state_conditioned_rejected():
         ),
     )
     with pytest.raises(ModelInitializationError, match="Gauss-Hermite"):
-        _model_with_income(income).solve(params=_params())
+        _model_with_income(income).solve(params=_params(), log_level="warning")
 
 
 def test_runtime_grid_param_rejected():
@@ -375,7 +375,7 @@ def test_runtime_grid_param_rejected():
         ),
     )
     with pytest.raises(ModelInitializationError, match="every grid parameter fixed"):
-        _model_with_income(income).solve(params=_params())
+        _model_with_income(income).solve(params=_params(), log_level="warning")
 
 
 @pytest.mark.parametrize("bad", [float("nan"), float("inf")])
@@ -392,7 +392,7 @@ def test_nonfinite_sigma_rejected(bad):
         ),
     )
     with pytest.raises(ModelInitializationError, match="finite positive sigmas"):
-        _model_with_income(income).solve(params=_params())
+        _model_with_income(income).solve(params=_params(), log_level="warning")
 
 
 def test_nonpositive_sigma_rejected():
@@ -408,4 +408,4 @@ def test_nonpositive_sigma_rejected():
         ),
     )
     with pytest.raises(ModelInitializationError, match="finite positive sigmas"):
-        _model_with_income(income).solve(params=_params())
+        _model_with_income(income).solve(params=_params(), log_level="warning")

tests/test_state_conditioned_shocks.py

@@ -166,7 +166,7 @@ def test_state_conditioned_dataclass_is_frozen():
     sc = StateConditioned(on="uncertainty", by={"low": 0.2, "high": 1.0})
     assert sc.on == "uncertainty"
     with pytest.raises((AttributeError, TypeError)):
-        sc.on = "other"  # type: ignore[misc]
+        sc.on = "other"  # ty: ignore[invalid-assignment]
 
 
 def test_conditioned_row_dispatch_sums_to_one():