#Week 4: Deep Multi-layer Perceptrons (MLPs)
**By Arash Ash, adapted from Neuromatch Academy contents**

---
# Tutorial Objectives
In this tutorial, we will dive deeper into MLPs and see more of their mathematical and practical aspects. Today we are going to see why MLPs:

* can model your problem as a classification task
* can be deep or wide and which is better?
* dependant on transfer functions
* sensitive to initialization

---
# Setup

This is a GPU free notebook!

```python
# Imports
import pathlib

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

import torch.nn as nn
import torch.optim as optim

from torchvision.utils import make_grid
import torchvision.transforms as transforms
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader, TensorDataset

from tqdm.auto import tqdm
from IPython.display import display
```

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

```python
# @title Plotting functions
def imshow(img):
  img = img / 2 + 0.5     # unnormalize
  npimg = img.numpy()
  plt.imshow(np.transpose(npimg, (1, 2, 0)))
  plt.axis(False)
  plt.show()

def sample_grid(M=500, x_max=2.0):
  ii, jj = torch.meshgrid(torch.linspace(-x_max, x_max,M),
                          torch.linspace(-x_max, x_max, M))
  X_all = torch.cat([ii.unsqueeze(-1),
                     jj.unsqueeze(-1)],
                     dim=-1).view(-1, 2)
  return X_all

def plot_decision_map(X_all, y_pred, X_test, y_test,
                      M=500, x_max=2.0, eps=1e-3):
  decision_map = torch.argmax(y_pred, dim=1)

  for i in range(len(X_test)):
    indeces = (X_all[:, 0] - X_test[i, 0])**2 + (X_all[:, 1] - X_test[i, 1])**2 < eps    # [TO-DO]
    decision_map[indeces] = (K + y_test[i]).long()

  decision_map = decision_map.view(M, M).cpu()
  plt.imshow(decision_map, extent=[-x_max, x_max, -x_max, x_max], cmap='jet')
  plt.plot()
```

```python
# @title Set random seed

# @markdown Executing `set_seed(seed=seed)` you are setting the seed

# for DL its critical to set the random seed so that students can have a
# baseline to compare their results to expected results.
# Read more here: https://pytorch.org/docs/stable/notes/randomness.html

# Call `set_seed` function in the exercises to ensure reproducibility.
import random
import torch

def set_seed(seed=None, seed_torch=True):
  if seed is None:
    seed = np.random.choice(2 ** 32)
  random.seed(seed)
  np.random.seed(seed)
  if seed_torch:
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True

  print(f'Random seed {seed} has been set.')

# In case that `DataLoader` is used
def seed_worker(worker_id):
  worker_seed = torch.initial_seed() % 2**32
  np.random.seed(worker_seed)
  random.seed(worker_seed)
```

```python
# @title Set device (GPU or CPU). Execute `set_device()`
# especially if torch modules used.

# inform the user if the notebook uses GPU or CPU.

def set_device():
  device = "cuda" if torch.cuda.is_available() else "cpu"
  if device != "cuda":
    print("GPU is not enabled in this notebook. \n"
          "If you want to enable it, in the menu under `Runtime` -> \n"
          "`Hardware accelerator.` and select `GPU` from the dropdown menu")
  else:
    print("GPU is enabled in this notebook. \n"
          "If you want to disable it, in the menu under `Runtime` -> \n"
          "`Hardware accelerator.` and select `None` from the dropdown menu")

  return device
```

```python
SEED = 2021
set_seed(seed=SEED)
DEVICE = set_device()
```

    Random seed 2021 has been set.
    GPU is not enabled in this notebook.
    If you want to enable it, in the menu under `Runtime` ->
    `Hardware accelerator.` and select `GPU` from the dropdown menu

```python
# @title Download of the Animal Faces dataset
# @markdown Animal faces consists of 16,130 32x32 images belonging to 3 classes
import requests, os
from zipfile import ZipFile

print("Start downloading and unzipping `AnimalFaces` dataset...")
name = 'AnimalFaces32x32'
fname = f"{name}.zip"
url = f"https://osf.io/kgfvj/download"
r = requests.get(url, allow_redirects=True)
with open(fname, 'wb') as fh:
  fh.write(r.content)

with ZipFile(fname, 'r') as zfile:
  zfile.extractall(f"./{name}")

if os.path.exists(fname):
  os.remove(fname)
else:
  print(f"The file {fname} does not exist")

os.chdir(name)
print("Download completed.")
```

    Start downloading and unzipping `AnimalFaces` dataset...
    Download completed.

## Section 1: Classification with MLPs

The main loss function we could use out of the box for multi-class classification for `N` samples and `C` number of classes is:
* CrossEntropyLoss:
This criterion expects a batch of predictions `x` with shape `(N, C)` and class index in the range $[0, C-1]$ as the target (label) for each `N` samples, hence a batch of `labels` with shape `(N, )`. There are other optional parameters like class weights and class ignores. Feel free to check the PyTorch documentation [here](https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html) for more detail. Additionally, [here](https://sparrow.dev/cross-entropy-loss-in-pytorch/) you can learn where is appropriate to use the CrossEntropyLoss.

To get CrossEntropyLoss of a sample $i$, we could first calculate $-\log(\text{softmax(x}))$ and then take the element corresponding to $\text { labels }_i$ as the loss. However, due to numerical stability, we implement this more stable equivalent form,

\begin{equation}
\operatorname{loss}(x_i, \text { labels }_i)=-\log \left(\frac{\exp (x[\text { labels }_i])}{\sum_{j} \exp (x[j])}\right)=-x_i[\text { labels }_i]+\log \left(\sum_{j=1}^C \exp (x_i[j])\right)
\end{equation}

### Coding Exercise 2.1: Implement Batch Cross Entropy Loss

To recap, since we will be doing batch learning, we'd like a loss function that given:
* a batch of predictions `x` with shape `(N, C)`
* a batch of `labels` with shape `(N, )` that ranges from `0` to `C-1`

returns the average loss $L$ calculated according to:

\begin{align}
loss(x_i, \text { labels }_i) &= -x_i[\text { labels }_i]+\log \left(\sum_{j=1}^C \exp (x_i[j])\right) \\
L &= \frac{1}{N} \sum_{i=1}^{N}{loss(x_i, \text { labels }_i)}
\end{align}

Steps:

1.   Use indexing operation to get predictions of class corresponding to the labels (i.e., $x_i[\text { labels }_i]$)
2.   Compute $loss(x_i, \text { labels }_i)$ vector (`losses`) using `torch.log()` and `torch.exp()` without Loops!
3. Return the average of the loss vector

```python
def cross_entropy_loss(x, labels):
  # x is the model predictions we'd like to evaluate using lables
  x_of_labels = torch.zeros(len(labels))
  ####################################################################
  # Fill in missing code below (...),
  # then remove or comment the line below to test your function
  raise NotImplementedError("Cross Entropy Loss")
  ####################################################################
  # 1. prediction for each class corresponding to the label
  for i, label in enumerate(labels):
    x_of_labels[i] = x[i, label]
  # 2. loss vector for the batch
  losses = ...
  # 3. Return the average of the loss vector
  avg_loss = ...

  return avg_loss

labels = torch.tensor([0, 1])
x = torch.tensor([[10.0, 1.0, -1.0, -20.0],  # correctly classified
                  [10.0, 10.0, 2.0, -10.0]])  # Not correctly classified
CE = nn.CrossEntropyLoss()
pytorch_loss = CE(x, labels).item()
## Uncomment below to test your function
# our_loss = cross_entropy_loss(x, labels).item()
# print(f'Our CE loss: {our_loss:0.8f}, Pytorch CE loss: {pytorch_loss:0.8f}')
# print(f'Difference: {np.abs(our_loss - pytorch_loss):0.8f}')
```

## Section 1.2: Spiral classification dataset
Before we could start optimizing these loss functions, we need a dataset!

Let's turn this fancy-looking equation into a classification dataset

\begin{equation}
\begin{array}{c}
X_{k}(t)=t\left(\begin{array}{c}
\sin \left[\frac{2 \pi}{K}\left(2 t+k-1\right)\right]+\mathcal{N}\left(0, \sigma\right) \\
\cos \left[\frac{2 \pi}{K}\left(2 t+k-1\right)\right]+\mathcal{N}\left(0, \sigma\right)
\end{array}\right)
\end{array}, \quad 0 \leq t \leq 1, \quad k=1, \ldots, K
\end{equation}

```python
def create_spiral_dataset(K, sigma, N):

  # Initialize t, X, y
  t = torch.linspace(0, 1, N)
  X = torch.zeros(K*N, 2)
  y = torch.zeros(K*N)

  # Create data
  for k in range(K):
    X[k*N:(k+1)*N, 0] = t*(torch.sin(2*np.pi/K*(2*t+k)) + sigma*torch.randn(N))
    X[k*N:(k+1)*N, 1] = t*(torch.cos(2*np.pi/K*(2*t+k)) + sigma*torch.randn(N))
    y[k*N:(k+1)*N] = k

  return X, y

# Set parameters
K = 4
sigma = 0.16
N = 1000

set_seed(seed=SEED)
X, y = create_spiral_dataset(K, sigma, N)
plt.scatter(X[:, 0], X[:, 1], c = y)
plt.show()
```

## Section 1.3: Training and Evaluation

### Coding Exercise 2.3: Implement it for a classfication task
Now that we have the Spiral dataset and a loss function, it's your turn to implement a simple train/test split for training and validation.

Steps to follow:
  * Dataset shuffle
  * Train/Test split (20% for test)
  * Dataloader definition
  * Training and Evaluation

```python
def shuffle_and_split_data(X, y, seed):
  # set seed for reproducibility
  torch.manual_seed(seed)
  # Number of samples
  N = X.shape[0]
  ####################################################################
  # Fill in missing code below (...),
  # then remove or comment the line below to test your function
  raise NotImplementedError("Shuffle & split data")
  ####################################################################
  # Shuffle data
  shuffled_indices = ...   # get indices to shuffle data, could use torch.randperm
  X = X[shuffled_indices]
  y = y[shuffled_indices]

  # Split data into train/test
  test_size = ...    # assign test datset size using 20% of samples
  X_test = X[:test_size]
  y_test = y[:test_size]
  X_train = X[test_size:]
  y_train = y[test_size:]

  return X_test, y_test, X_train, y_train

## Uncomment below to test your function
# X_test, y_test, X_train, y_train = shuffle_and_split_data(X, y, seed=SEED)
# plt.scatter(X_test[:, 0], X_test[:, 1], c=y_test)
# plt.title('Test data')
# plt.show()
```

```python
# to_remove solution
def shuffle_and_split_data(X, y, seed):
  # set seed for reproducibility
  torch.manual_seed(seed)
  # Number of samples
  N = X.shape[0]
  # Shuffle data
  shuffled_indices = torch.randperm(N)   # get indices to shuffle data, could use torch.randperm
  X = X[shuffled_indices]
  y = y[shuffled_indices]

  # Split data into train/test
  test_size = int(0.2 * N)    # assign test datset size using 20% of samples
  X_test = X[:test_size]
  y_test = y[:test_size]
  X_train = X[test_size:]
  y_train = y[test_size:]

  return X_test, y_test, X_train, y_train

## Uncomment below to test your function
X_test, y_test, X_train, y_train = shuffle_and_split_data(X, y, seed=SEED)
with plt.xkcd():
  plt.scatter(X_test[:, 0], X_test[:, 1], c=y_test)
  plt.title('Test data')
  plt.show()
```

And we need to make a Pytorch data loader out of it. Data loading in PyTorch can be separated in 2 parts:
* Data must be wrapped on a Dataset parent class where the methods __getitem__ and __len__ must be overrided. Note that at this point the data is not loaded on memory. PyTorch will only load what is needed to the memory. Here `TensorDataset` does this for us directly.
* Use a Dataloader that will actually read the data in batches and put into memory. Also, the option of `num_workers > 0` allows multithreading, which prepares multiple batches in the queue to speed things up.

```python
g_seed = torch.Generator()
g_seed.manual_seed(SEED)

batch_size = 128
test_data = TensorDataset(X_test, y_test)
test_loader = DataLoader(test_data, batch_size=batch_size,
                         shuffle=False, num_workers=2,
                         worker_init_fn=seed_worker,
                         generator=g_seed)

train_data = TensorDataset(X_train, y_train)
train_loader = DataLoader(train_data, batch_size=batch_size, drop_last=True,
                        shuffle=True, num_workers=2,
                         worker_init_fn=seed_worker,
                         generator=g_seed)
```

Let's write a general-purpose training and evaluation code and keep it in our pocket for next tutorial as well. So make sure you review it to see what it does.

Note that `model.train()` tells your model that you are training the model. So layers like dropout, batchnorm etc. which behave different on the train and test procedures know what is going on and hence can behave accordingly. And to turn off training mode we set `model.eval()`

```python
def train_test_classification(net, criterion, optimizer, train_loader,
                              test_loader, num_epochs=1, verbose=True,
                              training_plot=False, device='cpu'):

  net.train()
  training_losses = []
  for epoch in tqdm(range(num_epochs)):  # loop over the dataset multiple times
    running_loss = 0.0
    for i, data in enumerate(train_loader, 0):
      # get the inputs; data is a list of [inputs, labels]
      inputs, labels = data
      inputs = inputs.to(device).float()
      labels = labels.to(device).long()

      # zero the parameter gradients
      optimizer.zero_grad()

      # forward + backward + optimize
      outputs = net(inputs)

      loss = criterion(outputs, labels)
      loss.backward()
      optimizer.step()

      # print statistics
      if verbose:
        training_losses += [loss.item()]

  net.eval()
  def test(data_loader):
    correct = 0
    total = 0
    for data in data_loader:
      inputs, labels = data
      inputs = inputs.to(device).float()
      labels = labels.to(device).long()

      outputs = net(inputs)
      _, predicted = torch.max(outputs, 1)
      total += labels.size(0)
      correct += (predicted == labels).sum().item()

    acc = 100 * correct / total
    return total, acc

  train_total, train_acc = test(train_loader)
  test_total, test_acc = test(test_loader)

  if verbose:
    print(f"Accuracy on the {train_total} training samples: {train_acc:0.2f}")
    print(f"Accuracy on the {test_total} testing samples: {test_acc:0.2f}")

  if training_plot:
    plt.plot(training_losses)
    plt.xlabel('Batch')
    plt.ylabel('Training loss')
    plt.show()

  return train_acc, test_acc
```

### Think! 1.3.1: What's the point of .eval() and .train()?

Is it necessary to use `net.train()` and `net.eval()` for our MLP model? why?

Now let's put everything together and train your first deep-ish model!

```python
# MLP from Week 3

class Net(nn.Module):
  def __init__(self, actv, input_feature_num, hidden_unit_nums, output_feature_num):
    super(Net, self).__init__()
    self.input_feature_num = input_feature_num # save the input size for reshapinng later
    self.mlp = nn.Sequential() # Initialize layers of MLP

    in_num = input_feature_num # initialize the temporary input feature to each layer
    for i in range(len(hidden_unit_nums)): # Loop over layers and create each one
      out_num = hidden_unit_nums[i] # assign the current layer hidden unit from list
      layer = nn.Linear(in_num, out_num) # use nn.Linear to define the layer
      in_num = out_num # assign next layer input using current layer output
      self.mlp.add_module(f"Linear_{i}", layer) # append layer to the model with a name

      actv_layer = eval(f"nn.{actv}") # Assign activation function (eval allows us to instantiate object from string)
      self.mlp.add_module(f"Activation_{i}", actv_layer) # append activation to the model with a name

    out_layer = nn.Linear(in_num, output_feature_num) # Create final layer
    self.mlp.add_module('Output_Linear', out_layer) # append the final layer

  def forward(self, x):
    # reshape inputs to (batch_size, input_feature_num)
    # just in case the input vector is not 2D, like an image!
    x = x.view(-1, self.input_feature_num)

    logits = self.mlp(x) # forward pass of MLP
    return logits
```

```python
set_seed(SEED)
net = Net('ReLU()', X_train.shape[1], [128], K).to(DEVICE)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), lr=1e-3)
num_epochs = 100

_, _ = train_test_classification(net, criterion, optimizer, train_loader,
                                 test_loader, num_epochs=num_epochs,
                                 training_plot=True, device=DEVICE)
```

And finally, let's visualize the learned decision-map. We know you're probably running out of time, so we won't make you write code now! But make sure you have reviewed it since we'll start with another visualization technique next time.

```python
def sample_grid(M=500, x_max=2.0):
  ii, jj = torch.meshgrid(torch.linspace(-x_max, x_max, M),
                          torch.linspace(-x_max, x_max, M))
  X_all = torch.cat([ii.unsqueeze(-1),
                     jj.unsqueeze(-1)],
                     dim=-1).view(-1, 2)
  return X_all

def plot_decision_map(X_all, y_pred, X_test, y_test,
                      M=500, x_max=2.0, eps=1e-3):
  decision_map = torch.argmax(y_pred, dim=1)

  for i in range(len(X_test)):
    indices = (X_all[:, 0] - X_test[i, 0])**2 + (X_all[:, 1] - X_test[i, 1])**2 < eps
    decision_map[indices] = (K + y_test[i]).long()

  decision_map = decision_map.view(M, M)
  plt.imshow(decision_map, extent=[-x_max, x_max, -x_max, x_max], cmap='jet')
  plt.show()
```

```python
X_all = sample_grid()
y_pred = net(X_all)
plot_decision_map(X_all, y_pred, X_test, y_test)
```

### Think! 1.3.2: Does it generalize well?
Do you think this model is performing well outside its training distribution? Why?

---
# Section 2: Wider vs deeper networks

```python
# @title Video 1: Deep Expressivity
from ipywidgets import widgets

out2 = widgets.Output()
with out2:
  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=f"BV19f4y157vG", width=854, height=480, fs=1)
  print("Video available at https://www.bilibili.com/video/{0}".format(video.id))
  display(video)

out1 = widgets.Output()
with out1:
  from IPython.display import YouTubeVideo
  video = YouTubeVideo(id=f"g8JuGrNk9ag", width=854, height=480, fs=1, rel=0)
  print("Video available at https://youtube.com/watch?v=" + video.id)
  display(video)

out = widgets.Tab([out1, out2])
out.set_title(0, 'Youtube')
out.set_title(1, 'Bilibili')

# display(out)
```

## Coding Exercise 1: Wide vs. Deep while keeping number of parameters same
Let's find the optimal number of hidden layers under a fixed number of parameters constraint!

But first, we need a model parameter counter. You could iterate over the model layers by calling `.parameters()` and then use `.numel()` to count the layer parameters. Also, you can use [`requires_grad`](https://pytorch.org/docs/stable/notes/autograd.html) attribute to make sure it's a trainable parameter. E.g.,

```python
x = torch.ones(10, 5, requires_grad=True)
```

After defining the counter function, we will step by step increase the depth and then iterate over the possible number of hidden units (assuming same for all hidden layers); then using our parameter counter choose the number of hidden units that results in overall close to `max_par_count` parameters.

```python
def run_depth_optimizer(max_par_count, max_hidden_layer, device):
  ####################################################################
  # Fill in all missing code below (...),
  # then remove or comment the line below to test your function
  raise NotImplementedError("Define the depth optimizer function")
  ###################################################################

  def count_parameters(model):
    par_count = 0
    for p in model.parameters():
      if p.requires_grad:
        par_count += ...
    return par_count

  # list of hidden layers to try
  hidden_layers = range(1, max_hidden_layer+1)

  # test test score list
  test_scores = []

  for hidden_layer in hidden_layers:
    # Initialize the hidden units in each hidden layer to be 1
    hidden_units = np.ones(hidden_layer, dtype=np.int)

    # Define the the with hidden units equal to 1
    wide_net = Net('ReLU()', X_train.shape[1], hidden_units, K).to(device)
    par_count = count_parameters(wide_net)

    # increment hidden_units and repeat until the par_count reaches the desired count
    while par_count < max_par_count:
      hidden_units += 1
      wide_net = Net('ReLU()', X_train.shape[1], hidden_units, K).to(device)
      par_count = ...

    # Train it
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(wide_net.parameters(), lr=1e-3)
    _, test_acc = train_test_classification(wide_net, criterion, optimizer,
                                            train_loader, test_loader,
                                            num_epochs=100, device=device)
    test_scores += [test_acc]

  return hidden_layers, test_scores

set_seed(seed=SEED)
max_par_count = 100
max_hidden_layer = 5
## Uncomment below to test your function
# hidden_layers, test_scores = run_depth_optimizer(max_par_count, max_hidden_layer, DEVICE)
# plt.xlabel('# of hidden layers')
# plt.ylabel('Test accuracy')
# plt.plot(hidden_layers, test_scores)
# plt.show()
```

    Random seed 2021 has been set.

## Think! 1: Why the tradeoff?
Here we see that there is a particular number of hidden layers that is optimum. Why do you think increasing hidden layers after a certain point hurt in this scenario?

## Section 1.1: Where Wide Fails
Let's use the same Spiral dataset generated before with two features. And then add more polynomial features (which makes the first layer wider). And finally, train a single Linear layer. We could use the same MLP network with no hidden layers (though it would not be called an MLP anymore!).

Note that we will add polynomial terms upto $P=50$ which means that for every $x_1^n x_2^m$ term, $n+m\leq P$. Now it's fun math excercise to prove why the total number of polynomial features upto $P$ becomes,

\begin{equation}
\text{# of terms} = \frac{(P+1)(P+2)}{2}
\end{equation}

Also, we don't need the polynomial term with degree zero (which is the constatnt term) since `nn.Linear` layers have bias terms. Therefore we will have one fewer polynomial feature.

```python
def run_poly_clasification(poly_degree, device='cpu', seed=0):

  def make_poly_features(poly_degree, X):
    # Define the number of polynomial features except the bias term
    num_features = (poly_degree + 1)*(poly_degree + 2) // 2 - 1
    poly_X = torch.zeros((X.shape[0], num_features))
    count = 0
    for i in range(poly_degree+1):
      for j in range(poly_degree+1):
         # no need to add zero degree since model has biases
        if j + i > 0:
          if j + i <= poly_degree:
            # Define the polynomial term
            poly_X[:, count] = X[:, 0]**i * X [:, 1]**j
            count += 1
    return poly_X, num_features

  poly_X_test, num_features = make_poly_features(poly_degree, X_test)
  poly_X_train, _ = make_poly_features(poly_degree, X_train)

  batch_size = 128

  g_seed = torch.Generator()
  g_seed.manual_seed(seed)
  poly_test_data = TensorDataset(poly_X_test, y_test)
  poly_test_loader = DataLoader(poly_test_data,
                                batch_size=batch_size,
                                shuffle=False,
                                num_workers=1,
                                worker_init_fn=seed_worker,
                                generator=g_seed)

  poly_train_data = TensorDataset(poly_X_train, y_train)
  poly_train_loader = DataLoader(poly_train_data,
                                 batch_size=batch_size,
                                 shuffle=True,
                                 num_workers=1,
                                 worker_init_fn=seed_worker,
                                 generator=g_seed)

  # define a linear model using MLP class
  poly_net = Net('ReLU()', num_features, [], K).to(device)

  # Train it!
  criterion = nn.CrossEntropyLoss()
  optimizer = optim.Adam(poly_net.parameters(), lr=1e-3)
  _, _ = train_test_classification(poly_net, criterion, optimizer,
                                   poly_train_loader, poly_test_loader,
                                   num_epochs=100, device=DEVICE)
  # Test it
  X_all = sample_grid().to(device)
  poly_X_all, _ = make_poly_features(poly_degree, X_all)
  y_pred = poly_net(poly_X_all.to(device))

  # Plot it
  plot_decision_map(X_all.cpu(), y_pred.cpu(), X_test.cpu(), y_test.cpu())
  plt.show()

  return num_features

set_seed(seed=SEED)
max_poly_degree = 50
num_features = run_poly_clasification(max_poly_degree, DEVICE, SEED)
print(f'Number of features: {num_features}')
```

### Think! 1.1: Does it generalize well?
Do you think this model is performing well outside its training distribution? Why?

---
# Section 2: Deeper MLPs

```python
# @title Video 2: Case study
from ipywidgets import widgets

out2 = widgets.Output()
with out2:
  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=f"BV1FL411n7SH", width=854, height=480, fs=1)
  print("Video available at https://www.bilibili.com/video/{0}".format(video.id))
  display(video)

out1 = widgets.Output()
with out1:
  from IPython.display import YouTubeVideo
  video = YouTubeVideo(id=f"3g_OJ6dYE8E", width=854, height=480, fs=1, rel=0)
  print("Video available at https://youtube.com/watch?v=" + video.id)
  display(video)

out = widgets.Tab([out1, out2])
out.set_title(0, 'Youtube')
out.set_title(1, 'Bilibili')

# display(out)
```

## Coding Exercise 2: Dataloader on a real-world dataset
Let's build our first real-world dataset loader with Data Preprocessing and Augmentation! And we will use the Torchvision transforms to do it.

We'd like to have a simple data augmentation with the following steps:
* Random rotation with 10 degrees (`.RandomRotation`)
* Random horizontal flipping (`.RandomHorizontalFlip`)

and we'd like a preprocessing that:
* makes Pytorch tensors in the range [0, 1] (`.ToTensor`)
* normalizes the input in the range [-1, 1] (.`Normalize`)

**Hint:** For more info on transform, see the [official documentation](https://pytorch.org/vision/stable/transforms.html).

```python
def get_data_loaders(batch_size, seed):
  ####################################################################
  # Fill in all missing code below (...),
  # then remove or comment the line below to test your function
  raise NotImplementedError("Define the get data loaders function")
  ###################################################################

  # define the transform done only during training
  augmentation_transforms = ...

  # define the transform done in training and testing (after augmentation)
  preprocessing_transforms = ...

  # compose them together
  train_transform = transforms.Compose(augmentation_transforms + preprocessing_transforms)
  test_transform = transforms.Compose(preprocessing_transforms)

  # using pathlib to be compatible with all OS's
  data_path = pathlib.Path('.')/'afhq'

  # define the dataset objects (they can load one by one)
  img_train_dataset = ImageFolder(data_path/'train', transform=train_transform)
  img_test_dataset = ImageFolder(data_path/'val', transform=test_transform)

  g_seed = torch.Generator()
  g_seed.manual_seed(seed)
  # define the dataloader objects (they can load batch by batch)
  img_train_loader = DataLoader(img_train_dataset,
                                batch_size=batch_size,
                                shuffle=True,
                                worker_init_fn=seed_worker,
                                generator=g_seed)
  # num_workers can be set to higher if running on Colab Pro TPUs to speed up,
  # with more than one worker, it will do multithreading to queue batches
  img_test_loader = DataLoader(img_test_dataset,
                               batch_size=batch_size,
                               shuffle=False,
                               num_workers=1,
                               worker_init_fn=seed_worker,
                               generator=g_seed)

  return img_train_loader, img_test_loader

batch_size = 64
set_seed(seed=SEED)
## Uncomment below to test your function
# img_train_loader, img_test_loader = get_data_loaders(batch_size, SEED)
## get some random training images
# dataiter = iter(img_train_loader)
# images, labels = dataiter.next()
## show images
# imshow(make_grid(images, nrow=8))
```

    Random seed 2021 has been set.

```python
# Train it
set_seed(seed=SEED)
net = Net('ReLU()', 3*32*32, [64, 64, 64], 3).to(DEVICE)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), lr=3e-4)
_, _ = train_test_classification(net, criterion, optimizer,
                                img_train_loader, img_test_loader,
                                num_epochs=30, device=DEVICE)
```

```python
# visualize the feature map
fc1_weights = net.mlp[0].weight.view(64, 3, 32, 32).detach().cpu()
fc1_weights /= torch.max(torch.abs(fc1_weights))
imshow(make_grid(fc1_weights, nrow=8))
```

## Think! 2: why first layer features are high level?
Even though it's three layers deep, we see distinct animal faces in the first layer feature map. Do you think this MLP has a hierarchical feature representation? why?

---
# Summary

In the second tutorial of this day, we have dived deeper into MLPs and seen more of their mathematical and practical aspects. More specifically, we have learned about different architectures, i.e., deep, wide, and how they are dependent on the transfer function used. Also, we have learned about the importance of initialization, and we mathematically analyzed two methods for smart initialization.

---
# Bonus: The need for good initialization

In this section, we derive principles for initializing deep networks. We will see that if the weights are too large, then the forward propagation of signals will be chaotic, and the backpropagation of error gradients will explode. On the other hand, if the weights are too small, the forward propagation of signals will be ordered, and the backpropagation of error gradients will vanish. The key idea behind initialization is to choose the weights to be just right, i.e., at the edge between order and chaos. In this section, we derive this edge and show how to compute the correct initial variance of the weights.

Many of the typical initialization schemes in existing deep learning frameworks implicitly employ this principle of initialization at the edge of chaos. So this section can be safely skipped on a first pass and **is a bonus section**.

```python
# @title Video 5: Need for Good Initialization
from ipywidgets import widgets

out2 = widgets.Output()
with out2:
  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=f"BV1Qq4y1H7Px", width=854, height=480, fs=1)
  print("Video available at https://www.bilibili.com/video/{0}".format(video.id))
  display(video)

out1 = widgets.Output()
with out1:
  from IPython.display import YouTubeVideo
  video = YouTubeVideo(id=f"W0V2kwHSuUI", width=854, height=480, fs=1, rel=0)
  print("Video available at https://youtube.com/watch?v=" + video.id)
  display(video)

out = widgets.Tab([out1, out2])
out.set_title(0, 'Youtube')
out.set_title(1, 'Bilibili')

# display(out)
```

## Xavier initialization
Let us look at the scale distribution of an output (e.g., a hidden variable)  $o_i$  for some fully-connected layer without nonlinearities. With  $n_{in}$  inputs  ($x_j$)  and their associated weights  $w_{ij}$  for this layer. Then an output is given by,

\begin{equation}
o_{i} = \sum_{j=1}^{n_\mathrm{in}} w_{ij} x_j
\end{equation}

The weights  $w_{ij}$  are all drawn independently from the same distribution. Furthermore, let us assume that this distribution has zero mean and variance  $\sigma^2$ . Note that this does not mean that the distribution has to be Gaussian, just that the mean and variance need to exist. For now, let us assume that the inputs to the layer  $x_j$ also have zero mean and variance  $\gamma^2$  and that they are independent of $w_{ij}$ and independent of each other. In this case, we can compute the mean and variance of $o_i$ as follows:

\begin{split}
\begin{aligned}
    E[o_i] &= \sum_{j=1}^{n_\mathrm{in}} E[w_{ij} x_j] \\ \\
    &= \sum_{j=1}^{n_\mathrm{in}} E[w_{ij}] E[x_j] = 0, \\ \\ \\
    \mathrm{Var}[o_i] &= E[o_i^2] - (E[o_i])^2 \\ \\
    &= \sum_{j=1}^{n_\mathrm{in}} E[w^2_{ij} x^2_j] - 0 \\ \\
    &= \sum_{j=1}^{n_\mathrm{in}} E[w^2_{ij}] E[x^2_j] \\ \\
    &= n_\mathrm{in} \sigma^2 \gamma^2
\end{aligned}
\end{split}

One way to keep the variance fixed is to set $n_{in}\sigma^2=1$ . Now consider backpropagation. There we face a similar problem, albeit with gradients being propagated from the layers closer to the output. Using the same reasoning as for forward propagation, we see that the gradients’ variance can blow up unless $n_{out}\sigma^2=1$ , where  $n_{out}$ is the number of outputs of this layer. This leaves us in a dilemma: we cannot possibly satisfy both conditions simultaneously. Instead, we simply try to satisfy:

\begin{aligned}
\frac{1}{2} (n_\mathrm{in} + n_\mathrm{out}) \sigma^2 = 1 \text{ or equivalently }
\sigma = \sqrt{\frac{2}{n_\mathrm{in} + n_\mathrm{out}}}
\end{aligned}

This is the reasoning underlying the now-standard and practically beneficial Xavier initialization, named after the first author of its creators [Glorot & Bengio, 2010]. Typically, the Xavier initialization samples weights from a Gaussian distribution with zero mean and variance  $\sigma^2=\frac{2}{(n_{in}+n_{out})}$,

\begin{equation}
w_{ij} \sim \mathcal{N} \left (\mu=0, \sigma=\sqrt{\frac{2}{(n_{in}+n_{out})}} \right)
\end{equation}

We can also adapt Xavier’s intuition to choose the variance when sampling weights from a uniform distribution. Note that the uniform distribution $U(−a,a)$ has variance $\frac{a^2}{3}$. Plugging this into our condition on $\sigma^2$ yields the suggestion to initialize according to

\begin{equation}
w_{ij} \sim  \mathcal{U} \left(-\sqrt{\frac{6}{n_\mathrm{in} + n_\mathrm{out}}}, \sqrt{\frac{6}{n_\mathrm{in} + n_\mathrm{out}}}\right)
\end{equation}

This explanation is mainly taken from [here](https://d2l.ai/chapter_multilayer-perceptrons/numerical-stability-and-init.html).

If you want to see more about initializations and their differences see [here](https://www.deeplearning.ai/ai-notes/initialization/).

## Initialization with transfer function
Let's derive the optimal gain for LeakyReLU following similar steps.

LeakyReLU is described mathematically:

\begin{equation}
f(x)=\left\{
  \begin{array}{ll}
    \alpha \cdot x & \text { for } x<0 \\
    x & \text { for } x \geq 0
  \end{array}\right.
\end{equation}

where $\alpha$ controls the angle of the negative slope.

Considering a single layer with this activation function gives,

\begin{align}
o_{i} &= \sum_{j=1}^{n_\mathrm{in}} w_{ij} x_j\\
z_{i} &= f\left( o_{i} \right)
\end{align}

where $z_i$ denotes the activation of node $i$.

The expectation of the output is still zero, i.e., $\mathbb{E}[f(o_i)=0]$, but the variance changes, and assuming that the probability $P(x < 0) = 0.5$, we have that:

\begin{align}
\mathrm{Var}[f(o_i)] &= \mathbb{E}[f(o_i)^2] - \left( \mathbb{E}[f(o_i)] \right)^{2} \\ \\
&= \frac{\mathrm{Var}[o_i] + \alpha^2 \mathrm{Var}[o_i]}{2} \\ \\
&= \frac{1+\alpha^2}{2}n_\mathrm{in} \sigma^2 \gamma^2
\end{align}

where $\gamma$ is the variance of the distribution of the inputs $x_j$ and $\sigma$ is the variance of the distribution of weights $w_{ij}$, as before.

Therefore, following the rest of derivation as before,

\begin{equation}
\sigma = gain\sqrt{\frac{2}{n_\mathrm{in} + n_\mathrm{out}}}, \, \text{where} \,\, gain = \sqrt{\frac{2}{1+\alpha^2}}
\end{equation}

As we can see from the derived formula of $\sigma$, the transfer function we choose is related with the variance of the distribution of the weights. As the negative slope of the LeakyReLU $\alpha$ becomes larger, the $gain$ becomes smaller and thus, the distribution of the weights is narrower. On the other hand, as $\alpha$ becomes smaller and smaller, the distribution of the weights is wider. Recall that, we initialize our weights, for example, by sampling from a normal distribution with zero mean and variance $\sigma^2$,

## Best gain for Xavier Initialization with Leaky ReLU
You're probably running out of time, so let me explain what's happening here. We derived a theoretical gain for initialization. But the question is whether it holds in practice? Here we have a setup to confirm our finding. We will try a range of gains and see the empirical optimum and whether it matches our theoretical value!

If you have time left, you can change the distribution to sample the initial weights from a uniform distribution. Comment out line 11 and uncomment line 12.

```python
N = 10  # number of trials
gains = np.linspace(1/N, 3.0, N)
test_accs = []
train_accs = []
mode = 'uniform'
for gain in gains:
  print(f'\ngain: {gain}')

  def init_weights(m, mode='normal'):
    if type(m) == nn.Linear:
      torch.nn.init.xavier_normal_(m.weight, gain)
      # torch.nn.init.xavier_uniform_(m.weight, gain)

  negative_slope = 0.1
  actv = f'LeakyReLU({negative_slope})'
  set_seed(seed=SEED)
  net = Net(actv, 3*32*32, [128, 64, 32], 3).to(DEVICE)
  net.apply(init_weights)
  criterion = nn.CrossEntropyLoss()

  optimizer = optim.SGD(net.parameters(), lr=1e-2)
  train_acc, test_acc = train_test_classification(net, criterion, optimizer,
                                                  img_train_loader,
                                                  img_test_loader,
                                                  num_epochs=1,
                                                  verbose=True,
                                                  device=DEVICE)
  test_accs += [test_acc]
  train_accs += [train_acc]

best_gain = gains[np.argmax(train_accs)]
plt.plot(gains, test_accs, label='Test accuracy')
plt.plot(gains, train_accs, label='Train accuracy')
plt.scatter(best_gain, max(train_accs),
            label=f'best gain={best_gain:.1f}',
            c='k', marker ='x')
# calculate and plot the theoretical gain
theoretical_gain = np.sqrt(2.0 / (1 + negative_slope ** 2))
plt.scatter(theoretical_gain, max(train_accs),
            label=f'theoretical gain={theoretical_gain:.2f}',
            c='g', marker ='x')
plt.legend()
plt.plot()
```


```python
%%capture
# Load the quadrotor simulator from a previous notebook
%run quadsim.ipynb
```

```python
from IPython.core.display import HTML
HTML('<style>.impl-inner{background-color:#e4ffe6} .impl-outer{background-color:#e4ebff} .text_cell_render p { text-align: justify; text-justify: inter-word;}</style>')
```

<style>.impl-inner{background-color:#e4ffe6} .impl-outer{background-color:#e4ebff} .text_cell_render p { text-align: justify; text-justify: inter-word;}</style>

# Sliding Mode Control

This notebook discusses and implements a sliding mode controller for the attitude control of a quadrotor. The discussion follows the 2016 work of Shaik and Whidborne from Cranfield, UK [1]. An error representation is used to enable tracking of trajectories in $\phi$, $\theta$, $\psi$, or $r_z$. Note that there is currently no outer loop control of the quadrotor.

Compared to the [SMC implementation based on L'Afflitto](sliding_mode_control.ipynb), this implementation is more introductory and follows the motivating example of [2, Ch 14.1.1].

## Quadrotor Model

The derivation of the standard quadrotor model used herein can be found in the [Quadrotor Model](quadrotor_model.ipynb) notebook. We use the form with inertial translational kinematics and dynamics, and it is restated below, with the state defined as
\begin{align}
x &\triangleq
\begin{bmatrix} r_x & r_y & r_z & \dot{r}_x & \dot{r}_y & \dot{r}_z & \phi & \theta & \psi & p & q & r \end{bmatrix}^\top,
\end{align}

which evolves with the following simplified dynamics (i.e., $\Gamma\approx I_{3\times 3}$ so that we can easily perform control)

\begin{equation}
x =
\begin{bmatrix} \dot{r}_x \\ \dot{r}_y \\ \dot{r}_z \\ \ddot{r}_x \\ \ddot{r}_y \\ \ddot{r}_z \\ \dot{\phi} \\ \dot{\theta} \\ \dot{\psi} \\ \dot{p} \\ \dot{q} \\ \dot{r} \end{bmatrix} =
\begin{bmatrix}
\dot{r}_x \\ \dot{r}_y \\ \dot{r}_z \\
-\frac{u_1}{m}(\cos\phi\sin\theta\cos\psi+\sin\phi\sin\psi) \\
-\frac{u_1}{m}(\cos\phi\sin\theta\sin\psi-\sin\phi\sin\psi) \\
g-\frac{u_1}{m}(\cos\phi\cos\theta) \\
p \\ q \\ r \\
\frac{J_y - J_z}{J_x} qr + \frac{1}{J_x} u_2 \\
\frac{J_z - J_x}{J_y} pr + \frac{1}{J_y} u_3 \\
\frac{J_x - J_y}{J_z} pq + \frac{1}{J_z} u_4
\end{bmatrix}.
\end{equation}

<!--
\begin{equation}
x =
\begin{bmatrix} \dot{r_x} \\ \dot{r_y} \\ \dot{r_z} \\ \ddot{r_x} \\ \ddot{r_y} \\ \ddot{r_z} \\ \dot{\phi} \\ \dot{\theta} \\ \dot{\psi} \\ \dot{p} \\ \dot{q} \\ \dot{r} \end{bmatrix} =
\begin{bmatrix}
x_4 \\ x_5 \\ x_6 \\
-\frac{u_1}{m}(\cos x_7\sin x_8\cos x_9+\sin x_7\sin x_9) \\
-\frac{u_1}{m}(\cos x_7\sin x_8\sin x_9-\sin x_7\sin x_9) \\
g-\frac{u_1}{m}(\cos x_7\cos x_8) \\
x_{10} \\ x_{11} \\ x_{12} \\
\frac{J_y - J_z}{J_x} x_{11} x_{12} + \frac{1}{J_x} u_2 \\
\frac{J_z - J_x}{J_y} x_{10} x_{12} + \frac{1}{J_y} u_3 \\
\frac{J_x - J_y}{J_z} x_{10} x_{11} + \frac{1}{J_z} u_4
\end{bmatrix}.
\end{equation}
-->

The input to this system is defined as
\begin{align}
u &\triangleq
\begin{bmatrix} T & \tau_\phi & \tau_\theta & \tau_\psi \end{bmatrix}^\top.
\end{align}

## Shaik's Implementation

In Shaik [1], four scalar sliding surfaces are defined as
$$
\begin{equation}
s_\alpha = \dot{e}_\alpha + \lambda_\alpha e_\alpha = 0,
\end{equation}
$$
where $\lambda_\alpha$ is a tuning gain and the relevant tracking error is defined as $e_\alpha = \alpha^d - \alpha$. The variable $\alpha$ is used as a placeholder for one of the relevant states: $\phi$, $\theta$, $\psi$, or $r_z$. The time derivative of the sliding manifold is
$$
\begin{equation}
\dot{s}_\alpha = -f -u + \ddot\alpha^d + \lambda_\alpha \dot{e}_\alpha.
\end{equation}
$$

The objective of the sliding mode controller is to choose an input $u$ for each choice of $\alpha$ s.t. system trajectories are driven to the sliding manifold ($s_\alpha=0$) in finite time, which is then asymptotically driven to zero. This is done by choosing the following control inputs, with $k_\alpha$ a scalar tuning gain:
$$
\begin{align}
u_1 &= - \frac{m}{\cos\phi\cos\theta}\left[ -g +\ddot{r}_z^d +\lambda_{r_z}\dot{e}_{r_z} + k_{r_z}\text{sgn}(s_{r_z})\right] \label{eq:shaik-u1} \\
u_2 &= J_x\left[ -\frac{J_y - J_z}{J_x} qr + \ddot\phi^d + \lambda_\phi \dot{e}_\phi + k_\phi \text{sgn}(s_\phi)\right] \\
u_3 &= J_y\left[ -\frac{J_z - J_x}{J_y} pr + \ddot\theta^d + \lambda_\theta \dot{e}_\theta + k_\theta \text{sgn}(s_\theta)\right] \\
u_4 &= J_z\left[ -\frac{J_x - J_y}{J_z} pq + \ddot\psi^d + \lambda_\psi \dot{e}_\psi + k_\psi \text{sgn}(s_\psi)\right]. \label{eq:shaik-u4}
\end{align}
$$

## Discussion of Shaik's Implementation

Note that in [1], a right-handed coordinate frame with $z$ up is used (as opposed to NED, which is used in this series of notebooks). Therefore, the body coordinate frame is a front-left-up (FLU). As a result, the sign of equation \eqref{eq:shaik-u1} is different than $U_4$ in the paper. Also notice that in this notebook we have used $u_1$ for the force equation corresponding to thrust.

The sign issue can be more easily seen if fill in some of the gaps in Shaik's paper. Consider equation (15) of Shaik, which is a common sliding surface definition. Note how the sliding surface is $s_\alpha = \dot{e}_\alpha - \phi(e_\alpha) = 0$, where $\phi(e_\alpha) = -\lambda_\alpha e_\alpha$ is a linear control law, resulting in $s_\alpha = \dot{e}_\alpha + \lambda_\alpha e_\alpha$.

Now that we have identified a sliding surface such that $e_\alpha \to 0$ when the system is steered to the surface, we need to choose an input $u_\alpha$ to drive all system trajectories to the surface $s_\alpha=0$. Choosing the Lyapunov function $V = \frac{1}{2}s_\alpha^2$, we desire that $\dot{V} = s_\alpha\dot{s}_\alpha < 0$. The time-derivative of $s_\alpha$ is
\begin{align}
\dot{s}_\alpha &= \ddot{e}_\alpha + \lambda_\alpha \dot{e}_\alpha \\
&= \ddot{\alpha}^d - \ddot{\alpha} + \lambda_\alpha \dot{e}_\alpha \\
&= \ddot{\alpha}^d - f(\alpha, u_\alpha) + \lambda_\alpha \dot{e}_\alpha \\
&= \ddot{\alpha}^d - f_1(\alpha) - bu_\alpha + \lambda_\alpha \dot{e}_\alpha.  \quad \text{*} \label{eq:sdot}
\end{align}

Note that in step $\text{*}$ we have assumed that the dynamics $f(\alpha, u_\alpha)$ are linear in the control input $u_\alpha$. That is, $\;f(\alpha, u_\alpha) = f_1(\alpha) + bu_\alpha$. Here, we have a choice (see [2, p565]) on whether we design a control $u_\alpha$ that has only a switching component, or we may also add a continuous component that cancels any of the known quantities in \eqref{eq:sdot}. Given that we know the dynamics (albeit, with some uncertainty), we design a control that has the switching component $\beta(e_\alpha) sgn(s_\alpha)$ and a continuous component to cancel out what we know. Therefore, our designed control input is
$$
\begin{equation}
u_\alpha = \frac{1}{b}\left[\underbrace{\ddot{\alpha}^d - f_1(\alpha) + \lambda_\alpha \dot{e}_\alpha}_{continuous} + \underbrace{k_\alpha \text{sgn}(s_\alpha)}_{switching}\right].
\end{equation}
$$

The switching component is a fundamental aspect of sliding mode control. Its purpose is ensure that $\dot{V} = s_\alpha\dot{s}_\alpha < -\eta \left|s_\alpha\right| < 0$ for some value of $\eta$. See Khalil p552 [2] for more information.

```python
def sat(v):
    return v if np.abs(v) <= 1 else np.sign(v)

class ShaikSMC(Controller):
    """Sliding Mode Controller
    """
    def __init__(self):
        self.name = "Shaik SMC"

        # dirty derivative filters
        self.dZdt = DirtyDerivative(order=1, tau=0.05)
        self.dZdt2 = DirtyDerivative(order=2, tau=0.05)
        self.dphdt = DirtyDerivative(order=1, tau=0.05)
        self.dphdt2 = DirtyDerivative(order=2, tau=0.05)
        self.dthdt = DirtyDerivative(order=1, tau=0.05)
        self.dthdt2 = DirtyDerivative(order=2, tau=0.05)
        self.dpsdt = DirtyDerivative(order=1, tau=0.05)
        self.dpsdt2 = DirtyDerivative(order=2, tau=0.05)

        # estimates of the physical properties of the quadrotor
        self.g = 9.8
        self.mass = 4
#         Jxx = 0.060224; Jyy = 0.122198; Jzz = 0.132166
        Jxx = 0.1; Jyy = 0.1; Jzz = 0.1
        self.I = np.array([[Jxx,0,0],
                           [0,Jyy,0],
                           [0,0,Jzz]])

    def inner(self, commanded, state, Ts):
        # Transform the system into the form of HK (14.4) and (14.5)
        eta = np.array([state.flatten()[[2, 6, 7, 8]]]).T # r_z, ph, th, ps
        xi = np.array([state.flatten()[[5, 9, 10, 11]]]).T # \dot{r_z}, p, q, r

        ph = eta.flatten()[1]
        th = eta.flatten()[2]
        ps = eta.flatten()[3]

        Jx = self.I[0,0]
        Jy = self.I[1,1]
        Jz = self.I[2,2]

        #
        # r_z
        #

        # Desired
        des_Z = commanded[0]
        des_Zdot = self.dZdt.update(des_Z, Ts)
        des_Zddot = self.dZdt2.update(des_Zdot, Ts)

        # Gains
        lam_Z = 3
        kZ = 6

        # error
        e_Z = des_Z - eta[0]
        e_Zdot = des_Zdot - xi[0]

        sZ = e_Zdot + lam_Z*e_Z
#         u1 = -self.mass/(np.cos(ph)*np.cos(th))*(-self.g + des_Zddot + lam_Z*e_Zdot + kZ*np.sign(sZ))
        u1 = -self.mass/(np.cos(ph)*np.cos(th))*(-self.g + des_Zddot + lam_Z*e_Zdot + kZ*sat(sZ/0.5))

        # sign difference -- it seems that because of the way the tracking error was defined, it
        # it doesn't look quite the same as in Khalil [2] (Should be all positive -- see other nb).
        u1 = self.mass/(np.cos(ph)*np.cos(th))*(self.g + des_Zddot - lam_Z*e_Zdot) + kZ*sat(-sZ/0.5)

        #
        # \phi
        #

        # Desired
        des_ph = commanded[1]
        des_phdot = self.dphdt.update(des_ph, Ts)
        des_phddot = self.dphdt2.update(des_phdot, Ts)

        # Gains
        lam_ph = 3
        kph = 1

        # error
        e_ph = des_ph - eta[1]
        e_phdot = des_phdot - xi[1]

        sph = e_phdot + lam_ph*e_ph
        u2 = Jx*(-(Jy-Jz)/Jx*xi[2]*xi[3] + des_phddot + lam_ph*e_phdot + kph*np.sign(sph))

        #
        # \theta
        #

        # Desired
        des_th = commanded[2]
        des_thdot = self.dthdt.update(des_th, Ts)
        des_thddot = self.dthdt2.update(des_thdot, Ts)

        # Gains
        lam_th = 3
        kth = 1

        # error
        e_th = des_th - eta[2]
        e_thdot = des_thdot - xi[2]

        sth = e_thdot + lam_th*e_th
        u3 = Jy*(-(Jz-Jx)/Jy*xi[1]*xi[3] + des_thddot + lam_th*e_thdot + kth*np.sign(sth))

        #
        # \psi
        #

        # Desired
        des_ps = commanded[3]
        des_psdot = self.dpsdt.update(des_ps, Ts)
        des_psddot = self.dpsdt2.update(des_psdot, Ts)

        # Gains
        lam_ps = 3
        kps = 1

        # error
        e_ps = des_ps - eta[3]
        e_psdot = des_psdot - xi[3]

        sps = e_psdot + lam_ps*e_ps
        u4 = Jz*((Jx-Jy)/Jz*xi[1]*xi[2] + des_psddot + lam_ps*e_psdot + kps*np.sign(sps))

        u = np.hstack(np.array([u1,u2,u3,u4]))

        return u

    def outer(self):
        pass

    def update(self, commanded, state, pkt, Ts):

        #
        # Outer Loop
        #

        #
        # Inner Loop
        #

        cmd = commanded[[2,6,7,8]]

        u = self.inner(cmd, state, Ts)

        # actuator commands
        return u, commanded
```

```python
# Instantiate a quadrotor model with the given initial conditions
quad = Quadrotor(r=np.array([[0],[0],[10]]),
                 v=np.array([[0],[0],[0]]),
               Phi=np.array([[0.2],[-0.5],[-0.2]]))

# Instantiate the SMC controller
ctrl = ShaikSMC()

# Setup a setpoint commander
cmdr = Commander(default=True)

# Run the simulation
sim = Simulator(quad, ctrl, cmdr=cmdr)
sim.run(20, Ts=0.01)
sim.plot()

print(sim.quad)
```

## References

[1] M. K. Shaik and J. F. Whidborne, “Robust Sliding Mode Control of a Quadrotor,” 11th Int. Conf. Control, pp. 1–5, 2016.

[2] H. K. Khalil, Nonlinear Systems. Princeton, NJ: Prentice Hall, 2002.


# Geophysical Inversion

This application provides an interface to the open-source [SimPEG](https://simpeg.xyz/) inversion algorithms. The application currently supports

 - Electromagnetic (time or frequency) data using a Laterally Constrained 1D approach
 - Gravity and Magnetic Vector Inversion (MVI) (field and/or tensor data) with a tiled octree mesh approach.

New user? Visit the [Getting Started](../installation.rst) page.

**NOTE**

For gravity and magnetic 3D inversion, it is recommended to set the [Working Directory](#Output-panel) to a solid-state drive (SSD) to take full advantage of the [Dask](https://dask.org/) out-of-core parallelization.

## Application
The following sections provide details on the different parameters controling the application. Interactive widgets shown below are for demonstration purposes only.

```python
from geoapps.inversion import InversionApp

app = InversionApp(h5file=r"../../../assets/FlinFlon.geoh5")
app.main
```

    VBox(children=(VBox(children=(Label(value='Workspace', style=DescriptionStyle(description_width='initial')), H…

## Project Selection

Select and connect to an existing **geoh5** project file containing data.

```python
app.project_panel
```

    VBox(children=(Label(value='Workspace', style=DescriptionStyle(description_width='initial')), HBox(children=(F…

See the [Project Panel](base_application.ipynb#Project-Panel) page for more details.

## Object/Data Selection

List of objects available containing geophysical data to be inverted. Currently supports objects of type `Points`, `Curve`, `Surface` and `Grid2D`. Multiple data channels and/or data groups can be selected.

```python
app.data_panel
```

    VBox(children=(Dropdown(description='Object:', index=73, options=('', 'Data_TEM_pseudo3DEPSG_4326', 'Gravity_M…

## Window Selection

Manual selection of an area of interest and data resolution. Positive (solid) and negative (dash) contour lines are traced immediately after changes made to the contour values.

```python
app.window_selection
```

    VBox(children=(VBox(children=(FloatText(value=50.0, description='Grid Resolution (m)', style=DescriptionStyle(…

See the [Map View Selection](view_selection.ipynb#Map-View-Selection) page for more details.

## Data Channel Options

Define geophysical data.

### Survey Type

List of available survey types.

```python
app.system
```

    Dropdown(description='Survey Type: ', options=('MVI', 'Magnetics', 'Gravity', 'AeroTEM (2007)', 'AeroTEM (2010…

The application will attempt to assign the correct system based on
known data field name (e.g. *CPI* => *DIGHEM*).

#### Inducing field parameters

For magnetic data only, set the inducing field parameters:

```python
app.inducing_field

```

    Text(value='60000, 79, 11', description='Inducing Field [Amp, Inc, Dec]', style=DescriptionStyle(description_w…

### Channel Options

Set parameters to invividual [Data](#Object/Data-Selection) channels.

```python
app.data_channel_choices.data_channel_options['tmi']
```

    VBox(children=(Checkbox(value=True, description='Active', style=DescriptionStyle(description_width='initial'))…

#### Active

Checked if the channel is to be used in the inversion

#### Data channel:

Association between the "data channel" and one of the expected "system channel".

#### Error (%, floor)

Estimated data uncertainties associated with the data channel, expressed as a
percent (decimal) and floor value.

As example for magnetic data: 0.1, 1 => $0.1 * |d_{tmi}| + 1\;nT$

#### Offsets (EM systems)

Determine the constant offsets, in meters, between the receiver and the transmitter center location.

Used for the numerical forward calculations.

## Topography, Sensor and Line Options

```python
app.spatial_choices
```

    Dropdown(options=('Topography', 'Sensor'), value='Topography')

### Topography

Set the air/ground interface of the inversion model.

```python
app.topography.main
```

    VBox(children=(RadioButtons(description='Define by:', options=('Object', 'Relative to Sensor', 'Constant'), va…

#### Object

Set the topography based on an object stored in the target `geoh5` project. The Z value of the cells/vertices can be assigned based on a chosen `Data` field.

#### Relative to Sensor

Topography is defined by a fixed vertical offset from a selected object position (vertices or centroids), also referred to as a "draped height".

#### Constant

Topography is defined as a flat surface with constant elevation.

### Sensor

Defines the sensors position in 3D space.

```python
app.sensor.main
```

    VBox(children=(RadioButtons(description='Define by:', index=1, options=('sensor location + (dx, dy, dz)', 'top…

#### Sensor Location

Sensor position defined by a constant offset from the vertices of the selected
object. Typically used for towed system where the GPS receiver is on the
aircraft.

#### Topo and radar

Receiver locations defined by the vertices horizontal positions [x, y] and z value
interpolated from topography + clearance height. Typically used for gridded
data with constant draped height or for airborne survey with inaccurate GPS
elevation (radar height).

### Line ID

**(EM only)**: Line channel and line number to be inverted.

```python
app.lines.main
```

    VBox(children=(Dropdown(description='Lines field', options=('', '--- Channels ---', 'Airborne_Gxx', 'Airborne_…

## Inversion Options

List of parameters controlling the inversion.

```python
app.inversion_parameters.main
```

    VBox(children=(HBox(children=(Label(value='Inversion Options'),)), HBox(children=(Dropdown(options=('output na…

### Output Name

Name given to the inversion group added to the ANALYST project.

### Target Misfit

Target data misfit where 1 = number of data

```python
app.inversion_parameters.chi_factor
```

    FloatText(value=1.0, description='Target misfit')

### Uncertainty mode

(EM Only) Global change to the data uncertainties

```python
app.inversion_parameters.uncert_mode
```

    RadioButtons(index=1, options=('Estimated (%|data| + background)', 'User input (%|data| + floor)'), value='Use…

**Estimated** Floor uncertainties re-assgined based on the best-fitting halfspace model. Useful when dealing with fast time domain decays due to IP effects.

**User Input** (Default) Keep the same uncertainties as defined in the [Channel Options](#Channel-Options)

### Starting Model

Initial model used to begin the inversion.

```python
app.inversion_parameters.starting_model.main
```

    VBox(children=(Label(value='Starting effective susceptibility'), VBox(children=(RadioButtons(index=1, options=…

#### Model

Model object and values selected from any `Points`, `Curve`, `Surface`, `BlockModel` or `Octree` object.
Values are interpolated onto the inversion mesh using a nearest neighbor approach.

#### Constant

Constant background half-space value.

### Susceptibility model

**(FEM Only)**

Susceptibility values used in the forward calculations only.

```python
app.inversion_parameters.susceptibility_model.main
```

    VBox(children=(Label(value='Background susceptibility', style=DescriptionStyle(description_width='initial')), …

#### None

No susceptibility model used.

#### Model

Susceptibitlity model selected from any `Points`, `Curve`, `Surface`, `BlockModel` or `Octree` object.
Values are interpolated onto the inversion mesh using a nearest neighbor approach.

#### Constant

Constant susceptibility half-space value.

### Regularization

See the [Regularization Panel](#Regularization-Panel) section.

### Upper-Lower Bounds

Upper and lower bound constraints applied on physical property model.

```python
app.inversion_parameters.bound_panel
```

    VBox(children=(Text(value='', description='Upper bound value'), Text(value='', description='Lower bound value'…

If left empty, the inversion uses [$-\infty$, $\infty$] physical property bounds.

### Mesh

See the [Mesh Design](#Mesh-Design) section.

### Ignore Values

#### Data

Flag value ignored by the inversion by assigning $\infty$ uncertainties on the data points.

```python
app.inversion_parameters.ignore_values
```

    Text(value='-99999', description='Data (i.e. <0 = no negatives)')

#### Air cells

Value assigned to the output model cells above topography. Defaults to $10^{-8}$ S/m for EM inversions

```python
app.inversion_parameters.air_values
```

    FloatText(value=0.0, description='Air cells fill value')

### Optimization

Parameters controlling various aspects of the projected Gauss-Newton inversion algorithm.

```python
app.inversion_parameters.optimization
```

    VBox(children=(IntText(value=25, description='Max beta Iterations', style=DescriptionStyle(description_width='…

#### Max beta Iterations

Maximum number of $\beta$-iterations allowed.
Note that when applying sparse norms, the inversion may require $>20$ iterations to converge.

#### Target misfit

Target data misfit where $\chi=1$ corresponds to $\phi_d=N$ (number of data). (See [Data Misfit and Uncertainties](https://giftoolscookbook.readthedocs.io/en/latest/content/fundamentals/Uncertainties.html#data-misfit-and-uncertainties))

#### Starting trade-off ($\beta$)

**ratio**: Factor multiplying the initial $\beta$ defined by the ratio between the initial misfit and regularization:

\begin{equation}
\beta_0 = \gamma * \frac{\phi_d}{ \phi_m}
\end{equation}

**value**: Fixed $\beta$ value specified by the user.

#### Max CG Iterations

Maximum number of Conjugate Gradient (CG) iterations per Gauss-Newton solve.

#### CG Tolerance

Threshold on the minimum Conjugate Gradient (CG) step to end the Gauss-Newton solve.

#### Max CPUs

Maximum number of threads used for the parallelization. Defaults to half the system thread count.

#### Max RAM (Gb)

Approximate memory (RAM) used during the inversion. The value is used to calculate the chunk size for the storage of the sensitivity matrix on solid-state drive.
(See [dask.array.to_zarr](https://docs.dask.org/en/latest/array-api.html#dask.array.to_zarr))

## Regularization Panel

Parameters controlling the regularization function.

### Reference model

Reference model values used to contrain the inversion.

```python
app.inversion_parameters.reference_model.main
```

    VBox(children=(Label(value='Reference effective susceptibility'), VBox(children=(RadioButtons(index=2, options…

#### None

No reference model used. This is equivalent as setting the [Alphas](#Alpha-Parameters) `s` parameter to zero.

#### Model

Reference model selected from any `Points`, `Curve`, `Surface`, `BlockModel` or `Octree` object.
Values are interpolated onto the inversion mesh using a nearest neighbor approach.

#### Constant

Constant reference half-space value.

#### Best-fitting halfspace

**(EM only)** Run a preliminary inversion to determine a best-fitting halfspace
(single conductivity value) at each data station.

### Alpha Parameters

Scaling between the components of the regularization function.

```python
app.inversion_parameters.alphas

```

    Text(value='1.0, 1.0, 1.0, 1.0', description='Scaling alpha_(s, x, y, z)')

See the [SimPEG.API](https://docs.simpeg.xyz/content/api_core/api_Regularization.html#SimPEG.regularization.Simple) for technical details.

### $l_p$-norms

Approximated norms applied to the components of the regularization.

```python
app.inversion_parameters.norms

```

    Text(value='0, 2, 2, 2', continuous_update=False, description='Norms p_(s, x, y, z)', style=DescriptionStyle(d…

See notes on [Sparse and Blocky Norms](https://giftoolscookbook.readthedocs.io/en/latest/content/fundamentals/Norms.html#sparse-and-blocky-norms) for technical details.

## Mesh Design

Definition of the grid containing the physical property model.

### Octree Mesh

**(Gravity and magnetics)**

Parameters used for the construction of a global `Octree` mesh.

```python
app.inversion_parameters.mesh.main
```

    VBox(children=(Label(value='Octree Mesh'), Text(value='25, 25, 25', description='Smallest cells', style=Descri…

#### Smallest cells

Dimensions (x,y,z) of the smallest octree cells.

####  Layers below topography

Number of layers of cells at each octree level below the topography surface.
As topography generally extends further, it is common practice to use only
coarse cells once outside of the survey area.

####  Layers below data

Number of layers of cells at each octree level below the observation points.

####  Minimum depth

Minimum depth (m) of the mesh, rounded up to the next power of 2.

####  Padding [W,E,N,S,D,U]

Additional padding distance (m) along West, East, North, South, Down and Up,
chosen by default to be half of the data extent in each direction. Additional
padding can be added at depth, but it is recommended to set it through the
[Minimum depth](#Minimum-depth) parameter.

#### Max triangulation length

Parameter controlling the Delaunay triangulation used for the refinement of
the mesh around topography and observation points. Large triangles generated
between points far apart are filtered out, reducing the number of small cells
in regions without data.

### 1D Mesh

**(EM 1D only)**

For the laterally constrained 1D inversion (time or frequency), the earth is
discretized into layer models connected by the regularization function. The
following parameters define the 1D mesh repeated below each transmitter/receiver pairs

```python
app.mesh_1D.main
```

    VBox(children=(Label(value='1D Mesh'), FloatText(value=10.0, description='Smallest cell (m):', style=Descripti…

#### Smallest cell (m)

Smallest layer at the top of the 1D mesh.

#### Expansion factor

Rate of cell expansion from the smallest cell size.

#### Number of cells

Total number of layers below each station.

#### Max depth

Returns the full depth of the 1D mesh given the parameters selected above.

## Output panel

Specify the working directory where the inversion occurs. For gravity and magnetic data inversion, it is **highly recommended** to chose a location on a **solid-state drive (SSD)** to fully take advantage of the `Dask` parallelization.

```python
app.output_panel
```

    VBox(children=(FileChooser(path='C:\Users\dominiquef\Documents\GIT\mira\mirageoscience-apps\docs\content\appli…

**Write inpute**: Click to write the input parameters to a `*.json` file

**Run SimPEG**: Launch the inversion routine. Results will be written directly to the target `geoh5` project.

### Command line execution

Alternatively, the inversion can be launched on a different computer from command line:

`activate geoapps`

`python -m geoapps.pf_inversion [Name].json`

This assumes that the geoapps have been installed and that the reference `*.json` and `*.geoh5` project are accessible.

```python
import matplotlib.pyplot as plt
app.figure.savefig("images/inversion_thumbnail.png", dpi=300, bbox_inches="tight")
```


# IAF neurons singularity

This notebook describes how NEST handles the singularities appearing in the ODE's of integrate-and-fire model neurons with alpha- or exponentially-shaped current, when the membrane and the synaptic time-constants are identical.

```python
import sympy as sp
sp.init_printing(use_latex=True)
from sympy.matrices import zeros
tau_m, tau_s, C, h = sp.symbols('tau_m, tau_s, C, h')
```

For alpha-shaped currents we have:

```python
A = sp.Matrix([[-1/tau_s,0,0],[1,-1/tau_s,0],[0,1/C,-1/tau_m]])
```

## Non-singular case ($\tau_m\neq \tau_s$)
The propagator is:

```python
PA = sp.simplify(sp.exp(A*h))
PA
```

Note that the entry in the third line and the second column $A_{32}$ would also appear in the propagator matrix in case of an exponentially shaped current

## Singular case ($\tau_m = \tau_s$)
We have

```python
As = sp.Matrix([[-1/tau_m,0,0],[1,-1/tau_m,0],[0,1/C,-1/tau_m]])
As
```

The propagator is

```python
PAs = sp.simplify(sp.exp(As*h))
PAs
```

## Numeric stability of propagator elements
For the limes $ \tau_s\rightarrow\tau_m $ the entry $PA_{32}$ becomes numerically unstable, since denominator and enumerator go to zero.

$1.$ We show that $PAs_{32}$ is the limit of $PA_{32}(\tau_s)$ for $\tau_s\rightarrow\tau_m$.:

```python
PA_32 = PA.row(2).col(1)[0]
sp.limit(PA_32, tau_s, tau_m)
```

$2.$ The Taylor-series up to the second order of the function $PA_{32}(\tau_s)$ is:

```python
PA_32_series = PA_32.series(x=tau_s,x0=tau_m,n=2)
PA_32_series
```

Therefore we have

$T(PA_{32}(\tau_s,\tau_m))=PAs_{32}+PA_{32}^{lin}+O(2)$ where $PA_{32}^{lin}=h^2(-\tau_m + \tau_s)*exp(-h/\tau_m)/(2C\tau_m^2)$

$3.$ We define

$dev:=|PA_{32}-PAs_{32}|$

We also define $PA_{32}^{real}$ which is the correct value of P32 without misscalculation (instability).

In the following we assume $0<|\tau_s-\tau_m|<0.1$. We consider two different cases

a) When $dev \geq 2|PA_{32}^{lin}|$ we do not trust the numeric evaluation of $PA_{32}$, since it strongly deviates from the first order correction. In this case the error we make is

$|PAs_{32}-PA_{32}^{real}|\approx |P_{32}^{lin}|$

b) When $dev \le |2PA_{32}^{lin}|$ we trust the numeric evaluation of $PA_{32}$. In this case the maximal error occurs when $dev\approx 2 PA_{32}^{lin}$ due to numeric instabilities. The order of the error is again

$|PAs_{32}-PA_{32}^{real}|\approx |P_{32}^{lin}|$

The entry $A_{31}$ is numerically unstable, too and we treat it analogously.

## Tests and examples
We will now show that the stability criterion explained above leads to a reasonable behavior for $\tau_s\rightarrow\tau_m$

```python
import nest
import numpy as np
import pylab as pl
```

Neuron, simulation and plotting parameters

```python
taum = 10.
C_m = 250.
# array of distances between tau_m and tau_ex
epsilon_array = np.hstack(([0.],10.**(np.arange(-6.,1.,1.))))[::-1]
dt = 0.1
fig = pl.figure(1)
NUM_COLORS = len(epsilon_array)
cmap = pl.get_cmap('gist_ncar')
maxVs = []
```

Loop through epsilon array

```python
for i,epsilon in enumerate(epsilon_array):
    nest.ResetKernel() # reset simulation kernel
    nest.SetKernelStatus({'resolution':dt})

    # Current based alpha neuron
    neuron = nest.Create('iaf_psc_alpha')
    nest.SetStatus(neuron,{'C_m':C_m,'tau_m':taum,'t_ref':0.,'V_reset':-70.,'V_th':1e32,
                           'tau_syn_ex':taum+epsilon,'tau_syn_in':taum+epsilon,'I_e':0.})

    # create a spike generator
    spikegenerator_ex=nest.Create('spike_generator')
    nest.SetStatus(spikegenerator_ex,{'spike_times': [50.]})

    # create a voltmeter
    vm = nest.Create('voltmeter',params={'interval':dt})

    ## connect spike generator and voltmeter to the neuron

    nest.Connect(spikegenerator_ex, neuron,'all_to_all',{'weight':100.})
    nest.Connect(vm, neuron)

    # run simulation for 200ms
    nest.Simulate(200.)

    # read out recording time and voltage from voltmeter
    times=nest.GetStatus(vm)[0]['events']['times']
    voltage=nest.GetStatus(vm)[0]['events']['V_m']

    # store maximum value of voltage trace in array
    maxVs.append(np.max(voltage))

    # plot voltage trace
    if epsilon == 0.:
        pl.plot(times,voltage,'--',color='black',label='singular')
    else:
        pl.plot(times,voltage,color = cmap(1.*i/NUM_COLORS),label=str(epsilon))

pl.legend()
pl.xlabel('time t (ms)')
pl.ylabel('voltage V (mV)')
```

    <matplotlib.text.Text at 0x598b510>

Show maximum values of voltage traces

```python
fig = pl.figure(2)
pl.semilogx(epsilon_array,maxVs,color='red',label='maxV')
#show singular solution as horizontal line
pl.semilogx(epsilon_array,np.ones(len(epsilon_array))*maxVs[-1],color='black',label='singular')
pl.xlabel('epsilon')
pl.ylabel('max(voltage V) (mV)')
pl.legend()
```

    <matplotlib.legend.Legend at 0x6799350>

```python
pl.show()
```

The maximum of the voltage traces show that the non-singular case nicely converges to the singular one and no numeric instabilities occur.


Housekeeping...

```python
# Classic libraries
# from __future__ import division, print_function
import numpy as np
import sympy as sym
import pandas as pd
# Plotting libraries
import matplotlib.pyplot as plt
import seaborn as sn # library for pretty plots. See https://stanford.edu/~mwaskom/software/seaborn/
sn.set_style('whitegrid')
cp = sn.xkcd_rgb["pinkish"]
cb = "#3498db"
cr = "#1fa774"
# Specialist functions
from scipy.optimize import minimize, root, fsolve
from scipy.optimize import leastsq, least_squares
from scipy.interpolate import interp1d
from scipy.special import erf
# Core functions - solve the problem and plot the results
import capital_new as capital
# Other extensions
%matplotlib inline
%load_ext autoreload
%autoreload 2
```

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

Loading data from Adamopoulos and Restuccia (2014)

```python
CalUS = np.loadtxt('model_data.csv', delimiter=',').T
CalP = np.loadtxt('model_data_p.csv', delimiter=',').T
y_grid, y_prob, farm_grid, k_grid, output_m = CalUS
y_grid_p, y_prob_p, farm_grid_p, k_grid_p, output_m_p = CalP
```

# Replicating Adamopoulos and Restuccia (2014) and some more using Eeckhout and Kircher (2018)

This notebook shows how introducing heterogeneity in the quality of land changes the results or Adamopoulos and Restuccia (2014) (AR thereafter). This is based on the model of Assorative Martching with large firms of Eeckhout and Kircher (2018) (EK thereafter).

## Index

1. [Introduction](#sec1)
2. [Calculating $k^*$](#sec2)
3. [Solving original model](#sec3)
4. [Uneven spread of $x$](#sec4) \\ 4.1 [Generating the distributions](#sec41) \\
  4.2 [Solving the model](#sec42) \\
  4.3 [Results](#sec43)

5. [The PAM condition](#sec5)

## 1. Introduction

The original AR paper's production function is:

$$ p_aA\kappa(\eta(k)^{\rho} + (1 -\eta)(yl))^{\rho})^{\frac{\gamma}{\rho}} $$

Where:
- $y \to$ farmer skill
- $l \to$ farm size
- $k \to$ capital

The rest are parameters that are calibrated to match the farm size distribution of the US and the capital to land ratio.

EK introduce two more variables:
- $x \to$ land quality
- $r \to$ farmer *supervising intentisty*, normalized to 1.

Adding heterogenity in land quality ($x$) to the AR model requires a different solution method.

The optimal choice of $k$ can be calculated as a function of the other variables so that the problem can be solved as an initial value problem (IVP) in the standard EK way. That is, finding $k^*$ such that
$$k^*(x,y,l,r) = \mathrm{argmax}_{k}F(x,y,l,r,k)-rRk$$

There is a number of ways of introducing these $x$ and $r$, of which 4 were considered. Below is the list of all tehm, with their accompaning first order condiction (FOC) for solving for $k^*(x,y,l,r)$.

| # | name | function | FOC    |
|---| ---- |:--------:|:--------:|
| 0 | $\color{red}{x}$ with $l/r$ | $ rp_aA\kappa(\eta(k)^{1/4} + (1 -\eta)(y(l/r)^{\color{red}{x}})^{1/4})^{2} - rRk$ | $ K^{3} + -\frac{Ap_a\kappa\eta^{2}}{2R} K + \frac{Ap_a\kappa\eta(1 -\eta)(y(\frac{l}{r})^{x})^{1/4}}{2R}   = 0$ |
| 1 | $\color{red}{x}$ with $A$ | $ rp_a\color{red}{x}A\kappa [\eta (k)^{1/4} + (1-\eta)(y(l/r))^{1/4} ]^2 - rRk $ | $K^3 - \frac{xAp_a\kappa\eta^2}{2R}K - \frac{xAp_a\kappa\eta(1-\eta)(y\frac{l}{r})^{1/4}}{2R}= 0 $ |
| 2 | $\color{red}{x}$ with $k$ | $ rp_aA\kappa [\eta (k\color{red}{x})^{1/4} + (1-\eta)(y(l/r))^{1/4} ]^2 - rRk $ | $K^3 - \frac{Ap_a\kappa\eta^2 x^{1/2}}{2R}K - \frac{Ap_a\kappa\eta(1-\eta)(xy\frac{l}{r})^{1/4}}{2R}= 0 $ |
| 3 | $\color{red}{x}$ with $k$ (out of the bracket) | $ rp_aA\kappa [\eta \color{red}{x} k^{1/4} + (1-\eta)(y(l/r))^{1/4} ]^2 - rRk $ | $K^3 - \frac{Ap_a\kappa\eta^2 x^2}{2R}K - \frac{Ap_a\kappa\eta(1-\eta)(y\frac{l}{r})^{1/4} x}{2R}= 0 $
| 4 | $k^{\color{red}{x}}$ | $ rp_aA\kappa(\eta(k)^{\color{red}{x}/4} + (1 -\eta)(y(l/r))^{1/4})^{2} - rRk$   | None |
*Note: $K = k^{1/4}$*

For the case of $k^{\color{red}{x}}$, there is no closed-form solution and $k^*$ has to be solved numerically. Solving the model with this especification renders the shooting solver unstable, therefore it is not used. Nevertheless, `new_capital.py` includes a function to approximate this point by point.

Inside the code the problem is reformulated, using an equivalent production function

$$lp_aA\kappa \left( \eta(k)^{\rho} + (1 -\eta)(x(r/l)^{y})^{\rho} \right) ^{\frac{\gamma}{\rho}} - lRk$$.

The original formulation the shooting solver is unable to match the lowest values of $y$ (0.000003). This transformation does not alter the results, but requires some care when dealing with the results: farm size ($\theta = \frac{l}{r}$) coming from the solver would be the inverse of model $\theta$, and in the code below you'll find `xs` refering to farmer skill ($y$ in the model). Several times in this notebook there the change in variables will be made explicit.

For the rest of the functional forms, $k^*$ is found by solving a cubic root problem. This is done in section 2. `new_capital.py` approximates the solution with cubic or quadratic "tensor" polynomials - that is, polynomials with cross-products. This solves a multiple root problem and gives an easily differentiable expression. This approximation allows to use the shooting solver wihtout major complications - at the price of some numerical error.

Once $k^*$ has been found, we proceed to solve for the baseline case where $x=1$ (section 3) and then increase the spread up to $x\in[0.75,1.25]$ (section 4). In this application we use a shooting algorithm with logaritmic steps.

Worker skill $x$ is set to follow a lognormal distribution with mean 1. This distribution makes the results move *closest* to the data - that is, small farms increasing size, but not much. Section 4 provides the results for different distributions and spreads.

Finally, section 5 provides the symolic derivatives to the production function and checks numerically that PAM holds for any combination of $x,y,l,r$ and $k$ - given the parameters of AR. **Upadate**: The PAM condition has been corrected and it's checked for at each step numerically and graphically. The new condition is described in the Corrigendum to EK.

## 2 Calculating $k^*$

First step is to declare the parameters of the parameters of the rich and poor countries. Then these are fed to the `get_k` function - this calculates the optimal $k^*$ function for a given spread of x. The function also prints the sum or squared errors of the approximation.

The default option for the functional form is function (0) above: $ p_aA\kappa(\eta(k)^{\rho} + (1 -\eta)(y(l/r)^{x})^{\rho})^{\frac{\gamma}{\rho}} - Rk$.

Rich country (US calibration):

```python
pam_dic_R = {'A': 1.0, 'eta': 0.890204456766942, 'kappa': 0.3159, 'R':0.13099, 'rho': 0.25, 'gamma': 0.5}
pam_dic_P = {'A': 0.3987, 'eta': 0.890204456766942, 'kappa': 0.5209, 'R':0.3958, 'rho': 0.25, 'gamma': 0.5 }

k_star_R = capital.get_k(pam_dic_R, x_range=[0.9,1.1],print_residuals=True)
k_star_P = capital.get_k(pam_dic_P, x_range=[0.9,1.1],print_residuals=True)
```

    0.0006049062442879255
    0.00322205421125489

This cell defines the symbolical expression from PAM as defined in the Corrigendum:

```python
x, y, k = sym.var('x, y, k')
R, l, r, A, kappa, p_a, rho, gamma, eta = sym.var('R, l, r, A, kappa, p_a, rho, gamma, eta')

sym.init_printing()

F0 = r*A*kappa*(eta*k**rho + (1- eta)*(((l/r)**x)*y)**rho)**(gamma/rho)
F1 = r*A*kappa*x*(eta*k**rho + (1- eta)*((l/r)*y)**rho)**(gamma/rho)
F2 = r*A*kappa*(eta*(x*k)**rho + (1- eta)*((l/r)*y)**rho)**(gamma/rho)
F3 = r*A*kappa*(x*eta*k**rho + (1- eta)*((l/r)*y)**rho)**(gamma/rho)
F4 = r*A*kappa*(eta*k**(rho*x) + (1- eta)*((l/r)*y)**rho)**(gamma/rho)

lhs = F0.diff(x,y)*F0.diff(l,r)*F0.diff(k,k) - F0.diff(x,y)*F0.diff(l,k)*F0.diff(r,k) - F0.diff(x,k)*F0.diff(y,k)*F0.diff(l,r)
rhs = F0.diff(x,r)*F0.diff(y,l)*F0.diff(k,k) - F0.diff(x,r)*F0.diff(y,k)*F0.diff(l,k) - F0.diff(x,k)*F0.diff(y,l)*F0.diff(r,k)

PAM_R = sym.lambdify((x,y,l,r,k),sym.simplify(lhs-rhs).subs(pam_dic_R))
PAM_P = sym.lambdify((x,y,l,r,k),sym.simplify(lhs-rhs).subs(pam_dic_P))
```

## 3 Solving original model ($x=1$)

The following solves for the parameters of the rich country with no spread of $x$ - or $x \in (0.99999,1.00001)$.

```python
sol_R1 = capital.solve("R", k_star_R, ftype=0,assort='positive',
                       verbose=False,guess_overide=0.8)
```

    Success! All workers and firms are matched
    Warning: Approximated solution failed to satisfy required assortativity condition.

The warning about the PAM condition refers to the fact that for some points the PAM condition as defined in EK are not fulfilled:

```python
ders = capital.derivatives(sol_R1, "R")
capital.plot_PAM(ders, sol_R1['xs'])
```

However this is steeming from both (i) the approximation error of $k^*$ and (ii) the inversion of the problem.

If we check for the true PAM condition as defined from the Corrigendum the solution does satisfy this:

```python
plt.plot(sol_R1['xs'],PAM_R(sol_R1['xs'],sol_R1['ys'],sol_R1['thetas'],1.0,sol_R1['ks']))
plt.title('True PAM condition check - should be negative at all points')
plt.xscale("log")
plt.ylabel('PAM condition with $k$')
plt.xlabel('$x$')
plt.show()
```

```python
print("Points for which the PAM condition doesn't hold:",
      np.sum(PAM_R(sol_R1['ys'],sol_R1['xs'],sol_R1['thetas'],1.0,sol_R1['ks'])>=0))
```

    Points for which the PAM condition doesn't hold: 0

Solution plot:

```python
plt.figure(figsize=(9,3))
plt.subplot(131)
plt.title('land quality $x$')
plt.plot(sol_R1['xs'],sol_R1['ys'])
plt.ylim(0.9,1.1)
plt.subplot(132)
plt.title('farm size $l$')
plt.plot(sol_R1['xs'],1/sol_R1['thetas'])
plt.xlabel('farmer skill $y$')
plt.subplot(133)
plt.title('capital $k$')
plt.plot(sol_R1['xs'],sol_R1['ks'])
plt.show()
```

### Poor country parameters

Same procedure: solving first, checking PAM and plotting results.

```python
sol_P1 = capital.solve("P", k_star_P, ftype=0,assort='positive',
                       verbose=False,guess_overide=2.0) #guess_overide=2000.0, scaling_x=0.90607341901223637,
```

    Success! All workers and firms are matched
    Warning: Approximated solution failed to satisfy required assortativity condition.

```python
plt.plot(sol_P1['xs'],PAM_P(sol_P1['ys'],sol_P1['xs'],1/sol_P1['thetas'],1.0,sol_P1['ks']))
plt.title('True PAM condition check - should be negative at all points')
plt.xscale("log")
plt.ylabel('PAM condition with $k$')
plt.xlabel('$x$')
plt.show()
```

```python
print("Points for which the PAM condition doesn't hold:",
      np.sum(PAM_P(sol_P1['ys'],sol_P1['xs'],1/sol_P1['thetas'],1.0,sol_P1['ks'])>0))
```

    Points for which the PAM condition doesn't hold: 0

Histogram comparing with the results with AR (Figure 7 in the appendix of AR):

```python
capital.plot_histogram(sol_R1,sol_P1,CalUS,CalP)
```

## 4. Uneven spread of x

Here the spread of land quality $x$ increases. Although in the end only the lognormal distribution is used, illustrations with other distributions are shown as well to illustrate graphically how the distribution of $x$ could change the results.

### 4.1 Generating the distributions

Lognormal distribution:

```python
_modules = [{'ImmutableMatrix': np.array, 'erf': erf, 'sqrt': np.sqrt}, 'numpy']

x, loc2, mu2, sigma2 = sym.var('x, loc2, mu2, sigma2')
productivity_cdf = (0.5 + 0.5 * sym.erf((sym.log(x - loc2) - mu2) / sym.sqrt(2 * sigma2**2)))

productivity_params_lognorm = {'loc2': 0.0, 'mu2': 0.00, 'sigma2': 0.20}

pdf_lognorm = productivity_cdf.diff(x).subs(productivity_params_lognorm)
pdf_exe_lognorm = sym.lambdify(x, pdf_lognorm, 'numpy')
cdf_exe_lognorm = sym.lambdify(x, productivity_cdf.subs(productivity_params_lognorm), _modules)
```

Bimodal

```python
sig1, sig2, x, mu1, mu2= sym.var("sig1, sig2, x, mu1, mu2")
c1 = 0.5 * (1 + sym.erf((x-mu1)/ sym.sqrt(2 * sig1**2)))
c2 = 0.5 * (1 + sym.erf((x-mu2)/ sym.sqrt(2 * sig2**2)))
cdf = c1*0.5+c2*0.5

params_bimod = {'mu1': 0.93333, 'mu2': 1.06667, 'sig1': 0.033,'sig2': 0.033}

pdf_bimod = cdf.diff(x).subs(params_bimod)
cdf_bimod = cdf.subs(params_bimod)

pdf_exe_bimod = sym.lambdify(x, pdf_bimod,  _modules)
cdf_exe_bimod = sym.lambdify(x, cdf_bimod,  _modules)
```

Normal

```python
params_norm = {'mu1': 1.0, 'sig1': 0.06666 }

pdf_norm = c1.diff(x).subs(params_norm)
cdf_norm = c1.subs(params_norm)

pdf_exe_norm = sym.lambdify(x, pdf_norm,  _modules)
cdf_exe_norm = sym.lambdify(x, cdf_norm,  _modules)
```

Normalising constants - so total ammount of land integrates to 1.

```python
norm_bimod = 1/(cdf_exe_bimod(1.1)-cdf_exe_bimod(0.9))
norm_norm = 1/(cdf_exe_norm(1.1)-cdf_exe_norm(0.9))
norm_lognorm = 1/(cdf_exe_lognorm(1.1)-cdf_exe_lognorm(0.9))

c_bimod = cdf_exe_bimod(1.1)*norm_bimod-1
c_norm= cdf_exe_norm(1.1)*norm_norm-1
c_lognorm= cdf_exe_lognorm(1.1)*norm_lognorm-1
```

Normalising constants for different spreads, lognormal distribution

```python
norm_lognorm_05 = 1/(cdf_exe_lognorm(1.05)-cdf_exe_lognorm(0.95))
c_lognorm_05 = cdf_exe_lognorm(1.05)*norm_lognorm_05-1

norm_lognorm_10 = 1/(cdf_exe_lognorm(1.1)-cdf_exe_lognorm(0.9))
c_lognorm_10 = cdf_exe_lognorm(1.1)*norm_lognorm_10-1

norm_lognorm_15 = 1/(cdf_exe_lognorm(1.15)-cdf_exe_lognorm(0.85))
c_lognorm_15 = cdf_exe_lognorm(1.15)*norm_lognorm_15-1

norm_lognorm_20 = 1/(cdf_exe_lognorm(1.2)-cdf_exe_lognorm(0.8))
c_lognorm_20 = cdf_exe_lognorm(1.2)*norm_lognorm_20-1

norm_lognorm_25 = 1/(cdf_exe_lognorm(1.25)-cdf_exe_lognorm(0.75))
c_lognorm_25 = cdf_exe_lognorm(1.25)*norm_lognorm_25-1

norm_lognorm_30 = 1/(cdf_exe_lognorm(1.3)-cdf_exe_lognorm(0.7))
c_lognorm_30 = cdf_exe_lognorm(1.3)*norm_lognorm_30-1
```

Plots of the different distributions

```python
xss = np.linspace(0.9,1.1,6000)
plt.figure(figsize=(6,6))
sn.set_palette('colorblind', 5)
plt.axvline(1.0,label="x=1", c= 'black')
plt.plot(xss,5*np.ones(len(xss)), label="Uniform")
plt.plot(xss,pdf_exe_bimod(xss)*norm_bimod, label='Bimodal')
plt.plot(xss,pdf_exe_lognorm(xss)*norm_lognorm, label="Lognormal")
plt.xlim(0.9,1.1)
plt.ylim(0,8)
plt.title("Pdf of land quality", fontsize=16)
plt.plot(xss,pdf_exe_norm(xss)*norm_norm, label='Normal')

plt.tight_layout()
plt.legend(loc='best')
plt.show()
```

### 4.2 Solving the model

The distribution that used in the paper is the lognormal. This distribution delivers the results closer to the data provided in AR.

```python
# Execute this cell to store results
sol_log = {}
sol_logP = {}
```

Calculating $k^*(x,y,l,r)$ for different $x$ spreads:

```python
k_star_R25 = capital.get_k(pam_dic_R, x_range=[0.75,1.25],print_residuals=True)
k_star_P25 = capital.get_k(pam_dic_P, x_range=[0.75,1.25],print_residuals=True)

k_star_R30 = capital.get_k(pam_dic_R, x_range=[0.7,1.3],print_residuals=True)
k_star_P30 = capital.get_k(pam_dic_P, x_range=[0.7,1.3],print_residuals=True)
```

    0.0007566085962773021
    0.004719416672145595
    0.0008611003591574059
    0.00560531009701687

Lognormal

```python
sol_log[2605] = capital.solve("R", k_star_R, ftype=0,assort='positive',spread='lognorm',dispams=[0.00,0.2],
                             x_range=[0.95,1.05],scaling_x=norm_lognorm_05, normconst=c_lognorm_05)
sol_logP[2605]= capital.solve("P", k_star_P, ftype=0,assort='positive',spread='lognorm',dispams=[0.00,0.2],
                             x_range=[0.95,1.05],scaling_x=norm_lognorm_05, normconst=c_lognorm_05)
```

    Success! All workers and firms are matched
    Approximated solution failed to satisfy required assortativity condition.
    Success! All workers and firms are matched
    Approximated solution failed to satisfy required assortativity condition.

```python
sol_log[2610] = capital.solve("R", k_star_R, ftype=0,assort='positive',spread='lognorm',dispams=[0.00,0.2],
                             x_range=[0.9,1.1],scaling_x=norm_lognorm_10, normconst=c_lognorm_10)
sol_logP[2610]= capital.solve("P", k_star_P, ftype=0,assort='positive',spread='lognorm',dispams=[0.00,0.2],
                             x_range=[0.9,1.1],scaling_x=norm_lognorm_10, normconst=c_lognorm_10)
```

    Success! All workers and firms are matched
    Success! All workers and firms are matched
    Approximated solution failed to satisfy required assortativity condition.

```python
sol_log[2615] = capital.solve("R", k_star_R25, ftype=0,assort='positive',spread='lognorm',dispams=[0.00,0.2],
                             x_range=[0.85,1.15],scaling_x=norm_lognorm_15, normconst=c_lognorm_15,guess_overide=1/250)
sol_logP[2615]= capital.solve("P", k_star_P25, ftype=0,assort='positive',spread='lognorm',dispams=[0.00,0.2],
                             x_range=[0.85,1.15],scaling_x=norm_lognorm_15, normconst=c_lognorm_15,guess_overide=1/250)
```

    Success! All workers and firms are matched
    Success! All workers and firms are matched
    Approximated solution failed to satisfy required assortativity condition.

```python
sol_log[2620] = capital.solve("R", k_star_R25, ftype=0,assort='positive',spread='lognorm',dispams=[0.00,0.2],
                             x_range=[0.8,1.20],scaling_x=norm_lognorm_20, normconst=c_lognorm_20,guess_overide=1/250)
sol_logP[2620]= capital.solve("P", k_star_P25, ftype=0,assort='positive',spread='lognorm',dispams=[0.00,0.2],
                             x_range=[0.8,1.20],scaling_x=norm_lognorm_20, normconst=c_lognorm_20,guess_overide=1/250)
```

    Success! All workers and firms are matched
    Approximated solution failed to satisfy required assortativity condition.
    Success! All workers and firms are matched
    Approximated solution failed to satisfy required assortativity condition.

```python
sol_log[2625] = capital.solve("R", k_star_R25, ftype=0,assort='positive',spread='lognorm',dispams=[0.00,0.2],
                             x_range=[0.75,1.25],scaling_x=norm_lognorm_25, normconst=c_lognorm_25,guess_overide=1/250)
sol_logP[2625]= capital.solve("P", k_star_P25, ftype=0,assort='positive',spread='lognorm',dispams=[0.00,0.2],
                             x_range=[0.75,1.25],scaling_x=norm_lognorm_25, normconst=c_lognorm_25,guess_overide=1/250)
```

    Success! All workers and firms are matched
    Approximated solution failed to satisfy required assortativity condition.
    Success! All workers and firms are matched
    Approximated solution failed to satisfy required assortativity condition.

```python
sol_log[2630] = capital.solve("R", k_star_R30, ftype=0,assort='positive',spread='lognorm',dispams=[0.00,0.2],
                             x_range=[0.7,1.3],scaling_x=norm_lognorm_30, normconst=c_lognorm_30,guess_overide=1/250)
sol_logP[2630]= capital.solve("P", k_star_P30, ftype=0,assort='positive',spread='lognorm',dispams=[0.00,0.2],
                             x_range=[0.7,1.3],scaling_x=norm_lognorm_30, normconst=c_lognorm_30,guess_overide=1/250)
```

    Success! All workers and firms are matched
    Approximated solution failed to satisfy required assortativity condition.
    Success! All workers and firms are matched
    Approximated solution failed to satisfy required assortativity condition.

```python
# solsR = (sol_R1,sol_log[2605],sol_log[2610],sol_log[2615],sol_log[2620],sol_log[2625],sol_log[2630])
# solsP = (sol_P1,sol_logP[2605],sol_logP[2610],sol_logP[2615],sol_logP[2620],sol_logP[2625],sol_logP[2630])
# labs = ('$x=0$','$x\in[0.95,1.05]$ ','$x\in[0.90,1.10]$','$x\in[0.85,1.15]$','$x\in[0.80,1.20]$','$x\in[0.75,1.25]$','$x\in[0.70,1.30]$')

solsR = (sol_R1,sol_log[2610],sol_log[2620],sol_log[2630])
solsP = (sol_P1,sol_logP[2610],sol_logP[2620],sol_logP[2630])
labs = ('$x=1$','$x\in[0.90,1.10]$','$x\in[0.80,1.20]$','$x\in[0.70,1.30]$')
```

### Checking for PAM

As in the case of $x=1$ (all land quality is the same), we need to check ex-post for the PAM condition with capital. This needs to be negative at all points.

```python
plt.figure(figsize=(8,6))
plt.subplot(221)
plt.title('PAM Rich country')
for i,sol in enumerate(solsR):
    plt.plot(sol['xs'],PAM_R(sol['ys'],sol['xs'],1/sol['thetas'],1.0,sol['ks']), label=labs[i])
# plt.plot(sol_log[2625]['xs'],PAM_R(sol_log[2625]['ys'],sol_log[2625]['xs'],1/sol_log[2625]['thetas'],1.0,sol_log[2625]['ks']))
# plt.ylim(-0.2,0.2)
plt.legend(frameon=True)
plt.xscale("log")

plt.subplot(222)
plt.title('PAM$\geq$0 - Rich country')
for i,sol in enumerate(solsR):
    plt.plot(sol['xs'],PAM_R(sol['ys'],sol['xs'],1/sol['thetas'],1.0,sol['ks'])>=0, label=labs[i])
# plt.plot(sol_log[2625]['xs'],PAM_R(sol_log[2625]['ys'],sol_log[2625]['xs'],1/sol_log[2625]['thetas'],1.0,sol_log[2625]['ks'])>=0)
# plt.ylim(-0.2,0.2)
plt.xscale("log")

plt.subplot(223)
plt.title('PAM Poor country')
for i,sol in enumerate(solsP):
    plt.plot(sol['xs'],PAM_P(sol['ys'],sol['xs'],1/sol['thetas'],1.0,sol['ks']), label=labs[i])
# plt.ylim(-0.2,0.2)
plt.legend(frameon=True)
plt.xscale("log")

plt.subplot(224)
plt.title('PAM$\geq$0 - Poor country')
for i,sol in enumerate(solsP):
    plt.plot(sol['xs'],PAM_P(sol['ys'],sol['xs'],1/sol['thetas'],1.0,sol['ks'])>0, label=labs[i])

# plt.plot(sol_logP[2625]['xs'],PAM_P(sol_logP[2625]['ys'],sol_logP[2625]['xs'],1/sol_logP[2625]['thetas'],1.0,sol_logP[2625]['ks'])>=0)
# plt.ylim(-0.2,0.2)
plt.xscale("log")

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

```python
# Repackaging cell - uncomment if needed
# solsR = [sol_R1,] + list(sol_R10.values())
# solsP = [sol_P1] + list(sol_P10.values())
# labs = ('$x=1$','$x\in[0.975,1.025]$ ','$x\in[0.95,1.05]$ ','$x\in[0.925,1.075]$','$x\in[0.9,1.1]$')
```

```python
# Saving results cell - uncomment if needed
# for i,sol in enumerate(solsR):
#     data = np.vstack((sol['xs'],sol['ys'],1/sol['thetas'],sol['ks'])).T
#     np.savetxt('solR_'+labs[i]+'.csv', data, delimiter=',',header='xs, ys, ls, ks')

# for i,sol in enumerate(solsP):
#     data = np.vstack((sol['xs'],sol['ys'],1/sol['thetas'],sol['ks'])).T
#     np.savetxt('solP_'+labs[i]+'.csv', data, delimiter=',',header='xs, ys, ls, ks')
```

```python
# Loading results cell - uncomment if needed
# labs = ('$x=1$','$x\in[0.90,1.10]$','$x\in[0.80,1.20]$','$x\in[0.70,1.30]$')
# solsP = []
# for i,sol in enumerate(labs):
#     buffsol = pd.read_csv('solP_'+labs[i]+'.csv', header=0,names=['xs', 'ys', 'thetas', 'ks'])
#     buffsol['thetas'] = 1/buffsol['thetas']
#     solsP.append(buffsol)
# solsR = []
# for i,sol in enumerate(labs):
#     buffsol = pd.read_csv('solR_'+labs[i]+'.csv', header=0,names=['xs', 'ys', 'thetas', 'ks'])
#     buffsol['thetas'] = 1/buffsol['thetas']
#     solsR.append(buffsol)

```

## 4.3 Results

**Changes in farm size for widening spreads of $x$ **

In levels:

```python
capital.plot_comparative_land_choice(solsR,labs,solsP,labs, color=True, logs=False)
```

Focusing in the case of the poor country (Figure 1 in the corrigendum):

```python
capital.plot_land_choice(solsP,labs, color=True, yscale=300,logs=False,save_in='land_choice_corrigendum.eps')
```

```python
capital.plot_land_choice(solsP,labs, color=False, yscale=300,logs=False,save_in='land_choice_corrigendum_bw.eps')
```

In logs:

```python
capital.plot_comparative_land_choice(solsR,labs,solsP,labs, color=True, logs=True)
```

## 5. PAM condition

Finally, the formal, numerical check of the PAM condition. If lhs-rhs > 0, then PAM holds.

First without substituiting parameters:

```python
x, y, k = sym.var('x, y, k')
R, l, r, A, kappa, p_a, rho, gamma, eta = sym.var('R, l, r, A, kappa, p_a, rho, gamma, eta')
pam_dic_P['rho']=0.25
pam_dic_P['gamma']=0.5

pam_dic_R['rho']=0.25
pam_dic_R['gamma']=0.5

sym.init_printing()

F0 = r*A*kappa*(eta*k**rho + (1- eta)*(((l/r)**x)*y)**rho)**(gamma/rho)
F1 = r*A*kappa*x*(eta*k**rho + (1- eta)*((l/r)*y)**rho)**(gamma/rho)
F2 = r*A*kappa*(eta*(x*k)**rho + (1- eta)*((l/r)*y)**rho)**(gamma/rho)
F3 = r*A*kappa*(x*eta*k**rho + (1- eta)*((l/r)*y)**rho)**(gamma/rho)
F4 = r*A*kappa*(eta*k**(rho*x) + (1- eta)*((l/r)*y)**rho)**(gamma/rho)
```

Now substituting for the parameters of the rich country:

```python
lhs = F0.diff(x,y)*F0.diff(l,r)*F0.diff(k,k) - F0.diff(x,y)*F0.diff(l,k)*F0.diff(r,k) - F0.diff(x,k)*F0.diff(y,k)*F0.diff(l,r)
rhs = F0.diff(x,r)*F0.diff(y,l)*F0.diff(k,k) - F0.diff(x,r)*F0.diff(y,k)*F0.diff(l,k) - F0.diff(x,k)*F0.diff(y,l)*F0.diff(r,k)
print(sym.simplify(lhs-rhs).subs(pam_dic_R))
```

    0.00350791171645153*k**(-1.75)*r**2*x*(y*(l/r)**x)**0.5*(-0.00134143366352266*k**0.25 - 0.000496347002123088*(y*(l/r)**x)**0.25)*(0.890204456766942*k**0.25 + 0.109795543233058*(y*(l/r)**x)**0.25)**6.0/(l*y*(0.00471307000952237*k**0.25*(y*(l/r)**x)**0.75 + 0.0573192108348017*k**0.5*(y*(l/r)**x)**0.5 + 0.309823243220936*k**0.75*(y*(l/r)**x)**0.25 + 0.627999151431458*k**1.0 + 0.000145324503280975*(y*(l/r)**x)**1.0))

Now substituting for the parameters of the poor country:

```python
print(sym.simplify(lhs-rhs).subs(pam_dic_P))
```

    0.000996785143643158*k**(-1.75)*r**2*x*(y*(l/r)**x)**0.5*(-0.00134143366352266*k**0.25 - 0.000496347002123088*(y*(l/r)**x)**0.25)*(0.890204456766942*k**0.25 + 0.109795543233058*(y*(l/r)**x)**0.25)**6.0/(l*y*(0.00471307000952237*k**0.25*(y*(l/r)**x)**0.75 + 0.0573192108348017*k**0.5*(y*(l/r)**x)**0.5 + 0.309823243220936*k**0.75*(y*(l/r)**x)**0.25 + 0.627999151431458*k**1.0 + 0.000145324503280975*(y*(l/r)**x)**1.0))