Linear forward filtering backward smoothing/sampling with Zygote

sampler.jl

# Linear state space system forward filtering backward sampling with Zygote

using LinearAlgebra
using GaussianDistributions
using GaussianDistributions: ⊕, logpdf
using StaticArrays

using Statistics

using Zygote
using Zygote: @adjoint
using Test

using Random
Random.seed!(123)

lchol(a) = cholesky(a).L
pair(x) = (x...,)
sym(g::Gaussian) = Gaussian(g.μ, Symmetric(g.Σ))
sym(g::Gaussian{T,T}) where {T<:Number} = g

fit(::Type{Gaussian}, xs) = Gaussian(mean(xs), cov(xs))

"""
    correct(u, v, H) = u, yres, S
Joseph form correction step of a Kalman filter with u = Gaussian(x, P) state and
v = Gaussian(y, R) the observation with uncertainty R. H is the observation operator.
Returns corrected/conditional distribution u and the log-likelihood.
See https://en.wikipedia.org/wiki/Kalman_filter#Update.
"""
function correct(u, v, H, c = 0.0)
    x, Ppred = pair(u)
    y, R = pair(v)
    yres = y - H*x # innovation residual

    S = (H*Ppred*H' + R) # innovation covariance

    K = Ppred*H'*inv(S) # Kalman gain
    x = x + K*yres
    P = (I - K*H)*Ppred*(I - K*H)' + K*R*K'
    c = c - logpdf(Gaussian(zero(y), R), y)
    Gaussian(x, P), c
end


"""
Compute the law of Y = X + Z and return function v, x, z = v(x) which
samples X and Z conditional on Y = y.
Example: If x ~ N(3,1) and Z ~ N(0,1), J returns
N(3, 2), y -> v(y) = x, z
with v(y) sampling from X, Z conditional on Y = y
"""
function J(::typeof(+), x, z::Gaussian)
    # joint of x and x + z
    d = length(x.μ)
    x ⊕ z,
    function (y)
        x_ = rand(sym(GaussianDistributions.conditional(
            Gaussian([x.μ; x.μ + z.μ], ([x.Σ x.Σ; x.Σ (x.Σ + z.Σ)])),
            1:d,
            (d+1):2d,
            y,
        )))
        z_ = y - x_
        (x_, z_)
    end
end
@adjoint +(x::Gaussian, z::Gaussian) = J(+, x, z)
@adjoint *(x::Number, z::Gaussian) = x * z, y -> (x, x \ y) # what is z given x*z = y? just y/x (not random)
@adjoint *(x::Matrix, z::Gaussian) = x * z, y -> (x, x \ y) # what is z given x*z = y? just y/x (not random)

@adjoint *(z::Gaussian, x::Number) = x * z, y -> (y / x, x)

@adjoint +(x::Number, z::Gaussian) = x + z, y -> (x, y - x) # z given x + z = y, where x is not random
@adjoint +(x::Vector, z::Gaussian) = x + z, y -> (x, y - x) # z given x + z = y, where x is not random


"""
    myrand(u::Gaussian)
Normal: Sample Gaussian.
Forward pass: Pass through distribution.
Backward pass: Pass through.
"""
myrand(u) = rand(u)
@adjoint myrand(u) = u, y -> (y,)

"""
    observe!(data, x, H, R)
Normal call: Push observation `y = H*x + e` where `e ∼ N(0, R)` to `data`
Forward pass: condition `x` on the next observation `y = H*x + e` from `data`.
Backward pass: pass through.
"""
function observe!(data, x, H, R)
    push!(data, rand(Gaussian(H*x, R)))
    x
end
@adjoint observe!(data, x, H, R) = J(observe!, data, x, H, R)

function J(::typeof(observe!), data, x, H, R)
    n = nothing
    y = pop!(data)
    xf, c = correct(x, (y, R), H)
    xf, function (y)
        (n, y, n, n, n)
    end
end


# State space system
#
# x[0] ∼ N(x0, P0)
# x[k] = Φx[k−1] + w[k],    w[k] ∼ N(0, Q)
# y[k] = Hx[k] + v[k],    v[k] ∼ N(0, R)

x0 = [1., 0.]
P0 = Matrix(1.0I, 2, 2)

Φ = [0.8 0.5; -0.1 0.8]
Q = [0.2 0.0; 0.0 1.0]

yshadow = [0.0]
H = [1.0 0.0]
R = Matrix(1.0I, 1, 1)
INoise = Gaussian(zero(x0), Q)
Noise = Gaussian(zero(yshadow), R)


f(x) = Φ*x + myrand(INoise)
"""
    g(x, data = Any[]) -> x, data
Normal: create observation vector `data` from state space system.
Forward pass: Propagate uncertainty.
Backward pass: Sample smoothed trajectory.
"""
function g(x, data = Any[])
    x = observe!(data, x, H, R)
    for i in 1:2
        x = f(x)
        x = observe!(data, x, H, R)
    end
    x, data
end

x = Gaussian(x0, P0)

xsample = rand(x)

_, data = g(xsample)

(y, _), v = Zygote.pullback(g, x, reverse(data))
# y is Gaussian distribution of `x2` given the data.
# v(rand(y)) is a sample of `x0` given the data.

# Sample `x0` and fit Gaussian distribution
xhat = fit(Gaussian, [v((rand(y), []))[1] for i = 1:30000])


# Check that this is all correct:

# Write down joint distribution of x's and y's
# Define mean and covariance of the flattened vector of states and observations [x0 x1 x2 y0 y1 y2]
μ = [1.0, 0.0, 0.8, -0.1, 0.59, -0.16, 1.0, 0.8, 0.59]
Σ = [1.0   0.0   0.8    -0.1    0.59    -0.16     1.0   0.8    0.59
0.0   1.0   0.5     0.8    0.8      0.59     0.0   0.5    0.8
0.8   0.5   1.09    0.32   1.032    0.147    0.8   1.09   1.032
-0.1   0.8   0.32    1.65   1.081    1.288   -0.1   0.32   1.081
0.59  0.8   1.032   1.081  1.5661   0.7616   0.59  1.032  1.5661
-0.16  0.59  0.147   1.288  0.7616   2.0157  -0.16  0.147  0.7616
1.0   0.0   0.8    -0.1    0.59    -0.16     2.0   0.8    0.59
0.8   0.5   1.09    0.32   1.032    0.147    0.8   2.09   1.032
0.59  0.8   1.032   1.081  1.5661   0.7616   0.59  1.032  2.5661]

# Compute the conditional distribution of vector `x0` given data.
xtest = GaussianDistributions.conditional(Gaussian(μ, Σ), 1:2, 7:9, vcat(data...))

# agrees!
@test xhat ≈ xtest rtol=0.06