Move KaHyPar optimizer to package extension

Author: Sergio Sánchez Ramírez

Project.toml

@@ -9,7 +9,6 @@ Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
-KaHyPar = "2a6221f6-aa48-11e9-3542-2d9e0ef01880"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 PackageExtensionCompat = "65ce6f38-6b18-4e1d-a461-8949797d7930"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
@@ -19,18 +18,21 @@ Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 GraphMakie = "1ecd5474-83a3-4783-bb4f-06765db800d2"
+KaHyPar = "2a6221f6-aa48-11e9-3542-2d9e0ef01880"
 Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
 
 [extensions]
 EinExprsChainRulesCoreExt = "ChainRulesCore"
 EinExprsFiniteDifferencesExt = "FiniteDifferences"
 EinExprsGraphMakieExt = ["Makie", "GraphMakie"]
+EinExprsKaHyParExt = "KaHyPar"
 EinExprsMakieExt = "Makie"
 
 [extras]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 GraphMakie = "1ecd5474-83a3-4783-bb4f-06765db800d2"
+KaHyPar = "2a6221f6-aa48-11e9-3542-2d9e0ef01880"
 Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
 
 [compat]

ext/EinExprsKaHyParExt.jl

@@ -0,0 +1,46 @@
+module EinExprsKaHyParExt
+
+using EinExprs
+using AbstractTrees
+using SparseArrays
+using KaHyPar
+using Suppressor
+using Compat
+
+function EinExprs.einexpr(config::EinExprs.HyPar, path)
+    config.stop(path) && return path
+
+    inds = mapreduce(head, ∪, path.args)
+    indexmap = Dict(Iterators.map(@compat(splat(Pair)) ∘ reverse, enumerate(inds)))
+
+    I = EinExprs.flatmap(((i, tensor),) -> fill(i, ndims(tensor)), enumerate(path.args)) |> collect
+    J = EinExprs.flatmap(tensor -> Iterators.map(Base.Fix1(getindex, indexmap), head(tensor)), path.args) |> collect
+    V = fill(1, length(I))
+    incidence_matrix = sparse(I, J, V)
+
+    # NOTE indices in `inds` should be in the same order as unique indices appear by iterating on `path.args` because `∪` retains order
+    edge_weights = map(config.edge_scaler ∘ Base.Fix1(size, path), inds)
+    vertex_weights = map(config.vertex_scaler ∘ length, args(path))
+
+    hypergraph = KaHyPar.HyperGraph(incidence_matrix, vertex_weights, edge_weights)
+    KaHyPar.kahypar_set_seed(hypergraph.context, config.seed)
+
+    partitions = @suppress KaHyPar.partition(
+        hypergraph,
+        config.parts;
+        imbalance = config.imbalance,
+        configuration = config.configuration,
+    )
+
+    _args = map(unique(partitions)) do partition
+        selection = partitions .== partition
+        count(selection) == 1 && return only(args(path)[selection])
+
+        expr = sum(args(path)[selection], skip = path.head)
+        einexpr(config, expr)
+    end
+
+    return sum(_args, skip = path.head)
+end
+
+end

src/Optimizers/KaHyPar.jl

@@ -1,8 +1,4 @@
-using AbstractTrees
-using SparseArrays
-using KaHyPar
-using Suppressor
-using Compat
+using Base: @kwdef
 
 @kwdef struct HyPar <: Optimizer
     parts::Int = 2
@@ -13,39 +9,3 @@ using Compat
     vertex_scaler::Function = Base.Fix1(*, 1000) ∘ Int ∘ round ∘ log2
     seed::Int = 0
 end
-
-function EinExprs.einexpr(config::HyPar, path)
-    config.stop(path) && return path
-
-    inds = mapreduce(head, ∪, path.args)
-    indexmap = Dict(Iterators.map(@compat(splat(Pair)) ∘ reverse, enumerate(inds)))
-
-    I = flatmap(((i, tensor),) -> fill(i, ndims(tensor)), enumerate(path.args)) |> collect
-    J = flatmap(tensor -> Iterators.map(Base.Fix1(getindex, indexmap), head(tensor)), path.args) |> collect
-    V = fill(1, length(I))
-    incidence_matrix = sparse(I, J, V)
-
-    # NOTE indices in `inds` should be in the same order as unique indices appear by iterating on `path.args` because `∪` retains order
-    edge_weights = map(config.edge_scaler ∘ Base.Fix1(size, path), inds)
-    vertex_weights = map(config.vertex_scaler ∘ length, args(path))
-
-    hypergraph = KaHyPar.HyperGraph(incidence_matrix, vertex_weights, edge_weights)
-    KaHyPar.kahypar_set_seed(hypergraph.context, config.seed)
-
-    partitions = @suppress KaHyPar.partition(
-        hypergraph,
-        config.parts;
-        imbalance = config.imbalance,
-        configuration = config.configuration,
-    )
-
-    _args = map(unique(partitions)) do partition
-        selection = partitions .== partition
-        count(selection) == 1 && return only(args(path)[selection])
-
-        expr = sum(args(path)[selection], skip = path.head)
-        einexpr(config, expr)
-    end
-
-    return sum(_args, skip = path.head)
-end

test/KaHyPar_test.jl

@@ -1,4 +1,6 @@
 @testset "KaHyPar" begin
+    using KaHyPar
+
     @testset begin
         tensors = [
             EinExpr([:j, :b, :i, :h], Dict(i => 2 for i in [:j, :b, :i, :h])),

test/Project.toml

@@ -2,6 +2,7 @@
 CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 GraphMakie = "1ecd5474-83a3-4783-bb4f-06765db800d2"
+KaHyPar = "2a6221f6-aa48-11e9-3542-2d9e0ef01880"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 NetworkLayout = "46757867-2c16-5918-afeb-47bfcb05e46a"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"