Particle Filtering methods

Specific Domains Modelling & Simulations
6

For completeness, here’s an example that sets up the sparse jacobians. With this, I can push it to a 50x50 grid, but beyond that it becomes quite expensive. For large PDE problems, Ensemble Kalman filtering is commonly employed for this reason. There are still a lot of room for optimization of the performance here, for example, one may use in-place updates of the dynamics and measurement functions to prevent excessive memory allocations.

using LowLevelParticleFilters
using LinearAlgebra
using Random
using Statistics
using SparseConnectivityTracer, SparseMatrixColorings, ForwardDiff, DifferentiationInterface

# Grid and state
Nx, Ny = 50, 50
N = Nx * Ny

# Initial state
function initial_field()
    x = LinRange(0, 2π, Nx)
    y = LinRange(0, 2π, Ny)
    [sin(xi) * sin(2yi) for xi in x, yi in y] |> vec
end

# Dynamics: advection by 1 cell to the right
function dynamics(x, u, p, t)
    fmat = reshape(x, Nx, Ny)
    vec(circshift(fmat, (0, 1)))
end

# Observation: 100 random point samples
sensor_indices = rand(1:N, 100)
function measurement(x, u, p, t)
    x[sensor_indices]
end

# Covariances
R1 = 0.01^2 * I(N)  # small process noise
R2 = 0.1^2 * I(length(sensor_indices))  # observation noise

# Initialize EKF
x0 = initial_field()
P0 = I(N) * 0.1
d0 = LowLevelParticleFilters.SimpleMvNormal(x0, P0)

## Set up sparse jacobian using DifferentiationInterface for dynamics
f(x) = dynamics(x, nothing, nothing, 0)
backend = AutoSparse(AutoForwardDiff())
prep = prepare_jacobian(f, backend, zero(x0))
grad = similar(x0*x0')
Ajac(x,u,p,t) = jacobian!(f, grad, prep, backend, x)

## Set up sparse jacobian using DifferentiationInterface for measurement
h(x) = measurement(x, nothing, nothing, 0)
prep_meas = prepare_jacobian(h, backend, zero(x0))
grad_meas = similar(h(x0)*x0')
Cjac(x,u,p,t) = jacobian!(h, grad_meas, prep_meas, backend, x)

ekf = ExtendedKalmanFilter(
    dynamics, measurement, R1, R2, d0; nu=0, Ajac, Cjac,
)

# Simulate true dynamics
T = 50
true_states = Vector{Vector{Float64}}(undef, T)
observations = Vector{Vector{Float64}}(undef, T)

x = x0
for t in 1:T
    global x
    x = dynamics(x, nothing, nothing, t) + 0.01 * randn(N)  # system evolution
    y = measurement(x, nothing, nothing, t) + 0.1 * randn(length(sensor_indices))  # observation
    true_states[t] = x
    observations[t] = y
end

# Run EKF
estimates = Vector{Vector{Float64}}(undef, T)
for t in 1:T
    @show t
    update!(ekf, nothing, observations[t])
    estimates[t] = ekf.x
end

## A GIF animation of the evolution
using Plots
anim = @animate for t in 1:T
    plot(
        heatmap(reshape(true_states[t], Nx, Ny), title = "True State at t=$t"),
        heatmap(reshape(estimates[t], Nx, Ny), title = "EKF Estimate at t=$t"),
        size = (800, 400)
    )
end
gif(anim, "ekf_advection.gif", fps = 5)

ekf_advection

4 Likes