Output of Kalman smoother

[SamDuffield]

We had some discussion in #17 regarding what to output from the Kalman smoother in addition to the smoothing means and covariances, i.e. the content of KalmanSmootherInfo.

cuthbert/kalman/smoother.py

Lines 18 to 19 in e35e7e6

	class KalmanSmootherInfo(NamedTuple):
	gains: Array

Currently, it returns the smoother gains which allow you to easily construct the cross covariances $Cov[x_t, x_{t+1}]$ required for EM. But there may be other settings where more info is required and we also want to stay consistent to the square root formulation.

[SamDuffield]

Porting this comment over here #17 (comment)

One thing that might be useful (although I'm not sure when it would be strictly required) would be to reconstruct the joint covariance in square root form.

Concretely the joint smoothing covariance is

$$ Cov[x_{t+1} x_t \mid y_{1:T}] = \begin{pmatrix} P_{t+1|T} & P_{t+1|T} G_t^\intercal \\ G_t P_{t+1|T} & G_t P_{t+1|T}G_t^\intercal +\Omega_{t} \end{pmatrix}, $$

where $\Omega_t = Cov[x_t \mid x_{t+1}, y_{1:t}]$.

This has square root $Cov[x_{t+1} x_t \mid y_{1:T}] = L L^\intercal$ where

$$ L = \begin{pmatrix} P_{t+1|T}^{\frac12} & 0 \\ G_t P_{t+1|T}^{\frac12} & \Omega_{t}^{\frac12} \end{pmatrix}. $$

Which has components $P_{t+1|T}^{\frac12}, G_t, \Omega_{t}^{\frac12}$, which to my understanding are all calculated already:

• $P_{t+1|T}^{\frac12}$ is the square root smoother covariance - part of the output in KalmanState
• $G_t$ are the gain matrices associative_params.E[:-1] - now part of the output
• $\Omega_{t}^{\frac12}$ are the "conditional-reverse-one-step-filtering-covariances" which I believe are calculated in _sqrt_associative_params_single as the D elements.

  cuthbert/kalman/smoother.py

  Line 107 in e35e7e6

  D = Tria_Phi[nx:, nx:]

Relevant references are Section 3.2 of https://arxiv.org/pdf/2207.00426 and P256 of Särkkä & Svensson.

So the open question is whether it's worth also returning $\Omega_t$ in the output. At this stage we don't have a concrete application for the $\Omega_t$ so we can omit but we can now easily revisit in future.