def nearest_neighbors_noise_overlap(waveform_extractor: si.WaveformExtractor, this_unit_id: int, max_spikes_for_nn: int = 1000, n_neighbors: int = 5, n_components: int = 10, radius_um: float = 100, seed: int = 0): """Calculates unit noise overlap based on NearestNeighbors search in PCA space. Based on noise overlap metric described in Chung et al. (2017) Neuron 95: 1381-1394. Rough logic: ------------ 1) Generate a noise cluster by randomly sampling voltage snippets from recording. 2) Subtract projection onto the weighted average of noise snippets of both the target and noise clusters to correct for bias in sampling. 3) Compute the isolation score between the noise cluster and the target cluster. Implementation details: ----------------------- As with nn_isolation, the clusters that are compared (target and noise clusters) have the same number of spikes. See docstring for `_compute_isolation` for the definition of isolation score. Parameters: ----------- we: si.WaveformExtractor this_unit_id: int ID of unit for which this metric will be calculated max_spikes_for_nn: int max number of spikes to use per cluster n_neighbors: int number of neighbors to check membership of n_components: int number of PC components to project the snippets radius_um: float only the channels within this radius of the peak channel are used to compute the metric seed: int seed for random subsampling of spikes Outputs: -------- nn_noise_overlap : float """ # set random seed rng = np.random.default_rng(seed=seed) # get random snippets from the recording to create a noise cluster recording = waveform_extractor.recording noise_cluster = get_random_data_chunks( recording, return_scaled=waveform_extractor.return_scaled, num_chunks_per_segment=max_spikes_for_nn, chunk_size=waveform_extractor.nsamples, seed=seed) noise_cluster = np.reshape( noise_cluster, (max_spikes_for_nn, waveform_extractor.nsamples, -1)) # get waveforms for target cluster waveforms = waveform_extractor.get_waveforms(unit_id=this_unit_id) # adjust the size of the target and noise clusters to be equal if waveforms.shape[0] > max_spikes_for_nn: wf_ind = rng.choice(waveforms.shape[0], max_spikes_for_nn, replace=False) waveforms = waveforms[wf_ind] n_snippets = max_spikes_for_nn elif waveforms.shape[0] < max_spikes_for_nn: noise_ind = rng.choice(noise_cluster.shape[0], waveforms.shape[0], replace=False) noise_cluster = noise_cluster[noise_ind] n_snippets = waveforms.shape[0] else: n_snippets = max_spikes_for_nn # restrict to channels with significant signal closest_chans_idx = get_template_channel_sparsity(waveform_extractor, method='radius', outputs='index', peak_sign='both', radius_um=radius_um) waveforms = waveforms[:, :, closest_chans_idx[this_unit_id]] noise_cluster = noise_cluster[:, :, closest_chans_idx[this_unit_id]] # compute weighted noise snippet (Z) median_waveform = waveform_extractor.get_template(unit_id=this_unit_id, mode='median') median_waveform = median_waveform[:, closest_chans_idx[this_unit_id]] tmax, chmax = np.unravel_index(np.argmax(np.abs(median_waveform)), median_waveform.shape) weights = [noise_clip[tmax, chmax] for noise_clip in noise_cluster] weights = np.asarray(weights) weights = weights / np.sum(weights) weighted_noise_snippet = np.sum(weights * noise_cluster.swapaxes(0, 2), axis=2).swapaxes(0, 1) # subtract projection onto weighted noise snippet for snippet in range(n_snippets): waveforms[snippet, :, :] = _subtract_clip_component( waveforms[snippet, :, :], weighted_noise_snippet) noise_cluster[snippet, :, :] = _subtract_clip_component( noise_cluster[snippet, :, :], weighted_noise_snippet) # compute principal components after concatenation all_snippets = np.concatenate([ waveforms.reshape((n_snippets, -1)), noise_cluster.reshape((n_snippets, -1)) ], axis=0) pca = IncrementalPCA(n_components=n_components) pca.partial_fit(all_snippets) projected_snippets = pca.transform(all_snippets) # compute overlap nn_noise_overlap = 1 - _compute_isolation( projected_snippets[:n_snippets, :], projected_snippets[n_snippets:, :], n_neighbors) return nn_noise_overlap
def nearest_neighbors_isolation(waveform_extractor: si.WaveformExtractor, this_unit_id: int, max_spikes_for_nn: int = 1000, n_neighbors: int = 5, n_components: int = 10, radius_um: float = 100, seed: int = 0): """Calculates unit isolation based on NearestNeighbors search in PCA space Based on isolation metric described in Chung et al. (2017) Neuron 95: 1381-1394. Rough logic: ------------ 1) Choose a cluster 2) Compute the isolation score with every other cluster 3) Isolation score is defined as the min of (2) (i.e. 'worst-case measure') Implementation details: ----------------------- Let A and B be two clusters from sorting. We set |A| = |B|: If max_spikes_for_nn < |A| and max_spikes_for_nn < |B|, then randomly subsample max_spikes_for_nn samples from A and B. If max_spikes_for_nn > min(|A|, |B|) (e.g. |A| > max_spikes_for_nn > |B|), then randomly subsample min(|A|, |B|) samples from A and B. This is because the metric is affected by the size of the clusters being compared independently of how well-isolated they are. We also restrict the waveforms to channels with significant signal See docstring for `_compute_isolation` for the definition of isolation score. Parameters: ----------- all_pcs: array_like, (num_spikes, PCs) 2D array of PCs for all spikes all_labels: array_like, (num_spikes, ) 1D array of cluster labels for all spikes this_unit_id: int ID of unit for which this metric will be calculated max_spikes_for_nn: int max number of spikes to use per cluster n_neighbors: int number of neighbors to check membership of seed: int seed for random subsampling of spikes Outputs: -------- nn_isolation : float """ rng = np.random.default_rng(seed=seed) all_units_ids = waveform_extractor.sorting.get_unit_ids() other_units_ids = np.setdiff1d(all_units_ids, this_unit_id) # get waveforms for target cluster waveforms_target_unit = waveform_extractor.get_waveforms( unit_id=this_unit_id) n_spikes_target_unit = waveforms_target_unit.shape[0] # find units whose closest channels overlap with closest channels of target cluster closest_chans_all = get_template_channel_sparsity(waveform_extractor, method='radius', outputs='index', peak_sign='both', radius_um=radius_um) closest_chans_target_unit = closest_chans_all[this_unit_id] other_units_ids = [ unit_id for unit_id in other_units_ids if np.any( np.in1d(closest_chans_all[unit_id], closest_chans_target_unit)) ] # if no unit is within neighborhood of target unit, then just say isolation is 1 (best possible) if not other_units_ids: nn_isolation = 1 # if there are units to compare, then compute isolation with each else: isolation = np.zeros(len(other_units_ids), ) for other_unit_id in other_units_ids: waveforms_other_unit = waveform_extractor.get_waveforms( unit_id=other_unit_id) n_spikes_other_unit = waveforms_other_unit.shape[0] n_snippets = np.min( [n_spikes_target_unit, n_spikes_other_unit, max_spikes_for_nn]) # make the two clusters equal in terms of: # - number of spikes # - channels with signal waveforms_target_unit_idx = rng.choice(n_spikes_target_unit, size=n_snippets, replace=False) waveforms_target_unit_sampled = waveforms_target_unit[ waveforms_target_unit_idx] waveforms_target_unit_sampled = waveforms_target_unit_sampled[:, :, closest_chans_target_unit] waveforms_other_unit_idx = rng.choice(n_spikes_other_unit, size=n_snippets, replace=False) waveforms_other_unit_sampled = waveforms_other_unit[ waveforms_other_unit_idx] waveforms_other_unit_sampled = waveforms_other_unit_sampled[:, :, closest_chans_target_unit] # compute principal components after concatenation all_snippets = np.concatenate([ waveforms_target_unit_sampled.reshape((n_snippets, -1)), waveforms_other_unit_sampled.reshape((n_snippets, -1)) ], axis=0) pca = IncrementalPCA(n_components=n_components) pca.partial_fit(all_snippets) projected_snippets = pca.transform(all_snippets) # compute isolation isolation[other_unit_id == other_units_ids] = _compute_isolation( projected_snippets[:n_snippets, :], projected_snippets[n_snippets:, :], n_neighbors) nn_isolation = np.min(isolation) return nn_isolation