diff --git a/src/010-particle-smoother.jl b/src/010-particle-smoother.jl index 4b7c228..1459855 100644 --- a/src/010-particle-smoother.jl +++ b/src/010-particle-smoother.jl @@ -2,6 +2,7 @@ using Base.Threads: @threads using ProgressMeter: @showprogress export two_filter_smoother +export backward_smoother """ two_filter_smoother(; timeline::Vector{DateTime}, xfwd::Matrix, xbwd::Matrix, model_move::ModelMove, vmap, n_sim::Int) @@ -107,4 +108,152 @@ function two_filter_smoother(;timeline::Vector{DateTime}, xfwd::Matrix, xbwd::Ma trials = NaN ) +end + +""" + backward_smoother(;timeline::Vector{DateTime}, xfwd::Matrix, model_move::ModelMove, vmap, n_sim::Int) + +The backward particle smoother that samples from `f(X_t | {Y_1 ... Y_T}) for t ∈ 1:T`. + +# Arguments (keywords) + +- `timeline`: A `Vector{DateTime}` of ordered, regularly spaced time stamps that defines the time steps for the simulation; +- `xfwd`: A `Matrix` of [`State`](@ref)s from the forward filter (see [`particle_filter()`](@ref)); +- `model_move`: A [`ModelMove`](@ref) instance; +- `vmap`: (optional) A `GeoArray` that defines the 'validity map' (see [`logpdf_move()`](@ref)); +- `n_sim`: An integer that defines the number of Monte Carlo simulations (see [`logpdf_move()`](@ref)); + +# Details + +[`backward_smoother()`](@ref) smooths particles from the particle filter (see [`particle_filter()`](@ref)). The `timeline` from the particle filter should be supplied as well as a `Matrix` of particles from a forward run. The backward smoother works by iteratively resampling particles in line with the probability density of movement between the particles at time `t` and time `t+1`. [`logpdf_move()`](@ref) is an internal function that evaluates the log probability of a movement step between particles. This function wraps the [`logpdf_step()`](@ref) generic. Methods are provided for built-in [`State`](@ref) and [`ModelMove`](@ref) sub-types. To use custom sub-types, a corresponding [`logpdf_step()`](@ref) method should be provided. In [`backward_smoother`](@ref), the `vmap` and `n_sim` arguments support the calculate of probability densities (see [`logpdf_move()`](@ref)). + +# Returns + +- A `NamedTuple`, of the same format as returned by [`particle_filter()`](@ref), with the following fields: + - `timesteps` + - `timestamps` + - `direction`: `nothing` + - `state` + - `ess` + - `maxlp`: `NaN` + - `convergence`: `true` + - `trial`: NaN + +# See also + +- [`particle_filter()`](@ref) implements the particle filter; +- [`logpdf_step()`](@ref), [`logpdf_move_normalisation()`](@ref) and [`logpdf_move()`](@ref) evaluate the log probability (density) of movement between two [`State`](@ref)s; +- [`two_filter_smoother()`](@ref) implements the two-filter particle smoother; + +""" +function backward_smoother(;timeline::Vector{DateTime}, xfwd::Matrix, model_move::ModelMove, vmap = nothing, n_sim::Int = 100) + + #### Set up + # Identify dimension of input state + zdim = hasfield(typeof(xfwd[1]), :z) + # Use vmap, if supplied + # * vmap may be supplied for 'horizontal' (2D) movement models + # * If vmap is supplied, we update we update xfwd.map_value to 0.0 or 1.0 + # * If we are within mobility of the coastline, map_value -> 0.0 + # * map_value = 0.0 indicates that not all moves from that point are valid + # * This is a flag that we need to simulate the normalisation constant in logpdf_move() + if !isnothing(vmap) + xfwd = edit_map_value(xfwd, vmap) + end + # Define smoothed particles matrix + # (rows: particles; columns: time steps) + xout = similar(xfwd) + xout[:, end] .= xfwd[:, end] + np, nt = size(xout) + # Initialise weights + # * Particles from the forward filter have uniform weights (w) thanks to resampling + # * At the last (first smoothing) time step, smoothed weights (ws) are also uniform + lw = log.(ones(np) ./ np) + ws = ones(np) ./ np + # Initialise ESS vector + ess = zeros(Float64, nt) + ess[nt] = np # = 1 / sum(abs2, ws) + # Set up LRU cache + # * This caches movement normalisation constants for s_{i, t} in logpdf_move() + cache = LRU{eltype(xfwd), Float64}(maxsize = np) + + #### Run smoothing + @showprogress desc = "Running backward smoother..." for t in (nt-1):-1:1 + + # Compute pairwise densities f(s_{j, t + 1} | f(s_{i, t}) * w + # * We use a matrix as densities may be directional + w_i_to_j = zeros(Float64, np, np) + @threads for i in 1:np + for j in 1:np + # Compute density of move from s_{i, t} to s_{j, t + 1} * w + w_i_to_j[i, j] = exp(logpdf_move(xfwd[i, t], xfwd[j, t + 1], zdim, model_move, t, vmap, n_sim, cache) + lw[i]) + end + end + + # Compute k-sum for each particle j at time t: ∑ₖ f(sⱼ,ₜ₊₁ | sₖ,ₜ) * w + # * I.e., the summed density of moves from k (t) into each j (t + 1) * w, summed over k + # * This is a sum over rows (by column) for the w_i_to_j matrix + # * It is indexed by j + ksums = sum(w_i_to_j, dims = 1) + + # Validate ksums + # * It is possible that the subset of recorded particles at t is not connected to the subset at t + 1 + # * But ksums[j] should not = 0 (below ksums is used for division) + # * This check guards against ksums = 0 errors + if any(ksums .== 0) + error("The summed density of moves from particle k (t) into j (t + 1) * w, ksums[j], is 0 at time step $t. Boost `n_record` in the filter or the number of particles used in the smoother.") + end + + # Compute smoothed weights (ws) for each i (t) as a sum over j (t + 1) + # 1. For each particle j (matrix column) at t + 1 + # 1.1. Divide by jth ksum element + # 1.2. Multiply by jth smoothed weight ws element + # 2. Then sum over columns (j), by row + # * Note that ksums & ws must be one-row matrix for correct broadcasting + # * We return a _vector_ of weights for resampling + ws = reshape(ws, 1, :) + ws = vec(sum((w_i_to_j ./ ksums[:, 1:np]) .* ws[:, 1:np], dims = 2)) + + # Validate weights + if all(x -> x == 0 || isnan(x), ws) + error("All weights are zero or NaN at time step $t.") + end + if !isapprox(sum(ws), 1.0) + error("Weights do not sum to one at time step $t (sum: $(sum(ws))).") + end + + # Compute ESS + ess[t] = 1 / sum(abs2, ws) + + # Resample particles and store for time t using ws + idx = resample(ws, np) + xout[:, t] .= xfwd[idx, t] + + # Update & renormalise weights + # * ws is the smoothed weight of every particle (i) at time t + # * This becomes the smoothed every particle (j) at time t + 1 + # * ... before it is updated for every particle (i) at the next time step + # * Renormalise + ws = ws[idx] + ws = ws ./ sum(ws) + + end + + #### Clean up + # Reset map_value + xout = edit_map_value(xout, model_move.map) + + #### Return outputs + # Follow particle_filter() format + ( + timesteps = collect(1:nt), + timestamps = timeline, + state = xout, + direction = nothing, + ess = ess, + maxlp = NaN, + convergence = true, + trials = NaN + ) + end \ No newline at end of file