changed dc_dual to be a vector of the right size;

re-sizing implemented differently;
removing dependency ArrayInterface

Author: thomvet

Project.toml

@@ -4,12 +4,10 @@ authors = ["Chris Rackauckas <accounts@chrisrackauckas.com>"]
 version = "0.1.1"
 
 [deps]
-ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
-LabelledArrays = "2ee39098-c373-598a-b85f-a56591580800"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+LabelledArrays = "2ee39098-c373-598a-b85f-a56591580800"
 
 [compat]
-ArrayInterface = "2.6, 3.0"
 ForwardDiff = "0.10.3"
 LabelledArrays = "1"
 julia = "1.6"

src/PreallocationTools.jl

@@ -2,55 +2,49 @@ module PreallocationTools
 
 using ForwardDiff, ArrayInterface, LabelledArrays
 
-struct DiffCache{T<:AbstractArray, S<:AbstractArray}
+struct DiffCache{T<:AbstractArray,S<:AbstractArray}
     du::T
     dual_du::S
 end
 
+#= removed dependency on ArrayInterface, because it seemed not necessary anymore; 
+not sure whether it breaks things that are not in the testset; needs checking. =#
 function DiffCache(u::AbstractArray{T}, siz, ::Type{Val{chunk_size}}) where {T, chunk_size}
-    x = ArrayInterface.restructure(u,zeros(ForwardDiff.Dual{nothing,T,chunk_size}, siz...))
+    x = zeros(T,(chunk_size+1)*prod(siz)) 
     DiffCache(u, x)
 end
 
 """
 
 `dualcache(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
+Builds a `DualCache` object that stores versions of the cache for `u` and for the `Dual` version of `u` allowing use of pre-cached arrays with
 forward-mode automatic differentiation.
 
 """
 dualcache(u::AbstractArray, N=Val{ForwardDiff.pickchunksize(length(u))}) = DiffCache(u, size(u), N)
 
-chunksize(::Type{ForwardDiff.Dual{T,V,N}}) where {T,V,N} = N
+"""
+
+`get_tmp(dc::DiffCache, u)`
+
+Returns the `Dual` or normal cache array stored in `dc` based on the type of `u`. 
 
+"""
 function get_tmp(dc::DiffCache, 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
+  nelem = div(sizeof(T), sizeof(eltype(dc.dual_du)))*prod(size(dc.du))
+  reshape(reinterpret(T, view(dc.dual_du, 1:nelem)), size(dc.du))
 end
 
 function get_tmp(dc::DiffCache, 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
+    nelem = div(sizeof(T), sizeof(eltype(dc.dual_du)))*prod(size(dc.du))
+    reshape(reinterpret(T, view(dc.dual_du, 1:nelem)), size(dc.du))
 end
 
 function get_tmp(dc::DiffCache, u::LabelledArrays.LArray{T,N,D,Syms}) where {T,N,D,Syms}
-  x = reinterpret(T, dc.dual_du.__x)
-  _x = if chunksize(T) === chunksize(eltype(dc.dual_du))
-      x
-  else
-      @view x[axes(dc.du)...]
-  end
-  LabelledArrays.LArray{T,N,D,Syms}(_x)
+    nelem = div(sizeof(T), sizeof(eltype(dc.dual_du)))*prod(size(dc.du))
+    _x = reshape(reinterpret(T, view(dc.dual_du, 1:nelem)), size(dc.du))
+    LabelledArrays.LArray{T,N,D,Syms}(_x)
 end
 
 get_tmp(dc::DiffCache, u::Number) = dc.du
@@ -62,6 +56,7 @@ get_tmp(dc::DiffCache, u::AbstractArray) = dc.du
 A lazily allocated buffer object.  Given a vector `u`, `b[u]` returns a `Vector` of the
 same element type and length `f(length(u))` (defaulting to the same length), which is
 allocated as needed and then cached within `b` for subsequent usage.
+
 """
 struct LazyBufferCache{F<:Function}
     bufs::Dict # a dictionary mapping types to buffers

test/runtests.jl

@@ -1,5 +1,6 @@
-using LinearAlgebra, OrdinaryDiffEq, Test, PreallocationTools, ForwardDiff
+using LinearAlgebra, OrdinaryDiffEq, Test, PreallocationTools, ForwardDiff, LabelledArrays
 
+## Check ODE problem with specified chunk_size
 function foo(du, u, (A, tmp), t)
     tmp = get_tmp(tmp, u)
     mul!(tmp, A, u)
@@ -10,16 +11,17 @@ chunk_size = 5
 prob = ODEProblem(foo, ones(5, 5), (0., 1.0), (ones(5,5), dualcache(zeros(5,5), Val{chunk_size})))
 solve(prob, TRBDF2(chunk_size=chunk_size))
 
+## Check ODE problem with auto-detected chunk_size
 function foo(du, u, (A, tmp), t)
     tmp = get_tmp(tmp, u)
     mul!(tmp, A, u)
     @. du = u + tmp
     nothing
 end
-chunk_size = 5
 prob = ODEProblem(foo, ones(5, 5), (0., 1.0), (ones(5,5), dualcache(zeros(5,5))))
 solve(prob, TRBDF2())
 
+## Check ODE problem with a lazy buffer cache
 function foo(du, u, (A, lbc), t)
     tmp = lbc[u]
     mul!(tmp, A, u)
@@ -29,24 +31,43 @@ end
 prob = ODEProblem(foo, ones(5, 5), (0., 1.0), (ones(5,5), LazyBufferCache()))
 solve(prob, TRBDF2())
 
-## Check resizing
+## Check ODE problem with auto-detected chunk_size and LArray 
+A = LArray((2,2); a=1.0, b=1.0, c=1.0, d=1.0)
+u0 = LArray((2,2); a=1.0, b=1.0, c=1.0, d=1.0)
+function foo(du, u, (A, tmp), t)
+    tmp = get_tmp(tmp, u)
+    mul!(tmp, A, u)
+    @. du = u + tmp
+    nothing
+end
+prob = ODEProblem(foo, u0, (0., 1.0), (A, dualcache(A)))
+solve(prob, TRBDF2())
 
-randmat = rand(10, 2)
+## Check resizing
+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
 
     n = size(sto, 1)
-    for i in 1:n
-        v = sto[i, 2]
-        u = sto[i, 1]
-        sto[i, 2] = (1 - u^(-τ) + u^(-τ)*v^(-(τ/(1 + τ))))^(-1/τ)
-    end
+    println(size(sto))
+    v = @view sto[:, 2]
+    u = @view sto[:, 1]
+    sto[i, 1] = (1 - u^(-τ) + u^(-τ)*v^(-(τ/(1 + τ))))^(-1/τ)*α
+    sto[i, 2] = (1 - u^(-τ) + u^(-τ)*v^(-(τ/(1 + τ))))^(-1/τ)
+
     return sto
 end
 
-ForwardDiff.derivative(τ -> claytonsample!(stod, τ), 0.3)
+#taking the derivative of claytonsample! with respect to τ only
+df1 = ForwardDiff.derivative(τ -> claytonsample!(stod, τ, 0.0), 0.3)
+
+#calculating the jacobian of claytonsample! with respect to τ and α
+df2 = ForwardDiff.jacobian(x -> claytonsample!(stod, x[1], x[2]), [0.3; 0.0]) #should give a 15x2 array,
+#because ForwardDiff flattens the output of jacobian, see: https://juliadiff.org/ForwardDiff.jl/stable/user/api/#ForwardDiff.jacobian
+
+@test all(df1[1:5,2] ≈ df2[6:10,1])