Add initialization

Author: vyudu

lib/SimpleImplicitDiscreteSolve/Project.toml

@@ -20,11 +20,12 @@ 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"
-SimpleNonlinearSolve = "727e6d20-b764-4bd8-a329-72de5adea6c7"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["OrdinaryDiffEqSDIRK", "SimpleNonlinearSolve"]
+test = ["OrdinaryDiffEqSDIRK", "Test"]

lib/SimpleImplicitDiscreteSolve/src/alg_utils.jl

@@ -17,18 +17,3 @@ beta1_default(alg::SimpleIDSolve, beta2) = 0
 
 dt_required(alg::SimpleIDSolve) = false
 isdiscretealg(alg::SimpleIDSolve) = true
-
-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)
-    elseif !applicable(_initialize_dae!, integrator, prob,
-        BrownFullBasicInit(atol), x)
-        error("`OrdinaryDiffEqNonlinearSolve` is not loaded, which is required for the default initialization algorithm (`BrownFullBasicInit` or `ShampineCollocationInit`). To solve this problem, either do `using OrdinaryDiffEqNonlinearSolve` or pass `initializealg = CheckInit()` to the `solve` function. This second option requires consistent `u0`.")
-    else
-        _initialize_dae!(integrator, prob,
-            BrownFullBasicInit(atol), x)
-    end
-end

lib/SimpleImplicitDiscreteSolve/src/solve.jl

@@ -30,3 +30,52 @@ function initialize!(integrator, cache::SimpleIDSolveCache)
     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
+
+#### TODO: Implement real algorithm
+function _initialize_dae!(integrator, prob::ImplicitDiscreteProblem, alg::BrownFullBasicInit, isinplace::Val{true}) 
+    @unpack p, t, f = integrator
+    u0 = integrator.u
+    
+    nlequation! = (out, u, p) -> begin 
+        f(out, u, u0, p, t)
+    end
+
+    nlfunc = NonlinearFunction(nlequation!; jac_prototype = f.jac_prototype)
+    nlprob = NonlinearProblem(nlfunc, ifelse.(differential_vars, du, u), p)
+    nlsol = solve(nlprob, nlsolve; abstol = alg.abstol, reltol = integrator.opts.reltol)
+
+    @. du = ifelse(differential_vars, nlsol.u, du)
+    @. u = ifelse(differential_vars, u, nlsol.u)
+
+    recursivecopy!(integrator.uprev, integrator.u)
+    if alg_extrapolates(integrator.alg)
+        recursivecopy!(integrator.uprev2, integrator.uprev)
+    end
+
+    if nlsol.retcode != ReturnCode.Success
+        integrator.sol = SciMLBase.solution_new_retcode(integrator.sol,
+            ReturnCode.InitialFailure)
+    end
+    return
+end

lib/SimpleImplicitDiscreteSolve/test/runtests.jl

@@ -1,4 +1,5 @@
 #runtests
+using Test
 using SimpleImplicitDiscreteSolve
 using OrdinaryDiffEqCore
 using OrdinaryDiffEqSDIRK
@@ -24,7 +25,7 @@ using OrdinaryDiffEqSDIRK
     oprob = ODEProblem(lotkavolterra, u0, tspan)
     osol = solve(oprob, ImplicitEuler())
 
-    @test isapprox(idsol[end], osol[end], atol = 0.01)
+    @test isapprox(idsol[end], osol[end], atol = 0.1)
 
     ### free-fall
     # y, dy
@@ -38,16 +39,33 @@ using OrdinaryDiffEqSDIRK
         resid[2] = u_next[2] - u[2] - 0.01*f[2]
         nothing
     end
-    u0 = [100., 3.]
+    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.01)
+    @test isapprox(idsol[end], osol[end], atol = 0.1)
 end
 
-@testset "Solve respects initialization" begin
+@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