diff --git a/Project.toml b/Project.toml index 36563d1e..5d1e26f4 100644 --- a/Project.toml +++ b/Project.toml @@ -4,25 +4,33 @@ authors = ["Matthew Fishman and contributors"] version = "0.2.2" [deps] +Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa" Compat = "34da2185-b29b-5c13-b0c7-acf172513d20" DataGraphs = "b5a273c3-7e6c-41f6-98bd-8d7f1525a36a" Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4" DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" +GraphRecipes = "bd48cda9-67a9-57be-86fa-5b3c104eda73" Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6" +GraphsFlows = "06909019-6f44-4949-96fc-b9d9aaa02889" ITensors = "9136182c-28ba-11e9-034c-db9fb085ebd5" IsApprox = "28f27b66-4bd8-47e7-9110-e2746eb8bed7" IterTools = "c8e1da08-722c-5040-9ed9-7db0dc04731e" KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" NamedGraphs = "678767b0-92e7-4007-89e4-4527a8725b19" +OMEinsumContractionOrders = "6f22d1fd-8eed-4bb7-9776-e7d684900715" +Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" +Reexport = "189a3867-3050-52da-a836-e630ba90ab69" Observers = "338f10d5-c7f1-4033-a7d1-f9dec39bcaa0" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" Requires = "ae029012-a4dd-5104-9daa-d747884805df" +SimpleWeightedGraphs = "47aef6b3-ad0c-573a-a1e2-d07658019622" SimpleTraits = "699a6c99-e7fa-54fc-8d76-47d257e15c1d" SparseArrayKit = "a9a3c162-d163-4c15-8926-b8794fbefed2" SplitApplyCombine = "03a91e81-4c3e-53e1-a0a4-9c0c8f19dd66" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb" +SweepContractor = "75a5deae-e917-4509-af32-a989148c8d5f" TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f" [compat] @@ -37,7 +45,7 @@ Requires = "1.3" SimpleTraits = "0.9" SplitApplyCombine = "1.2" Suppressor = "0.2" -julia = "1.7" +julia = "1.6, 1.7, 1.8" [extras] Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" diff --git a/examples/approximate_contract/3dcube.jl b/examples/approximate_contract/3dcube.jl new file mode 100644 index 00000000..e0675ac2 --- /dev/null +++ b/examples/approximate_contract/3dcube.jl @@ -0,0 +1,366 @@ +using ITensors, TimerOutputs, Graphs +using KaHyPar +using ITensorNetworks +using ITensorNetworks: contraction_sequence, ITensorNetwork, ising_network, vertex_tag +using ITensorNetworks.ApproximateTNContraction: + approximate_contract, line_to_tree, timer, line_network + +INDEX = 0 + +function contract_log_norm(tn, seq) + global INDEX + if seq isa Vector + if length(seq) == 1 + return seq[1] + end + t1 = contract_log_norm(tn, seq[1]) + t2 = contract_log_norm(tn, seq[2]) + @info size(t1[1]), size(t2[1]) + INDEX += 1 + @info "INDEX", INDEX + out = t1[1] * t2[1] + nrm = norm(out) + out /= nrm + lognrm = log(nrm) + t1[2] + t2[2] + return (out, lognrm) + else + return tn[seq] + end +end + +function exact_contract(N; beta, sc_target) + ITensors.set_warn_order(1000) + reset_timer!(timer) + linkdim = 2 + network = ising_network(named_grid(N), beta) + tn = Array{ITensor,length(N)}(undef, N...) + for v in vertices(network) + tn[v...] = network[v...] + end + tn = vec(tn) + seq = contraction_sequence(tn; alg="kahypar_bipartite", sc_target=sc_target) + @info seq + tn = [(i, 0.0) for i in tn] + return contract_log_norm(tn, seq) +end + +function build_tntree(tn, N; env_size) + @assert length(N) == length(env_size) + n = [ceil(Int, N[i] / env_size[i]) for i in 1:length(N)] + tntree = nothing + for k in 1:n[3] + for j in 1:n[2] + for i in 1:n[1] + ii = (i - 1) * env_size[1] + jj = (j - 1) * env_size[2] + kk = (k - 1) * env_size[3] + ii_end = min(ii + env_size[1], N[1]) + jj_end = min(jj + env_size[2], N[2]) + kk_end = min(kk + env_size[3], N[3]) + sub_tn = tn[(ii + 1):ii_end, (jj + 1):jj_end, (kk + 1):kk_end] + sub_tn = vec(sub_tn) + if tntree == nothing + tntree = sub_tn + else + tntree = [tntree, sub_tn] + end + end + end + end + return tntree +end + +function build_recursive_tntree(tn, N; env_size) + @assert env_size == (3, 3, 1) + tn_tree1 = vec(tn[1:3, 1:3, 1]) + tn_tree1 = [vec(tn[1:3, 1:3, 2]), tn_tree1] + tn_tree1 = [vec(tn[1:3, 1:3, 3]), tn_tree1] + + tn_tree2 = vec(tn[1:3, 4:6, 1]) + tn_tree2 = [vec(tn[1:3, 4:6, 2]), tn_tree2] + tn_tree2 = [vec(tn[1:3, 4:6, 3]), tn_tree2] + + tn_tree3 = vec(tn[4:6, 1:3, 1]) + tn_tree3 = [vec(tn[4:6, 1:3, 2]), tn_tree3] + tn_tree3 = [vec(tn[4:6, 1:3, 3]), tn_tree3] + + tn_tree4 = vec(tn[4:6, 4:6, 1]) + tn_tree4 = [vec(tn[4:6, 4:6, 2]), tn_tree4] + tn_tree4 = [vec(tn[4:6, 4:6, 3]), tn_tree4] + + tn_tree5 = vec(tn[1:3, 1:3, 6]) + tn_tree5 = [vec(tn[1:3, 1:3, 5]), tn_tree5] + tn_tree5 = [vec(tn[1:3, 1:3, 4]), tn_tree5] + + tn_tree6 = vec(tn[1:3, 4:6, 6]) + tn_tree6 = [vec(tn[1:3, 4:6, 5]), tn_tree6] + tn_tree6 = [vec(tn[1:3, 4:6, 4]), tn_tree6] + + tn_tree7 = vec(tn[4:6, 1:3, 6]) + tn_tree7 = [vec(tn[4:6, 1:3, 5]), tn_tree7] + tn_tree7 = [vec(tn[4:6, 1:3, 4]), tn_tree7] + + tn_tree8 = vec(tn[4:6, 4:6, 6]) + tn_tree8 = [vec(tn[4:6, 4:6, 5]), tn_tree8] + tn_tree8 = [vec(tn[4:6, 4:6, 4]), tn_tree8] + return [ + [[tn_tree1, tn_tree2], [tn_tree3, tn_tree4]], + [[tn_tree5, tn_tree6], [tn_tree7, tn_tree8]], + ] +end + +# if ortho == true +# @info "orthogonalize tn towards the first vertex" +# itn = ITensorNetwork(named_grid(N); link_space=2) +# for i in 1:N[1] +# for j in 1:N[2] +# for k in 1:N[3] +# itn[i, j, k] = tn[i, j, k] +# end +# end +# end +# itn = orthogonalize(itn, (1, 1, 1)) +# @info itn[1, 1, 1] +# @info itn[1, 1, 1].tensor +# for i in 1:N[1] +# for j in 1:N[2] +# for k in 1:N[3] +# tn[i, j, k] = itn[i, j, k] +# end +# end +# end +# end +function build_tntree(N; block_size, beta, h, snake, env_size, szverts) + @info "beta is", beta + ITensors.set_warn_order(100) + network = ising_network(named_grid(N), beta, h; szverts=szverts) + tn = Array{ITensor,length(N)}(undef, N...) + for v in vertices(network) + tn[v...] = network[v...] + end + if snake == true + for k in 1:N[3] + rangej = iseven(k) ? reverse(1:N[2]) : 1:N[2] + tn[:, rangej, k] = tn[:, 1:N[2], k] + end + end + if block_size == (1, 1, 1) + return build_tntree(tn, N; env_size=env_size) + end + tn_reduced = ITensorNetwork() + reduced_N = ( + ceil(Int, N[1] / block_size[1]), + ceil(Int, N[2] / block_size[2]), + ceil(Int, N[3] / block_size[3]), + ) + for i in 1:reduced_N[1] + for j in 1:reduced_N[2] + for k in 1:reduced_N[3] + add_vertex!(tn_reduced, (i, j, k)) + ii = (i - 1) * block_size[1] + jj = (j - 1) * block_size[2] + kk = (k - 1) * block_size[3] + ii_end = min(ii + block_size[1], N[1]) + jj_end = min(jj + block_size[2], N[2]) + kk_end = min(kk + block_size[3], N[3]) + tn_reduced[(i, j, k)] = ITensors.contract( + tn[(ii + 1):ii_end, (jj + 1):jj_end, (kk + 1):kk_end]... + ) + end + end + end + for e in edges(tn_reduced) + v1, v2 = e.src, e.dst + C = combiner( + commoninds(tn_reduced[v1], tn_reduced[v2])...; + tags="$(vertex_tag(v1))↔$(vertex_tag(v2))", + ) + tn_reduced[v1] = tn_reduced[v1] * C + tn_reduced[v2] = tn_reduced[v2] * C + end + network_reduced = Array{ITensor,3}(undef, reduced_N...) + for v in vertices(tn_reduced) + network_reduced[v...] = tn_reduced[v...] + end + reduced_env = ( + ceil(Int, env_size[1] / block_size[1]), + ceil(Int, env_size[2] / block_size[2]), + ceil(Int, env_size[3] / block_size[3]), + ) + return build_tntree(network_reduced, reduced_N; env_size=reduced_env) +end + +function bench_3d_cube_lnZ( + N; + block_size, + beta, + h, + num_iter, + cutoff, + maxdim, + ansatz, + algorithm, + snake, + use_cache, + ortho, + env_size, +) + reset_timer!(timer) + tntree = build_tntree( + N; + block_size=block_size, + beta=beta, + h=h, + snake=snake, + env_size=env_size, + szverts=nothing, + ) + function _run() + out, log_acc_norm = approximate_contract( + tntree; + cutoff=cutoff, + maxdim=maxdim, + ansatz=ansatz, + algorithm=algorithm, + use_cache=use_cache, + orthogonalize=ortho, + ) + @info "out is", log(out[1][1]) + log_acc_norm + return log(out[1][1]) + log_acc_norm + end + out_list = [] + for _ in 1:num_iter + push!(out_list, _run()) + end + show(timer) + # after warmup, start to benchmark + reset_timer!(timer) + for _ in 1:num_iter + push!(out_list, _run()) + end + @info "lnZ results are", out_list, "mean is", sum(out_list) / (num_iter * 2) + return show(timer) +end + +function bench_3d_cube_magnetization( + N; + block_size, + beta, + h, + num_iter, + cutoff, + maxdim, + ansatz, + algorithm, + snake, + use_cache, + ortho, + env_size, + szverts, +) + reset_timer!(timer) + tntree1 = build_tntree( + N; + block_size=block_size, + beta=beta, + h=h, + snake=snake, + env_size=env_size, + szverts=szverts, + ) + tntree2 = build_tntree( + N; + block_size=block_size, + beta=beta, + h=h, + snake=snake, + env_size=env_size, + szverts=nothing, + ) + function _run() + out, log_acc_norm = approximate_contract( + tntree1; + cutoff=cutoff, + maxdim=maxdim, + ansatz=ansatz, + algorithm=algorithm, + use_cache=use_cache, + orthogonalize=ortho, + ) + lognorm1 = log(out[1][1]) + log_acc_norm + out, log_acc_norm = approximate_contract( + tntree2; + cutoff=cutoff, + maxdim=maxdim, + ansatz=ansatz, + algorithm=algorithm, + use_cache=use_cache, + orthogonalize=ortho, + ) + lognorm2 = log(out[1][1]) + log_acc_norm + return lognorm1 / lognorm2 + end + out_list = [] + for _ in 1:num_iter + push!(out_list, _run()) + end + show(timer) + # after warmup, start to benchmark + reset_timer!(timer) + for _ in 1:num_iter + push!(out_list, _run()) + end + @info "magnetization results are", out_list, "mean is", sum(out_list) / (num_iter * 2) + return show(timer) +end + +# exact_contract((4, 4, 10); beta=0.3, sc_target=28) +# TODO: (6, 6, 6), env_size=(2, 1, 1) is buggy (cutoff=1e-12, maxdim=256, ansatz="comb", algorithm="density_matrix",) +# TODO below is buggy +# @time bench_3d_cube_lnZ( +# (3, 8, 10); +# use_2D=false, +# beta=0.3, +# h=0.0, +# num_iter=2, +# cutoff=1e-20, +# maxdim=128, +# ansatz="mps", +# algorithm="density_matrix", +# snake=false, +# use_cache=true, +# ortho=false, +# env_size=(3, 1, 1), +# ) +@time bench_3d_cube_lnZ( + (6, 6, 6); + block_size=(6, 1, 1), + beta=0.3, + h=0.0, + num_iter=2, + cutoff=1e-12, + maxdim=128, + ansatz="mps", + algorithm="density_matrix", + snake=false, + use_cache=true, + ortho=false, + env_size=(6, 1, 1), +) + +# @time bench_3d_cube_magnetization( +# (1, 6, 6); +# use_2D=true, +# beta=0.44, +# h=0.0001, +# num_iter=2, +# cutoff=1e-20, +# maxdim=64, +# ansatz="mps", +# algorithm="density_matrix", +# snake=false, +# use_cache=true, +# ortho=false, +# env_size=(1, 6, 1), +# szverts=[(1, 3, 3)], +# ) diff --git a/src/ApproximateTNContraction/ApproximateTNContraction.jl b/src/ApproximateTNContraction/ApproximateTNContraction.jl new file mode 100644 index 00000000..b221366b --- /dev/null +++ b/src/ApproximateTNContraction/ApproximateTNContraction.jl @@ -0,0 +1,17 @@ +@reexport module ApproximateTNContraction + +using ITensors + +using ITensors: data, contract + +using NamedGraphs +using TimerOutputs + +const timer = TimerOutput() + +include("orthogonal_tensor.jl") +include("networks/itensor_network.jl") +include("interfaces/sweep_contractor.jl") +include("contract/contract.jl") + +end diff --git a/src/ApproximateTNContraction/contract/contract.jl b/src/ApproximateTNContraction/contract/contract.jl new file mode 100644 index 00000000..44a00086 --- /dev/null +++ b/src/ApproximateTNContraction/contract/contract.jl @@ -0,0 +1,1200 @@ +using Graphs, GraphsFlows, Combinatorics, SimpleWeightedGraphs +using GraphRecipes, Plots +using OMEinsumContractionOrders +using ITensorNetworks: contraction_sequence, TTN, IndsNetwork + +function Base.show(io::IO, tensor::ITensor) + return print(io, string(inds(tensor))) +end + +include("tree_utils.jl") +include("index_group.jl") +include("mincut_tree.jl") +include("tree_embedding.jl") + +function optcontract(t_list::Vector) + @timeit timer "optcontract" begin + if length(t_list) == 0 + return OrthogonalITensor(ITensor(1.0)) + end + @assert t_list isa Vector{OrthogonalITensor} + t_list = get_tensors(t_list) + # @info "start contract, size is", size(t_list) + # for t in t_list + # @info "size of t is", size(t) + # end + @timeit timer "contraction_sequence" begin + seq = contraction_sequence(t_list; alg="sa_bipartite") + end + @timeit timer "contract" begin + output = contract(t_list; sequence=seq) + end + # @info "finish contract" + return OrthogonalITensor(output) + end +end + +approximate_contract(tn::ITensor, inds_groups; kwargs...) = [tn], 0.0 + +approximate_contract(tn::OrthogonalITensor, inds_groups; kwargs...) = [tn], 0.0 + +function approximate_contract(tn::Vector{ITensor}, inds_btree=nothing; kwargs...) + out, log_norm = approximate_contract(orthogonal_tensors(tn), inds_btree; kwargs...) + return get_tensors(out), log_norm +end + +function approximate_contract(tn::Vector{OrthogonalITensor}, inds_btree=nothing; kwargs...) + ctree_to_tensor, log_root_norm = approximate_contract_ctree_to_tensor( + tn, inds_btree; kwargs... + ) + return Vector{OrthogonalITensor}(vcat(collect(values(ctree_to_tensor))...)), log_root_norm +end + +function approximate_contract_ctree_to_tensor( + tn::Vector{OrthogonalITensor}, + inds_btree=nothing; + cutoff, + maxdim, + maxsize=10^15, + ansatz="mps", + algorithm="density_matrix", +) + uncontract_inds = noncommoninds(tn...) + allinds = collect(Set(mapreduce(t -> collect(inds(t)), vcat, tn))) + innerinds = setdiff(allinds, uncontract_inds) + if length(uncontract_inds) <= 2 + if inds_btree == nothing + inds_btree = [[i] for i in uncontract_inds] + end + return Dict{Vector,OrthogonalITensor}(inds_btree => optcontract(tn)), 0.0 + end + # # cases where tn is a tree, or contains 2 disconnected trees + # if length(innerinds) <= length(tn) - 1 + # # TODO + # return tn + # end + # # TODO: may want to remove this + # if inds_groups != nothing + # deltainds = vcat(filter(g -> length(g) > 1, inds_groups)...) + # deltas, tnprime, _ = split_deltas(deltainds, tn) + # tn = Vector{ITensor}(vcat(deltas, tnprime)) + # end + if inds_btree == nothing + inds_btree = inds_binary_tree(get_tensors(tn), nothing; algorithm=ansatz) + end + embedding = tree_embedding(tn, inds_btree; algorithm=algorithm) + tn = Vector{OrthogonalITensor}(vcat(collect(values(embedding))...)) + i2 = noncommoninds(tn...) + @assert (length(uncontract_inds) == length(i2)) + @timeit timer "tree_approximation" begin + return tree_approximation( + embedding, + inds_btree; + cutoff=cutoff, + maxdim=maxdim, + maxsize=maxsize, + algorithm=algorithm, + ) + end +end + +function uncontractinds(tn) + if tn isa ITensor + return inds(tn) + else + return noncommoninds(vectorize(tn)...) + end +end + +# Note that the children ordering matters here. +mutable struct IndexAdjacencyTree + children::Union{Vector{IndexAdjacencyTree},Vector{IndexGroup}} + fixed_direction::Bool + fixed_order::Bool +end + +function Base.copy(tree::IndexAdjacencyTree) + node_to_copynode = Dict{IndexAdjacencyTree,IndexAdjacencyTree}() + for node in topo_sort(tree; type=IndexAdjacencyTree) + if node.children isa Vector{IndexGroup} + node_to_copynode[node] = IndexAdjacencyTree( + node.children, node.fixed_direction, node.fixed_order + ) + continue + end + copynode = IndexAdjacencyTree( + [node_to_copynode[n] for n in node.children], node.fixed_direction, node.fixed_order + ) + node_to_copynode[node] = copynode + end + return node_to_copynode[tree] +end + +function Base.show(io::IO, tree::IndexAdjacencyTree) + out_str = "\n" + stack = [tree] + node_to_level = Dict{IndexAdjacencyTree,Int}() + node_to_level[tree] = 0 + # pre-order traversal + while length(stack) != 0 + node = pop!(stack) + indent_vec = [" " for _ in 1:node_to_level[node]] + indent = string(indent_vec...) + if node.children isa Vector{IndexGroup} + for c in node.children + out_str = out_str * indent * string(c) * "\n" + end + else + out_str = + out_str * + indent * + "AdjTree: [fixed_direction]: " * + string(node.fixed_direction) * + " [fixed_order]: " * + string(node.fixed_order) * + "\n" + for c in node.children + node_to_level[c] = node_to_level[node] + 1 + push!(stack, c) + end + end + end + return print(io, out_str) +end + +function IndexAdjacencyTree(index_group::IndexGroup) + return IndexAdjacencyTree([index_group], false, false) +end + +function get_adj_tree_leaves(tree::IndexAdjacencyTree) + if tree.children isa Vector{IndexGroup} + return tree.children + end + leaves = [get_adj_tree_leaves(c) for c in tree.children] + return vcat(leaves...) +end + +function Base.contains(adj_tree::IndexAdjacencyTree, adj_igs::Set{IndexGroup}) + leaves = Set(get_adj_tree_leaves(adj_tree)) + return issubset(adj_igs, leaves) +end + +function Base.iterate(x::IndexAdjacencyTree) + return iterate(x, 1) +end + +function Base.iterate(x::IndexAdjacencyTree, index) + if index > length(x.children) + return nothing + end + return x.children[index], index + 1 +end + +function boundary_state(ancestor::IndexAdjacencyTree, adj_igs::Set{IndexGroup}) + if ancestor.children isa Vector{IndexGroup} + return "all" + end + if !ancestor.fixed_order + filter_children = filter(a -> contains(a, adj_igs), ancestor.children) + @assert length(filter_children) <= 1 + if length(filter_children) == 1 + return "middle" + elseif Set(get_adj_tree_leaves(ancestor)) == adj_igs + return "all" + else + return "invalid" + end + end + @assert length(ancestor.children) >= 2 + if contains(ancestor.children[1], adj_igs) + return "left" + elseif contains(ancestor.children[end], adj_igs) + return "right" + elseif Set(get_adj_tree_leaves(ancestor)) == adj_igs + return "all" + else + return "invalid" + end +end + +function reorder_to_right!( + ancestor::IndexAdjacencyTree, filter_children::Vector{IndexAdjacencyTree} +) + remain_children = setdiff(ancestor.children, filter_children) + @assert length(filter_children) >= 1 + @assert length(remain_children) >= 1 + if length(remain_children) == 1 + new_child1 = remain_children[1] + else + new_child1 = IndexAdjacencyTree(remain_children, false, false) + end + if length(filter_children) == 1 + new_child2 = filter_children[1] + else + new_child2 = IndexAdjacencyTree(filter_children, false, false) + end + ancestor.children = [new_child1, new_child2] + return ancestor.fixed_order = true +end + +""" +reorder adj_tree based on adj_igs +""" +function reorder!(adj_tree::IndexAdjacencyTree, adj_igs::Set{IndexGroup}; boundary="right") + @assert boundary in ["left", "right"] + if boundary_state(adj_tree, adj_igs) == "all" + return false + end + adj_trees = topo_sort(adj_tree; type=IndexAdjacencyTree) + ancestors = [tree for tree in adj_trees if contains(tree, adj_igs)] + ancestor_to_state = Dict{IndexAdjacencyTree,String}() + # get the boundary state + for ancestor in ancestors + state = boundary_state(ancestor, adj_igs) + if state == "invalid" + return false + end + ancestor_to_state[ancestor] = state + end + # update ancestors + for ancestor in ancestors + # reorder + if ancestor_to_state[ancestor] == "left" + ancestor.children = reverse(ancestor.children) + elseif ancestor_to_state[ancestor] == "middle" + @assert ancestor.fixed_order == false + filter_children = filter(a -> contains(a, adj_igs), ancestor.children) + reorder_to_right!(ancestor, filter_children) + end + # merge + if ancestor.fixed_order && ancestor.children isa Vector{IndexAdjacencyTree} + new_children = Vector{IndexAdjacencyTree}() + for child in ancestor.children + if !child.fixed_order + push!(new_children, child) + else + push!(new_children, child.children...) + end + end + ancestor.children = new_children + end + end + # check boundary + if boundary == "left" + for ancestor in ancestors + ancestor.children = reverse(ancestor.children) + end + end + return true +end + +# Update both keys and values in igs_to_adjacency_tree based on list_adjacent_igs +function update_igs_to_adjacency_tree!( + list_adjacent_igs::Vector, igs_to_adjacency_tree::Dict{Set{IndexGroup},IndexAdjacencyTree} +) + function update!(root_igs, adjacent_igs) + if !haskey(root_igs_to_adjacent_igs, root_igs) + root_igs_to_adjacent_igs[root_igs] = adjacent_igs + else + val = root_igs_to_adjacent_igs[root_igs] + root_igs_to_adjacent_igs[root_igs] = union(val, adjacent_igs) + end + end + @timeit timer "update_igs_to_adjacency_tree" begin + # get each root igs, get the adjacent igs needed. TODO: do we need to consider boundaries here? + root_igs_to_adjacent_igs = Dict{Set{IndexGroup},Set{IndexGroup}}() + for adjacent_igs in list_adjacent_igs + for root_igs in keys(igs_to_adjacency_tree) + if issubset(adjacent_igs, root_igs) + update!(root_igs, adjacent_igs) + end + end + end + if length(root_igs_to_adjacent_igs) == 1 + return nothing + end + # if at least 3: for now just put everything together + if length(root_igs_to_adjacent_igs) >= 3 + root_igs = keys(root_igs_to_adjacent_igs) + root = union(root_igs...) + igs_to_adjacency_tree[root] = IndexAdjacencyTree( + [igs_to_adjacency_tree[r] for r in root_igs], false, false + ) + for r in root_igs + delete!(igs_to_adjacency_tree, r) + end + return nothing + end + # if 2: assign adjacent_igs to boundary of root_igs (if possible), then concatenate + igs1, igs2 = collect(keys(root_igs_to_adjacent_igs)) + reordered_1 = reorder!( + igs_to_adjacency_tree[igs1], root_igs_to_adjacent_igs[igs1]; boundary="right" + ) + reordered_2 = reorder!( + igs_to_adjacency_tree[igs2], root_igs_to_adjacent_igs[igs2]; boundary="left" + ) + adj_tree_1 = igs_to_adjacency_tree[igs1] + adj_tree_2 = igs_to_adjacency_tree[igs2] + if (!reordered_1) && (!reordered_2) + out_adj_tree = IndexAdjacencyTree([adj_tree_1, adj_tree_2], false, false) + elseif (!reordered_1) + out_adj_tree = IndexAdjacencyTree([adj_tree_1, adj_tree_2.children...], false, true) + elseif (!reordered_2) + out_adj_tree = IndexAdjacencyTree([adj_tree_1.children..., adj_tree_2], false, true) + else + out_adj_tree = IndexAdjacencyTree( + [adj_tree_1.children..., adj_tree_2.children...], false, true + ) + end + root_igs = keys(root_igs_to_adjacent_igs) + root = union(root_igs...) + igs_to_adjacency_tree[root] = out_adj_tree + for r in root_igs + delete!(igs_to_adjacency_tree, r) + end + end +end + +# Generate the adjacency tree of a contraction tree +# Args: +# ========== +# ctree: the input contraction tree +# ancestors: ancestor ctrees of the input ctree +# ctree_to_igs: mapping each ctree to neighboring index groups +function generate_adjacency_tree(ctree, ancestors, ctree_to_igs) + @timeit timer "generate_adjacency_tree" begin + # mapping each index group to adjacent input igs + ig_to_adjacent_igs = Dict{IndexGroup,Set{IndexGroup}}() + # mapping each igs to an adjacency tree + igs_to_adjacency_tree = Dict{Set{IndexGroup},IndexAdjacencyTree}() + for ig in ctree_to_igs[ctree] + ig_to_adjacent_igs[ig] = Set([ig]) + igs_to_adjacency_tree[Set([ig])] = IndexAdjacencyTree(ig) + end + for (i, a) in ancestors + inter_igs = intersect(ctree_to_igs[a[1]], ctree_to_igs[a[2]]) + new_igs_index = (i == 1) ? 2 : 1 + new_igs = setdiff(ctree_to_igs[a[new_igs_index]], inter_igs) + # Tensor product is not considered for now + @assert length(inter_igs) >= 1 + list_adjacent_igs = [ig_to_adjacent_igs[ig] for ig in inter_igs] + update_igs_to_adjacency_tree!(list_adjacent_igs, igs_to_adjacency_tree) + for ig in new_igs + ig_to_adjacent_igs[ig] = union(list_adjacent_igs...) + end + if length(igs_to_adjacency_tree) == 1 + return collect(values(igs_to_adjacency_tree))[1] + end + end + if length(igs_to_adjacency_tree) >= 1 + @info "generate_adjacency_tree has ", length(igs_to_adjacency_tree), "outputs" + return IndexAdjacencyTree([collect(values(igs_to_adjacency_tree))...], false, false) + end + end +end + +function get_ancestors(ctree) + @timeit timer "get_ancestors" begin + ctree_to_ancestors = Dict{Vector,Vector}() + queue = [ctree] + ctree_to_ancestors[ctree] = [] + while queue != [] + node = popfirst!(queue) + if node isa Vector{ITensor} + continue + end + for (i, child) in enumerate(node) + queue = [queue..., child] + ctree_to_ancestors[child] = [(i, node), ctree_to_ancestors[node]...] + end + end + return ctree_to_ancestors + end +end + +# Mutates `v` by sorting elements `x[lo:hi]` using the insertion sort algorithm. +# This method is a copy-paste-edit of sort! in base/sort.jl, amended to return the bubblesort distance. +function _insertion_sort(v::Vector, lo::Int, hi::Int) + @timeit timer "_insertion_sort" begin + v = copy(v) + if lo == hi + return 0 + end + nswaps = 0 + for i in (lo + 1):hi + j = i + x = v[i] + while j > lo + if x < v[j - 1] + nswaps += 1 + v[j] = v[j - 1] + j -= 1 + continue + end + break + end + v[j] = x + end + return nswaps + end +end + +function insertion_sort(v1::Vector, v2::Vector) + value_to_index = Dict{Int,Int}() + for (i, v) in enumerate(v2) + value_to_index[v] = i + end + new_v1 = [value_to_index[v] for v in v1] + return _insertion_sort(new_v1, 1, length(new_v1)) +end + +function minswap_adjacency_tree!(adj_tree::IndexAdjacencyTree) + leaves = Vector{IndexGroup}(get_adj_tree_leaves(adj_tree)) + adj_tree.children = leaves + adj_tree.fixed_order = true + return adj_tree.fixed_direction = true +end + +function minswap_adjacency_tree!( + adj_tree::IndexAdjacencyTree, input_tree::IndexAdjacencyTree +) + nodes = input_tree.children + node_to_int = Dict{IndexGroup,Int}() + int_to_node = Dict{Int,IndexGroup}() + index = 1 + for node in nodes + node_to_int[node] = index + int_to_node[index] = node + index += 1 + end + for node in topo_sort(adj_tree; type=IndexAdjacencyTree) + if node.children isa Vector{IndexGroup} + continue + end + children_tree = [get_adj_tree_leaves(n) for n in node.children] + children_order = vcat(children_tree...) + input_int_order = [node_to_int[n] for n in nodes if n in children_order] + if node.fixed_order + perms = [children_tree, reverse(children_tree)] + else + perms = collect(permutations(children_tree)) + end + nswaps = [] + for perm in perms + int_order = [node_to_int[n] for n in vcat(perm...)] + push!(nswaps, insertion_sort(int_order, input_int_order)) + end + children_tree = perms[argmin(nswaps)] + node.children = vcat(children_tree...) + node.fixed_order = true + node.fixed_direction = true + end + int_order = [node_to_int[n] for n in adj_tree.children] + return _insertion_sort(int_order, 1, length(int_order)) +end + +function split_igs(igs::Vector{IndexGroup}, inter_igs::Vector{IndexGroup}) + igs_left = Vector{IndexGroup}() + igs_right = Vector{IndexGroup}() + target_array = igs_left + for i in igs + if i in inter_igs + target_array = igs_right + continue + end + push!(target_array, i) + end + return igs_left, igs_right +end + +function minswap_adjacency_tree( + adj_tree::IndexAdjacencyTree, + input_tree1::IndexAdjacencyTree, + input_tree2::IndexAdjacencyTree, +) + function merge(l1_left, l1_right, l2_left, l2_right) + if length(l1_left) < length(l2_left) + left_lists = [[l2_left..., l1_left...]] + elseif length(l1_left) > length(l2_left) + left_lists = [[l1_left..., l2_left...]] + else + left_lists = [[l2_left..., l1_left...], [l1_left..., l2_left...]] + end + if length(l1_right) < length(l2_right) + right_lists = [[l1_right..., l2_right...]] + elseif length(l1_right) > length(l2_right) + right_lists = [[l2_right..., l1_right...]] + else + right_lists = [[l2_right..., l1_right...], [l1_right..., l2_right...]] + end + out_lists = [] + for l in left_lists + for r in right_lists + push!(out_lists, IndexAdjacencyTree([l..., r...], true, true)) + end + end + return out_lists + end + @timeit timer "minswap_adjacency_tree" begin + leaves_1 = get_adj_tree_leaves(input_tree1) + leaves_2 = get_adj_tree_leaves(input_tree2) + inter_igs = intersect(leaves_1, leaves_2) + leaves_1_left, leaves_1_right = split_igs(leaves_1, inter_igs) + leaves_2_left, leaves_2_right = split_igs(leaves_2, inter_igs) + num_swaps_1 = + min(length(leaves_1_left), length(leaves_2_left)) + + min(length(leaves_1_right), length(leaves_2_right)) + num_swaps_2 = + min(length(leaves_1_left), length(leaves_2_right)) + + min(length(leaves_1_right), length(leaves_2_left)) + if num_swaps_1 == num_swaps_2 + inputs_1 = merge(leaves_1_left, leaves_1_right, leaves_2_left, leaves_2_right) + inputs_2 = merge( + leaves_1_left, leaves_1_right, reverse(leaves_2_right), reverse(leaves_2_left) + ) + inputs = [inputs_1..., inputs_2...] + elseif num_swaps_1 > num_swaps_2 + inputs = merge( + leaves_1_left, leaves_1_right, reverse(leaves_2_right), reverse(leaves_2_left) + ) + else + inputs = merge(leaves_1_left, leaves_1_right, leaves_2_left, leaves_2_right) + end + # TODO: may want to change this back + # leaves_1 = [i for i in leaves_1 if !(i in inter_igs)] + # leaves_2 = [i for i in leaves_2 if !(i in inter_igs)] + # input1 = IndexAdjacencyTree([leaves_1..., leaves_2...], true, true) + # input2 = IndexAdjacencyTree([leaves_1..., reverse(leaves_2)...], true, true) + # input3 = IndexAdjacencyTree([reverse(leaves_1)..., leaves_2...], true, true) + # input4 = IndexAdjacencyTree([reverse(leaves_1)..., reverse(leaves_2)...], true, true) + # inputs = [input1, input2, input3, input4] + # ====================================== + adj_tree_copies = [copy(adj_tree) for _ in 1:length(inputs)] + nswaps = [minswap_adjacency_tree!(t, i) for (t, i) in zip(adj_tree_copies, inputs)] + return adj_tree_copies[argmin(nswaps)] + end +end + +function _approximate_contract_pre_process(tn_leaves, ctrees) + @timeit timer "_approximate_contract_pre_process" begin + # mapping each contraction tree to its uncontracted index groups + ctree_to_igs = Dict{Vector,Vector{IndexGroup}}() + index_groups = get_index_groups(ctrees[end]) + for c in vcat(tn_leaves, ctrees) + ctree_to_igs[c] = neighbor_index_groups(c, index_groups) + end + ctree_to_ancestors = get_ancestors(ctrees[end]) + # mapping each contraction tree to its index adjacency tree + ctree_to_adj_tree = Dict{Vector,IndexAdjacencyTree}() + for leaf in tn_leaves + ctree_to_adj_tree[leaf] = generate_adjacency_tree( + leaf, ctree_to_ancestors[leaf], ctree_to_igs + ) + minswap_adjacency_tree!(ctree_to_adj_tree[leaf]) + end + for c in ctrees + ancestors = ctree_to_ancestors[c] + adj_tree = generate_adjacency_tree(c, ancestors, ctree_to_igs) + if adj_tree != nothing + ctree_to_adj_tree[c] = minswap_adjacency_tree( + adj_tree, ctree_to_adj_tree[c[1]], ctree_to_adj_tree[c[2]] + ) + end + end + # mapping each contraction tree to its contract igs + ctree_to_contract_igs = Dict{Vector,Vector{IndexGroup}}() + for c in ctrees + contract_igs = intersect(ctree_to_igs[c[1]], ctree_to_igs[c[2]]) + ctree_to_contract_igs[c[1]] = contract_igs + ctree_to_contract_igs[c[2]] = contract_igs + end + # special case when the network contains uncontracted inds + ctree_to_contract_igs[ctrees[end]] = ctree_to_igs[ctrees[end]] + # mapping each index group to a linear ordering + ig_to_linear_order = Dict{IndexGroup,Vector}() + for leaf in tn_leaves + for ig in ctree_to_igs[leaf] + if !haskey(ig_to_linear_order, ig) + ig_to_linear_order[ig] = inds_linear_order(leaf, ig.data) + end + end + end + return ctree_to_igs, ctree_to_adj_tree, ctree_to_contract_igs, ig_to_linear_order + end +end + +function ordered_igs_to_binary_tree(ordered_igs, contract_igs, ig_to_linear_order; ansatz) + @assert ansatz in ["comb", "mps"] + @timeit timer "ordered_igs_to_binary_tree" begin + @assert contract_igs != [] + left_igs, right_igs = split_igs(ordered_igs, contract_igs) + if ansatz == "comb" + return ordered_igs_to_binary_tree_comb( + left_igs, right_igs, contract_igs, ig_to_linear_order + ) + elseif ansatz == "mps" + return ordered_igs_to_binary_tree_mps( + left_igs, right_igs, contract_igs, ig_to_linear_order + ) + end + end +end + +function ordered_igs_to_binary_tree_mps( + left_igs, right_igs, contract_igs, ig_to_linear_order +) + left_order = get_leaves([ig_to_linear_order[ig] for ig in left_igs]) + right_order = get_leaves([ig_to_linear_order[ig] for ig in right_igs]) + contract_order = get_leaves([ig_to_linear_order[ig] for ig in contract_igs]) + if length(left_order) <= length(right_order) + left_order = [left_order..., contract_order...] + else + right_order = [contract_order..., right_order...] + end + return merge_tree(line_to_tree(left_order), line_to_tree(reverse(right_order))) +end + +function ordered_igs_to_binary_tree_comb( + left_igs, right_igs, contract_igs, ig_to_linear_order +) + tree_1 = line_to_tree([line_to_tree(ig_to_linear_order[ig]) for ig in left_igs]) + tree_contract = line_to_tree([ + line_to_tree(ig_to_linear_order[ig]) for ig in contract_igs + ]) + tree_2 = line_to_tree([line_to_tree(ig_to_linear_order[ig]) for ig in reverse(right_igs)]) + # make the binary tree more balanced to save tree approximation cost + if tree_1 == [] + return merge_tree(merge_tree(tree_1, tree_contract), tree_2) + end + if tree_2 == [] + return merge_tree(tree_1, merge_tree(tree_contract, tree_2)) + end + if length(vectorize(tree_1)) <= length(vectorize(tree_2)) + return merge_tree(merge_tree(tree_1, tree_contract), tree_2) + else + return merge_tree(tree_1, merge_tree(tree_contract, tree_2)) + end +end + +function get_igs_cache_info(igs_list, contract_igs_list) + function split_boundary(list1::Vector{IndexGroup}, list2::Vector{IndexGroup}) + index = 1 + boundary = Vector{IndexGroup}() + while list1[index] == list2[index] + push!(boundary, list2[index]) + index += 1 + if index > length(list1) || index > length(list2) + break + end + end + if index <= length(list1) + remain_list1 = list1[index:end] + else + remain_list1 = Vector{IndexGroup}() + end + return boundary, remain_list1 + end + function split_boundary(igs::Vector{IndexGroup}, lists::Vector{Vector{IndexGroup}}) + if length(igs) <= 1 + return Vector{IndexGroup}(), igs + end + for l in lists + if length(l) >= 2 && igs[1] == l[1] && igs[2] == l[2] + return split_boundary(igs, l) + end + end + return Vector{IndexGroup}(), igs + end + @timeit timer "get_igs_cache_info" begin + out, input1, input2 = igs_list + contract_out, contract_input1, contract_input2 = contract_igs_list + out_left, out_right = split_igs(out, contract_out) + out_right = reverse(out_right) + input1_left, input1_right = split_igs(input1, contract_input1) + input2_left, input2_right = split_igs(input2, contract_input2) + inputs = [input1_left, reverse(input1_right), input2_left, reverse(input2_right)] + boundary_left, remain_left = split_boundary(out_left, inputs) + boundary_right, remain_right = split_boundary(out_right, inputs) + return [remain_left..., contract_out..., reverse(remain_right)...], + boundary_left, + boundary_right + end +end + +function get_tn_cache_sub_info( + tn_tree::Dict{Vector,OrthogonalITensor}, cache_binary_trees::Vector +) + cached_tn = [] + cached_tn_tree = Dict{Vector,OrthogonalITensor}() + new_igs = [] + for binary_tree in cache_binary_trees + if binary_tree == [] || !haskey(tn_tree, binary_tree) + push!(new_igs, nothing) + else + binary_tree = Vector{Vector}(binary_tree) + nodes = topo_sort(binary_tree; type=Vector{<:Vector}) + sub_tn = [tn_tree[n] for n in nodes] + sub_tn_tree = Dict([n => tn_tree[n] for n in nodes]...) + index_leaves = vectorize(binary_tree) + new_indices = setdiff(noncommoninds(sub_tn...), index_leaves) + @assert length(new_indices) == 1 + new_indices = Vector{<:Index}(new_indices) + push!(new_igs, IndexGroup(new_indices)) + cached_tn = vcat(cached_tn, sub_tn) + cached_tn_tree = merge(cached_tn_tree, sub_tn_tree) + end + end + tn = vcat(collect(values(tn_tree))...) + uncached_tn = setdiff(tn, cached_tn) + return cached_tn_tree, uncached_tn, new_igs +end + +function get_tn_cache_info( + ctree_to_tn_tree, ctree_1::Vector, ctree_2::Vector, cache_binary_trees::Vector +) + @timeit timer "get_tn_cache_info" begin + if haskey(ctree_to_tn_tree, ctree_1) && ctree_to_tn_tree[ctree_1] isa Dict + tn_tree_1 = ctree_to_tn_tree[ctree_1] + cached_tn_tree1, uncached_tn1, new_igs_1 = get_tn_cache_sub_info( + tn_tree_1, cache_binary_trees + ) + else + cached_tn_tree1 = Dict{Vector,OrthogonalITensor}() + uncached_tn1 = get_child_tn(ctree_to_tn_tree, ctree_1) + new_igs_1 = [nothing, nothing] + end + if haskey(ctree_to_tn_tree, ctree_2) && ctree_to_tn_tree[ctree_2] isa Dict + tn_tree_2 = ctree_to_tn_tree[ctree_2] + cached_tn_tree2, uncached_tn2, new_igs_2 = get_tn_cache_sub_info( + tn_tree_2, cache_binary_trees + ) + else + cached_tn_tree2 = Dict{Vector,OrthogonalITensor}() + uncached_tn2 = get_child_tn(ctree_to_tn_tree, ctree_2) + new_igs_2 = [nothing, nothing] + end + uncached_tn = [uncached_tn1..., uncached_tn2...] + new_igs_left = [i for i in [new_igs_1[1], new_igs_2[1]] if i != nothing] + @assert length(new_igs_left) <= 1 + if length(new_igs_left) == 1 + new_ig_left = new_igs_left[1] + else + new_ig_left = nothing + end + new_igs_right = [i for i in [new_igs_1[2], new_igs_2[2]] if i != nothing] + @assert length(new_igs_right) <= 1 + if length(new_igs_right) == 1 + new_ig_right = new_igs_right[1] + else + new_ig_right = nothing + end + return merge(cached_tn_tree1, cached_tn_tree2), uncached_tn, new_ig_left, new_ig_right + end +end + +function update_tn_tree_keys!(tn_tree, inds_btree, pairs::Vector{Pair}) + @timeit timer "update_tn_tree_keys!" begin + current_to_update_key = Dict{Vector,Vector}(pairs...) + nodes = topo_sort(inds_btree; type=Vector{<:Vector}) + for n in nodes + @assert haskey(tn_tree, n) + new_key = n + if haskey(current_to_update_key, n[1]) + new_key = [current_to_update_key[n[1]], n[2]] + end + if haskey(current_to_update_key, n[2]) + new_key = [new_key[1], current_to_update_key[n[2]]] + end + if new_key != n + tn_tree[new_key] = tn_tree[n] + delete!(tn_tree, n) + current_to_update_key[n] = new_key + end + end + end +end + +function get_child_tn(ctree_to_tn_tree, ctree::Vector) + if !haskey(ctree_to_tn_tree, ctree) + @assert ctree isa Vector{ITensor} + return orthogonal_tensors(ctree) + elseif ctree_to_tn_tree[ctree] isa Vector{OrthogonalITensor} + return ctree_to_tn_tree[ctree] + else + return vcat(collect(values(ctree_to_tn_tree[ctree]))...) + end +end + +_index_less(a::Index, b::Index) = tags(a)[1] < tags(b)[1] + +function _replaceinds(t1::ITensor, t2::ITensor) + inds1 = sort(inds(t1); lt=_index_less) + inds2 = sort(inds(t2); lt=_index_less) + inds1_tags = [tags(i) for i in inds1] + inds2_tags = [tags(i) for i in inds2] + @assert inds1_tags == inds2_tags + return replaceinds(t1, inds1, inds2) +end + +function orthogonalize!(ctree_to_tn_tree::Dict, environments::Vector, c::Vector) + @timeit timer "orthogonalize" begin + index = [] + if c[1] in environments + push!(index, 1) + elseif c[2] in environments + push!(index, 2) + end + if index == [] + return nothing + end + environments = setdiff(environments, c) + if length(environments) == 0 + return nothing + end + @info "start orthogonalize with env size", length(environments) + network = vcat([get_child_tn(ctree_to_tn_tree, env) for env in environments]...) + network = get_tensors(network) + env_boundary = get_child_tn(ctree_to_tn_tree, c[index[1]]) + source_tensor = env_boundary[1].tensor + @assert !(source_tensor in network) + push!(network, source_tensor) + ctree_to_tn_tree[c[index[1]]] = env_boundary[2:end] + orth_tn = orthogonalize(ITensorNetwork(network), length(network)) + tensor_to_ortho_tensor = Dict{ITensor,ITensor}() + for i in 1:length(network) + new_tensor = _replaceinds(orth_tn[i], network[i]) + tensor_to_ortho_tensor[network[i]] = new_tensor + end + for env in environments + ortho_tensors = Vector{OrthogonalITensor}([ + OrthogonalITensor(tensor_to_ortho_tensor[t.tensor]) for + t in get_child_tn(ctree_to_tn_tree, env) + ]) + ctree_to_tn_tree[env] = ortho_tensors + end + push!( + ctree_to_tn_tree[c[index[1]]], + OrthogonalITensor(tensor_to_ortho_tensor[source_tensor]), + ) + end +end + +# ctree: contraction tree +# tn: vector of tensors representing a tensor network +# tn_tree: a dict maps each index tree in the tn to a tensor +# adj_tree: index adjacency tree +# ig: index group +# contract_ig: the index group to be contracted next +# ig_tree: an index group with a tree hierarchy +function approximate_contract( + ctree::Vector; + cutoff, + maxdim, + ansatz="mps", + use_cache=true, + orthogonalize=false, + algorithm="density_matrix", +) + @timeit timer "approximate_contract" begin + tn_leaves = get_leaves(ctree) + environments = tn_leaves + ctrees = topo_sort(ctree; leaves=tn_leaves) + ctree_to_igs, ctree_to_adj_tree, ctree_to_contract_igs, ig_to_linear_order = _approximate_contract_pre_process( + tn_leaves, ctrees + ) + # mapping each contraction tree to a tensor network + ctree_to_tn_tree = Dict{ + Vector,Union{Dict{Vector,OrthogonalITensor},Vector{OrthogonalITensor}} + }() + # accumulate norm + log_accumulated_norm = 0.0 + for (ii, c) in enumerate(ctrees) + @info "orthogonalize", orthogonalize + if orthogonalize == true + orthogonalize!(ctree_to_tn_tree, environments, c) + environments = setdiff(environments, c) + end + @info ii, "th tree approximation" + if ctree_to_igs[c] == [] + @assert c == ctrees[end] + tn1 = get_child_tn(ctree_to_tn_tree, c[1]) + tn2 = get_child_tn(ctree_to_tn_tree, c[2]) + tn = vcat(tn1, tn2) + return get_tensors([optcontract(tn)]), log_accumulated_norm + end + # caching is not used here + if use_cache == false + tn1 = get_child_tn(ctree_to_tn_tree, c[1]) + tn2 = get_child_tn(ctree_to_tn_tree, c[2]) + inds_btree = ordered_igs_to_binary_tree( + ctree_to_adj_tree[c].children, + ctree_to_contract_igs[c], + ig_to_linear_order; + ansatz=ansatz, + ) + ctree_to_tn_tree[c], log_root_norm = approximate_contract_ctree_to_tensor( + [tn1..., tn2...], inds_btree; cutoff=cutoff, maxdim=maxdim, algorithm=algorithm + ) + log_accumulated_norm += log_root_norm + continue + end + # caching + # Note: cache_igs_right has a reversed ordering + center_igs, cache_igs_left, cache_igs_right = get_igs_cache_info( + [ctree_to_adj_tree[i].children for i in [c, c[1], c[2]]], + [ctree_to_contract_igs[i] for i in [c, c[1], c[2]]], + ) + if ansatz == "comb" + cache_binary_tree_left = line_to_tree([ + line_to_tree(ig_to_linear_order[ig]) for ig in cache_igs_left + ]) + cache_binary_tree_right = line_to_tree([ + line_to_tree(ig_to_linear_order[ig]) for ig in cache_igs_right + ]) + elseif ansatz == "mps" + left_order = vcat([ig_to_linear_order[ig] for ig in cache_igs_left]...) + cache_binary_tree_left = line_to_tree(left_order) + right_order = vcat([ig_to_linear_order[ig] for ig in reverse(cache_igs_right)]...) + cache_binary_tree_right = line_to_tree(reverse(right_order)) + end + cached_tn_tree, uncached_tn, new_ig_left, new_ig_right = get_tn_cache_info( + ctree_to_tn_tree, c[1], c[2], [cache_binary_tree_left, cache_binary_tree_right] + ) + if new_ig_right == nothing && new_ig_left == nothing + @info "Caching is not used in this approximation" + @assert length(cached_tn_tree) == 0 + else + @info "Caching is used in this approximation", new_ig_left, new_ig_right + end + new_ig_to_binary_tree_pairs = Vector{Pair}() + new_igs = center_igs + new_ig_to_linear_order = ig_to_linear_order + if new_ig_left == nothing + new_igs = [cache_igs_left..., new_igs...] + else + new_ig_to_linear_order = merge( + new_ig_to_linear_order, Dict(new_ig_left => [new_ig_left.data]) + ) + new_igs = [new_ig_left, new_igs...] + push!(new_ig_to_binary_tree_pairs, new_ig_left.data => cache_binary_tree_left) + end + if new_ig_right == nothing + new_igs = [new_igs..., cache_igs_right...] + else + new_ig_to_linear_order = merge( + new_ig_to_linear_order, Dict(new_ig_right => [new_ig_right.data]) + ) + new_igs = [new_igs..., new_ig_right] + push!(new_ig_to_binary_tree_pairs, new_ig_right.data => cache_binary_tree_right) + end + inds_btree = ordered_igs_to_binary_tree( + new_igs, ctree_to_contract_igs[c], new_ig_to_linear_order; ansatz=ansatz + ) + new_tn_tree, log_root_norm = approximate_contract_ctree_to_tensor( + uncached_tn, inds_btree; cutoff=cutoff, maxdim=maxdim, algorithm=algorithm + ) + log_accumulated_norm += log_root_norm + if length(new_ig_to_binary_tree_pairs) != 0 + update_tn_tree_keys!(new_tn_tree, inds_btree, new_ig_to_binary_tree_pairs) + end + ctree_to_tn_tree[c] = merge(new_tn_tree, cached_tn_tree) + # release the memory + delete!(ctree_to_tn_tree, c[1]) + delete!(ctree_to_tn_tree, c[2]) + end + tn = vcat(collect(values(ctree_to_tn_tree[ctrees[end]]))...) + return get_tensors(tn), log_accumulated_norm + end +end + +function tree_approximation( + embedding::Dict, + inds_btree::Vector; + cutoff=1e-15, + maxdim=10000, + maxsize=10000, + algorithm="density_matrix", +) + @assert algorithm in ["density_matrix", "density_matrix_contract_first", "svd"] + if algorithm == "density_matrix" + return tree_approximation_density_matrix( + embedding, inds_btree; cutoff=cutoff, maxdim=maxdim, maxsize=maxsize + ) + end + if algorithm == "density_matrix_contract_first" + @info "density_matrix_contract_first" + btree_to_contracted_tn = Dict{Vector,Vector{OrthogonalITensor}}() + for (btree, ts) in embedding + btree_to_contracted_tn[btree] = [optcontract(ts)] + end + return tree_approximation_density_matrix( + btree_to_contracted_tn, inds_btree; cutoff=cutoff, maxdim=maxdim, maxsize=maxsize + ) + end + if algorithm == "svd" + return tree_approximation_svd( + embedding, inds_btree; cutoff=cutoff, maxdim=maxdim, maxsize=maxsize + ) + end +end + +function tree_approximation_svd( + embedding::Dict, inds_btree::Vector; cutoff=1e-15, maxdim=10000, maxsize=10000 +) + @info "start tree_approximation_svd", inds_btree + @info "cutoff", cutoff, "maxdim", maxdim + network = Vector{ITensor}() + btree_to_order = Dict{Vector,Int}() + root_vertex = nothing + for (btree, ts) in embedding + # use dense to convert Diag type to dense for QR decomposition TODO: raise an error in ITensors + push!(network, dense(optcontract(ts).tensor)) + btree_to_order[btree] = length(network) + if btree == inds_btree + root_vertex = length(network) + end + end + @assert root_vertex != nothing + ttn = TTN(ITensorNetwork(network)) + @timeit timer "truncate" begin + truncate_ttn = truncate(ttn; cutoff=cutoff, maxdim=maxdim, root_vertex=root_vertex) + end + out_network = [truncate_ttn[i] for i in 1:length(network)] + inds1 = sort(uncontractinds(out_network); lt=_index_less) + inds2 = sort(uncontractinds(network); lt=_index_less) + out_network = replaceinds(out_network, Dict(zip(inds1, inds2))) + root_norm = norm(out_network[root_vertex]) + out_network[root_vertex] /= root_norm + ctree_to_tensor = Dict{Vector,OrthogonalITensor}() + for node in topo_sort(inds_btree; type=Vector{<:Vector}) + children_tensors = [] + if !(node[1] isa Vector{<:Vector}) + push!(children_tensors, OrthogonalITensor(out_network[btree_to_order[node[1]]])) + end + if !(node[2] isa Vector{<:Vector}) + push!(children_tensors, OrthogonalITensor(out_network[btree_to_order[node[2]]])) + end + t = OrthogonalITensor(out_network[btree_to_order[node]]) + if children_tensors == [] + ctree_to_tensor[node] = t + else + ctree_to_tensor[node] = optcontract([t, children_tensors...]) + end + end + return ctree_to_tensor, log(root_norm) +end + +function _randomized_svd(t::ITensor, linds, rinds; maxdim) + cutoff = 1e-16 + @info "randomized svd", "cutoff", cutoff, "maxdim", maxdim + rand_dim = maxdim + 20 + new_index = Index(Integer(rand_dim), "rand_dim") + rand_mat = randomITensor(rinds..., new_index) + rand_t = contract(t, rand_mat) + Q, _ = factorize(rand_t, linds...; which_decomp="qr", ortho="left") + rand_t = contract(t, Q) + Q, _ = factorize(rand_t, rinds...; which_decomp="qr", ortho="left") + rand_t = contract(t, Q) + Q, _ = factorize(rand_t, linds...; which_decomp="qr", ortho="left") + rand_mat = contract(t, Q) + U, diag, _ = svd(rand_mat, rinds...; maxdim=maxdim, cutoff=cutoff, alg="qr_iteration") + return diag, U +end + +# interlaced HOSVD using density matrix and caching +function tree_approximation_density_matrix( + embedding::Dict, inds_btree::Vector; cutoff=1e-15, maxdim=10000, maxsize=10000 +) + @info "start tree_approximation_density matrix", inds_btree + ctree_to_tensor = Dict{Vector,OrthogonalITensor}() + # initialize sim_dict + network = vcat(collect(values(embedding))...) + uncontractinds = noncommoninds(network...) + innerinds = mapreduce(t -> [i for i in inds(t)], vcat, network) + innerinds = Vector(setdiff(innerinds, uncontractinds)) + siminner_dict = Dict([ind => sim(ind) for ind in innerinds]) + + function closednet(tree) + netbra = embedding[tree] + netket = replaceinds(netbra, siminner_dict) + if length(tree) == 1 + return optcontract(vcat(netbra, netket)) + end + tleft, tright = closednet(tree[1]), closednet(tree[2]) + return optcontract(vcat(netbra, netket, [tleft], [tright])) + end + + function insert_projectors(tree::Vector, env::OrthogonalITensor) + netbra = embedding[tree] + netket = replaceinds(netbra, siminner_dict) + if length(tree) == 1 + tensor_bra = optcontract(netbra) + tensor_ket = replaceinds([tensor_bra], siminner_dict)[1] + inds_pair = (tree[1], sim(tree[1])) + tensor_ket = replaceinds([tensor_ket], Dict([inds_pair[1] => inds_pair[2]]))[1] + return inds_pair, optcontract([netbra..., netket...]), [tensor_bra, tensor_ket] + end + # update children + subenvtensor = optcontract([env, netbra...]) + envnet = [subenvtensor, closednet(tree[2]), netket...] + ind1_pair, subnetsq1, subnet1 = insert_projectors(tree[1], optcontract(envnet)) + envnet = [subenvtensor, subnetsq1, netket...] + ind2_pair, _, subnet2 = insert_projectors(tree[2], optcontract(envnet)) + # compute the projector + rinds = (ind1_pair[1], ind2_pair[1]) + linds = (ind1_pair[2], ind2_pair[2]) + # to handle the corner cases where subnet1/subnet2 could be empty + netket = replaceinds( + netket, Dict([ind1_pair[1] => ind1_pair[2], ind2_pair[1] => ind2_pair[2]]) + ) + net = [subenvtensor, netket..., subnet1..., subnet2...] + tnormal = optcontract(net) + dim2 = floor(maxsize / (space(ind1_pair[1]) * space(ind2_pair[1]))) + max_dim = min(maxdim, dim2) + t00 = time() + @info "eigen input size", size(tnormal.tensor) + @timeit timer "eigen" begin + left_size = prod([dim(a) for a in linds]) + right_size = prod([dim(a) for a in rinds]) + if left_size < max_dim * 5 || right_size < max_dim * 5 + diag, U = eigen( + tnormal.tensor, linds, rinds; cutoff=cutoff, maxdim=max_dim, ishermitian=true + ) + else + diag, U = _randomized_svd(tnormal.tensor, linds, rinds; maxdim=max_dim) + end + end + t11 = time() - t00 + @info "size of U", size(U), "size of diag", size(diag), "costs", t11 + dr = commonind(diag, U) + Usim = replaceinds(U, rinds => linds) + ortho_U = OrthogonalITensor(U) + ortho_Usim = OrthogonalITensor(Usim) + net1 = [netbra..., subnet1[1], subnet2[1], ortho_U] + net2 = [netket..., subnet1[2], subnet2[2], ortho_Usim] + tensor1 = optcontract(net1) + tensor2 = replaceinds(tensor1, noncommoninds(net1...) => noncommoninds(net2...)) + subnetsq = optcontract([tensor1, tensor2]) + dr_pair = (dr, sim(dr)) + tensor2 = replaceinds(tensor2, [dr_pair[1]] => [dr_pair[2]]) + subnet = [tensor1, tensor2] + ctree_to_tensor[tree] = ortho_U + return dr_pair, subnetsq, subnet + end + + @assert (length(inds_btree) >= 2) + bra = embedding[inds_btree] + ket = replaceinds(bra, siminner_dict) + # update children + envnet = [closednet(inds_btree[2]), bra..., ket...] + _, netsq1, n1 = insert_projectors(inds_btree[1], optcontract(envnet)) + envnet = [netsq1, bra..., ket...] + _, _, n2 = insert_projectors(inds_btree[2], optcontract(envnet)) + # last tensor + envnet = [n1[1], n2[1], bra...] + last_tensor = optcontract(envnet) + root_norm = norm(last_tensor.tensor) + last_tensor.tensor /= root_norm + ctree_to_tensor[inds_btree] = last_tensor + return ctree_to_tensor, log(root_norm) +end diff --git a/src/ApproximateTNContraction/contract/index_group.jl b/src/ApproximateTNContraction/contract/index_group.jl new file mode 100644 index 00000000..db48ec19 --- /dev/null +++ b/src/ApproximateTNContraction/contract/index_group.jl @@ -0,0 +1,46 @@ +mutable struct IndexGroup + data::Vector + istree::Bool +end + +# TODO: general tags are not comparable +Base.isless(a::Index, b::Index) = id(a) < id(b) || (id(a) == id(b) && plev(a) < plev(b)) # && tags(a) < tags(b) + +function IndexGroup(indices::Vector{<:Index}) + return IndexGroup(sort(indices), false) +end + +function get_index_groups(tn_tree::Vector) + @timeit timer "get_index_groups" begin + tn_leaves = get_leaves(tn_tree) + tn = vcat(tn_leaves...) + uncontract_inds = noncommoninds(tn...) + igs = [] + for leaf in tn_leaves + inds = intersect(noncommoninds(leaf...), uncontract_inds) + if length(inds) >= 1 + push!(igs, IndexGroup(inds)) + end + end + for (t1, t2) in powerset(tn_leaves, 2, 2) + inds = intersect(noncommoninds(t1...), noncommoninds(t2...)) + if length(inds) >= 1 + push!(igs, IndexGroup(inds)) + end + end + return igs + end +end + +function neighbor_index_groups(contraction, index_groups) + @timeit timer "get_index_groups" begin + inds = noncommoninds(vectorize(contraction)...) + nigs = [] + for ig in index_groups + if issubset(ig.data, inds) + push!(nigs, ig) + end + end + return nigs + end +end diff --git a/src/ApproximateTNContraction/contract/mincut_tree.jl b/src/ApproximateTNContraction/contract/mincut_tree.jl new file mode 100644 index 00000000..6d51bd42 --- /dev/null +++ b/src/ApproximateTNContraction/contract/mincut_tree.jl @@ -0,0 +1,185 @@ +using ITensorNetworks: ITensorNetwork + +# a large number to prevent this edge being a cut +MAX_WEIGHT = 1e32 + +function _build_pseudo_network( + network::Vector{ITensor}, outinds::Union{Nothing,Vector{<:Index}} +) + # create a pseudo_tn (a tn without any data inside) whose outinds weights are MAX_WEIGHT + out_to_pseudo_ind = Dict{Index,Index}() + for ind in outinds + out_to_pseudo_ind[ind] = Index(MAX_WEIGHT, ind.tags) + end + pseudo_network = Vector{ITensor}() + for t in network + inds1 = [i for i in inds(t) if !(i in outinds)] + inds2 = [out_to_pseudo_ind[i] for i in inds(t) if i in outinds] + newt = ITensor(inds1..., inds2...) + push!(pseudo_network, newt) + end + return pseudo_network, out_to_pseudo_ind +end + +function inds_binary_tree( + network::Vector{ITensor}, outinds::Union{Nothing,Vector{<:Index}}; algorithm="mincut" +) + if outinds == nothing + outinds = noncommoninds(network...) + end + if length(outinds) == 1 + return outinds + end + p_network, out_to_pseudo_ind = _build_pseudo_network(network, outinds) + return _inds_binary_tree(network => p_network, out_to_pseudo_ind; algorithm=algorithm) +end + +function _inds_binary_tree( + network_pair::Pair{Vector{ITensor},Vector{ITensor}}, + out_to_pseudo_ind::Dict{Index,Index}; + algorithm="mincut", +) + @assert algorithm in ["mincut", "mps"] + if algorithm == "mincut" + return _mincut_inds(network_pair, out_to_pseudo_ind) + elseif algorithm == "mps" + return line_to_tree(_inds_linear_order(network_pair, out_to_pseudo_ind)) + end +end + +function inds_linear_order( + network::Vector{ITensor}, outinds::Union{Nothing,Vector{<:Index}} +) + if outinds == nothing + outinds = noncommoninds(network...) + end + if length(outinds) == 1 + return outinds + end + p_network, out_to_pseudo_ind = _build_pseudo_network(network, outinds) + return _inds_linear_order(network => p_network, out_to_pseudo_ind) +end + +function _inds_linear_order( + network_pair::Pair{Vector{ITensor},Vector{ITensor}}, out_to_pseudo_ind::Dict{Index,Index} +) + outinds = collect(keys(out_to_pseudo_ind)) + outinds = [[i] for i in outinds] + @assert length(outinds) >= 1 + if length(outinds) <= 2 + return outinds + end + new_edge = new_edge_mincut( + network_pair, out_to_pseudo_ind, collect(powerset(outinds, 1, 1)) + ) + new_edge = new_edge[1] + linear_order = [new_edge] + while length(outinds) > 2 + splitinds = [[new_edge, i] for i in outinds if i != new_edge] + new_edge = new_edge_mincut(network_pair, out_to_pseudo_ind, splitinds) + push!(linear_order, new_edge[2]) + outinds = setdiff(outinds, new_edge) + outinds = vcat([new_edge], outinds) + end + last_index = [i for i in outinds if i != new_edge] + @assert length(last_index) == 1 + push!(linear_order, last_index[1]) + return linear_order +end + +function _mincut_inds( + network_pair::Pair{Vector{ITensor},Vector{ITensor}}, out_to_pseudo_ind::Dict{Index,Index} +) + outinds = collect(keys(out_to_pseudo_ind)) + outinds = [[i] for i in outinds] + return __mincut_inds(network_pair, out_to_pseudo_ind, outinds) +end + +function __mincut_inds( + network_pair::Pair{Vector{ITensor},Vector{ITensor}}, + out_to_pseudo_ind::Dict{Index,Index}, + outinds::Vector{<:Vector}, +) + @assert length(outinds) >= 1 + if length(outinds) == 1 + return outinds[1] + end + if length(outinds) == 2 + return outinds + end + new_edge = new_edge_mincut( + network_pair, out_to_pseudo_ind, collect(powerset(outinds, 2, 2)) + ) + outinds = setdiff(outinds, new_edge) + outinds = vcat([new_edge], outinds) + return __mincut_inds(network_pair, out_to_pseudo_ind, outinds) +end + +function new_edge_mincut( + network_pair::Pair{Vector{ITensor},Vector{ITensor}}, + out_to_pseudo_ind::Dict{Index,Index}, + split_inds_list::Vector, +) + outinds = collect(keys(out_to_pseudo_ind)) + p_outinds = collect(values(out_to_pseudo_ind)) + mincuts, pseudo_mincuts, dists = [], [], [] + for split_inds in split_inds_list + if length(split_inds) == 2 + push!( + dists, + distance(network_pair.first, vectorize(split_inds[1]), vectorize(split_inds[2])), + ) + else + push!(dists, 0.0) + end + split_inds = vectorize(split_inds) + p_split_inds = [out_to_pseudo_ind[i] for i in split_inds] + push!(mincuts, mincut_value(network_pair.first, split_inds, outinds)[3]) + push!(pseudo_mincuts, mincut_value(network_pair.second, p_split_inds, p_outinds)[3]) + end + indices_min = [i for i in 1:length(mincuts) if mincuts[i] == min(mincuts...)] + selected_pseudo_mincuts = [pseudo_mincuts[i] for i in indices_min] + indices_min = [ + i for i in indices_min if pseudo_mincuts[i] == min(selected_pseudo_mincuts...) + ] + dists_min = [dists[i] for i in indices_min] + _, index = findmin(dists_min) + i = indices_min[index] + new_edge = split_inds_list[i] + return new_edge +end + +function distance(network::Vector{ITensor}, inds1::Vector{<:Index}, inds2::Vector{<:Index}) + new_t1 = ITensor(inds1...) + new_t2 = ITensor(inds2...) + tn = ITensorNetwork([network..., new_t1, new_t2]) + ds = dijkstra_shortest_paths(tn, length(network) + 1, weights(tn)) + return ds.dists[length(network) + 2] +end + +function mincut_value( + network::Vector{ITensor}, source_inds::Vector{<:Index}, out_inds::Vector{<:Index} +) + terminal_inds = setdiff(out_inds, source_inds) + new_t1 = ITensor(source_inds...) + new_t2 = ITensor(terminal_inds...) + tn = ITensorNetwork([network..., new_t1, new_t2]) + return GraphsFlows.mincut(tn, length(network) + 1, length(network) + 2, weights(tn)) +end + +function mincut_subnetwork( + network::Vector{ITensor}, source_inds::Vector{<:Index}, out_inds::Vector{<:Index} +) + @timeit timer "mincut_subnetwork" begin + if length(source_inds) == length(out_inds) + return network + end + p_network, out_to_pseudo_ind = _build_pseudo_network(network, out_inds) + p_source_inds = [out_to_pseudo_ind[i] for i in source_inds] + p_out_inds = [out_to_pseudo_ind[i] for i in out_inds] + part1, part2, val = mincut_value(p_network, p_source_inds, p_out_inds) + @assert length(part1) > 1 + @assert length(part2) > 1 + return [network[i] for i in part1 if i <= length(network)] + end +end diff --git a/src/ApproximateTNContraction/contract/tree_embedding.jl b/src/ApproximateTNContraction/contract/tree_embedding.jl new file mode 100644 index 00000000..db0ba6d5 --- /dev/null +++ b/src/ApproximateTNContraction/contract/tree_embedding.jl @@ -0,0 +1,147 @@ +include("union_find.jl") + +function tree_embedding(network::Vector{OrthogonalITensor}, inds_btree::Vector; algorithm) + # TODO: consider identity matrices + tensor_to_ortho_tensor = Dict{ITensor,OrthogonalITensor}() + for ortho_tensor in network + tensor_to_ortho_tensor[ortho_tensor.tensor] = ortho_tensor + end + tnets_dict = tree_embedding(get_tensors(network), inds_btree; algorithm=algorithm) + ortho_tnets_dict = Dict() + for (key, tensors) in tnets_dict + ortho_tensors = Vector{OrthogonalITensor}() + for t in tensors + if haskey(tensor_to_ortho_tensor, t) + push!(ortho_tensors, tensor_to_ortho_tensor[t]) + else + push!(ortho_tensors, OrthogonalITensor(t)) + end + end + ortho_tnets_dict[key] = ortho_tensors + end + return ortho_tnets_dict +end + +function mincut_subnetwork_insert_deltas( + network::Vector{ITensor}, source_inds::Vector{<:Index} +) + out_inds = noncommoninds(network...) + # terminal_inds = setdiff(out_inds, source_inds) + # tensors_to_add_delta = [] + # for t in network + # t_inds = inds(t) + # if length(intersect(source_inds, t_inds)) > 0 && length(intersect(terminal_inds, t_inds)) > 0 + # push!(tensors_to_add_delta, t) + # end + # end + # inds_to_add_delta = [] + # for t in tensors_to_add_delta + # uncontract_inds = intersect(inds(t), out_inds) + # inds_to_add_delta = [inds_to_add_delta..., uncontract_inds...] + # end + # deltas, networkprime, _ = split_deltas(inds_to_add_delta, network) + deltas, networkprime, _ = split_deltas(noncommoninds(network...), network) + network = Vector{ITensor}(vcat(deltas, networkprime)) + source_subnetwork = mincut_subnetwork(network, source_inds, out_inds) + remain_network = setdiff(network, source_subnetwork) + source_subnetwork = simplify_deltas(source_subnetwork) + remain_network = simplify_deltas(remain_network) + @assert ( + length(noncommoninds(network...)) == + length(noncommoninds(source_subnetwork..., remain_network...)) + ) + return source_subnetwork, remain_network +end + +function tree_embedding(network::Vector{ITensor}, inds_btree::Vector; algorithm) + btree_to_output_tn = Dict{Vector,Vector{ITensor}}() + btree_to_input_tn = Dict{Vector,Vector{ITensor}}() + btree_to_input_tn[inds_btree] = network + nodes = reverse(topo_sort(inds_btree; type=Vector{<:Vector})) + nodes = [nodes..., get_leaves(inds_btree)...] + for node in nodes + @assert haskey(btree_to_input_tn, node) + input_tn = btree_to_input_tn[node] + # @info "node", node + if length(node) == 1 + btree_to_output_tn[node] = input_tn + continue + end + net1, input_tn = mincut_subnetwork_insert_deltas(input_tn, vectorize(node[1])) + btree_to_input_tn[node[1]] = net1 + net1, input_tn = mincut_subnetwork_insert_deltas(input_tn, vectorize(node[2])) + btree_to_input_tn[node[2]] = net1 + btree_to_output_tn[node] = input_tn + # @info "btree_to_output_tn[node]", btree_to_output_tn[node] + end + if algorithm == "svd" + return btree_to_output_tn + else + return remove_deltas(btree_to_output_tn) + end +end + +is_delta(t) = (t.tensor.storage.data == 1.0) + +function simplify_deltas(network::Vector{ITensor}) + out_delta_inds = Vector{Pair}() + # outinds will always be the roots in union-find + outinds = noncommoninds(network...) + deltas = filter(t -> is_delta(t), network) + inds_list = map(t -> collect(inds(t)), deltas) + deltainds = collect(Set(vcat(inds_list...))) + uf = UF(deltainds) + for t in deltas + i1, i2 = inds(t) + if root(uf, i1) in outinds && root(uf, i2) in outinds + push!(out_delta_inds, root(uf, i1) => root(uf, i2)) + end + if root(uf, i1) in outinds + connect(uf, i2, i1) + else + connect(uf, i1, i2) + end + end + sim_dict = Dict([ind => root(uf, ind) for ind in deltainds]) + network = setdiff(network, deltas) + network = replaceinds(network, sim_dict) + out_delta = [delta(i.first, i.second) for i in out_delta_inds] + return Vector{ITensor}([network..., out_delta...]) +end + +# remove deltas to improve the performance +function remove_deltas(tnets_dict::Dict) + # only remove deltas in intermediate nodes + ks = filter(k -> (length(k) > 1), collect(keys(tnets_dict))) + network = vcat([tnets_dict[k] for k in ks]...) + # outinds will always be the roots in union-find + outinds = noncommoninds(network...) + + deltas = filter(t -> is_delta(t), network) + inds_list = map(t -> collect(inds(t)), deltas) + deltainds = collect(Set(vcat(inds_list...))) + uf = UF(deltainds) + for t in deltas + i1, i2 = inds(t) + if root(uf, i1) in outinds + connect(uf, i2, i1) + else + connect(uf, i1, i2) + end + end + sim_dict = Dict([ind => root(uf, ind) for ind in deltainds]) + for k in ks + net = tnets_dict[k] + net = setdiff(net, deltas) + tnets_dict[k] = replaceinds(net, sim_dict) + # @info "$(k), $(TreeTensor(net...))" + end + return tnets_dict +end + +function split_deltas(inds, subnet) + sim_dict = Dict([ind => sim(ind) for ind in inds]) + deltas = [delta(i, sim_dict[i]) for i in inds] + subnet = replaceinds(subnet, sim_dict) + return deltas, subnet, collect(values(sim_dict)) +end diff --git a/src/ApproximateTNContraction/contract/tree_embedding_.jl b/src/ApproximateTNContraction/contract/tree_embedding_.jl new file mode 100644 index 00000000..a6363c77 --- /dev/null +++ b/src/ApproximateTNContraction/contract/tree_embedding_.jl @@ -0,0 +1,119 @@ +include("union_find.jl") + +function tree_embedding(network::Vector{OrthogonalITensor}, inds_btree::Vector; algorithm) + # TODO: consider identity matrices + tensor_to_ortho_tensor = Dict{ITensor,OrthogonalITensor}() + for ortho_tensor in network + tensor_to_ortho_tensor[ortho_tensor.tensor] = ortho_tensor + end + tnets_dict = tree_embedding(get_tensors(network), inds_btree) + ortho_tnets_dict = Dict() + for (key, tensors) in tnets_dict + ortho_tensors = Vector{OrthogonalITensor}() + for t in tensors + if haskey(tensor_to_ortho_tensor, t) + push!(ortho_tensors, tensor_to_ortho_tensor[t]) + else + push!(ortho_tensors, OrthogonalITensor(t)) + end + end + ortho_tnets_dict[key] = ortho_tensors + end + return ortho_tnets_dict +end + +function tree_embedding(network::Vector{ITensor}, inds_btree::Vector) + deltas, networkprime, _ = split_deltas(noncommoninds(network...), network) + network = vcat(deltas, networkprime) + # tnets_dict map each inds_btree node to a tensor network + tnets_dict = Dict() + function embed(tree::Vector) + if length(tree) == 1 + # add delta to handle the case with two output edges neighboring to one tensor + # being split (MPS case with element grouping) + ind = tree[1] + sim_dict = Dict([ind => sim(ind)]) + tnets_dict[tree] = [delta(ind, sim_dict[ind])] + network = replaceinds(network, sim_dict) + return Tuple([sim_dict[ind]]) + end + ind1 = embed(tree[1]) + ind2 = embed(tree[2]) + deltas, splitinds, tnets_dict[tree[1]] = insert_deltas(ind1, ind2, tnets_dict[tree[1]]) + network = Vector{ITensor}(vcat(network, deltas)) + # use mincut to get the subnetwork + subnetwork = mincut_subnetwork(network, splitinds, noncommoninds(network...)) + subsplitinds = intersect(splitinds, noncommoninds(subnetwork...)) + remaininds = collect(setdiff(noncommoninds(subnetwork...), subsplitinds)) + network = collect(setdiff(network, subnetwork)) + # remaininds + deltas, subnetwork, _ = split_deltas(remaininds, subnetwork) + network = vcat(network, deltas) + # subsplitinds + inds = collect(setdiff(splitinds, subsplitinds)) + if length(inds) > 0 + inds = Vector{Index}(inds) + deltas, network, _ = split_deltas(inds, network) + subnetwork = vcat(subnetwork, deltas) + end + # @info "$(tree), $(TreeTensor(subnetwork...))" + tnets_dict[tree] = subnetwork + return Tuple(setdiff(noncommoninds(subnetwork...), splitinds)) + end + @assert (length(inds_btree) >= 2) + embed(inds_btree) + return remove_deltas(tnets_dict) +end + +is_delta(t) = (t.tensor.storage.data == 1.0) + +# remove deltas to improve the performance +function remove_deltas(tnets_dict) + # only remove deltas in intermediate nodes + ks = filter(k -> (length(k) > 1), collect(keys(tnets_dict))) + network = vcat([tnets_dict[k] for k in ks]...) + # outinds will always be the roots in union-find + outinds = noncommoninds(network...) + + deltas = filter(t -> is_delta(t), network) + inds_list = map(t -> collect(inds(t)), deltas) + deltainds = collect(Set(vcat(inds_list...))) + uf = UF(deltainds) + for t in deltas + i1, i2 = inds(t) + if root(uf, i1) in outinds + connect(uf, i2, i1) + else + connect(uf, i1, i2) + end + end + sim_dict = Dict([ind => root(uf, ind) for ind in deltainds]) + for k in ks + net = tnets_dict[k] + net = setdiff(net, deltas) + tnets_dict[k] = replaceinds(net, sim_dict) + # @info "$(k), $(TreeTensor(net...))" + end + return tnets_dict +end + +function split_deltas(inds, subnet) + sim_dict = Dict([ind => sim(ind) for ind in inds]) + deltas = [delta(i, sim_dict[i]) for i in inds] + subnet = replaceinds(subnet, sim_dict) + return deltas, subnet, collect(values(sim_dict)) +end + +function insert_deltas(ind1, ind2, subnet1) + intersect_inds = intersect(ind1, ind2) + ind1_unique = collect(setdiff(ind1, intersect_inds)) + ind2_unique = collect(setdiff(ind2, intersect_inds)) + outinds = vcat(ind1_unique, ind2_unique) + # look at intersect_inds + deltas = [] + if length(intersect_inds) >= 1 + deltas, subnet1, siminds = split_deltas(intersect_inds, subnet1) + outinds = vcat(outinds, intersect_inds, siminds) + end + return deltas, outinds, subnet1 +end diff --git a/src/ApproximateTNContraction/contract/tree_utils.jl b/src/ApproximateTNContraction/contract/tree_utils.jl new file mode 100644 index 00000000..211c7804 --- /dev/null +++ b/src/ApproximateTNContraction/contract/tree_utils.jl @@ -0,0 +1,103 @@ + +# merge two trees +# new tree: +# s +# / \ +# t1 t2 +function merge_tree(t1::Vector, t2::Vector; append=false) + if t2 == [] + return t1 + end + if t1 == [] + return t2 + end + if isleaf(t1) && isleaf(t2) + return [t1, t2] + end + if isleaf(t1) + return append ? [t1, t2...] : [t1, t2] + end + if isleaf(t2) + return append ? [t1..., t2] : [t1, t2] + end + return append ? [t1..., t2...] : [t1, t2] +end + +function isleaf(tree::Vector) + if tree == [] + @info "tree is empty" + return false + end + if all(v -> !(v isa Vector), tree) + return true + end + return false +end + +# get the subtree of tree that is in the subset +# example: +# subtree([[1, 2], [3, 4]], [1, 3]) = ([[1], [3]]) +function subtree(tree::Vector, subset::Union{Vector,Tuple}) + if tree == [] + return [] + end + if isleaf(tree) + return intersect(tree, subset) + end + tree = [subtree(i, subset) for i in tree] + tree = filter(t -> t != [], tree) + if length(tree) == 1 && tree[1] isa Vector + return tree[1] + end + return tree +end + +# vectorize a tree +# example: [[1,2], [3,4]] = [1, 2, 3, 4] +function vectorize(tree) + @assert tree != [] + if !(tree isa Vector) + return [tree] + end + return mapreduce(vectorize, vcat, tree) +end + +# example: [[[1,2], [3,4]], [[5,6], [7,8]]] = [[1,2], [3,4], [5,6], [7,8]] +function get_leaves(tree::Vector) + if tree == [] + return [] + end + if !(tree isa Vector{<:Vector}) + return [tree] + end + return mapreduce(get_leaves, vcat, tree) +end + +function line_to_tree(line::Vector) + if length(line) == 1 && line[1] isa Vector + return line[1] + end + if length(line) <= 2 + return line + end + return [line_to_tree(line[1:(end - 1)]), line[end]] +end + +function topo_sort(tn; type=Vector, leaves=[]) + @timeit timer "topo_sort" begin + topo_order = [] + topo_sort_dfs!(tn, topo_order, leaves, type) + return topo_order + end +end + +function topo_sort_dfs!(tn, topo_order, leaves, type) + #Post-order DFS + if (tn in leaves) || !(tn isa type) + return nothing + end + for subtn in tn + topo_sort_dfs!(subtn, topo_order, leaves, type) + end + return append!(topo_order, [tn]) +end diff --git a/src/ApproximateTNContraction/contract/union_find.jl b/src/ApproximateTNContraction/contract/union_find.jl new file mode 100644 index 00000000..0ae3a4dc --- /dev/null +++ b/src/ApproximateTNContraction/contract/union_find.jl @@ -0,0 +1,28 @@ +struct UF + parent_map::Dict +end + +function UF(values::Vector) + parent_map = Dict() + for value in values + parent_map[value] = value + end + return UF(parent_map) +end + +function root(uf::UF, n) + while uf.parent_map[n] != n + n = uf.parent_map[n] + end + return n +end + +function connect(uf, n1, n2) + rootn1 = root(uf, n1) + rootn2 = root(uf, n2) + if rootn1 == rootn2 + # Already connected + return nothing + end + return uf.parent_map[rootn1] = rootn2 +end diff --git a/src/ApproximateTNContraction/interfaces/sweep_contractor.jl b/src/ApproximateTNContraction/interfaces/sweep_contractor.jl new file mode 100644 index 00000000..e62a3a40 --- /dev/null +++ b/src/ApproximateTNContraction/interfaces/sweep_contractor.jl @@ -0,0 +1,33 @@ +using SweepContractor +using SweepContractor: TensorNetwork, LabelledTensorNetwork + +function ITensor_networks(TN::TensorNetwork) + index_dict = Dict() + function itensor(i, t) + # Construct indices + inds = [] + for (dim, j) in enumerate(t.adj) + label = sort([i, j]) + if !haskey(index_dict, label) + s = size(t.arr)[dim] + index_dict[label] = Index(s, string(label)) + end + push!(inds, index_dict[label]) + end + # build the tensor + return ITensor(t.arr, inds...) + end + inetwork = [itensor(i, t) for (i, t) in enumerate(TN)] + return inetwork +end + +function ITensor_networks(LTN::LabelledTensorNetwork) + return ITensor_networks(SweepContractor.delabel(LTN)) +end + +function line_network(network::Vector) + if length(network) <= 2 + return network + end + return [line_network(network[1:(end - 1)]), [network[end]]] +end diff --git a/src/ApproximateTNContraction/networks/itensor_network.jl b/src/ApproximateTNContraction/networks/itensor_network.jl new file mode 100644 index 00000000..152bc7fa --- /dev/null +++ b/src/ApproximateTNContraction/networks/itensor_network.jl @@ -0,0 +1,32 @@ +function ITensors.prime(indices::Array{<:Index,1}, network::Array{ITensor}, n::Integer=1) + function primeinds(tensor) + prime_inds = [ind for ind in inds(tensor) if ind in indices] + if (length(prime_inds) == 0) + return tensor + end + return replaceinds(tensor, prime_inds => prime(prime_inds, n)) + end + return map(x -> primeinds(x), network) +end + +function ITensors.replaceinds( + network::Union{Array{ITensor},Array{OrthogonalITensor}}, sim_dict::Dict +) + if length(network) == 0 + return network + end + indices = collect(keys(sim_dict)) + function siminds(tensor) + sim_inds = [ind for ind in inds(tensor) if ind in indices] + if (length(sim_inds) == 0) + return tensor + end + outinds = map(i -> sim_dict[i], sim_inds) + return replaceinds(tensor, sim_inds => outinds) + end + return map(x -> siminds(x), network) +end + +function ITensors.commoninds(n1::Array{ITensor}, n2::Array{ITensor}) + return mapreduce(a -> commoninds(a...), vcat, zip(n1, n2)) +end diff --git a/src/ApproximateTNContraction/orthogonal_tensor.jl b/src/ApproximateTNContraction/orthogonal_tensor.jl new file mode 100644 index 00000000..5354259c --- /dev/null +++ b/src/ApproximateTNContraction/orthogonal_tensor.jl @@ -0,0 +1,46 @@ +mutable struct OrthogonalITensor + tensor::ITensor + ortho_indices::Vector +end + +function OrthogonalITensor(tensor::ITensor) + return OrthogonalITensor(tensor, []) +end + +function orthogonal_tensors(tensors::Vector{ITensor}) + return [OrthogonalITensor(t) for t in tensors] +end + +function get_tensors(ortho_tensors::Vector{OrthogonalITensor}) + return [t.tensor for t in ortho_tensors] +end + +function ITensors.noncommoninds(ortho_tensors::OrthogonalITensor...) + if length(ortho_tensors) == 1 + return collect(inds(ortho_tensors[1])) + end + return noncommoninds(get_tensors([ortho_tensors...])...) +end + +function ITensors.inds(ortho_tensor::OrthogonalITensor) + return inds(ortho_tensor.tensor) +end + +function ITensors.replaceinds(ortho_tensor::OrthogonalITensor, change_inds::Pair) + input_inds = change_inds[1] + output_inds = change_inds[2] + input_to_output = Dict{Index,Index}() + for (i_in, i_out) in zip(input_inds, output_inds) + input_to_output[i_in] = i_out + end + new_ortho_indices = [] + for i in ortho_tensor.ortho_indices + if haskey(input_to_output, i) + push!(new_ortho_indices, input_to_output[i]) + else + push!(new_ortho_indices, i) + end + end + new_tensor = replaceinds(ortho_tensor.tensor, change_inds) + return OrthogonalITensor(new_tensor, new_ortho_indices) +end diff --git a/src/ITensorNetworks.jl b/src/ITensorNetworks.jl index db160341..7e41e69d 100644 --- a/src/ITensorNetworks.jl +++ b/src/ITensorNetworks.jl @@ -16,6 +16,7 @@ using NamedGraphs using Observers using Printf using Requires +using Reexport using SimpleTraits using SparseArrayKit using SplitApplyCombine @@ -124,4 +125,6 @@ function __init__() ) end +include("ApproximateTNContraction/ApproximateTNContraction.jl") + end diff --git a/src/specialitensornetworks.jl b/src/specialitensornetworks.jl index cd560d76..7af00894 100644 --- a/src/specialitensornetworks.jl +++ b/src/specialitensornetworks.jl @@ -1,3 +1,5 @@ +using LinearAlgebra: diagm + """ RETURN A TENSOR NETWORK WITH COPY TENSORS ON EACH VERTEX. Note that passing a link_space will mean the indices of the resulting network don't match those of the input indsnetwork @@ -49,6 +51,39 @@ function ising_network(eltype::Type, s::IndsNetwork, beta::Number; szverts=nothi return ψ end +function ising_network( + eltype::Type, s::IndsNetwork, beta::Number, h::Number; szverts=nothing +) + tn = delta_network(eltype, s) + if (szverts != nothing) + for v in szverts + tn[v] = diagITensor(eltype[1, -1], inds(tn[v])) + end + end + for edge in edges(tn) + v1 = src(edge) + v2 = dst(edge) + i = commoninds(tn[v1], tn[v2])[1] + deg_v1 = degree(tn, v1) + deg_v2 = degree(tn, v2) + f11 = exp(beta + h / deg_v1 + h / deg_v2) + f12 = exp(-beta + h / deg_v1 - h / deg_v2) + f21 = exp(-beta - h / deg_v1 + h / deg_v2) + f22 = exp(beta - h / deg_v1 - h / deg_v2) + q = eltype[f11 f12; f21 f22] + w, V = eigen(q) + w = map(sqrt, w) + sqrt_q = V * diagm(w) * inv(V) + t = itensor(sqrt_q, i, i') + tn[v1] = tn[v1] * t + tn[v1] = noprime!(tn[v1]) + t = itensor(sqrt_q, i', i) + tn[v2] = tn[v2] * t + tn[v2] = noprime!(tn[v2]) + end + return tn +end + function ising_network(s::IndsNetwork, beta::Number; szverts=nothing) return ising_network(typeof(beta), s, beta; szverts) end @@ -62,10 +97,20 @@ function ising_network(eltype::Type, g::NamedGraph, beta::Number; szverts=nothin return ising_network(eltype, IndsNetwork(g; link_space=2), beta; szverts) end +function ising_network( + eltype::Type, g::NamedGraph, beta::Number, h::Number; szverts=nothing +) + return ising_network(eltype, IndsNetwork(g; link_space=2), beta, h; szverts) +end + function ising_network(g::NamedGraph, beta::Number; szverts=nothing) return ising_network(eltype(beta), g, beta; szverts) end +function ising_network(g::NamedGraph, beta::Number, h::Number; szverts=nothing) + return ising_network(eltype(beta), g, beta, h; szverts) +end + """ Build an ITensor network on a graph specified by the inds network s. Bond_dim is given by link_space and entries are randomised (normal distribution, mean 0 std 1) """ @@ -79,7 +124,9 @@ function randomITensorNetwork(s::IndsNetwork; link_space=nothing) return randomITensorNetwork(Float64, s; link_space) end -@traitfn function randomITensorNetwork(eltype::Type, g::::IsUnderlyingGraph; link_space=nothing) +@traitfn function randomITensorNetwork( + eltype::Type, g::::IsUnderlyingGraph; link_space=nothing +) return randomITensorNetwork(eltype, IndsNetwork(g); link_space) end diff --git a/src/treetensornetworks/abstracttreetensornetwork.jl b/src/treetensornetworks/abstracttreetensornetwork.jl index 4b4e977e..ce305c94 100644 --- a/src/treetensornetworks/abstracttreetensornetwork.jl +++ b/src/treetensornetworks/abstracttreetensornetwork.jl @@ -114,7 +114,7 @@ end # function truncate( - ψ::AbstractTTN; root_vertex::Tuple=default_root_vertex(ψ), kwargs... + ψ::AbstractTTN; root_vertex::Union{Int,Tuple}=default_root_vertex(ψ), kwargs... ) for e in post_order_dfs_edges(ψ, root_vertex) # always orthogonalize towards source first to make truncations controlled diff --git a/test/Project.toml b/test/Project.toml index c4811c4e..b9fe9748 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -17,6 +17,8 @@ OMEinsumContractionOrders = "6f22d1fd-8eed-4bb7-9776-e7d684900715" Observers = "338f10d5-c7f1-4033-a7d1-f9dec39bcaa0" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +SweepContractor = "75a5deae-e917-4509-af32-a989148c8d5f" +TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f" SplitApplyCombine = "03a91e81-4c3e-53e1-a0a4-9c0c8f19dd66" Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" diff --git a/test/approximate_tn_contraction/cache.jl b/test/approximate_tn_contraction/cache.jl new file mode 100644 index 00000000..b82ec33a --- /dev/null +++ b/test/approximate_tn_contraction/cache.jl @@ -0,0 +1,17 @@ +using ITensors +using ITensorNetworks.ApproximateTNContraction: IndexGroup, get_igs_cache_info + +@testset "test get_igs_cache_info" begin + i = IndexGroup([Index(2, "i")]) + j = IndexGroup([Index(3, "j")]) + k = IndexGroup([Index(2, "k")]) + l = IndexGroup([Index(4, "l")]) + m = IndexGroup([Index(5, "m")]) + n = IndexGroup([Index(5, "n")]) + + type = Vector{IndexGroup} + igs_list = [type([l, k, i, m, n]), type([i, j, k]), type([l, j, m, n])] + contract_igs_list = [type([m]), type([j]), type([j])] + out = get_igs_cache_info(igs_list, contract_igs_list) + @test out == ([l, k, i, m, n], [], []) +end diff --git a/test/approximate_tn_contraction/contract.jl b/test/approximate_tn_contraction/contract.jl new file mode 100644 index 00000000..5bc3e75b --- /dev/null +++ b/test/approximate_tn_contraction/contract.jl @@ -0,0 +1,188 @@ +using ITensors, TimerOutputs +using ITensorNetworks.ApproximateTNContraction: + get_tensors, + OrthogonalITensor, + tree_approximation, + inds_binary_tree, + tree_embedding, + approximate_contract +using ITensorNetworks.ApproximateTNContraction: timer + +include("utils.jl") + +@testset "test tree approximation" begin + i = Index(2, "i") + j = Index(2, "j") + k = Index(2, "k") + l = Index(2, "l") + m = Index(2, "m") + n = Index(2, "n") + o = Index(2, "o") + p = Index(2, "p") + q = Index(2, "q") + r = Index(2, "r") + s = Index(2, "s") + t = Index(2, "t") + u = Index(2, "u") + A = OrthogonalITensor(randomITensor(i, n)) + B = OrthogonalITensor(randomITensor(j, o)) + AB = OrthogonalITensor(randomITensor(n, o, r)) + C = OrthogonalITensor(randomITensor(k, p)) + D = OrthogonalITensor(randomITensor(l, q)) + E = OrthogonalITensor(randomITensor(m, u)) + CD = OrthogonalITensor(randomITensor(p, q, s)) + ABCD = OrthogonalITensor(randomITensor(r, s, t)) + ABCDE = OrthogonalITensor(randomITensor(t, u)) + btree = [[[[i], [j]], [[k], [l]]], [m]] + tensors = [A, B, C, D, E, AB, CD, ABCD, ABCDE] + embedding = Dict([ + [i] => [A], + [j] => [B], + [k] => [C], + [l] => [D], + [m] => [E], + [[i], [j]] => [AB], + [[k], [l]] => [CD], + [[[i], [j]], [[k], [l]]] => [ABCD], + [[[[i], [j]], [[k], [l]]], [m]] => [ABCDE], + ]) + out, log_norm = tree_approximation(embedding, btree; algorithm="svd") + out[btree].tensor *= exp(log_norm) + out = get_tensors(collect(values(out))) + @test isapprox(contract(out...), contract(get_tensors(tensors)...)) +end + +@testset "test MPS times MPO" begin + N = (5, 3) + linkdim = 3 + cutoff = 1e-15 + network = randomITensorNetwork(IndsNetwork(named_grid(N)); link_space=linkdim) + tn = Array{ITensor,length(N)}(undef, N...) + for v in vertices(network) + tn[v...] = network[v...] + end + x, A = tn[:, 1], tn[:, 2] + out_true = contract(MPO(A), MPS(x); cutoff=cutoff, maxdim=linkdim * linkdim) + out2, log_norm = approximate_contract([A, x]; cutoff=cutoff, maxdim=linkdim * linkdim) + tsr_true = contract(out_true...) + tsr_nrmsquare = (tsr_true * tsr_true)[1] + @test isapprox(tsr_true, contract(out2...) * exp(log_norm)) + + maxdims = [2, 4, 6, 8] + for dim in maxdims + out = contract(MPO(A), MPS(x); cutoff=cutoff, maxdim=dim) + out2, log_norm = approximate_contract([A, x]; cutoff=cutoff, maxdim=dim) + residual1 = tsr_true - contract(out...) + residual2 = tsr_true - contract(out2...) * exp(log_norm) + error1 = sqrt((residual1 * residual1)[1] / tsr_nrmsquare) + error2 = sqrt((residual2 * residual2)[1] / tsr_nrmsquare) + print("maxdim, ", dim, ", error1, ", error1, ", error2, ", error2, "\n") + end +end + +@testset "test inds_binary_tree" begin + i = Index(2, "i") + j = Index(2, "j") + k = Index(2, "k") + l = Index(2, "l") + m = Index(2, "m") + n = Index(2, "n") + o = Index(2, "o") + p = Index(2, "p") + + T = randomITensor(i, j, k, l, m, n, o, p) + M = MPS(T, (i, j, k, l, m, n, o, p); cutoff=1e-5, maxdim=500) + network = M[:] + + out = inds_binary_tree(network, [i, j, k, l, m, n, o, p]; algorithm="mincut") + @test length(out) == 2 + out = inds_binary_tree(network, [i, j, k, l, m, n, o, p]; algorithm="mps") + @test length(out) == 2 +end + +@testset "test inds_binary_tree of a 2D network" begin + N = (3, 3, 3) + linkdim = 2 + network = randomITensorNetwork(IndsNetwork(named_grid(N)); link_space=linkdim) + tn = Array{ITensor,length(N)}(undef, N...) + for v in vertices(network) + tn[v...] = network[v...] + end + network = vec(tn[:, :, 1]) + out = inds_binary_tree(network, noncommoninds(network...); algorithm="mincut") + @test length(out) == 2 + out = inds_binary_tree(network, noncommoninds(network...); algorithm="mps") + @test length(out) == 2 +end + +@testset "test tree embedding" begin + i = Index(2, "i") + j = Index(2, "j") + k = Index(2, "k") + l = Index(2, "l") + m = Index(2, "m") + T = randomITensor(i, j, k, l, m) + M = MPS(T, (i, j, k, l, m); cutoff=1e-5, maxdim=5) + network = M[:] + out1 = contract(network...) + inds_btree = inds_binary_tree(network, [i, j, k, l, m]; algorithm="mincut") + tnet_dict = tree_embedding(network, inds_btree; algorithm="mincut") + network2 = vcat(collect(values(tnet_dict))...) + out2 = contract(network2...) + i1 = noncommoninds(network...) + i2 = noncommoninds(network2...) + @test (length(i1) == length(i2)) + @test isapprox(out1, out2) +end + +function benchmark_peps_contraction(tn; cutoff=1e-15, maxdim=1000) + out = peps_contraction_mpomps(tn; cutoff=cutoff, maxdim=maxdim, snake=false) + out2 = contract_line_group(tn; cutoff=cutoff, maxdim=maxdim) + return out[], out2 +end + +@testset "test PEPS" begin + N = (8, 8) + linkdim = 2 + cutoff = 1e-15 + network = randomITensorNetwork(IndsNetwork(named_grid(N)); link_space=linkdim) + tn = Array{ITensor,length(N)}(undef, N...) + for v in vertices(network) + tn[v...] = network[v...] + end + + ITensors.set_warn_order(100) + maxdim = linkdim^N[2] + out_true, out2 = benchmark_peps_contraction(tn; cutoff=cutoff, maxdim=maxdim) + print(out_true, out2) + @test abs((out_true - out2) / out_true) < 1e-3 + + maxdims = [i for i in 2:16] + for dim in maxdims + size = dim * dim * linkdim + out, out2 = benchmark_peps_contraction(tn; cutoff=cutoff, maxdim=dim) + error1 = abs((out - out_true) / out_true) + error2 = abs((out2 - out_true) / out_true) + print("maxdim, ", dim, ", error1, ", error1, ", error2, ", error2, "\n") + end +end + +@testset "benchmark PEPS" begin + N = (8, 8) + linkdim = 10 + cutoff = 1e-15 + network = randomITensorNetwork(IndsNetwork(named_grid(N)); link_space=linkdim) + tn = Array{ITensor,length(N)}(undef, N...) + for v in vertices(network) + tn[v...] = network[v...] + end + dim = 20 + # warmup + ITensors.set_warn_order(100) + benchmark_peps_contraction(tn; cutoff=cutoff, maxdim=dim) + + reset_timer!(timer) + ITensors.set_warn_order(100) + benchmark_peps_contraction(tn; cutoff=cutoff, maxdim=dim) + show(timer) +end diff --git a/test/approximate_tn_contraction/indexgroup.jl b/test/approximate_tn_contraction/indexgroup.jl new file mode 100644 index 00000000..d4648566 --- /dev/null +++ b/test/approximate_tn_contraction/indexgroup.jl @@ -0,0 +1,77 @@ +using ITensors +using ITensorNetworks +using ITensorNetworks.ApproximateTNContraction: + IndexGroup, get_index_groups, get_leaves, neighbor_index_groups +using ITensorNetworks.ApproximateTNContraction: + line_network, + IndexAdjacencyTree, + topo_sort, + get_ancestors, + generate_adjacency_tree, + minswap_adjacency_tree!, + minswap_adjacency_tree, + approximate_contract + +@testset "test generate_adjacency_tree" begin + N = (3, 3) + tn = Vector{ITensor}(randomITensorNetwork(IndsNetwork(named_grid(N)); link_space=2)) + ctree = line_network(tn) + tn_leaves = get_leaves(ctree) + ctrees = topo_sort(ctree; leaves=tn_leaves) + ctree_to_igs = Dict{Vector,Vector{IndexGroup}}() + index_groups = get_index_groups(ctree) + for c in vcat(tn_leaves, ctrees) + ctree_to_igs[c] = neighbor_index_groups(c, index_groups) + end + ctree_to_ancestors = get_ancestors(ctree) + adj_tree1 = generate_adjacency_tree( + tn_leaves[4], ctree_to_ancestors[tn_leaves[4]], ctree_to_igs + ) + adj_tree2 = generate_adjacency_tree( + ctrees[2], ctree_to_ancestors[ctrees[2]], ctree_to_igs + ) + for adj_tree in [adj_tree1, adj_tree1] + @test length(adj_tree.children) == 3 + @test adj_tree.fixed_order = true + c1, c2, c3 = adj_tree.children + @test length(c1.children) == 1 + @test length(c2.children) == 2 + @test length(c3.children) == 1 + end +end + +# @testset "test minswap_adjacency_tree!" begin +# i = IndexGroup([Index(2, "i")]) +# j = IndexGroup([Index(3, "j")]) +# k = IndexGroup([Index(2, "k")]) +# l = IndexGroup([Index(4, "l")]) +# m = IndexGroup([Index(5, "m")]) +# n = IndexGroup([Index(5, "n")]) +# I = IndexAdjacencyTree(i) +# J = IndexAdjacencyTree(j) +# K = IndexAdjacencyTree(k) +# L = IndexAdjacencyTree(l) +# M = IndexAdjacencyTree(m) +# N = IndexAdjacencyTree(n) +# JKL = IndexAdjacencyTree([J, K, L], false, false) +# tree = IndexAdjacencyTree([I, JKL, M], false, false) +# tree_copy = copy(tree) +# tree2 = IndexAdjacencyTree([i, k, m, j, l], true, true) +# nswaps = minswap_adjacency_tree!(tree, tree2) +# @test nswaps == 1 +# @test tree.children == [i, m, k, j, l] +# @test tree.fixed_direction && tree.fixed_order +# # test minswap_adjacency_tree +# tree3 = IndexAdjacencyTree([i, k, n, m], true, true) +# tree4 = IndexAdjacencyTree([l, n, j], true, true) +# out = minswap_adjacency_tree(tree_copy, tree3, tree4) +# @test out.children in [[i, m, k, j, l], [i, m, k, l, j], [m, i, k, j, l], [m, i, k, l, j]] +# end + +@testset "test approximate_contract" begin + N = (4, 4) + tn = Vector{ITensor}(randomITensorNetwork(IndsNetwork(named_grid(N)); link_space=2)) + ctree = line_network(tn) + approximate_contract(ctree; cutoff=1e-5, maxdim=20, ansatz="mps") + approximate_contract(ctree; cutoff=1e-5, maxdim=20, ansatz="comb") +end diff --git a/test/approximate_tn_contraction/interface.jl b/test/approximate_tn_contraction/interface.jl new file mode 100644 index 00000000..db5ed69d --- /dev/null +++ b/test/approximate_tn_contraction/interface.jl @@ -0,0 +1,163 @@ +using ITensors, Random, SweepContractor +using ITensorNetworks.ApproximateTNContraction: + timer, ITensor_networks, line_network, approximate_contract + +include("utils.jl") + +@testset "test the interface" begin + LTN = LabelledTensorNetwork{Char}() + LTN['A'] = Tensor(['D', 'B'], [i^2 - 2j for i in 0:2, j in 0:2], 0, 1) + LTN['B'] = Tensor(['A', 'D', 'C'], [-3^i * j + k for i in 0:2, j in 0:2, k in 0:2], 0, 0) + LTN['C'] = Tensor(['B', 'D'], [j for i in 0:2, j in 0:2], 1, 0) + LTN['D'] = Tensor(['A', 'B', 'C'], [i * j * k for i in 0:2, j in 0:2, k in 0:2], 1, 1) + + sweep = sweep_contract(LTN, 100, 100; fast=true) + out = ldexp(sweep...) + @test isapprox(out, contract(ITensor_networks(LTN))[]) + show(timer) +end + +function lattice(row, column, d) + function build_adj(i, j) + adj = Vector{Int64}() + i > 1 && push!(adj, delabel[(i - 1, j)]) + i < row && push!(adj, delabel[(i + 1, j)]) + j > 1 && push!(adj, delabel[(i, j - 1)]) + j < column && push!(adj, delabel[(i, j + 1)]) + return adj + end + TN = TensorNetwork() + delabel = Dict() + index = 1 + for i in 1:row + ranges = iseven(i) ? (column:-1:1) : (1:column) + for j in ranges + delabel[(i, j)] = index + index += 1 + end + end + for i in 1:row + ranges = iseven(i) ? (column:-1:1) : (1:column) + for j in ranges + adj = build_adj(i, j) + push!(TN, Tensor(adj, randn(d * ones(Int, length(adj))...), i, j)) + end + end + return TN +end + +function get_contracted_peps(LTN, rank, N) + tnet = ITensor_networks(LTN) + tnet_mat = reshape(tnet, N...) + out_mps = peps_contraction_mpomps(tnet_mat; cutoff=1e-15, maxdim=rank, snake=true) + out = contract_w_sweep(LTN, rank) + out2 = contract_element_group(tnet; cutoff=1e-15, maxdim=rank) + out3 = contract_line_group(tnet_mat; cutoff=1e-15, maxdim=rank) + return out, out2, out3, out_mps[] +end + +@testset "test on 2D grid" begin + Random.seed!(1234) + ITensors.set_warn_order(100) + row, column, d = 8, 8, 2 + LTN = lattice(row, column, d) + + out_true, out_element, out_line, out_mps = get_contracted_peps( + LTN, d^(Int(row / 2)), [row, column] + ) + @test abs((out_true - out_element) / out_true) < 1e-3 + @test abs((out_true - out_line) / out_true) < 1e-3 + @test abs((out_true - out_mps) / out_true) < 1e-3 + for rank in [2, 3, 4, 6, 8, 10, 12, 14, 15, 16] + out, out_element, out_line, out_mps = get_contracted_peps(LTN, rank, [row, column]) + error_sweepcontractor = abs((out - out_true) / out_true) + error_element = abs((out_element - out_true) / out_true) + error_line = abs((out_line - out_true) / out_true) + error_mps = abs((out_mps - out_true) / out_true) + print( + "maxdim, ", + rank, + ", error_sweepcontractor, ", + error_sweepcontractor, + ", error_element, ", + error_element, + ", error_line, ", + error_line, + ", error_mps, ", + error_mps, + "\n", + ) + end +end + +@testset "benchmark on 2D grid" begin + Random.seed!(1234) + ITensors.set_warn_order(100) + row, column, d, rank = 10, 10, 2, 10 + LTN = lattice(row, column, d) + # warm-up + get_contracted_peps(LTN, rank, [row, column]) + @info "start benchmark on 2D grid" + reset_timer!(timer) + LTN = lattice(row, column, d) + get_contracted_peps(LTN, rank, [row, column]) + show(timer) +end + +function cube_3d(L=3, d=2) + function build_adj(i, j, k) + adj = Vector{Int64}() + i > 1 && push!(adj, delabel[(i - 1, j, k)]) + i < L && push!(adj, delabel[(i + 1, j, k)]) + j > 1 && push!(adj, delabel[(i, j - 1, k)]) + j < L && push!(adj, delabel[(i, j + 1, k)]) + k > 1 && push!(adj, delabel[(i, j, k - 1)]) + k < L && push!(adj, delabel[(i, j, k + 1)]) + return adj + end + TN = TensorNetwork() + delabel = Dict() + index = 1 + for i in 1:L + ranges_j = iseven(i) ? (L:-1:1) : (1:L) + for j in ranges_j + ranges_k = iseven((i - 1) * L + j) ? (L:-1:1) : (1:L) + for k in ranges_k + delabel[(i, j, k)] = index + index += 1 + end + end + end + for i in 1:L + ranges_j = iseven(i) ? (L:-1:1) : (1:L) + for j in ranges_j + ranges_k = iseven((i - 1) * L + j) ? (L:-1:1) : (1:L) + for k in ranges_k + adj = build_adj(i, j, k) + newt = Tensor( + adj, randn(d * ones(Int, length(adj))...), i + 0.01 * randn(), j + 0.01 * randn() + ) + push!(TN, newt) + end + end + end + return TN +end + +@testset "test on 3D cube with element grouping" begin + Random.seed!(1234) + ITensors.set_warn_order(100) + L, d = 3, 2 + rank = 16 + TN = cube_3d(L, d) + out = contract_w_sweep(TN, rank) + tnet = ITensor_networks(TN) + out2 = contract_element_group(tnet; cutoff=1e-15, maxdim=rank) + + reset_timer!(timer) + TN = cube_3d(L, d) + out = contract_w_sweep(TN, rank) + tnet = ITensor_networks(TN) + out2 = contract_element_group(tnet; cutoff=1e-15, maxdim=rank) + show(timer) +end diff --git a/test/approximate_tn_contraction/runtests.jl b/test/approximate_tn_contraction/runtests.jl new file mode 100644 index 00000000..21464a3c --- /dev/null +++ b/test/approximate_tn_contraction/runtests.jl @@ -0,0 +1,9 @@ +using ITensorNetworks +using Test + +@testset "ApproximateTNContraction.jl" begin + for filename in ["tree.jl", "indexgroup.jl", "cache.jl", "contract.jl", "interface.jl"] + println("Running $filename in ApproximateTNContraction.jl") + include(filename) + end +end diff --git a/test/approximate_tn_contraction/tree.jl b/test/approximate_tn_contraction/tree.jl new file mode 100644 index 00000000..816de6c4 --- /dev/null +++ b/test/approximate_tn_contraction/tree.jl @@ -0,0 +1,26 @@ +using TimerOutputs +using ITensorNetworks.ApproximateTNContraction: + timer, merge_tree, subtree, vectorize, topo_sort, get_leaves + +@testset "test merge tree" begin + t1 = [[1], [2], [3]] + t2 = [4, 5, 6] + @test merge_tree(t1, t2; append=true) == [[1], [2], [3], [4, 5, 6]] + @test merge_tree(t1, t2; append=false) == [[[1], [2], [3]], [4, 5, 6]] + @test merge_tree([], [1, 2, 3]; append=false) == [1, 2, 3] +end + +@testset "test subtree and vectorize" begin + t1 = [[[1, 2], [3]], [4]] + subset = [1] + @test subtree(t1, subset) == [1] + @test vectorize(t1) == [1, 2, 3, 4] +end + +@testset "test find topo sort" begin + reset_timer!(timer) + tn = [[[1, 2], [3, 4]], [[5, 6], [7, 8]]] + @test length(topo_sort(tn)) == 7 + @test length(topo_sort(tn; leaves=get_leaves(tn))) == 3 + show(timer) +end diff --git a/test/approximate_tn_contraction/utils.jl b/test/approximate_tn_contraction/utils.jl new file mode 100644 index 00000000..0df755fc --- /dev/null +++ b/test/approximate_tn_contraction/utils.jl @@ -0,0 +1,41 @@ +using TimerOutputs +using ITensorNetworks.ApproximateTNContraction: timer, line_network, approximate_contract + +function peps_contraction_mpomps(tn; cutoff=1e-15, maxdim=1000, snake=false) + @timeit timer "peps_contraction_mpomps" begin + N = size(tn) + x = tn[:, 1] + for i in 2:(N[2] - 1) + A = (iseven(i) && snake) ? reverse(tn[:, i]) : tn[:, i] + x = contract(MPO(A), MPS(x); cutoff=cutoff, maxdim=maxdim)[:] + end + return contract(x..., tn[:, N[2]]...) + end +end + +function contract_w_sweep(tn, rank) + @timeit timer "contract_w_sweep" begin + sweep = sweep_contract(tn, rank, rank) + return ldexp(sweep...) + end +end + +function contract_element_group(tnet; cutoff, maxdim) + @timeit timer "contract_element_group" begin + element_grouping = line_network(tnet) + out, log_norm = approximate_contract(element_grouping; cutoff=cutoff, maxdim=maxdim) + return out[1][1] * exp(log_norm) + end +end + +function contract_line_group(tnet; cutoff, maxdim) + N = size(tnet) + @timeit timer "contract_line_group" begin + tntree = tnet[:, 1] + for i in 2:N[2] + tntree = [tntree, tnet[:, i]] + end + out, log_norm = approximate_contract(tntree; cutoff=cutoff, maxdim=maxdim) + return out[1][1] * exp(log_norm) + end +end diff --git a/test/runtests.jl b/test/runtests.jl index 58c64f41..01ab1c06 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -9,6 +9,11 @@ using ITensorNetworks println("Running test file $test_file") @time include(joinpath(root, test_file)) end + + for filename in ["approximate_tn_contraction/runtests.jl"] + println("Running $filename") + include(filename) + end end nothing