###### Text provided under a Creative Commons Attribution license, CC-BY. All code is made available under the FSF-approved MIT license. (c) Daniel Koehn based on Jupyter notebooks by Marc Spiegelman [Dynamical Systems APMA 4101](https://github.com/mspieg/dynamical-systems) and  Kyle Mandli from his course [Introduction to numerical methods](https://github.com/mandli/intro-numerical-methods), notebook style sheet by L.A. Barba, N.C. Clementi [Engineering Computations](https://github.com/engineersCode)

```python
# Execute this cell to load the notebook's style sheet, then ignore it
from IPython.core.display import HTML
css_file = '../style/custom.css'
HTML(open(css_file, "r").read())
```

<link href="https://fonts.googleapis.com/css?family=Merriweather:300,300i,400,400i,700,700i,900,900i" rel='stylesheet' >
<link href="https://fonts.googleapis.com/css?family=Source+Sans+Pro:300,300i,400,400i,700,700i" rel='stylesheet' >
<link href='http://fonts.googleapis.com/css?family=Source+Code+Pro:300,400' rel='stylesheet' >
<style>

@font-face {
    font-family: "Computer Modern";
    src: url('http://mirrors.ctan.org/fonts/cm-unicode/fonts/otf/cmunss.otf');
}

#notebook_panel { /* main background */
    background: rgb(245,245,245);
}

div.cell { /* set cell width */
    width: 800px;
}

div #notebook { /* centre the content */
    background: #fff; /* white background for content */
    width: 1000px;
    margin: auto;
    padding-left: 0em;
}

#notebook li { /* More space between bullet points */
margin-top:0.5em;
}

/* draw border around running cells */
div.cell.border-box-sizing.code_cell.running {
    border: 1px solid #111;
}

/* Put a solid color box around each cell and its output, visually linking them*/
div.cell.code_cell {
    background-color: rgb(256,256,256);
    border-radius: 0px;
    padding: 0.5em;
    margin-left:1em;
    margin-top: 1em;
}

div.text_cell_render{
    font-family: 'Source Sans Pro', sans-serif;
    line-height: 140%;
    font-size: 110%;
    width:680px;
    margin-left:auto;
    margin-right:auto;
}

/* Formatting for header cells */
.text_cell_render h1 {
    font-family: 'Merriweather', serif;
    font-style:regular;
    font-weight: bold;
    font-size: 250%;
    line-height: 100%;
    color: #004065;
    margin-bottom: 1em;
    margin-top: 0.5em;
    display: block;
}
.text_cell_render h2 {
    font-family: 'Merriweather', serif;
    font-weight: bold;
    font-size: 180%;
    line-height: 100%;
    color: #0096d6;
    margin-bottom: 0.5em;
    margin-top: 0.5em;
    display: block;
}

.text_cell_render h3 {
    font-family: 'Merriweather', serif;
	font-size: 150%;
    margin-top:12px;
    margin-bottom: 3px;
    font-style: regular;
    color: #008367;
}

.text_cell_render h4 {    /*Use this for captions*/
    font-family: 'Merriweather', serif;
    font-weight: 300;
    font-size: 100%;
    line-height: 120%;
    text-align: left;
    width:500px;
    margin-top: 1em;
    margin-bottom: 2em;
    margin-left: 80pt;
    font-style: regular;
}

.text_cell_render h5 {  /*Use this for small titles*/
    font-family: 'Source Sans Pro', sans-serif;
    font-weight: regular;
    font-size: 130%;
    color: #e31937;
    font-style: italic;
    margin-bottom: .5em;
    margin-top: 1em;
    display: block;
}

.text_cell_render h6 { /*use this for copyright note*/
    font-family: 'Source Code Pro', sans-serif;
    font-weight: 300;
    font-size: 9pt;
    line-height: 100%;
    color: grey;
    margin-bottom: 1px;
    margin-top: 1px;
}

    .CodeMirror{
            font-family: "Source Code Pro";
			font-size: 90%;
    }
/*    .prompt{
        display: None;
    }*/

    .warning{
        color: rgb( 240, 20, 20 )
        }
</style>

# Stationary Solutions of Time-Dependent Problems

Often, we do not want to compute the whole time evolution of a dynamical problem, but are only interested in **stationary solutions**, where the system reached an equilibrium state. These stationary solutions are estimated by setting all time derivatives in the differential equation to zero. As examples, we will compute the stationary solutions for streamlines from [Exercise: How to sail without wind](https://nbviewer.jupyter.org/github/daniel-koehn/Differential-equations-earth-system/blob/master/02_finite_difference_intro/3_fd_ODE_example_sailing_wo_wind.ipynb) and the Lorenz equations.

## Stationary Solutions of Streamlines

In [Exercise: How to sail without wind](https://nbviewer.jupyter.org/github/daniel-koehn/Differential-equations-earth-system/blob/master/02_finite_difference_intro/3_fd_ODE_example_sailing_wo_wind.ipynb), we computed streamlines for a given velocity vector field ${\bf{V}} = (v_x,v_y)^T$, starting from an initial position ${\bf{x_0}} = (x_0,y_0)^T$. The governing equations related componentwise the velocity field to the spatial coordinates ${\bf{x}} = (x,y)^T$ by

\begin{equation}
v_x = \frac{dx}{dt},\; v_y = \frac{dy}{dt} \tag{1}
\end{equation}

To find the stationary solutions of differential eqs. (1), we have to set the time derivatives on the RHS to zero $\frac{dx}{dt}=0$, $\frac{dy}{dt}=0$ leading to

\begin{equation}
v_x = 0,\; v_y = 0 \tag{2}
\end{equation}

Eqs. (2) simply state that a stationary solution is located at points, where the velocity field is zero or we can not change the system state, because no fluid flow exists. To explicitly calculate the stationary solutions of eqs.(1), also known as **Fixpoints**, we have to explicitly define the velocity field ${\bf{V}} = (v_x,v_y)^T$.

##### Exercise 1

Compute the fixpoints ${\bf{x_{fix}}} = (x_{fix},y_{fix})^T$ for the velocity fields

\begin{equation}
{\bf{V_1}} = (y,-x)^T \notag
\end{equation}

and

\begin{equation}
{\bf{V_2}}=(cos((x+1y)/500),sin((x-1y)/500))^T \notag
\end{equation}

- How many fixpoints exist for each velocity field $\bf{V_1}$, $\bf{V_2}$?
- Mark all fixpoint positions on the square $-1000 \le x \le 1000;\; -1000 \le y \le 1000$ by red dots in the vector plots below.

```python
import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
from pylab import rcParams

# Ignore Warning Messages
# -----------------------
import warnings
warnings.filterwarnings("ignore")

# define font size
FSize = 18
font = {'color':  'black',
        'weight': 'normal',
        'size': FSize}
mpl.rc('xtick', labelsize=FSize)
mpl.rc('ytick', labelsize=FSize)
```

```python
# define functions to compute the vector fields V1 and V2
# compute velocity components V = (vx,vy)^T at position x,y
def vel_xy_1(x,y):

    vx = y / 1000.
    vy = -x / 1000.

    return vx, vy

def vel_xy_2(x,y):

    vx = np.cos((x+1*y)/500)
    vy = np.sin((x-1*y)/500)

    return vx, vy
```

```python
# Plot velocity vector fields
# Define figure size
rcParams['figure.figsize'] = 9, 4

# Define coordinates
dh = 100.
x1 = -1000.
x2 = 1000.
X, Y = np.meshgrid(np.arange(x1, x2, dh), np.arange(x1, x2, dh))

# Plot vector field V1
# --------------------
ax1 = plt.subplot(1, 2, 1)

# Define vector field components for coordinates X,Y
VX1,VY1 = vel_xy_1(X,Y)

ax1.set_title(r'$V=(y,-x)^T/1000$')
plt.axis('equal')
Q = plt.streamplot(X,Y,VX1,VY1)
# PLOT FIXPOINT POSITIONS HERE!

plt.xlabel('x [m]')
plt.ylabel('y [m]')

# Plot vector field V2
# --------------------
ax2 = plt.subplot(1, 2, 2)

# Define vector field components for coordinates X,Y
VX2,VY2 = vel_xy_2(X,Y)

ax2.set_title(r'$V=(cos((x+1y)/500),sin((x-1y)/500))^T$')
plt.axis('equal')
Q = plt.quiver(X,Y,VX2,VY2)
# PLOT FIXPOINT POSITIONS HERE!

plt.xlabel('x [m]')
plt.ylabel('y [m]')

plt.tight_layout()
plt.show()
```

##### Exercise 2

Compute the fixpoints ${\bf{x_{fix}}} = (X_{fix},Y_{fix},Z_{fix})^T$ for the Lorenz equations:

\begin{equation}
\begin{split}
\frac{\partial X}{\partial t} &= \sigma( Y - X)\\
\frac{\partial Y}{\partial t} &= rX - Y - XZ \\
\frac{\partial Z}{\partial t} &= XY -b Z
\end{split}
\notag
\end{equation}

where $\sigma$ denotes the "Prandtl number",  $r = \mathrm{Ra}/\mathrm{Ra}_c$ is a scaled "Rayleigh number" and $b$ is a parameter that is related to the the aspect ratio of a convecting cell in the original derivation.

- How many fixpoints exist?
- Use the fixpoint solutions $X_{fix}$, $Y_{fix}$ and $Z_{fix}$ to plot the corresponding streamfunction $\psi(x,z)$, temperature field $T(x,z)$ and velocity vector field ${\bf{v}}$

\begin{equation}
\begin{split}
\psi(x,z) &= X_{fix} \sin(a\pi x)\sin(\pi z)\\
T(x,z) &= Y_{fix} \cos(a\pi x)\sin(\pi z) - Z_{fix} \sin(2\pi z) + (1-z)\\
{\bf{v}} &=(u,0,w) = {\bf{\nabla\times\psi\hat{j}}}=(-\frac{\partial\psi}{\partial z}, 0,  \frac{\partial\psi}{\partial x})^T\\
\end{split}
\notag
\end{equation}

for $b=8/3$ and scaled Rayleigh numbers $r=0.5$ and $r=28$, respectively.

- Describe and interpret the different Fixpoint solutions

```python
# Define coordinates
a = np.sqrt(0.5)
x0 = np.linspace(0,1./a,20)
z0 = np.linspace(0.,1.,20)
x,z = np.meshgrid(x0,z0)

# Define spatial part of streamfunction and temperature field
psi = np.sin(a*np.pi*x)*np.sin(np.pi*z)
theta0 = np.cos(a*np.pi*x)*np.sin(np.pi*z)
theta1 = -np.sin(2.*np.pi*z)
```

## Fixpoints of Lorenz Equations (r = 0.5)

PLOT VECTORFIELD $V = (U,V)^T$ AND TEMPERATURE $T$ FOR ALL FIXPOINT SOLUTIONS WITH $r=0.5$ HERE! PRODUCE A NEW PLOT FOR EACH FIXPOINT!

```python
# Define figure size
plt.figure()
plt.figure(figsize=(20,20))

# DEFINE FIXPOINT HERE!
b = 8/3
r = 0.5

Xfix =
Yfix =
Zfix =

# Initial Streamfunction psi and velocity field
plt.subplot(1,2,1)
plt.contourf(x,z,Xfix*psi,cmap='viridis_r')

# Velocity field
U = - Xfix * np.pi * np.sin(a*np.pi*x) * np.cos(np.pi*z)
V = Xfix * a * np.pi * np.cos(a*np.pi*x) * np.sin(np.pi*z)
plt.quiver(x,z,U,V)
plt.gca().set_aspect('equal')
plt.title('Streamfunction $\psi$', fontdict=font)

# Initial temperature field
plt.subplot(1,2,2)
plt.contourf(x,z,Yfix*theta0 + Zfix*theta1 + (1-z),cmap='magma')
plt.gca().set_aspect('equal')
plt.title(r'Temperature $T$', fontdict=font)

plt.tight_layout()
plt.show()
```

## Fixpoints of Lorenz Equations (r = 28)

PLOT VECTORFIELD $V = (U,V)^T$ AND TEMPERATURE $T$ FOR ALL FIXPOINT SOLUTIONS WITH $r=28$ HERE! PRODUCE A NEW PLOT FOR EACH FIXPOINT!

```python
# Define figure size
plt.figure()
plt.figure(figsize=(20,20))

# DEFINE FIXPOINT HERE!
b = 8/3
r = 28

Xfix =
Yfix =
Zfix =

# Initial Streamfunction psi and velocity field
plt.subplot(1,2,1)
plt.contourf(x,z,Xfix*psi,cmap='viridis_r')

# Velocity field
U = - Xfix * np.pi * np.sin(a*np.pi*x) * np.cos(np.pi*z)
V = Xfix * a * np.pi * np.cos(a*np.pi*x) * np.sin(np.pi*z)
plt.quiver(x,z,U,V)
plt.gca().set_aspect('equal')
plt.title('Streamfunction $\psi$', fontdict=font)

# Initial temperature field
plt.subplot(1,2,2)
plt.contourf(x,z,Yfix*theta0 + Zfix*theta1 + (1-z),cmap='magma')
plt.gca().set_aspect('equal')
plt.title(r'Temperature $T$', fontdict=font)

plt.tight_layout()
plt.show()
```

## What we learned:

- How to compute stationary solutions (fixpoints) of differential equations.


# Paralelismo en Julia

## 1. Introducción

Julia proporciona distintos mecanismos para paralelizar código. Algunas de las estrategias y desafíos para escribir algoritmos paralelos son los siguientes:

* Estrategias de paralelismo
     * SIMD
     * Multi-hilo
     * Tareas
     * Multiproceso
         * Memoria compartida
         * Memoria distribuida
     * Programación de GPU

* Desafíos de la computación paralela
     * Orden de ejecución
         * ejecución de fuera de orden de posibilidad
         * acceso y mutación simultáneos
     * Acceso y movimiento de datos
     * Código de acceso y movimiento
     * Adaptación adecuada de la estrategia de paralelismo a las capacidades de su máquina
     * Hacer coincidir adecuadamente la estrategia de paralelismo con el problema en cuestión

## ¿Qué es lo que está sucediendo con nuestras computadoras?

## Lo difícil de la computación paralela
   * No pensamos en paralelo
   * Aprendemos a escribir y razonar sobre programas en serie.
   * El deseo de paralelismo a menudo surge _después_ de haber escrito su algoritmo (¡y lo encontró demasiado lento!)

## Resumen:
   * Las arquitecturas computacionales actuales nos empujan hacia la programación paralela para un rendimiento máximo, ¡incluso si no estamos en un clúster!
   * Pero es difícil diseñar buenos algoritmos paralelos
   * Es difícil de expresar y razonar sobre esos algoritmos.

## 2. SIMD: El paralelismo que puede (a veces) suceder automáticamente

SIMD: Instrucción única, datos múltiples (Single Instruction Multiple Data)

**Nota:** También llamado confusamente vectorización

### Arquitectura

En lugar de calcular cuatro sumas secuencialmente:

\begin{align}
x_1 + y_1 &\rightarrow z_1 \\
x_2 + y_2 &\rightarrow z_2 \\
x_3 + y_3 &\rightarrow z_3 \\
x_4 + y_4 &\rightarrow z_4
\end{align}

Procesadores modernos tienen unidades de procesamiento vectorial que pueden hacer lo anterior a la vez:

$$
\left(\begin{array}{cc}
x_1 \\
x_2 \\
x_3 \\
x_4
\end{array}\right)
+
\left(\begin{array}{cc}
y_1 \\
y_2 \\
y_3 \\
y_4
\end{array}\right)
\rightarrow
\left(\begin{array}{cc}
z_1 \\
z_2 \\
z_3 \\
z_4
\end{array}\right)
$$

### ¿Cómo se logra?

```julia
using BenchmarkTools
```

```julia
A = rand(100_000)
function simplesum(A)
    result = zero(eltype(A))
    for i in eachindex(A)
        @inbounds result += A[i]
    end
    return result
end

simplesum(A)
```

Como muchos lenguajes de programación modernos, Julia utiliza la verificación de límites ([_boundchecking_](https://docs.julialang.org/en/v1/devdocs/boundscheck/)) para garantizar la seguridad del programa al acceder a arreglos.
En bucles internos u otras situaciones críticas de rendimiento, es posible que se desee omitir estas comprobaciones de límites para mejorar el rendimiento en tiempo de ejecución.

En consecuencia, Julia incluye la macro `@inbounds(...)` para decirle al compilador que omita dichas comprobaciones de límites dentro del bloque dado.

```julia
@btime simplesum($A)
```

¿ese tiempo es bueno?

```julia
@btime sum($A)
```

Diseñamos una función más lenta que la suma prediseñada `sum()`, ¡y también estamos obteniendo una respuesta diferente! Veamos qué sucede con un flotante de 32 bits en lugar de uno de 64 bits. Cada elemento tiene la mitad del número de bits, por lo que también permite duplicar la longitud (para que el número total de bits procesados permanezca constante).

```julia
A32 = rand(Float32, length(A)*2)
@btime simplesum($A32)
@btime sum($A32)
```

¡Eso es aun peor! ¿Que está pasando aqui?
Estamos viendo múltiples diferencias en el desempeño: ¿quizás la suma incorporada de Julia está usando algún paralelismo?

Intentemos usar SIMD nosotros mismos:

```julia
function simdsum(A)
    result = zero(eltype(A))
    @simd for i in eachindex(A)
        @inbounds result += A[i]
    end
    return result
end
@btime simdsum($A)
@btime simdsum($A32)
```

¿Qué hizo y por qué no siempre usamos (usa Julia pues) `@simd` para cada bucle **for** automáticamente?

Veamos los resultados:

```julia
simplesum(A), simdsum(A), sum(A)
```

```julia
simplesum(A32), simdsum(A32), sum(A32)
```

¿Por qué no son iguales?

Sin `@simd`, Julia está haciendo _exactamente_ lo que le dijimos que hiciera: está tomando cada elemento de nuestro arreglo y lo agrega a una gran pila secuencialmente. Nuestra respuesta es más pequeña de lo que la "suma" incorporada de Julia cree que es: eso es porque, como la pila se hace más grande, comenzamos a perder las partes inferiores de cada elemento que estamos sumando, ¡y esas pequeñas pérdidas comienzan a acumularse!

La macro `@simd` le dice a Julia que puede reorganizar las adiciones de punto flotante -
incluso si cambiara la respuesta. Dependiendo de su CPU, esto puede llevar a 2x o 4x
o incluso un paralelismo 8x. Básicamente, Julia está calculando sumas independientes para
los índices pares y los índices impares simultáneamente:

\begin{align}
odds &\leftarrow 0 \\
evens &\leftarrow 0 \\
\text{loop}&\ \text{odd}\ i: \\
    &\left(\begin{array}{cc}
odds \\
evens
\end{array}\right)
\leftarrow
\left(\begin{array}{cc}
odds \\
evens
\end{array}\right)
+
\left(\begin{array}{cc}
x_{i} \\
x_{i+1}
\end{array}\right) \\
total &\leftarrow evens + odds
\end{align}

En muchos casos, Julia puede y sabe que un bucle for puede ser vectorizado (SIMD-ed) y aprovechará esto por defecto.

```julia
B = rand(1:10, 100_000)
@btime simplesum($B)
@btime sum($B)
B32 = rand(Int32(1):Int32(10), length(B)*2)
@btime simplesum($B32)
@btime simdsum($B32)
```

¿Cómo inspeccionamos que se está vectorizando?

```julia
@code_llvm simdsum(A32)
```

Entonces, ¿cuáles son los desafíos?:

- El mayor obstáculo es que tienes que convencer a Julia y LLVM de que puede usar instrucciones SIMD para tu algoritmo dado. Eso no siempre es posible.
- Hay muchas limitaciones de lo que se puede y no se puede vectorizar
- Es necesario pensar en las consecuencias de reordenar su algoritmo

## Resumen

SIMD:
- Explota el paralelismo integrado en un procesador
- Ideal para bucles internos pequeños (y estrechos)
- A menudo ocurre automáticamente si tienes cuidado
    - Sigue las [mejores prácticas de rendimiento](https://docs.julialang.org/en/v1/manual/performance-tips/)
    - Usa `@inbounds` para cualquier acceso a un arreglo
    - Evita ramas o llamadas a funciones
- Dependiendo del procesador y los tipos involucrados, puede producir ganancias de 2-16x con una sobrecarga extraordinariamente pequeña

```julia

```

Curso Julia/Github - Material/src/Introduccion/S8-OperacionesVectorizadas.ipynb


# Models and Pricing of Financial Derivatives HW_02

**11510691 程远星$\DeclareMathOperator*{\argmin}{argmin}
\DeclareMathOperator*{\argmax}{argmax}
\DeclareMathOperator*{\plim}{plim}
\newcommand{\using}[1]{\stackrel{\mathrm{#1}}{=}}
\newcommand{\ffrac}{\displaystyle \frac}
\newcommand{\asim}{\overset{\text{a}}{\sim}}
\newcommand{\space}{\text{ }}
\newcommand{\bspace}{\;\;\;\;}
\newcommand{\QQQ}{\boxed{?\:}}
\newcommand{\void}{\left.\right.}
\newcommand{\Tran}[1]{{#1}^{\mathrm{T}}}
\newcommand{\d}[1]{\displaystyle{#1}}
\newcommand{\CB}[1]{\left\{ #1 \right\}}
\newcommand{\SB}[1]{\left[ #1 \right]}
\newcommand{\P}[1]{\left( #1 \right)}
\newcommand{\abs}[1]{\left| #1 \right|}
\newcommand{\norm}[1]{\left\| #1 \right\|}
\newcommand{\dd}{\mathrm{d}}
\newcommand{\Exp}{\mathrm{E}}
\newcommand{\RR}{\mathbb{R}}
\newcommand{\EE}{\mathbb{E}}
\newcommand{\NN}{\mathbb{N}}
\newcommand{\ZZ}{\mathbb{Z}}
\newcommand{\QQ}{\mathbb{Q}}
\newcommand{\AcA}{\mathscr{A}}
\newcommand{\FcF}{\mathscr{F}}
\newcommand{\Var}[2][\,\!]{\mathrm{Var}_{#1}\left[#2\right]}
\newcommand{\Avar}[2][\,\!]{\mathrm{Avar}_{#1}\left[#2\right]}
\newcommand{\Cov}[2][\,\!]{\mathrm{Cov}_{#1}\left(#2\right)}
\newcommand{\Corr}[2][\,\!]{\mathrm{Corr}_{#1}\left(#2\right)}
\newcommand{\I}[1]{\mathrm{I}\left( #1 \right)}
\newcommand{\N}[1]{\mathcal{N} \left( #1 \right)}
\newcommand{\ow}{\text{otherwise}}
\void^\dagger$**

## Question 1

A stock price is currently $\$40$. It is known that at the end of $3$ months it will be either $\$45$ or $\$35$. The risk-free rate of interest with ***quarterly compounding*** is $8\%$ per annum. Calculate the value of a $3$-month European put option on the stock with an exercise price of $\$40$. Verify that no-arbitrage arguments and risk-neutral valuation arguments give the same answers.

$\bspace Answer$

>Consider a portfolio with $x$ much bounds and $y$ units of asset which is used to replicate the put, then
>
>$$\begin{cases}
1.02x + 45y = \P{40 - 45}^+ = 0\\
1.02x + 35y = \P{40 - 35}^+ = 5
\end{cases}$$
>
>$$x = \ffrac{5\times45-0\times35} {1.02\times\P{45 - 35}}=22.059,y = \ffrac{0-5} {45-35} =-0.5$$
>
>Or, first find the up and down rate, $u=45/40 = 1.125$ and $d = 35/40 = 0.857$, then
>
>$$x = \ffrac{1.125\times5 - 0.875\times0} {\P{1.125-0.875} 1.02}=22.059,y = \ffrac{0-5} {\P{1.125-0.875}40} =-0.5$$
>
>So the price of the put option: $p = x+yS_0 = 22.059-0.5\times40 = 2.059$
>
>***
>
>And using the riskless hedging principle, suppose to hedge, we need to *short* $\Delta$ shares, then the terminal payoff would be either $-35\Delta+5$ or $-45\Delta$. If
>
>$$-35\Delta + 5 = -45\Delta$$
>
>the portfolio would be riskless and the terminal payoff would be $22.5$ with $\Delta = -0.5$. And discount that to time $0$, so we have
>
>$$-40\Delta + p = \ffrac{22.5}{1.02} \Rightarrow p = 22.5/1.02 + 40\P{-0.5} = 2.059$$
>
>Or by risk-neutral method, first we find the risk-neutral probability $q$:
>
>$$q = \ffrac{1.02-0.875}{1.125-0.875} =0.58$$
>
>Then the put option price: $p = \ffrac{1}{1.02}\P{0\times 0.58+5\times 0.42}=2.059$

$Remark$

>The risk-free rate of interest with quarterly compounding is $8\%$ per annum, meaning that $e^{rT} = 1.02$.

## Question 2

Consider the situation in which stock price movements during the life of a European option are governed by a *two-step binomial tree*. Explain why it is not possible to set up a position in the stock and the option that remains *riskless* for the *whole life* of the option.

$\bspace Answer$

>When the stock price enters the second phase, the stock price will change, usually, so that the riskless hedging principle fails after that. To continue doing so, we need to find the new value of the number of shares to hold.

## Question 3

A stock price is currently $\$50$. Over each of the next two $3$-month periods it is expected to go up by $6\%$ or down by $5\%$. The risk-free interest rate is $5\%$ per annum with ***continuous compounding***. What is the value of a $6$-month European call option with a strike price of $\$51$?

$\bspace Answer$

>The risk-neutral probability of an *up move*, $q$, is given by
>
>$$q = \ffrac{e^{rT} - d}{u-d} = \ffrac{e^{0.05\times 0.25-0.95}} {1.06-0.95} = 0.569$$
>
>And its value keeped for each perioed since the parameters never change. So at the end of the first phase, the option price would be
>
>$$c_1\P{1} = e^{-0.05\times0.25}\SB{\P{50\times1.06^2-51}\times q} = 2.911$$
>
>And then the first phase, we have
>
>$$c_0\P{0} = e^{-0.05\times0.25}\SB{c_1\P{1}\times q} = 1.636$$
>
>And about the formula for $q$, here's the proof:
>
>$$\begin{cases}
xe^{rT} + yuS_0 = \P{yuS_0 - K}^+ = c_u\\
xe^{rT} + ydS_0 = \P{ydS_0 - K}^+ = c_d
\end{cases}\Rightarrow \begin{cases}
x = \ffrac{uc_d - dc_u}{\P{u-d}e^{rT}}\leq 0\\
y = \ffrac{c_u-c_d}{\P{u-d}S_0}\geq 0
\end{cases}$$
>
>Then $c = x+yS_0 = e^{-rT}\P{q\cdot c_u + \P{1-q}c_d}$ where $q = \ffrac{e^{rT} - d}{u-d}$.

## Question 4

For the situation considered in HW_2.3, what is the value of a $6$-month European *put* option with a strike price of $\$51$? Verify that the *European call* and *European put* prices satisfy put-call parity. If the put option were *American*, would it ever be optimal to exercise it early at any of the nodes on the tree?

$\bspace Answer$

>We now share the $q$ with the last question so we can directly use the path probability method:
>
>$$p = e^{-0.05\times0.5}\SB{\P{1-q}^2\P{51 - 50\times0.95^2} + \binom{2} {1} q\P{1-q}\P{51-50 \times 0.95 \times 1.06}} = 1.375$$
>
>And to test the put-call parity:
>
>$$1.636 + 51 \times e^{-0.05\times0.5}-1.375-50 \approx 0$$
>
>That's it. And if it's American call options (without dividend), we won't early exercise cause it's never optimal while if it's American put option, first, we won't do that if the price goes up, or at the begining. Second, if the price goes down in the first period, at that time the payoff from early exercise is $51-50\times 0.95 = 3.5$ while the American put option value at that time is
>
>$$P_1\P{0} = e^{-0.05\times0.25}\SB{\P{1-q}\P{51 - 50\times0.95^2}+q\times\P{51-50 \times 0.95 \times 1.06}} = 2.866 < 3.5$$
>
>That's to say that it's better to early exercise the American Put option if the stock price goes down in the first phase.

## Question 5

A stock price is currently $25$ bucks. It is known that at the end of $2$ months it will be either $\$23$ or $\$27$. The risk-free interest rate is $10\%$ per annum with continuous compounding. Suppose $S_T$ is the stock price at the end of $2$ months. What is the value of a derivative that pays off $S^2_T$ at this time?

$\bspace Answer$

>We replicate that portfolio with a claim which is priced as $V_t^h$, with $x$ amount of cash and $y$ amount of the stocks. Then using no arbitrage pricing, we have
>
>$$\begin{cases}
x\cdot e^{0.1\times 1/6} + y \times 27 = 27^2 \\
x\cdot e^{0.1\times 1/6} + y \times 23 = 23^2
\end{cases} \\[2em]
\Rightarrow x= \ffrac{-27\times23} {e^{0.1 \times 1/6}},\bspace y = 50$$
>
>So that the price of that claim would be
>
>$$\Pi\P{\ffrac{1} {6}, U} = V_{1/6}^{h} = x + y\cdot S_0 = \ffrac{-27\times23} {e^{0.1 \times 1/6}} + 50 \times 25 = 639.264$$
>
>***
>
>Or using the riskless hedging principle, consider portfolio that longs this derivative and short $\Delta$ shares. The value at time $T$ is $27^2 - 27\Delta = 23^2 - 23\Delta$ and thus $\Delta = 50$, the value is $-621$. Let $f$ be the price of the derivative at time $0$, then
>
>$$\P{f - 25\Delta}e^{rT} = \P{f- 50\times 25}e^{0.1\times 1/6} = -621 \Rightarrow f = 639.264$$

## Question 6

A stock price is currently $\$40$. Over each of the next two $3$-month periods it is expected to go up by $10\%$ or down by $10\%$. The risk-free interest rate is $12\%$ per annum with continuous compounding.

$\P{\text a}$ What is the value of a $6$-month European put option with a strike price of $\$42$?

$\bspace Answer$

>$$q = \ffrac{e^{0.12 \times 0.25} - 0.9} {1.1-0.9} = 0.652$$
>
>The European put option can be find similarly:
>
>$$p = e^{-0.12\times 0.5} \SB{\P{1-q}^2 \P{42-40\times0.9^2} + \binom{2} {1} \P{1-q}q\P{42-40 \times 0.9 \times 0.1}}=2.121$$

$\P{\text b}$ What is the value of a $6$-month American put option with a strike price of $\$42$?

$\bspace Answer$

>We will calculate the American Put option backwards.
>
>$$\begin{align}
P_{1}\P{1} &= \max\P{e^{-0.12\times 0.25}\SB{\P{1-q}\P{42-40 \times 0.9 \times 1.1}},0} = 0.811 \\
P_{1}\P{0} &= \max\P{e^{-0.12\times 0.25}\SB{q\P{42-40 \times 0.9 \times 1.1} + \P{1-q}\P{42-40\times0.9^2}},42-40\times0.9} = 6 \\
P_{0}\P{0} &= \max\P{e^{-0.12\times 0.25}\SB{q\cdot P_{1}\P{1} + \P{1-q}\cdot P_{1}\P{0}},42-40} = 2.539
\end{align}
$$

## Question 7

A stock price is currently $30$ bucks. During each $2$-month period for the next $4$ months it will increase by $8\%$ or reduce by $10\%$. The risk-free interest rate is $5\%$. Use a two-step tree to calculate the value of a derivative that pays off $\P{\max\P{30 − S_T , 0}}^2$ , where $S_T$ is the stock price in $4$ months. If the derivative is American-style, should it be exercised early?

$\bspace Answer$

>It's not hard to find the value of the tree in the last column: $0$, $\P{30-29.16}^2 = 0.7056$ and $\P{30-24.3}^2 = 32.49$. Then we calculate the probability $q$:
>
>$$q = \ffrac{e^{0.05 \times 1/6}-0.9} {1.08-0.9} = 0.602$$
>
>Then for the European styled derivatives, the option value is:
>
>$$e^{-0.05\times1/3}\SB{\P{1-q}^2\times 32.49+\binom{2} {1}\P{1-q}q\times0.7056} = 5.394$$
>
>While for the American styled, firstly, like the same, it won't happen at the end of second phase, nor the begining. Then at the end of the first phase, if the stock price goes up in the first period, then still nothing should be done. If it goes down, then the payoff from early exercise is $\P{30-30\times0.9}^2=9$ however the value of American Put option at that time is
>
>$$P_{1}\P{0} = \max\P{e^{-0.05\times1/6}\SB{q\times 0.7056 + \P{1-q}\times 32.49},9} = 13.245>9$$
>
>So, even if it's an American styled derivatives, at no point of time should you early exercise it.

FinMath/Models and Pricing of Financial Derivatives/HW/HW_02_11510691_fixed.ipynb


# Create Spark Cluster

```python
from pyspark.sql import SparkSession
spark = (
    SparkSession
    .builder
    .config("spark.master", "local[*]")
    .config("spark.driver.memory", "120g")
    .config("spark.driver.maxResultSize", "0")
    .enableHiveSupport()
    .getOrCreate()
)
```

```python
import socket
port = spark.sparkContext.uiWebUrl.split(":")[-1]
hostname = socket.gethostname()
print(f"https://{hostname}/jupyter/user/stevengs/proxy/{port}/jobs/")
```

    https://epyc.astro.washington.edu/jupyter/user/stevengs/proxy/4043/jobs/

# Load current cut of lightcurves

```python
import axs
catalog = axs.AxsCatalog(spark)
wtf = catalog.load("stevengs_cut_wtf")
```

Look at the light curve/dib schema

```python
print(wtf.columns)
```

    ['ps1_objid', 'ra', 'dec', 'mean_mag_g', 'mean_mag_r', 'mean_mag_i', 'ra_stddev', 'dec_stddev', 'ps1_gMeanPSFMag', 'ps1_rMeanPSFMag', 'ps1_iMeanPSFMag', 'ra_detections', 'dec_detections', 'mjd_g', 'mag_g', 'magerr_g', 'psfflux_g', 'psffluxerr_g', 'catflags_g', 'expid_g', 'rcID_g', 'fieldID_g', 'xpos_g', 'ypos_g', 'nobs_g', 'mjd_r', 'mag_r', 'magerr_r', 'psfflux_r', 'psffluxerr_r', 'catflags_r', 'expid_r', 'rcID_r', 'fieldID_r', 'xpos_r', 'ypos_r', 'nobs_r', 'mjd_i', 'mag_i', 'magerr_i', 'psfflux_i', 'psffluxerr_i', 'catflags_i', 'expid_i', 'rcID_i', 'fieldID_i', 'xpos_i', 'ypos_i', 'nobs_i', 'dip', 'zone', 'dup']

```python
# dip schema
wtf.select("dip").head(1)[0]
```

    Row(dip=Row(center_mjd=58282.18359375, start_mjd=58271.09765625, end_mjd=58293.70703125, length=22.608280181884766, guess_start_mjd=58274.85546875, guess_end_mjd=58286.29296875, guess_significance=7.1239423751831055, window_start_mjd=58263.421875, window_end_mjd=58297.73046875, significant_observation_count=5, significant_length=9.91579818725586, integral=0.730926513671875, integral_uncertainty=0.13865260779857635, significance=5.271639347076416, max_gap=10.001840591430664, max_gap_fraction=0.44239720702171326, core_count=6, core_start_mjd=58276.35546875, core_end_mjd=58286.26953125, core_length=9.91579818725586, core_significant_count=5, core_not_significant_count=1, core_not_significant_fraction=0.1666666716337204, ref_observation_count_before=10, ref_observation_count_after=218, ref_observation_count=228, ref_pull_std=1.4057750701904297, ref_large_pull_fraction=0.048245612531900406, ref_length_after=573.8453369140625, ref_length_before=57.98050308227539, ref_length=631.8258056640625, ref_length_fraction_after=16.726078033447266, ref_length_fraction_before=1.689978837966919, ref_length_fraction=18.416057586669922, deep_count=22, deep_length=15.858842849731445, deep_gap_fraction=0.6306790709495544, flip_count=2))

```python
from fit_utils import (
    fit_band_around_dip, plot_fit_result, make_udf_from_annotated_function,
    evaluate_in_dip, evaluate_around_dip, evaluate
)
```

# Test simple quadratic model

```python
def quadratic(x, a, b, c):
    return a * x**2 + b * x + c

# fit model to data in window around the dip
# fit r-band data
# fit in window that is expanded by 4.0x the dip width from dip.start_mjd and dip.end_mjd
quadratic_fit_df = fit_band_around_dip(wtf, quadratic, "r", 4.0)

# select a small subset of columns
quadratic_fit_df_slim = quadratic_fit_df.select(
    "ps1_objid",
    "ra",
    "dec",
    "zone",
    "mjd_r",
    "mag_r",
    "magerr_r",
    "dip",
    "window_r",
    "fit_r",
)

# only include fits that converged
quadratic_fit_df_slim_good = quadratic_fit_df_slim.where(
    quadratic_fit_df_slim['fit_r']['info']['good']
)
# evaluate sum square error inside the dip
fits_evaluated_df = evaluate_in_dip(quadratic_fit_df_slim_good, quadratic, "r")
# around the dip
fits_evaluated_df = evaluate_around_dip(fits_evaluated_df, quadratic, "r", 4.0)
# over all data
fits_evaluated_df = evaluate(fits_evaluated_df, quadratic, "r")
```

```python
lc = fits_evaluated_df.head(1)[0]
```

## in-dip

```python
plot_fit_result(lc['dip_window_r']['x'], lc['dip_window_r']['y'], lc['dip_window_r']['yerr'], lc['fit_r'], quadratic)
print(lc['model_error_in_dip_r'])
```

## around-dip

```python
plot_fit_result(lc['window_r']['x'], lc['window_r']['y'], lc['window_r']['yerr'], lc['fit_r'], quadratic)
print(lc['model_error_around_dip_r'])
```

## all-data

```python
plot_fit_result(lc['mjd_r'], lc['mag_r'], lc['magerr_r'], lc['fit_r'], quadratic)
print(lc['model_error_r'])
```

# Skew Normal

Description: https://en.wikipedia.org/wiki/Skew_normal_distribution

A combination of the standard normal PDF $\phi(x)$ and the standard  normal CDF $\Phi(x)$:

$$
f(x) = \frac{2}{\omega} \phi\left(\frac{x - \xi}{\omega} \right) \Phi\left(\alpha \left(\frac{x - \xi}{\omega} \right) \right)
$$
$\xi$ is the location, $\omega$ is the scale, $\alpha$ is the skew.

\begin{align*}
\phi(x) &= \frac{1}{2\pi} \exp\left( -\frac{x^2}{2} \right) \\
\Phi(x) &= \int_{-\infty}^{x} \phi(t) dt = \frac{1}{2} \left[ 1 + \text{erf}\left( \frac{x}{\sqrt{2}} \right) \right]
\end{align*}

```python
from models import skew_normal, skew_normal_p0, skew_normal_p0_udf
```

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

loc = 0
xscale = 1
yscale = 1
offset = 0
plt.rc("figure", figsize=[10, 8])
plt.rc("font", size=18)
for skew in [-2.5, -1,  0, 1, 2.5]:
    _x = np.linspace(-5, 5, 100)
    _y = skew_normal(_x, skew, loc, xscale, yscale, offset)
    plt.plot(_x, _y, label=f"skew = {skew}")
plt.legend()
plt.show()
```

Fit r-band data

```python
band = "r"
wiggle = 4.0

skew_normal_p0_column = skew_normal_p0_udf(
    wtf['mjd_r'], wtf['mag_r'], wtf['magerr_r'], wtf['dip']['start_mjd'],
    wtf['dip']['end_mjd'], wtf['dip']['integral']
)

skew_norm_fit_df = fit_band_around_dip(wtf, skew_normal, band, wiggle, p0=skew_normal_p0_column)

# select a small subset of columns
skew_norm_fit_df_slim = skew_norm_fit_df.select(
    "ps1_objid",
    "ra",
    "dec",
    "zone",
    "mjd_r",
    "mag_r",
    "magerr_r",
    "dip",
    "window_r",
    "fit_r",
)

# only include fits that converged
skew_norm_fit_df_slim_good = skew_norm_fit_df_slim.where(
    skew_norm_fit_df_slim['fit_r']['info']['good']
)
# evaluate sum square error inside the dip
skew_norm_fits_evaluated_df = evaluate_in_dip(skew_norm_fit_df_slim_good, skew_normal, band)
# around the dip
skew_norm_fits_evaluated_df = evaluate_around_dip(skew_norm_fits_evaluated_df, skew_normal, band, wiggle)
# over all data
skew_norm_fits_evaluated_df = evaluate(skew_norm_fits_evaluated_df, skew_normal, band)
```

```python
_lc = skew_norm_fits_evaluated_df.head(1)[0]
```

```python
plot_fit_result(
    _lc['dip_window_r']['x'], _lc['dip_window_r']['y'], _lc['dip_window_r']['yerr'],
    _lc['fit_r'], skew_normal, with_p0=True
)
print(_lc['model_error_in_dip_r'])
print(_lc['model_error_around_dip_r'])
print(_lc['model_error_r'])
```

```python
catalog.save_axs_table(skew_norm_fits_evaluated_df, "6_4_20_stevengs_skew_normal_fits_r_band")
```

# Top hat

\begin{align}
f(x) = \begin{cases}
d + c&\text{for $x\in[l - w/2, l + w/2]$}\\
c&\text{otherwise}\\
\end{cases}
\end{align}
where $d$ is the depth, $c$ is the offset, $l$ is the location, and $w$ is the width

```python
from models import top_hat, top_hat_p0, top_hat_p0_udf
```

```python
loc = 0
width = 1
depth = 1
offset = 0
plt.rc("figure", figsize=[10, 8])
plt.rc("font", size=18)
for width in [1, 2, 5, 6]:
    _x = np.linspace(-5, 5, 100)
    _y = top_hat(_x, loc, width, depth, offset)
    plt.plot(_x, _y, label=f"width = {width}")
plt.legend()
plt.show()
```

```python
band = "r"
wiggle = 4.0

top_hat_p0_column = top_hat_p0_udf(
    wtf['mjd_r'], wtf['mag_r'], wtf['magerr_r'],
    wtf['dip']['start_mjd'], wtf['dip']['end_mjd'], wtf['dip']['integral']
)

top_hat_fit_df = fit_band_around_dip(wtf, top_hat, band, wiggle, p0=top_hat_p0_column)

# select a small subset of columns
top_hat_fit_df_slim = top_hat_fit_df.select(
    "ps1_objid",
    "ra",
    "dec",
    "zone",
    "mjd_r",
    "mag_r",
    "magerr_r",
    "dip",
    "window_r",
    "fit_r",
)

# only include fits that converged
top_hat_fit_df_slim_good = top_hat_fit_df_slim.where(
    top_hat_fit_df_slim['fit_r']['info']['good']
)
# evaluate sum square error inside the dip
top_hat_fits_evaluated_df = evaluate_in_dip(top_hat_fit_df_slim_good, top_hat, band)
# around the dip
top_hat_fits_evaluated_df = evaluate_around_dip(top_hat_fits_evaluated_df, top_hat, band, wiggle)
# over all data
top_hat_fits_evaluated_df = evaluate(top_hat_fits_evaluated_df, top_hat, band)
```

```python
catalog.save_axs_table(top_hat_fits_evaluated_df, "6_4_20_stevengs_top_hat_fits_r_band")
```

```python
top_hat_lcs = catalog.load("6_4_20_stevengs_top_hat_fits_r_band")
top_hat_lc = top_hat_lcs.head(1)[0]
```

```python
plot_fit_result(
    top_hat_lc['window_r']['x'], top_hat_lc['window_r']['y'], top_hat_lc['window_r']['yerr'],
    top_hat_lc['fit_r'], top_hat
)
print(top_hat_lc['model_error_in_dip_r'])
print(top_hat_lc['model_error_around_dip_r'])
print(top_hat_lc['model_error_r'])
```

```python
top_hat_lc_min_chi_square = top_hat_lcs.sort(
    top_hat_lcs['model_error_around_dip_r.reduced_sum_square_error'], ascending=True
).head(1)[0]
```

```python
plot_fit_result(
    top_hat_lc_min_chi_square['around_dip_window_r']['x'],
    top_hat_lc_min_chi_square['around_dip_window_r']['y'],
    top_hat_lc_min_chi_square['around_dip_window_r']['yerr'],
    top_hat_lc_min_chi_square['fit_r'], top_hat
)
print(top_hat_lc_min_chi_square['model_error_in_dip_r'])
print(top_hat_lc_min_chi_square['model_error_around_dip_r'])
print(top_hat_lc_min_chi_square['model_error_r'])
```


# **CS224W - Colab 4**

In Colab 2 we constructed GNN models by using PyTorch Geometric's built in GCN layer, `GCNConv`. In Colab 3 we implemented the  **GraphSAGE** ([Hamilton et al. (2017)](https://arxiv.org/abs/1706.02216)) layer. In this colab you'll use what you've learned and implement a more powerful layer: **GAT** ([Veličković et al. (2018)](https://arxiv.org/abs/1710.10903)). Then we will run our models on the CORA dataset, which is a standard citation network benchmark dataset.

**Note**: Make sure to **sequentially run all the cells in each section** so that the intermediate variables / packages will carry over to the next cell

Have fun and good luck on Colab 4 :)

# Device
We recommend using a GPU for this Colab.

Please click `Runtime` and then `Change runtime type`. Then set the `hardware accelerator` to **GPU**.

## Installation

```python
# Install torch geometric
import os
if 'IS_GRADESCOPE_ENV' not in os.environ:
  !pip install torch-scatter -f https://data.pyg.org/whl/torch-1.10.0+cu113.html
  !pip install torch-sparse -f https://data.pyg.org/whl/torch-1.10.0+cu113.html
  !pip install torch-geometric
  !pip install -q git+https://github.com/snap-stanford/deepsnap.git
```

    Looking in links: https://data.pyg.org/whl/torch-1.10.0+cu113.html
    Requirement already satisfied: torch-scatter in /usr/local/lib/python3.7/dist-packages (2.0.9)
    Looking in links: https://data.pyg.org/whl/torch-1.10.0+cu113.html
    Requirement already satisfied: torch-sparse in /usr/local/lib/python3.7/dist-packages (0.6.12)
    Requirement already satisfied: scipy in /usr/local/lib/python3.7/dist-packages (from torch-sparse) (1.4.1)
    Requirement already satisfied: numpy>=1.13.3 in /usr/local/lib/python3.7/dist-packages (from scipy->torch-sparse) (1.19.5)
    Requirement already satisfied: torch-geometric in /usr/local/lib/python3.7/dist-packages (2.0.3)
    Requirement already satisfied: networkx in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (2.6.3)
    Requirement already satisfied: PyYAML in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (3.13)
    Requirement already satisfied: rdflib in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (6.1.1)
    Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (1.19.5)
    Requirement already satisfied: pandas in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (1.1.5)
    Requirement already satisfied: pyparsing in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (3.0.6)
    Requirement already satisfied: yacs in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (0.1.8)
    Requirement already satisfied: scipy in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (1.4.1)
    Requirement already satisfied: tqdm in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (4.62.3)
    Requirement already satisfied: jinja2 in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (2.11.3)
    Requirement already satisfied: scikit-learn in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (1.0.2)
    Requirement already satisfied: googledrivedownloader in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (0.4)
    Requirement already satisfied: requests in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (2.23.0)
    Requirement already satisfied: MarkupSafe>=0.23 in /usr/local/lib/python3.7/dist-packages (from jinja2->torch-geometric) (2.0.1)
    Requirement already satisfied: pytz>=2017.2 in /usr/local/lib/python3.7/dist-packages (from pandas->torch-geometric) (2018.9)
    Requirement already satisfied: python-dateutil>=2.7.3 in /usr/local/lib/python3.7/dist-packages (from pandas->torch-geometric) (2.8.2)
    Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.7/dist-packages (from python-dateutil>=2.7.3->pandas->torch-geometric) (1.15.0)
    Requirement already satisfied: importlib-metadata in /usr/local/lib/python3.7/dist-packages (from rdflib->torch-geometric) (4.10.0)
    Requirement already satisfied: setuptools in /usr/local/lib/python3.7/dist-packages (from rdflib->torch-geometric) (57.4.0)
    Requirement already satisfied: isodate in /usr/local/lib/python3.7/dist-packages (from rdflib->torch-geometric) (0.6.1)
    Requirement already satisfied: typing-extensions>=3.6.4 in /usr/local/lib/python3.7/dist-packages (from importlib-metadata->rdflib->torch-geometric) (3.10.0.2)
    Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.7/dist-packages (from importlib-metadata->rdflib->torch-geometric) (3.7.0)
    Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests->torch-geometric) (2021.10.8)
    Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests->torch-geometric) (1.24.3)
    Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests->torch-geometric) (3.0.4)
    Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests->torch-geometric) (2.10)
    Requirement already satisfied: threadpoolctl>=2.0.0 in /usr/local/lib/python3.7/dist-packages (from scikit-learn->torch-geometric) (3.0.0)
    Requirement already satisfied: joblib>=0.11 in /usr/local/lib/python3.7/dist-packages (from scikit-learn->torch-geometric) (1.1.0)

```python
import torch_geometric
torch_geometric.__version__
```

    '2.0.3'

# 1) GNN Layers

## Implementing Layer Modules

In Colab 2, we implemented a GCN model for node and graph classification tasks. However, for that notebook we took advantage of PyG's built in GCN module. For Colabs 3 and 4, we provide a build upon a general Graph Neural Network Stack, into which we will be able to plugin our own module implementations: GraphSAGE and GAT.

We will then use our layer implemenations to complete node classification on the CORA dataset, a standard citation network benchmark. In this dataset, nodes correspond to documents and edges correspond to undirected citations. Each node or document in the graph is assigned a class label and features based on the documents binarized bag-of-words representation. Specifically, the Cora graph has 2708 nodes, 5429 edges, 7 prediction classes, and 1433 features per node.

## GNN Stack Module

Below is the implementation of a general GNN stack, where we can plugin any GNN layer, such as **GraphSage**, **GAT**, etc. This module is provided for you. Your implementations of the **GraphSage** and **GAT** layers will function as components in the GNNStack Module.

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

import torch_geometric.nn as pyg_nn
import torch_geometric.utils as pyg_utils

from torch import Tensor
from typing import Union, Tuple, Optional
from torch_geometric.typing import (OptPairTensor, Adj, Size, NoneType,
                                    OptTensor)

from torch.nn import Parameter, Linear
from torch_sparse import SparseTensor, set_diag
from torch_geometric.nn.conv import MessagePassing
from torch_geometric.utils import remove_self_loops, add_self_loops, softmax

class GNNStack(torch.nn.Module):

    def __init__(
        self,
        input_dim,
        hidden_dim,
        output_dim,
        args,
        emb=False
        ):

        # Parentizing
        super(GNNStack, self).__init__()

        # What type of GNN we are talking about
        conv_model = self.build_conv_model(args.model_type)

        # List of conv layers
        self.convs = nn.ModuleList()

        # Append input layer
        self.convs.append(conv_model(input_dim, hidden_dim))

        # Check if the number of layers is more than 1
        assert (args.num_layers >= 1), 'Number of layers is not >=1'

        # Go through all layers
        for l in range(args.num_layers-1):
            self.convs.append(conv_model(args.heads * hidden_dim, hidden_dim))

        # Post-message-passing
        self.post_mp = nn.Sequential(
            nn.Linear(args.heads * hidden_dim, hidden_dim),
            nn.Dropout(args.dropout),
            nn.Linear(hidden_dim, output_dim)
            )

        # Define dropout
        self.dropout = args.dropout

        # Store number of layers for future use into the methods
        self.num_layers = args.num_layers

        # Storing boolean about embeddings (True dont apply a classification
        # layer in the end)
        self.emb = emb

    def build_conv_model(self, model_type):

        if model_type == 'GraphSage':
            return GraphSage

        elif model_type == 'GAT':
            # When applying GAT with num heads > 1, you need to modify the
            # input and output dimension of the conv layers (self.convs),
            # to ensure that the input dim of the next layer is num heads
            # multiplied by the output dim of the previous layer.
            # HINT: In case you want to play with multiheads, you need to change the for-loop that builds up self.convs to be
            # self.convs.append(conv_model(hidden_dim * num_heads, hidden_dim)),
            # and also the first nn.Linear(hidden_dim * num_heads, hidden_dim) in post-message-passing.
            return GAT

    def forward(self, data):

        # Extract from batch
        x, edge_index, batch = data.x, data.edge_index, data.batch

        # Go through all layers
        for i in range(self.num_layers):
            x = self.convs[i](x, edge_index)
            x = F.relu(x)
            x = F.dropout(x, p=self.dropout, training=self.training)

        x = self.post_mp(x)

        if self.emb == True:
            return x

        return F.log_softmax(x, dim=1)

    def loss(self, pred, label):
        return F.nll_loss(pred, label)
```

## Creating Our Own Message Passing Layer

Now let's start implementing our own message passing layers! Working through this part will help us become acutely familiar with the behind the scenes work of implementing Pytorch Message Passing Layers, allowing us to build our own GNN models. To do so, we will work with and implement 3 critcal functions needed to define a PyG Message Passing Layer: `forward`, `message`, and `aggregate`.

Before diving head first into the coding details, let us quickly review the key components of the message passing process. To do so, we will focus on a single round of messsage passing with respect to a single central node $x$. Before message passing, $x$ is associated with a feature vector $x^{l-1}$, and the goal of message passing is to update this feature vector as $x^l$. To do so, we implement the following steps: 1) each neighboring node $v$ passes its current message $v^{l-1}$ across the edge $(x, v)$ - 2) for the node $x$, we aggregate all of the messages of the neighboring nodes (for example through a sum or mean) - and 3) we transform the aggregated information by for example applying linear and non-linear transformations. Altogether, the message passing process is applied such that every node $u$ in our graph updates its embedding by acting as the central node $x$ in step 1-3 described above.

Now, we extending this process to that of a single message passing layer, the job of a message passing layer is to update the current feature representation or embedding of each node in a graph by propagating and transforming information within the graph. Overall, the general paradigm of a message passing layers is: 1) pre-processing -> 2) **message passing** / propagation -> 3) post-processing.

The `forward` fuction that we will implement for our message passing layer captures this execution logic. Namely, the `forward` function handles the pre and post-processing of node features / embeddings, as well as initiates message passing by calling the `propagate` function.

The `propagate` function encapsulates the message passing process! It does so by calling three important functions: 1) `message`, 2) `aggregate`, and 3) `update`. Our implementation will vary slightly from this, as we will not explicitly implement `update`, but instead place the logic for updating node embeddings after message passing and within the `forward` function. To be more specific, after information is propagated (message passing), we can further transform the node embeddings outputed by `propagate`. Therefore, the output of `forward` is exactly the node embeddings after one GNN layer.

Lastly, before starting to implement our own layer, let us dig a bit deeper into each of the functions described above:

1.

```
def propagate(edge_index, x=(x_i, x_j), extra=(extra_i, extra_j), size=size):
```
Calling `propagate` initiates the message passing process. Looking at the function parameters, we highlight a couple of key parameters.

  - `edge_index` is passed to the forward function and captures the edge structure of the graph.
  - `x=(x_i, x_j)` represents the node features that will be used in message passing. In order to explain why we pass the tuple `(x_i, x_j)`, we first look at how our edges are represented. For every edge $(i, j) \in \mathcal{E}$, we can differentiate $i$ as the source or central node ($x_{central}$) and j as the neighboring node ($x_{neighbor}$).

    Taking the example of message passing above, for a central node $u$ we will aggregate and transform all of the messages associated with the nodes $v$ s.t. $(u, v) \in \mathcal{E}$ (i.e. $v \in \mathcal{N}_{u}$). Thus we see, the subscripts `_i` and `_j` allow us to specifcally differenciate features associated with central nodes (i.e. nodes  recieving message information) and neighboring nodes (i.e. nodes passing messages).

    This is definitely a somewhat confusing concept; however, one key thing to remember / wrap your head around is that depending on the perspective, a node $x$ acts as a central node or a neighboring node. In fact, in undirected graphs we store both edge directions (i.e. $(i, j)$ and $(j, i)$). From the central node perspective, `x_i`, x is collecting neighboring information to update its embedding. From a neighboring node perspective, `x_j`, x is passing its message information along the edge connecting it to a different central node.

  - `extra=(extra_i, extra_j)` represents additional information that we can associate with each node beyond its current feature embedding. In fact, we can include as many additional parameters of the form `param=(param_i, param_j)` as we would like. Again, we highlight that indexing with `_i` and `_j` allows us to differentiate central and neighboring nodes.

  The output of the `propagate` function is a matrix of node embeddings after the message passing process and has shape $[N, d]$.

2.
```
def message(x_j, ...):
```
The `message` function is called by propagate and constructs the messages from
neighboring nodes $j$ to central nodes $i$ for each edge $(i, j)$ in *edge_index*. This function can take any argument that was initially passed to `propagate`. Furthermore, we can again differentiate central nodes and neighboring nodes by appending `_i` or `_j` to the variable name, .e.g. `x_i` and `x_j`. Looking more specifically at the variables, we have:

  - `x_j` represents a matrix of feature embeddings for all neighboring nodes passing their messages along their respective edge (i.e. all nodes $j$ for edges $(i, j) \in \mathcal{E}$). Thus, its shape is $[|\mathcal{E}|, d]$!
  - In implementing GAT we will see how to access additional variables passed to propagate

  Critically, we see that the output of the `message` function is a matrix of neighboring node embeddings ready to be aggregated, having shape $[|\mathcal{E}|, d]$.

3.
```
def aggregate(self, inputs, index, dim_size = None):
```
Lastly, the `aggregate` function is used to aggregate the messages from neighboring nodes. Looking at the parameters we highlight:

  - `inputs` represents a matrix of the messages passed from neighboring nodes (i.e. the output of the `message` function).
  - `index` has the same shape as `inputs` and tells us the central node that corresponding to each of the rows / messages $j$ in the `inputs` matrix. Thus, `index` tells us which rows / messages to aggregate for each central node.

  The output of `aggregate` is of shape $[N, d]$.

For additional resources refer to the PyG documentation for implementing custom message passing layers: https://pytorch-geometric.readthedocs.io/en/latest/notes/create_gnn.html

## GAT Implementation

Attention mechanisms have become the state-of-the-art in many sequence-based tasks such as machine translation and learning sentence representations. One of the major benefits of attention-based mechanisms is their ability to focus on the most relevant parts of the input to make decisions. In this problem, we will see how attention mechanisms can be used to perform node classification over graph-structured data through the usage of Graph Attention Networks (GATs) ([Veličković et al. (2018)](https://arxiv.org/abs/1710.10903)).

The building block of the Graph Attention Network is the graph attention layer, which is a variant of the aggregation function. Let $N$ be the number of nodes and $F$ be the dimension of the feature vector for each node. The input to each graph attentional layer is a set of node features: $\mathbf{h} = \{\overrightarrow{h_1}, \overrightarrow{h_2}, \dots, \overrightarrow{h_N}$\}, $\overrightarrow{h_i} \in R^F$. The output of each graph attentional layer is a new set of node features, which may have a new dimension $F'$: $\mathbf{h'} = \{\overrightarrow{h_1'}, \overrightarrow{h_2'}, \dots, \overrightarrow{h_N'}\}$, with $\overrightarrow{h_i'} \in \mathbb{R}^{F'}$.

We will now describe how this transformation is performed for each graph attention layer. First, a shared linear transformation parametrized by the weight matrix $\mathbf{W} \in \mathbb{R}^{F' \times F}$ is applied to every node.

Next, we perform self-attention on the nodes. We use a shared attention function $a$:
\begin{equation}
a : \mathbb{R}^{F'} \times \mathbb{R}^{F'} \rightarrow \mathbb{R}.
\end{equation}

that computes the attention coefficients capturing the importance of node $j$'s features to node $i$:
\begin{equation}
e_{ij} = a(\mathbf{W_l}\overrightarrow{h_i}, \mathbf{W_r} \overrightarrow{h_j})
\end{equation}

The most general formulation of self-attention allows every node to attend to all other nodes which drops all structural information. However, to utilize graph structure in the attention mechanisms, we use **masked attention**. In masked attention, we only compute attention coefficients $e_{ij}$ for nodes $j \in \mathcal{N}_i$ where $\mathcal{N}_i$ is some neighborhood of node $i$ in the graph.

To easily compare coefficients across different nodes, we normalize the coefficients across $j$ using a softmax function:
\begin{equation}
\alpha_{ij} = \text{softmax}_j(e_{ij}) = \frac{\exp(e_{ij})}{\sum_{k \in \mathcal{N}_i} \exp(e_{ik})}
\end{equation}

For this problem, our attention mechanism $a$ will be a single-layer feedforward neural network parametrized by a weight vectors $\overrightarrow{a_l} \in \mathbb{R}^{F'}$ and $\overrightarrow{a_r} \in \mathbb{R}^{F'}$, followed by a LeakyReLU nonlinearity (with negative input slope 0.2). Let $\cdot^T$ represent transposition and $||$ represent concatenation. The coefficients computed by our attention mechanism may be expressed as:

\begin{equation}
\alpha_{ij} = \frac{\exp\Big(\text{LeakyReLU}\Big(\overrightarrow{a_l}^T \mathbf{W_l} \overrightarrow{h_i} + \overrightarrow{a_r}^T\mathbf{W_r}\overrightarrow{h_j}\Big)\Big)}{\sum_{k\in \mathcal{N}_i} \exp\Big(\text{LeakyReLU}\Big(\overrightarrow{a_l}^T \mathbf{W_l} \overrightarrow{h_i} + \overrightarrow{a_r}^T\mathbf{W_r}\overrightarrow{h_k}\Big)\Big)}
\end{equation}

For the following questions, we denote `alpha_l` = $\alpha_l = [...,\overrightarrow{a_l}^T \mathbf{W_l} \overrightarrow{h_i},...] \in \mathcal{R}^n$ and `alpha_r` = $\alpha_r = [..., \overrightarrow{a_r}^T \mathbf{W_r} \overrightarrow{h_j}, ...] \in \mathcal{R}^n$.

At every layer of GAT, after the attention coefficients are computed for that layer, the aggregation function can be computed by a weighted sum of neighborhood messages, where weights are specified by $\alpha_{ij}$.

Now, we use the normalized attention coefficients to compute a linear combination of the features corresponding to them. These aggregated features will serve as the final output features for every node.

\begin{equation}
h_i' = \sum_{j \in \mathcal{N}_i} \alpha_{ij} \mathbf{W_r} \overrightarrow{h_j}.
\end{equation}

At this point, we have covered a lot of information! Before reading further about multi-head attention, we encourage you to go again through the excersize of thinking about what components of the attention mechanism correspond with the different functions: 1) `forward`, 2) `message`, and 3 `aggregate`.

- Hint 1: Our aggregation is very similar to that of GraphSage except now we are using sum aggregation
- Hint 2: The terms we aggregate over again represent the individual message that each neighbor node j sends. Thus, we see that $\alpha_{ij}$ is part of the message each node sends and is thus computed during the message step. This makes sense since an attention weight is associated with each edge in the graph.
- Hint 3: Look at the terms in the definition of $\alpha_{ij}$. What values do we want to pre-process and pass as parameters to the `propagate` function. The parameters of `message(..., x_j, alpha_j, alpha_i, ...)` should give a good hint.

### Multi-Head Attention
To stabilize the learning process of self-attention, we use multi-head attention. To do this we use $K$ independent attention mechanisms, or ``heads'' compute output features as in the above equations. Then, we concatenate these output feature representations:

\begin{equation}
    \overrightarrow{h_i}' = ||_{k=1}^K \Big(\sum_{j \in \mathcal{N}_i} \alpha_{ij}^{(k)} \mathbf{W_r}^{(k)} \overrightarrow{h_j}\Big)
\end{equation}

where $||$ is concentation, $\alpha_{ij}^{(k)}$ are the normalized attention coefficients computed by the $k$-th attention mechanism $(a^k)$, and $\mathbf{W}^{(k)}$ is the corresponding input linear transformation's weight matrix. Note that for this setting, $\mathbf{h'} \in \mathbb{R}^{KF'}$.

```python
class GAT(MessagePassing):

    def __init__(
        self,
        in_channels,
        out_channels,
        heads = 2,
        negative_slope = 0.2,
        dropout = 0.,
        **kwargs):

        super(GAT, self).__init__(node_dim=0, **kwargs)

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.heads = heads
        self.negative_slope = negative_slope
        self.dropout = dropout

        self.lin_l = None
        self.lin_r = None
        self.att_l = None
        self.att_r = None

        ############################################################################
        # TODO: Your code here!
        # Define the layers needed for the message functions below.
        # self.lin_l is the linear transformation that you apply to embeddings
        # BEFORE message passing.
        #
        # Pay attention to dimensions of the linear layers, since we're using
        # multi-head attention.
        # Our implementation is ~1 lines, but don't worry if you deviate from this.
        self.lin_l = Linear(in_channels, heads * out_channels)

        ############################################################################

        self.lin_r = self.lin_l

        ############################################################################
        # TODO: Your code here!
        # Define the attention parameters \overrightarrow{a_l/r}^T in the above intro.
        # You have to deal with multi-head scenarios.
        # Use nn.Parameter instead of nn.Linear
        # Our implementation is ~2 lines, but don't worry if you deviate from this.

        # The idea here is to initialize all the parameters involving the attention
        # mechanism, which is described by the vector of weight a of dimension
        # out_channels.

        # However, if we stack several heads together, we need such a vector
        # for each one of these heads
        #
        # Therefore, following the documentation: https://pytorch.org/docs/1.9.1/generated/torch.nn.parameter.Parameter.html
        self.att_l = Parameter(torch.zeros(heads, out_channels))
        self.att_r = Parameter(torch.zeros(heads, out_channels))

        # PS: I almost used self.attention_l; but in the reset_parameters below
        # the guy is using att_l

        ############################################################################

        self.reset_parameters()

    def reset_parameters(self):
        nn.init.xavier_uniform_(self.lin_l.weight)
        nn.init.xavier_uniform_(self.lin_r.weight)
        nn.init.xavier_uniform_(self.att_l)
        nn.init.xavier_uniform_(self.att_r)

    def forward(self, x, edge_index, size=None):
        """
        Implement message passing, as well as any post-processing (our update rule).

        Parameters
        ----------

        self : GAT object
            GAT.

        x : tensor of shape (num_nodes, in_channels)
            Data associated to each node.

        edge_index : tensor of shape (2, num_edges)
            Describe graph structure via edge list.

        size : tuple with shape
            Needed if edge_index is a matrix representing connections as
            (num_nodes, num_connected_nodes).

        See: https://github.com/pyg-team/pytorch_geometric/blob/master/torch_geometric/nn/conv/message_passing.py

        Returns
        -------

        out : tensor of shape (num_nodes, out_channels)
            Embedded data associated to each node.
        """
        # Now what was called K is H, and F' is C... this guy is so confused
        H, C = self.heads, self.out_channels

        ############################################################################
        # TODO: Your code here!
        # Implement message passing, as well as any pre- and post-processing (our update rule).
        # 1. First apply linear transformation to node embeddings, and split that
        #    into multiple heads. We use the same representations for source and
        #    target nodes, but apply different linear weights (W_l and W_r)
        # 2. Calculate alpha vectors for central nodes (alpha_l) and neighbor nodes (alpha_r).
        # 3. Call propagate function to conduct the message passing.
        #    3.1 Remember to pass alpha = (alpha_l, alpha_r) as a parameter.
        #    3.2 See there for more information: https://pytorch-geometric.readthedocs.io/en/latest/notes/create_gnn.html
        # 4. Transform the output back to the shape of [N, H * C].
        # Our implementation is ~5 lines, but don't worry if you deviate from this.

        # 1. First apply linear transformation to node embeddings, and split that
        #    into multiple heads. We use the same representations for source and
        #    target nodes, but apply different linear weights (W_l and W_r)

        # Apply linear transformations
        x_i = self.lin_l(x)     # (num_nodes, heads * out_channels)
        x_j = self.lin_r(x)     # (num_nodes, heads * out_channels)

        # Split into multiple heads
        # I assume that split here means that we want to have a tensor of the
        # form
        #                 (num_nodes, heads, out_channels)
        # Then, we can treat heads independently, as if we are doing things in
        # parallel.
        #
        # Fortunately, tensor also has a .reshape method
        x_i = x_i.reshape(-1, H, C)     # (num_nodes, heads, out_channels)
        x_j = x_j.reshape(-1, H, C)     # (num_nodes, heads, out_channels)

        # 2. Calculate alpha vectors for central nodes (alpha_l) and neighbor nodes (alpha_r).
        # The guy defined it as
        #   `alpha_l` = $\alpha_l = [...,\overrightarrow{a_l}^T \mathbf{W_l} \overrightarrow{h_i},...]
        # We already have the W h part... now, we neet to apply the attention coeff
        alpha_l = self.att_l * x_i  # (num_nodes, heads, out_channels)
        alpha_r = self.att_r * x_j  # (num_nodes, heads, out_channels)

        # 3. Call propagate function to conduct the message passing.
        #    3.1 Remember to pass alpha = (alpha_l, alpha_r) as a parameter.
        #    3.2 See there for more information: https://pytorch-geometric.readthedocs.io/en/latest/notes/create_gnn.html
        out = self.propagate(edge_index, x=(x_i, x_j), alpha=(alpha_l, alpha_r), size=size)     # (num_nodes, heads, out_channels)

        # 4. Transform the output back to the shape of [N, H * C].
        out = out.reshape(-1, H * C)    # (num_edges, heads * out_channels)

        ############################################################################

        return out

    def message(self, x_j, alpha_j, alpha_i, index, ptr, size_i):
        """
        Messages from the neighbors.

        For GAT, the neighbors send their weighted raw data, where the weights
        are given by the attention mechanism.

        Parameters
        ----------

        x_j : tensor of shape (num_edges, out_channels)
            Data of the target node of each edge.

        alpha_j : tensor of shape (num_edges, heads, out_channels)
            Attention parameters for the target nodes.

        alpha_i : tensor of shape (num_edges, heads, out_channels)
            Attention parameters for the central nodes.

        index : tensor of shape (num_edges)
            List o indexes of the (target) nodes receiving the message:
                        index = edge_list[1],
            since we are using the flow="source-to-target".

        Returns
        -------

        out : tensor of shape (num_edges, head, out_channels)
            Messages sent to the central node from the target node of each edge.

        PS: Note how different is this message function from the one of Colab 3.

        """
        out = None

        ############################################################################
        # TODO: Your code here!
        # Implement your message function. Putting the attention in message
        # instead of in update is a little tricky.
        # 1. Calculate the final attention weights using alpha_i and alpha_j,
        #    and apply leaky Relu.
        # 2. Calculate softmax over the neighbor nodes for all the nodes. Use
        #    torch_geometric.utils.softmax instead of the one in Pytorch.
        # 3. Apply dropout to attention weights (alpha).
        # 4. Multiply embeddings and attention weights. As a sanity check, the output
        #    should be of shape [E, H, C].
        # 5. ptr (LongTensor, optional): If given, computes the softmax based on
        #    sorted inputs in CSR representation. You can simply pass it to softmax.
        # Our implementation is ~4-5 lines, but don't worry if you deviate from this.

        # 1. Calculate the final attention weights using alpha_i and alpha_j,
        #    and apply leaky Relu.
        #
        #   Documentation on Leak ReLU: https://pytorch.org/docs/1.9.1/generated/torch.nn.LeakyReLU.html
        e_vector = F.leaky_relu(alpha_j + alpha_i, negative_slope=self.negative_slope)  # (num_edges, heads, out_channels)

        # Following the notation of the guy, e_vector has the e_ij values before
        # normalization

        # 2. Calculate softmax over the neighbor nodes for all the nodes. Use
        #    torch_geometric.utils.softmax instead of the one in Pytorch.
        if ptr:
            alpha = F.softmax(e_vector, dim=ptr)    # (num_edges, heads, out_channels)
            # ptr figure outs how to compute the neighborhoods and so on
            # see: https://github.com/pyg-team/pytorch_geometric/blob/50b7bfc4a59b5b6f7ec547ff862985f3b2e22798/torch_geometric/nn/conv/message_passing.py#L336
        else:
            alpha = pyg_utils.softmax(e_vector, index=index)    # (num_edges, heads, out_channels)
            # THIS FUNCTION IS ALREADY FIGURING OUT THE NEIGHBOR SETS FOR US TO COMPUTE
            # THE SOFTMAX FUNCTION, see: https://pytorch-geometric.readthedocs.io/en/1.3.2/_modules/torch_geometric/utils/softmax.html
            # Unsurprisingly, it uses the scatter function.

        # 3. Apply dropout to attention weights (alpha).
        alpha = F.dropout(alpha, p=self.dropout)    # (num_edges, heads, out_channels)

        # 4. Multiply embeddings and attention weights. As a sanity check, the output
        #    should be of shape [E, H, C].
        out = alpha * x_j   # (num_edges, heads, out_channels)

        # This is denoted as h' in the explanation above, but without the sum
        # yet (aggreation); this will be realized in the aggreagation function

        # So, here we have this alpha Wr h

        # 5. ptr (LongTensor, optional): If given, computes the softmax based on
        #    sorted inputs in CSR representation. You can simply pass it to softmax.
        #
        # - OK, why he didnt say this before.. lets go back to the softmax
        # implementation

        ############################################################################

        return out

    def aggregate(self, inputs, index, dim_size=None):
        """
        Aggregates at the central node incoming messages from the neighbors.

        For GAT, the aggregation is a weighted sum over the data coming from
        the neighbors. The weights are defined by the attention mechanism.

        Parameters
        ----------

        inputs : tensor of shape (E, heads, out_channels)
            Messages sent to the central node from the target node of each edge.

        index : tensor of shape (num_edges)
            List o indexes of the (target) nodes receiving the message:
                        index = edge_list[1],
            since we are using the flow="source-to-target".

        dim_size : tuple with shape
            Tuple with shape of the aggregation.

        Returns
        -------

        out : tensor of shape (num_edges, heads, out_channels)
            Aggregated messages at the central nodes.
        """

        ############################################################################
        # TODO: Your code here!
        # Implement your aggregate function here.
        # See here as how to use torch_scatter.scatter: https://pytorch-scatter.readthedocs.io/en/latest/_modules/torch_scatter/scatter.html
        # Pay attention to "reduce" parameter is different from that in GraphSage.
        # Our implementation is ~1 lines, but don't worry if you deviate from this.

        # Now, we just need to sum over the neighboorhood sets, see the h' equation
        # above!!
        out = torch_scatter.scatter(src=inputs, index=index, dim=self.node_dim, dim_size=dim_size, reduce="sum")

        # dim needs to match the node dimension!! Using the node_dim
        # property that the MensagePassing class has; similar to Colab 3
        # see: https://github.com/pyg-team/pytorch_geometric/blob/50b7bfc4a59b5b6f7ec547ff862985f3b2e22798/torch_geometric/nn/conv/message_passing.py#L336

        ############################################################################

        return out
```

## Building Optimizers

This function has been implemented for you. **For grading purposes please use the default Adam optimizer**, but feel free to play with other types of optimizers on your own.

```python
import torch.optim as optim

def build_optimizer(args, params):
    weight_decay = args.weight_decay
    filter_fn = filter(lambda p : p.requires_grad, params)
    if args.opt == 'adam':
        optimizer = optim.Adam(filter_fn, lr=args.lr, weight_decay=weight_decay)
    elif args.opt == 'sgd':
        optimizer = optim.SGD(filter_fn, lr=args.lr, momentum=0.95, weight_decay=weight_decay)
    elif args.opt == 'rmsprop':
        optimizer = optim.RMSprop(filter_fn, lr=args.lr, weight_decay=weight_decay)
    elif args.opt == 'adagrad':
        optimizer = optim.Adagrad(filter_fn, lr=args.lr, weight_decay=weight_decay)
    if args.opt_scheduler == 'none':
        return None, optimizer
    elif args.opt_scheduler == 'step':
        scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=args.opt_decay_step, gamma=args.opt_decay_rate)
    elif args.opt_scheduler == 'cos':
        scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.opt_restart)
    return scheduler, optimizer
```

## Training and Testing

Here we provide you with the functions to train and test. **Please do not modify this part for grading purposes.**

```python
import time

import networkx as nx
import numpy as np
import torch
import torch.optim as optim
from tqdm import trange
import pandas as pd
import copy

from torch_geometric.datasets import TUDataset
from torch_geometric.datasets import Planetoid
from torch_geometric.data import DataLoader

import torch_geometric.nn as pyg_nn

import matplotlib.pyplot as plt

def train(dataset, args):

    print("Node task. test set size:", np.sum(dataset[0]['test_mask'].numpy()))
    print()
    test_loader = loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False)

    # build model
    model = GNNStack(dataset.num_node_features, args.hidden_dim, dataset.num_classes,
                            args)
    scheduler, opt = build_optimizer(args, model.parameters())

    # train
    losses = []
    test_accs = []
    best_acc = 0
    best_model = None
    for epoch in trange(args.epochs, desc="Training", unit="Epochs"):
        total_loss = 0
        model.train()
        for batch in loader:
            opt.zero_grad()
            pred = model(batch)
            label = batch.y
            pred = pred[batch.train_mask]
            label = label[batch.train_mask]
            loss = model.loss(pred, label)
            loss.backward()
            opt.step()
            total_loss += loss.item() * batch.num_graphs
        total_loss /= len(loader.dataset)
        losses.append(total_loss)

        if epoch % 10 == 0:
          test_acc = test(test_loader, model)
          test_accs.append(test_acc)
          if test_acc > best_acc:
            best_acc = test_acc
            best_model = copy.deepcopy(model)
        else:
          test_accs.append(test_accs[-1])

    return test_accs, losses, best_model, best_acc, test_loader

def test(loader, test_model, is_validation=False, save_model_preds=False, model_type=None):
    test_model.eval()

    correct = 0
    # Note that Cora is only one graph!
    for data in loader:
        with torch.no_grad():
            # max(dim=1) returns values, indices tuple; only need indices
            pred = test_model(data).max(dim=1)[1]
            label = data.y

        mask = data.val_mask if is_validation else data.test_mask
        # node classification: only evaluate on nodes in test set
        pred = pred[mask]
        label = label[mask]

        if save_model_preds:
          print ("Saving Model Predictions for Model Type", model_type)

          data = {}
          data['pred'] = pred.view(-1).cpu().detach().numpy()
          data['label'] = label.view(-1).cpu().detach().numpy()

          df = pd.DataFrame(data=data)
          # Save locally as csv
          df.to_csv('CORA-Node-' + model_type + '.csv', sep=',', index=False)

        correct += pred.eq(label).sum().item()

    total = 0
    for data in loader.dataset:
        total += torch.sum(data.val_mask if is_validation else data.test_mask).item()

    return correct / total

class objectview(object):
    def __init__(self, d):
        self.__dict__ = d

```

## Let's Start the Training!

We will be working on the CORA dataset on node-level classification.

This part is implemented for you. **For grading purposes, please do not modify the default parameters.** However, feel free to play with different configurations just for fun!

**Submit your best accuracy and loss on Gradescope.**

```python
if 'IS_GRADESCOPE_ENV' not in os.environ:
    for args in [{
        'model_type': 'GAT',
        'dataset': 'cora',
        'num_layers': 2,
        'heads': 1,
        'batch_size': 32,
        'hidden_dim': 32,
        'dropout': 0.5,
        'epochs': 500,
        'opt': 'adam',
        'opt_scheduler': 'none',
        'opt_restart': 0,
        'weight_decay': 5e-3,
        'lr': 0.01},
    ]:
        args = objectview(args)
        for model in ['GAT']:
            args.model_type = model

            # Match the dimension.
            if model == 'GAT':
              args.heads = 2
            else:
              args.heads = 1

            if args.dataset == 'cora':
                dataset = Planetoid(root='/tmp/cora', name='Cora')
            else:
                raise NotImplementedError("Unknown dataset")
            test_accs, losses, best_model, best_acc, test_loader = train(dataset, args)

            print("Maximum test set accuracy: {0}".format(max(test_accs)))
            print("Minimum loss: {0}".format(min(losses)))

            # Run test for our best model to save the predictions!
            test(test_loader, best_model, is_validation=False, save_model_preds=True, model_type=model)
            print()

            plt.title(dataset.name)
            plt.plot(losses, label="training loss" + " - " + args.model_type)
            plt.plot(test_accs, label="test accuracy" + " - " + args.model_type)
        plt.legend()
        plt.show()

```

## Question 1: What is the maximum accuracy obtained on test set for GAT? (10 points)

Running the training cell above will also save your best GAT model predictions as *CORA-Node-GAT.csv*.

When you sumbit your assignment, you will have to download this file and attatch it to your submission. As with the other colabs, please zip this file (DON'T CHANGE ITS NAME) and the .csv file that's generated!


```bash
%%bash -e
if ! [[ -f ./linkern ]]; then
  wget http://www.math.uwaterloo.ca/tsp/concorde/downloads/codes/src/co031219.tgz
  echo 'c3650a59c8d57e0a00e81c1288b994a99c5aa03e5d96a314834c2d8f9505c724  co031219.tgz' | sha256sum -c
  tar xf co031219.tgz
  (cd concorde && CFLAGS='-O3 -march=native -mtune=native -fPIC' ./configure --build=x86_64)
  (cd concorde/LINKERN && make -j && cp linkern ../../)
  rm -rf concorde co031219.tgz
fi
```

```python
from concorde.tsp import TSPSolver
from matplotlib import collections  as mc
import numpy as np
import pandas as pd
import time
import pylab as pl
```

```python
cities = pd.read_csv('cities.csv', index_col=['CityId'])
```

```python
cities1k = cities * 1000
```

```python
def write_tsp(cities, filename, name='traveling-santa-2018-prime-paths'):
    with open(filename, 'w') as f:
        f.write('NAME : %s\n' % name)
        f.write('COMMENT : %s\n' % name)
        f.write('TYPE : TSP\n')
        f.write('DIMENSION : %d\n' % len(cities))
        f.write('EDGE_WEIGHT_TYPE : EUC_2D\n')
        f.write('NODE_COORD_SECTION\n')
        for row in cities.itertuples():
            f.write('%d %.11f %.11f\n' % (row.Index+1, row.X, row.Y))
        f.write('EOF\n')

write_tsp(cities1k, 'cities1k.tsp')
```

```bash
%%bash -e
time ./linkern -s 42 -S linkern.tour -R 1000000000 -t 18000 ./cities1k.tsp >linkern.log
```

```python
!sed -Ene 's/([0-9]+) Steps.*Best: ([0-9]+).*/\1,\2/p' linkern.log >linkern.csv
pd.read_csv('linkern.csv', index_col=0, names=['TSP tour length']).plot();
```

```python
import sympy
def read_tour(filename):
    tour = open(filename).read().split()[1:]
    tour = list(map(int, tour))
    if tour[-1] == 0: tour.pop()
    return tour

def score_tour(tour):
    df = cities.reindex(tour + [0]).reset_index()
    primes = list(sympy.primerange(0, len(cities)))
    df['prime'] = df.CityId.isin(primes).astype(int)
    df['dist'] = np.hypot(df.X - df.X.shift(-1), df.Y - df.Y.shift(-1))
    df['penalty'] = df['dist'][9::10] * (1 - df['prime'][9::10]) * 0.1
    return df.dist.sum() + df.penalty.sum()

def write_submission(tour, filename):
    assert set(tour) == set(range(len(tour)))
    pd.DataFrame({'Path': list(tour) + [0]}).to_csv(filename, index=False)
```

```python
tour = read_tour('linkern.tour')
write_submission(tour, 'submission.csv')
```

```python
score_tour(tour)
```

    1517361.2529791819

```python
def flip(lst, k):
    return lst[0:1] + lst[k+1:] + lst[1:k+1]
```

```python
print("Tour path (0-5):",tour[0:5])
tourflip = tour[::-1]
print("Flipped tour path (0-5):", tourflip[0:5])
```

    ('Tour path (0-5):', [78934, 111804, 52086, 89712, 81072])
    ('Flipped tour path (0-5):', [0, 48816, 40230, 75405, 153911])

```python
originScore = score_tour(tour)
print("Score of original tour:", originScore)
tourflip = tour
for n in range(1,101):
    # And the flipped tour looks like:
    tourflip = flip(tourflip, 1)
    # The scores of our tours are:
    flippedScore = score_tour(tourflip)
    print("Score of flipped tour:", flippedScore)

    # If the flipped tour is quicker, change our tour:
    if flippedScore < originScore:
        print("The total improvement was:", abs(flippedScore - originScore))
        originScore = flippedScore
        tour = tourflip
        print("The better of the original/flipped tour is:", tour[0:5])
```

    ('Score of original tour:', 1517332.415616737)
    ('Score of flipped tour:', 1517356.0910889404)
    ('Score of flipped tour:', 1517356.9139951372)
    ('Score of flipped tour:', 1517413.2228083517)
    ('Score of flipped tour:', 1517403.1202388373)
    ('Score of flipped tour:', 1517403.8425916804)
    ('Score of flipped tour:', 1517390.3233206801)
    ('Score of flipped tour:', 1517426.1401942007)
    ('Score of flipped tour:', 1517471.5298787779)
    ('Score of flipped tour:', 1517440.7268941959)
    ('Score of flipped tour:', 1517424.4124984485)
    ('Score of flipped tour:', 1517426.253180478)
    ('Score of flipped tour:', 1517421.8727339443)
    ('Score of flipped tour:', 1517460.3421580733)
    ('Score of flipped tour:', 1517434.6743891994)
    ('Score of flipped tour:', 1517418.130972293)
    ('Score of flipped tour:', 1517393.934870807)
    ('Score of flipped tour:', 1517420.2643460205)
    ('Score of flipped tour:', 1517447.2864386647)
    ('Score of flipped tour:', 1517407.5226760462)
    ('Score of flipped tour:', 1517394.2942589528)
    ('Score of flipped tour:', 1517394.7993746104)
    ('Score of flipped tour:', 1517379.850304974)
    ('Score of flipped tour:', 1517414.5208543059)
    ('Score of flipped tour:', 1517387.1811606274)
    ('Score of flipped tour:', 1517381.9075844737)
    ('Score of flipped tour:', 1517369.7748633618)
    ('Score of flipped tour:', 1517413.2023676371)
    ('Score of flipped tour:', 1517456.4925552462)
    ('Score of flipped tour:', 1517417.9342675023)
    ('Score of flipped tour:', 1517390.2807072725)
    ('Score of flipped tour:', 1517375.7881616235)
    ('Score of flipped tour:', 1517354.4794941714)
    ('Score of flipped tour:', 1517391.7376516568)
    ('Score of flipped tour:', 1517373.4370702687)
    ('Score of flipped tour:', 1517382.8213495801)
    ('Score of flipped tour:', 1517374.1632063612)
    ('Score of flipped tour:', 1517415.894082329)
    ('Score of flipped tour:', 1517467.3961351814)
    ('Score of flipped tour:', 1517445.8441427345)
    ('Score of flipped tour:', 1517435.2475747936)
    ('Score of flipped tour:', 1517439.9246245627)
    ('Score of flipped tour:', 1517447.1559378831)
    ('Score of flipped tour:', 1517499.0819099913)
    ('Score of flipped tour:', 1517493.0927003776)
    ('Score of flipped tour:', 1517502.3384920147)
    ('Score of flipped tour:', 1517485.5076681615)
    ('Score of flipped tour:', 1517502.5371459988)
    ('Score of flipped tour:', 1517526.8949108373)
    ('Score of flipped tour:', 1517477.4368143654)
    ('Score of flipped tour:', 1517438.1595332872)
    ('Score of flipped tour:', 1517428.0130642096)
    ('Score of flipped tour:', 1517409.3595959663)
    ('Score of flipped tour:', 1517431.0013628749)
    ('Score of flipped tour:', 1517393.0333302217)
    ('Score of flipped tour:', 1517372.2586601183)
    ('Score of flipped tour:', 1517349.1398652336)
    ('Score of flipped tour:', 1517387.9553650559)
    ('Score of flipped tour:', 1517442.3627667939)
    ('Score of flipped tour:', 1517424.9399056523)
    ('Score of flipped tour:', 1517423.5400056166)
    ('Score of flipped tour:', 1517439.107014118)
    ('Score of flipped tour:', 1517455.0654943897)
    ('Score of flipped tour:', 1517519.386947802)
    ('Score of flipped tour:', 1517498.676528994)
    ('Score of flipped tour:', 1517498.4435672665)
    ('Score of flipped tour:', 1517492.0689992518)
    ('Score of flipped tour:', 1517525.8929229486)
    ('Score of flipped tour:', 1517580.3247003811)
    ('Score of flipped tour:', 1517556.868937073)
    ('Score of flipped tour:', 1517537.0343341671)
    ('Score of flipped tour:', 1517524.3500662614)
    ('Score of flipped tour:', 1517524.9141456622)
    ('Score of flipped tour:', 1517585.1937728256)
    ('Score of flipped tour:', 1517573.1172919387)
    ('Score of flipped tour:', 1517573.151211919)
    ('Score of flipped tour:', 1517550.4248782436)
    ('Score of flipped tour:', 1517581.9113448008)
    ('Score of flipped tour:', 1517626.8242998021)
    ('Score of flipped tour:', 1517599.5859399824)
    ('Score of flipped tour:', 1517586.2745984038)
    ('Score of flipped tour:', 1517579.5737715021)
    ('Score of flipped tour:', 1517570.7709022472)
    ('Score of flipped tour:', 1517621.9204671208)
    ('Score of flipped tour:', 1517610.5224881524)
    ('Score of flipped tour:', 1517608.9746056185)
    ('Score of flipped tour:', 1517594.5004874978)
    ('Score of flipped tour:', 1517625.8114939115)
    ('Score of flipped tour:', 1517667.3593493842)
    ('Score of flipped tour:', 1517638.9068630573)
    ('Score of flipped tour:', 1517618.783138324)
    ('Score of flipped tour:', 1517603.4920099739)
    ('Score of flipped tour:', 1517576.0349226668)
    ('Score of flipped tour:', 1517627.2630021616)
    ('Score of flipped tour:', 1517631.1658613831)
    ('Score of flipped tour:', 1517641.401856899)
    ('Score of flipped tour:', 1517611.6196710353)
    ('Score of flipped tour:', 1517620.6900758257)
    ('Score of flipped tour:', 1517642.3172794664)
    ('Score of flipped tour:', 1517612.9561278727)
    ('Score of flipped tour:', 1517606.5370181862)

```python
score_tour(tour)
```

    1517332.415616737

```python
write_submission(tour, 'submission.csv')
```

```python

```