Support missing (`nan`) emissions in `LinearGaussianConjugateSSM.fit_blocked_gibbs`

[hylkedonker]

Hi,

First off, thanks for building this awesome package. Well done!
I work in the hospital and I am interested in fitting a Kalman filter on sparsely measured blood levels where most observations are missing. Am I correct in assuming that missing values are currently not supported by LinearGaussianConjugateSSM?
If so, in theory this is not so difficult to implement in the fit_blocked_gibbs if an entire emission is missing, I think.

1. Update forward filtering backwards sampling algorithm with the missing observations integrated out.
2. Update the parameter sampling. Integrating out the missing observations would amount to removing the control-emission pairs from the emission summary statistics where the emission value is missing. More complex would be when part of the emission is missing (giving marginal emissions). That probably requires a generalization of the normal-inverse Wishart distribution. See also this paper: Missing observation analysis for matrix-variate time series data.

If you're interested, I can give the complete missing observation a try. Would you accept such a pull request?

Kind regards,

Hylke

[hylkedonker]

Kalman filter
For step 1, as far as I can tell, missing emissions don't affect the Kalman smoothing algorithm.
For the Kalman filtering step, I believe the following modifications need to made to take into account partially + complete missing emissions. If we represent a missing value by nan, then with is_missing = isnan( $y_t$ ) and $P_+$ = diag[1 - is_missing.astype(int)] then the algorithm remains unchanged upon replacing:
$H_t \rightarrow P_+ H_t;\quad R_t \rightarrow P_+ R_t P_+^T$

Details

Following the notation in [1], denote:

$$ p(\mathbf{z}) = \mathcal{N}(\mathbf{\breve{\mu}}, \mathbf{\breve{\Sigma}}) $$

$$ p(\mathbf{y}|\mathbf{z}) = \mathcal{N}(\mathbf{A} \mathbf{z} + \mathbf{b}, \mathbf{\Omega}) $$

Further, let $\mathbf{y} = [\mathbf{y}_{+}, \mathbf{y}_{-}]^T$ where $\mathbf{y}_{+}$ refer to observed emissions and $\mathbf{y}_{-}$ missing emissions.

Then joint distribution is given by

$p(\mathbf{y}, \mathbf{z}) = \mathcal{N}\left( \begin{pmatrix}\mathbf{\breve{\mu}}\\ \mathbf{\bar{\mu}_+} \\ \mathbf{\bar{\mu}_-} \end{pmatrix}, \begin{pmatrix} \mathbf{\breve{\Sigma}} & C_+ & C_-\\ C_+^T & S_{++} & S_{+-}\\ C_-^T & S_{-+}^T & S_{--}\end{pmatrix} \right)$

where $\bar{\mu} = \mathbf{A} \mathbf{\breve{\mu}} + \mathbf{b}$, $\mathbf{C} = \mathbf{\breve{\Sigma}} \mathbf{A}^T$ and $\mathbf{S} = \mathbf{A} \mathbf{\breve{\Sigma}} \mathbf{A}^T + \mathbf{\Omega}$. Integrating out the unobserved $\mathbf{y}_{-}$ gives:

$\int \mathrm{d}\mathbf{y}_- p(\mathbf{y}, \mathbf{z}) = p(\mathbf{y}_+, \mathbf{z}) = \mathcal{N}\left( \begin{pmatrix}\mathbf{\breve{\mu}}\\ \mathbf{\bar{\mu}_+} \end{pmatrix}, \begin{pmatrix} \mathbf{\breve{\Sigma}} & C_+ \\ C_+^T & S_{++}\end{pmatrix} \right)$

Thus, we do the update step using the block matrices corresponding to the observed values.
Using the definition $K_+ = C_+ S_{++}^{-1}$ the result is

$p(z|y) = \mathcal{N}(z|\breve{\mu} + K_+(y_+-\bar{\mu}_+)), \breve{\Sigma} - K_+ S_{++} K_+^T )$

Now, it is a bit of a nuisance to use different matrix shapes for each time step. However, we can get the correct update by zeroing out the weights and correlation matrix.

Let $P_+ = \begin{pmatrix} 1 & 0 \\ 0 & 0 \end{pmatrix}$ be the projection operator onto the block of observed values and $P_- = 1 - P_+$ the projector onto the missing value block.

Claim: Replacing $A \rightarrow A^\prime = P_+ A$ and $\Omega \rightarrow \Omega^\prime = P_+ \Omega P_+$ gives the correct update.

Proof:
Filling in $A^\prime$ and $\Omega^\prime$ and introducing a small $\epsilon$ we have
$S = P_+ A \Sigma A^T P_+^T + P_+ \Omega P_+ + \epsilon P_-$

$S = \begin{pmatrix} A_+ \\ 0 \end{pmatrix} \Sigma \begin{pmatrix} A_+^T & 0 \end{pmatrix} + \begin{pmatrix} \Omega_{++} & 0 \\ 0 & 0\end{pmatrix} + \epsilon \begin{pmatrix} 0 & 0 \\ 0 & 1\end{pmatrix}$

Next, since this is block-wise a diagonal matrix, the inverse remains block-wise diagonal

$S^{-1} = \begin{pmatrix} A_+ \Sigma A_+^T + \Omega_{++} & 0 \\ 0 & \epsilon \end{pmatrix}^{-1} = \begin{pmatrix} (A_+ \Sigma A_+^T + \Omega_{++})^{-1} & 0 \\ 0 & 1/\epsilon \end{pmatrix} = \begin{pmatrix} S_{++}^{-1} & 0\\ 0 & 1/\epsilon \end{pmatrix}$

The Kalman gain matrix

$K = C S^{-1} = \Sigma A^T P_+ S^{-1} = [C_+, 0]^T \begin{pmatrix} S_{++}^{-1} & 0\\ 0 & 1/\epsilon \end{pmatrix} = [ C_+ S_{++}^{-1}, 0]$

And therefore

$$ K S K^T = C_+ S_{++}^{-1} C_+^T $$

and similarly,

$\mu = \breve{\mu} + K(y - \bar{\mu}) = \breve{\mu} + [C_+ S_{++}^{-1}, 0]^T \begin{pmatrix} y_+ - \bar{\mu}_+\\ y_- - \bar{\mu}_-\end{pmatrix} = \breve{\mu} + C_+ S_{++}^{-1} (y_+ - \bar{\mu}_+)$
which is identical to the term that is obtained from integrating out the unobserved values.

[1] Murphy, Prob. ML: Adv. Topics