Merge pull request #123 from SciML/preallocationtools_dep

Pick up a PreallocationTools.jl dependency

Author: ChrisRackauckas

Project.toml

@@ -10,13 +10,17 @@ ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+PreallocationTools = "d236fae5-4411-538c-8e31-a6e3d9e00b46"
+RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 
 [compat]
 ArrayInterface = "2.8, 3.0, 4, 5, 6"
 ArrayInterfaceStaticArrays = "0.1"
 ChainRulesCore = "1"
 MacroTools = "0.5"
-StaticArrays = "0.10, 0.11, 0.12, 1.0"
+PreallocationTools = "0.4"
+RecursiveArrayTools = "2"
+StaticArrays = "1.0"
 julia = "1.6"
 
 [extras]

src/LabelledArrays.jl

@@ -1,6 +1,7 @@
 module LabelledArrays
 
 using LinearAlgebra, StaticArrays, ArrayInterface
+import RecursiveArrayTools, PreallocationTools
 
 include("slarray.jl")
 include("larray.jl")
@@ -63,6 +64,16 @@ function ArrayInterface.ismutable(::Type{<:LArray{T, N, Syms}}) where {T, N, Sym
 end
 ArrayInterface.can_setindex(::Type{<:SLArray}) = false
 
+function PreallocationTools.get_tmp(dc::PreallocationTools.DiffCache,
+                                    u::LArray{T, N, D, Syms}) where {T, N, D, Syms}
+    nelem = div(sizeof(T), sizeof(eltype(dc.dual_du))) * length(dc.du)
+    if nelem > length(dc.dual_du)
+        PreallocationTools.enlargedualcache!(dc, nelem)
+    end
+    _x = ArrayInterfaceCore.restructure(dc.du, reinterpret(T, view(dc.dual_du, 1:nelem)))
+    LabelledArrays.LArray{T, N, D, Syms}(_x)
+end
+
 export SLArray, LArray, SLVector, LVector, @SLVector, @LArray, @LVector, @SLArray
 
 export @SLSliced, @LSliced

src/larray.jl

@@ -366,3 +366,12 @@ function Base.vcat(x::LArray, y::LArray)
 end
 
 Base.elsize(::Type{<:LArray{T}}) where {T} = sizeof(T)
+
+function RecursiveArrayTools.recursive_unitless_eltype(a::Type{LArray{T, N, D, Syms}}) where {
+                                                                                              T,
+                                                                                              N,
+                                                                                              D,
+                                                                                              Syms
+                                                                                              }
+    LArray{typeof(one(T)), N, D, Syms}
+end

src/slarray.jl

@@ -39,6 +39,9 @@ function StaticArrays.similar_type(::Type{SLArray{S, T, N, L, Syms}}, T2,
                                    ::Size{S}) where {S, T, N, L, Syms}
     SLArray{S, T2, N, L, Syms}
 end
+function RecursiveArrayTools.recursive_unitless_eltype(a::Type{T}) where {T <: SLArray}
+    StaticArrays.similar_type(a, recursive_unitless_eltype(eltype(a)))
+end
 
 ## Named tuple to SLArray
 #=