Skip to content

Commit 0fb3a7e

Browse files
ChrisRackauckasChrisRackauckas
authored
Merge pull request #2946 from oscardssmith/os/cleanup-JET-simplification
Os/cleanup jet simplification
2 parents 948159c + 2ca958d commit 0fb3a7e

File tree

4 files changed

+35
-69
lines changed

4 files changed

+35
-69
lines changed

‎lib/OrdinaryDiffEqFIRK/src/firk_addsteps.jl‎

Lines changed: 16 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,6 @@ function _ode_addsteps!(integrator, cache::RadauIIA3ConstantCache)
4545
# Initialize variables for JET
4646
ndw = one(eltype(u))
4747
ndwprev = one(eltype(u))
48-
θ = one(eltype(u))
4948
η = max(cache.ηold, eps(eltype(integrator.opts.reltol)))^(0.8)
5049
fail_convergence = true
5150
iter = 0
@@ -76,7 +75,7 @@ function _ode_addsteps!(integrator, cache::RadauIIA3ConstantCache)
7675
dw2 = imag(dw12)
7776

7877
# compute norm of residuals
79-
iter > 1 && (ndwprev = ndw)
78+
ndwprev = ndw
8079
atmp1 = calculate_residuals(dw1, uprev, u, atol, rtol, internalnorm, t)
8180
atmp2 = calculate_residuals(dw2, uprev, u, atol, rtol, internalnorm, t)
8281
ndw = internalnorm(atmp1, t) + internalnorm(atmp2, t)
@@ -90,6 +89,7 @@ function _ode_addsteps!(integrator, cache::RadauIIA3ConstantCache)
9089
if diverge || veryslowconvergence
9190
break
9291
end
92+
η = θ / (1 - θ)
9393
end
9494

9595
w1 = @. w1 - dw1
@@ -100,7 +100,6 @@ function _ode_addsteps!(integrator, cache::RadauIIA3ConstantCache)
100100
z2 = @. T21 * w1 + T22 * w2
101101

102102
# check stopping criterion
103-
iter > 1 && (η = θ / (1 - θ))
104103
if η * ndw < κ && (iter > 1 || iszero(ndw) || !iszero(integrator.success_iter))
105104
# Newton method converges
106105
cache.status = η < alg.fast_convergence_cutoff ? FastConvergence :
@@ -173,7 +172,6 @@ function _ode_addsteps!(integrator, cache::RadauIIA3Cache, repeat_step = false)
173172
# Initialize variables for JET
174173
ndw = one(eltype(u))
175174
ndwprev = one(eltype(u))
176-
θ = one(eltype(u))
177175
η = max(cache.ηold, eps(eltype(integrator.opts.reltol)))^(0.8)
178176
fail_convergence = true
179177
iter = 0
@@ -225,7 +223,7 @@ function _ode_addsteps!(integrator, cache::RadauIIA3Cache, repeat_step = false)
225223
dw2 = imag(dw12)
226224

227225
# compute norm of residuals
228-
iter > 1 && (ndwprev = ndw)
226+
ndwprev = ndw
229227
calculate_residuals!(atmp, dw1, uprev, u, atol, rtol, internalnorm, t)
230228
ndw1 = internalnorm(atmp, t)
231229
calculate_residuals!(atmp, dw2, uprev, u, atol, rtol, internalnorm, t)
@@ -239,6 +237,7 @@ function _ode_addsteps!(integrator, cache::RadauIIA3Cache, repeat_step = false)
239237
if diverge
240238
break
241239
end
240+
η = θ / (1 - θ)
242241
end
243242

244243
@. w1 = w1 - dw1
@@ -249,7 +248,6 @@ function _ode_addsteps!(integrator, cache::RadauIIA3Cache, repeat_step = false)
249248
@. z2 = T21 * w1 + T22 * w2
250249

251250
# check stopping criterion
252-
iter > 1 && (η = θ / (1 - θ))
253251
if η * ndw < κ && (iter > 1 || iszero(ndw) || !iszero(integrator.success_iter))
254252
# Newton method converges
255253
cache.status = η < alg.fast_convergence_cutoff ? FastConvergence :
@@ -339,8 +337,6 @@ function _ode_addsteps!(integrator, cache::RadauIIA5ConstantCache,
339337
# Newton iteration
340338
# Initialize variables for JET
341339
ndw = one(eltype(u))
342-
ndwprev = one(eltype(u))
343-
θ = one(eltype(u))
344340
η = max(cache.ηold, eps(eltype(integrator.opts.reltol)))^(0.8)
345341
fail_convergence = true
346342
iter = 0
@@ -378,7 +374,7 @@ function _ode_addsteps!(integrator, cache::RadauIIA5ConstantCache,
378374
dw3 = imag(dw23)
379375

380376
# compute norm of residuals
381-
iter > 1 && (ndwprev = ndw)
377+
ndwprev = ndw
382378
atmp1 = calculate_residuals(dw1, uprev, u, atol, rtol, internalnorm, t)
383379
atmp2 = calculate_residuals(dw2, uprev, u, atol, rtol, internalnorm, t)
384380
atmp3 = calculate_residuals(dw3, uprev, u, atol, rtol, internalnorm, t)
@@ -392,6 +388,7 @@ function _ode_addsteps!(integrator, cache::RadauIIA5ConstantCache,
392388
if diverge || veryslowconvergence
393389
break
394390
end
391+
η = θ / (1 - θ)
395392
end
396393

397394
w1 = @.. w1-dw1
@@ -404,7 +401,6 @@ function _ode_addsteps!(integrator, cache::RadauIIA5ConstantCache,
404401
z3 = @.. T31 * w1+w2 # T32 = 1, T33 = 0
405402

406403
# check stopping criterion
407-
iter > 1 && (η = θ / (1 - θ))
408404
if η * ndw < κ && (iter > 1 || iszero(ndw) || !iszero(integrator.success_iter))
409405
# Newton method converges
410406
cache.status = η < alg.fast_convergence_cutoff ? FastConvergence :
@@ -497,8 +493,6 @@ function _ode_addsteps!(integrator, cache::RadauIIA5Cache, repeat_step = false)
497493
# Newton iteration
498494
# Initialize variables for JET
499495
ndw = one(eltype(u))
500-
ndwprev = one(eltype(u))
501-
θ = one(eltype(u))
502496
η = max(cache.ηold, eps(eltype(integrator.opts.reltol)))^(0.8)
503497
fail_convergence = true
504498
iter = 0
@@ -576,7 +570,7 @@ function _ode_addsteps!(integrator, cache::RadauIIA5Cache, repeat_step = false)
576570
@.. dw3=imag(dw23)
577571

578572
# compute norm of residuals
579-
iter > 1 && (ndwprev = ndw)
573+
ndwprev = ndw
580574
calculate_residuals!(atmp, dw1, uprev, u, atol, rtol, internalnorm, t)
581575
ndw1 = internalnorm(atmp, t)
582576
calculate_residuals!(atmp, dw2, uprev, u, atol, rtol, internalnorm, t)
@@ -594,6 +588,7 @@ function _ode_addsteps!(integrator, cache::RadauIIA5Cache, repeat_step = false)
594588
if diverge || veryslowconvergence
595589
break
596590
end
591+
η = θ / (1 - θ)
597592
end
598593

599594
@.. w1=w1 - dw1
@@ -606,7 +601,6 @@ function _ode_addsteps!(integrator, cache::RadauIIA5Cache, repeat_step = false)
606601
@.. z3=T31 * w1 + w2 # T32 = 1, T33 = 0
607602

608603
# check stopping criterion
609-
iter > 1 && (η = θ / (1 - θ))
610604
if η * ndw < κ && (iter > 1 || iszero(ndw) || !iszero(integrator.success_iter))
611605
# Newton method converges
612606
cache.status = η < alg.fast_convergence_cutoff ? FastConvergence :
@@ -736,8 +730,6 @@ function _ode_addsteps!(integrator, cache::RadauIIA9ConstantCache,
736730
# Newton iteration
737731
# Initialize variables for JET
738732
ndw = one(eltype(u))
739-
ndwprev = one(eltype(u))
740-
θ = one(eltype(u))
741733
η = max(cache.ηold, eps(eltype(integrator.opts.reltol)))^(0.8)
742734
fail_convergence = true
743735
iter = 0
@@ -788,7 +780,7 @@ function _ode_addsteps!(integrator, cache::RadauIIA9ConstantCache,
788780
dw5 = imag(dw45)
789781

790782
# compute norm of residuals
791-
iter > 1 && (ndwprev = ndw)
783+
ndwprev = ndw
792784
atmp1 = calculate_residuals(dw1, uprev, u, atol, rtol, internalnorm, t)
793785
atmp2 = calculate_residuals(dw2, uprev, u, atol, rtol, internalnorm, t)
794786
atmp3 = calculate_residuals(dw3, uprev, u, atol, rtol, internalnorm, t)
@@ -807,6 +799,7 @@ function _ode_addsteps!(integrator, cache::RadauIIA9ConstantCache,
807799
if diverge || veryslowconvergence
808800
break
809801
end
802+
η = θ / (1 - θ)
810803
end
811804

812805
w1 = @.. w1-dw1
@@ -823,7 +816,6 @@ function _ode_addsteps!(integrator, cache::RadauIIA9ConstantCache,
823816
z5 = @.. T51*w1+w2+w4#= T52=1, T53=0, T54=1, T55=0 =#
824817

825818
# check stopping criterion
826-
iter > 1 && (η = θ / (1 - θ))
827819
if η * ndw < κ && (iter > 1 || iszero(ndw) || !iszero(integrator.success_iter))
828820
# Newton method converges
829821
cache.status = η < alg.fast_convergence_cutoff ? FastConvergence :
@@ -970,8 +962,6 @@ function _ode_addsteps!(integrator, cache::RadauIIA9Cache, repeat_step = false)
970962
# Newton iteration
971963
# Initialize variables for JET
972964
ndw = one(eltype(u))
973-
ndwprev = one(eltype(u))
974-
θ = one(eltype(u))
975965
η = max(cache.ηold, eps(eltype(integrator.opts.reltol)))^(0.8)
976966
fail_convergence = true
977967
iter = 0
@@ -1088,7 +1078,7 @@ function _ode_addsteps!(integrator, cache::RadauIIA9Cache, repeat_step = false)
10881078
@.. dw5=imag(dw45)
10891079

10901080
# compute norm of residuals
1091-
iter > 1 && (ndwprev = ndw)
1081+
ndwprev = ndw
10921082
calculate_residuals!(atmp, dw1, uprev, u, atol, rtol, internalnorm, t)
10931083
ndw1 = internalnorm(atmp, t)
10941084
calculate_residuals!(atmp, dw2, uprev, u, atol, rtol, internalnorm, t)
@@ -1111,6 +1101,7 @@ function _ode_addsteps!(integrator, cache::RadauIIA9Cache, repeat_step = false)
11111101
if diverge || veryslowconvergence
11121102
break
11131103
end
1104+
η = θ / (1 - θ)
11141105
end
11151106

11161107
@.. w1=w1 - dw1
@@ -1127,8 +1118,6 @@ function _ode_addsteps!(integrator, cache::RadauIIA9Cache, repeat_step = false)
11271118
@.. z5=T51 * w1 + w2 + w4#= T52=1, T53=0, T54=1, T55=0 =#
11281119

11291120
# check stopping criterion
1130-
1131-
iter > 1 && (η = θ / (1 - θ))
11321121
if η * ndw < κ && (iter > 1 || iszero(ndw) || !iszero(integrator.success_iter))
11331122
# Newton method converges
11341123
cache.status = η < alg.fast_convergence_cutoff ? FastConvergence :
@@ -1247,8 +1236,6 @@ function _ode_addstep!(integrator, cache::AdaptiveRadauConstantCache, repeat_ste
12471236
# Newton iteration
12481237
# Initialize variables for JET
12491238
ndw = one(eltype(u))
1250-
ndwprev = one(eltype(u))
1251-
θ = one(eltype(u))
12521239
η = max(cache.ηold, eps(eltype(integrator.opts.reltol)))^(0.8)
12531240
fail_convergence = true
12541241
iter = 0
@@ -1300,7 +1287,7 @@ function _ode_addstep!(integrator, cache::AdaptiveRadauConstantCache, repeat_ste
13001287
integrator.stats.nsolve += (num_stages + 1) ÷ 2
13011288

13021289
# compute norm of residuals
1303-
iter > 1 && (ndwprev = ndw)
1290+
ndwprev = ndw
13041291
ndw = 0.0
13051292
for i in 1:num_stages
13061293
ndw += internalnorm(
@@ -1318,6 +1305,7 @@ function _ode_addstep!(integrator, cache::AdaptiveRadauConstantCache, repeat_ste
13181305
if diverge || veryslowconvergence
13191306
break
13201307
end
1308+
η = θ / (1 - θ)
13211309
end
13221310

13231311
for i in 1:num_stages
@@ -1340,7 +1328,6 @@ function _ode_addstep!(integrator, cache::AdaptiveRadauConstantCache, repeat_ste
13401328
end
13411329

13421330
# check stopping criterion
1343-
iter > 1 && (η = θ / (1 - θ))
13441331
if η * ndw < κ && (iter > 1 || iszero(ndw) || !iszero(integrator.success_iter))
13451332
# Newton method converges
13461333
cache.status = η < alg.fast_convergence_cutoff ? FastConvergence :
@@ -1484,8 +1471,6 @@ function _ode_addsteps!(integrator, cache::AdaptiveRadauCache, repeat_step = fal
14841471
# Newton iteration
14851472
# Initialize variables for JET
14861473
ndw = one(eltype(u))
1487-
ndwprev = one(eltype(u))
1488-
θ = one(eltype(u))
14891474
η = max(cache.ηold, eps(eltype(integrator.opts.reltol)))^(0.8)
14901475
fail_convergence = true
14911476
iter = 0
@@ -1577,7 +1562,7 @@ function _ode_addsteps!(integrator, cache::AdaptiveRadauCache, repeat_step = fal
15771562
end
15781563

15791564
# compute norm of residuals
1580-
iter > 1 && (ndwprev = ndw)
1565+
ndwprev = ndw
15811566
calculate_residuals!(atmp, dw1, uprev, u, atol, rtol, internalnorm, t)
15821567
ndw = internalnorm(atmp, t)
15831568
for i in 2:num_stages
@@ -1595,6 +1580,7 @@ function _ode_addsteps!(integrator, cache::AdaptiveRadauCache, repeat_step = fal
15951580
if diverge || veryslowconvergence
15961581
break
15971582
end
1583+
η = θ / (1 - θ)
15981584
end
15991585

16001586
@.. w[1] = w[1] - dw1
@@ -1618,7 +1604,6 @@ function _ode_addsteps!(integrator, cache::AdaptiveRadauCache, repeat_step = fal
16181604
end
16191605

16201606
# check stopping criterion
1621-
iter > 1 && (η = θ / (1 - θ))
16221607
if η * ndw < κ && (iter > 1 || iszero(ndw) || !iszero(integrator.success_iter))
16231608
# Newton method converges
16241609
cache.status = η < alg.fast_convergence_cutoff ? FastConvergence :

0 commit comments

Comments
 (0)