feat: SimpleImplicitDiscreteSolve

Author: vyudu

lib/SimpleImplicitDiscreteSolve/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 SciML
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

lib/SimpleImplicitDiscreteSolve/Project.toml

@@ -0,0 +1,27 @@
+name = "SimpleImplicitDiscreteSolve"
+uuid = "8b67ef88-54bd-43ff-aca0-8be02588656a"
+authors = ["vyudu <[email protected]>"]
+version = "0.1.0"
+
+[deps]
+DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
+Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
+SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
+SimpleNonlinearSolve = "727e6d20-b764-4bd8-a329-72de5adea6c7"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+
+[compat]
+DiffEqBase = "6.164.1"
+OrdinaryDiffEqSDIRK = "1.2.0"
+Reexport = "1.2.2"
+SciMLBase = "2.74.1"
+SimpleNonlinearSolve = "2.1.0"
+StaticArrays = "1.9.13"
+Test = "1.10"
+
+[extras]
+OrdinaryDiffEqSDIRK = "2d112036-d095-4a1e-ab9a-08536f3ecdbf"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[targets]
+test = ["OrdinaryDiffEqSDIRK", "Test"]

lib/SimpleImplicitDiscreteSolve/src/SimpleImplicitDiscreteSolve.jl

@@ -0,0 +1,96 @@
+module SimpleImplicitDiscreteSolve
+
+using SciMLBase
+using SimpleNonlinearSolve
+using Reexport
+using StaticArrays
+@reexport using DiffEqBase
+
+"""
+    SimpleIDSolve()
+
+Simple solver for `ImplicitDiscreteSystems`. Uses `SimpleNewtonRaphson` to solve for the next state at every timestep.
+"""
+struct SimpleIDSolve <: SciMLBase.AbstractODEAlgorithm end
+
+function DiffEqBase.__init(prob::ImplicitDiscreteProblem, alg::SimpleIDSolve; dt = 1)
+    u0 = prob.u0
+    p = prob.p
+    f = prob.f
+    t = prob.tspan[1]
+
+    nlf = isinplace(f) ? (out, u, p) -> f(out, u, u0, p, t) : (u, p) -> f(u, u0, p, t)
+    prob = NonlinearProblem{isinplace(f)}(nlf, u0, p)
+    sol = solve(prob, SimpleNewtonRaphson())
+    sol, (sol.retcode != ReturnCode.Success)
+end
+
+function DiffEqBase.solve(prob::ImplicitDiscreteProblem, alg::SimpleIDSolve;
+        dt = 1,
+        save_everystep = true,
+        save_start = true,
+        adaptive = false,
+        dense = false,
+        save_end = true,
+        kwargs...)
+    @assert !adaptive
+    @assert !dense
+    (initsol, initfail) = DiffEqBase.__init(prob, alg; dt)
+    if initfail
+        sol = DiffEqBase.build_solution(prob, alg, prob.tspan[1], u0, k = nothing,
+            stats = nothing, calculate_error = false)
+        return SciMLBase.solution_new_retcode(sol, ReturnCode.InitialFailure)
+    end
+
+    u0 = initsol.u
+    tspan = prob.tspan
+    f = prob.f
+    p = prob.p
+    t = tspan[1]
+    tf = prob.tspan[2]
+    ts = tspan[1]:dt:tspan[2]
+
+    l = save_everystep ? length(ts) - 1 : 1
+    save_start && (l = l + 1)
+    u0type = typeof(u0)
+    us = u0type <: StaticArray ? MVector{l, u0type}(undef) : Vector{u0type}(undef, l)
+
+    if save_start
+        us[1] = u0
+    end
+
+    u = u0
+    convfail = false
+    for i in 2:length(ts)
+        uprev = u
+        t = ts[i]
+        nlf = isinplace(f) ? (out, u, p) -> f(out, u, uprev, p, t) :
+              (u, p) -> f(u, uprev, p, t)
+        nlprob = NonlinearProblem{isinplace(f)}(nlf, uprev, p)
+        nlsol = solve(nlprob, SimpleNewtonRaphson())
+        u = nlsol.u
+        save_everystep && (us[i] = u)
+        convfail = (nlsol.retcode != ReturnCode.Success)
+
+        if convfail
+            sol = DiffEqBase.build_solution(prob, alg, ts[1:i], us[1:i], k = nothing,
+                stats = nothing, calculate_error = false)
+            sol = SciMLBase.solution_new_retcode(sol, ReturnCode.ConvergenceFailure)
+            return sol
+        end
+    end
+
+    !save_everystep && save_end && (us[end] = u)
+    sol = DiffEqBase.build_solution(prob, alg, ts, us,
+        k = nothing, stats = nothing,
+        calculate_error = false)
+
+    DiffEqBase.has_analytic(prob.f) &&
+        DiffEqBase.calculate_solution_errors!(
+            sol; timeseries_errors = true, dense_errors = false)
+    sol
+end
+
+export SimpleIDSolve
+
+end

lib/SimpleImplicitDiscreteSolve/test/runtests.jl

@@ -0,0 +1,71 @@
+#runtests
+using Test
+using SimpleImplicitDiscreteSolve
+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.0, 1.0]
+    tspan = (0.0, 0.5)
+
+    idprob = ImplicitDiscreteProblem(f!, u0, tspan, [])
+    idsol = solve(idprob, SimpleIDSolve(), dt = 0.01)
+
+    oprob = ODEProblem(lotkavolterra, u0, tspan)
+    osol = solve(oprob, ImplicitEuler())
+
+    @test isapprox(idsol[end - 1], 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, 0.0]
+    tspan = (0, 0.5)
+
+    idprob = ImplicitDiscreteProblem(g!, u0, tspan, [])
+    idsol = solve(idprob, SimpleIDSolve(); dt = 0.01)
+
+    oprob = ODEProblem(ff, u0, tspan)
+    osol = solve(oprob, ImplicitEuler())
+
+    @test isapprox(idsol[end - 1], 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.0, 3.0]
+    idprob = ImplicitDiscreteProblem(periodic!, u0, (0, tsteps), [])
+    initsol, initfail = DiffEqBase.__init(idprob, SimpleIDSolve())
+    @test initsol.u[1]^2 + initsol.u[2]^2 ≈ 16
+
+    idsol = solve(idprob, SimpleIDSolve())
+
+    for ts in 1:tsteps
+        step = idsol.u[ts]
+        @test step[1]^2 + step[2]^2 ≈ 16
+    end
+end