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

## Funciones variacionales lineales, método de Huckel
Se busca obtener los valores de $W$ resultantes de:
\begin{equation}
  \det(H_{ij}-S_{ij}W)=0
\end{equation}
En donde conocemos $H_{ij}$ como la integral de intercambio y $S_{ij}$ como la integral de solapamiento:
\begin{equation}
  H_{ij}=\int f_{i}^{*}\ H\ f_{j}\ d\tau \\
  S_{ij}=\int f_{i}^{*}\ f_{j}\ d\tau
\end{equation}

### Ejemplo de sección 8.5 de libro química cuántica de Levine edición 5

Se utilizan en este ejemplo las 4 funciones siguientes:
\begin{equation}
  f_{1} = x(l-x)\\
  f_{2} = x^{2}(l-x)^{2}\\
  f_{3} = x(l-x)\left(\frac{l}{2}-x\right)\\
  f_{4} = x^{2}(l-x)^{2}\left(\frac{l}{2}-x\right)
\end{equation}

```python
#Definicion de variables y funciones
var('x l m hbar W k') #Variables a utilizar

def hami(F): #Arreglo de funciones y cantidad de funciones como argumentos
    H=eval(input(str('Escribe el Hamiltoniano aplicado a f: ')))
    return H
```

```python
n = int(input('Indica la cantidad de funciones de prueba: '))
F=[]
for i in range (0,n):
    F.append(str(input('Escribre tu funcion %d: '  %(i+1))))

F=sympify(F)
F=Matrix(F)
#Funciones del ejemplo de sección 8.5
#x*(l-x)
#x**2*(l-x)**2
#x*(l-x)*(l/2-x)
#x**2*(l-x)**2*(l/2-x)
```

    Indica la cantidad de funciones de prueba: 4
    Escribre tu funcion 1: x*(l-x)
    Escribre tu funcion 2: x**2*(l-x)**2
    Escribre tu funcion 3: x*(l-x)*(l/2-x)
    Escribre tu funcion 4: x**2*(l-x)**2*(l/2-x)

```python
ff=F*F.T
fHf=F*hami(F).T
H=Integral(fHf, (x,0,l)).doit()
S=Integral(ff, (x,0,l)).doit()
H_SW=H-S*W
```

    Escribe el Hamiltoniano aplicado a f: (-hbar**2/(2*m))*diff(F, x, 2)

```python
H
```

$$\left[\begin{matrix}\frac{\hbar^{2} l^{3}}{6 m} & \frac{\hbar^{2} l^{5}}{30 m} & 0 & 0\\\frac{\hbar^{2} l^{5}}{30 m} & \frac{\hbar^{2} l^{7}}{105 m} & 0 & 0\\0 & 0 & \frac{\hbar^{2} l^{5}}{40 m} & \frac{\hbar^{2} l^{7}}{280 m}\\0 & 0 & \frac{\hbar^{2} l^{7}}{280 m} & \frac{\hbar^{2} l^{9}}{1260 m}\end{matrix}\right]$$

```python
S
```

$$\left[\begin{matrix}\frac{l^{5}}{30} & \frac{l^{7}}{140} & 0 & 0\\\frac{l^{7}}{140} & \frac{l^{9}}{630} & 0 & 0\\0 & 0 & \frac{l^{7}}{840} & \frac{l^{9}}{5040}\\0 & 0 & \frac{l^{9}}{5040} & \frac{l^{11}}{27720}\end{matrix}\right]$$

```python
H_SW
```

$$\left[\begin{matrix}- \frac{W l^{5}}{30} + \frac{\hbar^{2} l^{3}}{6 m} & - \frac{W l^{7}}{140} + \frac{\hbar^{2} l^{5}}{30 m} & 0 & 0\\- \frac{W l^{7}}{140} + \frac{\hbar^{2} l^{5}}{30 m} & - \frac{W l^{9}}{630} + \frac{\hbar^{2} l^{7}}{105 m} & 0 & 0\\0 & 0 & - \frac{W l^{7}}{840} + \frac{\hbar^{2} l^{5}}{40 m} & - \frac{W l^{9}}{5040} + \frac{\hbar^{2} l^{7}}{280 m}\\0 & 0 & - \frac{W l^{9}}{5040} + \frac{\hbar^{2} l^{7}}{280 m} & - \frac{W l^{11}}{27720} + \frac{\hbar^{2} l^{9}}{1260 m}\end{matrix}\right]$$

```python
#Soluciones para W
Wn=solve(H_SW.det(),W)
Wn
```

$$\left [ \frac{2 \hbar^{2}}{l^{2} m} \left(- \sqrt{133} + 14\right), \quad \frac{2 \hbar^{2}}{l^{2} m} \left(\sqrt{133} + 14\right), \quad \frac{2 \hbar^{2}}{l^{2} m} \left(- 9 \sqrt{5} + 30\right), \quad \frac{2 \hbar^{2}}{l^{2} m} \left(9 \sqrt{5} + 30\right)\right ]$$

### Obtención de $\phi_{1}$

Sabemos que:
\begin{equation}
  \phi = \sum _{i=1}^{n} c_{i} f_{i}
\end{equation}

Y además tenemos:
\begin{equation}
  \sum _{j=1}^{n} \left[\left(H_{ij}-S_{ij}W\right)c_{j}\right]=0, \qquad i=1,2,3,\cdots,n
\end{equation}

Por lo que para cada $W$ podemos calcular el valor de las $c_{i}$ correspondientes

```python
H_SW1 = zeros(n,n) #Matriz (Hij - Sij*W1)cj

c=[] #Matriz de coeficientes c
for i in range (0,n):
    c.append(str('c_%d'  %(i+1)))
c=sympify(c)

for i in range (0,n):
    for j in range (0,n):
        H_SW1[i,j]=((H[i,j]-S[i,j]*Wn[0])*c[j]) #Por lo pronto únicamente se realiza para W1
```

```python
H_SW1 #De esta matriz se obtienen las ecuaciones para obtener el valor de las c
```

$$\left[\begin{matrix}c_{1} \left(- \frac{\hbar^{2} l^{3}}{15 m} \left(- \sqrt{133} + 14\right) + \frac{\hbar^{2} l^{3}}{6 m}\right) & c_{2} \left(- \frac{\hbar^{2} l^{5}}{70 m} \left(- \sqrt{133} + 14\right) + \frac{\hbar^{2} l^{5}}{30 m}\right) & 0 & 0\\c_{1} \left(- \frac{\hbar^{2} l^{5}}{70 m} \left(- \sqrt{133} + 14\right) + \frac{\hbar^{2} l^{5}}{30 m}\right) & c_{2} \left(- \frac{\hbar^{2} l^{7}}{315 m} \left(- \sqrt{133} + 14\right) + \frac{\hbar^{2} l^{7}}{105 m}\right) & 0 & 0\\0 & 0 & c_{3} \left(- \frac{\hbar^{2} l^{5}}{420 m} \left(- \sqrt{133} + 14\right) + \frac{\hbar^{2} l^{5}}{40 m}\right) & c_{4} \left(- \frac{\hbar^{2} l^{7}}{2520 m} \left(- \sqrt{133} + 14\right) + \frac{\hbar^{2} l^{7}}{280 m}\right)\\0 & 0 & c_{3} \left(- \frac{\hbar^{2} l^{7}}{2520 m} \left(- \sqrt{133} + 14\right) + \frac{\hbar^{2} l^{7}}{280 m}\right) & c_{4} \left(- \frac{\hbar^{2} l^{9}}{13860 m} \left(- \sqrt{133} + 14\right) + \frac{\hbar^{2} l^{9}}{1260 m}\right)\end{matrix}\right]$$

```python
eq=[] #Matriz de ecuaciones para obtener las soluciones a las c
for i in range (0,n):
    eqn=0
    for j in range (0,n):
        eqn=eqn+H_SW1[i,j]
    eq.append(eqn)

eq.append(c[0]-k) #se agrega c1=k para establecer un valor de constante que después debe ser normalizado
eq=sympify(eq)

sol=solve(eq,c, dict=true)
#sol[0][c[0]] Forma de acceder a las soluciones
```

```python
eq
```

$$\left [ c_{1} \left(- \frac{\hbar^{2} l^{3}}{15 m} \left(- \sqrt{133} + 14\right) + \frac{\hbar^{2} l^{3}}{6 m}\right) + c_{2} \left(- \frac{\hbar^{2} l^{5}}{70 m} \left(- \sqrt{133} + 14\right) + \frac{\hbar^{2} l^{5}}{30 m}\right), \quad c_{1} \left(- \frac{\hbar^{2} l^{5}}{70 m} \left(- \sqrt{133} + 14\right) + \frac{\hbar^{2} l^{5}}{30 m}\right) + c_{2} \left(- \frac{\hbar^{2} l^{7}}{315 m} \left(- \sqrt{133} + 14\right) + \frac{\hbar^{2} l^{7}}{105 m}\right), \quad c_{3} \left(- \frac{\hbar^{2} l^{5}}{420 m} \left(- \sqrt{133} + 14\right) + \frac{\hbar^{2} l^{5}}{40 m}\right) + c_{4} \left(- \frac{\hbar^{2} l^{7}}{2520 m} \left(- \sqrt{133} + 14\right) + \frac{\hbar^{2} l^{7}}{280 m}\right), \quad c_{3} \left(- \frac{\hbar^{2} l^{7}}{2520 m} \left(- \sqrt{133} + 14\right) + \frac{\hbar^{2} l^{7}}{280 m}\right) + c_{4} \left(- \frac{\hbar^{2} l^{9}}{13860 m} \left(- \sqrt{133} + 14\right) + \frac{\hbar^{2} l^{9}}{1260 m}\right), \quad c_{1} - k\right ]$$

```python
sol
```

$$\left [ \left \{ c_{1} : k, \quad c_{2} : \frac{21 k}{l^{2} \left(7 + \sqrt{133}\right)}, \quad c_{3} : 0, \quad c_{4} : 0\right \}\right ]$$

```python
#Normalizar
Phi1=0
for i in range (0,n):
    Phi1=Phi1 + sol[0][c[i]]*F[i]

Norm=Integral(Phi1**2, (x,0,l)).doit()-1

K=(solve(Norm,k)[0]).evalf() #Valor positivo de k

```

```python
(Phi1.subs(k,K)).evalf()
```

$$- 4.40399751133633 x \left(l - x\right) \left(\frac{1}{l^{5}}\right)^{0.5} - \frac{4.99034859672658 x^{2}}{l^{2}} \left(l - x\right)^{2} \left(\frac{1}{l^{5}}\right)^{0.5}$$

### A continuación se busca generalizar el proceso para obtener todas las $\phi$

```python
#Generalizar para todas las Wn
import numpy as np
H_SWn = np.zeros((n,n,n), dtype=object) #Matriz (Hij - Sij*W1)cj

c_gral=zeros(n,n) #Matriz de coeficientes c
for i in range (0,n):
    for j in range (0,n):
        c_gral[i,j]=sympify((str('c_' +str(i+1) + '_' + str(j+1))))

for z in range (0,n):
    for i in range (0,n):
        for j in range (0,n):
            H_SWn[z,i,j]=sympify(((H[i,j]-S[i,j]*Wn[z])*c_gral[z,j])) #Por lo pronto únicamente se realiza para W1

eq_gral=zeros(n,n+1) #Matriz de ecuaciones para obtener las soluciones a las c
for z in range (0,n):
    for i in range (0,n):
        eqn=0
        for j in range (0,n):
            eqn=eqn+H_SWn[z,i,j]
        eq_gral[z,i]=eqn

#se agrega c1=k para establecer un valor de constante que después debe ser normalizado, se realizó de forma manual
#eq_gral[0,4]=c_gral[0,0]-k
#eq_gral[1,4]=c_gral[1,0]-k
#eq_gral[2,4]=c_gral[2,2]-k
#eq_gral[3,4]=c_gral[3,2]-k
eq_gral=sympify(eq_gral)

sol_gral=[]
for i in range (0,n):
    sol_gral.append(solve(eq_gral[i,:], c_gral[i,:], dict=true))
#sol[0][c[0]] Forma de acceder a las soluciones
```

```python
eq_gral
```

$$\left[\begin{matrix}c_{1 1} \left(- \frac{\hbar^{2} l^{3}}{15 m} \left(- \sqrt{133} + 14\right) + \frac{\hbar^{2} l^{3}}{6 m}\right) + c_{1 2} \left(- \frac{\hbar^{2} l^{5}}{70 m} \left(- \sqrt{133} + 14\right) + \frac{\hbar^{2} l^{5}}{30 m}\right) & c_{1 1} \left(- \frac{\hbar^{2} l^{5}}{70 m} \left(- \sqrt{133} + 14\right) + \frac{\hbar^{2} l^{5}}{30 m}\right) + c_{1 2} \left(- \frac{\hbar^{2} l^{7}}{315 m} \left(- \sqrt{133} + 14\right) + \frac{\hbar^{2} l^{7}}{105 m}\right) & c_{1 3} \left(- \frac{\hbar^{2} l^{5}}{420 m} \left(- \sqrt{133} + 14\right) + \frac{\hbar^{2} l^{5}}{40 m}\right) + c_{1 4} \left(- \frac{\hbar^{2} l^{7}}{2520 m} \left(- \sqrt{133} + 14\right) + \frac{\hbar^{2} l^{7}}{280 m}\right) & c_{1 3} \left(- \frac{\hbar^{2} l^{7}}{2520 m} \left(- \sqrt{133} + 14\right) + \frac{\hbar^{2} l^{7}}{280 m}\right) + c_{1 4} \left(- \frac{\hbar^{2} l^{9}}{13860 m} \left(- \sqrt{133} + 14\right) + \frac{\hbar^{2} l^{9}}{1260 m}\right) & 0\\c_{2 1} \left(- \frac{\hbar^{2} l^{3}}{15 m} \left(\sqrt{133} + 14\right) + \frac{\hbar^{2} l^{3}}{6 m}\right) + c_{2 2} \left(- \frac{\hbar^{2} l^{5}}{70 m} \left(\sqrt{133} + 14\right) + \frac{\hbar^{2} l^{5}}{30 m}\right) & c_{2 1} \left(- \frac{\hbar^{2} l^{5}}{70 m} \left(\sqrt{133} + 14\right) + \frac{\hbar^{2} l^{5}}{30 m}\right) + c_{2 2} \left(- \frac{\hbar^{2} l^{7}}{315 m} \left(\sqrt{133} + 14\right) + \frac{\hbar^{2} l^{7}}{105 m}\right) & c_{2 3} \left(- \frac{\hbar^{2} l^{5}}{420 m} \left(\sqrt{133} + 14\right) + \frac{\hbar^{2} l^{5}}{40 m}\right) + c_{2 4} \left(- \frac{\hbar^{2} l^{7}}{2520 m} \left(\sqrt{133} + 14\right) + \frac{\hbar^{2} l^{7}}{280 m}\right) & c_{2 3} \left(- \frac{\hbar^{2} l^{7}}{2520 m} \left(\sqrt{133} + 14\right) + \frac{\hbar^{2} l^{7}}{280 m}\right) + c_{2 4} \left(- \frac{\hbar^{2} l^{9}}{13860 m} \left(\sqrt{133} + 14\right) + \frac{\hbar^{2} l^{9}}{1260 m}\right) & 0\\c_{3 1} \left(- \frac{\hbar^{2} l^{3}}{15 m} \left(- 9 \sqrt{5} + 30\right) + \frac{\hbar^{2} l^{3}}{6 m}\right) + c_{3 2} \left(- \frac{\hbar^{2} l^{5}}{70 m} \left(- 9 \sqrt{5} + 30\right) + \frac{\hbar^{2} l^{5}}{30 m}\right) & c_{3 1} \left(- \frac{\hbar^{2} l^{5}}{70 m} \left(- 9 \sqrt{5} + 30\right) + \frac{\hbar^{2} l^{5}}{30 m}\right) + c_{3 2} \left(- \frac{\hbar^{2} l^{7}}{315 m} \left(- 9 \sqrt{5} + 30\right) + \frac{\hbar^{2} l^{7}}{105 m}\right) & c_{3 3} \left(- \frac{\hbar^{2} l^{5}}{420 m} \left(- 9 \sqrt{5} + 30\right) + \frac{\hbar^{2} l^{5}}{40 m}\right) + c_{3 4} \left(- \frac{\hbar^{2} l^{7}}{2520 m} \left(- 9 \sqrt{5} + 30\right) + \frac{\hbar^{2} l^{7}}{280 m}\right) & c_{3 3} \left(- \frac{\hbar^{2} l^{7}}{2520 m} \left(- 9 \sqrt{5} + 30\right) + \frac{\hbar^{2} l^{7}}{280 m}\right) + c_{3 4} \left(- \frac{\hbar^{2} l^{9}}{13860 m} \left(- 9 \sqrt{5} + 30\right) + \frac{\hbar^{2} l^{9}}{1260 m}\right) & 0\\c_{4 1} \left(- \frac{\hbar^{2} l^{3}}{15 m} \left(9 \sqrt{5} + 30\right) + \frac{\hbar^{2} l^{3}}{6 m}\right) + c_{4 2} \left(- \frac{\hbar^{2} l^{5}}{70 m} \left(9 \sqrt{5} + 30\right) + \frac{\hbar^{2} l^{5}}{30 m}\right) & c_{4 1} \left(- \frac{\hbar^{2} l^{5}}{70 m} \left(9 \sqrt{5} + 30\right) + \frac{\hbar^{2} l^{5}}{30 m}\right) + c_{4 2} \left(- \frac{\hbar^{2} l^{7}}{315 m} \left(9 \sqrt{5} + 30\right) + \frac{\hbar^{2} l^{7}}{105 m}\right) & c_{4 3} \left(- \frac{\hbar^{2} l^{5}}{420 m} \left(9 \sqrt{5} + 30\right) + \frac{\hbar^{2} l^{5}}{40 m}\right) + c_{4 4} \left(- \frac{\hbar^{2} l^{7}}{2520 m} \left(9 \sqrt{5} + 30\right) + \frac{\hbar^{2} l^{7}}{280 m}\right) & c_{4 3} \left(- \frac{\hbar^{2} l^{7}}{2520 m} \left(9 \sqrt{5} + 30\right) + \frac{\hbar^{2} l^{7}}{280 m}\right) + c_{4 4} \left(- \frac{\hbar^{2} l^{9}}{13860 m} \left(9 \sqrt{5} + 30\right) + \frac{\hbar^{2} l^{9}}{1260 m}\right) & 0\end{matrix}\right]$$

```python
sol_gral
```

$$\left [ \left [ \left \{ c_{1 1} : \frac{c_{1 2}}{3} l^{2} + \frac{c_{1 2}}{21} \sqrt{133} l^{2}, \quad c_{1 3} : 0, \quad c_{1 4} : 0\right \}\right ], \quad \left [ \left \{ c_{2 1} : - \frac{c_{2 2}}{21} \sqrt{133} l^{2} + \frac{c_{2 2}}{3} l^{2}, \quad c_{2 3} : 0, \quad c_{2 4} : 0\right \}\right ], \quad \left [ \left \{ c_{3 1} : 0, \quad c_{3 2} : 0, \quad c_{3 3} : \frac{c_{3 4}}{33} l^{2} + \frac{c_{3 4}}{11} \sqrt{5} l^{2}\right \}\right ], \quad \left [ \left \{ c_{4 1} : 0, \quad c_{4 2} : 0, \quad c_{4 3} : - \frac{c_{4 4}}{11} \sqrt{5} l^{2} + \frac{c_{4 4}}{33} l^{2}\right \}\right ]\right ]$$

```python
soluciones=zeros(n,n)
for i in range(0, n):
    for j in range (0,n):
        if sympify('c_'+str(i+1)+'_'+str(j+1)) in sol_gral[i][0]:
            soluciones[i,j]= sol_gral[i][0][c_gral[i,j]]
        else:
            #soluciones[i,j] = c_gral[i,j]
            soluciones[i,j] = k
            soluciones[i,:]=soluciones[i,:].subs(c_gral[i,j],k)
soluciones
```

$$\left[\begin{matrix}\frac{k l^{2}}{3} + \frac{\sqrt{133} k}{21} l^{2} & k & 0 & 0\\- \frac{\sqrt{133} k}{21} l^{2} + \frac{k l^{2}}{3} & k & 0 & 0\\0 & 0 & \frac{k l^{2}}{33} + \frac{\sqrt{5} k}{11} l^{2} & k\\0 & 0 & - \frac{\sqrt{5} k}{11} l^{2} + \frac{k l^{2}}{33} & k\end{matrix}\right]$$

```python
#Normalizar
Phi_n=zeros(n,1)
Norm_gral=zeros(n,1)
for j in range (0,n):
    for i in range (0,n):
        Phi_n[j,0]=Phi_n[j,0] + soluciones[j,i]*F[i]

for i in range (0,n):
    Norm_gral[i,0]=Integral(Phi_n[i,0]**2, (x,0,l)).doit()-1
    K=(solve(Norm_gral[i,0],k)[0]).evalf() #Valor positivo de k
    Phi_n[i,0]=(Phi_n[i,0].subs(k,K)).evalf()
```

```python
Phi_n
```

$$\left[\begin{matrix}- 4.40399751133633 l^{2} x \left(l - x\right) \left(\frac{1}{l^{9}}\right)^{0.5} - 4.99034859672658 x^{2} \left(l - x\right)^{2} \left(\frac{1}{l^{9}}\right)^{0.5}\\28.6462005494649 l^{2} x \left(l - x\right) \left(\frac{1}{l^{9}}\right)^{0.5} - 132.721876195613 x^{2} \left(l - x\right)^{2} \left(\frac{1}{l^{9}}\right)^{0.5}\\- 16.7823521557751 l^{2} x \left(0.5 l - x\right) \left(l - x\right) \left(\frac{1}{l^{11}}\right)^{0.5} - 71.8478164291383 x^{2} \left(0.5 l - x\right) \left(l - x\right)^{2} \left(\frac{1}{l^{11}}\right)^{0.5}\\98.9866286733697 l^{2} x \left(0.5 l - x\right) \left(l - x\right) \left(\frac{1}{l^{11}}\right)^{0.5} - 572.256840303692 x^{2} \left(0.5 l - x\right) \left(l - x\right)^{2} \left(\frac{1}{l^{11}}\right)^{0.5}\end{matrix}\right]$$


# 03 - Full-Waveform Inversion (FWI)

This notebook is the third in a series of tutorial highlighting various aspects of seismic inversion based on Devito operators. In this second example we aim to highlight the core ideas behind seismic inversion, where we create an image of the subsurface from field recorded data. This tutorial follows on the modelling tutorial and will reuse the modelling and velocity model.

## Inversion requirement

Seismic inversion relies on two known parameters:

- **Field data** - or also called **recorded data**. This is a shot record corresponding to the true velocity model. In practice this data is acquired as described in the first tutorial. In order to simplify this tutorial we will fake field data by modelling it with the true velocity model.

- **Initial velocity model**. This is a velocity model that has been obtained by processing the field data. This model is a rough and very smooth estimate of the velocity as an initial estimate for the inversion. This is a necessary requirement for any optimization (method).

## Inversion computational setup

In this tutorial, we will introduce the gradient operator. This operator corresponds to the imaging condition introduced in the previous tutorial with some minor modifications that are defined by the objective function (also referred to in the tutorial series as the *functional*, *f*) and its gradient, *g*. We will define these two terms in the tutorial too.

## Notes on the operators

As we have already described the creation of a forward modelling operator, we will only call a wrapper function here. This wrapper already contains all the necessary operators for seismic modeling, imaging and inversion. Operators introduced for the first time in this tutorial will be properly described.

```python
import numpy as np
%matplotlib inline

from devito import configuration
configuration['log-level'] = 'WARNING'
```

## Computational considerations

As we will see, FWI is computationally extremely demanding, even more than RTM. To keep this tutorial as lightwight as possible we therefore again use a very small demonstration model. We also define here a few parameters for the final example runs that can be changed to modify the overall runtime of the tutorial.

```python
nshots = 9  # Number of shots to create gradient from
nreceivers = 101  # Number of receiver locations per shot
fwi_iterations = 5  # Number of outer FWI iterations
```

# True and smooth velocity models

We will use a very simple model domain, consisting of a circle within a 2D domain. We will again use the "true" model to generate our synthetic shot data and use a "smooth" model as our initial guess. In this case the smooth model is very smooth indeed - it is simply a constant background velocity without any features.

```python
#NBVAL_IGNORE_OUTPUT
from examples.seismic import demo_model, plot_velocity, plot_perturbation

# Define true and initial model
shape = (101, 101)  # Number of grid point (nx, nz)
spacing = (10., 10.)  # Grid spacing in m. The domain size is now 1km by 1km
origin = (0., 0.)  # Need origin to define relative source and receiver locations

model = demo_model('circle-isotropic', vp=3.0, vp_background=2.5,
                    origin=origin, shape=shape, spacing=spacing, nbpml=40)

model0 = demo_model('circle-isotropic', vp=2.5, vp_background=2.5,
                     origin=origin, shape=shape, spacing=spacing, nbpml=40,
                    grid = model.grid)

plot_velocity(model)
plot_velocity(model0)
plot_perturbation(model0, model)
```

## Acquisition geometry

In this tutorial, we will use the easiest case for inversion, namely a transmission experiment. The sources are located on one side of the model and the receivers on the other side. This allows to record most of the information necessary for inversion, as reflections usually lead to poor inversion results.

```python
#NBVAL_IGNORE_OUTPUT
# Define acquisition geometry: source
from examples.seismic import AcquisitionGeometry

t0 = 0.
tn = 1000.
f0 = 0.010
# First, position source centrally in all dimensions, then set depth
src_coordinates = np.empty((1, 2))
src_coordinates[0, :] = np.array(model.domain_size) * .5
src_coordinates[0, 0] = 20.  # Depth is 20m

# Define acquisition geometry: receivers

# Initialize receivers for synthetic and imaging data
rec_coordinates = np.empty((nreceivers, 2))
rec_coordinates[:, 1] = np.linspace(0, model.domain_size[0], num=nreceivers)
rec_coordinates[:, 0] = 980.

# Geometry

geometry = AcquisitionGeometry(model, rec_coordinates, src_coordinates, t0, tn, f0=f0, src_type='Ricker')
# We can plot the time signature to see the wavelet
geometry.src.show()
```

```python
# Plot acquisition geometry
plot_velocity(model, source=geometry.src_positions,
              receiver=geometry.rec_positions[::4, :])
```

## True and smooth data

We can generate shot records for the true and smoothed initial velocity models, since the difference between them will again form the basis of our imaging procedure.

```python
# Compute synthetic data with forward operator
from examples.seismic.acoustic import AcousticWaveSolver

solver = AcousticWaveSolver(model, geometry, space_order=4)
true_d, _, _ = solver.forward(vp=model.vp)
```

```python
# Compute initial data with forward operator
smooth_d, _, _ = solver.forward(vp=model0.vp)
```

```python
#NBVAL_IGNORE_OUTPUT
from examples.seismic import plot_shotrecord

# Plot shot record for true and smooth velocity model and the difference
plot_shotrecord(true_d.data, model, t0, tn)
plot_shotrecord(smooth_d.data, model, t0, tn)
plot_shotrecord(smooth_d.data - true_d.data, model, t0, tn)
```

# Full-Waveform Inversion

## Formulation

Full-waveform inversion (FWI) aims to invert an accurate model of the discrete wave velocity, $\mathbf{c}$, or equivalently the square slowness of the wave, $\mathbf{m} = \frac{1}{\mathbf{c}^2}$, from a given set of measurements of the pressure wavefield $\mathbf{u}$. This can be expressed as the following optimization problem [1, 2]:

\begin{aligned}
	\mathop{\hbox{minimize}}_{\mathbf{m}} \Phi_s(\mathbf{m})&=\frac{1}{2}\left\lVert\mathbf{P}_r
	\mathbf{u} - \mathbf{d}\right\rVert_2^2 \\
	\mathbf{u} &= \mathbf{A}(\mathbf{m})^{-1} \mathbf{P}_s^T \mathbf{q}_s,
\end{aligned}

where $\mathbf{P}_r$ is the sampling operator at the receiver locations, $\mathbf{P}_s^T$ is the injection operator at the source locations, $\mathbf{A}(\mathbf{m})$ is the operator representing the discretized wave equation matrix, $\mathbf{u}$ is the discrete synthetic pressure wavefield, $\mathbf{q}_s$ is the corresponding pressure source and $\mathbf{d}$ is the measured data. It is worth noting that $\mathbf{m}$ is the unknown in this formulation and that multiple implementations of the wave equation operator $\mathbf{A}(\mathbf{m})$ are possible.

We have already defined a concrete solver scheme for $\mathbf{A}(\mathbf{m})$ in the first tutorial, including appropriate implementations of the sampling operator $\mathbf{P}_r$ and source term $\mathbf{q}_s$.

To solve this optimization problem using a gradient-based method, we use the
adjoint-state method to evaluate the gradient $\nabla\Phi_s(\mathbf{m})$:

\begin{align}
 \nabla\Phi_s(\mathbf{m})=\sum_{\mathbf{t} =1}^{n_t}\mathbf{u}[\mathbf{t}] \mathbf{v}_{tt}[\mathbf{t}] =\mathbf{J}^T\delta\mathbf{d}_s,
\end{align}

where $n_t$ is the number of computational time steps, $\delta\mathbf{d}_s = \left(\mathbf{P}_r \mathbf{u} - \mathbf{d} \right)$ is the data residual (difference between the measured data and the modelled data), $\mathbf{J}$ is the Jacobian operator and $\mathbf{v}_{tt}$ is the second-order time derivative of the adjoint wavefield solving:

\begin{align}
 \mathbf{A}^T(\mathbf{m}) \mathbf{v} = \mathbf{P}_r^T \delta\mathbf{d}.
\end{align}

We see that the gradient of the FWI function is the previously defined imaging condition with an extra second-order time derivative. We will therefore reuse the operators defined previously inside a Devito wrapper.

## FWI gradient operator

To compute a single gradient $\nabla\Phi_s(\mathbf{m})$ in our optimization workflow we again use `solver.forward` to compute the entire forward wavefield $\mathbf{u}$ and a similar pre-defined gradient operator to compute the adjoint wavefield `v`. The gradient operator provided by our `solver` utility also computes the correlation between the wavefields, allowing us to encode a similar procedure to the previous imaging tutorial as our gradient calculation:

- Simulate the forward wavefield with the background velocity model to get the synthetic data and save the full wavefield $\mathbf{u}$
- Compute the data residual
- Back-propagate the data residual and compute on the fly the gradient contribution at each time step.

This procedure is applied to multiple source positions and summed to obtain a gradient image of the subsurface. We again prepare the source locations for each shot and visualize them, before defining a single gradient computation over a number of shots as a single function.

```python
#NBVAL_IGNORE_OUTPUT

# Prepare the varying source locations sources
source_locations = np.empty((nshots, 2), dtype=np.float32)
source_locations[:, 0] = 30.
source_locations[:, 1] = np.linspace(0., 1000, num=nshots)

plot_velocity(model, source=source_locations)
```

```python
# Create FWI gradient kernel
from devito import Function, clear_cache, TimeFunction
from examples.seismic import Receiver

import scipy
def fwi_gradient(vp_in):
    # Create symbols to hold the gradient and residual
    grad = Function(name="grad", grid=model.grid)
    residual = Receiver(name='rec', grid=model.grid,
                        time_range=geometry.time_axis,
                        coordinates=geometry.rec_positions)
    objective = 0.

    # Creat forward wavefield to reuse to avoid memory overload
    u0 = TimeFunction(name='u', grid=model.grid, time_order=2, space_order=4,
                      save=geometry.nt)
    for i in range(nshots):
        # Important: We force previous wavefields to be destroyed,
        # so that we may reuse the memory.
        clear_cache()

        # Update source location
        geometry.src_positions[0, :] = source_locations[i, :]

        # Generate synthetic data from true model
        true_d, _, _ = solver.forward(vp=model.vp)

        # Compute smooth data and full forward wavefield u0
        u0.data.fill(0.)
        smooth_d, _, _ = solver.forward(vp=vp_in, save=True, u=u0)

        # Compute gradient from data residual and update objective function
        residual.data[:] = smooth_d.data[:] - true_d.data[:]

        objective += .5*np.linalg.norm(residual.data.flatten())**2
        solver.gradient(rec=residual, u=u0, vp=vp_in, grad=grad)

    return objective, -grad.data
```

Having defined our FWI gradient procedure we can compute the initial iteration from our starting model. This allows us to visualize the gradient alongside the model perturbation and the effect of the gradient update on the model.

```python
#NBVAL_IGNORE_OUTPUT

# Compute gradient of initial model
ff, update = fwi_gradient(model0.vp)
print('Objective value is %f ' % ff)
```

    Objective value is 56929.824697

```python
#NBVAL_IGNORE_OUTPUT
from examples.seismic import plot_image

# Plot the FWI gradient
plot_image(update, vmin=-1e4, vmax=1e4, cmap="jet")

# Plot the difference between the true and initial model.
# This is not known in practice as only the initial model is provided.
plot_image(model0.vp.data - model.vp.data, vmin=-1e-1, vmax=1e-1, cmap="jet")

# Show what the update does to the model
alpha = .5 / np.abs(update).max()
plot_image(model0.vp.data - alpha*update, vmin=2.5, vmax=3.0, cmap="jet")
```

We see that the gradient and the true perturbation have the same sign, therefore, with an appropriate scaling factor, we will update the model in the correct direction.

```python
# Define bounding box constraints on the solution.
def apply_box_constraint(vp):
    # Maximum possible 'realistic' velocity is 3.5 km/sec
    # Minimum possible 'realistic' velocity is 2 km/sec
    return np.clip(vp, 2.0, 3.5)
```

```python
#NBVAL_SKIP

# Run FWI with gradient descent
history = np.zeros((fwi_iterations, 1))
for i in range(0, fwi_iterations):
    # Compute the functional value and gradient for the current
    # model estimate
    phi, direction = fwi_gradient(model0.vp)

    # Store the history of the functional values
    history[i] = phi

    # Artificial Step length for gradient descent
    # In practice this would be replaced by a Linesearch (Wolfe, ...)
    # that would guarantee functional decrease Phi(m-alpha g) <= epsilon Phi(m)
    # where epsilon is a minimum decrease constant
    alpha = .05 / np.abs(direction).max()

    # Update the model estimate and enforce minimum/maximum values
    model0.vp = apply_box_constraint(model0.vp.data - alpha * direction)

    # Log the progress made
    print('Objective value is %f at iteration %d' % (phi, i+1))
```

    Objective value is 56929.824697 at iteration 1
    Objective value is 35754.861855 at iteration 2
    Objective value is 21001.778796 at iteration 3
    Objective value is 11548.171734 at iteration 4
    Objective value is 5790.152320 at iteration 5

```python
#NBVAL_IGNORE_OUTPUT

# Plot inverted velocity model
plot_velocity(model0)
```

```python
#NBVAL_SKIP
import matplotlib.pyplot as plt

# Plot objective function decrease
plt.figure()
plt.loglog(history)
plt.xlabel('Iteration number')
plt.ylabel('Misift value Phi')
plt.title('Convergence')
plt.show()
```

## References

[1] _Virieux, J. and Operto, S.: An overview of full-waveform inversion in exploration geophysics, GEOPHYSICS, 74, WCC1–WCC26, doi:10.1190/1.3238367, http://library.seg.org/doi/abs/10.1190/1.3238367, 2009._

[2] _Haber, E., Chung, M., and Herrmann, F. J.: An effective method for parameter estimation with PDE constraints with multiple right hand sides, SIAM Journal on Optimization, 22, http://dx.doi.org/10.1137/11081126X, 2012._

This notebook is part of the tutorial "Optimised Symbolic Finite Difference Computation with Devito" presented at the Intel® HPC Developer Conference 2017.


# FEEDBACK CONTROL EXAMPLE

This notebook explores feedback control for the glycolytic oscillations model of Jenna Wolf.

# Preliminaries

```python
from controlSBML.control_extensions.state_space_tf import StateSpaceTF
import controlSBML as ctl

import control
import numpy as np
import pandas as pd
import sympy
import tellurium as te
import matplotlib.pyplot as plt
```

```python
TIMES = [0.1*v for v in range(51)]
```

```python
CTLSB = ctl.ControlSBML("https://www.ebi.ac.uk/biomodels/model/download/BIOMD0000000206.2?filename=BIOMD0000000206_url.xml",
                       input_names=["at"])
```

```python
CTLSB.plotTrueModel(ylim=[0, 10], figsize=(5,5))
```

```python
# at: ATP
# na: NAD+
# s4: 3-Phospho-D-glycerate
# s5: pyruvic acid
print(CTLSB.antimony)
```

    // Created by libAntimony v2.12.0
    model *Wolf2000_Glycolytic_Oscillations()

      // Compartments and Species:
      compartment compartment_;
      species s1 in compartment_, at_ in compartment_, s2 in compartment_, s3 in compartment_;
      species na in compartment_, s4 in compartment_, s5 in compartment_, s6 in compartment_;
      species s6o in compartment_;

      // Reactions:
      v1: s1 + 2 at_ -> s2; compartment_*k1*s1*at_/(1 + (at_/ki)^n);
      v2: s2 -> 2 s3; compartment_*k2*s2;
      v3: s3 + na -> s4 + at_; compartment_*((k31*k32*s3*na*(atot - at_) - k33*k34*s4*at_*(ntot - na))/(k33*(ntot - na) + k32*(atot - at_)));
      v4: s4 -> s5 + at_; compartment_*k4*s4*(atot - at_);
      v5: s5 -> s6; compartment_*k5*s5;
      v7: at_ -> ; compartment_*k7*at_;
      v8: s3 -> na; compartment_*k8*s3*(ntot - na);
      v9: s6o -> ; compartment_*k9*s6o;
      v10: s6 -> 0.1 s6o; compartment_*k10*(s6 - s6o);
      v6: s6 -> na; compartment_*k6*s6*(ntot - na);
      v0:  -> s1; compartment_*k0;

      // Species initializations:
      s1 = 1;
      at_ = 2;
      s2 = 5;
      s3 = 0.6;
      na = 0.6;
      s4 = 0.7;
      s5 = 8;
      s6 = 0.08;
      s6o = 0.02;

      // Compartment initializations:
      compartment_ = 1;

      // Variable initializations:
      k0 = 50;
      k0 has mM_min_1;
      k1 = 550;
      k1 has mM_1_min_1;
      k2 = 9.8;
      k2 has min_1;
      k31 = 323.8;
      k31 has mM_1_min_1;
      k33 = 57823.1;
      k33 has mM_1_min_1;
      k32 = 76411.1;
      k32 has mM_1_min_1;
      k34 = 23.7;
      k34 has mM_1_min_1;
      k4 = 80;
      k4 has mM_1_min_1;
      k5 = 9.7;
      k5 has min_1;
      k6 = 2000;
      k6 has mM_1_min_1;
      k7 = 28;
      k7 has min_1;
      k8 = 85.7;
      k8 has mM_1_min_1;
      k9 = 80;
      k9 has min_1;
      k10 = 375;
      k10 has min_1;
      atot = 4;
      atot has mM;
      ntot = 1;
      ntot has mM;
      n = 4;
      n has dimensionless;
      ki = 1;
      ki has mM;

      // Other declarations:
      const compartment_, k0, k1, k2, k31, k33, k32, k34, k4, k5, k6, k7, k8;
      const k9, k10, atot, ntot, n, ki;

      // Unit definitions:
      unit substance = 1e-3 mole;
      unit time_unit = 60 second;
      unit mM = 1e-3 mole / litre;
      unit mM_min_1 = 1e-3 mole / (litre * 60 second);
      unit min_1 = 1 / 60 second;
      unit mM_1_min_1 = litre / (1e-3 mole * 60 second);

      // Display Names:
      substance is "milli mole";
      time_unit is "min";
      s1 is "Glucose";
      at_ is "ATP";
      s2 is "F16P";
      s3 is "Triose_Gly3Phos_DHAP";
      na is "NAD";
      s4 is "3PG";
      s5 is "Pyruvate";
      s6 is "Acetaldehyde";
      s6o is "extracellular acetaldehyde";

      // CV terms:
      s1 identity "http://identifiers.org/obo.chebi/CHEBI:17234",
                  "http://identifiers.org/kegg.compound/C00293"
      at_ identity "http://identifiers.org/obo.chebi/CHEBI:15422",
                   "http://identifiers.org/kegg.compound/C00002"
      s2 identity "http://identifiers.org/obo.chebi/CHEBI:16905",
                  "http://identifiers.org/kegg.compound/C05378"
      s3 part "http://identifiers.org/obo.chebi/CHEBI:16108",
              "http://identifiers.org/obo.chebi/CHEBI:29052",
              "http://identifiers.org/kegg.compound/C00111",
              "http://identifiers.org/kegg.compound/C00118"
      na identity "http://identifiers.org/obo.chebi/CHEBI:15846",
                  "http://identifiers.org/kegg.compound/C00003"
      s4 identity "http://identifiers.org/obo.chebi/CHEBI:17794",
                  "http://identifiers.org/kegg.compound/C00197"
      s5 identity "http://identifiers.org/kegg.compound/C00022",
                  "http://identifiers.org/chebi/CHEBI:32816"
      s5 identity "http://identifiers.org/obo.chebi/CHEBI:15361"
      s6 identity "http://identifiers.org/obo.chebi/CHEBI:15343",
                  "http://identifiers.org/kegg.compound/C00084"
      s6o identity "http://identifiers.org/obo.chebi/CHEBI:15343",
                   "http://identifiers.org/kegg.compound/C00084"
      v1 part "http://identifiers.org/kegg.reaction/R00756",
              "http://identifiers.org/kegg.reaction/R00299"
      v1 part "http://identifiers.org/ec-code/2.7.1.11",
              "http://identifiers.org/ec-code/5.3.1.9"
      v2 identity "http://identifiers.org/ec-code/4.1.2.13",
                  "http://identifiers.org/kegg.reaction/R01070"
      v3 part "http://identifiers.org/ec-code/2.7.2.3",
              "http://identifiers.org/ec-code/1.2.1.12"
      v4 part "http://identifiers.org/kegg.reaction/R00658",
              "http://identifiers.org/kegg.reaction/R00200"
      v4 part "http://identifiers.org/ec-code/2.7.1.40"
      v5 hypernym "http://identifiers.org/ec-code/4.1.1.1",
                  "http://identifiers.org/kegg.reaction/R00224"
      v7 hypernym "http://identifiers.org/obo.go/GO:0006754"
      v8 part "http://identifiers.org/kegg.reaction/R05679"
      v9 hypernym "http://identifiers.org/obo.go/GO:0046187"
      v6 hypernym "http://identifiers.org/ec-code/1.1.1.71",
                  "http://identifiers.org/kegg.reaction/R00754"
      v0 hypernym "http://identifiers.org/obo.go/GO:0046323"
    end

    Wolf2000_Glycolytic_Oscillations is "Wolf2000_Glycolytic_Oscillations"

    Wolf2000_Glycolytic_Oscillations model_entity_is "http://identifiers.org/biomodels.db/MODEL3352181362"
    Wolf2000_Glycolytic_Oscillations model_entity_is "http://identifiers.org/biomodels.db/BIOMD0000000206"
    Wolf2000_Glycolytic_Oscillations description "http://identifiers.org/pubmed/10692304"
    Wolf2000_Glycolytic_Oscillations taxon "http://identifiers.org/taxonomy/4932"
    Wolf2000_Glycolytic_Oscillations hypernym "http://identifiers.org/obo.go/GO:0006096"
    Wolf2000_Glycolytic_Oscillations identity "http://identifiers.org/kegg.pathway/sce00010"

# Controller Construction

## Reaction Network

    v1: s1 + 2 at_ -> s2; compartment_*k1*s1*at_/(1 + (at_/ki)^n);
    v2: s2 -> 2 s3; compartment_*k2*s2;
    v3: s3 + na -> s4 + at_; compartment_*((k31*k32*s3*na*(atot - at_)
                             - k33*k34*s4*at_*(ntot - na))/(k33*(ntot - na) + k32*(atot - at_)));
    v4: s4 -> s5 + at_; compartment_*k4*s4*(atot - at_);
    v5: s5 -> s6; compartment_*k5*s5;
    v7: at_ -> ; compartment_*k7*at_;
    v8: s3 -> na; compartment_*k8*s3*(ntot - na);
    v9: s6o -> ; compartment_*k9*s6o;
    v10: s6 -> 0.1 s6o; compartment_*k10*(s6 - s6o);
    v6: s6 -> na; compartment_*k6*s6*(ntot - na);
    v0:  -> s1; compartment_*k0;

## Regulating Pyruvate (``s5``)

```python
URL_206 = "https://www.ebi.ac.uk/biomodels/model/download/BIOMD0000000206.2?filename=BIOMD0000000206_url.xml"
CTLSB = ctl.ControlSBML(URL_206, input_names=["at"], output_names=["s5"])
CTLSB.plotTrueModel(names=["s5"], ylim=[0, 10], title="Uncontrolled", figsize=(5,2))
```

```python
CTLSB = ctl.ControlSBML(URL_206, input_names=["at"], output_names=["s5"])
sys = CTLSB.makeNonlinearIOSystem("sys")
```

```python
ref = [6]  # Desired value of pyruvate
def outfcn(t, x, u, _):
    return 0.01*(ref - u)
controller = control.NonlinearIOSystem(
  None,
  outfcn,
  inputs=['in'],
  outputs=['out'], name='controller')
```

```python
# Create the closed loop system
closed = control.interconnect(
  [sys, controller],       # systems
  connections=[
    ['sys.at', 'controller.out'],
    ['controller.in',  'sys.s5'],
  ],
  inplist=["controller.in"],
  outlist=["sys.s5"],
)
```

```python
times = TIMES
result = control.input_output_response(closed, times, X0=ctl.makeStateVector(closed))

ser = ctl.mat2DF(result.outputs, row_names=times)
_, ax = plt.subplots(1, figsize=(5,2))
ax.plot(ser.index, ser.values)
ax.set_ylim([0, 10])
ax.set_xlabel("time")
_ = ax.set_title("Controlled")
ax.legend(["s5"], fontsize=14)
```


# A Toolbox for Solving Dynamic Rational Inattention Problems

This notebook explains how to solve LQG Dynamic Rational Inattention models using the methods and the [solver](https://github.com/choongryulyang/dynamic_multivariate_RI) from [Afrouzi and Yang (2019)](http://www.afrouzi.com/dynamic_inattention.pdf) that was adapted for Julia by [Miguel Acosta](https://www.acostamiguel.com/home). To include the solver, execute:

```julia
include("DRIP.jl");
```

## A Structure for LQG Dynamic Rational Inattention Problems (D.R.I.P.)

### The Problem
A LQG Dynamic Rational Inattention Problem (**drip**) is defined as the following tracking problem, where at any point in time the agent chooses a vector of actions $\vec{a}_t\in\mathbb{R}^m$ to track a Gaussian stochastic process $\vec{x}_t\in\mathbb{R}^n$:
\begin{align}
    & \min_{\{\vec{a}_t\}_{t\geq 0}} \mathbb{E}\left[\sum_{t=0}^\infty \beta^t (\vec{a}_t - \vec{x}_t'\mathbf{H})'(\vec{a}_t - \vec{x}_t'\mathbf{H}) - \omega \mathbb{I}(\vec{a}^t;\vec{x}^t|\vec{a}^{t-1})\lvert \vec{a}^{-1}\right] \\
    s.t.\quad &
        \vec{x}_t=\mathbf{A}\vec{x}_{t-1}+\mathbf{Q}\vec{u}_t,\quad \vec{u}_t\sim \mathcal{N}(\mathbf{0},\mathbf{I}^{k\times k})
\end{align}
Here:
* $\vec{a}_t\in\mathbb{R}^m$ is the vector of the agent's actions at time $t$ (a firms choosing a price, or a househld choosing consumption and labor etc.) We denote the number of actions by $m$.
* $\vec{x}_t\in\mathbb{R}^n$ is the vector of the shocks that the agent faces at time $t$ that are exogenous to her decision, but could be endogenous to the GE model (marginal cost of a firm, real interest rates etc.) We denote the number of shocks by $n$.

#### The Parameters
The LQG-DRI problem is characterized by the following parameters:
* $\omega\in \mathbb{R}_+$: cost of 1 bit of information in units of the agent's payoff.
* $\beta\in[0,1]$: rate of discounting information.
* $\mathbf{A}\in \mathbb{R}^{n\times n}, \mathbf{Q}\in\mathbb{R}^{n\times k}$: Determine the state space representation of $\vec{x}_t$.
* $\mathbf{H}\in \mathbb{R}^{n\times m}$: interaction of payoffs with shocks. This comes from a second order approximation to the utility function and is such that under full information $\vec{a}^*=\mathbf{H}'\vec{x}$.

#### The Solution
The solution to the dynamic rational inattention problem is a joint stochastic process between the actions and the states: $\{(\vec{a}_t,\vec{x}_t):t\geq 0\}$. Moreover, in some economic applications, we are also interested in the law of motion for the agent's belief about $\vec{x}_t$ under the optimal information structure $\hat{x}_t=\mathbb{E}_t[\vec{x}_t]$ where the expectation is taken conditional on the agent's time $t$ information.

Theorem 2 and Proposition 3 in [Afrouzi and Yang (2019)](http://www.afrouzi.com/dynamic_inattention.pdf) characterize this joint distribution as a function of a tuple $(\mathbf{K},\mathbf{Y},\mathbf{\Sigma}_z)$ where
\begin{align}
    \vec{a}_t &= \mathbf{H}'\hat{x}_t = \mathbf{H}'\mathbf{A}\hat{x}_{t-1}+\mathbf{Y}'(\vec{x}_t-\mathbf{A}\hat{x}_{t-1})+\vec{z}_t \\
    \hat{x}_t &= \mathbf{A}\hat{x}_{t-1}+\mathbf{K}\mathbf{Y}'(\vec{x}_t-\mathbf{A}\hat{x}_{t-1})+\mathbf{K}\vec{z}_t,\quad \vec{z}\sim\mathcal{N}(0,\Sigma_z)
\end{align}

Here,

* $\mathbf{K}\in\mathbb{R}^{n\times m}$ is the Kalman-gain matrix of the agent in under optimal information acquisition in the stationary solution.
* $\mathbf{Y}\in\mathbb{R}^{m\times m}$ maps the state to the agent's action under rational inattetion and govern the covariance of the two.
* $\mathbf{\Sigma}_z\in\mathbb{R}^{m\times m}$ is the variance-covariance matrix of the agent's rational inattention error.

In addition to these, we might also be interested in the agent's prior and posterior subjective undertainty, along with the continution value that she assigns to information:
* $\mathbf{\Sigma}_{p}=\mathbb{V}ar(\vec{x}_t|\vec{a}^{t})\in\mathbb{R}^{n\times n}$.
* $\mathbf{\Sigma}_{-1}=\mathbb{V}ar(\vec{x}_t|\vec{a}^{t-1})\in\mathbb{R}^{n\times n}$,
* $\bar{\Omega}\in\mathbb{R}^{n\times n}$.

The matrix $\bar{\Omega}$, which is the continuation value of information, is also important for the number of iterations that the code needs to converge. When $n$ is large, accurate guesses for $\bar{\Omega}$ and $\mathbf{\Sigma}_{-1}$ speed up the convergence considerably.

### The Solver for the Steady State of D.R.I.P.

The solver function is `solve_drip(ω,β,A,Q,H;Σ0 = A*A'+Q*Q',Ω0 = H*H')`. It takes the primitives `(ω,β,A,Q,H;Σ0,Ω0)` as arguments, where the arguments `Σ0,Ω0` are the initial guesses for $\mathbf{\Sigma}_{-1}$ and $\bar{{\Omega}}$ and are optional for cases in which an accurate initial guess is at hand. If not specified, the default values are set as `Ω0=H*H'` and `Σ0=Q*Q'`. See the code for extra options regarding tolerance of convergence, maximum number of iterations, and updating weight for every iteration.

The function returns the solution of the model within a `drip` structure,  that contains all the primitives and the solution objects of the model:

```julia
mutable struct Drip
    ω; β; A; Q; H;                    # primitives
    K; Y; Σ_z; Σ_p; Σ_1; Ω;           # solution
    Drip()          = new();
    Drip(ω,β,A,Q,H) = new(ω,β,A,Q,H);
end
```

### Impulse Resonse Functions using Steady State Information Structure

Once the model is solved, one can generate the impulse response functions of actions and beliefs using the laws of motion stated above.

We have also included a built-in function that generates these IRFs. The function `dripirfs(P::drip,T::Int)` takes a `drip` structure as input along with the length of irfs `T` and returns the irfs of the state, beliefs and actions to all possible shocks in the following structure:

```julia
struct Dripirfs
    T     :: Int            # length of IRFs
    x     :: Array{Float64} # IRFs of the fundamental shocks
    x_hat :: Array{Float64} # IRFs of beliefs
    a     :: Array{Float64} # IRFs of actions
end
```

Here
* `x` is of dimension $n\times k \times T$ where `x[i,j,:]` is the IRF of the $i$'th element of the state vector ($\vec{x}_t$) to the $j$'th shock in the vector $\vec{u}_t$.

* `x_hat` is of dimension $n\times k \times T$ where `x_hat[i,j,:]` is the IRF of the agent's belief about $i$'th element of the state vector  ($\hat{x}_t$) to the $j$'th shock in the vector $\vec{u}_t$.

* `x_hat` is of dimension $m\times k \times T$ where `a[i,j,:]` is the IRF of the $i$'th element of the action vector ($\vec{a}_t$) to the $j$'th shock in the vector $\vec{u}_t$.

### The Solver for the Transition dynamics of Rational Inattention Problems (T.R.I.P.)

The Euler equation derived in [Afrouzi and Yang (2019)](http://www.afrouzi.com/dynamic_inattention.pdf) for the  also allows us to characterize the transition path of the information structure over time for an arbitrary initial prior. The function `solve_trip(P::drip; T=100)` takes the solution `P=solve_drip(.)` as an input and returns a `trip` structure that summarizes the transition path of information benefit matrics, priors and a vector for the marginal values of information over time (the input `T` determines the number of periods that it takes to reach to the steady state and is set to 100 by default):

```julia
struct Trip
    P::Drip;                # problem and solution in steady state
    T::Int;                 # length of T.R.I.P.
    Σ_1s; Σ_ps; Ωs;         # priors, posteriors and benefit matrices
    Ds;                     # eigenvalues of Σ_t^(0.5)Ω_tΣ_t^(0.5)
end
```


# Linear Gaussian filtering and smoothing

Provided are two examples of linear 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. [Linear, Discrete State-Space Model](#1.-Linear-Discrete-State-Space-Model:-Car-Tracking): Car Tracking
2. [Linear, Continuous-Discrete State-Space Model](#2.-Linear-Continuous-Discrete-State-Space-Model:-Ornstein-Uhlenbeck-Process): The Ornstein-Uhlenbeck Process

**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, statespace, 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")
```

## 1. **Linear Discrete** State-Space Model: Car Tracking

---

We begin showcasing the arguably most simple case in which we consider the following state-space model. Consider matrices $A \in \mathbb{R}^{d \times d}$ and $H \in \mathbb{R}^{m \times d}$ where $d$ is the state dimension and $m$ is the dimension of the measurements. Then we define the dynamics and the measurement model as follows:

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

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

This defines a dynamics model that assumes a state $\boldsymbol{x}_k$ in a **discrete** sequence of states arising from a linear projection of the previous state $x_{k-1}$ corrupted with additive Gaussian noise under a **process noise** covariance matrix $Q$.
Similarly, the measurements $\boldsymbol{y}_k$ are assumed to be linear projections of the latent state under additive Gaussian noise according to a **measurement noise** covariance $R$.
In the following example we consider projections and covariances that are constant over the state and measurement trajectories (linear time invariant, or **LTI**). Note that this can be generalized to a linear time-varying state-space model, as well. Then $A$ is a function $A: \mathbb{T} \rightarrow \mathbb{R}^{d \times d}$ and $H$ is a function $H: \mathbb{T} \rightarrow \mathbb{R}^{m \times d}$ where $\mathbb{T}$ is the "time dimension".

In other words, here, every relationship is linear and every distribution is a Gaussian distribution.
Under these simplifying assumptions it is possible to obtain a filtering posterior distribution over the state trajectory $(\boldsymbol{x}_k)_{k=1}^{K}$ by using a
**Kalman Filter**. The example is taken from Example 3.6 in [2].

### Define State-Space Model

#### I. Discrete Dynamics Model: Linear, Time-Invariant, Gaussian Transitions

```python
state_dim = 4
observation_dim = 2
```

```python
delta_t = 0.2
# Define linear transition operator
dynamics_transition_matrix = np.eye(state_dim) + delta_t * np.diag(np.ones(2), 2)
# Define process noise (covariance) matrix
process_noise_matrix = (
    np.diag(np.array([delta_t ** 3 / 3, delta_t ** 3 / 3, delta_t, delta_t]))
    + np.diag(np.array([delta_t ** 2 / 2, delta_t ** 2 / 2]), 2)
    + np.diag(np.array([delta_t ** 2 / 2, delta_t ** 2 / 2]), -2)
)
```

To create a discrete, LTI Gaussian dynamics model, `probnum` provides the `DiscreteLTIGaussian` class that takes
- `state_trans_mat` : the linear transition matrix (above: $A$)
- `shift_vec` : a force vector for _affine_ transformations of the state (here: zero)
- `proc_noise_cov_mat` : the covariance matrix for the Gaussian process noise

```python
# Create discrete, Linear Time-Invariant Gaussian dynamics model
dynamics_model = statespace.DiscreteLTIGaussian(
    state_trans_mat=dynamics_transition_matrix,
    shift_vec=np.zeros(state_dim),
    proc_noise_cov_mat=process_noise_matrix,
)
```

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

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

```python
measurement_model = statespace.DiscreteLTIGaussian(
    state_trans_mat=measurement_matrix,
    shift_vec=np.zeros(observation_dim),
    proc_noise_cov_mat=measurement_noise_matrix,
)
```

#### III. Initial State Random Variable

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

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

### Generate Data for the State-Space Model

`statespace.generate_samples()` 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=delta_t)
```

```python
latent_states, observations = statespace.generate_samples(
    rng=rng,
    dynmod=dynamics_model,
    measmod=measurement_model,
    initrv=initial_state_rv,
    times=time_grid,
)
```

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

### Kalman Filtering

#### I. Kalman Filter

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

#### II. 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
grid = state_posterior.locations
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=2, 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])
ax_10 = state_fig.add_subplot(state_fig_gs[1, 0])
ax_11 = state_fig.add_subplot(state_fig_gs[1, 1])

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

ax_00.plot(grid, mu_x_1, label="posterior mean")
ax_01.plot(grid, mu_x_2)
ax_10.plot(grid, mu_x_3)
ax_11.plot(grid, mu_x_4)

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

ax_00.fill_between(
    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(grid, mu_x_2 - 1.96 * std_x_2, mu_x_2 + 1.96 * std_x_2, alpha=0.2)
ax_10.fill_between(grid, mu_x_3 - 1.96 * std_x_3, mu_x_3 + 1.96 * std_x_3, alpha=0.2)
ax_11.fill_between(grid, mu_x_4 - 1.96 * std_x_4, mu_x_4 + 1.96 * std_x_4, alpha=0.2)

# Plot groundtruth
obs_x_1, obs_x_2 = [observations[:, i] for i in range(observation_dim)]

ax_00.scatter(time_grid, obs_x_1, marker=".", label="measurements")
ax_01.scatter(time_grid, obs_x_2, marker=".")

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

ax_00.set_title(r"$x_1$")
ax_01.set_title(r"$x_2$")
ax_10.set_title(r"$x_3$")
ax_11.set_title(r"$x_4$")
handles, labels = ax_00.get_legend_handles_labels()
state_fig.legend(handles, labels, loc="center left", bbox_to_anchor=(1, 0.5))

state_fig.tight_layout()
```

## 2. **Linear Continuous-Discrete** State-Space Model: Ornstein-Uhlenbeck Process

---

Now, consider we have a look at **continuous** dynamics. We assume that there is a continuous process that defines the dynamics of our latent space from which we collect discrete linear-Gaussian measurements (as above). Only the dynamics model becomes continuous. In particular, we formulate the dynamics as a stochastic process in terms of a linear time-invariant stochastic differential equation (LTISDE). We refer to [1] for more details.
Consider matrices $\boldsymbol{F} \in \mathbb{R}^{d \times d}$, $\boldsymbol{L} \in \mathbb{R}^{s \times d}$ and $H \in \mathbb{R}^{m \times d}$ where $d$ is the state dimension and $m$ is the dimension of the measurements.
We define the following **continuous-discrete** state-space model:

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

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

where $\boldsymbol{\omega} \in \mathbb{R}^s$ denotes a vector of driving forces (often Brownian Motion).

Note that this can be generalized to a linear time-varying state-space model, as well. Then $\boldsymbol{F}$ is a function $\mathbb{T} \rightarrow \mathbb{R}^{d \times d}$,
$\boldsymbol{L}$ is a function $\mathbb{T} \rightarrow \mathbb{R}^{s \times d}$, and $H$ is a function $\mathbb{T} \rightarrow \mathbb{R}^{m \times d}$ where $\mathbb{T}$ is the "time dimension". In the following example, however, we consider a LTI SDE, namely, the Ornstein-Uhlenbeck Process from which we observe discrete linear Gaussian measurements.

### Define State-Space Model

#### I. Continuous Dynamics Model: Linear, Time-Invariant Stochastic Differential Equation (LTISDE)

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

```python
delta_t = 0.2
# Define Linear, time-invariant Stochastic Differential Equation that models
# the (scalar) Ornstein-Uhlenbeck Process
drift_constant = 0.21
dispersion_constant = np.sqrt(0.5)
drift = -drift_constant * np.eye(state_dim)
force = np.zeros(state_dim)
dispersion = dispersion_constant * np.eye(state_dim)
```

The _continuous_ counterpart to the discrete LTI Gaussian model from above is provided via the `LTISDE` class.
It is initialized by the state space components
- `driftmat` : the drift matrix $\boldsymbol{F}$
- `forcevec` : a force vector that is added to the state (note that this is **not** $\boldsymbol{\omega}$.) Here: zero.
- `dispmat` : the dispersion matrix $\boldsymbol{L}$

```python
# Create dynamics model
dynamics_model = statespace.LTISDE(
    driftmat=drift,
    forcevec=force,
    dispmat=dispersion,
)
```

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

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

As above, the measurement model is discrete, LTI Gaussian. Only the dymanics are continuous (i.e. continuous-discrete).

```python
measurement_model = statespace.DiscreteLTIGaussian(
    state_trans_mat=measurement_matrix,
    shift_vec=np.zeros(observation_dim),
    proc_noise_cov_mat=measurement_noise_matrix,
)
```

#### III. Initial State Random Variable

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

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

### Generate Data for the State-Space Model

`statespace.generate_samples()` 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=delta_t)
```

```python
latent_states, observations = statespace.generate_samples(
    rng=rng,
    dynmod=dynamics_model,
    measmod=measurement_model,
    initrv=initial_state_rv,
    times=time_grid,
)
```

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

### Kalman Filtering

In fact, since we still consider a **linear** model, we can apply Kalman Filtering in this case again.
According to Section 10 in [1], the moments of the filtering posterior in the continuous-discrete case are solutions to linear differential equations, which `probnum` solves for us when invoking the `<Kalman_object>.filtsmooth(...)` method.

#### I. Kalman Filter

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

#### II. 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 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, )
posterior_state_covs = posterior_state_rvs.cov  # Shape: (num_time_points, )

samples = state_posterior.sample(rng=rng, size=3, t=grid)
```

### Visualize Results

```python
state_fig = plt.figure()

ax = state_fig.add_subplot()

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

# Plot samples
for smp in samples:
    ax.plot(
        grid,
        smp[:, 0],
        color="gray",
        alpha=0.75,
        linewidth=1,
        linestyle="dashed",
        label="sample",
    )

# 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",
)
ax.scatter(time_grid, observations, marker=".", label="measurements")
# Add labels etc.
ax.set_xlabel("t")
ax.set_title(r"$x$")

# These two lines just remove duplicate labels (caused by the samples) from the legend
handles, labels = ax.get_legend_handles_labels()
by_label = dict(zip(labels, handles))

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

state_fig.tight_layout()
```

```python

```


```python
import numpy as np
import math
import matplotlib.pyplot as plt
import scipy.integrate as inte
```

```python
def trajSin(t):
    return np.array([np.cos(np.pi*t),np.sin(np.pi*t)])

def trajSquare(t):
    spd = .2
    if t<2.5:
        x = np.array([.5,0])+spd*np.array([-1,1])*t
    elif t<5:
        x = np.array([0,.5])+spd*np.array([-1,-1])*(t-2.5)
    elif t<7.5:
        x = np.array([-.5,0])+spd*np.array([1,-1])*(t-5)
    else:
        x = np.array([0,-.5])+spd*np.array([1,1])*(t-7.5)

    return x
```

### State-space

\begin{align}
x &= \left[\begin{array}{c}x_1\\x_2\\x_3\\x_4\end{array}\right] = \left[\begin{array}{c}\theta_1\\ \theta_2\\ \dot{\theta_1}\\ \dot{\theta_2}\end{array}\right] \\
\dot{x} &= \left[\begin{array}{c}\dot{\theta_1}\\ \dot{\theta_2}\\\ddot{\theta_1}\\ \ddot{\theta_2}\end{array}\right] =\\
&\left[\begin{array}{c}\dot{\theta_1}\\ \dot{\theta_2}\\ \frac{-u_1+u_2-2glmc_{\theta_1}+u_2c_{\theta_2}+glmc_{\theta_2}c_{\theta_1+\theta_2}+2\dot{\theta_1}\dot{\theta_2}l^2ms_{\theta_2}+(\dot{\theta_2}l)^2ms_{\theta_2}+(\dot{\theta_1} l)^2ms_{\theta_2}(1+c_{\theta_2})}{ml^2(2-c_{\theta_2})} \\ \frac{-u_1+3u_2-2glmc_{\theta_1}-u_1c_{\theta_2}+2u_2c_{\theta_2}-2glmc_{\theta_2}c_{\theta_1}+2glmc_{\theta_1+\theta_2}+glmc_{\theta_2}c_{\theta_1+\theta_2}+(\dot{\theta_2}l)^2m(1+c_{\theta_2})s_{\theta_2}+\dot{\theta_1}\dot{\theta_2}l^2m\csc^2(\frac{\theta_2}{2})s^3_{\theta_2}+(\dot{\theta_1}l)^2m(3s_{\theta_2}+s_{2\theta_2})}{ml^2(2-\cos(\theta_2))} \end{array}\right]
\end{align}

### PD Control

\begin{align}
p &= l\left[\begin{array}{c}c_{\theta_1}+c_{\theta_1+\theta_2}\\ s_{\theta_1}+s_{\theta_1+\theta_2}\end{array}\right] \\
e &= p-p_d\\
J &= l\left[\begin{array}{cc}-s_{\theta_1}-s_{\theta_1+\theta_2}&-s_{\theta_1+\theta_2}\\c_{\theta_1}+c_{\theta_1+\theta_2}&c_{\theta_1+\theta_2}\end{array}\right] \\
u &= -k_pJ^\top e-k_d\left[\begin{array}{c}\dot{\theta_1}\\ \dot{\theta_2}\end{array}\right]
\end{align}

```python
def controller(params,t,X):
    l,kp,kd = .5,30000,100
    p = l*np.array([[np.cos(X[0])+np.cos(X[0]+X[1])],
                    [np.sin(X[0])+np.sin(X[0]+X[1])]])
    e = p.reshape(1,-1)-params['traj'](t)
    J = l*np.array([[-np.sin(X[0])-np.sin(X[0]+X[1]),-np.sin(X[0]+X[1])],
                    [np.cos(X[0])+np.cos(X[0]+X[1]),np.cos(X[0]+X[1])]])
    u = -1*kp*np.matmul(J.T,e.reshape(-1,1)).T-kd*X[2:]
    return u
```

```python
def robotSIm(traj):
    params = {}; params['m'] = 1
    params['g'] = 9.81; params['l'] = .5
    params['traj'] = traj
    t0,t1 = 0,10; t = np.linspace(0,10,250)
    y0 = np.array([0,1,0,0]); y = np.zeros((len(t),len(y0)))
    y[0,:] = y0
    r = inte.ode(lambda t,y:dynamics(params,t,y)).set_integrator("dopri5")
    r.set_initial_value(y0, t0)   # initial values
    for i in range(1, t.size):
        y[i, :] = r.integrate(t[i]) # get one more value, add it to the array
        if not r.successful():
           raise RuntimeError("Could not integrate")

    return t,y
```

```python
def dynamics(params,t,X):
    Xd = np.zeros(X.shape)
    th1,th2,dth1,dth2 = X[0],X[1],X[2],X[3]

    u = controller(params,t,X)
    u1,u2 = -u[0,0],u[0,1]
    m,l,g = params['m'],params['l'],params['g']

    Xd[:2] = X[2:]
    a=-u1+u2-2*g*l*m*np.cos(th1)+u2*np.cos(th2)+g*l*m*np.cos(th2)*np.cos(th1+th2)
    b=2*dth1*dth2*l**2*m*np.sin(th2)+(dth2*l)**2*m*np.sin(th2)+(dth1*l)**2*m*(1+np.cos(th2))*np.sin(th2)
    c=(2-np.cos(th2))*m*l**2
    d=-u1+3*u2-2*g*l*m*np.cos(th1)-u1*np.cos(th2)+u2*np.cos(th2)-2*g*l*m*np.cos(th1)*np.cos(th2)
    e=2*g*l*m*np.cos(th1+th2)+g*l*m*np.cos(th2)*np.cos(th1+th2)+(dth2*l)**2*m*(1+np.cos(th2))*np.sin(th2)
    f=dth1*dth2*l**2*m/np.sin(th2/2)**2*np.sin(th2)**3+(dth1*l)**2*m*(3*np.sin(th2)+np.sin(2*th2))
    tmp1 = (a+b)/c
    tmp2 = (d+e+f)/c
    Xd[2:] = np.array([tmp1,tmp2])

    return Xd
```

```python
t,y = robotSIm(lambda t:trajSin(t))
th1 = y[:,0]
th2 = y[:,1]
p1 = .5*np.array([np.cos(th1),np.sin(th1)])
p2 = p1+.5*np.array([np.cos(th1+th2),np.sin(th1+th2)])

fig = plt.figure(figsize=(15,15))

ax = plt.subplot(2,2,1)
plt.plot(t,y[:,0])
plt.plot(t,y[:,1])
ax = plt.subplot(2,2,2)
plt.plot(t,y[:,2])
plt.plot(t,y[:,3])

i = -1
ax = plt.subplot(2,2,3)
plt.ylim([-1,1]); plt.xlim([-1,1])
ax.plot(np.array([0,p1[0,i],p2[0,i]]),
       np.array([0,p1[1,i],p2[1,i]]),'k-')
ax.plot(p2[0,:i],p2[1,:i])
trajNow = trajSquare(t[i])
ax.plot(trajNow[0],trajNow[1],'r*')
```

```python
t,y = robotSIm(lambda t:trajSquare(t))
th1 = y[:,0]
th2 = y[:,1]
p1 = .5*np.array([np.cos(th1),np.sin(th1)])
p2 = p1+.5*np.array([np.cos(th1+th2),np.sin(th1+th2)])

fig = plt.figure(figsize=(15,15))

ax = plt.subplot(2,2,1)
plt.plot(t,y[:,0])
plt.plot(t,y[:,1])
ax = plt.subplot(2,2,2)
plt.plot(t,y[:,2])
plt.plot(t,y[:,3])

i = -1
ax = plt.subplot(2,2,3)
plt.ylim([-1,1]); plt.xlim([-1,1])
ax.plot(np.array([0,p1[0,i],p2[0,i]]),
       np.array([0,p1[1,i],p2[1,i]]),'k-')
ax.plot(p2[0,:i],p2[1,:i])
trajNow = trajSquare(t[i])
ax.plot(trajNow[0],trajNow[1],'r*')
```

```python

```

```python

```