This notebook is designed to demonstrate how to use the StateSpace.jl package to execute the Extended Kalman filter for a non-linear State Space model. The example that has been used here closely follows Markus Gesmann's Mages' blog titled Extended Kalman filter example in R which in turn was based on a blog post titled Fun with (Extended Kalman) Filters by Dominic Steinitz.
For those of you that do not need/want the explanation of the model and the code, you can skip right to the end of this notebook where the entire section of code required to run this example is given.
The problem that we will consider here is that of predicting the true value of a population (say of bacteria, but it could be anything) from some noisy measurements of the population. We also would like to predict the growth rate of the population, however we are faced with the problem that we do not observe the growth rate
It is assumed that the population evolves according to the logistic growth model. This is written mathematically as an ordinary differential equation:
$$\frac{\mathrm{d}p}{\mathrm{d}t} = rp\left(1 - \frac{p}{k} \right),$$where $p$ is the population number, $t$ is the time, and $k$ is the carrying capacity.
The solution to the above differetial equation is given by
where $p_0$ is the initial population.
This is almost in the form that we need. The only problem is that it needs to be formulated in a discrete manner whereby the population at the current time, $p_n$, is written in terms of the population and growth rate at the previous state, $p_{n-1}$ and $r_{n-1}$ respectively. This can be done like so:
where $\Delta t$ is the change in time from one measurement of the population to the next. This is our transistion (process) equation for the population.
As mentioned previously, we also want to know the growth rate, and hence it is both the population number and the growth rate value together that define the state of our system. Hence we need to also define how the growth rate evolves. Here we will simply assume that the growth rate stays constant. Hence the transition equation for the growth rate will be written as:
$$r_n = r_{n-1}.$$We will assume that these processes experience some noise. This will be in the form of zero mean Gaussian noise. We'll set the variance at 0.001 for both of these processes and hence we can write the full process equations as:
$$p_n = \frac{kp_{n-1}\exp(r_{n-1}\Delta t)}{k+p_{n-1}[\exp(r_{n-1}\Delta t) - 1]} + V_n \qquad V_n \sim \mathcal{N}(0, 0.001)$$$$r_n = r_{t-n} + W_n \qquad W_n \sim \mathcal{N}(0, 0.001)$$We assume that we observe the population directly but there is some zero mean Gaussian noise with variance = 25. Hence our observation model is:
$$y_n = p_n + Q_n \qquad Q_n \sim \mathcal{N}(0, 25),$$where $y_n$ is the observation at the current time step.
We do not observe the growth rate :(.
First we'll import the required modules
In [1]:
using StateSpace
using Distributions
using Gadfly
using DataFrames
using Colors
In [2]:
function logisticGrowth(r, p, k, t)
k * p * exp(r*t) / (k + p * (exp(r*t) - 1))
end
logisticGrowth(state) = logisticGrowth(state[1], state[2], k, Δt)
Out[2]:
The second definition of the logistic function will come in useful when we are defining the process function.
Let's set the values of the parameters that we're going to use for the logistic function
In [3]:
r = 0.2
k = 100.0
p0 = 0.1 * k
Δt = 0.1
Out[3]:
Now we'll set the variance of the observations (measurement noise)
In [4]:
measurement_noise_variance = 25.0
Out[4]:
We can now define the number of observations and generate the true and observed population values. Note that we don't actually observe the growth rate value so we'll set the corresponding values in the vector to 'NaN'. For more information on handling missing observations in this package check out the missing observations example.
In [5]:
numObs = 100
true_values = Vector{Float64}(numObs)
population_measurements = Vector{Float64}(numObs)
growth_rate_measurements = Vector{Float64}(numObs)
for i in 1:numObs
true_values[i] = logisticGrowth(r, p0, k, i*Δt)
population_measurements[i] = true_values[i] + randn() * sqrt(measurement_noise_variance)
growth_rate_measurements[i] = NaN
end
Finally we need to concatenate both population and growth rate arrays and ensure that have the correct dimensions
In [6]:
measurements = [growth_rate_measurements population_measurements]'
Out[6]:
In [7]:
function process_fcn(state)
predict_growth_rate = state[1]
predict_population = logisticGrowth(state[1], state[2], k, Δt)
new_state = [predict_growth_rate, predict_population]
return new_state
end
process_noise_mat = diagm([0.001, 0.001])
function observation_fcn(state)
growth_rate_observation = 0.0
population_observation = 1.0 * state[2]
observation = [growth_rate_observation, population_observation]
return observation
end
observation_noise_mat = diagm([1.0, measurement_noise_variance])
Out[7]:
some couple of things to note about the above code:
-The process of the growth and population are assumed uncorrelated, hence the off diagonal terms are zero. This is why the process noise matrix (and the observation noise matrix) are created as diagonal matrices.
-Although the growth rate is not observed, we've still given the observation noise for the growth rate a non-zero value. It could be any non-zero value, I've just picked 1.0 (It shouldn't change the overall results). The problem with setting this value to zero is that you end up getting a singularity and hence the matrix algebra doesn't work and it will crash. So it must be a non zero value.
We can now create an instance of the Nonlinear State Space Model
In [8]:
nonLinSSM = NonlinearGaussianSSM(process_fcn, process_noise_mat, observation_fcn, observation_noise_mat)
Out[8]:
In [9]:
initial_guess = MvNormal([0.5, 10], diagm([1.0,20.0]))
Out[9]:
In [10]:
filtered_state = filter(nonLinSSM, measurements, initial_guess, true)
Out[10]:
The last argument in the filter - 'true' - allows us to use the information we have about the observations that we actually observed. It also estimates the observations that we didn't observe using our process and observation model. For more information on handling missing observations in this package check out the missing observations example.
First we extract the filtered population values along with their corresponding $2\sigma$ values (the give the area that we would expect the true value to lie with 95% confidence)
In [11]:
x_data = 1:numObs
population_array = Vector{Float64}(numObs+1)
confidence_array = Vector{Float64}(numObs+1)
population_array[1] = initial_guess.μ[2]
confidence_array[1] = 2*sqrt(initial_guess.Σ.mat[2,2])
for i in x_data
current_state = filtered_state.state[i]
population_array[i+1] = current_state.μ[2]
confidence_array[i+1] = 2*sqrt(current_state.Σ.mat[2,2])
end
Next we will create a dataframe. This is simply so the syntax is simple for plotting the ribbon digram which will represent the state along with the confidence interval
In [12]:
df_fs = DataFrame(
x = [0;x_data],
y = population_array,
ymin = population_array - confidence_array,
ymax = population_array + confidence_array,
f = "Filtered values"
)
Out[12]:
Next we will create separate colours for the observations and the true value. Here we will generate the default distinguishable colors used by Gadfly.
In [13]:
n = 3 #We will require 3 different colors
getColors = distinguishable_colors(n, Color[LCHab(70, 60, 240)],
transform=c -> deuteranopic(c, 0.5),
lchoices=Float64[65, 70, 75, 80],
cchoices=Float64[0, 50, 60, 70],
hchoices=linspace(0, 330, 24))
Out[13]:
Finally we will plot the results of the filtered population estimates
In [14]:
population_state_plot = plot(
layer(x=0:numObs, y=[p0;measurements[2,:]'], Geom.point, Theme(default_color=getColors[2])),
layer(x=0:numObs, y=[p0;true_values], Geom.line, Theme(default_color=getColors[3])),
layer(df_fs, x=:x, y=:y, ymin=:ymin, ymax=:ymax, Geom.line, Geom.ribbon),
Guide.xlabel("Measurement Number"), Guide.ylabel("Population"),
Guide.manual_color_key("Colour Key",["Filtered Estimate", "Measurements","True Value "],[getColors[1],getColors[2],getColors[3]]),
Guide.title("Extended Kalman Filter Example")
)
display(population_state_plot)
We can use a similar approach to plot the filtered growth rate values
In [15]:
x_data = 1:numObs
growth_rate_array = Vector{Float64}(numObs+1)
confidence_array = Vector{Float64}(numObs+1)
growth_rate_array[1] = initial_guess.μ[1]
confidence_array[1] = initial_guess.Σ.mat[1,1]
for i in x_data
current_state = filtered_state.state[i]
growth_rate_array[i+1] = current_state.μ[1]
confidence_array[i+1] = 2*sqrt(current_state.Σ.mat[1,1])
end
df_fs = DataFrame(
x = [0;x_data],
y = growth_rate_array,
ymin = growth_rate_array - confidence_array,
ymax = growth_rate_array + confidence_array,
f = "Filtered values"
)
growth_rate_state_plot = plot(
layer(x=0:numObs, y=ones(numObs+1)*r, Geom.line, Theme(default_color=getColors[3])),
layer(df_fs, x=:x, y=:y, ymin=:ymin, ymax=:ymax, Geom.line, Geom.ribbon),
Guide.xlabel("Measurement Number"), Guide.ylabel("Growth Rate"),
Guide.manual_color_key("Colour Key",["Filtered Estimate", "Measurements","True Value "],[getColors[1],getColors[2],getColors[3]]),
Guide.title("Extended Kalman Filter Example")
)
display(growth_rate_state_plot)
For those that just want the code, without the explanation, you have come to the right place. Here it is:
In [1]:
using StateSpace
using Distributions
using Gadfly
using DataFrames
using Colors
r = 0.2 #r is the growth rate
k = 100.0 #k is the carrying capacity
p0 = 0.1 * k # p0 is the initial population
Δt = 0.1 # Δt is the change in time.
#Define the logistic growth function to set the observations
function logisticGrowth(r, p, k, t)
k * p * exp(r*t) / (k + p * (exp(r*t) - 1))
end
logisticGrowth(state) = logisticGrowth(state[1], state[2], k, Δt)
#Generate noisy measurements
measurement_noise_variance = 25.0
numObs = 100
true_values = Vector{Float64}(numObs)
population_measurements = Vector{Float64}(numObs)
growth_rate_measurements = Vector{Float64}(numObs)
for i in 1:numObs
true_values[i] = logisticGrowth(r, p0, k, i*Δt)
population_measurements[i] = true_values[i] + randn() * sqrt(measurement_noise_variance)
growth_rate_measurements[i] = NaN
end
measurements = [growth_rate_measurements population_measurements]'
#Section: Describe Extended Kalman Filter parameters
function process_fcn(state)
predict_growth_rate = state[1]
predict_population = logisticGrowth(state)
new_state = [predict_growth_rate, predict_population]
return new_state
end
process_noise_mat = diagm([0.001, 0.001])
function observation_fcn(state)
growth_rate_observation = 0.0
population_observation = 1.0 * state[2]
observation = [growth_rate_observation, population_observation]
return observation
end
observation_noise_mat = diagm([1.0, measurement_noise_variance])
#Create instance of our EKF model
nonLinSSM = NonlinearGaussianSSM(process_fcn, process_noise_mat, observation_fcn, observation_noise_mat)
initial_guess = MvNormal([0.5, 10], diagm([1.0,20.0])) #Set initial guess of the state
filtered_state = filter(nonLinSSM, measurements, initial_guess, true) #Execute the Extended Kalman Filter
#Plot Filtered results for population
x_data = 1:numObs
population_array = Vector{Float64}(numObs+1)
confidence_array = Vector{Float64}(numObs+1)
population_array[1] = initial_guess.μ[2]
confidence_array[1] = 2*sqrt(initial_guess.Σ.mat[2,2])
for i in x_data
current_state = filtered_state.state[i]
population_array[i+1] = current_state.μ[2]
confidence_array[i+1] = 2*sqrt(current_state.Σ.mat[2,2])
end
df_fs = DataFrame(
x = [0;x_data],
y = population_array,
ymin = population_array - confidence_array,
ymax = population_array + confidence_array,
f = "Filtered values"
)
n = 3
getColors = distinguishable_colors(n, Color[LCHab(70, 60, 240)],
transform=c -> deuteranopic(c, 0.5),
lchoices=Float64[65, 70, 75, 80],
cchoices=Float64[0, 50, 60, 70],
hchoices=linspace(0, 330, 24))
population_state_plot = plot(
layer(x=0:numObs, y=[p0;measurements[2,:]'], Geom.point, Theme(default_color=getColors[2])),
layer(x=0:numObs, y=[p0;true_values], Geom.line, Theme(default_color=getColors[3])),
layer(df_fs, x=:x, y=:y, ymin=:ymin, ymax=:ymax, Geom.line, Geom.ribbon),
Guide.xlabel("Measurement Number"), Guide.ylabel("Population"),
Guide.manual_color_key("Colour Key",["Filtered Estimate", "Measurements","True Value "],[getColors[1],getColors[2],getColors[3]]),
Guide.title("Extended Kalman Filter Example")
)
display(population_state_plot)
#Plot Filtered results for growth rate
x_data = 1:numObs
growth_rate_array = Vector{Float64}(numObs+1)
confidence_array = Vector{Float64}(numObs+1)
growth_rate_array[1] = initial_guess.μ[1]
confidence_array[1] = initial_guess.Σ.mat[1,1]
for i in x_data
current_state = filtered_state.state[i]
growth_rate_array[i+1] = current_state.μ[1]
confidence_array[i+1] = 2*sqrt(current_state.Σ.mat[1,1])
end
df_fs = DataFrame(
x = [0;x_data],
y = growth_rate_array,
ymin = growth_rate_array - confidence_array,
ymax = growth_rate_array + confidence_array,
f = "Filtered values"
)
growth_rate_state_plot = plot(
layer(x=0:numObs, y=ones(numObs+1)*r, Geom.line, Theme(default_color=getColors[3])),
layer(df_fs, x=:x, y=:y, ymin=:ymin, ymax=:ymax, Geom.line, Geom.ribbon),
Guide.xlabel("Measurement Number"), Guide.ylabel("Growth Rate"),
Guide.manual_color_key("Colour Key",["Filtered Estimate", "Measurements","True Value "],[getColors[1],getColors[2],getColors[3]]),
Guide.title("Extended Kalman Filter Example")
)
display(growth_rate_state_plot)