-
-
Notifications
You must be signed in to change notification settings - Fork 219
/
Copy pathgeneric_rosenbrock.jl
1850 lines (1659 loc) · 73.9 KB
/
generic_rosenbrock.jl
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
#! format: off
abstract type RosenbrockTableau{T,T2} end
struct RosenbrockFixedTableau{T,T2}<:RosenbrockTableau{T,T2}
a::Array{T,2}
C::Array{T,2}
b::Array{T,1}
gamma::T2
d::Array{T,1}
c::Array{T2,1}
end
struct RosenbrockAdaptiveTableau{T,T2}<:RosenbrockTableau{T,T2}
a::Array{T,2}
C::Array{T,2}
b::Array{T,1}
btilde::Array{T,1}
gamma::T2
d::Array{T,1}
c::Array{T2,1}
end
"""
@_bitarray2boolarray RosenbrockTableau(tab.a.!=0,...)
Transform BitArray (in the form of `xs.!=0` ) into 1D-Array of Bools by
`[i for i in xs.!=0]` to satisfy the type constraint of RosenbrockTableau
"""
macro _bitarray2boolarray(expr)
args=[:([i for i in $arg]) for arg in expr.args[2:end]]
args[end-2]=:(tab.gamma!=0)
esc(:($(expr.args[1])($(args...))))
end
"""
_masktab(tab)
Convert normal tableau into a dummy tableau consisting of Bools. We use dummy tableaus
where we only care about whether values in the tableau are zeros.
"""
_masktab(tab::RosenbrockFixedTableau)=@_bitarray2boolarray RosenbrockFixedTableau(tab.a.!=0,tab.C.!=0,tab.b.!=0,tab.gamma!=0,tab.d.!=0,tab.c.!=0)
_masktab(tab::RosenbrockAdaptiveTableau)=@_bitarray2boolarray RosenbrockAdaptiveTableau(tab.a.!=0,tab.C.!=0,tab.b.!=0,tab.btilde.!=0,tab.gamma!=0,tab.d.!=0,tab.c.!=0)
"""
_common_nonzero_vals(tab::RosenbrockTableau)
Return the common nonzero symbols in the tableau. Typical return value:
`[[:a21,:a31,:a32],[:C21,:C31,:C32],[:b1,:b2,:b3],:gamma,[:d1,:d2,:d3],[:c1,:c2,:c3]]`
"""
function _common_nonzero_vals(tab::RosenbrockTableau)
nzvals=[]
push!(nzvals,[Symbol(:a,ind[1],ind[2]) for ind in findall(!iszero,tab.a)])
push!(nzvals,[Symbol(:C,ind[1],ind[2]) for ind in findall(!iszero,tab.C)])
push!(nzvals,[Symbol(:b,ind) for ind in findall(!iszero,tab.b)])
push!(nzvals,:gamma)
push!(nzvals,[Symbol(:d,ind) for ind in findall(!iszero,tab.d)])
push!(nzvals,[Symbol(:c,ind) for ind in findall(!iszero,tab.c)])
nzvals
end
"""
_nonzero_vals(tab::RosenbrockFixedTableau)
Return all the nonzero symbols in the tableau. Typical return value:
`[:a21,:a31,:a32,:C21,:C31,:C32,:b1,:b2,:b3,:gamma,:d1,:d2,:d3,:c1,:c2,:c3]`
"""
function _nonzero_vals(tab::RosenbrockFixedTableau)
nzvals=_common_nonzero_vals(tab)
vcat(nzvals...)
end
"""
_nonzero_vals(tab::RosenbrockAdaptiveTableau)
Typical return value:
`[:a21,:a31,:a32,:C21,:C31,:C32,:b1,:b2,:b3,:btilde1,:btilde2,:btilde3,:gamma,:d1,:d2,:d3,:c1,:c2,:c3]`
"""
function _nonzero_vals(tab::RosenbrockAdaptiveTableau)
nzvals=_common_nonzero_vals(tab)
push!(nzvals,[Symbol(:btilde,ind) for ind in findall(!iszero,tab.btilde)])
vcat(nzvals...)
end
"""
_push_assigns!(valexprs,inds,name,type)
Insert a series of field statements like `[:(c2::T2),:(c3::T2)]` into the array `valexprs`.
# Arguments
- `valexpr::Array{Expr,1}`: the array to be inserted
- `inds`: an iterator that gives indices
- `name::Symbol`: the prefix name of the values
- `type::Symbol`: type in the statements
"""
function _push_assigns!(valexprs,inds,name,type::Symbol)
for ind in inds
push!(valexprs,:($(Symbol(name,"$(Tuple(ind)...)"))::$type))
end
end
"""
gen_tableau_struct(tab::RosenbrockTableau,tabstructname::Symbol)
Generate the tableau struct expression from a given tableau emulating those in
`tableaus/rosenbrock_tableaus.jl`. The statements of `aij`,`Cij` and `ci` are generated
according to the nonzero values of `tab.a`,`tab.C` and `tab.c` while others are generated
from their indices. One may choose to pass in a dummy tableau with type `<:RosenbrockTalbeau{Bool,Bool}`
to fully control the tableau struct.
"""
function gen_tableau_struct(tab::RosenbrockTableau,tabstructname::Symbol)
valexprs=Array{Expr,1}()
_push_assigns!(valexprs,findall(!iszero,tab.a),:a,:T)
_push_assigns!(valexprs,findall(!iszero,tab.C),:C,:T)
_push_assigns!(valexprs,eachindex(tab.b),:b,:T)
if typeof(tab)<:RosenbrockAdaptiveTableau
_push_assigns!(valexprs,eachindex(tab.btilde),:btilde,:T)
end
push!(valexprs,:(gamma::T2))
_push_assigns!(valexprs,eachindex(tab.d),:d,:T)
_push_assigns!(valexprs,findall(!iszero,tab.c),:c,:T2)
quote struct $tabstructname{T,T2}
$(valexprs...)
end
end
end
"""
gen_tableau(tab::RosenbrockTableau,tabstructexpr::Expr,tabname::Symbol)
Generate the tableau function expression emulating those in `tableaus/rosenbrock_tableaus.jl`.
It takes in the tableau struct expression (generated by gen_tableau_struct(...) or written by hand)
to make sure the actual values of the tableau are organized in the right order.
"""
function gen_tableau(tab::RosenbrockTableau,tabstructexpr::Expr,tabname::Symbol)
@capture(tabstructexpr, struct T_ fields__ end) || error("incorrect tableau expression")
tabstructname = namify(T)
valsym2tabdict=Dict("a"=>tab.a,"C"=>tab.C,"gamma"=>tab.gamma,"c"=>tab.c,"d"=>tab.d,"b"=>tab.b)
if typeof(tab)<:RosenbrockAdaptiveTableau
valsym2tabdict["btilde"]=tab.btilde
end
pattern=r"^([a-zA-Z]+)([1-9]{0,2})$"
assignexprs = Expr[]
valsyms = Symbol[]
for field in fields
if @capture(field, valsym_Symbol::valtype_)
push!(valsyms, valsym)
m = match(pattern, String(valsym))
val = valsym2tabdict[m[1]][(parse(Int, i) for i in m[2])...]
push!(assignexprs, :($valsym = convert($valtype, $val)))
end
end
quote function $tabname(T, T2)
$(assignexprs...)
$tabstructname($(valsyms...))
end
end
end
"""
gen_cache_struct(tab::RosenbrockTableau,cachename::Symbol,constcachename::Symbol)
Generate cache struct expression emulating those in `caches/rosenbrock_caches.jl`.
The length of k1,k2,... in the mutable cache struct is determined by the length of `tab.b`
because in the end of Rosenbrock's method, we have: `y_{n+1}=y_n+ki*bi`.
"""
function gen_cache_struct(tab::RosenbrockTableau,cachename::Symbol,constcachename::Symbol)
kstype=[:($(Symbol(:k,i))::rateType) for i in 1:length(tab.b)]
constcacheexpr=quote struct $constcachename{TF,UF,Tab,JType,WType,F} <: OrdinaryDiffEqConstantCache
tf::TF
uf::UF
tab::Tab
J::JType
W::WType
linsolve::F
end
end
cacheexpr=quote
@cache mutable struct $cachename{uType,rateType,uNoUnitsType,JType,WType,TabType,TFType,UFType,F,JCType,GCType} <: GenericRosenbrockMutableCache
u::uType
uprev::uType
du::rateType
du1::rateType
du2::rateType
$(kstype...)
fsalfirst::rateType
fsallast::rateType
dT::rateType
J::JType
W::WType
tmp::rateType
atmp::uNoUnitsType
weight::uNoUnitsType
tab::TabType
tf::TFType
uf::UFType
linsolve_tmp::rateType
linsolve::F
jac_config::JCType
grad_config::GCType
end
end
constcacheexpr,cacheexpr
end
"""
gen_algcache(cacheexpr::Expr,constcachename::Symbol,algname::Symbol,tabname::Symbol)
Generate expressions for `alg_cache(...)` emulating those in `caches/rosenbrock_caches.jl`.
"""
function gen_algcache(cacheexpr::Expr,constcachename::Symbol,algname::Symbol,tabname::Symbol)
@capture(cacheexpr, @cache mutable struct T_ fields__ end) || error("incorrect cache expression")
cachename = namify(T)
ksinit = Expr[]
valsyms = Symbol[]
for field in fields
if @capture(field, valsym_Symbol::valtype_)
push!(valsyms, valsym)
if match(r"^k[1-9]+$", String(valsym)) !== nothing
push!(ksinit, :($valsym = zero(rate_prototype)))
end
end
end
quote
function alg_cache(alg::$algname,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Val{false})
tf = TimeDerivativeWrapper(f,u,p)
uf = UDerivativeWrapper(f,t,p)
J,W = build_J_W(alg,u,uprev,p,t,dt,f,uEltypeNoUnits,Val(false))
$constcachename(tf,uf,$tabname(constvalue(uBottomEltypeNoUnits),constvalue(tTypeNoUnits)),J,W,nothing)
end
function alg_cache(alg::$algname,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Val{true})
du = zero(rate_prototype)
du1 = zero(rate_prototype)
du2 = zero(rate_prototype)
$(ksinit...)
fsalfirst = zero(rate_prototype)
fsallast = zero(rate_prototype)
dT = zero(rate_prototype)
J,W = build_J_W(alg,u,uprev,p,t,dt,f,uEltypeNoUnits,Val(true))
tmp = zero(rate_prototype)
atmp = similar(u, uEltypeNoUnits)
weight = similar(u, uEltypeNoUnits)
tab = $tabname(constvalue(uBottomEltypeNoUnits),constvalue(tTypeNoUnits))
tf = TimeGradientWrapper(f,uprev,p)
uf = UJacobianWrapper(f,t,p)
linsolve_tmp = zero(rate_prototype)
linprob = LinearProblem(W,_vec(linsolve_tmp); u0=_vec(tmp))
linsolve = init(linprob,alg.linsolve,alias_A=true,alias_b=true,
Pl = LinearSolve.InvPreconditioner(Diagonal(_vec(weight))),
Pr = Diagonal(_vec(weight)))
grad_config = build_grad_config(alg,f,tf,du1,t)
jac_config = build_jac_config(alg,f,uf,du1,uprev,u,tmp,du2)
$cachename($(valsyms...))
end
end
end
"""
gen_initialize(cachename::Symbol,constcachename::Symbol)
Generate expressions for `initialize!(...)` in `perform_step/rosenbrock_perform_step.jl`.
It only generates a default version of `initialize!` which support 3rd-order Hermite interpolation.
"""
function gen_initialize(cachename::Symbol,constcachename::Symbol)
quote
function initialize!(integrator, cache::$constcachename)
integrator.kshortsize = 2
integrator.k = typeof(integrator.k)(undef, integrator.kshortsize)
integrator.fsalfirst = integrator.f(integrator.uprev, integrator.p, integrator.t)
OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
# Avoid undefined entries if k is an array of arrays
integrator.fsallast = zero(integrator.fsalfirst)
integrator.k[1] = integrator.fsalfirst
integrator.k[2] = integrator.fsallast
end
function initialize!(integrator, cache::$cachename)
integrator.kshortsize = 2
@unpack fsalfirst,fsallast = cache
integrator.fsalfirst = fsalfirst
integrator.fsallast = fsallast
resize!(integrator.k, integrator.kshortsize)
integrator.k .= [fsalfirst,fsallast]
integrator.f(integrator.fsalfirst, integrator.uprev, integrator.p, integrator.t)
OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
end
end
end
"""
gen_constant_perform_step(tabmask::RosenbrockTableau{Bool,Bool},cachename::Symbol,n_normalstep::Int,specialstepexpr=:nothing)
Generate non-inplace version of `perform_step!` expression emulating those in `perform_step/rosenbrock_perform_step.jl`.
The `perform_step!` function calculates `k1,k2,k3,...` defined by `(-W)ki=f(u+aij*kj,t+ci*dt)+di*dt*dT+Cij*kj*dt, i=1,2,...,n_normalstep`
and then gives the result by `y_{n+1}=y_n+ki*bi`. Terms with 0s (according to tabmask) are skipped in the expressions.
Special steps can be added before calculating `y_{n+1}`. The non-inplace `perform_step!` assumes the mass_matrix==I.
"""
function gen_constant_perform_step(tabmask::RosenbrockTableau{Bool,Bool},cachename::Symbol,n_normalstep::Int,specialstepexpr=:nothing)
unpacktabexpr=:(@unpack ()=cache.tab)
unpacktabexpr.args[3].args[1].args=_nonzero_vals(tabmask)
dtCijexprs=[:($(Symbol(:dtC,Cind[1],Cind[2]))=$(Symbol(:C,Cind[1],Cind[2]))/dt) for Cind in findall(!iszero,tabmask.C)]
dtdiexprs=[:($(Symbol(:dtd,dind))=dt*$(Symbol(:d,dind))) for dind in eachindex(tabmask.d)]
iterexprs=[]
for i in 1:n_normalstep
aijkj=[:($(Symbol(:a,i+1,j))*$(Symbol(:k,j))) for j in findall(!iszero,tabmask.a[i+1,:])]
Cijkj=[:($(Symbol(:dtC,i+1,j))*$(Symbol(:k,j))) for j in findall(!iszero,tabmask.C[i+1,:])]
push!(iterexprs,
quote
$(Symbol(:k,i)) = _reshape(W\-_vec(linsolve_tmp), axes(uprev))
integrator.stats.nsolve += 1
u=+(uprev,$(aijkj...))
du = f(u, p, t+$(Symbol(:c,i+1))*dt)
OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
if mass_matrix === I
linsolve_tmp=+(du,$(Symbol(:dtd,i+1))*dT,$(Cijkj...))
else
linsolve_tmp=du+$(Symbol(:dtd,i+1))*dT+mass_matrix*(+($(Cijkj...)))
end
end)
end
push!(iterexprs,specialstepexpr)
n=length(tabmask.b)
biki=[:($(Symbol(:b,i))*$(Symbol(:k,i))) for i in 1:n]
push!(iterexprs,
quote
$(Symbol(:k,n))=_reshape(W\-_vec(linsolve_tmp), axes(uprev))
integrator.stats.nsolve += 1
u=+(uprev,$(biki...))
end)
adaptiveexpr=[]
if typeof(tabmask)<:RosenbrockAdaptiveTableau
btildeiki=[:($(Symbol(:btilde,i))*$(Symbol(:k,i))) for i in findall(!iszero,tabmask.btilde)]
push!(adaptiveexpr,quote
if integrator.opts.adaptive
utilde = +($(btildeiki...))
atmp = calculate_residuals(utilde, uprev, u, integrator.opts.abstol,
integrator.opts.reltol,integrator.opts.internalnorm,t)
integrator.EEst = integrator.opts.internalnorm(atmp,t)
end
end)
end
quote
@muladd function perform_step!(integrator, cache::$cachename, repeat_step=false)
@unpack t,dt,uprev,u,f,p = integrator
@unpack tf,uf = cache
$unpacktabexpr
$(dtCijexprs...)
$(dtdiexprs...)
dtgamma = dt*gamma
mass_matrix = integrator.f.mass_matrix
# Time derivative
tf.u = uprev
dT = calc_tderivative(integrator, cache)
W = calc_W(integrator, cache, dtgamma, repeat_step)
linsolve_tmp = integrator.fsalfirst + dtd1*dT #calc_rosenbrock_differentiation!
$(iterexprs...)
integrator.fsallast = f(u, p, t + dt)
OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
integrator.k[1] = integrator.fsalfirst
integrator.k[2] = integrator.fsallast
integrator.u = u
$(adaptiveexpr...)
end
end
end
"""
gen_perform_step(tabmask::RosenbrockTableau{Bool,Bool},cachename::Symbol,n_normalstep::Int,specialstepexpr=:nothing)
Generate inplace version of `perform_step!` expression emulating those in `perform_step/rosenbrock_perform_step.jl`.
The inplace `perform_step!` produces the same result as the non-inplace version except that it treats the mass_matrix appropriately.
"""
function gen_perform_step(tabmask::RosenbrockTableau{Bool,Bool},cachename::Symbol,n_normalstep::Int,specialstepexpr=:nothing)
unpacktabexpr=:(@unpack ()=cache.tab)
unpacktabexpr.args[3].args[1].args=_nonzero_vals(tabmask)
dtCij=[:($(Symbol(:dtC,"$(Cind[1])$(Cind[2])"))=$(Symbol(:C,"$(Cind[1])$(Cind[2])"))/dt) for Cind in findall(!iszero,tabmask.C)]
dtdi=[:($(Symbol(:dtd,dind[1]))=dt*$(Symbol(:d,dind[1]))) for dind in eachindex(tabmask.d)]
iterexprs=[]
for i in 1:n_normalstep
ki=Symbol(:k,i)
dtdj=Symbol(:dtd,i+1)
aijkj=[:($(Symbol(:a,i+1,j))*$(Symbol(:k,j))) for j in findall(!iszero,tabmask.a[i+1,:])]
dtCijkj=[:($(Symbol(:dtC,i+1,j))*$(Symbol(:k,j))) for j in findall(!iszero,tabmask.C[i+1,:])]
repeatstepexpr=[]
if i==1
repeatstepexpr=[:(!repeat_step)]
end
push!(iterexprs,quote
if $(i==1)
# Must be a part of the first linsolve for preconditioner step
linres = dolinsolve(integrator, linsolve; A = !repeat_step ? W : nothing, b = _vec(linsolve_tmp))
else
linres = dolinsolve(integrator, linsolve; b = _vec(linsolve_tmp))
end
linsolve = linres.cache
vecu = _vec(linres.u)
vecki = _vec($ki)
@.. broadcast=false vecki = -vecu
integrator.stats.nsolve += 1
@.. broadcast=false u = +(uprev,$(aijkj...))
f( du, u, p, t+$(Symbol(:c,i+1))*dt)
OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
if mass_matrix === I
@.. broadcast=false linsolve_tmp = +(du,$dtdj*dT,$(dtCijkj...))
else
@.. broadcast=false du1 = +($(dtCijkj...))
mul!(_vec(du2),mass_matrix,_vec(du1))
@.. broadcast=false linsolve_tmp = du + $dtdj*dT + du2
end
end)
end
push!(iterexprs,specialstepexpr)
n=length(tabmask.b)
ks=[Symbol(:k,i) for i in 1:n]
klast=Symbol(:k,n)
biki=[:($(Symbol(:b,i))*$(Symbol(:k,i))) for i in 1:n]
push!(iterexprs,quote
linres = dolinsolve(integrator, linsolve; b = _vec(linsolve_tmp))
linsolve = linres.cache
vecu = _vec(linres.u)
vecklast = _vec($klast)
@.. broadcast=false vecklast = -vecu
integrator.stats.nsolve += 1
@.. broadcast=false u = +(uprev,$(biki...))
end)
adaptiveexpr=[]
if typeof(tabmask)<:RosenbrockAdaptiveTableau
btildeiki=[:($(Symbol(:btilde,i))*$(Symbol(:k,i))) for i in findall(!iszero,tabmask.btilde)]
push!(adaptiveexpr,quote
utilde=du
if integrator.opts.adaptive
@.. broadcast=false utilde = +($(btildeiki...))
calculate_residuals!(atmp, utilde, uprev, u, integrator.opts.abstol,
integrator.opts.reltol,integrator.opts.internalnorm,t)
integrator.EEst = integrator.opts.internalnorm(atmp,t)
end
end)
end
quote
@muladd function perform_step!(integrator, cache::$cachename, repeat_step=false)
@unpack t,dt,uprev,u,f,p = integrator
@unpack du,du1,du2,fsallast,dT,J,W,uf,tf,$(ks...),linsolve_tmp,jac_config,atmp,weight = cache
$unpacktabexpr
# Assignments
sizeu = size(u)
uidx = eachindex(integrator.uprev)
mass_matrix = integrator.f.mass_matrix
# Precalculations
$(dtCij...)
$(dtdi...)
dtgamma = dt*gamma
calculate_residuals!(weight, fill!(weight, one(eltype(u))), uprev, uprev,
integrator.opts.abstol, integrator.opts.reltol, integrator.opts.internalnorm, t)
calc_rosenbrock_differentiation!(integrator, cache, dtd1, dtgamma, repeat_step)
linsolve = cache.linsolve
$(iterexprs...)
f( fsallast, u, p, t + dt)
OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
$(adaptiveexpr...)
end
end
end
"""
RosenbrockW6S4OSTableau()
Rahunanthan, A., & Stanescu, D. (2010). High-order W-methods.
Journal of computational and applied mathematics, 233(8), 1798-1811.
"""
function RosenbrockW6S4OSTableau()
a=[0 0 0 0 0;
0.5812383407115008 0 0 0 0;
0.9039624413714670 1.8615191555345010 0 0 0;
2.0765797196750000 0.1884255381414796 1.8701589674910320 0 0;
4.4355506384843120 5.4571817986101890 4.6163507880689300 3.1181119524023610 0;
10.791701698483260 -10.05691522584131 14.995644854284190 5.2743399543909430 1.4297308712611900]
C=[0 0 0 0 0;
-2.661294105131369 0 0 0 0;
-3.128450202373838 0.0000000000000000 0 0 0;
-6.920335474535658 -1.202675288266817 -9.733561811413620 0 0;
-28.09530629102695 20.371262954793770 -41.04375275302869 -19.66373175620895 0;
9.7998186780974000 11.935792886603180 3.6738749290132010 14.807828541095500 0.8318583998690680]
b=[6.4562170746532350,-4.853141317768053,9.7653183340692600,2.0810841772787230,0.6603936866352417,0.6000000000000000]
gamma=0.2500000000000000
d=[0.2500000000000000,0.0836691184292894,0.0544718623516351,-0.3402289722355864,0.0337651588339529,-0.0903074267618540]
c=[0 ,0.1453095851778752,0.3817422770256738,0.6367813704374599,0.7560744496323561,0.9271047239875670]
RosenbrockFixedTableau(a,C,b,gamma,d,c)
end
"""
@RosenbrockW6S4OS(part)
Generate code for the RosenbrockW6S4OS method.
`part` should be one of `:tableau`, `:cache`, `:init`, `:performstep`.
`@RosenbrockW6S4OS(:tableau)` should be placed in `tableaus/rosenbrock_tableaus.jl`.
`@RosenbrockW6S4OS(:cache)` should be placed in `caches/rosenbrock_caches.jl`.
`@RosenbrockW6S4OS(:init)` and `@RosenbrockW6S4OS(:performstep)` should be
placed in `perform_step/rosenbrock_perform_step.jl`.
"""
macro RosenbrockW6S4OS(part)
tab=RosenbrockW6S4OSTableau()
tabmask=_masktab(tab)
algname=:RosenbrockW6S4OS
tabname=:RosenbrockW6S4OSTableau
tabstructname=:RosenbrockW6STableau
cachename=:RosenbrockW6SCache
constcachename=:RosenbrockW6SConstantCache
n_normalstep=length(tab.b)-1
if part.value==:tableau
#println("Generating const cache")
tabstructexpr=gen_tableau_struct(tabmask,tabstructname)
tabexpr=gen_tableau(tab,tabstructexpr,tabname)
return esc(quote $([tabstructexpr,tabexpr]...) end)
elseif part.value==:cache
#println("Generating cache")
constcacheexpr,cacheexpr=gen_cache_struct(tabmask,cachename,constcachename)
algcacheexpr=gen_algcache(cacheexpr,constcachename,algname,tabname)
return esc(quote $([constcacheexpr,cacheexpr,algcacheexpr]...) end)
elseif part.value==:init
#println("Generating initialize")
return esc(gen_initialize(cachename,constcachename))
elseif part.value==:performstep
#println("Generating perform_step")
constperformstepexpr=gen_constant_perform_step(tabmask,constcachename,n_normalstep)
performstepexpr=gen_perform_step(tabmask,cachename,n_normalstep)
return esc(quote $([constperformstepexpr,performstepexpr]...) end)
else
throw(ArgumentError("Unknown parameter!"))
nothing
end
end
"""
Ros4dummyTableau()
Generate a dummy tableau for ROS4 methods. It can be considered as performing elementwise OR to the masks
of those specific tableaus: `Ros4dummyTableau()==_masktab(RosShamp4Tableau()) OR _masktab(Veldd4Tableau()) OR ...`
ROS4 methods have the property of a4j==a3j so a is a 3*3 matrix instead of a 4*4 matrix and c is a 1*3 vector instead of a 1*4 vector.
"""
function Ros4dummyTableau()#create a tabmask for all ROS4 methods where false->0,true->non-0
a=[false false false;
true false false;
true true false]
C=[false false false false;
true false false false;
true true false false;
true true true false]
b=[true,true,true,true]
btilde=[true,true,true,true]
gamma=true
c=[false,true,true]
d=[true,true,true,true]
RosenbrockAdaptiveTableau(a,C,b,btilde,gamma,d,c)
end
"""
RosShamp4Tableau()
L. F. Shampine, Implementation of Rosenbrock Methods,
ACM Transactions on Mathematical Software (TOMS), 8: 2, 93-113.
doi:10.1145/355993.355994
"""
function RosShamp4Tableau()
a=[0 0 0;
2 0 0;
48//25 6//25 0]
C=[ 0 0 0 0;
-8 0 0 0;
372//25 12//5 0 0;
-112//125 -54//125 -2//5 0]
b=[19//9,1//2,25//108,125//108]
btilde=[17//54,7//36,0,125//108]
gamma=1//2
c=[0,1,3//5]
d=[1//2,-3//2,242//100,116//1000]#2.42,0.116
RosenbrockAdaptiveTableau(a,C,b,btilde,gamma,d,c)
end
"""
Veldd4Tableau()
van Veldhuizen, D-stability and Kaps-Rentrop-methods,
M. Computing (1984) 32: 229. doi:10.1007/BF02243574
"""
function Veldd4Tableau()
a=[0 0 0;
2 0 0;
4.812234362695436 4.578146956747842 0]
C=[ 0 0 0 0;
-5.333333333333331 0 0 0;
6.100529678848254 1.804736797378427 0 0;
-2.540515456634749 -9.443746328915205 -1.988471753215993 0]
b=[4.289339254654537,5.036098482851414,0.6085736420673917,1.355958941201148]
btilde=[2.175672787531755,2.950911222575741,-0.7859744544887430,-1.355958941201148]
gamma=0.2257081148225682
c=[0,0.4514162296451364,0.8755928946018455]
d=[0.2257081148225682,-0.04599403502680582,0.5177590504944076,-0.03805623938054428]
RosenbrockAdaptiveTableau(a,C,b,btilde,gamma,d,c)
end
"""
Velds4Tableau()
van Veldhuizen, D-stability and Kaps-Rentrop-methods,
M. Computing (1984) 32: 229. doi:10.1007/BF02243574
"""
function Velds4Tableau()
a=[0 0 0;
2 0 0;
7//4 1//4 0]
C=[ 0 0 0 0;
-8 0 0 0;
-8 -1 0 0;
1//2 -1//2 2 0]
b=[4//3,2//3,-4//3,4//3]
btilde=[-1//3,-1//3,0,-4//3]
gamma=1//2
c=[0,1,1//2]
d=[1//2,-3//2,-3//4,1//4]
RosenbrockAdaptiveTableau(a,C,b,btilde,gamma,d,c)
end
"""
GRK4TTableau()
Kaps, P. & Rentrop, Generalized Runge-Kutta methods of order four
with stepsize control for stiff ordinary differential equations.
P. Numer. Math. (1979) 33: 55. doi:10.1007/BF01396495
"""
function GRK4TTableau()
a=[0 0 0;
2 0 0;
4.524708207373116 4.163528788597648 0]
C=[ 0 0 0 0;
-5.071675338776316 0 0 0;
6.020152728650786 0.1597506846727117 0 0;
-1.856343618686113 -8.505380858179826 -2.084075136023187 0]
b=[3.957503746640777,4.624892388363313,0.6174772638750108,1.282612945269037]
btilde=[2.302155402932996,3.073634485392623,-0.8732808018045032,-1.282612945269037]
gamma=0.231
c=[0,0.462,0.8802083333333334]
d=[0.2310000000000000,-0.03962966775244303,0.5507789395789127,-0.05535098457052764]
RosenbrockAdaptiveTableau(a,C,b,btilde,gamma,d,c)
end
"""
GRK4ATableau()
Kaps, P. & Rentrop, Generalized Runge-Kutta methods of order four
with stepsize control for stiff ordinary differential equations.
P. Numer. Math. (1979) 33: 55. doi:10.1007/BF01396495
"""
function GRK4ATableau()
a=[0 0 0;
1.108860759493671 0 0;
2.377085261983360 0.1850114988899692 0]
C=[ 0 0 0 0;
-4.920188402397641 0 0 0;
1.055588686048583 3.351817267668938 0 0;
3.846869007049313 3.427109241268180 -2.162408848753263 0]
b=[1.845683240405840,0.1369796894360503,0.7129097783291559,0.6329113924050632]
btilde=[0.04831870177201765,-0.6471108651049505,0.2186876660500240,-0.6329113924050632]
gamma=0.395
c=[0,0.438,0.87]
d=[0.395,-0.3726723954840920,0.06629196544571492,0.4340946962568634]
RosenbrockAdaptiveTableau(a,C,b,btilde,gamma,d,c)
end
"""
Ros4LSTableau()
E. Hairer, G. Wanner, Solving ordinary differential equations II,
stiff and differential-algebraic problems. Computational mathematics (2nd revised ed.), Springer (1996)
"""
function Ros4LSTableau()
a=[0 0 0;
2 0 0;
1.867943637803922 0.2344449711399156 0]
C=[ 0 0 0 0;
-7.137615036412310 0 0 0;
2.580708087951457 0.6515950076447975 0 0;
-2.137148994382534 -0.3214669691237626 -0.6949742501781779 0]
b=[2.255570073418735,0.2870493262186792,0.4353179431840180,1.093502252409163]
btilde=[-0.2815431932141155,-0.07276199124938920,-0.1082196201495311,-1.093502252409163]
gamma=0.5728200000000000
c=[0,1.145640000000000,0.6552168638155900]
d=[0.5728200000000000,-1.769193891319233,0.7592633437920482,-0.1049021087100450]
RosenbrockAdaptiveTableau(a,C,b,btilde,gamma,d,c)
end
"""
@Rosenbrock4(part)
Generate code for the Rosenbrock4 methods: RosShamp4, Veldd4, Velds4, GRK4A, GRK4T, Ros4LStab.
`part` should be one of `:tableau`, `:cache`, `:performstep`.
`@Rosenbrock4(:tableau)` should be placed in `tableaus/rosenbrock_tableaus.jl`.
`@Rosenbrock4(:cache)` should be placed in `caches/rosenbrock_caches.jl`.
`@Rosenbrock4(:performstep)` should be placed in `perform_step/rosenbrock_perform_step.jl`.
The `initialize!` function for Rosenbrock4 methods is already included in `rosenbrock_perform_step.jl`.
The special property of ROS4 methods that a4j==a3j requires a special step in `perform_step!` that
calculates `linsolve_tmp` from the previous `du` which reduces a function call.
"""
macro Rosenbrock4(part)
tabmask=Ros4dummyTableau()#_masktab(tab)
cachename=:Rosenbrock4Cache
constcachename=:Rosenbrock4ConstantCache
RosShamp4tabname=:RosShamp4Tableau
Veldd4tabname=:Veldd4Tableau
Velds4tabname=:Velds4Tableau
GRK4Ttabname=:GRK4TTableau
GRK4Atabname=:GRK4ATableau
Ros4LStabname=:Ros4LStabTableau
n_normalstep=2 #for the third step a4j=a3j which reduced one function call
if part.value==:tableau
#println("Generating tableau for Rosenbrock4")
tabstructexpr=gen_tableau_struct(tabmask,:Ros4Tableau)
tabexprs=Array{Expr,1}()
push!(tabexprs,tabstructexpr)
push!(tabexprs,gen_tableau(RosShamp4Tableau(),tabstructexpr,RosShamp4tabname))
push!(tabexprs,gen_tableau(Veldd4Tableau(),tabstructexpr,Veldd4tabname))
push!(tabexprs,gen_tableau(Velds4Tableau(),tabstructexpr,Velds4tabname))
push!(tabexprs,gen_tableau(GRK4TTableau(),tabstructexpr,GRK4Ttabname))
push!(tabexprs,gen_tableau(GRK4ATableau(),tabstructexpr,GRK4Atabname))
push!(tabexprs,gen_tableau(Ros4LSTableau(),tabstructexpr,Ros4LStabname))
return esc(quote $(tabexprs...) end)
elseif part.value==:cache
#println("Generating cache for Rosenbrock4")
constcacheexpr,cacheexpr=gen_cache_struct(tabmask,cachename,constcachename)
cacheexprs=Array{Expr,1}([constcacheexpr,cacheexpr])
push!(cacheexprs,gen_algcache(cacheexpr,constcachename,:RosShamp4,RosShamp4tabname))
push!(cacheexprs,gen_algcache(cacheexpr,constcachename,:Veldd4,Veldd4tabname))
push!(cacheexprs,gen_algcache(cacheexpr,constcachename,:Velds4,Velds4tabname))
push!(cacheexprs,gen_algcache(cacheexpr,constcachename,:GRK4T,GRK4Ttabname))
push!(cacheexprs,gen_algcache(cacheexpr,constcachename,:GRK4A,GRK4Atabname))
push!(cacheexprs,gen_algcache(cacheexpr,constcachename,:Ros4LStab,Ros4LStabname))
return esc(quote $(cacheexprs...) end)
elseif part.value==:performstep
#println("Generating perform_step for Rosenbrock4")
specialstepconst=quote
k3 = _reshape(W\-_vec(linsolve_tmp), axes(uprev))
integrator.stats.nsolve += 1
#u = uprev + a31*k1 + a32*k2 #a4j=a3j
#du = f(u, p, t+c3*dt) #reduced function call
if mass_matrix === I
linsolve_tmp = du + dtd4*dT + dtC41*k1 + dtC42*k2 + dtC43*k3
else
linsolve_tmp = du + dtd4*dT + mass_matrix * (dtC41*k1 + dtC42*k2 + dtC43*k3)
end
end
specialstep=quote
linres = dolinsolve(integrator, linsolve; b = _vec(linsolve_tmp))
linsolve = linres.cache
cache.linsolve = linsolve
vecu = _vec(linres.u)
veck3 = _vec(k3)
@.. broadcast=false veck3 = -vecu
integrator.stats.nsolve += 1
#@.. broadcast=false u = uprev + a31*k1 + a32*k2 #a4j=a3j
#f( du, u, p, t+c3*dt) #reduced function call
if mass_matrix === I
@.. broadcast=false linsolve_tmp = du + dtd4*dT + dtC41*k1 + dtC42*k2 + dtC43*k3
else
@.. broadcast=false du1 = dtC41*k1 + dtC42*k2 + dtC43*k3
mul!(du2,mass_matrix,du1)
@.. broadcast=false linsolve_tmp = du + dtd4*dT + du2
end
end
constperformstepexpr=gen_constant_perform_step(tabmask,constcachename,n_normalstep,specialstepconst)
performstepexpr=gen_perform_step(tabmask,cachename,n_normalstep,specialstep)
return esc(quote $([constperformstepexpr,performstepexpr]...) end)
else
throw(ArgumentError("Unknown parameter!"))
nothing
end
end
#ROS2, ROS23 and ROS34PW methods (Rang and Angermann, 2005)
"""
Ros34dummyTableau()
Generate a dummy tableau for ROS34W methods proposed by Rang and Angermann. This type of methods has 4 steps.
"""
function Ros34dummyTableau()
a=[false false false false;
true false false false;
true true false false;
true true true false]
C=[false false false false;
true false false false;
true true false false;
true true true false]
b=[true,true,true,true]
btilde=[true,true,true,true]
gamma=true
c=[false,true,true,true]
d=[true,true,true,true]
RosenbrockAdaptiveTableau(a,C,b,btilde,gamma,d,c)
end
"""
Ros23dummyTableau()
Generate a dummy tableau for ROS23 methods proposed by Rang. This type of methods has 3 steps.
"""
function Ros23dummyTableau()
a=[false false false;
true false false;
true true false]
C=[false false false;
true false false;
true true false;]
b=[true,true,true]
btilde=[true,true,true]
gamma=true
c=[false,true,true]
d=[true,true,true]
RosenbrockAdaptiveTableau(a,C,b,btilde,gamma,d,c)
end
"""
Ros2dummyTableau()
Generate a dummy tableau for ROS2 methods. This type of methods has 2 steps.
"""
function Ros2dummyTableau()
a=[false false;
true false]
C=[false false;
true false]
b=[true,true]
btilde=[true,true]
gamma=true
c=[false,true]
d=[true,true]
RosenbrockAdaptiveTableau(a,C,b,btilde,gamma,d,c)
end
"""
_transformtab(Alpha,Gamma,B,Bhat)
Transform the tableau from values in the paper into values used in OrdinaryDiffEq according to p112 in Hairer and Wanner.
E. Hairer, G. Wanner, Solving ordinary differential equations II, stiff and
differential-algebraic problems. Computational mathematics (2nd revised ed.), Springer (1996)
"""
function _transformtab(Alpha,Gamma,B,Bhat)
invGamma=inv(Gamma)
a=Alpha*invGamma
C=diagm(0=>diag(invGamma))-invGamma
b=[(transpose(B)*invGamma)...]# [2Darray...]=>1Darray
btilde=[(transpose(B-Bhat)*invGamma)...]
gamma=Gamma[1,1]#Gamma11==Gamma22==...==Gammass
d=[sum(Gamma,dims=2)...]#di=sum_j Gamma_ij
c=[sum(Alpha,dims=2)...]#ci=sum_j Alpha_ij
(a,C,b,btilde,d,c)
end
# 2 step ROS Methods
"""
ROS2Tableau()
2nd order stiffly accurate Rosenbrock method with 2 internal stages with (Rinf=0).
The embedded method is taken from Kinetic PreProcessor (KPP).
J. G. Verwer et al. (1999): A second-order Rosenbrock method applied to photochemical dispersion problems
https://doi.org/10.1137/S1064827597326651
"""
function ROS2Tableau() # 2nd order
gamma=1.7071067811865475 # 1+1/sqrt(2)
Alpha=[0 0;
1. 0]
Gamma=[gamma 0;
-3.414213562373095 gamma]
B=[0.5, 0.5]
Bhat=[1, 0]
a,C,b,btilde,d,c=_transformtab(Alpha,Gamma,B,Bhat)
RosenbrockAdaptiveTableau(a,C,b,btilde,gamma,d,c)
end
@doc rosenbrock_wanner_docstring(
"""
An Order 2/3 L-Stable Rosenbrock-W method which is good for very stiff equations with oscillations at low tolerances. 2nd order stiff-aware interpolation.
""",
"Rosenbrock23",
references = """
- Shampine L.F. and Reichelt M., (1997) The MATLAB ODE Suite, SIAM Journal of
Scientific Computing, 18 (1), pp. 1-22.
""",
with_step_limiter = true) Rosenbrock23
@doc rosenbrock_wanner_docstring(
"""
An Order 3/2 A-Stable Rosenbrock-W method which is good for mildly stiff equations without oscillations at low tolerances. Note that this method is prone to instability in the presence of oscillations, so use with caution. 2nd order stiff-aware interpolation.
""",
"Rosenbrock32",
references = """
- Shampine L.F. and Reichelt M., (1997) The MATLAB ODE Suite, SIAM Journal of
Scientific Computing, 18 (1), pp. 1-22.
""",
with_step_limiter = true) Rosenbrock32
@doc rosenbrock_docstring(
"""
3rd order A-stable and stiffly stable Rosenbrock method. Keeps high accuracy on discretizations of nonlinear parabolic PDEs.
""",
"ROS3P",
references = """
- Lang, J. & Verwer, ROS3P—An Accurate Third-Order Rosenbrock Solver Designed for
Parabolic Problems J. BIT Numerical Mathematics (2001) 41: 731. doi:10.1023/A:1021900219772
""",
with_step_limiter = true) ROS3P
@doc rosenbrock_wanner_docstring(
"""
An Order 2/3 L-Stable Rosenbrock-W method for stiff ODEs and DAEs in mass matrix form. 2nd order stiff-aware interpolation and additional error test for interpolation.
""",
"Rodas23W",
references = """
- Steinebach G., Rodas23W / Rodas32P - a Rosenbrock-type method for DAEs with additional error estimate for dense output and Julia implementation,
In progress.
""",
with_step_limiter = true) Rodas23W
@doc rosenbrock_wanner_docstring(
"""
A 4th order L-stable Rosenbrock-W method.
""",
"ROS34PW1a",
references = """
@article{rang2005new,
title={New Rosenbrock W-methods of order 3 for partial differential algebraic equations of index 1},
author={Rang, Joachim and Angermann, L},
journal={BIT Numerical Mathematics},
volume={45},
pages={761--787},
year={2005},
publisher={Springer}}
""") ROS34PW1a
@doc rosenbrock_wanner_docstring(
"""
A 4th order L-stable Rosenbrock-W method.
""",
"ROS34PW1b",
references = """
@article{rang2005new,
title={New Rosenbrock W-methods of order 3 for partial differential algebraic equations of index 1},
author={Rang, Joachim and Angermann, L},
journal={BIT Numerical Mathematics},
volume={45},
pages={761--787},
year={2005},
publisher={Springer}}
""") ROS34PW1b
@doc rosenbrock_wanner_docstring(
"""
A 4th order stiffy accurate Rosenbrock-W method for PDAEs.
""",
"ROS34PW2",
references = """
@article{rang2005new,
title={New Rosenbrock W-methods of order 3 for partial differential algebraic equations of index 1},
author={Rang, Joachim and Angermann, L},
journal={BIT Numerical Mathematics},
volume={45},