```python
%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.figure import Figure
import ipywidgets as widgets
import time
from scipy import linalg
from scipy import optimize
import sympy as sm
import pandas as pd

sm.init_printing(use_unicode=True)
```

```python
mu_t = sm.symbols("mu_t")
mu_t1 = sm.symbols("mu_t-1")
b_t = sm.symbols("b_t")
a_t = sm.symbols("a_t")
p_t = sm.symbols("P_t")
v = sm.symbols("v")
v_h = sm.symbols("v^H")
v_l = sm.symbols("v^L")
d_t = sm.symbols("d_t")
sd_t = sm.symbols("s(d_t)")
s_a = sm.symbols("s_t^a")
s_b = sm.symbols("s_t^b")
theta_t = sm.symbols("theta_t")
theta_t1 = sm.symbols("theta_t-1")
pi = sm.symbols("pi")
beta_b = sm.symbols("beta_B")
beta_s = sm.symbols("beta_S")
spib = sm.symbols("Pi_buy")
spis = sm.symbols("Pi_sell")
pvh=sm.symbols("p(v_h)")
pb=sm.symbols("p(Buy)")
pbvh= sm.symbols('p(Buy|v_h)\space')
pbvl= sm.symbols('p(Buy|v_l)\space')
pvhb= sm.symbols('p(v_h|Buy)\space')
pvlb=sm.symbols('p(v_l|Buy)\space')
theta_p=sm.symbols("theta^p")
theta_m=sm.symbols("theta^m")
```

```python
sm.Eq(mu_t,theta_t*v_h+(1-theta_t)*v_l)
```

```python
sm.Eq(pbvh,(1-pi)*beta_b+pi)
```

$\displaystyle p(Buy|v_h)\space = \beta_{B} \left(1 - \pi\right) + \pi$

```python
sm.Eq(pbvl,(1-pi)*beta_b)
```

$\displaystyle p(Buy|v_l)\space = \beta_{B} \left(1 - \pi\right)$

```python
sm.Eq(pvhb,(pvh*pbvh)/(pb))
```

$\displaystyle p(v_h|Buy)\space = \frac{p(Buy|v_h)\space p(v_{h)}}{p(Buy)}$

```python
sm.Eq(pvhb, (theta_t1*((1-pi)*beta_b+pi))/((1-pi)*beta_b+pi*theta_t1))
```

$\displaystyle p(v_h|Buy)\space = \frac{\theta_{t-1} \left(\beta_{B} \left(1 - \pi\right) + \pi\right)}{\beta_{B} \left(1 - \pi\right) + \pi \theta_{t-1}}$

```python
sm.Eq(pvlb,1-pvhb)
```

$\displaystyle p(v_l|Buy)\space = 1 - p(v_h|Buy)\space$

```python
sm.Eq(a_t,pvhb*v_h+pvlb*v_l)
```

$\displaystyle a_{t} = p(v_h|Buy)\space v^{H} + p(v_l|Buy)\space v^{L}$

```python
sm.Eq(a_t,mu_t1+s_a)
```

```python
sm.Eq(b_t,mu_t1-s_b)
```

```python
sm.Eq(theta_p,((1+pi)*beta_b)/(pi*theta_t1+(1-pi)*beta_b))
```

```python
sm.Eq(theta_m,((1-pi)*beta_s)/(pi*(1-theta_t1)+(1-pi)*beta_s))
```

```python
minimering = abs(v-p_t)

spib = theta_t1*pi*(a_t-v_h)+beta_b*(1-pi)*(a_t-mu_t1)
spis = (1-theta_t1)*pi*(v_l-b_t)+beta_s*(1-pi)*(mu_t1-b_t)

a_t = mu_t1 + s_a
b_t = mu_t1 + s_b

s_a = (pi*theta_t1*(1-theta_t1))/(pi*theta_t1+(1-pi)*beta_b)*(v_h-v_l)
s_b = (pi*(1-theta_t1)*(1-theta_t1))/(pi*theta_t1+(1-pi)*beta_s)*(v_h-v_l)

```

```python
sm.Eq(sm.symbols("s_t^a"), s_a)
```

draw.io

```python
def simulation(vh=1, vl=0, decision="v_h", ratio=0.2, uninformed=0.5, startvalue=0.5, iterations = 500, seed=5000, epsilon=10**-5, shockperiod = None, shock={}):
    #define constants
    v_h = vh
    v_l = vl
    pi = ratio
    beta_b = uninformed
    beta_s = 1-beta_b
    shock = shock

    #determine realized value of v
    v = decision

    #allocate space
    values=[]
    ratiovalues = []
    iteration = []
    thetavalues = np.empty(iterations)
    muvalues = np.empty(iterations)
    askvalues = np.empty(iterations)
    bidvalues = np.empty(iterations)
    gapvalues = np.empty(iterations)
    pivalues = np.empty(iterations)
    decisionvalues = np.empty(iterations)

    #simulation settings
    thetavalues[0]=startvalue
    theta_t1 = startvalue
    N = iterations
    np.random.seed(seed)

    break_index = 0

    for i in range(N):
        if v=="v_h":
            v=v_h
        elif v=="v_l":
            v=v_l

        if i==shockperiod:
            if shock != {}:
                if "Private" in shock:
                    if shock["Private"]==1:
                        v="v_h"
                    if shock["Private"]==0:
                        v="v_l"

                if "Public" in shock:
                    v_l, v_h = shock["Public"]

        mu_t1 = theta_t1*v_h+(1-theta_t1)*v_l
        muvalues[i] = mu_t1

        #calculate markup/discount
        s_a = (pi*theta_t1*(1-theta_t1))/(pi*theta_t1+(1-pi)*beta_b)*(v_h-v_l)
        s_b = (pi*theta_t1*(1-theta_t1))/(pi*(1-theta_t1)+(1-pi)*beta_s)*(v_h-v_l)

        #calculate ask/bid price
        askvalues[i] = a_t = mu_t1 + s_a
        bidvalues[i] = b_t = mu_t1 - s_b

        #calculate gap
        gapvalues[i] = gap_t = a_t - b_t

        #realize pi
        pivalues[i] = trader = np.random.binomial(1,pi)

        #if trader is informed
        if trader == 1:
            if v == v_h:
                if v_h>a_t:
                    d_t=1
            if v == v_l:
                if v_l<b_t:
                    d_t=-1

        #if trader is uninformed
        if trader == 0:
            buysell = np.random.binomial(1,beta_b)
            if buysell == 1:
                d_t = 1
            else:
                d_t = -1

        decisionvalues[i] = d_t

        #update theta
        if d_t == 1:
            theta_t = ((1+pi)*beta_b)/(pi*theta_t1+(1-pi)*beta_b)*theta_t1
            theta_t1 = theta_t
        if d_t == -1:
            theta_t = ((1-pi)*beta_b)/(pi*(1-theta_t1)+(1-pi)*beta_b)*theta_t1
            theta_t1 = theta_t

        if i<iterations-1:
            thetavalues[i+1] = theta_t
        ratiovalues.append(str(ratio))
        iteration.append(int(i))

        #off by one error
        break_index=i+1
        if gap_t<epsilon or i == N-1:
            values.extend(["Theta: "+str(theta_t), "Bid: "+str(b_t), "Ask: "+str(a_t), "Mu: "+str(mu_t1)])
            break

    dataframe = pd.DataFrame()
    dataframe["Iteration"]=iteration
    dataframe["ratio"] = ratiovalues
    dataframe["theta"] = thetavalues[0:break_index]
    dataframe["mu"] = muvalues[0:break_index]
    dataframe["ask"] = askvalues[0:break_index]
    dataframe["bid"] = bidvalues[0:break_index]
    dataframe["spread"] = gapvalues[0:break_index]
    dataframe["trader"] = pivalues[0:break_index]
    dataframe["order"] = decisionvalues[0:break_index]

    return dataframe, values
```

```python
run1 = simulation(ratio=0.15, iterations=1001,seed=404, shockperiod = 90, shock={"Private": 1, "Public": (0.5, 0.7)})
run2 = simulation(ratio=0.1501, iterations=1001,seed=404, shockperiod = 90, shock={"Private": 0})
run3 = simulation(ratio=0.1502, iterations=1001, seed=404, shockperiod = 90, shock={"Private": 0,"Public": (0.5, 0.7)})
run4 = simulation(ratio=0.5, iterations=201, seed=404)
run5 = simulation(ratio=0.9, iterations=201, seed=404)

dataframe1, values1 = run1
dataframe2, values2 = run2
dataframe3, values3 = run3
dataframe4, values4 = run4
dataframe5, values5 = run5

merged = pd.concat([dataframe1,dataframe2,dataframe3,dataframe4,dataframe5])
```

```python
print(values1,values2,values3)
```

    ['Theta: 0.9999148241152956', 'Bid: 0.6999829648230592', 'Ask: 0.6999906930881101', 'Mu: 0.6999874085024806'] ['Theta: 1.018638663218169e-05', 'Bid: 1.0186386632181688e-05', 'Ask: 1.865314333398436e-05', 'Mu: 1.3784352410597332e-05'] ['Theta: 4.591039783417812e-05', 'Bid: 0.5000091820795668', 'Ask: 0.5000168204549132', 'Mu: 0.500012427696542']

```python
fig = plt.figure(dpi=100)
ax = fig.add_subplot(1,1,1)
ax.plot(dataframe1["ask"])
ax.plot(dataframe1["bid"])
ax.plot(dataframe1["mu"])
ax.grid(True)
ax.set_xlabel("Iterations")
ax.legend()

fig1 = plt.figure(dpi=100)
ax1 = fig1.add_subplot(1,1,1)
ax1.plot(dataframe2["ask"])
ax1.plot(dataframe2["bid"])
ax1.plot(dataframe2["mu"])
ax1.grid(True)
ax1.set_xlabel("Iterations")
ax1.legend()
```

```python
run6 = simulation(ratio=0.15, iterations=201,seed=4795)
run7 = simulation(ratio=0.15, iterations=201,seed=6279)
run8 = simulation(ratio=0.15, iterations=201,seed=6130)
run9 = simulation(ratio=0.15, iterations=201,seed=9352)
run10 = simulation(ratio=0.15, iterations=201,seed=5059)

dataframe6, values6 = run6
dataframe7, values7 = run7
dataframe8, values8 = run8
dataframe9, values9 = run9
dataframe10, values10 = run10
```

```python
fig = plt.figure(dpi=100)
ax2 = fig.add_subplot(1,1,1)
ax2.plot(dataframe6["theta"], label="6")
ax2.plot(dataframe7["theta"], label="7")
ax2.plot(dataframe8["theta"], label="8")
ax2.plot(dataframe9["theta"], label="9")
ax2.plot(dataframe10["theta"], label="10")
ax2.grid(True)
ax2.legend()
ax2.set_title("Theta")
```

```python
from plotter import PlotterWindow
```

```python
graphwindow = PlotterWindow(data = merged, slicename = "ratio", xvariable = "Iteration", yvariablelist = ["mu","ask","bid"], lineplots=True)
graphwindow.start()
```

    C:\Users\Corfixen\Anaconda3\lib\site-packages\matplotlib\figure.py:98: MatplotlibDeprecationWarning:
    Adding an axes using the same arguments as a previous axes currently reuses the earlier instance.  In a future version, a new instance will always be created and returned.  Meanwhile, this warning can be suppressed, and the future behavior ensured, by passing a unique label to each axes instance.
      "Adding an axes using the same arguments as a previous axes "

```python
merged.head(20)
```

      Iteration
      ratio
      theta
      mu
      ask
      bid
      spread
      trader
      order

      0
      0
      0.15
      0.500000
      0.500000
      0.575000
      0.425000
      0.150000
      0.0
      -1.0

      1
      1
      0.15
      0.425000
      0.425000
      0.500000
      0.353301
      0.146699
      1.0
      1.0

      2
      2
      0.15
      0.500000
      0.500000
      0.575000
      0.425000
      0.150000
      0.0
      -1.0

      3
      3
      0.15
      0.425000
      0.425000
      0.500000
      0.353301
      0.146699
      1.0
      1.0

      4
      4
      0.15
      0.500000
      0.500000
      0.575000
      0.425000
      0.150000
      0.0
      -1.0

      5
      5
      0.15
      0.425000
      0.425000
      0.500000
      0.353301
      0.146699
      0.0
      1.0

      6
      6
      0.15
      0.500000
      0.500000
      0.575000
      0.425000
      0.150000
      0.0
      -1.0

      7
      7
      0.15
      0.425000
      0.425000
      0.500000
      0.353301
      0.146699
      0.0
      1.0

      8
      8
      0.15
      0.500000
      0.500000
      0.575000
      0.425000
      0.150000
      0.0
      1.0

      9
      9
      0.15
      0.575000
      0.575000
      0.646699
      0.500000
      0.146699
      0.0
      1.0

      10
      10
      0.15
      0.646699
      0.646699
      0.712354
      0.575000
      0.137354
      0.0
      1.0

      11
      11
      0.15
      0.712354
      0.712354
      0.770144
      0.646699
      0.123445
      0.0
      1.0

      12
      12
      0.15
      0.770144
      0.770144
      0.819269
      0.712354
      0.106916
      1.0
      1.0

      13
      13
      0.15
      0.819269
      0.819269
      0.859807
      0.770144
      0.089663
      0.0
      -1.0

      14
      14
      0.15
      0.770144
      0.770144
      0.819269
      0.712354
      0.106916
      0.0
      -1.0

      15
      15
      0.15
      0.712354
      0.712354
      0.770144
      0.646699
      0.123445
      0.0
      1.0

      16
      16
      0.15
      0.770144
      0.770144
      0.819269
      0.712354
      0.106916
      0.0
      1.0

      17
      17
      0.15
      0.819269
      0.819269
      0.859807
      0.770144
      0.089663
      0.0
      1.0

      18
      18
      0.15
      0.859807
      0.859807
      0.892445
      0.819269
      0.073176
      0.0
      1.0

      19
      19
      0.15
      0.892445
      0.892445
      0.918208
      0.859807
      0.058402
      0.0
      1.0

```python
from mpl_finance import candlestick_ochl
```

```python
mu_t
```

```python
fig = plt.figure(dpi=100)
ax = fig.add_subplot(1,1,1)

lower = 0
upper = len(list(dataframe1["bid"]))
ochl = []

iteration = list(dataframe1["Iteration"])
bid = list(dataframe1["bid"])
ask = list(dataframe1["ask"])

for lower in range(upper):
    append = iteration[lower], bid[lower], ask[lower], ask[lower], bid[lower]
    ochl.append(append)
candlestick_ochl(ax, ochl, colorup="g",colordown="r")
```

```python
candlestick_ochl(ax, ochl)
```

    ([<matplotlib.lines.Line2D at 0x1a9b49d3f60>,
      <matplotlib.lines.Line2D at 0x1a9b59b54e0>,
      <matplotlib.lines.Line2D at 0x1a9b5a09b00>,
      <matplotlib.lines.Line2D at 0x1a9b5a52a90>,
      <matplotlib.lines.Line2D at 0x1a9b59a2198>,
      <matplotlib.lines.Line2D at 0x1a9b5807278>,
      <matplotlib.lines.Line2D at 0x1a9b57c6160>,
      <matplotlib.lines.Line2D at 0x1a9b57c6940>,
      <matplotlib.lines.Line2D at 0x1a9b57cbeb8>,
      <matplotlib.lines.Line2D at 0x1a9b5a1c518>,
      <matplotlib.lines.Line2D at 0x1a9b4cfdbe0>,
      <matplotlib.lines.Line2D at 0x1a9b4d486d8>,
      <matplotlib.lines.Line2D at 0x1a9b4d484a8>,
      <matplotlib.lines.Line2D at 0x1a9b4d41470>,
      <matplotlib.lines.Line2D at 0x1a9b4d41048>,
      <matplotlib.lines.Line2D at 0x1a9b4cc01d0>,
      <matplotlib.lines.Line2D at 0x1a9b4cc0710>,
      <matplotlib.lines.Line2D at 0x1a9b4a59ef0>,
      <matplotlib.lines.Line2D at 0x1a9b4a59dd8>,
      <matplotlib.lines.Line2D at 0x1a9b4987668>,
      <matplotlib.lines.Line2D at 0x1a9b4d2a668>,
      <matplotlib.lines.Line2D at 0x1a9b4d2a7b8>,
      <matplotlib.lines.Line2D at 0x1a9b4d19eb8>,
      <matplotlib.lines.Line2D at 0x1a9b4d19390>,
      <matplotlib.lines.Line2D at 0x1a9b4d7fc18>,
      <matplotlib.lines.Line2D at 0x1a9b4d8af98>,
      <matplotlib.lines.Line2D at 0x1a9b4a17198>,
      <matplotlib.lines.Line2D at 0x1a9b4a17908>,
      <matplotlib.lines.Line2D at 0x1a9af69ef28>,
      <matplotlib.lines.Line2D at 0x1a9b4961c50>,
      <matplotlib.lines.Line2D at 0x1a9b493aa58>,
      <matplotlib.lines.Line2D at 0x1a9b49e0d30>,
      <matplotlib.lines.Line2D at 0x1a9b49e0358>,
      <matplotlib.lines.Line2D at 0x1a9b4973668>,
      <matplotlib.lines.Line2D at 0x1a9b49732e8>,
      <matplotlib.lines.Line2D at 0x1a9b4952588>,
      <matplotlib.lines.Line2D at 0x1a9b4952668>,
      <matplotlib.lines.Line2D at 0x1a9b490f6d8>,
      <matplotlib.lines.Line2D at 0x1a9afab36d8>,
      <matplotlib.lines.Line2D at 0x1a9b4a2fe10>,
      <matplotlib.lines.Line2D at 0x1a9b4a2f668>,
      <matplotlib.lines.Line2D at 0x1a9b4a27048>,
      <matplotlib.lines.Line2D at 0x1a9b4a27b70>,
      <matplotlib.lines.Line2D at 0x1a9b4a27128>,
      <matplotlib.lines.Line2D at 0x1a9b5ade668>,
      <matplotlib.lines.Line2D at 0x1a9b5adecc0>,
      <matplotlib.lines.Line2D at 0x1a9b5b0a358>,
      <matplotlib.lines.Line2D at 0x1a9b5b0a9b0>,
      <matplotlib.lines.Line2D at 0x1a9b5b0afd0>,
      <matplotlib.lines.Line2D at 0x1a9b5b136a0>,
      <matplotlib.lines.Line2D at 0x1a9b5b13cf8>,
      <matplotlib.lines.Line2D at 0x1a9b5b1c390>,
      <matplotlib.lines.Line2D at 0x1a9b5b1c9e8>,
      <matplotlib.lines.Line2D at 0x1a9b5b1ceb8>,
      <matplotlib.lines.Line2D at 0x1a9b5b266d8>,
      <matplotlib.lines.Line2D at 0x1a9b5b26d30>,
      <matplotlib.lines.Line2D at 0x1a9b5b313c8>,
      <matplotlib.lines.Line2D at 0x1a9b5b31a20>,
      <matplotlib.lines.Line2D at 0x1a9b5b31ef0>,
      <matplotlib.lines.Line2D at 0x1a9b5b3a710>,
      <matplotlib.lines.Line2D at 0x1a9b5b3ad68>,
      <matplotlib.lines.Line2D at 0x1a9b5b43400>,
      <matplotlib.lines.Line2D at 0x1a9b5b43a58>,
      <matplotlib.lines.Line2D at 0x1a9b5b43f28>,
      <matplotlib.lines.Line2D at 0x1a9b5b4b748>,
      <matplotlib.lines.Line2D at 0x1a9b5b4bda0>,
      <matplotlib.lines.Line2D at 0x1a9b5b56438>,
      <matplotlib.lines.Line2D at 0x1a9b5b56a90>,
      <matplotlib.lines.Line2D at 0x1a9b5b56f60>,
      <matplotlib.lines.Line2D at 0x1a9b5b60780>,
      <matplotlib.lines.Line2D at 0x1a9b5b60dd8>,
      <matplotlib.lines.Line2D at 0x1a9b5b6a470>,
      <matplotlib.lines.Line2D at 0x1a9b5b6aac8>,
      <matplotlib.lines.Line2D at 0x1a9b5b6af98>,
      <matplotlib.lines.Line2D at 0x1a9b5b747b8>,
      <matplotlib.lines.Line2D at 0x1a9b5b74e10>,
      <matplotlib.lines.Line2D at 0x1a9b5b7d4a8>,
      <matplotlib.lines.Line2D at 0x1a9b5b7db00>,
      <matplotlib.lines.Line2D at 0x1a9b5b7dfd0>,
      <matplotlib.lines.Line2D at 0x1a9b5b877f0>,
      <matplotlib.lines.Line2D at 0x1a9b5b87e48>,
      <matplotlib.lines.Line2D at 0x1a9b5b904e0>,
      <matplotlib.lines.Line2D at 0x1a9b5b90b38>,
      <matplotlib.lines.Line2D at 0x1a9b5b90fd0>,
      <matplotlib.lines.Line2D at 0x1a9b5b9a828>,
      <matplotlib.lines.Line2D at 0x1a9b5b9ae80>,
      <matplotlib.lines.Line2D at 0x1a9b5ba3518>,
      <matplotlib.lines.Line2D at 0x1a9b5ba3b70>,
      <matplotlib.lines.Line2D at 0x1a9b5bad208>,
      <matplotlib.lines.Line2D at 0x1a9b5bad860>,
      <matplotlib.lines.Line2D at 0x1a9b5badeb8>,
      <matplotlib.lines.Line2D at 0x1a9b5bb6550>,
      <matplotlib.lines.Line2D at 0x1a9b5bb6ba8>,
      <matplotlib.lines.Line2D at 0x1a9b5bc2240>,
      <matplotlib.lines.Line2D at 0x1a9b5bc2898>,
      <matplotlib.lines.Line2D at 0x1a9b5bc2ef0>,
      <matplotlib.lines.Line2D at 0x1a9b5bcd588>,
      <matplotlib.lines.Line2D at 0x1a9b5bcdbe0>,
      <matplotlib.lines.Line2D at 0x1a9b5bd6278>,
      <matplotlib.lines.Line2D at 0x1a9b5bd68d0>,
      <matplotlib.lines.Line2D at 0x1a9b5bd6f28>,
      <matplotlib.lines.Line2D at 0x1a9b5be15c0>,
      <matplotlib.lines.Line2D at 0x1a9b5be1c18>,
      <matplotlib.lines.Line2D at 0x1a9b5be82b0>,
      <matplotlib.lines.Line2D at 0x1a9b5be8908>,
      <matplotlib.lines.Line2D at 0x1a9b5be8f60>,
      <matplotlib.lines.Line2D at 0x1a9b5bf35f8>,
      <matplotlib.lines.Line2D at 0x1a9b5bf3c50>,
      <matplotlib.lines.Line2D at 0x1a9b5bfc2e8>,
      <matplotlib.lines.Line2D at 0x1a9b5bfc940>,
      <matplotlib.lines.Line2D at 0x1a9b5bfcf98>,
      <matplotlib.lines.Line2D at 0x1a9b5c08630>,
      <matplotlib.lines.Line2D at 0x1a9b5c08c88>,
      <matplotlib.lines.Line2D at 0x1a9b5c13320>,
      <matplotlib.lines.Line2D at 0x1a9b5c13978>,
      <matplotlib.lines.Line2D at 0x1a9b5c13fd0>,
      <matplotlib.lines.Line2D at 0x1a9b5c1c668>,
      <matplotlib.lines.Line2D at 0x1a9b5c1ccc0>,
      <matplotlib.lines.Line2D at 0x1a9b5c26358>,
      <matplotlib.lines.Line2D at 0x1a9b5c269b0>,
      <matplotlib.lines.Line2D at 0x1a9b5c26fd0>,
      <matplotlib.lines.Line2D at 0x1a9b5c2f6a0>,
      <matplotlib.lines.Line2D at 0x1a9b5c2fcf8>,
      <matplotlib.lines.Line2D at 0x1a9b5c39390>,
      <matplotlib.lines.Line2D at 0x1a9b5c399e8>,
      <matplotlib.lines.Line2D at 0x1a9b5c39eb8>,
      <matplotlib.lines.Line2D at 0x1a9b5c426d8>,
      <matplotlib.lines.Line2D at 0x1a9b5c42d30>,
      <matplotlib.lines.Line2D at 0x1a9b5c4d3c8>,
      <matplotlib.lines.Line2D at 0x1a9b5c4da20>,
      <matplotlib.lines.Line2D at 0x1a9b5c4def0>,
      <matplotlib.lines.Line2D at 0x1a9b5c57710>,
      <matplotlib.lines.Line2D at 0x1a9b5c57d68>,
      <matplotlib.lines.Line2D at 0x1a9b5c63400>,
      <matplotlib.lines.Line2D at 0x1a9b5c63a58>,
      <matplotlib.lines.Line2D at 0x1a9b5c63f28>,
      <matplotlib.lines.Line2D at 0x1a9b5c6c748>,
      <matplotlib.lines.Line2D at 0x1a9b5c6cda0>,
      <matplotlib.lines.Line2D at 0x1a9b5c75438>,
      <matplotlib.lines.Line2D at 0x1a9b5c75a90>,
      <matplotlib.lines.Line2D at 0x1a9b5c75f60>,
      <matplotlib.lines.Line2D at 0x1a9b5c7e780>,
      <matplotlib.lines.Line2D at 0x1a9b5c7edd8>],
     [<matplotlib.patches.Rectangle at 0x1a9b5a129e8>,
      <matplotlib.patches.Rectangle at 0x1a9b5999ac8>,
      <matplotlib.patches.Rectangle at 0x1a9b59e9da0>,
      <matplotlib.patches.Rectangle at 0x1a9b5a70be0>,
      <matplotlib.patches.Rectangle at 0x1a9b58074a8>,
      <matplotlib.patches.Rectangle at 0x1a9b58071d0>,
      <matplotlib.patches.Rectangle at 0x1a9b57c6400>,
      <matplotlib.patches.Rectangle at 0x1a9b57c6898>,
      <matplotlib.patches.Rectangle at 0x1a9b57cb828>,
      <matplotlib.patches.Rectangle at 0x1a9b59c1710>,
      <matplotlib.patches.Rectangle at 0x1a9b598f8d0>,
      <matplotlib.patches.Rectangle at 0x1a9b4d48438>,
      <matplotlib.patches.Rectangle at 0x1a9b4d41e48>,
      <matplotlib.patches.Rectangle at 0x1a9b4d41518>,
      <matplotlib.patches.Rectangle at 0x1a9b4cc0908>,
      <matplotlib.patches.Rectangle at 0x1a9b4cc0390>,
      <matplotlib.patches.Rectangle at 0x1a9b4a59160>,
      <matplotlib.patches.Rectangle at 0x1a9b4a59978>,
      <matplotlib.patches.Rectangle at 0x1a9b5a7bc88>,
      <matplotlib.patches.Rectangle at 0x1a9b4d2a160>,
      <matplotlib.patches.Rectangle at 0x1a9b4d2acf8>,
      <matplotlib.patches.Rectangle at 0x1a9b4d19780>,
      <matplotlib.patches.Rectangle at 0x1a9b4d19588>,
      <matplotlib.patches.Rectangle at 0x1a9b4da8f98>,
      <matplotlib.patches.Rectangle at 0x1a9b4a05a90>,
      <matplotlib.patches.Rectangle at 0x1a9b4a174e0>,
      <matplotlib.patches.Rectangle at 0x1a9b4a17748>,
      <matplotlib.patches.Rectangle at 0x1a9b49bceb8>,
      <matplotlib.patches.Rectangle at 0x1a9b4945128>,
      <matplotlib.patches.Rectangle at 0x1a9b49614e0>,
      <matplotlib.patches.Rectangle at 0x1a9b49e0b38>,
      <matplotlib.patches.Rectangle at 0x1a9b49e0b70>,
      <matplotlib.patches.Rectangle at 0x1a9b49731d0>,
      <matplotlib.patches.Rectangle at 0x1a9b4973160>,
      <matplotlib.patches.Rectangle at 0x1a9b494a940>,
      <matplotlib.patches.Rectangle at 0x1a9b4952f98>,
      <matplotlib.patches.Rectangle at 0x1a9b494d668>,
      <matplotlib.patches.Rectangle at 0x1a9b4914f28>,
      <matplotlib.patches.Rectangle at 0x1a9b4a2ff28>,
      <matplotlib.patches.Rectangle at 0x1a9b4a2fc50>,
      <matplotlib.patches.Rectangle at 0x1a9b4a2f390>,
      <matplotlib.patches.Rectangle at 0x1a9b4a27c88>,
      <matplotlib.patches.Rectangle at 0x1a9b4a27828>,
      <matplotlib.patches.Rectangle at 0x1a9b5ade128>,
      <matplotlib.patches.Rectangle at 0x1a9b5ade828>,
      <matplotlib.patches.Rectangle at 0x1a9b5adee80>,
      <matplotlib.patches.Rectangle at 0x1a9b5b0a518>,
      <matplotlib.patches.Rectangle at 0x1a9b5b0ab70>,
      <matplotlib.patches.Rectangle at 0x1a9b5b13208>,
      <matplotlib.patches.Rectangle at 0x1a9b5b13860>,
      <matplotlib.patches.Rectangle at 0x1a9b5b13eb8>,
      <matplotlib.patches.Rectangle at 0x1a9b5b1c550>,
      <matplotlib.patches.Rectangle at 0x1a9b5b1cba8>,
      <matplotlib.patches.Rectangle at 0x1a9b5b26240>,
      <matplotlib.patches.Rectangle at 0x1a9b5b26898>,
      <matplotlib.patches.Rectangle at 0x1a9b5b26ef0>,
      <matplotlib.patches.Rectangle at 0x1a9b5b31588>,
      <matplotlib.patches.Rectangle at 0x1a9b5b31be0>,
      <matplotlib.patches.Rectangle at 0x1a9b5b3a278>,
      <matplotlib.patches.Rectangle at 0x1a9b5b3a8d0>,
      <matplotlib.patches.Rectangle at 0x1a9b5b3af28>,
      <matplotlib.patches.Rectangle at 0x1a9b5b435c0>,
      <matplotlib.patches.Rectangle at 0x1a9b5b43c18>,
      <matplotlib.patches.Rectangle at 0x1a9b5b4b2b0>,
      <matplotlib.patches.Rectangle at 0x1a9b5b4b908>,
      <matplotlib.patches.Rectangle at 0x1a9b5b4bf60>,
      <matplotlib.patches.Rectangle at 0x1a9b5b565f8>,
      <matplotlib.patches.Rectangle at 0x1a9b5b56c50>,
      <matplotlib.patches.Rectangle at 0x1a9b5b602e8>,
      <matplotlib.patches.Rectangle at 0x1a9b5b60940>,
      <matplotlib.patches.Rectangle at 0x1a9b5b60f98>,
      <matplotlib.patches.Rectangle at 0x1a9b5b6a630>,
      <matplotlib.patches.Rectangle at 0x1a9b5b6ac88>,
      <matplotlib.patches.Rectangle at 0x1a9b5b74320>,
      <matplotlib.patches.Rectangle at 0x1a9b5b74978>,
      <matplotlib.patches.Rectangle at 0x1a9b5b74fd0>,
      <matplotlib.patches.Rectangle at 0x1a9b5b7d668>,
      <matplotlib.patches.Rectangle at 0x1a9b5b7dcc0>,
      <matplotlib.patches.Rectangle at 0x1a9b5b87358>,
      <matplotlib.patches.Rectangle at 0x1a9b5b879b0>,
      <matplotlib.patches.Rectangle at 0x1a9b5b90048>,
      <matplotlib.patches.Rectangle at 0x1a9b5b906a0>,
      <matplotlib.patches.Rectangle at 0x1a9b5b90cf8>,
      <matplotlib.patches.Rectangle at 0x1a9b5b9a390>,
      <matplotlib.patches.Rectangle at 0x1a9b5b9a9e8>,
      <matplotlib.patches.Rectangle at 0x1a9b5ba3080>,
      <matplotlib.patches.Rectangle at 0x1a9b5ba36d8>,
      <matplotlib.patches.Rectangle at 0x1a9b5ba3d30>,
      <matplotlib.patches.Rectangle at 0x1a9b5bad3c8>,
      <matplotlib.patches.Rectangle at 0x1a9b5bada20>,
      <matplotlib.patches.Rectangle at 0x1a9b5bb60b8>,
      <matplotlib.patches.Rectangle at 0x1a9b5bb6710>,
      <matplotlib.patches.Rectangle at 0x1a9b5bb6d68>,
      <matplotlib.patches.Rectangle at 0x1a9b5bc2400>,
      <matplotlib.patches.Rectangle at 0x1a9b5bc2a58>,
      <matplotlib.patches.Rectangle at 0x1a9b5bcd0f0>,
      <matplotlib.patches.Rectangle at 0x1a9b5bcd748>,
      <matplotlib.patches.Rectangle at 0x1a9b5bcdda0>,
      <matplotlib.patches.Rectangle at 0x1a9b5bd6438>,
      <matplotlib.patches.Rectangle at 0x1a9b5bd6a90>,
      <matplotlib.patches.Rectangle at 0x1a9b5be1128>,
      <matplotlib.patches.Rectangle at 0x1a9b5be1780>,
      <matplotlib.patches.Rectangle at 0x1a9b5be1dd8>,
      <matplotlib.patches.Rectangle at 0x1a9b5be8470>,
      <matplotlib.patches.Rectangle at 0x1a9b5be8ac8>,
      <matplotlib.patches.Rectangle at 0x1a9b5bf3160>,
      <matplotlib.patches.Rectangle at 0x1a9b5bf37b8>,
      <matplotlib.patches.Rectangle at 0x1a9b5bf3e10>,
      <matplotlib.patches.Rectangle at 0x1a9b5bfc4a8>,
      <matplotlib.patches.Rectangle at 0x1a9b5bfcb00>,
      <matplotlib.patches.Rectangle at 0x1a9b5c08198>,
      <matplotlib.patches.Rectangle at 0x1a9b5c087f0>,
      <matplotlib.patches.Rectangle at 0x1a9b5c08e48>,
      <matplotlib.patches.Rectangle at 0x1a9b5c134e0>,
      <matplotlib.patches.Rectangle at 0x1a9b5c13b38>,
      <matplotlib.patches.Rectangle at 0x1a9b5c1c1d0>,
      <matplotlib.patches.Rectangle at 0x1a9b5c1c828>,
      <matplotlib.patches.Rectangle at 0x1a9b5c1ce80>,
      <matplotlib.patches.Rectangle at 0x1a9b5c26518>,
      <matplotlib.patches.Rectangle at 0x1a9b5c26b70>,
      <matplotlib.patches.Rectangle at 0x1a9b5c2f208>,
      <matplotlib.patches.Rectangle at 0x1a9b5c2f860>,
      <matplotlib.patches.Rectangle at 0x1a9b5c2feb8>,
      <matplotlib.patches.Rectangle at 0x1a9b5c39550>,
      <matplotlib.patches.Rectangle at 0x1a9b5c39ba8>,
      <matplotlib.patches.Rectangle at 0x1a9b5c42240>,
      <matplotlib.patches.Rectangle at 0x1a9b5c42898>,
      <matplotlib.patches.Rectangle at 0x1a9b5c42ef0>,
      <matplotlib.patches.Rectangle at 0x1a9b5c4d588>,
      <matplotlib.patches.Rectangle at 0x1a9b5c4dbe0>,
      <matplotlib.patches.Rectangle at 0x1a9b5c57278>,
      <matplotlib.patches.Rectangle at 0x1a9b5c578d0>,
      <matplotlib.patches.Rectangle at 0x1a9b5c57f28>,
      <matplotlib.patches.Rectangle at 0x1a9b5c635c0>,
      <matplotlib.patches.Rectangle at 0x1a9b5c63c18>,
      <matplotlib.patches.Rectangle at 0x1a9b5c6c2b0>,
      <matplotlib.patches.Rectangle at 0x1a9b5c6c908>,
      <matplotlib.patches.Rectangle at 0x1a9b5c6cf60>,
      <matplotlib.patches.Rectangle at 0x1a9b5c755f8>,
      <matplotlib.patches.Rectangle at 0x1a9b5c75c50>,
      <matplotlib.patches.Rectangle at 0x1a9b5c7e2e8>,
      <matplotlib.patches.Rectangle at 0x1a9b5c7e940>,
      <matplotlib.patches.Rectangle at 0x1a9b5c7ef98>])

```python
len(list(dataframe1["bid"]))
```

```python

```

```python

```


```python
import numpy as np
import matplotlib.pyplot as plt
import math
from scipy import interpolate
from scipy import integrate
import time
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from math import sqrt
```

### Derive form for ODE in way that can be solved.
Starting with 16
\begin{equation}
    -\xi \frac{dT_e}{d\xi} = \frac{d^2 T_e^{n+4}}{d\xi^2} + \beta e \gamma_0 (T_i-T_e)
\end{equation}
define u
$$u(\xi) = T'(\xi)  $$
$$\frac{d^2 T_e^{n+4}}{d \xi} = (n+4) T_e(\xi)^{n+2} \left((n+3) u^2+T_e(\xi) u'(\xi)\right).$$

16 becomes
\begin{equation}
    -\xi u = (n+4) T_e(\xi)^{n+2} \left((n+3) u^2+T_e(\xi) u'(\xi)\right) + B_e(T_i-T_e)
\end{equation}
Solving for u'
$
\begin{equation}
     \frac{-\xi u - B_e(T_i-T_e)}{(n+4)T_e^{n+2}}-(n+3)u^2 = T_e(\xi)u'(\xi)
\end{equation}
Simplifying
\begin{equation}
    u' = -\frac{\frac{(\xi u + B_e(Ti-T_e))T_e^{-n} }{n+4} + (n+3)u^2T_e^2}{T_e^3}
\end{equation}
Finding the derivative of Ti
\begin{equation}\label{eq:17}
    \xi \frac{dT_i}{d\xi} = \beta_i\gamma_0(T_i-T_e)
\end{equation}
\begin{equation}
    Ti' = \frac{B_i(T_i-T_e)}{\xi}
\end{equation}

### Find initial conditions
Te initial conditions zero order
\begin{equation}
    T_e^0(\xi) = \bigg{(}\frac{(n+3)}{(n+4)}\xi_{\mathrm{max}}({\xi_{\mathrm{max}}}-\xi)\bigg{)}^{1/(n+3)}
\end{equation}
The derivative of this equation is
\begin{equation}
    \frac{dT_e^0}{d\xi} = -\bigg{(}\frac{(n+3)}{(n+4)}\xi_{\mathrm{max}}\bigg{)}^{1/(n+3)}\bigg{(}\frac{(\xi_{\mathrm{max}}-\xi)}{n+3}\bigg{)}^{(2-n)/(n+3)}
\end{equation}
Initial conditions for Ti from 33
\begin{equation}
    T^0_i(\xi)= B_i\frac{(n+3)}{(n+4)}(\xi_{\mathrm{max}}-\xi)T_e^0(\xi)
\end{equation}

### Test Problem
$$ B_e = \beta_e\gamma_0$$
$$\beta_e = \frac{1}{\rho C_{ve}}$$
$$C_{ve} = 0.1\rho, C_{vi} = 0.05\rho [GJcm^{-3}keV^{-1}]$$
$$\gamma_0 = 3.0GJ/(cm^3keV-ns)$$
$$\rho = 3.0 g cm^{-3}$$

### Timestep
$$ A = \frac{1}{2\alpha} $$
$$ \alpha = \frac{ac}{3} \frac{4}{n+4} \frac{1}{\kappa_0} \frac{1}{\rho C_{ve}} $$
Using the A defined in Petshek paper
$$ A^2 = \frac{1}{2K}$$
$$m = n+4$$
$$ K = \frac{ac}{3} \frac{4}{n+4} \frac{1}{\kappa_0} \frac{1}{ C_{ve}\rho^{2}} $$
$$ a = 0.01372 GJKeV^{-4} cm^{-3}$$-Wolbier paper
$$ c = 29.9792458 (cm ns^{-1}) $$
$$ \frac{\xi\sqrt{\tau}}{A}=x $$

```python
g = lambda x, xmax: ((n+3)*xmax*(-x+xmax)/(n+4))**(1/(n+3))
gprime = lambda x, xmax: -(((n+3)/(n+4))*xmax)**(1/(n+3))*1/(n+3)*(xmax-x)**((-2-n)/(n+3))
h = lambda x, xmax: Bi*((n+3)/(n+4))*(xmax-x)*((n+3)/(n+4)*xmax*(xmax-x))**(1/(n+3))
```

```python
#V[0] = Te, V[1] = u, v[2] = Ti
RHSfun = lambda t,v: np.array([v[1],
                               ((-t*v[1] - Be*(v[2]-v[0]))*(v[0]**(-n))/(n+4)-(n+3)*(v[1]**2)*(v[0]**2))/(v[0]**3),
                               Bi*(v[2]-v[0])/t])

```

```python
n = 3
c=299.792458
xmax = 1.12
a=0.01372
kappa0= 100
tol = 1e-3
Cve=0.1
Cvi=.05
rho = 3.0
gamma0 = 1.0
K = a*c*4/(3*(n+4)*kappa0*Cve*rho**2)
A = sqrt(1/(2*K))
print("A = ",A)
Betae=1/(rho*Cve)
Betai=1/(rho*Cvi)
Bi = Betae*gamma0
Be = Betai*gamma0
IC = [g(xmax-tol,xmax),gprime(xmax-tol,xmax), h(xmax-tol,xmax)]
output = integrate.DOP853(RHSfun, xmax,IC,0)
print(Bi, Be, Betai, Betae,n)
```

    A =  7.578766058136872
    3.333333333333333 6.666666666666666 6.666666666666666 3.333333333333333 3

```python

opt_fun = lambda xm: integrate.solve_ivp(RHSfun, (xm-tol,0), [g(xm-tol,xm),gprime(xm-tol,xm), h(xm-tol,xm)],
                                         method='DOP853', max_step = 1e-4).y[0][-1] - 1
```

```python
from scipy import optimize
minval = optimize.brentq(opt_fun, .5,1.5)
print(minval)
```

    /anaconda3/lib/python3.7/site-packages/ipykernel_launcher.py:4: RuntimeWarning: divide by zero encountered in double_scalars
      after removing the cwd from sys.path.
    /anaconda3/lib/python3.7/site-packages/ipykernel_launcher.py:4: RuntimeWarning: invalid value encountered in double_scalars
      after removing the cwd from sys.path.

    0.7647859303039829

```python
sol = integrate.solve_ivp(RHSfun, (minval-tol,0), IC, method='DOP853', max_step = 1e-4)
plt.figure(1)
plt.plot(sol.t, sol.y[0], label="Te")
plt.plot(sol.t, sol.y[2], label ="Ti")
plt.legend()
plt.xlabel(r"$\xi$")
plt.ylabel("T")
plt.title("Marshak wave n=%i" %(n))
plt.show
plt.figure(2)
for tau in [0.1,0.5,1]:#,5,10]:
    z=sqrt(tau)/A
    print(z*minval)
    plt.plot(sol.t*z, sol.y[0], label="Te t=%.1f ns"% (tau))
    plt.plot(sol.t*z, sol.y[2], label ="Ti t=%.1f ns"% (tau))
    plt.legend()
    plt.xlabel("x (cm)")
    plt.ylabel("T (KeV)")
    plt.title("Marshak wave n=%i at 3 timesteps" %(n))
    np.savetxt("MarshakTimes_n%d_t%1.2f_Bi%1.3f_Be%1.3f.csv" % (n,tau,Bi,Be), sol.t*z, delimiter=",")
    np.savetxt("MarshakTi_n%d_t%1.2f_Bi%1.3f_Be%1.3f.csv" % (n,tau,Bi,Be), sol.y[2], delimiter=",")
    np.savetxt("MarshakTe_n%d_t%1.2f_Bi%1.3f_Be%1.3f.csv" % (n,tau,Bi,Be), sol.y[0], delimiter=",")
plt.show()
```

```python

```

```python
"MarshakTimesn%d_t%1.2f_Bi%1.3f_Be%1.3f.csv" % (n,tau,Bi,Be)
```

    'MarshakTimesn2_t1.00_Bi0.17_Be0.33.csv'

```python

```

```python

```

```python

```

```python

```


# [2.3]() Generalized dynamic programming for multiple sequence alignment [edit]

**Table of Contents**
0. [Progressive alignment](#1)
    0. [Building the guide tree](#1.1)
    0. [Generalization of Needleman-Wunsch (with affine gap scoring) for progressive multiple sequence alignment](#1.2)
    0. [Putting it all together: progressive multiple sequence alignment](#1.3)
0. [Progressive alignment versus iterative alignment](#2)

Until now we worked with alignments between two sequences, but it is likely that you will want to align many sequences at the same time. For example, if you are trying to gain insight on the evolutionary relationships between all of the 16S bacterial genes in a given sample, it would be time consuming and very inefficient to compare them two at a time. It would be more efficient and useful to compare all of the 16S sequences from the bacteria in the same alignment.
In the pairwise sequence alignment chapter, we went over dynamic programming algorithms. It's possible to generalize Smith-Waterman and Needleman-Wunsch, the dynamic programming algorithms that we explored for pairwise sequence alignment, to identify the optimal alignment of more than two sequences. Remember that our scoring scheme for pairwise alignment with Smith-Waterman looked like the following:

$$
\begin{align}
& F(0, 0) = 0\\
& F(i, 0) = F(i-1, 0) - d\\
& F(0, j) = F(0, j-1) - d\\
\\
& F(i, j) = max \begin{pmatrix}
& F(i-1, j-1) + s(c_i, c_j)\\
& F(i-1, j) - d\\
& F(i, j-1) - d)\\
\end{pmatrix}
\end{align}
$$

To generalize this to three sequences, we could create $3 \times 3$ scoring, dynamic programming, and traceback matrices. Our scoring scheme would then look like the following:

$$
\begin{align}
& F(0, 0, 0) = 0\\
& F(i, 0, 0) = F(i-1, 0, 0) - d\\
& F(0, j, 0) = F(0, j-1, 0) - d\\
& F(0, 0, k) = F(0, 0, k-1) - d\\
\\
& F(i, j, k) = max \begin{pmatrix}
F(i-1, j-1, k-1) + s(c_i, c_j) + s(c_i, c_k) + s(c_j, c_k)\\
F(i, j-1, k-1) + s(c_j, c_k) - d\\
F(i-1, j, k-1) + s(c_i, c_k) - d\\
F(i-1, j-1, k) + s(c_i, c_j) - d\\
F(i, j, k-1) - 2d\\
F(i, j-1, k) - 2d\\
F(i-1, j, k) - 2d\\
\end{pmatrix}
\end{align}
$$

However the complexity of this algorithm is much worse than for pairwise alignment. For pairwise alignment, remember that if aligning two sequences of lengths $m$ and $n$, the runtime of the algorithm will be proportional to $m \times n$. If $n$ is longer than or as long as $m$, we simplify the statement to say that the runtime of the algorithm will be be proportional to $n^2$. This curve has a pretty scary trajectory: runtime for pairwise alignment with dynamic programming is said to scale quadratically.

```
%pylab inline
from __future__ import division, print_function
from functools import partial
from IPython.core import page
page.page = print
```

```
import matplotlib.pyplot as plt

seq_lengths = range(25)
s2_times = [t ** 2 for t in range(25)]

plt.plot(range(25), s2_times)
plt.xlabel('Sequence Length')
plt.ylabel('Runtime (s)')
```

The exponent in the $n^2$ term comes from the fact that, in pairwise alignment, if we assume our sequences are both of length $n$, there are $n \times n$ cells to fill in in the dynamic programming matrix. If we were to generalize either Smith-Waterman or Needleman-Wunsch to three sequences, we would need to create a 3 dimensional array to score and trace back the alignment. For sequences of length $n$, we would therefore have $n \times n \times n$ cells to fill in, and our runtime versus sequence length curve would look like the following.

```
s3_times = [t ** 3 for t in range(25)]

plt.plot(range(25), s3_times)
plt.xlabel('Sequence Length')
plt.ylabel('Runtime (s)')
```

That curve looks steeper than the curve for pairwise alignment, and the values on the y-axis are bigger, but it's not really clear how much of a problem this is until we plot runtime for three sequences in the context of the run times for pairwise alignment.

```
plt.plot(range(25), s2_times)
plt.plot(range(25), s3_times)
plt.xlabel('Sequence Length')
plt.ylabel('Runtime (s)')
```

And for four sequences:

```
s4_times = [t ** 4 for t in range(25)]

plt.plot(range(25), s2_times)
plt.plot(range(25), s3_times)
plt.plot(range(25), s4_times)
plt.xlabel('Sequence Length')
plt.ylabel('Runtime (s)')
```

We clearly have a problem here, and that is that the runtime for multiple sequence alignment using full dynamic programming algorithms grows exponentially with the number of sequences to be aligned. If $n$ is our sequence length, and $s$ is the number of sequences, that means that runtime is proportional to $n^s$. In pairwise alignment, $s$ is always equal to 2, so the problem is more manageable. However, for the general case of $s$ sequences, we really can't even consider Smith-Waterman or Needleman-Wunsch for more than just a few sequences. The pattern in the plots above should illustrate why.

As we explored with database searching, we need to figure out how to align fewer sequences. This is where *progressive alignment* comes in.

## [2.3.1](#1) Progressive alignment [edit]

**Table of Contents**
0. [Building the guide tree](#1.1)
0. [Generalization of Needleman-Wunsch (with affine gap scoring) for progressive multiple sequence alignment](#1.2)
0. [Putting it all together: progressive multiple sequence alignment](#1.3)

In progressive alignment, the problem of exponential growth of runtime and space is managed by selectively aligning pairs of sequences, and aligning alignments of sequences. What we typically do is identify a pair of closely related sequences, and align those. Then, we identify the next most closely related sequence to that initial pair, and align that sequence to the alignment. This concept of aligning a sequence to an alignment is new, and we'll come back to it in just a few minutes. The other concept of identifying the most closely related sequences, and then the next most closely related sequence, and so on should sound familiar. It effectively means that we're traversing a tree. And herein lies our problem: we need a tree to efficiently align multiple sequences, but we need an alignment to build a good tree.

You probably have two burning questions in your mind right now:

1. How do we build a tree to guide the alignment process, if we need an alignment to build a good tree?
2. How do we align a sequence to an alignment, or an alignment to an alignment?

We'll explore both of those through-out the rest of this notebook. First, let's cover the process of progressive multiple sequence alignment, just assuming for a moment that we know how to do both of those things.

The process of progressive multiple sequence alignment could look like the following. First, we start with some sequences and a tree representing the relationship between those sequences. We'll call this our guide tree, because it's going to guide us through the process of multiple sequence alignment. In progressive multiple sequence alignment, we build a multiple sequence alignment for each internal node of the tree, where the alignment at a given internal node contains all of the sequences in the clade defined by that node.

Starting from the root node, descend the bottom branch of the tree until you get to the an internal node. If an alignment hasn't been constructed for that node yet, continue descending the tree until to get to a pair of nodes. In this case, we follow the two branches to the tips. We then align the sequences at that pair of tips (usually with Needleman-Wunsch, for multiple sequence alignment), and assign that alignment to the node connecting those tips.

Next, we want to find what to align the resulting alignment to, so start from the root node and descend the top branch of the tree. When you get to the next node, determine if an alignment has already been created for that node. If not, our job is to build that alignment so we have something to align against. In this case, that means that we need to align `s1`, `s2`, and `s3`. We can achieve this by aligning `s1` and `s3` first, to get the alignment at the internal node connecting them.

We can next align the alignment of `s1` and `s3` with `s2`, to get the alignment at the internal node connecting those clades.

And finally, we can compute the alignment at the root node of the tree, by aligning the alignment of `s1`, `s2`, and `s3` with the alignment of `s4` and `s5`.

The alignment at the root node is our multiple sequence alignment.

### [2.3.1.1](#1.1) Building the guide tree [edit]

Let's address the first of our outstanding questions. I mentioned above that *we need an alignment to build a good tree*. The key word here is *good*. We can build a very rough tree - one that we would never want to present as representing the actual relationships between the sequences in question - without first aligning the sequences. Remember that building a UPGMA tree requires only a distance matrix, so if we can find a non-alignment-dependent way to compute distances between the sequences, we can build a rough UPGMA tree from them.

Let's compute distances between the sequences based on their *word* composition. We'll define a *word* here as `k` adjacent characters in the sequence. We can then define a function that will return all of the words in a sequence as follows. These words can be defined as being overlapping, or non-overlapping. We'll go with overlapping for this example, as the more words we have, the better our guide tree should be.

```
from skbio import DNA
%psource DNA.iter_kmers
```

```
for e in DNA("ACCGGTGACCAGTTGACCAGTA").iter_kmers(3):
    print(e)
```

```
for e in DNA("ACCGGTGACCAGTTGACCAGTA").iter_kmers(7):
    print(e)
```

```
for e in DNA("ACCGGTGACCAGTTGACCAGTA").iter_kmers(3, overlap=False):
    print(e)
```

If we then have two sequences, we can compute the word counts for each and define a distance between the sequences as the fraction of words that are unique to either sequence.

```
from iab.algorithms import kmer_distance
%psource kmer_distance
```

We can then use this as a distance function...

```
s1 = DNA("ACCGGTGACCAGTTGACCAGT")
s2 = DNA("ATCGGTACCGGTAGAAGT")
s3 = DNA("GGTACCAAATAGAA")

print(s1.distance(s2, kmer_distance))
print(s1.distance(s3, kmer_distance))
```

If we wanted to override the default to create (for example) a 5-mer distance function, we could use ``functools.partial``.

```
fivemer_distance = partial(kmer_distance, k=5)

s1 = DNA("ACCGGTGACCAGTTGACCAGT")
s2 = DNA("ATCGGTACCGGTAGAAGT")
s3 = DNA("GGTACCAAATAGAA")

print(s1.distance(s2, fivemer_distance))
print(s1.distance(s3, fivemer_distance))
```

We can now apply one of these functions to build a distance matrix for a set of sequences that we want to align.

```
query_sequences = [DNA("ACCGGTGACCAGTTGACCAGT", {"id": "s1"}),
                   DNA("ATCGGTACCGGTAGAAGT", {"id": "s2"}),
                   DNA("GGTACCAAATAGAA", {"id": "s3"}),
                   DNA("GGCACCAAACAGAA", {"id": "s4"}),
                   DNA("GGCCCACTGAT", {"id": "s5"})]
```

```
from skbio import DistanceMatrix

guide_dm = DistanceMatrix.from_iterable(query_sequences, metric=kmer_distance, key='id')
```

scikit-bio also has some basic visualization functionality for these objects. For example, we can easily visualize this object as a heatmap.

```
fig = guide_dm.plot(cmap='Greens')
```

We can next use some functionality from SciPy to cluster the sequences with UPGMA, and print out a dendrogram.

```
from scipy.cluster.hierarchy import average, dendrogram, to_tree

for q in query_sequences:
    print(q)

guide_lm = average(guide_dm.condensed_form())
guide_d = dendrogram(guide_lm, labels=guide_dm.ids, orientation='right',
                     link_color_func=lambda x: 'black')
guide_tree = to_tree(guide_lm)
```

```
from iab.algorithms import guide_tree_from_sequences
%psource guide_tree_from_sequences
```

```
t = guide_tree_from_sequences(query_sequences, display_tree=True)
```

We now have a guide tree, so we can move on to the next step of progressive alignment.

### [2.3.1.2](#1.2) Generalization of Needleman-Wunsch (with affine gap scoring) for progressive multiple sequence alignment [edit]

Next, we'll address our second burning question: aligning alignments. As illustrated above, there are basically three different types of pairwise alignment we need to support for progressive multiple sequence alignment with Needleman-Wunsch. These are:

1. Alignment of a pair of sequences.
2. Alignment of a sequence and an alignment.
3. Alignment of a pair of alignments.

Standard Needleman-Wunsch supports the first, and it is very easy to generalize it to support the latter two. The only change that is necessary is in how the alignment of two non-gap characters is scored. Recall that we previously scored an alignment of two characters by looking up the score of substitution from one to the other in a substitution matrix. To adapt this for aligning a sequence to an alignment, or for aligning an alignment to an alignment, we compute this substitution as the average score of aligning the pairs of characters.

For example, if we want to align the alignment column from $aln1$:

```
A
C
```

to the alignment column from $aln2$:

```
T
G
```

we could compute the substitution score using the matrix $m$ as:

$$
s = \frac{m[A][T] + m[A][G] + m[C][T] + m[C][G]}{aln1_{length} \times aln2_{length}}
$$

The following code adapts our implementation of Needleman-Wunsch to support aligning a sequence to an alignment, or aligning an alignment to an alignment.

```
from iab.algorithms import format_dynamic_programming_matrix, format_traceback_matrix
from skbio.alignment._pairwise import _compute_score_and_traceback_matrices

%psource _compute_score_and_traceback_matrices
```

```
from skbio.alignment._pairwise import _traceback
%psource _traceback
```

```
from skbio.alignment import global_pairwise_align_nucleotide
%psource global_pairwise_align_nucleotide
```

For the sake of the examples below, I'm going to override one of the ``global_pairwise_align_nucleotide`` defaults to penalize terminal gaps. This effectively tells the algorithm that we know we have a collection of sequences that are homologous from beginning to end.

```
global_pairwise_align_nucleotide = partial(global_pairwise_align_nucleotide, penalize_terminal_gaps=True)
```

For example, we can still use this code to align pairs of sequences (but note that we now need to pass those sequences in as a pair of one-item lists):

```
aln1, _, _ = global_pairwise_align_nucleotide(query_sequences[0], query_sequences[1])
print(aln1)
```

We can align that alignment to one of our other sequences.

```
aln1, _, _ = global_pairwise_align_nucleotide(aln1, query_sequences[2])
print(aln1)
```

Alternatively, we can align another pair of sequences:

```
aln2, _, _ = global_pairwise_align_nucleotide(query_sequences[2], query_sequences[3])
print(aln2)
```

And then align that alignment against our previous alignment:

```
aln3, _, _  = global_pairwise_align_nucleotide(aln1, aln2)
print(aln3)
```

### [2.3.1.3](#1.3) Putting it all together: progressive multiple sequence alignment [edit]

We can now combine all of these steps to take a set of query sequences, build a guide tree, perform progressive multiple sequence alignment, and return the guide tree (as a SciPy linkage matrix) and the alignment.

```
from skbio import TreeNode
guide_tree = TreeNode.from_linkage_matrix(guide_lm, guide_dm.ids)
```

We can view the guide tree in [Newick format](http://scikit-bio.org/docs/latest/generated/skbio.io.newick.html) as follows:

```
print(guide_tree)
```

```
from iab.algorithms import progressive_msa
%psource progressive_msa
```

```
msa = progressive_msa(query_sequences, pairwise_aligner=global_pairwise_align_nucleotide, guide_tree=guide_tree)
print(msa)
```

We can now build a (hopefully) improved tree from our multiple sequence alignment. First we'll look at our original distance matrix again, and then the distance matrix generated from the progressive multiple sequence alignment.

```
fig = guide_dm.plot(cmap='Greens')
```

```
msa_dm = DistanceMatrix.from_iterable(msa, metric=kmer_distance)
fig = msa_dm.plot(cmap='Greens')
```

The UPGMA trees that result from these alignments are very different. First we'll look at the guide tree, and then the tree resulting from the progressive multiple sequence alignment.

```
d = dendrogram(guide_lm, labels=guide_dm.ids, orientation='right',
               link_color_func=lambda x: 'black')
```

```
msa_lm = average(msa_dm.condensed_form())
d = dendrogram(msa_lm, labels=msa_dm.ids, orientation='right',
               link_color_func=lambda x: 'black')
```

And we can wrap this all up in a single convenience function:

```
from iab.algorithms import progressive_msa_and_tree
%psource progressive_msa_and_tree
```

```
msa = progressive_msa(query_sequences, pairwise_aligner=global_pairwise_align_nucleotide, guide_tree=guide_tree)
```

```
msa, tree = progressive_msa_and_tree(query_sequences, pairwise_aligner=global_pairwise_align_nucleotide,
                                     display_tree=True, display_aln=True)
```

## [2.3.2](#2) Progressive alignment versus iterative alignment [edit]

In an iterative alignment, the output tree from the above progressive alignment is used as a guide tree, and the full process repeated. This is performed to reduce errors that result from a low-quality guide tree.

```
from iab.algorithms import iterative_msa_and_tree
%psource iterative_msa_and_tree
```

```
msa, tree = iterative_msa_and_tree(query_sequences, pairwise_aligner=global_pairwise_align_nucleotide, num_iterations=1, display_aln=True, display_tree=True)
```

```
msa, tree = iterative_msa_and_tree(query_sequences, pairwise_aligner=global_pairwise_align_nucleotide, num_iterations=2, display_aln=True, display_tree=True)
```

```
msa, tree = iterative_msa_and_tree(query_sequences, pairwise_aligner=global_pairwise_align_nucleotide, num_iterations=3, display_aln=True, display_tree=True)
```

```
msa, tree = iterative_msa_and_tree(query_sequences, pairwise_aligner=global_pairwise_align_nucleotide, num_iterations=5, display_aln=True, display_tree=True)
```

Some references that I used in assembling these notes include [1](http://statweb.stanford.edu/~nzhang/345_web/sequence_slides3.pdf), [2](http://math.mit.edu/classes/18.417/Slides/alignment.pdf), [3](http://www.sciencedirect.com/science/article/pii/0378111988903307), [4](http://bioinformatics.oxfordjournals.org/content/23/21/2947.full), and [5](http://nar.oxfordjournals.org/content/32/5/1792.full).


# Non-linear Gaussian filtering and smoothing

Provided are two examples of nonlinear state-space models on which one can perform Bayesian filtering and smoothing in order to obtain
a posterior distribution over a latent state trajectory based on noisy observations.
In order to understand the theory behind these methods in detail we refer to [1] and [2].

We provide examples for two different types of state-space model:
1. [Nonlinear, Discrete State-Space Model](#1.-Non-Linear-Discrete-State-Space-Model:-Pendulum): Pendulum
2. [Nonlinear, Continuous-Discrete State-Space Model](#2.-Non-Linear-Continuous-Discrete-State-Space-Model:-Benes-Daum-Filter): Benes-Daum Filter

In order to perform Bayesian Filtering and Smoothing on non-linear models, we show how to use functionalities provided by ProbNum to
either linearise the respective model (Extended Kalman Filter, EKF) or to use an Unscented Transform (Unscented Kalman Filter, UKF) which does not require a Jacobian of the nonlinear model.

**References**:
> [1] Särkkä, Simo, and Solin, Arno. Applied Stochastic Differential Equations. Cambridge University Press, 2019.
>
> [2] Särkkä, Simo. Bayesian Filtering and Smoothing. Cambridge University Press, 2013.

```python
import numpy as np

import probnum as pn
from probnum import filtsmooth, randvars, randprocs
from probnum.problems import TimeSeriesRegressionProblem
```

```python
rng = np.random.default_rng(seed=123)
```

```python
# Make inline plots vector graphics instead of raster graphics
%matplotlib inline
from IPython.display import set_matplotlib_formats

set_matplotlib_formats("pdf", "svg")

# Plotting
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec

plt.style.use("../../probnum.mplstyle")
```

    /tmp/ipykernel_98060/236124620.py:5: DeprecationWarning: `set_matplotlib_formats` is deprecated since IPython 7.23, directly use `matplotlib_inline.backend_inline.set_matplotlib_formats()`
      set_matplotlib_formats("pdf", "svg")

## 1. **Non-Linear Discrete** State-Space Model: Pendulum

---

For nonlinear models, we assume that at least one of the components of our state-space model (i.e. either the dynamics, the measurement model, or both) is a nonlinear function of the latent state.

Consider nonlinear functions $f: \mathbb{R}^d \rightarrow \mathbb{R}^d$ and $h: \mathbb{R}^d \rightarrow \mathbb{R}^m$ where $d$ is the dimension of the latent state and
$m$ is the dimension of the measurements.

For $k = 1, \dots, K$ and $x_0 \sim \mathcal{N}(\mu_0, \Sigma_0)$:

$$
\begin{align}
    \boldsymbol{x}_k &\sim \mathcal{N}(f(\boldsymbol{x}_{k-1}), \boldsymbol{Q}) \\
    \boldsymbol{y}_k &\sim \mathcal{N}(h(\boldsymbol{x}_k), \boldsymbol{R})
\end{align}
$$

This defines a dynamics model that assumes a state $x_k$ in a **discrete** sequence of states arising from a nonlinear map of the previous state $x_{k-1}$ corrupted with additive Gaussian noise under a **process noise** covariance matrix $Q$.
Similarly, the measurements $y_k$ are assumed to be nonlinear transformations of the latent state under additive Gaussian noise according to a **measurement noise** covariance $R$.
In the described case, we cannot use the Kalman Filtering equations since these assume linear relationships. One method is to linearise the nonlinear parts of the model by a first-order Taylor Expansion (Extended Kalman Filter, EKF) or employing the Unscented Kalman Filter (UKF) which is based on the unscented transform and does not require the computation of a Jacobian.
The `probnum` packages provides both methods by wrapping linearizing "Components" around nonlinear dynamcis and/or measurement models, as we will see below.

Note that this can be generalized to a nonlinear time-varying state-space model, as well. Then we would have $f: \mathbb{T} \times \mathbb{R}^d \rightarrow \mathbb{R}^d$ and $h: \mathbb{T} \times \mathbb{R}^d \rightarrow \mathbb{R}^m$ where $\mathbb{T}$ is the "time dimension".

### Define State-Space Model

#### I. Discrete Dynamics Model: Nonlinear Gaussian Transitions

```python
state_dim = 2
observation_dim = 1
```

```python
# approx. gravitational constant
g = 9.81
# dt
delta_t = 0.0075

def pendulum_rhs(state):
    """Right-hand side of an ODE that defines pendulum dynamics"""
    x1, x2 = state
    y1 = x1 + x2 * delta_t
    y2 = x2 - g * np.sin(x1) * delta_t
    return np.array([y1, y2])

def pendulum_jacobian(state):
    """Jacobian of the pendulum ODE"""
    x1, x2 = state
    dy1_dx = [1.0, delta_t]
    dy2_dx = [-g * np.cos(x1) * delta_t, 1.0]
    return np.array([dy1_dx, dy2_dx])

dynamics_transition_function = lambda t, state: pendulum_rhs(state)
dynamics_transition_jacobian_function = lambda t, state: pendulum_jacobian(state)

dynamics_diffusion_matrix = 1.0 * (
    np.diag(np.array([delta_t ** 3 / 3, delta_t]))
    + np.diag(np.array([delta_t ** 2 / 2]), 1)
    + np.diag(np.array([delta_t ** 2 / 2]), -1)
)
```

We here define the non-linear dynamics of a pendulum. Note (in the cell directly above) that `probnum` expects `dynamicsfun` to be a function of time and state. Even though we here consider a nonlinear but time-**in**variant model, we still have to fulfill the contract that the interface requires.
```python
dynamics_transition_function = lambda t, state: pendulum_rhs(state)
```
In other words, we simply ignore the time variable `t` in case of time-invariant models.

To create a discrete, (non-linear) Gaussian dynamics model, `probnum` provides the `NonlinearGaussian` class that takes
- `input_dim` and `output_dim`
- `state_trans_fun` : the function that returns the next state given the current one
- `proc_noise_cov_mat_fun` : the function that returns the covariance matrix of the Gaussian process noise
- `jacob_state_trans_fun` : the function that returns the Jacobian matrix of the state transition function

```python
noise = randvars.Normal(
    mean=np.zeros(dynamics_diffusion_matrix.shape[0]), cov=dynamics_diffusion_matrix
)

# Create discrete, non-linear Gaussian dynamics model
dynamics_model = randprocs.markov.discrete.NonlinearGaussian(
    input_dim=state_dim,
    output_dim=state_dim,
    transition_fun=dynamics_transition_function,
    noise_fun=lambda t: noise,
    transition_fun_jacobian=dynamics_transition_jacobian_function,
)
```

#### II. Discrete Measurement Model: Nonlinear Gaussian Measurements

```python
def pendulum_measurement(state):
    x1, x2 = state
    return np.array([np.sin(x1)])

def pendulum_measurement_jacobian(state):
    x1, x2 = state
    return np.array([[np.cos(x1), 0.0]])

measurement_function = lambda t, state: pendulum_measurement(state)
measurement_jacobian_function = lambda t, state: pendulum_measurement_jacobian(state)

measurement_variance = 0.32 ** 2
measurement_covariance = measurement_variance * np.eye(observation_dim)
```

```python
measurement_jacobian_function(0, np.ones(2)) @ np.ones(2)
```

    array([0.54030231])

```python
# Create discrete, non-linear Gaussian measurement model
measurement_noise = randvars.Normal(mean=np.zeros(measurement_covariance.shape[0]), cov=measurement_covariance)
measurement_model = randprocs.markov.discrete.NonlinearGaussian(
    input_dim=state_dim,
    output_dim=observation_dim,
    transition_fun=measurement_function,
    noise_fun=lambda t: measurement_noise,
    transition_fun_jacobian=measurement_jacobian_function,
)

```

```python

```

#### III. Initial State Random Variable

```python
mu_0 = np.ones(state_dim)
sigma_0 = measurement_variance * np.eye(state_dim)
initial_state_rv = randvars.Normal(mean=mu_0, cov=sigma_0)
```

### Kalman Filtering

#### I. Linearize model (Extended Kalman Filter)

Since we could easily derive the Jacobian of our nonlinear model components above, we can use the Extended Kalman Filter (EKF) to obtain a filtering posterior over the states. Therefore, in `probnum` we just wrap the dynamics model and the measurement model in a `DiscreteEKFComponent` object.
In cases in which the Jacobian is not as easy to write down, `probnum` also provides the `DiscreteUKFComponent` class that implements the linearization by means of an
_unscented Kalman filter_ that does not require a Jacobian.
Alternatively, the Jacobian can be calculated via automatic differentiation.

```python
linearised_dynamics_model = filtsmooth.gaussian.approx.DiscreteEKFComponent(
    dynamics_model
)

```

```python
prior_process = randprocs.markov.MarkovProcess(
    transition=linearised_dynamics_model, initrv=initial_state_rv, initarg=0.0
)
```

#### II. Kalman Filter

Then, we simply pass these linearized models to the Kalman Filter (exactly the same interface as in the linear case) and proceed.

```python
kalman_filter = filtsmooth.gaussian.Kalman(prior_process)
```

#### III. Generate Data for the State-Space Model

`statespace.generate_artificial_measurements()` is used to sample both latent states and noisy observations from the specified state space model.

```python
time_grid = np.arange(0.0, 5.0, step=delta_t)
```

```python
latent_states, observations = randprocs.markov.utils.generate_artificial_measurements(
    rng=rng,
    prior_process=prior_process,
    measmod=measurement_model,
    times=time_grid,
)
regression_problem = TimeSeriesRegressionProblem(
    observations=observations,
    locations=time_grid,
    measurement_models=[
        filtsmooth.gaussian.approx.DiscreteEKFComponent(measurement_model)
    ]
    * len(time_grid),
)
```

#### IV. Perform Kalman Filtering + Rauch-Tung-Striebel Smoothing

```python
state_posterior, _ = kalman_filter.filtsmooth(regression_problem)
```

The method `filtsmooth` returns a `KalmanPosterior` object which provides convenience functions for e.g. sampling and interpolation.
We can also extract the just computed posterior smoothing state variables by querying the `.state_rvs` property.
This yields a list of Gaussian Random Variables from which we can extract the statistics in order to visualize them.

```python
posterior_state_rvs = (
    state_posterior.states
)  # List of <num_time_points> Normal Random Variables
posterior_state_means = posterior_state_rvs.mean  # Shape: (num_time_points, state_dim)
posterior_state_covs = (
    posterior_state_rvs.cov
)  # Shape: (num_time_points, state_dim, state_dim)
```

### Visualize Results

```python
state_fig = plt.figure()
state_fig_gs = gridspec.GridSpec(ncols=2, nrows=1, figure=state_fig)

ax_00 = state_fig.add_subplot(state_fig_gs[0, 0])
ax_01 = state_fig.add_subplot(state_fig_gs[0, 1])

# Plot means
mu_x_1, mu_x_2 = [posterior_state_means[:, i] for i in range(state_dim)]

ax_00.plot(time_grid, mu_x_1, label="posterior mean")
ax_01.plot(time_grid, mu_x_2)

# Plot marginal standard deviations
std_x_1, std_x_2 = [np.sqrt(posterior_state_covs[:, i, i]) for i in range(state_dim)]

ax_00.fill_between(
    time_grid,
    mu_x_1 - 1.96 * std_x_1,
    mu_x_1 + 1.96 * std_x_1,
    alpha=0.2,
    label="1.96 marginal stddev",
)
ax_01.fill_between(
    time_grid, mu_x_2 - 1.96 * std_x_2, mu_x_2 + 1.96 * std_x_2, alpha=0.2
)

# Plot 5 samples on a subsampled grid
samples = state_posterior.sample(rng=rng, t=state_posterior.locations[::5], size=5)
for smp in samples:
    ax_00.plot(
        state_posterior.locations[::5],
        smp[:, 0],
        color="gray",
        alpha=0.75,
        linewidth=1,
        linestyle="dashed",
        label="sample",
    )
    ax_01.plot(
        state_posterior.locations[::5],
        smp[:, 1],
        color="gray",
        alpha=0.75,
        linewidth=1,
        linestyle="dashed",
        label="sample",
    )

# Plot groundtruth
ax_00.scatter(time_grid, observations, marker=".", label="measurements")

# Add labels etc.
ax_00.set_xlabel("t")
ax_01.set_xlabel("t")

ax_00.set_title(r"$x_1$")
ax_01.set_title(r"$x_2$")

# The following two lines are just to remove duplicate labels (caused by samples) from the legend
handles, labels = ax_00.get_legend_handles_labels()
by_label = dict(zip(labels, handles))

state_fig.legend(
    by_label.values(), by_label.keys(), loc="center left", bbox_to_anchor=(1, 0.5)
)

state_fig.tight_layout()
```

## 2. **Non-Linear Continuous-Discrete** State-Space Model: Benes-Daum Filter

---

Now, consider we assume **continuous** dynamics. We assume that there is a continuous process that defines the nonlinear dynamics of our latent space from which we collect discrete nonlinear-Gaussian measurements (as above). Only the dynamics model changes. In particular, we formulate the dynamics as a stochastic process in terms of a Nonlinear Time-Invariant Stochastic Differential Equation. We refer to [1] for more details.
Again, consider non-linear functions $f: \mathbb{R}^d \rightarrow \mathbb{R}^d$ and $h: \mathbb{R}^d \rightarrow \mathbb{R}^m$ where $d$ is the dimension of the latent state and
$m$ is the dimension of the measurements.
We define the following nonlinear **continuous-discrete** state-space model:

Let $x(t_0) \sim \mathcal{N}(\mu_0, \Sigma_0)$.

$$
\begin{align}
    d\boldsymbol{x} &= f(\boldsymbol{x}) \, dt + \boldsymbol{L} \, d \boldsymbol{\omega} \\
    \boldsymbol{y}_k &\sim \mathcal{N}(h(\boldsymbol{x}(t_k)), \boldsymbol{R}), \qquad k = 1, \dots, K
\end{align}
$$

where $\boldsymbol{\omega}$ denotes a vector of driving forces (often Brownian Motion).

Note that this can be generalized to a nonlinear time-varying state-space model, as well. Then we would have $f: \mathbb{T} \times \mathbb{R}^d \rightarrow \mathbb{R}^d$ and $h: \mathbb{T} \times \mathbb{R}^d \rightarrow \mathbb{R}^m$ where $\mathbb{T}$ is the "time dimension".

### Define State-Space Model

#### I. Continuous Dynamics Model: Nonlinear Stochastic Differential Equation

```python
state_dim = 1
observation_dim = 1
```

```python
drift_function = lambda t, state: np.tanh(state)
drift_function_jacobian = lambda t, state: 1.0 - np.tanh(state) ** 2

dispersion_function = lambda t, x: np.ones(state_dim).reshape(1, 1)
```

For continous, non-linear dynamics, `probnum` provides a general stochastic differential equation (SDE) interface that takes
- `driftfun`: a (possibly non-linear) drift function,
- `dispersion_matrixfun`: a (possibly non-linear) dispersion function,
- `jacobfun`: the Jacobian of the drift function.

```python
# Create continuous, non-linear SDE dynamics model
dynamics_model = randprocs.markov.continuous.SDE(
    state_dimension=state_dim,
    wiener_process_dimension=1,
    drift_function=drift_function,
    dispersion_function=dispersion_function,
    drift_jacobian=drift_function_jacobian,
)
```

#### II. Discrete Measurement Model: Linear, Time-Invariant, Gaussian Measurements

```python
measurement_marginal_variance = 1.0
measurement_matrix = np.eye(observation_dim, state_dim)
measurement_noise_matrix = measurement_marginal_variance * np.eye(observation_dim)
```

```python
measurement_model = randprocs.markov.discrete.LTIGaussian(
    transition_matrix=measurement_matrix,
    noise = randvars.Normal(np.zeros(observation_dim), measurement_noise_matrix),
)
```

#### III. Initial State Random Variable

```python
mu_0 = np.zeros(state_dim)
sigma_0 = 3.0 * np.eye(state_dim)
initial_state_rv = randvars.Normal(mean=mu_0, cov=sigma_0)
```

### Kalman Filtering

#### I. Linearize model (Extended Kalman Filter)

```python
linearised_dynamics_model = filtsmooth.gaussian.approx.ContinuousEKFComponent(
    dynamics_model
)
```

#### II. Kalman Filter

```python
prior_process = randprocs.markov.MarkovProcess(
    transition=linearised_dynamics_model, initrv=initial_state_rv, initarg=0.0
)
kalman_filter = filtsmooth.gaussian.Kalman(prior_process)
```

#### III. Generate Data for the State-Space Model

`statespace.generate_artificial_measurements()` is used to sample both latent states and noisy observations from the specified state space model.

```python
time_grid = np.arange(0.0, 10.0, step=10 * delta_t)
```

```python
latent_states, observations = randprocs.markov.utils.generate_artificial_measurements(
    rng=rng,
    prior_process=prior_process,
    measmod=measurement_model,
    times=time_grid,
)
```

```python
regression_problem = TimeSeriesRegressionProblem(
    observations=observations,
    locations=time_grid,
    measurement_models=[measurement_model] * len(time_grid),
)
```

#### IV. Perform Kalman Filtering + Rauch-Tung-Striebel Smoothing

```python
state_posterior, _ = kalman_filter.filter(regression_problem)
```

```python
smoothed_state_posterior = kalman_filter.smooth(filter_posterior=state_posterior)
```

The method `filtsmooth` returns a `KalmanPosterior` object which provides convenience functions for e.g. sampling and prediction.
We can also extract the just computed posterior smoothing state variables by querying the `.state_rvs` property.
This yields a list of Gaussian Random Variables from which we can extract the statistics in order to visualize them.

```python
grid = np.linspace(0, 11, 500)

posterior_state_rvs = state_posterior(
    grid
)  # List of <num_time_points> Normal Random Variables
posterior_state_means = (
    posterior_state_rvs.mean.squeeze()
)  # Shape: (num_time_points, state_dim)
posterior_state_covs = (
    posterior_state_rvs.cov
)  # Shape: (num_time_points, state_dim, state_dim)
```

```python
smoothed_posterior_state_rvs = (
    smoothed_state_posterior.states
)  # List of <num_time_points> Normal Random Variables
smoothed_posterior_state_means = (
    smoothed_posterior_state_rvs.mean.squeeze()
)  # Shape: (num_time_points, state_dim)
smoothed_posterior_state_covs = (
    smoothed_posterior_state_rvs.cov
)  # Shape: (num_time_points, state_dim, state_dim)
```

### Visualize Results

```python
state_fig = plt.figure()

ax = state_fig.add_subplot()

# Plot means
ax.plot(grid, posterior_state_means, label="posterior mean")

# Plot marginal standard deviations
std_x = np.sqrt(np.abs(posterior_state_covs)).squeeze()
ax.fill_between(
    grid,
    posterior_state_means - 1.96 * std_x,
    posterior_state_means + 1.96 * std_x,
    alpha=0.2,
    label="1.96 marginal stddev for filter",
)

# Plot smoothed means
ax.plot(time_grid, smoothed_posterior_state_means, label="smoothed posterior mean")

# Plot marginal standard deviations
std_x = np.sqrt(smoothed_posterior_state_covs).squeeze()
ax.fill_between(
    time_grid,
    smoothed_posterior_state_means - 1.96 * std_x,
    smoothed_posterior_state_means + 1.96 * std_x,
    alpha=0.2,
    label="1.96 marginal stddev for smoother",
)

ax.scatter(time_grid, observations, marker=".", label="measurements")

# Add labels etc.
ax.set_xlabel("t")
ax.set_title(r"$x$")
ax.legend()

state_fig.tight_layout()
```

To get a closer look at filter and smoother, we inspect the interval $[0, 1]$ in greater detail.

```python
idx = time_grid < 1.0

ax.set_xlim([0, 1])
ax.set_ylim(np.min(observations[idx]), np.max(observations[idx]))

state_fig.tight_layout()
state_fig
```

```python

```

```python

```