# Learning disentangled representations in 'grid-world'

## 1. Preamble

```python
import os

import torch
import random
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt

import time
from IPython import display

try:
    import seaborn as sns
    sns.set()
    print("Set seaborn plotting defaults.")
except:
    print("seaborn not found.")
```

    Set seaborn plotting defaults.

```python
def mk_dir(export_dir, quite=False):
    if not os.path.exists(export_dir):
            try:
                os.makedirs(export_dir)
                print('created dir: ', export_dir)
            except OSError as exc: # Guard against race condition
                 if exc.errno != exc.errno.EEXIST:
                    raise
            except Exception:
                pass
    else:
        print('dir already exists: ', export_dir)
```

## 2. Defining objects

### GridWorld environment

First I'll define the GridWorld environment to work in.  The basic idea is that the ball is on a 2-d grid with $n$ discrete points in each dimension.  The available actions step the ball one step in each direction (i.e. there are 4 actions: up/down/left/right).

Initally, the observations returned are just an $n \times n$ matrix, with zeros everywhere except where the ball is, which is 1.

```python
class GridWorld():

    class action_space():
        def __init__(self,n_actions):
            self.n = n_actions

        def sample(self, k=1):
            return torch.randint(0,self.n,(k,))

    class observation_space():
        def __init__(self,n_features):
            self.shape = [n_features]

    def __init__(self,dim=5):

        self.dim = dim

        self.action_space = self.action_space(4)
        self.observation_space = self.observation_space(dim**2)

        self.reset()

    def reset(self):
        self.ball_coordinates = [random.randrange(self.dim), random.randrange(self.dim)]
        self.state = self.get_state()
        return self.get_observation()

    def get_state(self, ball_coordinates=None):
        if ball_coordinates is None:
            ball_coordinates=self.ball_coordinates

        state = torch.zeros((self.dim,self.dim))
        state[ball_coordinates[0],ball_coordinates[1]] = 1

        return state

    def get_observation(self):
        return self.state.flatten()

    def step(self,action):

        if action == 0:
            self.ball_coordinates[0] = (self.ball_coordinates[0] + 1)%self.dim
        elif action == 1:
            self.ball_coordinates[0] = (self.ball_coordinates[0] - 1)%self.dim
        elif action == 2:
            self.ball_coordinates[1] = (self.ball_coordinates[1] + 1)%self.dim
        elif action == 3:
            self.ball_coordinates[1] = (self.ball_coordinates[1] - 1)%self.dim
        else:
            raise Exception("Invalid action.")

        self.state = self.get_state()

        return self.get_observation()
```

Now show a few consecutive states from this gridworld

```python
def plot_state(state, ax):
    ax.pcolormesh(env.state, edgecolors='gray', linewidth=2, )
    ax.set_aspect('equal')
    ax.set_xticks([])
    ax.set_yticks([])

    return ax

n_steps = 5

env = GridWorld()

fig, axs = plt.subplots(1, n_steps+1,)

plot_state(env.state,axs[0])

for i in range(n_steps):
    env.step(np.random.randint(0,4))
    plot_state(env.state,axs[i+1])
```

### Latent space

**Encoder/Decoder**

Now we want to learn to represent this environment in some latent space (which we, for now, simply assume to be 4-dimensional).  We will require both an encoder and decoder.

```python
class Coder(nn.Module):

    def __init__(self, n_in, n_out, n_hid=[32], output_activation=nn.Sigmoid):
        super().__init__()

        if type(n_hid) != list:
            n_hid = [n_hid]
        n_layers = [n_in] + n_hid + [n_out]

        self.layers = []
        for i_layer, (n1, n2) in enumerate(zip(n_layers, n_layers[1:])):
            mods = [nn.Linear(n1, n2, bias=True)]
            act_fn = nn.ReLU if i_layer < len(n_layers) - 2 else output_activation
            if act_fn is not None:
                mods.append(act_fn())
            layer = nn.Sequential(*mods)
            self.layers.append(layer)

        self.layers = nn.ModuleList(self.layers)

    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

class nnNorm(nn.Module):

    def __init__(self, dim=-1):
        super().__init__()
        self.dim=dim

    def forward(self, x):
        return F.normalize(x, dim=self.dim)

class Decoder(Coder):

    def __init__(self, n_in, n_out, n_hid=[32]):
        super().__init__(n_in, n_out, n_hid, output_activation=nn.Sigmoid)

class Encoder(Coder):

    def __init__(self, n_in, n_out, n_hid=[32]):
        super().__init__(n_in, n_out, n_hid, output_activation=nnNorm)
```

```python
encoder = Encoder(n_in=5*5, n_out=4, n_hid=[32])
print(encoder)

decoder = Decoder(n_in=4, n_out=5*5, n_hid=[32])
print(decoder)
```

    Encoder(
      (layers): ModuleList(
        (0): Sequential(
          (0): Linear(in_features=25, out_features=32, bias=True)
          (1): ReLU()
        )
        (1): Sequential(
          (0): Linear(in_features=32, out_features=4, bias=True)
          (1): nnNorm()
        )
      )
    )
    Decoder(
      (layers): ModuleList(
        (0): Sequential(
          (0): Linear(in_features=4, out_features=32, bias=True)
          (1): ReLU()
        )
        (1): Sequential(
          (0): Linear(in_features=32, out_features=25, bias=True)
          (1): Sigmoid()
        )
      )
    )

Check to make sure dimensions are okay.

```python
x = torch.rand(25)
z = encoder(x)
x_recon = decoder(z)
print("Observable rep (shape = {}) : \n {}".format(x.shape, x))
print("Latent rep (shape = {}, norm = {}) : \n {}".format(z.shape, torch.norm(z), z))
print("Reconstructed Observable rep (shape = {}) : \n {}".format(x_recon.shape, x_recon))
```

    Observable rep (shape = torch.Size([25])) :
     tensor([0.2276, 0.8350, 0.4096, 0.0760, 0.1850, 0.5055, 0.7176, 0.2933, 0.2344,
            0.8434, 0.6855, 0.8067, 0.9070, 0.0959, 0.9930, 0.4419, 0.1822, 0.0209,
            0.8426, 0.0186, 0.0328, 0.8303, 0.1454, 0.4505, 0.5542])
    Latent rep (shape = torch.Size([4]), norm = 0.9999999403953552) :
     tensor([ 0.3447,  0.0625, -0.9034, -0.2474], grad_fn=<DivBackward0>)
    Reconstructed Observable rep (shape = torch.Size([25])) :
     tensor([0.4819, 0.5674, 0.5033, 0.4692, 0.4724, 0.3850, 0.4723, 0.4560, 0.5078,
            0.4804, 0.4849, 0.5639, 0.5185, 0.4518, 0.5257, 0.4698, 0.5441, 0.4963,
            0.5146, 0.4765, 0.5287, 0.4546, 0.4676, 0.5099, 0.5143],
           grad_fn=<SigmoidBackward>)

**Representation**

The crux of the matter is learning to 'represent' actions in the observation space with actions in latent space.  Here, we will do this by assuming every action is a generalized rotation in latent space, which we denote with a series of 2-dimensional rotations.

A 2-d rotation is given by:

\begin{pmatrix}
\cos(\theta) & \sin(\theta) \\
-\sin(\theta) & \cos(\theta)
\end{pmatrix}

and we denote a rotation in dimensions $i$ and $j$ of a higher dimensional space as $R_{i,j}(\theta)$.  For $i=1$, $j=4$, in a 4-dimensional space:

\begin{equation}
R_{1,4}(\theta) =
\begin{pmatrix}
\cos(\theta) & 0 & 0 & \sin(\theta) \\
0 & 1 & 0 & 0 \\
0 & 0 & 1 & 0 \\
-\sin(\theta) & 0 & 0 & \cos(\theta)
\end{pmatrix}
\end{equation}

An arbitrary rotation, denoted $g$ as I am subtly moving towards this being a group action, can then be written as:

\begin{equation}
    g(\theta_{1,2},\theta_{1,3},\dots,\theta_{n-1,n}) = \prod_{i=1}^{n-1} \prod_{j=1+1}^{n} R_{i,j}(\theta_{i,j})
\end{equation}

which has $n(n-1)/2$ free parameters (i.e. $\theta_{i,j}$'s).

```python
class Representation():

    def __init__(self, dim=4):
        self.dim = dim
        self.params = dim*(dim-1)//2
        self.thetas = torch.autograd.Variable(torch.rand(self.params)/dim, requires_grad=True)

        self.__matrix = None

    def set_thetas(self, thetas):
        self.thetas = thetas
        self.thetas.requires_grad = True
        self.clear_matrix()

    def clear_matrix(self):
        self.__matrix = None

    def get_matrix(self):
        if self.__matrix is None:
            k = 0
            mats = []
            for i in range(self.dim-1):
                for j in range(self.dim-1-i):
                    theta_ij = self.thetas[k]
                    k+=1
                    c, s = torch.cos(theta_ij), torch.sin(theta_ij)

                    rotation_i = torch.eye(self.dim, self.dim)
                    rotation_i[i, i] = c
                    rotation_i[i, i+j+1] = s
                    rotation_i[j+i+1, i] = -s
                    rotation_i[j+i+1, j+i+1] = c

                    mats.append(rotation_i)

            def chain_mult(l):
                if len(l)>=3:
                    return l[0]@l[1]@chain_mult(l[2:])
                elif len(l)==2:
                    return l[0]@l[1]
                else:
                    return l[0]

            self.__matrix = chain_mult(mats)

        return self.__matrix
```

**LatentWorld**

Now, for symmetry's sake, we'll also have a `LatentWorld` which acts as the environment in the latent space.

```python
class LatentWorld():

    class action_space():
        def __init__(self,n_actions):
            self.n = n_actions

        def sample(self, k=1):
            return torch.randint(0,self.n,(k,))

    class observation_space():
        def __init__(self,n_features):
            self.shape = [n_features]

    def __init__(self,
                 dim=4,
                 n_actions=4,
                 action_reps=None):

        self.dim = dim

        self.action_space = self.action_space(n_actions)
        self.observation_space = self.observation_space(dim)

        if action_reps is None:
            self.action_reps = [Representation(dim=self.dim) for _ in range(n_actions)]
        else:
            if len(action_reps)!=n_actions:
                raise Exception("Must pass an action representation for every action.")
            if not all([rep.dim==self.dim]):
                raise Exception("Action representations do not act on the dimension of the latent space.")
            self.action_reps = action_reps

    def reset(self, state_init):
        self.state = state_init
        return self.get_observation()

    def clear_representations(self):
        for rep in self.action_reps:
            rep.clear_matrix()

    def get_representation_params(self):
        params = []
        for rep in self.action_reps:
            params.append(rep.thetas)
        return params

    def save_representations(self, path):
        if os.path.splitext(path)[-1] != '.pth':
            path += '.pth'
        rep_thetas = [rep.thetas for rep in self.action_reps]
        return torch.save(rep_thetas, path)

    def load_reprentations(self, path):
        rep_thetas = torch.load(path)
        for rep in self.action_reps:
            rep.set_thetas(rep_thetas.pop(0))

    def get_observation(self):
        return self.state

    def step(self,action):
        self.state = torch.mv(self.action_reps[action].get_matrix(), self.state)
        obs = self.get_observation()
        return obs
```

## 3. Training

So the basic training loop is pretty straightfoward.  We simply play out episodes from random starting configurations, encoded by the `Encoder`, for `ep_steps` time-steps.  Each random action is executed in both the `GridWorld` and the `LatentWorld`, and then the latent state is transformed to the observation space by the `Decoder` where the loss function measures its deviation from the true state.

```python
obs_env = GridWorld(dim = 11)
lat_env = LatentWorld(dim = 4,
                      n_actions = obs_env.action_space.n)
decoder = Decoder(n_in = sum(lat_env.observation_space.shape),
                  n_out = sum(obs_env.observation_space.shape),
                  n_hid = [32])
encoder = Encoder(n_in = sum(obs_env.observation_space.shape),
                  n_out = sum(lat_env.observation_space.shape),
                  n_hid = [32])

optimizer_dec = optim.Adam(decoder.parameters(),
                           lr=2e-2,
                           weight_decay=0)

optimizer_enc = optim.Adam(encoder.parameters(),
                           lr=5e-2,
                           weight_decay=0)

optimizer_rep = optim.Adam(lat_env.get_representation_params(),
                           lr=2e-2,
                           weight_decay=0)

losses = []
```

```python
n_sgd_steps = 5000
ep_steps = 10
batch_eps = 16

i = 0

t_start = time.time()

temp = 0

while i < n_sgd_steps:

    loss = torch.zeros(1)

    for _ in range(batch_eps):
        t_ep = -1
        while t_ep < ep_steps:
            if t_ep == -1:
                obs_x = obs_env.reset()
                obs_z = lat_env.reset(encoder(obs_x))
            else:
                action = obs_env.action_space.sample().item()
                obs_x = obs_env.step(action)
                obs_z = lat_env.step(action)

            t_ep += 1

            obs_x_recon = decoder(obs_z)

            loss += F.binary_cross_entropy(obs_x_recon, obs_x)

    loss /= (ep_steps*batch_eps)

    losses.append(loss.item())

    optimizer_dec.zero_grad()
    optimizer_enc.zero_grad()
    optimizer_rep.zero_grad()
    loss.backward()
    optimizer_enc.step()
    optimizer_dec.step()
    optimizer_rep.step()

    # Rember to clear the cached action representations after we update the parameters!
    lat_env.clear_representations()

    i+=1

    if i%10==0:
        print("iter {} : loss={:.3f} : last 10 iters in {:.3f}s".format(i, loss.item(), time.time() - t_start),
              end="\r" if i%100 else "\n")
        t_start = time.time()
```

    iter 100 : loss=0.052 : last 10 iters in 0.570s
    iter 200 : loss=0.053 : last 10 iters in 0.700s
    iter 300 : loss=0.053 : last 10 iters in 0.590s
    iter 400 : loss=0.053 : last 10 iters in 0.603s
    iter 500 : loss=0.053 : last 10 iters in 0.812s
    iter 600 : loss=0.052 : last 10 iters in 0.631s
    iter 700 : loss=0.049 : last 10 iters in 0.575s
    iter 800 : loss=0.038 : last 10 iters in 0.729s
    iter 900 : loss=0.035 : last 10 iters in 0.647s
    iter 1000 : loss=0.033 : last 10 iters in 0.550s
    iter 1100 : loss=0.032 : last 10 iters in 0.692s
    iter 1200 : loss=0.034 : last 10 iters in 0.679s
    iter 1300 : loss=0.033 : last 10 iters in 0.623s
    iter 1400 : loss=0.031 : last 10 iters in 0.556s
    iter 1500 : loss=0.032 : last 10 iters in 0.582s
    iter 1600 : loss=0.029 : last 10 iters in 0.555s
    iter 1700 : loss=0.022 : last 10 iters in 0.561s
    iter 1800 : loss=0.017 : last 10 iters in 0.590s
    iter 1900 : loss=0.014 : last 10 iters in 0.606s
    iter 2000 : loss=0.009 : last 10 iters in 0.599s
    iter 2100 : loss=0.006 : last 10 iters in 0.667s
    iter 2200 : loss=0.006 : last 10 iters in 0.627s
    iter 2300 : loss=0.004 : last 10 iters in 0.801s
    iter 2400 : loss=0.002 : last 10 iters in 0.628s
    iter 2500 : loss=0.002 : last 10 iters in 0.610s
    iter 2600 : loss=0.001 : last 10 iters in 2.186s
    iter 2700 : loss=0.001 : last 10 iters in 0.845s
    iter 2800 : loss=0.000 : last 10 iters in 0.574s
    iter 2900 : loss=0.000 : last 10 iters in 0.797s
    iter 3000 : loss=0.000 : last 10 iters in 0.623s
    iter 3100 : loss=0.005 : last 10 iters in 0.567s
    iter 3200 : loss=0.008 : last 10 iters in 0.802s
    iter 3300 : loss=0.001 : last 10 iters in 0.601s
    iter 3400 : loss=0.000 : last 10 iters in 0.593s
    iter 3500 : loss=0.000 : last 10 iters in 0.872s
    iter 3600 : loss=0.000 : last 10 iters in 0.618s
    iter 3700 : loss=0.000 : last 10 iters in 0.669s
    iter 3800 : loss=0.000 : last 10 iters in 0.703s
    iter 3900 : loss=0.000 : last 10 iters in 0.571s
    iter 4000 : loss=0.000 : last 10 iters in 0.770s
    iter 4100 : loss=0.000 : last 10 iters in 0.606s
    iter 4200 : loss=0.000 : last 10 iters in 0.704s
    iter 4300 : loss=0.000 : last 10 iters in 0.778s
    iter 4400 : loss=0.000 : last 10 iters in 0.613s
    iter 4500 : loss=0.000 : last 10 iters in 0.585s
    iter 4600 : loss=0.000 : last 10 iters in 0.781s
    iter 4700 : loss=0.000 : last 10 iters in 0.609s
    iter 4800 : loss=0.000 : last 10 iters in 0.592s
    iter 4900 : loss=0.000 : last 10 iters in 0.616s
    iter 5000 : loss=0.000 : last 10 iters in 1.051s

**Plot training curve, then save the network**

```python
mk_dir("_data")

with plt.style.context('seaborn-paper', after_reset=True):

    fig, ax1 = plt.subplots(1, 1, figsize=(4, 2.5))

    window = 25
    avg_mask = np.ones(window) / window

    ax1.plot(np.convolve(range(len(losses)), avg_mask, 'valid'),
             np.convolve(np.log(losses), avg_mask, 'valid'),
             linewidth=0.75,
             alpha=0.8)
#     ax1.set_yscale("log")

    ax1.set_xlabel("Num. parameter updates")
    ax1.set_ylabel("Log Loss")

    fig_fname = "_data/training_loss"
    plt.savefig(fig_fname + ".pdf", bbox_inches='tight')
    plt.savefig(fig_fname + ".png", bbox_inches='tight')
```

```python
torch.save(decoder.state_dict(), "decoder.pth")
torch.save(encoder.state_dict(), "encoder.pth")
rep_thetas = [rep.thetas for rep in lat_env.action_reps]
torch.save(rep_thetas, "representations.pth")
torch.save(optimizer_dec.state_dict(), "optimizer_dec.pth")
torch.save(optimizer_enc.state_dict(), 'optimizer_enc.pth')
torch.save(optimizer_rep.state_dict(), "optimizer_rep.pth")
```

## 4. Testing

Testing is easy too, we just play out an episode and see how well the reconstructed image agrees with the ground truth!

```python
def plot_state(obs, ax, one_hot=True):
    if one_hot:
        obs=(obs == obs.max())
    ax.pcolormesh(obs.reshape(obs_env.dim, obs_env.dim), edgecolors='gray', linewidth=2)
    ax.set_aspect('equal')
    ax.set_xticks([])
    ax.set_yticks([])

    return ax

n_steps = 10

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

ax1.set_title("Ground truth")
ax2.set_title("Reconstruction")

for i in range(n_steps+1):

    if i==0:
        action = "N\A"
        obs_x = obs_env.reset()
        obs_z = lat_env.reset(encoder(obs_x))
    else:
        action = obs_env.action_space.sample().item()
        obs_x = obs_env.step(action)
        obs_z = lat_env.step(action)

    obs_x_recon = decoder(obs_z)

    fig.suptitle('step {} : last action = {}'.format(i, action), fontsize=16)

    plot_state(obs_x.detach().numpy(),ax1)
    plot_state(obs_x_recon.detach().numpy(),ax2, one_hot=False)

    fig_fname = "_data/testing_{}".format(i)
    plt.savefig(fig_fname + ".pdf", bbox_inches='tight')
    plt.savefig(fig_fname + ".png", bbox_inches='tight')

    display.clear_output(wait=True)
    display.display(plt.gcf())
    time.sleep(0.5)

display.clear_output(wait=False)
```

We will now have a look at the latent space, we will make a 2D projection of the 4D latent space for every possible frame (There are 25 for the 5x5 grid). Note that since we use random projections, in some cases the toric structure we find is more obvious than in others.

**1) Positions on the grid**

```python
from sklearn.random_projection import GaussianRandomProjection
import matplotlib as mpl
from mpl_toolkits.mplot3d import Axes3D

fig = plt.figure()
ax = fig.add_subplot(111)#, projection='3d')

num_frames = sum(obs_env.observation_space.shape)

latent_points = []

for i in range(num_frames):
    frame = torch.zeros(num_frames)
    frame[i] = 1
    latent = encoder(frame)
    latent_points.append(latent.detach().tolist())

latent_map = np.array(latent_points)

pca = GaussianRandomProjection(n_components=2)

latent_2d = pca.fit_transform(latent_map)

#color=["C0","C1","C2","C3","C4","C5","C6","C7","C8","C9","C10"]
color=['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', '#8c564b', '#e377c2', '#7f7f7f', '#bcbd22', '#17becf',"#2fa36b"]
marks=["1","2","3","4","+",">","<","^","v","x","d"]

s=[120]*5+[50]*6

for i in range (121):
    ax.scatter(x=i//11,
                y=i%11,
   #             zs=latent_2d.transpose()[2][i],
                c=color[i//11],
                s=s[i%11],
                marker=marks[i%11])

fig.show()

```

**2) Positions in Latent Space**

```python
fig = plt.figure()
ax = fig.add_subplot(111)#, projection='3d')
s=[120]*5+[50]*6
for i in range (121):
    ax.scatter(x=latent_2d.transpose()[0][i],
                y=latent_2d.transpose()[1][i],
             #   zs=latent_2d.transpose()[2][i],
                c=color[i//11],
                s=s[i%11],
                marker=marks[i%11])
 #   ax.set_xlim(-.6/1.4,.6/1.4)
  #  ax.set_ylim(-.8/1.4,.8/1.4)
  #  ax.set_zlim(-1./1.6,1./1.6)
    #ax.view_init(elev=45, azim=45)
fig.show()

```

**3) Action representations**

```python
width=0.75

test_folder="_data"

for rep in lat_env.action_reps:
    print(rep.get_matrix())
    print(torch.matrix_power(rep.get_matrix(), 5))
rep_thetas = [rep.thetas.detach().numpy() for rep in lat_env.action_reps]

plt_lim = max( 0.1, max([max(t) for t in rep_thetas])/(2*np.pi) )
titles = ["up", "down", "right", "left"]

with plt.style.context('seaborn-paper', after_reset=True):

    fig, axs = plt.subplots(1, len(rep_thetas), figsize=(8, 1.5), gridspec_kw={"wspace":0.4})

    for i, thetas in enumerate(rep_thetas):
        x = np.arange(len(thetas))
        axs[i].bar(x - width/2, thetas/(2*np.pi), width, label='Rep {}'.format(i))

        axs[i].set_xticks(x-0.25)
        axs[i].set_xticklabels(["12","13","14","23","24","34"])
        axs[i].set_ylabel("$ij$")

        axs[i].set_ylim(-plt_lim,plt_lim)

        axs[i].set_title(titles[i])

    axs[0].set_ylabel(r"$\theta / 2\pi$", fontsize=11)

    fig_fname = os.path.join(test_folder, "thetas")

    plt.savefig(fig_fname + ".pdf", bbox_inches='tight')
    plt.savefig(fig_fname + ".png", bbox_inches='tight')
```

## 5. Disentanglement

***Some jargon***

It's nice that it works, but the real point here is to try and learn a *disentangled* representation of the actions.

Before considering how best to do this, we want to define a metric of 'disentanglement'.  We consider the evolution of an observable (latent) vector, $x \in X$ ($z \in Z$), under the element $g \in G$ of the group of symmetries generating transformations of the object.  Then we are looking for a representation, $\rho:G \rightarrow GL(V)$, such that the transformation is linear in the latent space, i.e.
\begin{equation}
    z^{\prime} = \rho(g) \cdot z.
\end{equation}
Note, in our case, the representations are the rotation matrices we learn.

For this representation to be disentangled, it means that if there exists a subgroup decomposition of $G$
\begin{equation}
    G = G_1 \times G_2 \times \dots \times G_n,
\end{equation}
then we equivalently decompose the representation, $(\rho, G)$, into subrepresentations:
\begin{equation}
    V = V_1 \oplus V_2 \oplus \dots \oplus V_n
\end{equation}
such that the restricted subrepresentations $(\rho_{\vert G_i}, V_i)_i$ are non-trivial, and the restricted subrepresentations $(\rho_{\vert G_i}, V_j)_{j \neq i}$ are trivial.

In our context, a GridWorld with 5 points in each dimension is represented by $G = C_5 \times C_5$ (where $C_5$ is the cyclic group).  This is a subgroup of $\mathrm{SO}(2) \times \mathrm{SO}(2)$, therefore we hope to find the disentangled representation of the actions (up, down, left, right) that corresponds to this.

***Some practicalities***

Our intuition is that the disentangled representation acts as the identity on as many dimensions as possible.  We could attempt to enforce this with some regularization during training.  Normal weight decay won't cut it, as that tries to reduce all weights, where as what we really want to do is have all *but one* of our thetas (which corresponds to the rotation/coupling of two dimensions) to be zero.

**1. Entanglement regularisation**

So for $m$ parameters, ${\theta_1, \dots, \theta_m}$, we want to regularise with
\begin{equation}
    \sum_{i \neq j} \vert\theta_i\vert^2, \mathrm{where\ } \theta_j {=} \mathrm{max_k}({\vert\theta_k\vert}).
\end{equation}
We will also use this term as our metric of 'entanglement'.

```python
def calc_entanglement(params):
    params = params.abs().pow(2)
    return params.sum() - params.max()

params = torch.FloatTensor([1,1,0.5,0,0])
calc_entanglement(params)
```

    tensor(1.2500)

### Training with regularization

```python
obs_env = GridWorld(dim = 5)
lat_env = LatentWorld(dim = 4,
                      n_actions = obs_env.action_space.n)
decoder = Decoder(n_in = sum(lat_env.observation_space.shape),
                  n_out = sum(obs_env.observation_space.shape),
                  n_hid = [32])
encoder = Encoder(n_in = sum(obs_env.observation_space.shape),
                  n_out = sum(lat_env.observation_space.shape),
                  n_hid = [32])

optimizer_enc = optim.Adam(encoder.parameters(),
                           lr=2e-2,
                           weight_decay=0)

optimizer_dec = optim.Adam(decoder.parameters(),
                           lr=2e-2,
                           weight_decay=0)

optimizer_rep = optim.Adam(lat_env.get_representation_params(),
                           lr=2e-2,
                           weight_decay=0)

losses = []
entanglement = []
```

```python
n_sgd_steps = 5000
ep_steps = 5
batch_eps = 2

i = 0

t_start = time.time()

temp = 0

while i < n_sgd_steps:

    loss = torch.zeros(1)

    for _ in range(batch_eps):
        t_ep = -1
        while t_ep < ep_steps:
            if t_ep == -1:
                obs_x = obs_env.reset()
                obs_z = lat_env.reset(encoder(obs_x))
            else:
                action = obs_env.action_space.sample().item()
                obs_x = obs_env.step(action)
                obs_z = lat_env.step(action)

            t_ep += 1

            obs_x_recon = decoder(obs_z)
            #(obs_x_recon==obs_x_recon.max()).float()

            loss += F.binary_cross_entropy(obs_x_recon, obs_x)

    loss /= (ep_steps * batch_eps)
    raw_loss = loss.item()

    reg_loss = sum([calc_entanglement(r.thetas) for r in lat_env.action_reps])/4

    loss += reg_loss * 4e-2

    losses.append(raw_loss)
    entanglement.append(reg_loss.item())

    optimizer_dec.zero_grad()
    optimizer_enc.zero_grad()
    optimizer_rep.zero_grad()
    loss.backward()
    optimizer_enc.step()
    optimizer_dec.step()
    optimizer_rep.step()

    # Remember to clear the cached action representations after we update the parameters!
    lat_env.clear_representations()

    i+=1

    if i%10==0:
        print("iter {} : loss={:.3f} : entanglement={:.2e} : last 10 iters in {:.3f}s".format(
            i, raw_loss, reg_loss.item(), time.time() - t_start
            ), end="\r" if i%100 else "\n")
        t_start = time.time()
```

    iter 100 : loss=0.199 : entanglement=9.13e-04 : last 10 iters in 0.097s
    iter 200 : loss=0.215 : entanglement=2.20e-03 : last 10 iters in 0.165s
    iter 300 : loss=0.193 : entanglement=1.27e-02 : last 10 iters in 0.093s
    iter 400 : loss=0.200 : entanglement=4.60e-02 : last 10 iters in 0.110s
    iter 500 : loss=0.154 : entanglement=4.31e-02 : last 10 iters in 0.216s
    iter 600 : loss=0.155 : entanglement=1.46e-02 : last 10 iters in 0.112s
    iter 700 : loss=0.093 : entanglement=1.78e-02 : last 10 iters in 0.144s
    iter 800 : loss=0.036 : entanglement=1.52e-02 : last 10 iters in 0.166s
    iter 900 : loss=0.009 : entanglement=3.21e-03 : last 10 iters in 0.156s
    iter 1000 : loss=0.012 : entanglement=4.33e-03 : last 10 iters in 0.157s
    iter 1100 : loss=0.002 : entanglement=1.69e-04 : last 10 iters in 0.133s
    iter 1200 : loss=0.001 : entanglement=6.21e-05 : last 10 iters in 0.166s
    iter 1300 : loss=0.001 : entanglement=2.38e-05 : last 10 iters in 0.176s
    iter 1400 : loss=0.001 : entanglement=1.67e-05 : last 10 iters in 0.187s
    iter 1500 : loss=0.001 : entanglement=1.18e-05 : last 10 iters in 0.111s
    iter 1600 : loss=0.000 : entanglement=6.62e-06 : last 10 iters in 0.159s
    iter 1700 : loss=0.000 : entanglement=2.98e-06 : last 10 iters in 0.152s
    iter 1800 : loss=0.000 : entanglement=3.46e-06 : last 10 iters in 0.162s
    iter 1900 : loss=0.000 : entanglement=3.16e-06 : last 10 iters in 0.100s
    iter 2000 : loss=0.000 : entanglement=1.19e-06 : last 10 iters in 0.100s
    iter 2100 : loss=0.000 : entanglement=1.55e-06 : last 10 iters in 0.101s
    iter 2200 : loss=0.000 : entanglement=1.25e-06 : last 10 iters in 0.114s
    iter 2300 : loss=0.000 : entanglement=1.73e-06 : last 10 iters in 0.105s
    iter 2400 : loss=0.000 : entanglement=9.83e-07 : last 10 iters in 0.099s
    iter 2500 : loss=0.000 : entanglement=1.37e-06 : last 10 iters in 0.138s
    iter 2600 : loss=0.000 : entanglement=6.26e-07 : last 10 iters in 0.116s
    iter 2700 : loss=0.000 : entanglement=3.28e-07 : last 10 iters in 0.163s
    iter 2800 : loss=0.000 : entanglement=7.15e-07 : last 10 iters in 0.198s
    iter 2900 : loss=0.000 : entanglement=6.26e-07 : last 10 iters in 0.119s
    iter 3000 : loss=0.000 : entanglement=3.58e-07 : last 10 iters in 0.109s
    iter 3100 : loss=0.000 : entanglement=5.66e-07 : last 10 iters in 0.191s
    iter 3200 : loss=0.000 : entanglement=2.68e-07 : last 10 iters in 0.107s
    iter 3300 : loss=0.000 : entanglement=3.58e-07 : last 10 iters in 0.117s
    iter 3400 : loss=0.000 : entanglement=3.58e-07 : last 10 iters in 0.119s
    iter 3500 : loss=0.000 : entanglement=3.58e-07 : last 10 iters in 0.135s
    iter 3600 : loss=0.000 : entanglement=1.49e-07 : last 10 iters in 0.149s
    iter 3700 : loss=0.000 : entanglement=1.79e-07 : last 10 iters in 0.099s
    iter 3800 : loss=0.000 : entanglement=1.79e-07 : last 10 iters in 0.121s
    iter 3900 : loss=0.000 : entanglement=2.38e-07 : last 10 iters in 0.116s
    iter 4000 : loss=0.000 : entanglement=1.49e-07 : last 10 iters in 0.104s
    iter 4100 : loss=0.000 : entanglement=2.98e-07 : last 10 iters in 0.106s
    iter 4200 : loss=0.000 : entanglement=2.68e-07 : last 10 iters in 0.110s
    iter 4300 : loss=0.000 : entanglement=1.49e-07 : last 10 iters in 0.142s
    iter 4400 : loss=0.000 : entanglement=2.09e-07 : last 10 iters in 0.111s
    iter 4500 : loss=0.000 : entanglement=8.94e-08 : last 10 iters in 0.106s
    iter 4600 : loss=0.000 : entanglement=1.19e-07 : last 10 iters in 0.096s
    iter 4700 : loss=0.000 : entanglement=8.94e-08 : last 10 iters in 0.094s
    iter 4800 : loss=0.000 : entanglement=5.96e-08 : last 10 iters in 0.118s
    iter 4900 : loss=0.000 : entanglement=8.94e-08 : last 10 iters in 0.106s
    iter 5000 : loss=0.000 : entanglement=2.98e-08 : last 10 iters in 0.114s

```python
exp_name = "ent_reg_1e-2_orthNorm_reg_1e-2_v10"
# exp_name = "testing"

save_folder = os.path.join("_data/decoder_only", exp_name)
train_folder = os.path.join(save_folder,"train")
test_folder = os.path.join(save_folder,"test")

mk_dir(save_folder)
mk_dir(train_folder)
mk_dir(test_folder)
```

    created dir:  _data/decoder_only/ent_reg_1e-2_orthNorm_reg_1e-2_v10
    created dir:  _data/decoder_only/ent_reg_1e-2_orthNorm_reg_1e-2_v10/train
    created dir:  _data/decoder_only/ent_reg_1e-2_orthNorm_reg_1e-2_v10/test

**Show training curves**

```python
with plt.style.context('seaborn-paper', after_reset=True):

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 2.5), gridspec_kw={"wspace":0.3})

    window = 25
    avg_mask = np.ones(window) / window

    ax1.plot(np.convolve(range(len(losses)), avg_mask, 'valid'),
             np.convolve(np.log(losses), avg_mask, 'valid'),
             linewidth=0.75,
             alpha=0.8)
#     ax1.set_yscale("log")

    ax1.set_xlabel("Num. parameter updates")
    ax1.set_ylabel("Log Loss")

    ax2.plot(np.convolve(range(len(entanglement)), avg_mask, 'valid'),
             np.convolve(entanglement, avg_mask, 'valid'),
             linewidth=0.75,
             alpha=0.8)
    ax2.set_yscale("log")

    ax2.set_xlabel("Num. parameter updates")
    ax2.set_ylabel("Entanglement")

    fig_fname = os.path.join(train_folder, "training_curves")
    plt.savefig(fig_fname + ".pdf", bbox_inches='tight')
    plt.savefig(fig_fname + ".png", bbox_inches='tight')
```

```python
torch.save(decoder.state_dict(), os.path.join(save_folder,"decoder.pth"))
rep_thetas = [rep.thetas for rep in lat_env.action_reps]
torch.save(rep_thetas, os.path.join(save_folder,"representations.pth"))
torch.save(optimizer_dec.state_dict(), os.path.join(train_folder,"optimizer_dec.pth"))
torch.save(optimizer_rep.state_dict(), os.path.join(train_folder,"optimizer_rep.pth"))
```

### Testing: action representations

```python
width=0.5

rep_thetas = [rep.thetas.detach().numpy() for rep in lat_env.action_reps]

for rep in lat_env.action_reps:
    print(rep.get_matrix())
    print(torch.matrix_power(rep.get_matrix(), 5))

plt_lim = max( 0.22, max([max(t) for t in rep_thetas])/(2*np.pi) )
titles = ["up", "down", "right", "left"]

with plt.style.context('seaborn-paper', after_reset=True):

    fig, axs = plt.subplots(1, len(rep_thetas), figsize=(15, 3), gridspec_kw={"wspace":0.4})

    for i, thetas in enumerate(rep_thetas):
        x = np.arange(len(thetas))
        axs[i].bar(x - width/2, thetas/(2*np.pi), width, label='Rep {}'.format(i))
        axs[i].hlines((0.2,-0.2), -2., 7., linestyles="dashed")
        axs[i].hlines(0., -2., 7.)
        axs[i].set_xticks(x-0.25)
        axs[i].set_xticklabels(["12","13","14","23","24","34"], fontsize = 15)
        axs[i].set_xlabel("$ij$", fontsize = 15)

        axs[i].set_ylim(-plt_lim,plt_lim)
        axs[i].set_xlim(-.75, 5.75)
        axs[i].set_title(titles[i], fontsize = 15)

        axs[i].tick_params(labelsize=15)

    axs[0].set_ylabel(r"$\theta / 2\pi$", fontsize = 15)

    fig_fname = os.path.join(test_folder, "thetas")

    plt.savefig(fig_fname + ".pdf", bbox_inches='tight')
    plt.savefig(fig_fname + ".png", bbox_inches='tight')
```

**Show predictions made by trained network with disentangled representations**

```python
def plot_state(obs, ax, one_hot=True):
    if one_hot:
        obs=(obs == obs.max())
    ax.pcolormesh(obs.reshape(obs_env.dim, obs_env.dim), edgecolors='gray', linewidth=2)
    ax.set_aspect('equal')
    ax.set_xticks([])
    ax.set_yticks([])

    return ax

n_steps = 10

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

ax1.set_title("Ground truth")
ax2.set_title("Reconstruction")

for i in range(n_steps+1):

    if i==0:
        action = "N\A"
        obs_x = obs_env.reset()
        obs_z = lat_env.reset(encoder(obs_x))
    else:
        action = obs_env.action_space.sample().item()
        obs_x = obs_env.step(action)
        obs_z = lat_env.step(action)

    obs_x_recon = decoder(obs_z)

    fig.suptitle('step {} : last action = {}'.format(i, action), fontsize=16)

    plot_state(obs_x.detach().numpy(),ax1)
    plot_state(obs_x_recon.detach().numpy(),ax2, one_hot=False)

    fig_fname = os.path.join(test_folder, "step_{}".format(i))

    plt.savefig(fig_fname + ".pdf", bbox_inches='tight')
    plt.savefig(fig_fname + ".png", bbox_inches='tight')

    display.clear_output(wait=True)
    display.display(plt.gcf())
    time.sleep(0.5)

display.clear_output(wait=False)
```

```python

```


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

```python
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.linear_model import LinearRegression, Lasso, LogisticRegression
from sklearn.ensemble import GradientBoostingRegressor, GradientBoostingClassifier
from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier
from sklearn.preprocessing import PolynomialFeatures, StandardScaler
from sklearn.pipeline import Pipeline
import scipy.special
```

# NLSYM DATA

```python
# Preprocess data
df = pd.read_csv("data/card.csv")
data_filter = df['educ'].values >= 6
T = df['educ'].values[data_filter]
Z = df['nearc4'].values[data_filter]
y = df['lwage'].values[data_filter]

# Impute missing values with mean, add dummy columns
# I excluded the columns 'weights' as we don't know what it is
X_df = df[['exper', 'expersq']].copy()
X_df['fatheduc'] = df['fatheduc'].fillna(value=df['fatheduc'].mean())
X_df['fatheduc_nan'] = df['fatheduc'].isnull()*1
X_df['motheduc'] = df['motheduc'].fillna(value=df['motheduc'].mean())
X_df['motheduc_nan'] = df['motheduc'].isnull()*1
X_df[['momdad14', 'sinmom14', 'reg661', 'reg662',
        'reg663', 'reg664', 'reg665', 'reg666', 'reg667', 'reg668', 'reg669', 'south66']] = df[['momdad14', 'sinmom14',
        'reg661', 'reg662','reg663', 'reg664', 'reg665', 'reg666', 'reg667', 'reg668', 'reg669', 'south66']]
X_df[['black', 'smsa', 'south', 'smsa66']] = df[['black', 'smsa', 'south', 'smsa66']]
columns_to_scale = ['fatheduc', 'motheduc', 'exper', 'expersq']
scaler = StandardScaler()
X_raw = X_df.values[data_filter]
X_df[columns_to_scale] = scaler.fit_transform(X_df[columns_to_scale])
X = X_df.values[data_filter]
```

    C:\ProgramData\Anaconda3\lib\site-packages\sklearn\preprocessing\data.py:645: DataConversionWarning: Data with input dtype int64, float64 were all converted to float64 by StandardScaler.
      return self.partial_fit(X, y)
    C:\ProgramData\Anaconda3\lib\site-packages\sklearn\base.py:464: DataConversionWarning: Data with input dtype int64, float64 were all converted to float64 by StandardScaler.
      return self.fit(X, **fit_params).transform(X)

```python
X_df.columns
```

    Index(['exper', 'expersq', 'fatheduc', 'fatheduc_nan', 'motheduc',
           'motheduc_nan', 'momdad14', 'sinmom14', 'reg661', 'reg662', 'reg663',
           'reg664', 'reg665', 'reg666', 'reg667', 'reg668', 'reg669', 'south66',
           'black', 'smsa', 'south', 'smsa66'],
          dtype='object')

# 2. Semi-Synthetic Data with Binary Instrument and Continuous Treatment

Data generating process uning real covariates $X$ and instrument $Z$, but synthetic $y$ and $T$ based on the "intent to treat" instrument setting with non-compliance. The instrument corresponds to a fully randomized recommendation of treatment. Then each sample complies with the recommendation to some degree. This probability also depends on an unobserved confounder that has a direct effect on the outcome. Moreover, compliance also depends on the observed feature $X$.

\begin{align}
X \tag{ real features}\\
Z \tag{real instrument}\\
\nu \sim \; & \text{U}[0, 1] \tag{unobserved confounder}\\
C = \; & c\cdot X[i], \; c \;(const)\sim \text{U}[.2, .3] \tag{compliance degree}\\
T = \; & C\cdot Z + g(X) + \nu  \tag{treatment}\\
y \sim \; & \text{Normal}(\mu=\theta(X) \cdot (T + \nu) + f(X),\; \sigma=.1) \tag{outcome}
\end{align}

Moreover:
\begin{align}
\theta(X) = \; & \alpha + \beta \cdot X[i] \tag{CATE}\\
f(X) = \; & X[i] \tag{Nuissance function}\\
g(X) = \; & X[i] \tag{Nuissance function}
\end{align}

```python
def dgp_bin_Z_cont_T(X, Z, hetero_col, true_fn, random_seed=None):
    np.random.seed(random_seed)
    n, d = X.shape
    nu = np.random.uniform(-1, 1, size=(n,))
    c = np.random.uniform(0.2, 0.3)
    C = c * X[:, hetero_col] # Compliers when recomended
    T = C * Z + X[:, hetero_col] + nu # Treatment with compliance
    y = true_fn(X) * (T + nu)  + 0.05*X[:, hetero_col] + np.random.normal(0, .1, size=(n,))
    return y, T

hetero_col = 4 # Mother's education
hetero_col_2 = 7
true_fn = lambda X: 0.1 + 0.05*X[:, hetero_col] - 0.1*X[:, hetero_col_2]

np.random.seed(1237)
y, T = dgp_bin_Z_cont_T(X_raw, Z, hetero_col, true_fn)

plt.figure(figsize=(10, 2))
plt.subplot(1, 2, 1)
plt.hist(T[Z==0])
plt.title("T[Z=0]: Total: {}, Mean(T): {:.2f}".format(T[Z==0].shape[0], np.mean(T[Z==0])))
plt.subplot(1, 2, 2)
plt.hist(T[Z==1])
plt.title("T[Z=1]: Total: {}, Mean(T): {:.2f}".format(T[Z==1].shape[0], np.mean(T[Z==1])))
plt.show()
```

# ANALYSIS

### Defining some hyperparameters

```python
random_seed = 12345 # random seed for each experiment
N_SPLITS = 10 # number of splits for cross-fitting
COV_CLIP = 20/X.shape[0] # covariance clipping in driv
print(COV_CLIP)
```

    0.006686726847208292

### Defining some generic non-parametric regressors and classifiers

```python
from utilities import RegWrapper
from sklearn.model_selection import GridSearchCV
from sklearn.linear_model import LassoCV, LogisticRegressionCV
from xgboost import XGBClassifier, XGBRegressor
from xgb_utilities import XGBWrapper

# Linear models for regression and classification
model = lambda: Pipeline([('poly', PolynomialFeatures(degree=2, interaction_only=True, include_bias=False)),
                          ('lasso', LassoCV(cv=3, n_jobs=-1))])
model_clf = lambda: RegWrapper(Pipeline([('poly', PolynomialFeatures(degree=2, interaction_only=True, include_bias=False)),
                                    ('logistic', LogisticRegressionCV(cv=3, solver='liblinear',
                                                                      scoring='neg_log_loss', n_jobs=-1))]))
```

### Some utility functions

```python
def nuisance_diagnostic(cate, nuisance_model, property_name, property_fn,
                        index_names=None, statistic=np.std, threshold=None):
    std = statistic([property_fn(ns) for ns in cate.fitted_nuisances[nuisance_model]], axis=0)
    if hasattr(std, '__len__'):
        if threshold is None:
            coefs = np.argmax(std).flatten()
        else:
            coefs = np.argwhere(std >= threshold).flatten()
        if index_names is None:
            index_names = np.arange(std.shape[0])
        for high_var in coefs:
            plt.title("{}: {}[{}] Across Folds".format(nuisance_model, property_name, index_names[high_var]))
            plt.plot([property_fn(ns)[high_var] for ns in cate.fitted_nuisances[nuisance_model]])
            plt.xlabel('fold')
            plt.ylabel('property')
            plt.show()
    else:
        plt.title("{}: {} Across Folds".format(nuisance_model, property_name))
        plt.plot([property_fn(ns) for ns in cate.fitted_nuisances[nuisance_model]])
        plt.xlabel('fold')
        plt.ylabel('property')
        plt.show()

```

# ATE via DMLATEIV

```python
from dml_ate_iv import DMLATEIV

np.random.seed(random_seed)

# We need to specify models to be used for each of these residualizations
model_Y_X = lambda: model() # model for E[Y | X]
model_T_X = lambda: model() # model for E[T | X]. We use a regressor since T is continuous
model_Z_X = lambda: model_clf() # model for E[Z | X]. We use a classifier since Z is binary

dmlate = DMLATEIV(model_Y_X(), model_T_X(), model_Z_X(),
                  n_splits=N_SPLITS, # n_splits determines the number of splits to be used for cross-fitting.
                  binary_instrument=True, # a flag whether to stratify cross-fitting by instrument
                  binary_treatment=False # a flag whether to stratify cross-fitting by treatment
                 )
```

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

```python
# We fit DMLATEIV with these models
dmlate.fit(y, T, X, Z)
```

    <dml_ate_iv.DMLATEIV at 0x22503e76358>

```python
# We call effect() to get the ATE
ta_effect = dmlate.effect()
```

```python
# Comparison with true ATE
print("ATE Estimate: {:.3f}".format(ta_effect))
print("True ATE: {:.3f}".format(np.mean(true_fn(X_raw))))
# CATE MSE
print("CATE MSE: {:.2f}".format(np.mean((true_fn(X_raw) - ta_effect)**2)))
```

    ATE Estimate: 0.651
    True ATE: 0.609
    CATE MSE: 0.03

```python
# We can call normal_effect_interval to get confidence intervals based
# based on the asympotic normal approximation
ta_effect = dmlate.normal_effect_interval(lower=2.5, upper=97.5)
# Comparison with true ATE
print("ATE Estimate Interval: ({:.3f}, {:.3f})".format(ta_effect[0], ta_effect[1]))
print("True ATE: {:.3f}".format(np.mean(true_fn(X_raw))))
```

    ATE Estimate Interval: (0.607, 0.696)
    True ATE: 0.609

```python
def get_dmlateiv_coverage(true_effect, iteration):
    y, T = dgp_bin_Z_cont_T(X_raw, Z, hetero_col, true_fn, random_seed=iteration)
    dmlate = DMLATEIV(model_Y_X(), model_T_X(), model_Z_X(),
                  n_splits=N_SPLITS, # n_splits determines the number of splits to be used for cross-fitting.
                  binary_instrument=True, # a flag whether to stratify cross-fitting by instrument
                  binary_treatment=False # a flag whether to stratify cross-fitting by treatment
                 )
    dmlate.fit(y, T, X, Z)
    left, right = dmlate.normal_effect_interval(lower=2.5, upper=97.5)
    if true_effect >= left and true_effect <= right:
        return 1
    return 0

from joblib import Parallel, delayed
n_experiments=100
true_ate = np.mean(true_fn(X_raw))
if True:
    contains_truth = np.array(Parallel(n_jobs=-1, verbose=3)(
            delayed(get_dmlateiv_coverage)(true_ate, it) for it in range(n_experiments)))
    print("Coverage: {}".format(contains_truth.mean()))
```

    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 12 concurrent workers.
    [Parallel(n_jobs=-1)]: Done   8 tasks      | elapsed:  1.4min

    Coverage: 0.52

    [Parallel(n_jobs=-1)]: Done 100 out of 100 | elapsed:  8.4min finished

# ATE and CATE via DMLIV

```python
from dml_iv import DMLIV
from utilities import SelectiveLasso, SeparateModel
from sklearn.linear_model import LassoCV, LogisticRegressionCV
from econml.utilities import hstack

np.random.seed(random_seed)

# For DMLIV we also need a model for E[T | X, Z]. To allow for heterogeneity in the compliance, i.e.
# T = beta(X)*Z + gamma(X)
# we train a separate model for Z=1 and Z=0. The model for Z=1 learns the
# quantity beta(X) + gamma(X) and the model for Z=0 learns gamma(X).
model_T_XZ = lambda: SeparateModel(model(), model())

# We now specify the features to be used for heterogeneity. We will fit a CATE model of the form
#      theta(X) = <theta, phi(X)>
# for some set of features phi(X). The featurizer needs to support fit_transform, that takes
# X and returns phi(X). We need to include a bias if we also want a constant term.
dmliv_featurizer = lambda: PolynomialFeatures(degree=1, include_bias=True)

# Then we need to specify a model to be used for fitting the parameters theta in the linear form.
# This model will minimize the square loss:
#        (Y - E[Y|X] - <theta, phi(X)> * (E[T|X,Z] - E[T|X]))**2
#dmliv_model_effect = lambda: LinearRegression(fit_intercept=False)

# Potentially with some regularization on theta. Here we use an ell_1 penalty on theta
# If we also have a prior that there is no effect heterogeneity we can use a selective lasso
# that does not penalize the constant term in the CATE model
dmliv_model_effect = lambda: SelectiveLasso(np.arange(1, X.shape[1]+1), LassoCV(cv=5, fit_intercept=False))

# We initialize DMLIV with all these models and call fit
cate = DMLIV(model_Y_X(), model_T_X(), model_T_XZ(),
             dmliv_model_effect(), dmliv_featurizer(),
             n_splits=N_SPLITS, # number of splits to use for cross-fitting
             binary_instrument=True, # a flag whether to stratify cross-fitting by instrument
             binary_treatment=False # a flag whether to stratify cross-fitting by treatment
            )
```

```python
cate.fit(y, T, X, Z)
```

    <dml_iv.DMLIV at 0x225047062e8>

```python
# To get the CATE at every X we call effect(X)
dml_effect = cate.effect(X)
plt.hist(dml_effect, label='est')
plt.hist(true_fn(X_raw), alpha=.2, label='true')
plt.legend()
plt.show()
```

```python
# To get the parameter theta we call coef_. The first entry is the intercept of the CATE model
print(cate.coef_)
```

    [ 0.62317714  0.          0.          0.         -0.          0.10377389
     -0.          0.         -0.         -0.          0.          0.
     -0.         -0.         -0.         -0.          0.          0.
     -0.          0.          0.         -0.          0.        ]

```python
# We can average the CATE to get an ATE
print("ATE Estimate: {:.3f}".format(np.mean(dml_effect)))
print("True ATE: {:.3f}".format(np.mean(true_fn(X_raw))))
```

    ATE Estimate: 0.624
    True ATE: 0.609

```python
# We can also see how it compares to the true CATE at each target point and calculate MSE
plt.title("DMLIV CATE as Function of {}. MSE={:.3f}".format(X_df.columns[4], np.mean((dml_effect-true_fn(X_raw))**2)))
plt.scatter(X[:, 4], dml_effect, label='est')
plt.scatter(X[:, 4], true_fn(X_raw), label='true', alpha=.2)
plt.legend()
plt.show()
```

# ATE and Projected CATE via DRIV

```python
from dml_iv import DMLIV
from dr_iv import DRIV, ProjectedDRIV
from utilities import SubsetWrapper, StatsModelLinearRegression, ConstantModel
from sklearn.dummy import DummyRegressor

np.random.seed(random_seed)

# For DRIV we need a model for predicting E[T*Z | X]. We use a classifier
model_TZ_X = lambda: model()

# We also need a model for the final regression that will fit the function theta(X)
# If we want to fit an ATE, we simply fit a constant functin theta(X) = theta
# We can do this with a pipeline where the preprocessing step only creates a bias column
# and the regression step fits a linear regression with no intercept.
# To get normal confidence intervals easily we can use a statsmodels linear regression
# wrapped in an sklearn interface
const_driv_model_effect = lambda: ConstantModel()

# As in OrthoDMLIV we need a perliminary estimator of the CATE.
# We use a DMLIV estimator with no cross-fitting (n_splits=1)
dmliv_prel_model_effect = DMLIV(model_Y_X(), model_T_X(), model_T_XZ(),
                                dmliv_model_effect(), dmliv_featurizer(),
                                n_splits=1, binary_instrument=True, binary_treatment=False)

const_dr_cate = DRIV(model_Y_X(), model_T_X(), model_Z_X(), # same as in DMLATEIV
                        dmliv_prel_model_effect, # preliminary model for CATE, must support fit(y, T, X, Z) and effect(X)
                        model_TZ_X(), # model for E[T * Z | X]
                        const_driv_model_effect(), # model for final stage of fitting theta(X)
                        cov_clip=COV_CLIP, # covariance clipping to avoid large values in final regression from weak instruments
                        n_splits=N_SPLITS, # number of splits to use for cross-fitting
                        binary_instrument=True, # a flag whether to stratify cross-fitting by instrument
                        binary_treatment=False, # a flag whether to stratify cross-fitting by treatment
                        opt_reweighted=False # whether to optimally re-weight samples. Valid only for flexible final model
                       )
```

```python
const_dr_cate.fit(y, T, X, Z, store_final=True)
```

    C:\ProgramData\Anaconda3\lib\site-packages\sklearn\linear_model\coordinate_descent.py:492: ConvergenceWarning: Objective did not converge. You might want to increase the number of iterations. Fitting data with very small alpha may cause precision problems.
      ConvergenceWarning)
    C:\ProgramData\Anaconda3\lib\site-packages\sklearn\linear_model\coordinate_descent.py:492: ConvergenceWarning: Objective did not converge. You might want to increase the number of iterations. Fitting data with very small alpha may cause precision problems.
      ConvergenceWarning)
    C:\ProgramData\Anaconda3\lib\site-packages\sklearn\linear_model\coordinate_descent.py:492: ConvergenceWarning: Objective did not converge. You might want to increase the number of iterations. Fitting data with very small alpha may cause precision problems.
      ConvergenceWarning)
    C:\ProgramData\Anaconda3\lib\site-packages\sklearn\linear_model\coordinate_descent.py:492: ConvergenceWarning: Objective did not converge. You might want to increase the number of iterations. Fitting data with very small alpha may cause precision problems.
      ConvergenceWarning)
    C:\ProgramData\Anaconda3\lib\site-packages\sklearn\linear_model\coordinate_descent.py:492: ConvergenceWarning: Objective did not converge. You might want to increase the number of iterations. Fitting data with very small alpha may cause precision problems.
      ConvergenceWarning)

    <dr_iv.DRIV at 0x225048e3668>

```python
# To get the statsmodel summary we look at the effect_model, which is the pipeline, we then look
# at the reg step of the pipeline which is the statsmodel wrapper and then we look
# at the model attribute of the statsmodel wrapper and print the summary()
const_dr_cate.effect_model.summary()
```

    C:\ProgramData\Anaconda3\lib\site-packages\statsmodels\regression\linear_model.py:1554: RuntimeWarning: invalid value encountered in double_scalars
      return self.ess/self.df_model
    C:\ProgramData\Anaconda3\lib\site-packages\scipy\stats\_distn_infrastructure.py:877: RuntimeWarning: invalid value encountered in greater
      return (self.a < x) & (x < self.b)
    C:\ProgramData\Anaconda3\lib\site-packages\scipy\stats\_distn_infrastructure.py:877: RuntimeWarning: invalid value encountered in less
      return (self.a < x) & (x < self.b)
    C:\ProgramData\Anaconda3\lib\site-packages\scipy\stats\_distn_infrastructure.py:1831: RuntimeWarning: invalid value encountered in less_equal
      cond2 = cond0 & (x <= self.a)

OLS Regression Results

  Dep. Variable:	            y	          R-squared:         	    0.000

  Model:	                   OLS	         Adj. R-squared:    	    0.000

  Method:	             Least Squares	    F-statistic:       	      nan

  Date:	             Sat, 01 Jun 2019	   Prob (F-statistic):	     nan

  Time:	                 15:27:35	       Log-Likelihood:    	  -7831.2

  No. Observations:	        2991	        AIC:               	 1.566e+04

  Df Residuals:	            2990	        BIC:               	 1.567e+04

  Df Model:	                   0

  Covariance Type:	      nonrobust

    	       coef	     std err	      t	      P>|t|	  [0.025	    0.975]

  const	     0.5458	     0.061	     8.995	  0.000	     0.427	     0.665

  Omnibus:	       4190.040	   Durbin-Watson:     	     1.964

  Prob(Omnibus):	   0.000	    Jarque-Bera (JB):  	 8408463.090

  Skew:	           -7.266	    Prob(JB):          	      0.00

  Kurtosis:	       262.343	   Cond. No.          	      1.00

Warnings:
[1] Standard Errors assume that the covariance matrix of the errors is correctly specified.

```python
def get_driv_coverage(true_effect, iteration):
    y, T = dgp_bin_Z_cont_T(X_raw, Z, hetero_col, true_fn, random_seed=iteration)
    dmliv_prel_model_effect = DMLIV(model_Y_X(), model_T_X(), model_T_XZ(),
                                dmliv_model_effect(), dmliv_featurizer(),
                                n_splits=1, binary_instrument=True, binary_treatment=True)
    const_dr_cate = DRIV(model_Y_X(), model_T_X(), model_Z_X(), # same as in DMLATEIV
                            dmliv_prel_model_effect, # preliminary model for CATE, must support fit(y, T, X, Z) and effect(X)
                            model_TZ_X(), # model for E[T * Z | X]
                            const_driv_model_effect(), # model for final stage of fitting theta(X)
                            cov_clip=COV_CLIP, # covariance clipping to avoid large values in final regression from weak instruments
                            n_splits=N_SPLITS, # number of splits to use for cross-fitting
                            binary_instrument=True, # a flag whether to stratify cross-fitting by instrument
                            binary_treatment=False # a flag whether to stratify cross-fitting by treatment
                           )
    const_dr_cate.fit(y, T, X, Z, store_final=True)
    left, right = const_dr_cate.effect_model.est.model.conf_int(alpha=0.05)[0]
    if true_effect >= left and true_effect <= right:
        return 1
    return 0
```

```python
from joblib import Parallel, delayed
n_experiments=100
true_ate = np.mean(true_fn(X_raw))
if True:
    contains_truth = np.array(Parallel(n_jobs=-1, verbose=3)(
            delayed(get_driv_coverage)(true_ate, it) for it in range(n_experiments)))
    print("Coverage: {}".format(contains_truth.mean()))
```

    [Parallel(n_jobs=-1)]: Using backend LokyBackend with 12 concurrent workers.
    [Parallel(n_jobs=-1)]: Done   8 tasks      | elapsed:  1.3min

    Coverage: 0.98

    [Parallel(n_jobs=-1)]: Done 100 out of 100 | elapsed: 11.3min finished

## Projecting CATE to a pre-chosen subset of variables in final model

```python
from dml_iv import DMLIV
from dr_iv import DRIV, ProjectedDRIV
from utilities import SubsetWrapper

np.random.seed(random_seed)

# We could also fit a projection on a subset of the features by using the
# subset wrapper from our utilities.

# Example: including everything for expository purposes, but any array-like of indices would work
subset_names = set(['motheduc', 'sinmom14'])
# list of indices of features X to use in the final model
feature_inds = np.argwhere([(x in subset_names) for x in X_df.columns.values]).flatten() #[0] #np.arange(X.shape[1])
print(feature_inds)
# Because we are projecting to a low dimensional model space, we can
# do valid inference and we can use statsmodel linear regression to get all
# the hypothesis testing capability
proj_driv_model_effect = lambda: SubsetWrapper(StatsModelLinearRegression(),
                                          feature_inds # list of indices of features X to use in the final model
                                         )
```

    [4 7]

```python
X_df.columns
```

    Index(['exper', 'expersq', 'fatheduc', 'fatheduc_nan', 'motheduc',
           'motheduc_nan', 'momdad14', 'sinmom14', 'reg661', 'reg662', 'reg663',
           'reg664', 'reg665', 'reg666', 'reg667', 'reg668', 'reg669', 'south66',
           'black', 'smsa', 'south', 'smsa66'],
          dtype='object')

```python
proj_dr_cate = const_dr_cate.refit_final(proj_driv_model_effect())
```

```python
# To get the CATE at every X we call effect(X[:, feature_inds])
proj_dr_effect = proj_dr_cate.effect(X[:, feature_inds])
```

```python
# To get the statsmodel summary we look at the effect_model, which is
# an instance of SubsetWrapper, we look at the model of the SubsetWrapper which is
# and instance of the pipeline, we then look at the reg step of the pipeline which is the statsmodel wrapper and
# call summary() of the wrapper (most prob there is a better API for this, but we can go with this for now :)
proj_dr_cate.effect_model.summary(alpha=.05, xname=['const']+list(X_df.columns[feature_inds]))
```

OLS Regression Results

  Dep. Variable:	            y	          R-squared:         	    0.002

  Model:	                   OLS	         Adj. R-squared:    	    0.001

  Method:	             Least Squares	    F-statistic:       	    2.637

  Date:	             Sat, 01 Jun 2019	   Prob (F-statistic):	  0.0717

  Time:	                 15:38:55	       Log-Likelihood:    	  -7828.6

  No. Observations:	        2991	        AIC:               	 1.566e+04

  Df Residuals:	            2988	        BIC:               	 1.568e+04

  Df Model:	                   2

  Covariance Type:	      nonrobust

      	        coef	     std err	      t	      P>|t|	  [0.025	    0.975]

  const	        0.5561	     0.064	     8.684	  0.000	     0.431	     0.682

  motheduc	     0.1310	     0.062	     2.121	  0.034	     0.010	     0.252

  sinmom14	    -0.1149	     0.204	    -0.563	  0.573	    -0.515	     0.285

  Omnibus:	       4157.081	   Durbin-Watson:     	     1.966

  Prob(Omnibus):	   0.000	    Jarque-Bera (JB):  	 8324885.444

  Skew:	           -7.146	    Prob(JB):          	      0.00

  Kurtosis:	       261.061	   Cond. No.          	      3.41

Warnings:
[1] Standard Errors assume that the covariance matrix of the errors is correctly specified.

```python
# We can access the coefficient by looking at the coefficient attribute of the final step of the pipeline
print("Estimated Params: {}, {}".format(proj_dr_cate.intercept_, proj_dr_cate.coef_))
# True coefficients of projection
print("True Params: {}".format(
        LinearRegression(fit_intercept=False).fit(PolynomialFeatures(degree=1,
                                                                     include_bias=True).fit_transform(X[:, feature_inds]),
                                                  true_fn(X_raw)).coef_))
```

    Estimated Params: 0.5560648374345198, [ 0.13103506 -0.11489838]
    True Params: [ 0.61740685  0.14934234 -0.1       ]

```python
# We can also evaluate coverage and create prediction intervals using statsmodels attributes
from statsmodels.sandbox.regression.predstd import wls_prediction_std
res = proj_dr_cate.effect_model.model
predictions = res.get_prediction(PolynomialFeatures(degree=1, include_bias=True).fit_transform(X[:, feature_inds]))
frame = predictions.summary_frame(alpha=0.05)
pred = frame['mean']
iv_l = frame['mean_ci_lower']
iv_u = frame['mean_ci_upper']

# This is the true CATE functions
theta_true = true_fn(X_raw)
# This is the true projection of the CATE function on the subspace of linear functions of the
# subset of the features used in the projection
true_proj = LinearRegression().fit(X[:, feature_inds], theta_true).predict(X[:, feature_inds])

# Are we covering the true projection
covered = (true_proj <= iv_u) & (true_proj >= iv_l)
print("Coverage of True Projection: {:.2f}".format(np.mean(covered)))

fig, ax = plt.subplots(figsize=(8,6))

order = np.argsort(X[:, feature_inds[0]])
ax.plot(X[order, feature_inds[0]], iv_u[order], 'r--')
ax.plot(X[order, feature_inds[0]], iv_l[order], 'r--')
ax.plot(X[order, feature_inds[0]], pred[order], 'g--.', label="pred")
ax.plot(X[order, feature_inds[0]], theta_true[order], 'b-', label="True", alpha=.3)
ax.plot(X[order, feature_inds[0]], true_proj[order], 'b-', label="TrueProj", alpha=.3)
ax.legend(loc='best')
plt.show()
```

# CATE via Re-Weighted DRIV: DRIV-RW

## Lasso CATE

```python
from dml_iv import DMLIV
from dr_iv import DRIV, ProjectedDRIV
from utilities import SubsetWrapper, StatsModelLinearRegression, ConstantModel, WeightWrapper
from sklearn.dummy import DummyRegressor

np.random.seed(random_seed)

# For DRIV we need a model for predicting E[T*Z | X]. We use a classifier
model_TZ_X = lambda: model()

# We also need a model for the final regression that will fit the function theta(X)
# This model now needs to accept sample weights at fit time.
const_driv_model_effect = lambda: WeightWrapper(Pipeline([('bias', PolynomialFeatures(degree=1, include_bias=True)),
                                                          ('reg', SelectiveLasso(np.arange(1, X.shape[1]+1),
                                                                                LassoCV(cv=5, fit_intercept=False)))]))

# As in OrthoDMLIV we need a perliminary estimator of the CATE.
# We use a DMLIV estimator with no cross-fitting (n_splits=1)
dmliv_prel_model_effect = DMLIV(model_Y_X(), model_T_X(), model_T_XZ(),
                                dmliv_model_effect(), dmliv_featurizer(),
                                n_splits=1, binary_instrument=True, binary_treatment=False)

const_dr_cate = DRIV(model_Y_X(), model_T_X(), model_Z_X(), # same as in DMLATEIV
                        dmliv_prel_model_effect, # preliminary model for CATE, must support fit(y, T, X, Z) and effect(X)
                        model_TZ_X(), # model for E[T * Z | X]
                        const_driv_model_effect(), # model for final stage of fitting theta(X)
                        cov_clip=COV_CLIP, # covariance clipping to avoid large values in final regression from weak instruments
                        n_splits=N_SPLITS, # number of splits to use for cross-fitting
                        binary_instrument=True, # a flag whether to stratify cross-fitting by instrument
                        binary_treatment=False, # a flag whether to stratify cross-fitting by treatment
                        opt_reweighted=True # whether to optimally re-weight samples. Valid only for flexible final model
                       )
```

```python
const_dr_cate.fit(y, T, X, Z, store_final=True)
```

    C:\ProgramData\Anaconda3\lib\site-packages\sklearn\linear_model\coordinate_descent.py:492: ConvergenceWarning: Objective did not converge. You might want to increase the number of iterations. Fitting data with very small alpha may cause precision problems.
      ConvergenceWarning)
    C:\ProgramData\Anaconda3\lib\site-packages\sklearn\linear_model\coordinate_descent.py:492: ConvergenceWarning: Objective did not converge. You might want to increase the number of iterations. Fitting data with very small alpha may cause precision problems.
      ConvergenceWarning)
    C:\ProgramData\Anaconda3\lib\site-packages\sklearn\linear_model\coordinate_descent.py:492: ConvergenceWarning: Objective did not converge. You might want to increase the number of iterations. Fitting data with very small alpha may cause precision problems.
      ConvergenceWarning)
    C:\ProgramData\Anaconda3\lib\site-packages\sklearn\linear_model\coordinate_descent.py:492: ConvergenceWarning: Objective did not converge. You might want to increase the number of iterations. Fitting data with very small alpha may cause precision problems.
      ConvergenceWarning)
    C:\ProgramData\Anaconda3\lib\site-packages\sklearn\linear_model\coordinate_descent.py:492: ConvergenceWarning: Objective did not converge. You might want to increase the number of iterations. Fitting data with very small alpha may cause precision problems.
      ConvergenceWarning)

    <dr_iv.DRIV at 0x225048e3278>

```python
# We can average the CATE to get an ATE
dr_effect = const_dr_cate.effect(X)
print("ATE Estimate: {:.3f}".format(np.mean(dr_effect)))
print("True ATE: {:.3f}".format(np.mean(true_fn(X_raw))))
```

    ATE Estimate: 0.612
    True ATE: 0.609

```python
# We can also see how it compares to the true CATE at each target point and calculate MSE
plt.title("DMLIV CATE as Function of {}: MSE={:.3f}".format(X_df.columns[4], np.mean((dr_effect-true_fn(X_raw))**2)))
plt.scatter(X[:, 4], dr_effect, label='est')
plt.scatter(X[:, 4], true_fn(X_raw), label='true', alpha=.2)
plt.legend()
plt.show()
```

## Random Forest CATE

```python
from dml_iv import DMLIV
from dr_iv import DRIV, ProjectedDRIV
from utilities import SubsetWrapper
from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor

np.random.seed(random_seed)

rf_driv_model_effect = lambda: RandomForestRegressor(n_estimators=5000, max_depth=3, min_impurity_decrease=0.00001,
                                                     min_samples_leaf=100, bootstrap=True)
```

```python
rf_dr_cate = const_dr_cate.refit_final(rf_driv_model_effect())
```

```python
rf_dr_effect = rf_dr_cate.effect(X)
```

```python
print("ATE Estimate: {:.2f}".format(np.mean(rf_dr_effect)))
print("True ATE: {:.2f}".format(np.mean(true_fn(X_raw))))
```

    ATE Estimate: 0.63
    True ATE: 0.61

```python
plt.title("DRIV CATE: MSE {:.2}".format(np.mean((true_fn(X_raw) - rf_dr_effect)**2)))
plt.scatter(X[:, 4], rf_dr_effect, label='est')
plt.scatter(X[:, 4], true_fn(X_raw), label='true', alpha=.2)
plt.legend()
plt.show()
```

```python
import shap
import pandas as pd

Xdf = pd.DataFrame(X, columns=X_df.columns)
# explain the model's predictions using SHAP values
explainer = shap.TreeExplainer(rf_dr_cate.effect_model)
shap_values = explainer.shap_values(Xdf)

# visualize the first prediction's explanation (use matplotlib=True to avoid Javascript)
shap.force_plot(explainer.expected_value, shap_values[0,:], Xdf.iloc[0,:], matplotlib=True)
```

```python
shap.summary_plot(shap_values, Xdf)
```

```python
shap.summary_plot(shap_values, Xdf, plot_type='bar')
```

```python

```


## Teoría de perturbaciones

Consiste en resolver un sistema perturbado(se conoce la solución al no perturbado), y donde el interés es conocer la contribución de la parte perturbada $H'$ al nuevo sistema total.

$$ H = H^{0} + H'$$

Para sistemas no degenerados, la corrección a la energía a primer orden se calcula como

$$ E_{n}^{(1)} = \int\psi_{n}^{(0)*} H' \psi_{n}^{(0)}d\tau$$

Y la corrección a la función de onda, también a primer orden, se obtiene como:

 $$ \psi_{n}^{(1)} = \sum_{m\neq n} \frac{\langle\psi_{m}^{(0)} \left| H' \right| \psi_{n}^{(0)} \rangle}{E_{n}^{(0)} - E_{m}^{(0)}} \psi_{m}^{(0)}$$

** Tarea 1 : Programar ecuación para corrección de energía a primer orden si conoces $H^{0}$ y sus soluciones. **

**Tarea 2: Programar ecuación para correcciones a las funciones de onda a primer orden si conoces $H^{0}$ y sus soluciones. **

**Tarea 3: Investigue las soluciones a segundo orden y también programe las soluciones. **

La corrección de segundo orden a la energía se obtiene de la siguiente forma:

 $$ E_{n}^{(2)} = \sum_{m\neq n} \frac{\left| \langle\psi_{m}^{(0)} \left| H' \right| \psi_{n}^{(0)} \rangle\right|^{2}}{E_{n}^{(0)} - E_{m}^{(0)}}$$

```python
### Solución

from sympy import *
from sympy import init_printing; init_printing(use_latex = 'mathjax')
from sympy.physics.qho_1d import E_n, psi_n
from scipy import constants
import numpy as np
var('x a l psi') #Variables a utilizar
var('m n n1 n2 l1 l2 hbar omega c d q r r1 r2 k', positive=True, real=True)

def get_corrections(PType,Pert,NI,Mass,L,AngFreq, Range, PertLim1, PertLim2):
    E_correction=0
    Psi_correction_sym=0
    E_correction2=0
    Psin=zeros(1,Range[1]-Range[0]+1)
    En=zeros(1,Range[1]-Range[0]+1)

    #Problem type original energies and wavefuncions
    if PType == "Harmonic oscilator":
        Psi=psi_n(n,x,m,omega)
        Eb=E_n(n,omega)
        for i in range (0,Range[1]-Range[0]+1):
            Psin[i]=psi_n(Range[0]+i,x,eval(Mass),eval(AngFreq))
            En[i]=E_n(Range[0]+i,eval(AngFreq))

    elif PType == 'Particle in a one-dimensional box':
        Psi=sqrt(2/a)*sin(n*pi*x/a)
        Eb=hbar**2*pi**2*n**2/(2*m*a**2)
        for i in range (0,Range[1]-Range[0]+1):
            Psin[i]=Psi.subs(n,Range[0]+i).subs(a,eval(L))
            En[i]=Eb.subs(n,Range[0]+i).subs(m,eval(Mass)).subs(a,eval(L))

    #Solutions
    if (NI == 0 or Range[0]==0) and PType == 'Particle in a one-dimensional box':
        print("Error! n cannot be 0 for the particle in a box and helium atom problem")
        E_correction1=sympify('ERROR')
        E_correction2=sympify('ERROR')
        Psi_correction_sym=sympify('ERROR')

    else:
        PsiIntegrals=integrate(Psin[NI-Range[0]]*eval(Pert)*Psin, (x,eval(PertLim1),eval(PertLim2)))
        E_correction1=PsiIntegrals[NI-Range[0]]
        for i in range (0,Range[1]-Range[0]+1):
            if i!= NI-Range[0]:
                Psi_correction_sym=Psi_correction_sym+Psin[i]*PsiIntegrals[i]/(En[NI-Range[0]]-En[i])
                E_correction2=E_correction2+PsiIntegrals[i]**2/(En[NI-Range[0]]-En[i])

    return Psi, Eb, Psi_correction_sym, E_correction1, E_correction2

```

```python
import ipywidgets as wid
from IPython.display import display
from ipywidgets import Layout, Box, Text, Dropdown, Label, IntSlider, IntRangeSlider, HTML

form_item_layout = Layout(
    display='flex',
    flex_flow='row',
    justify_content='space-between'
)
PType=Dropdown(options=['Particle in a one-dimensional box', 'Harmonic oscilator'])
Pert=Text(value='0')
NI=IntSlider(value=0, min=0, max=9, step=1)
NI2=IntSlider(value=0, min=0, max=9, step=1)
lI1=IntSlider(value=0, min=0, max=9, step=1)
lI2=IntSlider(value=0, min=0, max=9, step=1)
Ran=IntRangeSlider(value=[0,9],min=0,max=9)
Mass=Text(value='m')
L=Text(value='a')
AngFreq=Text(value='omega')
Label1=HTML(value="Select <b>n</b> for the energy and wavefunction corrections, for particle in a box corrections the"
            " lowest <b>n</b> is <b>1</b>")
Label2=HTML(value="Select the function range for the wavefunction correction sum, for particle in a box the lowest"
            " <b>m</b> is <b>1</b>. Remember <b>n</b> must be contained in the selected <b>m</b> range")
Label3=HTML(value="The following indexes only work for the Helium Atom problem, <b>n</b> index above is used as <b>n1</b>")
Label4=HTML(value="Select <b>l</b> for the energy and wavefunction corrections to the helium atom, remember that <b>l</b>"
            " must be smaller than <b>n</b>")
PertLim1=Text(value='0')
PertLim2=Text(value='a')

form_items1 = [
    Box([Label(value='Problem'),
         PType], layout=form_item_layout),
    Box([Label(value='Perturbation'),
         Pert], layout=form_item_layout),
    Box([Label(value='P.Min.Lim'),
         PertLim1], layout=form_item_layout),
    Box([Label(value='P.May.Lim'),
         PertLim2], layout=form_item_layout),
    Box([Label(value='P. Mass'),
         Mass], layout=form_item_layout),
    Box([Label(value='Box size'),
         L], layout=form_item_layout),
    Box([Label(value='Ang. Freq.'),
         AngFreq], layout=form_item_layout),
    Box([Label(value=''),
         Label1], layout=form_item_layout),
    Box([Label(value='n Index'),
         NI], layout=form_item_layout),
    Box([Label(value=''),
         Label2], layout=form_item_layout),
    Box([Label(value='m Range'),
         Ran], layout=form_item_layout)]

form1 = Box(form_items1, layout=Layout(
    display='flex',
    flex_flow='column',
    border='solid 2px',
    align_items='stretch',
    width='40%'
))
form1

```

    A Jupyter Widget

```python
Psi,E,Psicorrection,Ecorrection,Ecorrection2=get_corrections(PType.value, Pert.value, NI.value, Mass.value, L.value,
                                                             AngFreq.value, Ran.value, PertLim1.value, PertLim2.value)
Correction=Matrix([[sympify('psi_n(0)'),Psi],
                   [sympify('E_n(0)'),E],
                   [sympify('psi_%d(1)'  %(NI.value)),factor(Psicorrection)],
                   [sympify('E_%d(1)'  %(NI.value)),factor(Ecorrection)],
                   [sympify('E_%d(2)'  %(NI.value)),factor(Ecorrection2)]])
Correction
```

$$\left[\begin{matrix}\psi_{n}{\left (0 \right )} & \frac{2^{- \frac{n}{2}} \sqrt[4]{m} \sqrt[4]{\omega}}{\sqrt[4]{\hbar} \sqrt[4]{\pi}} \sqrt{\frac{1}{n!}} e^{- \frac{m \omega x^{2}}{2 \hbar}} H_{n}\left(\frac{\sqrt{m} x}{\sqrt{\hbar}} \sqrt{\omega}\right)\\\operatorname{E_{n}}{\left (0 \right )} & \hbar \omega \left(n + \frac{1}{2}\right)\\\psi_{0}{\left (1 \right )} & - \frac{e^{- \frac{m \omega x^{2}}{2 \hbar}}}{48 \hbar^{\frac{5}{4}} \sqrt[4]{\pi} m^{\frac{7}{4}} \omega^{\frac{11}{4}}} \left(16 c m^{2} \omega^{2} x^{3} + 48 \hbar c m \omega x + 12 d m^{2} \omega^{2} x^{4} + 36 \hbar d m \omega x^{2} - 27 \hbar^{2} d\right)\\\operatorname{E_{0}}{\left (1 \right )} & \frac{3 \hbar^{2} d}{4 m^{2} \omega^{2}}\\\operatorname{E_{0}}{\left (2 \right )} & - \frac{\hbar^{2}}{8 m^{4} \omega^{5}} \left(11 c^{2} m \omega + 21 \hbar d^{2}\right)\end{matrix}\right]$$

**Tarea 4.  Resolver el átomo de helio aplicando los programas anteriores.**

Para resolver el átomo de helio se toma el átomo de hidrógeno como base en donde se considera que el Hamiltoniano original es:

$$ H_{0} = H_{1} + H_{2}$$

Y la función sin perturbar es:

$$ \psi_{n_1 l_1 n_2 l_2} = \psi_{n_1 l_1}(r_1)\psi_{n_2 l_2}(r_2) $$

Para lo cual la energía se obtiene de la siguiente forma:

$$ E_{n_1  n_2} = -\alpha \left( \frac{1}{n_1}+\frac{1}{n_2} \right) $$

Ahora agregaremos la perturbación que representa el término repulsivo:

$$ H'= \frac{e^2}{\left|r_1-r_2\right|}$$

```python
from sympy.physics.hydrogen import E_nl, R_nl
def Helium(N1,N2,L1,L2):

    Eb=E_nl(N1,1)+E_nl(N2,1)
    Psi1=R_nl(N1, L1, r1, Z=1)
    Psi2=R_nl(N2, L2, r2, Z=1)
    Psi=Psi1*Psi2
    E_correction1=integrate(integrate(r1**2*r2**2*conjugate(Psi1)*conjugate(Psi2)*q**2*Psi1*Psi2/abs(r1-r2), (r1,0,oo)), (r2,0,oo))
    E_correctionR1=q**2*integrate(r2**2*conjugate(Psi2)*Psi2*(integrate(r1**2*conjugate(Psi1)*Psi1/r2, (r1,0,r2))+integrate(r1**2*conjugate(Psi1)*Psi1/r1, (r1,r2,oo))), (r2,0,oo))

    return Psi, Eb, E_correction1, E_correctionR1
```

```python
NI1=IntSlider(value=1, min=0, max=9, step=1)
NI2=IntSlider(value=1, min=0, max=9, step=1)
LI1=IntSlider(value=0, min=0, max=9, step=1)
LI2=IntSlider(value=0, min=0, max=9, step=1)
Label3=HTML(value="Select <b>n1</b>, <b>n2</b>, <b>l1</b>, <b>l2</b> indexes to be used for the first order"
            " energy correction to the Helium Atom")

form_items2 = [
    Box([Label(value=''),
         Label3], layout=form_item_layout),
    Box([Label(value='n1 Index'),
         NI1], layout=form_item_layout),
    Box([Label(value='n2 Index'),
         NI2], layout=form_item_layout),
    Box([Label(value='l1 Index'),
         LI1], layout=form_item_layout),
    Box([Label(value='l2 Index'),
         LI2], layout=form_item_layout)]

form2 = Box(form_items2, layout=Layout(
    display='flex',
    flex_flow='column',
    border='solid 2px',
    align_items='stretch',
    width='40%'
))
form2

```

    A Jupyter Widget

```python
HPsi,HE,HEcorrection, HER=Helium(NI1.value, NI2.value, LI1.value, LI2.value)

HCorrection=Matrix([[sympify('psi_(0)'), HPsi],
                    [sympify('E_(0)'), HE],
                    [sympify('E_(1)'), factor(HEcorrection)],
                    [sympify('E_(1)'), HER]])
HCorrection
```

$$\left[\begin{matrix}\psi_{}{\left (0 \right )} & 4 e^{- r_{1}} e^{- r_{2}}\\\operatorname{E_{}}{\left (0 \right )} & -1\\\operatorname{E_{}}{\left (1 \right )} & 16 q^{2} \int_{0}^{\infty} r_{2}^{2} e^{- 2 r_{2}} \int_{0}^{\infty} \frac{r_{1}^{2} e^{- 2 r_{1}}}{\left|{r_{1} - r_{2}}\right|}\, dr_{1}\, dr_{2}\\\operatorname{E_{}}{\left (1 \right )} & \frac{5 q^{2}}{8}\end{matrix}\right]$$

**Tarea 5: Método variacional-perturbativo. **

Este método nos permite estimar de forma precisa $E^{(2)}$ y correcciones perturbativas de la energía de órdenes más elevados para el estado fundamental del sistema, sin evaluar sumas infinitas. Ver ecuación 9.38 del libro.

$$ \langle u \left| H^0 - E_g^{(0)} \right| u \rangle + \langle u \left| H^\prime - E_g^{(1)} \right| \psi_{g}^{(0)} \rangle
+  \langle \psi_{g}^{(0)} \left| H^\prime - E_g^{(1)} \right| u \rangle \geq E_g^{(2)} $$

```python
def var_pert(PType,u,Pert):
    #Problem type original energies and wavefuncions
    if PType == "Harmonic oscilator":
        Psin=psi_n(n,x,m,omega)
        En=E_n(n,omega)
        Psig=psi_n(0,x,m,omega)
        Eg=E_n(0,omega)
        H=(-hbar**2/(2*m))*diff(u, x, 2) + (m*omega**2*x**2/2)*u

        E_correction1=integrate(Psig*Pert*Psig, (x,-oo,oo))
        EQ=integrate(u*(H+Eg)*u, (x,-oo,oo))+integrate(u*(Pert-E_correction1)*Psig, (x,-oo,oo))+integrate(Psig*(Pert-E_correction1)*u, (x,-oo,oo))
        Opt=solve(diff(EQ,k,1),k, dict=true)
        EQ=EQ.subs(k,Opt)

    elif PType == 'Particle in a one-dimensional box':
        Psin=sqrt(2/a)*sin(n*pi*x/a)
        En=hbar**2*pi**2*n**2/(2*m*a**2)
        Psig=Psin.subs(n,1)
        Eg=En.subs(n,1)
        H=(-hbar**2/(2*m))*diff(u, x, 2)

        E_correction1=integrate(Psig*Pert*Psig, (x,0,a))
        EQ=integrate(u*H-u*Eg*u, (x,0,a))+integrate(u*(Pert-E_correction1)*Psig, (x,0,a))+integrate(Psig*(Pert-E_correction1)*u, (x,0,a))
        Opt=solve(diff(EQ,k,1),k, dict=true)
        EQ=EQ.subs(k,Opt[0][k])
    return Psin, En, E_correction1, EQ
```

```python
PType2=Dropdown(options=['Particle in a one-dimensional box', 'Harmonic oscilator'])
fPrueba=Text(value='x*k')
Pert2=Text(value='0')
Label4=HTML(value="Select the problem type and write the perturbation and test function, include a constant k for "
            "which the system will be optimized")

form_items3 = [
    Box([Label(value=''),
         Label4], layout=form_item_layout),
    Box([Label(value='Problem'),
         PType2], layout=form_item_layout),
    Box([Label(value='Test fun.'),
         fPrueba], layout=form_item_layout),
    Box([Label(value='Perturbation'),
         Pert2], layout=form_item_layout)]

form3 = Box(form_items3, layout=Layout(
    display='flex',
    flex_flow='column',
    border='solid 2px',
    align_items='stretch',
    width='40%'
))
form3

```

    A Jupyter Widget

```python
var_pert(PType2.value, eval(fPrueba.value), eval(Pert2.value))
```

$$\left ( \sqrt{2} \sqrt{\frac{1}{a}} \sin{\left (\frac{\pi x}{a} n \right )}, \quad \frac{\pi^{2} \hbar^{2} n^{2}}{2 a^{2} m}, \quad \frac{2}{a} \left(- \frac{a^{3}}{4 \pi^{2}} + \frac{a^{3}}{6}\right), \quad - \frac{a^{6} m \left(-33 + 4 \pi^{2}\right)^{2}}{3 \pi^{8} \hbar^{2}} + \frac{4 a^{3} m}{\pi^{5} a \hbar^{2}} \left(-33 + 4 \pi^{2}\right) \left(- \frac{11 a^{4}}{2 \pi^{3}} + \frac{2 a^{4}}{3 \pi}\right)\right )$$

**Tarea 6. Revisar sección 9.7. **

Inicialmente a mano, y sengunda instancia favor de intentar programar sección del problema, i.e. integral de Coulomb  e integral de intercambio.

```python

```

## Siguiente: Segunda parte, Octubre
Simetrías moleculares y Hartree-Fock

```python

```

$$NotProgramed$$

```python

```


# Tutorial 2: Linear regression with MLE
**Week 1, Day 3: Model Fitting**

**By Neuromatch Academy**

**Content creators**: Pierre-Étienne Fiquet, Anqi Wu, Alex Hyafil  with help from Byron Galbraith

**Content reviewers**: Lina Teichmann, Madineh Sarvestani, Patrick Mineault, Ella Batty, Michael Waskom

**Our 2021 Sponsors, including Presenting Sponsor Facebook Reality Labs**

---
# Tutorial Objectives

*Estimated timing of tutorial: 30 minutes*

This is Tutorial 2 of a series on fitting models to data. We start with simple linear regression, using least squares optimization (Tutorial 1) and Maximum Likelihood Estimation (Tutorial 2). We will use bootstrapping to build confidence intervals around the inferred linear model parameters (Tutorial 3). We'll finish our exploration of regression models by generalizing to multiple linear regression and polynomial regression (Tutorial 4). We end by learning how to choose between these various models. We discuss the bias-variance trade-off (Tutorial 5) and Cross Validation for model selection (Tutorial 6).

In this tutorial, we will use a different approach to fit linear models that incorporates the random 'noise' in our data.
- Learn about probability distributions and probabilistic models
- Learn how to calculate the likelihood of our model parameters
- Learn how to implement the maximum likelihood estimator, to find the model parameter with the maximum likelihood

```python
# @title Tutorial slides

# @markdown These are the slides for the videos in all tutorials today
from IPython.display import IFrame
IFrame(src=f"https://mfr.ca-1.osf.io/render?url=https://osf.io/2mkq4/?direct%26mode=render%26action=download%26mode=render", width=854, height=480)
```

---
# Setup

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

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

```python
# @title Plotting Functions
def plot_density_image(x, y, theta, sigma=1, ax=None):
  """ Plots probability distribution of y given x, theta, and sigma

  Args:

    x (ndarray): An array of shape (samples,) that contains the input values.
    y (ndarray): An array of shape (samples,) that contains the corresponding
      measurement values to the inputs.
    theta (float): Slope parameter
    sigma (float): standard deviation of Gaussian noise

  """

  # plot the probability density of p(y|x,theta)
  if ax is None:
    fig, ax = plt.subplots()

  xmin, xmax = np.floor(np.min(x)), np.ceil(np.max(x))
  ymin, ymax = np.floor(np.min(y)), np.ceil(np.max(y))
  xx = np.linspace(xmin, xmax, 50)
  yy = np.linspace(ymin, ymax, 50)

  surface = np.zeros((len(yy), len(xx)))
  for i, x_i in enumerate(xx):
    surface[:, i] = stats.norm(theta * x_i, sigma).pdf(yy)

  ax.set(xlabel='x', ylabel='y')

  return ax.imshow(surface, origin='lower', aspect='auto', vmin=0, vmax=None,
            cmap=plt.get_cmap('Wistia'),
            extent=[xmin, xmax, ymin, ymax])
```

---
# Section 1: Maximum Likelihood Estimation (MLE)

```python
# @title Video 1: Maximum Likelihood Estimation
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="BV1Mg4y1i7WH", 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="8mpNmzLKNfU", 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)
```

    Tab(children=(Output(), Output()), _titles={'0': 'Youtube', '1': 'Bilibili'})

This video covers Maximum Likelihood Estimation (MLE) in the context of a 1D linear regression.

## Section 1.1: Gaussian noise

 Click here for text recap of relevant part of video

In the previous tutorial we made the assumption that the data was drawn from a linear relationship with noise added, and found an effective approach for estimating model parameters based on minimizing the mean squared error.

In that case we treated the noise as simply a nuisance, but what if we factored it directly into our model?

Recall our linear model:

\begin{align}
y = \theta x + \epsilon.
\end{align}

The noise component $\epsilon$ is often modeled as a random variable drawn from a Gaussian distribution (also called the normal distribution).

The Gaussian distribution is described by its [probability density function](https://en.wikipedia.org/wiki/Probability_density_function) (pdf)
\begin{align}
\mathcal{N}(x; \mu, \sigma^2) = \frac{1}{\sqrt{2\pi\sigma^2}}e^{-\frac{1}{2\sigma^2}(x-\mu)^2}
\end{align}

and is dependent on two parameters: the mean $\mu$ and the variance $\sigma^2$. We often consider the noise signal to be Gaussian "white noise", with zero mean and unit variance:

\begin{align}
\epsilon \sim \mathcal{N}(0, 1).
\end{align}

### Interactive Demo 1.1: Gaussian Distribution Explorer

Use the explorer widget below to see how varying the $\mu$ and $\sigma$ parameters change the location and shape of the samples.

1.   What effect does varying $\mu$ have on the pdf?
2.   What effect does varying $\sigma$ have on the pdf?

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

@widgets.interact(mu=widgets.FloatSlider(0.0, min=-2.0, max=2.0),
                  sigma=widgets.FloatSlider(1.0, min=0.5, max=2.0))
def plot_normal_dist(mu=0, sigma=1):

  # Generate pdf & samples from normal distribution with mu/sigma
  rv = stats.norm(mu, sigma)
  x = np.linspace(-5, 5, 100)
  y = rv.pdf(x)
  samples = rv.rvs(1000)

  # Plot
  fig, ax = plt.subplots()
  ax.hist(samples, 20, density=True, color='g', histtype='stepfilled', alpha=0.8,
          label='histogram')
  ax.plot(x, y, color='orange', linewidth=3, label='pdf')
  ax.vlines(mu, 0, rv.pdf(mu), color='y', linewidth=3, label='$\mu$')
  ax.vlines([mu-sigma, mu+sigma], 0, rv.pdf([mu-sigma, mu+sigma]), colors='red',
            color='b', linewidth=3, label='$\sigma$')
  ax.set(xlabel='x', ylabel='probability density', xlim=[-5, 5], ylim=[0, 1.0])
  ax.legend()
```

    interactive(children=(FloatSlider(value=0.0, description='mu', max=2.0, min=-2.0), FloatSlider(value=1.0, desc…

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

## Section 1.2: Probabilistic Models

*Estimated timing to here from start of tutorial: 11 min*

Now that we have a model of our noise component $\epsilon$ as random variable, how do we incorporate this back into our original linear model from before? Consider again our simplified model $y = \theta x + \epsilon$ where the noise has zero mean and unit variance $\epsilon \sim \mathcal{N}(0, 1)$. We can now also treat $y$ as a random variable drawn from a Gaussian distribution where $\mu = \theta x$ and $\sigma^2 = 1$:

\begin{align}
y \sim \mathcal{N}(\theta x, 1)
\end{align}

which is to say that the probability of observing $y$ given $x$ and parameter $\theta$ is
\begin{align}
p(y|x,\theta) = \frac{1}{\sqrt{2\pi}}e^{-\frac{1}{2}(y-\theta x)^2}
\end{align}

Note that in this and the following sections, we will focus on a single data point (a single pairing of $x$ and $y$). We have dropped the subscript $i$ just for simplicity (that is, we use $x$ for a single data point, instead of $x_i$).

Now, let's revisit our original sample dataset where the true underlying model has $\theta = 1.2$.

```python
# @markdown Execute this cell to generate some simulated data

# setting a fixed seed to our random number generator ensures we will always
# get the same psuedorandom number sequence

np.random.seed(121)
theta = 1.2
n_samples = 30
x = 10 * np.random.rand(n_samples) # sample from a uniform distribution over [0,10)
noise = np.random.randn(n_samples) # sample from a standard normal distribution
y = theta * x + noise
```

This time we can plot the density of $p(y|x,\theta=1.2)$ and see how $p(y)$ changes for different values of $x$.

```python
# @markdown Execute this cell to visualize p(y|x, theta=1.2)

fig, (ax1, ax2) = plt.subplots(ncols=2, figsize=(10, 4))

# Invokes helper function to generate density image plots from data and parameters
im = plot_density_image(x, y, 1.2, ax=ax1)
plt.colorbar(im, ax=ax1)
ax1.axvline(8, color='k')
ax1.set(title=r'p(y | x, $\theta$=1.2)')

# Plot pdf for given x
ylim = ax1.get_ylim()
yy = np.linspace(ylim[0], ylim[1], 50)
ax2.plot(yy, stats.norm(theta * 8, 1).pdf(yy), color='orange', linewidth=2)
ax2.set(
    title=r'p(y|x=8, $\theta$=1.2)',
    xlabel='y',
    ylabel='probability density');
```

## Section 1.3: Likelihood Estimation

*Estimated timing to here from start of tutorial: 15 min*

Now that we have our probabilistic model, we turn back to our original challenge of finding a good estimate for $\theta$ that fits our data. Given the inherent uncertainty when dealing in probabilities, we talk about the [likelihood](https://en.wikipedia.org/wiki/Likelihood_function) that some estimate $\hat{\theta}$ fits our data. The likelihood function $\mathcal{L}(\theta)$ is equal to the probability density function parameterized by that $\theta$:

\begin{align}
\mathcal{L}(\theta|x,y) = p(y|x,\theta) = \frac{1}{\sqrt{2\pi\sigma^2}}e^{-\frac{1}{2\sigma^2}(y-\theta x)^2}
\end{align}

### Coding Exercise 1.3: Likelihood Function

In this exercise you will implement the likelihood function $\mathcal{L}(\theta|x, y)$ for our linear model where $\sigma = 1$.

After implementing this function, we can produce probabilities that our estimate $\hat{\theta}$ generated the provided observations. We will try with one of the samples from our dataset.

TIP: Use `np.exp` and `np.sqrt` for the exponential and square root functions, respectively.

```python
def likelihood(theta_hat, x, y):
  """The likelihood function for a linear model with noise sampled from a
    Gaussian distribution with zero mean and unit variance.

  Args:
    theta_hat (float): An estimate of the slope parameter.
    x (ndarray): An array of shape (samples,) that contains the input values.
    y (ndarray): An array of shape (samples,) that contains the corresponding
      measurement values to the inputs.

  Returns:
    ndarray: the likelihood values for the theta_hat estimate
  """
  sigma = 1

  ##############################################################################
  ## TODO for students: implement the likelihood function
  # Fill out function and remove
  #raise NotImplementedError("Student exercise: implement the likelihood function")
  ##############################################################################

  # Compute Gaussian likelihood
  pdf = 1/np.sqrt(2*np.pi)*np.exp(-(y-theta_hat*x)**2/2)

  return pdf

print(likelihood(1.0, x[1], y[1]))
```

    0.11344443599846923

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

We should see that $\mathcal{L}(\theta=1.0|x=2.1,y=3.7) \approx 0.11$. So far so good, but how does this tell us how this estimate is better than any others?

When dealing with a set of data points, as we are with our dataset, we are concerned with their joint probability -- the likelihood that all data points are explained by our parameterization. Since we have assumed that the noise affects each output independently, we can factorize the likelihood, and write:

\begin{align}
\mathcal{L}(\theta|\mathbf{x}, \mathbf{y}) = \prod_{i=1}^N \mathcal{L}(\theta|x_i,y_i),
\end{align}

where we have $N$ data points $\mathbf{x} = [x_1,...,x_N]$ and $\mathbf{y} = [y_1,...,y_N]$.

In practice, such a product can be numerically unstable. Indeed multiplying small values together can lead to [underflow](https://en.wikipedia.org/wiki/Arithmetic_underflow), the situation in which the digital representation of floating point number reaches its limit. This problem can be circumvented by taking the logarithm of the likelihood because the logarithm transforms products into sums:

\begin{align}
\log\mathcal{L}(\theta|\mathbf{x}, \mathbf{y}) = \sum_{i=1}^N \log\mathcal{L}(\theta|x_i,y_i)
\end{align}

We can take the sum of the log of the output of our `likelihood` method applied to the full dataset to get a better idea of how different $\hat{\theta}$ compare. We can also plot the different distribution densities over our dataset and see how they line up qualitatively.

```python
# @markdown Execute this cell to visualize different distribution densities
theta_hats = [0.5, 1.0, 2.2]
fig, axes = plt.subplots(ncols=3, figsize=(16, 4))
for theta_hat, ax in zip(theta_hats, axes):
  ll = np.sum(np.log(likelihood(theta_hat, x, y)))  # log likelihood
  im = plot_density_image(x, y, theta_hat, ax=ax)
  ax.scatter(x, y)
  ax.set(title=fr'$\hat{{\theta}}$ = {theta_hat}, log likelihood: {ll:.2f}')
plt.colorbar(im, ax=ax);
```

Using the log likelihood calculation, we see that $\mathcal{L}(\theta=1.0) > \mathcal{L}(\theta=0.5) > \mathcal{L}(\theta=2.2)$.

This is great: now we have a way to compare estimators based on likelihood. But like with the MSE approach, we want an analytic solution to find the best estimator. In this case, we want to find the estimator that maximizes the likelihood.

## Section 1.4: Finding the Maximum Likelihood Estimator

*Estimated timing to here from start of tutorial: 23 min*

 Click here for text recap of relevant part of video

We want to find the parameter value $\hat\theta$ that makes our data set most likely:

\begin{align}
\hat{\theta}_{\textrm{MLE}} = \underset{\theta}{\operatorname{argmax}} \mathcal{L}(\theta|X,Y)
\end{align}

We discussed how taking the logarithm of the likelihood helps with numerical stability, the good thing is that it does so without changing the parameter value that maximizes the likelihood. Indeed, the $\log()$ function is *monotonically increasing*, which means that it preserves the order of its inputs. So we have:

\begin{align}
\hat{\theta}_{\textrm{MLE}} = \underset{\theta}{\operatorname{argmax}} \sum_{i=1}^m \textrm{log} \mathcal{L}(\theta|x_i,y_i)
\end{align}

Now substituting our specific likelihood function and taking its logarithm, we get:
\begin{align}
\hat{\theta}_{\textrm{MLE}} = \underset{\theta}{\operatorname{argmax}} [-\frac{N}{2} \operatorname{log} 2\pi\sigma^2 - \frac{1}{2\sigma^2}\sum_{i=1}^N (y_i-\theta x_i)^2].
\end{align}

Note that maximizing the log likelihood is the same as minimizing the negative log likelihood (in practice optimization routines are developed to solve minimization not maximization problems). Because of the convexity of this objective function, we can take the derivative of our negative log likelihhood, set it to 0, and solve - just like our solution to minimizing MSE.

\begin{align}
\frac{\partial\operatorname{log}\mathcal{L}(\theta|x,y)}{\partial\theta}=\frac{1}{\sigma^2}\sum_{i=1}^N(y_i-\theta x_i)x_i = 0
\end{align}

This looks remarkably like the equation we had to solve for the optimal MSE estimator, and, in fact, we arrive to the exact same solution!

\begin{align}
\hat{\theta}_{\textrm{MLE}} = \hat{\theta}_{\textrm{MSE}} = \frac{\sum_{i=1}^N x_i y_i}{\sum_{i=1}^N x_i^2}
\end{align}

```python
# Compute theta_hat_MLE
theta_hat_mle = (x @ y) / (x @ x)
```

```python
#@title
#@markdown Execute this cell to visualize density with theta_hat_mle

# Plot the resulting distribution density
fig, ax = plt.subplots()
ll = np.sum(np.log(likelihood(theta_hat_mle, x, y))) # log likelihood
im = plot_density_image(x, y, theta_hat_mle, ax=ax)
plt.colorbar(im, ax=ax);
ax.scatter(x, y)
ax.set(title=fr'$\hat{{\theta}}$ = {theta_hat_mle:.2f}, log likelihood: {ll:.2f}');
```

---
# Summary

*Estimated timing of tutorial: 30 minutes*

Likelihood vs probability:

  - $\mathcal{L}(\theta|x, y) = p(y|x, \theta)$
  - $p(y|x, \theta)$ -> "probability of observing the response $y$ given parameter $\theta$ and input $x$"
  - $\mathcal{L}(\theta|x, y)$ -> "likelihood model that parameters $\theta$ produced response $y$ from input $x$"

Log-likelihood maximization:

  - We take the $\log$ of the likelihood function for computational convenience
  - The parameters $\theta$ that maximize $\log\mathcal{L}(\theta|x, y)$ are the model parameters that maximize the probability of observing the data.

**Key point**: the log-likelihood is a flexible cost function, and is often used to find model parameters that best fit the data.

---
# Notation

\begin{align}
x &\quad \text{input, independent variable}\\
y &\quad \text{response measurement, dependent variable}\\
\mathbf{x} &\quad \text{vector of input values}\\
\mathbf{y} &\quad \text{vector of measurements}\\
\epsilon &\quad \text{measurement error, noise contribution}\\
\epsilon \sim \mathcal{N}(\mu, \sigma^2) &\quad \text{random variable } \epsilon \text{ is distributed according to a Gaussian distribution, with mean } \mu \text{ and variance } \sigma^2\\
\mu &\quad \text{mean}\\
\sigma^2 &\quad \text{variance}\\
\sigma &\quad \text{standard deviation}\\
\theta &\quad \text{parameter}\\
\hat{\theta} &\quad \text{estimate of parameter}\\
\mathcal{L}(\theta|x, y) &\quad \text{likelihood of that parameter } \theta \text{ producing response } y \text{ from input } x \\
p(y|x, \theta) &\quad \text{probability of observing the response } y \text{ given input } x \text{ and parameter } \theta \\
\end{align}

---
# Bonus

We can also see $\mathrm{p}(\mathrm{y} | \mathrm{x}, \theta)$ as a function of $x$. This is the stimulus likelihood function, and it is useful in case we want to decode the input $x$ from observed responses $y$. This is what is relevant from the point of view of a neuron that does not have access to the outside world and tries to infer what's out there from the responses of other neurons!