Skip to content

Commit de4a21a

Browse files
gdallegdalle
committed
Add BlockBandedMatrices
1 parent e088877 commit de4a21a

8 files changed

Lines changed: 172 additions & 92 deletions

‎Project.toml‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,16 +15,18 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
1515

1616
[weakdeps]
1717
BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
18+
BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
19+
BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0"
1820

1921
[extensions]
2022
SparseMatrixColoringsBandedMatricesExt = "BandedMatrices"
21-
22-
[extras]
23-
BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
23+
SparseMatrixColoringsBlockBandedMatricesExt = ["BlockArrays", "BlockBandedMatrices"]
2424

2525
[compat]
2626
ADTypes = "1.2.1"
2727
BandedMatrices = "1.7.5"
28+
BlockArrays = "1.1.1"
29+
BlockBandedMatrices = "0.13.1"
2830
Compat = "3.46,4.2"
2931
DataStructures = "0.18"
3032
DocStringExtensions = "0.8,0.9"
@@ -33,3 +35,8 @@ Random = "<0.0.1, 1"
3335
Requires = "1.3.0"
3436
SparseArrays = "<0.0.1, 1"
3537
julia = "1.6"
38+
39+
[extras]
40+
BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
41+
BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
42+
BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0"

‎ext/SparseMatrixColoringsBandedMatricesExt.jl‎

Lines changed: 7 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
module SparseMatrixColoringsBandedMatricesExt
22

33
if isdefined(Base, :get_extension)
4-
using BandedMatrices: BandedMatrix, bandwidths, bandrange
4+
using BandedMatrices: BandedMatrix, bandrange, bandwidths, colrange, rowrange
55
using SparseMatrixColorings:
66
BipartiteGraph,
77
ColoringProblem,
@@ -13,7 +13,7 @@ if isdefined(Base, :get_extension)
1313
row_colors
1414
import SparseMatrixColorings as SMC
1515
else
16-
using ..BandedMatrices: BandedMatrix, bandwidths, bandrange
16+
using ..BandedMatrices: BandedMatrix, bandrange, bandwidths, colrange, rowrange
1717
using ..SparseMatrixColorings:
1818
BipartiteGraph,
1919
ColoringProblem,
@@ -27,7 +27,7 @@ else
2727
end
2828

2929
#=
30-
This code is partially taken from ArrayInterface.jl and FiniteDiff.jl
30+
This code is partly taken from ArrayInterface.jl and FiniteDiff.jl
3131
https://github.com/JuliaArrays/ArrayInterface.jl
3232
https://github.com/JuliaDiff/FiniteDiff.jl
3333
=#
@@ -38,8 +38,7 @@ function SMC.coloring(
3838
algo::GreedyColoringAlgorithm;
3939
kwargs...,
4040
)
41-
l, u = bandwidths(A)
42-
width = u + l + 1
41+
width = length(bandrange(A))
4342
color = cycle_range(width, size(A, 2))
4443
bg = BipartiteGraph(A)
4544
return ColumnColoringResult(A, bg, color)
@@ -51,20 +50,17 @@ function SMC.coloring(
5150
algo::GreedyColoringAlgorithm;
5251
kwargs...,
5352
)
54-
l, u = bandwidths(A)
55-
width = u + l + 1
53+
width = length(bandrange(A))
5654
color = cycle_range(width, size(A, 1))
5755
bg = BipartiteGraph(A)
5856
return RowColoringResult(A, bg, color)
5957
end
6058

6159
function SMC.decompress!(A::BandedMatrix, B::AbstractMatrix, result::ColumnColoringResult)
6260
color = column_colors(result)
63-
m, n = size(A)
64-
l, u = bandwidths(A)
6561
for j in axes(A, 2)
6662
c = color[j]
67-
for i in max(1, j - u):min(m, j + l)
63+
for i in colrange(A, j)
6864
A[i, j] = B[i, c]
6965
end
7066
end
@@ -73,11 +69,9 @@ end
7369

7470
function SMC.decompress!(A::BandedMatrix, B::AbstractMatrix, result::RowColoringResult)
7571
color = row_colors(result)
76-
m, n = size(A)
77-
l, u = bandwidths(A)
7872
for i in axes(A, 1)
7973
c = color[i]
80-
for j in max(1, i - l):min(n, i + u)
74+
for j in rowrange(A, i)
8175
A[i, j] = B[c, j]
8276
end
8377
end

‎ext/SparseMatrixColoringsBlockBandedMatricesExt.jl‎

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
module SparseMatrixColoringsBlockBandedMatricesExt
2+
3+
if isdefined(Base, :get_extension)
4+
using BlockArrays: blockaxes, blockfirsts, blocklasts, blocksize, blocklengths
5+
using BlockBandedMatrices:
6+
BlockBandedMatrix,
7+
blockbandrange,
8+
blockbandwidths,
9+
blocklengths,
10+
blocksize,
11+
subblockbandwidths
12+
using SparseMatrixColorings:
13+
BipartiteGraph,
14+
ColoringProblem,
15+
ColumnColoringResult,
16+
GreedyColoringAlgorithm,
17+
RowColoringResult,
18+
column_colors,
19+
cycle_range,
20+
row_colors
21+
import SparseMatrixColorings as SMC
22+
else
23+
using ..BlockArrays: blockaxes, blockfirsts, blocklasts, blocksize, blocklengths
24+
using ..BlockBandedMatrices:
25+
BlockBandedMatrix,
26+
blockbandrange,
27+
blockbandwidths,
28+
blocklengths,
29+
blocksize,
30+
subblockbandwidths
31+
using ..SparseMatrixColorings:
32+
BipartiteGraph,
33+
ColoringProblem,
34+
ColumnColoringResult,
35+
GreedyColoringAlgorithm,
36+
RowColoringResult,
37+
column_colors,
38+
cycle_range,
39+
row_colors
40+
import ..SparseMatrixColorings as SMC
41+
end
42+
43+
#=
44+
This code is partly taken from ArrayInterface.jl and FiniteDiff.jl
45+
https://github.com/JuliaArrays/ArrayInterface.jl
46+
https://github.com/JuliaDiff/FiniteDiff.jl
47+
=#
48+
49+
function SMC.coloring(
50+
A::BlockBandedMatrix,
51+
::ColoringProblem{:nonsymmetric,:column},
52+
algo::GreedyColoringAlgorithm;
53+
kwargs...,
54+
)
55+
# consider blocks of columns
56+
nb_blocks = blocksize(A, 2)
57+
nb_cols_in_block = blocklengths(axes(A, 2))
58+
first_col_in_block = blockfirsts(axes(A, 2))
59+
last_col_in_block = blocklasts(axes(A, 2))
60+
61+
# give a macroscopic color to each block, so that 2 blocks of columns with the same macro color do not intersect
62+
# same idea as for BandedMatrices
63+
nb_macrocolors = length(blockbandrange(A))
64+
macrocolor = cycle_range(nb_macrocolors, nb_blocks)
65+
66+
# for each macroscopic color, count how many microscopic colors will be needed
67+
# columns within a block are colored naively with all distinct micro colors
68+
nb_colors_in_macrocolor = [
69+
maximum(nb_cols_in_block[mc:nb_macrocolors:nb_blocks]; init=0) for
70+
mc in 1:nb_macrocolors
71+
]
72+
color_shift_in_macrocolor = vcat(0, cumsum(nb_colors_in_macrocolor)[1:(end - 1)])
73+
74+
# assign a microscopic color to each column as a function of its macroscopic color and its position within the block
75+
color = Vector{Int}(undef, size(A, 2))
76+
for b in 1:nb_blocks
77+
mc = macrocolor[b]
78+
shift = color_shift_in_macrocolor[mc]
79+
for j in first_col_in_block[b]:last_col_in_block[b]
80+
color[j] = shift + (j - first_col_in_block[b] + 1)
81+
end
82+
end
83+
84+
bg = BipartiteGraph(A)
85+
return ColumnColoringResult(A, bg, color)
86+
end
87+
88+
end

‎src/check.jl‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ function structurally_orthogonal_columns(
3636
group = group_by_color(color)
3737
for (c, g) in enumerate(group)
3838
Ag = view(A, :, g)
39-
nonzeros_per_row = dropdims(count(!iszero, Ag; dims=2); dims=2)
39+
nonzeros_per_row = only(eachcol(count(!iszero, Ag; dims=2)))
4040
max_nonzeros_per_row, i = findmax(nonzeros_per_row)
4141
if max_nonzeros_per_row > 1
4242
if verbose

‎src/structured.jl‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
#=
2-
This code is partially taken from ArrayInterface.jl
2+
This code is partly taken from ArrayInterface.jl
33
https://github.com/JuliaArrays/ArrayInterface.jl
44
=#
55

‎test/Project.toml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,9 @@
11
[deps]
22
ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
33
Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
4+
ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
45
BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
6+
BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
57
BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0"
68
CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
79
Chairmarks = "0ca39b1e-fe0b-4e98-acfc-b1656634c4de"

‎test/structured.jl‎

Lines changed: 26 additions & 73 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,11 @@
1+
using ArrayInterface: ArrayInterface
12
using BandedMatrices: BandedMatrix, brand
3+
using BlockBandedMatrices: BlockBandedMatrix
24
using LinearAlgebra
35
using SparseMatrixColorings
46
using SparseMatrixColorings: cycle_range
57
using Test
68

7-
column_problem = ColoringProblem(; structure=:nonsymmetric, partition=:column)
8-
row_problem = ColoringProblem(; structure=:nonsymmetric, partition=:row)
9-
10-
algo = GreedyColoringAlgorithm()
11-
129
@testset "Utils" begin
1310
@test cycle_range(2, 3) == [1, 2, 1]
1411
@test cycle_range(2, 4) == [1, 2, 1, 2]
@@ -18,87 +15,43 @@ algo = GreedyColoringAlgorithm()
1815
@test cycle_range(2, 1) == [1]
1916
@test cycle_range(3, 1) == [1]
2017
@test cycle_range(3, 2) == [1, 2]
21-
end
18+
end;
2219

2320
@testset "Diagonal" begin
24-
for n in (1, 2, 10, 100, 1000)
21+
@testset for n in (1, 2, 10, 100)
2522
A = Diagonal(rand(n))
26-
# column
27-
result = coloring(A, column_problem, algo)
28-
B = compress(A, result)
29-
D = decompress(B, result)
30-
@test size(B, 2) == 1
31-
@test D == A
32-
@test D isa Diagonal
33-
# row
34-
result = coloring(A, row_problem, algo)
35-
B = compress(A, result)
36-
D = decompress(B, result)
37-
@test size(B, 1) == 1
38-
@test D == A
39-
@test D isa Diagonal
23+
test_structured_coloring_decompression(A)
4024
end
41-
end
25+
end;
4226

4327
@testset "Bidiagonal" begin
44-
for n in (2, 10, 100, 1000)
28+
@testset for n in (2, 10, 100)
4529
A1 = Bidiagonal(rand(n), rand(n - 1), :U)
4630
A2 = Bidiagonal(rand(n), rand(n - 1), :L)
47-
for A in (A1, A2)
48-
# column
49-
result = coloring(A, column_problem, algo)
50-
B = compress(A, result)
51-
D = decompress(B, result)
52-
@test size(B, 2) == 2
53-
@test D == A
54-
@test D isa Bidiagonal
55-
# row
56-
result = coloring(A, row_problem, algo)
57-
B = compress(A, result)
58-
D = decompress(B, result)
59-
@test size(B, 1) == 2
60-
@test D == A
61-
@test D isa Bidiagonal
62-
end
31+
test_structured_coloring_decompression(A1)
32+
test_structured_coloring_decompression(A2)
6333
end
64-
end
34+
end;
6535

6636
@testset "Tridiagonal" begin
67-
for n in (2, 10, 100, 1000)
68-
A1 = Tridiagonal(rand(n - 1), rand(n), rand(n - 1))
69-
A2 = Tridiagonal(rand(n - 1), rand(n), rand(n - 1))
70-
for A in (A1, A2)
71-
# column
72-
result = coloring(A, column_problem, algo)
73-
B = compress(A, result)
74-
D = decompress(B, result)
75-
@test size(B, 2) == min(n, 3)
76-
@test D == A
77-
@test D isa Tridiagonal # row
78-
result = coloring(A, row_problem, algo)
79-
B = compress(A, result)
80-
D = decompress(B, result)
81-
@test size(B, 1) == min(n, 3)
82-
@test D == A
83-
@test D isa Tridiagonal
84-
end
37+
for n in (2, 10, 100)
38+
A = Tridiagonal(rand(n - 1), rand(n), rand(n - 1))
39+
test_structured_coloring_decompression(A)
8540
end
86-
end
41+
end;
8742

8843
@testset "BandedMatrices" begin
89-
for (m, n) in [(10, 20), (20, 10)], l in 0:5, u in 0:5
44+
@testset for (m, n) in [(10, 20), (20, 10)], l in 0:5, u in 0:5
9045
A = brand(m, n, l, u)
91-
# column
92-
result = coloring(A, column_problem, algo)
93-
B = compress(A, result)
94-
D = decompress(B, result)
95-
@test D == A
96-
@test D isa BandedMatrix
97-
# row
98-
result = coloring(A, row_problem, algo)
99-
B = compress(A, result)
100-
D = decompress(B, result)
101-
@test D == A
102-
@test D isa BandedMatrix
46+
test_structured_coloring_decompression(A)
47+
end
48+
end;
49+
50+
@testset "BlockBandedMatrices" begin
51+
@testset for (mb, nb) in [(10, 20), (20, 10)], lb in 0:3, ub in 0:3, bs in 1:3
52+
rows = rand(1:bs, mb)
53+
cols = rand(1:bs, nb)
54+
A = BlockBandedMatrix{Float64}(rand(sum(rows), sum(cols)), rows, cols, (l, u))
55+
test_structured_coloring_decompression(A)
10356
end
104-
end
57+
end;

‎test/utils.jl‎

Lines changed: 37 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,14 @@
1+
using ArrayInterface: ArrayInterface
2+
using BandedMatrices: BandedMatrix
3+
using BlockBandedMatrices: BlockBandedMatrix
14
using LinearAlgebra
25
using SparseMatrixColorings
36
using SparseMatrixColorings:
4-
AdjacencyGraph, LinearSystemColoringResult, matrix_versions, respectful_similar
7+
AdjacencyGraph,
8+
LinearSystemColoringResult,
9+
matrix_versions,
10+
respectful_similar,
11+
structurally_orthogonal_columns
512
using Test
613

714
function test_coloring_decompression(
@@ -112,3 +119,32 @@ function test_coloring_decompression(
112119
@test all(color_vec .== Ref(color_vec[1]))
113120
end
114121
end
122+
123+
OptimalColoringKnown = Union{Diagonal,Bidiagonal,Tridiagonal,BandedMatrix,BlockBandedMatrix}
124+
125+
function test_structured_coloring_decompression(A::AbstractMatrix)
126+
column_problem = ColoringProblem(; structure=:nonsymmetric, partition=:column)
127+
row_problem = ColoringProblem(; structure=:nonsymmetric, partition=:row)
128+
algo = GreedyColoringAlgorithm()
129+
130+
@testset "Column" begin
131+
result = coloring(A, column_problem, algo)
132+
color = column_colors(result)
133+
B = compress(A, result)
134+
D = decompress(B, result)
135+
@test D == A
136+
@test D isa typeof(A)
137+
@test structurally_orthogonal_columns(A, color)
138+
if A isa OptimalColoringKnown
139+
@test color == ArrayInterface.matrix_colors(A)
140+
end
141+
end
142+
143+
@testset "Row" begin
144+
result = coloring(A, row_problem, algo)
145+
B = compress(A, result)
146+
D = decompress(B, result)
147+
@test D == A
148+
@test D isa typeof(A)
149+
end
150+
end

0 commit comments

Comments
 (0)