Merge pull request #2626 from vyudu/ids

feat: ImplicitDiscreteSolve

Author: ChrisRackauckas

.github/workflows/CI.yml

@@ -56,6 +56,8 @@ jobs:
           - OrdinaryDiffEqTsit5
           - OrdinaryDiffEqVerner
 
+          - ImplicitDiscreteSolve
+
         version:
           - 'lts'
           - '1'

lib/ImplicitDiscreteSolve/Project.toml

@@ -0,0 +1,31 @@
+name = "ImplicitDiscreteSolve"
+uuid = "3263718b-31ed-49cf-8a0f-35a466e8af96"
+authors = ["vyudu <[email protected]>"]
+version = "0.1.0"
+
+[deps]
+DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
+OrdinaryDiffEqCore = "bbf590c4-e513-4bbe-9b18-05decba2e5d8"
+Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
+SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
+SimpleNonlinearSolve = "727e6d20-b764-4bd8-a329-72de5adea6c7"
+SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
+UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
+
+[compat]
+DiffEqBase = "6.164.1"
+OrdinaryDiffEqCore = "1.18.1"
+OrdinaryDiffEqSDIRK = "1.2.0"
+Reexport = "1.2.2"
+SciMLBase = "2.74.1"
+SimpleNonlinearSolve = "2.1.0"
+SymbolicIndexingInterface = "0.3.38"
+Test = "1.10.0"
+UnPack = "1.0.2"
+
+[extras]
+OrdinaryDiffEqSDIRK = "2d112036-d095-4a1e-ab9a-08536f3ecdbf"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[targets]
+test = ["OrdinaryDiffEqSDIRK", "Test"]

lib/ImplicitDiscreteSolve/src/ImplicitDiscreteSolve.jl

@@ -0,0 +1,25 @@
+module ImplicitDiscreteSolve
+
+using SciMLBase
+using SimpleNonlinearSolve
+using UnPack
+using SymbolicIndexingInterface: parameter_symbols
+import OrdinaryDiffEqCore: OrdinaryDiffEqAlgorithm, alg_cache, OrdinaryDiffEqMutableCache, OrdinaryDiffEqConstantCache, get_fsalfirstlast, isfsal, initialize!, perform_step!, isdiscretecache, isdiscretealg, alg_order, beta2_default, beta1_default, dt_required, _initialize_dae!, DefaultInit, BrownFullBasicInit, OverrideInit
+
+using Reexport
+@reexport using DiffEqBase
+
+"""
+    SimpleIDSolve()
+
+Simple solver for `ImplicitDiscreteSystems`. Uses `SimpleNewtonRaphson` to solve for the next state at every timestep.
+"""
+struct SimpleIDSolve <: OrdinaryDiffEqAlgorithm end
+
+include("cache.jl")
+include("solve.jl")
+include("alg_utils.jl")
+
+export SimpleIDSolve
+
+end

lib/ImplicitDiscreteSolve/src/alg_utils.jl

@@ -0,0 +1,19 @@
+function SciMLBase.isautodifferentiable(alg::SimpleIDSolve)
+    true
+end
+function SciMLBase.allows_arbitrary_number_types(alg::SimpleIDSolve)
+    true
+end
+function SciMLBase.allowscomplex(alg::SimpleIDSolve)
+    true
+end
+
+SciMLBase.isdiscrete(alg::SimpleIDSolve) = true
+
+isfsal(alg::SimpleIDSolve) = false
+alg_order(alg::SimpleIDSolve) = 0
+beta2_default(alg::SimpleIDSolve) = 0
+beta1_default(alg::SimpleIDSolve, beta2) = 0
+
+dt_required(alg::SimpleIDSolve) = false
+isdiscretealg(alg::SimpleIDSolve) = true

lib/ImplicitDiscreteSolve/src/cache.jl

@@ -0,0 +1,38 @@
+mutable struct ImplicitDiscreteState{uType, pType, tType}
+    u::Vector{uType}
+    p::pType
+    t_next::tType
+end
+
+mutable struct SimpleIDSolveCache{uType} <: OrdinaryDiffEqMutableCache
+    u::uType
+    uprev::uType
+    state::ImplicitDiscreteState
+    prob::Union{Nothing, SciMLBase.AbstractNonlinearProblem}
+end
+
+function alg_cache(alg::SimpleIDSolve, u, rate_prototype, ::Type{uEltypeNoUnits},
+        ::Type{uBottomEltypeNoUnits}, ::Type{tTypeNoUnits}, uprev, uprev2, f, t,
+        dt, reltol, p, calck,
+        ::Val{true}) where {uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits}
+
+    state = ImplicitDiscreteState(similar(u), p, t)
+    SimpleIDSolveCache(u, uprev, state, nothing)
+end
+
+isdiscretecache(cache::SimpleIDSolveCache) = true
+
+struct SimpleIDSolveConstantCache <: OrdinaryDiffEqConstantCache 
+    prob::Union{Nothing, SciMLBase.AbstractNonlinearProblem}
+end
+
+function alg_cache(alg::SimpleIDSolve, u, rate_prototype, ::Type{uEltypeNoUnits},
+        ::Type{uBottomEltypeNoUnits}, ::Type{tTypeNoUnits}, uprev, uprev2, f, t,
+        dt, reltol, p, calck,
+        ::Val{false}) where {uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits}
+
+    state = ImplicitDiscreteState(similar(u), p, t)
+    SimpleIDSolveCache(u, uprev, state, nothing)
+end
+
+get_fsalfirstlast(cache::SimpleIDSolveCache, rate_prototype) = (nothing, nothing)

lib/ImplicitDiscreteSolve/src/solve.jl

@@ -0,0 +1,53 @@
+# Remake the nonlinear problem, then update
+function perform_step!(integrator, cache::SimpleIDSolveCache, repeat_step = false)
+    @unpack alg, u, uprev, dt, t, f, p = integrator
+    @unpack state, prob = cache
+    state.u .= uprev
+    state.t_next = t
+    prob = remake(prob, p = state)
+
+    u = solve(prob, SimpleNewtonRaphson())
+    integrator.sol = SciMLBase.solution_new_retcode(integrator.sol, u.retcode)
+    integrator.u = u
+end
+
+function initialize!(integrator, cache::SimpleIDSolveCache)
+    cache.state.u .= integrator.u
+    cache.state.p = integrator.p
+    cache.state.t_next = integrator.t
+    f = integrator.f
+
+    _f = if isinplace(f)
+        (resid, u_next, p) -> f(resid, u_next, p.u, p.p, p.t_next)
+    else
+        (u_next, p) -> f(u_next, p.u, p.p, p.t_next)
+    end
+
+    prob = if isinplace(f)
+        NonlinearProblem{true}(_f, cache.state.u, cache.state)
+    else
+        NonlinearProblem{false}(_f, cache.state.u, cache.state)
+    end
+    cache.prob = prob
+end
+
+function _initialize_dae!(integrator, prob::ImplicitDiscreteProblem,
+        alg::DefaultInit, x::Union{Val{true}, Val{false}})
+    atol = one(eltype(prob.u0)) * 1e-12
+    if SciMLBase.has_initializeprob(prob.f)
+        _initialize_dae!(integrator, prob,
+                         OverrideInit(atol), x)
+    else
+        @unpack u, p, t, f = integrator
+        initstate = ImplicitDiscreteState(u, p, t)
+
+        _f = if isinplace(f)
+            (resid, u_next, p) -> f(resid, u_next, p.u, p.p, p.t_next)
+        else
+            (u_next, p) -> f(u_next, p.u, p.p, p.t_next)
+        end
+        prob = NonlinearProblem{isinplace(f)}(_f, u, initstate)
+        sol = solve(prob, SimpleNewtonRaphson())
+        integrator.u = sol
+    end
+end

lib/ImplicitDiscreteSolve/test/runtests.jl

@@ -0,0 +1,71 @@
+#runtests
+using Test
+using ImplicitDiscreteSolve
+using OrdinaryDiffEqCore
+using OrdinaryDiffEqSDIRK
+
+# Test implicit Euler using ImplicitDiscreteProblem
+@testset "Implicit Euler" begin
+    function lotkavolterra(u, p, t)
+        [1.5*u[1] - u[1]*u[2], -3.0*u[2] + u[1]*u[2]]
+    end
+
+    function f!(resid, u_next, u, p, t)
+        lv = lotkavolterra(u_next, p, t)
+        resid[1] = u_next[1] - u[1] - 0.01*lv[1]
+        resid[2] = u_next[2] - u[2] - 0.01*lv[2]
+        nothing
+    end
+    u0 = [1., 1.]
+    tspan = (0., 0.5)
+
+    idprob = ImplicitDiscreteProblem(f!, u0, tspan, []; dt = 0.01)
+    idsol = solve(idprob, SimpleIDSolve())
+
+    oprob = ODEProblem(lotkavolterra, u0, tspan)
+    osol = solve(oprob, ImplicitEuler())
+
+    @test isapprox(idsol[end], osol[end], atol = 0.1)
+
+    ### free-fall
+    # y, dy
+    function ff(u, p, t) 
+        [u[2], -9.8]
+    end
+
+    function g!(resid, u_next, u, p, t) 
+        f = ff(u_next, p, t) 
+        resid[1] = u_next[1] - u[1] - 0.01*f[1]
+        resid[2] = u_next[2] - u[2] - 0.01*f[2]
+        nothing
+    end
+    u0 = [10., 0.]
+    tspan = (0, 0.2)
+
+    idprob = ImplicitDiscreteProblem(g!, u0, tspan, []; dt = 0.01)
+    idsol = solve(idprob, SimpleIDSolve())
+
+    oprob = ODEProblem(ff, u0, tspan)
+    osol = solve(oprob, ImplicitEuler())
+
+    @test isapprox(idsol[end], osol[end], atol = 0.1)
+end
+
+@testset "Solver initializes" begin
+    function periodic!(resid, u_next, u, p, t) 
+        resid[1] = u_next[1] - u[1] - sin(t*π/4)
+        resid[2] = 16 - u_next[2]^2 - u_next[1]^2 
+    end
+
+    tsteps = 15
+    u0 = [1., 3.]
+    idprob = ImplicitDiscreteProblem(periodic!, u0, (0, tsteps), [])
+    integ = init(idprob, SimpleIDSolve())
+    @test integ.u[1]^2 + integ.u[2]^2 ≈ 16
+
+    for ts in 1:tsteps
+        step!(integ)
+        @show integ.u
+        @test integ.u[1]^2 + integ.u[2]^2 ≈ 16
+    end
+end

test/runtests.jl

@@ -26,7 +26,7 @@ end
 #Start Test Script
 
 @time begin
-    if contains(GROUP, "OrdinaryDiffEq")
+    if contains(GROUP, "OrdinaryDiffEq") || GROUP == "ImplicitDiscreteSolve"
         Pkg.develop(path = "../lib/$GROUP")
         Pkg.test(GROUP)
     elseif GROUP == "All" || GROUP == "InterfaceI" || GROUP == "Interface"