# Hindmarsh-Roseモデル

\begin{equation}
\left\{
    \begin{aligned}
        \dot{V} &= n - a V^3 + b V^2 - h + I,\\
        \dot{n} &= c - d V^2 -n.
    \end{aligned}
\right.
\end{equation}
リミットサイクルの追跡を有限差分法で行う

```python
import numpy as np
import pathfollowing as pf
from scipy.integrate import ode, solve_ivp
from scipy.linalg import solve, eigvals
import matplotlib.pyplot as plt
import seaborn as sns
%matplotlib inline
sns.set('poster', 'whitegrid', 'dark', rc={"lines.linewidth": 2, 'grid.linestyle': '-'})
```

```python
def f(x, a):
    return np.array([x[1] - x[0]**3 + 3 * x[0]**2 - a[0] + 2, 1 - 5*x[0]**2 - x[1]])

def fx(x, a):
    return np.array([[-3*x[0]**2+6*x[0], 1.0],[-10*x[0], -1.0]])

def fa(x, a):
    return np.array([-1.0, 0.0])
```

```python
def HR(t, x, a,b,c,d,I,h):
    return np.array([x[1] - a*x[0]**3 + b * x[0]**2 - h + I, c - d*x[0]**2 - x[1]])
```

## 平衡点の追跡

```python
h = 1.0
x0 = np.array([1.0, -4.52182801])
r = ode(HR).set_integrator('dop853')
r.set_initial_value(x0, 0.0).set_f_params(1.0, 3.0, 1.0, 5.0, 2.0, 5.0)
sol = []
while r.successful and r.t < 100:
    sol.append(r.integrate(r.t+h))
print(sol[-1])
```

    [ -2.3593041  -26.83157883]

```python
x=np.array([-2.35930409, -26.83157885])
a=np.array([5.0])
bdeq,bpeq,lpeq,hbeq=pf.pathfollow(x, a, f, fx, fa,nmax=175, h=0.1, epsr=1.0e-10, epsb=1.0e-10, problem='equilibria', quiet=True)
bdeq2,bpeq2,lpeq2,hbeq2=pf.pathfollow(x, a, f, fx, fa,nmax=550, h=-0.1, epsr=1.0e-10, epsb=1.0e-10, amin=-10.0,problem='equilibria',quiet=True)
```

    # parameter arrived at boundary

```python
print(hbeq, hbeq2)
print(bdeq2[-1]['a'])
for i in hbeq2:
    print(bdeq2[i])
```

    [] [290, 498]
    [-10.01899244]
    {'TY': 'H', 'x': array([0.18350342, 0.83163248]), 'a': array([2.92647379]), 'v': array([-0.40001257,  0.83188401,  0.3846543 ]), 'udim': 0, 'vr': array([-0.52073336,  0.78171337]), 'vi': array([-0.28559796, -0.19024926]), 'omg': 0.9138020522644607}
    {'TY': 'H', 'x': array([  1.81649658, -15.49829914]), 'a': array([-9.59314045]), 'v': array([-0.03995925,  0.72616588,  0.68635733]), 'udim': 2, 'vr': array([-0.05652179,  0.97347641]), 'vi': array([-0.22132291, -0.01285041]), 'omg': 4.143062370913243}

```python
bdeq_r = np.array([bdeq[m]['a'][0] for m in range(len(bdeq))])
bdeq_V = np.array([bdeq[m]['x'][0] for m in range(len(bdeq))])
bdeq2_r = np.array([bdeq2[m]['a'][0] for m in range(len(bdeq2))])
bdeq2_V = np.array([bdeq2[m]['x'][0] for m in range(len(bdeq2))])
```

```python
fig = plt.figure(figsize=(8, 5))
ax = fig.add_subplot(111)
# ax.set_xlim(0,3)
# ax.set_ylim(-3, 3)
ax.set_xlabel("$h$")
ax.set_ylabel("$V$")
ax.plot(bdeq_r, bdeq_V, '-k')
ax.plot(bdeq2_r, bdeq2_V, '-k')
# plt.savefig("bd_lrz.pdf", bbox_inches='tight')
```

## 有限差分による極限周期軌道追跡

```python
Npts = 128
def func(x, a):
    T = x[-1]
    h = 1.0 / Npts
    n = (len(x) - 1) // Npts
    ss = np.sin([2*np.pi*j*h for j in range(Npts)])
    y = np.zeros(len(x))
    # tau_0 = 0
    y[:n] = (x[n:2*n] - x[(Npts-1)*n:Npts*n])*0.5/h - T * f(x[:n], a)
    # tau_1 ... tau_{N-2}
    for j in range(1, Npts-1):
        y[j*n:(j+1)*n] = (x[(j+1)*n:(j+2)*n] - x[(j-1)*n:j*n])*0.5/h - T * f(x[j*n:(j+1)*n], a)
    # tau_{N-1} = (N-1) * h
    y[(Npts-1)*n:Npts*n] = (x[:n] - x[(Npts-2)*n:(Npts-1)*n])*0.5/h - T * f(x[(Npts-1)*n:Npts*n], a)
    # Phase condition
    y[-1] = x[:-1:n] @ ss
    return y

def dfdx(x, a):
    T = x[-1]
    n = (len(x)-1) // Npts
    h = 1.0 / Npts
    A = np.zeros((len(x), len(x)))
    I = np.identity(n)
    # j = 0
    A[:n, n:2*n] += I * 0.5 / h
    A[:n, (Npts-1)*n:Npts*n] -= I * 0.5 / h
    A[:n,:n] += -T * fx(x[:n], a)
    A[:n,-1] += -f(x[:n], a)
    # j = 1 ... N-2
    for j in range(1, Npts-1):
        A[j*n:(j+1)*n, (j+1)*n:(j+2)*n] += I * 0.5 / h
        A[j*n:(j+1)*n, (j-1)*n:j*n] -= I * 0.5 / h
        A[j*n:(j+1)*n, j*n:(j+1)*n] += -T * fx(x[j*n:(j+1)*n], a)
        A[j*n:(j+1)*n, -1] = -f(x[j*n:(j+1)*n], a)
    # j = N - 1
    A[(Npts-1)*n:Npts*n, :n] += I * 0.5 / h
    A[(Npts-1)*n:Npts*n, (Npts-2)*n:(Npts-1)*n] -= I * 0.5 / h
    A[(Npts-1)*n:Npts*n, (Npts-1)*n:Npts*n] += -T*fx(x[(Npts-1)*n:Npts*n], a)
    A[(Npts-1)*n:Npts*n, -1] += -f(x[(Npts-1)*n:Npts*n], a)
    # Phase condition
    A[-1, :-1:n] = np.sin([2*np.pi*j*h for j in range(Npts)])
    return A

def dfda(x, a):
    T = x[-1]
    n = (len(x) - 1) // Npts
    h = 1.0 / Npts
    b = np.zeros(len(x))
    for j in range(Npts):
        b[j*n:(j+1)*n] = -T*fa(x[j*n:(j+1)*n], a)
    return b
```

```python
vr = bdeq2[hbeq2[1]]['vr']
vi = bdeq2[hbeq2[1]]['vi']
vr /= np.sqrt(vr@vr)
vi /= np.sqrt(vi@vi)
rr = np.hypot(vr, vi)
tt = np.array([np.angle(vr[i]+1.j*vi[i]) for i in range(len(vr))])
ev = eigvals(fx(bdeq2[hbeq2[1]-1]['x'], bdeq2[hbeq2[1]-1]['a']))
```

```python
print(ev)
```

    [0.00790124+4.13724307j 0.00790124-4.13724307j]

```python
print(np.linalg.norm(vr), np.linalg.norm(vi))
```

    0.9999999999999999 1.0

```python
x = np.zeros(2*Npts+1)
da = np.abs(bdeq2[hbeq2[1]]['a'][0] - bdeq2[hbeq2[1]-1]['a'][0])
for j in range(Npts):
    x[2*j:2*(j+1)] = bdeq2[hbeq2[1]-1]['x'] + 2 * np.sqrt(da) * rr * np.cos(2*np.pi*j/Npts + tt - tt[0])
x[-1] = 2*np.pi/np.imag(ev[0])
```

```python
fig = plt.figure()
ax = fig.add_subplot(111)
ax.plot(x[:-1:2],'.')
ax.plot(x[1:-1:2],'.')
```

Newton法でリミットサイクルを捉える

```python
y = np.copy(x)
a = bdeq2[hbeq2[1]-1]['a']
for m in range(12):
    b = func(y, a)
    A = dfdx(y, a)
    y -= solve(A, b)
    print(np.linalg.norm(b))
print(y)
```

    83.39627205855717
    1.7751768822618856
    1.3369929065021189
    0.3483115208107321
    0.09531631392339272
    0.008491365270787178
    7.63802931758492e-05
    5.89808082141663e-09
    7.213397483126336e-13
    7.956177917048393e-13
    8.366934913085939e-13
    6.961110000650344e-13
    [  2.06834557 -15.64772672   2.06652367 -15.70390827   2.0640474
     -15.75889267   2.0609346  -15.81253637   2.05720431 -15.8647027
       2.05287658 -15.91526216   2.04797236 -15.96409267   2.04251331
     -16.01107976   2.0365217  -16.0561167    2.03002031 -16.09910459
       2.02303228 -16.13995242   2.01558103 -16.17857702   2.00769019
     -16.21490311   1.99938351 -16.24886314   1.9906848  -16.28039729
       1.98161786 -16.30945326   1.97220643 -16.33598618   1.96247419
     -16.35995843   1.95244465 -16.38133949   1.9421412  -16.40010568
       1.93158704 -16.41624003   1.92080518 -16.42973202   1.90981843
     -16.44057741   1.89864938 -16.44877795   1.88732042 -16.45434122
       1.8758537  -16.45728037   1.86427117 -16.45761389   1.85259457
     -16.45536541   1.84084544 -16.45056346   1.82904509 -16.44324123
       1.81721467 -16.43343642   1.80537513 -16.42119093   1.79354723
     -16.40655074   1.78175159 -16.38956567   1.77000864 -16.37028916
       1.75833867 -16.34877814   1.74676183 -16.32509278   1.7352981
     -16.29929637   1.72396733 -16.27145509   1.71278926 -16.2416379
       1.70178347 -16.20991634   1.69096939 -16.17636438   1.68036636
     -16.14105831   1.66999354 -16.10407655   1.65986996 -16.06549957
       1.65001449 -16.0254097    1.64044584 -15.98389107   1.63118256
     -15.94102946   1.62224299 -15.89691221   1.61364527 -15.85162808
       1.60540732 -15.80526721   1.59754683 -15.75792101   1.59008119
     -15.70968204   1.58302752 -15.66064398   1.57640262 -15.61090152
       1.57022292 -15.56055033   1.56450447 -15.50968694   1.5592629
     -15.45840873   1.55451337 -15.40681385   1.55027055 -15.35500118
       1.54654854 -15.30307026   1.54336085 -15.25112125   1.54072034
     -15.19925491   1.53863918 -15.14757251   1.53712875 -15.0961758
       1.53619965 -15.045167     1.53586156 -14.99464867   1.53612324
     -14.94472374   1.5369924  -14.89549541   1.5384757  -14.84706706
       1.54057859 -14.79954225   1.54330531 -14.75302458   1.54665876
     -14.70761759   1.55064044 -14.6634247    1.55525035 -14.62054906
       1.56048693 -14.57909343   1.56634694 -14.53915997   1.57282541
     -14.50085013   1.57991553 -14.46426442   1.58760856 -14.42950219
       1.59589377 -14.39666135   1.60475838 -14.36583815   1.61418744
     -14.33712682   1.6241638  -14.31061925   1.63466805 -14.28640465
       1.6456785  -14.2645691    1.65717108 -14.24519518   1.66911942
     -14.22836149   1.68149477 -14.21414214   1.69426606 -14.20260633
       1.70739994 -14.19381771   1.72086082 -14.18783396   1.73461098
     -14.18470614   1.74861066 -14.1844782    1.76281822 -14.18718639
       1.77719028 -14.19285874   1.79168191 -14.20151451   1.80624686
     -14.21316373   1.8208378  -14.22780668   1.83540653 -14.24543351
       1.84990433 -14.2660239    1.86428221 -14.28954668   1.87849123
     -14.31595968   1.89248281 -14.34520954   1.90620907 -14.37723165
       1.91962314 -14.41195019   1.93267945 -14.44927824   1.9453341
     -14.48911798   1.95754507 -14.53136106   1.96927255 -14.57588894
       1.98047914 -14.62257346   1.99113011 -14.67127738   2.00119356
     -14.72185508   2.01064058 -14.7741533    2.01944535 -14.82801194
       2.02758527 -14.88326493   2.03504097 -14.93974112   2.04179633
     -14.99726522   2.04783847 -15.05565875   2.05315772 -15.11474098
       2.0577475  -15.17432991   2.06160424 -15.23424317   2.06472728
     -15.29429897   2.06711869 -15.35431692   2.06878313 -15.41411891
       2.06972769 -15.47352989   2.06996171 -15.53237856   2.06949657
     -15.5904981    1.53045639]

```python
fig = plt.figure()
ax = fig.add_subplot(211)
ax2 = fig.add_subplot(212)
ax.plot(y[:-1:2],'o-')
ax2.plot(y[1:-1:2],'o-')
```

Hopf分岐で分岐したリミットサイクルの追跡

```python
x = np.copy(y)
bdc,bpc,lpc=pf.pathfollow(x, a, func, dfdx, dfda, nmax=400, h=0.5, epsr=1.0e-10, epsb=1.0e-10, amax=2.08, problem=None, quiet=True)
```

    # parameter arrived at boundary

```python
bdc[-1]['a']
```

    array([2.08374174])

```python
bdc_r = np.array([pt['a'][0] for pt in bdc])
bdc_T = np.array([pt['x'][-1] for pt in bdc])
bdc_Vmax = np.array([np.max(pt['x'][:-1:2]) for pt in bdc])
bdc_Vmin = np.array([np.min(pt['x'][:-1:2]) for pt in bdc])
```

```python
fig = plt.figure(figsize=(8, 5))
ax = fig.add_subplot(111)
ax.set_xlim(-10,2.5)
ax.set_xlabel("$h$")
ax.set_ylabel("$T$")
ax.plot(bdc_r, bdc_T, '-k')
# plt.savefig("bd_HR_T.pdf", bbox_inches='tight')
```

```python
fig = plt.figure(figsize=(8, 5))
ax = fig.add_subplot(111)
ax.set_xlim(-10,5)
ax.set_ylim(-3,3)
ax.set_xlabel("$h$")
ax.set_ylabel("$V$")
ax.plot(bdeq_r, bdeq_V, '--k')
ax.plot(bdeq2_r, bdeq2_V, '--k')
ax.plot(bdc_r, bdc_Vmax, '-k')
ax.plot(bdc_r, bdc_Vmin, '-k')
# plt.savefig("bd_HR_V.pdf", bbox_inches='tight')
```

```python
fig = plt.figure(figsize=(8,5))
ax = fig.add_subplot(111)
ax.plot(bdc[-40]['x'][:-1:2],'-')
```

# 多段射撃法

```python
Npts = 8
def func(x, a):
    T = x[-1]
    def f2(t, y):
        return T * f(y, a)
    r = ode(f2).set_integrator('dop853', atol=1.0e-14, rtol=1.0e-14)
    n = (len(x) - 1) // Npts
    h = 1.0 / Npts
    x1 = np.zeros(len(x))

    for k in range(Npts-1):
        y0 = np.copy(x[k*n:(k+1)*n])
        r.set_initial_value(y0, 0.0)
        y1 = r.integrate(r.t+h)
        x1[k*n:(k+1)*n] = y1 - x[(k+1)*n:(k+2)*n]

    y0 = np.copy(x[-(n+1):-1])
    r.set_initial_value(y0, 0.0)
    y1 = r.integrate(r.t+h)
    x1[-(n+1):-1] = y1 - x[:n]

    x1[-1] = x[:-1:n] @ np.array([np.sin(2*k*h*np.pi) for k in range(Npts)])
    return x1

def dfdx(x, a):
    def df2(t, y, n):
        z = np.zeros((n+1)*(n+2))
        z[:n] = y[n] * f(y[:n], a)
        z[n] = 0.0
        J = np.zeros((n+1, n+1))
        J[:n, :n] = y[n] * fx(y[:n], a)
        J[:n, n] = f(y[:n], a)
        for m in range(n+1):
            z[(n+1)*(m+1):(n+1)*(m+2)] = J @ y[(n+1)*(m+1):(n+1)*(m+2)]
        return z
    r = ode(df2).set_integrator('dop853', atol=1.0e-14, rtol=1.0e-14)
    n = (len(x)-1) // Npts
    h = 1.0 / Npts
    A = np.zeros((len(x), len(x)))
    I = np.identity(n+1)

    for k in range(Npts-1):
        y0 = np.zeros((n+1)*(n+2))
        y0[:n] = x[k*n:(k+1)*n]
        y0[n] = x[-1]
        for m in range(n+1):
            y0[(n+1)*(m+1):(n+1)*(m+2)] = I[:,m]
        r.set_initial_value(y0, 0.0).set_f_params(n)
        y1 = r.integrate(r.t+h)
        for m in range(n):
            A[k*n:(k+1)*n,k*n+m] = y1[(n+1)*(m+1):(n+1)*(m+1)+n]
        A[k*n:(k+1)*n, (k+1)*n:(k+2)*n] = -np.identity(n)
        A[k*n:(k+1)*n, -1] = y1[-(n+1):-1]

    y0 = np.zeros((n+1)*(n+2))
    y0[:n] = x[-(n+1):-1]
    y0[n] = x[-1]
    for m in range(n+1):
        y0[(n+1)*(m+1):(n+1)*(m+2)] = I[:,m]
    r.set_initial_value(y0, 0.0).set_f_params(n)
    y1 = r.integrate(r.t+h)
    for m in range(n):
        A[-(n+1):-1,-(n+1)+m] = y1[(n+1)*(m+1):(n+1)*(m+1)+n]
    A[-(n+1):-1, :n] = -np.identity(n)
    A[-(n+1):-1, -1] = y1[-(n+1):-1]

    for k in range(Npts):
        A[-1,k*n] = np.sin(2*np.pi*k*h)
    return A

def dfda(x, a):
    T = x[-1]
    def df2(t, y, n):
        z = np.zeros(2*(n+1))
        z[:n] = T * f(y[:n], np.array([y[n]]))
        z[n] = 0.0
        J = np.zeros((n+1, n+1))
        J[:n, :n] = fx(y[:n], np.array([y[n]]))
        J[:n, n] = fa(y[:n], np.array([y[n]]))
        z[n+1:] = T * J @ y[n+1:]
        return z
    n = (len(x) - 1) // Npts
    h = 1.0 / Npts
    r = ode(df2).set_integrator('dop853', atol=1e-14, rtol=1e-14)
    b = np.zeros(len(x))

    for k in range(Npts-1):
        y0 = np.zeros(2*(n+1))
        y0[:n] = np.copy(x[k*n:(k+1)*n])
        y0[n] = a[0]
        y0[-1] = 1.0
        r.set_initial_value(y0, 0.0).set_f_params(n)
        y1 = r.integrate(r.t+h)
        b[k*n:(k+1)*n] = y1[n+1:2*n+1]

    y0[:n] = np.copy(x[-(n+1):-1])
    y0[n] = a[0]
    y0[-1] = 1.0
    r.set_initial_value(y0, 0.0).set_f_params(n)
    y1 = r.integrate(r.t+h)
    b[-(n+1):-1] = y1[n+1:2*n+1]

    return b
```

```python
vr = bdeq2[hbeq2[1]]['vr']
vi = bdeq2[hbeq2[1]]['vi']
vr /= np.sqrt(vr@vr)
vi /= np.sqrt(vi@vi)
rr = np.hypot(vr, vi)
tt = np.array([np.angle(vr[i]+1.j*vi[i]) for i in range(len(vr))])
ev = eigvals(fx(bdeq2[hbeq2[1]-1]['x'], bdeq2[hbeq2[1]-1]['a']))
```

```python
x = np.zeros(2*Npts+1)
da = np.abs(bdeq2[hbeq2[1]]['a'][0] - bdeq2[hbeq2[1]-1]['a'][0])
for j in range(Npts):
    x[2*j:2*(j+1)] = bdeq2[hbeq2[1]-1]['x'] + 2 * np.sqrt(da) * rr * np.cos(2*np.pi*j/Npts + tt - tt[0])
x[-1] = 2*np.pi/np.imag(ev[0])
```

```python
y = np.copy(x)
a = bdeq2[hbeq2[1]-1]['a']
for m in range(12):
    b = func(y, a)
    A = dfdx(y, a)
    y -= solve(A, b)
    print(np.linalg.norm(b))
# x = np.copy(y)
print(y)
```

    2.2586504887251064
    0.07985898821767388
    0.007409036151628275
    0.003336334035679782
    0.00030515054468293457
    2.2452251675668243e-06
    1.0171411829740142e-10
    8.073390384088094e-15
    1.371114198290873e-14
    1.1214901273581282e-14
    6.3893310952859035e-15
    9.5608188662104e-15
    [  2.06901919 -15.64803573   1.97260051 -16.33826827   1.79344961
     -16.40899249   1.62167077 -15.89786852   1.53630407 -15.09488909
       1.59522576 -14.39344856   1.79158023 -14.19777262   2.00168924
     -14.71955072   1.53113366]

```python
fig = plt.figure()
ax = fig.add_subplot(211)
ax2 = fig.add_subplot(212)
ax.plot(y[:-1:2],'o-')
ax2.plot(y[1:-1:2],'o-')
```

```python
x = np.copy(y)
bdc,bpc,lpc=pf.pathfollow(x, a, func, dfdx, dfda, nmax=400, h=0.5, epsr=1.0e-10, epsb=1.0e-10, amax=2.08, problem=None, quiet=True)
```

    # parameter arrived at boundary

```python
bdc[-1]['a']
```

    array([2.080317])

```python
bdc_r = np.array([pt['a'][0] for pt in bdc])
bdc_T = np.array([pt['x'][-1] for pt in bdc])
```

```python
fig = plt.figure(figsize=(8, 5))
ax = fig.add_subplot(111)
ax.set_xlim(-10,2.5)
ax.set_ylim(0,40)
ax.set_xlabel("$h$")
ax.set_ylabel("$T$")
ax.plot(bdc_r, bdc_T, '-k')
# plt.savefig("bd_HR_T.pdf", bbox_inches='tight')
```

```python
N = 256
Vmax, Vmin = [], []
for pt in bdc:
    x0, a0, T = pt['x'][0:2], pt['a'][0], pt['x'][-1]
    dt = T / N
    r = ode(HR).set_integrator('dop853')
    r.set_initial_value(x0, 0.0).set_f_params(1.0, 3.0, 1.0, 5.0, 2.0, a0)
    V0, V1 = x0[0], x0[0]
    while r.successful and r.t < T:
        y0 = r.integrate(r.t+dt)
        V0 = y0[0] if y0[0] > V0 else V0
        V1 = y0[0] if y0[0] < V1 else V1
    Vmax.append(V0)
    Vmin.append(V1)
bdc_Vmax = np.array(Vmax)
bdc_Vmin = np.array(Vmin)
```

```python
fig = plt.figure(figsize=(8, 5))
ax = fig.add_subplot(111)
ax.set_xlim(-10,5)
ax.set_ylim(-3,3)
ax.set_xlabel("$h$")
ax.set_ylabel("$V$")
ax.plot(bdeq_r, bdeq_V, '-k')
ax.plot(bdeq2_r[:lpeq2[0]], bdeq2_V[:lpeq2[0]], '-k')
ax.plot(bdeq2_r[lpeq2[0]:lpeq2[1]], bdeq2_V[lpeq2[0]:lpeq2[1]], '--k')
ax.plot(bdeq2_r[lpeq2[1]:hbeq2[0]], bdeq2_V[lpeq2[1]:hbeq2[0]], '--k')
ax.plot(bdeq2_r[hbeq2[0]:hbeq2[1]], bdeq2_V[hbeq2[0]:hbeq2[1]], '--k')
ax.plot(bdeq2_r[hbeq2[1]:], bdeq2_V[hbeq2[1]:], '-k')
ax.plot(bdeq2_r[hbeq2[0]], bdeq2_V[hbeq2[0]], 'xk', markersize=10,label='HB')
ax.plot(bdeq2_r[hbeq2[1]], bdeq2_V[hbeq2[1]], 'xk', markersize=10)
ax.plot(bdc_r[::5], bdc_Vmax[::5], 'o-k')
ax.plot(bdc_r[::5], bdc_Vmin[::5], 'o-k')
ax.legend()
# plt.savefig("bd_HR_V.pdf", bbox_inches='tight')
```

```python
fig = plt.figure(figsize=(8,5))
ax = fig.add_subplot(111)
ax.plot(bdc[1]['x'][:-1:2], bdc[1]['x'][1:-1:2], '.')
```


 Exogeneous closure for the PKF applied to the Burgers dynamics

    Olivier Pannekoucke
 2020

The aim is to train a closure to predict the uncertainty dynamics for the Burgers dynamics.

The notebook present the situation where uknown physical processes are represented by an exogeneous NN to train.

---
 Table of contents

 1. [Introduction](#introduction)
 1. [The Burgers dynamics](#the-burgers-dynamics)
 1. [PKF for the Burgers dynamics](#pkf-forthe-burgers-dynamics)
 1. [Numerical application](#numerical-application)
   - [Generation of a database](#generation-of-a-database)
   - [Training of the closure](#training-of-the-closure)
   - [Comparison with the theoretically designed closure](#comparison-with-the-theoretically-designed-closure)
 1. [Conclusion](#conclusion)
 ---

```python
import numpy as np
import matplotlib.pyplot as plt
#%matplotlib inline
#%matplotlib notebook
```

## Introduction

The aim is to design a NN which merges known and unknown physics.

```python
import sympy
from sympy import (Function, symbols, init_printing, Derivative,
                   latex, Add, Mul, Pow,
                   Integer, Rational, Float, Symbol, symbol,
                   srepr, Tuple
                  )
init_printing()
```

## The Burgers dynamics

```python
from pdenetgen import NNModelBuilder, Eq
import tensorflow.keras as keras
```

```python
def display_system(system):
    print(50*'*')
    for equation in system:
        display(equation)
        print(50*'*')
```

#### Set of the function and symbols

```python
t, x = symbols('t x')

u = Function('u')(t,x)
closure = sympy.Function('closure')(t,x)
V = Function('{V_{u}}')(t,x)
nu = Function('{\\nu_{u,xx}}')(t,x)
Kappa = symbols('\\kappa')
```

#### Set constants for numerical experiments

```python
# Constant setting following Pannekoucke et al. (2018)
n = 241
kappa = 0.0025
dt = 0.002
```

#### Set of the Burgers equation

```python
burgers_dynamics = [
        Eq(
        Derivative(u,t),
        Kappa*Derivative(u,x,2)-u*Derivative(u,x)
      ),
]
display_system(burgers_dynamics)
```

```python
burgers_NN_builder = NNModelBuilder(burgers_dynamics, "Burgers")
print(burgers_NN_builder.code)
exec(burgers_NN_builder.code)
burgers = Burgers(shape=(n,), kappa=kappa)
```

    from pdenetgen.model import Model
    import numpy as np
    import tensorflow.keras as keras
    from pdenetgen.symbolic.nn_builder import DerivativeFactory, TrainableScalarLayerFactory

    class Burgers(Model):

        # Prognostic functions (sympy functions):
        prognostic_functions = (
                'u',    # Write comments on the function here
            )

        # Spatial coordinates
        coordinates = (
                'x',    # Write comments on the coordinate here
            )

        # Set constants
        constants = (
                'kappa',    # Writes comment on the constant here
            )

        def __init__(self, shape=None, lengths=None, **kwargs):

            super().__init__() # Time scheme is set from Model.__init__()

            #---------------------------------
            # Set index array from coordinates
            #---------------------------------

            # a) Set shape
            shape = len(self.coordinates)*(100,) if shape is None else shape
            if len(shape)!=len(self.coordinates):
                raise ValueError(f"len(shape) {len(shape)} is different from len(coordinates) {len(self.coordinates)}")
            else:
                self.shape = shape

            # b) Set input shape for coordinates
            self.input_shape_x = shape[0]

            # c) Set lengths
            lengths = len(self.coordinates)*(1.0,) if lengths is None else lengths
            if len(lengths)!=len(self.coordinates):
                raise ValueError(f"len(lengths) {len(lengths)} is different from len(coordinates) {len(self.coordinates)}")
            else:
                self.lengths = lengths

            # d) Set indexes
            self._index = {}
            for k,coord in enumerate(self.coordinates):
                self._index[(coord,0)] = np.arange(self.shape[k], dtype=int)

            # Set x/dx
            #-------------
            self.dx = tuple(length/shape for length, shape in zip(self.lengths, self.shape))
            self.x = tuple(self.index(coord,0)*dx for coord, dx in zip(self.coordinates, self.dx))
            self.X = np.meshgrid(*self.x)

            #---------------------------
            # Set constants of the model
            #---------------------------

            # Set a default nan value for constants
            self.kappa = np.nan # @@ set constant value @@

            # Set constant values from external **kwargs (when provided)
            for key in kwargs:
                if key in self.constants:
                    setattr(self, key, kwargs[key])

            # Alert when a constant is np.nan
            for constant in self.constants:
                if getattr(self, constant) is np.nan:
                    print(f"Warning: constant `{constant}` has to be set")

            # Set NN models
            self._trend_model = None
            self._exogenous_model = None

        def index(self, coord, step:int):
            """ Return int array of shift index associated with coordinate `coord` for shift `step` """
            # In this implementation, indexes are memory saved in a dictionary, feed at runtime
            if (coord,step) not in self._index:
                self._index[(coord,step)] = (self._index[(coord,0)]+step)%self.shape[self.coordinates.index(coord)]
            return self._index[(coord,step)]

        def _make_trend_model(self):
            """ Generate the NN used to compute the trend of the dynamics """

            # Alias for constants
            #--------------------
            kappa = self.kappa
            if kappa is np.nan:
                raise ValueError("Constant 'kappa' is not set")

            # Set input layers
            #------------------
            # Set Alias for coordinate input shapes
            input_shape_x = self.input_shape_x

            # Set input shape for prognostic functions
            u = keras.layers.Input(shape =(input_shape_x,1,))

            # Keras code

            # 2) Implementation of derivative as ConvNet
            # Compute derivative
            #-----------------------
            #
            #  Warning: might be modified to fit appropriate boundary conditions.
            #
            kernel_Du_x_o2 = np.asarray([self.dx[self.coordinates.index('x')]**(-2),
            -2/self.dx[self.coordinates.index('x')]**2,
            self.dx[self.coordinates.index('x')]**(-2)]).reshape((3,)+(1,1))
            Du_x_o2 = DerivativeFactory((3,),kernel=kernel_Du_x_o2,name='Du_x_o2')(u)

            kernel_Du_x_o1 = np.asarray([-1/(2*self.dx[self.coordinates.index('x')]),0.0,
            1/(2*self.dx[self.coordinates.index('x')])]).reshape((3,)+(1,1))
            Du_x_o1 = DerivativeFactory((3,),kernel=kernel_Du_x_o1,name='Du_x_o1')(u)

            # 3) Implementation of the trend as NNet

            #
            # Computation of trend_u
            #
            sc_mul_0 = keras.layers.Lambda(lambda x: kappa*x,name='ScalarMulLayer_0')(Du_x_o2)
            mul_0 = keras.layers.multiply([Du_x_o1,u],name='MulLayer_0')
            sc_mul_1 = keras.layers.Lambda(lambda x: -1.0*x,name='ScalarMulLayer_1')(mul_0)
            trend_u = keras.layers.add([sc_mul_0,sc_mul_1],name='AddLayer_0')

            # 4) Set 'input' of model
            inputs = [
                    # Prognostic functions
                    u,

                ]

            # 5) Set 'outputs' of model
            outputs = [
                trend_u,
                ]

            model = keras.models.Model(inputs=inputs, outputs=outputs)
            #model.trainable = False
            self._trend_model = model

        def trend(self, t, state):
            """ Trend of the dynamics """

            if self._trend_model is None:
                self._make_trend_model()

            # Init output state with pointer on data
            #-------------------------------------------

            #   a) Set the output array
            dstate = np.zeros(state.shape)

            #   b) Set pointers on output array `dstate` for the computation of the physical trend (alias only).
            du = dstate[0]

            # Load physical functions from state
            #------------------------------------
            u = state[0]

            # Compute the trend value from model.predict
            #-------------------------------------------
            inputs = [
                # Prognostic functions
                    u,
                    ]

            dstate = self._trend_model.predict( inputs )

            if not isinstance(dstate,list):
                dstate = [dstate]

            return np.array(dstate)

        def _make_dynamical_trend(self):
            """
            Computation of a trend model so to be used in a time scheme (as solving a dynamical system or an ODE)

            Description:
            ------------

            In the present implementation, the inputs of the trend `self._trend_model` is a list of fields, while
            entry of a time-scheme is a single array which contains all fields.

            The aims of `self._dynamical_trend` is to produce a Keras model which:
            1. takes a single array as input
            2. extract the `self._trend_model` input list from the input array
            3. compute the trends from `self._trend_model`
            4. outputs the trends as a single array

            Explaination of the code:
            -------------------------

            Should implement a code as the following, that is valid for the PKF-Burgers

            def _make_dynamical_trend(self):

                if self._trend_model is None:
                    self._make_trend_model()

                # 1. Extract the input of the model

                # 1.1 Set the input as an array
                state = keras.layers.Input(shape=(3,self.input_shape_x,1))

                # 1.2 Extract each components of the state
                u = keras.layers.Lambda(lambda x : x[:,0,:,:])(state)
                V = keras.layers.Lambda(lambda x : x[:,1,:,:])(state)
                nu_u_xx = keras.layers.Lambda(lambda x : x[:,2,:,:])(state)

                # 2. Compute the trend
                trend_u, trend_V, trend_nu = self._trend_model([u,V,nu_u_xx])

                # 3. Outputs the trend as a single array

                # 3.1 Reshape trends
                trend_u = keras.layers.Reshape((1,self.input_shape_x,1))(trend_u)
                trend_V = keras.layers.Reshape((1,self.input_shape_x,1))(trend_V)
                trend_nu = keras.layers.Reshape((1,self.input_shape_x,1))(trend_nu)

                # 3.2 Concatenates all trends
                trends = keras.layers.Concatenate(axis=1)([trend_u,trend_V,trend_nu])

                # 4. Set the dynamical_trend model
                self._dynamical_trend = keras.models.Model(inputs=state,outputs=trends)
            """

            if self._trend_model is None:
                self._make_trend_model()

            for exclude_case in ['constant_functions','exogenous_functions']:
                if hasattr(self,exclude_case):
                    raise NotImplementedError(f'Design of dynamical_model with {exclude_case} is not implemented')

            # Case 1 --  corresponds to the _trend_model if input is a single field
            if not isinstance(self._trend_model.input_shape, list):
                self._dynamical_trend = self._trend_model
                return

            # Case 2 -- Case where multiple list is used

            # 1. Extract the input of the model

            # 1.1 Set the input as an array
            """ from PKF-Burgers code:
            state = keras.layers.Input(shape=(3,self.input_shape_x,1))
            """

            # 1.1.1 Compute the input_shape from _trend_model
            shapes = []
            dimensions  = []
            for shape in self._trend_model.input_shape:
                shape = shape[1:] # Exclude batch_size (assumed to be at first)
                shapes.append(shape)
                dimensions.append(len(shape)-1)

            max_dimension = max(dimensions)
            if max_dimension!=1:
                if 1 in dimensions:
                    raise NotImplementedError('1D fields incompatible with 2D/3D fields')

            # todo: add test to check compatibility of shapes!!!!

            if max_dimension in [1,2]:
                input_shape = (len(shapes),)+shapes[0]
            elif max_dimension==3:

                # a. check the size of 2D fields: this is given by the first 2D field.
                for shape, dimension in zip(shapes, dimensions):
                    if dimension==2:
                        input_shape_2D = shape
                        break

                # b. Compute the numbers of 2D fields: this corresponds to the number of 3D layers and the number of 2D fields.

                for shape, dimension in zip(shapes, dimensions):
                    if dimension==2:
                        nb_outputs += 1
                    else:
                        nb_outputs += shape[0]

                input_shape = (nb_outputs,)+input_shape_2D

            # 1.1.2 Init the state of the dynamical_trend
            state = keras.layers.Input(shape=input_shape)

            # 1.2 Extract each components of the state
            """ From PKF-Burgers code:

            u = keras.layers.Lambda(lambda x : x[:,0,:,:])(state)
            V = keras.layers.Lambda(lambda x : x[:,1,:,:])(state)
            nu_u_xx = keras.layers.Lambda(lambda x : x[:,2,:,:])(state)

            inputs = [u, V, nu_u_xx]
            """

            def get_slice(dimension, k):
                def func(x):
                    if dimension == 1:
                        return x[:,k,:,:]
                    elif dimension == 2:
                        return x[:,k,:,:]
                return func

            def get_slice_3d(start,end):
                def func(x):
                    return x[:,start:end,:,:,:]
                return func

            inputs = []
            if max_dimension in [1,2]:
                for k in range(len(shapes)):
                    inputs.append(keras.layers.Lambda(get_slice(max_dimension,k))(state))
                    #if max_dimension==1:
                    #    inputs.append(keras.layers.Lambda(lambda x : x[:,k,:,:])(state))
                    #
                    #if max_dimension==2:
                    #    inputs.append(keras.layers.Lambda(lambda x : x[:,k,:,:,:])(state))
            else:
                k=0
                for shape, dimension in zip(shapes, dimensions):
                    if dimension==2:
                        #inputs.append(keras.layers.Lambda(lambda x : x[:,k,:,:,:])(state))
                        inputs.append(keras.layers.Lambda(get_slice(dimension,k))(state))
                        k += 1
                    if dimension==3:
                        start = k
                        end = start+shape[0]
                        inputs.append(keras.layers.Lambda(get_slice_3d(start,end))(state))
                        k = end

            # 2. Compute the trend
            """ From PKF-Burgers code
            trend_u, trend_V, trend_nu = self._trend_model([u,V,nu_u_xx])
            """
            trends = self._trend_model(inputs)

            # 3. Outputs the trend as a single array

            # 3.1 Reshape trends
            """ from PKF-Burgers code
            trend_u = keras.layers.Reshape((1,self.input_shape_x,1))(trend_u)
            trend_V = keras.layers.Reshape((1,self.input_shape_x,1))(trend_V)
            trend_nu = keras.layers.Reshape((1,self.input_shape_x,1))(trend_nu)
            """

            reshape_trends = []
            for trend, dimension in zip(trends, dimensions):

                #shape = tuple(dim.value for dim in trend.shape[1:])
                # update from keras -> tensorflow.keras
                shape = tuple(dim for dim in trend.shape[1:])

                if dimension==1 or dimension==2:
                    # for 1D fields like (128,1) transform into (1,128,1)
                    # for 2D fields like (128,128,1) transform into (1,128,128,1)
                    shape = (1,)+shape
                elif dimension==3:
                    # 3D fields can be compated: two fields (36,128,128,1) become the single field (72,128,128,1)
                    pass
                else:
                    raise NotImplementedError

                reshape_trends.append(keras.layers.Reshape(shape)(trend))

            # 3.2 Concatenates all trends
            """ From PKF-Burgers code:
            trends = keras.layers.Concatenate(axis=1)([trend_u,trend_V,trend_nu])
            """
            trends = keras.layers.Concatenate(axis=1)(reshape_trends)

            # 2.5 Compute the model
            self._dynamical_trend = keras.models.Model(inputs=state,outputs=trends)

##### Example of forecast from a given initial condition

```python
import matplotlib.pyplot as plt

def plot_results(data, label=None, labelx=True, title=None, save_file=None, normalisation=None,
                 selected_times=None,style=None, name=None, alpha=1., bolds=[0., 1.]):

    normalisation = 1. if normalisation is None else normalisation

    selected_times = [time for time in data] if selected_times is None else selected_times

    style = 'k' if style is None else style

    for time in selected_times:
        lalpha = alpha if time in bolds else 0.2
        lname = name if time==selected_times[-1] else None
        plt.plot(domain.x[0],data[time]/normalisation, style, alpha = lalpha, label=lname)

    if labelx:
        plt.xlabel('$x/D$', fontsize=15)
    if label:
        plt.ylabel(label, fontsize=15)
    if title:
        plt.title(title)
    if save_file:
        plt.savefig(save_file)
```

```python
domain = burgers
# Set initial condition for 'u'
U0=0.25*( 1+np.cos(2*np.pi/ domain.lengths[0]  *(domain.x[0]-0.25)) )
Umax = U0.max()
```

```python
burgers.set_dt(dt)
end_time_forecast = 1.
times = burgers.window(end_time_forecast)
saved_times = times[::50]
print('saved_times :' ,saved_times)
```

    saved_times : [0.  0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1. ]

```python
forecast = burgers.forecast(times, np.array([U0.reshape((1,)+U0.shape+(1,)) ]))
```

```python
for time in times:
    plt.plot(domain.x[0], forecast[time][0,0,:,0])
```

## PKF for the Burgers dynamics

#### Set of the PKF equations for the Burgers equation

```python
# From Pannekoucke et al. (2018)

pkf_dynamics = [
    # Trend of the expectation of 'u'
    Eq(
        Derivative(u,t),
        Kappa*Derivative(u,x,2)-u*Derivative(u,x)-Derivative(V,x)/Integer(2)
      ),
    # Trend of the variance
    Eq(
        Derivative(V,t),
        -Kappa*V/nu + Kappa*Derivative(V,x,2)-Kappa*Derivative(V,x)**Integer(2)/(Integer(2)*V)
        -u*Derivative(V,x)-Integer(2)*V*Derivative(u,x)
      ),
    # Trend of the diffusion
    Eq(
        Derivative(nu,t),
        Integer(4)*Kappa*nu**Integer(2)*closure
        -Integer(3)*Kappa*Derivative(nu,x,2)
        -Kappa
        +Integer(6)*Kappa*Derivative(nu,x)**Integer(2)/nu
        -Integer(2)*Kappa*nu*Derivative(V,x,2)/V
        +Kappa*Derivative(V,x)*Derivative(nu,x)/V
        +Integer(2)*Kappa*nu*Derivative(V,x)**Integer(2)/V**Integer(2)
        -u*Derivative(nu,x)
        +Integer(2)*nu*Derivative(u,x)
    )
]

display_system(pkf_dynamics)
```

```python
pkf_NN_builder = NNModelBuilder(pkf_dynamics,'NN_Unclosed_PKF_Burgers')
```

```python
print(pkf_NN_builder.code)
exec(pkf_NN_builder.code)
```

    from pdenetgen.model import Model
    import numpy as np
    import tensorflow.keras as keras
    from pdenetgen.symbolic.nn_builder import DerivativeFactory, TrainableScalarLayerFactory

    class NN_Unclosed_PKF_Burgers(Model):

        # Prognostic functions (sympy functions):
        prognostic_functions = (
                'u',    # Write comments on the function here
                'V_u',    # Write comments on the function here
                'nu_u_xx',    # Write comments on the function here
            )

        # Exogenous functions (sympy functions)
        exogenous_functions = (
                'closure',    # Write comments on the exogenous function here
            )

        # Spatial coordinates
        coordinates = (
                'x',    # Write comments on the coordinate here
            )

        # Set constants
        constants = (
                'kappa',    # Writes comment on the constant here
            )

        def __init__(self, shape=None, lengths=None, **kwargs):

            super().__init__() # Time scheme is set from Model.__init__()

            #---------------------------------
            # Set index array from coordinates
            #---------------------------------

            # a) Set shape
            shape = len(self.coordinates)*(100,) if shape is None else shape
            if len(shape)!=len(self.coordinates):
                raise ValueError(f"len(shape) {len(shape)} is different from len(coordinates) {len(self.coordinates)}")
            else:
                self.shape = shape

            # b) Set input shape for coordinates
            self.input_shape_x = shape[0]

            # c) Set lengths
            lengths = len(self.coordinates)*(1.0,) if lengths is None else lengths
            if len(lengths)!=len(self.coordinates):
                raise ValueError(f"len(lengths) {len(lengths)} is different from len(coordinates) {len(self.coordinates)}")
            else:
                self.lengths = lengths

            # d) Set indexes
            self._index = {}
            for k,coord in enumerate(self.coordinates):
                self._index[(coord,0)] = np.arange(self.shape[k], dtype=int)

            # Set x/dx
            #-------------
            self.dx = tuple(length/shape for length, shape in zip(self.lengths, self.shape))
            self.x = tuple(self.index(coord,0)*dx for coord, dx in zip(self.coordinates, self.dx))
            self.X = np.meshgrid(*self.x)

            #---------------------------
            # Set constants of the model
            #---------------------------

            # Set a default nan value for constants
            self.kappa = np.nan # @@ set constant value @@

            # Set constant values from external **kwargs (when provided)
            for key in kwargs:
                if key in self.constants:
                    setattr(self, key, kwargs[key])

            # Alert when a constant is np.nan
            for constant in self.constants:
                if getattr(self, constant) is np.nan:
                    print(f"Warning: constant `{constant}` has to be set")

            # Set NN models
            self._trend_model = None
            self._exogenous_model = None

        def index(self, coord, step:int):
            """ Return int array of shift index associated with coordinate `coord` for shift `step` """
            # In this implementation, indexes are memory saved in a dictionary, feed at runtime
            if (coord,step) not in self._index:
                self._index[(coord,step)] = (self._index[(coord,0)]+step)%self.shape[self.coordinates.index(coord)]
            return self._index[(coord,step)]

        def _make_trend_model(self):
            """ Generate the NN used to compute the trend of the dynamics """

            # Alias for constants
            #--------------------
            kappa = self.kappa
            if kappa is np.nan:
                raise ValueError("Constant 'kappa' is not set")

            # Set input layers
            #------------------
            # Set Alias for coordinate input shapes
            input_shape_x = self.input_shape_x

            # Set input shape for prognostic functions
            u = keras.layers.Input(shape =(input_shape_x,1,))
            V_u = keras.layers.Input(shape =(input_shape_x,1,))
            nu_u_xx = keras.layers.Input(shape =(input_shape_x,1,))

            # Set input shape for exogenous functions
            closure = keras.layers.Input(shape =(input_shape_x,1,))

            # Keras code

            # 2) Implementation of derivative as ConvNet
            # Compute derivative
            #-----------------------
            #
            #  Warning: might be modified to fit appropriate boundary conditions.
            #
            kernel_Dnu_u_xx_x_o2 = np.asarray([self.dx[self.coordinates.index('x')]**(-2),
            -2/self.dx[self.coordinates.index('x')]**2,
            self.dx[self.coordinates.index('x')]**(-2)]).reshape((3,)+(1,1))
            Dnu_u_xx_x_o2 = DerivativeFactory((3,),kernel=kernel_Dnu_u_xx_x_o2,name='Dnu_u_xx_x_o2')(nu_u_xx)

            kernel_DV_u_x_o1 = np.asarray([-1/(2*self.dx[self.coordinates.index('x')]),0.0,
            1/(2*self.dx[self.coordinates.index('x')])]).reshape((3,)+(1,1))
            DV_u_x_o1 = DerivativeFactory((3,),kernel=kernel_DV_u_x_o1,name='DV_u_x_o1')(V_u)

            kernel_Dnu_u_xx_x_o1 = np.asarray([-1/(2*self.dx[self.coordinates.index('x')]),0.0,
            1/(2*self.dx[self.coordinates.index('x')])]).reshape((3,)+(1,1))
            Dnu_u_xx_x_o1 = DerivativeFactory((3,),kernel=kernel_Dnu_u_xx_x_o1,name='Dnu_u_xx_x_o1')(nu_u_xx)

            kernel_Du_x_o2 = np.asarray([self.dx[self.coordinates.index('x')]**(-2),
            -2/self.dx[self.coordinates.index('x')]**2,
            self.dx[self.coordinates.index('x')]**(-2)]).reshape((3,)+(1,1))
            Du_x_o2 = DerivativeFactory((3,),kernel=kernel_Du_x_o2,name='Du_x_o2')(u)

            kernel_Du_x_o1 = np.asarray([-1/(2*self.dx[self.coordinates.index('x')]),0.0,
            1/(2*self.dx[self.coordinates.index('x')])]).reshape((3,)+(1,1))
            Du_x_o1 = DerivativeFactory((3,),kernel=kernel_Du_x_o1,name='Du_x_o1')(u)

            kernel_DV_u_x_o2 = np.asarray([self.dx[self.coordinates.index('x')]**(-2),
            -2/self.dx[self.coordinates.index('x')]**2,
            self.dx[self.coordinates.index('x')]**(-2)]).reshape((3,)+(1,1))
            DV_u_x_o2 = DerivativeFactory((3,),kernel=kernel_DV_u_x_o2,name='DV_u_x_o2')(V_u)

            # 3) Implementation of the trend as NNet

            #
            # Computation of trend_u
            #
            sc_mul_2 = keras.layers.Lambda(lambda x: -0.5*x,name='ScalarMulLayer_2')(DV_u_x_o1)
            sc_mul_3 = keras.layers.Lambda(lambda x: kappa*x,name='ScalarMulLayer_3')(Du_x_o2)
            mul_1 = keras.layers.multiply([Du_x_o1,u],name='MulLayer_1')
            sc_mul_4 = keras.layers.Lambda(lambda x: -1.0*x,name='ScalarMulLayer_4')(mul_1)
            trend_u = keras.layers.add([sc_mul_2,sc_mul_3,sc_mul_4],name='AddLayer_1')
            #
            # Computation of trend_V_u
            #
            sc_mul_5 = keras.layers.Lambda(lambda x: kappa*x,name='ScalarMulLayer_5')(DV_u_x_o2)
            mul_2 = keras.layers.multiply([DV_u_x_o1,u],name='MulLayer_2')
            sc_mul_6 = keras.layers.Lambda(lambda x: -1.0*x,name='ScalarMulLayer_6')(mul_2)
            mul_3 = keras.layers.multiply([Du_x_o1,V_u],name='MulLayer_3')
            sc_mul_7 = keras.layers.Lambda(lambda x: -2.0*x,name='ScalarMulLayer_7')(mul_3)
            div_0 = keras.layers.Lambda(lambda x: 1/x,name='DivLayer_0')(nu_u_xx)
            mul_4 = keras.layers.multiply([div_0,V_u],name='MulLayer_4')
            sc_mul_8 = keras.layers.Lambda(lambda x: -1.0*kappa*x,name='ScalarMulLayer_8')(mul_4)
            pow_0 = keras.layers.multiply([DV_u_x_o1,DV_u_x_o1,] ,name='PowLayer_0')
            div_1 = keras.layers.Lambda(lambda x: 1/x,name='DivLayer_1')(V_u)
            mul_5 = keras.layers.multiply([pow_0,div_1],name='MulLayer_5')
            sc_mul_9 = keras.layers.Lambda(lambda x: -0.5*kappa*x,name='ScalarMulLayer_9')(mul_5)
            trend_V_u = keras.layers.add([sc_mul_5,sc_mul_6,sc_mul_7,sc_mul_8,sc_mul_9],name='AddLayer_2')
            #
            # Computation of trend_nu_u_xx
            #
            mul_6 = keras.layers.multiply([Dnu_u_xx_x_o1,u],name='MulLayer_6')
            sc_mul_10 = keras.layers.Lambda(lambda x: -1.0*x,name='ScalarMulLayer_10')(mul_6)
            sc_mul_11 = keras.layers.Lambda(lambda x: -3.0*kappa*x,name='ScalarMulLayer_11')(Dnu_u_xx_x_o2)
            mul_7 = keras.layers.multiply([Du_x_o1,nu_u_xx],name='MulLayer_7')
            sc_mul_12 = keras.layers.Lambda(lambda x: 2.0*x,name='ScalarMulLayer_12')(mul_7)
            pow_1 = keras.layers.multiply([nu_u_xx,nu_u_xx,] ,name='PowLayer_1')
            mul_8 = keras.layers.multiply([pow_1,closure],name='MulLayer_8')
            sc_mul_13 = keras.layers.Lambda(lambda x: 4.0*kappa*x,name='ScalarMulLayer_13')(mul_8)
            pow_2 = keras.layers.multiply([Dnu_u_xx_x_o1,Dnu_u_xx_x_o1,] ,name='PowLayer_2')
            div_2 = keras.layers.Lambda(lambda x: 1/x,name='DivLayer_2')(nu_u_xx)
            mul_9 = keras.layers.multiply([pow_2,div_2],name='MulLayer_9')
            sc_mul_14 = keras.layers.Lambda(lambda x: 6.0*kappa*x,name='ScalarMulLayer_14')(mul_9)
            div_3 = keras.layers.Lambda(lambda x: 1/x,name='DivLayer_3')(V_u)
            mul_10 = keras.layers.multiply([div_3,DV_u_x_o1,Dnu_u_xx_x_o1],name='MulLayer_10')
            sc_mul_15 = keras.layers.Lambda(lambda x: kappa*x,name='ScalarMulLayer_15')(mul_10)
            div_4 = keras.layers.Lambda(lambda x: 1/x,name='DivLayer_4')(V_u)
            mul_11 = keras.layers.multiply([div_4,DV_u_x_o2,nu_u_xx],name='MulLayer_11')
            sc_mul_16 = keras.layers.Lambda(lambda x: -2.0*kappa*x,name='ScalarMulLayer_16')(mul_11)
            pow_3 = keras.layers.multiply([DV_u_x_o1,DV_u_x_o1,] ,name='PowLayer_3')
            div_5 = keras.layers.Lambda(lambda x: 1/x,name='DivLayer_5')(V_u)
            pow_4 = keras.layers.multiply([div_5,div_5,] ,name='PowLayer_4')
            mul_12 = keras.layers.multiply([pow_3,pow_4,nu_u_xx],name='MulLayer_12')
            sc_mul_17 = keras.layers.Lambda(lambda x: 2.0*kappa*x,name='ScalarMulLayer_17')(mul_12)
            add_3 = keras.layers.add([sc_mul_10,sc_mul_11,sc_mul_12,sc_mul_13,sc_mul_14,sc_mul_15,sc_mul_16,sc_mul_17],name='AddLayer_3')
            trend_nu_u_xx = keras.layers.Lambda(lambda x: x+-1.0*kappa,name='ScalarAddLayer_0')(add_3)

            # 4) Set 'input' of model
            inputs = [
                    # Prognostic functions
                    u,V_u,nu_u_xx,
                    # Exogenous functions
                    closure,

                ]

            # 5) Set 'outputs' of model
            outputs = [
                trend_u,
                trend_V_u,
                trend_nu_u_xx,
                ]

            model = keras.models.Model(inputs=inputs, outputs=outputs)
            #model.trainable = False
            self._trend_model = model

        def trend(self, t, state):
            """ Trend of the dynamics """

            if self._trend_model is None:
                self._make_trend_model()

            # Init output state with pointer on data
            #-------------------------------------------

            #   a) Set the output array
            dstate = np.zeros(state.shape)

            #   b) Set pointers on output array `dstate` for the computation of the physical trend (alias only).
            du = dstate[0]
            dV_u = dstate[1]
            dnu_u_xx = dstate[2]

            # Load physical functions from state
            #------------------------------------
            u = state[0]
            V_u = state[1]
            nu_u_xx = state[2]

            # Compute exogenous functions
            #-----------------------------
            exogenous = self.compute_exogenous(t, state) # None if no exogenous function exists
            closure = exogenous[0]

            # Compute the trend value from model.predict
            #-------------------------------------------
            inputs = [
                # Prognostic functions
                    u,
                    V_u,
                    nu_u_xx,
                    # Exogenous functions
                    closure,
                    ]

            dstate = self._trend_model.predict( inputs )

            if not isinstance(dstate,list):
                dstate = [dstate]

            return np.array(dstate)

        def _make_dynamical_trend(self):
            """
            Computation of a trend model so to be used in a time scheme (as solving a dynamical system or an ODE)

            Description:
            ------------

            In the present implementation, the inputs of the trend `self._trend_model` is a list of fields, while
            entry of a time-scheme is a single array which contains all fields.

            The aims of `self._dynamical_trend` is to produce a Keras model which:
            1. takes a single array as input
            2. extract the `self._trend_model` input list from the input array
            3. compute the trends from `self._trend_model`
            4. outputs the trends as a single array

            Explaination of the code:
            -------------------------

            Should implement a code as the following, that is valid for the PKF-Burgers

            def _make_dynamical_trend(self):

                if self._trend_model is None:
                    self._make_trend_model()

                # 1. Extract the input of the model

                # 1.1 Set the input as an array
                state = keras.layers.Input(shape=(3,self.input_shape_x,1))

                # 1.2 Extract each components of the state
                u = keras.layers.Lambda(lambda x : x[:,0,:,:])(state)
                V = keras.layers.Lambda(lambda x : x[:,1,:,:])(state)
                nu_u_xx = keras.layers.Lambda(lambda x : x[:,2,:,:])(state)

                # 2. Compute the trend
                trend_u, trend_V, trend_nu = self._trend_model([u,V,nu_u_xx])

                # 3. Outputs the trend as a single array

                # 3.1 Reshape trends
                trend_u = keras.layers.Reshape((1,self.input_shape_x,1))(trend_u)
                trend_V = keras.layers.Reshape((1,self.input_shape_x,1))(trend_V)
                trend_nu = keras.layers.Reshape((1,self.input_shape_x,1))(trend_nu)

                # 3.2 Concatenates all trends
                trends = keras.layers.Concatenate(axis=1)([trend_u,trend_V,trend_nu])

                # 4. Set the dynamical_trend model
                self._dynamical_trend = keras.models.Model(inputs=state,outputs=trends)
            """

            if self._trend_model is None:
                self._make_trend_model()

            for exclude_case in ['constant_functions','exogenous_functions']:
                if hasattr(self,exclude_case):
                    raise NotImplementedError(f'Design of dynamical_model with {exclude_case} is not implemented')

            # Case 1 --  corresponds to the _trend_model if input is a single field
            if not isinstance(self._trend_model.input_shape, list):
                self._dynamical_trend = self._trend_model
                return

            # Case 2 -- Case where multiple list is used

            # 1. Extract the input of the model

            # 1.1 Set the input as an array
            """ from PKF-Burgers code:
            state = keras.layers.Input(shape=(3,self.input_shape_x,1))
            """

            # 1.1.1 Compute the input_shape from _trend_model
            shapes = []
            dimensions  = []
            for shape in self._trend_model.input_shape:
                shape = shape[1:] # Exclude batch_size (assumed to be at first)
                shapes.append(shape)
                dimensions.append(len(shape)-1)

            max_dimension = max(dimensions)
            if max_dimension!=1:
                if 1 in dimensions:
                    raise NotImplementedError('1D fields incompatible with 2D/3D fields')

            # todo: add test to check compatibility of shapes!!!!

            if max_dimension in [1,2]:
                input_shape = (len(shapes),)+shapes[0]
            elif max_dimension==3:

                # a. check the size of 2D fields: this is given by the first 2D field.
                for shape, dimension in zip(shapes, dimensions):
                    if dimension==2:
                        input_shape_2D = shape
                        break

                # b. Compute the numbers of 2D fields: this corresponds to the number of 3D layers and the number of 2D fields.

                for shape, dimension in zip(shapes, dimensions):
                    if dimension==2:
                        nb_outputs += 1
                    else:
                        nb_outputs += shape[0]

                input_shape = (nb_outputs,)+input_shape_2D

            # 1.1.2 Init the state of the dynamical_trend
            state = keras.layers.Input(shape=input_shape)

            # 1.2 Extract each components of the state
            """ From PKF-Burgers code:

            u = keras.layers.Lambda(lambda x : x[:,0,:,:])(state)
            V = keras.layers.Lambda(lambda x : x[:,1,:,:])(state)
            nu_u_xx = keras.layers.Lambda(lambda x : x[:,2,:,:])(state)

            inputs = [u, V, nu_u_xx]
            """

            def get_slice(dimension, k):
                def func(x):
                    if dimension == 1:
                        return x[:,k,:,:]
                    elif dimension == 2:
                        return x[:,k,:,:]
                return func

            def get_slice_3d(start,end):
                def func(x):
                    return x[:,start:end,:,:,:]
                return func

            inputs = []
            if max_dimension in [1,2]:
                for k in range(len(shapes)):
                    inputs.append(keras.layers.Lambda(get_slice(max_dimension,k))(state))
                    #if max_dimension==1:
                    #    inputs.append(keras.layers.Lambda(lambda x : x[:,k,:,:])(state))
                    #
                    #if max_dimension==2:
                    #    inputs.append(keras.layers.Lambda(lambda x : x[:,k,:,:,:])(state))
            else:
                k=0
                for shape, dimension in zip(shapes, dimensions):
                    if dimension==2:
                        #inputs.append(keras.layers.Lambda(lambda x : x[:,k,:,:,:])(state))
                        inputs.append(keras.layers.Lambda(get_slice(dimension,k))(state))
                        k += 1
                    if dimension==3:
                        start = k
                        end = start+shape[0]
                        inputs.append(keras.layers.Lambda(get_slice_3d(start,end))(state))
                        k = end

            # 2. Compute the trend
            """ From PKF-Burgers code
            trend_u, trend_V, trend_nu = self._trend_model([u,V,nu_u_xx])
            """
            trends = self._trend_model(inputs)

            # 3. Outputs the trend as a single array

            # 3.1 Reshape trends
            """ from PKF-Burgers code
            trend_u = keras.layers.Reshape((1,self.input_shape_x,1))(trend_u)
            trend_V = keras.layers.Reshape((1,self.input_shape_x,1))(trend_V)
            trend_nu = keras.layers.Reshape((1,self.input_shape_x,1))(trend_nu)
            """

            reshape_trends = []
            for trend, dimension in zip(trends, dimensions):

                #shape = tuple(dim.value for dim in trend.shape[1:])
                # update from keras -> tensorflow.keras
                shape = tuple(dim for dim in trend.shape[1:])

                if dimension==1 or dimension==2:
                    # for 1D fields like (128,1) transform into (1,128,1)
                    # for 2D fields like (128,128,1) transform into (1,128,128,1)
                    shape = (1,)+shape
                elif dimension==3:
                    # 3D fields can be compated: two fields (36,128,128,1) become the single field (72,128,128,1)
                    pass
                else:
                    raise NotImplementedError

                reshape_trends.append(keras.layers.Reshape(shape)(trend))

            # 3.2 Concatenates all trends
            """ From PKF-Burgers code:
            trends = keras.layers.Concatenate(axis=1)([trend_u,trend_V,trend_nu])
            """
            trends = keras.layers.Concatenate(axis=1)(reshape_trends)

            # 2.5 Compute the model
            self._dynamical_trend = keras.models.Model(inputs=state,outputs=trends)

        def _make_exogenous_model(self):
            raise NotImplementedError

        def compute_exogenous(self, t, state):
            """ Computation of exogenous functions """
            if self._exogenous_model is None:
                self._make_exogenous_model()

            raise NotImplementedError # if no exogenous functions are needed

#### Construction of a closure as a NN from parameterized form

The aim is to compute the closure proposal
$$a\frac{\frac{\partial^{2}}{\partial x^{2}} \operatorname{{\nu_{u,xx}}}{\left(t,x \right)}}{\operatorname{{\nu_{u,xx}}}^{2}{\left(t,x \right)}} +b \frac{1}{ \operatorname{{\nu_{u,xx}}}^{2}{\left(t,x \right)}} +c\frac{ \left(\frac{\partial}{\partial x} \operatorname{{\nu_{u,xx}}}{\left(t,x \right)}\right)^{2}}{\operatorname{{\nu_{u,xx}}}^{3}{\left(t,x \right)}},$$
as an exogeneous neural network, where $(a,b,c)$ are trainable.

```python
class ClosedPKFBurgers(NN_Unclosed_PKF_Burgers):
    def _make_exogenous_model(self):

        u = keras.layers.Input(shape=(self.input_shape_x,1))
        V = keras.layers.Input(shape=(self.input_shape_x,1))
        nu_u_xx = keras.layers.Input(shape=(self.input_shape_x,1))

        #
        # Computation of the spatial derivatives
        #

        kernel_Dnu_u_xx_x_o2 = np.asarray([self.dx[self.coordinates.index('x')]**(-2),
        -2/self.dx[self.coordinates.index('x')]**2,
        self.dx[self.coordinates.index('x')]**(-2)]).reshape((3,)+(1,1))
        Dnu_u_xx_x_o2 = DerivativeFactory((3,),kernel=kernel_Dnu_u_xx_x_o2,name='Dnu_u_xx_x_o2')(nu_u_xx)

        kernel_Dnu_u_xx_x_o1 = np.asarray([-1/(2*self.dx[self.coordinates.index('x')]),0.0,
        1/(2*self.dx[self.coordinates.index('x')])]).reshape((3,)+(1,1))
        Dnu_u_xx_x_o1 = DerivativeFactory((3,),kernel=kernel_Dnu_u_xx_x_o1,name='Dnu_u_xx_x_o1')(nu_u_xx)

        #
        # Design of the unknown closure to train
        #

        # Terme 1
        div_14 = keras.layers.Lambda(lambda x: 1/x,name='DivLayer_14')(nu_u_xx)
        pow_12 = keras.layers.multiply([div_14,div_14,] ,name='PowLayer_12')
        term1 = keras.layers.multiply([pow_12,Dnu_u_xx_x_o2],name='MulLayer_25')

        # Terme 2
        div_13 = keras.layers.Lambda(lambda x: 1/x,name='DivLayer_13')(nu_u_xx)
        term2 = keras.layers.multiply([div_13,div_13,] ,name='PowLayer_11')

        # Terme 3
        pow_13 = keras.layers.multiply([Dnu_u_xx_x_o1,Dnu_u_xx_x_o1,] ,name='PowLayer_13')
        div_15 = keras.layers.Lambda(lambda x: 1/x,name='DivLayer_15')(nu_u_xx)
        pow_14 = keras.layers.multiply([div_15,div_15,div_15,] ,name='PowLayer_14')
        term3 = keras.layers.multiply([pow_13,pow_14],name='MulLayer_26')

        # Product by (a,b,c), implemented as Conv1D
        term1 = keras.layers.Conv1D(1,1,name='times_a',padding='same',use_bias=False,activation='linear')(term1)
        term2 = keras.layers.Conv1D(1,1,name='times_b',padding='same',use_bias=False,activation='linear')(term2)
        term3 = keras.layers.Conv1D(1,1,name='times_c',padding='same',use_bias=False,activation='linear')(term3)

        closure = keras.layers.add([term1, term2, term3],name='Closure')

        self._exogenous_model = keras.models.Model(inputs=[u,V,nu_u_xx], outputs=[closure])

    def compute_exogenous(self, t, state):

        if self._exogenous_model is None:
            self._make_exogenous_model()

        u,V,nu = state
        closure = self._exogenous_model.predict([u,V,nu])
        if not isinstance(closure, list):
            closure = [closure]
        return closure

    def _make_full_trend(self):

        if self._trend_model is None:
            self._make_trend_model()
        if self._exogenous_model is None:
            self._make_exogenous_model()

        state = keras.layers.Input(shape=(3,self.input_shape_x,1))

        u = keras.layers.Lambda(lambda x : x[:,0,:,:])(state)
        V = keras.layers.Lambda(lambda x : x[:,1,:,:])(state)
        nu_u_xx = keras.layers.Lambda(lambda x : x[:,2,:,:])(state)

        closure = self._exogenous_model([u,V,nu_u_xx])
        trend_u, trend_V, trend_nu = self._trend_model([u,V,nu_u_xx,closure])

        trend_u = keras.layers.Reshape((1,self.input_shape_x,1))(trend_u)
        trend_V = keras.layers.Reshape((1,self.input_shape_x,1))(trend_V)
        trend_nu = keras.layers.Reshape((1,self.input_shape_x,1))(trend_nu)

        trend = keras.layers.Concatenate(axis=1)([trend_u,trend_V,trend_nu])
        self._full_trend = keras.models.Model(inputs=state,outputs=trend)
```

```python
closed_burgers = ClosedPKFBurgers(shape=(241,),kappa=kappa)
```

```python
closed_burgers._make_full_trend()
```

**Set initial PKF fields**

```python
# Set initial condition for the variance parameter 'V_u'
V0 = (0.01*Umax)**2 + 0*U0

# Set the initial condition for the diffusion
# L**2 = 2nu t => nu = 0.5*L**2
lh = 0.02*domain.lengths[0]
nu0 = 0.5*lh**2 + 0*U0

state0 = np.asarray([U0, V0,nu0])
normalization = {
                'Velocity':U0.max(),
                'Variance':V0.max(),
                'Length-scale':lh
                }
```

```python
length_scale = lambda nu: np.sqrt(2*nu)
plt.figure(figsize=(12,12))
for k,field in enumerate(normalization):
    plt.subplot(221+k)
    if field=='Length-scale':
        data = {0:length_scale(state0[k])}
    else:
        data = {0:state0[k]}
    plot_results(data, label=field)

```

## Numerical application

```python
def plot_pkf_traj_ensemble(traj):
    plt.figure(figsize=(15,5))
    for k,field in enumerate(normalization):
        if field=='Length-scale':
            data = {time:length_scale(traj[time][k]) for time in traj}
        else:
            data = {time:traj[time][k] for time in traj}
        plt.subplot(131+k)
        plot_results(data,label=field,normalisation=normalization[field])
```

```python
def plot_pkf_traj_NN(traj):
    plt.figure(figsize=(15,5))
    for k,field in enumerate(normalization):
        if field=='Length-scale':
            data = {time:length_scale(traj[time][k][0,:,0]) for time in traj}
        else:
                data = {time:traj[time][k][0,:,0] for time in traj}
        plt.subplot(131+k)
        plot_results(data,label=field,normalisation=normalization[field])

```

```python
state0 = np.asarray([U0.reshape((1,)+U0.shape+(1,)),
                     V0.reshape((1,)+V0.shape+(1,)),
                     nu0.reshape((1,)+nu0.shape+(1,))])
```

### Generation of a database

**Gaussian random vector of Gaussian correlation function**

```python
# Création d'une matrice de covariance d'erreur de prévision initiale: $P_0$
#   Cette matrice est construite comme une matrice homogène de corrélation Gaussienne et de longueur de portée l_h

# 1) Définition de la fonction de corrélation homogène
gauss = lambda x : np.exp(-0.5*x**2/lh**2) # lh has been previously specified
correlation = gauss(domain.x[0]-domain.x[0][domain.shape[0]//2])
spectrum = np.abs(np.fft.fft(correlation))

# 2) Construction de B^(1/2)
std_spectrum = np.sqrt(spectrum)
def make_sample():
    zeta = np.random.normal(size=domain.shape)
    zeta = np.fft.fft(zeta)
    ef = np.fft.ifft(std_spectrum * zeta)
    ef = np.real(ef)
    return ef
```

```python
plt.figure(figsize=(12,5))
plt.subplot(121)
plt.plot(domain.x[0], correlation)
plt.title('Homogenous correlation function');
plt.subplot(122)
for k in range(10):
    plt.plot(domain.x[0], make_sample())
plt.title("Example of sample errors");
```

##### **Diagnosis tool for ensemble estimation of expectation/variance/diffusion tensor**

```python
def make_init_ensemble(Ne):
    return np.array([make_sample() for k in range(Ne)])

def estimate_covariance(ensemble):
    mean = ensemble.mean(0)
    error = (ensemble - mean)/np.sqrt(len(ensemble))
    return error.T @ error

class EnsembleDiagnosis(object):

    def __init__(self, ensemble, base_space):
        self.base_space = base_space

        if isinstance(ensemble, list):
            ensemble = np.array(ensemble)

        if len(ensemble.shape)==3:
            ensemble = np.array([elm[0] for elm in ensemble])

        # 1) Computation of the mean
        self.mean = ensemble.mean(axis=0)

        # 2) Computation of the variance
        self.std = ensemble.std(axis=0)
        self.variance = self.std*self.std

        # 3) Computation of the metric terms
        #  we use the formula g_ij = E[(D_i eps)(D_j eps)]

        #  a) Computation of the normalized error
        epsilon = (ensemble-self.mean)/self.std

        #  b) Computation of derivatives
        n = self.base_space.shape[0]
        K = np.arange(n)
        kp = (K+1)%n
        km = (K-1)%n
        dx = self.base_space.dx[0]
        Depsilon = np.array([(eps[kp]-eps[km])/(2*dx) for eps in epsilon])
        self.metric = (Depsilon*Depsilon).mean(axis=0)     # see Pannekoucke et al. (2018) for details

        # Computation of the diffusion tensor
        self.diffusion = 0.5*1/self.metric
        self.length_scale = np.sqrt(2*self.diffusion)
```

##### **Ensemble validation for the covariance setting**

```python
Ne = 1600

ensemble = make_init_ensemble(Ne)

mean = ensemble.mean(axis=0)
std = ensemble.std(axis=0)

print(f"Validation of the mean (=0): {mean.mean()} +/- {mean.std()}" )
print(f"Validation of the standard-deviation (=1): {std.mean()} +/- {std.std()}" )

ens_diagnosis = EnsembleDiagnosis(ensemble, domain)
nu_h = 0.5*lh**2

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

plt.subplot(131)
plt.plot(ens_diagnosis.mean)
plt.title('Moyenne')

plt.subplot(132)
plt.plot(ens_diagnosis.variance)
plt.title('Variance')

plt.subplot(133)
plt.plot(ens_diagnosis.diffusion/nu_h)
plt.title('diffusion (normalisée par $nu_h$)')
```

**Computation of a large ensemble (1600 members) to build a reference**

```python
# Standard deviation for the initial perturbation
sigma_f = 0.01*U0.max()
```

```python
# Set parameters for ensemble estimation
large_Ne = 1600

# 1. Set the initial background state
random_U0 = U0 + sigma_f*make_init_ensemble(1)[0]

# 2. Build an ensemble of initial perturbed state
ensemble = make_init_ensemble(large_Ne)
ensemble_ea = np.array([random_U0+sigma_f*ea for ea in ensemble])
ensemble_ea = ensemble_ea.reshape((1,)+ensemble_ea.shape+(1,))
print(f"shape of ensemble_ea: {ensemble_ea.shape}")

# 3. Build the ensemble of forecast using the NN architecture
ensemble_forecast = burgers.forecast(times,ensemble_ea)
```

    shape of ensemble_ea: (1, 1600, 241, 1)

```python
# 4. Compute diagnosis from ensemble
ensemble_traj = {}
for time in times[::50]:
    diagnosis = EnsembleDiagnosis(ensemble_forecast[time][0,:,:,0], domain)
    ensemble_traj[time] = [diagnosis.mean, diagnosis.variance, diagnosis.diffusion]
```

```python
plot_pkf_traj_ensemble(ensemble_traj)
```

#### **Generation of the training data set**

```python
def generate_data(k, Ne=400):
    # 1. Set the initial background state
    random_U0 = U0 + sigma_f*make_init_ensemble(1)[0]

    # 2. Build an ensemble of initial perturbed state
    ensemble = make_init_ensemble(Ne)
    ensemble_ea = np.array([random_U0+sigma_f*ea for ea in ensemble])
    ensemble_ea = ensemble_ea.reshape((1,)+ensemble_ea.shape+(1,))

    # 3. Compute the ensemble of forecasts
    ensemble_forecast = burgers.forecast(times,ensemble_ea)

    # 4. Compute the diagnosis
    diagnosis_list = []
    for time in times:
        diagnosis = EnsembleDiagnosis(ensemble_forecast[time][0,:,:,0], domain)
        diagnosis_list.append( np.array([diagnosis.mean, diagnosis.variance, diagnosis.diffusion]))

    return diagnosis_list
```

```python
data_size = 400  # for Ne=400, this takes 1h09'01'' so take care with this..
save_file = "pkf-dataset.npy"

generate_data_set = False
parallel_diagnosis = False

try:
    # load data
    data = np.load(save_file)
    data = data.reshape(data.shape+(1,))
except:
    # 1. Generate data
    #data = [generate_data(k) for k in range(data_size)]
    data = []
    for k in range(data_size):
        if k%5==0:
            print(k)
        data.append(generate_data(k))

    # 2. Save data
    data = np.array(data)
    np.save(save_file,data)

```

```python
data.shape
```

### Training of the closure

**Make a RK4 time scheme**

```python
# Schéma temporelle de type RK4
def make_time_scheme(dt, trend):
    """ Implémentation d'un schéma de RK4 sous forme de réseau de neurones """

    state = keras.layers.Input(shape = trend.input_shape[1:])

    # k1
    k1 = trend(state)
    # k2
    _tmp_1 = keras.layers.Lambda(lambda x : 0.5*dt*x)(k1)
    input_k2 = keras.layers.add([state,_tmp_1])
    k2 = trend(input_k2)
    # k3
    _tmp_2 = keras.layers.Lambda(lambda x : 0.5*dt*x)(k2)
    input_k3 = keras.layers.add([state,_tmp_2])
    k3 = trend(input_k3)
    # k4
    _tmp_3 = keras.layers.Lambda(lambda x : dt*x)(k3)
    input_k4 = keras.layers.add([state,_tmp_3])
    k4 = trend(input_k4)

    # output
    # k2+k3
    add_k2_k3 = keras.layers.add([k2,k3])
    add_k2_k3_mul2 = keras.layers.Lambda(lambda x:2.*x)(add_k2_k3)
    # Add k1,k4
    _sum = keras.layers.add([k1,add_k2_k3_mul2,k4])
    # *dt
    _sc_mul = keras.layers.Lambda(lambda x:dt/6.*x)(_sum)
    output = keras.layers.add([state, _sc_mul])

    time_scheme = keras.models.Model(inputs =[state], outputs=[output])
    return time_scheme
```

```python
closed_pkf_burgers = ClosedPKFBurgers(shape=(241,),kappa=kappa)

closed_pkf_burgers._make_full_trend()

trained = closed_pkf_burgers._full_trend.get_weights()
trained = np.array((trained[0], trained[1], trained[2])).flatten()
trained
```

    array([array([[[-120.5]],

           [[   0. ]],

           [[ 120.5]]], dtype=float32),
           array([[[  58081.]],

           [[-116162.]],

           [[  58081.]]], dtype=float32),
           array([[[0.8571569]]], dtype=float32)], dtype=object)

```python
closed_pkf_burgers._make_full_trend()
closed_pkf_burgers._full_trend.summary()
```

    Model: "model_7"
    __________________________________________________________________________________________________
    Layer (type)                    Output Shape         Param #     Connected to
    ==================================================================================================
    input_18 (InputLayer)           [(None, 3, 241, 1)]  0
    __________________________________________________________________________________________________
    lambda_6 (Lambda)               (None, 241, 1)       0           input_18[0][0]
    __________________________________________________________________________________________________
    lambda_7 (Lambda)               (None, 241, 1)       0           input_18[0][0]
    __________________________________________________________________________________________________
    lambda_8 (Lambda)               (None, 241, 1)       0           input_18[0][0]
    __________________________________________________________________________________________________
    model_5 (Model)                 (None, 241, 1)       9           lambda_6[0][0]
                                                                     lambda_7[0][0]
                                                                     lambda_8[0][0]
    __________________________________________________________________________________________________
    model_4 (Model)                 [(None, 241, 1), (No 18          lambda_6[0][0]
                                                                     lambda_7[0][0]
                                                                     lambda_8[0][0]
                                                                     model_5[2][0]
    __________________________________________________________________________________________________
    reshape_6 (Reshape)             (None, 1, 241, 1)    0           model_4[2][0]
    __________________________________________________________________________________________________
    reshape_7 (Reshape)             (None, 1, 241, 1)    0           model_4[2][1]
    __________________________________________________________________________________________________
    reshape_8 (Reshape)             (None, 1, 241, 1)    0           model_4[2][2]
    __________________________________________________________________________________________________
    concatenate_2 (Concatenate)     (None, 3, 241, 1)    0           reshape_6[0][0]
                                                                     reshape_7[0][0]
                                                                     reshape_8[0][0]
    ==================================================================================================
    Total params: 27
    Trainable params: 3
    Non-trainable params: 24
    __________________________________________________________________________________________________

```python
closed_pkf_burgers._exogenous_model.summary()
```

    Model: "model_5"
    __________________________________________________________________________________________________
    Layer (type)                    Output Shape         Param #     Connected to
    ==================================================================================================
    input_16 (InputLayer)           [(None, 241, 1)]     0
    __________________________________________________________________________________________________
    periodic1d_17 (Periodic1D)      (None, 243, 1)       0           input_16[0][0]
    __________________________________________________________________________________________________
    periodic1d_16 (Periodic1D)      (None, 243, 1)       0           input_16[0][0]
    __________________________________________________________________________________________________
    Dnu_u_xx_x_o1 (Conv1D)          (None, 243, 1)       3           periodic1d_17[0][0]
    __________________________________________________________________________________________________
    DivLayer_14 (Lambda)            (None, 241, 1)       0           input_16[0][0]
    __________________________________________________________________________________________________
    Dnu_u_xx_x_o2 (Conv1D)          (None, 243, 1)       3           periodic1d_16[0][0]
    __________________________________________________________________________________________________
    crop1d_17 (Crop1D)              (None, 241, 1)       0           Dnu_u_xx_x_o1[0][0]
    __________________________________________________________________________________________________
    DivLayer_15 (Lambda)            (None, 241, 1)       0           input_16[0][0]
    __________________________________________________________________________________________________
    PowLayer_12 (Multiply)          (None, 241, 1)       0           DivLayer_14[0][0]
                                                                     DivLayer_14[0][0]
    __________________________________________________________________________________________________
    crop1d_16 (Crop1D)              (None, 241, 1)       0           Dnu_u_xx_x_o2[0][0]
    __________________________________________________________________________________________________
    DivLayer_13 (Lambda)            (None, 241, 1)       0           input_16[0][0]
    __________________________________________________________________________________________________
    PowLayer_13 (Multiply)          (None, 241, 1)       0           crop1d_17[0][0]
                                                                     crop1d_17[0][0]
    __________________________________________________________________________________________________
    PowLayer_14 (Multiply)          (None, 241, 1)       0           DivLayer_15[0][0]
                                                                     DivLayer_15[0][0]
                                                                     DivLayer_15[0][0]
    __________________________________________________________________________________________________
    MulLayer_25 (Multiply)          (None, 241, 1)       0           PowLayer_12[0][0]
                                                                     crop1d_16[0][0]
    __________________________________________________________________________________________________
    PowLayer_11 (Multiply)          (None, 241, 1)       0           DivLayer_13[0][0]
                                                                     DivLayer_13[0][0]
    __________________________________________________________________________________________________
    MulLayer_26 (Multiply)          (None, 241, 1)       0           PowLayer_13[0][0]
                                                                     PowLayer_14[0][0]
    __________________________________________________________________________________________________
    times_a (Conv1D)                (None, 241, 1)       1           MulLayer_25[0][0]
    __________________________________________________________________________________________________
    times_b (Conv1D)                (None, 241, 1)       1           PowLayer_11[0][0]
    __________________________________________________________________________________________________
    times_c (Conv1D)                (None, 241, 1)       1           MulLayer_26[0][0]
    __________________________________________________________________________________________________
    input_14 (InputLayer)           [(None, 241, 1)]     0
    __________________________________________________________________________________________________
    input_15 (InputLayer)           [(None, 241, 1)]     0
    __________________________________________________________________________________________________
    Closure (Add)                   (None, 241, 1)       0           times_a[0][0]
                                                                     times_b[0][0]
                                                                     times_c[0][0]
    ==================================================================================================
    Total params: 9
    Trainable params: 3
    Non-trainable params: 6
    __________________________________________________________________________________________________

```python
time_scheme = make_time_scheme(dt, closed_pkf_burgers._full_trend)
#time_scheme.summary()
```

**Extract data for the training**

```python
data[0].shape
```

```python
select_from = 400 # 200
X = np.array([elm[select_from:-1] for elm in data])
Y = np.array([elm[select_from+1:] for elm in data])
```

```python
X = X.reshape((np.prod(X.shape[:2]),3,241,1))
Y = Y.reshape((np.prod(Y.shape[:2]),3,241,1))
```

```python
X.shape
```

**Training of the NN**

```python
trained = closed_pkf_burgers._full_trend.get_weights()
trained = np.array((trained[0], trained[1], trained[2])).flatten()
trained
```

    array([array([[[-120.5]],

           [[   0. ]],

           [[ 120.5]]], dtype=float32),
           array([[[  58081.]],

           [[-116162.]],

           [[  58081.]]], dtype=float32),
           array([[[0.8571569]]], dtype=float32)], dtype=object)

```python
# Expérience d'apprentissage:
# 2. Adam
lr = 0.1
epochs = 30

for iteration in range(3):
    # 1. Set the learning
    time_scheme.compile(optimizer=keras.optimizers.Adam(lr=lr),
        loss='mean_squared_error')

    # 2. Train
    history = time_scheme.fit(X,Y,epochs=epochs, batch_size=32,verbose=0)
    print(f"iteration {iteration} is complet")

    # 3. Plot history
    plt.figure()
    plt.plot(history.history['loss'])

    # 4. Update the learning rate for next iteration
    lr = lr/10

```

### Comparison with the theoretically designed closure

**Compare the weights with the previous theoretical closure**

The weights of the theoretical closure are : 1, 3/4, -2

```python
trained = closed_pkf_burgers._full_trend.get_weights()
trained = np.array((trained[0], trained[1], trained[2])).flatten()
trained
```

    array([array([[[-120.5]],

           [[   0. ]],

           [[ 120.5]]], dtype=float32),
           array([[[  58081.]],

           [[-116162.]],

           [[  58081.]]], dtype=float32),
           array([[[0.93310356]]], dtype=float32)], dtype=object)

```python
theoretical = np.array([1,3/4,-2])
```

```python
relative_error = (trained - theoretical)/theoretical
```

```python
relative_error*100
```

    array([array([[[-12150.]],

           [[  -100.]],

           [[ 11950.]]], dtype=float32),
           array([[[  7744033.5]],

           [[-15488367. ]],

           [[  7744033.5]]], dtype=float32),
           array([[[-146.65518]]], dtype=float32)], dtype=object)

**Exemple de prévision réalisée avec le modèle calibré**

```python
# default
closed_pkf_burgers.set_dt(dt)
times = closed_pkf_burgers.window(1)
saved_times = times[::50]
print('saved_times :' ,saved_times)
```

    saved_times : [0.  0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1. ]

```python
state0.shape
```

```python
trained_unclosed_traj = closed_pkf_burgers.forecast(times, state0, saved_times)
```

```python
normalization
```

    {'Velocity': 0.4999946897600651,
     'Variance': 2.4999468978826374e-05,
     'Length-scale': 0.02}

```python
# PKF using trained closure
plot_pkf_traj_NN(trained_unclosed_traj)
plt.savefig('./figures/burgers-exogeneous-b.pdf')
```

```python
# ensemble of forecast statistics
plot_pkf_traj_ensemble(ensemble_traj)
plt.savefig('./figures/burgers-exogeneous-a.pdf')
```

## Conclusion

In this notebook, we have considered the uncertainty prediction for the Burgers dynamics where an unclosed term is present.

A closure has been proposed and implemented as an exogeneous neural network with three unknowns.

A synthetic dataset has been  used to train the NN. The resulting closure has shown to be relevant to predict the uncertainty.