## Universidade Federal do Ceará - Departamento de Engenharia de Teleinformática

Prof. Lucas N. Ribeiro

Prof. João César M. Mota

### Processamento Digital de Sinais 2019.1 - Técnicas de projeto de filtros digitais

* Filtros seletivos em frequência: admissão e rejeição de componentes de frequências
* O que isso quer dizer?

## Projeto de filtro: passos

1. Especificação das propriedades desejadas;
2. Aproximação das especificações usando um sistema discreto e causal;
3. Realização do sistema.

## Exemplo

Projeto de filtro passa-baixa com frequência de corte $\omega_c$.

A resposta em frequência ideal com fase linear é dada por:

$$
\begin{equation}
    H_d(e^{j\omega}) = \begin{cases}
        e^{-j\alpha \omega} & |\omega| \leq \omega_c\\
        0 & c.c.
    \end{cases}
\end{equation}
$$

Resposta ao impulso:

$$
\begin{equation}
    h_d(n) = \frac{\sin(n-\alpha) \omega_c}{\pi (n-a)}, \quad n \in \mathbb{Z}
\end{equation}
$$

```python
import numpy as np
import matplotlib.pyplot as plt
from scipy import signal
from matplotlib.ticker import MaxNLocator

def plot_ideal_freq(wc, alpha):
    w = np.linspace(0, np.pi, 100)
    H = [np.exp(-1j*alpha*wi) if (np.abs(wi) <= wc) else 0 for wi in w]

    magH = [ np.abs(Hi)**2 for Hi in H ]
    phaH = [ np.angle(Hi) for Hi in H ]
    plt.figure()
    plt.plot(w/np.pi, magH)
    plt.xlabel('Frequência discreta [$\pi$ rad/s]')
    plt.ylabel('Magnitude')

def plot_ideal_tempo(wc, alpha):
    indices = np.arange(-20, 30, 1)
    h = np.zeros((len(indices), 1))
    for n in np.arange(0, len(indices)):
        if indices[n] != alpha:
            h[n] = (np.sin((indices[n] - alpha)*wc))/(np.pi*(indices[n]-alpha))
        else:
            h[n] = 1

    plt.figure()
    plt.stem(indices,h)
    plt.xlabel('Amostra')
    ax = plt.gca()
    ax.xaxis.set_major_locator(MaxNLocator(integer=True))

plot_ideal_freq(wc=0.5*np.pi, alpha=10)
plot_ideal_tempo(wc=0.5*np.pi, alpha=10)
```

## Problemas com esse filtro?

* Não-causal
* Instável

E se o filtro ideal for truncado?

## Janelamento

A resposta ao impulso ideal é multiplicada por uma janela causal de comprimento $N$ tal que

$h(n) = h_d(n) w(n)$

em que $0 \leq n \leq N-1$.

Tipos de janelas:

* Retangular
* Hanning
* Blackman-Harris
* ...

```python
def pb_retangular(N, wc):
    alpha = N/2
    h = np.zeros((N, 1))
    for n in np.arange(N):
        if n != alpha:
            h[n] = (np.sin((n - alpha)*wc))/(np.pi*(n-alpha))
        else:
            h[n] = 1
    return(h)

def calc_freqdomain(h, Nf = 100):
    N = len(h)
    w = np.linspace(0, np.pi, Nf)
    H = np.zeros((Nf,1)) + 1j*np.zeros((Nf,1))

    for nf in np.arange(Nf):
        for n in np.arange(N):
            H[nf] += h[n]*np.exp(-1j*n*w[nf])

    return(H, w)

def plot_freqdomain(H, w, ax = None):
    magH = [ 10*np.log10(np.abs(Hi)**2) for Hi in H ]
    phaH = [ np.angle(Hi) for Hi in H ]

    if ax is None:
        plt.figure()
        plt.plot(w/np.pi, magH)
    else:
        ax.plot(w/np.pi, magH)

h_pb1 = pb_retangular(5, 0.5*np.pi)
(H_pb1, w_pb1) = calc_freqdomain(h_pb1)

h_pb2 = pb_retangular(15, 0.5*np.pi)
(H_pb2, w_pb2) = calc_freqdomain(h_pb2)

h_pb3 = pb_retangular(55, 0.5*np.pi)
(H_pb3, w_pb3) = calc_freqdomain(h_pb3, Nf = 200)

fig = plt.figure()
plt.stem(h_pb1)
plt.title('Ordem 5')

fig = plt.figure()
plt.stem(h_pb2)
plt.title('Ordem 15')

fig = plt.figure()
plt.stem(h_pb3)
plt.title('Ordem 55')

fig = plt.figure()
ax = fig.gca()
plot_freqdomain(H_pb1, w_pb1, ax = ax)
plot_freqdomain(H_pb2, w_pb2, ax = ax)
plot_freqdomain(H_pb3, w_pb3, ax = ax)
plt.xlabel('Frequência discreta [$\pi$ rad/s]')
plt.ylabel('Magnitude')
plt.legend(['Ordem 5', 'Ordem 15', 'Ordem 55'])
```

## Parâmetros de projeto:

- Oscilação na faixa de passagem
- Oscilação na faixa de rejeitção
- Largura da faixa de transição
- etc.

## Projeto de filtros IIR

- Infinite impulse response (IIR)
- Filtros com retro-alimentação
- Pode ser mais eficiente que filtros com resposta finita ao impulso (FIR)
- Podem ser instáveis (por quê?); cuidado no projeto através de *análise*

Passos

1. Definir requerimentos do filtro: frequência de corte, ripples, etc
2. Obter protótipo analógico do filtro;
3. Determinar coeficientes do filtro IIR a partir do protótipo analógico:

- Invariância do impulso
- Transformação bilinear

## Protótipo de filtros analógicos

Tipos de filtros analógicos:

- Butterworth
- Chebyshev tipo I
- Chebyshev tipo II
- Elíptico

Para cada tipo de filtro, há um procedimento para a obtenção dos polos e zeros

Isso permite construir a função de transferência do filtro analógico protótipo

Em seguida, é preciso obter os coeficientes dos filtros **digitais**

```python
from scipy import signal

w = np.linspace(0, 1e3, 1000)

ordem = 10
ripple = 50 # dB
ripple_el_p = 5 # dB
ripple_el_s = 30 # dB
freqcorte = 100

# Butterworth
b_butter, a_butter = signal.butter(ordem, freqcorte, 'low', analog=True)
w_butter, h_butter = signal.freqs(b_butter, a_butter, worN=w)

# Chebyshev tipo I
b_ch1, a_ch1 = signal.cheby1(ordem, ripple, freqcorte, 'low', analog=True)
w_ch1, h_ch1 = signal.freqs(b_ch1, a_ch1, worN=w)

# Chebyshev tipo II
b_ch2, a_ch2 = signal.cheby2(ordem, ripple, freqcorte, 'low', analog=True)
w_ch2, h_ch2 = signal.freqs(b_ch2, a_ch2, worN=w)

b_el, a_el = signal.ellip(ordem, ripple_el_p, ripple_el_s, freqcorte, 'low', analog=True)
w_el, h_el = signal.freqs(b_el, a_el, worN=w)

plt.figure()
plt.semilogx(w, 20 * np.log10(abs(h_butter)))
plt.grid(which='both', axis='both')
plt.xlabel('Frequency [radians / second]')
plt.ylabel('Amplitude [dB]')
plt.legend(['Butterworth'])

plt.figure()
plt.grid(which='both', axis='both')
plt.semilogx(w, 20 * np.log10(abs(h_ch1)))
plt.semilogx(w, 20 * np.log10(abs(h_ch2)))
plt.legend(['Chebyshev tipo I', 'Chebyshev tipo II'])

plt.figure()
plt.grid(which='both', axis='both')
plt.semilogx(w, 20 * np.log10(abs(h_el)))
plt.xlabel('Frequency [radians / second]')
plt.ylabel('Amplitude [dB]')
plt.legend(['Elliptical'])
```

## Projeto de filtros IIR pela técnica da invariância do impulso

### Princípio###

A resposta ao impulso do protótipo é amostrada

$$
\begin{equation}
    h(n) = T_d h_c(nT_d)
\end{equation}
$$
em que $T_d$ é o período de amostragem

### Efeito no domínio da frequência

A resposta em frequência do filtro discreto é dada por:

$$
\begin{equation}
    H(e^{j\omega}) = \sum_{k=-\infty}^{\infty} H_c \left( j\frac{\omega}{T_d} + j\frac{2\pi}{T_d}k \right)
\end{equation}
$$

Se o filtro protótipo for __limitado em banda__, ou seja, $ H_c(j\Omega) = 0$ para $|\Omega| \geq \pi/Td$, então

$$
\begin{equation}
    H(e^{j\omega}) = H_c \left( j\frac{\omega}{T_d} \right), \quad |\omega| \leq \pi
\end{equation}
$$

Quer dizer, há um escalonamento linear $\omega = \Omega T_d$ no eixo das frequências para $|\omega| \leq \pi$ entre o filtro digital e seu protótipo analógico

Na prática, filtros não são limitados em banda, gerando *aliasing*. Porém, se as componentes de alta frequência forem fracas, então o *aliasing* torna-se desprezível.

### Obtenção da função de transferência discreta

Primeiro, expande-se a função de transferência do protótipo em frações parciais:

$$
\begin{equation}
    H_c(s) = \sum_{k=1}^{N} \frac{A_k}{s-s_k}
\end{equation}
$$

Aplicando a transformada de Laplace inversa, obtemos:

$$
\begin{equation}
    h_c(t) = \sum_{k=1}^N A_k e^{s_k t} u(t)
\end{equation}
$$

Amostra-se a resposta ao impulso do protótipo:

$$
    h(n) = T_d h_c(nT_d) = \sum_{k=1}^N T_d A_k e^{s_k n T_d} u(n T_s) = \sum_{k=1}^N T_d A_k (e^{s_k T_d})^n  u(n)
$$

A transformada Z da equação acima é dada por:

$$
    H(z) = \sum_{k=1}^N \frac{T_d A_k}{1 - e^{s_k T_d} z^{-1}}
$$

Observações:

* Um polo em $s = s_k$ se transforma num polo em $z=e^{s_kT_d}$;
* Coeficientes da expansão em frações parciais são os mesmos exceto pelo fator de escala $T_d$;
* Se o protótipo for estável, ou seja, $\text{Re}(s_k) < 0$, então $|e^{s_k T_d}| < 1$ e, logo, o filtro discreto será estável também.

## Projeto de filtros IIR pela técnica da transformação bilinear

### Princípio

Mapeia-se o eixo $j_\Omega$ do plano-s em uma revolução completa no círculo de raio unitário no plano-z.

Utiliza-se a seguinte transformação:

$$ s= \frac{2}{T_d} \left( \frac{1 - z^{-1}}{1 + z^{-1}} \right) $$

O que nos fornece:

$$ H(z) = H_c \left[\frac{2}{T_d} \left( \frac{1 - z^{-1}}{1 + z^{-1}} \right) \right] $$

### Análise

Podemos escrever:

$$ z = \frac{1 + (T_d/2) s}{1 - (T_d/2)s} $$

Substituindo $s = \sigma + j\Omega$:

$$ z = \frac{1 + (T_d/2) \sigma + j(T_d/2)\Omega}{1 - (T_d/2)\sigma - j(T_d/2)\Omega} $$

Ao calcularmos $|z| = z z^*$, e fazendo-se $\sigma < 0$, mostra-se que $|z| < 1$ para todo $\Omega$

De forma equivalente, mostra-se que $|z| > 1$ com $\sigma > 0$ para todo $\Omega$.

**Conclusão:** a transformação bilinear mapeia o semiplano $\text{Re}(s) \leq 0$ no círculo $|z| \leq 1$.

### Domínio da frequência

Vamos analisar a transformação sob o círculo de raio unitário, ou seja, utilizando $z=e^{j\omega}$.

Assim, temos:

$$
    \begin{align}
        s &= \frac{2}{T_d} \left( \frac{1 - e^{-j\omega}}{1 + e^{-j\omega}} \right)\\
          &= j \frac{2}{T_d} \tan(\omega/2)
    \end{align}
$$

Isso é, $\Omega = \frac{2}{T_d}\tan(\omega/2)$ e, portanto

$$\omega = 2\text{atan}(\Omega T_d/2)$$

```python
Td = 1e-4
omegazao = np.linspace(-1e6, 1e6, 1000)
omegazin = [ 2*np.math.atan(a*Td/2.0) for a in omegazao ]

plt.figure()
plt.plot(omegazao, omegazin)
plt.axhline(y=np.pi, xmin=-1e5, linestyle='--', color='black')
plt.axhline(y=-np.pi, xmin=-1e5, xmax=1e5, linestyle='--', color='black')
plt.grid(True)
plt.xlabel('$\Omega$ [rad/s]')
plt.ylabel('$\omega$ [rad/s]')
```

Observações:

* Transformação bilinear evita aliasing
* Compressão não-linear, distorcendo componentes de alta-frequência

## Conclusão

1. Especificações do filtro;
2. Função de transferência do protótipo $H_c(s)$
3. Função de transferência do filtro digital $H(z)$

## Trabalho computacional -- Projeto de um filtro digital

### Parte I

Determine os parâmetros do filtro Butterworth para as seguintes especificações do filtro digital:

* $0.99 \leq |H(e^{j\omega})| \leq 1,0$ para $|\omega| \leq 0.4\pi$
* $|H(e^{j\omega})| \leq 0.001$ para $0.6\pi \leq |\omega| \leq \pi$

Obtenha o filtro digital, utilizando as técnicas de invariância do impulso e transformação bilinear. Considere o filtro de Butterworth como protótipo analógico.

### Parte II

Determine o filtro RL para as características do filtro digital acima considerando transformação bilinear. Tome como referência a tensão sobre o resistor.

```python

```


#  NMA Model Types Tutorial 3: "Why" models

In this tutorial we will explore models and techniques that can potentially explain *why* the spiking data we have observed is produced the way it is.

Tutorial objectives:

- Write code to compute formula for entropy, a measure of information
- Compute the entropy of a number of toy distributions
- Compute the entropy of spiking activity from the Steinmetz dataset

```python
#@title Video: “Why” models
from IPython.display import YouTubeVideo
video = YouTubeVideo(id='fxbBJu258oE', width=854, height=480, fs=1)
print("Video available at https://youtube.com/watch?v=" + video.id)
video

```

    Video available at https://youtube.com/watch?v=fxbBJu258oE

## Setup

```python
#@title Imports
import io

import matplotlib.pyplot as plt
import numpy as np
import requests
import scipy.stats as stats

import ipywidgets as widgets

fig_w, fig_h = (6, 4)
plt.rcParams.update({'figure.figsize': (fig_w, fig_h)})

%matplotlib inline
```

```python
#@title Helper Functions
def histogram(counts, bins, vlines=(), ax=None, ax_args=None, **kwargs):
  """Plot a step histogram given counts over bins.

  TODO:
    * Passing the mean directly, but could estimate from the counts and bins.
  """
  if ax is None:
    _, ax = plt.subplots()

  # duplicate the first element of `counts` to match bin edges
  counts = np.insert(counts, 0, counts[0])

  ax.fill_between(bins, counts, step="pre", alpha=0.4, **kwargs)  # area shading
  ax.plot(bins, counts, drawstyle="steps", **kwargs)  # lines

  for x in vlines:
    ax.axvline(x, 0, ymax, colors='r', linestyles='dotted')  # vertical line

  if ax_args is None:
    ax_args = {}

  # heuristically set max y to leave a bit of room
  ymin, ymax = ax_args.get('ylim', [None, None])
  if ymax is None:
    ymax = np.max(counts)

  if ax_args.get('yscale', 'linear') == 'log':
    ymax *= 1.5
  else:
    ymax *= 1.1
    if ymin is None:
      ymin = 0
  ax_args['ylim'] = [ymin, ymax]

  ax.set(**ax_args)
  ax.autoscale(enable=True, axis='x', tight=True)

def eventplot(event_times, colos='k', ax=None, ax_args=None, **kwargs):
  if ax is None:
    _, ax = plt.subplots()

  p = ax.eventplot(event_times, colors='k', **kwargs)

  if ax_args is None:
    ax_args = {}

  # `eventplot` behaves differently with a single series? just blank the y-axis for now
  if len(p) == 1:
    ax.yaxis.set_visible(False)

  ax.set(**ax_args)
  ax.autoscale(enable=True, axis='x', tight=True)
```

```python
#@title Download Data
r = requests.get('https://osf.io/sy5xt/download')
if r.status_code != 200:
    print('oops')
steinmetz_spikes = np.load(io.BytesIO(r.content), allow_pickle=True)['spike_times']
```

## Optimization and Information

Neurons can only fire so often in a fixed period of time, as the act of emitting a spike consumes energy that is depleted and must eventually be replenished. To communicate effectively for downstream computation, the neuron would need to make good use of its limited spiking capability. This becomes an optimization problem:

What is the optimal way for a neuron to fire in order to maximize its ability to communicate information?

In order to explore this question, we first need to have a quantifiable measure for information. Shannon intoduced the concept of entropy to do just that, and defined it as

\begin{align}
  H_b(X) &= -\sum_{x\in X} p(x) \log_b p(x)
\end{align}

where $H$ is entropy measured in units of base $b$ and $p(x)$ is the probability of observing the event $x$ from the set of all possible events in $X$. See the Appendix for a more detailed look at how this equation was derived.

The most common base of measuring entropy is $b=2$, so we often talk about *bits* of information, though other bases are used as well e.g. when $b=e$ we call the units *nats*.

First, let's explore how entropy changes between some simple discrete probability distributions. In the rest of this tutorial we will refer to these as probability mass functions (PMF), where $p(x_i)$ equals the $i^{th}$ value in an array, and mass refers to how much of the distribution is contained at that value.

For our first PMF, we will choose one where all of the probability mass is located in the middle of the distribution.

```python
n_points = 50  # points supporting the distribution
bins = np.linspace(0, 1, n_points + 1)

pmf = np.zeros(n_points)
pmf[len(pmf) // 2] = 1.0  # middle point has all the mass
# Note: the middle point of `p` maps to the middle bin of `bins` because of shared indexing

histogram(pmf, bins, ax_args={
  'xlabel': "x",
  'ylabel': "p(x)",
  'xlim': [0, None],
})
```

If we were to draw a sample from this distribution, we know exactly what we would get every time. Distributions where all the mass is concentrated on a single event are known as *deterministic*.

How much entropy is contained in a deterministic distribution?

### Exercise: Computing Entropy

Your first exercise it to implement a method that computes the entropy of a discrete probability mass function. Remember that we are interested in bits, so be sure to use the correct log function. Also recall that $log(0)$ is undefined and will produce a `nan` (Not a Number) value if evaluated.

```python
def entropy(pmf):
  """Given a discrete distribution, return the Shannon entropy in bits.

  This is a measure of information in the distribution. For a totally
  deterministic distribution, where samples are always found in the same bin,
  then samples from the distribution give no more information and the entropy
  is 0.

  For now this assumes `pmf` arrives as a well-formed distribution (that is,
  `np.sum(pmf)==1` and `not np.any(pmf < 0)`)

  Args:
    pmf (np.ndarray): The probability mass function for a discrete distribution
      represented as an array of probabilities.
  Returns:
    h (number): The entropy of the distribution in `pmf`.

  """
  ###############################################################
  ## TODO for students: compute the entropy of the provided PMF #
  ###############################################################
  raise NotImplementedError("Student excercise: implement the equation for entropy")

# Uncomment once the entropy function is complete
# print(f"{entropy(pmf):.2f} bits")
```

We expect zero surprise from a deterministic distribution. If we had done this calculation by hand, it would simply be

$-1\log_2 1 = 0$

Note that changing the location of the peak (i.e. the point and bin on which all the mass rests) doesn't alter the entropy. The entropy is about how predictable a sample is with respect to a distribution. A single peak is deterministic regardless of which point it sits on.

```python
pmf = np.zeros(n_points)
pmf[2] = 1.0  # arbitrary point has all the mass

histogram(pmf, bins, ax_args={
  'title': f'Entropy = {entropy(pmf):.2f} bits',
  'xlabel': "x",
  'ylabel': "p(x)",
})
```

What about a distribution with mass split equally between two points?

```python
pmf = np.zeros(n_points)
pmf[len(pmf) // 3] = 0.5
pmf[2 * len(pmf) // 3] = 0.5

histogram(pmf, bins, ax_args={
  'title': f'Entropy = {entropy(pmf):.2f} bits',
  'xlabel': "x",
  'ylabel': "p(x)",
  'ylim': [0, 1],
})
```

Here, the entropy calculation is

$-(0.5 \log_2 0.5 + 0.5\log_2 0.5)=1$

There is 1 bit of entropy. This means that before we take a random sample, there is 1 bit of uncertainty about which point in the distribution the sample will fall on: it will either be the first peak or the second one.

Likewise, if we make one of the peaks taller (i.e. its point holds more of the probability mass) and the other one shorter, the entropy will decrease because of the increased certainty that the sample will fall on one point and not the other:

$-(0.2 \log_2 0.2 + 0.8\log_2 0.8)\approx 0.72$

Try changing the definition of the number and weighting of peaks, and see how the entropy varies.

If we split the probability mass among even more points, the entropy continues to increase. Let's derive the general form for $N$ points of equal mass, where $p_i=p=1/N$:

\begin{align}
 -\sum_i p_i \log_b p_i&= -\sum_i^N \frac{1}{N} \log_b \frac{1}{N}\\
                       &= -\log_b \frac{1}{N} \\
                       &= \log_b N
\end{align}
$$$$

If we have $N$ discrete points, the _uniform distribution_ (where points have equal mass) is the distribution with the highest entropy: $\log_b N$. This upper bound on entropy is useful when considering binning strategies, as any estimate of entropy over $N$ discrete points (or bins) must be in the interval $[0, \log_b N]$.

```python
pmf = np.ones(n_points) / n_points  # [1/N] * N

histogram(pmf, bins, ax_args={
  'title': f'Entropy = {entropy(pmf):.2f} bits',
  'xlabel': "x",
  'ylabel': "p(x)",
  'ylim': [0, 1],
})
```

Here, there are 50 bins and the entropy of the uniform distribution is $\log_2 50\approx 5.64$. If we construct _any_ discrete distribution $X$ over 50 bins and calculate an entropy of $H_2(X)>\log_2 50$, something must be wrong with our implementation of the discrete entropy computation.

## Information, neurons, and spikes

Recall the discussion of spike times and inter-spike intervals (ISIs) from Tutorial 1. What does the information content (or distributional entropy) of these measures say about our theory of nervous systems?

Let us consider three hypothetical neurons similar to those we explored in the Steinmetz data that all have the same mean ISIs but with different distributions:

1. Deterministic
2. Uniform
3. Exponential

How do we expect their entropies to differ?

```python
n_bins = 50
mean_isi = 0.025
bins = np.linspace(0, 0.25, n_bins + 1)
mean_idx = np.searchsorted(bins, mean_isi)

# 1. all mass concentrated on the ISI mean
pmf_single = np.zeros_like(bins[1:])
pmf_single[mean_idx] = 1.0

# 2. mass uniformly distributed about the ISI mean
pmf_uniform = np.zeros_like(bins[1:])
pmf_uniform[0:2*mean_idx] = 1 / (2 * mean_idx)

# 3. mass exponentially distributed about the ISI mean
pmf_exp = stats.expon.pdf(bins[1:], scale=mean_isi)
pmf_exp /= np.sum(pmf_exp)

fig, (ax1, ax2, ax3) = plt.subplots(ncols=3, figsize=(18, 5))
histogram(pmf_single, bins, ax=ax1, ax_args={
  'title': '(1) Deterministic',
  'xlabel': "Inter-spike interval (s)",
  'ylabel': "Probability mass",
  'ylim': [0, 1.0]
})
histogram(pmf_uniform, bins, ax=ax2, ax_args={
  'title': '(2) Uniform',
  'xlabel': "Inter-spike interval (s)",
  'ylabel': "Probability mass",
  'ylim': [0, 0.25]
})
histogram(pmf_exp, bins, ax=ax3, ax_args={
  'title': '(3) Exponential',
  'xlabel': "Inter-spike interval (s)",
  'ylabel': "Probability mass",
  'ylim': [0, 0.25]
})
```

```python
print(
  f"Deterministic: {entropy(pmf_single):.2f} bits",
  f"Uniform: {entropy(pmf_uniform):.2f} bits",
  f"Exponential: {entropy(pmf_exp):.2f} bits",
  sep="\n",
)
```

    Deterministic: 0.00 bits
    Uniform: 3.32 bits
    Exponential: 3.77 bits

In the previous example we created the PMFs by hand to illustrate idealized scenarios. How would we compute them from data recorded from actual neurons?

One way is to convert the ISI histograms we've previously computed and convert them into discrete probability distributions using the following equation:

\begin{align}
p(i) = \frac{n_i}{\sum\nolimits_{i}n_i}
\end{align}

where $p(i)$ is the probability of an ISI falling within a particular interval $i$ and $n_i$ is the count of how many ISIs were observed in that interval.

```python
#@title Video: Entropy of different distributions
from IPython.display import YouTubeVideo
video = YouTubeVideo(id='hySy-J51vcI', width=854, height=480, fs=1)
print("Video available at https://youtube.com/watch?v=" + video.id)
video
```

    Video available at https://youtube.com/watch?v=hySy-J51vcI

### Exercise: Probabilty Mass Function

Your second exercise is to implement a method that will produce a probability mass function from an array of ISI bin counts.

To verify your solution, we will compute the probability distribution of ISIs from real neural data taken from the Steinmetz dataset.

```python
neuron_idx = 283
isi = np.diff(steinmetz_spikes[neuron_idx])
bins = np.linspace(0, 0.25, n_bins + 1)
counts, _ = np.histogram(isi, bins)
```

```python
def pmf_from_counts(counts):
  """Given counts, normalize by the total to estimate probabilities."""
  ###########################################################################
  ## TODO for students: compute the probability mass function from ISI counts
  ###########################################################################
  raise NotImplementedError("Student excercise: compute the PMF from ISI counts")

# Uncomment once you complete the function
# pmf = pmf_from_counts(counts)
# histogram(pmf, bins, ax_args={
#   'title': f"Neuron {neuron_idx}",
#   'xlabel': "Inter-spike interval (s)",
#   'ylabel': "Probability mass",
# })
```

**Example output:**

```python
#@title Video: Calculating entropy of a probability distribution
from IPython.display import YouTubeVideo
video = YouTubeVideo(id='cyBu1pEGOh4', width=854, height=480, fs=1)
print("Video available at https://youtube.com/watch?v=" + video.id)
video
```

    Video available at https://youtube.com/watch?v=cyBu1pEGOh4

Now that we have the probability distribution for the actual neuron spiking activity, we can calculate its entropy.

```python
print(f"Entropy for Neuron {neuron_idx}: {entropy(pmf):.2f} bits")
```

    Entropy for Neuron 283: 3.36 bits

### Data Exploration

We can combine the above distribution plot and entropy calculation with an interactive widget to explore how the different neurons in the dataset vary in spiking activity and relative information.

```python
#@title Steinmetz Neuron Information Explorer

def _pmf_from_counts(counts):
  """Given counts, normalize by the total to estimate probabilities."""
  pmf = counts / np.sum(counts)
  return pmf

def _entropy(pmf):
  """Given a discrete distribution, return the Shannon entropy in bits."""
  # remove non-zero entries to avoid an error from log2(0)
  pmf = pmf[pmf > 0]
  h = -np.sum(pmf * np.log2(pmf))
  # absolute value applied to avoid getting a -0 result
  return np.abs(h)

@widgets.interact(neuron=widgets.IntSlider(0, min=0, max=(len(steinmetz_spikes)-1)))
def steinmetz_pmf(neuron):
  """ Given a neuron from the Steinmetz data, compute its PMF and entropy """
  isi = np.diff(steinmetz_spikes[neuron])
  bins = np.linspace(0, 0.25, n_bins + 1)
  counts, _ = np.histogram(isi, bins)
  pmf = _pmf_from_counts(counts)

  ymax = max(0.2, np.max(pmf))
  histogram(pmf, bins, ax_args={
    'title': f"Neuron {neuron}: H = {_entropy(pmf):.2f} bits",
    'xlabel': "Inter-spike interval (s)",
    'ylabel': "Probability mass",
    'ylim': [0, ymax]
  })
```

```python
#@title Video: Summary of “why” models
from IPython.display import YouTubeVideo
video = YouTubeVideo(id='-eGFd7E_smA', width=854, height=480, fs=1)
print("Video available at https://youtube.com/watch?v=" + video.id)
video
```

    Video available at https://youtube.com/watch?v=-eGFd7E_smA

## Appendix

### The foundations for Entropy

**Shannon** [ref] began with three criteria for a function $H$ defining the entropy of a discrete distribution of probability masses $p_i\in p(X)$ over the points $x_i\in X$:
1. $H$ should be continuous in the $p_i$.
  - That is, $H$ should change smoothly in response to smooth changes to the mass $p_i$ on each point $x_i$.
2. If all the points have equal shares of the probability mass, $p_i=1/N$, $H$ should be a non-decreasing function of $N$.
  - That is, if $X_N$ is the support with $N$ discrete points and $p(x\in X_N)$ assigns constant mass to each point, then $H(X_1) < H(X_2) < H(X_3) < \dots$
3. $H$ should be preserved by (invariant to) the equivalent (de)composition of distributions.
  - For example [from Shannon's paper] if we have a discrete distribution over three points with masses $(\frac{1}{2},\frac{1}{3},\frac{1}{6})$, then their entropy can be represented in terms of a direct choice between the three and calculated $H(\frac{1}{2},\frac{1}{3},\frac{1}{6})$. However, it could also be represented in terms of a series of two choices:
    1. either we sample the point with mass $1/2$ or not (_not_ is the other $1/2$, whose subdivisions are not given in the first choice),
    2. if (with probability $1/2$) we _don't_ sample the first point, we sample one of the two remaining points, masses $1/3$ and $1/6$.

    Thus in this case we require that $H(\frac{1}{2},\frac{1}{3},\frac{1}{6})=H(\frac{1}{2},\frac{1}{2}) + \frac{1}{2}H(\frac{1}{3}, \frac{1}{6})$

There is a unique function which satisfies these 3 requirements: [the function is unique up to scaling by a linear factor]

\begin{align}
  H_b(X) &= -\sum_{x\in X} p(x) \log_b p(x)
\end{align}

Where the base of the logarithm $b>1$ controls the units of entropy. The two most common cases are $b=2$ for units of _bits_, and $b=e$ for _nats_.

We can view this function as the expectation of the self-information over a distribution:

$$H_b(X) = \mathbb{E}_{x\in X} \left[I_b(x)\right]$$

$$I_b(x)=-\log_b p(x)$$

Self-information is just the negative logarithm of probability, and is a measure of how surprising an event sampled from the distribution would be. Events with $p(x)=1$ are certain to occur, and their self-information is zero [as is the entropy of the distribution they compose] meaning they are totally unsurprising. The smaller the probability of an event, the higher its self-information, and the more surprising the event would be to observe.


## Problem Statement

You will have to use PINNs to solve the fluid flow for the given geometry and flow parameters.

A 2D chip is placed inside a 2D channel. The flow enters inlet at $u=1.5\text{ m/s}$ and exits through the outlet which is a $0 Pa$. All the other walls are treated as no-slip. The kinematic viscosity $(\nu)$ for the flow is $0.02 \text{ }m^2/s$ and the density $(\rho)$ is $1 \text{ }kg/m^3$. The problem is shown in the figure below.

## Challenge

The main challenge in this problem is to correctly formulate the problem using PINNs. In order to achieve that, you will have to complete the following parts successfully:
1. Define the correct geometry for the problem
2. Set-up the correct boundary conditions and equations
3. Create the neural network and solve the problem

A successful completion of the problem should result in distribution of flow variables as shown below. Also, you should aim to achieve a relative $L_2$ error of less than 0.2 for all the variables w.r.t the given OpenFOAM solution.

In this template, we will have give you a skeleton code which you will fill in to define and solve the problem.

**Note: You need to edit the `chip_2d_template.py` script that is placed in the ../source_code/chip_2d/ directory.**

From the top menu, click on File, and Open `chip_2d_template.py` from the current directory at `../source_code/chip_2d` directory. Remember to SAVE your code after changes, before running below cells.

Let us start with importing the required packages. Pay attention to the various modules and packages that are being imported, especially the equations, geometry and architectures.

```python
from sympy import Symbol
import numpy as np
import tensorflow as tf
from simnet.solver import Solver
from simnet.dataset import TrainDomain, ValidationDomain
from simnet.data import Validation
from simnet.sympy_utils.geometry_2d import Rectangle, Line, Channel2D
from simnet.sympy_utils.functions import parabola
from simnet.csv_utils.csv_rw import csv_to_dict
from simnet.PDES.navier_stokes import IntegralContinuity, NavierStokes
from simnet.controller import SimNetController
from simnet.architecture import FourierNetArch
```

Now, define the simulation parameters and generate the geometry. You will be using the 2D geometry modules for this example. Please fill in the appropriate values for each geometry. Remember, `Channel2D` and `Rectangle` are defined by its two endpoints. The difference between a channel and rectangle in SimNet is that, the channel geometry does not have bounding curves in the x-direction. This is helpful in getting a signed distance field that it infinite in x-direction. This can be important when the signed distance field is used as a wall distance in some of the turbulence models (Refer *SimNet User Guide Chapter 3* for more details). Hence we will create the inlet and outlet using `Line` geometry (*Please note that this is a 2d line, as opposed to the* `Line1D` *that was used in the diffusion tutorial*)

```python
# simulation params
channel_length = (-2.5, 2.5)
channel_width = (-0.5, 0.5)
chip_pos = -1.0
chip_height = 0.6
chip_width = 1.0
inlet_vel = 1.5
```

```python
#TODO: Replace x1, y1, x2, y2, and X's with appropriate values

# define geometry
# define channel
channel = Channel2D((x1, y1), (x2, y2))
# define inlet and outlet
inlet = Line((x1, y1), (x1, y2), normal= X)
outlet = Line((x1, y1), (x1, y2), normal= X)
# define the chip
rec = Rectangle((x1, y1), (x2, y2))
# create a geometry for higher sampling of point cloud near the fin
flow_rec = Rectangle((chip_pos-0.25, channel_width[0]),
                     (chip_pos+chip_width+0.25, channel_width[1]))
# fluid area
geo = channel - rec
geo_hr = geo & flow_rec
geo_lr = geo - flow_rec
```

The current problem is a channel flow with an incompressible fluid. In such cases, the mass flow rate through each cross-section of the channel and in turn the volumetric flow is constant. This can be used as an additional constraint in the problem which we will help us in improving the speed of convergence.

Wherever, possible, using such additional knowledge about the problem can help in better and faster solutions. More examples of this can be found in the *SimNet User Guide Chapter 6*.

```python
# Optional integral continuity planes to speed up convergence
x_pos = Symbol('x_pos')
integral_line = Line((x_pos, channel_width[0]),
                     (x_pos, channel_width[1]),
                     1)
x_pos_range = {x_pos: lambda batch_size: np.full((batch_size, 1), np.random.uniform(channel_length[0], channel_length[1]))}
```

Now you will use the created geometry to define the training data for the problem. Implement the required boundary conditions and equations in the `Chip2DTrain` class below. Remember that you will have to create the training data points both for the boundary condition and to reduce the equation residuals. You can refer to the `NavierStokes` class from the PDEs module to check how the equations are defined and the keys required for each equation. For ease of access, we show the equations below.

```python
# Equation definitions of the NavierStokes class

# Note, this block is only for reference. Do not include this in your final script
# These equations are already present in the NavierStokes class definition

    # set equations
    self.equations = Variables()
    self.equations['continuity'] = rho.diff(t) + (rho*u).diff(x) + (rho*v).diff(y) + (rho*w).diff(z)
    self.equations['momentum_x'] = ((rho*u).diff(t)
                                  + (u*((rho*u).diff(x)) + v*((rho*u).diff(y)) + w*((rho*u).diff(z)) + rho*u*(curl))
                                  + p.diff(x)
                                  - (-2/3*mu*(curl)).diff(x)
                                  - (mu*u.diff(x)).diff(x)
                                  - (mu*u.diff(y)).diff(y)
                                  - (mu*u.diff(z)).diff(z)
                                  - (mu*(curl).diff(x)))
    self.equations['momentum_y'] = ((rho*v).diff(t)
                                  + (u*((rho*v).diff(x)) + v*((rho*v).diff(y)) + w*((rho*v).diff(z)) + rho*v*(curl))
                                  + p.diff(y)
                                  - (-2/3*mu*(curl)).diff(y)
                                  - (mu*v.diff(x)).diff(x)
                                  - (mu*v.diff(y)).diff(y)
                                  - (mu*v.diff(z)).diff(z)
                                  - (mu*(curl).diff(y)))
    self.equations['momentum_z'] = ((rho*w).diff(t)
                                  + (u*((rho*w).diff(x)) + v*((rho*w).diff(y)) + w*((rho*w).diff(z)) + rho*w*(curl))
                                  + p.diff(z)
                                  - (-2/3*mu*(curl)).diff(z)
                                  - (mu*w.diff(x)).diff(x)
                                  - (mu*w.diff(y)).diff(y)
                                  - (mu*w.diff(z)).diff(z)
                                  - (mu*(curl).diff(z)))

    if self.dim == 2:
      self.equations.pop('momentum_z')
```

Now use this understanding to define the problem. An example of the inlet boundary condition is shown. Also, the integral continuity constraint is already coded up for you.

```python
#TODO: Replace all the placeholders with appropriate values

# define sympy variables to parametrize domain curves
x, y = Symbol('x'), Symbol('y')

class Chip2DTrain(TrainDomain):
  def __init__(self, **config):
    super(Chip2DTrain, self).__init__()

    # inlet
    inlet_parabola = parabola(y, channel_width[0], channel_width[1], inlet_vel)
    inlet_bc = inlet.boundary_bc(outvar_sympy={'u': inlet_parabola, 'v': 0},
                                 batch_size_per_area=64)
    self.add(inlet_bc, name="Inlet")

    # outlet
    outlet_bc = outlet.boundary_bc(outvar_sympy={placeholder},
                                   batch_size_per_area=placeholder)
    self.add(outlet_bc, name="Outlet")

    # noslip
    noslip = geo.boundary_bc(outvar_sympy={placeholder},
                             batch_size_per_area=placeholder)
    self.add(noslip, name="ChipNS")

    # interior lr
    interior_lr = geo_lr.interior_bc(outvar_sympy={placeholder},
                                     bounds={placeholder},
                                     lambda_sympy={placeholder},
                                     batch_size_per_area=placeholder)
    self.add(interior_lr, name="InteriorLR")

    # interior hr
    interior_hr = geo_hr.interior_bc(outvar_sympy=placeholder,
                                     bounds=placeholder,
                                     lambda_sympy=placeholder,
                                     batch_size_per_area=placeholder)
    self.add(interior_hr, name="InteriorHR")

    # integral continuity
    for i in range(4):
      IC = integral_line.boundary_bc(outvar_sympy={'integral_continuity': 1.0},
                                     batch_size_per_area=512,
                                     lambda_sympy={'lambda_integral_continuity': 1.0},
                                     criteria=geo.sdf>0,
                                     param_ranges=x_pos_range,
                                     fixed_var=False)
      self.add(IC, name="IntegralContinuity_"+str(i))
```

Now, add validation data to the problem. The examples folder contains a `openfoam` directory that contains the solution of same problem using OpenFOAM solver. The CSV file is read in and converted to a dictionary for you. Now, you will have to complete the definition of `Chip2DVal` class below.

```python
#TODO: Set the appropriate normalization for the validation data
# The validation data has domain extents of (0,0) to (5,1). Normalize this based on your definition of the domain

# validation data
mapping = {'Points:0': 'x', 'Points:1': 'y',
           'U:0': 'u', 'U:1': 'v', 'p': 'p'}
openfoam_var = csv_to_dict('openfoam/2D_chip_fluid0.csv', mapping)
openfoam_var['x'] -= 2.5 #TODO Samle normalization of position. Edit based on your geometry definition
openfoam_var['y'] -= 0.5
openfoam_invar_numpy = {key: value for key, value in openfoam_var.items() if key in ['x', 'y']}
openfoam_outvar_numpy = {key: value for key, value in openfoam_var.items() if key in ['u', 'v', 'p']}

class Chip2DVal(ValidationDomain):
  def __init__(self, **config):
    super(Chip2DVal, self).__init__()
    val = Validation.from_numpy(placeholder)
    self.add(val, name='Val')
```

Now, complete the last part of the code by creating the `ChipSolver` to solve the problem. You will make use of an advanced architecture called the `FourierNetArch` in this problem. This architecture will help to reach the convergence faster. More details about this architecture can be found in *SimNet User Guide Section 1.6.1.1*. The important parameters of the neural network are defined for you. Feel free to tweak them and observe its behavior on the results and speed of convergence.

```python
#TODO: Replace all the placeholders with appropriate values
class ChipSolver(Solver):
  train_domain = placeholder
  val_domain = placeholder
  arch = FourierNetArch

  def __init__(self, **config):
    super(ChipSolver, self).__init__(**config)

    self.frequencies = ('axis,diagonal', [i/5. for i in range(25)])

    self.equations = (placeholder)
    flow_net = self.arch.make_node(name='flow_net',
                                   inputs=[placeholder],
                                   outputs=[placeholder])
    self.nets = [flow_net]

  @classmethod
  def update_defaults(cls, defaults):
    defaults.update({
        'network_dir': './network_checkpoint_chip_2d',
        'rec_results': True,
        'rec_results_freq': 5000,
        'max_steps': 10000,
        'decay_steps': 100,
        'xla': True
        })
if __name__ == '__main__':
  ctr = SimNetController(ChipSolver)
  ctr.run()
```

Once you have completed all the sections of the code, you can export the notebook as a python script and then execute it as we have seen in the earlier tutorials. Record your relative L2 errors with respect to the OpenFOAM solution and try to achieve errors for all the variables lower than 0.2. Also try to visualize your results using contour plots by reading in the `.npz` files created in the network checkpoint.

# Licensing
This material is released by NVIDIA Corporation under the Creative Commons Attribution 4.0 International (CC BY 4.0)


```python
import sys
sys.path.append('../../')
import os
import dill
import numpy as np
import scipy as sc
import random as rand
from sklearn import preprocessing, linear_model
import matplotlib.pyplot as plt

from core.controllers import PDController
from core.dynamics import LinearSystemDynamics, ConfigurationDynamics

from koopman_core.controllers import OpenLoopController, MPCController,BilinearFBLinController, PerturbedController, LinearLiftedController
from koopman_core.dynamics import LinearLiftedDynamics, BilinearLiftedDynamics
from koopman_core.learning import Edmd, BilinearEdmd, KoopDnn
from koopman_core.basis_functions import PlanarQuadBasis
from koopman_core.learning.utils import differentiate_vec
from core.systems import CartPole

class CartPolePdOutput(ConfigurationDynamics):
    def __init__(self, dynamics, xd, t_d, n, m):
        ConfigurationDynamics.__init__(self, dynamics, 1)
        self.xd = xd
        self.t_d = t_d
        self.n = n
        self.m = m

    def proportional(self, x, t):
        q, q_dot = x[:int(n/2)], x[int(n/2):]
        return self.y(q) - self.y_d(t)

    def derivative(self, x, t):
        q, q_dot = x[:int(n/2)], x[int(n/2):]
        return self.dydq(q)@q_dot - self.y_d_dot(t)

    def y(self, q):
        return q

    def dydq(self, q):
        return np.eye(int(self.n/2))

    def d2ydq2(self, q):
        return np.zeros((int(self.n/2), int(self.n/2), int(self.n/2)))

    def y_d(self, t):
        return self.desired_state_(t)[:int(self.n/2)]

    def y_d_dot(self, t):
        return self.desired_state_(t)[int(self.n/2):]

    def desired_state_(self, t):
        return [np.interp(t, self.t_d.flatten(),self.xd[:,ii].flatten()) for ii in range(self.xd.shape[1])]

class CartPoleDiscrete(CartPole):
    def __init__(self, m_c, m_p, l, g=9.81, dt=1e-2):
        CartPole.__init__(self, m_c, m_p, l, g=g)
        self.dt=dt

    def eval_dot(self, x, u, t):
        return x + self.dt*self.drift(x, t) + self.dt*np.dot(self.act(x, t),u)

    def get_linearization(self, q0, q1, u0, t):
        m_c, m_p, l, g = self.params
        x, th, x_dot, th_dot = q0
        F = u0[0]
        denom = m_p*np.sin(th)**2 + m_c
        A_lin = np.eye(self.n) + self.dt*np.array([[0, 0, 1, 0],
                                                   [0, 0, 0, 1],
                                                   [0,
                                                    (l**3*m_p**2*th_dot**2*np.cos(th) - g*l**2*m_p**2*np.cos(th)**2 + g*l**2*m_p**2*np.sin(th)**2)/(l**2*m_p*denom) - (2*np.cos(th)*np.sin(th)*(l**2*m_p*(l*m_p*np.sin(th)*th_dot**2 + F) - g*l**2*m_p**2*np.cos(th)*np.sin(th)))/(l**2*denom**2),
                                                    0,
                                                    2*l*m_p*th_dot*np.sin(th)/denom],
                                                   [ 0,
                                                    (- l*m_p*th_dot**2*np.cos(th)**2 + l*m_p*th_dot**2*np.sin(th)**2 + g*(m_c + m_p)*np.cos(th) + F*np.sin(th))/(l*(m_p*np.sin(th)**2 + m_c)) + (2*m_p*np.cos(th)*np.sin(th)*(l*m_p*np.cos(th)*np.sin(th)*th_dot**2 + F*np.cos(th) - g*np.sin(th)*(m_c + m_p)))/(l*denom**2),
                                                    0,
                                                    -(2*m_p*th_dot*np.cos(th)*np.sin(th))/denom]])

        B_lin = self.dt*np.array([[0],
                                  [0],
                                  [1/denom],
                                  [-np.cos(th)/(l*denom)]])

        if q1 is None:
            x1 = A_lin@q0 + B_lin@u0

        #f_d = self.eval_dot(q0, u0, t)
        #r_lin = f_d - q1
        r_lin=0

        return A_lin, B_lin, r_lin

```

## Cart pole

```python
mass_cart = .5
mass_pole = .2
length = .4
gravity = 9.81
sys_name = 'cart_pole'
system = CartPole(mass_cart, mass_pole, length, g=gravity)
dt = 1e-2
system_d = CartPoleDiscrete(mass_cart, mass_pole, length, g=gravity, dt=dt)

# Linearized system specification:
#denom = mass_cart+mass_pole + mass_pole*mass_cart*length**2
denom = mass_pole*mass_cart*length**2
n, m = 4, 1                                                       # Number of states, number of control inputs
A_nom, B_nom, _ = system_d.get_linearization(np.zeros(4), None, np.zeros(1), 0.)
A_nom = (A_nom - np.eye(n))/dt
B_nom = B_nom/dt
```

### Collect data for learning

To collect data, a nominal controller is designed with LQR on the dynamics's linearization around hover. However, any
controller can be used and the method does not require the knowledge of model's linearization. In addition, a
exploratory white noise is added to the controller to ensure that the data is sufficiently excited. Note that the system
is underactuated and that trajectory optimization is necessary to control the position of the vehicle. We use a
simplified trajectory generator based on a model predictive controller for the linearized dynamics. More careful design
of the desired trajectory may be necessary for more demanding applications and this is readily compatible with our method.

```python
q_dc, r_dc = 5e2, 1                                                 # State and actuation penalty values, data collection
Q_dc = q_dc * np.identity(n)                                        # State penalty matrix, data collection
R_dc = r_dc*np.identity(m)                                          # Actuation penalty matrix, data collection
P_dc = sc.linalg.solve_continuous_are(A_nom, B_nom, Q_dc, R_dc)     # Algebraic Ricatti equation solution, data collection
K_dc = np.linalg.inv(R_dc)@B_nom.T@P_dc                             # LQR feedback gain matrix, data collection
K_dc_p = K_dc[:,:int(n/2)]                                          # Proportional control gains, data collection
K_dc_d = K_dc[:,int(n/2):]                                          # Derivative control gains, data collection
nominal_sys = LinearLiftedDynamics(A_nom, B_nom, np.eye(n), lambda x: x)

# Data collection parameters:
collect_data = False
dt = 1.0e-2                                                         # Time step length
traj_length_dc = 2.                                                 # Trajectory length, data collection
n_pred_dc = int(traj_length_dc/dt)                                  # Number of time steps, data collection
t_train = dt * np.arange(n_pred_dc + 1)                             # Simulation time points
n_traj_train = 50                                                 # Number of trajectories to execute, data collection
n_traj_val = int(0.2*n_traj_train)
noise_var = 1.                                                      # Exploration noise to perturb controller, data collection

xmax = np.array([10, 10, 10, 10])                          # State constraints, trajectory generation
xmin = -xmax
umax = np.array([100])     # Actuation constraint, trajectory generation
umin = -umax
x0_max = np.array([1, 1, 1, 1])          # Initial value limits
Q_trajgen = sc.sparse.diags([0,0,0,0])                          # State penalty matrix, trajectory generation
QN_trajgen = sc.sparse.diags([5e1,5e1,1e1,1e1])             # Final state penalty matrix, trajectory generation
R_trajgen = sc.sparse.eye(m)                                        # Actuation penalty matrix, trajectory generation
sub_sample_rate = 1                                                 # Rate to subsample data for training
model_fname = 'examples/cart_pole_models'                         # Path to save learned models
n_cols = 10                                                         # Number of columns in training data plot
save_figures = False
directory = os.path.abspath("")                                     # Path to save learned models
#dropbox_folder = '/Users/carlaxelfolkestad/Dropbox/Apps/Overleaf/Koopman NMPC (ICRA21)/'
```

```python
if collect_data:
    xd = np.empty((n_traj_train + n_traj_val, n_pred_dc + 1, n))
    xs = np.empty((n_traj_train + n_traj_val, n_pred_dc + 1, n))
    us = np.empty((n_traj_train + n_traj_val, n_pred_dc, m))

    plt.figure(figsize=(12, 12 * (n_traj_train + n_traj_val) / (n_cols ** 2)))
    for ii in range(n_traj_train+n_traj_val):
        x0 = np.asarray([rand.uniform(l, u) for l, u in zip(-x0_max, x0_max)])
        set_pt_dc = np.asarray([rand.uniform(l, u) for l, u in zip(-x0_max, x0_max)])
        mpc_trajgen = MPCController(nominal_sys, n_pred_dc, dt, umin, umax, xmin, xmax, QN_trajgen, R_trajgen,
                                    QN_trajgen, set_pt_dc)
        mpc_trajgen.eval(x0, 0)
        xd[ii, :, :] = mpc_trajgen.parse_result().T

        while abs(x0[0]) + abs(x0[1]) < 1 or np.any(np.isnan(xd[ii, :, :])):
            x0 = np.asarray([rand.uniform(l, u) for l, u in zip(-x0_max, x0_max)])
            set_pt_dc = np.asarray([rand.uniform(l, u) for l, u in zip(-x0_max, x0_max)])
            mpc_trajgen = MPCController(nominal_sys, n_pred_dc, dt, umin, umax, xmin, xmax, QN_trajgen, R_trajgen,
                                        QN_trajgen, set_pt_dc)
            mpc_trajgen.eval(x0, 0)
            xd[ii, :, :] = mpc_trajgen.parse_result().T

        output = CartPolePdOutput(system, xd[ii, :, :], t_train, n, m)
        pd_controller = PDController(output, K_dc_p, K_dc_d)
        perturbed_pd_controller = PerturbedController(system, pd_controller, noise_var)
        xs[ii, :, :], us[ii, :, :] = system.simulate(x0, perturbed_pd_controller, t_train)

        plt.subplot(int(np.ceil((n_traj_train + n_traj_val) / n_cols)), n_cols, ii + 1)
        plt.plot(t_train, xs[ii, :, 0], 'b', label='$y$')
        plt.plot(t_train, xs[ii, :, 1], 'g', label='$z$')
        plt.plot(t_train, xs[ii, :, 2], 'r', label='$\\theta$')
        plt.plot(t_train, xd[ii, :, 0], '--b', label='$y_d$')
        plt.plot(t_train, xd[ii, :, 1], '--g', label='$z_d$')
        plt.plot(t_train, xd[ii, :, 2], '--r', label='$\\theta_d$')

    plt.suptitle(
        'Training data \nx-axis: time (sec), y-axis: state value, $x$ - blue, $xd$ - dotted blue, $\\theta$ - red, $\\theta_d$ - dotted red',
        y=0.94)
    plt.show()

    xs_train, us_train = xs[:n_traj_train,:,:], us[:n_traj_train, :, :]
    xs_val, us_val = xs[n_traj_train:,:,:], us[n_traj_train:, :, :]

    data_list = [xs_train, us_train, t_train, n_traj_train, xs_val, us_val, t_train, n_traj_val]
    outfile = open(directory + '/data/' + sys_name + '_data.pickle', 'wb')
    dill.dump(data_list, outfile)
    outfile.close()
else:
    infile = open(directory + '/data/' + sys_name + '_data.pickle', 'rb')
    xs_train, us_train, t_train, n_traj_train, xs_val, us_val, t_val, n_traj_val = dill.load(infile)
    infile.close()
```

### Learn a linear model with dynamic mode decomposition (DMD)

To compare our method with existing techniques, we first learn a linear state space model from data. This is dubbed
dynamic mode decomposition. I.e. we use linear regression with LASSO regularization to learn an approximate linear model
with model structure

\begin{equation}
    \mathbf{\dot{x}} = A_{dmd}\mathbf{x} + B_{dmd}\mathbf{u}
\end{equation}

```python
#DMD parameters:
alpha_dmd = 0#1.6e-3                                                  # Regularization strength (LASSO) DMD
tune_mdl_dmd = False
```

```python
basis = lambda x: x
C_dmd = np.eye(n)

optimizer_dmd = linear_model.MultiTaskLasso(alpha=alpha_dmd, fit_intercept=False, selection='random')
cv_dmd = linear_model.MultiTaskLassoCV(fit_intercept=False, n_jobs=-1, cv=3, selection='random')
standardizer_dmd = preprocessing.StandardScaler(with_mean=False)

model_dmd = Edmd(n, m, basis, n, n_traj_train, optimizer_dmd, cv=cv_dmd, standardizer=standardizer_dmd, C=C_dmd,
                 first_obs_const=False, continuous_mdl=False, dt=dt)
X_dmd, y_dmd = model_dmd.process(xs_train, us_train, np.tile(t_train,(n_traj_train,1)),
                                downsample_rate=sub_sample_rate)
model_dmd.fit(X_dmd, y_dmd, cv=tune_mdl_dmd, override_kinematics=True)
sys_dmd = LinearLiftedDynamics(model_dmd.A, model_dmd.B, model_dmd.C, model_dmd.basis, continuous_mdl=False, dt=dt)
if tune_mdl_dmd:
    print('$\\alpha$ DMD: ',model_dmd.cv.alpha_)
```

    /Users/carlaxelfolkestad/.conda/envs/keedmd/lib/python3.7/site-packages/sklearn/linear_model/_coordinate_descent.py:1790: UserWarning: Coordinate descent with l1_reg=0 may lead to unexpected results and is discouraged.
      check_random_state(self.random_state), random)
    /Users/carlaxelfolkestad/.conda/envs/keedmd/lib/python3.7/site-packages/sklearn/linear_model/_coordinate_descent.py:1790: ConvergenceWarning: Objective did not converge. You might want to increase the number of iterations. Duality gap: 12.528888988727509, tolerance: 0.03191050580215085
      check_random_state(self.random_state), random)

### Learn a lifted linear model with extended dynamic mode decomposition (EDMD)

In addition, we compare our method with the current state of the art of Koopman based learning, the extended dynamic mode
decomposition. We use a dictionary of nonlinear functions $\boldsymbol{\phi(x)}$ to lift the state variables and learn a lifted state space model
of the dynamics. I.e. we first lift and then use linear regression with LASSO regularization to learn an approximate
lifted linear model with model structure

\begin{equation}
    \mathbf{\dot{z}} = A_{edmd}\mathbf{z} + B_{edmd}\mathbf{u}, \qquad \mathbf{z} = \boldsymbol{\phi(x)}
\end{equation}

```python
#EDMD parameters:
alpha_edmd = 0#5.6e-3                                                 # Regularization strength (LASSO) EDMD
tune_mdl_edmd = False
```

```python
poly_basis = preprocessing.PolynomialFeatures(1)
poly_basis.fit(np.zeros((1,n)))
basis_edmd = lambda x: poly_basis.transform(x)
n_lift_edmd = basis_edmd(np.zeros((1,n))).shape[1]
C_edmd = np.zeros((n,n_lift_edmd))
C_edmd[:,1:n+1] = np.eye(n)

optimizer_edmd = linear_model.MultiTaskLasso(alpha=alpha_edmd, fit_intercept=False, selection='random',
                                             max_iter=2000)
cv_edmd = linear_model.MultiTaskLassoCV(fit_intercept=False, n_jobs=-1, cv=3, selection='random', max_iter=2000)
standardizer_edmd = preprocessing.StandardScaler(with_mean=False)

model_edmd = Edmd(n, m, basis_edmd, n_lift_edmd, n_traj_train, optimizer_edmd, cv=cv_edmd,
                  standardizer=standardizer_edmd, C=C_edmd, continuous_mdl=False, dt=dt)
X_edmd, y_edmd = model_edmd.process(xs_train, us_train, np.tile(t_train,(n_traj_train,1)),
                                    downsample_rate=sub_sample_rate)
model_edmd.fit(X_edmd, y_edmd, cv=tune_mdl_edmd, override_kinematics=True)
model_edmd.reduce_mdl()
sys_edmd = LinearLiftedDynamics(model_edmd.A, model_edmd.B, model_edmd.C, model_edmd.basis_reduced,
                                continuous_mdl=False, dt=dt)
if tune_mdl_edmd:
    print('$\\alpha$ EDMD: ',model_edmd.cv.alpha_)

```

    /Users/carlaxelfolkestad/.conda/envs/keedmd/lib/python3.7/site-packages/sklearn/linear_model/_coordinate_descent.py:1790: UserWarning: Coordinate descent with l1_reg=0 may lead to unexpected results and is discouraged.
      check_random_state(self.random_state), random)
    /Users/carlaxelfolkestad/.conda/envs/keedmd/lib/python3.7/site-packages/sklearn/linear_model/_coordinate_descent.py:1790: ConvergenceWarning: Objective did not converge. You might want to increase the number of iterations. Duality gap: 12.525049405737029, tolerance: 0.02663466947096737
      check_random_state(self.random_state), random)

### Learn a lifted bilinear model with bilinear extended mode decomposition (bEDMD)

Finally, we use the method developed in the paper to learn a lifted bilinear model of the dynamics, dubbed bilinear
extended mode decomposition (bEDMD). I.e. we first lift and then use linear regression with LASSO regularization to learn an approximate
lifted linear model with model structure

\begin{equation}
    \mathbf{\dot{z}}=F\mathbf{z}+\sum_{i=1}^m G_i\mathbf{z}\mathbf{u}_i, \qquad \mathbf{z} = \boldsymbol{\phi(x)}
\end{equation}

```python
#Bilinear EDMD parameters:
alpha_bedmd = 4.7e-4#1.8e-3  # Regularization strength (LASSO) bEDMD
tune_mdl_bedmd = False
```

```python
n_lift_bedmd = n_lift_edmd
C_bedmd = np.zeros((n,n_lift_bedmd))
C_bedmd[:,1:n+1] = np.eye(n)

basis_bedmd = basis_edmd
optimizer_bedmd = linear_model.MultiTaskLasso(alpha=alpha_bedmd, fit_intercept=False, selection='random')
cv_bedmd = linear_model.MultiTaskLassoCV(fit_intercept=False, n_jobs=-1, cv=3, selection='random')
standardizer_bedmd = preprocessing.StandardScaler(with_mean=False)

model_bedmd = BilinearEdmd(n, m, basis_bedmd, n_lift_bedmd, n_traj_train, optimizer_bedmd, cv=cv_bedmd,
                           standardizer=standardizer_bedmd, C=C_bedmd, continuous_mdl=False, dt=dt)
X_bedmd, y_bedmd = model_bedmd.process(xs_train, us_train, np.tile(t_train,(n_traj_train,1)),
                                       downsample_rate=sub_sample_rate)
model_bedmd.fit(X_bedmd, y_bedmd, cv=tune_mdl_bedmd, override_kinematics=False)
sys_bedmd = BilinearLiftedDynamics(model_bedmd.n_lift, m, model_bedmd.A, model_bedmd.B, model_bedmd.C,
                                   model_bedmd.basis, continuous_mdl=False, dt=dt)
if tune_mdl_bedmd:
    print('$\\alpha$ bilinear EDMD: ', model_bedmd.cv.alpha_)

```

### Learn a lifted bilinear model with Koopman NN

```python
import dill, os, torch

load_tuned_params = False

if load_tuned_params:
    infile = open(os.path.abspath('') + '/data/planar_quad_best_params.pickle', 'rb')
    net_params_lst, val_loss, test_loss, open_loop_mse, open_loop_std = dill.load(infile)
    infile.close()
    net_params = net_params_lst[-2]

    plt.figure()
    plt.plot(val_loss, label='Validation loss')
    plt.plot(test_loss, label='Test loss')
    plt.plot(open_loop_mse, label='Open loop mse')
    plt.legend()
    plt.show()

else:
    net_params = {}
    net_params['state_dim'] = n
    net_params['ctrl_dim'] = m
    net_params['encoder_hidden_width'] = 1
    net_params['encoder_hidden_depth'] = 1
    net_params['encoder_output_dim'] = 1
    net_params['optimizer'] = 'adam'
    net_params['activation_type'] = 'relu'
    net_params['lr'] = 1e-2
    net_params['epochs'] = 100
    net_params['batch_size'] = 256
    net_params['lin_loss_penalty'] = 0.
    net_params['l2_reg'] = alpha_bedmd
    net_params['l1_reg'] = 0
    net_params['first_obs_const'] = True
    net_params['override_kinematics'] = True  # TODO: Construct kinematics with normalized z...
    net_params['dt'] = dt

print(net_params)
```

    {'state_dim': 4, 'ctrl_dim': 1, 'encoder_hidden_width': 1, 'encoder_hidden_depth': 1, 'encoder_output_dim': 1, 'optimizer': 'adam', 'activation_type': 'relu', 'lr': 0.01, 'epochs': 100, 'batch_size': 256, 'lin_loss_penalty': 0.0, 'l2_reg': 0.00047, 'l1_reg': 0, 'first_obs_const': True, 'override_kinematics': True, 'dt': 0.01}

```python
from koopman_core.learning import KoopDnn, KoopmanNetCtrl
from koopman_core.util import fit_standardizer

standardizer_x_kdnn = fit_standardizer(xs_train, preprocessing.StandardScaler())
standardizer_u_kdnn = fit_standardizer(us_train, preprocessing.StandardScaler())
n_tot = n + net_params['encoder_output_dim'] + int(net_params['first_obs_const'])

net = KoopmanNetCtrl(net_params, standardizer_x=standardizer_x_kdnn, standardizer_u=standardizer_u_kdnn)
model_koop_dnn = KoopDnn(net)
model_koop_dnn.set_datasets(xs_train, t_train, u_train=us_train, x_val=xs_val, u_val=us_val, t_val=t_train)
model_koop_dnn.model_pipeline(net_params)
model_koop_dnn.construct_koopman_model()
sys_koop_dnn = BilinearLiftedDynamics(n_tot, m, model_koop_dnn.A, model_koop_dnn.B, model_koop_dnn.C, model_koop_dnn.basis_encode,
                                    continuous_mdl=False, dt=dt, standardizer_x=standardizer_x_kdnn,
                                      standardizer_u=standardizer_u_kdnn)
```

    Epoch   1: train loss: 80.68855286, validation loss: 65.86859846
    Epoch   2: train loss: 48.58369827, validation loss: 45.69984841
    Epoch   3: train loss: 31.17462540, validation loss: 30.48936164
    Epoch   4: train loss: 21.98208618, validation loss: 21.53788519
    Epoch   5: train loss: 16.98680496, validation loss: 15.80879676
    Epoch   6: train loss: 13.51797009, validation loss: 12.16691983
    Epoch   7: train loss: 11.52373695, validation loss: 10.42285490
    Epoch   8: train loss: 10.20478058, validation loss: 9.53328633
    Epoch   9: train loss: 9.27626514, validation loss: 8.96257347
    Epoch  10: train loss: 8.50618649, validation loss: 8.54730189
    Epoch  11: train loss: 7.82155466, validation loss: 8.15446168
    Epoch  12: train loss: 7.39987946, validation loss: 7.68488717
    Epoch  13: train loss: 6.77223492, validation loss: 7.32892382
    Epoch  14: train loss: 6.66629028, validation loss: 6.93586868
    Epoch  15: train loss: 6.21691036, validation loss: 6.46174043
    Epoch  16: train loss: 5.78240299, validation loss: 6.25879261
    Epoch  17: train loss: 5.66037035, validation loss: 5.99489212
    Epoch  18: train loss: 5.40725851, validation loss: 5.79376036
    Epoch  19: train loss: 5.20072174, validation loss: 5.66361129
    Epoch  20: train loss: 5.07510138, validation loss: 5.55239427
    Epoch  21: train loss: 5.12690973, validation loss: 5.50888568
    Epoch  22: train loss: 4.87418938, validation loss: 5.50315642
    Epoch  23: train loss: 4.86044407, validation loss: 5.33912179
    Epoch  24: train loss: 4.75271940, validation loss: 5.39429116
    Epoch  25: train loss: 4.73037052, validation loss: 5.41084015
    Epoch  26: train loss: 4.59696484, validation loss: 5.27784148
    Epoch  27: train loss: 4.60935497, validation loss: 5.08159077
    Epoch  28: train loss: 4.56286240, validation loss: 5.32517698
    Epoch  29: train loss: 4.86981869, validation loss: 5.29207334
    Epoch  30: train loss: 4.81983948, validation loss: 4.86354375
    Epoch  31: train loss: 4.51976681, validation loss: 5.30634946
    Epoch  32: train loss: 4.51075315, validation loss: 5.18649223
    Epoch  33: train loss: 4.81634188, validation loss: 5.20908684
    Epoch  34: train loss: 4.76180744, validation loss: 4.77165496
    Epoch  35: train loss: 4.67348051, validation loss: 5.25225675
    Epoch  36: train loss: 4.54818344, validation loss: 5.35542604
    Epoch  37: train loss: 4.56486654, validation loss: 5.27807534
    Epoch  38: train loss: 4.60901356, validation loss: 5.23050094
    Epoch  39: train loss: 4.52852440, validation loss: 5.07800835
    Epoch  40: train loss: 4.49372244, validation loss: 5.17713928
    Epoch  41: train loss: 4.56180143, validation loss: 5.36258638
    Epoch  42: train loss: 4.42932701, validation loss: 5.24117020
    Epoch  43: train loss: 4.44059706, validation loss: 5.15631345
    Epoch  44: train loss: 4.50531912, validation loss: 5.32103181
    Epoch  45: train loss: 4.80235386, validation loss: 4.99317873
    Epoch  46: train loss: 4.47430706, validation loss: 4.98011529
    Epoch  47: train loss: 4.43403721, validation loss: 5.07079938
    Epoch  48: train loss: 4.43539572, validation loss: 5.09540135
    Epoch  49: train loss: 4.43931818, validation loss: 5.25402033
    Epoch  50: train loss: 4.63382864, validation loss: 5.11100012
    Epoch  51: train loss: 4.52436066, validation loss: 5.42235619
    Epoch  52: train loss: 4.52489853, validation loss: 5.28379375
    Epoch  53: train loss: 4.46451283, validation loss: 5.25421080
    Epoch  54: train loss: 4.82569504, validation loss: 5.16878444
    Epoch  55: train loss: 4.49624729, validation loss: 4.87940678
    Epoch  56: train loss: 4.51541471, validation loss: 5.17004746
    Epoch  57: train loss: 4.43247414, validation loss: 5.11834484
    Epoch  58: train loss: 4.44207859, validation loss: 5.26790941
    Epoch  59: train loss: 4.46312857, validation loss: 5.20439234
    Epoch  60: train loss: 4.43391418, validation loss: 5.16899237
    Epoch  61: train loss: 4.68636513, validation loss: 5.18233898
    Epoch  62: train loss: 4.58862019, validation loss: 5.47781321
    Epoch  63: train loss: 4.46336317, validation loss: 5.20103630
    Epoch  64: train loss: 4.49433565, validation loss: 5.13986647
    Epoch  65: train loss: 4.42888403, validation loss: 5.26619828
    Epoch  66: train loss: 4.46454191, validation loss: 5.32365555
    Epoch  67: train loss: 4.63051891, validation loss: 5.23192072
    Epoch  68: train loss: 4.48505592, validation loss: 5.20033681
    Epoch  69: train loss: 4.48388863, validation loss: 5.29265881
    Epoch  70: train loss: 4.44096756, validation loss: 5.43157470
    Epoch  71: train loss: 4.60229206, validation loss: 5.27938345
    Epoch  72: train loss: 4.56974888, validation loss: 5.49110767
    Epoch  73: train loss: 4.43550301, validation loss: 5.14421219
    Epoch  74: train loss: 4.47516251, validation loss: 5.30382189
    Epoch  75: train loss: 4.50079823, validation loss: 5.30085221
    Epoch  76: train loss: 4.40063381, validation loss: 5.23422113
    Epoch  77: train loss: 4.41518784, validation loss: 5.29071924
    Epoch  78: train loss: 4.40364361, validation loss: 5.12983552
    Epoch  79: train loss: 4.56989908, validation loss: 5.09465232
    Epoch  80: train loss: 4.41352558, validation loss: 5.27207288
    Epoch  81: train loss: 4.41745424, validation loss: 5.35777545
    Epoch  82: train loss: 4.40817165, validation loss: 5.14946365
    Epoch  83: train loss: 4.44047165, validation loss: 5.29143304
    Epoch  84: train loss: 4.48641443, validation loss: 5.18243855
    Epoch  85: train loss: 4.48511505, validation loss: 5.06272173
    Epoch  86: train loss: 4.46955872, validation loss: 5.14922106
    Epoch  87: train loss: 4.48743057, validation loss: 5.20692608
    Epoch  88: train loss: 4.42195749, validation loss: 5.28306344
    Epoch  89: train loss: 4.46008778, validation loss: 5.20663133
    Epoch  90: train loss: 4.44343615, validation loss: 5.29376638
    Epoch  91: train loss: 4.48083973, validation loss: 5.19947386
    Epoch  92: train loss: 4.42796326, validation loss: 5.32726717
    Epoch  93: train loss: 4.46833706, validation loss: 5.24318728
    Epoch  94: train loss: 4.69247913, validation loss: 5.27703649
    Epoch  95: train loss: 4.56358051, validation loss: 5.16994485
    Epoch  96: train loss: 4.51683903, validation loss: 5.33654621
    Epoch  97: train loss: 4.48691845, validation loss: 5.17572734
    Epoch  98: train loss: 4.44535017, validation loss: 5.20803046
    Epoch  99: train loss: 4.40575218, validation loss: 5.20110521
    Epoch 100: train loss: 4.45091534, validation loss: 5.12370417
    Finished Training

### Evaluate open loop prediction performance

We first evaluate the open loop prediction performance of the proposed method.
This is done by generating a new data set in the same way as the training set, predicting the evolution of the system
with the control sequence of each trajectory executed in the data set with each of the models, and finally comparing
the mean and standard deviation of the error between the true and predicted evolution over the trajectories. The
experimental results support what is to be expected from the theory as the error in the $y$ and $z$ terms are
significantly lower for the bEDMD method than both DMD and EDMD. The reason for this
improvement is that the bEDMD method can capture the nonlinearities present in the actuation matrix of the
$(y,z)$-dynamics.

```python
# Prediction performance evaluation parameters:
folder_plots = 'figures/'                                  # Path to save plots
n_traj_ol = 50                                                     # Number of trajectories to execute, open loop
t_eval = t_train
```

```python
from koopman_core.util import evaluate_ol_pred
from tabulate import tabulate

xs_test = np.empty((n_traj_ol, t_eval.shape[0], n))
us_test = np.empty((n_traj_ol, t_eval.shape[0]-1, m))

for ii in range(n_traj_ol):
    x0 = np.asarray([rand.uniform(l, u) for l, u in zip(-x0_max, x0_max)])
    set_pt_dc = np.asarray([rand.uniform(l, u) for l, u in zip(-x0_max, x0_max)])
    mpc_trajgen = MPCController(nominal_sys, n_pred_dc, dt, umin, umax, xmin, xmax, QN_trajgen, R_trajgen,
                                QN_trajgen, set_pt_dc)
    mpc_trajgen.eval(x0, 0)
    xd = mpc_trajgen.parse_result().T

    while xd[0,0] is None:
        x0 = np.asarray([rand.uniform(l, u) for l, u in zip(-x0_max, x0_max)])
        set_pt_dc = np.asarray([rand.uniform(l, u) for l, u in zip(-x0_max, x0_max)])
        mpc_trajgen = MPCController(nominal_sys, n_pred_dc, dt, umin, umax, xmin, xmax, QN_trajgen, R_trajgen,
                                    QN_trajgen, set_pt_dc)
        mpc_trajgen.eval(x0, 0)
        xd = mpc_trajgen.parse_result().T

    output = CartPolePdOutput(system, xd, t_eval, n, m)
    pd_controller = PDController(output, K_dc_p, K_dc_d)
    perturbed_pd_controller = PerturbedController(system, pd_controller, noise_var)

    xs_test[ii,:,:], us_test[ii,:,:] = system.simulate(x0, perturbed_pd_controller, t_eval)

mdl_lst = [sys_dmd, sys_edmd, sys_bedmd, sys_koop_dnn]
#mdl_lst = [sys_bedmd, sys_koop_dnn]
u_data_lst = [us_test, us_test, us_test, us_test]
#u_data_lst = [us_test, us_test]
mdl_names = ['DMD', 'EDMD', 'bEDMD', 'Koop DNN']
#mdl_names = ['bEDMD', 'Koop DNN']
error, mse, std = [], [], []

for sys, u_data in zip(mdl_lst, u_data_lst):
    err_tmp, mse_tmp, std_tmp = evaluate_ol_pred(sys, xs_test, t_eval, us=u_data)
    error.append(err_tmp)
    mse.append(mse_tmp)
    std.append(std_tmp)

print('\nOpen loop performance statistics:')
table_data = []
for name, mse_mdl, std_mdl in zip(mdl_names, mse, std):
    table_data.append([name, "{:.5f}".format(mse_mdl), "{:.5f}".format(std_mdl)])

print(tabulate(table_data,
               headers=['Mean squared error', 'Standard deviation']))
```

    Open loop performance statistics:
                Mean squared error    Standard deviation
    --------  --------------------  --------------------
    DMD                     108959               329.196
    EDMD                    112885               334.907
    bEDMD                   296377               541.732
    Koop DNN                187130               430.219

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

figwidth = 16
lw = 2
fs = 14
y_lim_gain = 1.2
row = 2
col = n/row

#Plot open loop results:
plt.figure(figsize=(figwidth,4))
axs = [plt.subplot(row,col,jj+1) for jj in range(n)]

for ii, err in enumerate(error):
    err_mean = np.mean(err, axis=0)
    err_std = np.std(err, axis=0)

    for jj in range(n):
        axs[jj].plot(t_eval[1:], err_mean[:,jj], label=mdl_names[ii])
        axs[jj].fill_between(t_eval[1:], err_mean[:,jj]-err_std[:,jj], err_mean[:,jj]+err_std[:,jj], alpha=0.1)

for jj in range(n):
    axs[jj].set_xlabel('Time (sec)', fontsize=fs)
    axs[jj].set_ylabel('$x_'+ str(jj+1) + '$', fontsize=fs)
    axs[jj].set_ylim([-2,2])

plt.legend(frameon=False, fontsize=fs)
stitle=plt.suptitle('Open loop prediction performance of learned models', fontsize=fs+2)
matplotlib.rcParams['pdf.fonttype'] = 42
matplotlib.rcParams['ps.fonttype'] = 42
plt.savefig(folder_plots + 'koop_sys_prediction.pdf', format='pdf', dpi=2400, bbox_extra_artists=(stitle,), bbox_inches="tight")
plt.show()
```

```python
u_data_lst = [us_train, us_train]
#mdl_names = ['DMD', 'EDMD', 'bEDMD', 'Koop DNN']
mdl_names = ['bEDMD', 'Koop DNN']
error, mse, std = [], [], []

for sys, u_data in zip(mdl_lst, u_data_lst):
    err_tmp, mse_tmp, std_tmp = evaluate_ol_pred(sys, xs_train, t_eval, us=u_data)
    error.append(err_tmp)
    mse.append(mse_tmp)
    std.append(std_tmp)

print('\nOpen loop performance statistics:')
table_data = []
for name, mse_mdl, std_mdl in zip(mdl_names, mse, std):
    table_data.append([name, "{:.5f}".format(mse_mdl), "{:.5f}".format(std_mdl)])

print(tabulate(table_data,
               headers=['Mean squared error', 'Standard deviation']))
```

    Open loop performance statistics:
                Mean squared error    Standard deviation
    --------  --------------------  --------------------
    bEDMD                   136549               369.088
    Koop DNN                137764               370.829

We now study the closed loop performance of the control design.

```python
#Closed loop performance evaluation parameters:
traj_length=250
t_eval = dt * np.arange(traj_length+1)                       # Simulation time points, closed loop
Q_mpc = sc.sparse.diags([0,0,0,0,0,0])                       # State penalty matrix, trajectory generation
#QN_mpc = sc.sparse.diags([1e5,1e5,1e5,1e5,1e5,1e5])         # Final state penalty matrix, trajectory generation
QN_mpc = Q_mpc
R_mpc = sc.sparse.eye(m)                                     # Actuation penalty matrix, trajectory generation

# Design trajectory:
x0_cl = np.array([0., 0., 0., 0., 0., 0.])                   # Initial value, closed loop trajectory
set_pt_cl = np.array([1.5, 0.5, 0., 0., 0., 0.])              # Desired final value, closed loop trajectory
xmax = np.array([2, 2, np.pi/3, 1.,1.,1.])                          # State constraints, trajectory generation
xmin = -xmax
term_constraint=True

# Define initial solution for SQP algorithm:
x_init = np.linspace(x0_cl, set_pt_cl, int(traj_length)+1)
u_init = hover_thrust*np.ones((m,traj_length)).T
```

#### Design controllers for learned DMD, EDMD, and bEDMD models

```python
from koopman_core.controllers import MPCController, NonlinearMPCController, BilinearMPCController

# Define DMD-based controller:
controller_dmd = MPCController(sys_dmd, traj_length, dt, umin, umax, xmin, xmax, Q_mpc, R_mpc, QN_mpc, set_pt_cl, terminal_constraint=term_constraint, const_offset=ctrl_offset)

# Define EDMD-based controller:
controller_edmd = MPCController(sys_edmd, traj_length, dt, umin, umax, xmin, xmax, Q_mpc, R_mpc, QN_mpc, set_pt_cl, terminal_constraint=term_constraint, const_offset=ctrl_offset)

# Define bEDMD-based controller:
controller_bedmd = BilinearMPCController(sys_bedmd, traj_length, dt, umin+hover_thrust, umax+hover_thrust, xmin, xmax, Q_mpc, R_mpc, QN_mpc, set_pt_cl, terminal_constraint=term_constraint)
z0_cl = sys_bedmd.basis(x0_cl.reshape((1,-1))).squeeze()
z_init = sys_bedmd.basis(x_init)
controller_bedmd.construct_controller(z_init, u_init)

# Define KDNN-based controller:
controller_kdnn = BilinearMPCController(sys_koop_dnn, traj_length, dt, umin+hover_thrust, umax+hover_thrust, xmin, xmax, Q_mpc, R_mpc, QN_mpc, set_pt_cl,
                                        terminal_constraint=term_constraint, standardizer_x=standardizer_x_kdnn, standardizer_u=standardizer_u_kdnn)
z0_cl_kdnn = sys_koop_dnn.basis(x0_cl.reshape((1,-1))).squeeze()
z_init_kdnn = sys_koop_dnn.basis(x_init)
u_init_kdnn = standardizer_u_kdnn.transform(u_init)
controller_kdnn.construct_controller(z_init_kdnn, u_init_kdnn)
```

#### Design controller using full knowledge of nonlinear controller

```python
system_d = PlanarQuadrotorForceInputDiscrete(mass, inertia, prop_arm, g=gravity, dt=dt)
controller_nmpc = NonlinearMPCController(system_d, traj_length, dt, umin+hover_thrust, umax+hover_thrust, xmin, xmax, Q_mpc, R_mpc, QN_mpc, set_pt_cl, terminal_constraint=term_constraint)
controller_nmpc.construct_controller(x_init, u_init+hover_thrust)
```

#### Design trajectories with the contructed MPCs

```python
max_iter = 50

controller_dmd.eval(x0_cl, 0)
xr_dmd = controller_dmd.parse_result()
ur_dmd = controller_dmd.get_control_prediction() + hover_thrust

controller_edmd.eval(x0_cl, 0)
xr_edmd = sys_edmd.C@controller_edmd.parse_result()
ur_edmd = controller_edmd.get_control_prediction() + hover_thrust

controller_bedmd.solve_to_convergence(z0_cl, 0., z_init, u_init, max_iter=max_iter)
xr_bedmd = sys_bedmd.C@controller_bedmd.get_state_prediction().T
ur_bedmd = controller_bedmd.get_control_prediction().T

controller_kdnn.solve_to_convergence(z0_cl_kdnn, 0., z_init_kdnn, u_init_kdnn, max_iter=max_iter)
xr_kdnn = standardizer_x_kdnn.inverse_transform((sys_koop_dnn.C@controller_kdnn.get_state_prediction().T).T).T
ur_kdnn = controller_kdnn.get_control_prediction().T
```

```python
controller_nmpc.solve_to_convergence(x0_cl, 0., x_init, u_init + ctrl_offset.reshape(1,-1), max_iter=max_iter)
xr_nmpc = controller_nmpc.get_state_prediction().T
ur_nmpc = controller_nmpc.get_control_prediction().T
```

#### Simulate designed trajectories open loop

```python
ol_controller_dmd = OpenLoopController(system, ur_dmd.T, t_eval[:-1])
xs_dmd, us_dmd = system.simulate(x0_cl, ol_controller_dmd, t_eval)
xs_dmd, us_dmd = xs_dmd.T, us_dmd.T

ol_controller_edmd = OpenLoopController(system, ur_edmd.T, t_eval[:-1])
xs_edmd, us_edmd = system.simulate(x0_cl, ol_controller_edmd, t_eval)
xs_edmd, us_edmd = xs_edmd.T, us_edmd.T

ol_controller_bedmd = OpenLoopController(system, ur_bedmd.T, t_eval[:-1])
xs_bedmd, us_bedmd = system.simulate(x0_cl, ol_controller_bedmd, t_eval)
xs_bedmd, us_bedmd = xs_bedmd.T, us_bedmd.T

ol_controller_kdnn = OpenLoopController(system, ur_kdnn.T, t_eval[:-1])
xs_kdnn, us_kdnn = system.simulate(x0_cl, ol_controller_kdnn, t_eval)
xs_kdnn, us_kdnn = xs_kdnn.T, us_kdnn.T

ol_controller_nmpc = OpenLoopController(system, ur_nmpc.T, t_eval[:-1])
xs_nmpc, us_nmpc = system.simulate(x0_cl, ol_controller_nmpc, t_eval)
xs_nmpc, us_nmpc = xs_nmpc.T, us_nmpc.T
```

##### Compare performance

```python
import matplotlib.pyplot as plt
from matplotlib.ticker import MaxNLocator

plot_inds = [0, 1, 2, 3, 4, 5, 0, 1]
subplot_inds = [1, 2, 3, 5, 6, 7, 4, 8]
labels = ['$y$ (m)', '$z$ (m)', '$\\theta$ (rad)', '$\\dot{y}$ (m/s)','$\\dot{z}$ (m/s)', '$\\dot{\\theta}$', '$T_1$ (N)','$T_2$ (N)']
titles = ['y-coordinates', 'z-coordinates', '$\\theta$-coordinates', 'Control inputs']
colors = ['tab:blue', 'tab:orange', 'tab:brown', 'tab:green', 'tab:cyan']

plt.figure(figsize=(12,4))
#plt.suptitle('Trajectory designed with model predictive controllers\nsolid lines - designed trajectory | dashed lines - open loop simulated trajectory | black dotted lines - state/actuation bounds')
for ii in range(8):
    ind = plot_inds[ii]
    if ii < 6:
        ax = plt.subplot(2,4,subplot_inds[ii])
        plt.plot(t_eval, xr_dmd[ind,:], colors[0], label='DMD MPC')
        plt.plot(t_eval, xr_edmd[ind, :], colors[1], label='EDMD MPC')
        plt.plot(t_eval, xr_bedmd[ind, :], colors[2], label='K-MPC')
        plt.plot(t_eval, xr_kdnn[ind, :], colors[3], label='KDNN-MPC')
        plt.plot(t_eval, xr_nmpc[ind,:], colors[4], label='NMPC')

        plt.plot(t_eval, xs_dmd[ind,:], '--', color=colors[0], linewidth=1)
        plt.plot(t_eval, xs_edmd[ind, :], '--', color=colors[1], linewidth=1)
        plt.plot(t_eval, xs_bedmd[ind, :], '--', color=colors[2], linewidth=1)
        plt.plot(t_eval, xs_kdnn[ind, :], '--',  colors[3], linewidth=1)
        plt.plot(t_eval, xs_nmpc[ind,:], '--', color=colors[4], linewidth=1)

        plt.scatter(t_eval[0], x0_cl[ind], color='g')
        plt.scatter(t_eval[-1], set_pt_cl[ind], color='r')
        plt.ylabel(labels[ind])
        ax.yaxis.set_major_locator(MaxNLocator(integer=True))
        if ii >= 3:
            plt.plot([0, t_eval[-1]], [xmax[ind], xmax[ind]], ':k')
            plt.plot([0, t_eval[-1]], [xmin[ind], xmin[ind]], ':k')
            #plt.ylim(xmin[ind]-0.1,xmax[ind]+0.1)
        if subplot_inds[ii]==1:
            plt.legend(loc='upper left', frameon=False)
    elif ii < 8:
        ax = plt.subplot(2,4,subplot_inds[ii])
        plt.plot(t_eval[:-1],ur_dmd[ind,:], color=colors[0], label='DMD MPC')
        plt.plot(t_eval[:-1], ur_edmd[ind, :], color=colors[1], label='EDMD MPC')
        plt.plot(t_eval[:-1], ur_bedmd[ind, :], color=colors[2], label='K-NMPC')
        plt.plot(t_eval[:-1], ur_kdnn[ind, :], color=colors[3], label='KDNN-NMPC')
        plt.plot(t_eval[:-1],ur_nmpc[ind,:], color=colors[4], label='NMPC')
        plt.plot([0, t_eval[-1]], [umax[ind]+hover_thrust, umax[ind]+hover_thrust], ':k')
        plt.plot([0, t_eval[-1]], [umin[ind]+hover_thrust, umin[ind]+hover_thrust], ':k')
        plt.ylabel(labels[ii])
        ax.yaxis.set_major_locator(MaxNLocator(integer=True))

    if subplot_inds[ii] > 4:
        plt.xlabel('Time (sec)')
    else:
        plt.title(titles[subplot_inds[ii]-1])

if save_figures:
    matplotlib.rcParams['pdf.fonttype'] = 42
    matplotlib.rcParams['ps.fonttype'] = 42
    plt.tight_layout()
    #plt.savefig(dropbox_folder + 'planar_quad_trajectory.pdf', format='pdf', dpi=2400)
plt.show()

cost_ref_dmd = (xr_dmd[:,-1]-set_pt_cl).T@QN_mpc@(xr_dmd[:,-1]-set_pt_cl) + np.sum(np.diag(ur_dmd.T@R_mpc@ur_dmd))
cost_ref_edmd = (xr_edmd[:,-1]-set_pt_cl).T@QN_mpc@(xr_edmd[:,-1]-set_pt_cl) + np.sum(np.diag(ur_edmd.T@R_mpc@ur_edmd))
cost_ref_bedmd = (xr_bedmd[:,-1]-set_pt_cl).T@QN_mpc@(xr_bedmd[:,-1]-set_pt_cl) + np.sum(np.diag(ur_bedmd.T@R_mpc@ur_bedmd))
cost_ref_kdnn = (xr_kdnn[:,-1]-set_pt_cl).T@QN_mpc@(xr_kdnn[:,-1]-set_pt_cl) + np.sum(np.diag(ur_kdnn.T@R_mpc@ur_kdnn))
cost_ref_nmpc = (xr_nmpc[:,-1]-set_pt_cl).T@QN_mpc@(xr_nmpc[:,-1]-set_pt_cl) + np.sum(np.diag(ur_nmpc.T@R_mpc@ur_nmpc))

dist_ol_dmd = np.linalg.norm(xs_dmd[:,-1] - set_pt_cl)
dist_ol_edmd = np.linalg.norm(xs_edmd[:,-1] - set_pt_cl)
dist_ol_bedmd = np.linalg.norm(xs_bedmd[:,-1] - set_pt_cl)
dist_ol_kdnn = np.linalg.norm(xs_kdnn[:,-1] - set_pt_cl)
dist_ol_nmpc = np.linalg.norm(xs_nmpc[:,-1] - set_pt_cl)

print('Solution statistics:\n')
print(tabulate([['DMD MPC', "{:.4f}".format(cost_ref_dmd/cost_ref_nmpc), "{:.4f}".format(dist_ol_dmd), '-','-',sum(controller_dmd.comp_time)],
                ['EDMD MPC', "{:.4f}".format(cost_ref_edmd/cost_ref_nmpc), "{:.4f}".format(dist_ol_edmd),'-','-',sum(controller_edmd.comp_time)],
                ['bEDMD MPC', "{:.4f}".format(cost_ref_bedmd/cost_ref_nmpc), "{:.4f}".format(dist_ol_bedmd), len(controller_bedmd.x_iter), "{:.4f}".format(np.mean(controller_bedmd.comp_time)), sum(controller_bedmd.comp_time)],
                ['KDNN MPC', "{:.4f}".format(cost_ref_kdnn/cost_ref_nmpc), "{:.4f}".format(dist_ol_kdnn), len(controller_kdnn.x_iter), "{:.4f}".format(np.mean(controller_kdnn.comp_time)), sum(controller_kdnn.comp_time)],
                ['NMPC (benchmark)', 1, "{:.4f}".format(dist_ol_nmpc), len(controller_nmpc.x_iter), "{:.4f}".format(np.mean(controller_nmpc.comp_time)), sum(controller_nmpc.comp_time)]],
               headers=['Normalized cost,\ndesigned trajectory', 'Realized terminal,\nerror', '# of SQP\niterations','Mean comp. time\nper iteration (secs)', 'Total comp.\ntime (secs)']))
```

#### Study evolution of the solution after each iteration of the SQP-algorithm

```python
n_iter = min(len(controller_nmpc.x_iter),len(controller_bedmd.x_iter))

# Calculate cost after each iteration:
iter_cost_bedmd, iter_cost_nmpc = [], []

ol_controller_init = OpenLoopController(system, u_init, t_eval[:-1])
xs_init, _ = system.simulate(x0_cl, ol_controller_init, t_eval)
xs_init, us_init = xs_init.T, u_init.T
init_cost = (xs_init[:,-1]-set_pt_cl).T@QN_mpc@(xs_init[:,-1]-set_pt_cl) + np.sum(np.diag(us_init.T@R_mpc@us_init))
iter_cost_bedmd = [init_cost]
iter_cost_kdnn = [init_cost]
iter_cost_nmpc = [init_cost]
iter_norm_dist_bedmd = [np.linalg.norm(xs_init[:,-1]-set_pt_cl)]
iter_norm_dist_kdnn = [np.linalg.norm(xs_init[:,-1]-set_pt_cl)]
iter_norm_dist_nmpc = [np.linalg.norm(xs_init[:,-1]-set_pt_cl)]

for ii in range(len(controller_bedmd.x_iter)):
    ur_bedmd_iter = controller_bedmd.u_iter[ii].T
    ol_controller_bedmd_iter = OpenLoopController(system, ur_bedmd_iter, t_eval[:-1])
    xs_bedmd_iter, _ = system.simulate(x0_cl, ol_controller_bedmd_iter, t_eval)
    xs_bedmd_iter, us_bedmd_iter = xs_bedmd_iter.T, ur_bedmd_iter.T
    iter_cost_bedmd.append((xs_bedmd_iter[:,-1]-set_pt_cl).T@QN_mpc@(xs_bedmd_iter[:,-1]-set_pt_cl) + np.sum(np.diag(us_bedmd_iter.T@R_mpc@us_bedmd_iter)))
    iter_norm_dist_bedmd.append(np.linalg.norm(xs_bedmd_iter[:,-1]-set_pt_cl))

for ii in range(len(controller_kdnn.x_iter)):
    ur_kdnn_iter = standardizer_u_kdnn.inverse_transform(controller_kdnn.u_iter[ii].T)
    ol_controller_kdnn_iter = OpenLoopController(system, ur_kdnn_iter, t_eval[:-1])
    xs_kdnn_iter, _ = system.simulate(x0_cl, ol_controller_kdnn_iter, t_eval)
    xs_kdnn_iter, us_kdnn_iter = xs_kdnn_iter.T, ur_kdnn_iter.T
    iter_cost_kdnn.append((xs_kdnn_iter[:,-1]-set_pt_cl).T@QN_mpc@(xs_kdnn_iter[:,-1]-set_pt_cl) + np.sum(np.diag(us_kdnn_iter.T@R_mpc@us_kdnn_iter)))
    iter_norm_dist_kdnn.append(np.linalg.norm(xs_kdnn_iter[:,-1]-set_pt_cl))

for ii in range(len(controller_nmpc.x_iter)):
    ur_nmpc_iter = controller_nmpc.u_iter[ii].T
    ol_controller_nmpc_iter = OpenLoopController(system, ur_nmpc_iter, t_eval[:-1])
    xs_nmpc_iter, _ = system.simulate(x0_cl, ol_controller_nmpc_iter, t_eval)
    xs_nmpc_iter, us_nmpc_iter = xs_nmpc_iter.T, ur_nmpc_iter.T
    iter_cost_nmpc.append((xs_nmpc_iter[:,-1]-set_pt_cl).T@QN_mpc@(xs_nmpc_iter[:,-1]-set_pt_cl) + np.sum(np.diag(us_nmpc_iter.T@R_mpc@us_nmpc_iter)))
    iter_norm_dist_nmpc.append(np.linalg.norm(xs_nmpc_iter[:,-1]-set_pt_cl))
```

```python
plt.figure(figsize=(6,4))
#plt.suptitle('Control solution after each iteration of the SQP-algorithm for NMPC and K-NMPC')
plt.subplot(2,1,1)
plt.plot(np.arange(n_iter), iter_cost_bedmd[:n_iter]/iter_cost_nmpc[-1], color=colors[2], label='K-NMPC')
plt.plot(np.arange(n_iter), iter_cost_kdnn[:n_iter]/iter_cost_nmpc[-1], color=colors[3], label='KDNN-NMPC')
plt.plot(np.arange(n_iter), iter_cost_nmpc[:n_iter]/iter_cost_nmpc[-1], color=colors[4], label='NMPC')
plt.title('Control effort')
plt.ylabel('$||u||$')
plt.legend(loc='upper right', frameon=False)
plt.xlabel('SQP iteration')

plt.subplot(2,1,2)
plt.plot(np.arange(n_iter), iter_norm_dist_bedmd[:n_iter], color=colors[2], label=labels[2])
plt.plot(np.arange(n_iter), iter_norm_dist_kdnn[:n_iter], color=colors[3], label=labels[3])
plt.plot(np.arange(n_iter), iter_norm_dist_nmpc[:n_iter], color=colors[4], label=labels[4])
plt.ylim(0,5)
plt.title('Realized terminal distance from setpoint')
plt.ylabel('$||x_N - x_d||$')
plt.xlabel('SQP iteration')

if save_figures:
    matplotlib.rcParams['pdf.fonttype'] = 42
    matplotlib.rcParams['ps.fonttype'] = 42
    plt.tight_layout()
    #plt.savefig(dropbox_folder + 'planar_quad_sqp_iterations.pdf', format='pdf', dpi=2400)

plt.show()

print('Solution statistics\n')
print(tabulate([['Nonlinear MPC', len(controller_nmpc.x_iter), np.mean(controller_nmpc.comp_time), np.std(controller_nmpc.comp_time), sum(controller_nmpc.comp_time)],
                ['Koopman bilinear MPC', len(controller_bedmd.x_iter), np.mean(controller_bedmd.comp_time), np.std(controller_bedmd.comp_time), sum(controller_bedmd.comp_time)],
                ['Koopman DNN bilinear MPC', len(controller_kdnn.x_iter), np.mean(controller_kdnn.comp_time), np.std(controller_kdnn.comp_time), sum(controller_kdnn.comp_time)]],
               headers=['Number of SQP\niterations','Mean comp. time per\niteration (secs)', 'Std comp. time per\niteration (secs)', 'Total comp.\ntime (secs)']))
```

# Evaluate performance of controllers for closed-loop control

#### Design finite horizon controllers

```python
from koopman_core.controllers import PerturbedController

Q_mpc_cl = 5e2*np.diag([1, 1, 1, 2e-1, 2e-1, 2e-1])
QN_mpc_cl = Q_mpc_cl
R_mpc_cl = np.eye(m)
traj_duration = 1.
N_cl = int(traj_duration/dt)
t_eval_cl=np.arange(300)*dt

# Solver settings closed loop:
polish_osqp = False
max_iter_osqp = 10
linsys_solver_osqp = 'qdldl'
warm_start_osqp = True

controller_dmd_cl = MPCController(sys_dmd, N_cl, dt, umin, umax, xmin, xmax, Q_mpc_cl, R_mpc_cl, QN_mpc_cl, set_pt_cl, add_slack=True)
controller_dmd_cl = PerturbedController(sys_dmd,controller_dmd_cl,0.,const_offset=hover_thrust, umin=umin, umax=umax)

controller_edmd_cl = MPCController(sys_edmd, N_cl, dt, umin, umax, xmin, xmax, Q_mpc_cl, R_mpc_cl, QN_mpc_cl, set_pt_cl, add_slack=True)
controller_edmd_cl = PerturbedController(sys_edmd,controller_edmd_cl,0.,const_offset=hover_thrust, umin=umin, umax=umax)

controller_bedmd_cl = BilinearMPCController(sys_bedmd, N_cl, dt, umin+hover_thrust, umax+hover_thrust, xmin, xmax, Q_mpc_cl, R_mpc_cl, QN_mpc_cl, set_pt_cl, add_slack=True)
controller_bedmd_cl.construct_controller(controller_bedmd.cur_z[:N_cl+1,:], controller_bedmd.cur_u[:N_cl,:])
controller_bedmd_cl.solve_to_convergence(z0_cl, 0., controller_bedmd.cur_z[:N_cl+1,:], controller_bedmd.cur_u[:N_cl,:], max_iter=max_iter)
#controller_bedmd_cl = PerturbedController(sys_bedmd,controller_bedmd_cl,0.,const_offset=hover_thrust, umin=umin, umax=umax)
controller_bedmd_cl.update_solver_settings(polish=polish_osqp, max_iter=max_iter_osqp, linsys_solver=linsys_solver_osqp, warm_start=warm_start_osqp)

controller_kdnn_cl = BilinearMPCController(sys_koop_dnn, N_cl, dt, umin+hover_thrust, umax+hover_thrust, xmin, xmax, Q_mpc_cl, R_mpc_cl, QN_mpc_cl, set_pt_cl,
                                        add_slack=True, standardizer_x=standardizer_x_kdnn, standardizer_u=standardizer_u_kdnn)
controller_kdnn_cl.construct_controller(controller_kdnn.cur_z[:N_cl+1,:], controller_kdnn.cur_u[:N_cl,:])
controller_kdnn_cl.solve_to_convergence(z0_cl_kdnn, 0., controller_kdnn.cur_z[:N_cl+1,:], controller_kdnn.cur_u[:N_cl,:], max_iter=max_iter)
controller_kdnn_cl.update_solver_settings(polish=polish_osqp, max_iter=max_iter_osqp, linsys_solver=linsys_solver_osqp, warm_start=warm_start_osqp)

controller_nmpc_cl = NonlinearMPCController(system_d, N_cl, dt, umin+hover_thrust, umax+hover_thrust, xmin, xmax, Q_mpc_cl, R_mpc_cl, QN_mpc_cl, set_pt_cl, add_slack=True)
controller_nmpc_cl.construct_controller(controller_nmpc.cur_z[:N_cl+1,:], controller_nmpc.cur_u[:N_cl,:])
controller_nmpc_cl.solve_to_convergence(x0_cl, 0., controller_nmpc.cur_z[:N_cl+1,:], controller_nmpc.cur_u[:N_cl,:], max_iter=max_iter)
controller_nmpc_cl.update_solver_settings(polish=polish_osqp, max_iter=max_iter_osqp, linsys_solver=linsys_solver_osqp, warm_start=warm_start_osqp)
```

#### Simulate designed trajectories closed-loop

```python
xs_dmd_cl, us_dmd_cl = system.simulate(x0_cl, controller_dmd_cl, t_eval_cl)
xs_dmd_cl, us_dmd_cl = xs_dmd_cl.T, us_dmd_cl.T

xs_edmd_cl, us_edmd_cl = system.simulate(x0_cl, controller_edmd_cl, t_eval_cl)
xs_edmd_cl, us_edmd_cl = xs_edmd_cl.T, us_edmd_cl.T

controller_bedmd_cl.comp_time = []
xs_bedmd_cl, us_bedmd_cl = system.simulate(x0_cl, controller_bedmd_cl, t_eval_cl)
xs_bedmd_cl, us_bedmd_cl = xs_bedmd_cl.T, us_bedmd_cl.T

controller_kdnn_cl.comp_time = []
xs_kdnn_cl, us_kdnn_cl = system.simulate(x0_cl, controller_kdnn_cl, t_eval_cl)
xs_kdnn_cl, us_kdnn_cl = xs_kdnn_cl.T, us_kdnn_cl.T

controller_nmpc_cl.comp_time = []
xs_nmpc_cl, us_nmpc_cl = system.simulate(x0_cl, controller_nmpc_cl, t_eval_cl)
xs_nmpc_cl, us_nmpc_cl = xs_nmpc_cl.T, us_nmpc_cl.T
```

#### Plot/analyze the results

```python
plot_inds = [0, 1, 2, 0, 1]
subplot_inds = [1, 2, 3, 4, 8]

plt.figure(figsize=(12,2.5))
for ii in range(5):
    ind = plot_inds[ii]
    if ii < 3:
        ax = plt.subplot(1,4,subplot_inds[ii])
        plt.plot(t_eval_cl, xs_dmd_cl[ind,:], colors[0], label='DMD MPC')
        plt.plot(t_eval_cl, xs_edmd_cl[ind, :], colors[1], label='EDMD MPC')
        plt.plot(t_eval_cl, xs_bedmd_cl[ind, :], colors[2], label='K-NMPC')
        plt.plot(t_eval_cl, xs_kdnn_cl[ind, :], colors[3], label='KDNN-NMPC')
        plt.plot(t_eval_cl, xs_nmpc_cl[ind,:], colors[4], label='NMPC')

        plt.scatter(t_eval_cl[0], x0_cl[ind], color='g')
        plt.scatter(t_eval_cl[-1], set_pt_cl[ind], color='r')
        plt.ylabel(labels[ind])
        ax.yaxis.set_major_locator(MaxNLocator(integer=True))
        plt.title(titles[subplot_inds[ii]-1])
        plt.xlabel('Time (sec)')
        if subplot_inds[ii]==1:
            plt.legend(loc='upper left', frameon=False)
            plt.ylim(-0.15,2)
    else:
        bx = plt.subplot(2,4,subplot_inds[ii])
        plt.plot(t_eval_cl[:-1],us_dmd_cl[ind,:], color=colors[0], label='DMD MPC')
        plt.plot(t_eval_cl[:-1], us_edmd_cl[ind, :], color=colors[1], label='EDMD MPC')
        plt.plot(t_eval_cl[:-1], us_bedmd_cl[ind, :], color=colors[2], label='K-NMPC')
        plt.plot(t_eval_cl[:-1], us_kdnn_cl[ind, :], color=colors[3], label='KDNN-NMPC')
        plt.plot(t_eval_cl[:-1],us_nmpc_cl[ind,:], color=colors[4], label='NMPC')
        plt.plot([0, t_eval_cl[-1]], [umax[ind]+hover_thrust, umax[ind]+hover_thrust], ':k')
        plt.plot([0, t_eval_cl[-1]], [umin[ind]+hover_thrust, umin[ind]+hover_thrust], ':k')
        plt.ylabel(labels[ii+3])
        ax.yaxis.set_major_locator(MaxNLocator(integer=True))
        if subplot_inds[ii] == 4:
            plt.title('Control inputs')
        else:
            plt.xlabel('Time (sec)')
if save_figures:
    matplotlib.rcParams['pdf.fonttype'] = 42
    matplotlib.rcParams['ps.fonttype'] = 42
    plt.tight_layout()
    #plt.savefig(dropbox_folder + 'planar_quad_closed_loop.pdf', format='pdf', dpi=2400)
plt.show()

cost_cl_dmd = np.sum(np.diag((xs_dmd_cl[:,:-1]-set_pt_cl.reshape(-1,1)).T@Q_mpc_cl@(xs_dmd_cl[:,:-1]-set_pt_cl.reshape(-1,1)))) + (xs_dmd_cl[:,-1]-set_pt_cl).T@QN_mpc_cl@(xs_dmd_cl[:,-1]-set_pt_cl) + np.sum(np.diag(us_dmd_cl.T@R_mpc_cl@us_dmd_cl))
cost_cl_edmd = np.sum(np.diag((xs_edmd_cl[:,:-1]-set_pt_cl.reshape(-1,1)).T@Q_mpc_cl@(xs_edmd_cl[:,:-1]-set_pt_cl.reshape(-1,1)))) + (xs_edmd_cl[:,-1]-set_pt_cl).T@QN_mpc_cl@(xs_edmd_cl[:,-1]-set_pt_cl) + np.sum(np.diag(us_edmd_cl.T@R_mpc_cl@us_edmd_cl))
cost_cl_bedmd = np.sum(np.diag((xs_bedmd_cl[:,:-1]-set_pt_cl.reshape(-1,1)).T@Q_mpc_cl@(xs_bedmd_cl[:,:-1]-set_pt_cl.reshape(-1,1)))) + (xs_bedmd_cl[:,-1]-set_pt_cl).T@QN_mpc_cl@(xs_bedmd_cl[:,-1]-set_pt_cl) + np.sum(np.diag(us_bedmd_cl.T@R_mpc_cl@us_bedmd_cl))
cost_cl_kdnn = np.sum(np.diag((xs_kdnn_cl[:,:-1]-set_pt_cl.reshape(-1,1)).T@Q_mpc_cl@(xs_kdnn_cl[:,:-1]-set_pt_cl.reshape(-1,1)))) + (xs_kdnn_cl[:,-1]-set_pt_cl).T@QN_mpc_cl@(xs_kdnn_cl[:,-1]-set_pt_cl) + np.sum(np.diag(us_kdnn_cl.T@R_mpc_cl@us_kdnn_cl))
cost_cl_nmpc = np.sum(np.diag((xs_nmpc_cl[:,:-1]-set_pt_cl.reshape(-1,1)).T@Q_mpc_cl@(xs_nmpc_cl[:,:-1]-set_pt_cl.reshape(-1,1)))) + (xs_nmpc_cl[:,-1]-set_pt_cl).T@QN_mpc_cl@(xs_nmpc_cl[:,-1]-set_pt_cl) + np.sum(np.diag(us_nmpc_cl.T@R_mpc_cl@us_nmpc_cl))

print('Solution statistics:\n')
print(tabulate([['DMD MPC', "{:.4f}".format(cost_cl_dmd/cost_cl_nmpc), np.mean(controller_dmd_cl.nom_controller.comp_time), np.std(controller_dmd_cl.nom_controller.comp_time)],
                ['EDMD MPC', "{:.4f}".format(cost_cl_edmd/cost_cl_nmpc),np.mean(controller_edmd_cl.nom_controller.comp_time), np.std(controller_edmd_cl.nom_controller.comp_time)],
                ['bEDMD MPC', "{:.4f}".format(cost_cl_bedmd/cost_cl_nmpc), np.mean(controller_bedmd_cl.comp_time), np.std(controller_bedmd_cl.comp_time)],
                ['DNN bEDMD MPC', "{:.4f}".format(cost_cl_kdnn/cost_cl_nmpc), np.mean(controller_kdnn_cl.comp_time), np.std(controller_kdnn_cl.comp_time)],
                ['NMPC (benchmark, known model)',1, np.mean(controller_nmpc_cl.comp_time), np.std(controller_nmpc_cl.comp_time)]],
               headers=['Normalized cost,\nrealized trajectory', 'Mean comp. time (secs)', 'std comp. time (secs)']))
```

```python

```