# Neuromatch Academy: Week 1, Day 5, Tutorial 1
# Dimensionality Reduction: Geometric view of data

__Content creators:__ Alex Cayco Gajic, John Murray

__Content reviewers:__ Roozbeh Farhoudi, Matt Krause, Spiros Chavlis, Richard Gao, Michael Waskom

---
# Tutorial Objectives

In this notebook we'll explore how multivariate data can be represented in different orthonormal bases. This will help us build intuition that will be helpful in understanding PCA in the following tutorial.

Overview:
 - Generate correlated multivariate data.
 - Define an arbitrary orthonormal basis.
 - Project the data onto the new basis.

```python
# @title Video 1: Geometric view of data
from IPython.display import IFrame
class BiliVideo(IFrame):
    def __init__(self, id, page=1, width=400, height=300, **kwargs):
        self.id=id
        src = "https://player.bilibili.com/player.html?bvid={0}&page={1}".format(id, page)
        super(BiliVideo, self).__init__(src, width, height, **kwargs)

video = BiliVideo(id='BV1Af4y1R78w', width=854, height=480, fs=1)
print("Video available at https://www.bilibili.com/video/{0}".format(video.id))
video
```

    Video available at https://www.bilibili.com/video/BV1Af4y1R78w

---
# Setup

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

```python
# @title Figure Settings
import ipywidgets as widgets  # interactive display
%config InlineBackend.figure_format = 'retina'
plt.style.use("/share/dataset/COMMON/nma.mplstyle.txt")
```

```python
# @title Helper functions

def get_data(cov_matrix):
  """
  Returns a matrix of 1000 samples from a bivariate, zero-mean Gaussian.

  Note that samples are sorted in ascending order for the first random variable

  Args:
    cov_matrix (numpy array of floats): desired covariance matrix

  Returns:
    (numpy array of floats) : samples from the bivariate Gaussian, with each
                              column corresponding to a different random
                              variable
  """

  mean = np.array([0, 0])
  X = np.random.multivariate_normal(mean, cov_matrix, size=1000)
  indices_for_sorting = np.argsort(X[:, 0])
  X = X[indices_for_sorting, :]

  return X

def plot_data(X):
  """
  Plots bivariate data. Includes a plot of each random variable, and a scatter
  plot of their joint activity. The title indicates the sample correlation
  calculated from the data.

  Args:
    X (numpy array of floats) :   Data matrix each column corresponds to a
                                  different random variable

  Returns:
    Nothing.
  """

  fig = plt.figure(figsize=[8, 4])
  gs = fig.add_gridspec(2, 2)
  ax1 = fig.add_subplot(gs[0, 0])
  ax1.plot(X[:, 0], color='k')
  plt.ylabel('Neuron 1')
  plt.title('Sample var 1: {:.1f}'.format(np.var(X[:, 0])))
  ax1.set_xticklabels([])
  ax2 = fig.add_subplot(gs[1, 0])
  ax2.plot(X[:, 1], color='k')
  plt.xlabel('Sample Number')
  plt.ylabel('Neuron 2')
  plt.title('Sample var 2: {:.1f}'.format(np.var(X[:, 1])))
  ax3 = fig.add_subplot(gs[:, 1])
  ax3.plot(X[:, 0], X[:, 1], '.', markerfacecolor=[.5, .5, .5],
           markeredgewidth=0)
  ax3.axis('equal')
  plt.xlabel('Neuron 1 activity')
  plt.ylabel('Neuron 2 activity')
  plt.title('Sample corr: {:.1f}'.format(np.corrcoef(X[:, 0], X[:, 1])[0, 1]))
  plt.show()

def plot_basis_vectors(X, W):
  """
  Plots bivariate data as well as new basis vectors.

  Args:
    X (numpy array of floats) :   Data matrix each column corresponds to a
                                  different random variable
    W (numpy array of floats) :   Square matrix representing new orthonormal
                                  basis each column represents a basis vector

  Returns:
    Nothing.
  """

  plt.figure(figsize=[4, 4])
  plt.plot(X[:, 0], X[:, 1], '.', color=[.5, .5, .5], label='Data')
  plt.axis('equal')
  plt.xlabel('Neuron 1 activity')
  plt.ylabel('Neuron 2 activity')
  plt.plot([0, W[0, 0]], [0, W[1, 0]], color='r', linewidth=3,
           label='Basis vector 1')
  plt.plot([0, W[0, 1]], [0, W[1, 1]], color='b', linewidth=3,
           label='Basis vector 2')
  plt.legend()
  plt.show()

def plot_data_new_basis(Y):
  """
  Plots bivariate data after transformation to new bases.
  Similar to plot_data but with colors corresponding to projections onto
  basis 1 (red) and basis 2 (blue). The title indicates the sample correlation
  calculated from the data.

  Note that samples are re-sorted in ascending order for the first
  random variable.

  Args:
    Y (numpy array of floats):   Data matrix in new basis each column
                                 corresponds to a different random variable

  Returns:
    Nothing.
  """
  fig = plt.figure(figsize=[8, 4])
  gs = fig.add_gridspec(2, 2)
  ax1 = fig.add_subplot(gs[0, 0])
  ax1.plot(Y[:, 0], 'r')
  plt.xlabel
  plt.ylabel('Projection \n basis vector 1')
  plt.title('Sample var 1: {:.1f}'.format(np.var(Y[:, 0])))
  ax1.set_xticklabels([])
  ax2 = fig.add_subplot(gs[1, 0])
  ax2.plot(Y[:, 1], 'b')
  plt.xlabel('Sample number')
  plt.ylabel('Projection \n basis vector 2')
  plt.title('Sample var 2: {:.1f}'.format(np.var(Y[:, 1])))
  ax3 = fig.add_subplot(gs[:, 1])
  ax3.plot(Y[:, 0], Y[:, 1], '.', color=[.5, .5, .5])
  ax3.axis('equal')
  plt.xlabel('Projection basis vector 1')
  plt.ylabel('Projection basis vector 2')
  plt.title('Sample corr: {:.1f}'.format(np.corrcoef(Y[:, 0], Y[:, 1])[0, 1]))
  plt.show()
```

---
# Section 1: Generate correlated multivariate data

```python
# @title Video 2: Multivariate data
from IPython.display import IFrame
class BiliVideo(IFrame):
    def __init__(self, id, page=1, width=400, height=300, **kwargs):
        self.id=id
        src = "https://player.bilibili.com/player.html?bvid={0}&page={1}".format(id, page)
        super(BiliVideo, self).__init__(src, width, height, **kwargs)

video = BiliVideo(id='BV1xz4y1D7ES', width=854, height=480, fs=1)
print("Video available at https://www.bilibili.com/video/{0}".format(video.id))
video
```

    Video available at https://www.bilibili.com/video/BV1xz4y1D7ES

To gain intuition, we will first use a simple model to generate multivariate data. Specifically, we will draw random samples from a *bivariate normal distribution*. This is an extension of the one-dimensional normal distribution to two dimensions, in which each $x_i$ is marginally normal with mean $\mu_i$ and variance $\sigma_i^2$:

\begin{align}
x_i \sim \mathcal{N}(\mu_i,\sigma_i^2).
\end{align}

Additionally, the joint distribution for $x_1$ and $x_2$ has a specified correlation coefficient $\rho$. Recall that the correlation coefficient is a normalized version of the covariance, and ranges between -1 and +1:

\begin{align}
\rho = \frac{\text{cov}(x_1,x_2)}{\sqrt{\sigma_1^2 \sigma_2^2}}.
\end{align}

For simplicity, we will assume that the mean of each variable has already been subtracted, so that $\mu_i=0$. The remaining parameters can be summarized in the covariance matrix, which for two dimensions has the following form:

\begin{equation*}
{\bf \Sigma} =
\begin{pmatrix}
 \text{var}(x_1) & \text{cov}(x_1,x_2) \\
 \text{cov}(x_1,x_2) &\text{var}(x_2)
\end{pmatrix}.
\end{equation*}

In general, $\bf \Sigma$ is a symmetric matrix with the variances $\text{var}(x_i) = \sigma_i^2$ on the diagonal, and the covariances on the off-diagonal. Later, we will see that the covariance matrix plays a key role in PCA.

## Exercise 1: Draw samples from a distribution

We have provided code to draw random samples from a zero-mean bivariate normal distribution. Throughout this tutorial, we'll imagine these samples represent the activity (firing rates) of two recorded neurons on different trials. Fill in the function below to calculate the covariance matrix given the desired variances and correlation coefficient. The covariance can be found by rearranging the equation above:

\begin{align}
\text{cov}(x_1,x_2) = \rho \sqrt{\sigma_1^2 \sigma_2^2}.
\end{align}

Use these functions to generate and plot data while varying the parameters. You should get a feel for how changing the correlation coefficient affects the geometry of the simulated data.

**Steps**
* Fill in the function `calculate_cov_matrix` to calculate the desired covariance.
* Generate and plot the data for $\sigma_1^2 =1$, $\sigma_1^2 =1$, and $\rho = .8$. Try plotting the data for different values of the correlation coefficent: $\rho = -1, -.5, 0, .5, 1$.

```python
help(plot_data)
help(get_data)
```

    Help on function plot_data in module __main__:

    plot_data(X)
        Plots bivariate data. Includes a plot of each random variable, and a scatter
        plot of their joint activity. The title indicates the sample correlation
        calculated from the data.

        Args:
          X (numpy array of floats) :   Data matrix each column corresponds to a
                                        different random variable

        Returns:
          Nothing.

    Help on function get_data in module __main__:

    get_data(cov_matrix)
        Returns a matrix of 1000 samples from a bivariate, zero-mean Gaussian.

        Note that samples are sorted in ascending order for the first random variable

        Args:
          cov_matrix (numpy array of floats): desired covariance matrix

        Returns:
          (numpy array of floats) : samples from the bivariate Gaussian, with each
                                    column corresponding to a different random
                                    variable

```python
def calculate_cov_matrix(var_1, var_2, corr_coef):
  """
  Calculates the covariance matrix based on the variances and correlation
  coefficient.

  Args:
    var_1 (scalar)          : variance of the first random variable
    var_2 (scalar)          : variance of the second random variable
    corr_coef (scalar)      : correlation coefficient

  Returns:
    (numpy array of floats) : covariance matrix
  """

  #################################################
  ## TODO for students: calculate the covariance matrix
  # Fill out function and remove
  raise NotImplementedError("Student excercise: calculate the covariance matrix!")
  #################################################

  # Calculate the covariance from the variances and correlation
  cov = ...

  cov_matrix = np.array([[var_1, cov], [cov, var_2]])

  return cov_matrix

###################################################################
## TO DO for students: generate and plot bivariate Gaussian data with variances of 1
## and a correlation coefficients of: 0.8
## repeat while varying the correlation coefficient from -1 to 1
###################################################################
np.random.seed(2020)  # set random seed
variance_1 = 1
variance_2 = 1
corr_coef = 0.8

# Uncomment to test your code and plot
# cov_matrix = calculate_cov_matrix(variance_1, variance_2, corr_coef)
# X = get_data(cov_matrix)
# plot_data(X)
```

[*Click for solution*](https://github.com/erlichlab/course-content/tree/master//tutorials/W1D5_DimensionalityReduction/solutions/W1D5_Tutorial1_Solution_57497711.py)

*Example output:*

---
# Section 2: Define a new orthonormal basis

```python
# @title Video 3: Orthonormal bases
from IPython.display import IFrame
class BiliVideo(IFrame):
    def __init__(self, id, page=1, width=400, height=300, **kwargs):
        self.id=id
        src = "https://player.bilibili.com/player.html?bvid={0}&page={1}".format(id, page)
        super(BiliVideo, self).__init__(src, width, height, **kwargs)

video = BiliVideo(id='BV1wT4y1E71g', width=854, height=480, fs=1)
print("Video available at https://www.bilibili.com/video/{0}".format(video.id))
video
```

    Video available at https://www.bilibili.com/video/BV1wT4y1E71g

Next, we will define a new orthonormal basis of vectors ${\bf u} = [u_1,u_2]$ and ${\bf w} = [w_1,w_2]$. As we learned in the video, two vectors are orthonormal if:

1.   They are orthogonal (i.e., their dot product is zero):
\begin{equation}
{\bf u\cdot w} = u_1 w_1 + u_2 w_2 = 0
\end{equation}
2.   They have unit length:
\begin{equation}
||{\bf u} || = ||{\bf w} || = 1
\end{equation}

In two dimensions, it is easy to make an arbitrary orthonormal basis. All we need is a random vector ${\bf u}$, which we have normalized. If we now define the second basis vector to be ${\bf w} = [-u_2,u_1]$, we can check that both conditions are satisfied:
\begin{equation}
{\bf u\cdot w} = - u_1 u_2 + u_2 u_1 = 0
\end{equation}
and
\begin{equation}
{|| {\bf w} ||} = \sqrt{(-u_2)^2 + u_1^2} = \sqrt{u_1^2 + u_2^2} = 1,
\end{equation}
where we used the fact that ${\bf u}$ is normalized. So, with an arbitrary input vector, we can define an orthonormal basis, which we will write in matrix by stacking the basis vectors horizontally:

\begin{equation}
{{\bf W} } =
\begin{pmatrix}
 u_1 & w_1 \\
 u_2 & w_2
\end{pmatrix}.
\end{equation}

## Exercise 2: Find an orthonormal basis

In this exercise you will fill in the function below to define an orthonormal basis, given a single arbitrary 2-dimensional vector as an input.

**Steps**
* Modify the function `define_orthonormal_basis` to first normalize the first basis vector $\bf u$.
* Then complete the function by finding a basis vector $\bf w$ that is orthogonal to $\bf u$.
* Test the function using initial basis vector ${\bf u} = [3,1]$. Plot the resulting basis vectors on top of the data scatter plot using the function `plot_basis_vectors`. (For the data, use  $\sigma_1^2 =1$, $\sigma_2^2 =1$, and $\rho = .8$).

```python
help(plot_basis_vectors)
```

    Help on function plot_basis_vectors in module __main__:

    plot_basis_vectors(X, W)
        Plots bivariate data as well as new basis vectors.

        Args:
          X (numpy array of floats) :   Data matrix each column corresponds to a
                                        different random variable
          W (numpy array of floats) :   Square matrix representing new orthonormal
                                        basis each column represents a basis vector

        Returns:
          Nothing.

```python
def define_orthonormal_basis(u):
  """
  Calculates an orthonormal basis given an arbitrary vector u.

  Args:
    u (numpy array of floats) : arbitrary 2-dimensional vector used for new
                                basis

  Returns:
    (numpy array of floats)   : new orthonormal basis
                                columns correspond to basis vectors
  """

  #################################################
  ## TODO for students: calculate the orthonormal basis
  # Fill out function and remove
  raise NotImplementedError("Student excercise: implement the orthonormal basis function")
  #################################################

  # normalize vector u
  u = ...
  # calculate vector w that is orthogonal to w
  w = ...

  W = np.column_stack([u, w])

  return W

np.random.seed(2020)  # set random seed
variance_1 = 1
variance_2 = 1
corr_coef = 0.8

cov_matrix = calculate_cov_matrix(variance_1, variance_2, corr_coef)
X = get_data(cov_matrix)
u = np.array([3, 1])

# Uncomment and run below to plot the basis vectors
# W = define_orthonormal_basis(u)
# plot_basis_vectors(X, W)
```

[*Click for solution*](https://github.com/erlichlab/course-content/tree/master//tutorials/W1D5_DimensionalityReduction/solutions/W1D5_Tutorial1_Solution_7a9640ef.py)

*Example output:*

---
# Section 3: Project data onto new basis

```python
# @title Video 4: Change of basis
from IPython.display import IFrame
class BiliVideo(IFrame):
    def __init__(self, id, page=1, width=400, height=300, **kwargs):
        self.id=id
        src = "https://player.bilibili.com/player.html?bvid={0}&page={1}".format(id, page)
        super(BiliVideo, self).__init__(src, width, height, **kwargs)

video = BiliVideo(id='BV1LK411J7NQ', width=854, height=480, fs=1)
print("Video available at https://www.bilibili.com/video/{0}".format(video.id))
video
```

    Video available at https://www.bilibili.com/video/BV1LK411J7NQ

Finally, we will express our data in the new basis that we have just found. Since $\bf W$ is orthonormal, we can project the data into our new basis using simple matrix multiplication :

\begin{equation}
{\bf Y = X W}.
\end{equation}

We will explore the geometry of the transformed data $\bf Y$ as we vary the choice of basis.

## Exercise 3: Define an orthonormal basis
In this exercise you will fill in the function below to define an orthonormal basis, given a single arbitrary vector as an input.

**Steps**
* Complete the function `change_of_basis` to project the data onto the new basis.
* Plot the projected data using the function `plot_data_new_basis`.
* What happens to the correlation coefficient in the new basis? Does it increase or decrease?
* What happens to variance?

```python
def change_of_basis(X, W):
  """
  Projects data onto new basis W.

  Args:
    X (numpy array of floats) : Data matrix each column corresponding to a
                                different random variable
    W (numpy array of floats) : new orthonormal basis columns correspond to
                                basis vectors

  Returns:
    (numpy array of floats)    : Data matrix expressed in new basis
  """

  #################################################
  ## TODO for students: project the data onto o new basis W
  # Fill out function and remove
  raise NotImplementedError("Student excercise: implement change of basis")
  #################################################

  # project data onto new basis described by W
  Y = ...

  return Y

# Unomment below to transform the data by projecting it into the new basis
# Y = change_of_basis(X, W)
# plot_data_new_basis(Y)
```

[*Click for solution*](https://github.com/erlichlab/course-content/tree/master//tutorials/W1D5_DimensionalityReduction/solutions/W1D5_Tutorial1_Solution_a1124bbc.py)

*Example output:*

## Interactive Demo: Play with the basis vectors
To see what happens to the correlation as we change the basis vectors, run the cell below. The parameter $\theta$ controls the angle of $\bf u$ in degrees. Use the slider to rotate the basis vectors.

```python
# @title

# @markdown Make sure you execute this cell to enable the widget!

def refresh(theta=0):
  u = [1, np.tan(theta * np.pi / 180)]
  W = define_orthonormal_basis(u)
  Y = change_of_basis(X, W)
  plot_basis_vectors(X, W)
  plot_data_new_basis(Y)

_ = widgets.interact(refresh, theta=(0, 90, 5))
```

## Questions

* What happens to the projected data as you rotate the basis?
* How does the correlation coefficient change? How does the variance of the projection onto each basis vector change?
* Are you able to find a basis in which the projected data is **uncorrelated**?

---
# Summary

- In this tutorial, we learned that multivariate data can be visualized as a cloud of points in a high-dimensional vector space. The geometry of this cloud is shaped by the covariance matrix.

- Multivariate data can be represented in a new orthonormal basis using the dot product. These new basis vectors correspond to specific mixtures of the original variables - for example, in neuroscience, they could represent different ratios of activation  across a population of neurons.

- The projected data (after transforming into the new basis) will generally have a different geometry from the original data. In particular, taking basis vectors that are aligned with the spread of cloud of points decorrelates the data.

* These concepts - covariance, projections, and orthonormal bases - are key for understanding PCA, which we be our focus in the next tutorial.


Table of Contents

1  Instructions:

2  Things to Implement

3  Set Hyperparameters

4  Prepare the Dataset

5  Create a Neural Network

6  Train Classifier

# Optimization Assignment

## Instructions:

1. Clone this repository (or just pull changes if you already have it).
2. Start Jupyter (don't forget to activate conda).
3. Duplicate this file so that you can still pull changes without merging.
4. Complete the "Things to Implement."

## Things to Implement

1. Momentum

$$
\begin{align}
m_{t+1} &:= β_m m_t + (1 - β_m) \nabla_θ L_b(θ_t) \\
θ_{t+1} &:= θ_t - η m_{t+1}
\end{align}
$$

2. RMSProp

$$
\begin{align}
g_{t+1}^2 &:= β_g g_t^2 + (1 - β_g) (\nabla_θ L_b(θ_t))^2 \\
θ_{t+1} &:= θ_t - η \frac{\nabla_θ L_b(θ_t)}{\sqrt{g_{t+1}^2} + ε}
\end{align}
$$

3. Adam

$$
\begin{align}
m_{t+1} &:= β_m m_t + (1 - β_m) \nabla_θ L_b(θ_t) \\
\hat m_{t+1} &:= \frac{m_{t+1}}{1 - β_m^t} \\
g_{t+1}^2 &:= β_g g_t^2 + (1 - β_g) (\nabla_θ L_b(θ_t))^2 \\
\hat g_{t+1}^2 &:= \frac{g_{t+1}^2}{1 - β_g^t} \\
θ_{t+1} &:= θ_t - η \frac{\hat m_{t+1}}{\sqrt{\hat g_{t+1}^2} + ε}
\end{align}
$$

A few hints:

- Run the code all the way through without any changes and answer the first question on gradescope
- Adam combines momentum and RMSProp (and in this case add a bias correction)
- $t$ increments after each update (the actual value of $t$ is only used in Adam)
- You'll need to add code in two places
    1. At the top of the training cell (to initialize momentums and squared gradients)
    2. In the parameter update context manager (where you'll find `param -= ...`)
- The documentation for torch.optim.Adam will give good values for $\beta_m$ and $\beta_g$
- If momentum performs poorly, then it might be that you are not updating momentum values in-place

## Set Hyperparameters

```python
import torch
import torch.nn as nn
from torch.utils.data import DataLoader

from torchsummary import summary

from torchvision.datasets import FashionMNIST
from torchvision.transforms import Compose, Normalize, ToTensor

from fastprogress.fastprogress import master_bar, progress_bar

import matplotlib.pyplot as plt
from jupyterthemes import jtplot

jtplot.style(context="talk")
```

```python
# Let's use some shared space for the data (so that we don't have copies
# sitting around everywhere)
data_path = "/raid/cs152/cache/pytorch/data"

# Use the GPUs if they are available
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using '{device}' device.")

# Model hyperparameters
neurons_per_layer = [13, 17]

# Mini-Batch SGD hyperparameters
batch_size = 256
num_epochs = 10
learning_rate = 0.001

criterion = nn.CrossEntropyLoss()
```

## Prepare the Dataset

```python
def get_fmnist_data_loaders(path, batch_size, valid_batch_size=0):

    # Data specific transforms
    data_mean = (0.2860,)
    data_std = (0.3530,)
    xforms = Compose([ToTensor(), Normalize(data_mean, data_std)])

    # Training data loader
    train_dataset = FashionMNIST(root=path, train=True, download=True, transform=xforms)

    # Set the batch size to N if batch_size is 0
    tbs = len(train_dataset) if batch_size == 0 else batch_size
    train_loader = DataLoader(train_dataset, batch_size=tbs, shuffle=True)

    # Validation data loader
    valid_dataset = FashionMNIST(root=path, train=False, download=True, transform=xforms)

    # Set the batch size to N if batch_size is 0
    vbs = len(valid_dataset) if valid_batch_size == 0 else valid_batch_size
    valid_loader = DataLoader(valid_dataset, batch_size=vbs, shuffle=True)

    return train_loader, valid_loader
```

```python
# Computing normalization constants for Fashion-MNIST
# train_loader, valid_loader = get_fmnist_data_loaders(data_path, 0)
# X, _ = next(iter(train_loader))
# s, m = torch.std_mean(X)

train_loader, valid_loader = get_fmnist_data_loaders(data_path, batch_size)

print("Training dataset shape   :", train_loader.dataset.data.shape)
print("Validation dataset shape :", valid_loader.dataset.data.shape)
```

```python
# Let's plot a few images as an example
num_to_show = 8
images = train_loader.dataset.data[:num_to_show]
targets = train_loader.dataset.targets[:num_to_show]
labels = [train_loader.dataset.classes[t] for t in targets]

fig, axes = plt.subplots(1, num_to_show)

for axis, image, label in zip(axes, images, labels):
    axis.imshow(image.squeeze(), cmap="Greys")
    axis.tick_params(left=False, bottom=False, labelleft=False, labelbottom=False)
    axis.set_xticks([])
    axis.set_yticks([])
    axis.set_title(f"{label}")
```

## Create a Neural Network

```python
class NeuralNetwork(nn.Module):
    def __init__(self, layer_sizes):
        super(NeuralNetwork, self).__init__()

        # The first "layer" just rearranges the Nx28x28 input into Nx784
        first_layer = nn.Flatten()

        # The hidden layers include:
        # 1. a linear component (computing Z) and
        # 2. a non-linear comonent (computing A)
        hidden_layers = [
            nn.Sequential(nn.Linear(nlminus1, nl), nn.ReLU())
            for nl, nlminus1 in zip(layer_sizes[1:-1], layer_sizes)
        ]

        # The output layer must be Linear without an activation. See:
        #   https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html
        output_layer = nn.Linear(layer_sizes[-2], layer_sizes[-1])

        # Group all layers into the sequential container
        all_layers = [first_layer] + hidden_layers + [output_layer]
        self.layers = nn.Sequential(*all_layers)

    def forward(self, X):
        return self.layers(X)
```

```python
# The input layer size depends on the dataset
n0 = train_loader.dataset.data.shape[1:].numel()

# The output layer size depends on the dataset
nL = len(train_loader.dataset.classes)

# Preprend the input and append the output layer sizes
layer_sizes = [n0] + neurons_per_layer + [nL]
model = NeuralNetwork(layer_sizes).to(device)

summary(model);
```

## Train Classifier

```python
# Putting this here so that the model is recreated each time the cell is run
model = NeuralNetwork(layer_sizes).to(device)

t = 0

#
# TODO: Add your initialization code for momentum, RMSProp, and Adam
#
...

# A master bar for fancy output progress
mb = master_bar(range(num_epochs))

# Information for plots
mb.names = ["Train Loss", "Valid Loss"]
train_losses = []
valid_losses = []

for epoch in mb:

    #
    # Training
    #
    model.train()

    train_N = len(train_loader.dataset)
    num_train_batches = len(train_loader)
    train_dataiterator = iter(train_loader)

    train_loss_mean = 0

    for batch in progress_bar(range(num_train_batches), parent=mb):

        # Grab the batch of data and send it to the correct device
        train_X, train_Y = next(train_dataiterator)
        train_X, train_Y = train_X.to(device), train_Y.to(device)

        # Compute the output
        train_output = model(train_X)

        # Compute loss
        train_loss = criterion(train_output, train_Y)

        num_in_batch = len(train_X)
        tloss = train_loss.item() * num_in_batch / train_N
        train_loss_mean += tloss
        train_losses.append(train_loss.item())

        # Compute partial derivatives
        model.zero_grad()
        train_loss.backward()

        #
        # TODO: Add changes for momentum, RMSProp, and Adam
        #
        t += 1
        with torch.no_grad():
            # Original gradient descent
            for param in model.parameters():
                # θ_{t+1} := θ_t - η \nabla_θ L_b(θ_t)
                param -= learning_rate * param.grad

            # Gradient descent with momentum
            # for param, momentum in ...
                # m_{t+1} := β_m m_t + (1 - β_m) \nabla_θ L_b(θ_t)
                # θ_{t+1} := θ_t - η m_{t+1}
                ...

            # Gradient descent with RMSProp
            # for param, sq_grad in ...
                # g_{t+1}^2 := β_g g_t^2 + (1 - β_g) (\nabla_θ L_b(θ_t))^2
                # θ_{t+1} := θ_t - η \frac{\nabla_θ L_b(θ_t)}{\sqrt{g_{t+1}^2} + ε}
                ...

            # Gradient descent with Adam
            # for param, momentum, sq_grad in ...
                # m_{t+1} := β_m m_t + (1 - β_m) \nabla_θ L_b(θ_t)
                # \hat m_{t+1} := \frac{m_{t+1}}{1 - β_m^t}
                # g_{t+1}^2 := β_g g_t^2 + (1 - β_g) (\nabla_θ L_b(θ_t))^2
                # \hat g_{t+1}^2 := \frac{g_{t+1}^2}{1 - β_g^t}
                # θ_{t+1} := θ_t - η \frac{\hat m_{t+1}}{\sqrt{\hat g_{t+1}^2} + ε}
                ...

    #
    # Validation
    #
    model.eval()

    valid_N = len(valid_loader.dataset)
    num_valid_batches = len(valid_loader)

    valid_loss_mean = 0
    valid_correct = 0

    with torch.no_grad():

        # valid_loader is probably just one large batch, so not using progress bar
        for valid_X, valid_Y in valid_loader:

            valid_X, valid_Y = valid_X.to(device), valid_Y.to(device)

            valid_output = model(valid_X)

            valid_loss = criterion(valid_output, valid_Y)

            num_in_batch = len(valid_X)
            vloss = valid_loss.item() * num_in_batch / valid_N
            valid_loss_mean += vloss
            valid_losses.append(valid_loss.item())

            # Convert network output into predictions (one-hot -> number)
            predictions = valid_output.argmax(1)

            # Sum up total number that were correct
            valid_correct += (predictions == valid_Y).type(torch.float).sum().item()

    valid_accuracy = 100 * (valid_correct / valid_N)

    # Report information
    tloss = f"Train Loss = {train_loss_mean:.4f}"
    vloss = f"Valid Loss = {valid_loss_mean:.4f}"
    vaccu = f"Valid Accuracy = {(valid_accuracy):>0.1f}%"
    mb.write(f"[{epoch+1:>2}/{num_epochs}] {tloss}; {vloss}; {vaccu}")

    # Update plot data
    max_loss = max(max(train_losses), max(valid_losses))
    min_loss = min(min(train_losses), min(valid_losses))

    x_margin = 0.2
    x_bounds = [0 - x_margin, num_epochs + x_margin]

    y_margin = 0.1
    y_bounds = [min_loss - y_margin, max_loss + y_margin]

    valid_Xaxis = torch.linspace(0, epoch + 1, len(train_losses))
    valid_xaxis = torch.linspace(1, epoch + 1, len(valid_losses))
    graph_data = [[valid_Xaxis, train_losses], [valid_xaxis, valid_losses]]

    mb.update_graph(graph_data, x_bounds, y_bounds)
```

```python

```


# Applied Stochastic Processes HW02

**11510691 程远$\DeclareMathOperator*{\argmin}{argmin}
\newcommand{\using}[1]{\stackrel{\mathrm{#1}}{=}}
\newcommand{\ffrac}{\displaystyle \frac}
\newcommand{\space}{\text{ }}
\newcommand{\bspace}{\;\;\;\;}
\newcommand{\QQQ}{\boxed{?\:}}
\newcommand{\CB}[1]{\left\{ #1 \right\}}
\newcommand{\SB}[1]{\left[ #1 \right]}
\newcommand{\P}[1]{\left( #1 \right)}
\newcommand{\dd}{\mathrm{d}}
\newcommand{\Tran}[1]{{#1}^{\mathrm{T}}}
\newcommand{\d}[1]{\displaystyle{#1}}
\newcommand{\EE}[2][\,\!]{\mathbb{E}_{#1}\left[#2\right]}
\newcommand{\Var}[2][\,\!]{\mathrm{Var}_{#1}\left[#2\right]}
\newcommand{\Cov}[2][\,\!]{\mathrm{Cov}_{#1}\left(#2\right)}
\newcommand{\Corr}[2][\,\!]{\mathrm{Corr}_{#1}\left(#2\right)}
\newcommand{\I}[1]{\mathrm{I}\left( #1 \right)}
\newcommand{\N}[1]{\mathrm{N} \left( #1 \right)}
\newcommand{\ow}{\text{otherwise}}$星**

## Question 1

$\bspace \begin{align}
\EE{X \mid X > 1} &= \int_{0}^{\infty} x \cdot f_X\P{x \mid X > 1} \;\dd{x} \\
&= \int_{1}^{\infty} x \cdot \lambda e^{-\lambda x} \;\dd{x} \\
&= \left.-\P{x + \ffrac{1} {\lambda}}e^{-\lambda x}\right|_{1}^{\infty} = \P{1 + \ffrac{1} {\lambda}}e^{-\lambda}
\end{align}$

## Question 2

$\P{\text a}$

$\bspace \begin{align}
f_Y\P{y} &= \int_{-\infty}^{\infty} f\P{x,y} \;\dd{x} \\
&= \int_{0}^{\infty} \ffrac{e^{-x/y} e^{-y}} {y} \;\dd{x} \\
&= \ffrac{1} {y} e^{-y} \cdot \left. \P{-y\exp\P{-\ffrac{1} {y}x}}\right|_{0}^{\infty} \\
&\using{y>0} \ffrac{1} {y} e^{-y} \cdot y = e^{-y}, y > 0
\end{align}$

$\bspace$ and $f_Y\P{y} = 0$ otherwise.

$\P{\text b}$

$\bspace \begin{align}
\EE{X \mid Y = y} &= \int_{-\infty}^{\infty} x \cdot f_X\P{x \mid Y = y}\;\dd{x} \\
&= \int_{0}^{\infty} x \cdot \ffrac{f_{X,Y}\P{x,y}} {f_Y\P{y}}\;\dd{x} \\
&= \int_{0}^{\infty} x \cdot \ffrac{e^{-x/y} e^{-y}} {y} \cdot {e^{y}}\;\dd{x} \\
&= \left.-\P{x+y} \exp\P{-\ffrac{1} {y} x}\right|_{0}^{\infty} \using{y>0} y
\end{align}$

$\P{\text C}$

$\bspace \begin{align}
\EE{X} &= \EE{\EE{X\mid Y}} \\
&= \EE{Y} \\
&= \int_{-\infty}^{\infty} y\cdot f_Y\P{y} \;\dd{y} \\
&= \int_{0}^{\infty} y \cdot e^{-y} \;\dd{y} \\
&= \Big.-\P{y+1}e^{-y}\Big|_{0}^{\infty} = 1
\end{align}$

$\bspace$ And to calculate its variance we first condition $X^2$ on $Y$, giving:

$\bspace \begin{align}
\EE{X^2 \mid Y = y} &= \int_{0}^{\infty} x^2 \cdot \ffrac{e^{-x/y} e^{-y}} {y} \cdot {e^{y}}\;\dd{x} \\
&= \left.-\P{x^2+2yx + 2y^2} \exp\P{-\ffrac{1} {y} x}\right|_{0}^{\infty} \using{y>0} 2y^2
\end{align}$

Thus we can compute $\EE{X^2}$ like:

$\bspace \begin{align}
\EE{X^2} &= \EE{\EE{X^2\mid Y}} \\
&= \EE{2Y^2} \\
&= \int_{-\infty}^{\infty} 2y^2\cdot f_Y\P{y} \;\dd{y} \\
&= 2 \int_{0}^{\infty} y^2 \cdot e^{-y} \;\dd{y} \\
&= \Big.-\P{2y^2+ 4y + 4}e^{-y}\Big|_{0}^{\infty} = 4
\end{align}$

$\bspace \Var{X} = \EE{X^2} - \P{\EE{X}}^2 = 4 - 1^2 = 3$

## Question 3

$\bspace \begin{align}
\text{RHS} &= \Big(\EE{\EE{\P{X - \EE{X \mid Y}}^2 \mid Y}}\Big) + \Big(\EE{\P{\EE{X \mid Y}}^2} - \P{\EE{\EE{X \mid Y}}}^2 \Big) \\
&= \EE{\P{X - \EE{X \mid Y}}^2} + \EE{\P{\EE{X \mid Y}}^2} - \P{\EE{X}}^2 \\
&= \EE{X^2} - 2\cdot \EE{X \cdot \EE{X \mid Y}} + 2 \cdot \EE{\P{\EE{X \mid Y}}^2} - \P{\EE{X}}^2 \\
&= \EE{X^2} - \P{\EE{X}}^2 - 2 \cdot \EE{\EE{X \cdot \EE{X \mid Y} \mid Y}}+ 2 \cdot \EE{\P{\EE{X \mid Y}}^2} \\
&= \Var{X} - 2 \cdot \EE{\P{\EE{X \mid Y}}^2} + 2 \cdot \EE{\P{\EE{X \mid Y}}^2} = \Var{X} = \text{LHS}
\end{align}$

## Question 4

$\P{\text a}$

$\bspace \begin{align}
X &= \sum_{i=1}^{N}T_i
\end{align}$

$\P{\text b}$

$\bspace$ To find $\EE{N}$ we condition that on $Y$, the first choice:

$\bspace \begin{align}
\EE{N} &= \EE{N \mid Y = 1} \cdot P\CB{Y = 1} + \EE{N \mid Y = 2} \cdot P\CB{Y = 2} + \EE{N \mid Y = 3} \cdot P\CB{Y = 3} \\
&= \ffrac{1} {3} \P{1 + 1 + \EE{N} + 1 + \EE{N}}
\end{align}$

$\bspace$ Then we solve the preceding equation and find the solution that tis $\EE{N} = 3$

$\P{\text c}$

$\bspace$ Since $N$ is the total number of choices, or equivalently, there's no $N+1$ choice. Thus $T_N \equiv 2$ and $\EE{T_N} = 2$.

$\P{\text d}$

$\bspace \begin{align}
\EE{\sum_{i=1}^{N}T_i \mid N = n} &= \EE{\sum_{i=1}^{n-1}T_i \mid N = n} + \EE{T_n \mid N = n} \\
&= \EE{\sum_{i=1}^{n-1}T_i} + \EE{T_n} \\
&= \P{n-1} \cdot \ffrac{1} {2} \P{3+5} + 2 = 4n - 2
\end{align}$

$\P{\text e}$

$\bspace \begin{align}
\EE{X} &= \EE{\EE{\sum_{i=1}^{N}T_i \mid N} } \\
&= \EE{4N - 2} \\
&= 4\cdot \EE{N} - 2 = 4 \times 3 - 2 = 10
\end{align}$

## Question 5

$\P{\text a}$

$\bspace \begin{align}
\EE{X_1} = 1 \times \ffrac{r} {r+b} + 0 \times \ffrac{b} {r+b} = \ffrac{r} {r+b}
\end{align}$

$\P{\text b}$

$\bspace$ First we let $Y_i = 1$ if the $i\texttt{-th}$ draw is a red one and $0$ otherwise.

$\bspace \begin{align}
\EE{X_2} &= \EE{X_2 \mid Y_1 = 1} \cdot P\CB{Y_1 = 1} + \EE{X_2 \mid Y_1 = 0} \cdot P\CB{Y_1 = 0} \\
&= \P{2\times \ffrac{r+m} {r+b+m} \times \ffrac{r} {r+b} + \ffrac{b} {r+b + m} \cdot \ffrac{r} {r+b}} + \ffrac{r} {r+b + m} \cdot \ffrac{b} {r+b} = \ffrac{2r} {r+b}
\end{align}$

$\P{\text c}$

$\bspace \begin{align}
\EE{X_3} &= \EE{\EE{\EE{X_3 \mid Y_2}}\mid Y_1} \\
&= \EE{ \EE{X_3 \mid Y_2 = 1}\cdot P\CB{Y_2 = 1} + \EE{X_3 \mid Y_2 = 0}\cdot P\CB{Y_2 = 0}\mid Y_1} \\
&= \EE{X_3 \mid Y_2 = 1, Y_1 = 1} \cdot P\CB{Y_2 = 1 \mid Y_1 = 1} \cdot P\CB{Y_1 = 1} \\
&\bspace + \EE{X_3 \mid Y_2 = 1, Y_1 = 0} \cdot P\CB{Y_2 = 1 \mid Y_1 = 0} \cdot P\CB{Y_1 = 0} \\
&\bspace + \EE{X_3 \mid Y_2 = 0, Y_1 = 1} \cdot P\CB{Y_2 = 0 \mid Y_1 = 1} \cdot P\CB{Y_1 = 1} \\
&\bspace + \EE{X_3 \mid Y_2 = 0, Y_1 = 0} \cdot P\CB{Y_2 = 0 \mid Y_1 = 0} \cdot P\CB{Y_1 = 0} \\
&= \P{3 \times \ffrac{r+2m} {r+b+2m} + 2 \times \ffrac{b} {r+b+2m}}\cdot \ffrac{r+m} {r+b+m} \cdot \ffrac{r} {r+b} \\
&\bspace + \P{2\times\ffrac{r+m} {r+b+2m} + \ffrac{b+m} {r+b+2m}} \cdot \ffrac{r} {r+b+m} \cdot \ffrac{b} {r+b} \\
&\bspace + \P{2\times\ffrac{r+m} {r+b+2m} + \ffrac{b+m} {r+b+2m}} \cdot \ffrac{b} {r+b+m} \cdot \ffrac{r} {r+b} \\
&\bspace + \ffrac{r} {r+b+2m} \cdot \ffrac{b+m} {r+b+m} \cdot \ffrac{b} {r+b} \\
&= \ffrac{3r} {r+b}
\end{align}$

$\P{\text d}$

$\bspace$ Based on the preceding results, I conjecture that $\EE{X_k} = \ffrac{kr} {r+b}$ for $k = 1, 2, \dots$. To prove this we first let $r_i$ and $b_i$ denote the number of red balls and blue balls at $i\texttt{-th}$ draw. Then $r_1 = r$ and $b_1 = b$; by induction, we let this equation holds for $k =1,2,\dots,n-1$ and we condtion $X_n$ on $Y_{n-1}$ finding that:

$\bspace \begin{align}
\EE{X_n} &= \EE{\EE{X_n \mid Y_{n-1}}} \\
&= \EE{X_n \mid Y_{n-1} = 1} \cdot P\CB{Y_{n-1} = 1} + \EE{X_n \mid Y_{n-1} = 0} \cdot P\CB{Y_{n-1} = 0} \\
&= \P{\P{2 + \ffrac{\P{k-2}r} {r+b}}\cdot\ffrac{r_{n-1} + m} {r_{n-1} + b_{n-1} + m} + \P{1+\ffrac{\P{k-2}r} {r+b}} \cdot \ffrac{b_{n-1}} {r_{n-1} + b_{n-1} + m}} \cdot \ffrac{r_{n-1}} {r_{n-1} + b_{n-1}} \\
&\bspace + \P{\P{1+\ffrac{\P{k-2}r} {r+b}} \cdot \ffrac{r_{n-1}} {r_{n-1} + b_{n-1} + m} + \ffrac{\P{k-2}r} {r+b} \cdot \ffrac{b_{n-1} + m} {r_{n-1} + b_{n-1} + m}} \cdot \ffrac{b_{n-1}} {r_{n-1} + b_{n-1}} \\
&= \ffrac{nr} {r+b}
\end{align}$

$\bspace$ solved...😜😜😜

$\P{\text e}$

$\bspace$ To be intuitive, honestly, my first impression on this problem is:

$$\sum_{i=0}^{n} \P{\binom{n} {i} \cdot i \cdot \ffrac{\d{\prod_{j=0}^{i-1}\SB{r+jm}}\d{\prod_{j=0}^{n-i-1} \SB{b+jm}}} {\d{\prod_{k=0}^{n-1}\SB{r+b+km}}}}$$

$\bspace$ That's not something that's gonna help, so another thoughts, breaking $X_n$ apart into $Y_i$. After $\EE{Y_2}$ is given, since $\EE{Y_2} = \EE{Y_1}$, an intuitive thoughts could be that at this time the urn is in the same "**state**" with before. So that $\EE{Y_{n-1}} = \EE{Y_n}$ will always hold and thus $\EE{X_n} = \sum_n \EE{Y_i} = n\ffrac{r} {r+b}$

Probability and Statistics/Applied Random Process/HW/HW_02.ipynb


# Unrestricted Hartree-Fock

**This assignment is due on Friday, January 31 at 5PM by email.**

In carrying out the implementation of the closed-shell Hartree-Fock method we have, in fact, had to solve most of the design problems associated with open-shell Hartree-Fock.

\begin{align}
E_{elec}^{UHF} =&\ \sum_i^{n_\alpha} \ h_{ii}^{\alpha} + \sum_i^{n_\beta} \ h_{ii}^{\beta} \\
               &+ \sum_{i<j}^{n_\alpha} \biggl\{ (i^\alpha i^\alpha | j^\alpha j^\alpha) - (i^\alpha j^\alpha | i^\alpha j^\alpha) \biggr\} + \sum_{i<j}^{n_\beta} \biggl\{ (i^\beta i^\beta | j^\beta j^\beta) - (i^\beta j^\beta | i^\beta j^\beta) \biggr\} \\
               &+ \sum_i^{n_\alpha} \sum_j^{n_\beta} (i^\alpha i^\alpha | j^\beta j^\beta)\\
               =&\ \sum_i^{n_\alpha} \ h_{ii}^{\alpha} + \sum_i^{n_\beta} \ h_{ii}^{\beta} \\
               & + \sum_{i<j}^{n_\alpha} \biggl\{ J_{ij}^{\alpha} - K_{ij}^{\alpha} \biggr\} + \sum_{i<j}^{n_\beta} \biggl\{ J_{ij}^{\beta} - K_{ij}^{\beta} \biggr\} \\
               & + \sum_i^{n_\alpha} \sum_j^{n_\beta} J_{ij}^{\alpha\beta} \\
\end{align}

The density matrix
\begin{align}
D_{\mu\nu}^\alpha =&\ \sum_i^{n_{\alpha}} C_{\mu i}^\alpha C_{\nu i}^\alpha \\
D_{\mu\nu}^\beta =&\ \sum_i^{n_{\beta}} C_{\mu i}^\beta C_{\nu i}^\beta
\end{align}

### One-electron energy
\begin{align}
E_{one} =&\ \sum_i^{n_\alpha} h_{ii}^\alpha + \sum_i^{n_\beta} h_{ii}^\beta \\
        =&\ \sum_i^{n_\alpha} \biggl\{ \sum_{\mu \nu}^{AO} C_{\mu i}^\alpha C_{\nu i}^\alpha h_{\mu \nu} \biggr\} + \sum_i^{n_\beta} \biggl\{ \sum_{\mu \nu}^{AO} C_{\mu i}^\beta C_{\nu i}^\beta h_{\mu \nu} \biggr\} \\
        =&\ \sum_{\mu \nu}^{AO} \biggl\{ \sum_i^{n_\alpha} C_{\mu i}^\alpha C_{\nu i}^\alpha + C_{\mu i}^\beta C_{\nu i}^\beta \biggr\} h_{\mu \nu} \\
        =&\ \sum_{\mu \nu}^{AO} \biggl\{ D_{\mu \nu}^\alpha + D_{\mu \nu}^\beta \biggr\} h_{\mu \nu}
\end{align}
Notice that the one-electron energy is composed of an $\alpha$ and $\beta$ energy.

### Two-electron energy
\begin{align}
E_{two} =&\ \sum_{i<j}^{n_\alpha} \biggl\{ J_{ij}^{\alpha} - K_{ij}^{\alpha} \biggr\} + \sum_{i<j}^{n_\beta} \biggl\{ J_{ij}^{\beta} - K_{ij}^{\beta} \biggr\} + \sum_i^{n_\alpha} \sum_j^{n_\beta} J_{ij}^{\alpha\beta} + \sum_i^{n_\beta} \sum_j^{n_\alpha} J_{ij}^{\beta\alpha}  \\
\\
        =&\ \sum_{i<j}^{n_\alpha} \Biggl[ \sum_{\mu\nu\rho\sigma}^{AO} \biggl\{ C_{\mu i}^\alpha C_{\nu i}^\alpha C_{\rho j}^\alpha C_{\sigma j}^\alpha (\mu\nu | \rho\sigma) - C_{\mu i}^\alpha C_{\rho j}^\alpha C_{\nu i}^\alpha C_{\sigma j}^\alpha (\mu\rho | \nu\sigma)\biggr\} \Biggr] \\
        +&\ \sum_{i<j}^{n_\beta} \Biggl[ \sum_{\mu\nu\rho\sigma}^{AO} \biggl\{ C_{\mu i}^\beta C_{\nu i}^\beta C_{\rho j}^\beta C_{\sigma j}^\beta (\mu\nu | \rho\sigma) - C_{\mu i}^\beta C_{\rho j}^\beta C_{\nu i}^\beta C_{\sigma j}^\beta (\mu\rho | \nu\sigma)\biggr\} \Biggr] \\
        +&\ \sum_i^{n_\alpha} \sum_j^{n_\beta} \Biggl[ \sum_{\mu\nu\rho\sigma}^{AO} \biggl\{ C_{\mu i}^\alpha C_{\nu i}^\alpha C_{\rho j}^\beta C_{\sigma j}^\beta (\mu\nu | \rho\sigma) \biggr\} \Biggr] \\
        +&\ \sum_i^{n_\beta} \sum_j^{n_\alpha} \Biggl[ \sum_{\mu\nu\rho\sigma}^{AO} \biggl\{ C_{\mu i}^\beta C_{\nu i}^\beta C_{\rho j}^\alpha C_{\sigma j}^\alpha (\mu\nu | \rho\sigma) \biggr\} \Biggr] \\
\\
        =&\ \sum_{\mu\nu\rho\sigma}^{AO} \Biggl[ \biggl\{ \sum_i^{n_\alpha} C_{\mu i}^\alpha C_{\nu i}^\alpha \biggr\} \biggl\{ \sum_j^{n_\alpha} C_{\rho j}^\alpha C_{\sigma j}^\alpha \biggr\} \biggl\{ (\mu\nu|\rho\sigma) - (\mu\rho|\nu\sigma) \biggr\} \Biggr] \\
        +&\ \sum_{\mu\nu\rho\sigma}^{AO} \Biggl[ \biggl\{ \sum_i^{n_\beta} C_{\mu i}^\beta C_{\nu i}^\beta \biggr\} \biggl\{ \sum_j^{n_\beta} C_{\rho j}^\beta C_{\sigma j}^\beta \biggr\} \biggl\{ (\mu\nu|\rho\sigma) - (\mu\rho|\nu\sigma) \biggr\} \Biggr] \\
        +&\ \sum_{\mu\nu\rho\sigma}^{AO} \Biggl[ \biggl\{ \sum_i^{n_\alpha} C_{\mu i}^\alpha C_{\nu i}^\alpha \biggr\} \biggl\{ \sum_j^{n_\beta} C_{\rho j}^\beta C_{\sigma j}^\beta \biggr\} (\mu\nu|\rho\sigma) \Biggr] \\
        +&\ \sum_{\mu\nu\rho\sigma}^{AO} \Biggl[ \biggl\{ \sum_i^{n_\beta} C_{\mu i}^\beta C_{\nu i}^\beta \biggr\} \biggl\{ \sum_j^{n_\alpha} C_{\rho j}^\alpha C_{\sigma j}^\alpha \biggr\} (\mu\nu|\rho\sigma) \Biggr] \\
\\
        =&\ \sum_{\mu\nu\rho\sigma}^{AO} D_{\mu\nu}^\alpha D_{\rho\sigma}^\alpha \biggl\{ (\mu\nu|\rho\sigma) - (\mu\rho|\nu\sigma) \biggr\} \\
        +&\ \sum_{\mu\nu\rho\sigma}^{AO} D_{\mu\nu}^\beta D_{\rho\sigma}^\beta \biggl\{ (\mu\nu|\rho\sigma) - (\mu\rho|\nu\sigma) \biggr\} \\
        +&\ \sum_{\mu\nu\rho\sigma}^{AO} D_{\mu\nu}^\alpha D_{\rho\sigma}^\beta (\mu\nu|\rho\sigma) \\
        +&\ \sum_{\mu\nu\rho\sigma}^{AO} D_{\mu\nu}^\beta D_{\rho\sigma}^\alpha (\mu\nu|\rho\sigma) \\
\end{align}

### Fock Matrix
\begin{align}
F_{\mu\nu}^{\alpha} =&\ h_{\mu\nu}^\alpha + \sum_{\rho\sigma}^{AO} D_{\rho\sigma}^\alpha \biggl\{ (\mu\nu|\rho\sigma) - (\mu\rho|\nu\sigma) \biggr\}  + D_{\rho\sigma}^\beta (\mu\nu|\rho\sigma) \\
F_{\mu\nu}^{\alpha} =&\ h_{\mu\nu}^\alpha + J_{\mu\nu}^{\alpha} - K_{\mu\nu}^\alpha + J_{\mu\nu}^\beta\\
\\
F_{\mu\nu}^{\beta} =&\ h_{\mu\nu}^\beta + \sum_{\rho\sigma}^{AO} D_{\rho\sigma}^\beta \biggl\{ (\mu\nu|\rho\sigma) - (\mu\rho|\nu\sigma) \biggr\}  + D_{\rho\sigma}^\alpha (\mu\nu|\rho\sigma) \\
F_{\mu\nu}^{\beta} =&\ h_{\mu\nu}^\beta + J_{\mu\nu}^{\beta} - K_{\mu\nu}^\beta + J_{\mu\nu}^\alpha
\end{align}

### Electronic Energy
\begin{align}
E_{elec} = \frac{1}{2}\Biggl[ \sum_{\mu \nu}^{AO} \biggl\{ D_{\mu\nu}^\alpha + D_{\mu\nu}^\beta \biggr\} h_{\mu\nu} + \sum_{\mu \nu}^{AO} D_{\mu\nu}^\alpha F_{\mu\nu}^\alpha + \sum_{\mu\nu}^{AO} D_{\mu\nu}^\beta F_{\mu\nu}^{\beta} \Biggr]
\end{align}
The $\tfrac{1}{2}$ is to correct for double counting.

Additional insights into UHF can be found in Modern Quantum Chemistry by Szabo and Ostlund at Chapter 3 Section 8.

### Spin Contamination Metric

\begin{align}
S_{i j}^{\mathrm{metric}} =& C_{\mu i}^{\alpha} S_{\mu \nu} C_{\nu j}^\beta \\
S^{\mathrm{contamination}} =& \mathrm{min}(n_{\alpha}, n_{\beta}) - (S^{\mathrm{metric}} \cdot S^{\mathrm{metric}})
\end{align}

### Procedure
For this assignment, please create a separate `.py` file as this will form the base for the remaining assignments.

1. Load Psi4 and NumPy
2. Setup your initial conditions
   * Construct a new molecule in C1 symmetry
   * Load in the basis set
   * Create a `MintsHelper` object to compute the integrals.
3. Compute needed integrals
4. Form the one-electron Hamiltonian.
5. Construct the orthogonalizer $S^{-1/2}$
6. Construct the initial density matrices $D^\alpha$ and $D^\beta$
   * Form the initial (transformed) Fock matrices using a core Hamiltonian guess
   * Diagonalize the Fock matrices
   * Form the SCF eigenvectors $C^\alpha$ and $C^\beta$
   * Form the initial density matrices $D^\alpha$ and $D^\beta$
7. Perform SCF iterations until convergence is achieved
   * Construct new Fock matrices
   * Calculate this iteration's electronic and total energies
   * Transform the Fock matrices to the orthonormal basis
   * Diagonalize the Fock matrices
   * Construct the new SCF eigenvector matrices
   * Form the new density matrices
   * Test convergency of the density matrices and energy
8. Compute and print the spin-contamination.

My output:

    UHF iteration   1: energy   -73.50799290557293  dE -7.35080E+01
    UHF iteration   2: energy   -74.66361368523238  dE -1.15562E+00
    UHF iteration   3: energy   -74.66614383044545  dE -2.53015E-03
    UHF iteration   4: energy   -74.66638263555237  dE -2.38805E-04
    UHF iteration   5: energy   -74.66644299107340  dE -6.03555E-05
    UHF iteration   6: energy   -74.66646469386085  dE -2.17028E-05
    UHF iteration   7: energy   -74.66647352191376  dE -8.82805E-06
    UHF iteration   8: energy   -74.66647726970710  dE -3.74779E-06
    UHF iteration   9: energy   -74.66647888623531  dE -1.61653E-06
    UHF iteration  10: energy   -74.66647958779794  dE -7.01563E-07
    UHF iteration  11: energy   -74.66647989301480  dE -3.05217E-07
    UHF iteration  12: energy   -74.66648002592984  dE -1.32915E-07
    UHF iteration  13: energy   -74.66648008383399  dE -5.79041E-08
    UHF iteration  14: energy   -74.66648010906377  dE -2.52298E-08
    UHF iteration  15: energy   -74.66648012005750  dE -1.09937E-08
    UHF iteration  16: energy   -74.66648012484804  dE -4.79054E-09
    UHF iteration  17: energy   -74.66648012693554  dE -2.08750E-09
    UHF iteration  18: energy   -74.66648012784520  dE -9.09665E-10
    UHF iteration  19: energy   -74.66648012824157  dE -3.96369E-10
    UHF iteration  20: energy   -74.66648012841434  dE -1.72761E-10
    UHF iteration  21: energy   -74.66648012848960  dE -7.52607E-11

    Spin Contamination Metric: 6.40288E-03

### Testing against Psi4

```python
import psi4

psi4.set_options({'basis': 'sto-3g',
                 'scf_type': 'pk',
                 'e_convergence': 1e-10,
                 'reference': 'uhf'})
psi4.energy('scf')
```

```python
import numpy as np
from openshell import Settings

np.set_printoptions(suppress=True, linewidth=200)
```

```python
molecule = psi4.geometry(Settings['molecule'])
molecule.update_geometry()
```

```python
nalpha = Settings['nalpha']
nbeta = Settings['nbeta']

scf_max_iter = Settings['scf_max_iter']

basis = psi4.core.BasisSet.build(molecule, 'BASIS', Settings['basis'])
mints = psi4.core.MintsHelper(basis)
```

```python
S = mints.ao_overlap().np
T = mints.ao_kinetic().np
V = mints.ao_potential().np
I = mints.ao_eri().np

H = T + V
print(H)
```

    [[-32.67360938  -7.60160704  -0.01741536   0.           0.          -1.57818038  -1.57818038]
     [ -7.60160704  -9.2914758   -0.20781331   0.          -0.          -3.47334591  -3.47334591]
     [ -0.01741536  -0.20781331  -7.50695983   0.           0.          -1.55452724  -1.55452724]
     [  0.           0.           0.          -7.42214005   0.           0.           0.        ]
     [  0.          -0.           0.           0.          -7.56362101   1.92474988  -1.92474988]
     [ -1.57818038  -3.47334591  -1.55452724   0.           1.92474988  -4.90357867  -1.42265626]
     [ -1.57818038  -3.47334591  -1.55452724   0.          -1.92474988  -1.42265626  -4.90357867]]

```python
from scipy.linalg import fractional_matrix_power

A = fractional_matrix_power(S, -0.5)
print(A)
```

    [[ 1.02403592 -0.14128825 -0.00938732  0.          0.          0.02117023  0.02117023]
     [-0.14128825  1.21591378  0.09833518  0.          0.         -0.27163956 -0.27163956]
     [-0.00938732  0.09833518  1.04984867  0.         -0.         -0.1381764  -0.1381764 ]
     [ 0.          0.          0.          1.          0.          0.          0.        ]
     [-0.         -0.          0.          0.          1.10025748  0.2134062  -0.2134062 ]
     [ 0.02117023 -0.27163956 -0.1381764   0.          0.2134062   1.1823526  -0.08078394]
     [ 0.02117023 -0.27163956 -0.1381764   0.         -0.2134062  -0.08078394  1.1823526 ]]

```python
Fa = A.dot(H).dot(A)
Fb = A.dot(H).dot(A)

_, Ca = np.linalg.eigh(Fa)
_, Cb = np.linalg.eigh(Fb)

Ca = A.dot(Ca)
Cb = A.dot(Cb)

Cocca = Ca[:, :nalpha]
Coccb = Cb[:, :nbeta]

Da = np.einsum('pi,qi->pq', Cocca, Cocca)
Db = np.einsum('pi,qi->pq', Coccb, Coccb)

print(Da)
print(Db)
```

    [[ 1.06661727 -0.29755998 -0.02536219 -0.         -0.          0.04007352  0.04007352]
     [-0.29755998  1.32589965  0.15348429  0.          0.         -0.17684186 -0.17684186]
     [-0.02536219  0.15348429  1.08573619  0.          0.         -0.11682694 -0.11682694]
     [-0.          0.          0.          1.         -0.         -0.          0.        ]
     [-0.          0.          0.         -0.          1.17091368  0.1776736  -0.1776736 ]
     [ 0.04007352 -0.17684186 -0.11682694 -0.          0.1776736   0.05924138  0.00532125]
     [ 0.04007352 -0.17684186 -0.11682694  0.         -0.1776736   0.00532125  0.05924138]]
    [[ 1.06661727 -0.29755998 -0.02536219 -0.         -0.          0.04007352  0.04007352]
     [-0.29755998  1.32589965  0.15348429  0.          0.         -0.17684186 -0.17684186]
     [-0.02536219  0.15348429  1.08573619 -0.          0.         -0.11682694 -0.11682694]
     [-0.          0.         -0.          0.         -0.         -0.          0.        ]
     [-0.          0.          0.         -0.          1.17091368  0.1776736  -0.1776736 ]
     [ 0.04007352 -0.17684186 -0.11682694 -0.          0.1776736   0.05924138  0.00532125]
     [ 0.04007352 -0.17684186 -0.11682694  0.         -0.1776736   0.00532125  0.05924138]]

```python

```


# Neuromatch Academy: Week 1, Day 5, Tutorial 1
# Dimensionality Reduction: Geometric view of data

__Content creators:__ Alex Cayco Gajic, John Murray

__Content reviewers:__ Roozbeh Farhoudi, Matt Krause, Spiros Chavlis, Richard Gao, Michael Waskom

---
# Tutorial Objectives

In this notebook we'll explore how multivariate data can be represented in different orthonormal bases. This will help us build intuition that will be helpful in understanding PCA in the following tutorial.

Overview:
 - Generate correlated multivariate data.
 - Define an arbitrary orthonormal basis.
 - Project the data onto the new basis.

```python
# @title Video 1: Geometric view of data
from IPython.display import YouTubeVideo
video = YouTubeVideo(id="THu9yHnpq9I", width=854, height=480, fs=1)
print("Video available at https://youtube.com/watch?v=" + video.id)
video
```

    Video available at https://youtube.com/watch?v=THu9yHnpq9I

---
# Setup

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

```python
# @title Figure Settings
import ipywidgets as widgets  # interactive display
%config InlineBackend.figure_format = 'retina'
plt.style.use("https://raw.githubusercontent.com/NeuromatchAcademy/course-content/master/nma.mplstyle")
```

```python
# @title Helper functions

def get_data(cov_matrix):
  """
  Returns a matrix of 1000 samples from a bivariate, zero-mean Gaussian.

  Note that samples are sorted in ascending order for the first random variable

  Args:
    cov_matrix (numpy array of floats): desired covariance matrix

  Returns:
    (numpy array of floats) : samples from the bivariate Gaussian, with each
                              column corresponding to a different random
                              variable
  """

  mean = np.array([0, 0])
  X = np.random.multivariate_normal(mean, cov_matrix, size=1000)
  indices_for_sorting = np.argsort(X[:, 0])
  X = X[indices_for_sorting, :]

  return X

def plot_data(X):
  """
  Plots bivariate data. Includes a plot of each random variable, and a scatter
  plot of their joint activity. The title indicates the sample correlation
  calculated from the data.

  Args:
    X (numpy array of floats) :   Data matrix each column corresponds to a
                                  different random variable

  Returns:
    Nothing.
  """

  fig = plt.figure(figsize=[8, 4])
  gs = fig.add_gridspec(2, 2)
  ax1 = fig.add_subplot(gs[0, 0])
  ax1.plot(X[:, 0], color='k')
  plt.ylabel('Neuron 1')
  plt.title('Sample var 1: {:.1f}'.format(np.var(X[:, 0])))
  ax1.set_xticklabels([])
  ax2 = fig.add_subplot(gs[1, 0])
  ax2.plot(X[:, 1], color='k')
  plt.xlabel('Sample Number')
  plt.ylabel('Neuron 2')
  plt.title('Sample var 2: {:.1f}'.format(np.var(X[:, 1])))
  ax3 = fig.add_subplot(gs[:, 1])
  ax3.plot(X[:, 0], X[:, 1], '.', markerfacecolor=[.5, .5, .5],
           markeredgewidth=0)
  ax3.axis('equal')
  plt.xlabel('Neuron 1 activity')
  plt.ylabel('Neuron 2 activity')
  plt.title('Sample corr: {:.1f}'.format(np.corrcoef(X[:, 0], X[:, 1])[0, 1]))
  plt.show()

def plot_basis_vectors(X, W):
  """
  Plots bivariate data as well as new basis vectors.

  Args:
    X (numpy array of floats) :   Data matrix each column corresponds to a
                                  different random variable
    W (numpy array of floats) :   Square matrix representing new orthonormal
                                  basis each column represents a basis vector

  Returns:
    Nothing.
  """

  plt.figure(figsize=[4, 4])
  plt.plot(X[:, 0], X[:, 1], '.', color=[.5, .5, .5], label='Data')
  plt.axis('equal')
  plt.xlabel('Neuron 1 activity')
  plt.ylabel('Neuron 2 activity')
  plt.plot([0, W[0, 0]], [0, W[1, 0]], color='r', linewidth=3,
           label='Basis vector 1')
  plt.plot([0, W[0, 1]], [0, W[1, 1]], color='b', linewidth=3,
           label='Basis vector 2')
  plt.legend()
  plt.show()

def plot_data_new_basis(Y):
  """
  Plots bivariate data after transformation to new bases.
  Similar to plot_data but with colors corresponding to projections onto
  basis 1 (red) and basis 2 (blue). The title indicates the sample correlation
  calculated from the data.

  Note that samples are re-sorted in ascending order for the first
  random variable.

  Args:
    Y (numpy array of floats):   Data matrix in new basis each column
                                 corresponds to a different random variable

  Returns:
    Nothing.
  """
  fig = plt.figure(figsize=[8, 4])
  gs = fig.add_gridspec(2, 2)
  ax1 = fig.add_subplot(gs[0, 0])
  ax1.plot(Y[:, 0], 'r')
  plt.xlabel
  plt.ylabel('Projection \n basis vector 1')
  plt.title('Sample var 1: {:.1f}'.format(np.var(Y[:, 0])))
  ax1.set_xticklabels([])
  ax2 = fig.add_subplot(gs[1, 0])
  ax2.plot(Y[:, 1], 'b')
  plt.xlabel('Sample number')
  plt.ylabel('Projection \n basis vector 2')
  plt.title('Sample var 2: {:.1f}'.format(np.var(Y[:, 1])))
  ax3 = fig.add_subplot(gs[:, 1])
  ax3.plot(Y[:, 0], Y[:, 1], '.', color=[.5, .5, .5])
  ax3.axis('equal')
  plt.xlabel('Projection basis vector 1')
  plt.ylabel('Projection basis vector 2')
  plt.title('Sample corr: {:.1f}'.format(np.corrcoef(Y[:, 0], Y[:, 1])[0, 1]))
  plt.show()
```

---
# Section 1: Generate correlated multivariate data

```python
# @title Video 2: Multivariate data
video = YouTubeVideo(id="jcTq2PgU5Vw", width=854, height=480, fs=1)
print("Video available at https://youtube.com/watch?v=" + video.id)
video
```

    Video available at https://youtube.com/watch?v=jcTq2PgU5Vw

To gain intuition, we will first use a simple model to generate multivariate data. Specifically, we will draw random samples from a *bivariate normal distribution*. This is an extension of the one-dimensional normal distribution to two dimensions, in which each $x_i$ is marginally normal with mean $\mu_i$ and variance $\sigma_i^2$:

\begin{align}
x_i \sim \mathcal{N}(\mu_i,\sigma_i^2).
\end{align}

Additionally, the joint distribution for $x_1$ and $x_2$ has a specified correlation coefficient $\rho$. Recall that the correlation coefficient is a normalized version of the covariance, and ranges between -1 and +1:

\begin{align}
\rho = \frac{\text{cov}(x_1,x_2)}{\sqrt{\sigma_1^2 \sigma_2^2}}.
\end{align}

For simplicity, we will assume that the mean of each variable has already been subtracted, so that $\mu_i=0$. The remaining parameters can be summarized in the covariance matrix, which for two dimensions has the following form:

\begin{equation*}
{\bf \Sigma} =
\begin{pmatrix}
 \text{var}(x_1) & \text{cov}(x_1,x_2) \\
 \text{cov}(x_1,x_2) &\text{var}(x_2)
\end{pmatrix}.
\end{equation*}

In general, $\bf \Sigma$ is a symmetric matrix with the variances $\text{var}(x_i) = \sigma_i^2$ on the diagonal, and the covariances on the off-diagonal. Later, we will see that the covariance matrix plays a key role in PCA.

## Exercise 1: Draw samples from a distribution

We have provided code to draw random samples from a zero-mean bivariate normal distribution. Throughout this tutorial, we'll imagine these samples represent the activity (firing rates) of two recorded neurons on different trials. Fill in the function below to calculate the covariance matrix given the desired variances and correlation coefficient. The covariance can be found by rearranging the equation above:

\begin{align}
\text{cov}(x_1,x_2) = \rho \sqrt{\sigma_1^2 \sigma_2^2}.
\end{align}

Use these functions to generate and plot data while varying the parameters. You should get a feel for how changing the correlation coefficient affects the geometry of the simulated data.

**Steps**
* Fill in the function `calculate_cov_matrix` to calculate the desired covariance.
* Generate and plot the data for $\sigma_1^2 =1$, $\sigma_1^2 =1$, and $\rho = .8$. Try plotting the data for different values of the correlation coefficent: $\rho = -1, -.5, 0, .5, 1$.

```python
help(plot_data)
help(get_data)
```

    Help on function plot_data in module __main__:

    plot_data(X)
        Plots bivariate data. Includes a plot of each random variable, and a scatter
        plot of their joint activity. The title indicates the sample correlation
        calculated from the data.

        Args:
          X (numpy array of floats) :   Data matrix each column corresponds to a
                                        different random variable

        Returns:
          Nothing.

    Help on function get_data in module __main__:

    get_data(cov_matrix)
        Returns a matrix of 1000 samples from a bivariate, zero-mean Gaussian.

        Note that samples are sorted in ascending order for the first random variable

        Args:
          cov_matrix (numpy array of floats): desired covariance matrix

        Returns:
          (numpy array of floats) : samples from the bivariate Gaussian, with each
                                    column corresponding to a different random
                                    variable

```python
def calculate_cov_matrix(var_1, var_2, corr_coef):
  """
  Calculates the covariance matrix based on the variances and correlation
  coefficient.

  Args:
    var_1 (scalar)          : variance of the first random variable
    var_2 (scalar)          : variance of the second random variable
    corr_coef (scalar)      : correlation coefficient

  Returns:
    (numpy array of floats) : covariance matrix
  """

  #################################################
  ## TODO for students: calculate the covariance matrix
  # Fill out function and remove
#   raise NotImplementedError("Student excercise: calculate the covariance matrix!")
  #################################################

  # Calculate the covariance from the variances and correlation
  cov = corr_coef*np.sqrt(var_2*var_1)

  cov_matrix = np.array([[var_1, cov], [cov, var_2]])

  return cov_matrix

###################################################################
## TO DO for students: generate and plot bivariate Gaussian data with variances of 1
## and a correlation coefficients of: 0.8
## repeat while varying the correlation coefficient from -1 to 1
###################################################################
np.random.seed(2020)  # set random seed
variance_1 = 1
variance_2 = 1
corr_coef = 0.8

# Uncomment to test your code and plot
cov_matrix = calculate_cov_matrix(variance_1, variance_2, corr_coef)
X = get_data(cov_matrix)
plot_data(X)
```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W1D5_DimensionalityReduction/solutions/W1D5_Tutorial1_Solution_57497711.py)

*Example output:*

---
# Section 2: Define a new orthonormal basis

```python
# @title Video 3: Orthonormal bases
video = YouTubeVideo(id="PC1RZELnrIg", width=854, height=480, fs=1)
print("Video available at https://youtube.com/watch?v=" + video.id)
video
```

    Video available at https://youtube.com/watch?v=PC1RZELnrIg

Next, we will define a new orthonormal basis of vectors ${\bf u} = [u_1,u_2]$ and ${\bf w} = [w_1,w_2]$. As we learned in the video, two vectors are orthonormal if:

1.   They are orthogonal (i.e., their dot product is zero):
\begin{equation}
{\bf u\cdot w} = u_1 w_1 + u_2 w_2 = 0
\end{equation}
2.   They have unit length:
\begin{equation}
||{\bf u} || = ||{\bf w} || = 1
\end{equation}

In two dimensions, it is easy to make an arbitrary orthonormal basis. All we need is a random vector ${\bf u}$, which we have normalized. If we now define the second basis vector to be ${\bf w} = [-u_2,u_1]$, we can check that both conditions are satisfied:
\begin{equation}
{\bf u\cdot w} = - u_1 u_2 + u_2 u_1 = 0
\end{equation}
and
\begin{equation}
{|| {\bf w} ||} = \sqrt{(-u_2)^2 + u_1^2} = \sqrt{u_1^2 + u_2^2} = 1,
\end{equation}
where we used the fact that ${\bf u}$ is normalized. So, with an arbitrary input vector, we can define an orthonormal basis, which we will write in matrix by stacking the basis vectors horizontally:

\begin{equation}
{{\bf W} } =
\begin{pmatrix}
 u_1 & w_1 \\
 u_2 & w_2
\end{pmatrix}.
\end{equation}

## Exercise 2: Find an orthonormal basis

In this exercise you will fill in the function below to define an orthonormal basis, given a single arbitrary 2-dimensional vector as an input.

**Steps**
* Modify the function `define_orthonormal_basis` to first normalize the first basis vector $\bf u$.
* Then complete the function by finding a basis vector $\bf w$ that is orthogonal to $\bf u$.
* Test the function using initial basis vector ${\bf u} = [3,1]$. Plot the resulting basis vectors on top of the data scatter plot using the function `plot_basis_vectors`. (For the data, use  $\sigma_1^2 =1$, $\sigma_2^2 =1$, and $\rho = .8$).

```python
help(plot_basis_vectors)
```

    Help on function plot_basis_vectors in module __main__:

    plot_basis_vectors(X, W)
        Plots bivariate data as well as new basis vectors.

        Args:
          X (numpy array of floats) :   Data matrix each column corresponds to a
                                        different random variable
          W (numpy array of floats) :   Square matrix representing new orthonormal
                                        basis each column represents a basis vector

        Returns:
          Nothing.

```python
def define_orthonormal_basis(u):
  """
  Calculates an orthonormal basis given an arbitrary vector u.

  Args:
    u (numpy array of floats) : arbitrary 2-dimensional vector used for new
                                basis

  Returns:
    (numpy array of floats)   : new orthonormal basis
                                columns correspond to basis vectors
  """

  #################################################
  ## TODO for students: calculate the orthonormal basis
  # Fill out function and remove
#   raise NotImplementedError("Student excercise: implement the orthonormal basis function")
  #################################################

  # normalize vector u
  u = np.array(u)
  u = u/np.sqrt(u.T@u)
  # calculate vector w that is orthogonal to w
  w = [-u[1],u[0]]

  W = np.column_stack([u, w])

  return W

np.random.seed(2020)  # set random seed
variance_1 = 1
variance_2 = 1
corr_coef = 0.8

cov_matrix = calculate_cov_matrix(variance_1, variance_2, corr_coef)
X = get_data(cov_matrix)
u = np.array([3, 1])

# Uncomment and run below to plot the basis vectors
W = define_orthonormal_basis(u)
plot_basis_vectors(X, W)
```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W1D5_DimensionalityReduction/solutions/W1D5_Tutorial1_Solution_7a9640ef.py)

*Example output:*

---
# Section 3: Project data onto new basis

```python
# @title Video 4: Change of basis
video = YouTubeVideo(id="Mj6BRQPKKUc", width=854, height=480, fs=1)
print("Video available at https://youtube.com/watch?v=" + video.id)
video
```

    Video available at https://youtube.com/watch?v=Mj6BRQPKKUc

Finally, we will express our data in the new basis that we have just found. Since $\bf W$ is orthonormal, we can project the data into our new basis using simple matrix multiplication :

\begin{equation}
{\bf Y = X W}.
\end{equation}

We will explore the geometry of the transformed data $\bf Y$ as we vary the choice of basis.

## Exercise 3: Define an orthonormal basis
In this exercise you will fill in the function below to define an orthonormal basis, given a single arbitrary vector as an input.

**Steps**
* Complete the function `change_of_basis` to project the data onto the new basis.
* Plot the projected data using the function `plot_data_new_basis`.
* What happens to the correlation coefficient in the new basis? Does it increase or decrease?
* What happens to variance?

```python
def change_of_basis(X, W):
  """
  Projects data onto new basis W.

  Args:
    X (numpy array of floats) : Data matrix each column corresponding to a
                                different random variable
    W (numpy array of floats) : new orthonormal basis columns correspond to
                                basis vectors

  Returns:
    (numpy array of floats)    : Data matrix expressed in new basis
  """

  #################################################
  ## TODO for students: project the data onto o new basis W
  # Fill out function and remove
#   raise NotImplementedError("Student excercise: implement change of basis")
  #################################################

  # project data onto new basis described by W
  Y = X@W
  return Y

# Unomment below to transform the data by projecting it into the new basis
Y = change_of_basis(X, W)
plot_data_new_basis(Y)
```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W1D5_DimensionalityReduction/solutions/W1D5_Tutorial1_Solution_a1124bbc.py)

*Example output:*

## Interactive Demo: Play with the basis vectors
To see what happens to the correlation as we change the basis vectors, run the cell below. The parameter $\theta$ controls the angle of $\bf u$ in degrees. Use the slider to rotate the basis vectors.

```python
# @title

# @markdown Make sure you execute this cell to enable the widget!

def refresh(theta=0):
  u = [1, np.tan(theta * np.pi / 180)]
  W = define_orthonormal_basis(u)
  Y = change_of_basis(X, W)
  plot_basis_vectors(X, W)
  plot_data_new_basis(Y)

_ = widgets.interact(refresh, theta=(0, 90, 5))
```

    interactive(children=(IntSlider(value=0, description='theta', max=90, step=5), Output()), _dom_classes=('widge…

## Questions

* What happens to the projected data as you rotate the basis?
* How does the correlation coefficient change? How does the variance of the projection onto each basis vector change?
* Are you able to find a basis in which the projected data is **uncorrelated**?

---
# Summary

- In this tutorial, we learned that multivariate data can be visualized as a cloud of points in a high-dimensional vector space. The geometry of this cloud is shaped by the covariance matrix.

- Multivariate data can be represented in a new orthonormal basis using the dot product. These new basis vectors correspond to specific mixtures of the original variables - for example, in neuroscience, they could represent different ratios of activation  across a population of neurons.

- The projected data (after transforming into the new basis) will generally have a different geometry from the original data. In particular, taking basis vectors that are aligned with the spread of cloud of points decorrelates the data.

* These concepts - covariance, projections, and orthonormal bases - are key for understanding PCA, which we be our focus in the next tutorial.