# Inferring cosmology power spectrum parameters from 2D fields
(by Lucas Makinen and Tom Charnock)

A useful example to consider is the inference of the parameter values controlling the shape of the power spectrum used to generate 2D Gaussian random fields. This is interesting because we can analytically calculate the Fisher information to check our convergence and it has aspects which are very similar to cosmological field analysis. This example is split into two parts, first the analytic comparison when the power spectrum is of the form
$$P(k) = Ak^{-B}$$
and a second part where we use a jax implementation of a cosmology library (`pip install jax-cosmo`) to infer cosmological parameters from 2D Gaussian random fields.

In this example we will also show how inception blocks can be implemented in a stax-like neural model.

## Power law example - $P(k) = Ak^{-B}$

The likelihood for an $N_{\rm pix}\times N_{\rm pix}$ Gaussian field, $\boldsymbol{\delta}$, can be explicitly written down for the Fourier transformed data, $\boldsymbol{\Delta}$ as
$$\mathcal{L}(\boldsymbol{\Delta}|\boldsymbol{\theta}) = \frac{1}{(2\pi)^{N_{\rm pix}^2 / 2} |P({\bf k}, \boldsymbol{\theta})|^{1/2}}\exp{\left(-\frac{1}{2}\boldsymbol{\Delta}\left(P({\bf k}, \boldsymbol{\theta})\right)^{-1}\boldsymbol{\Delta}\right)}$$
Since the Fisher information can be calculated from the expectation value of the second derivative of the score, i.e. the log likelihood
$${\bf F}_{\alpha\beta} = - \left.\left\langle\frac{\partial^2\ln\mathcal{L}(\Delta|\boldsymbol{\theta})}{\partial\theta_\alpha\partial\theta_\beta}\right\rangle\right|_{\boldsymbol{\theta}=\boldsymbol{\theta}^\textrm{fid}}$$
Then we know that analytically the Fisher information must be
$${\bf F}_{\alpha\beta} = \frac{1}{2} {\rm Tr} \left(\frac{\partial P({\bf k}, \boldsymbol{\theta})}{\partial\theta_\alpha}\left(P({\bf k}, \boldsymbol{\theta})\right)^{-1}\frac{\partial P({\bf k}, \boldsymbol{\theta})}{\partial\theta_\beta}\left(P({\bf k}, \boldsymbol{\theta})\right)^{-1}\right)$$
where $\alpha$ and $\beta$ label the parameters $A$ and $B$ in the power spectrum. As each $k$-mode is uncoupled for this power law form then the derivatives  are
$$\begin{align}
\frac{\partial P({\bf k}, \boldsymbol{\theta})}{\partial A} = &~k^{-B}\\
\frac{\partial P({\bf k}, \boldsymbol{\theta})}{\partial B} = & -Ak^{-B}\ln k.\\
\end{align}$$
We can set up these functions. Note that for large fields there can be an issue with numerical overflow and as such we will use jax in 64 bit mode.

```python
cd ~/repositories/scratch/gitinstall/imnn
```

    /mnt/home/tmakinen/repositories/scratch/gitinstall/imnn

```python
!git reset --hard origin/dev
```

    HEAD is now at 5fbc399 2d field inference example added

```python
!git pull
```

    Already up-to-date.

```python
!python -m pip install .
```

    Processing /mnt/home/tmakinen/repositories/scratch/gitinstall/imnn
    Requirement already satisfied: jax>=0.2.10 in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from imnn==0.3.dev0) (0.2.10)
    Requirement already satisfied: jaxlib>=0.1.62 in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from imnn==0.3.dev0) (0.1.62+cuda101)
    Requirement already satisfied: tensorflow in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from imnn==0.3.dev0) (2.4.1)
    Requirement already satisfied: tqdm>=4.48.2 in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages/tqdm-4.59.0-py3.8.egg (from imnn==0.3.dev0) (4.59.0)
    Requirement already satisfied: tensorflow-probability[jax] in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages/tensorflow_probability-0.12.1-py3.8.egg (from imnn==0.3.dev0) (0.12.1)
    Requirement already satisfied: numpy in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from imnn==0.3.dev0) (1.19.5)
    Requirement already satisfied: scipy in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from imnn==0.3.dev0) (1.6.2)
    Requirement already satisfied: matplotlib in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages/matplotlib-3.4.1-py3.8-linux-x86_64.egg (from imnn==0.3.dev0) (3.4.1)
    Requirement already satisfied: opt-einsum in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from jax>=0.2.10->imnn==0.3.dev0) (3.3.0)
    Requirement already satisfied: absl-py in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from jax>=0.2.10->imnn==0.3.dev0) (0.12.0)
    Requirement already satisfied: flatbuffers in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from jaxlib>=0.1.62->imnn==0.3.dev0) (1.12)
    Requirement already satisfied: six in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from absl-py->jax>=0.2.10->imnn==0.3.dev0) (1.15.0)
    Requirement already satisfied: cycler>=0.10 in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages/cycler-0.10.0-py3.8.egg (from matplotlib->imnn==0.3.dev0) (0.10.0)
    Requirement already satisfied: kiwisolver>=1.0.1 in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages/kiwisolver-1.3.1-py3.8-linux-x86_64.egg (from matplotlib->imnn==0.3.dev0) (1.3.1)
    Requirement already satisfied: pillow>=6.2.0 in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages/Pillow-8.1.2-py3.8-linux-x86_64.egg (from matplotlib->imnn==0.3.dev0) (8.1.2)
    Requirement already satisfied: pyparsing>=2.2.1 in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from matplotlib->imnn==0.3.dev0) (2.4.7)
    Requirement already satisfied: python-dateutil>=2.7 in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from matplotlib->imnn==0.3.dev0) (2.8.1)
    Requirement already satisfied: wheel~=0.35 in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from tensorflow->imnn==0.3.dev0) (0.36.2)
    Requirement already satisfied: protobuf>=3.9.2 in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from tensorflow->imnn==0.3.dev0) (3.15.6)
    Requirement already satisfied: termcolor~=1.1.0 in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from tensorflow->imnn==0.3.dev0) (1.1.0)
    Requirement already satisfied: google-pasta~=0.2 in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from tensorflow->imnn==0.3.dev0) (0.2.0)
    Requirement already satisfied: gast==0.3.3 in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from tensorflow->imnn==0.3.dev0) (0.3.3)
    Requirement already satisfied: tensorboard~=2.4 in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from tensorflow->imnn==0.3.dev0) (2.4.1)
    Requirement already satisfied: typing-extensions~=3.7.4 in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from tensorflow->imnn==0.3.dev0) (3.7.4.3)
    Requirement already satisfied: astunparse~=1.6.3 in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from tensorflow->imnn==0.3.dev0) (1.6.3)
    Requirement already satisfied: keras-preprocessing~=1.1.2 in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from tensorflow->imnn==0.3.dev0) (1.1.2)
    Requirement already satisfied: wrapt~=1.12.1 in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from tensorflow->imnn==0.3.dev0) (1.12.1)
    Requirement already satisfied: tensorflow-estimator<2.5.0,>=2.4.0 in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from tensorflow->imnn==0.3.dev0) (2.4.0)
    Requirement already satisfied: grpcio~=1.32.0 in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from tensorflow->imnn==0.3.dev0) (1.32.0)
    Requirement already satisfied: h5py~=2.10.0 in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from tensorflow->imnn==0.3.dev0) (2.10.0)
    Requirement already satisfied: markdown>=2.6.8 in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from tensorboard~=2.4->tensorflow->imnn==0.3.dev0) (3.3.4)
    Requirement already satisfied: tensorboard-plugin-wit>=1.6.0 in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from tensorboard~=2.4->tensorflow->imnn==0.3.dev0) (1.8.0)
    Requirement already satisfied: werkzeug>=0.11.15 in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from tensorboard~=2.4->tensorflow->imnn==0.3.dev0) (1.0.1)
    Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from tensorboard~=2.4->tensorflow->imnn==0.3.dev0) (0.4.4)
    Requirement already satisfied: requests<3,>=2.21.0 in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from tensorboard~=2.4->tensorflow->imnn==0.3.dev0) (2.25.1)
    Requirement already satisfied: setuptools>=41.0.0 in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from tensorboard~=2.4->tensorflow->imnn==0.3.dev0) (54.2.0)
    Requirement already satisfied: google-auth<2,>=1.6.3 in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from tensorboard~=2.4->tensorflow->imnn==0.3.dev0) (1.28.0)
    Requirement already satisfied: rsa<5,>=3.1.4 in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from google-auth<2,>=1.6.3->tensorboard~=2.4->tensorflow->imnn==0.3.dev0) (4.7.2)
    Requirement already satisfied: cachetools<5.0,>=2.0.0 in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from google-auth<2,>=1.6.3->tensorboard~=2.4->tensorflow->imnn==0.3.dev0) (4.2.1)
    Requirement already satisfied: pyasn1-modules>=0.2.1 in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from google-auth<2,>=1.6.3->tensorboard~=2.4->tensorflow->imnn==0.3.dev0) (0.2.8)
    Requirement already satisfied: requests-oauthlib>=0.7.0 in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from google-auth-oauthlib<0.5,>=0.4.1->tensorboard~=2.4->tensorflow->imnn==0.3.dev0) (1.3.0)
    Requirement already satisfied: pyasn1<0.5.0,>=0.4.6 in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from pyasn1-modules>=0.2.1->google-auth<2,>=1.6.3->tensorboard~=2.4->tensorflow->imnn==0.3.dev0) (0.4.8)
    Requirement already satisfied: chardet<5,>=3.0.2 in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from requests<3,>=2.21.0->tensorboard~=2.4->tensorflow->imnn==0.3.dev0) (4.0.0)
    Requirement already satisfied: urllib3<1.27,>=1.21.1 in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from requests<3,>=2.21.0->tensorboard~=2.4->tensorflow->imnn==0.3.dev0) (1.26.4)
    Requirement already satisfied: certifi>=2017.4.17 in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from requests<3,>=2.21.0->tensorboard~=2.4->tensorflow->imnn==0.3.dev0) (2020.12.5)
    Requirement already satisfied: idna<3,>=2.5 in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from requests<3,>=2.21.0->tensorboard~=2.4->tensorflow->imnn==0.3.dev0) (2.10)
    Requirement already satisfied: oauthlib>=3.0.0 in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<0.5,>=0.4.1->tensorboard~=2.4->tensorflow->imnn==0.3.dev0) (3.1.0)
    Requirement already satisfied: cloudpickle>=1.3 in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages/cloudpickle-1.6.0-py3.8.egg (from tensorflow-probability[jax]->imnn==0.3.dev0) (1.6.0)
    Requirement already satisfied: decorator in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages (from tensorflow-probability[jax]->imnn==0.3.dev0) (4.4.2)
    Requirement already satisfied: dm-tree in /mnt/home/tmakinen/anaconda3/envs/pyjax9/lib/python3.8/site-packages/dm_tree-0.1.5-py3.8-linux-x86_64.egg (from tensorflow-probability[jax]->imnn==0.3.dev0) (0.1.5)
    Building wheels for collected packages: imnn
      Building wheel for imnn (setup.py) ... done
      Created wheel for imnn: filename=imnn-0.3.dev0-py3-none-any.whl size=105309 sha256=53025ec2938c87cfc09d3d552e5c1365c480eaaeea9e5e12d996605522a016d3
      Stored in directory: /tmp/pip-ephem-wheel-cache-qcew2oh_/wheels/df/88/04/7b2df4336511103974c1022b25ee71b9c7219512ec155c2ca8
    Successfully built imnn
    Installing collected packages: imnn
      Attempting uninstall: imnn
        Found existing installation: imnn 0.3.dev0
        Uninstalling imnn-0.3.dev0:
          Successfully uninstalled imnn-0.3.dev0
    Successfully installed imnn-0.3.dev0

```python
from jax.config import config
config.update('jax_enable_x64', True)
```

```python
import imnn
import imnn.lfi
import jax
import jax.numpy as np
import jax_cosmo as jc
import matplotlib.pyplot as plt
import tensorflow_probability
from jax.experimental import optimizers
from jax.experimental import stax
tfp = tensorflow_probability.substrates.jax
rng = jax.random.PRNGKey(0)
```

```python
from jax.lib import xla_bridge
print(xla_bridge.get_backend().platform)
```

    gpu

First we'll set up our $P(k)$

```python
def P(k, A=1, B=0):
    return A * k ** (-B)
```

And the log likelihood is then

```python
def log_likelihood(k, A, B, Δ, N):
    Δ = Δ[1:N // 2, 1:N // 2].flatten()
    k = k[1:N // 2, 1:N // 2].flatten()
    dlength = len(k)
    def fn(_A, _B):
        nrm = np.pad(np.ones(dlength - 2) * 2, (1, 1), constant_values=1.)
        nrm = jax.ops.index_update(
          nrm, np.array([[0],[(dlength - 2)]]), np.array([[1], [1]]))
        nrm = 1
        powers = P(k, A=_A, B=_B)
        C = powers * nrm
        invC = 1. / powers
        exponent = - 0.5 * np.sum(np.conj(Δ) * invC * Δ)
        norm = -0.5 * np.sum(np.log(C)) -0.5 * len(Δ) * np.log(np.pi * 2.)
        return (exponent + norm)
    return jax.vmap(fn)(A, B)
```

and the Fisher information matrix is given by

```python
def fisher(θ, k, N):
    A, B = θ
    k = k[1:N // 2, 1:N // 2].flatten()
    Pk = P(k, A, B)

    Cinv = np.diag(1. / Pk)
    C_A =  np.diag(k ** -B)
    C_B =  np.diag(- Pk * np.log(k))

    F_AA = 0.5 * np.trace((C_A @ Cinv @ C_A @ Cinv))
    F_AB = 0.5 * np.trace((C_A @ Cinv @ C_B @ Cinv))
    F_BA = 0.5 * np.trace((C_B @ Cinv @ C_A @ Cinv))
    F_BB = 0.5 * np.trace((C_B @ Cinv @ C_B @ Cinv))

    return np.array([[F_AA, F_AB], [F_BA, F_BB]])
```

Notice that we only take the unique modes to calculate the Fisher information. To actually set up our 2D $k$-modes we stack values from $0\to N_i/2 + 1$ to values from $-N_i/2+1\to0$ and normalising them by $2\pi/N_i$ where $N_i$ is the number of $k$-modes in each axis. If we choose that our fields have $128\times128$ pixels then we can calculate $k$ as

```python
N = 128
shape = (N, N)

k = np.sqrt(
    np.sum(
        np.array(
            np.meshgrid(
                *((np.hstack(
                    (np.arange(0, _shape // 2 + 1),
                     np.arange(-_shape // 2 + 1, 0)))
                   * 2 * np.pi / _shape)**2.
                  for _shape in shape))),
        axis=0))
```

For a fiducial $A^\textrm{fid}=1$ and $B^\textrm{fid}=1/2$ we get a Fisher information matrix equal to

```python
θ_fid = np.array([1., 0.5], dtype=np.float32)
n_params = θ_fid.shape[0]

F = fisher(θ_fid, k, N)
print(f"F = {F}")
print(f"det(F) = {np.linalg.det(F)}")
```

    F = [[ 1984.5        -1545.06379524]
     [-1545.06379524  1690.49264978]]
    det(F) = 967560.5321330648

### 2D Gaussian random field simulator in jax

To create a 2D Gaussian random field we can follow these steps:

1. Generate a $(N_\textrm{pix}\times N_\textrm{pix})$ white noise field $\varphi$ such that $\langle \varphi_k \varphi_{-k} \rangle' = 1$

2. Fourier Transform $\varphi$ to real space: $R_{\rm white}({\bf x}) \rightarrow R_{\rm white}({\bf k})$
   Note that NumPy's DFT Fourier convention is:
   $$\phi_{ab}^{\bf k} = \sum_{c,d = 0}^{N-1} \exp{(-i x_c k_a - i x_d k_b) \phi^{\bf x}_{cd}}$$
   $$\phi_{ab}^{\bf x} = \frac{1}{N^2}\sum_{c,d = 0}^{N-1} \exp{(-i x_c k_a - i x_d k_b) \phi^{\bf k}_{cd}}$$

3. Scale white noise $R_{\rm white}({\bf k})$ by the chosen power spectrum evaluated over a field of $k$ values:
   $$R_P({\bf k}) = P^{1/2}(k) R_{\rm white}({\bf k}) $$
   Here we need to ensure that this array of amplitudes are Hermitian, e.g. $\phi^{* {\bf k}}_{a(N/2 + b)} = \phi^{{\bf k}}_{a(N/2 - b)}$. This is accomplished by choosing indices $k_a = k_b = \frac{2\pi}{N} (0, \dots, N/2, -N/2+1, \dots, -1)$ (as above) and then evaluating the square root of the outer product of the meshgrid between the two: $k = \sqrt{k^2_a + k^2_b}$. We can then evaluate $P^{1/2}(k)$.

4. Fourier Transform $R_{P}({\bf k})$ to real space: $R_P({\bf x}) = \int d^d \tilde{k} e^{i{\bf k} \cdot {\bf x}} R_p({\bf k})$
   $$R_{ab}^{\bf x} = \frac{1}{N^2}\sum_{c,d = 0}^{N-1} \exp{(-i x_c k_a - i x_d k_b) R^{\bf k}_{cd}}$$

We are going to build a broadcastable jax simultor which takes in a variety of different shaped parameter arrays and vmaps them until a single parameter pair are passed. This is very efficient for performing the ABC for example.

```python
def simulator(rng, θ, simulator_args, foregrounds=None):
    def fn(rng, A, B):
        dim = len(simulator_args["shape"])
        L = simulator_args["L"]
        if np.isscalar(L):
            L = [L] * int(dim)
        Lk = ()
        shape = ()
        for i, _shape in enumerate(simulator_args["shape"]):
            Lk += (_shape / L[i],)
            if _shape % 2 == 0:
                shape += (_shape + 1,)
            else:
                shape += (_shape,)

        k = simulator_args["k"]
        k_shape = k.shape
        k = k.flatten()[1:]
        tpl = ()
        for _d in range(dim):
            tpl += (_d,)

        V = np.prod(np.array(L))
        scale = V**(1. / dim)
        fft_norm = np.prod(np.array(Lk))

        rng, key = jax.random.split(rng)

        mag = jax.random.normal(
            key, shape=shape)
        pha = 2. * np.pi * jax.random.uniform(
            key, shape=shape)

        # now make hermitian field (reality condition)
        revidx = (slice(None, None, -1),) * dim
        mag = (mag + mag[revidx]) / np.sqrt(2)
        pha = (pha - pha[revidx]) / 2 + np.pi
        dk = mag * (np.cos(pha) + 1j * np.sin(pha))
        cutidx = (slice(None, -1),) * dim
        dk = dk[cutidx]

        powers = np.concatenate(
            (np.zeros(1),
             np.sqrt(P(k, A=A, B=B)))).reshape(k_shape)

        if simulator_args['vol_norm']:
            powers /= V

        if simulator_args["log_normal"]:
            powers = np.real(
                np.fft.ifftshift(
                    np.fft.ifftn(
                        powers)
                    * fft_norm) * V)

            powers = np.log(1. + powers)
            powers = np.abs(np.fft.fftn(powers))

        fourier_field = powers * dk
        fourier_field = jax.ops.index_update(
            fourier_field,
            np.zeros(dim, dtype=int),
            np.zeros((1,)))

        if simulator_args["log_normal"]:
            field = np.real(np.fft.ifftn(fourier_field)) * fft_norm * np.sqrt(V)
            sg = np.var(field)
            field = np.exp(field - sg / 2.) - 1.

        else:
            field = np.real(np.fft.ifftn(fourier_field) * fft_norm * np.sqrt(V)**2)

        if simulator_args["N_scale"]:
            field *= scale

        if foregrounds is not None:
            rng, key = jax.random.split(key)
            foreground = foregrounds[
                jax.random.randint(
                    key,
                    minval=0,
                    maxval=foregrounds.shape[0],
                    shape=())]
            field = np.expand_dims(field + foreground, (0,))

        if not simulator_args["squeeze"]:
            field = np.expand_dims(field, (0, 1))

        return np.array(field, dtype='float32')

    if isinstance(θ, tuple):
        A, B = θ
    else:
        A = np.take(θ, 0, axis=-1)
        B = np.take(θ, 1, axis=-1)
    if A.shape == B.shape:
        if len(A.shape) == 0:
            return fn(rng, A, B)
        else:
            keys = jax.random.split(rng, num=A.shape[0] + 1)
            rng = keys[0]
            keys = keys[1:]
            return jax.vmap(
                lambda key, A, B: simulator(
                    key, (A, B), simulator_args=simulator_args))(
                keys, A, B)
    else:
        if len(A.shape) > 0:
            keys = jax.random.split(rng, num=A.shape[0] + 1)
            rng = keys[0]
            keys = keys[1:]
            return jax.vmap(
                lambda key, A: simulator(
                    key, (A, B), simulator_args=simulator_args))(
                keys, A)
        elif len(B.shape) > 0:
            keys = jax.random.split(rng, num=B.shape[0])
            return jax.vmap(
                lambda key, B: simulator(
                    key, (A, B), simulator_args=simulator_args))(
                keys, B)
```

We can now set the simulator arguments, i.e. the $k$-modes to evaluate, the length of the side of a box, the shape of the box and whether to normalise via the volume and squeeze the output dimensions

```python
simulator_args = dict(
    k=k,
    L=N,
    shape=shape,
    vol_norm=True,
    N_scale=True,
    squeeze=True,
    log_normal=False)
```

Now we can simulate some target data at, for example, $A^\textrm{target}=0.7$ and $B^\textrm{target}=0.8$:

```python
θ_target = np.array([0.7, 0.8])

rng, key = jax.random.split(rng)
δ_target = simulator(key, θ_target, simulator_args=simulator_args)
plt.imshow(δ_target)
plt.colorbar();
```

We can now define our prior distribution (in this case a uniform distribution over $A$ and $B$) with values between 0.1 and 1.25 for both parameters

```python
prior = tfp.distributions.Blockwise(
    [tfp.distributions.Uniform(low=low, high=high)
     for low, high in zip([0.1, 0.1], [1.25, 1.25])])
prior.low = np.array([0.1, 0.1])
prior.high = np.array([1.25, 1.25])
```

To evaluate the likelihood of this field we can now use (dividing the target $\delta$ by $N$ to remove added scaling)

```python
LFI = imnn.lfi.LikelihoodFreeInference(
    prior=prior,
    gridsize=100)
A, B = np.meshgrid(*LFI.ranges)
LFI.n_targets=1
LFI.put_marginals(
        jax.nn.softmax(
            np.real(
                log_likelihood(
                    k,
                    A.ravel(),
                    B.ravel(),
                    np.fft.fftn(δ_target / N),
                    N)
                ),axis=0
            ).reshape((100, 100)).T[np.newaxis]);
LFI.marginal_plot(
    known=θ_target,
    label="Analytic likelihood",
    axis_labels=["A", "B"]);
```

```python
# sanity check:
fig,ax = plt.subplots(2,2)

for i in range(20):
    θ_target = np.array([0.7, 0.8])

    rng, key = jax.random.split(rng)
    δ_target = simulator(key, θ_target, simulator_args=simulator_args)

    LFI = imnn.lfi.LikelihoodFreeInference(
        prior=prior,
        gridsize=100)
    A, B = np.meshgrid(*LFI.ranges)
    LFI.n_targets=1
    LFI.put_marginals(
        jax.nn.softmax(
            np.real(
                log_likelihood(
                    k,
                    A.ravel(),
                    B.ravel(),
                    np.fft.fftn(δ_target / N),
                    N)
                ), axis=0).reshape((100, 100)).T[np.newaxis]);
    ax = LFI.marginal_plot(
        ax=ax,
        known=θ_target,
        label="Analytic likelihood",
        axis_labels=["A", "B"]);

plt.show()
```

## Training an IMNN

Now lets train an IMNN to summaries such Gaussian random fields to see how much information we can extract an what sort of constraints we can get. We will use 5000 simulations to estimate the covariance and use all of their derivatives and we'll summarise the whole random Gaussian field by 2 summaries.

```python
n_s = 5000
n_d = n_s

n_summaries = 2
```

We're going to use a fully convolutional inception network built using stax with some custom designed blocks. The inception block itself is implemented as

```python
def InceptBlock(filters, strides, do_5x5=True, do_3x3=True):
    """InceptNet convolutional striding block.
    filters: tuple: (f1,f2,f3)
    filters1: for conv1x1
    filters2: for conv1x1,conv3x3
    filters3L for conv1x1,conv5x5"""

    filters1, filters2, filters3 = filters
    conv1x1 = stax.serial(stax.Conv(filters1, (1, 1), strides, padding="SAME"))

    filters4 = filters2
    conv3x3 = stax.serial(stax.Conv(filters2, (1, 1), strides=None, padding="SAME"),
                        stax.Conv(filters4, (3, 3), strides, padding="SAME"))

    filters5 = filters3
    conv5x5 = stax.serial(stax.Conv(filters3, (1, 1), strides=None, padding="SAME"),
                         stax.Conv(filters5, (5, 5), strides, padding="SAME"))

    maxpool = stax.serial(stax.MaxPool((3, 3), padding="SAME"),
                         stax.Conv(filters4, (1, 1), strides, padding="SAME"))

    if do_3x3:
        if do_5x5:
            return stax.serial(
                  stax.FanOut(4),
                  stax.parallel(conv1x1, conv3x3, conv5x5, maxpool),
                  stax.FanInConcat(),
                  stax.LeakyRelu)
        else:
            return stax.serial(
                  stax.FanOut(3),
                  stax.parallel(conv1x1, conv3x3, maxpool),
                  stax.FanInConcat(),
                  stax.LeakyRelu)
    else:
        return stax.serial(
              stax.FanOut(2),
              stax.parallel(conv1x1, maxpool),
              stax.FanInConcat(),
              stax.LeakyRelu)
```

We'll also want to make sure that the output of the network is the correct shape, for which we'll introduce a Reshaping layer

```python
def Reshape(shape):
    """Layer function for a reshape layer."""
    init_fun = lambda rng, input_shape: (shape,())
    apply_fun = lambda params, inputs, **kwargs: np.reshape(inputs, shape)
    return init_fun, apply_fun
```

Now we can build the network, with kernel sizes of 4 in each direction in each layer

```python
fs = 64

model = stax.serial(
        InceptBlock((fs, fs, fs), strides=(4, 4)),
        InceptBlock((fs, fs, fs), strides=(4, 4)),
        InceptBlock((fs, fs, fs), strides=(4, 4)),
        InceptBlock((fs, fs, fs), strides=(2, 2), do_5x5=False, do_3x3=False),
        stax.Conv(n_summaries, (1, 1), strides=(1, 1), padding="SAME"),
        stax.Flatten,
        Reshape((n_summaries,)))
```

We'll also grab an adam optimiser from jax.experimental.optimizers

```python
optimiser = optimizers.adam(step_size=1e-3)
```

Note that due to the form of the network we'll want to have simulations that have a "channel" dimension, which we can set up by not allowing for squeezing in the simulator.

### Initialise IMNN

Finally we can initialise the IMNN, letting the IMNN module decide what type of IMNN subclass will be used (we'll be using SimulatorIMNN)

```python
rng, key = jax.random.split(rng)
IMNN = imnn.IMNN(
        n_s=n_s,
        n_d=n_d,
        n_params=n_params,
        n_summaries=n_summaries,
        input_shape=(1, 1) + shape,
        θ_fid=θ_fid,
        model=model,
        optimiser=optimiser,
        key_or_state=key,
        simulator=lambda rng, θ: simulator(
            rng, θ, simulator_args={
                **simulator_args,
                **{"squeeze": False}}))
```

    `simulator` provided, using SimulatorIMNN

And finally we can fit the IMNN (we'll use generic regularisation parameters of $\lambda=10$ and $\epsilon=0.1$) and allow early stopping to determine the end of fitting.

```python
%%time
rng, key = jax.random.split(rng)
IMNN.fit(λ=10., ϵ=0.1, rng=key, print_rate=None, min_iterations=500, best=True)
```

    CPU times: user 13min 48s, sys: 6.06 s, total: 13min 54s
    Wall time: 14min 2s

```python
IMNN.plot(expected_detF=np.linalg.det(F));
```

```python
np.linalg.det(IMNN.F) / np.linalg.det(F)
```

    DeviceArray(0.97123594, dtype=float64)

```python
np.linalg.det(F)
```

    DeviceArray(967560.53213307, dtype=float64)

## Inference

We can now attempt to do inference of some target data using the IMNN. The first thing we should do is make a Gaussian approximation using a parameter estimate from the IMNN and the Fisher information reached at the end of fitting. Note that since the fiducial parameter values are far from the "target" that this estimate of the Fisher information as the covariance will likely be misleading.

```python
GA = imnn.lfi.GaussianApproximation(
    parameter_estimates=IMNN.get_estimate(np.expand_dims(δ_target, (0, 1, 2))),
    invF=np.expand_dims(np.linalg.inv(IMNN.F), 0),
    prior=prior,
    gridsize=100)
GA.marginal_plot(
    known=θ_target,
    label="Gaussian approximation",
    axis_labels=["A", "B"],
    colours="C1");
```

And finally we can do an approximate Bayesian computation

```python
ABC = imnn.lfi.ApproximateBayesianComputation(
    target_data=np.expand_dims(δ_target, (0, 1, 2)),
    prior=prior,
    simulator=lambda rng, θ: simulator(rng, θ, simulator_args={**simulator_args, **{'squeeze':False}}),
    compressor=IMNN.get_estimate,
    gridsize=100,
    F=np.expand_dims(IMNN.F, 0))
```

```python
rng, key = jax.random.split(rng)
ABC(ϵ=1., rng=key, n_samples=10000, min_accepted=1000,
    smoothing=1, max_iterations=1000);
```

    [1010] accepted in last  43 iterations  (430000 simulations done).

```python
ax = LFI.marginal_plot(
    known=θ_target,
    label="Analytic likelihood",
    axis_labels=["A", "B"])

GA.marginal_plot(
    ax=ax,
    label="Gaussian approximation",
    colours="C1",
    axis_labels=["A", "B"])

ABC.marginal_plot(
    ax=ax,
    label="Approximate Bayesian computation",
    colours="C2");

plt.show()

```

# Cosmological parameter inference of log normal fields

As a more realistic example of cosmological parameter inference from dark matter fields, albeit it one where we do not (yet) know the amount of information in the field, we can create a log normal field from a power spectrum generated with cosmological parameters.

For example lets say that our fiducial cosmology has $\Omega_c=0.40$ and $\sigma_8=0.75$, we can set

```python
cosmo_params = jc.Planck15(Omega_c=0.40, sigma8=0.75)
θ_fid = np.array(
    [cosmo_params.Omega_c,
     cosmo_params.sigma8],
    dtype=np.float32)
```

Our new $P(k)$ is simply the linear matter power spectrum defined as

```python
def P(k, A=0.40, B=0.75):
    cosmo_params = jc.Planck15(Omega_c=A, sigma8=B)
    return jc.power.linear_matter_power(cosmo_params, k)
```

```python
simulator_args = dict(
    k=k,
    L=250,
    shape=shape,
    vol_norm=True,
    N_scale=False,
    squeeze=True,
    log_normal=True)
```

Since our lognormal field simulator *and* power spectra code are differentiable via `Jax`, we can simulate a *differentiable* universe. We'll pull out a nice function to visualize fiducial example data and its derivatives with respect to the cosmological parameters.

```python
from imnn.utils import value_and_jacrev, value_and_jacfwd

def simulator_gradient(rng, θ, simulator_args=simulator_args):
    return value_and_jacrev(simulator, argnums=1, allow_int=True, holomorphic=True)(rng, θ, simulator_args=simulator_args)

simulation, simulation_gradient = value_and_jacfwd(simulator, argnums=1)(rng, θ_fid,
                                                        simulator_args=simulator_args)

from mpl_toolkits.axes_grid1 import make_axes_locatable

fig,ax = plt.subplots(nrows=1, ncols=3, figsize=(12,15))

im1 = ax[0].imshow(np.squeeze(simulation), extent=(0,1,0,1))
ax[0].title.set_text(r'example fiducial $\rm d$')
divider = make_axes_locatable(ax[0])
cax = divider.append_axes('right', size='5%', pad=0.05)
fig.colorbar(im1, cax=cax, orientation='vertical')

im1 = ax[1].imshow(np.squeeze(simulation_gradient).T[0].T, extent=(0,1,0,1))
ax[1].title.set_text(r'$\nabla_{\Omega_m} \rm d$')
divider = make_axes_locatable(ax[1])
cax = divider.append_axes('right', size='5%', pad=0.05)
fig.colorbar(im1, cax=cax, orientation='vertical')

im1 = ax[2].imshow(np.squeeze(simulation_gradient).T[1].T, extent=(0,1,0,1))
ax[2].title.set_text(r'$\nabla_{\sigma_8} \rm d$')
divider = make_axes_locatable(ax[2])
cax = divider.append_axes('right', size='5%', pad=0.05)
fig.colorbar(im1, cax=cax, orientation='vertical')

for a in ax:
    a.set_xticks([])
    a.set_yticks([])

plt.show()
```

We'll now make the target universe that we observe a little more realistic with $\Omega_c=0.35$ and $\sigma_8=0.8$

```python
θ_target = np.array([0.35, 0.8])

rng, key = jax.random.split(rng)
δ_target = simulator(
    key, θ_target, simulator_args=simulator_args,
)
plt.imshow(δ_target)
plt.colorbar();
```

We can now train an IMNN as before

```python
rng, key = jax.random.split(rng)
IMNN = imnn.IMNN(
        n_s=n_s,
        n_d=n_d,
        n_params=n_params,
        n_summaries=n_summaries,
        input_shape=(1, 1) + shape,
        θ_fid=θ_fid,
        model=model,
        optimiser=optimiser,
        key_or_state=key,
        simulator=lambda rng, θ: simulator(
            rng, θ, simulator_args={**simulator_args, **{"squeeze": False}}))
```

    `simulator` provided, using SimulatorIMNN

```python
βe
rng, key = jax.random.split(rng)
IMNN.fit(λ=10., ϵ=0.1, rng=key, print_rate=None, min_iterations=500, best=False)
```

    CPU times: user 19min 21s, sys: 3min 38s, total: 23min
    Wall time: 23min 6s

```python
IMNN.plot(expected_detF=None);
```

And finally we can do our inference. We'll first set the prior distribution

```python
prior = tfp.distributions.Blockwise(
    [tfp.distributions.Uniform(low=low, high=high)
     for low, high in zip([0., 0.], [1.5, 1.5])])
prior.low = np.array([0., 0.])
prior.high = np.array([1.5, 1.5])
```

And make the Gaussian approximation using the Fisher information

```python
GA = imnn.lfi.GaussianApproximation(
    parameter_estimates=IMNN.get_estimate(np.expand_dims(δ_target, (0, 1, 2))),
    invF=np.expand_dims(np.linalg.inv(IMNN.F), 0),
    prior=prior,
    gridsize=100)
```

And then run the ABC

```python
ABC = imnn.lfi.ApproximateBayesianComputation(
    target_data=np.expand_dims(δ_target, (0, 1, 2)),
    prior=prior,
    simulator=lambda rng, θ: simulator(rng, θ, simulator_args={**simulator_args, **{'squeeze':False}}),
    compressor=IMNN.get_estimate,
    gridsize=100,
    F=np.expand_dims(IMNN.F, 0))
```

```python
%%time
rng, key = jax.random.split(rng)
ABC(ϵ=0.1, rng=key, n_samples=10000, min_accepted=5000,
    smoothing=1., max_iterations=5000);
```

And then we can plot the constraints obtained using the IMNN and LFI

```python
ax = GA.marginal_plot(
    known=θ_target,
    label="Gaussian approximation",
    axis_labels=[r"$\Omega_c$", r"$\sigma_8$"],
    colours="C1")
ABC.marginal_plot(
    ax=ax,
    label="Approximate Bayesian computation",
    colours="C2");
```

```python

```


# 2022-01-26 Newton methods

* Office hours: Monday 9-10pm, Tuesday 2-3pm, Thursday 2-3pm
* This week will stay virtual. Plan is to start in-person the following Monday (Jan 31)

## Last time

* Discuss rootfinding as a modeling tool
* Limitations of bisection
* Convergence classes
* Intro to Newton methods

## Today

* Newton's method via Taylor series
* Convergence theory for fixed point methods
* Derive Newton's method via convergence theory
* Newton methods in computing culture
* Breaking Newton's method

```julia
using Plots
default(linewidth=4, legendfontsize=12)

f(x) = cos(x) - x
hasroot(f, a, b) = f(a) * f(b) < 0
function bisect_iter(f, a, b, tol)
    hist = Float64[]
    while abs(b - a) > tol
        mid = (a + b) / 2
        push!(hist, mid)
        if hasroot(f, a, mid)
            b = mid
        else
            a = mid
        end
    end
    hist
end
```

    bisect_iter (generic function with 1 method)

# Convergence classes

A convergent rootfinding algorithm produces a sequence of approximations $x_k$ such that $$\lim_{k \to \infty} x_k \to x_*$$ where $f(x_*) = 0$.  For analysis, it is convenient to define the errors $e_k = x_k - x_*$. We say that an iterative algorithm is **$q$-linearly convergent** if $$\lim_{k \to \infty} |e_{k+1}| / |e_k| = \rho < 1.$$  (The $q$ is for "quotient".)  A smaller convergence factor $\rho$ represents faster convergence.  A slightly weaker condition ($r$-linear convergence or just **linear convergence**) is that
$$ |e_k| \le \epsilon_k $$
for all sufficiently large $k$ when the sequence $\{\epsilon_k\}$ converges $q$-linearly to 0.

```julia
hist = bisect_iter(f, -1, 3, 1e-10)
r = hist[end] # What are we trusting?
hist = hist[1:end-1]
scatter( abs.(hist .- r), yscale=:log10)
ks = 1:length(hist)
ρ = 0.5
plot!(ks, 4 * (ρ .^ ks))
```

# Newton-Raphson Method

Much of numerical analysis reduces to [Taylor series](https://en.wikipedia.org/wiki/Taylor_series), the approximation
$$ f(x) = f(x_0) + f'(x_0) (x-x_0) + f''(x_0) (x - x_0)^2 / 2 + \underbrace{\dotsb}_{O((x-x_0)^3)} $$
centered on some reference point $x_0$.

In numerical computation, it is exceedingly rare to look beyond the first-order approximation
$$ \tilde f_{x_0}(x) = f(x_0) + f'(x_0)(x - x_0) . $$
Since $\tilde f_{x_0}(x)$ is a linear function, we can explicitly compute the unique solution of $\tilde f_{x_0}(x) = 0$ as
$$ x = x_0 - \frac{f(x_0)}{f'(x_0)} . $$
This is Newton's Method (aka Newton-Raphson or Newton-Raphson-Simpson) for finding the roots of differentiable functions.

# An implementation

```julia
function newton(f, fp, x0; tol=1e-8, verbose=false)
    x = x0
    for k in 1:100 # max number of iterations
        fx = f(x)
        fpx = fp(x)
        if verbose
            println("[$k] x=$x  f(x)=$fx  f'(x)=$fpx")
        end
        if abs(fx) < tol
            return x, fx, k
        end
        x = x - fx / fpx
    end
end

f(x) = cos(x) - x
fp(x) = -sin(x) - 1
newton(f, fp, 1; tol=1e-15, verbose=true)
```

    [1] x=1  f(x)=-0.45969769413186023  f'(x)=-1.8414709848078965
    [2] x=0.7503638678402439  f(x)=-0.018923073822117442  f'(x)=-1.6819049529414878
    [3] x=0.7391128909113617  f(x)=-4.6455898990771516e-5  f'(x)=-1.6736325442243012
    [4] x=0.739085133385284  f(x)=-2.847205804457076e-10  f'(x)=-1.6736120293089505
    [5] x=0.7390851332151607  f(x)=0.0  f'(x)=-1.6736120291832148

    (0.7390851332151607, 0.0, 5)

# That's really fast!

* 10 digits of accuracy in 4 iterations.
* How is this convergence test different from the one we used for bisection?
* How can this break down?

$$ x_{k+1} = x_k - \frac{f(x_k)}{f'(x_k)} $$

```julia
newton(f, fp, -pi/2+.1; verbose=true)
```

    [1] x=-1.4707963267948965  f(x)=1.5706297434417247  f'(x)=-0.0049958347219742905
    [2] x=312.9170549232224  f(x)=-312.59435002533314  f'(x)=-0.05350037037604283
    [3] x=-5529.927542752894  f(x)=5530.676391917825  f'(x)=-0.33725953180603474
    [4] x=10868.945936970244  f(x)=-10868.376227850798  f'(x)=-0.17815359146505727
    [5] x=-50136.70732252356  f(x)=50135.70777741902  f'(x)=-1.0301593101044748
    [6] x=-1468.7903453577164  f(x)=1468.8859787856973  f'(x)=-1.9954166200403913
    [7] x=-732.6603742863299  f(x)=731.8761094362295  f'(x)=-1.6204261800544972
    [8] x=-281.0038172368656  f(x)=280.8358913898124  f'(x)=-1.9857996296872256
    [9] x=-139.58174866993488  f(x)=139.79912372811944  f'(x)=-0.02391184615360531
    [10] x=5706.856156210999  f(x)=-5707.008659764705  f'(x)=-1.9883029222393844
    [11] x=2836.5648158265976  f(x)=-2837.5220962814674  f'(x)=-1.2891610809292957
    [12] x=635.503879177757  f(x)=-634.8839651181479  f'(x)=-1.784669713127157
    [13] x=279.7607629875442  f(x)=-280.7481464325292  f'(x)=-0.8416524942745961
    [14] x=-53.80700796583557  f(x)=52.88592082634005  f'(x)=-1.3893564966145653
    [15] x=-15.74196061822768  f(x)=14.742538472479499  f'(x)=-1.0339908015219368
    [16] x=-1.4840596284124175  f(x)=1.5706876106501757  f'(x)=-0.0037592697076601622
    [17] x=416.333158287511  f(x)=-416.4052274406877  f'(x)=-1.99739963763799
    [18] x=207.8594910282616  f(x)=-206.9888910802111  f'(x)=-1.4919915959184906
    [19] x=69.12620885264326  f(x)=-68.1262712417355  f'(x)=-1.0111702413616044
    [20] x=1.7525179991632598  f(x)=-1.933241162917233  f'(x)=-1.9835340045381016
    [21] x=0.7778731690296721  f(x)=-0.0654654835715871  f'(x)=-1.7017658368004631
    [22] x=0.7394040200105153  f(x)=-0.0005337303513481828  f'(x)=-1.6738476794194503
    [23] x=0.7390851556610822  f(x)=-3.756576461011463e-8  f'(x)=-1.6736120457726615
    [24] x=0.7390851332151608  f(x)=-2.220446049250313e-16  f'(x)=-1.6736120291832148

    (0.7390851332151608, -2.220446049250313e-16, 24)

# Convergence of fixed-point (by mean value theorem)

Consider the iteration
$$x_{k+1} = g(x_k)$$
where $g$ is a continuously differentiable function.
Suppose that there exists a fixed point $x_* = g(x_*)$.  By the [mean value theorem](https://en.wikipedia.org/wiki/Mean_value_theorem), we have that
$$ x_{k+1} - x_* = g(x_k) - g(x_*) = g'(c_k) (x_k - x_*) $$
for some $c_i$ between $x_k$ and $x_*$.

Taking absolute values, $$|e_{k+1}| = |g'(c_k)| |e_k|,$$ which converges to zero if $|g'(c_k)| < 1$.

# Convergence of fixed-point (by Taylor series)

Consider the iteration
$$x_{k+1} = g(x_k)$$
where $g$ is a continuously differentiable function.
Suppose that there exists a fixed point $x_* = g(x_*)$. There exists a Taylor series at $x_*$,
$$ g(x_k) = g(x_*) + g'(x_*)(x_k - x_*) + O((x_k-x_*)^2) $$
and thus
\begin{align}
x_{k+1} - x_* &= g(x_k) - g(x_*) \\
&= g'(x_*) (x_k - x_*) + O((x_k - x_*)^2).
\end{align}

In terms of the error $e_k = x_k - x_*$,
$$ \left\lvert \frac{e_{k+1}}{e_k} \right\rvert = \lvert g'(x_*) \rvert + O(e_k).$$

## Poll: Is this convergence A=q-linear, B=r-linear, C=neither?

Recall the definition of q-linear convergence
$$ \lim_{k\to\infty} \left\lvert \frac{e_{k+1}}{e_k} \right\rvert = \rho < 1. $$

# Aside: [Big $O$ ("big oh") notation](https://en.wikipedia.org/wiki/Big_O_notation)

## Limit $n\to\infty$

We'd say an algorithm costs $O(n^2)$ if its running time on input of size $n$ is less than $c n^2$ for some constant $c$ and sufficiently large $n$.

Sometimes we write $\operatorname{cost}(\texttt{algorithm}, n) = O(n^2)$ or (preferably) $\operatorname{cost}(\texttt{algorithm}) \in O(n^2)$.

Note that $O(\log n) \subset O(n) \subset O(n\log n) \subset O(n^2) \subset \dotsb$ so it's correct to say "binary search is in $O(n^2)$", even though a sharper statement is also true.

We say the algorithm is in $\Theta(n^2)$ ("big theta") if
$$ c_1 n^2 < \operatorname{cost}(\texttt{algorithm}) < c_2 n^2 $$
for some positive constants $c_1,c_2$ and sufficiently large $n$.

## Limit $h \to 0$

In numerical analysis, we often have a small real number, and now the definitions take the limit as the small number goes to zero. So we say a term in an expression is in $O(h^2)$ if

$$ \lim_{h\to 0} \frac{\operatorname{term}(h)}{h^2} < \infty . $$

Big $O$ terms can be manipulated as

\begin{align}
h O(h^k) &= O(h^{k+1}) \\
O(h^k)/h &= O(h^{k-1}) \\
c O(h^k) &= O(h^k) \\
O(h^k) - O(h^k) &= ?
\end{align}

# Example of a fixed point iteration

We wanted to solve $\cos x - x = 0$, which occurs when $g(x) = \cos x$ is a fixed point.

```julia
xstar, _ = newton(f, fp, 1.)
g(x) = cos(x)
gp(x) = -sin(x)
@show xstar
@show gp(xstar)
plot([x->x, g], xlims=(-2, 3))
scatter!([xstar], [xstar],
    label="\$x_*\$")
```

    xstar = 0.739085133385284
    gp(xstar) = -0.6736120293089505

```julia
function fixed_point(g, x, n)
    xs = [x]
    for k in 1:n
        x = g(x)
        append!(xs, x)
    end
    xs
end

xs = fixed_point(g, 2., 15)
plot!(xs, g.(xs), seriestype=:path, marker=:auto)
```

# Verifying fixed point convergence theory

$$ \left\lvert \frac{e_{k+1}}{e_k} \right\rvert \to \lvert g'(x_*) \rvert $$

```julia
@show gp(xstar)
es = xs .- xstar
es[2:end] ./ es[1:end-1]
```

    gp(xstar) = -0.6736120293089505

    15-element Vector{Float64}:
     -0.9161855415615605
     -0.15197657010596488
     -0.734870205299266
     -0.624132525531327
     -0.7026257933893496
     -0.6523498121376077
     -0.6870971782336925
     -0.664168570025122
     -0.6798044680427148
     -0.6693659427636027
     -0.6764378047956165
     -0.6716930541785153
     -0.6748976495459512
     -0.6727427617641084
     -0.6741962236114177

```julia
scatter(abs.(es), yscale=:log10, label="fixed point")
plot!(k -> abs(gp(xstar))^k, label="\$|g'|^k\$")
```

# Plotting Newton convergence

```julia
function newton_hist(f, fp, x0; tol=1e-12)
    x = x0
    hist = []
    for k in 1:100 # max number of iterations
        fx = f(x)
        fpx = fp(x)
        push!(hist, [x fx fpx])
        if abs(fx) < tol
            return vcat(hist...)
        end
        x = x - fx / fpx
    end
end
```

    newton_hist (generic function with 1 method)

```julia
xs = newton_hist(f, fp, 1.97)
@show x_star = xs[end,1]
plot(xs[1:end-1,1] .- x_star, yscale=:log10, marker=:auto)
```

    x_star = xs[end, 1] = 0.7390851332151607

## Poll: Is this convergence A=q-linear, B=r-linear, C=neither?

# Formulations are not unique (constants)

If $x = g(x)$ then
$$x = \underbrace{x + c(g(x) - x)}_{g_2}$$
for any constant $c \ne 0$. Can we choose $c$ to make $\lvert g_2'(x_*) \rvert$ small?

```julia
c = .5
g2(x) = x + c * (cos(x) - x)
g2p(x) = 1 + c * (-sin(x) - 1)
@show g2p(xstar)
plot([x->x, g, g2], ylims=(-5, 5), label=["x" "g" "g2"])
```

    g2p(xstar) = 0.16319398534552476

```julia
xs = fixed_point(g2, 1., 15)
xs .- xstar
```

    16-element Vector{Float64}:
      0.26091486661471597
      0.03106601954878585
      0.004893162344945079
      0.0007941171212053622
      0.00012947850276123773
      2.112687301181193e-5
      3.4475537732392425e-6
      5.62475483634195e-7
      9.16501970982253e-8
      1.4814399151852342e-8
      2.2752605355336186e-9
      2.2894852680366284e-10
     -1.0499723313017739e-10
     -1.594951948291623e-10
     -1.683889694348295e-10
     -1.6984036399492197e-10

# Formulations are not unique (functions)

If $x = g(x)$ then
$$x = \underbrace{x + h(x) \big(g(x) - x\big)}_{g_3(x)}$$
for any smooth $h(x) \ne 0$. Can we choose $h(x)$ to make $\lvert g_3'(x) \rvert$ small?

```julia
h(x) = -1 / (gp(x) - 1)
g3(x) = x + h(x) * (g(x) - x)
plot([x-> x, cos, g2, g3], ylims=(-5, 5))
```

* We don't know $g'(x_*)$ in advance because we don't know $x_*$ yet.
* This method converges very fast
* We actually just derived Newton's method.

# A fresh derivation of Newton's method

* A rootfinding problem $f(x) = 0$ can be converted to a fixed point problem $$x = x + f(x) =: g(x)$$ but there is no guarantee that $g'(x_*) = 1 + f'(x_*)$ will have magnitude less than 1.
* Problem-specific algebraic manipulation can be used to make $|g'(x_*)|$ small.
* $x = x + h(x) f(x)$ is also a valid formulation for any $h(x)$ bounded away from $0$.
* Can we choose $h(x)$ such that $$ g'(x) = 1 + h'(x) f(x) + h(x) f'(x) = 0$$ when $f(x) = 0$?

In other words,
$$ x_{k+1} = x_k + \underbrace{\frac{-1}{f'(x_k)}}_{h(x_k)} f(x_k)  . $$

# Quadratic convergence!

$$ \left\lvert \frac{e_{k+1}}{e_k} \right\rvert \to \lvert g'(x_*) \rvert $$

* What does it mean that $g'(x_*) = 0$?
* It turns out that Newton's method has _locally quadratic_ convergence to simple roots,
$$\lim_{k \to \infty} \frac{|e_{k+1}|}{|e_k|^2} < \infty.$$
* "The number of correct digits doubles each iteration."
* Now that we know how to make a good guess accurate, the effort lies in getting a good guess.

# Culture: fast inverse square root

The following code appeared literally (including comments) in the Quake III Arena source code (late 1990s).

```C
float Q_rsqrt( float number )
{
	long i;
	float x2, y;
	const float threehalfs = 1.5F;

	x2 = number * 0.5F;
	y  = number;
	i  = * ( long * ) &y;                       // evil floating point bit level hacking
	i  = 0x5f3759df - ( i >> 1 );               // what the fuck?
	y  = * ( float * ) &i;
    y  = y * ( threehalfs - ( x2 * y * y ) );   // 1st iteration
//  y  = y * ( threehalfs - ( x2 * y * y ) );   // 2nd iteration, this can be removed

	return y;
}
```

We now have [vector instructions](https://software.intel.com/sites/landingpage/IntrinsicsGuide/#text=rsqrt&expand=2989,1224,4470) for approximate inverse square root.
More at https://en.wikipedia.org/wiki/Fast_inverse_square_root

# How does it work?

Let's look at the last line
```c
y  = y * ( threehalfs - ( x2 * y * y ) );   // 1st iteration
```

We want a function $f(y)$ such that $f(1/\sqrt{x}) = 0$. One such function is
$$ f(y) = 1/y^2 - x, \quad f'(y) = -2/y^3.$$

There are others, e.g.,
$$f_1(y) = y^2 - 1/x,\quad f'(y) = 2 y,$$
but this would require a division.

Newton's method is
\begin{align}
y_{k+1} &= y_k - \frac{f(y_k)}{f'(y_k)} \\
&= y_k - \frac{1/y_k^2 - x}{-2/y_k^3} \\
&= y_k + \frac 1 2 (y_k - x y_k^3) \\
&= y_k \left(\frac 3 2 - \frac 1 2 x y_k^2\right)
\end{align}

# Rootfinding outlook

* Newton methods are immensely successful
  * Convergence theory is local; we need good initial guesses (activity)
  * Computing the derivative $f'(x)$ is *intrusive*
    * Avoided by secant methods (approximate the derivative; activity)
    * Algorithmic or numerical differentiation (future topics)
  * Bisection is robust when conditions are met
  * Line search (activity)
  * When does Newton diverge?

* More topics
  * Find *all* the roots
  * Use Newton-type methods with bounds
  * Times when Newton converges slowly


```python
import sympy
sympy.init_printing()
from sympy import I,pi,oo
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
```

```python
x = sympy.Symbol("x")
y = sympy.Symbol("y",real=True)
sympy.sqrt(x**3)
```

```python
x = sympy.Symbol("x")
expr = 1+2*x**2+3*x**3
expr
```

```python
expr.args
```

```python
expr.args[1]
```

```python
expr = 2*sympy.cos(x)*sympy.sin(x)
```

```python
expr
```

```python
sympy.simplify(expr)
```

```python
expr = sympy.exp(x) * sympy.exp(y)
```

```python
expr
```

```python
distances = [10,15,17,26,20]
times = [0.3,0.47,0.55,1.20,1.0]

```

```python
speeds = []
for i in range(len(distances)):
    speeds.append(distances[i]/times[i])

speeds
```

    [33.333333333333336,
     31.914893617021278,
     30.909090909090907,
     21.666666666666668,
     20.0]

```python
[d/t for d,t in zip(distances,times)]
```

    [33.333333333333336,
     31.914893617021278,
     30.909090909090907,
     21.666666666666668,
     20.0]

```python
product_quatities = [13,5,6,10,11]
prices = [1.2,6.5,1.0,4.8,5.0]
total = sum([q*p for q,p in zip(product_quatities,prices)])
total
```

    157.1

```python
distances = np.array(distances)
times = np.array(times)
product_quatities = np.array(product_quatities)
prices = np.array(prices)
```

```python
A = np.ones(shape=(1,3),dtype=float)
B = np.ones(shape=(3,1),dtype=float)
C = (A*B)
C
```

    array([[1., 1., 1.],
           [1., 1., 1.],
           [1., 1., 1.]])

```python
values = product_quatities*prices
total = values.sum()
print(values)
total
```

```python
two_dim = np.array([[3,5,2,4],[7,6,5,5],[1,6,-1,-1]])
two_dim
```

    array([[ 3,  5,  2,  4],
           [ 7,  6,  5,  5],
           [ 1,  6, -1, -1]])

```python
two_dim[:,1:3]
```

    array([[ 5,  2],
           [ 6,  5],
           [ 6, -1]])

```python
two_dim.flatten()
```

    array([ 3,  5,  2,  4,  7,  6,  5,  5,  1,  6, -1, -1])

```python
def f(t):
    return np.exp(-t) * np.sin(2*np.pi*t)

t1 = np.arange(0.0,5.0,0.1)
t2 = np.arange(0.0,5.0,0.02)
plt.figure()
plt.subplot(2,1,1)
plt.plot(t1,f(t1),'bo')
```

```python
x = np.linspace(start=-5,stop=5,num=150)
fig, ax = plt.subplots(figsize = (7,4))
ax.plot(x,x,label='linear')
ax.plot(x,x**2,label='quadratic')
ax.plot(x,x**3,label='cubic')
ax.plot(x,x**4,label='polynomial')
ax.legend()
```

```python
fig,axes = plt.subplots(nrows=2,ncols=5,figsize=(14,4.5))
for i, ax in enumerate(axes.flatten()):
    ax.set_title("x to the {:d}".format(i+1))
    ax.plot(x,x**(i+1))

fig.tight_layout();
```

```python

```

```python

```


# Lecture 1: Matrices & vectors

## Syllabus

**Week 2:** Matrices, vectors, norms, ranks

##What is a matrix

A matrix is a two-dimensional table.
Here is an example of a $3 \times 3$ matrix
\begin{equation}
A = \begin{pmatrix}
a_{11} & a_{12} & a_{13} \\
a_{21} & a_{22} & a_{23} \\
a_{31} & a_{32} & a_{33}
\end{pmatrix}
\end{equation}

A vector is a $n \times 1$ vector (there are **row** and **column** vectors).

## What we can do with matrices and vectors

1. Add them: $a = b + c$
2. Multiply by numbers

## Matrix as a linear operator
Matrix is typically used to encode a **linear operator**:
$$ y = A x,$$
called matrix-by-vector product, in the index form
$$y_i = \sum_{j=1}^m A_{ij} x_j, \quad i = 1, \ldots, n.$$

## Linear dependencies
Many physical models are formulated as **linear equations**:

- Newton law $F = ma$
- Hookes law $F = kx$

which is based on the fact that if the change is small, everything can be approximated by a linear function:
$$
   f(x + \delta x) \approx f(x) + \delta x.
$$
Of course, nonlinearities may come into play, but even in this case the numerical methods **linearize** the problem around the current approximation. But matrices and linear dependence also play important role in data analysis as well.

## Linear dependencies in real life
Matrix encodes linear dependence.
Linear dependence is the simplest and often very efficient model for the data. We will give two illustrations:
- Principal component analysis
- Independent component analysis

## Demo: principal component analysis
One of the basic factor models, factor analysis:
$y_1, \ldots, y_P$ are vectors (data points), that are observed. We think of the as **linear mixture**. We generate random points on a plane and then rotate them
by a certain mixture.

We generate a sequence of random points in two dimensions, and setup a rotation matrix $A$.

```
%matplotlib inline
import matplotlib.pyplot as plt
import prettyplotlib as ppl
from sklearn.decomposition import PCA
import numpy as np

P = 1000
points = np.random.randn(P,2)
A = [[2, 1], [0, 1]]
A = np.array(A)
ppl.plot(points[:, 0], points[:, 1], ls='', marker='o')
```

We can also plot the rotated points: look at how they are "skewed".

```
R = np.dot(points, A)
ppl.plot(R[:, 0], R[:, 1], ls='', marker='o')
```

We can also rotate them back by finding **principal components**, which is equivalent to singular value decomposition (SVD). We will talk about SVD later, now we just call it from **numpy.linalg** module. As we can see, the points are back to symmetric distribution, so we have recovered the matrix.
Principal component analysis is used for many applications, in particular, for the **classification problem**: after the rotation, the points, belonging to to different classes, can be visually separated.

```
%matplotlib inline
u, s, v = np.linalg.svd(R, full_matrices=False)
unrotated = R.dot(v.T)
ppl.plot(unrotated[:, 0], unrotated[:, 1], ls='', marker='o')
```

## Demo: Cocktail party problem
The linear models and the factors may have a real physical meaning.
One of the most interesting illustrations is the **cocktail party problem**, which is defined as follows.
We have a set of sources $x(t)$ (people talking) and a set of microphones.
At each microphone we record a **linear mixture**:
$$
    y = A x(t) + \eta(t),
$$
where $\eta(t)$ is some noise.
We do not know $A$ and want to recover it.

[Demo](bss.ipynb)

##Matrix-by-matrix product

Consider composition of two linear operators:

1. $y = Bx$
2. $z = Ay$

Then, $z = Ay =  A B x = C x$, where $C$ is the **matrix-by-matrix product**.

A product of an $n \times k$ matrix $A$ and a $k \times m$ matrix $B$ is a $n \times m$ matrix $C$ with the elements
$$
   c_{ij} = \sum_{s=1}^k a_{is} b_{sj}, \quad i = 1, \ldots, n, \quad j = 1, \ldots, m
$$

Complexity of a naive algorithm for MM is $\mathcal{O}(n^3)$.
Matrix-by-matrix product is the **core** for almost all efficient algorithms in linear algebra. Basically, all the NLA algorithms are reduced to a sequence of matrix-by-matrix products, so efficient implementation of MM reduces the complexity of numerical algorithms by the same factor. However, implementing MM is not easy at all!

## Efficient implementation for MM
Is it easy to multiply a matrix by a matrix?
The answer is: **no**, if you want it as fast as possible, using the computers that are at hand.

## Demo
Let us do a short demo and compare a `np.dot()` procedure which in my case uses MKL with a hand-written matrix-by-matrix routine in Python and also its Cython version (and also gives a very short introduction to Cython).

```
import numpy as np
def matmul(a, b):
    n = a.shape[0]
    k = a.shape[1]
    m = b.shape[1]
    c = np.zeros((n, m))
    for i in xrange(n):
        for j in xrange(m):
            for s in xrange(k):
                c[i, j] += a[i, s] * b[s, j]
```

`load_ext cythonmagic`   makes it possible to use `%%cython` magic for an automatic compilation of the [Cython code](http://cython.org/).
Cython is a Python compiler with additional directives that allow to specify the type of some variables. In many cases it may lead to significant speedup.
For other possibilities you may look at the [following comparison](http://nbviewer.ipython.org/url/jakevdp.github.io/downloads/notebooks/NumbaCython.ipynb)

```
%load_ext cythonmagic
```

```cython
%%cython
import numpy as np
def cython_matmul(double [:, :] a, double[:, :] b):
    cdef int n = a.shape[0]
    cdef int k = a.shape[1]
    cdef int m = b.shape[1]
    cdef int i
    cdef int j
    cdef int s
    c = np.zeros((n, m))
    cdef double[:, :] cview = c
    for i in xrange(n):
        for j in xrange(m):
            for s in xrange(k):
                c[i, j] += a[i, s] * b[s, j]
    return c
```

Then we just time three different routines. Guess the answer.

```
n = 100
a = np.random.randn(n, n)
b = np.random.randn(n, n)
%timeit c = matmul(a, b)
%timeit cf = cython_matmul(a, b)
%timeit c = np.dot(a, b)
```

    1 loops, best of 3: 1.57 s per loop
    1 loops, best of 3: 567 ms per loop
    10000 loops, best of 3: 110 µs per loop

Why it is so?
There are two important issues:

- Computers are more and more parallel (multicore, graphics processing units)
- The memory pyramid: there is a whole hierarchy of levels

## Memory architecture

Fast memory is small, bigger memory is slow.

- Data fits into the  fast memory: Load all data, do computations
- Data does not fit into the fast memory:  load data by chunks, do computations, load again

We need to reduce the number of read/write operations! This is typically achieved in efficient implementations of the BLAS libraries, one of which (Intel MKL) we now use.

## BLAS
Basic linear algebra operations (**BLAS**) have three levels:
1. BLAS-1, operations like $c = a + b$
2. BLAS-2, operations like matrix-by-vector product
3. BLAS-3, matrix-by-matrix product

What is the principal differences between those?

The main difference is the number of operations vs. the number of data!

1. BLAS-1: $\mathcal{O}(n)$ data, $\mathcal{O}(n)$ operations
2. BLAS-2: $\mathcal{O}(n^2)$ data, $\mathcal{O}(n^2)$ operations
3. BLAS-3: $\mathcal{O}(n^2)$ data, $\mathcal{O}(n^3)$ operations

**Remark**: a quest for $\mathcal{O}(n^2)$ matrix-by-matrix multiplication algorithm is not yet done.

Strassen gives $\mathcal{O}(n^{2.78...})$
World record $\mathcal{O}(n^{2.37})$ [Reference](http://arxiv.org/pdf/1401.7714v1.pdf)
The constant is unfortunately too big to make it practical!

## Memory hierarchy
How we can use memory hierarchy?

Break the matrix into blocks! ($2 \times 2$ is an **illustration**)
$$
   A = \begin{bmatrix}
         A_{11} & A_{12} \\
         A_{21} & A_{22}
        \end{bmatrix}, \quad B = \begin{bmatrix}
         B_{11} & B_{12} \\
         B_{21} & B_{22}
        \end{bmatrix}
$$

Then,
$$
   A B = \begin{bmatrix}
            A_{11} B_{11} + A_{12} B_{21} & A_{11} B_{12} + A_{12} B_{22} \\
            A_{21} B_{11} + A_{22} B_{21} & A_{21} B_{12} + A_{22} B_{22}
         \end{bmatrix}
$$

If $A_{11}, B_{11}$ and their product fit into the cache memory (which is 1024 Kb for the [Haswell Intel Chip](http://en.wikipedia.org/wiki/List_of_Intel_Core_i7_microprocessors#.22Haswell-H.22_.28MCP.2C_quad-core.2C_22_nm.29)), then we load them only once into the memory.

### Key point
The number of read/writes is reduced by a factor $\sqrt{M}$, where $M$ is the cache size.

- Have to do linear algebra in terms of blocks!
- So, you can not even do Gaussian elimination as usual (or just suffer 10x performance loss)

## Parallelization
The blocking has also deep connection with parallel computations.
Consider adding two vectors:
$$ c = a + b$$
and we have two processors. How fast can we go?  Of course, not faster then twice.

```
## This demo requires Anaconda distribution to be installed
import mkl
import numpy as np
n = 1000
a = np.random.randn(n)
mkl.set_num_threads(1)
%timeit a + a
mkl.set_num_threads(2)
%timeit a + a
```

    100000 loops, best of 3: 2.21 µs per loop
    100000 loops, best of 3: 2.13 µs per loop

```
## This demo requires Anaconda distribution to be installed
import mkl
n = 500
a = np.random.randn(n, n)
mkl.set_num_threads(1)
%timeit a.dot(a)
mkl.set_num_threads(2)
%timeit a.dot(a)
```

    100 loops, best of 3: 12.7 ms per loop
    100 loops, best of 3: 10.8 ms per loop

Typically, two cases are distinguished:
1. Shared memory (i.e., multicore on every desktop/smartphone)
2. Distributed memory (i.e. each processor has its own memory, can send information through a network)

In both cases, the efficiency is governed by a **bandwidth**:
I.e., for BLAS-1 routines (like sum of two vectors) reads/writes take all the time.
For BLAS-3 routines, the speedup can be obtained that is more noticable.
For large-scale clusters (>100 000 cores, see the [Top500 list](http://www.top500.org/lists/)) there is still scaling.

##### Questions?

```
from IPython.core.display import HTML
def css_styling():
    styles = open("./styles/custom.css", "r").read()
    return HTML(styles)
css_styling()
```

<link href='http://fonts.googleapis.com/css?family=Fenix' rel='stylesheet' type='text/css'>
<link href='http://fonts.googleapis.com/css?family=Alegreya+Sans:100,300,400,500,700,800,900,100italic,300italic,400italic,500italic,700italic,800italic,900italic' rel='stylesheet' type='text/css'>
<link href='http://fonts.googleapis.com/css?family=Source+Code+Pro:300,400' rel='stylesheet' type='text/css'>
<style>
    @font-face {
        font-family: "Computer Modern";
        src: url('http://mirrors.ctan.org/fonts/cm-unicode/fonts/otf/cmunss.otf');
    }
    div.cell{
        width:80%;
        margin-left:auto !important;
        margin-right:auto;
    }
    h1 {
        font-family: 'Alegreya Sans', sans-serif;
    }
    h2 {
        font-family: 'Fenix', serif;
    }
    h3{
		font-family: 'Fenix', serif;
        margin-top:12px;
        margin-bottom: 3px;
       }
	h4{
		font-family: 'Fenix', serif;
       }
    h5 {
        font-family: 'Alegreya Sans', sans-serif;
    }
    div.text_cell_render{
        font-family: 'Alegreya Sans',Computer Modern, "Helvetica Neue", Arial, Helvetica, Geneva, sans-serif;
        line-height: 1.2;
        font-size: 160%;
        width:70%;
        margin-left:auto;
        margin-right:auto;
    }
    .CodeMirror{
            font-family: "Source Code Pro";
			font-size: 90%;
    }
/*    .prompt{
        display: None;
    }*/
    .text_cell_render h1 {
        font-weight: 200;
        font-size: 50pt;
		line-height: 100%;
        color:#CD2305;
        margin-bottom: 0.5em;
        margin-top: 0.5em;
        display: block;
    }
    .text_cell_render h5 {
        font-weight: 300;
        font-size: 16pt;
        color: #CD2305;
        font-style: italic;
        margin-bottom: .5em;
        margin-top: 0.5em;
        display: block;
    }

    .warning{
        color: rgb( 240, 20, 20 )
        }
</style>

```

```


# STA208 Machine Learning Memo

(Contents are originally provided by James Sharpnack, annotated and revised by Gento Kato)

***THIS MEMO IS NOT COMPLETE***

## Lecture 8 (04/27/2017)

### Margin Based Methods

#### Logistic Regression Loss

 * **Model**

 DV: $y \in \{-1, 1\}$

 Prediction: $\hat{y} = 1\{x^\top \hat{\beta} + \hat{\beta}_0 \geq 0\}$

 * **Loss**

  * 0-1 loss (what we want to minimize, but nonconvex):

  $$\mathit{l}_{0/1}(y_i, x_i \beta) = 1\{y_i \cdot x_i^\top \beta \leq 0 \}$$

  * Logistic loss (actually minimized in logit):

  \begin{align}
  \mathit{l}_{logit} (y_i, x_i, \beta) &=
  \begin{cases}
  log(1 + e^{x_i^\top \beta}), & y_i = -1 \\
  log(1 + e^{-x_i^\top \beta}), & y_i = 1
  \end{cases}
  \\ &= log(1+e^{-y_i x_i^T \beta})
  \end{align}

  * Square error loss (we can also use but...):

  $$\mathit{l}_2(y_i, x_i, \beta) = (y_i - x_i^\top \beta)^2 = (1- y_i x_i^\top \beta)^2$$

  * Logistic loss approximates 0-1 loss more than square error loss, but **computationally inefficient**.

### Support Vector Machines (SVM)


```python
import sympy as sp
```

```python
s = sp.Symbol('s')
```

```python
V = sp.Symbol('V')
```

```python
top = (2*V*s)*(1/s+2*s/(4*s**2+1)) - 0
```

```python
bottom = (1/s+(6*s+3)/(3*s+1))*(1/s+2*s/(4*s**2+1)) - (1/s)**2
```

```python
sp.simplify(top/bottom)
```

    2*V*s*(18*s**3 + 6*s**2 + 3*s + 1)/(36*s**3 + 24*s**2 + 8*s + 3)

```python

```


# Data Visualization with PCA

## The Challenges of High-dimensional Data

Once we decide to measure more than three features per input vector, it can become challenging to understand how a network is learning to solve a problem since we can no longer generate a plot or visualization of the feature vector space to which the network is being exposed. One-, two-, or three-dimensional vectors are easy enough to plot, and we could even color the corresponding points based on their class assignments to see how groups of similar items are in similar parts of the vector space. If we can see straight boundaries between the differently colored point clouds, then we would understand why a linear network might be capable of producing a reasonable solution to the problem at hand. Also, if we see regions where a single line will not suffice for separating the different classes, then we might have reason to suspect that a linear network will fail to solve the problem, and probably attempt to use a multilayer network instead.

Given the advantages of visualizing such relationships, a common trick when exploring high-dimensional data sets is to **project** the high-dimensional data vectors onto a low-dimensional space (two or three dimensions) where we can see if such relationships exist. Such projection is _risky_ in the sense that we will be throwing information away to perform the projection (similar to how neural networks throw information away when performing regression or classification), and we may no longer see some important relationships in the low-dimensional projection of the data. However, it is sometimes possible that the projection **will** preserve relationships between the original data vectors that are important for making an accurate classification or regression while also enabling visualization.

In this assignment, we will explore a commonly-used **linear** projection known as Principal Component Analysis (PCA) to visualize some data sets and see how such projections might be useful for understanding why our single-layer networks were able to effectively learn the functions that these data sets represent. Since PCA is a linear method, it will be limited in its ability to produce useful _projections_ in a manner roughly analogous to how single-layer neural networks are limited in their ability to solve _linearly separable_ problems. Since we will be producing low-dimensional projections, relative to the original dimensionality of the vectors, this technique is also a form of **dimensionality reduction** in that the new two- or three- dimensional vectors that we produce will still share some of the relationships between one another that the higher-dimensional vectors possessed. This might even mean that we could use these new vectors as inputs for a neural network instead of the original, high-dimensional vectors. This could lead to a significant reduction in the number of neural units and connection weights in a network, and hence reduce its computation time. Also, some of the original features (or combinations of features) may just not be very useful for the problem at hand, and removing them would allow the network to focus on only the more relevant features in the data. In some cases (albeit rarely with PCA), this can even lead to superior performance of the trained neural network overall. While we will focus on two-dimensional projections in this assignment, PCA can be used to reduce the dimensionality of a given set of input vectors to any chosen number less than or equal to the original dimensionality. The smaller the dimensionality of the projection: the more information is being projected away. Thus, two-dimensional projections are often too low to be of any real use on many large data sets. However, some data sets might be reduced from millions of features to thousands, or thousands to hundreds, while still preserving the vast majority of the information that they encode. Such large reductions can make certain problems far more tractable to learn than they would otherwise be.

## Gathering the Iris data set

Let's start out by grabbing some data that we are already a little familiar with, and see if we can use PCA to better understand why a linear network can learn to solve this function effectively.

Let's start by importing some tools for the job...

```python
# For reading data sets from the web.
import pandas

# For lots of great things.
import numpy as np

# To make our plots.
import matplotlib.pyplot as plt
%matplotlib inline

# Because sympy and LaTeX make
# everything look wonderful!
from sympy import *
init_printing(use_latex=True)
from IPython.display import display

# We will use this to check our implementation...
from sklearn.decomposition import PCA

# We will grab another data set using Keras
# after we finish up with Iris...
import keras
```

    Using TensorFlow backend.

Now let's grab the Iris data set and start projecting!

```python
iris_data = np.array(pandas.read_table("https://www.cs.mtsu.edu/~jphillips/courses/CSCI4850-5850/public/iris-data.txt",
                                       delim_whitespace=True,
                                       header=None))
```

```python
# Remember the data is composed of feature
# vectors AND class labels...
X = iris_data[:,0:4] # 0,1,2,3
Y = iris_data[:,4]   # 4

# Pretty-print with display()!
display(X.shape)
display(Y.shape)
display(Matrix(np.unique(Y)).T)
```

The Iris data set consists of 150 four-dimensional feature vectors, each one assigned to one of three class labels (0,1,2) corresponding to an iris species.

We could potentially use **four** dimensions to plot and understand this data. Namely, we could make 3D plots using the first three dimensions, and sort the points along the fourth dimension so that we could play them in-sequence like a movie. However, that can still be tricky to visualize since we may miss some relationships between frames in our "movie" if they are far apart in time. Potentially more useful would be to plot the first three dimensions in one plot, then the last three dimensions in another, where the two plots now share the middle two dimensions. Still, if relationships between the first and fourth dimensions were the most important, we might not see them very clearly using this presentation.

Let's see if a PCA projection down to just two dimensions would be more effective.

To do this we will be using some linear algebra that we have already seen before, but we need to process the data just a little before we can use those tools.

First, we will _mean-center_ the values of _each_ feature in the data. That is, we will find the _mean_ value of the first feature across all examples in the data set, and then subtract the mean value from this feature for all examples. We will perform the same operation for all four features as well, so that each feature will have its mean value effectively set to zero. You can think of this as moving the entire set of data vectors so that the average of the data vectors now lies at the value zero in all four dimensions. In other words, it's a _translation_ operation on the original data. The relative distances between all of the points is maintained, so all of the relationships between the data vectors important for classification is maintained as well.

We will use a custom function for this, that we apply to each of the columns using the `apply_along_axis()` function:

```python
# Mean center a vector
def mean_center(x):
    return x - np.mean(x)

Xcentered = np.apply_along_axis(mean_center,0,X)
```

Now that we have a mean-centered data matrix, we will use singular value decomposition to extract the left-singular vectors and singular-values of this matrix.

```python
U,S,V = np.linalg.svd(Xcentered,full_matrices=True)

# Percent variance accounted for
plt.plot(100.0*S/np.sum(S))
plt.ylabel('% Var')
plt.xlabel('Singular Value')
plt.show()
```

Each of the singular values indicate some amount of variance present in the original data set that is captured by the corresponding left-singular vector (column of U). They are sorted in order from largest to smallest when returned by the `svd()` function so that you can see that the largest amount of variance is captured by the first left-singular vector, the second most variance by the second singular-vector, and so on. Often, the sum of the singular values is calculated to obtain the _total_ variance in the data, and then used to normalize the variance to obtain the percentage of variance captured by each left-singular vector. Given the data above, it is clear that the first two vectors alone will account for over 85% of the total variance in the data, and should form a reasonable projection for the data set.

```python
# Variance accounted for in the first two principal components
100.0*(S[0]+S[1])/np.sum(S)
```

The singular-values (S) are mapped into a rectangular, diagonal matrix which is then multiplied by the left-singular vectors (U). The vectors in U are all unit length, so this operation effectively scales the length of the first 4 vectors by each of the corresponding singular values. These operations produce a rotated version of our original data set where the major orthogonal directions capturing the largest variance in the data lie along the principal axes. Each of these so-called principal components is a linear combination of our original feature vectors, and allows us to produce a projection onto a smaller set of these components by simply throwing away vectors associated with small singular values. Thus, while we obtain all 4 of the principal components for the iris data set, we will throw away the last two as they capture less than 15% of the variance in the data.

```python
D = np.zeros([X.shape[0],X.shape[1]])
np.fill_diagonal(D,S)
Xrotated = np.dot(U,D)

# First two principal components!
PCs = Xrotated[:,0:2]
PCs.shape
```

Now that we have projected our data set into a low-dimensional space where we can better visualize it, let's make a plot to see how it looks. We will be careful to color the points by the associated class label.

```python
plt.scatter(PCs[:,0],PCs[:,1],
            color=[['red','green','blue'][i] for i in Y.astype(int)])
plt.xlabel("PC1")
plt.ylabel("PC2")
plt.show()
```

You can see in the above plot that the data forms three fairly distinct groups of points. The red group on the left is easily linearly separable from the other two groups. Even the green and blue groups are fairly distinct. While you can almost draw a straight line between them, it appears that a few data points from each of these groups would lie on the opposite side, and not allow for perfect classification with a linear network. Nevertheless, we can now see why a linear network might work well on this data, and (perhaps more importantly) that an additional feature measurement may be needed to completely separate the green and blue species.

We can perform the same analysis using SciKitLearn:

```python
pca = PCA(2)
PCs = pca.fit_transform(X)[:,0:2]
```

```python
plt.scatter(PCs[:,0],PCs[:,1],
            color=[['red','green','blue'][i] for i in Y.astype(int)])
plt.xlabel("PC1")
plt.ylabel("PC2")
plt.show()
```

The result is almost the same, but you will notice that this data has been "flipped" along the origin of the second principal component. This is common since there is an equally valid rotation of the data 180 degrees along any axis that still preserves all of the variance and internal relationships between the data points. Either way, the same amount of variance is accounted for on each component, and decision boundaries can still be explored.

Let's perform a similar analysis with a much larger data set.

## Exploring MNIST with PCA

```python
# Load the MNIST data set using Keras
# the data, shuffled and split between train and test sets
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()

# Subsample (there's a lot of data here!)
X = x_train[range(0,x_train.shape[0],10),:,:]
Y = y_train[range(0,y_train.shape[0],10)]

display(X.shape)
display(Y.shape)
```

We will only look at the training data for this data set for now. The training set consists of 60,000 images each 28x28 pixels in size. However, we have selected just 6000 of those images for this example to make the analysis more tractable (60,000 would take a long time and lots of memory to compute the PCs). Each pixel is represented by an integer intensity value between 0 and 255. For the sake of examining what these images look like, let's scale those intensities to be floating point values in the range [0,1]:

```python
X = X.astype('float32') / 255.0
```

```python
# Plot some of the images
for i in range(5):
    plt.figure()
    plt.imshow(X[i,:,:])
plt.show()
```

```python
display(Matrix(Y[0:5]))
```

You can see that each of these images corresponds to a hand-written digit, each labeled with the appropriate number in the category labels, Y.

We will now **flatten** these images so that we can perform principal component analysis, so that each pixel is treated as a single measurement.

```python
X = X.reshape(X.shape[0],X.shape[1]*X.shape[2])
X.shape
```

Each image has now been encoded as a feature vector in a 784-dimensional space. Even though the pixel intensities make sense to us when visualized as a 2D image, newly initialized neural networks only experience the vector space for the first time, and have to learn such relationships from scratch. However, let's see if PCA can provide some insight on the difficulty of this task.

We will apply the same approach as before:
1. mean-centering the features
2. calculating the SVD
3. examining the singular values
4. scaling the left-singular vectors
5. plotting the two-dimensional projection

NOTE: It may take a minute or two to compute the SVD for a data set of this size, so be patient on the steps below.

```python
# Mean-centering
Xcentered = np.apply_along_axis(mean_center,0,X)

# SVD
U,S,V = np.linalg.svd(Xcentered,full_matrices=True)

# Percent variance accounted for
plt.plot(100.0*S/np.sum(S))
plt.ylabel('% Var')
plt.xlabel('Singular Value')
plt.show()
```

```python
# Variance accounted for in the first two principal components
100.0*(S[0]+S[1])/np.sum(S)
```

You can see that the variance accounted for dips sharply (which is good, because have a few PCs which capture a lot of variance is a useful thing). However, the variance captured by the first two components in total is less then 5% of the total variance in the data set, so PCA might not be so useful for visualization.

However, let's take a quick look at one more thing before moving on:

```python
# Variance accounted for in the first two principal components
display(100.0*(np.sum(S[0:340]))/np.sum(S))

# Reduction?
display(100*340/len(S))
```

Notice that 90% of the **total** variance in the data can be captured using 340 principal components. 340 is only just over 43% of the original dimensionality, so 90% of the data set can be effectively represented using a space less than half of the original in size. Thus, PCA can also be thought of as a linear data compression technique as well. While we can't visualize a 340-dimensional space, a much smaller network would be required to process this data, possibly without sacrificing generalization accuracy (but that's for a later time). Let's just take a look at the first two principal components for now:

```python
D = np.zeros([X.shape[0],X.shape[1]])
np.fill_diagonal(D,S)
Xrotated = np.dot(U,D)

# First two principal components!
PCs = Xrotated[:,0:2]
PCs.shape
```

```python
# Need a lot of colors for this one!
plt.scatter(PCs[:,0],PCs[:,1],
            color=[['red','green','blue','cyan','magenta','yellow','black','brown','grey','purple'][i] for i in Y.astype(int)])
plt.xlabel("PC1")
plt.ylabel("PC2")
plt.show()
```

Note that the clutter of 6000 points is pretty bad in this space, and reducing the sampling may help in some ways. Let's try that now:

```python
plt.scatter(PCs[range(0,6000,10),0],PCs[range(0,6000,10),1],
            color=[['red','green','blue','cyan','magenta','yellow','black','brown','grey','purple'][i] for i in Y[range(0,6000,10)].astype(int)])
plt.xlabel("PC1")
plt.ylabel("PC2")
plt.show()
```

Even with the subsampling, some interesting divisions between classes can be found, but others are not so clear. For example, the red points now correspond to images of zeros and green dots now correspond to images of ones, and these seem to lie on opposite ends of the first principal component. Also, the brown sevens seem very distinct from the cyan threes. However, others are more cluttered, like the purple nines, brown sevens, and magenta fours. However, this makes some sense because these numbers share some common featural similarities.

Because PCA is a linear technique, it is somewhat limited in its ability to capture some of the non-linear relationships between the data vectors. However, two dimensions, while useful for visualization, is still often too low to capture all of the relevant relationships that allow for categorization. The 2D representation can be thought of as capturing the lower-bound on the distances between the data vectors, since adding more dimensions may cause points to move farther apart, but they could never cause them to move closer together. Even with its drawbacks, PCA is a useful technique for quickly creating projections of high-dimensional data onto lower-dimensional spaces.

```python

```


## 4.2 The Law of Large Numbers
Let $Z_i$ be $N$ independent samples from some probability distribution. According to the Law of Large numbers, so long as the expected value $E[Z]$ is finite, the following holds,

$$\frac{1}{N} \sum_{i=1}^N Z_i \rightarrow E[ Z ],  \;\;\; N \rightarrow \infty.$$

This may seem like a boring result, but it will be the most useful tool you use.

### 4.2.1 Intuition
If the above Law is somewhat surprising, it can be made more clear by examining a simple example.

Consider a random variable $Z$ that can take only two values, $c_1$ and $c_2$. Suppose we have a large number of samples of $Z$, denoting a specific sample $Z_i$. The Law says that we can approximate the expected value of $Z$ by averaging over all samples. Consider the average:

$$ \frac{1}{N} \sum_{i=1}^N \;Z_i $$
By construction, $Z_i$ can only take on $c_1$ or $c_2$, hence we can partition the sum over these two values:

$$\begin{align}
\frac{1}{N} \sum_{i=1}^N \;Z_i
=\frac{1}{N} \big(  \sum_{ Z_i = c_1}c_1 + \sum_{Z_i=c_2}c_2 \big) \\\\[5pt]
= c_1 \sum_{ Z_i = c_1}\frac{1}{N} + c_2 \sum_{ Z_i = c_2}\frac{1}{N} \\\\[5pt]
= c_1 \times \text{ (approximate frequency of $c_1$) } \\\\
\;\;\;\;\;\;\;\;\; + c_2 \times \text{ (approximate frequency of $c_2$) } \\\\[5pt]
\approx c_1 \times P(Z = c_1) + c_2 \times P(Z = c_2 ) \\\\[5pt]
= E[Z]
\end{align}$$

Equality holds in the limit, but we can get closer and closer by using more and more samples in the average. This Law holds for almost any distribution, minus some important cases we will encounter later.

### 4.2.2 Example
Below is a diagram of the Law of Large numbers in action for three different sequences of Poisson random variables.

We sample sample_size = 100000 Poisson random variables with parameter $\lambda = 4.5$. (Recall the expected value of a Poisson random variable is equal to its parameter.) We calculate the average for the first $n$ samples, for $n=1$ to sample_size.

```python
%matplotlib inline
import numpy as np
from IPython.core.pylabtools import figsize
import matplotlib.pyplot as plt

figsize( 12.5, 5 )

sample_size = 100000
expected_value = lambda_ = 4.5
poi = np.random.poisson
N_samples = range(1,sample_size,100)

for k in range(3):

    samples = poi( lambda_, sample_size )

    partial_average = [ samples[:i].mean() for i in N_samples ]

    plt.plot( N_samples, partial_average, lw=1.5,label="average \
of  $n$ samples; seq. %d"%k)

plt.plot( N_samples, expected_value*np.ones_like( partial_average),
    ls = "--", label = "true expected value", c = "k" )

plt.ylim( 4.35, 4.65)
plt.title( "Convergence of the average of \n random variables to its \
expected value" )
plt.ylabel( "average of $n$ samples" )
plt.xlabel( "# of samples, $n$")
plt.legend();
```

Looking at the above plot, it is clear that when the sample size is small, there is greater variation in the average (compare how jagged and jumpy the average is initially, then smooths out). All three paths approach the value 4.5, but just flirt with it as $N$ gets large. Mathematicians and statistician have another name for flirting: convergence.

Another very relevant question we can ask is how quickly am I converging to the expected value? Let's plot something new. For a specific $N$, let's do the above trials thousands of times and compute how far away we are from the true expected value, on average. But wait — compute on average? This is simply the law of large numbers again! For example, we are interested in, for a specific $N$, the quantity:

$$D(N) = \sqrt{ \;E\left[\;\; \left( \frac{1}{N}\sum_{i=1}^NZ_i  - 4.5 \;\right)^2 \;\;\right] \;\;}$$
The above formulae is interpretable as a distance away from the true value (on average), for some $N$. (We take the square root so the dimensions of the above quantity and our random variables are the same). As the above is an expected value, it can be approximated using the law of large numbers: instead of averaging $Z_i$, we calculate the following multiple times and average them:

$$ Y_k = \left( \;\frac{1}{N}\sum_{i=1}^NZ_i  - 4.5 \; \right)^2 $$
By computing the above many, $N_y$, times (remember, it is random), and averaging them:

$$ \frac{1}{N_Y} \sum_{k=1}^{N_Y} Y_k \rightarrow E[ Y_k ] = E\;\left[\;\; \left( \frac{1}{N}\sum_{i=1}^NZ_i  - 4.5 \;\right)^2 \right]$$
Finally, taking the square root:

$$ \sqrt{\frac{1}{N_Y} \sum_{k=1}^{N_Y} Y_k} \approx D(N) $$

```python
figsize( 12.5, 4)

N_Y = 250 #use this many to approximate D(N)
N_array = np.arange( 1000, 50000, 2500 ) #use this many samples in the approx. to the variance.
D_N_results = np.zeros( len( N_array ) )

lambda_ = 4.5
expected_value = lambda_ #for X ~ Poi(lambda) , E[ X ] = lambda

def D_N( n ):
    """
    This function approx. D_n, the average variance of using n samples.
    """
    Z = poi( lambda_, (n, N_Y) )
    average_Z = Z.mean(axis=0)
    return np.sqrt( (  (average_Z - expected_value)**2  ).mean() )

for i,n in enumerate(N_array):
    D_N_results[i] =  D_N(n)

plt.xlabel( "$N$" )
plt.ylabel( "expected squared-distance from true value" )
plt.plot(N_array, D_N_results, lw = 3,
            label="expected distance between\n\
expected value and \naverage of $N$ random variables.")
plt.plot( N_array, np.sqrt(expected_value)/np.sqrt(N_array), lw = 2, ls = "--",
        label = r"$\frac{\sqrt{\lambda}}{\sqrt{N}}$" )
plt.legend()
plt.title( "How 'fast' is the sample average converging? " );
```

As expected, the expected distance between our sample average and the actual expected value shrinks as $N$ grows large. But also notice that the rate of convergence decreases, that is, we need only 10 000 additional samples to move from 0.020 to 0.015, a difference of 0.005, but 20 000 more samples to again decrease from 0.015 to 0.010, again only a 0.005 decrease.

It turns out we can measure this rate of convergence. Above I have plotted a second line, the function $\sqrt{\lambda}/\sqrt{N}$. This was not chosen arbitrarily. In most cases, given a sequence of random variable distributed like $Z$, the rate of convergence to $E[Z]$ of the Law of Large Numbers is

$$ \frac{ \sqrt{ \; Var(Z) \; } }{\sqrt{N} }$$
This is useful to know: for a given large $N$, we know (on average) how far away we are from the estimate. On the other hand, in a Bayesian setting, this can seem like a useless result: Bayesian analysis is OK with uncertainty so what's the statistical point of adding extra precise digits? Though drawing samples can be so computationally cheap that having a larger $N$ is fine too.

### 4.2.3 How do we compute $Var(Z)$ though?
The variance is simply another expected value that can be approximated! Consider the following, once we have the expected value (by using the Law of Large Numbers to estimate it, denote it $\mu$), we can estimate the variance:

$$ \frac{1}{N}\sum_{i=1}^N \;(Z_i - \mu)^2 \rightarrow E[ \;( Z - \mu)^2 \;] = Var( Z )$$

### 4.2.4 Expected values and probabilities
There is an even less explicit relationship between expected value and estimating probabilities. Define the indicator function

$$\mathbb{1}_A(x) =
\begin{cases} 1  x \in A \\\\
              0  else
\end{cases}
$$
Then, by the law of large numbers, if we have many samples $X_i$, we can estimate the probability of an event $A$, denoted $P(A)$, by:

$$ \frac{1}{N} \sum_{i=1}^N \mathbb{1}_A(X_i) \rightarrow E[\mathbb{1}_A(X)] =  P(A) $$
Again, this is fairly obvious after a moments thought: the indicator function is only 1 if the event occurs, so we are summing only the times the event occurs and dividing by the total number of trials (consider how we usually approximate probabilities using frequencies). For example, suppose we wish to estimate the probability that a $Z \sim Exp(.5)$ is greater than 5, and we have many samples from a $Exp(.5)$ distribution.

$$ P( Z > 5 ) =  \frac{1}{N}\sum_{i=1}^N \mathbb{1}_{z > 5 }(Z_i) $$

```python
N = 10000
print( np.mean( [ np.random.exponential( 0.5 ) > 5 for i in range(N) ] ) )
```

    0.0

### 4.2.5 What does this all have to do with Bayesian statistics?
Point estimates, to be introduced in the next chapter, in Bayesian inference are computed using expected values. In more analytical Bayesian inference, we would have been required to evaluate complicated expected values represented as multi-dimensional integrals. No longer. If we can sample from the posterior distribution directly, we simply need to evaluate averages. Much easier. If accuracy is a priority, plots like the ones above show how fast you are converging. And if further accuracy is desired, just take more samples from the posterior.

When is enough enough? When can you stop drawing samples from the posterior? That is the practitioners decision, and also dependent on the variance of the samples (recall from above a high variance means the average will converge slower).

We also should understand when the Law of Large Numbers fails. As the name implies, and comparing the graphs above for small $N$, the Law is only true for large sample sizes. Without this, the asymptotic result is not reliable. Knowing in what situations the Law fails can give us confidence in how unconfident we should be. The next section deals with this issue.

## 4.3 The Disorder of Small Numbers
The Law of Large Numbers is only valid as $N$ gets infinitely large: never truly attainable. While the law is a powerful tool, it is foolhardy to apply it liberally. Our next example illustrates this.

### 4.3.1 Example: Aggregated geographic data
Often data comes in aggregated form. For instance, data may be grouped by state, county, or city level. Of course, the population numbers vary per geographic area. If the data is an average of some characteristic of each the geographic areas, we must be conscious of the Law of Large Numbers and how it can fail for areas with small populations.

We will observe this on a toy dataset. Suppose there are five thousand counties in our dataset. Furthermore, population number in each state are uniformly distributed between 100 and 1500. The way the population numbers are generated is irrelevant to the discussion, so we do not justify this. We are interested in measuring the average height of individuals per county. Unbeknownst to us, height does not vary across county, and each individual, regardless of the county he or she is currently living in, has the same distribution of what their height may be:

$$ \text{height} \sim \text{Normal}(150, 15 ) $$
We aggregate the individuals at the county level, so we only have data for the average in the county. What might our dataset look like?

```python
figsize( 12.5, 4)
std_height = 15
mean_height = 150

n_counties = 5000
pop_generator = np.random.randint
norm = np.random.normal

#generate some artificial population numbers
population = pop_generator(100, 1500, n_counties )

average_across_county = np.zeros( n_counties )
for i in range( n_counties ):
    #generate some individuals and take the mean
    average_across_county[i] = norm(mean_height, 1./std_height,
                                        population[i] ).mean()

#located the counties with the apparently most extreme average heights.
i_min = np.argmin( average_across_county )
i_max = np.argmax( average_across_county )

#plot population size vs. recorded average
plt.scatter( population, average_across_county, alpha = 0.5, c="#7A68A6")
plt.scatter( [ population[i_min], population[i_max] ],
           [average_across_county[i_min], average_across_county[i_max] ],
           s = 60, marker = "o", facecolors = "none",
           edgecolors = "#A60628", linewidths = 1.5,
            label="extreme heights")

plt.xlim( 100, 1500 )
plt.title( "Average height vs. County Population")
plt.xlabel("County Population")
plt.ylabel("Average height in county")
plt.plot( [100, 1500], [150, 150], color = "k", label = "true expected \
height", ls="--" )
plt.legend(scatterpoints = 1);
```

What do we observe? Without accounting for population sizes we run the risk of making an enormous inference error: if we ignored population size, we would say that the county with the shortest and tallest individuals have been correctly circled. But this inference is wrong for the following reason. These two counties do not necessarily have the most extreme heights. The error results from the calculated average of smaller populations not being a good reflection of the true expected value of the population (which in truth should be $\mu =150$). The sample size/population size/$N$, whatever you wish to call it, is simply too small to invoke the Law of Large Numbers effectively.

We provide more damning evidence against this inference. Recall the population numbers were uniformly distributed over 100 to 1500. Our intuition should tell us that the counties with the most extreme population heights should also be uniformly spread over 100 to 1500, and certainly independent of the county's population. Not so. Below are the population sizes of the counties with the most extreme heights.

```python
print("Population sizes of 10 'shortest' counties: ")
print(population[ np.argsort( average_across_county )[:10] ], '\n')
print("Population sizes of 10 'tallest' counties: ")
print(population[ np.argsort( -average_across_county )[:10] ])
```

    Population sizes of 10 'shortest' counties:
    [105 141 119 155 124 139 103 102 217 185]

    Population sizes of 10 'tallest' counties:
    [109 149 216 305 208 131 100 385 293 139]

### 4.3.2 Example: Kaggle's U.S. Census Return Rate Challenge
Below is data from the 2010 US census, which partitions populations beyond counties to the level of block groups (which are aggregates of city blocks or equivalents). The dataset is from a Kaggle machine learning competition some colleagues and I participated in. The objective was to predict the census letter mail-back rate of a group block, measured between 0 and 100, using census variables (median income, number of females in the block-group, number of trailer parks, average number of children etc.). Below we plot the census mail-back rate versus block group population:

```python
figsize( 12.5, 6.5 )
data = np.genfromtxt( "./data/census_data.csv", skip_header=1,
                        delimiter= ",")
plt.scatter( data[:,1], data[:,0], alpha = 0.5, c="#7A68A6")
plt.title("Census mail-back rate vs Population")
plt.ylabel("Mail-back rate")
plt.xlabel("population of block-group")
plt.xlim(-100, 15e3 )
plt.ylim( -5, 105)

i_min = np.argmin(  data[:,0] )
i_max = np.argmax(  data[:,0] )

plt.scatter( [ data[i_min,1], data[i_max, 1] ],
             [ data[i_min,0],  data[i_max,0] ],
             s = 60, marker = "o", facecolors = "none",
             edgecolors = "#A60628", linewidths = 1.5,
             label="most extreme points")

plt.legend(scatterpoints = 1);
```

The above is a classic phenomenon in statistics. I say classic referring to the "shape" of the scatter plot above. It follows a classic triangular form, that tightens as we increase the sample size (as the Law of Large Numbers becomes more exact).

I am perhaps overstressing the point and maybe I should have titled the book "You don't have big data problems!", but here again is an example of the trouble with small datasets, not big ones. Simply, small datasets cannot be processed using the Law of Large Numbers. Compare with applying the Law without hassle to big datasets (ex. big data). I mentioned earlier that paradoxically big data prediction problems are solved by relatively simple algorithms. The paradox is partially resolved by understanding that the Law of Large Numbers creates solutions that are stable, i.e. adding or subtracting a few data points will not affect the solution much. On the other hand, adding or removing data points to a small dataset can create very different results.

## 4.4 Conclusion
While the Law of Large Numbers is cool, it is only true so much as its name implies: with large sample sizes only. We have seen how our inference can be affected by not considering how the data is shaped.

1. By (cheaply) drawing many samples from the posterior distributions, we can ensure that the Law of Large Number applies as we approximate expected values (which we will do in the next chapter).

2. Bayesian inference understands that with small sample sizes, we can observe wild randomness. Our posterior distribution will reflect this by being more spread rather than tightly concentrated. Thus, our inference should be correctable.

3. There are major implications of not considering the sample size, and trying to sort objects that are unstable leads to pathological orderings. The method provided above solves this problem.