Skip to content

Commit 37dcb5a

Browse files
Merge pull request #2204 from oscardssmith/os/inf-norm-eigen-est
Switch RK methods to use an inf norm for `eigen_est`
2 parents 36ab509 + dbd9d84 commit 37dcb5a

File tree

6 files changed

+63
-162
lines changed

6 files changed

+63
-162
lines changed

src/composite_algs.jl

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ function is_stiff(integrator, alg, ntol, stol, is_stiffalg)
1515
stiffness = abs(eigen_est * dt / alg_stability_size(alg)) # `abs` here is just for safety
1616
tol = is_stiffalg ? stol : ntol
1717
os = oneunit(stiffness)
18-
bool = stiffness > os * tol
18+
bool = !(stiffness <= os * tol)
1919

2020
if !bool
2121
integrator.alg.choice_function.successive_switches += 1
@@ -111,7 +111,7 @@ current_nonstiff(current) = ifelse(current <= NUM_NONSTIFF, current, current - N
111111
function DefaultODEAlgorithm(; lazy = true, stiffalgfirst = false, kwargs...)
112112
nonstiff = (Tsit5(), Vern7(lazy = lazy))
113113
stiff = (Rosenbrock23(; kwargs...), Rodas5P(; kwargs...), FBDF(; kwargs...),
114-
FBDF(; linsolve = LinearSolve.KrylovJL_GMRES()))
114+
FBDF(; linsolve = LinearSolve.KrylovJL_GMRES(), kwargs...))
115115
AutoAlgSwitch(nonstiff, stiff; stiffalgfirst)
116116
end
117117

src/derivative_utils.jl

Lines changed: 29 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -140,7 +140,6 @@ function calc_J!(J, integrator, cache, next_step::Bool = false)
140140
if !(p isa DiffEqBase.NullParameters)
141141
uf.p = p
142142
end
143-
144143
jacobian!(J, uf, uprev, du1, integrator, jac_config)
145144
end
146145
end
@@ -294,7 +293,7 @@ SciMLBase.isinplace(::WOperator{IIP}, i) where {IIP} = IIP
294293
Base.eltype(W::WOperator) = eltype(W.J)
295294

296295
# In WOperator update_coefficients!, accept both missing u/p/t and missing dtgamma/transform and don't update them in that case.
297-
# This helps support partial updating logic used with Newton solvers.
296+
# This helps support partial updating logic used with Newton solvers.
298297
function SciMLOperators.update_coefficients!(W::WOperator,
299298
u = nothing,
300299
p = nothing,
@@ -483,16 +482,16 @@ end
483482
end
484483

485484
function jacobian2W!(
486-
W::AbstractMatrix, mass_matrix::MT, dtgamma::Number, J::AbstractMatrix,
487-
W_transform::Bool)::Nothing where {MT}
485+
W::AbstractMatrix, mass_matrix, dtgamma::Number, J::AbstractMatrix,
486+
W_transform::Bool)::Nothing
488487
# check size and dimension
489488
iijj = axes(W)
490489
@boundscheck (iijj == axes(J) && length(iijj) == 2) || _throwWJerror(W, J)
491490
mass_matrix isa UniformScaling ||
492491
@boundscheck axes(mass_matrix) == axes(W) || _throwWMerror(W, mass_matrix)
493492
@inbounds if W_transform
494493
invdtgamma = inv(dtgamma)
495-
if MT <: UniformScaling
494+
if mass_matrix isa UniformScaling
496495
copyto!(W, J)
497496
idxs = diagind(W)
498497
λ = -mass_matrix.λ
@@ -508,28 +507,9 @@ function jacobian2W!(
508507
@.. broadcast=false W=muladd(-mass_matrix, invdtgamma, J)
509508
end
510509
else
511-
if MT <: UniformScaling
510+
if mass_matrix isa UniformScaling
512511
λ = -mass_matrix.λ
513-
if W isa AbstractSparseMatrix && !(W isa SparseMatrixCSC)
514-
# This is specifically to catch the GPU sparse matrix cases
515-
# Which do not support diagonal indexing
516-
# https://github.com/JuliaGPU/CUDA.jl/issues/1395
517-
518-
Wn = nonzeros(W)
519-
Jn = nonzeros(J)
520-
521-
# I would hope to check this generically, but `CuSparseMatrixCSC` has `colPtr`
522-
# and `rowVal` while SparseMatrixCSC is colptr and rowval, and there is no
523-
# standard for checking sparsity patterns in general. So for now, write it for
524-
# the convention of CUDA.jl and handle the case of some other convention when
525-
# it comes up.
526-
527-
@assert J.colPtr == W.colPtr
528-
@assert J.rowVal == W.rowVal
529-
530-
@.. broadcast=false Wn=dtgamma * Jn
531-
W .= W + λ * I
532-
elseif W isa SparseMatrixCSC
512+
if W isa SparseMatrixCSC
533513
#=
534514
using LinearAlgebra, SparseArrays, FastBroadcast
535515
J = sparse(Diagonal(ones(4)))
@@ -553,6 +533,25 @@ function jacobian2W!(
553533
@.. broadcast=false W.nzval=dtgamma * J.nzval
554534
idxs = diagind(W)
555535
@.. broadcast=false @view(W[idxs])=@view(W[idxs]) + λ
536+
elseif W isa AbstractSparseMatrix
537+
# This is specifically to catch the GPU sparse matrix cases
538+
# Which do not support diagonal indexing
539+
# https://github.com/JuliaGPU/CUDA.jl/issues/1395
540+
541+
Wn = nonzeros(W)
542+
Jn = nonzeros(J)
543+
544+
# I would hope to check this generically, but `CuSparseMatrixCSC` has `colPtr`
545+
# and `rowVal` while SparseMatrixCSC is colptr and rowval, and there is no
546+
# standard for checking sparsity patterns in general. So for now, write it for
547+
# the convention of CUDA.jl and handle the case of some other convention when
548+
# it comes up.
549+
550+
@assert J.colPtr == W.colPtr
551+
@assert J.rowVal == W.rowVal
552+
553+
@.. broadcast=false Wn=dtgamma * Jn
554+
W .= W + λ * I
556555
else # Anything not a sparse matrix
557556
@.. broadcast=false W=dtgamma * J
558557
idxs = diagind(W)
@@ -565,16 +564,16 @@ function jacobian2W!(
565564
return nothing
566565
end
567566

568-
function jacobian2W!(W::Matrix, mass_matrix::MT, dtgamma::Number, J::Matrix,
569-
W_transform::Bool)::Nothing where {MT}
567+
function jacobian2W!(W::Matrix, mass_matrix, dtgamma::Number, J::Matrix,
568+
W_transform::Bool)::Nothing
570569
# check size and dimension
571570
iijj = axes(W)
572571
@boundscheck (iijj == axes(J) && length(iijj) == 2) || _throwWJerror(W, J)
573572
mass_matrix isa UniformScaling ||
574573
@boundscheck axes(mass_matrix) == axes(W) || _throwWMerror(W, mass_matrix)
575574
@inbounds if W_transform
576575
invdtgamma = inv(dtgamma)
577-
if MT <: UniformScaling
576+
if mass_matrix isa UniformScaling
578577
copyto!(W, J)
579578
idxs = diagind(W)
580579
λ = -mass_matrix.λ
@@ -587,7 +586,7 @@ function jacobian2W!(W::Matrix, mass_matrix::MT, dtgamma::Number, J::Matrix,
587586
end
588587
end
589588
else
590-
if MT <: UniformScaling
589+
if mass_matrix isa UniformScaling
591590
idxs = diagind(W)
592591
@inbounds @simd ivdep for i in eachindex(W)
593592
W[i] = dtgamma * J[i]

src/perform_step/explicit_rk_perform_step.jl

Lines changed: 6 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -48,10 +48,9 @@ end
4848
u = uprev + dt * accum
4949

5050
if integrator.alg isa CompositeAlgorithm
51-
# Hairer II, page 22
52-
ϱu = integrator.opts.internalnorm(kk[end] - kk[end - 1], t)
53-
ϱd = integrator.opts.internalnorm(u - u_beforefinal, t)
54-
integrator.eigen_est = ϱu / ϱd
51+
# Hairer II, page 22 modified to use Inf norm
52+
n = maximum(abs.((kk[end] .- kk[end - 1]) / (u .- u_beforefinal)))
53+
integrator.eigen_est = integrator.opts.internalnorm(n, t)
5554
end
5655

5756
if integrator.opts.adaptive
@@ -221,12 +220,9 @@ end
221220
end
222221

223222
if integrator.alg isa CompositeAlgorithm
224-
# Hairer II, page 22
225-
@.. broadcast=false utilde=kk[end] - kk[end - 1]
226-
ϱu = integrator.opts.internalnorm(utilde, t)
227-
@.. broadcast=false utilde=u - tmp
228-
ϱd = integrator.opts.internalnorm(utilde, t)
229-
integrator.eigen_est = ϱu / ϱd
223+
# Hairer II, page 22 modified to use Inf norm
224+
@.. broadcast=false utilde=abs((kk[end] - kk[end - 1]) / (u - tmp))
225+
integrator.eigen_est = integrator.opts.internalnorm(maximum(utilde), t)
230226
end
231227

232228
if integrator.opts.adaptive

src/perform_step/low_order_rk_perform_step.jl

Lines changed: 10 additions & 83 deletions
Original file line numberDiff line numberDiff line change
@@ -761,9 +761,8 @@ end
761761
integrator.stats.nf += 6
762762
if integrator.alg isa CompositeAlgorithm
763763
g7 = u
764-
# Hairer II, page 22
765-
integrator.eigen_est = integrator.opts.internalnorm(k7 - k6, t) /
766-
integrator.opts.internalnorm(g7 - g6, t)
764+
# Hairer II, page 22 modified to use the Inf norm
765+
integrator.eigen_est = integrator.opts.internalnorm(maximum(abs.(k7 .- k6) ./ (g7 .- g6)), t)
767766
end
768767
if integrator.opts.adaptive
769768
utilde = dt *
@@ -836,12 +835,9 @@ end
836835
if integrator.alg isa CompositeAlgorithm
837836
g7 = u
838837
g6 = tmp
839-
# Hairer II, page 22
840-
@.. broadcast=false thread=thread utilde=k7 - k6
841-
ϱu = integrator.opts.internalnorm(utilde, t)
842-
@.. broadcast=false thread=thread utilde=g7 - g6
843-
ϱd = integrator.opts.internalnorm(utilde, t)
844-
integrator.eigen_est = ϱu / ϱd
838+
# Hairer II, page 22 modified to use Inf norm
839+
@.. broadcast=false thread=thread utilde=abs((k7 - k6) / (g7 - g6))
840+
integrator.eigen_est = integrator.opts.internalnorm(maximum(utilde), t)
845841
end
846842
if integrator.opts.adaptive
847843
@.. broadcast=false thread=thread utilde=dt * (btilde1 * k1 + btilde2 * k2 +
@@ -889,9 +885,8 @@ end
889885
integrator.stats.nf += 6
890886
if integrator.alg isa CompositeAlgorithm
891887
g7 = u
892-
# Hairer II, page 22
893-
integrator.eigen_est = integrator.opts.internalnorm(k7 - k6, t) /
894-
integrator.opts.internalnorm(g7 - g6, t)
888+
# Hairer II, page 22 modified to use the Inf norm
889+
integrator.eigen_est = integrator.opts.internalnorm(maximum(abs.(k7 .- k6) ./ (g7 .- g6)), t)
895890
end
896891
if integrator.opts.adaptive
897892
utilde = dt *
@@ -957,12 +952,9 @@ end
957952
if integrator.alg isa CompositeAlgorithm
958953
g6 = tmp
959954
g7 = u
960-
# Hairer II, page 22
961-
@.. broadcast=false thread=thread g6=g7 - g6
962-
ϱd = integrator.opts.internalnorm(g6, t)
963-
@.. broadcast=false thread=thread tmp=k7 - k6
964-
ϱu = integrator.opts.internalnorm(tmp, t)
965-
integrator.eigen_est = (ϱu / ϱd) * oneunit(t)
955+
# Hairer II, page 22 modified to use Inf norm
956+
@.. broadcast=false thread=thread utilde=abs((k7 - k6) / (g7 - g6))
957+
integrator.eigen_est = integrator.opts.internalnorm(maximum(utilde) * oneunit(t), t)
966958
end
967959
if integrator.opts.adaptive
968960
@.. broadcast=false thread=thread utilde=dt * (btilde1 * k1 + btilde3 * k3 +
@@ -987,71 +979,6 @@ end
987979
return nothing
988980
end
989981

990-
#=
991-
@muladd function perform_step!(integrator, cache::DP5Cache, repeat_step=false)
992-
@unpack t,dt,uprev,u,f,p = integrator
993-
uidx = eachindex(integrator.uprev)
994-
@unpack a21,a31,a32,a41,a42,a43,a51,a52,a53,a54,a61,a62,a63,a64,a65,a71,a73,a74,a75,a76,btilde1,btilde3,btilde4,btilde5,btilde6,btilde7,c1,c2,c3,c4,c5,c6 = cache.tab
995-
@unpack k1,k2,k3,k4,k5,k6,k7,dense_tmp3,dense_tmp4,update,bspl,utilde,tmp,atmp = cache
996-
@unpack d1,d3,d4,d5,d6,d7 = cache.tab
997-
a = dt*a21
998-
@tight_loop_macros for i in uidx
999-
@inbounds tmp[i] = uprev[i]+a*k1[i]
1000-
end
1001-
f(k2, tmp, p, t+c1*dt)
1002-
@tight_loop_macros for i in uidx
1003-
@inbounds tmp[i] = uprev[i]+dt*(a31*k1[i]+a32*k2[i])
1004-
end
1005-
f(k3, tmp, p, t+c2*dt)
1006-
@tight_loop_macros for i in uidx
1007-
@inbounds tmp[i] = uprev[i]+dt*(a41*k1[i]+a42*k2[i]+a43*k3[i])
1008-
end
1009-
f(k4, tmp, p, t+c3*dt)
1010-
@tight_loop_macros for i in uidx
1011-
@inbounds tmp[i] = uprev[i]+dt*(a51*k1[i]+a52*k2[i]+a53*k3[i]+a54*k4[i])
1012-
end
1013-
f(k5, tmp, p, t+c4*dt)
1014-
@tight_loop_macros for i in uidx
1015-
@inbounds tmp[i] = uprev[i]+dt*(a61*k1[i]+a62*k2[i]+a63*k3[i]+a64*k4[i]+a65*k5[i])
1016-
end
1017-
f(k6, tmp, p, t+dt)
1018-
@tight_loop_macros for i in uidx
1019-
@inbounds update[i] = a71*k1[i]+a73*k3[i]+a74*k4[i]+a75*k5[i]+a76*k6[i]
1020-
@inbounds u[i] = uprev[i]+dt*update[i]
1021-
end
1022-
f(k7, u, p, t+dt)
1023-
integrator.stats.nf += 6
1024-
if integrator.alg isa CompositeAlgorithm
1025-
g6 = tmp
1026-
g7 = u
1027-
# Hairer II, page 22
1028-
ϱu, ϱd = zero(eltype(k7))^2, zero(eltype(g7))^2
1029-
@inbounds for i in eachindex(k7)
1030-
ϱu += (k7[i] - k6[i])^2
1031-
ϱd += (g7[i] - g6[i])^2
1032-
end
1033-
integrator.eigen_est = sqrt(ϱu/ϱd)*oneunit(t)
1034-
end
1035-
if integrator.opts.adaptive
1036-
@tight_loop_macros for i in uidx
1037-
@inbounds utilde[i] = dt*(btilde1*k1[i] + btilde3*k3[i] + btilde4*k4[i] + btilde5*k5[i] + btilde6*k6[i] + btilde7*k7[i])
1038-
end
1039-
calculate_residuals!(atmp, utilde, uprev, u, integrator.opts.abstol, integrator.opts.reltol,integrator.opts.internalnorm,t)
1040-
integrator.EEst = integrator.opts.internalnorm(atmp,t)
1041-
end
1042-
if integrator.opts.calck
1043-
@tight_loop_macros for i in uidx
1044-
#integrator.k[4] == k5
1045-
@inbounds integrator.k[4][i] = d1*k1[i]+d3*k3[i]+d4*k4[i]+d5*k5[i]+d6*k6[i]+d7*k7[i]
1046-
#bspl == k3
1047-
@inbounds bspl[i] = k1[i] - update[i]
1048-
# k6 === integrator.k[3] === k2
1049-
@inbounds integrator.k[3][i] = update[i] - k7[i] - bspl[i]
1050-
end
1051-
end
1052-
end
1053-
=#
1054-
1055982
function initialize!(integrator, cache::KYK2014DGSSPRK_3S2_ConstantCache)
1056983
integrator.fsalfirst = integrator.f(integrator.uprev, integrator.p, integrator.t) # Pre-start fsal
1057984
integrator.stats.nf += 1

src/perform_step/verner_rk_perform_step.jl

Lines changed: 12 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -40,9 +40,7 @@ end
4040
integrator.stats.nf += 8
4141
if integrator.alg isa CompositeAlgorithm
4242
g9 = u
43-
ϱu = integrator.opts.internalnorm(k9 - k8, t)
44-
ϱd = integrator.opts.internalnorm(g9 - g8, t)
45-
integrator.eigen_est = ϱu / ϱd
43+
integrator.eigen_est = integrator.opts.internalnorm(maximum(abs.((k9 .- k8) ./ (g9 .- g8))), t)
4644
end
4745
if integrator.opts.adaptive
4846
utilde = dt *
@@ -163,11 +161,8 @@ end
163161
if integrator.alg isa CompositeAlgorithm
164162
g9 = u
165163
g8 = tmp
166-
@.. broadcast=false thread=thread rtmp=k9 - k8
167-
ϱu = integrator.opts.internalnorm(rtmp, t)
168-
@.. broadcast=false thread=thread utilde=g9 - g8
169-
ϱd = integrator.opts.internalnorm(utilde, t)
170-
integrator.eigen_est = ϱu / ϱd
164+
@.. broadcast=false thread=thread rtmp=abs((k9 - k8) / (g9 - g8))
165+
integrator.eigen_est = integrator.opts.internalnorm(maximum(rtmp), t)
171166
end
172167
if integrator.opts.adaptive
173168
@.. broadcast=false thread=thread utilde=dt * (btilde1 * k1 + btilde4 * k4 +
@@ -252,9 +247,7 @@ end
252247
integrator.stats.nf += 10
253248
u = uprev + dt * (b1 * k1 + b4 * k4 + b5 * k5 + b6 * k6 + b7 * k7 + b8 * k8 + b9 * k9)
254249
if integrator.alg isa CompositeAlgorithm
255-
ϱu = integrator.opts.internalnorm(k10 - k9, t)
256-
ϱd = integrator.opts.internalnorm(g10 - g9, t)
257-
integrator.eigen_est = ϱu / ϱd
250+
integrator.eigen_est = integrator.opts.internalnorm(maximum(abs.((k10 .- k9) ./ (g10 .- g9))), t)
258251
end
259252
if integrator.opts.adaptive
260253
utilde = dt *
@@ -415,11 +408,8 @@ end
415408
if integrator.alg isa CompositeAlgorithm
416409
g10 = u
417410
g9 = tmp
418-
@.. broadcast=false thread=thread rtmp=k10 - k9
419-
ϱu = integrator.opts.internalnorm(rtmp, t)
420-
@.. broadcast=false thread=thread utilde=g10 - g9
421-
ϱd = integrator.opts.internalnorm(utilde, t)
422-
integrator.eigen_est = ϱu / ϱd
411+
@.. broadcast=false thread=thread rtmp=abs((k10 - k9) / (g10 - g9))
412+
integrator.eigen_est = integrator.opts.internalnorm(maximum(rtmp), t)
423413
end
424414
if integrator.opts.adaptive
425415
@.. broadcast=false thread=thread utilde=dt * (btilde1 * k1 + btilde4 * k4 +
@@ -547,9 +537,7 @@ end
547537
dt * (b1 * k1 + b6 * k6 + b7 * k7 + b8 * k8 + b9 * k9 + b10 * k10 + b11 * k11 +
548538
b12 * k12)
549539
if integrator.alg isa CompositeAlgorithm
550-
ϱu = integrator.opts.internalnorm(k13 - k12, t)
551-
ϱd = integrator.opts.internalnorm(g13 - g12, t)
552-
integrator.eigen_est = ϱu / ϱd
540+
integrator.eigen_est = integrator.opts.internalnorm(maximum(abs.((k13 .- k12) ./ (g13 .- g12))), t)
553541
end
554542
if integrator.opts.adaptive
555543
utilde = dt *
@@ -739,11 +727,8 @@ end
739727
if integrator.alg isa CompositeAlgorithm
740728
g13 = u
741729
g12 = tmp
742-
@.. broadcast=false thread=thread rtmp=k13 - k12
743-
ϱu = integrator.opts.internalnorm(rtmp, t)
744-
@.. broadcast=false thread=thread utilde=g13 - g12
745-
ϱd = integrator.opts.internalnorm(utilde, t)
746-
integrator.eigen_est = ϱu / ϱd
730+
@.. broadcast=false thread=thread rtmp = abs((k13 - k12) / (g13 - g12))
731+
integrator.eigen_est = integrator.opts.internalnorm(maximum(rtmp), t)
747732
end
748733
@.. broadcast=false thread=thread u=uprev +
749734
dt *
@@ -918,9 +903,7 @@ end
918903
dt * (b1 * k1 + b8 * k8 + b9 * k9 + b10 * k10 + b11 * k11 + b12 * k12 + b13 * k13 +
919904
b14 * k14 + b15 * k15)
920905
if integrator.alg isa CompositeAlgorithm
921-
ϱu = integrator.opts.internalnorm(k16 - k15, t)
922-
ϱd = integrator.opts.internalnorm(g16 - g15, t)
923-
integrator.eigen_est = ϱu / ϱd
906+
integrator.eigen_est = integrator.opts.internalnorm(maximum(abs.((k16 .- k15) ./ (g16 .- g15))), t)
924907
end
925908
if integrator.opts.adaptive
926909
utilde = dt * (btilde1 * k1 + btilde8 * k8 + btilde9 * k9 + btilde10 * k10 +
@@ -1148,11 +1131,8 @@ end
11481131
if integrator.alg isa CompositeAlgorithm
11491132
g16 = u
11501133
g15 = tmp
1151-
@.. broadcast=false thread=thread rtmp=k16 - k15
1152-
ϱu = integrator.opts.internalnorm(rtmp, t)
1153-
@.. broadcast=false thread=thread utilde=g16 - g15
1154-
ϱd = integrator.opts.internalnorm(utilde, t)
1155-
integrator.eigen_est = ϱu / ϱd
1134+
@.. broadcast=false thread=thread rtmp=abs((k16 - k15) / (g16 - g15))
1135+
integrator.eigen_est = integrator.opts.internalnorm(maximum(rtmp), t)
11561136
end
11571137
@.. broadcast=false thread=thread u=uprev +
11581138
dt *

0 commit comments

Comments
 (0)