From 52d06aa854378aac166e98aec33082c7bd60db92 Mon Sep 17 00:00:00 2001 From: FredericWantiez Date: Thu, 18 Jan 2024 20:41:00 +0000 Subject: [PATCH 1/7] EKF --- examples/extended-kalman-filter/Project.toml | 7 ++ examples/extended-kalman-filter/script.jl | 98 ++++++++++++++++++++ examples/extended-kalman-filter/sim.jl | 33 +++++++ 3 files changed, 138 insertions(+) create mode 100644 examples/extended-kalman-filter/Project.toml create mode 100644 examples/extended-kalman-filter/script.jl create mode 100644 examples/extended-kalman-filter/sim.jl diff --git a/examples/extended-kalman-filter/Project.toml b/examples/extended-kalman-filter/Project.toml new file mode 100644 index 0000000..6fe41ac --- /dev/null +++ b/examples/extended-kalman-filter/Project.toml @@ -0,0 +1,7 @@ +[deps] +Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +GaussianDistributions = "43dcc890-d446-5863-8d1a-14597580bb8d" +PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150" +Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" +SSMProblems = "26aad666-b158-4e64-9d35-0e672562fa48" diff --git a/examples/extended-kalman-filter/script.jl b/examples/extended-kalman-filter/script.jl new file mode 100644 index 0000000..574e078 --- /dev/null +++ b/examples/extended-kalman-filter/script.jl @@ -0,0 +1,98 @@ +# # Extended Kalman filter for a non-linear SSM: sine signal +using GaussianDistributions: correct, Gaussian +using LinearAlgebra +using Statistics +using Plots +using Random +using ForwardDiff: jacobian +using SSMProblems + +# Model definition +struct SineModel <: AbstractStateSpaceModel + """ + """ +end + + +f(x::AbstractArray, dt::Float64) = [x[1] + x[2] * dt; x[2] - sin(x[1])*dt] +jacob(x) = jacobian(state -> f(state, dt), x) + +f0(model::RangeBearingTracking) = Gaussian(model.z, model.P) +f(x::Vector{Float64}, model::RangeBearingTracking) = Gaussian(model.Φ * x + model.b, model.Q) +g(y::Vector{Float64}, model::RangeBearingTracking) = Gaussian(model.H * y, model.R) + +function transition!!(rng::AbstractRNG, model::RangeBearingTracking) + return Gaussian(model.z, model.P) +end + +function transition!!(rng::AbstractRNG, model::RangeBearingTracking, state::Gaussian) + let Φ = model.Φ, Q = model.Q, μ = state.μ, Σ = state.Σ + return Gaussian(Φ * μ, Φ * Σ * Φ' + Q) + end +end + +# Simulation parameters +SEED = 1 +T = 5 +nstep = 100 +dt = T / nstep +Q = [ + dt^3/3 dt^2/2; + dt^2/2 dt +] +R = 1.0*I + +model = NonLinearSSM() + +# 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::RangeBearingTracking, 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, +) diff --git a/examples/extended-kalman-filter/sim.jl b/examples/extended-kalman-filter/sim.jl new file mode 100644 index 0000000..9bfd3eb --- /dev/null +++ b/examples/extended-kalman-filter/sim.jl @@ -0,0 +1,33 @@ +using Random +using Distributions +using Plots +using LinearAlgebra + +sigma_e = 2 +sigma_t = 1e-7 +F = [ + 1. 0. 1. 0.; + 0. 1. 0. 1.; + 0. 0. 1. 0.; + 0. 0. 0. 1.; +] + +Z = sigma_t * 1.0I(2) +mu = zeros(2) + +function f(beta::AbstractArray) + atan(beta[2]/beta[1]) +end + +seed = 2 +N = 50 +rng = Random.MersenneTwister(seed) +beta = zeros(4, N) +beta[3:4, 1] = rand(rng, MvNormal(mu, Z)) +obs = zeros(N) +obs[1] = f(beta[:, 1]) + rand(rng, Normal(0., sigma_e)) +for t in 2:N + beta[:, t] = F * beta[:, t-1] + beta[3:4, t] += rand(rng, MvNormal(mu, Z)) + obs[t] = f(beta[:, t-1]) + rand(rng, Normal(0., sigma_e)) +end From fb20cc69860c413de600ecefc91d8d22c7b1493c Mon Sep 17 00:00:00 2001 From: FredericWantiez Date: Tue, 30 Jan 2024 22:17:46 +0000 Subject: [PATCH 2/7] Add EKF --- examples/extended-kalman-filter/script.jl | 120 +++++++++++----------- examples/extended-kalman-filter/sim.jl | 33 ------ 2 files changed, 59 insertions(+), 94 deletions(-) delete mode 100644 examples/extended-kalman-filter/sim.jl diff --git a/examples/extended-kalman-filter/script.jl b/examples/extended-kalman-filter/script.jl index 574e078..787e479 100644 --- a/examples/extended-kalman-filter/script.jl +++ b/examples/extended-kalman-filter/script.jl @@ -7,92 +7,90 @@ using Random using ForwardDiff: jacobian using SSMProblems -# Model definition -struct SineModel <: AbstractStateSpaceModel - """ - """ +struct PendulumModel + x0::Vector{Float64} + dt::Float64 + + Q::AbstractMatrix + R::AbstractMatrix end +# Simulation parameters +SEED = 4 +T = 5.0 +dt = 0.0125 +nstep = Int(T / dt) +g = 9.8 +r = 0.3 +qc = 1.0 -f(x::AbstractArray, dt::Float64) = [x[1] + x[2] * dt; x[2] - sin(x[1])*dt] -jacob(x) = jacobian(state -> f(state, dt), x) +x0 = [pi / 2; 0] +Q = qc .* [ + dt^3/3 dt^2/2 + dt^2/2 dt +] +model = PendulumModel(x0, dt, Q, r^2 * I(1)) -f0(model::RangeBearingTracking) = Gaussian(model.z, model.P) -f(x::Vector{Float64}, model::RangeBearingTracking) = Gaussian(model.Φ * x + model.b, model.Q) -g(y::Vector{Float64}, model::RangeBearingTracking) = Gaussian(model.H * y, model.R) +f(x::Array, model::PendulumModel) = + let dt = model.dt + [x[1] + x[2] * dt, x[2] - g * sin(x[1]) * dt] + end +h(x::Array, model::PendulumModel) = [sin(x[1])] -function transition!!(rng::AbstractRNG, model::RangeBearingTracking) - return Gaussian(model.z, model.P) +function transition!!(::AbstractRNG, model::PendulumModel) + return Gaussian(model.x0, 0.0) end -function transition!!(rng::AbstractRNG, model::RangeBearingTracking, state::Gaussian) - let Φ = model.Φ, Q = model.Q, μ = state.μ, Σ = state.Σ - return Gaussian(Φ * μ, Φ * Σ * Φ' + Q) - end +function transition!!(::AbstractRNG, model::PendulumModel, state::Gaussian) + # Jacobian - Linearization + Jf = jacobian(x -> f(x, model), state.μ) + Jh = jacobian(x -> h(x, model), state.μ) + pred = f(state.μ, model) + return Gaussian(pred, Jf * state.Σ * Jf' + model.Q) end -# Simulation parameters -SEED = 1 -T = 5 -nstep = 100 -dt = T / nstep -Q = [ - dt^3/3 dt^2/2; - dt^2/2 dt -] -R = 1.0*I - -model = NonLinearSSM() - # 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)) +x, y = Vector{Any}(undef, nstep), Vector{Any}(undef, nstep) +x[1] = x0 +for t in 1:nstep + y[t] = rand(rng, Gaussian(h(x[t], model), model.R)) + if t < nstep + x[t + 1] = rand(rng, Gaussian(f(x[t], model), model.Q)) end end -# Kalman filter -function filter(rng::Random.AbstractRNG, model::RangeBearingTracking, y::Vector{Any}) +function ekf_correct(obs, state::Gaussian, model::PendulumModel) + Jf = jacobian(x -> f(x, model), state.μ) # We should not have to recompute these jacobians + Jh = jacobian(x -> h(x, model), state.μ) + + S = model.R + Jh * state.Σ * Jh' + K = state.Σ * Jh' / S + pred = state.μ + K * (obs - h(state.μ, model)) + return Gaussian(pred, (I - K * Jh) * state.Σ) +end + +# Extended Kalman filter +function filter(rng::Random.AbstractRNG, model::PendulumModel, y::Vector) T = length(y) p = transition!!(rng, model) - ps = [p] + ps = [] for i in 1:T p = transition!!(rng, model, p) - p, yres, _ = correct(p, Gaussian(y[i], model.R), model.H) + p = ekf_correct(y[i], p, model) push!(ps, p) end return ps end -# Run filter and plot results ps = filter(rng, model, y) +ts = dt:dt:T +filtered_mean = first.(mean.(ps)) -p_mean = mean.(ps) -p_cov = sqrt.(cov.(ps)) +plot(ts, first.(x); color=:gray, label="Latent state") -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, +scatter!( + ts, first.(y); markersize=1, markerstrokealpha=0, label="Observations", color=:black ) -plot!( - p1, - 0:T, - last.(p_mean); - color="blue", - label="Filtered x2", - grid=false, - ribbon=getindex.(p_cov, 2, 2), - fillalpha=0.5, -) +plot!(ts, filtered_mean; label="Filtered mean", color=:red) diff --git a/examples/extended-kalman-filter/sim.jl b/examples/extended-kalman-filter/sim.jl deleted file mode 100644 index 9bfd3eb..0000000 --- a/examples/extended-kalman-filter/sim.jl +++ /dev/null @@ -1,33 +0,0 @@ -using Random -using Distributions -using Plots -using LinearAlgebra - -sigma_e = 2 -sigma_t = 1e-7 -F = [ - 1. 0. 1. 0.; - 0. 1. 0. 1.; - 0. 0. 1. 0.; - 0. 0. 0. 1.; -] - -Z = sigma_t * 1.0I(2) -mu = zeros(2) - -function f(beta::AbstractArray) - atan(beta[2]/beta[1]) -end - -seed = 2 -N = 50 -rng = Random.MersenneTwister(seed) -beta = zeros(4, N) -beta[3:4, 1] = rand(rng, MvNormal(mu, Z)) -obs = zeros(N) -obs[1] = f(beta[:, 1]) + rand(rng, Normal(0., sigma_e)) -for t in 2:N - beta[:, t] = F * beta[:, t-1] - beta[3:4, t] += rand(rng, MvNormal(mu, Z)) - obs[t] = f(beta[:, t-1]) + rand(rng, Normal(0., sigma_e)) -end From 0ace6bf8012793c0bdea4c1bee55f6fcf3bfa9b4 Mon Sep 17 00:00:00 2001 From: FredericWantiez Date: Thu, 1 Feb 2024 20:56:12 +0000 Subject: [PATCH 3/7] Add Literate --- examples/extended-kalman-filter/Project.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/examples/extended-kalman-filter/Project.toml b/examples/extended-kalman-filter/Project.toml index 6fe41ac..67e4c23 100644 --- a/examples/extended-kalman-filter/Project.toml +++ b/examples/extended-kalman-filter/Project.toml @@ -2,6 +2,7 @@ Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" GaussianDistributions = "43dcc890-d446-5863-8d1a-14597580bb8d" +Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306" PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150" Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" SSMProblems = "26aad666-b158-4e64-9d35-0e672562fa48" From 4791835494aae4133e8aa5042c33e8e59b4aec79 Mon Sep 17 00:00:00 2001 From: FredericWantiez Date: Thu, 1 Feb 2024 21:03:45 +0000 Subject: [PATCH 4/7] Comments --- examples/extended-kalman-filter/script.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/extended-kalman-filter/script.jl b/examples/extended-kalman-filter/script.jl index 787e479..6afc9ec 100644 --- a/examples/extended-kalman-filter/script.jl +++ b/examples/extended-kalman-filter/script.jl @@ -61,7 +61,7 @@ for t in 1:nstep end function ekf_correct(obs, state::Gaussian, model::PendulumModel) - Jf = jacobian(x -> f(x, model), state.μ) # We should not have to recompute these jacobians + Jf = jacobian(x -> f(x, model), state.μ) Jh = jacobian(x -> h(x, model), state.μ) S = model.R + Jh * state.Σ * Jh' From bc9828e0d75b71f6829811d40b92bb3926eccb53 Mon Sep 17 00:00:00 2001 From: FredericWantiez Date: Thu, 1 Feb 2024 21:10:06 +0000 Subject: [PATCH 5/7] More comments --- examples/extended-kalman-filter/script.jl | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/examples/extended-kalman-filter/script.jl b/examples/extended-kalman-filter/script.jl index 6afc9ec..402b3b0 100644 --- a/examples/extended-kalman-filter/script.jl +++ b/examples/extended-kalman-filter/script.jl @@ -42,7 +42,6 @@ function transition!!(::AbstractRNG, model::PendulumModel) end function transition!!(::AbstractRNG, model::PendulumModel, state::Gaussian) - # Jacobian - Linearization Jf = jacobian(x -> f(x, model), state.μ) Jh = jacobian(x -> h(x, model), state.μ) pred = f(state.μ, model) @@ -61,7 +60,7 @@ for t in 1:nstep end function ekf_correct(obs, state::Gaussian, model::PendulumModel) - Jf = jacobian(x -> f(x, model), state.μ) + Jf = jacobian(x -> f(x, model), state.μ) Jh = jacobian(x -> h(x, model), state.μ) S = model.R + Jh * state.Σ * Jh' From 136b006f3c6988c49d219bf0d90b0f1de6cb1dbc Mon Sep 17 00:00:00 2001 From: FredericWantiez Date: Tue, 5 Mar 2024 20:19:08 +0000 Subject: [PATCH 6/7] starting point --- examples/extended-kalman-filter/script.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/examples/extended-kalman-filter/script.jl b/examples/extended-kalman-filter/script.jl index 402b3b0..ba23efe 100644 --- a/examples/extended-kalman-filter/script.jl +++ b/examples/extended-kalman-filter/script.jl @@ -38,7 +38,7 @@ f(x::Array, model::PendulumModel) = h(x::Array, model::PendulumModel) = [sin(x[1])] function transition!!(::AbstractRNG, model::PendulumModel) - return Gaussian(model.x0, 0.0) + return Gaussian(model.x0, zeros(2,2)) end function transition!!(::AbstractRNG, model::PendulumModel, state::Gaussian) @@ -73,8 +73,8 @@ end function filter(rng::Random.AbstractRNG, model::PendulumModel, y::Vector) T = length(y) p = transition!!(rng, model) - ps = [] - for i in 1:T + ps = [p] + for i in 2:T p = transition!!(rng, model, p) p = ekf_correct(y[i], p, model) push!(ps, p) From 20a1e2a4f4fc5046c1a91ca92ed418a0e29e5f63 Mon Sep 17 00:00:00 2001 From: FredericWantiez Date: Tue, 5 Mar 2024 20:21:12 +0000 Subject: [PATCH 7/7] Format --- examples/extended-kalman-filter/script.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/extended-kalman-filter/script.jl b/examples/extended-kalman-filter/script.jl index ba23efe..88e6c6d 100644 --- a/examples/extended-kalman-filter/script.jl +++ b/examples/extended-kalman-filter/script.jl @@ -38,7 +38,7 @@ f(x::Array, model::PendulumModel) = h(x::Array, model::PendulumModel) = [sin(x[1])] function transition!!(::AbstractRNG, model::PendulumModel) - return Gaussian(model.x0, zeros(2,2)) + return Gaussian(model.x0, zeros(2, 2)) end function transition!!(::AbstractRNG, model::PendulumModel, state::Gaussian)