ManifoldUpdate callback (#79)

* Add a simple non-iip `update` in Joseph form

* Implement first simple version of the ManifoldUpdate callback

* Format the files with JuliaFormatter.jl

* Make the ManifoldUpdate callback an IEKF update

Project.toml

@@ -1,7 +1,7 @@
 name = "ProbNumDiffEq"
 uuid = "bf3e78b0-7d74-48a5-b855-9609533b56a5"
 authors = ["Nathanael Bosch"]
-version = "0.3.0"
+version = "0.3.1"
 
 [deps]
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"

src/ProbNumDiffEq.jl

@@ -122,6 +122,8 @@ include("ieks.jl")
 export IEKS, solve_ieks
 
 include("devtools.jl")
+include("callbacks.jl")
+export ManifoldUpdate
 
 # Do as they do here:
 # https://github.com/SciML/OrdinaryDiffEq.jl/blob/v5.61.1/src/OrdinaryDiffEq.jl#L175-L193

src/callbacks.jl

@@ -0,0 +1,31 @@
+function ManifoldUpdate(residualf; maxiters=100, ϵ₁=1e-25, ϵ₂=1e-15)
+    condition(u, t, integ) = true
+
+    function affect!(integ)
+        @unpack u = integ
+        @unpack x, SolProj = integ.cache
+
+        f(m) = residualf(SolProj * m)
+
+        m, C = x
+        m_i = copy(m)
+
+        local m_i_new, C_i_new
+        for i in 1:maxiters
+            z = f(m_i)
+            J = ForwardDiff.jacobian(f, m_i)
+            S = X_A_Xt(C, J)
+
+            m_i_new, C_i_new = update(x, Gaussian(z .+ (J * (m - m_i)), S), J)
+
+            if norm(m_i_new .- m_i) < ϵ₁ && norm(z) < ϵ₂
+                break
+            end
+            m_i = m_i_new
+        end
+        copy!(x, Gaussian(m_i_new, C_i_new))
+
+        return nothing
+    end
+    return DiscreteCallback(condition, affect!)
+end

src/filtering/update.jl

@@ -21,6 +21,17 @@ function update(x::Gaussian, measurement::Gaussian, H::AbstractMatrix)
 
     return Gaussian(m_new, C_new)
 end
+function update(x::SRGaussian, measurement::SRGaussian, H::AbstractMatrix)
+    """In Joseph form"""
+    m, C = x
+    z, S = measurement
+
+    K = C * H' * inv(S)
+    m_new = m - K * z
+    C_new = X_A_Xt(C, (I - K * H))
+
+    return Gaussian(m_new, C_new)
+end
 
 """
     update!(x_out, x_pred, measurement, H, R=0)

test/specific_problems.jl

@@ -114,6 +114,22 @@ end
     @test solve(prob, EK0(order=3), callback=Callback()) isa ProbNumDiffEq.ProbODESolution
 end
 
+@testset "ManifoldUpdate callback test" begin
+    # Again: Harmonic Oscillator with condition on E=2
+    u0 = ones(2)
+    function harmonic_oscillator(du, u, p, t)
+        du[1] = u[2]
+        return du[2] = -u[1]
+    end
+    prob = ODEProblem(harmonic_oscillator, u0, (0.0, 100.0))
+
+    E(u) = [dot(u, u) - 2]
+
+    @test solve(prob, EK0(order=3)) isa ProbNumDiffEq.ProbODESolution
+    @test solve(prob, EK0(order=3), callback=ManifoldUpdate(E)) isa
+          ProbNumDiffEq.ProbODESolution
+end
+
 @testset "Problem definition with ParameterizedFunctions.jl" begin
     f = @ode_def LotkaVolterra begin
         dx = a * x - b * x * y