Add BandedMatrices

Author: gdalle

.github/workflows/Test.yml

@@ -20,7 +20,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        julia-version: ['lts', '1', 'pre']
+        julia-version: ['1', 'pre']  # TODO: do we drop 'lts'?
 
     steps:
       - uses: actions/checkout@v4

Project.toml

@@ -10,14 +10,23 @@ DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
+[weakdeps]
+BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
+
+[extensions]
+SparseMatrixColoringsBandedMatricesExt = "BandedMatrices"
+
 [compat]
 ADTypes = "1.2.1"
+BandedMatrices = "1.7.5"
 Compat = "3.46,4.2"
-DocStringExtensions = "0.8,0.9"
 DataStructures = "0.18"
+DocStringExtensions = "0.8,0.9"
 LinearAlgebra = "<0.0.1, 1"
 Random = "<0.0.1, 1"
+Requires = "1.3.0"
 SparseArrays = "<0.0.1, 1"
 julia = "1.6"

ext/SparseMatrixColoringsBandedMatricesExt.jl

@@ -0,0 +1,87 @@
+module SparseMatrixColoringsBandedMatricesExt
+
+if isdefined(Base, :get_extension)
+    using BandedMatrices: BandedMatrix, bandwidths, bandrange
+    using SparseMatrixColorings:
+        BipartiteGraph,
+        ColoringProblem,
+        ColumnColoringResult,
+        GreedyColoringAlgorithm,
+        RowColoringResult,
+        column_colors,
+        cycle_range,
+        row_colors
+    import SparseMatrixColorings as SMC
+else
+    using ..BandedMatrices: BandedMatrix, bandwidths, bandrange
+    using ..SparseMatrixColorings:
+        BipartiteGraph,
+        ColoringProblem,
+        ColumnColoringResult,
+        GreedyColoringAlgorithm,
+        RowColoringResult,
+        column_colors,
+        cycle_range,
+        row_colors
+    import ..SparseMatrixColorings as SMC
+end
+
+#=
+This code is partially taken from ArrayInterface.jl and FiniteDiff.jl
+https://github.com/JuliaArrays/ArrayInterface.jl
+https://github.com/JuliaDiff/FiniteDiff.jl
+=#
+
+function SMC.coloring(
+    A::BandedMatrix,
+    ::ColoringProblem{:nonsymmetric,:column},
+    algo::GreedyColoringAlgorithm;
+    kwargs...,
+)
+    l, u = bandwidths(A)
+    width = u + l + 1
+    color = cycle_range(width, size(A, 2))
+    bg = BipartiteGraph(A)
+    return ColumnColoringResult(A, bg, color)
+end
+
+function SMC.coloring(
+    A::BandedMatrix,
+    ::ColoringProblem{:nonsymmetric,:row},
+    algo::GreedyColoringAlgorithm;
+    kwargs...,
+)
+    l, u = bandwidths(A)
+    width = u + l + 1
+    color = cycle_range(width, size(A, 1))
+    bg = BipartiteGraph(A)
+    return RowColoringResult(A, bg, color)
+end
+
+function SMC.decompress!(A::BandedMatrix, B::AbstractMatrix, result::ColumnColoringResult)
+    color = column_colors(result)
+    m, n = size(A)
+    l, u = bandwidths(A)
+    for j in axes(A, 2)
+        c = color[j]
+        for i in max(1, j - u):min(m, j + l)
+            A[i, j] = B[i, c]
+        end
+    end
+    return A
+end
+
+function SMC.decompress!(A::BandedMatrix, B::AbstractMatrix, result::RowColoringResult)
+    color = row_colors(result)
+    m, n = size(A)
+    l, u = bandwidths(A)
+    for i in axes(A, 1)
+        c = color[i]
+        for j in max(1, i - l):min(n, i + u)
+            A[i, j] = B[c, j]
+        end
+    end
+    return A
+end
+
+end

src/SparseMatrixColorings.jl

@@ -65,4 +65,15 @@ export column_colors, row_colors
 export column_groups, row_groups
 export compress, decompress, decompress!, decompress_single_color!
 
+if !isdefined(Base, :get_extension)
+    using Requires
+end
+
+@static if !isdefined(Base, :get_extension)
+    function __init__()
+        @require BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
+        return include("../ext/SparseMatrixColoringsBandedMatricesExt.jl")
+    end
+end
+
 end

src/structured.jl

@@ -1,8 +1,6 @@
 #=
 This code is partially taken from ArrayInterface.jl
 https://github.com/JuliaArrays/ArrayInterface.jl
-
-Question: when decompressing, should we always assume that the coloring was optimal?
 =#
 
 """
@@ -42,19 +40,15 @@ function coloring(
     return RowColoringResult(A, bg, color)
 end
 
-function decompress!(
-    A::Diagonal{R}, B::AbstractMatrix{R}, result::ColumnColoringResult
-) where {R<:Real}
+function decompress!(A::Diagonal, B::AbstractMatrix, result::ColumnColoringResult)
     color = column_colors(result)
     for j in axes(A, 2)
         A[j, j] = B[j, color[j]]
     end
     return A
 end
 
-function decompress!(
-    A::Diagonal{R}, B::AbstractMatrix{R}, result::RowColoringResult
-) where {R<:Real}
+function decompress!(A::Diagonal, B::AbstractMatrix, result::RowColoringResult)
     color = row_colors(result)
     for i in axes(A, 1)
         A[i, i] = B[color[i], i]
@@ -86,9 +80,7 @@ function coloring(
     return RowColoringResult(A, bg, color)
 end
 
-function decompress!(
-    A::Bidiagonal{R}, B::AbstractMatrix{R}, result::ColumnColoringResult
-) where {R<:Real}
+function decompress!(A::Bidiagonal, B::AbstractMatrix, result::ColumnColoringResult)
     color = column_colors(result)
     for j in axes(A, 2)
         c = color[j]
@@ -102,9 +94,7 @@ function decompress!(
     return A
 end
 
-function decompress!(
-    A::Bidiagonal{R}, B::AbstractMatrix{R}, result::RowColoringResult
-) where {R<:Real}
+function decompress!(A::Bidiagonal, B::AbstractMatrix, result::RowColoringResult)
     color = row_colors(result)
     for i in axes(A, 1)
         c = color[i]
@@ -142,9 +132,7 @@ function coloring(
     return RowColoringResult(A, bg, color)
 end
 
-function decompress!(
-    A::Tridiagonal{R}, B::AbstractMatrix{R}, result::ColumnColoringResult
-) where {R<:Real}
+function decompress!(A::Tridiagonal, B::AbstractMatrix, result::ColumnColoringResult)
     color = column_colors(result)
     for j in axes(A, 2)
         c = color[j]
@@ -159,9 +147,7 @@ function decompress!(
     return A
 end
 
-function decompress!(
-    A::Tridiagonal{R}, B::AbstractMatrix{R}, result::RowColoringResult
-) where {R<:Real}
+function decompress!(A::Tridiagonal, B::AbstractMatrix, result::RowColoringResult)
     color = row_colors(result)
     for i in axes(A, 1)
         c = color[i]

test/Project.toml

@@ -1,6 +1,7 @@
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
+BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
 CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
 Chairmarks = "0ca39b1e-fe0b-4e98-acfc-b1656634c4de"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"

test/structured.jl

@@ -1,3 +1,4 @@
+using BandedMatrices: BandedMatrix, brand
 using LinearAlgebra
 using SparseMatrixColorings
 using SparseMatrixColorings: cycle_range
@@ -25,15 +26,15 @@ end
         # column
         result = coloring(A, column_problem, algo)
         B = compress(A, result)
-        @test size(B, 2) == 1
         D = decompress(B, result)
+        @test size(B, 2) == 1
         @test D == A
         @test D isa Diagonal
         # row
         result = coloring(A, row_problem, algo)
         B = compress(A, result)
-        @test size(B, 1) == 1
         D = decompress(B, result)
+        @test size(B, 1) == 1
         @test D == A
         @test D isa Diagonal
     end
@@ -47,15 +48,15 @@ end
             # column
             result = coloring(A, column_problem, algo)
             B = compress(A, result)
-            @test size(B, 2) == 2
             D = decompress(B, result)
+            @test size(B, 2) == 2
             @test D == A
             @test D isa Bidiagonal
             # row
             result = coloring(A, row_problem, algo)
             B = compress(A, result)
-            @test size(B, 1) == 2
             D = decompress(B, result)
+            @test size(B, 1) == 2
             @test D == A
             @test D isa Bidiagonal
         end
@@ -70,16 +71,34 @@ end
             # column
             result = coloring(A, column_problem, algo)
             B = compress(A, result)
-            @test size(B, 2) == min(n, 3)
             D = decompress(B, result)
+            @test size(B, 2) == min(n, 3)
             @test D == A
             @test D isa Tridiagonal            # row
             result = coloring(A, row_problem, algo)
             B = compress(A, result)
-            @test size(B, 1) == min(n, 3)
             D = decompress(B, result)
+            @test size(B, 1) == min(n, 3)
             @test D == A
             @test D isa Tridiagonal
         end
     end
 end
+
+@testset "BandedMatrices" begin
+    for (m, n) in [(10, 20), (20, 10)], l in 0:5, u in 0:5
+        A = brand(m, n, l, u)
+        # column
+        result = coloring(A, column_problem, algo)
+        B = compress(A, result)
+        D = decompress(B, result)
+        @test D == A
+        @test D isa BandedMatrix
+        # row
+        result = coloring(A, row_problem, algo)
+        B = compress(A, result)
+        D = decompress(B, result)
+        @test D == A
+        @test D isa BandedMatrix
+    end
+end