Merge pull request #2640 from vyudu/rename

feat: allow solving NLLS in solve, rename SimpleIDSolve to IDSolve

Author: ChrisRackauckas

lib/ImplicitDiscreteSolve/src/ImplicitDiscreteSolve.jl

@@ -10,16 +10,16 @@ using Reexport
 @reexport using DiffEqBase
 
 """
-    SimpleIDSolve()
+    IDSolve()
 
 Simple solver for `ImplicitDiscreteSystems`. Uses `SimpleNewtonRaphson` to solve for the next state at every timestep.
 """
-struct SimpleIDSolve <: OrdinaryDiffEqAlgorithm end
+struct IDSolve <: OrdinaryDiffEqAlgorithm end
 
 include("cache.jl")
 include("solve.jl")
 include("alg_utils.jl")
 
-export SimpleIDSolve
+export IDSolve
 
 end

lib/ImplicitDiscreteSolve/src/alg_utils.jl

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

lib/ImplicitDiscreteSolve/src/cache.jl

@@ -4,35 +4,35 @@ mutable struct ImplicitDiscreteState{uType, pType, tType}
     t_next::tType
 end
 
-mutable struct SimpleIDSolveCache{uType} <: OrdinaryDiffEqMutableCache
+mutable struct IDSolveCache{uType} <: OrdinaryDiffEqMutableCache
     u::uType
     uprev::uType
     state::ImplicitDiscreteState
     prob::Union{Nothing, SciMLBase.AbstractNonlinearProblem}
 end
 
-function alg_cache(alg::SimpleIDSolve, u, rate_prototype, ::Type{uEltypeNoUnits},
+function alg_cache(alg::IDSolve, 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(isnothing(u) ? nothing : zero(u), p, t)
-    SimpleIDSolveCache(u, uprev, state, nothing)
+    IDSolveCache(u, uprev, state, nothing)
 end
 
-isdiscretecache(cache::SimpleIDSolveCache) = true
+isdiscretecache(cache::IDSolveCache) = true
 
-struct SimpleIDSolveConstantCache <: OrdinaryDiffEqConstantCache 
+struct IDSolveConstantCache <: OrdinaryDiffEqConstantCache 
     prob::Union{Nothing, SciMLBase.AbstractNonlinearProblem}
 end
 
-function alg_cache(alg::SimpleIDSolve, u, rate_prototype, ::Type{uEltypeNoUnits},
+function alg_cache(alg::IDSolve, 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(isnothing(u) ? nothing : zero(u), p, t)
-    SimpleIDSolveCache(u, uprev, state, nothing)
+    IDSolveCache(u, uprev, state, nothing)
 end
 
-get_fsalfirstlast(cache::SimpleIDSolveCache, rate_prototype) = (nothing, nothing)
+get_fsalfirstlast(cache::IDSolveCache, rate_prototype) = (nothing, nothing)

lib/ImplicitDiscreteSolve/src/solve.jl

@@ -1,5 +1,5 @@
 # Remake the nonlinear problem, then update
-function perform_step!(integrator, cache::SimpleIDSolveCache, repeat_step = false)
+function perform_step!(integrator, cache::IDSolveCache, repeat_step = false)
     @unpack alg, u, uprev, dt, t, f, p = integrator
     @unpack state, prob = cache
     state.u .= uprev
@@ -11,7 +11,7 @@ function perform_step!(integrator, cache::SimpleIDSolveCache, repeat_step = fals
     integrator.u = u
 end
 
-function initialize!(integrator, cache::SimpleIDSolveCache)
+function initialize!(integrator, cache::IDSolveCache)
     integrator.u isa AbstractVector && (cache.state.u .= integrator.u)
     cache.state.p = integrator.p
     cache.state.t_next = integrator.t
@@ -22,11 +22,13 @@ function initialize!(integrator, cache::SimpleIDSolveCache)
     else
         (u_next, p) -> f(u_next, p.u, p.p, p.t_next)
     end
+    u_len = isnothing(integrator.u) ? 0 : length(integrator.u)
+    nlls = !isnothing(f.resid_prototype) && (length(f.resid_prototype) != u_len)
 
-    prob = if isinplace(f)
-        NonlinearProblem{true}(_f, cache.state.u, cache.state)
+    prob = if nlls
+        NonlinearLeastSquaresProblem{isinplace(f)}(NonlinearFunction(_f; resid_prototype = f.resid_prototype), cache.state.u, cache.state)
     else
-        NonlinearProblem{false}(_f, cache.state.u, cache.state)
+        NonlinearProblem{isinplace(f)}(_f, cache.state.u, cache.state)
     end
     cache.prob = prob
 end
@@ -47,8 +49,18 @@ function _initialize_dae!(integrator, prob::ImplicitDiscreteProblem,
         else
             (u_next, p) -> f(u_next, p.u, p.p, p.t_next)
         end
-        prob = NonlinearProblem{isinplace(f)}(_f, u, initstate)
+
+        nlls = !isnothing(f.resid_prototype) && (length(f.resid_prototype) != length(integrator.u))
+        prob = if nlls
+            NonlinearLeastSquaresProblem{isinplace(f)}(NonlinearFunction(_f; resid_prototype = f.resid_prototype), u, initstate)
+        else
+            NonlinearProblem{isinplace(f)}(_f, u, initstate)
+        end
         sol = solve(prob, SimpleNewtonRaphson())
-        integrator.u = sol
+        if sol.retcode == ReturnCode.Success
+            integrator.u = sol
+        else
+            integrator.sol = SciMLBase.solution_new_retcode(integrator.sol, ReturnCode.InitialFailure)
+        end
     end
 end

lib/ImplicitDiscreteSolve/test/runtests.jl

@@ -3,6 +3,7 @@ using Test
 using ImplicitDiscreteSolve
 using OrdinaryDiffEqCore
 using OrdinaryDiffEqSDIRK
+using SciMLBase
 
 # Test implicit Euler using ImplicitDiscreteProblem
 @testset "Implicit Euler" begin
@@ -20,7 +21,7 @@ using OrdinaryDiffEqSDIRK
     tspan = (0., 0.5)
 
     idprob = ImplicitDiscreteProblem(f!, u0, tspan, []; dt = 0.01)
-    idsol = solve(idprob, SimpleIDSolve())
+    idsol = solve(idprob, IDSolve())
 
     oprob = ODEProblem(lotkavolterra, u0, tspan)
     osol = solve(oprob, ImplicitEuler())
@@ -43,7 +44,7 @@ using OrdinaryDiffEqSDIRK
     tspan = (0, 0.2)
 
     idprob = ImplicitDiscreteProblem(g!, u0, tspan, []; dt = 0.01)
-    idsol = solve(idprob, SimpleIDSolve())
+    idsol = solve(idprob, IDSolve())
 
     oprob = ODEProblem(ff, u0, tspan)
     osol = solve(oprob, ImplicitEuler())
@@ -60,7 +61,7 @@ end
     tsteps = 15
     u0 = [1., 3.]
     idprob = ImplicitDiscreteProblem(periodic!, u0, (0, tsteps), [])
-    integ = init(idprob, SimpleIDSolve())
+    integ = init(idprob, IDSolve())
     @test integ.u[1]^2 + integ.u[2]^2 ≈ 16
 
     for ts in 1:tsteps
@@ -77,5 +78,44 @@ end
     tsteps = 5
     u0 = nothing
     idprob = ImplicitDiscreteProblem(empty, u0, (0, tsteps), [])
-    @test_nowarn integ = init(idprob, SimpleIDSolve())
+    @test_nowarn integ = init(idprob, IDSolve())
+end
+
+@testset "Create NonlinearLeastSquaresProblem" begin
+    function over(u_next, u, p, t) 
+        [u_next[1] - 1, u_next[2] - 1, u_next[1] - u_next[2]]
+    end
+
+    tsteps = 5
+    u0 = [1., 1.]
+    idprob = ImplicitDiscreteProblem(ImplicitDiscreteFunction(over, resid_prototype = zeros(3)), u0, (0, tsteps), [])
+    integ = init(idprob, IDSolve())
+    @test integ.cache.prob isa NonlinearLeastSquaresProblem
+
+    function under(u_next, u, p, t) 
+        [u_next[1] - u_next[2] - 1] 
+    end
+    idprob = ImplicitDiscreteProblem(ImplicitDiscreteFunction(under; resid_prototype = zeros(1)), u0, (0, tsteps), [])
+    integ = init(idprob, IDSolve())
+    @test integ.cache.prob isa NonlinearLeastSquaresProblem
+
+    function full(u_next, u, p, t) 
+        [u_next[1]^2 - 3, u_next[2] - u[1]] 
+    end
+    idprob = ImplicitDiscreteProblem(ImplicitDiscreteFunction(full; resid_prototype = zeros(2)), u0, (0, tsteps), [])
+    integ = init(idprob, IDSolve())
+    @test integ.cache.prob isa NonlinearProblem
+end
+
+@testset "InitialFailure thrown" begin
+    function bad(u_next, u, p, t) 
+        [u_next[1] - u_next[2], u_next[1] - 3, u_next[2] - 4]
+    end
+
+    u0 = [3., 4.]
+    idprob = ImplicitDiscreteProblem(bad, u0, (0, 0), [])
+    integ = init(idprob, IDSolve())
+    @test check_error(integ) == ReturnCode.InitialFailure
+    sol = solve(idprob, IDSolve())
+    @test length(sol.u) == 1
 end

lib/OrdinaryDiffEqCore/src/solve.jl

@@ -563,7 +563,7 @@ function DiffEqBase.__init(
     if initialize_integrator
         if isdae || SciMLBase.has_initializeprob(prob.f) || prob isa SciMLBase.ImplicitDiscreteProblem
             DiffEqBase.initialize_dae!(integrator)
-            update_uprev!(integrator)
+            !isnothing(integrator.u) && update_uprev!(integrator)
         end
 
         if save_start