Pass the cache to nlsolve! in PDIRK44's out-of-place branch (#4352)

Author: singhharsh1708

lib/OrdinaryDiffEqNonlinearSolve/Project.toml

@@ -1,7 +1,7 @@
 name = "OrdinaryDiffEqNonlinearSolve"
 uuid = "127b3ac7-2247-4354-8eb6-78cf4e7c58e8"
 authors = ["Chris Rackauckas <accounts@chrisrackauckas.com>", "Yingbo Ma <mayingbo5@gmail.com>"]
-version = "2.9.0"
+version = "2.9.1"
 
 [deps]
 CommonSolve = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2"

lib/OrdinaryDiffEqNonlinearSolve/src/nlsolve.jl

@@ -112,6 +112,16 @@ function nlsolve!(
     # Initialize γW for JET
     γW = one(integrator.dt)
     if isnewton(nlsolver)
+        # Checking the type, not just `nothing`: passing something else entirely
+        # in this slot went unnoticed because `update_W!` happens not to read it
+        # on the out-of-place path.
+        cache isa Union{OrdinaryDiffEqCore.OrdinaryDiffEqCache, Nothing} ||
+            throw(
+            ArgumentError(
+                "`nlsolve!` expects the integrator cache in its third argument, got a " *
+                    "$(typeof(cache))"
+            )
+        )
         cache === nothing &&
             throw(ArgumentError("cache is not passed to `nlsolve!` when using NLNewton"))
         if nlsolver.method === DIRK

lib/OrdinaryDiffEqPDIRK/Project.toml

@@ -1,7 +1,7 @@
 name = "OrdinaryDiffEqPDIRK"
 uuid = "5dd0a6cf-3d4b-4314-aa06-06d4e299bc89"
 authors = ["ParamThakkar123 <paramthakkar864@gmail.com>"]
-version = "2.4.0"
+version = "2.4.1"
 
 [deps]
 FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898"

lib/OrdinaryDiffEqPDIRK/src/pdirk_perform_step.jl

@@ -45,28 +45,28 @@ function initialize!(integrator, cache::PDIRK44ConstantCache) end
         _nlsolver.γ = γs[1]
         _nlsolver.c = cs[1]
         markfirststage!(_nlsolver)
-        k11 = nlsolve!(_nlsolver, integrator, γs[1] * dt, repeat_step)
+        k11 = nlsolve!(_nlsolver, integrator, cache, repeat_step)
         nlsolvefail(_nlsolver) && return
         _nlsolver.z = zero(u)
         _nlsolver.tmp = uprev
         _nlsolver.γ = γs[2]
         _nlsolver.c = cs[2]
         markfirststage!(_nlsolver)
-        k12 = nlsolve!(_nlsolver, integrator, γs[2] * dt, repeat_step)
+        k12 = nlsolve!(_nlsolver, integrator, cache, repeat_step)
         nlsolvefail(_nlsolver) && return
         _nlsolver.z = zero(u)
         _nlsolver.tmp = uprev + α1[1] * k11 + α2[1] * k12
         _nlsolver.γ = γs[1]
         _nlsolver.c = cs[3]
         markfirststage!(_nlsolver)
-        k21 = nlsolve!(_nlsolver, integrator, γs[1] * dt, repeat_step)
+        k21 = nlsolve!(_nlsolver, integrator, cache, repeat_step)
         nlsolvefail(_nlsolver) && return
         _nlsolver.z = zero(u)
         _nlsolver.tmp = uprev + α1[2] * k11 + α2[2] * k12
         _nlsolver.γ = γs[2]
         _nlsolver.c = cs[4]
         markfirststage!(_nlsolver)
-        k22 = nlsolve!(_nlsolver, integrator, γs[2] * dt, repeat_step)
+        k22 = nlsolve!(_nlsolver, integrator, cache, repeat_step)
         nlsolvefail(_nlsolver) && return
         integrator.u = uprev + b1 * k11 + b2 * k21 + b3 * k12 + b4 * k22
     end

lib/OrdinaryDiffEqPDIRK/test/nlsolve_argument_tests.jl

@@ -0,0 +1,28 @@
+using OrdinaryDiffEqPDIRK
+using OrdinaryDiffEqCore
+using OrdinaryDiffEqNonlinearSolve: nlsolve!
+using SciMLBase
+using Test
+
+# The out-of-place branch used to pass `γ * dt` where `nlsolve!` takes the
+# integrator cache. Nothing complained: the `NLNewton` guard only rejected
+# `nothing`, and `update_W!` happens not to read that argument on the
+# out-of-place path. Both halves are covered here.
+@testset "PDIRK44 out-of-place agrees with in-place" begin
+    oop = ODEProblem((u, p, t) -> -u * (1 + 0.1u), 1.0, (0.0, 1.0))
+    iip = ODEProblem((du, u, p, t) -> (du[1] = -u[1] * (1 + 0.1u[1]); nothing), [1.0], (0.0, 1.0))
+    for threading in (false, true)
+        a = solve(oop, PDIRK44(; threading); dt = 0.05, adaptive = false)
+        b = solve(iip, PDIRK44(; threading); dt = 0.05, adaptive = false)
+        @test SciMLBase.successful_retcode(a)
+        @test a.u[end] ≈ b.u[end][1] rtol = 1.0e-10
+    end
+end
+
+@testset "nlsolve! rejects a third argument that is not a cache" begin
+    prob = ODEProblem((du, u, p, t) -> (du[1] = -u[1]; nothing), [1.0], (0.0, 1.0))
+    integrator = init(prob, PDIRK44(threading = false); dt = 0.1, adaptive = false)
+    nlsolver = first(integrator.cache.nlsolver)
+    @test_throws ArgumentError nlsolve!(nlsolver, integrator, 0.5, false)
+    @test_throws ArgumentError nlsolve!(nlsolver, integrator, nothing, false)
+end

lib/OrdinaryDiffEqPDIRK/test/runtests.jl

@@ -18,3 +18,4 @@ if (TEST_GROUP == "QA" || TEST_GROUP == "ALL") && isempty(VERSION.prerelease)
 end
 
 @time @safetestset "Convergence Tests" include("pdirk_convergence_tests.jl")
+@time @safetestset "nlsolve! Arguments" include("nlsolve_argument_tests.jl")