refactor: SimpleImplicitDiscreteSolve

Author: vyudu

lib/SimpleImplicitDiscreteSolve/Project.toml

@@ -5,23 +5,17 @@ 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.11.0"
-UnPack = "1.0.2"
 
 [extras]
 OrdinaryDiffEqSDIRK = "2d112036-d095-4a1e-ab9a-08536f3ecdbf"

lib/SimpleImplicitDiscreteSolve/src/SimpleImplicitDiscreteSolve.jl

@@ -2,10 +2,6 @@ module SimpleImplicitDiscreteSolve
 
 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
 
@@ -14,11 +10,85 @@ using Reexport
 
 Simple solver for `ImplicitDiscreteSystems`. Uses `SimpleNewtonRaphson` to solve for the next state at every timestep.
 """
-struct SimpleIDSolve <: OrdinaryDiffEqAlgorithm end
+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
 
-include("cache.jl")
-include("solve.jl")
-include("alg_utils.jl")
+    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]
+
+    if save_everystep && save_start
+        us = Vector{typeof(u0)}(undef, length(ts))
+        us[1] = u0
+    elseif save_everystep
+        us = Vector{typeof(u0)}(undef, length(ts) - 1)
+    elseif save_start
+        us = Vector{typeof(u0)}(undef, 2)
+        us[1] = u0
+    else
+        us = Vector{typeof(u0)}(undef, 1) # for interface compatibility
+    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
 

lib/SimpleImplicitDiscreteSolve/src/alg_utils.jl

@@ -1,7 +1,6 @@
 #runtests
 using Test
 using SimpleImplicitDiscreteSolve
-using OrdinaryDiffEqCore
 using OrdinaryDiffEqSDIRK
 
 # Test implicit Euler using ImplicitDiscreteProblem
@@ -19,13 +18,13 @@ using OrdinaryDiffEqSDIRK
     u0 = [1., 1.]
     tspan = (0., 0.5)
 
-    idprob = ImplicitDiscreteProblem(f!, u0, tspan, []; dt = 0.01)
-    idsol = solve(idprob, SimpleIDSolve())
+    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], osol[end], atol = 0.1)
+    @test isapprox(idsol[end-1], osol[end], atol = 0.1)
 
     ### free-fall
     # y, dy
@@ -40,15 +39,15 @@ using OrdinaryDiffEqSDIRK
         nothing
     end
     u0 = [10., 0.]
-    tspan = (0, 0.2)
+    tspan = (0, 0.5)
 
-    idprob = ImplicitDiscreteProblem(g!, u0, tspan, []; dt = 0.01)
-    idsol = solve(idprob, SimpleIDSolve())
+    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], osol[end], atol = 0.1)
+    @test isapprox(idsol[end-1], osol[end], atol = 0.1)
 end
 
 @testset "Solver initializes" begin
@@ -65,7 +64,6 @@ end
 
     for ts in 1:tsteps
         step!(integ)
-        @show integ.u
         @test integ.u[1]^2 + integ.u[2]^2 ≈ 16
     end
 end