Maximum Likelihood Demonstration

research/maximum_likelihood/Project.toml

@@ -0,0 +1,10 @@
+[deps]
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+GeneralisedFilters = "3ef92589-7ab8-43f9-b5b9-a3a0c86ecbb7"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
+Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+SSMProblems = "26aad666-b158-4e64-9d35-0e672562fa48"

research/maximum_likelihood/mle_demo.jl

@@ -0,0 +1,72 @@
+using GeneralisedFilters
+using SSMProblems
+using LinearAlgebra
+using Random
+
+## TOY MODEL ###############################################################################
+
+# this is taken from an example in Kalman.jl
+function toy_model(θ::T) where {T<:Real}
+    μ0 = T[1.0, 0.0]
+    Σ0 = Diagonal(ones(T, 2))
+
+    A = T[0.8 θ/2; -0.1 0.8]
+    Q = Diagonal(T[0.2, 1.0])
+    b = zeros(T, 2)
+
+    H = Matrix{T}(I, 1, 2)
+    R = Diagonal(T[0.2])
+    c = zeros(T, 1)
+
+    return create_homogeneous_linear_gaussian_model(μ0, Σ0, A, b, Q, H, c, R)
+end
+
+# data generation process
+rng = MersenneTwister(1234)
+true_model = toy_model(1.0)
+_, _, ys = sample(rng, true_model, 10000)
+
+# evaluate and return the log evidence
+function logℓ(θ, data)
+    rng = MersenneTwister(1234)
+    _, ll = GeneralisedFilters.filter(rng, toy_model(θ[]), KF(), data)
+    return -ll
+end
+
+# check type stability (important for use with Enzyme)
+@code_warntype logℓ([1.0], ys)
+
+## MLE #####################################################################################
+
+using DifferentiationInterface
+using ForwardDiff
+using Optimisers
+
+# initial value
+θ = [0.7]
+
+# setup optimiser (feel free to use other backends)
+state = Optimisers.setup(Optimisers.Descent(0.5), θ)
+backend = AutoForwardDiff()
+num_epochs = 1000
+
+# prepare gradients for faster AD
+grad_prep = prepare_gradient(logℓ, backend, θ, Constant(ys))
+hess_prep = prepare_hessian(logℓ, backend, θ, Constant(ys))
+
+for epoch in 1:num_epochs
+    # calculate gradients
+    val, ∇logℓ = DifferentiationInterface.value_and_gradient(
+        logℓ, grad_prep, backend, θ, Constant(ys)
+    )
+
+    # adjust the learning rate for a hacky Newton's method
+    H = DifferentiationInterface.hessian(logℓ, hess_prep, backend, θ, Constant(ys))
+    Optimisers.update!(state, θ, inv(H)*∇logℓ)
+
+    # stopping condition and printer
+    (epoch % 5) == 1 && println("$(epoch-1):\t $(θ[])")
+    if (∇logℓ'*∇logℓ) < 1e-12
+        break
+    end
+end

src/algorithms/kalman.jl

@@ -11,7 +11,7 @@ function initialise(
     rng::AbstractRNG, model::LinearGaussianStateSpaceModel, filter::KalmanFilter; kwargs...
 )
     μ0, Σ0 = calc_initial(model.dyn; kwargs...)
-    return Gaussian(μ0, Σ0)
+    return Gaussian(μ0, Matrix(Σ0))
 end
 
 function predict(