Fix cell position bug with drift

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "MEArec"
-version = "1.10.0"
+version = "1.11.0"
 authors = [
   { name="Alessio Buccino", email="alessiop.buccino@gmail.com" },
 ]

src/MEArec/tests/test_generators.py

@@ -659,6 +659,64 @@ def test_recording_custom_drifts(self):
             del recgen_drift
             recgen_mixed_file.unlink()
 
+    def test_min_dist(self):
+        print("Test recording generation - min_dist between cells")
+        ne = 5
+        ni = 3
+        n_neurons = ne + ni
+        min_dist = 30
+
+        rec_params = mr.get_default_recordings_params()
+
+        rec_params["spiketrains"]["n_exc"] = ne
+        rec_params["spiketrains"]["n_inh"] = ni
+        rec_params["spiketrains"]["duration"] = 2
+        rec_params["recordings"]["modulation"] = "none"
+        rec_params["recordings"]["filter"] = False
+        rec_params["templates"]["min_dist"] = min_dist
+
+        recgen = mr.gen_recordings(params=rec_params, tempgen=self.tempgen, verbose=False)
+
+        locs = np.asarray(recgen.template_locations)
+        # non-drifting: locations have shape (n_neurons, 3)
+        assert locs.shape == (n_neurons, 3)
+        for i in range(n_neurons):
+            for j in range(i + 1, n_neurons):
+                dist = np.linalg.norm(locs[i] - locs[j])
+                assert dist >= min_dist, f"Cells {i} and {j} are {dist} um apart (min_dist={min_dist})"
+        del recgen
+
+    def test_min_dist_drift(self):
+        print("Test recording generation - min_dist between cells with drift")
+        ne = 5
+        ni = 3
+        n_neurons = ne + ni
+        min_dist = 30
+
+        rec_params = mr.get_default_recordings_params()
+
+        rec_params["spiketrains"]["n_exc"] = ne
+        rec_params["spiketrains"]["n_inh"] = ni
+        rec_params["spiketrains"]["duration"] = 2
+        rec_params["recordings"]["modulation"] = "none"
+        rec_params["recordings"]["filter"] = False
+        rec_params["recordings"]["drifting"] = True
+        rec_params["templates"]["min_dist"] = min_dist
+
+        recgen = mr.gen_recordings(params=rec_params, tempgen=self.tempgen_drift, verbose=False)
+
+        locs = np.asarray(recgen.template_locations)
+        # drifting: locations have shape (n_neurons, n_drift_steps, 3)
+        assert locs.ndim == 3 and locs.shape[0] == n_neurons and locs.shape[-1] == 3
+        # the min_dist constraint must be satisfied at every drift step
+        for i in range(n_neurons):
+            for j in range(i + 1, n_neurons):
+                dist = np.linalg.norm(locs[i] - locs[j], axis=-1)
+                assert np.all(dist >= min_dist), (
+                    f"Cells {i} and {j} are {np.min(dist)} um apart at some drift step (min_dist={min_dist})"
+                )
+        del recgen
+
     def test_save_load_templates(self):
         tempgen = mr.load_templates(self.test_dir / "templates.h5", verbose=True)
         tempgen_drift = mr.load_templates(self.test_dir / "templates_drift.h5")

src/MEArec/tools.py

@@ -1167,7 +1167,7 @@ def select_templates(
             # excitatory cell
             if bcat == "E":
                 if n_sel_exc < n_exc:
-                    dist = np.array([np.linalg.norm(loc[id_cell] - p) for p in pos_sel])
+                    dist = np.array([np.linalg.norm(loc[id_cell] - p, axis=-1) for p in pos_sel])
                     if np.any(dist < min_dist):
                         if verbose:
                             print("Distance violation", np.min(dist), iter)
@@ -1300,7 +1300,7 @@ def select_templates(
             # inhibitory cell
             elif bcat == "I":
                 if n_sel_inh < n_inh:
-                    dist = np.array([np.linalg.norm(loc[id_cell] - p) for p in pos_sel])
+                    dist = np.array([np.linalg.norm(loc[id_cell] - p, axis=-1) for p in pos_sel])
                     if np.any(dist < min_dist):
                         if verbose:
                             print("Distance violation", np.min(dist), iter)
@@ -1373,6 +1373,7 @@ def select_templates(
                                 drift_angle = np.rad2deg(np.arccos(np.dot(drift_dir[id_cell], preferred_dir)))
                                 if drift_angle - angle_tol <= 0:
                                     if n_overlap_pairs is None:
+                                        pos_sel.append(loc[id_cell])
                                         selected_idxs.append(id_cell)
                                         n_sel += 1
                                         placed = True
@@ -1431,7 +1432,7 @@ def select_templates(
                         selected_cat.append("I")
         # unknown cell type
         else:
-            dist = np.array([np.linalg.norm(loc[id_cell] - p) for p in pos_sel])
+            dist = np.array([np.linalg.norm(loc[id_cell] - p, axis=-1) for p in pos_sel])
             if np.any(dist < min_dist):
                 if verbose:
                     print("Distance violation", np.min(dist), iter)