Merge pull request #53 from SciML/resizing

Add FixedSizeDiffCache

Author: ChrisRackauckas

src/PreallocationTools.jl

@@ -2,6 +2,73 @@ module PreallocationTools
 
 using ForwardDiff, ArrayInterfaceCore, Adapt
 
+struct FixedSizeDiffCache{T <: AbstractArray, S <: AbstractArray}
+    du::T
+    dual_du::S
+    any_du::Vector{Any}
+end
+
+function FixedSizeDiffCache(u::AbstractArray{T}, siz,
+                            ::Type{Val{chunk_size}}) where {T, chunk_size}
+    x = ArrayInterfaceCore.restructure(u,
+                                       zeros(ForwardDiff.Dual{nothing, T, chunk_size},
+                                             siz...))
+    xany = Any[]
+    FixedSizeDiffCache(deepcopy(u), x, xany)
+end
+
+"""
+
+`FixedSizeDiffCache(u::AbstractArray, N = Val{default_cache_size(length(u))})`
+
+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.
+"""
+function FixedSizeDiffCache(u::AbstractArray,
+                            ::Type{Val{N}} = Val{ForwardDiff.pickchunksize(length(u))}) where {
+                                                                                               N
+                                                                                               }
+    FixedSizeDiffCache(u, size(u), Val{N})
+end
+
+function FixedSizeDiffCache(u::AbstractArray, N::Integer)
+    FixedSizeDiffCache(u, size(u), Val{N})
+end
+
+chunksize(::Type{ForwardDiff.Dual{T, V, N}}) where {T, V, N} = N
+
+function get_tmp(dc::FixedSizeDiffCache, u::T) where {T <: ForwardDiff.Dual}
+    x = reinterpret(T, dc.dual_du)
+    if chunksize(T) === chunksize(eltype(dc.dual_du))
+        x
+    else
+        @view x[axes(dc.du)...]
+    end
+end
+
+function get_tmp(dc::FixedSizeDiffCache, u::AbstractArray{T}) where {T <: ForwardDiff.Dual}
+    x = reinterpret(T, dc.dual_du)
+    if chunksize(T) === chunksize(eltype(dc.dual_du))
+        x
+    else
+        @view x[axes(dc.du)...]
+    end
+end
+
+function get_tmp(dc::FixedSizeDiffCache, u::Union{Number, AbstractArray})
+    if promote_type(eltype(dc.du), eltype(u)) <: eltype(dc.du)
+        dc.du
+    else
+        if length(dc.du) > length(dc.any_du)
+            resize!(dc.any_du, length(dc.du))
+        end
+        _restructure(dc.du, dc.any_du)
+    end
+end
+
+# DiffCache
+
 struct DiffCache{T <: AbstractArray, S <: AbstractArray}
     du::T
     dual_du::S
@@ -16,25 +83,27 @@ function DiffCache(u::AbstractArray{T}, siz, chunk_sizes) where {T}
 end
 
 """
+`DiffCache(u::AbstractArray, N::Int = ForwardDiff.pickchunksize(length(u)); levels::Int = 1)`
+`DiffCache(u::AbstractArray; N::AbstractArray{<:Int})`
 
-`dualcache(u::AbstractArray, N::Int = ForwardDiff.pickchunksize(length(u)); levels::Int = 1)`
-`dualcache(u::AbstractArray; N::AbstractArray{<:Int})`
-
-Builds a `DualCache` object that stores both a version of the cache for `u`
+Builds a `DiffCache` 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. Supports nested AD via keyword `levels`
 or specifying an array of chunk_sizes.
 
 """
-function dualcache(u::AbstractArray, N::Int = ForwardDiff.pickchunksize(length(u));
+function DiffCache(u::AbstractArray, N::Int = ForwardDiff.pickchunksize(length(u));
                    levels::Int = 1)
     DiffCache(u, size(u), N * ones(Int, levels))
 end
-dualcache(u::AbstractArray, N::AbstractArray{<:Int}) = DiffCache(u, size(u), N)
-function dualcache(u::AbstractArray, ::Type{Val{N}}; levels::Int = 1) where {N}
-    dualcache(u, N; levels)
+DiffCache(u::AbstractArray, N::AbstractArray{<:Int}) = DiffCache(u, size(u), N)
+function DiffCache(u::AbstractArray, ::Type{Val{N}}; levels::Int = 1) where {N}
+    DiffCache(u, N; levels)
 end
-dualcache(u::AbstractArray, ::Val{N}; levels::Int = 1) where {N} = dualcache(u, N; levels)
+DiffCache(u::AbstractArray, ::Val{N}; levels::Int = 1) where {N} = dualcache(u, N; levels)
+
+# Legacy deprecate later
+const dualcache = DiffCache
 
 """
 
@@ -89,6 +158,8 @@ function enlargedualcache!(dc, nelem) #warning comes only once per dualcache.
     resize!(dc.dual_du, nelem)
 end
 
+# LazyBufferCache
+
 """
     b = LazyBufferCache(f=identity)
 
@@ -112,6 +183,7 @@ function Base.getindex(b::LazyBufferCache, u::T) where {T <: AbstractArray}
     return buf
 end
 
-export dualcache, get_tmp, LazyBufferCache
+export FixedSizeDiffCache, DiffCache, LazyBufferCache
+export get_tmp
 
 end

test/core_dispatch.jl

@@ -2,7 +2,7 @@ using LinearAlgebra, Test, PreallocationTools, ForwardDiff, LabelledArrays,
       RecursiveArrayTools
 
 function test(u0, dual, chunk_size)
-    cache = PreallocationTools.dualcache(u0, chunk_size)
+    cache = PreallocationTools.DiffCache(u0, chunk_size)
     allocs_normal1 = @allocated get_tmp(cache, u0)
     allocs_normal2 = @allocated get_tmp(cache, first(u0))
     allocs_dual1 = @allocated get_tmp(cache, dual)
@@ -41,6 +41,28 @@ results = test(u0, dual, chunk_size)
 @test eltype(results[5]) == eltype(u0)
 @test eltype(results[7]) == eltype(dual)
 
+chunk_size = 5
+u0_B = ones(5, 5)
+dual_B = zeros(ForwardDiff.Dual{ForwardDiff.Tag{typeof(something), Float64}, Float64,
+                                chunk_size}, 2, 2)
+cache_B = FixedSizeDiffCache(u0_B, chunk_size)
+tmp_du_BA = get_tmp(cache_B, u0_B)
+tmp_dual_du_BA = get_tmp(cache_B, dual_B)
+tmp_du_BN = get_tmp(cache_B, u0_B[1])
+tmp_dual_du_BN = get_tmp(cache_B, dual_B[1])
+@test size(tmp_du_BA) == size(u0_B)
+@test typeof(tmp_du_BA) == typeof(u0_B)
+@test eltype(tmp_du_BA) == eltype(u0_B)
+@test size(tmp_dual_du_BA) == size(u0_B)
+@test_broken typeof(tmp_dual_du_BA) == typeof(dual_B)
+@test eltype(tmp_dual_du_BA) == eltype(dual_B)
+@test size(tmp_du_BN) == size(u0_B)
+@test typeof(tmp_du_BN) == typeof(u0_B)
+@test eltype(tmp_du_BN) == eltype(u0_B)
+@test size(tmp_dual_du_BN) == size(u0_B)
+@test_broken typeof(tmp_dual_du_BN) == typeof(dual_B)
+@test eltype(tmp_dual_du_BN) == eltype(dual_B)
+
 #LArray tests
 chunk_size = 4
 u0 = LArray((2, 2); a = 1.0, b = 1.0, c = 1.0, d = 1.0)
@@ -94,3 +116,27 @@ results = test(u0, dual, chunk_size)
 #eltype tests
 @test eltype(results[5]) == eltype(u0)
 @test eltype(results[7]) == eltype(dual)
+
+u0_AP = ArrayPartition(ones(2, 2), ones(3, 3))
+dual_a = zeros(ForwardDiff.Dual{ForwardDiff.Tag{typeof(something), Float64}, Float64,
+                                chunk_size}, 2, 2)
+dual_b = zeros(ForwardDiff.Dual{ForwardDiff.Tag{typeof(something), Float64}, Float64,
+                                chunk_size}, 3, 3)
+dual_AP = ArrayPartition(dual_a, dual_b)
+cache_AP = FixedSizeDiffCache(u0_AP, chunk_size)
+tmp_du_APA = get_tmp(cache_AP, u0_AP)
+tmp_dual_du_APA = get_tmp(cache_AP, dual_AP)
+tmp_du_APN = get_tmp(cache_AP, u0_AP[1])
+tmp_dual_du_APN = get_tmp(cache_AP, dual_AP[1])
+@test size(tmp_du_APA) == size(u0_AP)
+@test typeof(tmp_du_APA) == typeof(u0_AP)
+@test eltype(tmp_du_APA) == eltype(u0_AP)
+@test size(tmp_dual_du_APA) == size(u0_AP)
+@test_broken typeof(tmp_dual_du_APA) == typeof(dual_AP)
+@test eltype(tmp_dual_du_APA) == eltype(dual_AP)
+@test size(tmp_du_APN) == size(u0_AP)
+@test typeof(tmp_du_APN) == typeof(u0_AP)
+@test eltype(tmp_du_APN) == eltype(u0_AP)
+@test size(tmp_dual_du_APN) == size(u0_AP)
+@test_broken typeof(tmp_dual_du_APN) == typeof(dual_AP)
+@test eltype(tmp_dual_du_APN) == eltype(dual_AP)

test/core_nesteddual.jl

@@ -19,27 +19,27 @@ function claytonsample!(sto, τ, α; randmat = randmat)
 end
 
 #= taking the second derivative of claytonsample! with respect to τ with manual chunk_sizes.
-In setting up the dualcache, we are setting chunk_size to [1, 1], because we differentiate 
+In setting up the DiffCache, we are setting chunk_size to [1, 1], because we differentiate
 only with respect to τ. This initializes the cache with the minimum memory needed. =#
-stod = dualcache(sto, [1, 1])
+stod = DiffCache(sto, [1, 1])
 df3 = ForwardDiff.derivative(τ -> ForwardDiff.derivative(ξ -> claytonsample!(stod, ξ, 0.0),
                                                          τ), 0.3)
 
-#= taking the second derivative of claytonsample! with respect to τ with auto-detected chunk-size. 
-For the given size of sto, ForwardDiff's heuristic chooses chunk_size = 8. Since this is greater 
+#= taking the second derivative of claytonsample! with respect to τ with auto-detected chunk-size.
+For the given size of sto, ForwardDiff's heuristic 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 dual numbers, even
-if we don't specify the keyword argument levels = 2. This should in general not be relied on to work, 
+if we don't specify the keyword argument levels = 2. This should in general not be relied on to work,
 especially if more levels of nesting occur (see optimization example below). =#
-stod = dualcache(sto)
+stod = DiffCache(sto)
 df4 = ForwardDiff.derivative(τ -> ForwardDiff.derivative(ξ -> claytonsample!(stod, ξ, 0.0),
                                                          τ), 0.3)
 
 @test df3 ≈ df4
 
-#= taking the second derivative of claytonsample! with respect to τ with auto-detected chunk-size. 
+#= taking the second derivative of claytonsample! with respect to τ with auto-detected chunk-size.
 For the given size of sto, ForwardDiff's heuristic chooses chunk_size = 8 and with keyword arg levels = 2,
 the created cache size is larger than what's needed (even more so than the last example). =#
-stod = dualcache(sto, levels = 2)
+stod = DiffCache(sto, levels = 2)
 df5 = ForwardDiff.derivative(τ -> ForwardDiff.derivative(ξ -> claytonsample!(stod, ξ, 0.0),
                                                          τ), 0.3)
 
@@ -58,7 +58,7 @@ end
 
 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)
+cache = DiffCache(zeros(ps, ps), levels = 3)
 prob = ODEProblem{true, SciMLBase.FullSpecialize}(foo, ones(ps, ps), (0.0, 1.0),
                                                   (coeffs, cache))
 realsol = solve(prob, TRBDF2(), saveat = 0.0:0.1:10.0, reltol = 1e-8)
@@ -68,7 +68,7 @@ function objfun(x, prob, realsol, 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))
+    if any((s.retcode != ReturnCode.Success for s in sol))
         ofv = 1e12
     else
         ofv = sum((sol .- realsol) .^ 2)
@@ -83,7 +83,7 @@ newtonsol = solve(optprob, Newton())
 @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), [4, 4, 2])
+cache = DiffCache(zeros(ps, ps), [4, 4, 2])
 prob = ODEProblem{true, SciMLBase.FullSpecialize}(foo, ones(ps, ps), (0.0, 1.0),
                                                   (coeffs, cache))
 realsol = solve(prob, TRBDF2(chunk_size = 2), saveat = 0.0:0.1:10.0, reltol = 1e-8)
@@ -93,7 +93,7 @@ function objfun(x, prob, realsol, cache)
     sol = solve(prob, TRBDF2(chunk_size = 2), saveat = 0.0:0.1:10.0, reltol = 1e-8)
 
     ofv = 0.0
-    if any((s.retcode != :Success for s in sol))
+    if any((s.retcode != ReturnCode.Success for s in sol))
         ofv = 1e12
     else
         ofv = sum((sol .- realsol) .^ 2)

test/core_odes.jl

@@ -1,6 +1,9 @@
 using LinearAlgebra, OrdinaryDiffEq, Test, PreallocationTools, LabelledArrays,
       RecursiveArrayTools
 
+# upstream
+OrdinaryDiffEq.DiffEqBase.anyeltypedual(x::FixedSizeDiffCache, counter = 0) = Any
+
 #Base array
 function foo(du, u, (A, tmp), t)
     tmp = get_tmp(tmp, u)
@@ -12,16 +15,26 @@ end
 chunk_size = 5
 u0 = ones(5, 5)
 A = ones(5, 5)
-cache = dualcache(zeros(5, 5), chunk_size)
+cache = DiffCache(zeros(5, 5), chunk_size)
+prob = ODEProblem{true, SciMLBase.FullSpecialize}(foo, u0, (0.0, 1.0), (A, cache))
+sol = solve(prob, TRBDF2(chunk_size = chunk_size))
+@test sol.retcode == ReturnCode.Success
+
+cache = FixedSizeDiffCache(zeros(5, 5), chunk_size)
 prob = ODEProblem{true, SciMLBase.FullSpecialize}(foo, u0, (0.0, 1.0), (A, cache))
 sol = solve(prob, TRBDF2(chunk_size = chunk_size))
-@test sol.retcode == :Success
+@test sol.retcode == ReturnCode.Success
 
 #with auto-detected chunk_size
-cache = dualcache(zeros(5, 5))
+cache = DiffCache(zeros(5, 5))
 prob = ODEProblem{true, SciMLBase.FullSpecialize}(foo, ones(5, 5), (0.0, 1.0), (A, cache))
 sol = solve(prob, TRBDF2())
-@test sol.retcode == :Success
+@test sol.retcode == ReturnCode.Success
+
+prob = ODEProblem(foo, ones(5, 5), (0.0, 1.0),
+                  (ones(5, 5), FixedSizeDiffCache(zeros(5, 5))))
+sol = solve(prob, TRBDF2())
+@test sol.retcode == ReturnCode.Success
 
 #Base array with LBC
 function foo(du, u, (A, lbc), t)
@@ -33,7 +46,7 @@ end
 prob = ODEProblem{true, SciMLBase.FullSpecialize}(foo, ones(5, 5), (0.0, 1.0),
                                                   (ones(5, 5), LazyBufferCache()))
 sol = solve(prob, TRBDF2())
-@test sol.retcode == :Success
+@test sol.retcode == ReturnCode.Success
 
 #LArray
 A = LArray((2, 2); a = 1.0, b = 1.0, c = 1.0, d = 1.0)
@@ -48,10 +61,10 @@ end
 #with specified chunk_size
 chunk_size = 4
 prob = ODEProblem{true, SciMLBase.FullSpecialize}(foo, u0, (0.0, 1.0),
-                                                  (A, dualcache(c, chunk_size)))
+                                                  (A, DiffCache(c, chunk_size)))
 sol = solve(prob, TRBDF2(chunk_size = chunk_size))
-@test sol.retcode == :Success
+@test sol.retcode == ReturnCode.Success
 #with auto-detected chunk_size
-prob = ODEProblem{true, SciMLBase.FullSpecialize}(foo, u0, (0.0, 1.0), (A, dualcache(c)))
+prob = ODEProblem{true, SciMLBase.FullSpecialize}(foo, u0, (0.0, 1.0), (A, DiffCache(c)))
 sol = solve(prob, TRBDF2())
-@test sol.retcode == :Success
+@test sol.retcode == ReturnCode.Success

test/core_resizing.jl

@@ -3,7 +3,7 @@ using Test, PreallocationTools, ForwardDiff, LinearAlgebra
 #test for downsizing cache
 randmat = rand(5, 3)
 sto = similar(randmat)
-stod = dualcache(sto)
+stod = DiffCache(sto)
 
 function claytonsample!(sto, τ, α; randmat = randmat)
     sto = get_tmp(sto, τ)
@@ -48,7 +48,7 @@ u = [3.0, 0.0]
 A = ones(2, 2)
 
 du = similar(u)
-_du = dualcache(du)
+_du = DiffCache(du)
 f = A -> loss(_du, u, A, 0.0)
 analyticalsolution = [3.0 0; 0 0]
 @test ForwardDiff.gradient(f, A) ≈ analyticalsolution