Implement queries mechanism for edges (#129)

CHANGELOG.rst

@@ -1,6 +1,13 @@
 Changelog
 =========
 
+Version v0.11.0
+---------------
+
+New Features
+~~~~~~~~~~~~~
+- Implement queries mechanism for edges
+
 Version v0.10.0
 ---------------
 

bluepysnap/edges.py

@@ -17,19 +17,21 @@
 
 """Edge population access."""
 import inspect
+from collections.abc import Mapping
 
 import libsonata
 import numpy as np
 import pandas as pd
 from cached_property import cached_property
 from more_itertools import first
 
+from bluepysnap import query
 from bluepysnap.network import NetworkObject
 from bluepysnap.exceptions import BluepySnapError
 from bluepysnap.circuit_ids import CircuitEdgeId, CircuitEdgeIds, CircuitNodeId, CircuitNodeIds
 from bluepysnap.sonata_constants import DYNAMICS_PREFIX, Edge, ConstContainer
 from bluepysnap import utils
-from bluepysnap.utils import Deprecate
+from bluepysnap.utils import Deprecate, IDS_DTYPE
 from bluepysnap._doctools import AbstractDocSubstitutionMeta
 
 
@@ -480,7 +482,37 @@ def _get(self, selection, properties=None):
 
         return result
 
-    def ids(self, group=None, limit=None, sample=None):
+    def _edge_ids_by_filter(self, queries, raise_missing_prop):
+        """Return edge IDs if their properties match the `queries` dict.
+
+        `props` values could be:
+            pairs (range match for floating dtype fields)
+            scalar or iterables (exact or "one of" match for other fields)
+
+        You can use the special operators '$or' and '$and' also to combine different queries
+        together.
+
+        Examples:
+            >>> self._edge_ids_by_filter({ Edge.POST_SECTION_ID: (0, 1),
+            >>>                            Edge.AXONAL_DELAY: (.5, 2.) })
+            >>> self._edge_ids_by_filter({'$or': [{ Edge.PRE_X_CENTER: [2, 3]},
+            >>>                              { Edge.POST_SECTION_POS: (0, 1),
+            >>>                              Edge.SYN_WEIGHT: (0.,1.4) }]})
+
+        """
+        properties = query.get_properties(queries)
+        unknown_props = properties - self.property_names
+        if raise_missing_prop and unknown_props:
+            raise BluepySnapError(f"Unknown edge properties: {unknown_props}")
+        res = []
+        ids = self.ids(None)
+        chunk_size = int(1e8)
+        for chunk in np.array_split(ids, 1 + len(ids) // chunk_size):
+            data = self.get(chunk, properties - unknown_props)
+            res.extend(chunk[query.resolve_ids(data, self.name, queries)])
+        return np.array(res, dtype=IDS_DTYPE)
+
+    def ids(self, group=None, limit=None, sample=None, raise_missing_property=True):
         """Edge IDs corresponding to edges ``edge_ids``.
 
         Args:
@@ -499,6 +531,9 @@ def ids(self, group=None, limit=None, sample=None):
                 IDs from the match result. If limit is greater than the size of the population
                 all node IDs are returned.
 
+            raise_missing_property (bool): if True, raises if a property is not listed in this
+                population. Otherwise the ids are just not selected if a property is missing.
+
         Returns:
             numpy.array: A numpy array of IDs.
         """
@@ -508,6 +543,9 @@ def ids(self, group=None, limit=None, sample=None):
             result = group.filter_population(self.name).get_ids()
         elif isinstance(group, np.ndarray):
             result = group
+        elif isinstance(group, Mapping):
+            result = self._edge_ids_by_filter(queries=group,
+                                              raise_missing_prop=raise_missing_property)
         else:
             result = utils.ensure_list(group)
             # test if first value is a CircuitEdgeId if yes then all values must be CircuitEdgeId

bluepysnap/nodes.py

@@ -27,10 +27,10 @@
 
 from cached_property import cached_property
 
+from bluepysnap import query
 from bluepysnap.network import NetworkObject
 from bluepysnap import utils
 from bluepysnap.exceptions import BluepySnapError
-import bluepysnap.query as query
 from bluepysnap.sonata_constants import (DYNAMICS_PREFIX, Node, ConstContainer)
 from bluepysnap.circuit_ids import CircuitNodeId, CircuitNodeIds
 from bluepysnap._doctools import AbstractDocSubstitutionMeta

tests/test_edges.py

@@ -185,6 +185,39 @@ def test_ids(self):
         expected = CircuitEdgeIds.from_dict({"default": [0, 1, 2, 3], "default2": [0]})
         assert tested == expected
 
+        with pytest.raises(BluepySnapError) as e:
+            self.test_obj.ids({'afferent_center_i': (10, 11)})
+        assert "Unknown edge properties: {'afferent_center_i'}" == e.value.args[0]
+
+        tested = self.test_obj.ids({"afferent_center_x": (1110, 1110.5)})
+        expected = CircuitEdgeIds.from_dict({"default": [0], "default2": [0]})
+        assert tested == expected
+
+        tested = self.test_obj.ids({"afferent_center_x": (1111, 1112), "efferent_center_z": (2132, 2134)})
+        expected = CircuitEdgeIds.from_dict({"default": [2], "default2": [2]})
+        assert tested == expected
+
+        tested = self.test_obj.ids({'$and': [{"@dynamics:param1": (0, 2)}, {"afferent_surface_x": (1211, 1211)}]})
+        expected = CircuitEdgeIds.from_dict({"default": [1], "default2": [1]})
+        assert tested == expected
+
+        tested = self.test_obj.ids({'$or': [{"@dynamics:param1": (0, 2)}, {"@source_node": [0]}]})
+        expected = CircuitEdgeIds.from_dict({"default": [0, 1, 2], "default2": [0, 1, 2]})
+        assert tested == expected
+
+        tested = self.test_obj.ids({"population": ["default2"], "afferent_center_x": (1113, 1114)})
+        expected = CircuitEdgeIds.from_dict({"default2": [3]})
+        assert tested == expected
+
+        tested = self.test_obj.ids({"population": ["default3"], "afferent_center_x": (1113, 1114)})
+        expected = CircuitEdgeIds.from_arrays([], [])
+        assert tested == expected
+
+        tested = self.test_obj.ids({"population": ["default", "default2"], "@target_node": [1]})
+        expected = CircuitEdgeIds.from_dict({"default": [1, 2, 3], "default2": [1, 2, 3]})
+        assert tested == expected
+
+
     def test_get(self):
         with pytest.raises(BluepySnapError):
             self.test_obj.get(properties=["other2", "unknown"])