probe geom tests

int-brain-lab · Aug 15, 2024 · ce394bd · ce394bd
1 parent 98a4bac
commit ce394bd
Showing 1 changed file with 70 additions and 0 deletions.
diff --git a/src/tests/unit/cpu/test_neuropixel.py b/src/tests/unit/cpu/test_neuropixel.py
@@ -15,3 +15,73 @@ def test_adc_shifts():
     h24 = neuropixel.trace_header(version=2.4)
     np.testing.assert_equal(h24["sample_shift"], h21["sample_shift"])
     np.testing.assert_equal(np.unique(h21["sample_shift"] * 16), np.arange(16))
+    # test ADC shifts uhd
+    hUHD = neuropixel.trace_header(version="uhd")
+    np.testing.assert_equal(hUHD["sample_shift"], h1["sample_shift"])
+
+
+def test_geom_np1():
+    gt = dict(
+        ind=np.arange(384),
+        shank=np.zeros(384),
+        row=np.repeat(np.arange(192), 2),
+        col=np.tile(np.array([2, 0, 3, 1]), 96),
+        x=np.tile(np.array([43, 11, 59, 27]), 96),
+        y=np.repeat(np.arange(0, 3840, 20), 2) + 20
+    )
+
+    h = neuropixel.trace_header(1)
+    for k, v in gt.items():
+        np.testing.assert_equal(v, h[k])
+
+
+def test_geom_np2_1shank():
+    gt = dict(
+        ind=np.arange(384),
+        shank=np.zeros(384),
+        row=np.repeat(np.arange(192), 2),
+        col=np.tile(np.array([0, 1]), 192),
+        x=np.tile(np.array([27, 59]), 192),
+        y=np.repeat(np.arange(0, 2880, 15), 2) + 20
+    )
+
+    h = neuropixel.trace_header(2, 1)
+    for k, v in gt.items():
+        np.testing.assert_equal(v, h[k])
+
+
+def test_geom_np2_4shank():
+    depth_blocks = np.vstack(
+        [
+            np.repeat(np.arange(24), 2),
+            np.repeat(np.arange(24, 48), 2)
+        ]
+    )
+    row_ind = np.concatenate([depth_blocks[i] for i in [0, 0, 1, 1, 0, 0, 1, 1]])
+    gt = dict(
+        ind=np.arange(384),
+        shank=np.repeat(np.array([0, 1, 0, 1, 2, 3, 2, 3]), 48),
+        row=row_ind,
+        col=np.tile(np.array([0, 1]), 192),
+        x=np.tile(np.array([27, 59]), 192),
+        y=row_ind * 15 + 20
+    )
+
+    h = neuropixel.trace_header(2, 4)
+    for k, v in gt.items():
+        np.testing.assert_equal(v, h[k])
+
+
+def test_geom_npultra():
+    gt = dict(
+        ind=np.arange(384),
+        shank=np.zeros(384),
+        row=np.repeat(np.arange(48), 8),
+        col=np.tile(np.arange(8), 48),
+        x=np.tile(np.arange(0, 48, 6), 48),
+        y=np.repeat(np.arange(0, 288, 6), 8)
+    )
+
+    h = neuropixel.trace_header("uhd")
+    for k, v in gt.items():
+        np.testing.assert_equal(v, h[k])