Support slicing (#5)

* support slicing

* support slicing basic

* use warn instead of errors

* bump version

* rm warning test

Author: GiggleLiu

Project.toml

@@ -28,7 +28,7 @@ FFTW = "1.4"
 Graphs = "1.4"
 Mods = "1.3"
 OMEinsum = "0.6.1"
-OMEinsumContractionOrders = "0.5"
+OMEinsumContractionOrders = "0.6"
 Polynomials = "2.0"
 Primes = "0.5"
 Requires = "1"

src/GraphTensorNetworks.jl

@@ -1,6 +1,6 @@
 module GraphTensorNetworks
 
-using OMEinsumContractionOrders: OMEinsum
+using OMEinsumContractionOrders: SlicedEinsum
 using Core: Argument
 using TropicalGEMM, TropicalNumbers
 using OMEinsum

src/arithematics.jl

@@ -130,7 +130,7 @@ end
 
 Base.length(x::ConfigEnumerator{N}) where N = length(x.data)
 Base.getindex(x::ConfigEnumerator, i) = x.data[i]
-Base.:(==)(x::ConfigEnumerator{N,S,C}, y::ConfigEnumerator{N,S,C}) where {N,S,C} = x.data == y.data
+Base.:(==)(x::ConfigEnumerator{N,S,C}, y::ConfigEnumerator{N,S,C}) where {N,S,C} = Set(x.data) == Set(y.data)
 
 function Base.:+(x::ConfigEnumerator{N,S,C}, y::ConfigEnumerator{N,S,C}) where {N,S,C}
     length(x) == 0 && return y

src/bounding.jl

@@ -1,5 +1,7 @@
 using OMEinsum: DynamicEinCode
 
+export AllConfigs, SingleConfig
+
 struct AllConfigs{K} end
 largest_k(::AllConfigs{K}) where K = K
 struct SingleConfig end
@@ -57,6 +59,12 @@ struct CacheTree{T}
     content::AbstractArray{T}
     siblings::Vector{CacheTree{T}}
 end
+function cached_einsum(se::SlicedEinsum, @nospecialize(xs), size_dict)
+    if length(se.slicing) != 0
+        @warn "Slicing is not supported for caching! Fallback to `NestedEinsum`."
+    end
+    return cached_einsum(se.eins, xs, size_dict)
+end
 function cached_einsum(code::NestedEinsum, @nospecialize(xs), size_dict)
     if OMEinsum.isleaf(code)
         y = xs[code.tensorindex]
@@ -69,6 +77,12 @@ function cached_einsum(code::NestedEinsum, @nospecialize(xs), size_dict)
 end
 
 # computed mask tree by back propagation
+function generate_masktree(mode, se::SlicedEinsum, cache, mask, size_dict)
+    if length(se.slicing) != 0
+        @warn "Slicing is not supported for generating masked tree! Fallback to `NestedEinsum`."
+    end
+    return generate_masktree(mode, se.eins, cache, mask, size_dict)
+end
 function generate_masktree(mode, code::NestedEinsum, cache, mask, size_dict)
     if OMEinsum.isleaf(code)
         return CacheTree(mask, CacheTree{Bool}[])
@@ -79,6 +93,12 @@ function generate_masktree(mode, code::NestedEinsum, cache, mask, size_dict)
 end
 
 # The masked einsum contraction
+function masked_einsum(se::SlicedEinsum, @nospecialize(xs), masks, size_dict)
+    if length(se.slicing) != 0
+        @warn "Slicing is not supported for masked contraction! Fallback to `NestedEinsum`."
+    end
+    return masked_einsum(se.eins, xs, masks, size_dict)
+end
 function masked_einsum(code::NestedEinsum, @nospecialize(xs), masks, size_dict)
     if OMEinsum.isleaf(code)
         y = copy(xs[code.tensorindex])
@@ -106,7 +126,7 @@ function bounding_contract(mode::AllConfigs, code::EinCode, @nospecialize(xsa),
     LT = OMEinsum.labeltype(code)
     bounding_contract(mode, NestedEinsum(NestedEinsum{DynamicEinCode{LT}}.(1:length(xsa)), code), xsa, ymask, xsb; size_info=size_info)
 end
-function bounding_contract(mode::AllConfigs, code::NestedEinsum, @nospecialize(xsa), ymask, @nospecialize(xsb); size_info=nothing)
+function bounding_contract(mode::AllConfigs, code::Union{NestedEinsum,SlicedEinsum}, @nospecialize(xsa), ymask, @nospecialize(xsb); size_info=nothing)
     size_dict = size_info===nothing ? Dict{OMEinsum.labeltype(code.eins),Int}() : copy(size_info)
     OMEinsum.get_size_dict!(code, xsa, size_dict)
     # compute intermediate tensors
@@ -125,7 +145,7 @@ function solution_ad(code::EinCode, @nospecialize(xsa), ymask; size_info=nothing
     solution_ad(NestedEinsum(NestedEinsum{DynamicEinCode{LT}}.(1:length(xsa)), code), xsa, ymask; size_info=size_info)
 end
 
-function solution_ad(code::NestedEinsum, @nospecialize(xsa), ymask; size_info=nothing)
+function solution_ad(code::Union{NestedEinsum,SlicedEinsum}, @nospecialize(xsa), ymask; size_info=nothing)
     size_dict = size_info===nothing ? Dict{OMEinsum.labeltype(code.eins),Int}() : copy(size_info)
     OMEinsum.get_size_dict!(code, xsa, size_dict)
     # compute intermediate tensors
@@ -138,6 +158,10 @@ function solution_ad(code::NestedEinsum, @nospecialize(xsa), ymask; size_info=no
     n, read_config!(code, mt, Dict())
 end
 
+# get the solution configuration from gradients.
+function read_config!(code::SlicedEinsum, mt, out)
+    read_config!(code.eins, mt, out)
+end
 function read_config!(code::NestedEinsum, mt, out)
     for (arg, ix, sibling) in zip(code.args, getixs(code.eins), mt.siblings)
         if OMEinsum.isleaf(arg)

src/configurations.jl

@@ -17,7 +17,7 @@ function best_solutions(gp::GraphProblem; all=false, usecuda=false)
     T = (all ? set_type : sampler_type)(CountingTropical{Int64}, length(syms), bondsize(gp))
     vertex_index = Dict([s=>i for (i, s) in enumerate(syms)])
     xst = generate_tensors(l->TropicalF64(1.0), gp)
-    ymask = trues(fill(2, length(_getiy(gp.code)))...)
+    ymask = trues(fill(2, length(getiyv(gp.code)))...)
     if usecuda
         xst = CuArray.(xst)
         ymask = CuArray(ymask)
@@ -63,7 +63,7 @@ function bestk_solutions(gp::GraphProblem, k::Int)
     syms = symbols(gp)
     vertex_index = Dict([s=>i for (i, s) in enumerate(syms)])
     xst = generate_tensors(l->TropicalF64(1.0), gp)
-    ymask = trues(fill(2, length(_getiy(gp.code)))...)
+    ymask = trues(fill(2, length(getiyv(gp.code)))...)
     T = set_type(TruncatedPoly{k,Float64,Float64}, length(syms), bondsize(gp))
     xs = generate_tensors(l->onehotv(T, vertex_index[l], 1), gp)
     return bounding_contract(AllConfigs{k}(), gp.code, xst, ymask, xs)

src/graph_polynomials.jl

@@ -62,13 +62,11 @@ function _polynomial_single(gp::GraphProblem, ::Type{T}; usecuda, maxorder) wher
     return res
 end
 
-_getiy(code::EinCode) = getiy(code)
-_getiy(code::NestedEinsum) = getiy(code.eins)
 function graph_polynomial(gp::GraphProblem, ::Val{:finitefield}; usecuda=false,
         maxorder=max_size(gp; usecuda=usecuda), max_iter=100)
     TI = Int32  # Int 32 is faster
     N = typemax(TI)
-    YS = fill(Any[], (fill(bondsize(gp), length(_getiy(gp.code)))...,))
+    YS = fill(Any[], (fill(bondsize(gp), length(getiyv(gp.code)))...,))
     local res, respre
     for k = 1:max_iter
 	    N = prevprime(N-TI(1))

src/networks.jl

@@ -1,5 +1,5 @@
 export Independence, MaximalIndependence, Matching, Coloring, optimize_code, set_packing, MaxCut
-const EinTypes = Union{EinCode,NestedEinsum}
+const EinTypes = Union{EinCode,NestedEinsum,SlicedEinsum}
 
 abstract type GraphProblem end
 

test/bounding.jl

@@ -7,7 +7,7 @@ using GraphTensorNetworks: cached_einsum, generate_masktree, masked_einsum, Cach
     size_dict = uniformsize(code, 2)
     c = cached_einsum(code, xs, size_dict)
     @test c.content == code(xs...)
-    mt = generate_masktree(code, c, rand(Bool,2,2,2), size_dict)
+    mt = generate_masktree(AllConfigs{1}(), code, c, rand(Bool,2,2,2), size_dict)
     @test mt isa CacheTree{Bool}
     y = masked_einsum(code, xs, mt, size_dict)
     @test y isa AbstractArray
@@ -19,7 +19,7 @@ end
         xs = map(x->TropicalF64.(x), [rand(1:5,2,2), rand(1:5,2), rand(1:5,2,2), rand(1:5,2,2), rand(1:5,2,2)])
         code = ein"((ij,j),jk, kl), ii->kli"
         y1 = code(xs...)
-        y2 = bounding_contract(code, xs, BitArray(ones(Bool,2,2,2)), xs)
+        y2 = bounding_contract(AllConfigs{1}(), code, xs, BitArray(ones(Bool,2,2,2)), xs)
         @test y1 ≈ y2
     end
     rawcode = Independence(random_regular_graph(10, 3); optimizer=nothing).code
@@ -28,10 +28,9 @@ end
         length(ix)==1 ? GraphTensorNetworks.misv(TropicalF64(1.0)) : GraphTensorNetworks.misb(TropicalF64)
     end
     y1 = rawcode(xs...)
-    y2 = bounding_contract(rawcode, xs, BitArray(fill(true)), xs)
+    y2 = bounding_contract(AllConfigs{1}(), rawcode, xs, BitArray(fill(true)), xs)
     @test y1 ≈ y2
     y1 = optcode(xs...)
-    y2 = bounding_contract(optcode, xs, BitArray(fill(true)), xs)
+    y2 = bounding_contract(AllConfigs{1}(), optcode, xs, BitArray(fill(true)), xs)
     @test y1 ≈ y2
 end
-

test/interfaces.jl

@@ -3,43 +3,45 @@ using Graphs, Test
 
 @testset "independence problem" begin
     g = Graphs.smallgraph("petersen")
-    gp = Independence(g; optimizer=GreedyMethod())
-    res1 = solve(gp, "size max")[]
-    res2 = solve(gp, "counting sum")[]
-    res3 = solve(gp, "counting max")[]
-    res4 = solve(gp, "counting max2")[]
-    res5 = solve(gp, "counting all")[]
-    res6 = solve(gp, "config max")[]
-    res7 = solve(gp, "configs max")[]
-    res8 = solve(gp, "configs max2")[]
-    res9 = solve(gp, "configs all")[]
-    res10 = solve(gp, "counting all (fft)")[]
-    res11 = solve(gp, "counting all (finitefield)")[]
-    res12 = solve(gp, "config max (bounded)")[]
-    res13 = solve(gp, "configs max (bounded)")[]
-    res14 = solve(gp, "counting max3")[]
-    res15 = solve(gp, "configs max3")[]
-    res16 = solve(gp, "configs max2 (bounded)")[]
-    res17 = solve(gp, "configs max3 (bounded)")[]
-    @test res1.n == 4
-    @test res2 == 76
-    @test res3.n == 4 && res3.c == 5
-    @test res4.maxorder == 4 && res4.coeffs[1] == 30 && res4.coeffs[2]==5
-    @test res5 == Polynomial([1.0, 10.0, 30, 30, 5])
-    @test res6.c.data ∈ res7.c.data
-    @test all(x->sum(x) == 4, res7.c.data)
-    @test all(x->sum(x) == 3, res8.coeffs[1].data) && all(x->sum(x) == 4, res8.coeffs[2].data) && length(res8.coeffs[1].data) == 30 && length(res8.coeffs[2].data) == 5
-    @test all(x->all(c->sum(c) == x[1]-1, x[2].data), enumerate(res9.coeffs))
-    @test res10 ≈ res5
-    @test res11 == res5
-    @test res12.c.data ∈ res13.c.data
-    @test res13.c.data == res7.c.data
-    @test res14.maxorder == 4 && res14.coeffs[1]==30 && res14.coeffs[2] == 30 && res14.coeffs[3]==5
-    @test all(x->sum(x) == 2, res15.coeffs[1].data) && all(x->sum(x) == 3, res15.coeffs[2].data) && all(x->sum(x) == 4, res15.coeffs[3].data) &&
-            length(res15.coeffs[1].data) == 30 && length(res15.coeffs[2].data) == 30 && length(res15.coeffs[3].data) == 5
-    @test all(x->sum(x) == 3, res16.coeffs[1].data) && all(x->sum(x) == 4, res16.coeffs[2].data) && length(res16.coeffs[1].data) == 30 && length(res16.coeffs[2].data) == 5
-    @test all(x->sum(x) == 2, res17.coeffs[1].data) && all(x->sum(x) == 3, res17.coeffs[2].data) && all(x->sum(x) == 4, res17.coeffs[3].data) &&
-            length(res17.coeffs[1].data) == 30 && length(res17.coeffs[2].data) == 30 && length(res17.coeffs[3].data) == 5
+    for optimizer in (GreedyMethod(), TreeSA(ntrials=1))
+        gp = Independence(g; optimizer=optimizer)
+        res1 = solve(gp, "size max")[]
+        res2 = solve(gp, "counting sum")[]
+        res3 = solve(gp, "counting max")[]
+        res4 = solve(gp, "counting max2")[]
+        res5 = solve(gp, "counting all")[]
+        res6 = solve(gp, "config max")[]
+        res7 = solve(gp, "configs max")[]
+        res8 = solve(gp, "configs max2")[]
+        res9 = solve(gp, "configs all")[]
+        res10 = solve(gp, "counting all (fft)")[]
+        res11 = solve(gp, "counting all (finitefield)")[]
+        res12 = solve(gp, "config max (bounded)")[]
+        res13 = solve(gp, "configs max (bounded)")[]
+        res14 = solve(gp, "counting max3")[]
+        res15 = solve(gp, "configs max3")[]
+        res16 = solve(gp, "configs max2 (bounded)")[]
+        res17 = solve(gp, "configs max3 (bounded)")[]
+        @test res1.n == 4
+        @test res2 == 76
+        @test res3.n == 4 && res3.c == 5
+        @test res4.maxorder == 4 && res4.coeffs[1] == 30 && res4.coeffs[2]==5
+        @test res5 == Polynomial([1.0, 10.0, 30, 30, 5])
+        @test res6.c.data ∈ res7.c.data
+        @test all(x->sum(x) == 4, res7.c.data)
+        @test all(x->sum(x) == 3, res8.coeffs[1].data) && all(x->sum(x) == 4, res8.coeffs[2].data) && length(res8.coeffs[1].data) == 30 && length(res8.coeffs[2].data) == 5
+        @test all(x->all(c->sum(c) == x[1]-1, x[2].data), enumerate(res9.coeffs))
+        @test res10 ≈ res5
+        @test res11 == res5
+        @test res12.c.data ∈ res13.c.data
+        @test res13.c == res7.c
+        @test res14.maxorder == 4 && res14.coeffs[1]==30 && res14.coeffs[2] == 30 && res14.coeffs[3]==5
+        @test all(x->sum(x) == 2, res15.coeffs[1].data) && all(x->sum(x) == 3, res15.coeffs[2].data) && all(x->sum(x) == 4, res15.coeffs[3].data) &&
+                length(res15.coeffs[1].data) == 30 && length(res15.coeffs[2].data) == 30 && length(res15.coeffs[3].data) == 5
+        @test all(x->sum(x) == 3, res16.coeffs[1].data) && all(x->sum(x) == 4, res16.coeffs[2].data) && length(res16.coeffs[1].data) == 30 && length(res16.coeffs[2].data) == 5
+        @test all(x->sum(x) == 2, res17.coeffs[1].data) && all(x->sum(x) == 3, res17.coeffs[2].data) && all(x->sum(x) == 4, res17.coeffs[3].data) &&
+                length(res17.coeffs[1].data) == 30 && length(res17.coeffs[2].data) == 30 && length(res17.coeffs[3].data) == 5
+    end
 end
 
 @testset "save load" begin
@@ -79,4 +81,46 @@ end
     save_configs("_test.txt", m; format=:text)
     mb = load_configs("_test.txt"; format=:text, nflavors=3)
     @test mb == m
-end
+end
+
+@testset "slicing" begin
+    g = Graphs.smallgraph("petersen")
+    gp = Independence(g; optimizer=TreeSA(nslices=5, ntrials=1))
+    res1 = solve(gp, "size max")[]
+    res2 = solve(gp, "counting sum")[]
+    res3 = solve(gp, "counting max")[]
+    res4 = solve(gp, "counting max2")[]
+    res5 = solve(gp, "counting all")[]
+    res6 = solve(gp, "config max")[]
+    res7 = solve(gp, "configs max")[]
+    res8 = solve(gp, "configs max2")[]
+    res9 = solve(gp, "configs all")[]
+    res10 = solve(gp, "counting all (fft)")[]
+    res11 = solve(gp, "counting all (finitefield)")[]
+    res12 = solve(gp, "config max (bounded)")[]
+    res13 = solve(gp, "configs max (bounded)")[]
+    res14 = solve(gp, "counting max3")[]
+    res15 = solve(gp, "configs max3")[]
+    res16 = solve(gp, "configs max2 (bounded)")[]
+    res17 = solve(gp, "configs max3 (bounded)")[]
+    @test res1.n == 4
+    @test res2 == 76
+    @test res3.n == 4 && res3.c == 5
+    @test res4.maxorder == 4 && res4.coeffs[1] == 30 && res4.coeffs[2]==5
+    @test res5 == Polynomial([1.0, 10.0, 30, 30, 5])
+    @test res6.c.data ∈ res7.c.data
+    @test all(x->sum(x) == 4, res7.c.data)
+    @test all(x->sum(x) == 3, res8.coeffs[1].data) && all(x->sum(x) == 4, res8.coeffs[2].data) && length(res8.coeffs[1].data) == 30 && length(res8.coeffs[2].data) == 5
+    @test all(x->all(c->sum(c) == x[1]-1, x[2].data), enumerate(res9.coeffs))
+    @test res10 ≈ res5
+    @test res11 == res5
+    @test res12.c.data ∈ res13.c.data
+    @test res13.c == res7.c
+    @test res14.maxorder == 4 && res14.coeffs[1]==30 && res14.coeffs[2] == 30 && res14.coeffs[3]==5
+    @test all(x->sum(x) == 2, res15.coeffs[1].data) && all(x->sum(x) == 3, res15.coeffs[2].data) && all(x->sum(x) == 4, res15.coeffs[3].data) &&
+            length(res15.coeffs[1].data) == 30 && length(res15.coeffs[2].data) == 30 && length(res15.coeffs[3].data) == 5
+    @test all(x->sum(x) == 3, res16.coeffs[1].data) && all(x->sum(x) == 4, res16.coeffs[2].data) && length(res16.coeffs[1].data) == 30 && length(res16.coeffs[2].data) == 5
+    @test all(x->sum(x) == 2, res17.coeffs[1].data) && all(x->sum(x) == 3, res17.coeffs[2].data) && all(x->sum(x) == 4, res17.coeffs[3].data) &&
+            length(res17.coeffs[1].data) == 30 && length(res17.coeffs[2].data) == 30 && length(res17.coeffs[3].data) == 5
+end
+