1+ using LinearAlgebra, OrdinaryDiffEq, Test, PreallocationTools, ForwardDiff, GalacticOptim, Optim
2+
3+ randmat = rand (5 , 3 )
4+ sto = similar (randmat)
5+ function claytonsample! (sto, τ, α; randmat= randmat)
6+ sto = get_tmp (sto, τ)
7+ sto .= randmat
8+ τ == 0 && return sto
9+
10+ n = size (sto, 1 )
11+ for i in 1 : n
12+ v = sto[i, 2 ]
13+ u = sto[i, 1 ]
14+ sto[i, 1 ] = (1 - u^ (- τ) + u^ (- τ)* v^ (- (τ/ (1 + τ))))^ (- 1 / τ)* α
15+ sto[i, 2 ] = (1 - u^ (- τ) + u^ (- τ)* v^ (- (τ/ (1 + τ))))^ (- 1 / τ)
16+ end
17+ return sto
18+ end
19+
20+ #= taking the second derivative of claytonsample! with respect to τ with manual chunk_sizes.
21+ In setting up the dualcache, we are setting chunk_size to [1, 1], because we differentiate
22+ only with respect to τ. This initializes the cache with the minimum memory needed. =#
23+ stod = dualcache (sto, [1 , 1 ])
24+ df3 = ForwardDiff. derivative (τ -> ForwardDiff. derivative (ξ -> claytonsample! (stod, ξ, 0.0 ), τ), 0.3 )
25+
26+ #= taking the second derivative of claytonsample! with respect to τ with auto-detected chunk-size.
27+ For the given size of sto, ForwardDiff's heuristic chooses chunk_size = 8. Since this is greater
28+ than what's needed (1+1), the auto-allocated cache is big enough to handle the nested dual numbers, even
29+ if we don't specify the keyword argument levels = 2. This should in general not be relied on to work,
30+ especially if more levels of nesting occur (see optimization example below). =#
31+ stod = dualcache (sto)
32+ df4 = ForwardDiff. derivative (τ -> ForwardDiff. derivative (ξ -> claytonsample! (stod, ξ, 0.0 ), τ), 0.3 )
33+
34+ @test df3 ≈ df4
35+
36+ #= taking the second derivative of claytonsample! with respect to τ with auto-detected chunk-size.
37+ For the given size of sto, ForwardDiff's heuristic chooses chunk_size = 8 and with keyword arg levels = 2,
38+ the created cache size is larger than what's needed (even more so than the last example). =#
39+ stod = dualcache (sto, levels = 2 )
40+ df5 = ForwardDiff. derivative (τ -> ForwardDiff. derivative (ξ -> claytonsample! (stod, ξ, 0.0 ), τ), 0.3 )
41+
42+ @test df3 ≈ df5
43+
44+ #= Checking nested dual numbers using optimization problem involving Optim.jl's Newton() (involving Hessians);
45+ so, we will need one level of AD for the ODE solver (TRBDF2) and two more to calculate the Hessian =#
46+ function foo (du, u, p, t)
47+ tmp = p[2 ]
48+ A = reshape (p[1 ], size (tmp. du))
49+ tmp = get_tmp (tmp, u)
50+ mul! (tmp, A, u)
51+ @. du = u + tmp
52+ nothing
53+ end
54+
55+ ps = 3 # use to specify problem size; don't go crazy on this, because of the compilation time...
56+ coeffs = - rand (ps,ps)
57+ cache = dualcache (zeros (ps,ps), levels = 3 )
58+ prob = ODEProblem (foo, ones (ps, ps), (0. , 1.0 ), (coeffs, cache))
59+ realsol = solve (prob, TRBDF2 (), saveat = 0.0 : 0.1 : 10.0 , reltol = 1e-8 )
60+ u0 = rand (length (coeffs))
61+
62+ function objfun (x, prob, realsol, cache)
63+ prob = remake (prob, u0 = eltype (x).(prob. u0), p = (x, cache))
64+ sol = solve (prob, TRBDF2 (), saveat = 0.0 : 0.1 : 10.0 , reltol = 1e-8 )
65+
66+ ofv = 0.0
67+ if any ((s. retcode != :Success for s in sol))
68+ ofv = 1e12
69+ else
70+ ofv = sum ((sol.- realsol). ^ 2 )
71+ end
72+ return ofv
73+ end
74+ fn (x,p) = objfun (x, p[1 ], p[2 ], p[3 ])
75+ optfun = OptimizationFunction (fn, GalacticOptim. AutoForwardDiff ())
76+ optprob = OptimizationProblem (optfun, - rand (length (coeffs)), (prob, realsol, cache), chunk_size = 2 )
77+ newtonsol = solve (optprob, Newton ())
78+
79+ @test all (abs .(coeffs[:] .- newtonsol. u) .< 1e-2 )
80+
81+ # an example where chunk_sizes are not the same on all differentiation levels:
82+ cache = dualcache (zeros (ps,ps), [9 , 9 , 2 ])
83+ realsol = solve (prob, TRBDF2 (chunk_size = 2 ), saveat = 0.0 : 0.1 : 10.0 , reltol = 1e-8 )
84+
85+ function objfun (x, prob, realsol, cache)
86+ prob = remake (prob, u0 = eltype (x).(prob. u0), p = (x, cache))
87+ sol = solve (prob, TRBDF2 (chunk_size = 2 ), saveat = 0.0 : 0.1 : 10.0 , reltol = 1e-8 )
88+
89+ ofv = 0.0
90+ if any ((s. retcode != :Success for s in sol))
91+ ofv = 1e12
92+ else
93+ ofv = sum ((sol.- realsol). ^ 2 )
94+ end
95+ return ofv
96+ end
97+
98+ fn (x,p) = objfun (x, p[1 ], p[2 ], p[3 ])
99+
100+ optfun = OptimizationFunction (fn, GalacticOptim. AutoForwardDiff ())
101+ optprob = OptimizationProblem (optfun, - rand (length (coeffs)), (prob, realsol, cache))
102+ newtonsol2 = solve (optprob, Newton ())
103+
104+ @test all (abs .(coeffs[:] .- newtonsol2. u) .< 1e-2 )
0 commit comments