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)