From 4e733b048fe80118ee24eaef5392fb4436edb521 Mon Sep 17 00:00:00 2001
From: Michel Schanen <mschanen@anl.gov>
Date: Mon, 21 Feb 2022 11:16:41 -0600
Subject: [PATCH] Column and row partitioning helper

---
 Project.toml     |  5 +++--
 src/ColPack.jl   |  8 ++++---
 src/utils.jl     | 55 ++++++++++++++++++++++++++++++++++++++++++++++++
 test/runtests.jl | 15 +++++++++++++
 4 files changed, 78 insertions(+), 5 deletions(-)
 create mode 100644 src/utils.jl

diff --git a/Project.toml b/Project.toml
index be48f96..fc56fd0 100644
--- a/Project.toml
+++ b/Project.toml
@@ -11,12 +11,13 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [compat]
-ColPack_jll = "0.2"
+ColPack_jll = "0.3.0"
 MatrixMarket = "0.3"
 julia = "1.6"
 
 [extras]
+SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test"]
+test = ["SparseDiffTools", "Test"]
diff --git a/src/ColPack.jl b/src/ColPack.jl
index 9125f97..050bba2 100644
--- a/src/ColPack.jl
+++ b/src/ColPack.jl
@@ -5,7 +5,7 @@ using LinearAlgebra
 using MatrixMarket
 using SparseArrays
 
-export ColPackColoring, get_colors
+export ColPackColoring, get_colors, matrix2adjmatrix
 
 abstract type AbstractColoring end
 include("colorings.jl")
@@ -15,6 +15,8 @@ abstract type AbstractOrdering end
 include("orderings.jl")
 export AbstractOrdering
 
+include("utils.jl")
+
 mutable struct ColPackColoring
     refColPack::Vector{Ptr{Cvoid}}
     coloring::Vector{Cint}
@@ -45,7 +47,7 @@ function ColPackColoring(filename::AbstractString, method::AbstractColoring, ord
     end
 
     g = ColPackColoring(refColPack, zeros(Int, reflen[1]), method, order, nothing)
-    # finalizer(free_coloring, g)
+    finalizer(free_coloring, g)
     return g
 end
 
@@ -74,7 +76,7 @@ function ColPackColoring(M::SparseMatrixCSC{VT,IT}, method::AbstractColoring, or
     end
 
     g = ColPackColoring(refColPack, zeros(Int, reflen[1]), method, order, csr)
-    # finalizer(free_coloring, g)
+    finalizer(free_coloring, g)
     return g
 end
 
diff --git a/src/utils.jl b/src/utils.jl
new file mode 100644
index 0000000..d9e841d
--- /dev/null
+++ b/src/utils.jl
@@ -0,0 +1,55 @@
+"""
+        _cols_by_rows(rows_index,cols_index)
+Returns a vector of rows where each row contains
+a vector of its column indices.
+"""
+function _cols_by_rows(rows_index,cols_index)
+    nrows = maximum(rows_index)
+    cols_by_rows = [eltype(rows_index)[] for _ in 1:nrows]
+    for (i,j) in zip(rows_index,cols_index)
+        push!(cols_by_rows[i],j)
+    end
+    return cols_by_rows
+end
+
+
+"""
+        _rows_by_cols(rows_index,cols_index)
+Returns a vector of columns where each column contains
+a vector of its row indices.
+"""
+function _rows_by_cols(rows_index,cols_index)
+    return _cols_by_rows(cols_index,rows_index)
+end
+
+function matrix2adjmatrix(M; partition_by_rows = true)
+    (rows_index, cols_index, _) = findnz(M)
+    if partition_by_rows
+        A = SparseMatrixCSC{Float64, Cuint}(size(M,1), size(M,1), ones(Cuint, size(M,1)+1), Vector{Cuint}(), Vector{Float64}())
+        rows_by_cols = _rows_by_cols(rows_index,cols_index)
+        @inbounds for (cur_row,cur_col) in zip(rows_index,cols_index)
+            if !isempty(rows_by_cols[cur_col])
+                for next_row in rows_by_cols[cur_col]
+                    if next_row < cur_row 
+                        A[cur_row, next_row] = 1.0
+                        A[next_row, cur_row] = 1.0
+                    end
+                end
+            end
+        end
+    else
+        A = SparseMatrixCSC{Float64, Cuint}(size(M,2), size(M,2), ones(Cuint, size(M,2)+1), Vector{Cuint}(), Vector{Float64}())
+        cols_by_rows = _cols_by_rows(rows_index,cols_index)
+        @inbounds for (cur_row,cur_col) in zip(rows_index,cols_index)
+            if !isempty(cols_by_rows[cur_row])
+                for next_col in cols_by_rows[cur_row]
+                    if next_col < cur_col 
+                        A[cur_col, next_col] = 1.0
+                        A[next_col, cur_col] = 1.0
+                    end
+                end
+            end
+        end
+    end
+    return A
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index 8a90839..45947cc 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -3,6 +3,7 @@ using LinearAlgebra
 using MatrixMarket
 using Random
 using SparseArrays
+using SparseDiffTools
 using Test
 
 orderings = Vector{AbstractOrdering}()
@@ -50,4 +51,18 @@ verbose = false
             @test length(unique(get_colors(coloring))) == ncolors[i]
         end
     end
+    @testset "Test ColPack Columns Coloring" begin
+        A = [
+        [1.0 1.0 0.0 0.0 0.0];
+        [0.0 0.0 1.0 0.0 0.0];
+        [0.0 1.0 0.0 1.0 0.0];
+        [0.0 0.0 0.0 1.0 1.0];
+        ]
+
+        A = sparse(A)
+        adjA = ColPack.matrix2adjmatrix(A; partition_by_rows=false)
+        @test issymmetric(adjA)
+        coloring = ColPackColoring(adjA, d1_coloring(), natural_ordering(); verbose=true)
+        @test get_colors(coloring) == SparseDiffTools.matrix_colors(A)
+    end
 end