coverage code accounts for 4 shank probes

needles2/probe_model.py

@@ -1,5 +1,6 @@
 import time
 import copy
+import re
 import numpy as np
 import pandas as pd
 from scipy.signal import fftconvolve
@@ -317,16 +318,20 @@ def compute_coverage(self, trajs, dist_fcn=[50, 100], limit=True, coverage=None,
         """
 
         ba = self.ba
-        ACTIVE_LENGTH_UM = 3.84 * 1e3  # This is the length of the NP1 probe with electrodes
+        ACTIVE_LENGTH_UM_1shank = 3.84 * 1e3  # This is the length of the NP1 probe with electrodes
+        ACTIVE_LENGTH_UM_4shank = 705  # This is the length of the NP2 4 shank probe with electrodes
         MAX_DIST_UM = dist_fcn[1]  # max distance around the probe to be searched for
 
         # Covered_length_um are two values which indicate on the path from tip to entry of a given insertion
         # where the region that is considered to be covered by this insertion begins and ends in micro meter
         # Note that the second value is negative, because the covered regions extends beyond the tip of the probe
         # We multiply by sqrt(2) to translate the radius given by dist_fcn[1] into the side length of a square
         # that is contained in the circle with radius dist_fcn[1]
-        covered_length_um = TIP_SIZE_UM + np.array([ACTIVE_LENGTH_UM + MAX_DIST_UM * np.sqrt(2),
-                                                    -MAX_DIST_UM * np.sqrt(2)])
+        covered_length_um_1shank = TIP_SIZE_UM + np.array([ACTIVE_LENGTH_UM_1shank + MAX_DIST_UM * np.sqrt(2),
+                                                           -MAX_DIST_UM * np.sqrt(2)])
+        covered_length_um_4shank = TIP_SIZE_UM + np.array([ACTIVE_LENGTH_UM_4shank + MAX_DIST_UM * np.sqrt(2),
+                                                           -MAX_DIST_UM * np.sqrt(2)])
+
 
         # Horizontal slice of voxels to be considered around each trajectory is only dependent on MAX_DIST_UM
         # and the voxel resolution, so can be defined here. We translate max dist in voxels and add 1 for safety
@@ -361,6 +366,19 @@ def crawl_up_from_tip(ins, covered_length):
             if len(trajs) > 20 and self.verbose is True:
                 if p % 20 == 0:
                     print(p / len(trajs))
+
+            # Here we find out if this is trajectory is from a 1shank or 4shank probe
+            # In an ideal world we would read in the metadata and find this out, but that would require
+            # downloading this dataset for all trajectories. Instead we go based on naming convention. We know
+            # probes with the label probe00a, probe00b etc. are 4 shank probes that have been split.
+            pname = traj['probe_name']
+            if len(re.findall("[a-d]", pname[-1])) == 1:
+                covered_length_um = covered_length_um_4shank
+                ACTIVE_LENGTH_UM = ACTIVE_LENGTH_UM_4shank
+            else:
+                covered_length_um = covered_length_um_1shank
+                ACTIVE_LENGTH_UM = ACTIVE_LENGTH_UM_1shank
+
             # Get one trajectory from the list and create an insertion in the brain atlas
             # x and y coordinates of entry are translated to the atlas voxel space
             # z is locked to surface of the brain at these x,y coordinates (disregarding actual z value of trajectory)
@@ -513,16 +531,19 @@ def compute_coverage_dict(self, trajs, dist_fcn=[50, 100], limit=True, coverage=
         """
 
         ba = self.ba
-        ACTIVE_LENGTH_UM = 3.84 * 1e3  # This is the length of the NP1 probe with electrodes
+        ACTIVE_LENGTH_UM_1shank = 3.84 * 1e3  # This is the length of the NP1 probe with electrodes
+        ACTIVE_LENGTH_UM_4shank = 705  # This is the length of the NP2 4 shank probe with electrodes
         MAX_DIST_UM = dist_fcn[1]  # max distance around the probe to be searched for
 
         # Covered_length_um are two values which indicate on the path from tip to entry of a given insertion
         # where the region that is considered to be covered by this insertion begins and ends in micro meter
         # Note that the second value is negative, because the covered regions extends beyond the tip of the probe
         # We multiply by sqrt(2) to translate the radius given by dist_fcn[1] into the side length of a square
         # that is contained in the circle with radius dist_fcn[1]
-        covered_length_um = TIP_SIZE_UM + np.array([ACTIVE_LENGTH_UM + MAX_DIST_UM * np.sqrt(2),
-                                                    -MAX_DIST_UM * np.sqrt(2)])
+        covered_length_um_1shank = TIP_SIZE_UM + np.array([ACTIVE_LENGTH_UM_1shank + MAX_DIST_UM * np.sqrt(2),
+                                                           -MAX_DIST_UM * np.sqrt(2)])
+        covered_length_um_4shank = TIP_SIZE_UM + np.array([ACTIVE_LENGTH_UM_4shank + MAX_DIST_UM * np.sqrt(2),
+                                                           -MAX_DIST_UM * np.sqrt(2)])
 
         # Horizontal slice of voxels to be considered around each trajectory is only dependent on MAX_DIST_UM
         # and the voxel resolution, so can be defined here. We translate max dist in voxels and add 1 for safety
@@ -542,10 +563,25 @@ def crawl_up_from_tip(ins, covered_length):
             if len(trajs) > 20 and self.verbose is True:
                 if p % 20 == 0:
                     print(p / len(trajs))
-            # Get one trajectory from the list and create an insertion in the brain atlas
+
+            # Get one trajectory from the list and
+            traj = trajs[p]
+
+            # Here we find out if this is trajectory is from a 1shank or 4shank probe
+            # In an ideal world we would read in the metadata and find this out, but that would require
+            # downloading this dataset for all trajectories. Instead we go based on naming convention. We know
+            # probes with the label probe00a, probe00b etc. are 4 shank probes that have been split.
+            pname = traj['probe_name']
+            if len(re.findall("[a-d]", pname[-1])) == 1:
+                covered_length_um = covered_length_um_4shank
+                ACTIVE_LENGTH_UM = ACTIVE_LENGTH_UM_4shank
+            else:
+                covered_length_um = covered_length_um_1shank
+                ACTIVE_LENGTH_UM = ACTIVE_LENGTH_UM_1shank
+
+            # Create an insertion in the brain atlas
             # x and y coordinates of entry are translated to the atlas voxel space
             # z is locked to surface of the brain at these x,y coordinates (disregarding actual z value of trajectory)
-            traj = trajs[p]
             ins = atlas.Insertion.from_dict(traj, brain_atlas=ba)
             # Don't use probes that have same entry and tip, something is wrong
             set_nan = False