-
-
Notifications
You must be signed in to change notification settings - Fork 202
/
derivative_utils.jl
600 lines (545 loc) · 18.4 KB
/
derivative_utils.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
function calc_tderivative!(integrator, cache, dtd1, repeat_step)
@inbounds begin
@unpack t,dt,uprev,u,f,p = integrator
@unpack du2,fsalfirst,dT,tf,linsolve_tmp = cache
# Time derivative
if !repeat_step # skip calculation if step is repeated
if DiffEqBase.has_tgrad(f)
f.tgrad(dT, uprev, p, t)
else
tf.uprev = uprev
tf.p = p
derivative!(dT, tf, t, du2, integrator, cache.grad_config)
end
end
f(fsalfirst, uprev, p, t)
integrator.destats.nf += 1
@.. linsolve_tmp = fsalfirst + dtd1*dT
end
end
function calc_tderivative(integrator, cache)
@unpack t,dt,uprev,u,f,p = integrator
# Time derivative
if DiffEqBase.has_tgrad(f)
dT = f.tgrad(uprev, p, t)
else
tf = cache.tf
tf.u = uprev
tf.p = p
dT = derivative(tf, t, integrator)
end
dT
end
"""
calc_J(integrator, cache)
Return a new Jacobian object.
If `integrator.f` has a custom Jacobian update function, then it will be called. Otherwise,
either automatic or finite differencing will be used depending on the `uf` object of the
cache.
"""
function calc_J(integrator, cache)
@unpack t,uprev,f,p,alg = integrator
if alg isa DAEAlgorithm
if DiffEqBase.has_jac(f)
J = f.jac(duprev, uprev, p, t)
else
@unpack uf = cache
x = zero(uprev)
J = jacobian(uf, x, integrator)
end
else
if DiffEqBase.has_jac(f)
J = f.jac(uprev, p, t)
else
@unpack uf = cache
uf.f = nlsolve_f(f, alg)
uf.p = p
uf.t = t
J = jacobian(uf, uprev, integrator)
end
integrator.destats.njacs += 1
if alg isa CompositeAlgorithm
integrator.eigen_est = constvalue(opnorm(J, Inf))
end
end
J
end
"""
calc_J!(J, integrator, cache) -> J
Update the Jacobian object `J`.
If `integrator.f` has a custom Jacobian update function, then it will be called. Otherwise,
either automatic or finite differencing will be used depending on the `cache`.
"""
function calc_J!(J, integrator, cache)
@unpack t,uprev,f,p,alg = integrator
if alg isa DAEAlgorithm
if DiffEqBase.has_jac(f)
duprev = integrator.duprev
uf = cache.uf
f.jac(J, duprev, uprev, p, uf.α * uf.invγdt, t)
else
@unpack du1, uf, jac_config = cache
# using `dz` as temporary array
x = cache.dz
fill!(x, zero(eltype(x)))
jacobian!(J, uf, x, du1, integrator, jac_config)
end
else
if DiffEqBase.has_jac(f)
f.jac(J, uprev, p, t)
else
@unpack du1, uf, jac_config = cache
uf.f = nlsolve_f(f, alg)
uf.t = t
uf.p = p
jacobian!(J, uf, uprev, du1, integrator, jac_config)
end
integrator.destats.njacs += 1
if alg isa CompositeAlgorithm
integrator.eigen_est = constvalue(opnorm(J, Inf))
end
end
return nothing
end
"""
WOperator(mass_matrix,gamma,J[;transform=false])
A linear operator that represents the W matrix of an ODEProblem, defined as
```math
W = MM - \\gamma J
```
or, if `transform=true`:
```math
W = \\frac{1}{\\gamma}MM - J
```
where `MM` is the mass matrix (a regular `AbstractMatrix` or a `UniformScaling`),
`γ` is a real number proportional to the time step, and `J` is the Jacobian
operator (must be a `AbstractDiffEqLinearOperator`). A `WOperator` can also be
constructed using a `*DEFunction` directly as
WOperator(f,gamma[;transform=false])
`f` needs to have a jacobian and `jac_prototype`, but the prototype does not need
to be a diffeq operator --- it will automatically be converted to one.
`WOperator` supports lazy `*` and `mul!` operations, the latter utilizing an
internal cache (can be specified in the constructor; default to regular `Vector`).
It supports all of `AbstractDiffEqLinearOperator`'s interface.
"""
mutable struct WOperator{IIP,T,
MType,
GType,
JType <: DiffEqBase.AbstractDiffEqLinearOperator,
F,
C,
} <: DiffEqBase.AbstractDiffEqLinearOperator{T}
mass_matrix::MType
gamma::GType
J::JType
transform::Bool # true => W = mm/gamma - J; false => W = mm - gamma*J
_func_cache::F # cache used in `mul!`
_concrete_form::C # non-lazy form (matrix/number) of the operator
function WOperator{IIP}(mass_matrix, gamma, J, u; transform=false) where IIP
# TODO: there is definitely a missing interface.
# Tentative interface: `has_concrete` and `concertize(A)`
if J isa Union{Number,DiffEqScalar}
if transform
_concrete_form = -mass_matrix / gamma + convert(Number,J)
else
_concrete_form = -mass_matrix + gamma * convert(Number,J)
end
_func_cache = nothing
else
AJ = J isa DiffEqArrayOperator ? convert(AbstractMatrix, J) : J
if AJ isa AbstractMatrix
mm = mass_matrix isa DiffEqArrayOperator ? convert(AbstractMatrix, mass_matrix) : mass_matrix
if transform
_concrete_form = -mm / gamma + AJ
else
_concrete_form = -mm + gamma * AJ
end
else
_concrete_form = nothing
end
_func_cache = zero(u)
end
T = eltype(_concrete_form)
MType = typeof(mass_matrix)
GType = typeof(gamma)
JType = typeof(J)
F = typeof(_func_cache)
C = typeof(_concrete_form)
return new{IIP,T,MType,GType,JType,F,C}(mass_matrix,gamma,J,transform,_func_cache,_concrete_form)
end
end
function WOperator{IIP}(f, u, gamma; transform=false) where IIP
@assert DiffEqBase.has_jac(f) "f needs to have an associated jacobian"
if isa(f, Union{SplitFunction, DynamicalODEFunction})
error("WOperator does not support $(typeof(f)) yet")
end
mass_matrix = f.mass_matrix
# TODO: does this play nicely with time-state dependent mass matrix?
if !isa(mass_matrix, Union{AbstractMatrix,UniformScaling})
mass_matrix = convert(AbstractMatrix, mass_matrix)
end
# Convert jacobian, if needed
J = deepcopy(f.jac_prototype)
if !isa(J, DiffEqBase.AbstractDiffEqLinearOperator)
J = DiffEqArrayOperator(J; update_func=f.jac)
end
return WOperator{IIP}(mass_matrix, gamma, J, u; transform=transform)
end
set_gamma!(W::WOperator, gamma) = (W.gamma = gamma; W)
DiffEqBase.update_coefficients!(W::WOperator,u,p,t) = (update_coefficients!(W.J,u,p,t); update_coefficients!(W.mass_matrix,u,p,t); W)
function Base.convert(::Type{AbstractMatrix}, W::WOperator{IIP}) where IIP
if !IIP
# Allocating
if W.transform
W._concrete_form = -W.mass_matrix / W.gamma + convert(AbstractMatrix,W.J)
else
W._concrete_form = -W.mass_matrix + W.gamma * convert(AbstractMatrix,W.J)
end
else
# Non-allocating
_W = W._concrete_form
J = convert(AbstractMatrix,W.J)
if W.transform
if _W isa Diagonal # axpby doesn't specialize on Diagonal matrix
@inbounds for i in axes(W._concrete_form, 1)
_W[i, i] = J[i, i] - inv(W.gamma) * W.mass_matrix[i, i]
end
else
copyto!(_W, W.mass_matrix)
axpby!(one(W.gamma), J, -inv(W.gamma), _W)
end
else
if _W isa Diagonal # axpby doesn't specialize on Diagonal matrix
@inbounds for i in axes(W._concrete_form, 1)
_W[i, i] = W.gamma*J[i, i] - W.mass_matrix[i, i]
end
else
copyto!(_W, W.mass_matrix)
axpby!(W.gamma, J, -one(W.gamma), W._concrete_form)
end
end
end
return W._concrete_form
end
function Base.convert(::Type{Number}, W::WOperator)
if W.transform
W._concrete_form = -W.mass_matrix / W.gamma + convert(Number,W.J)
else
W._concrete_form = -W.mass_matrix + W.gamma * convert(Number,W.J)
end
return W._concrete_form
end
Base.size(W::WOperator, args...) = size(W.J, args...)
function Base.getindex(W::WOperator, i::Int)
if W.transform
-W.mass_matrix[i] / W.gamma + W.J[i]
else
-W.mass_matrix[i] + W.gamma * W.J[i]
end
end
function Base.getindex(W::WOperator, I::Vararg{Int,N}) where {N}
if W.transform
-W.mass_matrix[I...] / W.gamma + W.J[I...]
else
-W.mass_matrix[I...] + W.gamma * W.J[I...]
end
end
function Base.:*(W::WOperator, x::Union{AbstractVecOrMat,Number})
if W.transform
(W.mass_matrix*x) / -W.gamma + W.J*x
else
-W.mass_matrix*x + W.gamma * (W.J*x)
end
end
function Base.:\(W::WOperator, x::AbstractVecOrMat)
if size(W) == () # scalar operator
convert(Number,W) \ x
else
convert(AbstractMatrix,W) \ x
end
end
function Base.:\(W::WOperator, x::Number)
if size(W) == () # scalar operator
convert(Number,W) \ x
else
convert(AbstractMatrix,W) \ x
end
end
function LinearAlgebra.mul!(Y::AbstractVecOrMat, W::WOperator, B::AbstractVecOrMat)
vec = Base.vec
if W.transform
# Compute mass_matrix * B
if isa(W.mass_matrix, UniformScaling)
a = -W.mass_matrix.λ / W.gamma
@.. Y = a * B
else
mul!(vec(Y), W.mass_matrix, vec(B))
lmul!(-1/W.gamma, Y)
end
# Compute J * B and add
mul!(vec(W._func_cache), W.J, vec(B))
vec(Y) .+= vec(W._func_cache)
else
# Compute mass_matrix * B
if isa(W.mass_matrix, UniformScaling)
vY = vec(Y)
vB = vec(B)
@.. vY = W.mass_matrix.λ * vB
else
mul!(vec(Y), W.mass_matrix, vec(B))
end
# Compute J * B
mul!(vec(W._func_cache), W.J, vec(B))
# Add result
axpby!(W.gamma, vec(W._func_cache), -one(W.gamma), vec(Y))
end
end
"""
islinearfunction(integrator) -> Tuple{Bool,Bool}
return the tuple `(is_linear_wrt_odealg, islinearodefunction)`.
"""
islinearfunction(integrator) = islinearfunction(integrator.f, integrator.alg)
"""
islinearfunction(f, alg) -> Tuple{Bool,Bool}
return the tuple `(is_linear_wrt_odealg, islinearodefunction)`.
"""
function islinearfunction(f, alg)::Tuple{Bool,Bool}
isode = f isa ODEFunction && islinear(f.f)
islin = isode || (alg isa SplitAlgorithms && f isa SplitFunction && islinear(f.f1.f))
return islin, isode
end
function do_newJW(integrator, alg, nlsolver, repeat_step)::NTuple{2,Bool}
integrator.iter <= 1 && return true, true # at least one JW eval at the start
repeat_step && return false, false
islin, _ = islinearfunction(integrator)
islin && return false, false # no further JW eval when it's linear
alg isa DAEAlgorithm && return true, true
isnewton(nlsolver) || return true, true
isfirstcall(nlsolver) && return true, true
isfs = isfirststage(nlsolver)
iszero(nlsolver.fast_convergence_cutoff) && return isfs, isfs
W_iγdt = inv(nlsolver.cache.W_γdt)
iγdt = inv(nlsolver.γ * integrator.dt)
smallstepchange = abs(iγdt/W_iγdt - 1) <= get_new_W_γdt_cutoff(nlsolver)
jbad = nlsolver.status === TryAgain && smallstepchange
errorfail = integrator.EEst > one(integrator.EEst)
return jbad, (jbad || (!smallstepchange) || (isfs && errorfail))
end
@noinline _throwWJerror(W, J) = throw(DimensionMismatch("W: $(axes(W)), J: $(axes(J))"))
@noinline _throwWMerror(W, mass_matrix) = throw(DimensionMismatch("W: $(axes(W)), mass matrix: $(axes(mass_matrix))"))
@noinline _throwJMerror(J, mass_matrix) = throw(DimensionMismatch("J: $(axes(J)), mass matrix: $(axes(mass_matrix))"))
@inline function jacobian2W!(W::AbstractMatrix, mass_matrix::MT, dtgamma::Number, J::AbstractMatrix, W_transform::Bool)::Nothing where MT
# check size and dimension
iijj = axes(W)
@boundscheck (iijj === axes(J) && length(iijj) === 2) || _throwWJerror(W, J)
mass_matrix isa UniformScaling || @boundscheck axes(mass_matrix) === axes(W) || _throwWMerror(W, mass_matrix)
@inbounds if W_transform
invdtgamma = inv(dtgamma)
if MT <: UniformScaling
copyto!(W, J)
idxs = diagind(W)
λ = -mass_matrix.λ
@.. @view(W[idxs]) = muladd(λ, invdtgamma, @view(J[idxs]))
else
@.. W = muladd(-mass_matrix, invdtgamma, J)
end
else
if MT <: UniformScaling
idxs = diagind(W)
@.. W = dtgamma*J
λ = -mass_matrix.λ
@.. @view(W[idxs]) = @view(W[idxs]) + λ
else
@.. W = muladd(dtgamma, J, -mass_matrix)
end
end
return nothing
end
@inline function jacobian2W(mass_matrix::MT, dtgamma::Number, J::AbstractMatrix, W_transform::Bool)::Nothing where MT
# check size and dimension
mass_matrix isa UniformScaling || @boundscheck axes(mass_matrix) === axes(J) || _throwJMerror(J, mass_matrix)
@inbounds if W_transform
invdtgamma = inv(dtgamma)
if MT <: UniformScaling
λ = -mass_matrix.λ
W = J + (λ * invdtgamma)*I
else
W = muladd(-mass_matrix, invdtgamma, J)
end
else
if MT <: UniformScaling
λ = -mass_matrix.λ
W = dtgamma*J + λ*I
else
W = muladd(dtgamma, J, -mass_matrix)
end
end
return W
end
function calc_W!(W, integrator, nlsolver::Union{Nothing,AbstractNLSolver}, cache, dtgamma, repeat_step, W_transform=false)
@unpack t,dt,uprev,u,f,p = integrator
lcache = nlsolver === nothing ? cache : nlsolver.cache
@unpack J = lcache
isdae = integrator.alg isa DAEAlgorithm
alg = unwrap_alg(integrator, true)
if !isdae
mass_matrix = integrator.f.mass_matrix
end
is_compos = integrator.alg isa CompositeAlgorithm
# handle Wfact
if W_transform && DiffEqBase.has_Wfact_t(f)
f.Wfact_t(W, u, p, dtgamma, t)
isnewton(nlsolver) && set_W_γdt!(nlsolver, dtgamma)
is_compos && (integrator.eigen_est = constvalue(opnorm(LowerTriangular(W), Inf)) + inv(dtgamma)) # TODO: better estimate
return nothing
elseif !W_transform && DiffEqBase.has_Wfact(f)
f.Wfact(W, u, p, dtgamma, t)
isnewton(nlsolver) && set_W_γdt!(nlsolver, dtgamma)
if is_compos
opn = opnorm(LowerTriangular(W), Inf)
integrator.eigen_est = (constvalue(opn) + one(opn)) / dtgamma # TODO: better estimate
end
return nothing
end
# check if we need to update J or W
new_jac, new_W = do_newJW(integrator, alg, nlsolver, repeat_step)
if new_jac && isnewton(lcache)
lcache.J_t = t
if isdae
lcache.uf.α = nlsolver.α
lcache.uf.invγdt = inv(dtgamma)
lcache.uf.tmp = nlsolver.tmp
end
end
# calculate W
if W isa WOperator
isnewton(nlsolver) || DiffEqBase.update_coefficients!(W,uprev,p,t) # we will call `update_coefficients!` in NLNewton
W.transform = W_transform; set_gamma!(W, dtgamma)
else # concrete W using jacobian from `calc_J!`
islin, isode = islinearfunction(integrator)
islin ? (J = isode ? f.f : f.f1.f) : ( new_jac && (calc_J!(J, integrator, lcache)) )
!isdae && update_coefficients!(mass_matrix,uprev,p,t)
new_W && !isdae && jacobian2W!(W, mass_matrix, dtgamma, J, W_transform)
end
if isnewton(nlsolver)
set_new_W!(nlsolver, new_W)
if new_jac && isdae
set_W_γdt!(nlsolver, nlsolver.α * inv(dtgamma))
elseif new_W && !isdae
set_W_γdt!(nlsolver, dtgamma)
end
end
new_W && (integrator.destats.nw += 1)
return new_W
end
@noinline function calc_W(integrator, cache, dtgamma, repeat_step, W_transform=false)
@unpack t,uprev,p,f = integrator
isdae = integrator.alg isa DAEAlgorithm
if !isdae
mass_matrix = integrator.f.mass_matrix
end
isarray = uprev isa AbstractArray
# calculate W
is_compos = integrator.alg isa CompositeAlgorithm
islin, isode = islinearfunction(integrator)
!isdae && update_coefficients!(mass_matrix,uprev,p,t)
if islin
J = isode ? f.f : f.f1.f # unwrap the Jacobian accordingly
W = WOperator{false}(mass_matrix, dtgamma, J, uprev; transform=W_transform)
elseif DiffEqBase.has_jac(f)
J = f.jac(uprev, p, t)
if !isa(J, DiffEqBase.AbstractDiffEqLinearOperator)
J = DiffEqArrayOperator(J)
end
W = WOperator{false}(mass_matrix, dtgamma, J, uprev; transform=W_transform)
integrator.destats.nw += 1
else
integrator.destats.nw += 1
J = calc_J(integrator, cache)
if isdae
W = J
else
W_full = W_transform ? J - mass_matrix*inv(dtgamma) :
dtgamma*J - mass_matrix
W = W_full isa Number ? W_full : DiffEqBase.default_factorize(W_full)
end
end
(W isa WOperator && unwrap_alg(integrator, true) isa NewtonAlgorithm) && (W = DiffEqBase.update_coefficients!(W,uprev,p,t)) # we will call `update_coefficients!` in NLNewton
is_compos && (integrator.eigen_est = isarray ? constvalue(opnorm(J, Inf)) : integrator.opts.internalnorm(J, t))
return W
end
function calc_rosenbrock_differentiation!(integrator, cache, dtd1, dtgamma, repeat_step, W_transform)
calc_tderivative!(integrator, cache, dtd1, repeat_step)
nlsolver = nothing
# we need to skip calculating `W` when a step is repeated
new_W = false
if !repeat_step
new_W = calc_W!(cache.W, integrator, nlsolver, cache, dtgamma, repeat_step, W_transform)
end
return new_W
end
# update W matrix (only used in Newton method)
update_W!(integrator, cache, dtgamma, repeat_step) =
update_W!(cache.nlsolver, integrator, cache, dtgamma, repeat_step)
function update_W!(nlsolver::AbstractNLSolver, integrator, cache::OrdinaryDiffEqMutableCache, dtgamma, repeat_step)
if isnewton(nlsolver)
calc_W!(get_W(nlsolver), integrator, nlsolver, cache, dtgamma, repeat_step, true)
end
nothing
end
function update_W!(nlsolver::AbstractNLSolver, integrator, cache, dtgamma, repeat_step)
if isnewton(nlsolver)
isdae = integrator.alg isa DAEAlgorithm
new_jac, new_W = true, true
if isdae && new_jac
lcache = nlsolver.cache
lcache.uf.α = nlsolver.α
lcache.uf.invγdt = inv(dtgamma)
lcache.uf.tmp = @. nlsolver.tmp
lcache.uf.uprev = @. integrator.uprev
end
nlsolver.cache.W = calc_W(integrator, nlsolver.cache, dtgamma, repeat_step, true)
#TODO: jacobian reuse for oop
new_jac && (nlsolver.cache.J_t = integrator.t)
set_new_W!(nlsolver, new_W)
if new_jac && isdae
set_W_γdt!(nlsolver, nlsolver.α * inv(dtgamma))
elseif new_W && !isdae
set_W_γdt!(nlsolver, dtgamma)
end
end
nothing
end
function build_J_W(alg,u,uprev,p,t,dt,f,uEltypeNoUnits,::Val{IIP}) where IIP
islin, isode = islinearfunction(f, alg)
if f.jac_prototype isa DiffEqBase.AbstractDiffEqLinearOperator
W = WOperator{IIP}(f, u, dt)
J = W.J
elseif IIP && f.jac_prototype !== nothing
J = similar(f.jac_prototype)
W = similar(J)
elseif islin || (!IIP && DiffEqBase.has_jac(f))
J = islin ? (isode ? f.f : f.f1.f) : f.jac(uprev, p, t) # unwrap the Jacobian accordingly
if !isa(J, DiffEqBase.AbstractDiffEqLinearOperator)
J = DiffEqArrayOperator(J)
end
W = WOperator{IIP}(f.mass_matrix, dt, J, u)
else
J = if f.jac_prototype === nothing
ArrayInterface.zeromatrix(u)
else
deepcopy(f.jac_prototype)
end
isdae = alg isa DAEAlgorithm
W = if isdae
J
elseif IIP
similar(J)
else
ArrayInterface.lu_instance(J)
end
end
return J, W
end
build_uf(alg,nf,t,p,::Val{true}) = UJacobianWrapper(nf,t,p)
build_uf(alg,nf,t,p,::Val{false}) = UDerivativeWrapper(nf,t,p)