|
| 1 | +# Optimized tolerance checking that avoids allocations for scalar abstol |
| 2 | +@inline function check_dae_tolerance(internalnorm, err, abstol, t) |
| 3 | + if abstol isa Number |
| 4 | + return internalnorm(err, t) / abstol <= 1 |
| 5 | + else |
| 6 | + return internalnorm(err ./ abstol, t) <= 1 |
| 7 | + end |
| 8 | +end |
| 9 | + |
1 | 10 | function default_nlsolve( |
2 | 11 | ::Nothing, isinplace::Val{true}, u, ::AbstractNonlinearProblem, autodiff = false) |
3 | 12 | FastShortcutNonlinearPolyalg(; |
@@ -57,7 +66,8 @@ function _initialize_dae!(integrator, prob::ODEProblem, alg::ShampineCollocation |
57 | 66 | f(tmp, u0, p, t) |
58 | 67 | tmp .= ArrayInterface.restructure(tmp, algebraic_eqs .* _vec(tmp)) |
59 | 68 |
|
60 | | - integrator.opts.internalnorm(tmp ./ integrator.opts.abstol, t) <= 1 && return |
| 69 | + check_dae_tolerance(integrator.opts.internalnorm, tmp, integrator.opts.abstol, t) && |
| 70 | + return |
61 | 71 |
|
62 | 72 | if isdefined(integrator.cache, :nlsolver) && !isnothing(alg.nlsolve) |
63 | 73 | # backward Euler |
@@ -171,7 +181,8 @@ function _initialize_dae!(integrator, prob::ODEProblem, alg::ShampineCollocation |
171 | 181 | du = f(u0, p, t) |
172 | 182 | resid = _vec(du)[algebraic_eqs] |
173 | 183 |
|
174 | | - integrator.opts.internalnorm(resid ./ integrator.opts.abstol, t) <= 1 && return |
| 184 | + check_dae_tolerance(integrator.opts.internalnorm, resid, integrator.opts.abstol, t) && |
| 185 | + return |
175 | 186 |
|
176 | 187 | if isdefined(integrator.cache, :nlsolver) && !isnothing(alg.nlsolve) |
177 | 188 | # backward Euler |
@@ -239,7 +250,8 @@ function _initialize_dae!(integrator, prob::DAEProblem, |
239 | 250 | dt = t != 0 ? min(t / 1000, dtmax / 10) : dtmax / 10 # Haven't implemented norm reduction |
240 | 251 |
|
241 | 252 | f(resid, integrator.du, u0, p, t) |
242 | | - integrator.opts.internalnorm(resid ./ integrator.opts.abstol, t) <= 1 && return |
| 253 | + check_dae_tolerance(integrator.opts.internalnorm, resid, integrator.opts.abstol, t) && |
| 254 | + return |
243 | 255 |
|
244 | 256 | # _du and _u should be non-dual since NonlinearSolve does not differentiate the solver |
245 | 257 | # These non-dual values are thus used to make the caches |
@@ -320,7 +332,8 @@ function _initialize_dae!(integrator, prob::DAEProblem, |
320 | 332 | nlequation = (u, _) -> nlequation_oop(u) |
321 | 333 |
|
322 | 334 | resid = f(integrator.du, u0, p, t) |
323 | | - integrator.opts.internalnorm(resid ./ integrator.opts.abstol, t) <= 1 && return |
| 335 | + check_dae_tolerance(integrator.opts.internalnorm, resid, integrator.opts.abstol, t) && |
| 336 | + return |
324 | 337 |
|
325 | 338 | jac = if isnothing(f.jac) |
326 | 339 | f.jac |
@@ -385,7 +398,7 @@ function _initialize_dae!(integrator, prob::ODEProblem, |
385 | 398 |
|
386 | 399 | tmp .= ArrayInterface.restructure(tmp, algebraic_eqs .* _vec(tmp)) |
387 | 400 |
|
388 | | - integrator.opts.internalnorm(tmp ./ alg.abstol, t) <= 1 && return |
| 401 | + check_dae_tolerance(integrator.opts.internalnorm, tmp, alg.abstol, t) && return |
389 | 402 | alg_u = @view u[algebraic_vars] |
390 | 403 |
|
391 | 404 | # These non-dual values are thus used to make the caches |
@@ -464,7 +477,7 @@ function _initialize_dae!(integrator, prob::ODEProblem, |
464 | 477 | du = f(u0, p, t) |
465 | 478 | resid = _vec(du)[algebraic_eqs] |
466 | 479 |
|
467 | | - integrator.opts.internalnorm(resid ./ alg.abstol, t) <= 1 && return |
| 480 | + check_dae_tolerance(integrator.opts.internalnorm, resid, alg.abstol, t) && return |
468 | 481 |
|
469 | 482 | isAD = alg_autodiff(integrator.alg) isa AutoForwardDiff |
470 | 483 | if isAD |
@@ -543,7 +556,7 @@ function _initialize_dae!(integrator, prob::DAEProblem, |
543 | 556 | normtmp = get_tmp_cache(integrator)[1] |
544 | 557 | f(normtmp, du, u, p, t) |
545 | 558 |
|
546 | | - if integrator.opts.internalnorm(normtmp ./ alg.abstol, t) <= 1 |
| 559 | + if check_dae_tolerance(integrator.opts.internalnorm, normtmp, alg.abstol, t) |
547 | 560 | return |
548 | 561 | elseif differential_vars === nothing |
549 | 562 | error("differential_vars must be set for DAE initialization to occur. Either set consistent initial conditions, differential_vars, or use a different initialization algorithm.") |
@@ -604,8 +617,7 @@ function _initialize_dae!(integrator, prob::DAEProblem, |
604 | 617 | @unpack p, t, f = integrator |
605 | 618 | differential_vars = prob.differential_vars |
606 | 619 |
|
607 | | - if integrator.opts.internalnorm(f(integrator.du, integrator.u, p, t) ./ alg.abstol, t) <= |
608 | | - 1 |
| 620 | + if check_dae_tolerance(integrator.opts.internalnorm, f(integrator.du, integrator.u, p, t), alg.abstol, t) |
609 | 621 | return |
610 | 622 | elseif differential_vars === nothing |
611 | 623 | error("differential_vars must be set for DAE initialization to occur. Either set consistent initial conditions, differential_vars, or use a different initialization algorithm.") |
|
0 commit comments