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)