Merge pull request #35 from pkj-m/chunksize

Limit size of partials by chunk size

Author: ChrisRackauckas

src/differentiation/compute_jacobian_ad.jl

@@ -44,20 +44,34 @@ end
 generate_chunked_partials(x,color,N::Integer) = generate_chunked_partials(x,color,Val(N))
 function generate_chunked_partials(x,color,::Val{N}) where N
 
-    # TODO: should only go up to the chunksize each time, and should
-    # generate p[i] different parts, each with less than the chunksize
+    chunksize = getsize(default_chunk_size(maximum(color)))
+    num_of_chunks = Int64(ceil(maximum(color) / chunksize))
+
+    padding_size = (chunksize - (maximum(color) % chunksize)) % chunksize
+
+    partials = BitMatrix(undef, length(x), maximum(color))
+    partial = BitMatrix(undef, length(x), chunksize)
+    chunked_partials = Array{Array{NTuple,1},1}(undef, num_of_chunks)
 
-    partials_array = BitMatrix(undef, length(x), maximum(color))
     for color_i in 1:maximum(color)
-        for i in eachindex(x)
-            if color[i]==color_i
-                partials_array[i,color_i] = true
-            else
-                partials_array[i,color_i] = false
-            end
+        for j in 1:length(x)
+            partials[j, color_i] = color[j] == color_i
+        end
+    end
+
+    padding_matrix = BitMatrix(undef, length(x), padding_size)
+    partials = hcat(partials, padding_matrix)
+
+    for i in 1:num_of_chunks
+        partial[:,1] .= partials[:,(i-1)*chunksize+1]
+        for j in 2:chunksize
+            partial[:,j] .= partials[:,(i-1)*chunksize+j]
         end
+        chunked_partials[i] = Tuple.(eachrow(partial))
     end
-    p = Tuple.(eachrow(partials_array))
+
+    chunked_partials
+
 end
 
 function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number},
@@ -79,23 +93,34 @@ function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number},
     p = jac_cache.p
     color = jac_cache.color
 
-    # TODO: Should compute on each p[i] and decompress
-    t .= Dual{typeof(f)}.(x, p)
-    f(fx, t)
-
-    if J isa SparseMatrixCSC
-        rows_index, cols_index, val = findnz(J)
-        for color_i in 1:maximum(color)
-            dx .= partials.(fx,color_i)
-            for i in 1:length(cols_index)
-                if color[cols_index[i]]==color_i
-                    J[rows_index[i],cols_index[i]] = dx[rows_index[i]]
+    color_i = 1
+    chunksize = getsize(default_chunk_size(maximum(color)))
+
+    for i in 1:length(p)
+        partial_i = p[i]
+        t .= Dual{typeof(f)}.(x, partial_i)
+        f(fx,t)
+
+        if J isa SparseMatrixCSC
+            rows_index, cols_index, val = findnz(J)
+            for j in 1:chunksize
+                dx .= partials.(fx, j)
+                for k in 1:length(cols_index)
+                    if color[cols_index[k]] == color_i
+                        J[rows_index[k], cols_index[k]] = dx[rows_index[k]]
+                    end
+                end
+                color_i += 1
+            end
+        else
+            for j in 1:chunksize
+                col_index = (i-1)*chunksize + j
+                J[:, col_index] .= partials.(fx, color_i)
+                color_i += 1
+                if color_i == maximum(color) + 1
+                    color_i = 1
                 end
             end
-        end
-    else # Compute the compressed version
-        for color_i in 1:maximum(color)
-            J[:,i] .= partials.(fx,color_i)
         end
     end
     nothing