From a26c079e6a3a60cf48b4c73f5ce579ac81b0703c Mon Sep 17 00:00:00 2001 From: Langwen Huang Date: Sun, 6 Oct 2019 19:34:30 +0200 Subject: [PATCH 1/2] correction for nonsquare jacobian --- README.md | 11 ++- src/differentiation/compute_jacobian_ad.jl | 35 +++++---- test/test_ad.jl | 91 +++++++++++++++++++++- 3 files changed, 120 insertions(+), 17 deletions(-) diff --git a/README.md b/README.md index bb7b3bbd..ab605f9e 100644 --- a/README.md +++ b/README.md @@ -73,10 +73,17 @@ 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)) +#OR +jacout = forwarddiff_color_jacobian(g, x, dx = similar(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 +One can specify the type and shape 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 `jac_prototype` and `sparsity` is not specified (both default to `nothing`), the oop jacobian will +assume that the function has a *square* jacobian matrix. If it is not the case, please specify +the shape of output by giving `dx` or giving `dx = nothing` where the shape is determined +automatically at the cost of a function evaluation `f(x)`. +If not necessary (e.g `StaticArrays` involved), please turn to the inplace jacobian function +`forwarddiff_color_jacobian!` which is faster in general cases. If one only need to compute products, one can use the operators. For example, diff --git a/src/differentiation/compute_jacobian_ad.jl b/src/differentiation/compute_jacobian_ad.jl index cf293251..3a286120 100644 --- a/src/differentiation/compute_jacobian_ad.jl +++ b/src/differentiation/compute_jacobian_ad.jl @@ -24,7 +24,7 @@ function ForwardColorJacCache(f,x,_chunksize = nothing; colorvec=1:length(x), sparsity::Union{AbstractArray,Nothing}=nothing) - if _chunksize === nothing + if _chunksize isa Nothing chunksize = default_chunk_size(maximum(colorvec)) else chunksize = _chunksize @@ -33,11 +33,17 @@ function ForwardColorJacCache(f,x,_chunksize = nothing; p = adapt.(typeof(x),generate_chunked_partials(x,colorvec,chunksize)) t = reshape(Dual{typeof(f)}.(vec(x),first(p)),size(x)...) - if dx === nothing + if dx isa Nothing fx = similar(t) _dx = similar(x) else - fx = reshape(Dual{typeof(f)}.(vec(dx),first(p)),size(x)...) + pi = first(p) #perform trim (length(dx)length(x)) to first(p) + if length(dx)>length(x) + pi = vcat(pi,reshape(mapslices(Tuple,zeros(Bool,length(first(pi)),length(dx)-length(x)),dims=1),:)) + else + pi = pi[1:length(dx)] + end + fx = reshape(Dual{typeof(f)}.(vec(dx),pi),size(dx)...) _dx = dx end @@ -60,16 +66,18 @@ end function forwarddiff_color_jacobian(f, x::AbstractArray{<:Number}; - dx = nothing, + dx = similar(x), #if fx is nothing, we will estimate fx at the cost of a function call colorvec = 1:length(x), sparsity = nothing, jac_prototype = nothing) + if dx isa Nothing + dx = f(x) + end forwarddiff_color_jacobian(f,x,ForwardColorJacCache(f,x,dx=dx,colorvec=colorvec,sparsity=sparsity),jac_prototype) 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 @@ -81,7 +89,7 @@ function forwarddiff_color_jacobian(f,x::AbstractArray{<:Number},jac_cache::Forw vecx = vec(x) ncols=length(x) - J = jac_prototype isa Nothing ? false .* x .* x' : zero(jac_prototype) + J = jac_prototype isa Nothing ? (sparsity isa Nothing ? false .* dx .* x' : zeros(eltype(x),size(sparsity))) : zero(jac_prototype) if !(sparsity isa Nothing) rows_index, cols_index = ArrayInterface.findstructralnz(sparsity) @@ -93,6 +101,7 @@ function forwarddiff_color_jacobian(f,x::AbstractArray{<:Number},jac_cache::Forw partial_i = p[i] t = reshape(Dual{typeof(f)}.(vecx, partial_i),size(t)) fx = f(t) + nrows=length(fx) if !(sparsity isa Nothing) for j in 1:chunksize dx = vec(partials.(fx, j)) @@ -100,10 +109,10 @@ function forwarddiff_color_jacobian(f,x::AbstractArray{<:Number},jac_cache::Forw 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) + unused_rows = setdiff(1:nrows,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] + cols_index_c = vcat(cols_index_c,zeros(Int,nrows-len_rows))[perm_rows] + Ji = [j==cols_index_c[i] ? dx[i] : false for i in 1:nrows, j in 1:ncols] J = J + Ji color_i += 1 (color_i > maxcolor) && return J @@ -111,8 +120,8 @@ function forwarddiff_color_jacobian(f,x::AbstractArray{<:Number},jac_cache::Forw 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) + (col_index > ncols) && return J + J = J + mapreduce(i -> i==col_index ? partials.(vec(fx), j) : zeros(nrows), hcat, 1:ncols) end end end @@ -122,7 +131,7 @@ end function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number}, f, x::AbstractArray{<:Number}; - dx = nothing, + dx = Array{eltype(x)}(undef,size(J,1)), 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)) @@ -187,7 +196,7 @@ function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number}, else for j in 1:chunksize col_index = (i-1)*chunksize + j - (col_index > maxcolor) && return + (col_index > ncols) && return J[:, col_index] .= partials.(vecfx, j) end end diff --git a/test/test_ad.jl b/test/test_ad.jl index db8402de..32ba2e0e 100644 --- a/test/test_ad.jl +++ b/test/test_ad.jl @@ -27,11 +27,53 @@ function oopf(x) dx end +function nsqf(x)#length(dx)length(x) + global fcalls +=1 + dx = zeros(eltype(x),length(x)*2) + 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 + +function nsqf!(dx,x) + global fcalls +=1 + for i in 2:length(dx) + dx[i] = x[i-1] - 2x[i] + x[i+1] + end + dx[1] = -2x[1] + x[2] + nothing +end + +function nsqf2!(dx,x) + global fcalls +=1 + for i in 2:length(x)-1 + dx[i] = x[i-1] - 2x[i] + x[i+1] + end + dx[1] = -2x[1] + x[2] + nothing +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 staticnsqf(x,N=div(length(x),2)) + global fcalls += 1 + SVector{N}(vcat([-2x[1]+x[2]],[x[i-1]-2x[i]+x[i+1] for i in 2:N])) +end function second_derivative_stencil(N) A = zeros(N,N) @@ -73,11 +115,56 @@ _J1 = forwarddiff_color_jacobian(staticf, SVector{30}(x), colorvec = repeat(1:3, _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) +_J1 = forwarddiff_color_jacobian(oopf, x, jac_prototype = similar(_J)) @test _J1 ≈ J _J1 = forwarddiff_color_jacobian(oopf, x) @test _J1 ≈ J +#Non-square Jacobian +#length(dx)length(x) +nsqJ = jacobian(nsqf2,x) +spnsqJ = sparse(nsqJ) +_nsqJ = forwarddiff_color_jacobian(nsqf2, x, dx = nothing) +@test _nsqJ ≈ nsqJ +_nsqJ = forwarddiff_color_jacobian(nsqf2, x, colorvec = repeat(1:3,10), sparsity = spnsqJ) +@test _nsqJ ≈ nsqJ +_nsqJ = forwarddiff_color_jacobian(nsqf2, x, jac_prototype = similar(nsqJ)) +@test _nsqJ ≈ nsqJ +_nsqJ = forwarddiff_color_jacobian(nsqf2, x, colorvec = repeat(1:3,10), sparsity = spnsqJ, jac_prototype = similar(nsqJ)) +@test _nsqJ ≈ nsqJ +_nsqJ = forwarddiff_color_jacobian(nsqf2, x, jac_prototype = SMatrix{60,30}(nsqJ)) +@test _nsqJ ≈ nsqJ +_nsqJ = similar(nsqJ) +forwarddiff_color_jacobian!(_nsqJ, nsqf2!, x) +@test _nsqJ ≈ nsqJ +_nsqJ = similar(nsqJ) +forwarddiff_color_jacobian!(_nsqJ, nsqf2!, x, colorvec = repeat(1:3,10), sparsity = spnsqJ ) +@test _nsqJ ≈ nsqJ + fcalls = 0 _J1 = similar(_J) jac_cache = ForwardColorJacCache(f,x,colorvec = repeat(1:3,10), sparsity = _J1) From b314d45e893e013abdc077488397d246bfdc63a0 Mon Sep 17 00:00:00 2001 From: Langwen Huang Date: Sun, 6 Oct 2019 19:38:46 +0200 Subject: [PATCH 2/2] del fx --- src/differentiation/compute_jacobian_ad.jl | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/differentiation/compute_jacobian_ad.jl b/src/differentiation/compute_jacobian_ad.jl index 3a286120..49fea45a 100644 --- a/src/differentiation/compute_jacobian_ad.jl +++ b/src/differentiation/compute_jacobian_ad.jl @@ -66,7 +66,7 @@ end function forwarddiff_color_jacobian(f, x::AbstractArray{<:Number}; - dx = similar(x), #if fx is nothing, we will estimate fx at the cost of a function call + dx = similar(x), #if dx is nothing, we will estimate dx at the cost of a function call colorvec = 1:length(x), sparsity = nothing, jac_prototype = nothing) @@ -88,8 +88,8 @@ function forwarddiff_color_jacobian(f,x::AbstractArray{<:Number},jac_cache::Forw vecx = vec(x) - ncols=length(x) J = jac_prototype isa Nothing ? (sparsity isa Nothing ? false .* dx .* x' : zeros(eltype(x),size(sparsity))) : zero(jac_prototype) + nrows,ncols = size(J) if !(sparsity isa Nothing) rows_index, cols_index = ArrayInterface.findstructralnz(sparsity) @@ -101,7 +101,6 @@ function forwarddiff_color_jacobian(f,x::AbstractArray{<:Number},jac_cache::Forw partial_i = p[i] t = reshape(Dual{typeof(f)}.(vecx, partial_i),size(t)) fx = f(t) - nrows=length(fx) if !(sparsity isa Nothing) for j in 1:chunksize dx = vec(partials.(fx, j))