Maximum Likelihood Demonstration

GeneralisedFilters/src/GeneralisedFilters.jl

@@ -12,23 +12,16 @@ using StatsBase
 using CUDA
 using NNlib
 
+# Filtering utilities
+include("callbacks.jl")
+include("containers.jl")
+include("resamplers.jl")
+
 ## FILTERING BASE ##########################################################################
 
 abstract type AbstractFilter <: AbstractSampler end
 abstract type AbstractBatchFilter <: AbstractFilter end
 
-"""
-    instantiate(model, alg, initial; kwargs...)
-
-Create an intermediate storage object to store the proposed/filtered states at each step.
-"""
-function instantiate end
-
-# Default method
-function instantiate(model, alg, initial; kwargs...)
-    return Intermediate(initial, initial)
-end
-
 """
     initialise([rng,] model, alg; kwargs...)
 
@@ -37,9 +30,9 @@ Propose an initial state distribution.
 function initialise end
 
 """
-    step([rng,] model, alg, iter, intermediate, observation; kwargs...)
+    step([rng,] model, alg, iter, state, observation; kwargs...)
 
-Perform a combined predict and update call of the filtering on the intermediate storage.
+Perform a combined predict and update call of the filtering on the state.
 """
 function step end
 
@@ -70,24 +63,22 @@ function filter(
     model::AbstractStateSpaceModel,
     alg::AbstractFilter,
     observations::AbstractVector;
-    callback=nothing,
+    callback::Union{AbstractCallback,Nothing}=nothing,
     kwargs...,
 )
-    initial = initialise(rng, model, alg; kwargs...)
-    intermediate = instantiate(model, alg, initial; kwargs...)
-    isnothing(callback) || callback(model, alg, intermediate, observations; kwargs...)
+    state = initialise(rng, model, alg; kwargs...)
+    isnothing(callback) || callback(model, alg, state, observations, PostInit; kwargs...)
+
     log_evidence = initialise_log_evidence(alg, model)
 
     for t in eachindex(observations)
-        intermediate, ll_increment = step(
-            rng, model, alg, t, intermediate, observations[t]; callback, kwargs...
+        state, ll_increment = step(
+            rng, model, alg, t, state, observations[t]; callback, kwargs...
         )
         log_evidence += ll_increment
-        isnothing(callback) ||
-            callback(model, alg, t, intermediate, observations; kwargs...)
     end
 
-    return intermediate.filtered, log_evidence
+    return state, log_evidence
 end
 
 function initialise_log_evidence(::AbstractFilter, model::AbstractStateSpaceModel)
@@ -112,16 +103,20 @@ function step(
     model::AbstractStateSpaceModel,
     alg::AbstractFilter,
     iter::Integer,
-    intermediate,
+    state,
     observation;
+    callback::Union{AbstractCallback,Nothing}=nothing,
     kwargs...,
 )
-    intermediate.proposed = predict(rng, model, alg, iter, intermediate.filtered; kwargs...)
-    intermediate.filtered, ll_increment = update(
-        model, alg, iter, intermediate.proposed, observation; kwargs...
-    )
+    state = predict(rng, model, alg, iter, state; kwargs...)
+    isnothing(callback) ||
+        callback(model, alg, iter, state, observation, PostPredict; kwargs...)
 
-    return intermediate, ll_increment
+    state, ll_increment = update(model, alg, iter, state, observation; kwargs...)
+    isnothing(callback) ||
+        callback(model, alg, iter, state, observation, PostUpdate; kwargs...)
+
+    return state, ll_increment
 end
 
 ## SMOOTHING BASE ##########################################################################
@@ -131,11 +126,6 @@ abstract type AbstractSmoother <: AbstractSampler end
 # function smooth end
 # function backward end
 
-# Filtering utilities
-include("callbacks.jl")
-include("containers.jl")
-include("resamplers.jl")
-
 # Model types
 include("models/linear_gaussian.jl")
 include("models/discrete.jl")

GeneralisedFilters/src/algorithms/bootstrap.jl

@@ -7,28 +7,31 @@ function step(
     model::AbstractStateSpaceModel,
     alg::AbstractParticleFilter,
     iter::Integer,
-    intermediate,
+    state,
     observation;
     ref_state::Union{Nothing,AbstractVector}=nothing,
+    callback::Union{AbstractCallback,Nothing}=nothing,
     kwargs...,
 )
-    intermediate.proposed, intermediate.ancestors = resample(
-        rng, alg.resampler, intermediate.filtered
-    )
+    state = resample(rng, alg.resampler, state)
+    isnothing(callback) ||
+        callback(model, alg, iter, state, observation, PostResample; kwargs...)
 
-    intermediate.proposed = predict(
-        rng, model, alg, iter, intermediate.proposed; ref_state=ref_state, kwargs...
-    )
-    # TODO: this is quite inelegant and should be refactored
+    state = predict(rng, model, alg, iter, state; ref_state=ref_state, kwargs...)
+
+    # TODO: this is quite inelegant and should be refactored. It also might introduce bugs
+    # with callbacks that track the ancestry (and use PostResample)
     if !isnothing(ref_state)
-        CUDA.@allowscalar intermediate.ancestors[1] = 1
+        CUDA.@allowscalar state.ancestors[1] = 1
     end
+    isnothing(callback) ||
+        callback(model, alg, iter, state, observation, PostPredict; kwargs...)
 
-    intermediate.filtered, ll_increment = update(
-        model, alg, iter, intermediate.proposed, observation; kwargs...
-    )
+    state, ll_increment = update(model, alg, iter, state, observation; kwargs...)
+    isnothing(callback) ||
+        callback(model, alg, iter, state, observation, PostUpdate; kwargs...)
 
-    return intermediate, ll_increment
+    return state, ll_increment
 end
 
 struct BootstrapFilter{RS<:AbstractResampler} <: AbstractParticleFilter
@@ -46,13 +49,6 @@ function BootstrapFilter(
     return BootstrapFilter{ESSResampler}(N, conditional_resampler)
 end
 
-function instantiate(
-    ::StateSpaceModel{T}, filter::BootstrapFilter, initial; kwargs...
-) where {T}
-    N = filter.N
-    return ParticleIntermediate(initial, initial, Vector{Int}(undef, N))
-end
-
 function initialise(
     rng::AbstractRNG,
     model::StateSpaceModel{T},
@@ -71,36 +67,51 @@ function predict(
     model::StateSpaceModel,
     filter::BootstrapFilter,
     step::Integer,
-    filtered::ParticleDistribution;
+    state::ParticleDistribution;
     ref_state::Union{Nothing,AbstractVector}=nothing,
     kwargs...,
 )
-    new_particles = map(
-        x -> SSMProblems.simulate(rng, model.dyn, step, x; kwargs...), collect(filtered)
+    state.particles = map(
+        x -> SSMProblems.simulate(rng, model.dyn, step, x; kwargs...), collect(state)
     )
-    # Don't need to deep copy weights as filtered will be overwritten in the update step
-    proposed = ParticleDistribution(new_particles, filtered.log_weights)
 
-    return update_ref!(proposed, ref_state, step)
+    return update_ref!(state, ref_state, step)
 end
 
 function update(
     model::StateSpaceModel{T},
     filter::BootstrapFilter,
     step::Integer,
-    proposed::ParticleDistribution,
+    state::ParticleDistribution,
     observation;
     kwargs...,
 ) where {T}
+    old_ll = logsumexp(state.log_weights)
+
     log_increments = map(
         x -> SSMProblems.logdensity(model.obs, step, x, observation; kwargs...),
-        collect(proposed),
+        collect(state),
     )
 
-    new_weights = proposed.log_weights + log_increments
-    filtered = ParticleDistribution(deepcopy(proposed.particles), new_weights)
+    state.log_weights += log_increments
 
-    ll_increment = logsumexp(filtered.log_weights) - logsumexp(proposed.log_weights)
+    ll_increment = logsumexp(state.log_weights) - old_ll
 
-    return filtered, ll_increment
+    return state, ll_increment
+end
+
+# Application of bootstrap filter to hierarchical models
+function filter(
+    rng::AbstractRNG,
+    model::HierarchicalSSM,
+    alg::BootstrapFilter,
+    observations::AbstractVector;
+    ref_state::Union{Nothing,AbstractVector}=nothing,
+    kwargs...,
+)
+    ssm = StateSpaceModel(
+        HierarchicalDynamics(model.outer_dyn, model.inner_model.dyn),
+        HierarchicalObservations(model.inner_model.obs),
+    )
+    return filter(rng, ssm, alg, observations; ref_state=ref_state, kwargs...)
 end