Merge pull request #2868 from ChrisRackauckas-Claude/fix-vector-abstol-dae-2820

Fix DAE mass matrix initialization with vector abstol (fixes #2820)

Author: ChrisRackauckas

lib/OrdinaryDiffEqNonlinearSolve/src/initialize_dae.jl

@@ -1,3 +1,21 @@
+# Optimized tolerance checking that avoids allocations
+@inline function check_dae_tolerance(integrator, err, abstol, t, ::Val{true})
+    if abstol isa Number
+        return integrator.opts.internalnorm(err, t) / abstol <= 1
+    else
+        @. err = err / abstol  # Safe for in-place functions
+        return integrator.opts.internalnorm(err, t) <= 1
+    end
+end
+
+@inline function check_dae_tolerance(integrator, err, abstol, t, ::Val{false})
+    if abstol isa Number
+        return integrator.opts.internalnorm(err, t) / abstol <= 1
+    else
+        return integrator.opts.internalnorm(err ./ abstol, t) <= 1  # Allocates for out-of-place
+    end
+end
+
 function default_nlsolve(
         ::Nothing, isinplace::Val{true}, u, ::AbstractNonlinearProblem, autodiff = false)
     FastShortcutNonlinearPolyalg(;
@@ -57,22 +75,24 @@ function _initialize_dae!(integrator, prob::ODEProblem, alg::ShampineCollocation
     f(tmp, u0, p, t)
     tmp .= ArrayInterface.restructure(tmp, algebraic_eqs .* _vec(tmp))
 
-    integrator.opts.internalnorm(tmp, t) <= integrator.opts.abstol && return
+    check_dae_tolerance(integrator, tmp, integrator.opts.abstol, t, isinplace) && return
 
     if isdefined(integrator.cache, :nlsolver) && !isnothing(alg.nlsolve)
         # backward Euler
         nlsolver = integrator.cache.nlsolver
-        oldγ, oldc, oldmethod, olddt = nlsolver.γ, nlsolver.c, nlsolver.method,
+        oldγ, oldc, oldmethod,
+        olddt = nlsolver.γ, nlsolver.c, nlsolver.method,
         integrator.dt
         nlsolver.tmp .= integrator.uprev
         nlsolver.γ, nlsolver.c = 1, 1
         nlsolver.method = DIRK
         integrator.dt = dt
         z = nlsolve!(nlsolver, integrator, integrator.cache)
-        nlsolver.γ, nlsolver.c, nlsolver.method, integrator.dt = oldγ, oldc, oldmethod,
+        nlsolver.γ, nlsolver.c, nlsolver.method,
+        integrator.dt = oldγ, oldc, oldmethod,
         olddt
         failed = nlsolvefail(nlsolver)
-        @.. broadcast=false integrator.u=integrator.uprev + z
+        @.. broadcast=false integrator.u=integrator.uprev+z
     else
 
         # _u0 should be non-dual since NonlinearSolve does not differentiate the solver
@@ -169,22 +189,24 @@ function _initialize_dae!(integrator, prob::ODEProblem, alg::ShampineCollocation
     du = f(u0, p, t)
     resid = _vec(du)[algebraic_eqs]
 
-    integrator.opts.internalnorm(resid, t) <= integrator.opts.abstol && return
+    check_dae_tolerance(integrator, resid, integrator.opts.abstol, t, isinplace) && return
 
     if isdefined(integrator.cache, :nlsolver) && !isnothing(alg.nlsolve)
         # backward Euler
         nlsolver = integrator.cache.nlsolver
-        oldγ, oldc, oldmethod, olddt = nlsolver.γ, nlsolver.c, nlsolver.method,
+        oldγ, oldc, oldmethod,
+        olddt = nlsolver.γ, nlsolver.c, nlsolver.method,
         integrator.dt
         nlsolver.tmp .= integrator.uprev
         nlsolver.γ, nlsolver.c = 1, 1
         nlsolver.method = DIRK
         integrator.dt = dt
         z = nlsolve!(nlsolver, integrator, integrator.cache)
-        nlsolver.γ, nlsolver.c, nlsolver.method, integrator.dt = oldγ, oldc, oldmethod,
+        nlsolver.γ, nlsolver.c, nlsolver.method,
+        integrator.dt = oldγ, oldc, oldmethod,
         olddt
         failed = nlsolvefail(nlsolver)
-        @.. broadcast=false integrator.u=integrator.uprev + z
+        @.. broadcast=false integrator.u=integrator.uprev+z
     else
         nlequation_oop = @closure (u, _) -> begin
             update_coefficients!(M, u, p, t)
@@ -235,7 +257,7 @@ function _initialize_dae!(integrator, prob::DAEProblem,
     dt = t != 0 ? min(t / 1000, dtmax / 10) : dtmax / 10 # Haven't implemented norm reduction
 
     f(resid, integrator.du, u0, p, t)
-    integrator.opts.internalnorm(resid, t) <= integrator.opts.abstol && return
+    check_dae_tolerance(integrator, resid, integrator.opts.abstol, t, isinplace) && return
 
     # _du and _u should be non-dual since NonlinearSolve does not differentiate the solver
     # These non-dual values are thus used to make the caches
@@ -316,7 +338,7 @@ function _initialize_dae!(integrator, prob::DAEProblem,
     nlequation = (u, _) -> nlequation_oop(u)
 
     resid = f(integrator.du, u0, p, t)
-    integrator.opts.internalnorm(resid, t) <= integrator.opts.abstol && return
+    check_dae_tolerance(integrator, resid, integrator.opts.abstol, t, isinplace) && return
 
     jac = if isnothing(f.jac)
         f.jac
@@ -381,7 +403,7 @@ function _initialize_dae!(integrator, prob::ODEProblem,
 
     tmp .= ArrayInterface.restructure(tmp, algebraic_eqs .* _vec(tmp))
 
-    integrator.opts.internalnorm(tmp, t) <= alg.abstol && return
+    check_dae_tolerance(integrator, tmp, alg.abstol, t, isinplace) && return
     alg_u = @view u[algebraic_vars]
 
     # These non-dual values are thus used to make the caches
@@ -460,7 +482,7 @@ function _initialize_dae!(integrator, prob::ODEProblem,
     du = f(u0, p, t)
     resid = _vec(du)[algebraic_eqs]
 
-    integrator.opts.internalnorm(resid, t) <= alg.abstol && return
+    check_dae_tolerance(integrator, resid, alg.abstol, t, isinplace) && return
 
     isAD = alg_autodiff(integrator.alg) isa AutoForwardDiff
     if isAD
@@ -539,7 +561,7 @@ function _initialize_dae!(integrator, prob::DAEProblem,
     normtmp = get_tmp_cache(integrator)[1]
     f(normtmp, du, u, p, t)
 
-    if integrator.opts.internalnorm(normtmp, t) <= alg.abstol
+    if check_dae_tolerance(integrator, normtmp, alg.abstol, t, isinplace)
         return
     elseif differential_vars === nothing
         error("differential_vars must be set for DAE initialization to occur. Either set consistent initial conditions, differential_vars, or use a different initialization algorithm.")
@@ -600,7 +622,8 @@ function _initialize_dae!(integrator, prob::DAEProblem,
     @unpack p, t, f = integrator
     differential_vars = prob.differential_vars
 
-    if integrator.opts.internalnorm(f(integrator.du, integrator.u, p, t), t) <= alg.abstol
+    if check_dae_tolerance(
+        integrator, f(integrator.du, integrator.u, p, t), alg.abstol, t, isinplace)
         return
     elseif differential_vars === nothing
         error("differential_vars must be set for DAE initialization to occur. Either set consistent initial conditions, differential_vars, or use a different initialization algorithm.")