# 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
tfpj = 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 = tfpj.distributions.Blockwise(
    [tfpj.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 = 84

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=1000, best=True)
```

    CPU times: user 15min 35s, sys: 779 ms, total: 15min 35s
    Wall time: 15min 32s

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

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

    DeviceArray(0.91142211, dtype=float64)

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

    DeviceArray(967560.53213306, dtype=float64)

## Inference with Normalising Flows

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.

We will then proceed to doing inference using a Conditional Autoregressive Flow.

```python
import tensorflow as tf
import tensorflow_probability as tfp
tfd = tfp.distributions
tfb = tfp.bijectors
```

```python
import tensorflow_probability.substrates.jax as tfpj
tfdj = tfpj.distributions
tfbj = tfpj.bijectors

import tensorflow.keras.backend as K
```

```python
from tqdm import trange
import numpy as onp

def affine_sample(log_prob, n_params, n_walkers, n_steps, walkers1, walkers2):

    # initialize current state
    current_state1 = tf.Variable(walkers1)
    current_state2 = tf.Variable(walkers2)

    # initial target log prob for the walkers (and set any nans to -inf)...
    logp_current1 = log_prob(current_state1)
    logp_current2 = log_prob(current_state2)
    logp_current1 = tf.where(tf.math.is_nan(logp_current1), tf.ones_like(logp_current1)*tf.math.log(0.), logp_current1)
    logp_current2 = tf.where(tf.math.is_nan(logp_current2), tf.ones_like(logp_current2)*tf.math.log(0.), logp_current2)

    # holder for the whole chain
    chain = [tf.concat([current_state1, current_state2], axis=0)]

    # MCMC loop
    with trange(1, n_steps) as t:
        for epoch in t:

            # first set of walkers:

            # proposals
            partners1 = tf.gather(current_state2, onp.random.randint(0, n_walkers, n_walkers))
            z1 = 0.5*(tf.random.uniform([n_walkers], minval=0, maxval=1)+1)**2
            proposed_state1 = partners1 + tf.transpose(z1*tf.transpose(current_state1 - partners1))

            # target log prob at proposed points
            logp_proposed1 = log_prob(proposed_state1)
            logp_proposed1 = tf.where(tf.math.is_nan(logp_proposed1), tf.ones_like(logp_proposed1)*tf.math.log(0.), logp_proposed1)

            # acceptance probability
            p_accept1 = tf.math.minimum(tf.ones(n_walkers), z1**(n_params-1)*tf.exp(logp_proposed1 - logp_current1) )

            # accept or not
            accept1_ = (tf.random.uniform([n_walkers], minval=0, maxval=1) <= p_accept1)
            accept1 = tf.cast(accept1_, tf.float32)

            # update the state
            current_state1 = tf.transpose( tf.transpose(current_state1)*(1-accept1) + tf.transpose(proposed_state1)*accept1)
            logp_current1 = tf.where(accept1_, logp_proposed1, logp_current1)

            # second set of walkers:

            # proposals
            partners2 = tf.gather(current_state1, onp.random.randint(0, n_walkers, n_walkers))
            z2 = 0.5*(tf.random.uniform([n_walkers], minval=0, maxval=1)+1)**2
            proposed_state2 = partners2 + tf.transpose(z2*tf.transpose(current_state2 - partners2))

            # target log prob at proposed points
            logp_proposed2 = log_prob(proposed_state2)
            logp_proposed2 = tf.where(tf.math.is_nan(logp_proposed2), tf.ones_like(logp_proposed2)*tf.math.log(0.), logp_proposed2)

            # acceptance probability
            p_accept2 = tf.math.minimum(tf.ones(n_walkers), z2**(n_params-1)*tf.exp(logp_proposed2 - logp_current2) )

            # accept or not
            accept2_ = (tf.random.uniform([n_walkers], minval=0, maxval=1) <= p_accept2)
            accept2 = tf.cast(accept2_, tf.float32)

            # update the state
            current_state2 = tf.transpose( tf.transpose(current_state2)*(1-accept2) + tf.transpose(proposed_state2)*accept2)
            logp_current2 = tf.where(accept2_, logp_proposed2, logp_current2)

            # append to chain
            chain.append(tf.concat([current_state1, current_state2], axis=0))

    # stack up the chain
    chain = tf.stack(chain, axis=0)

    return chain

```

```python
class ConditionalMaskedAutoregressiveFlow(tf.Module):
    def __init__(self, n_dimensions=None, n_conditionals=None, n_mades=1, n_hidden=[50,50], input_order="random",
    activation=tf.keras.layers.LeakyReLU(0.01),
    all_layers=True,
    kernel_initializer=tf.keras.initializers.RandomNormal(mean=0.0, stddev=1e-5, seed=None),
    bias_initializer=tf.keras.initializers.RandomNormal(mean=0.0, stddev=1e-5, seed=None),
    kernel_regularizer=None, bias_regularizer=None, kernel_constraint=None,
    bias_constraint=None):
        super(ConditionalMaskedAutoregressiveFlow, self).__init__()
        # extract init parameters
        self.n_dimensions = n_dimensions
        self.n_conditionals = n_conditionals
        self.n_mades = n_mades
        # construct the base (normal) distribution
        self.base_distribution = tfd.MultivariateNormalDiag(loc=tf.zeros(self.n_dimensions), scale_diag=tf.ones(self.n_dimensions))
        # put the conditional inputs to all layers, or just the first layer?
        if all_layers == True:
            all_layers = "all_layers"
        else:
            all_layers = "first_layer"
        # construct stack of conditional MADEs
        self.MADEs = [tfb.AutoregressiveNetwork(
                        params=2,
                        hidden_units=n_hidden,
                        activation=activation,
                        event_shape=[n_dimensions],
                        conditional=True,
                        conditional_event_shape=[n_conditionals],
                        conditional_input_layers=all_layers,
                        input_order=input_order,
                        kernel_initializer=kernel_initializer,
                        bias_initializer=bias_initializer,
                        kernel_regularizer=kernel_regularizer,
                        bias_regularizer=bias_regularizer,
                        kernel_constraint=kernel_constraint,
                        bias_constraint=bias_constraint,
                        name="MADE_{}".format(i)) for i in range(n_mades)]
    # bijector for x | y (chain the conditional MADEs together)
    def bijector(self, y):
        # start with an empty bijector
        MAF = tfb.Identity()
        # pass through the MADE layers (passing conditional inputs each time)
        for i in range(self.n_mades):
            MAF = tfb.MaskedAutoregressiveFlow(shift_and_log_scale_fn=lambda x: self.MADEs[i](x, conditional_input=y))(MAF)
        return MAF
    # construct distribution P(x | y)
    def __call__(self, y):
        return tfd.TransformedDistribution(
            self.base_distribution,
            bijector=self.bijector(y))
    # log probability ln P(x | y)
    def log_prob(self, x, y):
        return self.__call__(y).log_prob(x)
    # sample n samples from P(x | y)
    def sample(self, n, y):
        # base samples
        base_samples = self.base_distribution.sample(n)
        # biject the samples
        return self.bijector(y).forward(base_samples)
```

```python
num_models = 2
cmaf_models = [ConditionalMaskedAutoregressiveFlow(n_dimensions=2,
                                                   n_conditionals=2, n_hidden=[50,50,50]) for i in range(num_models)]

cmaf_models += [ConditionalMaskedAutoregressiveFlow(n_dimensions=2,
                                                    n_conditionals=2, n_hidden=[50,50]) for i in range(2)]
```

```python
maf_optimizers = [tf.keras.optimizers.Adam(learning_rate=1e-3) for i in range(len(cmaf_models))]
```

```python
# prior for sim draws in jax
prior = tfpj.distributions.Blockwise(
    [tfpj.distributions.Uniform(low=low, high=high)
     for low, high in zip([0.4, 0.4], [1.0, 1.1])])
prior.low = np.array([0.4, 0.4])
prior.high = np.array([1.0, 1.1])

# set up prior in non-jax tfp
samp_prior = tfp.distributions.Blockwise(
    [tfp.distributions.Uniform(low=low, high=high)
     for low, high in zip([0.4, 0.4], [1.0, 1.1])])
prior.low = np.array([0.4, 0.4])
prior.high = np.array([1.0, 1.1])

@tf.function
def train_step(x, y):
    _ls = []
    for m in range(len(cmaf_models)):
        with tf.GradientTape() as tape:
            _l = K.mean(-cmaf_models[m].log_prob(x, y) - samp_prior.log_prob(y))
            _ls.append(_l)

        grads = tape.gradient(_l, cmaf_models[m].trainable_variables)
        maf_optimizers[m].apply_gradients(zip(grads, cmaf_models[m].trainable_variables))
        #train_acc_metric.update_state(y, logits)
    return _ls

@tf.function
def val_step(x, y):
    _val_l = []
    for m in range(len(cmaf_models)):
        _val_l.append(K.mean(-cmaf_models[m].log_prob(x, y) - samp_prior.log_prob(y)))
    return _val_l

@tf.function
def loss(x, y):
    return K.mean(-cmaf_model.log_prob(x, y))
```

```python
from sklearn.preprocessing import StandardScaler
summary_scaler = StandardScaler()
```

```python
# create dataset for this demo

num_sims = 1000

keys = jax.random.split(rng, num=num_sims)

def get_params_summaries(key, n_samples, ϑ_samp):
    keys = np.array(jax.random.split(key, num=n_samples))
    sim = lambda rng, θ: simulator(
            rng, θ, simulator_args={
                **simulator_args,
                **{"squeeze": False}})
    d = jax.vmap(sim)(keys, ϑ_samp)

    t = IMNN.get_estimate(d)

    t1 = t.flatten()[~np.isnan(t.flatten())]
    ϑ_samp = ϑ_samp.flatten()[~np.isnan(t.flatten())]

    # scale all summaries
    t1 = t1.reshape((len(t1)//2, 2))

    #_tstd = np.std(t1, axis=0)
    #_tmu = np.mean(t1, axis=0)

    #t1 = (t1 - _tmu) / _tstd

    # return x,y; e.g. t and ϑ (conditional)
    # learn p(x | y) = p(t | ϑ)
    return t1, ϑ_samp.reshape(len(ϑ_samp)//2, 2) #, _tmu, _tstd

def get_dataset(data, batch_size=20, buffer_size=1000, split=0.75):
    x,y = data

    # shufle data right off the bat
#     shuffle_idx = onp.random.choice(np.arange(len(x)), size=len(x))
#     x = x[shuffle_idx]
#     y = y[shuffle_idx]

    idx = int(len(x)*split)
    x_train = x[:idx]
    y_train = y[:idx]
    x_val = x[idx:]
    y_val = y[idx:]

    # Prepare the training dataset.
    train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
    train_dataset = train_dataset.shuffle(buffer_size=buffer_size).batch(batch_size)

    # Prepare the validation dataset.
    val_dataset = tf.data.Dataset.from_tensor_slices((x_val, y_val))
    val_dataset = val_dataset.batch(batch_size)

    return train_dataset, val_dataset

def proposal_distribution(prior, MAF):
    """return geometric mean of proposal distribution"""

```

```python
n_samples = 1000
batch_size = 100
buffer_size = n_samples
key1,key2 = jax.random.split(rng)

ϑ_samp = prior.sample(sample_shape=(n_samples,), seed=key1)
ts, ϑ_samp = get_params_summaries(key2, n_samples, ϑ_samp)
data = (ts, ϑ_samp)
```

```python
train_dataset, val_dataset = get_dataset(data, batch_size=batch_size, buffer_size=buffer_size)
```

```python
import time
from tqdm import tqdm

epochs = 2000

train_losses = []
val_losses = []

for epoch in tqdm(range(epochs)):
    #print("\nStart of epoch %d" % (epoch,))
    start_time = time.time()

#     ϑ_samp = prior.sample(sample_shape=(n_samples,), seed=key)
#     data = get_params_summaries(key, n_samples, ϑ_samp)
#    train_dataset, val_dataset = get_dataset(data, batch_size=100)

    # shuffle training data anew every 100th epoch
    if epoch % 100 == 0:
        train_dataset = train_dataset.shuffle(buffer_size=buffer_size)

    # Iterate over the batches of the dataset.
    for step, (x_batch_train, y_batch_train) in enumerate(train_dataset):
        loss_values = np.array(train_step(x_batch_train, y_batch_train))

    train_losses.append(loss_values)

    # Run a validation loop at the end of each epoch.
    for x_batch_val, y_batch_val in val_dataset:
        val_loss = np.array(val_step(x_batch_val, y_batch_val))

    val_losses.append(val_loss)
```

    100%|██████████| 2000/2000 [01:02<00:00, 32.16it/s]

```python
import seaborn as sns
%matplotlib inline

sns.set_theme()
plt.figure(figsize=(8,3.5))
plt.subplot(131)
plt.plot(np.array(train_losses).T[0], label='train')
plt.plot(np.array(val_losses).T[0], label='val')
plt.ylabel(r'$p(t\ |\ \vartheta, w)$')

plt.subplot(132)
plt.plot(np.array(train_losses).T[1], label='train')
plt.plot(np.array(val_losses).T[1], label='val')
plt.ylabel(r'$p(t\ |\ \vartheta, w)$')

plt.subplot(133)
plt.plot(np.array(train_losses).T[4], label='train')
plt.plot(np.array(val_losses).T[4], label='val')
plt.ylabel(r'$p(t\ |\ \vartheta, w)$')
plt.legend()

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

```python
# get summaries for observed data from IMNN
# δ_target is the 128x128 simulation
estimates = IMNN.get_estimate(np.expand_dims(δ_target, (0, 1, 2)))

#_estimates = (estimates - tmu) / tstd

# set up prior in non-jax tfp
samp_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])

# define cmaf model log prob with prior defined above
@tf.function
def my_log_prob(y, x=estimates):

    # here all cmafs are trained to draw from p(x|y) <=> p(t|ϑ)
    # take the mean of the data likelihood p(x|y) over all cmafs
    _like = cmaf_models[0].log_prob(x,y)

    _prior = samp_prior.log_prob(x)

    return _like + _prior

# sample using affine
n_steps = 2000
n_walkers = 500
burnin_steps = 1800
n_params = 2

# initialize walkers...
walkers1 = tf.random.normal([n_walkers, 2], 0.5, 0.5)
walkers2 = tf.random.normal([n_walkers, 2], 0.5, 0.5)

# walkers1 = tf.random.uniform([n_walkers, 2], 0.1, 1.25)
# walkers2 = tf.random.uniform([n_walkers, 2], 0.1, 1.25)

chain = affine_sample(my_log_prob, n_params, n_walkers, n_steps, walkers1, walkers2)

skip = 4

post = np.stack([chain.numpy()[burnin_steps::skip,:,0].flatten(), chain.numpy()[burnin_steps::skip,:,1].flatten()], axis=-1)

# make ABC object for the marginal plot
_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))

margs = _ABC.get_marginals(accepted_parameters = [np.array(post)], smoothing=None, gridsize=400)

MAF = imnn.lfi.LikelihoodFreeInference(
    prior=prior,
    gridsize=100)

MAF.n_targets=2
MAF.put_marginals(
                margs[1][0]);

MAF.marginal_plot(
    known=θ_target,
    label="Delfi estimate",
    axis_labels=["A", "B"]);
```

# next sample from the new posterior and re-train:

```python
n_samples = 1000
key1,key2 = jax.random.split(rng)

idx = np.arange(len(post))
ϑ_samp = post[45000:][onp.random.choice(idx, size=n_samples)]

# ensure not < 0
idx = np.where(ϑ_samp > 0)[0]
ϑ_samp = ϑ_samp[idx]
n_samples = len(ϑ_samp)

ts, ϑ_samp = get_params_summaries(key2, n_samples, ϑ_samp)

new_data = (ts, ϑ_samp)

# this should shuffle the dataset
new_train_dataset, new_val_dataset = get_dataset(new_data, batch_size=batch_size, buffer_size=len(new_data[0]))

# concatenate datasets
train_dataset = train_dataset.concatenate(new_train_dataset)
#val_dataset = val_dataset.concatenate(new_val_dataset)
```

```python
epochs = 1000

new_train_losses = []
new_val_losses = []

for epoch in tqdm(range(epochs)):
    #print("\nStart of epoch %d" % (epoch,))
    start_time = time.time()

    # generate data anew each epoch
#     ϑ_samp = prior.sample(sample_shape=(n_samples,), seed=key)
#     data = get_params_summaries(key, n_samples, ϑ_samp)
#     train_dataset, val_dataset = get_dataset(data, batch_size=100)

    # Iterate over the batches of the dataset.
    for step, (x_batch_train, y_batch_train) in enumerate(train_dataset):
        loss_values = np.array(train_step(x_batch_train, y_batch_train))

    new_train_losses.append(loss_values)

    # Run a validation loop at the end of each epoch.
    for x_batch_val, y_batch_val in val_dataset:
        val_loss = np.array(val_step(x_batch_val, y_batch_val))

    new_val_losses.append(val_loss)
```

    100%|██████████| 1000/1000 [01:01<00:00, 16.16it/s]

```python
plt.figure(figsize=(8,3.5))
plt.subplot(131)
plt.plot(np.array(new_train_losses).T[0], label='train')
plt.plot(np.array(new_val_losses).T[0], label='val')
plt.ylabel(r'$p(t\ |\ \vartheta, w)$')

plt.subplot(132)
plt.plot(np.array(new_train_losses).T[1], label='train')
plt.plot(np.array(new_val_losses).T[1], label='val')
plt.ylabel(r'$p(t\ |\ \vartheta, w)$')

plt.subplot(133)
plt.plot(np.array(new_train_losses).T[4], label='train')
plt.plot(np.array(new_val_losses).T[4], label='val')
plt.ylabel(r'$p(t\ |\ \vartheta, w)$')
plt.legend()

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

```python
# get summaries for observed data from IMNN
# δ_target is the 128x128 simulation

# do for each model
MAFs = []

for m,cmaf_model in enumerate(cmaf_models):

    @tf.function
    def my_log_prob(y, x=estimates):

        # here all cmafs are trained to draw from p(x|y) <=> p(t|ϑ)
        # take the mean of the data likelihood p(x|y) over all cmafs
        _like = cmaf_models[m].log_prob(x,y)
        _prior = samp_prior.log_prob(y)

        return _like + _prior

    # sample using affine
    n_steps = 2000
    n_walkers = 500
    burnin_steps = 1800
    n_params = 2

    # initialize walkers...
    walkers1 = tf.random.normal([n_walkers, 2], 0.5, 0.5)
    walkers2 = tf.random.normal([n_walkers, 2], 0.5, 0.5)

    # walkers1 = tf.random.uniform([n_walkers, 2], 0.1, 1.25)
    # walkers2 = tf.random.uniform([n_walkers, 2], 0.1, 1.25)

    chain = affine_sample(my_log_prob, n_params, n_walkers, n_steps, walkers1, walkers2)

    skip = 4

    post = np.stack([chain.numpy()[burnin_steps::skip,:,0].flatten(), chain.numpy()[burnin_steps::skip,:,1].flatten()], axis=-1)

    margs = ABC.get_marginals(accepted_parameters = [np.array(post)], smoothing=None, gridsize=400)

    _MAF = imnn.lfi.LikelihoodFreeInference(
        prior=prior,
        gridsize=100)

    _MAF.n_targets=2
    _MAF.put_marginals(
                    margs[1][0]);

    MAFs.append(_MAF)

    _MAF.marginal_plot(
        known=θ_target,
        label="Delfi estimate",
        axis_labels=["A", "B"]);
```

```python
post = np.stack([chain.numpy()[burnin_steps:,:,0].flatten(), chain.numpy()[burnin_steps:,:,1].flatten()], axis=-1)
```

```python
post.shape
```

    (600000, 2)

# now get the MAF's posterior estimate

```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=2000,
    smoothing=1, max_iterations=2000);
```

    [2019] accepted in last  87 iterations  (870000 simulations done).

```python
ABC.parameters.accepted[0].shape
```

    (2019, 2)

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

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

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

for m,_MAF in enumerate(MAFs[2:3]):
    if m == 0:
        label = 'Delfi Estimate'
    else:
        label=None

    _MAF.marginal_plot(
        ax=ax,
        label=label,
        colours="C1", linestyle='solid');

plt.savefig('/mnt/home/tmakinen/repositories/field-plots/gauss_delfi_vs_abc.png', dpi=800)
plt.legend()
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

```


# Three-dimensional rigid body kinetics

> Renato Naville Watanabe

> Laboratory of Biomechanics and Motor Control ([http://pesquisa.ufabc.edu.br/bmclab](http://pesquisa.ufabc.edu.br/bmclab))
> Federal University of ABC, Brazil

Table of Contents

1  The Newton-Euler laws

1.1  Resultant force and moments

2  Deduction of the angular momentum in three-dimensional movement

2.1  Angular momentum in three-dimensional movement

3  The matrix of inertia is different depending on the body orientation

3.1  The solution is to attach a frame reference to the body

3.2  The fixed frame is chosen in a way that the matrix of inertia is a diagonal matrix

4  The angular velocity in the fixed frame

4.1  The angular momentum in the fixed frame

5  Derivative of the angular momentum

5.1  Angular velocity

5.2  The derivative of the versors

5.3  The derivative of the angular momentum

6  Newton-Euler laws

7  Examples

7.1  3D pendulum bar

7.2  Data from postural control

8  Problems

9  References

## The Newton-Euler laws

The Newton-Euler laws remain valid in the three-dimensional motion of a rigid body (for revision on Newton-Euler laws, [see this notebook about these laws](newton_euler_equations.ipynb)).

\begin{equation}
    \vec{F} = m\vec{a_{cm}}
\end{equation}

\begin{equation}
    \vec{M_O} = \frac{d\vec{H_O}}{dt}
\end{equation}

### Resultant force and moments

The resultant force and the resultant moment around a point O are computed in the same way of the two-dimensional case.

\begin{equation}
    \vec{F} = \sum\limits_{i=1}^n \vec{F_n}
\end{equation}

\begin{equation}
    \vec{M_O} = \sum\limits_{i=1}^n \vec{r_{i/O}} \times \vec{F_i}
\end{equation}

## Deduction of the angular momentum in three-dimensional movement

The major difference that appears when we are dealing with three-dimensional motions is to compute the derivative of the angular momentum. For the sake of simplicity, the analysis will only consider the point O as the center of mass of the body.

We begin computing the angular momentum of a rigid body around its center of mass.

\begin{equation}
    \vec{H_{cm}} =  \int_B \vec{r_{/cm}} \times \vec{v}\,dm = \int  \vec{r_{/cm}} \times (\vec{\omega}\times\vec{r_{/cm}})\,dm
\end{equation}

where $\vec{\boldsymbol{r_{/cm}}}$ is the vector from the point O to the position of the infinitesimal mass considered in the integral. For simplicity of the notation, we will use $\vec{\boldsymbol{r}} = \vec{\boldsymbol{r_{/cm}}}$. The triple vector product above can be computed using the rule:

    $$\vec{a}\times(\vec{b}\times\vec{c}) = (\vec{a}.\vec{c})\vec{b} - (\vec{a}.\vec{b})\vec{c}$$.

So, the angular momentum is:

\begin{equation}
 \vec{H_{cm}} = \int  \vec{r} \times (\vec{\omega}\times\vec{r})\,dm =
      \int  (r_x\hat{\boldsymbol{i}} + r_y\hat{\boldsymbol{j}}+r_z\hat{\boldsymbol{k}}) \times \left[(\omega_x\hat{\boldsymbol{i}} + \omega_y\hat{\boldsymbol{j}}+\omega_z\hat{\boldsymbol{k}})\times(r_x\hat{\boldsymbol{i}} + r_y\hat{\boldsymbol{j}}+r_z\hat{\boldsymbol{k}})\right]\,dm = \int \left[(r_x\hat{\boldsymbol{i}} + r_y\hat{\boldsymbol{j}}+r_z\hat{\boldsymbol{k}}).(r_x\hat{\boldsymbol{i}} + r_y\hat{\boldsymbol{j}}+r_z\hat{\boldsymbol{k}})\right](\omega_x\hat{\boldsymbol{i}} + \omega_y\hat{\boldsymbol{j}}+\omega_z\hat{\boldsymbol{k}}) -\left[(r_x\hat{\boldsymbol{i}} + r_y\hat{\boldsymbol{j}}+r_z\hat{\boldsymbol{k}}).(\omega_x\hat{\boldsymbol{i}} + \omega_y\hat{\boldsymbol{j}}+\omega_z\hat{\boldsymbol{k}})\right](r_x\hat{\boldsymbol{i}} + r_y\hat{\boldsymbol{j}}+r_z\hat{\boldsymbol{k}}) \,dm = \int (r_x^2 + r_y^2+r_z^2)(\omega_x\hat{\boldsymbol{i}} + \omega_y\hat{\boldsymbol{j}}+\omega_z\hat{\boldsymbol{k}}) - (r_x\omega_x + r_y\omega_y+r_z\omega_z)(r_x\hat{\boldsymbol{i}} + r_y\hat{\boldsymbol{j}}+r_z\hat{\boldsymbol{k}})\,dm
\end{equation}

\begin{equation}
 \vec{H_{cm}} = \int (r_x^2 + r_y^2+r_z^2)\omega_x\hat{\boldsymbol{i}} - (r_x\omega_x + r_y\omega_y+r_z\omega_z)r_x\hat{\boldsymbol{i}} + (r_x^2 + r_y^2+r_z^2)\omega_y\hat{\boldsymbol{j}} - (r_x\omega_x + r_y\omega_y+r_z\omega_z)r_y\hat{\boldsymbol{j}}+(r_x^2 + r_y^2+r_z^2)\omega_z\hat{\boldsymbol{k}} - (r_x\omega_x + r_y\omega_y+r_z\omega_z)r_z\hat{\boldsymbol{k}}\,dm
\end{equation}

\begin{equation}
\vec{H_{cm}}=\left(\int \omega_x(r_y^2+r_z^2)\,dm+  \int-\omega_yr_xr_y\,dm + \int-\omega_z r_xr_z\,dm\right)\;\hat{\boldsymbol{i}} + \left(\int-\omega_x r_xr_y\,dm +\int\omega_y (r_x^2 +r_z^2)\,dm + \int- \omega_zr_yr_z\,dm\right)\hat{\boldsymbol{j}} + \left(\int-\omega_x r_xr_z\,dm  + \int-\omega_yr_yr_z\,dm + \int \omega_z(r_x^2+r_y^2) \,dm\right)\hat{\boldsymbol{k}}
\end{equation}

\begin{equation}
      \vec{H_{cm}}=\bigg(\omega_x\int (r_y^2+r_z^2)\,dm+  \omega_y\int-r_xr_y\,dm + \omega_z\int- r_xr_z\,dm\bigg)\;\hat{\boldsymbol{i}}+\bigg(\omega_x\int- r_xr_y\,dm +\omega_y\int (r_x^2 +r_z^2)\,dm +  \omega_z\int-r_yr_z\,dm\bigg)\hat{\boldsymbol{j}} + \bigg(\omega_x\int- r_xr_z\,dm  + \omega_y\int-r_yr_z\,dm + \omega_z\int (r_x^2+r_y^2) \,dm\bigg)\hat{\boldsymbol{k}}
\end{equation}

\begin{equation}
\vec{H_{cm}}=\left(\omega_xI_{xx}^{cm}+  \omega_yI_{xy}^{cm} + \omega_zI_{xz}^{cm}\right)\;\hat{\boldsymbol{i}} + \left(\omega_xI_{xy}^{cm} +\omega_yI_{yy}^{cm} +  \omega_zI_{yz}^{cm}\right)\hat{\boldsymbol{j}}+\left(\omega_xI_{yz}^{cm}  + \omega_yI_{yz}^{cm} + \omega_zI_{zz}^{cm}\right)\hat{\boldsymbol{k}}
\end{equation}

### Angular momentum in three-dimensional movement

\begin{align}
\begin{split}
      \vec{H_{cm}}=\left[\begin{array}{ccc}I_{xx}^{cm}&I_{xy}^{cm}&I_{xz}^{cm}\\
                                  I_{xy}^{cm}&I_{yy}^{cm}&I_{yz}^{cm}\\
                                  I_{xz}^{cm}&I_{yz}^{cm}&I_{zz}^{cm}\end{array}\right]\cdot
               \left[\begin{array}{c}\omega_x\\\omega_y\\\omega_z \end{array}\right] = I\vec{\omega}
      \end{split}
      \label{eq:angmom}
 \end{align}

 where this matrix is the Matrix of Inertia (or more rigorously, Tensor of Inertia) as defined previously in [this notebook about moment of inertia](CenterOfMassAndMomentOfInertia.ipynb#matrixinertia).

## The matrix of inertia is different depending on the body orientation

So, to compute the angular momentum of a body we have to multiply the matrix of inertia $I$ of the body by its angular velocity $\vec{\omega}$. The problem on this approach is that the moments and products of inertia depends on the orientation of the body relative to the frame of reference. As they depend on the distances of each point of the body to the axes, if the body is rotating, the matrix of inertia $I$ will be different at each instant.

### The solution is to attach a frame reference to the body

The solution to this problem is to attach a frame of reference to the body. We will denote this frame of reference as $\hat{\boldsymbol{e_1}}$, $\hat{\boldsymbol{e_2}}$ and $\hat{\boldsymbol{e_3}}$, with origin in the center of mass of the body. As can be noted from the figure below, the frame of reference moves along the body. Now the matrix of inertia $I$ will be constant relative to this new basis $\hat{\boldsymbol{e_1}}$, $\hat{\boldsymbol{e_2}}$ and $\hat{\boldsymbol{e_3}}$.

### The fixed frame is chosen in a way that the matrix of inertia is a diagonal matrix

As we can choose the basis versors as we want, we can choose them so as to the products of inertia be equal to zero. This can always be done. If the body has axes of symmetry,  we can choose these axes (principal axes) as the direction of the basis to the products of inertia be equal to zero. Having the basis $\hat{\boldsymbol{e_1}}$, $\hat{\boldsymbol{e_2}}$ and $\hat{\boldsymbol{e_3}}$, the matrix of inertia will be a diagonal matrix:

\begin{equation}
    I = \left[\begin{array}{ccc}I_1&0&0\\
          0&I_2&0\\
          0&0&I_3\end{array}\right]
\end{equation}

## The angular velocity in the fixed frame

So, we can write the angular momentum vector relative to this new basis. To do this, we must write the angular velocity in this new basis:

\begin{equation}
    \vec{\omega} = \omega_1\hat{\boldsymbol{e_1}} + \omega_2\hat{\boldsymbol{e_2}} + \omega_3\hat{\boldsymbol{e_3}}
\end{equation}

Note that this angular velocity is the same vector that we used previously. We are just describing it in a basis attached to the body (local basis).

### The angular momentum in the fixed frame

So, using the basis in the direction of the principal axes of the body the angular momentum simplifies to:

\begin{equation}
 \vec{H_{cm}} = I\vec{\omega} = I_1\omega_1 \hat{\boldsymbol{e_1}} + I_2\omega_2 \hat{\boldsymbol{e_2}} +I_3\omega_3 \hat{\boldsymbol{e_3}}
 \label{eq:angmomprinc}
\end{equation}

## Derivative of the angular momentum

For the second Newton-Euler law, we must compute the derivative of the angular momentum. So, we derive the angular momentum in Eq. \eqref{eq:angmomprinc}. As the versors $\hat{\boldsymbol{e_1}}$, $\hat{\boldsymbol{e_2}}$ and $\hat{\boldsymbol{e_3}}$ are varying in time, we must consider their derivatives.

\begin{equation}
    \frac{d\vec{H_{cm}}}{dt} = I_1\dot{\omega_1}\hat{\boldsymbol{e_1}} + I_2\dot{\omega_2}\hat{\boldsymbol{e_2}}+I_3\dot{\omega_3}\hat{\boldsymbol{e_3}} + I_1\omega_1\frac{d\hat{\boldsymbol{e_1}}}{dt} + I_2\omega_2\frac{d\hat{\boldsymbol{e_2}}}{dt}+I_3\omega_3\frac{d\hat{\boldsymbol{e_3}}}{dt}
    \label{eq:derivangmom}
\end{equation}

Now it only remains to find an expression for the derivative of the versors $\frac{d\hat{\boldsymbol{e_1}}}{dt}$, $\frac{d\hat{\boldsymbol{e_2}}}{dt}$ and $\frac{d\hat{\boldsymbol{e_3}}}{dt}$.

### Angular velocity

The angular velocity can be written as ([here you can learn more about the computation of the angular velocity](AngularVelocity3D.ipynb)):

\begin{equation}
    \vec{\omega} =  \left(\frac{d\hat{\boldsymbol{e_2}}}{dt}\cdot \hat{\boldsymbol{e_3}}\right) \hat{\boldsymbol{e_1}} + \left(\frac{d\hat{\boldsymbol{e_3}}}{dt}\cdot \hat{\boldsymbol{e_1}}\right)
    \hat{\boldsymbol{e_2}} + \left(\frac{d\hat{\boldsymbol{e_1}}}{dt}\cdot \hat{\boldsymbol{e_2}}\right) \hat{\boldsymbol{e_3}}
    \label{eq:angvel}
\end{equation}

### The derivative of the versors

Now, we must isolate the derivative of the versors to substitute them in the Eq. \eqref{eq:derivangmom}. To isolate the derivative of the versor $\hat{\boldsymbol{e_1}}$, first we cross multiply both sides of the equation above by $\hat{\boldsymbol{e_1}}$:

\begin{equation}
    \vec{\omega} \times \hat{\boldsymbol{e_1}} = - \left(\frac{d\hat{\boldsymbol{e_3}}}{dt}\cdot \hat{\boldsymbol{e_1}}\right) \hat{\boldsymbol{e_3}} + \left(\frac{d\hat{\boldsymbol{e_1}}}{dt}\cdot \hat{\boldsymbol{e_2}}\right) \hat{\boldsymbol{e_2}}
\end{equation}

If we note that the term multipliying $\hat{\boldsymbol{e_3}}$ in the right side of the identity can be obtained by $\frac{d\hat{\boldsymbol{e_1}}\cdot\hat{\boldsymbol{e_3}} }{dt} = \frac{d\hat{\boldsymbol{e_1}}}{dt}\cdot\hat{\boldsymbol{e_3}} + \frac{d\hat{\boldsymbol{e_3}}}{dt}\cdot\hat{\boldsymbol{e_1}} \rightarrow 0 = \frac{d\hat{\boldsymbol{e_1}}}{dt}\cdot\hat{\boldsymbol{e_3}} + \frac{d\hat{\boldsymbol{e_3}}}{dt}\cdot\hat{\boldsymbol{e_1}} \rightarrow \frac{d\hat{\boldsymbol{e_3}}}{dt}\cdot\hat{\boldsymbol{e_1}}  = - \frac{d\hat{\boldsymbol{e_1}}}{dt}\cdot\hat{\boldsymbol{e_3}} $ (the scalar product $\hat{\boldsymbol{e_1}}\cdot\hat{\boldsymbol{e_3}}$ is zero because these vectors are orthogonal), we can write the equation above becomes:

\begin{equation}
    \vec{\omega} \times \hat{\boldsymbol{e_1}} = \left(\frac{d\hat{\boldsymbol{e_1}}}{dt}\cdot \hat{\boldsymbol{e_3}}\right) \hat{\boldsymbol{e_3}} + \left(\frac{d\hat{\boldsymbol{e_1}}}{dt}\cdot \hat{\boldsymbol{e_2}}\right) \hat{\boldsymbol{e_2}}
\end{equation}

Finally, we can note that $\frac{d\hat{\boldsymbol{e_1}}\cdot\hat{\boldsymbol{e_1}} }{dt} = \frac{d\hat{\boldsymbol{e_1}}}{dt}\cdot\hat{\boldsymbol{e_1}} + \frac{d\hat{\boldsymbol{e_1}}}{dt}\cdot\hat{\boldsymbol{e_1}} \rightarrow \frac{d(1)}{dt} = 2\frac{d\hat{\boldsymbol{e_1}}}{dt}\cdot\hat{\boldsymbol{e_1}} \rightarrow \frac{d\hat{\boldsymbol{e_1}}}{dt}\cdot\hat{\boldsymbol{e_1}}  = 0 $. As this term is equal to zero, we can add it to the expression above:

\begin{equation}
    \vec{\boldsymbol{\omega}} \times \hat{\boldsymbol{e_1}} = \left(\frac{d\hat{\boldsymbol{e_1}}}{dt}\cdot\hat{\boldsymbol{e_1}}\right)\hat{\boldsymbol{e_1}} +  \left(\frac{d\hat{\boldsymbol{e_1}}}{dt}\cdot \hat{\boldsymbol{e_3}}\right) \hat{\boldsymbol{e_3}} + \left(\frac{d\hat{\boldsymbol{e_1}}}{dt}\cdot \hat{\boldsymbol{e_2}}\right) \hat{\boldsymbol{e_2}}
\end{equation}

Note that the expression above is just another manner to write the vector $\frac{d\hat{\boldsymbol{e_1}}}{dt}$, as any vector can be described by the sum of the projections on each of the versors forming a basis.

So, the derivative of the versor $\hat{\boldsymbol{e_1}}$  can be written as:

\begin{equation}
    \frac{d\hat{\boldsymbol{e_1}}}{dt} = \vec{\boldsymbol{\omega}} \times \hat{\boldsymbol{e_1}}
\end{equation}

Similarly, the derivatives of the versors $\hat{\boldsymbol{e_2}}$ and $\hat{\boldsymbol{e_3}}$ can be written as:

\begin{equation}
    \frac{d\hat{\boldsymbol{e_2}}}{dt} = \vec{\boldsymbol{\omega}} \times \hat{\boldsymbol{e_2}} ~~~~~~~~\text{and} ~~~~~~ \frac{d\hat{\boldsymbol{e_3}}}{dt} = \vec{\boldsymbol{\omega}} \times \hat{\boldsymbol{e_3}}
\end{equation}

### The derivative of the angular momentum

Now we can get back to the equation describing the derivative of the angular momentum (Eq.\eqref{eq:derivangmom}):

\begin{align}
    \begin{split}
    \frac{d\vec{H_{cm}}}{dt} =&I_1\dot{\omega_1}\hat{\boldsymbol{e_1}} + I_2\dot{\omega_2}\hat{\boldsymbol{e_2}}+I_3\dot{\omega_3}\hat{\boldsymbol{e_3}} + I_1\omega_1\frac{d\hat{\boldsymbol{e_1}}}{dt} + I_2\omega_2\frac{d\hat{\boldsymbol{e_2}}}{dt}+I_3\omega_3\frac{d\hat{\boldsymbol{e_3}}}{dt}=\\
    =& I_1\dot{\omega_1}\hat{\boldsymbol{e_1}} + I_2\dot{\omega_2}\hat{\boldsymbol{e_2}}+I_3\dot{\omega_3}\hat{\boldsymbol{e_3}} + I_1\omega_1(\vec{\boldsymbol{\omega}} \times \hat{\boldsymbol{e_1}}) + I_2\omega_2(\vec{\boldsymbol{\omega}} \times \hat{\boldsymbol{e_2}})+I_3\omega_3(\vec{\boldsymbol{\omega}} \times \hat{\boldsymbol{e_3}}) = \\
   =& I_1\dot{\omega_1}\hat{\boldsymbol{e_1}} + I_2\dot{\omega_2}\hat{\boldsymbol{e_2}}+I_3\dot{\omega_3}\hat{\boldsymbol{e_3}} + \vec{\boldsymbol{\omega}} \times I_1\omega_1\hat{\boldsymbol{e_1}} + \vec{\boldsymbol{\omega}} \times I_2\omega_2\hat{\boldsymbol{e_2}}+\vec{\boldsymbol{\omega}} \times I_3\omega_3\hat{\boldsymbol{e_3}} = \\
   =& I_1\dot{\omega_1}\hat{\boldsymbol{e_1}} + I_2\dot{\omega_2}\hat{\boldsymbol{e_2}}+I_3\dot{\omega_3}\hat{\boldsymbol{e_3}} + \vec{\boldsymbol{\omega}} \times (I_1\omega_1\hat{\boldsymbol{e_1}} + I_2\omega_2\hat{\boldsymbol{e_2}} +  I_3\omega_3\hat{\boldsymbol{e_3}})=\\
   =&I\vec{\dot{\omega}} + \vec{\omega} \times (I\vec{\omega})
   \end{split}
   \label{eq:derivangmomVec}
\end{align}

Performing the cross products, we can get the expressions for each of the coordinates attached to the body:

\begin{align}
    \begin{split}
    \frac{d\vec{H_{cm}}}{dt} =\left[\begin{array}{c}I_1\dot{\omega_1}\\I_2\dot{\omega_2}\\I_3\dot{\omega_3}\end{array}\right] + \left[\begin{array}{c}\omega_1\\\omega_2\\\omega_3\end{array}\right]  \times \left[\begin{array}{c}I_1\omega_1\\I_2\omega_2\\I_3\omega_3\end{array}\right]  = \left[\begin{array}{c}I_1\dot{\omega_1} + \omega_2\omega_3(I_3-I_2)\\I_2\dot{\omega_2}+\omega_1\omega_3(I_1-I_3)\\I_3\dot{\omega_3}+\omega_1\omega_2(I_2-I_1)\end{array}\right]
   \end{split}
\end{align}

## Newton-Euler laws

Having computed the derivative of the angular momentum, we have the final forms of the Newton-Euler laws:

\begin{equation}
    F_x = ma_{cm_x}
\end{equation}

\begin{equation}
    F_y = ma_{cm_y}
\end{equation}

\begin{equation}
    F_z = ma_{cm_z}
\end{equation}

\begin{equation}
    M_{cm_1} = I_1\dot{\omega_1} + \omega_2\omega_3(I_3-I_2)
    \label{eq:M1}
\end{equation}

\begin{equation}
    M_{cm_2} = I_2\dot{\omega_2}+\omega_1\omega_3(I_1-I_3)
    \label{eq:M2}
\end{equation}

\begin{equation}
    M_{cm_3} = I_3\dot{\omega_3}+\omega_1\omega_2(I_2-I_1)
    \label{eq:M3}
\end{equation}

Or, in the vectorial form:

\begin{equation}
    \vec{F} = m\vec{a_{cm}}
\end{equation}

\begin{equation}
    \vec{M_{cm}} = I\dot{\vec{\omega}}+\vec{\omega}\times(I\vec{\omega}) \text{            (vectors in local frame)}
\end{equation}

Note that the equations of the forces are written in the global frame of reference and the equations of the moments are described in the frame of reference of the body. So, before using Eq.~\eqref{eq:derivangmomVec} or the equations Eq.\eqref{eq:M1},\eqref{eq:M2} and \eqref{eq:M3} you must transform all the forces and moment-arms to the frame of reference of the body by using rotation matrices.

Below are shown some examples with three-dimensional kinematic data to find the forces and moments acting on the body.

## Examples

###  3D pendulum bar

 At the file '../data/3Dpendulum.txt' there are 3 seconds of data of 3 points of the three-dimensional cylindrical pendulum, very similar to the pendulum shown in the [notebook about free-body diagrams](FreeBodyDiagramForRigidBodies.ipynb#pendulum), except that it can move in every direction. Also it has a motor at the upper part of the cylindrical bar producing torques to move the bar. It has mass $m=1$ kg, length $l=1$ m and radius $r=0.1$ m.

 The point m1 is at the upper part of the cylinder and is the origin of the system.

 The point m2 is at the center of mass of the cylinder.

 The point m3 is a point at the surface of the cylinder.

The free-body diagram of the 3d pendulum is depicted below. There is the gravitational force acting at the center of mass of gravity of the body and the torque $\vec{M_1}$ due to the motor acting on the pendulum and the force $\vec{F_1}$ due to the restraint at the upper part of the cylinder. Together with the forces, the local basis $\hat{\boldsymbol{e_1}}$, $\hat{\boldsymbol{e_2}}$ and $\hat{\boldsymbol{e_3}}$ in the direction of the principal axes an origin at the center of mass of the body is also shown.

The resultant forces acting on the cylinder is:

\begin{equation}
    \vec{F} = \vec{F_O} - mg\hat{\boldsymbol{k}}
\end{equation}

So, the first Newton-Euler law, at each component of the global  basis, is written as:

\begin{align}
    \begin{split}
        F_{O_x} &=  ma_{cm_x} &\rightarrow  F_{O_x} &=  ma_{cm_x} \\
        F_{O_y} &= ma_{cm_y} &\rightarrow  F_{O_y} &=  ma_{cm_y}\\
        F_{O_z} - mg &= ma_{cm_z} &\rightarrow  F_{O_z} &=  ma_{cm_z} + mg
    \end{split}
    \label{eq:fnependulum}
\end{align}

Now, the resultant moment applied to the body, computed relative to the center of mass, is:

\begin{equation}
    \vec{M} = \vec{M_O} + \vec{r_{O/cm}} \times \vec{F_O}
\end{equation}

So, the second Newton-Euler law, at each of the components at the local basis of the body, is written as:

\begin{align}
    \begin{split}
        M_{O_1} + MFocm_1 &= I_1\dot{\omega_1} + \omega_2\omega_3(I_3-I_2) \rightarrow M_{O_1} &= I_1\dot{\omega_1} + \omega_2\omega_3(I_3-I_2) -  MFocm_1\\
        M_{O_2} + MFocm_2 &= I_2\dot{\omega_2} + \omega_1\omega_3(I_1-I_3) \rightarrow M_{O_2} &= I_2\dot{\omega_2} + \omega_1\omega_3(I_1-I_3) -  MFocm_2\\
        M_{O_3} + MFocm_3 &= I_3\dot{\omega_3} + \omega_1\omega_2(I_2-I_1) \rightarrow M_{O_3} &= I_3\dot{\omega_3} + \omega_1\omega_2(I_2-I_1) -  MFocm_3
    \end{split}
\end{align}

where $\vec{MFocm} = \vec{r_{O/cm}} \times \vec{F_O}$.

The moments of inertia at the directions of $\hat{\boldsymbol{e_1}}$, $\hat{\boldsymbol{e_2}}$ and $\hat{\boldsymbol{e_3}}$ are, $I_1 = \frac{mR^2}{12}$ and $I_2=I_3=\frac{m(3R^2+l^2)}{12}$. Now, to compute the moment $\vec{M_O}$ and the force  $\vec{F_O}$, we will need to find the acceleration of the center of mass $\vec{a_{cm}}$, the angular velocity  $\vec{\omega}$, the time-derivatives of each component of the angular velocity, and the moment-arm $\vec{r_{O/cm}}$ to compute the torque due to the force $\vec{F_O}$. These signals will come from the kinematic data file.

First,  we need to open the file with the data:

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

np.set_printoptions(precision=2)
f = open('../data/3dPendulum.txt')
data = np.loadtxt('../data/3dPendulum.txt',
                  skiprows=1, delimiter = ',')
header = f.readline()
print(header)
print(data)
```

    t(s),m1_X (m),m1_Y (m),m1_Z(m),m2_X (m),m2_Y (m),m2_Z(m),m3_X (m),m3_Y (m),m3_Z(m)

    [[ 0.    0.    0.   ...  0.22  0.12 -0.44]
     [ 0.01  0.    0.   ...  0.22  0.12 -0.44]
     [ 0.02  0.    0.   ...  0.22  0.12 -0.44]
     ...
     [ 2.97  0.    0.   ...  0.2   0.12 -0.45]
     [ 2.98  0.    0.   ...  0.2   0.12 -0.45]
     [ 2.99  0.    0.   ...  0.2   0.12 -0.45]]

```python
m = 1
g= 9.81
l = 1
r = 0.1

I1 = m*r**2/12
I2 = m*(3*r**2+l**2)/12
I3 = I2
```

Now, we assign the proper columns to variables:

```python
t = data[:,0]
m1 = data[:,1:4]
m2 = data[:,4:7]
m3 = data[:,7:]
```

As the center of mass is the data contained in m2, we can find the acceleration of the center of mass $\vec{\boldsymbol{a_{cm}}}$. This will be performed by deriving numerically the position of the center of mass twice. The numerical derivative of a function $f(t)$ with the samples $f(i)$ can be performed by taking the forward difference of the values $f(i)$:

\begin{equation}
    \frac{df}{dt}(i) = \frac{f(i+1)-f(i)}{\Delta t}
\end{equation}

The numerical derivative could be obtained by taking the backward differences as well:

\begin{equation}
    \frac{df}{dt}(i) = \frac{f(i)-f(i-1)}{\Delta t}
\end{equation}

A better estimation of the derivative of the time derivative of the function $f(t)$ would be obtained by the average value between the estimations using the backward and forward differences (this subject is treated [in this notebook about data filtering](DataFiltering.ipynb#numdiff)):

\begin{equation}
    \frac{df}{dt}(i) = \frac{\frac{f(i+1)-f(i)}{\Delta t} + \frac{f(i)-f(i-1)}{\Delta t}}{2} = \frac{f(i+1)-f(i-1)}{2\Delta t}
    \label{eq:centralderiv}
\end{equation}

So, the acceleration of the center of mass, is:

```python
dt = t[1]-t[0]

rcm = m2
vcm = (rcm[2:,:]-rcm[0:-2,:])/(2*dt)
acm = (vcm[2:,:]-vcm[0:-2,:])/(2*dt)
print(rcm.shape)
print(acm.shape)
```

    (300, 3)
    (296, 3)

Now we can find the force $\vec{\boldsymbol{F_O}}$  using the Eq. \eqref{eq:fnependulum}.

```python
Fox = m*acm[:,0]
Foy = m*acm[:,1]
Foz = m*acm[:,2] + m*g
Fo=np.hstack((Fox.reshape(-1,1),Foy.reshape(-1,1),Foz.reshape(-1,1)))

plt.figure()
plt.plot(t[0:acm.shape[0]], Fox)
plt.plot(t[0:acm.shape[0]], Foy, '--')
plt.plot(t[0:acm.shape[0]], Foz)
plt.legend(('x','y','z'))
plt.title('Force (N)')
plt.show()

```

Now, to find the moment being applied to the body, we need to compute a basis attached to the body

```python
e1 = m2 - m1
e1 = e1/np.linalg.norm(e1,axis=1,keepdims=True)

e2 = m3-m2
e2 = e2/np.linalg.norm(e2,axis=1,keepdims=True)

e3 = np.cross(e1,e2,axis=1)
e3 = e3/np.linalg.norm(e3,axis=1,keepdims=True)

e2 = np.cross(e3,e1, axis=1)
e2 = e2/np.linalg.norm(e2,axis=1,keepdims=True)

```

 To compute the moment applied to the body, we need the angular velocity described in the basis attached to the body. The easiest way to find this angular velocity is to use Eq. \eqref{eq:angvel}, repeated here.

 \begin{equation}
    \vec{\omega} =  \left(\frac{d\hat{\boldsymbol{e_2}}}{dt}\cdot \hat{\boldsymbol{e_3}}\right) \hat{\boldsymbol{e_1}} + \left(\frac{d\hat{\boldsymbol{e_3}}}{dt}\cdot \hat{\boldsymbol{e_1}}\right) \hat{\boldsymbol{e_2}} + \left(\frac{d\hat{\boldsymbol{e_1}}}{dt}\cdot \hat{\boldsymbol{e_2}}\right) \hat{\boldsymbol{e_3}}
\end{equation}

 To do this we need the derivatives of the basis versors. This will also be performed with Eq. \eqref{eq:centralderiv}.

 To perform the computation of the angular velocity remember that the scalar product between two vectors is given by:

 \begin{equation}
         \vec{v}\cdot\vec{w} = \left[\begin{array}{c}v_x\\v_y\\v_z \end{array}\right]\cdot \left[\begin{array}{c}w_x\\w_y\\w_z \end{array}\right] =  v_x.w_x+v_yw_y+v_zw_z
 \end{equation}

```python
de1dt = (e1[2:,:]-e1[0:-2,:])/(2*dt)
de2dt = (e2[2:,:]-e2[0:-2,:])/(2*dt)
de3dt = (e3[2:,:]-e3[0:-2,:])/(2*dt)

omega = np.hstack((np.sum(de2dt*e3[1:-1,:], axis = 1).reshape(-1,1),
                   np.sum(de3dt*e1[1:-1,:], axis = 1).reshape(-1,1),
                   np.sum(de1dt*e2[1:-1,:], axis = 1).reshape(-1,1)))
```

From the angular velocity vector we can obtain the derivatives of each component of it, also needed to compute the moment applied to the body. To do this we will use Eq. \eqref{eq:centralderiv}:

```python
alpha = (omega[2:,:]-omega[0:-2,:])/(2*dt)
```

It remains to find the moment caused by the force $\vec{F_O}$, $\vec{MFocm} = \vec{r_{O/cm}} \times \vec{F_O}$. The moment-arm $\vec{r_{O/cm}} =-\vec{r_{cm}}$.

```python
MFocm = np.cross(-rcm[2:-2], Fo)
```

The problem is that this moment is in the global basis. We need to transform it to the local basis. This will be performed using the rotation matrix of the bar. Each row of this matrix is one of the basis versors. Note that at each instant the matrix of rotation $R$ will be different.  After the matrix is formed, we can find the components of the moment $\vec{MFocm}$ in the local basis by multiplying the matrix of rotation $R$ by the vector $\vec{MFocm}$.

```python
MFocmLocal = np.zeros_like(MFocm)
for i in range(MFocm.shape[0]):
    R = np.vstack((e1[i,:],e2[i,:],e3[i,:]))
    MFocmLocal[i,:] = R@MFocm[i,:]

```

```python
Mo1 = I1*alpha[:,0] + omega[0:alpha.shape[0],1]*omega[0:alpha.shape[0],2]*(I3-I2) - MFocmLocal[:,0]
Mo2 = I2*alpha[:,1] + omega[0:alpha.shape[0],0]*omega[0:alpha.shape[0],2]*(I1-I3) - MFocmLocal[:,1]
Mo3 = I3*alpha[:,2] + omega[0:alpha.shape[0],0]*omega[0:alpha.shape[0],1]*(I2-I1) - MFocmLocal[:,2]
plt.figure()
plt.plot(t[2:-2], Mo1)
plt.plot(t[2:-2], Mo2)
plt.plot(t[2:-2], Mo3)
plt.legend(('$e_1$','$e_2$','$e_3$'))
plt.show()
```

We could also have used the vectorial form of the derivative of the angular momentum (Eq. \eqref{eq:derivangmomVec}) and instead of writing three lines of code, write only one. The result is the same.

```python
I = np.array([[I1,0,0],[0,I2,0],[0,0,I3]])
Mo = (I@alpha.T).T  + np.cross(omega[0:alpha.shape[0],:], (I@omega[0:alpha.shape[0],:].T).T,axis=1) - MFocmLocal
plt.figure()
plt.plot(t[2:-2], Mo)
plt.legend(('$e_1$','$e_2$','$e_3$'))
plt.show()
```

### Data from postural control

This example will use real data from a subject during quiet standing during 60 seconds. This data is from the database freely available at [https://github.com/demotu/datasets/tree/master/PDS](https://github.com/demotu/datasets/tree/master/PDS). The data of this subject is in the file '../data/postureData.txt'.

The mass of the subject was $m = 53$ kg and her height was $h= 1.65 $ m.

   The free-body diagram is very similar to the free-body diagram shown [in the notebook about free-body diagram](FreeBodyDiagramForRigidBodies.ipynb#quietstanding), except that the force $\vec{F_A}$ and the moment $\vec{M_A}$ have components at all 3 directions.

   So, the first Newton-Euler law, at each component of the global  basis, is written as (note that in these data, the vertical direction is the y coordinate):

\begin{align}
    \begin{split}
        F_{A_x} &=  ma_{cm_x} &\rightarrow  F_{A_x} &=  ma_{cm_x} \\
        F_{A_y}  - mg &= ma_{cm_y} &\rightarrow  F_{A_y} &=  ma_{cm_y} + mg\\
        F_{A_z} &= ma_{cm_z} &\rightarrow  F_{A_z} &=  ma_{cm_z}
    \end{split}
    \label{eq:fnequiet}
\end{align}

Now, the resultant moment applied to the body, computed relative to the center of mass, is:

\begin{equation}
    \vec{M} = \vec{M_A} + \vec{r_{A/cm}} \times \vec{F_A}
\end{equation}

So, the second Newton-Euler law, at each of the components at the local basis of the body, is written as:

\begin{align}
    \begin{split}
        \vec{M_A} + \vec{MFacm} &= I\left[\begin{array}{c}\dot{\omega_1}\\\dot{\omega_2}\\\dot{\omega_3}\end{array}\right] + \vec{\omega}\times  I\vec{\omega} \rightarrow \vec{M_A} = I\left[\begin{array}{c}\dot{\omega_1}\\\dot{\omega_2}\\\dot{\omega_3}\end{array}\right] + \vec{\omega} \times  I\vec{\omega}- \vec{MFacm}
    \end{split}
\end{align}

where  $\vec{MFAcm} = \vec{r_{A/cm}} \times \vec{F_A}$.

Now we open the data and assign the coordinates of each marker to a variable.

```python
data = np.loadtxt('../data/postureData.txt', skiprows=1, delimiter = ',')

t = data[:,0]
dt = t[2]-t[1]
rcm = data[:,1:4] #center of mass
rrA = data[:,4:7] # Right lateral malleolus
rlA = data[:,7:] # Left lateral maleolus

```

 The body will be approximated by a cylinder with the height of the subject and radius equal to half of the mean distances between the right and left medial malleoli.

```python
m = 53
h = 1.65
r = np.mean(np.linalg.norm(rrA-rlA, axis = 1))/2
I1 = m*r**2/12 # longitudnal
I2 = m*(3*r**2+h**2)/12 # sagittal
I3 = I2 # transversal

```

```python
# acceleration of the center of mass by deriving the center of mass twice
vcm = (rcm[2:,:]-rcm[0:-2,:])/(2*dt)
acm = (vcm[2:,:]-vcm[0:-2,:])/(2*dt)

FAx = m*acm[:,0]
FAy = m*acm[:,1] + m*g
FAz = m*acm[:,2]
FA=np.hstack((FAx.reshape(-1,1),FAy.reshape(-1,1),FAz.reshape(-1,1)))

```

Now we form the basis attached to the body. The first versor $\hat{\boldsymbol{e_1}}$ will be a versor from the midpoint between the medial malleoli and the center of mass of the body. The second versor $\hat{\boldsymbol{e_2}}$ will be a versor from  the right to the left malleolus. The third versor $\hat{\boldsymbol{e_3}}$ will be a cross product between $\hat{\boldsymbol{e_1}}$ and $\hat{\boldsymbol{e_2}}$.

```python
e1 = rcm - (rlA+rrA)/2
e1 = e1/np.linalg.norm(e1,axis=1,keepdims=True)

e2 = rlA-rrA
e2 = e2/np.linalg.norm(e2,axis=1,keepdims=True)

e3 = np.cross(e1,e2,axis=1)
e3 = e3/np.linalg.norm(e3,axis=1,keepdims=True)

e2 = np.cross(e3,e1, axis=1)
e2 = e2/np.linalg.norm(e2,axis=1,keepdims=True)

```

Now we can find the angular velocity $\vec{\omega}$ at the basis attached to the body using Eq.\eqref{eq:angvel} and the time derivatives of its components.

```python
de1dt = (e1[2:,:]-e1[0:-2,:])/(2*dt)
de2dt = (e2[2:,:]-e2[0:-2,:])/(2*dt)
de3dt = (e3[2:,:]-e3[0:-2,:])/(2*dt)

omega = np.hstack((np.sum(de2dt*e3[1:-1,:], axis = 1).reshape(-1,1),
                   np.sum(de3dt*e1[1:-1,:], axis = 1).reshape(-1,1),
                   np.sum(de1dt*e2[1:-1,:], axis = 1).reshape(-1,1)))

alpha = (omega[2:,:]-omega[0:-2,:])/(2*dt)

```

Now we need to find the moment caused by the force at the ankles $\vec{F_A}$, $\vec{MFAcm} = \vec{r_{A/cm}} \times \vec{F_A}$. The moment-arm $\vec{r_{A/cm}}$ is the vector from the center of mass to the midpoint of the lateral malleoli.

Besides the description of the moment due to the force $\vec{F_A}$ in the basis attached to the body, we will also describe the force $\vec{F_A}$ at the local basis. This is useful because it has an anatomical meaning. After this we can use the equations of Newton-Euler to obtain the moment at the ankle.

After having all signals described in the basis of the body, the moment being applied by the muscles of the ankle is computed using Eq. \eqref{eq:derivangmomVec}.

```python

#computing the moment due to the ankle force
racm = (rlA+rrA)/2-rcm
MFAcm = np.cross(racm[0:FA.shape[0],:], FA)

MFAcmLocal = np.zeros_like(MFAcm)
FALocal = np.zeros_like(MFAcm)

# rotation matrix and description of the ankle force and moment due to the ankle force
for i in range(MFAcm.shape[0]):
    R = np.vstack((e1[i,:],e2[i,:],e3[i,:]))
    MFAcmLocal[i,:]=R@MFAcm[i,:]
    FALocal[i,:]=R@FA[i,:]

# Second Newton-Euler law to obtain the ankle moment

I = np.array([[I1,0,0],[0,I2,0],[0,0,I3]])
MA = (I@alpha.T).T  + np.cross(omega[0:alpha.shape[0],:],
                               (I@omega[0:alpha.shape[0],:].T).T,axis=1) - MFAcmLocal

plt.figure()
plt.plot(t[2:-2], MA)
plt.legend(('longitudinal','sagittal','mediolateral'))
plt.title('Ankle Torque')
plt.xlabel('t (s)')
plt.ylabel('M (N.m)')

plt.show()

plt.figure()
plt.plot(t[2:-2], FALocal[:,0])
plt.plot(t[2:-2], FALocal[:,1])
plt.plot(t[2:-2], FALocal[:,2])
plt.title('Ankle Force')
plt.legend(('longitudinal','sagittal','mediolateral'))
plt.xlabel('t (s)')
plt.ylabel('F (N)')
plt.show()
```

## Problems

1) Compute the derivative of the angular momentum of the foot and the leg using one of  the following data acquired during the gait of a subject: ['../data/BiomecII2018_gait_d.txt'](../data/BiomecII2018_gait_d.txt) or ['../data/BiomecII2018_gait_n.txt'](../data/BiomecII2018_gait_n.txt).

2) Problem 20.2.7 from Ruina and Rudra book.

## References

- Beer, F P; Johnston, E R; Cornwell, P. J.(2010) Vector Mechanics for Enginners: Dynamics.
- Kane T, Levinson D (1985) [Dynamics: Theory and Applications](https://ecommons.cornell.edu/handle/1813/638). McGraw-Hill, Inc
- Hibbeler, R. C. (2005) Engineering Mechanics: Dynamics.
- Taylor, J, R (2005) Classical Mechanics
- Winter D. A., (2009) Biomechanics and motor control of human movement. John Wiley and Sons.
- Santos DA, Fukuchi CA, Fukuchi RK, Duarte M. (2017) A data set with kinematic and ground reaction forces of human balance. PeerJ Preprints.
- Ruina A., Rudra P. (2019) [Introduction to Statics and Dynamics](http://ruina.tam.cornell.edu/Book/index.html). Oxford University Press.