[WIP] Introduce partitioned_contract

examples/approx_itensornetwork/bench_mps_times_mpos.jl

@@ -0,0 +1,60 @@
+using ITensors, Graphs, NamedGraphs, TimerOutputs
+using Distributions, Random
+using KaHyPar, Metis
+using ITensorNetworks
+using ITensorNetworks: ising_network, contract, _nested_vector_to_digraph
+using OMEinsumContractionOrders
+using AbstractTrees
+
+include("utils.jl")
+
+function mps_times_mpos(len_mps, num_mpo; link_space, physical_space)
+  inds_net = IndsNetwork(named_grid((len_mps, num_mpo)))
+  for j in 1:num_mpo
+    for i in 1:(len_mps - 1)
+      inds_net[(i, j) => (i + 1, j)] = [Index(link_space, "$(i)x$(j),$(i+1)x$(j)")]
+    end
+  end
+  for i in 1:len_mps
+    for j in 1:(num_mpo - 1)
+      inds_net[(i, j) => (i, j + 1)] = [Index(physical_space, "$(i)x$(j),$(i)x$(j+1)")]
+    end
+  end
+  inds_leaves = [Index(physical_space, "$(i)x$(num_mpo),out") for i in 1:len_mps]
+  for i in 1:len_mps
+    inds_net[(i, num_mpo)] = [inds_leaves[i]]
+  end
+  distribution = Uniform{Float64}(-1.0, 1.0)
+  return randomITensorNetwork(distribution, inds_net), inds_leaves
+end
+
+Random.seed!(1234)
+TimerOutputs.enable_debug_timings(ITensorNetworks)
+# TimerOutputs.enable_debug_timings(@__MODULE__)
+ITensors.set_warn_order(100)
+
+len_mps = 30
+num_mpo = 2
+link_space = 80 # the largest seems 80, beyond which would get too slow.
+maxdim = link_space
+physical_space = 2
+tree = "comb"
+
+tn, inds_leaves = mps_times_mpos(
+  len_mps, num_mpo; link_space=link_space, physical_space=physical_space
+)
+if tree == "mps"
+  btree = _nested_vector_to_digraph(linear_sequence(inds_leaves))
+else
+  btree = _nested_vector_to_digraph(bipartite_sequence(inds_leaves))
+end
+# @info "tn", tn
+# @info "inds_leaves is", inds_leaves
+for alg in ["ttn_svd", "density_matrix"]
+  @info "alg is", alg
+  reset_timer!(ITensors.timer)
+  out = @time bench(tn, btree; alg=alg, maxdim=maxdim)
+  @info "out norm is", out[2]
+  show(ITensors.timer)
+  @info ""
+end

examples/approx_itensornetwork/bench_peps.jl

@@ -0,0 +1,48 @@
+using ITensors, Graphs, NamedGraphs, TimerOutputs
+using Distributions, Random
+using KaHyPar, Metis
+using ITensorNetworks
+using ITensorNetworks: ising_network, contract, _nested_vector_to_digraph, _ansatz_tree
+using OMEinsumContractionOrders
+using AbstractTrees
+
+include("utils.jl")
+
+function peps(N; link_space)
+  inds_net = IndsNetwork(named_grid(N); link_space=link_space)
+  inds_leaves = []
+  for i in 1:N[1]
+    push!(inds_leaves, [Index(link_space, "$(i)x$(j),out") for j in 1:N[2]])
+  end
+  for i in 1:N[1]
+    for j in 1:N[2]
+      inds_net[(i, j)] = [inds_leaves[i][j]]
+    end
+  end
+  distribution = Uniform{Float64}(-1.0, 1.0)
+  return randomITensorNetwork(distribution, inds_net), inds_leaves
+end
+
+Random.seed!(1234)
+TimerOutputs.enable_debug_timings(ITensorNetworks)
+# TimerOutputs.enable_debug_timings(@__MODULE__)
+ITensors.set_warn_order(100)
+
+N = (10, 10) # (9, 9) is the largest for comb
+link_space = 2
+maxdim = 100
+ansatz = "mps"
+
+tn, inds_leaves = peps(N; link_space=link_space)
+ortho_center = div(length(inds_leaves), 2, RoundDown)
+btree = _ansatz_tree(inds_leaves, ansatz, ortho_center)
+@info "tn", tn
+@info "inds_leaves is", inds_leaves
+for alg in ["density_matrix", "ttn_svd"]
+  @info "alg is", alg
+  reset_timer!(ITensors.timer)
+  out = @time bench(tn, btree; alg=alg, maxdim=maxdim)
+  @info "out norm is", out[2]
+  show(ITensors.timer)
+  @info ""
+end

examples/approx_itensornetwork/bench_tree_truncate.jl

@@ -0,0 +1,76 @@
+using ITensors, Graphs, NamedGraphs, TimerOutputs
+using Distributions, Random
+using KaHyPar, Metis
+using ITensorNetworks
+using ITensorNetworks: ising_network, contract, _nested_vector_to_digraph
+using OMEinsumContractionOrders
+using AbstractTrees
+
+include("utils.jl")
+
+function balanced_binary_tree_tn(num_leaves; link_space, physical_spaces)
+  sequence = bipartite_sequence(collect(1:num_leaves))
+  g = NamedGraph()
+  for v in PostOrderDFS(sequence)
+    add_vertex!(g, collect(Leaves(v)))
+    if v isa Vector
+      for c in v
+        add_edge!(g, collect(Leaves(c)) => collect(Leaves(v)))
+      end
+    end
+  end
+  inds_net = IndsNetwork(g; link_space=link_space)
+  inds_leaves = [Index(physical_spaces[i], "$(i)") for i in 1:num_leaves]
+  for i in 1:num_leaves
+    inds_net[[i]] = [inds_leaves[i]]
+  end
+  distribution = Uniform{Float64}(-1.0, 1.0)
+  return randomITensorNetwork(distribution, inds_net), inds_leaves
+end
+
+function unbalanced_binary_tree_tn(num_leaves; link_space, physical_spaces)
+  inds_net = IndsNetwork(named_grid((num_leaves)); link_space=link_space)
+  inds_leaves = [Index(physical_spaces[i], "$(i)") for i in 1:num_leaves]
+  for i in 1:num_leaves
+    inds_net[i] = [inds_leaves[i]]
+  end
+  distribution = Uniform{Float64}(-1.0, 1.0)
+  return randomITensorNetwork(distribution, inds_net), inds_leaves
+end
+
+Random.seed!(1234)
+TimerOutputs.enable_debug_timings(ITensorNetworks)
+# TimerOutputs.enable_debug_timings(@__MODULE__)
+ITensors.set_warn_order(100)
+
+# balanced binary tree
+# n = 16
+# link_space = 200
+# physical_spaces = [200 for i in 1:16]
+# maxdim = 100
+# tn, inds_leaves = balanced_binary_tree_tn(
+#   n; link_space=link_space, physical_spaces=physical_spaces
+# )
+# btree = _nested_vector_to_digraph(bipartite_sequence(inds_leaves))
+
+# unbalanced binary tree
+n = 16
+link_space = 500
+physical_spaces = [2 for i in 1:16]
+physical_spaces[1] = link_space
+maxdim = 200
+tn, inds_leaves = unbalanced_binary_tree_tn(
+  n; link_space=link_space, physical_spaces=physical_spaces
+)
+btree = _nested_vector_to_digraph(linear_sequence(inds_leaves))
+
+@info inds_leaves
+# @info tn
+for alg in ["ttn_svd", "density_matrix", "density_matrix_direct_eigen"]
+  @info "alg is", alg
+  reset_timer!(ITensors.timer)
+  out = bench(tn, btree; alg=alg, maxdim=maxdim)
+  @info "out norm is", out[2]
+  show(ITensors.timer)
+  @info ""
+end

examples/approx_itensornetwork/utils.jl

@@ -0,0 +1,31 @@
+function bipartite_sequence(vec::Vector)
+  if length(vec) == 1
+    return vec[1]
+  end
+  if length(vec) == 2
+    return vec
+  end
+  middle = floor(Int64, length(vec) / 2)
+  return [bipartite_sequence(vec[1:middle]), bipartite_sequence(vec[(middle + 1):end])]
+end
+
+function linear_sequence(vec::Vector)
+  if length(vec) <= 2
+    return vec
+  end
+  return [linear_sequence(vec[1:(end - 1)]), vec[end]]
+end
+
+function bench(tn, btree; alg, maxdim)
+  @timeit ITensors.timer "approx_itensornetwork" begin
+    return approx_itensornetwork(
+      tn,
+      btree;
+      alg,
+      cutoff=1e-15,
+      maxdim=maxdim,
+      contraction_sequence_alg="sa_bipartite",
+      contraction_sequence_kwargs=(;),
+    )
+  end
+end

examples/partitioned_contract/3d_cube.jl

@@ -0,0 +1,102 @@
+using ITensorNetworks: vertex_tag
+
+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
+  ranges = [1:i for i in n]
+  for coord in Iterators.product(ranges...)
+    @info coord
+    block_coord_floor = [(i - 1) * s for (i, s) in zip(coord, env_size)]
+    block_coord_ceil = [min(s + f, n) for (s, f, n) in zip(env_size, block_coord_floor, N)]
+    sub_tn = tn[collect(
+      (f + 1):c for (f, c) in zip(block_coord_floor, block_coord_ceil)
+    )...]
+    sub_tn = vec(sub_tn)
+    if tntree == nothing
+      tntree = sub_tn
+    else
+      tntree = [tntree, sub_tn]
+    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
+
+function build_tntree(N, network::ITensorNetwork; block_size::Tuple, env_size::Tuple)
+  @assert length(block_size) == length(env_size)
+  order = length(block_size)
+  tn = Array{ITensor,length(N)}(undef, N...)
+  for v in vertices(network)
+    tn[v...] = network[v...]
+  end
+  if block_size == Tuple(1 for _ in 1:order)
+    return build_tntree(tn, N; env_size=env_size)
+  end
+  tn_reduced = ITensorNetwork()
+  reduced_N = [ceil(Int, x / y) for (x, y) in zip(N, block_size)]
+  ranges = [1:n for n in reduced_N]
+  for coord in Iterators.product(ranges...)
+    add_vertex!(tn_reduced, coord)
+    block_coord_floor = [(i - 1) * s for (i, s) in zip(coord, block_size)]
+    block_coord_ceil = [
+      min(s + f, n) for (s, f, n) in zip(block_size, block_coord_floor, N)
+    ]
+    tn_reduced[coord] = ITensors.contract(
+      tn[collect((f + 1):c for (f, c) in zip(block_coord_floor, block_coord_ceil))...]...
+    )
+  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,length(reduced_N)}(undef, reduced_N...)
+  for v in vertices(tn_reduced)
+    network_reduced[v...] = tn_reduced[v...]
+  end
+  reduced_env = [ceil(Int, x / y) for (x, y) in zip(env_size, block_size)]
+  return build_tntree(network_reduced, reduced_N; env_size=reduced_env)
+end