Get neighbors (#588)

* Added get_neighbors method

* documentation

* added bottom space

Author: Chase Roberts

docs/network.rst

@@ -33,4 +33,5 @@ Common Functions
 .. autofunction:: tensornetwork.flatten_all_edges
 .. autofunction:: tensornetwork.flatten_edges_between
 .. autofunction:: tensornetwork.get_parallel_edges
-.. autofunction:: tensornetwork.get_shared_edges
+.. autofunction:: tensornetwork.get_shared_edges
+.. autofunction:: tensornetwork.get_neighbors

tensornetwork/__init__.py

@@ -1,7 +1,7 @@
 #pylint: disable=line-too-long
 from tensornetwork.network_components import Node, Edge, CopyNode, BaseNode, NodeCollection
 #pylint: disable=line-too-long
-from tensornetwork.network_operations import norm, conj, copy, replicate_nodes, transpose, split_node, split_node_qr, split_node_rq, split_node_full_svd, reachable, check_connected, check_correct, get_all_nodes, get_all_edges, remove_node, contract_trace_edges, get_subgraph_dangling, reduced_density, switch_backend
+from tensornetwork.network_operations import norm, conj, copy, replicate_nodes, transpose, split_node, split_node_qr, split_node_rq, split_node_full_svd, reachable, check_connected, check_correct, get_all_nodes, get_all_edges, remove_node, contract_trace_edges, get_subgraph_dangling, reduced_density, switch_backend, get_neighbors
 #pylint: disable=line-too-long
 from tensornetwork.network_components import contract, contract_copy_node, contract_between, outer_product, outer_product_final_nodes, contract_parallel, flatten_edges, split_edge, slice_edge, get_all_nondangling, get_all_dangling, flatten_all_edges, flatten_edges_between, get_parallel_edges, get_shared_edges
 from tensornetwork.backends.base_backend import BaseBackend

tensornetwork/network_operations.py

@@ -874,3 +874,24 @@ def switch_backend(nodes: Iterable[BaseNode], new_backend: Text) -> None:
                                 "is '{}'".format(node.backend))
     node.tensor = backend.convert_to_tensor(node.tensor)
     node.backend = backend
+
+def get_neighbors(node: BaseNode) -> List[Node]:
+  """Get all of the neighbors that are directly connected to the given node.
+
+  Note: `node` will never be in the returned list, even if `node` has a
+  trace edge.
+
+  Args:
+    node: A node.
+
+  Returns:
+    All of the neighboring edges that share an `Edge` with `node`.
+  """
+  neighbors = []
+  for edge in node.edges:
+    if not edge.is_dangling() and not edge.is_trace():
+      if edge.node1 is node:
+        neighbors.append(edge.node2)
+      else:
+        neighbors.append(edge.node1)
+  return neighbors

tensornetwork/tests/network_operations_test.py

@@ -500,3 +500,16 @@ def test_split_node_qr_orig_shape(backend):
   n1 = tn.Node(np.random.rand(3, 4, 5), backend=backend)
   tn.split_node_qr(n1, [n1[0], n1[2]], [n1[1]])
   np.testing.assert_allclose(n1.shape, (3, 4, 5))
+
+def test_get_neighbors(backend):
+  with tn.DefaultBackend(backend):
+    a = tn.Node(np.ones((2, 2)))
+    b = tn.Node(np.ones((2, 2, 2, 2)))
+    c = tn.Node(np.ones((2, 2, 2)))
+    d = tn.Node(np.ones((2, 2)))
+    b[0] ^ a[1]
+    b[3] ^ c[2]
+    a[0] ^ d[1]
+    b[1] ^ b[2]
+    result = tn.get_neighbors(b)
+    assert result == [a, c]