From e42fdedc5c61418c86841f1003d181df072c0d3d Mon Sep 17 00:00:00 2001
From: Fredrik Ekre <ekrefredrik@gmail.com>
Date: Tue, 12 Jul 2022 01:40:28 +0200
Subject: [PATCH] Add support for edge weights from sparse matrix input.

This patch adds the weights keyword argument to Metis.graph to construct
a Metis.Graph with edge weights based on the matrix entries, i.e.:

    g = Metis.graph(A; weights=true)

Note that the weights must be integers.
---
 src/Metis.jl     | 18 ++++++++++++++----
 test/runtests.jl | 26 ++++++++++++++++++++++----
 2 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/src/Metis.jl b/src/Metis.jl
index d2f980c..8cb7b62 100644
--- a/src/Metis.jl
+++ b/src/Metis.jl
@@ -22,17 +22,21 @@ struct Graph
     xadj::Vector{idx_t}
     adjncy::Union{Vector{idx_t}, Ptr{Nothing}}
     vwgt::Union{Vector{idx_t}, Ptr{Nothing}}
-    Graph(nvtxs, xadj, adjncy, vwgt=C_NULL) = new(nvtxs, xadj, adjncy, vwgt)
+    adjwgt::Union{Vector{idx_t}, Ptr{Nothing}}
+    function Graph(nvtxs, xadj, adjncy, vwgt=C_NULL, adjwgt=C_NULL)
+        return new(nvtxs, xadj, adjncy, vwgt, adjwgt)
+    end
 end
 
 """
-    Metis.graph(G::SparseMatrixCSC; check_hermitian=true)
+    Metis.graph(G::SparseMatrixCSC; weights=false, check_hermitian=true)
 
 Construct the 1-based CSR representation of the sparse matrix `G`.
 If `check_hermitian` is `false` the matrix is not checked for being hermitian
 before constructing the graph.
+If `weights=true` the entries of the matrix are used as edge weights.
 """
-function graph(G::SparseMatrixCSC; check_hermitian=true)
+function graph(G::SparseMatrixCSC; weights::Bool=false, check_hermitian::Bool=true)
     if check_hermitian
         ishermitian(G) || throw(ArgumentError("matrix must be Hermitian"))
     end
@@ -40,6 +44,8 @@ function graph(G::SparseMatrixCSC; check_hermitian=true)
     xadj = Vector{idx_t}(undef, N+1)
     xadj[1] = 1
     adjncy = Vector{idx_t}(undef, nnz(G))
+    vwgt = C_NULL # TODO: Vertex weights could be passed as input argument
+    adjwgt = weights ? Vector{idx_t}(undef, nnz(G)) : C_NULL
     adjncy_i = 0
     @inbounds for j in 1:N
         n_rows = 0
@@ -49,12 +55,16 @@ function graph(G::SparseMatrixCSC; check_hermitian=true)
                 n_rows += 1
                 adjncy_i += 1
                 adjncy[adjncy_i] = i
+                if weights
+                    adjwgt[adjncy_i] = G.nzval[k]
+                end
             end
         end
         xadj[j+1] = xadj[j] + n_rows
     end
     resize!(adjncy, adjncy_i)
-    return Graph(idx_t(N), xadj, adjncy)
+    weights && resize!(adjwgt, adjncy_i)
+    return Graph(idx_t(N), xadj, adjncy, vwgt, adjwgt)
 end
 
 """
diff --git a/test/runtests.jl b/test/runtests.jl
index ae33753..4f2fc8a 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -4,9 +4,25 @@ using Test
 using SparseArrays
 import LightGraphs, Graphs
 
+@testset "Metis.graph(::SparseMatrixCSC)" begin
+    rng = MersenneTwister(0)
+    S = sprand(rng, Int, 10, 10, 0.2); S = S + S'; fill!(S.nzval, 1)
+    foreach(i -> S[i, i] = 0, 1:10); dropzeros!(S)
+    g = Metis.graph(S)
+    gw = Metis.graph(S; weights=true)
+    @test g.xadj == gw.xadj
+    @test g.adjncy == gw.adjncy
+    @test g.adjwgt == C_NULL
+    @test gw.adjwgt == ones(Int, length(gw.adjncy))
+    G  = SparseMatrixCSC(10, 10, g.xadj, g.adjncy, ones(Int, length(g.adjncy)))
+    GW = SparseMatrixCSC(10, 10, gw.xadj, gw.adjncy, gw.adjwgt)
+    @test iszero(S - G)
+    @test iszero(S - GW)
+end
+
 @testset "Metis.permutation" begin
     rng = MersenneTwister(0)
-    S = sprand(rng, 10, 10, 0.5); S = S + S'; fill(S.nzval, 1)
+    S = sprand(rng, 10, 10, 0.5); S = S + S'; fill!(S.nzval, 1)
     perm, iperm = Metis.permutation(S)
     @test isperm(perm)
     @test isperm(iperm)
@@ -15,10 +31,11 @@ end
 
 @testset "Metis.partition" begin
     rng = MersenneTwister(0)
-    S = sprand(rng, 10, 10, 0.5); S = S + S'; fill(S.nzval, 1)
+    S = sprand(rng, Int, 10, 10, 0.5); S = S + S'; fill!(S.nzval, 1)
+    SG = Metis.graph(S; weights=true)
     T = Graphs.smallgraph(:tutte)
     TL = LightGraphs.smallgraph(:tutte)
-    for G in (S, T, TL), alg in (:RECURSIVE, :KWAY), nparts in (3, 4)
+    for G in (S, T, TL, SG), alg in (:RECURSIVE, :KWAY), nparts in (3, 4)
         partition = Metis.partition(G, nparts, alg = alg)
         @test extrema(partition) == (1, nparts)
         @test all(x -> findfirst(==(x), partition) !== nothing, 1:nparts)
@@ -27,7 +44,8 @@ end
 
 @testset "Metis.separator" begin
     rng = MersenneTwister(0)
-    S = sprand(rng, 10, 10, 0.5); S = S + S'; fill(S.nzval, 1)
+    S = sprand(rng, 10, 10, 0.5); S = S + S'; fill!(S.nzval, 1)
+    SG = Metis.graph(S; weights=true)
     T = Graphs.smallgraph(:tutte)
     TL = LightGraphs.smallgraph(:tutte)
     for G in (S, T, TL)