def save_nodes(nodes: List[AbstractNode], path: Union[str, BinaryIO]) -> None:
"""Save an iterable of nodes into hdf5 format.
Args:
nodes: An iterable of connected nodes. All nodes have to connect within
`nodes`.
path: path to file where network is saved.
"""
if reachable(nodes) > set(nodes):
raise ValueError(
"Some nodes in `nodes` are connected to nodes not contained in `nodes`."
" Saving not possible.")
if len(set(nodes)) < len(list(nodes)):
raise ValueError(
'Some nodes in `nodes` appear more than once. This is not supported')
#we need to iterate twice and order matters
edges = list(get_all_edges(nodes))
nodes = list(nodes)
old_edge_names = {n: edge.name for n, edge in enumerate(edges)}
old_node_names = {n: node.name for n, node in enumerate(nodes)}
#generate unique names for nodes and edges
#for saving them
for n, node in enumerate(nodes):
node.set_name('node{}'.format(n))
for e, edge in enumerate(edges):
edge.set_name('edge{}'.format(e))
with h5py.File(path, 'w') as net_file:
nodes_group = net_file.create_group('nodes')
node_names_group = net_file.create_group('node_names')
node_names_group.create_dataset(
'names',
dtype=string_type,
data=np.array(list(old_node_names.values()), dtype=object))
edges_group = net_file.create_group('edges')
edge_names_group = net_file.create_group('edge_names')
edge_names_group.create_dataset(
'names',
dtype=string_type,
data=np.array(list(old_edge_names.values()), dtype=object))
for n, node in enumerate(nodes):
node_group = nodes_group.create_group(node.name)
node._save_node(node_group)
for edge in node.edges:
if edge.node1 == node and edge in edges:
edge_group = edges_group.create_group(edge.name)
edge._save_edge(edge_group)
edges.remove(edge)
#name edges and nodes back to their original names
for n, node in enumerate(nodes):
nodes[n].set_name(old_node_names[n])
for n, edge in enumerate(edges):
edges[n].set_name(old_edge_names[n])
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 auto(
nodes: BaseNode,
output_edge_order: Optional[Sequence[Edge]] = None,
memory_limit: Optional[int] = None,
ignore_edge_order: bool = False) -> BaseNode:
"""Chooses one of the above algorithms according to network size.
Default behavior is based on `opt_einsum`'s `auto` contractor.
Args:
nodes: A collection of connected nodes.
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.
ignore_edge_order: An option to ignore the output edge order.
Returns:
Final node after full contraction.
"""
n = len(list(nodes)) #pytype thing
_nodes = nodes
if n <= 0:
raise ValueError("Cannot contract empty tensor network.")
if n == 1:
if not ignore_edge_order:
if output_edge_order is None:
output_edge_order = list(
(get_all_edges(_nodes) - get_all_nondangling(_nodes)))
if len(output_edge_order) > 1:
raise ValueError("The final node after contraction has more than "
"one dangling edge. In this case `output_edge_order` "
"has to be provided.")
edges = get_all_nondangling(_nodes)
if edges:
final_node = contract_parallel(edges.pop())
else:
final_node = list(_nodes)[0]
final_node.reorder_edges(output_edge_order)
if not ignore_edge_order:
final_node.reorder_edges(output_edge_order)
return final_node
if n < 5:
return optimal(nodes, output_edge_order, memory_limit, ignore_edge_order)
if n < 7:
return branch(nodes, output_edge_order, memory_limit, ignore_edge_order)
if n < 9:
return branch(nodes, output_edge_order, memory_limit, nbranch=2, ignore_edge_order=ignore_edge_order)
if n < 15:
return branch(nodes, output_edge_order, nbranch=1, ignore_edge_order=ignore_edge_order)
return greedy(nodes, output_edge_order, memory_limit, ignore_edge_order)
def _get_path_nodes(
nodes: Iterable[BaseNode],
algorithm: Algorithm) -> Tuple[List[Tuple[int, int]], List[BaseNode]]:
"""Calculates the contraction paths using `opt_einsum` methods.
Args:
nodes: An iterable of nodes.
algorithm: `opt_einsum` method to use for calculating the contraction path.
Returns:
The optimal contraction path as returned by `opt_einsum`.
"""
sorted_nodes = sorted(nodes, key=lambda n: n.signature)
input_sets = [set(node.edges) for node in sorted_nodes]
output_set = get_all_edges(nodes) - get_all_nondangling(nodes)
size_dict = {edge: edge.dimension for edge in get_all_edges(nodes)}
return algorithm(input_sets, output_set, size_dict), sorted_nodes
def contract_path(path: Tuple[List[Tuple[int,
int]]], nodes: Iterable[AbstractNode],
output_edge_order: Sequence[Edge]) -> AbstractNode:
"""Contract `nodes` using `path`.
Args:
path: The contraction path as returned from `path_solver`.
nodes: A collection of connected nodes.
output_edge_order: A list of edges. Edges of the
final node in `nodes`
are reordered into `output_edge_order`;
Returns:
Final node after full contraction.
"""
edges = get_all_edges(nodes)
for edge in edges:
if not edge.is_disabled: #if its disabled we already contracted it
if edge.is_trace():
contract_parallel(edge)
if len(nodes) == 1:
newnode = nodes[0].copy()
for edge in nodes[0].edges:
redirect_edge(edge, newnode, nodes[0])
return newnode.reorder_edges(output_edge_order)
if len(path) == 0:
return nodes
for p in path:
if len(p) > 1:
a, b = p
new_node = contract_between(nodes[a],
nodes[b],
allow_outer_product=True)
nodes.append(new_node)
nodes = utils.multi_remove(nodes, [a, b])
elif len(p) == 1:
a = p[0]
node = nodes.pop(a)
new_node = contract_trace_edges(node)
nodes.append(new_node)
# 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
#some contractors miss trace edges
final_node = contract_trace_edges(final_node)
final_node.reorder_edges(output_edge_order)
return final_node
def get_path(
nodes: Iterable[AbstractNode], algorithm: Algorithm
) -> Tuple[List[Tuple[int, int]], List[AbstractNode]]:
"""Calculates the contraction paths using `opt_einsum` methods.
Args:
nodes: An iterable of nodes.
algorithm: `opt_einsum` method to use for calculating the contraction path.
Returns:
The optimal contraction path as returned by `opt_einsum`.
"""
nodes = list(nodes)
input_sets = [set(node.edges) for node in nodes]
output_set = get_subgraph_dangling(nodes)
size_dict = {edge: edge.dimension for edge in get_all_edges(nodes)}
return algorithm(input_sets, output_set, size_dict), nodes
def get_all_edges(self):
"""Return the set of all edges."""
return network_operations.get_all_edges(self.nodes_set)
