fix ruff

Author: Kleyt0n

src/dynaris/core/nonlinear.py

@@ -2,8 +2,8 @@
 
 from __future__ import annotations
 
+from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Callable
 
 import jax
 import jax.numpy as jnp

src/dynaris/filters/ekf.py

@@ -50,8 +50,8 @@ def predict(state: GaussianState, model: NonlinearSSM) -> GaussianState:
     R_t = F_t @ C_{t-1} @ F_t' + Q
     """
     mean = model.f(state.mean)
-    F_jac = jax.jacfwd(model.f)(state.mean)  # (n, n)
-    cov = F_jac @ state.cov @ F_jac.T + model.Q
+    f_jac = jax.jacfwd(model.f)(state.mean)  # (n, n)
+    cov = f_jac @ state.cov @ f_jac.T + model.Q
     return GaussianState(mean=mean, cov=cov)
 
 
@@ -70,22 +70,22 @@ def update(
     """
     y = observation
     y_pred = model.h(predicted.mean)  # (m,)
-    H_jac = jax.jacfwd(model.h)(predicted.mean)  # (m, n)
+    h_jac = jax.jacfwd(model.h)(predicted.mean)  # (m, n)
 
     e = y - y_pred  # innovation (m,)
-    S = H_jac @ predicted.cov @ H_jac.T + model.R  # innovation covariance (m, m)
+    s = h_jac @ predicted.cov @ h_jac.T + model.R  # innovation covariance (m, m)
 
     # Kalman gain: K = P @ H' @ S^{-1}
-    K = jnp.linalg.solve(S.T, (predicted.cov @ H_jac.T).T).T  # (n, m)
+    k_gain = jnp.linalg.solve(s.T, (predicted.cov @ h_jac.T).T).T  # (n, m)
 
-    filtered_mean = predicted.mean + K @ e
+    filtered_mean = predicted.mean + k_gain @ e
     identity = jnp.eye(predicted.mean.shape[-1])
-    filtered_cov = (identity - K @ H_jac) @ predicted.cov
+    filtered_cov = (identity - k_gain @ h_jac) @ predicted.cov
 
     # Log-likelihood: log N(e; 0, S)
     m = observation.shape[-1]
-    log_det = jnp.linalg.slogdet(S)[1]
-    mahal = e @ jnp.linalg.solve(S, e)
+    log_det = jnp.linalg.slogdet(s)[1]
+    mahal = e @ jnp.linalg.solve(s, e)
     ll = -0.5 * (m * jnp.log(2.0 * jnp.pi) + log_det + mahal)
 
     # Handle missing observations: if any element is NaN, skip update
@@ -181,9 +181,7 @@ def h(x):
         log_likelihood=jnp.array(0.0),
     )
 
-    def _scan_step(
-        carry: _ScanCarry, obs: Array
-    ) -> tuple[_ScanCarry, _ScanOutput]:
+    def _scan_step(carry: _ScanCarry, obs: Array) -> tuple[_ScanCarry, _ScanOutput]:
         predicted = predict(carry.filtered, model)
         filtered, ll = update(predicted, obs, model)
         new_carry = _ScanCarry(

src/dynaris/filters/ukf.py

@@ -79,14 +79,17 @@ def sigma_points(state: GaussianState, lam: Array) -> Array:
     """
     n = state.mean.shape[0]
     scaled_cov = (n + lam) * state.cov
-    L = jnp.linalg.cholesky(scaled_cov)  # (n, n)
+    chol = jnp.linalg.cholesky(scaled_cov)  # (n, n)
 
     # Build sigma points: [mean, mean + L_i, mean - L_i]
-    offsets = jnp.concatenate([
-        jnp.zeros((1, n)),
-        L,    # rows of L as positive offsets
-        -L,   # rows of L as negative offsets
-    ], axis=0)  # (2n+1, n)
+    offsets = jnp.concatenate(
+        [
+            jnp.zeros((1, n)),
+            chol,  # rows of L as positive offsets
+            -chol,  # rows of L as negative offsets
+        ],
+        axis=0,
+    )  # (2n+1, n)
 
     return state.mean[None, :] + offsets
 
@@ -162,29 +165,29 @@ def update(
     # Predicted observation mean
     y_pred = jnp.sum(weights.wm[:, None] * pts_obs, axis=0)  # (m,)
 
-    # Innovation covariance S = sum wc * (y_diff)(y_diff)' + R
+    # Innovation covariance s = sum wc * (y_diff)(y_diff)' + R
     y_diff = pts_obs - y_pred[None, :]  # (2n+1, m)
-    S = jnp.sum(weights.wc[:, None, None] * (y_diff[:, :, None] * y_diff[:, None, :]), axis=0)
-    S = S + model.R  # (m, m)
+    s = jnp.sum(weights.wc[:, None, None] * (y_diff[:, :, None] * y_diff[:, None, :]), axis=0)
+    s = s + model.R  # (m, m)
 
-    # Cross-covariance P_xy = sum wc * (x_diff)(y_diff)'
+    # Cross-covariance p_xy = sum wc * (x_diff)(y_diff)'
     x_diff = pts - predicted.mean[None, :]  # (2n+1, n)
-    P_xy = jnp.sum(weights.wc[:, None, None] * (x_diff[:, :, None] * y_diff[:, None, :]), axis=0)
+    p_xy = jnp.sum(weights.wc[:, None, None] * (x_diff[:, :, None] * y_diff[:, None, :]), axis=0)
     # (n, m)
 
-    # Kalman gain K = P_xy @ S^{-1}
-    K = jnp.linalg.solve(S.T, P_xy.T).T  # (n, m)
+    # Kalman gain k = p_xy @ s^{-1}
+    k_gain = jnp.linalg.solve(s.T, p_xy.T).T  # (n, m)
 
     # Innovation
     e = y - y_pred  # (m,)
 
-    filtered_mean = predicted.mean + K @ e
-    filtered_cov = predicted.cov - K @ S @ K.T
+    filtered_mean = predicted.mean + k_gain @ e
+    filtered_cov = predicted.cov - k_gain @ s @ k_gain.T
 
-    # Log-likelihood: log N(e; 0, S)
+    # Log-likelihood: log N(e; 0, s)
     m = observation.shape[-1]
-    log_det = jnp.linalg.slogdet(S)[1]
-    mahal = e @ jnp.linalg.solve(S, e)
+    log_det = jnp.linalg.slogdet(s)[1]
+    mahal = e @ jnp.linalg.solve(s, e)
     ll = -0.5 * (m * jnp.log(2.0 * jnp.pi) + log_det + mahal)
 
     # Handle missing observations
@@ -248,8 +251,12 @@ def scan(
     ) -> FilterResult:
         """Run full forward UKF via jax.lax.scan."""
         return _ukf_filter_impl(
-            model, observations, initial_state,
-            self.alpha, self.beta, self.kappa,
+            model,
+            observations,
+            initial_state,
+            self.alpha,
+            self.beta,
+            self.kappa,
         )
 
 
@@ -326,9 +333,7 @@ def _ukf_scan(
         log_likelihood=jnp.array(0.0),
     )
 
-    def _scan_step(
-        carry: _ScanCarry, obs: Array
-    ) -> tuple[_ScanCarry, _ScanOutput]:
+    def _scan_step(carry: _ScanCarry, obs: Array) -> tuple[_ScanCarry, _ScanOutput]:
         predicted = predict(carry.filtered, model, weights)
         filtered, ll = update(predicted, obs, model, weights)
         new_carry = _ScanCarry(

tests/test_filters/test_ekf.py

@@ -23,12 +23,10 @@
 # ---------------------------------------------------------------------------
 
 
-def _linear_nonlinear_model(
-    sigma_level: float = 1.0, sigma_obs: float = 1.0
-) -> NonlinearSSM:
+def _linear_nonlinear_model(sigma_level: float = 1.0, sigma_obs: float = 1.0) -> NonlinearSSM:
     """Local-level model as a NonlinearSSM (identity transition/observation)."""
-    Q = jnp.array([[sigma_level**2]])
-    R = jnp.array([[sigma_obs**2]])
+    q = jnp.array([[sigma_level**2]])
+    r = jnp.array([[sigma_obs**2]])
 
     def f(x: Array) -> Array:
         return x
@@ -39,8 +37,8 @@ def h(x: Array) -> Array:
     return NonlinearSSM(
         transition_fn=f,
         observation_fn=h,
-        transition_cov=Q,
-        observation_cov=R,
+        transition_cov=q,
+        observation_cov=r,
         state_dim=1,
         obs_dim=1,
     )
@@ -72,15 +70,15 @@ def test_predict_identity_transition() -> None:
 
 def test_predict_nonlinear_transition() -> None:
     """Test with a nonlinear transition: f(x) = x + 0.1 * sin(x)."""
-    Q = jnp.array([[0.5]])
+    q = jnp.array([[0.5]])
 
     def f(x: Array) -> Array:
         return x + 0.1 * jnp.sin(x)
 
     model = NonlinearSSM(
         transition_fn=f,
         observation_fn=lambda x: x,
-        transition_cov=Q,
+        transition_cov=q,
         observation_cov=jnp.array([[1.0]]),
         state_dim=1,
         obs_dim=1,
@@ -140,15 +138,11 @@ def test_ekf_matches_kalman_on_linear_model() -> None:
     ekf_result = ekf_filter(nl_model, observations, initial_state=init)
     kf_result = kalman_filter(lin_model, observations, initial_state=init)
 
-    np.testing.assert_allclose(
-        ekf_result.filtered_states, kf_result.filtered_states, atol=1e-4
-    )
+    np.testing.assert_allclose(ekf_result.filtered_states, kf_result.filtered_states, atol=1e-4)
     np.testing.assert_allclose(
         ekf_result.filtered_covariances, kf_result.filtered_covariances, atol=1e-3
     )
-    np.testing.assert_allclose(
-        ekf_result.log_likelihood, kf_result.log_likelihood, atol=1e-2
-    )
+    np.testing.assert_allclose(ekf_result.log_likelihood, kf_result.log_likelihood, atol=1e-2)
 
 
 # ---------------------------------------------------------------------------
@@ -198,9 +192,7 @@ def test_ekf_filter_with_missing_obs() -> None:
     assert jnp.isfinite(result.log_likelihood)
 
     # At NaN points, predicted == filtered
-    np.testing.assert_allclose(
-        result.filtered_states[10], result.predicted_states[10], atol=1e-5
-    )
+    np.testing.assert_allclose(result.filtered_states[10], result.predicted_states[10], atol=1e-5)
 
 
 # ---------------------------------------------------------------------------
@@ -280,7 +272,7 @@ def h(x: Array) -> Array:
     # Simulate observations from a known trajectory
     true_state = jnp.array([3.0, 4.0])
     obs_list = []
-    for t in range(50):
+    for _t in range(50):
         true_state = f(true_state) + jax.random.normal(key, (2,)) * 0.01
         key, _ = jax.random.split(key)
         obs = h(true_state) + jax.random.normal(key, (2,)) * 0.1
@@ -315,13 +307,13 @@ def test_grad_through_ekf() -> None:
     observations = NILE[:20].reshape(-1, 1)
 
     def neg_ll(log_sigma_level: Array, log_sigma_obs: Array) -> Array:
-        Q = jnp.exp(log_sigma_level) * jnp.eye(1)
-        R = jnp.exp(log_sigma_obs) * jnp.eye(1)
+        q_mat = jnp.exp(log_sigma_level) * jnp.eye(1)
+        r_mat = jnp.exp(log_sigma_obs) * jnp.eye(1)
         model = NonlinearSSM(
             transition_fn=lambda x: x,
             observation_fn=lambda x: x,
-            transition_cov=Q,
-            observation_cov=R,
+            transition_cov=q_mat,
+            observation_cov=r_mat,
             state_dim=1,
             obs_dim=1,
         )

tests/test_filters/test_ukf.py

@@ -16,9 +16,9 @@
 from dynaris.filters.ukf import (
     UnscentedKalmanFilter,
     compute_weights,
+    predict,
     sigma_points,
     ukf_filter,
-    predict,
     update,
 )
 
@@ -30,9 +30,7 @@
 # ---------------------------------------------------------------------------
 
 
-def _linear_nonlinear_model(
-    sigma_level: float = 1.0, sigma_obs: float = 1.0
-) -> NonlinearSSM:
+def _linear_nonlinear_model(sigma_level: float = 1.0, sigma_obs: float = 1.0) -> NonlinearSSM:
     """Local-level model as a NonlinearSSM."""
     return NonlinearSSM(
         transition_fn=lambda x: x,
@@ -96,9 +94,7 @@ def test_sigma_points_symmetric() -> None:
     w = compute_weights(n=1)
     pts = sigma_points(state, w.lam)
     # Points 1 and 2 should be equidistant from the mean
-    np.testing.assert_allclose(
-        pts[1] - state.mean, -(pts[2] - state.mean), atol=1e-6
-    )
+    np.testing.assert_allclose(pts[1] - state.mean, -(pts[2] - state.mean), atol=1e-6)
 
 
 def test_sigma_points_weighted_mean_recovers_mean() -> None:
@@ -202,9 +198,7 @@ def test_ukf_matches_kalman_on_linear_model() -> None:
     np.testing.assert_allclose(
         ukf_result.filtered_states[10:], kf_result.filtered_states[10:], atol=0.5
     )
-    np.testing.assert_allclose(
-        ukf_result.log_likelihood, kf_result.log_likelihood, atol=5.0
-    )
+    np.testing.assert_allclose(ukf_result.log_likelihood, kf_result.log_likelihood, atol=5.0)
 
 
 # ---------------------------------------------------------------------------
@@ -251,9 +245,7 @@ def test_ukf_filter_with_missing_obs() -> None:
     result = ukf_filter(model, observations)
     assert jnp.all(jnp.isfinite(result.filtered_states))
     assert jnp.isfinite(result.log_likelihood)
-    np.testing.assert_allclose(
-        result.filtered_states[10], result.predicted_states[10], atol=1e-5
-    )
+    np.testing.assert_allclose(result.filtered_states[10], result.predicted_states[10], atol=1e-5)
 
 
 # ---------------------------------------------------------------------------
@@ -366,13 +358,13 @@ def test_grad_through_ukf() -> None:
     observations = NILE[:20].reshape(-1, 1)
 
     def neg_ll(log_sigma_level: Array, log_sigma_obs: Array) -> Array:
-        Q = jnp.exp(log_sigma_level) * jnp.eye(1)
-        R = jnp.exp(log_sigma_obs) * jnp.eye(1)
+        q = jnp.exp(log_sigma_level) * jnp.eye(1)
+        r = jnp.exp(log_sigma_obs) * jnp.eye(1)
         model = NonlinearSSM(
             transition_fn=lambda x: x,
             observation_fn=lambda x: x,
-            transition_cov=Q,
-            observation_cov=R,
+            transition_cov=q,
+            observation_cov=r,
             state_dim=1,
             obs_dim=1,
         )