give access to spatial (#177)

* node_populations now have a "spatial_segment_index" property which gives access to the segment index
* edge_populations now have a "spatial_synapse_index" property which gives access to the synapse index

.gitignore

@@ -1,4 +1,5 @@
 doc/build
+build
 doc/source/submodules
 coverage.xml
 .coverage

CHANGELOG.rst

@@ -14,6 +14,7 @@ Version v1.0.1
 New Features
 ~~~~~~~~~~~~
 - Access the population configs for node/edge populations via population_config property
+- Spatial index for segments and synapses
 
 Version v1.0.0
 --------------

bluepysnap/edges/edge_population.py

@@ -576,3 +576,21 @@ def iter_connections(
     def h5_filepath(self):
         """Get the H5 edges file associated with population."""
         return self.config["edges_file"]
+
+    @cached_property
+    def spatial_synapse_index(self):
+        """Access to edges spatial index."""
+        try:
+            from spatial_index import open_index
+        except ImportError as e:
+            raise BluepySnapError(
+                (
+                    "Spatial index is for now only available internally to BBP. "
+                    "It requires `spatial_index`, an internal package."
+                )
+            ) from e
+
+        properties = self._circuit.to_libsonata.edge_population_properties(self.name)
+        if not properties.spatial_synapse_index_dir:
+            raise BluepySnapError(f"It appears {self.name} does not have synapse indices")
+        return open_index(properties.spatial_synapse_index_dir)

bluepysnap/nodes/node_population.py

@@ -580,3 +580,23 @@ def models(self):
     def h5_filepath(self):
         """Get the H5 nodes file associated with population."""
         return self.config["nodes_file"]
+
+    @cached_property
+    def spatial_segment_index(self):
+        """Access to edges spatial index."""
+        try:
+            from spatial_index import open_index
+        except ImportError as e:
+            raise BluepySnapError(
+                (
+                    "Spatial index is for now only available internally to BBP. ",
+                    "It requires `spatial_index`, an internal package.",
+                )
+            ) from e
+
+        properties = self._circuit.to_libsonata.node_population_properties(self.name)
+
+        if not properties.spatial_segment_index_dir:
+            raise BluepySnapError(f"It appears {self.name} does not have segment indices")
+
+        return open_index(properties.spatial_segment_index_dir)

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.16,<1.0.0",
+        "libsonata>=0.1.17,<1.0.0",
         "morphio>=3.0.0,<4.0.0",
         "morph-tool>=2.4.3,<3.0.0",
         "neurom>=1.6",
@@ -59,7 +59,7 @@ def __init__(self, *args, **kwargs):
     extras_require={
         "docs": ["sphinx", "sphinx-bluebrain-theme"],
         "plots": ["matplotlib>=3.0.0"],
-        "tools": ["bglibpy>=4.4.26"],
+        "spatial-index": ["spatial-index>=1.2.1,<2.0"],
     },
     packages=find_packages(),
     package_data={

tests/data/circuit_config.json

@@ -16,9 +16,10 @@
         "populations": {
             "default": {
                 "type": "biophysical"
-            },
+              },
             "default2": {
-                "type": "biophysical"
+                "type": "biophysical",
+                "spatial_segment_index_dir": "path/to/node/dir"
             }
         }
       }
@@ -29,9 +30,10 @@
         "populations": {
             "default":{
                 "type": "chemical"
-            },
+              },
             "default2":{
-                "type": "chemical"
+                "type": "chemical",
+                "spatial_synapse_index_dir": "path/to/edge/dir"
             }
         }
       }

tests/test_circuit.py

@@ -15,7 +15,10 @@ def test_all():
     circuit = test_module.Circuit(str(TEST_DATA_DIR / "circuit_config.json"))
     assert circuit.config["networks"]["nodes"][0] == {
         "nodes_file": str(TEST_DATA_DIR / "nodes.h5"),
-        "populations": {"default": {"type": "biophysical"}, "default2": {"type": "biophysical"}},
+        "populations": {
+            "default": {"type": "biophysical"},
+            "default2": {"type": "biophysical", "spatial_segment_index_dir": "path/to/node/dir"},
+        },
     }
     assert isinstance(circuit.nodes, Nodes)
     assert isinstance(circuit.edges, Edges)

tests/test_edges/test_edge_population.py

@@ -1,4 +1,6 @@
+import sys
 from pathlib import Path
+from unittest import mock
 
 import numpy as np
 import numpy.testing as npt
@@ -613,3 +615,41 @@ def test_iter_connection_unique(self):
 
     def test_h5_filepath_from_config(self):
         assert self.test_obj.h5_filepath == str(TEST_DATA_DIR / "edges.h5")
+
+    @pytest.mark.skip(reason="Until spatial-index is released publicly")
+    def test_spatial_synapse_index(self):
+        with mock.patch("spatial_index.open_index") as mock_open_index:
+            self.test_obj.spatial_synapse_index
+        mock_open_index.assert_called_once_with("path/to/edge/dir")
+
+    @mock.patch.dict(sys.modules, {"spatial_index": mock.Mock()})
+    def test_spatial_synapse_index_call(self):
+        with pytest.raises(
+            BluepySnapError,
+            match="It appears default does not have synapse indices",
+        ):
+            self.test_obj.spatial_synapse_index
+
+    def test_spatial_synapse_index_error(self):
+        with pytest.raises(
+            BluepySnapError,
+            match=(
+                "Spatial index is for now only available internally to BBP. "
+                "It requires `spatial_index`, an internal package."
+            ),
+        ):
+            self.test_obj.spatial_synapse_index
+
+
+class TestEdgePopulationSpatialIndex:
+    def setup_method(self):
+        self.test_obj = TestEdgePopulation.get_edge_population(
+            str(TEST_DATA_DIR / "circuit_config.json"), "default2"
+        )
+
+    @mock.patch.dict(sys.modules, {"spatial_index": mock.Mock()})
+    def test_spatial_synapse_index_call(self):
+        self.test_obj.spatial_synapse_index
+        mock = sys.modules["spatial_index"].open_index
+        assert mock.call_count == 1
+        assert mock.call_args[0][0].endswith("path/to/edge/dir")

tests/test_nodes/test_node_population.py

@@ -1,4 +1,6 @@
 import json
+import sys
+from unittest import mock
 
 import numpy as np
 import numpy.testing as npt
@@ -596,3 +598,36 @@ def test_models(self):
 
     def test_h5_filepath_from_config(self):
         assert self.test_obj.h5_filepath == str(TEST_DATA_DIR / "nodes.h5")
+
+    @pytest.mark.skip(reason="Until spatial-index is released publicly")
+    def test_spatial_segment_index(self):
+        with mock.patch("spatial_index.open_index") as mock_open_index:
+            self.test_obj.spatial_segment_index
+        mock_open_index.assert_called_once_with("path/to/node/dir")
+
+    @mock.patch.dict(sys.modules, {"spatial_index": mock.Mock()})
+    def test_spatial_segment_index_call(self):
+        with pytest.raises(
+            BluepySnapError,
+            match="It appears default does not have segment indices",
+        ):
+            self.test_obj.spatial_segment_index
+
+    def test_spatial_segment_index_error(self):
+        with pytest.raises(
+            BluepySnapError, match="Spatial index is for now only available internally to BBP."
+        ):
+            self.test_obj.spatial_segment_index
+
+
+class TestNodePopulationSpatialIndex:
+    def setup_method(self):
+        self.test_obj = Circuit(str(TEST_DATA_DIR / "circuit_config.json")).nodes["default2"]
+
+    @mock.patch.dict(sys.modules, {"spatial_index": mock.Mock()})
+    def test_spatial_segment_index_call(self):
+        self.test_obj.spatial_segment_index
+
+        mock = sys.modules["spatial_index"].open_index
+        assert mock.call_count == 1
+        assert mock.call_args[0][0].endswith("path/to/node/dir")