Non allocating versions for StaticArrays

avik-pal · avik-pal · commit 110a076018f4 · 2023-11-29T18:07:00.000-05:00
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "SparseDiffTools"
 uuid = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
 authors = ["Pankaj Mishra <pankajmishra1511@gmail.com>", "Chris Rackauckas <contact@chrisrackauckas.com>"]
-version = "2.13.0"
+version = "2.14.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
diff --git a/src/SparseDiffTools.jl b/src/SparseDiffTools.jl
@@ -16,7 +16,7 @@ import ForwardDiff: Dual, jacobian, partials, DEFAULT_CHUNK_THRESHOLD
 using ArrayInterface, SparseArrays
 import ArrayInterface: matrix_colors
 import StaticArrays
-import StaticArrays: StaticArray
+import StaticArrays: StaticArray, SArray, MArray, Size
 # Others
 using SciMLOperators, LinearAlgebra, Random
 import DataStructures: DisjointSets, find_root!, union!
diff --git a/src/highlevel/common.jl b/src/highlevel/common.jl
@@ -173,6 +173,13 @@ A cache for computing the Jacobian of type `AbstractMaybeSparseJacobianCache`.
 """
 function sparse_jacobian_cache end
 
+function sparse_jacobian_static_array(ad, cache, f, x::SArray)
+    # Not the most performant fallback
+    J = init_jacobian(cache)
+    sparse_jacobian!(J, ad, cache, f, MArray(x))
+    return J
+end
+
 """
     sparse_jacobian(ad::AbstractADType, sd::AbstractMaybeSparsityDetection, f, x; fx=nothing)
     sparse_jacobian(ad::AbstractADType, sd::AbstractMaybeSparsityDetection, f!, fx, x)
@@ -189,20 +196,29 @@ function sparse_jacobian(ad::AbstractADType, sd::AbstractMaybeSparsityDetection,
     sparse_jacobian!(J, ad, cache, args...)
     return J
 end
+function sparse_jacobian(ad::AbstractADType, sd::AbstractMaybeSparsityDetection, f,
+        x::SArray; kwargs...)
+    cache = sparse_jacobian_cache(ad, sd, f, x; kwargs...)
+    return sparse_jacobian_static_array(ad, cache, f, x)
+end
 
 """
     sparse_jacobian(ad::AbstractADType, cache::AbstractMaybeSparseJacobianCache, f, x)
     sparse_jacobian(ad::AbstractADType, cache::AbstractMaybeSparseJacobianCache, f!, fx, x)
 
 Use the sparsity detection `cache` for computing the sparse Jacobian. This allocates a new
-Jacobian at every function call
+Jacobian at every function call.
 """
 function sparse_jacobian(ad::AbstractADType, cache::AbstractMaybeSparseJacobianCache,
         args...)
     J = init_jacobian(cache)
     sparse_jacobian!(J, ad, cache, args...)
     return J
 end
+function sparse_jacobian(ad::AbstractADType, cache::AbstractMaybeSparseJacobianCache, f,
+        x::SArray)
+    return sparse_jacobian_static_array(ad, cache, f, x)
+end
 
 """
     sparse_jacobian!(J::AbstractMatrix, ad::AbstractADType, sd::AbstractSparsityDetection,
@@ -247,14 +263,18 @@ function __chunksize(::Union{AutoSparseForwardDiff{C}, AutoForwardDiff{C}}, x) w
     C isa ForwardDiff.Chunk && return C
     return __chunksize(Val(C), x)
 end
-__chunksize(::Val{nothing}, x) = ForwardDiff.Chunk(x)
+__chunksize(::Val{nothing}, x) = __chunksize(x)
 function __chunksize(::Val{C}, x) where {C}
     if C isa Integer && !(C isa Bool)
-        return C ≤ 0 ? ForwardDiff.Chunk(x) : ForwardDiff.Chunk{C}()
+        return C ≤ 0 ? __chunksize(x) : ForwardDiff.Chunk{C}()
     else
         error("$(C)::$(typeof(C)) is not a valid chunksize!")
     end
 end
+
+__chunksize(x) = ForwardDiff.Chunk(x)
+__chunksize(x::StaticArray) = ForwardDiff.Chunk{ForwardDiff.pickchunksize(prod(Size(x)))}()
+
 function __chunksize(::Union{AutoSparseForwardDiff{C}, AutoForwardDiff{C}}) where {C}
     C === nothing && return nothing
     C isa Integer && !(C isa Bool) && return C ≤ 0 ? nothing : Val(C)
@@ -273,18 +293,36 @@ end
     return :(nothing)
 end
 
-function init_jacobian(c::AbstractMaybeSparseJacobianCache)
+"""
+    init_jacobian(cache::AbstractMaybeSparseJacobianCache;
+        preserve_immutable::Val = Val(false))
+
+Initialize the Jacobian based on the cache. Uses sparse jacobians if possible.
+
+If `preserve_immutable` is `true`, then the Jacobian returned might be immutable, this is
+relevant if the inputs are immutable like `StaticArrays`.
+"""
+function init_jacobian(c::AbstractMaybeSparseJacobianCache;
+        preserve_immutable::Val = Val(false))
     T = promote_type(eltype(c.fx), eltype(c.x))
-    return init_jacobian(__getfield(c, Val(:jac_prototype)), T, c.fx, c.x)
+    return init_jacobian(__getfield(c, Val(:jac_prototype)), T, c.fx, c.x;
+        preserve_immutable)
 end
-init_jacobian(::Nothing, ::Type{T}, fx, x) where {T} = similar(fx, T, length(fx), length(x))
-function init_jacobian(::Nothing, ::Type{T}, fx::StaticArray, x::StaticArray) where {T}
-    # We want to construct a MArray to preserve types
-    J = StaticArrays.MArray{Tuple{length(fx), length(x)}, T}(undef)
-    return J
+function init_jacobian(::Nothing, ::Type{T}, fx, x; kwargs...) where {T}
+    return similar(fx, T, length(fx), length(x))
+end
+function init_jacobian(::Nothing, ::Type{T}, fx::StaticArray, x::StaticArray;
+        preserve_immutable::Val{PI} = Val(true)) where {T, PI}
+    if PI
+        return StaticArrays.SArray{Tuple{length(fx), length(x)}, T}(I)
+    else
+        return StaticArrays.MArray{Tuple{length(fx), length(x)}, T}(undef)
+    end
+end
+function init_jacobian(J, ::Type{T}, fx, x; kwargs...) where {T}
+    return similar(J, T, size(J, 1), size(J, 2))
 end
-init_jacobian(J, ::Type{T}, _, _) where {T} = similar(J, T, size(J, 1), size(J, 2))
-init_jacobian(J::SparseMatrixCSC, ::Type{T}, _, _) where {T} = T.(J)
+init_jacobian(J::SparseMatrixCSC, ::Type{T}, fx, x; kwargs...) where {T} = T.(J)
 
 __maybe_copy_x(_, x) = x
 __maybe_copy_x(_, ::Nothing) = nothing
diff --git a/src/highlevel/forward_mode.jl b/src/highlevel/forward_mode.jl
@@ -71,3 +71,11 @@ function sparse_jacobian!(J::AbstractMatrix, _, cache::ForwardDiffJacobianCache,
     end
     return J
 end
+
+function sparse_jacobian_static_array(_, cache::ForwardDiffJacobianCache, f, x::SArray)
+    if cache.cache isa ForwardColorJacCache
+        return forwarddiff_color_jacobian(f, x, cache.cache)
+    else
+        return ForwardDiff.jacobian(f, x, cache.cache)
+    end
+end
