UnitMatchPy 2.20__tar.gz → 2.31__tar.gz

{UnitMatchPy-2.20/UnitMatchPy.egg-info → UnitMatchPy-2.31}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: UnitMatchPy
-Version: 2.20
+Version: 2.31
 Author: Enny van Beest, Celian Bimbard and Sam Dodgson
 Author-email: e.beest@ucl.ac.uk
 Requires-Python: >=3.7

{UnitMatchPy-2.20 → UnitMatchPy-2.31}/UnitMatchPy/default_params.py

@@ -23,8 +23,7 @@ def get_default_param(param = None):
             'stepsz' : 0.01, # size of of a step in probability distributions
             'smooth_prob' : 9, # probability smoothing size
             'min_angle_dist' : 0.1, # smallest distance for and angle to be consider
-            'no_shanks' : 4,  # No. probe shanks -default neuropixels 2
-            'shank_dist' : 175, # distance between shanks
+            'min_new_shank_distance' : 100, #The smallest distance which separates 2 shanks
             'units_per_shank_thrs' : 15, # threshold for doing per shank drift correction
             'match_threshold' : 0.5 # probability threshold to consider as a match
         }

{UnitMatchPy-2.20 → UnitMatchPy-2.31}/UnitMatchPy/metric_functions.py

@@ -547,7 +547,7 @@ def apply_drift_correction_basic(pairs, sid, session_switch, avg_centroid, avg_w
     ##need to add the drift to the location
     avg_waveform_per_tp[0,session_switch[sid+1]:session_switch[sid+2],:,:] += drift[0]
     avg_waveform_per_tp[1,session_switch[sid+1]:session_switch[sid+2],:,:] += drift[1]
-    avg_waveform_per_tp[2,session_switch[sid+1]:session_switch[sid]+2,:,:] += drift[2]
+    avg_waveform_per_tp[2,session_switch[sid+1]:session_switch[sid+2],:,:] += drift[2]
 
     avg_centroid[0,session_switch[sid+1]:session_switch[sid+2],:] += drift[0]
     avg_centroid[1,session_switch[sid+1]:session_switch[sid+2],:] += drift[1]
@@ -555,7 +555,7 @@ def apply_drift_correction_basic(pairs, sid, session_switch, avg_centroid, avg_w
 
     return drift, avg_waveform_per_tp, avg_centroid
 
-def appply_drift_correction_per_shank(pairs, sid, session_switch, avg_centroid, avg_waveform_per_tp, param):
+def apply_drift_correction_per_shank(pairs, sid, session_switch, avg_centroid, avg_waveform_per_tp, param):
     """
     This is the same as "basic" drift correction, however treats each shank separately, if there is enough units per shank
 
@@ -587,25 +587,22 @@ def appply_drift_correction_per_shank(pairs, sid, session_switch, avg_centroid,
 
     max_dist = 0
     min_dist = 0
-    drift_per_shank = np.zeros([4,3])
+    drift_per_shank = np.zeros([n_shanks,3])
 
     for i in range(n_shanks):
         max_dist += shank_dist
 
-        #test to see if a centroid is within the area of that shank
+        #get which shank each unit is on
         correct_shank_a = np.logical_and(centroid_a[1,:] < max_dist,  centroid_a[1,:] > min_dist)
         correct_shank_b = np.logical_and(centroid_b[1,:] < max_dist,  centroid_b[1,:] > min_dist)
 
-        if np.all(correct_shank_a == correct_shank_b) != True:
-            print(f'These pairs may be bad {np.argwhere(correct_shank_a != correct_shank_b)}')
-
         drifts = centroid_a[:,correct_shank_a] - centroid_b[:,correct_shank_b]
-        drift =  np.nanmedian(drifts, axis = 1)
+        drift = np.nanmedian(drifts, axis = 1)
         drift_per_shank[i,:] = drift
 
         #need to get idx for each shank, to apply correct drift correction
 
-        shank_session_idx = session_switch[sid] + np.argwhere( shank_id == i) 
+        shank_session_idx = session_switch[sid] + np.argwhere(shank_id == i) 
 
         avg_waveform_per_tp[0,shank_session_idx,:,:] += drift[0]
         avg_waveform_per_tp[1,shank_session_idx,:,:] += drift[1]
@@ -694,25 +691,30 @@ def test_matches_per_shank(pairs, avg_centroid, sid, param):
     n_shanks = param['no_shanks']
     shank_dist = param['shank_dist']
 
-    for i in range(n_shanks):
-        max_dist += shank_dist
+    if n_shanks != 1:
+        for i in range(n_shanks):
+            max_dist += shank_dist
 
-        correct_shank_a = np.logical_and(centroid_a[1,:] < max_dist,  centroid_a[1,:] > min_dist)
-        correct_shank_b = np.logical_and(centroid_b[1,:] < max_dist,  centroid_b[1,:] > min_dist)
+            correct_shank_a = np.logical_and(centroid_a[1,:] < max_dist,  centroid_a[1,:] > min_dist)
+            correct_shank_b = np.logical_and(centroid_b[1,:] < max_dist,  centroid_b[1,:] > min_dist)
 
-        if np.all(correct_shank_a == correct_shank_b) != True:
-            print(f'These pairs may be bad {np.argwhere(correct_shank_a != correct_shank_b)}')
-            
-        shank_id_tmp[correct_shank_a] = i
+            if np.all(correct_shank_a == correct_shank_b) != True:
+                #print(f'These pairs may be bad {np.argwhere(correct_shank_a != correct_shank_b)}')
+                #delete pairs which are one different shanks
+                correct_shank_a = np.delete(correct_shank_a, np.argwhere(correct_shank_a != correct_shank_b))
+                correct_shank_b = np.delete(correct_shank_b, np.argwhere(correct_shank_a != correct_shank_b))
+                
+            shank_id_tmp[correct_shank_a] = i
 
-        min_dist += shank_dist
+            min_dist += shank_dist
 
-    __, counts = np.unique(shank_id_tmp, return_counts=True)
+        __, counts = np.unique(shank_id_tmp, return_counts=True)
 
-    if np.any(counts < units_per_shank_thrs):
+        if np.any(counts < units_per_shank_thrs):
+            do_per_shank_correction = False
+            print(f'Session pair {sid+1}/{sid+2} has {counts} matches per shank, which is below threshold to do per shank drift correction')
+    else:
         do_per_shank_correction = False
-        print(f'Session pair {sid+1}/{sid+2} has {counts} matches per shank, which is below threshold to do per shank drift correction')
-
     return do_per_shank_correction
 
 
@@ -765,8 +767,8 @@ def drift_n_sessions(candidate_pairs, session_switch, avg_centroid, avg_waveform
 
             #Test to see if there are enough matches to do drift correction per shank
             if test_matches_per_shank(pairs, avg_centroid, did, param) == True and best_drift == True:
-                drifts = np.zeros( (n_sessions - 1, param['no_shanks'], 3)) 
-                drifts[did,:,:], avg_waveform_per_tp, avg_centroid = appply_drift_correction_per_shank(pairs, did, session_switch, avg_centroid, avg_waveform_per_tp, param)
+                drifts = np.zeros( (n_sessions - 1, param['no_shanks'], 3))
+                drifts[did,:,:], avg_waveform_per_tp, avg_centroid = apply_drift_correction_per_shank(pairs, did, session_switch, avg_centroid, avg_waveform_per_tp, param)
                 print(f'Done drift correction per shank for session pair {did+1} and {did+2}')
             else:
                 drifts = np.zeros( (n_sessions - 1, 3))

{UnitMatchPy-2.20 → UnitMatchPy-2.31}/UnitMatchPy/utils.py

@@ -124,6 +124,7 @@ def load_good_waveforms(wave_paths, unit_label_paths, param, good_units_only = T
 
     good_units = []
     n_units_per_session_all = []
+    all_units = []
 
     for i in range(len(unit_label_paths)):
     #see if bombcell unit labels
@@ -139,6 +140,7 @@ def load_good_waveforms(wave_paths, unit_label_paths, param, good_units_only = T
         n_units_per_session_all.append(unit_label.shape[0])
         good_unit_idx = unit_label[tmp_idx, 0]
         good_units.append(good_unit_idx)
+        all_units.append(unit_label[:,0]) #get all the unit labels 
 
     waveforms = []
     if good_units_only:
@@ -162,16 +164,17 @@ def load_good_waveforms(wave_paths, unit_label_paths, param, good_units_only = T
     else:
         for ls in range(len(wave_paths)):
             #load in the first good unit, to get the shape of each waveform
-            p_file = os.path.join(wave_paths[ls],f'Unit{int(good_units[ls][0].squeeze())}_RawSpikes.npy')
+            p_file = os.path.join(wave_paths[ls],f'Unit{int(all_units[ls][0])}_RawSpikes.npy')
             tmp = np.load(p_file)
             tmp_waveform = np.zeros( (len(os.listdir(wave_paths[ls])), tmp.shape[0], tmp.shape[1], tmp.shape[2]))
 
             for i in range(len(os.listdir(wave_paths[ls]))):
                 #loads in all GoodUnits for that session
-                p_file_good = os.path.join(wave_paths[ls], f'Unit{int(good_units[ls][0].squeeze())}_RawSpikes.npy')
+                p_file_good = os.path.join(wave_paths[ls], f'Unit{int(all_units[ls][i])}_RawSpikes.npy')
                 tmp_waveform[i] = np.load(p_file_good)
             #adds that session to the list
             waveforms.append(tmp_waveform)
+            print(f'UnitMatch is treating all the units as good and including all units from {wave_paths[ls]}, we recommended using curated data!')
 
         del tmp_waveform
         del tmp
@@ -605,4 +608,50 @@ def paths_from_KS(KS_dirs):
         else:
             raise Exception('Could not find RawWaveforms folder')
     
-    return wave_paths, unit_label_paths, channel_pos
+    return wave_paths, unit_label_paths, channel_pos
+
+def get_probe_geometry(channel_pos, param, verbose = False):
+    """
+    From the channel positions, estimate the number of shanks and there spacing used to 
+    identify a unit to a probe.
+    needmin_new_shank_distance from param.
+
+    Parameters
+    ----------
+    channel_pos : ndarray
+        The channel positions of a session 
+    param : dict
+        The param dict
+    verbose : bool, optional
+        If True will print the calculated results, by default False
+
+    Returns
+    -------
+    param
+        The param dictionary updated with the calculated params
+    """
+    min_new_shank_distance = param['min_new_shank_distance']
+    x_val = np.unique(channel_pos[:,0])
+    x_val = np.sort(x_val) #make sure they are in ascending order 
+    x_spacing = np.diff(x_val)
+    too_close = np.argwhere(x_spacing < min_new_shank_distance)
+
+    #remove when shanks are two close
+    x_val = np.delete(x_val, too_close + 1)
+
+    n_shanks = x_val.size
+    if n_shanks == 1:
+        shank_spacing = 100 #make it cover the full possible range
+    else:
+        shank_spacing = np.abs(np.diff(x_val))[0]
+
+    #NOTE this shank_dist is the distance within a centroid would be assigned to a shank!
+    shank_dist = shank_spacing * 0.9 #the area from the start which is assigned to each probe
+    #e.g 0 -> 0.9to shank 1, 0.9-1.8 to shank 2...
+    #NOTE current shank identification doesn't work for 9+ shanks
+
+    param['no_shanks'] = n_shanks
+    param['shank_dist'] = shank_dist
+    if verbose == True:
+        print(f'We have found {n_shanks} with spacing ~ {shank_spacing}')
+    return param

{UnitMatchPy-2.20 → UnitMatchPy-2.31/UnitMatchPy.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: UnitMatchPy
-Version: 2.20
+Version: 2.31
 Author: Enny van Beest, Celian Bimbard and Sam Dodgson
 Author-email: e.beest@ucl.ac.uk
 Requires-Python: >=3.7

{UnitMatchPy-2.20 → UnitMatchPy-2.31}/setup.py

@@ -5,7 +5,7 @@ with open('README.md', 'r') as f:
 
 setup(
     name="UnitMatchPy",
-    version="2.20",
+    version="2.31",
     descriptin="Description of the package",
     author="Enny van Beest, Celian Bimbard and Sam Dodgson",
     author_email="e.beest@ucl.ac.uk",