UnitMatchPy 2.42__tar.gz → 3.0__tar.gz

{unitmatchpy-2.42/UnitMatchPy.egg-info → unitmatchpy-3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: UnitMatchPy
-Version: 2.42
+Version: 3.0
 Summary: Automatic matching of neurons across sessions
 Author: Enny van Beest
 License: CC BY-NC-SA 4.0 International

{unitmatchpy-2.42 → unitmatchpy-3.0}/UnitMatchPy/bayes_functions.py

@@ -26,7 +26,7 @@ def get_parameter_kernels(scores_to_include, labels, cond, param, add_one = 1):
     ndarray
         The probability distributions for each metric
     """
-    print('Calculating the probability distributions of the metric scores')
+    # print('Calculating the probability distributions of the metric scores')
     score_vector = param['score_vector']
     bins = param['bins']
     smooth_prob = param['smooth_prob']
@@ -73,27 +73,43 @@ def apply_naive_bayes(parameter_kernels, priors, predictors, param, cond):
     ndarray
         The probability the unit is or is not a match
     """
-    print('Calculating the match probabilities')
+    # print('Calculating the match probabilities')
     
     score_vector = param['score_vector']
     n_pairs = predictors.shape[0] ** 2
+    n_metrics = predictors.shape[2]
 
-    unravel = np.reshape(predictors , (predictors.shape[0] * predictors.shape[1], predictors.shape[2],1))
-    x1 = np.tile(unravel, ( 1, 1, len(score_vector)))
-    tmp = np.expand_dims(score_vector, axis = (0,1))
-    x2 = np.tile(tmp, (x1.shape[0], x1.shape[1], 1))
-    min_idx = np.argmin( np.abs(x1 - x2), axis = 2)
+    predictors_flat = predictors.reshape(n_pairs, n_metrics)
+    
+    # Process in chunks
+    chunk_size = min(1000, n_pairs)
+    min_idx = np.zeros((n_pairs, n_metrics), dtype=int)
+    
+    for start_idx in range(0, n_pairs, chunk_size):
+        end_idx = min(start_idx + chunk_size, n_pairs)
+        chunk = predictors_flat[start_idx:end_idx]  # (chunk_size, n_metrics)
+        
+        # For each metric, find closest score indices for this chunk
+        for metric_idx in range(n_metrics):
+            metric_values = chunk[:, metric_idx]  # (chunk_size,)
+            # Broadcast: (chunk_size, 1) - (1, n_scores) = (chunk_size, n_scores)
+            distances = np.abs(metric_values[:, np.newaxis] - score_vector[np.newaxis, :])
+            min_idx[start_idx:end_idx, metric_idx] = np.argmin(distances, axis=1)
 
     likelihood = np.full((n_pairs, len(cond)), np.nan)
     for ck in range(len(cond)):
-        tmp_prob = np.zeros_like(min_idx, np.float64)
-        for yy in range(min_idx.shape[1]):
-            tmp_prob[:,yy] = parameter_kernels[min_idx[:,yy],yy,ck]
-        likelihood[:,ck] = np.prod(tmp_prob, axis=1)
 
+        metric_indices = np.arange(n_metrics)[np.newaxis, :]  # (1, n_metrics)
+        
+        # Extract probabilities for all pairs and metrics at once
+        tmp_prob = parameter_kernels[min_idx, metric_indices, ck]  # (n_pairs, n_metrics)
+        
+        # Compute product across metrics for each pair
+        likelihood[:, ck] = np.prod(tmp_prob, axis=1)
 
-    prob = np.full((n_pairs,2), np.nan )
-    for ck in range(len(cond)):
-        prob[:,ck] = priors[ck] * likelihood[:,ck] / np.nansum((priors * likelihood), axis =1)
+    denominator = np.nansum(priors[np.newaxis, :] * likelihood, axis=1)  # (n_pairs,)
+    
+    # Compute final probabilities for all conditions at once
+    prob = (priors[np.newaxis, :] * likelihood) / denominator[:, np.newaxis]  # (n_pairs, len(cond))
     
     return prob

{unitmatchpy-2.42 → unitmatchpy-3.0}/UnitMatchPy/metric_functions.py

@@ -168,24 +168,35 @@ def get_Euclidean_dist(avg_waveform_per_tp_flip,param):
     ndarray
         The euclidean distance between each unit for each time point
     """
-    # NOTE these arrays could get very large, may needed calculate in batches or use dask
-    # if arrays become too large
-
+    # Memory-efficient implementation using loops instead of massive tiling
     waveidx = param['waveidx']
     n_units = param['n_units']
-
-    x1 = np.tile( np.expand_dims(avg_waveform_per_tp_flip[:,:,waveidx,0,:], axis = -1), (1,1,1,1,n_units)).squeeze()
-    x2 = np.swapaxes(np.tile( np.expand_dims(avg_waveform_per_tp_flip[:,:,waveidx,1,:], axis = -1), (1,1,1,1,n_units)).squeeze(), 1, 4)
-
-    w = np.isnan( np.abs(x1[0,:,:,:,:] - x2[0,:,:,:,:])).squeeze()
-
-    tmp_euclid = np.linalg.norm(x1-x2, axis = 0)
-    tmp_euclid[w] = np.nan
-    euclid_dist = tmp_euclid
-    del x1
-    del x2
-    del w
-    del tmp_euclid
+    
+    # Extract the relevant slices once
+    data_cv0 = avg_waveform_per_tp_flip[:,:,waveidx,0,:]  # (3, n_units, len(waveidx), n_flips)
+    data_cv1 = avg_waveform_per_tp_flip[:,:,waveidx,1,:]  # (3, n_units, len(waveidx), n_flips)
+    
+    # Initialise output array to match original
+    euclid_dist = np.full((n_units, len(waveidx), data_cv0.shape[-1], n_units), np.nan)
+    
+    # Process unit by unit
+    for i in range(n_units):
+        for j in range(n_units):
+            # Get data for units i and j
+            unit_i_data = data_cv0[:, i, :, :]  # (3, len(waveidx), n_flips)
+            unit_j_data = data_cv1[:, j, :, :]  # (3, len(waveidx), n_flips)
+            
+            # Check for NaN values
+            w = np.isnan(np.abs(unit_i_data[0,:,:] - unit_j_data[0,:,:]))
+            
+            # Compute Euclidean distance between units i and j
+            diff = unit_i_data - unit_j_data  # (3, len(waveidx), n_flips)
+            tmp_euclid = np.linalg.norm(diff, axis=0)  # (len(waveidx), n_flips)
+            tmp_euclid[w] = np.nan
+            
+            # Store result
+            euclid_dist[i, :, :, j] = tmp_euclid
+    
     return euclid_dist
 
 def centroid_metrics(euclid_dist, param):
@@ -323,7 +334,7 @@ def dist_angle(avg_waveform_per_tp_flip, param):
         for dim_id2 in np.arange(1,avg_waveform_per_tp_flip.shape[0]):
             if dim_id2 <= dim_id1:
                 continue
-            ang = np.abs( x1[dim_id1,:,:,:,:] - x2[dim_id1,:,:,:,:]) / np.abs(x1[dim_id2,:,:,:,:] - x2[dim_id2,:,:,:,:])
+            ang = np.abs( x1[dim_id1,:,:,:,:] - x2[dim_id1,:,:,:,:]) / (np.abs(x1[dim_id2,:,:,:,:] - x2[dim_id2,:,:,:,:]) + 1e-10)
             
             loc_angle[:,:,:,:,count_id] = np.arctan(ang) * good_angle # only selects angles for units where there is sufficient distance between time poitns
             count_id +=1
@@ -766,23 +777,25 @@ def drift_n_sessions(candidate_pairs, session_switch, avg_centroid, avg_waveform
     drifts = np.zeros( (n_sessions - 1, 3))
 
     for did in range(n_sessions - 1):
-            idx = np.argwhere( ( (best_pairs[:,0] >= session_switch[did]) * (best_pairs[:,0] < session_switch[did + 1]) *
-                                (best_pairs[:,1] >= session_switch[did + 1]) * (best_pairs[:,1] < session_switch[did + 2]) ) == True)
-
-            pairs = best_pairs[idx,:].squeeze()
-            if best_match == True:
-                pairs = get_good_matches(pairs, total_score)
-
-            #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 = 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}')
-            elif len(pairs)>0: #if there exist pairs across sessions:
-                drifts = np.zeros( (n_sessions - 1, 3))
-                drifts[did,:], avg_waveform_per_tp, avg_centroid = apply_drift_correction_basic(pairs, did, session_switch, avg_centroid, avg_waveform_per_tp)
-            elif len(pairs)==0:
-                print('No pairs across sessions to perform drift correction.')
+        idx = np.argwhere( ( (best_pairs[:,0] >= session_switch[did]) * (best_pairs[:,0] < session_switch[did + 1]) *
+                            (best_pairs[:,1] >= session_switch[did + 1]) * (best_pairs[:,1] < session_switch[did + 2]) ) == True)
+
+        pairs = best_pairs[idx,:].squeeze()
+        if pairs.ndim == 1:
+            pairs = np.expand_dims(pairs, axis = 0)
+        if best_match == True:
+            pairs = get_good_matches(pairs, total_score)
+
+        #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 = 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}')
+        elif len(pairs)>0: #if there exist pairs across sessions:
+            drifts = np.zeros( (n_sessions - 1, 3))
+            drifts[did,:], avg_waveform_per_tp, avg_centroid = apply_drift_correction_basic(pairs, did, session_switch, avg_centroid, avg_waveform_per_tp)
+        elif len(pairs)==0:
+            print('No pairs across sessions to perform drift correction.')
     return drifts, avg_centroid, avg_waveform_per_tp
 
 

{unitmatchpy-2.42 → unitmatchpy-3.0}/UnitMatchPy/overlord.py

@@ -38,7 +38,7 @@ def extract_parameters(waveform, channel_pos, clus_info, param):
     
     return extracted_wave_properties
 
-def extract_metric_scores(extracted_wave_properties, session_switch, within_session, param, niter  = 2):
+def extract_metric_scores(extracted_wave_properties, session_switch, within_session, param, niter  = 2, to_use = None):
     """
     This function runs all of the metric calculations and drift correction to calculate the probability
     distribution needed for UnitMatch.
@@ -56,6 +56,8 @@ def extract_metric_scores(extracted_wave_properties, session_switch, within_sess
     niter : int, optional
         The number of pass through the function, 1 mean no drift correction
             2 is one pass of drift correction, by default 2
+    to_use : list, optional
+        A list of scores to include in the total score calculation. If None (default), all scores are included.
 
     Returns
     -------
@@ -106,7 +108,11 @@ def extract_metric_scores(extracted_wave_properties, session_switch, within_sess
 
         scores_to_include = {'amp_score' : amp_score, 'spatial_decay_score' : spatial_decay_score, 'centroid_overlord_score' : centroid_overlord_score,
                         'centroid_dist' : centroid_dist, 'waveform_score' : waveform_score, 'trajectory_score': trajectory_score }
-
+        if to_use is not None:
+            to_exclude = [key for key in scores_to_include if key not in to_use]
+            for key in to_exclude:
+                scores_to_include.pop(key, None)
+        
         total_score, predictors = mf.get_total_score(scores_to_include, param)
 
         #Initial thresholding
@@ -120,7 +126,6 @@ def extract_metric_scores(extracted_wave_properties, session_switch, within_sess
 
             drifts, avg_centroid, avg_waveform_per_tp = mf.drift_n_sessions(candidate_pairs, session_switch, avg_centroid, avg_waveform_per_tp, total_score, param)
 
-
     thrs_opt = mf.get_threshold(total_score, within_session, euclid_dist, param, is_first_pass = False)
     param['n_expected_matches'] = np.sum( (total_score > thrs_opt).astype(int))
     prior_match = 1 - ( param['n_expected_matches'] / len(include_these_pairs))

unitmatchpy-3.0/UnitMatchPy/run.py

@@ -0,0 +1,114 @@
+# This is an example usage script. It will need to be modified to run on your data.
+# The script processes data from lots of mice concurrently so it is fast.
+# The end result of the script is to write MatchProb values to a SQL database with a match table for each mouse/probe/location combination.
+# It also assumes that the waveform snippets are stored in a specific directory structure: root/mouse/probe/location/experiment_id/processed_waveforms
+
+import os, sys
+sys.path.insert(0, os.path.dirname(os.getcwd()))
+sys.path.insert(0, os.getcwd())
+sys.path.insert(0, os.path.join(os.path.dirname(os.getcwd()), 'UnitMatchPy'))
+import bayes_functions as bf
+import utils as util
+import overlord as ov
+import numpy as np
+import default_params as default_params
+from concurrent.futures import ProcessPoolExecutor
+import pandas as pd
+
+
+def process_single_location(location_data, path_lookup, param):
+
+    try:
+        mouse, probe, loc, mt_path = location_data
+
+        print(f"Processing {mouse}, {probe}, {loc}")
+
+        # Check if already processed (with retry logic)
+        table_name = f'{mouse}_{probe}_{loc}'
+
+        # Create robust database connection with retry logic
+        conn, cursor = util.create_robust_db_connection(r"C:\Users\suyash\UCL\DeepUnitMatch\DeepMatch\matchtables.db")
+        
+        def _check_if_processed():
+            cursor.execute(f"PRAGMA table_info({table_name})")
+            columns = [column[1] for column in cursor.fetchall()]
+            if "MatchProbNew" in columns:
+                # additionally check if the column contains NULLs
+                cursor.execute(f"SELECT COUNT(*) FROM {table_name} WHERE MatchProbNew IS NULL")
+                count_null = cursor.fetchone()[0]
+                if count_null == 0:
+                    return True
+            return False
+        
+        if util.retry_db_operation(_check_if_processed):
+            print(f"Skipping {mouse}, {probe}, {loc} as it is already processed.")
+            conn.close()
+            return None
+
+        KS_dirs = path_lookup.loc[(path_lookup['mouse'] == mouse) & (path_lookup['probe'] == probe) & (path_lookup['loc'] == loc), 'recordings'].to_list()[0]
+        param['KS_dirs'] = KS_dirs
+        wave_paths, unit_label_paths, channel_pos = util.paths_from_KS(KS_dirs)
+        param = util.get_probe_geometry(channel_pos[0], param)
+
+        waveform, session_id, session_switch, within_session, good_units, param = util.load_good_waveforms(wave_paths, unit_label_paths, param, good_units_only = False)
+        waveform, session_switch, within_session, session_id, good_units, param = util.filter_good_units_and_merge(mt_path, mouse, KS_dirs, waveform, session_switch, param)
+        # waveform, session_switch, within_session, session_id, good_units, param = util.filter_good_units(mouse, probe, loc, conn, waveform, session_switch, param)
+
+        mt = pd.read_sql_query(f"SELECT RecSes1,RecSes2,ID1,ID2 FROM {table_name}", conn)
+        ids = {}
+        for session in mt['RecSes1'].unique():
+            df = mt.loc[mt['RecSes1'] == session]
+            ids[session] = df['ID1'].unique()
+
+        clus_info = {'session_switch' : session_switch, 'session_id' : session_id}
+
+        # STEP 1
+        # Extract parameters from waveform
+        extracted_wave_properties = ov.extract_parameters(waveform, channel_pos, clus_info, param)
+
+        # STEP 2, 3, 4
+        # Extract metric scores
+        total_score, candidate_pairs, scores_to_include, predictors  = ov.extract_metric_scores(extracted_wave_properties, session_switch, within_session, param, niter  = 2)
+
+        prior_match = 1 - (param['n_expected_matches'] / param['n_units']**2 )
+        priors = np.array((prior_match, 1-prior_match))
+
+        labels = candidate_pairs.astype(int)
+        cond = np.unique(labels)
+        parameter_kernels = bf.get_parameter_kernels(scores_to_include, labels, cond, param, add_one = 1)
+        probability = bf.apply_naive_bayes(parameter_kernels, priors, predictors, param, cond)
+
+        util.add_col_to_sql(conn, cursor, mouse, probe, loc, "MatchProbNew", probability[:, 1], overwrite=True)
+        conn.close()
+
+        print(f"Finished processing {mouse}, {probe}, {loc}")
+        return probability[:, 1]
+    
+    except Exception as e:
+        print(f"Error processing {mt_path}: {e}")
+        try:
+            if 'conn' in locals():
+                conn.close()
+        except:
+            pass
+        return None
+
+
+if __name__ == "__main__":
+    param = default_params.get_default_param()
+    root = r"path/to/data/root"  # Root directory containing all data folders - data should be organised under this as mouse/probe/location
+
+    paths = pd.DataFrame(util.read_datapaths(os.listdir(root), r'/server/data/base'))
+    valid_locs = util.get_valid_locations(root)
+    # paths is a pandas dataframe with columns mouse, probe, loc, recordings, where recordings is a list of KS_dirs for that mouse/probe/location
+    # valid_locs is a list of tuples (mouse, probe, loc, mt_path) for each location with a match table
+    
+    max_workers = min(os.cpu_count() - 1, 12)
+
+    with ProcessPoolExecutor(max_workers=max_workers) as executor:
+        # Submit all jobs
+        future_to_location = {
+            executor.submit(process_single_location, location, paths, param): location
+            for location in valid_locs
+        }
+