From e4e6dd9502e3054dd428cf0060415b8bf61da7ce Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Wed, 2 Oct 2024 09:22:11 +0200
Subject: [PATCH 01/33] Structured decompression

---
 src/SparseMatrixColorings.jl |   3 +
 src/result.jl                |  24 +++--
 src/structured.jl            | 168 +++++++++++++++++++++++++++++++++++
 test/allocations.jl          |  38 +++++++-
 test/runtests.jl             |   3 +
 test/structured.jl           |  62 +++++++++++++
 test/type_stability.jl       |  48 +++++++++-
 7 files changed, 332 insertions(+), 14 deletions(-)
 create mode 100644 src/structured.jl
 create mode 100644 test/structured.jl

diff --git a/src/SparseMatrixColorings.jl b/src/SparseMatrixColorings.jl
index 0fe6e6d2..d712ea31 100644
--- a/src/SparseMatrixColorings.jl
+++ b/src/SparseMatrixColorings.jl
@@ -16,11 +16,13 @@ using DataStructures: DisjointSets, find_root!, root_union!, num_groups
 using DocStringExtensions: README, EXPORTS, SIGNATURES, TYPEDEF, TYPEDFIELDS
 using LinearAlgebra:
     Adjoint,
+    Bidiagonal,
     Diagonal,
     Hermitian,
     LowerTriangular,
     Symmetric,
     Transpose,
+    Tridiagonal,
     UpperTriangular,
     adjoint,
     checksquare,
@@ -51,6 +53,7 @@ include("matrices.jl")
 include("interface.jl")
 include("constant.jl")
 include("decompression.jl")
+include("structured.jl")
 include("check.jl")
 include("examples.jl")
 
diff --git a/src/result.jl b/src/result.jl
index 62ff4864..3668016d 100644
--- a/src/result.jl
+++ b/src/result.jl
@@ -24,7 +24,7 @@ Combination between the type parameters of [`ColoringProblem`](@ref) and [`Greed
 !!! warning
     Unlike the methods above, the concrete subtypes of `AbstractColoringResult` are not part of the public API and may change without notice.
 """
-abstract type AbstractColoringResult{structure,partition,decompression,M<:SparseMatrixCSC} end
+abstract type AbstractColoringResult{structure,partition,decompression,M} end
 
 """
     column_colors(result::AbstractColoringResult)
@@ -99,7 +99,8 @@ $TYPEDFIELDS
 
 - [`AbstractColoringResult`](@ref)
 """
-struct ColumnColoringResult{M} <: AbstractColoringResult{:nonsymmetric,:column,:direct,M}
+struct ColumnColoringResult{M,C<:Union{Nothing,Vector{Int}}} <:
+       AbstractColoringResult{:nonsymmetric,:column,:direct,M}
     "matrix that was colored"
     S::M
     "one integer color for each column or row (depending on `partition`)"
@@ -107,7 +108,12 @@ struct ColumnColoringResult{M} <: AbstractColoringResult{:nonsymmetric,:column,:
     "color groups for columns or rows (depending on `partition`)"
     group::Vector{Vector{Int}}
     "flattened indices mapping the compressed matrix `B` to the uncompressed matrix `A` when `A isa SparseMatrixCSC`. They satisfy `nonzeros(A)[k] = vec(B)[compressed_indices[k]]`"
-    compressed_indices::Vector{Int}
+    compressed_indices::C
+end
+
+function ColumnColoringResult(A::AbstractMatrix, color::Vector{Int})
+    group = group_by_color(color)
+    return ColumnColoringResult(A, color, group, nothing)
 end
 
 function ColumnColoringResult(S::SparseMatrixCSC, color::Vector{Int})
@@ -141,12 +147,18 @@ $TYPEDFIELDS
 
 - [`AbstractColoringResult`](@ref)
 """
-struct RowColoringResult{M} <: AbstractColoringResult{:nonsymmetric,:row,:direct,M}
+struct RowColoringResult{M,Mᵀ,C<:Union{Nothing,Vector{Int}}} <:
+       AbstractColoringResult{:nonsymmetric,:row,:direct,M}
     S::M
-    Sᵀ::M
+    Sᵀ::Mᵀ
     color::Vector{Int}
     group::Vector{Vector{Int}}
-    compressed_indices::Vector{Int}
+    compressed_indices::C
+end
+
+function RowColoringResult(A::AbstractMatrix, color::Vector{Int})
+    group = group_by_color(color)
+    return RowColoringResult(A, nothing, color, group, nothing)
 end
 
 function RowColoringResult(S::SparseMatrixCSC, color::Vector{Int})
diff --git a/src/structured.jl b/src/structured.jl
new file mode 100644
index 00000000..7a2f7a53
--- /dev/null
+++ b/src/structured.jl
@@ -0,0 +1,168 @@
+#=
+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?
+=#
+
+"""
+    cycle_until(iterator, max_length::Integer)
+
+Concatenate copies of `iterator` to fill a vector of length `max_length` (with one partial copy allowed at the end).
+"""
+function cycle_until(iterator, max_length::Integer)
+    a = repeat(iterator, div(max_length, length(iterator)) + 1)
+    return resize!(a, max_length)
+end
+
+## Diagonal
+
+function coloring(
+    A::Diagonal,
+    ::ColoringProblem{:nonsymmetric,:column},
+    algo::GreedyColoringAlgorithm;
+    kwargs...,
+)
+    color = fill(1, size(A, 2))
+    return ColumnColoringResult(A, color)
+end
+
+function coloring(
+    A::Diagonal,
+    ::ColoringProblem{:nonsymmetric,:row},
+    algo::GreedyColoringAlgorithm;
+    kwargs...,
+)
+    color = fill(1, size(A, 1))
+    return RowColoringResult(A, color)
+end
+
+function decompress!(
+    A::Diagonal{R}, B::AbstractMatrix{R}, result::ColumnColoringResult
+) where {R<:Real}
+    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}
+    color = row_colors(result)
+    for i in axes(A, 1)
+        A[i, i] = B[color[i], i]
+    end
+    return A
+end
+
+## Bidiagonal
+
+function coloring(
+    A::Bidiagonal,
+    ::ColoringProblem{:nonsymmetric,:column},
+    algo::GreedyColoringAlgorithm;
+    kwargs...,
+)
+    color = cycle_until(1:2, size(A, 2))
+    return ColumnColoringResult(A, color)
+end
+
+function coloring(
+    A::Bidiagonal,
+    ::ColoringProblem{:nonsymmetric,:row},
+    algo::GreedyColoringAlgorithm;
+    kwargs...,
+)
+    color = cycle_until(1:2, size(A, 1))
+    return RowColoringResult(A, color)
+end
+
+function decompress!(
+    A::Bidiagonal{R}, B::AbstractMatrix{R}, result::ColumnColoringResult
+) where {R<:Real}
+    color = column_colors(result)
+    for j in axes(A, 2)
+        c = color[j]
+        A[j, j] = B[j, c]
+        if A.uplo == 'U' && j > 1  # above
+            A[j - 1, j] = B[j - 1, c]
+        elseif A.uplo == 'L' && j < size(A, 2)  # below
+            A[j + 1, j] = B[j + 1, c]
+        end
+    end
+    return A
+end
+
+function decompress!(
+    A::Bidiagonal{R}, B::AbstractMatrix{R}, result::RowColoringResult
+) where {R<:Real}
+    color = row_colors(result)
+    for i in axes(A, 1)
+        c = color[i]
+        A[i, i] = B[c, i]
+        if A.uplo == 'U' && i < size(A, 1)  # right
+            A[i, i + 1] = B[c, i + 1]
+        elseif A.uplo == 'L' && i > 1  # left
+            A[i, i - 1] = B[c, i - 1]
+        end
+    end
+    return A
+end
+
+## Tridiagonal
+
+function coloring(
+    A::Tridiagonal,
+    ::ColoringProblem{:nonsymmetric,:column},
+    algo::GreedyColoringAlgorithm;
+    kwargs...,
+)
+    color = cycle_until(1:3, size(A, 2))
+    return ColumnColoringResult(A, color)
+end
+
+function coloring(
+    A::Tridiagonal,
+    ::ColoringProblem{:nonsymmetric,:row},
+    algo::GreedyColoringAlgorithm;
+    kwargs...,
+)
+    color = cycle_until(1:3, size(A, 1))
+    return RowColoringResult(A, color)
+end
+
+function decompress!(
+    A::Tridiagonal{R}, B::AbstractMatrix{R}, result::ColumnColoringResult
+) where {R<:Real}
+    color = column_colors(result)
+    for j in axes(A, 2)
+        c = color[j]
+        A[j, j] = B[j, c]
+        if j > 1  # above
+            A[j - 1, j] = B[j - 1, c]
+        end
+        if j < size(A, 2)  # below
+            A[j + 1, j] = B[j + 1, c]
+        end
+    end
+    return A
+end
+
+function decompress!(
+    A::Tridiagonal{R}, B::AbstractMatrix{R}, result::RowColoringResult
+) where {R<:Real}
+    color = row_colors(result)
+    for i in axes(A, 1)
+        c = color[i]
+        A[i, i] = B[c, i]
+        if i < size(A, 1)  # right
+            A[i, i + 1] = B[c, i + 1]
+        end
+        if i > 1  # left
+            A[i, i - 1] = B[c, i - 1]
+        end
+    end
+    return A
+end
diff --git a/test/allocations.jl b/test/allocations.jl
index 36e67959..f0d3bbeb 100644
--- a/test/allocations.jl
+++ b/test/allocations.jl
@@ -2,7 +2,7 @@ using Chairmarks
 using LinearAlgebra
 using SparseArrays
 using SparseMatrixColorings
-using SparseMatrixColorings: partial_distance2_coloring!
+using SparseMatrixColorings: BipartiteGraph, partial_distance2_coloring!
 using StableRNGs
 using Test
 
@@ -21,7 +21,7 @@ end
     test_noallocs_distance2_coloring(1000)
 end;
 
-function test_noallocs_decompression(
+function test_noallocs_sparse_decompression(
     n::Integer; structure::Symbol, partition::Symbol, decompression::Symbol
 )
     A = sparse(Symmetric(sprand(rng, n, n, 5 / n)))
@@ -75,7 +75,27 @@ function test_noallocs_decompression(
     end
 end
 
-@testset "Decompression" begin
+function test_noallocs_structured_decompression(
+    n::Integer; structure::Symbol, partition::Symbol, decompression::Symbol
+)
+    @testset "$(nameof(typeof(A)))" for A in [
+        Diagonal(rand(n)),
+        Bidiagonal(rand(n), rand(n - 1), 'U'),
+        Bidiagonal(rand(n), rand(n - 1), 'L'),
+        Tridiagonal(rand(n - 1), rand(n), rand(n - 1)),
+    ]
+        result = coloring(
+            A,
+            ColoringProblem(; structure, partition),
+            GreedyColoringAlgorithm(; decompression),
+        )
+        B = compress(A, result)
+        bench = @be similar(A) decompress!(_, B, result) evals = 1
+        @test minimum(bench).allocs == 0
+    end
+end
+
+@testset "Sparse decompression" begin
     @testset "$structure - $partition - $decompression" for (
         structure, partition, decompression
     ) in [
@@ -84,6 +104,16 @@ end
         (:symmetric, :column, :direct),
         (:symmetric, :column, :substitution),
     ]
-        test_noallocs_decompression(1000; structure, partition, decompression)
+        test_noallocs_sparse_decompression(1000; structure, partition, decompression)
+    end
+end;
+
+@testset "Structured decompression" begin
+    @testset "$structure - $partition - $decompression" for (
+        structure, partition, decompression
+    ) in [
+        (:nonsymmetric, :column, :direct), (:nonsymmetric, :row, :direct)
+    ]
+        test_noallocs_structured_decompression(1000; structure, partition, decompression)
     end
 end;
diff --git a/test/runtests.jl b/test/runtests.jl
index 51ed33a8..ca0a09af 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -53,6 +53,9 @@ include("utils.jl")
         @testset "Random instances" begin
             include("random.jl")
         end
+        @testset "Structured matrices" begin
+            include("structured.jl")
+        end
         @testset "Instances with known colorings" begin
             include("theory.jl")
         end
diff --git a/test/structured.jl b/test/structured.jl
new file mode 100644
index 00000000..5ee336b1
--- /dev/null
+++ b/test/structured.jl
@@ -0,0 +1,62 @@
+using LinearAlgebra
+using SparseMatrixColorings
+using Test
+
+column_problem = ColoringProblem(; structure=:nonsymmetric, partition=:column)
+row_problem = ColoringProblem(; structure=:nonsymmetric, partition=:row)
+
+algo = GreedyColoringAlgorithm()
+
+@testset "Diagonal" begin
+    for n in (1, 2, 10, 100, 1000)
+        A = Diagonal(rand(n))
+        # column
+        result = coloring(A, column_problem, algo)
+        B = compress(A, result)
+        @test size(B, 2) == 1
+        @test decompress(B, result) == A
+        # row
+        result = coloring(A, row_problem, algo)
+        B = compress(A, result)
+        @test size(B, 1) == 1
+        @test decompress(B, result) == A
+    end
+end
+
+@testset "Bidiagonal" begin
+    for n in (2, 10, 100, 1000)
+        A1 = Bidiagonal(rand(n), rand(n - 1), :U)
+        A2 = Bidiagonal(rand(n), rand(n - 1), :L)
+        for A in (A1, A2)
+            # column
+            result = coloring(A, column_problem, algo)
+            B = compress(A, result)
+            @test size(B, 2) == 2
+            @test decompress(B, result) == A
+            # row
+            result = coloring(A, row_problem, algo)
+            B = compress(A, result)
+            @test size(B, 1) == 2
+            @test decompress(B, result) == A
+        end
+    end
+end
+
+@testset "Tridiagonal" begin
+    for n in (2, 10, 100, 1000)
+        A1 = Tridiagonal(rand(n - 1), rand(n), rand(n - 1))
+        A2 = Tridiagonal(rand(n - 1), rand(n), rand(n - 1))
+        for A in (A1, A2)
+            # column
+            result = coloring(A, column_problem, algo)
+            B = compress(A, result)
+            @test size(B, 2) == min(n, 3)
+            @test decompress(B, result) == A
+            # row
+            result = coloring(A, row_problem, algo)
+            B = compress(A, result)
+            @test size(B, 1) == min(n, 3)
+            @test decompress(B, result) == A
+        end
+    end
+end
diff --git a/test/type_stability.jl b/test/type_stability.jl
index a62d2e8c..81848ab2 100644
--- a/test/type_stability.jl
+++ b/test/type_stability.jl
@@ -3,13 +3,13 @@ using JET
 using LinearAlgebra
 using SparseArrays
 using SparseMatrixColorings
-using SparseMatrixColorings: respectful_similar
+using SparseMatrixColorings: matrix_versions, respectful_similar
 using StableRNGs
 using Test
 
 rng = StableRNG(63)
 
-@testset "Coloring" begin
+@testset "Sparse coloring" begin
     n = 10
     A = sprand(rng, n, n, 5 / n)
 
@@ -40,9 +40,28 @@ rng = StableRNG(63)
             GreedyColoringAlgorithm(; decompression),
         )
     end
-end
+end;
 
-@testset "Decompression" begin
+@testset "Structured coloring" begin
+    n = 10
+    @testset "$(nameof(typeof(A))) - $structure - $partition - $decompression" for A in [
+            Diagonal(rand(n)),
+            Bidiagonal(rand(n), rand(n - 1), 'U'),
+            Bidiagonal(rand(n), rand(n - 1), 'L'),
+            Tridiagonal(rand(n - 1), rand(n), rand(n - 1)),
+        ],
+        (structure, partition, decompression) in
+        [(:nonsymmetric, :column, :direct), (:nonsymmetric, :row, :direct)]
+
+        @test_opt target_modules = (SparseMatrixColorings,) coloring(
+            A,
+            ColoringProblem(; structure, partition),
+            GreedyColoringAlgorithm(; decompression),
+        )
+    end
+end;
+
+@testset "Sparse decompression" begin
     n = 10
     A0 = sparse(Symmetric(sprand(rng, n, n, 5 / n)))
 
@@ -92,3 +111,24 @@ end
         end
     end
 end;
+
+@testset "Structured decompression" begin
+    n = 10
+    @testset "$(nameof(typeof(A))) - $structure - $partition - $decompression" for A in [
+            Diagonal(rand(n)),
+            Bidiagonal(rand(n), rand(n - 1), 'U'),
+            Bidiagonal(rand(n), rand(n - 1), 'L'),
+            Tridiagonal(rand(n - 1), rand(n), rand(n - 1)),
+        ],
+        (structure, partition, decompression) in
+        [(:nonsymmetric, :column, :direct), (:nonsymmetric, :row, :direct)]
+
+        result = coloring(
+            A,
+            ColoringProblem(; structure, partition),
+            GreedyColoringAlgorithm(; decompression);
+        )
+        B = compress(A, result)
+        @test_opt decompress(B, result)
+    end
+end;

From c3a7242267085db983b41aa252ceae202959aa7d Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Wed, 2 Oct 2024 10:33:26 +0200
Subject: [PATCH 02/33] Store graph in result to allow generic matrices

---
 src/constant.jl      |   8 +--
 src/decompression.jl | 114 ++++++++++++++++++++-----------------------
 src/graph.jl         |  27 ++++++++++
 src/interface.jl     |  29 +++++------
 src/matrices.jl      |  13 ++---
 src/result.jl        |  74 ++++++++++++++++++----------
 test/graph.jl        |  12 +++--
 test/utils.jl        |   6 ++-
 8 files changed, 165 insertions(+), 118 deletions(-)

diff --git a/src/constant.jl b/src/constant.jl
index 046a1b0f..dd10cf60 100644
--- a/src/constant.jl
+++ b/src/constant.jl
@@ -75,8 +75,8 @@ end
 function ConstantColoringAlgorithm{:column}(
     matrix_template::AbstractMatrix, color::Vector{Int}
 )
-    S = convert(SparseMatrixCSC, matrix_template)
-    result = ColumnColoringResult(S, color)
+    bg = BipartiteGraph(matrix_template)
+    result = ColumnColoringResult(matrix_template, bg, color)
     M, R = typeof(matrix_template), typeof(result)
     return ConstantColoringAlgorithm{:column,M,R}(matrix_template, color, result)
 end
@@ -84,8 +84,8 @@ end
 function ConstantColoringAlgorithm{:row}(
     matrix_template::AbstractMatrix, color::Vector{Int}
 )
-    S = convert(SparseMatrixCSC, matrix_template)
-    result = RowColoringResult(S, color)
+    bg = BipartiteGraph(matrix_template)
+    result = RowColoringResult(matrix_template, bg, color)
     M, R = typeof(matrix_template), typeof(result)
     return ConstantColoringAlgorithm{:row,M,R}(matrix_template, color, result)
 end
diff --git a/src/decompression.jl b/src/decompression.jl
index 1d9e92a4..209df1d7 100644
--- a/src/decompression.jl
+++ b/src/decompression.jl
@@ -115,9 +115,8 @@ true
 - [`ColoringProblem`](@ref)
 - [`AbstractColoringResult`](@ref)
 """
-function decompress(B::AbstractMatrix{R}, result::AbstractColoringResult) where {R<:Real}
-    @compat (; S) = result
-    A = respectful_similar(S, R)
+function decompress(B::AbstractMatrix, result::AbstractColoringResult)
+    A = respectful_similar(result.A, eltype(B))
     return decompress!(A, B, result)
 end
 
@@ -264,12 +263,11 @@ end
 
 ## ColumnColoringResult
 
-function decompress!(
-    A::AbstractMatrix{R}, B::AbstractMatrix{R}, result::ColumnColoringResult
-) where {R<:Real}
-    @compat (; S, color) = result
+function decompress!(A::AbstractMatrix, B::AbstractMatrix, result::ColumnColoringResult)
+    @compat (; color) = result
+    S = result.bg.S2
     check_same_pattern(A, S)
-    A .= zero(R)
+    fill!(A, zero(eltype(A)))
     rvS = rowvals(S)
     for j in axes(S, 2)
         cj = color[j]
@@ -282,9 +280,10 @@ function decompress!(
 end
 
 function decompress_single_color!(
-    A::AbstractMatrix{R}, b::AbstractVector{R}, c::Integer, result::ColumnColoringResult
-) where {R<:Real}
-    @compat (; S, group) = result
+    A::AbstractMatrix, b::AbstractVector, c::Integer, result::ColumnColoringResult
+)
+    @compat (; group) = result
+    S = result.bg.S2
     check_same_pattern(A, S)
     rvS = rowvals(S)
     for j in group[c]
@@ -296,10 +295,9 @@ function decompress_single_color!(
     return A
 end
 
-function decompress!(
-    A::SparseMatrixCSC{R}, B::AbstractMatrix{R}, result::ColumnColoringResult
-) where {R<:Real}
-    @compat (; S, compressed_indices) = result
+function decompress!(A::SparseMatrixCSC, B::AbstractMatrix, result::ColumnColoringResult)
+    @compat (; compressed_indices) = result
+    S = result.bg.S2
     check_same_pattern(A, S)
     nzA = nonzeros(A)
     for k in eachindex(nzA, compressed_indices)
@@ -309,9 +307,10 @@ function decompress!(
 end
 
 function decompress_single_color!(
-    A::SparseMatrixCSC{R}, b::AbstractVector{R}, c::Integer, result::ColumnColoringResult
-) where {R<:Real}
-    @compat (; S, group) = result
+    A::SparseMatrixCSC, b::AbstractVector, c::Integer, result::ColumnColoringResult
+)
+    @compat (; group) = result
+    S = result.bg.S2
     check_same_pattern(A, S)
     rvS = rowvals(S)
     nzA = nonzeros(A)
@@ -326,12 +325,11 @@ end
 
 ## RowColoringResult
 
-function decompress!(
-    A::AbstractMatrix{R}, B::AbstractMatrix{R}, result::RowColoringResult
-) where {R<:Real}
-    @compat (; S, color) = result
+function decompress!(A::AbstractMatrix, B::AbstractMatrix, result::RowColoringResult)
+    @compat (; color) = result
+    S = result.bg.S2
     check_same_pattern(A, S)
-    A .= zero(R)
+    fill!(A, zero(eltype(A)))
     rvS = rowvals(S)
     for j in axes(S, 2)
         for k in nzrange(S, j)
@@ -344,9 +342,10 @@ function decompress!(
 end
 
 function decompress_single_color!(
-    A::AbstractMatrix{R}, b::AbstractVector{R}, c::Integer, result::RowColoringResult
-) where {R<:Real}
-    @compat (; S, Sᵀ, group) = result
+    A::AbstractMatrix, b::AbstractVector, c::Integer, result::RowColoringResult
+)
+    @compat (; group) = result
+    S, Sᵀ = result.bg.S2, result.bg.S1
     check_same_pattern(A, S)
     rvSᵀ = rowvals(Sᵀ)
     for i in group[c]
@@ -358,10 +357,9 @@ function decompress_single_color!(
     return A
 end
 
-function decompress!(
-    A::SparseMatrixCSC{R}, B::AbstractMatrix{R}, result::RowColoringResult
-) where {R<:Real}
-    @compat (; S, compressed_indices) = result
+function decompress!(A::SparseMatrixCSC, B::AbstractMatrix, result::RowColoringResult)
+    @compat (; compressed_indices) = result
+    S = result.bg.S2
     check_same_pattern(A, S)
     nzA = nonzeros(A)
     for k in eachindex(nzA, compressed_indices)
@@ -373,15 +371,13 @@ end
 ## StarSetColoringResult
 
 function decompress!(
-    A::AbstractMatrix{R},
-    B::AbstractMatrix{R},
-    result::StarSetColoringResult,
-    uplo::Symbol=:F,
-) where {R<:Real}
-    @compat (; S, color, star_set) = result
+    A::AbstractMatrix, B::AbstractMatrix, result::StarSetColoringResult, uplo::Symbol=:F
+)
+    @compat (; color, star_set) = result
     @compat (; star, hub, spokes) = star_set
+    S = result.ag.S
     uplo == :F && check_same_pattern(A, S)
-    A .= zero(R)
+    A .= zero(eltype(A))
     for i in axes(A, 1)
         if !iszero(S[i, i])
             A[i, i] = B[i, color[i]]
@@ -403,14 +399,15 @@ function decompress!(
 end
 
 function decompress_single_color!(
-    A::AbstractMatrix{R},
-    b::AbstractVector{R},
+    A::AbstractMatrix,
+    b::AbstractVector,
     c::Integer,
     result::StarSetColoringResult,
     uplo::Symbol=:F,
-) where {R<:Real}
-    @compat (; S, color, group, star_set) = result
+)
+    @compat (; color, group, star_set) = result
     @compat (; hub, spokes) = star_set
+    S = result.ag.S
     uplo == :F && check_same_pattern(A, S)
     for i in axes(A, 1)
         if !iszero(S[i, i]) && color[i] == c
@@ -434,12 +431,10 @@ function decompress_single_color!(
 end
 
 function decompress!(
-    A::SparseMatrixCSC{R},
-    B::AbstractMatrix{R},
-    result::StarSetColoringResult,
-    uplo::Symbol=:F,
-) where {R<:Real}
-    @compat (; S, compressed_indices) = result
+    A::SparseMatrixCSC, B::AbstractMatrix, result::StarSetColoringResult, uplo::Symbol=:F
+)
+    @compat (; compressed_indices) = result
+    S = result.ag.S
     nzA = nonzeros(A)
     if uplo == :F
         check_same_pattern(A, S)
@@ -468,13 +463,12 @@ end
 # TODO: add method for A::SparseMatrixCSC
 
 function decompress!(
-    A::AbstractMatrix{R},
-    B::AbstractMatrix{R},
-    result::TreeSetColoringResult,
-    uplo::Symbol=:F,
-) where {R<:Real}
-    @compat (; S, color, vertices_by_tree, reverse_bfs_orders, buffer) = result
+    A::AbstractMatrix, B::AbstractMatrix, result::TreeSetColoringResult, uplo::Symbol=:F
+)
+    @compat (; color, vertices_by_tree, reverse_bfs_orders, buffer) = result
+    S = result.ag.S
     uplo == :F && check_same_pattern(A, S)
+    R = eltype(A)
     A .= zero(R)
 
     if eltype(buffer) == R
@@ -513,19 +507,19 @@ end
 ## MatrixInverseColoringResult
 
 function decompress!(
-    A::AbstractMatrix{R},
-    B::AbstractMatrix{R},
+    A::AbstractMatrix,
+    B::AbstractMatrix,
     result::LinearSystemColoringResult,
     uplo::Symbol=:F,
-) where {R<:Real}
-    @compat (;
-        S, color, strict_upper_nonzero_inds, T_factorization, strict_upper_nonzeros_A
-    ) = result
+)
+    @compat (; color, strict_upper_nonzero_inds, T_factorization, strict_upper_nonzeros_A) =
+        result
+    S = result.ag.S
     uplo == :F && check_same_pattern(A, S)
 
     # TODO: for some reason I cannot use ldiv! with a sparse QR
     strict_upper_nonzeros_A = T_factorization \ vec(B)
-    A .= zero(R)
+    A .= zero(eltype(A))
     for i in axes(A, 1)
         if !iszero(S[i, i])
             A[i, i] = B[i, color[i]]
diff --git a/src/graph.jl b/src/graph.jl
index 9d13b827..ec53aa6a 100644
--- a/src/graph.jl
+++ b/src/graph.jl
@@ -24,6 +24,19 @@ end
 SparsityPatternCSC(A::SparseMatrixCSC) = SparsityPatternCSC(A.m, A.n, A.colptr, A.rowval)
 
 Base.size(S::SparsityPatternCSC) = (S.m, S.n)
+
+function Base.size(S::SparsityPatternCSC, d::Integer)
+    if d == 1
+        return S.m
+    elseif d == 2
+        return S.n
+    else
+        return 1
+    end
+end
+
+Base.axes(S::SparsityPatternCSC, d::Integer) = Base.OneTo(size(S, d))
+
 SparseArrays.nnz(S::SparsityPatternCSC) = length(S.rowval)
 SparseArrays.rowvals(S::SparsityPatternCSC) = S.rowval
 SparseArrays.nzrange(S::SparsityPatternCSC, j::Integer) = S.colptr[j]:(S.colptr[j + 1] - 1)
@@ -81,6 +94,15 @@ function Base.transpose(S::SparsityPatternCSC{T}) where {T}
     return SparsityPatternCSC{T}(n, m, B_colptr, B_rowval)
 end
 
+# copied from SparseArrays.jl
+function Base.getindex(S::SparsityPatternCSC, i0::Integer, i1::Integer)
+    r1 = Int(S.colptr[i1])
+    r2 = Int(S.colptr[i1 + 1] - 1)
+    (r1 > r2) && return false
+    r1 = searchsortedfirst(rowvals(S), i0, r1, r2, Base.Order.Forward)
+    return ((r1 > r2) || (rowvals(S)[r1] != i0)) ? false : true
+end
+
 ## Adjacency graph
 
 """
@@ -109,6 +131,7 @@ struct AdjacencyGraph{T}
     S::SparsityPatternCSC{T}
 end
 
+AdjacencyGraph(A::AbstractMatrix) = AdjacencyGraph(SparseMatrixCSC(A))
 AdjacencyGraph(A::SparseMatrixCSC) = AdjacencyGraph(SparsityPatternCSC(A))
 
 pattern(g::AdjacencyGraph) = g.S
@@ -183,6 +206,10 @@ struct BipartiteGraph{T<:Integer}
     S2::SparsityPatternCSC{T}
 end
 
+function BipartiteGraph(A::AbstractMatrix; symmetric_pattern::Bool=false)
+    return BipartiteGraph(SparseMatrixCSC(A); symmetric_pattern)
+end
+
 function BipartiteGraph(A::SparseMatrixCSC; symmetric_pattern::Bool=false)
     S2 = SparsityPatternCSC(A)  # columns to rows
     if symmetric_pattern
diff --git a/src/interface.jl b/src/interface.jl
index 5a559ffc..e943d2d1 100644
--- a/src/interface.jl
+++ b/src/interface.jl
@@ -180,12 +180,11 @@ function coloring(
     decompression_eltype::Type=Float64,
     symmetric_pattern::Bool=false,
 )
-    S = convert(SparseMatrixCSC, A)
     bg = BipartiteGraph(
-        S; symmetric_pattern=symmetric_pattern || A isa Union{Symmetric,Hermitian}
+        A; symmetric_pattern=symmetric_pattern || A isa Union{Symmetric,Hermitian}
     )
     color = partial_distance2_coloring(bg, Val(2), algo.order)
-    return ColumnColoringResult(S, color)
+    return ColumnColoringResult(A, bg, color)
 end
 
 function coloring(
@@ -195,12 +194,11 @@ function coloring(
     decompression_eltype::Type=Float64,
     symmetric_pattern::Bool=false,
 )
-    S = convert(SparseMatrixCSC, A)
     bg = BipartiteGraph(
-        S; symmetric_pattern=symmetric_pattern || A isa Union{Symmetric,Hermitian}
+        A; symmetric_pattern=symmetric_pattern || A isa Union{Symmetric,Hermitian}
     )
     color = partial_distance2_coloring(bg, Val(1), algo.order)
-    return RowColoringResult(S, color)
+    return RowColoringResult(A, bg, color)
 end
 
 function coloring(
@@ -209,10 +207,9 @@ function coloring(
     algo::GreedyColoringAlgorithm{:direct};
     decompression_eltype::Type=Float64,
 )
-    S = convert(SparseMatrixCSC, A)
-    ag = AdjacencyGraph(S)
+    ag = AdjacencyGraph(A)
     color, star_set = star_coloring(ag, algo.order)
-    return StarSetColoringResult(S, color, star_set)
+    return StarSetColoringResult(A, ag, color, star_set)
 end
 
 function coloring(
@@ -221,31 +218,27 @@ function coloring(
     algo::GreedyColoringAlgorithm{:substitution};
     decompression_eltype::Type=Float64,
 )
-    S = convert(SparseMatrixCSC, A)
-    ag = AdjacencyGraph(S)
+    ag = AdjacencyGraph(A)
     color, tree_set = acyclic_coloring(ag, algo.order)
-    return TreeSetColoringResult(S, color, tree_set, decompression_eltype)
+    return TreeSetColoringResult(A, ag, color, tree_set, decompression_eltype)
 end
 
 ## ADTypes interface
 
 function ADTypes.column_coloring(A::AbstractMatrix, algo::GreedyColoringAlgorithm)
-    S = convert(SparseMatrixCSC, A)
-    bg = BipartiteGraph(S; symmetric_pattern=A isa Union{Symmetric,Hermitian})
+    bg = BipartiteGraph(A; symmetric_pattern=A isa Union{Symmetric,Hermitian})
     color = partial_distance2_coloring(bg, Val(2), algo.order)
     return color
 end
 
 function ADTypes.row_coloring(A::AbstractMatrix, algo::GreedyColoringAlgorithm)
-    S = convert(SparseMatrixCSC, A)
-    bg = BipartiteGraph(S; symmetric_pattern=A isa Union{Symmetric,Hermitian})
+    bg = BipartiteGraph(A; symmetric_pattern=A isa Union{Symmetric,Hermitian})
     color = partial_distance2_coloring(bg, Val(1), algo.order)
     return color
 end
 
 function ADTypes.symmetric_coloring(A::AbstractMatrix, algo::GreedyColoringAlgorithm)
-    S = convert(SparseMatrixCSC, A)
-    ag = AdjacencyGraph(S)
+    ag = AdjacencyGraph(A)
     color, star_set = star_coloring(ag, algo.order)
     return color
 end
diff --git a/src/matrices.jl b/src/matrices.jl
index 2ccfed12..795983c8 100644
--- a/src/matrices.jl
+++ b/src/matrices.jl
@@ -62,22 +62,23 @@ function respectful_similar(A::Union{Symmetric,Hermitian}, ::Type{T}) where {T}
 end
 
 """
-    same_pattern(A::AbstractMatrix, B::AbstractMatrix)
+    same_pattern(A, B)
 
 Perform a partial equality check on the sparsity patterns of `A` and `B`:
 
 - if the return is `true`, they might have the same sparsity pattern but we're not sure
 - if the return is `false`, they definitely don't have the same sparsity pattern
 """
-function same_pattern(A::AbstractMatrix, B::AbstractMatrix)
-    return size(A) == size(B)
-end
+same_pattern(A, B) = size(A) == size(B)
 
-function same_pattern(A::SparseMatrixCSC, B::SparseMatrixCSC)
+function same_pattern(
+    A::Union{SparseMatrixCSC,SparsityPatternCSC},
+    B::Union{SparseMatrixCSC,SparsityPatternCSC},
+)
     return size(A) == size(B) && nnz(A) == nnz(B)
 end
 
-function check_same_pattern(A::AbstractMatrix, S::AbstractMatrix)
+function check_same_pattern(A, S)
     if !same_pattern(A, S)
         throw(DimensionMismatch("`A` and `S` must have the same sparsity pattern."))
     end
diff --git a/src/result.jl b/src/result.jl
index 62ff4864..1c35bcdb 100644
--- a/src/result.jl
+++ b/src/result.jl
@@ -24,7 +24,7 @@ Combination between the type parameters of [`ColoringProblem`](@ref) and [`Greed
 !!! warning
     Unlike the methods above, the concrete subtypes of `AbstractColoringResult` are not part of the public API and may change without notice.
 """
-abstract type AbstractColoringResult{structure,partition,decompression,M<:SparseMatrixCSC} end
+abstract type AbstractColoringResult{structure,partition,decompression} end
 
 """
     column_colors(result::AbstractColoringResult)
@@ -99,9 +99,12 @@ $TYPEDFIELDS
 
 - [`AbstractColoringResult`](@ref)
 """
-struct ColumnColoringResult{M} <: AbstractColoringResult{:nonsymmetric,:column,:direct,M}
+struct ColumnColoringResult{M<:AbstractMatrix,G<:BipartiteGraph} <:
+       AbstractColoringResult{:nonsymmetric,:column,:direct}
     "matrix that was colored"
-    S::M
+    A::M
+    "bipartite graph that was used for coloring"
+    bg::G
     "one integer color for each column or row (depending on `partition`)"
     color::Vector{Int}
     "color groups for columns or rows (depending on `partition`)"
@@ -110,7 +113,8 @@ struct ColumnColoringResult{M} <: AbstractColoringResult{:nonsymmetric,:column,:
     compressed_indices::Vector{Int}
 end
 
-function ColumnColoringResult(S::SparseMatrixCSC, color::Vector{Int})
+function ColumnColoringResult(A::AbstractMatrix, bg::BipartiteGraph, color::Vector{Int})
+    S = bg.S2
     group = group_by_color(color)
     n = size(S, 1)
     rv = rowvals(S)
@@ -123,7 +127,7 @@ function ColumnColoringResult(S::SparseMatrixCSC, color::Vector{Int})
             compressed_indices[k] = (c - 1) * n + i
         end
     end
-    return ColumnColoringResult(S, color, group, compressed_indices)
+    return ColumnColoringResult(A, bg, color, group, compressed_indices)
 end
 
 """
@@ -141,16 +145,17 @@ $TYPEDFIELDS
 
 - [`AbstractColoringResult`](@ref)
 """
-struct RowColoringResult{M} <: AbstractColoringResult{:nonsymmetric,:row,:direct,M}
-    S::M
-    Sᵀ::M
+struct RowColoringResult{M<:AbstractMatrix,G<:BipartiteGraph} <:
+       AbstractColoringResult{:nonsymmetric,:row,:direct}
+    A::M
+    bg::G
     color::Vector{Int}
     group::Vector{Vector{Int}}
     compressed_indices::Vector{Int}
 end
 
-function RowColoringResult(S::SparseMatrixCSC, color::Vector{Int})
-    Sᵀ = convert(SparseMatrixCSC, transpose(S))
+function RowColoringResult(A::AbstractMatrix, bg::BipartiteGraph, color::Vector{Int})
+    S = bg.S2
     group = group_by_color(color)
     C = length(group)  # ncolors
     rv = rowvals(S)
@@ -163,7 +168,7 @@ function RowColoringResult(S::SparseMatrixCSC, color::Vector{Int})
             compressed_indices[k] = (j - 1) * C + c
         end
     end
-    return RowColoringResult(S, Sᵀ, color, group, compressed_indices)
+    return RowColoringResult(A, bg, color, group, compressed_indices)
 end
 
 """
@@ -181,15 +186,20 @@ $TYPEDFIELDS
 
 - [`AbstractColoringResult`](@ref)
 """
-struct StarSetColoringResult{M} <: AbstractColoringResult{:symmetric,:column,:direct,M}
-    S::M
+struct StarSetColoringResult{M<:AbstractMatrix,G<:AdjacencyGraph} <:
+       AbstractColoringResult{:symmetric,:column,:direct}
+    A::M
+    ag::G
     color::Vector{Int}
     group::Vector{Vector{Int}}
     star_set::StarSet
     compressed_indices::Vector{Int}
 end
 
-function StarSetColoringResult(S::SparseMatrixCSC, color::Vector{Int}, star_set::StarSet)
+function StarSetColoringResult(
+    A::AbstractMatrix, ag::AdjacencyGraph, color::Vector{Int}, star_set::StarSet
+)
+    S = ag.S
     group = group_by_color(color)
     n = size(S, 1)
     rv = rowvals(S)
@@ -202,7 +212,7 @@ function StarSetColoringResult(S::SparseMatrixCSC, color::Vector{Int}, star_set:
             compressed_indices[k] = (c - 1) * n + l
         end
     end
-    return StarSetColoringResult(S, color, group, star_set, compressed_indices)
+    return StarSetColoringResult(A, ag, color, group, star_set, compressed_indices)
 end
 
 """
@@ -220,9 +230,10 @@ $TYPEDFIELDS
 
 - [`AbstractColoringResult`](@ref)
 """
-struct TreeSetColoringResult{M,R} <:
-       AbstractColoringResult{:symmetric,:column,:substitution,M}
-    S::M
+struct TreeSetColoringResult{M<:AbstractMatrix,G<:AdjacencyGraph,R} <:
+       AbstractColoringResult{:symmetric,:column,:substitution}
+    A::M
+    ag::G
     color::Vector{Int}
     group::Vector{Vector{Int}}
     vertices_by_tree::Vector{Vector{Int}}
@@ -231,8 +242,13 @@ struct TreeSetColoringResult{M,R} <:
 end
 
 function TreeSetColoringResult(
-    S::SparseMatrixCSC, color::Vector{Int}, tree_set::TreeSet, decompression_eltype::Type{R}
+    A::AbstractMatrix,
+    ag::AdjacencyGraph,
+    color::Vector{Int},
+    tree_set::TreeSet,
+    decompression_eltype::Type{R},
 ) where {R}
+    S = ag.S
     nvertices = length(color)
     group = group_by_color(color)
 
@@ -347,7 +363,7 @@ function TreeSetColoringResult(
     buffer = Vector{R}(undef, nvertices)
 
     return TreeSetColoringResult(
-        S, color, group, vertices_by_tree, reverse_bfs_orders, buffer
+        A, ag, color, group, vertices_by_tree, reverse_bfs_orders, buffer
     )
 end
 
@@ -368,9 +384,10 @@ $TYPEDFIELDS
 
 - [`AbstractColoringResult`](@ref)
 """
-struct LinearSystemColoringResult{M,R,F} <:
-       AbstractColoringResult{:symmetric,:column,:substitution,M}
-    S::M
+struct LinearSystemColoringResult{M<:AbstractMatrix,G<:AdjacencyGraph,R,F} <:
+       AbstractColoringResult{:symmetric,:column,:substitution}
+    A::M
+    ag::G
     color::Vector{Int}
     group::Vector{Vector{Int}}
     strict_upper_nonzero_inds::Vector{Tuple{Int,Int}}
@@ -379,10 +396,11 @@ struct LinearSystemColoringResult{M,R,F} <:
 end
 
 function LinearSystemColoringResult(
-    S::SparseMatrixCSC, color::Vector{Int}, decompression_eltype::Type{R}
+    A::AbstractMatrix, ag::AdjacencyGraph, color::Vector{Int}, decompression_eltype::Type{R}
 ) where {R}
     group = group_by_color(color)
     C = length(group)  # ncolors
+    S = ag.S
     rv = rowvals(S)
 
     # build T such that T * strict_upper_nonzeros(A) = B
@@ -411,6 +429,12 @@ function LinearSystemColoringResult(
     strict_upper_nonzeros_A = Vector{float(R)}(undef, size(T, 2))
 
     return LinearSystemColoringResult(
-        S, color, group, strict_upper_nonzero_inds, strict_upper_nonzeros_A, T_factorization
+        A,
+        ag,
+        color,
+        group,
+        strict_upper_nonzero_inds,
+        strict_upper_nonzeros_A,
+        T_factorization,
     )
 end
diff --git a/test/graph.jl b/test/graph.jl
index 31d56f03..9f5be1a6 100644
--- a/test/graph.jl
+++ b/test/graph.jl
@@ -15,13 +15,19 @@ using Test
 
 @testset "SparsityPatternCSC" begin
     @testset "Transpose" begin
-        for _ in 1:1000
+        @test all(1:1000) do _
             A = sprand(rand(100:1000), rand(100:1000), 0.1)
             S = SparsityPatternCSC(A)
             Sᵀ = transpose(S)
             Sᵀ_true = SparsityPatternCSC(sparse(transpose(A)))
-            @test Sᵀ.colptr == Sᵀ_true.colptr
-            @test Sᵀ.rowval == Sᵀ_true.rowval
+            Sᵀ.colptr == Sᵀ_true.colptr && Sᵀ.rowval == Sᵀ_true.rowval
+        end
+    end
+    @testset "getindex" begin
+        A = sprand(Bool, 100, 100, 0.1)
+        S = SparsityPatternCSC(A)
+        @test all(zip(axes(S, 1), axes(S, 2))) do (i, j)
+            A[i, j] == S[i, j]
         end
     end
 end
diff --git a/test/utils.jl b/test/utils.jl
index 89a4749c..2ebad2a4 100644
--- a/test/utils.jl
+++ b/test/utils.jl
@@ -1,6 +1,7 @@
 using LinearAlgebra
 using SparseMatrixColorings
-using SparseMatrixColorings: LinearSystemColoringResult, matrix_versions, respectful_similar
+using SparseMatrixColorings:
+    AdjacencyGraph, LinearSystemColoringResult, matrix_versions, respectful_similar
 using Test
 
 function test_coloring_decompression(
@@ -99,7 +100,8 @@ function test_coloring_decompression(
 
         @testset "Linear system decompression" begin
             if structure == :symmetric && count(!iszero, A) > 0  # sparse factorization cannot handle empty matrices
-                linresult = LinearSystemColoringResult(sparse(A), color, Float64)
+                ag = AdjacencyGraph(A)
+                linresult = LinearSystemColoringResult(A, ag, color, Float64)
                 @test decompress(float.(B), linresult) ≈ A0
                 @test decompress!(respectful_similar(float.(A)), float.(B), linresult) ≈ A0
             end

From d5837b93fe7289b19e3dd35ea00f423ca3e6f7e8 Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Wed, 2 Oct 2024 10:42:51 +0200
Subject: [PATCH 03/33] Use `fill!` whenever possible

---
 src/decompression.jl | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/decompression.jl b/src/decompression.jl
index 209df1d7..42c5c576 100644
--- a/src/decompression.jl
+++ b/src/decompression.jl
@@ -377,7 +377,7 @@ function decompress!(
     @compat (; star, hub, spokes) = star_set
     S = result.ag.S
     uplo == :F && check_same_pattern(A, S)
-    A .= zero(eltype(A))
+    fill!(A, zero(eltype(A)))
     for i in axes(A, 1)
         if !iszero(S[i, i])
             A[i, i] = B[i, color[i]]
@@ -469,7 +469,7 @@ function decompress!(
     S = result.ag.S
     uplo == :F && check_same_pattern(A, S)
     R = eltype(A)
-    A .= zero(R)
+    fill!(A, zero(R))
 
     if eltype(buffer) == R
         buffer_right_type = buffer
@@ -519,7 +519,7 @@ function decompress!(
 
     # TODO: for some reason I cannot use ldiv! with a sparse QR
     strict_upper_nonzeros_A = T_factorization \ vec(B)
-    A .= zero(eltype(A))
+    fill!(A, zero(eltype(A)))
     for i in axes(A, 1)
         if !iszero(S[i, i])
             A[i, i] = B[i, color[i]]

From 12dcf60e027c124628f36c078aa9f431a3baec18 Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Wed, 2 Oct 2024 10:47:56 +0200
Subject: [PATCH 04/33] More tests

---
 test/graph.jl | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/test/graph.jl b/test/graph.jl
index 9f5be1a6..cb3a4467 100644
--- a/test/graph.jl
+++ b/test/graph.jl
@@ -23,6 +23,16 @@ using Test
             Sᵀ.colptr == Sᵀ_true.colptr && Sᵀ.rowval == Sᵀ_true.rowval
         end
     end
+    @testset "size" begin
+        A = spzeros(10, 20)
+        S = SparsityPatternCSC(A)
+        @test size(A) == size(S)
+        @test size(A, 1) == size(S, 1)
+        @test size(A, 2) == size(S, 2)
+        @test size(A, 3) == size(S, 3)
+        @test axes(A, 1) == axes(S, 1)
+        @test axes(A, 2) == axes(S, 2)
+    end
     @testset "getindex" begin
         A = sprand(Bool, 100, 100, 0.1)
         S = SparsityPatternCSC(A)

From 0e9adcd7ca05c9e9c520ff7a8d46a07f70a41fff Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Wed, 2 Oct 2024 10:58:46 +0200
Subject: [PATCH 05/33] Record useless BipartiteGraph

---
 src/result.jl     | 14 ++------------
 src/structured.jl | 18 ++++++++++++------
 2 files changed, 14 insertions(+), 18 deletions(-)

diff --git a/src/result.jl b/src/result.jl
index a0e5dc3d..1c35bcdb 100644
--- a/src/result.jl
+++ b/src/result.jl
@@ -110,12 +110,7 @@ struct ColumnColoringResult{M<:AbstractMatrix,G<:BipartiteGraph} <:
     "color groups for columns or rows (depending on `partition`)"
     group::Vector{Vector{Int}}
     "flattened indices mapping the compressed matrix `B` to the uncompressed matrix `A` when `A isa SparseMatrixCSC`. They satisfy `nonzeros(A)[k] = vec(B)[compressed_indices[k]]`"
-    compressed_indices::C
-end
-
-function ColumnColoringResult(A::AbstractMatrix, color::Vector{Int})
-    group = group_by_color(color)
-    return ColumnColoringResult(A, color, group, nothing)
+    compressed_indices::Vector{Int}
 end
 
 function ColumnColoringResult(A::AbstractMatrix, bg::BipartiteGraph, color::Vector{Int})
@@ -156,12 +151,7 @@ struct RowColoringResult{M<:AbstractMatrix,G<:BipartiteGraph} <:
     bg::G
     color::Vector{Int}
     group::Vector{Vector{Int}}
-    compressed_indices::C
-end
-
-function RowColoringResult(A::AbstractMatrix, color::Vector{Int})
-    group = group_by_color(color)
-    return RowColoringResult(A, nothing, color, group, nothing)
+    compressed_indices::Vector{Int}
 end
 
 function RowColoringResult(A::AbstractMatrix, bg::BipartiteGraph, color::Vector{Int})
diff --git a/src/structured.jl b/src/structured.jl
index 7a2f7a53..cdff6c6b 100644
--- a/src/structured.jl
+++ b/src/structured.jl
@@ -24,7 +24,8 @@ function coloring(
     kwargs...,
 )
     color = fill(1, size(A, 2))
-    return ColumnColoringResult(A, color)
+    bg = BipartiteGraph(A)
+    return ColumnColoringResult(A, bg, color)
 end
 
 function coloring(
@@ -34,7 +35,8 @@ function coloring(
     kwargs...,
 )
     color = fill(1, size(A, 1))
-    return RowColoringResult(A, color)
+    bg = BipartiteGraph(A)
+    return RowColoringResult(A, bg, color)
 end
 
 function decompress!(
@@ -66,7 +68,8 @@ function coloring(
     kwargs...,
 )
     color = cycle_until(1:2, size(A, 2))
-    return ColumnColoringResult(A, color)
+    bg = BipartiteGraph(A)
+    return ColumnColoringResult(A, bg, color)
 end
 
 function coloring(
@@ -76,7 +79,8 @@ function coloring(
     kwargs...,
 )
     color = cycle_until(1:2, size(A, 1))
-    return RowColoringResult(A, color)
+    bg = BipartiteGraph(A)
+    return RowColoringResult(A, bg, color)
 end
 
 function decompress!(
@@ -120,7 +124,8 @@ function coloring(
     kwargs...,
 )
     color = cycle_until(1:3, size(A, 2))
-    return ColumnColoringResult(A, color)
+    bg = BipartiteGraph(A)
+    return ColumnColoringResult(A, bg, color)
 end
 
 function coloring(
@@ -130,7 +135,8 @@ function coloring(
     kwargs...,
 )
     color = cycle_until(1:3, size(A, 1))
-    return RowColoringResult(A, color)
+    bg = BipartiteGraph(A)
+    return RowColoringResult(A, bg, color)
 end
 
 function decompress!(

From fe178d704a5cbf3bbe9ce35fd9780b079bf81a47 Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Wed, 2 Oct 2024 11:02:06 +0200
Subject: [PATCH 06/33] Fix docs

---
 docs/src/dev.md | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/docs/src/dev.md b/docs/src/dev.md
index b8aa2792..abded68b 100644
--- a/docs/src/dev.md
+++ b/docs/src/dev.md
@@ -65,3 +65,9 @@ SparseMatrixColorings.what_fig_61
 SparseMatrixColorings.efficient_fig_1
 SparseMatrixColorings.efficient_fig_4
 ```
+
+## Misc
+
+```@docs
+SparseMatrixColorings.cycle_until
+```
\ No newline at end of file

From 3937800471bf29b0de17fa942512e5b4ce4c77d7 Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Wed, 2 Oct 2024 12:26:22 +0200
Subject: [PATCH 07/33] Fix cycling inferrability

---
 docs/src/dev.md    |  4 ++--
 src/structured.jl  | 21 ++++++++++++---------
 test/structured.jl | 12 ++++++++++++
 3 files changed, 26 insertions(+), 11 deletions(-)

diff --git a/docs/src/dev.md b/docs/src/dev.md
index abded68b..cd30a9b6 100644
--- a/docs/src/dev.md
+++ b/docs/src/dev.md
@@ -69,5 +69,5 @@ SparseMatrixColorings.efficient_fig_4
 ## Misc
 
 ```@docs
-SparseMatrixColorings.cycle_until
-```
\ No newline at end of file
+SparseMatrixColorings.cycle_range
+```
diff --git a/src/structured.jl b/src/structured.jl
index cdff6c6b..9320e438 100644
--- a/src/structured.jl
+++ b/src/structured.jl
@@ -6,13 +6,16 @@ Question: when decompressing, should we always assume that the coloring was opti
 =#
 
 """
-    cycle_until(iterator, max_length::Integer)
+    cycle_range(k::Integer, n::Integer)
 
-Concatenate copies of `iterator` to fill a vector of length `max_length` (with one partial copy allowed at the end).
+Concatenate copies of `1:k` to fill a vector of length `n` (with one partial copy allowed at the end).
 """
-function cycle_until(iterator, max_length::Integer)
-    a = repeat(iterator, div(max_length, length(iterator)) + 1)
-    return resize!(a, max_length)
+function cycle_range(k::Integer, n::Integer)
+    color = Vector{Int}(undef, n)
+    for i in eachindex(color)
+        color[i] = 1 + (i - 1) % k
+    end
+    return color
 end
 
 ## Diagonal
@@ -67,7 +70,7 @@ function coloring(
     algo::GreedyColoringAlgorithm;
     kwargs...,
 )
-    color = cycle_until(1:2, size(A, 2))
+    color = cycle_range(1:2, size(A, 2))
     bg = BipartiteGraph(A)
     return ColumnColoringResult(A, bg, color)
 end
@@ -78,7 +81,7 @@ function coloring(
     algo::GreedyColoringAlgorithm;
     kwargs...,
 )
-    color = cycle_until(1:2, size(A, 1))
+    color = cycle_range(1:2, size(A, 1))
     bg = BipartiteGraph(A)
     return RowColoringResult(A, bg, color)
 end
@@ -123,7 +126,7 @@ function coloring(
     algo::GreedyColoringAlgorithm;
     kwargs...,
 )
-    color = cycle_until(1:3, size(A, 2))
+    color = cycle_range(1:3, size(A, 2))
     bg = BipartiteGraph(A)
     return ColumnColoringResult(A, bg, color)
 end
@@ -134,7 +137,7 @@ function coloring(
     algo::GreedyColoringAlgorithm;
     kwargs...,
 )
-    color = cycle_until(1:3, size(A, 1))
+    color = cycle_range(1:3, size(A, 1))
     bg = BipartiteGraph(A)
     return RowColoringResult(A, bg, color)
 end
diff --git a/test/structured.jl b/test/structured.jl
index 5ee336b1..e20adaae 100644
--- a/test/structured.jl
+++ b/test/structured.jl
@@ -1,5 +1,6 @@
 using LinearAlgebra
 using SparseMatrixColorings
+using SparseMatrixColorings: cycle_range
 using Test
 
 column_problem = ColoringProblem(; structure=:nonsymmetric, partition=:column)
@@ -7,6 +8,17 @@ row_problem = ColoringProblem(; structure=:nonsymmetric, partition=:row)
 
 algo = GreedyColoringAlgorithm()
 
+@testset "Utils" begin
+    @test cycle_range(2, 3) == [1, 2, 1]
+    @test cycle_range(2, 4) == [1, 2, 1, 2]
+    @test cycle_range(2, 5) == [1, 2, 1, 2, 1]
+    @test cycle_range(3, 5) == [1, 2, 3, 1, 2]
+    @test cycle_range(3, 6) == [1, 2, 3, 1, 2, 3]
+    @test cycle_range(2, 1) == [1]
+    @test cycle_range(3, 1) == [1]
+    @test cycle_range(3, 2) == [1, 2]
+end
+
 @testset "Diagonal" begin
     for n in (1, 2, 10, 100, 1000)
         A = Diagonal(rand(n))

From 45b16a39b31dbaf4360ff9a7c7982b13b183ba68 Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Wed, 2 Oct 2024 12:31:22 +0200
Subject: [PATCH 08/33] More tests

---
 src/structured.jl  |  8 ++++----
 test/structured.jl | 25 ++++++++++++++++++-------
 2 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/src/structured.jl b/src/structured.jl
index 9320e438..99fdea3f 100644
--- a/src/structured.jl
+++ b/src/structured.jl
@@ -70,7 +70,7 @@ function coloring(
     algo::GreedyColoringAlgorithm;
     kwargs...,
 )
-    color = cycle_range(1:2, size(A, 2))
+    color = cycle_range(2, size(A, 2))
     bg = BipartiteGraph(A)
     return ColumnColoringResult(A, bg, color)
 end
@@ -81,7 +81,7 @@ function coloring(
     algo::GreedyColoringAlgorithm;
     kwargs...,
 )
-    color = cycle_range(1:2, size(A, 1))
+    color = cycle_range(2, size(A, 1))
     bg = BipartiteGraph(A)
     return RowColoringResult(A, bg, color)
 end
@@ -126,7 +126,7 @@ function coloring(
     algo::GreedyColoringAlgorithm;
     kwargs...,
 )
-    color = cycle_range(1:3, size(A, 2))
+    color = cycle_range(3, size(A, 2))
     bg = BipartiteGraph(A)
     return ColumnColoringResult(A, bg, color)
 end
@@ -137,7 +137,7 @@ function coloring(
     algo::GreedyColoringAlgorithm;
     kwargs...,
 )
-    color = cycle_range(1:3, size(A, 1))
+    color = cycle_range(3, size(A, 1))
     bg = BipartiteGraph(A)
     return RowColoringResult(A, bg, color)
 end
diff --git a/test/structured.jl b/test/structured.jl
index e20adaae..7e699a3e 100644
--- a/test/structured.jl
+++ b/test/structured.jl
@@ -26,12 +26,16 @@ end
         result = coloring(A, column_problem, algo)
         B = compress(A, result)
         @test size(B, 2) == 1
-        @test decompress(B, result) == A
+        D = decompress(B, result)
+        @test D == A
+        @test D isa Diagonal
         # row
         result = coloring(A, row_problem, algo)
         B = compress(A, result)
         @test size(B, 1) == 1
-        @test decompress(B, result) == A
+        D = decompress(B, result)
+        @test D == A
+        @test D isa Diagonal
     end
 end
 
@@ -44,12 +48,16 @@ end
             result = coloring(A, column_problem, algo)
             B = compress(A, result)
             @test size(B, 2) == 2
-            @test decompress(B, result) == A
+            D = decompress(B, result)
+            @test D == A
+            @test D isa Bidiagonal
             # row
             result = coloring(A, row_problem, algo)
             B = compress(A, result)
             @test size(B, 1) == 2
-            @test decompress(B, result) == A
+            D = decompress(B, result)
+            @test D == A
+            @test D isa Bidiagonal
         end
     end
 end
@@ -63,12 +71,15 @@ end
             result = coloring(A, column_problem, algo)
             B = compress(A, result)
             @test size(B, 2) == min(n, 3)
-            @test decompress(B, result) == A
-            # row
+            D = decompress(B, result)
+            @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)
-            @test decompress(B, result) == A
+            D = decompress(B, result)
+            @test D == A
+            @test D isa Tridiagonal
         end
     end
 end

From 361e53eaae68bb95f60a8ff0f1a141d21909cbd7 Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Wed, 2 Oct 2024 13:37:08 +0200
Subject: [PATCH 09/33] Add BandedMatrices

---
 .github/workflows/Test.yml                    |  2 +-
 Project.toml                                  | 11 ++-
 ext/SparseMatrixColoringsBandedMatricesExt.jl | 87 +++++++++++++++++++
 src/SparseMatrixColorings.jl                  | 11 +++
 src/structured.jl                             | 26 ++----
 test/Project.toml                             |  1 +
 test/structured.jl                            | 31 +++++--
 7 files changed, 141 insertions(+), 28 deletions(-)
 create mode 100644 ext/SparseMatrixColoringsBandedMatricesExt.jl

diff --git a/.github/workflows/Test.yml b/.github/workflows/Test.yml
index e9dc089e..f81a67d9 100644
--- a/.github/workflows/Test.yml
+++ b/.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
diff --git a/Project.toml b/Project.toml
index 9e0676a0..17f2b999 100644
--- a/Project.toml
+++ b/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"
diff --git a/ext/SparseMatrixColoringsBandedMatricesExt.jl b/ext/SparseMatrixColoringsBandedMatricesExt.jl
new file mode 100644
index 00000000..f827f32a
--- /dev/null
+++ b/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
diff --git a/src/SparseMatrixColorings.jl b/src/SparseMatrixColorings.jl
index d712ea31..7a9250e8 100644
--- a/src/SparseMatrixColorings.jl
+++ b/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
diff --git a/src/structured.jl b/src/structured.jl
index 99fdea3f..3acfb8b3 100644
--- a/src/structured.jl
+++ b/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,9 +40,7 @@ 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]]
@@ -52,9 +48,7 @@ function decompress!(
     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]
diff --git a/test/Project.toml b/test/Project.toml
index 820476d6..178f1d93 100644
--- a/test/Project.toml
+++ b/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"
diff --git a/test/structured.jl b/test/structured.jl
index 7e699a3e..6b976029 100644
--- a/test/structured.jl
+++ b/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

From 4df06f31efce9a83293c2f6c7e9da873f96e4342 Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Wed, 2 Oct 2024 13:37:56 +0200
Subject: [PATCH 10/33] Reactivate lts

---
 .github/workflows/Test.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/Test.yml b/.github/workflows/Test.yml
index f81a67d9..e9dc089e 100644
--- a/.github/workflows/Test.yml
+++ b/.github/workflows/Test.yml
@@ -20,7 +20,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        julia-version: ['1', 'pre']  # TODO: do we drop 'lts'?
+        julia-version: ['lts', '1', 'pre']
 
     steps:
       - uses: actions/checkout@v4

From 2d71688f4c7d39d1aedb18811e0ec1fe8d694620 Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Wed, 2 Oct 2024 13:44:20 +0200
Subject: [PATCH 11/33] Extras

---
 Project.toml | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/Project.toml b/Project.toml
index 17f2b999..7f533928 100644
--- a/Project.toml
+++ b/Project.toml
@@ -19,6 +19,9 @@ BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
 [extensions]
 SparseMatrixColoringsBandedMatricesExt = "BandedMatrices"
 
+[extras]
+BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
+
 [compat]
 ADTypes = "1.2.1"
 BandedMatrices = "1.7.5"

From eefd1eebeb0d6ee22749ed0b43e816ace76441e6 Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Wed, 2 Oct 2024 13:48:11 +0200
Subject: [PATCH 12/33] Fix LTS

---
 src/SparseMatrixColorings.jl | 5 +++--
 test/Project.toml            | 1 +
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/SparseMatrixColorings.jl b/src/SparseMatrixColorings.jl
index 7a9250e8..6c636bc2 100644
--- a/src/SparseMatrixColorings.jl
+++ b/src/SparseMatrixColorings.jl
@@ -71,8 +71,9 @@ end
 
 @static if !isdefined(Base, :get_extension)
     function __init__()
-        @require BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
-        return include("../ext/SparseMatrixColoringsBandedMatricesExt.jl")
+        @require BandedMatrices = "aae01518-5342-5314-be14-df237901396f" include(
+            "../ext/SparseMatrixColoringsBandedMatricesExt.jl"
+        )
     end
 end
 
diff --git a/test/Project.toml b/test/Project.toml
index 178f1d93..20509c33 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -2,6 +2,7 @@
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
+BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0"
 CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
 Chairmarks = "0ca39b1e-fe0b-4e98-acfc-b1656634c4de"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"

From e088877c81e107e7671f6140bb24c1c80c8e61aa Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Wed, 2 Oct 2024 13:57:21 +0200
Subject: [PATCH 13/33] Ignore unloaded Requires

---
 test/runtests.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/runtests.jl b/test/runtests.jl
index ca0a09af..7c04c6c1 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -11,7 +11,7 @@ include("utils.jl")
     @testset verbose = true "Code quality" begin
         if VERSION >= v"1.10"
             @testset "Aqua" begin
-                Aqua.test_all(SparseMatrixColorings)
+                Aqua.test_all(SparseMatrixColorings; stale_deps=(; ignore=[:Requires],))
             end
             @testset "JET" begin
                 JET.test_package(SparseMatrixColorings; target_defined_modules=true)

From de4a21aae710871df1fb9d93a2c9c7331d317b49 Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Wed, 2 Oct 2024 16:12:13 +0200
Subject: [PATCH 14/33] Add BlockBandedMatrices

---
 Project.toml                                  | 13 ++-
 ext/SparseMatrixColoringsBandedMatricesExt.jl | 20 ++--
 ...seMatrixColoringsBlockBandedMatricesExt.jl | 88 +++++++++++++++++
 src/check.jl                                  |  2 +-
 src/structured.jl                             |  2 +-
 test/Project.toml                             |  2 +
 test/structured.jl                            | 99 +++++--------------
 test/utils.jl                                 | 38 ++++++-
 8 files changed, 172 insertions(+), 92 deletions(-)
 create mode 100644 ext/SparseMatrixColoringsBlockBandedMatricesExt.jl

diff --git a/Project.toml b/Project.toml
index 7f533928..f76324b3 100644
--- a/Project.toml
+++ b/Project.toml
@@ -15,16 +15,18 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [weakdeps]
 BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
+BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
+BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0"
 
 [extensions]
 SparseMatrixColoringsBandedMatricesExt = "BandedMatrices"
-
-[extras]
-BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
+SparseMatrixColoringsBlockBandedMatricesExt = ["BlockArrays", "BlockBandedMatrices"]
 
 [compat]
 ADTypes = "1.2.1"
 BandedMatrices = "1.7.5"
+BlockArrays = "1.1.1"
+BlockBandedMatrices = "0.13.1"
 Compat = "3.46,4.2"
 DataStructures = "0.18"
 DocStringExtensions = "0.8,0.9"
@@ -33,3 +35,8 @@ Random = "<0.0.1, 1"
 Requires = "1.3.0"
 SparseArrays = "<0.0.1, 1"
 julia = "1.6"
+
+[extras]
+BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
+BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
+BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0"
\ No newline at end of file
diff --git a/ext/SparseMatrixColoringsBandedMatricesExt.jl b/ext/SparseMatrixColoringsBandedMatricesExt.jl
index f827f32a..85ee4b88 100644
--- a/ext/SparseMatrixColoringsBandedMatricesExt.jl
+++ b/ext/SparseMatrixColoringsBandedMatricesExt.jl
@@ -1,7 +1,7 @@
 module SparseMatrixColoringsBandedMatricesExt
 
 if isdefined(Base, :get_extension)
-    using BandedMatrices: BandedMatrix, bandwidths, bandrange
+    using BandedMatrices: BandedMatrix, bandrange, bandwidths, colrange, rowrange
     using SparseMatrixColorings:
         BipartiteGraph,
         ColoringProblem,
@@ -13,7 +13,7 @@ if isdefined(Base, :get_extension)
         row_colors
     import SparseMatrixColorings as SMC
 else
-    using ..BandedMatrices: BandedMatrix, bandwidths, bandrange
+    using ..BandedMatrices: BandedMatrix, bandrange, bandwidths, colrange, rowrange
     using ..SparseMatrixColorings:
         BipartiteGraph,
         ColoringProblem,
@@ -27,7 +27,7 @@ else
 end
 
 #=
-This code is partially taken from ArrayInterface.jl and FiniteDiff.jl
+This code is partly taken from ArrayInterface.jl and FiniteDiff.jl
 https://github.com/JuliaArrays/ArrayInterface.jl
 https://github.com/JuliaDiff/FiniteDiff.jl
 =#
@@ -38,8 +38,7 @@ function SMC.coloring(
     algo::GreedyColoringAlgorithm;
     kwargs...,
 )
-    l, u = bandwidths(A)
-    width = u + l + 1
+    width = length(bandrange(A))
     color = cycle_range(width, size(A, 2))
     bg = BipartiteGraph(A)
     return ColumnColoringResult(A, bg, color)
@@ -51,8 +50,7 @@ function SMC.coloring(
     algo::GreedyColoringAlgorithm;
     kwargs...,
 )
-    l, u = bandwidths(A)
-    width = u + l + 1
+    width = length(bandrange(A))
     color = cycle_range(width, size(A, 1))
     bg = BipartiteGraph(A)
     return RowColoringResult(A, bg, color)
@@ -60,11 +58,9 @@ 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)
+        for i in colrange(A, j)
             A[i, j] = B[i, c]
         end
     end
@@ -73,11 +69,9 @@ 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)
+        for j in rowrange(A, i)
             A[i, j] = B[c, j]
         end
     end
diff --git a/ext/SparseMatrixColoringsBlockBandedMatricesExt.jl b/ext/SparseMatrixColoringsBlockBandedMatricesExt.jl
new file mode 100644
index 00000000..f1fe46c7
--- /dev/null
+++ b/ext/SparseMatrixColoringsBlockBandedMatricesExt.jl
@@ -0,0 +1,88 @@
+module SparseMatrixColoringsBlockBandedMatricesExt
+
+if isdefined(Base, :get_extension)
+    using BlockArrays: blockaxes, blockfirsts, blocklasts, blocksize, blocklengths
+    using BlockBandedMatrices:
+        BlockBandedMatrix,
+        blockbandrange,
+        blockbandwidths,
+        blocklengths,
+        blocksize,
+        subblockbandwidths
+    using SparseMatrixColorings:
+        BipartiteGraph,
+        ColoringProblem,
+        ColumnColoringResult,
+        GreedyColoringAlgorithm,
+        RowColoringResult,
+        column_colors,
+        cycle_range,
+        row_colors
+    import SparseMatrixColorings as SMC
+else
+    using ..BlockArrays: blockaxes, blockfirsts, blocklasts, blocksize, blocklengths
+    using ..BlockBandedMatrices:
+        BlockBandedMatrix,
+        blockbandrange,
+        blockbandwidths,
+        blocklengths,
+        blocksize,
+        subblockbandwidths
+    using ..SparseMatrixColorings:
+        BipartiteGraph,
+        ColoringProblem,
+        ColumnColoringResult,
+        GreedyColoringAlgorithm,
+        RowColoringResult,
+        column_colors,
+        cycle_range,
+        row_colors
+    import ..SparseMatrixColorings as SMC
+end
+
+#=
+This code is partly taken from ArrayInterface.jl and FiniteDiff.jl
+https://github.com/JuliaArrays/ArrayInterface.jl
+https://github.com/JuliaDiff/FiniteDiff.jl
+=#
+
+function SMC.coloring(
+    A::BlockBandedMatrix,
+    ::ColoringProblem{:nonsymmetric,:column},
+    algo::GreedyColoringAlgorithm;
+    kwargs...,
+)
+    # consider blocks of columns
+    nb_blocks = blocksize(A, 2)
+    nb_cols_in_block = blocklengths(axes(A, 2))
+    first_col_in_block = blockfirsts(axes(A, 2))
+    last_col_in_block = blocklasts(axes(A, 2))
+
+    # give a macroscopic color to each block, so that 2 blocks of columns with the same macro color do not intersect
+    # same idea as for BandedMatrices
+    nb_macrocolors = length(blockbandrange(A))
+    macrocolor = cycle_range(nb_macrocolors, nb_blocks)
+
+    # for each macroscopic color, count how many microscopic colors will be needed
+    # columns within a block are colored naively with all distinct micro colors
+    nb_colors_in_macrocolor = [
+        maximum(nb_cols_in_block[mc:nb_macrocolors:nb_blocks]; init=0) for
+        mc in 1:nb_macrocolors
+    ]
+    color_shift_in_macrocolor = vcat(0, cumsum(nb_colors_in_macrocolor)[1:(end - 1)])
+
+    # assign a microscopic color to each column as a function of its macroscopic color and its position within the block
+    color = Vector{Int}(undef, size(A, 2))
+    for b in 1:nb_blocks
+        mc = macrocolor[b]
+        shift = color_shift_in_macrocolor[mc]
+        for j in first_col_in_block[b]:last_col_in_block[b]
+            color[j] = shift + (j - first_col_in_block[b] + 1)
+        end
+    end
+
+    bg = BipartiteGraph(A)
+    return ColumnColoringResult(A, bg, color)
+end
+
+end
diff --git a/src/check.jl b/src/check.jl
index 1d7ce3c4..25d027f5 100644
--- a/src/check.jl
+++ b/src/check.jl
@@ -36,7 +36,7 @@ function structurally_orthogonal_columns(
     group = group_by_color(color)
     for (c, g) in enumerate(group)
         Ag = view(A, :, g)
-        nonzeros_per_row = dropdims(count(!iszero, Ag; dims=2); dims=2)
+        nonzeros_per_row = only(eachcol(count(!iszero, Ag; dims=2)))
         max_nonzeros_per_row, i = findmax(nonzeros_per_row)
         if max_nonzeros_per_row > 1
             if verbose
diff --git a/src/structured.jl b/src/structured.jl
index 3acfb8b3..78b2db69 100644
--- a/src/structured.jl
+++ b/src/structured.jl
@@ -1,5 +1,5 @@
 #=
-This code is partially taken from ArrayInterface.jl
+This code is partly taken from ArrayInterface.jl
 https://github.com/JuliaArrays/ArrayInterface.jl
 =#
 
diff --git a/test/Project.toml b/test/Project.toml
index 20509c33..f2cec836 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -1,7 +1,9 @@
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
+ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
+BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
 BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0"
 CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
 Chairmarks = "0ca39b1e-fe0b-4e98-acfc-b1656634c4de"
diff --git a/test/structured.jl b/test/structured.jl
index 6b976029..5e6b9db1 100644
--- a/test/structured.jl
+++ b/test/structured.jl
@@ -1,14 +1,11 @@
+using ArrayInterface: ArrayInterface
 using BandedMatrices: BandedMatrix, brand
+using BlockBandedMatrices: BlockBandedMatrix
 using LinearAlgebra
 using SparseMatrixColorings
 using SparseMatrixColorings: cycle_range
 using Test
 
-column_problem = ColoringProblem(; structure=:nonsymmetric, partition=:column)
-row_problem = ColoringProblem(; structure=:nonsymmetric, partition=:row)
-
-algo = GreedyColoringAlgorithm()
-
 @testset "Utils" begin
     @test cycle_range(2, 3) == [1, 2, 1]
     @test cycle_range(2, 4) == [1, 2, 1, 2]
@@ -18,87 +15,43 @@ algo = GreedyColoringAlgorithm()
     @test cycle_range(2, 1) == [1]
     @test cycle_range(3, 1) == [1]
     @test cycle_range(3, 2) == [1, 2]
-end
+end;
 
 @testset "Diagonal" begin
-    for n in (1, 2, 10, 100, 1000)
+    @testset for n in (1, 2, 10, 100)
         A = Diagonal(rand(n))
-        # column
-        result = coloring(A, column_problem, algo)
-        B = compress(A, result)
-        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)
-        D = decompress(B, result)
-        @test size(B, 1) == 1
-        @test D == A
-        @test D isa Diagonal
+        test_structured_coloring_decompression(A)
     end
-end
+end;
 
 @testset "Bidiagonal" begin
-    for n in (2, 10, 100, 1000)
+    @testset for n in (2, 10, 100)
         A1 = Bidiagonal(rand(n), rand(n - 1), :U)
         A2 = Bidiagonal(rand(n), rand(n - 1), :L)
-        for A in (A1, A2)
-            # column
-            result = coloring(A, column_problem, algo)
-            B = compress(A, result)
-            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)
-            D = decompress(B, result)
-            @test size(B, 1) == 2
-            @test D == A
-            @test D isa Bidiagonal
-        end
+        test_structured_coloring_decompression(A1)
+        test_structured_coloring_decompression(A2)
     end
-end
+end;
 
 @testset "Tridiagonal" begin
-    for n in (2, 10, 100, 1000)
-        A1 = Tridiagonal(rand(n - 1), rand(n), rand(n - 1))
-        A2 = Tridiagonal(rand(n - 1), rand(n), rand(n - 1))
-        for A in (A1, A2)
-            # column
-            result = coloring(A, column_problem, algo)
-            B = compress(A, result)
-            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)
-            D = decompress(B, result)
-            @test size(B, 1) == min(n, 3)
-            @test D == A
-            @test D isa Tridiagonal
-        end
+    for n in (2, 10, 100)
+        A = Tridiagonal(rand(n - 1), rand(n), rand(n - 1))
+        test_structured_coloring_decompression(A)
     end
-end
+end;
 
 @testset "BandedMatrices" begin
-    for (m, n) in [(10, 20), (20, 10)], l in 0:5, u in 0:5
+    @testset 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
+        test_structured_coloring_decompression(A)
+    end
+end;
+
+@testset "BlockBandedMatrices" begin
+    @testset for (mb, nb) in [(10, 20), (20, 10)], lb in 0:3, ub in 0:3, bs in 1:3
+        rows = rand(1:bs, mb)
+        cols = rand(1:bs, nb)
+        A = BlockBandedMatrix{Float64}(rand(sum(rows), sum(cols)), rows, cols, (l, u))
+        test_structured_coloring_decompression(A)
     end
-end
+end;
diff --git a/test/utils.jl b/test/utils.jl
index 2ebad2a4..d19133f6 100644
--- a/test/utils.jl
+++ b/test/utils.jl
@@ -1,7 +1,14 @@
+using ArrayInterface: ArrayInterface
+using BandedMatrices: BandedMatrix
+using BlockBandedMatrices: BlockBandedMatrix
 using LinearAlgebra
 using SparseMatrixColorings
 using SparseMatrixColorings:
-    AdjacencyGraph, LinearSystemColoringResult, matrix_versions, respectful_similar
+    AdjacencyGraph,
+    LinearSystemColoringResult,
+    matrix_versions,
+    respectful_similar,
+    structurally_orthogonal_columns
 using Test
 
 function test_coloring_decompression(
@@ -112,3 +119,32 @@ function test_coloring_decompression(
         @test all(color_vec .== Ref(color_vec[1]))
     end
 end
+
+OptimalColoringKnown = Union{Diagonal,Bidiagonal,Tridiagonal,BandedMatrix,BlockBandedMatrix}
+
+function test_structured_coloring_decompression(A::AbstractMatrix)
+    column_problem = ColoringProblem(; structure=:nonsymmetric, partition=:column)
+    row_problem = ColoringProblem(; structure=:nonsymmetric, partition=:row)
+    algo = GreedyColoringAlgorithm()
+
+    @testset "Column" begin
+        result = coloring(A, column_problem, algo)
+        color = column_colors(result)
+        B = compress(A, result)
+        D = decompress(B, result)
+        @test D == A
+        @test D isa typeof(A)
+        @test structurally_orthogonal_columns(A, color)
+        if A isa OptimalColoringKnown
+            @test color == ArrayInterface.matrix_colors(A)
+        end
+    end
+
+    @testset "Row" begin
+        result = coloring(A, row_problem, algo)
+        B = compress(A, result)
+        D = decompress(B, result)
+        @test D == A
+        @test D isa typeof(A)
+    end
+end

From c743b22cd75554861fd3359006e0fe5ce391cc9e Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Wed, 2 Oct 2024 16:19:20 +0200
Subject: [PATCH 15/33] Row coloring

---
 ...seMatrixColoringsBlockBandedMatricesExt.jl | 51 ++++++++++++-------
 1 file changed, 33 insertions(+), 18 deletions(-)

diff --git a/ext/SparseMatrixColoringsBlockBandedMatricesExt.jl b/ext/SparseMatrixColoringsBlockBandedMatricesExt.jl
index f1fe46c7..23eab3ad 100644
--- a/ext/SparseMatrixColoringsBlockBandedMatricesExt.jl
+++ b/ext/SparseMatrixColoringsBlockBandedMatricesExt.jl
@@ -46,43 +46,58 @@ https://github.com/JuliaArrays/ArrayInterface.jl
 https://github.com/JuliaDiff/FiniteDiff.jl
 =#
 
-function SMC.coloring(
-    A::BlockBandedMatrix,
-    ::ColoringProblem{:nonsymmetric,:column},
-    algo::GreedyColoringAlgorithm;
-    kwargs...,
-)
-    # consider blocks of columns
-    nb_blocks = blocksize(A, 2)
-    nb_cols_in_block = blocklengths(axes(A, 2))
-    first_col_in_block = blockfirsts(axes(A, 2))
-    last_col_in_block = blocklasts(axes(A, 2))
+function blockbanded_coloring(A::BlockBandedMatrix, dim::Integer)
+    # consider blocks of columns or rows (let's call them vertices) depending on `dim`
+    nb_blocks = blocksize(A, dim)
+    nb_in_block = blocklengths(axes(A, dim))
+    first_in_block = blockfirsts(axes(A, dim))
+    last_in_block = blocklasts(axes(A, dim))
+    color = zeros(Int, size(A, dim))
 
-    # give a macroscopic color to each block, so that 2 blocks of columns with the same macro color do not intersect
+    # give a macroscopic color to each block, so that 2 blocks with the same macro color are orthogonal
     # same idea as for BandedMatrices
     nb_macrocolors = length(blockbandrange(A))
     macrocolor = cycle_range(nb_macrocolors, nb_blocks)
 
     # for each macroscopic color, count how many microscopic colors will be needed
-    # columns within a block are colored naively with all distinct micro colors
+    # vertices within a block are colored naively with all distinct micro colors
     nb_colors_in_macrocolor = [
-        maximum(nb_cols_in_block[mc:nb_macrocolors:nb_blocks]; init=0) for
-        mc in 1:nb_macrocolors
+        maximum(nb_in_block[mc:nb_macrocolors:nb_blocks]; init=0) for mc in 1:nb_macrocolors
     ]
     color_shift_in_macrocolor = vcat(0, cumsum(nb_colors_in_macrocolor)[1:(end - 1)])
 
     # assign a microscopic color to each column as a function of its macroscopic color and its position within the block
-    color = Vector{Int}(undef, size(A, 2))
     for b in 1:nb_blocks
         mc = macrocolor[b]
         shift = color_shift_in_macrocolor[mc]
-        for j in first_col_in_block[b]:last_col_in_block[b]
-            color[j] = shift + (j - first_col_in_block[b] + 1)
+        for k in first_in_block[b]:last_in_block[b]
+            color[k] = shift + (k - first_in_block[b] + 1)
         end
     end
 
+    return color
+end
+
+function SMC.coloring(
+    A::BlockBandedMatrix,
+    ::ColoringProblem{:nonsymmetric,:column},
+    algo::GreedyColoringAlgorithm;
+    kwargs...,
+)
+    color = blockbanded_coloring(A, 2)
     bg = BipartiteGraph(A)
     return ColumnColoringResult(A, bg, color)
 end
 
+function SMC.coloring(
+    A::BlockBandedMatrix,
+    ::ColoringProblem{:nonsymmetric,:row},
+    algo::GreedyColoringAlgorithm;
+    kwargs...,
+)
+    color = blockbanded_coloring(A, 1)
+    bg = BipartiteGraph(A)
+    return RowColoringResult(A, bg, color)
+end
+
 end

From 4f4075820af85815e6601fbd19e3354f03c608e2 Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Wed, 2 Oct 2024 16:25:54 +0200
Subject: [PATCH 16/33] Fix LTS

---
 ext/SparseMatrixColoringsBlockBandedMatricesExt.jl | 4 ++++
 src/SparseMatrixColorings.jl                       | 3 +++
 2 files changed, 7 insertions(+)

diff --git a/ext/SparseMatrixColoringsBlockBandedMatricesExt.jl b/ext/SparseMatrixColoringsBlockBandedMatricesExt.jl
index 23eab3ad..7e30de66 100644
--- a/ext/SparseMatrixColoringsBlockBandedMatricesExt.jl
+++ b/ext/SparseMatrixColoringsBlockBandedMatricesExt.jl
@@ -46,6 +46,8 @@ https://github.com/JuliaArrays/ArrayInterface.jl
 https://github.com/JuliaDiff/FiniteDiff.jl
 =#
 
+## BlockBandedMatrix
+
 function blockbanded_coloring(A::BlockBandedMatrix, dim::Integer)
     # consider blocks of columns or rows (let's call them vertices) depending on `dim`
     nb_blocks = blocksize(A, dim)
@@ -100,4 +102,6 @@ function SMC.coloring(
     return RowColoringResult(A, bg, color)
 end
 
+
+
 end
diff --git a/src/SparseMatrixColorings.jl b/src/SparseMatrixColorings.jl
index 6c636bc2..44eab3ad 100644
--- a/src/SparseMatrixColorings.jl
+++ b/src/SparseMatrixColorings.jl
@@ -74,6 +74,9 @@ end
         @require BandedMatrices = "aae01518-5342-5314-be14-df237901396f" include(
             "../ext/SparseMatrixColoringsBandedMatricesExt.jl"
         )
+        @require BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0" include(
+            "../ext/SparseMatrixColoringsBlockBandedMatricesExt.jl"
+        )
     end
 end
 

From 116691dfe60d142c63baa66d12e01798352a6958 Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Wed, 2 Oct 2024 16:27:20 +0200
Subject: [PATCH 17/33] No fail fast

---
 .github/workflows/Test.yml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.github/workflows/Test.yml b/.github/workflows/Test.yml
index e9dc089e..a2841d90 100644
--- a/.github/workflows/Test.yml
+++ b/.github/workflows/Test.yml
@@ -19,6 +19,7 @@ jobs:
   test:
     runs-on: ubuntu-latest
     strategy:
+      fail-fast: false
       matrix:
         julia-version: ['lts', '1', 'pre']
 

From 66acbca48b501007028c8b20279f7cc180b33bd7 Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Wed, 2 Oct 2024 17:32:05 +0200
Subject: [PATCH 18/33] BandedBlockBandedMatrices

---
 ...seMatrixColoringsBlockBandedMatricesExt.jl | 42 ++++++++++++-------
 test/structured.jl                            | 23 +++++++---
 test/utils.jl                                 | 36 +++++++---------
 3 files changed, 61 insertions(+), 40 deletions(-)

diff --git a/ext/SparseMatrixColoringsBlockBandedMatricesExt.jl b/ext/SparseMatrixColoringsBlockBandedMatricesExt.jl
index 7e30de66..2274b8f6 100644
--- a/ext/SparseMatrixColoringsBlockBandedMatricesExt.jl
+++ b/ext/SparseMatrixColoringsBlockBandedMatricesExt.jl
@@ -3,6 +3,7 @@ module SparseMatrixColoringsBlockBandedMatricesExt
 if isdefined(Base, :get_extension)
     using BlockArrays: blockaxes, blockfirsts, blocklasts, blocksize, blocklengths
     using BlockBandedMatrices:
+        BandedBlockBandedMatrix,
         BlockBandedMatrix,
         blockbandrange,
         blockbandwidths,
@@ -22,6 +23,7 @@ if isdefined(Base, :get_extension)
 else
     using ..BlockArrays: blockaxes, blockfirsts, blocklasts, blocksize, blocklengths
     using ..BlockBandedMatrices:
+        BandedBlockBandedMatrix,
         BlockBandedMatrix,
         blockbandrange,
         blockbandwidths,
@@ -46,9 +48,14 @@ https://github.com/JuliaArrays/ArrayInterface.jl
 https://github.com/JuliaDiff/FiniteDiff.jl
 =#
 
-## BlockBandedMatrix
+function subblockbandrange(A::BandedBlockBandedMatrix)
+    u, l = subblockbandwidths(A)
+    return (-l):u
+end
 
-function blockbanded_coloring(A::BlockBandedMatrix, dim::Integer)
+function blockbanded_coloring(
+    A::Union{BlockBandedMatrix,BandedBlockBandedMatrix}, dim::Integer
+)
     # consider blocks of columns or rows (let's call them vertices) depending on `dim`
     nb_blocks = blocksize(A, dim)
     nb_in_block = blocklengths(axes(A, dim))
@@ -61,27 +68,34 @@ function blockbanded_coloring(A::BlockBandedMatrix, dim::Integer)
     nb_macrocolors = length(blockbandrange(A))
     macrocolor = cycle_range(nb_macrocolors, nb_blocks)
 
+    width = if A isa BandedBlockBandedMatrix
+        # vertices within a block are colored cleverly using bands
+        length(subblockbandrange(A))
+    else
+        # vertices within a block are colored naively with distinct micro colors (~ infinite band width)
+        minimum(size(A))
+    end
+
     # for each macroscopic color, count how many microscopic colors will be needed
-    # vertices within a block are colored naively with all distinct micro colors
-    nb_colors_in_macrocolor = [
-        maximum(nb_in_block[mc:nb_macrocolors:nb_blocks]; init=0) for mc in 1:nb_macrocolors
-    ]
+    nb_colors_in_macrocolor = zeros(Int, nb_macrocolors)
+    for mc in 1:nb_macrocolors
+        largest_nb_in_macrocolor = maximum(nb_in_block[mc:nb_macrocolors:nb_blocks]; init=0)
+        nb_colors_in_macrocolor[mc] = min(width, largest_nb_in_macrocolor)
+    end
     color_shift_in_macrocolor = vcat(0, cumsum(nb_colors_in_macrocolor)[1:(end - 1)])
 
     # assign a microscopic color to each column as a function of its macroscopic color and its position within the block
     for b in 1:nb_blocks
-        mc = macrocolor[b]
-        shift = color_shift_in_macrocolor[mc]
-        for k in first_in_block[b]:last_in_block[b]
-            color[k] = shift + (k - first_in_block[b] + 1)
-        end
+        block_color = cycle_range(width, nb_in_block[b])
+        shift = color_shift_in_macrocolor[macrocolor[b]]
+        color[first_in_block[b]:last_in_block[b]] .= shift .+ block_color
     end
 
     return color
 end
 
 function SMC.coloring(
-    A::BlockBandedMatrix,
+    A::Union{BlockBandedMatrix,BandedBlockBandedMatrix},
     ::ColoringProblem{:nonsymmetric,:column},
     algo::GreedyColoringAlgorithm;
     kwargs...,
@@ -92,7 +106,7 @@ function SMC.coloring(
 end
 
 function SMC.coloring(
-    A::BlockBandedMatrix,
+    A::Union{BlockBandedMatrix,BandedBlockBandedMatrix},
     ::ColoringProblem{:nonsymmetric,:row},
     algo::GreedyColoringAlgorithm;
     kwargs...,
@@ -102,6 +116,4 @@ function SMC.coloring(
     return RowColoringResult(A, bg, color)
 end
 
-
-
 end
diff --git a/test/structured.jl b/test/structured.jl
index 5e6b9db1..a88e565c 100644
--- a/test/structured.jl
+++ b/test/structured.jl
@@ -18,14 +18,14 @@ using Test
 end;
 
 @testset "Diagonal" begin
-    @testset for n in (1, 2, 10, 100)
+    for n in (1, 2, 10, 100)
         A = Diagonal(rand(n))
         test_structured_coloring_decompression(A)
     end
 end;
 
 @testset "Bidiagonal" begin
-    @testset for n in (2, 10, 100)
+    for n in (2, 10, 100)
         A1 = Bidiagonal(rand(n), rand(n - 1), :U)
         A2 = Bidiagonal(rand(n), rand(n - 1), :L)
         test_structured_coloring_decompression(A1)
@@ -48,10 +48,23 @@ end;
 end;
 
 @testset "BlockBandedMatrices" begin
-    @testset for (mb, nb) in [(10, 20), (20, 10)], lb in 0:3, ub in 0:3, bs in 1:3
-        rows = rand(1:bs, mb)
-        cols = rand(1:bs, nb)
+    for (mb, nb) in [(10, 20), (20, 10)], lb in 0:3, ub in 0:3, _ in 1:10
+        rows = rand(1:5, mb)
+        cols = rand(1:5, nb)
         A = BlockBandedMatrix{Float64}(rand(sum(rows), sum(cols)), rows, cols, (l, u))
         test_structured_coloring_decompression(A)
     end
 end;
+
+@testset "BandedBlockBandedMatrices" begin
+    for (mb, nb) in [(10, 20), (20, 10)], lb in 0:3, ub in 0:3, _ in 1:10
+        rows = rand(5:10, mb)
+        cols = rand(5:10, nb)
+        λ = rand(0:5)
+        μ = rand(0:5)
+        A = BandedBlockBandedMatrix{Float64}(
+            rand(sum(rows), sum(cols)), rows, cols, (lb, ub), (λ, μ)
+        )
+        test_structured_coloring_decompression(A)
+    end
+end;
diff --git a/test/utils.jl b/test/utils.jl
index d19133f6..3f2a78c3 100644
--- a/test/utils.jl
+++ b/test/utils.jl
@@ -127,24 +127,20 @@ function test_structured_coloring_decompression(A::AbstractMatrix)
     row_problem = ColoringProblem(; structure=:nonsymmetric, partition=:row)
     algo = GreedyColoringAlgorithm()
 
-    @testset "Column" begin
-        result = coloring(A, column_problem, algo)
-        color = column_colors(result)
-        B = compress(A, result)
-        D = decompress(B, result)
-        @test D == A
-        @test D isa typeof(A)
-        @test structurally_orthogonal_columns(A, color)
-        if A isa OptimalColoringKnown
-            @test color == ArrayInterface.matrix_colors(A)
-        end
-    end
-
-    @testset "Row" begin
-        result = coloring(A, row_problem, algo)
-        B = compress(A, result)
-        D = decompress(B, result)
-        @test D == A
-        @test D isa typeof(A)
-    end
+    # Column
+    result = coloring(A, column_problem, algo)
+    color = column_colors(result)
+    B = compress(A, result)
+    D = decompress(B, result)
+    @test D == A
+    @test nameof(typeof(D)) == nameof(typeof(A))
+    @test structurally_orthogonal_columns(A, color)
+    @test color == ArrayInterface.matrix_colors(A)
+
+    # Row
+    result = coloring(A, row_problem, algo)
+    B = compress(A, result)
+    D = decompress(B, result)
+    @test D == A
+    @test nameof(typeof(D)) == nameof(typeof(A))
 end

From 4c838ff6f01d3db56c20de1984d89f9f9766e057 Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Wed, 2 Oct 2024 17:37:24 +0200
Subject: [PATCH 19/33] Infinite width

---
 ext/SparseMatrixColoringsBlockBandedMatricesExt.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ext/SparseMatrixColoringsBlockBandedMatricesExt.jl b/ext/SparseMatrixColoringsBlockBandedMatricesExt.jl
index 2274b8f6..5baf26af 100644
--- a/ext/SparseMatrixColoringsBlockBandedMatricesExt.jl
+++ b/ext/SparseMatrixColoringsBlockBandedMatricesExt.jl
@@ -73,7 +73,7 @@ function blockbanded_coloring(
         length(subblockbandrange(A))
     else
         # vertices within a block are colored naively with distinct micro colors (~ infinite band width)
-        minimum(size(A))
+        typemax(Int)
     end
 
     # for each macroscopic color, count how many microscopic colors will be needed

From 59f4d422f30973ad8bd24b859b2a940dd29e6e2e Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Wed, 2 Oct 2024 17:38:01 +0200
Subject: [PATCH 20/33] Import

---
 test/structured.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/structured.jl b/test/structured.jl
index a88e565c..399ac874 100644
--- a/test/structured.jl
+++ b/test/structured.jl
@@ -1,6 +1,6 @@
 using ArrayInterface: ArrayInterface
 using BandedMatrices: BandedMatrix, brand
-using BlockBandedMatrices: BlockBandedMatrix
+using BlockBandedMatrices: BandedBlockBandedMatrix, BlockBandedMatrix
 using LinearAlgebra
 using SparseMatrixColorings
 using SparseMatrixColorings: cycle_range

From f3c3776fdd2d9eb0ef546f29696f724a48bb204f Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Wed, 2 Oct 2024 17:42:19 +0200
Subject: [PATCH 21/33] Fix

---
 test/structured.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/structured.jl b/test/structured.jl
index 399ac874..b04c624b 100644
--- a/test/structured.jl
+++ b/test/structured.jl
@@ -51,7 +51,7 @@ end;
     for (mb, nb) in [(10, 20), (20, 10)], lb in 0:3, ub in 0:3, _ in 1:10
         rows = rand(1:5, mb)
         cols = rand(1:5, nb)
-        A = BlockBandedMatrix{Float64}(rand(sum(rows), sum(cols)), rows, cols, (l, u))
+        A = BlockBandedMatrix{Float64}(rand(sum(rows), sum(cols)), rows, cols, (lb, ub))
         test_structured_coloring_decompression(A)
     end
 end;

From f55c09ab47e0d18cee37bdac7877a36c6b951498 Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Wed, 2 Oct 2024 17:55:53 +0200
Subject: [PATCH 22/33] Fix version

---
 test/utils.jl | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/test/utils.jl b/test/utils.jl
index 3f2a78c3..deeee6c4 100644
--- a/test/utils.jl
+++ b/test/utils.jl
@@ -120,8 +120,6 @@ function test_coloring_decompression(
     end
 end
 
-OptimalColoringKnown = Union{Diagonal,Bidiagonal,Tridiagonal,BandedMatrix,BlockBandedMatrix}
-
 function test_structured_coloring_decompression(A::AbstractMatrix)
     column_problem = ColoringProblem(; structure=:nonsymmetric, partition=:column)
     row_problem = ColoringProblem(; structure=:nonsymmetric, partition=:row)
@@ -135,7 +133,10 @@ function test_structured_coloring_decompression(A::AbstractMatrix)
     @test D == A
     @test nameof(typeof(D)) == nameof(typeof(A))
     @test structurally_orthogonal_columns(A, color)
-    @test color == ArrayInterface.matrix_colors(A)
+    if VERSION >= v"1.10" || A isa Union{Diagonal,Bidiagonal,Tridiagonal}
+        # banded matrices not supported by ArrayInterface on Julia 1.6
+        @test color == ArrayInterface.matrix_colors(A)
+    end
 
     # Row
     result = coloring(A, row_problem, algo)

From 92a3e2dbce78561b2a1af8bf511e3e213c3da178 Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Thu, 3 Oct 2024 09:49:15 +0200
Subject: [PATCH 23/33] Remove optimized structured implementations

---
 Project.toml                                  |  19 --
 docs/src/dev.md                               |   6 -
 ext/SparseMatrixColoringsBandedMatricesExt.jl |  81 ---------
 ...seMatrixColoringsBlockBandedMatricesExt.jl | 119 -------------
 src/SparseMatrixColorings.jl                  |  16 --
 src/structured.jl                             | 163 ------------------
 test/utils.jl                                 |   2 +-
 7 files changed, 1 insertion(+), 405 deletions(-)
 delete mode 100644 ext/SparseMatrixColoringsBandedMatricesExt.jl
 delete mode 100644 ext/SparseMatrixColoringsBlockBandedMatricesExt.jl
 delete mode 100644 src/structured.jl

diff --git a/Project.toml b/Project.toml
index 1fe45849..87257096 100644
--- a/Project.toml
+++ b/Project.toml
@@ -10,33 +10,14 @@ 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"
-BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
-BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0"
-
-[extensions]
-SparseMatrixColoringsBandedMatricesExt = "BandedMatrices"
-SparseMatrixColoringsBlockBandedMatricesExt = ["BlockArrays", "BlockBandedMatrices"]
-
 [compat]
 ADTypes = "1.2.1"
-BandedMatrices = "1.7.5"
-BlockArrays = "1.1.1"
-BlockBandedMatrices = "0.13.1"
 Compat = "3.46,4.2"
 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"
-
-[extras]
-BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
-BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
-BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0"
\ No newline at end of file
diff --git a/docs/src/dev.md b/docs/src/dev.md
index cd30a9b6..b8aa2792 100644
--- a/docs/src/dev.md
+++ b/docs/src/dev.md
@@ -65,9 +65,3 @@ SparseMatrixColorings.what_fig_61
 SparseMatrixColorings.efficient_fig_1
 SparseMatrixColorings.efficient_fig_4
 ```
-
-## Misc
-
-```@docs
-SparseMatrixColorings.cycle_range
-```
diff --git a/ext/SparseMatrixColoringsBandedMatricesExt.jl b/ext/SparseMatrixColoringsBandedMatricesExt.jl
deleted file mode 100644
index 85ee4b88..00000000
--- a/ext/SparseMatrixColoringsBandedMatricesExt.jl
+++ /dev/null
@@ -1,81 +0,0 @@
-module SparseMatrixColoringsBandedMatricesExt
-
-if isdefined(Base, :get_extension)
-    using BandedMatrices: BandedMatrix, bandrange, bandwidths, colrange, rowrange
-    using SparseMatrixColorings:
-        BipartiteGraph,
-        ColoringProblem,
-        ColumnColoringResult,
-        GreedyColoringAlgorithm,
-        RowColoringResult,
-        column_colors,
-        cycle_range,
-        row_colors
-    import SparseMatrixColorings as SMC
-else
-    using ..BandedMatrices: BandedMatrix, bandrange, bandwidths, colrange, rowrange
-    using ..SparseMatrixColorings:
-        BipartiteGraph,
-        ColoringProblem,
-        ColumnColoringResult,
-        GreedyColoringAlgorithm,
-        RowColoringResult,
-        column_colors,
-        cycle_range,
-        row_colors
-    import ..SparseMatrixColorings as SMC
-end
-
-#=
-This code is partly 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...,
-)
-    width = length(bandrange(A))
-    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...,
-)
-    width = length(bandrange(A))
-    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)
-    for j in axes(A, 2)
-        c = color[j]
-        for i in colrange(A, j)
-            A[i, j] = B[i, c]
-        end
-    end
-    return A
-end
-
-function SMC.decompress!(A::BandedMatrix, B::AbstractMatrix, result::RowColoringResult)
-    color = row_colors(result)
-    for i in axes(A, 1)
-        c = color[i]
-        for j in rowrange(A, i)
-            A[i, j] = B[c, j]
-        end
-    end
-    return A
-end
-
-end
diff --git a/ext/SparseMatrixColoringsBlockBandedMatricesExt.jl b/ext/SparseMatrixColoringsBlockBandedMatricesExt.jl
deleted file mode 100644
index 5baf26af..00000000
--- a/ext/SparseMatrixColoringsBlockBandedMatricesExt.jl
+++ /dev/null
@@ -1,119 +0,0 @@
-module SparseMatrixColoringsBlockBandedMatricesExt
-
-if isdefined(Base, :get_extension)
-    using BlockArrays: blockaxes, blockfirsts, blocklasts, blocksize, blocklengths
-    using BlockBandedMatrices:
-        BandedBlockBandedMatrix,
-        BlockBandedMatrix,
-        blockbandrange,
-        blockbandwidths,
-        blocklengths,
-        blocksize,
-        subblockbandwidths
-    using SparseMatrixColorings:
-        BipartiteGraph,
-        ColoringProblem,
-        ColumnColoringResult,
-        GreedyColoringAlgorithm,
-        RowColoringResult,
-        column_colors,
-        cycle_range,
-        row_colors
-    import SparseMatrixColorings as SMC
-else
-    using ..BlockArrays: blockaxes, blockfirsts, blocklasts, blocksize, blocklengths
-    using ..BlockBandedMatrices:
-        BandedBlockBandedMatrix,
-        BlockBandedMatrix,
-        blockbandrange,
-        blockbandwidths,
-        blocklengths,
-        blocksize,
-        subblockbandwidths
-    using ..SparseMatrixColorings:
-        BipartiteGraph,
-        ColoringProblem,
-        ColumnColoringResult,
-        GreedyColoringAlgorithm,
-        RowColoringResult,
-        column_colors,
-        cycle_range,
-        row_colors
-    import ..SparseMatrixColorings as SMC
-end
-
-#=
-This code is partly taken from ArrayInterface.jl and FiniteDiff.jl
-https://github.com/JuliaArrays/ArrayInterface.jl
-https://github.com/JuliaDiff/FiniteDiff.jl
-=#
-
-function subblockbandrange(A::BandedBlockBandedMatrix)
-    u, l = subblockbandwidths(A)
-    return (-l):u
-end
-
-function blockbanded_coloring(
-    A::Union{BlockBandedMatrix,BandedBlockBandedMatrix}, dim::Integer
-)
-    # consider blocks of columns or rows (let's call them vertices) depending on `dim`
-    nb_blocks = blocksize(A, dim)
-    nb_in_block = blocklengths(axes(A, dim))
-    first_in_block = blockfirsts(axes(A, dim))
-    last_in_block = blocklasts(axes(A, dim))
-    color = zeros(Int, size(A, dim))
-
-    # give a macroscopic color to each block, so that 2 blocks with the same macro color are orthogonal
-    # same idea as for BandedMatrices
-    nb_macrocolors = length(blockbandrange(A))
-    macrocolor = cycle_range(nb_macrocolors, nb_blocks)
-
-    width = if A isa BandedBlockBandedMatrix
-        # vertices within a block are colored cleverly using bands
-        length(subblockbandrange(A))
-    else
-        # vertices within a block are colored naively with distinct micro colors (~ infinite band width)
-        typemax(Int)
-    end
-
-    # for each macroscopic color, count how many microscopic colors will be needed
-    nb_colors_in_macrocolor = zeros(Int, nb_macrocolors)
-    for mc in 1:nb_macrocolors
-        largest_nb_in_macrocolor = maximum(nb_in_block[mc:nb_macrocolors:nb_blocks]; init=0)
-        nb_colors_in_macrocolor[mc] = min(width, largest_nb_in_macrocolor)
-    end
-    color_shift_in_macrocolor = vcat(0, cumsum(nb_colors_in_macrocolor)[1:(end - 1)])
-
-    # assign a microscopic color to each column as a function of its macroscopic color and its position within the block
-    for b in 1:nb_blocks
-        block_color = cycle_range(width, nb_in_block[b])
-        shift = color_shift_in_macrocolor[macrocolor[b]]
-        color[first_in_block[b]:last_in_block[b]] .= shift .+ block_color
-    end
-
-    return color
-end
-
-function SMC.coloring(
-    A::Union{BlockBandedMatrix,BandedBlockBandedMatrix},
-    ::ColoringProblem{:nonsymmetric,:column},
-    algo::GreedyColoringAlgorithm;
-    kwargs...,
-)
-    color = blockbanded_coloring(A, 2)
-    bg = BipartiteGraph(A)
-    return ColumnColoringResult(A, bg, color)
-end
-
-function SMC.coloring(
-    A::Union{BlockBandedMatrix,BandedBlockBandedMatrix},
-    ::ColoringProblem{:nonsymmetric,:row},
-    algo::GreedyColoringAlgorithm;
-    kwargs...,
-)
-    color = blockbanded_coloring(A, 1)
-    bg = BipartiteGraph(A)
-    return RowColoringResult(A, bg, color)
-end
-
-end
diff --git a/src/SparseMatrixColorings.jl b/src/SparseMatrixColorings.jl
index 65fd8c4a..330d662a 100644
--- a/src/SparseMatrixColorings.jl
+++ b/src/SparseMatrixColorings.jl
@@ -53,7 +53,6 @@ include("matrices.jl")
 include("interface.jl")
 include("constant.jl")
 include("decompression.jl")
-include("structured.jl")
 include("check.jl")
 include("examples.jl")
 
@@ -66,19 +65,4 @@ export column_groups, row_groups
 export sparsity_pattern
 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" include(
-            "../ext/SparseMatrixColoringsBandedMatricesExt.jl"
-        )
-        @require BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0" include(
-            "../ext/SparseMatrixColoringsBlockBandedMatricesExt.jl"
-        )
-    end
-end
-
 end
diff --git a/src/structured.jl b/src/structured.jl
deleted file mode 100644
index 78b2db69..00000000
--- a/src/structured.jl
+++ /dev/null
@@ -1,163 +0,0 @@
-#=
-This code is partly taken from ArrayInterface.jl
-https://github.com/JuliaArrays/ArrayInterface.jl
-=#
-
-"""
-    cycle_range(k::Integer, n::Integer)
-
-Concatenate copies of `1:k` to fill a vector of length `n` (with one partial copy allowed at the end).
-"""
-function cycle_range(k::Integer, n::Integer)
-    color = Vector{Int}(undef, n)
-    for i in eachindex(color)
-        color[i] = 1 + (i - 1) % k
-    end
-    return color
-end
-
-## Diagonal
-
-function coloring(
-    A::Diagonal,
-    ::ColoringProblem{:nonsymmetric,:column},
-    algo::GreedyColoringAlgorithm;
-    kwargs...,
-)
-    color = fill(1, size(A, 2))
-    bg = BipartiteGraph(A)
-    return ColumnColoringResult(A, bg, color)
-end
-
-function coloring(
-    A::Diagonal,
-    ::ColoringProblem{:nonsymmetric,:row},
-    algo::GreedyColoringAlgorithm;
-    kwargs...,
-)
-    color = fill(1, size(A, 1))
-    bg = BipartiteGraph(A)
-    return RowColoringResult(A, bg, color)
-end
-
-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, B::AbstractMatrix, result::RowColoringResult)
-    color = row_colors(result)
-    for i in axes(A, 1)
-        A[i, i] = B[color[i], i]
-    end
-    return A
-end
-
-## Bidiagonal
-
-function coloring(
-    A::Bidiagonal,
-    ::ColoringProblem{:nonsymmetric,:column},
-    algo::GreedyColoringAlgorithm;
-    kwargs...,
-)
-    color = cycle_range(2, size(A, 2))
-    bg = BipartiteGraph(A)
-    return ColumnColoringResult(A, bg, color)
-end
-
-function coloring(
-    A::Bidiagonal,
-    ::ColoringProblem{:nonsymmetric,:row},
-    algo::GreedyColoringAlgorithm;
-    kwargs...,
-)
-    color = cycle_range(2, size(A, 1))
-    bg = BipartiteGraph(A)
-    return RowColoringResult(A, bg, color)
-end
-
-function decompress!(A::Bidiagonal, B::AbstractMatrix, result::ColumnColoringResult)
-    color = column_colors(result)
-    for j in axes(A, 2)
-        c = color[j]
-        A[j, j] = B[j, c]
-        if A.uplo == 'U' && j > 1  # above
-            A[j - 1, j] = B[j - 1, c]
-        elseif A.uplo == 'L' && j < size(A, 2)  # below
-            A[j + 1, j] = B[j + 1, c]
-        end
-    end
-    return A
-end
-
-function decompress!(A::Bidiagonal, B::AbstractMatrix, result::RowColoringResult)
-    color = row_colors(result)
-    for i in axes(A, 1)
-        c = color[i]
-        A[i, i] = B[c, i]
-        if A.uplo == 'U' && i < size(A, 1)  # right
-            A[i, i + 1] = B[c, i + 1]
-        elseif A.uplo == 'L' && i > 1  # left
-            A[i, i - 1] = B[c, i - 1]
-        end
-    end
-    return A
-end
-
-## Tridiagonal
-
-function coloring(
-    A::Tridiagonal,
-    ::ColoringProblem{:nonsymmetric,:column},
-    algo::GreedyColoringAlgorithm;
-    kwargs...,
-)
-    color = cycle_range(3, size(A, 2))
-    bg = BipartiteGraph(A)
-    return ColumnColoringResult(A, bg, color)
-end
-
-function coloring(
-    A::Tridiagonal,
-    ::ColoringProblem{:nonsymmetric,:row},
-    algo::GreedyColoringAlgorithm;
-    kwargs...,
-)
-    color = cycle_range(3, size(A, 1))
-    bg = BipartiteGraph(A)
-    return RowColoringResult(A, bg, color)
-end
-
-function decompress!(A::Tridiagonal, B::AbstractMatrix, result::ColumnColoringResult)
-    color = column_colors(result)
-    for j in axes(A, 2)
-        c = color[j]
-        A[j, j] = B[j, c]
-        if j > 1  # above
-            A[j - 1, j] = B[j - 1, c]
-        end
-        if j < size(A, 2)  # below
-            A[j + 1, j] = B[j + 1, c]
-        end
-    end
-    return A
-end
-
-function decompress!(A::Tridiagonal, B::AbstractMatrix, result::RowColoringResult)
-    color = row_colors(result)
-    for i in axes(A, 1)
-        c = color[i]
-        A[i, i] = B[c, i]
-        if i < size(A, 1)  # right
-            A[i, i + 1] = B[c, i + 1]
-        end
-        if i > 1  # left
-            A[i, i - 1] = B[c, i - 1]
-        end
-    end
-    return A
-end
diff --git a/test/utils.jl b/test/utils.jl
index 250d16aa..0cec5373 100644
--- a/test/utils.jl
+++ b/test/utils.jl
@@ -157,7 +157,7 @@ function test_structured_coloring_decompression(A::AbstractMatrix)
     @test structurally_orthogonal_columns(A, color)
     if VERSION >= v"1.10" || A isa Union{Diagonal,Bidiagonal,Tridiagonal}
         # banded matrices not supported by ArrayInterface on Julia 1.6
-        @test color == ArrayInterface.matrix_colors(A)
+        # @test color == ArrayInterface.matrix_colors(A)  # TODO: uncomment
     end
 
     # Row

From f8cfd6adae7d29db20f5898551f9eba7a50bff85 Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Thu, 3 Oct 2024 09:50:59 +0200
Subject: [PATCH 24/33] Min diff

---
 .github/workflows/Test.yml   | 1 -
 Project.toml                 | 2 +-
 src/SparseMatrixColorings.jl | 2 --
 3 files changed, 1 insertion(+), 4 deletions(-)

diff --git a/.github/workflows/Test.yml b/.github/workflows/Test.yml
index a2841d90..e9dc089e 100644
--- a/.github/workflows/Test.yml
+++ b/.github/workflows/Test.yml
@@ -19,7 +19,6 @@ jobs:
   test:
     runs-on: ubuntu-latest
     strategy:
-      fail-fast: false
       matrix:
         julia-version: ['lts', '1', 'pre']
 
diff --git a/Project.toml b/Project.toml
index 87257096..42da01ea 100644
--- a/Project.toml
+++ b/Project.toml
@@ -16,8 +16,8 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 ADTypes = "1.2.1"
 Compat = "3.46,4.2"
 DataStructures = "0.18"
-DocStringExtensions = "0.8,0.9"
 LinearAlgebra = "<0.0.1, 1"
+DocStringExtensions = "0.8,0.9"
 Random = "<0.0.1, 1"
 SparseArrays = "<0.0.1, 1"
 julia = "1.6"
diff --git a/src/SparseMatrixColorings.jl b/src/SparseMatrixColorings.jl
index 330d662a..b91b699a 100644
--- a/src/SparseMatrixColorings.jl
+++ b/src/SparseMatrixColorings.jl
@@ -16,13 +16,11 @@ using DataStructures: DisjointSets, find_root!, root_union!, num_groups
 using DocStringExtensions: README, EXPORTS, SIGNATURES, TYPEDEF, TYPEDFIELDS
 using LinearAlgebra:
     Adjoint,
-    Bidiagonal,
     Diagonal,
     Hermitian,
     LowerTriangular,
     Symmetric,
     Transpose,
-    Tridiagonal,
     UpperTriangular,
     adjoint,
     checksquare,

From f3531afb1ebe39c2a274c790cc977730a3596566 Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Thu, 3 Oct 2024 09:56:54 +0200
Subject: [PATCH 25/33] Remove cycle range tests

---
 test/structured.jl | 11 -----------
 1 file changed, 11 deletions(-)

diff --git a/test/structured.jl b/test/structured.jl
index b04c624b..e038b98d 100644
--- a/test/structured.jl
+++ b/test/structured.jl
@@ -6,17 +6,6 @@ using SparseMatrixColorings
 using SparseMatrixColorings: cycle_range
 using Test
 
-@testset "Utils" begin
-    @test cycle_range(2, 3) == [1, 2, 1]
-    @test cycle_range(2, 4) == [1, 2, 1, 2]
-    @test cycle_range(2, 5) == [1, 2, 1, 2, 1]
-    @test cycle_range(3, 5) == [1, 2, 3, 1, 2]
-    @test cycle_range(3, 6) == [1, 2, 3, 1, 2, 3]
-    @test cycle_range(2, 1) == [1]
-    @test cycle_range(3, 1) == [1]
-    @test cycle_range(3, 2) == [1, 2]
-end;
-
 @testset "Diagonal" begin
     for n in (1, 2, 10, 100)
         A = Diagonal(rand(n))

From a430140a533f607fae738fe1dc092e33fc3bf440 Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Thu, 3 Oct 2024 10:06:26 +0200
Subject: [PATCH 26/33] Optimized implementation for structured

---
 Project.toml                                  |  21 ++-
 docs/src/dev.md                               |   6 +
 ext/SparseMatrixColoringsBandedMatricesExt.jl |  81 +++++++++
 ...seMatrixColoringsBlockBandedMatricesExt.jl | 119 +++++++++++++
 src/SparseMatrixColorings.jl                  |  18 ++
 src/structured.jl                             | 163 ++++++++++++++++++
 test/utils.jl                                 |   2 +-
 7 files changed, 408 insertions(+), 2 deletions(-)
 create mode 100644 ext/SparseMatrixColoringsBandedMatricesExt.jl
 create mode 100644 ext/SparseMatrixColoringsBlockBandedMatricesExt.jl
 create mode 100644 src/structured.jl

diff --git a/Project.toml b/Project.toml
index 42da01ea..1fe45849 100644
--- a/Project.toml
+++ b/Project.toml
@@ -10,14 +10,33 @@ 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"
+BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
+BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0"
+
+[extensions]
+SparseMatrixColoringsBandedMatricesExt = "BandedMatrices"
+SparseMatrixColoringsBlockBandedMatricesExt = ["BlockArrays", "BlockBandedMatrices"]
+
 [compat]
 ADTypes = "1.2.1"
+BandedMatrices = "1.7.5"
+BlockArrays = "1.1.1"
+BlockBandedMatrices = "0.13.1"
 Compat = "3.46,4.2"
 DataStructures = "0.18"
-LinearAlgebra = "<0.0.1, 1"
 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"
+
+[extras]
+BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
+BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
+BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0"
\ No newline at end of file
diff --git a/docs/src/dev.md b/docs/src/dev.md
index b8aa2792..cd30a9b6 100644
--- a/docs/src/dev.md
+++ b/docs/src/dev.md
@@ -65,3 +65,9 @@ SparseMatrixColorings.what_fig_61
 SparseMatrixColorings.efficient_fig_1
 SparseMatrixColorings.efficient_fig_4
 ```
+
+## Misc
+
+```@docs
+SparseMatrixColorings.cycle_range
+```
diff --git a/ext/SparseMatrixColoringsBandedMatricesExt.jl b/ext/SparseMatrixColoringsBandedMatricesExt.jl
new file mode 100644
index 00000000..85ee4b88
--- /dev/null
+++ b/ext/SparseMatrixColoringsBandedMatricesExt.jl
@@ -0,0 +1,81 @@
+module SparseMatrixColoringsBandedMatricesExt
+
+if isdefined(Base, :get_extension)
+    using BandedMatrices: BandedMatrix, bandrange, bandwidths, colrange, rowrange
+    using SparseMatrixColorings:
+        BipartiteGraph,
+        ColoringProblem,
+        ColumnColoringResult,
+        GreedyColoringAlgorithm,
+        RowColoringResult,
+        column_colors,
+        cycle_range,
+        row_colors
+    import SparseMatrixColorings as SMC
+else
+    using ..BandedMatrices: BandedMatrix, bandrange, bandwidths, colrange, rowrange
+    using ..SparseMatrixColorings:
+        BipartiteGraph,
+        ColoringProblem,
+        ColumnColoringResult,
+        GreedyColoringAlgorithm,
+        RowColoringResult,
+        column_colors,
+        cycle_range,
+        row_colors
+    import ..SparseMatrixColorings as SMC
+end
+
+#=
+This code is partly 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...,
+)
+    width = length(bandrange(A))
+    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...,
+)
+    width = length(bandrange(A))
+    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)
+    for j in axes(A, 2)
+        c = color[j]
+        for i in colrange(A, j)
+            A[i, j] = B[i, c]
+        end
+    end
+    return A
+end
+
+function SMC.decompress!(A::BandedMatrix, B::AbstractMatrix, result::RowColoringResult)
+    color = row_colors(result)
+    for i in axes(A, 1)
+        c = color[i]
+        for j in rowrange(A, i)
+            A[i, j] = B[c, j]
+        end
+    end
+    return A
+end
+
+end
diff --git a/ext/SparseMatrixColoringsBlockBandedMatricesExt.jl b/ext/SparseMatrixColoringsBlockBandedMatricesExt.jl
new file mode 100644
index 00000000..5baf26af
--- /dev/null
+++ b/ext/SparseMatrixColoringsBlockBandedMatricesExt.jl
@@ -0,0 +1,119 @@
+module SparseMatrixColoringsBlockBandedMatricesExt
+
+if isdefined(Base, :get_extension)
+    using BlockArrays: blockaxes, blockfirsts, blocklasts, blocksize, blocklengths
+    using BlockBandedMatrices:
+        BandedBlockBandedMatrix,
+        BlockBandedMatrix,
+        blockbandrange,
+        blockbandwidths,
+        blocklengths,
+        blocksize,
+        subblockbandwidths
+    using SparseMatrixColorings:
+        BipartiteGraph,
+        ColoringProblem,
+        ColumnColoringResult,
+        GreedyColoringAlgorithm,
+        RowColoringResult,
+        column_colors,
+        cycle_range,
+        row_colors
+    import SparseMatrixColorings as SMC
+else
+    using ..BlockArrays: blockaxes, blockfirsts, blocklasts, blocksize, blocklengths
+    using ..BlockBandedMatrices:
+        BandedBlockBandedMatrix,
+        BlockBandedMatrix,
+        blockbandrange,
+        blockbandwidths,
+        blocklengths,
+        blocksize,
+        subblockbandwidths
+    using ..SparseMatrixColorings:
+        BipartiteGraph,
+        ColoringProblem,
+        ColumnColoringResult,
+        GreedyColoringAlgorithm,
+        RowColoringResult,
+        column_colors,
+        cycle_range,
+        row_colors
+    import ..SparseMatrixColorings as SMC
+end
+
+#=
+This code is partly taken from ArrayInterface.jl and FiniteDiff.jl
+https://github.com/JuliaArrays/ArrayInterface.jl
+https://github.com/JuliaDiff/FiniteDiff.jl
+=#
+
+function subblockbandrange(A::BandedBlockBandedMatrix)
+    u, l = subblockbandwidths(A)
+    return (-l):u
+end
+
+function blockbanded_coloring(
+    A::Union{BlockBandedMatrix,BandedBlockBandedMatrix}, dim::Integer
+)
+    # consider blocks of columns or rows (let's call them vertices) depending on `dim`
+    nb_blocks = blocksize(A, dim)
+    nb_in_block = blocklengths(axes(A, dim))
+    first_in_block = blockfirsts(axes(A, dim))
+    last_in_block = blocklasts(axes(A, dim))
+    color = zeros(Int, size(A, dim))
+
+    # give a macroscopic color to each block, so that 2 blocks with the same macro color are orthogonal
+    # same idea as for BandedMatrices
+    nb_macrocolors = length(blockbandrange(A))
+    macrocolor = cycle_range(nb_macrocolors, nb_blocks)
+
+    width = if A isa BandedBlockBandedMatrix
+        # vertices within a block are colored cleverly using bands
+        length(subblockbandrange(A))
+    else
+        # vertices within a block are colored naively with distinct micro colors (~ infinite band width)
+        typemax(Int)
+    end
+
+    # for each macroscopic color, count how many microscopic colors will be needed
+    nb_colors_in_macrocolor = zeros(Int, nb_macrocolors)
+    for mc in 1:nb_macrocolors
+        largest_nb_in_macrocolor = maximum(nb_in_block[mc:nb_macrocolors:nb_blocks]; init=0)
+        nb_colors_in_macrocolor[mc] = min(width, largest_nb_in_macrocolor)
+    end
+    color_shift_in_macrocolor = vcat(0, cumsum(nb_colors_in_macrocolor)[1:(end - 1)])
+
+    # assign a microscopic color to each column as a function of its macroscopic color and its position within the block
+    for b in 1:nb_blocks
+        block_color = cycle_range(width, nb_in_block[b])
+        shift = color_shift_in_macrocolor[macrocolor[b]]
+        color[first_in_block[b]:last_in_block[b]] .= shift .+ block_color
+    end
+
+    return color
+end
+
+function SMC.coloring(
+    A::Union{BlockBandedMatrix,BandedBlockBandedMatrix},
+    ::ColoringProblem{:nonsymmetric,:column},
+    algo::GreedyColoringAlgorithm;
+    kwargs...,
+)
+    color = blockbanded_coloring(A, 2)
+    bg = BipartiteGraph(A)
+    return ColumnColoringResult(A, bg, color)
+end
+
+function SMC.coloring(
+    A::Union{BlockBandedMatrix,BandedBlockBandedMatrix},
+    ::ColoringProblem{:nonsymmetric,:row},
+    algo::GreedyColoringAlgorithm;
+    kwargs...,
+)
+    color = blockbanded_coloring(A, 1)
+    bg = BipartiteGraph(A)
+    return RowColoringResult(A, bg, color)
+end
+
+end
diff --git a/src/SparseMatrixColorings.jl b/src/SparseMatrixColorings.jl
index b91b699a..65fd8c4a 100644
--- a/src/SparseMatrixColorings.jl
+++ b/src/SparseMatrixColorings.jl
@@ -16,11 +16,13 @@ using DataStructures: DisjointSets, find_root!, root_union!, num_groups
 using DocStringExtensions: README, EXPORTS, SIGNATURES, TYPEDEF, TYPEDFIELDS
 using LinearAlgebra:
     Adjoint,
+    Bidiagonal,
     Diagonal,
     Hermitian,
     LowerTriangular,
     Symmetric,
     Transpose,
+    Tridiagonal,
     UpperTriangular,
     adjoint,
     checksquare,
@@ -51,6 +53,7 @@ include("matrices.jl")
 include("interface.jl")
 include("constant.jl")
 include("decompression.jl")
+include("structured.jl")
 include("check.jl")
 include("examples.jl")
 
@@ -63,4 +66,19 @@ export column_groups, row_groups
 export sparsity_pattern
 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" include(
+            "../ext/SparseMatrixColoringsBandedMatricesExt.jl"
+        )
+        @require BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0" include(
+            "../ext/SparseMatrixColoringsBlockBandedMatricesExt.jl"
+        )
+    end
+end
+
 end
diff --git a/src/structured.jl b/src/structured.jl
new file mode 100644
index 00000000..78b2db69
--- /dev/null
+++ b/src/structured.jl
@@ -0,0 +1,163 @@
+#=
+This code is partly taken from ArrayInterface.jl
+https://github.com/JuliaArrays/ArrayInterface.jl
+=#
+
+"""
+    cycle_range(k::Integer, n::Integer)
+
+Concatenate copies of `1:k` to fill a vector of length `n` (with one partial copy allowed at the end).
+"""
+function cycle_range(k::Integer, n::Integer)
+    color = Vector{Int}(undef, n)
+    for i in eachindex(color)
+        color[i] = 1 + (i - 1) % k
+    end
+    return color
+end
+
+## Diagonal
+
+function coloring(
+    A::Diagonal,
+    ::ColoringProblem{:nonsymmetric,:column},
+    algo::GreedyColoringAlgorithm;
+    kwargs...,
+)
+    color = fill(1, size(A, 2))
+    bg = BipartiteGraph(A)
+    return ColumnColoringResult(A, bg, color)
+end
+
+function coloring(
+    A::Diagonal,
+    ::ColoringProblem{:nonsymmetric,:row},
+    algo::GreedyColoringAlgorithm;
+    kwargs...,
+)
+    color = fill(1, size(A, 1))
+    bg = BipartiteGraph(A)
+    return RowColoringResult(A, bg, color)
+end
+
+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, B::AbstractMatrix, result::RowColoringResult)
+    color = row_colors(result)
+    for i in axes(A, 1)
+        A[i, i] = B[color[i], i]
+    end
+    return A
+end
+
+## Bidiagonal
+
+function coloring(
+    A::Bidiagonal,
+    ::ColoringProblem{:nonsymmetric,:column},
+    algo::GreedyColoringAlgorithm;
+    kwargs...,
+)
+    color = cycle_range(2, size(A, 2))
+    bg = BipartiteGraph(A)
+    return ColumnColoringResult(A, bg, color)
+end
+
+function coloring(
+    A::Bidiagonal,
+    ::ColoringProblem{:nonsymmetric,:row},
+    algo::GreedyColoringAlgorithm;
+    kwargs...,
+)
+    color = cycle_range(2, size(A, 1))
+    bg = BipartiteGraph(A)
+    return RowColoringResult(A, bg, color)
+end
+
+function decompress!(A::Bidiagonal, B::AbstractMatrix, result::ColumnColoringResult)
+    color = column_colors(result)
+    for j in axes(A, 2)
+        c = color[j]
+        A[j, j] = B[j, c]
+        if A.uplo == 'U' && j > 1  # above
+            A[j - 1, j] = B[j - 1, c]
+        elseif A.uplo == 'L' && j < size(A, 2)  # below
+            A[j + 1, j] = B[j + 1, c]
+        end
+    end
+    return A
+end
+
+function decompress!(A::Bidiagonal, B::AbstractMatrix, result::RowColoringResult)
+    color = row_colors(result)
+    for i in axes(A, 1)
+        c = color[i]
+        A[i, i] = B[c, i]
+        if A.uplo == 'U' && i < size(A, 1)  # right
+            A[i, i + 1] = B[c, i + 1]
+        elseif A.uplo == 'L' && i > 1  # left
+            A[i, i - 1] = B[c, i - 1]
+        end
+    end
+    return A
+end
+
+## Tridiagonal
+
+function coloring(
+    A::Tridiagonal,
+    ::ColoringProblem{:nonsymmetric,:column},
+    algo::GreedyColoringAlgorithm;
+    kwargs...,
+)
+    color = cycle_range(3, size(A, 2))
+    bg = BipartiteGraph(A)
+    return ColumnColoringResult(A, bg, color)
+end
+
+function coloring(
+    A::Tridiagonal,
+    ::ColoringProblem{:nonsymmetric,:row},
+    algo::GreedyColoringAlgorithm;
+    kwargs...,
+)
+    color = cycle_range(3, size(A, 1))
+    bg = BipartiteGraph(A)
+    return RowColoringResult(A, bg, color)
+end
+
+function decompress!(A::Tridiagonal, B::AbstractMatrix, result::ColumnColoringResult)
+    color = column_colors(result)
+    for j in axes(A, 2)
+        c = color[j]
+        A[j, j] = B[j, c]
+        if j > 1  # above
+            A[j - 1, j] = B[j - 1, c]
+        end
+        if j < size(A, 2)  # below
+            A[j + 1, j] = B[j + 1, c]
+        end
+    end
+    return A
+end
+
+function decompress!(A::Tridiagonal, B::AbstractMatrix, result::RowColoringResult)
+    color = row_colors(result)
+    for i in axes(A, 1)
+        c = color[i]
+        A[i, i] = B[c, i]
+        if i < size(A, 1)  # right
+            A[i, i + 1] = B[c, i + 1]
+        end
+        if i > 1  # left
+            A[i, i - 1] = B[c, i - 1]
+        end
+    end
+    return A
+end
diff --git a/test/utils.jl b/test/utils.jl
index 0cec5373..250d16aa 100644
--- a/test/utils.jl
+++ b/test/utils.jl
@@ -157,7 +157,7 @@ function test_structured_coloring_decompression(A::AbstractMatrix)
     @test structurally_orthogonal_columns(A, color)
     if VERSION >= v"1.10" || A isa Union{Diagonal,Bidiagonal,Tridiagonal}
         # banded matrices not supported by ArrayInterface on Julia 1.6
-        # @test color == ArrayInterface.matrix_colors(A)  # TODO: uncomment
+        @test color == ArrayInterface.matrix_colors(A)
     end
 
     # Row

From 61da959ef3f1cda0dbf20be1e9c4761f92e6bc68 Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Thu, 3 Oct 2024 10:07:43 +0200
Subject: [PATCH 27/33] Fix import

---
 test/structured.jl | 1 -
 1 file changed, 1 deletion(-)

diff --git a/test/structured.jl b/test/structured.jl
index e038b98d..23ba967f 100644
--- a/test/structured.jl
+++ b/test/structured.jl
@@ -3,7 +3,6 @@ using BandedMatrices: BandedMatrix, brand
 using BlockBandedMatrices: BandedBlockBandedMatrix, BlockBandedMatrix
 using LinearAlgebra
 using SparseMatrixColorings
-using SparseMatrixColorings: cycle_range
 using Test
 
 @testset "Diagonal" begin

From 692ff43fa8596157b19e481fe7803d1a24a08748 Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Sun, 6 Oct 2024 17:03:43 +0200
Subject: [PATCH 28/33] Fix

---
 Project.toml | 1 -
 1 file changed, 1 deletion(-)

diff --git a/Project.toml b/Project.toml
index 9e3b1012..7661daaa 100644
--- a/Project.toml
+++ b/Project.toml
@@ -29,7 +29,6 @@ BlockArrays = "1.1.1"
 BlockBandedMatrices = "0.13.1"
 Compat = "3.46,4.2"
 DataStructures = "0.18"
-DocStringExtensions = "0.8,0.9"
 LinearAlgebra = "<0.0.1, 1"
 DocStringExtensions = "0.8,0.9"
 Random = "<0.0.1, 1"

From 8e25d9c74f05cf60c9478d005e026d80d69d4448 Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Tue, 1 Jul 2025 16:07:24 +0200
Subject: [PATCH 29/33] Remove Requires

---
 Project.toml                                  |  4 +-
 ext/SparseMatrixColoringsBandedMatricesExt.jl | 36 ++++-------
 ...seMatrixColoringsBlockBandedMatricesExt.jl | 60 ++++++-------------
 src/SparseMatrixColorings.jl                  | 15 -----
 4 files changed, 32 insertions(+), 83 deletions(-)

diff --git a/Project.toml b/Project.toml
index 71129082..5ce3e0c0 100644
--- a/Project.toml
+++ b/Project.toml
@@ -28,8 +28,8 @@ SparseMatrixColoringsColorsExt = "Colors"
 
 [compat]
 ADTypes = "1.2.1"
-BandedMatrices = "1.7.5"
-BlockArrays = "1.1.1"
+BandedMatrices = "1.9.4"
+BlockArrays = "1.6.3"
 BlockBandedMatrices = "0.13.1"
 CUDA = "5.8.2"
 CliqueTrees = "1"
diff --git a/ext/SparseMatrixColoringsBandedMatricesExt.jl b/ext/SparseMatrixColoringsBandedMatricesExt.jl
index 85ee4b88..01fec06f 100644
--- a/ext/SparseMatrixColoringsBandedMatricesExt.jl
+++ b/ext/SparseMatrixColoringsBandedMatricesExt.jl
@@ -1,30 +1,16 @@
 module SparseMatrixColoringsBandedMatricesExt
 
-if isdefined(Base, :get_extension)
-    using BandedMatrices: BandedMatrix, bandrange, bandwidths, colrange, rowrange
-    using SparseMatrixColorings:
-        BipartiteGraph,
-        ColoringProblem,
-        ColumnColoringResult,
-        GreedyColoringAlgorithm,
-        RowColoringResult,
-        column_colors,
-        cycle_range,
-        row_colors
-    import SparseMatrixColorings as SMC
-else
-    using ..BandedMatrices: BandedMatrix, bandrange, bandwidths, colrange, rowrange
-    using ..SparseMatrixColorings:
-        BipartiteGraph,
-        ColoringProblem,
-        ColumnColoringResult,
-        GreedyColoringAlgorithm,
-        RowColoringResult,
-        column_colors,
-        cycle_range,
-        row_colors
-    import ..SparseMatrixColorings as SMC
-end
+using BandedMatrices: BandedMatrix, bandrange, bandwidths, colrange, rowrange
+using SparseMatrixColorings:
+    BipartiteGraph,
+    ColoringProblem,
+    ColumnColoringResult,
+    GreedyColoringAlgorithm,
+    RowColoringResult,
+    column_colors,
+    cycle_range,
+    row_colors
+import SparseMatrixColorings as SMC
 
 #=
 This code is partly taken from ArrayInterface.jl and FiniteDiff.jl
diff --git a/ext/SparseMatrixColoringsBlockBandedMatricesExt.jl b/ext/SparseMatrixColoringsBlockBandedMatricesExt.jl
index 5baf26af..89fa80c4 100644
--- a/ext/SparseMatrixColoringsBlockBandedMatricesExt.jl
+++ b/ext/SparseMatrixColoringsBlockBandedMatricesExt.jl
@@ -1,46 +1,24 @@
 module SparseMatrixColoringsBlockBandedMatricesExt
 
-if isdefined(Base, :get_extension)
-    using BlockArrays: blockaxes, blockfirsts, blocklasts, blocksize, blocklengths
-    using BlockBandedMatrices:
-        BandedBlockBandedMatrix,
-        BlockBandedMatrix,
-        blockbandrange,
-        blockbandwidths,
-        blocklengths,
-        blocksize,
-        subblockbandwidths
-    using SparseMatrixColorings:
-        BipartiteGraph,
-        ColoringProblem,
-        ColumnColoringResult,
-        GreedyColoringAlgorithm,
-        RowColoringResult,
-        column_colors,
-        cycle_range,
-        row_colors
-    import SparseMatrixColorings as SMC
-else
-    using ..BlockArrays: blockaxes, blockfirsts, blocklasts, blocksize, blocklengths
-    using ..BlockBandedMatrices:
-        BandedBlockBandedMatrix,
-        BlockBandedMatrix,
-        blockbandrange,
-        blockbandwidths,
-        blocklengths,
-        blocksize,
-        subblockbandwidths
-    using ..SparseMatrixColorings:
-        BipartiteGraph,
-        ColoringProblem,
-        ColumnColoringResult,
-        GreedyColoringAlgorithm,
-        RowColoringResult,
-        column_colors,
-        cycle_range,
-        row_colors
-    import ..SparseMatrixColorings as SMC
-end
+using BlockArrays: blockaxes, blockfirsts, blocklasts, blocksize, blocklengths
+using BlockBandedMatrices:
+    BandedBlockBandedMatrix,
+    BlockBandedMatrix,
+    blockbandrange,
+    blockbandwidths,
+    blocklengths,
+    blocksize,
+    subblockbandwidths
+using SparseMatrixColorings:
+    BipartiteGraph,
+    ColoringProblem,
+    ColumnColoringResult,
+    GreedyColoringAlgorithm,
+    RowColoringResult,
+    column_colors,
+    cycle_range,
+    row_colors
+import SparseMatrixColorings as SMC
 
 #=
 This code is partly taken from ArrayInterface.jl and FiniteDiff.jl
diff --git a/src/SparseMatrixColorings.jl b/src/SparseMatrixColorings.jl
index 2e641aec..bf8f3211 100644
--- a/src/SparseMatrixColorings.jl
+++ b/src/SparseMatrixColorings.jl
@@ -73,19 +73,4 @@ export column_groups, row_groups
 export sparsity_pattern
 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" include(
-            "../ext/SparseMatrixColoringsBandedMatricesExt.jl"
-        )
-        @require BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0" include(
-            "../ext/SparseMatrixColoringsBlockBandedMatricesExt.jl"
-        )
-    end
-end
-
 end

From d1f0a517f9569bb251505cda3773fa078417e439 Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Tue, 1 Jul 2025 16:08:16 +0200
Subject: [PATCH 30/33] Fix test

---
 test/utils.jl | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/test/utils.jl b/test/utils.jl
index ad8addcb..7c2329b7 100644
--- a/test/utils.jl
+++ b/test/utils.jl
@@ -231,10 +231,7 @@ function test_structured_coloring_decompression(A::AbstractMatrix)
     @test D == A
     @test nameof(typeof(D)) == nameof(typeof(A))
     @test structurally_orthogonal_columns(A, color)
-    if VERSION >= v"1.10" || A isa Union{Diagonal,Bidiagonal,Tridiagonal}
-        # banded matrices not supported by ArrayInterface on Julia 1.6
-        @test color == ArrayInterface.matrix_colors(A)
-    end
+    @test color == ArrayInterface.matrix_colors(A)
 
     # Row
     result = coloring(A, row_problem, algo)

From 69c2d96713a2cc547b9de5f6aaaabc04ded54283 Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Tue, 1 Jul 2025 16:25:58 +0200
Subject: [PATCH 31/33] Define StructuredColoringAlgorithm

---
 Project.toml                                  |  2 +-
 docs/make.jl                                  |  6 +-
 docs/src/api.md                               |  6 ++
 ext/SparseMatrixColoringsBandedMatricesExt.jl |  6 +-
 ...seMatrixColoringsBlockBandedMatricesExt.jl |  6 +-
 src/SparseMatrixColorings.jl                  |  2 +-
 src/structured.jl                             | 30 ++++++--
 test/structured.jl                            | 68 +++++++++++--------
 test/utils.jl                                 |  9 ++-
 9 files changed, 89 insertions(+), 46 deletions(-)

diff --git a/Project.toml b/Project.toml
index 5ce3e0c0..662685ed 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "SparseMatrixColorings"
 uuid = "0a514795-09f3-496d-8182-132a7b665d35"
 authors = ["Guillaume Dalle", "Alexis Montoison"]
-version = "0.4.21"
+version = "0.4.22"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
diff --git a/docs/make.jl b/docs/make.jl
index 78fba268..718b8f15 100644
--- a/docs/make.jl
+++ b/docs/make.jl
@@ -2,7 +2,11 @@ using Documenter
 using DocumenterInterLinks
 using SparseMatrixColorings
 
-links = InterLinks("ADTypes" => "https://sciml.github.io/ADTypes.jl/stable/")
+links = InterLinks(
+    "ADTypes" => "https://sciml.github.io/ADTypes.jl/stable/",
+    "BandedMatrices" => "https://julialinearalgebra.github.io/BandedMatrices.jl/stable/",
+    "LinearAlgebra" => "https://docs.julialang.org/en/v1/stdlib/LinearAlgebra/",
+)
 
 cp(joinpath(@__DIR__, "..", "README.md"), joinpath(@__DIR__, "src", "index.md"); force=true)
 
diff --git a/docs/src/api.md b/docs/src/api.md
index c121d1e3..4fd6d28e 100644
--- a/docs/src/api.md
+++ b/docs/src/api.md
@@ -17,8 +17,14 @@ SparseMatrixColorings
 coloring
 fast_coloring
 ColoringProblem
+```
+
+## Coloring algorithms
+
+```@docs
 GreedyColoringAlgorithm
 ConstantColoringAlgorithm
+StructuredColoringAlgorithm
 ```
 
 ## Result analysis
diff --git a/ext/SparseMatrixColoringsBandedMatricesExt.jl b/ext/SparseMatrixColoringsBandedMatricesExt.jl
index 01fec06f..e4dd32f7 100644
--- a/ext/SparseMatrixColoringsBandedMatricesExt.jl
+++ b/ext/SparseMatrixColoringsBandedMatricesExt.jl
@@ -5,7 +5,7 @@ using SparseMatrixColorings:
     BipartiteGraph,
     ColoringProblem,
     ColumnColoringResult,
-    GreedyColoringAlgorithm,
+    StructuredColoringAlgorithm,
     RowColoringResult,
     column_colors,
     cycle_range,
@@ -21,7 +21,7 @@ https://github.com/JuliaDiff/FiniteDiff.jl
 function SMC.coloring(
     A::BandedMatrix,
     ::ColoringProblem{:nonsymmetric,:column},
-    algo::GreedyColoringAlgorithm;
+    ::StructuredColoringAlgorithm;
     kwargs...,
 )
     width = length(bandrange(A))
@@ -33,7 +33,7 @@ end
 function SMC.coloring(
     A::BandedMatrix,
     ::ColoringProblem{:nonsymmetric,:row},
-    algo::GreedyColoringAlgorithm;
+    ::StructuredColoringAlgorithm;
     kwargs...,
 )
     width = length(bandrange(A))
diff --git a/ext/SparseMatrixColoringsBlockBandedMatricesExt.jl b/ext/SparseMatrixColoringsBlockBandedMatricesExt.jl
index 89fa80c4..aaecb0ad 100644
--- a/ext/SparseMatrixColoringsBlockBandedMatricesExt.jl
+++ b/ext/SparseMatrixColoringsBlockBandedMatricesExt.jl
@@ -13,7 +13,7 @@ using SparseMatrixColorings:
     BipartiteGraph,
     ColoringProblem,
     ColumnColoringResult,
-    GreedyColoringAlgorithm,
+    StructuredColoringAlgorithm,
     RowColoringResult,
     column_colors,
     cycle_range,
@@ -75,7 +75,7 @@ end
 function SMC.coloring(
     A::Union{BlockBandedMatrix,BandedBlockBandedMatrix},
     ::ColoringProblem{:nonsymmetric,:column},
-    algo::GreedyColoringAlgorithm;
+    ::StructuredColoringAlgorithm;
     kwargs...,
 )
     color = blockbanded_coloring(A, 2)
@@ -86,7 +86,7 @@ end
 function SMC.coloring(
     A::Union{BlockBandedMatrix,BandedBlockBandedMatrix},
     ::ColoringProblem{:nonsymmetric,:row},
-    algo::GreedyColoringAlgorithm;
+    ::StructuredColoringAlgorithm;
     kwargs...,
 )
     color = blockbanded_coloring(A, 1)
diff --git a/src/SparseMatrixColorings.jl b/src/SparseMatrixColorings.jl
index bf8f3211..532ef192 100644
--- a/src/SparseMatrixColorings.jl
+++ b/src/SparseMatrixColorings.jl
@@ -66,7 +66,7 @@ export NaturalOrder, RandomOrder, LargestFirst
 export DynamicDegreeBasedOrder, SmallestLast, IncidenceDegree, DynamicLargestFirst
 export PerfectEliminationOrder
 export ColoringProblem, GreedyColoringAlgorithm, AbstractColoringResult
-export ConstantColoringAlgorithm
+export ConstantColoringAlgorithm, StructuredColoringAlgorithm
 export coloring, fast_coloring
 export column_colors, row_colors, ncolors
 export column_groups, row_groups
diff --git a/src/structured.jl b/src/structured.jl
index 78b2db69..24692196 100644
--- a/src/structured.jl
+++ b/src/structured.jl
@@ -1,3 +1,21 @@
+"""
+    StructuredColoringAlgorithm  <: ADTypes.AbstractColoringAlgorithm
+
+Coloring algorithm which leverages specific matrix structures to produce optimal or near-optimal solutions.
+
+The following matrix types are supported:
+
+- From the standard library `LinearAlgebra`: [`Diagonal`](@extref LinearAlgebra.Diagonal), [`Bidiagonal`](@extref LinearAlgebra.Bidiagonal), [`Tridiagonal`](@extref LinearAlgebra.Tridiagonal)
+- From [BandedMatrices.jl](https://github.com/JuliaLinearAlgebra/BandedMatrices.jl): [`BandedMatrix`](@extref BandedMatrices.BandedMatrix)
+
+!!! warning
+    Only `:nonsymmetric` structures with `:row` or `:column` partitions (aka unidirectional Jacobian colorings) are supported by this algorithm at the moment.
+
+!!! tip
+    To request support for a new type of structured matrix, open an issue on the SparseMatrixColorings.jl GitHub repository!
+"""
+struct StructuredColoringAlgorithm <: ADTypes.AbstractColoringAlgorithm end
+
 #=
 This code is partly taken from ArrayInterface.jl
 https://github.com/JuliaArrays/ArrayInterface.jl
@@ -21,7 +39,7 @@ end
 function coloring(
     A::Diagonal,
     ::ColoringProblem{:nonsymmetric,:column},
-    algo::GreedyColoringAlgorithm;
+    ::StructuredColoringAlgorithm;
     kwargs...,
 )
     color = fill(1, size(A, 2))
@@ -32,7 +50,7 @@ end
 function coloring(
     A::Diagonal,
     ::ColoringProblem{:nonsymmetric,:row},
-    algo::GreedyColoringAlgorithm;
+    ::StructuredColoringAlgorithm;
     kwargs...,
 )
     color = fill(1, size(A, 1))
@@ -61,7 +79,7 @@ end
 function coloring(
     A::Bidiagonal,
     ::ColoringProblem{:nonsymmetric,:column},
-    algo::GreedyColoringAlgorithm;
+    ::StructuredColoringAlgorithm;
     kwargs...,
 )
     color = cycle_range(2, size(A, 2))
@@ -72,7 +90,7 @@ end
 function coloring(
     A::Bidiagonal,
     ::ColoringProblem{:nonsymmetric,:row},
-    algo::GreedyColoringAlgorithm;
+    ::StructuredColoringAlgorithm;
     kwargs...,
 )
     color = cycle_range(2, size(A, 1))
@@ -113,7 +131,7 @@ end
 function coloring(
     A::Tridiagonal,
     ::ColoringProblem{:nonsymmetric,:column},
-    algo::GreedyColoringAlgorithm;
+    ::StructuredColoringAlgorithm;
     kwargs...,
 )
     color = cycle_range(3, size(A, 2))
@@ -124,7 +142,7 @@ end
 function coloring(
     A::Tridiagonal,
     ::ColoringProblem{:nonsymmetric,:row},
-    algo::GreedyColoringAlgorithm;
+    ::StructuredColoringAlgorithm;
     kwargs...,
 )
     color = cycle_range(3, size(A, 1))
diff --git a/test/structured.jl b/test/structured.jl
index 23ba967f..d61c7735 100644
--- a/test/structured.jl
+++ b/test/structured.jl
@@ -6,53 +6,65 @@ using SparseMatrixColorings
 using Test
 
 @testset "Diagonal" begin
-    for n in (1, 2, 10, 100)
-        A = Diagonal(rand(n))
-        test_structured_coloring_decompression(A)
+    @testset for algo in [GreedyColoringAlgorithm(), StructuredColoringAlgorithm()]
+        for n in (1, 2, 10, 100)
+            A = Diagonal(rand(n))
+            test_structured_coloring_decompression(A, algo)
+        end
     end
 end;
 
 @testset "Bidiagonal" begin
-    for n in (2, 10, 100)
-        A1 = Bidiagonal(rand(n), rand(n - 1), :U)
-        A2 = Bidiagonal(rand(n), rand(n - 1), :L)
-        test_structured_coloring_decompression(A1)
-        test_structured_coloring_decompression(A2)
+    @testset for algo in [GreedyColoringAlgorithm(), StructuredColoringAlgorithm()]
+        for n in (2, 10, 100)
+            A1 = Bidiagonal(rand(n), rand(n - 1), :U)
+            A2 = Bidiagonal(rand(n), rand(n - 1), :L)
+            test_structured_coloring_decompression(A1, algo)
+            test_structured_coloring_decompression(A2, algo)
+        end
     end
 end;
 
 @testset "Tridiagonal" begin
-    for n in (2, 10, 100)
-        A = Tridiagonal(rand(n - 1), rand(n), rand(n - 1))
-        test_structured_coloring_decompression(A)
+    @testset for algo in [GreedyColoringAlgorithm(), StructuredColoringAlgorithm()]
+        for n in (2, 10, 100)
+            A = Tridiagonal(rand(n - 1), rand(n), rand(n - 1))
+            test_structured_coloring_decompression(A, algo)
+        end
     end
 end;
 
 @testset "BandedMatrices" begin
-    @testset for (m, n) in [(10, 20), (20, 10)], l in 0:5, u in 0:5
-        A = brand(m, n, l, u)
-        test_structured_coloring_decompression(A)
+    @testset for algo in [GreedyColoringAlgorithm(), StructuredColoringAlgorithm()]
+        @testset for (m, n) in [(10, 20), (20, 10)], l in 0:5, u in 0:5
+            A = brand(m, n, l, u)
+            test_structured_coloring_decompression(A, algo)
+        end
     end
 end;
 
 @testset "BlockBandedMatrices" begin
-    for (mb, nb) in [(10, 20), (20, 10)], lb in 0:3, ub in 0:3, _ in 1:10
-        rows = rand(1:5, mb)
-        cols = rand(1:5, nb)
-        A = BlockBandedMatrix{Float64}(rand(sum(rows), sum(cols)), rows, cols, (lb, ub))
-        test_structured_coloring_decompression(A)
+    @testset for algo in [GreedyColoringAlgorithm(), StructuredColoringAlgorithm()]
+        for (mb, nb) in [(10, 20), (20, 10)], lb in 0:3, ub in 0:3, _ in 1:10
+            rows = rand(1:5, mb)
+            cols = rand(1:5, nb)
+            A = BlockBandedMatrix{Float64}(rand(sum(rows), sum(cols)), rows, cols, (lb, ub))
+            test_structured_coloring_decompression(A, algo)
+        end
     end
 end;
 
 @testset "BandedBlockBandedMatrices" begin
-    for (mb, nb) in [(10, 20), (20, 10)], lb in 0:3, ub in 0:3, _ in 1:10
-        rows = rand(5:10, mb)
-        cols = rand(5:10, nb)
-        λ = rand(0:5)
-        μ = rand(0:5)
-        A = BandedBlockBandedMatrix{Float64}(
-            rand(sum(rows), sum(cols)), rows, cols, (lb, ub), (λ, μ)
-        )
-        test_structured_coloring_decompression(A)
+    @testset for algo in [GreedyColoringAlgorithm(), StructuredColoringAlgorithm()]
+        for (mb, nb) in [(10, 20), (20, 10)], lb in 0:3, ub in 0:3, _ in 1:10
+            rows = rand(5:10, mb)
+            cols = rand(5:10, nb)
+            λ = rand(0:5)
+            μ = rand(0:5)
+            A = BandedBlockBandedMatrix{Float64}(
+                rand(sum(rows), sum(cols)), rows, cols, (lb, ub), (λ, μ)
+            )
+            test_structured_coloring_decompression(A, algo)
+        end
     end
 end;
diff --git a/test/utils.jl b/test/utils.jl
index 7c2329b7..200d07d5 100644
--- a/test/utils.jl
+++ b/test/utils.jl
@@ -218,10 +218,11 @@ function test_bicoloring_decompression(
     end
 end
 
-function test_structured_coloring_decompression(A::AbstractMatrix)
+function test_structured_coloring_decompression(
+    A::AbstractMatrix, algo=StructuredColoringAlgorithm()
+)
     column_problem = ColoringProblem(; structure=:nonsymmetric, partition=:column)
     row_problem = ColoringProblem(; structure=:nonsymmetric, partition=:row)
-    algo = GreedyColoringAlgorithm()
 
     # Column
     result = coloring(A, column_problem, algo)
@@ -231,7 +232,9 @@ function test_structured_coloring_decompression(A::AbstractMatrix)
     @test D == A
     @test nameof(typeof(D)) == nameof(typeof(A))
     @test structurally_orthogonal_columns(A, color)
-    @test color == ArrayInterface.matrix_colors(A)
+    if algo isa StructuredColoringAlgorithm
+        @test color == ArrayInterface.matrix_colors(A)
+    end
 
     # Row
     result = coloring(A, row_problem, algo)

From 96f5662d064ea71f9ba8326c1f7192f4d2b41f98 Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Tue, 1 Jul 2025 16:31:13 +0200
Subject: [PATCH 32/33] Fix doc lunk

---
 docs/make.jl      | 1 -
 src/structured.jl | 2 +-
 2 files changed, 1 insertion(+), 2 deletions(-)

diff --git a/docs/make.jl b/docs/make.jl
index 718b8f15..46536a16 100644
--- a/docs/make.jl
+++ b/docs/make.jl
@@ -5,7 +5,6 @@ using SparseMatrixColorings
 links = InterLinks(
     "ADTypes" => "https://sciml.github.io/ADTypes.jl/stable/",
     "BandedMatrices" => "https://julialinearalgebra.github.io/BandedMatrices.jl/stable/",
-    "LinearAlgebra" => "https://docs.julialang.org/en/v1/stdlib/LinearAlgebra/",
 )
 
 cp(joinpath(@__DIR__, "..", "README.md"), joinpath(@__DIR__, "src", "index.md"); force=true)
diff --git a/src/structured.jl b/src/structured.jl
index 24692196..812256bf 100644
--- a/src/structured.jl
+++ b/src/structured.jl
@@ -5,7 +5,7 @@ Coloring algorithm which leverages specific matrix structures to produce optimal
 
 The following matrix types are supported:
 
-- From the standard library `LinearAlgebra`: [`Diagonal`](@extref LinearAlgebra.Diagonal), [`Bidiagonal`](@extref LinearAlgebra.Bidiagonal), [`Tridiagonal`](@extref LinearAlgebra.Tridiagonal)
+- From the standard library `LinearAlgebra`: `Diagonal`, `Bidiagonal`, `Tridiagonal`
 - From [BandedMatrices.jl](https://github.com/JuliaLinearAlgebra/BandedMatrices.jl): [`BandedMatrix`](@extref BandedMatrices.BandedMatrix)
 
 !!! warning

From b24c181985a101be4682e397b4c9d3cbea29f7f7 Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Tue, 1 Jul 2025 17:00:30 +0200
Subject: [PATCH 33/33] No optimized decompression

---
 ext/SparseMatrixColoringsBandedMatricesExt.jl | 22 -------------------
 1 file changed, 22 deletions(-)

diff --git a/ext/SparseMatrixColoringsBandedMatricesExt.jl b/ext/SparseMatrixColoringsBandedMatricesExt.jl
index e4dd32f7..05d9677c 100644
--- a/ext/SparseMatrixColoringsBandedMatricesExt.jl
+++ b/ext/SparseMatrixColoringsBandedMatricesExt.jl
@@ -42,26 +42,4 @@ function SMC.coloring(
     return RowColoringResult(A, bg, color)
 end
 
-function SMC.decompress!(A::BandedMatrix, B::AbstractMatrix, result::ColumnColoringResult)
-    color = column_colors(result)
-    for j in axes(A, 2)
-        c = color[j]
-        for i in colrange(A, j)
-            A[i, j] = B[i, c]
-        end
-    end
-    return A
-end
-
-function SMC.decompress!(A::BandedMatrix, B::AbstractMatrix, result::RowColoringResult)
-    color = row_colors(result)
-    for i in axes(A, 1)
-        c = color[i]
-        for j in rowrange(A, i)
-            A[i, j] = B[c, j]
-        end
-    end
-    return A
-end
-
 end