def base(
net: network.TensorNetwork,
algorithm: Callable[[List[Set[int]], Set[int], Dict[int, int]], List],
output_edge_order: Optional[Sequence[network_components.Edge]] = None
) -> network.TensorNetwork:
"""Base method for all `opt_einsum` contractors.
Args:
net: a TensorNetwork object. Should be connected.
algorithm: `opt_einsum` contraction method to use.
output_edge_order: An optional list of edges. Edges of the
final node in `nodes_set`
are reordered into `output_edge_order`;
if final node has more than one edge,
`output_edge_order` must be provided.
Returns:
The network after full contraction.
"""
net.check_connected()
# First contract all trace edges
edges = net.get_all_nondangling()
for edge in edges:
if edge in net and edge.is_trace():
net.contract_parallel(edge)
if not net.get_all_nondangling():
# There's nothing to contract.
return net
# Then apply `opt_einsum`'s algorithm
nodes = sorted(net.nodes_set)
input_sets = utils.get_input_sets(net)
output_set = utils.get_output_set(net)
size_dict = utils.get_size_dict(net)
path = algorithm(input_sets, output_set, size_dict)
for a, b in path:
new_node = nodes[a] @ nodes[b]
nodes.append(new_node)
nodes = utils.multi_remove(nodes, [a, b])
# if the final node has more than one edge,
# output_edge_order has to be specified
final_node = net.get_final_node()
if (len(final_node.edges) <= 1) and (output_edge_order is None):
output_edge_order = list(
(net.get_all_edges() - net.get_all_nondangling()))
elif (len(final_node.edges) > 1) and (output_edge_order is None):
raise ValueError("if the final node has more than one dangling edge"
" `output_edge_order` has to be provided")
if set(output_edge_order) != (net.get_all_edges() -
net.get_all_nondangling()):
raise ValueError("output edges are not all dangling.")
final_node.reorder_edges(output_edge_order)
return net
def _base_network(net: TensorNetwork,
algorithm: utils.Algorithm,
output_edge_order: Optional[Sequence[Edge]] = None,
ignore_edge_order: bool = False) -> TensorNetwork:
"""Base method for all `opt_einsum` contractors.
Args:
net: a TensorNetwork object. Should be connected.
algorithm: `opt_einsum` contraction method to use.
output_edge_order: An optional list of edges. Edges of the
final node in `nodes_set`
are reordered into `output_edge_order`;
if final node has more than one edge,
`output_edge_order` must be provided.
ignore_edge_order: An option to ignore the output edge
order.
Returns:
The network after full contraction.
"""
net.check_connected()
# First contract all trace edges
edges = net.get_all_nondangling()
for edge in edges:
if edge in net and edge.is_trace():
net.contract_parallel(edge)
if not net.get_all_nondangling():
# There's nothing to contract.
return net
# Then apply `opt_einsum`'s algorithm
path, nodes = utils.get_path(net, algorithm)
for a, b in path:
new_node = nodes[a] @ nodes[b]
nodes.append(new_node)
nodes = utils.multi_remove(nodes, [a, b])
# if the final node has more than one edge,
# output_edge_order has to be specified
final_node = net.get_final_node()
if not ignore_edge_order:
if (len(final_node.edges) <= 1) and (output_edge_order is None):
output_edge_order = list(
(net.get_all_edges() - net.get_all_nondangling()))
elif (len(final_node.edges) > 1) and (output_edge_order is None):
raise ValueError(
"The final node after contraction has more than "
"one dangling edge. In this case `output_edge_order` "
"has to be provided.")
if set(output_edge_order) != (net.get_all_edges() -
net.get_all_nondangling()):
raise ValueError("output edges are not all dangling.")
final_node.reorder_edges(output_edge_order)
return net
def _get_path_network(net: TensorNetwork, algorithm: Algorithm
) -> Tuple[List[Tuple[int, int]], List[BaseNode]]:
"""Calculates the contraction paths using `opt_einsum` methods.
Args:
net: TensorNetwork object to contract.
algorithm: `opt_einsum` method to use for calculating the contraction path.
Returns:
The optimal contraction path as returned by `opt_einsum`.
"""
sorted_nodes = sorted(net.nodes_set, key=lambda n: n.signature)
input_sets = [set(node.edges) for node in sorted_nodes]
output_set = net.get_all_edges() - net.get_all_nondangling()
print(output_set)
size_dict = {edge: edge.dimension for edge in net.get_all_edges()}
return algorithm(input_sets, output_set, size_dict), sorted_nodes
def auto(net: network.TensorNetwork,
output_edge_order: Sequence[network_components.Edge] = None,
memory_limit: Optional[int] = None) -> network.TensorNetwork:
"""Chooses one of the above algorithms according to network size.
Default behavior is based on `opt_einsum`'s `auto` contractor.
Args:
net: a TensorNetwork object.
output_edge_order: An optional list of edges.
Edges of the final node in `nodes_set`
are reordered into `output_edge_order`;
if final node has more than one edge,
`output_edge_order` must be provided.
memory_limit: Maximum number of elements in an array during contractions.
Returns:
The network after full contraction.
"""
n = len(net.nodes_set)
if n <= 0:
raise ValueError("Cannot contract empty tensor network.")
if n == 1:
edges = net.get_all_nondangling()
net.contract_parallel(edges.pop())
final_node = net.get_final_node()
if (len(final_node.edges) <= 1) and (output_edge_order is None):
output_edge_order = list(
(net.get_all_edges() - net.get_all_nondangling()))
elif (len(final_node.edges) > 1) and (output_edge_order is None):
raise ValueError(
"if the final node has more than one dangling edge"
", `output_edge_order` has to be provided")
final_node.reorder_edges(output_edge_order)
return net
if n < 5:
return optimal(net, output_edge_order, memory_limit)
if n < 7:
return branch(net, output_edge_order, memory_limit)
if n < 9:
return branch(net, output_edge_order, memory_limit, nbranch=2)
if n < 15:
return branch(net, output_edge_order, nbranch=1)
return greedy(net, output_edge_order, memory_limit)
def get_size_dict(net: network.TensorNetwork) -> Dict[int, int]:
return {edge: edge.dimension for edge in net.get_all_edges()}
def get_output_set(net: network.TensorNetwork) -> Set[int]:
dangling_edges = net.get_all_edges() - net.get_all_nondangling()
return set(dangling_edges)