fix types

Author: Kleyt0n

src/dynaris/core/ssm.py

@@ -193,13 +193,15 @@ def fit(
         elif self._filter_name == "ekf":
             self._filter_result = ekf_filter(self._model, obs, initial_state=initial_state)
         elif self._filter_name == "ukf":
+            assert isinstance(self._model, NonlinearSSM)
             self._filter_result = ukf_filter(
                 self._model,
                 obs,
                 initial_state=initial_state,
                 **self._filter_kwargs,
             )
         elif self._filter_name == "particle":
+            assert isinstance(self._model, NonlinearSSM)
             key = self._key if self._key is not None else jax.random.PRNGKey(0)
             self._filter_result = particle_filter(
                 self._model,
@@ -209,6 +211,7 @@ def fit(
                 **self._filter_kwargs,
             )
         elif self._filter_name == "hamilton":
+            assert isinstance(self._model, MarkovSwitchingSSM)
             self._filter_result = hamilton_filter(self._model, obs, initial_state=initial_state)
 
         self._is_fitted = True
@@ -229,9 +232,11 @@ def residuals(self) -> Array:
         if isinstance(self._model, StateSpaceModel):
             from dynaris.estimation.diagnostics import standardized_residuals
 
+            assert isinstance(fr, FilterResult)
             return standardized_residuals(fr, self._model)
 
         # Nonlinear: compute y - h(predicted_state)
+        assert isinstance(self._model, NonlinearSSM)
         predicted_obs = jax.vmap(self._model.h)(fr.predicted_states)
         return fr.observations - predicted_obs
 
@@ -289,11 +294,13 @@ def _plot_filtered(self, **kwargs: Any) -> Any:
         if isinstance(self._model, StateSpaceModel):
             from dynaris.plotting.plots import plot_filtered
 
+            assert isinstance(fr, FilterResult)
             return plot_filtered(fr, self._model, **kwargs)
 
         # Nonlinear: compute observation-space predictions
         import matplotlib.pyplot as plt
 
+        assert isinstance(self._model, NonlinearSSM)
         filtered_obs = jax.vmap(self._model.h)(fr.filtered_states)
         obs = np.asarray(fr.observations)
         filt = np.asarray(filtered_obs)

src/dynaris/core/switching.py

@@ -125,22 +125,22 @@ def obs_dim(self) -> int:
     @property
     def G_stack(self) -> Array:  # noqa: N802
         """Stacked system matrices, shape (K, n, n)."""
-        return self._G_stack  # type: ignore[attr-defined]
+        return self._G_stack  # type: ignore[attr-defined, no-any-return]
 
     @property
     def F_stack(self) -> Array:  # noqa: N802
         """Stacked observation matrices, shape (K, m, n)."""
-        return self._F_stack  # type: ignore[attr-defined]
+        return self._F_stack  # type: ignore[attr-defined, no-any-return]
 
     @property
     def W_stack(self) -> Array:  # noqa: N802
         """Stacked evolution covariances, shape (K, n, n)."""
-        return self._W_stack  # type: ignore[attr-defined]
+        return self._W_stack  # type: ignore[attr-defined, no-any-return]
 
     @property
     def V_stack(self) -> Array:  # noqa: N802
         """Stacked observation covariances, shape (K, m, m)."""
-        return self._V_stack  # type: ignore[attr-defined]
+        return self._V_stack  # type: ignore[attr-defined, no-any-return]
 
     # --- Factory methods ---
 
@@ -185,7 +185,7 @@ def tree_flatten(self) -> tuple[list[Array], dict[str, object]]:
             self.transition_matrix,
             self.initial_probs,
         ]
-        aux = {
+        aux: dict[str, object] = {
             "n_regimes": self.n_regimes,
             "state_dim": self.state_dim,
             "obs_dim": self.obs_dim,

src/dynaris/estimation/bayesian.py

@@ -125,7 +125,7 @@ def model_fn(params):
     def _log_density(params: Array) -> Array:
         model = model_fn(params)
         fr = kalman_filter(model, observations)
-        return fr.log_likelihood + log_prior_fn(params)
+        return fr.log_likelihood + log_prior_fn(params)  # type: ignore[no-any-return]
 
     # NumPyro model: sample unconstrained params, factor by log-density
     def _numpyro_model() -> None:
@@ -149,7 +149,7 @@ def _numpyro_model() -> None:
     @jax.jit
     def _compute_ll(params: Array) -> Array:
         model = model_fn(params)
-        return kalman_filter(model, observations).log_likelihood
+        return kalman_filter(model, observations).log_likelihood  # type: ignore[no-any-return]
 
     log_lls = jax.vmap(_compute_ll)(samples)
 

src/dynaris/estimation/comparison.py

@@ -71,7 +71,7 @@ def compute_waic(
 
     def _pw_ll(params: Array) -> Array:
         model = model_fn(params)
-        return _pointwise_log_likelihood(model, observations)
+        return _pointwise_log_likelihood(model, observations)  # type: ignore[no-any-return]
 
     pw_lls = jax.vmap(_pw_ll)(result.samples)  # (n_samples, T)
 
@@ -109,7 +109,7 @@ def compute_loo(
 
     def _pw_ll(params: Array) -> Array:
         model = model_fn(params)
-        return _pointwise_log_likelihood(model, observations)
+        return _pointwise_log_likelihood(model, observations)  # type: ignore[no-any-return]
 
     pw_lls = jax.vmap(_pw_ll)(result.samples)  # (n_samples, T)
     pw_lls_np = np.asarray(pw_lls)[np.newaxis, :, :]  # (1, n_samples, T)
@@ -158,7 +158,7 @@ def to_arviz(
 
     def _pw_ll(params: Array) -> Array:
         model = model_fn(params)
-        return _pointwise_log_likelihood(model, observations)
+        return _pointwise_log_likelihood(model, observations)  # type: ignore[no-any-return]
 
     pw_lls = jax.vmap(_pw_ll)(result.samples)
 

src/dynaris/estimation/parallel_mcmc.py

@@ -86,7 +86,7 @@ def fit_bayesian_parallel(
     def _log_density(params: Array) -> Array:
         model = model_fn(params)
         fr = kalman_filter(model, observations)
-        return fr.log_likelihood + log_prior_fn(params)
+        return fr.log_likelihood + log_prior_fn(params)  # type: ignore[no-any-return]
 
     def _numpyro_model() -> None:
         params = numpyro.sample(
@@ -119,7 +119,7 @@ def _numpyro_model() -> None:
     @jax.jit
     def _compute_ll(params: Array) -> Array:
         model = model_fn(params)
-        return kalman_filter(model, observations).log_likelihood
+        return kalman_filter(model, observations).log_likelihood  # type: ignore[no-any-return]
 
     log_lls = jax.vmap(_compute_ll)(all_samples)
 

src/dynaris/estimation/predictive.py

@@ -102,7 +102,7 @@ def _forecast_one(params: Array) -> Array:
             cov=fr.filtered_covariances[-1],
         )
         fc = forecast(model, last_state, steps)
-        return fc.mean  # (steps, obs_dim)
+        return fc.mean  # type: ignore[no-any-return]  # (steps, obs_dim)
 
     all_forecasts = jax.vmap(_forecast_one)(samples)  # (n, steps, obs_dim)
 

src/dynaris/filters/ekf.py

@@ -131,7 +131,7 @@ def scan(
         initial_state: GaussianState | None = None,
     ) -> FilterResult:
         """Run full forward EKF via jax.lax.scan."""
-        return ekf_filter(model, observations, initial_state)
+        return ekf_filter(model, observations, initial_state)  # type: ignore[no-any-return]
 
 
 @jax.jit

src/dynaris/filters/hamilton.py

@@ -230,7 +230,7 @@ def hamilton_filter(
     if initial_state is None:
         initial_state = model.initial_state()
 
-    return _hamilton_scan(model, observations, initial_state)
+    return _hamilton_scan(model, observations, initial_state)  # type: ignore[no-any-return]
 
 
 @jax.jit

src/dynaris/filters/particle.py

@@ -181,7 +181,7 @@ def compute_log_weights(particles: Array, observation: Array, model: NonlinearSS
     log_det = jnp.linalg.slogdet(model.R)[1]
     r_inv_diff = jnp.linalg.solve(model.R, diff.T).T  # (N, m)
     mahal = jnp.sum(diff * r_inv_diff, axis=-1)  # (N,)
-    return -0.5 * (m * jnp.log(2.0 * jnp.pi) + log_det + mahal)
+    return -0.5 * (m * jnp.log(2.0 * jnp.pi) + log_det + mahal)  # type: ignore[no-any-return]
 
 
 def _normalize_log_weights(log_weights: Array) -> tuple[Array, Array]:

src/dynaris/filters/ukf.py

@@ -317,7 +317,7 @@ def _ukf_filter_impl(
         initial_state = model.initial_state()
 
     weights = compute_weights(model.state_dim, alpha, beta, kappa)
-    return _ukf_scan(model, observations, initial_state, weights)
+    return _ukf_scan(model, observations, initial_state, weights)  # type: ignore[no-any-return]
 
 
 @jax.jit