def test_path_optimal(params):
algorithm_name, network_name, correct_path = params
net = globals()[network_name]()
path_algorithm = getattr(opt_einsum.paths, algorithm_name)
calculated_path, _ = utils.get_path(net, path_algorithm)
assert check_path(calculated_path, correct_path)
def path_solver(
algorithm: Text,
nodes: Iterable[AbstractNode],
memory_limit: Optional[int] = None,
nbranch: Optional[int] = None
) -> Tuple[List[Tuple[int, int]], List[AbstractNode]]:
"""Calculates the contraction paths using `opt_einsum` methods.
Args:
algorithm: `opt_einsum` method to use for calculating the contraction path.
nodes: an iterable of `AbstractNode` objects to contract.
memory_limit: Maximum number of elements in an array during contractions.
Only relevant for `algorithm in (optimal, greedy)`
nbranch: Number of best contractions to explore.
If None it explores all inner products starting with those that
have the best cost heuristic. Only relevant for `algorithm=branch`.
Returns:
The optimal contraction path as returned by `opt_einsum`.
"""
if algorithm == "optimal":
alg = functools.partial(opt_einsum.paths.dynamic_programming,
memory_limit=memory_limit)
elif algorithm == "branch":
alg = functools.partial(opt_einsum.paths.branch,
memory_limit=memory_limit,
nbranch=nbranch)
elif algorithm == "greedy":
alg = functools.partial(opt_einsum.paths.greedy,
memory_limit=memory_limit)
elif algorithm == "auto":
n = len(list(nodes)) #pytype thing
_nodes = nodes
if n <= 1:
return []
if n < 5:
alg = functools.partial(opt_einsum.paths.dynamic_programming,
memory_limit=memory_limit)
if n < 7:
alg = functools.partial(opt_einsum.paths.branch,
memory_limit=memory_limit,
nbranch=None)
if n < 9:
alg = functools.partial(opt_einsum.paths.branch,
memory_limit=memory_limit,
nbranch=2)
if n < 15:
alg = functools.partial(opt_einsum.paths.branch,
memory_limit=memory_limit,
nbranch=1)
else:
alg = functools.partial(opt_einsum.paths.greedy,
memory_limit=memory_limit)
else:
raise ValueError("algorithm {algorithm} not implemented")
path, _ = utils.get_path(nodes, alg)
return path
def _base_nodes(
nodes: Iterable[BaseNode],
algorithm: utils.Algorithm,
output_edge_order: Optional[Sequence[Edge]] = None) -> BaseNode:
"""Base method for all `opt_einsum` contractors.
Args:
nodes: A collection of connected nodes.
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 pronvided.
Returns:
Final node after full contraction.
"""
nodes_set = set(nodes)
check_connected(nodes_set)
edges = get_all_edges(nodes_set)
#output edge order has to be determinded before any contraction
#(edges are refreshed after contractions)
if output_edge_order is None:
output_edge_order = list(get_subgraph_dangling(nodes))
if len(output_edge_order) > 1:
raise ValueError(
"The final node after contraction has more than "
"one remaining edge. In this case `output_edge_order` "
"has to be provided.")
if set(output_edge_order) != get_subgraph_dangling(nodes):
raise ValueError("output edges are not equal to the remaining "
"non-contracted edges of the final node.")
for edge in edges:
if not edge.is_disabled: #if its disabled we already contracted it
if edge.is_trace():
nodes_set.remove(edge.node1)
nodes_set.add(contract_parallel(edge))
if len(nodes_set) == 1:
# There's nothing to contract.
return list(nodes_set)[0].reorder_edges(output_edge_order)
# Then apply `opt_einsum`'s algorithm
path, nodes = utils.get_path(nodes_set, 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 = nodes[0] # nodes were connected, we checked this
final_node.reorder_edges(output_edge_order)
return final_node
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
