Co-authored-by: Songchen Tan <[email protected]>

Author: hanqizhu65

lib/OrdinaryDiffEqTaylorSeries/Project.toml

@@ -4,11 +4,14 @@ authors = ["ParamThakkar123 <[email protected]>"]
 version = "1.1.0"
 
 [deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MuladdMacro = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221"
 OrdinaryDiffEqCore = "bbf590c4-e513-4bbe-9b18-05decba2e5d8"
+OrdinaryDiffEqDifferentiation = "4302a76b-040a-498a-8c04-15b101fed76b"
+OrdinaryDiffEqNonlinearSolve = "127b3ac7-2247-4354-8eb6-78cf4e7c58e8"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
@@ -20,12 +23,15 @@ TaylorDiff = "b36ab563-344f-407b-a36a-4f200bebf99c"
 TruncatedStacktraces = "781d530d-4396-4725-bb49-402e4bee1e77"
 
 [compat]
+ADTypes = "1.14.0"
 DiffEqBase = "6.152.2"
 DiffEqDevTools = "2.44.4"
 FastBroadcast = "0.3.5"
 LinearAlgebra = "<0.0.1, 1"
 MuladdMacro = "0.2.4"
 OrdinaryDiffEqCore = "1.1"
+OrdinaryDiffEqDifferentiation = "1.6.0"
+OrdinaryDiffEqNonlinearSolve = "1.6.0"
 PrecompileTools = "1.2.1"
 Preferences = "1.4.3"
 Random = "<0.0.1, 1"

lib/OrdinaryDiffEqTaylorSeries/src/OrdinaryDiffEqTaylorSeries.jl

@@ -2,7 +2,9 @@ module OrdinaryDiffEqTaylorSeries
 
 import OrdinaryDiffEqCore: alg_order, get_current_adaptive_order, get_current_alg_order,
                            alg_stability_size, explicit_rk_docstring,
-                           OrdinaryDiffEqAdaptiveAlgorithm, OrdinaryDiffEqMutableCache,
+                           OrdinaryDiffEqAdaptiveAlgorithm,
+                           OrdinaryDiffEqNewtonAdaptiveAlgorithm,
+                           OrdinaryDiffEqNewtonAlgorithm,OrdinaryDiffEqMutableCache,
                            alg_cache,
                            OrdinaryDiffEqConstantCache, @fold, trivial_limiter!,
                            step_accept_controller!,
@@ -15,6 +17,11 @@ import OrdinaryDiffEqCore: alg_order, get_current_adaptive_order, get_current_al
                            CompositeAlgorithm, _ode_addsteps!, copyat_or_push!,
                            AutoAlgSwitch, get_fsalfirstlast,
                            full_cache, DerivativeOrderNotPossibleError
+using OrdinaryDiffEqDifferentiation: UJacobianWrapper, dolinsolve
+using OrdinaryDiffEqNonlinearSolve: du_alias_or_new, markfirststage!, build_nlsolver,
+                                    nlsolve!, nlsolvefail, isnewton, get_W, set_new_W!,
+                                    NLNewton, COEFFICIENT_MULTISTEP
+import ADTypes: AutoForwardDiff
 import Static: False
 import MuladdMacro: @muladd
 import FastBroadcast: @..
@@ -59,6 +66,6 @@ PrecompileTools.@compile_workload begin
     solver_list = nothing
 end
 
-export ExplicitTaylor2, ExplicitTaylor, ExplicitTaylorAdaptiveOrder
+export ExplicitTaylor2, ExplicitTaylor, ExplicitTaylorAdaptiveOrder, ImplicitTaylor2
 
 end

lib/OrdinaryDiffEqTaylorSeries/src/TaylorSeries_caches.jl

@@ -150,3 +150,17 @@ function alg_cache(
     ExplicitTaylorAdaptiveOrderConstantCache(
         alg.min_order, alg.max_order, current_order, jets)
 end
+
+struct ImplicitTaylor2ConstantCache{N} <: OrdinaryDiffEqConstantCache
+    nlsolver::N
+end
+
+function alg_cache(alg::ImplicitTaylor2, u, rate_prototype, ::Type{uEltypeNoUnits},
+        ::Type{uBottomEltypeNoUnits},
+        ::Type{tTypeNoUnits}, uprev, uprev2, f, t, dt, reltol, p, calck,
+        ::Val{false}) where {uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits}
+    γ, c = 1, 1
+    nlsolver = build_nlsolver(alg, u, uprev, p, t, dt, f, rate_prototype, uEltypeNoUnits,
+        uBottomEltypeNoUnits, tTypeNoUnits, γ, c, Val(false))
+    ImplicitTaylor2ConstantCache(nlsolver)
+end

lib/OrdinaryDiffEqTaylorSeries/src/TaylorSeries_perform_step.jl

@@ -207,3 +207,71 @@ end
     end
     return nothing
 end
+
+function initialize!(integrator, cache::ImplicitTaylor2ConstantCache)
+    integrator.kshortsize = 2
+    integrator.k = typeof(integrator.k)(undef, integrator.kshortsize)
+    integrator.fsalfirst = integrator.f(integrator.uprev, integrator.p, integrator.t) # Pre-start fsal
+    OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
+
+    # Avoid undefined entries if k is an array of arrays
+    integrator.fsallast = zero(integrator.fsalfirst)
+    integrator.k[1] = integrator.fsalfirst
+    integrator.k[2] = integrator.fsallast
+end
+
+@muladd function perform_step!(integrator, cache::ImplicitTaylor2ConstantCache,
+        repeat_step = false)
+    @unpack t, dt, uprev, u, f, p = integrator
+    nlsolver = cache.nlsolver
+    alg = unwrap_alg(integrator, true)
+    markfirststage!(nlsolver)
+
+    # initial guess
+    if alg.extrapolant == :linear
+        nlsolver.z = dt * integrator.fsalfirst
+    else # :constant
+        nlsolver.z = zero(u)
+    end
+
+    nlsolver.tmp = uprev
+    nlsolver.γ = 1
+    z = nlsolve!(nlsolver, integrator, cache, repeat_step)
+    nlsolvefail(nlsolver) && return
+    u = nlsolver.tmp + z
+
+    if integrator.opts.adaptive && integrator.success_iter > 0
+        # local truncation error (LTE) bound by dt^2/2*max|y''(t)|
+        # use 2nd divided differences (DD) a la SPICE and Shampine
+
+        # TODO: check numerical stability
+        uprev2 = integrator.uprev2
+        tprev = integrator.tprev
+
+        dt1 = dt * (t + dt - tprev)
+        dt2 = (t - tprev) * (t + dt - tprev)
+        c = 7 / 12 # default correction factor in SPICE (LTE overestimated by DD)
+        r = c * dt^2 # by mean value theorem 2nd DD equals y''(s)/2 for some s
+
+        tmp = r *
+              integrator.opts.internalnorm.((u - uprev) / dt1 - (uprev - uprev2) / dt2, t)
+        atmp = calculate_residuals(tmp, uprev, u, integrator.opts.abstol,
+            integrator.opts.reltol, integrator.opts.internalnorm, t)
+        integrator.EEst = integrator.opts.internalnorm(atmp, t)
+    else
+        integrator.EEst = 1
+    end
+
+    integrator.fsallast = f(u, p, t + dt)
+
+    if integrator.opts.adaptive && integrator.differential_vars !== nothing
+        atmp = @. ifelse(!integrator.differential_vars, integrator.fsallast, false) ./
+                  integrator.opts.abstol
+        integrator.EEst += integrator.opts.internalnorm(atmp, t)
+    end
+
+    OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
+    integrator.k[1] = integrator.fsalfirst
+    integrator.k[2] = integrator.fsallast
+    integrator.u = u
+end

lib/OrdinaryDiffEqTaylorSeries/src/algorithms.jl

@@ -35,3 +35,31 @@ Base.@kwdef struct ExplicitTaylorAdaptiveOrder{StageLimiter, StepLimiter, Thread
     step_limiter!::StepLimiter = trivial_limiter!
     thread::Thread = False()
 end
+
+@doc explicit_rk_docstring(
+    "An implicit Taylor series method.",
+    "ImplicitTaylor")
+struct ImplicitTaylor2{CS, AD, F, F2, P, FDT, ST, CJ, StepLimiter} <:
+       OrdinaryDiffEqNewtonAdaptiveAlgorithm{CS, AD, FDT, ST, CJ}
+    linsolve::F
+    nlsolve::F2
+    precs::P
+    extrapolant::Symbol
+    controller::Symbol
+    step_limiter!::StepLimiter
+    autodiff::AD
+end
+
+function ImplicitTaylor2(; chunk_size = Val{0}(), autodiff = AutoForwardDiff(),
+        standardtag = Val{true}(), concrete_jac = nothing,
+        diff_type = Val{:forward}(),
+        linsolve = nothing, precs = DEFAULT_PRECS, nlsolve = NLNewton(),
+        extrapolant = :constant,
+        controller = :PI, step_limiter! = trivial_limiter!)
+    AD_choice, chunk_size, diff_type = _process_AD_choice(autodiff, chunk_size, diff_type)
+
+    ImplicitTaylor2{_unwrap_val(chunk_size), typeof(AD_choice), typeof(linsolve),
+        typeof(nlsolve), typeof(precs), diff_type, _unwrap_val(standardtag),
+        _unwrap_val(concrete_jac), typeof(step_limiter!)}(linsolve,
+        nlsolve, precs, extrapolant, controller, step_limiter!, AD_choice)
+end

lib/OrdinaryDiffEqTaylorSeries/test/runtests.jl

@@ -33,3 +33,8 @@ end
     sol = solve(prob_ode_linear, ExplicitTaylorAdaptiveOrder())
     @test length(sol) < 20
 end
+
+@testset "Implicit Taylor Tests" begin
+    sol = solve(prob_ode_linear, ImplicitTaylor2())
+    @test length(sol) < 100
+end