[NSETM-2226] Add NodeSets adapter to support custom node_sets_file (#22)

Support custom node_sets_file in extraction, neuron_classes and trial_steps

CHANGELOG.rst

@@ -1,6 +1,15 @@
 Changelog
 =========
 
+Version 0.8.0
+-------------
+
+New Features
+~~~~~~~~~~~~
+
+- Support custom node_sets_file in extraction, neuron_classes and trial_steps [NSETM-2226]
+
+
 Version 0.7.1
 -------------
 

pyproject.toml

@@ -123,7 +123,7 @@ branch = true
 parallel = false
 omit = [
     "*/blueetl/_version.py",
-    "*/blueetl/adapters/bluepy/*.py",
+    "*/blueetl/adapters/impl/bluepy/*.py",
     "*/blueetl/external/**/*.py",
 ]
 

src/blueetl/adapters/base.py

@@ -1,13 +1,14 @@
 """Base Adapter."""
 
 from abc import ABC, abstractmethod
+from pathlib import Path
 from typing import Any, Generic, Optional, TypeVar
 
 from blueetl.adapters.interfaces.circuit import CircuitInterface
+from blueetl.adapters.interfaces.node_sets import NodeSetsInterface
 from blueetl.adapters.interfaces.simulation import SimulationInterface
-from blueetl.types import StrOrPath
 
-InterfaceT = TypeVar("InterfaceT", CircuitInterface, SimulationInterface)
+InterfaceT = TypeVar("InterfaceT", CircuitInterface, SimulationInterface, NodeSetsInterface)
 BaseAdapterT = TypeVar("BaseAdapterT", bound="BaseAdapter")
 
 
@@ -18,13 +19,14 @@ class AdapterError(Exception):
 class BaseAdapter(Generic[InterfaceT], ABC):
     """Base Adapter to be subclassed by other adapters."""
 
-    def __init__(self, config: StrOrPath) -> None:
-        """Init the adapter from the specified config."""
-        self._impl: Optional[InterfaceT] = self._load_impl(str(config))
+    def __init__(self, _impl: Optional[InterfaceT]) -> None:
+        """Init the adapter from the specified implementation."""
+        self._impl: Optional[InterfaceT] = _impl
 
+    @classmethod
     @abstractmethod
-    def _load_impl(self, config: str) -> Optional[InterfaceT]:
-        """Load and return the implementation object, or None if the config file doesn't exist."""
+    def from_file(cls, filepath: Optional[Path]) -> "BaseAdapter":
+        """Load and return a new object from file."""
 
     @property
     def _ensure_impl(self) -> InterfaceT:
@@ -41,10 +43,3 @@ def exists(self) -> bool:
     def instance(self) -> Any:
         """Return the wrapped instance, or None if it doesn't exist."""
         return self._impl.instance if self._impl is not None else None
-
-    @classmethod
-    def from_impl(cls: type[BaseAdapterT], impl: InterfaceT) -> BaseAdapterT:
-        """Return a new adapter with the specified implementation."""
-        result = object.__new__(cls)
-        result._impl = impl
-        return result

src/blueetl/adapters/circuit.py

@@ -6,22 +6,25 @@
 
 from blueetl.adapters.base import BaseAdapter
 from blueetl.adapters.interfaces.circuit import CircuitInterface, NodePopulationInterface
+from blueetl.adapters.node_sets import NodeSetsAdapter
 
 
 class CircuitAdapter(BaseAdapter[CircuitInterface]):
     """Circuit Adapter."""
 
-    def _load_impl(self, config: str) -> Optional[CircuitInterface]:
-        """Load and return the implementation object, or None if the config file doesn't exist."""
+    @classmethod
+    def from_file(cls, filepath: Optional[Path]) -> "CircuitAdapter":
+        """Load and return a new object from file."""
         # pylint: disable=import-outside-toplevel
-        if not config or not Path(config).exists():
-            return None
+        if not filepath or not filepath.exists():
+            return cls(None)
         CircuitImpl: type[CircuitInterface]
-        if config.endswith(".json"):
-            from blueetl.adapters.bluepysnap.circuit import Circuit, CircuitImpl
+        if filepath.suffix == ".json":
+            from blueetl.adapters.impl.bluepysnap.circuit import Circuit, CircuitImpl
         else:
-            from blueetl.adapters.bluepy.circuit import Circuit, CircuitImpl
-        return CircuitImpl(Circuit(config))
+            from blueetl.adapters.impl.bluepy.circuit import Circuit, CircuitImpl
+        impl = CircuitImpl(Circuit(str(filepath)))
+        return cls(impl)
 
     def checksum(self) -> str:
         """Return a checksum of the relevant keys in the circuit configuration."""
@@ -31,3 +34,13 @@ def checksum(self) -> str:
     def nodes(self) -> Mapping[Optional[str], NodePopulationInterface]:
         """Return the nodes as a dict: population -> nodes."""
         return self._ensure_impl.nodes
+
+    @property
+    def node_sets_file(self) -> Optional[Path]:
+        """Returns the NodeSets file used by the circuit."""
+        return self._ensure_impl.node_sets_file
+
+    @property
+    def node_sets(self) -> NodeSetsAdapter:
+        """Returns the NodeSets file used by the circuit."""
+        return NodeSetsAdapter(self._ensure_impl.node_sets)

src/blueetl/adapters/impl/__init__.py

@@ -0,0 +1 @@
+"""Implementations."""

src/blueetl/adapters/bluepy/circuit.py → src/blueetl/adapters/impl/bluepy/circuit.py

@@ -1,11 +1,13 @@
 """Bluepy circuit implementation."""
 
 from collections.abc import Mapping
+from pathlib import Path
 from typing import Optional
 
 from bluepy import Circuit
 
 from blueetl.adapters.interfaces.circuit import CircuitInterface, NodePopulationInterface
+from blueetl.adapters.interfaces.node_sets import NodeSetsInterface
 from blueetl.utils import checksum_json
 
 
@@ -40,3 +42,13 @@ def nodes(self) -> Mapping[Optional[str], NodePopulationInterface]:
         The population name in the returned dict is always None, because undefined in the config.
         """
         return {None: self._circuit.cells}
+
+    @property
+    def node_sets_file(self) -> Optional[Path]:
+        """Returns the NodeSets file used by the circuit."""
+        raise NotImplementedError
+
+    @property
+    def node_sets(self) -> NodeSetsInterface:
+        """Returns the NodeSets used by the circuit."""
+        raise NotImplementedError

src/blueetl/adapters/bluepy/simulation.py → src/blueetl/adapters/impl/bluepy/simulation.py

@@ -12,7 +12,7 @@
 from bluepy.impl.spike_report import SpikeReport
 from bluepy.simulation import PathHelpers
 
-from blueetl.adapters.bluepy.circuit import CircuitImpl
+from blueetl.adapters.impl.bluepy.circuit import CircuitImpl
 from blueetl.adapters.interfaces.circuit import CircuitInterface
 from blueetl.adapters.interfaces.simulation import (
     PopulationReportInterface,

src/blueetl/adapters/bluepysnap/circuit.py → src/blueetl/adapters/impl/bluepysnap/circuit.py

@@ -1,11 +1,14 @@
 """Bluepysnap circuit implementation."""
 
 from collections.abc import Mapping
+from pathlib import Path
 from typing import Optional
 
 from bluepysnap import Circuit
 
+from blueetl.adapters.impl.bluepysnap.node_sets import NodeSetsImpl
 from blueetl.adapters.interfaces.circuit import CircuitInterface, NodePopulationInterface
+from blueetl.adapters.interfaces.node_sets import NodeSetsInterface
 from blueetl.utils import checksum_json
 
 
@@ -23,3 +26,14 @@ def checksum(self) -> str:
     def nodes(self) -> Mapping[Optional[str], NodePopulationInterface]:
         """Return the nodes as a dict: population -> nodes."""
         return self._circuit.nodes
+
+    @property
+    def node_sets_file(self) -> Optional[Path]:
+        """Returns the NodeSets file used by the circuit."""
+        path = self._circuit.to_libsonata.node_sets_path
+        return Path(path) if path else None
+
+    @property
+    def node_sets(self) -> NodeSetsInterface:
+        """Returns the NodeSets used by the circuit."""
+        return NodeSetsImpl(self._circuit.node_sets)

src/blueetl/adapters/impl/bluepysnap/node_sets.py

@@ -0,0 +1,13 @@
+"""Bluepysnap NodeSets implementation."""
+
+from bluepysnap.node_sets import NodeSets
+
+from blueetl.adapters.interfaces.node_sets import NodeSetsInterface
+
+
+class NodeSetsImpl(NodeSetsInterface[NodeSets]):
+    """Bluepysnap NodeSets implementation."""
+
+    def update(self, other: NodeSetsInterface) -> None:
+        """Update the wrapped node sets."""
+        self._node_sets.update(other.instance)

src/blueetl/adapters/bluepysnap/simulation.py → src/blueetl/adapters/impl/bluepysnap/simulation.py

@@ -7,7 +7,7 @@
 
 from bluepysnap import Simulation
 
-from blueetl.adapters.bluepysnap.circuit import CircuitImpl
+from blueetl.adapters.impl.bluepysnap.circuit import CircuitImpl
 from blueetl.adapters.interfaces.circuit import CircuitInterface
 from blueetl.adapters.interfaces.simulation import (
     PopulationReportInterface,

src/blueetl/adapters/interfaces/__init__.py

@@ -1 +1 @@
-"""Interfaces for Simulation and Circuit."""
+"""Interfaces."""

src/blueetl/adapters/interfaces/circuit.py

@@ -2,11 +2,13 @@
 
 from abc import ABC, abstractmethod
 from collections.abc import Mapping
+from pathlib import Path
 from typing import Generic, Optional, TypeVar
 
 import numpy as np
 import pandas as pd
 
+from blueetl.adapters.interfaces.node_sets import NodeSetsInterface
 from blueetl.utils import CachedPropertyMixIn
 
 CircuitT = TypeVar("CircuitT")
@@ -44,3 +46,13 @@ def checksum(self) -> str:
     @abstractmethod
     def nodes(self) -> Mapping[Optional[str], NodePopulationInterface]:
         """Return the nodes as a dict: population -> nodes."""
+
+    @property
+    @abstractmethod
+    def node_sets_file(self) -> Optional[Path]:
+        """Returns the NodeSets file used by the circuit."""
+
+    @property
+    @abstractmethod
+    def node_sets(self) -> NodeSetsInterface:
+        """Returns the NodeSets used by the circuit."""

src/blueetl/adapters/interfaces/node_sets.py

@@ -0,0 +1,23 @@
+"""Interfaces for NodeSets."""
+
+from abc import ABC, abstractmethod
+from typing import Generic, TypeVar
+
+NodeSetsT = TypeVar("NodeSetsT")
+
+
+class NodeSetsInterface(Generic[NodeSetsT], ABC):
+    """NodeSets Interface."""
+
+    def __init__(self, node_sets: NodeSetsT) -> None:
+        """Init the NodeSets interface with the specified NodeSetsT."""
+        self._node_sets = node_sets
+
+    @property
+    def instance(self) -> NodeSetsT:
+        """Return the wrapped instance."""
+        return self._node_sets
+
+    @abstractmethod
+    def update(self, other: "NodeSetsInterface") -> None:
+        """Update the wrapped node sets."""

src/blueetl/adapters/node_sets.py

@@ -0,0 +1,33 @@
+"""Adapters for NodeSets."""
+
+from pathlib import Path
+from typing import Optional
+
+from blueetl.adapters.base import BaseAdapter
+from blueetl.adapters.interfaces.node_sets import NodeSetsInterface
+
+
+class NodeSetsAdapter(BaseAdapter[NodeSetsInterface]):
+    """NodeSets Adapter."""
+
+    @classmethod
+    def from_file(cls, filepath: Optional[Path]) -> "NodeSetsAdapter":
+        """Load and return a new object from file."""
+        # pylint: disable=import-outside-toplevel
+        from blueetl.adapters.impl.bluepysnap.node_sets import NodeSets, NodeSetsImpl
+
+        if not filepath:
+            impl = NodeSetsImpl(NodeSets.from_dict({}))
+        else:
+            impl = NodeSetsImpl(NodeSets.from_file(str(filepath)))
+        return cls(impl)
+
+    def update(self, other: "NodeSetsAdapter") -> None:
+        """Update the node sets."""
+        # pylint: disable=protected-access
+        self._ensure_impl.update(other._ensure_impl)
+
+    def __ior__(self, other: "NodeSetsAdapter") -> "NodeSetsAdapter":
+        """Support ``A |= B``."""
+        self.update(other)
+        return self

src/blueetl/adapters/simulation.py

@@ -16,17 +16,19 @@
 class SimulationAdapter(BaseAdapter[SimulationInterface]):
     """Simulation Adapter."""
 
-    def _load_impl(self, config: str) -> Optional[SimulationInterface]:
-        """Load and return the implementation object, or None if the config file doesn't exist."""
+    @classmethod
+    def from_file(cls, filepath: Optional[Path]) -> "SimulationAdapter":
+        """Load and return a new object from file."""
         # pylint: disable=import-outside-toplevel
-        if not config or not Path(config).exists():
-            return None
+        if not filepath or not filepath.exists():
+            return cls(None)
         SimulationImpl: type[SimulationInterface]
-        if config.endswith(".json"):
-            from blueetl.adapters.bluepysnap.simulation import Simulation, SimulationImpl
+        if filepath.suffix == ".json":
+            from blueetl.adapters.impl.bluepysnap.simulation import Simulation, SimulationImpl
         else:
-            from blueetl.adapters.bluepy.simulation import Simulation, SimulationImpl
-        return SimulationImpl(Simulation(config))
+            from blueetl.adapters.impl.bluepy.simulation import Simulation, SimulationImpl
+        impl = SimulationImpl(Simulation(str(filepath)))
+        return cls(impl)
 
     def is_complete(self) -> bool:
         """Return True if the simulation is complete, False otherwise."""
@@ -35,7 +37,7 @@ def is_complete(self) -> bool:
     @property
     def circuit(self) -> CircuitAdapter:
         """Return the circuit used for the simulation."""
-        return CircuitAdapter.from_impl(self._ensure_impl.circuit)
+        return CircuitAdapter(self._ensure_impl.circuit)
 
     @property
     def spikes(self) -> Mapping[Optional[str], PopulationSpikesReportInterface]:

src/blueetl/config/analysis.py

@@ -14,27 +14,41 @@
     WindowConfig,
 )
 from blueetl.constants import CHECKSUM_SEP
-from blueetl.utils import checksum_json, dict_product
+from blueetl.utils import checksum_json, dict_product, load_json
 from blueetl.validation import read_schema, validate_config
 
 L = logging.getLogger(__name__)
 
 
 def _resolve_paths(global_config: MultiAnalysisConfig, base_path: Path) -> None:
-    """Resolve any relative path."""
+    """Resolve any relative path at the top level."""
     base_path = base_path or Path()
     global_config.output = base_path / global_config.output
     global_config.simulation_campaign = base_path / global_config.simulation_campaign
 
 
-def _resolve_trial_steps(global_config: MultiAnalysisConfig):
+def _resolve_neuron_classes(global_config: MultiAnalysisConfig, base_path: Path):
+    """Resolve the relative paths in neuron_classes."""
+    for config in global_config.analysis.values():
+        for neuron_classes_config in config.extraction.neuron_classes.values():
+            if path := neuron_classes_config.node_sets_file:
+                path = base_path / path
+                neuron_classes_config.node_sets_file = path
+                neuron_classes_config.node_sets_checksum = checksum_json(load_json(path))
+
+
+def _resolve_trial_steps(global_config: MultiAnalysisConfig, base_path: Path):
     """Set trial_steps_config.base_path to the same value as global_config.output.
 
     In this way, the custom function can use it as the base path to save any figure.
     """
     for config in global_config.analysis.values():
         for trial_steps_config in config.extraction.trial_steps.values():
             trial_steps_config.base_path = str(global_config.output)
+            if path := trial_steps_config.node_sets_file:
+                path = base_path / path
+                trial_steps_config.node_sets_file = path
+                trial_steps_config.node_sets_checksum = checksum_json(load_json(path))
 
 
 def _resolve_windows(global_config: MultiAnalysisConfig) -> None:
@@ -155,7 +169,8 @@ def init_multi_analysis_configuration(global_config: dict, base_path: Path) -> M
     validate_config(global_config, schema=read_schema("analysis_config"))
     config = MultiAnalysisConfig(**global_config)
     _resolve_paths(config, base_path=base_path)
-    _resolve_trial_steps(config)
+    _resolve_neuron_classes(config, base_path=base_path)
+    _resolve_trial_steps(config, base_path=base_path)
     _resolve_windows(config)
     _resolve_analysis_configs(config)
     return config