Add Kalman filter example

examples/kalman-filter/Project.toml

@@ -0,0 +1,6 @@
+[deps]
+GaussianDistributions = "43dcc890-d446-5863-8d1a-14597580bb8d"
+Kalman = "d59c0ba6-2ef2-5409-8dc5-1fd9a2b46832"
+Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
+Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
+SSMProblems = "26aad666-b158-4e64-9d35-0e672562fa48"

examples/kalman-filter/script.jl

@@ -0,0 +1,105 @@
+## Kalman filter using Kalman.jl
+using GaussianDistributions: correct, Gaussian
+using LinearAlgebra
+using Statistics
+using Plots
+using Random
+using SSMProblems
+
+# Model definition
+struct LinearGaussianSSM <: AbstractStateSpaceModel
+ """
+ A state space model with linear dynamics and Gaussian noise.
+ The model is defined by the following equations:
+ x[0] = z + ϵ, ϵ ∼ N(0, P)
+ x[k] = Φx[k-1] + b + w[k], w[k] ∼ N(0, Q)
+ y[k] = Hx[k] + v[k], v[k] ∼ N(0, R)
+ """
+ z::Vector{Float64}
+ P::Matrix{Float64}
+ Φ::Matrix{Float64}
+ b::Vector{Float64}
+ Q::Matrix{Float64}
+ H::Matrix{Float64}
+ R::Matrix{Float64}
+end
+
+f0(model::LinearGaussianSSM) = Gaussian(model.z, model.P)
+f(x::Vector{Float64}, model::LinearGaussianSSM) = Gaussian(model.Φ * x + model.b, model.Q)
+g(y::Vector{Float64}, model::LinearGaussianSSM) = Gaussian(model.H * y, model.R)
+
+function transition!!(rng::AbstractRNG, model::LinearGaussianSSM)
+ return Gaussian(model.z, model.P)
+end
+
+function transition!!(rng::AbstractRNG, model::LinearGaussianSSM, state::Gaussian)
+ let Φ = model.Φ, Q = model.Q, μ = state.μ, Σ = state.Σ
+ return Gaussian(Φ * μ, Φ * Σ * Φ' + Q)
+ end
+end
+
+# Simulation parameters
+SEED = 1
+T = 100
+z = [-1.0, 1.0]
+P = Matrix(1.0I, 2, 2)
+Φ = [0.8 0.2; -0.1 0.8]
+b = zeros(2)
+Q = [0.2 0.0; 0.0 0.5]
+H = [1.0 0.0;]
+R = Matrix(0.3I, 1, 1)
+
+model = LinearGaussianSSM(z, P, Φ, b, Q, H, R)
+
+# Generate synthetic data
+rng = MersenneTwister(SEED)
+x, y = Vector{Any}(undef, T), Vector{Any}(undef, T)
+x[1] = rand(rng, f0(model))
+for t in 1:T
+ y[t] = rand(rng, g(x[t], model))
+ if t < T
+ x[t + 1] = rand(rng, f(x[t], model))
+ end
+end
+
+# Kalman filter
+function filter(rng::Random.AbstractRNG, model::LinearGaussianSSM, y::Vector{Any})
+ T = length(y)
+ p = transition!!(rng, model)
+ ps = [p]
+ for i in 1:T
+ p = transition!!(rng, model, p)
+ p, yres, _ = correct(p, Gaussian(y[i], model.R), model.H)
+ push!(ps, p)
+ end
+ return ps
+end
+
+# Run filter and plot results
+ps = filter(rng, model, y)
+
+p_mean = mean.(ps)
+p_cov = sqrt.(cov.(ps))
+
+p1 = scatter(1:T, first.(y); color="red", label="Observations")
+plot!(
+ p1,
+ 0:T,
+ first.(p_mean);
+ color="orange",
+ label="Filtered x1",
+ grid=false,
+ ribbon=getindex.(p_cov, 1, 1),
+ fillalpha=0.5,
+)
+
+plot!(
+ p1,
+ 0:T,
+ last.(p_mean);
+ color="blue",
+ label="Filtered x2",
+ grid=false,
+ ribbon=getindex.(p_cov, 2, 2),
+ fillalpha=0.5,
+)