update sa with new greedy and sub_optimizer

src/sa.jl

@@ -20,14 +20,14 @@ Then finds the contraction order inside each group with the greedy search algori
 ### References
 * [Hyper-optimized tensor network contraction](https://arxiv.org/abs/2002.01935)
 """
-Base.@kwdef struct SABipartite{RT,BT,GM} <: CodeOptimizer
+Base.@kwdef struct SABipartite{RT,BT,SO} <: CodeOptimizer
     sc_target::RT = 25
     ntrials::Int = 50  # number of trials
     βs::BT = 0.1:0.2:15.0  # temperatures
     niters::Int = 1000  # number of iterations in each temperature
     max_group_size::Int = 40
     # configure greedy algorithm
-    greedy_config::GM = GreedyMethod()
+    sub_optimizer::SO = GreedyMethod()
     initializer::Symbol = :random
 end
 
@@ -181,7 +181,7 @@ function initialize_greedy(adj, vertices, log2_sizes)
     incidence_list = IncidenceList(v2e; openedges=openedges)
     log2_edge_sizes = Dict([i=>log2_sizes[i] for i=1:length(log2_sizes)])
     # nrepeat=3 because there are overheads
-    tree, _, _ = tree_greedy(incidence_list, log2_edge_sizes; method=MinSpaceOut(), nrepeat=3)
+    tree, _, _ = tree_greedy(incidence_list, log2_edge_sizes; nrepeat=3)
 
     # build configuration from the tree
     res = ones(Int, size(adj, 1))
@@ -213,10 +213,10 @@ Check the docstring of `SABipartite` for detailed explaination of other input ar
 * [Hyper-optimized tensor network contraction](https://arxiv.org/abs/2002.01935)
 """
 function optimize_sa(code::EinCode, size_dict; sc_target, max_group_size=40,
-             βs=0.01:0.02:15.0, niters=1000, ntrials=50, greedy_method=MinSpaceOut(), greedy_nrepeat=10,
+             βs=0.01:0.02:15.0, niters=1000, ntrials=50, sub_optimizer=GreedyMethod(),
              initializer=:random)
     bipartiter = SABipartite(; sc_target=sc_target, βs=βs, niters=niters, ntrials=ntrials,
-        greedy_config=GreedyMethod(method=greedy_method, nrepeat=greedy_nrepeat),
+    sub_optimizer=sub_optimizer,
         max_group_size=max_group_size, initializer=initializer)
     recursive_bipartite_optimize(bipartiter, code, size_dict)
 end

src/simplify.jl

@@ -36,7 +36,7 @@ function merge_greedy(code::EinCode{LT}, size_dict; threshhold=-1e-12) where LT
         return collect(vertices(incidence_list))[1]
     end
     tree = Dict{Int,NestedEinsum}([v=>NestedEinsum{LT}(v) for v in vertices(incidence_list)])
-    cost_values = evaluate_costs(MinSpaceDiff(), incidence_list, log2_edge_sizes)
+    cost_values = evaluate_costs(1.0, incidence_list, log2_edge_sizes)
     while true
         if length(cost_values) == 0
             return _buildsimplifier(tree, incidence_list)
@@ -48,7 +48,7 @@ function merge_greedy(code::EinCode{LT}, size_dict; threshhold=-1e-12) where LT
             if nv(incidence_list) <= 1
                 return _buildsimplifier(tree, incidence_list)
             end
-            update_costs!(cost_values, pair..., MinSpaceDiff(), incidence_list, log2_edge_sizes)
+            update_costs!(cost_values, pair..., 1.0, incidence_list, log2_edge_sizes)
         else
             return _buildsimplifier(tree, incidence_list)
         end

test/kahypar.jl

@@ -81,7 +81,7 @@ end
     Random.seed!(2)
     code = random_regular_eincode(220, 3)
     codeg_auto = optimize_kahypar_auto(code, uniformsize(code, 2), sub_optimizer=GreedyMethod())
-    codet_auto = optimize_kahypar_auto(code, uniformsize(code, 2), sub_optimizer=TreeSA(ntrials = 4, sc_weight = 0.1))
+    codet_auto = optimize_kahypar_auto(code, uniformsize(code, 2), sub_optimizer=TreeSA(ntrials = 1, sc_weight = 0.1))
     ccg = contraction_complexity(codeg_auto, uniformsize(code, 2))
     @show ccg.sc, ccg.tc
     cct = contraction_complexity(codet_auto, uniformsize(code, 2))

test/runtests.jl

@@ -5,26 +5,26 @@ using Test
     include("greedy.jl")
 end
 
+@testset "sa" begin
+    include("sa.jl")
+end
+
 @testset "kahypar" begin
     include("kahypar.jl")
 end
 
-# @testset "sa" begin
-#     include("sa.jl")
-# end
-
 @testset "treesa" begin
     include("treesa.jl")
 end
 
-# @testset "simplify" begin
-#     include("simplify.jl")
-# end
+@testset "simplify" begin
+    include("simplify.jl")
+end
 
-# @testset "interfaces" begin
-#     include("interfaces.jl")
-# end
+@testset "interfaces" begin
+    include("interfaces.jl")
+end
 
-# @testset "json" begin
-#     include("json.jl")
-# end
+@testset "json" begin
+    include("json.jl")
+end