diff --git a/src/iteration_utils.jl b/src/iteration_utils.jl index cc06c05..cdcbda7 100644 --- a/src/iteration_utils.jl +++ b/src/iteration_utils.jl @@ -22,84 +22,85 @@ _use_findstructralnz(::SparseMatrixCSC) = false function __init__() @require BlockBandedMatrices="ffab5731-97b5-5995-9138-79e8c1846df0" begin - _use_findstructralnz(::BlockBandedMatrices.BandedBlockBandedMatrix) = false - _use_findstructralnz(::BlockBandedMatrices.BlockBandedMatrix) = false + @require BlockArrays="8e7c35d0-a365-5155-bbbb-fb81a777f24e" begin + _use_findstructralnz(::BlockBandedMatrices.BandedBlockBandedMatrix) = false + _use_findstructralnz(::BlockBandedMatrices.BlockBandedMatrix) = 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) + @inline function _colorediteration!(Jac::BlockBandedMatrices.BandedBlockBandedMatrix, + sparsity::BlockBandedMatrices.BandedBlockBandedMatrix, + rows_index,cols_index,vfx,colorvec,color_i,ncols) + λ,μ = BlockBandedMatrices.subblockbandwidths(Jac) + rs = BlockArrays.blocklengths(BlockArrays.axes(Jac,1)) # column block sizes + cs = BlockArrays.blocklengths(BlockArrays.axes(Jac,1)) + b = BlockBandedMatrices.BlockArray(vfx,rs) + c = BlockBandedMatrices.BlockArray(colorvec,cs) + @inbounds for J=BlockArrays.blockaxes(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 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] + @inline function _colorediteration!(Jac::BlockBandedMatrices.BlockBandedMatrix, + sparsity::BlockBandedMatrices.BlockBandedMatrix, + rows_index,cols_index,vfx,colorvec,color_i,ncols) + rs = BlockArrays.blocklengths(BlockArrays.axes(Jac,1)) # column block sizes + cs = BlockArrays.blocklengths(BlockArrays.axes(Jac,1)) + b = BlockBandedMatrices.BlockArray(vfx,rs) + c = BlockBandedMatrices.BlockArray(colorvec,cs) + @inbounds for J=BlockArrays.blockaxes(Jac,2) + c_v = c.blocks[J.n[1]] + blockcolrange = BlockBandedMatrices.blockcolrange(Jac,J) + _,n = length.(getindex.(axes(Jac), (blockcolrange[1], J))) + @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 - end + end - @require BandedMatrices = "aae01518-5342-5314-be14-df237901396f" begin + @require BandedMatrices = "aae01518-5342-5314-be14-df237901396f" begin - _use_findstructralnz(::BandedMatrices.BandedMatrix) = false + _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] + @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 - end diff --git a/test/coloring_tests.jl b/test/coloring_tests.jl index 0d3a5b2..230ff93 100644 --- a/test/coloring_tests.jl +++ b/test/coloring_tests.jl @@ -106,13 +106,13 @@ function f(out, x) return vec(out) end x = rand(10000) -Jbbb = BandedBlockBandedMatrix(Ones(10000, 10000), (fill(100, 100), fill(100, 100)), (1, 1), (1, 1)) +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)); +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