Fixing the iter_connection with unique_node_ids (#57)

* Fixing the iter_connection with unique_node_ids

* fixing the function
* adding comments and explanation in docstrings
* fixing the pylint in tox due to the multiprocessing error
  pylint-dev/pylint#3524

* Change node_id to node ID in doc

bluepysnap/edges.py

@@ -34,6 +34,7 @@
 
 class EdgeStorage(object):
     """Edge storage access."""
+
     def __init__(self, config, circuit):
         """Initializes a EdgeStorage object from a edge config and a Circuit.
 
@@ -369,6 +370,7 @@ def _optimal_direction():
             if target_node_ids is None:
                 return 'source'
             else:
+                # Checking the indexing 'direction'. One direction has contiguous indices.
                 range_size_source = _estimate_range_size(
                     self._population.efferent_edges, source_node_ids
                 )
@@ -399,6 +401,7 @@ def _optimal_direction():
 
         for key_node_id in primary_node_ids:
             connected_node_ids = get_connected_node_ids(key_node_id, unique=False)
+            # [[secondary_node_id, count], ...]
             connected_node_ids_with_count = np.stack(
                 np.unique(connected_node_ids, return_counts=True)
             ).transpose()
@@ -410,7 +413,10 @@ def _optimal_direction():
                 connected_node_ids_with_count = connected_node_ids_with_count[mask]
             if shuffle:
                 np.random.shuffle(connected_node_ids_with_count)
+
             for conn_node_id, edge_count in connected_node_ids_with_count:
+                if unique_node_ids and (conn_node_id in secondary_node_ids_used):
+                    continue
                 if direction == 'target':
                     yield conn_node_id, key_node_id, edge_count
                 else:
@@ -428,7 +434,9 @@ def iter_connections(
         Args:
             source: source node group
             target: target node group
-            unique_node_ids: if True, no node ID would be used more than once
+            unique_node_ids: if True, no node ID will be used more than once as source or
+            target for edges. Careful, this flag does not provide unique (source, target)
+            pairs but unique node IDs.
             shuffle: if True, result order would be (somewhat) randomized
             return_edge_count: if True, edge count is added to yield result
             return_edge_ids: if True, edge ID list is added to yield result

tests/test_edges.py

@@ -420,3 +420,27 @@ def test_iter_connections_10(self):
             self.test_obj.iter_connections(
                 [0, 2], [1], return_edge_ids=True, return_edge_count=True
             )
+
+    def test_iter_connection_unique(self):
+        test_obj = TestEdgePopulation.create_population(
+            str(TEST_DATA_DIR / "edges_complete_graph.h5"), 'default')
+        it = test_obj.iter_connections([0, 1, 2], [0, 1, 2])
+        assert sorted(it) == [(0, 1), (0, 2), (1, 0), (1, 2), (2, 0), (2, 1)]
+
+        it = test_obj.iter_connections([0, 1, 2], [0, 1, 2], unique_node_ids=True)
+        assert sorted(it) == [(0, 1), (1, 0)]
+
+        it = test_obj.iter_connections([0, 1, 2], [0, 2], unique_node_ids=True)
+        assert sorted(it) == [(0, 2), (1, 0)]
+
+        it = test_obj.iter_connections([0, 2], [0, 2], unique_node_ids=True)
+        assert sorted(it) == [(0, 2), (2, 0)]
+
+        it = test_obj.iter_connections([0, 1, 2], [0, 2, 1], unique_node_ids=True)
+        assert sorted(it) == [(0, 1), (1, 0)]
+
+        it = test_obj.iter_connections([1, 2], [0, 1, 2], unique_node_ids=True)
+        assert sorted(it) == [(1, 0), (2, 1)]
+
+        it = test_obj.iter_connections([0, 1, 2], [1, 2], unique_node_ids=True)
+        assert sorted(it) == [(0, 1), (1, 2)]

tox.ini

@@ -19,7 +19,7 @@ basepython=python3.6
 deps =
     pycodestyle
     pydocstyle
-    pylint
+    pylint==2.4.4  # see https://github.com/PyCQA/pylint/issues/3524
 commands =
     pycodestyle {[base]name}
     pydocstyle {[base]name}