Unscented Kalman Filter (nonlinear version of the classic Kalman filter) in F#

Ukf.fs

module Drone.Control.Ukf

open System
open Drone.Control.Matrix


type IDiscreteModel =
    abstract Process : Matrix -> Matrix
    abstract Observe : Matrix -> Matrix
    abstract InitialState : Matrix
    abstract InitialCovariance : Matrix
    abstract ProcessNoiseMean : Matrix
    abstract ProcessNoiseCovariance : Matrix
    abstract ObservationNoiseMean : Matrix
    abstract ObservationNoiseCovariance : Matrix




/// From http://en.wikipedia.org/wiki/Kalman_filter#Unscented_Kalman_filter
type Ukf (model : IDiscreteModel) =

    let alpha = 0.001
    let beta = 2.0
    let kappa = 0.0

    let mutable estimatedState = model.InitialState
    let mutable covariance = model.InitialCovariance

    member this.EstimatedState = estimatedState

    member this.Predict () =
        // Augment the state vector with the process noise
        let xh_k1k1 = estimatedState
        let p_k1k1 = covariance
        let w_k = model.ProcessNoiseMean
        let q_k = model.ProcessNoiseCovariance
        let xa_k1k1 = mat [[xh_k1k1 |> mt; w_k |> mt]] |> mt |> meval
        let pa_k1k1 = mat [[p_k1k1;                          MZero (mrows p_k1k1, mcols q_k)]
                           [MZero (mrows q_k, mcols p_k1k1); q_k]] |> meval
        let nl = mrows model.ProcessNoiseMean
        let l = mrows xh_k1k1
        let lambda = alpha*alpha*(float l + kappa) - float l

        // Create "unscented" sampling points (sigma points)
        let sigmaPointOffset = (float l + lambda) * pa_k1k1 |> msqrt |> mcol
        let sigmas_k1k1 = Array.init (2*l + 1) (fun i ->
            if i = 0 then xa_k1k1
            else if i <= l then xa_k1k1 + sigmaPointOffset (i - 1) |> meval
            else xa_k1k1 - sigmaPointOffset ((i - 1) - l) |> meval)

        // Advance those samples and combine them for the next prediction
        let sigmas_kk1 = sigmas_k1k1 |> Array.map model.Process
        let wi = 1.0 / (2.0 * (float l + lambda))
        let ws0 = lambda / (float l + lambda)
        let wc0 = lambda / (float l + lambda) + (1.0 - alpha*alpha + beta)
        let xh_kk1 =
            sigmas_kk1
            |> Array.mapi (fun i s ->
                let w = if i = 0 then ws0 else wi
                w * s)
            |> Array.fold (fun a s -> a + s) (MZero (l, 1))
            |> meval
        let p_kk1 =
            sigmas_kk1
            |> Array.mapi (fun i s ->
                let w = if i = 0 then wc0 else wi
                let d = s - xh_kk1
                w * d * mt d |> meval)
            |> Array.fold (fun a s -> a + s) (MZero (l, l))
            |> meval
        xh_kk1, p_kk1

    member this.Update (observation : Matrix) =
        let xh_kk1, p_kk1 = this.Predict ()

        let n = mrows xh_kk1

        let r_k = model.ObservationNoiseCovariance

        // Augment the predicted state vector with the output noise
        let xa_kk1 = mat [[xh_kk1 |> mt; model.ObservationNoiseMean |> mt]] |> mt |> meval
        let pa_kk1 = mat [[p_kk1;                          MZero (mrows p_kk1, mcols r_k)]
                          [MZero (mrows r_k, mcols p_kk1); r_k]] |> meval

        // Sample around the predicted state
        let l = mrows xa_kk1
        let lambda = alpha*alpha*(float l + kappa) - float l
        //printfn "pa_kk1 = %O" pa_kk1
        let paSqrt = (float l + lambda) * pa_kk1 |> msqrt
        //printfn "paSqrt = %O" paSqrt
        let sigmaPointOffset = paSqrt |> mcol
        let sigma_kk1 = Array.init (2*l + 1) (fun i ->
            if i = 0 then xa_kk1
            else if i <= l then xa_kk1 + sigmaPointOffset (i - 1) |> meval
            else xa_kk1 - sigmaPointOffset ((i - 1) - l) |> meval)
        
        // Predict the observation for those samples
        let gamma_k = sigma_kk1 |> Array.map model.Observe
        let numMeas = mrows gamma_k.[0]

        // Calculate statistics for the observation
        let wi = 1.0 / (2.0 * (float l + lambda))
        let ws0 = lambda / (float l + lambda)
        let wc0 = lambda / (float l + lambda) + (1.0 - alpha*alpha + beta)
        let zh_k =
            gamma_k
            |> Array.mapi (fun i s ->
                let w = if i = 0 then ws0 else wi
                w * s |> meval)
            |> Array.fold (fun a s -> a + s) (MZero (numMeas, 1))
            |> meval
        let p_zz =
            gamma_k
            |> Array.mapi (fun i s ->
                let w = if i = 0 then wc0 else wi
                let d = s - zh_k
                w * d * mt d |> meval)
            |> Array.fold (fun a s -> a + s) (MZero (numMeas, numMeas))
            |> meval

        // Find the cross-covariance between state and output
        let p_xz =
            gamma_k
            |> Array.mapi (fun i s ->
                let w = if i = 0 then wc0 else wi
                let d1 = msub 0 0 n 1 sigma_kk1.[i] - xh_kk1
                let d2 = s - zh_k
                w * d1 * mt d2 |> meval)
            |> Array.fold (fun a s -> a + s) (MZero (n, numMeas))
            |> meval

        // Finally, the Kalman gain
        let k_k = p_xz * minv p_zz
        let z_k = observation
        let xh_kk = xh_kk1 + k_k * (z_k - zh_k)
        let p_kk = p_kk1 - k_k * p_zz * (k_k |> mt)

        // Save it
        estimatedState <- xh_kk
        covariance <- p_kk