diff --git a/Project.toml b/Project.toml index 1909a306..1f699f75 100644 --- a/Project.toml +++ b/Project.toml @@ -30,6 +30,7 @@ SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" SparsityDetection = "684fba80-ace3-11e9-3d08-3bc7ed6f96df" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" +StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" [targets] -test = ["Test", "IterativeSolvers", "Random", "SafeTestsets", "Zygote", "SparsityDetection"] +test = ["Test", "IterativeSolvers", "Random", "SafeTestsets", "Zygote", "SparsityDetection", "StaticArrays"] diff --git a/README.md b/README.md index 42b1a856..bb7b3bbd 100644 --- a/README.md +++ b/README.md @@ -19,7 +19,7 @@ Suppose we had the function ```julia fcalls = 0 -function f(dx,x) +function f(dx,x) # in-place global fcalls += 1 for i in 2:length(x)-1 dx[i] = x[i-1] - 2x[i] + x[i+1] @@ -28,6 +28,17 @@ function f(dx,x) dx[end] = x[end-1] - 2x[end] nothing end + +function g(x) # out-of-place + global fcalls += 1 + dx = zero(x) + for i in 2:length(x)-1 + dx[i] = x[i-1] - 2x[i] + x[i+1] + end + dx[1] = -2x[1] + x[2] + dx[end] = x[end-1] - 2x[end] + dx +end ``` For this function, we know that the sparsity pattern of the Jacobian is a @@ -62,7 +73,10 @@ In addition, a faster forward-mode autodiff call can be utilized as well: ```julia forwarddiff_color_jacobian!(jac, f, x, colorvec = colors) +jacout = forwarddiff_color_jacobian(g, x, colorvec = colors, sparsity = similar(jac), jac_prototype = similar(jac)) ``` +One can specify the type of the output jacobian by giving an additional `jac_prototype` to +the out-of place `forwarddiff_color_jacobian` function, otherwise it will become a dense matrix. If one only need to compute products, one can use the operators. For example, diff --git a/src/SparseDiffTools.jl b/src/SparseDiffTools.jl index e99ace8b..fd4e0e41 100644 --- a/src/SparseDiffTools.jl +++ b/src/SparseDiffTools.jl @@ -25,6 +25,7 @@ export contract_color, matrix2graph, matrix_colors, forwarddiff_color_jacobian!, + forwarddiff_color_jacobian, ForwardColorJacCache, auto_jacvec,auto_jacvec!, num_jacvec,num_jacvec!, diff --git a/src/differentiation/compute_jacobian_ad.jl b/src/differentiation/compute_jacobian_ad.jl index 51fa92c0..cf293251 100644 --- a/src/differentiation/compute_jacobian_ad.jl +++ b/src/differentiation/compute_jacobian_ad.jl @@ -41,48 +41,89 @@ function ForwardColorJacCache(f,x,_chunksize = nothing; _dx = dx end - ForwardColorJacCache(t,fx,_dx,p,colorvec,sparsity,getsize(chunksize)) end generate_chunked_partials(x,colorvec,N::Integer) = generate_chunked_partials(x,colorvec,Val(N)) function generate_chunked_partials(x,colorvec,::Val{chunksize}) where chunksize + maxcolor = maximum(colorvec) + num_of_chunks = Int(ceil(maxcolor / chunksize)) + padding_size = (chunksize - (maxcolor % chunksize)) % chunksize + partials = colorvec .== (1:maxcolor)' + padding_matrix = BitMatrix(undef, length(x), padding_size) + partials = hcat(partials, padding_matrix) - num_of_chunks = Int(ceil(maximum(colorvec) / chunksize)) + chunked_partials = map(i -> Tuple.(eachrow(partials[:,(i-1)*chunksize+1:i*chunksize])),1:num_of_chunks) + chunked_partials - padding_size = (chunksize - (maximum(colorvec) % chunksize)) % chunksize +end - 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) +function forwarddiff_color_jacobian(f, + x::AbstractArray{<:Number}; + dx = nothing, + colorvec = 1:length(x), + sparsity = nothing, + jac_prototype = nothing) + forwarddiff_color_jacobian(f,x,ForwardColorJacCache(f,x,dx=dx,colorvec=colorvec,sparsity=sparsity),jac_prototype) +end - for color_i in 1:maximum(colorvec) - for j in 1:length(x) - partials[j,color_i] = colorvec[j]==color_i - end - end +function forwarddiff_color_jacobian(f,x::AbstractArray{<:Number},jac_cache::ForwardColorJacCache,jac_prototype=nothing) + t = jac_cache.t + fx = jac_cache.fx + dx = jac_cache.dx + p = jac_cache.p + colorvec = jac_cache.colorvec + sparsity = jac_cache.sparsity + chunksize = jac_cache.chunksize + color_i = 1 + maxcolor = maximum(colorvec) - padding_matrix = BitMatrix(undef, length(x), padding_size) - partials = hcat(partials, padding_matrix) + vecx = vec(x) + + ncols=length(x) + J = jac_prototype isa Nothing ? false .* x .* x' : zero(jac_prototype) - 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)) + if !(sparsity isa Nothing) + rows_index, cols_index = ArrayInterface.findstructralnz(sparsity) + rows_index = [rows_index[i] for i in 1:length(rows_index)] + cols_index = [cols_index[i] for i in 1:length(cols_index)] end - chunked_partials - + for i in eachindex(p) + partial_i = p[i] + t = reshape(Dual{typeof(f)}.(vecx, partial_i),size(t)) + fx = f(t) + if !(sparsity isa Nothing) + for j in 1:chunksize + dx = vec(partials.(fx, j)) + pick_inds = [i for i in 1:length(rows_index) if colorvec[cols_index[i]] == color_i] + rows_index_c = rows_index[pick_inds] + cols_index_c = cols_index[pick_inds] + len_rows = length(pick_inds) + unused_rows = setdiff(1:ncols,rows_index_c) + perm_rows = sortperm(vcat(rows_index_c,unused_rows)) + cols_index_c = vcat(cols_index_c,zeros(Int,ncols-len_rows))[perm_rows] + Ji = [j==cols_index_c[i] ? dx[i] : false for i in 1:ncols, j in 1:ncols] + J = J + Ji + color_i += 1 + (color_i > maxcolor) && return J + end + else + for j in 1:chunksize + col_index = (i-1)*chunksize + j + (col_index > maxcolor) && return J + J = J + mapreduce(i -> i==col_index ? partials.(vec(fx), j) : zeros(ncols), hcat, 1:ncols) + end + end + end + J end function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number}, f, x::AbstractArray{<:Number}; dx = nothing, - colorvec = eachindex(x), + colorvec = 1:length(x), sparsity = ArrayInterface.has_sparsestruct(J) ? J : nothing) forwarddiff_color_jacobian!(J,f,x,ForwardColorJacCache(f,x,dx=dx,colorvec=colorvec,sparsity=sparsity)) end @@ -100,6 +141,7 @@ function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number}, sparsity = jac_cache.sparsity chunksize = jac_cache.chunksize color_i = 1 + maxcolor = maximum(colorvec) fill!(J, zero(eltype(J))) if DiffEqDiffTools._use_findstructralnz(sparsity) @@ -124,6 +166,7 @@ function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number}, for j in 1:chunksize dx .= partials.(fx, j) if ArrayInterface.fast_scalar_indexing(dx) + #dx is implicitly used in vecdx DiffEqDiffTools._colorediteration!(J,sparsity,rows_index,cols_index,vecdx,colorvec,color_i,ncols) else #= @@ -139,12 +182,12 @@ function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number}, end end color_i += 1 - (color_i > maximum(colorvec)) && return + (color_i > maxcolor) && return end else for j in 1:chunksize col_index = (i-1)*chunksize + j - (col_index > maximum(colorvec)) && return + (col_index > maxcolor) && return J[:, col_index] .= partials.(vecfx, j) end end diff --git a/test/test_ad.jl b/test/test_ad.jl index 2d851099..db8402de 100644 --- a/test/test_ad.jl +++ b/test/test_ad.jl @@ -3,6 +3,7 @@ using ForwardDiff: Dual, jacobian using SparseArrays, Test using LinearAlgebra using BlockBandedMatrices +using StaticArrays fcalls = 0 function f(dx,x) @@ -15,6 +16,23 @@ function f(dx,x) nothing end +function oopf(x) + global fcalls += 1 + dx = zero(x) + for i in 2:length(x)-1 + dx[i] = x[i-1] - 2x[i] + x[i+1] + end + dx[1] = -2x[1] + x[2] + dx[end] = x[end-1] - 2x[end] + dx +end + +function staticf(x,N=length(x)) + global fcalls += 1 + SVector{N}([i == 1 ? -2x[1]+x[2] : (i == N ? x[N-1]-2x[N] : x[i-1]-2x[i]+x[i+1]) for i in 1:N]) +end + + function second_derivative_stencil(N) A = zeros(N,N) for i in 1:N, j in 1:N @@ -39,6 +57,29 @@ forwarddiff_color_jacobian!(_J1, f, x, colorvec = repeat(1:3,10)) @test fcalls == 1 fcalls = 0 +_J1 = forwarddiff_color_jacobian(oopf, x, colorvec = repeat(1:3,10), sparsity = _J, jac_prototype = _J) +@test _J1 ≈ J +@test fcalls == 1 + +fcalls = 0 +_J1 = forwarddiff_color_jacobian(oopf, x, colorvec = repeat(1:3,10), sparsity = _J) +@test _J1 ≈ J +@test fcalls == 1 + +fcalls = 0 +_J1 = forwarddiff_color_jacobian(staticf, SVector{30}(x), colorvec = repeat(1:3,10), sparsity = _J, jac_prototype = SMatrix{30,30}(_J)) +@test _J1 ≈ J +@test fcalls == 1 + +_J1 = forwarddiff_color_jacobian(staticf, SVector{30}(x), jac_prototype = SMatrix{30,30}(_J)) +@test _J1 ≈ J +_J1 = forwarddiff_color_jacobian(oopf, x, jac_prototype = _J) +@test _J1 ≈ J +_J1 = forwarddiff_color_jacobian(oopf, x) +@test _J1 ≈ J + +fcalls = 0 +_J1 = similar(_J) jac_cache = ForwardColorJacCache(f,x,colorvec = repeat(1:3,10), sparsity = _J1) forwarddiff_color_jacobian!(_J1, f, x, jac_cache) @test _J1 ≈ J