diff --git a/Project.toml b/Project.toml index de25f4a67..74ede93c0 100644 --- a/Project.toml +++ b/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" diff --git a/src/caches.jl b/src/caches.jl index d9569fe8e..89fe6dadc 100644 --- a/src/caches.jl +++ b/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 diff --git a/src/perform_step.jl b/src/perform_step.jl index 1f61f1b53..afc5dd820 100644 --- a/src/perform_step.jl +++ b/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)) diff --git a/src/solution.jl b/src/solution.jl index a03514de1..c2fe1a4db 100644 --- a/src/solution.jl +++ b/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) diff --git a/src/state_initialization.jl b/src/state_initialization.jl index 2f5b52bf7..df922908a 100644 --- a/src/state_initialization.jl +++ b/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.μ diff --git a/test/specific_problems.jl b/test/specific_problems.jl index ba6fdb7a5..449660064 100644 --- a/test/specific_problems.jl +++ b/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 diff --git a/test/state_init.jl b/test/state_init.jl index 18c53606e..b5228fc74 100644 --- a/test/state_init.jl +++ b/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