reorganized and added Mooncake MWE

Author: charlesknipp

Diff for: research/maximum_likelihood/Project.toml

@@ -1,7 +1,11 @@
 [deps]
 DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
+DistributionsAD = "ced4e74d-a319-5a8a-b0ac-84af2272839c"
+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"

Diff for: research/maximum_likelihood/mle_demo.jl

@@ -36,10 +36,10 @@ end
 # check type stability (important for use with Enzyme)
 @code_warntype logℓ([1.0], ys)
 
-## MLE #####################################################################################
+## NEWTONS METHOD ##########################################################################
 
 using DifferentiationInterface
-using ForwardDiff
+import ForwardDiff
 using Optimisers
 
 # initial value

Diff for: research/maximum_likelihood/mooncake_test.jl

@@ -0,0 +1,63 @@
+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 with small sample
+rng = MersenneTwister(1234)
+true_model = toy_model(1.0)
+_, _, ys = sample(rng, true_model, 20)
+
+## RUN MOONCKAE TESTS ######################################################################
+
+using DifferentiationInterface
+import Mooncake
+using DistributionsAD
+
+function build_objective(rng, θ, algo, data)
+    _, ll = GeneralisedFilters.filter(rng, toy_model(θ[]), algo, data)
+    return -ll
+end
+
+# kalman filter likelihood testing (works, but is slow)
+logℓ1 = θ -> build_objective(rng, θ, KF(), ys)
+Mooncake.TestUtils.test_rule(rng, logℓ1, [0.7]; is_primitive=false, debug_mode=true)
+
+# bootstrap filter likelihood testing (shouldn't work)
+logℓ2 = θ -> build_objective(rng, θ, BF(512), ys)
+Mooncake.TestUtils.test_rule(rng, logℓ2, [0.7]; is_primitive=false, debug_mode=true)
+
+## FOR USE WITH DIFFERENTIATION INTERFACE ##################################################
+
+# data should be part of the objective, but be held constant by DifferentiationInterface
+logℓ3 = (θ, data) -> build_objective(rng, θ, KF(), data)
+
+# set the backend with default configuration
+backend = AutoMooncake(; config=nothing)
+
+# prepare the gradient for faster subsequent iteration
+grad_prep = prepare_gradient(logℓ3, backend, [0.7], Constant(ys))
+
+# evaluate gradients and iterate to show proof of concept
+DifferentiationInterface.gradient(logℓ3, grad_prep, backend, [0.7], Constant(ys))
+DifferentiationInterface.gradient(logℓ3, grad_prep, backend, [0.8], Constant(ys))
+DifferentiationInterface.gradient(logℓ3, grad_prep, backend, [0.9], Constant(ys))
+DifferentiationInterface.gradient(logℓ3, grad_prep, backend, [1.0], Constant(ys))