Merge pull request #11 from thomvet/Support-nested-duals

removed random numbers from tests; updated readme

Author: ChrisRackauckas

README.md

@@ -24,13 +24,17 @@ needed.
 ### Using dualcache
 
 ```julia
-dualcache(u::AbstractArray, N = Val{default_cache_size(length(u))})
+dualcache(u::AbstractArray, N::Int=ForwardDiff.pickchunksize(length(u)); levels::Int = 1)
+dualcache(u::AbstractArray, N::AbstractArray{<:Int})
 ```
 
 The `dualcache` function builds a `DualCache` object that stores both a version
 of the cache for `u` and for the `Dual` version of `u`, allowing use of
-pre-cached vectors with forward-mode automatic differentiation. To access the
-caches, one uses:
+pre-cached vectors with forward-mode automatic differentiation. Note that
+`dualcache`, due to its design, is only compatible with arrays that contain concretely
+typed elements.
+
+To access the caches, one uses:
 
 ```julia
 get_tmp(tmp::DualCache, u)
@@ -41,19 +45,29 @@ version of the cache. Otherwise it returns the standard cache (for use in the
 calls without automatic differentiation).
 
 In order to preallocate to the right size, the `dualcache` needs to be specified
-to have the corrent `N` matching the chunk size of the dual numbers or larger.
+to have the correct `N` matching the chunk size of the dual numbers or larger. 
+If the chunk size `N` specified is too large, `get_tmp` will automatically resize 
+when dispatching; this remains type-stable and non-allocating, but comes at the 
+expense of additional memory.
+
 In a differential equation, optimization, etc., the default chunk size is computed
 from the state vector `u`, and thus if one creates the `dualcache` via
 `dualcache(u)` it will match the default chunking of the solver libraries.
 
+`dualcache` is also compatible with nested automatic differentiation calls through
+the `levels` keyword (`N` for each level computed using based on the size of the 
+state vector) or by specifying `N` as an array of integers of chunk sizes, which
+enables full control of chunk sizes on all differentation levels.
+
 ### dualcache Example 1: Direct Usage
 
 ```julia
-randmat = rand(10, 2)
+using ForwardDiff, PreallocationTools
+randmat = rand(5, 3)
 sto = similar(randmat)
 stod = dualcache(sto)
 
-function claytonsample!(sto, τ; randmat=randmat)
+function claytonsample!(sto, τ, α; randmat=randmat)
     sto = get_tmp(sto, τ)
     sto .= randmat
     τ == 0 && return sto
@@ -62,14 +76,23 @@ function claytonsample!(sto, τ; randmat=randmat)
     for i in 1:n
         v = sto[i, 2]
         u = sto[i, 1]
+        sto[i, 1] = (1 - u^(-τ) + u^(-τ)*v^(-(τ/(1 + τ))))^(-1/τ)*α
         sto[i, 2] = (1 - u^(-τ) + u^(-τ)*v^(-(τ/(1 + τ))))^(-1/τ)
     end
     return sto
 end
 
-ForwardDiff.derivative(τ -> claytonsample!(stod, τ), 0.3)
+ForwardDiff.derivative(τ -> claytonsample!(stod, τ, 0.0), 0.3)
+ForwardDiff.jacobian(x -> claytonsample!(stod, x[1], x[2]), [0.3; 0.0])
 ```
 
+In the above, the chunk size of the dual numbers has been selected based on the size
+of `randmat`, resulting in a chunk size of 8 in this case. However, since the derivative 
+is calculated with respect to τ and the Jacobian is calculated with respect to τ and α, 
+specifying the `dualcache` with `stod = dualcache(sto, 1)` or `stod = dualcache(sto, 2)`, 
+respectively, would have been the most memory efficient way of performing these calculations
+(only really relevant for much larger problems).
+
 ### dualcache Example 2: ODEs
 
 ```julia
@@ -80,7 +103,7 @@ function foo(du, u, (A, tmp), t)
     nothing
 end
 prob = ODEProblem(foo, ones(5, 5), (0., 1.0), (ones(5,5), zeros(5,5)))
-solve(prob, Rosenbrock23())
+solve(prob, TRBDF2())
 ```
 
 fails because `tmp` is only real numbers, but during automatic differentiation
@@ -96,7 +119,7 @@ function foo(du, u, (A, tmp), t)
     nothing
 end
 chunk_size = 5
-prob = ODEProblem(foo, ones(5, 5), (0., 1.0), (ones(5,5), dualcache(zeros(5,5), Val{chunk_size})))
+prob = ODEProblem(foo, ones(5, 5), (0., 1.0), (ones(5,5), dualcache(zeros(5,5), chunk_size)))
 solve(prob, TRBDF2(chunk_size=chunk_size))
 ```
 
@@ -114,6 +137,46 @@ chunk_size = 5
 prob = ODEProblem(foo, ones(5, 5), (0., 1.0), (ones(5,5), dualcache(zeros(5,5))))
 solve(prob, TRBDF2())
 ```
+### dualcache Example 3: Nested AD calls in an optimization problem involving a Hessian matrix
+
+```julia
+using LinearAlgebra, OrdinaryDiffEq, PreallocationTools, Optim, GalacticOptim
+function foo(du, u, p, t)
+    tmp = p[2]
+    A = reshape(p[1], size(tmp.du))
+    tmp = get_tmp(tmp, u)
+    mul!(tmp, A, u)
+    @. du = u + tmp
+    nothing
+end
+
+coeffs = -collect(0.1:0.1:0.4)
+cache = dualcache(zeros(2,2), levels = 3)
+prob = ODEProblem(foo, ones(2, 2), (0., 1.0), (coeffs, cache))
+realsol = solve(prob, TRBDF2(), saveat = 0.0:0.1:10.0, reltol = 1e-8)
+
+function objfun(x, prob, realsol, cache)
+    prob = remake(prob, u0 = eltype(x).(prob.u0), p = (x, cache))
+    sol = solve(prob, TRBDF2(), saveat = 0.0:0.1:10.0, reltol = 1e-8)
+
+    ofv = 0.0
+    if any((s.retcode != :Success for s in sol))
+        ofv = 1e12
+    else
+        ofv = sum((sol.-realsol).^2)
+    end    
+    return ofv
+end
+fn(x,p) = objfun(x, p[1], p[2], p[3])
+optfun = OptimizationFunction(fn, GalacticOptim.AutoForwardDiff())
+optprob = OptimizationProblem(optfun, zeros(length(coeffs)), (prob, realsol, cache))
+solve(optprob, Newton())
+```
+Solves an optimization problem for the coefficients, `coeffs`, appearing in a differential equation.
+The optimization is done with [Optim.jl](https://github.com/JuliaNLSolvers/Optim.jl)'s `Newton()` 
+algorithm. Since this involves automatic differentiation in the ODE solver and the calculation 
+of Hessians, three automatic differentiations are nested within each other. Therefore, the `dualcache` 
+is specified with `levels = 3`. 
 
 ## LazyBufferCache
 

test/core_nesteddual.jl

@@ -52,34 +52,34 @@ function foo(du, u, p, t)
     nothing
 end
 
-ps = 3 #use to specify problem size; don't go crazy on this, because of the compilation time...
-coeffs = -rand(ps,ps)
+ps = 2 #use to specify problem size; don't go crazy on this, because of the compilation time...
+coeffs = -collect(0.1:0.1:(ps^2/10))
 cache = dualcache(zeros(ps,ps), levels = 3)
 prob = ODEProblem(foo, ones(ps, ps), (0., 1.0), (coeffs, cache))
 realsol = solve(prob, TRBDF2(), saveat = 0.0:0.1:10.0, reltol = 1e-8)
-u0 = rand(length(coeffs))
 
 function objfun(x, prob, realsol, cache)
     prob = remake(prob, u0 = eltype(x).(prob.u0), p = (x, cache))
     sol = solve(prob, TRBDF2(), saveat = 0.0:0.1:10.0, reltol = 1e-8)
 
     ofv = 0.0
     if any((s.retcode != :Success for s in sol))
-      ofv = 1e12
+        ofv = 1e12
     else
-      ofv = sum((sol.-realsol).^2)
+        ofv = sum((sol.-realsol).^2)
     end    
     return ofv
 end
 fn(x,p) = objfun(x, p[1], p[2], p[3])
 optfun = OptimizationFunction(fn, GalacticOptim.AutoForwardDiff())
-optprob = OptimizationProblem(optfun, -rand(length(coeffs)), (prob, realsol, cache), chunk_size = 2)
+optprob = OptimizationProblem(optfun, zeros(length(coeffs)), (prob, realsol, cache))
 newtonsol = solve(optprob, Newton())
 
-@test all(abs.(coeffs[:] .- newtonsol.u) .< 1e-2)
+@test all(abs.(coeffs .- newtonsol.u) .< 1e-3)
 
 #an example where chunk_sizes are not the same on all differentiation levels:
-cache = dualcache(zeros(ps,ps), [9, 9, 2])
+cache = dualcache(zeros(ps,ps), [4, 4, 2])
+prob = ODEProblem(foo, ones(ps, ps), (0., 1.0), (coeffs, cache))
 realsol = solve(prob, TRBDF2(chunk_size = 2), saveat = 0.0:0.1:10.0, reltol = 1e-8)
 
 function objfun(x, prob, realsol, cache)
@@ -88,17 +88,17 @@ function objfun(x, prob, realsol, cache)
 
     ofv = 0.0
     if any((s.retcode != :Success for s in sol))
-      ofv = 1e12
+        ofv = 1e12
     else
-      ofv = sum((sol.-realsol).^2)
+        ofv = sum((sol.-realsol).^2)
     end    
     return ofv
 end
 
 fn(x,p) = objfun(x, p[1], p[2], p[3])
 
 optfun = OptimizationFunction(fn, GalacticOptim.AutoForwardDiff())
-optprob = OptimizationProblem(optfun, -rand(length(coeffs)), (prob, realsol, cache))
+optprob = OptimizationProblem(optfun, zeros(length(coeffs)), (prob, realsol, cache))
 newtonsol2 = solve(optprob, Newton())
 
-@test all(abs.(coeffs[:] .- newtonsol2.u) .< 1e-2)
+@test all(abs.(coeffs .- newtonsol2.u) .< 1e-3)