```python
import numpy as np
import pandas as pd
import random as random
from collections import Counter
from math import factorial
import matplotlib.pyplot as plt
%matplotlib inline
```

```python
def coin():
    toss = [] # empty list
    for i in range(11):   # 10 coins are tossed
        s = random.randint(0, 1) # 0 for tail, 1 for head
        toss.append(s)
    h=Counter(toss)
    d=dict(h)
    d1=pd.DataFrame(d.items(), columns=['Head', 'Frequency'])
    d1.sort_values("Head", axis = 0, ascending = True,inplace = True)
    nd=d1.Frequency[1]
    return nd
```

```python
N=100
data=[]
for i in range(N):
    dd=coin()
    data.append(dd)
cc=Counter(data)
D=dict(cc)
DF=pd.DataFrame(D.items(), columns=['Head', 'Frequency'])
DF
```

      Head
      Frequency

      0
      7
      9

      1
      3
      13

      2
      4
      21

      3
      6
      18

      4
      5
      32

      5
      8
      5

      6
      2
      1

      7
      9
      1

```python
DF.sort_values("Head", axis = 0, ascending = True,inplace = True)
DF
```

      Head
      Frequency

      6
      2
      1

      1
      3
      13

      2
      4
      21

      4
      5
      32

      3
      6
      18

      0
      7
      9

      5
      8
      5

      7
      9
      1

```python
plt.bar(DF['Head'],DF['Frequency'],color='cyan',edgecolor='black')
plt.plot(DF['Head'],DF['Frequency'],color='magenta',linewidth=3,)
plt.xlabel('No. of head', fontsize=14)
plt.ylabel('Frequency',fontsize=14)
plt.title('Tossing of coin', fontsize = 18)
plt.show()
```

```python
DF['Cum.Freq.'] = DF['Frequency'].cumsum()
DF
```

      Head
      Frequency
      Cum.Freq.

      6
      2
      1
      1

      1
      3
      13
      14

      2
      4
      21
      35

      4
      5
      32
      67

      3
      6
      18
      85

      0
      7
      9
      94

      5
      8
      5
      99

      7
      9
      1
      100

```python
plt.plot(DF['Head'], DF['Cum.Freq.'], c='blue',label='Less than')
plt.plot(DF['Head'], N-DF['Cum.Freq.'], c='green',label='Greater than')
plt.xlabel('no. of Head', fontsize=14)
plt.ylabel('Cum.freq.',fontsize=14)
plt.legend()
plt.title('Ogive Curve', fontsize = 18 ,color='red')
plt.tight_layout()
plt.show()
```

```python
rdata=sorted(data) #arranging data
mean=np.mean(rdata)
sd=np.sqrt(np.var(rdata))
```

```python
print("Mean : ", mean)
print("standard deviation  : ", sd)
```

    Mean :  5.05
    standard deviation  :  1.4026760139105536

```python
sk= pd.DataFrame(rdata).skew()[0]           # skewness
ku= pd.DataFrame(rdata).kurtosis()[0]      # kurtosis
```

```python
print("skewness : ",sk)
print("kurtosis : ",ku)
```

    skewness :  0.3953910767397892
    kurtosis :  -0.08088534328441499

```python
if(sk>0):
    print('Distribution is positively tail')
elif(sk<0):
    print('Distribution is negatively tail')
elif(sk1==0):
    print('Distribution is symmetric ')
```

    Distribution is positively tail

```python
if(ku>0):
    print('Distribution is Leptokurtic')
elif(ku<0):
    print('Distribution is Platykurtic')
elif(ku==0):
    print('Distribution is Mesokurtic')
```

    Distribution is Platykurtic

```python
Q=np.percentile(rdata, [25, 50, 75])      # percentile
print('first quartiles is ',Q[0])
print('Median is ',Q[1])
print('third quartiles is ',Q[2])
```

    first quartiles is  4.0
    Median is  5.0
    third quartiles is  6.0

Binomial distribution
\begin{equation}
prob(x)=\frac{n!}{x!(n-x)!}p^xq^{n-x}
\end{equation}

```python
def binomial(x, n, p):

    return float(factorial(n)/(factorial(x) * factorial(n - x))) * p ** x * (1 - p)**(n - x)
```

```python
n=max(DF['Head'])+1
ef=[] #expected frequency
for x in range(min(DF['Head']), n):
    y =binomial(x, n, 0.5)
    ef.append(round(N*y))
```

```python
DF['Exp. Freq.'] = ef
DF
```

      Head
      Frequency
      Cum.Freq.
      Exp. Freq.

      6
      2
      1
      1
      4

      1
      3
      13
      14
      12

      2
      4
      21
      35
      21

      4
      5
      32
      67
      25

      3
      6
      18
      85
      21

      0
      7
      9
      94
      12

      5
      8
      5
      99
      4

      7
      9
      1
      100
      1

```python
plt.plot(DF['Head'], ef,'b-',label='Exp.freq.')
plt.plot(DF['Head'], DF['Frequency'],'g-.', label='Obs.freq.')
plt.xlabel("No. of head", fontsize=14)
plt.ylabel("Freqyency", fontsize=14)
plt.legend()
plt.title('Binomial Distribution', fontsize = 18 ,color='red')
plt.tight_layout()
plt.show()
```

```python
chi2=sum((ef-DF['Frequency'])**2/ef)
print('chi square is ',chi2)
```

    chi square is  5.721904761904762

```python

```


# **Introduction to PyTorch** - Logistic Regression, MLP, and MNIST

**Duke Machine Learning Winter School: Computer Vision 2022**

**Code by Billy Carson**

In this notebook, we're going to gain familiarity with the PyTorch deep learning framework by building PyTorch models to perform handwritten digit classification.

## **1. Setting Up the Enivronment**

### 1.1. Import libraries, modules, and functions

First, we need to import necessary libraries, modules, and functions we'll use throughout this tutorial. To give ourselves access to the offerings of the PyTorch library, we'll import the [torch](https://pytorch.org/docs/stable/torch.html) package and the [torch.nn](https://pytorch.org/docs/stable/nn.html) module. Later, you can cut down on import bloat by only importing specific modules, functions, and objects that you'll need for your project. For now, we'll keep it simple by just using these two import statements.

```python
# Define random state
random_state = 0

# Import libraries, modules, and functions
import time
from idx2numpy import convert_from_file
import numpy as np
import matplotlib.pyplot as plt

# PyTorch imports
import torch
import torch.nn as nn

```

### 1.2. GPU/CUDA availability

When setting up your coding environment for a project or experiment, it's important that we are aware of the available computing resources. The resources we have access to will determine what models we can or can't use and the amount of data we can load into memory, among other considerations. For this tutorial, since we are dealing with simpler machine learning models, we won't need access to a [Graphics Processing Unit](https://en.wikipedia.org/wiki/Graphics_processing_unit) (GPU). GPUs help speed up the training of deep learning models by parallelizing many computations intead of computing sequentially. This parallelization is critical when training more complex models like [Convolutional Neural Networks](https://en.wikipedia.org/wiki/Convolutional_neural_network) (CNNs), and without it these models would take prohibitively long to train.

We can check whether or not our environment has access to a GPU by calling the [`torch.cuda.is_available()`](https://pytorch.org/docs/stable/cuda.html) function. [CUDA](https://developer.nvidia.com/about-cuda) is a parallel computing platform and programming model developed by Nvidia for general computing on its own GPUs. This function will tell us if CUDA is available in this computing environment. Since we don't need/have access to a GPU, `torch.cuda.is_available()` should return `False`.

```python
# Display GPU/CUDA availability
print('\nCUDA availability:  %s\n' % (torch.cuda.is_available()))

```

    CUDA availability:  False

## **2. Loading and Preparing the Data**

One of the first (and often most important) steps in a machine learning pipeline is preparing the data. Here, we will load the MNIST dataset of images of handwritten digits and format this data so it can be used with PyTorch models.

### 2.1. Load MNIST dataset

The [MNIST dataset](http://yann.lecun.com/exdb/mnist/) is very popular machine learning dataset, consisting of 70,000 grayscale, 28x28 images of handwritten digits.  We'll be using it as our example dataset for this section of the tutorial, with the goal being to predict which digit is depicted in each image.

First, let's load the training and test sets of MNIST. The MNIST dataset needs to be converted to a type that can be operated on by PyTorch models. MNIST data will first be loaded from IDX file format and converted into NumPy arrays using the [idx2numpy](https://pypi.org/project/idx2numpy/) package. [NumPy](https://numpy.org) is a commonly-used Python library which provides support for multi-dimensional arrays and matrices, along with a variety of mathematical functions to operate on these arrays. In Python data science and machine learning pipelines, [NumPy arrays](https://numpy.org/doc/stable/reference/generated/numpy.array.html) are commonly used to store data and facilitate quick and efficient operations on array-like data.

```python
# Load NumPy arrays of MNIST dataset
mnist_data_dir = '/datasets/mnist/'
X_raw_train_npy = convert_from_file(mnist_data_dir + 'train-images-idx3-ubyte')
y_train_npy = convert_from_file(mnist_data_dir + 'train-labels-idx1-ubyte')
X_raw_test_npy = convert_from_file(mnist_data_dir + 't10k-images-idx3-ubyte')
y_test_npy = convert_from_file(mnist_data_dir + 't10k-labels-idx1-ubyte')

# Display class type
print('\nClass types:')
print('  Images:  %s' % (type(X_raw_train_npy)))
print('  Labels:  %s' % (type(y_train_npy)))

```

    Class types:
      Images:  <class 'numpy.ndarray'>
      Labels:  <class 'numpy.ndarray'>

### 2.2 Convert data from NumPy arrays to PyTorch tensors

Now, MNIST data stored as NumPy arrays must be converted to [PyTorch tensors](https://pytorch.org/docs/stable/tensors.html) so that this data can be operated on by PyTorch models.

```python
# Convert NumPy arrays to PyTorch tensors
X_raw_train = torch.from_numpy(X_raw_train_npy)
y_train = torch.from_numpy(y_train_npy)
X_raw_test = torch.from_numpy(X_raw_test_npy)
y_test = torch.from_numpy(y_test_npy)

```

### 2.3. Explore data/PyTorch tensors

Tensors are a specialized data structure very similar to arrays and matrices that house data in $N$ dimensions. Tensors can be thought of as generalizations of matrices (which are specifically 2-dimensional tensors) to $N$-dimensional space.

In PyTorch, tensors are used to encode data, model inputs, model outputs, and model parameters. [PyTorch tensors](https://pytorch.org/docs/stable/tensors.html) are objects similar to [NumPy ndarrays](https://numpy.org/doc/stable/reference/generated/numpy.ndarray.html), except that tensors can run on GPUs or other hardware accelerators. Tensors are also optimized for automatic differentiation with [`autograd`](https://pytorch.org/tutorials/beginner/blitz/autograd_tutorial.html), PyTorch’s automatic differentiation engine that powers neural network training. Here, where call the tensor method `shape()` to get an idea of the size and dimensionality of our data.

```python
# Display training data dimensionality
print('\nTraining data dimensionality:')
print('  Images:  (%d, %d, %d)' % (X_raw_train.shape))
print('  Labels:  (%s,)' % (y_train.shape))

# Display test data dimensionality
print('\nTest data dimensionality:')
print('  Images:  (%d, %d, %d)' % (X_raw_test.shape))
print('  Labels:  (%d,)\n' % (y_test.shape))

# Display different labels
print('Training labels:')
print(np.unique(y_train.numpy()))
print('\nTest labels:')
print(np.unique(y_test.numpy()))
print()

```

    Training data dimensionality:
      Images:  (60000, 28, 28)
      Labels:  (60000,)

    Test data dimensionality:
      Images:  (10000, 28, 28)
      Labels:  (10000,)

    Training labels:
    [0 1 2 3 4 5 6 7 8 9]

    Test labels:
    [0 1 2 3 4 5 6 7 8 9]

### 2.4. Format tensors

For our PyTorch models and loss functions to be able to operate on the tensors of MNIST data, they first need to be of a specific [tensor type](https://pytorch.org/docs/stable/tensors.html). Specifically, the image data tensors need to be 32-bit floating point tensors (tensor type `FloatTensor`) and the labels tensors need to be 64-bit signed integers (tensor type `LongTensor`).

We can display the tensor type by calling the tensor method `type()`, which will return the tensor type. As we'll see, image data tensors need to be converted from type `ByteTensor` to type `FloatTensor` and labels from type `ByteTensor` to type `LongTensor`. We can make this conversion by calling the tensor method `float()`.

```python
# Display original tensor type
print('\nTensor type (before conversion):')
print('  Images:  %s' % (X_raw_train.type()))
print('  Labels:  %s\n' % (y_train.type()))

```

    Tensor type (before conversion):
      Images:  torch.ByteTensor
      Labels:  torch.ByteTensor

```python
# Convert image tensor type to float
X_raw_train = X_raw_train.float()
X_raw_test = X_raw_test.float()

# Convert labels tensor type to long
y_train = y_train.long()
y_test = y_test.long()

```

```python
# Display tensor type after conversion
print('\nTensor type (after conversion):')
print('  Images:  %s' % (X_raw_train.type()))
print('  Labels:  %s\n' % (y_train.type()))

```

    Tensor type (after conversion):
      Images:  torch.FloatTensor
      Labels:  torch.LongTensor

### 2.5. Scale data

Often in many machine learning and data science applications, it helps to first scale the data. Scaling our data can ensure that all of our features are on approximately the same scale. Significant differences in the scales across input variables may increase the difficulty of our optimization problem [1].

For this modelling problem, we'll scale our data to be between 0 and 1. Since MNIST data are black and white images, data elements correspond to pixel intensities with values between 0 and 255. Therefore, to scale our data between 0 and 1 we just need to divide our image data tensors by 255.

```python
# Display data max and min
print('\nData range pre-scaling (min, max)):')
print('  Train images:  %.1f, %.1f' % (torch.min(X_raw_train), torch.max(X_raw_train)))
print('  Test images:  %.1f, %.1f\n' % (torch.min(X_raw_test), torch.max(X_raw_test)))

```

    Data range pre-scaling (min, max)):
      Train images:  0.0, 255.0
      Test images:  0.0, 255.0

```python
# Scale image pixel intensities between 0 and 1
X_train = X_raw_train / 255.0
X_test = X_raw_test / 255.0

```

```python
# Display data max and min
print('\nData range post-scaling (min, max)):')
print('  Train images:  %.1f, %.1f' % (torch.min(X_train), torch.max(X_train)))
print('  Test images:  %.1f, %.1f\n' % (torch.min(X_test), torch.max(X_test)))

```

    Data range post-scaling (min, max)):
      Train images:  0.0, 1.0
      Test images:  0.0, 1.0

### 2.6. Reshape data

Sometimes we need to reshape the data so that our chosen model can operate on it. Computer vision models such as CNNs operate on directly 2-dimensional image data, so if we were using a CNN we would not need to reshape our MNIST data. However, for this tutorial we will be working with logistic regression and MLP models. Both of these models require input data to be a 1-dimensional vector, meanining that we'll have to "flatten" our image data.

When flattening the image data, we want to preserve the first dimension of our tensors, which represents the samples or observations dimension. We can reshape our training and test data tensors using the [`view()`](https://pytorch.org/docs/stable/tensor_view.html) method. Additionally, this method avoids explicit data copy, thus facilitating fast and memory efficient reshaping, slicing, and element-wise operations on tensors [2].

```python
# Reshape and "flatten" image data
# Hint: -1 tells view() to infer the size of a given dimension
X_train = X_train.view(X_train.shape[0], -1)
X_test = X_test.view(X_test.shape[0], -1)

```

```python
# Display training data dimensionality
print('\nTraining data dimensionality:')
print('  Images:  (%d, %d)' % (X_train.shape))

# Display test data dimensionality
print('\nTest data dimensionality:')
print('  Images:  (%d, %d)\n' % (X_test.shape))

```

    Training data dimensionality:
      Images:  (60000, 784)

    Test data dimensionality:
      Images:  (10000, 784)

### 2.7. Define PyTorch dataset object

[`torch.utils.data.Dataset`](https://pytorch.org/docs/stable/data.html#torch.utils.data.Dataset) is an abstract class representing a dataset. It allows us to treat the dataset as an object of a class, rather than a set of data and labels. Creating a custom dataset will allow us to have greater control over how our model loads and interacts with the data, as well as permit built-in batching, shuffling, and loading data in parallel later on [3]. A custom `Dataset` class must implement three functions: ` __init__`, `__len__`, and `__getitem__`:

* The `__init__` function is run once when instantiating the `Dataset` object.
* The `__len__` function returns the number of samples in our dataset.
* The `__getitem__` function loads and returns a sample from the dataset at the given index idx.

```python
# MNIST PyTorch dataset object
class MNISTDataset(torch.utils.data.Dataset):
    """
    PyTorch dataset object for the MNIST dataset.

    Parameters
    ----------
    X : torch.Tensor
        Flattened image data of size (n_samples x 784).
    y : torch.Tensor
        Image labels, integer values 0, 1, 2, ..., 9.

    Attributes
    ----------
    X : torch.Tensor
        Flattened image data of size (n_samples x 784).
    y : torch.Tensor
        Image labels, integer values 0, 1, 2, ..., 9.
    len : int
        Length of dataset/number of observations in dataset.
    """

    # MNIST PyTorch dataset instantiation method
    def __init__(self, X, y):
        """
        MNIST dataset instantiation method.

        Parameters
        ----------
        X : torch.Tensor
            Flattened image data of size (n_samples x 784).
        y : torch.Tensor
            Image labels, integer values 0, 1, 2, ..., 9.

        Attributes
        ----------
        X : torch.Tensor
            Flattened image data of size (n_samples x 784).
        y : torch.Tensor
            Image labels, integer values 0, 1, 2, ..., 9.
        len : int
            Length of dataset/number of observations in dataset.
        """

        # Assign attributes
        self.X = X         # flattened image data
        self.y = y         # image labels
        self.len = len(X)  # length of dataset

    # Dataset length method
    def __len__(self):
        """
        Method that returns length of dataset/number of observations in dataset.

        Returns
        -------
        self.len : int
            Length of dataset/number of observations in dataset.
        """

        return self.len

    # Dataset indexing method
    def __getitem__(self, idx):
        """
        Method that retrieves samples and corresponding labels from the dataset.

        Parameters
        ----------
        idx : int
            Integer index of dataset sample to be retrieved.

        Returns
        -------
        X_samp : torch.Tensor
            Image sample.
        y_samp : torch.Tensor
            Label corresponding to image sample.
        """

        # Retrieve sample of data indexed by idx
        X_samp = self.X[idx, :]
        y_samp = self.y[idx]

        # Return data sample
        return X_samp, y_samp

```

### 2.8. Instantiate PyTorch datasets

Here, we instantiate a training and test version of the `MNISTDataset` class we defined above. Instantiating a class creates a copy or instance of the class which inherits all class attributes and methods. To instantiate an object in Python, we simply call the class as if it were a function, passing the requisite arguments to the `__init__` method. The newly created object will then be assigned to whatever variable name we choose.

```python
# Instantiate MNIST training and test PyTorch dataset objects
train_dataset = MNISTDataset(X_train, y_train)
test_dataset = MNISTDataset(X_test, y_test)

```

### 2.9. Instantiate PyTorch data loader objects

While we could interface directly with the data as a PyTorch `torch.utils.data.Dataset` object, it's often easier to use a PyTorch [`torch.utils.data.DataLoader`](https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader) object. The PyTorch `DataLoader` is an iterator that automatically takes care of:
* batching the data
* shuffling the data
* loading the data in parallel using multiprocessing workers

Similar to how we created both training and test set instances of PyTorch `Dataset` objects, we will now create instances of training and test set `DataLoader` objects.

```python
# Instantiate MNIST training and test PyTorch data loader objects
batch_size = 100
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

```

## **3. Building a Logistic Regression Model**

### 3.1. Define logistic regression model

For the first part of this tutorial, we'll be building a logistic regression model, which is essentially a fully-connected neural network without any hidden layers. Although basic, logistic regression can perform quite well on many simple classification tasks.

PyTorch has higher level abstractions to help speed up implementation and improve model organization. While there are many ways to organize PyTorch code, one common paradigm is the `torch.nn.Module`. PyTorch uses modules to represent neural networks. Modules are:

* Building blocks of stateful computation. PyTorch provides a robust library of pre-defined modules and makes it simple to define new custom modules, facilitating construction of many different deep neural network architecturess.
* Tightly integrated with PyTorch’s autograd system. Modules make it simple to specify learnable parameters for PyTorch’s Optimizers to update.
* Easy to work with and transform. Modules are straightforward to save, restore, transfer between CPU/GPU/TPU devices, and more [4].

For more on `torch.nn.Module`, see [this](https://pytorch.org/tutorials/beginner/nn_tutorial.html) tutorial. To create our LogisticRegression model that inherits from `torch.nn.Module`, we need to define two methods: `__init__()` and `forward()`. `__init__()` initializes the model parameters and assigns layers to object attributes. The `forward()` method computes the forward pass of the model by passing minibatches of input data through the different layers of the model.

To get an idea of what types of PyTorch layers to include in our logistic regression models, let's take a look at the mathematical formulation of logistic regression. Multiclass logistic regression can be expressed in matrix form as:

\begin{align}
y = x W + b
\end{align}

To take advantage of parallel computation, we commonly process multiple inputs $x$ at once, in a minibatch. We can individual observations or samples of our data $x$ into a batch matrix of $X$. This can be expressed in matrix form as:

\begin{align}
Y = X W + b
\end{align}

In our specific example, the minibatch size $m$ is 100, the dimension of the data is 28 $\times$ 28 = 784, and the number of classes $c$ is 10. Below is an image depicting the transformation of a batch of data $X$ into a set of corresponding predictions $Y$:

PyTorch offers a pre-defined [`torch.nn.Linear`](https://pytorch.org/docs/stable/generated/torch.nn.Linear.html) layer that computes a linear transformation of the input data. We can use this layer to define the parameters of our logistic regression model. Typically, after the `Linear` layer we would implement a softmax activation function that converts our output to a multinomial probability distribution. However, PyTorch has a way of building this calculation into its loss functions. This helps model stability during training. Therefore, we won't add a softmax activation function for our models and instead leave that calculation to PyTorch's loss functions.

```python
# Logistic regression model
class LogisticRegression(nn.Module):
    """
    PyTorch implementation of a logistic regression model.

    Parameters
    ----------
    input_dim : int
        Size/dimensionality of the model input data.
    output_dim : int
        Size/dimensionality of the model output.

    Attributes
    ----------
    input_dim : int
        Size/dimensionality of the model input data.
    output_dim : int
        Size/dimensionality of the model output.
    linear : torch.nn.Linear
        Fully-connected/linear layer.
    """

    # Logistic regression instantiation method
    def __init__(self, input_dim, output_dim):
        """
        Logistic regression model instantiation method.

        Parameters
        ----------
        input_dim : int
            Size/dimensionality of the model input data.
        output_dim : int
            Size/dimensionality of the model output.

        Attributes
        ----------
        input_dim : int
            Size/dimensionality of the model input data.
        output_dim : int
            Size/dimensionality of the model output.
        linear : torch.nn.Linear
            Fully-connected/linear layer.
        """

        # Inherit from torch.nn.Module
        super(LogisticRegression, self).__init__()

        # Assign logistic regression parameters to model attributes
        self.input_dim = input_dim
        self.output_dim = output_dim

        self.linear = nn.Linear(in_features=input_dim,
                                out_features=output_dim,
                                bias=True)

    # Logistic regression forward pass method
    def forward(self, x):
        """
        Logistic regression forward pass method.

        Parameters
        ----------
        x : torch.Tensor()
            Tensor of input data.

        Returns
        -------
        logits : torch.Tensor()
             Raw model predictions (not passed through sigmoid or softmax function).
        """

        # Calculate model predictions/logits
        logits = self.linear(x)

        # Return model output
        return logits

```

### 3.2. Instantiate model

To instantiate our logistic regression model, we need to pass two parameters to the instantiation method `__init__()`: the model `input_dim` and model `output_dim`. `input_dim` will be the dimensionality of our flattened image data, which is 784. `output_dim` corresponds to the number of digits we are trying to classify, which is 10.

```python
# Define model input and output dimensions from data
input_dim = X_train.shape[1]             # second dimension is data dimensionality
output_dim = len(np.unique(y_train.numpy()))  # output dim. is equal to number of unique labels

# Instantiate logistic regression model
lr_model = LogisticRegression(input_dim=input_dim, output_dim=output_dim)

```

## **4. Training the Logistic Regression Model**

### 4.1. Defining a loss function

We need a way to evaluate our model predictions. This evaluation is done via a loss function, which takes the model's predictions and returns a single number summarizing model performance. The magnitude of this loss will inform how much or how little we should update the model parameters. The loss we commonly use in classification is cross-entropy, a concept from information theory. Cross-entropy not only captures whether or not model predictions are correct, it also accounts for how *confident* these predictions are. This encourages the model to produce very high probabilities for correct answers while driving down the probabilities for the wrong answers.

To define our loss function object, we can use PyTorch's implementation of cross entropy [`torch.nn.CrossEntropyLoss`](https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html#torch.nn.CrossEntropyLoss). In addition to computing the cross entropy loss, this loss function takes logits as inputs computes the log-softmax internally, which is why we did not have to add a softmax activation at the end of our logistic regression implementation.

```python
# Instantiate cross-entropy loss function
loss_func = nn.CrossEntropyLoss()

```

### 4.2. Defining optimizer and learning rate

Now that we've defined our loss function as a way of quantifying how good or bad our model predictions are, we can improve our model by changing the parameters in a way that minimizes the loss. For neural networks, this is done via [backpropagation](https://en.wikipedia.org/wiki/Backpropagation): we take the gradient of the loss with respect to model parameters and take a "step" in the direction that reduces our loss. If we were not using a deep learning framework like PyTorch, we would have to go through and derive all the gradients ourselves by hand, then code them into our training pipeline. Luckily, Pytorch has automated gradient calculation with [`torch.autograd`](https://pytorch.org/docs/stable/autograd.html): an "automatic differentiation engine that powers neural network training" [5].

Next, we define our optimizer and the learning rate we'll use for training the model. [`torch.optim`](https://pytorch.org/docs/stable/optim.html) is a package of various, commonly-used optimization algorithms implementations. A `torch.optim` optimizer object will keep track of the current state of model parameters and update the parameters based on the gradients computed by `torch.autograd`. To construct an optimizer you must provide an iterable containing the model parameters to optimize. This can be done by calling the `parameters()` method of our model (a method that is inherited from `torch.nn.Module`). Then, you can specify optimizer-specific options such as the learning rate and weight decay. For many deep learning applications, standard [stochastic gradient descent](https://scikit-learn.org/stable/modules/sgd.html) (SGD) optimizes models during training quite well. We can use PyTorch's `torch.optim` implementation of SGD as our optimizer.

```python
# Instantiate stochastic gradient descent (SGD) optimizer
lr = 0.1
optimizer = torch.optim.SGD(lr_model.parameters(), lr=lr)

```

### 4.3. Define number of training epochs

```python
# Number of training epochs (epochs = # of passes through data)
n_epochs = 10

```

### Training logistic regression on MNIST

To train our model, we need to repeat the following steps:

1. Draw a minibatch.
2. Reset all gradients to zero in the buffers.
3. Perform the forward pass (compute prediction, calculate loss).
4. Perform the backward pass (compute gradients with respects to loss, optimize parameters with SGD).

Performing these steps for the entirety of a dataset once is referred to as an epoch. We'll accomplish the above steps using the following PyTorch functionalities:

1. Draw a minibatches by iterating over our training `DataLoader` instance. Remember, our custom PyTorch data loader takes care of batching the data, shuffling the data, and loading the data in parallel.
2. Reset all gradients to zero in the buffers using the `zero_grad()` method of our optimizer. We need to set the gradients to zero before performing backpropragation because PyTorch accumulates the gradients on subsequent backward passes if `zero_grad()` is not called. Accumulation of gradients is convenient for models such as recurrent neural networks (RNNs); however, for this tutorial we'll zero out gradients before processing each minibatch of data.
3. Perform the forward pass by providing the minibatch of data to our model's `forward()` method.
4. Perform the backward pass by calling the `backward()` method of our model's computed loss and the `step()` method of our optimizer. The `backward()` method of our model computes the parameter gradients with respects to the calculated loss, and the `step()` method of our optimizer updates the parameters according to the calculated gradients.

```python
# Get weights
weights_before = lr_model.linear.weight.detach().numpy().copy()

# Time model training
print('\nTraining logistic regression model...\n')
time_start = time.time()

# Place model in training mode
# .train() method affects operations such as dropout and batch normalization
lr_model.train()

# Train model/iterate over epochs
for epoch in range(n_epochs):
    # Initialize epoch metrics variables
    n_obs = 0
    loss_sum = 0
    n_correct = 0

    # Iterate through training data mini-batches
    for images_batch, labels_batch in train_loader:
        # Zero out the gradients
        optimizer.zero_grad()

        # Forward pass
        y_pred = lr_model(images_batch)         # model predictions
        loss = loss_func(y_pred, labels_batch)  # loss function evaluation

        # Backward pass
        loss.backward()    # backpropagation
        optimizer.step()   # update parameters according to learning rate, gradients

        # Update epoch metrics variables
        n_batch = len(labels_batch)
        n_obs += n_batch
        loss_sum += n_batch * loss.item()
        n_correct += (y_pred.argmax(1)==labels_batch.long()).float().sum().item()

    # Calculate epoch training loss and training accuracy
    loss_train = loss_sum / n_obs
    acc_train = n_correct / n_obs

    # Display training progress
    prog_disp_freq = 1   # frequency of training progress display
    if (((epoch + 1) % prog_disp_freq) == 0) | (epoch == 0) | ((epoch + 1) == n_epochs):
        print('  E%02d | train loss: %s | train acc.: %s' %
              (epoch + 1, '{:.4f}'.format(loss_train), '{:.4f}'.format(acc_train)))

# Time model training
time_end = time.time()
print('\nLogistic regression model training complete.\n')
print('Time to train logistic regression model:  %.1f s\n' % (time_end - time_start))

# Get weights after training
weights_after = lr_model.linear.weight.detach().numpy().copy()

```

    Training logistic regression model...

      E01 | train loss: 0.5373 | train acc.: 0.8662
      E02 | train loss: 0.3596 | train acc.: 0.9005
      E03 | train loss: 0.3314 | train acc.: 0.9079
      E04 | train loss: 0.3167 | train acc.: 0.9115
      E05 | train loss: 0.3070 | train acc.: 0.9141
      E06 | train loss: 0.3003 | train acc.: 0.9163
      E07 | train loss: 0.2952 | train acc.: 0.9171
      E08 | train loss: 0.2906 | train acc.: 0.9186
      E09 | train loss: 0.2871 | train acc.: 0.9198
      E10 | train loss: 0.2842 | train acc.: 0.9215

    Logistic regression model training complete.

    Time to train logistic regression model:  214.6 s

### 4.4. Visualize logistic regression weights

One of the nice things about logistic regression is that since it is a linear model we can interpret what the model has learned by plotting the weights before and after training.

Before training, the model weights for each digit/output look like a static TV screen. Since the weights were randomly initialized, it is expected that the model parameters will appear random before training on any data. After training, we can see the model learned a general template for each digit. Remember that logistic regression takes a dot product between the weights of each digit and input. Therefore, the more the input matches the template or weights for a given digit, the higher the value of the dot product for that digit will be which makes the model more likely to predict that digit.

```python
# Plot logistic regression weights before and after training
fig1, ax1 = plt.subplots(2, 10, figsize=(14, 2.6))
for digit in range(10):
    ax1[0, digit].imshow(np.reshape(weights_before[digit, :], newshape=(28, 28)), cmap='gray')
    ax1[1, digit].imshow(np.reshape(weights_after[digit, :], newshape=(28, 28)), cmap='gray')
    ax1[0, digit].axis('off')
    ax1[1, digit].axis('off')
plt.show()

```

## **5. Evaluating Logistic Regression Model on Test Set**

### 5.1. Evaluate model on test set

Once our model has been trained, we now need to evaluate it on a set of hold-out data, or a test set. Evaluating trained models on a test set is a critical step of machine learning pipelines, as this gives us an idea of how well our model will generalize in the real world to new, never-seen-before data.

First, we need to put our model in evaluation mode using the `eval()` method. This method affects operations with different training and test-time behavior such as dropout and batch normalization. Although we don't use any layers in this tutorial that would be affected by this method, it's a good thing to get into the habit of doing.

Additionally, since we aren't training our model any more via backpropagation, we don't need to compute a loss or calculate gradients. We can use the [`torch.no_grad()`](https://pytorch.org/docs/stable/generated/torch.no_grad.html) context-manager that disables gradient calculation. Disabling gradient calculation is useful for model inference, when you are sure that you will not call Tensor.backward(). Using `torch.no_grad()` will help reduce memory consumption for computations that would otherwise require gradients during training [6].

```python
# Place model in evaluation mode
# .eval() method affects operations such as dropout and batch normalization
lr_model.eval()

# Initialize test set metrics variables
n_obs = 0
n_correct = 0

# Disable gradient calculation
with torch.no_grad():
    # Iterate through test data mini-batches
    for images_batch, labels_batch in test_loader:
        # Forward pass
        y_pred = lr_model(images_batch)  # model predictions

        # Update test set metrics variables
        n_batch = len(labels_batch)
        n_obs += n_batch
        n_correct += (y_pred.argmax(1)==labels_batch.long()).float().sum().item()

# Calculate test accuracy
acc_test = n_correct / n_obs

# Display test accuracy
print('\nLogistic regression model MNIST test acc.:  %.4f\n' % (acc_test))

```

    Logistic regression model MNIST test acc.:  0.9220

## **6. Training and Evaluating Multi-Layer Perceptron on MNIST Dataset**

### 6.1. Defining MLP model

Similarly to how we previously defined our logistic regression model by using PyTorch's object-oriented style, we can implement a neural network/[multi-layer perceptron](https://machinelearningmastery.com/neural-networks-crash-course/) (MLP) with just a few tweaks. Instead of just one `Linear` layer, this time we'll use two. Additionally, after the first `Linear` layer we'll apply a non-linear [activation function](https://machinelearningmastery.com/choose-an-activation-function-for-deep-learning/). Activation functions are a critical part of neural network design, and help give our neural network model added expressivity for modelling complex, non-linear relationships. The de facto nonlinearity used in neural networks is the [rectified linear unit](https://en.wikipedia.org/wiki/Rectifier_(neural_networks)) (ReLU). We can implement this ReLU activation function as a layer in our MLP by using PyTorch's [`torch.nn.ReLU`](https://pytorch.org/docs/stable/generated/torch.nn.ReLU.html) layer.

*Note: multi-layer perceptron and artifical neural network (ANN) are often used interchangeably to describe a feed-forward, fully-connected neural network with many layers.*

```python
# Multi-layer perceptron model
class MLP(nn.Module):
    """
    PyTorch implementation of a multi-layer perceptron model.

    Parameters
    ----------
    input_dim : int
        Size/dimensionality of the model input data.
    output_dim : int
        Size/dimensionality of the model output.
    hidden_dim : int
        Size/dimensionality of the hidden layer.

    Attributes
    ----------
    input_dim : int
        Size/dimensionality of the model input data.
    output_dim : int
        Size/dimensionality of the model output.
    hidden_dim : int
        Size/dimensionality of the hidden layer.
    linear_1 : torch.nn.Linear
        Fully-connected/linear neural network layer.
    linear_2 : torch.nn.Linear
        Fully-connected/linear neural network layer.
    relu : torch.nn.ReLU
        Rectified linear unit activation function.
    """

    # MLP instantiation method
    def __init__(self, input_dim, output_dim, hidden_dim=100):
        """
        Model instantiation method.

        Parameters
        ----------
        input_dim : int
            Size/dimensionality of the model input data.
        output_dim : int
            Size/dimensionality of the model output.
        hidden_dim : int
            Size/dimensionality of the hidden layer.

        Attributes
        ----------
        input_dim : int
            Size/dimensionality of the model input data.
        output_dim : int
            Size/dimensionality of the model output.
        hidden_dim : int
            Size/dimensionality of the hidden layer.
        linear_1 : torch.nn.Linear
            Fully-connected/linear neural network layer.
        linear_2 : torch.nn.Linear
            Fully-connected/linear neural network layer.
        relu : torch.nn.ReLU
            Rectified linear unit activation function.
        """

        # Inherit from torch.nn.Module
        super(MLP, self).__init__()

        # Assign MLP parameters to model attributes
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.output_dim = output_dim

        # Define MLP layers
        self.linear_1 = nn.Linear(in_features=input_dim,
                                  out_features=hidden_dim,
                                  bias=True)
        self.linear_2 = nn.Linear(in_features=hidden_dim,
                                  out_features=output_dim,
                                  bias=True)

        # Define MLP activation function
        self.relu = nn.ReLU()

    # MLP forward pass method
    def forward(self, x):
        """
        MLP forward pass method.

        Parameters
        ----------
        x : torch.Tensor()
            Tensor of input data.

        Returns
        -------
        logits : torch.Tensor()
             Raw model predictions (not passed through sigmoid or softmax function).
        """

        # Forward pass through model
        x = self.linear_1(x)       # first fully-connected layer
        x = self.relu(x)           # non-linear activation transformation
        logits = self.linear_2(x)  # second fully-connected layer

        # Return model output
        return logits

```

### 6.2. Instantiate MLP model

```python
# Define model input and output dimensions from data
input_dim = X_train.shape[1]                  # second dimension is data dimensionality
output_dim = len(np.unique(y_train.numpy()))  # output dim. is equal to number of unique labels

# Instantiate multi-layer perceptron model
mlp_model = MLP(input_dim=input_dim, output_dim=output_dim)

```

### 6.3. Re-define optimizer with MLP model parameters

Since a PyTorch optimizer is an object that keeps track of and updates model parameters, we need to re-define our optimizer with the parameters of our newly defined MLP model. We can keep the same learning rate of 0.1 as last time.

```python
# Instantiate stochastic gradient descent (SGD) optimizer
lr = 0.1
optimizer = torch.optim.SGD(mlp_model.parameters(), lr=lr)

```

### 6.4. Train MLP model on training data

Now we train our MLP model on MNIST. Ideally, since we have now defined a non-linear model with added expressivity by including a hidden layer and a non-linear activation function following this hidden layer, we would like to see better classification performance than that achieved by logistic regression. As we will see, these modifications do indeed result in better hand-written digit classification performance.

Even though we changed the type of model, we can keep the same exact training loop we developed for training logistic regression on MNIST.

```python
# Number of training epochs (epochs = # of passes through data)
n_epochs = 10

# Time model training
print('\nTraining MLP model...\n')
time_start = time.time()

# Place model in training mode
# .train() method affects operations such as dropout and batch normalization
mlp_model.train()

# Train model/iterate over epochs
for epoch in range(n_epochs):
    # Initialize epoch metrics variables
    n_obs = 0
    loss_sum = 0
    n_correct = 0

    # Iterate through training data mini-batches
    for images_batch, labels_batch in train_loader:
        # Zero out the gradients
        optimizer.zero_grad()

        # Forward pass
        y_pred = mlp_model(images_batch)         # model predictions
        loss = loss_func(y_pred, labels_batch)  # loss function evaluation

        # Backward pass
        loss.backward()    # backpropagation
        optimizer.step()   # update parameters according to learning rate, gradients

        # Update epoch metrics variables
        n_batch = len(labels_batch)
        n_obs += n_batch
        loss_sum += n_batch * loss.item()
        n_correct += (y_pred.argmax(1)==labels_batch.long()).float().sum().item()

    # Calculate epoch training loss and training accuracy
    loss_train = loss_sum / n_obs
    acc_train = n_correct / n_obs

    # Display training progress
    prog_disp_freq = 1   # frequency of training progress display
    if (((epoch + 1) % prog_disp_freq) == 0) | (epoch == 0) | ((epoch + 1) == n_epochs):
        print('  E%02d | train loss: %s | train acc.: %s' %
              (epoch + 1, '{:.4f}'.format(loss_train), '{:.4f}'.format(acc_train)))

# Time model training
time_end = time.time()
print('\nMLP model training complete.\n')
print('Time to train MLP model:  %.1f s\n' % (time_end - time_start))

```

    Training MLP model...

      E01 | train loss: 0.5237 | train acc.: 0.8653
      E02 | train loss: 0.2728 | train acc.: 0.9215
      E03 | train loss: 0.2193 | train acc.: 0.9373
      E04 | train loss: 0.1843 | train acc.: 0.9469
      E05 | train loss: 0.1588 | train acc.: 0.9539
      E06 | train loss: 0.1392 | train acc.: 0.9604
      E07 | train loss: 0.1240 | train acc.: 0.9648
      E08 | train loss: 0.1109 | train acc.: 0.9686
      E09 | train loss: 0.1003 | train acc.: 0.9717
      E10 | train loss: 0.0918 | train acc.: 0.9742

    MLP model training complete.

    Time to train MLP model:  272.2 s

### 6.5. Evaluate MLP model on test set

We'll now evaluate our MLP model on the test set. The added expressivity and complexity of our model can be a double-edged sword. While this added complexity gives our model the capacity to learn more complex, non-linear relationships, we must now also be weary of [overfitting](https://www.ibm.com/cloud/learn/overfitting). Overfitting occurs when a model learns the information represented in training data to the point where it's performance on hold-out data is negatively influenced. Deep learning models can have millions (sometimes [billions](https://developer.nvidia.com/blog/openai-presents-gpt-3-a-175-billion-parameters-language-model/)!) of trainable parameters, which means that the number of model parameters can be greater than the number of data points on which the model is trained. These "over-parametrized" models can be prone to overfitting, especially in cases of small data. Here, we're hoping that our test accuracy is similar to our training accuracy of around 96%-97%. As we'll see, our model did not overfit the training set and achieves a comparable test set accuracy.

```python
# Place model in evaluation mode
# .eval() method affects operations such as dropout and batch normalization
mlp_model.eval()

# Initialize test set metrics variables
n_obs = 0
n_correct = 0

# Disable gradient calculation
with torch.no_grad():
    # Iterate through test data mini-batches
    for images_batch, labels_batch in test_loader:
        # Forward pass
        y_pred = mlp_model(images_batch)  # model predictions

        # Update test set metrics variables
        n_batch = len(labels_batch)
        n_obs += n_batch
        n_correct += (y_pred.argmax(1)==labels_batch.long()).float().sum().item()

# Calculate test accuracy
acc_test = n_correct / n_obs

# Display test accuracy
print('\nMLP model MNIST test acc.:  %.4f\n' % (acc_test))

```

    MLP model MNIST test acc.:  0.9704

## **7. Full Pipeline** - Training PyTorch Model on MNIST Dataset

Finally, we can combine our entire deep learning pipeline into a single cell.

In practice, when using Jupyter notebooks you shouldn't combine all of your code into a single cell, as this defeats the purpose of the modular nature of Jupyter notebooks. However, it can be helpful to see the entirety of a basic deep learning pipeline from beginning to end in code (without the extra bloat of my writing).

```python
# -----------------
# IMPORT STATEMENTS
# -----------------

# Define random state
random_state = 0

# Import libraries, modules, and functions
import time
import numpy as np
import matplotlib.pyplot as plt

# PyTorch imports
import torch
import torch.nn as nn

# -----------------------------------------------
# CLASS DEFINITIONS - MNIST DATASET AND MLP MODEL
# -----------------------------------------------

# MNIST PyTorch dataset object
class MNISTDataset(torch.utils.data.Dataset):
    # MNIST PyTorch dataset instantiation method
    def __init__(self, X, y):
        # Assign attributes
        self.X = X         # flattened image data
        self.y = y         # image labels
        self.len = len(X)  # length of dataset

    # Dataset length method
    def __len__(self):
        # Return length of dataset
        return self.len

    # Dataset indexing method
    def __getitem__(self, idx):
        # Retrieve sample of data indexed by idx
        X_samp = self.X[idx, :]
        y_samp = self.y[idx]

        # Return data sample
        return X_samp, y_samp

# Multi-layer perceptron model
class MLP(nn.Module):
    # MLP instantiation method
    def __init__(self, input_dim, output_dim, hidden_dim=100):
        # Inherit from torch.nn.Module
        super(MLP, self).__init__()

        # Assign MLP parameters to model attributes
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.output_dim = output_dim

        # Define MLP layers
        self.linear_1 = nn.Linear(in_features=input_dim,
                                  out_features=hidden_dim,
                                  bias=True)
        self.linear_2 = nn.Linear(in_features=hidden_dim,
                                  out_features=output_dim,
                                  bias=True)

        # Define MLP activation function
        self.relu = nn.ReLU()

    # MLP forward pass method
    def forward(self, x):
        # Forward pass through model
        x = self.linear_1(x)       # first fully-connected layer
        x = self.relu(x)           # non-linear activation transformation
        logits = self.linear_2(x)  # second fully-connected layer

        # Return model output
        return logits

# ---------------------------
# LOADING AND FORMATTING DATA
# ---------------------------

# Load NumPy arrays of MNIST dataset
mnist_data_dir = '/datasets/mnist/'
X_raw_train_npy = convert_from_file(mnist_data_dir + 'train-images-idx3-ubyte')
y_train_npy = convert_from_file(mnist_data_dir + 'train-labels-idx1-ubyte')
X_raw_test_npy = convert_from_file(mnist_data_dir + 't10k-images-idx3-ubyte')
y_test_npy = convert_from_file(mnist_data_dir + 't10k-labels-idx1-ubyte')

# Convert NumPy arrays to PyTorch tensors
X_raw_train = torch.from_numpy(X_raw_train_npy)
y_train = torch.from_numpy(y_train_npy)
X_raw_test = torch.from_numpy(X_raw_test_npy)
y_test = torch.from_numpy(y_test_npy)

# Convert image tensor type to float and labels tensor type to long
X_raw_train = X_raw_train.float()
X_raw_test = X_raw_test.float()
y_train = y_train.long()
y_test = y_test.long()

# Scale image pixel intensities between 0 and 1
X_train = X_raw_train / 255.0
X_test = X_raw_test / 255.0

# Reshape and "flatten" image data
X_train = X_train.view(X_train.shape[0], -1)
X_test = X_test.view(X_test.shape[0], -1)

# -------------------------------------------
# INSTANTIATE DATASET AND DATA LOADER OBJECTS
# -------------------------------------------

# Instantiate MNIST training and test PyTorch dataset objects
train_dataset = MNISTDataset(X_train, y_train)
test_dataset = MNISTDataset(X_test, y_test)

# Load the data
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=100, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=100, shuffle=False)

# -----------------------------------------------
# INSTANTIATE MODEL, LOSS FUNCTION, AND OPTIMIZER
# -----------------------------------------------

# Define model input and output dimensions from data
input_dim = X_train.shape[1]                  # second dimension is data dimensionality
output_dim = len(np.unique(y_train.numpy()))  # output dim. is equal to number of unique labels

# Instantiate model
# model = LogisticRegression(input_dim=input_dim, output_dim=output_dim)
model = MLP(input_dim=input_dim, output_dim=output_dim)

# Loss and Optimizer
loss_func = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

# ----------------------------
# TRAIN MODEL ON TRAINING DATA
# ----------------------------

# Define number of training epochs
n_epochs = 10

# Time model training
print('\nTraining model...\n')
time_start = time.time()

# Place model in training mode
# .train() method affects operations such as dropout and batch normalization
model.train()

# Train model/iterate over epochs
for epoch in range(n_epochs):
    # Initialize epoch metrics variables
    n_obs = 0
    loss_sum = 0
    n_correct = 0

    # Iterate through training data mini-batches
    for images_batch, labels_batch in train_loader:
        # Zero out the gradients
        optimizer.zero_grad()

        # Forward pass
        y_pred = model(images_batch)            # model predictions
        loss = loss_func(y_pred, labels_batch)  # loss function evaluation

        # Backward pass
        loss.backward()   # backpropagation
        optimizer.step()  # update parameters according to learning rate, gradients

        # Update epoch metrics variables
        n_batch = len(labels_batch)
        n_obs += n_batch
        loss_sum += n_batch * loss.item()
        n_correct += (y_pred.argmax(1)==labels_batch.long()).float().sum().item()

    # Calculate epoch training loss and training accuracy
    loss_train = loss_sum / n_obs
    acc_train = n_correct / n_obs

    # Display training progress
    prog_disp_freq = 1   # frequency of training progress display
    if (((epoch + 1) % prog_disp_freq) == 0) | (epoch == 0) | ((epoch + 1) == n_epochs):
        print('  E%02d | train loss: %s | train acc.: %s' %
              (epoch + 1, '{:.4f}'.format(loss_train), '{:.4f}'.format(acc_train)))

# Time model training
time_end = time.time()
print('\nModel training complete.\n')
print('Time to train model:  %.1f s\n' % (time_end - time_start))

# ---------------------------
# EVALUATE MODEL ON TEST DATA
# ---------------------------

# Place model in evaluation mode
# .eval() method affects operations such as dropout and batch normalization
model.eval()

# Initialize test set metrics variables
n_obs = 0
n_correct = 0

# Disable gradient calculation
with torch.no_grad():
    # Iterate through test data mini-batches
    for images_batch, labels_batch in test_loader:
        # Forward pass
        y_pred = model(images_batch)  # model predictions

        # Update test set metrics variables
        n_batch = len(labels_batch)
        n_obs += n_batch
        n_correct += (y_pred.argmax(1)==labels_batch.long()).float().sum().item()

# Calculate test accuracy
acc_test = n_correct / n_obs

# Display test accuracy
print('Model MNIST test acc.:  %.4f\n' % (acc_test))

```

    Training model...

      E01 | train loss: 0.5182 | train acc.: 0.8657
      E02 | train loss: 0.2750 | train acc.: 0.9216
      E03 | train loss: 0.2208 | train acc.: 0.9375
      E04 | train loss: 0.1844 | train acc.: 0.9479
      E05 | train loss: 0.1575 | train acc.: 0.9551
      E06 | train loss: 0.1372 | train acc.: 0.9615
      E07 | train loss: 0.1211 | train acc.: 0.9657
      E08 | train loss: 0.1088 | train acc.: 0.9695
      E09 | train loss: 0.0987 | train acc.: 0.9722
      E10 | train loss: 0.0907 | train acc.: 0.9745

    Model training complete.

    Time to train model:  268.6 s

    Model MNIST test acc.:  0.9682

## **References**

[1] https://machinelearningmastery.com/how-to-improve-neural-network-stability-and-modeling-performance-with-data-scaling/

[2] https://pytorch.org/docs/stable/tensor_view.html

[3] https://pytorch.org/tutorials/beginner/data_loading_tutorial.html

[4] https://pytorch.org/docs/stable/notes/modules.html

[5] https://pytorch.org/docs/stable/autograd.html, https://pytorch.org/docs/stable/optim.html

[6] https://pytorch.org/docs/stable/generated/torch.no_grad.html


# Econometrics HW_12

**11510691 程$
% color
% Aquamarine, black, blue, brown, cyan, darkgray, gray, green, lightgray, lime, magenta, olive, orange, pink, purple, red, teal, violet, white, yellow
\DeclareMathOperator*{\argmin}{argmin}
\DeclareMathOperator*{\argmax}{argmax}
\DeclareMathOperator*{\plim}{plim}
\newcommand{\ffrac}{\displaystyle \frac}
\newcommand{\d}[1]{\displaystyle{#1}}
\newcommand{\space}{\text{ }}
\newcommand{\bspace}{\;\;\;\;}
\newcommand{\bbspace}{\;\;\;\;\;\;\;\;}
\newcommand{\QQQ}{\boxed{?\:}}
\newcommand{\void}{\left.\right.}
\newcommand{\Tran}[1]{{#1}^{\mathrm{T}}}
\newcommand{\CB}[1]{\left\{ #1 \right\}}
\newcommand{\SB}[1]{\left[ #1 \right]}
\newcommand{\P}[1]{\left( #1 \right)}
\newcommand{\abs}[1]{\left| #1 \right|}
\newcommand{\norm}[1]{\left\| #1 \right\|}
\newcommand{\given}[1]{\left. #1 \right|}
\newcommand{\using}[1]{\stackrel{\mathrm{#1}}{=}}
\newcommand{\asim}{\overset{\text{a}}{\sim}}
\newcommand{\RR}{\mathbb{R}}
\newcommand{\EE}{\mathbb{E}}
\newcommand{\II}{\mathbb{I}}
\newcommand{\NN}{\mathbb{N}}
\newcommand{\ZZ}{\mathbb{Z}}
\newcommand{\QQ}{\mathbb{Q}}
\newcommand{\PP}{\mathbb{P}}
\newcommand{\AcA}{\mathcal{A}}
\newcommand{\FcF}{\mathcal{F}}
\newcommand{\AsA}{\mathscr{A}}
\newcommand{\FsF}{\mathscr{F}}
\newcommand{\dd}{\mathrm{d}}
\newcommand{\I}[1]{\mathrm{I}\left( #1 \right)}
\newcommand{\N}[1]{\mathcal{N}\left( #1 \right)}
\newcommand{\Exp}[1]{\mathrm{E}\left[ #1 \right]}
\newcommand{\Var}[1]{\mathrm{Var}\left[ #1 \right]}
\newcommand{\Avar}[1]{\mathrm{Avar}\left[ #1 \right]}
\newcommand{\Cov}[1]{\mathrm{Cov}\left( #1 \right)}
\newcommand{\Corr}[1]{\mathrm{Corr}\left( #1 \right)}
\newcommand{\ExpH}{\mathrm{E}}
\newcommand{\VarH}{\mathrm{Var}}
\newcommand{\AVarH}{\mathrm{Avar}}
\newcommand{\CovH}{\mathrm{Cov}}
\newcommand{\CorrH}{\mathrm{Corr}}
\newcommand{\ow}{\text{otherwise}}
\newcommand{\wp}{\text{with probability }}
\newcommand{\FSD}{\text{FSD}}
\newcommand{\SSD}{\text{SSD}}$远星$^\dagger$**

## Question 15.e1

$\bspace$To prove the two methods lead to the same estimation, is equivalent to prove the following three

$$\left\{\begin{align}
\sum_{i=1}^n \hat\beta_1 y_{i2} &= \sum_{i=1}^n \hat\beta_1 \hat y_{i2}\\
\sum_{i=1}^n z_{i1}\hat \beta_1 y_{i2} &= \sum_{i=1}^n z_{i1}\hat\beta_1 \hat y_{i2}\\
\sum_{i=1}^n \hat y_{i2} \hat \beta_1 y_{i2} &= \sum_{i=1}^n \hat y_{i2}\hat\beta_1 \hat y_{i2}\\
\end{align}\right.$$

$\bspace$The first one is true because $\sum_{i=1}^n \P{y_{i2} - \hat y_{i2}} = \sum \hat v_{i2}=0$. The second one is true because $\sum_{i=1}^n z_{i1}\P{y_{i2} - \hat y_{i2}} = \Cov{z_1,\hat v_2} = 0$. The third one is true because $\sum_{i=1}^n \hat y_{i2}\P{y_{i2} - \hat y_{i2}} = \Cov{\hat y_2,\hat v_2} = 0$

***
$\\[20em]$

$\\[10em]$


#  Neuromatch Academy: Week 1, Day 1, Tutorial 3
# Model Types: "Why" models
__Content creators:__ Matt Laporte, Byron Galbraith, Konrad Kording

__Content reviewers:__ Dalin Guo, Aishwarya Balwani, Madineh Sarvestani, Maryam Vaziri-Pashkam, Michael Waskom

We would like to acknowledge [Steinmetz _et al._ (2019)](https://www.nature.com/articles/s41586-019-1787-x) for sharing their data, a subset of which is used here.

___
# Tutorial Objectives
This is tutorial 3 of a 3-part series on different flavors of models used to understand neural data. In parts 1 and 2 we explored mechanisms that would produce the data. In this tutorial we will explore models and techniques that can potentially explain *why* the spiking data we have observed is produced the way it is.

To understand why different spiking behaviors may be beneficial, we will learn about the concept of entropy. Specifically, we will:

- Write code to compute formula for entropy, a measure of information
- Compute the entropy of a number of toy distributions
- Compute the entropy of spiking activity from the Steinmetz dataset

```python
#@title Video 1: “Why” models
from IPython.display import YouTubeVideo
video = YouTubeVideo(id='OOIDEr1e5Gg', 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=OOIDEr1e5Gg

```python
# Q1: Circular argument b/c in "How" Model assumes Poisson input to start with
# Q2: In "typical" models of neural coding, we assume a rate code and ISI is treated as nusiance or nosie
```

# Setup

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

```python
#@title Figure Settings
import ipywidgets as widgets #interactive display

%matplotlib inline
%config InlineBackend.figure_format = 'retina'
plt.style.use("https://raw.githubusercontent.com/NeuromatchAcademy/course-content/master/nma.mplstyle")

```

```python
#@title Helper Functions

def plot_pmf(pmf,isi_range):
  """Plot the probability mass function."""
  ymax = max(0.2, 1.05 * np.max(pmf))
  pmf_ = np.insert(pmf, 0, pmf[0])
  plt.plot(bins, pmf_, drawstyle="steps")
  plt.fill_between(bins, pmf_, step="pre", alpha=0.4)
  plt.title(f"Neuron {neuron_idx}")
  plt.xlabel("Inter-spike interval (s)")
  plt.ylabel("Probability mass")
  plt.xlim(isi_range);
  plt.ylim([0, ymax])
```

```python
#@title Download Data
import io
import requests
r = requests.get('https://osf.io/sy5xt/download')
if r.status_code != 200:
  print('Could not download data')
else:
  steinmetz_spikes = np.load(io.BytesIO(r.content), allow_pickle=True)['spike_times']
```

# Section 1: Optimization and Information

Neurons can only fire so often in a fixed period of time, as the act of emitting a spike consumes energy that is depleted and must eventually be replenished. To communicate effectively for downstream computation, the neuron would need to make good use of its limited spiking capability. This becomes an optimization problem:

What is the optimal way for a neuron to fire in order to maximize its ability to communicate information?

In order to explore this question, we first need to have a quantifiable measure for information. Shannon introduced the concept of entropy to do just that, and defined it as

\begin{align}
  H_b(X) &= -\sum_{x\in X} p(x) \log_b p(x)
\end{align}

where $H$ is entropy measured in units of base $b$ and $p(x)$ is the probability of observing the event $x$ from the set of all possible events in $X$. See the Appendix for a more detailed look at how this equation was derived.

The most common base of measuring entropy is $b=2$, so we often talk about *bits* of information, though other bases are used as well e.g. when $b=e$ we call the units *nats*.

First, let's explore how entropy changes between some simple discrete probability distributions. In the rest of this tutorial we will refer to these as probability mass functions (PMF), where $p(x_i)$ equals the $i^{th}$ value in an array, and mass refers to how much of the distribution is contained at that value.

For our first PMF, we will choose one where all of the probability mass is located in the middle of the distribution.

```python
n_bins = 50  # number of points supporting the distribution
x_range = (0, 1)  # will be subdivided evenly into bins corresponding to points

bins = np.linspace(*x_range, n_bins + 1)  # bin edges

pmf = np.zeros(n_bins)
pmf[len(pmf) // 2] = 1.0  # middle point has all the mass

# Since we already have a PMF, rather than un-binned samples, `plt.hist` is not
# suitable. Instead, we directly plot the PMF as a step function to visualize
# the histogram:
pmf_ = np.insert(pmf, 0, pmf[0])  # this is necessary to align plot steps with bin edges
plt.plot(bins, pmf_, drawstyle="steps")
# `fill_between` provides area shading
plt.fill_between(bins, pmf_, step="pre", alpha=0.4)
plt.xlabel("x")
plt.ylabel("p(x)")
plt.xlim(x_range)
plt.ylim(0, 1);
```

If we were to draw a sample from this distribution, we know exactly what we would get every time. Distributions where all the mass is concentrated on a single event are known as *deterministic*.

How much entropy is contained in a deterministic distribution?

## Exercise 1: Computing Entropy

Your first exercise is to implement a method that computes the entropy of a discrete probability distribution, given its mass function. Remember that we are interested in entropy in units of _bits_, so be sure to use the correct log function.

Recall that $\log(0)$ is undefined. When evaluated at $0$, NumPy log functions (such as `np.log2`) return `np.nan` ("Not a Number"). By convention, these undefined terms— which correspond to points in the distribution with zero mass—are excluded from the sum that computes the entropy.

```python
def entropy(pmf):
  """Given a discrete distribution, return the Shannon entropy in bits.

  This is a measure of information in the distribution. For a totally
  deterministic distribution, where samples are always found in the same bin,
  then samples from the distribution give no more information and the entropy
  is 0.

  For now this assumes `pmf` arrives as a well-formed distribution (that is,
  `np.sum(pmf)==1` and `not np.any(pmf < 0)`)

  Args:
    pmf (np.ndarray): The probability mass function for a discrete distribution
      represented as an array of probabilities.
  Returns:
    h (number): The entropy of the distribution in `pmf`.

  """
  p_val = pmf[pmf > 0]

  # implement the equation for Shannon entropy (in bits)
  h = -np.sum(p_val * np.log2(p_val))

  # return the absolute value (avoids getting a -0 result)
  return np.abs(h)

# Uncomment to test your entropy function
print(f"{entropy(pmf):.2f} bits")
```

    0.00 bits

We expect zero surprise from a deterministic distribution. If we had done this calculation by hand, it would simply be

$-1\log_2 1 = -0=0$

Note that changing the location of the peak (i.e. the point and bin on which all the mass rests) doesn't alter the entropy. The entropy is about how predictable a sample is with respect to a distribution. A single peak is deterministic regardless of which point it sits on.

```python
pmf = np.zeros(n_bins)
pmf[2] = 1.0  # arbitrary point has all the mass

pmf_ = np.insert(pmf, 0, pmf[0])
plt.plot(bins, pmf_, drawstyle="steps")
plt.fill_between(bins, pmf_, step="pre", alpha=0.4)
plt.xlabel("x")
plt.ylabel("p(x)")
plt.xlim(x_range)
plt.ylim(0, 1);
```

What about a distribution with mass split equally between two points?

```python
pmf = np.zeros(n_bins)
pmf[len(pmf) // 3] = 0.5
pmf[2 * len(pmf) // 3] = 0.5

pmf_ = np.insert(pmf, 0, pmf[0])
plt.plot(bins, pmf_, drawstyle="steps")
plt.fill_between(bins, pmf_, step="pre", alpha=0.4)
plt.xlabel("x")
plt.ylabel("p(x)")
plt.xlim(x_range)
plt.ylim(0, 1);

print("Entropy: %.2f" % entropy(pmf))
```

Here, the entropy calculation is

$-(0.5 \log_2 0.5 + 0.5\log_2 0.5)=1$

There is 1 bit of entropy. This means that before we take a random sample, there is 1 bit of uncertainty about which point in the distribution the sample will fall on: it will either be the first peak or the second one.

Likewise, if we make one of the peaks taller (i.e. its point holds more of the probability mass) and the other one shorter, the entropy will decrease because of the increased certainty that the sample will fall on one point and not the other:

$-(0.2 \log_2 0.2 + 0.8\log_2 0.8)\approx 0.72$

Try changing the definition of the number and weighting of peaks, and see how the entropy varies.

If we split the probability mass among even more points, the entropy continues to increase. Let's derive the general form for $N$ points of equal mass, where $p_i=p=1/N$:

\begin{align}
 -\sum_i p_i \log_b p_i&= -\sum_i^N \frac{1}{N} \log_b \frac{1}{N}\\
                       &= -\log_b \frac{1}{N} \\
                       &= \log_b N
\end{align}
$$$$

If we have $N$ discrete points, the _uniform distribution_ (where all points have equal mass) is the distribution with the highest entropy: $\log_b N$. This upper bound on entropy is useful when considering binning strategies, as any estimate of entropy over $N$ discrete points (or bins) must be in the interval $[0, \log_b N]$.

```python
pmf = np.ones(n_bins) / n_bins  # [1/N] * N

pmf_ = np.insert(pmf, 0, pmf[0])
plt.plot(bins, pmf_, drawstyle="steps")
plt.fill_between(bins, pmf_, step="pre", alpha=0.4)
plt.xlabel("x")
plt.ylabel("p(x)")
plt.xlim(x_range);
plt.ylim(0, 1);
```

Here, there are 50 points and the entropy of the uniform distribution is $\log_2 50\approx 5.64$. If we construct _any_ discrete distribution $X$ over 50 points (or bins) and calculate an entropy of $H_2(X)>\log_2 50$, something must be wrong with our implementation of the discrete entropy computation.

# Section 2: Information, neurons, and spikes

```python
#@title Video 2: Entropy of different distributions
from IPython.display import YouTubeVideo
video = YouTubeVideo(id='o6nyrx3KH20', 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=o6nyrx3KH20

Recall the discussion of spike times and inter-spike intervals (ISIs) from Tutorial 1. What does the information content (or distributional entropy) of these measures say about our theory of nervous systems?

We'll consider three hypothetical neurons that all have the same mean ISI, but with different distributions:

1. Deterministic
2. Uniform
3. Exponential

Fixing the mean of the ISI distribution is equivalent to fixing its inverse: the neuron's mean firing rate. If a neuron has a fixed energy budget and each of its spikes has the same energy cost, then by fixing the mean firing rate, we are normalizing for energy expenditure. This provides a basis for comparing the entropy of different ISI distributions. In other words: if our neuron has a fixed budget, what ISI distribution should it express (all else being equal) to maximize the information content of its outputs?

Let's construct our three distributions and see how their entropies differ.

```python
n_bins = 50
mean_isi = 0.025
isi_range = (0, 0.25)

bins = np.linspace(*isi_range, n_bins + 1)
mean_idx = np.searchsorted(bins, mean_isi)

# 1. all mass concentrated on the ISI mean
pmf_single = np.zeros(n_bins)
pmf_single[mean_idx] = 1.0

# 2. mass uniformly distributed about the ISI mean
pmf_uniform = np.zeros(n_bins)
pmf_uniform[0:2*mean_idx] = 1 / (2 * mean_idx)

# 3. mass exponentially distributed about the ISI mean
pmf_exp = stats.expon.pdf(bins[1:], scale=mean_isi)
pmf_exp /= np.sum(pmf_exp)
```

```python
#@title
#@markdown Run this cell to plot the three PMFs
fig, axes = plt.subplots(ncols=3, figsize=(18, 5))

dists =  [# (subplot title, pmf, ylim)
          ("(1) Deterministic", pmf_single, (0, 1.05)),
          ("(1) Uniform", pmf_uniform, (0, 1.05)),
          ("(1) Exponential", pmf_exp, (0, 1.05))]

for ax, (label, pmf_, ylim) in zip(axes, dists):
  pmf_ = np.insert(pmf_, 0, pmf_[0])
  ax.plot(bins, pmf_, drawstyle="steps")
  ax.fill_between(bins, pmf_, step="pre", alpha=0.4)
  ax.set_title(label)
  ax.set_xlabel("Inter-spike interval (s)")
  ax.set_ylabel("Probability mass")
  ax.set_xlim(isi_range);
  ax.set_ylim(ylim);
```

```python
print(
  f"Deterministic: {entropy(pmf_single):.2f} bits",
  f"Uniform: {entropy(pmf_uniform):.2f} bits",
  f"Exponential: {entropy(pmf_exp):.2f} bits",
  sep="\n",
)
```

    Deterministic: 0.00 bits
    Uniform: 3.32 bits
    Exponential: 3.77 bits

```python
#@title Video 3: Probabilities from histogram
from IPython.display import YouTubeVideo
video = YouTubeVideo(id='e2U_-07O9jo', 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=e2U_-07O9jo

In the previous example we created the PMFs by hand to illustrate idealized scenarios. How would we compute them from data recorded from actual neurons?

One way is to convert the ISI histograms we've previously computed into discrete probability distributions using the following equation:

\begin{align}
p_i = \frac{n_i}{\sum\nolimits_{i}n_i}
\end{align}

where $p_i$ is the probability of an ISI falling within a particular interval $i$ and $n_i$ is the count of how many ISIs were observed in that interval.

### Exercise 2: Probabilty Mass Function

Your second exercise is to implement a method that will produce a probability mass function from an array of ISI bin counts.

To verify your solution, we will compute the probability distribution of ISIs from real neural data taken from the Steinmetz dataset.

```python
neuron_idx = 283

isi = np.diff(steinmetz_spikes[neuron_idx])
bins = np.linspace(*isi_range, n_bins + 1)
counts, _ = np.histogram(isi, bins)
```

```python
def pmf_from_counts(counts):
  """Given counts, normalize by the total to estimate probabilities."""

  pmf = counts / np.sum(counts)

  return pmf

# Uncomment when ready to test your function
pmf = pmf_from_counts(counts)
plot_pmf(pmf,isi_range)
```

# Section 3: Calculate entropy from a PMF

```python
#@title Video 4: Calculating entropy from pmf
from IPython.display import YouTubeVideo
video = YouTubeVideo(id='Xjy-jj-6Oz0', 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=Xjy-jj-6Oz0

Now that we have the probability distribution for the actual neuron spiking activity, we can calculate its entropy.

```python
print(f"Entropy for Neuron {neuron_idx}: {entropy(pmf):.2f} bits")
```

    Entropy for Neuron 283: 3.36 bits

## Interactive Demo: Entropy of neurons

We can combine the above distribution plot and entropy calculation with an interactive widget to explore how the different neurons in the dataset vary in spiking activity and relative information. Note that the mean firing rate across neurons is not fixed, so some neurons with a uniform ISI distribution may have higher entropy than neurons with a more exponential distribution.

```python
#@title
#@markdown **Run the cell** to enable the sliders.

def _pmf_from_counts(counts):
  """Given counts, normalize by the total to estimate probabilities."""
  pmf = counts / np.sum(counts)
  return pmf

def _entropy(pmf):
  """Given a discrete distribution, return the Shannon entropy in bits."""
  # remove non-zero entries to avoid an error from log2(0)
  pmf = pmf[pmf > 0]
  h = -np.sum(pmf * np.log2(pmf))
  # absolute value applied to avoid getting a -0 result
  return np.abs(h)

@widgets.interact(neuron=widgets.IntSlider(0, min=0, max=(len(steinmetz_spikes)-1)))
def steinmetz_pmf(neuron):
  """ Given a neuron from the Steinmetz data, compute its PMF and entropy """
  isi = np.diff(steinmetz_spikes[neuron])
  bins = np.linspace(*isi_range, n_bins + 1)
  counts, _ = np.histogram(isi, bins)
  pmf = _pmf_from_counts(counts)

  plot_pmf(pmf,isi_range)
  plt.title(f"Neuron {neuron}: H = {_entropy(pmf):.2f} bits")

```

    interactive(children=(IntSlider(value=0, description='neuron', max=733), Output()), _dom_classes=('widget-inte…

---
# Summary

```python
#@title Video 5: Summary of model types
from IPython.display import YouTubeVideo
video = YouTubeVideo(id='X4K2RR5qBK8', 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=X4K2RR5qBK8

Congratulations! You've finished your first NMA tutorial. In this 3 part tutorial series, we used different types of models to understand the spiking behavior of neurons recorded in the Steinmetz data set.

 - We used "what" models to discover that the ISI distribution of real neurons is closest to an exponential distribution
 - We used "how" models to discover that balanced excitatory and inbhitiory inputs, coupled with a leaky membrane, can give rise to neuronal spiking with exhibiting such an exponential ISI distribution
 - We used "why" models to discover that exponential ISI distributions contain the most information when the mean spiking is constrained

# Bonus

### The foundations for Entropy

In his foundational [1948 paper](https://en.wikipedia.org/wiki/A_Mathematical_Theory_of_Communication) on information theory, Claude Shannon began with three criteria for a function $H$ defining the entropy of a discrete distribution of probability masses $p_i\in p(X)$ over the points $x_i\in X$:
1. $H$ should be continuous in the $p_i$.
  - That is, $H$ should change smoothly in response to smooth changes to the mass $p_i$ on each point $x_i$.
2. If all the points have equal shares of the probability mass, $p_i=1/N$, $H$ should be a non-decreasing function of $N$.
  - That is, if $X_N$ is the support with $N$ discrete points and $p(x\in X_N)$ assigns constant mass to each point, then $H(X_1) < H(X_2) < H(X_3) < \dots$
3. $H$ should be preserved by (invariant to) the equivalent (de)composition of distributions.
  - For example (from Shannon's paper) if we have a discrete distribution over three points with masses $(\frac{1}{2},\frac{1}{3},\frac{1}{6})$, then their entropy can be represented in terms of a direct choice between the three and calculated $H(\frac{1}{2},\frac{1}{3},\frac{1}{6})$. However, it could also be represented in terms of a series of two choices:
    1. either we sample the point with mass $1/2$ or not (_not_ is the other $1/2$, whose subdivisions are not given in the first choice),
    2. if (with probability $1/2$) we _don't_ sample the first point, we sample one of the two remaining points, masses $1/3$ and $1/6$.

    Thus in this case we require that $H(\frac{1}{2},\frac{1}{3},\frac{1}{6})=H(\frac{1}{2},\frac{1}{2}) + \frac{1}{2}H(\frac{1}{3}, \frac{1}{6})$

There is a unique function (up to a linear scaling factor) which satisfies these 3 requirements:

\begin{align}
  H_b(X) &= -\sum_{x\in X} p(x) \log_b p(x)
\end{align}

Where the base of the logarithm $b>1$ controls the units of entropy. The two most common cases are $b=2$ for units of _bits_, and $b=e$ for _nats_.

We can view this function as the expectation of the self-information over a distribution:

$$H_b(X) = \mathbb{E}_{x\in X} \left[I_b(x)\right]$$

$$I_b(x)=-\log_b p(x)$$

Self-information is just the negative logarithm of probability, and is a measure of how surprising an event sampled from the distribution would be. Events with $p(x)=1$ are certain to occur, and their self-information is zero (as is the entropy of the distribution they compose) meaning they are totally unsurprising. The smaller the probability of an event, the higher its self-information, and the more surprising the event would be to observe.


# STIRAP in a 3-level system
STIRAP (STImulated Raman Adiabatic Passage, see e.g. [Shore1998](https://journals.aps.org/rmp/pdf/10.1103/RevModPhys.70.1003)) is a method for adiabatically transferring the population of a quantum system from one state to another by using two drive fields coupled to an intermediate state without actually ever populating the intermediate state. The benefits over e.g. two Rabi pulses, are that since STIRAP is an adiabatic process, it is relatively easy (I've been told) to make it highly efficient. The other key benefit is that the intermediate state can be an unstable state, yet there is no population loss since it is never populated.

This notebook sets up a 3-level system and relevant couplings using the `toy_models` package and then time evolves the system using `QuTiP` to simulate STIRAP. I'll be following the definitions of Shore1998 as best as I can. The level diagram from the paper is shown below.

## Imports

```python
%load_ext autoreload
%autoreload 2

import matplotlib.pyplot as plt
plt.style.use("ggplot")
import numpy as np
import qutip
from sympy import Symbol

from toy_systems.couplings import ToyCoupling, ToyEnergy
from toy_systems.decays import ToyDecay
from toy_systems.hamiltonian import Hamiltonian
from toy_systems.quantum_system import QuantumSystem
from toy_systems.states import Basis, BasisState, ToyQuantumNumbers
from toy_systems.visualization import Visualizer
```

## Set up states and basis
We start by defining the three states of the system: we'll have two ground states (i.e. states that don't decay) $|1\rangle$ and $|3\rangle$, and one excited state $|2\rangle$, which we will later set to have a decay to an additional state $|4\rangle$ representing all decays out of the system:

```python
# Define states
s1 = BasisState(qn=ToyQuantumNumbers(label="1"))
s2 = BasisState(qn=ToyQuantumNumbers(label="2"))
s3 = BasisState(qn=ToyQuantumNumbers(label="3"))
s4 = BasisState(qn=ToyQuantumNumbers(label="4")) # A target state for decays from |2>

# Define basis
basis = Basis((s1, s2, s3, s4))
basis.print()
```

    |0> = |1>
    |1> = |2>
    |2> = |3>
    |3> = |4>

## Define energies, couplings and decays
I'm going to define the system in the rotating frame as given in [Shore1998](https://journals.aps.org/rmp/pdf/10.1103/RevModPhys.70.1003) so that the Hamiltonian doesn't have any quickly rotating terms of the form $e^{i\omega t}$.

The Hamiltonian I'm trying to produce is shown below (with $\hbar = 1$):

### Energies

```python
Δp = Symbol('Delta_p') # Detuning for pump beam
Δs = Symbol('Delta_s') # Detuning for Stokes beam
E1 = ToyEnergy([s1], 0)
E2 = ToyEnergy([s2], Δp)

# The energy for state |3> needs to be defined in two parts since it contains two sympy.Symbols
E3p = ToyEnergy([s3], Δp)
E3s = ToyEnergy([s3], -Δs)
```

### Couplings

```python
Ωp = Symbol('Omega_p') # Drive field Rabi rate for pump beam
Ωs = Symbol('Omega_s') # Drive field Rabi rate for Stokes beam

coupling_p = ToyCoupling(s1,s2,Ωp/2, time_dep = "exp(-(t+t_p)**2/(2*sigma_p**2))", time_args= {"t_p":-1, "sigma_p":1})
coupling_s = ToyCoupling(s2,s3,Ωs/2, time_dep = "exp(-(t+t_s)**2/(2*sigma_s**2))", time_args= {"t_s":1, "sigma_s":1})
```

### Decays
Defining a decay from $|2\rangle$ to $|4\rangle$ :

```python
decay = ToyDecay(s2, ground = s4, gamma = Symbol("Gamma"))
```

### Define a QuantumSystem
The QuantumSystem object combines the basis, Hamiltonian and decays to make setting parameters for time evolution using QuTiP more convenient.

```python
# Define the system
system = QuantumSystem(
    basis=basis,
    couplings=[E1, E2, E3p, E3s, coupling_p, coupling_s],
    decays=[decay],
)

# Get representations of the Hamiltonian and the decays that will be accepted by qutip
Hqobj, c_qobj = system.get_qobjs()

visualizer = Visualizer(system, vertical={"label":10}, horizontal={"label":50})
```

## Time-evolution using `QuTiP`
We can now see if time evolving the system results in something resembling STIRAP. The key to success is to choose the parameters well. Shore gives us the rule of thumb that we should have $\sqrt{\Omega_p^2 + \Omega_s^2}\tau > 10$ where $\tau$ is proportional to the time overlap of the Stokes and pump pulse. In practice it seems that taking the centers of the Gaussians to be separated by $2\sigma$ works pretty well. The broader the Gaussians are (i.e. larger $\sigma$), the more adiabatic the process, which results in less population in the intermediate state and therefore less loss. I'm taking both pulses to have the same parameters for simplicity (except they occur at different times of course).

```python
# Get a pointer to the time-evolution arguments
args = Hqobj.args
print("Keys for setting arguments:")
print(f"args = {args}")
```

    Keys for setting arguments:
    args = {'Delta_p': 1, 'Delta_s': 1, 't_p': -1, 'sigma_p': 1, 'Omega_p': 1, 't_s': 1, 'sigma_s': 1, 'Omega_s': 1, 'Gamma': 1}

```python
# Generate a Qobj representing the initial state
psi0 = (1*s1).qobj(basis)

# Make operators for getting the probability of being in each state
P_1_op = qutip.Qobj((1*s1).density_matrix(basis), type = "oper")
P_2_op = qutip.Qobj((1*s2).density_matrix(basis), type = "oper")
P_3_op = qutip.Qobj((1*s3).density_matrix(basis), type = "oper")
P_4_op = qutip.Qobj((1*s4).density_matrix(basis), type = "oper")

# Set the parameters for the system

# Good STIRAP
Omega = 10
t0 = 10
sigma = 10
Delta = 0

# Bad STIRAP
# Omega = 5
# t0 = 1
# sigma = 1
# Delta = 0

args["Delta_p"] = Delta
args["Omega_p"] = Omega
args["sigma_p"] = sigma
args["t_p"] = -t0

args["Delta_s"] = Delta
args["Omega_s"] = Omega
args["sigma_s"] = sigma
args["t_s"] = t0

# Times at which result is requested
times = np.linspace(-5*sigma,5*sigma,1001)

# Setting the max_step is sometimes necessary
options = qutip.solver.Options(method = 'adams', nsteps=10000, max_step=1e0)

# Setup a progress bar
pb = qutip.ui.progressbar.EnhancedTextProgressBar()

# Run the time-evolution
result = qutip.mesolve(Hqobj, psi0, times, c_ops = c_qobj, e_ops = [P_1_op, P_2_op, P_3_op, P_4_op],
                      progress_bar=pb, options = options)
```

     Total run time:   0.48s*] Elapsed 0.48s / Remaining 00:00:00:00[*********40%           ] Elapsed 0.16s / Remaining 00:00:00:00

```python
fig, ax = plt.subplots(figsize = (16,9))
ln = []
ln+=ax.plot(times, result.expect[0], label = "P_1")
ln+=ax.plot(times, result.expect[1], label = "P_2")
ln+=ax.plot(times, result.expect[2], label = "P_3")
ln+=ax.plot(times, result.expect[3], label = "P_4")
ax.set_title("STIRAP", fontsize = 18)
ax.set_xlabel("Time / (1/Γ)", fontsize = 16)
ax.set_ylabel("Population in each state", fontsize = 16)

axc = ax.twinx()
ln+=coupling_p.plot_time_dep(times, args, ax=axc, ls = '--', c = 'k', lw = 1, label = 'Pump')
ln+=coupling_s.plot_time_dep(times, args, ax=axc, ls = ':', c = 'k', lw = 1, label = 'Stokes')

ax.legend(ln, [l.get_label() for l in ln], fontsize = 16)

print(f"Transfer efficiency: {result.expect[2][-1]*100:.1f} %")
```

```python

```

examples/STIRAP in a 3-level system.ipynb