TuringLang / SSMProblems.jl

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

Consistency of filtering method syntax #30

Closed charlesknipp closed 1 month ago

charlesknipp commented 9 months ago

The recent addition of the Kalman filter is nice, but also introduces some problems when comparing outputs to the particle filter. The syntax introduced in #26 lacks the Turing style function call sample(rng,model,data). Even if it did, multiple dispatch couldn't identify whether to use a Kalman filter or a particle filter.

I propose the following structure to add context for the sampler given the filtering methods provided:

abstract type AbstractFilter end

struct ParticleFilter <: AbstractFilter
    num_particles::Int
    resampling_algorithm
    ess_threshold::Float
end

function sample(
        rng::AbstractRNG,
        model::AbstractStateSpaceModel,
        observations::AbstractVector,
        filter::ParticleFilter
    )
    particles = map(1:N) do i
        state = transition!!(rng, model)
        SSMProblems.Utils.Particle(state)
    end

    states, log_likelihood = sweep!(rng, model, particles, observations, filter)
    return samples, log_likelihood 
end

struct KalmanFilter <: AbstractFilter end

function sample(
        rng::AbstractRNG,
        model::LinearGaussianSSM,
        observations::AbstractVector,
        filter::KalmanFilter
    )
    states, log_likelihood = filter(rng, model, observations)
    return states, log_likelihood
end

Clearly, some modifications need to be made to sweep!() to allow arguments of type ParticleFilter, which is trivial.

I also think this could be taken a step further to address #27 with the function filter_step!(rng, model, observation, filter) to return an iterative step of the filtering process for both Kalman and particle filters. This could be especially useful for implementations of online filtering algorithms such as SMC^2

THargreaves commented 1 month ago

This is now a de facto part of the interface and will soon be formalised.