From 5bbda06f61eec2948fadd78cb5b8cfaa0988da53 Mon Sep 17 00:00:00 2001 From: Langwen Huang Date: Fri, 2 Aug 2019 22:00:46 +0800 Subject: [PATCH 1/9] switch to column based iteration --- src/jacobians.jl | 60 +++++++++++++++++++----------------------------- 1 file changed, 24 insertions(+), 36 deletions(-) diff --git a/src/jacobians.jl b/src/jacobians.jl index 07eb5e4..8388e41 100644 --- a/src/jacobians.jl +++ b/src/jacobians.jl @@ -217,12 +217,10 @@ function finite_difference_jacobian!( @. vfx1 = (vfx1 - vfx) / epsilon if ArrayInterface.fast_scalar_indexing(x1) - for i in 1:length(cols_index) - if color[cols_index[i]] == color_i - if J isa SparseMatrixCSC - J.nzval[i] = vfx1[rows_index[i]] - else - J[rows_index[i],cols_index[i]] = vfx1[rows_index[i]] + for col_index in 1:n + if color[col_index] == color_i + for row_index in ArrayInterface.findstructralnz(J,col_index) + J[row_index,col_index]=vfx1[row_index] end end end @@ -252,12 +250,10 @@ function finite_difference_jacobian!( _vfx1 = (vfx1 - vfx) / epsilon if ArrayInterface.fast_scalar_indexing(x1) - for i in 1:length(cols_index) - if color[cols_index[i]] == color_i - if J isa SparseMatrixCSC - J.nzval[i] = vfx1[rows_index[i]] - else - J[rows_index[i],cols_index[i]] = vfx1[rows_index[i]] + for col_index in 1:n + if color[col_index] == color_i + for row_index in ArrayInterface.findstructralnz(J,col_index) + J[row_index,col_index]=vfx1[row_index] end end end @@ -329,12 +325,10 @@ function finite_difference_jacobian!( @. vfx1 = (vfx1 - vfx) / 2epsilon if ArrayInterface.fast_scalar_indexing(x1) - for i in 1:length(cols_index) - if color[cols_index[i]] == color_i - if J isa SparseMatrixCSC - J.nzval[i] = vfx1[rows_index[i]] - else - J[rows_index[i],cols_index[i]] = vfx1[rows_index[i]] + for col_index in 1:n + if color[col_index] == color_i + for row_index in ArrayInterface.findstructralnz(J,col_index) + J[row_index,col_index]=vfx1[row_index] end end end @@ -369,12 +363,10 @@ function finite_difference_jacobian!( # vfx1 is the compressed Jacobian column if ArrayInterface.fast_scalar_indexing(x1) - for i in 1:length(cols_index) - if color[cols_index[i]] == color_i - if J isa SparseMatrixCSC - J.nzval[i] = vfx1[rows_index[i]] - else - J[rows_index[i],cols_index[i]] = vfx1[rows_index[i]] + for col_index in 1:n + if color[col_index] == color_i + for row_index in ArrayInterface.findstructralnz(J,col_index) + J[row_index,col_index]=vfx1[row_index] end end end @@ -435,12 +427,10 @@ function finite_difference_jacobian!( @. vfx = imag(vfx) / epsilon if ArrayInterface.fast_scalar_indexing(x1) - for i in 1:length(cols_index) - if color[cols_index[i]] == color_i - if J isa SparseMatrixCSC - J.nzval[i] = vfx[rows_index[i]] - else - J[rows_index[i],cols_index[i]] = vfx[rows_index[i]] + for col_index in 1:n + if color[col_index] == color_i + for row_index in ArrayInterface.findstructralnz(J,col_index) + J[row_index,col_index]=vfx[row_index] end end end @@ -471,12 +461,10 @@ function finite_difference_jacobian!( vfx = imag(vfx) / epsilon if ArrayInterface.fast_scalar_indexing(x1) - for i in 1:length(cols_index) - if color[cols_index[i]] == color_i - if J isa SparseMatrixCSC - J.nzval[i] = vfx1[rows_index[i]] - else - J[rows_index[i],cols_index[i]] = vfx1[rows_index[i]] + for col_index in 1:n + if color[col_index] == color_i + for row_index in ArrayInterface.findstructralnz(J,col_index) + J[row_index,col_index]=vfx1[row_index] end end end From 28d05d81fea2485287dd959ceb1590d6972680bf Mon Sep 17 00:00:00 2001 From: Langwen Huang Date: Fri, 2 Aug 2019 22:51:32 +0800 Subject: [PATCH 2/9] add fast_column_indexing trait --- src/jacobians.jl | 98 ++++++++++++++++++++++++++++++++++++------------ 1 file changed, 73 insertions(+), 25 deletions(-) diff --git a/src/jacobians.jl b/src/jacobians.jl index 8388e41..90513c8 100644 --- a/src/jacobians.jl +++ b/src/jacobians.jl @@ -161,7 +161,7 @@ function finite_difference_jacobian!( end vfx = vec(fx) - if ArrayInterface.has_sparsestruct(sparsity) + if ArrayInterface.has_sparsestruct(sparsity) & !ArrayInterface.fast_column_indexing(sparsity) rows_index, cols_index = ArrayInterface.findstructralnz(sparsity) end @@ -217,10 +217,18 @@ function finite_difference_jacobian!( @. vfx1 = (vfx1 - vfx) / epsilon if ArrayInterface.fast_scalar_indexing(x1) - for col_index in 1:n - if color[col_index] == color_i - for row_index in ArrayInterface.findstructralnz(J,col_index) - J[row_index,col_index]=vfx1[row_index] + if ArrayInterface.fast_column_indexing(J) + for col_index in 1:n + if color[col_index] == color_i + for row_index in ArrayInterface.findstructralnz(J,col_index) + J[row_index,col_index]=vfx1[row_index] + end + end + end + else + for i in 1:length(cols_index) + if color[cols_index[i]] == color_i + J[rows_index[i],cols_index[i]] = vfx1[rows_index[i]] end end end @@ -250,10 +258,18 @@ function finite_difference_jacobian!( _vfx1 = (vfx1 - vfx) / epsilon if ArrayInterface.fast_scalar_indexing(x1) - for col_index in 1:n - if color[col_index] == color_i - for row_index in ArrayInterface.findstructralnz(J,col_index) - J[row_index,col_index]=vfx1[row_index] + if ArrayInterface.fast_column_indexing(J) + for col_index in 1:n + if color[col_index] == color_i + for row_index in ArrayInterface.findstructralnz(J,col_index) + J[row_index,col_index]=vfx1[row_index] + end + end + end + else + for i in 1:length(cols_index) + if color[cols_index[i]] == color_i + J[rows_index[i],cols_index[i]] = vfx1[rows_index[i]] end end end @@ -325,10 +341,18 @@ function finite_difference_jacobian!( @. vfx1 = (vfx1 - vfx) / 2epsilon if ArrayInterface.fast_scalar_indexing(x1) - for col_index in 1:n - if color[col_index] == color_i - for row_index in ArrayInterface.findstructralnz(J,col_index) - J[row_index,col_index]=vfx1[row_index] + if ArrayInterface.fast_column_indexing(J) + for col_index in 1:n + if color[col_index] == color_i + for row_index in ArrayInterface.findstructralnz(J,col_index) + J[row_index,col_index]=vfx1[row_index] + end + end + end + else + for i in 1:length(cols_index) + if color[cols_index[i]] == color_i + J[rows_index[i],cols_index[i]] = vfx1[rows_index[i]] end end end @@ -363,10 +387,18 @@ function finite_difference_jacobian!( # vfx1 is the compressed Jacobian column if ArrayInterface.fast_scalar_indexing(x1) - for col_index in 1:n - if color[col_index] == color_i - for row_index in ArrayInterface.findstructralnz(J,col_index) - J[row_index,col_index]=vfx1[row_index] + if ArrayInterface.fast_column_indexing(J) + for col_index in 1:n + if color[col_index] == color_i + for row_index in ArrayInterface.findstructralnz(J,col_index) + J[row_index,col_index]=vfx1[row_index] + end + end + end + else + for i in 1:length(cols_index) + if color[cols_index[i]] == color_i + J[rows_index[i],cols_index[i]] = vfx1[rows_index[i]] end end end @@ -427,10 +459,18 @@ function finite_difference_jacobian!( @. vfx = imag(vfx) / epsilon if ArrayInterface.fast_scalar_indexing(x1) - for col_index in 1:n - if color[col_index] == color_i - for row_index in ArrayInterface.findstructralnz(J,col_index) - J[row_index,col_index]=vfx[row_index] + if ArrayInterface.fast_column_indexing(J) + for col_index in 1:n + if color[col_index] == color_i + for row_index in ArrayInterface.findstructralnz(J,col_index) + J[row_index,col_index]=vfx[row_index] + end + end + end + else + for i in 1:length(cols_index) + if color[cols_index[i]] == color_i + J[rows_index[i],cols_index[i]] = vfx[rows_index[i]] end end end @@ -461,10 +501,18 @@ function finite_difference_jacobian!( vfx = imag(vfx) / epsilon if ArrayInterface.fast_scalar_indexing(x1) - for col_index in 1:n - if color[col_index] == color_i - for row_index in ArrayInterface.findstructralnz(J,col_index) - J[row_index,col_index]=vfx1[row_index] + if ArrayInterface.fast_column_indexing(J) + for col_index in 1:n + if color[col_index] == color_i + for row_index in ArrayInterface.findstructralnz(J,col_index) + J[row_index,col_index]=vfx1[row_index] + end + end + end + else + for i in 1:length(cols_index) + if color[cols_index[i]] == color_i + J[rows_index[i],cols_index[i]] = vfx1[rows_index[i]] end end end From f7822445dc22de12d457219fcdcc89beefb229bc Mon Sep 17 00:00:00 2001 From: Langwen Huang Date: Sat, 3 Aug 2019 21:39:24 +0800 Subject: [PATCH 3/9] findnz on sparsity instead of J --- src/jacobians.jl | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/src/jacobians.jl b/src/jacobians.jl index 90513c8..6c404f6 100644 --- a/src/jacobians.jl +++ b/src/jacobians.jl @@ -217,10 +217,10 @@ function finite_difference_jacobian!( @. vfx1 = (vfx1 - vfx) / epsilon if ArrayInterface.fast_scalar_indexing(x1) - if ArrayInterface.fast_column_indexing(J) + if ArrayInterface.fast_column_indexing(sparsity) for col_index in 1:n if color[col_index] == color_i - for row_index in ArrayInterface.findstructralnz(J,col_index) + for row_index in ArrayInterface.findstructralnz(sparsity,col_index) J[row_index,col_index]=vfx1[row_index] end end @@ -258,10 +258,10 @@ function finite_difference_jacobian!( _vfx1 = (vfx1 - vfx) / epsilon if ArrayInterface.fast_scalar_indexing(x1) - if ArrayInterface.fast_column_indexing(J) + if ArrayInterface.fast_column_indexing(sparsity) for col_index in 1:n if color[col_index] == color_i - for row_index in ArrayInterface.findstructralnz(J,col_index) + for row_index in ArrayInterface.findstructralnz(sparsity,col_index) J[row_index,col_index]=vfx1[row_index] end end @@ -341,10 +341,10 @@ function finite_difference_jacobian!( @. vfx1 = (vfx1 - vfx) / 2epsilon if ArrayInterface.fast_scalar_indexing(x1) - if ArrayInterface.fast_column_indexing(J) + if ArrayInterface.fast_column_indexing(sparsity) for col_index in 1:n if color[col_index] == color_i - for row_index in ArrayInterface.findstructralnz(J,col_index) + for row_index in ArrayInterface.findstructralnz(sparsity,col_index) J[row_index,col_index]=vfx1[row_index] end end @@ -387,10 +387,10 @@ function finite_difference_jacobian!( # vfx1 is the compressed Jacobian column if ArrayInterface.fast_scalar_indexing(x1) - if ArrayInterface.fast_column_indexing(J) + if ArrayInterface.fast_column_indexing(sparsity) for col_index in 1:n if color[col_index] == color_i - for row_index in ArrayInterface.findstructralnz(J,col_index) + for row_index in ArrayInterface.findstructralnz(sparsity,col_index) J[row_index,col_index]=vfx1[row_index] end end @@ -459,10 +459,10 @@ function finite_difference_jacobian!( @. vfx = imag(vfx) / epsilon if ArrayInterface.fast_scalar_indexing(x1) - if ArrayInterface.fast_column_indexing(J) + if ArrayInterface.fast_column_indexing(sparsity) for col_index in 1:n if color[col_index] == color_i - for row_index in ArrayInterface.findstructralnz(J,col_index) + for row_index in ArrayInterface.findstructralnz(sparsity,col_index) J[row_index,col_index]=vfx[row_index] end end @@ -501,10 +501,10 @@ function finite_difference_jacobian!( vfx = imag(vfx) / epsilon if ArrayInterface.fast_scalar_indexing(x1) - if ArrayInterface.fast_column_indexing(J) + if ArrayInterface.fast_column_indexing(sparsity) for col_index in 1:n if color[col_index] == color_i - for row_index in ArrayInterface.findstructralnz(J,col_index) + for row_index in ArrayInterface.findstructralnz(sparsity,col_index) J[row_index,col_index]=vfx1[row_index] end end From 5d3302b34d83905d6b20d2b47f72d08f3bf1e6c3 Mon Sep 17 00:00:00 2001 From: Langwen Huang Date: Sun, 4 Aug 2019 15:29:03 +0800 Subject: [PATCH 4/9] simplify using macro --- src/jacobians.jl | 108 ++++++++++++++++------------------------------- 1 file changed, 36 insertions(+), 72 deletions(-) diff --git a/src/jacobians.jl b/src/jacobians.jl index eb4097c..4fce27a 100644 --- a/src/jacobians.jl +++ b/src/jacobians.jl @@ -124,6 +124,30 @@ function finite_difference_jacobian(f, x::AbstractArray{<:Number}, finite_difference_jacobian(f, x, cache, f_in; relstep=relstep, absstep=absstep, colorvec=colorvec, sparsity=sparsity, dir=dir) end +macro _colorediteration(mode,vfx) + if mode.value==:columniteration + return esc(quote + for col_index in 1:n + if colorvec[col_index] == color_i + for row_index in ArrayInterface.findstructralnz(sparsity,col_index) + J[row_index,col_index]=$vfx[row_index] + end + end + end + end) + elseif mode.value==:indexing + return esc(quote + for i in 1:length(cols_index) + if colorvec[cols_index[i]] == color_i + J[rows_index[i],cols_index[i]] = $vfx[rows_index[i]] + end + end + end) + else + throw(ErrorException("Can't recognize mode: $mode")) + end +end + function finite_difference_jacobian( f, x, @@ -218,19 +242,9 @@ function finite_difference_jacobian!( if ArrayInterface.fast_scalar_indexing(x1) if ArrayInterface.fast_column_indexing(sparsity) - for col_index in 1:n - if colorvec[col_index] == color_i - for row_index in ArrayInterface.findstructralnz(sparsity,col_index) - J[row_index,col_index]=vfx1[row_index] - end - end - end + @_colorediteration(:columniteration,vfx1) else - for i in 1:length(cols_index) - if colorvec[cols_index[i]] == color_i - J[rows_index[i],cols_index[i]] = vfx1[rows_index[i]] - end - end + @_colorediteration(:indexing,vfx1) end else #= @@ -259,19 +273,9 @@ function finite_difference_jacobian!( if ArrayInterface.fast_scalar_indexing(x1) if ArrayInterface.fast_column_indexing(sparsity) - for col_index in 1:n - if colorvec[col_index] == color_i - for row_index in ArrayInterface.findstructralnz(sparsity,col_index) - J[row_index,col_index]=vfx1[row_index] - end - end - end + @_colorediteration(:columniteration,vfx1) else - for i in 1:length(cols_index) - if colorvec[cols_index[i]] == color_i - J[rows_index[i],cols_index[i]] = vfx1[rows_index[i]] - end - end + @_colorediteration(:indexing,vfx1) end else #= @@ -342,19 +346,9 @@ function finite_difference_jacobian!( if ArrayInterface.fast_scalar_indexing(x1) if ArrayInterface.fast_column_indexing(sparsity) - for col_index in 1:n - if colorvec[col_index] == color_i - for row_index in ArrayInterface.findstructralnz(sparsity,col_index) - J[row_index,col_index]=vfx1[row_index] - end - end - end + @_colorediteration(:columniteration,vfx1) else - for i in 1:length(cols_index) - if colorvec[cols_index[i]] == color_i - J[rows_index[i],cols_index[i]] = vfx1[rows_index[i]] - end - end + @_colorediteration(:indexing,vfx1) end else #= @@ -388,19 +382,9 @@ function finite_difference_jacobian!( if ArrayInterface.fast_scalar_indexing(x1) if ArrayInterface.fast_column_indexing(sparsity) - for col_index in 1:n - if colorvec[col_index] == color_i - for row_index in ArrayInterface.findstructralnz(sparsity,col_index) - J[row_index,col_index]=vfx1[row_index] - end - end - end + @_colorediteration(:columniteration,vfx1) else - for i in 1:length(cols_index) - if colorvec[cols_index[i]] == color_i - J[rows_index[i],cols_index[i]] = vfx1[rows_index[i]] - end - end + @_colorediteration(:indexing,vfx1) end else #= @@ -460,19 +444,9 @@ function finite_difference_jacobian!( if ArrayInterface.fast_scalar_indexing(x1) if ArrayInterface.fast_column_indexing(sparsity) - for col_index in 1:n - if colorvec[col_index] == color_i - for row_index in ArrayInterface.findstructralnz(sparsity,col_index) - J[row_index,col_index]=vfx[row_index] - end - end - end + @_colorediteration(:columniteration,vfx) else - for i in 1:length(cols_index) - if colorvec[cols_index[i]] == color_i - J[rows_index[i],cols_index[i]] = vfx[rows_index[i]] - end - end + @_colorediteration(:indexing,vfx) end else #= @@ -502,19 +476,9 @@ function finite_difference_jacobian!( if ArrayInterface.fast_scalar_indexing(x1) if ArrayInterface.fast_column_indexing(sparsity) - for col_index in 1:n - if colorvec[col_index] == color_i - for row_index in ArrayInterface.findstructralnz(sparsity,col_index) - J[row_index,col_index]=vfx1[row_index] - end - end - end + @_colorediteration(:columniteration,vfx1) else - for i in 1:length(cols_index) - if colorvec[cols_index[i]] == color_i - J[rows_index[i],cols_index[i]] = vfx1[rows_index[i]] - end - end + @_colorediteration(:indexing,vfx1) end else #= From bffa23df3cb09c8402efc4cb768b268b4836ced3 Mon Sep 17 00:00:00 2001 From: Langwen Huang Date: Wed, 7 Aug 2019 23:09:20 +0800 Subject: [PATCH 5/9] optimize for BBB matrix --- Project.toml | 1 + src/DiffEqDiffTools.jl | 2 +- src/jacobians.jl | 101 +++++++++++++++++++++-------------------- test/coloring_tests.jl | 23 +++++++++- 4 files changed, 76 insertions(+), 51 deletions(-) diff --git a/Project.toml b/Project.toml index 1e2c289..41d618d 100644 --- a/Project.toml +++ b/Project.toml @@ -7,6 +7,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9" +BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0" [compat] julia = "1" diff --git a/src/DiffEqDiffTools.jl b/src/DiffEqDiffTools.jl index 64ef94a..7d1c31b 100644 --- a/src/DiffEqDiffTools.jl +++ b/src/DiffEqDiffTools.jl @@ -2,7 +2,7 @@ __precompile__() module DiffEqDiffTools -using LinearAlgebra, SparseArrays, StaticArrays, ArrayInterface +using LinearAlgebra, SparseArrays, StaticArrays, ArrayInterface, BlockBandedMatrices import Base: resize! diff --git a/src/jacobians.jl b/src/jacobians.jl index 4fce27a..9b9f7bb 100644 --- a/src/jacobians.jl +++ b/src/jacobians.jl @@ -124,27 +124,49 @@ function finite_difference_jacobian(f, x::AbstractArray{<:Number}, finite_difference_jacobian(f, x, cache, f_in; relstep=relstep, absstep=absstep, colorvec=colorvec, sparsity=sparsity, dir=dir) end -macro _colorediteration(mode,vfx) - if mode.value==:columniteration - return esc(quote - for col_index in 1:n - if colorvec[col_index] == color_i - for row_index in ArrayInterface.findstructralnz(sparsity,col_index) - J[row_index,col_index]=$vfx[row_index] +@inline function _colorediteration!(J,sparsity,rows_index,cols_index,vfx,colorvec,color_i,ncols) + @inbounds for i in 1:length(cols_index) + if colorvec[cols_index[i]] == color_i + J[rows_index[i],cols_index[i]] = vfx[rows_index[i]] + end + end +end + +@inline function _colorediteration!(J,sparsity::SparseMatrixCSC,rows_index,cols_index,vfx,colorvec,color_i,ncols) + @inbounds for col_index in 1:ncols + if colorvec[col_index] == color_i + @inbounds for row_index in view(sparsity.rowval,sparsity.colptr[col_index]:sparsity.colptr[col_index+1]-1) + J[row_index,col_index]=vfx[row_index] + end + end + end +end + +@inline function _colorediteration!(Jac::BlockBandedMatrices.BandedBlockBandedMatrix, + sparsity::BlockBandedMatrices.BandedBlockBandedMatrix, + rows_index,cols_index,vfx,colorvec,color_i,ncols) + λ,μ = subblockbandwidths(Jac) + rs = BlockBandedMatrices.BlockSizes((BlockBandedMatrices.cumulsizes(Jac,1),)) # column block sizes + cs = BlockBandedMatrices.BlockSizes((BlockBandedMatrices.cumulsizes(Jac,2),)) + b = BlockBandedMatrices.BlockArray(vfx,rs) + c = BlockBandedMatrices.BlockArray(colorvec,cs) + @inbounds for J=Block.(1:BlockBandedMatrices.nblocks(Jac,2)) + c_v = c.blocks[J.n[1]] + @inbounds for K=BlockBandedMatrices.blockcolrange(Jac,J) + V = view(Jac,K,J) + b_v = b.blocks[K.n[1]] + data = BlockBandedMatrices.bandeddata(V) + p = pointer(data) + st = stride(data,2) + m,n = size(V) + @inbounds for j=1:n + if c_v[j] == color_i + @inbounds for k=max(1,j-μ):min(m,j+λ) + unsafe_store!(p, b_v[k], (j-1)*st + μ + k - j + 1) end end end - end) - elseif mode.value==:indexing - return esc(quote - for i in 1:length(cols_index) - if colorvec[cols_index[i]] == color_i - J[rows_index[i],cols_index[i]] = $vfx[rows_index[i]] - end end - end) - else - throw(ErrorException("Can't recognize mode: $mode")) end end @@ -164,6 +186,11 @@ function finite_difference_jacobian( _J isa SMatrix ? SArray(J) : J end +#override default setting of using findstructralnz +_use_findstructralnz(sparsity) = ArrayInterface.has_sparsestruct(sparsity) +_use_findstructralnz(::SparseMatrixCSC) = false +_use_findstructralnz(::BlockBandedMatrices.BandedBlockBandedMatrix) = false + function finite_difference_jacobian!( J::AbstractMatrix{<:Number}, f, @@ -185,7 +212,9 @@ function finite_difference_jacobian!( end vfx = vec(fx) - if ArrayInterface.has_sparsestruct(sparsity) & !ArrayInterface.fast_column_indexing(sparsity) + rows_index = nothing + cols_index = nothing + if _use_findstructralnz(sparsity) rows_index, cols_index = ArrayInterface.findstructralnz(sparsity) end @@ -241,11 +270,7 @@ function finite_difference_jacobian!( @. vfx1 = (vfx1 - vfx) / epsilon if ArrayInterface.fast_scalar_indexing(x1) - if ArrayInterface.fast_column_indexing(sparsity) - @_colorediteration(:columniteration,vfx1) - else - @_colorediteration(:indexing,vfx1) - end + _colorediteration!(J,sparsity,rows_index,cols_index,vfx1,colorvec,color_i,n) else #= J.nzval[rows_index] .+= (colorvec[cols_index] .== color_i) .* vfx1[rows_index] @@ -272,11 +297,7 @@ function finite_difference_jacobian!( _vfx1 = (vfx1 - vfx) / epsilon if ArrayInterface.fast_scalar_indexing(x1) - if ArrayInterface.fast_column_indexing(sparsity) - @_colorediteration(:columniteration,vfx1) - else - @_colorediteration(:indexing,vfx1) - end + _colorediteration!(J,sparsity,rows_index,cols_index,vfx1,colorvec,color_i,n) else #= J.nzval[rows_index] .+= (colorvec[cols_index] .== color_i) .* vfx1[rows_index] @@ -345,11 +366,7 @@ function finite_difference_jacobian!( @. vfx1 = (vfx1 - vfx) / 2epsilon if ArrayInterface.fast_scalar_indexing(x1) - if ArrayInterface.fast_column_indexing(sparsity) - @_colorediteration(:columniteration,vfx1) - else - @_colorediteration(:indexing,vfx1) - end + _colorediteration!(J,sparsity,rows_index,cols_index,vfx1,colorvec,color_i,n) else #= J.nzval[rows_index] .+= (colorvec[cols_index] .== color_i) .* vfx1[rows_index] @@ -381,11 +398,7 @@ function finite_difference_jacobian!( # vfx1 is the compressed Jacobian column if ArrayInterface.fast_scalar_indexing(x1) - if ArrayInterface.fast_column_indexing(sparsity) - @_colorediteration(:columniteration,vfx1) - else - @_colorediteration(:indexing,vfx1) - end + _colorediteration!(J,sparsity,rows_index,cols_index,vfx1,colorvec,color_i,n) else #= J.nzval[rows_index] .+= (colorvec[cols_index] .== color_i) .* vfx1[rows_index] @@ -443,11 +456,7 @@ function finite_difference_jacobian!( @. vfx = imag(vfx) / epsilon if ArrayInterface.fast_scalar_indexing(x1) - if ArrayInterface.fast_column_indexing(sparsity) - @_colorediteration(:columniteration,vfx) - else - @_colorediteration(:indexing,vfx) - end + _colorediteration!(J,sparsity,rows_index,cols_index,vfx,colorvec,color_i,n) else #= J.nzval[rows_index] .+= (colorvec[cols_index] .== color_i) .* vfx[rows_index] @@ -475,11 +484,7 @@ function finite_difference_jacobian!( vfx = imag(vfx) / epsilon if ArrayInterface.fast_scalar_indexing(x1) - if ArrayInterface.fast_column_indexing(sparsity) - @_colorediteration(:columniteration,vfx1) - else - @_colorediteration(:indexing,vfx1) - end + _colorediteration!(J,sparsity,rows_index,cols_index,vfx1,colorvec,color_i,n) else #= J.nzval[rows_index] .+= (colorvec[cols_index] .== color_i) .* vfx1[rows_index] diff --git a/test/coloring_tests.jl b/test/coloring_tests.jl index dd9603e..c10298c 100644 --- a/test/coloring_tests.jl +++ b/test/coloring_tests.jl @@ -1,4 +1,4 @@ -using DiffEqDiffTools, LinearAlgebra, SparseArrays, Test, LinearAlgebra +using DiffEqDiffTools, LinearAlgebra, SparseArrays, Test, LinearAlgebra, BlockBandedMatrices, ArrayInterface fcalls = 0 function f(dx,x) @@ -80,4 +80,23 @@ _J2 = similar(_J2) fcalls = 0 DiffEqDiffTools.finite_difference_jacobian!(_J2,f,rand(30),Val{:complex},colorvec=repeat(1:3,10)) @test fcalls == 3 -@test _J2 ≈ _J \ No newline at end of file +@test _J2 ≈ _J + +#https://github.com/JuliaDiffEq/DiffEqDiffTools.jl/issues/67#issuecomment-516871956 +function f(out, x) + x = reshape(x, 100, 100) + out = reshape(out, 100, 100) + for i in 1:100 + for j in 1:100 + out[i, j] = x[i, j] + x[max(i -1, 1), j] + x[min(i+1, size(x, 1)), j] + x[i, max(j-1, 1)] + x[i, min(j+1, size(x, 2))] + end + end + return vec(out) +end +x = rand(10000) +J = BandedBlockBandedMatrix(Ones(10000, 10000), (fill(100, 100), fill(100, 100)), (1, 1), (1, 1)) +Jsparse = sparse(J) +colors = ArrayInterface.matrix_colors(J) +DiffEqDiffTools.finite_difference_jacobian!(J, f, x, colorvec=colors) +DiffEqDiffTools.finite_difference_jacobian!(Jsparse, f, x, colorvec=colors) +@test J ≈ Jsparse \ No newline at end of file From ddcd86330b146e168104477d79896ea574078557 Mon Sep 17 00:00:00 2001 From: Langwen Huang Date: Thu, 8 Aug 2019 12:43:06 +0800 Subject: [PATCH 6/9] move BBB dep to require block --- Project.toml | 5 ++-- src/DiffEqDiffTools.jl | 2 +- src/jacobians.jl | 57 +++++++++++++++++++++++------------------- 3 files changed, 35 insertions(+), 29 deletions(-) diff --git a/Project.toml b/Project.toml index 41d618d..d7edd65 100644 --- a/Project.toml +++ b/Project.toml @@ -7,7 +7,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9" -BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0" +Requires = "ae029012-a4dd-5104-9daa-d747884805df" [compat] julia = "1" @@ -15,6 +15,7 @@ ArrayInterface = "1.1" [extras] Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" +BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0" [targets] -test = ["Test"] +test = ["Test","BlockBandedMatrices"] diff --git a/src/DiffEqDiffTools.jl b/src/DiffEqDiffTools.jl index 7d1c31b..0485fe9 100644 --- a/src/DiffEqDiffTools.jl +++ b/src/DiffEqDiffTools.jl @@ -2,7 +2,7 @@ __precompile__() module DiffEqDiffTools -using LinearAlgebra, SparseArrays, StaticArrays, ArrayInterface, BlockBandedMatrices +using LinearAlgebra, SparseArrays, StaticArrays, ArrayInterface, Requires import Base: resize! diff --git a/src/jacobians.jl b/src/jacobians.jl index 9b9f7bb..0ab462c 100644 --- a/src/jacobians.jl +++ b/src/jacobians.jl @@ -142,27 +142,37 @@ end end end -@inline function _colorediteration!(Jac::BlockBandedMatrices.BandedBlockBandedMatrix, - sparsity::BlockBandedMatrices.BandedBlockBandedMatrix, - rows_index,cols_index,vfx,colorvec,color_i,ncols) - λ,μ = subblockbandwidths(Jac) - rs = BlockBandedMatrices.BlockSizes((BlockBandedMatrices.cumulsizes(Jac,1),)) # column block sizes - cs = BlockBandedMatrices.BlockSizes((BlockBandedMatrices.cumulsizes(Jac,2),)) - b = BlockBandedMatrices.BlockArray(vfx,rs) - c = BlockBandedMatrices.BlockArray(colorvec,cs) - @inbounds for J=Block.(1:BlockBandedMatrices.nblocks(Jac,2)) - c_v = c.blocks[J.n[1]] - @inbounds for K=BlockBandedMatrices.blockcolrange(Jac,J) - V = view(Jac,K,J) - b_v = b.blocks[K.n[1]] - data = BlockBandedMatrices.bandeddata(V) - p = pointer(data) - st = stride(data,2) - m,n = size(V) - @inbounds for j=1:n - if c_v[j] == color_i - @inbounds for k=max(1,j-μ):min(m,j+λ) - unsafe_store!(p, b_v[k], (j-1)*st + μ + k - j + 1) +#override default setting of using findstructralnz +_use_findstructralnz(sparsity) = ArrayInterface.has_sparsestruct(sparsity) +_use_findstructralnz(::SparseMatrixCSC) = false + +function __init__() + @require BlockBandedMatrices="ffab5731-97b5-5995-9138-79e8c1846df0" begin + _use_findstructralnz(::BlockBandedMatrices.BandedBlockBandedMatrix) = false + + @inline function _colorediteration!(Jac::BlockBandedMatrices.BandedBlockBandedMatrix, + sparsity::BlockBandedMatrices.BandedBlockBandedMatrix, + rows_index,cols_index,vfx,colorvec,color_i,ncols) + λ,μ = BlockBandedMatrices.subblockbandwidths(Jac) + rs = BlockBandedMatrices.BlockSizes((BlockBandedMatrices.cumulsizes(Jac,1),)) # column block sizes + cs = BlockBandedMatrices.BlockSizes((BlockBandedMatrices.cumulsizes(Jac,2),)) + b = BlockBandedMatrices.BlockArray(vfx,rs) + c = BlockBandedMatrices.BlockArray(colorvec,cs) + @inbounds for J=BlockBandedMatrices.Block.(1:BlockBandedMatrices.nblocks(Jac,2)) + c_v = c.blocks[J.n[1]] + @inbounds for K=BlockBandedMatrices.blockcolrange(Jac,J) + V = view(Jac,K,J) + b_v = b.blocks[K.n[1]] + data = BlockBandedMatrices.bandeddata(V) + p = pointer(data) + st = stride(data,2) + m,n = size(V) + @inbounds for j=1:n + if c_v[j] == color_i + @inbounds for k=max(1,j-μ):min(m,j+λ) + unsafe_store!(p, b_v[k], (j-1)*st + μ + k - j + 1) + end + end end end end @@ -186,11 +196,6 @@ function finite_difference_jacobian( _J isa SMatrix ? SArray(J) : J end -#override default setting of using findstructralnz -_use_findstructralnz(sparsity) = ArrayInterface.has_sparsestruct(sparsity) -_use_findstructralnz(::SparseMatrixCSC) = false -_use_findstructralnz(::BlockBandedMatrices.BandedBlockBandedMatrix) = false - function finite_difference_jacobian!( J::AbstractMatrix{<:Number}, f, From b07dee9eeefc3cdb20c3289a09024ac19cee3b75 Mon Sep 17 00:00:00 2001 From: Langwen Huang Date: Thu, 8 Aug 2019 17:59:07 +0800 Subject: [PATCH 7/9] add optimization for BlockBanded and BandedMatrix --- Project.toml | 3 ++- src/jacobians.jl | 50 ++++++++++++++++++++++++++++++++++++++++++ test/coloring_tests.jl | 22 +++++++++++++------ 3 files changed, 67 insertions(+), 8 deletions(-) diff --git a/Project.toml b/Project.toml index d7edd65..f51009f 100644 --- a/Project.toml +++ b/Project.toml @@ -16,6 +16,7 @@ ArrayInterface = "1.1" [extras] Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0" +BandedMatrices="aae01518-5342-5314-be14-df237901396f" [targets] -test = ["Test","BlockBandedMatrices"] +test = ["Test","BlockBandedMatrices","BandedMatrices"] diff --git a/src/jacobians.jl b/src/jacobians.jl index 0ab462c..a7bb988 100644 --- a/src/jacobians.jl +++ b/src/jacobians.jl @@ -149,6 +149,7 @@ _use_findstructralnz(::SparseMatrixCSC) = false function __init__() @require BlockBandedMatrices="ffab5731-97b5-5995-9138-79e8c1846df0" begin _use_findstructralnz(::BlockBandedMatrices.BandedBlockBandedMatrix) = false + _use_findstructralnz(::BlockBandedMatrices.BlockBandedMatrix) = false @inline function _colorediteration!(Jac::BlockBandedMatrices.BandedBlockBandedMatrix, sparsity::BlockBandedMatrices.BandedBlockBandedMatrix, @@ -177,7 +178,56 @@ function __init__() end end end + + @inline function _colorediteration!(Jac::BlockBandedMatrices.BlockBandedMatrix, + sparsity::BlockBandedMatrices.BlockBandedMatrix, + rows_index,cols_index,vfx,colorvec,color_i,ncols) + rs = BlockBandedMatrices.BlockSizes((BlockBandedMatrices.cumulsizes(Jac,1),)) # column block sizes + cs = BlockBandedMatrices.BlockSizes((BlockBandedMatrices.cumulsizes(Jac,2),)) + b = BlockBandedMatrices.BlockArray(vfx,rs) + c = BlockBandedMatrices.BlockArray(colorvec,cs) + @inbounds for J=BlockBandedMatrices.Block.(1:BlockBandedMatrices.nblocks(Jac,2)) + c_v = c.blocks[J.n[1]] + blockcolrange = BlockBandedMatrices.blockcolrange(Jac,J) + _,n = BlockBandedMatrices.blocksize(Jac,(blockcolrange[1].n[1],J.n[1])) + @inbounds for j = 1:n + if c_v[j] == color_i + @inbounds for K = blockcolrange + V = view(Jac,K,J) + b_v = b.blocks[K.n[1]] + m = size(V,1) + @inbounds for k = 1:m + V[k,j] = b_v[k] + end + end + end + end + end + end + + end + + @require BandedMatrices = "aae01518-5342-5314-be14-df237901396f" begin + + _use_findstructralnz(::BandedMatrices.BandedMatrix) = false + + @inline function _colorediteration!(Jac::BandedMatrices.BandedMatrix, + sparsity::BandedMatrices.BandedMatrix, + rows_index,cols_index,vfx,colorvec,color_i,ncols) + nrows = size(Jac,1) + l,u = BandedMatrices.bandwidths(Jac) + #data = BandedMatrices.bandeddata(Jac) + @inbounds for col_index in max(1,1-l):min(ncols,ncols+u) + if colorvec[col_index] == color_i + @inbounds for row_index in max(1,col_index-u):min(nrows,col_index+l) + #data[u+row_index-col_index+1,col_index] = vfx[row_index] + Jac[row_index,col_index]=vfx[row_index] + end + end + end + end end + end function finite_difference_jacobian( diff --git a/test/coloring_tests.jl b/test/coloring_tests.jl index c10298c..83fbf80 100644 --- a/test/coloring_tests.jl +++ b/test/coloring_tests.jl @@ -1,4 +1,4 @@ -using DiffEqDiffTools, LinearAlgebra, SparseArrays, Test, LinearAlgebra, BlockBandedMatrices, ArrayInterface +using DiffEqDiffTools, LinearAlgebra, SparseArrays, Test, LinearAlgebra, BlockBandedMatrices, ArrayInterface, BandedMatrices fcalls = 0 function f(dx,x) @@ -94,9 +94,17 @@ function f(out, x) return vec(out) end x = rand(10000) -J = BandedBlockBandedMatrix(Ones(10000, 10000), (fill(100, 100), fill(100, 100)), (1, 1), (1, 1)) -Jsparse = sparse(J) -colors = ArrayInterface.matrix_colors(J) -DiffEqDiffTools.finite_difference_jacobian!(J, f, x, colorvec=colors) -DiffEqDiffTools.finite_difference_jacobian!(Jsparse, f, x, colorvec=colors) -@test J ≈ Jsparse \ No newline at end of file +Jbbb = BandedBlockBandedMatrix(Ones(10000, 10000), (fill(100, 100), fill(100, 100)), (1, 1), (1, 1)) +Jsparse = sparse(Jbbb) +colorsbbb = ArrayInterface.matrix_colors(Jbbb) +DiffEqDiffTools.finite_difference_jacobian!(Jbbb, f, x, colorvec=colorsbbb) +DiffEqDiffTools.finite_difference_jacobian!(Jsparse, f, x, colorvec=colorsbbb) +@test Jbbb ≈ Jsparse +Jbb = BlockBandedMatrix(similar(Jsparse),(fill(100, 100), fill(100, 100)),(1,1)); +colorsbb = ArrayInterface.matrix_colors(Jbb) +DiffEqDiffTools.finite_difference_jacobian!(Jbb, f, x, colorvec=colorsbb) +@test Jbb ≈ Jsparse +Jb = BandedMatrices.BandedMatrix(similar(Jsparse),(1,1)) +colorsb = ArrayInterface.matrix_colors(Jb) +DiffEqDiffTools.finite_difference_jacobian!(Jb, f, x, colorvec=colorsb) +@test Jb ≈ Jsparse \ No newline at end of file From c3bcbd770875e2b7b8016a3457bde9bab6ca9cb9 Mon Sep 17 00:00:00 2001 From: Langwen Huang Date: Thu, 8 Aug 2019 18:18:34 +0800 Subject: [PATCH 8/9] fix tests --- test/coloring_tests.jl | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/test/coloring_tests.jl b/test/coloring_tests.jl index 83fbf80..c7b162c 100644 --- a/test/coloring_tests.jl +++ b/test/coloring_tests.jl @@ -82,6 +82,14 @@ DiffEqDiffTools.finite_difference_jacobian!(_J2,f,rand(30),Val{:complex},colorve @test fcalls == 3 @test _J2 ≈ _J +_Jb = BandedMatrices.BandedMatrix(similar(_J2),(1,1)) +DiffEqDiffTools.finite_difference_jacobian!(_Jb, f, rand(30), colorvec=colorvec=repeat(1:3,10)) +@test _Jb ≈ _J + +_Jtri = Tridiagonal(similar(_J2)) +DiffEqDiffTools.finite_difference_jacobian!(_Jtri, f, rand(30), colorvec=colorvec=repeat(1:3,10)) +@test _Jtri ≈ _J + #https://github.com/JuliaDiffEq/DiffEqDiffTools.jl/issues/67#issuecomment-516871956 function f(out, x) x = reshape(x, 100, 100) @@ -103,8 +111,4 @@ DiffEqDiffTools.finite_difference_jacobian!(Jsparse, f, x, colorvec=colorsbbb) Jbb = BlockBandedMatrix(similar(Jsparse),(fill(100, 100), fill(100, 100)),(1,1)); colorsbb = ArrayInterface.matrix_colors(Jbb) DiffEqDiffTools.finite_difference_jacobian!(Jbb, f, x, colorvec=colorsbb) -@test Jbb ≈ Jsparse -Jb = BandedMatrices.BandedMatrix(similar(Jsparse),(1,1)) -colorsb = ArrayInterface.matrix_colors(Jb) -DiffEqDiffTools.finite_difference_jacobian!(Jb, f, x, colorvec=colorsb) -@test Jb ≈ Jsparse \ No newline at end of file +@test Jbb ≈ Jsparse \ No newline at end of file From 8321080451e67a7ac1a2128b8a4666be9233c908 Mon Sep 17 00:00:00 2001 From: Langwen Huang Date: Thu, 8 Aug 2019 18:24:09 +0800 Subject: [PATCH 9/9] update dependency --- Project.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/Project.toml b/Project.toml index f51009f..1ebd64d 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "DiffEqDiffTools" uuid = "01453d9d-ee7c-5054-8395-0335cb756afa" -version = "1.2.0" +version = "1.3.0" [deps] LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" @@ -11,7 +11,7 @@ Requires = "ae029012-a4dd-5104-9daa-d747884805df" [compat] julia = "1" -ArrayInterface = "1.1" +ArrayInterface = ">= 1.1" [extras] Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"