TuringLang / SSMProblems.jl

Common abstractions for state-space models
http://turinglang.org/SSMProblems.jl/
MIT License
2 stars 2 forks source link

EKF #32

Closed FredericWantiez closed 1 month ago

THargreaves commented 7 months ago

What do you think about this as a way to generalise the interface?

abstract type GaussianSSM end

struct PendulumModel <: GaussianSSM
    x0::Float64
    dt::Float64
    Q::Matrix{Float64}
    R::Matrix{Float64}
end

function f(x::Vector{Float64}, model::PendulumModel)
    return [x[1] + x[2] * dt, x[2] - g * sin(x[1]) * model.dt]
end

# Pre-defined methods need to expect a covariance matrix.
# This could be defined completely differently for different concrete models.
# E.g. the covariance matrix could be an arbitrary function of the step number.
function process_covariance(model::PendulumModel)
    return model.Q
end

# Generic methods for any Gaussian SSM that do not need to
# be specified by the user

# Note, these don't use model.Q since this may not exist.
# Instead `process_covariance`` is called.

# Used to run EKF
function transition!!(model::GaussianSSM, state::Gaussian)
    Jf = ForwardDiff.jacobian(x -> f(x, model), state.μ)
    pred = f(state.μ, model)
    return Gaussian(pred, Jf * state.Σ * Jf' + process_covariance(model))
end

# Used to run particle filter
function transition!!(model::GaussianSSM, x::Vector{Float64})
    return f(x, model) + rand(MvNormal(process_covariance(model)))
end
THargreaves commented 7 months ago

This approach only has inheritance in one direction.

E.g. A LinearGaussianModel knows it's a GaussianSSM which knows it's a SSM

Hence a KF, EKF or PF can be used on a LinearGaussianModel.

The reverse is not true. Even if an GaussianSSM has affine f, as far as I can tell, there is no way for the KF to know that. I don't think this is a massive issue though.

FredericWantiez commented 6 months ago

What do you think about this as a way to generalise the interface?

abstract type GaussianSSM end

struct PendulumModel <: GaussianSSM
    x0::Float64
    dt::Float64
    Q::Matrix{Float64}
    R::Matrix{Float64}
end

function f(x::Vector{Float64}, model::PendulumModel)
    return [x[1] + x[2] * dt, x[2] - g * sin(x[1]) * model.dt]
end

# Pre-defined methods need to expect a covariance matrix.
# This could be defined completely differently for different concrete models.
# E.g. the covariance matrix could be an arbitrary function of the step number.
function process_covariance(model::PendulumModel)
    return model.Q
end

# Generic methods for any Gaussian SSM that do not need to
# be specified by the user

# Note, these don't use model.Q since this may not exist.
# Instead `process_covariance`` is called.

# Used to run EKF
function transition!!(model::GaussianSSM, state::Gaussian)
    Jf = ForwardDiff.jacobian(x -> f(x, model), state.μ)
    pred = f(state.μ, model)
    return Gaussian(pred, Jf * state.Σ * Jf' + process_covariance(model))
end

# Used to run particle filter
function transition!!(model::GaussianSSM, x::Vector{Float64})
    return f(x, model) + rand(MvNormal(process_covariance(model)))
end

In general the API here is really tailored to the particle filter point of view. To implement the unscented kalman filter we would need a few more functions, like your process_covariance as well as emission

THargreaves commented 6 months ago

In general the API here is really tailored to the particle filter point of view. To implement the unscented kalman filter we would need a few more functions, like your process_covariance as well as emission

I would agree with this as a general sentiment. It feels quite awkward writing these specific models in such a way that they can be generalised to a generic SSM that a particle filter can run on. At least all of this would be hidden from the user.

For the Kalman Filter the burden becomes even greater, you would need to define functions for:

as listed here.

I question how much benefit is gained from this added complexity.