From 8c7ca57ce87523997c031d8e56b1fe4236c58f53 Mon Sep 17 00:00:00 2001
From: Kyle Johnsen <johnsenkyle13@gmail.com>
Date: Wed, 8 May 2024 12:04:05 -0400
Subject: [PATCH] patch numpy unit stripping inconsistency (#49)

---
 cleo/coords.py              | 6 +++---
 cleo/ephys/spiking.py       | 5 +++--
 cleo/imaging/scope.py       | 2 +-
 pyproject.toml              | 2 +-
 tests/imaging/test_scope.py | 2 +-
 5 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/cleo/coords.py b/cleo/coords.py
index 7206568..6968647 100644
--- a/cleo/coords.py
+++ b/cleo/coords.py
@@ -196,9 +196,9 @@ def coords_from_xyz(x: Quantity, y: Quantity, z: Quantity) -> Quantity:
     return (
         np.concatenate(
             [
-                np.reshape(x, (*x.shape, 1)),
-                np.reshape(y, (*y.shape, 1)),
-                np.reshape(z, (*z.shape, 1)),
+                np.reshape(x / meter, (*x.shape, 1)),
+                np.reshape(y / meter, (*y.shape, 1)),
+                np.reshape(z / meter, (*z.shape, 1)),
             ],
             axis=-1,
         )
diff --git a/cleo/ephys/spiking.py b/cleo/ephys/spiking.py
index e0c8054..75c4deb 100644
--- a/cleo/ephys/spiking.py
+++ b/cleo/ephys/spiking.py
@@ -94,8 +94,9 @@ def connect_to_neuron_group(
                 getattr(self.probe, f"{dim}s"),
                 indexing="ij",
             )
-            dist2 += (dim_ng - dim_probe) ** 2
-        distances = np.sqrt(dist2) * meter  # since units stripped
+            # proactively strip units to avoid numpy maybe doing so
+            dist2 += (dim_ng / mm - dim_probe / mm) ** 2
+        distances = np.sqrt(dist2) * mm
         # probs is n_neurons by n_channels
         probs = self._detection_prob_for_distance(distances)
         # cut off to get indices of neurons to monitor
diff --git a/cleo/imaging/scope.py b/cleo/imaging/scope.py
index ca4f8ff..c4688c6 100644
--- a/cleo/imaging/scope.py
+++ b/cleo/imaging/scope.py
@@ -329,7 +329,7 @@ def add_self_to_plot(
         coords = (
             np.concatenate(
                 [
-                    coords_from_ng(ng)[i_targets]
+                    coords_from_ng(ng)[i_targets] / meter
                     for ng, i_targets in zip(
                         self.neuron_groups, self.i_targets_per_injct
                     )
diff --git a/pyproject.toml b/pyproject.toml
index 271a08f..0ad8682 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "cleosim"
-version = "0.14.2"
+version = "0.14.3"
 description = "Cleo: the Closed-Loop, Electrophysiology, and Optogenetics experiment simulation testbed"
 authors = [
     "Kyle Johnsen <kyle@kjohnsen.org>",
diff --git a/tests/imaging/test_scope.py b/tests/imaging/test_scope.py
index d7974cc..d69cade 100644
--- a/tests/imaging/test_scope.py
+++ b/tests/imaging/test_scope.py
@@ -102,7 +102,7 @@ def test_target_neurons_in_plane(rand_seed):
     assert np.sum(np.sqrt(ng.x**2 + ng.y**2) < img_width / 2) == len(i_targets)
 
     # noise_focus_factor (nff) smaller for bigger soma
-    ng.z = rng.uniform(focus_depth - 30 * um, focus_depth + 30 * um, ng.N) * meter
+    ng.z = rng.uniform(focus_depth / um - 30, focus_depth / um + 30, ng.N) * um
     i_targets, nff, cop = target_neurons_in_plane(ng, focus_depth, img_width)
     i_targets_bigger, nff_bigger, cop_bigger = target_neurons_in_plane(
         ng, focus_depth, img_width, soma_radius=20 * um