From 45e0c4eb0e8ecb3728f5ef3ba17a7c012327ab41 Mon Sep 17 00:00:00 2001
From: BadPrograms <johnliaskas@gmail.com>
Date: Thu, 18 Dec 2025 16:12:52 +0100
Subject: [PATCH 1/2] added compoute neighbours

---
 src/lineagetree/_mixins/spatial_mixin.py |  2 +
 src/lineagetree/measure/spatial.py       | 51 +++++++++++++++++++-----
 2 files changed, 42 insertions(+), 11 deletions(-)

diff --git a/src/lineagetree/_mixins/spatial_mixin.py b/src/lineagetree/_mixins/spatial_mixin.py
index e22edeb..e2d581c 100644
--- a/src/lineagetree/_mixins/spatial_mixin.py
+++ b/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
diff --git a/src/lineagetree/measure/spatial.py b/src/lineagetree/measure/spatial.py
index 417ac05..0b78879 100644
--- a/src/lineagetree/measure/spatial.py
+++ b/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,9 +136,7 @@ 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:
@@ -146,11 +144,42 @@ def compute_spatial_density(
     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 / 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(

From a82cdb19f81360bf88537f162d363df5e39aaaea Mon Sep 17 00:00:00 2001
From: BadPrograms <johnliaskas@gmail.com>
Date: Thu, 18 Dec 2025 16:23:06 +0100
Subject: [PATCH 2/2] fix tests and density

---
 src/lineagetree/measure/spatial.py | 2 +-
 tests/test_lineageTree.py          | 6 ++++--
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/src/lineagetree/measure/spatial.py b/src/lineagetree/measure/spatial.py
index 0b78879..3e9e352 100644
--- a/src/lineagetree/measure/spatial.py
+++ b/src/lineagetree/measure/spatial.py
@@ -176,7 +176,7 @@ def compute_spatial_density(
     """
     s_vol = 4 / 3.0 * np.pi * th**3
     spatial_density = {
-        k: v / s_vol
+        k: (v + 1) / s_vol
         for k, v in lT.compute_neighbours_in_radius(t_b, t_e, th).items()
     }
     return spatial_density
diff --git a/tests/test_lineageTree.py b/tests/test_lineageTree.py
index bd0d54e..df5a9f2 100644
--- a/tests/test_lineageTree.py
+++ b/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():