LGSSM notebook Latex and Code Mismatch and other improvements

[canyon289]

I believe the order of the indices in the code versus the latex code are different for the lgssm docs, where in the code the positions are the first two indices as

Latex

\underbrace{\begin{pmatrix} u_t\\ \dot{u}_t \\ v_t \\ \dot{v}_t \end{pmatrix}}_{z_t}

Code

\underbrace{\begin{pmatrix} u_t\\  v_t \\  \dot{u}_t \\\dot{v}_t \end{pmatrix}}_{z_t}

In addition to this the docs could generally be improved by labeling the axes figures so its more apparent what is what, and adding comments that certain covariances were scaled to be quite small so the true movement would be "straight' whereas others were multiplied by a float so jax won't complain.

Happy to take this one if I

1. Am indeed correct that the notebook disagrees with itself
2. That additional fixes to the documentation are welcome

Screenshots

[gileshd]

Hi @canyon289! Nice catch!

Yes I think there is a problem in the equation in the docs which should read:

\underbrace{\begin{pmatrix} u_t\\  v_t \\  \dot{u}_t \\\dot{v}_t \end{pmatrix}}_{z_t}

this would involve also updating the following equation to read:

\begin{align*}
\underbrace{\begin{pmatrix}  y_{1,t} \\  y_{2,t} \end{pmatrix}}_{y_t}
  &=
    \underbrace{
    \begin{pmatrix}
1 & 0 & 0 & 0 \\
0 & 1 & 0 & 0
    \end{pmatrix}
    }_{H}

\underbrace{\begin{pmatrix} u_t\\ v_t \\ \dot{u}_t \\ \dot{v}_t \end{pmatrix}}_{z_t}    
 + r_t
\end{align*}

If you are happy to give this a shot go ahead and open a PR 👍 and feel free to also add axis labels.

re multiplying by a float so that jax doesn't complain are you referring to:

emission_covariance = jnp.eye(emission_dim) * 1.0

If so I think that the * 1.0 is probably superfluous because, as far as I know, jnp.eye() returns floats anyway.

[canyon289]

Thanks for confirming. Ill be happy to open a PR. Please give me till Sunday evening

[canyon289]

My apologies give me a couple more days. The week got away from me