(nm_lu_decomposition)=
# LU decomposition
```{index} LU decomposition
```
## Motivation

The Gaussian elimination solves a matrix system with one RHS vector \\(\pmb{b}\\). We often have to deal with problems where we have multiple RHS vectors, all with the same matrix \\(A\\).

We could call the same code (e.g. `upper_triangle` and `back_substitution` from {ref}`Gaussian elimination <nm_gaussian_elimination_code>` notebook) multiple times to solve all of these corresponding linear systems, but note that as the elimination algorithm is actually performing operations based upon (the same) \\(A\\) each time, we would actually be wasting time repeating exactly the same operations - this is therefore clearly not an efficient solution to the problem.

We could easily generalise our Gaussian elimination/back substitution algorithms to include multiple RHS column vectors in the augmented system, perform the same sequence of row operations (but now only once) to transform the matrix to upper-triangular form, and then perform back substitution on each of the transformed RHS vectors from the augmented system - cf. the use of Gaussian elimination to compute the inverse to a matrix by placing the identity on the right of the augmented system.

However, it is often the case that each RHS vector depends on the solutions to the matrix systems obtained from some or all of the earlier RHS vectors, and so this generalisation would not work in this case. For example, our discrete system could be of the form

\\[A\pmb{x}^{n+1} = \pmb{b}(\pmb{x}^n, \ldots)\\]

```{margin} Note
This is true if our model is linear, in the nonlinear case \\(A\\) would be updated evey time step.
```
where \\(\pmb{x}^{n+1}\\) is the numerical solution of an ordinary or partial differential equation at time level \\(n+1\\), the RHS is a function of the the solution at the previous time level \\(n\\) (and possibly other things such as forcing terms, represented by the \\(\ldots\\) above. \\(A\\) is then a discretisation of the differential equation written in the form that it maps the old solution to the new one. Time stepping this problem invovles solving the same linear system multiple times with different RHSs.

## Theory

To deal with this situation efficiently we decompose or factorise the matrix \\(A\\) in such a way that it is cheap to compute a new solution vector \\(\pmb{x}\\) for any given RHS vector \\(\pmb{b}\\). This decomposition involves a lower- and an upper-triangular matrix, hence the name **LU decomposition**. These matrices encode the steps conducted in Gaussian elimination, so we don't have to explicilty conduct all of the operations again and again.

Mathematically, let's assume that we have already found/constructed a lower-triangular matrix (\\(L\\) - where all entries above the diagonal are zero) and an upper-triangular matrix (\\(U\\) - where all entries below the diagonal are zero) such that we can write

\\[ A = LU.\\]

In this case the matrix system we need to solve for \\(\pmb{x}\\) becomes

\\[ A\pmb{x} = \pmb{b} \iff (LU)\pmb{x} = L(U\pmb{x}) = \pmb{b}.\\]

Notice that the matrix-vector product \\(U\pmb{x}\\) is itself a vector, let's call it \\(\pmb{c}\\) for the time-being (i.e. \\(\pmb{c}=U\pmb{x}\\)).

The above system then reads

\\[ L\pmb{c} = \pmb{b}, \\]

where \\(L\\) is a matrix and \\(\pmb{c}\\) is an unknown.

As \\(L\\) is in the lower-triangular form, we can use forward substitution (generalising the back subsitution algorithm/code) to very easily find \\(\pmb{c}\\) in relatively few operations (we don't need to call the entire Gaussian elimination algorithm).

Once we know \\(\pmb{c}\\) we then solve the second linear system

\\[ U\pmb{x} = \pmb{c},\\]

where now we can use the fact that \\(U\\) is upper-triangular to use our back substitution algorithm again very efficiently to give the solution \\(\pmb{x}\\) we require.

So for a given \\(\pmb{b}\\) we can find the corresponding \\(\pmb{x}\\) very efficiently, we can therefore do this repeatedly as each new \\(\pmb{b}\\) is given to us.

Our challenge is therefore to find the matrices \\(L\\) and \\(U\\) allowing us to perform the decomposition \\(A=LU\\).

## Doolittle algorithm

Recall the comment above on the \\(L\\) and \\(U\\) matrices encoding the steps taken in Gaussian elimination. Let's see how this works through the development of the so-called **Doolittle algorithm**.

Let's consider an example matrix:

\\[
  A=\begin{bmatrix}
    {\color{black}5} & {\color{black}7} & {\color{black}5} & {\color{black}9}\\\\\\
{\color{black}5} & {\color{black}14} & {\color{black}7} & {\color{black}10}\\\\\\
{\color{black}20} & {\color{black}77} & {\color{black}41} & {\color{black}48}\\\\\\
{\color{black}25} & {\color{black}91} & {\color{black}55} & {\color{black}67}\\\\\\
    \end{bmatrix}
\\]

The first step of Gaussian elimination is to set the sub-diagonal elements in the first column to zero by subtracting multiples of the first row from each of the subsequent rows.

For this example, this requires the row operations

\\[
\begin{align}
Eq. (2) &\leftarrow Eq. (2) - 1\times Eq. (1)\\\\\\
Eq. (3) &\leftarrow Eq. (3) - 4\times Eq. (1)\\\\\\
Eq. (4) &\leftarrow Eq. (4) - 5\times Eq. (1)\\\\\\
\end{align}
\\]

or mathematically, and for each element of the matrix (remembering that we are adding rows together - while one of
the entries of a row will end up being zero, this also has the consequence of updating the rest of the values in that row, hence the iteration over \\(j\\) below):

\\[
\begin{align}
A_{2j} &\leftarrow A_{2j} - \frac{A_{21}}{A_{11}} A_{1j} = A_{2j} - \frac{5}{5} \times A_{1j}, \quad j=1,2,3,4\\\\\\
A_{3j} &\leftarrow A_{3j} - \frac{A_{31}}{A_{11}} A_{1j} = A_{3j} - \frac{20}{5} \times A_{1j}, \quad j=1,2,3,4\\\\\\
A_{4j} &\leftarrow A_{4j} - \frac{A_{41}}{A_{11}} A_{1j} = A_{4j} - \frac{25}{5} \times A_{1j}, \quad j=1,2,3,4\\\\\\
\end{align}
\\]

Notice that we can also write these exact operations on elements in terms of multiplication by a carefully chosen lower-triangular matrix where the non-zero's below the diagonal are restricted to a single column, e.g. for the example above

\\[
  \begin{bmatrix}
    {\color{black}1} & {\color{black}0} & {\color{black}0} & {\color{black}0}\\\\\\
    {\color{Orange}{-1}} & {\color{black}1} & {\color{black}0} & {\color{black}0}\\\\\\
    {\color{Orange}{-4}} & {\color{black}0} & {\color{black}1} & {\color{black}0}\\\\\\
    {\color{Orange}{-5}} & {\color{black}0} & {\color{black}0} & {\color{black}1}\\\\\\
  \end{bmatrix}\qquad\times\qquad\begin{bmatrix}
    {\color{black}5} & {\color{black}7} & {\color{black}5} & {\color{black}9}\\\\\\
    {\color{black}5} & {\color{black}14} & {\color{black}7} & {\color{black}10}\\\\\\
    {\color{black}20} & {\color{black}77} & {\color{black}41} & {\color{black}48}\\\\\\
    {\color{black}25} & {\color{black}91} & {\color{black}55} & {\color{black}67}\\\\\\
  \end{bmatrix}\qquad=\qquad\begin{bmatrix}
    {\color{black}5} & {\color{black}7} & {\color{black}5} & {\color{black}9}\\\\\\
    {\color{blue}{0}} & {\color{blue}{7}} & {\color{blue}{2}} & {\color{blue}{1}}\\\\\\
    {\color{blue}{0}} & {\color{blue}{49}} & {\color{blue}{21}} & {\color{blue}{12}}\\\\\\
    {\color{blue}{0}} & {\color{blue}{56}} & {\color{blue}{30}} & {\color{blue}{22}}\\\\\\
  \end{bmatrix}
\\]

The lower-triangular matrix (let's call this one \\(L_0\\)) is thus encoding the first step in Gaussian elimination.

The next step involves taking the second row of the updated matrix as the new pivot (we will ignore partial pivoting for simplicity), and subtracting multiples of this row from those below in order to set the zeros below the diagonal in the second column to zero.

Note that this step can be achieved here with the multiplication by the following lower-triangular matrix (call this one \\(L_1\\))

\\[
  \begin{bmatrix}
    {\color{black}1} & {\color{black}0} & {\color{black}0} & {\color{black}0}\\\\\\
    {\color{black}0} & {\color{black}1} & {\color{black}0} & {\color{black}0}\\\\\\
    {\color{black}0} & {\color{Orange}{-7}} & {\color{black}1} & {\color{black}0}\\\\\\
    {\color{black}0} & {\color{Orange}{-8}} & {\color{black}0} & {\color{black}1}\\\\\\
  \end{bmatrix}\qquad\times\qquad\begin{bmatrix}
    {\color{black}5} & {\color{black}7} & {\color{black}5} & {\color{black}9}\\\\\\
    {\color{black}0} & {\color{black}7} & {\color{black}2} & {\color{black}1}\\\\\\
    {\color{black}0} & {\color{black}49} & {\color{black}21} & {\color{black}12}\\\\\\
    {\color{black}0} & {\color{black}56} & {\color{black}30} & {\color{black}22}\\\\\\
  \end{bmatrix}\qquad=\qquad\begin{bmatrix}
    {\color{black}5} & {\color{black}7} & {\color{black}5} & {\color{black}9}\\\\\\
    {\color{black}0} & {\color{black}7} & {\color{black}2} & {\color{black}1}\\\\\\
    {\color{black}0} & {\color{blue}{0}} & {\color{blue}{7}} & {\color{blue}{5}}\\\\\\
    {\color{black}0} & {\color{blue}{0}} & {\color{blue}{14}} & {\color{blue}{14}}\\\\\\
  \end{bmatrix}
\\]

and finally for this example to get rid the of the leading 14 on the final row:

\\[
  \begin{bmatrix}
    {\color{black}1} & {\color{black}0} & {\color{black}0} & {\color{black}0}\\\\\\
    {\color{black}0} & {\color{black}1} & {\color{black}0} & {\color{black}0}\\\\\\
    {\color{black}0} & {\color{black}0} & {\color{black}1} & {\color{black}0}\\\\\\
    {\color{black}0} & {\color{black}0} & {\color{Orange}{-2}} & {\color{black}{1}}\\\\\\
  \end{bmatrix}\qquad\times\qquad\begin{bmatrix}
    {\color{black}5} & {\color{black}7} & {\color{black}5} & {\color{black}9}\\\\\\
    {\color{black}0} & {\color{black}7} & {\color{black}2} & {\color{black}1}\\\\\\
    {\color{black}0} & {\color{black}0} & {\color{black}7} & {\color{black}5}\\\\\\
    {\color{black}0} & {\color{black}0} & {\color{black}14} & {\color{black}14}\\\\\\
  \end{bmatrix}\qquad=\qquad\begin{bmatrix}
    {\color{black}5} & {\color{black}7} & {\color{black}5} & {\color{black}9}\\\\\\
    {\color{black}0} & {\color{black}7} & {\color{black}2} & {\color{black}1}\\\\\\
    {\color{black}0} & {\color{black}0} & {\color{black}7} & {\color{black}5}\\\\\\
    {\color{black}0} & {\color{black}0} & {\color{blue}{0}} & {\color{blue}{4}}\\\\\\
  \end{bmatrix}
\\]

where this lower triangular matrix we will call \\(L_2\\), and the RHS matrix is now in upper-triangular form as we expect from Gaussian elimination (call this \\(U\\)).

In summary, the above operations (starting from \\(A\\), first multipling on the left by \\(L_0\\), the multiplying the result on the left by \\(L_1\\) and so on to arrive at an upper trianglualr matrix \\(U\\)) can be written as

\\[ L_2(L_1(L_0A)) = U.\\]

Note that these special lower-triangular matrices \\(L_0, \ldots\\) are all examples of what is known as *atomic* lower-triangular matrices: a special form of unitriangular matrix - the diagonals are unity, where the off-diagonal entries are all zero apart from in a single column.

Notice that for these special, simple matrices, their inverse is simply the original with the sign of those off-diagnonals changed:

$$
  \begin{bmatrix}
    1      & 0      & \cdots      &        &              &        &         &   & 0 \\
    0      & 1      & 0           & \cdots &              &        &         &   & 0 \\
    0      & \ddots & \ddots      & \ddots &              &        &         &   &  \vdots \\
    \vdots & \ddots & \ddots      & \ddots &              &        &         &   &  \\
           &        &             &   0    &   1          &  0     &         & &  &  \\
           &        &             &   0    &   l_{i+1,i}  &  1     &  \ddots &   &  &  \\
           &        &             &   0    &   l_{i+2,i}  &  0     &  \ddots &   & &  \\
           &        &             & \vdots &   \vdots     & \vdots &  \ddots &   & 0 &  \\
    0      & \cdots &             & 0      &   l_{n,i}    & 0      &  \cdots & 0 & 1 &  \\
\end{bmatrix}^{-1}
=
\begin{bmatrix}
    1      & 0      & \cdots      &        &              &        &         &   & 0 \\
    0      & 1      & 0           & \cdots &              &        &         &   & 0 \\
    0      & \ddots & \ddots      & \ddots &              &        &         &   &  \vdots \\
    \vdots & \ddots & \ddots      & \ddots &              &        &         &   &  \\
           &        &             &   0    &   1          &  0     &         & &  &  \\
           &        &             &   0    &   -l_{i+1,i}  &  1     &  \ddots &   &  &  \\
           &        &             &   0    &   -l_{i+2,i}  &  0     &  \ddots &   & &  \\
           &        &             & \vdots &   \vdots     & \vdots &  \ddots &   & 0 &  \\
    0      & \cdots &             & 0      &   -l_{n,i}    & 0      &  \cdots & 0 & 1 &  \\
\end{bmatrix}
$$

### Exercise

Convince yourselves of the following facts:

1. the above statement on what the inverse of the special \\(L\\) matrices is - check on a small example.
2. the multiplication of arbitrary lower-triangular square matrices is also lower-triangular.
3. \\(L_2(L_1(L_0A)) = U \implies A = L_0^{-1}(L_1^{-1}(L_2^{-1}U))\\)
4. and hence that \\(A=LU\\) where \\(U\\) is the upper-triangular matrix found at the end of Guassian elimination, and where \\(L\\) is the following matrix
    \\[L = L_0^{-1}L_1^{-1}L_2^{-1}.\\]
5. finally, compute this product of these lower-triangular matrices to show that
    \\[L =
    \begin{bmatrix}
      {\color{black}1} & {\color{black}0} & {\color{black}0} & {\color{black}0}\\\\\\
      {\color{black}{1}} & {\color{black}1} & {\color{black}0} & {\color{black}0}\\\\\\
      {\color{black}{4}} & {\color{black}7} & {\color{black}1} & {\color{black}0}\\\\\\
      {\color{black}{5}} & {\color{black}8} & {\color{black}2} & {\color{black}1}\\\\\\
    \end{bmatrix}
    \\]
    i.e. that the multiplication of these individual atomic matrices (importantly in this order) simply merges the entries from the non-zero columns of each atomic matrix, and hence is both lower-triangular, as well as trivial to compute.

```python
import numpy as np
from scipy import linalg
from pprint import pprint

# as above: A matrix,
A = np.array([[5., 7. , 5., 9.],
            [5., 14.,   7.,  10.],
            [20., 77., 41., 48.],
            [25., 91.,  55., 67.]])

# lower triangular matrices,
L0 = np.array([[1,0,0,0],[-1,1,0,0],[-4,0,1,0],[-5,0,0,1]])
L1 = np.array([[1,0,0,0],[0,1,0,0],[0,-7,1,0],[0,-8,0,1]])
L2 = np.array([[1,0,0,0],[0,1,0,0],[0,0,1,0],[0,0,-2,1]])

# and their inverse matrices
L0_ = linalg.inv(L0)
L1_ = linalg.inv(L1)
L2_ = linalg.inv(L2)

# our base matrix
pprint(A)

# ad: The multiplication of arbitrary lower-triangular
# square matrices is also lower-triangular
print("Question 1:")
pprint(np.dot(L0, L1))
pprint(np.dot(L2, L1))

# ad: L2(L1(L0 A)) = U
print("Question 2:")
U = np.dot(L2, np.dot(L1, np.dot(L0, A)))
pprint(U)

# ad: A = L0^-1(L1^-1(L2^-1 U))
print("Question 3:")
A = np.dot(L0_, np.dot(L1_, np.dot(L2_, U)))
pprint(A)

# ad: L = L0^-1 L1^-1 L2^-1
print("Question 4:")
L = np.dot(L0_,np.dot(L1_,L2_))
pprint(L)

# ad: A = LU
# numpy.allclose takes two numpy arrays and checks
# whether all elements are identical to some tolerance
print("Question 5:")
print(np.allclose(A, np.dot(L, U)))
```

    array([[ 5.,  7.,  5.,  9.],
           [ 5., 14.,  7., 10.],
           [20., 77., 41., 48.],
           [25., 91., 55., 67.]])
    Question 1:
    array([[ 1,  0,  0,  0],
           [-1,  1,  0,  0],
           [-4, -7,  1,  0],
           [-5, -8,  0,  1]])
    array([[ 1,  0,  0,  0],
           [ 0,  1,  0,  0],
           [ 0, -7,  1,  0],
           [ 0,  6, -2,  1]])
    Question 2:
    array([[5., 7., 5., 9.],
           [0., 7., 2., 1.],
           [0., 0., 7., 5.],
           [0., 0., 0., 4.]])
    Question 3:
    array([[ 5.,  7.,  5.,  9.],
           [ 5., 14.,  7., 10.],
           [20., 77., 41., 48.],
           [25., 91., 55., 67.]])
    Question 4:
    array([[1., 0., 0., 0.],
           [1., 1., 0., 0.],
           [4., 7., 1., 0.],
           [5., 8., 2., 1.]])
    Question 5:
    True

## Implementation
We can build an LU code easily from {ref}`Gaussian elimination code <nm_gaussian_elimination>`. The final \\(U\\) matrix we need here is as was already constructed through Gaussian elimination, the entries of \\(L\\) we need are simply the \\({A_{ik}}/{A_{kk}}\\) multipliers we computed as part of the elimination, but threw away previously.

For a given pivot row \\(k\\), for each of these multipliers (for every row below the pivot), as we compute them we know that we are going to transform the augmented matrix in order to achieve a new zero below the diagonal - we can store each multiplier in this position before moving on to the following row, we implicitly know that the diagonals of \\(L\\) will be unity and so don't need to store these (and noting that we don't actually have a space for them anyway!). We then move on to the following pivot row, replacing the zeros in the corresponding column we are zero'ing, but again using the now spare space to store the multipliers.

For example, for the case above

\\[ A =
  \begin{bmatrix}
    {\color{black}5} & {\color{black}7} & {\color{black}5} & {\color{black}9}\\\\\\
    {\color{black}5} & {\color{black}14} & {\color{black}7} & {\color{black}10}\\\\\\
    {\color{black}20} & {\color{black}77} & {\color{black}41} & {\color{black}48}\\\\\\
    {\color{black}25} & {\color{black}91} & {\color{black}55} & {\color{black}67}\\\\\\
  \end{bmatrix}\quad\rightarrow\quad
  \begin{bmatrix}
    {\color{black}5} & {\color{black}7} & {\color{black}5} & {\color{black}9}\\\\\\
    {\color{blue}{1}} & {\color{black}{7}} & {\color{black}{2}} & {\color{black}{1}}\\\\\\
    {\color{blue}{4}} & {\color{black}{49}} & {\color{black}{21}} & {\color{black}{12}}\\\\\\
    {\color{blue}{5}} & {\color{black}{56}} & {\color{black}{30}} & {\color{black}{22}}\\\\\\
  \end{bmatrix}\quad\rightarrow\quad
  \begin{bmatrix}
    {\color{black}5} & {\color{black}7} & {\color{black}5} & {\color{black}9}\\\\\\
    {\color{blue}1} & {\color{black}7} & {\color{black}2} & {\color{black}1}\\\\\\
    {\color{blue}4} & {\color{blue}{7}} & {\color{black}{7}} & {\color{black}{5}}\\\\\\
    {\color{blue}5} & {\color{blue}{8}} & {\color{black}{14}} & {\color{black}{14}}\\\\\\
  \end{bmatrix}\quad\rightarrow\quad
  \begin{bmatrix}
    {\color{black}5} & {\color{black}7} & {\color{black}5} & {\color{black}9}\\\\\\
    {\color{blue}1} & {\color{black}7} & {\color{black}2} & {\color{black}1}\\\\\\
    {\color{blue}4} & {\color{blue}7} & {\color{black}7} & {\color{black}5}\\\\\\
    {\color{blue}5} & {\color{blue}8} & {\color{blue}{2}} & {\color{black}{4}}\\\\\\
  \end{bmatrix}
  = [\color{blue}L\backslash U]
\\]

Starting from {ref}`Gaussian elimination code <nm_gaussian_elimination>` we will compute the LU decomposition of a matrix. First, we will store \\(L\\) and \\(U\\) in two different matrices:

```python
def LU_dec(A):
    # Upper triangular matrix contains gaussian elimination result
    # We won't change A in-place but create a local copy
    A = A.copy()
    m, n = A.shape
    # Lower triangular matrix has identity diagonal and scaling factors
    # Start from the identity:
    L = np.identity(n)
    # Loop over each pivot row.
    for k in range(n):
        # Loop over each equation below the pivot row.
        for i in range(k+1, n):
            # Define the scaling factor outside the innermost
            # loop otherwise its value gets changed.
            s = (A[i,k]/A[k,k])
            for j in range(k, n):
                A[i,j] = A[i,j] - s*A[k,j]
            # Scaling factors make up the lower matrix
            L[i,k] = s
    # A now is the upper triangular matrix U
    return L, A
```

Test the code:

```python
A = np.array([[ 5., 7.,   5.,  9.],
              [ 5., 14.,  7., 10.],
              [20., 77., 41., 48.],
              [25., 91. ,55., 67.]])

L, U = LU_dec(A)

pprint(A)
pprint(L)
pprint(U)
print(np.allclose(np.dot(L,U), A))
```

    array([[ 5.,  7.,  5.,  9.],
           [ 5., 14.,  7., 10.],
           [20., 77., 41., 48.],
           [25., 91., 55., 67.]])
    array([[1., 0., 0., 0.],
           [1., 1., 0., 0.],
           [4., 7., 1., 0.],
           [5., 8., 2., 1.]])
    array([[5., 7., 5., 9.],
           [0., 7., 2., 1.],
           [0., 0., 7., 5.],
           [0., 0., 0., 4.]])
    True

## Partial pivoting

Problems can occur in a numerical implementation of our algorithm if $A_{kk}$ divided through by in the Gaussian elimination and/or back substitution algorithms might be zero (or close to zero).

Using Gaussian elimination as an example, let's again consider the algorithm mid-way working on an arbitrary matrix system, i.e. assume that the first \\(k\\) rows have already been transformed into upper-triangular form, while the equations/rows below are not yet in this form:

\\[
\begin{bmatrix}
    A_{11} & A_{12} & A_{13} & \cdots & A_{1k}  & \cdots & A_{1n} & b_1 \\\\\\
    0      & A_{22} & A_{23} & \cdots & A_{2k} & \cdots & A_{2n} & b_2 \\\\\\
    0      & 0      & A_{33} & \cdots & A_{3k}  & \cdots & A_{3n} & b_3 \\\\\\
    \vdots & \vdots & \vdots & \ddots & \vdots  & \ddots & \vdots & \vdots \\\\\\
\hdashline
    0      & 0      & 0      & \cdots & A_{kk}  & \cdots & A_{kn} & b_k \\\\\\
    \vdots & \vdots & \vdots & \ddots & \vdots  & \ddots & \vdots & \vdots \\\\\\
    0      & 0      & 0      & \cdots & A_{nk}  & \cdots & A_{nn} & b_n \\\\\\
\end{bmatrix}
\\]

We have drawn the horizontal dashed line one row higher, as we are not going to blindly asssume that it is wise to take the current row \\(k\\) as the pivot row, and \\(A_{kk}\\) as the so-called pivot element.

**Partial pivoting** selects the best pivot (row or element) as the one where the \\(A_{ik}\\) (for \\(i\ge k\\)) value is largest (relative to the other values of components in its own row \\(i\\)), and then swaps this row with the current \\(k\\) row.

To generalise our codes above we would simply need to search for this row, and perform the row swap operation.

Python's `scipy.linalg` library has its own implementation of the LU decomposition, that uses partial pivoting.

```{admonition} Example of partial pivoting
:class: dropdown

Suppose that we have the matrix \\(\pmb{A}\\) from before

\\[
A =
  \begin{bmatrix}
    5 & 7 & 5 &  9 \\\\\\
    5 & 14 & 7 & 10 \\\\\\
    20 & 77 &41 & 48 \\\\\\
    25 & 91 & 55 & 67 \\\\\\
  \end{bmatrix}
\\]

and then we do some timeskipping and skip to partway through

\\[
A =
  \begin{bmatrix}
    5 & 7 & 5 &  9 \\\\\\
    0 & 7 & 2 & 1 \\\\\\
    0 & 49 &21 & 12 \\\\\\
    0 & 56 & 30 & 22 \\\\\\
  \end{bmatrix}.
\\]

Now, magic happens for the sake of illustrating the partial pivoting - "magic magic poof poof and magically \\(7\\) becomes \\(0\\)":

\\[
A =
  \begin{bmatrix}
    5 & 7 & 5 &  9 \\\\\\
    0 & 0 & 2 & 1 \\\\\\
    0 & 49 &21 & 12 \\\\\\
    0 & 56 & 30 & 22 \\\\\\
\end{bmatrix}
\\]

Suppose that the matrix storing our unknowns is

\\[
\pmb{x} =
\begin{bmatrix}
  x \\\\\\
  y \\\\\\
  z \\\\\\
  \text{letter after }z
\end{bmatrix}
\\]

And the matrix storing the results of our equations, after all of the Gaussian elimination previously that have modified the matrix \\(\pmb{b}\\)

\\[
\pmb{b} =
\begin{bmatrix}
  b_1 \\\\\\
  b_2 \\\\\\
  b_3 \\\\\\
  b_4
\end{bmatrix}
\\]

Writing this in equation this would be

\\[
5x + 7y + 5z + 9(\text{letter after }z) = b_1 \\\\\\
0x + 0y + 2z + 1(\text{letter after }z) = b_2 \\\\\\
0x + 49y + 21z + 12(\text{letter after }z) = b_3  \\\\\\
0x + 56y + 30z + 22(\text{letter after }z) = b_4
\\]

Now, we are going to use equation 2 and equation 3

\\[
A =
  \begin{bmatrix}
    5 & 7 & 5 &  9 \\\\\\
    0 & 0 & 2 & 1 \\\\\\
    0 & 49 &21 & 12 \\\\\\
    0 & 56 & 30 & 22 \\\\\\
  \end{bmatrix}
\\]

The coefficient of the second unknown in the second equation is 0. The coefficient of the second unknown in the third equation is 49. Thus, I will divide the coefficient of the 3rd equation with the coefficient of the 2nd equation, and then I would multiply the coefficient of the 3rd equation by the quotient. As such, I will divide \\(49/0\\), and then, I will get an error since you cannot divide by \\(0\\). It does not neccesarily have to be \\(0\\) to give errors, even if it is very close to \\(0\\), it will still give an error.

Thus what can I do?

Well, I need two equations. I cannot use the equation 2 because there is a \\(0\\). Then, I will use equation 3 and equation 4. Therefore, let's swap the position of equation 4 and equation 2.

\\[
A =
  \begin{bmatrix}
    5 & 7 & 5 &  9 \\\\\\
    0 & 56 & 30 & 22 \\\\\\
    0 & 49 &21 & 12 \\\\\\
    0 & 0 & 2 & 1 \\\\\\
  \end{bmatrix}
\\]

Remember when you swap rows in matrix \\(\pmb{A}\\), you must also swap rows in the matrix \\(\pmb{b}\\). Matrix \\(\pmb{x}\\) does not have to swapped since the order of the variables has not changed.

The matrix storing our unknowns remains

\\[
\pmb{x} =
\begin{bmatrix}
  x \\\\\\
  y \\\\\\
  z \\\\\\
  \text{letter after }z
\end{bmatrix}
\\]

And the matrix storing the results of our equations

\\[
\pmb{b} =
\begin{bmatrix}
  b_1 \\\\\\
  b_4 \\\\\\
  b_3 \\\\\\
  b_2
\end{bmatrix}
\\]

Writing this in equation this would be

\\[
5x + 7y + 5z + 9(\text{letter after }z) = b_1 \\\\\\
0x + 56y + 30z + 22(\text{letter after }z) = b_4 \\\\\\
0x + 49y + 21z + 12(\text{letter after }z) = b_3  \\\\\\
0x + 0y + 2z + 1(\text{letter after }z) = b_2
\\]

Inspecting this equation, we see that the equation is essentially the same as it was before the swap, so we can say that our action to swap it is a valid action. I have rewritten the equation before the swap as reference

\\[
5x + 7y + 5z + 9(\text{letter after }z) = b_1 \\\\\\
0x + 0y + 2z + 1(\text{letter after }z) = b_2 \\\\\\
0x + 49y + 21z + 12(\text{letter after }z) = b_3  \\\\\\
0x + 56y + 30z + 22(\text{letter after }z) = b_4
\\]

Now let's continue, after we finished swapping, and we will focus only on getting to the upper triangle form, and only on matrix \\(\pmb{A}\\). The matrix \\(\pmb{b}\\) will also change accordingly as detailed in the previous lecture, but let's focus on matrix \\(\pmb{A}\\):

\\[
A =
  \begin{bmatrix}
    5 & 7 & 5 &  9 \\\\\\
    0 & 56 & 30 & 22 \\\\\\
    0 & 56 & 24 & 96/7 \\\\\\
    0 & 0 & 2 & 1 \\\\\\
  \end{bmatrix}
\\]

\\[
A =
  \begin{bmatrix}
    5 & 7 & 5 &  9 \\\\\\
    0 & 56 & 30 & 22 \\\\\\
    0 & 0 & -6 & 96/7 -22 \\\\\\
    0 & 0 & 2 & 1 \\\\\\
  \end{bmatrix}
\\]

We will not encounter the problem with dividing by zero anymore, at least for this matrix.

This matrix was nice and only required one swap. But if your matrix has lots of \\(0\\) or very small numbers, then maybe you might need more swaps. You might also want to nondimensionalize that matrix. That is, of course assuming that the equations contain enough information to actually allow you to solve the equation. You could use the determinant, or the {ref}`Gauss-Jordan elimination <nm_gauss_jordan_elimination>` to check if there is actually a solution. You should always check the solvability of an equation before solving it, since trying to solve an equation which cannot be solved is wasted effort.

You could use conditionals to control the swapping or the skipping of rows. Although the example here used \\(0\\), it could be a very small number too, and not neccesary \\(0\\) that could cause the problem, so maybe it would be nice to compare the conditional to a certain threshold to decided whether to swap or not.

```

Example partial pivoting using `scipy.linalg`:

```python
import scipy.linalg as sl

A = np.array([[ 5., 7.,   5.,  9.],
               [ 5., 14.,  7., 10.],
               [20., 77., 41., 48.],
               [25., 91. ,55., 67.]])
print("A =", A)

P, L, U = sl.lu(A)

# P here is a 'permutation matrix' that performs
# swaps based upon partial pivoting
print("P =", P)

# The lower-triangular matrix
print(L)

# The upper-triangular matrix
print(U)

# Double check that P*L*U does indeed equal A
print("Does P*L*U = A?", np.allclose(P @ L @ U, A))
```

    A = [[ 5.  7.  5.  9.]
     [ 5. 14.  7. 10.]
     [20. 77. 41. 48.]
     [25. 91. 55. 67.]]
    P = [[0. 1. 0. 0.]
     [0. 0. 0. 1.]
     [0. 0. 1. 0.]
     [1. 0. 0. 0.]]
    [[ 1.          0.          0.          0.        ]
     [ 0.2         1.          0.          0.        ]
     [ 0.8        -0.375       1.          0.        ]
     [ 0.2         0.375       0.33333333  1.        ]]
    [[ 25.          91.          55.          67.        ]
     [  0.         -11.2         -6.          -4.4       ]
     [  0.           0.          -5.25        -7.25      ]
     [  0.           0.           0.           0.66666667]]
    Does P*L*U = A? True

Looking at the form of \\(P\\) above, we can re-order the rows in advance and consider the LU decomposition of the matrix where \\(P=I\\), as below. As we haven't bothered implementing pivoting ourselves, check that your LU implementation recreates the \\(A\\), \\(L\\) and \\(U\\) below.

```python
A = np.array([[25. ,91. ,55. ,67.],
               [ 5.,  7.,  5.,  9.],
               [20., 77., 41., 48.],
               [ 5., 14., 7.,  10.]])
print("A =", A)

P, L, U = sl.lu(A)

# P here is a 'permutation matrix' that performs
# swaps based upon partial pivoting
print("P =", P)

# The lower-triangular matrix
print(L)

# The upper-triangular matrix
print(U)

# Double check that P*L*U does indeed equal A
print("Does P*L*U = A?", np.allclose(P @ L @ U, A))
```

    A = [[25. 91. 55. 67.]
     [ 5.  7.  5.  9.]
     [20. 77. 41. 48.]
     [ 5. 14.  7. 10.]]
    P = [[1. 0. 0. 0.]
     [0. 1. 0. 0.]
     [0. 0. 1. 0.]
     [0. 0. 0. 1.]]
    [[ 1.          0.          0.          0.        ]
     [ 0.2         1.          0.          0.        ]
     [ 0.8        -0.375       1.          0.        ]
     [ 0.2         0.375       0.33333333  1.        ]]
    [[ 25.          91.          55.          67.        ]
     [  0.         -11.2         -6.          -4.4       ]
     [  0.           0.          -5.25        -7.25      ]
     [  0.           0.           0.           0.66666667]]
    Does P*L*U = A? True

### Implementation

See below the code implementation to partial pivoting:

```python
# This function is not necessary but makes the LU_dec_pp function less cluttered
# Making general operations into functions is good way of improving readability
# and reducing your workload later!
def swap_rows(A,j,k):
    B = np.copy(A[j,:])
    C = np.copy(A[k,:])
    # second label, BUT NOT CREATING A 2ND ARRAY.
    # Try it for yourself if you want.
    A[k,:] = B
    A[j,:] = C
    return A

# A function to perform LU decomposition with partial pivoting
def LU_dec_pp(A):
    m, n = A.shape
    A = A.copy() # we won't modify in place but create local copy
    P_ = np.identity(m)
    L  = np.identity(m)
    for k in range(m):
        j = np.argmax(abs(A[k:,k]))
        # Find the index of the largest ABSOLUTE value. np.argmax will return
        # the index of the largest element in an array
        j+= k               # A[1+2,2] = A[3,2] = 3!
        A  = swap_rows(A,j,k)
        P_ = swap_rows(P_,j,k)
        for j in range(k+1,m):
            s = A[j,k]/A[k,k]
            A[j,k:] -= A[k,k:]*s
            L[j,k] = s
    U = A
    return P_.T, L, U

A = np.array([[ 5., 7.,   5.,  9.],
              [ 5., 14.,  7., 10.],
              [20., 77., 41., 48.],
              [25., 91. ,55., 67.]])

P, L, U = LU_dec_pp(A)
pprint(A)
pprint(L)
pprint(U)
pprint(np.dot(P, np.dot(L,U)))
print("Have we succeeded in creating partial pivoting?", np.allclose(np.dot(P, np.dot(L,U)), A))
```

    array([[ 5.,  7.,  5.,  9.],
           [ 5., 14.,  7., 10.],
           [20., 77., 41., 48.],
           [25., 91., 55., 67.]])
    array([[ 1.        ,  0.        ,  0.        ,  0.        ],
           [ 0.2       ,  1.        ,  0.        ,  0.        ],
           [ 0.8       , -0.375     ,  1.        ,  0.        ],
           [ 0.2       ,  0.375     ,  0.33333333,  1.        ]])
    array([[ 25.        ,  91.        ,  55.        ,  67.        ],
           [  0.        , -11.2       ,  -6.        ,  -4.4       ],
           [  0.        ,   0.        ,  -5.25      ,  -7.25      ],
           [  0.        ,   0.        ,   0.        ,   0.66666667]])
    array([[ 5.,  7.,  5.,  9.],
           [ 5., 14.,  7., 10.],
           [20., 77., 41., 48.],
           [25., 91., 55., 67.]])
    Have we succeeded in creating partial pivoting? True


## Online system identification in Duffing oscillator by free energy minimisation

This project considers a [Duffing oscillator](hkps://en.wikipedia.org/wiki/Duffing_equation), a driven damped harmonic oscillator with a cubic nonlinearity in its spring stiffness component. State-space model description of the system:

$$\begin{align}
m \frac{d^2 x(t)}{dt^2} + c \frac{d x(t)}{dt} + a x(t) + b x^3(t) =&\ u(t) + w(t) \\
y(t) =&\ x(t) + e(t)
\end{align}$$

where
$$\begin{align}
m     =&\ \text{mass} \\
c     =&\ \text{damping} \\
a     =&\ \text{linear stiffness} \\
b     =&\ \text{nonlinear stiffness} \\
y(t)    =&\ \text{observation (displacement)} \\
x(t)    =&\ \text{state (displacement)} \\
u(t)    =&\ \text{force} \\
v(t)    =&\ \text{measurement noise} \\
w(t)    =&\ \text{process noise}
\end{align}$$

The process noise is a Wiener process, where the increment is Gaussian distributed $w(t) \sim \mathcal{N}(0, \tau^{-1}dt)$. The parameter $\tau$ represents the precision of the process. The measurement noise is also a Wiener process, $v(t) \sim \mathcal{N}(0, \xi^{-1}dt)$.

## Experiment: simulation error

In this notebook, we will perform a simulation error experiment: the model will simulate future outputs for a large time horizon using only inputs and inferred parameters. The predictions will be evaluated and compared to a few benchmark methods.

### Data

There is an electronic implementation of the Duffing oscillator on the Nonlinear System Identification Benchmark website: http://nonlinearbenchmark.org/#Silverbox. It's called the Silverbox setup.

```julia
using Pkg
Pkg.activate(".")
Pkg.instantiate()
```

     Activating environment at `~/syndr/Wouter/Onderzoek/Projecten/tueindhoven/control-nsi/IWAI2020-onlinesysid/Project.toml`

```julia
using CSV
using DataFrames
```

```julia
using Plots
pyplot();
viz = true;
```

```julia
# Read data from CSV file
df = CSV.read("data/SNLS80mV.csv", ignoreemptylines=true)
df = select(df, [:V1, :V2])

# Sampling frequency
fs = 610.35

# Shorthand
input = df[:,1]
output = df[:,2]

# Time horizon
T = size(df, 1);
```

    ┌ Warning: `CSV.read(input; kw...)` is deprecated in favor of `using DataFrames; CSV.read(input, DataFrame; kw...)
    │   caller = ip:0x0
    └ @ Core :-1

```julia
# Zoomed in versions of signals
tt = 50000:1:50500

p01 = plot(tt, input[tt], color="red", label="", markersize=2, xlabel="", ylabel="input", xticks=:none, ylims=[-0.22, 0.22], tickfontsize=14, legendfontsize=12, guidefontsize=16)
p02 = plot(tt, output[tt], color="blue", label="", markersize=2, xlabel="time (t)", ylabel="output", ylims=[-0.22, 0.22], tickfontsize=14, legendfontsize=12, guidefontsize=16)
p00 = plot(p01, p02, layout=(2,1), size=(800,400))
savefig(p00, "figures/signals_zoomed.png")
savefig(p00, "figures/signals_zoomed.pdf")
```

```julia
# Select training set
ix_trn = collect(40101:131072)
input_trn = input[ix_trn]
output_trn = output[ix_trn]
T_trn = length(ix_trn);

# Select validation set
ix_val = 1:40100
input_val = input[ix_val]
output_val = output[ix_val]
T_val = length(ix_val);
```

```julia
# Plot entire series with training split
ss = 80
ix = 1:ss:T
p11 = plot(ix, output[ix], color="blue", linewidth=2, xlabel="time (t)", ylabel="output", label="", ylim=[-.24, .24], tickfontsize=14, legendfontsize=12, guidefontsize=16)
vline!([40100], color="black", linewidth=4, label="")
p12 = plot(ix, input[ix], color="red", linewidth=2, xticks=:none, xlabel="", ylabel="input", label="", ylim=[-.24, .24], tickfontsize=14, legendfontsize=12, guidefontsize=16)
vline!([40100], color="black", linewidth=4, label="")
p10 = plot(p12, p11, layout=(2,1), size=(1200,400))
```

```julia
savefig(p10, "figures/dataset_split.png")
```

```julia
# Plot example of signals
ss = 4
ix = 126000:ss:127500
p31 = Plots.plot(ix, output[ix], color="blue", linewidth=3, xlabel="time (t)", label="output")
Plots.plot!(ix, input[ix], color="red", linewidth=3, xlabel="time (t)", label="input", size=(1200,300), ylim=[-.16, .21], legend=:topright, tickfontsize=14, legendfontsize=12, ylabel="signal", guidefontsize=16)
```

```julia
savefig(p31, "figures/input-output_seq1.png")
savefig(p31, "figures/input-output_seq1.pdf")
```

## Solution steps

### 1. Discretize

I'm using a central difference for the second derivative and a forward difference for the first derivative. Let $w_t$ be a sample from $\mathcal{N}(0, \tau^{-1})$. The state transition can now be wriken as the following discrete-time system:

$$\begin{align}
m (x_{t+1} - 2x_{t} + x_{t-1}) + c (x_{t+1} - x_{t}) + a x_t + b x_t^3 =&\ u_t + w_t
\end{align}$$
Re-writing this as a function of $x_{t+1}$ yields:
$$\begin{align}
% (m + c) x_{t+1}&\ + (-2m - c + a) x_{t} + bx_t^3 + m x_{t-1} = u_t + w_t \\
x_{t+1}&\ = \frac{2m + c - a}{m + c} x_{t} + \frac{-b}{m + c}x_t^3 + \frac{-m}{m + c} x_{t-1} + \frac{1}{m + c} u_t + \frac{1}{m + c} w_t \, .
\end{align}$$

### 2. Substitute variables and reduce order

I can cast the above system into matrix form:

$$ \underbrace{\begin{bmatrix} x_{t+1} \\ x_{t} \end{bmatrix}}_{z_t} = \underbrace{\begin{bmatrix} 0 & 0 \\ 1 & 0 \end{bmatrix}}_{S} \underbrace{\begin{bmatrix} x_{t} \\ x_{t-1} \end{bmatrix}}_{z_{t-1}} + \underbrace{\begin{bmatrix} 1 \\ 0 \end{bmatrix}}_{s} g(\theta, z_{t-1}) + \begin{bmatrix} 1 \\ 0 \end{bmatrix} \eta u_t + \begin{bmatrix} 1 \\ 0 \end{bmatrix} \tilde{w}_t \, ,$$

where

$$\begin{align}
\theta_1 = \frac{2m+c-a}{m+c} \ , \quad
\theta_2 = \frac{-b}{m+c} \ , \quad
\theta_3 = \frac{-m}{m+c} \ , \quad
\eta = \frac{1}{m+c} \ , \quad
\gamma^{-1} = \frac{\tau^{-1}}{(m+c)^2} \, ,
\end{align}$$

with $g(\theta, z_{t-1}) = \theta_1 x_t + \theta_2 x_t^3 + \theta_3 x_{t-1}$ and $\tilde{w}_t \sim \mathcal{N}(0, \gamma^{-1})$. In total, I have five unknowns $m,c,a,b,\tau$ and five equations. I can invert the mapping between $\phi = (m, c, a, b, \tau)$ and $\psi = (\theta_1, \theta_2, \theta_3, \eta, \gamma)$ to recover MAP estimates for the physical parameters. An additional advantage of variable substitution is that it allows for more freedom in choosing priors.

The system is now a nonlinear autoregressive process:

$$z_t = f(\theta, z_{t-1}, \eta, u_t) + \tilde{w}_t$$

where $f(\theta, z_{t-1}, \eta, u_t) = Sz_{t-1} + s g(\theta, z_{t-1}) + s \eta u_t$. Note that the states are two-dimensional now.

### 3. Convert to Gaussian probability

Integrating out $\tilde{w}_t$ and $v_t$ produces a Gaussian state transition node:

$$\begin{align}
z_t \sim&\ \mathcal{N}(f(\theta, z_{t-1}, \eta, u_t), V) \\
y_t \sim&\ \mathcal{N}(s^{\top} z_t, \xi^{-1}) \, ,
\end{align}$$

where $V = \begin{bmatrix} \gamma^{-1} & 0 \\ 0 & \epsilon \end{bmatrix}$ and $W = V^{-1} = \begin{bmatrix} \gamma & 0 \\ 0 & \epsilon^{-1} \end{bmatrix}$.

### 4. Approximating the nonlinearity

The nonlinearity is approximated using a first-order Taylor expansion. The work here revolves around working out the expectations for $g(x,\theta)$:

$$ g(\theta, x) = g(m_{\theta}, m_x) + J_{x}(m_{\theta}, m_x)^{\top}(x - m_x) + J_{\theta}(m_{\theta}, m_x)^{\top}(\theta - m_{\theta}) \, ,$$

where $J_x$ denotes the partial derivative of $g$ with respect to $x$ and $J_{\theta}$ w.r.t. $\theta$. Note that our current $g$ is linear in $\theta$ and one could argue that the approximation is unnecessary. However, this form is more general and the first-order T_valaylor is exact anyway.

### 5. Choose priors

We know that mass $m$ and process precision $\gamma$ are strictly positive parameters and that the damping and stiffness coefficients can be both positive and negative. By examing the nonlinear transform $\psi = G(\phi)$, we realize that $\theta_1$, $\theta_2$, $\theta_3$ and $\eta$ can be both positive and negative, but $\gamma$ can only be positive. As such, we choose the following priors:

$$\begin{align}
\theta \sim \text{Normal}(m^{0}_{\theta}, V^{0}_{\theta}) \ , \quad
\eta \sim \text{Normal}(m^{0}_{\eta}, v^{0}_{\eta}) \ , \quad
\gamma \sim \text{Gamma}(a^{0}_\gamma, b^{0}_\gamma) \, .
\end{align}$$

### 6. Choose recognition model

We do not introduce any independencies; the recognition model follows the generative model:

$$\begin{align}
q(\theta) \sim \text{Normal}(m_{\theta}, V_{\theta}) \ , \quad
q(\eta) \sim \text{Normal}(m_{\eta}, v_{\eta}) \ , \quad
q(\gamma) \sim \text{Gamma}(a_\gamma, b_\gamma) \, .
\end{align}$$

## Implementation

The procedure described above was implemented using [ForneyLab.jl](hkps://github.com/biaslab/ForneyLab.jl) with a custom node called "NLARX". It contains a Nonlinear Latent Autoregressive model with eXogenous input to model the state transition.

```julia
using LinearAlgebra
using ForneyLab
using ForneyLab: unsafeMean, unsafeCov, unsafeVar, unsafePrecision
using ProgressMeter

include("NLARX-node/NLARX.jl")
include("NLARX-node/util.jl")
using .NLARX
```

```julia
# System identification graph
graph1 = FactorGraph()

# Static parameters
@RV θ ~ GaussianMeanPrecision(placeholder(:m_θ, dims=(3,)), placeholder(:w_θ, dims=(3,3)))
@RV η ~ GaussianMeanPrecision(placeholder(:m_η), placeholder(:w_η))
@RV γ ~ Gamma(placeholder(:a_γ), placeholder(:b_γ))
@RV ξ ~ Gamma(placeholder(:a_ξ), placeholder(:b_ξ))

# Nonlinearity
g(θ, x) = θ[1]*x[1] + θ[2]*x[1]^3 + θ[3]*x[2]

# State prior
@RV z_tmin1 ~ GaussianMeanPrecision(placeholder(:m_z, dims=(2,)), placeholder(:w_z, dims=(2, 2)), id=:z_tmin1)

# Autoregressive node
@RV z_t ~ NLatentAutoregressiveX(θ, z_tmin1, η, placeholder(:u_t), γ, g=g, id=:z_t)

# Specify likelihood
@RV y_t ~ GaussianMeanPrecision(dot([1. , 0.], z_t), ξ, id=:y_t)

# Placeholder for observation
placeholder(y_t, :y_t)

# Draw time-slice subgraph
ForneyLab.draw(graph1)
```

<?xml version="1.0" encoding="UTF-8" standalone="no"?>
<!DOCTYPE svg PUBLIC "-//W3C//DTD SVG 1.1//EN"
 "http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd">
<!-- Generated by graphviz version 2.40.1 (20161225.0304)
 -->
<!-- Title: G Pages: 1 -->
<svg width="1432pt" height="629pt"
 viewBox="0.00 0.00 1432.00 629.00" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink">
<g id="graph0" class="graph" transform="scale(1 1) rotate(0) translate(4 625)">
<title>G</title>
<polygon fill="#ffffff" stroke="transparent" points="-4,4 -4,-625 1428,-625 1428,4 -4,4"/>
<!-- 2867327560033567407 -->
<g id="node1" class="node">
<title>2867327560033567407</title>
<polygon fill="#d3d3d3" stroke="#000000" points="1232,-54 1155,-54 1155,0 1232,0 1232,-54"/>
<text text-anchor="middle" x="1193.5" y="-24.8" font-family="Times,serif" font-size="9.00" fill="#000000">placeholder_m_η</text>
</g>
<!-- 2127061049655424227 -->
<g id="node2" class="node">
<title>2127061049655424227</title>
<polygon fill="#d3d3d3" stroke="#000000" points="292,-279 217,-279 217,-225 292,-225 292,-279"/>
<text text-anchor="middle" x="254.5" y="-249.8" font-family="Times,serif" font-size="9.00" fill="#000000">placeholder_b_ξ</text>
</g>
<!-- 3361969956495173598 -->
<g id="node3" class="node">
<title>3361969956495173598</title>
<polygon fill="#d3d3d3" stroke="#000000" points="77,-54 0,-54 0,0 77,0 77,-54"/>
<text text-anchor="middle" x="38.5" y="-24.8" font-family="Times,serif" font-size="9.00" fill="#000000">placeholder_w_z</text>
</g>
<!-- 7590602929731448283 -->
<g id="node4" class="node">
<title>7590602929731448283</title>
<polygon fill="#d3d3d3" stroke="#000000" points="572,-621 499,-621 499,-567 572,-567 572,-621"/>
<text text-anchor="middle" x="535.5" y="-591.8" font-family="Times,serif" font-size="9.00" fill="#000000">placeholder_y_t</text>
</g>
<!-- 6335428470881007641 -->
<g id="node8" class="node">
<title>6335428470881007641</title>
<polygon fill="none" stroke="#000000" points="571.5,-522 499.5,-522 499.5,-450 571.5,-450 571.5,-522"/>
<text text-anchor="middle" x="535.5" y="-488.8" font-family="Times,serif" font-size="9.00" fill="#000000">𝒩</text>
<text text-anchor="middle" x="535.5" y="-478.8" font-family="Times,serif" font-size="9.00" fill="#000000">y_t</text>
</g>
<!-- 7590602929731448283--6335428470881007641 -->
<g id="edge10" class="edge">
<title>7590602929731448283--6335428470881007641</title>
<path fill="none" stroke="#000000" d="M535.5,-566.747C535.5,-553.2495 535.5,-536.7693 535.5,-522.2253"/>
<text text-anchor="start" x="535.5" y="-542.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">y_t</text>
<text text-anchor="start" x="516.5" y="-524.8253" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="516.5" y="-560.347" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
</g>
<!-- 18176827165368624414 -->
<g id="node5" class="node">
<title>18176827165368624414</title>
<polygon fill="#d3d3d3" stroke="#000000" points="654.5,-279 600.5,-279 600.5,-225 654.5,-225 654.5,-279"/>
<text text-anchor="middle" x="627.5" y="-249.8" font-family="Times,serif" font-size="9.00" fill="#000000">clamp_1</text>
</g>
<!-- 12774672834151751891 -->
<g id="node6" class="node">
<title>12774672834151751891</title>
<polygon fill="#d3d3d3" stroke="#000000" points="1424,-54 1347,-54 1347,0 1424,0 1424,-54"/>
<text text-anchor="middle" x="1385.5" y="-24.8" font-family="Times,serif" font-size="9.00" fill="#000000">placeholder_w_η</text>
</g>
<!-- 5547258369390740305 -->
<g id="node7" class="node">
<title>5547258369390740305</title>
<polygon fill="#d3d3d3" stroke="#000000" points="843,-54 768,-54 768,0 843,0 843,-54"/>
<text text-anchor="middle" x="805.5" y="-24.8" font-family="Times,serif" font-size="9.00" fill="#000000">placeholder_a_γ</text>
</g>
<!-- 7883531339165132651 -->
<g id="node10" class="node">
<title>7883531339165132651</title>
<polygon fill="none" stroke="#000000" points="478.5,-405 406.5,-405 406.5,-333 478.5,-333 478.5,-405"/>
<text text-anchor="middle" x="442.5" y="-371.8" font-family="Times,serif" font-size="9.00" fill="#000000">Gam</text>
<text text-anchor="middle" x="442.5" y="-361.8" font-family="Times,serif" font-size="9.00" fill="#000000">gamma_2</text>
</g>
<!-- 6335428470881007641--7883531339165132651 -->
<g id="edge1" class="edge">
<title>6335428470881007641--7883531339165132651</title>
<path fill="none" stroke="#000000" d="M499.3308,-453.9578C492.3847,-447.0402 485.4345,-439.5323 479.5,-432 472.9845,-423.7303 466.8046,-414.2006 461.4275,-405.0972"/>
<text text-anchor="start" x="479.5" y="-425.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">ξ</text>
<text text-anchor="start" x="442.4275" y="-407.6972" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="484.3308" y="-456.5578" font-family="Times,serif" font-size="8.00" fill="#000000">3 w </text>
</g>
<!-- 14810396312109737190 -->
<g id="node14" class="node">
<title>14810396312109737190</title>
<polygon fill="none" stroke="#000000" points="665.5,-405 593.5,-405 593.5,-333 665.5,-333 665.5,-405"/>
<text text-anchor="middle" x="629.5" y="-371.8" font-family="Times,serif" font-size="9.00" fill="#000000">dot</text>
<text text-anchor="middle" x="629.5" y="-361.8" font-family="Times,serif" font-size="9.00" fill="#000000">dotproduct_1</text>
</g>
<!-- 6335428470881007641--14810396312109737190 -->
<g id="edge13" class="edge">
<title>6335428470881007641--14810396312109737190</title>
<path fill="none" stroke="#000000" d="M564.4701,-449.9415C575.9063,-435.7071 589.0211,-419.3833 600.4637,-405.1409"/>
<text text-anchor="start" x="585.5" y="-425.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">variable_1</text>
<text text-anchor="start" x="581.4637" y="-407.7409" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="549.4701" y="-443.5415" font-family="Times,serif" font-size="8.00" fill="#000000">2 m </text>
</g>
<!-- 4485287826409161871 -->
<g id="node9" class="node">
<title>4485287826409161871</title>
<polygon fill="none" stroke="#000000" points="266.5,-171 194.5,-171 194.5,-99 266.5,-99 266.5,-171"/>
<text text-anchor="middle" x="230.5" y="-137.8" font-family="Times,serif" font-size="9.00" fill="#000000">𝒩</text>
<text text-anchor="middle" x="230.5" y="-127.8" font-family="Times,serif" font-size="9.00" fill="#000000">z_tmin1</text>
</g>
<!-- 4485287826409161871--3361969956495173598 -->
<g id="edge19" class="edge">
<title>4485287826409161871--3361969956495173598</title>
<path fill="none" stroke="#000000" d="M194.3863,-114.686C160.9287,-95.8662 111.3131,-67.9574 77.0871,-48.7052"/>
<text text-anchor="start" x="134.5" y="-74.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">w_z</text>
<text text-anchor="start" x="77.0871" y="-51.3052" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="179.3863" y="-117.286" font-family="Times,serif" font-size="8.00" fill="#000000">3 w </text>
</g>
<!-- 11222458692561156744 -->
<g id="node20" class="node">
<title>11222458692561156744</title>
<polygon fill="#d3d3d3" stroke="#000000" points="269,-54 192,-54 192,0 269,0 269,-54"/>
<text text-anchor="middle" x="230.5" y="-24.8" font-family="Times,serif" font-size="9.00" fill="#000000">placeholder_m_z</text>
</g>
<!-- 4485287826409161871--11222458692561156744 -->
<g id="edge6" class="edge">
<title>4485287826409161871--11222458692561156744</title>
<path fill="none" stroke="#000000" d="M230.5,-98.7371C230.5,-84.1526 230.5,-67.6262 230.5,-54.1146"/>
<text text-anchor="start" x="230.5" y="-74.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">m_z</text>
<text text-anchor="start" x="211.5" y="-56.7146" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="215.5" y="-92.3371" font-family="Times,serif" font-size="8.00" fill="#000000">2 m </text>
</g>
<!-- 7883531339165132651--2127061049655424227 -->
<g id="edge17" class="edge">
<title>7883531339165132651--2127061049655424227</title>
<path fill="none" stroke="#000000" d="M406.2517,-346.4412C373.4947,-326.0552 325.461,-296.1619 292.2479,-275.492"/>
<text text-anchor="start" x="354.5" y="-308.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">b_ξ</text>
<text text-anchor="start" x="292.2479" y="-278.092" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="393.2517" y="-349.0412" font-family="Times,serif" font-size="8.00" fill="#000000">3 b </text>
</g>
<!-- 6955477092239402756 -->
<g id="node12" class="node">
<title>6955477092239402756</title>
<polygon fill="#d3d3d3" stroke="#000000" points="481.5,-279 407.5,-279 407.5,-225 481.5,-225 481.5,-279"/>
<text text-anchor="middle" x="444.5" y="-249.8" font-family="Times,serif" font-size="9.00" fill="#000000">placeholder_a_ξ</text>
</g>
<!-- 7883531339165132651--6955477092239402756 -->
<g id="edge14" class="edge">
<title>7883531339165132651--6955477092239402756</title>
<path fill="none" stroke="#000000" d="M443.1164,-332.9415C443.4125,-315.6185 443.7615,-295.201 444.0349,-279.2064"/>
<text text-anchor="start" x="444.5" y="-308.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">a_ξ</text>
<text text-anchor="start" x="425.0349" y="-281.8064" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="430.1164" y="-326.5415" font-family="Times,serif" font-size="8.00" fill="#000000">2 a </text>
</g>
<!-- 17270202296712975879 -->
<g id="node11" class="node">
<title>17270202296712975879</title>
<polygon fill="#d3d3d3" stroke="#000000" points="461,-54 384,-54 384,0 461,0 461,-54"/>
<text text-anchor="middle" x="422.5" y="-24.8" font-family="Times,serif" font-size="9.00" fill="#000000">placeholder_w_θ</text>
</g>
<!-- 7302414247694462375 -->
<g id="node13" class="node">
<title>7302414247694462375</title>
<polygon fill="none" stroke="#000000" points="841.5,-288 769.5,-288 769.5,-216 841.5,-216 841.5,-288"/>
<text text-anchor="middle" x="805.5" y="-254.8" font-family="Times,serif" font-size="9.00" fill="#000000">NLARX</text>
<text text-anchor="middle" x="805.5" y="-244.8" font-family="Times,serif" font-size="9.00" fill="#000000">z_t</text>
</g>
<!-- 7302414247694462375--4485287826409161871 -->
<g id="edge16" class="edge">
<title>7302414247694462375--4485287826409161871</title>
<path fill="none" stroke="#000000" d="M769.4139,-236.064C752.4172,-229.0707 731.7,-221.2707 712.5,-216 627.6047,-192.6951 363.6653,-153.9537 266.7041,-140.1109"/>
<text text-anchor="start" x="622.5" y="-191.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">z_tmin1</text>
<text text-anchor="start" x="266.7041" y="-142.7109" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="756.4139" y="-238.664" font-family="Times,serif" font-size="8.00" fill="#000000">2 x </text>
</g>
<!-- 9673951270592102218 -->
<g id="node16" class="node">
<title>9673951270592102218</title>
<polygon fill="none" stroke="#000000" points="654,-171 547,-171 547,-99 654,-99 654,-171"/>
<text text-anchor="middle" x="600.5" y="-137.8" font-family="Times,serif" font-size="9.00" fill="#000000">𝒩</text>
<text text-anchor="middle" x="600.5" y="-127.8" font-family="Times,serif" font-size="9.00" fill="#000000">gaussianmeanprecision_1</text>
</g>
<!-- 7302414247694462375--9673951270592102218 -->
<g id="edge15" class="edge">
<title>7302414247694462375--9673951270592102218</title>
<path fill="none" stroke="#000000" d="M769.322,-231.3521C737.3675,-213.1146 690.3423,-186.2759 654.033,-165.553"/>
<text text-anchor="start" x="708.5" y="-191.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">θ</text>
<text text-anchor="start" x="654.033" y="-168.153" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="756.322" y="-224.9521" font-family="Times,serif" font-size="8.00" fill="#000000">3 θ </text>
</g>
<!-- 13615283093068107238 -->
<g id="node18" class="node">
<title>13615283093068107238</title>
<polygon fill="none" stroke="#000000" points="841.5,-171 769.5,-171 769.5,-99 841.5,-99 841.5,-171"/>
<text text-anchor="middle" x="805.5" y="-137.8" font-family="Times,serif" font-size="9.00" fill="#000000">Gam</text>
<text text-anchor="middle" x="805.5" y="-127.8" font-family="Times,serif" font-size="9.00" fill="#000000">gamma_1</text>
</g>
<!-- 7302414247694462375--13615283093068107238 -->
<g id="edge18" class="edge">
<title>7302414247694462375--13615283093068107238</title>
<path fill="none" stroke="#000000" d="M805.5,-215.9415C805.5,-201.7071 805.5,-185.3833 805.5,-171.1409"/>
<text text-anchor="start" x="805.5" y="-191.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">γ</text>
<text text-anchor="start" x="786.5" y="-173.7409" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="792.5" y="-209.5415" font-family="Times,serif" font-size="8.00" fill="#000000">6 γ </text>
</g>
<!-- 3687048895192590763 -->
<g id="node19" class="node">
<title>3687048895192590763</title>
<polygon fill="#d3d3d3" stroke="#000000" points="1030,-162 957,-162 957,-108 1030,-108 1030,-162"/>
<text text-anchor="middle" x="993.5" y="-132.8" font-family="Times,serif" font-size="9.00" fill="#000000">placeholder_u_t</text>
</g>
<!-- 7302414247694462375--3687048895192590763 -->
<g id="edge20" class="edge">
<title>7302414247694462375--3687048895192590763</title>
<path fill="none" stroke="#000000" d="M841.7483,-229.4412C874.8254,-208.856 923.4792,-178.5768 956.7215,-157.8887"/>
<text text-anchor="start" x="904.5" y="-191.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">u_t</text>
<text text-anchor="start" x="937.7215" y="-160.4887" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="841.7483" y="-232.0412" font-family="Times,serif" font-size="8.00" fill="#000000">5 u </text>
</g>
<!-- 3990627805482885988 -->
<g id="node21" class="node">
<title>3990627805482885988</title>
<polygon fill="none" stroke="#000000" points="1252,-171 1145,-171 1145,-99 1252,-99 1252,-171"/>
<text text-anchor="middle" x="1198.5" y="-137.8" font-family="Times,serif" font-size="9.00" fill="#000000">𝒩</text>
<text text-anchor="middle" x="1198.5" y="-127.8" font-family="Times,serif" font-size="9.00" fill="#000000">gaussianmeanprecision_2</text>
</g>
<!-- 7302414247694462375--3990627805482885988 -->
<g id="edge12" class="edge">
<title>7302414247694462375--3990627805482885988</title>
<path fill="none" stroke="#000000" d="M841.6731,-242.0331C895.4449,-227.1216 999.5733,-197.8777 1087.5,-171 1106.2928,-165.2554 1126.7805,-158.6956 1144.9805,-152.7631"/>
<text text-anchor="start" x="1025.5" y="-191.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">η</text>
<text text-anchor="start" x="1125.9805" y="-155.3631" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="841.6731" y="-244.6331" font-family="Times,serif" font-size="8.00" fill="#000000">4 η </text>
</g>
<!-- 14810396312109737190--18176827165368624414 -->
<g id="edge3" class="edge">
<title>14810396312109737190--18176827165368624414</title>
<path fill="none" stroke="#000000" d="M621.8583,-332.8181C620.8779,-326.8892 620.0317,-320.7916 619.5,-315 618.4214,-303.2523 619.642,-290.319 621.4283,-279.2434"/>
<text text-anchor="start" x="619.5" y="-308.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">clamp_1</text>
<text text-anchor="start" x="602.4283" y="-281.8434" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="602.8583" y="-326.4181" font-family="Times,serif" font-size="8.00" fill="#000000">3 in2 </text>
</g>
<!-- 14810396312109737190--7302414247694462375 -->
<g id="edge8" class="edge">
<title>14810396312109737190--7302414247694462375</title>
<path fill="none" stroke="#000000" d="M665.5409,-345.041C695.8121,-324.9175 738.863,-296.2985 769.2053,-276.1277"/>
<text text-anchor="start" x="722.5" y="-308.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">z_t</text>
<text text-anchor="start" x="756.2053" y="-278.7277" font-family="Times,serif" font-size="8.00" fill="#000000">1 y </text>
<text text-anchor="start" x="665.5409" y="-347.641" font-family="Times,serif" font-size="8.00" fill="#000000">2 in1 </text>
</g>
<!-- 12092874215945484629 -->
<g id="node15" class="node">
<title>12092874215945484629</title>
<polygon fill="#d3d3d3" stroke="#000000" points="653,-54 576,-54 576,0 653,0 653,-54"/>
<text text-anchor="middle" x="614.5" y="-24.8" font-family="Times,serif" font-size="9.00" fill="#000000">placeholder_m_θ</text>
</g>
<!-- 9673951270592102218--17270202296712975879 -->
<g id="edge4" class="edge">
<title>9673951270592102218--17270202296712975879</title>
<path fill="none" stroke="#000000" d="M546.6111,-102.3034C519.0859,-85.6027 486.2246,-65.6644 461.2322,-50.5004"/>
<text text-anchor="start" x="511.5" y="-74.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">w_θ</text>
<text text-anchor="start" x="461.2322" y="-53.1004" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="531.6111" y="-104.9034" font-family="Times,serif" font-size="8.00" fill="#000000">3 w </text>
</g>
<!-- 9673951270592102218--12092874215945484629 -->
<g id="edge11" class="edge">
<title>9673951270592102218--12092874215945484629</title>
<path fill="none" stroke="#000000" d="M605.2007,-98.7371C607.0913,-84.1526 609.2336,-67.6262 610.9851,-54.1146"/>
<text text-anchor="start" x="608.5" y="-74.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">m_θ</text>
<text text-anchor="start" x="591.9851" y="-56.7146" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="590.2007" y="-92.3371" font-family="Times,serif" font-size="8.00" fill="#000000">2 m </text>
</g>
<!-- 13413068550611593338 -->
<g id="node17" class="node">
<title>13413068550611593338</title>
<polygon fill="#d3d3d3" stroke="#000000" points="1034.5,-54 958.5,-54 958.5,0 1034.5,0 1034.5,-54"/>
<text text-anchor="middle" x="996.5" y="-24.8" font-family="Times,serif" font-size="9.00" fill="#000000">placeholder_b_γ</text>
</g>
<!-- 13615283093068107238--5547258369390740305 -->
<g id="edge5" class="edge">
<title>13615283093068107238--5547258369390740305</title>
<path fill="none" stroke="#000000" d="M805.5,-98.7371C805.5,-84.1526 805.5,-67.6262 805.5,-54.1146"/>
<text text-anchor="start" x="805.5" y="-74.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">a_γ</text>
<text text-anchor="start" x="786.5" y="-56.7146" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="792.5" y="-92.3371" font-family="Times,serif" font-size="8.00" fill="#000000">2 a </text>
</g>
<!-- 13615283093068107238--13413068550611593338 -->
<g id="edge7" class="edge">
<title>13615283093068107238--13413068550611593338</title>
<path fill="none" stroke="#000000" d="M841.8752,-114.4318C875.2119,-95.5817 924.3956,-67.771 958.2966,-48.6019"/>
<text text-anchor="start" x="911.5" y="-74.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">b_γ</text>
<text text-anchor="start" x="939.2966" y="-51.2019" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="841.8752" y="-117.0318" font-family="Times,serif" font-size="8.00" fill="#000000">3 b </text>
</g>
<!-- 3990627805482885988--2867327560033567407 -->
<g id="edge2" class="edge">
<title>3990627805482885988--2867327560033567407</title>
<path fill="none" stroke="#000000" d="M1196.8212,-98.7371C1196.146,-84.1526 1195.3808,-67.6262 1194.7553,-54.1146"/>
<text text-anchor="start" x="1195.5" y="-74.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">m_η</text>
<text text-anchor="start" x="1175.7553" y="-56.7146" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="1181.8212" y="-92.3371" font-family="Times,serif" font-size="8.00" fill="#000000">2 m </text>
</g>
<!-- 3990627805482885988--12774672834151751891 -->
<g id="edge9" class="edge">
<title>3990627805482885988--12774672834151751891</title>
<path fill="none" stroke="#000000" d="M1252.0867,-104.0515C1282.306,-86.5987 1319.3732,-65.1909 1346.7606,-49.3736"/>
<text text-anchor="start" x="1302.5" y="-74.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">w_η</text>
<text text-anchor="start" x="1327.7606" y="-51.9736" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="1252.0867" y="-106.6515" font-family="Times,serif" font-size="8.00" fill="#000000">3 w </text>
</g>
</g>
</svg>

```julia
# Specify recognition model
q1 = PosteriorFactorization(z_t, z_tmin1, θ, η, γ, ξ, ids=[:z_t, :z_tmin1, :θ, :η, :γ, :ξ])
algo1 = variationalAlgorithm(q1, free_energy=true)

# Compile inference algorithm
source_code1 = algorithmSourceCode(algo1, free_energy=true)
eval(Meta.parse(source_code1));
# println(source_code1)
```

### Infer parameters on training data

```julia
# Inference parameters
num_iterations = 10

# Initialize marginal distribution and observed data dictionaries
data = Dict()
marginals = Dict()

# Initialize free energy tracking array
free_energy_trn = zeros(T_trn, num_iterations)

# Initialize arrays of parameterizations
params_z = (zeros(2,T_trn+1), repeat(.1 .*float(eye(2)), outer=(1,1,T_trn+1)))
params_θ = (ones(3,T_trn+1), repeat(.1 .*float(eye(3)), outer=(1,1,T_trn+1)))
params_η = (2*ones(1,T_trn+1), 1e2 *ones(1,T_trn+1))
params_γ = (1e8*ones(1,T_trn+1), 1e3*ones(1,T_trn+1))
params_ξ = (1e8*ones(1,T_trn+1), 1e1*ones(1,T_trn+1))

# Start progress bar
p = Progress(T_trn, 1, "At time ")

# Perform inference at each time-step
for t = 1:T_trn

    # Update progress bar
    update!(p, t)

    # Initialize marginals
    marginals[:z_tmin1] = ProbabilityDistribution(Multivariate, GaussianMeanPrecision, m=params_z[1][:,t], w=params_z[2][:,:,t])
    marginals[:z_t] = ProbabilityDistribution(Multivariate, GaussianMeanPrecision, m=params_z[1][:,t], w=params_z[2][:,:,t])
    marginals[:θ] = ProbabilityDistribution(Multivariate, GaussianMeanPrecision, m=params_θ[1][:,t], w=params_θ[2][:,:,t])
    marginals[:η] = ProbabilityDistribution(Univariate, GaussianMeanPrecision, m=params_η[1][1,t], w=params_η[2][1,t])
    marginals[:γ] = ProbabilityDistribution(Univariate, Gamma, a=params_γ[1][1,t], b=params_γ[2][1,t])
    marginals[:ξ] = ProbabilityDistribution(Univariate, Gamma, a=params_ξ[1][1,t], b=params_ξ[2][1,t])

    data = Dict(:y_t => output_trn[t],
                :u_t => input_trn[t],
                :m_z => params_z[1][:,t],
                :w_z => params_z[2][:,:,t],
                :m_θ => params_θ[1][:,t],
                :w_θ => params_θ[2][:,:,t],
                :m_η => params_η[1][1,t],
                :w_η => params_η[2][1,t],
                :a_γ => params_γ[1][1,t],
                :b_γ => params_γ[2][1,t],
                :a_ξ => params_ξ[1][1,t],
                :b_ξ => params_ξ[2][1,t])

    # Iterate variational parameter updates
    for i = 1:num_iterations

        # Update parameters
        stepη!(data, marginals)
        stepθ!(data, marginals)

        # Update states
        stepz_t!(data, marginals)
        stepz_tmin1!(data, marginals)

        # Update noise
        stepγ!(data, marginals)
        stepξ!(data, marginals)

        # Compute free energy
        free_energy_trn[t, i] = freeEnergy(data, marginals)

    end

    # Store current parameterizations of marginals
    params_z[1][:,t+1] = unsafeMean(marginals[:z_t])
    params_z[2][:,:,t+1] = marginals[:z_t].params[:w]
    params_θ[1][:,t+1] = unsafeMean(marginals[:θ])
    params_θ[2][:,:,t+1] = marginals[:θ].params[:w]
    params_η[1][1,t+1] = unsafeMean(marginals[:η])
    params_η[2][1,t+1] = marginals[:η].params[:w]
    params_γ[1][1,t+1] = marginals[:γ].params[:a]
    params_γ[2][1,t+1] = marginals[:γ].params[:b]
    params_ξ[1][1,t+1] = marginals[:ξ].params[:a]
    params_ξ[2][1,t+1] = marginals[:ξ].params[:b]

end
```

    At time 100%|███████████████████████████████████████████| Time: 0:03:5939m

### Simulate validation data

```julia
# Prediction graph
graph2 = FactorGraph()

# Autoregressive node
@RV z_pred ~ NLatentAutoregressiveX(placeholder(:θ, dims=(3,)), placeholder(:z_tmin1, dims=(2,)), placeholder(:η), placeholder(:u_t), placeholder(:γ), g=g, id=:z_pred_t)

# Draw time-slice subgraph
ForneyLab.draw(graph2)

# Inference algorithm
q2 = PosteriorFactorization(z_pred, ids=[:z_pred])
algo2 = variationalAlgorithm(q2, free_energy=true)
source_code2 = algorithmSourceCode(algo2, free_energy=true)
eval(Meta.parse(source_code2));
# println(source_code2)
```

<?xml version="1.0" encoding="UTF-8" standalone="no"?>
<!DOCTYPE svg PUBLIC "-//W3C//DTD SVG 1.1//EN"
 "http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd">
<!-- Generated by graphviz version 2.40.1 (20161225.0304)
 -->
<!-- Title: G Pages: 1 -->
<svg width="832pt" height="278pt"
 viewBox="0.00 0.00 831.50 278.00" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink">
<g id="graph0" class="graph" transform="scale(1 1) rotate(0) translate(4 274)">
<title>G</title>
<polygon fill="#ffffff" stroke="transparent" points="-4,4 -4,-274 827.5,-274 827.5,4 -4,4"/>
<!-- 16433011899932135561 -->
<g id="node1" class="node">
<title>16433011899932135561</title>
<polygon fill="#d3d3d3" stroke="#000000" points="66,-54 0,-54 0,0 66,0 66,-54"/>
<text text-anchor="middle" x="33" y="-24.8" font-family="Times,serif" font-size="9.00" fill="#000000">placeholder_η</text>
</g>
<!-- 5163914471287504424 -->
<g id="node2" class="node">
<title>5163914471287504424</title>
<polygon fill="none" stroke="#000000" points="432,-171 360,-171 360,-99 432,-99 432,-171"/>
<text text-anchor="middle" x="396" y="-137.8" font-family="Times,serif" font-size="9.00" fill="#000000">NLARX</text>
<text text-anchor="middle" x="396" y="-127.8" font-family="Times,serif" font-size="9.00" fill="#000000">z_pred_t</text>
</g>
<!-- 5163914471287504424--16433011899932135561 -->
<g id="edge4" class="edge">
<title>5163914471287504424--16433011899932135561</title>
<path fill="none" stroke="#000000" d="M359.7786,-124.2909C307.4843,-108.8211 207.8003,-79.3003 123,-54 104.2415,-48.4034 83.2937,-42.121 66.3308,-37.0254"/>
<text text-anchor="start" x="214" y="-74.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">η</text>
<text text-anchor="start" x="66.3308" y="-39.6254" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="346.7786" y="-126.8909" font-family="Times,serif" font-size="8.00" fill="#000000">4 η </text>
</g>
<!-- 7294770454988156620 -->
<g id="node3" class="node">
<title>7294770454988156620</title>
<polygon fill="#d3d3d3" stroke="#000000" points="247,-54 181,-54 181,0 247,0 247,-54"/>
<text text-anchor="middle" x="214" y="-24.8" font-family="Times,serif" font-size="9.00" fill="#000000">placeholder_θ</text>
</g>
<!-- 5163914471287504424--7294770454988156620 -->
<g id="edge6" class="edge">
<title>5163914471287504424--7294770454988156620</title>
<path fill="none" stroke="#000000" d="M359.607,-113.4042C326.8268,-93.9522 279.045,-65.5981 247.2658,-46.7402"/>
<text text-anchor="start" x="305" y="-74.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">θ</text>
<text text-anchor="start" x="247.2658" y="-49.3402" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="346.607" y="-116.0042" font-family="Times,serif" font-size="8.00" fill="#000000">3 θ </text>
</g>
<!-- 11160641670622635050 -->
<g id="node4" class="node">
<title>11160641670622635050</title>
<polygon fill="#d3d3d3" stroke="#000000" points="429.5,-54 362.5,-54 362.5,0 429.5,0 429.5,-54"/>
<text text-anchor="middle" x="396" y="-24.8" font-family="Times,serif" font-size="9.00" fill="#000000">placeholder_γ</text>
</g>
<!-- 5163914471287504424--11160641670622635050 -->
<g id="edge2" class="edge">
<title>5163914471287504424--11160641670622635050</title>
<path fill="none" stroke="#000000" d="M396,-98.7371C396,-84.1526 396,-67.6262 396,-54.1146"/>
<text text-anchor="start" x="396" y="-74.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">γ</text>
<text text-anchor="start" x="377" y="-56.7146" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="383" y="-92.3371" font-family="Times,serif" font-size="8.00" fill="#000000">6 γ </text>
</g>
<!-- 15668819891996249435 -->
<g id="node5" class="node">
<title>15668819891996249435</title>
<polygon fill="#d3d3d3" stroke="#000000" points="635,-54 545,-54 545,0 635,0 635,-54"/>
<text text-anchor="middle" x="590" y="-24.8" font-family="Times,serif" font-size="9.00" fill="#000000">placeholder_z_tmin1</text>
</g>
<!-- 5163914471287504424--15668819891996249435 -->
<g id="edge1" class="edge">
<title>5163914471287504424--15668819891996249435</title>
<path fill="none" stroke="#000000" d="M432.0353,-114.9391C463.9221,-97.1877 510.5903,-71.2075 544.9371,-52.0865"/>
<text text-anchor="start" x="505" y="-74.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">z_tmin1</text>
<text text-anchor="start" x="525.9371" y="-54.6865" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="432.0353" y="-117.5391" font-family="Times,serif" font-size="8.00" fill="#000000">2 x </text>
</g>
<!-- 207787082123420713 -->
<g id="node6" class="node">
<title>207787082123420713</title>
<polygon fill="#d3d3d3" stroke="#000000" points="823.5,-54 750.5,-54 750.5,0 823.5,0 823.5,-54"/>
<text text-anchor="middle" x="787" y="-24.8" font-family="Times,serif" font-size="9.00" fill="#000000">placeholder_u_t</text>
</g>
<!-- 5163914471287504424--207787082123420713 -->
<g id="edge5" class="edge">
<title>5163914471287504424--207787082123420713</title>
<path fill="none" stroke="#000000" d="M432.0274,-125.2633C487.7888,-110.1657 598.3421,-80.1233 692,-54 711.3235,-48.6102 732.8275,-42.5074 750.4826,-37.469"/>
<text text-anchor="start" x="619" y="-74.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">u_t</text>
<text text-anchor="start" x="731.4826" y="-40.069" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="432.0274" y="-127.8633" font-family="Times,serif" font-size="8.00" fill="#000000">5 u </text>
</g>
<!-- 161314157077685787422 -->
<g id="node7" class="node">
<title>161314157077685787422</title>
</g>
<!-- 161314157077685787422--5163914471287504424 -->
<g id="edge3" class="edge">
<title>161314157077685787422--5163914471287504424</title>
<path fill="none" stroke="#000000" d="M396,-215.747C396,-202.2495 396,-185.7693 396,-171.2253"/>
<text text-anchor="start" x="396" y="-191.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">z_pred</text>
<text text-anchor="start" x="383" y="-173.8253" font-family="Times,serif" font-size="8.00" fill="#000000">1 y </text>
</g>
</g>
</svg>

```julia
# Initialize free energy tracking array
free_energy_pred = zeros(T_val, num_iterations)

# Initialize future state arrays
params_preds = (zeros(2, T_val), repeat(.1 .*float(eye(2)), outer=(1,1,T_val)))

# Start simulation with known output
params_preds[1][1,2] = output_val[2]
params_preds[1][2,2] = output_val[1]

# Start progress bar
p = Progress(T_val, 1, "At time ")

for t = 3:T_val
    update!(p, t)

    # Initialize marginals
    marginals[:z_pred] = ProbabilityDistribution(Multivariate, GaussianMeanPrecision, m=params_preds[1][:,t], w=params_preds[2][:,:,t])

    # Clamp data
    data = Dict(:u_t => input[t],
                :z_tmin1 => params_preds[1][:,t-1],
                :θ => params_θ[1][:,end],
                :η => params_η[1][end],
                :γ => params_γ[1][end]/params_γ[2][end])

    # Iterate variational parameter updates
    for i = 1:num_iterations

        # Make prediction
        stepz_pred!(data, marginals)

        # Compute free energy
        free_energy_pred[t, i] = freeEnergy(data, marginals)
    end

    # Store current parameterizations of marginals
    params_preds[1][:,t] = unsafeMean(marginals[:z_pred])
    params_preds[2][:,:,t] = marginals[:z_pred].params[:w]

end
```

    At time 100%|███████████████████████████████████████████| Time: 0:00:13

```julia
# Store predictions for later comparisons
results_NLARX = Dict()
results_NLARX["preds"] = params_preds
results_NLARX["params_z"] = params_z;
results_NLARX["params_θ"] = params_θ;
results_NLARX["params_η"] = params_η;
results_NLARX["params_γ"] = params_γ;
results_NLARX["params_ξ"] = params_ξ;
results_NLARX["FE_pred"] = free_energy_pred;
results_NLARX["FE_trn"] = free_energy_trn;
```

### Visualize results

```julia
# Mean and std dev of predictions
predictions_mean = params_preds[1][1,:]
predictions_std = sqrt.(inv.(params_preds[2][1,1,:]))

# Subsample for visualization
ss = 10
# viz_ix = 50000:ss:60000
viz_ix = 1024:ss:40000

# Plot predictions
p23 = scatter(viz_ix, output[viz_ix], label="observations", xlabel="time (t)", ylims=[-.4, .4], color="black")
plot!(viz_ix, predictions_mean[viz_ix], ribbon=[predictions_std[viz_ix], predictions_std[viz_ix]], label="predictions", color="red")
```

```julia
Plots.savefig(p23, "figures/simulations_nlarx.png")
```

```julia
# Compute prediction error
pred_error = (predictions_mean[2:end] .- output_val[2:end]).^2

# Subsample for visualization
ss = 10
viz_ix = 1:ss:40000

# Scatter error over time
p24 = scatter(viz_ix, pred_error[viz_ix], color="black", xlabel="time (t)", ylabel="Prediction error", label="", ylims=[1e-10, 1e-2], yscale=:log10)
```

```julia
Plots.savefig(p24, "figures/pred-error_nlarx.png")
```

```julia
# Subsample for visualization
ss = 10
viz_ix = 1:ss:40000

# Scatter error over time
p24 = plot(viz_ix, free_energy_trn[viz_ix,end], color="black", xlabel="time (t)", ylabel="F[q]", label="", title="Free energy at training time")
```

```julia
# Subsample for visualization
ss = 10
viz_ix = 1:ss:40000

# Scatter error over time
p24 = plot(viz_ix, free_energy_pred[viz_ix,end], color="black", xlabel="time (t)", ylabel="F[q]", label="", title="Free energy of predictions")
```

## Baseline: linear autoregression

```julia
# System identification graph
graph3 = FactorGraph()

# Static parameters
@RV θ ~ GaussianMeanPrecision(placeholder(:m_θ, dims=(2,)), placeholder(:w_θ, dims=(2,2)))
@RV η ~ GaussianMeanPrecision(placeholder(:m_η), placeholder(:w_η))
@RV γ ~ Gamma(placeholder(:a_γ), placeholder(:b_γ))
@RV ξ ~ Gamma(placeholder(:a_ξ), placeholder(:b_ξ))

# Linear autoregression function
g(θ, x) = θ[1]*x[1] + θ[2]*x[2]

# State prior
@RV z_tmin1 ~ GaussianMeanPrecision(placeholder(:m_z, dims=(2,)), placeholder(:w_z, dims=(2, 2)), id=:z_tmin1)

# Autoregressive node
@RV z_t ~ NLatentAutoregressiveX(θ, z_tmin1, η, placeholder(:u_t), γ, g=g, id=:z_t)

# Specify likelihood
@RV y_t ~ GaussianMeanPrecision(dot([1. , 0.], z_t), ξ, id=:y_t)

# Placeholder for observation
placeholder(y_t, :y_t)

# Specify recognition model
q3 = PosteriorFactorization(z_t, z_tmin1, θ, η, γ, ξ, ids=[:z_t, :z_tmin1, :θ, :η, :γ, :ξ])
algo3 = variationalAlgorithm(q3, free_energy=true)

# Compile inference algorithm
source_code3 = algorithmSourceCode(algo3, free_energy=true)
eval(Meta.parse(source_code3));
# println(source_code3)
```

```julia
# Inference parameters
num_iterations = 10

# Initialize marginal distribution and observed data dictionaries
data = Dict()
marginals = Dict()

# Initialize free energy tracking array
free_energy_trn = zeros(T_trn, num_iterations)

# Initialize arrays of parameterizations
params_z = (zeros(2,T_trn+1), repeat(.1 .*float(eye(2)), outer=(1,1,T_trn+1)))
params_θ = (ones(2,T_trn+1), repeat(.1 .*float(eye(2)), outer=(1,1,T_trn+1)))
params_η = (2*ones(1,T_trn+1), 1e2 *ones(1,T_trn+1))
params_γ = (1e8*ones(1,T_trn+1), 1e3*ones(1,T_trn+1))
params_ξ = (1e8*ones(1,T_trn+1), 1e1*ones(1,T_trn+1))

# Start progress bar
p = Progress(T_trn, 1, "At time ")

# Perform inference at each time-step
for t = 1:T_trn

    # Update progress bar
    update!(p, t)

    # Initialize marginals
    marginals[:z_tmin1] = ProbabilityDistribution(Multivariate, GaussianMeanPrecision, m=params_z[1][:,t], w=params_z[2][:,:,t])
    marginals[:z_t] = ProbabilityDistribution(Multivariate, GaussianMeanPrecision, m=params_z[1][:,t], w=params_z[2][:,:,t])
    marginals[:θ] = ProbabilityDistribution(Multivariate, GaussianMeanPrecision, m=params_θ[1][:,t], w=params_θ[2][:,:,t])
    marginals[:η] = ProbabilityDistribution(Univariate, GaussianMeanPrecision, m=params_η[1][1,t], w=params_η[2][1,t])
    marginals[:γ] = ProbabilityDistribution(Univariate, Gamma, a=params_γ[1][1,t], b=params_γ[2][1,t])
    marginals[:ξ] = ProbabilityDistribution(Univariate, Gamma, a=params_ξ[1][1,t], b=params_ξ[2][1,t])

    data = Dict(:y_t => output_trn[t],
                :u_t => input_trn[t],
                :m_z => params_z[1][:,t],
                :w_z => params_z[2][:,:,t],
                :m_θ => params_θ[1][:,t],
                :w_θ => params_θ[2][:,:,t],
                :m_η => params_η[1][1,t],
                :w_η => params_η[2][1,t],
                :a_γ => params_γ[1][1,t],
                :b_γ => params_γ[2][1,t],
                :a_ξ => params_ξ[1][1,t],
                :b_ξ => params_ξ[2][1,t])

    # Iterate variational parameter updates
    for i = 1:num_iterations

        # Update parameters
        stepη!(data, marginals)
        stepθ!(data, marginals)

        # Update states
        stepz_t!(data, marginals)
        stepz_tmin1!(data, marginals)

        # Update noise
        stepγ!(data, marginals)
        stepξ!(data, marginals)

        # Compute free energy
        free_energy_trn[t, i] = freeEnergy(data, marginals)

    end

    # Store current parameterizations of marginals
    params_z[1][:,t+1] = unsafeMean(marginals[:z_t])
    params_z[2][:,:,t+1] = marginals[:z_t].params[:w]
    params_θ[1][:,t+1] = unsafeMean(marginals[:θ])
    params_θ[2][:,:,t+1] = marginals[:θ].params[:w]
    params_η[1][1,t+1] = unsafeMean(marginals[:η])
    params_η[2][1,t+1] = marginals[:η].params[:w]
    params_γ[1][1,t+1] = marginals[:γ].params[:a]
    params_γ[2][1,t+1] = marginals[:γ].params[:b]
    params_ξ[1][1,t+1] = marginals[:ξ].params[:a]
    params_ξ[2][1,t+1] = marginals[:ξ].params[:b]

end
```

    At time 100%|███████████████████████████████████████████| Time: 0:03:29

```julia
# Prediction graph
graph4 = FactorGraph()

# Autoregressive node
@RV z_pred ~ NLatentAutoregressiveX(placeholder(:θ, dims=(3,)), placeholder(:z_tmin1, dims=(2,)), placeholder(:η), placeholder(:u_t), placeholder(:γ), g=g, id=:z_pred_t)

# Inference algorithm
q4 = PosteriorFactorization(z_pred, ids=[:z_pred])
algo4 = variationalAlgorithm(q4, free_energy=true)
source_code4 = algorithmSourceCode(algo4, free_energy=true)
eval(Meta.parse(source_code4));
```

```julia
# Initialize free energy tracking array
free_energy_pred = zeros(T_val, num_iterations)

# Initialize future state arrays
params_preds = (zeros(2, T_val), repeat(.1 .*float(eye(2)), outer=(1,1,T_val)))

# Start simulation with known output
params_preds[1][1,2] = output_val[2]
params_preds[1][2,2] = output_val[1]

# Start progress bar
p = Progress(T_val, 1, "At time ")

for t = 3:T_val
    update!(p, t)

    # Initialize marginals
    marginals[:z_pred] = ProbabilityDistribution(Multivariate, GaussianMeanPrecision, m=params_preds[1][:,t], w=params_preds[2][:,:,t])

    # Clamp data
    data = Dict(:u_t => input[t],
                :z_tmin1 => params_preds[1][:,t-1],
                :θ => params_θ[1][:,end],
                :η => params_η[1][end],
                :γ => params_γ[1][end]/params_γ[2][end])

    # Iterate variational parameter updates
    for i = 1:num_iterations

        # Make prediction
        stepz_pred!(data, marginals)

        # Compute free energy
        free_energy_pred[t, i] = freeEnergy(data, marginals)
    end

    # Store current parameterizations of marginals
    params_preds[1][:,t] = unsafeMean(marginals[:z_pred])
    params_preds[2][:,:,t] = marginals[:z_pred].params[:w]

end
```

    At time 100%|███████████████████████████████████████████| Time: 0:00:11

```julia
# Store predictions for later comparisons
results_LARX = Dict()
results_LARX["preds"] = params_preds
results_LARX["params_z"] = params_z;
results_LARX["params_θ"] = params_θ;
results_LARX["params_η"] = params_η;
results_LARX["params_γ"] = params_γ;
results_LARX["params_ξ"] = params_ξ;
results_LARX["FE_pred"] = free_energy_pred;
results_LARX["FE_trn"] = free_energy_trn;
```

### Visualize results

```julia
# Mean and std dev of predictions
predictions_mean = params_preds[1][1,:]
predictions_std = sqrt.(inv.(params_preds[2][1,1,:]))

# Subsample for visualization
ss = 40
# viz_ix = 50000:ss:60000
viz_ix = 1024:ss:40000

# Plot predictions
p230 = scatter(viz_ix, output[viz_ix], color="black", label="observations", xlabel="time (t)", ylims=[-.4, .4], legend=:topleft)
plot!(viz_ix, predictions_mean[viz_ix], ribbon=[predictions_std[viz_ix], predictions_std[viz_ix]], color="red",  label="predictions")
```

```julia
Plots.savefig(p230, "figures/simulations_larx.png")
```

```julia
# Compute prediction error
pred_error = (predictions_mean[3:end] .- output_val[3:end]).^2

# Subsample for visualization
ss = 10
viz_ix = 1:ss:40000

# Scatter error over time
p240 = scatter(viz_ix, pred_error[viz_ix], color="black", xlabel="time (t)", ylabel="Prediction error", ylims=[1e-10, 1e0], label="", yscale=:log10)
```

```julia
Plots.savefig(p240, "figures/sim-error_larx.png")
```

```julia
# Subsample for visualization
ss = 10
viz_ix = 1:ss:40000

# Scatter error over time
p24 = plot(viz_ix, free_energy_trn[viz_ix,end], color="black", xlabel="time (t)", ylabel="F[q]", label="", title="Free energy at training time")
```

```julia
# Subsample for visualization
ss = 10
viz_ix = 1:ss:40000

# Scatter error over time
p24 = plot(viz_ix, free_energy_pred[viz_ix,end], color="black", xlabel="time (t)", ylabel="F[q]", label="", title="Free energy of predictions")
```

### Baseline: offline sigmoid network NARX model

As baseline, I ran a nonlinear ARX model using Matlab's System Identification Toolbox. Parameters are estimated using a Prediction Error Minimisation (PEM) procedure. It simultaneously estimates a static nonlinearity that I'm modelling with a sigmoid network. I've used 4 units to keep the comparison with the 4-coefficient NLARX model fair.

```julia
using MAT
results_NARX = matread("results/results_narx_sigmoidnet4_simulation.mat")
```

    Dict{String,Any} with 2 entries:
      "pred_states" => [0.0; 0.0; … ; 0.0048962; 0.0449377]
      "pred_error"  => [8.83187e-5; 0.000168246; … ; 0.000351895; 2.3589e-5]

```julia
# Subsample for visualization
ss = 10
viz_ix = 1:ss:40000

# Plot predictions
p402 = scatter(viz_ix, output[viz_ix], label="observations", ylims=[-.4, .4], color="black")
plot!(viz_ix, results_NARX["pred_states"][viz_ix], label="predictions", xlabel="time (t)", color="red")
```

```julia
# Subsample for visualization
ss = 40
viz_ix = 1:ss:40000

# Scatter error over time
p403 = scatter(viz_ix, results_NARX["pred_error"][viz_ix], color="black", xlabel="time (t)", ylims=[1e-10, 1e0], ylabel="Prediction error", label="", yscale=:log10)
```

## Comparison

```julia
# Store prediction error for each step ahead
prederror_LARX = (results_LARX["preds"][1][1,3:end] .- output_val[3:end]).^2
prederror_NLARX = (results_NLARX["preds"][1][1,3:end] .- output_val[3:end]).^2
prederror_NARX = results_NARX["pred_error"]

println("MSE FEM-LARX = "*string(mean(prederror_LARX)))
println("MSE FEM-NLARX = "*string(mean(prederror_NLARX)))
println("MSE PEM-NARX = "*string(mean(prederror_NARX)))
```

    MSE FEM-LARX = 0.0010204335589997665
    MSE FEM-NLARX = 0.0009256401174690808
    MSE PEM-NARX = 0.001001367833873263

```julia
# Joint plot of prediction errors

# Subsample for visualization
ss = 40
viz_ix = 1:ss:40000

# Scatter error over time
p71 = plot(viz_ix, output[viz_ix], label="", color="blue", alpha=0.5, ylims=[-0.2 0.2])
plot!(viz_ix, results_NLARX["preds"][1][1,viz_ix], color="black", alpha=0.9, xlabel="", ylabel="", label="FEM-NLARX", legend=:topleft)
p72 = plot(viz_ix, output[viz_ix], label="", color="blue", alpha=0.5, ylims=[-0.2 0.2])
plot!(viz_ix, results_LARX["preds"][1][1,viz_ix], color="black", alpha=0.9, xlabel="", ylabel="", label="FEM-LARX", legend=:topleft)
p73 = plot(viz_ix, output[viz_ix], label="", color="blue", alpha=0.5, ylims=[-0.2 0.2])
plot!(viz_ix, results_NARX["pred_states"][viz_ix], color="black", alpha=0.9, xlabel="", ylabel="output", label="PEM-NARX", legend=:topleft)
p70 = plot(p73, p72, p71, layout=(1,3), size=(1200,300))
```

```julia
savefig(p70, "figures/sim-states-comparison.png")
```

```julia
# Joint plot of prediction errors

prederror_LARX = (results_LARX["preds"][1][1,3:end] .- output_val[3:end]).^2
prederror_NLARX = (results_NLARX["preds"][1][1,3:end] .- output_val[3:end]).^2
prederror_NARX = results_NARX["pred_error"]

# Subsample for visualization
ss = 40
viz_ix = 1:ss:40000

# Scatter error over time
p81 = plot(viz_ix, results_NLARX["preds"][1][1,viz_ix], color="purple", alpha=0.9, xlabel="", ylabel="", label="predictions", title="FEM-NLARX", legend=:topleft)
plot!(viz_ix, prederror_NLARX[viz_ix], color="black", alpha=0.9, xlabel="time (t)", ylabel="", label="error", yticks=:none)
p82 = plot(viz_ix, results_LARX["preds"][1][1,viz_ix], color="purple", alpha=0.9, xlabel="", label="predictions", title="FEM-LARX", legend=:topleft)
plot!(viz_ix, prederror_LARX[viz_ix], color="black", alpha=0.9, xlabel="time (t)", ylabel="", label="error", yticks=:none)
p83 = plot(viz_ix, results_NARX["pred_states"][viz_ix], color="purple", alpha=0.9, xlabel="", ylabel="", label="predictions", title="PEM-NARX", legend=:topleft)
plot!(viz_ix, prederror_NARX[viz_ix], color="black", alpha=0.9, xlabel="time (t)", ylabel="", label="error")
p80 = plot(p83, p82, p81, layout=(1,3), size=(1200,300))
```

```julia
savefig(p80, "figures/simulation-comparison.png")
```

```julia
# Joint plot of prediction errors

# Subsample for visualization
ss = 40
viz_ix = 1:ss:40000

# Scatter error over time
p61 = plot(viz_ix, prederror_NLARX[viz_ix], color="black", alpha=0.9, xlabel="time (t)", title="", ylabel="", label="error", yticks=:none, ylims=[10e-12, 10e-2], yscale=:log10)
p62 = plot(viz_ix, prederror_LARX[viz_ix], color="black", alpha=0.9, xlabel="time (t)", title="", ylabel="", label="error", yticks=:none, ylims=[10e-12, 10e-2], yscale=:log10)
p63 = plot(viz_ix, prederror_NARX[viz_ix], color="black", alpha=0.9, xlabel="time (t)", title="", ylabel="", label="error", ylims=[10e-12, 10e-2], yscale=:log10)
p60 = plot(p63, p62, p61, layout=(1,3), size=(1200,300))
```

```julia
Plots.savefig(p60, "figures/sim-error-comparison.png")
```

```julia

```


# Optimizing the check for overlapping features
---

### Overview
Currently checking for overlap by checking if each feature's bounding box overlaps with every other feature's bounding box. Although we are only checking unique combinations/pairs of bounding boxes, the algorithm is still $O(n^2)$ given that the number of combinations scales with

\begin{equation}
{}_n C_2 = \frac{n!}{2(n-2)!}
\end{equation}

Here we'll look at implementing a more efficient method, cell listing [1] ($O(n)$). Basically involves subdividing the spatial domain into cells (wider than the threshold distance) and then checking for overlap among only the features in each cell (and the neighboring cells).

[1] https://en.wikipedia.org/wiki/Cell_lists

```python
from itertools import combinations

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle

%matplotlib inline
```

### Check if two boxes overlap
(This part is fine I think)

```python
def bboxes_overlap(bbox_1, bbox_2):
    """Determines if two bounding boxes overlap or coincide

    Parameters
    ----------
    bbox_1 : array-like (or 4-tuple)
        1st bounding box
        convention: (x_min, y_min, x_max, y_max)
    bbox_2 : array-like (or 4-tuple)
        2nd bounding box
        convention: (x_min, y_min, x_max, y_max)

    Returns
    -------
    overlap : bool
        True if bounding boxes overlap / coincide
        False otherwise

    References
    ----------
    [1] https://stackoverflow.com/a/20925869/5285918
    """
    # 2 tiles overlap iff their projections onto both x and y axis overlap
    # Overlap in 1D iff box1_max > box2_min AND box1_min < box2_max
    overlap = ((bbox_1[2] >= bbox_2[0]) & (bbox_1[0] <= bbox_2[2])) & \
              ((bbox_1[3] >= bbox_2[1]) & (bbox_1[1] <= bbox_2[3]))
    return overlap
```

### Create features and boxes

```python
# Set plane dimensions
Nx, Ny = (100, 100)

# Set bbox size
wx = 10
wy = 10

# Create a random distribution of features
N = 20
X = np.random.RandomState(15).uniform(0, Nx, size=N)
Y = np.random.RandomState(38).uniform(0, Ny, size=N)

# Create DataFrame of features and their bboxes
features = pd.DataFrame({'x': X, 'y': Y})
features['x_min'] = X - wx/2
features['y_min'] = Y - wy/2
features['x_max'] = X + wx/2
features['y_max'] = Y + wy/2

# Sample features
features.head()
```

      x
      y
      x_min
      y_min
      x_max
      y_max

      0
      84.881770
      38.477312
      79.881770
      33.477312
      89.881770
      43.477312

      1
      17.889592
      85.970785
      12.889592
      80.970785
      22.889592
      90.970785

      2
      5.436321
      94.419964
      0.436321
      89.419964
      10.436321
      99.419964

      3
      36.153845
      70.282489
      31.153845
      65.282489
      41.153845
      75.282489

      4
      27.540093
      63.363410
      22.540093
      58.363410
      32.540093
      68.363410

```python
# Make figure
fig, ax = plt.subplots(figsize=(5, 5))
ax.plot(features['x'], features['y'], 'o', color='C1', ms=3, lw=0)
for i, row in features.iterrows():
    p = Rectangle((features.loc[i, 'x_min'], features.loc[i, 'y_min']),
                  wx, wy, facecolor='none', lw=1, edgecolor='C0')
    ax.add_patch(p)
    ax.text(features.loc[i, 'x_min'], features.loc[i, 'y_max'],
            f'{i}', ha='left', va='top')
ax.grid(ls=':')
```

### Cell listing method

```python
# Keep track of overlapping features
overlapping = []
# Define cell parameters
cw = 2*wx  # cell width

# Loop through a grid in x, y to create cells
for x in np.arange(0, Nx, cw):
    for y in np.arange(0, Ny, cw):
        # Create cell
        cell = [x-cw, y-cw, x+2*cw, y+2*cw]
        # Get features in cell
        in_cell = features[((cell[0] < X) & (X < cell[2]) &\
                            (cell[1] < Y) & (Y < cell[3]))]
        # Combinations
        pairs = list(combinations(in_cell.reset_index().values, 2))

        # Loop through pairs of bboxes
        for (bbox_i, bbox_j) in pairs:
            if bboxes_overlap(bbox_i[-4:], bbox_j[-4:]):
                overlapping.append(bbox_i[0])
                overlapping.append(bbox_j[0])

overlapping = np.unique(overlapping)
overlapping
```

    array([ 1.,  3.,  4.,  7., 11., 12., 14., 15., 16., 19.])

### Plot overlapping features

```python
def plot_overlapping_features(features, overlapping, plot_feature_index=False):
    """Plot features showing where bboxes overlap"""
    # Make bboxes
    features['x_min'] = features['x'] - wx/2
    features['y_min'] = features['y'] - wy/2
    features['x_max'] = features['x'] + wx/2
    features['y_max'] = features['y'] + wy/2

    # Make figure
    fig, ax = plt.subplots(figsize=(5, 5))
    ax.plot(features['x'], features['y'], 'o', color='C1', ms=3, lw=0)
    for i, row in features.iterrows():
        if i not in overlapping:
            p = Rectangle((features.loc[i, 'x_min'], features.loc[i, 'y_min']),
                          wx, wy, facecolor='none', lw=1, edgecolor='C0')
        else:
            p = Rectangle((features.loc[i, 'x_min'], features.loc[i, 'y_min']),
                          wx, wy, facecolor='none', lw=1, edgecolor='C3')
        ax.add_patch(p)
        if plot_feature_index:
            ax.text(features.loc[i, 'x_min'], features.loc[i, 'y_max'],
                    f'{i}', ha='left', va='top')
    ax.grid(ls=':')
```

```python
plot_overlapping_features(features, overlapping, plot_feature_index=True)
```

## Timing comparison
---

```python
from tqdm.notebook import tqdm
from psf_extractor import remove_overlapping_features, remove_overlapping_features_bruteforce
```

### Bigger feature set

```python
# Set plane dimensions
Nx, Ny = (2000, 2000)

# Create a random distribution of features
N = 300
X = np.random.RandomState(15).uniform(0, Nx, size=N)
Y = np.random.RandomState(38).uniform(0, Ny, size=N)
features = pd.DataFrame({'x': X, 'y': Y})

# Create a bounding box for each bead
wx = 80
wy = 80
```

### Old method (brute-force)

```python
_, overlapping = remove_overlapping_features_bruteforce(features, wx=wx, wy=wy, return_indices=True)

out = f"Overlapping features... {len(overlapping)} / {N}"
print(out)
```

      0%|          | 0/44850 [00:00<?, ?it/s]

    Overlapping features... 250 / 300

```python
plot_overlapping_features(features, overlapping)
```

### New method (cell listing)

```python
_, overlapping = remove_overlapping_features(features, wx=wx, wy=wy, return_indices=True)

out = f"Overlapping features... {len(overlapping)} / {N}"
print(out)
```

      0%|          | 0/14 [00:00<?, ?it/s]

    Overlapping features... 250 / 300

```python
plot_overlapping_features(features, overlapping)
```


```python
import control
import sympy
import numpy as np
import matplotlib.pyplot as plt
sympy.init_printing()
plt.rcParams['figure.figsize']=(10, 5)
```

## Original System

```python
s = control.tf([1, 0], [0, 1])
```

```python
G = 1/((s+2)*s)
G
```

$$\frac{1}{s^2 + 2 s}$$

```python
control.rlocus(G, xlim=[-5, 0.1], ylim=[-5, 5]);
plt.title('uncompensated root locus')
sd = -2 + 2j

plt.plot(np.real(sd), np.imag(sd), 's')
plt.plot(np.real(sd), -np.imag(sd), 's')
```

```python
Gc = G/(1+G)
t, y = control.step_response(Gc/s, T=np.linspace(0, 20));
plt.plot(t, t)
plt.plot(t, y)
plt.title('uncompensated step response')

Kv = np.real((control.minreal(s*G).horner(0))[0, 0])
Kv

plt.figure()
t, y = control.step_response(Gc/s, T=np.linspace(0, 20));
plt.plot(t, t-y)
plt.hlines(1/Kv, 0, 20)
```

## Lead Compensator

```python
H_lead = 8*(s+2)/(s+4)
H_lead
```

$$\frac{8 s + 16}{s + 4}$$

```python
control.rlocus(G*H_lead, ylim=[-3, 3]);
plt.plot(-2, 2, 's')
plt.plot(-2, -2, 's')
plt.title('lead compensated root locus');
```

```python
Gc_lead = G*H_lead/(1+G*H_lead)
t = np.linspace(0, 10, 1000)
_, y = control.step_response(Gc_lead/s, T=t);
plt.plot(t, t)
plt.plot(t, y)
plt.title('lead compensated ramp response')

plt.figure()
_, y = control.step_response(Gc_lead/s, T=t);
plt.plot(t, t-y)
plt.hlines(0.5, 0, 20)
```

## Lead-Lag Compensator

```python
H_lag = (s+0.05)/(s+0.01)
H_lag
```

$$\frac{s + 0.05}{s + 0.01}$$

```python
H = H_lag*H_lead
```

```python
control.rlocus(G*H);
plt.plot(-2, 2, 's')
plt.plot(-2, -2, 's')
```

```python
control.minreal(G*H_lead*s).horner(0)
```

    2 states have been removed from the model

    array([[[2.+0.j]]])

```python
control.minreal(G*H*s).horner(0)
```

    2 states have been removed from the model

    array([[[10.+0.j]]])

```python
Gc_leadlag = G*H/(1+G*H)
t = np.linspace(0, 100, 1000)
_, y = control.step_response(Gc_leadlag/s, T=t);
plt.plot(t, t)
plt.plot(t, y)
plt.title('lead-lag compensated ramp response')

plt.figure()
plt.plot(t, t - y)
plt.hlines(0.1, t[0], t[-1])
plt.title('lead-lag compensated ramp response error')
```


### Usar este ejemplo como base para resolver el problema completo.

```python
from sympy import *
from sympy import init_printing; init_printing(use_latex = 'mathjax')
```

```python
var('x l m hbar')
```

$$\left ( x, \quad l, \quad m, \quad \hbar\right )$$

```python
f_1 = x*(l - x)
```

```python
f_1
```

$$x \left(l - x\right)$$

```python
diff(f_1, x, 2)
```

$$-2$$

```python
integrando = (-hbar**2/(2*m))*f_1*diff(f_1, x, 2)
integrando
```

$$\frac{\hbar^{2} x}{m} \left(l - x\right)$$

```python
H_11 = Integral(integrando, (x, 0, l))
```

```python
H_11.doit()
```

$$\frac{\hbar^{2} l^{3}}{6 m}$$

```python
s_11 = integrate(f_1**2, (x,0,l))
s_11
```

$$\frac{l^{5}}{30}$$

```python

```