Repository: m-beau/NeuroPyxels (language: python)

Issue:
Make mplp() optional in every npyx function through 'prettify' argument.



Discussion / hints:

Fix patch:
diff --git a/npyx/__init__.py b/npyx/__init__.py
index 1c7af653..17c8bcfd 100644
--- a/npyx/__init__.py
+++ b/npyx/__init__.py
@@ -14,6 +14,7 @@
 from .merger import *
 from .circuitProphyler import *
 from .feat import *
+from .metrics import *
 from .info import *
 from .h5 import *
 from .testing import *
diff --git a/npyx/behav.py b/npyx/behav.py
index 4023d24b..a3ef1d4a 100644
--- a/npyx/behav.py
+++ b/npyx/behav.py
@@ -9,6 +9,10 @@
 from pathlib import Path
 import pickle
 
+from joblib import Memory
+cachedir = Path(op.expanduser("~")) / ".NeuroPyxels"
+cache_memory = Memory(cachedir, verbose=0)
+
 import numpy as np
 import scipy as sp
 import pandas as pd
@@ -21,10 +25,11 @@
 
 import cv2
 
-from npyx.utils import npa, thresh, thresh_consec, smooth, sign, align_timeseries, assert_int, get_bins
+from npyx.utils import npa, thresh, thresh_consec, smooth, sign, assert_int
 
 from npyx.inout import read_metadata, get_npix_sync, paq_read, list_files
-from npyx.spk_t import mean_firing_rate
+from npyx.gl import get_rec_len
+from npyx.spk_t import mean_firing_rate, get_common_good_sections
 from npyx.corr import crosscorr_cyrille, frac_pop_sync
 from npyx.merger import assert_multi, get_ds_table
 
@@ -647,6 +652,134 @@ def get_wheelturn_df_dic(dp, paqdic, include_wheel_data=False, add_spont_licks=F
 
     return df, wheel_turn_dic
 
+
+#%% Recording period extraction
+
+# Find periods without motion
+
+def load_baseline_periods(dp = None, behavdic = None, rec_len = None, dataset_with_opto = True, 
+                          speed_th_steer = 0.1, speed_th_run = 0.5, light_buffer = 0.5,
+                          again = False, verbose = True):
+    """
+    Function to calculate periods of undisturbed neural activity
+    (no (monitored) behaviour or optostims).
+
+    Arguments:
+        - dp: str, path to Neuropixels dataset (only provide if no behavdic) 
+        - behavdic: dict, output of behav_dic() (only provide if no dp)
+        - rec_len: float, recording length (in seconds). Only useful is dp is None.
+        - dataset_with_opto: whether the dataset had functioning channel rhodopsin and optostims were delivered
+        - speed_th_steer: threhsold to consider a 'steering' period (cm/s)
+        - speed_th_run: threshold to consider a 'running' period (mm/s)
+        - light_buffer: buffer to add before light onsets/after light offsets (s)
+        - again: bool, whether to reload behaviour and recalculate periods
+        - verbose: bool, whether to print details.
+    
+    Returns:
+        - baseline_periods: 2D array of shape (n_periods, 2), [[t1,t2], ...] in seconds.
+    """
+
+    assert dp is None or behavdic is None, "Only one of dp or behavdic should be provided"
+
+    if behavdic is None:
+        behavdic = behav_dic(dp, f_behav=None, vid_path=None, again=again, again_align=again, again_rawpaq=again,
+                     lick_ili_th=0.075, n_ticks=1024, diam=200, gsd=25,
+                     plot=False, drop_raw=False, cam_paqi_to_use=None)
+
+    if rec_len is None:
+        assert dp is not None, "You must provide the recording length in seconds!"
+        rec_len = get_rec_len(dp, unit='seconds')
+
+    if 'wheel_position_mm' in behavdic.keys():
+        behaviour = "steering"
+    elif 'ROT_SPEED' in behavdic.keys():
+        behaviour = "running"
+    else:
+        raise ValueError("Neither 'wheel_position_mm' nor 'ROT_SPEED' was found in behavdic!")
+
+    move_periods = light_periods = None
+    if dp is not None and not again:
+        dp = Path(dp)
+        th_str = speed_th_steer if behaviour == "steering" else speed_th_run
+        fn_move = f"periods_{th_str}_{behaviour}.npy"
+        fn_nomove = fn_move.replace(f"{th_str}_", f"{th_str}_no_")
+        if (dp/fn_move).exists():
+            move_periods = np.load(dp/fn_move)
+            nomove_periods = np.load(dp/fn_nomove)
+        if dataset_with_opto:
+            fn_light = f"periods_{light_buffer}_opto.npy"
+            fn_nolight = fn_light.replace(f"_opto.npy", f"_no_opto.npy")
+            if (dp/fn_light).exists() and dataset_with_opto:
+                light_periods = np.load(dp/fn_light)
+                no_light_periods = np.load(dp/fn_nolight)
+
+    if move_periods is None:
+        if behaviour == "steering":
+            wheel_speed = np.abs(behavdic['wheel_speed_cm/s']) 
+            span=1 # s
+            # because paqIO oversamples at 5000Hz whereas VR runs at 30Hz, 
+            # only onsets are meaningful -> cannot use onest to offset as periods
+            # instead, take all periods outside of +/-1second around any threshold crossing
+            steer_th = thresh(wheel_speed, speed_th_steer, 1)
+            steer_spans = np.hstack([(steer_th-span*behavdic['paq_fs'])[:,None], (steer_th+span*behavdic['paq_fs'])[:,None]])
+            steer_mask = np.zeros(wheel_speed.shape).astype(np.int8)
+            for steer_span in steer_spans:
+                steer_mask[steer_span[0]:steer_span[1]]=1
+            
+            steer_values   = thresh_consec(steer_mask, 0.5, 1)
+            move_periods   = np.round(npa([[r[0,0],r[0,-1]] for r in steer_values])*behavdic['a']+behavdic['b'])/behavdic['npix_fs']
+            nomove_periods = np.concatenate([[0], move_periods.ravel(), [rec_len]]).reshape((-1,2))
+            
+            if np.any(move_periods) and verbose:
+                print(f'{np.sum(np.diff(move_periods, axis=1))}s of steering in total.')
+                print(f'{np.sum(np.diff(nomove_periods, axis=1))}s of no steering in total.')
+                
+        elif behaviour == "running":
+            v=behavdic['ROT_SPEED']
+            consec=1 # s
+            run_values = thresh_consec(v, speed_th_run, consec, consec*behavdic['paq_fs'], ret_values=True)
+            move_periods = np.round(npa([[r[0,0],r[0,-1]] for r in run_values])*behavdic['a']+behavdic['b'])/behavdic['npix_fs']
+            nomove_periods = np.concatenate([[0], move_periods.ravel(), [rec_len]]).reshape((-1,2))
+            
+            if np.any(move_periods) and verbose:
+                print(f'{np.sum(np.diff(move_periods, axis=1))}s of running in total.')
+                print(f'{np.sum(np.diff(nomove_periods, axis=1))}s of no running in total.')
+            
+    # find periods without light
+    if dataset_with_opto and light_periods is None:
+        light_periods = np.array(list(zip(behavdic['opto_stims_ON_npix']/behavdic['npix_fs'], behavdic['opto_stims_OFF_npix']/behavdic['npix_fs'])))
+        light_periods_buffered = light_periods.copy()
+        light_periods_buffered[0] -= light_buffer
+        light_periods_buffered[1] += light_buffer
+        no_light_periods = np.concatenate([[0], light_periods_buffered.ravel(), [rec_len]])
+        no_light_periods = no_light_periods.reshape(-1, 2)
+
+        print(f'{np.sum(np.diff(light_periods, axis=1))}s of periods with light in total.')
+        print(f'{np.sum(np.diff(no_light_periods, axis=1))}s of periods with no light in total (buffer before onset/after offset: {light_buffer}s).')
+
+    # save arrays to disk
+    if dp is not None:
+        np.save(dp/fn_move, move_periods)
+        np.save(dp/fn_nomove, nomove_periods)
+        if dataset_with_opto:
+            np.save(dp/fn_light, light_periods)
+            np.save(dp/fn_nolight, no_light_periods)
+
+    if dataset_with_opto:
+        # get intersection of both
+        fs = 30_000
+        no_light_periods = (no_light_periods*fs).astype(int)
+        nomove_periods = (nomove_periods*fs).astype(int)
+        no_light_nor_move_periods = get_common_good_sections([list(no_light_periods), list(nomove_periods)])
+        no_light_nor_move_periods = np.array(no_light_nor_move_periods)/fs
+        print(f'{np.sum(np.diff(no_light_nor_move_periods, axis=1))}s of periods with no light nor running in total (buffer before onset/after light offset: {light_buffer}s).')
+        baseline_recording_periods = no_light_nor_move_periods
+    else:
+        baseline_recording_periods = nomove_periods
+
+    return baseline_recording_periods
+
+
 #%% Alignement, binning and processing of time series
 
 def align_variable(events, variable_t, variable, b=2, window=[-1000,1000], remove_empty_trials=False):
@@ -727,9 +860,6 @@ def align_times(times, events, b=2, window=[-1000,1000], remove_empty_trials=Fal
 
     return aligned_t, aligned_tb
 
-def time_to_phase():
-
-    return
 
 def align_times_manyevents(times, events, b=2, window=[-1000,1000], fs=30000):
     '''
@@ -815,6 +945,7 @@ def jPSTH(spikes1, spikes2, events, b=2, window=[-1000,1000], convolve=False, me
 
     return jpsth, jpsth_ccg, coincidence_psth
 
+@cache_memory.cache
 def get_ifr(times, events, b=2, window=[-1000,1000], remove_empty_trials=False):
     '''
     Parameters:
@@ -876,6 +1007,7 @@ def process_2d_trials_array(y, y_bsl, zscore=False, zscoretype='within',
 
     return y, y_p, y_p_var
 
+
 def get_processed_ifr(times, events, b=10, window=[-1000,1000], remove_empty_trials=False,
                       zscore=False, zscoretype='within',
                       convolve=False, gsd=1, method='gaussian',
@@ -1041,6 +1173,8 @@ def filter_allneurons_active(M, p=0, return_masks=False):
     return M[:,all_active_mask,:]
 
 
+#%% Population synchrony analysis
+
 def get_processed_popsync(trains, events, psthb=10, window=[-1000,1000],
                           events_tiling_frac=0.1, sync_win=2, fs=30000, t_end=None,
                           b=1, sd=1000, th=0.02,
diff --git a/npyx/corr.py b/npyx/corr.py
index cfff4948..a36315b2 100644
--- a/npyx/corr.py
+++ b/npyx/corr.py
@@ -41,8 +41,7 @@
 
 def make_phy_like_spikeClustersTimes(dp, U, periods='all', verbose=False, trains=None, enforced_rp=0):
     '''
-
-    - trains: dict, of the form {unit1:train1InSamples, unit2:...}'''
+    - trains: list of spike trains, in samples.'''
     trains_dic={}
     if trains is None:
         for iu, u in enumerate(U):
@@ -85,7 +84,7 @@ def crosscorrelate_cyrille(dp, bin_size, win_size, U, fs=30000, symmetrize=True,
        - U (list of integers): list of units indices.
        - fs: sampling rate (Hertz). Default 30000.
        - symmetrize (bool): symmetrize the semi correlograms. Default=True.
-       - trains: dictionnary of trains, to calculate the CCG of an arbitrary list of trains in SAMPLES for fs=30kHz.'''
+       - trains: list of spike trains, in samples.'''
 
     #### Get clusters and times
     U=list(U)
@@ -1020,7 +1019,7 @@ def ccg_sig_stack(dp, U_src, U_trg, cbin=0.5, cwin=100, name=None,
                 features=pd.read_csv(feat_path)
                 return sigstack, sigustack, features
             features=pd.DataFrame(columns=feat_columns)
-            for i,c in enumerate(tqdm(sigstack, desc=f'Computing features of significant CCGs overs {num_cores} cores')):
+            for i,c in enumerate(tqdm(sigstack, desc=f'Computing features of significant CCGs over {num_cores} cores')):
                 pks=get_ccg_sig(c, cbin, cwin, p_th, n_consec_bins, sgn,
                                 fract_baseline, W_sd, test, ret_features=ret_features, only_max=only_max)
                 for p in pks:
@@ -1056,7 +1055,7 @@ def ccg_sig_stack(dp, U_src, U_trg, cbin=0.5, cwin=100, name=None,
 
     ccgsig_results = Parallel(n_jobs=num_cores)(\
         delayed(get_ccg_sig)(*ccgsig_args[i]) for i in tqdm(range(len(ccgsig_args)),
-        desc=f'Looking for significant CCGs overs {num_cores} cores'))
+        desc=f'Looking for significant CCGs over {num_cores} cores'))
 
     sigustack=[]
     sigstack=[]
diff --git a/npyx/metrics.py b/npyx/metrics.py
new file mode 100644
index 00000000..b793721f
--- /dev/null
+++ b/npyx/metrics.py
@@ -0,0 +1,1044 @@
+## Code borrowed from Allen institute spikemetrics (jan 2022):
+## https://github.com/AllenInstitute/ecephys_spike_sorting/tree/master/ecephys_spike_sorting/modules/quality_metrics
+
+from pathlib import Path
+import hashlib
+
+import numpy as np
+import pandas as pd
+from collections import OrderedDict
+
+import warnings
+
+from tqdm.auto import tqdm
+
+from sklearn.discriminant_analysis import LinearDiscriminantAnalysis as LDA
+from sklearn.neighbors import NearestNeighbors
+from sklearn.metrics import silhouette_score
+
+from scipy.spatial.distance import cdist
+from scipy.stats import chi2
+from scipy.ndimage.filters import gaussian_filter1d
+
+
+
+def quality_metrics(dp, units=None, params=None, save=True, again=False):
+    """
+    Wrapper of calculate_quality_metrics to easily run on a kilosort directory.
+    
+    """
+    # preprocess parameters
+    dp = Path(dp)
+    if params is None:
+        params = {
+            "isi_threshold" : 0.0015,
+            "min_isi" : 0.000166,
+            "num_channels_to_compare" : 4,
+            "max_spikes_for_unit" : 500,
+            "max_spikes_for_nn" : 10000,
+            "n_neighbors" : 4,
+            'n_silhouette' : 10000,
+            "drift_metrics_interval_s" : 51,
+            "drift_metrics_min_spikes_per_interval" : 10,
+            "include_pc_metrics" : True,
+            }
+
+    params_keys = ['isi_threshold', 'min_isi', 'num_channels_to_compare',
+                   'max_spikes_for_unit', 'max_spikes_for_nn', 'n_neighbors',
+                   'n_silhouette']
+    params_str = str([params[k] for k in params_keys]).replace(' ', '')
+
+    spike_clusters = np.load(dp / 'spike_clusters.npy').squeeze()
+    if units is not None:
+        units = list(np.sort(units).astype(int))
+        units_hash = hashlib.sha1(str(units).encode()).hexdigest()[:20]
+    else:
+        units_hash = hashlib.sha1(str(np.unique(spike_clusters)).encode()).hexdigest()[:20]
+
+    fn = f"metrics_{params_str}_{units_hash}.csv"
+
+    # Load metrics if found
+    if (dp/fn).exists() and not again:
+        return pd.read_csv(dp/fn, index_col=0)
+
+    # Load kilosort data
+    spike_times = np.load(dp / 'spike_times.npy').squeeze()
+    spike_templates = np.load(dp / 'spike_templates.npy').squeeze()
+    amplitudes = np.load(dp / 'amplitudes.npy').squeeze()
+    channel_map = np.load(dp / 'channel_map.npy').squeeze()
+    pc_features = np.load(dp / 'pc_features.npy').squeeze()
+    pc_feature_ind = np.load(dp / 'pc_feature_ind.npy').squeeze()
+
+    # Filter unwanted units
+    if units is not None:
+        print("WARNING if you only subselect some units, \
+            relative quality metrics (such as isolation distances) cannot be interpreted.")
+        assert np.all(np.isin(units, spike_clusters))
+        units_mask = np.isin(spike_clusters, units)
+        spike_times = spike_times[units_mask]
+        spike_clusters = spike_clusters[units_mask]
+        spike_templates = spike_templates[units_mask]
+        amplitudes = amplitudes[units_mask]
+        pc_features = pc_features[units_mask,:,:]
+        pc_feature_ind = pc_feature_ind[units_mask,:]
+
+    # calculate quality metrics
+    metrics =  calculate_quality_metrics(
+                      spike_times,
+                      spike_clusters,
+                      spike_templates,
+                      amplitudes,
+                      channel_map,
+                      pc_features,
+                      pc_feature_ind,
+                      fs = 30_000,
+                      periods = None,
+                      randseed = None,
+                      **params)
+
+    if save:
+        metrics.to_csv(dp/fn)
+
+    return metrics
+
+
+
+def calculate_quality_metrics(spike_times,
+                      spike_clusters,
+                      spike_templates,
+                      amplitudes,
+                      channel_map,
+                      pc_features,
+                      pc_feature_ind,
+                      fs = 30_000,
+                      periods = None,
+                      randseed = None,
+                      **params):
+
+    """ Calculate metrics for all units on one probe
+
+    Inputs:
+    ------
+    spike_times : numpy.ndarray (num_spikes x 0)
+        Spike times in samples
+    spike_clusters : numpy.ndarray (num_spikes x 0)
+        Cluster IDs for each spike
+    spike_templates : numpy.ndarray (num_spikes x 0)
+        Original template IDs for each spike time
+    amplitudes : numpy.ndarray (num_spikes x 0)
+        Amplitude value for each spike time
+    channel_map : numpy.ndarray (num_units x 0)
+        Original data channel for pc_feature_ind array
+    pc_features : numpy.ndarray (num_spikes x num_pcs x num_channels)
+        Pre-computed PCs for blocks of channels around each spike
+        If 'None', PC-based metrics will not be computed
+    pc_feature_ind : numpy.ndarray (num_units x num_channels)
+        Channel indices of PCs for each unit
+    periods : list of recording periods in seconds [[t1, t2],...]
+
+    randseed : int to seed random number generator, optional
+    
+    params : dict of parameters
+        - "isi_threshold" : minimum time for isi violations (default 0.0015)
+        - "min_isi" : (default 0.000166)
+        - "num_channels_to_compare" : (default 7)
+        - "max_spikes_for_unit" : (default 500)
+        - "max_spikes_for_nn" : (default 10000)
+        - "n_neighbors" : (default 4)
+        - 'n_silhouette' : (default 10000)
+        - "drift_metrics_interval_s" : (default 51)
+        - "drift_metrics_min_spikes_per_interval" : (default 10)
+        - "include_pc_metrics" : (default True)
+
+    Outputs:
+    --------
+    metrics : pandas.DataFrame
+        one column for each metric
+        one row per unit per epoch
+
+    """
+
+    # preprocess parameters
+    if randseed is not None:
+        np.random.seed(randseed)
+    spike_times = spike_times/fs
+
+    metrics = pd.DataFrame()
+
+    cluster_ids = np.unique(spike_clusters)
+    total_units = len(cluster_ids)
+
+    if periods is None: periods = [[0,spike_times[-1]]]
+
+    for period in periods:
+
+        period_mask = (spike_times > period[0])&(spike_times <= period[1])
+
+        print("Calculating isi violations...")
+        isi_viol = calculate_isi_violations(spike_times[period_mask], spike_clusters[period_mask], total_units, params['isi_threshold'], params['min_isi'])
+
+        print("Calculating presence ratio...")
+        presence_ratio = calculate_presence_ratio(spike_times[period_mask], spike_clusters[period_mask], total_units)
+
+        print("Calculating firing rate...")
+        firing_rate = calculate_firing_rate(spike_times[period_mask], spike_clusters[period_mask], total_units)
+
+        print("Calculating amplitude cutoff...")
+        amplitude_cutoff = calculate_amplitude_cutoff(spike_clusters[period_mask], amplitudes[period_mask], total_units)
+
+        if pc_features is not None:
+
+            print("Calculating PC-based metrics...")
+            isolation_distance, l_ratio, d_prime, nn_hit_rate, nn_miss_rate = calculate_pc_metrics(spike_clusters[period_mask],
+                                                                                                    spike_templates[period_mask],
+                                                                                                    total_units,
+                                                                                                    pc_features[period_mask,:,:],
+                                                                                                    pc_feature_ind,
+                                                                                                    params['num_channels_to_compare'],
+                                                                                                    params['max_spikes_for_unit'],
+                                                                                                    params['max_spikes_for_nn'],
+                                                                                                    params['n_neighbors'])
+
+
+
+            print("Calculating silhouette score")
+            the_silhouette_score = calculate_silhouette_score(spike_clusters[period_mask],
+                                                        spike_templates[period_mask],
+                                                       total_units,
+                                                       pc_features[period_mask,:,:],
+                                                       pc_feature_ind,
+                                                       params['n_silhouette'])
+
+            print("Calculating drift metrics")
+            # max_drift, cumulative_drift = calculate_drift_metrics(spike_times[period_mask],
+            #                                            spike_clusters[period_mask],
+            #                                            spike_templates[period_mask],
+            #                                            total_units,
+            #                                            pc_features[period_mask,:,:],
+            #                                            pc_feature_ind,
+            #                                            params['drift_metrics_interval_s'],
+            #                                            params['drift_metrics_min_spikes_per_interval'])
+
+            max_drift, cumulative_drift = np.nan*the_silhouette_score, np.nan*the_silhouette_score
+
+        if pc_features is not None:
+            metrics = pd.concat((metrics, pd.DataFrame(data= OrderedDict((('cluster_id', cluster_ids),
+                                ('firing_rate' , firing_rate),
+                                ('presence_ratio' , presence_ratio),
+                                ('isi_viol' , isi_viol),
+                                ('amplitude_cutoff' , amplitude_cutoff),
+                                ('isolation_distance' , isolation_distance),
+                                ('l_ratio' , l_ratio),
+                                ('d_prime' , d_prime),
+                                ('nn_hit_rate' , nn_hit_rate),
+                                ('nn_miss_rate' , nn_miss_rate),
+                                ('silhouette_score', the_silhouette_score),
+                                ('max_drift', max_drift),
+                                ('cumulative_drift', cumulative_drift),
+                                ('period' , [period] * total_units),
+                                )))))
+
+        else:
+            metrics = pd.concat((metrics,
+                        pd.DataFrame(data = OrderedDict((
+                                ('cluster_id', cluster_ids),
+                                ('firing_rate' , firing_rate),
+                                ('presence_ratio' , presence_ratio),
+                                ('isi_viol' , isi_viol),
+                                ('amplitude_cutoff' , amplitude_cutoff),
+                                ('epoch_name' , [period] * total_units),
+                                ))
+                                )
+                                ))
+
+    return metrics
+
+# ===============================================================
+
+# HELPER FUNCTIONS TO LOOP THROUGH CLUSTERS:
+
+# ===============================================================
+
+def calculate_isi_violations(spike_times, spike_clusters, total_units, isi_threshold, min_isi):
+
+    cluster_ids = np.unique(spike_clusters)
+
+    viol_rates = np.zeros((total_units,))
+
+    for idx, cluster_id in enumerate(tqdm(cluster_ids)):
+
+        for_this_cluster = (spike_clusters == cluster_id)
+        viol_rates[idx], num_violations = isi_violations(spike_times[for_this_cluster],
+                                                       min_time = np.min(spike_times),
+                                                       max_time = np.max(spike_times),
+                                                       isi_threshold=isi_threshold,
+                                                       min_isi = min_isi)
+
+    return viol_rates
+
+def calculate_presence_ratio(spike_times, spike_clusters, total_units):
+
+    cluster_ids = np.unique(spike_clusters)
+
+    ratios = np.zeros((total_units,))
+
+    for idx, cluster_id in enumerate(tqdm(cluster_ids)):
+
+        for_this_cluster = (spike_clusters == cluster_id)
+        ratios[idx] = presence_ratio(spike_times[for_this_cluster],
+                                                       min_time = np.min(spike_times),
+                                                       max_time = np.max(spike_times))
+
+    return ratios
+
+
+
+def calculate_firing_rate(spike_times, spike_clusters, total_units):
+
+    cluster_ids = np.unique(spike_clusters)
+
+    firing_rates = np.zeros((total_units,))
+
+    min_time = np.min(spike_times)
+    max_time = np.max(spike_times)
+
+    for idx, cluster_id in enumerate(tqdm(cluster_ids)):
+
+        for_this_cluster = (spike_clusters == cluster_id)
+        firing_rates[idx] = firing_rate(spike_times[for_this_cluster],
+                                        min_time = np.min(spike_times),
+                                        max_time = np.max(spike_times))
+
+    return firing_rates
+
+
+def calculate_amplitude_cutoff(spike_clusters, amplitudes, total_units):
+
+    cluster_ids = np.unique(spike_clusters)
+
+    amplitude_cutoffs = np.zeros((total_units,))
+
+    for idx, cluster_id in enumerate(tqdm(cluster_ids)):
+
+        for_this_cluster = (spike_clusters == cluster_id)
+        amplitude_cutoffs[idx] = amplitude_cutoff(amplitudes[for_this_cluster])
+
+    return amplitude_cutoffs
+
+
+def calculate_pc_metrics_one_cluster(cluster_peak_channels, idx, cluster_id,cluster_ids,
+                                         half_spread, pc_features, pc_feature_ind,
+                                         spike_clusters, spike_templates,
+                                         max_spikes_for_cluster, max_spikes_for_nn, n_neighbors):
+
+    peak_channel = cluster_peak_channels[idx]
+    num_spikes_in_cluster = np.sum(spike_clusters == cluster_id)
+
+    half_spread_down = peak_channel \
+        if peak_channel < half_spread \
+        else half_spread
+
+    half_spread_up = np.max(pc_feature_ind) - peak_channel \
+        if peak_channel + half_spread > np.max(pc_feature_ind) \
+        else half_spread
+
+    channels_to_use = np.arange(peak_channel - half_spread_down, peak_channel + half_spread_up + 1)
+    units_in_range = cluster_ids[np.isin(cluster_peak_channels, channels_to_use)]
+
+    spike_counts = np.zeros(units_in_range.shape)
+
+    for idx2, cluster_id2 in enumerate(units_in_range):
+        spike_counts[idx2] = np.sum(spike_clusters == cluster_id2)
+
+    if num_spikes_in_cluster > max_spikes_for_cluster:
+        relative_counts = spike_counts / num_spikes_in_cluster * max_spikes_for_cluster
+    else:
+        relative_counts = spike_counts
+
+    all_pcs = np.zeros((0, pc_features.shape[1], channels_to_use.size))
+    all_labels = np.zeros((0,))
+
+    for idx2, cluster_id2 in enumerate(units_in_range):
+
+        subsample = int(relative_counts[idx2])
+
+        pcs = get_unit_pcs(cluster_id2, spike_clusters, spike_templates,
+                           pc_feature_ind, pc_features, channels_to_use,
+                           subsample)
+
+        if pcs is not None and len(pcs.shape) == 3:
+
+            labels = np.ones((pcs.shape[0],)) * cluster_id2
+
+            all_pcs = np.concatenate((all_pcs, pcs),0)
+            all_labels = np.concatenate((all_labels, labels),0)
+
+    all_pcs = np.reshape(all_pcs, (all_pcs.shape[0], pc_features.shape[1]*channels_to_use.size))
+    if ((all_pcs.shape[0] > 10)
+            and not (all_labels == cluster_id).all()  # Not all labels are this cluster
+            and (sum(all_labels == cluster_id) > 20)  # No fewer than 20 spikes in this cluster
+            and (len(channels_to_use) > 0)):
+        isolation_distance, l_ratio = mahalanobis_metrics(all_pcs, all_labels, cluster_id)
+
+        d_prime = lda_metrics(all_pcs, all_labels, cluster_id)
+
+        nn_hit_rate, nn_miss_rate = nearest_neighbors_metrics(all_pcs, all_labels,
+                                                                             cluster_id,
+                                                                             max_spikes_for_nn,
+                                                                             n_neighbors)
+    else:  # Too few spikes or cluster doesnt exist
+        isolation_distance = np.nan
+        d_prime = np.nan
+        nn_miss_rate = np.nan
+        nn_hit_rate = np.nan
+        l_ratio = np.nan
+    return isolation_distance, d_prime, nn_miss_rate, nn_hit_rate, l_ratio
+
+
+def calculate_pc_metrics(spike_clusters,
+                         spike_templates,
+                         total_units,
+                         pc_features,
+                         pc_feature_ind,
+                         num_channels_to_compare,
+                         max_spikes_for_cluster,
+                         max_spikes_for_nn,
+                         n_neighbors,
+                         do_parallel=True):
+    """
+
+    :param spike_clusters:
+    :param total_units:
+    :param pc_features:
+    :param pc_feature_ind:
+    :param num_channels_to_compare:
+    :param max_spikes_for_cluster:
+    :param max_spikes_for_nn:
+    :param n_neighbors:
+    :return:
+    """
+
+    assert (num_channels_to_compare % 2 == 1)
+    half_spread = int((num_channels_to_compare - 1) / 2)
+
+    cluster_ids = np.unique(spike_clusters)
+    template_ids = np.unique(spike_templates)
+
+    template_peak_channels = np.zeros((len(template_ids),), dtype='uint16')
+    cluster_peak_channels = np.zeros((len(cluster_ids),), dtype='uint16')
+
+    for idx, template_id in enumerate(template_ids):
+        for_template = np.squeeze(spike_templates == template_id)
+        pc_max = np.argmax(np.mean(pc_features[for_template, 0, :], 0))
+        template_peak_channels[idx] = pc_feature_ind[template_id, pc_max]
+
+    for idx, cluster_id in enumerate(cluster_ids):
+        for_unit = np.squeeze(spike_clusters == cluster_id)
+        templates_for_unit = np.unique(spike_templates[for_unit])
+        template_positions = np.where(np.isin(template_ids, templates_for_unit))[0]
+        cluster_peak_channels[idx] = np.median(template_peak_channels[template_positions])
+
+    # Loop over clusters:
+    if do_parallel:
+        from joblib import Parallel, delayed
+        meas = Parallel(n_jobs=10, verbose=3)(  # -1 means use all cores
+            delayed(calculate_pc_metrics_one_cluster)
+            (cluster_peak_channels, idx, cluster_id, cluster_ids,
+             half_spread, pc_features, pc_feature_ind,
+             spike_clusters, spike_templates,
+             max_spikes_for_cluster, max_spikes_for_nn, n_neighbors
+             )
+            for idx, cluster_id in enumerate(tqdm(cluster_ids)))
+    else:
+        meas = []
+        for idx, cluster_id in enumerate(tqdm(cluster_ids, total=cluster_ids.max(), desc='PC metrics')):
+            meas.append(calculate_pc_metrics_one_cluster(
+                cluster_peak_channels, idx, cluster_id, cluster_ids,
+                half_spread, pc_features, pc_feature_ind,
+                spike_clusters, spike_templates,
+                max_spikes_for_cluster, max_spikes_for_nn, n_neighbors))
+
+    # Unpack:
+    isolation_distances = []
+    l_ratios = []
+    d_primes = []
+    nn_hit_rates = []
+    nn_miss_rates = []
+    for mea in meas:
+        isolation_distance, d_prime, nn_miss_rate, nn_hit_rate, l_ratio = mea
+        isolation_distances.append(isolation_distance)
+        d_primes.append(d_prime)
+        nn_miss_rates.append(nn_miss_rate)
+        nn_hit_rates.append(nn_hit_rate)
+        l_ratios.append(l_ratio)
+
+    return (np.array(isolation_distances), np.array(l_ratios), np.array(d_primes),
+            np.array(nn_hit_rates), np.array(nn_miss_rates))
+
+
+def calculate_silhouette_score(spike_clusters,
+                               spike_templates,
+                               total_units,
+                               pc_features,
+                               pc_feature_ind,
+                               total_spikes,
+                               do_parallel=True):
+
+    def score_inner_loop(i, cluster_ids):
+        """
+        Helper to loop over cluster_ids in one dimension. We dont want to loop over both dimensions in parallel-
+        that will create too much worker overhead
+        Args:
+            i: index of first dimension
+            cluster_ids: iterable of cluster ids
+
+        Returns: scores for dimension j
+
+        """
+        scores_1d = []
+        for j in cluster_ids:
+            if j > i:
+                inds = np.in1d(cluster_labels, np.array([i, j]))
+                X = all_pcs[inds, :]
+                labels = cluster_labels[inds]
+
+                # len(np.unique(labels))=1 Can happen if total_spikes is low:
+                if (len(labels) > 2) and (len(np.unique(labels)) > 1):
+                    scores_1d.append(silhouette_score(X, labels))
+                else:
+                    scores_1d.append(np.nan)
+            else:
+                scores_1d.append(np.nan)
+        return scores_1d
+
+    cluster_ids = np.unique(spike_clusters)
+
+    random_spike_inds = np.random.permutation(spike_clusters.size)
+    random_spike_inds = random_spike_inds[:total_spikes]
+    num_pc_features = pc_features.shape[1]
+    num_channels = np.max(pc_feature_ind) + 1
+
+    all_pcs = np.zeros((total_spikes, num_channels * num_pc_features))
+
+    for idx, i in enumerate(random_spike_inds):
+
+        unit_id = spike_templates[i]
+        channels = pc_feature_ind[unit_id,:]
+
+        for j in range(0,num_pc_features):
+            all_pcs[idx, channels + num_channels * j] = pc_features[i,j,:]
+
+    cluster_labels = np.squeeze(spike_clusters[random_spike_inds])
+
+    SS = np.empty((total_units, total_units))
+    SS[:] = np.nan
+
+
+    # Build lists
+    if do_parallel:
+        from joblib import Parallel, delayed
+        scores = Parallel(n_jobs=10, verbose=2)(delayed(score_inner_loop)(i, cluster_ids) for i in tqdm(cluster_ids))
+    else:
+        scores = [score_inner_loop(i, cluster_ids) for i in tqdm(cluster_ids)]
+
+    # Fill the 2d array
+    for i, col_score in enumerate(scores):
+        for j, one_score in enumerate(col_score):
+            SS[i, j] = one_score
+
+    with warnings.catch_warnings():
+      warnings.simplefilter("ignore")
+      a = np.nanmin(SS, 0)
+      b = np.nanmin(SS, 1)
+
+    return np.array([np.nanmin([a,b]) for a, b in zip(a,b)])
+
+
+def calculate_drift_metrics(spike_times,
+                            spike_clusters,
+                            spike_templates,
+                            total_units,
+                            pc_features,
+                            pc_feature_ind,
+                            interval_length,
+                            min_spikes_per_interval,
+                            do_parallel=True):
+    def calc_one_cluster(cluster_id):
+        """
+        Helper to calculate drift for one cluster
+        Args:
+            cluster_id:
+
+        Returns:
+            max_drift, cumulative_drift
+        """
+        in_cluster = spike_clusters == cluster_id
+        times_for_cluster = spike_times[in_cluster]
+        depths_for_cluster = depths[in_cluster]
+
+        median_depths = []
+
+        for t1, t2 in zip(interval_starts, interval_ends):
+
+            in_range = (times_for_cluster > t1) * (times_for_cluster < t2)
+
+            if np.sum(in_range) >= min_spikes_per_interval:
+                median_depths.append(np.median(depths_for_cluster[in_range]))
+            else:
+                median_depths.append(np.nan)
+
+        median_depths = np.array(median_depths)
+        max_drift = np.around(np.nanmax(median_depths) - np.nanmin(median_depths), 2)
+        cumulative_drift = np.around(np.nansum(np.abs(np.diff(median_depths))), 2)
+        return max_drift, cumulative_drift
+
+
+    max_drifts = []
+    cumulative_drifts = []
+
+    depths = get_spike_depths(spike_templates, pc_features, pc_feature_ind)
+
+    interval_starts = np.arange(np.min(spike_times), np.max(spike_times), interval_length)
+    interval_ends = interval_starts + interval_length
+
+    cluster_ids = np.unique(spike_clusters)
+
+    if do_parallel:
+        from joblib import Parallel, delayed
+        meas = Parallel(n_jobs=10, verbose=2)(delayed(calc_one_cluster)(cluster_id)
+                                              for cluster_id in tqdm(cluster_ids))
+    else:
+        meas = [calc_one_cluster(cluster_id) for cluster_id in tqdm(cluster_ids)]
+
+    for max_drift, cumulative_drift in meas:
+        max_drifts.append(max_drift)
+        cumulative_drifts.append(cumulative_drift)
+    return np.array(max_drifts), np.array(cumulative_drifts)
+
+
+
+# ==========================================================
+
+# IMPLEMENTATION OF ACTUAL METRICS:
+
+# ==========================================================
+
+
+def isi_violations(spike_train, min_time, max_time, isi_threshold, min_isi=0):
+    """Calculate ISI violations for a spike train.
+
+    Based on metric described in Hill et al. (2011) J Neurosci 31: 8699-8705
+
+    modified by Dan Denman from cortex-lab/sortingQuality GitHub by Nick Steinmetz
+
+    Inputs:
+    -------
+    spike_train : array of spike times, in seconds
+    min_time : minimum time for potential spikes, in seconds
+    max_time : maximum time for potential spikes, in seconds
+    isi_threshold : threshold for isi violation, in seconds
+    min_isi : threshold for duplicate spikes, in seconds
+
+    Outputs:
+    --------
+    fpRate : rate of contaminating spikes as a fraction of overall rate
+        A perfect unit has a fpRate = 0
+        A unit with some contamination has a fpRate < 0.5
+        A unit with lots of contamination has a fpRate > 1.0
+    num_violations : total number of violations
+
+    """
+    spike_train_chunk = spike_train[(spike_train > min_time) & (spike_train < max_time)]
+    duplicate_spikes = np.where(np.diff(spike_train_chunk) <= min_isi)[0]
+    spike_train_chunk = np.delete(spike_train_chunk, duplicate_spikes + 1)
+    if len(spike_train_chunk) < 2 :
+        return 0, 0
+
+    isis = np.diff(spike_train_chunk)
+
+    num_spikes = len(spike_train_chunk)
+    num_violations = sum(isis < isi_threshold)
+    violation_time = 2*num_spikes*(isi_threshold - min_isi)
+    total_rate = mean_firing_rate(spike_train, exclusion_quantile=0.005, fs=1) #firing_rate(spike_train, min_time, max_time)
+    violation_rate = num_violations/violation_time
+    fpRate = violation_rate/total_rate
+
+    return fpRate, num_violations
+
+
+
+def presence_ratio(spike_train, min_time, max_time, num_bins=100):
+    """Calculate fraction of time the unit is present within an epoch.
+
+    Inputs:
+    -------
+    spike_train : array of spike times
+    min_time : minimum time for potential spikes
+    max_time : maximum time for potential spikes
+
+    Outputs:
+    --------
+    presence_ratio : fraction of time bins in which this unit is spiking
+
+    """
+
+    h, b = np.histogram(spike_train, np.linspace(min_time, max_time, num_bins))
+
+    return np.sum(h > 0) / num_bins
+
+
+def firing_rate(spike_train, min_time = None, max_time = None):
+    """Calculate firing rate for a spike train.
+
+    If no temporal bounds are specified, the first and last spike time are used.
+
+    Inputs:
+    -------
+    spike_train : numpy.ndarray
+        Array of spike times in seconds
+    min_time : float
+        Time of first possible spike (optional)
+    max_time : float
+        Time of last possible spike (optional)
+
+    Outputs:
+    --------
+    fr : float
+        Firing rate in Hz
+
+    """
+
+    if min_time is not None and max_time is not None:
+        duration = max_time - min_time
+    else:
+        duration = np.max(spike_train) - np.min(spike_train)
+
+    fr = spike_train.size / duration
+
+    return fr
+
+
+def amplitude_cutoff(amplitudes, num_histogram_bins = 500, histogram_smoothing_value = 3):
+
+    """ Calculate approximate fraction of spikes missing from a distribution of amplitudes
+
+    Assumes the amplitude histogram is symmetric (not valid in the presence of drift)
+
+    Inspired by metric described in Hill et al. (2011) J Neurosci 31: 8699-8705
+
+    Input:
+    ------
+    amplitudes : numpy.ndarray
+        Array of amplitudes (don't need to be in physical units)
+
+    Output:
+    -------
+    fraction_missing : float
+        Fraction of missing spikes (0-0.5)
+        If more than 50% of spikes are missing, an accurate estimate isn't possible
+
+    """
+
+
+    h,b = np.histogram(amplitudes, num_histogram_bins, density=True)
+
+    pdf = gaussian_filter1d(h,histogram_smoothing_value)
+    support = b[:-1]
+
+    peak_index = np.argmax(pdf)
+    G = np.argmin(np.abs(pdf[peak_index:] - pdf[0])) + peak_index
+
+    bin_size = np.mean(np.diff(support))
+    fraction_missing = np.sum(pdf[G:])*bin_size
+
+    fraction_missing = np.min([fraction_missing, 0.5])
+
+    return fraction_missing
+
+
+def mahalanobis_metrics(all_pcs, all_labels, this_unit_id):
+
+    """ Calculates isolation distance and L-ratio (metrics computed from Mahalanobis distance)
+
+    Based on metrics described in Schmitzer-Torbert et al. (2005) Neurosci 131: 1-11
+
+    Inputs:
+    -------
+    all_pcs : numpy.ndarray (num_spikes x PCs)
+        2D array of PCs for all spikes
+    all_labels : numpy.ndarray (num_spikes x 0)
+        1D array of cluster labels for all spikes
+    this_unit_id : Int
+        number corresponding to unit for which these metrics will be calculated
+
+    Outputs:
+    --------
+    isolation_distance : float
+        Isolation distance of this unit
+    l_ratio : float
+        L-ratio for this unit
+
+    """
+
+    pcs_for_this_unit = all_pcs[all_labels == this_unit_id,:]
+    pcs_for_other_units = all_pcs[all_labels != this_unit_id, :]
+
+    mean_value = np.expand_dims(np.mean(pcs_for_this_unit,0),0)
+
+    try:
+        VI = np.linalg.inv(np.cov(pcs_for_this_unit.T))
+    except np.linalg.linalg.LinAlgError: # case of singular matrix
+        return np.nan, np.nan
+
+    mahalanobis_other = np.sort(cdist(mean_value,
+                       pcs_for_other_units,
+                       'mahalanobis', VI = VI)[0])
+
+    mahalanobis_self = np.sort(cdist(mean_value,
+                             pcs_for_this_unit,
+                             'mahalanobis', VI = VI)[0])
+
+    n = np.min([pcs_for_this_unit.shape[0], pcs_for_other_units.shape[0]]) # number of spikes
+
+    if n >= 2:
+
+        dof = pcs_for_this_unit.shape[1] # number of features
+
+        l_ratio = np.sum(1 - chi2.cdf(pow(mahalanobis_other,2), dof)) / \
+                mahalanobis_self.shape[0] # normalize by size of cluster, not number of other spikes
+        isolation_distance = pow(mahalanobis_other[n-1],2)
+
+    else:
+        l_ratio = np.nan
+        isolation_distance = np.nan
+
+    return isolation_distance, l_ratio
+
+
+
+
+def lda_metrics(all_pcs, all_labels, this_unit_id):
+
+    """ Calculates d-prime based on Linear Discriminant Analysis
+
+    Based on metric described in Hill et al. (2011) J Neurosci 31: 8699-8705
+
+    Inputs:
+    -------
+    all_pcs : numpy.ndarray (num_spikes x PCs)
+        2D array of PCs for all spikes
+    all_labels : numpy.ndarray (num_spikes x 0)
+        1D array of cluster labels for all spikes
+    this_unit_id : Int
+        number corresponding to unit for which these metrics will be calculated
+
+    Outputs:
+    --------
+    d_prime : float
+        Isolation distance of this unit
+    l_ratio : float
+        L-ratio for this unit
+
+    """
+
+    X = all_pcs
+
+    y = np.zeros((X.shape[0],),dtype='bool')
+    y[all_labels == this_unit_id] = True
+
+    lda = LDA(n_components=1)
+
+    X_flda = lda.fit_transform(X, y)
+
+    flda_this_cluster  = X_flda[np.where(y)[0]]
+    flda_other_cluster = X_flda[np.where(np.invert(y))[0]]
+
+    d_prime = (np.mean(flda_this_cluster) - np.mean(flda_other_cluster))/np.sqrt(0.5*(np.std(flda_this_cluster)**2+np.std(flda_other_cluster)**2))
+
+    return d_prime
+
+
+
+def nearest_neighbors_metrics(all_pcs, all_labels, this_unit_id, max_spikes_for_nn, n_neighbors):
+
+    """ Calculates unit contamination based on NearestNeighbors search in PCA space
+
+    Based on metrics described in Chung, Magland et al. (2017) Neuron 95: 1381-1394
+
+    Inputs:
+    -------
+    all_pcs : numpy.ndarray (num_spikes x PCs)
+        2D array of PCs for all spikes
+    all_labels : numpy.ndarray (num_spikes x 0)
+        1D array of cluster labels for all spikes
+    this_unit_id : Int
+        number corresponding to unit for which these metrics will be calculated
+    max_spikes_for_nn : Int
+        number of spikes to use (calculation can be very slow when this number is >20000)
+    n_neighbors : Int
+        number of neighbors to use
+
+    Outputs:
+    --------
+    hit_rate : float
+        Fraction of neighbors for target cluster that are also in target cluster
+    miss_rate : float
+        Fraction of neighbors outside target cluster that are in target cluster
+
+    """
+
+    total_spikes = all_pcs.shape[0]
+    ratio = max_spikes_for_nn / total_spikes
+    this_unit = all_labels == this_unit_id
+
+    X = np.concatenate((all_pcs[this_unit,:], all_pcs[np.invert(this_unit),:]),0)
+
+    n = np.sum(this_unit)
+
+    if ratio < 1:
+        inds = np.arange(0,X.shape[0]-1,1/ratio).astype('int')
+        X = X[inds,:]
+        n = int(n * ratio)
+
+
+    nbrs = NearestNeighbors(n_neighbors=n_neighbors, algorithm='ball_tree').fit(X)
+    distances, indices = nbrs.kneighbors(X)
+
+    this_cluster_inds = np.arange(n)
+
+    this_cluster_nearest = indices[:n,1:].flatten()
+    other_cluster_nearest = indices[n:,1:].flatten()
+
+    hit_rate = np.mean(this_cluster_nearest < n)
+    miss_rate = np.mean(other_cluster_nearest < n)
+
+    return hit_rate, miss_rate
+
+# ==========================================================
+
+# HELPER FUNCTIONS:
+
+# ==========================================================
+
+def features_intersect(pc_feature_ind, these_channels):
+    """
+    # Take only the channels that have calculated features out of the ones we are interested in:
+    # This should reduce the occurence of 'except IndexError' below
+
+    Args:
+        pc_feature_ind
+        these_channels: channels_to_use or units_for_channel
+
+    Returns:
+        channels_to_use: intersect of what's available in PCs and what's needed
+    """
+    intersect = set(pc_feature_ind[these_channels[0], :])  # Initialize
+    for cluster_id2 in these_channels:
+        # Make a running intersect of what is available and what is needed
+        intersect = intersect & set(pc_feature_ind[cluster_id2, :])
+    return np.array(list(intersect))
+
+
+def get_unit_pcs(unit_id,
+                 spike_clusters,
+                 spike_templates,
+                 pc_feature_ind,
+                 pc_features,
+                 channels_to_use,
+                 subsample):
+
+    """ Return PC features for one unit
+
+    Inputs:
+    -------
+    unit_id : Int
+        ID for this unit
+    spike_clusters : np.ndarray
+        Cluster labels for each spike
+    spike_templates : np.ndarry
+        Template labels for each spike
+    pc_feature_ind : np.ndarray
+        Channels used for PC calculation for each unit
+    pc_features : np.ndarray
+        Array of all PC features
+    channels_to_use : np.ndarray
+        Channels to use for calculating metrics
+    subsample : Int
+        maximum number of spikes to return
+
+    Output:
+    -------
+    unit_PCs : numpy.ndarray (float)
+        PCs for one unit (num_spikes x num_PCs x num_channels)
+
+    """
+
+
+    inds_for_unit = np.where(spike_clusters == unit_id)[0]
+
+    spikes_to_use = np.random.permutation(inds_for_unit)[:subsample]
+
+    unique_template_ids = np.unique(spike_templates[spikes_to_use])
+
+    unit_PCs = []
+
+    for template_id in unique_template_ids:
+
+        index_mask = spikes_to_use[np.squeeze(spike_templates[spikes_to_use]) == template_id]
+        these_inds = pc_feature_ind[template_id, :]
+
+        pc_array = []
+
+        for i in channels_to_use:
+
+            if np.isin(i, these_inds):
+                channel_index = np.argwhere(these_inds == i)[0][0]
+                pc_array.append(pc_features[index_mask, :, channel_index])
+            else:
+                return None
+
+        unit_PCs.append(np.stack(pc_array, axis=-1))
+
+    if len(unit_PCs) > 0:
+
+        return np.concatenate(unit_PCs)
+    else:
+        return None
+
+
+# ===============================================================
+
+# UTILITIES
+
+# ===============================================================
+
+def get_spike_depths(spike_templates, pc_features, pc_feature_ind):
+
+    """
+    Calculates the distance (in microns) of individual spikes from the probe tip
+    This implementation is based on Matlab code from github.com/cortex-lab/spikes
+    Input:
+    -----
+    spike_templates : numpy.ndarray (N x 0)
+        Template IDs for N spikes
+    pc_features : numpy.ndarray (N x channels x num_PCs)
+        PC features for each spike
+    pc_feature_ind  : numpy.ndarray (M x channels)
+        Channels used for PC calculation for each unit
+    Output:
+    ------
+    spike_depths : numpy.ndarray (N x 0)
+        Distance (in microns) from each spike waveform from the probe tip
+    """
+
+    pc_features_copy = np.copy(pc_features)
+    pc_features_copy = np.squeeze(pc_features_copy[:,0,:])
+    pc_features_copy[pc_features_copy < 0] = 0
+    pc_power = pow(pc_features_copy, 2)
+
+    spike_feat_ind = pc_feature_ind[spike_templates, :]
+    spike_depths = np.sum(spike_feat_ind * pc_power, 1) / np.sum(pc_power,1)
+
+    return spike_depths * 10
+
+from npyx.spk_t import mean_firing_rate
\ No newline at end of file
diff --git a/npyx/plot.py b/npyx/plot.py
index 1c21f4b1..8681b36c 100644
--- a/npyx/plot.py
+++ b/npyx/plot.py
@@ -27,7 +27,7 @@
 mpl.rcParams['pdf.fonttype'] = 42 # necessary to make the text editable
 mpl.rcParams['ps.fonttype'] = 42
 
-from npyx.utils import phyColorsDic, npa, zscore, isnumeric, assert_iterable, save_np_array
+from npyx.utils import phyColorsDic, npa, zscore, isnumeric, assert_iterable, save_np_array, pprint_dic, docstring_decorator
 from npyx.stats import fractile_normal, fractile_poisson
 
 from npyx.inout import read_metadata, extract_rawChunk, assert_chan_in_dataset, chan_map, predefined_chanmap
@@ -40,27 +40,79 @@
 
 #%% plotting utilities ##############################################################################################
 
-def mplp(fig=None, ax=None, figsize=None,
+default_mplp_params = dict(
+            # title default parameters
+            title_w='regular',
+            title_s=22,
+
+            # axes labels default parameters
+            axlab_w='regular',
+            axlab_s=22,
+
+            # tick labels default parameters
+            ticklab_w='regular',
+            ticklab_s=22,
+            ticks_direction='out',
+
+            # ticks default parameters
+            xtickrot=0,
+            ytickrot=0,
+            xtickha='center',
+            xtickva='top',
+            ytickha='right',
+            ytickva='center',
+
+            # spines and layout default parameters
+            lw=1,
+            hide_top_right=True,
+            hide_axis=False,
+            tight_layout=False,
+
+            # legend default parameters
+            show_legend=False,
+            hide_legend=False,
+            legend_loc=(1,1),
+
+            # figure saving default parameters
+            saveFig=False,
+            saveDir = "~/Downloads",
+            figname="figure",
+            _format="pdf",
+
+            # colorbar default parameters
+            colorbar=False,
+            cbar_w=0.03,
+            cbar_h=0.4,
+            clabel=None,
+            clabel_w='regular',
+            clabel_s=22,
+            cticks_s=22,
+)
+
+@docstring_decorator(pprint_dic(default_mplp_params))
+def mplp(fig=None, ax=None, figsize=None, axsize=None,
          xlim=None, ylim=None, xlabel=None, ylabel=None,
-         xticks=None, yticks=None, xtickslabels=None, ytickslabels=None, reset_xticks=False, reset_yticks=False,
-         xtickrot=0, ytickrot=0, xtickha='center', xtickva='top', ytickha='right', ytickva='center',
-         axlab_w='regular', axlab_s=22,
-         ticklab_w='regular', ticklab_s=22, ticks_direction='out', lw=1,
-         title=None, title_w='bold', title_s=22,
-         hide_top_right=True, hide_axis=False,
-         tight_layout=False, hspace=None, wspace=None,
-         show_legend=False, hide_legend=False, legend_loc=(1,1),
-         saveFig=False, saveDir = "~/Downloads", figname="figure", _format="pdf",
-         colorbar=False, vmin=None, vmax=None, cmap=None, cticks=None,
-         cbar_w=0.03, cbar_h=0.4, clabel=None, clabel_w='regular', clabel_s=22, cticks_s=22,
-         axsize=None):
-    '''
+         xticks=None, yticks=None, xtickslabels=None, ytickslabels=None, reset_xticks=None, reset_yticks=None,
+         xtickrot=None, ytickrot=None, xtickha=None, xtickva=None, ytickha=None, ytickva=None,
+         axlab_w=None, axlab_s=None,
+         ticklab_w=None, ticklab_s=None, ticks_direction=None,
+         title=None, title_w=None, title_s=None,
+         lw=None, hide_top_right=None, hide_axis=None,
+         tight_layout=None, hspace=None, wspace=None,
+         show_legend=None, hide_legend=None, legend_loc=None,
+         saveFig=None, saveDir = None, figname=None, _format=None,
+         colorbar=None, vmin=None, vmax=None, cmap=None, cticks=None,
+         cbar_w=None, cbar_h=None, clabel=None, clabel_w=None, clabel_s=None, cticks_s=None,
+         prettify=True):
+    """
     make plots pretty
-    matplotlib plots
+    matplotlib plotter
 
     Awesome utility to format matplotlib plots.
     Simply add mplp() at the end of any plotting script, feeding it with your fav parameters!
 
+    IMPORTANT If you set prettify = False, it will only reset the parameters that you provide actively, and leave the rest as is.
+
     In a breeze,
         - change the weight/size/alignment/rotation of the axis labels, ticks labels, title
         - edit the x, y and colorbar axis ticks and ticks labels
@@ -73,9 +125,74 @@ def mplp(fig=None, ax=None, figsize=None,
 
     How it works: it will grab the currently active figure and axis (plt.gcf() and plt.gca()).
     Alternatively, you can pass a matplotlib figure and specific axes as arguments.
-    '''
+
+    Default Arguments:
+        {0}
+    """
+
     if fig is None: fig=plt.gcf()
     if ax is None: ax=plt.gca()
+
+    # if prettify is set to True (default),
+    # mplp() will change the plot parameters in the background,
+    # even if not actively passed
+    if prettify:
+
+        # limits default parameters
+        if xlim is None: xlim = ax.get_xlim()
+        if ylim is None: ylim = ax.get_ylim()
+
+        # title default parameters
+        if title is None: title = ax.get_title()
+        if title_w is None: title_w = default_mplp_params['title_w']
+        if title_s is None: title_s = default_mplp_params['title_s']
+
+        # axes labels default parameters
+        if ylabel is None: ylabel = ax.get_ylabel()
+        if xlabel is None: xlabel = ax.get_xlabel()
+        if axlab_w is None: axlab_w = default_mplp_params['axlab_w']
+        if axlab_s is None: axlab_s = default_mplp_params['axlab_s']
+
+        # tick labels default parameters
+        if ticklab_w is None: ticklab_w = default_mplp_params['ticklab_w']
+        if ticklab_s is None: ticklab_s = default_mplp_params['ticklab_s']
+        if ticks_direction is None: ticks_direction = default_mplp_params['ticks_direction']
+
+        # ticks default parameters
+        if xtickrot is None: xtickrot = default_mplp_params['xtickrot']
+        if ytickrot is None: ytickrot = default_mplp_params['ytickrot']
+        if xtickha is None: xtickha = default_mplp_params['xtickha']
+        if xtickva is None: xtickva = default_mplp_params['xtickva']
+        if ytickha is None: ytickha = default_mplp_params['ytickha']
+        if ytickva is None: ytickva = default_mplp_params['ytickva']
+
+        # spines and layout default parameters
+        if lw is None: lw = default_mplp_params['lw']
+        if hide_top_right is None: hide_top_right = default_mplp_params['hide_top_right']
+        if hide_axis is None: hide_axis = default_mplp_params['hide_axis']
+        if tight_layout is None: tight_layout = default_mplp_params['tight_layout']
+
+        # legend default parameters
+        if show_legend is None: show_legend = default_mplp_params['show_legend']
+        if hide_legend is None: hide_legend = default_mplp_params['hide_legend']
+        if legend_loc is None: legend_loc = default_mplp_params['legend_loc']
+
+        # figure saving default parameters
+        if saveFig is None: saveFig = default_mplp_params['saveFig']
+        if saveDir is None: saveDir = default_mplp_params['saveDir']
+        if figname is None: figname = default_mplp_params['figname']
+        if _format is None: _format = default_mplp_params['_format']
+
+        # colorbar default parameters
+        if colorbar is None: colorbar = default_mplp_params['colorbar']
+        if cbar_w is None: cbar_w = default_mplp_params['cbar_w']
+        if cbar_h is None: cbar_h = default_mplp_params['cbar_h']
+        if clabel is None: clabel = default_mplp_params['clabel']
+        if clabel_w is None: clabel_w = default_mplp_params['clabel_w']
+        if clabel_s is None: clabel_s = default_mplp_params['clabel_s']
+        if cticks_s is None: cticks_s = default_mplp_params['cticks_s']
+
+
     hfont = {'fontname':'Arial'}
     if figsize is not None:
         assert axsize is  None,\
@@ -86,85 +203,100 @@ def mplp(fig=None, ax=None, figsize=None,
         assert figsize is  None,\
             "You cannot set both the axes and figure size - the axes size is based on the figure size."
         set_ax_size(ax, *axsize)
-    # Opportunity to easily hide everything
-    if hide_axis:
-        ax.axis('off')
 
-    else: ax.axis('on')
+    # Opportunity to easily hide everything
+    if hide_axis is not None:
+        if hide_axis:
+            ax.axis('off')
+        else: ax.axis('on')
 
     # Axis labels
-    if ylabel is None:ylabel=ax.get_ylabel()
-    if xlabel is None:xlabel=ax.get_xlabel()
-    ax.set_ylabel(ylabel, weight=axlab_w, size=axlab_s, **hfont)
-    ax.set_xlabel(xlabel, weight=axlab_w, size=axlab_s, **hfont)
+    if ylabel is not None: ax.set_ylabel(ylabel, weight=axlab_w, size=axlab_s, **hfont)
+    if xlabel is not None: ax.set_xlabel(xlabel, weight=axlab_w, size=axlab_s, **hfont)
 
-    # Setup limits BEFORE altering the ticks
+    # Setup x/y limits BEFORE altering the ticks
     # since the limits will alter the ticks
-    if xlim is None: xlim=ax.get_xlim()
-    if ylim is None: ylim=ax.get_ylim()
-    ax.set_xlim(xlim)
-    ax.set_ylim(ylim)
+    if xlim is not None: ax.set_xlim(xlim)
+    if ylim is not None: ax.set_ylim(ylim)
 
     # Tick values
-    if xticks is None:
+    if prettify and xticks is None:
         if reset_xticks:
             ax.xaxis.set_major_locator(AutoLocator())
-        xticks=ax.get_xticks()
-        ax.set_xticks(xticks)
-    else: ax.set_xticks(xticks)
-    if yticks is None:
+        xticks = ax.get_xticks()
+    if xticks is not None: ax.set_xticks(xticks)
+
+    if prettify and yticks is None:
         if reset_yticks:
             ax.yaxis.set_major_locator(AutoLocator())
-        yticks=ax.get_yticks()
-        ax.set_yticks(yticks)
-    else: ax.set_yticks(yticks)
+        yticks = ax.get_yticks()
+    if yticks is not None: ax.set_yticks(yticks)
 
-    # Tick labels and x/y limits
+    # Tick labels
     fig.canvas.draw() # To force setting of ticklabels
-    if xtickslabels is None:
+    if xtickslabels is None and prettify:
         if any(ax.get_xticklabels()):
-            if isnumeric(ax.get_xticklabels()[0].get_text()): xtickslabels,x_nflt=get_labels_from_ticks(xticks)
-            else: xtickslabels = ax.get_xticklabels()
-    if ytickslabels is None:
+            if isnumeric(ax.get_xticklabels()[0].get_text()):
+                xtickslabels, x_nflt = get_labels_from_ticks(xticks)
+            else:
+                xtickslabels = ax.get_xticklabels()
+    if ytickslabels is None and prettify:
         if any(ax.get_yticklabels()):
-            if isnumeric(ax.get_yticklabels()[0].get_text()): ytickslabels,y_nflt=get_labels_from_ticks(yticks)
-            else: ytickslabels = ax.get_yticklabels()
-    ax.set_xlim(xlim)
-    ax.set_ylim(ylim)
+            if isnumeric(ax.get_yticklabels()[0].get_text()):
+                ytickslabels, y_nflt = get_labels_from_ticks(yticks)
+            else:
+                ytickslabels = ax.get_yticklabels()
+
     if xtickslabels is not None:
-        assert len(xtickslabels)==len(xticks), 'WARNING you provided too many/few xtickslabels! Make sure that the default/provided xticks match them.'
+        if xticks is not None:
+            assert len(xtickslabels)==len(xticks), 'WARNING you provided too many/few xtickslabels! Make sure that the default/provided xticks match them.'
+        if xtickha is None: xtickha = ax.xaxis.get_ticklabels()[0].get_ha()
+        if xtickva is None: xtickva = ax.xaxis.get_ticklabels()[0].get_va()
         ax.set_xticklabels(xtickslabels, fontsize=ticklab_s, fontweight=ticklab_w, color=(0,0,0), **hfont, rotation=xtickrot, ha=xtickha, va=xtickva)
     if ytickslabels is not None:
-        assert len(ytickslabels)==len(yticks), 'WARNING you provided too many/few ytickslabels! Make sure that the default/provided yticks match them.'
+        if yticks is not None:
+            assert len(ytickslabels)==len(yticks), 'WARNING you provided too many/few ytickslabels! Make sure that the default/provided yticks match them.'
+        if ytickha is None: ytickha = ax.yaxis.get_ticklabels()[0].get_ha()
+        if ytickva is None: ytickva = ax.yaxis.get_ticklabels()[0].get_va()
         ax.set_yticklabels(ytickslabels, fontsize=ticklab_s, fontweight=ticklab_w, color=(0,0,0), **hfont, rotation=ytickrot, ha=ytickha, va=ytickva)
 
+    # Reset x/y limits a second time
+    if xlim is not None: ax.set_xlim(xlim)
+    if ylim is not None: ax.set_ylim(ylim)
+
     # Title
-    if title is None: title=ax.get_title()
-    ax.set_title(title, size=title_s, weight=title_w)
+    if title is not None: ax.set_title(title, size=title_s, weight=title_w)
 
     # Ticks and spines aspect
-    ax.tick_params(axis='both', bottom=1, left=1, top=0, right=0, width=lw, length=4, direction=ticks_direction)
-    spine_keys = list(ax.spines.keys())
-    hide_spine_keys = ['polar'] if 'polar' in spine_keys else ['top', 'right']
-    lw_spine_keys = ['polar'] if 'polar' in spine_keys else ['left', 'bottom', 'top', 'right']
-    if hide_top_right and 'top' in hide_spine_keys: [ax.spines[sp].set_visible(False) for sp in hide_spine_keys]
-    else: [ax.spines[sp].set_visible(True) for sp in hide_spine_keys]
-    for sp in lw_spine_keys:
-        ax.spines[sp].set_lw(lw)
-
-
-    # Alignement and spacing elements
-    if tight_layout:fig.tight_layout(rect=[0, 0.03, 1, 0.95])
-    if (hspace is not None): fig.subplots_adjust(hspace=hspace)
-    if (wspace is not None): fig.subplots_adjust(wspace=wspace)
-    axis_to_align = [AX for AX in fig.axes if 'AxesSubplot' in AX.__repr__()]
-    fig.align_ylabels(axis_to_align)
-    fig.align_xlabels(axis_to_align)
+    if prettify:
+        ax.tick_params(axis='both', bottom=1, left=1, top=0, right=0, width=lw, length=4, direction=ticks_direction)
+    elif lw is not None or ticks_direction is not None:
+        ax.tick_params(axis='both', width=lw, direction=ticks_direction)
+
+    if hide_top_right is not None:
+        spine_keys = list(ax.spines.keys())
+        hide_spine_keys = ['polar'] if 'polar' in spine_keys else ['top', 'right']
+        lw_spine_keys = ['polar'] if 'polar' in spine_keys else ['left', 'bottom', 'top', 'right']
+        if hide_top_right and 'top' in hide_spine_keys: [ax.spines[sp].set_visible(False) for sp in hide_spine_keys]
+        else: [ax.spines[sp].set_visible(True) for sp in hide_spine_keys]
+        for sp in lw_spine_keys:
+            ax.spines[sp].set_lw(lw)
+
+
+    # Aligning and spacing axes and labels
+    if tight_layout: fig.tight_layout(rect=[0, 0.03, 1, 0.95])
+    if hspace is not None: fig.subplots_adjust(hspace=hspace)
+    if wspace is not None: fig.subplots_adjust(wspace=wspace)
+    if prettify:
+        axis_to_align = [AX for AX in fig.axes if 'AxesSubplot' in AX.__repr__()]
+        fig.align_ylabels(axis_to_align)
+        fig.align_xlabels(axis_to_align)
 
     assert not (show_legend and hide_legend),\
         "You instructed to both show and hide the legend...?"
-    assert len(legend_loc)==2 or len(legend_loc)==4,\
-        "legend_loc must comply to the bbox_to_anchor format ( (x,y) or (x,y,width,height))."
+    if legend_loc is not None:
+        assert len(legend_loc)==2 or len(legend_loc)==4,\
+            "legend_loc must comply to the bbox_to_anchor format ( (x,y) or (x,y,width,height))."
     if show_legend: plt.legend(bbox_to_anchor=legend_loc)
     elif hide_legend: plt.legend([],[], frameon=False)
 
@@ -232,7 +364,7 @@ def hvshow(hvobject, backend='matplotlib', return_mpl=False):
     - for static instanciation, refinement and data exploitation:
         in matplotlib current backend, pops up a holoview object as a matplotlib figure
         and eventually returns it for further tweaking.
-    Parameters:
+    Arguments:
         - hvobject: a Holoviews object e.g. Element, Overlay or Layout.
         - backend: 'bokeh' or 'matplotlib', which backend to use to show figure
         - return_mpl: bool, returns a matplotlib figure
@@ -529,7 +661,28 @@ def set_ax_size(ax,w,h):
 def hist_MB(arr, a=None, b=None, s=None,
             title='Histogram', xlabel='', ylabel='',
             ax=None, color=None, alpha=1, figsize=None, xlim=None,
-            saveFig=False, saveDir='', _format='pdf'):
+            saveFig=False, saveDir='', _format='pdf', prettify=True, **mplp_kwargs):
+    """
+    Plot histogram of array arr.
+    Arguments:
+        - arr: array, data to plot
+        - a: float, lower bound of histogram
+        - b: float, upper bound of histogram
+        - s: float, bin size
+        - title: str, title of plot
+        - xlabel: str, label of x axis
+        - ylabel: str, label of y axis
+        - ax: matplotlib axis, axis to plot on (new figure is created if none is provided)
+        - color: str, color of bars
+        - alpha: float, opacity of bars
+        - figsize: tuple, size of figure
+        - xlim: tuple, limits of x axis
+        - saveFig: bool, whether to save figure or not
+        - saveDir: str, directory to save figure to
+        - _format: str, format to save figure to
+        - prettify: bool, whether to apply mplp() prettification or not
+        - **mplp_kwargs: any additional formatting parameters, passed to mplp()
+    """
     if a is None: a=np.min(arr)
     if b is None: b=np.max(arr)
     if s is None: s=(b-a)/100
@@ -546,7 +699,7 @@ def hist_MB(arr, a=None, b=None, s=None,
     ax.set_ylabel(ylabel) if len(ylabel)>0 else ax.set_ylabel('Counts')
 
     if xlim is None: xlim = [a-s/2,b+s/2]
-    fig,ax=mplp(fig, ax, xlim=xlim, figsize=figsize)
+    fig, ax = mplp(fig, ax, xlim=xlim, figsize=figsize, prettify=prettify, **mplp_kwargs)
 
     if saveFig: save_mpl_fig(fig, title, saveDir, _format)
       
@@ -554,10 +707,21 @@ def hist_MB(arr, a=None, b=None, s=None,
 
 def paired_plot(df, cats, ylabel=None, xlabel=None, color="grey",
                dodge=True, dodge_range=0.1,
-               figsize=(3,5)):
+               figsize=(3,5), prettify=True, **mplp_kwargs):
     """
-    df: pandas dataframe
-    cats: list of str, columns of dataframe df
+    Plot scatter plot of each column of dataframe df against the first column,
+    in a paired-plot fashion (rows are linked by a line).
+    Arguments:
+        - df: pandas dataframe
+        - cats: list of str, columns of dataframe df
+        - ylabel: str, label of y axis
+        - xlabel: str, label of x axis
+        - color: str, color of points
+        - dodge: bool, whether to randomly offset points so that they do not overlap
+        - dodge_range: float, range of random offset
+        - figsize: (x,y) tuple, size of figure in inches
+        - prettify: bool, whether to apply mplp() prettification or not
+        - **mplp_kwargs: any additional formatting parameters, passed to mplp()
     """
 
     plt.figure()
@@ -580,19 +744,30 @@ def paired_plot(df, cats, ylabel=None, xlabel=None, color="grey",
     plt.plot(lines_x, lines, color=color, alpha=0.5, zorder=-1)
 
     mplp(figsize=figsize, ylabel=ylabel, xlabel=xlabel,
-         xlim=[-0.5, i+0.5], xticks=np.arange(i+1), xtickslabels=cats)
+         xlim=[-0.5, i+0.5], xticks=np.arange(i+1), xtickslabels=cats,
+         prettify=prettify, **mplp_kwargs)
 
 
 #%% Stats plots ##############################################################################################
 
 def plot_pval_borders(Y, p, dist='poisson', Y_pred=None, gauss_baseline_fract=1, x=None, ax=None, color=None,
-                      ylabel=None, xlabel=None, title=None):
+                      ylabel=None, xlabel=None, title=None, prettify=True, **mplp_kwargs):
     '''
     Function to plot array X and the upper and lower borders for a given p value.
-    Parameters:
-        - X: np array
-        - p:float, p value [0-1]
+    Arguments:
+        - Y: np array
+        - p:float, p value to plot threshold [0-1]
         - dist: whether to assume Poisson or Normal distribution
+        - Y_pred: np array or same size as Y, predictor for distribution (If none is provided, the mean of X is used)
+        - gauss_baseline_fract: float, fraction of data to use as baseline for normal distribution
+        - x: np array, x axis values
+        - ax: matplotlib axis, axis to plot on (new figure is created if none is provided)
+        - color: str, color of bars
+        - ylabel: str, label of y axis
+        - xlabel: str, label of x axis
+        - title: str, title of plot
+        - prettify: bool, whether to apply mplp() prettification or not
+        - **mplp_kwargs: any additional formatting parameters, passed to mplp()
     '''
     Y=npa(Y)
     assert 0<p<1
@@ -617,7 +792,8 @@ def plot_pval_borders(Y, p, dist='poisson', Y_pred=None, gauss_baseline_fract=1,
     ax.plot(x,fp2, c='r', ls='--')
     ax.legend(fontsize=14)
 
-    fig, ax = mplp(fig, ax, ylabel=ylabel, xlabel=xlabel, title=title)
+    fig, ax = mplp(fig, ax, ylabel=ylabel, xlabel=xlabel,
+                   title=title, prettify=prettify, **mplp_kwargs)
 
     return fig
 
@@ -695,7 +871,7 @@ def plot_wvf(dp, u=None, Nchannels=12, chStart=None, n_waveforms=300, t_waveform
              as_heatmap=False, use_dsmatch=False, verbose=False):
     '''
     To plot main channel alone: use Nchannels=1, chStart=None
-    Parameters:
+    Arguments:
         - dp: string, datapath to kilosort directory
         - u: int, unit index
         - Nchannels: int, number of channels where waveform is plotted
@@ -882,7 +1058,7 @@ def plt_wvf(waveforms, subcm=None, waveforms_std=None,
                 xvalues=x, yvalues=subcm[::-1,0], xticks=hm_xticks, yticks=hm_yticks,
                 xticklabels=hm_xticks, yticklabels=hm_yticks, xlabel='Time (ms)', ylabel='Channel', xtickrot=0, title=title,
                 cmapstr="RdBu_r", vmin=ylim1*0.5, vmax=ylim2*0.5, center=0, colorseq='linear',
-                clabel='Voltage (\u03bcV)', extend_cmap='neither', cticks=None,
+                clabel='Voltage (\u03bcV)', cticks=None,
                 figsize=(figw_inch/2,figh_inch), aspect='auto', function='imshow',
                 ax=None)
     else:
@@ -1018,7 +1194,7 @@ def plot_raw(dp, times=None, alignement_events=None, window=None, channels=np.ar
     '''
     Plot raw data over a specified window of time, over a specified range of channels.
 
-    Parameters:
+    Arguments:
     - dp: binary path (files must ends in .bin, typically ap.bin)
     - times: list of boundaries of the time window, in seconds [t1, t2].
     - alignement_events: list of events to align the stimulus to compute an average, in seconds
@@ -1036,7 +1212,7 @@ def plot_raw(dp, times=None, alignement_events=None, window=None, channels=np.ar
     - saveDir: directory where to save either the figure or the data (default: ~/Downloads)
     - saveFig: save the figure at saveDir
     - _format: format of the figure to save | default: pdf
-    - figsize: (x_inches, y_inches) figure size
+    - figsize: (x,y) tuple, size of figure in inches
     - saveData: bool, whether to save data used to make the plot (n_channels x time array, where time is in samples (30kHz))
     - again: bool, whether to recompute data rather than loading it from disc
 
@@ -1153,7 +1329,7 @@ def plot_raw(dp, times=None, alignement_events=None, window=None, channels=np.ar
                         xvalues=None, yvalues=None, xticks=xticks-xticks[0], yticks=y_ticks,
                         xticklabels=xticklabels, yticklabels=y_ticks_labels, xlabel=None, ylabel=None,
                         cmapstr="RdBu_r", vmin=vmin, vmax=vmax, center=center, colorseq='nonlinear',
-                        clabel='Voltage (\u03BCV)', extend_cmap='neither', cticks=None,
+                        clabel='Voltage (\u03BCV)', cticks=None,
                         figsize=(4,10), aspect='auto', function='imshow', ax=None)
         ax=fig.axes[0]
         ax.set_ylabel('Depth (\u03BCm)', size=14, weight='bold')
@@ -1220,7 +1396,7 @@ def plot_raw_units(dp, times, units=[], channels=np.arange(384), offset=450,
     f'''
     Plot raw traces with colored overlaid spike times of specified units.
 
-    Parameters:
+    Arguments:
     - most parameters from plot_raw (see below)
     - units: list/array of units (if they do not spike within 'times', will be ignored)
     - Nchan_plot: int, number of channels over which to plot colored unit spikes
@@ -1359,7 +1535,7 @@ def psth_plot(times, events, psthb=5, psthw=[-1000, 1000], remove_empty_trials=T
            saveDir='~/Downloads', saveFig=0, ret_data=0, _format='pdf',
            zscore=False, bsl_subtract=False, bsl_window=[-2000,-1000], ylim=None,
            convolve=True, gsd=2, xticks=None, xticklabels=None, xlabel='Time (ms)', ylabel=None,
-           ax=None, figsize=None, tight_layout=True, hspace=None, wspace=None):
+           ax=None, figsize=None, tight_layout=True, hspace=None, wspace=None, prettify=True, **mplp_kwargs):
 
     x, y, y_p, y_p_var = get_processed_ifr(times, events, b=psthb, window=psthw, remove_empty_trials=remove_empty_trials,
                                       zscore=zscore, zscoretype='within',
@@ -1372,7 +1548,7 @@ def psth_plot(times, events, psthb=5, psthw=[-1000, 1000], remove_empty_trials=T
            zscore, bsl_subtract, ylim,
            convolve, xticks, xticklabels,
            xlabel, ylabel, legend_label, legend,
-           ax, figsize, tight_layout, hspace, wspace)
+           ax, figsize, tight_layout, hspace, wspace, prettify, **mplp_kwargs)
 
     return (x,y,y_p,y_p_var) if ret_data else fig
 
@@ -1382,8 +1558,40 @@ def psth_plt(x, y_p, y_p_var, psthw, events_toplot=[0], events_color='r',
            zscore=False, bsl_subtract=False, ylim=None,
            convolve=True, xticks=None, xticklabels=None,
            xlabel='Time (ms)', ylabel='IFR (spk/s)', legend_label=None, legend=False,
-           ax=None, figsize=None, tight_layout=True, hspace=None, wspace=None):
-
+           ax=None, figsize=None, tight_layout=True, hspace=None, wspace=None,
+           prettify=True, **mplp_kwargs):
+    """
+    Plots peri-event PSTHs
+    Arguments:
+        - x: time vector
+        - y_p: mean PSTH
+        - y_p_var: std of PSTH
+        - psthw: peri-stimulus time window
+        - events_toplot: list of event indices to plot
+        - events_color: color of event lines
+        - title: plot title
+        - color: color of PSTH
+        - saveDir: directory to save figure
+        - saveFig: bool, whether to save figure or not
+        - _format: format to save figure in
+        - zscore: bool, PSTH was zscored or not
+        - bsl_subtract: bool, whether to subtract baseline from PSTH or not
+        - ylim: y-axis limits
+        - convolve: bool, whether PSTH was convolved or not
+        - xticks: x-axis ticks
+        - xticklabels: x-axis tick labels
+        - xlabel: x-axis label
+        - ylabel: y-axis label
+        - legend_label: label for legend
+        - legend: bool, whether to plot legend or not
+        - ax: axis to plot on
+        - figsize: figure size
+        - tight_layout: bool, whether to apply tight_layout() or not
+        - hspace: horizontal space between subplots
+        - wspace: vertical space between subplots
+        - prettify: bool, whether to apply mplp() prettification or not
+        - **mplp_kwargs: any additional formatting parameters, passed to mplp()
+    """
     if ax is None:
         fig, ax = plt.subplots()
     else:
@@ -1433,7 +1641,8 @@ def psth_plt(x, y_p, y_p_var, psthw, events_toplot=[0], events_color='r',
      axlab_w='bold', axlab_s=16,
      ticklab_w='regular',ticklab_s=14, lw=1,
      title=title, title_w='bold', title_s=16,
-     hide_top_right=True, tight_layout=tight_layout, hspace=hspace, wspace=wspace)
+     hide_top_right=True, tight_layout=tight_layout, hspace=hspace, wspace=wspace,
+     prettify=prettify, **mplp_kwargs)
 
     if saveFig:
         figname=title
@@ -1449,27 +1658,55 @@ def raster_plot(times, events, window=[-1000, 1000], events_toplot=[0], events_c
                 as_heatmap=False, vmin=None, center=None, vmax=None, cmap_str=None,
                 show_psth=False, psthb=10,
                 zscore=False, bsl_subtract=False, bsl_window=[-2000,-1000], ylim_psth=None,
-                convolve=True, gsd=2):
+                convolve=True, gsd=2, prettify=True, **mplp_kwargs):
     '''
     Make a raster plot of the provided 'times' aligned on the provided 'events', from window[0] to window[1].
     By default, there will be len(events) lines. you can pick a subset of events to plot
     by providing their indices as a list.array with 'events_toplot'.
 
-    Parameters:
+    Arguments:
         - times: list/array of spike times, in seconds. If list of lists/arrays,
                  each item of the list is considered an individual spike train.
         - events: list/array of events, in seconds. TRIALS WILL BE PLOTTED ACORDING TO EVENTS ORDER.
-        - events_toplot: list/array of events indices to display on the raster | Default: None (plots everything)
         - window: list/array of shape (2,): the raster will be plotted from events-window[0] to events-window[1] | Default: [-1000,1000]
+        - events_toplot: list/array of events indices to display on the raster | Default: None (plots everything)
+        - events_color: string or list of strings of same size as events_toplot (in cases of several events)
+        - trials_toplot: list/array of trials indices to display on the raster | Default: None (plots everything)
         - remove_empty_trials: boolean, if True does not use empty trials to compute psth
+
         - title: string, title of the plot + if saved file name will be raster_title._format.
-        - color: string or list of strings of same size as times (in cases with several cells)
-        - figsize: tuple, (x,y) figure size
+        - color: string or list of strings of same length as 'times' (in cases of several cells)
+        - palette: string, name of the color palette to use instead of directly passing colors | Default: 'batlow'
+        - marker: string, marker to use for the raster | Default: '|'
+        - malpha: float, opacity for the raster | Default: 0.9
+        - size: float, size of the raster | Default: None (uses lw)
+        - lw: float, line width of the raster | Default: 3
+        - sparseylabels: boolean, if True, only displays ylabels for the first and last trials | Default: True
+
+        - figsize: tuple, (x,y) figure size in inches
         - saveDir: save directory to save data and figure
         - saveFig: boolean, if 1 saves figure with name raster_title._format at saveDir
         - ret_data: boolean, whether to return data (x,y,y_p,y_p_var) instead of matplotlib figure.
         - _format: string, format used to save figure if saveFig=1 | Default: 'pdf'
 
+        - as_heatmap: boolean, if True, plots a heatmap instead of a raster | Default: False
+        - vmin: float, minimum value for the heatmap | Default: None (uses min of data)
+        - center: float, center value for the heatmap | Default: None (uses mean of data)
+        - vmax: float, maximum value for the heatmap | Default: None (uses max of data)
+        - cmap_str: string, name of the colormap to use | Default: None (uses 'viridis' if as_heatmap else 'Greys')
+
+        - show_psth: boolean, if True, plots a psth below the raster | Default: False
+        - psthb: float, bin size for the psth | Default: 10
+        - zscore: boolean, if True, zscores the psth | Default: False
+        - bsl_subtract: boolean, if True, subtracts the baseline from the psth | Default: False
+        - bsl_window: list/array of shape (2,): baseline window for the psth | Default: [-2000,-1000]
+        - ylim_psth: list/array of shape (2,): y limits for the psth | Default: None (uses min/max of data)
+        - convolve: boolean, if True, convolves the psth with a gaussian kernel | Default: True
+        - gsd: float, gaussian standard deviation for the psth | Default: 2
+
+        - prettify: bool, whether to apply mplp() prettification or not
+        - **mplp_kwargs: any additional formatting parameters, passed to mplp()
+
     Returns:
         - fig: matplotlib figure.
     '''
@@ -1545,7 +1782,7 @@ def raster_plot(times, events, window=[-1000, 1000], events_toplot=[0], events_c
                         xticks=xticks, yticks=y_ticks,
                         xticklabels=None, yticklabels=y_ticks_labels, xlabel=xlabel_plot, ylabel='Trials', title=title,
                         cmapstr=cmap_str, vmin=vmin, vmax=vmax, center=center, colorseq='nonlinear',
-                        clabel=clab, extend_cmap='neither', cticks=None,
+                        clabel=clab, cticks=None,
                         figsize=figsize, aspect='auto', function='imshow', ax=ax)
 
         else:
@@ -1563,7 +1800,8 @@ def raster_plot(times, events, window=[-1000, 1000], events_toplot=[0], events_c
                  axlab_w='bold', axlab_s=20,
                  ticklab_w='regular',ticklab_s=16, lw=1,
                  title=title, title_w='bold', title_s=24,
-                 hide_top_right=True, hide_axis=False)
+                 hide_top_right=True, hide_axis=False,
+                 prettify=prettify, **mplp_kwargs)
     print(f'{ntrials} trials.')
     xl=ax.get_xlim()
     yl=ax.get_ylim()
@@ -1590,7 +1828,8 @@ def raster_plot(times, events, window=[-1000, 1000], events_toplot=[0], events_c
                        zscore=zscore, bsl_subtract=bsl_subtract, bsl_window=bsl_window, ylim=ylim_psth,
                        convolve=convolve, gsd=gsd,
                        xticks=xticks, xticklabels=xticklabels_subplot, xlabel=xlabel_subplot, ylabel=None,
-                       ax=ax_psth, figsize=None, tight_layout=True, hspace=None, wspace=None)
+                       ax=ax_psth, figsize=None, tight_layout=True, hspace=None, wspace=None,
+                       prettify=True, **mplp_kwargs)
 
     if saveFig:
         figname=title
@@ -1607,7 +1846,7 @@ def summary_psth(trains, trains_str, events, events_str, psthb=5, psthw=[-1000,1
     '''
     Function to plot a bunch of PSTHs, all trains aligned to all sets of events, in a single summary figure.
 
-    Parameters:
+    Arguments:
         Related to PSTH data:
             - trains: list of np arrays (s), spike trains
             - trains_str: list of str, name of trains units
@@ -1768,7 +2007,7 @@ def summary_psth(trains, trains_str, events, events_str, psthb=5, psthw=[-1000,1
                     xticklabels=None, yticklabels=y_ticks_labels, xlabel=xlab, xtickrot=0,
                     ylabel=ylab, title=None,
                     cmapstr=cmap_str, vmin=vmin1, vmax=vmax1, center=center1, colorseq='nonlinear',
-                    clabel=clab, extend_cmap='neither', cticks=None,
+                    clabel=clab, cticks=None,
                     figsize=figsize, aspect='auto', function='imshow', ax=ax_im, tight_layout=False,
                     cmap_h=0.6/nmaps)
 
@@ -1782,15 +2021,38 @@ def summary_psth(trains, trains_str, events, events_str, psthb=5, psthw=[-1000,1
 def plt_ccg(uls, CCG, cbin=0.04, cwin=5, bChs=None, fs=30000, saveDir='~/Downloads', saveFig=True,
             _format='pdf', labels=True, title=None, color=None,
             ylim=None, normalize='Hertz', ccg_mn=None, ccg_std=None,
-            figsize=None, show_hz=False, style='line', hide_axis=False, show_ttl=True):
-    '''Plots acg and saves it given the acg array.
-    unit: int.
-    ACG: acg array in non normalized counts.
-    cwin and cbin: full window and bin in ms.
-    phycolor: index (0 to 5) of the phy colorchart.
-    savedir: plot saving destination.
-    save: boolean, to save the figure or not.
-    '''
+            figsize=None, show_hz=False, style='line', hide_axis=False, show_ttl=True,
+            prettify=True, **mplp_kwargs):
+    """
+    Plots precomputed crosscorrelogram between a pair of units.
+    Arguments:
+        - uls: list of unit indices
+        - CCG: np array, crosscorrelogram (n_bins,) array
+        - cbin: float, bin size in ms
+        - cwin: float, full window size in ms
+        - bChs: list of peak channels for each unit
+        - fs: int, sampling frequency (to convert cbincorreclty)
+
+        - saveDir: str, directory to save figure
+        - saveFig: bool, whether to save figure
+        - _format: str, format to save figure
+
+        - labels: bool, whether to show figure labels and splines
+        - title: title of figure
+        - color: color of CCG
+        - ylim: list of floats, y axis limits
+        - normalize: str, normalization method (Hertz, Counts, Pearson or zscore)
+        - ccg_mn: np array, mean CCG
+        - ccg_std: np array, std CCG
+        - figsize: tuple of floats, figure size
+        - show_hz: bool, whether to show Hz on a second y axis
+        - style: str 'line' or 'bar', style of plot
+        - hide_axis: bool, whether to hide axis
+        - show_ttl: bool, whether to show title
+
+        - prettify: bool, whether to apply mplp() prettification or not
+        - **mplp_kwargs: any additional formatting parameters, passed to mplp()
+    """
     global phyColorsDic
     assert style in ['line', 'bar']
 
@@ -1855,10 +2117,9 @@ def plt_ccg(uls, CCG, cbin=0.04, cwin=5, bChs=None, fs=30000, saveDir='~/Downloa
     mplp(fig, ax, figsize=figsize,
          title=title,
          xlabel='Time (ms)', ylabel=ylabel,
-         #xticks = get_bestticks(-cwin/2, cwin/2, light=True),
-         #yticks = get_bestticks(ylim[0], ylim[1], light=True),
          title_s=20, axlab_s=20, ticklab_s=20, axlab_w = 'regular',
-         xlim=[-cwin*1./2, cwin*1./2], ylim=ylim, hide_axis=hide_axis)
+         xlim=[-cwin*1./2, cwin*1./2], ylim=ylim, hide_axis=hide_axis,
+         prettify=prettify, **mplp_kwargs)
 
     # optional second y axis
     if ccg_mn is not None and ccg_std is not None and show_hz:
@@ -1877,21 +2138,48 @@ def plt_ccg(uls, CCG, cbin=0.04, cwin=5, bChs=None, fs=30000, saveDir='~/Downloa
 
     return fig
 
-def plt_acg(unit, ACG, cbin=0.2, cwin=80, bChs=None, color=0, fs=30000, saveDir='~/Downloads', saveFig=False,
-            _format='pdf', labels=True, title=None, ref_per=True,
+def plt_acg(unit, ACG, cbin=0.2, cwin=80, bChs=None, color=0, fs=30000,
+            saveDir='~/Downloads', saveFig=False, _format='pdf',
+            labels=True, title=None, ref_per=True, ax=None,
             ylim1=0, ylim2=0, normalize='Hertz', acg_mn=None, acg_std=None, figsize=None,
-            hide_axis=False):
-    '''Plots acg and saves it given the acg array.
-    unit: int.
-    ACG: acg array in non normalized counts.
-    cwin and cbin: full window and bin in ms.
-    phycolor: index (0 to 5) of the phy colorchart.
-    savedir: plot saving destination.
-    saveFig: boolean, to save the figure or not.
-    '''
+            hide_axis=False, prettify=True, **mplp_kwargs):
+    """
+    Plots precomputed autocorrelogram.
+    Arguments:
+        - unit: unit index
+        - ACG: np array, crosscorrelogram (n_bins,) array
+        - cbin: float, bin size in ms
+        - cwin: float, full window size in ms
+        - bChs: peak channels of unit
+
+        - color: color of ACG
+        - fs: int, sampling frequency (to convert cbincorreclty)
+
+        - saveDir: str, directory to save figure
+        - saveFig: bool, whether to save figure
+        - _format: str, format to save figure
+
+        - labels: bool, whether to show figure labels and splines
+        - title: title of figure
+        - ref_per: bool, whether to plot refractory period (+/- 1ms)
+        - ax: matplotlib axis, axis to plot on (if None, creates new figure)
+        - ylim1: float, lower y axis limit
+        - ylim2: float, upper y axis limit
+        - normalize: str, normalization method (Hertz, Counts, Pearson or zscore)
+        - acg_mn: float, mean of ACG if was normalized
+        - acg_std: np array, std of ACG if was normalized
+
+        - figsize: tuple of floats (x, y), figure size in inches
+        - hide_axis: bool, whether to hide axis
+
+        - prettify: bool, whether to apply mplp() prettification or not
+        - **mplp_kwargs: any additional formatting parameters, passed to mplp()
+    """
     global phyColorsDic
-
-    fig, ax = plt.subplots()
+    if ax is None:
+        fig, ax = plt.subplots()
+    else:
+        fig = ax.figure
 
     # x axis
     cbin = np.clip(cbin, 1000*1./fs, 1e8)
@@ -1952,11 +2240,10 @@ def plt_acg(unit, ACG, cbin=0.2, cwin=80, bChs=None, color=0, fs=30000, saveDir=
 
     if figsize is None: figsize = (4.5,4)
     mplp(fig, figsize=figsize,
-         title=title,
-         xlabel='Time (ms)', ylabel=ylabel,
-         title_s=20, axlab_s=20, ticklab_s=20,
-         xlim=[-cwin*1./2, cwin*1./2], ylim=[ylim1, ylim2], hide_axis=hide_axis)
-
+    title=title, xlabel='Time (ms)', ylabel=ylabel,
+    title_s=20, axlab_s=20, ticklab_s=20,
+    xlim=[-cwin*1./2, cwin*1./2], ylim=[ylim1, ylim2],
+    hide_axis=hide_axis, prettify=prettify, **mplp_kwargs)
 
     # Eventually save figure
     if saveFig:
@@ -1969,7 +2256,8 @@ def plt_ccg_subplots(units, CCGs, cbin=0.2, cwin=80, bChs=None, saveDir='~/Downl
                      saveFig=False, _format='pdf', figsize=None,
                      labels=True, show_ttl=True, title=None,
                      ylim_acg=None, ylim_ccg=None, share_y=0, normalize='zscore',
-                     acg_color=None, ccg_color='black', hide_axis=False, pad=0):
+                     acg_color=None, ccg_color='black', hide_axis=False, pad=0,
+                     prettify=True, **mplp_kwargs):
 
     ## format parameters
     if acg_color is not None and not isinstance(acg_color, str) and not isinstance(acg_color, int):
@@ -2098,7 +2386,8 @@ def plt_ccg_subplots(units, CCGs, cbin=0.2, cwin=80, bChs=None, saveDir='~/Downl
                 title_s=8, title_w='regular',
                 axlab_s=16, axlab_w='regular',
                 ticklab_s=18, ticklab_w='regular',
-            tight_layout=False, hide_axis=hide_axis)
+                tight_layout=False, hide_axis=hide_axis,
+                prettify=prettify, **mplp_kwargs)
             i+=1
             pbar.update(1)
 
@@ -2110,26 +2399,27 @@ def plt_ccg_subplots(units, CCGs, cbin=0.2, cwin=80, bChs=None, saveDir='~/Downl
 
     return fig
 
+
 def plot_acg(dp, unit, cbin=0.2, cwin=80, normalize='Hertz', periods='all',
              saveDir='~/Downloads', saveFig=False, _format='pdf', figsize=None, verbose=False,
-             color=0, labels=True, title=None, ref_per=True, ylim=[0,0],
+             color=0, labels=True, title=None, ref_per=True, ylim=[0, 0], ax=None,
              acg_mn=None, acg_std=None, again=False,
-              train=None, hide_axis=False):
+             train=None, hide_axis=False, prettify=True, **mplp_kwargs):
     """
-    Parameters:
+    Plots precomputed autocorrelogram.
+    Arguments:
         - dp: str, data path
         - unit: float, unit index
         - cbin: float, binsize (ms)
         - cwin: float, full window size (ms)
         - normalize: str in ['Counts', 'Hertz', 'Pearson', 'zscore', 'mixte'], unit of y axis
         - periods: 'all' or [(t1,t2), (t3,t4)...], periods to use to compute CCG (in SECONDS)
-        
 
         - saveDir: str, save directory for figure
         - saveFig: bool, whether to dave Figure at saveDir
         - _format: str, format to save fig (pdf, svg, eps, png, jpeg...)
-        - figsize: (float, float), figure size ((x, y) in inches)
-        - verbose: bool, if True prints more information
+        - figsize: (x,y) tuple, size of figure in inches
+        - verbose: bool, if True prints information
 
         - color: string, self explanatory (can also use 0-5 as keys of npyx.utils.phyColorsDic)
         - labels: bool, whether to plot axis labels/title
@@ -2137,68 +2427,92 @@ def plot_acg(dp, unit, cbin=0.2, cwin=80, normalize='Hertz', periods='all',
         - ref_per: bool, if True plot vertical lines highlighting 2ms refractory period
 
         - ylim: [float, float], ylim for autocorrelograms in case as_grid is True
+        - ax: matplotlib axis, axis to plot on (if None, creates new figure)
         - acg_mn: float, optionally feed externally calculated mean to zscore the CCG
         - acg_std: float, optionally feed externally calculated std to zscore the CCG
 
         - again: bool, whether to recompute the CCG
-        - train: [array1, array2...], optional externally fed train to compute ACG (in samples, not seconds).
+        - train: np array (n_spikes,), optional externally fed train to compute ACG (in samples, not seconds).
                  If used, use any integers as 'units'.
 
+        - hide_axis: bool, whether to hide axis
+        - prettify: bool, whether to apply mplp() prettification or not
+        - **mplp_kwargs: any additional formatting parameters, passed to mplp()
+
     Returns:
         - fig: matplotlib figure object
     """
-    saveDir=op.expanduser(saveDir)
+    saveDir = op.expanduser(saveDir)
     if train is not None:
-        bChs=None
+        bChs = None
     else:
-        bChs=get_depthSort_peakChans(dp, units=[unit])[:,1].flatten()
+        bChs = get_depthSort_peakChans(dp, units=[unit])[:, 1].flatten()
     ylim1, ylim2 = ylim[0], ylim[1]
-    ACG=acg(dp, unit, cbin, cwin, fs=30000, normalize=normalize, verbose=verbose, periods=periods, again=again, train=train)
-    if normalize=='zscore':
-        ACG_hertz=acg(dp, unit, cbin, cwin, fs=30000, normalize='Hertz', verbose=verbose, periods=periods)
-        acg25, acg35 = ACG_hertz[:int(len(ACG_hertz)*2./5)], ACG_hertz[int(len(ACG_hertz)*3./5):]
-        acg_std=np.std(np.append(acg25, acg35))
-        acg_mn=np.mean(np.append(acg25, acg35))
-    fig=plt_acg(unit, ACG, cbin, cwin, bChs, color, 30000, saveDir, saveFig, _format=_format,
-            labels=labels, title=title, ref_per=ref_per, ylim1=ylim1, ylim2=ylim2,
-            normalize=normalize, acg_mn=acg_mn, acg_std=acg_std, figsize=figsize, hide_axis=hide_axis)
+
+    ACG = acg(dp, unit, cbin, cwin, fs=30000, normalize=normalize,
+              verbose=verbose, periods=periods, again=again, train=train)
+    if normalize == 'zscore':
+        ACG_hertz = acg(dp, unit, cbin, cwin, fs=30000, normalize='Hertz', verbose=verbose, periods=periods)
+        acg25, acg35 = ACG_hertz[:int(len(ACG_hertz) * 2. / 5)], ACG_hertz[int(len(ACG_hertz) * 3. / 5):]
+        acg_std = np.std(np.append(acg25, acg35))
+        acg_mn = np.mean(np.append(acg25, acg35))
+    fig = plt_acg(unit, ACG, cbin, cwin, bChs, color, 30000, saveDir, saveFig, _format=_format,
+                  labels=labels, title=title, ref_per=ref_per, ylim1=ylim1, ylim2=ylim2, ax=ax,
+                  normalize=normalize, acg_mn=acg_mn, acg_std=acg_std, figsize=figsize, hide_axis=hide_axis,
+                  prettify=prettify, **mplp_kwargs)
 
     return fig
 
-def plot_ccg(dp, units, cbin=0.2, cwin=80, normalize='mixte', saveDir='~/Downloads', saveFig=False,
-             _format='pdf', figsize=None, periods='all', labels=True,
-             title=None, show_ttl=True, color=None, CCG=None,
+
+def plot_ccg(dp, units, cbin=0.2, cwin=80, normalize='mixte',
+             saveDir='~/Downloads', saveFig=False, _format='pdf', figsize=None, periods='all',
+             labels=True, title=None, show_ttl=True, color=None, CCG=None,
              ylim_acg=None, ylim_ccg=None, share_y=False,
              ccg_mn=None, ccg_std=None, again=False, trains=None, as_grid=False, show_hz=False,
-             use_template=True, enforced_rp=0, style='line', hide_axis=False, pad=0):
+             use_template=True, enforced_rp=0, style='line', hide_axis=False, pad=0,
+             prettify=True, **mplp_kwargs):
     """
-    Parameters:
+    Arguments:
         - dp: str, data path
         - units: [float, float], list of 2 unit indices
         - cbin: float, binsize (ms)
         - cwin: float, full window size (ms)
         - normalize: str in ['Counts', 'Hertz', 'Pearson', 'zscore', 'mixte'], unit of y axis
+
         - saveDir: str, save directory for figure
         - saveFig: bool, whether to dave Figure at saveDir
         - _format: str, format to save fig (pdf, svg, eps, png, jpeg...)
-        - figsize: (float, float), figure size ((x, y) in inches)
+        - figsize: (x,y) tuple, size of figure in inches
+
         - periods: 'all' or [(t1,t2), (t3,t4)...], periods to use to compute CCG (in SECONDS)
+
         - labels: bool, whether to plot axis labels/title
         - title: str, figure title
         - show_ttl: bool, whether to show title
-        - color: int, classical phy ACG colors (-1 black, 0 blue, 1 red...)
+        - color: int, classical phy ACG colors (-1 black, 0 blue, 1 red...
+
         - CCG: array, optionnaly feed externally computed CCG
+
         - ylim_acg: [float, float], ylim for autocorrelograms in case as_grid is True
         - ylim_ccg: [float, float], ylim for crosscorrelogram(s)
         - share_y: bool, whether to use the same y limit for all CCGs if as_grid is True
+
         - ccg_mn: float, optionally feed externally calculated mean to zscore the CCG
         - ccg_std: float, optionally feed externally calculated std to zscore the CCG
+
         - again: bool, whether to recompute the CCG
         - trains: [array1, array2...], optional externally fed list of trains to compute ACGs/CCGs (in samples). Then use any integers as 'units'.
         - as_grid: bool, also plot units autocorrelograms along the diagonal (only relevant when plotting 2 units)
+        - show_hz: bool, whether to add a second y axis and show the values corresponding to z-score units in Hz (only applies if normalize='zscore')
+
         - use_template: bool, whether to use the template files to find the peak channel
         - enforced_rp: float, enforced refractory period (will remove spikes happening within enforced_rp ms of another) | Default 0
         - style: str, 'line' or 'bar' (to plot ccg as a line or a histogram)
+        - hide_axis: bool, whether to hide axis
+        - pad: float, padding between subplots (in inches)
+
+        - prettify: bool, whether to apply mplp() prettification or not
+        - **mplp_kwargs: any additional formatting parameters, passed to mplp()
 
     Returns:
         - fig: matplotlib figure object
@@ -2231,12 +2545,13 @@ def plot_ccg(dp, units, cbin=0.2, cwin=80, normalize='mixte', saveDir='~/Downloa
         fig = plt_ccg(units, CCG[0,1,:], cbin, cwin, bChs, 30000, saveDir, saveFig, _format,
                       labels=labels, title=title, color=color, ylim=ylim_ccg,
                       normalize=normalize, ccg_mn=ccg_mn, ccg_std=ccg_std,
-                      figsize=figsize, style=style, hide_axis=hide_axis, show_hz=show_hz, show_ttl=show_ttl)
+                      figsize=figsize, style=style, hide_axis=hide_axis, show_hz=show_hz, show_ttl=show_ttl,
+                      prettify=prettify, **mplp_kwargs)
     else:
         fig = plt_ccg_subplots(units, CCG, cbin, cwin, bChs, saveDir, saveFig, _format, figsize,
                                labels=labels, show_ttl=show_ttl,title=title,
                                ylim_acg=ylim_acg, ylim_ccg=ylim_ccg, share_y=share_y, normalize=normalize,
-                               acg_color=color, hide_axis=hide_axis, pad=pad)
+                               acg_color=color, hide_axis=hide_axis, pad=pad, prettify=prettify, **mplp_kwargs)
 
     return fig
 
@@ -2295,17 +2610,29 @@ def imshow_cbar(im, origin='top', xevents_toplot=[], yevents_toplot=[], events_c
                 cmapstr="RdBu_r", vmin=None, vmax=None, center=None, colorseq='nonlinear',
                 clabel='', cticks=None,
                 figsize=(6,4), aspect='auto', function='imshow',
-                ax=None, tight_layout=True, cmap_h=0.3, **kwargs):
+                ax=None, tight_layout=True, cmap_h=0.3, prettify=True,
+                **kwargs):
     '''
     Essentially plt.imshow(im, cmap=cmapstr), but with a nicer and actually customizable colorbar.
 
-    Parameters:
+    Arguments:
         - im: 2D array def to matplotlib.pyplot.imshow
         - origin: y axis origin, either top or bottom | Default: top
+
+        - xevents_toplot: list of events to plot as vertical dashed lines
+        - yevents_toplot: list of events to plot as horizontal dashed lines
+        - events_color: color of the dashed lines
+        - events_lw: linewidth of the dashed lines
+
         - xvalues, yvalues: lists/arrays of lengths im.shape[1] and im.shape[0], respectively.
                             Allows to alter the value to which pixel positions are mapped (which are dumb pixel ranks by default).
-        - xticks, yticks: allows to alter the position of the ticks (in [0,npixels] space by default, in xvalues/yvalues space if they are provided)
-        - xticklabels, yticklabels: allows to alter the label of the ticks - should have the same size as xticks/yticks
+        - xticks, yticks: lists/array, allows to alter the position of the ticks (in [0,npixels] space by default, in xvalues/yvalues space if they are provided)
+        - xticklabels, yticklabels: list of str, allows to alter the label of the ticks - should have the same size as xticks/yticks
+        - xlabel, ylabel: str, labels for the x and y axes
+
+        - xtickrot: int, rotation of the xticklabels (degrees)
+        - title: str, figure title
+
         - cmapstr: string, colormap name from matplotlib ('RdBu_r') or Fabio Crameri package ('batlow')
         - vmin: value to which the lower boundary of the colormap corresponds
         - vmax: value to which the upper boundary of the colormap corresponds
@@ -2313,9 +2640,17 @@ def imshow_cbar(im, origin='top', xevents_toplot=[], yevents_toplot=[], events_c
         - colorseq: string, {'linear', 'nonlinear'}, whether to shrink or not the colormap between the center and the closest boundary
                     when 'center' is not None and isn't equidistant between vmax and vmin
         - clabel: string, colormap label
-        - extend_cmap: tring, {'neither', 'both', 'min', 'max'}. If not 'neither', make pointed end(s) for out-of- range values.
         - cticks: list of ticks to show
-        - aspect: {'equal', 'auto'}, see imshow documentation
+
+        - figsize: (x,y) tuple, size of figure in inches
+        - aspect: {'equal', 'auto'}, see matplotlib.pyplot.imshow documentation
+        - function: {'imshow', 'pcolormesh'}, whether to use imshow or pcolormesh to plot the image
+        - ax: matplotlib axis, if None, a new figure is created
+        - tight_layout: bool, whether to use plt.tight_layout() or not
+        - cmap_h: float, height of the colorbar in inches
+
+        - prettify: bool, whether to apply mplp() prettification or not
+        - **kwargs: additional arguments to be passed to the plotting function imshow or pcolormesh (e.g. interpolation='nearest')
     '''
     assert colorseq in ['linear', 'nonlinear']
     if im.ndim==1:
@@ -2380,7 +2715,8 @@ def imshow_cbar(im, origin='top', xevents_toplot=[], yevents_toplot=[], events_c
           axlab_w='bold', axlab_s=16,
           ticklab_w='regular', ticklab_s=12, ticks_direction='out', lw=1,
           title=title, title_w='regular', title_s=12,
-          hide_top_right=False, hide_axis=False, tight_layout=False)
+          hide_top_right=False, hide_axis=False, tight_layout=False,
+          prettify=prettify)
 
     if tight_layout: fig.tight_layout(rect=[0,0,0.8,1])
 
@@ -2501,7 +2837,7 @@ def plot_sfcm(dp, corr_type='connections', metric='amp_z', cbin=0.5, cwin=100,
     Visually represents the connectivity matrix sfcm computed with npyx.corr.gen_sfc().
     Each line/row is a unit, sorted by depth, and the colormap corresponds to the 'metric' parameter.
 
-    Parameters:
+    Arguments:
         - all parameters of npyx.corr.gen_sfc():
             {gen_sfc.__doc__}
     Returns:
@@ -2543,7 +2879,7 @@ def plot_sfcm(dp, corr_type='connections', metric='amp_z', cbin=0.5, cwin=100,
                 xvalues=None, yvalues=None, xticks=tks, yticks=tks, title=ttl,
                 xticklabels=labs, yticklabels=labs, xlabel=lab, ylabel=lab,
                 cmapstr="RdBu_r", vmin=vminmax[0], vmax=vminmax[1], center=0, colorseq='nonlinear',
-                clabel='Crosscorr. (zscore)', extend_cmap='neither', cticks=None,
+                clabel='Crosscorr. (zscore)', cticks=None,
                 figsize=figsize, aspect='auto', function='imshow',
                 ax=None)
 
@@ -2794,4 +3130,4 @@ def plot_filtered_times(dp, unit, first_n_minutes=20, consecutive_n_seconds = 18
     for i in gauss_sec:
         s_time, e_time = i ,(i+1)
         plt.hlines(2, s_time, e_time, color = 'blue')
-    plt.show()
\ No newline at end of file
+    plt.show()
diff --git a/npyx/spk_t.py b/npyx/spk_t.py
index 2e213fd7..f8408efe 100644
--- a/npyx/spk_t.py
+++ b/npyx/spk_t.py
@@ -8,8 +8,9 @@
 opj=op.join
 from pathlib import Path, PosixPath, WindowsPath
 
-# from itertools import groupby
-# from operator import itemgetter
+from joblib import Memory
+cachedir = Path(op.expanduser("~")) / ".NeuroPyxels"
+cache_memory = Memory(cachedir, verbose=0)
 
 import matplotlib.pyplot as plt
 
@@ -18,7 +19,7 @@
 from scipy.stats import iqr
 from scipy.optimize import curve_fit
 from scipy.stats import norm
-from npyx.utils import smooth, thresh_consec, npa, assert_int, assert_float
+from npyx.utils import smooth, thresh_consec, npa, assert_int, assert_float, docstring_decorator
 from npyx.gl import get_units, get_npyx_memory, check_periods
 from npyx.inout import read_metadata
 
@@ -212,12 +213,16 @@ def inst_cv2(t):
 
     return cv2
 
+@cache_memory.cache
 def mean_firing_rate(t, exclusion_quantile=0.005, fs=30000):
-    i = np.diff(t) if len(t)>1 else None
-    if i is None: return 0
-    # Remove outlyers
-    i=i[(i>=np.quantile(i, exclusion_quantile))&(i<=np.quantile(i, 1-exclusion_quantile))]/fs
-    return np.round(1./np.mean(i),2)
+    isint = np.diff(t) if len(t)>1 else None
+    if isint is None: return 0
+    # Remove outliers
+    quantile_high = np.quantile(isint, 1-exclusion_quantile)
+    quantil_low = np.quantile(isint, exclusion_quantile)
+    isint = isint[(isint >= quantil_low) & (isint <= quantile_high)]
+    isint = isint/fs
+    return np.round(1./np.mean(isint),2)
 
 def mfr(dp=None, U=None, exclusion_quantile=0.005, enforced_rp=0,
         periods='all', again=False, train=None, fs=None):
@@ -387,7 +392,6 @@ def train_quality(dp, unit, period_m=[0,20],
                   violations_ms = 0.8, fp_threshold = 0.05, fn_threshold = 0.05,
                   again = False, save = True, verbose = False, plot_debug = False,
                   enforced_rp = 0):
-
     """
     Subselect spike times which meet two criteria:
         low number of 'missed spikes' (false negatives)
@@ -442,8 +446,9 @@ def train_quality(dp, unit, period_m=[0,20],
         - enforced_rp: float, enforced refractory period in ms (if 2 spikes are closer than enforced_rp ms, only the first one  is kept.)
     
     Returns:
-    - goodsec, acgsec, gausssec
-
+        - good_spikes_m: mask of spikes belonging to the intersection of the chunks with low enough fp/fn rates.
+        - good_fp_start_end: list of [start, end] of chunks with low enough fp rate (in seconds).
+        - good_fn_start_end: list of [start, end] of chunks with low enough fn rate (in seconds).
     """
     
     dp = Path(dp)
@@ -477,6 +482,9 @@ def train_quality(dp, unit, period_m=[0,20],
     unit_amp = load_amplitudes(dp, unit, verbose, 'all', again, enforced_rp)
     
     unit_train = trn(dp, unit, enforced_rp=enforced_rp, again=again, verbose=verbose)/fs
+
+    if period_m is None:
+        period_m = [0, unit_train[-1]/60]
     period_s=[period_m[0]*60, period_m[1]*60]
 
     # Attempt to reload if precomputed
@@ -543,11 +551,11 @@ def train_quality(dp, unit, period_m=[0,20],
         # Compute denominator of FP on total period
         acg_tot = acg(dp, unit, c_bin, c_win, verbose = False,  periods=[period_s])
         x_block = np.round(np.arange(-c_win/2, c_win/2 + c_bin, c_bin), 8) # round to fix binary imprecisions
-        rp_mask = (x_block >= -violations_ms) & (x_block <= violations_ms)
+        # rp_mask = (x_block >= -violations_ms) & (x_block <= violations_ms)
         baseline_mask = (acg_tot*0).astype(bool)
         baseline_mask[:n_bins_acg_baseline] = True
         baseline_mask[-n_bins_acg_baseline:] = True
-        baseline_mean = np.mean(acg_tot[baseline_mask])
+        # baseline_mean = np.mean(acg_tot[baseline_mask])
         
         # False positive estimation
         for i, (t1,t2) in enumerate(fp_chunks):
@@ -556,12 +564,16 @@ def train_quality(dp, unit, period_m=[0,20],
             n_spikes_chunk=np.sum(chunk_mask)
 
             if n_spikes_chunk > 15:
-                ACG = acg(dp, unit, c_bin, c_win, verbose = False,  periods=[(t1, t2)])
-                violations_mean = np.mean(ACG[rp_mask])
-                rpv_ratio_acg = round(violations_mean / baseline_mean, 4)
-                fp_toplot.append(rpv_ratio_acg)
+
+                # ACG = acg(dp, unit, c_bin, c_win, verbose = False,  periods=[(t1, t2)])
+                # violations_mean = np.mean(ACG[rp_mask])
+                #rpv_ratio_acg = round(violations_mean / baseline_mean, 4)
+                # fp_toplot.append(rpv_ratio_acg)
+                fp_rate = isi_violations(unit_train, min_time=t1, max_time=t2,
+                                         isi_threshold=violations_ms/1000, min_isi=0.0005)[0]
+                fp_toplot.append(fp_rate)
                 chunk_fp_t.append(t1+(t2-t1)/2)
-                if (rpv_ratio_acg <= fp_threshold):
+                if (fp_rate <= fp_threshold):
                     passed_fp[i] = [t1, t2, 1]
 
     # Across all chunks, if at least 1 good chunk for both (else no spike can be called good)
@@ -645,23 +657,7 @@ def train_quality(dp, unit, period_m=[0,20],
             
         return good_spikes_m, [0], [0]
 
-# def get_consec_sections(seconds):
-#         """
-#         Given an array with seconds as entries (with 1 sec increments)
-
-#         Return: list of consecutive sections start and end times
-#         """
-#         sec_all = []
-#         for k, g in groupby(enumerate(seconds), lambda ix: ix[0]-ix[1]):
-#             sec_all.append(list( map(itemgetter(1), g)))
-#         start_end = []
-
-#         # get the start and end times of these section
-#         for good_section in sec_all:
-#             start_end.append([good_section[0], good_sectienforced_rp
-#         return start_end
-
-
+@docstring_decorator(train_quality.__doc__)
 def trn_filtered(dp, unit, period_m=[0,20],
                   fp_chunk_span=3, fp_chunk_size = 10,
                   fn_chunk_span = 3, fn_chunk_size = 10,
@@ -670,7 +666,7 @@ def trn_filtered(dp, unit, period_m=[0,20],
                   use_consecutive = False, consecutive_n_seconds = 180,
                   again = False, save = True, verbose = False, plot_debug = False,
                   enforced_rp=0):
-    f"""
+    """
     Returns spike times (in sample) meeting the false positive and false negative criteria.
     Mainly wrapper of train_quality().
     
@@ -688,7 +684,7 @@ def trn_filtered(dp, unit, period_m=[0,20],
                           
     train_quality docstring:
     
-    {train_quality.__doc__}
+    {0}
     """
     dp = Path(dp)
     t = trn(dp,unit, enforced_rp=enforced_rp)
@@ -700,10 +696,7 @@ def trn_filtered(dp, unit, period_m=[0,20],
 
     # use spike times themselves to define beginning and end of good Sections
     # as the FP and FN sections do not necessarily overlap
-    edges = np.diff([0]+list(good_spikes_m.astype(int))+[0])
-    good_left = np.nonzero(edges==1)[0]
-    good_right = np.nonzero(edges==-1)[0]-1
-    good_sections = [[t_s[l], t_s[r]] for l,r in zip(good_left, good_right)]
+    good_sections = good_sections_from_mask(good_spikes_m, t_s)
     
     if len(good_sections)>0:#
         total_good_sections = sum([s[1]-s[0] for s in good_sections])
@@ -727,6 +720,48 @@ def trn_filtered(dp, unit, period_m=[0,20],
     if verbose: print('No consecutive section passed the filters')
     return np.array([0]), (t*0).astype(bool)
 
+def good_sections_from_mask(good_times_m, time_series=None):
+    """
+    Returns a list of good sections [[t1,t2], ...] in units of 'time_series'
+    from a boolean mask of time stamps good_times_m.
+    """
+    edges = np.diff([0]+list(good_times_m.astype(int))+[0])
+    good_left = np.nonzero(edges==1)[0]
+    good_right = np.nonzero(edges==-1)[0]-1
+    if time_series is not None:
+        return [[time_series[l], time_series[r]] for l,r in zip(good_left, good_right) if r>l]
+    else:
+        return [[l,r] for l,r in zip(good_left, good_right) if r>l]
+
+@cache_memory.cache
+def get_common_good_sections(good_sections_list):
+    """
+    good_sections: list of lists of periods
+    [
+    [[t11, t12], [t13,t14], ...],
+    [[t21, t22], [t23,t24], ...],
+    ]
+    Must be in SAMPLES, INTEGERS.
+
+    Returns:
+        - common_good_sections: list of sections common to all the lists of periods.
+    """
+
+    periods = [sec[i] for sec in good_sections_list for i in range(len(sec))]
+    end = int(np.max(periods))
+    bad_sections = []
+    for sec in good_sections_list:
+        sec = np.array(sec)
+        assert sec.shape[1]==2, "All periods must be of size 2!"
+        bad_sec = np.concatenate([[0], np.array(sec).ravel(), [end]]).astype(int)
+        bad_sec = bad_sec.reshape(-1,2)
+        bad_sections.append(bad_sec)
+    m = np.full((end,), True)
+    for sec in bad_sections:
+        for sec_per in sec:
+            m[slice(sec_per[0]+1, sec_per[1])] = False
+
+    return good_sections_from_mask(m)
 
 def gaussian_cut(x, a, mu, sigma, x_cut):
     g = a * np.exp(-(x - mu) ** 2 / (2 * sigma ** 2))
@@ -839,3 +874,4 @@ def estimate_bins(x, rule):
 from  npyx.corr import acg
 from npyx.merger import assert_multi, get_dataset_id, get_ds_table, get_source_dp_u
 from npyx.plot import plot_fp_fn_rates
+from npyx.metrics import isi_violations
\ No newline at end of file
diff --git a/npyx/utils.py b/npyx/utils.py
index de850be0..801b92b8 100644
--- a/npyx/utils.py
+++ b/npyx/utils.py
@@ -196,6 +196,26 @@ def list_files(directory, extension, full_path=False):
         return [Path('/'.join([directory,f])) for f in files]
     return files
 
+
+## Text formatting utilities
+
+def pprint_dic(dic):
+    kv = "".join([f"{k}: {v},\n" for k,v in dic.items()])
+    return "{\n"+kv+"}"
+
+def docstring_decorator(*args):
+    """
+    Feed as many arguments as wished to incorporate into the function f's docstring.
+    Edit f' docstring with {0}, {1}... acordingly.
+    """
+    def decorate(f):
+        f.__doc__ = f.__doc__.format(*args)
+        return f
+    return decorate
+
+
+## Threholding utilities
+
 def any_n_consec(X, n_consec, where=False):
     '''
     The trick below finds whether there are n_consec consecutive ones in the array comp
@@ -377,6 +397,9 @@ def thresh_cons(arr, th, sgn=1, n_consec=0, exclude_edges=True, ret_values=True)
 
     return crosses
 
+
+## Zscoring, smoothing
+
 def zscore(arr, frac=4./5, mn_ext=None, sd_ext=None):
     '''
     Returns z-scored (centered, reduced) array using outer edges of array to compute mean and std.
@@ -421,7 +444,6 @@ def smooth(arr, method='gaussian_causal', sd=5, axis=1, gamma_a=5):
     pad_width = [[C,C] if i==axis else [0,0] for i in range(arr.ndim)]
     padarr=np.pad(arr, pad_width, 'symmetric')
 
-
     ## Compute the kernel
     if method in ['gaussian', 'gaussian_causal']:
         X=np.arange(-4*sd, 4*sd+1)
@@ -452,7 +474,7 @@ def smooth(arr, method='gaussian_causal', sd=5, axis=1, gamma_a=5):
 
 
     ## Remove padding
-    sarr = sarr[slice_along_axis(C,-C,axis=axis)]
+    sarr = sarr[slice_along_axis(C+1,-C+1, axis=axis)]
     assert np.all(sarr.shape==arr.shape)
 
 
diff --git a/requirements_dev.txt b/requirements_dev.txt
index 55df9879..128fe084 100644
--- a/requirements_dev.txt
+++ b/requirements_dev.txt
@@ -1,4 +1,4 @@
 pip>=19.2
-wheel==0.33.1
+wheel==0.38.1
 twine==1.13.0
 python<=3.10


Test patch:
diff --git a/npyx/testing.py b/npyx/testing.py
index 336ef669..92141ddd 100644
--- a/npyx/testing.py
+++ b/npyx/testing.py
@@ -1,12 +1,14 @@
 import traceback
 
+import numpy as np
+
 import npyx
 from npyx.inout import get_npix_sync
 from npyx.gl import get_units, read_metadata
 from npyx.spk_t import ids, trn, trn_filtered
 from npyx.spk_wvf import wvf, wvf_dsmatch, get_peak_chan, templates
 from npyx.corr import ccg
-from npyx.plot import plot_ccg, plot_wvf, plot_raw
+from npyx.plot import plot_acg, plot_ccg, plot_wvf, plot_raw
 
 prefix = "\033[34;1m--- "
 red_prefix = "\033[91;1m--- "
@@ -32,8 +34,7 @@ def test_npyx(dp, raise_error=False):
     units = test_function(get_units, raise_error, True, dp=dp)
     if units is None:
         raise ValueError("Something went really wrong, get_units did not run. Fix code and try again before testing next function.")
-    u = units[-1]
-    u1 = units[-3]
+    u, u1 = units[np.random.randint(0, len(units)-1, 2)]
 
     test_function(ids, raise_error, dp=dp, unit=u, again=1)
 
@@ -48,11 +49,13 @@ def test_npyx(dp, raise_error=False):
     test_function(get_peak_chan, raise_error, dp=dp, unit=u, again=1)
 
     test_function(templates, raise_error, dp=dp, u=u)
-
+    
     test_function(ccg, raise_error, dp=dp, U=[u,u1], bin_size=0.2, win_size=40, again=1)
 
     test_function(plot_wvf, raise_error, dp=dp, u=u, again=1, use_dsmatch=True)
 
+    test_function(plot_acg, raise_error, dp=dp, unit=u, again=1)
+    
     test_function(plot_ccg, raise_error, dp=dp, units=[u,u1], as_grid=True, again=1)
 
     test_function(plot_raw, raise_error, dp=dp, times=[0.1,0.15], channels=list(range(50)), again=1)