Add support for edge weights using SimpleWeightedGraphs (#51)

This patch adds support for edge weights based on the
SimpleWeightedGraphs package. This is implemented as a package
extension, and thus requires Julia 1.9 or newer. Closes #37.

Co-authored-by: Rodolfo Carvajal <rodolfo.carvajal@gmail.com>
Co-authored-by: Fredrik Ekre <ekrefredrik@gmail.com>

Project.toml

@@ -13,23 +13,27 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 [weakdeps]
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 LightGraphs = "093fc24a-ae57-5d10-9952-331d41423f4d"
+SimpleWeightedGraphs = "47aef6b3-ad0c-573a-a1e2-d07658019622"
 
 [extensions]
 MetisGraphs = "Graphs"
 MetisLightGraphs = "LightGraphs"
+MetisSimpleWeightedGraphs = ["SimpleWeightedGraphs", "Graphs"]
 
 [compat]
 CEnum = "0.4"
 Graphs = "1"
 LightGraphs = "1"
 METIS_jll = "5.1"
+SimpleWeightedGraphs = "1"
 julia = "1.6"
 
 [extras]
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 LightGraphs = "093fc24a-ae57-5d10-9952-331d41423f4d"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+SimpleWeightedGraphs = "47aef6b3-ad0c-573a-a1e2-d07658019622"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Graphs", "LightGraphs", "Random", "Test"]
+test = ["Graphs", "LightGraphs", "Random", "SimpleWeightedGraphs", "Test"]

ext/GraphsCommon.jl

@@ -6,11 +6,13 @@
 using Metis.LibMetis: idx_t
 using Metis: Graph
 
-function graph(G)
+function graph(G; weights::Bool=false, kwargs...)
     N = nv(G)
     xadj = Vector{idx_t}(undef, N+1)
     xadj[1] = 1
     adjncy = Vector{idx_t}(undef, 2*ne(G))
+    vwgt = C_NULL # TODO: Vertex weights could be passed as an input argument
+    adjwgt = weights ? Vector{idx_t}(undef, 2*ne(G)) : C_NULL
     adjncy_i = 0
     for j in 1:N
         ne = 0
@@ -19,9 +21,17 @@ function graph(G)
             ne += 1
             adjncy_i += 1
             adjncy[adjncy_i] = i
+            if weights
+                weight = get_weight(G, i, j)
+                if !(isinteger(weight) && weight > 0)
+                    error("weights must be positive integers, got weight $(weight) for edge ($(i), $(j))")
+                end
+                adjwgt[adjncy_i] = weight
+            end
         end
         xadj[j+1] = xadj[j] + ne
     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

ext/MetisSimpleWeightedGraphs.jl

@@ -0,0 +1,20 @@
+# SPDX-License-Identifier: MIT
+
+module MetisSimpleWeightedGraphs
+
+using Graphs: ne, nv, outneighbors
+using Metis: Metis
+using SimpleWeightedGraphs: AbstractSimpleWeightedGraph, get_weight
+
+"""
+    graph(G::SimpleWeightedGraphs.AbstractSimpleGraph; weights=true) :: Metis.Graph
+
+Construct the 1-based CSR representation of the weighted graph `G`.
+"""
+function Metis.graph(G::AbstractSimpleWeightedGraph; weights::Bool=true, kwargs...)
+    return graph(G; weights=weights, kwargs...)
+end
+
+include("GraphsCommon.jl")
+
+end # module MetisSimpleWeightedGraphs

src/Metis.jl

@@ -102,10 +102,10 @@ function partition(G::Graph, nparts::Integer; alg = :KWAY)
     nparts == 1 && return fill!(part, 1) # https://github.com/JuliaSparse/Metis.jl/issues/49
     edgecut = fill(idx_t(0), 1)
     if alg === :RECURSIVE
-        @check METIS_PartGraphRecursive(Ref{idx_t}(G.nvtxs), Ref{idx_t}(1), G.xadj, G.adjncy, G.vwgt, C_NULL, C_NULL,
+        @check METIS_PartGraphRecursive(Ref{idx_t}(G.nvtxs), Ref{idx_t}(1), G.xadj, G.adjncy, G.vwgt, C_NULL, G.adjwgt,
                                         Ref{idx_t}(nparts), C_NULL, C_NULL, options, edgecut, part)
     elseif alg === :KWAY
-        @check METIS_PartGraphKway(Ref{idx_t}(G.nvtxs), Ref{idx_t}(1), G.xadj, G.adjncy, G.vwgt, C_NULL, C_NULL,
+        @check METIS_PartGraphKway(Ref{idx_t}(G.nvtxs), Ref{idx_t}(1), G.xadj, G.adjncy, G.vwgt, C_NULL, G.adjwgt,
                                    Ref{idx_t}(nparts), C_NULL, C_NULL, options, edgecut, part)
     else
         throw(ArgumentError("unknown algorithm $(repr(alg))"))

test/runtests.jl

@@ -40,6 +40,26 @@ end
         @test extrema(partition) == (1, nparts)
         @test all(x -> findfirst(==(x), partition) !== nothing, 1:nparts)
     end
+    # With weights
+    if isdefined(Base, :get_extension)
+        import SimpleWeightedGraphs
+        G1 = SimpleWeightedGraphs.SimpleWeightedGraph(Graphs.nv(T))
+        G2 = SimpleWeightedGraphs.SimpleWeightedGraph(Graphs.nv(T))
+        for edge in Graphs.edges(T)
+            i, j = Tuple(edge)
+            SimpleWeightedGraphs.add_edge!(G1, i, j, 1)
+            SimpleWeightedGraphs.add_edge!(G2, i, j, i+j)
+        end
+        for alg in (:RECURSIVE, :KWAY), nparts in (3, 4)
+            unwpartition = Metis.partition(T, nparts, alg = alg)
+            partition = Metis.partition(G1, nparts, alg = alg)
+            @test partition == unwpartition
+            partition = Metis.partition(G2, nparts, alg = alg)
+            @test partition != unwpartition
+            @test extrema(partition) == (1, nparts)
+            @test all(x -> findfirst(==(x), partition) !== nothing, 1:nparts)
+        end
+    end
 end
 
 @testset "Metis.separator" begin