Add SecondOrderODEProblem support (only IIP for the time being) (#40)

* Added functionality to solve (IIP) SecondOrderODEs

* Fixed a bug, where the EK1 assumed a specific vector-field property

* Better test structure

* Bump version number

* Adjusted the tests for correct state initialization

Project.toml

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

src/caches.jl

@@ -3,7 +3,7 @@
 ########################################################################################
 abstract type ODEFiltersCache <: OrdinaryDiffEq.OrdinaryDiffEqCache end
 mutable struct GaussianODEFilterCache{
-    RType, ProjType, SolProjType, FP, uType, xType, AType, QType, matType, diffusionType, diffModelType,
+    RType, ProjType, SolProjType, FP, uType, duType, xType, AType, QType, matType, diffusionType, diffModelType,
     measType, llType,
 } <: ODEFiltersCache
     # Constants
@@ -28,14 +28,14 @@ mutable struct GaussianODEFilterCache{
     x_tmp2::xType
     measurement::measType
     H::matType
-    du::uType
+    du::duType
     ddu::matType
     K::matType
     G::matType
     covmatcache::matType
     local_diffusion::diffusionType
     global_diffusion::diffusionType
-    err_tmp::uType
+    err_tmp::duType
     log_likelihood::llType
 end
 
@@ -48,8 +48,13 @@ function OrdinaryDiffEq.alg_cache(
               "or a matrix) are currently not supported")
     end
 
+    is_secondorder_ode = f isa DynamicalODEFunction
+    if is_secondorder_ode
+        @warn "Assuming that the given ODE is a SecondOrderODE. If this is not the case, e.g. because it is some other dynamical ODE, the solver will probably run into errors!"
+    end
+
     q = alg.order
-    d = length(u)
+    d = is_secondorder_ode ? length(u[1, :]) : length(u)
     D = d*(q+1)
 
     u0 = u
@@ -63,7 +68,7 @@ function OrdinaryDiffEq.alg_cache(
     Proj(deriv) = deriv > q ? error("Projection called for non-modeled derivative") :
         kron([i==(deriv+1) ? 1 : 0 for i in 1:q+1]', diagm(0 => ones(uElType, d)))
     @assert f isa AbstractODEFunction
-    SolProj = f isa DynamicalODEFunction ? [Proj(0); Proj(1)] : Proj(0)
+    SolProj = f isa DynamicalODEFunction ? [Proj(1); Proj(0)] : Proj(0)
 
     # Prior dynamics
     @assert alg.prior == :ibm "Only the ibm prior is implemented so far"
@@ -98,7 +103,7 @@ function OrdinaryDiffEq.alg_cache(
 
     return GaussianODEFilterCache{
         typeof(R), typeof(Proj), typeof(SolProj), typeof(Precond),
-        uType, typeof(x0), typeof(A), typeof(Q), matType, typeof(initdiff),
+        uType, typeof(du), typeof(x0), typeof(A), typeof(Q), matType, typeof(initdiff),
         typeof(diffmodel), typeof(measurement), uEltypeNoUnits,
     }(
         # Constants
@@ -108,7 +113,7 @@ function OrdinaryDiffEq.alg_cache(
         copy(x0), copy(x0), copy(x0), copy(x0), copy(x0),
         measurement,
         H, du, ddu, K, G, covmatcache, initdiff, initdiff,
-        copy(u0),
+        copy(du),
         zero(uEltypeNoUnits),
     )
 end

src/perform_step.jl

@@ -77,9 +77,17 @@ function OrdinaryDiffEq.perform_step!(integ, cache::GaussianODEFilterCache, repe
     # Estimate error for adaptive steps
     if integ.opts.adaptive
         err_est_unscaled = estimate_errors(integ, integ.cache)
-        DiffEqBase.calculate_residuals!(
-            err_tmp, dt * err_est_unscaled, integ.u, u_filt,
-            integ.opts.abstol, integ.opts.reltol, integ.opts.internalnorm, t)
+        if integ.f isa DynamicalODEFunction # second-order ODE
+            DiffEqBase.calculate_residuals!(
+                err_tmp, dt * err_est_unscaled,
+                integ.u[1, :], u_filt[1, :],
+                integ.opts.abstol, integ.opts.reltol, integ.opts.internalnorm, t)
+        else # regular first-order ODE
+            DiffEqBase.calculate_residuals!(
+                err_tmp, dt * err_est_unscaled,
+                integ.u, u_filt,
+                integ.opts.abstol, integ.opts.reltol, integ.opts.internalnorm, t)
+        end
         integ.EEst = integ.opts.internalnorm(err_tmp, t) # scalar
     end
 
@@ -92,7 +100,7 @@ function OrdinaryDiffEq.perform_step!(integ, cache::GaussianODEFilterCache, repe
     end
 end
 
-function measure!(integ, x_pred, t)
+function measure!(integ, x_pred, t, second_order::Val{false})
     @unpack f, p, dt, alg = integ
     @unpack u_pred, du, ddu, Proj, Precond, measurement, R, H = integ.cache
     @assert iszero(R)
@@ -131,6 +139,51 @@ function measure!(integ, x_pred, t)
     return measurement
 end
 
+function measure!(integ, x_pred, t, second_order::Val{true})
+    @unpack f, p, dt, alg = integ
+    @unpack d, u_pred, du, ddu, Proj, Precond, measurement, R, H = integ.cache
+    @assert iszero(R)
+    du2 = du
+
+    PI = inv(Precond(dt))
+    E0, E1, E2 = Proj(0), Proj(1), Proj(2)
+
+    z, S = measurement.μ, measurement.Σ
+
+    # Mean
+    # _u_pred = E0 * PI * x_pred.μ
+    # _du_pred = E1 * PI * x_pred.μ
+    @assert isinplace(f) "Currently the code only supports IIP `SecondOrderProblem`s"
+    f.f1(du2, view(u_pred, 1:d), view(u_pred, d+1:2d), p, t)
+    integ.destats.nf += 1
+    z .= E2*PI*x_pred.μ .- du2
+
+    # Cov
+    if alg isa EK1
+        @assert !(alg isa IEKS)
+
+        J0 = copy(ddu)
+        ForwardDiff.jacobian!(J0, (du2, u) -> f.f1(du2, view(u_pred, 1:d), u, p, t), du2,
+                              u_pred[d+1:2d])
+
+        J1 = copy(ddu)
+        ForwardDiff.jacobian!(J1, (du2, du) -> f.f1(du2, du, view(u_pred, d+1:2d),
+                                                   p, t), du2,
+                              u_pred[1:d])
+
+        integ.destats.njacs += 1
+        mul!(H, (E2 .- J0 * E0 .- J1 * E1), PI)
+    else
+        mul!(H, E2, PI)
+    end
+
+    copy!(S, Matrix(X_A_Xt(x_pred.Σ, H)))
+
+    return measurement
+end
+measure!(integ, x_pred, t) = measure!(
+    integ, x_pred, t, Val(integ.f isa DynamicalODEFunction))
+
 # The following functions are just there to handle both IIP and OOP easily
 _eval_f!(du, u, p, t, f::AbstractODEFunction{true}) = f(du, u, p, t)
 _eval_f!(du, u, p, t, f::AbstractODEFunction{false}) = (du .= f(u, p, t))

src/solution.jl

@@ -149,11 +149,12 @@ set_smooth(p::GaussianODEFilterPosterior) = GaussianODEFilterPosterior(
     p.d, p.q, p.SolProj, p.A, p.Q, p.Precond, true)
 function GaussianODEFilterPosterior(alg, u0)
     uElType = eltype(u0)
-    d = length(u0)
+    d = u0 isa ArrayPartition ? length(u0) ÷ 2 : length(u0)
     q = alg.order
 
     Proj(deriv) = kron([i==(deriv+1) ? 1 : 0 for i in 1:q+1]', diagm(0 => ones(uElType, d)))
-    SolProj = Proj(0)
+    SolProj = u0 isa ArrayPartition ? [Proj(1); Proj(0)] : Proj(0)
+
 
     A, Q = ibm(d, q, uElType)
     Precond = preconditioner(uElType, d, q)

src/state_initialization.jl

@@ -4,14 +4,17 @@ function initial_update!(integ)
     @unpack d, x, Proj = integ.cache
     q = integ.alg.order
 
-    condition_on!(x, Proj(0), u)
-
     f_derivatives = get_derivatives(u, f, p, t, q)
-    @assert length(1:q) == length(f_derivatives)
-    for (o, df) in zip(1:q, f_derivatives)
+    @assert length(0:q) == length(f_derivatives)
+    for (o, df) in zip(0:q, f_derivatives)
         condition_on!(x, Proj(o), df)
     end
 end
+
+
+"""
+    Compute initial derivatives of an ODEProblem with TaylorSeries.jl
+"""
 function get_derivatives(u, f, p, t, q)
     d = length(u)
 
@@ -38,9 +41,46 @@ function get_derivatives(u, f, p, t, q)
         push!(f_derivatives, df)
     end
 
-    return evaluate.(f_derivatives)
+    return [u, evaluate.(f_derivatives)...]
+end
+
+
+"""
+    Compute initial derivatives of a SecondOrderODE with TaylorSeries.jl
+"""
+function get_derivatives(u::ArrayPartition, f::DynamicalODEFunction, p, t, q)
+
+    d = length(u[1,:])
+    Proj(deriv) = deriv > q ? error("Projection called for non-modeled derivative") :
+        kron([i==(deriv+1) ? 1 : 0 for i in 1:q+1]', diagm(0 => ones(d)))
+
+    f_oop(du, u, p, t) = (ddu = copy(du); f.f1(ddu, du, u, p, t); return ddu)
+
+    # Make sure that the vector field f does not depend on t
+    f_t_taylor = taylor_expand(_t -> f_oop(u[1:d], u[d+1:end], p, _t), t)
+    @assert !(eltype(f_t_taylor) <: TaylorN) "The vector field depends on t; The code might not yet be able to handle these (but it should be easy to implement)"
+
+    set_variables("u", numvars=2d, order=q+1)
+
+    fp1 = taylor_expand(u -> f_oop(u[1:d], u[d+1:end], p, t), u[:])
+    fp2 = taylor_expand(u -> u[1:d], u[:])
+    f_derivatives = [fp1]
+    for o in 3:q
+        _curr_f_deriv = f_derivatives[end]
+        dfdu1 = stack([derivative.(_curr_f_deriv, i) for i in 1:d])'
+        dfdu2 = stack([derivative.(_curr_f_deriv, i) for i in d+1:2d])'
+        df = dfdu1 * fp1 + dfdu2 * fp2
+        push!(f_derivatives, df)
+    end
+
+    return [u[2,:], u[1,:], evaluate.(f_derivatives)...]
 end
 
+
+
+
+
+
 # TODO Either name texplicitly for the initial update, or think about how to use this in general
 function condition_on!(x::SRGaussian, H::AbstractMatrix, data::AbstractVector)
     z = H*x.μ

test/specific_problems.jl

@@ -6,6 +6,7 @@ using Test
 using LinearAlgebra
 using UnPack
 using ParameterizedFunctions
+using OrdinaryDiffEq
 
 
 using DiffEqProblemLibrary.ODEProblemLibrary: importodeproblems; importodeproblems()
@@ -110,17 +111,46 @@ end
 end
 
 
-@testset "2nd Order ODE" begin
-    function vanderpol!(ddu, du, u, p, t)
-        μ = p[1]
-        @. ddu = μ * ((1-u^2) * du - u)
-    end
+@testset "SecondOrderODEProblem" begin
+
     du0 = [0.0]
     u0 = [2.0]
     tspan = (0.0, 6.3)
     p = [1e1]
-    prob = SecondOrderODEProblem(vanderpol!, du0, u0, tspan, p)
-    @test_broken solve(prob, EK0(order=3)) isa ProbNumDiffEq.ProbODESolution
+
+    function vanderpol!(ddu, du, u, p, t)
+        μ = p[1]
+        @. ddu = μ * ((1-u^2) * du - u)
+    end
+    prob_iip = SecondOrderODEProblem(vanderpol!, du0, u0, tspan, p)
+
+    function vanderpol(du, u, p, t)
+        μ = p[1]
+        ddu = μ .* ((1 .- u .^ 2) .* du .- u)
+        return ddu
+    end
+    prob_oop = SecondOrderODEProblem(vanderpol, du0, u0, tspan, p)
+
+    appxsol = solve(prob_iip, Tsit5(), abstol=1e-7, reltol=1e-7)
+
+    @testset "IIP" begin
+        for alg in (EK0(), EK1())
+            @testset "$alg" begin
+                @test solve(prob_iip, alg) isa ProbNumDiffEq.ProbODESolution
+                @test solve(prob_iip, alg).u[end] ≈ appxsol.u[end] rtol=1e-3
+            end
+        end
+    end
+
+    @testset "OOP" begin
+        appxsol = solve(prob, Tsit5())
+        for alg in (EK0(), EK1())
+            @testset "$alg" begin
+                @test_broken solve(prob_oop, alg) isa ProbNumDiffEq.ProbODESolution
+                @test_broken solve(prob_oop, alg).u[end] ≈ appxsol.u[end] rtol=1e-3
+            end
+        end
+    end
 end
 
 

test/state_init.jl

@@ -30,8 +30,8 @@ true_init_states = [u(t0); du(t0); ddu(t0); dddu(t0); ddddu(t0); dddddu(t0); ddd
 
 @testset "OOP state init" begin
     dfs = ProbNumDiffEq.get_derivatives(prob.u0, prob.f, prob.p, prob.tspan[1], q)
-    @test length(dfs) == q
-    @test true_init_states[d+1:end] ≈ vcat(dfs...)
+    @test length(dfs) == q+1
+    @test true_init_states ≈ vcat(dfs...)
 end
 
 
@@ -40,6 +40,6 @@ end
     prob = ODEProblem(f!, u0, tspan)
 
     dfs = ProbNumDiffEq.get_derivatives(prob.u0, prob.f, prob.p, prob.tspan[1], q)
-    @test length(dfs) == q
-    @test true_init_states[d+1:end] ≈ vcat(dfs...)
+    @test length(dfs) == q+1
+    @test true_init_states ≈ vcat(dfs...)
 end