Use libsonata to resolve node sets (#213)

* Node set resolution done by libsonata

* `NodeSets` object can be instantiated with three methods: `from_file`, `from_string`, `from_dict`

* `Circuit.node_sets` and `Simulation.node_sets` return `NodeSets` object initialized with empty dict when node sets file is not present

* `NodeSet.resolved` is no longer available

* `FrameReport.node_set` returns node_set name instead of resolved node set query

---------

Co-authored-by: Gianluca Ficarelli <26835404+GianlucaFicarelli@users.noreply.github.com>

CHANGELOG.rst

@@ -1,6 +1,24 @@
 Changelog
 =========
 
+Version v1.1.0
+--------------
+
+New Features
+~~~~~~~~~~~~
+- ``NodeSets`` object can be instantiated with three methods: ``from_file``, ``from_string``, ``from_dict``
+
+Improvements
+~~~~~~~~~~~~
+- Node set resolution is done by libsonata
+
+Breaking Changes
+~~~~~~~~~~~~~~~~
+- ``Circuit.node_sets``, ``Simulation.node_sets`` returns ``NodeSets`` object initialized with empty dict when node sets file is not present
+- ``NodeSet.resolved`` is no longer available
+- ``FrameReport.node_set`` returns node_set name instead of resolved node set query
+
+
 Version v1.0.7
 --------------
 

bluepysnap/circuit.py

@@ -77,9 +77,8 @@ def get_edge_population_config(self, name):
     @cached_property
     def node_sets(self):
         """Returns the NodeSets object bound to the circuit."""
-        if "node_sets_file" in self.config:
-            return NodeSets(self.config["node_sets_file"])
-        return {}
+        path = self.to_libsonata.node_sets_path
+        return NodeSets.from_file(path) if path else NodeSets.from_dict({})
 
     @cached_property
     def nodes(self):

bluepysnap/frame_report.py

@@ -245,8 +245,8 @@ def data_units(self):
 
     @property
     def node_set(self):
-        """Returns the node set for the report."""
-        return self.simulation.node_sets[self.to_libsonata.cells]
+        """Returns the name of the node set for the report."""
+        return self.to_libsonata.cells
 
     @property
     def simulation(self):

bluepysnap/node_sets.py

@@ -20,88 +20,45 @@
 For more information see:
 https://github.com/AllenInstitute/sonata/blob/master/docs/SONATA_DEVELOPER_GUIDE.md#node-sets-file
 """
-from collections.abc import Mapping
-from copy import deepcopy
 
-import numpy as np
+import json
+
+import libsonata
 
 from bluepysnap import utils
 from bluepysnap.exceptions import BluepySnapError
 
 
-def _sanitize(node_set):
-    """Sanitize standard node set (not compounds).
-
-    Set a single value instead of a one element list.
-    Sorted and unique values for the lists of values.
-
-    Args:
-        node_set (Mapping): A standard non compound node set.
-
-    Return:
-         map: The sanitized node set.
-    """
-    for key, values in node_set.items():
-        if isinstance(values, list):
-            if len(values) == 1:
-                node_set[key] = values[0]
-            else:
-                # sorted unique value list
-                node_set[key] = np.unique(np.asarray(values)).tolist()
-    return node_set
-
-
-def _resolve_set(content, resolved, node_set_name):
-    """Resolve the node set 'node_set_name' from content.
-
-    The resolved node set is returned and the resolved dict is updated in place with the
-    resolved node set.
-
-    Args:
-        content (dict): the global dictionary containing all unresolved node sets.
-        resolved (dict): the global resolved dictionary containing the already resolved node sets.
-        node_set_name (str): the name of the current node set to resolve.
+class NodeSet:
+    """Access to single node set."""
 
-    Returns:
-        dict: the resolved node set.
+    def __init__(self, node_sets, name):
+        """Initializes a single node set object.
 
-    Notes:
-        If the node set is a compound node set then all the sub node sets are also resolved and
-        stored inside the resolved dictionary.
-    """
-    if node_set_name in resolved:
-        # return already resolved node_sets
-        return resolved[node_set_name]
-
-    # keep the content intact
-    set_value = deepcopy(content.get(node_set_name))
-    if set_value is None:
-        raise BluepySnapError(f"Missing node_set: '{node_set_name}'")
-    if not isinstance(set_value, (Mapping, list)) or not set_value:
-        raise BluepySnapError(f"Ambiguous node_set: { {node_set_name: set_value} }")
-    if isinstance(set_value, Mapping):
-        resolved[node_set_name] = _sanitize(set_value)
-        return resolved[node_set_name]
-
-    # compounds only
-    res = [_resolve_set(content, resolved, sub_set_name) for sub_set_name in set_value]
-
-    resolved[node_set_name] = {"$or": res}
-    return resolved[node_set_name]
+        Args:
+            node_sets (libsonata.NodeSets): libsonata NodeSets instance.
+            name (str): name of the node set.
 
+        Returns:
+            NodeSet: A NodeSet object.
+        """
+        self._node_sets = node_sets
+        self._name = name
 
-def _resolve(content):
-    """Resolve all node sets in content."""
-    resolved = {}
-    for set_name in content:
-        _resolve_set(content, resolved, set_name)
-    return resolved
+    def get_ids(self, population, raise_missing_property=True):
+        """Get the resolved node set as ids."""
+        try:
+            return self._node_sets.materialize(self._name, population).flatten()
+        except libsonata.SonataError as e:
+            if not raise_missing_property and "No such attribute" in e.args[0]:
+                return []
+            raise BluepySnapError(*e.args) from e
 
 
 class NodeSets:
     """Access to node sets data."""
 
-    def __init__(self, filepath):
+    def __init__(self, content, instance):
         """Initializes a node set object from a node sets file.
 
         Args:
@@ -110,13 +67,41 @@ def __init__(self, filepath):
         Returns:
             NodeSets: A NodeSets object.
         """
-        self.content = utils.load_json(filepath)
-        self.resolved = _resolve(self.content)
-
-    def __getitem__(self, node_set_name):
-        """Get the resolved node set using name as key."""
-        return self.resolved[node_set_name]
+        self.content = content
+        self._instance = instance
+
+    @classmethod
+    def from_file(cls, filepath):
+        """Create NodeSets instance from a file."""
+        content = utils.load_json(filepath)
+        instance = libsonata.NodeSets.from_file(filepath)
+        return cls(content, instance)
+
+    @classmethod
+    def from_string(cls, content):
+        """Create NodeSets instance from a JSON string."""
+        instance = libsonata.NodeSets(content)
+        content = json.loads(content)
+        return cls(content, instance)
+
+    @classmethod
+    def from_dict(cls, content):
+        """Create NodeSets instance from a dict."""
+        return cls.from_string(json.dumps(content))
+
+    def __contains__(self, name):
+        """Check if node set exists."""
+        if isinstance(name, str):
+            return name in self._instance.names
+
+        raise BluepySnapError(f"Unexpected type: '{type(name).__name__}' (expected: 'str')")
+
+    def __getitem__(self, name):
+        """Return a node set instance for the given node set name."""
+        if name not in self:
+            raise BluepySnapError(f"Undefined node set: '{name}'")
+        return NodeSet(self._instance, name)
 
     def __iter__(self):
         """Iter through the different node sets names."""
-        return iter(self.resolved)
+        return iter(self._instance.names)

bluepysnap/nodes/node_population.py

@@ -70,6 +70,7 @@
 from bluepysnap import query, utils
 from bluepysnap.circuit_ids import CircuitNodeId, CircuitNodeIds
 from bluepysnap.exceptions import BluepySnapError
+from bluepysnap.node_sets import NodeSet
 from bluepysnap.sonata_constants import DYNAMICS_PREFIX, ConstContainer, Node
 
 
@@ -347,18 +348,15 @@ def _check_properties(self, properties):
         if unknown_props:
             raise BluepySnapError(f"Unknown node properties: {sorted(unknown_props)}")
 
-    def _get_node_set(self, node_set_name):
-        """Returns the node set named 'node_set_name'."""
-        if node_set_name not in self._node_sets:
-            raise BluepySnapError(f"Undefined node set: '{node_set_name}'")
-        return self._node_sets[node_set_name]
-
-    def _resolve_nodesets(self, queries):
+    def _resolve_nodesets(self, queries, raise_missing_prop):
         def _resolve(queries, queries_key):
             if queries_key == query.NODE_SET_KEY:
                 if query.AND_KEY not in queries:
                     queries[query.AND_KEY] = []
-                queries[query.AND_KEY].append(self._get_node_set(queries[queries_key]))
+                node_set = self._node_sets[queries[queries_key]]
+                queries[query.AND_KEY].append(
+                    {query.NODE_ID_KEY: node_set.get_ids(self._population, raise_missing_prop)}
+                )
                 del queries[queries_key]
 
         resolved_queries = deepcopy(queries)
@@ -382,7 +380,7 @@ def _node_ids_by_filter(self, queries, raise_missing_prop):
             >>>                              { Node.X: (0, 1), Node.MTYPE: 'L1_SLAC' }]})
 
         """
-        queries = self._resolve_nodesets(queries)
+        queries = self._resolve_nodesets(queries, raise_missing_prop)
         properties = query.get_properties(queries)
         if raise_missing_prop:
             self._check_properties(properties)
@@ -414,16 +412,16 @@ def ids(self, group=None, limit=None, sample=None, raise_missing_property=True):
         # pylint: disable=too-many-branches
         preserve_order = False
         if isinstance(group, str):
-            group = self._get_node_set(group)
+            group = self._node_sets[group]
         elif isinstance(group, CircuitNodeIds):
             group = group.filter_population(self.name).get_ids()
 
         if group is None:
             result = np.arange(self.size)
+        elif isinstance(group, NodeSet):
+            result = group.get_ids(self._population, raise_missing_property)
         elif isinstance(group, Mapping):
-            result = self._node_ids_by_filter(
-                queries=group, raise_missing_prop=raise_missing_property
-            )
+            result = self._node_ids_by_filter(group, raise_missing_property)
         elif isinstance(group, np.ndarray):
             result = group
             self._check_ids(result)

bluepysnap/simulation.py

@@ -125,8 +125,8 @@ def simulator(self):
     @cached_property
     def node_sets(self):
         """Returns the NodeSets object bound to the simulation."""
-        node_sets_file = self.to_libsonata.node_sets_file
-        return NodeSets(node_sets_file) if node_sets_file else {}
+        path = self.to_libsonata.node_sets_file
+        return NodeSets.from_file(path) if path else NodeSets.from_dict({})
 
     @cached_property
     def spikes(self):

setup.py

@@ -46,7 +46,7 @@ def __init__(self, *args, **kwargs):
         "cached_property>=1.0",
         "h5py>=3.0.1,<4.0.0",
         "jsonschema>=4.0.0,<5.0.0",
-        "libsonata>=0.1.20,<1.0.0",
+        "libsonata>=0.1.21,<1.0.0",
         "morphio>=3.0.0,<4.0.0",
         "morph-tool>=2.4.3,<3.0.0",
         "numpy>=1.8",

tests/data/node_sets_extra.json

@@ -0,0 +1,5 @@
+{
+  "ExtraLayer2": {
+    "layer": 2
+  }
+}

tests/data/node_sets_file.json

@@ -11,5 +11,8 @@
     "population": "default",
     "mtype": "L6_Y"
   },
-  "combined": ["Node2_L6_Y", "Layer23"]
-}
+  "combined": ["Node2_L6_Y", "Layer23"],
+  "failing": {
+      "unknown_property": [0]
+  }
+}

tests/test_circuit.py

@@ -69,7 +69,7 @@ def test_no_node_set():
         with edit_config(config_path) as config:
             config.pop("node_sets_file")
         circuit = test_module.Circuit(config_path)
-        assert circuit.node_sets == {}
+        assert circuit.node_sets.content == {}
 
 
 def test_integration():

tests/test_frame_report.py

@@ -64,8 +64,8 @@ def test_sim(self):
         assert isinstance(self.test_obj_info.simulation, Simulation)
 
     def test_node_set(self):
-        assert self.test_obj.node_set == {"layer": [2, 3]}
-        assert self.test_obj_info.node_set == {"layer": [2, 3]}
+        assert self.test_obj.node_set == "Layer23"
+        assert self.test_obj_info.node_set == "Layer23"
 
     def test_population_names(self):
         assert sorted(self.test_obj.population_names) == ["default", "default2"]