From 10cb5bf96445ed0a28468c466abadba9c18eba35 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jes=C3=BAs=20Royeth?= Date: Tue, 1 Sep 2026 15:22:09 -0400 Subject: [PATCH 1/5] Parallelize coherence-based bad channel detection --- .../preprocessing/detect_bad_channels.py | 98 ++++++++++++++----- .../tests/test_detect_bad_channels.py | 65 ++++++++++++ 2 files changed, 141 insertions(+), 22 deletions(-) diff --git a/src/spikeinterface/preprocessing/detect_bad_channels.py b/src/spikeinterface/preprocessing/detect_bad_channels.py index b255c04a9e..9bf4b17728 100644 --- a/src/spikeinterface/preprocessing/detect_bad_channels.py +++ b/src/spikeinterface/preprocessing/detect_bad_channels.py @@ -4,6 +4,8 @@ from typing import Literal from spikeinterface.core.core_tools import define_function_handling_dict_from_class +from spikeinterface.core.job_tools import TimeSeriesChunkExecutor, fix_job_kwargs +from spikeinterface.core.time_series_tools import get_random_sample_slices from .filter import highpass_filter from spikeinterface.core import get_random_data_chunks, order_channels_by_depth, BaseRecording from spikeinterface.core.channelslice import ChannelSliceRecording @@ -75,6 +77,12 @@ The random seed to extract chunks channel_filters : set | None, default: None For coherence+psd - only return `bad_channel_ids` whose labels are in the set `channel_filter`. +job_kwargs : dict | None, default: None + Keyword arguments for parallel processing. Only used for the "coherence+psd" method. Only the + execution-related keys (`pool_engine`, `n_jobs`, `progress_bar`, `mp_context`, + `max_threads_per_worker`) apply; the chunking size is fixed by `chunk_duration_s` and + `num_random_chunks` above, so `chunk_size`, `chunk_memory`, `total_memory` and `chunk_duration` + are not used here. """ @@ -153,6 +161,39 @@ def _get_all_detect_bad_channel_kwargs(detect_bad_channels_kwargs): return all_detect_bad_channels_kwargs +def _detect_bad_channels_chunk_init(recording, method_kwargs): + return {"recording": recording, "method_kwargs": method_kwargs} + + +def _detect_bad_channels_chunk(segment_index, start_frame, end_frame, worker_context): + recording = worker_context["recording"] + method_kwargs = worker_context["method_kwargs"] + + random_chunk = recording.get_traces( + start_frame=start_frame, + end_frame=end_frame, + segment_index=segment_index, + return_in_uV=True, + ) + + order_f = method_kwargs["order_f"] + order_r = method_kwargs["order_r"] + random_chunk_sorted = random_chunk[:, order_f] if order_f is not None else random_chunk + chunk_labels = detect_bad_channels_ibl( + raw=random_chunk_sorted, + fs=recording.sampling_frequency, + psd_hf_threshold=method_kwargs["psd_hf_threshold"], + dead_channel_thr=method_kwargs["dead_channel_threshold"], + noisy_channel_thr=method_kwargs["noisy_channel_threshold"], + outside_channel_thr=method_kwargs["outside_channel_threshold"], + n_neighbors=method_kwargs["n_neighbors"], + nyquist_threshold=method_kwargs["nyquist_threshold"], + welch_window_ms=method_kwargs["welch_window_ms"], + outside_channels_location=method_kwargs["outside_channels_location"], + ) + return chunk_labels[order_r] if order_r is not None else chunk_labels + + def detect_bad_channels( recording: BaseRecording, method: str = "coherence+psd", @@ -173,6 +214,7 @@ def detect_bad_channels( neighborhood_r2_radius_um: float = 30.0, seed: int | None = None, channel_filters: set | None = None, + job_kwargs: dict | None = None, ): """ Perform bad channel detection. @@ -225,14 +267,12 @@ def detect_bad_channels( if method in ("std", "mad"): random_chunk_kwargs["return_in_uV"] = False random_chunk_kwargs["concatenated"] = True - elif method == "coherence+psd": - random_chunk_kwargs["return_in_uV"] = True - random_chunk_kwargs["concatenated"] = False elif method == "neighborhood_r2": random_chunk_kwargs["return_in_uV"] = False random_chunk_kwargs["concatenated"] = False - random_data = get_random_data_chunks(recording_hp, **random_chunk_kwargs) + if method != "coherence+psd": + random_data = get_random_data_chunks(recording_hp, **random_chunk_kwargs) channel_labels = np.zeros(recording.get_num_channels(), dtype="U5") channel_labels[:] = "good" @@ -248,6 +288,14 @@ def detect_bad_channels( channel_labels[mask] = "noise" elif method == "coherence+psd": + if job_kwargs is None: + job_kwargs = {"progress_bar": False} + job_kwargs = fix_job_kwargs(job_kwargs) + executor_job_kwargs = { + key: job_kwargs[key] + for key in ("pool_engine", "n_jobs", "progress_bar", "mp_context", "max_threads_per_worker") + } + # some checks assert recording.has_scaleable_traces(), ( "The 'coherence+psd' method uses thresholds assuming the traces are in uV, " @@ -267,24 +315,30 @@ def detect_bad_channels( order_f = None order_r = None - # Create empty channel labels and fill with bad-channel detection estimate for each chunk - chunk_channel_labels = np.zeros((recording.get_num_channels(), len(random_data)), dtype=np.int8) - - for i, random_chunk in enumerate(random_data): - random_chunk_sorted = random_chunk[:, order_f] if order_f is not None else random_chunk - chunk_labels = detect_bad_channels_ibl( - raw=random_chunk_sorted, - fs=recording.sampling_frequency, - psd_hf_threshold=psd_hf_threshold, - dead_channel_thr=dead_channel_threshold, - noisy_channel_thr=noisy_channel_threshold, - outside_channel_thr=outside_channel_threshold, - n_neighbors=n_neighbors, - nyquist_threshold=nyquist_threshold, - welch_window_ms=welch_window_ms, - outside_channels_location=outside_channels_location, - ) - chunk_channel_labels[:, i] = chunk_labels[order_r] if order_r is not None else chunk_labels + method_kwargs = dict( + order_f=order_f, + order_r=order_r, + psd_hf_threshold=psd_hf_threshold, + dead_channel_threshold=dead_channel_threshold, + noisy_channel_threshold=noisy_channel_threshold, + outside_channel_threshold=outside_channel_threshold, + n_neighbors=n_neighbors, + nyquist_threshold=nyquist_threshold, + welch_window_ms=welch_window_ms, + outside_channels_location=outside_channels_location, + ) + random_slices = get_random_sample_slices(recording_hp, **random_chunk_kwargs) + executor = TimeSeriesChunkExecutor( + recording_hp, + _detect_bad_channels_chunk, + _detect_bad_channels_chunk_init, + (recording_hp, method_kwargs), + handle_returns=True, + chunk_size=random_chunk_kwargs["chunk_size"], + job_name="detect_bad_channels", + **executor_job_kwargs, + ) + chunk_channel_labels = np.stack(executor.run(slices=random_slices), axis=1) # Take the mode of the chunk estimates as final result. Convert to binary good / bad channel output. mode_channel_labels, _ = mode(chunk_channel_labels, axis=1, keepdims=False) diff --git a/src/spikeinterface/preprocessing/tests/test_detect_bad_channels.py b/src/spikeinterface/preprocessing/tests/test_detect_bad_channels.py index 05fbabdf54..e73493b326 100644 --- a/src/spikeinterface/preprocessing/tests/test_detect_bad_channels.py +++ b/src/spikeinterface/preprocessing/tests/test_detect_bad_channels.py @@ -101,6 +101,71 @@ def test_detect_bad_channels_std_mad(): ), "wrong channels locations." +@pytest.mark.parametrize("pool_engine", ["thread", "process"]) +def test_detect_bad_channels_parallel(pool_engine): + recording = generate_recording(num_channels=16, durations=[1, 1], seed=0) + recording.set_channel_gains(1) + recording.set_channel_offsets(0) + method_kwargs = dict( + method="coherence+psd", + num_random_chunks=4, + chunk_duration_s=0.05, + seed=0, + ) + + expected_bad_channel_ids, expected_channel_labels = detect_bad_channels(recording, **method_kwargs) + job_kwargs = dict(n_jobs=2, pool_engine=pool_engine, max_threads_per_worker=1) + if pool_engine == "process": + job_kwargs["mp_context"] = "spawn" + + bad_channel_ids, channel_labels = detect_bad_channels( + recording, + **method_kwargs, + job_kwargs=job_kwargs, + ) + + np.testing.assert_array_equal(bad_channel_ids, expected_bad_channel_ids) + np.testing.assert_array_equal(channel_labels, expected_channel_labels) + + +def test_detect_bad_channels_parallel_unfiltered_spawn(): + """ + generate_recording() marks its output as already filtered, so the parallel test above never + exercises the highpass_filter() wrapper that detect_bad_channels builds internally for an + unfiltered recording. Use a plain NumpyRecording (is_filtered() defaults to False) with a + spawned process pool, so that wrapper has to survive cross-process serialization. + """ + num_channels = 16 + sampling_frequency = 30000.0 + rng = np.random.default_rng(0) + traces_list = [rng.standard_normal((int(sampling_frequency), num_channels)).astype("float32") for _ in range(2)] + recording = NumpyRecording(traces_list, sampling_frequency) + recording.set_channel_gains(1) + recording.set_channel_offsets(0) + probe = generate_linear_probe(num_elec=num_channels) + probe.set_device_channel_indices(np.arange(num_channels)) + recording.set_probe(probe) + assert not recording.is_filtered() + + method_kwargs = dict( + method="coherence+psd", + num_random_chunks=4, + chunk_duration_s=0.05, + seed=0, + ) + + expected_bad_channel_ids, expected_channel_labels = detect_bad_channels(recording, **method_kwargs) + + bad_channel_ids, channel_labels = detect_bad_channels( + recording, + **method_kwargs, + job_kwargs=dict(n_jobs=2, pool_engine="process", mp_context="spawn", max_threads_per_worker=1), + ) + + np.testing.assert_array_equal(bad_channel_ids, expected_bad_channel_ids) + np.testing.assert_array_equal(channel_labels, expected_channel_labels) + + @pytest.mark.parametrize("outside_channels_location", ["bottom", "top", "both"]) def test_detect_bad_channels_extremes(outside_channels_location): num_channels = 64 From f2e2934a55688f3c201a95aa92a25eb5c9111c96 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jes=C3=BAs=20Royeth?= <42451234+JESUSROYETH@users.noreply.github.com> Date: Thu, 3 Sep 2026 08:04:49 -0400 Subject: [PATCH 2/5] Apply batched suggestions from code review Co-authored-by: Alessio Buccino --- .../preprocessing/detect_bad_channels.py | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/src/spikeinterface/preprocessing/detect_bad_channels.py b/src/spikeinterface/preprocessing/detect_bad_channels.py index 9bf4b17728..28b445fc4a 100644 --- a/src/spikeinterface/preprocessing/detect_bad_channels.py +++ b/src/spikeinterface/preprocessing/detect_bad_channels.py @@ -288,13 +288,8 @@ def detect_bad_channels( channel_labels[mask] = "noise" elif method == "coherence+psd": - if job_kwargs is None: - job_kwargs = {"progress_bar": False} - job_kwargs = fix_job_kwargs(job_kwargs) - executor_job_kwargs = { - key: job_kwargs[key] - for key in ("pool_engine", "n_jobs", "progress_bar", "mp_context", "max_threads_per_worker") - } + job_kwargs = {} if job_kwargs is None else job_kwargs + job_kwargs = fix_job_kwargs(job_kwargs) # some checks assert recording.has_scaleable_traces(), ( @@ -336,7 +331,7 @@ def detect_bad_channels( handle_returns=True, chunk_size=random_chunk_kwargs["chunk_size"], job_name="detect_bad_channels", - **executor_job_kwargs, + **job_kwargs, ) chunk_channel_labels = np.stack(executor.run(slices=random_slices), axis=1) From 7f84adbd44ba5b1cd2d7a163eaef4514c148a532 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 3 Sep 2026 12:05:17 +0000 Subject: [PATCH 3/5] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/spikeinterface/preprocessing/detect_bad_channels.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/spikeinterface/preprocessing/detect_bad_channels.py b/src/spikeinterface/preprocessing/detect_bad_channels.py index 28b445fc4a..794f9c0fcc 100644 --- a/src/spikeinterface/preprocessing/detect_bad_channels.py +++ b/src/spikeinterface/preprocessing/detect_bad_channels.py @@ -288,8 +288,8 @@ def detect_bad_channels( channel_labels[mask] = "noise" elif method == "coherence+psd": - job_kwargs = {} if job_kwargs is None else job_kwargs - job_kwargs = fix_job_kwargs(job_kwargs) + job_kwargs = {} if job_kwargs is None else job_kwargs + job_kwargs = fix_job_kwargs(job_kwargs) # some checks assert recording.has_scaleable_traces(), ( From 3c28bd374dfa5f8ccbf5253da289b1f28b687d1a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jes=C3=BAs=20Royeth?= Date: Fri, 4 Sep 2026 11:59:19 -0400 Subject: [PATCH 4/5] Refactor coherence-based bad channel detection --- .../preprocessing/detect_bad_channels.py | 143 +++++++++--------- 1 file changed, 71 insertions(+), 72 deletions(-) diff --git a/src/spikeinterface/preprocessing/detect_bad_channels.py b/src/spikeinterface/preprocessing/detect_bad_channels.py index 794f9c0fcc..e29d1889dd 100644 --- a/src/spikeinterface/preprocessing/detect_bad_channels.py +++ b/src/spikeinterface/preprocessing/detect_bad_channels.py @@ -161,11 +161,11 @@ def _get_all_detect_bad_channel_kwargs(detect_bad_channels_kwargs): return all_detect_bad_channels_kwargs -def _detect_bad_channels_chunk_init(recording, method_kwargs): +def _detect_bad_channels_coherence_psd_chunk_init(recording, method_kwargs): return {"recording": recording, "method_kwargs": method_kwargs} -def _detect_bad_channels_chunk(segment_index, start_frame, end_frame, worker_context): +def _detect_bad_channels_coherence_psd_chunk(segment_index, start_frame, end_frame, worker_context): recording = worker_context["recording"] method_kwargs = worker_context["method_kwargs"] @@ -263,31 +263,10 @@ def detect_bad_channels( else: recording_hp = recording - # Adjust random chunk kwargs based on method - if method in ("std", "mad"): - random_chunk_kwargs["return_in_uV"] = False - random_chunk_kwargs["concatenated"] = True - elif method == "neighborhood_r2": - random_chunk_kwargs["return_in_uV"] = False - random_chunk_kwargs["concatenated"] = False - - if method != "coherence+psd": - random_data = get_random_data_chunks(recording_hp, **random_chunk_kwargs) - channel_labels = np.zeros(recording.get_num_channels(), dtype="U5") channel_labels[:] = "good" - if method in ("std", "mad"): - if method == "std": - deviations = np.std(random_data, axis=0) - else: - deviations = median_abs_deviation(random_data, axis=0) - thresh = std_mad_threshold * np.median(deviations) - mask = deviations > thresh - bad_channel_ids = recording.channel_ids[mask] - channel_labels[mask] = "noise" - - elif method == "coherence+psd": + if method == "coherence+psd": job_kwargs = {} if job_kwargs is None else job_kwargs job_kwargs = fix_job_kwargs(job_kwargs) @@ -325,8 +304,8 @@ def detect_bad_channels( random_slices = get_random_sample_slices(recording_hp, **random_chunk_kwargs) executor = TimeSeriesChunkExecutor( recording_hp, - _detect_bad_channels_chunk, - _detect_bad_channels_chunk_init, + _detect_bad_channels_coherence_psd_chunk, + _detect_bad_channels_coherence_psd_chunk_init, (recording_hp, method_kwargs), handle_returns=True, chunk_size=random_chunk_kwargs["chunk_size"], @@ -371,52 +350,72 @@ def detect_bad_channels( filtered_bad_channel_mask = np.isin(channel_labels, list(channel_filters)) bad_channel_ids = recording.channel_ids[filtered_bad_channel_mask] - elif method == "neighborhood_r2": - # make neighboring channels structure. this should probably be a function in core. - geom = recording.get_channel_locations() - num_channels = recording.get_num_channels() - chan_distances = np.linalg.norm(geom[:, None, :] - geom[None, :, :], axis=2) - np.fill_diagonal(chan_distances, neighborhood_r2_radius_um + 1) - neighbors_mask = chan_distances < neighborhood_r2_radius_um - if neighbors_mask.sum(axis=1).min() < 1: - warnings.warn( - f"neighborhood_r2_radius_um={neighborhood_r2_radius_um} led " - "to channels with no neighbors for this geometry, which has " - f"minimal channel distance {chan_distances.min()}um. These " - "channels will not be marked as bad, but you might want to " - "check them." - ) - max_neighbors = neighbors_mask.sum(axis=1).max() - channel_index = np.full((num_channels, max_neighbors), num_channels) - for c in range(num_channels): - my_neighbors = np.flatnonzero(neighbors_mask[c]) - channel_index[c, : my_neighbors.size] = my_neighbors - - # get the correlation of each channel with its neighbors' median inside each chunk - # note that we did not concatenate the chunks here - correlations = [] - for chunk in random_data: - chunk = chunk.astype(np.float32, copy=False) - chunk = chunk - np.median(chunk, axis=0, keepdims=True) - padded_chunk = np.pad(chunk, [(0, 0), (0, 1)], constant_values=np.nan) - # channels with no neighbors will get a pure-nan median trace here - neighbmeans = np.nanmedian( - padded_chunk[:, channel_index], - axis=2, - ) - denom = np.sqrt(np.nanmean(np.square(chunk), axis=0) * np.nanmean(np.square(neighbmeans), axis=0)) - denom[denom == 0] = 1 - # channels with no neighbors will get a nan here - chunk_correlations = np.nanmean(chunk * neighbmeans, axis=0) / denom - correlations.append(chunk_correlations) - - # now take the median over chunks and threshold to finish - median_correlations = np.nanmedian(correlations, 0) - r2s = median_correlations**2 - # channels with no neighbors will have r2==nan, and nan thresh + bad_channel_ids = recording.channel_ids[mask] + channel_labels[mask] = "noise" + + else: # neighborhood_r2 + # make neighboring channels structure. this should probably be a function in core. + geom = recording.get_channel_locations() + num_channels = recording.get_num_channels() + chan_distances = np.linalg.norm(geom[:, None, :] - geom[None, :, :], axis=2) + np.fill_diagonal(chan_distances, neighborhood_r2_radius_um + 1) + neighbors_mask = chan_distances < neighborhood_r2_radius_um + if neighbors_mask.sum(axis=1).min() < 1: + warnings.warn( + f"neighborhood_r2_radius_um={neighborhood_r2_radius_um} led " + "to channels with no neighbors for this geometry, which has " + f"minimal channel distance {chan_distances.min()}um. These " + "channels will not be marked as bad, but you might want to " + "check them." + ) + max_neighbors = neighbors_mask.sum(axis=1).max() + channel_index = np.full((num_channels, max_neighbors), num_channels) + for c in range(num_channels): + my_neighbors = np.flatnonzero(neighbors_mask[c]) + channel_index[c, : my_neighbors.size] = my_neighbors + + # get the correlation of each channel with its neighbors' median inside each chunk + # note that we did not concatenate the chunks here + correlations = [] + for chunk in random_data: + chunk = chunk.astype(np.float32, copy=False) + chunk = chunk - np.median(chunk, axis=0, keepdims=True) + padded_chunk = np.pad(chunk, [(0, 0), (0, 1)], constant_values=np.nan) + # channels with no neighbors will get a pure-nan median trace here + neighbmeans = np.nanmedian( + padded_chunk[:, channel_index], + axis=2, + ) + denom = np.sqrt(np.nanmean(np.square(chunk), axis=0) * np.nanmean(np.square(neighbmeans), axis=0)) + denom[denom == 0] = 1 + # channels with no neighbors will get a nan here + chunk_correlations = np.nanmean(chunk * neighbmeans, axis=0) / denom + correlations.append(chunk_correlations) + + # now take the median over chunks and threshold to finish + median_correlations = np.nanmedian(correlations, 0) + r2s = median_correlations**2 + # channels with no neighbors will have r2==nan, and nan Date: Fri, 4 Sep 2026 12:08:41 -0400 Subject: [PATCH 5/5] Revert refactor of bad channel detection --- .../preprocessing/detect_bad_channels.py | 143 +++++++++--------- 1 file changed, 72 insertions(+), 71 deletions(-) diff --git a/src/spikeinterface/preprocessing/detect_bad_channels.py b/src/spikeinterface/preprocessing/detect_bad_channels.py index e29d1889dd..794f9c0fcc 100644 --- a/src/spikeinterface/preprocessing/detect_bad_channels.py +++ b/src/spikeinterface/preprocessing/detect_bad_channels.py @@ -161,11 +161,11 @@ def _get_all_detect_bad_channel_kwargs(detect_bad_channels_kwargs): return all_detect_bad_channels_kwargs -def _detect_bad_channels_coherence_psd_chunk_init(recording, method_kwargs): +def _detect_bad_channels_chunk_init(recording, method_kwargs): return {"recording": recording, "method_kwargs": method_kwargs} -def _detect_bad_channels_coherence_psd_chunk(segment_index, start_frame, end_frame, worker_context): +def _detect_bad_channels_chunk(segment_index, start_frame, end_frame, worker_context): recording = worker_context["recording"] method_kwargs = worker_context["method_kwargs"] @@ -263,10 +263,31 @@ def detect_bad_channels( else: recording_hp = recording + # Adjust random chunk kwargs based on method + if method in ("std", "mad"): + random_chunk_kwargs["return_in_uV"] = False + random_chunk_kwargs["concatenated"] = True + elif method == "neighborhood_r2": + random_chunk_kwargs["return_in_uV"] = False + random_chunk_kwargs["concatenated"] = False + + if method != "coherence+psd": + random_data = get_random_data_chunks(recording_hp, **random_chunk_kwargs) + channel_labels = np.zeros(recording.get_num_channels(), dtype="U5") channel_labels[:] = "good" - if method == "coherence+psd": + if method in ("std", "mad"): + if method == "std": + deviations = np.std(random_data, axis=0) + else: + deviations = median_abs_deviation(random_data, axis=0) + thresh = std_mad_threshold * np.median(deviations) + mask = deviations > thresh + bad_channel_ids = recording.channel_ids[mask] + channel_labels[mask] = "noise" + + elif method == "coherence+psd": job_kwargs = {} if job_kwargs is None else job_kwargs job_kwargs = fix_job_kwargs(job_kwargs) @@ -304,8 +325,8 @@ def detect_bad_channels( random_slices = get_random_sample_slices(recording_hp, **random_chunk_kwargs) executor = TimeSeriesChunkExecutor( recording_hp, - _detect_bad_channels_coherence_psd_chunk, - _detect_bad_channels_coherence_psd_chunk_init, + _detect_bad_channels_chunk, + _detect_bad_channels_chunk_init, (recording_hp, method_kwargs), handle_returns=True, chunk_size=random_chunk_kwargs["chunk_size"], @@ -350,72 +371,52 @@ def detect_bad_channels( filtered_bad_channel_mask = np.isin(channel_labels, list(channel_filters)) bad_channel_ids = recording.channel_ids[filtered_bad_channel_mask] - else: - if method in ("std", "mad"): - random_chunk_kwargs["return_in_uV"] = False - random_chunk_kwargs["concatenated"] = True - else: # neighborhood_r2 - random_chunk_kwargs["return_in_uV"] = False - random_chunk_kwargs["concatenated"] = False - - random_data = get_random_data_chunks(recording_hp, **random_chunk_kwargs) - - if method in ("std", "mad"): - if method == "std": - deviations = np.std(random_data, axis=0) - else: - deviations = median_abs_deviation(random_data, axis=0) - thresh = std_mad_threshold * np.median(deviations) - mask = deviations > thresh - bad_channel_ids = recording.channel_ids[mask] - channel_labels[mask] = "noise" - - else: # neighborhood_r2 - # make neighboring channels structure. this should probably be a function in core. - geom = recording.get_channel_locations() - num_channels = recording.get_num_channels() - chan_distances = np.linalg.norm(geom[:, None, :] - geom[None, :, :], axis=2) - np.fill_diagonal(chan_distances, neighborhood_r2_radius_um + 1) - neighbors_mask = chan_distances < neighborhood_r2_radius_um - if neighbors_mask.sum(axis=1).min() < 1: - warnings.warn( - f"neighborhood_r2_radius_um={neighborhood_r2_radius_um} led " - "to channels with no neighbors for this geometry, which has " - f"minimal channel distance {chan_distances.min()}um. These " - "channels will not be marked as bad, but you might want to " - "check them." - ) - max_neighbors = neighbors_mask.sum(axis=1).max() - channel_index = np.full((num_channels, max_neighbors), num_channels) - for c in range(num_channels): - my_neighbors = np.flatnonzero(neighbors_mask[c]) - channel_index[c, : my_neighbors.size] = my_neighbors - - # get the correlation of each channel with its neighbors' median inside each chunk - # note that we did not concatenate the chunks here - correlations = [] - for chunk in random_data: - chunk = chunk.astype(np.float32, copy=False) - chunk = chunk - np.median(chunk, axis=0, keepdims=True) - padded_chunk = np.pad(chunk, [(0, 0), (0, 1)], constant_values=np.nan) - # channels with no neighbors will get a pure-nan median trace here - neighbmeans = np.nanmedian( - padded_chunk[:, channel_index], - axis=2, - ) - denom = np.sqrt(np.nanmean(np.square(chunk), axis=0) * np.nanmean(np.square(neighbmeans), axis=0)) - denom[denom == 0] = 1 - # channels with no neighbors will get a nan here - chunk_correlations = np.nanmean(chunk * neighbmeans, axis=0) / denom - correlations.append(chunk_correlations) - - # now take the median over chunks and threshold to finish - median_correlations = np.nanmedian(correlations, 0) - r2s = median_correlations**2 - # channels with no neighbors will have r2==nan, and nan