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. In setting up the dualcache,
21+ we are setting chunk_size to [1, 1], because we differentiate only with respect to τ.
22+ 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 τ, auto-detect. For the given size of sto, ForwardDiff's heuristic
27+ chooses chunk_size = 8. Since this is greater than what's needed (1+1), the auto-allocated cache is big enough to handle the nested
28+ dual numbers. This should in general not be relied on to work, especially if more levels of nesting occurs (as below). =#
29+ stod = dualcache (sto)
30+ df4 = ForwardDiff. derivative (τ -> ForwardDiff. derivative (ξ -> claytonsample! (stod, ξ, 0.0 ), τ), 0.3 )
31+
32+ @test df3 ≈ df4
33+
34+ # # Checking nested dual numbers: Checking an optimization problem inspired by the above tests
35+ # # (using Optim.jl's Newton() (involving Hessians) and BFGS() (involving gradients))
36+ function foo (du, u, p, t)
37+ tmp = p[2 ]
38+ A = reshape (p[1 ], size (tmp. du))
39+ tmp = get_tmp (tmp, u)
40+ mul! (tmp, A, u)
41+ @. du = u + tmp
42+ nothing
43+ end
44+
45+ ps = 2 # use to specify problem size (ps ∈ {1,2})
46+ coeffs = rand (ps^ 2 )
47+ cache = dualcache (zeros (ps,ps), [4 , 4 , 4 ])
48+ prob = ODEProblem (foo, ones (ps, ps), (0. , 1.0 ), (coeffs, cache))
49+ realsol = solve (prob, TRBDF2 (), saveat = 0.0 : 0.01 : 1.0 , reltol = 1e-8 )
50+
51+ function objfun (x, prob, realsol, cache)
52+ prob = remake (prob, u0 = eltype (x).(ones (ps, ps)), p = (x, cache))
53+ sol = solve (prob, TRBDF2 (), saveat = 0.0 : 0.01 : 1.0 , reltol = 1e-8 )
54+
55+ ofv = 0.0
56+ if any ((s. retcode != :Success for s in sol))
57+ ofv = 1e12
58+ else
59+ ofv = sum ((sol.- realsol). ^ 2 )
60+ end
61+ return ofv
62+ end
63+
64+ fn (x,p) = objfun (x, p[1 ], p[2 ], p[3 ])
65+
66+ optfun = OptimizationFunction (fn, GalacticOptim. AutoForwardDiff ())
67+ optprob = OptimizationProblem (optfun, rand (size (coeffs)... ), (prob, realsol, cache))
68+ newtonsol = solve (optprob, Newton ())
69+ bfgssol = solve (optprob, BFGS ()) # since only gradients are used here, we could go with a slim dualcache(zeros(ps,ps), [4,4]) as well.
70+
71+ @test all (abs .(coeffs .- newtonsol. u) .< 1e-3 )
72+ @test all (abs .(coeffs .- bfgssol. u) .< 1e-3 )
73+
74+ # an example where chunk_sizes are not the same on all differentiation levels:
75+ function foo (du, u, p, t)
76+ tmp = p[2 ]
77+ A = ones (size (tmp. du)).* p[1 ]
78+ tmp = get_tmp (tmp, u)
79+ mul! (tmp, A, u)
80+ @. du = u + tmp
81+ nothing
82+ end
83+
84+ ps = 2 # use to specify problem size (ps ∈ {1,2})
85+ coeffs = rand (1 )
86+ cache = dualcache (zeros (ps,ps), [1 , 1 , 4 ])
87+ prob = ODEProblem (foo, ones (ps, ps), (0. , 1.0 ), (coeffs, cache))
88+ realsol = solve (prob, TRBDF2 (), saveat = 0.0 : 0.01 : 1.0 , reltol = 1e-8 )
89+
90+ function objfun (x, prob, realsol, cache)
91+ prob = remake (prob, u0 = eltype (x).(ones (ps, ps)), p = (x, cache))
92+ sol = solve (prob, TRBDF2 (), saveat = 0.0 : 0.01 : 1.0 , reltol = 1e-8 )
93+
94+ ofv = 0.0
95+ if any ((s. retcode != :Success for s in sol))
96+ ofv = 1e12
97+ else
98+ ofv = sum ((sol.- realsol). ^ 2 )
99+ end
100+ return ofv
101+ end
102+
103+ fn (x,p) = objfun (x, p[1 ], p[2 ], p[3 ])
104+
105+ optfun = OptimizationFunction (fn, GalacticOptim. AutoForwardDiff ())
106+ optprob = OptimizationProblem (optfun, rand (size (coeffs)... ), (prob, realsol, cache))
107+ newtonsol2 = solve (optprob, Newton ())
108+
109+ @test all (abs .(coeffs .- newtonsol2. u) .< 1e-3 )
0 commit comments