## Calculate local finite-amplitude wave activity
**22 March 2022 by MiniUFO**

---
[TOC]

---
### 1. Introduction
Wave play an important role of rearranging angular momentum in the atmosphere.  This is expressed by the generalized Eliassen-Palm (E-P) relation:
$$
\begin{align}
\frac{\partial A}{\partial t}+\nabla\cdot\mathbf F=D+O(\alpha^3) \tag{1}
\end{align}
$$
where $A$ is the density of wave activity (negative angular pseudomomentum), $\mathbf F$ the generalized E-P flux for a slowly modulated **small-amplitude** wave, and $D$ the nonconservative effects on wave density.  For small-amplitude wave, $A$ is given as:
$$
\begin{align}
A=\frac{1}{2}\frac{\overline{q'^2}}{\partial \overline{q}/\partial y} \tag{2}
\end{align}
$$

[Nakamura and Zhu (2010)](https://journals.ametsoc.org/view/journals/atsc/67/9/2010jas3432.1.xml) have proposed a new definition of wave activity to account for **finite-amplitude wave** (FAWA) and [Huang and Nakamura (2016)](https://journals.ametsoc.org/view/journals/atsc/73/1/jas-d-15-0194.1.xml) extended FAWA to be defined locally, which is called local finite-amplitude wave activity (LWA).  A simple zonal mean of LWA recovers FAWA exactly.  The key to LWA is to first sorting the PV in the meridional direction conservatively and obtain a sorted PV profile instead of a traditional zonal mean state to define eddy portion.  After some email discussions with Prof. Nakamura, I recently realize that **LWA defined in the horizontal $x-y$ plane is conceptually equivalent to local available potential energy (APE) in the vertical $x-z$ plane** (see this [notebook for local APE]()).  Here we will demonstrate how to calculate LWA under the help of [xcontour](https://github.com/miniufo/xcontour).  It will show that the whole calculation procedure is very clean and clear.

---

### 2. Theoretical guide
The LWA is defined, according to [Huang and Nakamura (2016)](https://journals.ametsoc.org/view/journals/atsc/73/1/jas-d-15-0194.1.xml), as:
$$
\begin{align}
\tilde A\left(x, y, t\right)=-\int_0^{\eta\left(x,y,t\right)}q\left(x,y+y',t\right)-Q\left(y,t\right)dy' \tag{3}
\end{align}
$$
where $\eta(x,y,t)$ represents the displacement of a tracer contour from its sorted place.

---

### 3. Calculation procedure
#### 3.1 Load data
Here for demonstrating purpose, we repeat the result of [this notebooke from Clare Huang's github repo](https://github.com/csyhuang/hn2016_falwa/blob/master/examples/simple/Example_barotropic.ipynb).  The data, barotropic (absolute) vorticity over the globe, are also downloaded from there.

First, we read in the data and calculate the conservative sorted state:

```python
import xarray as xr
import numpy as np
import sys
sys.path.append('../../')
from xcontour.xcontour import Contour2D, latitude_lengths_at, add_latlon_metrics

##################################  load data  #######################################
dset = xr.open_dataset('../Data/barotropic_vorticity.nc')

# add metrics for xgcm.Grid
dset, grid = add_latlon_metrics(dset)
# get vorticity as a tracer
tracer = dset.absolute_vorticity

###########################  calculate sorted state  #################################
# Initialize equally-spaced contours from minimum value to maximum
# value (within lat/lon dims).  Here will implicitly loop over each isentropic level
%matplotlib inline
N  = 121           # increase the contour number may get non-monotonic A(q) relation
increase = True    # Y-index increases with latitude (sometimes not)
lt = True          # northward of PV contours (larger than) is inside the contour
                   # change this should not change the result of Keff, but may alter
                   # the values at boundaries
dtype = np.float32 # use float32 to save memory
undef = -9.99e8    # for maskout topography if present

# initialize a Contour2D analysis class using grid and tracer
analysis = Contour2D(grid, tracer,
                     dims={'X':'longitude','Y':'latitude'},
                     dimEq={'Y':'latitude'},
                     increase=increase,
                     lt=lt)
# evenly-spaced contours
ctr = analysis.cal_contours(N)

# Mask for A(q) relation table.
mask = xr.where(tracer!=undef, 1, 0).astype(dtype)

# calculate related quantities for
table   = analysis.cal_area_eqCoord_table(mask) # A(Yeq) table
area    = analysis.cal_integral_within_contours(ctr).rename('intArea')
latEq   = table.lookup_coordinates(area).rename('latEq')

#####################  interpolat to equivalent latitudes  ###########################
# combined the results
ds_contour = xr.merge([ctr, area, latEq])

# interpolate from contour space to original latitudes
preLats = dset.latitude.astype(dtype)
# results in latEq space
ds_latEq = analysis.interp_to_dataset(preLats, latEq, ds_contour)
```

    case 1: increase & lt

When we get the sorted profile of vorticity, we can calculate the LWA.

```python
#%% calculate local finite-amplitude wave activity
lwa, contours, masks = analysis.cal_local_wave_activity(tracer,     # the original vorticiy
                                       ds_latEq.absolute_vorticity, # the sorted vorticity
                                       mask_idx=[37,125,170,213])   # select these mask to see the area of LWA
```

Now we can take a look at the results.

```python
import proplot as pplt

fig, axes = pplt.subplots(nrows=2, ncols=2, figsize=(9, 6), sharex=0, sharey=3)

fontsize = 12

ax = axes[0, 0]
m1=ax.pcolormesh(tracer*1e5, levels=np.linspace(-17, 17, 69), cmap='jet')
ax.colorbar(m1, loc='r', ticks=4, label='')
ax.set_title('absolute vorticity $q\\times 10^5$', fontsize=fontsize)
ax.set_xlabel('longitude', fontsize=fontsize-1)
ax.set_ylabel('latitude', fontsize=fontsize-1)
ax.set_xticks([0, 60, 120, 180, 240, 300, 360])
ax.set_yticks([-90, -60, -30, 0, 30, 60, 90])

ax = axes[0, 1]
m1=ax.pcolormesh(lwa, levels=np.linspace(0, 28, 57), cmap='viridis')
ax.colorbar(m1, loc='r', ticks=4, label='')
ax.set_title('LWA $\\tilde{A}$', fontsize=fontsize)
ax.set_xlabel('longitude', fontsize=fontsize-1)
ax.set_ylabel('latitude', fontsize=fontsize-1)
ax.set_xticks([0, 60, 120, 180, 240, 300, 360])
ax.set_yticks([-90, -60, -30, 0, 30, 60, 90])

ax = axes[1, 0]
m1=ax.contourf(sum(masks).where(sum(masks)!=0), cmap='bwr')
ax.contour(tracer, levels=xr.concat(contours, 'latitude').values, color='k', lw=0.6)
ax.colorbar(m1, loc='r', ticks=1, label='')
ax.set_title('masks for 4 contours of LWA $\\tilde{A}$', fontsize=fontsize)
ax.set_xlabel('longitude', fontsize=fontsize-1)
ax.set_ylabel('latitude', fontsize=fontsize-1)
ax.set_xticks([0, 60, 120, 180, 240, 300, 360])
ax.set_yticks([-90, -60, -30, 0, 30, 60, 90])

ax = axes[1, 1]
m1=ax.plot(tracer.mean('longitude')*1e5, tracer.latitude, label='zonal mean $\\overline{q}$')
m2=ax.plot(ds_latEq.absolute_vorticity*1e5, tracer.latitude, label='sorted $Q$')
m3=ax.plot(lwa.mean('longitude'), tracer.latitude, label='FAWA $A$')
ax.set_title('meridional profiles of $\\overline{q}$, $Q$, and $A$', fontsize=fontsize)
ax.set_xlabel('vorticity', fontsize=fontsize-1)
ax.set_ylabel('latitude', fontsize=fontsize-1)
ax.set_xticks([-12, -8, -4, 0, 4, 8, 12])
ax.legend([m1,m2,m3], loc='lr', ncols=1, fontsize=fontsize-1)

axes.format(abc='(a)')
```

It is clear that the vorticity field in (a) is characterized by six Kelvin cat's eyes.  The maxima of LWA in (b) is associated with these eyes.  The conservatively-sorted vorticity profile $Q$ in (d) is monotonic increasing with latitude, indicating barotropic stable profile.  In contrast, the zonal-mean profile $\overline{q}$ does not.  It has two obvious vorticity inversion at about 30°N and 60°N, where wave goes to finite amplitude.  The zonal-mean of LWA, that is FAWA, has two peaks at Northern Hemisphere (60°N, 30°N), a peak at equator, and a very weak one in Southern Hemisphere (62°S).  Four contours related to these four peaks and assoicated masks are also plotted in (c) for a clear view of the area (shading) used for calculating LWA.

#### References
1. Nakamura, N., and D. Zhu, 2010: [Finite-amplitude wave activity and diffusive flux of potential vorticity in eddy-mean flow interaction](https://journals.ametsoc.org/view/journals/atsc/67/9/2010jas3432.1.xml). *J. Atmos. Sci.*, **67**, 2701-2716.
2. Huang, C. S. Y., and N. Nakamura, 2016: [Local finite-amplitude wave activity as a diagnostic of anomalous weather events](https://journals.ametsoc.org/view/journals/atsc/73/1/jas-d-15-0194.1.xml). *J. Atmos. Sci.*, **73**, 211-229.
3. Clare Huang's github repo: https://github.com/csyhuang/hn2016_falwa


# Implementation

[Video here](mov-wave/mayavi/wave2D_u0_gaussian/movie.ogg)

We shall now describe in detail various Python implementations
for solving a standard 2D, linear wave equation with constant
wave velocity and $u=0$ on the
boundary. The wave equation is to be solved
in the space-time domain $\Omega\times (0,T]$,
where $\Omega = (0,L_x)\times (0,L_y)$ is a rectangular spatial
domain. More precisely,
the complete initial-boundary value problem is defined by

$$
\begin{equation}
u_{tt} = c^2(u_{xx} + u_{yy}) + f(x,y,t),\quad (x,y)\in \Omega,\ t\in (0,T],
\label{_auto1} \tag{1}
\end{equation}
$$

$$
\begin{equation}
u(x,y,0) = I(x,y),\quad (x,y)\in\Omega,
\label{_auto2} \tag{2}
\end{equation}
$$

$$
\begin{equation}
u_t(x,y,0) = V(x,y),\quad (x,y)\in\Omega,
\label{_auto3} \tag{3}
\end{equation}
$$

$$
\begin{equation}
u = 0,\quad (x,y)\in\partial\Omega,\ t\in (0,T],
\label{_auto4} \tag{4}
\end{equation}
$$

where $\partial\Omega$ is the boundary of $\Omega$, in this case
the four sides of the rectangle $\Omega = [0,L_x]\times [0,L_y]$:
$x=0$, $x=L_x$, $y=0$, and $y=L_y$.

The PDE is discretized as

$$
[D_t D_t u = c^2(D_xD_x u + D_yD_y u) + f]^n_{i,j},
$$

which leads to an explicit updating formula to be implemented in a
program:

$$
u^{n+1}_{i,j} = -u^{n-1}_{i,j} + 2u^n_{i,j} + \nonumber
$$

$$
\begin{equation}
\quad C_x^2(
u^{n}_{i+1,j} - 2u^{n}_{i,j} + u^{n}_{i-1,j}) + C_y^2
(u^{n}_{i,j+1} - 2u^{n}_{i,j} + u^{n}_{i,j-1}) + \Delta t^2 f_{i,j}^n,
\label{wave:2D3D:impl1:2Du0:ueq:discrete} \tag{5}
\end{equation}
$$

for all interior mesh points $i\in\seti{\mathcal{I}_x}$ and
$j\in\seti{\Iy}$, for $n\in\setr{\mathcal{I}_t}$.
The constants $C_x$ and $C_y$ are defined as

$$
C_x = c\frac{\Delta t}{\Delta x},\quad C_y = c\frac{\Delta t}{\Delta y}
\thinspace .
$$

At the boundary, we simply set $u^{n+1}_{i,j}=0$ for
$i=0$, $j=0,\ldots,N_y$; $i=N_x$, $j=0,\ldots,N_y$;
$j=0$, $i=0,\ldots,N_x$; and $j=N_y$, $i=0,\ldots,N_x$.
For the first step, $n=0$, ([5](#wave:2D3D:impl1:2Du0:ueq:discrete))
is combined with the discretization of the initial condition $u_t=V$,
$[D_{2t} u = V]^0_{i,j}$ to obtain a special formula for
$u^1_{i,j}$ at the interior mesh points:

$$
u^{1}_{i,j} = u^0_{i,j} + \Delta t V_{i,j} + \nonumber
$$

$$
\quad {\frac{1}{2}}C_x^2(
u^{0}_{i+1,j} - 2u^{0}_{i,j} + u^{0}_{i-1,j}) + {\frac{1}{2}}C_y^2
(u^{0}_{i,j+1} - 2u^{0}_{i,j} + u^{0}_{i,j-1}) +\nonumber
$$

$$
\begin{equation}
\quad \frac{1}{2}\Delta t^2f_{i,j}^n,
\label{wave:2D3D:impl1:2Du0:ueq:n0:discrete} \tag{6}
\end{equation}
$$

The algorithm is very similar to the one in 1D:

1. Set initial condition $u^0_{i,j}=I(x_i,y_j)$

2. Compute $u^1_{i,j}$ from ([5](#wave:2D3D:impl1:2Du0:ueq:discrete))

3. Set $u^1_{i,j}=0$ for the boundaries $i=0,N_x$, $j=0,N_y$

4. For $n=1,2,\ldots,N_t$:

a. Find $u^{n+1}_{i,j}$ from ([5](#wave:2D3D:impl1:2Du0:ueq:discrete))
    for all internal mesh points, $i\in\seti{\mathcal{I}_x}$, $j\in\seti{\Iy}$

b. Set $u^{n+1}_{i,j}=0$ for the boundaries $i=0,N_x$, $j=0,N_y$

## Scalar computations

The `solver` function for a 2D case with constant wave velocity and
boundary condition $u=0$ is analogous to the 1D case with similar parameter
values (see `wave1D_u0.py`), apart from a few necessary
extensions. The code is found in the program
[`wave2D_u0.py`](${src_wave}/wave2D_u0/wave2D_u0.py).

### Domain and mesh

The spatial domain is now $[0,L_x]\times [0,L_y]$, specified
by the arguments `Lx` and `Ly`. Similarly, the number of mesh
points in the $x$ and $y$ directions,
$N_x$ and $N_y$, become the arguments `Nx` and `Ny`.
In multi-dimensional problems it makes less sense to specify a
Courant number since the wave velocity is a vector and mesh spacings
may differ in the various spatial directions.
We therefore give $\Delta t$ explicitly. The signature of
the `solver` function is then

```python
def solver(I, V, f, c, Lx, Ly, Nx, Ny, dt, T,
           user_action=None, version='scalar'):
```

Key parameters used in the calculations are created as

```python
x = linspace(0, Lx, Nx+1)                  # mesh points in x dir
y = linspace(0, Ly, Ny+1)                  # mesh points in y dir
dx = x[1] - x[0]
dy = y[1] - y[0]
Nt = int(round(T/float(dt)))
t = linspace(0, N*dt, N+1)                 # mesh points in time
Cx2 = (c*dt/dx)**2;  Cy2 = (c*dt/dy)**2    # help variables
dt2 = dt**2
```

### Solution arrays

We store $u^{n+1}_{i,j}$, $u^{n}_{i,j}$, and
$u^{n-1}_{i,j}$ in three two-dimensional arrays,

```python
u   = zeros((Nx+1,Ny+1))   # solution array
u_n = [zeros((Nx+1,Ny+1)), zeros((Nx+1,Ny+1))]  # t-dt, t-2*dt
```

where $u^{n+1}_{i,j}$ corresponds to `u[i,j]`,
$u^{n}_{i,j}$ to `u_n[i,j]`, and
$u^{n-1}_{i,j}$ to `u_nm1[i,j]`.

### Index sets

mathcal{I}_t is also convenient to introduce the index sets (cf. the section [wave:indexset](#wave:indexset))

```python
mathcal{I}_x = range(0, u.shape[0])
Iy = range(0, u.shape[1])
mathcal{I}_t = range(0, t.shape[0])
```

### Computing the solution

Inserting the initial
condition `I` in `u_n` and making a callback to the user in terms of
the `user_action` function is a straightforward generalization of
the 1D code from the section [wave:string:impl](#wave:string:impl):

```python
for i in mathcal{I}_x:
    for j in Iy:
        u_n[i,j] = I(x[i], y[j])

if user_action is not None:
    user_action(u_n, x, xv, y, yv, t, 0)
```

The `user_action` function has additional arguments compared to the
1D case. The arguments `xv` and `yv` will be commented
upon in the section [Vectorized computations](#wave2D3D:impl:vectorized).

The key finite difference formula ([wave:2D3D:models:unp1](#wave:2D3D:models:unp1))
for updating the solution at
a time level is implemented in a separate function as

```python
def advance_scalar(u, u_n, u_nm1, f, x, y, t, n, Cx2, Cy2, dt2,
                   V=None, step1=False):
    mathcal{I}_x = range(0, u.shape[0]);  Iy = range(0, u.shape[1])
    if step1:
        dt = sqrt(dt2)  # save
        Cx2 = 0.5*Cx2;  Cy2 = 0.5*Cy2; dt2 = 0.5*dt2  # redefine
        D1 = 1;  D2 = 0
    else:
        D1 = 2;  D2 = 1
    for i in mathcal{I}_x[1:-1]:
        for j in Iy[1:-1]:
            u_xx = u_n[i-1,j] - 2*u_n[i,j] + u_n[i+1,j]
            u_yy = u_n[i,j-1] - 2*u_n[i,j] + u_n[i,j+1]
            u[i,j] = D1*u_n[i,j] - D2*u_nm1[i,j] + \
                     Cx2*u_xx + Cy2*u_yy + dt2*f(x[i], y[j], t[n])
            if step1:
                u[i,j] += dt*V(x[i], y[j])
    # Boundary condition u=0
    j = Iy[0]
    for i in mathcal{I}_x: u[i,j] = 0
    j = Iy[-1]
    for i in mathcal{I}_x: u[i,j] = 0
    i = mathcal{I}_x[0]
    for j in Iy: u[i,j] = 0
    i = mathcal{I}_x[-1]
    for j in Iy: u[i,j] = 0
    return u
```

The `step1` variable has been introduced to allow the formula to be
reused for the first step, computing $u^1_{i,j}$:

```python
u = advance_scalar(u, u_n, f, x, y, t,
                   n, Cx2, Cy2, dt, V, step1=True)
```

Below, we will make many alternative implementations of the
`advance_scalar` function to speed up the code since most of
the CPU time in simulations is spent in this function.

**Remark: How to use the solution.**

The `solver` function in the `wave2D_u0.py` code
updates arrays for the next time step by switching references as
described in the section [wave:pde1:impl:ref:switch](#wave:pde1:impl:ref:switch). Any use of `u` on the
user's side is assumed to take place in the user action function. However,
should the code be changed such that `u` is returned and used as solution,
have in mind that you must return `u_n` after the time limit, otherwise
a `return u` will actually return `u_nm1` (due to the switching of array
indices in the loop)!

## Vectorized computations

The scalar code above turns out to be extremely slow for large 2D
meshes, and probably useless in 3D beyond debugging of small test cases.
Vectorization is therefore a must for multi-dimensional
finite difference computations in Python. For example,
with a mesh consisting of $30\times 30$ cells, vectorization
brings down the CPU time by a factor of 70 (!). Equally important,
vectorized code can also easily be parallelized to take (usually)
optimal advantage of parallel computer platforms.

In the vectorized case, we must be able to evaluate user-given
functions like $I(x,y)$ and $f(x,y,t)$ for the entire mesh in one
operation (without loops). These user-given functions are provided as
Python functions `I(x,y)` and `f(x,y,t)`, respectively.  Having the
one-dimensional coordinate arrays `x` and `y` is not sufficient when
calling `I` and `f` in a vectorized way.  We must extend `x` and `y`
to their vectorized versions `xv` and `yv`:

```python
from numpy import newaxis
xv = x[:,newaxis]
yv = y[newaxis,:]
# or
xv = x.reshape((x.size, 1))
yv = y.reshape((1, y.size))
```

This is a standard required technique when evaluating functions over
a 2D mesh, say `sin(xv)*cos(xv)`, which then gives a result with shape
`(Nx+1,Ny+1)`. Calling `I(xv, yv)` and `f(xv, yv, t[n])` will now
return `I` and `f` values for the entire set of mesh points.

With the `xv` and `yv` arrays for vectorized computing,
setting the initial condition is just a matter of

```python
u_n[:,:] = I(xv, yv)
```

One could also have written `u_n = I(xv, yv)` and let `u_n` point to a
new object, but vectorized operations often make use of direct
insertion in the original array through `u_n[:,:]`, because sometimes
not all of the array is to be filled by such a function
evaluation. This is the case with the computational scheme for
$u^{n+1}_{i,j}$:

```python
def advance_vectorized(u, u_n, u_nm1, f_a, Cx2, Cy2, dt2,
                       V=None, step1=False):
    if step1:
        dt = np.sqrt(dt2)  # save
        Cx2 = 0.5*Cx2;  Cy2 = 0.5*Cy2; dt2 = 0.5*dt2  # redefine
        D1 = 1;  D2 = 0
    else:
        D1 = 2;  D2 = 1
    u_xx = u_n[:-2,1:-1] - 2*u_n[1:-1,1:-1] + u_n[2:,1:-1]
    u_yy = u_n[1:-1,:-2] - 2*u_n[1:-1,1:-1] + u_n[1:-1,2:]
    u[1:-1,1:-1] = D1*u_n[1:-1,1:-1] - D2*u_nm1[1:-1,1:-1] + \
                   Cx2*u_xx + Cy2*u_yy + dt2*f_a[1:-1,1:-1]
    if step1:
        u[1:-1,1:-1] += dt*V[1:-1, 1:-1]
    # Boundary condition u=0
    j = 0
    u[:,j] = 0
    j = u.shape[1]-1
    u[:,j] = 0
    i = 0
    u[i,:] = 0
    i = u.shape[0]-1
    u[i,:] = 0
    return u
```

Array slices in 2D are more complicated to understand than those in
1D, but the logic from 1D applies to each dimension separately.
For example, when doing $u^{n}_{i,j} - u^{n}_{i-1,j}$ for $i\in\setr{\mathcal{I}_x}$,
we just keep `j` constant and make a slice in the first index:
`u_n[1:,j] - u_n[:-1,j]`, exactly as in 1D. The `1:` slice
specifies all the indices $i=1,2,\ldots,N_x$ (up to the last
valid index),
while `:-1` specifies the relevant indices for the second term:
$0,1,\ldots,N_x-1$ (up to, but not including the last index).

In the above code segment, the situation is slightly more complicated,
because each displaced slice in one direction is
accompanied by a `1:-1` slice in the other direction. The reason is
that we only work with the internal points for the index that is
kept constant in a difference.

The boundary conditions along the four sides make use of
a slice consisting of all indices along a boundary:

```python
u[: ,0] = 0
u[:,Ny] = 0
u[0 ,:] = 0
u[Nx,:] = 0
```

In the vectorized update of `u` (above), the function `f` is first computed
as an array over all mesh points:

        f_a = f(xv, yv, t[n])

We could, alternatively, have used the call `f(xv, yv, t[n])[1:-1,1:-1]`
in the last term of the update statement, but other implementations
in compiled languages benefit from having `f` available in an array
rather than calling our Python function `f(x,y,t)` for
every point.

Also in the `advance_vectorized` function we have introduced a
boolean `step1` to reuse the formula for the first time step
in the same way as we did with `advance_scalar`.
We refer to the `solver` function in `wave2D_u0.py`
for the details on how the overall algorithm is implemented.

The callback function now has the arguments
`u, x, xv, y, yv, t, n`. The inclusion of `xv` and `yv` makes it
easy to, e.g., compute an exact 2D solution in the callback function
and compute errors, through an expression like
`u - u_exact(xv, yv, t[n])`.

## Verification

### Testing a quadratic solution

The 1D solution from the section [wave:pde2:fd:verify:quadratic](#wave:pde2:fd:verify:quadratic) can be
generalized to multi-dimensions and provides a test case where the
exact solution also fulfills the discrete equations, such that we know
(to machine precision) what numbers the solver function should
produce. In 2D we use the following generalization of
([wave:pde2:fd:verify:quadratic:uex](#wave:pde2:fd:verify:quadratic:uex)):

$$
\begin{equation}
\uex(x,y,t) = x(L_x-x)y(L_y-y)(1+{\frac{1}{2}}t)
\thinspace .
\label{wave2D3D:impl:verify:quadratic} \tag{7}
\end{equation}
$$

This solution fulfills the PDE problem if $I(x,y)=\uex(x,y,0)$,
$V=\frac{1}{2}\uex(x,y,0)$, and $f=2c^2(1+{\frac{1}{2}}t)(y(L_y-y) +
x(L_x-x))$. To show that $\uex$ also solves the discrete equations,
we start with the general results $[D_t D_t 1]^n=0$, $[D_t D_t t]^n=0$,
and $[D_t D_t t^2]=2$, and use these to compute

$$
\begin{align*}
[D_xD_x \uex]^n_{i,j} &= [y(L_y-y)(1+{\frac{1}{2}}t) D_xD_x x(L_x-x)]^n_{i,j}\\
&= y_j(L_y-y_j)(1+{\frac{1}{2}}t_n)(-2)\thinspace .
\end{align*}
$$

A similar calculation must be carried out for the $[D_yD_y
\uex]^n_{i,j}$ and $[D_tD_t \uex]^n_{i,j}$ terms.  One must also show
that the quadratic solution fits the special formula for
$u^1_{i,j}$. The details are left as [Exercise 1: Check that a solution fulfills the discrete model](#wave:exer:quadratic:2D).
The `test_quadratic` function in the
[`wave2D_u0.py`](${src_wave}/wave2D_u0/wave2D_u0.py)
program implements this verification as a proper test function
for the pytest and nose frameworks.

## Visualization

Eventually, we are ready for a real application with our code!
Look at the `wave2D_u0.py` and the `gaussian` function. mathcal{I}_t
starts with a Gaussian function to see how it propagates in a square
with $u=0$ on the boundaries:

```python
def gaussian(plot_method=2, version='vectorized', save_plot=True):
    """
    Initial Gaussian bell in the middle of the domain.
    plot_method=1 applies mesh function,
    =2 means surf, =3 means Matplotlib, =4 means mayavi,
    =0 means no plot.
    """
    # Clean up plot files
    for name in glob('tmp_*.png'):
        os.remove(name)

    Lx = 10
    Ly = 10
    c = 1.0

    from numpy import exp

    def I(x, y):
        """Gaussian peak at (Lx/2, Ly/2)."""
        return exp(-0.5*(x-Lx/2.0)**2 - 0.5*(y-Ly/2.0)**2)

    def plot_u(u, x, xv, y, yv, t, n):
        """User action function for plotting."""
	...

    Nx = 40; Ny = 40; T = 20
    dt, cpu = solver(I, None, None, c, Lx, Ly, Nx, Ny, -1, T,
                     user_action=plot_u, version=version)
```

### Matplotlib

We want to animate a 3D surface in Matplotlib, but this is a really
slow process and not recommended, so we consider Matplotlib not an
option as long as on-screen animation is desired. One can use the
recipes for single shots of $u$, where it does produce high-quality
3D plots.

### Gnuplot

Let us look at different ways for visualization. We import SciTools as
`st` and can access `st.mesh` and `st.surf` in Matplotlib or Gnuplot,
but this is not supported except for the Gnuplot package, where it
works really well ([Figure](#wave2D3D:impl:viz:fig:gnuplot1)).  Then
we choose `plot_method=2` (or less relevant `plot_method=1`) and force
the backend for SciTools to be Gnuplot (if you have the C package
Gnuplot and the `Gnuplot.py` Python interface module installed):

        Terminal> python wave2D_u0.py --SCITOOLS_easyviz_backend gnuplot

mathcal{I}_t gives a nice visualization with lifted surface and contours beneath.
[Figure](#wave2D3D:impl:viz:fig:gnuplot1) shows four plots of $u$.

Snapshots of the surface plotted by Gnuplot.

Video files can be made of the PNG frames:

        Terminal> ffmpeg -i tmp_%04d.png -r 25 -vcodec flv movie.flv
        Terminal> ffmpeg -i tmp_%04d.png -r 25 -vcodec linx264 movie.mp4
        Terminal> ffmpeg -i tmp_%04d.png -r 25 -vcodec libvpx movie.webm
        Terminal> ffmpeg -i tmp_%04d.png -r 25 -vcodec libtheora movie.ogg

mathcal{I}_t is wise to use a high frame rate - a low one will just skip many
frames. There may also be considerable quality differences between the
different formats.

```python
from IPython.display import HTML
_s = """
<div>

</div>
<p><em></em></p>

<!-- Issue warning if in a Safari browser -->

"""
HTML(_s)
```

### Mayavi

The best option for doing visualization of 2D and 3D scalar and vector fields
in Python programs is Mayavi, which is an interface to the high-quality
package VTK in C++. There is good online documentation and also
an introduction in Chapter 5 of [[Langtangen_2012]](#Langtangen_2012).

To obtain Mayavi on Ubuntu platforms you can write

        pip install mayavi --upgrade

For Mac OS X and Windows, we recommend using Anaconda.
To obtain Mayavi for Anaconda you can write

        conda install mayavi

Mayavi has a MATLAB-like interface called `mlab`. We can do

```python
import mayavi.mlab as plt
# or
from mayavi import mlab
```

and have `plt` (as usual) or `mlab`
as a kind of MATLAB visualization access inside our program (just
more powerful and with higher visual quality).

The official documentation of the `mlab` module is provided in two
places, one for the [basic functionality](http://docs.enthought.com/mayavi/mayavi/auto/mlab_helper_functions.html)
and one for [further functionality](http://docs.enthought.com/mayavi/mayavi/auto/mlab_other_functions.html).
Basic [figure
handling](http://docs.enthought.com/mayavi/mayavi/auto/mlab_figure.html)
is very similar to the one we know from Matplotlib.  Just as for
Matplotlib, all plotting commands you do in `mlab` will go into the
same figure, until you manually change to a new figure.

Back to our application, the following code for the user action
function with plotting in Mayavi is relevant to add.

```python
%matplotlib inline

# Top of the file
try:
    import mayavi.mlab as mlab
except:
    # We don't have mayavi
    pass

def solver(...):
    ...

def gaussian(...):
    ...
    if plot_method == 3:
        from mpl_toolkits.mplot3d import axes3d
        import matplotlib.pyplot as plt
        from matplotlib import cm
        plt.ion()
        fig = plt.figure()
        u_surf = None

    def plot_u(u, x, xv, y, yv, t, n):
        """User action function for plotting."""
        if t[n] == 0:
            time.sleep(2)
        if plot_method == 1:
            # Works well with Gnuplot backend, not with Matplotlib
            st.mesh(x, y, u, title='t=%g' % t[n], zlim=[-1,1],
                    caxis=[-1,1])
        elif plot_method == 2:
            # Works well with Gnuplot backend, not with Matplotlib
            st.surfc(xv, yv, u, title='t=%g' % t[n], zlim=[-1, 1],
                  colorbar=True, colormap=st.hot(), caxis=[-1,1],
                  shading='flat')
        elif plot_method == 3:
            print 'Experimental 3D matplotlib...not recommended'
        elif plot_method == 4:
	    # Mayavi visualization
            mlab.clf()
            extent1 = (0, 20, 0, 20,-2, 2)
            s = mlab.surf(x , y, u,
                          colormap='Blues',
                          warp_scale=5,extent=extent1)
            mlab.axes(s, color=(.7, .7, .7), extent=extent1,
                      ranges=(0, 10, 0, 10, -1, 1),
                      xlabel='', ylabel='', zlabel='',
                      x_axis_visibility=False,
                      z_axis_visibility=False)
            mlab.outline(s, color=(0.7, .7, .7), extent=extent1)
            mlab.text(6, -2.5, '', z=-4, width=0.14)
            mlab.colorbar(object=None, title=None,
                          orientation='horizontal',
                          nb_labels=None, nb_colors=None,
                          label_fmt=None)
            mlab.title('Gaussian t=%g' % t[n])
            mlab.view(142, -72, 50)
            f = mlab.gcf()
            camera = f.scene.camera
            camera.yaw(0)

        if plot_method > 0:
            time.sleep(0) # pause between frames
            if save_plot:
                filename = 'tmp_%04d.png' % n
		if plot_method == 4:
                    mlab.savefig(filename)  # time consuming!
		elif plot_method in (1,2):
                    st.savefig(filename)  # time consuming!
```

This is a point to get started - visualization is as always a very
time-consuming and experimental discipline. With the PNG files we
can use `ffmpeg` to create videos.

Plot with Mayavi.

```python
# from IPython.display import HTML

# _s = """
# <div>
#
# </div>
# <p><em></em></p>

# <!-- Issue warning if in a Safari browser -->
#

# """
# #     <source src='https://github.com/hplgit/fdm-book/blob/master/doc/pub/book/html/mov-wave/mayavi/wave2D_u0_gaussian/movie.ogg'  type='video/ogg;  codecs="theora, vorbis"'>
# HTML(_s)
```

[Video here](mov-wave/mayavi/wave2D_u0_gaussian/movie.mp4)

# Exercises

## Exercise 1: Check that a solution fulfills the discrete model

Carry out all mathematical details to show that
([7](#wave2D3D:impl:verify:quadratic)) is indeed a solution of the
discrete model for a 2D wave equation with $u=0$ on the boundary.
One must check the boundary conditions, the initial conditions,
the general discrete equation at a time level and the special
version of this equation for the first time level.
Filename: `check_quadratic_solution`.

## Project 2: Calculus with 2D mesh functions

The goal of this project is to redo
[wave:exer:mesh1D:calculus](#wave:exer:mesh1D:calculus) with 2D
mesh functions ($f_{i,j}$).

**Differentiation.**
The differentiation results in a discrete gradient
function, which in the 2D case can be represented by a three-dimensional
array `df[d,i,j]` where `d` represents the direction of
the derivative, and `i,j` is a mesh point in 2D.
Use centered differences for
the derivative at inner points and one-sided forward or backward
differences at the boundary points. Construct unit tests and
write a corresponding test function.

**Integration.**
The integral of a 2D mesh function $f_{i,j}$ is defined as

$$
F_{i,j} = \int_{y_0}^{y_j} \int_{x_0}^{x_i} f(x,y)dxdy,
$$

where $f(x,y)$ is a function that takes on the values of the
discrete mesh function $f_{i,j}$ at the mesh points, but can also
be evaluated in between the mesh points. The particular variation
between mesh points can be taken as bilinear, but this is not
important as we will use a product Trapezoidal rule to approximate
the integral over a cell in the mesh and then we only need to
evaluate $f(x,y)$ at the mesh points.

Suppose $F_{i,j}$ is computed. The calculation of $F_{i+1,j}$
is then

$$
\begin{align*}
F_{i+1,j} &= F_{i,j} + \int_{x_i}^{x_{i+1}}\int_{y_0}^{y_j} f(x,y)dydx\\
& \approx \Delta x \frac{1}{2}\left(
\int_{y_0}^{y_j} f(x_{i},y)dy
+ \int_{y_0}^{y_j} f(x_{i+1},y)dy\right)
\end{align*}
$$

The integrals in the $y$ direction can be approximated by a Trapezoidal
rule. A similar idea can be used to compute $F_{i,j+1}$. Thereafter,
$F_{i+1,j+1}$ can be computed by adding the integral over the final
corner cell to $F_{i+1,j} + F_{i,j+1} - F_{i,j}$. Carry out the
details of these computations and implement a function that can
return $F_{i,j}$ for all mesh indices $i$ and $j$. Use the
fact that the Trapezoidal rule is exact for linear functions and
write a test function.
Filename: `mesh_calculus_2D`.

## Exercise 3: Implement Neumann conditions in 2D

Modify the [`wave2D_u0.py`](${src_wave}/wave2D_u0/wave2D_u0.py)
program, which solves the 2D wave equation $u_{tt}=c^2(u_{xx}+u_{yy})$
with constant wave velocity $c$ and $u=0$ on the boundary, to have
Neumann boundary conditions: $\partial u/\partial n=0$.
Include both scalar code (for debugging and reference) and
vectorized code (for speed).

To test the code, use $u=1.2$ as solution ($I(x,y)=1.2$, $V=f=0$, and
$c$ arbitrary), which should be exactly reproduced with any mesh
as long as the stability criterion is satisfied.
Another test is to use the plug-shaped pulse
in the `pulse` function from the section [wave:pde2:software](#wave:pde2:software)
and the [`wave1D_dn_vc.py`](${src_wave}/wave1D/wave1D_dn_vc.py)
program. This pulse
is exactly propagated in 1D if $c\Delta t/\Delta x=1$. Check
that also the 2D program can propagate this pulse exactly
in $x$ direction ($c\Delta t/\Delta x=1$, $\Delta y$ arbitrary)
and $y$ direction ($c\Delta t/\Delta y=1$, $\Delta x$ arbitrary).
Filename: `wave2D_dn`.

## Exercise 4: Test the efficiency of compiled loops in 3D

Extend the `wave2D_u0.py` code and the Cython, Fortran, and C versions to 3D.
Set up an efficiency experiment to determine the relative efficiency of
pure scalar Python code, vectorized code, Cython-compiled loops,
Fortran-compiled loops, and C-compiled loops.
Normalize the CPU time for each mesh by the fastest version.
Filename: `wave3D_u0`.


```python
%matplotlib inline
```

# The Johnson-Lindenstrauss bound for embedding with random projections

The `Johnson-Lindenstrauss lemma`_ states that any high dimensional
dataset can be randomly projected into a lower dimensional Euclidean
space while controlling the distortion in the pairwise distances.

Theoretical bounds
==================

The distortion introduced by a random projection `p` is asserted by
the fact that `p` is defining an eps-embedding with good probability
as defined by:

\begin{align}(1 - eps) \|u - v\|^2 < \|p(u) - p(v)\|^2 < (1 + eps) \|u - v\|^2\end{align}

Where u and v are any rows taken from a dataset of shape [n_samples,
n_features] and p is a projection by a random Gaussian N(0, 1) matrix
with shape [n_components, n_features] (or a sparse Achlioptas matrix).

The minimum number of components to guarantees the eps-embedding is
given by:

\begin{align}n\_components >= 4 log(n\_samples) / (eps^2 / 2 - eps^3 / 3)\end{align}

The first plot shows that with an increasing number of samples ``n_samples``,
the minimal number of dimensions ``n_components`` increased logarithmically
in order to guarantee an ``eps``-embedding.

The second plot shows that an increase of the admissible
distortion ``eps`` allows to reduce drastically the minimal number of
dimensions ``n_components`` for a given number of samples ``n_samples``

Empirical validation
====================

We validate the above bounds on the digits dataset or on the 20 newsgroups
text document (TF-IDF word frequencies) dataset:

- for the digits dataset, some 8x8 gray level pixels data for 500
  handwritten digits pictures are randomly projected to spaces for various
  larger number of dimensions ``n_components``.

- for the 20 newsgroups dataset some 500 documents with 100k
  features in total are projected using a sparse random matrix to smaller
  euclidean spaces with various values for the target number of dimensions
  ``n_components``.

The default dataset is the digits dataset. To run the example on the twenty
newsgroups dataset, pass the --twenty-newsgroups command line argument to this
script.

For each value of ``n_components``, we plot:

- 2D distribution of sample pairs with pairwise distances in original
  and projected spaces as x and y axis respectively.

- 1D histogram of the ratio of those distances (projected / original).

We can see that for low values of ``n_components`` the distribution is wide
with many distorted pairs and a skewed distribution (due to the hard
limit of zero ratio on the left as distances are always positives)
while for larger values of n_components the distortion is controlled
and the distances are well preserved by the random projection.

Remarks
=======

According to the JL lemma, projecting 500 samples without too much distortion
will require at least several thousands dimensions, irrespective of the
number of features of the original dataset.

Hence using random projections on the digits dataset which only has 64 features
in the input space does not make sense: it does not allow for dimensionality
reduction in this case.

On the twenty newsgroups on the other hand the dimensionality can be decreased
from 56436 down to 10000 while reasonably preserving pairwise distances.

```python
print(__doc__)

import sys
from time import time
import numpy as np
import matplotlib.pyplot as plt
from sklearn.random_projection import johnson_lindenstrauss_min_dim
from sklearn.random_projection import SparseRandomProjection
from sklearn.datasets import fetch_20newsgroups_vectorized
from sklearn.datasets import load_digits
from sklearn.metrics.pairwise import euclidean_distances

# Part 1: plot the theoretical dependency between n_components_min and
# n_samples

# range of admissible distortions
eps_range = np.linspace(0.1, 0.99, 5)
colors = plt.cm.Blues(np.linspace(0.3, 1.0, len(eps_range)))

# range of number of samples (observation) to embed
n_samples_range = np.logspace(1, 9, 9)

plt.figure()
for eps, color in zip(eps_range, colors):
    min_n_components = johnson_lindenstrauss_min_dim(n_samples_range, eps=eps)
    plt.loglog(n_samples_range, min_n_components, color=color)

plt.legend(["eps = %0.1f" % eps for eps in eps_range], loc="lower right")
plt.xlabel("Number of observations to eps-embed")
plt.ylabel("Minimum number of dimensions")
plt.title("Johnson-Lindenstrauss bounds:\nn_samples vs n_components")

# range of admissible distortions
eps_range = np.linspace(0.01, 0.99, 100)

# range of number of samples (observation) to embed
n_samples_range = np.logspace(2, 6, 5)
colors = plt.cm.Blues(np.linspace(0.3, 1.0, len(n_samples_range)))

plt.figure()
for n_samples, color in zip(n_samples_range, colors):
    min_n_components = johnson_lindenstrauss_min_dim(n_samples, eps=eps_range)
    plt.semilogy(eps_range, min_n_components, color=color)

plt.legend(["n_samples = %d" % n for n in n_samples_range], loc="upper right")
plt.xlabel("Distortion eps")
plt.ylabel("Minimum number of dimensions")
plt.title("Johnson-Lindenstrauss bounds:\nn_components vs eps")

# Part 2: perform sparse random projection of some digits images which are
# quite low dimensional and dense or documents of the 20 newsgroups dataset
# which is both high dimensional and sparse

if '--twenty-newsgroups' in sys.argv:
    # Need an internet connection hence not enabled by default
    data = fetch_20newsgroups_vectorized().data[:500]
else:
    data = load_digits().data[:500]

n_samples, n_features = data.shape
print("Embedding %d samples with dim %d using various random projections"
      % (n_samples, n_features))

n_components_range = np.array([300, 1000, 10000])
dists = euclidean_distances(data, squared=True).ravel()

# select only non-identical samples pairs
nonzero = dists != 0
dists = dists[nonzero]

for n_components in n_components_range:
    t0 = time()
    rp = SparseRandomProjection(n_components=n_components)
    projected_data = rp.fit_transform(data)
    print("Projected %d samples from %d to %d in %0.3fs"
          % (n_samples, n_features, n_components, time() - t0))
    if hasattr(rp, 'components_'):
        n_bytes = rp.components_.data.nbytes
        n_bytes += rp.components_.indices.nbytes
        print("Random matrix with size: %0.3fMB" % (n_bytes / 1e6))

    projected_dists = euclidean_distances(
        projected_data, squared=True).ravel()[nonzero]

    plt.figure()
    plt.hexbin(dists, projected_dists, gridsize=100, cmap=plt.cm.PuBu)
    plt.xlabel("Pairwise squared distances in original space")
    plt.ylabel("Pairwise squared distances in projected space")
    plt.title("Pairwise distances distribution for n_components=%d" %
              n_components)
    cb = plt.colorbar()
    cb.set_label('Sample pairs counts')

    rates = projected_dists / dists
    print("Mean distances rate: %0.2f (%0.2f)"
          % (np.mean(rates), np.std(rates)))

    plt.figure()
    plt.hist(rates, bins=50, normed=True, range=(0., 2.), edgecolor='k')
    plt.xlabel("Squared distances rate: projected / original")
    plt.ylabel("Distribution of samples pairs")
    plt.title("Histogram of pairwise distance rates for n_components=%d" %
              n_components)

    # TODO: compute the expected value of eps and add them to the previous plot
    # as vertical lines / region

plt.show()
```

```python

```


# Trading Lists and Trading Trajectories

### Introduction

[Almgren and Chriss](https://cims.nyu.edu/~almgren/papers/optliq.pdf) provided a solution to the optimal liquidation problem by assuming the that stock prices follow a discrete arithmetic random walk, and that the permanent and temporary market impact functions are linear functions of the trading rate.

Almgren and Chriss showed that for each value of risk aversion there is a unique optimal execution strategy. This optimal execution strategy is determined by a trading trajectory and its associated trading list. The optimal trading trajectory is given by:

\begin{equation}
x_j = \frac{\sinh \left( \kappa \left( T-t_j\right)\right)}{ \sinh (\kappa T)}X, \hspace{1cm}\text{ for } j=0,...,N
\end{equation}

and the associated trading list is given by:

\begin{equation}
n_j = \frac{2 \sinh \left(\frac{1}{2} \kappa \tau \right)}{ \sinh \left(\kappa T\right) } \cosh \left(\kappa \left(T - t_{j-\frac{1}{2}}\right)\right) X, \hspace{1cm}\text{ for } j=1,...,N
\end{equation}

where $t_{j-1/2} = (j-\frac{1}{2}) \tau$.

Given some initial parameters, such as the number of shares, the liquidation time, the trader's risk aversion, etc..., the trading list will tell us how many shares we should sell at each trade to minimize our transaction costs.

In this notebook, we will see how the trading list varies according to some initial trading parameters.

## Visualizing Trading Lists and Trading Trajectories

Let's assume we have 1,000,000 shares that we wish to liquidate. In the code below, we will plot the optimal trading trajectory and its associated trading list for different trading parameters, such as trader's risk aversion, number of trades, and liquidation time.

```python
%matplotlib inline

import matplotlib.pyplot as plt

import utils

# We set the default figure size
plt.rcParams['figure.figsize'] = [17.0, 7.0]

# Set the number of days to sell all shares (i.e. the liquidation time)
l_time = 60

# Set the number of trades
n_trades = 60

# Set the trader's risk aversion
t_risk = 1e-6

# Plot the trading list and trading trajectory. If show_trl = True, the data frame containing the values of the
# trading list and trading trajectory is printed
utils.plot_trade_list(lq_time = l_time, nm_trades = n_trades, tr_risk = t_risk, show_trl = True)
```

# Implementing a Trading List

Once we have the trading list for a given set of initial parameters, we can actually implement it. That is, we can sell our shares in the stock market according to the trading list and see how much money we made or lost. To do this, we are going to simulate the stock market with a simple trading environment. This simulated trading environment uses the same price dynamics and market impact functions as the Almgren and Chriss model. That is, stock price movements evolve according to a discrete arithmetic random walk and the permanent and temporary market impact functions are linear functions of the trading rate. We are going to use the same environment to train our Deep Reinforcement Learning algorithm later on.

We will describe the details of the trading environment in another notebook, for now we will just take a look at its default parameters. We will distinguish between financial parameters, such the annual volatility in stock price, and the parameters needed to calculate the trade list using the Almgren and Criss model, such as the trader's risk aversion.

```python
import utils

# Get the default financial and AC Model parameters
financial_params, ac_params = utils.get_env_param()
```

### Default Financial Parameters

```python
financial_params
```

Financial Parameters

  Annual Volatility:	  12%	   Bid-Ask Spread:    	     0.125

  Daily Volatility:	  0.8%	   Daily Trading Volume:	 5,000,000

### Parameters for the Almgren and Chriss Model

```python
ac_params
```

Almgren and Chriss Model Parameters

  Total Number of Shares to Sell:	                  1,000,000	   Fixed Cost of Selling per Share:	 $0.062

  Starting Price per Share:	                         $50.00	     Trader's Risk Aversion:	           1e-06

  Price Impact for Each 1% of Daily Volume Traded:	 $2.5e-06	    Permanent Impact Constant:	       2.5e-07

  Number of Days to Sell All the Shares:	              60	       Single Step Variance:	             0.144

  Number of Trades:	                                   60	       Time Interval between trades:	      1.0

The code below implements the trading list resulting from different trading parameters, such as trader's risk aversion, number of trades, and liquidation time. All other parameters, such as total number shares to sell, are taken from the simulated trading environment (see above).

```python
%matplotlib inline

import matplotlib.pyplot as plt

import utils

# We set the default figure size
plt.rcParams['figure.figsize'] = [17.0, 7.0]

# Set the random seed
sd = 0

# Set the number of days to sell all shares (i.e. the liquidation time)
l_time = 60

# Set the number of trades
n_trades = 60

# Set the trader's risk aversion
t_risk = 1e-6

# Implement the trading list for the given parameters
utils.implement_trade_list(seed = sd, lq_time = l_time, nm_trades = n_trades, tr_risk = t_risk)
```

```python

```

```python

```