From b256f54c2f9f7e9b0951745c02adc71651e01703 Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Wed, 21 Jun 2023 22:54:38 -0500 Subject: [PATCH 01/35] Introduce contraction_sequence_to_directed_graph and simplify contraction_sequence_to_graph --- src/contraction_tree_to_graph.jl | 73 ++++++++++------------ test/test_contraction_sequence_to_graph.jl | 15 +++-- 2 files changed, 44 insertions(+), 44 deletions(-) diff --git a/src/contraction_tree_to_graph.jl b/src/contraction_tree_to_graph.jl index 94e74939..8c79bbda 100644 --- a/src/contraction_tree_to_graph.jl +++ b/src/contraction_tree_to_graph.jl @@ -1,3 +1,29 @@ +""" +Take a contraction sequence and return a directed graph. +""" +function contraction_sequence_to_directed_graph(contract_sequence) + g = NamedDiGraph() + leaves = collect(Leaves(contract_sequence)) + seq_to_v = Dict() + for seq in PostOrderDFS(contract_sequence) + if !(seq isa Array) + v = ([seq], setdiff(leaves, [seq])) + add_vertex!(g, v) + else + group1 = collect(Leaves(seq[1])) + group2 = collect(Leaves(seq[2])) + remaining_verts = setdiff(leaves, vcat(group1, group2)) + v = (group1, group2, remaining_verts) + add_vertex!(g, v) + c1 = get(seq_to_v, seq[1], nothing) + c2 = get(seq_to_v, seq[2], nothing) + add_edge!(g, v => c1) + add_edge!(g, v => c2) + end + seq_to_v[seq] = v + end + return g +end """ Take a contraction_sequence and return a graphical representation of it. The leaves of the graph represent the leaves of the sequence whilst the internal_nodes of the graph @@ -5,49 +31,18 @@ define a tripartition of the graph and thus are named as an n = 3 element tuples Edges connect parents/children within the contraction sequence. """ function contraction_sequence_to_graph(contract_sequence) - g = fill_contraction_sequence_graph_vertices(contract_sequence) - - #Now we have the vertices we need to figure out the edges - for v in vertices(g) - #Only add edges from a parent (which defines a tripartition and thus has length 3) to its children - if (length(v) == 3) - #Work out which vertices it connects to - concat1, concat2, concat3 = [v[1]..., v[2]...], [v[2]..., v[3]...], [v[1]..., v[3]...] - for vn in setdiff(vertices(g), [v]) - vn_set = [Set(vni) for vni in vn] - if (Set(concat1) ∈ vn_set || Set(concat2) ∈ vn_set || Set(concat3) ∈ vn_set) - add_edge!(g, v => vn) - end - end - end + direct_g = contraction_sequence_to_directed_graph(contract_sequence) + g = NamedGraph(vertices(direct_g)) + for e in edges(direct_g) + add_edge!(g, e) end - - return g -end - -function fill_contraction_sequence_graph_vertices(contract_sequence) - g = NamedGraph() - leaves = collect(Leaves(contract_sequence)) - fill_contraction_sequence_graph_vertices!(g, contract_sequence[1], leaves) - fill_contraction_sequence_graph_vertices!(g, contract_sequence[2], leaves) + root = _root(direct_g) + c1, c2 = child_vertices(direct_g, root) + rem_vertex!(g, root) + add_edge!(g, c1 => c2) return g end -"""Given a contraction sequence which is a subsequence of some larger sequence (with leaves `leaves`) which is being built on g -Spawn `contract sequence' as a vertex on `current_g' and continue on with its children """ -function fill_contraction_sequence_graph_vertices!(g, contract_sequence, leaves) - if (isa(contract_sequence, Array)) - group1 = collect(Leaves(contract_sequence[1])) - group2 = collect(Leaves(contract_sequence[2])) - remaining_verts = setdiff(leaves, vcat(group1, group2)) - add_vertex!(g, (group1, group2, remaining_verts)) - fill_contraction_sequence_graph_vertices!(g, contract_sequence[1], leaves) - fill_contraction_sequence_graph_vertices!(g, contract_sequence[2], leaves) - else - add_vertex!(g, ([contract_sequence], setdiff(leaves, [contract_sequence]))) - end -end - """Get the vertex bi-partition that a given edge represents""" function contraction_tree_leaf_bipartition(g::AbstractGraph, e) if (!is_leaf_edge(g, e)) diff --git a/test/test_contraction_sequence_to_graph.jl b/test/test_contraction_sequence_to_graph.jl index 5fd6bc7d..2f2522de 100644 --- a/test/test_contraction_sequence_to_graph.jl +++ b/test/test_contraction_sequence_to_graph.jl @@ -1,10 +1,12 @@ using ITensorNetworks using ITensorNetworks: + contraction_sequence_to_directed_graph, contraction_sequence_to_graph, internal_edges, contraction_tree_leaf_bipartition, distance_to_leaf, - leaf_vertices + leaf_vertices, + _root using Test using ITensors using NamedGraphs @@ -20,12 +22,15 @@ using NamedGraphs seq = contraction_sequence(ψψ) - g_seq = contraction_sequence_to_graph(seq) - - #Get all leaf nodes (should match number of tensors in original network) - g_seq_leaves = leaf_vertices(g_seq) + g_directed_seq = contraction_sequence_to_directed_graph(seq) + g_seq_leaves = leaf_vertices(g_directed_seq) + @test length(g_seq_leaves) == n * n + @test 2 * length(g_seq_leaves) - 1 == length(vertices(g_directed_seq)) + @test _root(g_directed_seq)[3] == [] + g_seq = contraction_sequence_to_graph(seq) @test length(g_seq_leaves) == n * n + @test 2 * length(g_seq_leaves) - 2 == length(vertices(g_seq)) for eb in internal_edges(g_seq) vs = contraction_tree_leaf_bipartition(g_seq, eb) From 1e9794831b6d35f8c94ed742fa74e53c5b5cc829 Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Thu, 22 Jun 2023 16:55:20 -0500 Subject: [PATCH 02/35] Introduce approx_itensornetwork with algorithm ttn_svd --- src/ITensorNetworks.jl | 5 +- .../approx_itensornetwork.jl | 178 ++++++++++++++++ .../density_matrix.jl} | 201 +----------------- src/approx_itensornetwork/ttn_svd.jl | 31 +++ src/approx_itensornetwork/utils.jl | 43 ++++ src/contract.jl | 2 +- test/test_binary_tree_partition.jl | 23 +- 7 files changed, 269 insertions(+), 214 deletions(-) create mode 100644 src/approx_itensornetwork/approx_itensornetwork.jl rename src/{approx_itensornetwork.jl => approx_itensornetwork/density_matrix.jl} (65%) create mode 100644 src/approx_itensornetwork/ttn_svd.jl create mode 100644 src/approx_itensornetwork/utils.jl diff --git a/src/ITensorNetworks.jl b/src/ITensorNetworks.jl index 3fe84083..30d1d6b2 100644 --- a/src/ITensorNetworks.jl +++ b/src/ITensorNetworks.jl @@ -84,7 +84,10 @@ include("itensornetwork.jl") include("mincut.jl") include("contract_deltas.jl") include("binary_tree_partition.jl") -include("approx_itensornetwork.jl") +include(joinpath("approx_itensornetwork", "utils.jl")) +include(joinpath("approx_itensornetwork", "density_matrix.jl")) +include(joinpath("approx_itensornetwork", "ttn_svd.jl")) +include(joinpath("approx_itensornetwork", "approx_itensornetwork.jl")) include("contract.jl") include("utility.jl") include("specialitensornetworks.jl") diff --git a/src/approx_itensornetwork/approx_itensornetwork.jl b/src/approx_itensornetwork/approx_itensornetwork.jl new file mode 100644 index 00000000..5041a8a5 --- /dev/null +++ b/src/approx_itensornetwork/approx_itensornetwork.jl @@ -0,0 +1,178 @@ +# Density matrix algorithm and ttn_svd algorithm +""" +Approximate a `binary_tree_partition` into an output ITensorNetwork +with the same binary tree structure. `root` is the root vertex of the +pre-order depth-first-search traversal used to perform the truncations. +""" +function approx_itensornetwork( + ::Algorithm"density_matrix", + binary_tree_partition::DataGraph; + root, + cutoff=1e-15, + maxdim=10000, + contraction_sequence_alg, + contraction_sequence_kwargs, +) + @assert is_tree(binary_tree_partition) + @assert root in vertices(binary_tree_partition) + @assert _is_rooted_directed_binary_tree(dfs_tree(binary_tree_partition, root)) + # The `binary_tree_partition` may contain multiple delta tensors to make sure + # the partition has a binary tree structure. These delta tensors could hurt the + # performance when computing density matrices so we remove them first. + partition_wo_deltas = _contract_deltas_ignore_leaf_partitions( + binary_tree_partition; root=root + ) + return _approx_itensornetwork_density_matrix!( + partition_wo_deltas, + underlying_graph(binary_tree_partition); + root, + cutoff, + maxdim, + contraction_sequence_alg, + contraction_sequence_kwargs, + ) +end + +function approx_itensornetwork( + ::Algorithm"ttn_svd", + binary_tree_partition::DataGraph; + root, + cutoff=1e-15, + maxdim=10000, + contraction_sequence_alg, + contraction_sequence_kwargs, +) + @assert is_tree(binary_tree_partition) + @assert root in vertices(binary_tree_partition) + @assert _is_rooted_directed_binary_tree(dfs_tree(binary_tree_partition, root)) + return _approx_itensornetwork_ttn_svd!( + binary_tree_partition; + root, + cutoff, + maxdim, + contraction_sequence_alg, + contraction_sequence_kwargs, + ) +end + +""" +Approximate a given ITensorNetwork `tn` into an output ITensorNetwork +with a binary tree structure. The binary tree structure is defined based +on `inds_btree`, which is a directed binary tree DataGraph of indices. +""" +function approx_itensornetwork( + alg::Union{Algorithm"density_matrix",Algorithm"ttn_svd"}, + tn::ITensorNetwork, + inds_btree::DataGraph; + cutoff=1e-15, + maxdim=10000, + contraction_sequence_alg="optimal", + contraction_sequence_kwargs=(;), +) + par = partition(tn, inds_btree; alg="mincut_recursive_bisection") + output_tn, log_root_norm = approx_itensornetwork( + alg, + par; + root=_root(inds_btree), + cutoff=cutoff, + maxdim=maxdim, + contraction_sequence_alg=contraction_sequence_alg, + contraction_sequence_kwargs=contraction_sequence_kwargs, + ) + # Each leaf vertex in `output_tn` is adjacent to one output index. + # We remove these leaf vertices so that each non-root vertex in `output_tn` + # is an order 3 tensor. + _rem_leaf_vertices!( + output_tn; + root=_root(inds_btree), + contraction_sequence_alg=contraction_sequence_alg, + contraction_sequence_kwargs=contraction_sequence_kwargs, + ) + return output_tn, log_root_norm +end + +""" +Approximate a given ITensorNetwork `tn` into an output ITensorNetwork with `output_structure`. +`output_structure` outputs a directed binary tree DataGraph defining the desired graph structure. +""" +function approx_itensornetwork( + alg::Union{Algorithm"density_matrix",Algorithm"ttn_svd"}, + tn::ITensorNetwork, + output_structure::Function=path_graph_structure; + cutoff=1e-15, + maxdim=10000, + contraction_sequence_alg="optimal", + contraction_sequence_kwargs=(;), +) + inds_btree = output_structure(tn) + return approx_itensornetwork( + alg, + tn, + inds_btree; + cutoff=cutoff, + maxdim=maxdim, + contraction_sequence_alg=contraction_sequence_alg, + contraction_sequence_kwargs=contraction_sequence_kwargs, + ) +end + +# interface +function approx_itensornetwork( + partitioned_tn::DataGraph; + alg::String, + root, + cutoff=1e-15, + maxdim=10000, + contraction_sequence_alg="optimal", + contraction_sequence_kwargs=(;), +) + return approx_itensornetwork( + Algorithm(alg), + partitioned_tn; + root, + cutoff, + maxdim, + contraction_sequence_alg, + contraction_sequence_kwargs, + ) +end + +function approx_itensornetwork( + tn::ITensorNetwork, + inds_btree::DataGraph; + alg::String, + cutoff=1e-15, + maxdim=10000, + contraction_sequence_alg="optimal", + contraction_sequence_kwargs=(;), +) + return approx_itensornetwork( + Algorithm(alg), + tn, + inds_btree; + cutoff, + maxdim, + contraction_sequence_alg, + contraction_sequence_kwargs, + ) +end + +function approx_itensornetwork( + tn::ITensorNetwork, + output_structure::Function=path_graph_structure; + alg::String, + cutoff=1e-15, + maxdim=10000, + contraction_sequence_alg="optimal", + contraction_sequence_kwargs=(;), +) + return approx_itensornetwork( + Algorithm(alg), + tn, + output_structure; + cutoff, + maxdim, + contraction_sequence_alg, + contraction_sequence_kwargs, + ) +end diff --git a/src/approx_itensornetwork.jl b/src/approx_itensornetwork/density_matrix.jl similarity index 65% rename from src/approx_itensornetwork.jl rename to src/approx_itensornetwork/density_matrix.jl index 5cccbc66..e5009067 100644 --- a/src/approx_itensornetwork.jl +++ b/src/approx_itensornetwork/density_matrix.jl @@ -142,27 +142,6 @@ function _DensityMartrixAlgGraph(partition::DataGraph, out_tree::NamedGraph, roo ) end -""" -Contract of a vector of tensors, `network`, with a contraction sequence generated via sa_bipartite -""" -function _optcontract( - network::Vector; contraction_sequence_alg="optimal", contraction_sequence_kwargs=(;) -) - @timeit_debug ITensors.timer "[approx_binary_tree_itensornetwork]: _optcontract" begin - if length(network) == 0 - return ITensor(1.0) - end - @assert network isa Vector{ITensor} - @timeit_debug ITensors.timer "[approx_binary_tree_itensornetwork]: contraction_sequence" begin - seq = contraction_sequence( - network; alg=contraction_sequence_alg, contraction_sequence_kwargs... - ) - end - output = contract(network; sequence=seq) - return output - end -end - function _get_low_rank_projector(tensor, inds1, inds2; cutoff, maxdim) @assert length(inds(tensor)) <= 4 @timeit_debug ITensors.timer "[approx_binary_tree_itensornetwork]: eigen" begin @@ -356,34 +335,11 @@ function _rem_vertex!( return U end -""" -For a given ITensorNetwork `tn` and a `root` vertex, remove leaf vertices in the directed tree -with root `root` without changing the tensor represented by tn. -In particular, the tensor of each leaf vertex is contracted with the tensor of its parent vertex -to keep the tensor unchanged. -""" -function _rem_leaf_vertices!( - tn::ITensorNetwork; - root=first(vertices(tn)), - contraction_sequence_alg, - contraction_sequence_kwargs, -) - dfs_t = dfs_tree(tn, root) - leaves = leaf_vertices(dfs_t) - parents = [parent_vertex(dfs_t, leaf) for leaf in leaves] - for (l, p) in zip(leaves, parents) - tn[p] = _optcontract( - [tn[p], tn[l]]; contraction_sequence_alg, contraction_sequence_kwargs - ) - rem_vertex!(tn, l) - end -end - """ Approximate a `partition` into an output ITensorNetwork with the binary tree structure defined by `out_tree`. """ -function _approx_binary_tree_itensornetwork!( +function _approx_itensornetwork_density_matrix!( input_partition::DataGraph, out_tree::NamedGraph; root=first(vertices(partition)), @@ -419,158 +375,3 @@ function _approx_binary_tree_itensornetwork!( output_tn[root] = root_tensor return output_tn, log(root_norm) end - -""" -Approximate a `binary_tree_partition` into an output ITensorNetwork -with the same binary tree structure. `root` is the root vertex of the -pre-order depth-first-search traversal used to perform the truncations. -""" -function approx_itensornetwork( - ::Algorithm"density_matrix", - binary_tree_partition::DataGraph; - root, - cutoff=1e-15, - maxdim=10000, - contraction_sequence_alg, - contraction_sequence_kwargs, -) - @assert is_tree(binary_tree_partition) - @assert root in vertices(binary_tree_partition) - @assert _is_rooted_directed_binary_tree(dfs_tree(binary_tree_partition, root)) - # The `binary_tree_partition` may contain multiple delta tensors to make sure - # the partition has a binary tree structure. These delta tensors could hurt the - # performance when computing density matrices so we remove them first. - partition_wo_deltas = _contract_deltas_ignore_leaf_partitions( - binary_tree_partition; root=root - ) - return _approx_binary_tree_itensornetwork!( - partition_wo_deltas, - underlying_graph(binary_tree_partition); - root, - cutoff, - maxdim, - contraction_sequence_alg, - contraction_sequence_kwargs, - ) -end - -""" -Approximate a given ITensorNetwork `tn` into an output ITensorNetwork with `output_structure`. -`output_structure` outputs a directed binary tree DataGraph defining the desired graph structure. -""" -function approx_itensornetwork( - ::Algorithm"density_matrix", - tn::ITensorNetwork, - output_structure::Function=path_graph_structure; - cutoff=1e-15, - maxdim=10000, - contraction_sequence_alg="optimal", - contraction_sequence_kwargs=(;), -) - inds_btree = output_structure(tn) - return approx_itensornetwork( - tn, - inds_btree; - alg="density_matrix", - cutoff=cutoff, - maxdim=maxdim, - contraction_sequence_alg=contraction_sequence_alg, - contraction_sequence_kwargs=contraction_sequence_kwargs, - ) -end - -""" -Approximate a given ITensorNetwork `tn` into an output ITensorNetwork -with a binary tree structure. The binary tree structure is defined based -on `inds_btree`, which is a directed binary tree DataGraph of indices. -""" -function approx_itensornetwork( - ::Algorithm"density_matrix", - tn::ITensorNetwork, - inds_btree::DataGraph; - cutoff=1e-15, - maxdim=10000, - contraction_sequence_alg="optimal", - contraction_sequence_kwargs=(;), -) - par = partition(tn, inds_btree; alg="mincut_recursive_bisection") - output_tn, log_root_norm = approx_itensornetwork( - par; - alg="density_matrix", - root=_root(inds_btree), - cutoff=cutoff, - maxdim=maxdim, - contraction_sequence_alg=contraction_sequence_alg, - contraction_sequence_kwargs=contraction_sequence_kwargs, - ) - # Each leaf vertex in `output_tn` is adjacent to one output index. - # We remove these leaf vertices so that each non-root vertex in `output_tn` - # is an order 3 tensor. - _rem_leaf_vertices!( - output_tn; - root=_root(inds_btree), - contraction_sequence_alg=contraction_sequence_alg, - contraction_sequence_kwargs=contraction_sequence_kwargs, - ) - return output_tn, log_root_norm -end - -function approx_itensornetwork( - partitioned_tn::DataGraph; - alg::String, - root, - cutoff=1e-15, - maxdim=10000, - contraction_sequence_alg="optimal", - contraction_sequence_kwargs=(;), -) - return approx_itensornetwork( - Algorithm(alg), - partitioned_tn; - root, - cutoff, - maxdim, - contraction_sequence_alg, - contraction_sequence_kwargs, - ) -end - -function approx_itensornetwork( - tn::ITensorNetwork, - inds_btree::DataGraph; - alg::String, - cutoff=1e-15, - maxdim=10000, - contraction_sequence_alg="optimal", - contraction_sequence_kwargs=(;), -) - return approx_itensornetwork( - Algorithm(alg), - tn, - inds_btree; - cutoff, - maxdim, - contraction_sequence_alg, - contraction_sequence_kwargs, - ) -end - -function approx_itensornetwork( - tn::ITensorNetwork, - output_structure::Function=path_graph_structure; - alg::String, - cutoff=1e-15, - maxdim=10000, - contraction_sequence_alg="optimal", - contraction_sequence_kwargs=(;), -) - return approx_itensornetwork( - Algorithm(alg), - tn, - output_structure; - cutoff, - maxdim, - contraction_sequence_alg, - contraction_sequence_kwargs, - ) -end diff --git a/src/approx_itensornetwork/ttn_svd.jl b/src/approx_itensornetwork/ttn_svd.jl new file mode 100644 index 00000000..5ccc6019 --- /dev/null +++ b/src/approx_itensornetwork/ttn_svd.jl @@ -0,0 +1,31 @@ +""" +Approximate a `partition` into an output ITensorNetwork +with the binary tree structure defined by `out_tree` by +first transforming the partition into a TTN, then truncating +the ttn using a sequence of SVDs. +""" +function _approx_itensornetwork_ttn_svd!( + input_partition::DataGraph; + root=first(vertices(partition)), + cutoff=1e-15, + maxdim=10000, + contraction_sequence_alg, + contraction_sequence_kwargs, +) + tn = ITensorNetwork() + for v in vertices(input_partition) + add_vertex!(tn, v) + tn[v] = _optcontract( + Vector{ITensor}(input_partition[v]); + contraction_sequence_alg=contraction_sequence_alg, + contraction_sequence_kwargs=contraction_sequence_kwargs, + ) + end + truncate_ttn = truncate(TTN(tn); cutoff=cutoff, maxdim=maxdim, root_vertex=root) + out_tn = ITensorNetwork(truncate_ttn) + root_tensor = out_tn[root] + root_norm = norm(root_tensor) + root_tensor /= root_norm + out_tn[root] = root_tensor + return out_tn, log(root_norm) +end diff --git a/src/approx_itensornetwork/utils.jl b/src/approx_itensornetwork/utils.jl new file mode 100644 index 00000000..e650ac49 --- /dev/null +++ b/src/approx_itensornetwork/utils.jl @@ -0,0 +1,43 @@ +""" +For a given ITensorNetwork `tn` and a `root` vertex, remove leaf vertices in the directed tree +with root `root` without changing the tensor represented by tn. +In particular, the tensor of each leaf vertex is contracted with the tensor of its parent vertex +to keep the tensor unchanged. +""" +function _rem_leaf_vertices!( + tn::ITensorNetwork; + root=first(vertices(tn)), + contraction_sequence_alg, + contraction_sequence_kwargs, +) + dfs_t = dfs_tree(tn, root) + leaves = leaf_vertices(dfs_t) + parents = [parent_vertex(dfs_t, leaf) for leaf in leaves] + for (l, p) in zip(leaves, parents) + tn[p] = _optcontract( + [tn[p], tn[l]]; contraction_sequence_alg, contraction_sequence_kwargs + ) + rem_vertex!(tn, l) + end +end + +""" +Contract of a vector of tensors, `network`, with a contraction sequence generated via sa_bipartite +""" +function _optcontract( + network::Vector; contraction_sequence_alg="optimal", contraction_sequence_kwargs=(;) +) + @timeit_debug ITensors.timer "[approx_binary_tree_itensornetwork]: _optcontract" begin + if length(network) == 0 + return ITensor(1.0) + end + @assert network isa Vector{ITensor} + @timeit_debug ITensors.timer "[approx_binary_tree_itensornetwork]: contraction_sequence" begin + seq = contraction_sequence( + network; alg=contraction_sequence_alg, contraction_sequence_kwargs... + ) + end + output = contract(network; sequence=seq) + return output + end +end diff --git a/src/contract.jl b/src/contract.jl index 4d1d5f9e..97265d77 100644 --- a/src/contract.jl +++ b/src/contract.jl @@ -10,7 +10,7 @@ function contract( end function contract( - alg::Algorithm"density_matrix", + alg::Union{Algorithm"density_matrix",Algorithm"ttn_svd"}, tn::AbstractITensorNetwork; output_structure::Function=path_graph_structure, kwargs..., diff --git a/test/test_binary_tree_partition.jl b/test/test_binary_tree_partition.jl index f00ef278..32942b4a 100644 --- a/test/test_binary_tree_partition.jl +++ b/test/test_binary_tree_partition.jl @@ -88,18 +88,17 @@ end @test isapprox(out1, out2) # test approx_itensornetwork (here we call `contract` to test the interface) for structure in [path_graph_structure, binary_tree_structure] - approx_tn, lognorm = contract( - tn; - alg="density_matrix", - output_structure=structure, - contraction_sequence_alg="sa_bipartite", - ) - network3 = Vector{ITensor}(approx_tn) - out3 = contract(network3...) * exp(lognorm) - i1 = noncommoninds(network...) - i3 = noncommoninds(network3...) - @test (length(i1) == length(i3)) - @test isapprox(out1, out3) + for alg in ["density_matrix", "ttn_svd"] + approx_tn, lognorm = contract( + tn; alg=alg, output_structure=structure, contraction_sequence_alg="sa_bipartite" + ) + network3 = Vector{ITensor}(approx_tn) + out3 = contract(network3...) * exp(lognorm) + i1 = noncommoninds(network...) + i3 = noncommoninds(network3...) + @test (length(i1) == length(i3)) + @test isapprox(out1, out3) + end end end end From 5505385fc3ad1a924f3190eda55bdbfc941c6e0d Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Thu, 22 Jun 2023 21:23:54 -0500 Subject: [PATCH 03/35] save changes --- src/ITensorNetworks.jl | 1 + src/contract.jl | 9 ++ src/partitioned_contract/adjacency_tree.jl | 0 .../partitioned_contract.jl | 144 ++++++++++++++++++ test/runtests.jl | 2 +- 5 files changed, 155 insertions(+), 1 deletion(-) create mode 100644 src/partitioned_contract/adjacency_tree.jl create mode 100644 src/partitioned_contract/partitioned_contract.jl diff --git a/src/ITensorNetworks.jl b/src/ITensorNetworks.jl index 30d1d6b2..e6409f13 100644 --- a/src/ITensorNetworks.jl +++ b/src/ITensorNetworks.jl @@ -87,6 +87,7 @@ include("binary_tree_partition.jl") include(joinpath("approx_itensornetwork", "utils.jl")) include(joinpath("approx_itensornetwork", "density_matrix.jl")) include(joinpath("approx_itensornetwork", "ttn_svd.jl")) +include(joinpath("approx_itensornetwork", "partitioned_contraction.jl")) include(joinpath("approx_itensornetwork", "approx_itensornetwork.jl")) include("contract.jl") include("utility.jl") diff --git a/src/contract.jl b/src/contract.jl index 97265d77..0b82b7fa 100644 --- a/src/contract.jl +++ b/src/contract.jl @@ -17,3 +17,12 @@ function contract( ) return approx_itensornetwork(alg, tn, output_structure; kwargs...) end + +function contract( + alg::Algorithm"partitioned_contraction", + tn::AbstractITensorNetwork; + output_structure::Function=path_graph_structure, + kwargs..., +) + return approx_itensornetwork(alg, tn, output_structure; kwargs...) +end diff --git a/src/partitioned_contract/adjacency_tree.jl b/src/partitioned_contract/adjacency_tree.jl new file mode 100644 index 00000000..e69de29b diff --git a/src/partitioned_contract/partitioned_contract.jl b/src/partitioned_contract/partitioned_contract.jl new file mode 100644 index 00000000..b9da8e0c --- /dev/null +++ b/src/partitioned_contract/partitioned_contract.jl @@ -0,0 +1,144 @@ +function partitioned_contract( + partition::DataGraph, + contraction_tree::NamedDiGraph; + ansatz="mps", + approx_itensornetwork_alg="density_matrix", + cutoff=1e-15, + maxdim=10000, + contraction_sequence_alg, + contraction_sequence_kwargs, +) + @timeit_debug ITensors.timer "partitioned_contract" begin + tn_leaves = get_leaves(ctree) + ctrees = topo_sort(ctree; leaves=tn_leaves) + @info "start _approximate_contract_pre_process" + ctree_to_igs, ctree_to_adj_tree, ig_to_linear_order = _approximate_contract_pre_process( + tn_leaves, ctrees + ) + ctree_to_reference_order = Dict{Vector,Vector}() + ctree_to_edgeset_order = Dict{Vector,Vector}() + for leaf in tn_leaves + reference_order = _mps_partition_inds_set_order( + ITensorNetwork(leaf), ctree_to_igs[leaf] + ) + order, _ = mindist_ordering(ctree_to_adj_tree[leaf], reference_order, vectorize(leaf)) + ctree_to_edgeset_order[leaf] = order + end + for (ii, c) in enumerate(ctrees) + @info "$(ii)/$(length(ctrees))", "th pre-process" + if ctree_to_igs[c] == [] + continue + end + ctree_to_reference_order[c], ctree_to_edgeset_order[c] = mindist_ordering( + ctree_to_adj_tree[c], + ctree_to_edgeset_order[c[1]], + ctree_to_edgeset_order[c[2]], + c[1] in tn_leaves, + c[2] in tn_leaves, + vectorize(c), + ) + # @info "reference ordering", ctree_to_reference_order[c] + # @info "edge set ordering", ctree_to_edgeset_order[c] + end + # return [ITensor(1.0)], 1.0 + # mapping each contraction tree to its contract igs + ctree_to_contract_igs = Dict{Vector,Vector{IndexGroup}}() + for c in ctrees + if c[1] in tn_leaves + contract_igs = intersect(ctree_to_edgeset_order[c[2]], ctree_to_edgeset_order[c[1]]) + l_igs_c1, r_igs_c1 = split_igs(ctree_to_edgeset_order[c[1]], contract_igs) + ctree_to_edgeset_order[c[1]] = Vector{IndexGroup}([ + l_igs_c1..., contract_igs..., r_igs_c1... + ]) + else + contract_igs = intersect(ctree_to_edgeset_order[c[1]], ctree_to_edgeset_order[c[2]]) + l_igs_c2, r_igs_c2 = split_igs(ctree_to_edgeset_order[c[2]], contract_igs) + ctree_to_edgeset_order[c[2]] = Vector{IndexGroup}([ + l_igs_c2..., contract_igs..., r_igs_c2... + ]) + end + 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 + if haskey(ctree_to_edgeset_order, ctrees[end]) + ctree_to_contract_igs[ctrees[end]] = ctree_to_edgeset_order[ctrees[end]] + end + # mapping each contraction tree to a tensor network + ctree_to_tn_tree = Dict{Vector,Union{Dict{Vector,ITensor},Vector{ITensor}}}() + # accumulate norm + log_accumulated_norm = 0.0 + for (ii, c) in enumerate(ctrees) + t00 = time() + @info "$(ii)/$(length(ctrees))", "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) + out = _optcontract(tn) + out_nrm = norm(out) + out /= out_nrm + return [out], log_accumulated_norm + log(out_nrm) + end + tn1 = get_child_tn(ctree_to_tn_tree, c[1]) + tn2 = get_child_tn(ctree_to_tn_tree, c[2]) + # TODO: change new_igs into a vector of igs + inds_btree = ordered_igs_to_binary_tree( + ctree_to_edgeset_order[c], + 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 + # release the memory + delete!(ctree_to_tn_tree, c[1]) + delete!(ctree_to_tn_tree, c[2]) + t11 = time() - t00 + @info "time of this contraction is", t11 + end + tn = vcat(collect(values(ctree_to_tn_tree[ctrees[end]]))...) + return tn, log_accumulated_norm + end +end + +function _approximate_contract_pre_process(tn_leaves, ctrees) + @timeit_debug ITensors.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) + # TODO: the order here is not optimized + ctree_to_igs[c] = neighbor_index_groups(c, index_groups) + end + ctree_to_path = _get_paths(ctrees[end]) + # mapping each contraction tree to its index adjacency tree + ctree_to_adj_tree = Dict{Vector,NamedDiGraph{Tuple{Tuple,String}}}() + for leaf in tn_leaves + ctree_to_adj_tree[leaf] = _generate_adjacency_tree( + leaf, ctree_to_path[leaf], ctree_to_igs + ) + end + for c in ctrees + adj_tree = _generate_adjacency_tree(c, ctree_to_path[c], ctree_to_igs) + if adj_tree != nothing + ctree_to_adj_tree[c] = adj_tree + # @info "ctree_to_adj_tree[c]", ctree_to_adj_tree[c] + end + 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) + inds_order = _mps_partition_inds_order(ITensorNetwork(leaf), ig.data) + ig_to_linear_order[ig] = [[i] for i in inds_order] + end + end + end + return ctree_to_igs, ctree_to_adj_tree, ig_to_linear_order + end +end \ No newline at end of file diff --git a/test/runtests.jl b/test/runtests.jl index f02b8f2b..e6d073a0 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -6,7 +6,7 @@ using ITensorNetworks @testset "ITensorNetworks.jl, test directory $root" for (root, dirs, files) in walkdir( joinpath(pkgdir(ITensorNetworks), "test") ) - test_files = filter!(f -> occursin(Glob.FilenameMatch("test_*.jl"), f), files) + test_files = ["test_binary_tree_partition.jl"]# filter!(f -> occursin(Glob.FilenameMatch("test_*.jl"), f), files) @testset "Test file $test_file" for test_file in test_files println("Running test file $test_file") @time include(joinpath(root, test_file)) From b36d05d07c26a0b59a858f0723b1e80a591e8e7b Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Thu, 22 Jun 2023 21:30:09 -0500 Subject: [PATCH 04/35] remove contraction_sequence_to_directed_graph --- src/contraction_tree_to_graph.jl | 22 +++-------------- test/test_contraction_sequence_to_graph.jl | 28 +++++++--------------- 2 files changed, 12 insertions(+), 38 deletions(-) diff --git a/src/contraction_tree_to_graph.jl b/src/contraction_tree_to_graph.jl index 8c79bbda..2feb1df4 100644 --- a/src/contraction_tree_to_graph.jl +++ b/src/contraction_tree_to_graph.jl @@ -1,7 +1,7 @@ """ Take a contraction sequence and return a directed graph. """ -function contraction_sequence_to_directed_graph(contract_sequence) +function contraction_sequence_to_graph(contract_sequence) g = NamedDiGraph() leaves = collect(Leaves(contract_sequence)) seq_to_v = Dict() @@ -25,24 +25,7 @@ function contraction_sequence_to_directed_graph(contract_sequence) return g end -""" -Take a contraction_sequence and return a graphical representation of it. The leaves of the graph represent the leaves of the sequence whilst the internal_nodes of the graph -define a tripartition of the graph and thus are named as an n = 3 element tuples, which each element specifying the keys involved. -Edges connect parents/children within the contraction sequence. -""" -function contraction_sequence_to_graph(contract_sequence) - direct_g = contraction_sequence_to_directed_graph(contract_sequence) - g = NamedGraph(vertices(direct_g)) - for e in edges(direct_g) - add_edge!(g, e) - end - root = _root(direct_g) - c1, c2 = child_vertices(direct_g, root) - rem_vertex!(g, root) - add_edge!(g, c1 => c2) - return g -end - +#= """Get the vertex bi-partition that a given edge represents""" function contraction_tree_leaf_bipartition(g::AbstractGraph, e) if (!is_leaf_edge(g, e)) @@ -75,3 +58,4 @@ function contraction_tree_leaf_bipartition(g::AbstractGraph, e) return left_bipartition, right_bipartition end +=# diff --git a/test/test_contraction_sequence_to_graph.jl b/test/test_contraction_sequence_to_graph.jl index 2f2522de..0733a492 100644 --- a/test/test_contraction_sequence_to_graph.jl +++ b/test/test_contraction_sequence_to_graph.jl @@ -1,12 +1,6 @@ using ITensorNetworks using ITensorNetworks: - contraction_sequence_to_directed_graph, - contraction_sequence_to_graph, - internal_edges, - contraction_tree_leaf_bipartition, - distance_to_leaf, - leaf_vertices, - _root + contraction_sequence_to_graph, internal_edges, distance_to_leaf, leaf_vertices, _root using Test using ITensors using NamedGraphs @@ -22,21 +16,17 @@ using NamedGraphs seq = contraction_sequence(ψψ) - g_directed_seq = contraction_sequence_to_directed_graph(seq) - g_seq_leaves = leaf_vertices(g_directed_seq) - @test length(g_seq_leaves) == n * n - @test 2 * length(g_seq_leaves) - 1 == length(vertices(g_directed_seq)) - @test _root(g_directed_seq)[3] == [] - g_seq = contraction_sequence_to_graph(seq) + g_seq_leaves = leaf_vertices(g_seq) @test length(g_seq_leaves) == n * n - @test 2 * length(g_seq_leaves) - 2 == length(vertices(g_seq)) + @test 2 * length(g_seq_leaves) - 1 == length(vertices(g_seq)) + @test _root(g_seq)[3] == [] - for eb in internal_edges(g_seq) - vs = contraction_tree_leaf_bipartition(g_seq, eb) - @test length(vs) == 2 - @test Set([v.I for v in vcat(vs[1], vs[2])]) == Set(vertices(ψψ)) - end + # for eb in internal_edges(g_seq) + # vs = contraction_tree_leaf_bipartition(g_seq, eb) + # @test length(vs) == 2 + # @test Set([v.I for v in vcat(vs[1], vs[2])]) == Set(vertices(ψψ)) + # end #Check all internal vertices define a correct tripartition and all leaf vertices define a bipartition (tensor on that leafs vs tensor on rest of tree) for v in vertices(g_seq) if (!is_leaf(g_seq, v)) From e2175375d58fb2907c66288dcf6702b5b3ee01c4 Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Fri, 23 Jun 2023 12:53:03 -0500 Subject: [PATCH 05/35] revert changes --- src/contraction_tree_to_graph.jl | 22 ++++++++++++++--- test/test_contraction_sequence_to_graph.jl | 28 +++++++++++++++------- 2 files changed, 38 insertions(+), 12 deletions(-) diff --git a/src/contraction_tree_to_graph.jl b/src/contraction_tree_to_graph.jl index 2feb1df4..8156e88d 100644 --- a/src/contraction_tree_to_graph.jl +++ b/src/contraction_tree_to_graph.jl @@ -1,7 +1,7 @@ """ Take a contraction sequence and return a directed graph. """ -function contraction_sequence_to_graph(contract_sequence) +function contraction_sequence_to_digraph(contract_sequence) g = NamedDiGraph() leaves = collect(Leaves(contract_sequence)) seq_to_v = Dict() @@ -25,7 +25,24 @@ function contraction_sequence_to_graph(contract_sequence) return g end -#= +""" +Take a contraction_sequence and return a graphical representation of it. The leaves of the graph represent the leaves of the sequence whilst the internal_nodes of the graph +define a tripartition of the graph and thus are named as an n = 3 element tuples, which each element specifying the keys involved. +Edges connect parents/children within the contraction sequence. +""" +function contraction_sequence_to_graph(contract_sequence) + direct_g = contraction_sequence_to_directed_graph(contract_sequence) + g = NamedGraph(vertices(direct_g)) + for e in edges(direct_g) + add_edge!(g, e) + end + root = _root(direct_g) + c1, c2 = child_vertices(direct_g, root) + rem_vertex!(g, root) + add_edge!(g, c1 => c2) + return g +end + """Get the vertex bi-partition that a given edge represents""" function contraction_tree_leaf_bipartition(g::AbstractGraph, e) if (!is_leaf_edge(g, e)) @@ -58,4 +75,3 @@ function contraction_tree_leaf_bipartition(g::AbstractGraph, e) return left_bipartition, right_bipartition end -=# diff --git a/test/test_contraction_sequence_to_graph.jl b/test/test_contraction_sequence_to_graph.jl index 0733a492..caee8993 100644 --- a/test/test_contraction_sequence_to_graph.jl +++ b/test/test_contraction_sequence_to_graph.jl @@ -1,6 +1,12 @@ using ITensorNetworks using ITensorNetworks: - contraction_sequence_to_graph, internal_edges, distance_to_leaf, leaf_vertices, _root + contraction_sequence_to_digraph, + contraction_sequence_to_graph, + internal_edges, + contraction_tree_leaf_bipartition, + distance_to_leaf, + leaf_vertices, + _root using Test using ITensors using NamedGraphs @@ -16,17 +22,21 @@ using NamedGraphs seq = contraction_sequence(ψψ) + g_directed_seq = contraction_sequence_to_digraph(seq) + g_seq_leaves = leaf_vertices(g_directed_seq) + @test length(g_seq_leaves) == n * n + @test 2 * length(g_seq_leaves) - 1 == length(vertices(g_directed_seq)) + @test _root(g_directed_seq)[3] == [] + g_seq = contraction_sequence_to_graph(seq) - g_seq_leaves = leaf_vertices(g_seq) @test length(g_seq_leaves) == n * n - @test 2 * length(g_seq_leaves) - 1 == length(vertices(g_seq)) - @test _root(g_seq)[3] == [] + @test 2 * length(g_seq_leaves) - 2 == length(vertices(g_seq)) - # for eb in internal_edges(g_seq) - # vs = contraction_tree_leaf_bipartition(g_seq, eb) - # @test length(vs) == 2 - # @test Set([v.I for v in vcat(vs[1], vs[2])]) == Set(vertices(ψψ)) - # end + for eb in internal_edges(g_seq) + vs = contraction_tree_leaf_bipartition(g_seq, eb) + @test length(vs) == 2 + @test Set([v.I for v in vcat(vs[1], vs[2])]) == Set(vertices(ψψ)) + end #Check all internal vertices define a correct tripartition and all leaf vertices define a bipartition (tensor on that leafs vs tensor on rest of tree) for v in vertices(g_seq) if (!is_leaf(g_seq, v)) From b875840f9a3905ff8a1156e73332485d48ca458e Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Mon, 26 Jun 2023 12:59:01 -0500 Subject: [PATCH 06/35] save changes --- src/approx_itensornetwork/ttn_svd.jl | 2 + src/partitioned_contract/adjacency_tree.jl | 1 + .../partitioned_contract.jl | 120 +++++++++++++++++- 3 files changed, 119 insertions(+), 4 deletions(-) diff --git a/src/approx_itensornetwork/ttn_svd.jl b/src/approx_itensornetwork/ttn_svd.jl index 5ccc6019..7821f612 100644 --- a/src/approx_itensornetwork/ttn_svd.jl +++ b/src/approx_itensornetwork/ttn_svd.jl @@ -15,6 +15,8 @@ function _approx_itensornetwork_ttn_svd!( tn = ITensorNetwork() for v in vertices(input_partition) add_vertex!(tn, v) + # TODO: use `dense(optcontract(ts))` to convert Diag type to dense + # for QR decomposition TODO: raise an error in ITensors tn[v] = _optcontract( Vector{ITensor}(input_partition[v]); contraction_sequence_alg=contraction_sequence_alg, diff --git a/src/partitioned_contract/adjacency_tree.jl b/src/partitioned_contract/adjacency_tree.jl index e69de29b..8b137891 100644 --- a/src/partitioned_contract/adjacency_tree.jl +++ b/src/partitioned_contract/adjacency_tree.jl @@ -0,0 +1 @@ + diff --git a/src/partitioned_contract/partitioned_contract.jl b/src/partitioned_contract/partitioned_contract.jl index e768feb0..97b2de5b 100644 --- a/src/partitioned_contract/partitioned_contract.jl +++ b/src/partitioned_contract/partitioned_contract.jl @@ -62,7 +62,7 @@ function partitioned_contract( inds_orderings = [p_edge_to_ordered_inds[e] for e in p_edges] v_to_tn[v], log_root_norm = approx_itensornetwork( tn, - _ansatz_tree(inds_orderings, ansatz); + _ansatz_tree(inds_orderings, _ortho_center(v_to_ordered_p_edges[v]), ansatz); alg=approx_itensornetwork_alg, cutoff=cutoff, maxdim=maxdim, @@ -74,12 +74,43 @@ function partitioned_contract( end end -# TODO: replace the subarray of `v1` with `v2` +function _is_neighbored_subset(v::Vector, set::Set) + if set == Set() + return true + end + @assert issubset(set, Set(v)) + i_begin = 1 + while !(i_begin in set) + i_begin += 1 + end + i_end = length(v) + while !(i_end in set) + i_end -= 1 + end + return Set(v[i_begin:i_end]) == set +end + +# replace the subarray of `v1` with `v2` function _replace_subarray(v1::Vector, v2::Vector) + if v2 == [] + return v1 + end + v1 = copy(v1) + num = 0 + for i in 1:length(v1) + if v1[i] in v2 + num += 1 + v1[i] = v2[num] + end + end + @assert num == length(v2) + return v1 end function _neighbor_edges(graph, vs) - return filter(e -> (e.src in vs and !(e.dst in vs)) || (e.dst in vs and !(e.src in vs)), edges(graph)) + return filter( + e -> (e.src in vs && !(e.dst in vs)) || (e.dst in vs && !(e.src in vs)), edges(graph) + ) end function _ind_orderings(partition::DataGraph) @@ -90,6 +121,87 @@ function _constrained_mincost_inds_ordering(inds_set::Set, tn::ITensorNetwork, p # TODO: edge set ordering of tn end -function _ansatz_tree(inds_orderings::Vector, ansatz::String) +function _ansatz_tree(inds_orderings::Vector, ortho_center::Integer, ansatz::String) # TODO end + +function _ortho_center(edges_to_contract_edges) + # TODO +end + +# function ordered_igs_to_binary_tree(ordered_igs, contract_igs, ig_to_linear_order; ansatz) +# @assert ansatz in ["comb", "mps"] +# @timeit_debug ITensors.timer "ordered_igs_to_binary_tree" begin +# if contract_igs == [] +# @info "contract_igs is empty vector" +# end +# # @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(igs, ig_to_linear_order; ansatz, direction) +# @assert ansatz in ["comb", "mps"] +# @assert direction in ["left", "right"] +# if ansatz == "comb" +# return line_to_tree([line_to_tree(ig_to_linear_order[ig]) for ig in igs]) +# end +# if direction == "left" +# order = vcat([ig_to_linear_order[ig] for ig in igs]...) +# return line_to_tree(order) +# else +# # First reverse get the order from middle to boundary, +# # and second reverse get the overall inds order from boundary to middle. +# order = vcat([ig_to_linear_order[ig] for ig in reverse(igs)]...) +# return line_to_tree(reverse(order)) +# 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 = ordered_igs_to_binary_tree( +# left_igs, ig_to_linear_order; ansatz="comb", direction="left" +# ) +# tree_contract = ordered_igs_to_binary_tree( +# contract_igs, ig_to_linear_order; ansatz="comb", direction="left" +# ) +# tree_2 = ordered_igs_to_binary_tree( +# reverse(right_igs), ig_to_linear_order; ansatz="comb", direction="left" +# ) +# # 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 From 81dc16ad31410d7bdb39e8439055319328ae4b95 Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Sat, 1 Jul 2023 11:54:59 -0500 Subject: [PATCH 07/35] api change --- src/mincut.jl | 21 ++++++++++++--------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/src/mincut.jl b/src/mincut.jl index 067dc045..9e6a0f14 100644 --- a/src/mincut.jl +++ b/src/mincut.jl @@ -4,31 +4,31 @@ MAX_WEIGHT = 1e32 """ Outputs a maximimally unbalanced directed binary tree DataGraph defining the desired graph structure """ -function path_graph_structure(tn::ITensorNetwork) - return path_graph_structure(tn, noncommoninds(Vector{ITensor}(tn)...)) +function path_graph_structure(tn::ITensorNetwork; alg="top_down") + return path_graph_structure(tn, noncommoninds(Vector{ITensor}(tn)...); alg=alg) end """ Given a `tn` and `outinds` (a subset of noncommoninds of `tn`), outputs a maximimally unbalanced directed binary tree DataGraph of `outinds` defining the desired graph structure """ -function path_graph_structure(tn::ITensorNetwork, outinds::Vector{<:Index}) - return _binary_tree_structure(tn, outinds; maximally_unbalanced=true) +function path_graph_structure(tn::ITensorNetwork, outinds::Vector{<:Index}; alg="top_down") + return _binary_tree_structure(Algorithm(alg), tn, outinds; maximally_unbalanced=true) end """ Outputs a directed binary tree DataGraph defining the desired graph structure """ -function binary_tree_structure(tn::ITensorNetwork) - return binary_tree_structure(tn, noncommoninds(Vector{ITensor}(tn)...)) +function binary_tree_structure(tn::ITensorNetwork; alg="top_down") + return binary_tree_structure(tn, noncommoninds(Vector{ITensor}(tn)...); alg=alg) end """ Given a `tn` and `outinds` (a subset of noncommoninds of `tn`), outputs a directed binary tree DataGraph of `outinds` defining the desired graph structure """ -function binary_tree_structure(tn::ITensorNetwork, outinds::Vector{<:Index}) - return _binary_tree_structure(tn, outinds; maximally_unbalanced=false) +function binary_tree_structure(tn::ITensorNetwork, outinds::Vector{<:Index}; alg="top_down") + return _binary_tree_structure(Algorithm(alg), tn, outinds; maximally_unbalanced=false) end """ @@ -126,7 +126,10 @@ Example: # TODO """ function _binary_tree_structure( - tn::ITensorNetwork, outinds::Vector{<:Index}; maximally_unbalanced::Bool=false + ::Algorithm"bottom_up", + tn::ITensorNetwork, + outinds::Vector{<:Index}; + maximally_unbalanced::Bool=false, ) inds_tree_vector = _binary_tree_partition_inds( tn, outinds; maximally_unbalanced=maximally_unbalanced From e5d73ed536ec55e79d13b7890e23c6f0bc166be9 Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Sun, 2 Jul 2023 10:26:32 -0500 Subject: [PATCH 08/35] Add binary tree structure with top down algorithm --- src/mincut.jl | 83 ++++++++++++++++++++++++++++-- test/test_binary_tree_partition.jl | 39 ++++++++------ 2 files changed, 102 insertions(+), 20 deletions(-) diff --git a/src/mincut.jl b/src/mincut.jl index 9e6a0f14..721f8781 100644 --- a/src/mincut.jl +++ b/src/mincut.jl @@ -116,6 +116,66 @@ function _maxweightoutinds_tn(tn::ITensorNetwork, outinds::Union{Nothing,Vector{ return maxweight_tn, out_to_maxweight_ind end +function _binary_tree_structure( + alg::Algorithm"top_down", + tn::ITensorNetwork, + outinds::Vector{<:Index}; + maximally_unbalanced::Bool=false, + backend="Metis", +) + nested_vector = _binary_tree_partition_inds_recursive_bisection( + tn, outinds; backend=backend + ) + if maximally_unbalanced + ordering = collect(Leaves(nested_vector)) + nested_vector = line_to_tree(ordering) + end + return _nested_vector_to_directed_tree(nested_vector) +end + +function _binary_tree_partition_inds_recursive_bisection( + tn::ITensorNetwork, + outinds::Vector{<:Index}; + backend="Metis", + left_inds=Set(), + right_inds=Set(), +) + inds = intersect(outinds, noncommoninds(Vector{ITensor}(tn)...)) + if length(inds) <= 1 + return length(inds) == 1 ? inds[1] : nothing + end + g_parts = partition(tn; npartitions=2, backend=backend) + # order g_parts + outinds_1 = noncommoninds(Vector{ITensor}(g_parts[1])...) + outinds_2 = noncommoninds(Vector{ITensor}(g_parts[2])...) + left_inds_1 = intersect(left_inds, outinds_1) + left_inds_2 = intersect(left_inds, outinds_2) + right_inds_1 = intersect(right_inds, outinds_1) + right_inds_2 = intersect(right_inds, outinds_2) + if length(left_inds_2) + length(right_inds_1) > length(left_inds_1) + length(right_inds_2) + g_parts[1], g_parts[2] = g_parts[2], g_parts[1] + end + intersect_inds = intersect(outinds_1, outinds_2) + inds1 = _binary_tree_partition_inds_recursive_bisection( + g_parts[1], + outinds; + backend=backend, + left_inds=left_inds, + right_inds=union(right_inds, intersect_inds), + ) + inds2 = _binary_tree_partition_inds_recursive_bisection( + g_parts[2], + outinds; + backend=backend, + left_inds=union(left_inds, intersect_inds), + right_inds=right_inds, + ) + if inds1 == nothing || inds2 == nothing + return inds1 == nothing ? inds2 : inds1 + end + return [inds1, inds2] +end + """ Given a tn and outinds (a subset of noncommoninds of tn), get a `DataGraph` with binary tree structure of outinds that will be used in the binary tree partition. @@ -149,7 +209,7 @@ function _binary_tree_partition_inds( return _binary_tree_partition_inds_mincut(tn_pair, out_to_maxweight_ind) else return line_to_tree( - _binary_tree_partition_inds_maximally_unbalanced(tn_pair, out_to_maxweight_ind) + _binary_tree_partition_inds_order_maximally_unbalanced(tn_pair, out_to_maxweight_ind) ) end end @@ -182,19 +242,32 @@ end Given a tn and outinds, returns a vector of indices representing MPS inds ordering. """ function _mps_partition_inds_order( - tn::ITensorNetwork, outinds::Union{Nothing,Vector{<:Index}} + tn::ITensorNetwork, + outinds::Union{Nothing,Vector{<:Index}}; + alg="top_down", + backend="Metis", ) + @assert alg in ["top_down", "bottom_up"] if outinds == nothing outinds = noncommoninds(Vector{ITensor}(tn)...) end if length(outinds) == 1 return outinds end - tn2, out_to_maxweight_ind = _maxweightoutinds_tn(tn, outinds) - return _binary_tree_partition_inds_maximally_unbalanced(tn => tn2, out_to_maxweight_ind) + if alg == "bottom_up" + tn2, out_to_maxweight_ind = _maxweightoutinds_tn(tn, outinds) + return _binary_tree_partition_inds_order_maximally_unbalanced( + tn => tn2, out_to_maxweight_ind + ) + else + nested_vector = _binary_tree_partition_inds_recursive_bisection( + tn, outinds; backend=backend + ) + return collect(Leaves(nested_vector)) + end end -function _binary_tree_partition_inds_maximally_unbalanced( +function _binary_tree_partition_inds_order_maximally_unbalanced( tn_pair::Pair{<:ITensorNetwork,<:ITensorNetwork}, out_to_maxweight_ind::Dict{Index,Index} ) outinds = collect(keys(out_to_maxweight_ind)) diff --git a/test/test_binary_tree_partition.jl b/test/test_binary_tree_partition.jl index f00ef278..4e0d14a5 100644 --- a/test/test_binary_tree_partition.jl +++ b/test/test_binary_tree_partition.jl @@ -1,5 +1,6 @@ using ITensors, OMEinsumContractionOrders using Graphs, NamedGraphs +using KaHyPar, Metis using ITensors: contract using ITensorNetworks: _root, @@ -10,7 +11,7 @@ using ITensorNetworks: _rem_vertex!, _DensityMartrixAlgGraph -@testset "test mincut functions on top of MPS" begin +@testset "test mincut and partition functions on top of MPS" begin i = Index(2, "i") j = Index(2, "j") k = Index(2, "k") @@ -23,13 +24,15 @@ using ITensorNetworks: 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) tn = ITensorNetwork(M[:]) - for out in [binary_tree_structure(tn), path_graph_structure(tn)] - @test out isa DataGraph - @test _is_rooted_directed_binary_tree(out) - @test length(vertex_data(out).values) == 8 + for alg in ["top_down", "bottom_up"] + for out in [binary_tree_structure(tn; alg=alg), path_graph_structure(tn; alg=alg)] + @test out isa DataGraph + @test _is_rooted_directed_binary_tree(out) + @test length(vertex_data(out).values) == 8 + end + out = _mps_partition_inds_order(tn, [o, p, i, j, k, l, m, n]; alg=alg) + @test out in [[i, j, k, l, m, n, o, p], [p, o, n, m, l, k, j, i]] end - out = _mps_partition_inds_order(tn, [o, p, i, j, k, l, m, n]) - @test out in [[i, j, k, l, m, n, o, p], [p, o, n, m, l, k, j, i]] p1, p2 = _mincut_partitions(tn, [k, l], [m, n]) # When MPS bond dimensions are large, the partition will not across internal inds @test (length(p1) == 0) || (length(p2) == 0) @@ -51,20 +54,26 @@ end tn[v...] = network[v...] end tn = ITensorNetwork(vec(tn[:, :, 1])) - for out in [binary_tree_structure(tn), path_graph_structure(tn)] - @test out isa DataGraph - @test _is_rooted_directed_binary_tree(out) - @test length(vertex_data(out).values) == 9 + for alg in ["top_down", "bottom_up"] + for out in [binary_tree_structure(tn; alg=alg), path_graph_structure(tn; alg=alg)] + @test out isa DataGraph + @test _is_rooted_directed_binary_tree(out) + @test length(vertex_data(out).values) == 9 + end end end @testset "test partition with mincut_recursive_bisection alg of disconnected tn" begin inds = [Index(2, "$i") for i in 1:5] tn = ITensorNetwork([randomITensor(i) for i in inds]) - par = partition(tn, binary_tree_structure(tn); alg="mincut_recursive_bisection") - networks = [Vector{ITensor}(par[v]) for v in vertices(par)] - network = vcat(networks...) - @test isapprox(contract(Vector{ITensor}(tn)), contract(network...)) + for alg in ["top_down", "bottom_up"] + par = partition( + tn, binary_tree_structure(tn; alg=alg); alg="mincut_recursive_bisection" + ) + networks = [Vector{ITensor}(par[v]) for v in vertices(par)] + network = vcat(networks...) + @test isapprox(contract(Vector{ITensor}(tn)), contract(network...)) + end end @testset "test partition with mincut_recursive_bisection alg and approx_itensornetwork" begin From 2f3b2b445d4a36f967bc3df5d35567da08fd5947 Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Sun, 2 Jul 2023 17:46:49 -0500 Subject: [PATCH 09/35] Let binary_tree_structure support a vector of inds vectors --- src/mincut.jl | 67 ++++++++++++++++++++---------- test/test_binary_tree_partition.jl | 14 +++++++ 2 files changed, 59 insertions(+), 22 deletions(-) diff --git a/src/mincut.jl b/src/mincut.jl index 721f8781..532317dd 100644 --- a/src/mincut.jl +++ b/src/mincut.jl @@ -130,19 +130,42 @@ function _binary_tree_structure( ordering = collect(Leaves(nested_vector)) nested_vector = line_to_tree(ordering) end - return _nested_vector_to_directed_tree(nested_vector) + return _nested_vector_to_digraph(nested_vector) +end + +function _map_nested_vector(dict::Dict, nested_vector) + if haskey(dict, nested_vector) + return dict[nested_vector] + end + return map(v -> _map_nested_vector(dict, v), nested_vector) end function _binary_tree_partition_inds_recursive_bisection( + tn::ITensorNetwork, outinds::Union{Vector{<:Vector},Vector{<:Index}}; backend="Metis" +) + tn = copy(tn) + tensor_to_inds = Dict() + ts = Vector{ITensor}() + for is in outinds + new_t = ITensor(is) + push!(ts, new_t) + tensor_to_inds[new_t] = is + end + new_tn = disjoint_union(tn, ITensorNetwork(ts)) + ts_nested_vector = _tensor_order_recursive_bisection(new_tn, ts; backend=backend) + return _map_nested_vector(tensor_to_inds, ts_nested_vector) +end + +function _tensor_order_recursive_bisection( tn::ITensorNetwork, - outinds::Vector{<:Index}; + tensors::Vector{ITensor}; backend="Metis", left_inds=Set(), right_inds=Set(), ) - inds = intersect(outinds, noncommoninds(Vector{ITensor}(tn)...)) - if length(inds) <= 1 - return length(inds) == 1 ? inds[1] : nothing + ts = intersect(tensors, Vector{ITensor}(tn)) + if length(ts) <= 1 + return length(ts) == 1 ? ts[1] : nothing end g_parts = partition(tn; npartitions=2, backend=backend) # order g_parts @@ -156,24 +179,24 @@ function _binary_tree_partition_inds_recursive_bisection( g_parts[1], g_parts[2] = g_parts[2], g_parts[1] end intersect_inds = intersect(outinds_1, outinds_2) - inds1 = _binary_tree_partition_inds_recursive_bisection( + ts1 = _tensor_order_recursive_bisection( g_parts[1], - outinds; + tensors; backend=backend, left_inds=left_inds, right_inds=union(right_inds, intersect_inds), ) - inds2 = _binary_tree_partition_inds_recursive_bisection( + ts2 = _tensor_order_recursive_bisection( g_parts[2], - outinds; + tensors; backend=backend, left_inds=union(left_inds, intersect_inds), right_inds=right_inds, ) - if inds1 == nothing || inds2 == nothing - return inds1 == nothing ? inds2 : inds1 + if ts1 == nothing || ts2 == nothing + return ts1 == nothing ? ts2 : ts1 end - return [inds1, inds2] + return [ts1, ts2] end """ @@ -194,7 +217,7 @@ function _binary_tree_structure( inds_tree_vector = _binary_tree_partition_inds( tn, outinds; maximally_unbalanced=maximally_unbalanced ) - return _nested_vector_to_directed_tree(inds_tree_vector) + return _nested_vector_to_digraph(inds_tree_vector) end function _binary_tree_partition_inds( @@ -214,19 +237,19 @@ function _binary_tree_partition_inds( end end -function _nested_vector_to_directed_tree(inds_tree_vector::Vector) - if length(inds_tree_vector) == 1 && inds_tree_vector[1] isa Index - inds_btree = DataGraph(NamedDiGraph([1]), Index) - inds_btree[1] = inds_tree_vector[1] +function _nested_vector_to_digraph(nested_vector::Vector) + if length(nested_vector) == 1 && !(nested_vector[1] isa Vector) + inds_btree = DataGraph(NamedDiGraph([1]), Any) + inds_btree[1] = nested_vector[1] return inds_btree end treenode_to_v = Dict{Union{Vector,Index},Int}() - graph = DataGraph(NamedDiGraph(), Index) + graph = DataGraph(NamedDiGraph(), Any) v = 1 - for treenode in PostOrderDFS(inds_tree_vector) + for treenode in PostOrderDFS(nested_vector) add_vertex!(graph, v) treenode_to_v[treenode] = v - if treenode isa Index + if !(treenode isa Vector) graph[v] = treenode else @assert length(treenode) == 2 @@ -243,7 +266,7 @@ Given a tn and outinds, returns a vector of indices representing MPS inds orderi """ function _mps_partition_inds_order( tn::ITensorNetwork, - outinds::Union{Nothing,Vector{<:Index}}; + outinds::Union{Nothing,Vector{<:Index},Vector{<:Vector}}; alg="top_down", backend="Metis", ) @@ -263,7 +286,7 @@ function _mps_partition_inds_order( nested_vector = _binary_tree_partition_inds_recursive_bisection( tn, outinds; backend=backend ) - return collect(Leaves(nested_vector)) + return filter(v -> v in outinds, collect(PreOrderDFS(nested_vector))) end end diff --git a/test/test_binary_tree_partition.jl b/test/test_binary_tree_partition.jl index 4e0d14a5..e6ea82af 100644 --- a/test/test_binary_tree_partition.jl +++ b/test/test_binary_tree_partition.jl @@ -45,6 +45,20 @@ using ITensorNetworks: @test sort(p2) == [5, 6, 7, 8] end +@testset "test _mps_partition_inds_order of a sub 2D grid" begin + N = (8, 8) + 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 + tensors = vec(tn)[1:20] + tn = ITensorNetwork(tensors) + @info _mps_partition_inds_order(tn, noncommoninds(tensors...); alg="top_down") + # TODO: think about how to check this +end + @testset "test _binary_tree_partition_inds of a 2D network" begin N = (3, 3, 3) linkdim = 2 From 9beee04464637a77f6c25a0c5fe21c9158bc5b15 Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Sun, 2 Jul 2023 21:36:48 -0500 Subject: [PATCH 10/35] implement reference ordering --- src/contract.jl | 3 +- .../partitioned_contract.jl | 60 ++++++++++++------- 2 files changed, 41 insertions(+), 22 deletions(-) diff --git a/src/contract.jl b/src/contract.jl index 0b82b7fa..aff4a022 100644 --- a/src/contract.jl +++ b/src/contract.jl @@ -19,10 +19,11 @@ function contract( end function contract( - alg::Algorithm"partitioned_contraction", + alg::Algorithm"partitioned_contract", tn::AbstractITensorNetwork; output_structure::Function=path_graph_structure, kwargs..., ) + # TODO return approx_itensornetwork(alg, tn, output_structure; kwargs...) end diff --git a/src/partitioned_contract/partitioned_contract.jl b/src/partitioned_contract/partitioned_contract.jl index 97b2de5b..a0d01c25 100644 --- a/src/partitioned_contract/partitioned_contract.jl +++ b/src/partitioned_contract/partitioned_contract.jl @@ -13,35 +13,51 @@ function partitioned_contract( traversal = post_order_dfs_vertices(contraction_tree, _root(contraction_tree)) contractions = setdiff(traversal, leaves) p_edge_to_ordered_inds = _ind_orderings(partition) - # build the orderings used for the ansatz tree - # For each pair in `v_to_ordered_p_edges`, the first item - # is the ordering of uncontracted edges for the contraction `v`, - # and the second item is the ordering of part of the uncontracted edges - # that are to be contracted in the next contraction (first contraction - # in the path of `v`). - v_to_ordered_p_edges = Dict{Tuple,Pair}() - for (ii, v) in enumerate(contractions) - @info "$(ii)/$(length(contractions)) th ansatz tree construction" - p_leaves = [v[1]..., v[2]...] + # Build the orderings used for the ansatz tree. + # For each tuple in `v_to_ordered_p_edges`, the first item is the + # reference ordering of uncontracted edges for the contraction `v`, + # the second item is the target ordering, and the third item is the + # ordering of part of the uncontracted edges that are to be contracted + # in the next contraction (first contraction in the path of `v`). + v_to_ordered_p_edges = Dict{Tuple,Tuple}() + # TODO: children + for (ii, v) in enumerate(traversal) + @info "$(ii)/$(length(traversal)) th ansatz tree construction" + p_leaves = vcat(v[1:(end - 1)]...) tn = ITensorNetwork(mapreduce(l -> Vector{ITensor}(partition[l]), vcat, p_leaves)) path = filter(u -> issubset(p_leaves, u[1]) || issubset(p_leaves, u[2]), contractions) p_edges = _neighbor_edges(partition, p_leaves) - inds_set = [Set(p_edge_to_ordered_inds[e]) for e in p_edges] - ordering = _constrained_minswap_inds_ordering(inds_set, tn, path) - p_edges = p_edges[sortperm(ordering)] - # update the contracted ordering and `v_to_ordered_p_edges[v]`. - # path[1] is the vertex to be contracted with `v` next. - if haskey(v_to_ordered_p_edges, path[1]) - contract_edges = v_to_ordered_p_edges[path[1]].second + # Get the reference ordering and target ordering of `v` + v_inds = map(e -> Set(p_edge_to_ordered_inds[e]), p_edges) + if v in leaves + ref_ordering, ordering = _constrained_minswap_inds_ordering(v_inds, tn, path) + else + c1, c2 = child_vertices(contraction_tree, v) + c1_inds_ordering = map( + e -> Set(p_edge_to_ordered_inds[e]), v_to_ordered_p_edges[c1][2] + ) + c2_inds_ordering = map( + e -> Set(p_edge_to_ordered_inds[e]), v_to_ordered_p_edges[c2][2] + ) + ref_ordering, ordering = _constrained_minswap_inds_ordering( + v_inds, tn, path, c1_inds_ordering, c2_inds_ordering + ) + end + ref_p_edges = p_edges[ref_ordering] + p_edges = p_edges[ordering] + # Update the contracted ordering and `v_to_ordered_p_edges[v]`. + # `sibling` is the vertex to be contracted with `v` next. + sibling = setdiff(child_vertices(contractions, path[1]), [v])[1] + if haskey(v_to_ordered_p_edges, sibling) + contract_edges = v_to_ordered_p_edges[sibling][3] @assert(_is_neighbored_subset(p_edges, Set(contract_edges))) p_edges = _replace_subarray(p_edges, contract_edges) else - p_leaves_2 = [path[1][1]..., path[1][2]...] + p_leaves_2 = vcat(sibling[1:(end - 1)]...) p_edges_2 = _neighbor_edges(partition, p_leaves_2) contract_edges = intersect(p_edges, p_edges_2) - contract_edges = filter(e -> e in contract_edges, p_edges) end - v_to_ordered_p_edges[v] = Pair(p_edges, contract_edges) + v_to_ordered_p_edges[v] = (ref_p_edges, p_edges, contract_edges) end # start approx_itensornetwork v_to_tn = Dict{Tuple,ITensorNetwork}() @@ -55,7 +71,8 @@ function partitioned_contract( c1, c2 = child_vertices(contraction_tree, v) tn = disjoint_union(v_to_tn[c1], v_to_tn[c2]) # TODO: rename tn since the names will be too long. - p_edges = v_to_ordered_p_edges[v].first + # TODO: swap size + p_edges = v_to_ordered_p_edges[v][2] if p_edges == [] # TODO: edge case with output being a scalar end @@ -115,6 +132,7 @@ end function _ind_orderings(partition::DataGraph) input_tn = # TODO + # TODO: use `_mps_partition_inds_order` with default `alg` and `backend` end function _constrained_mincost_inds_ordering(inds_set::Set, tn::ITensorNetwork, path::Vector) From bfd6f686197d68f78f6a590cf1b06ac9a6ecc667 Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Sun, 2 Jul 2023 21:49:18 -0500 Subject: [PATCH 11/35] Start on adjacency tree --- src/partitioned_contract/adjacency_tree.jl | 239 ++++++++++++++++++ .../partitioned_contract.jl | 7 +- 2 files changed, 241 insertions(+), 5 deletions(-) diff --git a/src/partitioned_contract/adjacency_tree.jl b/src/partitioned_contract/adjacency_tree.jl index 8b137891..7ed34fb5 100644 --- a/src/partitioned_contract/adjacency_tree.jl +++ b/src/partitioned_contract/adjacency_tree.jl @@ -1 +1,240 @@ +function boundary_state(v::Tuple{Tuple,String}, adj_igs::Set) + if Set(Leaves(v[1])) == adj_igs + return "all" + end + if v[2] == "unordered" + filter_children = filter(c -> issubset(adj_igs, Set(Leaves(c))), v[1]) + # length(filter_children) < 1 means adj_igs is distributed in multiple children + @assert length(filter_children) <= 1 + if length(filter_children) == 1 + return "middle" + end + # TODO: if more than 1 children contain adj_igs, currently we don't reorder the + # leaves. This may need to be optimized later. + return "invalid" + end + @assert length(v[1]) >= 2 + for i in 1:(length(v[1]) - 1) + leaves = vcat([Set(Leaves(c)) for c in v[1][1:i]]...) + if Set(leaves) == adj_igs + return "left" + end + end + for i in 2:length(v[1]) + leaves = vcat([Set(Leaves(c)) for c in v[1][i:end]]...) + if Set(leaves) == adj_igs + return "right" + end + end + return "invalid" +end +function reorder_to_boundary!( + adj_tree::NamedDiGraph{Tuple{Tuple,String}}, + v::Tuple{Tuple,String}, + target_child::Tuple{Tuple,String}; + direction="right", +) + new_v = v + children = child_vertices(adj_tree, v) + remain_children = setdiff(children, [target_child]) + @assert length(remain_children) >= 1 + if length(remain_children) == 1 + remain_child = remain_children[1] + if direction == "right" + new_v = ((remain_child[1], target_child[1]), "ordered") + else + new_v = ((target_child[1], remain_child[1]), "ordered") + end + if new_v != v + _add_vertex_edges!( + adj_tree, new_v; children=children, parent=parent_vertex(adj_tree, v) + ) + rem_vertex!(adj_tree, v) + end + else + new_child = (Tuple([v[1] for v in remain_children]), "unordered") + _add_vertex_edges!(adj_tree, new_child; children=remain_children, parent=v) + if direction == "right" + new_v = ((new_child[1], target_child[1]), "ordered") + else + new_v = ((target_child[1], new_child[1]), "ordered") + end + _add_vertex_edges!( + adj_tree, new_v; children=[new_child, target_child], parent=parent_vertex(adj_tree, v) + ) + rem_vertex!(adj_tree, v) + end + return new_v +end + +function _add_vertex_edges!( + adj_tree::NamedDiGraph{Tuple{Tuple,String}}, v; children=[], parent=nothing +) + add_vertex!(adj_tree, v) + if parent != nothing + add_edge!(adj_tree, parent => v) + end + for c in children + add_edge!(adj_tree, v => c) + end +end + +""" +reorder adj_tree based on adj_igs +""" +function reorder!( + adj_tree::NamedDiGraph{Tuple{Tuple,String}}, + root::Tuple{Tuple,String}, + adj_igs::Set; + boundary="right", +) + @assert boundary in ["left", "right"] + if boundary_state(root, adj_igs) == "all" + return false, root + end + sub_tree = subgraph(v -> issubset(Set(Leaves(v[1])), Set(Leaves(root[1]))), adj_tree) + traversal = post_order_dfs_vertices(sub_tree, root) + path = [v for v in traversal if issubset(adj_igs, Set(Leaves(v[1])))] + new_root = root + # get the boundary state + v_to_state = Dict{Tuple{Tuple,String},String}() + for v in path + state = boundary_state(v, adj_igs) + if state == "invalid" + return false, root + end + v_to_state[v] = state + end + for v in path + children = child_vertices(adj_tree, v) + # reorder + if v_to_state[v] in ["left", "right"] && v_to_state[v] != boundary + @assert v[2] == "ordered" + new_v = (reverse(v[1]), v[2]) + new_root = (v == root) ? new_v : new_root + _add_vertex_edges!( + adj_tree, new_v; children=children, parent=parent_vertex(adj_tree, v) + ) + rem_vertex!(adj_tree, v) + elseif v_to_state[v] == "middle" + @assert v[2] == "unordered" + target_child = filter(c -> issubset(adj_igs, Set(Leaves(c[1]))), children) + @assert length(target_child) == 1 + new_v = reorder_to_boundary!(adj_tree, v, target_child[1]; direction=boundary) + new_root = (v == root) ? new_v : new_root + end + end + return true, new_root +end + +# Update both keys and values in igs_to_adjacency_tree based on adjacent_igs +# NOTE: keys of `igs_to_adjacency_tree` are target igs, not those adjacent to ancestors +function update_adjacency_tree!( + adjacency_tree::NamedDiGraph{Tuple{Tuple,String}}, adjacent_igs::Set +) + @timeit_debug ITensors.timer "update_adjacency_tree" begin + root_v_to_adjacent_igs = Dict{Tuple{Tuple,String},Set}() + for r in _roots(adjacency_tree) + root_igs = Set(Leaves(r[1])) + common_igs = intersect(adjacent_igs, root_igs) + if common_igs != Set() + root_v_to_adjacent_igs[r] = common_igs + end + end + if length(root_v_to_adjacent_igs) == 1 + return nothing + end + # if at least 3: for now just put everything together + if length(root_v_to_adjacent_igs) >= 3 + __roots = keys(root_v_to_adjacent_igs) + new_v = (Tuple([r[1] for r in __roots]), "unordered") + _add_vertex_edges!(adjacency_tree, new_v; children=__roots) + return nothing + end + # if 2: assign adjacent_igs to boundary of root_igs (if possible), then concatenate + v1, v2 = collect(keys(root_v_to_adjacent_igs)) + reordered_1, update_v1 = reorder!( + adjacency_tree, v1, root_v_to_adjacent_igs[v1]; boundary="right" + ) + reordered_2, update_v2 = reorder!( + adjacency_tree, v2, root_v_to_adjacent_igs[v2]; boundary="left" + ) + cs1 = child_vertices(adjacency_tree, update_v1) + cs2 = child_vertices(adjacency_tree, update_v2) + if (!reordered_1) && (!reordered_2) + new_v = ((update_v1[1], update_v2[1]), "unordered") + _add_vertex_edges!(adjacency_tree, new_v; children=[update_v1, update_v2]) + elseif (reordered_2) + new_v = ((update_v1[1], update_v2[1]...), "ordered") + _add_vertex_edges!(adjacency_tree, new_v; children=[update_v1, cs2...]) + rem_vertex!(adjacency_tree, update_v2) + elseif (reordered_1) + new_v = ((update_v1[1]..., update_v2[1]), "ordered") + _add_vertex_edges!(adjacency_tree, new_v; children=[update_v2, cs1...]) + rem_vertex!(adjacency_tree, update_v1) + else + new_v = ((update_v1[1]..., update_v2[1]...), "ordered") + _add_vertex_edges!(adjacency_tree, new_v; children=[cs1..., cs2...]) + rem_vertex!(adjacency_tree, update_v1) + rem_vertex!(adjacency_tree, update_v2) + end + end +end + +# Generate the adjacency tree of a contraction tree +# Args: +# ========== +# ctree: the input contraction tree +# path: the path containing ancestor ctrees of the input ctree +# ctree_to_igs: mapping each ctree to neighboring index groups +function _generate_adjacency_tree(ctree, path, ctree_to_open_edges) + @timeit_debug ITensors.timer "_generate_adjacency_tree" begin + # mapping each index group to adjacent input igs + ig_to_input_adj_igs = Dict{Any,Set}() + # mapping each igs to an adjacency tree + adjacency_tree = NamedDiGraph{Tuple{Tuple,String}}() + for ig in ctree_to_open_edges[ctree] + ig_to_input_adj_igs[ig] = Set([ig]) + v = ((ig,), "unordered") + add_vertex!(adjacency_tree, v) + end + for (i, a) in path + inter_igs = intersect(ctree_to_open_edges[a[1]], ctree_to_open_edges[a[2]]) + new_igs_index = (i == 1) ? 2 : 1 + new_igs = setdiff(ctree_to_open_edges[a[new_igs_index]], inter_igs) + adjacent_igs = union([ig_to_input_adj_igs[ig] for ig in inter_igs]...) + # `inter_igs != []` means it's a tensor product + if inter_igs != [] + update_adjacency_tree!(adjacency_tree, adjacent_igs) + end + for ig in new_igs + ig_to_input_adj_igs[ig] = adjacent_igs + end + # @info "adjacency_tree", adjacency_tree + if length(_roots(adjacency_tree)) == 1 + return adjacency_tree + end + end + __roots = _roots(adjacency_tree) + if length(__roots) > 1 + new_v = (Tuple([r[1] for r in __roots]), "unordered") + _add_vertex_edges!(adjacency_tree, new_v; children=__roots) + end + return adjacency_tree + end +end + +function _constrained_mincost_inds_ordering(inds_set::Vector{Set}, tn::ITensorNetwork, path::Vector) + # TODO: edge set ordering of tn +end + +function _constrained_mincost_inds_ordering( + inds_set::Vector{Set}, + tn::ITensorNetwork, + path::Vector, + input_order_1::Vector{Set}, + input_order_2::Vector{Set}, +) + # TODO: edge set ordering of tn +end diff --git a/src/partitioned_contract/partitioned_contract.jl b/src/partitioned_contract/partitioned_contract.jl index a0d01c25..aaafd352 100644 --- a/src/partitioned_contract/partitioned_contract.jl +++ b/src/partitioned_contract/partitioned_contract.jl @@ -20,7 +20,6 @@ function partitioned_contract( # ordering of part of the uncontracted edges that are to be contracted # in the next contraction (first contraction in the path of `v`). v_to_ordered_p_edges = Dict{Tuple,Tuple}() - # TODO: children for (ii, v) in enumerate(traversal) @info "$(ii)/$(length(traversal)) th ansatz tree construction" p_leaves = vcat(v[1:(end - 1)]...) @@ -30,6 +29,7 @@ function partitioned_contract( # Get the reference ordering and target ordering of `v` v_inds = map(e -> Set(p_edge_to_ordered_inds[e]), p_edges) if v in leaves + # TODO ref_ordering, ordering = _constrained_minswap_inds_ordering(v_inds, tn, path) else c1, c2 = child_vertices(contraction_tree, v) @@ -39,6 +39,7 @@ function partitioned_contract( c2_inds_ordering = map( e -> Set(p_edge_to_ordered_inds[e]), v_to_ordered_p_edges[c2][2] ) + # TODO ref_ordering, ordering = _constrained_minswap_inds_ordering( v_inds, tn, path, c1_inds_ordering, c2_inds_ordering ) @@ -135,10 +136,6 @@ function _ind_orderings(partition::DataGraph) # TODO: use `_mps_partition_inds_order` with default `alg` and `backend` end -function _constrained_mincost_inds_ordering(inds_set::Set, tn::ITensorNetwork, path::Vector) - # TODO: edge set ordering of tn -end - function _ansatz_tree(inds_orderings::Vector, ortho_center::Integer, ansatz::String) # TODO end From 4936d79f54d4f3dc07d0f36c7356bea9366081cc Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Sun, 2 Jul 2023 23:20:09 -0500 Subject: [PATCH 12/35] Finish _adjacency_tree, tests to be added --- src/mincut.jl | 6 +-- src/partitioned_contract/adjacency_tree.jl | 38 ++++++++++++------- .../partitioned_contract.jl | 9 +++-- 3 files changed, 34 insertions(+), 19 deletions(-) diff --git a/src/mincut.jl b/src/mincut.jl index 532317dd..9e223f3d 100644 --- a/src/mincut.jl +++ b/src/mincut.jl @@ -141,13 +141,13 @@ function _map_nested_vector(dict::Dict, nested_vector) end function _binary_tree_partition_inds_recursive_bisection( - tn::ITensorNetwork, outinds::Union{Vector{<:Vector},Vector{<:Index}}; backend="Metis" + tn::ITensorNetwork, outinds::Union{Vector{Set},Vector{<:Index}}; backend="Metis" ) tn = copy(tn) tensor_to_inds = Dict() ts = Vector{ITensor}() for is in outinds - new_t = ITensor(is) + new_t = ITensor(collect(is)...) push!(ts, new_t) tensor_to_inds[new_t] = is end @@ -266,7 +266,7 @@ Given a tn and outinds, returns a vector of indices representing MPS inds orderi """ function _mps_partition_inds_order( tn::ITensorNetwork, - outinds::Union{Nothing,Vector{<:Index},Vector{<:Vector}}; + outinds::Union{Nothing,Vector{<:Index},Vector{Set}}; alg="top_down", backend="Metis", ) diff --git a/src/partitioned_contract/adjacency_tree.jl b/src/partitioned_contract/adjacency_tree.jl index 7ed34fb5..5f9ffb28 100644 --- a/src/partitioned_contract/adjacency_tree.jl +++ b/src/partitioned_contract/adjacency_tree.jl @@ -183,26 +183,34 @@ function update_adjacency_tree!( end # Generate the adjacency tree of a contraction tree -# Args: -# ========== -# ctree: the input contraction tree -# path: the path containing ancestor ctrees of the input ctree -# ctree_to_igs: mapping each ctree to neighboring index groups -function _generate_adjacency_tree(ctree, path, ctree_to_open_edges) +# TODO: add test +function _adjacency_tree(v::Tuple, path::Vector, partition::DataGraph, p_edge_to_inds::Dict) @timeit_debug ITensors.timer "_generate_adjacency_tree" begin # mapping each index group to adjacent input igs ig_to_input_adj_igs = Dict{Any,Set}() # mapping each igs to an adjacency tree adjacency_tree = NamedDiGraph{Tuple{Tuple,String}}() - for ig in ctree_to_open_edges[ctree] + p_leaves = vcat(v[1:(end - 1)]...) + p_edges = _neighbor_edges(partition, p_leaves) + for ig in map(e -> Set(p_edge_to_inds[e]), p_edges) ig_to_input_adj_igs[ig] = Set([ig]) v = ((ig,), "unordered") add_vertex!(adjacency_tree, v) end - for (i, a) in path - inter_igs = intersect(ctree_to_open_edges[a[1]], ctree_to_open_edges[a[2]]) - new_igs_index = (i == 1) ? 2 : 1 - new_igs = setdiff(ctree_to_open_edges[a[new_igs_index]], inter_igs) + for contraction in path + children = child_vertices(contractions, path[1]) + ancester = filter(u -> p_leaves in vcat(u[1:(end - 1)]...), children)[1] + sibling = setdiff(children, [ancester])[1] + ancester_igs = map( + e -> Set(p_edge_to_inds[e]), + _neighbor_edges(partition, vcat(ancester[1:(end - 1)]...)), + ) + sibling_igs = map( + e -> Set(p_edge_to_inds[e]), + _neighbor_edges(partition, vcat(sibling[1:(end - 1)]...)), + ) + inter_igs = intersect(ancester_igs, sibling_igs) + new_igs = setdiff(sibling_igs, inter_igs) adjacent_igs = union([ig_to_input_adj_igs[ig] for ig in inter_igs]...) # `inter_igs != []` means it's a tensor product if inter_igs != [] @@ -225,14 +233,18 @@ function _generate_adjacency_tree(ctree, path, ctree_to_open_edges) end end -function _constrained_mincost_inds_ordering(inds_set::Vector{Set}, tn::ITensorNetwork, path::Vector) +function _constrained_mincost_inds_ordering( + inds_set::Vector{Set}, + constraint_tree::NamedDiGraph{Tuple{Tuple,String}}, + tn::ITensorNetwork, +) # TODO: edge set ordering of tn end function _constrained_mincost_inds_ordering( inds_set::Vector{Set}, + constraint_tree::NamedDiGraph{Tuple{Tuple,String}}, tn::ITensorNetwork, - path::Vector, input_order_1::Vector{Set}, input_order_2::Vector{Set}, ) diff --git a/src/partitioned_contract/partitioned_contract.jl b/src/partitioned_contract/partitioned_contract.jl index aaafd352..2060b071 100644 --- a/src/partitioned_contract/partitioned_contract.jl +++ b/src/partitioned_contract/partitioned_contract.jl @@ -28,9 +28,12 @@ function partitioned_contract( p_edges = _neighbor_edges(partition, p_leaves) # Get the reference ordering and target ordering of `v` v_inds = map(e -> Set(p_edge_to_ordered_inds[e]), p_edges) + constraint_tree = _adjacency_tree(v, path, partition, p_edge_to_ordered_inds) if v in leaves # TODO - ref_ordering, ordering = _constrained_minswap_inds_ordering(v_inds, tn, path) + ref_ordering, ordering = _constrained_minswap_inds_ordering( + v_inds, constraint_tree, tn + ) else c1, c2 = child_vertices(contraction_tree, v) c1_inds_ordering = map( @@ -41,7 +44,7 @@ function partitioned_contract( ) # TODO ref_ordering, ordering = _constrained_minswap_inds_ordering( - v_inds, tn, path, c1_inds_ordering, c2_inds_ordering + v_inds, constraint_tree, tn, c1_inds_ordering, c2_inds_ordering ) end ref_p_edges = p_edges[ref_ordering] @@ -132,7 +135,7 @@ function _neighbor_edges(graph, vs) end function _ind_orderings(partition::DataGraph) - input_tn = # TODO + # input_tn = # TODO # TODO: use `_mps_partition_inds_order` with default `alg` and `backend` end From 7c12d4425c1e410390d020566e7f501a5c74baf2 Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Mon, 3 Jul 2023 11:46:46 -0500 Subject: [PATCH 13/35] Implement the first version of _constrained_minswap_inds_ordering --- src/mincut.jl | 57 +++++++++++---- src/partitioned_contract/adjacency_tree.jl | 18 ----- .../constrained_ordering.jl | 64 ++++++++++++++++ .../partitioned_contract.jl | 54 ++------------ src/partitioned_contract/utils.jl | 73 +++++++++++++++++++ 5 files changed, 185 insertions(+), 81 deletions(-) create mode 100644 src/partitioned_contract/constrained_ordering.jl create mode 100644 src/partitioned_contract/utils.jl diff --git a/src/mincut.jl b/src/mincut.jl index 9e223f3d..227956cb 100644 --- a/src/mincut.jl +++ b/src/mincut.jl @@ -334,22 +334,7 @@ function _binary_tree_partition_inds_mincut( return outinds end -""" -Find a vector of indices within sourceinds_list yielding the mincut of given tn_pair. -Args: - tn_pair: a pair of tns (tn1 => tn2), where tn2 is generated via _maxweightoutinds_tn(tn1) - out_to_maxweight_ind: a dict mapping each out ind in tn1 to out ind in tn2 - sourceinds_list: a list of vector of indices to be considered -Note: - For each sourceinds in sourceinds_list, we consider its mincut within both tns (tn1, tn2) given in tn_pair. - The mincut in tn1 represents the rank upper bound when splitting sourceinds with other inds in outinds. - The mincut in tn2 represents the rank upper bound when the weights of outinds are very large. - The first mincut upper_bounds the number of non-zero singular values, while the second empirically reveals the - singular value decay. - We output the sourceinds where the first mincut value is the minimum, the secound mincut value is also - the minimum under the condition that the first mincut is optimal, and the sourceinds have the lowest all-pair shortest path. -""" -function _mincut_inds( +function _mincut_partitions_costs( tn_pair::Pair{<:ITensorNetwork,<:ITensorNetwork}, out_to_maxweight_ind::Dict{Index,Index}, sourceinds_list::Vector{<:Vector{<:Index}}, @@ -377,6 +362,46 @@ function _mincut_inds( weights, _get_weights(source_inds, outinds, maxweight_source_inds, maxweight_outinds) ) end + return weights +end + +""" +Find a vector of indices within sourceinds_list yielding the mincut of given tn_pair. +Args: + tn_pair: a pair of tns (tn1 => tn2), where tn2 is generated via _maxweightoutinds_tn(tn1) + out_to_maxweight_ind: a dict mapping each out ind in tn1 to out ind in tn2 + sourceinds_list: a list of vector of indices to be considered +Note: + For each sourceinds in sourceinds_list, we consider its mincut within both tns (tn1, tn2) given in tn_pair. + The mincut in tn1 represents the rank upper bound when splitting sourceinds with other inds in outinds. + The mincut in tn2 represents the rank upper bound when the weights of outinds are very large. + The first mincut upper_bounds the number of non-zero singular values, while the second empirically reveals the + singular value decay. + We output the sourceinds where the first mincut value is the minimum, the secound mincut value is also + the minimum under the condition that the first mincut is optimal, and the sourceinds have the lowest all-pair shortest path. +""" +function _mincut_inds( + tn_pair::Pair{<:ITensorNetwork,<:ITensorNetwork}, + out_to_maxweight_ind::Dict{Index,Index}, + sourceinds_list::Vector{<:Vector{<:Index}}, +) + weights = _partition_mincuts_dist(tn_pair, out_to_maxweight_ind, sourceinds_list) i = argmin(weights) return sourceinds_list[i], i end + +function _mps_mincut_partition_cost(tn::ITensorNetwork, inds_vector::Vector{Set}) + @timeit_debug ITensors.timer "_comb_mincuts_dist" begin + inds_vector = map(inds -> collect(inds), inds_vector) + outinds = vcat(inds_vector...) + maxweight_tn, out_to_maxweight_ind = _maxweightoutinds_tn(tn, outinds) + sourceinds_list = [vcat(inds_vector[1:i]...) for i in 1:(length(inds_vector) - 1)] + weights = _partition_mincuts_dist( + tn => maxweight_tn, out_to_maxweight_ind, sourceinds_list + ) + mincut_val = sum([w[1] for w in weights]) + maxweight_mincut_val = sum([w[2] for w in weights]) + dist = sum([w[3] for w in weights]) + return (mincut_val, maxweight_mincut_val, dist) + end +end diff --git a/src/partitioned_contract/adjacency_tree.jl b/src/partitioned_contract/adjacency_tree.jl index 5f9ffb28..1c1a46d0 100644 --- a/src/partitioned_contract/adjacency_tree.jl +++ b/src/partitioned_contract/adjacency_tree.jl @@ -232,21 +232,3 @@ function _adjacency_tree(v::Tuple, path::Vector, partition::DataGraph, p_edge_to return adjacency_tree end end - -function _constrained_mincost_inds_ordering( - inds_set::Vector{Set}, - constraint_tree::NamedDiGraph{Tuple{Tuple,String}}, - tn::ITensorNetwork, -) - # TODO: edge set ordering of tn -end - -function _constrained_mincost_inds_ordering( - inds_set::Vector{Set}, - constraint_tree::NamedDiGraph{Tuple{Tuple,String}}, - tn::ITensorNetwork, - input_order_1::Vector{Set}, - input_order_2::Vector{Set}, -) - # TODO: edge set ordering of tn -end diff --git a/src/partitioned_contract/constrained_ordering.jl b/src/partitioned_contract/constrained_ordering.jl new file mode 100644 index 00000000..97c42c37 --- /dev/null +++ b/src/partitioned_contract/constrained_ordering.jl @@ -0,0 +1,64 @@ +function _constrained_minswap_inds_ordering( + constraint_tree::NamedDiGraph{Tuple{Tuple,String}}, + ref_ordering::Vector{Set}, + tn::ITensorNetwork, +) + leaves = leaf_vertices(constraint_tree) + root = _root(constraint_tree) + v_to_order = Dict{Tuple{Tuple,String},Vector{IndexGroup}}() + for v in post_order_dfs_vertices(constraint_tree, root) + if v in leaves + v_to_order[v] = [v[1]...] + continue + end + child_orders = Vector{Vector{IndexGroup}}() + children = child_vertices(constraint_tree, v) + for inds_tuple in v[1] + cs = filter(c -> c[1] == inds_tuple, children) + @assert length(cs) == 1 + push!(child_orders, v_to_order[cs[1]]) + end + input_order = [n for n in ref_ordering if n in vcat(child_orders...)] + # Optimize the ordering in child_orders + if v[2] == "ordered" + perms = [child_orders, reverse(child_orders)] + nswaps = [length(_bubble_sort(vcat(p...), input_order)) for p in perms] + perms = [perms[i] for i in 1:length(perms) if nswaps[i] == min(nswaps...)] + output_order = _mincut_permutation(perms, tn) + else + output_order = _best_perm_greedy(child_orders, input_order, tn) + end + v_to_order[v] = vcat(output_order...) + end + nswap = length(_bubble_sort(v_to_order[root], ref_ordering)) + return v_to_order[root], nswap +end + +function _constrained_minswap_inds_ordering( + constraint_tree::NamedDiGraph{Tuple{Tuple,String}}, + input_order_1::Vector{Set}, + input_order_2::Vector{Set}, + tn::ITensorNetwork, +) + # TODO: edge set ordering of tn +end + +function _mincut_permutation(perms::Vector{<:Vector}, tn::ITensorNetwork) + if length(perms) == 1 + return perms[1] + end + mincuts_dist = map(p -> _mps_mincut_partition_cost(tn, p), perms) + return perms[argmin(mincuts_dist)] +end + +function _best_perm_greedy(vs::Vector{<:Vector}, order::Vector, tn::ITensorNetwork) + ordered_vs = [vs[1]] + for v in vs[2:end] + perms = [insert!(copy(ordered_vs), i, v) for i in 1:(length(ordered_vs) + 1)] + suborder = filter(n -> n in vcat(perms[1]...), order) + nswaps = map(p -> length(_bubble_sort(vcat(p...), suborder)), perms) + perms = [perms[i] for i in 1:length(perms) if nswaps[i] == min(nswaps...)] + ordered_vs = _mincut_permutation(perms, tn) + end + return ordered_vs +end diff --git a/src/partitioned_contract/partitioned_contract.jl b/src/partitioned_contract/partitioned_contract.jl index 2060b071..4f5a4617 100644 --- a/src/partitioned_contract/partitioned_contract.jl +++ b/src/partitioned_contract/partitioned_contract.jl @@ -30,9 +30,9 @@ function partitioned_contract( v_inds = map(e -> Set(p_edge_to_ordered_inds[e]), p_edges) constraint_tree = _adjacency_tree(v, path, partition, p_edge_to_ordered_inds) if v in leaves - # TODO - ref_ordering, ordering = _constrained_minswap_inds_ordering( - v_inds, constraint_tree, tn + ref_inds_ordering = _mps_partition_inds_order(tn, v_inds; alg="top_down") + inds_ordering = _constrained_minswap_inds_ordering( + constraint_tree, ref_inds_ordering, tn ) else c1, c2 = child_vertices(contraction_tree, v) @@ -42,13 +42,12 @@ function partitioned_contract( c2_inds_ordering = map( e -> Set(p_edge_to_ordered_inds[e]), v_to_ordered_p_edges[c2][2] ) - # TODO - ref_ordering, ordering = _constrained_minswap_inds_ordering( - v_inds, constraint_tree, tn, c1_inds_ordering, c2_inds_ordering + ref_inds_ordering, inds_ordering = _constrained_minswap_inds_ordering( + constraint_tree, c1_inds_ordering, c2_inds_ordering, tn ) end - ref_p_edges = p_edges[ref_ordering] - p_edges = p_edges[ordering] + ref_p_edges = p_edges[_findperm(v_inds, ref_inds_ordering)] + p_edges = p_edges[_findperm(v_inds, inds_ordering)] # Update the contracted ordering and `v_to_ordered_p_edges[v]`. # `sibling` is the vertex to be contracted with `v` next. sibling = setdiff(child_vertices(contractions, path[1]), [v])[1] @@ -95,45 +94,6 @@ function partitioned_contract( end end -function _is_neighbored_subset(v::Vector, set::Set) - if set == Set() - return true - end - @assert issubset(set, Set(v)) - i_begin = 1 - while !(i_begin in set) - i_begin += 1 - end - i_end = length(v) - while !(i_end in set) - i_end -= 1 - end - return Set(v[i_begin:i_end]) == set -end - -# replace the subarray of `v1` with `v2` -function _replace_subarray(v1::Vector, v2::Vector) - if v2 == [] - return v1 - end - v1 = copy(v1) - num = 0 - for i in 1:length(v1) - if v1[i] in v2 - num += 1 - v1[i] = v2[num] - end - end - @assert num == length(v2) - return v1 -end - -function _neighbor_edges(graph, vs) - return filter( - e -> (e.src in vs && !(e.dst in vs)) || (e.dst in vs && !(e.src in vs)), edges(graph) - ) -end - function _ind_orderings(partition::DataGraph) # input_tn = # TODO # TODO: use `_mps_partition_inds_order` with default `alg` and `backend` diff --git a/src/partitioned_contract/utils.jl b/src/partitioned_contract/utils.jl new file mode 100644 index 00000000..ebbccef9 --- /dev/null +++ b/src/partitioned_contract/utils.jl @@ -0,0 +1,73 @@ +function _is_neighbored_subset(v::Vector, set::Set) + if set == Set() + return true + end + @assert issubset(set, Set(v)) + i_begin = 1 + while !(i_begin in set) + i_begin += 1 + end + i_end = length(v) + while !(i_end in set) + i_end -= 1 + end + return Set(v[i_begin:i_end]) == set +end + +# replace the subarray of `v1` with `v2` +function _replace_subarray(v1::Vector, v2::Vector) + if v2 == [] + return v1 + end + v1 = copy(v1) + num = 0 + for i in 1:length(v1) + if v1[i] in v2 + num += 1 + v1[i] = v2[num] + end + end + @assert num == length(v2) + return v1 +end + +function _neighbor_edges(graph, vs) + return filter( + e -> (e.src in vs && !(e.dst in vs)) || (e.dst in vs && !(e.src in vs)), edges(graph) + ) +end + +# Find the permutation to change `v1` into `v2` +function _findperm(v1, v2) + index_to_number = Dict() + for (i, v) in enumerate(v2) + index_to_number[v] = i + end + v1_num = [index_to_number[v] for v in v1] + return sortperm(v1_num) +end + +function _bubble_sort(v::Vector) + @timeit_debug ITensors.timer "bubble_sort" begin + permutations = [] + n = length(v) + for i in 1:n + for j in 1:(n - i) + if v[j] > v[j + 1] + v[j], v[j + 1] = v[j + 1], v[j] + push!(permutations, j) + end + end + end + return permutations + end +end + +function _bubble_sort(v1::Vector, v2::Vector) + index_to_number = Dict() + for (i, v) in enumerate(v2) + index_to_number[v] = i + end + v1_num = [index_to_number[v] for v in v1] + return _bubble_sort(v1_num) +end From 6c2998f14c699e9efb1974a44330219f9110ff32 Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Mon, 3 Jul 2023 12:14:03 -0500 Subject: [PATCH 14/35] Finish _constrained_minswap_inds_ordering --- .../constrained_ordering.jl | 74 ++++++++++++++++++- 1 file changed, 73 insertions(+), 1 deletion(-) diff --git a/src/partitioned_contract/constrained_ordering.jl b/src/partitioned_contract/constrained_ordering.jl index 97c42c37..dc5371f0 100644 --- a/src/partitioned_contract/constrained_ordering.jl +++ b/src/partitioned_contract/constrained_ordering.jl @@ -40,7 +40,37 @@ function _constrained_minswap_inds_ordering( input_order_2::Vector{Set}, tn::ITensorNetwork, ) - # TODO: edge set ordering of tn + inter_igs = intersect(input_order_1, input_order_2) + left_1, right_1 = _split_vector(input_order_1, inter_igs) + left_2, right_2 = _split_vector(input_order_2, inter_igs) + # @info "lengths of the input partitions", + # sort([length(left_1), length(right_2), length(left_2), length(right_2)]) + num_swaps_1 = min(length(left_1), length(left_2)) + min(length(right_1), length(right_2)) + num_swaps_2 = min(length(left_1), length(right_2)) + min(length(right_1), length(left_2)) + if num_swaps_1 == num_swaps_2 + inputs_1 = _low_swap_merge(left_1, right_1, left_2, right_2) + inputs_2 = _low_swap_merge(left_1, right_1, reverse(right_2), reverse(left_2)) + inputs = [inputs_1..., inputs_2...] + elseif num_swaps_1 > num_swaps_2 + inputs = _low_swap_merge(left_1, right_1, reverse(right_2), reverse(left_2)) + else + inputs = _low_swap_merge(left_1, right_1, left_2, right_2) + end + adj_tree_copies = [copy(adj_tree) for _ in 1:length(inputs)] + outputs = [] + nswaps_list = [] + for (t, i) in zip(adj_tree_copies, inputs) + output, nswaps = _constrained_minswap_inds_ordering(t, i, tn) + push!(outputs, output) + push!(nswaps_list, nswaps) + end + inputs = [inputs[i] for i in 1:length(inputs) if nswaps_list[i] == min(nswaps_list...)] + outputs = [outputs[i] for i in 1:length(outputs) if nswaps_list[i] == min(nswaps_list...)] + if length(inputs) == 1 + return inputs[1], outputs[1] + end + mincuts = map(o -> _mps_mincut_partition_cost(tn, o), outputs) + return inputs[argmin(mincuts)], outputs[argmin(mincuts)] end function _mincut_permutation(perms::Vector{<:Vector}, tn::ITensorNetwork) @@ -62,3 +92,45 @@ function _best_perm_greedy(vs::Vector{<:Vector}, order::Vector, tn::ITensorNetwo end return ordered_vs end + +function _make_list(left_lists, right_lists) + out_lists = [] + for l in left_lists + for r in right_lists + push!(out_lists, [l..., r...]) + end + end + return out_lists +end + +function _low_swap_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 + return _make_list(left_lists, right_lists) +end + +function _split_vector(v::Vector, sub_v::Vector) + left_v = [] + right_v = [] + target_array = left_v + for i in v + if i in sub_v + target_array = right_v + continue + end + push!(target_array, i) + end + return left_v, right_v +end From 29db3b3a90ee6623de1c7365ff623cadd0439b43 Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Mon, 3 Jul 2023 15:40:53 -0500 Subject: [PATCH 15/35] Finish a pass of partitioned_contract --- .../partitioned_contract.jl | 89 ++++++++++++++----- 1 file changed, 69 insertions(+), 20 deletions(-) diff --git a/src/partitioned_contract/partitioned_contract.jl b/src/partitioned_contract/partitioned_contract.jl index 4f5a4617..6146a6de 100644 --- a/src/partitioned_contract/partitioned_contract.jl +++ b/src/partitioned_contract/partitioned_contract.jl @@ -5,6 +5,7 @@ function partitioned_contract( approx_itensornetwork_alg="density_matrix", cutoff=1e-15, maxdim=10000, + swap_size=1, contraction_sequence_alg, contraction_sequence_kwargs, ) @@ -50,6 +51,9 @@ function partitioned_contract( p_edges = p_edges[_findperm(v_inds, inds_ordering)] # Update the contracted ordering and `v_to_ordered_p_edges[v]`. # `sibling` is the vertex to be contracted with `v` next. + # Note: the contracted ordering in `ref_p_edges` is not changed, + # since `ref_p_edges` focuses on minimizing the number of swaps + # rather than making later contractions efficient. sibling = setdiff(child_vertices(contractions, path[1]), [v])[1] if haskey(v_to_ordered_p_edges, sibling) contract_edges = v_to_ordered_p_edges[sibling][3] @@ -62,35 +66,55 @@ function partitioned_contract( end v_to_ordered_p_edges[v] = (ref_p_edges, p_edges, contract_edges) end - # start approx_itensornetwork + # start approximate contraction v_to_tn = Dict{Tuple,ITensorNetwork}() for v in leaves @assert length(v[1]) == 1 v_to_tn[v] = partition[v[1][1]] end - log_accumulated_norm = 0.0 + log_acc_norm = 0.0 for (ii, v) in enumerate(contractions) @info "$(ii)/$(length(contractions)) th tree approximation" c1, c2 = child_vertices(contraction_tree, v) - tn = disjoint_union(v_to_tn[c1], v_to_tn[c2]) - # TODO: rename tn since the names will be too long. - # TODO: swap size - p_edges = v_to_ordered_p_edges[v][2] + tn = ITensorNetwork() + ts = vcat(Vector{ITensor}(v_to_tn[c1]), Vector{ITensor}(v_to_tn[c2])) + for (i, t) in enumerate(ts) + add_vertex!(tn, i) + tn[i] = t + end + ref_p_edges, p_edges, contract_edges = v_to_ordered_p_edges[v] if p_edges == [] - # TODO: edge case with output being a scalar + @assert v == contractions[end] + out = _optcontract( + Vector{ITensor}(tn); + contraction_sequence_alg=contraction_sequence_alg, + contraction_sequence_kwargs=contraction_sequence_kwargs, + ) + out_nrm = norm(out) + out /= out_nrm + return ITensorNetwork(out), log_acc_norm + log(out_nrm) + end + for edge_order in _intervals(ref_p_edges, p_edges; size=swap_size) + inds_ordering = map(e -> p_edge_to_ordered_inds[e], edge_order) + ortho_center = if edge_order == p_edges + _ortho_center(p_edges, contract_edges) + else + div(length(edge_order), 2, RoundUp) + end + tn, log_norm = approx_itensornetwork( + tn, + _ansatz_tree(inds_ordering, ansatz, ortho_center); + alg=approx_itensornetwork_alg, + cutoff=cutoff, + maxdim=maxdim, + contraction_sequence_alg=contraction_sequence_alg, + contraction_sequence_kwargs=contraction_sequence_kwargs, + ) + log_acc_norm += log_norm end - inds_orderings = [p_edge_to_ordered_inds[e] for e in p_edges] - v_to_tn[v], log_root_norm = approx_itensornetwork( - tn, - _ansatz_tree(inds_orderings, _ortho_center(v_to_ordered_p_edges[v]), ansatz); - alg=approx_itensornetwork_alg, - cutoff=cutoff, - maxdim=maxdim, - contraction_sequence_alg=contraction_sequence_alg, - contraction_sequence_kwargs=contraction_sequence_kwargs, - ) - log_accumulated_norm += log_root_norm + v_to_tn[v] = tn end + return v_to_tn[contractions[end]], log_acc_norm end end @@ -99,14 +123,39 @@ function _ind_orderings(partition::DataGraph) # TODO: use `_mps_partition_inds_order` with default `alg` and `backend` end -function _ansatz_tree(inds_orderings::Vector, ortho_center::Integer, ansatz::String) +function _ansatz_tree(inds_orderings::Vector, ansatz::String, ortho_center::Integer) # TODO end -function _ortho_center(edges_to_contract_edges) +function _ortho_center(edges::Vector, contract_edges::Vector) # TODO end +function _permute(v::Vector, perms) + v = copy(v) + for p in perms + temp = v[p] + v[p] = v[p + 1] + v[p + 1] = temp + end + return v +end + +function _intervals(v1::Vector, v2::Vector; size) + if v1 == v2 + return [v2] + end + perms_list = collect(Iterators.partition(_bubble_sort(v1, v2), size)) + out = [v1] + current = v1 + for perms in perms_list + v = _permute(current, perms) + push!(out, v) + current = v + end + return out +end + # function ordered_igs_to_binary_tree(ordered_igs, contract_igs, ig_to_linear_order; ansatz) # @assert ansatz in ["comb", "mps"] # @timeit_debug ITensors.timer "ordered_igs_to_binary_tree" begin From d685eb5e2aeba46ff57d4797956e336e9849f2c2 Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Mon, 3 Jul 2023 16:35:17 -0500 Subject: [PATCH 16/35] Finish _ansatz_tree --- .../constrained_ordering.jl | 18 +-- .../partitioned_contract.jl | 118 ++++++------------ src/partitioned_contract/utils.jl | 14 +++ 3 files changed, 54 insertions(+), 96 deletions(-) diff --git a/src/partitioned_contract/constrained_ordering.jl b/src/partitioned_contract/constrained_ordering.jl index dc5371f0..72c98b87 100644 --- a/src/partitioned_contract/constrained_ordering.jl +++ b/src/partitioned_contract/constrained_ordering.jl @@ -41,8 +41,8 @@ function _constrained_minswap_inds_ordering( tn::ITensorNetwork, ) inter_igs = intersect(input_order_1, input_order_2) - left_1, right_1 = _split_vector(input_order_1, inter_igs) - left_2, right_2 = _split_vector(input_order_2, inter_igs) + left_1, right_1 = _split_array(input_order_1, inter_igs) + left_2, right_2 = _split_array(input_order_2, inter_igs) # @info "lengths of the input partitions", # sort([length(left_1), length(right_2), length(left_2), length(right_2)]) num_swaps_1 = min(length(left_1), length(left_2)) + min(length(right_1), length(right_2)) @@ -120,17 +120,3 @@ function _low_swap_merge(l1_left, l1_right, l2_left, l2_right) end return _make_list(left_lists, right_lists) end - -function _split_vector(v::Vector, sub_v::Vector) - left_v = [] - right_v = [] - target_array = left_v - for i in v - if i in sub_v - target_array = right_v - continue - end - push!(target_array, i) - end - return left_v, right_v -end diff --git a/src/partitioned_contract/partitioned_contract.jl b/src/partitioned_contract/partitioned_contract.jl index 6146a6de..fa880dfb 100644 --- a/src/partitioned_contract/partitioned_contract.jl +++ b/src/partitioned_contract/partitioned_contract.jl @@ -96,10 +96,12 @@ function partitioned_contract( end for edge_order in _intervals(ref_p_edges, p_edges; size=swap_size) inds_ordering = map(e -> p_edge_to_ordered_inds[e], edge_order) + # `ortho_center` denotes the number of edges arranged at + # the left of the center vertex. ortho_center = if edge_order == p_edges _ortho_center(p_edges, contract_edges) else - div(length(edge_order), 2, RoundUp) + div(length(edge_order), 2, RoundDown) end tn, log_norm = approx_itensornetwork( tn, @@ -124,11 +126,44 @@ function _ind_orderings(partition::DataGraph) end function _ansatz_tree(inds_orderings::Vector, ansatz::String, ortho_center::Integer) - # TODO + @assert ansatz in ["comb", "mps"] + nested_vec_left, nested_vec_right = _nested_vector_pair( + Algorithm(ansatz), inds_orderings, ortho_center + ) + if nested_vec_left == [] || nested_vec_right == [] + nested_vec = nested_vec_left == [] ? nested_vec_right : nested_vec_left + else + nested_vec = [nested_vec_left, nested_vec_right] + end + return _nested_vector_to_digraph(nested_vec) +end + +function _nested_vector_pair( + ::Algorithm"mps", inds_orderings::Vector, ortho_center::Integer +) + nested_vec_left = line_to_tree(vcat(inds_orderings[1:ortho_center]...)) + nested_vec_right = line_to_tree(reverse(vcat(inds_orderings[(ortho_center + 1):end]...))) + return nested_vec_left, nested_vec_right +end + +function _nested_vector_pair( + ::Algorithm"comb", inds_orderings::Vector, ortho_center::Integer +) + nested_vec_left = line_to_tree( + map(ig -> line_to_tree(ig), inds_orderings[1:ortho_center]) + ) + nested_vec_right = line_to_tree( + map(ig -> line_to_tree(ig), inds_orderings[end:-1:(ortho_center + 1)]) + ) + return nested_vec_left, nested_vec_right end function _ortho_center(edges::Vector, contract_edges::Vector) - # TODO + left_edges, right_edges = _split_array(edges, contract_edges) + if length(left_edges) < length(right_edges) + return length(left_edges) + length(contract_edges) + end + return length(left_edges) end function _permute(v::Vector, perms) @@ -155,80 +190,3 @@ function _intervals(v1::Vector, v2::Vector; size) end return out end - -# function ordered_igs_to_binary_tree(ordered_igs, contract_igs, ig_to_linear_order; ansatz) -# @assert ansatz in ["comb", "mps"] -# @timeit_debug ITensors.timer "ordered_igs_to_binary_tree" begin -# if contract_igs == [] -# @info "contract_igs is empty vector" -# end -# # @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(igs, ig_to_linear_order; ansatz, direction) -# @assert ansatz in ["comb", "mps"] -# @assert direction in ["left", "right"] -# if ansatz == "comb" -# return line_to_tree([line_to_tree(ig_to_linear_order[ig]) for ig in igs]) -# end -# if direction == "left" -# order = vcat([ig_to_linear_order[ig] for ig in igs]...) -# return line_to_tree(order) -# else -# # First reverse get the order from middle to boundary, -# # and second reverse get the overall inds order from boundary to middle. -# order = vcat([ig_to_linear_order[ig] for ig in reverse(igs)]...) -# return line_to_tree(reverse(order)) -# 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 = ordered_igs_to_binary_tree( -# left_igs, ig_to_linear_order; ansatz="comb", direction="left" -# ) -# tree_contract = ordered_igs_to_binary_tree( -# contract_igs, ig_to_linear_order; ansatz="comb", direction="left" -# ) -# tree_2 = ordered_igs_to_binary_tree( -# reverse(right_igs), ig_to_linear_order; ansatz="comb", direction="left" -# ) -# # 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 diff --git a/src/partitioned_contract/utils.jl b/src/partitioned_contract/utils.jl index ebbccef9..60a4d24e 100644 --- a/src/partitioned_contract/utils.jl +++ b/src/partitioned_contract/utils.jl @@ -31,6 +31,20 @@ function _replace_subarray(v1::Vector, v2::Vector) return v1 end +function _split_array(v::Vector, sub_v::Vector) + left_v = [] + right_v = [] + target_array = left_v + for i in v + if i in sub_v + target_array = right_v + continue + end + push!(target_array, i) + end + return left_v, right_v +end + function _neighbor_edges(graph, vs) return filter( e -> (e.src in vs && !(e.dst in vs)) || (e.dst in vs && !(e.src in vs)), edges(graph) From 7c4405ec5e340118ab93f18737a373559a50fda6 Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Mon, 3 Jul 2023 22:46:54 -0500 Subject: [PATCH 17/35] finish _ind_orderings --- .../partitioned_contract.jl | 34 ++++++++++++++++--- 1 file changed, 30 insertions(+), 4 deletions(-) diff --git a/src/partitioned_contract/partitioned_contract.jl b/src/partitioned_contract/partitioned_contract.jl index fa880dfb..eeb58cfd 100644 --- a/src/partitioned_contract/partitioned_contract.jl +++ b/src/partitioned_contract/partitioned_contract.jl @@ -9,11 +9,16 @@ function partitioned_contract( contraction_sequence_alg, contraction_sequence_kwargs, ) + input_tn = ITensorNetwork( + mapreduce(l -> Vector{ITensor}(partition[l]), vcat, vertices(partition)) + ) + # TODO: currently we assume that the tn represented by 'partition' is closed. + @assert noncommoninds(Vector{ITensor}(input_tn)...) == [] @timeit_debug ITensors.timer "partitioned_contract" begin leaves = leaf_vertices(contraction_tree) traversal = post_order_dfs_vertices(contraction_tree, _root(contraction_tree)) contractions = setdiff(traversal, leaves) - p_edge_to_ordered_inds = _ind_orderings(partition) + p_edge_to_ordered_inds = _ind_orderings(partition, contractions) # Build the orderings used for the ansatz tree. # For each tuple in `v_to_ordered_p_edges`, the first item is the # reference ordering of uncontracted edges for the contraction `v`, @@ -120,9 +125,30 @@ function partitioned_contract( end end -function _ind_orderings(partition::DataGraph) - # input_tn = # TODO - # TODO: use `_mps_partition_inds_order` with default `alg` and `backend` +# Note: currently this function assumes that the input tn represented by +# 'partition' is closed +function _ind_orderings(partition::DataGraph, contractions::Vector) + p_edge_to_ordered_inds = Dict{Any,Vector}() + for v in contractions + contract_edges = intersect( + _neighbor_edges(partition, v[1]), _neighbor_edges(partition, v[2]) + ) + tn1 = ITensorNetwork(mapreduce(l -> Vector{ITensor}(partition[l]), vcat, v[1])) + tn2 = ITensorNetwork(mapreduce(l -> Vector{ITensor}(partition[l]), vcat, v[2])) + tn = length(vertices(tn1)) >= length(vertices(tn2)) ? tn1 : tn2 + tn_inds = noncommoninds(Vector{ITensor}(tn)...) + for e in contract_edges + inds1 = noncommoninds(Vector{ITensor}(partition[e.src])...) + inds2 = noncommoninds(Vector{ITensor}(partition[e.dst])...) + source_inds = intersect(inds1, inds2) + @assert issubset(source_inds, tn_inds) + terminal_inds = setdiff(tn_inds, source_inds) + p, _ = _mincut_partitions(tn, source_inds, terminal_inds) + sub_tn = subgraph(u -> u in p, tn) + p_edge_to_ordered_inds[e] = _mps_partition_inds_order(sub_tn, source_inds) + end + end + return p_edge_to_ordered_inds end function _ansatz_tree(inds_orderings::Vector, ansatz::String, ortho_center::Integer) From d8d2fcda9699a0dede15ef19566099c77ac8d681 Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Tue, 4 Jul 2023 11:41:04 -0500 Subject: [PATCH 18/35] A new partitioned_contract interface --- src/ITensorNetworks.jl | 5 +- src/exports.jl | 3 + src/partitioned_contract/adjacency_tree.jl | 10 ++- .../partitioned_contract.jl | 62 ++++++++++++------- 4 files changed, 52 insertions(+), 28 deletions(-) diff --git a/src/ITensorNetworks.jl b/src/ITensorNetworks.jl index 2f578df3..cd7468b9 100644 --- a/src/ITensorNetworks.jl +++ b/src/ITensorNetworks.jl @@ -87,8 +87,11 @@ include("binary_tree_partition.jl") include(joinpath("approx_itensornetwork", "utils.jl")) include(joinpath("approx_itensornetwork", "density_matrix.jl")) include(joinpath("approx_itensornetwork", "ttn_svd.jl")) -# include(joinpath("approx_itensornetwork", "partitioned_contraction.jl")) include(joinpath("approx_itensornetwork", "approx_itensornetwork.jl")) +include(joinpath("partitioned_contract", "utils.jl")) +include(joinpath("partitioned_contract", "adjacency_tree.jl")) +include(joinpath("partitioned_contract", "constrained_ordering.jl")) +include(joinpath("partitioned_contract", "partitioned_contract.jl")) include("contract.jl") include("utility.jl") include("specialitensornetworks.jl") diff --git a/src/exports.jl b/src/exports.jl index 5f08950e..9305aabd 100644 --- a/src/exports.jl +++ b/src/exports.jl @@ -104,6 +104,9 @@ export path_graph_structure, binary_tree_structure # ITensorNetworks: approx_itensornetwork.jl export approx_itensornetwork +# ITensorNetworks: partitioned_contract.jl +export partitioned_contract + # ITensorNetworks: lattices.jl # TODO: DELETE export hypercubic_lattice_graph, square_lattice_graph, chain_lattice_graph diff --git a/src/partitioned_contract/adjacency_tree.jl b/src/partitioned_contract/adjacency_tree.jl index 1c1a46d0..39a500c6 100644 --- a/src/partitioned_contract/adjacency_tree.jl +++ b/src/partitioned_contract/adjacency_tree.jl @@ -184,14 +184,14 @@ end # Generate the adjacency tree of a contraction tree # TODO: add test -function _adjacency_tree(v::Tuple, path::Vector, partition::DataGraph, p_edge_to_inds::Dict) +function _adjacency_tree(v::Tuple, path::Vector, par::DataGraph, p_edge_to_inds::Dict) @timeit_debug ITensors.timer "_generate_adjacency_tree" begin # mapping each index group to adjacent input igs ig_to_input_adj_igs = Dict{Any,Set}() # mapping each igs to an adjacency tree adjacency_tree = NamedDiGraph{Tuple{Tuple,String}}() p_leaves = vcat(v[1:(end - 1)]...) - p_edges = _neighbor_edges(partition, p_leaves) + p_edges = _neighbor_edges(par, p_leaves) for ig in map(e -> Set(p_edge_to_inds[e]), p_edges) ig_to_input_adj_igs[ig] = Set([ig]) v = ((ig,), "unordered") @@ -202,12 +202,10 @@ function _adjacency_tree(v::Tuple, path::Vector, partition::DataGraph, p_edge_to ancester = filter(u -> p_leaves in vcat(u[1:(end - 1)]...), children)[1] sibling = setdiff(children, [ancester])[1] ancester_igs = map( - e -> Set(p_edge_to_inds[e]), - _neighbor_edges(partition, vcat(ancester[1:(end - 1)]...)), + e -> Set(p_edge_to_inds[e]), _neighbor_edges(par, vcat(ancester[1:(end - 1)]...)) ) sibling_igs = map( - e -> Set(p_edge_to_inds[e]), - _neighbor_edges(partition, vcat(sibling[1:(end - 1)]...)), + e -> Set(p_edge_to_inds[e]), _neighbor_edges(par, vcat(sibling[1:(end - 1)]...)) ) inter_igs = intersect(ancester_igs, sibling_igs) new_igs = setdiff(sibling_igs, inter_igs) diff --git a/src/partitioned_contract/partitioned_contract.jl b/src/partitioned_contract/partitioned_contract.jl index eeb58cfd..59c35410 100644 --- a/src/partitioned_contract/partitioned_contract.jl +++ b/src/partitioned_contract/partitioned_contract.jl @@ -1,5 +1,30 @@ +function partitioned_contract(nested_vector::Vector; kwargs...) + leaves = filter( + v -> v isa Vector && v[1] isa ITensor, collect(PreOrderDFS(nested_vector)) + ) + tn = ITensorNetwork() + subgraph_vs = Vector{<:Vector}() + i = 1 + for ts in leaves + vs = [] + for t in ts + add_vertex!(tn, i) + tn[i] = t + push!(vs, i) + i += 1 + end + push!(subgraph_vs, vs) + end + par = partition(tn, subgraph_vs) + vs_nested_vector = _map_nested_vector( + Dict(leaves .=> collect(1:length(leaves))), nested_vector + ) + contraction_tree = contraction_sequence_to_digraph(vs_nested_vector) + return partitioned_contract(par, contraction_tree; kwargs...) +end + function partitioned_contract( - partition::DataGraph, + par::DataGraph, contraction_tree::NamedDiGraph; ansatz="mps", approx_itensornetwork_alg="density_matrix", @@ -9,16 +34,14 @@ function partitioned_contract( contraction_sequence_alg, contraction_sequence_kwargs, ) - input_tn = ITensorNetwork( - mapreduce(l -> Vector{ITensor}(partition[l]), vcat, vertices(partition)) - ) - # TODO: currently we assume that the tn represented by 'partition' is closed. + input_tn = ITensorNetwork(mapreduce(l -> Vector{ITensor}(par[l]), vcat, vertices(par))) + # TODO: currently we assume that the tn represented by 'par' is closed. @assert noncommoninds(Vector{ITensor}(input_tn)...) == [] @timeit_debug ITensors.timer "partitioned_contract" begin leaves = leaf_vertices(contraction_tree) traversal = post_order_dfs_vertices(contraction_tree, _root(contraction_tree)) contractions = setdiff(traversal, leaves) - p_edge_to_ordered_inds = _ind_orderings(partition, contractions) + p_edge_to_ordered_inds = _ind_orderings(par, contractions) # Build the orderings used for the ansatz tree. # For each tuple in `v_to_ordered_p_edges`, the first item is the # reference ordering of uncontracted edges for the contraction `v`, @@ -29,12 +52,12 @@ function partitioned_contract( for (ii, v) in enumerate(traversal) @info "$(ii)/$(length(traversal)) th ansatz tree construction" p_leaves = vcat(v[1:(end - 1)]...) - tn = ITensorNetwork(mapreduce(l -> Vector{ITensor}(partition[l]), vcat, p_leaves)) + tn = ITensorNetwork(mapreduce(l -> Vector{ITensor}(par[l]), vcat, p_leaves)) path = filter(u -> issubset(p_leaves, u[1]) || issubset(p_leaves, u[2]), contractions) - p_edges = _neighbor_edges(partition, p_leaves) + p_edges = _neighbor_edges(par, p_leaves) # Get the reference ordering and target ordering of `v` v_inds = map(e -> Set(p_edge_to_ordered_inds[e]), p_edges) - constraint_tree = _adjacency_tree(v, path, partition, p_edge_to_ordered_inds) + constraint_tree = _adjacency_tree(v, path, par, p_edge_to_ordered_inds) if v in leaves ref_inds_ordering = _mps_partition_inds_order(tn, v_inds; alg="top_down") inds_ordering = _constrained_minswap_inds_ordering( @@ -66,7 +89,7 @@ function partitioned_contract( p_edges = _replace_subarray(p_edges, contract_edges) else p_leaves_2 = vcat(sibling[1:(end - 1)]...) - p_edges_2 = _neighbor_edges(partition, p_leaves_2) + p_edges_2 = _neighbor_edges(par, p_leaves_2) contract_edges = intersect(p_edges, p_edges_2) end v_to_ordered_p_edges[v] = (ref_p_edges, p_edges, contract_edges) @@ -75,7 +98,7 @@ function partitioned_contract( v_to_tn = Dict{Tuple,ITensorNetwork}() for v in leaves @assert length(v[1]) == 1 - v_to_tn[v] = partition[v[1][1]] + v_to_tn[v] = par[v[1][1]] end log_acc_norm = 0.0 for (ii, v) in enumerate(contractions) @@ -125,21 +148,18 @@ function partitioned_contract( end end -# Note: currently this function assumes that the input tn represented by -# 'partition' is closed -function _ind_orderings(partition::DataGraph, contractions::Vector) +# Note: currently this function assumes that the input tn represented by 'par' is closed +function _ind_orderings(par::DataGraph, contractions::Vector) p_edge_to_ordered_inds = Dict{Any,Vector}() for v in contractions - contract_edges = intersect( - _neighbor_edges(partition, v[1]), _neighbor_edges(partition, v[2]) - ) - tn1 = ITensorNetwork(mapreduce(l -> Vector{ITensor}(partition[l]), vcat, v[1])) - tn2 = ITensorNetwork(mapreduce(l -> Vector{ITensor}(partition[l]), vcat, v[2])) + contract_edges = intersect(_neighbor_edges(par, v[1]), _neighbor_edges(par, v[2])) + tn1 = ITensorNetwork(mapreduce(l -> Vector{ITensor}(par[l]), vcat, v[1])) + tn2 = ITensorNetwork(mapreduce(l -> Vector{ITensor}(par[l]), vcat, v[2])) tn = length(vertices(tn1)) >= length(vertices(tn2)) ? tn1 : tn2 tn_inds = noncommoninds(Vector{ITensor}(tn)...) for e in contract_edges - inds1 = noncommoninds(Vector{ITensor}(partition[e.src])...) - inds2 = noncommoninds(Vector{ITensor}(partition[e.dst])...) + inds1 = noncommoninds(Vector{ITensor}(par[e.src])...) + inds2 = noncommoninds(Vector{ITensor}(par[e.dst])...) source_inds = intersect(inds1, inds2) @assert issubset(source_inds, tn_inds) terminal_inds = setdiff(tn_inds, source_inds) From 95eed58d344d240ba2f2378daad7dd79b4be72e6 Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Tue, 4 Jul 2023 15:28:18 -0500 Subject: [PATCH 19/35] Rewrite _recursive_bisection to better handle the case with target_set being a vector of inds --- src/mincut.jl | 48 +++++++++++++-------------- test/test_binary_tree_partition.jl | 53 +++++++++++++++++------------- 2 files changed, 53 insertions(+), 48 deletions(-) diff --git a/src/mincut.jl b/src/mincut.jl index 227956cb..bc362600 100644 --- a/src/mincut.jl +++ b/src/mincut.jl @@ -123,9 +123,7 @@ function _binary_tree_structure( maximally_unbalanced::Bool=false, backend="Metis", ) - nested_vector = _binary_tree_partition_inds_recursive_bisection( - tn, outinds; backend=backend - ) + nested_vector = _recursive_bisection(tn, outinds; backend=backend) if maximally_unbalanced ordering = collect(Leaves(nested_vector)) nested_vector = line_to_tree(ordering) @@ -140,9 +138,7 @@ function _map_nested_vector(dict::Dict, nested_vector) return map(v -> _map_nested_vector(dict, v), nested_vector) end -function _binary_tree_partition_inds_recursive_bisection( - tn::ITensorNetwork, outinds::Union{Vector{Set},Vector{<:Index}}; backend="Metis" -) +function _recursive_bisection(tn::ITensorNetwork, outinds::Vector{Set}; backend="Metis") tn = copy(tn) tensor_to_inds = Dict() ts = Vector{ITensor}() @@ -152,20 +148,24 @@ function _binary_tree_partition_inds_recursive_bisection( tensor_to_inds[new_t] = is end new_tn = disjoint_union(tn, ITensorNetwork(ts)) - ts_nested_vector = _tensor_order_recursive_bisection(new_tn, ts; backend=backend) + ts_nested_vector = _recursive_bisection(new_tn, ts; backend=backend) return _map_nested_vector(tensor_to_inds, ts_nested_vector) end -function _tensor_order_recursive_bisection( +function _recursive_bisection( tn::ITensorNetwork, - tensors::Vector{ITensor}; + target_set::Union{Vector{ITensor},Vector{<:Index}}; backend="Metis", left_inds=Set(), right_inds=Set(), ) - ts = intersect(tensors, Vector{ITensor}(tn)) - if length(ts) <= 1 - return length(ts) == 1 ? ts[1] : nothing + if target_set isa Vector{ITensor} + set = intersect(target_set, Vector{ITensor}(tn)) + else + set = intersect(target_set, noncommoninds(Vector{ITensor}(tn)...)) + end + if length(set) <= 1 + return length(set) == 1 ? set[1] : nothing end g_parts = partition(tn; npartitions=2, backend=backend) # order g_parts @@ -179,24 +179,24 @@ function _tensor_order_recursive_bisection( g_parts[1], g_parts[2] = g_parts[2], g_parts[1] end intersect_inds = intersect(outinds_1, outinds_2) - ts1 = _tensor_order_recursive_bisection( + set1 = _recursive_bisection( g_parts[1], - tensors; + target_set; backend=backend, left_inds=left_inds, right_inds=union(right_inds, intersect_inds), ) - ts2 = _tensor_order_recursive_bisection( + set2 = _recursive_bisection( g_parts[2], - tensors; + target_set; backend=backend, left_inds=union(left_inds, intersect_inds), right_inds=right_inds, ) - if ts1 == nothing || ts2 == nothing - return ts1 == nothing ? ts2 : ts1 + if set1 == nothing || set2 == nothing + return set1 == nothing ? set2 : set1 end - return [ts1, ts2] + return [set1, set2] end """ @@ -283,9 +283,7 @@ function _mps_partition_inds_order( tn => tn2, out_to_maxweight_ind ) else - nested_vector = _binary_tree_partition_inds_recursive_bisection( - tn, outinds; backend=backend - ) + nested_vector = _recursive_bisection(tn, outinds; backend=backend) return filter(v -> v in outinds, collect(PreOrderDFS(nested_vector))) end end @@ -385,18 +383,18 @@ function _mincut_inds( out_to_maxweight_ind::Dict{Index,Index}, sourceinds_list::Vector{<:Vector{<:Index}}, ) - weights = _partition_mincuts_dist(tn_pair, out_to_maxweight_ind, sourceinds_list) + weights = _mincut_partitions_costs(tn_pair, out_to_maxweight_ind, sourceinds_list) i = argmin(weights) return sourceinds_list[i], i end function _mps_mincut_partition_cost(tn::ITensorNetwork, inds_vector::Vector{Set}) - @timeit_debug ITensors.timer "_comb_mincuts_dist" begin + @timeit_debug ITensors.timer "_mps_mincut_partition_cost" begin inds_vector = map(inds -> collect(inds), inds_vector) outinds = vcat(inds_vector...) maxweight_tn, out_to_maxweight_ind = _maxweightoutinds_tn(tn, outinds) sourceinds_list = [vcat(inds_vector[1:i]...) for i in 1:(length(inds_vector) - 1)] - weights = _partition_mincuts_dist( + weights = _mincut_partitions_costs( tn => maxweight_tn, out_to_maxweight_ind, sourceinds_list ) mincut_val = sum([w[1] for w in weights]) diff --git a/test/test_binary_tree_partition.jl b/test/test_binary_tree_partition.jl index 50287ad5..5160983d 100644 --- a/test/test_binary_tree_partition.jl +++ b/test/test_binary_tree_partition.jl @@ -12,34 +12,40 @@ using ITensorNetworks: _DensityMartrixAlgGraph @testset "test mincut and partition functions on top of MPS" 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) - tn = ITensorNetwork(M[:]) - for alg in ["top_down", "bottom_up"] - for out in [binary_tree_structure(tn; alg=alg), path_graph_structure(tn; alg=alg)] - @test out isa DataGraph - @test _is_rooted_directed_binary_tree(out) - @test length(vertex_data(out).values) == 8 + function _build_tn_inds(num) + is = [Index(2, string(i)) for i in 1:num] + T = randomITensor(is...) + M = MPS(T, Tuple(is); cutoff=1e-5, maxdim=500) + tn = ITensorNetwork(M[:]) + return tn, is + end + for num in [6, 8] + tn, is = _build_tn_inds(num) + for alg in ["top_down", "bottom_up"] + for out in [binary_tree_structure(tn; alg=alg), path_graph_structure(tn; alg=alg)] + @test out isa DataGraph + @test _is_rooted_directed_binary_tree(out) + @test length(vertex_data(out).values) == num + end + out = _mps_partition_inds_order(tn, is; alg=alg) + @test out in [is, reverse(is)] end - out = _mps_partition_inds_order(tn, [o, p, i, j, k, l, m, n]; alg=alg) - @test out in [[i, j, k, l, m, n, o, p], [p, o, n, m, l, k, j, i]] end - p1, p2 = _mincut_partitions(tn, [k, l], [m, n]) +end + +@testset "test _mincut_partitions" begin + is = [Index(2, string(i)) for i in 1:8] + T = randomITensor(is...) + + M = MPS(T, Tuple(is); cutoff=1e-5, maxdim=500) + tn = ITensorNetwork(M[:]) + p1, p2 = _mincut_partitions(tn, [is[3], is[4]], [is[5], is[6]]) # When MPS bond dimensions are large, the partition will not across internal inds @test (length(p1) == 0) || (length(p2) == 0) - M = MPS(T, (i, j, k, l, m, n, o, p); cutoff=1e-5, maxdim=2) + M = MPS(T, Tuple(is); cutoff=1e-5, maxdim=2) tn = ITensorNetwork(M[:]) - p1, p2 = _mincut_partitions(tn, [k, l], [m, n]) + p1, p2 = _mincut_partitions(tn, [is[3], is[4]], [is[5], is[6]]) # When MPS bond dimensions are small, the partition will across internal inds @test sort(p1) == [1, 2, 3, 4] @test sort(p2) == [5, 6, 7, 8] @@ -56,7 +62,8 @@ end tensors = vec(tn)[1:20] tn = ITensorNetwork(tensors) @info _mps_partition_inds_order(tn, noncommoninds(tensors...); alg="top_down") - # TODO: think about how to check this + # TODO: this case is not supported for now, since two indices are adjacent to + # one tensor. end @testset "test _binary_tree_partition_inds of a 2D network" begin From 37cdd5f9b6f410b8a6d5667c9ab29265bebb97cf Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Tue, 4 Jul 2023 15:38:46 -0500 Subject: [PATCH 20/35] Add _roots --- src/Graphs/abstractgraph.jl | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/src/Graphs/abstractgraph.jl b/src/Graphs/abstractgraph.jl index 91c0a5b1..80a7450c 100644 --- a/src/Graphs/abstractgraph.jl +++ b/src/Graphs/abstractgraph.jl @@ -38,17 +38,17 @@ end Return the root vertex of a rooted directed graph """ @traitfn function _root(graph::AbstractGraph::IsDirected) - @assert _is_rooted(graph) "the input $(graph) has to be rooted" - v = vertices(graph)[1] - while parent_vertex(graph, v) != nothing - v = parent_vertex(graph, v) - end - return v + __roots = _roots(graph) + @assert length(__roots) == 1 "the input $(graph) has to be rooted" + return __roots[1] +end + +@traitfn function _roots(graph::AbstractGraph::IsDirected) + return [v for v in vertices(graph) if parent_vertex(graph, v) == nothing] end @traitfn function _is_rooted(graph::AbstractGraph::IsDirected) - roots = [v for v in vertices(graph) if parent_vertex(graph, v) == nothing] - return length(roots) == 1 + return length(_roots(graph)) == 1 end @traitfn function _is_rooted_directed_binary_tree(graph::AbstractGraph::IsDirected) From a638eae7365311d083309e1be93f442e5a875b4f Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Tue, 4 Jul 2023 15:43:05 -0500 Subject: [PATCH 21/35] Fix multiple bugs in partitioned_contract --- src/partitioned_contract/adjacency_tree.jl | 13 ++++-------- .../constrained_ordering.jl | 18 ++++++++-------- .../partitioned_contract.jl | 21 +++++++++++++++---- src/partitioned_contract/utils.jl | 2 +- 4 files changed, 31 insertions(+), 23 deletions(-) diff --git a/src/partitioned_contract/adjacency_tree.jl b/src/partitioned_contract/adjacency_tree.jl index 39a500c6..0395ecd8 100644 --- a/src/partitioned_contract/adjacency_tree.jl +++ b/src/partitioned_contract/adjacency_tree.jl @@ -198,15 +198,10 @@ function _adjacency_tree(v::Tuple, path::Vector, par::DataGraph, p_edge_to_inds: add_vertex!(adjacency_tree, v) end for contraction in path - children = child_vertices(contractions, path[1]) - ancester = filter(u -> p_leaves in vcat(u[1:(end - 1)]...), children)[1] - sibling = setdiff(children, [ancester])[1] - ancester_igs = map( - e -> Set(p_edge_to_inds[e]), _neighbor_edges(par, vcat(ancester[1:(end - 1)]...)) - ) - sibling_igs = map( - e -> Set(p_edge_to_inds[e]), _neighbor_edges(par, vcat(sibling[1:(end - 1)]...)) - ) + ancester_leaves = filter(u -> issubset(p_leaves, u), contraction[1:2])[1] + sibling_leaves = setdiff(contraction[1:2], [ancester_leaves])[1] + ancester_igs = map(e -> Set(p_edge_to_inds[e]), _neighbor_edges(par, ancester_leaves)) + sibling_igs = map(e -> Set(p_edge_to_inds[e]), _neighbor_edges(par, sibling_leaves)) inter_igs = intersect(ancester_igs, sibling_igs) new_igs = setdiff(sibling_igs, inter_igs) adjacent_igs = union([ig_to_input_adj_igs[ig] for ig in inter_igs]...) diff --git a/src/partitioned_contract/constrained_ordering.jl b/src/partitioned_contract/constrained_ordering.jl index 72c98b87..4ffe98cc 100644 --- a/src/partitioned_contract/constrained_ordering.jl +++ b/src/partitioned_contract/constrained_ordering.jl @@ -1,17 +1,18 @@ +# TODO: test needed function _constrained_minswap_inds_ordering( constraint_tree::NamedDiGraph{Tuple{Tuple,String}}, - ref_ordering::Vector{Set}, + ref_ordering::Vector, tn::ITensorNetwork, ) leaves = leaf_vertices(constraint_tree) root = _root(constraint_tree) - v_to_order = Dict{Tuple{Tuple,String},Vector{IndexGroup}}() + v_to_order = Dict{Tuple{Tuple,String},Vector{Set}}() for v in post_order_dfs_vertices(constraint_tree, root) if v in leaves v_to_order[v] = [v[1]...] continue end - child_orders = Vector{Vector{IndexGroup}}() + child_orders = Vector{Vector{Set}}() children = child_vertices(constraint_tree, v) for inds_tuple in v[1] cs = filter(c -> c[1] == inds_tuple, children) @@ -30,14 +31,13 @@ function _constrained_minswap_inds_ordering( end v_to_order[v] = vcat(output_order...) end - nswap = length(_bubble_sort(v_to_order[root], ref_ordering)) - return v_to_order[root], nswap + return v_to_order[root] end function _constrained_minswap_inds_ordering( constraint_tree::NamedDiGraph{Tuple{Tuple,String}}, - input_order_1::Vector{Set}, - input_order_2::Vector{Set}, + input_order_1::Vector, + input_order_2::Vector, tn::ITensorNetwork, ) inter_igs = intersect(input_order_1, input_order_2) @@ -60,9 +60,9 @@ function _constrained_minswap_inds_ordering( outputs = [] nswaps_list = [] for (t, i) in zip(adj_tree_copies, inputs) - output, nswaps = _constrained_minswap_inds_ordering(t, i, tn) + output = _constrained_minswap_inds_ordering(t, i, tn) push!(outputs, output) - push!(nswaps_list, nswaps) + push!(nswaps_list, length(_bubble_sort(output, i))) end inputs = [inputs[i] for i in 1:length(inputs) if nswaps_list[i] == min(nswaps_list...)] outputs = [outputs[i] for i in 1:length(outputs) if nswaps_list[i] == min(nswaps_list...)] diff --git a/src/partitioned_contract/partitioned_contract.jl b/src/partitioned_contract/partitioned_contract.jl index 59c35410..254c019f 100644 --- a/src/partitioned_contract/partitioned_contract.jl +++ b/src/partitioned_contract/partitioned_contract.jl @@ -3,7 +3,7 @@ function partitioned_contract(nested_vector::Vector; kwargs...) v -> v isa Vector && v[1] isa ITensor, collect(PreOrderDFS(nested_vector)) ) tn = ITensorNetwork() - subgraph_vs = Vector{<:Vector}() + subgraph_vs = [] i = 1 for ts in leaves vs = [] @@ -33,6 +33,7 @@ function partitioned_contract( swap_size=1, contraction_sequence_alg, contraction_sequence_kwargs, + linear_ordering_alg, ) input_tn = ITensorNetwork(mapreduce(l -> Vector{ITensor}(par[l]), vcat, vertices(par))) # TODO: currently we assume that the tn represented by 'par' is closed. @@ -41,7 +42,7 @@ function partitioned_contract( leaves = leaf_vertices(contraction_tree) traversal = post_order_dfs_vertices(contraction_tree, _root(contraction_tree)) contractions = setdiff(traversal, leaves) - p_edge_to_ordered_inds = _ind_orderings(par, contractions) + p_edge_to_ordered_inds = _ind_orderings(par, contractions; linear_ordering_alg) # Build the orderings used for the ansatz tree. # For each tuple in `v_to_ordered_p_edges`, the first item is the # reference ordering of uncontracted edges for the contraction `v`, @@ -59,7 +60,7 @@ function partitioned_contract( v_inds = map(e -> Set(p_edge_to_ordered_inds[e]), p_edges) constraint_tree = _adjacency_tree(v, path, par, p_edge_to_ordered_inds) if v in leaves - ref_inds_ordering = _mps_partition_inds_order(tn, v_inds; alg="top_down") + ref_inds_ordering = _mps_partition_inds_order(tn, v_inds; alg=linear_ordering_alg) inds_ordering = _constrained_minswap_inds_ordering( constraint_tree, ref_inds_ordering, tn ) @@ -149,7 +150,8 @@ function partitioned_contract( end # Note: currently this function assumes that the input tn represented by 'par' is closed -function _ind_orderings(par::DataGraph, contractions::Vector) +# TODO: test needed: +function _ind_orderings(par::DataGraph, contractions::Vector; linear_ordering_alg) p_edge_to_ordered_inds = Dict{Any,Vector}() for v in contractions contract_edges = intersect(_neighbor_edges(par, v[1]), _neighbor_edges(par, v[2])) @@ -162,10 +164,21 @@ function _ind_orderings(par::DataGraph, contractions::Vector) inds2 = noncommoninds(Vector{ITensor}(par[e.dst])...) source_inds = intersect(inds1, inds2) @assert issubset(source_inds, tn_inds) + # Note: another more efficient implementation is below, where we first get + # `sub_tn` of `tn` that is closely connected to `source_inds`, and then compute + # the ordering only on top of `sub_tn`. The problem is that for multiple graphs + # `sub_tn` will be empty. We keep the implementation below for reference. + #= terminal_inds = setdiff(tn_inds, source_inds) p, _ = _mincut_partitions(tn, source_inds, terminal_inds) sub_tn = subgraph(u -> u in p, tn) p_edge_to_ordered_inds[e] = _mps_partition_inds_order(sub_tn, source_inds) + =# + p_edge_to_ordered_inds[e] = _mps_partition_inds_order( + tn, source_inds; alg=linear_ordering_alg + ) + # @info "tn has size", length(vertices(tn)) + # @info "out ordering is", p_edge_to_ordered_inds[e] end end return p_edge_to_ordered_inds diff --git a/src/partitioned_contract/utils.jl b/src/partitioned_contract/utils.jl index 60a4d24e..00b0afc9 100644 --- a/src/partitioned_contract/utils.jl +++ b/src/partitioned_contract/utils.jl @@ -52,7 +52,7 @@ function _neighbor_edges(graph, vs) end # Find the permutation to change `v1` into `v2` -function _findperm(v1, v2) +function _findperm(v1::Vector, v2::Vector) index_to_number = Dict() for (i, v) in enumerate(v2) index_to_number[v] = i From 54fa283dda307071a357149a6235eb5a34085b0b Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Tue, 4 Jul 2023 23:09:03 -0500 Subject: [PATCH 22/35] Fix bugs --- src/mincut.jl | 6 +++--- src/partitioned_contract/constrained_ordering.jl | 8 ++++---- src/partitioned_contract/partitioned_contract.jl | 10 ++++++++-- src/partitioned_contract/utils.jl | 4 ++-- 4 files changed, 17 insertions(+), 11 deletions(-) diff --git a/src/mincut.jl b/src/mincut.jl index bc362600..a8d3152e 100644 --- a/src/mincut.jl +++ b/src/mincut.jl @@ -138,7 +138,7 @@ function _map_nested_vector(dict::Dict, nested_vector) return map(v -> _map_nested_vector(dict, v), nested_vector) end -function _recursive_bisection(tn::ITensorNetwork, outinds::Vector{Set}; backend="Metis") +function _recursive_bisection(tn::ITensorNetwork, outinds::Vector{<:Set}; backend="Metis") tn = copy(tn) tensor_to_inds = Dict() ts = Vector{ITensor}() @@ -266,7 +266,7 @@ Given a tn and outinds, returns a vector of indices representing MPS inds orderi """ function _mps_partition_inds_order( tn::ITensorNetwork, - outinds::Union{Nothing,Vector{<:Index},Vector{Set}}; + outinds::Union{Nothing,Vector{<:Index},Vector{<:Set}}; alg="top_down", backend="Metis", ) @@ -388,7 +388,7 @@ function _mincut_inds( return sourceinds_list[i], i end -function _mps_mincut_partition_cost(tn::ITensorNetwork, inds_vector::Vector{Set}) +function _mps_mincut_partition_cost(tn::ITensorNetwork, inds_vector::Vector{<:Set}) @timeit_debug ITensors.timer "_mps_mincut_partition_cost" begin inds_vector = map(inds -> collect(inds), inds_vector) outinds = vcat(inds_vector...) diff --git a/src/partitioned_contract/constrained_ordering.jl b/src/partitioned_contract/constrained_ordering.jl index 4ffe98cc..acb3a263 100644 --- a/src/partitioned_contract/constrained_ordering.jl +++ b/src/partitioned_contract/constrained_ordering.jl @@ -56,10 +56,10 @@ function _constrained_minswap_inds_ordering( else inputs = _low_swap_merge(left_1, right_1, left_2, right_2) end - adj_tree_copies = [copy(adj_tree) for _ in 1:length(inputs)] + constraint_tree_copies = [copy(constraint_tree) for _ in 1:length(inputs)] outputs = [] nswaps_list = [] - for (t, i) in zip(adj_tree_copies, inputs) + for (t, i) in zip(constraint_tree_copies, inputs) output = _constrained_minswap_inds_ordering(t, i, tn) push!(outputs, output) push!(nswaps_list, length(_bubble_sort(output, i))) @@ -73,11 +73,11 @@ function _constrained_minswap_inds_ordering( return inputs[argmin(mincuts)], outputs[argmin(mincuts)] end -function _mincut_permutation(perms::Vector{<:Vector}, tn::ITensorNetwork) +function _mincut_permutation(perms::Vector{<:Vector{<:Vector}}, tn::ITensorNetwork) if length(perms) == 1 return perms[1] end - mincuts_dist = map(p -> _mps_mincut_partition_cost(tn, p), perms) + mincuts_dist = map(p -> _mps_mincut_partition_cost(tn, vcat(p...)), perms) return perms[argmin(mincuts_dist)] end diff --git a/src/partitioned_contract/partitioned_contract.jl b/src/partitioned_contract/partitioned_contract.jl index 254c019f..a425e5ec 100644 --- a/src/partitioned_contract/partitioned_contract.jl +++ b/src/partitioned_contract/partitioned_contract.jl @@ -56,11 +56,17 @@ function partitioned_contract( tn = ITensorNetwork(mapreduce(l -> Vector{ITensor}(par[l]), vcat, p_leaves)) path = filter(u -> issubset(p_leaves, u[1]) || issubset(p_leaves, u[2]), contractions) p_edges = _neighbor_edges(par, p_leaves) + if p_edges == [] + v_to_ordered_p_edges[v] = ([], [], []) + continue + end # Get the reference ordering and target ordering of `v` v_inds = map(e -> Set(p_edge_to_ordered_inds[e]), p_edges) constraint_tree = _adjacency_tree(v, path, par, p_edge_to_ordered_inds) if v in leaves - ref_inds_ordering = _mps_partition_inds_order(tn, v_inds; alg=linear_ordering_alg) + # TODO: the "bottom_up" algorithm currently doesn't support `v_inds` + # being a vector of set of indices. + ref_inds_ordering = _mps_partition_inds_order(tn, v_inds; alg="top_down") inds_ordering = _constrained_minswap_inds_ordering( constraint_tree, ref_inds_ordering, tn ) @@ -83,7 +89,7 @@ function partitioned_contract( # Note: the contracted ordering in `ref_p_edges` is not changed, # since `ref_p_edges` focuses on minimizing the number of swaps # rather than making later contractions efficient. - sibling = setdiff(child_vertices(contractions, path[1]), [v])[1] + sibling = setdiff(child_vertices(contraction_tree, path[1]), [v])[1] if haskey(v_to_ordered_p_edges, sibling) contract_edges = v_to_ordered_p_edges[sibling][3] @assert(_is_neighbored_subset(p_edges, Set(contract_edges))) diff --git a/src/partitioned_contract/utils.jl b/src/partitioned_contract/utils.jl index 00b0afc9..f63c41e0 100644 --- a/src/partitioned_contract/utils.jl +++ b/src/partitioned_contract/utils.jl @@ -4,11 +4,11 @@ function _is_neighbored_subset(v::Vector, set::Set) end @assert issubset(set, Set(v)) i_begin = 1 - while !(i_begin in set) + while !(v[i_begin] in set) i_begin += 1 end i_end = length(v) - while !(i_end in set) + while !(v[i_end] in set) i_end -= 1 end return Set(v[i_begin:i_end]) == set From 1889d992f004cce5c5247310f90c33499cde2325 Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Tue, 4 Jul 2023 23:56:44 -0500 Subject: [PATCH 23/35] 3d cube running working, performance better than the prev implementation --- examples/approx_contract/3d_cube.jl | 123 +++++++++++++++ examples/approx_contract/bench.jl | 110 ++++++++++++++ examples/approx_contract/exact_contract.jl | 34 +++++ examples/approx_contract/randregulargraph.jl | 39 +++++ examples/approx_contract/run.jl | 143 ++++++++++++++++++ examples/approx_contract/sweep_contractor.jl | 17 --- src/approx_itensornetwork/density_matrix.jl | 3 + .../partitioned_contract.jl | 4 +- 8 files changed, 455 insertions(+), 18 deletions(-) create mode 100644 examples/approx_contract/3d_cube.jl create mode 100644 examples/approx_contract/bench.jl create mode 100644 examples/approx_contract/exact_contract.jl create mode 100644 examples/approx_contract/randregulargraph.jl create mode 100644 examples/approx_contract/run.jl diff --git a/examples/approx_contract/3d_cube.jl b/examples/approx_contract/3d_cube.jl new file mode 100644 index 00000000..6d5a8572 --- /dev/null +++ b/examples/approx_contract/3d_cube.jl @@ -0,0 +1,123 @@ +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 + +# 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, 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 diff --git a/examples/approx_contract/bench.jl b/examples/approx_contract/bench.jl new file mode 100644 index 00000000..a5487b20 --- /dev/null +++ b/examples/approx_contract/bench.jl @@ -0,0 +1,110 @@ +using ITensorNetworks + +function bench_lnZ( + tntree; + num_iter, + cutoff, + maxdim, + ansatz, + approx_itensornetwork_alg, + swap_size, + contraction_sequence_alg, + contraction_sequence_kwargs, + linear_ordering_alg, +) + reset_timer!(ITensors.timer) + function _run() + out, log_acc_norm = partitioned_contract( + tntree; + cutoff, + maxdim, + ansatz, + approx_itensornetwork_alg, + swap_size, + contraction_sequence_alg, + contraction_sequence_kwargs, + linear_ordering_alg, + ) + out = Vector{ITensor}(out) + log_acc_norm = log(norm(out)) + log_acc_norm + @info "out value is", out[1].tensor + @info "out norm is", log_acc_norm + return log_acc_norm + end + out_list = [] + for _ in 1:num_iter + push!(out_list, _run()) + end + show(ITensors.timer) + # after warmup, start to benchmark + reset_timer!(ITensors.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(ITensors.timer) +end + +function bench_magnetization( + tntree_pair; + num_iter, + cutoff, + maxdim, + ansatz, + approx_itensornetwork_alg, + swap_size, + contraction_sequence_alg, + contraction_sequence_kwargs, + linear_ordering_alg, + warmup, +) + reset_timer!(ITensors.timer) + tntree1, tntree2 = tntree_pair + function _run() + out, log_acc_norm = approximate_contract( + tntree1; + cutoff, + maxdim, + ansatz, + approx_itensornetwork_alg, + swap_size, + contraction_sequence_alg, + contraction_sequence_kwargs, + linear_ordering_alg, + ) + out = Vector{ITensor}(out) + lognorm1 = log(norm(out)) + log_acc_norm + out, log_acc_norm = approximate_contract( + tntree2; + cutoff, + maxdim, + ansatz, + approx_itensornetwork_alg, + swap_size, + contraction_sequence_alg, + contraction_sequence_kwargs, + linear_ordering_alg, + ) + out = Vector{ITensor}(out) + lognorm2 = log(norm(out)) + log_acc_norm + @info "magnetization is", exp(lognorm1 - lognorm2) + return exp(lognorm1 - lognorm2) + end + out_list = [] + total_iter = num_iter + if warmup + for _ in 1:num_iter + push!(out_list, _run()) + end + show(ITensors.timer) + total_iter += num_iter + end + @info "warmup finished" + # after warmup, start to benchmark + reset_timer!(ITensors.timer) + for _ in 1:num_iter + push!(out_list, _run()) + end + @info "magnetization results are", out_list, "mean is", sum(out_list) / total_iter + return show(ITensors.timer) +end diff --git a/examples/approx_contract/exact_contract.jl b/examples/approx_contract/exact_contract.jl new file mode 100644 index 00000000..0c4e31d5 --- /dev/null +++ b/examples/approx_contract/exact_contract.jl @@ -0,0 +1,34 @@ +using ITensorNetworks: contraction_sequence + +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(network::ITensorNetwork; sc_target=30) + ITensors.set_warn_order(1000) + reset_timer!(ITensors.timer) + tn = Vector{ITensor}(network) + seq = contraction_sequence(tn; alg="tree_sa")#alg="kahypar_bipartite", sc_target=sc_target) + @info seq + tn = [(i, 0.0) for i in tn] + return contract_log_norm(tn, seq) +end diff --git a/examples/approx_contract/randregulargraph.jl b/examples/approx_contract/randregulargraph.jl new file mode 100644 index 00000000..9f868b5a --- /dev/null +++ b/examples/approx_contract/randregulargraph.jl @@ -0,0 +1,39 @@ +function build_tntree_unbalanced( + tn::ITensorNetwork; nvertices_per_partition=2, backend="KaHyPar" +) + @assert is_connected(tn) + g_parts = partition(tn; nvertices_per_partition=nvertices_per_partition, backend=backend) + @assert is_connected(g_parts) + root = 1 + tree = bfs_tree(g_parts, root) + tntree = nothing + queue = [root] + while queue != [] + v = popfirst!(queue) + queue = vcat(queue, child_vertices(tree, v)) + if tntree == nothing + tntree = Vector{ITensor}(g_parts[v]) + else + tntree = [tntree, Vector{ITensor}(g_parts[v])] + end + end + return tntree +end + +function build_tntree_balanced( + tn::ITensorNetwork; nvertices_per_partition=2, backend="KaHyPar" +) + # @assert is_connected(tn) + g_parts = partition(tn; npartitions=2, backend=backend) + if nv(g_parts[1]) >= max(nvertices_per_partition, 2) + tntree_1 = build_tntree_balanced(g_parts[1]; nvertices_per_partition, backend) + else + tntree_1 = Vector{ITensor}(g_parts[1]) + end + if nv(g_parts[2]) >= max(nvertices_per_partition, 2) + tntree_2 = build_tntree_balanced(g_parts[2]; nvertices_per_partition, backend=backend) + else + tntree_2 = Vector{ITensor}(g_parts[2]) + end + return [tntree_1, tntree_2] +end diff --git a/examples/approx_contract/run.jl b/examples/approx_contract/run.jl new file mode 100644 index 00000000..156d43c6 --- /dev/null +++ b/examples/approx_contract/run.jl @@ -0,0 +1,143 @@ +using ITensors, Graphs, NamedGraphs, TimerOutputs +using Distributions, Random +using KaHyPar, Metis +using ITensorNetworks +using ITensorNetworks: ising_network +using OMEinsumContractionOrders + +include("exact_contract.jl") +include("3d_cube.jl") +include("sweep_contractor.jl") +include("randregulargraph.jl") +include("bench.jl") + +Random.seed!(1234) +TimerOutputs.enable_debug_timings(ITensorNetworks) +TimerOutputs.enable_debug_timings(@__MODULE__) +ITensors.set_warn_order(100) + +#= +Exact contraction of ising_network +=# +# N = (3, 3, 3) +# beta = 0.3 +# network = ising_network(named_grid(N), beta=beta) +# exact_contract(network; sc_target=28) + +#= +Exact contraction of randomITensorNetwork +=# +# N = (5, 5, 5) +# distribution = Uniform{Float64}(-0.4, 1.0) +# network = randomITensorNetwork(named_grid(N); link_space=2, distribution=distribution) +# exact_contract(network; sc_target=30) # 47.239753244708396 + +#= +Bugs +=# +# 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); +# block_size=(1, 1, 1), +# beta=0.3, +# h=0.0, +# num_iter=2, +# cutoff=1e-20, +# maxdim=128, +# ansatz="mps", +# algorithm="density_matrix", +# use_cache=true, +# ortho=false, +# env_size=(3, 1, 1), +# ) + +#= +bench_3d_cube_lnZ +=# +N = (6, 6, 6) +beta = 0.3 +network = ising_network(named_grid(N), beta; h=0.0, szverts=nothing) +tntree = build_tntree(N, network; block_size=(1, 1, 1), env_size=(6, 1, 1)) +@time bench_lnZ( + tntree; + num_iter=1, + cutoff=1e-12, + maxdim=128, + ansatz="mps", + approx_itensornetwork_alg="density_matrix", + swap_size=10000, + contraction_sequence_alg="sa_bipartite", + contraction_sequence_kwargs=(;), + linear_ordering_alg="bottom_up", +) + +#= +bench_3d_cube_magnetization +=# +# N = (6, 6, 6) +# beta = 0.22 +# h = 0.050 +# szverts = [(3, 3, 3)] +# network1 = ising_network(named_grid(N), beta; h=h, szverts=szverts) +# network2 = ising_network(named_grid(N), beta; h=h, szverts=nothing) +# # e1 = exact_contract(network1; sc_target=28) +# # e2 = exact_contract(network2; sc_target=28) +# # @info exp(e1[2] - e2[2]) +# tntree1 = build_tntree(N, network1; block_size=(1, 1, 1), env_size=(3, 1, 1)) +# tntree2 = build_tntree(N, network2; block_size=(1, 1, 1), env_size=(3, 1, 1)) +# @time bench_magnetization( +# tntree1 => tntree2; +# num_iter=1, +# cutoff=1e-12, +# maxdim=256, +# ansatz="mps", +# algorithm="density_matrix", +# use_cache=true, +# ortho=false, +# swap_size=10000, +# warmup=false, +# ) + +#= +SweepContractor +=# +# reset_timer!(ITensors.timer) +# # N = (5, 5, 5) +# # beta = 0.3 +# # network = ising_network(named_grid(N), beta; h=0.0, szverts=nothing) +# N = (3, 3, 3) +# distribution = Uniform{Float64}(-0.4, 1.0) +# network = randomITensorNetwork(named_grid(N); link_space=2, distribution=distribution) +# ltn = sweep_contractor_tensor_network( +# network, (i, j, k) -> (j + 0.01 * randn(), k + 0.01 * randn()) +# ) +# @time lnz = contract_w_sweep(ltn; rank=256) +# @info "lnZ of SweepContractor is", lnz +# show(ITensors.timer) + +#= +random regular graph +=# +# nvertices = 220 +# deg = 3 +# distribution = Uniform{Float64}(-0.2, 1.0) +# network = randomITensorNetwork( +# random_regular_graph(nvertices, deg); link_space=2, distribution=distribution +# ) +# # exact_contract(network; sc_target=30) +# # -26.228887728408008 (-1.0, 1.0) +# # 5.633462619348083 (-0.2, 1.0) +# # nvertices_per_partition=10 works 15/20 not work +# tntree = build_tntree_balanced(network; nvertices_per_partition=10) +# @time bench_lnZ( +# tntree; +# num_iter=2, +# cutoff=1e-12, +# maxdim=64, +# ansatz="mps", +# algorithm="density_matrix", +# use_cache=true, +# ortho=false, +# swap_size=4, +# ) diff --git a/examples/approx_contract/sweep_contractor.jl b/examples/approx_contract/sweep_contractor.jl index 6ee9f704..708ee823 100644 --- a/examples/approx_contract/sweep_contractor.jl +++ b/examples/approx_contract/sweep_contractor.jl @@ -33,20 +33,3 @@ function contract_w_sweep(ltn::SweepContractor.LabelledTensorNetwork; rank) return log(abs(sweep[1])) + sweep[2] * log(2) end end - -Random.seed!(1234) -TimerOutputs.enable_debug_timings(@__MODULE__) -reset_timer!(ITensors.timer) - -N = (3, 3, 3) -beta = 0.3 -network = ising_network(named_grid(N), beta; h=0.0, szverts=nothing) -# N = (5, 5, 5) -# distribution = Uniform{Float64}(-0.4, 1.0) -# network = randomITensorNetwork(named_grid(N); link_space=2, distribution=distribution) -ltn = sweep_contractor_tensor_network( - network, (i, j, k) -> (j + 0.01 * randn(), k + 0.01 * randn()) -) -@time lnz = contract_w_sweep(ltn; rank=256) -@info "lnZ of SweepContractor is", lnz -show(ITensors.timer) diff --git a/src/approx_itensornetwork/density_matrix.jl b/src/approx_itensornetwork/density_matrix.jl index e5009067..f33cfc48 100644 --- a/src/approx_itensornetwork/density_matrix.jl +++ b/src/approx_itensornetwork/density_matrix.jl @@ -144,9 +144,12 @@ end function _get_low_rank_projector(tensor, inds1, inds2; cutoff, maxdim) @assert length(inds(tensor)) <= 4 + # t00 = time() @timeit_debug ITensors.timer "[approx_binary_tree_itensornetwork]: eigen" begin F = eigen(tensor, inds1, inds2; cutoff=cutoff, maxdim=maxdim, ishermitian=true) end + # t11 = time() - t00 + # @info "size of U", size(F.V), "size of diag", size(F.D), "costs", t11 return F.Vt end diff --git a/src/partitioned_contract/partitioned_contract.jl b/src/partitioned_contract/partitioned_contract.jl index a425e5ec..847d0af0 100644 --- a/src/partitioned_contract/partitioned_contract.jl +++ b/src/partitioned_contract/partitioned_contract.jl @@ -129,7 +129,9 @@ function partitioned_contract( out /= out_nrm return ITensorNetwork(out), log_acc_norm + log(out_nrm) end - for edge_order in _intervals(ref_p_edges, p_edges; size=swap_size) + intervals = _intervals(ref_p_edges, p_edges; size=swap_size) + @info "number of approx_itensornetwork calls, $(length(intervals))" + for edge_order in intervals inds_ordering = map(e -> p_edge_to_ordered_inds[e], edge_order) # `ortho_center` denotes the number of edges arranged at # the left of the center vertex. From d60c21eef3fef03d81621ea170e91dfec4f58e6e Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Wed, 5 Jul 2023 12:14:31 -0500 Subject: [PATCH 24/35] Implement contract with partitioned_contract --- src/contract.jl | 31 ++++++++++++++++++++++++++++--- 1 file changed, 28 insertions(+), 3 deletions(-) diff --git a/src/contract.jl b/src/contract.jl index aff4a022..9c100a12 100644 --- a/src/contract.jl +++ b/src/contract.jl @@ -21,9 +21,34 @@ end function contract( alg::Algorithm"partitioned_contract", tn::AbstractITensorNetwork; - output_structure::Function=path_graph_structure, + nvertices_per_partition::Integer=2, + backend::String="KaHyPar", kwargs..., ) - # TODO - return approx_itensornetwork(alg, tn, output_structure; kwargs...) + sequence = _partitioned_contraction_sequence( + tn; nvertices_per_partition=nvertices_per_partition, backend=backend + ) + return partitioned_contract(sequence; kwargs...) +end + +function _partitioned_contraction_sequence( + tn::ITensorNetwork; nvertices_per_partition=2, backend="KaHyPar" +) + # @assert is_connected(tn) + g_parts = partition(tn; npartitions=2, backend=backend) + if nv(g_parts[1]) >= max(nvertices_per_partition, 2) + tntree_1 = _partitioned_contraction_sequence( + g_parts[1]; nvertices_per_partition, backend + ) + else + tntree_1 = Vector{ITensor}(g_parts[1]) + end + if nv(g_parts[2]) >= max(nvertices_per_partition, 2) + tntree_2 = _partitioned_contraction_sequence( + g_parts[2]; nvertices_per_partition, backend=backend + ) + else + tntree_2 = Vector{ITensor}(g_parts[2]) + end + return [tntree_1, tntree_2] end From 53348f3ca286cb01aeeca19b6927ead06ddef854 Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Wed, 5 Jul 2023 12:16:24 -0500 Subject: [PATCH 25/35] bug fixes & logging & renaming for partitioned_contract --- src/approx_itensornetwork/utils.jl | 1 + src/mincut.jl | 2 +- .../partitioned_contract.jl | 24 ++++++++++++++----- 3 files changed, 20 insertions(+), 7 deletions(-) diff --git a/src/approx_itensornetwork/utils.jl b/src/approx_itensornetwork/utils.jl index e650ac49..35c653ae 100644 --- a/src/approx_itensornetwork/utils.jl +++ b/src/approx_itensornetwork/utils.jl @@ -23,6 +23,7 @@ end """ Contract of a vector of tensors, `network`, with a contraction sequence generated via sa_bipartite +# TODO: rewrite using `contract` """ function _optcontract( network::Vector; contraction_sequence_alg="optimal", contraction_sequence_kwargs=(;) diff --git a/src/mincut.jl b/src/mincut.jl index a8d3152e..b5a1288e 100644 --- a/src/mincut.jl +++ b/src/mincut.jl @@ -390,7 +390,7 @@ end function _mps_mincut_partition_cost(tn::ITensorNetwork, inds_vector::Vector{<:Set}) @timeit_debug ITensors.timer "_mps_mincut_partition_cost" begin - inds_vector = map(inds -> collect(inds), inds_vector) + inds_vector = map(inds -> Vector{Index}(collect(inds)), inds_vector) outinds = vcat(inds_vector...) maxweight_tn, out_to_maxweight_ind = _maxweightoutinds_tn(tn, outinds) sourceinds_list = [vcat(inds_vector[1:i]...) for i in 1:(length(inds_vector) - 1)] diff --git a/src/partitioned_contract/partitioned_contract.jl b/src/partitioned_contract/partitioned_contract.jl index 847d0af0..3dc20362 100644 --- a/src/partitioned_contract/partitioned_contract.jl +++ b/src/partitioned_contract/partitioned_contract.jl @@ -1,6 +1,6 @@ -function partitioned_contract(nested_vector::Vector; kwargs...) +function partitioned_contract(contraction_sequence::Vector; kwargs...) leaves = filter( - v -> v isa Vector && v[1] isa ITensor, collect(PreOrderDFS(nested_vector)) + v -> v isa Vector && v[1] isa ITensor, collect(PreOrderDFS(contraction_sequence)) ) tn = ITensorNetwork() subgraph_vs = [] @@ -16,10 +16,10 @@ function partitioned_contract(nested_vector::Vector; kwargs...) push!(subgraph_vs, vs) end par = partition(tn, subgraph_vs) - vs_nested_vector = _map_nested_vector( - Dict(leaves .=> collect(1:length(leaves))), nested_vector + vs_contraction_sequence = _map_nested_vector( + Dict(leaves .=> collect(1:length(leaves))), contraction_sequence ) - contraction_tree = contraction_sequence_to_digraph(vs_nested_vector) + contraction_tree = contraction_sequence_to_digraph(vs_contraction_sequence) return partitioned_contract(par, contraction_tree; kwargs...) end @@ -35,9 +35,18 @@ function partitioned_contract( contraction_sequence_kwargs, linear_ordering_alg, ) + @info "ansatz: $(ansatz)" + @info "approx_itensornetwork_alg: $(approx_itensornetwork_alg)" + @info "cutoff: $(cutoff)" + @info "maxdim: $(maxdim)" + @info "swap_size: $(swap_size)" + @info "contraction_sequence_alg: $(contraction_sequence_alg)" + @info "contraction_sequence_kwargs: $(contraction_sequence_kwargs)" + @info "linear_ordering_alg: $(linear_ordering_alg)" input_tn = ITensorNetwork(mapreduce(l -> Vector{ITensor}(par[l]), vcat, vertices(par))) # TODO: currently we assume that the tn represented by 'par' is closed. @assert noncommoninds(Vector{ITensor}(input_tn)...) == [] + @timeit_debug ITensors.timer "partitioned_contract" begin leaves = leaf_vertices(contraction_tree) traversal = post_order_dfs_vertices(contraction_tree, _root(contraction_tree)) @@ -130,7 +139,7 @@ function partitioned_contract( return ITensorNetwork(out), log_acc_norm + log(out_nrm) end intervals = _intervals(ref_p_edges, p_edges; size=swap_size) - @info "number of approx_itensornetwork calls, $(length(intervals))" + @info "number of approx_itensornetwork calls: $(length(intervals))" for edge_order in intervals inds_ordering = map(e -> p_edge_to_ordered_inds[e], edge_order) # `ortho_center` denotes the number of edges arranged at @@ -200,6 +209,9 @@ function _ansatz_tree(inds_orderings::Vector, ansatz::String, ortho_center::Inte if nested_vec_left == [] || nested_vec_right == [] nested_vec = nested_vec_left == [] ? nested_vec_right : nested_vec_left else + nested_vec_left = length(nested_vec_left) == 1 ? nested_vec_left[1] : nested_vec_left + nested_vec_right = + length(nested_vec_right) == 1 ? nested_vec_right[1] : nested_vec_right nested_vec = [nested_vec_left, nested_vec_right] end return _nested_vector_to_digraph(nested_vec) From 2711bbdc035a07e4fa37935d60b73136023629fe Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Wed, 5 Jul 2023 12:37:46 -0500 Subject: [PATCH 26/35] randregulargraph works --- examples/approx_contract/3d_cube.jl | 21 --------- examples/approx_contract/bench.jl | 49 +++++++++++--------- examples/approx_contract/randregulargraph.jl | 39 ---------------- examples/approx_contract/run.jl | 20 ++++---- 4 files changed, 37 insertions(+), 92 deletions(-) delete mode 100644 examples/approx_contract/randregulargraph.jl diff --git a/examples/approx_contract/3d_cube.jl b/examples/approx_contract/3d_cube.jl index 6d5a8572..bf5140f1 100644 --- a/examples/approx_contract/3d_cube.jl +++ b/examples/approx_contract/3d_cube.jl @@ -61,27 +61,6 @@ function build_recursive_tntree(tn, N; env_size) ] 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, network::ITensorNetwork; block_size::Tuple, env_size::Tuple) @assert length(block_size) == length(env_size) order = length(block_size) diff --git a/examples/approx_contract/bench.jl b/examples/approx_contract/bench.jl index a5487b20..703f5851 100644 --- a/examples/approx_contract/bench.jl +++ b/examples/approx_contract/bench.jl @@ -1,30 +1,33 @@ using ITensorNetworks -function bench_lnZ( - tntree; - num_iter, - cutoff, - maxdim, - ansatz, - approx_itensornetwork_alg, - swap_size, - contraction_sequence_alg, - contraction_sequence_kwargs, - linear_ordering_alg, -) +function bench_lnZ(tntree::Vector; num_iter, kwargs...) reset_timer!(ITensors.timer) function _run() - out, log_acc_norm = partitioned_contract( - tntree; - cutoff, - maxdim, - ansatz, - approx_itensornetwork_alg, - swap_size, - contraction_sequence_alg, - contraction_sequence_kwargs, - linear_ordering_alg, - ) + out, log_acc_norm = partitioned_contract(tntree; kwargs...) + out = Vector{ITensor}(out) + log_acc_norm = log(norm(out)) + log_acc_norm + @info "out value is", out[1].tensor + @info "out norm is", log_acc_norm + return log_acc_norm + end + out_list = [] + for _ in 1:num_iter + push!(out_list, _run()) + end + show(ITensors.timer) + # after warmup, start to benchmark + reset_timer!(ITensors.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(ITensors.timer) +end + +function bench_lnZ(tn::ITensorNetwork; num_iter, kwargs...) + reset_timer!(ITensors.timer) + function _run() + out, log_acc_norm = contract(tn; alg="partitioned_contract", kwargs...) out = Vector{ITensor}(out) log_acc_norm = log(norm(out)) + log_acc_norm @info "out value is", out[1].tensor diff --git a/examples/approx_contract/randregulargraph.jl b/examples/approx_contract/randregulargraph.jl deleted file mode 100644 index 9f868b5a..00000000 --- a/examples/approx_contract/randregulargraph.jl +++ /dev/null @@ -1,39 +0,0 @@ -function build_tntree_unbalanced( - tn::ITensorNetwork; nvertices_per_partition=2, backend="KaHyPar" -) - @assert is_connected(tn) - g_parts = partition(tn; nvertices_per_partition=nvertices_per_partition, backend=backend) - @assert is_connected(g_parts) - root = 1 - tree = bfs_tree(g_parts, root) - tntree = nothing - queue = [root] - while queue != [] - v = popfirst!(queue) - queue = vcat(queue, child_vertices(tree, v)) - if tntree == nothing - tntree = Vector{ITensor}(g_parts[v]) - else - tntree = [tntree, Vector{ITensor}(g_parts[v])] - end - end - return tntree -end - -function build_tntree_balanced( - tn::ITensorNetwork; nvertices_per_partition=2, backend="KaHyPar" -) - # @assert is_connected(tn) - g_parts = partition(tn; npartitions=2, backend=backend) - if nv(g_parts[1]) >= max(nvertices_per_partition, 2) - tntree_1 = build_tntree_balanced(g_parts[1]; nvertices_per_partition, backend) - else - tntree_1 = Vector{ITensor}(g_parts[1]) - end - if nv(g_parts[2]) >= max(nvertices_per_partition, 2) - tntree_2 = build_tntree_balanced(g_parts[2]; nvertices_per_partition, backend=backend) - else - tntree_2 = Vector{ITensor}(g_parts[2]) - end - return [tntree_1, tntree_2] -end diff --git a/examples/approx_contract/run.jl b/examples/approx_contract/run.jl index 156d43c6..399e4290 100644 --- a/examples/approx_contract/run.jl +++ b/examples/approx_contract/run.jl @@ -2,13 +2,12 @@ using ITensors, Graphs, NamedGraphs, TimerOutputs using Distributions, Random using KaHyPar, Metis using ITensorNetworks -using ITensorNetworks: ising_network +using ITensorNetworks: ising_network, contract using OMEinsumContractionOrders include("exact_contract.jl") include("3d_cube.jl") include("sweep_contractor.jl") -include("randregulargraph.jl") include("bench.jl") Random.seed!(1234) @@ -123,21 +122,24 @@ random regular graph # deg = 3 # distribution = Uniform{Float64}(-0.2, 1.0) # network = randomITensorNetwork( -# random_regular_graph(nvertices, deg); link_space=2, distribution=distribution +# distribution, random_regular_graph(nvertices, deg); link_space=2 # ) # # exact_contract(network; sc_target=30) # # -26.228887728408008 (-1.0, 1.0) # # 5.633462619348083 (-0.2, 1.0) # # nvertices_per_partition=10 works 15/20 not work -# tntree = build_tntree_balanced(network; nvertices_per_partition=10) +# # tntree = partitioned_contraction_sequence(network; nvertices_per_partition=10) # @time bench_lnZ( -# tntree; +# network; # num_iter=2, +# nvertices_per_partition=10, +# backend="KaHyPar", # cutoff=1e-12, -# maxdim=64, +# maxdim=512, # ansatz="mps", -# algorithm="density_matrix", -# use_cache=true, -# ortho=false, +# approx_itensornetwork_alg="density_matrix", # swap_size=4, +# contraction_sequence_alg="sa_bipartite", +# contraction_sequence_kwargs=(;), +# linear_ordering_alg="bottom_up", # ) From 0567f1bbc407e8972b2c4ea00e6230e0490f6d05 Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Wed, 5 Jul 2023 12:42:45 -0500 Subject: [PATCH 27/35] rename approx_contract in examples to partitioned_contract --- examples/{approx_contract => partitioned_contract}/3d_cube.jl | 0 examples/{approx_contract => partitioned_contract}/bench.jl | 0 .../{approx_contract => partitioned_contract}/exact_contract.jl | 0 examples/{approx_contract => partitioned_contract}/run.jl | 0 .../{approx_contract => partitioned_contract}/sweep_contractor.jl | 0 5 files changed, 0 insertions(+), 0 deletions(-) rename examples/{approx_contract => partitioned_contract}/3d_cube.jl (100%) rename examples/{approx_contract => partitioned_contract}/bench.jl (100%) rename examples/{approx_contract => partitioned_contract}/exact_contract.jl (100%) rename examples/{approx_contract => partitioned_contract}/run.jl (100%) rename examples/{approx_contract => partitioned_contract}/sweep_contractor.jl (100%) diff --git a/examples/approx_contract/3d_cube.jl b/examples/partitioned_contract/3d_cube.jl similarity index 100% rename from examples/approx_contract/3d_cube.jl rename to examples/partitioned_contract/3d_cube.jl diff --git a/examples/approx_contract/bench.jl b/examples/partitioned_contract/bench.jl similarity index 100% rename from examples/approx_contract/bench.jl rename to examples/partitioned_contract/bench.jl diff --git a/examples/approx_contract/exact_contract.jl b/examples/partitioned_contract/exact_contract.jl similarity index 100% rename from examples/approx_contract/exact_contract.jl rename to examples/partitioned_contract/exact_contract.jl diff --git a/examples/approx_contract/run.jl b/examples/partitioned_contract/run.jl similarity index 100% rename from examples/approx_contract/run.jl rename to examples/partitioned_contract/run.jl diff --git a/examples/approx_contract/sweep_contractor.jl b/examples/partitioned_contract/sweep_contractor.jl similarity index 100% rename from examples/approx_contract/sweep_contractor.jl rename to examples/partitioned_contract/sweep_contractor.jl From 00bc6cd06ef6b58a63fee96286b189d0acb24e51 Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Fri, 7 Jul 2023 10:48:46 -0500 Subject: [PATCH 28/35] Add quantum circuit simulation. Performance can be improved --- .../partitioned_contract/random_circuit.jl | 105 ++++++++++++++++++ examples/partitioned_contract/run.jl | 52 ++++++--- .../constrained_ordering.jl | 4 + .../partitioned_contract.jl | 7 ++ 4 files changed, 152 insertions(+), 16 deletions(-) create mode 100644 examples/partitioned_contract/random_circuit.jl diff --git a/examples/partitioned_contract/random_circuit.jl b/examples/partitioned_contract/random_circuit.jl new file mode 100644 index 00000000..38bbbead --- /dev/null +++ b/examples/partitioned_contract/random_circuit.jl @@ -0,0 +1,105 @@ +using ITensors, Graphs, NamedGraphs, TimerOutputs +using Random +using PastaQ +using KaHyPar +using ITensorNetworks + +# Return all the tensors represented by `gates`, and the `hilbert` +# after applying these gates. The `hilbert` is a vector if indices. +function gate_tensors(hilbert::Vector{<:Index}, gates; applydag=false) + hilbert = copy(hilbert) + Us = ITensor[] + for g in gates + if applydag == true + t = dag(gate(hilbert, g)) + inds1 = [hilbert[n] for n in g[2]] + t = swapinds(t, inds1, [i' for i in inds1]) + else + t = gate(hilbert, g) + end + if length(g[2]) == 2 + ind = hilbert[g[2][1]] + L, R = factorize(t, ind, ind'; cutoff=1e-16) + push!(Us, L, R) + else + push!(Us, t) + end + for n in g[2] + hilbert[n] = hilbert[n]' + end + end + # @info "gates", gates + # @info "hilbert", hilbert + return Us, hilbert +end + +function line_partition_gates(N, gates; order=1) + @assert order in [1, 2, 3, 0] + nrow, ncol = N + if order in [0, 1] + # each partition contains a column + partition_qubits = [[(c - 1) * nrow + r for r in 1:nrow] for c in 1:ncol] + elseif order == 2 + # Add columns [1], [2, 3], [4, 5], ... + partition_qubits = [collect(1:nrow)] + push!( + partition_qubits, + [ + [(c1 - 1) * nrow + r for r in 1:nrow for c1 in [c, c + 1]] for c in 2:2:(ncol - 1) + ]..., + ) + if iseven(ncol) + push!(partition_qubits, [(ncol - 1) * nrow + r for r in 1:nrow]) + end + else + # Add columns [1, 2], [3, 4], [5, 6], ... + partition_qubits = [ + [(c1 - 1) * nrow + r for r in 1:nrow for c1 in [c, c + 1]] for c in 1:2:ncol + ] + if isodd(ncol) + push!(partition_qubits, [(ncol - 1) * nrow + r for r in 1:nrow]) + end + end + @info partition_qubits + partition_gates = [filter(g -> all(q -> q in p, g[2]), gates) for p in partition_qubits] + return partition_gates +end + +function line_network(network::Vector) + tntree = network[1] + for i in 2:length(network) + if network[i] isa Vector + tntree = [tntree, network[i]] + else + tntree = [tntree, [network[i]]] + end + end + return tntree +end + +function random_circuit_line_partition_sequence( + N, depth; twoqubitgates="CX", onequbitgates="Ry" +) + @assert N isa Number || length(N) <= 2 + # Each layer contains multiple two-qubit gates, followed by one-qubit + # gates, each applying on one qubit. + layers = randomcircuit( + N; depth=depth, twoqubitgates=twoqubitgates, onequbitgates=onequbitgates + ) + partition_gates = [] + for (i, l) in enumerate(layers) + push!(partition_gates, line_partition_gates(N, l; order=i % 4)...) + end + hilbert = qubits(prod(N)) + tensor_partitions = [productstate(hilbert)[:]] + for gates in partition_gates + tensors, hilbert = gate_tensors(hilbert, gates) + push!(tensor_partitions, tensors) + end + for gates in reverse(partition_gates) + tensors, hilbert = gate_tensors(hilbert, reverse(gates); applydag=true) + push!(tensor_partitions, tensors) + end + push!(tensor_partitions, productstate(hilbert)[:]) + return line_network(tensor_partitions) +end diff --git a/examples/partitioned_contract/run.jl b/examples/partitioned_contract/run.jl index 399e4290..a88d8dee 100644 --- a/examples/partitioned_contract/run.jl +++ b/examples/partitioned_contract/run.jl @@ -8,6 +8,7 @@ using OMEinsumContractionOrders include("exact_contract.jl") include("3d_cube.jl") include("sweep_contractor.jl") +include("random_circuit.jl") include("bench.jl") Random.seed!(1234) @@ -54,22 +55,22 @@ Bugs #= bench_3d_cube_lnZ =# -N = (6, 6, 6) -beta = 0.3 -network = ising_network(named_grid(N), beta; h=0.0, szverts=nothing) -tntree = build_tntree(N, network; block_size=(1, 1, 1), env_size=(6, 1, 1)) -@time bench_lnZ( - tntree; - num_iter=1, - cutoff=1e-12, - maxdim=128, - ansatz="mps", - approx_itensornetwork_alg="density_matrix", - swap_size=10000, - contraction_sequence_alg="sa_bipartite", - contraction_sequence_kwargs=(;), - linear_ordering_alg="bottom_up", -) +# N = (6, 6, 6) +# beta = 0.3 +# network = ising_network(named_grid(N), beta; h=0.0, szverts=nothing) +# tntree = build_tntree(N, network; block_size=(1, 1, 1), env_size=(6, 1, 1)) +# @time bench_lnZ( +# tntree; +# num_iter=1, +# cutoff=1e-12, +# maxdim=128, +# ansatz="mps", +# approx_itensornetwork_alg="density_matrix", +# swap_size=10000, +# contraction_sequence_alg="sa_bipartite", +# contraction_sequence_kwargs=(;), +# linear_ordering_alg="bottom_up", +# ) #= bench_3d_cube_magnetization @@ -143,3 +144,22 @@ random regular graph # contraction_sequence_kwargs=(;), # linear_ordering_alg="bottom_up", # ) + +#= +Simulation of random quantum circuit +=# +N = (6, 6) +depth = 6 +sequence = random_circuit_line_partition_sequence(N, depth) +@time bench_lnZ( + sequence; + num_iter=1, + cutoff=1e-12, + maxdim=256, + ansatz="mps", + approx_itensornetwork_alg="density_matrix", + swap_size=8, + contraction_sequence_alg="sa_bipartite", + contraction_sequence_kwargs=(;), + linear_ordering_alg="bottom_up", +) diff --git a/src/partitioned_contract/constrained_ordering.jl b/src/partitioned_contract/constrained_ordering.jl index acb3a263..332e7846 100644 --- a/src/partitioned_contract/constrained_ordering.jl +++ b/src/partitioned_contract/constrained_ordering.jl @@ -56,6 +56,10 @@ function _constrained_minswap_inds_ordering( else inputs = _low_swap_merge(left_1, right_1, left_2, right_2) end + # #### + # mincuts = map(o -> _mps_mincut_partition_cost(tn, o), inputs) + # inputs = [inputs[i] for i in 1:length(inputs) if mincuts[i] == mincuts[argmin(mincuts)]] + # #### constraint_tree_copies = [copy(constraint_tree) for _ in 1:length(inputs)] outputs = [] nswaps_list = [] diff --git a/src/partitioned_contract/partitioned_contract.jl b/src/partitioned_contract/partitioned_contract.jl index 3dc20362..41e86dc7 100644 --- a/src/partitioned_contract/partitioned_contract.jl +++ b/src/partitioned_contract/partitioned_contract.jl @@ -52,6 +52,10 @@ function partitioned_contract( traversal = post_order_dfs_vertices(contraction_tree, _root(contraction_tree)) contractions = setdiff(traversal, leaves) p_edge_to_ordered_inds = _ind_orderings(par, contractions; linear_ordering_alg) + for (p_edge, ordered_inds) in p_edge_to_ordered_inds + @info "edge", p_edge + @info "ordered_inds", ordered_inds + end # Build the orderings used for the ansatz tree. # For each tuple in `v_to_ordered_p_edges`, the first item is the # reference ordering of uncontracted edges for the contraction `v`, @@ -109,6 +113,8 @@ function partitioned_contract( contract_edges = intersect(p_edges, p_edges_2) end v_to_ordered_p_edges[v] = (ref_p_edges, p_edges, contract_edges) + @info "ref_p_edges is", ref_p_edges + @info "p_edges is", p_edges end # start approximate contraction v_to_tn = Dict{Tuple,ITensorNetwork}() @@ -169,6 +175,7 @@ end # Note: currently this function assumes that the input tn represented by 'par' is closed # TODO: test needed: function _ind_orderings(par::DataGraph, contractions::Vector; linear_ordering_alg) + @info "number of contractions", length(contractions) p_edge_to_ordered_inds = Dict{Any,Vector}() for v in contractions contract_edges = intersect(_neighbor_edges(par, v[1]), _neighbor_edges(par, v[2])) From a14267e7f48497f675f979bb7c44e2c63d01ea00 Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Fri, 7 Jul 2023 19:25:24 -0500 Subject: [PATCH 29/35] Add timer to _mps_order --- src/mincut.jl | 32 +++++++++++++++++--------------- 1 file changed, 17 insertions(+), 15 deletions(-) diff --git a/src/mincut.jl b/src/mincut.jl index b5a1288e..6d02f00c 100644 --- a/src/mincut.jl +++ b/src/mincut.jl @@ -270,21 +270,23 @@ function _mps_partition_inds_order( alg="top_down", backend="Metis", ) - @assert alg in ["top_down", "bottom_up"] - if outinds == nothing - outinds = noncommoninds(Vector{ITensor}(tn)...) - end - if length(outinds) == 1 - return outinds - end - if alg == "bottom_up" - tn2, out_to_maxweight_ind = _maxweightoutinds_tn(tn, outinds) - return _binary_tree_partition_inds_order_maximally_unbalanced( - tn => tn2, out_to_maxweight_ind - ) - else - nested_vector = _recursive_bisection(tn, outinds; backend=backend) - return filter(v -> v in outinds, collect(PreOrderDFS(nested_vector))) + @timeit_debug ITensors.timer "_mps_partition_inds_order" begin + @assert alg in ["top_down", "bottom_up"] + if outinds == nothing + outinds = noncommoninds(Vector{ITensor}(tn)...) + end + if length(outinds) == 1 + return outinds + end + if alg == "bottom_up" + tn2, out_to_maxweight_ind = _maxweightoutinds_tn(tn, outinds) + return _binary_tree_partition_inds_order_maximally_unbalanced( + tn => tn2, out_to_maxweight_ind + ) + else + nested_vector = _recursive_bisection(tn, outinds; backend=backend) + return filter(v -> v in outinds, collect(PreOrderDFS(nested_vector))) + end end end From 2da8714dae8f708fc1137cf89dd14f1458ff01ed Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Fri, 7 Jul 2023 22:57:12 -0500 Subject: [PATCH 30/35] Make _ind_orderings more efficient, and add c_in_leaves to _constrained_minswap_inds_ordering --- .../constrained_ordering.jl | 34 ++++++++++---- .../partitioned_contract.jl | 46 +++++++++++++++---- 2 files changed, 61 insertions(+), 19 deletions(-) diff --git a/src/partitioned_contract/constrained_ordering.jl b/src/partitioned_contract/constrained_ordering.jl index 332e7846..849385d2 100644 --- a/src/partitioned_contract/constrained_ordering.jl +++ b/src/partitioned_contract/constrained_ordering.jl @@ -39,22 +39,38 @@ function _constrained_minswap_inds_ordering( input_order_1::Vector, input_order_2::Vector, tn::ITensorNetwork, + c1_in_leaves::Bool, + c2_in_leaves::Bool, ) inter_igs = intersect(input_order_1, input_order_2) left_1, right_1 = _split_array(input_order_1, inter_igs) left_2, right_2 = _split_array(input_order_2, inter_igs) # @info "lengths of the input partitions", # sort([length(left_1), length(right_2), length(left_2), length(right_2)]) - num_swaps_1 = min(length(left_1), length(left_2)) + min(length(right_1), length(right_2)) - num_swaps_2 = min(length(left_1), length(right_2)) + min(length(right_1), length(left_2)) - if num_swaps_1 == num_swaps_2 - inputs_1 = _low_swap_merge(left_1, right_1, left_2, right_2) - inputs_2 = _low_swap_merge(left_1, right_1, reverse(right_2), reverse(left_2)) - inputs = [inputs_1..., inputs_2...] - elseif num_swaps_1 > num_swaps_2 - inputs = _low_swap_merge(left_1, right_1, reverse(right_2), reverse(left_2)) + + # TODO: this conditions seems not that useful, except the quantum + # circuit simulation experiment. We should consider removing this + # once confirming that this is only useful for special cases. + if c1_in_leaves && !c2_in_leaves + inputs = collect(permutations([left_1..., right_1...])) + inputs = [[left_2..., i..., right_2...] for i in inputs] + elseif !c1_in_leaves && c2_in_leaves + inputs = collect(permutations([left_2..., right_2...])) + inputs = [[left_1..., i..., right_1...] for i in inputs] else - inputs = _low_swap_merge(left_1, right_1, left_2, right_2) + num_swaps_1 = + min(length(left_1), length(left_2)) + min(length(right_1), length(right_2)) + num_swaps_2 = + min(length(left_1), length(right_2)) + min(length(right_1), length(left_2)) + if num_swaps_1 == num_swaps_2 + inputs_1 = _low_swap_merge(left_1, right_1, left_2, right_2) + inputs_2 = _low_swap_merge(left_1, right_1, reverse(right_2), reverse(left_2)) + inputs = [inputs_1..., inputs_2...] + elseif num_swaps_1 > num_swaps_2 + inputs = _low_swap_merge(left_1, right_1, reverse(right_2), reverse(left_2)) + else + inputs = _low_swap_merge(left_1, right_1, left_2, right_2) + end end # #### # mincuts = map(o -> _mps_mincut_partition_cost(tn, o), inputs) diff --git a/src/partitioned_contract/partitioned_contract.jl b/src/partitioned_contract/partitioned_contract.jl index 41e86dc7..d8ce500e 100644 --- a/src/partitioned_contract/partitioned_contract.jl +++ b/src/partitioned_contract/partitioned_contract.jl @@ -51,7 +51,7 @@ function partitioned_contract( leaves = leaf_vertices(contraction_tree) traversal = post_order_dfs_vertices(contraction_tree, _root(contraction_tree)) contractions = setdiff(traversal, leaves) - p_edge_to_ordered_inds = _ind_orderings(par, contractions; linear_ordering_alg) + p_edge_to_ordered_inds = _ind_orderings(par, contraction_tree; linear_ordering_alg) for (p_edge, ordered_inds) in p_edge_to_ordered_inds @info "edge", p_edge @info "ordered_inds", ordered_inds @@ -91,8 +91,14 @@ function partitioned_contract( c2_inds_ordering = map( e -> Set(p_edge_to_ordered_inds[e]), v_to_ordered_p_edges[c2][2] ) + # TODO: consider removing `c1 in leaves` and `c2 in leaves` below. ref_inds_ordering, inds_ordering = _constrained_minswap_inds_ordering( - constraint_tree, c1_inds_ordering, c2_inds_ordering, tn + constraint_tree, + c1_inds_ordering, + c2_inds_ordering, + tn, + c1 in leaves, + c2 in leaves, ) end ref_p_edges = p_edges[_findperm(v_inds, ref_inds_ordering)] @@ -174,19 +180,39 @@ end # Note: currently this function assumes that the input tn represented by 'par' is closed # TODO: test needed: -function _ind_orderings(par::DataGraph, contractions::Vector; linear_ordering_alg) +function _ind_orderings(par::DataGraph, contraction_tree::NamedDiGraph; linear_ordering_alg) + leaves = leaf_vertices(contraction_tree) + traversal = post_order_dfs_vertices(contraction_tree, _root(contraction_tree)) + contractions = setdiff(traversal, leaves) @info "number of contractions", length(contractions) p_edge_to_ordered_inds = Dict{Any,Vector}() + # mapping each vector of par vertices to a given tensor network. + # This dict is used so that the construnctor `ITensorNetwork` + # does not need to be called too many times. + p_vs_to_tn = Dict{Any,ITensorNetwork}() + # This dict is used so that `noncommoninds` does not need to be + # called too many times. + p_vs_to_noncommon_inds = Dict{Any,Vector}() + for v in leaves + @assert length(v[1]) == 1 + p_vs_to_tn[v[1]] = par[v[1][1]] + p_vs_to_noncommon_inds[v[1]] = noncommoninds(Vector{ITensor}(par[v[1][1]])...) + end for v in contractions contract_edges = intersect(_neighbor_edges(par, v[1]), _neighbor_edges(par, v[2])) - tn1 = ITensorNetwork(mapreduce(l -> Vector{ITensor}(par[l]), vcat, v[1])) - tn2 = ITensorNetwork(mapreduce(l -> Vector{ITensor}(par[l]), vcat, v[2])) - tn = length(vertices(tn1)) >= length(vertices(tn2)) ? tn1 : tn2 - tn_inds = noncommoninds(Vector{ITensor}(tn)...) + tn1, tn2 = p_vs_to_tn[v[1]], p_vs_to_tn[v[2]] + inds1, inds2 = p_vs_to_noncommon_inds[v[1]], p_vs_to_noncommon_inds[v[2]] + p_vs_to_tn[union(v[1], v[2])] = union(tn1, tn2) + p_vs_to_noncommon_inds[union(v[1], v[2])] = union( + setdiff(inds1, inds2), setdiff(inds2, inds1) + ) + # The inds ordering for each edge is selected based on a + # larger partition of the tn. + tn, tn_inds = length(vertices(tn1)) >= length(vertices(tn2)) ? (tn1, inds1) : (tn2, inds2) for e in contract_edges - inds1 = noncommoninds(Vector{ITensor}(par[e.src])...) - inds2 = noncommoninds(Vector{ITensor}(par[e.dst])...) - source_inds = intersect(inds1, inds2) + source_inds = intersect( + p_vs_to_noncommon_inds[[e.src]], p_vs_to_noncommon_inds[[e.dst]] + ) @assert issubset(source_inds, tn_inds) # Note: another more efficient implementation is below, where we first get # `sub_tn` of `tn` that is closely connected to `source_inds`, and then compute From 6ff94cf05d825c90a394a29613e59886678fabc3 Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Sat, 15 Jul 2023 16:16:29 -0500 Subject: [PATCH 31/35] Add qr-svd option for density matrix algorithm --- src/approx_itensornetwork/density_matrix.jl | 116 ++++++++++++++---- .../partitioned_contract.jl | 3 +- 2 files changed, 97 insertions(+), 22 deletions(-) diff --git a/src/approx_itensornetwork/density_matrix.jl b/src/approx_itensornetwork/density_matrix.jl index f33cfc48..6ee118f4 100644 --- a/src/approx_itensornetwork/density_matrix.jl +++ b/src/approx_itensornetwork/density_matrix.jl @@ -270,30 +270,15 @@ function _rem_vertex!( contraction_sequence_kwargs, ) caches = alg_graph.caches - outinds_root_to_sim = _densitymatrix_outinds_to_sim(alg_graph.partition, root) - inds_to_sim = merge(alg_graph.innerinds_to_sim, outinds_root_to_sim) dm_dfs_tree = dfs_tree(alg_graph.out_tree, root) - @assert length(child_vertices(dm_dfs_tree, root)) == 1 - for v in post_order_dfs_vertices(dm_dfs_tree, root) - children = sort(child_vertices(dm_dfs_tree, v)) - @assert length(children) <= 2 - network = Vector{ITensor}(alg_graph.partition[v]) - _update!( - caches, - NamedEdge(parent_vertex(dm_dfs_tree, v), v), - children, - Vector{ITensor}(network), - inds_to_sim; - contraction_sequence_alg, - contraction_sequence_kwargs, - ) - end - U = _get_low_rank_projector( - caches.e_to_dm[NamedEdge(nothing, root)], - collect(keys(outinds_root_to_sim)), - collect(values(outinds_root_to_sim)); + U = _update_cache_w_low_rank_projector!( + Algorithm("qr_svd"), + alg_graph, + root; cutoff, maxdim, + contraction_sequence_alg, + contraction_sequence_kwargs, ) # update partition and out_tree root_tensor = _optcontract( @@ -309,12 +294,16 @@ function _rem_vertex!( rem_vertex!(alg_graph.out_tree, root) # update es_to_pdm truncate_dfs_tree = dfs_tree(alg_graph.out_tree, alg_graph.root) + # Remove all partial density matrices that contain `root`, + # since `root` has been removed from the graph. for es in filter(es -> dst(first(es)) == root, keys(caches.es_to_pdm)) delete!(caches.es_to_pdm, es) end for es in filter(es -> dst(first(es)) == new_root, keys(caches.es_to_pdm)) parent_edge = NamedEdge(parent_vertex(truncate_dfs_tree, new_root), new_root) edge_to_remove = NamedEdge(root, new_root) + # The pdm represented by `new_es` will be used to + # update the sibling vertex of `root`. if intersect(es, Set([parent_edge])) == Set() new_es = setdiff(es, [edge_to_remove]) if new_es == Set() @@ -338,6 +327,91 @@ function _rem_vertex!( return U end +function _update_cache_w_low_rank_projector!( + ::Algorithm"direct_eigen", + alg_graph::_DensityMartrixAlgGraph, + root; + cutoff, + maxdim, + contraction_sequence_alg, + contraction_sequence_kwargs, +) + dm_dfs_tree = dfs_tree(alg_graph.out_tree, root) + outinds_root_to_sim = _densitymatrix_outinds_to_sim(alg_graph.partition, root) + # For keys that appear in both dicts, the value in + # `outinds_root_to_sim` is used. + inds_to_sim = merge(alg_graph.innerinds_to_sim, outinds_root_to_sim) + @assert length(child_vertices(dm_dfs_tree, root)) == 1 + for v in post_order_dfs_vertices(dm_dfs_tree, root) + children = sort(child_vertices(dm_dfs_tree, v)) + @assert length(children) <= 2 + network = Vector{ITensor}(alg_graph.partition[v]) + _update!( + alg_graph.caches, + NamedEdge(parent_vertex(dm_dfs_tree, v), v), + children, + Vector{ITensor}(network), + inds_to_sim; + contraction_sequence_alg, + contraction_sequence_kwargs, + ) + end + return U = _get_low_rank_projector( + alg_graph.caches.e_to_dm[NamedEdge(nothing, root)], + collect(keys(outinds_root_to_sim)), + collect(values(outinds_root_to_sim)); + cutoff, + maxdim, + ) +end + +function _update_cache_w_low_rank_projector!( + ::Algorithm"qr_svd", + alg_graph::_DensityMartrixAlgGraph, + root; + cutoff, + maxdim, + contraction_sequence_alg, + contraction_sequence_kwargs, +) + dm_dfs_tree = dfs_tree(alg_graph.out_tree, root) + outinds_root_to_sim = _densitymatrix_outinds_to_sim(alg_graph.partition, root) + # For keys that appear in both dicts, the value in + # `outinds_root_to_sim` is used. + inds_to_sim = merge(alg_graph.innerinds_to_sim, outinds_root_to_sim) + @assert length(child_vertices(dm_dfs_tree, root)) == 1 + # Note: here we are not updating the density matrix of `root` + traversal = post_order_dfs_vertices(dm_dfs_tree, root)[1:(end - 1)] + for v in traversal + children = sort(child_vertices(dm_dfs_tree, v)) + @assert length(children) <= 2 + network = Vector{ITensor}(alg_graph.partition[v]) + _update!( + alg_graph.caches, + NamedEdge(parent_vertex(dm_dfs_tree, v), v), + children, + Vector{ITensor}(network), + inds_to_sim; + contraction_sequence_alg, + contraction_sequence_kwargs, + ) + end + root_t = _optcontract( + Vector{ITensor}(alg_graph.partition[root]); + contraction_sequence_alg, + contraction_sequence_kwargs, + ) + Q, R = factorize( + root_t, collect(keys(outinds_root_to_sim))...; which_decomp="qr", ortho="left" + ) + dm = alg_graph.caches.e_to_dm[NamedEdge( + parent_vertex(dm_dfs_tree, traversal[end]), traversal[end] + )] + R = _optcontract([R, dm]; contraction_sequence_alg, contraction_sequence_kwargs) + U2, _, _ = svd(R, commoninds(Q, R)...; maxdim=maxdim, cutoff=cutoff) + return _optcontract([Q, U2]; contraction_sequence_alg, contraction_sequence_kwargs) +end + """ Approximate a `partition` into an output ITensorNetwork with the binary tree structure defined by `out_tree`. diff --git a/src/partitioned_contract/partitioned_contract.jl b/src/partitioned_contract/partitioned_contract.jl index d8ce500e..0f12a3ce 100644 --- a/src/partitioned_contract/partitioned_contract.jl +++ b/src/partitioned_contract/partitioned_contract.jl @@ -208,7 +208,8 @@ function _ind_orderings(par::DataGraph, contraction_tree::NamedDiGraph; linear_o ) # The inds ordering for each edge is selected based on a # larger partition of the tn. - tn, tn_inds = length(vertices(tn1)) >= length(vertices(tn2)) ? (tn1, inds1) : (tn2, inds2) + tn, tn_inds = + length(vertices(tn1)) >= length(vertices(tn2)) ? (tn1, inds1) : (tn2, inds2) for e in contract_edges source_inds = intersect( p_vs_to_noncommon_inds[[e.src]], p_vs_to_noncommon_inds[[e.dst]] From 27c57791d03a831868e7b941a8a84249d17f09af Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Tue, 18 Jul 2023 14:50:32 -0500 Subject: [PATCH 32/35] Fix a big in density matrix alg with qr svd, add args to choose the alg --- .../approx_itensornetwork.jl | 13 ++++++++--- src/approx_itensornetwork/density_matrix.jl | 23 +++++++++++++++---- src/contract.jl | 2 +- test/test_binary_tree_partition.jl | 2 +- 4 files changed, 31 insertions(+), 9 deletions(-) diff --git a/src/approx_itensornetwork/approx_itensornetwork.jl b/src/approx_itensornetwork/approx_itensornetwork.jl index 5041a8a5..5b818fc3 100644 --- a/src/approx_itensornetwork/approx_itensornetwork.jl +++ b/src/approx_itensornetwork/approx_itensornetwork.jl @@ -5,7 +5,7 @@ with the same binary tree structure. `root` is the root vertex of the pre-order depth-first-search traversal used to perform the truncations. """ function approx_itensornetwork( - ::Algorithm"density_matrix", + alg::Union{Algorithm"density_matrix",Algorithm"density_matrix_direct_eigen"}, binary_tree_partition::DataGraph; root, cutoff=1e-15, @@ -22,6 +22,8 @@ function approx_itensornetwork( partition_wo_deltas = _contract_deltas_ignore_leaf_partitions( binary_tree_partition; root=root ) + density_matrix_alg = + alg isa Algorithm"density_matrix_direct_eigen" ? "direct_eigen" : "qr_svd" return _approx_itensornetwork_density_matrix!( partition_wo_deltas, underlying_graph(binary_tree_partition); @@ -30,6 +32,7 @@ function approx_itensornetwork( maxdim, contraction_sequence_alg, contraction_sequence_kwargs, + density_matrix_alg, ) end @@ -61,7 +64,9 @@ with a binary tree structure. The binary tree structure is defined based on `inds_btree`, which is a directed binary tree DataGraph of indices. """ function approx_itensornetwork( - alg::Union{Algorithm"density_matrix",Algorithm"ttn_svd"}, + alg::Union{ + Algorithm"density_matrix",Algorithm"density_matrix_direct_eigen",Algorithm"ttn_svd" + }, tn::ITensorNetwork, inds_btree::DataGraph; cutoff=1e-15, @@ -96,7 +101,9 @@ Approximate a given ITensorNetwork `tn` into an output ITensorNetwork with `outp `output_structure` outputs a directed binary tree DataGraph defining the desired graph structure. """ function approx_itensornetwork( - alg::Union{Algorithm"density_matrix",Algorithm"ttn_svd"}, + alg::Union{ + Algorithm"density_matrix",Algorithm"density_matrix_direct_eigen",Algorithm"ttn_svd" + }, tn::ITensorNetwork, output_structure::Function=path_graph_structure; cutoff=1e-15, diff --git a/src/approx_itensornetwork/density_matrix.jl b/src/approx_itensornetwork/density_matrix.jl index 6ee118f4..a695c838 100644 --- a/src/approx_itensornetwork/density_matrix.jl +++ b/src/approx_itensornetwork/density_matrix.jl @@ -268,11 +268,12 @@ function _rem_vertex!( maxdim, contraction_sequence_alg, contraction_sequence_kwargs, + density_matrix_alg="qr_svd", ) caches = alg_graph.caches dm_dfs_tree = dfs_tree(alg_graph.out_tree, root) U = _update_cache_w_low_rank_projector!( - Algorithm("qr_svd"), + Algorithm(density_matrix_alg), alg_graph, root; cutoff, @@ -404,11 +405,17 @@ function _update_cache_w_low_rank_projector!( Q, R = factorize( root_t, collect(keys(outinds_root_to_sim))...; which_decomp="qr", ortho="left" ) + qr_commonind = commoninds(Q, R)[1] + R_sim = replaceinds(R, inds_to_sim) + R_sim = replaceinds(R_sim, qr_commonind => sim(qr_commonind)) dm = alg_graph.caches.e_to_dm[NamedEdge( parent_vertex(dm_dfs_tree, traversal[end]), traversal[end] )] - R = _optcontract([R, dm]; contraction_sequence_alg, contraction_sequence_kwargs) - U2, _, _ = svd(R, commoninds(Q, R)...; maxdim=maxdim, cutoff=cutoff) + R = _optcontract([R, dm, R_sim]; contraction_sequence_alg, contraction_sequence_kwargs) + U2 = _get_low_rank_projector( + R, [qr_commonind], setdiff(inds(R), [qr_commonind]); cutoff, maxdim + ) + # U2, _, _ = svd(R, commoninds(Q, R)...; maxdim=maxdim, cutoff=cutoff) return _optcontract([Q, U2]; contraction_sequence_alg, contraction_sequence_kwargs) end @@ -424,7 +431,9 @@ function _approx_itensornetwork_density_matrix!( maxdim=10000, contraction_sequence_alg, contraction_sequence_kwargs, + density_matrix_alg="direct_eigen", ) + @assert density_matrix_alg in ["direct_eigen", "qr_svd"] # Change type of each partition[v] since they will be updated # with potential data type chage. partition = DataGraph() @@ -437,7 +446,13 @@ function _approx_itensornetwork_density_matrix!( output_tn = ITensorNetwork() for v in post_order_dfs_vertices(out_tree, root)[1:(end - 1)] U = _rem_vertex!( - alg_graph, v; cutoff, maxdim, contraction_sequence_alg, contraction_sequence_kwargs + alg_graph, + v; + cutoff, + maxdim, + contraction_sequence_alg, + contraction_sequence_kwargs, + density_matrix_alg, ) add_vertex!(output_tn, v) output_tn[v] = U diff --git a/src/contract.jl b/src/contract.jl index 9c100a12..3b628ca9 100644 --- a/src/contract.jl +++ b/src/contract.jl @@ -10,7 +10,7 @@ function contract( end function contract( - alg::Union{Algorithm"density_matrix",Algorithm"ttn_svd"}, + alg::Union{Algorithm"density_matrix",Algorithm"density_matrix_direct_eigen",Algorithm"ttn_svd"}, tn::AbstractITensorNetwork; output_structure::Function=path_graph_structure, kwargs..., diff --git a/test/test_binary_tree_partition.jl b/test/test_binary_tree_partition.jl index 5160983d..14112ce5 100644 --- a/test/test_binary_tree_partition.jl +++ b/test/test_binary_tree_partition.jl @@ -61,7 +61,7 @@ end end tensors = vec(tn)[1:20] tn = ITensorNetwork(tensors) - @info _mps_partition_inds_order(tn, noncommoninds(tensors...); alg="top_down") + # @info _mps_partition_inds_order(tn, noncommoninds(tensors...); alg="top_down") # TODO: this case is not supported for now, since two indices are adjacent to # one tensor. end From 55dd8b7d58c812a7b929617c640e398ef849c4ad Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Tue, 18 Jul 2023 14:52:09 -0500 Subject: [PATCH 33/35] Add bench for tree tn truncate --- .../approx_itensornetwork/bench_truncate.jl | 106 ++++++++++++++++++ 1 file changed, 106 insertions(+) create mode 100644 examples/approx_itensornetwork/bench_truncate.jl diff --git a/examples/approx_itensornetwork/bench_truncate.jl b/examples/approx_itensornetwork/bench_truncate.jl new file mode 100644 index 00000000..a69b0ce8 --- /dev/null +++ b/examples/approx_itensornetwork/bench_truncate.jl @@ -0,0 +1,106 @@ +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 + +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 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 + +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 + +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 From d9ef21246b231a74a5017f8deca83e66e2fce10f Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Tue, 18 Jul 2023 20:52:06 -0500 Subject: [PATCH 34/35] Add benchmarks for approx contract --- .../bench_mps_times_mpos.jl | 60 +++++++++++++++++++ examples/approx_itensornetwork/bench_peps.jl | 48 +++++++++++++++ ...nch_truncate.jl => bench_tree_truncate.jl} | 32 +--------- examples/approx_itensornetwork/utils.jl | 31 ++++++++++ 4 files changed, 140 insertions(+), 31 deletions(-) create mode 100644 examples/approx_itensornetwork/bench_mps_times_mpos.jl create mode 100644 examples/approx_itensornetwork/bench_peps.jl rename examples/approx_itensornetwork/{bench_truncate.jl => bench_tree_truncate.jl} (76%) create mode 100644 examples/approx_itensornetwork/utils.jl diff --git a/examples/approx_itensornetwork/bench_mps_times_mpos.jl b/examples/approx_itensornetwork/bench_mps_times_mpos.jl new file mode 100644 index 00000000..61b9b7f8 --- /dev/null +++ b/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 diff --git a/examples/approx_itensornetwork/bench_peps.jl b/examples/approx_itensornetwork/bench_peps.jl new file mode 100644 index 00000000..e96e9aa4 --- /dev/null +++ b/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 diff --git a/examples/approx_itensornetwork/bench_truncate.jl b/examples/approx_itensornetwork/bench_tree_truncate.jl similarity index 76% rename from examples/approx_itensornetwork/bench_truncate.jl rename to examples/approx_itensornetwork/bench_tree_truncate.jl index a69b0ce8..9686403b 100644 --- a/examples/approx_itensornetwork/bench_truncate.jl +++ b/examples/approx_itensornetwork/bench_tree_truncate.jl @@ -6,23 +6,7 @@ using ITensorNetworks: ising_network, contract, _nested_vector_to_digraph using OMEinsumContractionOrders using AbstractTrees -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 +include("utils.jl") function balanced_binary_tree_tn(num_leaves; link_space, physical_spaces) sequence = bipartite_sequence(collect(1:num_leaves)) @@ -54,20 +38,6 @@ function unbalanced_binary_tree_tn(num_leaves; link_space, physical_spaces) return randomITensorNetwork(distribution, inds_net), inds_leaves 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 - Random.seed!(1234) TimerOutputs.enable_debug_timings(ITensorNetworks) # TimerOutputs.enable_debug_timings(@__MODULE__) diff --git a/examples/approx_itensornetwork/utils.jl b/examples/approx_itensornetwork/utils.jl new file mode 100644 index 00000000..6eb8c639 --- /dev/null +++ b/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 From 7e9ce10ecb053c86cebf7e9a3110f3047eb8b809 Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Tue, 18 Jul 2023 21:06:00 -0500 Subject: [PATCH 35/35] change runtest back --- test/runtests.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/runtests.jl b/test/runtests.jl index e6d073a0..f02b8f2b 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -6,7 +6,7 @@ using ITensorNetworks @testset "ITensorNetworks.jl, test directory $root" for (root, dirs, files) in walkdir( joinpath(pkgdir(ITensorNetworks), "test") ) - test_files = ["test_binary_tree_partition.jl"]# filter!(f -> occursin(Glob.FilenameMatch("test_*.jl"), f), files) + test_files = filter!(f -> occursin(Glob.FilenameMatch("test_*.jl"), f), files) @testset "Test file $test_file" for test_file in test_files println("Running test file $test_file") @time include(joinpath(root, test_file))