viz: split row and column colors in bicoloring visualization (#195)

gdalle · web-flow · commit bb695faf7474 · 2025-02-20T14:51:57.000+01:00
* viz: split row and column colors in bicoloring visualization

* Add border in docs

* Typo

* Min diff
diff --git a/docs/src/vis.md b/docs/src/vis.md
@@ -77,8 +77,9 @@ result_bi = coloring(S, problem_bi, algo_bi)
 A_img, Br_img, Bc_img = show_colors(
     result_bi;
     colorscheme=ColorSchemes.progress,
-    background=RGB(1, 1, 1),  # white
+    background_color=RGB(1, 1, 1),  # white
     scale=10,
+    border=1,
     pad=2
 )
 ```
diff --git a/ext/SparseMatrixColoringsColorsExt.jl b/ext/SparseMatrixColoringsColorsExt.jl
@@ -26,88 +26,153 @@ using SparseMatrixColorings:
     compress
 using Colors: Colorant, RGB, RGBA, distinguishable_colors
 
-const DEFAULT_BACKGROUND = RGBA(0, 0, 0, 0)
-const DEFAULT_SCALE = 1   # update docstring in src/images.jl when changing this default
-const DEFAULT_PAD = 0 # update docstring in src/images.jl when changing this default
-
-# Sample n distinguishable colors, excluding the background color
-default_colorscheme(n, background) = distinguishable_colors(n, background; dropseed=true)
+# update docstring in src/images.jl when changing this default
+const DEFAULT_BACKGROUND_COLOR = RGBA(0, 0, 0, 0)
+const DEFAULT_BORDER_COLOR = RGB(0, 0, 0)
+const DEFAULT_SCALE = 1
+const DEFAULT_BORDER = 0
+const DEFAULT_PAD = 0
 
 ## Top-level function that handles argument errors, eagerly promotes types and allocates output buffer
 
 function SparseMatrixColorings.show_colors(
     res::AbstractColoringResult;
     colorscheme=nothing,
-    background::Colorant=DEFAULT_BACKGROUND, # color used for zero matrix entries and pad
+    background_color::Colorant=DEFAULT_BACKGROUND_COLOR, # color used for zero matrix entries and pad
+    border_color::Colorant=DEFAULT_BORDER_COLOR, # color used for zero matrix entries and pad
     scale::Int=DEFAULT_SCALE, # scale size of matrix entries to `scale × scale` pixels
+    border::Int=DEFAULT_BORDER,  # border around matrix entries
     pad::Int=DEFAULT_PAD, # pad between matrix entries
     warn::Bool=true,
 )
     scale < 1 && throw(ArgumentError("`scale` has to be ≥ 1."))
+    border < 0 && throw(ArgumentError("`border` has to be ≥ 0."))
     pad < 0 && throw(ArgumentError("`pad` has to be ≥ 0."))
 
     if !isnothing(colorscheme)
         if warn && ncolors(res) > length(colorscheme)
             @warn "`show_colors` will reuse colors since the provided `colorscheme` has $(length(colorscheme)) colors and the matrix needs $(ncolors(res)). You can turn off this warning via the keyword argument `warn = false`, or choose a larger `colorscheme` from ColorSchemes.jl."
         end
-        colorscheme, background = promote_colors(colorscheme, background)
+        colorscheme, background_color, border_color = promote_colors(
+            colorscheme, background_color, border_color
+        )
     else
-        colorscheme = default_colorscheme(ncolors(res), convert(RGB, background))
+        # Sample n distinguishable colors, excluding the background and border color
+        colorscheme = distinguishable_colors(
+            ncolors(res),
+            [convert(RGB, background_color), convert(RGB, border_color)];
+            dropseed=true,
+        )
     end
-    outs = allocate_outputs(res, background, scale, pad)
-    return show_colors!(outs..., res, colorscheme, scale, pad)
+    outs = allocate_outputs(res, background_color, border_color, scale, border, pad)
+    return show_colors!(outs..., res, colorscheme, scale, border, pad)
 end
 
-function promote_colors(colorscheme, background)
+function promote_colors(colorscheme, background_color, border_color)
     # eagerly promote colors to same type
-    T = promote_type(eltype(colorscheme), typeof(background))
+    T = promote_type(eltype(colorscheme), typeof(background_color), typeof(border_color))
     colorscheme = convert.(T, colorscheme)
-    background = convert(T, background)
-    return colorscheme, background
+    background_color = convert(T, background_color)
+    border_color = convert(T, border_color)
+    return colorscheme, background_color, border_color
+end
+
+# Given a CartesianIndex I of an entry in the original matrix, 
+# this function returns the corresponding area in the output image as CartesianIndices.
+function matrix_entry_area(I::CartesianIndex, scale, border, pad)
+    stencil = CartesianIndices((1:scale, 1:scale))
+    return CartesianIndex(1, 1) * (border + pad) +
+           (I - CartesianIndex(1, 1)) * (scale + 2border + pad) .+ stencil
+end
+
+function matrix_entry_plus_border_area(I::CartesianIndex, scale, border, pad)
+    stencil = CartesianIndices((1:(scale + 2border), 1:(scale + 2border)))
+    return CartesianIndex(1, 1) * pad +
+           (I - CartesianIndex(1, 1)) * (scale + 2border + pad) .+ stencil
 end
 
 function allocate_outputs(
     res::Union{AbstractColoringResult{s,:column},AbstractColoringResult{s,:row}},
-    background::Colorant,
+    background_color::Colorant,
+    border_color::Colorant,
     scale::Int,
+    border::Int,
     pad::Int,
 ) where {s}
     A = sparsity_pattern(res)
     B = compress(A, res)
     Base.require_one_based_indexing(A)
     Base.require_one_based_indexing(B)
-    hA, wA = size(A) .* (scale + pad) .+ pad
-    hB, wB = size(B) .* (scale + pad) .+ pad
-    A_img = fill(background, hA, wA)
-    B_img = fill(background, hB, wB)
+    hA, wA = size(A) .* (scale + 2border + pad) .+ (pad)
+    hB, wB = size(B) .* (scale + 2border + pad) .+ (pad)
+    A_img = fill(background_color, hA, wA)
+    B_img = fill(background_color, hB, wB)
+    for I in CartesianIndices(A)
+        if !iszero(A[I])
+            area = matrix_entry_area(I, scale, border, pad)
+            barea = matrix_entry_plus_border_area(I, scale, border, pad)
+            A_img[barea] .= border_color
+            A_img[area] .= background_color
+        end
+    end
+    for I in CartesianIndices(B)
+        if !iszero(B[I])
+            area = matrix_entry_area(I, scale, border, pad)
+            barea = matrix_entry_plus_border_area(I, scale, border, pad)
+            B_img[barea] .= border_color
+            B_img[area] .= background_color
+        end
+    end
     return A_img, B_img
 end
 
 function allocate_outputs(
     res::AbstractColoringResult{s,:bidirectional},
-    background::Colorant,
+    background_color::Colorant,
+    border_color::Colorant,
     scale::Int,
+    border::Int,
     pad::Int,
 ) where {s}
     A = sparsity_pattern(res)
     Br, Bc = compress(A, res)
     Base.require_one_based_indexing(A)
     Base.require_one_based_indexing(Br)
     Base.require_one_based_indexing(Bc)
-    hA, wA = size(A) .* (scale + pad) .+ pad
-    hBr, wBr = size(Br) .* (scale + pad) .+ pad
-    hBc, wBc = size(Bc) .* (scale + pad) .+ pad
-    A_img = fill(background, hA, wA)
-    Br_img = fill(background, hBr, wBr)
-    Bc_img = fill(background, hBc, wBc)
-    return A_img, Br_img, Bc_img
-end
-
-# Given a CartesianIndex I of an entry in the original matrix, 
-# this function returns the corresponding area in the output image as CartesianIndices.
-function matrix_entry_area(I::CartesianIndex, scale, pad)
-    stencil = CartesianIndices((1:scale, 1:scale))
-    return CartesianIndex(pad, pad) + (I - CartesianIndex(1, 1)) * (scale + pad) .+ stencil
+    hA, wA = size(A) .* (scale + 2border + pad) .+ (pad)
+    hBr, wBr = size(Br) .* (scale + 2border + pad) .+ (pad)
+    hBc, wBc = size(Bc) .* (scale + 2border + pad) .+ (pad)
+    Ar_img = fill(background_color, hA, wA)
+    Ac_img = fill(background_color, hA, wA)
+    Br_img = fill(background_color, hBr, wBr)
+    Bc_img = fill(background_color, hBc, wBc)
+    for I in CartesianIndices(A)
+        if !iszero(A[I])
+            area = matrix_entry_area(I, scale, border, pad)
+            barea = matrix_entry_plus_border_area(I, scale, border, pad)
+            Ar_img[barea] .= border_color
+            Ac_img[barea] .= border_color
+            Ar_img[area] .= background_color
+            Ac_img[area] .= background_color
+        end
+    end
+    for I in CartesianIndices(Br)
+        if !iszero(Br[I])
+            area = matrix_entry_area(I, scale, border, pad)
+            barea = matrix_entry_plus_border_area(I, scale, border, pad)
+            Br_img[barea] .= border_color
+            Br_img[area] .= background_color
+        end
+    end
+    for I in CartesianIndices(Bc)
+        if !iszero(Bc[I])
+            area = matrix_entry_area(I, scale, border, pad)
+            barea = matrix_entry_plus_border_area(I, scale, border, pad)
+            Bc_img[barea] .= border_color
+            Bc_img[area] .= background_color
+        end
+    end
+    return Ar_img, Ac_img, Br_img, Bc_img
 end
 
 ## Implementations for different AbstractColoringResult types start here
@@ -117,8 +182,9 @@ function show_colors!(
     B_img::AbstractMatrix{<:Colorant},
     res::AbstractColoringResult{s,:column},
     colorscheme,
-    scale,
-    pad,
+    scale::Int,
+    border::Int,
+    pad::Int,
 ) where {s}
     # cycle color indices if necessary
     A_color_indices = mod1.(column_colors(res), length(colorscheme))
@@ -130,7 +196,7 @@ function show_colors!(
     for I in CartesianIndices(A)
         if !iszero(A[I])
             r, c = Tuple(I)
-            area = matrix_entry_area(I, scale, pad)
+            area = matrix_entry_area(I, scale, border, pad)
             if column_colors(res)[c] > 0
                 A_img[area] .= A_colors[c]
             end
@@ -139,7 +205,7 @@ function show_colors!(
     for I in CartesianIndices(B)
         if !iszero(B[I])
             r, c = Tuple(I)
-            area = matrix_entry_area(I, scale, pad)
+            area = matrix_entry_area(I, scale, border, pad)
             B_img[area] .= B_colors[c]
         end
     end
@@ -151,8 +217,9 @@ function show_colors!(
     B_img::AbstractMatrix{<:Colorant},
     res::AbstractColoringResult{s,:row},
     colorscheme,
-    scale,
-    pad,
+    scale::Int,
+    border::Int,
+    pad::Int,
 ) where {s}
     # cycle color indices if necessary
     A_color_indices = mod1.(row_colors(res), length(colorscheme))
@@ -164,7 +231,7 @@ function show_colors!(
     for I in CartesianIndices(A)
         if !iszero(A[I])
             r, c = Tuple(I)
-            area = matrix_entry_area(I, scale, pad)
+            area = matrix_entry_area(I, scale, border, pad)
             if row_colors(res)[r] > 0
                 A_img[area] .= A_colors[r]
             end
@@ -173,21 +240,23 @@ function show_colors!(
     for I in CartesianIndices(B)
         if !iszero(B[I])
             r, c = Tuple(I)
-            area = matrix_entry_area(I, scale, pad)
+            area = matrix_entry_area(I, scale, border, pad)
             B_img[area] .= B_colors[r]
         end
     end
     return A_img, B_img
 end
 
 function show_colors!(
-    A_img::AbstractMatrix{<:Colorant},
+    Ar_img::AbstractMatrix{<:Colorant},
+    Ac_img::AbstractMatrix{<:Colorant},
     Br_img::AbstractMatrix{<:Colorant},
     Bc_img::AbstractMatrix{<:Colorant},
     res::AbstractColoringResult{s,:bidirectional},
     colorscheme,
-    scale,
-    pad,
+    scale::Int,
+    border::Int,
+    pad::Int,
 ) where {s}
     scale < 3 && throw(ArgumentError("`scale` has to be ≥ 3 to visualize bicoloring"))
     # cycle color indices if necessary
@@ -206,35 +275,30 @@ function show_colors!(
     for I in CartesianIndices(A)
         if !iszero(A[I])
             r, c = Tuple(I)
-            area = matrix_entry_area(I, scale, pad)
-            for i in axes(area, 1), j in axes(area, 2)
-                if j > i
-                    if column_colors(res)[c] > 0
-                        A_img[area[i, j]] = A_ccolors[c]
-                    end
-                elseif i > j
-                    if row_colors(res)[r] > 0
-                        A_img[area[i, j]] = A_rcolors[r]
-                    end
-                end
+            area = matrix_entry_area(I, scale, border, pad)
+            if column_colors(res)[c] > 0
+                Ac_img[area] .= A_ccolors[c]
+            end
+            if row_colors(res)[r] > 0
+                Ar_img[area] .= A_rcolors[r]
             end
         end
     end
     for I in CartesianIndices(Br)
         if !iszero(Br[I])
             r, c = Tuple(I)
-            area = matrix_entry_area(I, scale, pad)
+            area = matrix_entry_area(I, scale, border, pad)
             Br_img[area] .= B_rcolors[r]
         end
     end
     for I in CartesianIndices(Bc)
         if !iszero(Bc[I])
             r, c = Tuple(I)
-            area = matrix_entry_area(I, scale, pad)
+            area = matrix_entry_area(I, scale, border, pad)
             Bc_img[area] .= B_ccolors[c]
         end
     end
-    return A_img, Br_img, Bc_img
+    return Ar_img, Ac_img, Br_img, Bc_img
 end
 
 end # module
diff --git a/src/show_colors.jl b/src/show_colors.jl
@@ -5,19 +5,21 @@
 
 Create a visualization for an [`AbstractColoringResult`](@ref), with the help of the [JuliaImages](https://juliaimages.org) ecosystem.
 
-- For `:column` or `:row` colorings, it returns a tuple `(A_img, B_img)`.
-- For `:bidirectional` colorings, it returns a tuple `(A_img, Br_img, Bc_img)`.
+- For `:column` or `:row` colorings, it returns a couple `(A_img, B_img)`.
+- For `:bidirectional` colorings, it returns a 4-tuple `(Ar_img, Ac_img, Br_img, Bc_img)`.
 
 !!! warning
     This function is implemented in a package extension, using it requires loading [Colors.jl](https://github.com/JuliaGraphics/Colors.jl).
 
 # Keyword arguments
 
 - `colorscheme`: colors used for non-zero matrix entries. This can be a vector of `Colorant`s or a subsampled scheme from [ColorSchemes.jl](https://github.com/JuliaGraphics/ColorSchemes.jl).
-- `background::Colorant`: color used for zero matrix entries and pad. Defaults to `RGBA(0,0,0,0)`, a transparent background.
-- `scale::Int`: scale the size of matrix entries to `scale × scale` pixels. Defaults to `1`. 
+- `background_color::Colorant`: color used for zero matrix entries and pad. Defaults to `RGBA(0,0,0,0)`, a transparent background.
+- `border_color::Colorant`: color used around matrix entries. Defaults to `RGB(0, 0, 0)`, a black border.
+- `scale::Int`: scale the size of matrix entries to `scale × scale` pixels. Defaults to `1`.
+- `border::Int`: set border width around matrix entries, in pixles. Defaults to `0`. 
 - `pad::Int`: set padding between matrix entries, in pixels. Defaults to `0`. 
 
-For a matrix of size `(m, n)`, the resulting output will be of size `(m * (scale + pad) + pad, n * (scale + pad) + pad)`.
+For a matrix of size `(m, n)`, the resulting output will be of size `(m * (scale + 2border + pad) + pad, n * (scale + 2border + pad) + pad)`.
 """
 function show_colors end
diff --git a/test/show_colors.jl b/test/show_colors.jl
@@ -32,9 +32,10 @@ algo = GreedyColoringAlgorithm(; decompression=:direct)
         @test A_img isa Matrix{<:Colorant}
 
         pad = 2
-        A_img, B_img = show_colors(result; scale=scale, pad=pad)
-        @test size(A_img) == size(A) .* (scale + pad) .+ pad
-        @test size(B_img) == size(B) .* (scale + pad) .+ pad
+        border = 3
+        A_img, B_img = show_colors(result; scale=scale, border=border, pad=pad)
+        @test size(A_img) == size(A) .* (scale + 2border + pad) .+ pad
+        @test size(B_img) == size(B) .* (scale + 2border + pad) .+ pad
         @test A_img isa Matrix{<:Colorant}
 
         @testset "color cycling" begin
@@ -51,11 +52,13 @@ algo = GreedyColoringAlgorithm(; decompression=:direct)
         Br, Bc = compress(A, result)
 
         scale = 3
-        A_img, Br_img, Bc_img = show_colors(result; scale=scale)
-        @test size(A_img) == size(A) .* scale
+        Ar_img, Ac_img, Br_img, Bc_img = show_colors(result; scale=scale)
+        @test size(Ar_img) == size(A) .* scale
+        @test size(Ac_img) == size(A) .* scale
         @test size(Br_img) == size(Br) .* scale
         @test size(Bc_img) == size(Bc) .* scale
-        @test A_img isa Matrix{<:Colorant}
+        @test Ar_img isa Matrix{<:Colorant}
+        @test Ac_img isa Matrix{<:Colorant}
         @test Br_img isa Matrix{<:Colorant}
         @test Bc_img isa Matrix{<:Colorant}
     end
@@ -72,4 +75,8 @@ end
         @test_throws ArgumentError show_colors(result; pad=-1)
         @test_nowarn show_colors(result; pad=0)
     end
+    @testset "border too small" begin
+        @test_throws ArgumentError show_colors(result; border=-1)
+        @test_nowarn show_colors(result; border=0)
+    end
 end
