# FastEMRIWaveforms Tutorial
## ICERM Workshop: Waveform acceleration with machine learning and GPUs

### Michael Katz, Max Planck Institute for Gravitational Physics (Albert Einstein Institute), Lead developer for FastEMRIWaveforms

In this tutorial, you will learn the basics of building an accelereated EMRI waveform. We encourage participants to see our paper ([arxiv.org/2008.06071](https://arxiv.org/abs/2008.06071)) and the FastEMRIWaveforms [package documentation](https://bhptoolkit.org/FastEMRIWaveforms/) for more information, as well as our forthcoming paper that will describe the waveform in much more detail than our PRL submission.

Import packages:

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

import h5py

from few.amplitude import romannet
from few.utils.utility import check_for_file_download, p_to_y
from few.trajectory.flux import RunSchwarzEccFluxInspiral
from few.amplitude.romannet import RomanAmplitude
from few.amplitude.interp2dcubicspline import Interp2DAmplitude
from few.waveform import FastSchwarzschildEccentricFlux
```

    Matplotlib is building the font cache; this may take a moment.

## A Quick Introduction to FastEMRIWaveforms

### Collaborators: Alvin Chua, Niels Warburton, Scott Hughes, Lorenzo Speri

```python
generate = FastSchwarzschildEccentricFlux()

M = 1e6  # large mass
mu = 1e1  # small mass
p0 = 12.0  # separation
e0 = 0.4  # eccentricity
theta = np.pi/3.  # polar viewing angle
phi = np.pi/4.  # azimuthal viewing angle
dist = 1.0  # distance in Gpc

T = 1/365.  # in years
dt = 10.0  # time spacing of data stream

wave = generate(M, mu, p0, e0, theta, phi, dist, T=T, dt=dt)

plt.plot(wave.real)
plt.plot(wave.imag)
```

FastEMRIWaveforms is the first fully relativistic template generation tool for extreme mass ratio inspirals. We think of it more as a framework than any specific waveform. In this tutorial, we will discuss the main pieces of this framework and how they relate to our first fully relativistic waveform model shown above.

As a basic primer, our five key points we want to achieve with this framework are:

* Accuracy: Our fast waveforms must be accurate when compared to slow and accurate waveforms generated by the waveform modeling community.

* Modularity: These waveforms are to be built out of a set of modules. These modules are to be easily interchangeable, as well as stand alone tools available for more in-depth analysis.

* Flexibility: This framework must be easily adaptable to new computational methods and/or improvements in EMRI physics.

* Easy User Interface: All modules and complete waveforms are to have a front-facing python interface with clear and extensive documentation, as well as many examples.

* Parallelization: These waveforms must take advantage of parallelization techniques such as OpenMP, as well as accelerator hardware such as GPUs.

## Basics of an EMRI Waveform

The guiding equation for building an EMRI waveform is given by

\begin{equation}
h_+-ih_x = \frac{1}{r}\sum_{lmkn}\left(-\frac{Z_{lmkn}}{\omega_{mkn}^2}\right)\left(S_{lmkn}(\theta)e^{-im\phi}\right)e^{i\Phi_{mkn}} = \frac{1}{r}\sum_{lmkn}A_{lmkn}\Theta_{lmkn}e^{i\Phi_{mkn}}.
\end{equation}

Here we are concerned with generating fast and accurate waveforms. These waveforms are created with a sequence of modules. We will discuss the three main modules used to produce these waveforms. The first module is the Trajectory module which takes initial parameters and produces arrays for all of the evolving quantities of concern. These include the phase evolution, {$\Phi_\varphi,\Phi_\theta, \Phi_r$}, from the start to the end of the EMRI orbit. With these quantities, we can produce:

\begin{equation}
\Phi_{mkn} = m\Phi_\varphi + k\Phi_\theta + n\Phi_r.
\end{equation}

The trajectory also produces orbital quantities over time. These include the separation ($p$), eccentricity ($e$), and the inclination angle of the orbit ($\iota$). With these arrays containing the orbital evolution of these values, we generate the amplitudes, $A_{lmkn}$. While the evolution is over time, the various phasing and amplitude computations are produced within a frequency decomposition into $(l,m,k,n)$ modes.

With phases and amplitdues in hand, we combine this with the angular harmonic information ($\Theta_{lmkn}$) to produce the final waveform. In this step, we calculate the sum of all modes at each time point in the template data stream.

## Overall Waveform Strategy

* We calculate our trajectories and amplitudes as sparse arrays in time by using a large adaptive stepping integrator. This produces arrays with $\sim100$ points. All of these quantities vary slowly and smoothly. After these calculations are complete, we scale this up to the actual data stream cadence.

* The speed of EMRI waveforms is strongly determined by the amount of harmonic content. Higher eccentricities require more modes to produce a high fidelity waveform. In order to make our waveforms as efficient as possible, we perform an online mode content calculation that removes modes from consideration if they do not contribute to the waveform power determined by a user-defined threshold.

## Current Waveform Model: Schwarzschild Eccentric

* No $k$ modes
* Orbit is equatorial
* $S_{lmkn}(\theta)e^{-im\phi}$ reduces to $_{-2}Y_{lm}(\theta,\phi)$
* $l:\{2,10\}$, $m:\{-l,l\}$, $n:\{-30,30\}$ $\rightarrow$ 3843 modes.

## Fast Trajectories: $\{p, e, \Phi_\varphi, \Phi_r\}$

We are not going to spend too much time on the trajectory part. However, we need to generate it in order to build the rest of our waveform. To build the trajectory, we integrate with large steps using an RK8 integrator.

```python
traj = RunSchwarzEccFluxInspiral()
```

```python
p0 = 16.0  # initial separation
e0 = 0.4  # initial eccentricity
mu = 180.  # iniital small mass in solar masses, produces approximately 1 yr waveform
M = 1e6   # initial large mass in solar masses
dt = 10.0 # sets initial step size
T = 1.0 # in years

t, p, e, Phi_phi, Phi_r, flux = traj(M, mu, p0, e0, T=T, dt=dt)
print("length:", len(t), "duration:", t[-1])
```

    length: 122 duration: 29164387.553433623

```python
fig, axes = plt.subplots(2, 3)
plt.subplots_adjust(wspace=0.3)
fig.set_size_inches(14, 8)
axes = axes.ravel()

ylabels = [r'$e$', r'$p$', r'$e$', r'$\Phi_\phi$', r'$\Phi_r$', r'Flux']
xlabels = [r'$p$', r'$t$', r'$t$', r'$t$', r'$t$', r'$t$', r'$t$', r'$t$']
ys = [e, p, e, Phi_phi, Phi_r, flux]
xs = [p, t, t, t, t, t]

for i, (ax, x, y, xlab, ylab) in enumerate(zip(axes, xs, ys, xlabels, ylabels)):
    ax.plot(x, y, lw=0.5)
    ax.scatter(x, y, s=5)
    ax.set_xlabel(xlab, fontsize=16)
    ax.set_ylabel(ylab, fontsize=16)
```

## RomanNet Amplitudes: $A_{lmn}$

To generate the amplitudes, we use a RomanNet ([arXiv:1811.05491](https://arxiv.org/abs/1811.05491)). Roman stands for Reduced Order Modelling with Artificial Neurons. When training a neural network, it can be crucial to hand it data that is distilled to highlight the most pertinent information. A nice tool for doing this is Reduced Order Modeling. Reduced order modelling projects the information with lossless compression down to a lower dimensional space. We take our complex mode amplitude vectors containing 3843 modes and project this down to a real-valued space with 198 values. The neural network is then trained with inputs given by the $p$ and $e$ values and outputs given by the reduced order coefficients ($\alpha_i$):

\begin{equation}
A_{lmn}\in\mathbb{C}^{3843}\xrightarrow{\mathit{ROM}}\alpha_i\in\mathbb{R}^{198}\xrightarrow{\mathit{train}} f(p,e)=\alpha_i
\end{equation}

The neural network itself is extremely simple. It is a fully connected network with a LeakyReLU activation on all layers but the final layer. This means the neural network can be built simply with a sequence of linear matrix multiplications followed by a pass through the activation function. Once the neural network is trained, we will have a set of weights. This is where we will start in this tutorial.

During online evaluation of the waveform, we perform the reverse process:

\begin{equation}
f(p,e)\xrightarrow{\mathit{eval}}\alpha_i\in\mathbb{R}^{198}\xrightarrow{\mathit{project}}A_{lmn}\in\mathbb{C}^{3843}
\end{equation}

This method has pros and cons:

Pros:

* This is more of a global fit, rather than individual fits to given modes. This generally means storage of less information in memory. If we do individual fits of every mode, the memory necessary to store this information would scale badly with mode content. We are currently working in Schwarzschild Eccentric. As we go to generic Kerr, we expect the number of modes to increase by a factor of $\sim10$.

* Due to its global fit nature, this method is generally faster to evaluate than individual interpolants.

* Since it is a neural network and a linear projection, this is very suitable to GPUs.

* We expect this method, or methods similar to this, to scale better with dimensionality as we move towards the end goal of generic Kerr orbits.

Cons:

* For extremely quiet modes, this method can be less accurate. However, as these modes are quiet, this does not result in a significant loss of accuracy in the final waveform. Caution must be taken when using this method to analyze individual mode amplitudes. There is an approximate floor in the amplitude values at $\sim10^{-5}$. For reference, the loudest modes at a given $p$ and $e$ value are usually $\sim0.1-1.0$. (These amplitudes are not scaled for distance.|)

* Training these neural networks can be more of an art than a science. It takes a lot of trial and error to get this right.

### Constructing the neural network from trained weights

```python

# prepare to load the weights

path_to_few_dir = romannet.__file__[:-25]

weight_file = fp = "SchwarzschildEccentricInput.hdf5"
check_for_file_download(fp, path_to_few_dir)

weights = []
bias = []
dim1 = []
dim2 = []

# get highest layer number
num_layers = 0

# extract all necessary information from the file
with h5py.File(path_to_few_dir + "few/files/" + weight_file, "r") as fp:

    # get basic information
    num_teuk_modes = fp.attrs["num_teuk_modes"]
    transform_factor = fp.attrs["transform_factor"]
    break_index = fp.attrs["break_index"]

    # determine layer arrangement
    for key, value in fp.items():
        if key == "reduced_basis":
            continue

        layer_num = int(key[1:])

        if layer_num > num_layers:
            num_layers = layer_num

    # get weights and bias
    for i in range(1, num_layers + 1):
        temp = {}
        for let in ["w", "b"]:
            mat = fp.get(let + str(i))[:]
            temp[let] = np.asarray(mat)

        weights.append(temp["w"])
        bias.append(temp["b"])
        dim1.append(temp["w"].shape[0])
        dim2.append(temp["w"].shape[1])

    # get the post network transform matrix
    transform_matrix = np.asarray(fp["reduced_basis"])

# activation function
# we use a factor of 0.2 for negative values
def LeakyReLU(x):
    out = (x >= 0.0) * x +  (x < 0.0) * 0.2*x
    return out

# build the neural network
def RomanNet(p, e):

    p = np.atleast_1d(p)
    e = np.atleast_1d(e)

    # convert from the p coordinate to a special y coordinate
    # see the documentation for more details
    y = p_to_y(p, e)

    # prepare input
    x = np.array([y, e])

    # basic fully connected network
    for layer_i in range(num_layers):

        # linear transformation
        x = np.dot(weights[layer_i].T, x) + bias[layer_i][:, np.newaxis]

        # do not want to activate last layer
        if layer_i < num_layers - 1:
            # non-linear activatation
            x = LeakyReLU(x)

    # separate real and imaginary
    x = x[:break_index] + 1j * x[break_index:]

    # project back to amplitude basis
    out = np.dot(transform_matrix.T, x)/transform_factor
    return out.T

```

```python
# test it
p_test = np.array([11.0, 10.0, 10.0])
e_test = np.array([0.2, 0.3, 0.1])

RomanNet(p_test, e_test)
```

    array([[ 6.61875897e-06-9.45689250e-06j,  6.69607160e-06-1.13351771e-05j,
             6.41141133e-06-1.40002164e-05j, ...,
             7.50732732e-05+2.44753603e-05j,  6.73201363e-05-2.51759069e-06j,
            -2.62936078e-07-1.61971784e-05j],
           [-3.03984515e-05+2.23504926e-06j, -3.46638688e-05+2.11040541e-06j,
            -3.99084181e-05+1.67172621e-06j, ...,
             8.48364501e-05-3.51356321e-05j,  1.27022103e-04-6.69149191e-06j,
             9.19366706e-05+1.92716832e-05j],
           [-1.78544871e-06+1.54745952e-05j, -1.96676157e-06+1.75539377e-05j,
             7.20295633e-07+2.05609755e-05j, ...,
            -1.38437626e-05+1.46714870e-05j, -2.75508914e-05+9.37792116e-06j,
             1.12870278e-05+7.38750376e-06j]])

### Produce amplitudes associated with our trajectories

```python
# generate amplitudes with roman net
amps = RomanNet(p, e)

# check against actual code
RomanNetTrue = RomanAmplitude()
amps_check_1 = RomanNetTrue(p, e)

assert np.allclose(amps, amps_check_1)
```

### Check our RomanNet global fit against accurate values

```python
# get accurate values
# each mode is fitted with a bicubic spline
BicubicAmps = Interp2DAmplitude()
amps_check_2 = BicubicAmps(p, e)
```

#### Print vectors

```python
print("romannet:", amps_check_1[0][0:10], "\nBicubic Spline:", amps_check_2[0][0:10])
```

    romannet: [7.69192924e-06-5.39154229e-06j 9.99753291e-06-6.31307778e-06j
     1.35171826e-05-8.06391248e-06j 1.78239290e-05-7.33261441e-06j
     2.14797124e-05-6.31329971e-06j 2.57309168e-05-3.37761067e-06j
     2.89725596e-05+1.17404671e-06j 3.27106651e-05+7.79482012e-06j
     3.64016950e-05+1.64240791e-05j 3.63310870e-05+2.77218824e-05j]
    Bicubic Spline: [0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j]

We see if we take a quick examination of our results (a small subset) we see the problem with the global fit. It cannot handle modes with a very small amplitude. Therefore, you might naively think that this method might not work.

#### Compare via the cosine between the vectors

```python
cos = (np.dot(amps_check_2.conj()[0].T, amps_check_1[0]) /
       np.sqrt(np.dot(amps_check_2.conj()[0], amps_check_2[0])
               * np.dot(amps_check_1.conj()[0], amps_check_1[0]))).real

print('Cos:', cos)
```

    Cos: 0.9999733790390088

We now see that if we compare the results as a whole, the results match very well. If we look at individual modes that have high power, they are likely to strongly match when the bicubic spline is compared to the RomanNet method.

## GPU-accelerated Waveform Build: $\sum_{lmn}$

### Quick primer on GPUs

GPUs run code in parallel in a configuration of grids, blocks, and threads.

* **Threads** are the the actual software units that run the code. Threads run independently of one another. Threads are referenced in code using `threadIdx.x`.
* **Blocks** are an array of threads. This array can be 1, 2, or 3 dimensions. In more than one dimension, thread indices are referenced using `x`, `y`, and `z` (e.g. `threadIdx.z`). Blocks, similar to threads, are referenced using `blockIdx.x`, `blockIdx.y`, or `blockIdx.z`. The size of a block, i.e. the number of threads along a given dimension is given as `blockDim.x`. **Note**: In applications I have worked on, I rarely ever use more than 1 dimension of threads.
* **Grids** are an array of blocks. This array can be 1, 2, or 3 dimensions. There is no reference for grids as they efffectively represent the entire GPU kernel. The size of a grid, i.e. the number of blocks along a dimension, is determined with `gridDim.x`.

Organizing your code properly into blocks and threads is a key component of maximizing your efficiency.

Another equally important aspect to the maximization of efficiency is the proper use of the GPU memory structures. There are three main GPU memory structures to consider (there are a few more but they are for more specific uses). These are global memory, shared memory, and local memory.

* **Global memory** is contained on the host (or off-chip). Global memory is accessible by every block and thread in the entire grid. It contains many GBs of RAM. (The GPUs I use currently, which are Tesla V100s, have ~16GB. The new A100s have 40-80GB.) Because this memory is contained off-chip, it is slow to access. A key to accessing global memory efficiently is to use so-called "memory coalescence." This effectively means that neighboring threads access neighboring addresses in memory. This allows the compiler to make up to 32 memory reads at the same time, rather than 32 separate memory reads. We will see this in a simple case below.

* **Local memory** is also contained on the host. Local memory consists of any arrays allocated within the kernel specific to each thread. Therefore, this memory is only accessible by the thread it is created on. Usually, around 512 KB are available for local memory for each thread. Since, this memory is off-chip, it is also slow to access. Local memory is, however, always accessed in a memory coalescing manner.

* **Shared memory** is different. Shared memory is located on-chip. It is, therefore, much faster to read from (~100x faster). Shared memory is accessible by all the threads on a given block. The catch is that only ~48 KB are available for shared memory. Therefore, leveraging shared memory effectively is key to the efficiency of GPU code.

To sum up, there are two main points to consider when beginning to program on GPUs:

* Layout your grid effectively for your given problem.
* Use memory effectively: leverage the availability of shared memory and read from global memory in a coalescing fashion.

Some parting thoughts on GPUs:

* For maximal efficiency and stability, I usually code in C++/CUDA so that everything is precompiled and tested.
* There are python libraries that leverage the power of GPUs. Check out numba, CuPy, PyTorch, Tensorflow, PyCUDA.
* Lately, I have really focused on writing CPU/GPU agnostic code. This means that the source code is ~99% the same between the two. On the python side I usually sub in CuPy for Numpy. We will see some basic pointers on this below. In C++, I use short compiler directives to make minimal changes.
* Generally speaking it is optimal to store all quantities in a 1 dimensional array when working with GPUs. Let's say you have a two dimensional array that has dimensions (dim1, dim2) and is referenced with (i, j). You can turn this into a 1D array that references each value with (i * dim2 + j).

### Why are GPUs important for EMRI Waveforms?

The waveform summation is the key bottleneck. This is an operation that is uniquely suited to GPUs. You can see the improvement below.

### Basic example coded in python

We are going to write code in python that will reflect what we will do on the GPU. It will obviously still be on the CPU, but will give you a chance to see how to write some basic GPU code and understand how it works. **Make sure to read the comments in the code.**

We will do two examples.

#### Multiply two arrays

Here we will multiply two arrays as we would on the GPU.

```python
# setup our GPU quantities that would come in if we were actually on GPUs.

NUM_THREADS = 64  # needs to be a multiple of 32
blockDim = NUM_THREADS  # blockDim.x
```

```python

# this initial piece is a basic of CPU/GPU agnostic code in python
try:
    import cupy as xp

except ModuleNotFoundError:
    import numpy as xp

# __global__
def multiply_arrays(array_out, array1, array2, n):
    """
    // what this would like like in CUDA

    // if (threadIdx.x + blockDim.x * blockIdx.x >= n) return;
    for (int i = threadIdx.x + blockDim.x * blockIdx.x; i < n; i += blockDim.x * gridDim.x)
    {
        array_out[i] = array1[i] * array2[i];
    }
    """
    # The loops here are to simulate the GPU.
    # In reality the GPU will run all threads and blocks in parallel

    # begin simulation
    for block in range(num_blocks):
        for thread in range(NUM_THREADS):
    # end simulation

            # get the overall index in the grid
            # based on the thread and block
            i = thread + block * blockDim

            # since the GPU runs threads in parallel
            # we need to make sure the GPU does not index a value
            # on the last block that goes over the array length
            if i >= n:
                continue

            # use our index to get values out of the array
            # notice this simple statement uses memory coalescence
            # neighboring threads will have consecutive indices
            # therefore, they will access consecutive addresses in each array
            array_out[i] = array1[i] * array2[i]

length = int(2 ** 14)

# initialize arrays
array1 = xp.random.rand(length)
array2 = xp.random.rand(length)

# prepare output array
array_out = xp.zeros_like(array1)

# get the number of blocks
# the number of blocks multiplied by the number of threads per block
# must be greater than the length of the array
num_blocks = int(np.ceil((length + NUM_THREADS -1)/NUM_THREADS))

# this would actually be called as
"""
multiply_arrays<<<num_blocks, NUM_THREADS>>>(array_out, array1, array2, length);
cudaDeviceSynchronize();
"""

multiply_arrays(array_out, array1, array2, length)

# confirm it
assert np.allclose(array_out, array1 * array2)
```

#### Linear Interpolation to scale up an array

For our second example it will be a bit more complicated. This way we can see how to use shared memory.

Here we will scale up the size of an array using linear interpolation. The key here is that we need every block to read in the original array to shared memory. From there, we perform the interpolation. This allows us to achieve memory coalescence with all global reads and ensure that all reads when actually interpolating are done from shared memory.

```python
# in CUDA, when you statically allocate shared arrays,
# you cannot use a variable length
# therefore, in C++/CUDA you need to declare a max value up front

#define MAX_INPUT 1000

def linear_interpolation(array_out, array_in, n_out, n_in, dx_out, dx_in):

    # we want to read in to shared memory
    # want to do this with memory coalescence
    # since this is only on a specific block
    # we need to use only the threads on this block

    """
    // the below simulated code block in CUDA would look like this:

    // in CUDA, you have to declare shared arrays

    __shared__ double shared_array_in[MAX_INPUT];

    for (int i = threadIdx.x; i < n_in; i += blockDim.x)
    {
        // memory coalescence only needed from array_in (global memory)
        shared_array_in[i] = array_in[i];
    }

    // we need to make sure the threads all finish this operation before we move on
    // Therefore, we force the threads to synchronize

    __syncthreads();

    """

    # we will just do the following to simulate in python easily
    shared_array_in = xp.zeros(n_in)

    for thread in range(NUM_THREADS):
        i = thread

        # this while statement simulates the above commented loop
        while (i < n_in):
            # here is the memory coalesced reads
            shared_array_in[i] = array_in[i]

            # just for simulation
            i += blockDim

    """
    // this is what the below would really like it in CUDA
    // if (threadIdx.x + blockDim.x * blockIdx.x >= n_out) return;
    for (int i = threadIdx.x + blockDim.x * blockIdx.x; i < n_out; i += blockDim.x * gridDim.x)
    {
        // get the new out value assuming equal spacing
        double x_new = dx_out * i;

        // get index of the point in the original array below the new point
        int ind_in = (int) (x_new / dx_in);

        // get the below x value
        double x_old = ind_in * dx_in;

        // slope of segment
        double m = (shared_array_in[ind_in + 1] - shared_array_in[ind_in])/dx_in;

        // interpolate
        double new_value = m * (x_new - x_old) + shared_array_in[ind_in];
        array_out[i] = new_value;
    }

    """
    # The loops here are to simulate the GPU.
    # In reality the GPU will run all threads and blocks in parallel

    # begin simulation
    for block in range(num_blocks):
        for thread in range(NUM_THREADS):
    # end simulation

            # get the overall index in the grid
            # based on the thread and block
            i = thread + block * blockDim

            # since the GPU runs threads in parallel
            # we need to make sure the GPU does not index a value
            # on the last block that goes over the array length
            if i >= n_out:
                continue

            x_new = dx_out * i

            ind_in = int(x_new / dx_in)
            x_old = ind_in * dx_in

            m = (shared_array_in[ind_in + 1] - shared_array_in[ind_in])/dx_in

            new_value = m * (x_new - x_old) + shared_array_in[ind_in]
            array_out[i] = new_value

# original length
length_in = 100

# prepare input arrays
x_in = xp.arange(length_in)
y_in = (x_in ** 2).astype(xp.float64)

# set length out
length_out = int(2 ** 14)

# setup the new x_values
x_new = xp.linspace(x_in[0], x_in[-1], length_out + 1)[:-1]

# change in original x values
dx_in = 1.0

# we will actually only use the spacing to find new points
dx_out = x_new[1] - x_new[0]

y_out = xp.zeros(length_out)

num_blocks = int(np.ceil((length_out + NUM_THREADS -1)/NUM_THREADS))

# this would actually be called as
# linear_interpolation<<<num_blocks, NUM_THREADS>>>(y_out, y_in, length_out, length_in, dx_out, dx_in);
# cudaDeviceSynchronize();
linear_interpolation(y_out, y_in, length_out, length_in, dx_out, dx_in)

# confirm it
plt.plot(x_new, y_out, lw=6, label='out')
plt.plot(x_in, y_in, '.', label='in')
```

### Comment on CPU/GPU Agnostic code in C++/CUDA and Python

Making CPU/GPU agnostic code consists of 3 main parts:

* Sub CuPy for NumPy
* Compiler directives in C++/CUDA
* Easy transition from Python to C++ through an augmented Cython process

Above, we saw an example of how to deal with the Python side. The Cython functionality is effectively a decorator function that gets the pointer of a Numpy or Cupy array and sends that into the C++ code. Below is a basic example of how to use basic compiler directives to make your code more CPU/GPU agnostic.

**Note**: I generally argue that we should build codes for GPUs and then adapt them to CPUs, not the other way around which is the typical direction. In my experience, the CPU codes adapted from GPU codes are just as fast or within a small percentage of the speed of CPU designed codes. And, generally, it is much harder to optimize going from CPU->GPU rather than GPU->CPU.

## Future Plans

* Build fast trajectories in Kerr under the NIT framework
* Generate amplitudes in Kerr regime
* Further analyze how to determine best methods for mode content inclusion

We need your help! If anyone is interested in working on these types of issues, please let myself or any of my collaborators (listed above) know!


# Interpolating chemical equilibrium abundances

The atmospheric composition is determined by its elemental composition and the chemcical reactions that link these elements and chemical reactant species. In the atmosphere of an exoplanet, which is inherently 3-d and dynamic, mixing and advection processes complicate this picture, because the chemical abundances inferred from spectra are not necessarily representative of the pressures and temperatures probed locally by the observation. In addition, also the radiation field of the host star may influence the abundance and opacity structure of the atmosphere, by dissociating or ionizing chemical species or forming photochemical hazes. In this picture, using chemical equilibrium abundances, which means the steady state chemical composition of the atmosphere when treating its layers as isolated and independent (no advection/mixing/photochemistry), appears unwise.

Yet, calculating chemical equilibrium abundances may still serve as a useful starting point for initializing abundances in planetary atmospheres. In certain conditions, for example for planets hot enough such that mixing is negligible, chemical equilibrium abundances can even be sufficient to accurately describe the abundances in the atmospheric region probed by observations.

Here we show how to use our ``poor_mans_nonequ_chem()`` function. This function interpolates the abundances of the most important species from a big chemical equilibrium table as a function of pressure, temperature, metallicity [Fe/H] (where [Fe/H]=0 is solar) and cabon-to-oxygen number ratio (C/O). A C/O $\sim 0.55$ is the solar value, see, for example, [Asplund et al. (2009)](https://ui.adsabs.harvard.edu/abs/2009ARA%26A..47..481A/abstract).

**poor_mans_nonequ_chem() also contains a simple quenching implementation**. This means that a quench pressure can be specified above which the abundances of H$_2$O, CH$_4$ and CO are taken to be constant, following the reasoning in, for example, [Zahnle & Marley (2014)](https://ui.adsabs.harvard.edu/abs/2014ApJ...797...41Z/abstract).

The chemical grid has been calculated with easyCHEM, which is our self-written Gibbs free energy minimizer that is described in [Mollière et al. 2017](https://ui.adsabs.harvard.edu/abs/2017A%26A...600A..10M/abstract), and essentially a CEA clone, see [Gordon & McBride (1994)](https://ntrs.nasa.gov/citations/19950013764). easyCHEM has been benchmarked as the chemical equilibrium tool of petitCODE in [Baudino et al. (2017)](https://ui.adsabs.harvard.edu/abs/2017ApJ...850..150B/abstract). The grid dimensions are $T \in [60, 4000]~{\rm K}$, with 100 equidistant temperature points, $P \in [10^{-8}, 1000]~{\rm bar}$, with 100 pressure points spaced equidistantly in log space,  ${\rm C/O} \in [0.1, 1.6]$, with 20 equidistant points, and ${\rm [Fe/H]} \in [-2,3]$, with 40 equidistant points. C/O is varied by varing the oxygen abundance, once the total elemental composition has been set by [Fe/H]. The total list of reactant species included was H, H$_2$, He, O, C, N, Mg, Si, Fe, S, AL, Ca, Na, Ni, P, K, Ti, CO, OH, SH, N$_2$, O$_2$, SiO, TiO, SiS, H$_2$O, C$_2$, CH, CN, CS, SiC, NH, SiH, NO, SN, SiN, SO, S$_2$, C$_2$H, HCN, C$_2$H$_2$,acetylene, CH$_4$, ALH, ALOH, AL$_2$O, CaOH, MgH, Mg, OH, PH$_3$, CO$_2$, TiO$_2$, Si$_2$C, SiO$_2$, FeO, NH$_2$, NH$_3$, CH$_2$, CH$_3$, H$_2$S, VO, VO$_2$, NaCL, KCL, e-, H+, H-, Na+, K+, PH$_2$, P$_2$, PS, PO, P$_4$O$_6$, PH, V, FeH, VO(c), VO(L), MgSiO$_3$(c), SiC(c), Fe(c), Na$_2$S(c), KCL(c), Fe(L), SiC(L), MgSiO$_3$(L), H$_2$O(L), H$_2$O(c), TiO(c), TiO(L), TiO$_2$(c), TiO$_2$(L), H$_3$PO$_4$(c), H$_3$PO$_4$(L), where (c) stands for solid and (L) for liquid species. To conserve space only the mass fractions of the following species are tabulated for use in ``poor_mans_nonequ_chem()``: H$_2$, He, CO, H$_2$O, HCN, C$_2$H$_2$,acetylene, CH$_4$, PH$_3$, CO$_2$, NH$_3$, H$_2$S, VO, TiO, Na, K, SiO, e-, H-, H, FeH, MMW, nabla_ad. MMW denotes the mean molcular weight in the atmosphere, in units of amu. nabla_ad is the moist adiabatic lapse rate $\nabla_{\rm ad} = (\partial {\rm ln}T/\partial {\rm ln}P)_{\rm ad}$, determined as described in [Mollière et al. (2020)](https://arxiv.org/abs/2006.09394).

**Abundances in poor_mans_nonequ_chem():** abundances in poor_mans_nonequ_chem() are in units of **mass fractions**, not number fractions (aka volume mixing ratio, VMR). You can convert between mass fractions and VMRs by using
\begin{equation}
X_i = \frac{\mu_i}{\mu}n_i,
\end{equation}
where $X_i$ is the mass fraction of species $i$, $\mu_i$ the mass of a single molecule/atom/ion/... of species $i$, $\mu$ is the atmospheric mean molecular weight, and $n_i$ is the VMR of species $i$.

## Example

Here we will give some examples for how to interpolate chemical abundances using ``poor_mans_nonequ_chem()``. The installation of it is described immediately below.

We begin by loading the relevant packages:

```python
import numpy as np
import pylab as plt
plt.rcParams['figure.figsize'] = (10, 6)
from poor_mans_nonequ_chem import poor_mans_nonequ_chem as pm
from petitRADTRANS import nat_cst as nc
```

Next, we define an atmospheric temperature and abundance structure, following an example similar to what is shown in ["Getting Started"](getting_started.html).

```python
pressures = np.logspace(-6, 2, 100)
gravity = 1e1**2.45
kappa_IR = 0.01
gamma = 0.4
T_int = 200.
T_equ = 800.

temperature = nc.guillot_global(pressures, kappa_IR, gamma, gravity, T_int, T_equ)

plt.plot(temperature, pressures)
plt.yscale('log')
plt.ylim([1e2, 1e-6])
plt.xlabel('T (K)')
plt.ylabel('P (bar)')
plt.show()
plt.clf()
```

No we calculate the abundances. Note that the input parameters temperature, pressure, [Fe/H] and C/O are all arrays, and must be defined at every layer, but usually C/O and [Fe/H] are not vertically varying in the atmosphere. Here we chose a solar metallicity and C/O.

```python
COs = 0.55 * np.ones_like(pressures)
FeHs = 0. * np.ones_like(pressures)

mass_fractions = pm.interpol_abundances(COs, \
            FeHs, \
            temperature, \
            pressures)
```

Lets plot the abundances:

```python
plt.yscale('log')
plt.xscale('log')
plt.ylim([100,1e-6])
plt.xlim([1e-10, 1])

for species in mass_fractions.keys():
    if species not in ['MMW', 'nabla_ad']:
        plt.plot(mass_fractions[species], pressures, label = species)

plt.legend(loc='best', ncol=3)
plt.xlabel('Mass fraction')
plt.ylabel('P (bar)')
plt.show()
plt.clf()
```

Also the MMW and adiabatic temperature gradient are accessible, which turns out to be vertically constant for this atmosphere:

```python
plt.yscale('log')
plt.ylim([100,1e-6])

plt.plot(mass_fractions['MMW'], pressures, label = species)

plt.xlabel('Mean molecular weight (MMW)')
plt.ylabel('P (bar)')
plt.show()
plt.clf()
```

```python
plt.yscale('log')
plt.ylim([100,1e-6])

plt.plot(mass_fractions['nabla_ad'], pressures, label = species)

plt.xlabel(r'Moist adiabatic temperature gradient $\nabla_{\rm ad}$')
plt.ylabel('P (bar)')
plt.show()
plt.clf()
```

As a test, we can also plot abundances at a given pressure, temperature, C/O, as a function of metallicity:

```python
FeHs = np.linspace(-2, 3, 100)
p_const = np.ones_like(FeHs)*1e-1
t_const = np.ones_like(FeHs)*900.
COs = 0.55 * np.ones_like(FeHs)

mass_fractions = pm.interpol_abundances(COs, \
            FeHs, \
            t_const, \
            p_const)
```

```python
plt.yscale('log')
plt.ylim([1e-10, 1])
plt.xlim([-2, 3])

for species in mass_fractions.keys():
    if species not in ['MMW', 'nabla_ad']:
        plt.plot(FeHs, mass_fractions[species], label = species)

plt.legend(loc='best', ncol=3)
plt.ylabel('Mass fraction')
plt.xlabel('[Fe/H]')
plt.title('T = 900 K, P = 0.1 bar, C/O = 0.55')
plt.show()
plt.clf()
```

This is the mean molecular weight as a function of metallicity:

```python
plt.xlim([-2, 3])
plt.ylim([0, 15])

plt.plot(FeHs, mass_fractions['MMW'])

plt.ylabel('Mean molecular weight')
plt.xlabel('[Fe/H]')
plt.title('T = 900 K, P = 0.1 bar, C/O = 0.55')
plt.show()
plt.clf()
```

Here the mass fractions as a function of the C/O ratio, at solar metallicity:

```python
COs = np.linspace(0.1, 1.6, 100)
FeHs = 0. * np.ones_like(COs)
p_const = np.ones_like(COs)*1e-1
t_const = np.ones_like(COs)*1400.

mass_fractions = pm.interpol_abundances(COs, \
            FeHs, \
            t_const, \
            p_const)
```

```python
plt.yscale('log')
plt.ylim([1e-10, 1])
plt.xlim([0.1, 1.6])

for species in mass_fractions.keys():
    if species not in ['MMW', 'nabla_ad']:
        plt.plot(COs, mass_fractions[species], label = species)

plt.legend(loc='lower right', ncol=3)
plt.ylabel('Mass fraction')
plt.xlabel('C/O')
plt.title('T = 1400 K, P = 0.1 bar, [Fe/H] = 0')
plt.show()
plt.clf()
```

Finally, we show an example with quenching here, assuming that the CH$_4$, H$_2$O and CO abundances are constant in for pressures below 1 bar. We will use the temperature structure from above again.

```python
# Nominal case:
COs = 0.55 * np.ones_like(pressures)
FeHs = 0. * np.ones_like(pressures)

mass_fractions = pm.interpol_abundances(COs, \
            FeHs, \
            temperature, \
            pressures)

# Quenching case:
mass_fractions_quench = pm.interpol_abundances(COs, \
            FeHs, \
            temperature, \
            pressures, \
            Pquench_carbon = 1.)
```

```python
plt.yscale('log')
plt.xscale('log')
plt.ylim([100,1e-6])
plt.xlim([1e-10, 1])

plt.plot(mass_fractions['H2'], pressures, linestyle = '--', color = 'black')
plt.plot(mass_fractions_quench['H2'], pressures, label = 'H2', linestyle = '-', color = 'black')
plt.plot(mass_fractions['He'], pressures, linestyle = '--', color = 'gray')
plt.plot(mass_fractions_quench['He'], pressures, label = 'He', linestyle = '-', color = 'gray')
plt.plot(mass_fractions['CO'], pressures, linestyle = '--', color = 'red')
plt.plot(mass_fractions_quench['CO'], pressures, label = 'CO', linestyle = '-', color = 'red')
plt.plot(mass_fractions['H2O'], pressures, linestyle = '--', color = 'purple')
plt.plot(mass_fractions_quench['H2O'], pressures, label = 'H2O', linestyle = '-', color = 'purple')
plt.plot(mass_fractions['CH4'], pressures, linestyle = '--', color = 'orange')
plt.plot(mass_fractions_quench['CH4'], pressures, label = 'CH4', linestyle = '-', color = 'orange')

plt.axhline(1., color = 'blue', linestyle = ':', label = 'Quench pressure')

plt.legend(loc='best', ncol=1)
plt.xlabel('Mass fraction')
plt.ylabel('P (bar)')
plt.title("Poor man's non-equilibrium chemistry: quenching")
plt.show()
plt.clf()
```

## Installation of poor_mans_nonequ_chem()

Download the ``poor_mans_nonequ_chem()`` [here](https://www.dropbox.com/s/6r3djmzu8y1qplf/poor_mans_nonequ_chem.zip?dl=0). If relevant, activate the python environment that you want to run your code in (similar to the [installation tips when installing pRT](../installation.html)). Then, in the terminal, switch into the poor_mans_nonequ_chem/chem_fortran_util folder and type ``./make.sh``, again similar to the [pRT installation](../installation.html). Lastly, [like when installing pRT](../installation.html), enter the following line to your .bashrch or .bashprofile (depending on your operating system):

``export PYTHONPATH=Path of the folder containing the highest level poor_mans_nonequ_chem folder/:$PYTHONPATH``

After this close and reopen the terminal such that it will set the Python path correctly, or source the  .bashrc/.bashprofile file. Alternatively, you can place the following line at the top of your scripts that use ``poor_mans_nonequ_chem``, berfore importing ``poor_mans_nonequ_chem``:

```python
import sys
sys.path.insert(0,'Path to folder containing the highest level poor_mans_nonequ_chem folder/')
```

You are done!


```python
# %load /Users/facai/Study/book_notes/preconfig.py
%matplotlib inline

import matplotlib.pyplot as plt
import seaborn as sns
sns.set(color_codes=True)
sns.set(font='SimHei')
plt.rcParams['axes.grid'] = False

from IPython.display import SVG

def show_image(filename, figsize=None):
    if figsize:
        plt.figure(figsize=figsize)

    plt.imshow(plt.imread(filename))
```

GBDT(Gradient Boosting Decision Tree) 原理简介
=============================================

### 0. 前言

我最开始了解 GBDT 时，死活不理解决策树这种分段函数，怎么可能算出一阶导数。读了论文 Friedman - Greedy Function Approximation: A Gradient Boosting Machine 后，才发现自己完全误解了决策树在GBDT中的作用。论文总是倾向把简单事情描述复杂，博客又常常过于简练地只给具体算法。写这篇文章的想法，就是想用自己的理解，既照顾到数学框架，又落实到具体算法实现，把GBDT说清楚。这里可能有谬误，读者应审慎地阅读。

### 1. Gradient Boosting 与最优化
ref:

+ [机器学习中的算法(1)-决策树模型组合之随机森林与GBDT](http://www.cnblogs.com/LeftNotEasy/archive/2011/03/07/random-forest-and-gbdt.html)

在说GBDT前，我们先说下它的俩前缀Gradient Boosting：

+ Boosting: 这是一种迭代算法，每一次训练都是在前面已有模型的预测基础上进行。
  最简单地说，先训练一个初始模型，对比真实值和预测值的残差；用残差再训练一个模型，再计算残差；再训练……。这样，每一个模型都专注于修正前面模型最不给力的效果方面。于是，通过这种方式联合多个弱分类器，就能得到一个强分类器。
+ Gradient: 梯度。对于单个变量来说，就是一阶导数。
  前面Boosting要计算残差，方法就很多了，你可以直接相减，也可以定义个函数。这里Gradient就是说用梯度来计算残差。

所以说Gradient Boosting是一种算法框架，它是用Grdient计算残差的Boosting过程，而GBDT只是用决策树来训练的一种具体实现。说到这，学习过最优化的童鞋一定会对Gradient Boosting有种强烈的既视感，是的，那就是我们耳熟能详的最速梯度下降法。

所以，我们先从最速梯度下降法作引，一步步走到Gradeint Boosting框架，然后再落到决策树上，最后如果有精力，介绍下进一步的优化TreeBoost。

### 2. 最速梯度下降法

我简要说下基本思路，详细可见[梯度下降法 - 维基百科](https://zh.wikipedia.org/wiki/梯度下降法)。

假设有个有界函数 $| f(x) | < M$，我们需要知道它的最小值。日常的例子中，这个函数可能是一项工程所用的花销，一个项目占用的时间，我们想找到损失最小的方案，即：$\operatorname{arg min}_x \, f(x) $。 找到最小解的过程就叫寻优，函数$f(x)$常被称为损失函数。

#### 2.0 理论

梯度下降法是一种数值寻优算法。什么是数值寻优呢？就是现实中$f(x)$很多时候难以用理论求解出最值，或者$f(x)$就是个黑盒子，我们没有办法直接找到最值，只能去有限的数值解中搜索，所以叫数值寻优。

梯度下降法就是按函数梯度来步进到极值点的过程，其数学表达式形如：
\begin{equation}
    x_{m+1} = x_m - \lambda \frac{\partial f(x)}{\partial x}
\end{equation}

具体过程如下图所示，取自[可扩展机器学习——梯度下降(Gradient Descent)](http://blog.csdn.net/google19890102/article/details/50276775)：

```python
show_image("./res/gradient_descent.jpg", figsize=(12,8))
```

```python
show_image("./res/iterator.jpg")
```

#### 2.1 增量寻优

梯度下降法是利用了导数做为步进参考值，它要求函数是已知且可导的。问题是，现实中，有时这个函数是未知的，比如说对应的是一台机器。我不知道怎么对它建模，唯一的信息就是喂它输入，它就会返回输出。那怎么办呢？

首先，我们回到 $\frac{\partial f(x)}{\partial x}$ 来看，因为步进大小还受 $\lambda$ 调控。这里梯度的作用，其实主要是指示步进方向。

然后，我们换个思路找步进方向。

假设机器的真实模型是 $f(x)$，对我们来说它是未知的。尝试如下：

+ 第一次试探，随机给个初始值 $x_0$，得到 $f(x_0)$。
+ 第二次试探，再随机输入 $x_1$，得到 $f(x_1)$。

这时，有如下三种情况，可以指导第三次试探：

1. $f(x_1) < f(x_0)$，也就是 $x_0 \to x_1$的方向是对的。
2. $f(x_1) = f(x_0)$，没有有用信息。
3. $f(x_1) > f(x_0)$，方向反了。

```python
show_image("./res/incr_opt.png", figsize=(10,5))
```

因为我们始终是找极小值点，这个过程就始终如上图所示"U"形。那么每次的步进方向就可用 $f(x_1) - f(x_0)$ 来指示。也就是说，虽然我们无法对机器$f(x)$建模，但我们可以对寻优的过程 $z = f(x_i) - f(x_{i+1})$ 建模 $g(z)$。而 $g(z)$ 只要满足单调且过零点，理论上它可以是任意的，不同函数影响的只是收敛速度，但都会收敛到 $w^*$ 附近。

#### 2.2 逼近

再换个角度，我们不再想寻优。而是假设这个极小值$w^*$已经知道了，如何从初始值$w_0 = f(x_0)$一步步逼进到这个极小值$w^*$。其实可以完全借用上节的增量寻优建模手法，我们定义 $z = f(x_i) - w^*$，从而对逼近过程建模 $g(z)$。

```python
show_image("./res/approx.png", figsize=(10,5))
```

同样地，$g(z)$只要满足一定条件的函数。在这里，就是过零点，正数区间单增，负数区间单减，最常见的代表就是二次函数 $g(z) = z^2$。

我们梳理下整个逼近过程，如下图：

```python
show_image("./res/model.png", figsize=(10,5))
```

从初始点$w_0$，沿着$w_1, w_2, \dotsc, w_i$到达$w^*$。其中每步的步长是$\lambda \frac{\partial g(z)}{\partial z} |_{z = z_i}$。$\lambda$是控制步长的字习率，不用管。我们重点看下$z_i$:

\begin{align}
    z_i &= w_i - w^* \\
        &= f(x_i) - w^*
\end{align}

上面，每一步对应的$f(x_i)$是一个恒定的模型，变动的是$x$。我们反着想，假设$x$是固定的，而这个模型其实是变化的 $f_i(x)$，再把$w^*$用$y$表示。那么

\begin{align}
    z_i &= f(x_i) - w^* \\
        &= f_i(x) - y
\end{align}

是不是开始有点眼熟了，如果我们再令 $g(z) = z^2$，展开看看：

\begin{align}
    g(z_i) &= z_i^2 \\
           &= (y - f_i(x))^2
\end{align}

好了，我们总结下上述函数的意义：
给定自变量$x$、目标值$y$和初始位置$w_0$，可以用规定的$g(z)$和模型$f_i(x)$来构建出一个迭代序列$\{w_0, w_1, \dotsc, y\}$，其中$w_{i+1} = w_i - \lambda \frac{\partial g(z)}{\partial z} \, |_{z = y - f_i(x)}$。

而：
+ 对给定$x$，$y$是常数，所以$\partial z = \partial f_i(x)$。
+ $g(z) | z = y - f_i(x)$表述的是$y$和$f_i(x)$的距离相近度，可以更泛性推广记为$L(y, f_i(x))$。

所以记$w_{i+1} = w_i - \lambda \frac{\partial L(y, f_i(x))}{\partial f_i(x)}$

用机器学习的术语翻译就是，如果有特征$x$、目标值$y$和初始模型$f_0$，对某个损失函数$L(y, f_i(x))$，就有模型序列使得$f_i(x) \to y$。

而Gradient Boosting就是回答了上述的两个规定：

1. 怎样构建模型序列$f_i$?
2. 什么样的损失函数$L(y, f_i(x))$?

### 3. Gradient Boosting

#### 3.1 怎样构建模型序列$f_i$
Gradient Boosting的答案就是使用Boosting来训练Gradient(梯度）值的加法模型。

具体来说，已知有特征$x$、标签$y$和初始模型$f_0$。
然后，我们选定一种机器学习算法$\text{Model}(\text{feature}, \text{value})$，输入数据后就会给出模型$F(x)$。

我们规定模型生成序列方法如下：

\begin{align}
    & F_0 = f_0 \\
    & \hat{y_i} = - \frac{\partial L(y, f_{i-1}(x)}{\partial f_{i-1}(x)} \\
    & F_i = \text{Model}(x, \hat{y_i}) \\
    & f_{i} = f_{i-1} + F_i  = \displaystyle \sum_{k=0}^{i} F_k \\
\end{align}

上式第二行就是Gradient（梯度），第三行就是用梯度训练的单个模型$F_i$，而第四行$f_i$就是Boosting加法。请注意区分$f_i$和$F_i$。

很多时候，我们期望初始模型自动产生，并且和其它模型是同种类型的，上面的生成序列方法就可改写为：

\begin{align}
    & F_0 = \text{Model}(x, y) \\
    & \hat{y_i} = - \frac{\partial L(y, f_{i-1}(x)}{\partial f_{i-1}(x)} \\
    & F_i = \text{Model}(x, \hat{y_i}) \\
    & f_{i} = f_{i-1} + F_i  = \displaystyle \sum_{k=0}^{i} F_k \\
\end{align}

总结下流程：

1. 用$x$和$y$训练初始模型$F_0$，并令累加模型$f_0 = F_0$。
2. 用梯度计算出残差 $\hat{y_i}$。
3. 用$x$和残差训练新模型$F_i$。
4. 将新模型加入累加模型$f_i$。
5. 迭代直到达到终止条件。

#### 3.2 什么样的损失函数 $L(y, f_i(x))$

正如2.2节所说，损失函数$L(y, f_i(x))$定义的是从初始值到目标值的逼近过程。常理来说，离得越远，我们期望步距越大，调整越剧烈；离得越近，期望步距越小，调整越缓合。常用的有L1, L2和logistic。当然，理论上，你也可以用其它千奇百怪的函数，只要它符合$x_1 \to x_2 \implies f(x_1) \to f(x_2)$。

### 4. GBDT(Gradient Boosting Decision Tree)

将第3节中的机器模型Model换成决策树，就得到了GBDT。嗯，就是这么简单。

Treeboost是具体对决策树参与的计算做了公式推导，作用是简化运算量，打算后续开文专门讲。

### 总结
我们最后总结下，GBDT是Gradient Boosting框架用决策树完成的一个具体实现。Gradeint Boosting框架的本质思路是一个从基准值到目标值的逼近过程，也是一个基准弱模型向强模型的进化过程。逼近过程是用损失函数来描述和控制的，具体的步进量（残差）由损失函数的梯度来指示。机器模型的作用是用来拟合和预测残差。

所以，我以前不理解决策树怎么求导，而事实是决策树压根就和求导没关系。希望这篇文章能帮助到其它人的理解，也请指正错点漏点疑点。

```python

```


```python
%matplotlib inline
```

Sequence-to-Sequence Modeling with nn.Transformer and TorchText
===============================================================

This is a tutorial on how to train a sequence-to-sequence model
that uses the
`nn.Transformer <https://pytorch.org/docs/master/nn.html?highlight=nn%20transformer#torch.nn.Transformer>`__ module.

PyTorch 1.2 release includes a standard transformer module based on the
paper `Attention is All You
Need <https://arxiv.org/pdf/1706.03762.pdf>`__. The transformer model
has been proved to be superior in quality for many sequence-to-sequence
problems while being more parallelizable. The ``nn.Transformer`` module
relies entirely on an attention mechanism (another module recently
implemented as `nn.MultiheadAttention <https://pytorch.org/docs/master/nn.html?highlight=multiheadattention#torch.nn.MultiheadAttention>`__) to draw global dependencies
between input and output. The ``nn.Transformer`` module is now highly
modularized such that a single component (like `nn.TransformerEncoder <https://pytorch.org/docs/master/nn.html?highlight=nn%20transformerencoder#torch.nn.TransformerEncoder>`__
in this tutorial) can be easily adapted/composed.

Define the model
----------------

In this tutorial, we train ``nn.TransformerEncoder`` model on a
language modeling task. The language modeling task is to assign a
probability for the likelihood of a given word (or a sequence of words)
to follow a sequence of words. A sequence of tokens are passed to the embedding
layer first, followed by a positional encoding layer to account for the order
of the word (see the next paragraph for more details). The
``nn.TransformerEncoder`` consists of multiple layers of
`nn.TransformerEncoderLayer <https://pytorch.org/docs/master/nn.html?highlight=transformerencoderlayer#torch.nn.TransformerEncoderLayer>`__. Along with the input sequence, a square
attention mask is required because the self-attention layers in
``nn.TransformerEncoder`` are only allowed to attend the earlier positions in
the sequence. For the language modeling task, any tokens on the future
positions should be masked. To have the actual words, the output
of ``nn.TransformerEncoder`` model is sent to the final Linear
layer, which is followed by a log-Softmax function.

```python
import math
import torch
import torch.nn as nn
import torch.nn.functional as F

class TransformerModel(nn.Module):

    def __init__(self, ntoken, ninp, nhead, nhid, nlayers, dropout=0.5):
        super(TransformerModel, self).__init__()
        from torch.nn import TransformerEncoder, TransformerEncoderLayer
        self.model_type = 'Transformer'
        self.src_mask = None
        self.pos_encoder = PositionalEncoding(ninp, dropout)
        encoder_layers = TransformerEncoderLayer(ninp, nhead, nhid, dropout)
        self.transformer_encoder = TransformerEncoder(encoder_layers, nlayers)
        self.encoder = nn.Embedding(ntoken, ninp)
        self.ninp = ninp
        self.decoder = nn.Linear(ninp, ntoken)

        self.init_weights()

    def _generate_square_subsequent_mask(self, sz):
        mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
        return mask

    def init_weights(self):
        initrange = 0.1
        self.encoder.weight.data.uniform_(-initrange, initrange)
        self.decoder.bias.data.zero_()
        self.decoder.weight.data.uniform_(-initrange, initrange)

    def forward(self, src):
        if self.src_mask is None or self.src_mask.size(0) != len(src):
            device = src.device
            mask = self._generate_square_subsequent_mask(len(src)).to(device)
            self.src_mask = mask

        src = self.encoder(src) * math.sqrt(self.ninp)
        src = self.pos_encoder(src)
        output = self.transformer_encoder(src, self.src_mask)
        output = self.decoder(output)
        return output
```

``PositionalEncoding`` module injects some information about the
relative or absolute position of the tokens in the sequence. The
positional encodings have the same dimension as the embeddings so that
the two can be summed. Here, we use ``sine`` and ``cosine`` functions of
different frequencies.

```python
class PositionalEncoding(nn.Module):

    def __init__(self, d_model, dropout=0.1, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)

        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0).transpose(0, 1)
        self.register_buffer('pe', pe)

    def forward(self, x):
        x = x + self.pe[:x.size(0), :]
        return self.dropout(x)
```

Load and batch data
-------------------

The training process uses Wikitext-2 dataset from ``torchtext``. The
vocab object is built based on the train dataset and is used to numericalize
tokens into tensors. Starting from sequential data, the ``batchify()``
function arranges the dataset into columns, trimming off any tokens remaining
after the data has been divided into batches of size ``batch_size``.
For instance, with the alphabet as the sequence (total length of 26)
and a batch size of 4, we would divide the alphabet into 4 sequences of
length 6:

\begin{align}\begin{bmatrix}
  \text{A} & \text{B} & \text{C} & \ldots & \text{X} & \text{Y} & \text{Z}
  \end{bmatrix}
  \Rightarrow
  \begin{bmatrix}
  \begin{bmatrix}\text{A} \\ \text{B} \\ \text{C} \\ \text{D} \\ \text{E} \\ \text{F}\end{bmatrix} &
  \begin{bmatrix}\text{G} \\ \text{H} \\ \text{I} \\ \text{J} \\ \text{K} \\ \text{L}\end{bmatrix} &
  \begin{bmatrix}\text{M} \\ \text{N} \\ \text{O} \\ \text{P} \\ \text{Q} \\ \text{R}\end{bmatrix} &
  \begin{bmatrix}\text{S} \\ \text{T} \\ \text{U} \\ \text{V} \\ \text{W} \\ \text{X}\end{bmatrix}
  \end{bmatrix}\end{align}

These columns are treated as independent by the model, which means that
the dependence of ``G`` and ``F`` can not be learned, but allows more
efficient batch processing.

```python
!pwd
```

    /Users/xuming06/Codes/python-tutorial

```python
import torchtext
from torchtext.data.utils import get_tokenizer
TEXT = torchtext.data.Field(init_token='<sos>',
                            eos_token='<eos>',
                            lower=True)
import os
train_txt, val_txt, test_txt = torchtext.datasets.WikiText2.splits(TEXT,root='./20pytorch/.data',)
TEXT.build_vocab(train_txt)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def batchify(data, bsz):
    data = TEXT.numericalize([data.examples[0].text])
    # Divide the dataset into bsz parts.
    nbatch = data.size(0) // bsz
    # Trim off any extra elements that wouldn't cleanly fit (remainders).
    data = data.narrow(0, 0, nbatch * bsz)
    # Evenly divide the data across the bsz batches.
    data = data.view(bsz, -1).t().contiguous()
    return data.to(device)

batch_size = 20
eval_batch_size = 10
train_data = batchify(train_txt, batch_size)
val_data = batchify(val_txt, eval_batch_size)
test_data = batchify(test_txt, eval_batch_size)
```

Functions to generate input and target sequence
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

``get_batch()`` function generates the input and target sequence for
the transformer model. It subdivides the source data into chunks of
length ``bptt``. For the language modeling task, the model needs the
following words as ``Target``. For example, with a ``bptt`` value of 2,
we’d get the following two Variables for ``i`` = 0:

It should be noted that the chunks are along dimension 0, consistent
with the ``S`` dimension in the Transformer model. The batch dimension
``N`` is along dimension 1.

```python
bptt = 35
def get_batch(source, i):
    seq_len = min(bptt, len(source) - 1 - i)
    data = source[i:i+seq_len]
    target = source[i+1:i+1+seq_len].view(-1)
    return data, target
```

Initiate an instance
--------------------

The model is set up with the hyperparameter below. The vocab size is
equal to the length of the vocab object.

```python
ntokens = len(TEXT.vocab.stoi) # the size of vocabulary
emsize = 200 # embedding dimension
nhid = 200 # the dimension of the feedforward network model in nn.TransformerEncoder
nlayers = 2 # the number of nn.TransformerEncoderLayer in nn.TransformerEncoder
nhead = 2 # the number of heads in the multiheadattention models
dropout = 0.2 # the dropout value
model = TransformerModel(ntokens, emsize, nhead, nhid, nlayers, dropout).to(device)
```

Run the model
-------------

`CrossEntropyLoss <https://pytorch.org/docs/master/nn.html?highlight=crossentropyloss#torch.nn.CrossEntropyLoss>`__
is applied to track the loss and
`SGD <https://pytorch.org/docs/master/optim.html?highlight=sgd#torch.optim.SGD>`__
implements stochastic gradient descent method as the optimizer. The initial
learning rate is set to 5.0. `StepLR <https://pytorch.org/docs/master/optim.html?highlight=steplr#torch.optim.lr_scheduler.StepLR>`__ is
applied to adjust the learn rate through epochs. During the
training, we use
`nn.utils.clip_grad_norm\_ <https://pytorch.org/docs/master/nn.html?highlight=nn%20utils%20clip_grad_norm#torch.nn.utils.clip_grad_norm_>`__
function to scale all the gradient together to prevent exploding.

```python
criterion = nn.CrossEntropyLoss()
lr = 5.0 # learning rate
optimizer = torch.optim.SGD(model.parameters(), lr=lr)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 1.0, gamma=0.95)

import time
def train():
    model.train() # Turn on the train mode
    total_loss = 0.
    start_time = time.time()
    ntokens = len(TEXT.vocab.stoi)
    for batch, i in enumerate(range(0, train_data.size(0) - 1, bptt)):
        data, targets = get_batch(train_data, i)
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output.view(-1, ntokens), targets)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
        optimizer.step()

        total_loss += loss.item()
        log_interval = 200
        if batch % log_interval == 0 and batch > 0:
            cur_loss = total_loss / log_interval
            elapsed = time.time() - start_time
            print('| epoch {:3d} | {:5d}/{:5d} batches | '
                  'lr {:02.2f} | ms/batch {:5.2f} | '
                  'loss {:5.2f} | ppl {:8.2f}'.format(
                    epoch, batch, len(train_data) // bptt, scheduler.get_lr()[0],
                    elapsed * 1000 / log_interval,
                    cur_loss, math.exp(cur_loss)))
            total_loss = 0
            start_time = time.time()

def evaluate(eval_model, data_source):
    eval_model.eval() # Turn on the evaluation mode
    total_loss = 0.
    ntokens = len(TEXT.vocab.stoi)
    with torch.no_grad():
        for i in range(0, data_source.size(0) - 1, bptt):
            data, targets = get_batch(data_source, i)
            output = eval_model(data)
            output_flat = output.view(-1, ntokens)
            total_loss += len(data) * criterion(output_flat, targets).item()
    return total_loss / (len(data_source) - 1)
```

Loop over epochs. Save the model if the validation loss is the best
we've seen so far. Adjust the learning rate after each epoch.

```python
best_val_loss = float("inf")
epochs = 3 # The number of epochs
best_model = None
MODEL_PATH = 'transformer_lm.pth'
for epoch in range(1, epochs + 1):
    epoch_start_time = time.time()
    train()
    val_loss = evaluate(model, val_data)
    print('-' * 89)
    print('| end of epoch {:3d} | time: {:5.2f}s | valid loss {:5.2f} | '
          'valid ppl {:8.2f}'.format(epoch, (time.time() - epoch_start_time),
                                     val_loss, math.exp(val_loss)))
    print('-' * 89)

    if val_loss < best_val_loss:
        best_val_loss = val_loss
        best_model = model
        torch.save(best_model.state_dict(), MODEL_PATH)

    scheduler.step()

best_model.load_state_dict(torch.load(MODEL_PATH))
```

    | epoch   1 |   200/ 2983 batches | lr 5.00 | ms/batch 616.60 | loss  8.05 | ppl  3133.19
    | epoch   1 |   400/ 2983 batches | lr 5.00 | ms/batch 576.70 | loss  6.81 | ppl   905.95
    | epoch   1 |   600/ 2983 batches | lr 5.00 | ms/batch 571.79 | loss  6.40 | ppl   601.20
    | epoch   1 |   800/ 2983 batches | lr 5.00 | ms/batch 623.83 | loss  6.27 | ppl   530.66
    | epoch   1 |  1000/ 2983 batches | lr 5.00 | ms/batch 560.21 | loss  6.17 | ppl   477.33
    | epoch   1 |  1200/ 2983 batches | lr 5.00 | ms/batch 559.55 | loss  6.14 | ppl   466.38
    | epoch   1 |  1400/ 2983 batches | lr 5.00 | ms/batch 563.76 | loss  6.08 | ppl   438.14
    | epoch   1 |  1600/ 2983 batches | lr 5.00 | ms/batch 571.45 | loss  6.10 | ppl   444.22
    | epoch   1 |  1800/ 2983 batches | lr 5.00 | ms/batch 566.21 | loss  5.98 | ppl   395.20
    | epoch   1 |  2000/ 2983 batches | lr 5.00 | ms/batch 572.62 | loss  5.98 | ppl   397.28
    | epoch   1 |  2200/ 2983 batches | lr 5.00 | ms/batch 578.50 | loss  5.90 | ppl   364.42
    | epoch   1 |  2400/ 2983 batches | lr 5.00 | ms/batch 775.22 | loss  5.93 | ppl   374.94
    | epoch   1 |  2600/ 2983 batches | lr 5.00 | ms/batch 651.06 | loss  5.93 | ppl   377.25
    | epoch   1 |  2800/ 2983 batches | lr 5.00 | ms/batch 737.84 | loss  5.84 | ppl   342.14

Evaluate the model with the test dataset
-------------------------------------

Apply the best model to check the result with the test dataset.

```python
test_loss = evaluate(best_model, test_data)
print('=' * 89)
print('| End of training | test loss {:5.2f} | test ppl {:8.2f}'.format(
    test_loss, math.exp(test_loss)))
print('=' * 89)
```