# Tutorial Notebook 2. Inference and Planning

In this notebook, we will continue on from the last notebook to build a fully fledged active inference agent capable of performing inference and planning using Active Inference in the simple grid-world environment. We will also begin to use some aspects of `pymdp`, although this will mostly be helpful functions for building and sampling from discrete distributions while we will implement the core functionality of the agent ourselves.

First, we simply start out by defining our generative model as we did last time.

## Add `pymdp` module

```python
# This is needed (on my machine at least) due to weird python import issues
import os
import sys
from pathlib import Path

path = Path(os.getcwd())
print(path)
module_path = str(path.parent) + '/'
sys.path.append(module_path)
```

    /Users/conor/pymdp/examples

## Imports

```python
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
from copy import deepcopy

from pymdp import maths, utils
from pymdp.maths import spm_log_single as log_stable # @NOTE: we use the `spm_log_single` helper function from the `maths` sub-library of pymdp. This is a numerically stable version of np.log()
from pymdp import control
print("imports loaded")
```

    imports loaded

## Plotting

```python
state_mapping = {0: (0,0), 1: (1,0), 2: (2,0), 3: (0,1), 4: (1,1), 5:(2,1), 6: (0,2), 7:(1,2), 8:(2,2)}

A = np.eye(9)
def plot_beliefs(Qs, title=""):
    #values = Qs.values[:, 0]
    plt.grid(zorder=0)
    plt.bar(range(Qs.shape[0]), Qs, color='r', zorder=3)
    plt.xticks(range(Qs.shape[0]))
    plt.title(title)
    plt.show()

labels = [state_mapping[i] for i in range(A.shape[1])]
def plot_likelihood(A):
    fig = plt.figure(figsize = (6,6))
    ax = sns.heatmap(A, xticklabels = labels, yticklabels = labels, cbar = False)
    plt.title("Likelihood distribution (A)")
    plt.show()

def plot_empirical_prior(B):
    fig, axes = plt.subplots(3,2, figsize=(8, 10))
    actions = ['UP', 'RIGHT', 'DOWN', 'LEFT', 'STAY']
    count = 0
    for i in range(3):
        for j in range(2):
            if count >= 5:
                break

            g = sns.heatmap(B[:,:,count], cmap="OrRd", linewidth=2.5, cbar=False, ax=axes[i,j])

            g.set_title(actions[count])
            count += 1
    fig.delaxes(axes.flatten()[5])
    plt.tight_layout()
    plt.show()

def plot_transition(B):
    fig, axes = plt.subplots(2,3, figsize = (15,8))
    a = list(actions.keys())
    count = 0
    for i in range(dim-1):
        for j in range(dim):
            if count >= 5:
                break
            g = sns.heatmap(B[:,:,count], cmap = "OrRd", linewidth = 2.5, cbar = False, ax = axes[i,j], xticklabels=labels, yticklabels=labels)
            g.set_title(a[count])
            count +=1
    fig.delaxes(axes.flatten()[5])
    plt.tight_layout()
    plt.show()
```

## Generative model

Here, we setup our generative model which is the same as in the last notebook. This is formed of a likelihood distribution $P(o_t|s_t)$, denoted `A`, and a empirical prior (transition) distribution $P(s_t|s_{t-1},a_{t-1})$, denoted `B`.

Since this was covered in more detail in the previous tutorial, we quickly skip over the details here

```python
# A matrix
A = np.eye(9)
plot_likelihood(A)
```

```python
# construct B matrix

P = {}
dim = 3
actions = {'UP':0, 'RIGHT':1, 'DOWN':2, 'LEFT':3, 'STAY':4}

for state_index, xy_coordinates in state_mapping.items():
    P[state_index] = {a : [] for a in range(len(actions))}
    x, y = xy_coordinates

    '''if your y-coordinate is all the way at the top (i.e. y == 0), you stay in the same place -- otherwise you move one upwards (achieved by subtracting 3 from your linear state index'''
    P[state_index][actions['UP']] = state_index if y == 0 else state_index - dim

    '''f your x-coordinate is all the way to the right (i.e. x == 2), you stay in the same place -- otherwise you move one to the right (achieved by adding 1 to your linear state index)'''
    P[state_index][actions["RIGHT"]] = state_index if x == (dim -1) else state_index+1

    '''if your y-coordinate is all the way at the bottom (i.e. y == 2), you stay in the same place -- otherwise you move one down (achieved by adding 3 to your linear state index)'''
    P[state_index][actions['DOWN']] = state_index if y == (dim -1) else state_index + dim

    ''' if your x-coordinate is all the way at the left (i.e. x == 0), you stay at the same place -- otherwise, you move one to the left (achieved by subtracting 1 from your linear state index)'''
    P[state_index][actions['LEFT']] = state_index if x == 0 else state_index -1

    ''' Stay in the same place (self explanatory) '''
    P[state_index][actions['STAY']] = state_index

num_states = 9
B = np.zeros([num_states, num_states, len(actions)])
for s in range(num_states):
    for a in range(len(actions)):
        ns = int(P[s][a])
        B[ns, s, a] = 1

plot_transition(B)
```

# Create Environment Class

To make things simple we will parcel up the $A$ and $B$ matrices into a class which represents the environment. The environment has two functions `step` which when given an action will update the environment a single step and `reset` which resets the environment back to its initial condition. The API of our simple environment class is similar to the `Env` base class used by `pymdp`, although the `pymdp` version has many more features than we use here.

```python
class GridWorldEnv():

    def __init__(self,A,B):
        self.A = deepcopy(A)
        self.B = deepcopy(B)
        print("B:", B.shape)
        self.state = np.zeros(9)
        # start at state 3
        self.state[2] = 1

    def step(self,a):
        self.state = np.dot(self.B[:,:,a], self.state)
        obs = utils.sample(np.dot(self.A, self.state))
        return obs

    def reset(self):
        self.state =np.zeros(9)
        self.state[2] =1
        obs = utils.sample(np.dot(self.A, self.state))
        return obs

env = GridWorldEnv(A,B)
```

    B: (9, 9, 5)

# Inference

Now that we have the generative model setup, we turn to the behaviour of the active inference agent itself. To recap, we assume that this agent receives observations from the environment and can emit actions. Moreover, we assume that this agent has some kind of goal or preferences over the state of the environment it wants to create, and will choose actions in order to increase the probability of observing itself in its preferred state. For the time being, we will not deal with the problem of action selection but only with inference.

The agent receives observations $o_t$ from the environment but does not naturally know the environments true state $x_t$. Thus, the agent must *infer* this state by computing the posterior distribution $p(x_t | o_t)$. It can do this by Bayesian inference using Bayes rule but, as we discussed last time, explicitly computing Bayes rule is often intractable because the marginal likelihood requires the averaging over an infinite number of hypotheses. We therefore need some other way to compute or approximate this posterior. Active inference assumes that this posterior can be approximated through a family of methods called *variational inference* which only approximate the posterior, but are fast and computationally efficient.

# Variational Inference

Variational inference is a set of inference methods which can rapidly and efficiently compute *approximate* posteriors for Bayesian inference problems. The key idea behind variational inference is that instead of trying to compute the true posterior $p(x_t | o_t)$ which may be extremely complex, is that we instead optimize an *approximate posterior*. Specifically, we will define another distribution $q(x_t ; \phi)$ which has some parameters $\phi$ which we then optimize so as to make $q(x_t ; \phi)$ as close as possible to the true distribution. Typically, we choose this $q$ distribution to be some simple distribution which is easy to work with mathematically. If the process works, then we can get the $q$ distribution very close to the true posterior, and as such get a good estimate of the posterior without ever explicitly computing it using Bayes rule.

Mathematically, we can do this by setting up an optimization problem. We have the true posterior $p(x_t | o_t)$, which is unknown, and we have our $q$ distribution which we do know. We then want to optimize the $q$ distribution to make it as *close as possible* to the true posterior $p(x_t | o_t)$. To do this, we first need a way to quantify *how close* two probability distributions are. The way we do this is by using a quantity known as the *Kullback-Leibler (KL) divergence*. This is a metric derived from information theory which lets us quantify the distance (in bits) of two distributions. The KL divergence between two distributions $q(x)$ and $p(x)$ is defined as,
$$
\begin{align}
KL[q(x) || p(x)] = \sum_x q(x) (\ln q(x) - \ln p(x))
\end{align}
$$
Mathematically it can be thought of as the average of the difference of the logarithms of the probabilities assigned by $q$ and $p$ to the states $x$. The KL divergence is smallest when $q(x) = p(x)$ when it is equal to 0, and can grow to be infinitely large which happens wherever $q(x)$ assigns a nonzero probability but $p(x)$ doesn't. In code, we can compute the KL divergence as:

```python
def KL_divergence(q,p):
    return np.sum(q * (log_stable(q) - log_stable(p)))
```

Now that we know about the KL divergence, we can express our variational problem of making our approximate posterior $q(x_t ; \varphi)$ as close as possible to the true posterior as simply minimizing the KL divergence between the two distributions. That is, we can define the optimal approximate distribution as,
$$
\begin{align}
q^*(x_t ; \varphi) = \operatorname{argmin}_{\varphi} \, KL[q(x_t ; \varphi) || p(x_t | o_t)]
\end{align}
$$

And then simply try to optimize this objective so as to find the setting of the variational parameters $\varphi$ that make $q(x_t ; \varphi)$ as close to $p(x_t |o_t)$ as possible. The trouble with this is that our objective actually explicitly contains the true posterior in it and so, since we can't conpute the true posterior, we can't compute this objective either -- so we are stuck!

Variational inference provides a clever way to get around this problem by instead minimizing an *upper bound* on this divergence called the *variational free energy*. Importantly this bound is computable so we can actually optimize it and moreover since it is an upper bound, if we minimize it, we can make $q$ as close as possible to the real bound, thus still managing to obtain a good approximate posterior distribution. Deriving the variational free energy is very simple, we first take our initial objective and apply Bayes rule to the true posterior, and then take out the marginal likelihood term separately
<!-- $$
\begin{align}
 KL[q(x_t | o_t) || p(x_t |  o_t)] &= KL[q(x_t | o_t) || \frac{p(o_t,x_t)}{p(o_t)}] \\
 &= KL[q(x_t | o_t) || p(o_t, x_t)] + \sum_x q(x_t | o_t) \ln p(o_t) \\
 &= KL[q(x_t | o_t) || p(o_t, x_t)] +  \ln p(o_t)
\end{align}
$$ -->

$$
\begin{align}
 KL[q(x_t ; \varphi) || p(x_t |  o_t)] &= KL[q(x_t ; \varphi) || \frac{p(o_t,x_t)}{p(o_t)}] \\
 &= KL[q(x_t ; \varphi) || p(o_t, x_t)] + \sum_x q(x_t ; \varphi) \ln p(o_t) \\
 &= KL[q(x_t ; \varphi) || p(o_t, x_t)] +  \ln p(o_t)
\end{align}
$$

Where in the final line we have used the fact that the sum is over a different variable than the distribution, and the sum of a probability distribution is $1$ -- i.e $\sum_x q(x_t ; \varphi) \ln p(o_t) = \ln p(o_t) * \sum_x q(x_t ; \varphi)$ and $\sum_x q(x_t ; \varphi) = 1$. Specifically, since $ \ln p(o_t)$ is the log of a probability distribution it is always negative, since the probability of a state is always between 0 and 1. This means that we know that this term we have devised $KL[q(x_t ; \varphi) || p(o_t, x_t)]$ is always necessarily greater than our original divergence between the approximate posterior and the true posterior, so it is an *upper bound*. We call this term the *variational free energy* and denote it by $\mathcal{F}$.
$$
\begin{align}
\mathcal{F} = KL[q(x_t ; \varphi) || p(o_t, x_t)]
\end{align}
$$

The free energy here is simply the divergence between the approximate posterior and the *generative model* of the agent. Since we know both the approximate posterior (as we defined it in the first place!) and the generative model, then both terms of this divergence are computable. We thus have our algorithm to approximate the posterior! Since the free energy is an upper bound, if we minimize the free energy, we also implicitly minimize the true divergence between the true and approximate posteriors, which will force the approximate posterior to be close to the true posterior and thus a good approximation! Moreover, since we can compute the free energy, we can actually perform this optimization!

In many cases, we typically perform variational inference by taking the gradients of the free energy with respect to the variational parameteres $\varphi$ and then doing gradient descent on the parameters $\varphi$ that define $q(x_t ; \varphi)$. However, when the distributions are discrete (i.e. Categorical distributions), the parameters of the approximate distribution are simply the probability values for each state (the elements of the vector $q(x)$. For some simple generative models (e.g. the grid-world described here), we can actually solve this optimization problem directly to obtain $q(x)$ in a single step, instead of as a gradient descent.

# Directly solving variational inference in the case of a simple discrete model

To recap, remember that we have turned the problem of computing the posterior distribution $p(x_t | o_t)$ into that of minimizing the variational free energy: $\mathcal{F} = KL[q(x_t) || p(o_t, x_t)]$ with respect to an approximate posterior distribution $q(x_t)$. From now on, we will leave out the variational parameters $\varphi$ when referring to $q(x_t)$, since we are dealing with a single Categorical distribution $q(x_t)$ whose vector elements are identical to the variational parameters, i.e. $\forall_{i} \varphi_i = q(x_t = i)$.

The optimal distribution is simply that particular $q^*(x_t)$ that minimizes the KL divergence. Now, remember from high-school calculus that we can explicitly compute the minimum of a function by taking its derivative and setting it to 0 (i.e. at the minimum the first derivative of the function is 0) (if you don't remember this from calculus, trust me on this). This means that to solve this problem all we need to do is take the derivative of the free energy and set it to 0 and rearrange. First, let's write out the free energy explicitly.

<!-- $$
\begin{align}
\mathcal{F} &= KL[q(x_t | o_t) || p(o_t, x_t)] \\
&= \sum_x q(x_t | o_t) (\ln q(x_t | o_t) - \ln p(o_t,x_t))
\end{align}
$$ -->

$$
\begin{align}
\mathcal{F} &= KL[q(x_t) || p(o_t, x_t)] \\
&= \sum_x q(x_t)(\ln q(x_t) - \ln p(o_t,x_t))
\end{align}
$$
If we then split the generative model up into a likelihood and a prior, we can write it as,

$$
\begin{align}
\mathcal{F} = \sum_x q(x_t) \big[ \ln q(x_t) - \ln p(o_t | x_t)  - \ln p(x_t) \big]
\end{align}
$$
Recall that we have explicitly defined the observation and transition likelihood distributions as the $\textbf{A}$ and $\textbf{B}$ matrices, which play the role of the likelihood and prior distributions, respectively. In particular, we can define the prior over the current timestep $p(x_t)$ to be the "expected prior", given the beliefs about the state at the last timesep $q(x_{t-1})$ the beliefs about the transition dynamics and the past action, i.e.:

$$
\begin{align}
p(x_t) = \mathcal{E}_{q(x_{t-1})}\big[p(x_t | x_{t-1}, a_t) \big]
\end{align}
$$

which can be expressed as a simple matrix vector product

$$
\begin{align}
p(x_t) = \textbf{B}_{a_t}q(x_{t-1})
\end{align}
$$

where $\textbf{B}_{a_t}$ is the component of the $\textbf{B}$ matrix that is conditioned on action $a_t$. In other words, for this simple Markovian model, we assume that "yesterday's posterior is today's prior." For simplicity, we refer to this entire prior term as $\textbf{B}$ below, i.e. we temporarily define $\textbf{B} \equiv \textbf{B}_{a_t}q(x_{t-1})$.

As mentioned above, the posterior beliefs $q(x_t)$ are a vector of probabilities (the variational parameters) which we denote $\textbf{q} = [q_1, q_2, q_3 \dots]$. With this all defined, we can write out the free energy as,

$$
\begin{align}
\mathcal{F} = \sum_x \textbf{q} * \big[ \ln \textbf{q} - \ln \textbf{A}  - \ln \textbf{B} \big]
\end{align}
$$

And for fun, we can explicitly compute it in code:

```python
def compute_free_energy(q,A, B):
    return np.sum(q * (log_stable(q) - log_stable(A) - log_stable(B)))
```

Then, all we need to do is take the derivative of the free energy with respect to the approximate posterior distribution $\textbf{q}$ as follows,
$$
\begin{align}
\frac{\partial \mathcal{F}}{\partial \textbf{q}} = \ln \textbf{q} - \ln \textbf{A} - \ln \textbf{B} - \textbf{1}
\end{align}
$$

Where $\textbf{1}$ is just a vector of ones of equal length to $\textbf{q}$ and comes from the $q \frac{\partial \ln q}{\partial q} = q * \frac{1}{q} = 1$. Thus, if we set this derivative to 0 and rearrange, we can get,
$$
\begin{align}
0 &= \ln \textbf{q} - \ln \textbf{A} - \ln \textbf{B} - \textbf{1} \\
&\implies \textbf{q}^* = \sigma(\ln \textbf{A} + \ln \textbf{B})
\end{align}
$$

Where $\sigma$ is a softmax function $\sigma(x) = \frac{e^x}{\sum_x e^x}$ which ensures that the resulting probability distribution is normalized. This expression lets us compute the optimal approximate posterior instantly as a straightforward function of the current observation $o_t$, $\textbf{A}$, and $\textbf{B}$. We can thus quickly right the code for inference:

```python
def softmax(x):
    return np.exp(x) / np.sum(np.exp(x))

def perform_inference(likelihood, prior):
    return softmax(log_stable(likelihood) + log_stable(prior))
```

Note that the likelihood term is not the entire $\textbf{A}$ matrix, but just the 'row' of the $\textbf{A}$ matrix corresponding to the current observation, i.e. $P(o_t = o \mid x_t)$.

Inference for simple discrete state-space models like these is therefore very simple. All we need to do is have some initial set of beliefs $\textbf{q}_0$ and then update them according to these rules for every observation we get, using the $\textbf{A}$, $\textbf{B}$, the past action $a_t$ and the past posterior $q(x_{t-1})$ to provide the likelihood and prior terms.

# Planning through Active Inference

So far we just have an agent that can perform perform inference in the discrete state space, but inference by itself isn't really that useful. Instead what we really want to do is *planning*. That is, the agent needs to be able to figure out how to emit a series of actions which will take it to a certain goal state. A key part of Active Inference is that this process of planning, or more broadly action selection, can also be solved as a process of variational inference. This is why it is called *Active* Inference, after all.

<!-- However, when starting to think about this, it is not immediately obvious how to turn planning into an inference problem. What are the hypotheses? What are the observations? To turn the problem of planning into an inference problem, we need to introduce two additional concepts. The first is the idea of a *goal state*. To make planning useful, the agent has to *want something*. This is different from just performing objective inference about the state of the world in which there is no goal except to infer correctly. In Active Inference we define the goal state as a separate goal vector denoted $\textbf{C}$. When performing planning we then modify the generative model of the agent so that it no longer reflects the true distribution of observations in the environment, but rather includes the goal vector. We denote this new generative model $\tilde{p}(o_t, x_t) = p(o_t | x_t)\tilde{p}(x_t)$ where we use $\tilde{p}$ to say that this distribution is not a true distribution describing the agents model of the world, but is instead a *desired distribution*. Here, we set the desired distribution to be over the state of the environment $x_t$ such that $\tilde{p}(x_t) = \textbf{C}$. -->

However, when starting to think about this, it is not immediately obvious how to turn planning into an inference problem. What are the hypotheses? What are the observations? To turn the problem of planning into an inference problem, we need to introduce two additional concepts. The first is the idea of a *preferred observations*. To make planning useful, the agent has to *want something*. This is different from just performing objective inference about the state of the world in which there is no goal except to infer correctly. In Active Inference, we define the preferred observations as a separate goal vector denoted $\textbf{C}$. When performing planning we then modify the generative model of the agent so that it no longer reflects the true distribution of observations in the environment, but rather includes the goal vector, which is a *prior over observations*. We denote this new generative model $\tilde{p}(o_t, x_t) = p(x_t | o_t)\tilde{p}(o_t)$ where we use $\tilde{p}$ to say that this distribution is not a true distribution describing the agent's model of the world, but is instead a *preference distribution*. Here, we set the preference distribution to be over the state of the environment $o_t$ such that $\tilde{p}(o_t) = \textbf{C}$.

By changing the generative model in this way, we have effectively changed the inference problem from: *infer the most likely states and actions given the true generative model of the world* to *infer the most likely states and actions given a false model of the world, in which I achieve my goals*. Perhaps more intuitively, we can think of this inference problem as answering the question: *Given that I have achieved my goals, what actions must have I taken to get there?*.

The second thing we need to do to perform planning is to also extend the inference problem to *actions in the future*, since these are the fundamental things that the agent controls which it can use to adjust the environment. We call a sequence of future actions from now (time $t$) until some set future time $T$ a *policy* and denote it $\pi = [a_t, a_{t+1}, a_{t+2} \dots a_T]$. The goal is then to infer the optimal policy $\pi^*$ given the preferences $\textbf{C}$.

However here there is a problem. Typically we would perform variational inference to solve this inference problem, but the variational free energy is not defined over future trajectories of observations, which are not yet known. Instead, we define a new objective which can handle this -- the *Expected Free Energy (EFE)*, which is defined over policies and which we denote as $\mathcal{G}(\pi)$. We define the EFE for a particular timepoint $\tau$ and policy $\pi$ as:

$$
\begin{align}
\mathcal{G}(\pi)_{\tau} = \sum_{o_{\tau}, x_{\tau}} q(x_{\tau} | \pi)q(o_{\tau} | \pi) \big[ \ln q(x_{\tau} | \pi) - \ln \tilde{p}(o_{\tau}, x_{\tau}) \big]
\end{align}
$$

The expected free energy is defined for a single time-step of a trajectory *in the future*, i.e. prior to receiving any observations. The key difference between the standard variational free energy and the expected free energy is that the expected free energy also averages over the *expected observations* $q(o_{\tau} | pi)$ and *expected states* $q(x_{\tau} | \pi)$, where the expectations are conditioned on some policy $\pi$. This is necessary because with the expected free energy, we are evaluating possible future trajectories without a given observation, unlike the variational free energy where we can assume that we have already received the observation.

To get an intuitive handle on what the expected free energy *means* we can decompose it into two more intuitive quantities.

$$
\begin{align}
\mathcal{G}(\pi)_{\tau}&= \sum_{o_{\tau}, x_{\tau}} q(x_{\tau} | \pi)q(o_{\tau} | \pi) \big[ \ln q(x_{\tau} | \pi) - \ln \tilde{p}(o_{\tau}, x_{\tau}) \big] \\
&= -\underbrace{\sum_{o_{\tau}, x_{\tau}} q(x_{\tau} | \pi)q(o_{\tau} | \pi) \big[ \ln p(o_{\tau} | x_{\tau}) \big]}_{\text{Uncertainty}} + \underbrace{KL[q(x_{\tau} | \pi) || \tilde{p}(x_{\tau}) ]}_{\text{Divergence}}
\end{align}
$$

The first term is called *expected uncertainty* or sometimes *novelty* and represents essentially the spread of the observations expected in the future. Since we are choosing actions that *minimize* the whole quantity, we want to take actions that *maximize* novelty. We can think of this as a bonus to aid exploration since active inference agents will preferentially pursue resolveable uncertainty. The second term is the *divergence* term which is the KL divergence between the states expected under a given policy (also known as the posterior predictive density $q(x_\tau | \pi)$) and the goal distribution of the generative model, $\tilde{p}(x)$. This term scores how far away (in an informational sense) the agent expects it will be from the goal, if it were to pursue that policy. Since we are minimizing the expected free energy this term is positive and so is also minimized -- that is by minimizing the expected free energy, we are trying to choose trajectories which will bring the expected states in the future close to the desired states.

Now that we have the expected free energy to score possible trajectories, we now need to infer the optimal policy. A simple approach (which can be derived explicitly although it is somewhat complex) is to say that the posterior probability of a policy is proportional to the (exponentiated) sum of the expected free energy accumulated along the trajectory through the environment created by that policy. While this sounds complex, mathematically we can express it very simply as,
$$
\begin{align}
q(\pi) = \sigma( \sum_{\tau}^T \mathcal{G}(\pi)_{\tau})
\end{align}
$$
<!-- We can then choose which policy we implement for the next timestep by just sampling a policy from $q(\pi)$ and then emitting the first action of that policy. -->
We can then choose which action we choose for the next timestep by computing the marginal probability of each action $P(a)$, given the posterior over actions $q(\pi)$ and some mapping between policies and actions $P(a | \pi)$:

$$
\begin{align}
P(a) = \sum_{\pi}P(a | \pi)q(\pi)
\end{align}
$$

We can sample from this distribution over actions to choose our action for the next timestep.

In our simple grid world environment, policies are identical to actions since we are only planning ahead 1-step in the future, so $\forall_{i} P(a = i | \pi = i) = 1$, which implies $P(a) == q(\pi)$.

While all this may seem exceptionally long and complex, it results in an algorithm which is actually remarkably simple. The algorithm is:

1.) There is an agent with a generative model of the environment ($\textbf{A}$ and $\textbf{B}$ matrices), some initial set of approximate posterior beliefs $q(x_t | o_t)$ and a desired state vector $\textbf{C}$.

2.) The agent receives an observation $o_t$ and computes its posterior beliefs as we did earlier by minimizing the free energy.

3.) The agent now needs to choose what action to make to achieve its goals. It does this by:

 3.1.) First creating a set of potential policies to evaluate.

   3.2.) For each policy in this set, use the generative model to simulate the agent's trajectory in the environment *as if* it had emitted the actions prescribed by the policy

   3.3.) For each future timestep of each future trajectory, compute the expected free energy of that time-step

   3.4.) Sum the expected free energies for each timestep of each trajectory to get a total expected free energy for each possible policy.

   3.5.) Use these total expected free energies to compute the posterior distribution $q(\pi)$ as done above.

   3.6 ) Compute the marginal probability of each action, expected under the policies that include them, and then sample from this "action marginal" to generate an action at the current timestep.

4.) Execute the sampled action to 'step forward' the environmental dynamics and get a new observation. Go back to step 1.

And that's it! We're done. We have the full algorithm to create an active inference agent. Now all we do is show how to translate this algorithm into code for our specific case.

## Beliefs

First we need to setup an initial belief distribution which we will then update according to the observations we will receive.

```python
# setup initial prior beliefs -- uncertain -- completely unknown which state it is in
Qs = np.ones(9) * 1/9
plot_beliefs(Qs)
```

# Preferences

Now we have to encode the agent's preferences, so that it can learn to go to its reward state. In the current context, the agent wants (i.e. expects) to be in the reward location 7.

```python
# C matrix -- desires

REWARD_LOCATION = 7
reward_state = state_mapping[REWARD_LOCATION]
print(reward_state)

C = np.zeros(num_states)
C[REWARD_LOCATION] = 1.
print(C)
plot_beliefs(C)
```

The C matrix is a 1x9 matrix, where each value represents the preference to occupy a given state. We will create a one-hot C matrix, so that the agent only has a preference to be in state 7.

# Implementing the Active Inference Agent

## Evaluate policy

This helper function we evaluate the negative expected free energy for a given policy $-\mathcal{G}(\pi)$. To do this we need to calculate the cumulative expected free energy for that policy. All that entails is looping through the timesteps of the policy, simulate what the environment would do, using our generative model, evolving posterior beliefs and the actions entailed by the policy, and then compute the expected free energy of our (policy-dependent) expectations about the hidden states and observations.

```python
def evaluate_policy(policy, Qs, A, B, C):
    # initialize expected free energy at 0
    G = 0

    # loop over policy
    for t in range(len(policy)):

        # get action entailed by the policy at timestep `t`
        u = int(policy[t])

        # work out expected state, given the action
        Qs_pi = B[:,:,u].dot(Qs)

        # work out expected observations, given the action
        Qo_pi = A.dot(Qs_pi)

        # get entropy
        H = - (A * log_stable(A)).sum(axis = 0)

        # get predicted divergence
        # divergence = np.sum(Qo_pi * (log_stable(Qo_pi) - log_stable(C)), axis=0)
        divergence = KL_divergence(Qo_pi, C)

        # compute the expected uncertainty or ambiguity
        uncertainty = H.dot(Qs_pi)

        # increment the expected free energy counter for the policy, using the expected free energy at this timestep
        G += (divergence + uncertainty)

    return -G
```

## Infer action

This helper function will infer the most likely action. Specifically, it computes steps 3.1 to 3.5 in the active inference algorithm. First, it constructs all possible policies for a given policy length and set of actions. Then it loops through every possible policy and computes the expected free energy of that policy using our previous function, and then computing the policy distribution $q(\pi)$ using the softmax over the expected free energies.

```python
def infer_action(Qs, A, B, C, n_actions, policies):

    # initialize the negative expected free energy
    neg_G = np.zeros(len(policies))

    # loop over every possible policy and compute the EFE of each policy
    for i, policy in enumerate(policies):
        neg_G[i] = evaluate_policy(policy, Qs, A, B, C)

    # get distribution over policies
    Q_pi = maths.softmax(neg_G)

    # initialize probabilites of control states (convert from policies to actions)
    Qu = np.zeros(n_actions)

    # sum probabilites of control states or actions
    for i, policy in enumerate(policies):
        # control state specified by policy
        u = int(policy[0])
        # add probability of policy
        Qu[u] += Q_pi[i]

    # normalize action marginal
    utils.norm_dist(Qu)

    # sample control from action marginal
    u = utils.sample(Qu)

    return u
```

## Main loop

Here we simply implement the main loop of the active inference agent interacting with the environment. Specifically, this essentially implements steps 1-5 of the MDP "program" we discussed in notebook 1. Specifically, for each timestep, the agent infers an action, it emits that action to the environment, the environment is updated and returns an observation to the agent. The agent then infers the new state of the environment given that observation.

```python
# number of time steps
T = 10

#n_actions = env.n_control
n_actions = 5

# length of policies we consider
policy_len = 4

# this function generates all possible combinations of policies
policies = control.construct_policies([B.shape[0]], [n_actions], policy_len)

# reset environment
o = env.reset()

# loop over time
for t in range(T):

    # infer which action to take
    a = infer_action(Qs, A, B, C, n_actions, policies)

    # perform action in the environment and update the environment
    o = env.step(int(a))

    # infer new hidden state (this is the same equation as above but with PyMDP functions)
    likelihood = A[o,:]
    prior = B[:,:,int(a)].dot(Qs)

    Qs = maths.softmax(log_stable(likelihood) + log_stable(prior))

    print(Qs.round(3))
    plot_beliefs(Qs, "Beliefs (Qs) at time {}".format(t))
```

And that's it! In the last two notebooks, we have implemented a basic active inference agent which can successfully navigate around a 3x3 gridworld using active inference. Moreover, we have implemented this all from scratch using mostly basic numpy functions and not using much of the functionality of `pymdp`.

Hopefully after going through this you now understand roughly what active inference is and how it works, as well as ideally have some intuitions about how inference as well as policy selection work "under the hood", as well as learnt a lot about Bayesian and specifically variational inference. In the next notebook, we will focus more on the `pymdp` library itself and demonstrate how `pymdp` provides a useful set of abstractions that allows us to easily create active inference agents, as well as perform inference and policy selection in considerably more complex environments than the one described here.

We will discuss the high level structure of the library and show how its possible to replicate these notebooks in a much smaller amount of code using the `pymdp` abstractions.


# Applied Stochastic Processes HW_04

**11510691 程远星$\DeclareMathOperator*{\argmin}{argmin}
\DeclareMathOperator*{\argmax}{argmax}
\newcommand{\using}[1]{\stackrel{\mathrm{#1}}{=}}
\newcommand{\ffrac}{\displaystyle \frac}
\newcommand{\space}{\text{ }}
\newcommand{\bspace}{\;\;\;\;}
\newcommand{\QQQ}{\boxed{?\:}}
\newcommand{\void}{\left.\right.}
\newcommand{\CB}[1]{\left\{ #1 \right\}}
\newcommand{\SB}[1]{\left[ #1 \right]}
\newcommand{\P}[1]{\left( #1 \right)}
\newcommand{\dd}{\mathrm{d}}
\newcommand{\Tran}[1]{{#1}^{\mathrm{T}}}
\newcommand{\d}[1]{\displaystyle{#1}}
\newcommand{\EE}[2][\,\!]{\mathbb{E}_{#1}\left[#2\right]}
\newcommand{\Var}[2][\,\!]{\mathrm{Var}_{#1}\left[#2\right]}
\newcommand{\Cov}[2][\,\!]{\mathrm{Cov}_{#1}\left(#2\right)}
\newcommand{\Corr}[2][\,\!]{\mathrm{Corr}_{#1}\left(#2\right)}
\newcommand{\I}[1]{\mathrm{I}\left( #1 \right)}
\newcommand{\N}[1]{\mathrm{N} \left( #1 \right)}
\newcommand{\ow}{\text{otherwise}}\void^\dagger$**

## Question 1

$$\begin{align}
\Cov{N,S} &= \EE{NS} -\EE{N}\EE{S}\\
&= \EE{\EE{NS\mid N}} - \EE{N}\EE{N}\EE{X_1} \\
&= \EE{N\EE{S \mid N}}- \EE{N}\EE{N}\EE{X_1} \\
&= \EE{N^2 \EE{X_1}}- \EE{N}\EE{N}\EE{X_1} \\
&\using{\text{independency}}\P{\EE{N^2} - \SB{\EE{N}}^2}\EE{X_1}\\
&= \EE{X_1}\Var{N}
\end{align}$$

## Question 2

$\P{1}$

$\bspace P\CB{N = k+r}$ means that during the $k+r$ times of trial, the last $k$ trials are all successful while the $k+1$-th from the last is failed. Besides, the first $N-1$ trials cannot complete the experiment. Combine all conditions above we can conclude that

$$P\CB{N = k+r} = P\CB{N > r-1} q p^k,r>0$$

$\P{2}$

$\bspace$The mean value of $N$ is given in class, which is also in one example from the textbook that

$$M_k = \ffrac{1} {p} + \ffrac{1}{p}+\cdots + \ffrac{1} {p^k}$$

$\bspace$Here's how to prove this:

>Let $N_k$ denote the number of necessary trials to obtain $k$ consecutive success and $M_k = \EE{N_k}$. Then we might write: $N_k = N_{k-1} + \texttt{something}$... And actually that something is the number of additional trials needed to go from have $k-1$ successes in a row to having $k$ in a row. We denote that by $A_{k-1,k}$. Then we take the expectation and obtain:
>
>$\bspace M_k = M_{k-1} + \EE{A_{k-1,k}}$. Wanna know what's $\EE{A_{k-1,k}}$? There're only two possible result, the $k\texttt{-th}$ trial is a success, or not, then:
>
>$\bspace \EE{A_{k-1,k}} = 1 \cdot p + \P{1+M_k}\cdot\P{1-p} = 1+\P{1-p}M_k \Rightarrow M_k = \ffrac{1} {p} + \ffrac{M_{k-1}} {p}$
>
>Well, what's $M_1$? Obviously $N_1$ is a geometric with parameter $p$, thus $M_1 = 1/p$ and recursively we have:
>
>$\bspace M_k = \ffrac{1} {p} + \ffrac{1} {p^2} + \cdots \ffrac{1} {p^k}$

$\bspace$After this the equation to be proven is quite simple now,

$$\begin{align}
qp^k\EE{N} &= \P{1-p}p^k\P{\ffrac{1} {p} + \ffrac{1} {p^2} + \cdots \ffrac{1} {p^k}} \\
&= \P{1-p}\P{1+p+p^2+\cdots+p^{k-1}}\\
&= 1-p^k
\end{align}$$

## Question 3

$$\begin{array}{rll}
& P\CB{X_{n+1} = j\mid X_n =i,X_{n-1} = i_{n-1},\cdots, X_0 = i_0} = &\!\!\!\!\! P\CB{X_{n+1}= j \mid X_n = i} \\[0.7em]
\Leftrightarrow\!\!\! & \ffrac{P\CB{X_{n+1} = j, X_n =i,X_{n-1} = i_{n-1},\cdots, X_0 = i_0}} {P\CB{X_n =i,X_{n-1} = i_{n-1},\cdots, X_0 = i_0}} = &\!\!\!\!\! \ffrac{P\CB{X_{n+1} = j,X_{n} = i}} {P\CB{X_n = i}} \\[0.7em]
\Leftrightarrow\!\!\! & \ffrac{P\CB{X_{n+1} = j, X_n =i,X_{n-1} = i_{n-1},\cdots, X_0 = i_0}} {P\CB{X_n = i}} = &\!\!\!\!\! \ffrac{P\CB{X_{n+1} = j,X_{n} = i}} {P\CB{X_n = i}} \\
&&\!\!\!\!\!\bspace\times \ffrac{P\CB{X_n =i,X_{n-1} = i_{n-1},\cdots, X_0 = i_0}} {P\CB{X_n = i}}\\[0.7em]
\Leftrightarrow\!\!\! & P\CB{X_{n+1} = j ,X_{n-1} = i_{n-1},\cdots, X_0 = i_0\mid X_n =i} =&\!\!\!\!\! P\CB{X_{n+1}= j \mid X_n = i}\cdot P\CB{X_{n-1} = i_{n-1},\cdots, X_0 = i_0\mid X_n =i}
\end{array}$$

## Question 4

$$
\begin{array}{rlrl}
P_{00} &= P\CB{X_{n+1} = 0\mid X_n = 0} & P_{11} \!\!\!\!&= P\CB{X_{n+1} = 1\mid X_n = 1}\\
\!\!\!\!&= P\CB{\max\P{X_n,\xi_{n+1}} = 0\mid X_n = 0} & \!\!\!\!&= P\CB{\max\P{X_n,\xi_{n+1}} = 1\mid X_n = 1}\\
\!\!\!\!&= P\CB{\max\P{0,\xi_{n+1}} = 0\mid X_n = 0}  & \!\!\!\!&= P\CB{\max\P{1,\xi_{n+1}} = 1\mid X_n = 1} \\
\!\!\!\!&= P\CB{\max\P{0,\xi_{n+1}} = 0} & \!\!\!\!&= P\CB{\max\P{1,\xi_{n+1}} = 1}\\
\!\!\!\!&= P\CB{\xi_{n+1} \leq 0} & \!\!\!\!&= P\CB{\xi_{n+1} \leq 1} = 0.3 + 0.4 \\
\!\!\!\!&= 0.3 & \!\!\!\!&= 0.7\\
P_{22} &= P\CB{X_{n+1} = 2\mid X_n = 2} & P_{12} \!\!\!\!&= P\CB{X_{n+1} = 2\mid X_n = 1}\\
\!\!\!\!&= P\CB{\max\P{X_n,\xi_{n+1}} = 2\mid X_n = 2} & \!\!\!\!&= P\CB{\max\P{X_n,\xi_{n+1}} = 2\mid X_n = 1}\\
\!\!\!\!&= P\CB{\max\P{2,\xi_{n+1}} = 2\mid X_n = 2}  & \!\!\!\!&= P\CB{\max\P{1,\xi_{n+1}} = 2\mid X_n = 1} \\
\!\!\!\!&= P\CB{\max\P{2,\xi_{n+1}} = 2} & \!\!\!\!&= P\CB{\max\P{1,\xi_{n+1}} = 2}\\
\!\!\!\!&= P\CB{\xi_{n+1} \leq 2} = 0.3 + 0.4 + 0.3 & \!\!\!\!&= P\CB{\xi_{n+1} = 2}\\
\!\!\!\!&= 1 & \!\!\!\!&= 0.3\\
P_{02} &= P\CB{X_{n+1} = 2\mid X_n = 0} & P_{01} \!\!\!\!&= P\CB{X_{n+1} = 1\mid X_n = 0}\\
\!\!\!\!&= P\CB{\max\P{X_n,\xi_{n+1}} = 2\mid X_n = 0} & \!\!\!\!&= P\CB{\max\P{X_n,\xi_{n+1}} = 1\mid X_n = 0}\\
\!\!\!\!&= P\CB{\max\P{0,\xi_{n+1}} = 2\mid X_n = 0}  & \!\!\!\!&= P\CB{\max\P{0,\xi_{n+1}} = 1\mid X_n = 0} \\
\!\!\!\!&= P\CB{\max\P{0,\xi_{n+1}} = 2} & &= P\CB{\max\P{0,\xi_{n+1}} = 1}\\
\!\!\!\!&= P\CB{\xi_{n+1} = 2} & \!\!\!\!&= P\CB{\xi_{n+1} = 1}\\
\!\!\!\!&= 0.3 & \!\!\!\!&= 0.4\\
\end{array}$$

$$\mathbf{P} = \begin{Vmatrix}
0.3 & 0.4 & 0.3\\
0  & 0.7 & 0.3 \\
0  & 0   & 1
\end{Vmatrix}$$

## Question 5

$\P{1}$

$$\begin{align}
P\CB{X_2 = 1, X_3 = 1\mid X_1 = 0} &= P\CB{X_2 = 1\mid X_1 = 0}\cdot P\CB{X_3 = 1\mid X_2 = 1, X_1 = 0}\\
&= P\CB{X_2 = 1\mid X_1 = 0}\cdot P\CB{X_3 = 1\mid X_2 = 1}\\
&= P_{01} \times P_{11} = 0.12
\end{align}$$

$\P{2}$

$$\begin{align}
P\CB{X_1 = 1, X_2 = 1\mid X_0 = 0} &= P\CB{X_1 = 1\mid X_0 = 0}\cdot P\CB{X_2 = 1\mid X_1 = 1, X_0 = 0}\\
&= P_{01} \times P_{11} = 0.12
\end{align}$$

## Question 6

$\P{1}$

$\bspace$To apply the markov chain, we need $2^3 = 8$ states. Each is a triad containing all information of the past $3$ days. If to rain is to success and denote that as $1$, all $8$ possible states then are $\P{0,0,0}$, $\P{0,0,1}$, $\P{0,1,0}$, $\P{1,0,0}$, $\P{0,1,1}$, $\P{1,0,1}$, $\P{1,1,0}$, $\P{1,1,1}$.

$\P{2}$

$$\mathbf{P} = \begin{Vmatrix}
0.8 & 0.2 & 0 & 0 & 0 & 0 & 0 & 0\\
0 & 0 & 0.4 & 0.6 & 0 & 0 & 0 & 0\\
0 & 0 & 0 & 0 & 0.6 & 0.4 & 0 & 0\\
0 & 0 & 0 & 0 & 0 & 0 & 0.4 & 0.6\\
0.6 & 0.4 & 0 & 0 & 0 & 0 & 0 & 0\\
0 & 0 & 0.4 & 0.6 & 0 & 0 & 0 & 0\\
0 & 0 & 0 & 0 & 0.6 & 0.4 & 0 & 0\\
0 & 0 & 0 & 0 & 0 & 0 & 0.2 & 0.8
\end{Vmatrix}$$

Probability and Statistics/Applied Random Process/HW/HW_04.ipynb


Trusted Notebook" align="middle">

# _*Qiskit Aqua: Generating Random Variates*_

The latest version of this notebook is available on https://github.com/Qiskit/qiskit-tutorials.

***
### Contributors
Albert Akhriev[1], Jakub Marecek[1]

### Affliation
- [1]IBMQ

## Introduction

While classical computers use only pseudo-random routines, quantum computers
can generate true random variates.
For example, the measurement of a quantum superposition is intrinsically random,
as suggested by Born's rule.
Consequently, some of the
best random-number generators are based on such quantum-mechanical effects. (See the
Further, with a logarithmic amount of random bits, quantum computers can produce
linearly many more bits, which is known as
randomness expansion protocols.

In practical applications, one wishes to use random variates of well-known
distributions, rather than random bits.
In this notebook, we illustrate ways of generating random variates of several popular
distributions on IBM Q.

## Random Bits and the Bernoulli distribution

It is clear that there are many options for generating random bits (i.e., Bernoulli-distributed scalars, taking values either 0 or 1). Starting from a simple circuit such as a Hadamard gate followed by measurement, one can progress to vectors of Bernoulli-distributed elements. By addition of such random variates, we could get binomial distributions. By multiplication we could get geometric distributions, although perhaps leading to a circuit depth that may be impratical at the moment, though.

Let us start by importing the basic modules and creating a quantum circuit for generating random bits:

```python
import matplotlib.pyplot as plt
%matplotlib inline
import numpy as np
import sys, math, time
import warnings
warnings.filterwarnings("ignore", category=DeprecationWarning)

from qiskit import BasicAer
from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister, execute

# In this example we use 'qasm_simulator' backend.
glo_backend = BasicAer.get_backend("qasm_simulator")
```

In the next step we create a quantum circuit, which will be used for generation:

```python
# Number of qubits utilised simultaneously.
glo_num_qubits = 5

def create_circuit(num_target_qubits: int) -> QuantumCircuit:
    """
    Creates and returns quantum circuit for random variate generation.
    :param num_target_qubits: number of qubits to be used.
    :return: quantum curcuit.
    """
    assert isinstance(num_target_qubits, int) and num_target_qubits > 0
    q = QuantumRegister(num_target_qubits)
    c = ClassicalRegister(num_target_qubits)
    circuit = QuantumCircuit(q, c)
    circuit.h(q)
    circuit.barrier()
    circuit.measure(q, c)
    return circuit

# Create and plot generating quantum circuit.
circuit = create_circuit(glo_num_qubits)
#print(circuit)
circuit.draw(output='mpl')
```

## Uniformly-distributed scalars and vectors

It is clear that there are many options for approximating uniformly-distributed scalars by the choice of an integer from a finite range uniformly at random, e.g., by a binary-code construction from the Bernoulli-distributed vectors. In the following snippet, we generate random bits, which we then convert using the binary-code construction, up to the machine precision of a classical computer.

```python
def uniform_rand_float64(circuit: QuantumCircuit, num_target_qubits: int,
                         size: int, vmin: float, vmax: float) -> np.ndarray:
    """
    Generates a vector of random float64 values in the range [vmin, vmax].
    :param circuit: quantum circuit for random variate generation.
    :param num_target_qubits: number of qubits to be used.
    :param size: length of the vector.
    :param vmin: lower bound.
    :param vmax: upper bound.
    :return: vector of random values.
    """
    assert sys.maxsize == np.iinfo(np.int64).max    # sizeof(int) == 64 bits
    assert isinstance(size, int) and size > 0
    assert isinstance(vmin, float) and isinstance(vmax, float) and vmin <= vmax
    nbits = 7 * 8                                   # nbits > mantissa of float64
    bit_str_len = (nbits * size + num_target_qubits - 1) // num_target_qubits
    job = execute(circuit, glo_backend, shots=bit_str_len, memory=True)
    bit_str = ''.join(job.result().get_memory())
    scale = float(vmax - vmin) / float(2**nbits - 1)
    return np.array([vmin + scale * float(int(bit_str[i:i+nbits], 2))
                     for i in range(0, nbits * size, nbits)], dtype=np.float64)

def uniform_rand_int64(circuit: QuantumCircuit, num_target_qubits: int,
                       size: int, vmin: int, vmax: int) -> np.ndarray:
    """
    Generates a vector of random int64 values in the range [vmin, vmax].
    :param circuit: quantum circuit for random variate generation.
    :param num_target_qubits: number of qubits to be used.
    :param size: length of the vector.
    :param vmin: lower bound.
    :param vmax: upper bound.
    :return: vector of random values.
    """
    assert sys.maxsize == np.iinfo(np.int64).max        # sizeof(int) == 64 bits
    assert isinstance(size, int) and size > 0
    assert isinstance(vmin, int) and isinstance(vmax, int) and vmin <= vmax
    assert abs(vmin) <= 2**52 and abs(vmax) <= 2**52    # 52 == mantissa of float64
    return np.rint(uniform_rand_float64(circuit, num_target_qubits,
                                        size, float(vmin), float(vmax))).astype(np.int64)
```

### Uniform distribution over floating point numbers.
In this example we draw a random vector of floating-point values uniformly distributed within some arbitrary selected interval:

```python
# Draw a sample from uniform distribution.
start_time = time.time()
sample = uniform_rand_float64(circuit, glo_num_qubits, size=54321, vmin=-7.67, vmax=19.52)
sampling_time = time.time() - start_time

# Print out some details.
print("Uniform distribution over floating point numbers:")
print("  sample type:", type(sample), ", element type:", sample.dtype, ", shape:", sample.shape)
print("  sample min: {:.4f}, max: {:.4f}".format(np.amin(sample), np.amax(sample)))
print("  sampling time: {:.2f} secs".format(sampling_time))

# Plotting the distribution.
plt.hist(sample.ravel(),
         bins=min(int(np.ceil(np.sqrt(sample.size))), 100),
         density=True, facecolor='b', alpha=0.75)
plt.xlabel("value", size=12)
plt.ylabel("probability", size=12)
plt.title("Uniform distribution over float64 numbers in [{:.2f} ... {:.2f}]".format(
            np.amin(sample), np.amax(sample)), size=12)
plt.grid(True)
# plt.savefig("uniform_distrib_float.png", bbox_inches="tight")
plt.show()
```

### Uniform distribution over integers.
Our next example is similar to the previous one, but here we generate a random vector of integers:

```python
# Draw a sample from uniform distribution.
start_time = time.time()
sample = uniform_rand_int64(circuit, glo_num_qubits, size=54321, vmin=37, vmax=841)
sampling_time = time.time() - start_time

# Print out some details.
print("Uniform distribution over bounded integer numbers:")
print("  sample type:", type(sample), ", element type:", sample.dtype, ", shape:", sample.shape)
print("  sample min: {:d}, max: {:d}".format(np.amin(sample), np.amax(sample)))
print("  sampling time: {:.2f} secs".format(sampling_time))

# Plotting the distribution.
plt.hist(sample.ravel(),
         bins=min(int(np.ceil(np.sqrt(sample.size))), 100),
         density=True, facecolor='g', alpha=0.75)
plt.xlabel("value", size=12)
plt.ylabel("probability", size=12)
plt.title("Uniform distribution over int64 numbers in [{:d} ... {:d}]".format(
            np.amin(sample), np.amax(sample)), size=12)
plt.grid(True)
# plt.savefig("uniform_distrib_int.png", bbox_inches="tight")
plt.show()
```

## Normal distribution

To generate random variates with a standard normal distribution using two independent
samples $u_1, u_2$ of the uniform distribution on the unit interval [0, 1], one can
consider the Box-Muller transform to obtain a 2-vector:

\begin{align}
\begin{bmatrix}
%R\cos(\Theta )=
{\sqrt {-2\ln u_{1}}}\cos(2\pi u_{2}) \\
% R\sin(\Theta )=
{\sqrt {-2\ln u_{1}}}\sin(2\pi u_{2})
\end{bmatrix},
\end{align}

wherein we have two independent samples of the standard normal distribution.
In IBM Q, this is implemented as follows:

```python
def normal_rand_float64(circuit: QuantumCircuit, num_target_qubits: int,
                        size: int, mu: float, sigma: float) -> np.ndarray:
    """
    Draws a sample vector from the normal distribution given the mean and standard
    deviation, using the Box-Muller method.
    """
    TINY = np.sqrt(np.finfo(np.float64).tiny)
    assert isinstance(size, int) and size > 0
    rand_vec = np.zeros((size,), dtype=np.float64)

    # Generate array of uniformly distributed samples, factor 1.5 longer that
    # actually needed.
    n = (3 * size) // 2
    x = np.reshape(uniform_rand_float64(circuit, num_target_qubits,
                                        2*n, 0.0, 1.0), (-1, 2))

    x1 = 0.0                # first sample in a pair
    c = 0                   # counter
    for d in range(size):
        r2 = 2.0
        while r2 >= 1.0 or r2 < TINY:
            # Regenerate array of uniformly distributed samples upon shortage.
            if c >= n:
                c = 0
                n = max(size // 10, 1)
                x = np.reshape(uniform_rand_float64(circuit, num_target_qubits,
                                                    2*n, 0.0, 1.0), (-1, 2))

            x1 = 2.0 * x[c, 0] - 1.0        # first sample in a pair
            x2 = 2.0 * x[c, 1] - 1.0        # second sample in a pair
            r2 = x1 * x1 + x2 * x2
            c += 1

        f = np.sqrt(np.abs(-2.0 * np.log(r2) / r2))
        rand_vec[d] = f * x1

    return (rand_vec * sigma + mu)
```

The following example demonstrates how to draw a random vector of normally distributed variates:

```python
# Mean and standard deviation.
mu = 2.4
sigma = 5.1

# Draw a sample from the normal distribution.
start_time = time.time()
sample = normal_rand_float64(circuit, glo_num_qubits, size=4321, mu=mu, sigma=sigma)
sampling_time = time.time() - start_time

# Print out some details.
print("Normal distribution (mu={:.3f}, sigma={:.3f}):".format(mu, sigma))
print("  sample type:", type(sample), ", element type:", sample.dtype, ", shape:", sample.shape)
print("  sample min: {:.4f}, max: {:.4f}".format(np.amin(sample), np.amax(sample)))
print("  sampling time: {:.2f} secs".format(sampling_time))

# Plotting the distribution.
x = np.linspace(mu - 4.0 * sigma, mu + 4.0 * sigma, 1000)
analyt = np.exp(-0.5 * ((x - mu) / sigma)**2) / (sigma * math.sqrt(2.0 * math.pi))
plt.hist(sample.ravel(),
         bins=min(int(np.ceil(np.sqrt(sample.size))), 100),
         density=True, facecolor='r', alpha=0.75)
plt.plot(x, analyt, '-b', lw=1)
plt.xlabel("value", size=12)
plt.ylabel("probability", size=12)
plt.title("Normal distribution: empirical vs analytic", size=12)
plt.grid(True)
# plt.savefig("normal_distrib.png", bbox_inches="tight")
plt.show()
```

There is a substantial amount of further work needed to either certify the quality of the source of random numbers (cf. NIST SP 800-90B, Recommendation for the Entropy Sources Used for Random Bit Generation) or to use random variates within quantum algorithms (cf. uncertainty_models within Qiskit Aqua).

```python
import qiskit.tools.jupyter
%qiskit_version_table
%qiskit_copyright
```

Version Information

Qiskit Software	Version

Qiskit	None

Terra	0.9.0

Aer	0.3.0

Ignis	0.2.0

Aqua	0.5.6

IBM Q Provider	0.3.2rc1

System information

Python	3.7.4 (default, Aug 13 2019, 15:17:50)
[Clang 4.0.1 (tags/RELEASE_401/final)]

OS	Darwin

CPUs	4

Memory (Gb)	16.0

Wed Aug 21 21:39:36 2019 EDT

This code is a part of Qiskit

© Copyright IBM 2017, 2019.

This code is licensed under the Apache License, Version 2.0. You may
obtain a copy of this license in the LICENSE.txt file in the root directory
 of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
Any modifications or derivative works of this code must retain this
copyright notice, and modified files need to carry a notice indicating
that they have been altered from the originals.

```python

```


[Table of Contents](./table_of_contents.ipynb)

# The Extended Kalman Filter

```python
from __future__ import division, print_function
%matplotlib inline
```

```python
#format the book
import book_format
book_format.set_style()
```

<style>
.output_wrapper, .output {
    height:auto !important;
    max-height:100000px;
}
.output_scroll {
    box-shadow:none !important;
    webkit-box-shadow:none !important;
}
</style>

We have developed the theory for the linear Kalman filter. Then, in the last two chapters we broached the topic of using Kalman filters for nonlinear problems. In this chapter we will learn the Extended Kalman filter (EKF). The EKF handles nonlinearity by linearizing the system at the point of the current estimate, and then the linear Kalman filter is used to filter this linearized system. It was one of the very first techniques used for nonlinear problems, and it remains the most common technique.

The EKF provides significant mathematical challenges to the designer of the filter; this is the most challenging chapter of the book. I do everything I can to avoid the EKF in favor of other techniques that have been developed to filter nonlinear problems. However, the topic is unavoidable; all classic papers and a majority of current papers in the field use the EKF. Even if you do not use the EKF in your own work you will need to be familiar with the topic to be able to read the literature.

## Linearizing the Kalman Filter

The Kalman filter uses linear equations, so it does not work with nonlinear problems. Problems can be nonlinear in two ways. First, the process model might be nonlinear. An object falling through the atmosphere encounters drag which reduces its acceleration. The drag coefficient varies based on the velocity the object. The resulting behavior is nonlinear - it cannot be modeled with linear equations. Second, the measurements could be nonlinear. For example, a radar gives a range and bearing to a target. We use trigonometry, which is nonlinear, to compute the position of the target.

For the linear filter we have these equations for the process and measurement models:

$$\begin{aligned}\dot{\mathbf x} &= \mathbf{Ax} + w_x\\
\mathbf z &= \mathbf{Hx} + w_z
\end{aligned}$$

Where $\mathbf A$ is the systems dynamic matrix. Using the state space methods covered in the **Kalman Filter Math** chapter these equations can be transformed into
$$\begin{aligned}\bar{\mathbf x} &= \mathbf{Fx} \\
\mathbf z &= \mathbf{Hx}
\end{aligned}$$

where $\mathbf F$ is the *fundamental matrix*. The noise $w_x$ and $w_z$ terms are incorporated into the matrices $\mathbf R$ and $\mathbf Q$. This form of the equations allow us to compute the state at step $k$ given a measurement at step $k$ and the state estimate at step $k-1$. In earlier chapters I built your intuition and minimized the math by using problems describable with Newton's equations. We know how to design $\mathbf F$ based on high school physics.

For the nonlinear model the linear expression $\mathbf{Fx} + \mathbf{Bu}$ is replaced by a nonlinear function $f(\mathbf x, \mathbf u)$, and the linear expression $\mathbf{Hx}$ is replaced by a nonlinear function $h(\mathbf x)$:

$$\begin{aligned}\dot{\mathbf x} &= f(\mathbf x, \mathbf u) + w_x\\
\mathbf z &= h(\mathbf x) + w_z
\end{aligned}$$

You might imagine that we could proceed by finding a new set of Kalman filter equations that optimally solve these equations. But if you remember the charts in the **Nonlinear Filtering** chapter you'll recall that passing a Gaussian through a nonlinear function results in a probability distribution that is no longer Gaussian. So this will not work.

The EKF does not alter the Kalman filter's linear equations. Instead, it *linearizes* the nonlinear equations at the point of the current estimate, and uses this linearization in the linear Kalman filter.

*Linearize* means what it sounds like. We find a line that most closely matches the curve at a defined point. The graph below linearizes the parabola $f(x)=x^2-2x$ at $x=1.5$.

```python
import kf_book.ekf_internal as ekf_internal
ekf_internal.show_linearization()
```

If the curve above is the process model, then the dotted lines shows the linearization of that curve for the estimate $x=1.5$.

We linearize systems by taking the derivative, which finds the slope of a curve:

$$\begin{aligned}
f(x) &= x^2 -2x \\
\frac{df}{dx} &= 2x - 2
\end{aligned}$$

and then evaluating it at $x$:

$$\begin{aligned}m &= f'(x=1.5) \\&= 2(1.5) - 2 \\&= 1\end{aligned}$$

Linearizing systems of differential equations is similar. We linearize $f(\mathbf x, \mathbf u)$, and $h(\mathbf x)$ by taking the partial derivatives of each to evaluate $\mathbf F$ and $\mathbf H$ at the point $\mathbf x_t$ and $\mathbf u_t$. We call the partial derivative of a matrix the [*Jacobian*](https://en.wikipedia.org/wiki/Jacobian_matrix_and_determinant). This gives us the the discrete state transition matrix and measurement model matrix:

$$
\begin{aligned}
\mathbf F
&= {\frac{\partial{f(\mathbf x_t, \mathbf u_t)}}{\partial{\mathbf x}}}\biggr|_{{\mathbf x_t},{\mathbf u_t}} \\
\mathbf H &= \frac{\partial{h(\bar{\mathbf x}_t)}}{\partial{\bar{\mathbf x}}}\biggr|_{\bar{\mathbf x}_t}
\end{aligned}
$$

This leads to the following equations for the EKF. I put boxes around the differences from the linear filter:

$$\begin{array}{l|l}
\text{linear Kalman filter} & \text{EKF} \\
\hline
& \boxed{\mathbf F = {\frac{\partial{f(\mathbf x_t, \mathbf u_t)}}{\partial{\mathbf x}}}\biggr|_{{\mathbf x_t},{\mathbf u_t}}} \\
\mathbf{\bar x} = \mathbf{Fx} + \mathbf{Bu} & \boxed{\mathbf{\bar x} = f(\mathbf x, \mathbf u)}  \\
\mathbf{\bar P} = \mathbf{FPF}^\mathsf{T}+\mathbf Q  & \mathbf{\bar P} = \mathbf{FPF}^\mathsf{T}+\mathbf Q \\
\hline
& \boxed{\mathbf H = \frac{\partial{h(\bar{\mathbf x}_t)}}{\partial{\bar{\mathbf x}}}\biggr|_{\bar{\mathbf x}_t}} \\
\textbf{y} = \mathbf z - \mathbf{H \bar{x}} & \textbf{y} = \mathbf z - \boxed{h(\bar{x})}\\
\mathbf{K} = \mathbf{\bar{P}H}^\mathsf{T} (\mathbf{H\bar{P}H}^\mathsf{T} + \mathbf R)^{-1} & \mathbf{K} = \mathbf{\bar{P}H}^\mathsf{T} (\mathbf{H\bar{P}H}^\mathsf{T} + \mathbf R)^{-1} \\
\mathbf x=\mathbf{\bar{x}} +\mathbf{K\textbf{y}} & \mathbf x=\mathbf{\bar{x}} +\mathbf{K\textbf{y}} \\
\mathbf P= (\mathbf{I}-\mathbf{KH})\mathbf{\bar{P}} & \mathbf P= (\mathbf{I}-\mathbf{KH})\mathbf{\bar{P}}
\end{array}$$

We don't normally use $\mathbf{Fx}$ to propagate the state for the EKF as the linearization causes inaccuracies. It is typical to compute $\bar{\mathbf x}$ using a suitable numerical integration technique such as Euler or Runge Kutta. Thus I wrote $\mathbf{\bar x} = f(\mathbf x, \mathbf u)$. For the same reasons we don't use $\mathbf{H\bar{x}}$ in the computation for the residual, opting for the more accurate $h(\bar{\mathbf x})$.

I think the easiest way to understand the EKF is to start off with an example. Later you may want to come back and reread this section.

## Example: Tracking a Airplane

This example tracks an airplane using ground based radar. We implemented a UKF for this problem in the last chapter. Now we will implement an EKF for the same problem so we can compare both the filter performance and the level of effort required to implement the filter.

Radars work by emitting a beam of radio waves and scanning for a return bounce. Anything in the beam's path will reflects some of the signal back to the radar. By timing how long it takes for the reflected signal to get back to the radar the system can compute the *slant distance* - the straight line distance from the radar installation to the object.

The relationship between the radar's slant range distance $r$ and elevation angle $\epsilon$ with the horizontal position $x$ and altitude $y$ of the aircraft is illustrated in the figure below:

```python
ekf_internal.show_radar_chart()
```

This gives us the equalities:

$$\begin{aligned}
\epsilon &= \tan^{-1} \frac y x\\
r^2 &= x^2 + y^2
\end{aligned}$$

### Design the State Variables

We want to track the position of an aircraft assuming a constant velocity and altitude, and measurements of the slant distance to the aircraft. That means we need 3 state variables - horizontal distance, horizonal velocity, and altitude:

$$\mathbf x = \begin{bmatrix}\mathtt{distance} \\\mathtt{velocity}\\ \mathtt{altitude}\end{bmatrix}=    \begin{bmatrix}x \\ \dot x\\ y\end{bmatrix}$$

### Design the Process Model

We assume a Newtonian, kinematic system for the aircraft. We've used this model in previous chapters, so by inspection you may recognize that we want

$$\mathbf F = \left[\begin{array}{cc|c} 1 & \Delta t & 0\\
0 & 1 & 0 \\ \hline
0 & 0 & 1\end{array}\right]$$

I've partioned the matrix into blocks to show the upper left block is a constant velocity model for $x$, and the lower right block is a constant position model for $y$.

However, let's practice finding these matrices. We model systems with a set of differential equations. We need an equation in the form

$$\dot{\mathbf x} = \mathbf{Ax} + \mathbf{w}$$
where $\mathbf{w}$ is the system noise.

The variables $x$ and $y$ are independent so we can compute them separately. The differential equations for motion in one dimension are:

$$\begin{aligned}v &= \dot x \\
a &= \ddot{x} = 0\end{aligned}$$

Now we put the differential equations into state-space form. If this was a second or greater order differential system we would have to first reduce them to an equivalent set of first degree equations. The equations are first order, so we put them in state space matrix form as

$$\begin{aligned}\begin{bmatrix}\dot x \\ \ddot{x}\end{bmatrix} &= \begin{bmatrix}0&1\\0&0\end{bmatrix} \begin{bmatrix}x \\
\dot x\end{bmatrix} \\ \dot{\mathbf x} &= \mathbf{Ax}\end{aligned}$$
where $\mathbf A=\begin{bmatrix}0&1\\0&0\end{bmatrix}$.

Recall that $\mathbf A$ is the *system dynamics matrix*. It describes a set of linear differential equations. From it we must compute the state transition matrix $\mathbf F$. $\mathbf F$ describes a discrete set of linear equations which compute $\mathbf x$ for a discrete time step $\Delta t$.

A common way to compute $\mathbf F$ is to use the power series expansion of the matrix exponential:

$$\mathbf F(\Delta t) = e^{\mathbf A\Delta t} = \mathbf{I} + \mathbf A\Delta t  + \frac{(\mathbf A\Delta t)^2}{2!} + \frac{(\mathbf A \Delta t)^3}{3!} + ... $$

$\mathbf A^2 = \begin{bmatrix}0&0\\0&0\end{bmatrix}$, so all higher powers of $\mathbf A$ are also $\mathbf{0}$. Thus the power series expansion is:

$$
\begin{aligned}
\mathbf F &=\mathbf{I} + \mathbf At + \mathbf{0} \\
&= \begin{bmatrix}1&0\\0&1\end{bmatrix} + \begin{bmatrix}0&1\\0&0\end{bmatrix}\Delta t\\
\mathbf F &= \begin{bmatrix}1&\Delta t\\0&1\end{bmatrix}
\end{aligned}$$

This is the same result used by the kinematic equations! This exercise was unnecessary other than to illustrate finding the state transition matrix from linear differential equations. We will conclude the chapter with an example that will require the use of this technique.

### Design the Measurement Model

The measurement function takes the state estimate of the prior $\bar{\mathbf x}$ and turn it into a measurement of the slant range distance. We use the Pythagorean theorem to derive:

$$h(\bar{\mathbf x}) = \sqrt{x^2 + y^2}$$

The relationship between the slant distance and the position on the ground is nonlinear due to the square root. We linearize it by evaluating its partial derivative at $\mathbf x_t$:

$$
\mathbf H = \frac{\partial{h(\bar{\mathbf x})}}{\partial{\bar{\mathbf x}}}\biggr|_{\bar{\mathbf x}_t}
$$

The partial derivative of a matrix is called a Jacobian, and takes the form

$$\frac{\partial \mathbf H}{\partial \bar{\mathbf x}} =
\begin{bmatrix}
\frac{\partial h_1}{\partial x_1} & \frac{\partial h_1}{\partial x_2} &\dots \\
\frac{\partial h_2}{\partial x_1} & \frac{\partial h_2}{\partial x_2} &\dots \\
\vdots & \vdots
\end{bmatrix}
$$

In other words, each element in the matrix is the partial derivative of the function $h$ with respect to the $x$ variables. For our problem we have

$$\mathbf H = \begin{bmatrix}{\partial h}/{\partial x} & {\partial h}/{\partial \dot{x}} & {\partial h}/{\partial y}\end{bmatrix}$$

Solving each in turn:

$$\begin{aligned}
\frac{\partial h}{\partial x} &= \frac{\partial}{\partial x} \sqrt{x^2 + y^2} \\
&= \frac{x}{\sqrt{x^2 + y^2}}
\end{aligned}$$

and

$$\begin{aligned}
\frac{\partial h}{\partial \dot{x}} &=
\frac{\partial}{\partial \dot{x}} \sqrt{x^2 + y^2} \\
&= 0
\end{aligned}$$

and

$$\begin{aligned}
\frac{\partial h}{\partial y} &= \frac{\partial}{\partial y} \sqrt{x^2 + y^2} \\
&= \frac{y}{\sqrt{x^2 + y^2}}
\end{aligned}$$

giving us

$$\mathbf H =
\begin{bmatrix}
\frac{x}{\sqrt{x^2 + y^2}} &
0 &
&
\frac{y}{\sqrt{x^2 + y^2}}
\end{bmatrix}$$

This may seem daunting, so step back and recognize that all of this math is doing something very simple. We have an equation for the slant range to the airplane which is nonlinear. The Kalman filter only works with linear equations, so we need to find a linear equation that approximates $\mathbf H$. As we discussed above, finding the slope of a nonlinear equation at a given point is a good approximation. For the Kalman filter, the 'given point' is the state variable $\mathbf x$ so we need to take the derivative of the slant range with respect to $\mathbf x$. For the linear Kalman filter $\mathbf H$ was a constant that we computed prior to running the filter. For the EKF $\mathbf H$ is updated at each step as the evaluation point $\bar{\mathbf x}$ changes at each epoch.

To make this more concrete, let's now write a Python function that computes the Jacobian of $h$ for this problem.

```python
from math import sqrt
def HJacobian_at(x):
    """ compute Jacobian of H matrix at x """

    horiz_dist = x[0]
    altitude   = x[2]
    denom = sqrt(horiz_dist**2 + altitude**2)
    return array ([[horiz_dist/denom, 0., altitude/denom]])
```

Finally, let's provide the code for $h(\bar{\mathbf x})$:

```python
def hx(x):
    """ compute measurement for slant range that
    would correspond to state x.
    """

    return (x[0]**2 + x[2]**2) ** 0.5
```

Now let's write a simulation for our radar.

```python
from numpy.random import randn
import math

class RadarSim:
    """ Simulates the radar signal returns from an object
    flying at a constant altityude and velocity in 1D.
    """

    def __init__(self, dt, pos, vel, alt):
        self.pos = pos
        self.vel = vel
        self.alt = alt
        self.dt = dt

    def get_range(self):
        """ Returns slant range to the object. Call once
        for each new measurement at dt time from last call.
        """

        # add some process noise to the system
        self.vel = self.vel  + .1*randn()
        self.alt = self.alt + .1*randn()
        self.pos = self.pos + self.vel*self.dt

        # add measurement noise
        err = self.pos * 0.05*randn()
        slant_dist = math.sqrt(self.pos**2 + self.alt**2)

        return slant_dist + err
```

### Design Process and Measurement Noise

The radar measures the range to a target. We will use $\sigma_{range}= 5$ meters for the noise. This gives us

$$\mathbf R = \begin{bmatrix}\sigma_{range}^2\end{bmatrix} = \begin{bmatrix}25\end{bmatrix}$$

The design of $\mathbf Q$ requires some discussion. The state $\mathbf x= \begin{bmatrix}x & \dot x & y\end{bmatrix}^\mathtt{T}$. The first two elements are position (down range distance) and velocity, so we can use `Q_discrete_white_noise` noise to compute the values for the upper left hand side of $\mathbf Q$. The third element of  $\mathbf x$ is altitude, which we are assuming is independent of the down range distance. That leads us to a block design of $\mathbf Q$ of:

$$\mathbf Q = \begin{bmatrix}\mathbf Q_\mathtt{x} & 0 \\ 0 & \mathbf Q_\mathtt{y}\end{bmatrix}$$

### Implementation

`FilterPy` provides the class `ExtendedKalmanFilter`. It works similarly to the `KalmanFilter` class we have been using, except that it allows you to provide a function that computes the Jacobian of $\mathbf H$ and the function $h(\mathbf x)$.

We start by importing the filter and creating it. The dimension of `x` is 3 and `z` has dimension 1.

```python
from filterpy.kalman import ExtendedKalmanFilter

rk = ExtendedKalmanFilter(dim_x=3, dim_z=1)
```
We create the radar simulator:
```python
radar = RadarSim(dt, pos=0., vel=100., alt=1000.)
```
We will initialize the filter near the airplane's actual position:

```python
rk.x = array([radar.pos, radar.vel-10, radar.alt+100])
```

We assign the system matrix using the first term of the Taylor series expansion we computed above:

```python
dt = 0.05
rk.F = eye(3) + array([[0, 1, 0],
                       [0, 0, 0],
                       [0, 0, 0]])*dt
```

After assigning reasonable values to $\mathbf R$, $\mathbf Q$, and $\mathbf P$ we can run the filter with a simple loop. We pass the functions for computing the Jacobian of $\mathbf  H$ and $h(x)$ into the `update` method.

```python
for i in range(int(20/dt)):
    z = radar.get_range()
    rk.update(array([z]), HJacobian_at, hx)
    rk.predict()
```

Adding some boilerplate code to save and plot the results we get:

```python
from filterpy.common import Q_discrete_white_noise
from filterpy.kalman import ExtendedKalmanFilter
from numpy import eye, array, asarray
import numpy as np

dt = 0.05
rk = ExtendedKalmanFilter(dim_x=3, dim_z=1)
radar = RadarSim(dt, pos=0., vel=100., alt=1000.)

# make an imperfect starting guess
rk.x = array([radar.pos-100, radar.vel+100, radar.alt+1000])

rk.F = eye(3) + array([[0, 1, 0],
                       [0, 0, 0],
                       [0, 0, 0]]) * dt

range_std = 5. # meters
rk.R = np.diag([range_std**2])
rk.Q[0:2, 0:2] = Q_discrete_white_noise(2, dt=dt, var=0.1)
rk.Q[2,2] = 0.1
rk.P *= 50

xs, track = [], []
for i in range(int(20/dt)):
    z = radar.get_range()
    track.append((radar.pos, radar.vel, radar.alt))

    rk.update(array([z]), HJacobian_at, hx)
    xs.append(rk.x)
    rk.predict()

xs = asarray(xs)
track = asarray(track)
time = np.arange(0, len(xs)*dt, dt)
ekf_internal.plot_radar(xs, track, time)
```

## Using SymPy to compute Jacobians

Depending on your experience with derivatives you may have found the computation of the Jacobian difficult. Even if you found it easy, a slightly more difficult problem easily leads to very difficult computations.

As explained in Appendix A, we can use the SymPy package to compute the Jacobian for us.

```python
import sympy
from IPython.display import display
sympy.init_printing(use_latex='mathjax')

x, x_vel, y = sympy.symbols('x, x_vel y')

H = sympy.Matrix([sympy.sqrt(x**2 + y**2)])

state = sympy.Matrix([x, x_vel, y])
J = H.jacobian(state)

display(state)
display(J)
```

$\displaystyle \left[\begin{matrix}x\\x_{vel}\\y\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}\frac{x}{\sqrt{x^{2} + y^{2}}} & 0 & \frac{y}{\sqrt{x^{2} + y^{2}}}\end{matrix}\right]$

This result is the same as the result we computed above, and with much less effort on our part!

## Robot Localization

It's time to try a real problem. I warn you that this section is difficult. However, most books choose simple, textbook problems with simple answers, and you are left wondering how to solve a real world problem.

We will consider the problem of robot localization. We already implemented this in the **Unscented Kalman Filter** chapter, and I recommend you read it now if you haven't already.  In this scenario we have a robot that is moving through a landscape using a sensor to detect landmarks. This could be a self driving car using computer vision to identify trees, buildings, and other landmarks. It might be one of those small robots that vacuum your house, or a robot in a warehouse.

The robot has 4 wheels in the same configuration used by automobiles. It maneuvers by pivoting the front wheels. This causes the robot to pivot around the rear axle while moving forward. This is nonlinear behavior which we will have to model.

The robot has a sensor that measures the range and bearing to known targets in the landscape. This is nonlinear because computing a position from a range and bearing requires square roots and trigonometry.

Both the process model and measurement models are nonlinear. The EKF accommodates both, so we provisionally conclude that the EKF is a viable choice for this problem.

### Robot Motion Model

At a first approximation an automobile steers by pivoting the front tires while moving forward. The front of the car moves in the direction that the wheels are pointing while pivoting around the rear tires. This simple description is complicated by issues such as slippage due to friction, the differing behavior of the rubber tires at different speeds, and the need for the outside tire to travel a different radius than the inner tire. Accurately modeling steering requires a complicated set of differential equations.

For lower speed robotic applications a simpler *bicycle model* has been found to perform well. This is a depiction of the model:

```python
ekf_internal.plot_bicycle()
```

In the **Unscented Kalman Filter** chapter we derived these equations:

$$\begin{aligned}
\beta &= \frac d w \tan(\alpha) \\
x &= x - R\sin(\theta) + R\sin(\theta + \beta) \\
y &= y + R\cos(\theta) - R\cos(\theta + \beta) \\
\theta &= \theta + \beta
\end{aligned}
$$

where $\theta$ is the robot's heading.

You do not need to understand this model in detail if you are not interested in steering models. The important thing to recognize is that our motion model is nonlinear, and we will need to deal with that with our Kalman filter.

### Design the State Variables

For our filter we will maintain the position $x,y$ and orientation $\theta$ of the robot:

$$\mathbf x = \begin{bmatrix}x \\ y \\ \theta\end{bmatrix}$$

Our control input $\mathbf u$ is the velocity $v$ and steering angle $\alpha$:

$$\mathbf u = \begin{bmatrix}v \\ \alpha\end{bmatrix}$$

### Design the System Model

We model our system as a nonlinear motion model plus noise.

$$\bar x = f(x, u) + \mathcal{N}(0, Q)$$

Using the motion model for a robot that we created above, we can expand this to

$$\bar{\begin{bmatrix}x\\y\\\theta\end{bmatrix}} = \begin{bmatrix}x\\y\\\theta\end{bmatrix} +
\begin{bmatrix}- R\sin(\theta) + R\sin(\theta + \beta) \\
R\cos(\theta) - R\cos(\theta + \beta) \\
\beta\end{bmatrix}$$

We find The $\mathbf F$ by taking the Jacobian of $f(x,u)$.

$$\mathbf F = \frac{\partial f(x, u)}{\partial x} =\begin{bmatrix}
\frac{\partial f_1}{\partial x} &
\frac{\partial f_1}{\partial y} &
\frac{\partial f_1}{\partial \theta}\\
\frac{\partial f_2}{\partial x} &
\frac{\partial f_2}{\partial y} &
\frac{\partial f_2}{\partial \theta} \\
\frac{\partial f_3}{\partial x} &
\frac{\partial f_3}{\partial y} &
\frac{\partial f_3}{\partial \theta}
\end{bmatrix}
$$

When we calculate these we get

$$\mathbf F = \begin{bmatrix}
1 & 0 & -R\cos(\theta) + R\cos(\theta+\beta) \\
0 & 1 & -R\sin(\theta) + R\sin(\theta+\beta) \\
0 & 0 & 1
\end{bmatrix}$$

We can double check our work with SymPy.

```python
import sympy
from sympy.abc import alpha, x, y, v, w, R, theta
from sympy import symbols, Matrix
sympy.init_printing(use_latex="mathjax", fontsize='16pt')
time = symbols('t')
d = v*time
beta = (d/w)*sympy.tan(alpha)
r = w/sympy.tan(alpha)

fxu = Matrix([[x-r*sympy.sin(theta) + r*sympy.sin(theta+beta)],
              [y+r*sympy.cos(theta)- r*sympy.cos(theta+beta)],
              [theta+beta]])
F = fxu.jacobian(Matrix([x, y, theta]))
F
```

$\displaystyle \left[\begin{matrix}1 & 0 & - \frac{w \cos{\left(\theta \right)}}{\tan{\left(\alpha \right)}} + \frac{w \cos{\left(\frac{t v \tan{\left(\alpha \right)}}{w} + \theta \right)}}{\tan{\left(\alpha \right)}}\\0 & 1 & - \frac{w \sin{\left(\theta \right)}}{\tan{\left(\alpha \right)}} + \frac{w \sin{\left(\frac{t v \tan{\left(\alpha \right)}}{w} + \theta \right)}}{\tan{\left(\alpha \right)}}\\0 & 0 & 1\end{matrix}\right]$

That looks a bit complicated. We can use SymPy to substitute terms:

```python
# reduce common expressions
B, R = symbols('beta, R')
F = F.subs((d/w)*sympy.tan(alpha), B)
F.subs(w/sympy.tan(alpha), R)
```

$\displaystyle \left[\begin{matrix}1 & 0 & - R \cos{\left(\theta \right)} + R \cos{\left(\beta + \theta \right)}\\0 & 1 & - R \sin{\left(\theta \right)} + R \sin{\left(\beta + \theta \right)}\\0 & 0 & 1\end{matrix}\right]$

This form  verifies that the computation of the Jacobian is correct.

Now we can turn our attention to the noise. Here, the noise is in our control input, so it is in *control space*. In other words, we command a specific velocity and steering angle, but we need to convert that into errors in $x, y, \theta$. In a real system this might vary depending on velocity, so it will need to be recomputed for every prediction. I will choose this as the noise model; for a real robot you will need to choose a model that accurately depicts the error in your system.

$$\mathbf{M} = \begin{bmatrix}\sigma_{vel}^2 & 0 \\ 0 & \sigma_\alpha^2\end{bmatrix}$$

If this was a linear problem we would convert from control space to state space using the by now familiar $\mathbf{FMF}^\mathsf T$ form. Since our motion model is nonlinear we do not try to find a closed form solution to this, but instead linearize it with a Jacobian which we will name $\mathbf{V}$.

$$\mathbf{V} = \frac{\partial f(x, u)}{\partial u} \begin{bmatrix}
\frac{\partial f_1}{\partial v} & \frac{\partial f_1}{\partial \alpha} \\
\frac{\partial f_2}{\partial v} & \frac{\partial f_2}{\partial \alpha} \\
\frac{\partial f_3}{\partial v} & \frac{\partial f_3}{\partial \alpha}
\end{bmatrix}$$

These partial derivatives become very difficult to work with. Let's compute them with SymPy.

```python
V = fxu.jacobian(Matrix([v, alpha]))
V = V.subs(sympy.tan(alpha)/w, 1/R)
V = V.subs(time*v/R, B)
V = V.subs(time*v, 'd')
V
```

$\displaystyle \left[\begin{matrix}t \cos{\left(\beta + \theta \right)} & \frac{d \left(\tan^{2}{\left(\alpha \right)} + 1\right) \cos{\left(\beta + \theta \right)}}{\tan{\left(\alpha \right)}} - \frac{w \left(- \tan^{2}{\left(\alpha \right)} - 1\right) \sin{\left(\theta \right)}}{\tan^{2}{\left(\alpha \right)}} + \frac{w \left(- \tan^{2}{\left(\alpha \right)} - 1\right) \sin{\left(\beta + \theta \right)}}{\tan^{2}{\left(\alpha \right)}}\\t \sin{\left(\beta + \theta \right)} & \frac{d \left(\tan^{2}{\left(\alpha \right)} + 1\right) \sin{\left(\beta + \theta \right)}}{\tan{\left(\alpha \right)}} + \frac{w \left(- \tan^{2}{\left(\alpha \right)} - 1\right) \cos{\left(\theta \right)}}{\tan^{2}{\left(\alpha \right)}} - \frac{w \left(- \tan^{2}{\left(\alpha \right)} - 1\right) \cos{\left(\beta + \theta \right)}}{\tan^{2}{\left(\alpha \right)}}\\\frac{t}{R} & \frac{d \left(\tan^{2}{\left(\alpha \right)} + 1\right)}{w}\end{matrix}\right]$

This should give you an appreciation of how quickly the EKF become mathematically intractable.

This gives us the final form of our prediction equations:

$$\begin{aligned}
\mathbf{\bar x} &= \mathbf x +
\begin{bmatrix}- R\sin(\theta) + R\sin(\theta + \beta) \\
R\cos(\theta) - R\cos(\theta + \beta) \\
\beta\end{bmatrix}\\
\mathbf{\bar P} &=\mathbf{FPF}^{\mathsf T} + \mathbf{VMV}^{\mathsf T}
\end{aligned}$$

This form of linearization is not the only way to predict $\mathbf x$. For example, we could use a numerical integration technique such as *Runge Kutta* to compute the movement
of the robot. This will be required if the time step is relatively large. Things are not as cut and dried with the EKF as for the Kalman filter. For a real problem you have to carefully model your system with differential equations and then determine the most appropriate way to solve that system. The correct approach depends on the accuracy you require, how nonlinear the equations are, your processor budget, and numerical stability concerns.

### Design the Measurement Model

The robot's sensor provides a noisy bearing and range measurement to multiple known locations in the landscape. The measurement model must convert the state $\begin{bmatrix}x & y&\theta\end{bmatrix}^\mathsf T$ into a range and bearing to the landmark. If $\mathbf p$
is the position of a landmark, the range $r$ is

$$r = \sqrt{(p_x - x)^2 + (p_y - y)^2}$$

The sensor provides bearing relative to the orientation of the robot, so we must subtract the robot's orientation from the bearing to get the sensor reading, like so:

$$\phi = \arctan(\frac{p_y - y}{p_x - x}) - \theta$$

Thus our measurement model $h$ is

$$\begin{aligned}
\mathbf z& = h(\bar{\mathbf x}, \mathbf p) &+ \mathcal{N}(0, R)\\
&= \begin{bmatrix}
\sqrt{(p_x - x)^2 + (p_y - y)^2} \\
\arctan(\frac{p_y - y}{p_x - x}) - \theta
\end{bmatrix} &+ \mathcal{N}(0, R)
\end{aligned}$$

This is clearly nonlinear, so we need linearize $h$  at $\mathbf x$ by taking its Jacobian. We compute that with SymPy below.

```python
px, py = symbols('p_x, p_y')
z = Matrix([[sympy.sqrt((px-x)**2 + (py-y)**2)],
            [sympy.atan2(py-y, px-x) - theta]])
z.jacobian(Matrix([x, y, theta]))
```

$\displaystyle \left[\begin{matrix}\frac{- p_{x} + x}{\sqrt{\left(p_{x} - x\right)^{2} + \left(p_{y} - y\right)^{2}}} & \frac{- p_{y} + y}{\sqrt{\left(p_{x} - x\right)^{2} + \left(p_{y} - y\right)^{2}}} & 0\\- \frac{- p_{y} + y}{\left(p_{x} - x\right)^{2} + \left(p_{y} - y\right)^{2}} & - \frac{p_{x} - x}{\left(p_{x} - x\right)^{2} + \left(p_{y} - y\right)^{2}} & -1\end{matrix}\right]$

Now we need to write that as a Python function. For example we might write:

```python
from math import sqrt

def H_of(x, landmark_pos):
    """ compute Jacobian of H matrix where h(x) computes
    the range and bearing to a landmark for state x """

    px = landmark_pos[0]
    py = landmark_pos[1]
    hyp = (px - x[0, 0])**2 + (py - x[1, 0])**2
    dist = sqrt(hyp)

    H = array(
        [[-(px - x[0, 0]) / dist, -(py - x[1, 0]) / dist, 0],
         [ (py - x[1, 0]) / hyp,  -(px - x[0, 0]) / hyp, -1]])
    return H
```

We also need to define a function that converts the system state into a measurement.

```python
from math import atan2

def Hx(x, landmark_pos):
    """ takes a state variable and returns the measurement
    that would correspond to that state.
    """
    px = landmark_pos[0]
    py = landmark_pos[1]
    dist = sqrt((px - x[0, 0])**2 + (py - x[1, 0])**2)

    Hx = array([[dist],
                [atan2(py - x[1, 0], px - x[0, 0]) - x[2, 0]]])
    return Hx
```

### Design Measurement Noise

It is reasonable to assume that the noise of the range and bearing measurements are independent, hence

$$\mathbf R=\begin{bmatrix}\sigma_{range}^2 & 0 \\ 0 & \sigma_{bearing}^2\end{bmatrix}$$

### Implementation

We will use `FilterPy`'s `ExtendedKalmanFilter` class to implement the filter. Its `predict()` method uses the standard linear equations for the process model. Ours is nonlinear, so we will have to override `predict()` with our own implementation.  I'll want to also use this class to simulate the robot, so I'll add a method `move()` that computes the position of the robot which both `predict()` and my simulation can call.

The matrices for the prediction step are quite large. While writing this code I made several errors before I finally got it working. I only found my errors by using SymPy's `evalf` function. `evalf` evaluates a SymPy `Matrix` with specific values for the variables. I decided to demonstrate this technique to you, and used `evalf` in the Kalman filter code. You'll need to understand a couple of points.

First, `evalf` uses a dictionary to specify the values. For example, if your matrix contains an `x` and `y`, you can write

```python
    M.evalf(subs={x:3, y:17})
```

to evaluate the matrix for `x=3` and `y=17`.

Second, `evalf` returns a `sympy.Matrix` object. Use `numpy.array(M).astype(float)` to convert it to a NumPy array. `numpy.array(M)` creates an array of  type `object`, which is not what you want.

Here is the code for the EKF:

```python
from filterpy.kalman import ExtendedKalmanFilter as EKF
from numpy import array, sqrt
class RobotEKF(EKF):
    def __init__(self, dt, wheelbase, std_vel, std_steer):
        EKF.__init__(self, 3, 2, 2)
        self.dt = dt
        self.wheelbase = wheelbase
        self.std_vel = std_vel
        self.std_steer = std_steer

        a, x, y, v, w, theta, time = symbols(
            'a, x, y, v, w, theta, t')
        d = v*time
        beta = (d/w)*sympy.tan(a)
        r = w/sympy.tan(a)

        self.fxu = Matrix(
            [[x-r*sympy.sin(theta)+r*sympy.sin(theta+beta)],
             [y+r*sympy.cos(theta)-r*sympy.cos(theta+beta)],
             [theta+beta]])

        self.F_j = self.fxu.jacobian(Matrix([x, y, theta]))
        self.V_j = self.fxu.jacobian(Matrix([v, a]))

        # save dictionary and it's variables for later use
        self.subs = {x: 0, y: 0, v:0, a:0,
                     time:dt, w:wheelbase, theta:0}
        self.x_x, self.x_y, = x, y
        self.v, self.a, self.theta = v, a, theta

    def predict(self, u):
        self.x = self.move(self.x, u, self.dt)

        self.subs[self.theta] = self.x[2, 0]
        self.subs[self.v] = u[0]
        self.subs[self.a] = u[1]

        F = array(self.F_j.evalf(subs=self.subs)).astype(float)
        V = array(self.V_j.evalf(subs=self.subs)).astype(float)

        # covariance of motion noise in control space
        M = array([[self.std_vel*u[0]**2, 0],
                   [0, self.std_steer**2]])

        self.P = np.dot(F, self.P).dot(F.T) + np.dot(V, M).dot(V.T)

    def move(self, x, u, dt):
        hdg = x[2, 0]
        vel = u[0]
        steering_angle = u[1]
        dist = vel * dt

        if abs(steering_angle) > 0.001: # is robot turning?
            beta = (dist / self.wheelbase) * tan(steering_angle)
            r = self.wheelbase / tan(steering_angle) # radius

            dx = np.array([[-r*sin(hdg) + r*sin(hdg + beta)],
                           [r*cos(hdg) - r*cos(hdg + beta)],
                           [beta]])
        else: # moving in straight line
            dx = np.array([[dist*cos(hdg)],
                           [dist*sin(hdg)],
                           [0]])
        return x + dx
```

Now we have another issue to handle. The residual is notionally computed as $y = z - h(x)$ but this will not work because our measurement contains an angle in it. Suppose z has a bearing of $1^\circ$ and $h(x)$ has a bearing of $359^\circ$. Naively subtracting them would yield a angular difference of $-358^\circ$, whereas the correct value is $2^\circ$. We have to write code to correctly compute the bearing residual.

```python
def residual(a, b):
    """ compute residual (a-b) between measurements containing
    [range, bearing]. Bearing is normalized to [-pi, pi)"""
    y = a - b
    y[1] = y[1] % (2 * np.pi)    # force in range [0, 2 pi)
    if y[1] > np.pi:             # move to [-pi, pi)
        y[1] -= 2 * np.pi
    return y
```

The rest of the code runs the simulation and plots the results, and shouldn't need too much comment by now. I create a variable `landmarks` that contains the landmark coordinates. I update the simulated robot position 10 times a second, but run the EKF only once per second. This is for two reasons. First, we are not using Runge Kutta to integrate the differental equations of motion, so a narrow time step allows our simulation to be more accurate. Second, it is fairly normal in embedded systems to have limited processing speed. This forces you to run your Kalman filter only as frequently as absolutely needed.

```python
from filterpy.stats import plot_covariance_ellipse
from math import sqrt, tan, cos, sin, atan2
import matplotlib.pyplot as plt

dt = 1.0

def z_landmark(lmark, sim_pos, std_rng, std_brg):
    x, y = sim_pos[0, 0], sim_pos[1, 0]
    d = np.sqrt((lmark[0] - x)**2 + (lmark[1] - y)**2)
    a = atan2(lmark[1] - y, lmark[0] - x) - sim_pos[2, 0]
    z = np.array([[d + randn()*std_rng],
                  [a + randn()*std_brg]])
    return z

def ekf_update(ekf, z, landmark):
    ekf.update(z, HJacobian=H_of, Hx=Hx,
               residual=residual,
               args=(landmark), hx_args=(landmark))

def run_localization(landmarks, std_vel, std_steer,
                     std_range, std_bearing,
                     step=10, ellipse_step=20, ylim=None):
    ekf = RobotEKF(dt, wheelbase=0.5, std_vel=std_vel,
                   std_steer=std_steer)
    ekf.x = array([[2, 6, .3]]).T # x, y, steer angle
    ekf.P = np.diag([.1, .1, .1])
    ekf.R = np.diag([std_range**2, std_bearing**2])

    sim_pos = ekf.x.copy() # simulated position
    # steering command (vel, steering angle radians)
    u = array([1.1, .01])

    plt.figure()
    plt.scatter(landmarks[:, 0], landmarks[:, 1],
                marker='s', s=60)

    track = []
    for i in range(200):
        sim_pos = ekf.move(sim_pos, u, dt/10.) # simulate robot
        track.append(sim_pos)

        if i % step == 0:
            ekf.predict(u=u)

            if i % ellipse_step == 0:
                plot_covariance_ellipse(
                    (ekf.x[0,0], ekf.x[1,0]), ekf.P[0:2, 0:2],
                     std=6, facecolor='k', alpha=0.3)

            x, y = sim_pos[0, 0], sim_pos[1, 0]
            for lmark in landmarks:
                z = z_landmark(lmark, sim_pos,
                               std_range, std_bearing)
                ekf_update(ekf, z, lmark)

            if i % ellipse_step == 0:
                plot_covariance_ellipse(
                    (ekf.x[0,0], ekf.x[1,0]), ekf.P[0:2, 0:2],
                    std=6, facecolor='g', alpha=0.8)
    track = np.array(track)
    plt.plot(track[:, 0], track[:,1], color='k', lw=2)
    plt.axis('equal')
    plt.title("EKF Robot localization")
    if ylim is not None: plt.ylim(*ylim)
    plt.show()
    return ekf
```

```python
landmarks = array([[5, 10], [10, 5], [15, 15]])

ekf = run_localization(
    landmarks, std_vel=0.1, std_steer=np.radians(1),
    std_range=0.3, std_bearing=0.1)
print('Final P:', ekf.P.diagonal())
```

I have plotted the landmarks as solid squares. The path of the robot is drawn with a black line. The covariance ellipses for the predict step are light gray, and the covariances of the update are shown in green. To make them visible at this scale I have set the ellipse boundary at 6$\sigma$.

We can see that there is a lot of uncertainty added by our motion model, and that most of the error in in the direction of motion. We determine that from the shape of the blue ellipses. After a few steps we can see that the filter incorporates the landmark measurements and the errors improve.

I used the same initial conditions and landmark locations in the UKF chapter. The UKF achieves much better accuracy in terms of the error ellipse. Both perform roughly as well as far as their estimate for $\mathbf x$ is concerned.

Now let's add another landmark.

```python
landmarks = array([[5, 10], [10, 5], [15, 15], [20, 5]])

ekf = run_localization(
    landmarks, std_vel=0.1, std_steer=np.radians(1),
    std_range=0.3, std_bearing=0.1)
plt.show()
print('Final P:', ekf.P.diagonal())
```

The uncertainly in the estimates near the end of the track are smaller.  We can see the effect that multiple landmarks have on our uncertainty by only using the first two landmarks.

```python
ekf = run_localization(
    landmarks[0:2], std_vel=1.e-10, std_steer=1.e-10,
    std_range=1.4, std_bearing=.05)
print('Final P:', ekf.P.diagonal())
```

The estimate quickly diverges from the robot's path after passing the landmarks. The covariance also grows quickly. Let's see what happens with only one landmark:

```python
ekf = run_localization(
    landmarks[0:1], std_vel=1.e-10, std_steer=1.e-10,
    std_range=1.4, std_bearing=.05)
print('Final P:', ekf.P.diagonal())
```

As you probably suspected, one landmark produces a very bad result. Conversely, a large number of landmarks allows us to make very accurate estimates.

```python
landmarks = array([[5, 10], [10,  5], [15, 15], [20,  5], [15, 10],
                   [10,14], [23, 14], [25, 20], [10, 20]])

ekf = run_localization(
    landmarks, std_vel=0.1, std_steer=np.radians(1),
    std_range=0.3, std_bearing=0.1, ylim=(0, 21))
print('Final P:', ekf.P.diagonal())
```

### Discussion

I said that this was a real problem, and in some ways it is. I've seen alternative presentations that used robot motion models that led to simpler Jacobians. On the other hand, my model of the movement is also simplistic in several ways. First, it uses a bicycle model. A real car has two sets of tires, and each travels on a different radius. The wheels do not grip the surface perfectly. I also assumed that the robot responds instantaneously to the control input. Sebastian Thrun writes in *Probabilistic Robots* that this simplified model is justified because the filters perform well when used to track real vehicles. The lesson here is that while you have to have a reasonably accurate nonlinear model, it does not need to be perfect to operate well. As a designer you will need to balance the fidelity of your model with the difficulty of the math and the CPU time required to perform the linear algebra.

Another way in which this problem was simplistic is that we assumed that we knew the correspondance between the landmarks and measurements. But suppose we are using radar - how would we know that a specific signal return corresponded to a specific building in the local scene? This question hints at SLAM algorithms - simultaneous localization and mapping. SLAM is not the point of this book, so I will not elaborate on this topic.

## UKF vs EKF

In the last chapter I used the UKF to solve this problem. The difference in implementation should be very clear. Computing the Jacobians for the state and measurement models was not trivial despite a rudimentary motion model. A different problem could result in a Jacobian which is difficult or impossible to derive analytically. In contrast, the UKF only requires you to provide a function that computes the system motion model and another for the measurement model.

There are many cases where the Jacobian cannot be found analytically. The details are beyond the scope of this book, but you will have to use numerical methods to compute the Jacobian. That undertaking is not trivial, and you will spend a significant portion of a master's degree at a STEM school learning techniques to handle such situations. Even then you'll likely only be able to solve problems related to your field - an aeronautical engineer learns a lot about Navier Stokes equations, but not much about modelling chemical reaction rates.

So, UKFs are easy. Are they accurate? In practice they often perform better than the EKF. You can find plenty of research papers that prove that the UKF outperforms the EKF in various problem domains. It's not hard to understand why this would be true. The EKF works by linearizing the system model and measurement model at a single point, and the UKF uses $2n+1$ points.

Let's look at a specific example. Take $f(x) = x^3$ and pass a Gaussian distribution through it. I will compute an accurate answer using a monte carlo simulation. I generate 50,000 points randomly distributed according to the Gaussian, pass each through $f(x)$, then compute the mean and variance of the result.

The EKF linearizes the function by taking the derivative to find the slope at the evaluation point $x$. This slope becomes the linear function that we use to transform the Gaussian. Here is a plot of that.

```python
import kf_book.nonlinear_plots as nonlinear_plots
nonlinear_plots.plot_ekf_vs_mc()
```

The EKF computation is rather inaccurate. In contrast, here is the performance of the UKF:

```python
nonlinear_plots.plot_ukf_vs_mc(alpha=0.001, beta=3., kappa=1.)
```

Here we can see that the computation of the UKF's mean is accurate to 2 decimal places. The standard deviation is slightly off, but you can also fine tune how the UKF computes the distribution by using the $\alpha$, $\beta$, and $\gamma$ parameters for generating the sigma points. Here I used $\alpha=0.001$, $\beta=3$, and $\gamma=1$. Feel free to modify them to see the result. You should be able to get better results than I did. However, avoid over-tuning the UKF for a specific test. It may perform better for your test case, but worse in general.