# Learning disentangled representations in 'sphere-world' with an object that can also change color

## 1. Preamble

```python
import os

import torch
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 random

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)

mk_dir("_data/")
```

    created dir:  _data/

## 2. Defining objects

### SphereWorld environment

We use a cubic 3D world in which a ball evolves on a sphere. Observations are a 3-dimensional tensor with values 0 except at the voxel occupied by the ball which has value 1. When the ball is not at the center of the voxel (which is most of the time) its density is distributed around neighboring voxels.

The ball can also "change color"; in our implementation color is encoded as the numerical value of a voxel in a corner; we note that other ways of encoding color could be used.

Actions are discrete rotations around the three axes and changes in color.

```python
class SphereWorld():

    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,num_colors=5,radius=3.5):

        self.size = dim
        self.num_colors = num_colors
        self.radius = radius
        self.action_space = self.action_space(8)
        self.observation_space = self.observation_space((2*dim)**3)

        self.reset()

    def reset(self):

        self.theta = random.uniform(0, np.pi)
        self.phi = random.uniform(0, np.pi*2)
        self.color = random.randrange(self.num_colors)

        self.ball_coordinates = [self.radius*np.sin(self.theta)*np.cos(self.phi), \
                                 self.radius*np.sin(self.theta)*np.sin(self.phi), \
                                 self.radius*np.cos(self.theta)]

        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((2*self.size, 2*self.size, 2*self.size))

        ball_low_x = ball_coordinates[0] + self.size - .5
        ball_low_y = ball_coordinates[1] + self.size - .5
        ball_low_z = ball_coordinates[2] + self.size - .5

        fill_x = (int(ball_low_x + 1) - ball_low_x)
        fill_y = (int(ball_low_y + 1) - ball_low_y)
        fill_z = (int(ball_low_z + 1) - ball_low_z)

        state[int(ball_low_x), int(ball_low_y), int(ball_low_z)] \
        = fill_x * fill_y * fill_z
        state[(int(ball_low_x) + 1), int(ball_low_y), int(ball_low_z)] \
        = (1 - fill_x) * fill_y * fill_z
        state[int(ball_low_x), (int(ball_low_y) + 1), int(ball_low_z)] \
        = fill_x * (1 - fill_y) * fill_z
        state[(int(ball_low_x) + 1), (int(ball_low_y) +1), int(ball_low_z)] \
        = (1 - fill_x) * (1 - fill_y) * fill_z
        state[int(ball_low_x), int(ball_low_y), int(ball_low_z) + 1] \
        = fill_x * fill_y * (1 - fill_z)
        state[(int(ball_low_x) + 1), int(ball_low_y), int(ball_low_z) + 1] \
        = (1 - fill_x) * fill_y * (1 - fill_z)
        state[int(ball_low_x), (int(ball_low_y) + 1), int(ball_low_z) + 1] \
        = fill_x * (1 - fill_y) * (1 - fill_z)
        state[(int(ball_low_x) + 1), (int(ball_low_y) +1), int(ball_low_z) + 1] \
        = (1 - fill_x) * (1 - fill_y) * (1 - fill_z)

        state[0, 0, 0] = self.color / (self.num_colors - 1)

        self.state = state
        return self.state

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

    def step(self, action, angle):

        ball_coordinates = self.ball_coordinates

        if action == 0 or action == 1:
            if action == 1:
                angle = - angle
            new_ball_coordinates = [ball_coordinates[0]*np.cos(angle) + ball_coordinates[1]*np.sin(angle), \
                                    - ball_coordinates[0]*np.sin(angle) + ball_coordinates[1]*np.cos(angle), \
                                    ball_coordinates[2]]
        elif action == 2 or action == 3:
            if action == 3:
                angle = -angle
            new_ball_coordinates = [ball_coordinates[0]*np.cos(angle) + ball_coordinates[2]*np.sin(angle), \
                                    ball_coordinates[1], \
                                    - ball_coordinates[0]*np.sin(angle) + ball_coordinates[2]*np.cos(angle)]
        elif action == 4 or action == 5:
            if action == 5:
                angle = -angle
            new_ball_coordinates = [ball_coordinates[0], \
                                    ball_coordinates[1]*np.cos(angle) + ball_coordinates[2]*np.sin(angle), \
                                    - ball_coordinates[1]*np.sin(angle) + ball_coordinates[2]*np.cos(angle)]
        elif action == 6 or action == 7:
            if action == 6:
                self.color = (self.color + 1) % self.num_colors
            if action == 7:
                self.color = (self.color - 1) % self.num_colors
            new_ball_coordinates = ball_coordinates
        else:
            raise Exception("INVALID ACTION")

        self.ball_coordinates = new_ball_coordinates
        self.state = self.get_state()
        return self.get_observation()
```

Show observations from the sphere world

```python
def plot_state(state, ax):
    ax[0].pcolormesh(state.sum(axis=0), edgecolors='gray', \
                     linewidth=2, cmap=colors[round(state[0,0,0].item()*4)])
    ax[0].set_aspect("equal")
    ax[0].set_xticks([])
    ax[0].set_yticks([])

    ax[1].pcolormesh(state.sum(axis=1), edgecolors='gray', \
                     linewidth=2, cmap=colors[round(state[0,0,0].item()*4)])
    ax[1].set_aspect("equal")
    ax[1].set_xticks([])
    ax[1].set_yticks([])

    ax[2].pcolormesh(state.sum(axis=2), edgecolors='gray', \
                     linewidth=2, cmap=colors[round(state[0,0,0].item()*4)])
    ax[2].set_aspect("equal")
    ax[2].set_xticks([])
    ax[2].set_yticks([])

    return ax

n_steps = 5

env = SphereWorld()
colors = ['prism', 'ocean', 'gist_earth', 'terrain', 'jet']
fig, axs = plt.subplots(3, n_steps+1, figsize=(30, 10))

plot_state(env.state,axs[:,0])

for i in range(n_steps):
    action = 7
    print(action)
    env.step(action, 2*np.pi/5)
    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 5-dimensional).  We will require both an encoder and decoder.

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

    def __init__(self, n_in, n_out, n_hid=[16], 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*5, n_out=6, n_hid=[16])
print(encoder)

decoder = Decoder(n_in=6, n_out=5*5*5, n_hid=[16])
print(decoder)

def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

count_parameters(encoder)
```

    Encoder(
      (layers): ModuleList(
        (0): Sequential(
          (0): Linear(in_features=125, out_features=16, bias=True)
          (1): ReLU()
        )
        (1): Sequential(
          (0): Linear(in_features=16, out_features=6, bias=True)
          (1): nnNorm()
        )
      )
    )
    Decoder(
      (layers): ModuleList(
        (0): Sequential(
          (0): Linear(in_features=6, out_features=16, bias=True)
          (1): ReLU()
        )
        (1): Sequential(
          (0): Linear(in_features=16, out_features=125, bias=True)
          (1): Sigmoid()
        )
      )
    )

    2118

Check dimensions

```python
x = torch.rand(125)
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([125])) :
     tensor([0.6761, 0.4820, 0.7408, 0.2386, 0.3486, 0.2950, 0.8632, 0.8790, 0.8269,
            0.4200, 0.7829, 0.6518, 0.4770, 0.2077, 0.7223, 0.8875, 0.8376, 0.6497,
            0.1558, 0.7254, 0.2747, 0.7799, 0.4912, 0.8103, 0.0544, 0.9080, 0.1421,
            0.9176, 0.6991, 0.6408, 0.5825, 0.3857, 0.5358, 0.0152, 0.4510, 0.5514,
            0.7959, 0.1915, 0.1922, 0.7259, 0.5406, 0.2116, 0.6574, 0.7666, 0.1712,
            0.3441, 0.0542, 0.6058, 0.6850, 0.7394, 0.4468, 0.1122, 0.3258, 0.2841,
            0.7519, 0.5401, 0.0196, 0.9722, 0.5374, 0.9505, 0.8633, 0.8376, 0.3554,
            0.5921, 0.5202, 0.5923, 0.6955, 0.0742, 0.9440, 0.4858, 0.4068, 0.6569,
            0.8637, 0.4485, 0.4143, 0.9760, 0.3303, 0.5757, 0.3738, 0.8466, 0.1429,
            0.5652, 0.3835, 0.3634, 0.6160, 0.5427, 0.9429, 0.2058, 0.2615, 0.4748,
            0.9359, 0.4714, 0.7497, 0.2833, 0.5276, 0.2132, 0.9034, 0.4006, 0.7558,
            0.6108, 0.1836, 0.3713, 0.1677, 0.9924, 0.7160, 0.4070, 0.8322, 0.7397,
            0.2043, 0.1117, 0.4084, 0.9744, 0.4957, 0.8936, 0.3244, 0.2283, 0.9148,
            0.6934, 0.1835, 0.2860, 0.8170, 0.7777, 0.8354, 0.2647, 0.2892])
    Latent rep (shape = torch.Size([6]), norm = 1.0) :
     tensor([ 0.4450, -0.3121, -0.8203, -0.0913,  0.1152,  0.1006],
           grad_fn=<DivBackward0>)
    Reconstructed Observable rep (shape = torch.Size([125])) :
     tensor([0.4317, 0.5300, 0.4662, 0.4641, 0.4741, 0.5925, 0.5151, 0.4596, 0.5708,
            0.4913, 0.5302, 0.4450, 0.4219, 0.4561, 0.5029, 0.4770, 0.6102, 0.4330,
            0.4871, 0.6013, 0.4541, 0.3889, 0.4015, 0.6460, 0.5246, 0.4306, 0.4311,
            0.4858, 0.4467, 0.3812, 0.5108, 0.4627, 0.4441, 0.4698, 0.6368, 0.5235,
            0.5743, 0.4840, 0.4748, 0.5002, 0.4177, 0.5188, 0.4505, 0.5266, 0.4684,
            0.4768, 0.4641, 0.4411, 0.4641, 0.4193, 0.5670, 0.5425, 0.4793, 0.4458,
            0.5203, 0.4904, 0.4852, 0.5162, 0.5413, 0.5340, 0.5184, 0.5467, 0.5072,
            0.5066, 0.4123, 0.4753, 0.4902, 0.5263, 0.5678, 0.5391, 0.5060, 0.4970,
            0.4364, 0.4427, 0.4733, 0.4427, 0.4827, 0.4897, 0.5100, 0.5074, 0.4745,
            0.4689, 0.4961, 0.5624, 0.5184, 0.5460, 0.5717, 0.3963, 0.5178, 0.5655,
            0.5850, 0.4252, 0.4954, 0.5148, 0.5041, 0.5039, 0.5243, 0.5690, 0.5419,
            0.5085, 0.5790, 0.6220, 0.4009, 0.5038, 0.5568, 0.5146, 0.6718, 0.4251,
            0.4165, 0.4897, 0.5056, 0.5288, 0.4663, 0.5104, 0.5399, 0.4680, 0.4149,
            0.5248, 0.5181, 0.4919, 0.5986, 0.4356, 0.4749, 0.5328, 0.5199],
           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=5):
        self.dim = dim
        self.params = dim*(dim-1)//2
        self.thetas = torch.autograd.Variable(torch.rand(self.params), 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=5,
                 n_actions=8,
                 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):
        self.state = state
        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
```

## 5. Disentanglement

***Some jargon***

We focus on learning 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 SphereWorld can be represented as a subgroup of $\mathrm{SO}(3) \times \mathrm{SO}(2)$, therefore we hope to find the disentangled representation of 6 rotations and 2 changes in color that correspond 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

Note: undesirable local minima are known to exist (in particular the loss can sometimes get stuck at 8e-3), which can sometimes lead to learning poor representations

```python
obs_env = SphereWorld(dim = 5, radius = 4., num_colors=5)
lat_env = LatentWorld(dim = 5,
                      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 = [64])

encoder = Encoder(n_in = sum(obs_env.observation_space.shape),
                  n_out = sum(lat_env.observation_space.shape),
                  n_hid = [64])

optimizer_dec = optim.Adam(decoder.parameters(),
                           lr=3e-2,
#                            betas=(0.9, 0.99),
                           weight_decay=0)

optimizer_enc = optim.Adam(encoder.parameters(),
                           lr=4e-2,
#                            betas=(0.9, 0.99),
                           weight_decay=0)

optimizer_rep = optim.Adam(lat_env.get_representation_params(),
                           lr=3e-2,
#                            betas=(0.9, 0.99),
                           weight_decay=0)

losses = []
entanglement = []
orthogonality = []
```

```python
n_sgd_steps = 15000
ep_steps = 5
batch_eps = 16

angle = np.pi * 2/5

i = 1

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, angle)
                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 += (obs_x_recon[0] - obs_x[0]).abs() / 100 #Insist on the color

    loss /= (ep_steps * batch_eps)
    loss_raw = loss

    reg_loss = sum([calc_entanglement(r.thetas) for r in lat_env.action_reps])*1e-4
    loss += reg_loss

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

    optimizer_dec.zero_grad()
    optimizer_rep.zero_grad()
    optimizer_enc.zero_grad()
    loss.backward()
    optimizer_dec.step()
    optimizer_rep.step()
    optimizer_enc.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={:.3e} : entanglement={:.2e} : last 10 iters in {:.3f}s".format(
            i, loss.item(), reg_loss.item(), time.time() - t_start
            ), end="\r" if i%100 else "\n")
        t_start = time.time()
```

    iter 100 : loss=1.413e-02 : entanglement=4.17e-04 : last 10 iters in 0.913s
    iter 200 : loss=1.025e-02 : entanglement=9.16e-05 : last 10 iters in 0.926s
    iter 300 : loss=1.047e-02 : entanglement=4.18e-05 : last 10 iters in 1.172s
    iter 400 : loss=9.737e-03 : entanglement=4.85e-05 : last 10 iters in 0.903s
    iter 500 : loss=9.817e-03 : entanglement=5.36e-05 : last 10 iters in 0.906s
    iter 600 : loss=9.832e-03 : entanglement=4.15e-05 : last 10 iters in 1.038s
    iter 700 : loss=9.425e-03 : entanglement=3.88e-05 : last 10 iters in 0.910s
    iter 800 : loss=9.323e-03 : entanglement=3.58e-05 : last 10 iters in 0.996s
    iter 900 : loss=9.213e-03 : entanglement=3.25e-05 : last 10 iters in 0.964s
    iter 1000 : loss=9.013e-03 : entanglement=2.87e-05 : last 10 iters in 0.925s
    iter 1100 : loss=9.338e-03 : entanglement=4.03e-05 : last 10 iters in 0.896s
    iter 1200 : loss=9.229e-03 : entanglement=2.93e-05 : last 10 iters in 0.909s
    iter 1300 : loss=8.654e-03 : entanglement=1.86e-05 : last 10 iters in 0.925s
    iter 1400 : loss=8.736e-03 : entanglement=1.12e-05 : last 10 iters in 1.031s
    iter 1500 : loss=9.302e-03 : entanglement=1.55e-05 : last 10 iters in 0.937s
    iter 1600 : loss=8.273e-03 : entanglement=1.21e-05 : last 10 iters in 0.977s
    iter 1700 : loss=9.641e-03 : entanglement=1.78e-05 : last 10 iters in 1.070s
    iter 1800 : loss=8.381e-03 : entanglement=3.20e-05 : last 10 iters in 1.088s
    iter 1900 : loss=8.553e-03 : entanglement=2.37e-05 : last 10 iters in 1.078s
    iter 2000 : loss=8.242e-03 : entanglement=2.76e-05 : last 10 iters in 0.927s
    iter 2100 : loss=8.756e-03 : entanglement=5.27e-05 : last 10 iters in 0.892s
    iter 2200 : loss=8.191e-03 : entanglement=2.52e-05 : last 10 iters in 0.977s
    iter 2300 : loss=8.468e-03 : entanglement=2.74e-05 : last 10 iters in 0.928s
    iter 2400 : loss=8.265e-03 : entanglement=2.82e-05 : last 10 iters in 0.932s
    iter 2500 : loss=8.004e-03 : entanglement=5.72e-05 : last 10 iters in 0.936s
    iter 2600 : loss=6.925e-03 : entanglement=4.53e-05 : last 10 iters in 0.914s
    iter 2700 : loss=6.640e-03 : entanglement=2.14e-05 : last 10 iters in 0.919s
    iter 2800 : loss=5.999e-03 : entanglement=1.59e-05 : last 10 iters in 0.895s
    iter 2900 : loss=5.755e-03 : entanglement=1.58e-05 : last 10 iters in 0.912s
    iter 3000 : loss=5.840e-03 : entanglement=1.24e-05 : last 10 iters in 0.941s
    iter 3100 : loss=5.392e-03 : entanglement=6.65e-06 : last 10 iters in 0.949s
    iter 3200 : loss=6.008e-03 : entanglement=1.60e-05 : last 10 iters in 1.007s
    iter 3300 : loss=4.956e-03 : entanglement=3.61e-06 : last 10 iters in 0.904s
    iter 3400 : loss=4.831e-03 : entanglement=7.39e-06 : last 10 iters in 0.947s
    iter 3500 : loss=4.991e-03 : entanglement=1.25e-05 : last 10 iters in 0.932s
    iter 3600 : loss=5.360e-03 : entanglement=8.10e-06 : last 10 iters in 0.903s
    iter 3700 : loss=4.709e-03 : entanglement=4.43e-06 : last 10 iters in 0.959s
    iter 3800 : loss=4.304e-03 : entanglement=2.49e-06 : last 10 iters in 0.920s
    iter 3900 : loss=5.200e-03 : entanglement=9.96e-06 : last 10 iters in 0.953s
    iter 4000 : loss=4.792e-03 : entanglement=6.95e-06 : last 10 iters in 0.926s
    iter 4100 : loss=4.522e-03 : entanglement=3.11e-06 : last 10 iters in 0.958s
    iter 4200 : loss=4.516e-03 : entanglement=7.16e-06 : last 10 iters in 0.990s
    iter 4300 : loss=4.724e-03 : entanglement=6.29e-06 : last 10 iters in 0.943s
    iter 4400 : loss=4.377e-03 : entanglement=3.62e-06 : last 10 iters in 1.035s
    iter 4500 : loss=4.955e-03 : entanglement=7.52e-06 : last 10 iters in 1.563s
    iter 4600 : loss=4.886e-03 : entanglement=4.16e-06 : last 10 iters in 0.953s
    iter 4700 : loss=4.283e-03 : entanglement=5.05e-06 : last 10 iters in 0.912s
    iter 4800 : loss=4.080e-03 : entanglement=3.88e-06 : last 10 iters in 0.890s
    iter 4900 : loss=4.438e-03 : entanglement=2.01e-06 : last 10 iters in 0.907s
    iter 5000 : loss=3.991e-03 : entanglement=2.23e-06 : last 10 iters in 0.950s
    iter 5100 : loss=4.517e-03 : entanglement=5.03e-06 : last 10 iters in 0.933s
    iter 5200 : loss=4.436e-03 : entanglement=4.95e-06 : last 10 iters in 1.118s
    iter 5300 : loss=4.098e-03 : entanglement=3.93e-06 : last 10 iters in 1.046s
    iter 5400 : loss=3.779e-03 : entanglement=3.09e-06 : last 10 iters in 1.181s
    iter 5500 : loss=4.209e-03 : entanglement=7.97e-06 : last 10 iters in 1.087s
    iter 5600 : loss=4.373e-03 : entanglement=2.65e-06 : last 10 iters in 1.099s
    iter 5700 : loss=3.727e-03 : entanglement=2.77e-06 : last 10 iters in 1.136s
    iter 5800 : loss=4.096e-03 : entanglement=2.12e-06 : last 10 iters in 1.284s
    iter 5900 : loss=4.421e-03 : entanglement=5.37e-06 : last 10 iters in 1.219s
    iter 6000 : loss=4.593e-03 : entanglement=7.83e-06 : last 10 iters in 1.068s
    iter 6100 : loss=3.780e-03 : entanglement=3.30e-06 : last 10 iters in 1.220s
    iter 6200 : loss=4.091e-03 : entanglement=2.63e-06 : last 10 iters in 1.105s
    iter 6300 : loss=4.537e-03 : entanglement=2.64e-06 : last 10 iters in 1.133s
    iter 6400 : loss=4.738e-03 : entanglement=7.06e-06 : last 10 iters in 1.143s
    iter 6500 : loss=5.066e-03 : entanglement=8.21e-06 : last 10 iters in 1.084s
    iter 6600 : loss=4.277e-03 : entanglement=2.51e-06 : last 10 iters in 1.233s
    iter 6700 : loss=4.308e-03 : entanglement=3.53e-06 : last 10 iters in 1.062s
    iter 6800 : loss=4.121e-03 : entanglement=2.02e-06 : last 10 iters in 0.972s
    iter 6900 : loss=4.075e-03 : entanglement=1.67e-06 : last 10 iters in 1.045s
    iter 7000 : loss=4.583e-03 : entanglement=2.15e-06 : last 10 iters in 1.384s
    iter 7100 : loss=4.608e-03 : entanglement=1.21e-05 : last 10 iters in 1.491s
    iter 7200 : loss=4.418e-03 : entanglement=6.30e-06 : last 10 iters in 1.197s
    iter 7300 : loss=4.034e-03 : entanglement=1.33e-06 : last 10 iters in 1.105s
    iter 7400 : loss=3.932e-03 : entanglement=2.26e-06 : last 10 iters in 0.982s
    iter 7500 : loss=4.022e-03 : entanglement=4.16e-06 : last 10 iters in 1.076s
    iter 7600 : loss=3.966e-03 : entanglement=5.05e-06 : last 10 iters in 1.106s
    iter 7700 : loss=4.221e-03 : entanglement=3.03e-06 : last 10 iters in 1.048s
    iter 7800 : loss=5.230e-03 : entanglement=1.71e-05 : last 10 iters in 1.084s
    iter 7900 : loss=4.315e-03 : entanglement=3.25e-06 : last 10 iters in 1.366s
    iter 8000 : loss=4.493e-03 : entanglement=1.56e-06 : last 10 iters in 1.148s
    iter 8100 : loss=4.248e-03 : entanglement=2.54e-06 : last 10 iters in 1.511s
    iter 8200 : loss=4.166e-03 : entanglement=1.90e-06 : last 10 iters in 1.184s
    iter 8300 : loss=4.382e-03 : entanglement=4.76e-06 : last 10 iters in 1.065s
    iter 8400 : loss=4.241e-03 : entanglement=1.94e-06 : last 10 iters in 1.017s
    iter 8500 : loss=3.934e-03 : entanglement=1.47e-06 : last 10 iters in 1.361s
    iter 8600 : loss=4.225e-03 : entanglement=3.91e-06 : last 10 iters in 1.145s
    iter 8700 : loss=4.054e-03 : entanglement=1.66e-06 : last 10 iters in 1.188s
    iter 8800 : loss=4.907e-03 : entanglement=6.63e-06 : last 10 iters in 1.095s
    iter 8900 : loss=4.075e-03 : entanglement=6.63e-06 : last 10 iters in 1.038s
    iter 9000 : loss=4.984e-03 : entanglement=8.87e-06 : last 10 iters in 0.955s
    iter 9100 : loss=3.860e-03 : entanglement=2.12e-06 : last 10 iters in 1.127s
    iter 9200 : loss=4.209e-03 : entanglement=4.24e-06 : last 10 iters in 1.404s
    iter 9300 : loss=3.879e-03 : entanglement=2.85e-06 : last 10 iters in 1.263s
    iter 9400 : loss=4.142e-03 : entanglement=2.56e-06 : last 10 iters in 1.051s
    iter 9500 : loss=3.917e-03 : entanglement=2.78e-06 : last 10 iters in 1.244s
    iter 9600 : loss=4.708e-03 : entanglement=9.57e-06 : last 10 iters in 1.036s
    iter 9700 : loss=5.169e-03 : entanglement=9.32e-06 : last 10 iters in 1.088s
    iter 9800 : loss=4.446e-03 : entanglement=2.13e-06 : last 10 iters in 1.154s
    iter 9900 : loss=4.108e-03 : entanglement=5.23e-06 : last 10 iters in 1.216s
    iter 10000 : loss=3.914e-03 : entanglement=1.31e-06 : last 10 iters in 1.064s
    iter 10100 : loss=3.935e-03 : entanglement=3.22e-06 : last 10 iters in 1.069s
    iter 10200 : loss=4.080e-03 : entanglement=2.34e-06 : last 10 iters in 1.075s
    iter 10300 : loss=3.707e-03 : entanglement=1.73e-06 : last 10 iters in 1.093s
    iter 10400 : loss=3.917e-03 : entanglement=2.00e-06 : last 10 iters in 1.379s
    iter 10500 : loss=3.937e-03 : entanglement=2.65e-06 : last 10 iters in 1.089s
    iter 10600 : loss=4.206e-03 : entanglement=5.94e-06 : last 10 iters in 1.089s
    iter 10700 : loss=4.691e-03 : entanglement=1.87e-05 : last 10 iters in 1.063s
    iter 10800 : loss=4.010e-03 : entanglement=2.17e-06 : last 10 iters in 1.294s
    iter 10900 : loss=4.046e-03 : entanglement=1.81e-06 : last 10 iters in 0.976s
    iter 11000 : loss=4.244e-03 : entanglement=2.73e-06 : last 10 iters in 0.975s
    iter 11100 : loss=3.734e-03 : entanglement=1.48e-06 : last 10 iters in 0.960s
    iter 11200 : loss=3.819e-03 : entanglement=2.85e-06 : last 10 iters in 1.088s
    iter 11300 : loss=4.627e-03 : entanglement=6.61e-06 : last 10 iters in 1.225s
    iter 11400 : loss=4.015e-03 : entanglement=5.34e-06 : last 10 iters in 1.045s
    iter 11500 : loss=4.027e-03 : entanglement=2.13e-06 : last 10 iters in 1.079s
    iter 11600 : loss=4.084e-03 : entanglement=2.79e-06 : last 10 iters in 1.052s
    iter 11700 : loss=4.131e-03 : entanglement=5.95e-06 : last 10 iters in 1.050s
    iter 11800 : loss=4.189e-03 : entanglement=6.10e-06 : last 10 iters in 1.098s
    iter 11900 : loss=3.964e-03 : entanglement=2.41e-06 : last 10 iters in 0.956s
    iter 12000 : loss=4.004e-03 : entanglement=3.22e-06 : last 10 iters in 0.976s
    iter 12100 : loss=3.814e-03 : entanglement=4.90e-06 : last 10 iters in 1.057s
    iter 12200 : loss=3.990e-03 : entanglement=1.68e-06 : last 10 iters in 0.970s
    iter 12300 : loss=4.068e-03 : entanglement=1.45e-06 : last 10 iters in 1.336s
    iter 12400 : loss=5.012e-03 : entanglement=1.30e-05 : last 10 iters in 1.169s
    iter 12500 : loss=4.852e-03 : entanglement=8.47e-06 : last 10 iters in 1.359s
    iter 12600 : loss=4.157e-03 : entanglement=2.17e-06 : last 10 iters in 1.077s
    iter 12700 : loss=3.950e-03 : entanglement=2.74e-06 : last 10 iters in 1.314s
    iter 12800 : loss=4.048e-03 : entanglement=3.68e-06 : last 10 iters in 1.124s
    iter 12900 : loss=3.875e-03 : entanglement=2.11e-06 : last 10 iters in 1.159s
    iter 13000 : loss=4.225e-03 : entanglement=2.76e-06 : last 10 iters in 1.127s
    iter 13100 : loss=4.599e-03 : entanglement=8.96e-06 : last 10 iters in 1.164s
    iter 13200 : loss=4.326e-03 : entanglement=2.19e-06 : last 10 iters in 0.956s
    iter 13300 : loss=4.023e-03 : entanglement=5.88e-06 : last 10 iters in 0.959s
    iter 13400 : loss=4.102e-03 : entanglement=2.75e-06 : last 10 iters in 0.934s
    iter 13500 : loss=3.915e-03 : entanglement=3.60e-06 : last 10 iters in 1.072s
    iter 13600 : loss=3.797e-03 : entanglement=1.83e-06 : last 10 iters in 0.941s
    iter 13700 : loss=3.552e-03 : entanglement=1.35e-06 : last 10 iters in 0.954s
    iter 13800 : loss=3.968e-03 : entanglement=1.27e-06 : last 10 iters in 1.095s
    iter 13900 : loss=4.242e-03 : entanglement=3.15e-06 : last 10 iters in 0.943s
    iter 14000 : loss=4.400e-03 : entanglement=3.52e-06 : last 10 iters in 0.916s
    iter 14100 : loss=4.055e-03 : entanglement=2.65e-06 : last 10 iters in 1.399s
    iter 14200 : loss=3.895e-03 : entanglement=4.00e-06 : last 10 iters in 0.955s
    iter 14300 : loss=3.775e-03 : entanglement=1.37e-06 : last 10 iters in 1.020s
    iter 14400 : loss=4.202e-03 : entanglement=2.32e-06 : last 10 iters in 1.048s
    iter 14500 : loss=5.215e-03 : entanglement=1.08e-05 : last 10 iters in 0.988s
    iter 14600 : loss=4.252e-03 : entanglement=9.86e-06 : last 10 iters in 1.344s
    iter 14700 : loss=4.185e-03 : entanglement=6.25e-06 : last 10 iters in 1.079s
    iter 14800 : loss=3.908e-03 : entanglement=2.49e-06 : last 10 iters in 0.972s
    iter 14900 : loss=4.084e-03 : entanglement=1.15e-06 : last 10 iters in 0.982s
    iter 15000 : loss=3.819e-03 : entanglement=1.60e-06 : last 10 iters in 1.087s

```python
exp_name = "ent_reg_1e-2_orthNorm_reg_1e-2_20steps_16batch"
# 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)
```

    dir already exists:  _data/decoder_only/ent_reg_1e-2_orthNorm_reg_1e-2_20steps_16batch
    dir already exists:  _data/decoder_only/ent_reg_1e-2_orthNorm_reg_1e-2_20steps_16batch/train
    dir already exists:  _data/decoder_only/ent_reg_1e-2_orthNorm_reg_1e-2_20steps_16batch/test

**Show training curve**

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

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 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 reconstruction 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

**1) Show action representations**

```python
width=0.5

rep_thetas = [rep.thetas.detach().numpy() for rep in lat_env.action_reps]
print(rep_thetas)
plt_lim = max( 0.22, max([max(t) for t in rep_thetas])/(2*np.pi) )
titles = ["transformation +", "transformation -"]
cols=["r","b","g","black"]
labels=["x","y","z","color"]

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

    fig, axs = plt.subplots(1, 2, figsize=(12, 4), gridspec_kw={"wspace":0.5, "hspace":0.5})

    for i, thetas in enumerate(rep_thetas):
        x = np.arange(len(thetas))
        axs[i%2].bar(x - width/2, thetas/(2*np.pi), width, label=labels[i//2], color=cols[i//2])
    for i in range(2):
        axs[i].hlines([0.2,-0.2], xmin=-10, xmax=10, linestyles="dashed")
        axs[i].hlines(0., xmin=-10, xmax=10)
        axs[i].set_yticks([-0.2, 0., 0.2])
        axs[i].set_xticks(x-0.25)
        axs[i].set_xticklabels(["12","13","14","15","23","24","25","34","35","45"], fontsize=10)
        axs[i].set_xlabel("$ij$", fontsize=20)

        axs[i].set_ylim(-plt_lim,plt_lim)
        axs[i].set_xlim(-0.5,9)
        axs[i].set_title(titles[i], fontsize=20)
        axs[i].tick_params(labelsize=15)

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

    axs[1].legend(loc="center right", bbox_to_anchor=(1.5,0.5), 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')
```

**2) Show latent space**

Note: this isn't particularly informative, but the general spherical shape of the latent space is visible

```python
from sklearn.random_projection import GaussianRandomProjection

latent_points = []
radius = 4.

for i in range(20):
    for j in range(20):
        obs = obs_env.reset()
        latent_points.append(encoder(obs).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']
marks=["+","1","2","3","4",">","<","^","v","x"]

for i in range (400):
    plt.scatter(latent_map[i][0],
                latent_map[i][1],
                c=color[i//40],
                marker=marks[i%10])
plt.plot()

```

**3) Show predictions**

```python
def plot_state(obs, ax):
    ax.pcolormesh(obs.reshape(2*obs_env.size, 2*obs_env.size, 2*obs_env.size).sum(axis=2), edgecolors='gray', linewidth=2)
    ax.set_aspect('equal')
    ax.set_xticks([])
    ax.set_yticks([])

    return ax

n_steps = 50

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, angle)
        obs_z = lat_env.step(action)

    obs_x_recon = decoder(obs_z)

    fig.suptitle('step {} : last action = {} in the x-y plane'.format(i, action), fontsize=16)

    plot_state(obs_x.detach().numpy(),ax1)
    plot_state(obs_x_recon.detach().numpy(),ax2)

    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(.2)

display.clear_output(wait=False)
```

```python

```


```python
from IPython.display import Image
Image('../../../python_for_probability_statistics_and_machine_learning.jpg')
```

[Python for Probability, Statistics, and Machine Learning](https://www.springer.com/fr/book/9783319307152)

```python
from __future__ import division
%pylab  inline
```

    Populating the interactive namespace from numpy and matplotlib

# Worked Examples of Conditional Expectation and Mean Square Error Optimization

Brzezniak [[brzezniak1999basic]](#brzezniak1999basic) is a great book because it approaches
conditional expectation through a sequence of exercises, which is what we are
trying to do here. The main difference is that Brzezniak takes a more abstract
measure-theoretic approach to the same problems. Note that you *do* need to
grasp  measure theory for advanced areas in  probability, but for what we have
covered so far, working the same problems in his text using our methods is
illuminating. It always helps to have more than one way to solve *any* problem.
I urge you to get a copy of his book or at least look at some pages on Google
Books. I have numbered the examples corresponding to the book and tried to
follow its notation.

## Example

This is Example 2.1 from Brzezniak.  Three coins, 10p, 20p and 50p are tossed.
The values of those coins that land heads up are totaled. What is the expected
total given that two coins have landed heads up?  In this case we have
we want to compute $\mathbb{E}(\xi|\eta)$ where

$$
\xi := 10 X_{10} + 20 X_{20} +50 X_{50}
$$

 where $X_i \in \{0,1\} $ and where $X_{10}$ is the
Bernoulli-distributed random variable corresponding to the 10p coin (and so
on). Thus, $\xi$ represents the total value of the heads-up coins. The $\eta$
represents the condition that only two of the three coins are heads-up,

$$
\eta := X_{10} X_{20} (1-X_{50})+ (1-X_{10}) X_{20} X_{50}+ X_{10} (1-X_{20}) X_{50}
$$

 and is a function that is non-zero *only* when two of the three coins lands
heads-up. Each triple term catches each of these three possibilities. For example,
the first term equals one when the 10p and 20p are heads up  and the 50p is
heads down. The  the remaining terms are zero.

To compute the conditional expectation, we want to find a function $h$ of
$\eta$ that minimizes the mean-squared-error (MSE),

$$
\mbox{MSE}= \sum_{X\in\{0,1\}^3} \frac{1}{2^3} (\xi-h(\eta))^2
$$

 where the sum is taken over all possible triples of outcomes for
$\{X_{10},X_{20} ,X_{50}\}$ because each
of the three coins has a $\frac{1}{2}$ chance of coming up heads.

Now, the question boils down to how can we characterize the function $h(\eta)$?
Note that $\eta \mapsto \{0,1\}$ so $h$ takes on only two values.  So, the
orthogonal inner product condition is the following:

$$
\langle \xi -h(\eta), \eta \rangle = 0
$$

 But, because are only interested in $\eta=1$, this simplifies to

$$
\begin{align*}
\langle \xi -h(1), 1 \rangle &= 0  \\\
\langle \xi,1  \rangle &=\langle h(1),1  \rangle
\end{align*}
$$

 This doesn't look so hard to evaluate but we have to compute the
integral over the set where $\eta=1$. In other words, we need the set of
triples $\{X_{10},X_{20},X_{50}\}$ where $\eta=1$. That is, we can
compute

$$
\int_{\{\eta=1\}} \xi dX = h(1)  \int_{\{\eta=1\}} dX
$$

  which is what Brzezniak does. Instead, we can define
$h(\eta)=\alpha \eta$ and then find $\alpha$. Re-writing the
orthogonal condition gives

$$
\begin{align*}
\langle \xi -\eta, \alpha\eta \rangle &= 0  \\\
\langle \xi, \eta \rangle &= \alpha \langle \eta,\eta \rangle  \\\
\alpha &= \frac{\langle \xi, \eta \rangle}{\langle \eta,\eta \rangle}
\end{align*}
$$

 where

$$
\langle \xi, \eta \rangle =\sum_{X\in\{0,1\}^3} \frac{1}{2^3}(\xi\eta)
$$

 Note that we can just sweep over all triples
$\{X_{10},X_{20},X_{50}\}$ because the definition of $h(\eta)$ zeros out when
$\eta=0$ anyway.  All we have to do is plug everything in and solve. This
tedious job is perfect for Sympy.

```python
import sympy as S
X10,X20,X50 = S.symbols('X10,X20,X50',real=True)
xi  = 10*X10+20*X20+50*X50
eta = X10*X20*(1-X50)+X10*(1-X20)*(X50)+(1-X10)*X20*(X50)
num=S.summation(xi*eta,(X10,0,1),(X20,0,1),(X50,0,1))
den=S.summation(eta*eta,(X10,0,1),(X20,0,1),(X50,0,1))
alpha=num/den
print alpha # alpha=160/3
```

    160/3

  This means that

$$
\mathbb{E}(\xi|\eta) = \frac{160}{3} \eta
$$

 which we can check with a quick simulation

```python
>>> import pandas as pd
>>> d = pd.DataFrame(columns=['X10','X20','X50'])
>>> d.X10 = np.random.randint(0,2,1000)
>>> d.X10 = np.random.randint(0,2,1000)
>>> d.X20 = np.random.randint(0,2,1000)
>>> d.X50 = np.random.randint(0,2,1000)
```

**Programming Tip.**

The code above creates an empty Pandas data frame with the named columns.
The next four lines assigns values to each of the columns.

 The code above simulates flipping the three coins 1000
times. Each column of the dataframe is either `0` or `1`
corresponding to heads-down or heads-up, respectively. The
condition is that two of the three coins have landed heads-up.
Next, we can group the columns according to their sums. Note that
the sum can only be in $\{0,1,2,3\}$ corresponding to `0`
heads-up, `1` heads-up, and so on.

```python
>>> grp=d.groupby(d.eval('X10+X20+X50'))
```

**Programming Tip.**

The `eval` function of the Pandas data frame takes the named
columns and evaluates the given formula. At the time of this
writing, only simple formulas involving primitive operations are
possible.

 Next, we can get the `2` group, which corresponds to
exactly two coins having landed heads up and then evaluate
the sum of the values of the coins. Finally, we can take the mean
of these sums.

```python
>>> grp.get_group(2).eval('10*X10+20*X20+50*X50').mean()
52.719999999999999
```

    52.72

 The result is close to `160/3=53.33` which supports
the analytic result. The following code shows that we
can accomplish the same simulation using pure Numpy.

```python
>>> import numpy as np
>>> from numpy import array
>>> x=np.random.randint(0,2,(3,1000))
>>> print np.dot(x[:,x.sum(axis=0)==2].T,array([10,20,50])).mean()
52.698998418555611
```

    54.3817204301

    52.69899841855561

 In this case, we used the Numpy dot product to compute
the value of the heads-up coins. The `sum(axis=0)==2` part selects
the columns that correspond to two heads-up coins.

Still another way to get at the same problem is to forego the
random sampling part and just consider all possibilities
exhaustively using the `itertools` module in Python's standard
library.

```python
>>> import itertools as it
>>> list(it.product((0,1),(0,1),(0,1)))
[(0, 0, 0),
 (0, 0, 1),
 (0, 1, 0),
 (0, 1, 1),
 (1, 0, 0),
 (1, 0, 1),
 (1, 1, 0),
 (1, 1, 1)]
```

    [(0, 0, 0),
     (0, 0, 1),
     (0, 1, 0),
     (0, 1, 1),
     (1, 0, 0),
     (1, 0, 1),
     (1, 1, 0),
     (1, 1, 1)]

 Note that we need to call `list` above in order to trigger the
iteration in `it.product`. This is because the `itertools` module is
generator-based so does not actually *do* the iteration until it is iterated
over (by `list` in this case).   This shows all possible triples
$(X_{10},X_{20},X_{50})$ where `0` and `1` indicate heads-down and heads-up,
respectively. The next step is to filter out the cases that correspond to two
heads-up coins.

```python
>>> list(it.ifilter(lambda i:sum(i)==2,it.product((0,1),(0,1),(0,1))))
[(0, 1, 1), (1, 0, 1), (1, 1, 0)]
```

    [(0, 1, 1), (1, 0, 1), (1, 1, 0)]

 Next, we need to compute the sum of the coins and combine
the prior code.

```python
>>> map(lambda k:10*k[0]+20*k[1]+50*k[2],
...                   it.ifilter(lambda i:sum(i)==2,
...                              it.product((0,1),(0,1),(0,1))))
[70, 60, 30]
```

    [70, 60, 30]

 The mean of the output is `53.33`, which is yet another way to get
the same result. For this example, we demonstrated the full spectrum of
approaches made possible using Sympy, Numpy, and Pandas. It is always valuable
to have multiple ways of approaching the same problem and cross-checking
the result.

## Example

This is Example 2.2 from Brzezniak. Three coins, 10p, 20p and 50p are tossed as
before. What is the conditional expectation of the total amount shown by the
three coins given the total amount shown by the 10p and 20p coins only?  For
this problem,

$$
\begin{align*}
\xi := & 10 X_{10} + 20 X_{20} +50 X_{50}  \\\
\eta :=& 30 X_{10} X_{20} + 20 (1-X_{10}) X_{20} + 10 X_{10} (1-X_{20})
\end{align*}
$$

 which takes on four values $\eta \mapsto \{0,10,20,30\}$ and only
considers the 10p and 20p coins.  In contrast to the last problem, here we are
interested in $h(\eta)$ for all of the values of $\eta$. Naturally, there are
only four values for $h(\eta)$ corresponding to each of these four values.
Let's first consider $\eta=10$. The orthogonal condition is then

$$
\langle\xi-h(10),10\rangle = 0
$$

 The domain for $\eta=10$ is $\{X_{10}=1,X_{20}=0,X_{50}\}$ which we
can integrate out of the expectation below,

$$
\begin{align*}
\mathbb{E}_{\{X_{10}=1,X_{20}=0,X_{50}\}}(\xi-h(10)) 10  &=0 \\\
\mathbb{E}_{\{X_{50}\}}(10-h(10)+50 X_{50}) &=0  \\\
10-h(10) + 25 &=0
\end{align*}
$$

 which gives $h(10)=35$. Repeating the same process for $\eta \in
\{20,30\}$ gives $h(20)=45$ and $h(30)=55$, respectively.  This is the approach
Brzezniak takes. On the other hand, we can just look at affine functions,
$h(\eta) = a \eta + b $ and use brute-force calculus.

```python
>>> from sympy.abc import a,b
>>> h = a*eta + b
>>> eta = X10*X20*30 + X10*(1-X20)*(10)+ (1-X10)*X20*(20)
>>> MSE=S.summation((xi-h)**2*S.Rational(1,8),(X10,0,1),
...                 (X20,0,1),
...                 (X50,0,1))
>>> sol=S.solve([S.diff(MSE,a),S.diff(MSE,b)],(a,b))
>>> print sol
{b: 25, a: 1}
```

    {b: 32, a: 64/3}

    {b: 25, a: 1}

**Programming Tip.**

The `Rational` function from Sympy code expresses a rational number that Sympy
is able to manipulate as such. This is different that specifying a fraction
like `1/8.`, which Python would automatically compute as a floating point
number (i.e., `0.125`). The advantage of using `Rational` is that Sympy can
later produce rational numbers as output, which are sometimes easier to make
sense of.

  This means that

$$
\begin{equation}
\mathbb{E}(\xi|\eta) = 25+\eta
\label{_auto1} \tag{1}
\end{equation}
$$

  since $\eta$ takes on only four values, $\{0,10,20,30\}$, we can
write this out explicitly as

$$
\begin{equation}
\mathbb{E}(\xi|\eta) =
\begin{cases}
25 & \text{for}\: \eta=0  \\\
35 & \text{for}\: \eta=10 \\\
45 & \text{for}\: \eta=20 \\\
55 & \text{for}\: \eta=30
\end{cases}
\end{equation}
\label{eq:ex21sol} \tag{2}
$$

 Alternatively, we can use orthogonal inner products to write out
the following conditions:

$$
\begin{equation}
\label{eq:ex22a} \tag{3}
\langle \xi-h(\eta), \eta \rangle = 0
\end{equation}
$$

$$
\begin{equation}
\label{eq:ex22b} \tag{4}
\langle \xi-h(\eta),1\rangle = 0
\end{equation}
$$

 Writing these out and solving for $a$ and  $b$ is tedious and
a perfect job for Sympy. Starting with Equation ref{eq:ex22a},

```python
>>> expr = expr=S.expand((xi-h)*eta)
>>> print expr
-100*X10**2*a + 100*X10**2 - 400*X10*X20*a + 400*X10*X20 + 500*X10*X50
- 10*X10*b - 400*X20**2*a + 400*X20**2 + 1000*X20*X50 - 20*X20*b
```

    30*X10**2*X20*X50*a - 10*X10**2*X20*a - 10*X10**2*X50*a + 100*X10**2 + 60*X10*X20**2*X50*a - 20*X10*X20**2*a - 30*X10*X20*X50*a + 400*X10*X20 + 500*X10*X50 - 10*X10*b - 20*X20**2*X50*a + 400*X20**2 + 1000*X20*X50 - 20*X20*b

    -10*X10*b - 400*X20**2*a + 400*X20**2 + 1000*X20*X50 - 20*X20*b

 and then because $\mathbb{E}(X_i^2)=1/2=\mathbb{E}(X_i)$, we make the
following substitutions

```python
>>> expr.xreplace({X10**2:0.5, X20**2:0.5,X10:0.5,X20:0.5,X50:0.5})
-350.0*a - 15.0*b + 725.0
```

    -350.0*a - 15.0*b + 725.0

 We can do this for the other orthogonal inner product in Equation
ref{eq:ex22b} as follows,

**Programming Tip.**

Because Sympy symbols are hashable, they can be used as keys in Python
dictionaries as in the `xreplace` function above.

```python
>>> print S.expand((xi-h)*1).xreplace({X10**2:0.5,
...                                    X20**2:0.5,
...                                    X10:0.5,
...                                    X20:0.5,
...                                    X50:0.5})
-15.0*a - b + 40.0
```

    -0.375*a - b + 40.0

    -15.0*a - b + 40.0

 Then, combining this result with the previous one and solving
for `a` and `b` gives,

```python
>>> print S.solve([-350.0*a-15.0*b+725.0,-15.0*a-b+40.0])
{b: 25.0000000000000, a: 1.00000000000000}
```

    {b: 25.0000000000000, a: 1.00000000000000}

    {b: 25.0, a: 1.0}

 which again gives us the final solution,

$$
\mathbb{E}(\xi|\eta) = 25+ \eta
$$

  The following is a quick simulation to demonstrate this. We can
build on the Pandas dataframe we used for the last example and create
a new column for the sum of the 10p and 20p coins, as shown below.

```python
>>> d['sm'] = d.eval('X10*10+X20*20')
```

 We can group this by the values of this sum,

```python
>>> d.groupby('sm').mean()

```

      X10
      X20
      X50

      sm

      0
      0
      0
      0.464730

      10
      1
      0
      0.533597

      20
      0
      1
      0.509960

      30
      1
      1
      0.498039

 But we want the expectation of the value of the coins

```python
>>> d.groupby('sm').mean().eval('10*X10+20*X20+50*X50')

```

    sm
    0     23.236515
    10    36.679842
    20    45.498008
    30    54.901961
    dtype: float64

 which is very close to our analytical result in Equation ref{eq:ex21sol}.

## Example

This is Example 2.3 paraphrased from Brzezniak. Given $X$ uniformly distributed
on $[0,1]$, find $\mathbb{E}(\xi|\eta)$ where

$$
\xi(x) = 2 x^2
$$

$$
\eta(x) =
\begin{cases}
1 & \mbox{if } x \in [0,1/3]  \\\
2 & \mbox{if } x \in (1/3,2/3)  \\\
0 & \mbox{if } x \in (2/3,1]
\end{cases}
$$

 Note that this problem is different from the previous two because the
sets that characterize $\eta$ are intervals instead of discrete points.
Nonetheless, we will eventually have three values for $h(\eta)$ because $\eta
\mapsto \{0,1,2\}$. For $\eta=1$, we have the orthogonal conditions,

$$
\langle \xi-h(1),1\rangle = 0
$$

 which boils down to

$$
\mathbb{E}_{\{x \in [0,1/3]\}}(\xi-h(1))=0
$$

$$
\int_0^{\frac{1}{3}}(2 x^2-h(1))dx = 0
$$

 and then by solving this for $h(1)$ gives $h(1)=2/24$. This is the way
Brzezniak works this problem. Alternatively, we can use $h(\eta) =  a + b\eta
+ c\eta^2$ and brute force calculus.  Note the the `Piecewise` object in
`sympy` is not complete at this point in its development, so we'll have to be
exceptionally verbose in the following,

```python
x,c,b,a=S.symbols('x,c,b,a')
xi = 2*x**2

eta=S.Piecewise((1,S.And(S.Gt(x,0),
                         S.Lt(x,S.Rational(1,3)))),  #  0 < x < 1/3
                (2,S.And(S.Gt(x,S.Rational(1,3)),
                         S.Lt(x,S.Rational(2,3)))), # 1/3 < x < 2/3,
                (0,S.And(S.Gt(x,S.Rational(2,3)),
                         S.Lt(x,1)))) # 1/3 < x < 2/3
h = a + b*eta + c*eta**2
J=S.integrate((xi-h)**2,(x,0,1))
sol=S.solve([S.diff(J,a),
             S.diff(J,b),
             S.diff(J,c),
            ],
            (a,b,c))
```

```python
>>> print sol
```

    {c: 8/9, b: -20/9, a: 38/27}

```python
>>> print S.piecewise_fold(h.subs(sol))
```

    Piecewise((2/27, And(x < 1/3, x > 0)), (14/27, And(x < 2/3, x > 1/3)), (38/27, And(x < 1, x > 2/3)))

Piecewise((2/27,And(x<1/3,x>0)),
          (14/27,And(x<2/3,x>1/3)),(38/27,And(x<1,x>2/3)))
  Thus, collecting this result gives:

$$
\mathbb{E}(\xi|\eta) = \frac{38}{27} - \frac{20}{9}\eta + \frac{8}{9} \eta^2
$$

  which can be re-written as a piecewise function of x,

$$
\begin{equation}
\mathbb{E}(\xi|\eta(x)) =\begin{cases} \frac{2}{27} & \text{for}\: 0 < x < \frac{1}{3} \\\frac{14}{27} & \text{for}\: \frac{1}{3} < x < \frac{2}{3} \\\frac{38}{27} & \text{for}\: \frac{2}{3}<x < 1 \end{cases}
\end{equation}
\label{eq:ex23a} \tag{5}
$$

Alternatively, we can use the orthogonal inner product conditions directly by
choosing $h(\eta)=c+\eta b +\eta^2 a$,

$$
\begin{align*}
\langle \xi-h(\eta),1\rangle = 0 \\\
\langle \xi-h(\eta),\eta\rangle = 0 \\\
\langle \xi-h(\eta),\eta^2\rangle = 0
\end{align*}
\label{eq:ex23b} \tag{6}
$$

 and then solving for $a$,$b$, and $c$.

```python
>>> x,a,b,c,eta = S.symbols('x,a,b,c,eta',real=True)
>>> xi  = 2*x**2
>>> eta=S.Piecewise((1,S.And(S.Gt(x,0),
...                      S.Lt(x,S.Rational(1,3)))),  #  0 < x < 1/3
...             (2,S.And(S.Gt(x,S.Rational(1,3)),
...                      S.Lt(x,S.Rational(2,3)))), # 1/3 < x < 2/3,
...             (0,S.And(S.Gt(x,S.Rational(2,3)),
...                      S.Lt(x,1)))) # 1/3 < x < 2/3

>>> h = c+b*eta+a*eta**2
```

 Then, the orthogonal conditions become,

```python
>>> S.integrate((xi-h)*1,(x,0,1))
-5*a/3 - b - c + 2/3
>>> S.integrate((xi-h)*eta,(x,0,1))
-3*a - 5*b/3 - c + 10/27
>>> S.integrate((xi-h)*eta**2,(x,0,1))
-17*a/3 - 3*b - 5*c/3 + 58/81
```

    -17*a/3 - 3*b - 5*c/3 + 0.716049382716049

 Now, we just combine the three equations and solve
for the parameters,

```python
>>> eqs=[ -5*a/3 - b - c + 2/3,
...   -3*a - 5*b/3 - c + 10/27,
...   -17*a/3 - 3*b - 5*c/3 + 58/81]
>>> sol=S.solve(eqs)
>>> print sol
{a: 0.888888888888889, c: 1.40740740740741, b: -2.22222222222222}
```

    {a: 0.888888888888889, c: 1.40740740740741, b: -2.22222222222222}

    {a: 0.888888888888889, c: 1.40740740740741, b: -2.22222222222222}

 We can assemble the final result by substituting in the solution,

```python
>>> print S.piecewise_fold(h.subs(sol))
```

    Piecewise((0.0740740740740740, And(x < 1/3, x > 0)), (0.518518518518518, And(x < 2/3, x > 1/3)), (1.40740740740741, And(x < 1, x > 2/3)))

Piecewise((0.0740740740740740, And (x < 1/3, x > 0)),
           (0.518518518518518, And (x < 2/3, x > 1/3)),
            (1.40740740740741, And (x < 1, x > 2/3)))
 which is the same as our analytic result in Equation ref{eq:ex23a},
just in decimal format.

**Programming Tip.**

The definition of Sympy's piecewise function is verbose because of the way
Python parses inequality statements. As of this writing, this has not been
reconciled in Sympy, so we have to use the verbose declaration.

To reinforce our result, let's do a quick simulation using Pandas.

```python
>>> d = pd.DataFrame(columns=['x','eta','xi'])
>>> d.x = np.random.rand(1000)
>>> d.xi = 2*d.x**2
```

 Now, we can use the `pd.cut` function to group the `x`
values in the following,

```python
>>> pd.cut(d.x,[0,1/3,2/3,1]).head()
```

    0        (0.667, 1]
    1    (0.333, 0.667]
    2        (0, 0.333]
    3        (0, 0.333]
    4        (0, 0.333]
    Name: x, dtype: category
    Categories (3, object): [(0, 0.333] < (0.333, 0.667] < (0.667, 1]]

0        (0.667, 1]
1        (0, 0.333]
2        (0.667, 1]
3    (0.333, 0.667]
4    (0.333, 0.667]
Name: x, dtype: category
Categories (3, object): [(0, 0.333] < (0.333, 0.667] < (0.667, 1]]
 Note that the `head()` call above is only to limit the printout shown.
The categories listed are each of the intervals for `eta` that we specified
using the `[0,1/3,2/3,1]` list. Now that we know how to use `pd.cut`, we
can just compute the mean on each group as shown below,

```python
>>> d.groupby(pd.cut(d.x,[0,1/3,2/3,1])).mean()['xi']
```

    x
    (0, 0.333]        0.076390
    (0.333, 0.667]    0.521691
    (0.667, 1]        1.410679
    Name: xi, dtype: float64

(0, 0.333]        0.069240
(0.333, 0.667]    0.520154
(0.667, 1]        1.409747
Name: xi, dtype: float64
 which is pretty close to our analytic result in Equation
ref{eq:ex23a}. Alternatively, `sympy.stats` has some limited tools for the same
calculation.

```python
>>> from sympy.stats import E, Uniform
>>> x=Uniform('x',0,1)
>>> E(2*x**2,S.And(x < S.Rational(1,3), x > 0))
2/27
>>> E(2*x**2,S.And(x < S.Rational(2,3), x > S.Rational(1,3)))
14/27
>>> E(2*x**2,S.And(x < 1, x > S.Rational(2,3)))
38/27
```

    1.4074074074074074

 which again gives the same result still another way.

## Example

This is Example 2.4 from Brzezniak. Find $\mathbb{E}(\xi|\eta)$ for

$$
\xi(x) = 2 x^2
$$

$$
\eta =
\begin{cases}2 & \mbox{if } 0 \le x < \frac{1}{2} \\ x & \mbox{if } \frac{1}{2} < x \le 1 \end{cases}

\label{eq:ex24} \tag{7}
$$

 Once again, $X$ is uniformly distributed on the unit interval. Note
that $\eta$ is no longer discrete for every domain. For the domain $0 <x <
1/2$, $h(2)$ takes on only one value, say, $h_0$. For this domain, the
orthogonal condition becomes,

$$
\mathbb{E}_{\{\eta=2\}}((\xi(x)-h_0)2)=0
$$

 which simplifies to,

$$
\begin{align*}
\int_0^{1/2} 2 x^2-h_0 dx &= 0 \\\
\int_0^{1/2} 2 x^2 dx &= \int_0^{1/2} h_0 dx \\\
h_0 &= 2 \int_0^{1/2} 2 x^2 dx \\\
h_0 &= \frac{1}{6}
\end{align*}
$$

For the other domain where $\{\eta=x\}$ in Equation ref{eq:ex24}, we again
use the orthogonal condition,

$$
\begin{align*}
\mathbb{E}_{\{\eta=x\}}((\xi(x)-h(x))x)&=0 \\\
\int_{1/2}^1 (2x^2-h(x)) x dx &=0 \\\
h(x) &= 2x^2
\end{align*}
$$

 Assembling the solution gives,

$$
\mathbb{E}(\xi|\eta(x)) =\begin{cases} \frac{1}{6} & \text{for}\: 0 \le x < \frac{1}{2} \\ 2 x^2 & \text{for}\: \frac{1}{2} < x \le 1 \end{cases}
$$

 although  this result is not explicitly written as a function of $\eta$.

## Example

This is Exercise 2.6 in Brzezniak. Find $\mathbb{E}(\xi|\eta)$ where

$$
\xi(x) = 2 x^2
$$

$$
\eta(x) = 1 - \lvert 2 x-1 \rvert
$$

 and $X$ is uniformly distributed in the unit interval. We
can write this out as a piecewise function in the following,

$$
\eta =\begin{cases} 2 x & \text{for}\: 0 \le x < \frac{1}{2} \\ 2 -2x & \text{for}\: \frac{1}{2} < x \le 1 \end{cases}
$$

 The discontinuity is at $x=1/2$. Let's start with the $\{\eta=2x\}$ domain.

$$
\begin{align*}
\mathbb{E}_{\{\eta=2x\}}((2 x^2-h(2 x)) 2 x)& = 0 \\\
\int_{0}^{1/2} (2x^2-h(2 x) ) 2 x dx &=0
\end{align*}
$$

 We can make this explicitly a function of $\eta$ by a change
of variables ($\eta=2x$) which gives

$$
\int_{0}^{1} (\eta^2/2-h(\eta))\frac{\eta}{2} d\eta =0
$$

 Thus, for this domain, $h(\eta)=\eta^2/2$. Note that due to the
change of variables, $h(\eta)$ is valid defined over $\eta\in[0,1]$.

For the other domain where $\{\eta=2-2x\}$, we have

$$
\begin{align*}
\mathbb{E}_{\{\eta=2-2x\}}((2 x^2-h(2-2x)) (2-2x))& = 0 \\\
\int_{1/2}^{1} (2 x^2-h(2-2x) ) (2-2x) dx &=0
\end{align*}
$$

 Once again, a change of variables makes the $ \eta$ dependency
explicit using $\eta=2-2x$ which gives

$$
\begin{align*}
\int_{0}^{1} ((2-\eta)^2/2-h(\eta) ) \frac{\eta}{2} d\eta &=0 \\\
h(\eta) &= (2-\eta)^2/2
\end{align*}
$$

 Once again, the change of variables means this solution is valid
over $\eta\in[0,1]$. Thus, because both pieces are valid over the
same domain ($\eta\in[0,1]$), we can just add them to get the final solution,

$$
h(\eta) = \eta^2-2\eta+2
$$

 A quick simulation can help bear this out.

```python
from pandas import DataFrame
import numpy as np
d = DataFrame(columns=['xi','eta','x','h','h1','h2'])
# 100 random samples
d.x = np.random.rand(100)
d.xi = d.eval('2*x**2')
d.eta =1-abs(2*d.x-1)
d.h1=d[(d.x<0.5)].eval('eta**2/2')
d.h2=d[(d.x>=0.5)].eval('(2-eta)**2/2')
d.fillna(0,inplace=True)
d.h = d.h1+d.h2
d.head()
```

      xi
      eta
      x
      h
      h1
      h2

      0
      0.346459
      0.832416
      0.416208
      0.346459
      0.346459
      0.000000

      1
      0.432610
      0.930172
      0.465086
      0.432610
      0.432610
      0.000000

      2
      0.001143
      0.047811
      0.023905
      0.001143
      0.001143
      0.000000

      3
      0.096686
      0.439741
      0.219870
      0.096686
      0.096686
      0.000000

      4
      1.679429
      0.167281
      0.916359
      1.679429
      0.000000
      1.679429

 Note that we have to be careful where we apply the individual
solutions using the slice `(d.x<0.5)` index. The `fillna` part ensures that the
default `NaN` that fills out the empty row-etries is replaced with zero before
combining the individual solutions. Otherwise, the `NaN` values would circulate
through the rest of the computation. The following is the
essential code that draws [Figure](#fig:Conditional_expectation_MSE_005).

```python
from matplotlib.pyplot import subplots
fig,ax=subplots()
ax.plot(d.xi,d.eta,'.',alpha=.3,label='$\eta$')
ax.plot(d.xi,d.h,'k.',label='$h(\eta)$')
ax.legend(loc=0,fontsize=18)
ax.set_xlabel('$2 x^2$',fontsize=18)
ax.set_ylabel('$h(\eta)$',fontsize=18)
```

**Programming Tip.**

Basic \LaTeX{} formatting works for the labels in [Figure](#fig:Conditional_expectation_MSE_005).  The `loc=0` in the `legend`
function is the code for the *best* placement for the labels in the legend. The
individual labels should be specified when the elements are drawn individually,
otherwise they will be hard to separate out later. This is accomplished using
the `label` keyword in the `plot` commands.

```python
from matplotlib.pyplot import subplots
from pandas import DataFrame
import numpy as np

d = DataFrame(columns=['xi','eta','x','h','h1','h2'])
# 100 random samples
d.x = np.random.rand(100)
d.xi = d.eval('2*x**2')
d.eta =1-abs(2*d.x-1)
d.h1=d[(d.x<0.5)].eval('eta**2/2')
d.h2=d[(d.x>=0.5)].eval('(2-eta)**2/2')
d.fillna(0,inplace=True)
d.h = d.h1+d.h2

fig,ax=subplots()
_=ax.plot(d.xi,d.eta,'.k',alpha=.3,label=r'$\eta$')
_=ax.plot(d.xi,d.h,'ks',label=r'$h(\eta)$',alpha=.3)
_=ax.set_aspect(1)
_=ax.legend(loc=0,fontsize=18)
_=ax.set_xlabel(r'$\xi=2 x^2$',fontsize=24)
_=ax.set_ylabel(r'$h(\eta),\eta$',fontsize=24)
fig.tight_layout()
#fig.savefig('fig-probability/Conditional_expectation_MSE_Ex_005.png')
```

The diagonal line shows where the conditional expectation equals the $\xi$ function.

[Figure](#fig:Conditional_expectation_MSE_005) shows the $\xi$ data plotted
against $\eta$ and $h(\eta) = \mathbb{E}(\xi|\eta)$. Points on the diagonal
are points where $\xi$ and $\mathbb{E}(\xi|\eta)$ match. As shown by the
dots, there is no agreement between the raw $\eta$ data and $\xi$. Thus, one
way to think about the conditional expectation is as a functional transform
that bends the curve onto the diagonal line. The black dots plot $\xi$
versus $\mathbb{E}(\xi|\eta)$ and the two match everywhere along the diagonal
line. This is to be expected because the conditional expectation is the MSE
best estimate for $\xi$ among all functions of $\eta$.

## Example

This is Exercise 2.14 from Brzezniak. Find $\mathbb{E}(\xi|\eta)$ where

$$
\xi(x) = 2 x^2
$$

$$
\eta =
\begin{cases} 2x & \mbox{if } 0 \le x < \frac{1}{2} \\ 2x-1 & \mbox{if } \frac{1}{2} < x \le 1 \end{cases}
$$

 and $X$ is uniformly distributed in the unit interval.  This is the
same as the last example and the only difference here is that $\eta$ is not
continuous at $x=\frac{1}{2}$, as before. The first part is exactly the same as
the first part of the prior example so we will skip it here. The second part
follows the same reasoning as the last example, so we will just write the
answer for the $\{\eta = 2x-1\}$ case as the following

$$
h(\eta)=\frac{(1+\eta)^2}{2} , \: \forall \eta \: \in [0,1]
$$

  and then adding these up as before gives the full solution:

$$
h(\eta)= \frac{1}{2} +\eta + \eta^2
$$

 The interesting part about this example is shown in [Figure](#fig:Conditional_expectation_MSE_006). The dots show where $\eta$ is
discontinuous and yet the $h(\eta)=\mathbb{E}(\xi|\eta)$ solution is equal to
$\xi$ (i.e., matches the diagonal). This illustrates the power of the orthogonal
inner product technique, which does not need continuity or complex
set-theoretic arguments to calculate solutions. By contrast, I urge you to
consider Brzezniak's solution to this problem which requires such  methods.

```python
d = DataFrame(columns=['xi','eta','x','h','h1','h2'])
d.x = np.random.rand(100) # 100 random samples
d.xi = d.eval('2*x**2')
d['eta']=(d.x<0.5)*(2*d.x)+(d.x>=0.5)*(2*d.x-1)
d.h1=d[(d.x<0.5)].eval('eta**2/2')
d.h2=d[(d.x>=0.5)].eval('(1+eta)**2/2')
d.fillna(0,inplace=True)
d.h = d.h1+d.h2
fig,ax=subplots()
_=ax.plot(d.xi,d.eta,'.k',alpha=.3,label='$\eta$')
_=ax.plot(d.xi,d.h,'ks',label='$h(\eta)$',alpha=0.3)
_=ax.set_aspect(1)
_=ax.legend(loc=0,fontsize=18)
_=ax.set_xlabel('$2 x^2$',fontsize=24)
_=ax.set_ylabel('$h(\eta),\eta$',fontsize=24)
fig.tight_layout()
#fig.savefig('fig-probability/Conditional_expectation_MSE_Ex_006.png')
```

The diagonal line shows where the conditional expectation equals the $\xi$ function.

Extending projection methods to random variables provides multiple ways for
calculating solutions to conditional expectation problems. In this section, we
also worked out corresponding simulations using a variety of Python modules. It
is always advisable to have more than one technique at hand to cross-check
potential solutions. We worked out some of the examples in Brzezniak's
book using our methods as a way to show multiple ways to solve the same
problem. Comparing Brzezniak's measure-theoretic methods to our less abstract
techniques is a great way to get a handle on both concepts, which are important
for advanced study in stochastic process.


# The Space-Arrow of Space-time

People have claimed Nature does not have an arrow for time. I don't think the question, as stated, is well-formed. Any analysis of the arrow of time for one observer will look like the arrow of space-time to another one moving relative to the first.

Two different problems are often cited on this subject. The first is that the fundamental forces of the standard model - EM, the weak force, and the strong force - are unaltered by a change in the arrow of (space-)time. Second, that the Lorentz group can be used an arbitrary number of times and still flip the value of time precisely.

In this short talk, I will show how keeping the space terms next to the time ones leads to the world we know, where space-time reversal does not come easy.

## Space-time Reversal in EM

Let me provide a quick sketch of how to derive the Maxwell source equations, not just write them down like Minkowski did. The goal is to calculate the Lorentz invariant quantity $B^2 - E^2$, the difference of two squares. Based on high school algebra, this should be obtained by the product of the sum and difference of $E$ and $B$. Once one has this difference of squares, the Euler-Lagrange equations to derive Gauss's and Ampere's laws.

Start all the ways a 4-potential can change in space-time (doesn't that sound general?):

The first term is a gauge field. Since photons are being studied, this gauge field has to be set to zero. For anything that is not a photon, this field will be non-zero. It is interesting to me that it is just "naturally" here. It can be zeroed out easily enough:

By changing the order of the differential acting on the potential, only the magnetic field $B$ will change signs.

Form the product of these two:

The first observation to make that everyone makes is that reversing time will never, ever change the value of the difference between the square of the magnetic field and the square of the electric field. That is what a square does. This is why one can conclude with confidence that the fundamental forces in physics are unaltered by changes in time. The weak and strong forces are mere variations on EM using other gauge groups.

The second observation that no one makes is that the Poynting vector is sitting right next door. The Poynting vector will change signs under time reversal. A complete analysis of EM must include the part that responds to time reversal. By omitting it, a mystery is claimed.

## Space-time Reversal is Locally Irreversible

In this notebook, I will show two different ways how when one thinks about space-time reversal instead of just time reversal using space-time numbers, that such a system cannot be reversed an arbitrarily large number of times as is the case for the Lorentz group. First the tools need to work with space-time numbers have to loaded.

```python
%%capture
%matplotlib inline
import numpy as np
import sympy as sp
import matplotlib.pyplot as plt

# To get equations the look like, well, equations, use the following.
from sympy.interactive import printing
printing.init_printing(use_latex=True)
from IPython.display import display

# Tools for manipulating quaternions.
import Q_tools as qt;
```

The member of the Lorentz group that reverses time is remarkably simple: it is a matrix that has minus one in the upper diagonal position, positive one for the other diagonal positions, and zeros elsewhere.

```python
TimeReversal = np.array([[-1,0,0,0],[0,1,0,0],[0,0,1,0],[0,0,0,1]])
display(TimeReversal)
```

    array([[-1,  0,  0,  0],
           [ 0,  1,  0,  0],
           [ 0,  0,  1,  0],
           [ 0,  0,  0,  1]])

Create a 4-vector.

```python
t, x, y, z = sp.symbols("t x y z")
```

```python
Vector = np.array([t, x, y, z])
display(Vector)
```

    array([t, x, y, z], dtype=object)

```python
display(Vector * TimeReversal)
```

    array([[-t, 0, 0, 0],
           [0, x, 0, 0],
           [0, 0, y, 0],
           [0, 0, 0, z]], dtype=object)

Do the time reversal a bunch of times.

```python
display(Vector * TimeReversal * TimeReversal)
display(Vector * TimeReversal * TimeReversal * TimeReversal)
```

    array([[t, 0, 0, 0],
           [0, x, 0, 0],
           [0, 0, y, 0],
           [0, 0, 0, z]], dtype=object)

    array([[-t, 0, 0, 0],
           [0, x, 0, 0],
           [0, 0, y, 0],
           [0, 0, 0, z]], dtype=object)

It comes as no surprise that if the time-reversal member of the Lorentz group is done an even number of times, then nothing changes, but an odd number of times reverses time exactly.

For the quaternion approach, one solves a pretty simple algebraic expression instead of using a global matrix. The equation to solve is:
$$ P T_r = - P^* $$
Solve for $T_r$:
$$ T_r = - P^{-1} P^* $$

```python
P = qt.QH([t, x, y, z])
Tr = P.flip_signs().invert().product(P.conj())
display(Tr.t)
display(Tr.x)
display(Tr.y)
display(Tr.z)
```

Does this considerably more complicated expression than the Lorentz group do its job? Of course is **should**, but let's just show this is the case:

```python
PFlip = P.product(Tr)
display(sp.simplify(PFlip.t))
display(sp.simplify(PFlip.x))
display(sp.simplify(PFlip.y))
display(sp.simplify(PFlip.z))
```

Apply Tr twice to see if one gets back to the start point.

```python
PFlipFlip = P.product(Tr).product(Tr)
display(sp.simplify(PFlipFlip.t))
display(sp.simplify(PFlipFlip.x))
display(sp.simplify(PFlipFlip.y))
display(sp.simplify(PFlipFlip.z))
```

This may not look "right" to the eye, so test it. Use "classical" values, meaning time $t >>> x, y, z$.

```python
Classical_subs = {t:1, x:0.0000000002, y:0.00000000012, z:-0.0000000003}
display(sp.simplify(PFlip.t.subs(Classical_subs)))
display(sp.simplify(PFlip.x.subs(Classical_subs)))
display(sp.simplify(PFlip.y.subs(Classical_subs)))
display(sp.simplify(PFlip.z.subs(Classical_subs)))
print()
display(sp.simplify(PFlipFlip.t.subs(Classical_subs)))
display(sp.simplify(PFlipFlip.x.subs(Classical_subs)))
display(sp.simplify(PFlipFlip.y.subs(Classical_subs)))
display(sp.simplify(PFlipFlip.z.subs(Classical_subs)))
```

The value for t returned to unity as it should, but the same cannot be said for the spatial terms. This is due to the cross product. See what happens if one does this many, many times. Define a function to do the work.

```python
def reverse_n_times(P1, T1, sub_1, n):
    """Given a symbolic expression P, applies symbolic space-time reversal using a dictionary of values n times."""

    P1_t, P1_x, P1_y, P1_z = P1.t.subs(sub_1), P1.x.subs(sub_1), P1.y.subs(sub_1), P1.z.subs(sub_1)
    P_result = qt.QH([P1_t, P1_x, P1_y, P1_z])

    T1_t, T1_x, T1_y, T1_z = T1.t.subs(sub_1), T1.x.subs(sub_1), T1.y.subs(sub_1), T1.z.subs(sub_1)
    T_sub = qt.QH([T1_t, T1_x, T1_y, T1_z])

    for i in range(n):
        P_result = P_result.product(T_sub)

    return P_result
```

```python
print(reverse_n_times(P, Tr, Classical_subs, 100))
print(reverse_n_times(P, Tr, Classical_subs, 101))
print(reverse_n_times(P, Tr, Classical_subs, 1000))
print(reverse_n_times(P, Tr, Classical_subs, 1001))
```

    (1.00000000000000, -3.97999999999999E-8, -2.38800000000001E-8, 5.96999999999999E-8) ...
    (-1.00000000000000, 4.01999999999999E-8, 2.41200000000001E-8, -6.02999999999999E-8) ...
    (0.999999999999712, -3.99799999999961E-7, -2.39879999999978E-7, 5.99699999999939E-7) ...
    (-0.999999999999712, 4.00199999999961E-7, 2.40119999999978E-7, -6.00299999999938E-7) ...

```python
print(reverse_n_times(P, Tr, Classical_subs, 10000))
print(reverse_n_times(P, Tr, Classical_subs, 100000))
```

    (0.999999999971126, -0.00000399979999996151, -0.00000239987999997690, 0.00000599969999994228) ...
    (0.999999997112059, -0.0000399997999614952, -0.0000239998799768971, 0.0000599996999422425) ...

```python
print(reverse_n_times(P, Tr, Classical_subs, 1000000))
print(reverse_n_times(P, Tr, Classical_subs, 10000000))
```

    (0.999999711200594, -0.000399999761493507, -0.000239999856896105, 0.000599999642240260) ...


# Neuromatch Academy: Week 1, Day 5, Tutorial 3
# Dimensionality Reduction and reconstruction

__Content creators:__ Alex Cayco Gajic, John Murray

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

---
# Tutorial Objectives

In this notebook we'll learn to apply PCA for dimensionality reduction, using a classic dataset that is often used to benchmark machine learning algorithms: MNIST. We'll also learn how to use PCA for reconstruction and denoising.

Overview:
- Perform PCA on MNIST
- Calculate the variance explained
- Reconstruct data with different numbers of PCs
- (Bonus) Examine denoising using PCA

You can learn more about MNIST dataset [here](https://en.wikipedia.org/wiki/MNIST_database).

```python
# @title Video 1: PCA for dimensionality reduction
from IPython.display import YouTubeVideo
video = YouTubeVideo(id="oO0bbInoO_0", width=854, height=480, fs=1)
print("Video available at https://youtube.com/watch?v=" + video.id)
video
```

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

---
# Setup
Run these cells to get the tutorial started.

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

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

```python
# @title Helper Functions

def plot_variance_explained(variance_explained):
  """
  Plots eigenvalues.

  Args:
    variance_explained (numpy array of floats) : Vector of variance explained
                                                 for each PC

  Returns:
    Nothing.

  """

  plt.figure()
  plt.plot(np.arange(1, len(variance_explained) + 1), variance_explained,
           '--k')
  plt.xlabel('Number of components')
  plt.ylabel('Variance explained')
  plt.show()

def plot_MNIST_reconstruction(X, X_reconstructed):
  """
  Plots 9 images in the MNIST dataset side-by-side with the reconstructed
  images.

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

  Returns:
    Nothing.
  """

  plt.figure()
  ax = plt.subplot(121)
  k = 0
  for k1 in range(3):
    for k2 in range(3):
      k = k + 1
      plt.imshow(np.reshape(X[k, :], (28, 28)),
                 extent=[(k1 + 1) * 28, k1 * 28, (k2 + 1) * 28, k2 * 28],
                 vmin=0, vmax=255)
  plt.xlim((3 * 28, 0))
  plt.ylim((3 * 28, 0))
  plt.tick_params(axis='both', which='both', bottom=False, top=False,
                  labelbottom=False)
  ax.set_xticks([])
  ax.set_yticks([])
  plt.title('Data')
  plt.clim([0, 250])
  ax = plt.subplot(122)
  k = 0
  for k1 in range(3):
    for k2 in range(3):
      k = k + 1
      plt.imshow(np.reshape(np.real(X_reconstructed[k, :]), (28, 28)),
                 extent=[(k1 + 1) * 28, k1 * 28, (k2 + 1) * 28, k2 * 28],
                 vmin=0, vmax=255)
  plt.xlim((3 * 28, 0))
  plt.ylim((3 * 28, 0))
  plt.tick_params(axis='both', which='both', bottom=False, top=False,
                  labelbottom=False)
  ax.set_xticks([])
  ax.set_yticks([])
  plt.clim([0, 250])
  plt.title('Reconstructed')
  plt.tight_layout()

def plot_MNIST_sample(X):
  """
  Plots 9 images in the MNIST dataset.

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

  Returns:
    Nothing.

  """

  fig, ax = plt.subplots()
  k = 0
  for k1 in range(3):
    for k2 in range(3):
      k = k + 1
      plt.imshow(np.reshape(X[k, :], (28, 28)),
                 extent=[(k1 + 1) * 28, k1 * 28, (k2+1) * 28, k2 * 28],
                 vmin=0, vmax=255)
  plt.xlim((3 * 28, 0))
  plt.ylim((3 * 28, 0))
  plt.tick_params(axis='both', which='both', bottom=False, top=False,
                  labelbottom=False)
  plt.clim([0, 250])
  ax.set_xticks([])
  ax.set_yticks([])
  plt.show()

def plot_MNIST_weights(weights):
  """
  Visualize PCA basis vector weights for MNIST. Red = positive weights,
  blue = negative weights, white = zero weight.

  Args:
     weights (numpy array of floats) : PCA basis vector

  Returns:
     Nothing.
  """

  fig, ax = plt.subplots()
  cmap = plt.cm.get_cmap('seismic')
  plt.imshow(np.real(np.reshape(weights, (28, 28))), cmap=cmap)
  plt.tick_params(axis='both', which='both', bottom=False, top=False,
                  labelbottom=False)
  plt.clim(-.15, .15)
  plt.colorbar(ticks=[-.15, -.1, -.05, 0, .05, .1, .15])
  ax.set_xticks([])
  ax.set_yticks([])
  plt.show()

def add_noise(X, frac_noisy_pixels):
  """
  Randomly corrupts a fraction of the pixels by setting them to random values.

  Args:
     X (numpy array of floats)  : Data matrix
     frac_noisy_pixels (scalar) : Fraction of noisy pixels

  Returns:
     (numpy array of floats)    : Data matrix + noise

  """

  X_noisy = np.reshape(X, (X.shape[0] * X.shape[1]))
  N_noise_ixs = int(X_noisy.shape[0] * frac_noisy_pixels)
  noise_ixs = np.random.choice(X_noisy.shape[0], size=N_noise_ixs,
                               replace=False)
  X_noisy[noise_ixs] = np.random.uniform(0, 255, noise_ixs.shape)
  X_noisy = np.reshape(X_noisy, (X.shape[0], X.shape[1]))

  return X_noisy

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

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

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

  Y = np.matmul(X, W)

  return Y

def get_sample_cov_matrix(X):
  """
  Returns the sample covariance matrix of data X.

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

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

  X = X - np.mean(X, 0)
  cov_matrix = 1 / X.shape[0] * np.matmul(X.T, X)
  return cov_matrix

def sort_evals_descending(evals, evectors):
  """
  Sorts eigenvalues and eigenvectors in decreasing order. Also aligns first two
  eigenvectors to be in first two quadrants (if 2D).

  Args:
    evals (numpy array of floats)    :   Vector of eigenvalues
    evectors (numpy array of floats) :   Corresponding matrix of eigenvectors
                                         each column corresponds to a different
                                         eigenvalue

  Returns:
    (numpy array of floats)          : Vector of eigenvalues after sorting
    (numpy array of floats)          : Matrix of eigenvectors after sorting
  """

  index = np.flip(np.argsort(evals))
  evals = evals[index]
  evectors = evectors[:, index]
  if evals.shape[0] == 2:
    if np.arccos(np.matmul(evectors[:, 0],
                           1 / np.sqrt(2) * np.array([1, 1]))) > np.pi / 2:
      evectors[:, 0] = -evectors[:, 0]
    if np.arccos(np.matmul(evectors[:, 1],
                           1 / np.sqrt(2)*np.array([-1, 1]))) > np.pi / 2:
      evectors[:, 1] = -evectors[:, 1]

  return evals, evectors

def pca(X):
  """
  Performs PCA on multivariate data. Eigenvalues are sorted in decreasing order

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

  Returns:
    (numpy array of floats)    : Data projected onto the new basis
    (numpy array of floats)    : Vector of eigenvalues
    (numpy array of floats)    : Corresponding matrix of eigenvectors

  """

  X = X - np.mean(X, 0)
  cov_matrix = get_sample_cov_matrix(X)
  evals, evectors = np.linalg.eigh(cov_matrix)
  evals, evectors = sort_evals_descending(evals, evectors)
  score = change_of_basis(X, evectors)

  return score, evectors, evals

def plot_eigenvalues(evals, limit=True):
  """
  Plots eigenvalues.

  Args:
     (numpy array of floats) : Vector of eigenvalues

  Returns:
    Nothing.

  """

  plt.figure()
  plt.plot(np.arange(1, len(evals) + 1), evals, 'o-k')
  plt.xlabel('Component')
  plt.ylabel('Eigenvalue')
  plt.title('Scree plot')
  if limit:
    plt.show()
```

---
# Section 1: Perform PCA on MNIST

The MNIST dataset consists of a 70,000 images of individual handwritten digits. Each image is a 28x28 pixel grayscale image. For convenience, each 28x28 pixel image is often unravelled into a single 784 (=28*28) element vector, so that the whole dataset is represented as a 70,000 x 784 matrix. Each row represents a different image, and each column represents a different pixel.

Enter the following cell to load the MNIST dataset and plot the first nine images.

```python
from sklearn.datasets import fetch_openml
mnist = fetch_openml(name='mnist_784')
X = mnist.data
plot_MNIST_sample(X)
```

The MNIST dataset has an extrinsic dimensionality of 784, much higher than the 2-dimensional examples used in the previous tutorials! To make sense of this data, we'll use dimensionality reduction. But first, we need to determine the intrinsic dimensionality $K$ of the data. One way to do this is to look for an "elbow" in the scree plot, to determine which eigenvalues are signficant.

## Exercise 1: Scree plot of MNIST

In this exercise you will examine the scree plot in the MNIST dataset.

**Steps:**
- Perform PCA on the dataset and examine the scree plot.
- When do the eigenvalues appear (by eye) to reach zero? (**Hint:** use `plt.xlim` to zoom into a section of the plot).

```python
help(pca)
help(plot_eigenvalues)
```

    Help on function pca in module __main__:

    pca(X)
        Performs PCA on multivariate data. Eigenvalues are sorted in decreasing order

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

        Returns:
          (numpy array of floats)    : Data projected onto the new basis
          (numpy array of floats)    : Vector of eigenvalues
          (numpy array of floats)    : Corresponding matrix of eigenvectors

    Help on function plot_eigenvalues in module __main__:

    plot_eigenvalues(evals, limit=True)
        Plots eigenvalues.

        Args:
           (numpy array of floats) : Vector of eigenvalues

        Returns:
          Nothing.

```python
#################################################
## TO DO for students: perform PCA and plot the eigenvalues
#################################################

# perform PCA
score, evectors, evals = pca(X)
# plot the eigenvalues
plot_eigenvalues(evals, limit=False)
plt.xlim([0, 50])  # limit x-axis up to 100 for zooming
```

---
# Section 2: Calculate the variance explained

The scree plot suggests that most of the eigenvalues are near zero, with fewer than 100 having large values. Another common way to determine the intrinsic dimensionality is by considering the variance explained. This can be examined with a cumulative plot of the fraction of the total variance explained by the top $K$ components, i.e.,

\begin{equation}
\text{var explained} = \frac{\sum_{i=1}^K \lambda_i}{\sum_{i=1}^N \lambda_i}
\end{equation}

The intrinsic dimensionality is often quantified by the $K$ necessary to explain a large proportion of the total variance of the data (often a defined threshold, e.g., 90%).

## Exercise 2: Plot the explained variance

In this exercise you will plot the explained variance.

**Steps:**
- Fill in the function below to calculate the fraction variance explained as a function of the number of principal componenets. **Hint:** use `np.cumsum`.
- Plot the variance explained using `plot_variance_explained`.

**Questions:**
- How many principal components are required to explain 90% of the variance?
- How does the intrinsic dimensionality of this dataset compare to its extrinsic dimensionality?

```python
help(plot_variance_explained)
```

    Help on function plot_variance_explained in module __main__:

    plot_variance_explained(variance_explained)
        Plots eigenvalues.

        Args:
          variance_explained (numpy array of floats) : Vector of variance explained
                                                       for each PC

        Returns:
          Nothing.

```python
def get_variance_explained(evals):
  """
  Calculates variance explained from the eigenvalues.

  Args:
    evals (numpy array of floats) : Vector of eigenvalues

  Returns:
    (numpy array of floats)       : Vector of variance explained

  """

  # cumulatively sum the eigenvalues
  csum = np.cumsum(evals)
  # normalize by the sum of eigenvalues
  variance_explained = csum / np.sum(evals)

  return variance_explained

#################################################
## TO DO for students: call the function and plot the variance explained
#################################################

# calculate the variance explained
variance_explained = get_variance_explained(evals)

# Uncomment to plot the variance explained
plot_variance_explained(variance_explained)
```

---
# Section 3: Reconstruct data with different numbers of PCs

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

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

Now we have seen that the top 100 or so principal components of the data can explain most of the variance. We can use this fact to perform *dimensionality reduction*, i.e., by storing the data using only 100 components rather than the samples of all 784 pixels. Remarkably, we will be able to reconstruct much of the structure of the data using only the top 100 components. To see this, recall that to perform PCA we projected the data $\bf X$ onto the eigenvectors of the covariance matrix:
\begin{equation}
\bf S = X W
\end{equation}
Since $\bf W$ is an orthogonal matrix, ${\bf W}^{-1} = {\bf W}^T$. So by multiplying by ${\bf W}^T$ on each side we can rewrite this equation as
\begin{equation}
{\bf X = S W}^T.
\end{equation}
This now gives us a way to reconstruct the data matrix from the scores and loadings. To reconstruct the data from a low-dimensional approximation, we just have to truncate these matrices.  Let's call ${\bf S}_{1:K}$ and ${\bf W}_{1:K}$ as keeping only the first $K$ columns of this matrix. Then our reconstruction is:
\begin{equation}
{\bf \hat X = S}_{1:K} ({\bf W}_{1:K})^T.
\end{equation}

## Exercise 3: Data reconstruction

Fill in the function below to reconstruct the data using different numbers of principal components.

**Steps:**

* Fill in the following function to reconstruct the data based on the weights and scores. Don't forget to add the mean!
* Make sure your function works by reconstructing the data with all $K=784$ components. The two images should look identical.

```python
help(plot_MNIST_reconstruction)
```

    Help on function plot_MNIST_reconstruction in module __main__:

    plot_MNIST_reconstruction(X, X_reconstructed)
        Plots 9 images in the MNIST dataset side-by-side with the reconstructed
        images.

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

        Returns:
          Nothing.

```python
def reconstruct_data(score, evectors, X_mean, K):
  """
  Reconstruct the data based on the top K components.

  Args:
    score (numpy array of floats)    : Score matrix
    evectors (numpy array of floats) : Matrix of eigenvectors
    X_mean (numpy array of floats)   : Vector corresponding to data mean
    K (scalar)                       : Number of components to include

  Returns:
    (numpy array of floats)          : Matrix of reconstructed data

  """

  # Reconstruct the data from the score and eigenvectors
  # Don't forget to add the mean!!
  X_reconstructed = score[:, :K] @ evectors[:, :K].T + X_mean

  return X_reconstructed

K = 784

# Reconstruct the data based on all components
X_mean = np.mean(X, 0)
X_reconstructed = reconstruct_data(score, evectors, X_mean, K)

# Plot the data and reconstruction
plot_MNIST_reconstruction(X, X_reconstructed)
```

## Interactive Demo: Reconstruct the data matrix using different numbers of PCs

Now run the code below and experiment with the slider to reconstruct the data matrix using different numbers of principal components.

**Steps**
* How many principal components are necessary to reconstruct the numbers (by eye)? How does this relate to the intrinsic dimensionality of the data?
* Do you see any information in the data with only a single principal component?

```python
# @title

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

def refresh(K=100):
  X_reconstructed = reconstruct_data(score, evectors, X_mean, K)
  plot_MNIST_reconstruction(X, X_reconstructed)
  plt.title('Reconstructed, K={}'.format(K))

_ = widgets.interact(refresh, K=(1, 784, 10))
```

    interactive(children=(IntSlider(value=100, description='K', max=784, min=1, step=10), Output()), _dom_classes=…

## Exercise 4: Visualization of the weights

Next, let's take a closer look at the first principal component by visualizing its corresponding weights.

**Steps:**

* Enter `plot_MNIST_weights` to visualize the weights of the first basis vector.
* What structure do you see? Which pixels have a strong positive weighting? Which have a strong negative weighting? What kinds of images would this basis vector differentiate?
* Try visualizing the second and third basis vectors. Do you see any structure? What about the 100th basis vector? 500th? 700th?

```python
help(plot_MNIST_weights)
```

    Help on function plot_MNIST_weights in module __main__:

    plot_MNIST_weights(weights)
        Visualize PCA basis vector weights for MNIST. Red = positive weights,
        blue = negative weights, white = zero weight.

        Args:
           weights (numpy array of floats) : PCA basis vector

        Returns:
           Nothing.

```python
# Plot the weights of the first principal component
plot_MNIST_weights(evectors[:, 0])
```

---
# Summary
* In this tutorial, we learned how to use PCA for dimensionality reduction by selecting the top principal components. This can be useful as the intrinsic dimensionality ($K$) is often less than the extrinsic dimensionality ($N$) in neural data. $K$ can be inferred by choosing the number of eigenvalues necessary to capture some fraction of the variance.
* We also learned how to reconstruct an approximation of the original data using the top $K$ principal components. In fact, an alternate formulation of PCA is to find the $K$ dimensional space that minimizes the reconstruction error.
* Noise tends to inflate the apparent intrinsic dimensionality, however the higher components reflect noise rather than new structure in the data. PCA can be used for denoising data by removing noisy higher components.
* In MNIST, the weights corresponding to the first principal component appear to discriminate between a 0 and 1. We will discuss the implications of this for data visualization in the following tutorial.

---
# Bonus: Examine denoising using PCA

In this lecture, we saw that PCA finds an optimal low-dimensional basis to minimize the reconstruction error. Because of this property, PCA can be useful for denoising corrupted samples of the data.

## Exercise 5: Add noise to the data
In this exercise you will add salt-and-pepper noise to the original data and see how that affects the eigenvalues.

**Steps:**
- Use the function `add_noise` to add noise to 20% of the pixels.
- Then, perform PCA and plot the variance explained. How many principal components are required to explain 90% of the variance? How does this compare to the original data?

```python
help(add_noise)
```

    Help on function add_noise in module __main__:

    add_noise(X, frac_noisy_pixels)
        Randomly corrupts a fraction of the pixels by setting them to random values.

        Args:
           X (numpy array of floats)  : Data matrix
           frac_noisy_pixels (scalar) : Fraction of noisy pixels

        Returns:
           (numpy array of floats)    : Data matrix + noise

```python
np.random.seed(2020)  # set random seed
X_noisy = add_noise(X, .2)
score_noisy, evectors_noisy, evals_noisy = pca(X_noisy)
variance_explained_noisy = get_variance_explained(evals_noisy)

with plt.xkcd():
  plot_MNIST_sample(X_noisy)
  plot_variance_explained(variance_explained_noisy)
```

## Exercise 6: Denoising

Next, use PCA to perform denoising by projecting the noise-corrupted data onto the basis vectors found from the original dataset. By taking the top K components of this projection, we can reduce noise in dimensions orthogonal to the K-dimensional latent space.

**Steps:**
- Subtract the mean of the noise-corrupted data.
- Project the data onto the basis found with the original dataset (`evectors`, not `evectors_noisy`) and take the top $K$ components.
- Reconstruct the data as normal, using the top 50 components.
- Play around with the amount of noise and K to build intuition.

```python
X_noisy_mean = np.mean(X_noisy, 0)
projX_noisy = np.matmul(X_noisy - X_noisy_mean, evectors)
X_reconstructed = reconstruct_data(projX_noisy, evectors, X_noisy_mean, 50)

with plt.xkcd():
  plot_MNIST_reconstruction(X_noisy, X_reconstructed)
```