Guided Filter Implementation

GeneralisedFilters/src/GeneralisedFilters.jl

@@ -54,7 +54,7 @@ function initialise(model, alg; kwargs...)
     return initialise(default_rng(), model, alg; kwargs...)
 end
 
-function predict(model, alg, step, filtered; kwargs...)
+function predict(model, alg, step, filtered, observation; kwargs...)
     return predict(default_rng(), model, alg, step, filtered; kwargs...)
 end
 
@@ -108,7 +108,7 @@ function step(
     callback::Union{AbstractCallback,Nothing}=nothing,
     kwargs...,
 )
-    state = predict(rng, model, alg, iter, state; kwargs...)
+    state = predict(rng, model, alg, iter, state, observation; kwargs...)
     isnothing(callback) ||
         callback(model, alg, iter, state, observation, PostPredict; kwargs...)
 
@@ -132,7 +132,7 @@ include("models/discrete.jl")
 include("models/hierarchical.jl")
 
 # Filtering/smoothing algorithms
-include("algorithms/bootstrap.jl")
+include("algorithms/particles.jl")
 include("algorithms/kalman.jl")
 include("algorithms/forward.jl")
 include("algorithms/rbpf.jl")

GeneralisedFilters/src/algorithms/forward.jl

@@ -14,7 +14,8 @@ function predict(
     model::DiscreteStateSpaceModel{T},
     filter::ForwardAlgorithm,
     step::Integer,
-    states::AbstractVector;
+    states::AbstractVector,
+    observation;
     kwargs...,
 ) where {T}
     P = calc_P(model.dyn, step; kwargs...)

GeneralisedFilters/src/algorithms/kalman.jl

@@ -1,6 +1,7 @@
 export KalmanFilter, filter, BatchKalmanFilter
 using GaussianDistributions
 using CUDA: i32
+import LinearAlgebra: hermitianpart
 
 export KalmanFilter, KF, KalmanSmoother, KS
 
@@ -18,42 +19,40 @@ end
 function predict(
     rng::AbstractRNG,
     model::LinearGaussianStateSpaceModel,
-    filter::KalmanFilter,
-    step::Integer,
-    filtered::Gaussian;
+    algo::KalmanFilter,
+    iter::Integer,
+    state::Gaussian,
+    observation=nothing;
     kwargs...,
 )
-    μ, Σ = GaussianDistributions.pair(filtered)
-    A, b, Q = calc_params(model.dyn, step; kwargs...)
+    μ, Σ = GaussianDistributions.pair(state)
+    A, b, Q = calc_params(model.dyn, iter; kwargs...)
     return Gaussian(A * μ + b, A * Σ * A' + Q)
 end
 
 function update(
     model::LinearGaussianStateSpaceModel,
-    filter::KalmanFilter,
-    step::Integer,
-    proposed::Gaussian,
-    obs::AbstractVector;
+    algo::KalmanFilter,
+    iter::Integer,
+    state::Gaussian,
+    observation::AbstractVector;
     kwargs...,
 )
-    μ, Σ = GaussianDistributions.pair(proposed)
-    H, c, R = calc_params(model.obs, step; kwargs...)
+    μ, Σ = GaussianDistributions.pair(state)
+    H, c, R = calc_params(model.obs, iter; kwargs...)
 
     # Update state
     m = H * μ + c
-    y = obs - m
-    S = H * Σ * H' + R
+    y = observation - m
+    S = hermitianpart(H * Σ * H' + R)
     K = Σ * H' / S
 
-    # HACK: force the covariance to be positive definite
-    S = (S + S') / 2
-
-    filtered = Gaussian(μ + K * y, Σ - K * H * Σ)
+    state = Gaussian(μ + K * y, Σ - K * H * Σ)
 
     # Compute log-likelihood
-    ll = logpdf(MvNormal(m, S), obs)
+    ll = logpdf(MvNormal(m, S), observation)
 
-    return filtered, ll
+    return state, ll
 end
 
 struct BatchKalmanFilter <: AbstractBatchFilter
@@ -74,12 +73,13 @@ function predict(
     rng::AbstractRNG,
     model::LinearGaussianStateSpaceModel{T},
     algo::BatchKalmanFilter,
-    step::Integer,
-    state::BatchGaussianDistribution;
+    iter::Integer,
+    state::BatchGaussianDistribution,
+    observation;
     kwargs...,
 ) where {T}
     μs, Σs = state.μs, state.Σs
-    As, bs, Qs = batch_calc_params(model.dyn, step, algo.batch_size; kwargs...)
+    As, bs, Qs = batch_calc_params(model.dyn, iter, algo.batch_size; kwargs...)
     μ̂s = NNlib.batched_vec(As, μs) .+ bs
     Σ̂s = NNlib.batched_mul(NNlib.batched_mul(As, Σs), NNlib.batched_transpose(As)) .+ Qs
     return BatchGaussianDistribution(μ̂s, Σ̂s)
@@ -88,17 +88,17 @@ end
 function update(
     model::LinearGaussianStateSpaceModel{T},
     algo::BatchKalmanFilter,
-    step::Integer,
+    iter::Integer,
     state::BatchGaussianDistribution,
-    obs;
+    observation;
     kwargs...,
 ) where {T}
     μs, Σs = state.μs, state.Σs
-    Hs, cs, Rs = batch_calc_params(model.obs, step, algo.batch_size; kwargs...)
-    D = size(obs, 1)
+    Hs, cs, Rs = batch_calc_params(model.obs, iter, algo.batch_size; kwargs...)
+    D = size(observation, 1)
 
     m = NNlib.batched_vec(Hs, μs) .+ cs
-    y_res = cu(obs) .- m
+    y_res = cu(observation) .- m
     S = NNlib.batched_mul(Hs, NNlib.batched_mul(Σs, NNlib.batched_transpose(Hs))) .+ Rs
 
     ΣH_T = NNlib.batched_mul(Σs, NNlib.batched_transpose(Hs))
@@ -151,7 +151,7 @@ function (callback::StateCallback)(
     algo::KalmanFilter,
     iter::Integer,
     state,
-    obs,
+    observation,
     ::PostPredictCallback;
     kwargs...,
 )
@@ -164,7 +164,7 @@ function (callback::StateCallback)(
     algo::KalmanFilter,
     iter::Integer,
     state,
-    obs,
+    observation,
     ::PostUpdateCallback;
     kwargs...,
 )
@@ -175,7 +175,7 @@ end
 function smooth(
     rng::AbstractRNG,
     model::LinearGaussianStateSpaceModel{T},
-    alg::KalmanSmoother,
+    algo::KalmanSmoother,
     observations::AbstractVector;
     t_smooth=1,
     callback=nothing,
@@ -190,7 +190,7 @@ function smooth(
     back_state = filtered
     for t in (length(observations) - 1):-1:t_smooth
         back_state = backward(
-            rng, model, alg, t, back_state, observations[t]; states_cache=cache, kwargs...
+            rng, model, algo, t, back_state, observations[t]; states_cache=cache, kwargs...
         )
     end
 
@@ -200,7 +200,7 @@ end
 function backward(
     rng::AbstractRNG,
     model::LinearGaussianStateSpaceModel{T},
-    alg::KalmanSmoother,
+    algo::KalmanSmoother,
     iter::Integer,
     back_state,
     obs;