# 02 - Reverse Time Migration

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

## Imaging requirement

Seismic imaging relies on two known parameters:

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

- **Background velocity model**. This is a velocity model that has been obtained by processing and inverting the field data. We will look at this methods in the following tutorial as it relies on the method we are describing here. This velocity model is usually a **smooth version** of the true velocity model.

## Imaging computational setup

In this tutorial, we will introduce the back-propagation operator. This operator simulates the adjoint wave-equation, that is a wave-equation solved in a reversed time order. This time reversal led to the naming of the method we present here, called Reverse Time Migration. The notion of adjoint in exploration geophysics is fundamental as most of the wave-equation based imaging and inversion methods rely on adjoint based optimization methods.

## Notes on the operators

As we already describe the creation of a forward modelling operator, we will use a thin wrapper function instead. This wrapper is provided by a utility class called `AcousticWaveSolver`, which provides all the necessary operators for seismic modeling, imaging and inversion. The `AcousticWaveSolver` provides a more concise API for common wave propagation operators and caches the the Devito `Operator` objects to avoid unnecessary recompilation. However, any newly introduced operators will be fully described and only used from the wrapper in the next tutorials.

As before we initialize printing and import some utilities. We also raise the Devito log level to avoid excessive logging for repeated operator invocations.

```python
import numpy as np
%matplotlib inline

from devito import configuration
configuration['log_level'] = 'WARNING'
```

## Computational considerations

Seismic inversion algorithms are generally very computationally demanding and require a large amount of memory to store the forward wavefield. In order to keep this tutorial as light-weight as possible we are using a very simple
velocity model that requires low temporal and special resolution. For a more realistic model, a second set of preset parameters for a reduced version of the 2D Marmousi data set [1] is provided below in comments. This can be run to create some more realistic subsurface images. However, this second present is more computationally demanding and requires a slightly more powerful workstation.

```python
# Configure model presets
from examples.seismic import demo_model

# Enable model presets here:
preset = 'twolayer-isotropic'  # A simple but cheap model (recommended)
# preset = 'marmousi2d'  # A larger more realistic model

# Standard preset with a simple two-layer model
if preset == 'twolayer-isotropic':
    def create_model():
        return demo_model('twolayer-isotropic', origin=(0., 0.), shape=(101, 101),
                          spacing=(10., 10.), nbpml=20)
    filter_sigma = (1, 1)
    nshots = 21
    nreceivers = 101
    t0 = 0.
    tn = 1000.  # Simulation last 1 second (1000 ms)
    f0 = 0.010  # Source peak frequency is 10Hz (0.010 kHz)

# A more computationally demanding preset based on the 2D Marmousi model
if preset == 'marmousi2d-isotropic':
    def create_model():
        return demo_model('marmousi2d-isotropic', data_path='../../../../opesci-data/')
    filter_sigma = (6, 6)
    nshots = 301  # Need good covergae in shots, one every two grid points
    nreceivers = 601  # One recevier every grid point
    t0 = 0.
    tn = 3500.  # Simulation last 3.5 second (3500 ms)
    f0 = 0.025  # Source peak frequency is 25Hz (0.025 kHz)
```

# True and smooth velocity models

First, we create the model data for the "true" model from a given demonstration preset. This model represents the subsurface topology for the purposes of this example and we will later use it to generate our synthetic data readings. We also generate a second model and apply a smoothing filter to it, which represents our initial model for the imaging algorithm. The perturbation between these two models can be thought of as the image we are trying to recover.

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

# Create true model from a preset
model = create_model()

# Create initial model and smooth the boundaries
model0 = create_model()
model0.vp = ndimage.gaussian_filter(model0.vp, sigma=filter_sigma, order=0)

# Plot the true and initial model and the perturbation between them
plot_velocity(model)
plot_velocity(model0)
plot_perturbation(model0, model)
```

## Acquisition geometry

Next we define the positioning and the wave signal of our source and the location of our receivers. To generate the wavelet for our source we require the discretized values of time that we are going to use to model a single "shot",,
which again depends on the grid spacing used in our model. For consistency this initial setup will look exactly as in the previous modelling tutorial, although we will vary the position of our source later on during the actual imaging algorithm.

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

# Define time discretization according to grid spacing
dt = model.critical_dt  # Time step from model grid spacing
nt = int(1 + (tn-t0) / dt)  # Discrete time axis length
time = np.linspace(t0, tn, nt)  # Discrete modelling time

src = RickerSource(name='src', grid=model.grid, f0=f0, time=np.linspace(t0, tn, nt))

# First, position source centrally in all dimensions, then set depth
src.coordinates.data[0, :] = np.array(model.domain_size) * .5
src.coordinates.data[0, -1] = 20.  # Depth is 20m

# We can plot the time signature to see the wavelet
src.show()
```

```python
# Define acquisition geometry: receivers
from examples.seismic import Receiver

# Initialize receivers for synthetic and imaging data
rec = Receiver(name='rec', grid=model.grid, npoint=nreceivers, ntime=nt)
rec.coordinates.data[:, 0] = np.linspace(0, model.domain_size[0], num=nreceivers)
rec.coordinates.data[:, 1] = 30.
```

# True and smooth data

We can now generate the shot record (receiver readings) corresponding to our true and initial models. The difference between these two records will be the basis of the imaging procedure.

For this purpose we will use the same forward modelling operator that was introduced in the previous tutorial, provided by the `WaveSolver` utility class. This object instantiates a set of pre-defined operators according to an initial definition of the acquisition geometry, consisting of source and receiver symbols. The solver objects caches the individual operators and provides a slightly more high-level API that allows us to invoke the modelling modelling operators from the initial tutorial in a single line. In the following cells we use this to generate shot data by only specifying the respective model symbol `m` to use, and the solver will create and return a new `Receiver` object the represents the readings at the previously defined receiver coordinates.

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

solver = AcousticWaveSolver(model, src, rec, space_order=4)
true_d , _, _ = solver.forward(src=src, m=model.m)
```

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

```python
#NBVAL_IGNORE_OUTPUT
# Plot shot record for true and smooth velocity model and the difference
from examples.seismic import plot_shotrecord

plot_shotrecord(true_d.data, model, t0, tn)
plot_shotrecord(smooth_d.data, model, t0, tn)
plot_shotrecord(smooth_d.data - true_d.data, model, t0, tn)
```

# Imaging with back-propagation

As we explained in the introduction of this tutorial, this method is based on back-propagation.

## Adjoint wave equation

If we go back to the modelling part, we can rewrite the simulation as a linear system solve:

\begin{equation}
\mathbf{A}(\mathbf{m}) \mathbf{u} = \mathbf{q}
\end{equation}

where $\mathbf{m}$ is the discretized square slowness, $\mathbf{q}$ is the discretized source and $\mathbf{A}(\mathbf{m})$ is the discretized wave-equation. The discretized wave-equation matricial representation is a lower triangular matrix that can be solve with forward substitution. The pointwise writing or the forward substitution leads to the time-stepping stencil.

On a small problem one could form the matrix explicitly and transpose it to obtain the adjoint discrete wave-equation:

\begin{equation}
\mathbf{A}(\mathbf{m})^T \mathbf{v} = \delta \mathbf{d}
\end{equation}

where $\mathbf{v}$ is the discrete **adjoint wavefield** and  $\delta \mathbf{d}$ is the data residual defined as the difference between the field/observed data and the synthetic data $\mathbf{d}_s = \mathbf{P}_r \mathbf{u}$. In our case we derive the discrete adjoint wave-equation from the discrete forward wave-equation to get its stencil.

## Imaging

Wave-equation based imaging relies on one simple concept:

- If the background velocity model is cinematically correct, the forward wavefield $\mathbf{u}$ and the adjoint wavefield $\mathbf{v}$ meet at the reflectors position at zero time offset.

The sum over time of the zero time-offset correlation of these two fields then creates an image of the subsurface. Mathematically this leads to the simple imaging condition:

\begin{equation}
  \text{Image} = \sum_{t=1}^{n_t} \mathbf{u}[t] \mathbf{v}[t]
\end{equation}

In the following tutorials we will describe a more advanced imaging condition that produces shaper and more accurate results.

## Operator

We will now define the imaging operator that computes the adjoint wavefield $\mathbf{v}$ and correlates it with the forward wavefield $\mathbf{u}$. This operator essentially consist of three components:
* Stencil update of the adjoint wavefield `v`
* Injection of the data residual at the adjoint source (forward receiver) location
* Correlation of `u` and `v` to compute the image contribution at each timestep

```python
# Define gradient operator for imaging
from devito import TimeFunction, Backward, Operator
from examples.seismic import PointSource

from sympy import solve, Eq

def ImagingOperator(model, image):
    # Define the wavefield with the size of the model and the time dimension
    v = TimeFunction(name='v', grid=model.grid, time_order=2, space_order=4)

    u = TimeFunction(name='u', grid=model.grid, time_order=2, space_order=4,
                     save=True, time_dim=nt)

    # Define the wave equation, but with a negated damping term
    eqn = model.m * v.dt2 - v.laplace - model.damp * v.dt

    # Use SymPy to rearranged the equation into a stencil expression
    stencil = Eq(v.backward, solve(eqn, v.backward)[0])

    # Define residual injection at the location of the forward receivers
    dt = model.critical_dt
    residual = PointSource(name='residual', grid=model.grid,
                           ntime=nt, coordinates=rec.coordinates.data)
    res_term = residual.inject(field=v, expr=residual * dt**2 / model.m, offset=model.nbpml)

    # Correlate u and v for the current time step and add it to the image
    image_update = Eq(image, image - u * v)

    return Operator([stencil] + res_term + [image_update],
                    subs=model.spacing_map, time_axis=Backward)
```

## Implementation of the imaging loop

As we just explained, the forward wave-equation is solved forward in time while the adjoint wave-equation is solved in a reversed time order. Therefore, the correlation of these two fields over time requires to store one of the two fields. The computational procedure for imaging follows:

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

This procedure is applied to multiple source positions (shots) and summed to obtain the full image of the subsurface. We can first visualize the varying locations of the sources that we will use.

```python
#NBVAL_IGNORE_OUTPUT

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

plot_velocity(model, source=source_locations)
```

```python
# Run imaging loop over shots
from devito import Function

# Create image symbol and instantiate the previously defined imaging operator
image = Function(name='image', grid=model.grid)
op_imaging = ImagingOperator(model, image)

for i in range(nshots):
    print('Imaging source %d out of %d' % (i+1, nshots))

    # Update source location
    src.coordinates.data[0, :] = source_locations[i, :]

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

    # Compute smooth data and full forward wavefield u0
    smooth_d, u0, _ = solver.forward(src=src, m=model0.m, save=True)

    # Compute gradient from the data residual
    v = TimeFunction(name='v', grid=model.grid, time_order=2, space_order=4)
    op_imaging(u=u0, v=v, m=model0.m, dt=model0.critical_dt, residual=smooth_d.data - true_d.data)
```

    Imaging source 1 out of 21
    Imaging source 2 out of 21
    Imaging source 3 out of 21
    Imaging source 4 out of 21
    Imaging source 5 out of 21
    Imaging source 6 out of 21
    Imaging source 7 out of 21
    Imaging source 8 out of 21
    Imaging source 9 out of 21
    Imaging source 10 out of 21
    Imaging source 11 out of 21
    Imaging source 12 out of 21
    Imaging source 13 out of 21
    Imaging source 14 out of 21
    Imaging source 15 out of 21
    Imaging source 16 out of 21
    Imaging source 17 out of 21
    Imaging source 18 out of 21
    Imaging source 19 out of 21
    Imaging source 20 out of 21
    Imaging source 21 out of 21

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

# Plot the inverted image
plot_image(np.diff(image.data, axis=1))
```

And we have an image of the subsurface with a strong reflector at the original location.

## References

[1] _Versteeg, R.J. & Grau, G. (eds.) (1991): The Marmousi experience. Proc. EAGE workshop on Practical Aspects of Seismic Data Inversion (Copenhagen, 1990), Eur. Assoc. Explor. Geophysicists, Zeist._


 ILI286 - Computación Científica II

 Ecuaciones Diferenciales Parciales: Elípticas

  [S]cientific [C]omputing [T]eam

 Version: 1.14

# Tabla de Contenidos
* [Introducción](#intro)
* [Marco Teórico](#teo)
* [Ecuación de Poisson](#poisson)
* [Ecuación de Helmotz](#helmotz)
* [Acknowledgements](#acknowledgements)

```python
import numpy as np
from mpl_toolkits.mplot3d import axes3d
from matplotlib import pyplot as plt
from ipywidgets import interact
from ipywidgets import IntSlider
import sympy as sym
import matplotlib as mpl
mpl.rcParams['font.size'] = 14
mpl.rcParams['axes.labelsize'] = 20
mpl.rcParams['xtick.labelsize'] = 14
mpl.rcParams['ytick.labelsize'] = 14
sym.init_printing()
%matplotlib inline

def plot(x,y,w,elev=40,azim=230):
    # Plot the solution
    X,Y = np.meshgrid(y,x)
    W = w.reshape(X.shape)
    fig = plt.figure(figsize=(10,10))
    ax = fig.add_subplot(111, projection='3d')
    ax.plot_wireframe(X, Y, W)
    #ax.plot_surface(X, Y, W, alpha=0.25)
    plt.xlabel("y")
    plt.ylabel("x")
    #ax.set_zlim(0.,1.)
    #plt.savefig("sol%dx%d.png"%(Nx+1,Ny+1))
    ax.view_init(elev,azim)
    plt.show()
```

# Introducción

En el siguiente notebook se estudia la resolución numérica de **ecuaciones diferenciales parciales elípticas**. La resolución de estas tiene gran importancia, ya que aparecen repetidas veces en diferentes modelos físicos relacionados con los potenciales de energía. Por ejemplo, el potencial de carga eléctrica de acuerdo a las ecuaciones de _Maxwell_ puede ser escrito como:
$$
\Delta u = -\frac{\rho}{\epsilon},
$$
con $u$ el potencial eléctrico, $\rho$ la densidad de carga, y $\epsilon$ el la permitividad eléctrica.

El método numérico que estudiaremos se basa en **diferencias finitas**, del mismo modo en que se hizo para resolver EDOs anteriormente, esto es, las derivadas son aproximadas por gradientes calculados numéricamente en base a los valores vecinos de cada punto. Sin embargo ahora *las derivadas son parciales!*, pero como veremos esto no es gran problema.

# Marco Teórico

Si consideramos una función $u(x,y)$ dos veces diferenciable, entonces se define el operador *Laplaciano* como:

$$
\Delta u(x,y) = u_{xx}(x,y) + u_{yy}(x,y),
$$

si se considera además una función $f(x,y)$, entonces es posible definir:

$$
\Delta u(x,y) = u_{xx}(x,y) + u_{yy}(x,y) = f(x,y), \ \ \  \text{con } \ x,y \in \Omega \ \text{ y condiciones de borde en } \ \partial \Omega
$$

como la **ecuación de Poisson**, la cual es una de las más conocidas dentro de la clase _ Elípticas _. El caso particular donde $f(x,y) = 0$ se conoce como **ecuación de Laplace**.

### Formulación numérica

Sea la ecuación de Laplace $\Delta u(x,y) = 0$ sobre un dominio rectangular $[x_a, x_b] \times [y_a, y_b]$, con condiciones de borde de Dirichlet:

\begin{align}
    u(x,y_a) &= g_1(x) \\
    u(x,y_b) &= g_2(x) \\
    u(x_a,y) &= g_3(y) \\
    u(x_b,y) &= g_4(y) .
\end{align}

Para resolver este problema por diferencias finitas, es necesario discretizar el dominio $\Omega$ sobre el que se define la función. Tal discretización se puede ver gráficamente como se muestra a continuación:

```python
x = np.linspace(0., 1., 10)
y = np.linspace(0., 1., 10)
xgrid, ygrid = np.meshgrid(x, y, sparse=False)
plt.figure(figsize=(8,8))
plt.scatter(xgrid.ravel(), ygrid.ravel())
plt.plot((0,1),(0,0), 'g--' ,lw=2)
plt.plot((1,1),(0,1), 'g--' ,lw=2)
plt.plot((1,0),(1,1), 'g--' ,lw=2)
plt.plot((0,0),(1,0), 'g--' ,lw=2)
plt.title('Esquema de discretizacion uniforme')
plt.xlim(-0.1,1.1)
plt.ylim(-0.1,1.1)
plt.axis('equal')
plt.show()
```

en donde cada punto interior (*azul*), representa un punto donde queremos conocer el valor de la función $u(x,y)$. Consideraremos además que  $h_x$ y $h_y$ el *space step* de la malla.

Dado que las derivadas del Laplaciano son *parciales*, se puede utilizar diferencias finitas sin ninguna modificación. Para aproximar cada segunda derivada se utiliza *diferencias centradas* (**centered difference formula**), que posee un error de orden cuadrático en $h$.

$$
\underbrace{ \frac{u(x-h_x,y) - 2 u(x,y) + u(x+h_x, y)}{h_x^2} + O(h_x^2) }_{= u_{xx}(x,y)} + \underbrace{ \frac{u(x,y-h_y) - 2 u(x,y) + u(x, y+h_y)}{h_y^2} + O(h_y^2) }_{= u_{yy}(x,y)} = f(x,y)
$$

Ocupando la notación $u(x_i,y_j) = w_{ij}$ para el punto $(x_i, y_j)$ de la malla, la ecuación de discretización es finalmente:

$$
\frac{w_{i-1,j} - 2 w_{i,j} + w_{i+1,j}}{h_x^2} + \frac{w_{i,j-1} - 2 w_{i,j} + w_{i,j+1}}{h_y^2} \approx f(x_i, y_j),
$$

válida para todos los puntos $(x_i, y_j) \in \Omega - \partial \Omega$ (puntos interiores). Hay una aspecto muy importante a notar: **Esta ecuación es lineal en los $\mathbf{w_{i,j}}$!**. El procedimiento es entonces como sigue:

* Evaluar la ecuación discretizada para cada punto interior $(x_i, y_j)$, utilizando la información de los valores en la frontera.
* Construir un sistema lineal $A \mathbf{w} = b$, en donde $\mathbf{w}$ es el vector con todos los $w_{i,j}$. (Para su construcción revisar el libro guía: *Numerical Analysis, Timothy Sauer*).
* Resolver el sistema anterior para $\mathbf{w}$, utilizando cualquier `solver` de sistemas lineales (LU, PALU, Jacobi, Gauss-Seidel, etc).

# Ecuación de Poisson

Buscaremos resolver la ecuación de Poisson con condición de frontera de Dirichlet:

\begin{align*}
\Delta u(x,y) &= f(x,y) \ , \ (x, y) \in \Omega=[0,1]\times[0,1] \\
u(x,0) &= b(x)\\
u(x,1) &= t(x)\\
u(0, y) &= l(y)\\
u(1,y) &= r(y)
\end{align*}

A continuación se definen tres problemas a resolver.

```python
# Problema 1
xmin, xmax = -1, 1.
ymin, ymax = -1, 1.
f  = lambda x,y : x*y
bottom = lambda x : 0# np.sin(np.pi*x)
top = lambda x : 0 #np.sin(np.pi*x)
left = lambda y : 0
right = lambda y: 0
P1 = {"f":f, "b":bottom, "t":top, "l":left, "r":right,
      "xmin":xmin, "xmax":xmax, "ymin":ymin, "ymax":ymax}

# Problema 2
xmin, xmax = 0., 1.
ymin, ymax = 0., 1.
f  = lambda x,y : x
bottom = lambda x : np.sin(np.pi*x)
top = lambda x : np.sin(np.pi*x)
left = lambda y : 0
right = lambda y: 0
P2 = {"f":f, "b":bottom, "t":top, "l":left, "r":right,
      "xmin":xmin, "xmax":xmax, "ymin":ymin, "ymax":ymax}

# Problema 3
xmin, xmax = -1, 1.
ymin, ymax = -1, 1.
f  = lambda x,y : 0
bottom = lambda x : np.sin(np.pi*x)
top = lambda x : -np.sin(np.pi*x)
left = lambda y : 0
right = lambda y: 0
P3 = {"f":f, "b":bottom, "t":top, "l":left, "r":right,
      "xmin":xmin, "xmax":xmax, "ymin":ymin, "ymax":ymax}

# Problema 4
xmin, xmax = 0, 1.
ymin, ymax = 0, 1.
f  = lambda x,y : x*np.exp(y)
bottom = lambda x : x
top = lambda x : x*np.exp(1)
left = lambda y : 0*y
right = lambda y: np.exp(y)
P4 = {"f":f, "b":bottom, "t":top, "l":left, "r":right,
      "xmin":xmin, "xmax":xmax, "ymin":ymin, "ymax":ymax}

P_Poisson=[('P1', P1), ('P2', P2),('P3', P3),('P4', P4)]
```

La función `solve_laplace()` es la encargada de construir el sistema lineal correspondiente para al problema `P` a resolver.

**Q:** _¿Podría usted explicar esta construcción en base a lo visto en la formulación teórica?_.

```python
def solve_laplace(P, Nx=30, Ny=30,flag_plot=False,elev=40,azim=230):
    # Discretize x and y
    x = np.linspace(P["xmin"], P["xmax"], Nx+1)
    y = np.linspace(P["ymin"], P["ymax"], Ny+1)

    # Define the discretization parameters
    dx = x[1]-x[0]
    dy = y[1]-y[0]

    # Create the matrix and the right hand size vector
    A = np.zeros([(Nx+1)*(Ny+1), (Nx+1)*(Ny+1)])
    b = np.zeros([(Nx+1)*(Ny+1), 1])

    # Define global indexing
    def index(i, j, nCols=(Ny+1)):
        return j + i*nCols

    # Fill up the matrix and right hand side vector
    for i in range(Nx+1):
        for j in range(Ny+1):
            k = index(i,j)
            if j==0: # y=ymin, bottom
                A[k,k] =  1.
                b[k] = P["b"](x[i])
            elif i==Nx: # x=xmax, right
                A[k,k] =  1.
                b[k] = P["r"](y[j])
            elif j==Ny: # y=ymax, top
                A[k,k] =  1.
                b[k] = P["t"](x[i])
            elif i==0:  # x=xmin, left
                A[k,k] =  1.
                b[k] = P["l"](y[j])
            else:
                A[k, k] = -2./dx**2 - 2./dy**2
                A[k,index(i+1,j)] =  1./dx**2
                A[k,index(i-1,j)] =  1./dx**2
                A[k,index(i,j-1)] =  1./dy**2
                A[k,index(i,j+1)] =  1./dy**2
                b[k] = P["f"](x[i], y[j])

    # Solve the linear system
    w = np.linalg.solve(A, b)

    if flag_plot:
        plot(x,y,w,elev,azim)
        return

    return x, y, w
```

```python
elev_widget = IntSlider(min=0, max=180, step=10, value=40)
azim_widget = IntSlider(min=0, max=360, step=10, value=230)

interact(solve_laplace,P=P_Poisson,Nx=(5,50,5),Ny=(5,50,5),flag_plot=[True],elev=elev_widget,azim=azim_widget)
```

    interactive(children=(Dropdown(description='P', options=(('P1', {'f': <function <lambda> at 0x121eec0e0>, 'b':…

    <function __main__.solve_laplace(P, Nx=30, Ny=30, flag_plot=False, elev=40, azim=230)>

# Ecuación de Helmholtz

Buscaremos resolver la ecuación de Helmholtz con condición de frontera de Dirichlet:

\begin{align*}
\Delta u(x,y) -\lambda \, u(x,y)&= f(x,y) \ , \ (x, y) \in \Omega=[0,1]\times[0,1] \\
u(x,0) &= b(x)\\
u(x,1) &= t(x)\\
u(0, y) &= l(y)\\
u(1,y) &= r(y)
\end{align*}

A continuación se definen tres problemas a resolver.

```python
# Problem 1
Lambda = 0.1
xmin, xmax = 0., 1.
ymin, ymax = 0., 1.
f  = lambda x,y : 0
bottom = lambda x : 0
top = lambda x : 0
left = lambda y : 1
right = lambda y: 1
P1 = {"Lambda":Lambda, "f":f, "b":bottom, "t":top, "l":left, "r":right,
      "xmin":xmin, "xmax":xmax, "ymin":ymin, "ymax":ymax}

# Problem 2
Lambda = 2.0
xmin, xmax = 0., 1.
ymin, ymax = 0., 1.
f  = lambda x,y : 0
bottom = lambda x : 0
top = lambda x : 0
left = lambda y : 1
right = lambda y: 1
P2 = {"Lambda":Lambda, "f":f, "b":bottom, "t":top, "l":left, "r":right,
      "xmin":xmin, "xmax":xmax, "ymin":ymin, "ymax":ymax}

# Problem 3
Lambda = 0.0
xmin, xmax = -1, 1.
ymin, ymax = -1, 1.
f  = lambda x,y : 0
bottom = lambda x : 0
top = lambda x : 0
left = lambda y : np.sin(np.pi*y)
right = lambda y: -np.sin(np.pi*y)
P3 = {"Lambda":Lambda, "f":f, "b":bottom, "t":top, "l":left, "r":right,
      "xmin":xmin, "xmax":xmax, "ymin":ymin, "ymax":ymax}

P_Helmholtz=[('P1', P1), ('P2', P2),('P3', P3)]
```

La función `solve_herlmotz()` es la encargada de construir el sistema lineal correspondiente para al problema `P` a resolver.

**Q:** _¿Podría usted explicar esta construcción en base a lo visto en la formulación teórica?_.

```python
def solve_helmholtz(P, Nx=30, Ny=30,flag_plot=False,elev=40,azim=230):
    # Discretize x and y
    x = np.linspace(P["xmin"], P["xmax"], Nx+1)
    y = np.linspace(P["ymin"], P["ymax"], Ny+1)
    L = P["Lambda"]

    # Define the discretization parameters
    dx = x[1]-x[0]
    dy = y[1]-y[0]

    # Create the matrix and the right hand size vector
    A = np.zeros([(Nx+1)*(Ny+1), (Nx+1)*(Ny+1)])
    b = np.zeros([(Nx+1)*(Ny+1), 1])

    # Define global indexing
    def index(i, j, nCols=(Ny+1)):
        return j + i*nCols

    # Fill up the matrix and right hand side vector
    for i in range(Nx+1):
        for j in range(Ny+1):
            k = index(i,j)
            if j==0: # y=ymin, bottom
                A[k,k] =  -1.5/dy
                A[k,index(i,j+1)] = 2.0/dy
                A[k,index(i,j+2)] =-0.5/dy
                b[k] = P["b"](x[i])
            elif i==Nx: # x=xmax, right
                A[k,k] =  1.
                b[k] = P["r"](y[j])
            elif j==Ny: # y=ymax, top
                A[k,k] =  1.5/dy
                A[k,index(i,j-1)] =  -2.0/dy
                A[k,index(i,j-2)] =  +0.5/dy
                b[k] = P["t"](x[i])
            elif i==0:  # x=xmin, left
                A[k,k] =  1.
                b[k] = P["l"](y[j])
            else:
                A[k, k] = -2./dx**2 - 2./dy**2 - L
                A[k,index(i+1,j)] =  1./dx**2
                A[k,index(i-1,j)] =  1./dx**2
                A[k,index(i,j-1)] =  1./dy**2
                A[k,index(i,j+1)] =  1./dy**2
                b[k] = P["f"](x[i], y[j])

    # Solve the linear system
    w = np.linalg.solve(A, b)

    if flag_plot:
        plot(x,y,w,elev,azim)
        return

    return x, y, w
```

```python
elev_widget = IntSlider(min=0, max=180, step=10, value=40)
azim_widget = IntSlider(min=0, max=360, step=10, value=230)

interact(solve_helmholtz,P=P_Helmholtz,Nx=(5,50,5),Ny=(5,50,5),flag_plot=[True],elev=elev_widget,azim=azim_widget)
```

    interactive(children=(Dropdown(description='P', options=(('P1', {'Lambda': 0.1, 'f': <function <lambda> at 0x1…

    <function __main__.solve_helmholtz(P, Nx=30, Ny=30, flag_plot=False, elev=40, azim=230)>

# Acknowledgements

* _Material creado por profesor Claudio Torres_ (`ctorres@inf.utfsm.cl`) _y ayudantes: Alvaro Salinas y Martín Villanueva. DI UTFSM. Abril 2016._

***
### DISCLAIMER ###

El presente notebook ha sido creado para el curso **ILI286 - Computación Científica 2**, del [Departamento de Informática](http://www.inf.utfsm.cl/), [Universidad Técnica Federico Santa María](http://www.utfsm.cl/).

El material ha sido creado por Claudio Torres <ctorres@inf.utfsm.cl> y ayudantes, y es distribuido sin restricciones. En caso de encontrar un error, por favor no dude en contactarnos.

[Update 2016] V1.1 (Martín) Integrados los notebook de Poisson y Helmholtz en un sólo notebook. Se agregó contexto y marco teórico.

[Update 2019] (C. Torres) Fixing titles. Adding ipython widget and adding notation. Fixing typos.
***

```python

```


# Episodic Mountain Car with function appoximation and control

This Notebook is intended to solve the Episodic Mountain car problem using Semi-gradient sarsa and Tile Coding.

The description of the problem is given below:

"A car is on a one-dimensional track, positioned between two "mountains". The goal is to drive up the mountain on the right; however, the car's engine is not strong enough to scale the mountain in a single pass. Therefore, the only way to succeed is to drive back and forth to build up momentum."

An extensive description and solution of the problem can be seen here [Section 10.1 Reinforment Learning an Introduction](http://www.incompleteideas.net/book/RLbook2018.pdf#page=267)

Image and Text taken from Taken from [Official documentaiton Mountain car](https://gym.openai.com/envs/MountainCar-v0/).

```python
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import tiles3 as tc
from tqdm import tqdm

import gym
from gym.wrappers import Monitor
from utils import *

%matplotlib inline
```

## Undestanding the Workflow of OpenAI

The following variables are used at each timestep and they are returned by the Mountain Car environment.

- **observation** (object): an environment-specific object representing your observation of the environment. For example, pixel data from a camera, joint angles and joint velocities of a robot, or the board state in a board game.
- **reward** (float): amount of reward achieved by the previous action. The scale varies between environments, but the goal is always to increase your total reward.
- **done** (boolean): whether it’s time to reset the environment again. Most (but not all) tasks are divided up into well-defined episodes, and done being True indicates the episode has terminated. (For example, perhaps the pole tipped too far, or you lost your last life.)
- **info** (dict): diagnostic information useful for debugging. It can sometimes be useful for learning (for example, it might contain the raw probabilities behind the environment’s last state change). However, official evaluations of your agent are not allowed to use this for learning.

As a quick recap, the diagram below explains the workflow of a Markov Decision Process (MDP)

Image taken from [Section 3.1 Reinforment Learning an Introduction](http://www.incompleteideas.net/book/RLbook2018.pdf#page=70)

## Environment and Agent specifications

Below are presented the main features of the environment and agent. Overall, the action space of the problem is discrete with three posible actions. The observations or state space is continuios, therefore it is necessary to use a function approximation technique to solve this challenge. The agent receives a reward of -1 at each timestep unless it reaches the goal. The episode ends if the agent reaches the goal or a specific number of iterations are done. Additionally, the agent will always start at a random position between $-0.6$ and $-0.4$ with zero velocity.

**Observation**:

     Type:  Box(2)
     Num    Observation               Min            Max
     0      Car Position              -1.2           0.6
     1      Car Velocity              -0.07          0.07

**Actions**:

     Type: Discrete(3)
     Num    Action
     0      Accelerate to the Left
     1      Don't accelerate
     2      Accelerate to the Right

     Note: This does not affect the amount of velocity affected by the gravitational pull acting on the car

**Reward**:

     Reward of 0 is awarded if the agent reached the flag(position = 0.5) on top of the mountain
     Reward of -1 is awarded if the position of the agent is less than 0.5

**Starting State**:

     The position of the car is assigned a uniform random value in [-0.6 , -0.4]
     The velocity of the car is always assigned to 0

**Episode Termination**:

     The car position is more than 0.5
     Episode length is greater than 200

For further information see [Github source code](https://github.com/openai/gym/blob/master/gym/envs/classic_control/mountain_car.py)

The next cell aims to show how to iterate with the action and observation space of the agent and extract relevant information from it

```python
env = gym.make("MountainCar-v0")
observation = env.reset()

# Object's type in the action Space
print("The Action Space is an object of type: {0}\n".format(env.action_space))
# Shape of the action Space
print("The shape of the action space is: {0}\n".format(env.action_space.n))
# Object's type in the Observation Space
print("The Environment Space is an object of type: {0}\n".format(env.observation_space))
# Shape of the observation space
print("The Shape of the dimension Space are: {0}\n".format(env.observation_space.shape))
# The high and low values in the observation space
print("The High values in the observation space are {0}, the low values are {1}\n".format(
    env.observation_space.high, env.observation_space.low))
# Minimum and Maximum car position
print("The minimum and maximum car's position are: {0}, {1}\n".format(
    env.observation_space.low[0], env.observation_space.high[0]))
# Minimum and Maximum car velocity
print("The minimum and maximum car's velocity are: {0}, {1}\n".format(
    env.observation_space.low[1], env.observation_space.high[1]))
# Example of observation
print("The Observations at a given timestep are {0}\n".format(env.observation_space.sample()))
```

    The Action Space is an object of type: Discrete(3)

    The shape of the action space is: 3

    The Environment Space is an object of type: Box(2,)

    The Shape of the dimension Space are: (2,)

    The High values in the observation space are [0.6  0.07], the low values are [-1.2  -0.07]

    The minimum and maximum car's position are: -1.2000000476837158, 0.6000000238418579

    The minimum and maximum car's velocity are: -0.07000000029802322, 0.07000000029802322

    The Observations at a given timestep are [-1.0707569  0.0590123]

# Tile Coding Class

For a complete explanation about what is tile coding and how it works, see [Section 9.5.4 of Reinforment Learning an Introduction](http://www.incompleteideas.net/book/RLbook2018.pdf#page=239). Overall, this is a way to create features that can both provide good generalization and discrimination for value function approximation. Tile coding consists of multiple overlapping tiling, where each tiling is a partitioning of the space into tiles.

**Note**: Tile coding can be only be used with 2d observation spaces.

This technique is implemented using Tiles3, which is a python library written by Richard S. Sutton. For the full documentation see [Tiles3 documentation](http://incompleteideas.net/tiles/tiles3.html)

Image taken from [Section 9.5.4 of Reinforment Learning an Introduction](http://www.incompleteideas.net/book/RLbook2018.pdf#page=239)

```python
# Tile Coding Class
class MountainCarTileCoder:
    def __init__(self, iht_size=4096, num_tilings=8, num_tiles=8):
        """
        Initializes the MountainCar Tile Coder
        Initializers:
        iht_size -- int, the size of the index hash table, typically a power of 2
        num_tilings -- int, the number of tilings
        num_tiles -- int, the number of tiles. Here both the width and height of the
                     tile coder are the same
        Class Variables:
        self.iht -- tc.IHT, the index hash table that the tile coder will use
        self.num_tilings -- int, the number of tilings the tile coder will use
        self.num_tiles -- int, the number of tiles the tile coder will use
        """
        self.iht = tc.IHT(iht_size)
        self.num_tilings = num_tilings
        self.num_tiles = num_tiles

    def get_tiles(self, position, velocity):
        """
        Takes in a position and velocity from the mountaincar environment
        and returns a numpy array of active tiles.

        Arguments:
        position -- float, the position of the agent between -1.2 and 0.5
        velocity -- float, the velocity of the agent between -0.07 and 0.07
        returns:
        tiles - np.array, active tiles
        """
        # Set the max and min of position and velocity to scale the input
        # The max position is set to 0.5 as this is the position to end the experiment
        POSITION_MIN = -1.2
        POSITION_MAX = 0.5
        VELOCITY_MIN = -0.07
        VELOCITY_MAX = 0.07

        # Scale position and velocity by multiplying the inputs of each by their scale
        position_scale = self.num_tiles / (POSITION_MAX - POSITION_MIN)
        velocity_scale = self.num_tiles / (VELOCITY_MAX - VELOCITY_MIN)

        # Obtain active tiles for current position and velocity
        tiles = tc.tiles(self.iht, self.num_tilings, [position * position_scale,
                                                      velocity * velocity_scale])

        return np.array(tiles)
```

```python
# Test the TileCoder class
mctc = MountainCarTileCoder(iht_size = 1024, num_tilings = 8, num_tiles = 8)
tiles = mctc.get_tiles(position = -1.0, velocity = 0.01)
# Tiles obtained at a random pos and vel
print("The Tiles obtained are: {0}\n".format(tiles))
```

    The Tiles obtained are: [0 1 2 3 4 5 6 7]

# Implementing Sarsa Agent

To solve the Mountain Car problem, Value Function approximation and control will be used (Owing to the continuous state space). As a quick recap, Action-values can be computed using value function approximation giving the following equation.

\begin{equation}
q_\pi(s) \approx \hat{q}(s, a, w) \doteq w^T x(s,a)
\end{equation}

Where $w$ are a set of weights and $x(s,a)$ are the features vector which are computed using tile coding.

Using the Tile coder implemented above it is possible to compute the action-values $\hat{q}(s, a, w)$ and solve this RL task.

The equation to update the weights using the Sarsa algorithm is given below. Here, $\nabla \hat{q}(S_t, A_t, w)$ is the gradient of the action-values approximation but as $x(s,a)$ is a linear function, the gradient is one only for the active features.

\begin{equation}
w \leftarrow w + \alpha[R_{t+1} + \gamma \hat{q}(S_{t+1}, A_{t+1}, w)- \hat{q}(S_t, A_t, w)]\nabla \hat{q}(S_t, A_t, w)
\end{equation}

Additionally, the update "Target" is composed of the following terms:

\begin{equation}
\delta \leftarrow R_{t+1} + \gamma \hat{q}(S_{t+1}, A_{t+1}, w)
\end{equation}

\begin{equation}
w \leftarrow w + \alpha[\delta - \hat{q}(S_t, A_t, w)]\nabla \hat{q}(S_t, A_t, w)
\end{equation}

The Pseudo-code implementation of this algorithm is given below.

For further details, see [Section 9.5.4 of Reinforment Learning an Introduction](http://www.incompleteideas.net/book/RLbook2018.pdf#page=266). Image taken from the last reference.

```python
# SARSA
class SarsaAgent():
    """
    Initialization of Sarsa Agent. All values are set to None so they can
    be initialized in the agent_init method.
    """
    def __init__(self, agent_info={}):
        """Setup for the agent called when the experiment first starts."""
        self.last_action = None
        self.last_state = None
        self.epsilon = None
        self.gamma = None
        self.iht_size = None
        self.w = None
        self.alpha = None
        self.num_tilings = None
        self.num_tiles = None
        self.mctc = None
        self.initial_weights = None
        self.num_actions = None
        self.previous_tiles = None

    def agent_init(self, agent_info={}):
        """Setup for the agent called when the experiment first starts."""
        self.num_tilings = agent_info.get("num_tilings", 8)
        self.num_tiles = agent_info.get("num_tiles", 8)
        self.iht_size = agent_info.get("iht_size", 4096)
        self.epsilon = agent_info.get("epsilon", 0.0)
        self.gamma = agent_info.get("gamma", 1.0)
        self.alpha = agent_info.get("alpha", 0.5) / self.num_tilings
        self.initial_weights = agent_info.get("initial_weights", 0.0)
        self.num_actions = agent_info.get("num_actions", 3)

        # Initialize self.w to three times the iht_size. Recall this is because
        # we need to have one set of weights for each action (Stacked values).
        self.w = np.ones((self.num_actions, self.iht_size)) * self.initial_weights

        # Initialize self.mctc to the mountaincar verions of the  tile coder created
        self.mctc = MountainCarTileCoder(iht_size = self.iht_size,
                                         num_tilings = self.num_tilings,
                                         num_tiles = self.num_tiles)

    def select_action(self, tiles):
        """
        Selects an action using epsilon greedy
        Args:
        tiles - np.array, an array of active tiles
        Returns:
        (chosen_action, action_value) - (int, float), tuple of the chosen action
                                        and it's value
        """
        action_values = []
        chosen_action = None

        # Obtain action values for all actions (sum through rows)
        action_values = np.sum(self.w[:, tiles], axis = 1)

        # Epsilon Greedy action selecion
        if np.random.random() < self.epsilon:
            # Select random action among the three posible actions
            chosen_action = np.random.randint(self.num_actions)
        else:
            # Select the greedy action
            chosen_action = argmax(action_values)

        return chosen_action, action_values[chosen_action]

    def agent_start(self, state):
        """The first method called when the experiment starts, called after
        the environment starts.
        Args:
            state (Numpy array): the state observation from the
                environment's env.reset() function.
        Returns:
            The first action the agent takes.
        """
        # Current state
        position, velocity = state

        # Obtain tiles activated at state cero
        active_tiles = self.mctc.get_tiles(position = position, velocity = velocity)
        # Select an action and obtain action values of the state
        current_action, action_value = self.select_action(active_tiles)

        # Save action as last action
        self.last_action = current_action
        # Save tiles as previous tiles
        self.previous_tiles = np.copy(active_tiles)

        return self.last_action

    def agent_step(self, reward, state):
        """A step taken by the agent.
        Args:
            reward (float): the reward received for taking the last action taken
            state (Numpy array): the state observation from the
                environment's step based, where the agent ended up after the
                last step
        Returns:
            The action the agent is taking.
        """
        # Current state
        position, velocity = state

        # Compute current tiles
        active_tiles = self.mctc.get_tiles(position = position, velocity = velocity)
        # Obtain new action and action value before updating actition values
        current_action, action_value = self.select_action(active_tiles)

        # Update the Sarsa Target (delta)
        target = reward + (self.gamma * action_value)

        # Compute last action values to update weights
        last_action_val = np.sum(self.w[self.last_action][self.previous_tiles])

        # As we are using tile coding, which is a variant of linear function approximation
        # The gradient of the active tiles are one, otherwise cero.
        grad = 1
        self.w[self.last_action][self.previous_tiles] = self.w[self.last_action][self.previous_tiles] + \
            self.alpha * (target - last_action_val) * grad

        self.last_action = current_action
        self.previous_tiles = np.copy(active_tiles)
        return self.last_action

    def agent_end(self, reward):
        """Run when the agent terminates.
        Args:
            reward (float): the reward the agent received for entering the
                terminal state.
        """

        # There is no action_value used here because this is the end
        # of the episode.

        # Compute delta
        target = reward
        # Compute last action value
        last_action_val = np.sum(self.w[self.last_action][self.previous_tiles])
        grad = 1
        # Update weights
        self.w[self.last_action][self.previous_tiles] = self.w[self.last_action][self.previous_tiles] + \
            self.alpha * (target - last_action_val) * grad

    def return_action_value(self, state):
        """Run to obtain action-values for a given state.
        Args:
            state (Numpy array): the state observation
        Returns:
            The max action-value
        """

        # Current state
        position, velocity = state

        # Obtain tiles activated at state cero
        active_tiles = self.mctc.get_tiles(position = position, velocity = velocity)
        # Obtain action values for all actions (sum through rows)
        action_values = np.sum(self.w[:, active_tiles], axis = 1)
        # Obtain max action value
        max_action_value = np.max(action_values)

        return max_action_value

```

## Running the experiment

The following lines solves the Mountain Car problem and plot the average reward obtained over episodes and steps taken to solve the challenge at a specific episode.

```python
# Test Sarsa Agent
num_runs = 10
num_episodes = 100
agent_info_options = {"num_tilings": 8, "num_tiles": 8, "iht_size": 4096,
                      "epsilon": 0.0, "gamma": 1.0, "alpha": 0.5,
                      "initial_weights": 0.0, "num_actions": 3}

# Variable to store the amount of steps taken to solve the challeng
all_steps = []
# Variable to save the rewards in an episode
all_rewards = []

# Agent
agent = SarsaAgent(agent_info_options)

# Environment
env = gym.make('MountainCar-v0')
env.reset()
# Maximum number of possible iterations (default was 200)
env._max_episode_steps = 10000

# Number of runs are the times the experiment will start again (a.k.a episode)
for n_runs in tqdm(range(num_runs)):

    # Resets environment
    observation = env.reset()
    # Reset agent
    agent.agent_init(agent_info_options)
    # Generate last state and action in the agent
    last_action = agent.agent_start(observation)
    # Steps taken at each episode to solve the challenge
    steps_per_episode = []
    rewards_per_episode = []

    # Times the environment will start again without resetting the agent
    for t in range(num_episodes):
        # Store number of steps taken to solve experiment
        n_steps = 0
        rewards = 0
        # Reset done flag
        done = False
        # Reset environment
        observation = env.reset()
        # Run until the experiment is over
        while not done:

            # Take a step with the environment
            observation, reward, done, info = env.step(last_action)

            # Number of steps the agent take to solve the challenge
            n_steps += 1
            # Accumulate reward
            rewards += reward

            # If the goal has been reached stop
            if done:
                # Last step with the agent
                agent.agent_end(reward)

            else:
                # Take a step with the agent
                last_action = agent.agent_step(reward, observation)

        # Save the amount of steps needed to complete the experiment
        # Without rebooting the agent
        steps_per_episode.append(n_steps)
        # Save the amount of award obtained at each episode
        rewards_per_episode.append(rewards)

    # Save the list of steeps needed to finish the experiment
    # in all the Episodes
    all_steps.append(np.array(steps_per_episode))
    # Awards obtained in every episode
    all_rewards.append(np.array(rewards_per_episode))

env.close()
```

    100%|██████████| 10/10 [00:47<00:00,  4.74s/it]

```python
steps_average = np.mean(np.array(all_steps), axis=0)
plt.plot(steps_average, label = 'Steps')
plt.xlabel("Episodes")
plt.ylabel("Iterations",rotation=0, labelpad=40)
plt.xlim(-0.2, num_episodes)
plt.ylim(steps_average.min(), steps_average.max())
plt.title("Average iterations to solve the experiment over runs")
plt.legend()
plt.show()
print("The Minimum number of iterations used to solve the experiment were: {0}\n".format(np.array(all_steps).max()))
print("The Maximum number of iterations used to solve the experiment were: {0}\n".format(np.array(all_steps).min()))
```

```python
rewards_average = np.mean(all_rewards, axis=0)
plt.plot(rewards_average, label = 'Average Reward')
plt.xlabel("Episodes")
plt.ylabel("Sum of\n rewards\n during\n episode" ,rotation=0, labelpad=40)
plt.xlim(-0.2, num_episodes)
plt.ylim(rewards_average.min(), rewards_average.max())
plt.title("Average iterations to solve the experiment over runs")
plt.legend()
plt.show()
print("The best reward obtained solving the experiment was: {0}\n".format(np.array(all_rewards).max()))
print("The Wordt reward obtained solving the experiment was: {0}\n".format(np.array(all_rewards).min()))
```

## Using the last trained Agent

This lines shows in a video the performance of the last trained agent and save a video with the results.

```python
# Test Sarsa Agent
num_runs = 1
num_episodes = 1000

# Environment
env_to_wrap = gym.make('MountainCar-v0')
# Maximum number of possible iterations (default was 200)
env_to_wrap._max_episode_steps = 1500
env = Monitor(env_to_wrap, "./videos/mountainCar", video_callable=lambda episode_id: True, force=True)

# Number of runs are the times the experiment will start again (a.k.a episode)
for n_runs in tqdm(range(num_runs)):

    # Resets environment
    observation = env.reset()
    # Generate last state and action in the agent
    last_action = agent.agent_start(observation)

    # Times the environment will start again without resetting the agent
    for t in tqdm(range(num_episodes)):

        # View environment
        env.render()

        # Take a step with the environment
        observation, reward, done, info = env.step(last_action)

        # If the goal has been reached stop
        if done:
            # Last step with the agent
            agent.agent_end(reward)
            break

        else:
            # Take a step with the agent
            last_action = agent.agent_step(reward, observation)

env.close()
env_to_wrap.close()

print("Episode finished after {} timesteps".format(t+1))
```

      0%|          | 0/1 [00:00<?, ?it/s]
      0%|          | 0/1000 [00:00<?, ?it/s]
      1%|▏         | 13/1000 [00:00<00:07, 124.86it/s]
      3%|▎         | 26/1000 [00:00<00:07, 123.89it/s]
      4%|▍         | 39/1000 [00:00<00:07, 123.81it/s]
      5%|▌         | 52/1000 [00:00<00:07, 123.93it/s]
      6%|▋         | 65/1000 [00:00<00:07, 123.64it/s]
      8%|▊         | 79/1000 [00:00<00:07, 126.19it/s]
      9%|▉         | 92/1000 [00:00<00:07, 127.29it/s]
     10%|█         | 105/1000 [00:00<00:07, 126.48it/s]
     12%|█▏        | 119/1000 [00:00<00:06, 128.02it/s]
     13%|█▎        | 132/1000 [00:01<00:06, 128.56it/s]
     14%|█▍        | 145/1000 [00:01<00:06, 127.71it/s]
     16%|█▌        | 158/1000 [00:01<00:06, 126.67it/s]
     17%|█▋        | 171/1000 [00:01<00:06, 125.91it/s]
     18%|█▊        | 184/1000 [00:01<00:06, 125.88it/s]
     20%|█▉        | 197/1000 [00:01<00:06, 124.24it/s]
     22%|██▏       | 221/1000 [00:01<00:06, 125.14it/s]
    100%|██████████| 1/1 [00:02<00:00,  2.15s/it]

    Episode finished after 222 timesteps

## Plotting the Action-Values of the agent

This final plot aims to show the action-values learned by the agent with Sarsa. The action value for a given state was calculated using: -$max_a\hat{q}(s, a, w)$

```python
# Resolution
values = 500
# Vector of positions
pos_vals = np.linspace(-1.2, 0.5, num = values)
# Vector of velocities
vel_vals = np.linspace(-0.07, 0.07, num = values)

# Z grid values
av_grid = np.zeros((values, values))

# Compute Action-values for each pos - vel pair
for ix in range(len(pos_vals)):
    for iy in range(len(vel_vals)):
        av_grid[ix][iy] = -1 * agent.return_action_value([pos_vals[ix], vel_vals[iy]])
```

```python
# Plot the 3D surface
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
Px, Vy = np.meshgrid(pos_vals, vel_vals)
ax.plot_surface(Vy, Px, av_grid, color = 'gray')
ax.set_title("Cost-to-go function learned", y = 1.1)
ax.set_xlabel('Velocity')
ax.set_ylabel('Position')
ax.set_zlabel('Iterations')
ax.view_init(45, azim=30)
plt.tight_layout()
plt.show()
```

```python

```