...

Author: GiggleLiu

src/GraphTensorNetworks.jl

@@ -4,7 +4,7 @@ using OMEinsumContractionOrders: OMEinsum
 using Core: Argument
 using TropicalGEMM, TropicalNumbers
 using OMEinsum
-using OMEinsum: flatten, timespace_complexity
+using OMEinsum: timespace_complexity
 using LightGraphs
 
 export timespace_complexity, @ein_str

src/bounding.jl

@@ -68,7 +68,7 @@ function generate_masktree(code::Int, cache, mask, size_dict, mode=:all)
     CacheTree(mask, CacheTree{Bool}[])
 end
 function generate_masktree(code::NestedEinsum, cache, mask, size_dict, mode=:all)
-    submasks = backward_tropical(mode, OMEinsum.getixs(code.eins), (getfield.(cache.siblings, :content)...,), OMEinsum.getiy(code.eins), cache.content, mask, size_dict)
+    submasks = backward_tropical(mode, getixs(code.eins), (getfield.(cache.siblings, :content)...,), OMEinsum.getiy(code.eins), cache.content, mask, size_dict)
     return CacheTree(mask, generate_masktree.(code.args, cache.siblings, submasks, Ref(size_dict), mode))
 end
 
@@ -96,7 +96,7 @@ function bounding_contract(@nospecialize(code::EinCode), @nospecialize(xsa), yma
     bounding_contract(NestedEinsum((1:length(xsa)), code), xsa, ymask, xsb; size_info=size_info)
 end
 function bounding_contract(code::NestedEinsum, @nospecialize(xsa), ymask, @nospecialize(xsb); size_info=nothing)
-    size_dict = OMEinsum.get_size_dict(getixs(flatten(code)), xsa, size_info)
+    size_dict = OMEinsum.get_size_dict(collect_ixs(code), xsa, size_info)
     # compute intermediate tensors
     @debug "caching einsum..."
     c = cached_einsum(code, xsa, size_dict)
@@ -113,7 +113,7 @@ function solution_ad(@nospecialize(code::EinCode), @nospecialize(xsa), ymask; si
 end
 
 function solution_ad(code::NestedEinsum, @nospecialize(xsa), ymask; size_info=nothing)
-    size_dict = OMEinsum.get_size_dict(getixs(flatten(code)), xsa, size_info)
+    size_dict = OMEinsum.get_size_dict(collect_ixs(code), xsa, size_info)
     # compute intermediate tensors
     @debug "caching einsum..."
     c = cached_einsum(code, xsa, size_dict)
@@ -125,7 +125,7 @@ function solution_ad(code::NestedEinsum, @nospecialize(xsa), ymask; size_info=no
 end
 
 function read_config!(code::NestedEinsum, mt, out)
-    for (arg, ix, sibling) in zip(code.args, OMEinsum.getixs(code.eins), mt.siblings)
+    for (arg, ix, sibling) in zip(code.args, getixs(code.eins), mt.siblings)
         if arg isa Int
             assign = convert(Array, sibling.content)  # note: the content can be CuArray
             if length(ix) == 1

src/configurations.jl

@@ -16,7 +16,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(OMEinsum.getiy(flatten(gp.code))))...)
+    ymask = trues(fill(2, length(_getiy(gp.code)))...)
     if usecuda
         xst = CuArray.(xst)
         ymask = CuArray(ymask)
@@ -89,4 +89,4 @@ end
 for GP in [:Independence, :Matching, :MaximalIndependence, :Coloring]
     @eval symbols(gp::$GP) = labels(gp.code)
 end
-symbols(gp::MaxCut) = collect(OMEinsum.getixs(OMEinsum.flatten(gp.code)))
+symbols(gp::MaxCut) = collect_ixs(gp.code)

src/graph_polynomials.jl

@@ -62,11 +62,13 @@ 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(flatten(gp.code))))...,))
+    YS = fill(Any[], (fill(bondsize(gp), length(_getiy(gp.code)))...,))
     local res, respre
     for k = 1:max_iter
 	    N = prevprime(N-TI(1))
@@ -103,9 +105,8 @@ end
 ############### Problem specific implementations ################
 ### independent set ###
 function generate_tensors(fx, gp::Independence)
-    flatten_code = flatten(gp.code)
-    ixs = getixs(flatten_code)
-    n = length(labels(flatten_code))
+    ixs = collect_ixs(gp.code)
+    n = length(unique!(vcat(ixs...)))
     T = typeof(fx(ixs[1][1]))
     return Tuple(map(enumerate(ixs)) do (i, ix)
         if i <= n
@@ -115,6 +116,23 @@ function generate_tensors(fx, gp::Independence)
         end
     end)
 end
+
+function collect_ixs(ne::NestedEinsum)
+    d = collect_ixs!(ne, Dict{Int,Vector{OMEinsum.labeltype(ne.eins)}}())
+    return [d[i] for i=1:length(d)]
+end
+
+function collect_ixs!(ne::NestedEinsum, d::Dict)
+    for i=1:length(ne.args)
+        if ne.args[i] isa Integer
+            d[ne.args[i]] = collect(OMEinsum.getixs(ne.eins)[i])
+        else
+            collect_ixs!(ne.args[i], d)
+        end
+    end
+    return d
+end
+
 function misb(::Type{T}, n::Integer=2) where T
     res = zeros(T, fill(2, n)...)
     res[1] = one(T)
@@ -127,7 +145,7 @@ misv(val::T) where T = [one(T), val]
 
 ### coloring ###
 function generate_tensors(fx, c::Coloring{K}) where K
-    ixs = getixs(flatten(c.code))
+    ixs = collect_ixs(c.code)
     T = eltype(fx(ixs[1][1]))
     return map(ixs) do ix
         # if the tensor rank is 1, create a vertex tensor.
@@ -149,9 +167,9 @@ coloringv(vals::Vector{T}) where T = vals
 
 ### matching ###
 function generate_tensors(fx, m::Matching)
-    ixs = OMEinsum.getixs(flatten(m.code))
+    ixs = collect_ixs(m.code)
     T = typeof(fx(ixs[1][1]))
-    n = length(unique(vcat(collect.(ixs)...)))  # number of vertices
+    n = length(unique!(vcat(ixs...)))  # number of vertices
     tensors = []
     for i=1:length(ixs)
         if i<=n
@@ -176,7 +194,7 @@ end
 
 ### maximal independent set ###
 function generate_tensors(fx, mi::MaximalIndependence)
-    ixs = OMEinsum.getixs(flatten(mi.code))
+    ixs = collect_ixs(mi.code)
     T = eltype(fx(ixs[1][end]))
 	return map(ixs) do ix
         neighbortensor(fx(ix[end]), length(ix))
@@ -193,11 +211,10 @@ end
 
 ### max cut/spin glass problem ###
 function generate_tensors(fx, gp::MaxCut)
-    flatten_code = flatten(gp.code)
-    ixs = getixs(flatten_code)
-    return Tuple(map(enumerate(ixs)) do (i, ix)
+    ixs = collect_ixs(gp.code)
+    return map(enumerate(ixs)) do (i, ix)
         maxcutb(fx(ix))
-    end)
+    end
 end
 function maxcutb(expJ::T) where T
     return T[one(T) expJ; expJ one(T)]

src/networks.jl

@@ -5,52 +5,53 @@ abstract type GraphProblem end
 
 """
     Independence{CT<:EinTypes} <: GraphProblem
-    Independence(graph; kwargs...)
+    Independence(graph; openvertices=(), kwargs...)
 
 Independent set problem. `kwargs` is forwarded to `optimize_code`.
 """
 struct Independence{CT<:EinTypes} <: GraphProblem
     code::CT
 end
 
-function Independence(g::SimpleGraph; outputs=(), kwargs...)
+function Independence(g::SimpleGraph; openvertices=(), kwargs...)
     rawcode = EinCode(([(i,) for i in LightGraphs.vertices(g)]..., # labels for vertex tensors
-                    [minmax(e.src,e.dst) for e in LightGraphs.edges(g)]...), outputs)  # labels for edge tensors
-    Independence(optimize_code(rawcode; kwargs...))
+                    [minmax(e.src,e.dst) for e in LightGraphs.edges(g)]...), openvertices)  # labels for edge tensors
+    code = optimize_code(rawcode; kwargs...)
+    Independence(code)
 end
 
 """
     MaxCut{CT<:EinTypes} <: GraphProblem
-    MaxCut(graph; kwargs...)
+    MaxCut(graph; openvertices=(), kwargs...)
 
 Max cut problem (or spin glass problem). `kwargs` is forwarded to `optimize_code`.
 """
 struct MaxCut{CT<:EinTypes} <: GraphProblem
     code::CT
 end
-function MaxCut(g::SimpleGraph; outputs=(), kwargs...)
-    rawcode = EinCode(([minmax(e.src,e.dst) for e in LightGraphs.edges(g)]...,), outputs)  # labels for edge tensors
+function MaxCut(g::SimpleGraph; openvertices=(), kwargs...)
+    rawcode = EinCode(([minmax(e.src,e.dst) for e in LightGraphs.edges(g)]...,), openvertices)  # labels for edge tensors
     MaxCut(optimize_code(rawcode; kwargs...))
 end
 
 """
     MaximalIndependence{CT<:EinTypes} <: GraphProblem
-    MaximalIndependence(graph; kwargs...)
+    MaximalIndependence(graph; openvertices=(), kwargs...)
 
 Maximal independent set problem. `kwargs` is forwarded to `optimize_code`.
 """
 struct MaximalIndependence{CT<:EinTypes} <: GraphProblem
     code::CT
 end
 
-function MaximalIndependence(g::SimpleGraph; outputs=(), kwargs...)
-    rawcode = EinCode(([(LightGraphs.neighbors(g, v)..., v) for v in LightGraphs.vertices(g)]...,), outputs)
+function MaximalIndependence(g::SimpleGraph; openvertices=(), kwargs...)
+    rawcode = EinCode(([(LightGraphs.neighbors(g, v)..., v) for v in LightGraphs.vertices(g)]...,), openvertices)
     MaximalIndependence(optimize_code(rawcode; kwargs...))
 end
 
 """
     Matching{CT<:EinTypes} <: GraphProblem
-    Matching(graph; kwargs...)
+    Matching(graph; openvertices=(), kwargs...)
 
 Vertex matching problem. `kwargs` is forwarded to `optimize_code`.
 The matching polynomial adopts the first definition in wiki page: https://en.wikipedia.org/wiki/Matching_polynomial
@@ -63,15 +64,15 @@ struct Matching{CT<:EinTypes} <: GraphProblem
     code::CT
 end
 
-function Matching(g::SimpleGraph; outputs=(), kwargs...)
+function Matching(g::SimpleGraph; openvertices=(), kwargs...)
     rawcode = EinCode(([(minmax(e.src,e.dst),) for e in LightGraphs.edges(g)]..., # labels for edge tensors
-                    [([minmax(i,j) for j in neighbors(g, i)]...,) for i in LightGraphs.vertices(g)]...,), outputs)       # labels for vertex tensors
+                    [([minmax(i,j) for j in neighbors(g, i)]...,) for i in LightGraphs.vertices(g)]...,), openvertices)       # labels for vertex tensors
     Matching(optimize_code(rawcode; kwargs...))
 end
 
 """
     Coloring{K,CT<:EinTypes} <: GraphProblem
-    Coloring{K}(graph; kwargs...)
+    Coloring{K}(graph; openvertices=(), kwargs...)
 
 K-Coloring problem. `kwargs` is forwarded to `optimize_code`.
 """
@@ -80,7 +81,7 @@ struct Coloring{K,CT<:EinTypes} <: GraphProblem
 end
 Coloring{K}(code::ET) where {K,ET<:EinTypes} = Coloring{K,ET}(code)
 # same network layout as independent set.
-Coloring{K}(g::SimpleGraph; outputs=(), kwargs...) where K = Coloring{K}(Independence(g; outputs=outputs, kwargs...).code)
+Coloring{K}(g::SimpleGraph; openvertices=(), kwargs...) where K = Coloring{K}(Independence(g; openvertices=openvertices, kwargs...).code)
 
 """
     labels(code)
@@ -89,7 +90,7 @@ Return a vector of unique labels in an Einsum token.
 """
 function labels(code::EinTypes)
     res = []
-    for ix in OMEinsum.getixs(OMEinsum.flatten(code))
+    for ix in collect_ixs(code)
         for l in ix
             if l ∉ res
                 push!(res, l)

src/viz.jl

@@ -2,14 +2,14 @@ using Viznet
 export vizeinsum, vizconfig
 using Compose
 
-function vizconfig(g::SimpleGraph; locs, config=zeros(Int, length(locs)), unit=1.0, graphsize=12cm)
-    vizconfig([string(v)=>locs[v] for v in LightGraphs.vertices(g)], [(e.src, e.dst) for e in edges(g)]; config=config, unit=unit, graphsize=graphsize)
+function vizconfig(g::SimpleGraph; locs, config=zeros(Int, length(locs)), unit=1.0, graphsize=12cm, radius=0.03)
+    vizconfig([string(v)=>locs[v] for v in LightGraphs.vertices(g)], [(e.src, e.dst) for e in edges(g)]; config=config, unit=unit, graphsize=graphsize, radius=radius)
 end
 
-function vizconfig(nodes, edges; config=zeros(Int, length(nodes)), unit=1.0, graphsize=12cm)
+function vizconfig(nodes, edges; config=zeros(Int, length(nodes)), unit=1.0, graphsize=12cm, radius=0.03)
 	tb = textstyle(:default, fill("white"), fontsize(10pt*unit))
-	nb = nodestyle(:circle, fill("black"), r=0.03*unit)
-	nb2 = nodestyle(:circle, fill("red"),r=0.03*unit)
+	nb = nodestyle(:circle, fill("black"), r=radius*unit)
+	nb2 = nodestyle(:circle, fill("red"),r=radius*unit)
 	eb = bondstyle(:default, linewidth(0.4mm*unit))
 	img = canvas() do
 		for (i, (t, p)) in enumerate(nodes)
@@ -69,11 +69,11 @@ function vizeinsum(::EinCode{ixs, iy}, locs::AbstractVector{<:Pair}; kwargs...)
 	vizeinsum(ixs, iy, Dict(locs); kwargs...)
 end
 function vizeinsum(code::NestedEinsum, locs::AbstractVector{<:Pair}; kwargs...)
-	vizeinsum(flatten(code), locs; kwargs...)
+	vizeinsum(OMEinsum.flatten(code), locs; kwargs...)
 end
 function vizeinsum(ixs::Tuple, iy::Tuple, locs::Dict; kwargs...)
 	legs = unique!([Iterators.flatten(ixs)..., iy...])
 	nodes = [l=>locs[l] for l in legs]
 	edges = [map(i->findfirst(==(i), legs), ix) for ix in ixs]
 	vizeinsum(nodes, edges; config=[l ∈ iy for l in legs], kwargs...)
-end
+end