added compute_neighbours_in_radius

src/lineagetree/_mixins/spatial_mixin.py

@@ -4,6 +4,7 @@
     compute_spatial_edges,
     get_gabriel_graph,
     get_idx3d,
+    compute_neighbours_in_radius,
 )
 
 from ._methodize import AutoMethodizeMeta
@@ -17,3 +18,4 @@ class SpatialMixin(metaclass=AutoMethodizeMeta):
     compute_k_nearest_neighbours = compute_k_nearest_neighbours
     compute_spatial_edges = compute_spatial_edges
     compute_spatial_density = compute_spatial_density
+    compute_neighbours_in_radius = compute_neighbours_in_radius

src/lineagetree/measure/spatial.py

@@ -111,14 +111,14 @@ def get_gabriel_graph(lT: LineageTree, t: int) -> dict[int, set[int]]:
     return lT.Gabriel_graph[t]
 
 
-def compute_spatial_density(
+def compute_neighbours_in_radius(
     lT: LineageTree,
     t_b: int | None = None,
     t_e: int | None = None,
     th: float = 50,
 ) -> dict[int, float]:
-    """Computes the spatial density of nodes between `t_b` and `t_e`.
-    The results is stored in `lT.spatial_density` and returned.
+    """Computes the number of neighbours for nodes between `t_b` and `t_e`.
+    The results is stored in `lT.neighbours` and returned.
 
     Parameters
     ----------
@@ -136,21 +136,50 @@ def compute_spatial_density(
     dict mapping int to float
         dictionary that maps a node id to its spatial density
     """
-    if not hasattr(lT, "spatial_density"):
-        lT.spatial_density = {}
-    s_vol = 4 / 3.0 * np.pi * th**3
+    neighbours = {}
     if t_b is None:
         t_b = lT.t_b
     if t_e is None:
         t_e = lT.t_e
     time_range = set(range(t_b, t_e)).intersection(lT._time.values())
     for t in time_range:
         idx3d, nodes = lT.get_idx3d(t)
-        nb_ni = [
-            (len(ni) - 1) / s_vol for ni in idx3d.query_ball_tree(idx3d, th)
-        ]
-        lT.spatial_density.update(dict(zip(nodes, nb_ni, strict=True)))
-    return lT.spatial_density
+        nb_ni = [(len(ni) - 1) for ni in idx3d.query_ball_tree(idx3d, th)]
+        neighbours.update(dict(zip(nodes, nb_ni, strict=True)))
+    return neighbours
+
+
+def compute_spatial_density(
+    lT: LineageTree,
+    t_b: int | None = None,
+    t_e: int | None = None,
+    th: float = 50,
+) -> dict[int, float]:
+    """Computes the spatial density of nodes between `t_b` and `t_e`.
+    The results is stored in `lT.spatial_density` and returned.
+
+    Parameters
+    ----------
+    lT : LineageTree
+        The LineageTree instance.
+    t_b : int, optional
+        starting time to look at, default first time point
+    t_e : int, optional
+        ending time to look at, default last time point
+    th : float, default=50
+        size of the neighbourhood
+
+    Returns
+    -------
+    dict mapping int to float
+        dictionary that maps a node id to its spatial density
+    """
+    s_vol = 4 / 3.0 * np.pi * th**3
+    spatial_density = {
+        k: (v + 1) / s_vol
+        for k, v in lT.compute_neighbours_in_radius(t_b, t_e, th).items()
+    }
+    return spatial_density
 
 
 def compute_k_nearest_neighbours(

tests/test_lineageTree.py

@@ -618,8 +618,10 @@ def test_get_subtree_nodes():
 
 
 def test_spatial_density():
-    density = list(lT1.compute_spatial_density(0, th=40).values())
-    assert np.count_nonzero(density) == 1669
+    assert np.isclose(
+        lT1.compute_spatial_density(0, th=40)[110832], 7.460387957432594e-06
+    )
+    assert lT1.compute_neighbours_in_radius(0, th=40)[110832] == 1
 
 
 def test_compute_k_nearest_neighbours():