# Pricing Uncertainty Induced by Climate Change
by [Michael Barnett](https://sites.google.com/site/michaelduglasbarnett/home), [William Brock](https://www.ssc.wisc.edu/~wbrock/) and [Lars Peter Hansen](https://larspeterhansen.org/).

The latest draft of the paper can be found [here](https://larspeterhansen.org/research/papers/).

Notebook by: Jiaming Wang

## Overview

This notebook provides the source code and explanations for how we solve the model setting with __climate damages to growth__. The computational procedures are similar to the model setting with __climate damages to consumption__. Users should refer to the notebook for the [Consumption Damages Model](ConsumptionModel.ipynb) for more computational details.

Table of Contents

1  Overview

2  Code and Illustration

2.1  Choosing key parameters

2.2  Solving HJB for Growth Damage

2.2.1  Remark:  Accounting for uncertainty about growth damages

2.2.1.1  Selecting the values of $\gamma_1, \gamma_2$

2.2.1.2  Varying damage models

2.2.1.3  Distorted probabilities

2.2.2  Remark: Tolerance level and conergence criteria for HJB

2.3  Simulation

2.4  SCC Calculation Feyman Kac

2.4.1  Remark: Convergence criteria for Feyman Kac

2.5  Probabilities

3  Results

3.1  Worst case probabilities for 9 growth damage models

3.2  Social Cost of Carbon Decomposition

3.3  Emission trajectory

## Code and Illustration

```python
# Required packages
import os
import sys
sys.path.append('./src')
from supportfunctions import *
sys.stdout.flush()
```

### Choosing key parameters

```python
ξp =  1  / 175  # Ambiguity Aversion Paramter
# We stored solutions for ξp =  1 / 175 to which we referred as "Ambiguity Averse" and ξp = 1000 as “Ambiguity Neutral” in the paper
# Sensible choices are from 0.0002 to 4000, while for parameters input over 0.01 the final results won't alter much

if ξp < 1:
    aversespec = "Averse"
else:
    aversespec = 'Neutral'

smart_guess = False
```

```python
# Don't run this cell if you want to solve without smart guess
# Loading Smart guesses
guess = pickle.load(open('./data/Growth{}guess.pickle'.format(aversespec), "rb", -1))
v0_guess = guess['v0']
q_guess = guess['q']
e_guess = guess['e']
base_guess = guess['base']
worst_guess = guess['worst']
smart_guess = True
```

### Solving HJB for Growth Damage

In order to solve the HJB corresponding to the planner's problem for the version of the model with climate damages to growth, we use the numerical method as we did for the model version with climate damages to consumption. The main difference comes from the expressions for the coefficients $A(x), B(x), C(x)$, and $D(x)$, which are now given by

\begin{eqnarray}
A(x) &=& - \delta \\
B(x) &=& {\widehat \mu}(x)\\
C(x) &=& {\frac 1 2} \sigma_X(x)'I\sigma_X(x) \\
D(x) &=& \delta (1 - \kappa) \left( \log \left[ {\alpha}  - {\widehat i}(x)  - {\widehat j} (x)  \right] + k  \right) + \delta \kappa \left[ \log {\widehat e}(x)   +  r \right] + {\frac {\xi_m} 2} {\widehat g}(x) \cdot {\widehat g}(x)   + \xi_p {\widehat {\mathbb I}}(x) \\
\end{eqnarray}

where $D(x)$ no longer contains the climate damage impact, but instead ${\widehat \mu}(x)$ incorporates climate damages through its impact on the drift of capital.

#### Remark:  Accounting for uncertainty about growth damages

In calculating growth damamage uncertainty, the term of interest is the contribution of smooth ambiguity to our planner's HJB equation:

\begin{equation}
-\Gamma_i(\beta F)v_{k},
\end{equation}

where the function $\Gamma_i(\beta f)$ is given by

\begin{eqnarray*}
\Gamma_i(\beta F)	=\gamma_{0,i}+\gamma_{1,i}\beta (F+\tilde{F}/\beta)+\frac{\gamma_{2,i}}{2}(\beta (F+\tilde{F}/\beta))^{2}.
\end{eqnarray*}

$\tilde{F}$ is the pre-industrial global mean temperature before climate damages negatively affect growth based on empirical measurements and estimates from \cite{BurkeDavisDiffenbaugh:2018}.

##### Selecting the values of $\gamma_1, \gamma_2$

In this setting, we take a draw of n values of $\gamma = (\gamma_1, \gamma_2)$ based on the distributional structure provided \cite{BurkeDavisDiffenbaugh:2018}:

\begin{eqnarray*}
\gamma \sim \mathcal{N}(\mu, \Sigma), \\
\mu = [\mu_1, \mu_2], \\
\Sigma = \left[ \begin{matrix}
\sigma_1^2 & \rho \\
\rho & \sigma_2^2
\end{matrix} \right].
\end{eqnarray*}

We select points based on the Hermite-Gaussian quadrature abscissas and give probability weights to these outcomes based on the pdf for $\gamma$.

##### Varying damage models

We must proceed differently for the high damage model with numerical computation.
Form
\begin{equation}
{\rm num}_i(\beta)  = \exp\left( - {\frac 1 \xi_p} \left[ -\Gamma_i(\beta F)v_{k}  \right] \right),
\end{equation}

compute:
\begin{equation}
{\rm scale}_i = \int_\beta {\rm num}_i(\beta) p(\beta) d\beta
\end{equation}

via Gauss-Hermite quadrature,
and form a density conditioned on model two:
\begin{equation}
{\widetilde q}_i(\beta)  = {\frac {{\rm num}_i(\beta)}{{\rm scale}_i}}
\end{equation}

relative to $p(\beta)$.

Also compute:
\begin{equation}
{\mathcal  I}_i =  -\xi_p \log {\rm scale}_i,
\end{equation}

along with:
\begin{equation}
{\mathcal J}_i = {\frac {-\Gamma_i(\beta F)v_{k}}{{\rm scale}_i}}
\end{equation}

using Gaussian-Hermite quadrature for the numerator integral.
Notice that ${\mathcal J}_i$ can be computed for an arbitrary $e$ since the numerator integral does not depend on $e$.
 With these computations in hand, form relative entropy conditioned on model two:
\begin{equation}
{\mathcal R}_i = {\frac 1 {\xi_p}} \left[{\mathcal I}_i - {\mathcal J}_i \right] .
\end{equation}

Notice that ${\mathcal J}_i(e)$ can be computed for an arbitrary $e$ since the numerator integral does not depend on $e$.

##### Distorted probabilities

Worst-case probabilities over models satisfy:
\begin{equation}
{\widetilde \pi}_i \  \ \propto  \  \ {\widehat \pi}_i \exp\left( -{\frac 1 {\xi_p}} {\mathcal I}_i \right)
\end{equation}

where the constant of proportionality scales the right-hand side so that the resulting ${\tilde \pi}_i$ sum to one.

```python
start_time = time.time()

McD = np.loadtxt('./data/TCRE_MacDougallEtAl2017_update.txt')
par_lambda_McD = McD / 1000

β𝘧 = np.mean(par_lambda_McD)  # Climate sensitivity parameter, MacDougall (2017)
σᵦ = np.var(par_lambda_McD, ddof=1)  # Variance of climate sensitivity parameters
λ = 1.0 / σᵦ

# Parameters are defined the same as they are in the paper
δ = 0.01
κ = 0.032
σ𝘨 = 0.02
σ𝘬 = 0.0161
σ𝘳 = 0.0339
α = 0.115000000000000
ϕ0 = 0.0600
ϕ1 = 16.666666666666668
μk = -0.034977443912449
ψ0 = 0.112733407891680
ψ1 = 0.142857142857143
F̄ = 13

# False trasient time step
ε = 0.3
# Specifying tolerance level
tol = 1e-12
# Cobweb learning rate
η = 0.05

# Solving HJB
R_min = 0
R_max = 9
F_min = 0
F_max = 750
K_min = 0
K_max = 18

# nR = 30
# nF = 30
# nK = 25

# R = np.linspace(R_min, R_max, nR)
# F = np.linspace(F_min, F_max, nF)
# K = np.linspace(K_min, K_max, nK)

# hR = R[1] - R[0]
# hF = F[1] - F[0]
# hK = K[1] - K[0]

hR = 0.05
hF = 25
hK = 0.15

R = np.arange(R_min, R_max + hR, hR)
nR = len(R)
F = np.arange(F_min, F_max + hF, hF)
nF = len(F)
K = np.arange(K_min, K_max + hK, hK)
nK = len(K)

# Here's how we discretize our state space formulating PDE, see Remark 2 in consumption setting
(R_mat, F_mat, K_mat) = np.meshgrid(R,F,K, indexing = 'ij')
stateSpace = np.hstack([R_mat.reshape(-1,1,order = 'F'),F_mat.reshape(-1,1,order = 'F'),K_mat.reshape(-1,1,order = 'F')])

μ1 = 1.272e-02
μ2 = -4.871e-04
σ1 = 3.248e-03
σ2 = 1.029e-04
ρ12 = -2.859133e-07

σ2 = σ2 * 2
ρ12 = ρ12 * 2

μ = np.matrix([-μ1, -μ2 * 2])
σ = np.matrix([[σ1 ** 2, ρ12], [ρ12, σ2 ** 2]])
Σ = np.matrix([[σᵦ,0,0],[0,σ1 ** 2, ρ12],[0, ρ12, σ2 ** 2]])

[gam1,w1] = quad_points_hermite(3)
gamm1 = np.sqrt(2) * 1 * gam1 + 0
[gam2,w2] = quad_points_hermite(3)
gamm2 = np.sqrt(2) * 1 * gam2 + 0

At = np.linalg.cholesky(σ)
x = np.zeros([2,9])
tmp = [-μ1,-μ2 * 2]
x[:,0] = tmp + At.dot([gamm1[0], gamm2[0]])
x[:,1] = tmp + At.dot([gamm1[0], gamm2[1]])
x[:,2] = tmp + At.dot([gamm1[0], gamm2[2]])
x[:,3] = tmp + At.dot([gamm1[1], gamm2[0]])
x[:,4] = tmp + At.dot([gamm1[1], gamm2[1]])
x[:,5] = tmp + At.dot([gamm1[1], gamm2[2]])
x[:,6] = tmp + At.dot([gamm1[2], gamm2[0]])
x[:,7] = tmp + At.dot([gamm1[2], gamm2[1]])
x[:,8] = tmp + At.dot([gamm1[2], gamm2[2]])

w = np.array([[w1[0],w1[0],w1[0],w1[1],w1[1],w1[1],w1[2],w1[2],w1[2]],
               [w2[0],w2[1],w2[2],w2[0],w2[1],w2[2],w2[0],w2[1],w2[2]]])
gamma1 = x[0,:]
gamma2 = x[1,:]
wgt1 = w[0,:]
wgt2 = w[1,:]

vals = np.linspace(0,30,100)

weight = np.zeros(9)
total_weight = 0
for ite in range(9):
    weight[ite] = wgt1[ite] * wgt2[ite]
    total_weight += weight[ite]

weight = weight / total_weight

gamma0 = np.zeros(9)
for ite in range(9):
    gamma0[ite] = max(-(gamma1[ite] * vals + 0.5 * gamma2[ite] * vals ** 2))

v0 = κ * R_mat + (1-κ) * K_mat

FC_Err = 1
episode = 0

if smart_guess:
    v0 = v0_guess
    q = q_guess
    e_star = e_guess
    episode = 1
    ε = 0.1

PDE_errs = []
FC_errs = []

while episode == 0 or FC_Err > tol:
    vold = v0.copy()
    # Applying finite difference scheme to the value function
    v0_dr = finiteDiff(v0,0,1,hR)
    v0_df = finiteDiff(v0,1,1,hF)
    v0_dk = finiteDiff(v0,2,1,hK)

    v0_drr = finiteDiff(v0,0,2,hR)
    v0_drr[v0_dr < 1e-16] = 0
    v0_dr[v0_dr < 1e-16] = 1e-16
    v0_dff = finiteDiff(v0,1,2,hF)
    v0_dkk = finiteDiff(v0,2,2,hK)
    if episode > 2000:
        ε = 0.1
    elif episode > 1000:
        ε = 0.2
    else:
        pass

    if episode == 0:
        Acoeff = np.ones(R_mat.shape)
        Bcoeff = ((δ * (1 - κ) * ϕ1 + ϕ0 * ϕ1 * v0_dk) * δ * (1 - κ) / (v0_dr * ψ0 * 0.5) * np.exp(0.5 * (R_mat - K_mat))) / (δ * (1 - κ) * ϕ1)
        Ccoeff = -α  - 1 / ϕ1
        j = ((-Bcoeff + np.sqrt(Bcoeff ** 2 - 4 * Acoeff * Ccoeff)) / (2 * Acoeff)) ** 2
        i = α - j - (δ * (1 - κ)) / (v0_dr * ψ0 * 0.5) * j ** 0.5 * np.exp(0.5 * (R_mat - K_mat))
        q = δ * (1 - κ) / (α - i - j)
    else:
        e_hat = e_star

        # Cobeweb scheme to update controls i and j; q is an intermediary variable that determines i and j
        Converged = 0
        nums = 0
        while Converged == 0:
            i_star = (ϕ0 * ϕ1 * v0_dk / q - 1) / ϕ1
            j_star = (q * np.exp(ψ1 * (R_mat - K_mat)) / (v0_dr * ψ0 * ψ1)) ** (1 / (ψ1 - 1))
            if α > np.max(i_star + j_star):
                q_star = η * δ * (1 - κ) / (α - i_star - j_star) + (1 - η) * q
            else:
                q_star = 2 * q
            if np.max(abs(q - q_star) / η) <= 1e-5:
                Converged = 1
                q = q_star
                i = i_star
                j = j_star
            else:
                q = q_star
                i = i_star
                j = j_star

            nums += 1

        if episode % 100 == 0:
            print('Cobweb Passed, iterations: {}, i error: {:10f}, j error: {:10f}'.format(nums, np.max(i - i_star), np.max(j - j_star)))

    # calculating growth damage uncertainties, see remark 2.1.1 for details
    a_ = []
    b_ = []
    c_ = []
    λ̃_ = []
    β̃_ = []
    I_ = []
    R_ = []
    π̃_ = []
    J_ = []

    for ite in range(9):
        a_.append( -v0_dk * (gamma0[ite] + gamma1[ite] * F̄ + 0.5 * gamma2[ite] * F̄ ** 2) )
        b_.append( -v0_dk * F_mat * (gamma1[ite] + gamma2[ite] * F̄) )
        c_.append( -v0_dk * gamma2[ite] * F_mat ** 2 )
        λ̃_.append( λ + c_[ite] / ξp )
        β̃_.append( β𝘧 - c_[ite] / ξp / λ̃_[ite] * β𝘧 - b_[ite] / (ξp * λ̃_[ite]))
        I_.append( a_[ite] - 0.5 * np.log(λ) * ξp + 0.5 * np.log(λ̃_[ite]) * ξp + 0.5 * λ * β𝘧 ** 2 * ξp - 0.5 * λ̃_[ite] * (β̃_[ite]) ** 2 * ξp )
        π̃_.append( weight[ite] * np.exp(-1 / ξp * I_[ite]) )
        J_.append( a_[ite] + b_[ite] * β̃_[ite] + 0.5 * c_[ite] * β̃_[ite] ** 2 + 0.5 * c_[ite] / λ̃_[ite] )
        R_.append((I_[ite] - J_[ite]) / ξp)

    π̃_total = sum(π̃_)
    π̃_norm_ = π̃_ / π̃_total

    B1 = v0_dr - v0_df * np.exp(R_mat)
    C1 = δ * κ
    e = C1 / B1
    e_star = e

    I_term = -1 * ξp * np.log(sum(π̃_))
    drift_distort = sum([x*y for (x,y) in zip(π̃_norm_, J_)])
    RE = sum(x * y + x * np.log(x / z) for (x,y,z) in zip(π̃_norm_, R_, weight))
    RE_total = ξp * RE

    # Formulating coefficients for the HJB False Transient problem, See remark 4 for more details
    A = -δ * np.ones(R_mat.shape)
    B_r = -e_star + ψ0 * (j ** ψ1) * np.exp(ψ1 * (K_mat - R_mat)) - 0.5 * (σ𝘳 ** 2)
    B_k = μk + ϕ0 * np.log(1 + i * ϕ1) - 0.5 * (σ𝘬 ** 2)
    B_f = e_star * np.exp(R_mat)
    C_rr = 0.5 * σ𝘳 ** 2 * np.ones(R_mat.shape)
    C_kk = 0.5 * σ𝘬 ** 2 * np.ones(R_mat.shape)
    C_ff = np.zeros(R_mat.shape)

    D = δ * κ * np.log(e_star) + δ * κ * R_mat + δ * (1 - κ) * (np.log(α - i - j) + K_mat) + I_term #  + drift_distort + RE_total

    # Solve linear system using a conjugate-gradient method in C++, see remark 5, 6 for more details
    out = PDESolver(stateSpace, A, B_r, B_f, B_k, C_rr, C_ff, C_kk, D, v0, ε, solverType='False Transient')

    out_comp = out[2].reshape(v0.shape,order = "F")

    # Calculating PDE Error and False Transient (lhs) Error
    PDE_rhs = A * v0 + B_r * v0_dr + B_f * v0_df + B_k * v0_dk + C_rr * v0_drr + C_kk * v0_dkk + C_ff * v0_dff + D
    PDE_Err = np.max(abs(PDE_rhs))
    PDE_errs.append(PDE_Err)
    FC_Err = np.max(abs((out_comp - v0)))
    FC_errs.append(FC_Err)
    if episode % 100 == 0:
        print("Episode {:d}: PDE Error: {:.10f}; False Transient Error: {:.10f}; Iterations: {:d}; CG Error: {:.10f}" .format(episode, PDE_Err, FC_Err, out[0], out[1]))
    episode += 1
    v0 = out_comp

print("Episode {:d}: PDE Error: {:.10f}; False Transient Error: {:.10f}; Iterations: {:d}; CG Error: {:.10f}" .format(episode, PDE_Err, FC_Err, out[0], out[1]))
print("--- %s seconds ---" % (time.time() - start_time))

```

#### Remark: Tolerance level and conergence criteria for HJB

Just as we did for the damages to consumption model, we use variable step sizes $\epsilon$ to reduce the time needed for solving the model while ensuring convergence. We use $\epsilon= 0.3$ for the first 1000 iterations, $\epsilon= 0.2$ for iterations 1001-2000, and $\epsilon= 0.1$ thereafter.

The table below lists the convergence criteria and final errors for solving the HJB equations for the different growth damages settings. For solving the model with damages to growth, we apply a two-stage convergence approach: we first solve the problem at a looser convergence tolerance (CG tolerance at 1e-10 and outer loop tolerance at 1e-8); then we import this solution and solve again using a tighter convergence criteria as listed in the table

| ambiguity    |    stage    | Cobweb tolerance |  outer loop tolerance  | CG tolerance |  PDE Error | dt | iterations | time|
|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|
|averse | 1  | 1e-5 | 1e-08 | 1e-10 | 1.55e-5 |variable dt| 10587 | 2.5 hrs
|averse | 2  | 1e-5 | 1e-12 | 1e-14 | 1.55e-5 | 0.1 | 20819(including stage 1) | 2 hrs
|neutral  | 1  | 1e-5 | 1e-08 | 1e-10 | 1.27e-5 | 0.1 | 12855 | 3 hrs
|neutral  | 2  | 1e-5 | 1e-12 | 1e-14 | 1.28e-5 | 0.1 |24857(including stage 1) | 2.5 hrs

### Simulation

```python
class GridInterp():

    def __init__(self, grids, values, method = 'Linear'):

        # unpacking
        self.grids = grids
        (self.xs, self.ys, self.zs) = grids
        self.nx = len(self.xs)
        self.ny = len(self.ys)
        self.nz = len(self.zs)

        self.values = values

        assert (self.nx, self.ny, self.nz) == values.shape, "ValueError: Dimensions not match"
        self.method = method

    def get_value(self, x, y, z):

        if self.method == 'Linear':

            func = RegularGridInterpolator(self.grids, self.values)
            return func([x,y,z])[0]

        elif self.method == 'Spline':

            func1 = CubicSpline(self.xs, self.values)
            yzSpace = func1(x)

            func2 = CubicSpline(self.ys, yzSpace)
            zSpace = func2(y)

            func3 = CubicSpline(self.zs, zSpace)
            return func3(z)

        else:
            raise ValueError('Method Not Supported')
```

```python
method = 'Linear'
T = 100
pers = 4 * T
dt = T / pers
nDims = 4
its = 1

gridpoints = (R, F, K)

# emissions
e_func_r = GridInterp(gridpoints, e, method)
def e_func(x):
    return e_func_r.get_value(np.log(x[0]), x[2], np.log(x[1]))

# investment in reserves
j_func_r = GridInterp(gridpoints, j, method)
def j_func(x):
    return max(j_func_r.get_value(np.log(x[0]), x[2], np.log(x[1])), 0)

# investment in productive capital
i_func_r = GridInterp(gridpoints, i, method)
def i_func(x):
    return i_func_r.get_value(np.log(x[0]), x[2], np.log(x[1]))

# first derivative of Value function w.r.t R
v_drfunc_r = GridInterp(gridpoints, v0_dr, method)
def v_drfunc(x):
    return v_drfunc_r.get_value(np.log(x[0]), x[2], np.log(x[1]))

# first derivative of Value function w.r.t F
v_dffunc_r = GridInterp(gridpoints, v0_df, method)
def v_dffunc(x):
    return v_dffunc_r.get_value(np.log(x[0]), x[2], np.log(x[1]))

# first derivative of Value function w.r.t k
v_dkfunc_r = GridInterp(gridpoints, v0_dk, method)
def v_dkfunc(x):
    return v_dkfunc_r.get_value(np.log(x[0]), x[2], np.log(x[1]))

# Value function
v_func_r = GridInterp(gridpoints, v0, method)
def v_func(x):
    return v_func_r.get_value(np.log(x[0]), x[2], np.log(x[1]))

# Distorted model probabilities
π̃_norm_func = []
π̃_norm_func_r = []

for ite in range(len(π̃_norm_)):
    π̃_norm_func_r.append(GridInterp(gridpoints, π̃_norm_[ite], method))
    π̃_norm_func.append(lambda x, ite = ite: π̃_norm_func_r[ite].get_value(np.log(x[0]), x[2], np.log(x[1])))

# Relaive entropies
RE_func_r = GridInterp(gridpoints, RE, method)
def RE_func(x):
    return RE_func_r.get_value(np.log(x[0]), x[2], np.log(x[1]))

a = []
b = []
c = []
dmg_tilt_ = []
dmg_ = []
base_driftK_r = []
base_driftK = []
tilt_driftK_r = []
tilt_driftK = []

for ite in range(len(π̃_norm_)):
    a.append(gamma0[ite] + gamma1[ite] * F̄ + 0.5 * gamma2[ite] * F̄ ** 2)
    b.append(F_mat * (gamma1[ite] + gamma2[ite] * F̄))
    c.append(gamma2[ite] * F_mat ** 2)
    dmg_tilt_.append(a[ite] + b[ite] * β̃_[ite] + 0.5 * c[ite] *
                     (β̃_[ite] ** 2) + 0.5 * c[ite] / λ̃_[ite])
    dmg_.append(a[ite] + b[ite] * βf + 0.5 * c[ite] * βf ** 2 + 0.5 * c[ite] / λ)
    # drift distortions
    base_driftK_r.append(GridInterp(gridpoints, dmg_[ite], method))
    base_driftK.append(lambda x, ite = ite: base_driftK_r[ite].get_value(np.log(x[0]), x[2], np.log(x[1])))
    tilt_driftK_r.append(GridInterp(gridpoints, dmg_tilt_[ite], method))
    tilt_driftK.append(lambda x, ite = ite: tilt_driftK_r[ite].get_value(np.log(x[0]), x[2], np.log(x[1])))

# Weighted sum of drift distortion
def Gamma_base(x):
    res = 0
    for ite in range(len(weight)):
        res += weight[ite] * base_driftK[ite](x)
    return res

def Gamma_tilted(x):
    res = 0
    for ite in range(len(weight)):
        res += π̃_norm_func[ite](x) * tilt_driftK[ite](x)
    return res

# drifts for each diffusion process
def muR(x):
    return -e_func(x) + ψ0 * (j_func(x) * x[1] / x[0]) ** ψ1
def muK_tilted(x):
    return (μk + ϕ0 * np.log(1 + i_func(x) * ϕ1)- Gamma_tilted(x))
def muK_base(x):
    return (μk + ϕ0 * np.log(1 + i_func(x) * ϕ1)- Gamma_base(x))
def muF(x):
    return e_func(x) * x[0]

# volatilities for each diffusion process
def sigmaR(x):
    return np.zeros(x[:4].shape)
def sigmaK(x):
    return np.zeros(x[:4].shape)
def sigmaF(x):
    return np.zeros(x[:4].shape)

# initial points
R_0 = 650
K_0 = 80 / α
F_0 = 870 - 580
initial_val = np.array([R_0, K_0, F_0])

# Set bounds
R_max_sim = np.exp(max(R))
K_max_sim = np.exp(max(K))
F_max_sim = max(F)

R_min_sim = np.exp(min(R))
K_min_sim = np.exp(min(K))
F_min_sim = min(F)

upperbounds = np.array([R_max_sim, K_max_sim, F_max_sim, K_max_sim])
lowerbounds = np.array([R_min_sim, K_min_sim, F_min_sim, K_min_sim])

hists = np.zeros([pers, nDims, its])
e_hists = np.zeros([pers,its])
j_hists = np.zeros([pers,its])
i_hists = np.zeros([pers,its])
RE_hists = np.zeros([pers,its])

pi_tilde_hists = []
for ite in range(len(π̃_norm_)):
    pi_tilde_hists.append(np.zeros([pers,its]))

for iters in range(its):
    hist = np.zeros([pers,nDims])
    e_hist = np.zeros([pers,1])
    i_hist = np.zeros([pers,1])
    j_hist = np.zeros([pers,1])

    RE_hist = np.zeros([pers,1])
    pi_tilde_hist = []
    for ite in range(len(π̃_norm_)):
        pi_tilde_hist.append(np.zeros([pers,1]))

    hist[0,:] = [R_0, K_0, F_0, K_0]
    e_hist[0] = e_func(hist[0,:]) * hist[0,0]
    i_hist[0] = i_func(hist[0,:]) * hist[0,1]
    j_hist[0] = j_func(hist[0,:]) * hist[0,0]
    RE_hist[0] = RE_func(hist[0,:])

    for ite in range(len(π̃_norm_)):
        pi_tilde_hist[ite][0] = π̃_norm_func[ite](hist[0,:])

    for tm in range(1,pers):
        shock = norm.rvs(0,np.sqrt(dt),nDims)
        hist[tm,0] = cap(hist[tm-1,0] * np.exp((muR(hist[tm-1,:])- 0.5 * sum((sigmaR(hist[tm-1,:])) ** 2))* dt + sigmaR(hist[tm-1,:]).dot(shock)),lowerbounds[0], upperbounds[0])
        hist[tm,1] = cap(hist[tm-1,1] * np.exp((muK_base(hist[tm-1,:])- 0.5 * sum((sigmaK(hist[tm-1,:])) ** 2))* dt + sigmaK(hist[tm-1,:]).dot(shock)),lowerbounds[1], upperbounds[1])
        hist[tm,2] = cap(hist[tm-1,2] + muF(hist[tm-1,:]) * dt + sigmaF(hist[tm-1,:]).dot(shock), lowerbounds[2], upperbounds[2])
        hist[tm,3] = cap(hist[tm-1,3] * np.exp((muK_tilted(hist[tm-1,:])- 0.5 * sum((sigmaK(hist[tm-1,:])) ** 2))* dt + sigmaK(hist[tm-1,:]).dot(shock)),lowerbounds[3], upperbounds[3])

        e_hist[tm] = e_func(hist[tm-1,:]) * hist[tm-1,0]
        i_hist[tm] = i_func(hist[tm-1,:]) * hist[tm-1,1]
        j_hist[tm] = j_func(hist[tm-1,:]) * hist[tm-1,1]
        RE_hist[tm] = RE_func(hist[tm-1, :])

        for ite in range(len(π̃_norm_)):
            pi_tilde_hist[ite][tm] = π̃_norm_func[ite](hist[tm-1,:])

    hists[:,:,iters] = hist
    e_hists[:,[iters]] = e_hist
    i_hists[:,[iters]] = i_hist
    j_hists[:,[iters]] = j_hist

    RE_hists[:,[iters]] =  RE_hist

    for ite in range(len(π̃_norm_)):
        pi_tilde_hists[ite][:,[iters]] = pi_tilde_hist[ite]
```

### SCC Calculation Feyman Kac

The growth damages model setting is solved using the same numerical method as the consumption damage model, with a few minor adjustments made to the equations and calculations. First, the expressions for $\Psi^*(x)$ and $\bar{\Psi}(x)$ are now given by

\begin{align*}
 \Psi^*(x) & = -  \int_{\Theta} V_k(x) \left[ \nabla \Gamma(\beta f) \beta \right] q^*(\theta | x) P(d\theta) \\
 \bar{\Psi}(x) & = -  \int_{\Theta} V_k(x) \left[ \nabla \Gamma(\beta f) \beta \right] P(d\theta).
\end{align*}

Second, these integral expressions can be calculated using entirely analytical simplifictions as will be outlined below.

Finally, there are no longer direct contributions to the damage evolution that have integral expressions that need to be numerically calculated and added to the PDEs derived from application of the F-K formula.

With these differences we can apply the analogous approach to the consumption damage setting to derive $ecc^*$, $\bar{ecc}$, and $ucc^*$ for the growth damages model setting. As before, after imposing the same finite difference scheme the PDEs are characterized by unconditionally linear equations that we can solve directly by approximately inverting the linear systems using the conjugate gradient method without the use of the false transient or other iterative schemes. These PDEs can also be alternatively be solved using the implicit, finite-difference scheme with the conjugate gradient solver augmented by a false transient where the coefficients are definied as in the consumption damage case, appropriately adjusted for the $\mu^*_X[x,a^*(x)]$, $\mu_X[x,a^*(x)]$, $\Psi^*(x)$, and $\bar{\Psi}(x)$ that correspond to the growth damages model.

```python
a, b, c, dmg_ = ([] for ite in range(4))
for ite in range(len(π̃_norm_)):
    a.append(np.zeros(F_mat.shape) )
    b.append(v0_dk * (gamma1[ite] + gamma2[ite] * F̄))
    c.append(2 * v0_dk * gamma2[ite] * F_mat)
    dmg_.append(a[ite] + b[ite] * βf + 0.5 * c[ite] * βf ** 2 + 0.5 * c[ite] / λ)

flow_base = sum(w * d for w, d in zip(weight, dmg_))

a, b, c, dmg_ = ([] for ite in range(4))
for ite in range(len(π̃_norm_)):
    a.append(v0_dk * (gamma0[ite] + gamma1[ite] * F̄ + 0.5 * gamma2[ite] * F̄ ** 2))
    b.append(v0_dk * F_mat * (gamma1[ite] + gamma2[ite] * F̄))
    c.append(v0_dk * gamma2[ite] * F_mat ** 2)
    dmg_.append(a[ite] + b[ite] * βf + 0.5 * c[ite] * βf ** 2 + 0.5 * c[ite] / λ)

Gamma_base = sum(w * d for w, d in zip(weight, dmg_))

A = -δ * np.ones(R_mat.shape)
# B_r = -e + ψ0 * (j ** ψ1) - 0.5 * (σ𝘳 ** 2)
B_r = -e + ψ0 * (j ** ψ1) * np.exp(ψ1 * (K_mat - R_mat)) - 0.5 * (σ𝘳 ** 2)
B_k = μk + ϕ0 * np.log(1 + i * ϕ1) - 0.5 * (σ𝘬 ** 2) - Gamma_base
B_f = e * np.exp(R_mat)
C_rr = 0.5 * σ𝘳 ** 2 * np.ones(R_mat.shape)
C_kk = 0.5 * σ𝘬 ** 2 * np.ones(R_mat.shape)
C_ff = np.zeros(R_mat.shape)
D = flow_base

if smart_guess:
     out = PDESolver(stateSpace, A, B_r, B_f, B_k, C_rr, C_ff, C_kk, D, base_guess, ε, smartguess = True, solverType='Feyman Kac')

else:
    base_guess = np.zeros(R_mat.shape)
    out = PDESolver(stateSpace, A, B_r, B_f, B_k, C_rr, C_ff, C_kk, D, base_guess, ε, smartguess = False, solverType='Feyman Kac')

v0_base = out[2].reshape(v0.shape, order="F")
v0_base = v0_base

v0_dr_base = finiteDiff(v0_base,0,1,hR,1e-16)
v0_df_base = finiteDiff(v0_base,1,1,hF)
v0_dk_base = finiteDiff(v0_base,2,1,hK)

v0_drr_base = finiteDiff(v0_base,0,2,hR)
v0_dff_base = finiteDiff(v0_base,1,2,hF)
v0_dkk_base = finiteDiff(v0_base,2,2,hK)

PDE_rhs = A * v0_base + B_r * v0_dr_base + B_f * v0_df_base + B_k * v0_dk_base + C_rr * v0_drr_base + C_kk * v0_dkk_base + C_ff * v0_dff_base + D
PDE_Err = np.max(abs(PDE_rhs))
print("Feyman Kac Base Model Solved. PDE Error: {:.10f}; Iterations: {:d}; CG Error: {:.10f}".format(PDE_Err, out[0], out[1]))

a, b, c, dmg_tilt_ = ([] for ite in range(4))
for ite in range(len(π̃_norm_)):
    a.append(np.zeros(F_mat.shape) )
    b.append(v0_dk * (gamma1[ite] + gamma2[ite] * F̄))
    c.append(2 * v0_dk * gamma2[ite] * F_mat)
    dmg_tilt_.append(a[ite] + b[ite] * β̃_[ite] + 0.5 * c[ite] * β̃_[ite] ** 2 + 0.5 * c[ite] / λ̃_[ite])

flow_tilted = sum(w * d for w, d in zip(π̃_norm_, dmg_tilt_))

a, b, c, dmg_tilt_ = ([] for ite in range(4))
for ite in range(len(π̃_norm_)):
    a.append(v0_dk * (gamma0[ite] + gamma1[ite] * F̄ + 0.5 * gamma2[ite] * F̄ ** 2))
    b.append(v0_dk * F_mat * (gamma1[ite] + gamma2[ite] * F̄))
    c.append(v0_dk * gamma2[ite] * F_mat ** 2)
    dmg_tilt_.append(a[ite] + b[ite] * β̃_[ite] + 0.5 * c[ite] * β̃_[ite] ** 2 + 0.5 * c[ite] / λ̃_[ite])

Gamma_tilted = sum(w * d for w, d in zip(π̃_norm_, dmg_tilt_))

A = -δ * np.ones(R_mat.shape)
# B_r = -e + ψ0 * (j ** ψ1) - 0.5 * (σ𝘳 ** 2)
B_r = -e + ψ0 * (j ** ψ1) * np.exp(ψ1 * (K_mat - R_mat)) - 0.5 * (σ𝘳 ** 2)
B_k = μk + ϕ0 * np.log(1 + i * ϕ1) - 0.5 * (σ𝘬 ** 2) - Gamma_tilted
B_f = e * np.exp(R_mat)
C_rr = 0.5 * σ𝘳 ** 2 * np.ones(R_mat.shape)
C_kk = 0.5 * σ𝘬 ** 2 * np.ones(R_mat.shape)
C_ff = np.zeros(R_mat.shape)
D = flow_tilted

if smart_guess:
    out = PDESolver(stateSpace, A, B_r, B_f, B_k, C_rr, C_ff, C_kk, D, worst_guess, ε, smartguess = True, solverType='Feyman Kac')

else:
    worst_guess = np.zeros(R_mat.shape)
    out = PDESolver(stateSpace, A, B_r, B_f, B_k, C_rr, C_ff, C_kk, D, worst_guess, ε, smartguess = False, solverType='Feyman Kac')

v0_worst = out[2].reshape(v0.shape, order="F")
v0_worst = v0_worst

v0_dr_worst = finiteDiff(v0_worst,0,1,hR,1e-16)
v0_df_worst = finiteDiff(v0_worst,1,1,hF)
v0_dk_worst = finiteDiff(v0_worst,2,1,hK)

v0_drr_worst = finiteDiff(v0_worst,0,2,hR)
v0_dff_worst = finiteDiff(v0_worst,1,2,hF)
v0_dkk_worst = finiteDiff(v0_worst,2,2,hK)

PDE_rhs = A * v0_worst + B_r * v0_dr_worst + B_f * v0_df_worst + B_k * v0_dk_worst + C_rr * v0_drr_worst + C_kk * v0_dkk_worst + C_ff * v0_dff_worst + D
PDE_Err = np.max(abs(PDE_rhs))
print("Feyman Kac Worst Model Solved. PDE Error: {:.10f}; Iterations: {:d}; CG Error: {:.10f}".format(PDE_Err, out[0], out[1]))
```

```python
# Total SCC
MC = δ * (1-κ) / (α * np.exp(K_mat) - i * np.exp(K_mat) - j * np.exp(K_mat))
ME = δ * κ / (e * np.exp(R_mat))
SCC = 1000 * ME / MC
SCC_func_r = GridInterp(gridpoints, SCC, method)

def SCC_func(x):
#     return SCC_func_r.get_value(np.log(x[0]), x[2], np.log(x[3]))
    return SCC_func_r.get_value(np.log(x[0]), x[2], np.log(x[1]))

# Private SCC
ME1 = v0_dr * np.exp(-R_mat)
SCC1 = 1000 * ME1 / MC
SCC1_func_r = GridInterp(gridpoints, SCC1, method)
def SCC1_func(x):
#     return SCC1_func_r.get_value(np.log(x[0]), x[2], np.log(x[3]))
    return SCC1_func_r.get_value(np.log(x[0]), x[2], np.log(x[1]))

# Feyman Kac solution under original unadjusted probabilities
ME2_base = v0_base
SCC2_base = 1000 * ME2_base / MC
SCC2_base_func_r = GridInterp(gridpoints, SCC2_base, method)
def SCC2_base_func(x):
#     return SCC2_base_func_r.get_value(np.log(x[0]), x[2], np.log(x[3]))
    return SCC2_base_func_r.get_value(np.log(x[0]), x[2], np.log(x[1]))

ME2_base_a = -v0_df
SCC2_base_a = 1000 * ME2_base_a / MC
SCC2_base_a_func_r = GridInterp(gridpoints, SCC2_base_a, method)
def SCC2_base_a_func(x):
#     return SCC2_base_a_func_r.get_value(np.log(x[0]), x[2], np.log(x[3]))
    return SCC2_base_a_func_r.get_value(np.log(x[0]), x[2], np.log(x[1]))

# Feyman Kac solution under ambiguity-adjusted probability
ME2_tilt = v0_worst
SCC2_tilt = 1000 * ME2_tilt / MC
SCC2_tilt_func_r = GridInterp(gridpoints, SCC2_tilt, method)
def SCC2_tilt_func(x):
#     return SCC2_tilt_func_r.get_value(np.log(x[0]), x[2], np.log(x[3]))
    return SCC2_tilt_func_r.get_value(np.log(x[0]), x[2], np.log(x[1]))

# Interpolating via simulated trajectories
SCC_values = np.zeros([pers,its])
SCC1_values = np.zeros([pers,its])
SCC2_base_values = np.zeros([pers,its])
SCC2_tilt_values = np.zeros([pers,its])
SCC2_base_a_values = np.zeros([pers,its])

for tm in range(pers):   # time horizons
    for path in range(its):    # Number of trajectories, currently it's 1
        SCC_values[tm, path] = SCC_func(hists[tm,:,path])
        SCC1_values[tm, path] = SCC1_func(hists[tm,:,path])
        SCC2_base_values[tm, path] = SCC2_base_func(hists[tm,:,path])
        SCC2_tilt_values[tm, path] = SCC2_tilt_func(hists[tm,:,path])
        SCC2_base_a_values[tm, path] = SCC2_base_a_func(hists[tm,:,path])

SCC_total = np.mean(SCC_values,axis = 1)
SCC_private = np.mean(SCC1_values,axis = 1)
SCC2_FK_base = np.mean(SCC2_base_values,axis = 1)
SCC2_FK_tilt = np.mean(SCC2_tilt_values,axis = 1)

uncertain = SCC2_FK_tilt - SCC2_FK_base

uncertain = SCC2_FK_tilt - SCC2_FK_base

SCCs = {}
SCCs['SCC'] = SCC_total   # Total SCC
SCCs['SCC1'] = SCC_private   # private SCC
SCCs['SCC2'] = SCC2_FK_base   # external SCC
SCCs['SCC3'] = SCC2_FK_tilt - SCC2_FK_base  # Uncertainty SCC
```

#### Remark: Convergence criteria for Feyman Kac

Below table listed the convergence criteria and final errors for solving HJB equations at different growth damages setting which we used in the paper. As the discretized state space contains over 3 million points, even the Conjugate Gradient approximation is quite slow to converge to a small error. Therefore, we cap the max iterations for the conjugate gradient solver to be 400,000.

| ambiguity    |   damage specification   | iterations | CG Error |  PDE Error | Approx time
|:-:|:-:|:-:|:-:|:-:|:-:|
|averse | base  | 340,113 | 1e-14 | 4.50e-6|2 hrs
|averse | worst | 391,635 | 1e-14| 2.21e-6| 2.5 hrs
|neutral | base  | 577,167 | 1e-14 | 2.21e-6|3.5 hrs
|neutral | worst | 556,956 | 1e-14 | 2.21e-6|3.5 hrs

### Probabilities

Here we caculated and plotted the ${\widetilde q}_2(\beta)$ implied Ambiguity-Adjusted Probabilities for 9 growth damage models

```python
Dists = {}

a = β𝘧 - 5 * np.sqrt(σᵦ)
b = β𝘧 + 5 * np.sqrt(σᵦ)
a_10std = β𝘧 - 10 * np.sqrt(σᵦ)
b_10std = β𝘧 + 10 * np.sqrt(σᵦ)
beta_f_space = np.linspace(a_10std,b_10std,200)

original_dist = norm.pdf(beta_f_space, β𝘧, np.sqrt(σᵦ))
Dists['Original'] = original_dist

R_value = np.mean(hists[:,0,:], axis = 1)
K_value = np.mean(hists[:,1,:], axis = 1)
F_value = np.mean(hists[:,2,:], axis = 1)

R_func_r = []
R_func = []
β̃_func_r = []
β̃_func = []
λ̃_func_r = []
λ̃_func = []
for ite in range(len(R_)):
    R_func_r.append(GridInterp(gridpoints, R_[ite], method))
    R_func.append(lambda x, ite = ite: R_func_r[ite].get_value(np.log(x[0]), x[2], np.log(x[1])))
    β̃_func_r.append(GridInterp(gridpoints, β̃_[ite], method))
    β̃_func.append(lambda x, ite = ite: β̃_func_r[ite].get_value(np.log(x[0]), x[2], np.log(x[1])))
    λ̃_func_r.append(GridInterp(gridpoints, λ̃_[ite], method))
    λ̃_func.append(lambda x, ite = ite: λ̃_func_r[ite].get_value(np.log(x[0]), x[2], np.log(x[1])))

hists_mean = np.mean(hists, axis = 2)
RE_plot = np.zeros(pers)
weight_plot = [np.zeros([pers,1]) for ite in range(len(π̃_norm_func))]
for tm in range(pers):
    RE_plot[tm] = RE_func(hists_mean[tm,:])
    for ite in range(len(π̃_norm_func)):
        weight_plot[ite][tm] = π̃_norm_func[ite](hists_mean[tm,:])

for tm in [1,100,200,300,400]:
    R0 = hists[tm-1,0,0]
    K0 = hists[tm-1,1,0]
    F0 = hists[tm-1,2,0]
    weights_prob = []
    mean_distort_ = []
    lambda_tilde_ = []
    tilt_dist_ = []

    for ite in range(len(π̃_norm_func)):
        weights_prob.append(π̃_norm_func[ite]([R0, K0, F0]))
        mean_distort_.append(β̃_func[ite]([R0, K0, F0]) - βf)
        lambda_tilde_.append(λ̃_func[ite]([R0, K0, F0]))
        tilt_dist_.append(norm.pdf(beta_f_space, mean_distort_[ite] + βf, 1 / np.sqrt(lambda_tilde_[ite])))

    weighted = sum(w * til for w, til in zip(weights_prob, tilt_dist_) )
    Dists['Year' + str(int((tm) / 4))] = dict(tilt_dist = tilt_dist_, weighted = weighted, weights = weights_prob)
```

## Results

### Worst case probabilities for 9 growth damage models

```python
growthdensityPlot(beta_f_space, Dists)
```

### Social Cost of Carbon Decomposition

```python
growthSCCDecomposePlot(SCCs, ξp)
```

### Emission trajectory

```python
growthemissionPlot(ξp, e_hists)
```

```python

```


```python
%matplotlib inline
```

    Bad key "text.kerning_factor" on line 4 in
    C:\Users\sensio\miniconda3\lib\site-packages\matplotlib\mpl-data\stylelib\_classic_test_patch.mplstyle.
    You probably need to get an updated matplotlibrc file from
    https://github.com/matplotlib/matplotlib/blob/v3.1.3/matplotlibrc.template
    or from the matplotlib source distribution

# How nangs works

In this notebook we show how nangs works. We will solve the following PDE:

\begin{equation}
    \frac{\partial \phi}{\partial t} + u \frac{\partial \phi}{\partial x} = 0
\end{equation}

The independent variables (i.e, $x$ and $t$) are used as input values for the NN, and the solution (i.e. $\phi$) is the output. In order to find the solution, at each step the NN outputs are derived w.r.t the inputs. Then, a loss function that matches the PDE is built and the weights are updated accordingly. If the loss function goes to zero, we can assume that our NN is indeed the solution to our PDE.

```python
# imports

import numpy as np
import matplotlib.pyplot as plt
import torch

device = "cuda" if torch.cuda.is_available() else "cpu"

torch.__version__
```

    '1.5.0'

## Define your data

To solve the PDE we first need a set of points to evaluate it, we will use this points as the dataset for training de NN.

```python
# define the mesh
x = np.linspace(0,1,20)
t = np.linspace(0,1,30)

# combine all points
m = np.meshgrid(x, t)
m = np.stack(m, -1).reshape(-1, 2)
m.shape
```

    (600, 2)

```python
class Dataset(torch.utils.data.Dataset):
    def __init__(self, X):
        self.X = torch.from_numpy(X).to(device).float()
    def __len__(self):
        return len(self.X)
    def __getitem__(self, ix):
        return self.X[ix]

dataset = Dataset(m)

len(dataset)
```

    600

## Define your solution topology

We use a Multilayer Perceptron to approximate the solution to the PDE

```python
# solution topology

class Sine(torch.nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self, x):
        return torch.sin(x)

def block(i, o):
    fc = torch.nn.Linear(i, o)
    return torch.nn.Sequential(
        Sine(),
        torch.nn.Linear(i, o)
    )

class MLP(torch.nn.Module):
    def __init__(self, inputs, outputs, layers, neurons):
        super().__init__()
        fc_in = torch.nn.Linear(inputs, neurons)
        fc_hidden = [
            block(neurons, neurons)
            for layer in range(layers-1)
        ]
        fc_out = block(neurons, outputs)

        self.mlp = torch.nn.Sequential(
            fc_in,
            *fc_hidden,
            fc_out
        )

    def forward(self, x):
        return self.mlp(x)

mlp = MLP(2, 1, 3, 100)
mlp
```

    MLP(
      (mlp): Sequential(
        (0): Linear(in_features=2, out_features=100, bias=True)
        (1): Sequential(
          (0): Sine()
          (1): Linear(in_features=100, out_features=100, bias=True)
        )
        (2): Sequential(
          (0): Sine()
          (1): Linear(in_features=100, out_features=100, bias=True)
        )
        (3): Sequential(
          (0): Sine()
          (1): Linear(in_features=100, out_features=1, bias=True)
        )
      )
    )

```python
# check output shape

mlp(torch.randn(100,2)).shape
```

    torch.Size([100, 1])

## Boundary conditions

We can attempt to solve our PDE at this points, but we would obtain a trivial solution. Instead, we need to specify a set of boundary conditions.

```python
# initial condition (t = 0)

t0 = np.array([0.])
m0 = np.meshgrid(x, t0)
m0 = np.stack(m0, -1).reshape(-1, 2)

p0 = np.sin(2*np.pi*x)

plt.plot(x, p0)
plt.grid(True)
plt.xlabel('x')
plt.ylabel('$p_0$')
plt.show()
```

```python
class DirichletDataset(torch.utils.data.Dataset):
    def __init__(self, X, Y):
        assert len(X) == len(Y)
        self.X = torch.from_numpy(X).to(device).float()
        self.Y = torch.from_numpy(Y).to(device).float()
    def __len__(self):
        return len(self.X)
    def __getitem__(self, ix):
        return self.X[ix], self.Y[ix]

# we use the names to indicate the order of the variables in the data
initial_condition_dataset = DirichletDataset(m0, p0.reshape(-1, 1))

len(initial_condition_dataset)
```

    20

```python
# boundary conditions (peridic conditions at x = 0 and x = 1)

xb0 = np.array([0])
mb0 = np.meshgrid(xb0, t)
mb0 = np.stack(mb0, -1).reshape(-1, 2)
mb0.shape
```

    (30, 2)

```python
xb1 = np.array([1])
mb1 = np.meshgrid(xb1, t)
mb1 = np.stack(mb1, -1).reshape(-1, 2)
mb1.shape
```

    (30, 2)

```python
boco_dataset = DirichletDataset(mb0, mb1)

len(boco_dataset)
```

    30

## Solve the PDE

We can now proceed with solving the PDE

```python
BATCH_SIZE = 32
EPOCHS = 50
U = 1

dataloaders = {
    'inner': torch.utils.data.DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True),
    'initial': torch.utils.data.DataLoader(initial_condition_dataset, batch_size=BATCH_SIZE, shuffle=True),
    'periodic': torch.utils.data.DataLoader(boco_dataset, batch_size=BATCH_SIZE, shuffle=True)
}

mlp = MLP(2, 1, 3, 128)
mlp.to(device)
optimizer = torch.optim.Adam(mlp.parameters())
criterion = torch.nn.MSELoss()

hist = []
for epoch in range(1, EPOCHS+1):
    # iterate over the internal points in batches
    for batch in dataloaders['inner']:
        X = batch
        optimizer.zero_grad()
        # optimize for boundary points
        losses = 0
        for batch in dataloaders['initial']:
            x, y = batch
            p = mlp(x)
            loss = criterion(p, y)
            loss.backward()
            losses += loss.item()
        for batch in dataloaders['periodic']:
            x1, x2 = batch
            p1 = mlp(x1)
            p2 = mlp(x2)
            loss = criterion(p1, p2)
            loss.backward()
            losses += loss.item()
        # optimize for internal points
        X.requires_grad = True
        p = mlp(X)
        grads, = torch.autograd.grad(p, X,
                    grad_outputs=p.data.new(p.shape).fill_(1),
                    create_graph=True, only_inputs=True)
        dpdx, dpdt = grads[:,0], grads[:,1]
        pde = dpdt + U*dpdx
        loss = pde.pow(2).mean()
        loss.backward()
        optimizer.step()
        losses += loss.item()
        hist.append(losses)
    print(f"Epoch {epoch}/{EPOCHS} loss {losses:.5f}")
```

    Epoch 1/50 loss 0.44584
    Epoch 2/50 loss 0.43346
    Epoch 3/50 loss 0.42094
    Epoch 4/50 loss 0.40049
    Epoch 5/50 loss 0.35735
    Epoch 6/50 loss 0.17384
    Epoch 7/50 loss 0.11475
    Epoch 8/50 loss 0.09106
    Epoch 9/50 loss 0.07281
    Epoch 10/50 loss 0.06439
    Epoch 11/50 loss 0.05049
    Epoch 12/50 loss 0.02657
    Epoch 13/50 loss 0.00799
    Epoch 14/50 loss 0.00553
    Epoch 15/50 loss 0.00346
    Epoch 16/50 loss 0.00295
    Epoch 17/50 loss 0.00208
    Epoch 18/50 loss 0.00168
    Epoch 19/50 loss 0.00208
    Epoch 20/50 loss 0.00168
    Epoch 21/50 loss 0.00113
    Epoch 22/50 loss 0.00112
    Epoch 23/50 loss 0.00136
    Epoch 24/50 loss 0.00091
    Epoch 25/50 loss 0.00076
    Epoch 26/50 loss 0.00076
    Epoch 27/50 loss 0.00062
    Epoch 28/50 loss 0.00096
    Epoch 29/50 loss 0.00104
    Epoch 30/50 loss 0.00077
    Epoch 31/50 loss 0.00039
    Epoch 32/50 loss 0.00048
    Epoch 33/50 loss 0.00067
    Epoch 34/50 loss 0.00038
    Epoch 35/50 loss 0.00036
    Epoch 36/50 loss 0.00030
    Epoch 37/50 loss 0.00043
    Epoch 38/50 loss 0.00046
    Epoch 39/50 loss 0.00067
    Epoch 40/50 loss 0.00049
    Epoch 41/50 loss 0.00140
    Epoch 42/50 loss 0.00033
    Epoch 43/50 loss 0.00036
    Epoch 44/50 loss 0.00031
    Epoch 45/50 loss 0.00135
    Epoch 46/50 loss 0.00040
    Epoch 47/50 loss 0.00054
    Epoch 48/50 loss 0.00025
    Epoch 49/50 loss 0.00031
    Epoch 50/50 loss 0.00045

```python
plt.plot(hist)
plt.xlabel("update step")
plt.ylabel("loss")
plt.yscale("log")
plt.grid(True)
plt.show()
```

## Evaluate your solution

```python
def build_mesh(N, t):
    x = np.linspace(0,1,N)
    _t = np.array([t])
    m = np.meshgrid(x, _t)
    m = np.stack(m, -1).reshape(-1, 2)
    return x, t, m

x, t, m = build_mesh(20, 0)
```

```python
def eval_model(m):
    mlp.eval()
    with torch.no_grad():
        p = mlp(torch.tensor(m).float().to(device)).cpu().numpy()
    return p

p = eval_model(m)
```

```python
def plot_model(x, p, t):
    pe = np.sin(2.*np.pi*(x-U*t)).reshape(-1,1)
    plt.plot(x, pe, label="exact")
    plt.plot(x, p, '.k', label="solution")
    plt.legend()
    plt.grid(True)
    l2 = np.sqrt(np.sum((p-pe)**2))
    plt.title(f"t = {t:.3f} (L2 = {l2:.5f})")
    plt.show()

plot_model(x, p, 0)
```

```python
def eval_solution(N, t):
    x, t, m = build_mesh(N, t)
    p = eval_model(m)
    plot_model(x, p, t)
```

```python
from matplotlib import animation, rc
rc('animation', html='html5')

t = np.linspace(0,1,10)
x, t, m = build_mesh(30, t)
p = eval_model(m).reshape(len(t), -1)

fig = plt.figure()
ax = plt.subplot(111)
def update_plot(i):
    ax.clear()
    pe = np.sin(2.*np.pi*(x-U*t[i]))
    ax.plot(x, pe, label=f"exact (u = {U})")
    ax.plot(x, p[i], '.k', label="solution")
    ax.set_xlabel("x", fontsize=14)
    ax.set_ylabel("p", fontsize=14, rotation=np.pi/2)
    ax.legend(loc="upper right")
    ax.grid(True)
    ax.set_xlim([0, 1])
    ax.set_ylim([-1.2, 1.2])
    l2 = np.sqrt(np.sum((p[i]-pe)**2))
    ax.set_title(f"t = {t[i]:.3f} (L2 = {l2:.5f})")
    return ax

anim = animation.FuncAnimation(fig, update_plot, frames=len(t), interval=300)
plt.close()

```

```python
anim
```

## What nangs provides you

Nangs will provideo you with classes and functions to asbtract some of the concepts explained here in order to enable fast experimentation while allowing you to customize it to your needs.


```
%nbtoc
```

Contents [hide]

# Introduction

## Abstract

This notebook demonstrates the application of geometric algebra to projective geometry. The calculations are performed using the symbolic `galgebra` module from  [sympy](www.sympy.org), which is found in Alan Bromborsky's [GitHub](https://github.com/brombo/) repo's.

## Motivation

A common theme in applied mathematics is that a  complicated problem in a simple space can be traded for a simple problem in a more complicated space. Think of using the Fourier transform to linearize convolution.

```
from IPython.display import *
Image('pics/generic commutative diagram.png')
```

Similarly, Projective Geometry is slightly more complicated than Euclidean space, but it makes certain transformations, like translation and shearing,  simpler.  Additionally, it removes the perculiarities of  zero and infinity. Computer graphics people use projective geometry regularly.

While these features provide sufficient  motivation, there are more fundamental reason to study alternative geometries. One reason is to illustrate the freedom of algebraic interpretation. In Eucliean space a vector is a vector is a vector. In projective geometry a vector may be represented as a bivector. We can choose to think of a vector as a vector or as a bivector to suite our needs.

So, having more ways to view a given problem is never a bad thing, and can only add to the richness of possible solutions. Lets roll.

More information can be found in the [references](#References) below.

## References

* [1]C. Doran, Geometric Algebra for Physicists, 1st Pbk. Ed. with Corr edition. Cambridge; New York: Cambridge University Press, 2007.
* [2]J. Vince, Geometric Algebra for Computer Graphics, 2008 edition. London: Springer, 2008.
* [3]C. Perwass, Geometric Algebra with Applications in Engineering, Softcover reprint of hardcover 1st ed. 2009. Springer, 2010.

# The Projective Plane

## Theory

The projective plane describes the application of projective geomtry to 2D space, illustrated in the image below (taken from [[2]](#References) ). In this example, the orginal 2D space created by the basis $e_1, e_2$ is extended by a third orthonormal vector $n$. A vector in the plane (  $q$ ) is represented by the vector in 3D pace ( $p$ ). From this geometric construction we can relate vectors in both spaces.

```
Image('pics/projective plane.png')
```

### Projecting Down onto the plane

First we need to define the mapping between a vector in G3, and the
projective plane. From the illustration above it is clear that

\begin{eqnarray*}
n+q & = & \lambda p
\end{eqnarray*}

Taking the inner product with $n$, solving for $\lambda$

\begin{eqnarray*}
n^{2}+n\cdot q & = & \lambda p\cdot n\\
n^{2} & = & \lambda p\cdot n\\
\lambda & = & \frac{n^{2}}{p\cdot n}
\end{eqnarray*}

Combining this with the equation above, solving for $q$

\begin{eqnarray*}
q & = & \frac{n^{2}}{p\cdot n}p-n
\end{eqnarray*}

This can be re-arranged, to provide

\begin{eqnarray*}
q & = & \frac{n^{2}}{p\cdot n}p-n\\
 & = & \frac{n^{2}p-nn\cdot p}{p\cdot n}\\
 & = & n\frac{np-n\cdot p}{p\cdot n}\\
 & = & n\frac{n\wedge p}{n\cdot p}\\
 q & = & \frac{p\wedge n}{p\cdot n} n\\
\end{eqnarray*}

This form illustrates that the vector q can be put into correspondence with the bivector $P= \frac{p \wedge n}{p\cdot n}$.  This is a nice visual interpretation: the bivector $P$ is formed by extending the $n$ vector to the $p$ vector, which sweeps out $q$ in the projection plane.

So a vector is the extension of the zero  point ( $n$ ) to a point $p$.

It also demonstrates that the magnitude of $p$ is irrelevant to the vector $q$, as it will get factored out of the fraction.

### Projecting up into G3

Projecting a vector from the plane into 3D space is trivial, due to the homogeneity of the perspective relationship. Setting $\lambda =1$

$$p= q+n$$

note that this implies,

$$\lambda n\wedge p= n\wedge q$$

## Ok lets do it

First we build the geometric algebra of 3D space  using an orthonormal basis of $\boldsymbol{(e_1, e_2, n)}$, using `Ga.build()`. This returns the object representing the algebra, and the basis vectors.

 * (For an intruduction to the `sympy.galgebra` module, see the ** [[ Sould make an Intro Tutorial ]] **)

```
from sympy import *
from galgebra.ga import Ga

# create G3 from the  orthonoormal basis: e1, e2,n.
(ga,e1,e2,n) = Ga.build('e1 e2 n',g=[1,1,1])
```

```
# assign some more commonly used variables
e12 = e1^e2 # projection plane
ZERO = e1|e2 # this is dumb
I = ga.I()
```

```
# projection of vector a onto Blade B
proj = lambda a,B:  (a|B)*B.inv()
# rejection  of vector a onto Blade B
rej  = lambda a,B: (a^B)*B.inv()

# up project a vector from  original space to projection space
up = lambda q: q+n
# down project a vector from  projection space to the original space
down = lambda p: n*(n^p)*(n|p).inv()

# normalize projection vector so that p_n = 1
norm = lambda p:  p*(p|n).inv()

# 'normalized projection vector': creates a projection vector with n = 1
npv = lambda s: proj(ga.mv(s,'vector'),e12)+n
```

### Verify Projection Functions

Given a projection vector

```
p = ga.mv('p','vector')
p
```

\begin{equation*} p = p^{1} e_{1} + p^{2} e_{2} + p^{3} n \end{equation*}

Ensure that projecting up, and then down gives us what we started with.

```
q = proj(p,e12)
q == down(up(q))
```

    True

###  Vectors  $ \rightarrow $ Points

The projection vector $p$ can be thought of as representing the point in the projection plane $q$. Vectors in projective space are homogeneous, meaning that two  projection vectors with different magnitudes represent the same point in the original space.

```
p
```

\begin{equation*} p = p^{1} e_{1} + p^{2} e_{2} + p^{3} n \end{equation*}

```
q = (p^n*(n|p).inv())*n
q

```

\begin{equation*} \frac{p^{1}}{p^{3}} e_{1} + \frac{p^{2}}{p^{3}} e_{2} \end{equation*}

So its clear that the magnitude of $p$ will cancel out in the fraction, hence its homogeneous.

### Points at $0$ and $\infty$

The values  of $0$ and $\infty$  are represented by projective vectors parallel and perpendicular to $n$.

** [ THESE CAUSE EXCEPTIONS] **

    AttributeError: 'int' object has no attribute 'is_commutative'

```
#down(proj(a,n)) # a vector parallel  to a
```

```
#down(rej(a, n)) # a vector perpendicular to n
```

### Bivectors $\rightarrow$ Vectors, or Lines

Recall that the relation between a vector in the projection space and orginal space is given by

$$
q  =  \underbrace{\frac{p\wedge n}{p\cdot n}}_{P} n =  Pn\\
$$

```
P = p^n*(n|p).inv()
P
```

\begin{equation*} \frac{p^{1}}{p^{3}} e_{1}\wedge n + \frac{p^{2}}{p^{3}} e_{2}\wedge n \end{equation*}

```
P*n
```

\begin{equation*} \frac{p^{1}}{p^{3}} e_{1} + \frac{p^{2}}{p^{3}} e_{2} \end{equation*}

```

```

Given that a  line in the original space is represented as a bivector in the projection space. The advantage of this is that the geometrical  entities ( like lines, planes, etc)  become objects which can be operated on like vectors, ie rotated, reflected, etc.

Below is an illustration (taken from [[3]](#References)), which shows two projection  vector $p_h$ and $q_h$, and their bivector $p_h\wedge q_h$

```
Image('pics/projective bivector.png')
```

Taking the outer product of two projection vectors. (Note i dropped the '$_h$' subscript because its a little too much ).

```
p, q = [ga.mv(k, 'vector') for k in ['p','q']]
p,q = npv('p'), npv('q') # would be nice to have a cleaner way to do this

Z = p^q
Z
```

\begin{equation*} \left ( p^{1} q^{2} - p^{2} q^{1}\right ) e_{1}\wedge e_{2} + \left ( p^{1} - q^{1}\right ) e_{1}\wedge n + \left ( p^{2} - q^{2}\right ) e_{2}\wedge n \end{equation*}

```

```

The coefficients  of this bivector are $\Delta_y$, $\Delta_x$ and $x_0\Delta-y_0\Delta_x$, which define the parameters of a line.

The dual of this is

```
z=Z.dual()
z
```

\begin{equation*} \left ( - p^{2} + q^{2}\right ) e_{1} + \left ( p^{1} - q^{1}\right ) e_{2} + \left ( - p^{1} q^{2} + p^{2} q^{1}\right ) n \end{equation*}

### Intersections (the *meet* operator)

```
# the meet of A and B
meet = lambda A,B:  A.dual()|B
```

To  demonstrate, we can show that meet'ing our basis bivectors together yields their vector of intersection

```
e12 = e1^e2
e2n = e2^n
en1 = n^e1
```

```
meet(e12, e2n)
```

\begin{equation*}  e_{2} \end{equation*}

```
meet(en1, e2n)
```

\begin{equation*} - n \end{equation*}

In projective space, a line is represented by joining two vectors. This is just the outerproduct.

```

```

```

```

This yields a nice visualization; a bivector in G3 will intersect the projection plane in a line, defined by the points $p$ and $q$.

### Numerical Example

Here is a numerical exmaples which demonstrates that the outer product in projective space determines a line in the original space. Given the two  points $p$ and $q$

```
p = 1*e1 + 2*e2
q = 3*e1 + 3*e2
```

The line $PQ$ is represented by $( p_h\wedge q_h)^*$

```
PQ =(up(p)^up(q)).dual()
PQ
```

\begin{equation*}  e_{1} -2 e_{2} + 3 n \end{equation*}

To see where two arbitrary line interset we can use  the meet

```
rv = lambda : randint(-10,10)*e1 + randint(-10,10)*e2
p,q,r,s = [rv() for k in range(4)]
```

## Scratch

```
p, q = [ga.mv(k, 'vector') for k in ['p','q']]
p,q = npv('p'), npv('q')

p_,q_ = down(p),down(q)
```

```
e0
```

```

```

```

```


# Single Experiment

In this notebook we run a single experiment and display the estimates of the dynamic effects based on our dynamic DML algorithm. We also display some performance of alternative benchmark approaches.

## 1. Data Generation from a Markovian Treatment Model

We consider the following DGP:

\begin{align}
    X_t =& (\pi'X_{t-1} + 1) \cdot A\, T_{t-1} + B X_{t-1} + \epsilon_t\\
    T_t =& \gamma\, T_{t-1} + (1-\gamma) \cdot D X_t + \zeta_t\\
    Y_t =& (\sigma' X_{t} + 1) \cdot e\, T_{t} + f' X_t + \eta_t
\end{align}

with $X_0, T_0 = 0$ and $\epsilon_t, \zeta_t, \eta_t$ normal $N(0, \sigma^2)$ r.v.'s. Moreover, $X_t \in R^{n_x}$, $B[:, 0:s_x] \neq 0$ and $B[:, s_x:-1] = 0$, $\gamma\in [0, 1]$, $D[:, 0:s_x] \neq 0$, $D[:, s_x:-1]=0$, $f[0:s_x]\neq 0$, $f[s_x:-1]=0$. We draw a single time series of samples of length $n\_samples$.

```python
%load_ext autoreload
%autoreload 2
```

    The autoreload extension is already loaded. To reload it, use:
      %reload_ext autoreload

```python
import numpy as np
from dynamic_panel_dgp import DynamicPanelDGP, LongRangeDynamicPanelDGP
n_units = 400
n_periods = 3
n_treatments = 1
n_x = 100
s_x = 10
s_t = 10
sigma_x = .5
sigma_t = .5
sigma_y = .5
gamma = .0
autoreg = .5
state_effect = .5
conf_str = 6
hetero_strength = 0
hetero_inds = None

#dgp_class = LongRangeDynamicPanelDGP
dgp_class = DynamicPanelDGP
dgp = dgp_class(n_periods, n_treatments, n_x).create_instance(s_x, sigma_x, sigma_y,
                                                              conf_str, hetero_strength, hetero_inds,
                                                              autoreg, state_effect,
                                                              random_seed=39)
```

    C:\ProgramData\Anaconda3\lib\site-packages\numba\errors.py:105: UserWarning: Insufficiently recent colorama version found. Numba requires colorama >= 0.3.9
      warnings.warn(msg)

```python
Y, T, X, groups = dgp.observational_data(n_units, gamma, s_t, sigma_t, random_seed=1234)
true_effect = dgp.true_effect
```

```python
import matplotlib.pyplot as plt
plt.figure(figsize=(20, 5))
plt.subplot(1, 2, 1)
plt.plot(Y, label="Outcome")
plt.plot(T, label="Treatment")
plt.legend()
plt.subplot(1, 2, 2)
for it in range(3):
    plt.plot(X[:, it], label="X[{}]".format(it))
plt.legend()
plt.show()
```

### 1. 1 True Parameters for Dynamic Effects for 3 periods

```python
import matplotlib.pyplot as plt

true_effect = true_effect.flatten()
plt.plot(true_effect, 'o')
plt.show()
```

## 2. Dynamic DML

```python
from sklearn.linear_model import LinearRegression, LassoCV, Lasso, MultiTaskLasso, MultiTaskLassoCV
from sklearn.model_selection import GroupKFold
import warnings
warnings.simplefilter('ignore')

np.random.seed(123)

alpha_regs = [1e-4, 1e-3, 1e-2, 5e-2, .1, 1]
lasso_model = lambda : LassoCV(cv=3, alphas=alpha_regs, max_iter=500)
mlasso_model = lambda : MultiTaskLassoCV(cv=3, alphas=alpha_regs, max_iter=500)
```

```python
from panel_dynamic_dml import DynamicPanelDML

est = DynamicPanelDML(model_t=mlasso_model(),
                      model_y=lasso_model(),
                      n_cfit_splits=3).fit(Y, T, X, groups)
```

### 2.1 Parameter Recovery and Confidence Intervals

```python
param_hat = est.param
conf_ints = est.param_interval(alpha=.05)
for kappa in range(n_periods):
    for t in range(n_treatments):
        param_ind = kappa*n_treatments + t
        print("Effect Lag={}, T={}: {:.3f} ({:.3f}, {:.6f}), (Truth={:.6f})".format(kappa, t,
                                                                                param_hat[param_ind],
                                                                                *conf_ints[param_ind],
                                                                                true_effect[param_ind]))
```

    Effect Lag=0, T=0: 0.777 (0.687, 0.868074), (Truth=0.866165)
    Effect Lag=1, T=0: 1.147 (0.919, 1.375321), (Truth=1.200000)
    Effect Lag=2, T=0: 0.896 (0.536, 1.257250), (Truth=0.641424)

```python
plt.figure(figsize=(15, 5))
plt.errorbar(np.arange(n_periods*n_treatments)-.04, param_hat, yerr=(conf_ints[:, 1] - param_hat,
                                                    param_hat - conf_ints[:, 0]), fmt='o', label='dyn-dml')
plt.errorbar(np.arange(n_periods*n_treatments), true_effect.flatten(), fmt='o', alpha=.6, label='true')
for t in np.arange(1, n_periods):
    plt.axvline(x=t * n_treatments - .5, linestyle='--', alpha=.4)
plt.xticks([t * n_treatments - .5 + n_treatments/2 for t in range(n_periods)],
           ["$\\theta_{}$".format(t) for t in range(n_periods)])
plt.gca().set_xlim([-.5, n_periods*n_treatments - .5])
plt.legend()
plt.show()
```

### 2.2 Benchmark Method Comparison

```python
panelX = X.reshape(-1, n_periods, n_x)
panelT = T.reshape(-1, n_periods, n_treatments)
panelY = Y.reshape(-1, n_periods)
```

#### 2.2.1 Regressing Y on all T

```python
est_lr = LinearRegression().fit(panelT[:, ::-1, :].reshape(-1, n_periods*n_treatments), panelY[:, -1]).coef_
```

#### 2.2.2 Regressing Y on all T and either final or initial States

```python
est_lr_x0 = lasso_model().fit(np.hstack([panelT[:, ::-1, :].reshape(-1, n_periods*n_treatments),
                                  panelX[:, 0, :]]), panelY[:, -1]).coef_[:n_periods*n_treatments]
est_lr_xfinal = lasso_model().fit(np.hstack([panelT[:, ::-1, :].reshape(-1, n_periods*n_treatments),
                                  panelX[:, -1, :]]), panelY[:, -1]).coef_[:n_periods*n_treatments]
```

#### 2.2.3 Performing DML with Y and all T and controlling for either final or initial States

```python
from econml.dml import LinearDMLCateEstimator
dml_model = lambda : LinearDMLCateEstimator(model_y=lasso_model(), model_t=mlasso_model(),
                                      n_splits=3, linear_first_stages=False)
est_dml_x0 = dml_model().fit(panelY[:, -1], T=panelT[:, ::-1, :].reshape(-1, n_periods*n_treatments),
                      X=None, W=panelX[:, 0, :]).intercept_
est_dml_xfinal = dml_model().fit(panelY[:, -1], T=panelT[:, ::-1, :].reshape(-1, n_periods*n_treatments),
                      X=None, W=panelX[:, -1, :]).intercept_
```

#### 2.2.4 Running a Direct version of Dynamic DML

Where a direct lasso is performed at each stage, regressing the calibrated outcome on the current period treatment and state and reading the coefficient in front of the treatment as the causal effect.

```python
Y_cal = panelY[:, -1].copy()
direct_theta = np.zeros((n_periods, n_treatments))
for t in np.arange(n_periods):
    direct_theta[t, :] = lasso_model().fit(np.hstack([panelT[:, n_periods - 1 - t, :],
                                                      panelX[:, n_periods - 1 - t, :]]), Y_cal).coef_[:n_treatments]
    Y_cal -= np.dot(panelT[:, n_periods - 1 - t, :], direct_theta[t, :])
est_direct = direct_theta.flatten()
```

#### 2.2.5 Plot all estimates

```python
plt.figure(figsize=(15, 5))
plt.errorbar(np.arange(n_periods*n_treatments)-.04, param_hat, yerr=(conf_ints[:, 1] - param_hat,
                                                    param_hat - conf_ints[:, 0]), fmt='o', label='dyn-dml')
plt.errorbar(np.arange(n_periods*n_treatments)-.02, est_lr, fmt='o', alpha=.6, label='no-ctrls')
plt.errorbar(np.arange(n_periods*n_treatments), true_effect.flatten(), fmt='*', alpha=.6, label='true')
plt.errorbar(np.arange(n_periods*n_treatments)+.02, est_lr_x0, fmt='o', alpha=.6, label='init-ctrls')
plt.errorbar(np.arange(n_periods*n_treatments)+.04, est_dml_x0, fmt='o', alpha=.6, label='init-ctrls-dml')
plt.errorbar(np.arange(n_periods*n_treatments)+.1, est_lr_xfinal, fmt='o', alpha=.6, label='fin-ctrls')
plt.errorbar(np.arange(n_periods*n_treatments)+.12, est_dml_xfinal, fmt='o', alpha=.6, label='fin-ctrls-dml')
plt.errorbar(np.arange(n_periods*n_treatments)+.14, est_direct, fmt='o', alpha=.6, label='dyn-direct')
for t in np.arange(1, n_periods):
    plt.axvline(x=t * n_treatments - .5, linestyle='--', alpha=.4)
plt.xticks([t * n_treatments - .5 + n_treatments/2 for t in range(n_periods)],
           ["$\\theta_{}$".format(t) for t in range(n_periods)])
plt.gca().set_xlim([-.5, n_periods*n_treatments - .5])
plt.legend()
plt.show()
```

```python
import matplotlib.pyplot as plt
import numpy as np
from scipy.stats import multivariate_normal

id1, id2 = 0, 1
length = 2 * max(est.param_stderr[id1], est.param_stderr[id2])
xlin = np.linspace(param_hat[id1]-length, param_hat[id1] + length, 500)
ylin = np.linspace(param_hat[id2]-length, param_hat[id2] + length, 500)
grX,grY = np.meshgrid(xlin, ylin)
pos = np.array([grX.flatten(), grY.flatten()]).T

rv = multivariate_normal(param_hat[[id1,id2]], est._cov[np.ix_([id1, id2], [id1, id2])]/n_units)

fig = plt.figure(figsize=(4,4))
ax0 = fig.add_subplot(111)
ax0.contourf(grX, grY, rv.pdf(pos).reshape(500,500))
ax0.scatter(true_effect[id1], true_effect[id2])
plt.xlabel("$\\theta_{{{}}}$".format(id1))
plt.ylabel("$\\theta_{{{}}}$".format(id2))
plt.show()
```

## 3. Policy Effect

```python
tau = np.random.binomial(1, .5, size=(n_periods, n_treatments))
true_policy_effect = dgp.static_policy_effect(tau, mc_samples=1000)
policy_effect_hat = est.policy_effect(tau)
policy_ints = est.policy_effect_interval(tau)
print("Policy effect for treatment seq: \n {}\n {:.3f} ({:.3f}, {:.3f}) (truth={:.3f})".format(tau,
                                                                                policy_effect_hat,
                                                                                *policy_ints,
                                                                                true_policy_effect))
```

    Policy effect for treatment seq:
     [[0]
     [1]
     [0]]
     1.147 (0.919, 1.375) (truth=1.200)

```python
test_policies = np.random.binomial(1, .5, size=(16, n_periods, n_treatments))

plt.figure(figsize=(15, 15))
for t, tau in enumerate(test_policies):
    true_policy_effect = np.dot(true_effect, tau[::-1].flatten())
    policy_effect_hat = est.policy_effect(tau)
    policy_ints = est.policy_effect_interval(tau)
    plt.subplot(4, 4, t + 1)
    plt.errorbar([t -.04], [policy_effect_hat], yerr=([policy_ints[1] - policy_effect_hat],
                                                  [policy_effect_hat - policy_ints[0]]), fmt='o', label='dyn-dml')
    plt.errorbar([t -.02], [np.dot(est_lr, tau[::-1].flatten())], fmt='o', alpha=.6, label='no-ctrls')
    plt.errorbar([t], [true_policy_effect], fmt='o', alpha=.6, label='true')
    plt.hlines([true_policy_effect], t - .06, t + .14, linestyles='--', alpha=.4)
    plt.errorbar([t + .02], [np.dot(est_lr_x0, tau[::-1].flatten())], fmt='o', alpha=.6, label='init-ctrls')
    plt.errorbar([t + .04], [np.dot(est_dml_x0, tau[::-1].flatten())], fmt='o', alpha=.6, label='init-ctrls-dml')
    plt.errorbar([t + .1], [np.dot(est_lr_xfinal, tau[::-1].flatten())], fmt='o', alpha=.6, label='fin-ctrls')
    plt.errorbar([t + .12], [np.dot(est_dml_xfinal, tau[::-1].flatten())], fmt='o', alpha=.6, label='fin-ctrls-dml')
    plt.errorbar([t +.14], [np.dot(est_direct, tau[::-1].flatten())], fmt='o', alpha=.6, label='dyn-direct')
    plt.title("{}".format(tau.flatten()))
#plt.legend()
plt.tight_layout()
plt.show()
```

```python
# Optimal Contextual Binary Treatment Policy
def adaptive_policy(t, x, period):
    return 1.*(dgp.hetero_effect_fn(n_periods - 1 - period, x) > 0)
dgp.adaptive_policy_effect(adaptive_policy)
```

    2.707589092135044

```python
est.adaptive_policy_effect(X, groups, adaptive_policy)
```

    (2.820942109125225, (2.547666040491572, 3.094218177758878))

## 4. Estimation Diagnostics

```python
import matplotlib.pyplot as plt
plt.plot(est.param_stderr)
plt.show()
```

```python
plt.imshow(est._M)
plt.colorbar()
plt.show()
```

```python
plt.imshow(est._Sigma)
plt.colorbar()
plt.show()
```

```python
plt.imshow(est._cov)
plt.colorbar()
plt.show()
```

```python

```


# **CS224W - Colab 3**

In Colab 2 we constructed GNN models by using PyTorch Geometric's built in GCN layer, `GCNConv`. In this Colab we will go a step deeper and implement the **GraphSAGE** ([Hamilton et al. (2017)](https://arxiv.org/abs/1706.02216)) layer directly. 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 3 :)

# 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
!pip install --upgrade torch-scatter
!pip install --upgrade torch-sparse
!pip install --upgrade torch-cluster
!pip install --upgrade torch-spline-conv
!pip install torch-geometric
```

    Requirement already satisfied: torch-scatter in /usr/local/lib/python3.7/dist-packages (2.0.9)
    Collecting torch-sparse
      Using cached torch_sparse-0.6.12.tar.gz (43 kB)
    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)
    Building wheels for collected packages: torch-sparse
      Building wheel for torch-sparse (setup.py) ... done
      Created wheel for torch-sparse: filename=torch_sparse-0.6.12-cp37-cp37m-linux_x86_64.whl size=1586647 sha256=3d96e5bd3d6d71ac62471a8170c244b3d46ea826f6010b5697ef02973839f238
      Stored in directory: /root/.cache/pip/wheels/fb/e2/2f/44956c61e3299573ffe12da9d1374c7576ca0c5fb1fe1ed38c
    Successfully built torch-sparse
    Installing collected packages: torch-sparse
    Successfully installed torch-sparse-0.6.12
    Collecting torch-cluster
      Using cached torch_cluster-1.5.9.tar.gz (38 kB)
    Building wheels for collected packages: torch-cluster
      Building wheel for torch-cluster (setup.py) ... done
      Created wheel for torch-cluster: filename=torch_cluster-1.5.9-cp37-cp37m-linux_x86_64.whl size=920366 sha256=417135ff33bb8fd1913ad4c7816f75c0b4ab85aedbc672bf21f11c5043e8046f
      Stored in directory: /root/.cache/pip/wheels/a6/60/d8/8bb27f58d8578ba8046f7ea0aadbae89a731db884a644ba361
    Successfully built torch-cluster
    Installing collected packages: torch-cluster
    Successfully installed torch-cluster-1.5.9
    Collecting torch-spline-conv
      Downloading torch_spline_conv-1.2.1.tar.gz (13 kB)
    Building wheels for collected packages: torch-spline-conv
      Building wheel for torch-spline-conv (setup.py) ... done
      Created wheel for torch-spline-conv: filename=torch_spline_conv-1.2.1-cp37-cp37m-linux_x86_64.whl size=400782 sha256=74398caa8e524a5677c77356a29ad956271a2b9c312d69847a6514311dd50326
      Stored in directory: /root/.cache/pip/wheels/9c/33/73/780370b7c7bdf2340c0a7b971e915643f14795b4caa7a9a31f
    Successfully built torch-spline-conv
    Installing collected packages: torch-spline-conv
    Successfully installed torch-spline-conv-1.2.1
    Collecting torch-geometric
      Downloading torch_geometric-2.0.3.tar.gz (370 kB)
         |████████████████████████████████| 370 kB 5.3 MB/s
    Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (1.19.5)
    Requirement already satisfied: tqdm in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (4.62.3)
    Requirement already satisfied: scipy in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (1.4.1)
    Requirement already satisfied: networkx in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (2.6.3)
    Requirement already satisfied: scikit-learn in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (1.0.1)
    Requirement already satisfied: requests in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (2.23.0)
    Requirement already satisfied: pandas in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (1.1.5)
    Collecting rdflib
      Downloading rdflib-6.1.1-py3-none-any.whl (482 kB)
         |████████████████████████████████| 482 kB 45.1 MB/s
    Requirement already satisfied: googledrivedownloader in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (0.4)
    Requirement already satisfied: jinja2 in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (2.11.3)
    Requirement already satisfied: pyparsing in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (3.0.6)
    Collecting yacs
      Downloading yacs-0.1.8-py3-none-any.whl (14 kB)
    Requirement already satisfied: PyYAML in /usr/local/lib/python3.7/dist-packages (from torch-geometric) (3.13)
    Requirement already satisfied: MarkupSafe>=0.23 in /usr/local/lib/python3.7/dist-packages (from jinja2->torch-geometric) (2.0.1)
    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: pytz>=2017.2 in /usr/local/lib/python3.7/dist-packages (from pandas->torch-geometric) (2018.9)
    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: setuptools in /usr/local/lib/python3.7/dist-packages (from rdflib->torch-geometric) (57.4.0)
    Collecting isodate
      Downloading isodate-0.6.1-py2.py3-none-any.whl (41 kB)
         |████████████████████████████████| 41 kB 597 kB/s
    Requirement already satisfied: importlib-metadata in /usr/local/lib/python3.7/dist-packages (from rdflib->torch-geometric) (4.8.2)
    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.6.0)
    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: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests->torch-geometric) (2021.10.8)
    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)
    Building wheels for collected packages: torch-geometric
      Building wheel for torch-geometric (setup.py) ... done
      Created wheel for torch-geometric: filename=torch_geometric-2.0.3-py3-none-any.whl size=581969 sha256=c349970a4e27ac0b332170bdba751ad78db1abd89458429ae7472cf7deb20901
      Stored in directory: /root/.cache/pip/wheels/c3/2a/58/87ce0508964d4def1aafb92750c4f3ac77038efd1b9a89dcf5
    Successfully built torch-geometric
    Installing collected packages: isodate, yacs, rdflib, torch-geometric
    Successfully installed isodate-0.6.1 rdflib-6.1.1 torch-geometric-2.0.3 yacs-0.1.8

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

```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 Colab 3, 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** (Colab 4) 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):
        super(GNNStack, self).__init__()
        conv_model = self.build_conv_model(args.model_type)
        self.convs = nn.ModuleList()
        self.convs.append(conv_model(input_dim, hidden_dim))
        assert (args.num_layers >= 1), 'Number of layers is not >=1'
        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))

        self.dropout = args.dropout
        self.num_layers = args.num_layers

        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):
        x, edge_index, batch = data.x, data.edge_index, data.batch

        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

## GraphSage Implementation

For our first GNN layer, we will implement the well known GraphSage ([Hamilton et al. (2017)](https://arxiv.org/abs/1706.02216)) layer!

For a given *central* node $v$ with current embedding $h_v^{l-1}$, the message passing update rule to tranform $h_v^{l-1} \rightarrow h_v^l$ is as follows:

\begin{equation}
h_v^{(l)} = W_l\cdot h_v^{(l-1)} + W_r \cdot AGG(\{h_u^{(l-1)}, \forall u \in N(v) \})
\end{equation}

where $W_1$ and $W_2$ are learanble weight matrices and the nodes $u$ are *neighboring* nodes. Additionally, we use mean aggregation for simplicity:

\begin{equation}
AGG(\{h_u^{(l-1)}, \forall u \in N(v) \}) = \frac{1}{|N(v)|} \sum_{u\in N(v)} h_u^{(l-1)}
\end{equation}

One thing to note is that we're adding a **skip connection** to our GraphSage implementation through the term $W_l\cdot h_v^{(l-1)}$.

Before implementing this update rule, we encourage you to think about how different parts of the formulas above correspond with the functions outlined earlier: 1) `forward`, 2) `message`, and 3) `aggregate`. As a hint, we are given what the aggregation function is (i.e. mean aggregation)! Now the question remains, what are the messages passed by each neighbor nodes and when do we call the `propagate` function?

Note: in this case the message function or messages are actually quite simple. Additionally, remember that the `propagate` function encapsulates the operations of / the outputs of the combined `message` and `aggregate` functions.

Lastly, $\ell$-2 normalization of the node embeddings is applied after each iteration.

<font color='red'>For the following questions, DON'T refer to any existing implementations online.</font>

```python
class GraphSage(MessagePassing):

    def __init__(self, in_channels, out_channels, normalize = True,
                 bias = False, heads=1, **kwargs):
        super(GraphSage, self).__init__(**kwargs)

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.normalize = normalize

        self.lin_l = None
        self.lin_r = None

        ############################################################################
        # TODO: Your code here!
        # Define the layers needed for the message and update functions below.
        # self.lin_l is the linear transformation that you apply to embedding
        #            for central node.
        # self.lin_r is the linear transformation that you apply to aggregated
        #            message from neighbors.
        # Our implementation is ~2 lines, but don't worry if you deviate from this.
        self.lin_l = nn.Linear(self.in_channels, self.out_channels)
        self.lin_r = nn.Linear(self.in_channels, self.out_channels)
        ############################################################################

        self.reset_parameters()

    def reset_parameters(self):
        self.lin_l.reset_parameters()
        self.lin_r.reset_parameters()

    def forward(self, x, edge_index, size = None):
        """"""

        out = None

        ############################################################################
        # TODO: Your code here!
        # Implement message passing, as well as any post-processing (our update rule).
        # 1. First call propagate function to conduct the message passing.
        #    1.1 See there for more information:
        #        https://pytorch-geometric.readthedocs.io/en/latest/notes/create_gnn.html
        #    1.2 We use the same representations for central (x_central) and
        #        neighbor (x_neighbor) nodes, which means you'll pass x=(x, x)
        #        to propagate.
        # 2. Update our node embedding with skip connection.
        # 3. If normalize is set, do L-2 normalization (defined in
        #    torch.nn.functional)
        # Our implementation is ~5 lines, but don't worry if you deviate from this.
        x_prop = self.propagate(edge_index, x=(x,x), size=size)
        x = self.lin_l(x) + x_prop
        if self.normalize:
            out = F.normalize(x)
        ############################################################################

        return out

    def message(self, x_j):

        out = None

        ############################################################################
        # TODO: Your code here!
        # Implement your message function here.
        # Our implementation is ~1 lines, but don't worry if you deviate from this.
        out = self.lin_r(x_j)
        ############################################################################

        return out

    def aggregate(self, inputs, index, dim_size = None):

        out = None

        # The axis along which to index number of nodes.
        node_dim = self.node_dim

        ############################################################################
        # 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/functions/scatter.html#torch_scatter.scatter
        # Our implementation is ~1 lines, but don't worry if you deviate from this.
        out = torch_scatter.scatter(inputs, index, dim=node_dim, reduce="mean")
        ############################################################################

        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
import os
if 'IS_GRADESCOPE_ENV' not in os.environ:
    for args in [
        {'model_type': 'GraphSage', '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 ['GraphSage']:
            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.1: What is the maximum accuracy obtained on the test set for GraphSage? (10 points)

Running the cell above will show the results of your best model and save your best model's predictions to a file named *CORA-Node-GraphSage.csv*.

As we have seen before you can view this file by clicking on the *Folder* icon on the left side pannel. When you sumbit your assignment, you will have to download this file and attatch it to your submission.