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/36] 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/36] 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/36] 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/36] 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/36] 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/36] 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/36] 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/36] 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/36] 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/36] 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/36] 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/36] 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/36] 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/36] 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/36] 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/36] 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/36] 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/36] 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/36] 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/36] 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/36] 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/36] 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/36] 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/36] 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/36] 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/36] 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/36] 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/36] 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/36] 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/36] 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/36] 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/36] 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/36] 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

From 135e425ef18f19480f98e5519290860a1569b219 Mon Sep 17 00:00:00 2001
From: Qingyu Qu <52615090+ErikQQY@users.noreply.github.com>
Date: Tue, 28 Apr 2026 20:30:44 +0800
Subject: [PATCH 34/36] Merge main into gd/structured3 (#316)

* Merge main into gd/structured3

* Add the missing docs
---
 .buildkite/pipeline.yml                  |  18 ++
 .github/workflows/Benchmark.yml          |  10 +-
 .github/workflows/CompatHelper.yml       |   2 +-
 .github/workflows/Documentation.yml      |   6 +-
 .github/workflows/Format.yml             |  12 +
 .github/workflows/Test-GPU.yml           |  46 ----
 .github/workflows/Test.yml               |  14 +-
 CITATION.cff                             |   2 +-
 Project.toml                             |  22 +-
 README.md                                |  10 +-
 docs/make.jl                             |   4 +-
 docs/src/api.md                          |   1 +
 docs/src/dev.md                          |   6 +-
 docs/src/tutorial.md                     |   2 +-
 ext/SparseMatrixColoringsCUDAExt.jl      |  18 +-
 ext/SparseMatrixColoringsColorsExt.jl    |   5 +-
 ext/SparseMatrixColoringsGPUArraysExt.jl |  29 +++
 ext/SparseMatrixColoringsJuMPExt.jl      |  80 +++++++
 src/SparseMatrixColorings.jl             |   3 +
 src/adtypes.jl                           |  46 ++--
 src/check.jl                             | 250 ++++++++++++++++++++-
 src/coloring.jl                          | 271 ++++++++---------------
 src/constant.jl                          |  91 ++++----
 src/decompression.jl                     | 124 +++++++----
 src/graph.jl                             |  58 ++---
 src/interface.jl                         | 153 +++++++++----
 src/matrices.jl                          |  59 +++--
 src/optimal.jl                           |  34 +++
 src/order.jl                             |   2 +-
 src/postprocessing.jl                    | 168 ++++++++++++++
 src/result.jl                            |  28 ++-
 test/Project.toml                        |   5 +-
 test/adtypes.jl                          |  61 +++--
 test/allocations.jl                      |   2 +-
 test/check.jl                            | 135 +++++++++++
 test/constant.jl                         |  39 +++-
 test/cuda.jl                             |   4 +-
 test/graph.jl                            |   5 +-
 test/matrices.jl                         |  41 +++-
 test/optimal.jl                          |  46 ++++
 test/order.jl                            | 114 +++++++++-
 test/random.jl                           |   8 +-
 test/result.jl                           |   3 +-
 test/runtests.jl                         |  20 +-
 test/structured.jl                       |  17 ++
 test/type_stability.jl                   |  10 +
 test/utils.jl                            |  60 ++++-
 47 files changed, 1601 insertions(+), 543 deletions(-)
 create mode 100644 .buildkite/pipeline.yml
 create mode 100644 .github/workflows/Format.yml
 delete mode 100644 .github/workflows/Test-GPU.yml
 create mode 100644 ext/SparseMatrixColoringsGPUArraysExt.jl
 create mode 100644 ext/SparseMatrixColoringsJuMPExt.jl
 create mode 100644 src/optimal.jl
 create mode 100644 src/postprocessing.jl
 create mode 100644 test/optimal.jl

diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
new file mode 100644
index 00000000..a1a04cb0
--- /dev/null
+++ b/.buildkite/pipeline.yml
@@ -0,0 +1,18 @@
+env:
+  SECRET_CODECOV_TOKEN: "l7PxpMgxHEiz2l3UDDj5122Xirv5du1QFQiTf09YISC8VOz/nXeQC2yYY6iTMbccQG3xrumfWke9vtLbAc9LPH75kw3zQODuWMbz0eGSf9wtw7foqO6KAo5neMTeZJ0ZFgFMa2Q89T5SNIQsi4b5Zs7BVto6qB9Z/3QEs/BpsR25cYkY4Y6JBU6XuqZ6GRAc6BtlB4OmBEp3BBsXatx6y64zF0qbp4rocmcBQEoeQkcXtxf0dfA0KNAHpweWPrzIAMZ0aUYp6iEikxwLY5TjVFhOIcpVXUlxIdhl2qaaDF6b6WlgGXiGLpBjsLQfvlOgIlXH59Ddg0IVxboF2a37OA==;U2FsdGVkX19MveWMe1aasYZoJwaeiKO5XqwMZ/utOFj7h1CE9UX4nduBMXTyJ77tpNmCBYsQNJ/PadbX2224hQ=="
+
+steps:
+  - label: "Julia v1"
+    plugins:
+      - JuliaCI/julia#v1:
+          version: "1"
+      - JuliaCI/julia-test#v1: ~
+      - JuliaCI/julia-coverage#v1:
+          codecov: true
+    agents:
+      queue: "juliagpu"
+      cuda: "*"
+    if: build.message !~ /\[skip tests\]/
+    timeout_in_minutes: 60
+    env:
+      JULIA_SMC_TEST_GROUP: 'GPU'
\ No newline at end of file
diff --git a/.github/workflows/Benchmark.yml b/.github/workflows/Benchmark.yml
index fda51ed6..4a4a444b 100644
--- a/.github/workflows/Benchmark.yml
+++ b/.github/workflows/Benchmark.yml
@@ -12,11 +12,11 @@ jobs:
         runs-on: ubuntu-latest
         if: contains(github.event.pull_request.labels.*.name, 'benchmark')
         steps:
-            - uses: actions/checkout@v4
-            - uses: julia-actions/setup-julia@v2
+            - uses: actions/checkout@v6
+            - uses: julia-actions/setup-julia@v3
               with:
                 version: "1"
-            - uses: julia-actions/cache@v2
+            - uses: julia-actions/cache@v3
             - name: Extract Package Name from Project.toml
               id: extract-package-name
               run: |
@@ -46,14 +46,14 @@ jobs:
                 echo '' >> body.md
                 echo '' >> body.md
             - name: Find Comment
-              uses: peter-evans/find-comment@v3
+              uses: peter-evans/find-comment@v4
               id: fcbenchmark
               with:
                 issue-number: ${{ github.event.pull_request.number }}
                 comment-author: 'github-actions[bot]'
                 body-includes: Benchmark Results
             - name: Comment on PR
-              uses: peter-evans/create-or-update-comment@v4
+              uses: peter-evans/create-or-update-comment@v5
               with:
                 comment-id: ${{ steps.fcbenchmark.outputs.comment-id }}
                 issue-number: ${{ github.event.pull_request.number }}
diff --git a/.github/workflows/CompatHelper.yml b/.github/workflows/CompatHelper.yml
index 1f74cd19..492a6c0b 100644
--- a/.github/workflows/CompatHelper.yml
+++ b/.github/workflows/CompatHelper.yml
@@ -15,7 +15,7 @@ jobs:
         run: which julia
         continue-on-error: true
       - name: Install Julia, but only if it is not already available in the PATH
-        uses: julia-actions/setup-julia@v2
+        uses: julia-actions/setup-julia@v3
         with:
           version: '1'
           arch: ${{ runner.arch }}
diff --git a/.github/workflows/Documentation.yml b/.github/workflows/Documentation.yml
index a8d1110e..2f2e57df 100644
--- a/.github/workflows/Documentation.yml
+++ b/.github/workflows/Documentation.yml
@@ -19,11 +19,11 @@ jobs:
       statuses: write
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v4
-      - uses: julia-actions/setup-julia@v2
+      - uses: actions/checkout@v6
+      - uses: julia-actions/setup-julia@v3
         with:
           version: '1'
-      - uses: julia-actions/cache@v2
+      - uses: julia-actions/cache@v3
       - name: Install dependencies
         run: julia --project=docs/ -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()'
       - name: Build and deploy
diff --git a/.github/workflows/Format.yml b/.github/workflows/Format.yml
new file mode 100644
index 00000000..dd2b2c50
--- /dev/null
+++ b/.github/workflows/Format.yml
@@ -0,0 +1,12 @@
+name: Format suggestions
+on:
+  pull_request:
+    # this argument is not required if you don't use the `suggestion-label` input
+    types: [opened, reopened, synchronize, labeled, unlabeled]
+jobs:
+  code-style:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: julia-actions/julia-format@v4
+        with:
+          version: '1' # Set `version` to '1.0.54' if you need to use JuliaFormatter.jl v1.0.54 (default: '1')
diff --git a/.github/workflows/Test-GPU.yml b/.github/workflows/Test-GPU.yml
deleted file mode 100644
index cd67c6e1..00000000
--- a/.github/workflows/Test-GPU.yml
+++ /dev/null
@@ -1,46 +0,0 @@
-name: Test-GPU
-
-on:
-  push:
-    branches:
-      - main
-  pull_request:
-
-concurrency:
-  group: ${{ github.workflow }}-${{ github.ref }}
-  cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }}
-
-# needed to allow julia-actions/cache to delete old caches that it has created
-permissions:
-  actions: write
-  contents: read
-
-jobs:
-  test:
-    runs-on: self-hosted
-    env:
-      CUDA_VISIBLE_DEVICES: 1
-      JULIA_DEPOT_PATH: /scratch/github-actions/julia_depot_alexis
-      JULIA_SMC_TEST_GROUP: "GPU"
-    strategy:
-      matrix:
-        julia-version: ['1.10', '1']
-    steps:
-      - uses: actions/checkout@v4
-      - uses: julia-actions/setup-julia@v2
-        with:
-          version: ${{ matrix.julia-version }}
-          arch: x64
-      - uses: julia-actions/julia-downgrade-compat@v1
-        if: ${{ matrix.version == '1.10' }}
-        with:
-          skip: LinearAlgebra, Random, SparseArrays
-      - uses: julia-actions/cache@v2
-      - uses: julia-actions/julia-buildpkg@v1
-      - uses: julia-actions/julia-runtest@v1
-      - uses: julia-actions/julia-processcoverage@v1
-      - uses: codecov/codecov-action@v5
-        with:
-          files: lcov.info
-          token: ${{ secrets.CODECOV_TOKEN }}
-          fail_ci_if_error: false
diff --git a/.github/workflows/Test.yml b/.github/workflows/Test.yml
index 1f0046a1..6e8c6af5 100644
--- a/.github/workflows/Test.yml
+++ b/.github/workflows/Test.yml
@@ -20,24 +20,24 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        julia-version: ['1.10', '1']
+        julia-version: ['1.10', '1.11', '1.12']
 
     steps:
-      - uses: actions/checkout@v4
-      - uses: julia-actions/setup-julia@v2
+      - uses: actions/checkout@v6
+      - uses: julia-actions/setup-julia@v3
         with:
           version: ${{ matrix.julia-version }}
           arch: x64
-      - uses: julia-actions/julia-downgrade-compat@v1
+      - uses: julia-actions/julia-downgrade-compat@v2
         if: ${{ matrix.version == '1.10' }}
         with:
           skip: LinearAlgebra, Random, SparseArrays
-      - uses: julia-actions/cache@v2
+      - uses: julia-actions/cache@v3
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1
       - uses: julia-actions/julia-processcoverage@v1
-      - uses: codecov/codecov-action@v5
+      - uses: codecov/codecov-action@v6
         with:
           files: lcov.info
           token: ${{ secrets.CODECOV_TOKEN }}
-          fail_ci_if_error: false
\ No newline at end of file
+          fail_ci_if_error: false
diff --git a/CITATION.cff b/CITATION.cff
index c2c0a8fb..a36bc54e 100644
--- a/CITATION.cff
+++ b/CITATION.cff
@@ -18,7 +18,7 @@ identifiers:
   - type: doi
     value: 10.5281/zenodo.11314275
     description: Zenodo
-repository-code: 'https://github.com/gdalle/SparseMatrixColorings.jl'
+repository-code: 'https://github.com/JuliaDiff/SparseMatrixColorings.jl'
 abstract: >-
   Coloring algorithms for sparse Jacobian and Hessian
   matrices
diff --git a/Project.toml b/Project.toml
index 662685ed..72a5fad2 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "SparseMatrixColorings"
 uuid = "0a514795-09f3-496d-8182-132a7b665d35"
+version = "0.4.27"
 authors = ["Guillaume Dalle", "Alexis Montoison"]
-version = "0.4.22"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -18,25 +18,35 @@ BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 CliqueTrees = "60701a23-6482-424a-84db-faee86b9b1f8"
 Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
+GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
+JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
+MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
+cuSPARSE = "b26da814-b3bc-49ef-b0ee-c816305aa060"
 
 [extensions]
 SparseMatrixColoringsBandedMatricesExt = "BandedMatrices"
 SparseMatrixColoringsBlockBandedMatricesExt = ["BlockArrays", "BlockBandedMatrices"]
-SparseMatrixColoringsCUDAExt = "CUDA"
+SparseMatrixColoringsCUDAExt = ["CUDA", "cuSPARSE"]
 SparseMatrixColoringsCliqueTreesExt = "CliqueTrees"
 SparseMatrixColoringsColorsExt = "Colors"
+SparseMatrixColoringsGPUArraysExt = "GPUArrays"
+SparseMatrixColoringsJuMPExt = ["JuMP", "MathOptInterface"]
 
 [compat]
 ADTypes = "1.2.1"
 BandedMatrices = "1.9.4"
 BlockArrays = "1.6.3"
 BlockBandedMatrices = "0.13.1"
-CUDA = "5.8.2"
+CUDA = "6.0.0"
 CliqueTrees = "1"
 Colors = "0.12.11, 0.13"
 DocStringExtensions = "0.8,0.9"
-LinearAlgebra = "<0.0.1, 1"
+GPUArrays = "11.5.0"
+JuMP = "1.29.1"
+LinearAlgebra = "1"
+MathOptInterface = "1.45.0"
 PrecompileTools = "1.2.1"
-Random = "<0.0.1, 1"
-SparseArrays = "<0.0.1, 1"
+Random = "1"
+SparseArrays = "1"
+cuSPARSE = "6.0.0"
 julia = "1.10"
diff --git a/README.md b/README.md
index 66b10dd3..b37b0ca8 100644
--- a/README.md
+++ b/README.md
@@ -1,9 +1,11 @@
 # SparseMatrixColorings.jl
 
-[![Build Status](https://github.com/gdalle/SparseMatrixColorings.jl/actions/workflows/Test.yml/badge.svg?branch=main)](https://github.com/gdalle/SparseMatrixColorings.jl/actions/workflows/Test.yml?query=branch%3Amain)
-[![Stable Documentation](https://img.shields.io/badge/docs-stable-blue.svg)](https://gdalle.github.io/SparseMatrixColorings.jl/stable/)
-[![Dev Documentation](https://img.shields.io/badge/docs-dev-blue.svg)](https://gdalle.github.io/SparseMatrixColorings.jl/dev/)
-[![Coverage](https://codecov.io/gh/gdalle/SparseMatrixColorings.jl/branch/main/graph/badge.svg)](https://app.codecov.io/gh/gdalle/SparseMatrixColorings.jl)
+[![Build Status](https://github.com/juliadiff/SparseMatrixColorings.jl/actions/workflows/Test.yml/badge.svg?branch=main)](https://github.com/juliadiff/SparseMatrixColorings.jl/actions/workflows/Test.yml?query=branch%3Amain)
+[![GPU build status](https://badge.buildkite.com/7d8ed289d7bdb5a25ae48b2c778a202ce4990b7ee558cdfef8.svg?branch=main)](https://buildkite.com/julialang/sparsematrixcolorings-dot-jl)
+[![Coverage](https://codecov.io/gh/juliadiff/SparseMatrixColorings.jl/branch/main/graph/badge.svg)](https://app.codecov.io/gh/juliadiff/SparseMatrixColorings.jl)
+
+[![Stable Documentation](https://img.shields.io/badge/docs-stable-blue.svg)](https://juliadiff.org/SparseMatrixColorings.jl/stable/)
+[![Dev Documentation](https://img.shields.io/badge/docs-dev-blue.svg)](https://juliadiff.org/SparseMatrixColorings.jl/dev/)
 [![Code Style: Blue](https://img.shields.io/badge/code%20style-blue-4495d1.svg)](https://github.com/JuliaDiff/BlueStyle)
 [![arXiv](https://img.shields.io/badge/arXiv-2505.07308-b31b1b.svg)](https://arxiv.org/abs/2505.07308)
 [![DOI](https://zenodo.org/badge/801999408.svg)](https://zenodo.org/doi/10.5281/zenodo.11314275)
diff --git a/docs/make.jl b/docs/make.jl
index 46536a16..bbf3f859 100644
--- a/docs/make.jl
+++ b/docs/make.jl
@@ -24,5 +24,7 @@ makedocs(;
 )
 
 deploydocs(;
-    repo="github.com/gdalle/SparseMatrixColorings.jl", push_preview=true, devbranch="main"
+    repo="github.com/JuliaDiff/SparseMatrixColorings.jl",
+    push_preview=true,
+    devbranch="main",
 )
diff --git a/docs/src/api.md b/docs/src/api.md
index 4fd6d28e..640198d3 100644
--- a/docs/src/api.md
+++ b/docs/src/api.md
@@ -24,6 +24,7 @@ ColoringProblem
 ```@docs
 GreedyColoringAlgorithm
 ConstantColoringAlgorithm
+OptimalColoringAlgorithm
 StructuredColoringAlgorithm
 ```
 
diff --git a/docs/src/dev.md b/docs/src/dev.md
index 0f2510f4..cd294a17 100644
--- a/docs/src/dev.md
+++ b/docs/src/dev.md
@@ -41,6 +41,7 @@ SparseMatrixColorings.TreeSetColoringResult
 SparseMatrixColorings.LinearSystemColoringResult
 SparseMatrixColorings.BicoloringResult
 SparseMatrixColorings.remap_colors
+SparseMatrixColorings.decompress_csc!
 ```
 
 ## Testing
@@ -50,6 +51,9 @@ SparseMatrixColorings.directly_recoverable_columns
 SparseMatrixColorings.symmetrically_orthogonal_columns
 SparseMatrixColorings.structurally_orthogonal_columns
 SparseMatrixColorings.structurally_biorthogonal
+SparseMatrixColorings.substitutable_columns
+SparseMatrixColorings.substitutable_bidirectional
+SparseMatrixColorings.rank_nonzeros_from_trees
 SparseMatrixColorings.valid_dynamic_order
 ```
 
@@ -58,7 +62,7 @@ SparseMatrixColorings.valid_dynamic_order
 ```@docs
 SparseMatrixColorings.respectful_similar
 SparseMatrixColorings.matrix_versions
-SparseMatrixColorings.same_pattern
+SparseMatrixColorings.compatible_pattern
 ```
 
 ## Visualization
diff --git a/docs/src/tutorial.md b/docs/src/tutorial.md
index fec54f13..6e803511 100644
--- a/docs/src/tutorial.md
+++ b/docs/src/tutorial.md
@@ -31,7 +31,7 @@ problem = ColoringProblem()
 
 The algorithm defines how you want to solve it. It can be either a [`GreedyColoringAlgorithm`](@ref) or a [`ConstantColoringAlgorithm`](@ref). For `GreedyColoringAlgorithm`, you can select options such as
 
-- the order in which vertices are processed (a subtype of [`AbstractOrder`](@ref SparseMatrixColorings.AbstractOrder))
+- the order in which vertices are processed (a subtype of [`AbstractOrder`](@ref SparseMatrixColorings.AbstractOrder) , or a tuple of such objects)
 - the type of decompression you want (`:direct` or `:substitution`)
 
 ```@example tutorial
diff --git a/ext/SparseMatrixColoringsCUDAExt.jl b/ext/SparseMatrixColoringsCUDAExt.jl
index 2750964d..ed33eece 100644
--- a/ext/SparseMatrixColoringsCUDAExt.jl
+++ b/ext/SparseMatrixColoringsCUDAExt.jl
@@ -3,23 +3,7 @@ module SparseMatrixColoringsCUDAExt
 import SparseMatrixColorings as SMC
 using SparseArrays: SparseMatrixCSC, rowvals, nnz, nzrange
 using CUDA: CuVector, CuMatrix
-using CUDA.CUSPARSE: AbstractCuSparseMatrix, CuSparseMatrixCSC, CuSparseMatrixCSR
-
-SMC.matrix_versions(A::AbstractCuSparseMatrix) = (A,)
-
-## Compression (slow, through CPU)
-
-function SMC.compress(
-    A::AbstractCuSparseMatrix, result::SMC.AbstractColoringResult{structure,:column}
-) where {structure}
-    return CuMatrix(SMC.compress(SparseMatrixCSC(A), result))
-end
-
-function SMC.compress(
-    A::AbstractCuSparseMatrix, result::SMC.AbstractColoringResult{structure,:row}
-) where {structure}
-    return CuMatrix(SMC.compress(SparseMatrixCSC(A), result))
-end
+using cuSPARSE: AbstractCuSparseMatrix, CuSparseMatrixCSC, CuSparseMatrixCSR
 
 ## CSC Result
 
diff --git a/ext/SparseMatrixColoringsColorsExt.jl b/ext/SparseMatrixColoringsColorsExt.jl
index 7f844e89..c934b5e9 100644
--- a/ext/SparseMatrixColoringsColorsExt.jl
+++ b/ext/SparseMatrixColoringsColorsExt.jl
@@ -271,9 +271,8 @@ function show_colors!(
     A_ccolor_indices = mod1.(column_colors(res), length(colorscheme))
     A_rcolor_indices = mod1.(row_shift .+ row_colors(res), length(colorscheme))
     B_ccolor_indices = mod1.(1:maximum(column_colors(res)), length(colorscheme))
-    B_rcolor_indices = mod1.(
-        (row_shift + 1):(row_shift + maximum(row_colors(res))), length(colorscheme)
-    )
+    B_rcolor_indices =
+        mod1.((row_shift + 1):(row_shift + maximum(row_colors(res))), length(colorscheme))
     A_ccolors = colorscheme[A_ccolor_indices]
     A_rcolors = colorscheme[A_rcolor_indices]
     B_ccolors = colorscheme[B_ccolor_indices]
diff --git a/ext/SparseMatrixColoringsGPUArraysExt.jl b/ext/SparseMatrixColoringsGPUArraysExt.jl
new file mode 100644
index 00000000..37574e73
--- /dev/null
+++ b/ext/SparseMatrixColoringsGPUArraysExt.jl
@@ -0,0 +1,29 @@
+module SparseMatrixColoringsGPUArraysExt
+
+using GPUArrays: AbstractGPUSparseMatrix, dense_array_type
+using SparseArrays: SparseMatrixCSC
+import SparseMatrixColorings as SMC
+
+SMC.matrix_versions(A::AbstractGPUSparseMatrix) = (A,)
+
+## Compression (slow, through CPU)
+
+function SMC.compress(
+    A::AbstractGPUSparseMatrix, result::SMC.AbstractColoringResult{structure,:column}
+) where {structure}
+    A_cpu = SparseMatrixCSC(A)
+    B_cpu = SMC.compress(A_cpu, result)
+    B = dense_array_type(A)(B_cpu)
+    return B
+end
+
+function SMC.compress(
+    A::AbstractGPUSparseMatrix, result::SMC.AbstractColoringResult{structure,:row}
+) where {structure}
+    A_cpu = SparseMatrixCSC(A)
+    B_cpu = SMC.compress(A_cpu, result)
+    B = dense_array_type(A)(B_cpu)
+    return B
+end
+
+end
diff --git a/ext/SparseMatrixColoringsJuMPExt.jl b/ext/SparseMatrixColoringsJuMPExt.jl
new file mode 100644
index 00000000..77b74dbd
--- /dev/null
+++ b/ext/SparseMatrixColoringsJuMPExt.jl
@@ -0,0 +1,80 @@
+module SparseMatrixColoringsJuMPExt
+
+using ADTypes: ADTypes
+using JuMP:
+    Model,
+    assert_is_solved_and_feasible,
+    optimize!,
+    primal_status,
+    set_silent,
+    set_start_value,
+    value,
+    @variable,
+    @constraint,
+    @objective
+using JuMP
+import MathOptInterface as MOI
+using SparseMatrixColorings:
+    BipartiteGraph, OptimalColoringAlgorithm, nb_vertices, neighbors, pattern, vertices
+
+function optimal_distance2_coloring(
+    bg::BipartiteGraph,
+    ::Val{side},
+    optimizer::O;
+    silent::Bool=true,
+    assert_solved::Bool=true,
+) where {side,O}
+    other_side = 3 - side
+    n = nb_vertices(bg, Val(side))
+    model = Model(optimizer)
+    silent && set_silent(model)
+    # one variable per vertex to color, removing some renumbering symmetries
+    @variable(model, 1 <= color[i=1:n] <= i, Int)
+    # one variable to count the number of distinct colors
+    @variable(model, ncolors, Int)
+    @constraint(model, [ncolors; color] in MOI.CountDistinct(n + 1))
+    # distance-2 coloring: neighbors of the same vertex must have distinct colors
+    for i in vertices(bg, Val(other_side))
+        neigh = neighbors(bg, Val(other_side), i)
+        @constraint(model, color[neigh] in MOI.AllDifferent(length(neigh)))
+    end
+    # minimize the number of distinct colors (can't use maximum because they are not necessarily numbered contiguously)
+    @objective(model, Min, ncolors)
+    # actual solving step where time is spent
+    optimize!(model)
+    if assert_solved
+        # assert feasibility and optimality
+        assert_is_solved_and_feasible(model)
+    else
+        # only assert feasibility
+        @assert primal_status(model) == MOI.FEASIBLE_POINT
+    end
+    # native solver solutions are floating point numbers
+    color_int = round.(Int, value.(color))
+    # remap to 1:cmax in case they are not contiguous
+    true_ncolors = 0
+    remap = fill(0, maximum(color_int))
+    for c in color_int
+        if remap[c] == 0
+            true_ncolors += 1
+            remap[c] = true_ncolors
+        end
+    end
+    return remap[color_int]
+end
+
+function ADTypes.column_coloring(A::AbstractMatrix, algo::OptimalColoringAlgorithm)
+    bg = BipartiteGraph(A)
+    return optimal_distance2_coloring(
+        bg, Val(2), algo.optimizer; algo.silent, algo.assert_solved
+    )
+end
+
+function ADTypes.row_coloring(A::AbstractMatrix, algo::OptimalColoringAlgorithm)
+    bg = BipartiteGraph(A)
+    return optimal_distance2_coloring(
+        bg, Val(1), algo.optimizer; algo.silent, algo.assert_solved
+    )
+end
+
+end
diff --git a/src/SparseMatrixColorings.jl b/src/SparseMatrixColorings.jl
index 532ef192..74486d80 100644
--- a/src/SparseMatrixColorings.jl
+++ b/src/SparseMatrixColorings.jl
@@ -49,6 +49,7 @@ include("graph.jl")
 include("forest.jl")
 include("order.jl")
 include("coloring.jl")
+include("postprocessing.jl")
 include("result.jl")
 include("matrices.jl")
 include("interface.jl")
@@ -59,6 +60,7 @@ include("structured.jl")
 include("check.jl")
 include("examples.jl")
 include("show_colors.jl")
+include("optimal.jl")
 
 include("precompile.jl")
 
@@ -67,6 +69,7 @@ export DynamicDegreeBasedOrder, SmallestLast, IncidenceDegree, DynamicLargestFir
 export PerfectEliminationOrder
 export ColoringProblem, GreedyColoringAlgorithm, AbstractColoringResult
 export ConstantColoringAlgorithm, StructuredColoringAlgorithm
+export OptimalColoringAlgorithm
 export coloring, fast_coloring
 export column_colors, row_colors, ncolors
 export column_groups, row_groups
diff --git a/src/adtypes.jl b/src/adtypes.jl
index 7500b1e2..7c464433 100644
--- a/src/adtypes.jl
+++ b/src/adtypes.jl
@@ -1,21 +1,33 @@
-function coloring(
-    A::AbstractMatrix,
-    ::ColoringProblem{:nonsymmetric,:column},
-    algo::ADTypes.NoColoringAlgorithm;
-    kwargs...,
-)
-    bg = BipartiteGraph(A)
-    color = convert(Vector{eltype(bg)}, ADTypes.column_coloring(A, algo))
-    return ColumnColoringResult(A, bg, color)
-end
+## From ADTypes to SMC
 
 function coloring(
     A::AbstractMatrix,
-    ::ColoringProblem{:nonsymmetric,:row},
-    algo::ADTypes.NoColoringAlgorithm;
-    kwargs...,
-)
-    bg = BipartiteGraph(A)
-    color = convert(Vector{eltype(bg)}, ADTypes.row_coloring(A, algo))
-    return RowColoringResult(A, bg, color)
+    problem::ColoringProblem{structure,partition},
+    algo::ADTypes.AbstractColoringAlgorithm;
+    decompression_eltype::Type{R}=Float64,
+    symmetric_pattern::Bool=false,
+) where {structure,partition,R}
+    symmetric_pattern = symmetric_pattern || A isa Union{Symmetric,Hermitian}
+    if structure == :nonsymmetric
+        if partition == :column
+            forced_colors = ADTypes.column_coloring(A, algo)
+        elseif partition == :row
+            forced_colors = ADTypes.row_coloring(A, algo)
+        else
+            # TODO: improve once https://github.com/SciML/ADTypes.jl/issues/69 is done
+            A_and_Aᵀ, _ = bidirectional_pattern(A; symmetric_pattern)
+            forced_colors = ADTypes.symmetric_coloring(A_and_Aᵀ, algo)
+        end
+    else
+        forced_colors = ADTypes.symmetric_coloring(A, algo)
+    end
+    return _coloring(
+        WithResult(),
+        A,
+        problem,
+        GreedyColoringAlgorithm(),
+        R,
+        symmetric_pattern;
+        forced_colors,
+    )
 end
diff --git a/src/check.jl b/src/check.jl
index f47f0958..db5c0a34 100644
--- a/src/check.jl
+++ b/src/check.jl
@@ -86,11 +86,15 @@ A partition of the columns of a symmetric matrix `A` is _symmetrically orthogona
 1. the group containing the column `A[:, j]` has no other column with a nonzero in row `i`
 2. the group containing the column `A[:, i]` has no other column with a nonzero in row `j`
 
+It is equivalent to a __star coloring__.
+
 !!! warning
     This function is not coded with efficiency in mind, it is designed for small-scale tests.
 
 # References
 
+> [_On the Estimation of Sparse Hessian Matrices_](https://doi.org/10.1137/0716078), Powell and Toint (1979)
+> [_Estimation of sparse hessian matrices and graph coloring problems_](https://doi.org/10.1007/BF02612334), Coleman and Moré (1984)
 > [_What Color Is Your Jacobian? Graph Coloring for Computing Derivatives_](https://epubs.siam.org/doi/10.1137/S0036144504444711), Gebremedhin et al. (2005)
 """
 function symmetrically_orthogonal_columns(
@@ -102,7 +106,7 @@ function symmetrically_orthogonal_columns(
     end
     issymmetric(A) || return false
     group = group_by_color(color)
-    for i in axes(A, 2), j in axes(A, 2)
+    for i in axes(A, 1), j in axes(A, 2)
         iszero(A[i, j]) && continue
         ci, cj = color[i], color[j]
         check = _bilateral_check(
@@ -126,6 +130,8 @@ A bipartition of the rows and columns of a matrix `A` is _structurally biorthogo
 1. the group containing the column `A[:, j]` has no other column with a nonzero in row `i`
 2. the group containing the row `A[i, :]` has no other row with a nonzero in column `j`
 
+It is equivalent to a __star bicoloring__.
+
 !!! warning
     This function is not coded with efficiency in mind, it is designed for small-scale tests.
 """
@@ -138,8 +144,8 @@ function structurally_biorthogonal(
     if !proper_length_bicoloring(A, row_color, column_color; verbose)
         return false
     end
-    column_group = group_by_color(column_color)
     row_group = group_by_color(row_color)
+    column_group = group_by_color(column_color)
     for i in axes(A, 1), j in axes(A, 2)
         iszero(A[i, j]) && continue
         ci, cj = row_color[i], column_color[j]
@@ -261,6 +267,174 @@ function directly_recoverable_columns(
     return true
 end
 
+"""
+    substitutable_columns(
+        A::AbstractMatrix, rank_nonzeros::AbstractMatrix, color::AbstractVector{<:Integer};
+        verbose=false
+    )
+
+Return `true` if coloring the columns of the symmetric matrix `A` with the vector `color` results in a partition that is substitutable, and `false` otherwise.
+For all nonzeros `A[i, j]`, `rank_nonzeros[i, j]` provides its rank of recovery.
+
+A partition of the columns of a symmetric matrix `A` is _substitutable_ if, for every nonzero element `A[i, j]`, either of the following statements holds:
+
+1. the group containing the column `A[:, j]` has all nonzeros in row `i` ordered before `A[i, j]`
+2. the group containing the column `A[:, i]` has all nonzeros in row `j` ordered before `A[i, j]`
+
+It is equivalent to an __acyclic coloring__.
+
+!!! warning
+    This function is not coded with efficiency in mind, it is designed for small-scale tests.
+
+# References
+
+> [_On the Estimation of Sparse Hessian Matrices_](https://doi.org/10.1137/0716078), Powell and Toint (1979)
+> [_The Cyclic Coloring Problem and Estimation of Sparse Hessian Matrices_](https://doi.org/10.1137/0607026), Coleman and Cai (1986)
+> [_What Color Is Your Jacobian? Graph Coloring for Computing Derivatives_](https://epubs.siam.org/doi/10.1137/S0036144504444711), Gebremedhin et al. (2005)
+"""
+function substitutable_columns(
+    A::AbstractMatrix,
+    rank_nonzeros::AbstractMatrix,
+    color::AbstractVector{<:Integer};
+    verbose::Bool=false,
+)
+    checksquare(A)
+    if !proper_length_coloring(A, color; verbose)
+        return false
+    end
+    issymmetric(A) || return false
+    group = group_by_color(color)
+    for i in axes(A, 1), j in axes(A, 2)
+        iszero(A[i, j]) && continue
+        ci, cj = color[i], color[j]
+        check = _substitutable_check(
+            A, rank_nonzeros; i, j, ci, cj, row_group=group, column_group=group, verbose
+        )
+        !check && return false
+    end
+    return true
+end
+
+"""
+    substitutable_bidirectional(
+        A::AbstractMatrix, rank_nonzeros::AbstractMatrix, row_color::AbstractVector{<:Integer}, column_color::AbstractVector{<:Integer};
+        verbose=false
+    )
+
+Return `true` if bicoloring of the matrix `A` with the vectors `row_color` and `column_color` results in a bipartition that is substitutable, and `false` otherwise.
+For all nonzeros `A[i, j]`, `rank_nonzeros[i, j]` provides its rank of recovery.
+
+A bipartition of the rows and columns of a matrix `A` is _substitutable_ if, for every nonzero element `A[i, j]`, either of the following statements holds:
+
+1. the group containing the column `A[:, j]` has all nonzeros in row `i` ordered before `A[i, j]`
+2. the group containing the row `A[i, :]` has all nonzeros in column `j` ordered before `A[i, j]`
+
+It is equivalent to an __acyclic bicoloring__.
+
+!!! warning
+    This function is not coded with efficiency in mind, it is designed for small-scale tests.
+"""
+function substitutable_bidirectional(
+    A::AbstractMatrix,
+    rank_nonzeros::AbstractMatrix,
+    row_color::AbstractVector{<:Integer},
+    column_color::AbstractVector{<:Integer};
+    verbose::Bool=false,
+)
+    if !proper_length_bicoloring(A, row_color, column_color; verbose)
+        return false
+    end
+    row_group = group_by_color(row_color)
+    column_group = group_by_color(column_color)
+    for i in axes(A, 1), j in axes(A, 2)
+        iszero(A[i, j]) && continue
+        ci, cj = row_color[i], column_color[j]
+        check = _substitutable_check(
+            A, rank_nonzeros; i, j, ci, cj, row_group, column_group, verbose
+        )
+        !check && return false
+    end
+    return true
+end
+
+function _substitutable_check(
+    A::AbstractMatrix,
+    rank_nonzeros::AbstractMatrix;
+    i::Integer,
+    j::Integer,
+    ci::Integer,
+    cj::Integer,
+    row_group::AbstractVector,
+    column_group::AbstractVector,
+    verbose::Bool,
+)
+    order_ij = rank_nonzeros[i, j]
+    k_row = 0
+    k_column = 0
+    if ci != 0
+        for k in row_group[ci]
+            (k == i) && continue
+            if !iszero(A[k, j])
+                order_kj = rank_nonzeros[k, j]
+                @assert !iszero(order_kj)
+                if order_kj > order_ij
+                    k_row = k
+                end
+            end
+        end
+    end
+    if cj != 0
+        for k in column_group[cj]
+            (k == j) && continue
+            if !iszero(A[i, k])
+                order_ik = rank_nonzeros[i, k]
+                @assert !iszero(order_ik)
+                if order_ik > order_ij
+                    k_column = k
+                end
+            end
+        end
+    end
+    if ci == 0 && cj == 0
+        if verbose
+            @warn """
+            For coefficient (i=$i, j=$j) with colors (ci=$ci, cj=$cj):
+            - Row color ci=$ci is neutral.
+            - Column color cj=$cj is neutral.
+            """
+        end
+        return false
+    elseif ci == 0 && !iszero(k_column)
+        if verbose
+            @warn """
+            For coefficient (i=$i, j=$j) with colors (ci=$ci, cj=$cj):
+            - Row color ci=$ci is neutral.
+            - For the column $k_column in column color cj=$cj, A[$i, $k_column] is ordered after A[$i, $j].
+            """
+        end
+        return false
+    elseif cj == 0 && !iszero(k_row)
+        if verbose
+            @warn """
+            For coefficient (i=$i, j=$j) with colors (ci=$ci, cj=$cj):
+            - For the row $k_row in row color ci=$ci, A[$k_row, $j] is ordered after A[$i, $j].
+            - Column color cj=$cj is neutral.
+            """
+        end
+        return false
+    elseif !iszero(k_row) && !iszero(k_column)
+        if verbose
+            @warn """
+            For coefficient (i=$i, j=$j) with colors (ci=$ci, cj=$cj):
+            - For the row $k_row in row color ci=$ci, A[$k_row, $j] is ordered after A[$i, $j].
+            - For the column $k_column in column color cj=$cj, A[$i, $k_column] is ordered after A[$i, $j].
+            """
+        end
+        return false
+    end
+    return true
+end
+
 """
     valid_dynamic_order(g::AdjacencyGraph, π::AbstractVector{<:Integer}, order::DynamicDegreeBasedOrder)
     valid_dynamic_order(bg::AdjacencyGraph, ::Val{side}, π::AbstractVector{<:Integer}, order::DynamicDegreeBasedOrder)
@@ -326,3 +500,75 @@ function valid_dynamic_order(
     end
     return true
 end
+
+"""
+    rank_nonzeros_from_trees(result::TreeSetColoringResult)
+    rank_nonzeros_from_trees(result::BicoloringResult)
+
+Construct a sparse matrix `rank_nonzeros` that assigns a unique recovery rank
+to each nonzero coefficient associated with an acyclic coloring or bicoloring.
+
+For every nonzero entry `result.A[i, j]`, `rank_nonzeros[i, j]` stores a strictly positive
+integer representing the order in which this coefficient is recovered during the decompression.
+A larger value means the coefficient is recovered later.
+
+This ranking is used to test substitutability (acyclicity) of colorings:
+for a given nonzero `result.A[i, j]`, the ranks allow one to check whether all competing
+nonzeros in the same row or column (within a color group) are recovered before it.
+"""
+function rank_nonzeros_from_trees end
+
+function rank_nonzeros_from_trees(result::TreeSetColoringResult)
+    (; A, ag, reverse_bfs_orders, diagonal_indices, tree_edge_indices, nt) = result
+    (; S) = ag
+    m, n = size(A)
+    nnzS = nnz(S)
+    nzval = zeros(Int, nnzS)
+    rank_nonzeros = SparseMatrixCSC(n, n, S.colptr, S.rowval, nzval)
+    counter = 0
+    for i in diagonal_indices
+        counter += 1
+        rank_nonzeros[i, i] = counter
+    end
+    for k in 1:nt
+        first = tree_edge_indices[k]
+        last = tree_edge_indices[k + 1] - 1
+        for pos in first:last
+            (i, j) = reverse_bfs_orders[pos]
+            counter += 1
+            rank_nonzeros[i, j] = counter
+            rank_nonzeros[j, i] = counter
+        end
+    end
+    return rank_nonzeros
+end
+
+function rank_nonzeros_from_trees(result::BicoloringResult)
+    (; A, abg, row_color, column_color, symmetric_result, large_colptr, large_rowval) =
+        result
+    @assert symmetric_result isa TreeSetColoringResult
+    (; ag, reverse_bfs_orders, tree_edge_indices, nt) = symmetric_result
+    (; S) = ag
+    m, n = size(A)
+    nnzA = nnz(S) ÷ 2
+    nzval = zeros(Int, nnzA)
+    colptr = large_colptr[1:(n + 1)]
+    rowval = large_rowval[1:nnzA]
+    rowval .-= n
+    rank_nonzeros = SparseMatrixCSC(m, n, colptr, rowval, nzval)
+    counter = 0
+    for k in 1:nt
+        first = tree_edge_indices[k]
+        last = tree_edge_indices[k + 1] - 1
+        for pos in first:last
+            (i, j) = reverse_bfs_orders[pos]
+            counter += 1
+            if i > j
+                rank_nonzeros[i - n, j] = counter
+            else
+                rank_nonzeros[j - n, i] = counter
+            end
+        end
+    end
+    return rank_nonzeros
+end
diff --git a/src/coloring.jl b/src/coloring.jl
index 7eec55c0..699dbc64 100644
--- a/src/coloring.jl
+++ b/src/coloring.jl
@@ -1,5 +1,10 @@
+struct InvalidColoringError <: Exception end
+
 """
-    partial_distance2_coloring(bg::BipartiteGraph, ::Val{side}, vertices_in_order::AbstractVector)
+    partial_distance2_coloring(
+        bg::BipartiteGraph, ::Val{side}, vertices_in_order::AbstractVector;
+        forced_colors::Union{AbstractVector{<:Integer},Nothing}=nothing
+    )
 
 Compute a distance-2 coloring of the given `side` (`1` or `2`) in the bipartite graph `bg` and return a vector of integer colors.
 
@@ -7,6 +12,8 @@ A _distance-2 coloring_ is such that two vertices have different colors if they
 
 The vertices are colored in a greedy fashion, following the order supplied.
 
+The optional `forced_colors` keyword argument is used to enforce predefined vertex colors (e.g. coming from another optimization algorithm) but still run the distance-2 coloring procedure to verify correctness.
+
 # See also
 
 - [`BipartiteGraph`](@ref)
@@ -17,11 +24,16 @@ The vertices are colored in a greedy fashion, following the order supplied.
 > [_What Color Is Your Jacobian? Graph Coloring for Computing Derivatives_](https://epubs.siam.org/doi/10.1137/S0036144504444711), Gebremedhin et al. (2005), Algorithm 3.2
 """
 function partial_distance2_coloring(
-    bg::BipartiteGraph{T}, ::Val{side}, vertices_in_order::AbstractVector{<:Integer}
+    bg::BipartiteGraph{T},
+    ::Val{side},
+    vertices_in_order::AbstractVector{<:Integer};
+    forced_colors::Union{AbstractVector{<:Integer},Nothing}=nothing,
 ) where {T,side}
     color = Vector{T}(undef, nb_vertices(bg, Val(side)))
     forbidden_colors = Vector{T}(undef, nb_vertices(bg, Val(side)))
-    partial_distance2_coloring!(color, forbidden_colors, bg, Val(side), vertices_in_order)
+    partial_distance2_coloring!(
+        color, forbidden_colors, bg, Val(side), vertices_in_order; forced_colors
+    )
     return color
 end
 
@@ -30,7 +42,8 @@ function partial_distance2_coloring!(
     forbidden_colors::AbstractVector{<:Integer},
     bg::BipartiteGraph,
     ::Val{side},
-    vertices_in_order::AbstractVector{<:Integer},
+    vertices_in_order::AbstractVector{<:Integer};
+    forced_colors::Union{AbstractVector{<:Integer},Nothing}=nothing,
 ) where {side}
     color .= 0
     forbidden_colors .= 0
@@ -44,17 +57,32 @@ function partial_distance2_coloring!(
                 end
             end
         end
-        for i in eachindex(forbidden_colors)
-            if forbidden_colors[i] != v
-                color[v] = i
-                break
+        if isnothing(forced_colors)
+            for i in eachindex(forbidden_colors)
+                if forbidden_colors[i] != v
+                    color[v] = i
+                    break
+                end
+            end
+        else
+            f = forced_colors[v]
+            if (
+                (f == 0 && length(neighbors(bg, Val(side), v)) > 0) ||
+                (f > 0 && forbidden_colors[f] == v)
+            )
+                throw(InvalidColoringError())
+            else
+                color[v] = f
             end
         end
     end
 end
 
 """
-    star_coloring(g::AdjacencyGraph, vertices_in_order::AbstractVector, postprocessing::Bool)
+    star_coloring(
+        g::AdjacencyGraph, vertices_in_order::AbstractVector, postprocessing::Bool;
+        forced_colors::Union{AbstractVector,Nothing}=nothing
+    )
 
 Compute a star coloring of all vertices in the adjacency graph `g` and return a tuple `(color, star_set)`, where
 
@@ -67,6 +95,8 @@ The vertices are colored in a greedy fashion, following the order supplied.
 
 If `postprocessing=true`, some colors might be replaced with `0` (the "neutral" color) as long as they are not needed during decompression.
 
+The optional `forced_colors` keyword argument is used to enforce predefined vertex colors (e.g. coming from another optimization algorithm) but still run the star coloring procedure to verify correctness and build auxiliary data structures, useful during decompression.
+
 # See also
 
 - [`AdjacencyGraph`](@ref)
@@ -77,7 +107,10 @@ If `postprocessing=true`, some colors might be replaced with `0` (the "neutral"
 > [_New Acyclic and Star Coloring Algorithms with Application to Computing Hessians_](https://epubs.siam.org/doi/abs/10.1137/050639879), Gebremedhin et al. (2007), Algorithm 4.1
 """
 function star_coloring(
-    g::AdjacencyGraph{T}, vertices_in_order::AbstractVector{<:Integer}, postprocessing::Bool
+    g::AdjacencyGraph{T},
+    vertices_in_order::AbstractVector{<:Integer},
+    postprocessing::Bool;
+    forced_colors::Union{AbstractVector{<:Integer},Nothing}=nothing,
 ) where {T<:Integer}
     # Initialize data structures
     nv = nb_vertices(g)
@@ -91,7 +124,7 @@ function star_coloring(
 
     for v in vertices_in_order
         for (w, index_vw) in neighbors_with_edge_indices(g, v)
-            !has_diagonal(g) || (v == w && continue)
+            augmented_graph(g) || (v == w && continue)
             color_w = color[w]
             iszero(color_w) && continue
             forbidden_colors[color_w] = v
@@ -106,7 +139,7 @@ function star_coloring(
             else
                 first_neighbor[color[w]] = (v, w, index_vw)
                 for (x, index_wx) in neighbors_with_edge_indices(g, w)
-                    !has_diagonal(g) || (w == x && continue)
+                    augmented_graph(g) || (w == x && continue)
                     color_x = color[x]
                     (x == v || iszero(color_x)) && continue
                     if x == hub[star[index_wx]]  # potential Case 2 (which is always false for trivial stars with two vertices, since the associated hub is negative)
@@ -115,10 +148,18 @@ function star_coloring(
                 end
             end
         end
-        for i in eachindex(forbidden_colors)
-            if forbidden_colors[i] != v
-                color[v] = i
-                break
+        if isnothing(forced_colors)
+            for i in eachindex(forbidden_colors)
+                if forbidden_colors[i] != v
+                    color[v] = i
+                    break
+                end
+            end
+        else
+            if forbidden_colors[forced_colors[v]] == v  # TODO: handle forced_colors[v] == 0
+                throw(InvalidColoringError())
+            else
+                color[v] = forced_colors[v]
             end
         end
         _update_stars!(star, hub, g, v, color, first_neighbor)
@@ -143,7 +184,7 @@ function _treat!(
     color::AbstractVector{<:Integer},
 )
     for x in neighbors(g, w)
-        !has_diagonal(g) || (w == x && continue)
+        augmented_graph(g) || (w == x && continue)
         color_x = color[x]
         iszero(color_x) && continue
         forbidden_colors[color_x] = v
@@ -163,12 +204,12 @@ function _update_stars!(
     first_neighbor::AbstractVector{<:Tuple},
 )
     for (w, index_vw) in neighbors_with_edge_indices(g, v)
-        !has_diagonal(g) || (v == w && continue)
+        augmented_graph(g) || (v == w && continue)
         color_w = color[w]
         iszero(color_w) && continue
         x_exists = false
         for (x, index_wx) in neighbors_with_edge_indices(g, w)
-            !has_diagonal(g) || (w == x && continue)
+            augmented_graph(g) || (w == x && continue)
             if x != v && color[x] == color[v]  # vw, wx ∈ E
                 star_wx = star[index_wx]
                 hub[star_wx] = w  # this may already be true
@@ -245,23 +286,22 @@ function acyclic_coloring(
 
     for v in vertices_in_order
         for w in neighbors(g, v)
-            !has_diagonal(g) || (v == w && continue)
+            augmented_graph(g) || (v == w && continue)
             color_w = color[w]
             iszero(color_w) && continue
             forbidden_colors[color_w] = v
         end
         for w in neighbors(g, v)
-            !has_diagonal(g) || (v == w && continue)
+            augmented_graph(g) || (v == w && continue)
             iszero(color[w]) && continue
             for (x, index_wx) in neighbors_with_edge_indices(g, w)
-                !has_diagonal(g) || (w == x && continue)
+                augmented_graph(g) || (w == x && continue)
                 color_x = color[x]
                 iszero(color_x) && continue
                 if forbidden_colors[color_x] != v
                     _prevent_cycle!(
                         v,
                         w,
-                        x,
                         index_wx,
                         color_x,
                         first_visit_to_tree,
@@ -271,6 +311,7 @@ function acyclic_coloring(
                 end
             end
         end
+        # TODO: handle forced colors
         for i in eachindex(forbidden_colors)
             if forbidden_colors[i] != v
                 color[v] = i
@@ -278,20 +319,20 @@ function acyclic_coloring(
             end
         end
         for (w, index_vw) in neighbors_with_edge_indices(g, v)  # grow two-colored stars around the vertex v
-            !has_diagonal(g) || (v == w && continue)
+            augmented_graph(g) || (v == w && continue)
             color_w = color[w]
             iszero(color_w) && continue
             _grow_star!(v, w, index_vw, color_w, first_neighbor, forest)
         end
         for (w, index_vw) in neighbors_with_edge_indices(g, v)
-            !has_diagonal(g) || (v == w && continue)
+            augmented_graph(g) || (v == w && continue)
             iszero(color[w]) && continue
             for (x, index_wx) in neighbors_with_edge_indices(g, w)
-                !has_diagonal(g) || (w == x && continue)
+                augmented_graph(g) || (w == x && continue)
                 color_x = color[x]
                 (x == v || iszero(color_x)) && continue
                 if color_x == color[v]
-                    _merge_trees!(v, w, x, index_vw, index_wx, forest)  # merge trees T₁ ∋ vw and T₂ ∋ wx if T₁ != T₂
+                    _merge_trees!(index_vw, index_wx, forest)  # merge trees T₁ ∋ vw and T₂ ∋ wx if T₁ != T₂
                 end
             end
         end
@@ -312,7 +353,6 @@ function _prevent_cycle!(
     # not modified
     v::Integer,
     w::Integer,
-    x::Integer,
     index_wx::Integer,
     color_x::Integer,
     # modified
@@ -354,9 +394,6 @@ end
 
 function _merge_trees!(
     # not modified
-    v::Integer,
-    w::Integer,
-    x::Integer,
     index_vw::Integer,
     index_wx::Integer,
     # modified
@@ -380,9 +417,23 @@ Encode a set of 2-colored trees resulting from the [`acyclic_coloring`](@ref) al
 $TYPEDFIELDS
 """
 struct TreeSet{T}
+    """
+    For a tree index `1 <= k <= nt`, the list
+    `list = reverse_bfs_order[tree_edge_indices[k]:(tree_edge_indices[k+1]-1)]` is a list of edges
+    `list[i] = (leaf, inner)` of the `k`th tree such that `leaf` is a leaf of the tree containing
+    the edges `list[i:end]`.
+    From an other point of view, `reverse(list)` contains the edges in the order of a breadth first
+    (BFS) traversal order of the `k`th tree starting from a depth-minimizing root.
+    """
     reverse_bfs_orders::Vector{Tuple{T,T}}
+    "For a tree index `1 <= k <= nt`, `is_star[k]` indicates whether the `k`th three is a star."
     is_star::Vector{Bool}
+    """
+    `tree_edge_indices[1]` is one and `tree_edge_indices[k+1] - tree_edge_indices[k]` is the number of edges in the `k`th tree.
+    One can think of it as a kind of fused vector of offsets (similar to the `colptr` field of `SparseMatrixCSC`) of all trees together.
+    """
     tree_edge_indices::Vector{T}
+    "numbers of 2-colored trees for which trees sharing the same 2 colors have disjoint edges"
     nt::T
 end
 
@@ -390,6 +441,7 @@ function TreeSet(
     g::AdjacencyGraph{T},
     forest::Forest{T},
     buffer::AbstractVector{T},
+    # The value of `reverse_bfs_orders` is ignored, we just provide the storage for it (or reuse memory allocated during acyclic coloring)
     reverse_bfs_orders::Vector{Tuple{T,T}},
     ne::Integer,
 ) where {T}
@@ -524,7 +576,15 @@ function TreeSet(
     # Number of edges treated
     num_edges_treated = zero(T)
 
-    # reverse_bfs_orders contains the reverse breadth first (BFS) traversal order for each tree in the forest
+    # The `rank` of the `k`th tree encoded in `forest` does not correspond
+    # to the depth of the tree rooted as the root encoded in `forest` because
+    # `forest.parents[u] = v` only needs a path to exists from `u` to `v` but
+    # there may not be an edge `(u, v)`.
+    # We also want a root `r` that minimizes the depth of the tree rooted at
+    # `r`. To achieve this, we start from each leaf and remove the corresponding
+    # edges. We then look at all leaves of the corresponding graphs and repeat
+    # the process until there is only one vertex left. This vertex will then be
+    # a depth-minimizing root.
     for k in 1:nt
         # Initialize the queue to store the leaves
         queue_start = 1
@@ -601,150 +661,3 @@ function TreeSet(
 
     return TreeSet(reverse_bfs_orders, is_star, tree_edge_indices, nt)
 end
-
-## Postprocessing, mirrors decompression code
-
-function postprocess!(
-    color::AbstractVector{<:Integer},
-    star_or_tree_set::Union{StarSet,TreeSet},
-    g::AdjacencyGraph,
-    offsets::AbstractVector{<:Integer},
-)
-    S = pattern(g)
-    edge_to_index = edge_indices(g)
-    # flag which colors are actually used during decompression
-    nb_colors = maximum(color)
-    color_used = zeros(Bool, nb_colors)
-
-    # nonzero diagonal coefficients force the use of their respective color (there can be no neutral colors if the diagonal is fully nonzero)
-    if has_diagonal(g)
-        for i in axes(S, 1)
-            if !iszero(S[i, i])
-                color_used[color[i]] = true
-            end
-        end
-    end
-
-    if star_or_tree_set isa StarSet
-        # only the colors of the hubs are used
-        (; star, hub) = star_or_tree_set
-        nb_trivial_stars = 0
-
-        # Iterate through all non-trivial stars
-        for s in eachindex(hub)
-            h = hub[s]
-            if h > 0
-                color_used[color[h]] = true
-            else
-                nb_trivial_stars += 1
-            end
-        end
-
-        # Process the trivial stars (if any)
-        if nb_trivial_stars > 0
-            rvS = rowvals(S)
-            for j in axes(S, 2)
-                for k in nzrange(S, j)
-                    i = rvS[k]
-                    if i > j
-                        index_ij = edge_to_index[k]
-                        s = star[index_ij]
-                        h = hub[s]
-                        if h < 0
-                            h = abs(h)
-                            spoke = h == j ? i : j
-                            if color_used[color[spoke]]
-                                # Switch the hub and the spoke to possibly avoid adding one more used color
-                                hub[s] = spoke
-                            else
-                                # Keep the current hub
-                                color_used[color[h]] = true
-                            end
-                        end
-                    end
-                end
-            end
-        end
-    else
-        # only the colors of non-leaf vertices are used
-        (; reverse_bfs_orders, is_star, tree_edge_indices, nt) = star_or_tree_set
-        nb_trivial_trees = 0
-
-        # Iterate through all non-trivial trees
-        for k in 1:nt
-            # Position of the first edge in the tree
-            first = tree_edge_indices[k]
-
-            # Total number of edges in the tree
-            ne_tree = tree_edge_indices[k + 1] - first
-
-            # Check if we have more than one edge in the tree (non-trivial tree)
-            if ne_tree > 1
-                # Determine if the tree is a star
-                if is_star[k]
-                    # It is a non-trivial star and only the color of the hub is needed
-                    (_, hub) = reverse_bfs_orders[first]
-                    color_used[color[hub]] = true
-                else
-                    # It is not a star and both colors are needed during the decompression
-                    (i, j) = reverse_bfs_orders[first]
-                    color_used[color[i]] = true
-                    color_used[color[j]] = true
-                end
-            else
-                nb_trivial_trees += 1
-            end
-        end
-
-        # Process the trivial trees (if any)
-        if nb_trivial_trees > 0
-            for k in 1:nt
-                # Position of the first edge in the tree
-                first = tree_edge_indices[k]
-
-                # Total number of edges in the tree
-                ne_tree = tree_edge_indices[k + 1] - first
-
-                # Check if we have exactly one edge in the tree
-                if ne_tree == 1
-                    (i, j) = reverse_bfs_orders[first]
-                    if color_used[color[i]]
-                        # Make i the root to avoid possibly adding one more used color
-                        # Switch it with the (only) leaf
-                        reverse_bfs_orders[first] = (j, i)
-                    else
-                        # Keep j as the root
-                        color_used[color[j]] = true
-                    end
-                end
-            end
-        end
-    end
-
-    # if at least one of the colors is useless, modify the color assignments of vertices
-    if any(!, color_used)
-        num_colors_useless = 0
-
-        # determine what are the useless colors and compute the offsets
-        for ci in 1:nb_colors
-            if color_used[ci]
-                offsets[ci] = num_colors_useless
-            else
-                num_colors_useless += 1
-            end
-        end
-
-        # assign the neutral color to every vertex with a useless color and remap the colors
-        for i in eachindex(color)
-            ci = color[i]
-            if !color_used[ci]
-                # assign the neutral color
-                color[i] = 0
-            else
-                # remap the color to not have any gap
-                color[i] -= offsets[ci]
-            end
-        end
-    end
-    return color
-end
diff --git a/src/constant.jl b/src/constant.jl
index d79caab0..56bf4101 100644
--- a/src/constant.jl
+++ b/src/constant.jl
@@ -4,20 +4,24 @@
 Coloring algorithm which always returns the same precomputed vector of colors.
 Useful when the optimal coloring of a matrix can be determined a priori due to its specific structure (e.g. banded).
 
-It is passed as an argument to the main function [`coloring`](@ref), but will only work if the associated `problem` has `:nonsymmetric` structure.
-Indeed, for symmetric coloring problems, we need more than just the vector of colors to allow fast decompression.
+It is passed as an argument to the main function [`coloring`](@ref), but will only work if the associated `problem` has a `:column` or `:row` partition.
 
 # Constructors
 
     ConstantColoringAlgorithm{partition}(matrix_template, color)
-    ConstantColoringAlgorithm(matrix_template, color; partition=:column)
+    ConstantColoringAlgorithm{partition,structure}(matrix_template, color)
+    ConstantColoringAlgorithm(
+        matrix_template, color;
+        structure=:nonsymmetric, partition=:column
+    )
 
 - `partition::Symbol`: either `:row` or `:column`.
+- `structure::Symbol`: either `:nonsymmetric` or `:symmetric`.
 - `matrix_template::AbstractMatrix`: matrix for which the vector of colors was precomputed (the algorithm will only accept matrices of the exact same size).
 - `color::Vector{<:Integer}`: vector of integer colors, one for each row or column (depending on `partition`).
 
 !!! warning
-    The second constructor (based on keyword arguments) is type-unstable.
+    The constructor based on keyword arguments is type-unstable if these arguments are not compile-time constants.
 
 We do not necessarily verify consistency between the matrix template and the vector of colors, this is the responsibility of the user.
 
@@ -63,71 +67,68 @@ julia> column_colors(result)
 
 - [`ADTypes.column_coloring`](@extref ADTypes.column_coloring)
 - [`ADTypes.row_coloring`](@extref ADTypes.row_coloring)
+- [`ADTypes.symmetric_coloring`](@extref ADTypes.symmetric_coloring)
 """
-struct ConstantColoringAlgorithm{
-    partition,
-    M<:AbstractMatrix,
-    T<:Integer,
-    R<:AbstractColoringResult{:nonsymmetric,partition,:direct},
-} <: ADTypes.AbstractColoringAlgorithm
+struct ConstantColoringAlgorithm{partition,structure,M<:AbstractMatrix,T<:Integer} <:
+       ADTypes.AbstractColoringAlgorithm
     matrix_template::M
     color::Vector{T}
-    result::R
-end
 
-function ConstantColoringAlgorithm{:column}(
-    matrix_template::AbstractMatrix, color::Vector{<:Integer}
-)
-    bg = BipartiteGraph(matrix_template)
-    result = ColumnColoringResult(matrix_template, bg, color)
-    T, M, R = eltype(bg), typeof(matrix_template), typeof(result)
-    return ConstantColoringAlgorithm{:column,M,T,R}(matrix_template, color, result)
+    function ConstantColoringAlgorithm{partition,structure}(
+        matrix_template::AbstractMatrix, color::Vector{<:Integer}
+    ) where {partition,structure}
+        check_valid_problem(structure, partition)
+        return new{partition,structure,typeof(matrix_template),eltype(color)}(
+            matrix_template, color
+        )
+    end
 end
 
-function ConstantColoringAlgorithm{:row}(
+function ConstantColoringAlgorithm{partition}(
     matrix_template::AbstractMatrix, color::Vector{<:Integer}
-)
-    bg = BipartiteGraph(matrix_template)
-    result = RowColoringResult(matrix_template, bg, color)
-    T, M, R = eltype(bg), typeof(matrix_template), typeof(result)
-    return ConstantColoringAlgorithm{:row,M,T,R}(matrix_template, color, result)
+) where {partition}
+    return ConstantColoringAlgorithm{partition,:nonsymmetric}(matrix_template, color)
 end
 
 function ConstantColoringAlgorithm(
-    matrix_template::AbstractMatrix, color::Vector{<:Integer}; partition::Symbol=:column
+    matrix_template::AbstractMatrix,
+    color::Vector{<:Integer};
+    structure::Symbol=:nonsymmetric,
+    partition::Symbol=:column,
 )
-    return ConstantColoringAlgorithm{partition}(matrix_template, color)
+    return ConstantColoringAlgorithm{partition,structure}(matrix_template, color)
 end
 
-function coloring(
-    A::AbstractMatrix,
-    ::ColoringProblem{:nonsymmetric,partition},
-    algo::ConstantColoringAlgorithm{partition};
-    decompression_eltype::Type=Float64,
-    symmetric_pattern::Bool=false,
-) where {partition}
-    (; matrix_template, result) = algo
+function check_template(algo::ConstantColoringAlgorithm, A::AbstractMatrix)
+    (; matrix_template) = algo
     if size(A) != size(matrix_template)
         throw(
             DimensionMismatch(
                 "`ConstantColoringAlgorithm` expected matrix of size $(size(matrix_template)) but got matrix of size $(size(A))",
             ),
         )
-    else
-        return result
     end
 end
 
 function ADTypes.column_coloring(
-    A::AbstractMatrix, algo::ConstantColoringAlgorithm{:column}
+    A::AbstractMatrix, algo::ConstantColoringAlgorithm{:column,:nonsymmetric}
+)
+    check_template(algo, A)
+    return algo.color
+end
+
+function ADTypes.row_coloring(
+    A::AbstractMatrix, algo::ConstantColoringAlgorithm{:row,:nonsymmetric}
 )
-    problem = ColoringProblem{:nonsymmetric,:column}()
-    result = coloring(A, problem, algo)
-    return column_colors(result)
+    check_template(algo, A)
+    return algo.color
 end
 
-function ADTypes.row_coloring(A::AbstractMatrix, algo::ConstantColoringAlgorithm)
-    problem = ColoringProblem{:nonsymmetric,:row}()
-    result = coloring(A, problem, algo)
-    return row_colors(result)
+function ADTypes.symmetric_coloring(
+    A::AbstractMatrix, algo::ConstantColoringAlgorithm{:column,:symmetric}
+)
+    check_template(algo, A)
+    return algo.color
 end
+
+# TODO: handle bidirectional once https://github.com/SciML/ADTypes.jl/issues/69 is done
diff --git a/src/decompression.jl b/src/decompression.jl
index a3dd471a..849f0f6b 100644
--- a/src/decompression.jl
+++ b/src/decompression.jl
@@ -106,7 +106,7 @@ function compress(
     return Br, Bc
 end
 
-struct UnsupportedDecompressionError
+struct UnsupportedDecompressionError <: Exception
     msg::String
 end
 
@@ -339,9 +339,9 @@ end
 ## ColumnColoringResult
 
 function decompress!(A::AbstractMatrix, B::AbstractMatrix, result::ColumnColoringResult)
-    (; color) = result
-    S = result.bg.S2
-    check_same_pattern(A, S)
+    (; bg, color) = result
+    S = bg.S2
+    check_compatible_pattern(A, bg)
     fill!(A, zero(eltype(A)))
     rvS = rowvals(S)
     for j in axes(S, 2)
@@ -357,9 +357,9 @@ end
 function decompress_single_color!(
     A::AbstractMatrix, b::AbstractVector, c::Integer, result::ColumnColoringResult
 )
-    (; group) = result
-    S = result.bg.S2
-    check_same_pattern(A, S)
+    (; bg, group) = result
+    S = bg.S2
+    check_compatible_pattern(A, bg)
     rvS = rowvals(S)
     for j in group[c]
         for k in nzrange(S, j)
@@ -371,9 +371,9 @@ function decompress_single_color!(
 end
 
 function decompress!(A::SparseMatrixCSC, B::AbstractMatrix, result::ColumnColoringResult)
-    (; compressed_indices) = result
-    S = result.bg.S2
-    check_same_pattern(A, S)
+    (; bg, compressed_indices) = result
+    S = bg.S2
+    check_compatible_pattern(A, bg)
     nzA = nonzeros(A)
     for k in eachindex(nzA, compressed_indices)
         nzA[k] = B[compressed_indices[k]]
@@ -384,9 +384,9 @@ end
 function decompress_single_color!(
     A::SparseMatrixCSC, b::AbstractVector, c::Integer, result::ColumnColoringResult
 )
-    (; group) = result
-    S = result.bg.S2
-    check_same_pattern(A, S)
+    (; bg, group) = result
+    S = bg.S2
+    check_compatible_pattern(A, bg)
     rvS = rowvals(S)
     nzA = nonzeros(A)
     for j in group[c]
@@ -401,9 +401,9 @@ end
 ## RowColoringResult
 
 function decompress!(A::AbstractMatrix, B::AbstractMatrix, result::RowColoringResult)
-    (; color) = result
-    S = result.bg.S2
-    check_same_pattern(A, S)
+    (; bg, color) = result
+    S = bg.S2
+    check_compatible_pattern(A, bg)
     fill!(A, zero(eltype(A)))
     rvS = rowvals(S)
     for j in axes(S, 2)
@@ -419,9 +419,9 @@ end
 function decompress_single_color!(
     A::AbstractMatrix, b::AbstractVector, c::Integer, result::RowColoringResult
 )
-    (; group) = result
-    S, Sᵀ = result.bg.S2, result.bg.S1
-    check_same_pattern(A, S)
+    (; bg, group) = result
+    S, Sᵀ = bg.S2, bg.S1
+    check_compatible_pattern(A, bg)
     rvSᵀ = rowvals(Sᵀ)
     for i in group[c]
         for k in nzrange(Sᵀ, i)
@@ -433,9 +433,9 @@ function decompress_single_color!(
 end
 
 function decompress!(A::SparseMatrixCSC, B::AbstractMatrix, result::RowColoringResult)
-    (; compressed_indices) = result
-    S = result.bg.S2
-    check_same_pattern(A, S)
+    (; bg, compressed_indices) = result
+    S = bg.S2
+    check_compatible_pattern(A, bg)
     nzA = nonzeros(A)
     for k in eachindex(nzA, compressed_indices)
         nzA[k] = B[compressed_indices[k]]
@@ -450,7 +450,7 @@ function decompress!(
 )
     (; ag, compressed_indices) = result
     (; S) = ag
-    uplo == :F && check_same_pattern(A, S)
+    check_compatible_pattern(A, ag, uplo)
     fill!(A, zero(eltype(A)))
 
     rvS = rowvals(S)
@@ -474,7 +474,7 @@ function decompress_single_color!(
 )
     (; ag, compressed_indices, group) = result
     (; S) = ag
-    uplo == :F && check_same_pattern(A, S)
+    check_compatible_pattern(A, ag, uplo)
 
     lower_index = (c - 1) * S.n + 1
     upper_index = c * S.n
@@ -508,8 +508,8 @@ function decompress!(
     (; ag, compressed_indices) = result
     (; S) = ag
     nzA = nonzeros(A)
+    check_compatible_pattern(A, ag, uplo)
     if uplo == :F
-        check_same_pattern(A, S)
         for k in eachindex(nzA, compressed_indices)
             nzA[k] = B[compressed_indices[k]]
         end
@@ -534,9 +534,10 @@ end
 function decompress!(
     A::AbstractMatrix, B::AbstractMatrix, result::TreeSetColoringResult, uplo::Symbol=:F
 )
-    (; ag, color, reverse_bfs_orders, tree_edge_indices, nt, buffer) = result
+    (; ag, color, reverse_bfs_orders, tree_edge_indices, nt, diagonal_indices, buffer) =
+        result
     (; S) = ag
-    uplo == :F && check_same_pattern(A, S)
+    check_compatible_pattern(A, ag, uplo)
     R = eltype(A)
     fill!(A, zero(R))
 
@@ -547,11 +548,9 @@ function decompress!(
     end
 
     # Recover the diagonal coefficients of A
-    if has_diagonal(ag)
-        for i in axes(S, 1)
-            if !iszero(S[i, i])
-                A[i, i] = B[i, color[i]]
-            end
+    if !augmented_graph(ag)
+        for i in diagonal_indices
+            A[i, i] = B[i, color[i]]
         end
     end
 
@@ -586,8 +585,23 @@ function decompress!(
     return A
 end
 
-function decompress!(
-    A::SparseMatrixCSC{R},
+"""
+   decompress_csc!(
+        nzA::AbstractVector{R},
+        A_colptr::AbstractVector,
+        B::AbstractMatrix{R},
+        result::TreeSetColoringResult,
+        uplo::Symbol=:F,
+    ) where {R<:Real}
+
+Decompress the values of `B` into the vector of nonzero entries `nzA` of a
+sparse matrix with column pointers `colptr`. This function assumes that the
+row indices are sorted in increasing order and are the same as those of the
+sparse matrix given to the `coloring` function that returned `result`.
+"""
+function decompress_csc!(
+    nzA::AbstractVector{R},
+    A_colptr::AbstractVector{<:Integer},
     B::AbstractMatrix{R},
     result::TreeSetColoringResult,
     uplo::Symbol=:F,
@@ -604,10 +618,6 @@ function decompress!(
         upper_triangle_offsets,
         buffer,
     ) = result
-    (; S) = ag
-    A_colptr = A.colptr
-    nzA = nonzeros(A)
-    uplo == :F && check_same_pattern(A, S)
 
     if eltype(buffer) == R
         buffer_right_type = buffer
@@ -616,7 +626,7 @@ function decompress!(
     end
 
     # Recover the diagonal coefficients of A
-    if has_diagonal(ag)
+    if !augmented_graph(ag)
         if uplo == :L
             for i in diagonal_indices
                 # A[i, i] is the first element in column i
@@ -697,6 +707,17 @@ function decompress!(
             #! format: on
         end
     end
+    return nothing
+end
+
+function decompress!(
+    A::SparseMatrixCSC{R},
+    B::AbstractMatrix{R},
+    result::TreeSetColoringResult,
+    uplo::Symbol=:F,
+) where {R<:Real}
+    check_compatible_pattern(A, result.ag, uplo)
+    decompress_csc!(nonzeros(A), A.colptr, B, result, uplo)
     return A
 end
 
@@ -708,9 +729,10 @@ function decompress!(
     result::LinearSystemColoringResult,
     uplo::Symbol=:F,
 )
-    (; color, strict_upper_nonzero_inds, M_factorization, strict_upper_nonzeros_A) = result
-    S = result.ag.S
-    uplo == :F && check_same_pattern(A, S)
+    (; ag, color, strict_upper_nonzero_inds, M_factorization, strict_upper_nonzeros_A) =
+        result
+    S = ag.S
+    check_compatible_pattern(A, ag, uplo)
 
     # TODO: for some reason I cannot use ldiv! with a sparse QR
     strict_upper_nonzeros_A = M_factorization \ vec(B)
@@ -770,10 +792,22 @@ end
 function decompress!(
     A::AbstractMatrix, Br::AbstractMatrix, Bc::AbstractMatrix, result::BicoloringResult
 )
+    (; large_colptr, large_rowval, symmetric_result) = result
     m, n = size(A)
+    R = eltype(A)
+    fill!(A, zero(R))
+    nzval = Vector{R}(undef, length(large_rowval))
+    A_and_noAᵀ = SparseMatrixCSC(m + n, m + n, large_colptr, large_rowval, nzval)
     Br_and_Bc = _join_compressed!(result, Br, Bc)
-    A_and_Aᵀ = decompress(Br_and_Bc, result.symmetric_result)
-    copyto!(A, A_and_Aᵀ[(n + 1):(n + m), 1:n])  # original matrix in bottom left corner
+    decompress!(A_and_noAᵀ, Br_and_Bc, symmetric_result, :L)
+    rvA = rowvals(A_and_noAᵀ)
+    nzA = nonzeros(A_and_noAᵀ)
+    for j in 1:n
+        for k in nzrange(A_and_noAᵀ, j)
+            i = rvA[k]
+            A[i - n, j] = nzA[k]
+        end
+    end
     return A
 end
 
@@ -782,10 +816,10 @@ function decompress!(
 )
     (; large_colptr, large_rowval, symmetric_result) = result
     m, n = size(A)
-    Br_and_Bc = _join_compressed!(result, Br, Bc)
     # pretend A is larger
     A_and_noAᵀ = SparseMatrixCSC(m + n, m + n, large_colptr, large_rowval, A.nzval)
     # decompress lower triangle only
+    Br_and_Bc = _join_compressed!(result, Br, Bc)
     decompress!(A_and_noAᵀ, Br_and_Bc, symmetric_result, :L)
     return A
 end
diff --git a/src/graph.jl b/src/graph.jl
index 9b53a36d..94954602 100644
--- a/src/graph.jl
+++ b/src/graph.jl
@@ -23,7 +23,7 @@ end
 
 SparsityPatternCSC(A::SparseMatrixCSC) = SparsityPatternCSC(A.m, A.n, A.colptr, A.rowval)
 
-Base.eltype(::SparsityPatternCSC{T}) where {T} = T
+SparseArrays.indtype(::SparsityPatternCSC{T}) where {T} = T
 Base.size(S::SparsityPatternCSC) = (S.m, S.n)
 Base.size(S::SparsityPatternCSC, d::Integer) = d::Integer <= 2 ? size(S)[d] : 1
 Base.axes(S::SparsityPatternCSC, d::Integer) = Base.OneTo(size(S, d))
@@ -32,6 +32,11 @@ 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)
 
+# Needed if using `coloring(::SparsityPatternCSC, ...)`
+function Base.similar(A::SparsityPatternCSC, ::Type{T}) where {T}
+    return SparseArrays.SparseMatrixCSC(A.m, A.n, A.colptr, A.rowval, similar(A.rowval, T))
+end
+
 """
     transpose(S::SparsityPatternCSC)
 
@@ -179,6 +184,7 @@ function build_edge_to_index(S::SparsityPatternCSC{T}) where {T}
     # edge_to_index gives an index for each edge
     edge_to_index = Vector{T}(undef, nnz(S))
     offsets = zeros(T, S.n)
+    nb_self_loops = 0
     counter = 0
     rvS = rowvals(S)
     for j in axes(S, 2)
@@ -193,16 +199,17 @@ function build_edge_to_index(S::SparsityPatternCSC{T}) where {T}
             elseif i == j
                 # this should never be used, make sure it errors
                 edge_to_index[k] = 0
+                nb_self_loops += 1
             end
         end
     end
-    return edge_to_index
+    return edge_to_index, nb_self_loops
 end
 
 ## Adjacency graph
 
 """
-    AdjacencyGraph{T,has_diagonal}
+    AdjacencyGraph{T,augmented_graph}
 
 Undirected graph without self-loops representing the nonzeros of a symmetric matrix (typically a Hessian matrix).
 
@@ -213,45 +220,52 @@ The adjacency graph of a symmetric matrix `A ∈ ℝ^{n × n}` is `G(A) = (V, E)
 
 # Constructors
 
-    AdjacencyGraph(A::SparseMatrixCSC; has_diagonal::Bool=true)
+    AdjacencyGraph(A::SparseMatrixCSC; augmented_graph::Bool=false)
 
 # Fields
 
-- `S::SparsityPatternCSC{T}`: Underlying sparsity pattern, whose diagonal is empty whenever `has_diagonal` is `false`
+- `S::SparsityPatternCSC{T}`: Underlying sparsity pattern, which represents an augmented graph whenever `augmented_graph` is `true`. Here, "augmented graph" means the sparsity pattern of the augmented matrix `H = [0 Jᵀ; J 0]`.
 - `edge_to_index::Vector{T}`: A vector mapping each nonzero of `S` to a unique edge index (ignoring diagonal and accounting for symmetry, so that `(i, j)` and `(j, i)` get the same index)
 
 # References
 
 > [_What Color Is Your Jacobian? SparsityPatternCSC Coloring for Computing Derivatives_](https://epubs.siam.org/doi/10.1137/S0036144504444711), Gebremedhin et al. (2005)
 """
-struct AdjacencyGraph{T<:Integer,has_diagonal}
+struct AdjacencyGraph{T<:Integer,augmented_graph}
     S::SparsityPatternCSC{T}
     edge_to_index::Vector{T}
+    nb_self_loops::Int
 end
 
 Base.eltype(::AdjacencyGraph{T}) where {T} = T
 
 function AdjacencyGraph(
     S::SparsityPatternCSC{T},
-    edge_to_index::Vector{T}=build_edge_to_index(S);
-    has_diagonal::Bool=true,
+    edge_to_index::Vector{T},
+    nb_self_loops::Int;
+    augmented_graph::Bool=false,
 ) where {T}
-    return AdjacencyGraph{T,has_diagonal}(S, edge_to_index)
+    return AdjacencyGraph{T,augmented_graph}(S, edge_to_index, nb_self_loops)
+end
+
+function AdjacencyGraph(S::SparsityPatternCSC; augmented_graph::Bool=false)
+    edge_to_index, nb_self_loops = build_edge_to_index(S)
+    return AdjacencyGraph(S, edge_to_index, nb_self_loops; augmented_graph)
 end
 
-function AdjacencyGraph(A::SparseMatrixCSC; has_diagonal::Bool=true)
-    return AdjacencyGraph(SparsityPatternCSC(A); has_diagonal)
+function AdjacencyGraph(A::SparseMatrixCSC; augmented_graph::Bool=false)
+    return AdjacencyGraph(SparsityPatternCSC(A); augmented_graph)
 end
 
-function AdjacencyGraph(A::AbstractMatrix; has_diagonal::Bool=true)
-    return AdjacencyGraph(SparseMatrixCSC(A); has_diagonal)
+function AdjacencyGraph(A::AbstractMatrix; augmented_graph::Bool=false)
+    return AdjacencyGraph(SparseMatrixCSC(A); augmented_graph)
 end
 
 pattern(g::AdjacencyGraph) = g.S
 edge_indices(g::AdjacencyGraph) = g.edge_to_index
 nb_vertices(g::AdjacencyGraph) = pattern(g).n
 vertices(g::AdjacencyGraph) = 1:nb_vertices(g)
-has_diagonal(::AdjacencyGraph{T,hd}) where {T,hd} = hd
+augmented_graph(::AdjacencyGraph{T,ag}) where {T,ag} = ag
 
 function neighbors(g::AdjacencyGraph, v::Integer)
     S = pattern(g)
@@ -267,30 +281,22 @@ function neighbors_with_edge_indices(g::AdjacencyGraph, v::Integer)
     return zip(neighbors_v, edges_indices_v)
 end
 
-degree(g::AdjacencyGraph{T,false}, v::Integer) where {T} = g.S.colptr[v + 1] - g.S.colptr[v]
+degree(g::AdjacencyGraph{T,true}, v::Integer) where {T} = g.S.colptr[v + 1] - g.S.colptr[v]
 
-function degree(g::AdjacencyGraph{T,true}, v::Integer) where {T}
+function degree(g::AdjacencyGraph{T,false}, v::Integer) where {T}
     neigh = neighbors(g, v)
     has_selfloop = insorted(v, neigh)
     return g.S.colptr[v + 1] - g.S.colptr[v] - has_selfloop
 end
 
-nb_edges(g::AdjacencyGraph{T,false}) where {T} = nnz(g.S) ÷ 2
-
-function nb_edges(g::AdjacencyGraph{T,true}) where {T}
-    ne = 0
-    for v in vertices(g)
-        ne += degree(g, v)
-    end
-    return ne ÷ 2
-end
+nb_edges(g::AdjacencyGraph) = (nnz(g.S) - g.nb_self_loops) ÷ 2
 
 maximum_degree(g::AdjacencyGraph) = maximum(Base.Fix1(degree, g), vertices(g))
 minimum_degree(g::AdjacencyGraph) = minimum(Base.Fix1(degree, g), vertices(g))
 
 function has_neighbor(g::AdjacencyGraph, v::Integer, u::Integer)
     for w in neighbors(g, v)
-        !has_diagonal(g) || (v == w && continue)
+        augmented_graph(g) || (v == w && continue)
         if w == u
             return true
         end
diff --git a/src/interface.jl b/src/interface.jl
index 79fe71e9..0d183c9c 100644
--- a/src/interface.jl
+++ b/src/interface.jl
@@ -72,7 +72,7 @@ It is passed as an argument to the main function [`coloring`](@ref).
     GreedyColoringAlgorithm{decompression}(order=NaturalOrder(); postprocessing=false)
     GreedyColoringAlgorithm(order=NaturalOrder(); postprocessing=false, decompression=:direct)
 
-- `order::AbstractOrder`: the order in which the columns or rows are colored, which can impact the number of colors.
+- `order::Union{AbstractOrder,Tuple}`: the order in which the columns or rows are colored, which can impact the number of colors. Can also be a tuple of different orders to try out, from which the best order (the one with the lowest total number of colors) will be used.
 - `postprocessing::Bool`: whether or not the coloring will be refined by assigning the neutral color `0` to some vertices.
 - `decompression::Symbol`: either `:direct` or `:substitution`. Usually `:substitution` leads to fewer colors, at the cost of a more expensive coloring (and decompression). When `:substitution` is not applicable, it falls back on `:direct` decompression.
 
@@ -94,26 +94,31 @@ See their respective docstrings for details.
 - [`AbstractOrder`](@ref)
 - [`decompress`](@ref)
 """
-struct GreedyColoringAlgorithm{decompression,O<:AbstractOrder} <:
+struct GreedyColoringAlgorithm{decompression,N,O<:NTuple{N,AbstractOrder}} <:
        ADTypes.AbstractColoringAlgorithm
-    order::O
+    orders::O
     postprocessing::Bool
-end
 
-function GreedyColoringAlgorithm{decompression}(
-    order::AbstractOrder=NaturalOrder(); postprocessing::Bool=false
-) where {decompression}
-    check_valid_algorithm(decompression)
-    return GreedyColoringAlgorithm{decompression,typeof(order)}(order, postprocessing)
+    function GreedyColoringAlgorithm{decompression}(
+        order_or_orders::Union{AbstractOrder,Tuple}=NaturalOrder();
+        postprocessing::Bool=false,
+    ) where {decompression}
+        check_valid_algorithm(decompression)
+        if order_or_orders isa AbstractOrder
+            orders = (order_or_orders,)
+        else
+            orders = order_or_orders
+        end
+        return new{decompression,length(orders),typeof(orders)}(orders, postprocessing)
+    end
 end
 
 function GreedyColoringAlgorithm(
-    order::AbstractOrder=NaturalOrder();
+    order_or_orders::Union{AbstractOrder,Tuple}=NaturalOrder();
     postprocessing::Bool=false,
     decompression::Symbol=:direct,
 )
-    check_valid_algorithm(decompression)
-    return GreedyColoringAlgorithm{decompression,typeof(order)}(order, postprocessing)
+    return GreedyColoringAlgorithm{decompression}(order_or_orders; postprocessing)
 end
 
 ## Coloring
@@ -225,12 +230,16 @@ function _coloring(
     ::ColoringProblem{:nonsymmetric,:column},
     algo::GreedyColoringAlgorithm,
     decompression_eltype::Type,
-    symmetric_pattern::Bool,
+    symmetric_pattern::Bool;
+    forced_colors::Union{AbstractVector{<:Integer},Nothing}=nothing,
 )
     symmetric_pattern = symmetric_pattern || A isa Union{Symmetric,Hermitian}
     bg = BipartiteGraph(A; symmetric_pattern)
-    vertices_in_order = vertices(bg, Val(2), algo.order)
-    color = partial_distance2_coloring(bg, Val(2), vertices_in_order)
+    color_by_order = map(algo.orders) do order
+        vertices_in_order = vertices(bg, Val(2), order)
+        return partial_distance2_coloring(bg, Val(2), vertices_in_order; forced_colors)
+    end
+    color = argmin(maximum, color_by_order)
     if speed_setting isa WithResult
         return ColumnColoringResult(A, bg, color)
     else
@@ -244,12 +253,16 @@ function _coloring(
     ::ColoringProblem{:nonsymmetric,:row},
     algo::GreedyColoringAlgorithm,
     decompression_eltype::Type,
-    symmetric_pattern::Bool,
+    symmetric_pattern::Bool;
+    forced_colors::Union{AbstractVector{<:Integer},Nothing}=nothing,
 )
     symmetric_pattern = symmetric_pattern || A isa Union{Symmetric,Hermitian}
     bg = BipartiteGraph(A; symmetric_pattern)
-    vertices_in_order = vertices(bg, Val(1), algo.order)
-    color = partial_distance2_coloring(bg, Val(1), vertices_in_order)
+    color_by_order = map(algo.orders) do order
+        vertices_in_order = vertices(bg, Val(1), order)
+        return partial_distance2_coloring(bg, Val(1), vertices_in_order; forced_colors)
+    end
+    color = argmin(maximum, color_by_order)
     if speed_setting isa WithResult
         return RowColoringResult(A, bg, color)
     else
@@ -263,11 +276,15 @@ function _coloring(
     ::ColoringProblem{:symmetric,:column},
     algo::GreedyColoringAlgorithm{:direct},
     decompression_eltype::Type,
-    symmetric_pattern::Bool,
+    symmetric_pattern::Bool;
+    forced_colors::Union{AbstractVector{<:Integer},Nothing}=nothing,
 )
-    ag = AdjacencyGraph(A; has_diagonal=true)
-    vertices_in_order = vertices(ag, algo.order)
-    color, star_set = star_coloring(ag, vertices_in_order, algo.postprocessing)
+    ag = AdjacencyGraph(A; augmented_graph=false)
+    color_and_star_set_by_order = map(algo.orders) do order
+        vertices_in_order = vertices(ag, order)
+        return star_coloring(ag, vertices_in_order, algo.postprocessing; forced_colors)
+    end
+    color, star_set = argmin(maximum ∘ first, color_and_star_set_by_order)
     if speed_setting isa WithResult
         return StarSetColoringResult(A, ag, color, star_set)
     else
@@ -283,9 +300,18 @@ function _coloring(
     decompression_eltype::Type{R},
     symmetric_pattern::Bool,
 ) where {R}
-    ag = AdjacencyGraph(A; has_diagonal=true)
-    vertices_in_order = vertices(ag, algo.order)
-    color, tree_set = acyclic_coloring(ag, vertices_in_order, algo.postprocessing)
+    ag = AdjacencyGraph(A; augmented_graph=false)
+    color_and_tree_set_by_order = map(algo.orders) do order
+        vertices_in_order = vertices(ag, order)
+        return acyclic_coloring(ag, vertices_in_order, algo.postprocessing)
+    end
+    # if `color` is empty, `maximum` will fail but `color_and_tree_set_by_order`
+    # is also one so we can just add a special case for this
+    if length(color_and_tree_set_by_order) == 1
+        color, tree_set = only(color_and_tree_set_by_order)
+    else
+        color, tree_set = argmin(maximum ∘ first, color_and_tree_set_by_order)
+    end
     if speed_setting isa WithResult
         return TreeSetColoringResult(A, ag, color, tree_set, R)
     else
@@ -299,19 +325,44 @@ function _coloring(
     ::ColoringProblem{:nonsymmetric,:bidirectional},
     algo::GreedyColoringAlgorithm{:direct},
     decompression_eltype::Type{R},
-    symmetric_pattern::Bool,
+    symmetric_pattern::Bool;
+    forced_colors::Union{AbstractVector{<:Integer},Nothing}=nothing,
 ) where {R}
     A_and_Aᵀ, edge_to_index = bidirectional_pattern(A; symmetric_pattern)
-    ag = AdjacencyGraph(A_and_Aᵀ, edge_to_index; has_diagonal=false)
-    vertices_in_order = vertices(ag, algo.order)
-    color, star_set = star_coloring(ag, vertices_in_order, algo.postprocessing)
+    ag = AdjacencyGraph(A_and_Aᵀ, edge_to_index, 0; augmented_graph=true)
+    outputs_by_order = map(algo.orders) do order
+        vertices_in_order = vertices(ag, order)
+        _color, _star_set = star_coloring(
+            ag, vertices_in_order, algo.postprocessing; forced_colors
+        )
+        (_row_color, _column_color, _symmetric_to_row, _symmetric_to_column) = remap_colors(
+            eltype(ag), _color, maximum(_color), size(A)...
+        )
+        return (
+            _color,
+            _star_set,
+            _row_color,
+            _column_color,
+            _symmetric_to_row,
+            _symmetric_to_column,
+        )
+    end
+    (color, star_set, row_color, column_color, symmetric_to_row, symmetric_to_column) = argmin(
+        t -> maximum(t[3]) + maximum(t[4]), outputs_by_order
+    )  # can't use ncolors without computing the full result
     if speed_setting isa WithResult
         symmetric_result = StarSetColoringResult(A_and_Aᵀ, ag, color, star_set)
-        return BicoloringResult(A, ag, symmetric_result, R)
-    else
-        row_color, column_color, _ = remap_colors(
-            eltype(ag), color, maximum(color), size(A)...
+        return BicoloringResult(
+            A,
+            ag,
+            symmetric_result,
+            row_color,
+            column_color,
+            symmetric_to_row,
+            symmetric_to_column,
+            R,
         )
+    else
         return row_color, column_color
     end
 end
@@ -325,16 +376,38 @@ function _coloring(
     symmetric_pattern::Bool,
 ) where {R}
     A_and_Aᵀ, edge_to_index = bidirectional_pattern(A; symmetric_pattern)
-    ag = AdjacencyGraph(A_and_Aᵀ, edge_to_index; has_diagonal=false)
-    vertices_in_order = vertices(ag, algo.order)
-    color, tree_set = acyclic_coloring(ag, vertices_in_order, algo.postprocessing)
+    ag = AdjacencyGraph(A_and_Aᵀ, edge_to_index, 0; augmented_graph=true)
+    outputs_by_order = map(algo.orders) do order
+        vertices_in_order = vertices(ag, order)
+        _color, _tree_set = acyclic_coloring(ag, vertices_in_order, algo.postprocessing)
+        (_row_color, _column_color, _symmetric_to_row, _symmetric_to_column) = remap_colors(
+            eltype(ag), _color, maximum(_color), size(A)...
+        )
+        return (
+            _color,
+            _tree_set,
+            _row_color,
+            _column_color,
+            _symmetric_to_row,
+            _symmetric_to_column,
+        )
+    end
+    (color, tree_set, row_color, column_color, symmetric_to_row, symmetric_to_column) = argmin(
+        t -> maximum(t[3]) + maximum(t[4]), outputs_by_order
+    )  # can't use ncolors without computing the full result
     if speed_setting isa WithResult
         symmetric_result = TreeSetColoringResult(A_and_Aᵀ, ag, color, tree_set, R)
-        return BicoloringResult(A, ag, symmetric_result, R)
-    else
-        row_color, column_color, _ = remap_colors(
-            eltype(ag), color, maximum(color), size(A)...
+        return BicoloringResult(
+            A,
+            ag,
+            symmetric_result,
+            row_color,
+            column_color,
+            symmetric_to_row,
+            symmetric_to_column,
+            R,
         )
+    else
         return row_color, column_color
     end
 end
diff --git a/src/matrices.jl b/src/matrices.jl
index eaf79882..453e3062 100644
--- a/src/matrices.jl
+++ b/src/matrices.jl
@@ -62,24 +62,57 @@ function respectful_similar(A::Union{Symmetric,Hermitian}, ::Type{T}) where {T}
 end
 
 """
-    same_pattern(A, B)
+    compatible_pattern(A::AbstractMatrix, bg::BipartiteGraph)
+    compatible_pattern(A::AbstractMatrix, ag::AdjacencyGraph, uplo::Symbol)
 
-Perform a partial equality check on the sparsity patterns of `A` and `B`:
+Perform a coarse compatibility check between the sparsity pattern of `A`
+and the reference sparsity pattern encoded in a graph structure.
 
-- 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
+This function only checks necessary conditions for the two sparsity patterns to match.
+In particular, it may return `true` even if the patterns are not identical.
+
+When A is a `SparseMatrixCSC`, additional checks on the sparsity structure are performed.
+
+# Return value
+- `true`  : the sparsity patterns are potentially compatible
+- `false` : the sparsity patterns are definitely incompatible
 """
-same_pattern(A, B) = size(A) == size(B)
+compatible_pattern(A::AbstractMatrix, bg::BipartiteGraph) = size(A) == size(bg.S2)
+function compatible_pattern(A::SparseMatrixCSC, bg::BipartiteGraph)
+    return size(A) == size(bg.S2) && nnz(A) == nnz(bg.S2)
+end
+
+function compatible_pattern(A::AbstractMatrix, ag::AdjacencyGraph, uplo::Symbol)
+    return size(A) == size(ag.S)
+end
+
+function compatible_pattern(A::SparseMatrixCSC, ag::AdjacencyGraph, uplo::Symbol)
+    nnzS = (uplo == :L || uplo == :U) ? (nb_edges(ag) + ag.nb_self_loops) : nnz(ag.S)
+    return size(A) == size(ag.S) && nnz(A) == nnzS
+end
 
-function same_pattern(
-    A::Union{SparseMatrixCSC,SparsityPatternCSC},
-    B::Union{SparseMatrixCSC,SparsityPatternCSC},
-)
-    return size(A) == size(B) && nnz(A) == nnz(B)
+function check_compatible_pattern(A::AbstractMatrix, bg::BipartiteGraph)
+    if !compatible_pattern(A, bg)
+        throw(DimensionMismatch("`A` and `bg.S2` must have the same sparsity pattern."))
+    end
 end
 
-function check_same_pattern(A, S)
-    if !same_pattern(A, S)
-        throw(DimensionMismatch("`A` and `S` must have the same sparsity pattern."))
+function check_compatible_pattern(A::AbstractMatrix, ag::AdjacencyGraph, uplo::Symbol)
+    if !compatible_pattern(A, ag, uplo)
+        if uplo == :L
+            throw(
+                DimensionMismatch(
+                    "`A` and `tril(ag.S)` must have the same sparsity pattern."
+                ),
+            )
+        elseif uplo == :U
+            throw(
+                DimensionMismatch(
+                    "`A` and `triu(ag.S)` must have the same sparsity pattern."
+                ),
+            )
+        else  # uplo == :F
+            throw(DimensionMismatch("`A` and `ag.S` must have the same sparsity pattern."))
+        end
     end
 end
diff --git a/src/optimal.jl b/src/optimal.jl
new file mode 100644
index 00000000..2a70a59d
--- /dev/null
+++ b/src/optimal.jl
@@ -0,0 +1,34 @@
+"""
+    OptimalColoringAlgorithm
+
+Coloring algorithm that relies on mathematical programming with [JuMP](https://jump.dev/) to find an optimal coloring.
+
+!!! warning
+    This algorithm is only available when JuMP is loaded. If you encounter a method error, run `import JuMP` in your REPL and try again.
+    It only works for nonsymmetric, unidirectional colorings problems.
+
+!!! danger
+    The coloring problem is NP-hard, so it is unreasonable to expect an optimal solution in reasonable time for large instances.
+
+# Constructor
+
+    OptimalColoringAlgorithm(optimizer; silent::Bool=true, assert_solved::Bool=true)
+
+The `optimizer` argument can be any JuMP-compatible optimizer.
+However, the problem formulation is best suited to CP-SAT optimizers like [MiniZinc](https://github.com/jump-dev/MiniZinc.jl).
+You can use [`optimizer_with_attributes`](https://jump.dev/JuMP.jl/stable/api/JuMP/#optimizer_with_attributes) to set solver-specific parameters.
+
+# Keyword arguments
+
+- `silent`: whether to suppress solver output
+- `assert_solved`: whether to check that the solver found an optimal solution (as opposed to running out of time for example)
+"""
+struct OptimalColoringAlgorithm{O} <: ADTypes.AbstractColoringAlgorithm
+    optimizer::O
+    silent::Bool
+    assert_solved::Bool
+end
+
+function OptimalColoringAlgorithm(optimizer; silent::Bool=true, assert_solved::Bool=true)
+    return OptimalColoringAlgorithm(optimizer, silent, assert_solved)
+end
diff --git a/src/order.jl b/src/order.jl
index aab5bb23..20884c6a 100644
--- a/src/order.jl
+++ b/src/order.jl
@@ -334,7 +334,7 @@ function vertices(
 
         π[index] = u
         for v in neighbors(g, u)
-            !has_diagonal(g) || (u == v && continue)
+            augmented_graph(g) || (u == v && continue)
             dv = degrees[v]
             dv == -1 && continue
             update_bucket!(db, v, dv; degtype, direction)
diff --git a/src/postprocessing.jl b/src/postprocessing.jl
new file mode 100644
index 00000000..7b2b58f5
--- /dev/null
+++ b/src/postprocessing.jl
@@ -0,0 +1,168 @@
+## Postprocessing
+
+function postprocess!(
+    color::AbstractVector{<:Integer},
+    star_or_tree_set::Union{StarSet,TreeSet},
+    g::AdjacencyGraph,
+    offsets::AbstractVector{<:Integer},
+)
+    S = pattern(g)
+    edge_to_index = edge_indices(g)
+    # flag which colors are actually used during decompression
+    nb_colors = maximum(color)
+    color_used = zeros(Bool, nb_colors)
+
+    # nonzero diagonal coefficients force the use of their respective color (there can be no neutral colors if the diagonal is fully nonzero)
+    if !augmented_graph(g)
+        for i in axes(S, 1)
+            if !iszero(S[i, i])
+                color_used[color[i]] = true
+            end
+        end
+    end
+
+    if star_or_tree_set isa StarSet
+        # star_or_tree_set is a StarSet
+        postprocess_with_star_set!(g, color_used, color, star_or_tree_set)
+    else
+        # star_or_tree_set is a TreeSet
+        postprocess_with_tree_set!(color_used, color, star_or_tree_set)
+    end
+
+    # if at least one of the colors is useless, modify the color assignments of vertices
+    if any(!, color_used)
+        num_colors_useless = 0
+
+        # determine what are the useless colors and compute the offsets
+        for ci in 1:nb_colors
+            if color_used[ci]
+                offsets[ci] = num_colors_useless
+            else
+                num_colors_useless += 1
+            end
+        end
+
+        # assign the neutral color to every vertex with a useless color and remap the colors
+        for i in eachindex(color)
+            ci = color[i]
+            if !color_used[ci]
+                # assign the neutral color
+                color[i] = 0
+            else
+                # remap the color to not have any gap
+                color[i] -= offsets[ci]
+            end
+        end
+    end
+    return color
+end
+
+function postprocess_with_star_set!(
+    g::AdjacencyGraph,
+    color_used::Vector{Bool},
+    color::AbstractVector{<:Integer},
+    star_set::StarSet,
+)
+    S = pattern(g)
+    edge_to_index = edge_indices(g)
+
+    # only the colors of the hubs are used
+    (; star, hub) = star_set
+    nb_trivial_stars = 0
+
+    # Iterate through all non-trivial stars
+    for s in eachindex(hub)
+        h = hub[s]
+        if h > 0
+            color_used[color[h]] = true
+        else
+            nb_trivial_stars += 1
+        end
+    end
+
+    # Process the trivial stars (if any)
+    if nb_trivial_stars > 0
+        rvS = rowvals(S)
+        for j in axes(S, 2)
+            for k in nzrange(S, j)
+                i = rvS[k]
+                if i > j
+                    index_ij = edge_to_index[k]
+                    s = star[index_ij]
+                    h = hub[s]
+                    if h < 0
+                        h = abs(h)
+                        spoke = h == j ? i : j
+                        if color_used[color[spoke]]
+                            # Switch the hub and the spoke to possibly avoid adding one more used color
+                            hub[s] = spoke
+                        else
+                            # Keep the current hub
+                            color_used[color[h]] = true
+                        end
+                    end
+                end
+            end
+        end
+    end
+    return color_used
+end
+
+function postprocess_with_tree_set!(
+    color_used::Vector{Bool}, color::AbstractVector{<:Integer}, tree_set::TreeSet
+)
+    # only the colors of non-leaf vertices are used
+    (; reverse_bfs_orders, is_star, tree_edge_indices, nt) = tree_set
+    nb_trivial_trees = 0
+
+    # Iterate through all non-trivial trees
+    for k in 1:nt
+        # Position of the first edge in the tree
+        first = tree_edge_indices[k]
+
+        # Total number of edges in the tree
+        ne_tree = tree_edge_indices[k + 1] - first
+
+        # Check if we have more than one edge in the tree (non-trivial tree)
+        if ne_tree > 1
+            # Determine if the tree is a star
+            if is_star[k]
+                # It is a non-trivial star and only the color of the hub is needed
+                (_, hub) = reverse_bfs_orders[first]
+                color_used[color[hub]] = true
+            else
+                # It is not a star and both colors are needed during the decompression
+                (i, j) = reverse_bfs_orders[first]
+                color_used[color[i]] = true
+                color_used[color[j]] = true
+            end
+        else
+            nb_trivial_trees += 1
+        end
+    end
+
+    # Process the trivial trees (if any)
+    if nb_trivial_trees > 0
+        for k in 1:nt
+            # Position of the first edge in the tree
+            first = tree_edge_indices[k]
+
+            # Total number of edges in the tree
+            ne_tree = tree_edge_indices[k + 1] - first
+
+            # Check if we have exactly one edge in the tree
+            if ne_tree == 1
+                (i, j) = reverse_bfs_orders[first]
+                if color_used[color[i]]
+                    # Make i the root to avoid possibly adding one more used color
+                    # Switch it with the (only) leaf
+                    reverse_bfs_orders[first] = (j, i)
+                else
+                    # Keep j as the root
+                    color_used[color[j]] = true
+                end
+            end
+        end
+    end
+    return color_used
+end
diff --git a/src/result.jl b/src/result.jl
index b5b9cdd2..0d47cf13 100644
--- a/src/result.jl
+++ b/src/result.jl
@@ -82,6 +82,9 @@ Create a color-indexed vector `group` such that `i ∈ group[c]` iff `color[i] =
 Assumes the colors are contiguously numbered from `0` to some `cmax`.
 """
 function group_by_color(::Type{T}, color::AbstractVector) where {T<:Integer}
+    if isempty(color)
+        return typeof(view(T[], 1:0))[]
+    end
     cmin, cmax = extrema(color)
     @assert cmin >= 0
     # Compute group sizes and offsets for a joint storage
@@ -401,31 +404,31 @@ function TreeSetColoringResult(
     decompression_eltype::Type{R},
 ) where {T<:Integer,R}
     (; reverse_bfs_orders, tree_edge_indices, nt) = tree_set
-    (; S) = ag
+    (; S, nb_self_loops) = ag
     nvertices = length(color)
     group = group_by_color(T, color)
     rv = rowvals(S)
 
     # Vector for the decompression of the diagonal coefficients
-    diagonal_indices = T[]
-    diagonal_nzind = T[]
-    ndiag = 0
+    diagonal_indices = Vector{T}(undef, nb_self_loops)
+    diagonal_nzind = Vector{T}(undef, nb_self_loops)
 
-    if has_diagonal(ag)
+    if !augmented_graph(ag)
+        l = 0
         for j in axes(S, 2)
             for k in nzrange(S, j)
                 i = rv[k]
                 if i == j
-                    push!(diagonal_indices, i)
-                    push!(diagonal_nzind, k)
-                    ndiag += 1
+                    l += 1
+                    diagonal_indices[l] = i
+                    diagonal_nzind[l] = k
                 end
             end
         end
     end
 
     # Vectors for the decompression of the off-diagonal coefficients
-    nedges = (nnz(S) - ndiag) ÷ 2
+    nedges = nb_edges(ag)
     lower_triangle_offsets = Vector{T}(undef, nedges)
     upper_triangle_offsets = Vector{T}(undef, nedges)
 
@@ -686,14 +689,15 @@ function BicoloringResult(
     A::AbstractMatrix,
     ag::AdjacencyGraph{T},
     symmetric_result::AbstractColoringResult{:symmetric,:column},
+    row_color::Vector{T},
+    column_color::Vector{T},
+    symmetric_to_row::Vector{T},
+    symmetric_to_column::Vector{T},
     decompression_eltype::Type{R},
 ) where {T,R}
     m, n = size(A)
     symmetric_color = column_colors(symmetric_result)
     num_sym_colors = maximum(symmetric_color)
-    row_color, column_color, symmetric_to_row, symmetric_to_column = remap_colors(
-        T, symmetric_color, num_sym_colors, m, n
-    )
     column_group = group_by_color(T, column_color)
     row_group = group_by_color(T, row_color)
     Br_and_Bc = Matrix{R}(undef, n + m, num_sym_colors)
diff --git a/test/Project.toml b/test/Project.toml
index ab595aba..5e41ceb2 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -12,12 +12,15 @@ CliqueTrees = "60701a23-6482-424a-84db-faee86b9b1f8"
 Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+HiGHS = "87dc4568-4c63-4d18-b0c0-bb2238e4078b"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
-JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899"
+JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MatrixDepot = "b51810bb-c9f3-55da-ae3c-350fc1fbce05"
+MiniZinc = "a7f392d2-6c35-496e-b8cc-0974fbfcbf91"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+cuSPARSE = "b26da814-b3bc-49ef-b0ee-c816305aa060"
diff --git a/test/adtypes.jl b/test/adtypes.jl
index bbb124a6..da0807d3 100644
--- a/test/adtypes.jl
+++ b/test/adtypes.jl
@@ -1,26 +1,49 @@
 using ADTypes: ADTypes
 using SparseArrays
+using LinearAlgebra
 using SparseMatrixColorings
 using Test
 
-@testset "Column coloring" begin
-    problem = ColoringProblem(; structure=:nonsymmetric, partition=:column)
-    algo = ADTypes.NoColoringAlgorithm()
-    A = sprand(10, 20, 0.1)
-    result = coloring(A, problem, algo)
-    B = compress(A, result)
-    @test size(B) == size(A)
-    @test column_colors(result) == ADTypes.column_coloring(A, algo)
-    @test decompress(B, result) == A
-end
+@testset "NoColoringAlgorithm" begin
+    @testset "Column coloring" begin
+        problem = ColoringProblem(; structure=:nonsymmetric, partition=:column)
+        algo = ADTypes.NoColoringAlgorithm()
+        A = sprand(10, 20, 0.3)
+        result = coloring(A, problem, algo)
+        B = compress(A, result)
+        @test size(B) == size(A)
+        @test column_colors(result) == ADTypes.column_coloring(A, algo)
+        @test decompress(B, result) == A
+    end
+
+    @testset "Row coloring" begin
+        problem = ColoringProblem(; structure=:nonsymmetric, partition=:row)
+        algo = ADTypes.NoColoringAlgorithm()
+        A = sprand(10, 20, 0.3)
+        result = coloring(A, problem, algo)
+        B = compress(A, result)
+        @test size(B) == size(A)
+        @test row_colors(result) == ADTypes.row_coloring(A, algo)
+        @test decompress(B, result) == A
+    end
+
+    @testset "Symmetric coloring" begin
+        problem = ColoringProblem(; structure=:symmetric, partition=:column)
+        algo = ADTypes.NoColoringAlgorithm()
+        A = Symmetric(sprand(20, 20, 0.3))
+        result = coloring(A, problem, algo)
+        B = compress(A, result)
+        @test size(B) == size(A)
+        @test column_colors(result) == ADTypes.column_coloring(A, algo)
+        @test decompress(B, result) == A
+    end
 
-@testset "Row coloring" begin
-    problem = ColoringProblem(; structure=:nonsymmetric, partition=:row)
-    algo = ADTypes.NoColoringAlgorithm()
-    A = sprand(10, 20, 0.1)
-    result = coloring(A, problem, algo)
-    B = compress(A, result)
-    @test size(B) == size(A)
-    @test row_colors(result) == ADTypes.row_coloring(A, algo)
-    @test decompress(B, result) == A
+    @testset "Bicoloring" begin
+        problem = ColoringProblem(; structure=:nonsymmetric, partition=:bidirectional)
+        algo = ADTypes.NoColoringAlgorithm()
+        A = sprand(10, 20, 0.3)
+        result = coloring(A, problem, algo)
+        Br, Bc = compress(A, result)
+        @test decompress(Br, Bc, result) == A
+    end
 end
diff --git a/test/allocations.jl b/test/allocations.jl
index caed10f8..31305e63 100644
--- a/test/allocations.jl
+++ b/test/allocations.jl
@@ -154,7 +154,7 @@ end
             b64 = @b fast_coloring(A64, problem, algo)
             b32 = @b fast_coloring(A32, problem, algo)
             # check that we allocate no more than 50% + epsilon with Int32
-            @test b32.bytes < 0.6 * b64.bytes
+            @test b32.bytes < 0.7 * b64.bytes
         end
     end
 end;
diff --git a/test/check.jl b/test/check.jl
index 4e9e5a61..8cf2c495 100644
--- a/test/check.jl
+++ b/test/check.jl
@@ -1,9 +1,12 @@
 using LinearAlgebra
+using SparseArrays
 using SparseMatrixColorings:
     structurally_orthogonal_columns,
     symmetrically_orthogonal_columns,
     structurally_biorthogonal,
     directly_recoverable_columns,
+    substitutable_columns,
+    substitutable_bidirectional,
     what_fig_41,
     efficient_fig_1
 using Test
@@ -184,3 +187,135 @@ For coefficient (i=1, j=1) with colors (ci=0, cj=0):
 """,
     ) structurally_biorthogonal(A, [0, 2, 2, 3], [0, 2, 2, 2, 3], verbose=true)
 end
+
+@testset "Substitutable columns" begin
+    A1 = [
+        1 1 1 1 1
+        1 1 0 0 0
+        1 0 1 0 0
+        1 0 0 1 0
+        1 0 0 0 1
+    ]
+    B1 = [
+        1 6 7 8 9
+        6 2 0 0 0
+        7 0 3 0 0
+        8 0 0 4 0
+        9 0 0 0 5
+    ]
+    A2 = [
+        1 1 0 0 0
+        1 1 1 0 0
+        0 1 1 1 0
+        0 0 1 1 1
+        0 0 0 1 1
+    ]
+    B2 = [
+        5 1 0 0 0
+        1 6 2 0 0
+        0 2 7 3 0
+        0 0 3 8 4
+        0 0 0 4 9
+    ]
+    A3 = [
+        0 1 1 1 1
+        1 0 1 1 1
+        1 1 0 1 1
+        1 1 1 0 1
+        1 1 1 1 0
+    ]
+    B3 = [
+        0 1 2 3 4
+        1 0 5 6 7
+        2 5 0 8 9
+        3 6 8 0 10
+        4 7 9 10 0
+    ]
+
+    # success
+
+    @test substitutable_columns(A1, B1, [1, 2, 2, 2, 2])
+    @test substitutable_columns(A2, B2, [1, 2, 3, 1, 2])
+    @test substitutable_columns(A3, B3, [1, 2, 3, 4, 0])
+
+    # failure
+
+    @test !substitutable_columns(A1, B1, [1, 1, 1, 1])
+    log = (:warn, "4 colors provided for 5 columns.")
+    @test_logs log substitutable_columns(A1, B1, [1, 1, 1, 1]; verbose=true)
+
+    @test !substitutable_columns(A1, B1, [1, 1, 1, 1, 1])
+    @test_logs (
+        :warn,
+        """
+For coefficient (i=1, j=1) with colors (ci=1, cj=1):
+- For the row 5 in row color ci=1, A[5, 1] is ordered after A[1, 1].
+- For the column 5 in column color cj=1, A[1, 5] is ordered after A[1, 1].
+""",
+    ) substitutable_columns(A1, B1, [1, 1, 1, 1, 1]; verbose=true)
+
+    @test !substitutable_columns(A2, B2, [1, 2, 0, 1, 2])
+    @test_logs (
+        :warn,
+        """
+For coefficient (i=3, j=3) with colors (ci=0, cj=0):
+- Row color ci=0 is neutral.
+- Column color cj=0 is neutral.
+""",
+    ) substitutable_columns(A2, B2, [1, 2, 0, 1, 2]; verbose=true)
+
+    @test !substitutable_columns(A3, B3, [0, 1, 2, 3, 3])
+    @test_logs (
+        :warn,
+        """
+For coefficient (i=1, j=4) with colors (ci=0, cj=3):
+- Row color ci=0 is neutral.
+- For the column 5 in column color cj=3, A[1, 5] is ordered after A[1, 4].
+""",
+    ) substitutable_columns(A3, B3, [0, 1, 2, 3, 3]; verbose=true)
+
+    @test !substitutable_columns(A3, B3, [1, 2, 3, 3, 0])
+    @test_logs (
+        :warn,
+        """
+For coefficient (i=3, j=5) with colors (ci=3, cj=0):
+- For the row 4 in row color ci=3, A[4, 5] is ordered after A[3, 5].
+- Column color cj=0 is neutral.
+""",
+    ) substitutable_columns(A3, B3, [1, 2, 3, 3, 0]; verbose=true)
+end
+
+@testset "Substitutable bidirectional" begin
+    A = [
+        1 0 0
+        0 1 0
+        0 0 1
+    ]
+    B = [
+        1 0 0
+        0 2 0
+        0 0 3
+    ]
+
+    # success
+
+    @test substitutable_bidirectional(A, B, [1, 0, 0], [0, 1, 1])
+
+    # failure
+
+    log = (:warn, "2 colors provided for 3 columns.")
+    @test_logs log !substitutable_bidirectional(A, B, [1, 0, 0], [0, 1]; verbose=true)
+
+    log = (:warn, "4 colors provided for 3 rows.")
+    @test_logs log !substitutable_bidirectional(A, B, [1, 0, 0, 1], [0, 1, 1]; verbose=true)
+end
+
+# See https://github.com/JuliaDiff/SparseMatrixColorings.jl/pull/300
+@testset "Empty matrix" begin
+    problem = ColoringProblem(; structure=:symmetric, partition=:column)
+    algo = GreedyColoringAlgorithm(; decompression=:substitution)
+    S = spzeros(Int, 0, 0)
+    result = coloring(S, problem, algo)
+    @test isempty(result.color)
+    @test isempty(result.group)
+end
diff --git a/test/constant.jl b/test/constant.jl
index 5de881b8..97c7f238 100644
--- a/test/constant.jl
+++ b/test/constant.jl
@@ -1,16 +1,20 @@
 using ADTypes: ADTypes
 using SparseMatrixColorings
+using SparseMatrixColorings: InvalidColoringError
 using Test
 
-matrix_template = ones(100, 200)
+matrix_template = ones(Bool, 10, 20)
+sym_matrix_template = ones(Bool, 10, 10)
 
 @testset "Column coloring" begin
     problem = ColoringProblem(; structure=:nonsymmetric, partition=:column)
-    color = rand(1:5, size(matrix_template, 2))
+    color = collect(1:20)
     algo = ConstantColoringAlgorithm(matrix_template, color; partition=:column)
-    wrong_algo = ConstantColoringAlgorithm(matrix_template, color; partition=:row)
+    wrong_algo = ConstantColoringAlgorithm{:row}(matrix_template, color)
+    wrong_color = ConstantColoringAlgorithm{:column}(matrix_template, ones(Int, 20))
     @test_throws DimensionMismatch coloring(transpose(matrix_template), problem, algo)
     @test_throws MethodError coloring(matrix_template, problem, wrong_algo)
+    @test_throws InvalidColoringError coloring(matrix_template, problem, wrong_color)
     result = coloring(matrix_template, problem, algo)
     @test column_colors(result) == color
     @test ADTypes.column_coloring(matrix_template, algo) == color
@@ -19,9 +23,13 @@ end
 
 @testset "Row coloring" begin
     problem = ColoringProblem(; structure=:nonsymmetric, partition=:row)
-    color = rand(1:5, size(matrix_template, 1))
+    color = collect(1:10)
     algo = ConstantColoringAlgorithm(matrix_template, color; partition=:row)
+    wrong_algo = ConstantColoringAlgorithm{:column}(matrix_template, color)
+    wrong_color = ConstantColoringAlgorithm{:row}(matrix_template, ones(Int, 10))
     @test_throws DimensionMismatch coloring(transpose(matrix_template), problem, algo)
+    @test_throws MethodError coloring(matrix_template, problem, wrong_algo)
+    @test_throws InvalidColoringError coloring(matrix_template, problem, wrong_color)
     result = coloring(matrix_template, problem, algo)
     @test row_colors(result) == color
     @test ADTypes.row_coloring(matrix_template, algo) == color
@@ -29,9 +37,22 @@ end
 end
 
 @testset "Symmetric coloring" begin
-    wrong_problem = ColoringProblem(; structure=:symmetric, partition=:column)
-    color = rand(1:5, size(matrix_template, 2))
-    algo = ConstantColoringAlgorithm(matrix_template, color; partition=:column)
-    @test_throws MethodError coloring(matrix_template, wrong_problem, algo)
-    @test_throws MethodError ADTypes.symmetric_coloring(matrix_template, algo)
+    problem = ColoringProblem(; structure=:symmetric, partition=:column)
+    color = collect(1:10)
+    algo = ConstantColoringAlgorithm(
+        sym_matrix_template, color; partition=:column, structure=:symmetric
+    )
+    wrong_algo = ConstantColoringAlgorithm{:column,:nonsymmetric}(
+        sym_matrix_template, color
+    )
+    wrong_color = ConstantColoringAlgorithm{:column,:symmetric}(
+        sym_matrix_template, ones(Int, 20)
+    )
+    @test_throws DimensionMismatch coloring(matrix_template, problem, algo)
+    @test_throws MethodError coloring(sym_matrix_template, problem, wrong_algo)
+    @test_throws InvalidColoringError coloring(sym_matrix_template, problem, wrong_color)
+    result = coloring(sym_matrix_template, problem, algo)
+    @test column_colors(result) == color
+    @test ADTypes.symmetric_coloring(sym_matrix_template, algo) == color
+    @test_throws MethodError ADTypes.column_coloring(sym_matrix_template, algo)
 end
diff --git a/test/cuda.jl b/test/cuda.jl
index 6549b168..8b27e07b 100644
--- a/test/cuda.jl
+++ b/test/cuda.jl
@@ -1,4 +1,4 @@
-using CUDA.CUSPARSE: CuSparseMatrixCSC, CuSparseMatrixCSR
+using cuSPARSE: CuSparseMatrixCSC, CuSparseMatrixCSR
 using LinearAlgebra
 using SparseArrays
 using SparseMatrixColorings
@@ -6,8 +6,6 @@ import SparseMatrixColorings as SMC
 using StableRNGs
 using Test
 
-include("utils.jl")
-
 rng = StableRNG(63)
 
 asymmetric_params = vcat(
diff --git a/test/graph.jl b/test/graph.jl
index 9c8f66c1..773a0206 100644
--- a/test/graph.jl
+++ b/test/graph.jl
@@ -15,7 +15,8 @@ using Test
 ## SparsityPatternCSC
 
 @testset "SparsityPatternCSC" begin
-    @test eltype(SparsityPatternCSC(sprand(10, 10, 0.1))) == Int
+    @test eltype(SparsityPatternCSC(sprand(10, 10, 0.1))) == Bool
+    @test SparseArrays.indtype(SparsityPatternCSC(sprand(10, 10, 0.1))) == Int
     @testset "Transpose" begin
         for _ in 1:1000
             m, n = rand(100:1000), rand(100:1000)
@@ -166,7 +167,7 @@ end;
     @test degree(g, 7) == 6
     @test degree(g, 8) == 6
 
-    g = AdjacencyGraph(transpose(A) * A; has_diagonal=false)
+    g = AdjacencyGraph(transpose(A) * A; augmented_graph=true)
     # wrong degree
     @test degree(g, 1) == 4
     @test degree(g, 2) == 4
diff --git a/test/matrices.jl b/test/matrices.jl
index 1571497b..fa76946c 100644
--- a/test/matrices.jl
+++ b/test/matrices.jl
@@ -1,7 +1,12 @@
 using LinearAlgebra
 using SparseArrays
 using SparseMatrixColorings:
-    check_same_pattern, matrix_versions, respectful_similar, same_pattern
+    BipartiteGraph,
+    AdjacencyGraph,
+    matrix_versions,
+    respectful_similar,
+    compatible_pattern,
+    check_compatible_pattern
 using StableRNGs
 using Test
 
@@ -33,7 +38,7 @@ same_view(::Adjoint, ::Adjoint) = true
     end
 end
 
-@testset "Sparsity pattern" begin
+@testset "Compatible sparsity pattern -- BipartiteGraph" begin
     S = sparse([
         0 1 1
         0 1 0
@@ -44,9 +49,33 @@ end
     A2 = copy(S)
     A2[1, 1] = 1
 
-    @test same_pattern(A1, S)
-    @test !same_pattern(A2, S)
-    @test same_pattern(Matrix(A2), S)
+    bg1 = BipartiteGraph(A1)
+    bg2 = BipartiteGraph(A2)
+    @test compatible_pattern(S, bg1)
+    @test !compatible_pattern(S, bg2)
+    @test compatible_pattern(Matrix(S), bg1)
 
-    @test_throws DimensionMismatch check_same_pattern(A2, S)
+    @test_throws DimensionMismatch check_compatible_pattern(S, bg2)
+end
+
+@testset "Compatible sparsity pattern -- AdjacencyGraph" begin
+    S = sparse([
+        1 0 1
+        0 1 1
+        1 1 0
+    ])
+
+    A1 = copy(S)
+    A2 = copy(S)
+    A2[3, 3] = 1
+
+    ag1 = AdjacencyGraph(A1)
+    ag2 = AdjacencyGraph(A2)
+    for (op, uplo) in ((tril, :L), (triu, :U), (identity, :F))
+        @test compatible_pattern(op(S), ag1, uplo)
+        @test !compatible_pattern(op(S), ag2, uplo)
+        @test compatible_pattern(Matrix(op(S)), ag1, uplo)
+
+        @test_throws DimensionMismatch check_compatible_pattern(op(S), ag2, uplo)
+    end
 end
diff --git a/test/optimal.jl b/test/optimal.jl
new file mode 100644
index 00000000..36bf91c6
--- /dev/null
+++ b/test/optimal.jl
@@ -0,0 +1,46 @@
+using SparseArrays
+using SparseMatrixColorings
+using StableRNGs
+using Test
+using JuMP
+using MiniZinc
+using HiGHS
+
+rng = StableRNG(0)
+
+asymmetric_params = vcat(
+    [(10, 20, p) for p in (0.0:0.1:0.5)], [(20, 10, p) for p in (0.0:0.1:0.5)]
+)
+
+algo = GreedyColoringAlgorithm()
+optalgo = OptimalColoringAlgorithm(() -> MiniZinc.Optimizer{Float64}("highs"); silent=false)
+
+# TODO: reactivate tests once https://github.com/jump-dev/MiniZinc.jl/issues/103 is fixed
+
+@testset "Column coloring" begin
+    problem = ColoringProblem(; structure=:nonsymmetric, partition=:column)
+    for (m, n, p) in asymmetric_params
+        A = sprand(rng, m, n, p)
+        result = coloring(A, problem, algo)
+        @test_skip ncolors(result) >= ncolors(coloring(A, problem, optalgo))
+    end
+end
+
+@testset "Row coloring" begin
+    problem = ColoringProblem(; structure=:nonsymmetric, partition=:row)
+    for (m, n, p) in asymmetric_params
+        A = sprand(rng, m, n, p)
+        result = coloring(A, problem, algo)
+        @test_skip ncolors(result) >= ncolors(coloring(A, problem, optalgo))
+    end
+end
+
+@testset "Too big" begin
+    A = sprand(rng, Bool, 100, 100, 0.1)
+    optalgo_timelimit = OptimalColoringAlgorithm(
+        optimizer_with_attributes(HiGHS.Optimizer, "time_limit" => 10.0); # 1 second
+        silent=false,
+        assert_solved=false,
+    )
+    @test_throws AssertionError coloring(A, ColoringProblem(), optalgo_timelimit)
+end
diff --git a/test/order.jl b/test/order.jl
index 9bb9e6dc..9cae16f1 100644
--- a/test/order.jl
+++ b/test/order.jl
@@ -61,16 +61,14 @@ end;
 end;
 
 @testset "LargestFirst" begin
-    for has_diagonal in (false, true)
-        A = sparse([
-            0 1 0 0
-            1 0 1 1
-            0 1 0 1
-            0 1 1 0
-        ])
-        ag = AdjacencyGraph(A; has_diagonal)
-        @test vertices(ag, LargestFirst()) == [2, 3, 4, 1]
-    end
+    A = sparse([
+        0 1 0 0
+        1 0 1 1
+        0 1 0 1
+        0 1 1 0
+    ])
+    ag = AdjacencyGraph(A)
+    @test vertices(ag, LargestFirst()) == [2, 3, 4, 1]
 
     A = sparse([
         1 1 0 0
@@ -146,3 +144,99 @@ end;
         @test isperm(π)
     end
 end
+
+@testset "Multiple orders" begin
+    # I used brute force to find examples where LargestFirst is *strictly* better than NaturalOrder, just to check that the best order is indeed selected when multiple orders are provided
+    @testset "Column coloring" begin
+        A = [
+            0 0 1 1
+            0 1 0 1
+            0 0 1 1
+            1 1 0 0
+        ]
+        problem = ColoringProblem{:nonsymmetric,:column}()
+        algo = GreedyColoringAlgorithm(NaturalOrder())
+        better_algo = GreedyColoringAlgorithm((NaturalOrder(), LargestFirst()))
+        @test ncolors(coloring(A, problem, better_algo)) <
+            ncolors(coloring(A, problem, algo))
+    end
+    @testset "Row coloring" begin
+        A = [
+            1 0 0 0
+            0 0 1 0
+            0 1 1 1
+            1 0 0 1
+        ]
+        problem = ColoringProblem{:nonsymmetric,:row}()
+        algo = GreedyColoringAlgorithm(NaturalOrder())
+        better_algo = GreedyColoringAlgorithm((NaturalOrder(), LargestFirst()))
+        @test ncolors(coloring(A, problem, better_algo)) <
+            ncolors(coloring(A, problem, algo))
+    end
+    @testset "Star coloring" begin
+        A = [
+            0 1 0 1 1
+            1 1 0 1 0
+            0 0 1 0 1
+            1 1 0 1 0
+            1 0 1 0 0
+        ]
+        problem = ColoringProblem{:symmetric,:column}()
+        algo = GreedyColoringAlgorithm(NaturalOrder())
+        better_algo = GreedyColoringAlgorithm((NaturalOrder(), LargestFirst()))
+        @test ncolors(coloring(A, problem, better_algo)) <
+            ncolors(coloring(A, problem, algo))
+    end
+    @testset "Acyclic coloring" begin
+        A = [
+            1 0 0 0 0 1 0
+            0 0 0 1 0 0 0
+            0 0 0 1 0 0 0
+            0 1 1 1 0 1 1
+            0 0 0 0 0 0 1
+            1 0 0 1 0 0 1
+            0 0 0 1 1 1 1
+        ]
+        problem = ColoringProblem{:symmetric,:column}()
+        algo = GreedyColoringAlgorithm{:substitution}(NaturalOrder())
+        better_algo = GreedyColoringAlgorithm{:substitution}((
+            NaturalOrder(), LargestFirst()
+        ))
+        @test ncolors(coloring(A, problem, better_algo)) <
+            ncolors(coloring(A, problem, algo))
+    end
+    @testset "Star bicoloring" begin
+        A = [
+            0 1 0 0 0
+            1 0 1 0 0
+            0 1 0 0 1
+            0 0 0 0 0
+            0 0 1 0 1
+        ]
+        problem = ColoringProblem{:nonsymmetric,:bidirectional}()
+        algo = GreedyColoringAlgorithm(NaturalOrder())
+        better_algo = GreedyColoringAlgorithm((NaturalOrder(), LargestFirst()))
+        @test ncolors(coloring(A, problem, better_algo)) <
+            ncolors(coloring(A, problem, algo))
+    end
+    @testset "Acyclic bicoloring" begin
+        A = [
+            0 1 0 1 1 0 1 0 1
+            1 0 0 0 0 0 0 0 1
+            0 0 0 0 0 0 0 0 0
+            1 0 0 1 1 0 1 0 0
+            1 0 0 1 0 0 0 0 0
+            0 0 0 0 0 0 0 0 0
+            1 0 0 1 0 0 0 0 0
+            0 0 0 0 0 0 0 0 0
+            1 1 0 0 0 0 0 0 0
+        ]
+        problem = ColoringProblem{:nonsymmetric,:bidirectional}()
+        algo = GreedyColoringAlgorithm{:substitution}(NaturalOrder())
+        better_algo = GreedyColoringAlgorithm{:substitution}((
+            NaturalOrder(), LargestFirst()
+        ))
+        @test ncolors(coloring(A, problem, better_algo)) <
+            ncolors(coloring(A, problem, algo))
+    end
+end
diff --git a/test/random.jl b/test/random.jl
index 02c1a3a1..406dfea9 100644
--- a/test/random.jl
+++ b/test/random.jl
@@ -81,10 +81,10 @@ end;
     problem = ColoringProblem(; structure=:nonsymmetric, partition=:bidirectional)
     @testset for algo in (
         GreedyColoringAlgorithm(
-            RandomOrder(rng); postprocessing=false, decompression=:direct
+            RandomOrder(StableRNG(0), 0); postprocessing=false, decompression=:direct
         ),
         GreedyColoringAlgorithm(
-            RandomOrder(rng); postprocessing=true, decompression=:direct
+            RandomOrder(StableRNG(0), 0); postprocessing=true, decompression=:direct
         ),
     )
         @testset "$((; m, n, p))" for (m, n, p) in asymmetric_params
@@ -102,10 +102,10 @@ end;
     problem = ColoringProblem(; structure=:nonsymmetric, partition=:bidirectional)
     @testset for algo in (
         GreedyColoringAlgorithm(
-            RandomOrder(rng); postprocessing=false, decompression=:substitution
+            RandomOrder(StableRNG(0), 0); postprocessing=false, decompression=:substitution
         ),
         GreedyColoringAlgorithm(
-            RandomOrder(rng); postprocessing=true, decompression=:substitution
+            RandomOrder(StableRNG(0), 0); postprocessing=true, decompression=:substitution
         ),
     )
         @testset "$((; m, n, p))" for (m, n, p) in asymmetric_params
diff --git a/test/result.jl b/test/result.jl
index 0bea3353..49344aab 100644
--- a/test/result.jl
+++ b/test/result.jl
@@ -1,3 +1,4 @@
+using SparseMatrixColorings
 using SparseMatrixColorings: group_by_color, UnsupportedDecompressionError
 using Test
 
@@ -20,7 +21,7 @@ using Test
 end
 
 @testset "Empty compression" begin
-    A = rand(10, 10)
+    A = zeros(Bool, 10, 10)
     color = zeros(Int, 10)
     problem = ColoringProblem{:nonsymmetric,:column}()
     algo = ConstantColoringAlgorithm(A, color; partition=:column)
diff --git a/test/runtests.jl b/test/runtests.jl
index 81c4ecc9..ec679d4e 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1,7 +1,6 @@
 using Aqua
 using Documenter
 using JET
-using JuliaFormatter
 using SparseMatrixColorings
 using Test
 
@@ -13,20 +12,17 @@ include("utils.jl")
 @testset verbose = true "SparseMatrixColorings" begin
     if get(ENV, "JULIA_SMC_TEST_GROUP", nothing) == "GPU"
         @testset "CUDA" begin
-            using CUDA
+            using CUDA, cuSPARSE
             include("cuda.jl")
         end
     else
         @testset verbose = true "Code quality" begin
             @testset "Aqua" begin
-                Aqua.test_all(SparseMatrixColorings; stale_deps=(; ignore=[:Requires],))
+                Aqua.test_all(SparseMatrixColorings; undocumented_names=true)
             end
             @testset "JET" begin
-                JET.test_package(SparseMatrixColorings; target_defined_modules=true)
-            end
-            @testset "JuliaFormatter" begin
-                @test JuliaFormatter.format(
-                    SparseMatrixColorings; verbose=false, overwrite=false
+                JET.test_package(
+                    SparseMatrixColorings; target_modules=(SparseMatrixColorings,)
                 )
             end
             @testset "Doctests" begin
@@ -58,6 +54,9 @@ include("utils.jl")
             @testset "Constant coloring" begin
                 include("constant.jl")
             end
+            @testset "Optimal coloring" begin
+                include("optimal.jl")
+            end
             @testset "ADTypes coloring algorithms" begin
                 include("adtypes.jl")
             end
@@ -84,7 +83,10 @@ include("utils.jl")
         end
         @testset verbose = true "Performance" begin
             @testset "Type stability" begin
-                include("type_stability.jl")
+                if VERSION < v"1.12"
+                    # TODO: fix JET misbehaving
+                    include("type_stability.jl")
+                end
             end
             @testset "Allocations" begin
                 include("allocations.jl")
diff --git a/test/structured.jl b/test/structured.jl
index d61c7735..857ebbfc 100644
--- a/test/structured.jl
+++ b/test/structured.jl
@@ -2,6 +2,7 @@ using ArrayInterface: ArrayInterface
 using BandedMatrices: BandedMatrix, brand
 using BlockBandedMatrices: BandedBlockBandedMatrix, BlockBandedMatrix
 using LinearAlgebra
+using SparseArrays
 using SparseMatrixColorings
 using Test
 
@@ -68,3 +69,19 @@ end;
         end
     end
 end;
+
+# See https://github.com/JuliaDiff/SparseMatrixColorings.jl/pull/299
+@testset "SparsityPatternCSC $T" for T in [Int, Float32]
+    S = sparse(T[
+        0 0 1 1 0 1
+        1 0 0 0 1 0
+        0 1 0 0 1 0
+        0 1 1 0 0 0
+    ])
+    P = SparseMatrixColorings.SparsityPatternCSC(S)
+    problem = ColoringProblem()
+    algo = GreedyColoringAlgorithm()
+    result = coloring(P, problem, algo)
+    B = compress(S, result)
+    @test decompress(B, result) isa SparseMatrixCSC{T,Int}
+end;
diff --git a/test/type_stability.jl b/test/type_stability.jl
index 5a488596..3c6f9f02 100644
--- a/test/type_stability.jl
+++ b/test/type_stability.jl
@@ -40,11 +40,21 @@ rng = StableRNG(63)
                 ColoringProblem(; structure, partition),
                 GreedyColoringAlgorithm(order; decompression),
             )
+            @test_opt coloring(
+                A,
+                ColoringProblem(; structure, partition),
+                GreedyColoringAlgorithm((NaturalOrder(), order); decompression),
+            )
             @inferred coloring(
                 A,
                 ColoringProblem(; structure, partition),
                 GreedyColoringAlgorithm(order; decompression),
             )
+            @inferred coloring(
+                A,
+                ColoringProblem(; structure, partition),
+                GreedyColoringAlgorithm((NaturalOrder(), order); decompression),
+            )
         end
     end
 end;
diff --git a/test/utils.jl b/test/utils.jl
index 200d07d5..15aa3f9f 100644
--- a/test/utils.jl
+++ b/test/utils.jl
@@ -12,9 +12,16 @@ using SparseMatrixColorings:
     respectful_similar,
     structurally_orthogonal_columns,
     symmetrically_orthogonal_columns,
-    structurally_biorthogonal
+    structurally_biorthogonal,
+    substitutable_columns,
+    substitutable_bidirectional,
+    rank_nonzeros_from_trees
 using Test
 
+const _ALL_ORDERS = (
+    NaturalOrder(), LargestFirst(), SmallestLast(), IncidenceDegree(), DynamicLargestFirst()
+)
+
 function test_coloring_decompression(
     A0::AbstractMatrix,
     problem::ColoringProblem{structure,partition},
@@ -73,7 +80,6 @@ function test_coloring_decompression(
         end
 
         @testset "Recoverability" begin
-            # TODO: find tests for recoverability for substitution decompression
             if decompression == :direct
                 if structure == :nonsymmetric
                     if partition == :column
@@ -93,6 +99,15 @@ function test_coloring_decompression(
             end
         end
 
+        @testset "Substitutable" begin
+            if decompression == :substitution
+                if structure == :symmetric
+                    rank_nonzeros = rank_nonzeros_from_trees(result)
+                    @test substitutable_columns(A0, rank_nonzeros, color)
+                end
+            end
+        end
+
         @testset "Single-color decompression" begin
             if decompression == :direct  # TODO: implement for :substitution too
                 A2 = respectful_similar(A, eltype(B))
@@ -169,6 +184,22 @@ function test_coloring_decompression(
             @show color_vec
         end
     end
+
+    @testset "More orders is better" begin
+        more_orders = (algo.orders..., _ALL_ORDERS...)
+        better_algo = GreedyColoringAlgorithm{decompression}(
+            more_orders; algo.postprocessing
+        )
+        all_algos = [
+            GreedyColoringAlgorithm{decompression}(order; algo.postprocessing) for
+            order in more_orders
+        ]
+        result = coloring(A0, problem, algo)
+        better_result = coloring(A0, problem, better_algo)
+        all_results = [coloring(A0, problem, _algo) for _algo in all_algos]
+        @test ncolors(better_result) <= ncolors(result)
+        @test ncolors(better_result) == minimum(ncolors, all_results)
+    end
 end
 
 function test_bicoloring_decompression(
@@ -215,6 +246,31 @@ function test_bicoloring_decompression(
                 @test structurally_biorthogonal(A0, row_color, column_color)
             end
         end
+
+        if decompression == :substitution
+            @testset "Substitutable" begin
+                rank_nonzeros = rank_nonzeros_from_trees(result)
+                @test substitutable_bidirectional(
+                    A0, rank_nonzeros, row_color, column_color
+                )
+            end
+        end
+    end
+
+    @testset "More orders is better" begin
+        more_orders = (algo.orders..., _ALL_ORDERS...)
+        better_algo = GreedyColoringAlgorithm{decompression}(
+            more_orders; algo.postprocessing
+        )
+        all_algos = [
+            GreedyColoringAlgorithm{decompression}(order; algo.postprocessing) for
+            order in more_orders
+        ]
+        result = coloring(A0, problem, algo)
+        better_result = coloring(A0, problem, better_algo)
+        all_results = [coloring(A0, problem, _algo) for _algo in all_algos]
+        @test ncolors(better_result) <= ncolors(result)
+        @test ncolors(better_result) == minimum(ncolors, all_results)
     end
 end
 

From 52abd519ea674cdd181775b2c4c9da97350654b7 Mon Sep 17 00:00:00 2001
From: Qingyu Qu <52615090+ErikQQY@users.noreply.github.com>
Date: Tue, 28 Apr 2026 22:12:24 +0800
Subject: [PATCH 35/36] Merge main into gd/structured3 again (#317)


From f8ced9b0dfbbeea165a1343d53131d95915825d7 Mon Sep 17 00:00:00 2001
From: Qingyu Qu <52615090+ErikQQY@users.noreply.github.com>
Date: Tue, 28 Apr 2026 23:27:56 +0800
Subject: [PATCH 36/36] Merge main branch again (#318)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* Bump julia-actions/julia-downgrade-compat from 1 to 2 (#258)

Bumps [julia-actions/julia-downgrade-compat](https://github.com/julia-actions/julia-downgrade-compat) from 1 to 2.
- [Release notes](https://github.com/julia-actions/julia-downgrade-compat/releases)
- [Commits](https://github.com/julia-actions/julia-downgrade-compat/compare/v1...v2)

---
updated-dependencies:
- dependency-name: julia-actions/julia-downgrade-compat
  dependency-version: '2'
  dependency-type: direct:production
  update-type: version-update:semver-major
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>

* Bump actions/checkout from 4 to 5 (#259)

* ensure `UnsupportedDecompressionError <: Exception` (#260)

Not sure if there is any benefit, but I suppose the intention is for
this type, as an exception type, to subtype `Exception`.

* Bump peter-evans/create-or-update-comment from 4 to 5 (#261)

* Bump peter-evans/find-comment from 3 to 4 (#262)

Bumps [peter-evans/find-comment](https://github.com/peter-evans/find-comment) from 3 to 4.
- [Release notes](https://github.com/peter-evans/find-comment/releases)
- [Commits](https://github.com/peter-evans/find-comment/compare/v3...v4)

---
updated-dependencies:
- dependency-name: peter-evans/find-comment
  dependency-version: '4'
  dependency-type: direct:production
  update-type: version-update:semver-major
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>

* Allow picking the best of several orders for GreedyColoringAlgorithm (#265)

* Allow picking the best of several orders for GreedyColoringAlgorithm

* Fix type params

* Tuple

* Better tests

* Better test

* Fix foc

* Fix type inference inside closure

* Fix seed

* Test on 1.11

* Avoid duplicate remap_colors

* chore: bump version to 0.4.22 (#268)

* chore: fix typo in package UUID (#269)

* Make any coloring algorithm compatible with SMC (#263)

* Make any coloring algorithm compatible with SMC

* Fix StackOverflow

* Run tests on 1.11

* Start working on optimal coloring

* Fix tests

* Improve feasibility check

* Format

* Fix

* Fix

* Typo

* Use HiGHS

* Remove OptimalColoringAlgorithm

* Remove test deps

* Don't handle forced colors in acyclic

* Run GPU CI on 1.11 too

* Actually use forced colors

* Format

* Deactivate JuliaFormatter (can't figure out why it fails)

* Optimal coloring algorithm with JuMP formulation (#271)

* Add a file postprocessing.jl (#275)

* Rename has_diagonal into augmented_graph (#273)

* Rename has_diagonal into augmented_graph

* Fix a typo in  postprocessing.jl

* Update src/graph.jl

Co-authored-by: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>

---------

Co-authored-by: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>

* Try tests on Julia 1.12 (#276)

* Bump SMC to v0.4.23 (#277)

* Patch failing alloc test (#278)

* Add postprocess_with_star_set! and postprocess_with_tree_set! (#279)

* Bump actions/checkout from 5 to 6 (#282)

Bumps [actions/checkout](https://github.com/actions/checkout) from 5 to 6.
- [Release notes](https://github.com/actions/checkout/releases)
- [Changelog](https://github.com/actions/checkout/blob/main/CHANGELOG.md)
- [Commits](https://github.com/actions/checkout/compare/v5...v6)

---
updated-dependencies:
- dependency-name: actions/checkout
  dependency-version: '6'
  dependency-type: direct:production
  update-type: version-update:semver-major
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>

* Replace is_solved_and_feasible with assert_is_solved_and_feasible (#284)

* Fix JET tests (#285)

* Fix JET tests

* Move JuliaFormatter to action

* Fix formatting

* Use diagonal_indices in the general decompress! for acyclic coloring (#287)

* Enhance decompress! for bicoloring (#288)

* Add a compat entry for CUDA.jl in test/Project.toml (#293)

* Add nb_self_loops in AdjacencyGraph (#290)

* Remove unused arguments from internal functions (#295)

* Add test functions substitutable_columns and substitutable_bidirectional (#297)

* Document fields of TreeSet (#296)

* Fix respectful_similar with SparsityPatternCSC (#299)

* Fix respectful_similar with SparsityPatternCSC

* Add test

* Fix format

* Add test

---------

Co-authored-by: Alexis Montoison <35051714+amontoison@users.noreply.github.com>

* Fix coloring with empty matrix as input (#300)

* Fix coloring with empty matrix as input

* Update src/result.jl

Co-authored-by: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>

* Add check

* Fix format

* Change check

* Apply suggestion from @gdalle

---------

Co-authored-by: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>

* test: skip failing MiniZinc tests (#305)

* Fix eltype invalidation for SparsityPatternCSC (#304)

SparsityPatternCSC{Ti} <: AbstractMatrix{Bool} but the custom
Base.eltype method was returning Ti (the index type) instead of Bool.
This caused 24,906 method invalidations when loading the package,
as it invalidated the backedge from Base.eltype(::AbstractArray).

Change the custom method to SparseArrays.indtype instead, since
that's what the type parameter Ti actually represents. The eltype
is now correctly inherited from AbstractMatrix{Bool}.

Co-authored-by: ChrisRackauckas-Claude <accounts@chrisrackauckas.com>
Co-authored-by: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>

* chore: bump version (#306)

* Bump julia-actions/cache from 2 to 3 (#307)

* Update repo owner to JuliaDiff org (#309)

* Update repo owner

* Bump version

* Switch GPU CI to buildkite

* Add badge

* Fix buildkite badge [skip tests] (#310)

* Accept AbstractSparseMatrixCSC for decompression (#298)

* Accept AbstractSparseMatrixCSC for decompression

* Define decompress_csc

* Remove unused import of AbstractSparseMatrixCSC

* Apply suggestion from @blegat

* Change return statement to return nothing

* Apply suggestion from @blegat

* Add to docs

* Apply suggestion from @gdalle

Co-authored-by: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>

* Move to internals

---------

Co-authored-by: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>

* Bump version (#311)

* Bump codecov/codecov-action from 5 to 6 (#312)

Bumps [codecov/codecov-action](https://github.com/codecov/codecov-action) from 5 to 6.
- [Release notes](https://github.com/codecov/codecov-action/releases)
- [Changelog](https://github.com/codecov/codecov-action/blob/main/CHANGELOG.md)
- [Commits](https://github.com/codecov/codecov-action/compare/v5...v6)

---
updated-dependencies:
- dependency-name: codecov/codecov-action
  dependency-version: '6'
  dependency-type: direct:production
  update-type: version-update:semver-major
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>

* chore: bump CUDA compat to v6 (cuSPARSE now a separate package) (#314)

* chore: bump CUDA compat to v6 (cuSPARSE now a separate package)

* More renaming

* Format

* Fix compression?

* Generic compression and result

* Import

* No generic result

* Fix ambiguity

* Bump julia-actions/setup-julia from 2 to 3 (#315)

Bumps [julia-actions/setup-julia](https://github.com/julia-actions/setup-julia) from 2 to 3.
- [Release notes](https://github.com/julia-actions/setup-julia/releases)
- [Commits](https://github.com/julia-actions/setup-julia/compare/v2...v3)

---
updated-dependencies:
- dependency-name: julia-actions/setup-julia
  dependency-version: '3'
  dependency-type: direct:production
  update-type: version-update:semver-major
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>

---------

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: Neven Sajko <4944410+nsajko@users.noreply.github.com>
Co-authored-by: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Co-authored-by: Alexis Montoison <35051714+amontoison@users.noreply.github.com>
Co-authored-by: Benoît Legat <benoit.legat@gmail.com>
Co-authored-by: Chris Rackauckas - Beep Boop Edition <admin@chrisrackauckas.com>
Co-authored-by: ChrisRackauckas-Claude <accounts@chrisrackauckas.com>