Add extension for SparseConnectivityTracer

Author: ErikQQY

Project.toml

@@ -1,7 +1,7 @@
 name = "PreallocationTools"
 uuid = "d236fae5-4411-538c-8e31-a6e3d9e00b46"
 authors = ["Chris Rackauckas <[email protected]>"]
-version = "0.4.24"
+version = "0.4.25"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -10,9 +10,11 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 
 [weakdeps]
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
+SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
 
 [extensions]
 PreallocationToolsReverseDiffExt = "ReverseDiff"
+PreallocationToolsSparseConnectivityTracerExt = "SparseConnectivityTracer"
 
 [compat]
 Adapt = "3.4, 4"
@@ -30,6 +32,7 @@ RecursiveArrayTools = "3.2"
 ReverseDiff = "1"
 SafeTestsets = "0.1"
 SparseArrays = "1"
+SparseConnectivityTracer = "0.6.12"
 Symbolics = "5.12"
 Test = "1"
 julia = "1.10"
@@ -47,8 +50,9 @@ RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "Random", "LabelledArrays", "LinearAlgebra", "OrdinaryDiffEq", "Test", "RecursiveArrayTools", "Pkg", "SafeTestsets", "Optimization", "OptimizationOptimJL", "SparseArrays", "Symbolics"]
+test = ["Aqua", "Random", "LabelledArrays", "LinearAlgebra", "OrdinaryDiffEq", "Test", "RecursiveArrayTools", "Pkg", "SafeTestsets", "Optimization", "OptimizationOptimJL", "SparseArrays", "Symbolics", "SparseConnectivityTracer"]

ext/PreallocationToolsSparseConnectivityTracerExt.jl

@@ -0,0 +1,46 @@
+module PreallocationToolsSparseConnectivityTracerExt
+
+using PreallocationTools
+isdefined(Base, :get_extension) ? (import SparseConnectivityTracer) :
+(import ..SparseConnectivityTracer)
+
+function PreallocationTools.get_tmp(
+        dc::DiffCache, u::T) where {T <: SparseConnectivityTracer.Dual}
+    if isbitstype(T)
+        nelem = div(sizeof(T), sizeof(eltype(dc.dual_du))) * length(dc.du)
+        if nelem > length(dc.dual_du)
+            PreallocationTools.enlargediffcache!(dc, nelem)
+        end
+        PreallocationTools._restructure(dc.du, reinterpret(T, view(dc.dual_du, 1:nelem)))
+    else
+        PreallocationTools._restructure(dc.du, zeros(T, size(dc.du)))
+    end
+end
+
+function PreallocationTools.get_tmp(
+        dc::DiffCache, ::Type{T}) where {T <: SparseConnectivityTracer.Dual}
+    if isbitstype(T)
+        nelem = div(sizeof(T), sizeof(eltype(dc.dual_du))) * length(dc.du)
+        if nelem > length(dc.dual_du)
+            PreallocationTools.enlargediffcache!(dc, nelem)
+        end
+        PreallocationTools._restructure(dc.du, reinterpret(T, view(dc.dual_du, 1:nelem)))
+    else
+        PreallocationTools._restructure(dc.du, zeros(T, size(dc.du)))
+    end
+end
+
+function PreallocationTools.get_tmp(
+        dc::DiffCache, u::AbstractArray{T}) where {T <: SparseConnectivityTracer.Dual}
+    if isbitstype(T)
+        nelem = div(sizeof(T), sizeof(eltype(dc.dual_du))) * length(dc.du)
+        if nelem > length(dc.dual_du)
+            PreallocationTools.enlargediffcache!(dc, nelem)
+        end
+        PreallocationTools._restructure(dc.du, reinterpret(T, view(dc.dual_du, 1:nelem)))
+    else
+        PreallocationTools._restructure(dc.du, zeros(T, size(dc.du)))
+    end
+end
+
+end