# System Identification - NLARX

Silverbox refers to one of the nonlinear system identification benchmarks on http://nonlinearbenchmark.org/#Silverbox.
It is a simulation of a [Duffing oscillator](https://en.wikipedia.org/wiki/Duffing_equation), ocurring for instance in nonlinear spring pendulums.

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)$.

### Forecasting experiment

In this notebook, we will perform a forecasting experiment. The time-series will be split into a training and a validation part. In the training part, we will provide both input and output, and then infer parameters. In the forecasting part, we will only provide input. The parameters will be fixed to their inferred values and we will let the model make predictions for the output. Those predictions are evaluated against the true output.

### Data

Let's first have a look at the data.

```julia
using Revise
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])

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

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

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

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

```julia
# Plot every n-th time-point to avoid figure size exploding
n = 100

if viz
    p1a = Plots.plot(1:n:T_trn, input_trn[1:n:end], color="black", label="", markersize=2, xlabel="", ylabel="control", size=(1200,400), title="training")
    p1b = Plots.plot(1:n:T_trn, output_trn[1:n:end], color="black", label="", markersize=2, xlabel="time (t)", ylabel="output (displacement)", size=(1200,400))
    p1 = plot(p1a, p1b, layout=(2,1))
    Plots.savefig(p1, "viz/training_set.png")
end
```

```julia
if viz
    p2a = Plots.plot(1:n:T_val, input_val[1:n:end], color="black", label="", markersize=2, xlabel="", ylabel="control", size=(1200,400), title="validation")
    p2b = Plots.plot(1:n:T_val, output_val[1:n:end], color="black", label="", markersize=2, xlabel="time (t)", ylabel="output (displacement)", size=(1200,400))
    p2 = plot(p2a, p2b, layout=(2,1))
    Plots.savefig(p2, "viz/validation_set.png")
end
```

## 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 written 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}}_{e} g(z_{t-1}, \theta) + \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} \\
\theta_2 =&\ \frac{-b}{m+c} \\
\theta_3 =&\ \frac{-m}{m+c} \\
\eta =&\ \frac{1}{m+c} \\
\gamma^{-1} =& \tau^{-1} \frac{1}{(m+c)^2} \, ,
\end{align}$$

and $g(z_{t-1}, \theta) = \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. Variable substitution allows for more freedom in choosing priors.

The system is now a nonlinear autoregressive process:

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

where $f(z_{t-1}, \theta, \eta, u_t) = Sz_{t-1} + e g(z_{t-1}, \theta) + e \eta u_t$. Note that we need a two-dimensional state prior now (reminiscent of adding an initial condition on the velocity).

### 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(z_{t-1}, \theta, \eta, u_t), V) \\
y_t \sim&\ \mathcal{N}(e^{\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(x,\theta) = g(m_x, m_{\theta}) + J_{x}(m_x, m_{\theta})^{\top}(x - m_x) + J_{\theta}(m_x, m_{\theta})^{\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 you could argue that you don't need to approximate it. However, the first-order Taylor is exact in that case.

### 6. 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}) \\
\eta \sim&\ \text{Normal}(m^{0}_{\eta}, v^{0}_{\eta}) \\
\gamma \sim&\ \text{Gamma}(a^{0}_\gamma, b^{0}_\gamma) \, .
\end{align}$$

### 7. Recover physical variables from substituted ones

Since we have five equations and five unknowns, we can perfectly recover point estimates of the physical variables from substituted ones. However, we don't want point estimates, we want posteriors.

If we approximate $\gamma$ with a log-Normal distribution and then map it to a Gaussian distribution (i.e. model $\tilde{\gamma} = \log(\gamma)$), we end up with a Gaussian distributed random vector $\psi = [\theta_1, \theta_2, \theta_3, \eta, \tilde{\gamma}]$. We can perform a Gaussian approximation of the inverse mapping $G^{-1}(\psi)$ using a first-order Taylor expansion:

$$\begin{align}
m_{\phi} \triangleq \mathbb{E}[G^{-1}(\psi)] =&\ G^{-1}(m_{\psi}) \\
V_{\phi} \triangleq \mathbb{V}[G^{-1}(\psi)] =&\ J_{\psi}(m_{\psi}) G^{-1}(m_{\psi}) J_{\psi}(m_{\psi})^{\top} \, ,
\end{align}$$

The inverse mapping has the form:

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

The Jacobian can be obtained automatically using Julia packages such as ForwardDiff.jl or Zygote.jl.

## Implementation

We implemented the above solution with ForneyLab and a custom node called "NLARX". It contains a Nonlinear Latent Autoregressive model with eXogenous input (hence NLARX), that will model the state transition.

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

```julia
using NLARX
import NLARX: Jacobian
```

```julia
# Start 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
# function g(arg::Array{Float64,1}; num_coeffs::Int64=3)
#     θ = arg[1:num_coeffs]
#     x = arg[num_coeffs+1:end]
#     return θ[1]*x[1] + θ[2]*x[1]^3 + θ[3]*x[2]
# end
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)

# Inference algorithm
q1 = PosteriorFactorization(z_t, z_tmin1, θ, η, γ, ξ, ids=[:z_t, :z_tmin1, :θ, :η, :γ, :ξ])
algo1 = messagePassingAlgorithm([z_t, z_tmin1, θ, η, γ, ξ], q1, free_energy=false)
source_code1 = algorithmSourceCode(algo1, free_energy=false)
eval(Meta.parse(source_code1));
println(source_code1)
```

<?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="1427pt" height="629pt"
 viewBox="0.00 0.00 1427.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 1423,-625 1423,4 -4,4"/>
<!-- 2639251214800624216 -->
<g id="node1" class="node">
<title>2639251214800624216</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_m_η</text>
</g>
<!-- 3086109183766523893 -->
<g id="node2" class="node">
<title>3086109183766523893</title>
<polygon fill="none" stroke="#000000" points="284,-171 177,-171 177,-99 284,-99 284,-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">gaussianmeanprecision_2</text>
</g>
<!-- 3086109183766523893--2639251214800624216 -->
<g id="edge13" class="edge">
<title>3086109183766523893--2639251214800624216</title>
<path fill="none" stroke="#000000" d="M176.762,-104.7724C145.1446,-86.9876 105.8859,-64.9046 77.2895,-48.8191"/>
<text text-anchor="start" x="134.5" y="-74.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">m_η</text>
<text text-anchor="start" x="77.2895" y="-51.4191" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="161.762" y="-107.3724" font-family="Times,serif" font-size="8.00" fill="#000000">2 m </text>
</g>
<!-- 6969439574663989185 -->
<g id="node10" class="node">
<title>6969439574663989185</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_w_η</text>
</g>
<!-- 3086109183766523893--6969439574663989185 -->
<g id="edge4" class="edge">
<title>3086109183766523893--6969439574663989185</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">w_η</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">3 w </text>
</g>
<!-- 7697944740034904758 -->
<g id="node3" class="node">
<title>7697944740034904758</title>
<polygon fill="none" stroke="#000000" points="921.5,-405 849.5,-405 849.5,-333 921.5,-333 921.5,-405"/>
<text text-anchor="middle" x="885.5" y="-371.8" font-family="Times,serif" font-size="9.00" fill="#000000">dot</text>
<text text-anchor="middle" x="885.5" y="-361.8" font-family="Times,serif" font-size="9.00" fill="#000000">dotproduct_1</text>
</g>
<!-- 14220016728939314403 -->
<g id="node15" class="node">
<title>14220016728939314403</title>
<polygon fill="none" stroke="#000000" points="745.5,-288 673.5,-288 673.5,-216 745.5,-216 745.5,-288"/>
<text text-anchor="middle" x="709.5" y="-254.8" font-family="Times,serif" font-size="9.00" fill="#000000">NLARX</text>
<text text-anchor="middle" x="709.5" y="-244.8" font-family="Times,serif" font-size="9.00" fill="#000000">z_t</text>
</g>
<!-- 7697944740034904758--14220016728939314403 -->
<g id="edge1" class="edge">
<title>7697944740034904758--14220016728939314403</title>
<path fill="none" stroke="#000000" d="M849.4591,-345.041C819.1879,-324.9175 776.137,-296.2985 745.7947,-276.1277"/>
<text text-anchor="start" x="803.5" y="-308.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">z_t</text>
<text text-anchor="start" x="745.7947" y="-278.7277" font-family="Times,serif" font-size="8.00" fill="#000000">1 y </text>
<text text-anchor="start" x="830.4591" y="-347.641" font-family="Times,serif" font-size="8.00" fill="#000000">2 in1 </text>
</g>
<!-- 16777235638894054080 -->
<g id="node18" class="node">
<title>16777235638894054080</title>
<polygon fill="#d3d3d3" stroke="#000000" points="914.5,-279 860.5,-279 860.5,-225 914.5,-225 914.5,-279"/>
<text text-anchor="middle" x="887.5" y="-249.8" font-family="Times,serif" font-size="9.00" fill="#000000">clamp_1</text>
</g>
<!-- 7697944740034904758--16777235638894054080 -->
<g id="edge2" class="edge">
<title>7697944740034904758--16777235638894054080</title>
<path fill="none" stroke="#000000" d="M886.1164,-332.9415C886.4125,-315.6185 886.7615,-295.201 887.0349,-279.2064"/>
<text text-anchor="start" x="887.5" y="-308.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">clamp_1</text>
<text text-anchor="start" x="868.0349" y="-281.8064" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="867.1164" y="-326.5415" font-family="Times,serif" font-size="8.00" fill="#000000">3 in2 </text>
</g>
<!-- 13007345344252160785 -->
<g id="node4" class="node">
<title>13007345344252160785</title>
<polygon fill="#d3d3d3" stroke="#000000" points="1108,-279 1033,-279 1033,-225 1108,-225 1108,-279"/>
<text text-anchor="middle" x="1070.5" y="-249.8" font-family="Times,serif" font-size="9.00" fill="#000000">placeholder_b_ξ</text>
</g>
<!-- 17542400908337638584 -->
<g id="node5" class="node">
<title>17542400908337638584</title>
<polygon fill="#d3d3d3" stroke="#000000" points="1227,-54 1150,-54 1150,0 1227,0 1227,-54"/>
<text text-anchor="middle" x="1188.5" y="-24.8" font-family="Times,serif" font-size="9.00" fill="#000000">placeholder_w_θ</text>
</g>
<!-- 11344522920545235127 -->
<g id="node6" class="node">
<title>11344522920545235127</title>
<polygon fill="none" stroke="#000000" points="557.5,-171 485.5,-171 485.5,-99 557.5,-99 557.5,-171"/>
<text text-anchor="middle" x="521.5" y="-137.8" font-family="Times,serif" font-size="9.00" fill="#000000">𝒩</text>
<text text-anchor="middle" x="521.5" y="-127.8" font-family="Times,serif" font-size="9.00" fill="#000000">z_tmin1</text>
</g>
<!-- 16116054223711630911 -->
<g id="node9" class="node">
<title>16116054223711630911</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_z</text>
</g>
<!-- 11344522920545235127--16116054223711630911 -->
<g id="edge17" class="edge">
<title>11344522920545235127--16116054223711630911</title>
<path fill="none" stroke="#000000" d="M488.259,-98.7371C474.8899,-84.1526 459.7407,-67.6262 447.3551,-54.1146"/>
<text text-anchor="start" x="472.5" y="-74.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">w_z</text>
<text text-anchor="start" x="428.3551" y="-56.7146" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="473.259" y="-92.3371" font-family="Times,serif" font-size="8.00" fill="#000000">3 w </text>
</g>
<!-- 1308902856082378519 -->
<g id="node11" class="node">
<title>1308902856082378519</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_z</text>
</g>
<!-- 11344522920545235127--1308902856082378519 -->
<g id="edge3" class="edge">
<title>11344522920545235127--1308902856082378519</title>
<path fill="none" stroke="#000000" d="M552.7264,-98.7371C565.2853,-84.1526 579.5164,-67.6262 591.1513,-54.1146"/>
<text text-anchor="start" x="574.5" y="-74.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">m_z</text>
<text text-anchor="start" x="572.1513" y="-56.7146" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="537.7264" y="-92.3371" font-family="Times,serif" font-size="8.00" fill="#000000">2 m </text>
</g>
<!-- 13896333478482913612 -->
<g id="node7" class="node">
<title>13896333478482913612</title>
<polygon fill="#d3d3d3" stroke="#000000" points="1015,-621 942,-621 942,-567 1015,-567 1015,-621"/>
<text text-anchor="middle" x="978.5" y="-591.8" font-family="Times,serif" font-size="9.00" fill="#000000">placeholder_y_t</text>
</g>
<!-- 6378483230063897781 -->
<g id="node19" class="node">
<title>6378483230063897781</title>
<polygon fill="none" stroke="#000000" points="1014.5,-522 942.5,-522 942.5,-450 1014.5,-450 1014.5,-522"/>
<text text-anchor="middle" x="978.5" y="-488.8" font-family="Times,serif" font-size="9.00" fill="#000000">𝒩</text>
<text text-anchor="middle" x="978.5" y="-478.8" font-family="Times,serif" font-size="9.00" fill="#000000">y_t</text>
</g>
<!-- 13896333478482913612--6378483230063897781 -->
<g id="edge8" class="edge">
<title>13896333478482913612--6378483230063897781</title>
<path fill="none" stroke="#000000" d="M978.5,-566.747C978.5,-553.2495 978.5,-536.7693 978.5,-522.2253"/>
<text text-anchor="start" x="978.5" y="-542.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">y_t</text>
<text text-anchor="start" x="959.5" y="-524.8253" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="959.5" y="-560.347" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
</g>
<!-- 1402612114742527793 -->
<g id="node8" class="node">
<title>1402612114742527793</title>
<polygon fill="#d3d3d3" stroke="#000000" points="746,-162 673,-162 673,-108 746,-108 746,-162"/>
<text text-anchor="middle" x="709.5" y="-132.8" font-family="Times,serif" font-size="9.00" fill="#000000">placeholder_u_t</text>
</g>
<!-- 11759791787336724266 -->
<g id="node12" class="node">
<title>11759791787336724266</title>
<polygon fill="none" stroke="#000000" points="933.5,-171 861.5,-171 861.5,-99 933.5,-99 933.5,-171"/>
<text text-anchor="middle" x="897.5" y="-137.8" font-family="Times,serif" font-size="9.00" fill="#000000">Gam</text>
<text text-anchor="middle" x="897.5" y="-127.8" font-family="Times,serif" font-size="9.00" fill="#000000">gamma_1</text>
</g>
<!-- 14553083819736600033 -->
<g id="node14" class="node">
<title>14553083819736600033</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>
<!-- 11759791787336724266--14553083819736600033 -->
<g id="edge6" class="edge">
<title>11759791787336724266--14553083819736600033</title>
<path fill="none" stroke="#000000" d="M861.2163,-103.1067C854.0301,-96.119 846.7757,-88.5461 840.5,-81 833.6421,-72.7539 827.0058,-63.0987 821.3892,-54.2507"/>
<text text-anchor="start" x="840.5" y="-74.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">a_γ</text>
<text text-anchor="start" x="802.3892" y="-56.8507" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="848.2163" y="-105.7067" font-family="Times,serif" font-size="8.00" fill="#000000">2 a </text>
</g>
<!-- 275140831312593084 -->
<g id="node16" class="node">
<title>275140831312593084</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>
<!-- 11759791787336724266--275140831312593084 -->
<g id="edge11" class="edge">
<title>11759791787336724266--275140831312593084</title>
<path fill="none" stroke="#000000" d="M930.741,-98.7371C944.1101,-84.1526 959.2593,-67.6262 971.6449,-54.1146"/>
<text text-anchor="start" x="953.5" y="-74.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">b_γ</text>
<text text-anchor="start" x="952.6449" y="-56.7146" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="917.741" y="-92.3371" font-family="Times,serif" font-size="8.00" fill="#000000">3 b </text>
</g>
<!-- 18008076840631499642 -->
<g id="node13" class="node">
<title>18008076840631499642</title>
<polygon fill="#d3d3d3" stroke="#000000" points="1297.5,-279 1223.5,-279 1223.5,-225 1297.5,-225 1297.5,-279"/>
<text text-anchor="middle" x="1260.5" y="-249.8" font-family="Times,serif" font-size="9.00" fill="#000000">placeholder_a_ξ</text>
</g>
<!-- 14220016728939314403--3086109183766523893 -->
<g id="edge19" class="edge">
<title>14220016728939314403--3086109183766523893</title>
<path fill="none" stroke="#000000" d="M673.4971,-243.206C589.9494,-222.7987 382.5374,-172.1365 284.0531,-148.0808"/>
<text text-anchor="start" x="484.5" y="-191.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">η</text>
<text text-anchor="start" x="284.0531" y="-150.6808" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="660.4971" y="-245.806" font-family="Times,serif" font-size="8.00" fill="#000000">4 η </text>
</g>
<!-- 14220016728939314403--11344522920545235127 -->
<g id="edge10" class="edge">
<title>14220016728939314403--11344522920545235127</title>
<path fill="none" stroke="#000000" d="M673.2517,-229.4412C639.9612,-208.7232 590.8923,-178.1856 557.637,-157.4895"/>
<text text-anchor="start" x="621.5" y="-191.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">z_tmin1</text>
<text text-anchor="start" x="557.637" y="-160.0895" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="660.2517" y="-232.0412" font-family="Times,serif" font-size="8.00" fill="#000000">2 x </text>
</g>
<!-- 14220016728939314403--1402612114742527793 -->
<g id="edge5" class="edge">
<title>14220016728939314403--1402612114742527793</title>
<path fill="none" stroke="#000000" d="M709.5,-215.9415C709.5,-198.6185 709.5,-178.201 709.5,-162.2064"/>
<text text-anchor="start" x="709.5" y="-191.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">u_t</text>
<text text-anchor="start" x="690.5" y="-164.8064" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="696.5" y="-209.5415" font-family="Times,serif" font-size="8.00" fill="#000000">5 u </text>
</g>
<!-- 14220016728939314403--11759791787336724266 -->
<g id="edge14" class="edge">
<title>14220016728939314403--11759791787336724266</title>
<path fill="none" stroke="#000000" d="M745.7483,-229.4412C779.0388,-208.7232 828.1077,-178.1856 861.363,-157.4895"/>
<text text-anchor="start" x="809.5" y="-191.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">γ</text>
<text text-anchor="start" x="842.363" y="-160.0895" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="745.7483" y="-232.0412" font-family="Times,serif" font-size="8.00" fill="#000000">6 γ </text>
</g>
<!-- 8265780755932945244 -->
<g id="node20" class="node">
<title>8265780755932945244</title>
<polygon fill="none" stroke="#000000" points="1242,-171 1135,-171 1135,-99 1242,-99 1242,-171"/>
<text text-anchor="middle" x="1188.5" y="-137.8" font-family="Times,serif" font-size="9.00" fill="#000000">𝒩</text>
<text text-anchor="middle" x="1188.5" y="-127.8" font-family="Times,serif" font-size="9.00" fill="#000000">gaussianmeanprecision_1</text>
</g>
<!-- 14220016728939314403--8265780755932945244 -->
<g id="edge20" class="edge">
<title>14220016728939314403--8265780755932945244</title>
<path fill="none" stroke="#000000" d="M745.6864,-236.4155C762.7043,-229.4962 783.4102,-221.6569 802.5,-216 864.2946,-197.6883 1043.934,-162.5076 1134.6887,-145.1691"/>
<text text-anchor="start" x="913.5" y="-191.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">θ</text>
<text text-anchor="start" x="1115.6887" y="-147.7691" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="745.6864" y="-239.0155" font-family="Times,serif" font-size="8.00" fill="#000000">3 θ </text>
</g>
<!-- 14590098818513227933 -->
<g id="node17" class="node">
<title>14590098818513227933</title>
<polygon fill="none" stroke="#000000" points="1108.5,-405 1036.5,-405 1036.5,-333 1108.5,-333 1108.5,-405"/>
<text text-anchor="middle" x="1072.5" y="-371.8" font-family="Times,serif" font-size="9.00" fill="#000000">Gam</text>
<text text-anchor="middle" x="1072.5" y="-361.8" font-family="Times,serif" font-size="9.00" fill="#000000">gamma_2</text>
</g>
<!-- 14590098818513227933--13007345344252160785 -->
<g id="edge16" class="edge">
<title>14590098818513227933--13007345344252160785</title>
<path fill="none" stroke="#000000" d="M1071.8836,-332.9415C1071.5875,-315.6185 1071.2385,-295.201 1070.9651,-279.2064"/>
<text text-anchor="start" x="1072.5" y="-308.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">b_ξ</text>
<text text-anchor="start" x="1051.9651" y="-281.8064" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="1058.8836" y="-326.5415" font-family="Times,serif" font-size="8.00" fill="#000000">3 b </text>
</g>
<!-- 14590098818513227933--18008076840631499642 -->
<g id="edge15" class="edge">
<title>14590098818513227933--18008076840631499642</title>
<path fill="none" stroke="#000000" d="M1108.7483,-346.4412C1141.7187,-325.9224 1190.1655,-295.772 1223.3993,-275.0892"/>
<text text-anchor="start" x="1172.5" y="-308.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">a_ξ</text>
<text text-anchor="start" x="1204.3993" y="-277.6892" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="1108.7483" y="-349.0412" font-family="Times,serif" font-size="8.00" fill="#000000">2 a </text>
</g>
<!-- 6378483230063897781--7697944740034904758 -->
<g id="edge18" class="edge">
<title>6378483230063897781--7697944740034904758</title>
<path fill="none" stroke="#000000" d="M942.3926,-474.7478C922.1636,-466.3984 898.6288,-452.7983 886.5,-432 881.816,-423.9679 880.1322,-414.3243 879.9727,-405.0247"/>
<text text-anchor="start" x="886.5" y="-425.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">variable_1</text>
<text text-anchor="start" x="860.9727" y="-407.6247" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="927.3926" y="-477.3478" font-family="Times,serif" font-size="8.00" fill="#000000">2 m </text>
</g>
<!-- 6378483230063897781--14590098818513227933 -->
<g id="edge7" class="edge">
<title>6378483230063897781--14590098818513227933</title>
<path fill="none" stroke="#000000" d="M1007.4701,-449.9415C1018.9063,-435.7071 1032.0211,-419.3833 1043.4637,-405.1409"/>
<text text-anchor="start" x="1028.5" y="-425.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">ξ</text>
<text text-anchor="start" x="1024.4637" y="-407.7409" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="992.4701" y="-443.5415" font-family="Times,serif" font-size="8.00" fill="#000000">3 w </text>
</g>
<!-- 8265780755932945244--17542400908337638584 -->
<g id="edge9" class="edge">
<title>8265780755932945244--17542400908337638584</title>
<path fill="none" stroke="#000000" d="M1188.5,-98.7371C1188.5,-84.1526 1188.5,-67.6262 1188.5,-54.1146"/>
<text text-anchor="start" x="1188.5" y="-74.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">w_θ</text>
<text text-anchor="start" x="1169.5" y="-56.7146" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="1173.5" y="-92.3371" font-family="Times,serif" font-size="8.00" fill="#000000">3 w </text>
</g>
<!-- 7399300043099135832 -->
<g id="node21" class="node">
<title>7399300043099135832</title>
<polygon fill="#d3d3d3" stroke="#000000" points="1419,-54 1342,-54 1342,0 1419,0 1419,-54"/>
<text text-anchor="middle" x="1380.5" y="-24.8" font-family="Times,serif" font-size="9.00" fill="#000000">placeholder_m_θ</text>
</g>
<!-- 8265780755932945244--7399300043099135832 -->
<g id="edge12" class="edge">
<title>8265780755932945244--7399300043099135832</title>
<path fill="none" stroke="#000000" d="M1242.238,-104.7724C1273.8554,-86.9876 1313.1141,-64.9046 1341.7105,-48.8191"/>
<text text-anchor="start" x="1296.5" y="-74.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">m_θ</text>
<text text-anchor="start" x="1322.7105" y="-51.4191" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="1242.238" y="-107.3724" font-family="Times,serif" font-size="8.00" fill="#000000">2 m </text>
</g>
</g>
</svg>

    begin

    function stepγ!(data::Dict, marginals::Dict=Dict(), messages::Vector{Message}=Array{Message}(undef, 2))

    messages[1] = ruleVBGammaOut(nothing, ProbabilityDistribution(Univariate, PointMass, m=data[:a_γ]), ProbabilityDistribution(Univariate, PointMass, m=data[:b_γ]))
    messages[2] = ruleVariationalNLARXIn5PPPPPN(g, marginals[:z_t], marginals[:z_tmin1], marginals[:θ], marginals[:η], ProbabilityDistribution(Univariate, PointMass, m=data[:u_t]), nothing)

    marginals[:γ] = messages[1].dist * messages[2].dist

    return marginals

    end

    function stepξ!(data::Dict, marginals::Dict=Dict(), messages::Vector{Message}=Array{Message}(undef, 2))

    messages[1] = ruleVBGammaOut(nothing, ProbabilityDistribution(Univariate, PointMass, m=data[:a_ξ]), ProbabilityDistribution(Univariate, PointMass, m=data[:b_ξ]))
    messages[2] = ruleVBGaussianMeanPrecisionW(ProbabilityDistribution(Univariate, PointMass, m=data[:y_t]), marginals[:variable_1], nothing)

    marginals[:ξ] = messages[1].dist * messages[2].dist

    return marginals

    end

    function stepz_tmin1!(data::Dict, marginals::Dict=Dict(), messages::Vector{Message}=Array{Message}(undef, 2))

    messages[1] = ruleVBGaussianMeanPrecisionOut(nothing, ProbabilityDistribution(Multivariate, PointMass, m=data[:m_z]), ProbabilityDistribution(MatrixVariate, PointMass, m=data[:w_z]))
    messages[2] = ruleVariationalNLARXIn1PNPPPP(g, marginals[:z_t], nothing, marginals[:θ], marginals[:η], ProbabilityDistribution(Univariate, PointMass, m=data[:u_t]), marginals[:γ])

    marginals[:z_tmin1] = messages[1].dist * messages[2].dist

    return marginals

    end

    function stepz_t!(data::Dict, marginals::Dict=Dict(), messages::Vector{Message}=Array{Message}(undef, 4))

    messages[1] = ruleVariationalNLARXOutNPPPPP(g, nothing, marginals[:z_tmin1], marginals[:θ], marginals[:η], ProbabilityDistribution(Univariate, PointMass, m=data[:u_t]), marginals[:γ])
    messages[2] = ruleVBGaussianMeanPrecisionM(ProbabilityDistribution(Univariate, PointMass, m=data[:y_t]), nothing, marginals[:ξ])
    messages[3] = ruleSPDotProductIn1GNP(messages[2], nothing, Message(Multivariate, PointMass, m=[1.0, 0.0]))
    messages[4] = ruleSPDotProductOutNGP(nothing, messages[1], Message(Multivariate, PointMass, m=[1.0, 0.0]))

    marginals[:variable_1] = messages[4].dist * messages[2].dist
    marginals[:z_t] = messages[1].dist * messages[3].dist

    return marginals

    end

    function stepθ!(data::Dict, marginals::Dict=Dict(), messages::Vector{Message}=Array{Message}(undef, 2))

    messages[1] = ruleVBGaussianMeanPrecisionOut(nothing, ProbabilityDistribution(Multivariate, PointMass, m=data[:m_θ]), ProbabilityDistribution(MatrixVariate, PointMass, m=data[:w_θ]))
    messages[2] = ruleVariationalNLARXIn2PPNPPP(g, marginals[:z_t], marginals[:z_tmin1], nothing, marginals[:η], ProbabilityDistribution(Univariate, PointMass, m=data[:u_t]), marginals[:γ])

    marginals[:θ] = messages[1].dist * messages[2].dist

    return marginals

    end

    function stepη!(data::Dict, marginals::Dict=Dict(), messages::Vector{Message}=Array{Message}(undef, 2))

    messages[1] = ruleVBGaussianMeanPrecisionOut(nothing, ProbabilityDistribution(Univariate, PointMass, m=data[:m_η]), ProbabilityDistribution(Univariate, PointMass, m=data[:w_η]))
    messages[2] = ruleVariationalNLARXIn3PPPNPP(g, marginals[:z_t], marginals[:z_tmin1], marginals[:θ], nothing, ProbabilityDistribution(Univariate, PointMass, m=data[:u_t]), marginals[:γ])

    marginals[:η] = messages[1].dist * messages[2].dist

    return marginals

    end

    end # block

### Infer parameters

```julia
# Inference parameters
num_iterations = 10

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

# 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_η = (ones(1,T_trn+1), .1*ones(1,T_trn+1))
params_γ = (1e4*ones(1,T_trn+1), 1e1*ones(1,T_trn+1))
params_ξ = (1e2*ones(1,T_trn+1), 1e0*ones(1,T_trn+1))

# Initialize physical coefficient estimate arrays
params_ϕ = (zeros(5,T_trn), zeros(5,5,T_trn))
params_ψ = (zeros(5,T_trn), zeros(5,5,T_trn))

# Transformations between physical and substituted variables: ψ = G(ϕ) => ϕ = G_inv(ψ)
G(ϕ) = [(2*ϕ[1] + ϕ[2] - ϕ[3])/(ϕ[1]+ϕ[2]), -ϕ[4]/(ϕ[1]+ϕ[2]), -ϕ[1]/(ϕ[1]+ϕ[2]), 1/(ϕ[1]+ϕ[2]), ϕ[5]*(ϕ[1]+ϕ[2])^2]
G_inv(ψ) = [-ψ[3]/ψ[4], (1+ψ[3])/ψ[4], (1 - ψ[1] - ψ[3])/ψ[4], -ψ[2]/ψ[4], ψ[5]*ψ[4]^2]

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

    "Filtering"

    # 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

        stepz_t!(data, marginals)
        stepθ!(data, marginals)
        stepη!(data, marginals)
        stepγ!(data, marginals)
        stepξ!(data, marginals)
#         stepz_tmin1!(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]

    "Map substituted to physical variables via first-order Taylor"

    # Approximate gamma with log-normal via moment-matching
    Eγ = unsafeMean(marginals[:γ])
    Vγ = unsafeVar(marginals[:γ])
    m_γ = log(Eγ^2/sqrt(Vγ + Eγ^2))
    v_γ = log(Vγ/Eγ^2 + 1)

    # Construct vector of parameter estimates ψ
    m_ψ = [unsafeMean(marginals[:θ])[1], unsafeMean(marginals[:θ])[2], unsafeMean(marginals[:θ])[3], unsafeMean(marginals[:η])[1], m_γ]
    V_ψ = [unsafeCov(marginals[:θ]) zeros(3,2); zeros(2,3) [unsafeCov(marginals[:η])[1,1] 0;0 v_γ]]

    # Store psi
    params_ψ[1][:,t] = m_ψ
    params_ψ[2][:,:,t] = V_ψ

    # Compute Jacobian of transformation
    J_ψ = Jacobian(G_inv, m_ψ)

    # Compute moments of transformed Gaussian using first-order Taylor approx
    m_ϕ = G_inv(m_ψ)
    V_ϕ = J_ψ*V_ψ*J_ψ'

    # Store phi
    params_ϕ[1][:,t] = m_ϕ
    params_ϕ[2][:,:,t] = V_ϕ

end
```

    Progress: 100%|█████████████████████████████████████████| Time: 0:01:48m

```julia
# Extract mean of state marginals
estimated_states = params_z[1][1,:]

# Plot every n-th time-point to avoid figure size exploding
n = 50

p3 = Plots.scatter(1:n:T_trn, output_trn[1:n:T_trn], color="black", label="output", markersize=2, size=(700,400), xlabel="time (t)", ylabel="response", legend=:topleft)
Plots.plot!(1:n:T_trn, estimated_states[1:n:T_trn], color="red", linewidth=1, label="estimated")
```

### Forecast on validation data

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

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

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

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

# Inference algorithm
q2 = PosteriorFactorization(y_t, z_t, ids=[:yv, :zv])
algo2 = messagePassingAlgorithm([y_t, z_t], q2, free_energy=false)
source_code2 = algorithmSourceCode(algo2, free_energy=false)
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="831pt" height="512pt"
 viewBox="0.00 0.00 830.50 512.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 508)">
<title>G</title>
<polygon fill="#ffffff" stroke="transparent" points="-4,4 -4,-508 826.5,-508 826.5,4 -4,4"/>
<!-- 1299751371715971027 -->
<g id="node1" class="node">
<title>1299751371715971027</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>
<!-- 16992314313014175888 -->
<g id="node2" class="node">
<title>16992314313014175888</title>
<polygon fill="#d3d3d3" stroke="#000000" points="171,-279 105,-279 105,-225 171,-225 171,-279"/>
<text text-anchor="middle" x="138" y="-249.8" font-family="Times,serif" font-size="9.00" fill="#000000">placeholder_ξ</text>
</g>
<!-- 1347555674150171148 -->
<g id="node3" class="node">
<title>1347555674150171148</title>
<polygon fill="#d3d3d3" stroke="#000000" points="271,-54 181,-54 181,0 271,0 271,-54"/>
<text text-anchor="middle" x="226" y="-24.8" font-family="Times,serif" font-size="9.00" fill="#000000">placeholder_z_tmin1</text>
</g>
<!-- 16209196617579284110 -->
<g id="node4" class="node">
<title>16209196617579284110</title>
<polygon fill="#d3d3d3" stroke="#000000" points="452,-54 386,-54 386,0 452,0 452,-54"/>
<text text-anchor="middle" x="419" y="-24.8" font-family="Times,serif" font-size="9.00" fill="#000000">placeholder_θ</text>
</g>
<!-- 5301233074523136995 -->
<g id="node5" class="node">
<title>5301233074523136995</title>
<polygon fill="#d3d3d3" stroke="#000000" points="640.5,-54 567.5,-54 567.5,0 640.5,0 640.5,-54"/>
<text text-anchor="middle" x="604" y="-24.8" font-family="Times,serif" font-size="9.00" fill="#000000">placeholder_u_t</text>
</g>
<!-- 4599565527555012186 -->
<g id="node6" class="node">
<title>4599565527555012186</title>
<polygon fill="none" stroke="#000000" points="236,-405 164,-405 164,-333 236,-333 236,-405"/>
<text text-anchor="middle" x="200" y="-371.8" font-family="Times,serif" font-size="9.00" fill="#000000">𝒩</text>
<text text-anchor="middle" x="200" y="-361.8" font-family="Times,serif" font-size="9.00" fill="#000000">y_t</text>
</g>
<!-- 4599565527555012186--16992314313014175888 -->
<g id="edge2" class="edge">
<title>4599565527555012186--16992314313014175888</title>
<path fill="none" stroke="#000000" d="M163.8266,-340.5985C156.4535,-333.0663 149.5646,-324.3866 145,-315 139.658,-304.0147 137.6325,-290.7115 137.0518,-279.1657"/>
<text text-anchor="start" x="145" y="-308.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">ξ</text>
<text text-anchor="start" x="118.0518" y="-281.7657" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="148.8266" y="-343.1985" font-family="Times,serif" font-size="8.00" fill="#000000">3 w </text>
</g>
<!-- 8179469243439215026 -->
<g id="node7" class="node">
<title>8179469243439215026</title>
<polygon fill="none" stroke="#000000" points="358,-288 286,-288 286,-216 358,-216 358,-288"/>
<text text-anchor="middle" x="322" y="-254.8" font-family="Times,serif" font-size="9.00" fill="#000000">dot</text>
<text text-anchor="middle" x="322" y="-244.8" font-family="Times,serif" font-size="9.00" fill="#000000">dotproduct_1</text>
</g>
<!-- 4599565527555012186--8179469243439215026 -->
<g id="edge4" class="edge">
<title>4599565527555012186--8179469243439215026</title>
<path fill="none" stroke="#000000" d="M236.2737,-334.2129C251.8667,-319.259 270.0545,-301.8166 285.6554,-286.8551"/>
<text text-anchor="start" x="265" y="-308.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">variable_1</text>
<text text-anchor="start" x="266.6554" y="-289.4551" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="236.2737" y="-336.8129" font-family="Times,serif" font-size="8.00" fill="#000000">2 m </text>
</g>
<!-- 1913842529661092658 -->
<g id="node8" class="node">
<title>1913842529661092658</title>
<polygon fill="#d3d3d3" stroke="#000000" points="268,-162 214,-162 214,-108 268,-108 268,-162"/>
<text text-anchor="middle" x="241" y="-132.8" font-family="Times,serif" font-size="9.00" fill="#000000">clamp_1</text>
</g>
<!-- 8179469243439215026--1913842529661092658 -->
<g id="edge6" class="edge">
<title>8179469243439215026--1913842529661092658</title>
<path fill="none" stroke="#000000" d="M285.5556,-239.92C266.2511,-231.4476 244.2594,-218.0014 233,-198 226.9425,-187.2394 227.7166,-173.8358 230.5041,-162.1614"/>
<text text-anchor="start" x="233" y="-191.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">clamp_1</text>
<text text-anchor="start" x="211.5041" y="-164.7614" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="266.5556" y="-242.52" font-family="Times,serif" font-size="8.00" fill="#000000">3 in2 </text>
</g>
<!-- 12483274340840929235 -->
<g id="node10" class="node">
<title>12483274340840929235</title>
<polygon fill="none" stroke="#000000" points="455,-171 383,-171 383,-99 455,-99 455,-171"/>
<text text-anchor="middle" x="419" y="-137.8" font-family="Times,serif" font-size="9.00" fill="#000000">NLARX</text>
<text text-anchor="middle" x="419" y="-127.8" font-family="Times,serif" font-size="9.00" fill="#000000">z_t</text>
</g>
<!-- 8179469243439215026--12483274340840929235 -->
<g id="edge10" class="edge">
<title>8179469243439215026--12483274340840929235</title>
<path fill="none" stroke="#000000" d="M351.8947,-215.9415C363.6958,-201.7071 377.2292,-185.3833 389.037,-171.1409"/>
<text text-anchor="start" x="373" y="-191.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">z_t</text>
<text text-anchor="start" x="376.037" y="-173.7409" font-family="Times,serif" font-size="8.00" fill="#000000">1 y </text>
<text text-anchor="start" x="332.8947" y="-209.5415" font-family="Times,serif" font-size="8.00" fill="#000000">2 in1 </text>
</g>
<!-- 387448956254057190 -->
<g id="node9" class="node">
<title>387448956254057190</title>
<polygon fill="#d3d3d3" stroke="#000000" points="822.5,-54 755.5,-54 755.5,0 822.5,0 822.5,-54"/>
<text text-anchor="middle" x="789" y="-24.8" font-family="Times,serif" font-size="9.00" fill="#000000">placeholder_γ</text>
</g>
<!-- 12483274340840929235--1299751371715971027 -->
<g id="edge9" class="edge">
<title>12483274340840929235--1299751371715971027</title>
<path fill="none" stroke="#000000" d="M382.7842,-119.5136C365.7602,-112.6149 345.0575,-104.7647 326,-99 237.5459,-72.2434 212.4287,-77.2929 123,-54 104.0565,-49.0659 83.0971,-42.8249 66.1741,-37.5864"/>
<text text-anchor="start" x="253" y="-74.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">η</text>
<text text-anchor="start" x="66.1741" y="-40.1864" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="369.7842" y="-122.1136" font-family="Times,serif" font-size="8.00" fill="#000000">4 η </text>
</g>
<!-- 12483274340840929235--1347555674150171148 -->
<g id="edge8" class="edge">
<title>12483274340840929235--1347555674150171148</title>
<path fill="none" stroke="#000000" d="M382.6982,-114.686C351.1411,-97.0271 305.2838,-71.3661 271.3242,-52.3628"/>
<text text-anchor="start" x="323" y="-74.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">z_tmin1</text>
<text text-anchor="start" x="271.3242" y="-54.9628" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="369.6982" y="-108.286" font-family="Times,serif" font-size="8.00" fill="#000000">2 x </text>
</g>
<!-- 12483274340840929235--16209196617579284110 -->
<g id="edge1" class="edge">
<title>12483274340840929235--16209196617579284110</title>
<path fill="none" stroke="#000000" d="M419,-98.7371C419,-84.1526 419,-67.6262 419,-54.1146"/>
<text text-anchor="start" x="419" y="-74.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">θ</text>
<text text-anchor="start" x="400" y="-56.7146" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="406" y="-92.3371" font-family="Times,serif" font-size="8.00" fill="#000000">3 θ </text>
</g>
<!-- 12483274340840929235--5301233074523136995 -->
<g id="edge3" class="edge">
<title>12483274340840929235--5301233074523136995</title>
<path fill="none" stroke="#000000" d="M455.109,-113.9201C487.4962,-95.013 534.7978,-67.3991 567.3498,-48.3958"/>
<text text-anchor="start" x="522" y="-74.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">u_t</text>
<text text-anchor="start" x="548.3498" y="-50.9958" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="455.109" y="-116.5201" font-family="Times,serif" font-size="8.00" fill="#000000">5 u </text>
</g>
<!-- 12483274340840929235--387448956254057190 -->
<g id="edge7" class="edge">
<title>12483274340840929235--387448956254057190</title>
<path fill="none" stroke="#000000" d="M455.1357,-124.5824C508.3434,-109.2265 610.8872,-79.5685 698,-54 716.9748,-48.4307 738.1558,-42.15 755.3055,-37.0486"/>
<text text-anchor="start" x="629" y="-74.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">γ</text>
<text text-anchor="start" x="736.3055" y="-39.6486" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
<text text-anchor="start" x="455.1357" y="-127.1824" font-family="Times,serif" font-size="8.00" fill="#000000">6 γ </text>
</g>
<!-- 43705316805150733622 -->
<g id="node11" class="node">
<title>43705316805150733622</title>
</g>
<!-- 43705316805150733622--4599565527555012186 -->
<g id="edge5" class="edge">
<title>43705316805150733622--4599565527555012186</title>
<path fill="none" stroke="#000000" d="M200,-449.747C200,-436.2495 200,-419.7693 200,-405.2253"/>
<text text-anchor="start" x="200" y="-425.6" font-family="Times,serif" font-size="8.00" fill="#ff0000">y_t</text>
<text text-anchor="start" x="181" y="-407.8253" font-family="Times,serif" font-size="8.00" fill="#000000">1 out </text>
</g>
</g>
</svg>

    begin

    function stepyv!(data::Dict, marginals::Dict=Dict(), messages::Vector{Message}=Array{Message}(undef, 1))

    messages[1] = ruleVBGaussianMeanPrecisionOut(nothing, marginals[:variable_1], ProbabilityDistribution(Univariate, PointMass, m=data[:ξ]))

    marginals[:y_t] = messages[1].dist

    return marginals

    end

    function stepzv!(data::Dict, marginals::Dict=Dict(), messages::Vector{Message}=Array{Message}(undef, 4))

    messages[1] = ruleVariationalNLARXOutNPPPPP(g, nothing, ProbabilityDistribution(Multivariate, PointMass, m=data[:z_tmin1]), ProbabilityDistribution(Multivariate, PointMass, m=data[:θ]), ProbabilityDistribution(Univariate, PointMass, m=data[:η]), ProbabilityDistribution(Univariate, PointMass, m=data[:u_t]), ProbabilityDistribution(Univariate, PointMass, m=data[:γ]))
    messages[2] = ruleVBGaussianMeanPrecisionM(marginals[:y_t], nothing, ProbabilityDistribution(Univariate, PointMass, m=data[:ξ]))
    messages[3] = ruleSPDotProductIn1GNP(messages[2], nothing, Message(Multivariate, PointMass, m=[1.0, 0.0]))
    messages[4] = ruleSPDotProductOutNGP(nothing, messages[1], Message(Multivariate, PointMass, m=[1.0, 0.0]))

    marginals[:variable_1] = messages[4].dist * messages[2].dist
    marginals[:z_t] = messages[1].dist * messages[3].dist

    return marginals

    end

    end # block

```julia
# Initialize arrays of parameterizations
params_zv = (zeros(2,T_val), repeat(.1 .*Matrix{Float64}(I,2,2), outer=(1,1,T_val)))
params_yv = (zeros(1,T_val), zeros(1,T_val))

# State 1
params_zv[1][1,1] = output_val[1]
params_zv[1][2,1] = output_val[1]
params_zv[2][:,:,1] = 1e8*Matrix{Float64}(I,2,2)
params_zv[2][:,:,1] = 1e8*Matrix{Float64}(I,2,2)

# State 2
params_zv[1][1,2] = output_val[2]
params_zv[1][2,2] = output_val[1]
params_zv[2][:,:,2] = 1e8*Matrix{Float64}(I,2,2)
params_zv[2][:,:,2] = 1e8*Matrix{Float64}(I,2,2)

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

# Perform inference at each time-step
for t = 3:T_val

    # Update progress bar
    update!(p, t)

    # Initialize marginals
    marginals[:z_t] = ProbabilityDistribution(Multivariate, GaussianMeanPrecision, m=params_zv[1][:,t-1], w=params_zv[2][:,:,t-1])
    marginals[:y_t] = ProbabilityDistribution(Univariate, GaussianMeanPrecision, m=params_yv[1][1,t-1], w=params_yv[2][1,t-1])

    data = Dict(:u_t => input_val[t],
                :z_tmin1 => params_zv[1][:,t-1],
                :θ => params_θ[1][:,end],
                :η => params_η[1][end],
                :γ => params_γ[1][end]/params_γ[2][end],
                :ξ => params_ξ[1][end]/params_ξ[2][end])

    # Iterate variational parameter updates
    for i = 1:num_iterations
        stepzv!(data, marginals)
        stepyv!(data, marginals)
    end

    # Store current parameterizations of marginals
    params_zv[1][:,t] = unsafeMean(marginals[:z_t])
    params_zv[2][:,:,t] = marginals[:z_t].params[:w]
    params_yv[1][1,t] = unsafeMean(marginals[:y_t])
    params_yv[2][1,t] = marginals[:y_t].params[:w]

end
```

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

### Visualize forecasted states

```julia
# Extract mean of state marginals
predicted_states = params_yv[1][1,:]

# Plot every n-th time-point to avoid figure size exploding
n = 50

p3 = Plots.scatter(1:n:T_val, output_val[1:n:T_val], color="black", label="output", markersize=2, size=(700,400), xlabel="time (t)", ylabel="response", legend=:topleft)
Plots.plot!(1:n:T_val, predicted_states[1:n:T_val], color="red", linewidth=1, label="estimated")
```

```julia
Plots.savefig(p3, "viz/forecasted_states.png")
```

```julia
# Error plot
pred_error = (predicted_states .- output_val).^2
p4 = Plots.scatter(3:n:T_val, pred_error[3:n:end], color="black", label="", markersize=2, size=(700,400), xlabel="time (t)", yscale=:log10, ylabel="prediction error")
```

```julia
Plots.savefig(p4, "viz/forecasted_error.png")
```

```julia
println("MSE = "*string(mean(pred_error)))
```

    MSE = 0.0017565303307327284

### Baseline: neural network

We have implemented a neural network in Matlab as a baseline method.

```julia
using MAT
```

```julia
results = matread("../sota-baselines/training_results.mat")
```

    Dict{String,Any} with 10 entries:
      "options"    => Dict{String,Any}("viewnet"=>0.0,"nIter"=>1000.0,"feedthrough"…
      "details"    => Dict{String,Any}("trainingDetails"=>Dict{String,Any}("states"…
      "fs"         => 610.35
      "yVal"       => [0.0093978 0.012971 … -0.00605201 -0.00722471]
      "error"      => [0.00402748; 0.000125919; … ; 4.0117e-9; 3.17628e-8]
      "model"      => Dict{String,Any}("gradientFcn"=>"defaultderiv","derivFcn"=>"d…
      "yMod"       => [-0.0540646; 0.00174961; … ; -0.00611535; -0.00740293]
      ""           => UInt8[0x00 0x01 … 0x00 0x00]
      "timeTrain"  => [0.0 0.0016384 … 148.392 148.394]
      "modelParam" => Dict{String,Any}("nu"=>1.0,"nx"=>2.0,"transferFcn"=>"tansig",…

```julia
# Error plot
p5 = Plots.scatter(3:n:T_val, results["error"][3:n:end], color="black", label="", markersize=2, size=(700,400), xlabel="time (t)", yscale=:log10, ylabel="prediction error")
```

```julia
println("MSE = "*string(mean(results["error"])))
```

    MSE = 4.461002659364319e-7

## Visualize parameter estimates

```julia
# Extract mean of coefficient marginals
estimated_θ1_mean = params_θ[1][1,2:end]
estimated_θ2_mean = params_θ[1][2,2:end]
estimated_θ3_mean = params_θ[1][3,2:end]
estimated_θ1_std = sqrt.(inv.(params_θ[2][1,1,2:end]))
estimated_θ2_std = sqrt.(inv.(params_θ[2][2,2,2:end]))
estimated_θ3_std = sqrt.(inv.(params_θ[2][3,3,2:end]))

# Plot both coefficients next to each other
p2a = Plots.plot(1:n:T_trn, estimated_θ1_mean[1:n:T_trn], ribbon=[estimated_θ1_std[1:n:T_trn], estimated_θ1_std[1:n:T_trn]], color="red", label="θ_1", xlabel="time (t)", ylim=[1.45, 1.55], legend=:bottomright)
p2b = Plots.plot(1:n:T_trn, estimated_θ2_mean[1:n:T_trn], ribbon=[estimated_θ2_std[1:n:T_trn], estimated_θ2_std[1:n:T_trn]], color="blue", label="θ_2", xlabel="time (t)", legend=:topright)
p2c = Plots.plot(1:n:T_trn, estimated_θ3_mean[1:n:T_trn], ribbon=[estimated_θ3_std[1:n:T_trn], estimated_θ3_std[1:n:T_trn]], color="green", label="θ_3", xlabel="time (t)", ylim=[-1., -0.9], legend=:topright)
p2 = plot(p2a, p2b, p2c, layout=(3,1), size=(800,1200))
# Plots.savefig(p2, "viz/forecasted_estimatedθ.png")
```

```julia
p2a = Plots.plot(1:n:T_trn, estimated_θ1_mean[1:n:T_trn], ribbon=[estimated_θ1_std[1:n:T_trn], estimated_θ1_std[1:n:T_trn]], color="red", label="θ_1", xlabel="time (t)", ylim=[1.45, 1.55], legend=:bottomright, size=(600,400))
```

```julia
Plots.savefig(p2a, "viz/forecasted_estimatedθ1.png")
```

```julia
p2b = Plots.plot(1:n:T_trn, estimated_θ2_mean[1:n:T_trn], ribbon=[estimated_θ2_std[1:n:T_trn], estimated_θ2_std[1:n:T_trn]], color="blue", label="θ_2", xlabel="time (t)", legend=:topright, size=(600,400))
```

```julia
Plots.savefig(p2b, "viz/forecasted_estimatedθ2.png")
```

```julia
p2c = Plots.plot(1:n:T_trn, estimated_θ3_mean[1:n:T_trn], ribbon=[estimated_θ3_std[1:n:T_trn], estimated_θ3_std[1:n:T_trn]], color="green", label="θ_3", xlabel="time (t)", ylim=[-1., -0.9], legend=:topright, size=(600,400))
```

```julia
Plots.savefig(p2c, "viz/forecasted_estimatedθ3.png")
```

```julia
# Extract mean of control coefficient marginals
estimated_η_mean = params_η[1][1,2:end]
estimated_η_std = sqrt.(inv.(params_η[2][1,2:end]))

# Plot both coefficients next to each other
p3 = Plots.plot(1:n:T_trn, estimated_η_mean[1:n:T_trn], ribbon=[estimated_η_std[1:n:T_trn], estimated_η_std[1:n:T_trn]], color="blue", label="η", xlabel="time (t)", ylim=[0.15, .32], size=(600,400))
```

```julia
Plots.savefig(p3, "viz/forecasted_estimatedη.png")
```

```julia
# Extract mean of process precision marginals
estimated_γ_mean = params_γ[1][1,2:end] ./ params_γ[2][1,2:end]
estimated_γ_std = sqrt.(params_γ[1][1,2:end] ./ params_γ[2][1,2:end].^2)

# Plot both coefficients next to each other
p4 = Plots.plot(1:n:T_trn, estimated_γ_mean[1:n:T_trn], ribbon=[estimated_γ_std[1:n:T_trn], estimated_γ_std[1:n:T_trn]], color="blue", label="γ", xlabel="time (t)", size=(600,400), legend=:bottomright)
```

```julia
Plots.savefig(p4, "viz/forecasted_estimatedγ.png")
```

```julia
# Extract mean of measurement precision marginals
estimated_ξ_mean = params_ξ[1][1,2:end] ./ params_ξ[2][1,2:end]
estimated_ξ_std = sqrt.(params_ξ[1][1,2:end] ./ params_ξ[2][1,2:end].^2)

# Plot both coefficients next to each other
p8 = Plots.plot(1:n:T_trn, estimated_ξ_mean[1:n:T_trn], ribbon=[estimated_ξ_std[1:n:T_trn], estimated_ξ_std[1:n:T_trn]], color="blue", label="measurement-precision", xlabel="time (t)", size=(600,400), legend=:topleft)
```

```julia
Plots.savefig(p8, "viz/forecasted_measurement-precision.png")
```

```julia

```