예제 #1
0
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)
예제 #2
0
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
예제 #3
0
 def get_all_nondangling(self):
     """Return the set of all non-dangling edges."""
     return network_components.get_all_nondangling(self.nodes_set)