colorvec

Author: ChrisRackauckas

Project.toml

@@ -1,7 +1,7 @@
 name = "SparseDiffTools"
 uuid = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
 authors = ["Pankaj Mishra <[email protected]>", "Chris Rackauckas <[email protected]>"]
-version = "0.7.0"
+version = "0.8.0"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/differentiation/compute_jacobian_ad.jl

@@ -3,7 +3,7 @@ struct ForwardColorJacCache{T,T2,T3,T4,T5,T6}
     fx::T2
     dx::T3
     p::T4
-    color::T5
+    colorvec::T5
     sparsity::T6
     chunksize::Int
 end
@@ -21,16 +21,16 @@ getsize(N::Integer) = N
 
 function ForwardColorJacCache(f,x,_chunksize = nothing;
                               dx = nothing,
-                              color=1:length(x),
+                              colorvec=1:length(x),
                               sparsity::Union{AbstractArray,Nothing}=nothing)
 
     if _chunksize === nothing
-        chunksize = default_chunk_size(maximum(color))
+        chunksize = default_chunk_size(maximum(colorvec))
     else
         chunksize = _chunksize
     end
 
-    p = adapt.(typeof(x),generate_chunked_partials(x,color,chunksize))
+    p = adapt.(typeof(x),generate_chunked_partials(x,colorvec,chunksize))
     t = Dual{typeof(f)}.(x,first(p))
 
     if dx === nothing
@@ -42,24 +42,24 @@ function ForwardColorJacCache(f,x,_chunksize = nothing;
     end
 
 
-    ForwardColorJacCache(t,fx,_dx,p,color,sparsity,getsize(chunksize))
+    ForwardColorJacCache(t,fx,_dx,p,colorvec,sparsity,getsize(chunksize))
 end
 
-generate_chunked_partials(x,color,N::Integer) = generate_chunked_partials(x,color,Val(N))
-function generate_chunked_partials(x,color,::Val{chunksize}) where chunksize
+generate_chunked_partials(x,colorvec,N::Integer) = generate_chunked_partials(x,colorvec,Val(N))
+function generate_chunked_partials(x,colorvec,::Val{chunksize}) where chunksize
 
-    num_of_chunks = Int(ceil(maximum(color) / chunksize))
+    num_of_chunks = Int(ceil(maximum(colorvec) / chunksize))
 
-    padding_size = (chunksize - (maximum(color) % chunksize)) % chunksize
+    padding_size = (chunksize - (maximum(colorvec) % chunksize)) % chunksize
 
-    partials = BitMatrix(undef, length(x), maximum(color))
+    partials = BitMatrix(undef, length(x), maximum(colorvec))
     partial = BitMatrix(undef, length(x), chunksize)
     chunked_partials = Array{Array{Tuple{Vararg{Bool,chunksize}},1},1}(
                                                           undef, num_of_chunks)
 
-    for color_i in 1:maximum(color)
+    for color_i in 1:maximum(colorvec)
         for j in 1:length(x)
-            partials[j,color_i] = color[j]==color_i
+            partials[j,color_i] = colorvec[j]==color_i
         end
     end
 
@@ -82,9 +82,9 @@ function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number},
                 f,
                 x::AbstractArray{<:Number};
                 dx = nothing,
-                color = eachindex(x),
+                colorvec = eachindex(x),
                 sparsity = ArrayInterface.has_sparsestruct(J) ? J : nothing)
-    forwarddiff_color_jacobian!(J,f,x,ForwardColorJacCache(f,x,dx=dx,color=color,sparsity=sparsity))
+    forwarddiff_color_jacobian!(J,f,x,ForwardColorJacCache(f,x,dx=dx,colorvec=colorvec,sparsity=sparsity))
 end
 
 function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number},
@@ -96,7 +96,7 @@ function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number},
     fx = jac_cache.fx
     dx = jac_cache.dx
     p = jac_cache.p
-    color = jac_cache.color
+    colorvec = jac_cache.colorvec
     sparsity = jac_cache.sparsity
     chunksize = jac_cache.chunksize
     color_i = 1
@@ -112,7 +112,7 @@ function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number},
                 dx .= partials.(fx, j)
                 if ArrayInterface.fast_scalar_indexing(dx)
                     for k in 1:length(cols_index)
-                        if color[cols_index[k]] == color_i
+                        if colorvec[cols_index[k]] == color_i
                             if J isa SparseMatrixCSC
                                 J.nzval[k] = dx[rows_index[k]]
                             else
@@ -122,24 +122,24 @@ function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number},
                     end
                 else
                     #=
-                    J.nzval[rows_index] .+= (color[cols_index] .== color_i) .* dx[rows_index]
+                    J.nzval[rows_index] .+= (colorvec[cols_index] .== color_i) .* dx[rows_index]
                     or
-                    J[rows_index, cols_index] .+= (color[cols_index] .== color_i) .* dx[rows_index]
+                    J[rows_index, cols_index] .+= (colorvec[cols_index] .== color_i) .* dx[rows_index]
                     += means requires a zero'd out start
                     =#
                     if J isa SparseMatrixCSC
-                        @. setindex!((J.nzval,),getindex((J.nzval,),rows_index) + (getindex((color,),cols_index) == color_i) * getindex((dx,),rows_index),rows_index)
+                        @. setindex!((J.nzval,),getindex((J.nzval,),rows_index) + (getindex((colorvec,),cols_index) == color_i) * getindex((dx,),rows_index),rows_index)
                     else
-                        @. setindex!((J,),getindex((J,),rows_index, cols_index) + (getindex((color,),cols_index) == color_i) * getindex((dx,),rows_index),rows_index, cols_index)
+                        @. setindex!((J,),getindex((J,),rows_index, cols_index) + (getindex((colorvec,),cols_index) == color_i) * getindex((dx,),rows_index),rows_index, cols_index)
                     end
                 end
                 color_i += 1
-                (color_i > maximum(color)) && continue
+                (color_i > maximum(colorvec)) && continue
             end
         else
             for j in 1:chunksize
                 col_index = (i-1)*chunksize + j
-                (col_index > maximum(color)) && continue
+                (col_index > maximum(colorvec)) && continue
                 J[:, col_index] .= partials.(fx, j)
             end
         end