Simplify environment

Author: cwindolf

environment.yml

@@ -9,3 +9,4 @@ dependencies:
   - scikit-learn
   - hdbscan
   - colorcet
+  - pytest

requirements.txt

@@ -1,4 +1 @@
-pytest
-ibl-neuropixel
 spikeinterface
-cloudpickle

tests/test_detect.py

@@ -2,7 +2,7 @@
 import torch
 from dartsort.detect.detect import detect_and_deduplicate
 from dartsort.util.waveform_util import make_channel_index
-from neuropixel import dense_layout
+from test_util import dense_layout
 
 
 def test_detect_and_deduplicate():

tests/test_enforce_decrease.py

@@ -3,7 +3,7 @@
 from dartsort.transform.enforce_decrease import (EnforceDecrease,
                                                  make_parents_index)
 from dartsort.util.waveform_util import make_channel_index
-from neuropixel import dense_layout
+from test_util import dense_layout
 
 
 def test_make_parents_index():
@@ -13,7 +13,5 @@ def test_make_parents_index():
     parents_index = make_parents_index(geom, channel_index)
     assert parents_index.shape[0] == channel_index.shape[0]
 
-    enfdec = EnforceDecrease(
-        torch.as_tensor(channel_index), torch.as_tensor(geom)
-    )
+    enfdec = EnforceDecrease(torch.as_tensor(channel_index), torch.as_tensor(geom))
     assert enfdec is not None

tests/test_localize_torch.py

@@ -2,7 +2,7 @@
 import torch
 from dartsort.localize import localize_torch
 from dartsort.util.waveform_util import make_channel_index
-from neuropixel import dense_layout
+from test_util import dense_layout
 
 
 def test_localize_torch():

tests/test_relocate.py

@@ -5,7 +5,7 @@
 import numpy as np
 from dartsort.cluster import relocate
 from dartsort.util import drift_util, waveform_util
-from neuropixel import dense_layout
+from test_util import dense_layout
 
 h = dense_layout()
 geom = np.c_[h["x"], h["y"]]

tests/test_subtract.py

@@ -9,7 +9,7 @@
 from dartsort.localize.localize_torch import point_source_amplitude_at
 from dartsort.main import subtract
 from dartsort.util import waveform_util
-from neuropixel import dense_layout
+from test_util import dense_layout
 
 
 def test_fakedata_nonn():

tests/test_transform.py

@@ -1,14 +1,13 @@
-import neuropixel
 import numpy as np
-import torch
 from dartsort.transform import WaveformPipeline, transformers_by_class_name
 from dartsort.util.waveform_util import make_channel_index
+from test_util import dense_layout
 
 
 def test_all_transformers():
     # make a bunch of fake waveforms, put all of the transformers into
     # one long pipeline, and try running its fit and forward
-    h = neuropixel.dense_layout()
+    h = dense_layout()
     geom = np.c_[h["x"], h["y"]]
     channel_index = make_channel_index(geom, 100)
     rg = np.random.default_rng(0)

tests/test_util.py

@@ -24,3 +24,64 @@ def no_overlap_recording_sorting(templates, fs=30000, trough_offset_samples=42,
         ),
     )
     return rec, sorting
+
+
+def rc2xy(row, col, version=1):
+    """
+    converts the row/col indices to um coordinates.
+    :param row: row index on the probe
+    :param col: col index on the probe
+    :param version: neuropixel major version 1 or 2
+    :return: dictionary with keys x and y
+    """
+    if version == 1:
+        x = col * 16 + 11
+        y = (row * 20) + 20
+    elif np.floor(version) == 2:
+        x = col * 32
+        y = row * 15
+    return {"x": x, "y": y}
+
+
+def dense_layout(version=1, nshank=1, NC=384):
+    """Copied from ibl-neuropixel
+
+    Returns a dense layout indices map for neuropixel, as used at IBL
+    :param version: major version number: 1 or 2 or 2.4
+    :return: dictionary with keys 'ind', 'col', 'row', 'x', 'y'
+    """
+    ch = {
+        "ind": np.arange(NC),
+        "row": np.floor(np.arange(NC) / 2),
+        "shank": np.zeros(NC),
+    }
+
+    if version == 1:  # version 1 has a dense layout, checkerboard pattern
+        ch.update({"col": np.tile(np.array([2, 0, 3, 1]), int(NC / 4))})
+    elif (
+        np.floor(version) == 2 and nshank == 1
+    ):  # single shank NP1 has 2 columns in a dense patter
+        ch.update({"col": np.tile(np.array([0, 1]), int(NC / 2))})
+    elif (
+        np.floor(version) == 2 and nshank == 4
+    ):  # the 4 shank version default is rather complicated
+        shank_row = np.tile(np.arange(NC / 16), (2, 1)).T[:, np.newaxis].flatten()
+        shank_row = np.tile(shank_row, 8)
+        shank_row += (
+            np.tile(
+                np.array([0, 0, 1, 1, 0, 0, 1, 1])[:, np.newaxis], (1, int(NC / 8))
+            ).flatten()
+            * 24
+        )
+        ch.update(
+            {
+                "col": np.tile(np.array([0, 1]), int(NC / 2)),
+                "shank": np.tile(
+                    np.array([0, 1, 0, 1, 2, 3, 2, 3])[:, np.newaxis], (1, int(NC / 8))
+                ).flatten(),
+                "row": shank_row,
+            }
+        )
+    # for all, get coordinates
+    ch.update(rc2xy(ch["row"], ch["col"], version=version))
+    return ch

tests/test_waveform_util.py

@@ -2,7 +2,7 @@
 import torch
 import torch.nn.functional as F
 from dartsort.util import waveform_util
-from neuropixel import dense_layout
+from test_util import dense_layout
 
 
 def test_make_channel_index():