Steady state kalman filtering

[AdrienCorenflos]

No description provided.

[AdrienCorenflos]

The computational improvements exist by are not as good as I expected. I am wondering if perhaps I missed some simple improvement, or maybe it's a feature of testing with a very high dim state which triggers an expensive QR decomposition in the prediction anyway. Not sure, I would appreciate a second look.

[SamDuffield]

Instead of the extra arg could we have get_observation_params possibly return SteadyStateFilterParams if using steady state?

[SamDuffield]

The computational improvements exist by are not as good as I expected. I am wondering if perhaps I missed some simple improvement, or maybe it's a feature of testing with a very high dim state which triggers an expensive QR decomposition in the prediction anyway. Not sure, I would appreciate a second look.

Seems to me that this PR reduces the number of tria calls from 4 or 5 (4 in the example) to 3 per step? On my laptop we get a 1.16x/1.52x speedup (which is more than in your doc). Seems some discrepancy across devices. But I don't have a gut feel if we should expect more than that

[Sahel13]

This doesn't handle missing observation dimensions correctly. The following test fails when (and only when) y has a nan in it:

@pytest.mark.parametrize("seed", [0, 42, 99])
def test_update_steady_state_matches_qr_path_with_missing_observations(seed):
    """Steady-state update must preserve the standard NaN-observation behavior."""
    nx, ny = 4, 3
    F, _, chol_Q, H, d, chol_R, *_ = _make_lgssm(seed, nx, ny)
    ss = compute_steady_state_filter_params(
        jnp.array(F),
        jnp.array(chol_Q),
        jnp.array(H),
        jnp.array(chol_R),
    )

    rng = np.random.default_rng(seed + 10)
    m = jnp.array(rng.standard_normal(nx))
    y = jnp.array(rng.standard_normal(ny)).at[1].set(jnp.nan)

    (m_qr, chol_P_qr), ell_qr = update(
        m,
        jnp.array(chol_Q),
        jnp.array(H),
        jnp.array(d),
        jnp.array(chol_R),
        y,
    )
    (m_ss, chol_P_ss), ell_ss = update(
        m,
        jnp.array(chol_Q),
        jnp.array(H),
        jnp.array(d),
        jnp.array(chol_R),
        y,
        steady_state_params=ss,
    )

    chex.assert_trees_all_close(m_ss, m_qr, atol=1e-10)
    chex.assert_trees_all_close(
        chol_P_ss @ chol_P_ss.T,
        chol_P_qr @ chol_P_qr.T,
        atol=1e-10,
    )
    chex.assert_trees_all_close(ell_ss, ell_qr, atol=1e-10)

[Sahel13]

Suggested change

	Fields:
	Attributes:

[Sahel13]

these can have type ArrayLike