Blankenbach Benchmark Case 1
======

Steady isoviscous thermal convection
----

Two-dimensional, incompressible, bottom heated, steady isoviscous thermal convection in a 1 x 1 box, see case 1 of Blankenbach *et al.* 1989 for details.

**This example introduces:**
1. Loading/Saving variables to disk.
2. Defining analysis tools.
3. Finding a steady state.

**Keywords:** Stokes system, advective diffusive systems, analysis tools

**References**

B. Blankenbach, F. Busse, U. Christensen, L. Cserepes, D. Gunkel, U. Hansen, H. Harder, G. Jarvis, M. Koch, G. Marquart, D. Moore, P. Olson, H. Schmeling and T. Schnaubelt. A benchmark comparison for mantle convection codes. Geophysical Journal International, 98, 1, 23–38, 1989
http://onlinelibrary.wiley.com/doi/10.1111/j.1365-246X.1989.tb05511.x/abstract

```python
import underworld as uw
from underworld import function as fn
import underworld.visualisation as vis
import math
import numpy as np
```

Setup parameters
-----

```python
boxHeight = 1.0
boxLength = 1.0
# Set grid resolution.
res = 128
# Set max & min temperautres
tempMin = 0.0
tempMax = 1.0
```

Choose which Rayleigh number, see case 1 of Blankenbach *et al.* 1989 for details.

```python
case = "a"
if(case=="a"):
    Ra=1.e4
    eta0=1.e23
elif(case=="b"):
    Ra=1.e5
    eta0=1.e22
else:
    Ra=1.e6
    eta0=1.e21
```

Set input and output file directory

```python
inputPath  = 'input/03_BlankenbachBenchmark/'
outputPath = 'output/'
# Make output directory if necessary.
if uw.mpi.rank==0:
    import os
    if not os.path.exists(outputPath):
        os.makedirs(outputPath)
```

Create mesh and variables
------

```python
mesh = uw.mesh.FeMesh_Cartesian( elementType = ("Q1/dQ0"),
                                 elementRes  = (res, res),
                                 minCoord    = (0., 0.),
                                 maxCoord    = (boxLength, boxHeight))

velocityField       = mesh.add_variable(         nodeDofCount=2 )
pressureField       = mesh.subMesh.add_variable( nodeDofCount=1 )
temperatureField    = mesh.add_variable(         nodeDofCount=1 )
temperatureDotField = mesh.add_variable(         nodeDofCount=1 )

# initialise velocity, pressure and temperatureDot field
velocityField.data[:]       = [0.,0.]
pressureField.data[:]       = 0.
temperatureField.data[:]    = 0.
temperatureDotField.data[:] = 0.
```

Set up material parameters and functions
-----

Set values and functions for viscosity, density and buoyancy force.

```python
# Set a constant viscosity.
viscosity = 1.

# Create our density function.
densityFn = Ra * temperatureField

# Define our vertical unit vector using a python tuple (this will be automatically converted to a function).
z_hat = ( 0.0, 1.0 )

# A buoyancy function.
buoyancyFn = densityFn * z_hat
```

Set initial temperature field
-----
The initial temperature field can be loaded from a pre-run steady state data set ( ``LoadFromFile = True`` ) or set to a sinusodial perterbation ( ``LoadFromFile = False`` ).

```python
# Steady state temperature field to be loaded from data file.
LoadFromFile = True
```

**If loading steady state data set**

Data is stored in h5 format from a 64\*64 grid resolution model.  Data has been saved for 3 different Rayleigh numbers, $Ra = 10^4$, $10^5$ or $10^6$.

Once loaded the data will need to be re-meshed onto a new grid, unless the new resolution is also 64\*64.

For more information on using meshes see the user guide.

```python
if(LoadFromFile == True):
    # Setup mesh and temperature field for 64*64 data file.
    mesh64 = uw.mesh.FeMesh_Cartesian( elementType = ("Q1/dQ0"),
                                       elementRes  = (64, 64),
                                       minCoord    = (0., 0.),
                                       maxCoord    = (boxLength, boxHeight),
                                       partitioned = False )
    temperatureField64  = mesh64.add_variable( nodeDofCount=1 )

    # read in saved steady state temperature field data
    if( case == "a" ):
        temperatureField64.load(inputPath+'tempfield_inp_64_Ra1e4.h5')
    elif( case == "b" ):
        temperatureField64.load(inputPath+'tempfield_inp_64_Ra1e5.h5')
    else:
        temperatureField64.load(inputPath+'tempfield_inp_64_Ra1e6.h5')

    if( res==64 ): # no remeshing needed, copy directly
        temperatureField.data[:] = temperatureField64.data[:]
    else: # remeshing needed
        temperatureField.data[:] = temperatureField64.evaluate(mesh)
```

**If using sinusodial perturbation**

```python
if(LoadFromFile == False):
    temperatureField.data[:] = 0.
    pertStrength = 0.1
    deltaTemp = tempMax - tempMin
    for index, coord in enumerate(mesh.data):
        pertCoeff = math.cos( math.pi * coord[0]/boxLength ) * math.sin( math.pi * coord[1]/boxLength )
        temperatureField.data[index] = tempMin + deltaTemp*(boxHeight - coord[1]) + pertStrength * pertCoeff
        temperatureField.data[index] = max(tempMin, min(tempMax, temperatureField.data[index]))
```

**Show initial temperature field**

```python
fig = vis.Figure()
fig.append( vis.objects.Surface(mesh, temperatureField) )
fig.show()
```

Create boundary conditions
----------

Set temperature boundary conditions on the bottom ( ``MinJ`` ) and top ( ``MaxJ`` ).

```python
for index in mesh.specialSets["MinJ_VertexSet"]:
    temperatureField.data[index] = tempMax
for index in mesh.specialSets["MaxJ_VertexSet"]:
    temperatureField.data[index] = tempMin
```

Construct sets for the both horizontal and vertical walls. Combine the sets of vertices to make the ``I`` (left and right side walls) and ``J`` (top and bottom walls) sets.

```python
iWalls = mesh.specialSets["MinI_VertexSet"] + mesh.specialSets["MaxI_VertexSet"]
jWalls = mesh.specialSets["MinJ_VertexSet"] + mesh.specialSets["MaxJ_VertexSet"]

freeslipBC = uw.conditions.DirichletCondition( variable        = velocityField,
                                               indexSetsPerDof = (iWalls, jWalls) )
tempBC     = uw.conditions.DirichletCondition( variable        = temperatureField,
                                               indexSetsPerDof = (jWalls,) )
```

System setup
-----

**Setup a Stokes system**

```python
stokes = uw.systems.Stokes(    velocityField = velocityField,
                               pressureField = pressureField,
                               conditions    = [freeslipBC,],
                               fn_viscosity  = viscosity,
                               fn_bodyforce  = buoyancyFn )
# get the default stokes equation solver
solver = uw.systems.Solver( stokes )
```

**Create an advection diffusion system**

```python
advDiff = uw.systems.AdvectionDiffusion( phiField       = temperatureField,
                                         phiDotField    = temperatureDotField,
                                         velocityField  = velocityField,
                                         fn_diffusivity = 1.0,
                                         conditions     = [tempBC,] )
```

Analysis tools
-----

**Nusselt number**

The Nusselt number is the ratio between convective and conductive heat transfer

\\[
Nu = -h \frac{ \int_0^l \partial_z T (x, z=h) dx}{ \int_0^l T (x, z=0) dx}
\\]

```python
nuTop    = uw.utils.Integral( fn=temperatureField.fn_gradient[1],
                              mesh=mesh, integrationType='Surface',
                              surfaceIndexSet=mesh.specialSets["MaxJ_VertexSet"])

nuBottom = uw.utils.Integral( fn=temperatureField,
                              mesh=mesh, integrationType='Surface',
                              surfaceIndexSet=mesh.specialSets["MinJ_VertexSet"])
```

```python
nu = - nuTop.evaluate()[0]/nuBottom.evaluate()[0]
if uw.mpi.rank == 0 : print('Nusselt number = {0:.6f}'.format(nu))
```

    Nusselt number = 4.856902

**RMS velocity**

The root mean squared velocity is defined by intergrating over the entire simulation domain via

\\[
\begin{aligned}
v_{rms}  =  \sqrt{ \frac{ \int_V (\mathbf{v}.\mathbf{v}) dV } {\int_V dV} }
\end{aligned}
\\]

where $V$ denotes the volume of the box.

```python
intVdotV = uw.utils.Integral( fn.math.dot( velocityField, velocityField ), mesh )

vrms = math.sqrt( intVdotV.evaluate()[0] )
if uw.mpi.rank == 0 : print('Initial vrms = {0:.3f}'.format(vrms))
```

    Initial vrms = 0.000

Main simulation loop
-----

If the initial conditions are loaded from file then this loop will only take a single step. If you would like to run the entire simulation from a small perturbation then change the ``LoadFromFile`` variable above to equal ``False``. Warning: the simulation will take a long time to get to steady state.

```python
#initialise time, step, output arrays
time = 0.
step = 0
timeVal = []
vrmsVal = []

# starting from steady state == True
if(LoadFromFile == True):
    step_end = 1
else:
    step_end = 5000

# output frequency
step_output = max(1,min(100, step_end/10))
epsilon    = 1.e-8

velplotmax = 0.0
nuLast     = -1.0
```

```python
# define an update function
def update():
    # Determining the maximum timestep for advancing the a-d system.
    dt = advDiff.get_max_dt()
    # Advect using this timestep size.
    advDiff.integrate(dt)
    return time+dt, step+1
```

```python
# Perform steps.
while step<=step_end:
    # Solving the Stokes system.
    solver.solve()

    # Calculate & store the RMS velocity and Nusselt number.
    vrms = math.sqrt( intVdotV.evaluate()[0] )
    nu   = - nuTop.evaluate()[0]/nuBottom.evaluate()[0]
    vrmsVal.append(vrms)
    timeVal.append(time)
    velplotmax = max(vrms, velplotmax)

    # print output statistics
    if step%(step_end/step_output) == 0:
        if(uw.mpi.rank==0):
            print('steps = {0:6d}; time = {1:.3e}; v_rms = {2:.3f}; Nu = {3:.3f}; Rel change = {4:.3e}'
              .format(step, time, vrms, nu, abs((nu - nuLast)/nu)))
    # Check loop break conditions.
    if(abs((nu - nuLast)/nu) < epsilon):
        if(uw.mpi.rank==0):
            print('steps = {0:6d}; time = {1:.3e}; v_rms = {2:.3f}; Nu = {3:.3f}; Rel change = {4:.3e}'
              .format(step, time, vrms, nu, abs((nu - nuLast)/nu)))
        break
    nuLast = nu

    # update
    time, step = update()
```

    steps =      0; time = 0.000e+00; v_rms = 42.869; Nu = 4.857; Rel change = 1.206e+00
    steps =      1; time = 3.052e-05; v_rms = 42.869; Nu = 4.864; Rel change = 1.416e-03

Post analysis
-----

**Benchmark values**

The time loop above outputs $v_{rms}$ and $Nu$ as general statistics for the system. For comparison, the benchmark values for the RMS velocity and Nusselt number are shown below for different Rayleigh numbers. All benchmark values shown below were determined in Blankenbach *et al.* 1989 by extroplation of numerical results.

| $Ra$          | $v_{rms}$   | $Nu$  |   $q_1$  | $q_2$  |
| ------------- |:-------------:|:-----:|:-----:|:-----:|
| 10$^4$ | 42.865 |  4.884 | 8.059 | 0.589 |
| 10$^5$ | 193.215 | 10.535 | 19.079 | 0.723 |
| 10$^6$ | 833.990 | 21.972 | 45.964 | 0.877 |

```python
# Let's add a test to ensure things are working as expected
if case == "a":
    if not np.isclose(nu,4.884,rtol=1.e-2):
        raise RuntimeError("Model did not produce the expected Nusselt number.")
    if not np.isclose(vrms,42.865,rtol=1.e-2):
        raise RuntimeError("Model did not produce the expected Nusselt number.")
```

**Resulting pressure field**

Use the same method as above to plot the new temperature field. This can also be used to plot the pressure field, or any other data structures of interest.

```python
figtemp = vis.Figure()
figtemp.append( vis.objects.Surface( mesh, pressureField ) )
figtemp.show()
```

**Plot the velocity vector field**

For this example the velocity field is interesting to see. This is visualised in two ways, firstly plotting a surface colour map of the velocity magnitude, and secondly the velocity vectors at points on the mesh. For aesthetics the vector arrows are scaled by a little more than the maximum $v_{rms}$ value found in the time loop above.

```python
fig2 = vis.Figure()
velmagfield = uw.function.math.sqrt( uw.function.math.dot( velocityField, velocityField ) )
fig2.append( vis.objects.VectorArrows(mesh, velocityField/(2.5*velplotmax), arrowHead=0.2, scaling=0.1) )
fig2.append( vis.objects.Surface(mesh, temperatureField) )
fig2.show()
```

Analysis
----

The final benchmarks in the Blankenbach paper involve the temperature gradient in the vertical direction ($\frac{\partial T}{\partial z}$). This is easy to find using the underworld functions, as shown below.

```python
tempgradField = temperatureField.fn_gradient
vertTGradField = - boxHeight * tempgradField[1] / tempMax # scaled for direct benchmarking below
```

**More benchmark values**

The vertical temperature gradient (above) is set up to be non-dimensional as per Blankenbach et al 1989. To compare to the benchmark values in their work the gradient is compared at the corners of the simulation box: $q_1$ at $x = 0$, $z = h$; $q_2$ at $x = l$, $z = h$; $q_3$ at $x = l$, $z = 0$; $q_4$ at $x = 0$, $z = 0$. Where $h$ = Box_Height and $l$ = Box_Length and the non-dimensional gradient field is given by
\\[
q = \frac{-h}{\Delta T} \left( \frac{\partial T}{\partial z} \right)
\\]

Provided the simulation is run to steady-state with sufficent resolution then the $q$ values should be close to the benchmark values given again below for different Rayleigh numbers.

| $Ra$          |    $q_1$  | $q_2$  |
| ------------- |:-----:|:-----:|
| 10$^4$ | 8.059 | 0.589 |
| 10$^5$ | 19.079 | 0.723 |
| 10$^6$ | 45.964 | 0.877 |

```python
q1 = vertTGradField.evaluate_global( (       0., boxHeight) )
q2 = vertTGradField.evaluate_global( (boxLength, boxHeight) )
q3 = vertTGradField.evaluate_global( (boxLength,        0.) )
q4 = vertTGradField.evaluate_global( (       0.,        0.) )

if(uw.mpi.rank==0):
    print('Rayleigh number = {0:.1e}'.format(Ra))
    print('q1 = {0:.3f}; q2 = {1:.3f}'.format(q1[0][0], q2[0][0]))
    print('q3 = {0:.3f}; q4 = {1:.3f}'.format(q3[0][0], q4[0][0]))

    # Let's add a test to ensure things are working as expected
    if case == "a":
        if not np.isclose(q1[0][0],8.020,rtol=1.e-2):
            raise RuntimeError("Model did not produce the expected q1.")
        if not np.isclose(q2[0][0],0.589,rtol=1.e-2):
            raise RuntimeError("Model did not produce the expected q2.")
```

    Rayleigh number = 1.0e+04
    q1 = 8.021; q2 = 0.589
    q3 = 8.021; q4 = 0.589

**Save time and rms values**

```python
if(uw.mpi.rank==0):
    np.savetxt(outputPath+'vrms.txt', np.c_[timeVal, vrmsVal], header="Time, VRMS" )
```

**Calculate stress values for benchmark comparison**

Determine stress field for whole box in dimensionless units (King 2009)
\begin{equation}
\tau_{ij} = \eta \frac{1}{2} \left[ \frac{\partial v_j}{\partial x_i} + \frac{\partial v_i}{\partial x_j}\right]
\end{equation}
which for vertical normal stress becomes
\begin{equation}
\tau_{zz} = \eta \frac{1}{2} \left[ \frac{\partial v_z}{\partial z} + \frac{\partial v_z}{\partial z}\right] = \eta \frac{\partial v_z}{\partial z}
\end{equation}
which is implemented for the whole box in the functions defined below.

```python
# get topography
stresstensorFn = 2.* stokes.fn_viscosity*fn.tensor.symmetric( velocityField.fn_gradient ) - (1.,1.,0.)*pressureField
verticalStressFn  = stresstensorFn[1]

# subtract the average value for benchmark
top_verts = mesh.specialSets["MaxJ_VertexSet"]
mean_sigma_zz_top = uw.utils.Integral(-verticalStressFn,
                                      mesh,
                                      'surface',
                                      top_verts).evaluate()/boxLength
sigma_zz_top = -verticalStressFn - mean_sigma_zz_top
```

Dimensionalise the stress from the vertical normal stress at the top of the box (King 2009)

$$
    \sigma_{t} = \frac{\eta_0 \kappa}{\rho g h^2}\tau _{zz} \left( x, z=h\right)
$$

where all constants have been defined above. Finally calculate the topography, defined using $h = \sigma_{top} / (\rho g)$.

```python
# Set parameters in SI units
grav = 10        # m.s^-2
height = 1.e6    # m
rho  = 4.0e3     # g.m^-3
kappa  = 1.0e-6  # m^2.s^-1

# dimensionalise
dim_sigma_zz_top  = (eta0 * kappa / (height*height)) * sigma_zz_top

# find topography in [m]
topography        = dim_sigma_zz_top / (rho * grav)
```

**Calculate x-coordinate at zero stress**

Calculate the zero point for the stress along the x-axis at the top of the box.

```python
# First construct abs function
abs_topo = fn.math.abs(topography)
# Now create a `min_max` view function.
# These functions don't modify results, they simply observe
# calculated results as functions are being evaluated.
# We will use this to watch out for the minimum value as
# we traverse function evaluates.
# We also set `fn_auxiliary` which allows us to also capture
# the coordinate of the min value.
abs_topo_min_max = fn.view.min_max(abs_topo,fn_auxiliary=fn.coord())   # create min_max
ignore = abs_topo_min_max.evaluate(mesh.specialSets["MaxJ_VertexSet"]) # traverse function evaluations
min_abs_topo_coord = abs_topo_min_max.min_global_auxiliary()[0][0]     # get x coord of global min
```

**Topography comparison**

Topography of the top boundary calculated in the left and right corners as given in Table 9 of Blankenbach et al 1989.

| $Ra$          |    $\xi_1$  | $\xi_2$  |  $x$ ($\xi = 0$) |
| ------------- |:-----------:|:--------:|:--------------:|
| 10$^4$  | 2254.02   | -2903.23  | 0.539372          |
| 10$^5$  | 1460.99   | -2004.20  | 0.529330          |
| 10$^6$  | 931.96   | -1283.80  | 0.506490          |

```python
e1 = topography.evaluate_global( (       0.,boxHeight) )
e2 = topography.evaluate_global( (boxLength,boxHeight) )

if(uw.mpi.rank==0):
    e1 = e1[0][0]  # pull values out of list of results
    e2 = e2[0][0]  # pull values out of list of results
    print('Rayleigh number = {0:.1e}'.format(Ra))
    print('Topography[x=0],[x=max] = {0:.2f}, {1:.2f}'.format(e1, e2))
    print('x(topo=0) = {0:.6f}'.format(min_abs_topo_coord))
    # output a summary file with benchmark values (useful for parallel runs)
    np.savetxt(outputPath+'summary.txt', [Ra, e1, e2, min_abs_topo_coord, q1, q2, q3, q4])

    # Let's add a test to ensure things are working as expected
    if case == "a":
        if not np.isclose(e1,2254.02,rtol=1.e-3):
            raise RuntimeError("Model did not produce the expected xi1.")
        if not np.isclose(e2,-2903.23,rtol=1.e-3):
            raise RuntimeError("Model did not produce the expected xi2.")
        if not np.isclose(min_abs_topo_coord,0.539372,rtol=1.e-3):
            raise RuntimeError("Model did not produce the expected xi2.")
```

    Rayleigh number = 1.0e+04
    Topography[x=0],[x=max] = 2252.36, -2900.67
    x(topo=0) = 0.539062


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

import torchsde
import math
import matplotlib.pyplot as plt

from cfollmer.objectives import log_g, relative_entropy_control_cost
from cfollmer.sampler_utils import FollmerSDE
import numpy as np

from tqdm.notebook import tqdm

from torch import _vmap_internals
```

## Generating Toy 2D Dataset.

```python
import torch
from sklearn.datasets import make_blobs
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

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

X, y = make_blobs(n_samples=300,
                  cluster_std=[0.5, 0.5, 0.5],
                  random_state=random_state)

X = X[y==0,:]
y = y[y==0]
# # Scale data to have mean 0 and variance 1
# # which is importance for convergence of the neural network
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

# Split the data set into training and testing
X_train, X_test, y_train, y_test = train_test_split(
    X_scaled, y, test_size=0.2, random_state=2)

X_train, X_test, y_train, y_test = \
    torch.tensor(X_train, dtype=torch.float32, device=device), \
    torch.tensor(X_test, dtype=torch.float32, device=device), \
    torch.tensor(y_train, dtype=torch.float32, device=device), \
    torch.tensor(y_test, dtype=torch.float32, device=device)
```

```python
y
```

    array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
           0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])

```python
fig, ax1 = plt.subplots(1, 1, figsize=(16, 6))

ax1.plot(X_scaled[:, 0], X_scaled[:, 1],
         linestyle='none',
         marker='o')

ax1.axis('equal');
```

$$\DeclareMathOperator*{\argmin}{arg\,min}$$
$$\def\E{{\mathbb{E}}}$$
$$\def\rvu{{\mathbf{u}}}$$
$$\def\rvTheta{{\bm{\Theta}}}$$
$$\def\gU{{\mathcal{U}}}$$
$$\def\mX{{\mathbf{X}}}$$

## Controlled Schrodinger Follmer Sampler

The objevtive we are trying to implement is:

\begin{align}
  \mathbf{u}_t^{*}=  \argmin_{\rvu_t \in \mathcal{U}}\mathbb{E}\left[\frac{1}{2\gamma}\int_0^1||\rvu(t, \Theta_t)||^2 dt - \ln\left(\frac{ p(\mX | \Theta_1)p(\Theta_1)}{\mathcal{N}(\Theta_1|\mathbf{0}, \gamma \mathbb{I} )}\right)\right] \
\end{align}

Where:
\begin{align}
d\Theta_t = \rvu(t, \Theta_t)dt + \sqrt{\gamma} dB_t
\end{align}

To do so we use the EM discretisation.

```python
def log_gaussian(x, mean=0):
    """
    Returns the density of x under the supplied gaussian. Defaults to
    standard gaussian N(0, I)
    :param x: (*) torch.Tensor
    :param mean: float or torch.FloatTensor with dimensions (*)
    :param logvar: float or torch.FloatTensor with dimensions (*)
    :return: (*) elementwise log density
    """

    log_norm_constant = -0.5 * np.log(2 * np.pi)

    var =  torch.tensor(0.1)
    logvar = torch.log(var).reshape(1,1).to(device)

    A = (x - mean) ** 2
    log_p = -0.5 * (logvar + A / logvar.exp())
    log_p = log_p + log_norm_constant
#     import pdb; pdb.set_trace()
    return log_p.sum(dim=-1)

def ln_prior(Θ, σ_w=1):
    """
    Prior for means in Bayesian GMM
    """
    return -0.5 * (Θ**2).sum(axis=1) / σ_w

def log_likelihood_single(μ, X, log=True):
    """
    :param X: design matrix (examples, features)
    :param mu: the component means (K, features)
    :param logvar: the component log-variances (K, features)
    :param log: return value in log domain?

    """

    n, d = X.shape
    # get feature-wise log-likelihoods (K, examples, features)
    log_likelihoods = log_gaussian(
        X[None, :, :],
        μ.reshape(1, d)[:, None, :],
    )

    return log_likelihoods.sum()

def log_likelihood(Θ, X, y=None):
    """
    batching the above (hopefully its right)
    """

    loss_ = lambda μ: log_likelihood_single(μ, X)

    batched_loss =  torch._vmap_internals.vmap(loss_)

    return batched_loss(Θ)
```

```python
Δt=0.05
t_size = int(math.ceil(1.0/Δt))
dim =  2

ts = torch.linspace(0, 1, t_size).to(device)
no_posterior_samples = 50

sde = FollmerSDE(dim, dim, no_posterior_samples, 1.0, device=device).to(device)
Θ_0 = torch.zeros((no_posterior_samples, dim)).to(device) # Θ_0 ~ δ_0

# Initial state y0, the SDE is solved over the interval [ts[0], ts[-1]].
# ys will have shape (t_size, batch_size, state_size)
ys = torchsde.sdeint(sde, Θ_0, ts, dt=Δt)
```

```python
relative_entropy_control_cost(sde, Θ_0, X_train, y_train,
                              ln_prior, log_likelihood, γ=1.0, device=device)
```

    C:\Users\vargf\anaconda3\lib\site-packages\torch\_vmap_internals.py:252: UserWarning: torch.vmap is an experimental prototype that is subject to change and/or deletion. Please use at your own risk.
      warnings.warn(
    <ipython-input-5-bd9d1ba955c7>:47: UserWarning: Batching rule not implemented for aten::sum falling back to slow (for loop and stack) implementation (Triggered internally at  ..\aten\src\ATen\BatchedFallback.cpp:63.)
      return log_likelihoods.sum()

    tensor(1546.1788, device='cuda:0', grad_fn=<MeanBackward0>)

```python
γ = 1.0
Δt=0.05
t_size = int(math.ceil(1.0/Δt))
print(t_size)
ts = torch.linspace(0, 1, t_size).to(device)

sde = FollmerSDE(dim, dim, no_posterior_samples  , γ=γ, device=device).to(device)
optimizer = torch.optim.Adam(sde.μ.parameters(), lr=0.01, weight_decay =1)
#     optimizer = torch.optim.LBFGS(gpr.parameters(), lr=0.01)
losses = []
num_steps = 300

train = True
if train:
    for i in tqdm(range(num_steps)):
        optimizer.zero_grad()

        if isinstance(optimizer, torch.optim.LBFGS):
            def closure():
                loss = relative_entropy_control_cost(
                    sde, Θ_0.float(),
                    X_train.float(), y_train.float(),
                    ln_prior, log_likelihood, γ=γ
                )
                optimizer.zero_grad()
                loss.backward()
                return loss

            optimizer.step(closure)
            losses.append(closure().item())
        else:
            loss = relative_entropy_control_cost(
                sde, Θ_0,
                X_train, y_train,
                ln_prior, log_likelihood, γ=γ
            )
            optimizer.zero_grad()
            loss.backward()

            optimizer.step()
            losses.append(loss.item())
```

    20

      0%|          | 0/300 [00:00<?, ?it/s]

    C:\Users\vargf\anaconda3\lib\site-packages\torch\_vmap_internals.py:252: UserWarning: torch.vmap is an experimental prototype that is subject to change and/or deletion. Please use at your own risk.
      warnings.warn(
    <ipython-input-5-bd9d1ba955c7>:47: UserWarning: Batching rule not implemented for aten::sum falling back to slow (for loop and stack) implementation (Triggered internally at  ..\aten\src\ATen\BatchedFallback.cpp:63.)
      return log_likelihoods.sum()

```python
losses
```

    [1451.687744140625,
     1518.5587158203125,
     1454.130615234375,
     1447.2540283203125,
     1291.2786865234375,
     1316.039306640625,
     1184.6658935546875,
     1233.3717041015625,
     1317.087158203125,
     1158.0584716796875,
     1094.017333984375,
     1094.5360107421875,
     1084.8631591796875,
     1022.4356079101562,
     1071.853759765625,
     1036.6337890625,
     981.9197998046875,
     986.9926147460938,
     970.450439453125,
     934.8773193359375,
     925.0805053710938,
     907.3358154296875,
     918.404052734375,
     907.4785766601562,
     901.797119140625,
     931.2405395507812,
     895.2091064453125,
     883.7435913085938,
     897.390625,
     884.6437377929688,
     893.3555297851562,
     865.6217651367188,
     904.397705078125,
     882.1803588867188,
     914.0098266601562,
     901.625,
     894.4642944335938,
     898.7965698242188,
     887.6434936523438,
     905.745849609375,
     878.3790283203125,
     899.3026123046875,
     887.8787841796875,
     912.1152954101562,
     894.7665405273438,
     890.3003540039062,
     912.908203125,
     917.9296875,
     900.7048950195312,
     892.1951904296875,
     887.3540649414062,
     870.8582763671875,
     887.481689453125,
     901.2720336914062,
     893.16357421875,
     867.3097534179688,
     910.703369140625,
     905.1739501953125,
     896.3814086914062,
     902.2969360351562,
     913.9576416015625,
     887.4429321289062,
     886.411865234375,
     878.4768676757812,
     883.8621826171875,
     882.857177734375,
     893.8919067382812,
     875.2954711914062,
     891.3561401367188,
     872.8623046875,
     902.867431640625,
     882.5993041992188,
     864.7813720703125,
     907.2366943359375,
     893.711181640625,
     904.2862548828125,
     920.5433959960938,
     881.212158203125,
     898.2505493164062,
     881.405029296875,
     906.5316772460938,
     876.1931762695312,
     912.2235107421875,
     919.2583618164062,
     881.1099853515625,
     884.4055786132812,
     909.7791137695312,
     910.9274291992188,
     895.3540649414062,
     909.9713745117188,
     899.049072265625,
     878.7960205078125,
     921.9085083007812,
     895.3351440429688,
     865.0919189453125,
     874.3062133789062,
     881.4456176757812,
     872.3929443359375,
     886.61669921875,
     881.1976928710938,
     867.465087890625,
     888.9425048828125,
     876.0540161132812,
     883.9166870117188,
     888.1195678710938,
     914.775634765625,
     886.3356323242188,
     891.9719848632812,
     885.8046875,
     865.2760620117188,
     878.1293334960938,
     892.6008911132812,
     904.58984375,
     894.122802734375,
     906.2154541015625,
     905.689453125,
     874.1265258789062,
     872.7413940429688,
     866.2747192382812,
     918.149169921875,
     872.7871704101562,
     889.385009765625,
     911.3499145507812,
     889.5477905273438,
     893.0620727539062,
     893.3655395507812,
     880.0462036132812,
     870.8622436523438,
     865.2344970703125,
     860.7862548828125,
     883.2098999023438,
     899.572998046875,
     896.296142578125,
     893.7417602539062,
     892.6500854492188,
     869.04296875,
     887.6395263671875,
     883.0781860351562,
     897.3348999023438,
     879.8421630859375,
     890.9437255859375,
     879.3034057617188,
     894.408447265625,
     882.348876953125,
     886.04638671875,
     893.4269409179688,
     901.9030151367188,
     893.0975341796875,
     885.6099243164062,
     911.9644165039062,
     885.121826171875,
     881.714599609375,
     890.2120971679688,
     909.7028198242188,
     879.0960083007812,
     883.9276123046875,
     901.1802368164062,
     894.373046875,
     895.1536254882812,
     889.4401245117188,
     889.339111328125,
     868.6571655273438,
     889.7194213867188,
     890.079345703125,
     898.2642211914062,
     888.8846435546875,
     899.5350341796875,
     894.9535522460938,
     918.594482421875,
     875.18359375,
     883.8409423828125,
     903.809814453125,
     887.6345825195312,
     875.0167236328125,
     914.5338745117188,
     875.4165649414062,
     895.240478515625,
     884.982177734375,
     904.2759399414062,
     897.9696044921875,
     872.4069213867188,
     902.766357421875,
     886.1785888671875,
     875.642578125,
     879.141845703125,
     889.8800048828125,
     872.7417602539062,
     885.2280883789062,
     873.5779418945312,
     884.8536376953125,
     889.5856323242188,
     875.0771484375,
     889.975830078125,
     879.3343505859375,
     881.7213745117188,
     913.1499633789062,
     903.7865600585938,
     914.137451171875,
     892.0629272460938,
     895.4060668945312,
     911.0921630859375,
     873.1602783203125,
     913.5433349609375,
     880.6787109375,
     875.4392700195312,
     883.5231323242188,
     885.645751953125,
     872.55859375,
     881.4219970703125,
     891.16552734375,
     865.645263671875,
     874.4321899414062,
     876.830078125,
     889.818603515625,
     877.8781127929688,
     911.778076171875,
     886.58740234375,
     906.398193359375,
     877.4483642578125,
     912.5648193359375,
     890.63232421875,
     882.0387573242188,
     905.4183959960938,
     889.5732421875,
     884.8413696289062,
     866.6030883789062,
     886.377685546875,
     890.9879760742188,
     904.787841796875,
     911.3005981445312,
     901.88623046875,
     870.806396484375,
     889.0690307617188,
     889.22607421875,
     903.8233642578125,
     910.7587280273438,
     886.8912353515625,
     924.6859741210938,
     915.8907470703125,
     879.3692016601562,
     908.4287109375,
     896.89599609375,
     898.6245727539062,
     886.71875,
     919.2901611328125,
     880.2071533203125,
     883.0101928710938,
     900.515380859375,
     862.9461669921875,
     877.8587036132812,
     880.078125,
     915.202880859375,
     913.90380859375,
     891.61669921875,
     892.9073486328125,
     854.7891845703125,
     908.4188842773438,
     897.5305786132812,
     882.6728515625,
     891.8377075195312,
     890.3117065429688,
     875.2115478515625,
     880.3904418945312,
     898.313232421875,
     908.2223510742188,
     896.8568725585938,
     895.6785888671875,
     878.617919921875,
     909.171875,
     887.0452880859375,
     869.6325073242188,
     892.7060546875,
     890.8649291992188,
     869.7185668945312,
     896.6862182617188,
     867.33984375,
     874.4204711914062,
     885.5037841796875,
     916.707763671875,
     941.8524780273438,
     885.7215576171875,
     897.1398315429688,
     910.3076171875,
     869.0447998046875,
     910.656982421875,
     869.6954345703125,
     888.329833984375,
     889.3101196289062,
     883.7896728515625,
     881.4376220703125,
     895.952392578125,
     891.0249633789062,
     905.899169921875,
     908.4364624023438,
     888.7196655273438,
     887.660888671875,
     867.0065307617188,
     891.4090576171875,
     884.4114379882812,
     886.5270385742188]

```python
plt.plot(losses[:])
```

```python
 torchsde.sdeint(sde, Θ_0, ts, method="euler", dt=Δt).shape
```

    torch.Size([20, 50, 2])

```python
torch.eye(3,2)
```

    tensor([[1., 0.],
            [0., 1.],
            [0., 0.]])

```python
# Mannualy doing one step of Euler

Θ_1 = torchsde.sdeint(sde, Θ_0, ts, dt=Δt)[-1,...]
```

```python
Θ_1 = Θ_1.reshape(50,2)
```

```python
fig, (ax1,ax2) = plt.subplots(1,2)

ax1.hist(Θ_1[:,0].flatten().cpu().detach().numpy())
ax2.hist(Θ_1[:,1].flatten().cpu().detach().numpy())
```

```python

# Manually doing 2 steps of euler
Θ_2 = sde.f(0,
        sde.f(0, Θ_0) * Δt + torch.normal(0,math.sqrt(Δt), (Θ_0.shape[0], Θ_0.shape[1]) ).to(device)
    ) * Δt +  torch.normal(0,math.sqrt(Δt), (Θ_0.shape[0], Θ_0.shape[1]) ).to(device)
```

```python
Θ_2
```

    tensor([[-0.0966,  0.1767],
            [ 0.0321, -0.0365],
            [-0.1502, -0.3579],
            [-0.2261, -0.1561],
            [-0.4054, -0.5228],
            [-0.0458, -0.0394],
            [-0.1069,  0.0584],
            [-0.2941,  0.1297],
            [-0.0049,  0.0908],
            [-0.0542,  0.1678],
            [ 0.3236, -0.0216],
            [-0.0798,  0.6549],
            [-0.2652,  0.3210],
            [ 0.3513, -0.0988],
            [-0.0220,  0.1065],
            [ 0.3536, -0.0596],
            [-0.1342, -0.1040],
            [ 0.1000, -0.1490],
            [-0.2025, -0.1363],
            [-0.2584,  0.1880],
            [ 0.0523, -0.4530],
            [ 0.3523,  0.0691],
            [ 0.0939,  0.1024],
            [ 0.4845, -0.0576],
            [ 0.2731, -0.3017],
            [-0.0184, -0.0646],
            [-0.3032, -0.3065],
            [ 0.2559,  0.3414],
            [ 0.0284, -0.0778],
            [ 0.2468,  0.0886],
            [ 0.3418,  0.2607],
            [-0.2585,  0.1147],
            [ 0.0282,  0.3719],
            [ 0.0214,  0.1145],
            [-0.0436,  0.2373],
            [ 0.0860, -0.3432],
            [ 0.2875,  0.5089],
            [ 0.0214,  0.3297],
            [ 0.3125, -0.2462],
            [-0.0202, -0.0354],
            [-0.0735,  0.1105],
            [-0.0358, -0.0526],
            [ 0.0853, -0.3932],
            [ 0.3369,  0.0478],
            [ 0.1126,  0.0190],
            [ 0.3073,  0.1938],
            [-0.3438,  0.4824],
            [ 0.2818,  0.3452],
            [ 0.0444, -0.2270],
            [-0.1801, -0.1227]], device='cuda:0', grad_fn=<AddBackward0>)

```python
fig, ax_1 = plt.subplots(1, 1, figsize=(16, 6))
Θ_plot = Θ_1.cpu().detach() #.reshape(50*3,2)

Θ_plot2 = Θ_2.cpu().detach()

ax_1.plot(Θ_plot[:, 0], Θ_plot[:, 1],
         linestyle='none',
         marker='o', color="red", label="Euler done by torche SDE")

# ax_1.plot(Θ_plot2[:, 0], Θ_plot2[:, 1],
#          linestyle='none',
#          marker='o', color="blue", label="Euler done manually")

ax_1.plot(X_scaled[:, 0], X_scaled[:, 1],
         linestyle='none',
         marker='o', label="Observations $X$")

ax_1.legend()
ax_1.axis('equal');
```

```python
Θ_plot[:, 0]-Θ_plot[:, 1]
```

    tensor([ 0.1793,  0.1312,  0.0910,  0.2549,  0.1410,  0.0585,  0.1621,  0.1637,
             0.2204,  0.0843,  0.0656,  0.0881,  0.1349,  0.1894,  0.1220,  0.2094,
             0.2215,  0.1651,  0.0495,  0.1466,  0.1212,  0.0617,  0.1152,  0.1255,
             0.0891,  0.1635,  0.1823,  0.2343,  0.1153,  0.1324,  0.1912,  0.1300,
             0.1652,  0.2043,  0.2035,  0.3616,  0.2387,  0.0995,  0.2329,  0.1757,
             0.1981,  0.1266,  0.1408, -0.0010,  0.2378,  0.2390,  0.0642,  0.2355,
             0.1346,  0.1499])


# Batch Normalization
One way to make deep networks easier to train is to use more sophisticated optimization procedures such as SGD+momentum, RMSProp, or Adam. Another strategy is to change the architecture of the network to make it easier to train.
One idea along these lines is batch normalization which was proposed by [1] in 2015.

The idea is relatively straightforward. Machine learning methods tend to work better when their input data consists of uncorrelated features with zero mean and unit variance. When training a neural network, we can preprocess the data before feeding it to the network to explicitly decorrelate its features; this will ensure that the first layer of the network sees data that follows a nice distribution. However, even if we preprocess the input data, the activations at deeper layers of the network will likely no longer be decorrelated and will no longer have zero mean or unit variance since they are output from earlier layers in the network. Even worse, during the training process the distribution of features at each layer of the network will shift as the weights of each layer are updated.

The authors of [1] hypothesize that the shifting distribution of features inside deep neural networks may make training deep networks more difficult. To overcome this problem, [1] proposes to insert batch normalization layers into the network. At training time, a batch normalization layer uses a minibatch of data to estimate the mean and standard deviation of each feature. These estimated means and standard deviations are then used to center and normalize the features of the minibatch. A running average of these means and standard deviations is kept during training, and at test time these running averages are used to center and normalize features.

It is possible that this normalization strategy could reduce the representational power of the network, since it may sometimes be optimal for certain layers to have features that are not zero-mean or unit variance. To this end, the batch normalization layer includes learnable shift and scale parameters for each feature dimension.

[1] [Sergey Ioffe and Christian Szegedy, "Batch Normalization: Accelerating Deep Network Training by Reducing
Internal Covariate Shift", ICML 2015.](https://arxiv.org/abs/1502.03167)

```python
# As usual, a bit of setup
import time
import numpy as np
import matplotlib.pyplot as plt
from cs231n.classifiers.fc_net import *
from cs231n.data_utils import get_CIFAR10_data
from cs231n.gradient_check import eval_numerical_gradient, eval_numerical_gradient_array
from cs231n.solver import Solver

%matplotlib inline
plt.rcParams['figure.figsize'] = (10.0, 8.0) # set default size of plots
plt.rcParams['image.interpolation'] = 'nearest'
plt.rcParams['image.cmap'] = 'gray'

# for auto-reloading external modules
# see http://stackoverflow.com/questions/1907993/autoreload-of-modules-in-ipython
%load_ext autoreload
%autoreload 2

def rel_error(x, y):
    """ returns relative error """
    return np.max(np.abs(x - y) / (np.maximum(1e-8, np.abs(x) + np.abs(y))))

def print_mean_std(x,axis=0):
    print('  means: ', x.mean(axis=axis))
    print('  stds:  ', x.std(axis=axis))
    print()
```

    run the following from the cs231n directory and try again:
    python setup.py build_ext --inplace
    You may also need to restart your iPython kernel

```python
# Load the (preprocessed) CIFAR10 data.
data = get_CIFAR10_data()
for k, v in data.items():
  print('%s: ' % k, v.shape)
```

    X_train:  (49000, 3, 32, 32)
    y_train:  (49000,)
    X_val:  (1000, 3, 32, 32)
    y_val:  (1000,)
    X_test:  (1000, 3, 32, 32)
    y_test:  (1000,)

## Batch normalization: forward
In the file `cs231n/layers.py`, implement the batch normalization forward pass in the function `batchnorm_forward`. Once you have done so, run the following to test your implementation.

Referencing the paper linked to above in [1] may be helpful!

```python
# Check the training-time forward pass by checking means and variances
# of features both before and after batch normalization

# Simulate the forward pass for a two-layer network
np.random.seed(231)
N, D1, D2, D3 = 200, 50, 60, 3
X = np.random.randn(N, D1)
W1 = np.random.randn(D1, D2)
W2 = np.random.randn(D2, D3)
a = np.maximum(0, X.dot(W1)).dot(W2)

print('Before batch normalization:')
print_mean_std(a,axis=0)

gamma = np.ones((D3,))
beta = np.zeros((D3,))
# Means should be close to zero and stds close to one
print('After batch normalization (gamma=1, beta=0)')
a_norm, _ = batchnorm_forward(a, gamma, beta, {'mode': 'train'})
print_mean_std(a_norm,axis=0)

gamma = np.asarray([1.0, 2.0, 3.0])
beta = np.asarray([11.0, 12.0, 13.0])
# Now means should be close to beta and stds close to gamma
print('After batch normalization (gamma=', gamma, ', beta=', beta, ')')
a_norm, _ = batchnorm_forward(a, gamma, beta, {'mode': 'train'})
print_mean_std(a_norm,axis=0)
```

    Before batch normalization:
      means:  [ -2.3814598  -13.18038246   1.91780462]
      stds:   [27.18502186 34.21455511 37.68611762]

    After batch normalization (gamma=1, beta=0)
      means:  [3.33066907e-17 7.04991621e-17 1.85962357e-17]
      stds:   [0.99999999 1.         1.        ]

    After batch normalization (gamma= [1. 2. 3.] , beta= [11. 12. 13.] )
      means:  [11. 12. 13.]
      stds:   [0.99999999 1.99999999 2.99999999]

```python
# Check the test-time forward pass by running the training-time
# forward pass many times to warm up the running averages, and then
# checking the means and variances of activations after a test-time
# forward pass.

np.random.seed(231)
N, D1, D2, D3 = 200, 50, 60, 3
W1 = np.random.randn(D1, D2)
W2 = np.random.randn(D2, D3)

bn_param = {'mode': 'train'}
gamma = np.ones(D3)
beta = np.zeros(D3)

for t in range(50):
  X = np.random.randn(N, D1)
  a = np.maximum(0, X.dot(W1)).dot(W2)
  batchnorm_forward(a, gamma, beta, bn_param)

bn_param['mode'] = 'test'
X = np.random.randn(N, D1)
a = np.maximum(0, X.dot(W1)).dot(W2)
a_norm, _ = batchnorm_forward(a, gamma, beta, bn_param)

# Means should be close to zero and stds close to one, but will be
# noisier than training-time forward passes.
print('After batch normalization (test-time):')
print_mean_std(a_norm,axis=0)
```

    After batch normalization (test-time):
      means:  [-0.03927354 -0.04349152 -0.10452688]
      stds:   [1.01531428 1.01238373 0.97819988]

## Batch normalization: backward
Now implement the backward pass for batch normalization in the function `batchnorm_backward`.

To derive the backward pass you should write out the computation graph for batch normalization and backprop through each of the intermediate nodes. Some intermediates may have multiple outgoing branches; make sure to sum gradients across these branches in the backward pass.

Once you have finished, run the following to numerically check your backward pass.

```python
# Gradient check batchnorm backward pass
np.random.seed(231)
N, D = 4, 5
x = 5 * np.random.randn(N, D) + 12
gamma = np.random.randn(D)
beta = np.random.randn(D)
dout = np.random.randn(N, D)

bn_param = {'mode': 'train'}
fx = lambda x: batchnorm_forward(x, gamma, beta, bn_param)[0]
fg = lambda a: batchnorm_forward(x, a, beta, bn_param)[0]
fb = lambda b: batchnorm_forward(x, gamma, b, bn_param)[0]

dx_num = eval_numerical_gradient_array(fx, x, dout)
da_num = eval_numerical_gradient_array(fg, gamma.copy(), dout)
db_num = eval_numerical_gradient_array(fb, beta.copy(), dout)

_, cache = batchnorm_forward(x, gamma, beta, bn_param)
dx, dgamma, dbeta = batchnorm_backward(dout, cache)
#You should expect to see relative errors between 1e-13 and 1e-8
print('dx error: ', rel_error(dx_num, dx))
print('dgamma error: ', rel_error(da_num, dgamma))
print('dbeta error: ', rel_error(db_num, dbeta))
```

    dx error:  1.6674604875341426e-09
    dgamma error:  7.378954934067588e-13
    dbeta error:  2.6853898000928757e-12

## Batch normalization: alternative backward
In class we talked about two different implementations for the sigmoid backward pass. One strategy is to write out a computation graph composed of simple operations and backprop through all intermediate values. Another strategy is to work out the derivatives on paper. For example, you can derive a very simple formula for the sigmoid function's backward pass by simplifying gradients on paper.

Surprisingly, it turns out that you can do a similar simplification for the batch normalization backward pass too!

In the forward pass, given a set of inputs $X=\begin{bmatrix}x_1\\x_2\\...\\x_N\end{bmatrix}$,

we first calculate the mean $\mu$ and variance $v$.
With $\mu$ and $v$ calculated, we can calculate the standard deviation $\sigma$  and normalized data $Y$.
The equations and graph illustration below describe the computation ($y_i$ is the i-th element of the vector $Y$).

\begin{align}
& \mu=\frac{1}{N}\sum_{k=1}^N x_k  &  v=\frac{1}{N}\sum_{k=1}^N (x_k-\mu)^2 \\
& \sigma=\sqrt{v+\epsilon}         &  y_i=\frac{x_i-\mu}{\sigma}
\end{align}

The meat of our problem during backpropagation is to compute $\frac{\partial L}{\partial X}$, given the upstream gradient we receive, $\frac{\partial L}{\partial Y}.$ To do this, recall the chain rule in calculus gives us $\frac{\partial L}{\partial X} = \frac{\partial L}{\partial Y} \cdot \frac{\partial Y}{\partial X}$.

The unknown/hart part is $\frac{\partial Y}{\partial X}$. We can find this by first deriving step-by-step our local gradients at
$\frac{\partial v}{\partial X}$, $\frac{\partial \mu}{\partial X}$,
$\frac{\partial \sigma}{\partial v}$,
$\frac{\partial Y}{\partial \sigma}$, and $\frac{\partial Y}{\partial \mu}$,
and then use the chain rule to compose these gradients (which appear in the form of vectors!) appropriately to compute $\frac{\partial Y}{\partial X}$.

If it's challenging to directly reason about the gradients over $X$ and $Y$ which require matrix multiplication, try reasoning about the gradients in terms of individual elements $x_i$ and $y_i$ first: in that case, you will need to come up with the derivations for $\frac{\partial L}{\partial x_i}$, by relying on the Chain Rule to first calculate the intermediate $\frac{\partial \mu}{\partial x_i}, \frac{\partial v}{\partial x_i}, \frac{\partial \sigma}{\partial x_i},$ then assemble these pieces to calculate $\frac{\partial y_i}{\partial x_i}$.

You should make sure each of the intermediary gradient derivations are all as simplified as possible, for ease of implementation.

After doing so, implement the simplified batch normalization backward pass in the function `batchnorm_backward_alt` and compare the two implementations by running the following. Your two implementations should compute nearly identical results, but the alternative implementation should be a bit faster.

```python
np.random.seed(231)
N, D = 100, 500
x = 5 * np.random.randn(N, D) + 12
gamma = np.random.randn(D)
beta = np.random.randn(D)
dout = np.random.randn(N, D)

bn_param = {'mode': 'train'}
out, cache = batchnorm_forward(x, gamma, beta, bn_param)

t1 = time.time()
dx1, dgamma1, dbeta1 = batchnorm_backward(dout, cache)
t2 = time.time()
dx2, dgamma2, dbeta2 = batchnorm_backward_alt(dout, cache)
t3 = time.time()

print('dx difference: ', rel_error(dx1, dx2))
print('dgamma difference: ', rel_error(dgamma1, dgamma2))
print('dbeta difference: ', rel_error(dbeta1, dbeta2))
print('speedup: %.2fx' % ((t2 - t1) / (t3 - t2)))
```

    dx difference:  7.666703093511459e-13
    dgamma difference:  0.0
    dbeta difference:  0.0
    speedup: 1.12x

## Fully Connected Nets with Batch Normalization
Now that you have a working implementation for batch normalization, go back to your `FullyConnectedNet` in the file `cs231n/classifiers/fc_net.py`. Modify your implementation to add batch normalization.

Concretely, when the `normalization` flag is set to `"batchnorm"` in the constructor, you should insert a batch normalization layer before each ReLU nonlinearity. The outputs from the last layer of the network should not be normalized. Once you are done, run the following to gradient-check your implementation.

HINT: You might find it useful to define an additional helper layer similar to those in the file `cs231n/layer_utils.py`. If you decide to do so, do it in the file `cs231n/classifiers/fc_net.py`.

```python
np.random.seed(231)
N, D, H1, H2, C = 2, 15, 20, 30, 10
X = np.random.randn(N, D)
y = np.random.randint(C, size=(N,))

# You should expect losses between 1e-4~1e-10 for W,
# losses between 1e-08~1e-10 for b,
# and losses between 1e-08~1e-09 for beta and gammas.
for reg in [0, 3.14]:
  print('Running check with reg = ', reg)
  model = FullyConnectedNet([H1, H2], input_dim=D, num_classes=C,
                            reg=reg, weight_scale=5e-2, dtype=np.float64,
                            normalization='batchnorm')

  loss, grads = model.loss(X, y)
  print('Initial loss: ', loss)

  for name in sorted(grads):
    f = lambda _: model.loss(X, y)[0]
    grad_num = eval_numerical_gradient(f, model.params[name], verbose=False, h=1e-5)
    print('%s relative error: %.2e' % (name, rel_error(grad_num, grads[name])))
  if reg == 0: print()
```

    Running check with reg =  0
    Initial loss:  2.2611955101340957
    W1 relative error: 1.10e-04
    W2 relative error: 2.85e-06
    W3 relative error: 3.75e-10
    b1 relative error: 1.39e-09
    b2 relative error: 2.22e-08
    b3 relative error: 1.01e-10
    beta1 relative error: 6.94e-09
    beta2 relative error: 4.84e-09
    gamma1 relative error: 6.96e-09
    gamma2 relative error: 2.41e-09

    Running check with reg =  3.14
    Initial loss:  6.996533220108303
    W1 relative error: 1.98e-06
    W2 relative error: 2.29e-06
    W3 relative error: 2.79e-08
    b1 relative error: 5.55e-09
    b2 relative error: 2.22e-08
    b3 relative error: 2.10e-10
    beta1 relative error: 6.65e-09
    beta2 relative error: 3.39e-09
    gamma1 relative error: 6.27e-09
    gamma2 relative error: 5.28e-09

# Batchnorm for deep networks
Run the following to train a six-layer network on a subset of 1000 training examples both with and without batch normalization.

```python
np.random.seed(231)
# Try training a very deep net with batchnorm
hidden_dims = [100, 100, 100, 100, 100]

num_train = 1000
small_data = {
  'X_train': data['X_train'][:num_train],
  'y_train': data['y_train'][:num_train],
  'X_val': data['X_val'],
  'y_val': data['y_val'],
}

weight_scale = 2e-2
bn_model = FullyConnectedNet(hidden_dims, weight_scale=weight_scale, normalization='batchnorm')
model = FullyConnectedNet(hidden_dims, weight_scale=weight_scale, normalization=None)

print('Solver with batch norm:')
bn_solver = Solver(bn_model, small_data,
                num_epochs=10, batch_size=50,
                update_rule='adam',
                optim_config={
                  'learning_rate': 1e-3,
                },
                verbose=True,print_every=20)
bn_solver.train()

print('\nSolver without batch norm:')
solver = Solver(model, small_data,
                num_epochs=10, batch_size=50,
                update_rule='adam',
                optim_config={
                  'learning_rate': 1e-3,
                },
                verbose=True, print_every=20)
solver.train()
```

    Solver with batch norm:
    (Iteration 1 / 200) loss: 2.340974
    (Epoch 0 / 10) train acc: 0.107000; val_acc: 0.115000
    (Epoch 1 / 10) train acc: 0.313000; val_acc: 0.266000
    (Iteration 21 / 200) loss: 2.039365
    (Epoch 2 / 10) train acc: 0.383000; val_acc: 0.278000
    (Iteration 41 / 200) loss: 2.041103
    (Epoch 3 / 10) train acc: 0.492000; val_acc: 0.308000
    (Iteration 61 / 200) loss: 1.753902
    (Epoch 4 / 10) train acc: 0.529000; val_acc: 0.304000
    (Iteration 81 / 200) loss: 1.241217
    (Epoch 5 / 10) train acc: 0.594000; val_acc: 0.310000
    (Iteration 101 / 200) loss: 1.371031
    (Epoch 6 / 10) train acc: 0.645000; val_acc: 0.325000
    (Iteration 121 / 200) loss: 1.127299
    (Epoch 7 / 10) train acc: 0.689000; val_acc: 0.323000
    (Iteration 141 / 200) loss: 1.185819
    (Epoch 8 / 10) train acc: 0.741000; val_acc: 0.305000
    (Iteration 161 / 200) loss: 0.825489
    (Epoch 9 / 10) train acc: 0.776000; val_acc: 0.332000
    (Iteration 181 / 200) loss: 0.896460
    (Epoch 10 / 10) train acc: 0.787000; val_acc: 0.331000

    Solver without batch norm:
    (Iteration 1 / 200) loss: 2.302332
    (Epoch 0 / 10) train acc: 0.129000; val_acc: 0.131000
    (Epoch 1 / 10) train acc: 0.283000; val_acc: 0.250000
    (Iteration 21 / 200) loss: 2.041970
    (Epoch 2 / 10) train acc: 0.316000; val_acc: 0.277000
    (Iteration 41 / 200) loss: 1.900473
    (Epoch 3 / 10) train acc: 0.373000; val_acc: 0.282000
    (Iteration 61 / 200) loss: 1.713156
    (Epoch 4 / 10) train acc: 0.390000; val_acc: 0.310000
    (Iteration 81 / 200) loss: 1.662209
    (Epoch 5 / 10) train acc: 0.434000; val_acc: 0.300000
    (Iteration 101 / 200) loss: 1.696062
    (Epoch 6 / 10) train acc: 0.536000; val_acc: 0.346000
    (Iteration 121 / 200) loss: 1.550785
    (Epoch 7 / 10) train acc: 0.530000; val_acc: 0.310000
    (Iteration 141 / 200) loss: 1.436308
    (Epoch 8 / 10) train acc: 0.622000; val_acc: 0.342000
    (Iteration 161 / 200) loss: 1.000868
    (Epoch 9 / 10) train acc: 0.654000; val_acc: 0.328000
    (Iteration 181 / 200) loss: 0.925455
    (Epoch 10 / 10) train acc: 0.726000; val_acc: 0.335000

Run the following to visualize the results from two networks trained above. You should find that using batch normalization helps the network to converge much faster.

```python
def plot_training_history(title, label, baseline, bn_solvers, plot_fn, bl_marker='.', bn_marker='.', labels=None):
    """utility function for plotting training history"""
    plt.title(title)
    plt.xlabel(label)
    bn_plots = [plot_fn(bn_solver) for bn_solver in bn_solvers]
    bl_plot = plot_fn(baseline)
    num_bn = len(bn_plots)
    for i in range(num_bn):
        label='with_norm'
        if labels is not None:
            label += str(labels[i])
        plt.plot(bn_plots[i], bn_marker, label=label)
    label='baseline'
    if labels is not None:
        label += str(labels[0])
    plt.plot(bl_plot, bl_marker, label=label)
    plt.legend(loc='lower center', ncol=num_bn+1)

plt.subplot(3, 1, 1)
plot_training_history('Training loss','Iteration', solver, [bn_solver], \
                      lambda x: x.loss_history, bl_marker='o', bn_marker='o')
plt.subplot(3, 1, 2)
plot_training_history('Training accuracy','Epoch', solver, [bn_solver], \
                      lambda x: x.train_acc_history, bl_marker='-o', bn_marker='-o')
plt.subplot(3, 1, 3)
plot_training_history('Validation accuracy','Epoch', solver, [bn_solver], \
                      lambda x: x.val_acc_history, bl_marker='-o', bn_marker='-o')

plt.gcf().set_size_inches(15, 15)
plt.show()
```

# Batch normalization and initialization
We will now run a small experiment to study the interaction of batch normalization and weight initialization.

The first cell will train 8-layer networks both with and without batch normalization using different scales for weight initialization. The second layer will plot training accuracy, validation set accuracy, and training loss as a function of the weight initialization scale.

```python
np.random.seed(231)
# Try training a very deep net with batchnorm
hidden_dims = [50, 50, 50, 50, 50, 50, 50]
num_train = 1000
small_data = {
  'X_train': data['X_train'][:num_train],
  'y_train': data['y_train'][:num_train],
  'X_val': data['X_val'],
  'y_val': data['y_val'],
}

bn_solvers_ws = {}
solvers_ws = {}
weight_scales = np.logspace(-4, 0, num=20)
for i, weight_scale in enumerate(weight_scales):
  print('Running weight scale %d / %d' % (i + 1, len(weight_scales)))
  bn_model = FullyConnectedNet(hidden_dims, weight_scale=weight_scale, normalization='batchnorm')
  model = FullyConnectedNet(hidden_dims, weight_scale=weight_scale, normalization=None)

  bn_solver = Solver(bn_model, small_data,
                  num_epochs=10, batch_size=50,
                  update_rule='adam',
                  optim_config={
                    'learning_rate': 1e-3,
                  },
                  verbose=False, print_every=200)
  bn_solver.train()
  bn_solvers_ws[weight_scale] = bn_solver

  solver = Solver(model, small_data,
                  num_epochs=10, batch_size=50,
                  update_rule='adam',
                  optim_config={
                    'learning_rate': 1e-3,
                  },
                  verbose=False, print_every=200)
  solver.train()
  solvers_ws[weight_scale] = solver
```

    Running weight scale 1 / 20
    Running weight scale 2 / 20
    Running weight scale 3 / 20
    Running weight scale 4 / 20
    Running weight scale 5 / 20
    Running weight scale 6 / 20
    Running weight scale 7 / 20
    Running weight scale 8 / 20
    Running weight scale 9 / 20
    Running weight scale 10 / 20
    Running weight scale 11 / 20
    Running weight scale 12 / 20
    Running weight scale 13 / 20
    Running weight scale 14 / 20
    Running weight scale 15 / 20
    Running weight scale 16 / 20
    Running weight scale 17 / 20
    Running weight scale 18 / 20
    Running weight scale 19 / 20
    Running weight scale 20 / 20

```python
# Plot results of weight scale experiment
best_train_accs, bn_best_train_accs = [], []
best_val_accs, bn_best_val_accs = [], []
final_train_loss, bn_final_train_loss = [], []

for ws in weight_scales:
  best_train_accs.append(max(solvers_ws[ws].train_acc_history))
  bn_best_train_accs.append(max(bn_solvers_ws[ws].train_acc_history))

  best_val_accs.append(max(solvers_ws[ws].val_acc_history))
  bn_best_val_accs.append(max(bn_solvers_ws[ws].val_acc_history))

  final_train_loss.append(np.mean(solvers_ws[ws].loss_history[-100:]))
  bn_final_train_loss.append(np.mean(bn_solvers_ws[ws].loss_history[-100:]))

plt.subplot(3, 1, 1)
plt.title('Best val accuracy vs weight initialization scale')
plt.xlabel('Weight initialization scale')
plt.ylabel('Best val accuracy')
plt.semilogx(weight_scales, best_val_accs, '-o', label='baseline')
plt.semilogx(weight_scales, bn_best_val_accs, '-o', label='batchnorm')
plt.legend(ncol=2, loc='lower right')

plt.subplot(3, 1, 2)
plt.title('Best train accuracy vs weight initialization scale')
plt.xlabel('Weight initialization scale')
plt.ylabel('Best training accuracy')
plt.semilogx(weight_scales, best_train_accs, '-o', label='baseline')
plt.semilogx(weight_scales, bn_best_train_accs, '-o', label='batchnorm')
plt.legend()

plt.subplot(3, 1, 3)
plt.title('Final training loss vs weight initialization scale')
plt.xlabel('Weight initialization scale')
plt.ylabel('Final training loss')
plt.semilogx(weight_scales, final_train_loss, '-o', label='baseline')
plt.semilogx(weight_scales, bn_final_train_loss, '-o', label='batchnorm')
plt.legend()
plt.gca().set_ylim(1.0, 3.5)

plt.gcf().set_size_inches(15, 15)
plt.show()
```

## Inline Question 1:
Describe the results of this experiment. How does the scale of weight initialization affect models with/without batch normalization differently, and why?

## Answer:
Without batch normalization, if weights are too small or too high network will immediately saturate and the performance are poor. With batch normalization wee normalize the input to the relu, so the network doesn't saturate any more.

# Batch normalization and batch size
We will now run a small experiment to study the interaction of batch normalization and batch size.

The first cell will train 6-layer networks both with and without batch normalization using different batch sizes. The second layer will plot training accuracy and validation set accuracy over time.

```python
def run_batchsize_experiments(normalization_mode):
    np.random.seed(231)
    # Try training a very deep net with batchnorm
    hidden_dims = [100, 100, 100, 100, 100]
    num_train = 1000
    small_data = {
      'X_train': data['X_train'][:num_train],
      'y_train': data['y_train'][:num_train],
      'X_val': data['X_val'],
      'y_val': data['y_val'],
    }
    n_epochs=10
    weight_scale = 2e-2
    batch_sizes = [5,10,50]
    lr = 10**(-3.5)
    solver_bsize = batch_sizes[0]

    print('No normalization: batch size = ',solver_bsize)
    model = FullyConnectedNet(hidden_dims, weight_scale=weight_scale, normalization=None)
    solver = Solver(model, small_data,
                    num_epochs=n_epochs, batch_size=solver_bsize,
                    update_rule='adam',
                    optim_config={
                      'learning_rate': lr,
                    },
                    verbose=False)
    solver.train()

    bn_solvers = []
    for i in range(len(batch_sizes)):
        b_size=batch_sizes[i]
        print('Normalization: batch size = ',b_size)
        bn_model = FullyConnectedNet(hidden_dims, weight_scale=weight_scale, normalization=normalization_mode)
        bn_solver = Solver(bn_model, small_data,
                        num_epochs=n_epochs, batch_size=b_size,
                        update_rule='adam',
                        optim_config={
                          'learning_rate': lr,
                        },
                        verbose=False)
        bn_solver.train()
        bn_solvers.append(bn_solver)

    return bn_solvers, solver, batch_sizes

batch_sizes = [5,10,50]
bn_solvers_bsize, solver_bsize, batch_sizes = run_batchsize_experiments('batchnorm')
```

    No normalization: batch size =  5
    Normalization: batch size =  5
    Normalization: batch size =  10
    Normalization: batch size =  50

```python
plt.subplot(2, 1, 1)
plot_training_history('Training accuracy (Batch Normalization)','Epoch', solver_bsize, bn_solvers_bsize, \
                      lambda x: x.train_acc_history, bl_marker='-^', bn_marker='-o', labels=batch_sizes)
plt.subplot(2, 1, 2)
plot_training_history('Validation accuracy (Batch Normalization)','Epoch', solver_bsize, bn_solvers_bsize, \
                      lambda x: x.val_acc_history, bl_marker='-^', bn_marker='-o', labels=batch_sizes)

plt.gcf().set_size_inches(15, 10)
plt.show()
```

## Inline Question 2:
Describe the results of this experiment. What does this imply about the relationship between batch normalization and batch size? Why is this relationship observed?

## Answer:
training network with and without batch normalization with batch size of 5, we see that the training accuracy is lower, while the validation accuracy remains similar. This suggests that the batch normalization helps regulazing the network when we use small batches.

# Layer Normalization
Batch normalization has proved to be effective in making networks easier to train, but the dependency on batch size makes it less useful in complex networks which have a cap on the input batch size due to hardware limitations.

Several alternatives to batch normalization have been proposed to mitigate this problem; one such technique is Layer Normalization [2]. Instead of normalizing over the batch, we normalize over the features. In other words, when using Layer Normalization, each feature vector corresponding to a single datapoint is normalized based on the sum of all terms within that feature vector.
Also, we can notice that validation accuracy reamins similar for different variations in batch size.

[2] [Ba, Jimmy Lei, Jamie Ryan Kiros, and Geoffrey E. Hinton. "Layer Normalization." stat 1050 (2016): 21.](https://arxiv.org/pdf/1607.06450.pdf)

## Inline Question 3:
Which of these data preprocessing steps is analogous to batch normalization, and which is analogous to layer normalization?

1. Scaling each image in the dataset, so that the RGB channels for each row of pixels within an image sums up to 1.
2. Scaling each image in the dataset, so that the RGB channels for all pixels within an image sums up to 1.
3. Subtracting the mean image of the dataset from each image in the dataset.
4. Setting all RGB values to either 0 or 1 depending on a given threshold.

## Answer:
1. layer
2. layer
3. batch
4. batch

# Layer Normalization: Implementation

Now you'll implement layer normalization. This step should be relatively straightforward, as conceptually the implementation is almost identical to that of batch normalization. One significant difference though is that for layer normalization, we do not keep track of the moving moments, and the testing phase is identical to the training phase, where the mean and variance are directly calculated per datapoint.

Here's what you need to do:

* In `cs231n/layers.py`, implement the forward pass for layer normalization in the function `layernorm_backward`.

Run the cell below to check your results.
* In `cs231n/layers.py`, implement the backward pass for layer normalization in the function `layernorm_backward`.

Run the second cell below to check your results.
* Modify `cs231n/classifiers/fc_net.py` to add layer normalization to the `FullyConnectedNet`. When the `normalization` flag is set to `"layernorm"` in the constructor, you should insert a layer normalization layer before each ReLU nonlinearity.

Run the third cell below to run the batch size experiment on layer normalization.

```python
# Check the training-time forward pass by checking means and variances
# of features both before and after layer normalization

# Simulate the forward pass for a two-layer network
np.random.seed(231)
N, D1, D2, D3 =4, 50, 60, 3
X = np.random.randn(N, D1)
W1 = np.random.randn(D1, D2)
W2 = np.random.randn(D2, D3)
a = np.maximum(0, X.dot(W1)).dot(W2)

print('Before layer normalization:')
print_mean_std(a,axis=1)

gamma = np.ones(D3)
beta = np.zeros(D3)
# Means should be close to zero and stds close to one
print('After layer normalization (gamma=1, beta=0)')
a_norm, _ = layernorm_forward(a, gamma, beta, {'mode': 'train'})
print_mean_std(a_norm,axis=1)

gamma = np.asarray([3.0,3.0,3.0])
beta = np.asarray([5.0,5.0,5.0])
# Now means should be close to beta and stds close to gamma
print('After layer normalization (gamma=', gamma, ', beta=', beta, ')')
a_norm, _ = layernorm_forward(a, gamma, beta, {'mode': 'train'})
print_mean_std(a_norm,axis=1)
```

    Before layer normalization:
      means:  [-59.06673243 -47.60782686 -43.31137368 -26.40991744]
      stds:   [10.07429373 28.39478981 35.28360729  4.01831507]

    After layer normalization (gamma=1, beta=0)
      means:  [ 4.81096644e-16 -7.40148683e-17  2.22044605e-16 -5.92118946e-16]
      stds:   [0.99999995 0.99999999 1.         0.99999969]

    After layer normalization (gamma= [3. 3. 3.] , beta= [5. 5. 5.] )
      means:  [5. 5. 5. 5.]
      stds:   [2.99999985 2.99999998 2.99999999 2.99999907]

```python
# Gradient check batchnorm backward pass
np.random.seed(231)
N, D = 4, 5
x = 5 * np.random.randn(N, D) + 12
gamma = np.random.randn(D)
beta = np.random.randn(D)
dout = np.random.randn(N, D)

ln_param = {}
fx = lambda x: layernorm_forward(x, gamma, beta, ln_param)[0]
fg = lambda a: layernorm_forward(x, a, beta, ln_param)[0]
fb = lambda b: layernorm_forward(x, gamma, b, ln_param)[0]

dx_num = eval_numerical_gradient_array(fx, x, dout)
da_num = eval_numerical_gradient_array(fg, gamma.copy(), dout)
db_num = eval_numerical_gradient_array(fb, beta.copy(), dout)

_, cache = layernorm_forward(x, gamma, beta, ln_param)
dx, dgamma, dbeta = layernorm_backward(dout, cache)

#You should expect to see relative errors between 1e-12 and 1e-8
print('dx error: ', rel_error(dx_num, dx))
print('dgamma error: ', rel_error(da_num, dgamma))
print('dbeta error: ', rel_error(db_num, dbeta))
```

    dx error:  1.3050018584741255e-09
    dgamma error:  4.759518365681667e-12
    dbeta error:  2.6853898000928757e-12

# Layer Normalization and batch size

We will now run the previous batch size experiment with layer normalization instead of batch normalization. Compared to the previous experiment, you should see a markedly smaller influence of batch size on the training history!

```python
ln_solvers_bsize, solver_bsize, batch_sizes = run_batchsize_experiments('layernorm')

plt.subplot(2, 1, 1)
plot_training_history('Training accuracy (Layer Normalization)','Epoch', solver_bsize, ln_solvers_bsize, \
                      lambda x: x.train_acc_history, bl_marker='-^', bn_marker='-o', labels=batch_sizes)
plt.subplot(2, 1, 2)
plot_training_history('Validation accuracy (Layer Normalization)','Epoch', solver_bsize, ln_solvers_bsize, \
                      lambda x: x.val_acc_history, bl_marker='-^', bn_marker='-o', labels=batch_sizes)

plt.gcf().set_size_inches(15, 10)
plt.show()
```

## Inline Question 4:
When is layer normalization likely to not work well, and why?

1. Using it in a very deep network
2. Having a very small dimension of features
3. Having a high regularization term

## Answer:
1. It works, we just used kind a deep network for training.
2. Yes, this is a problem if the dimension of sample is very small. The statistics will be noisy.
On the other hand, problems with small dimensionality can be solved much easily with other machine learning methods with
 quite smaller structure.