diff --git a/src/spikeinterface/core/base.py b/src/spikeinterface/core/base.py index eae31791a3..0f11d1c00e 100644 --- a/src/spikeinterface/core/base.py +++ b/src/spikeinterface/core/base.py @@ -287,6 +287,28 @@ def get_annotation(self, key: str, copy: bool = True) -> Any: def get_annotation_keys(self) -> list: return list(self._annotations.keys()) + def _invalidate_noise_levels_for_calibration(self, key): + """ + Invalidate scaled noise-level estimates when calibration properties change. + """ + if key == "gain_to_uV": + noise_keys = [ + "noise_level_mad_scaled", + "noise_level_std_scaled", + "noise_level_rms_scaled", + ] + elif key == "offset_to_uV": + noise_keys = ["noise_level_rms_scaled"] + else: + return + + noise_levels_params = getattr(self, "_noise_levels_params", None) + + for noise_key in noise_keys: + self._properties.pop(noise_key, None) + if noise_levels_params is not None: + noise_levels_params.pop(noise_key, None) + def set_property( self, key, @@ -318,6 +340,13 @@ def set_property( The `missing_value` is required for types int and unsigned int. """ + # Invalidate noise-level sampling metadata when the property is updated. + noise_levels_params = getattr(self, "_noise_levels_params", None) + if noise_levels_params is not None: + noise_levels_params.pop(key, None) + + self._invalidate_noise_levels_for_calibration(key) + if values is None: if key in self._properties: self._properties.pop(key) @@ -406,6 +435,11 @@ def get_property_keys(self) -> list: def delete_property(self, key) -> None: if key in self._properties: del self._properties[key] + noise_levels_params = getattr(self, "_noise_levels_params", None) + if noise_levels_params is not None: + noise_levels_params.pop(key, None) + + self._invalidate_noise_levels_for_calibration(key) else: raise Exception(f"{key} is not a property key") diff --git a/src/spikeinterface/core/recording_tools.py b/src/spikeinterface/core/recording_tools.py index fb37cb44b4..177d889d56 100644 --- a/src/spikeinterface/core/recording_tools.py +++ b/src/spikeinterface/core/recording_tools.py @@ -10,6 +10,7 @@ from probeinterface import ProbeGroup from .job_tools import ( + chunk_duration_to_chunk_size, ensure_chunk_size, divide_segment_into_chunks, fix_job_kwargs, @@ -470,6 +471,72 @@ def get_closest_channels(recording, channel_ids=None, num_channels=None): return np.array(closest_channels_inds), np.array(dists) +_NOISE_LEVEL_SAMPLE_PARAMS_DEFAULTS = { + "method": "full_random", + "num_chunks_per_segment": 20, + "chunk_duration": "500ms", + "chunk_size": None, + "margin_frames": 0, + "seed": None, +} + + +def _normalize_noise_level_params(recording, random_slices_kwargs): + """Return the effective sampling parameters used by ``get_random_sample_slices``. + + ``chunk_duration`` is converted to samples, since it is only an alternative way + of specifying ``chunk_size``. Chunk sizes are also clipped in the same way as + ``get_random_sample_slices`` for recordings with short segments. + """ + params = _NOISE_LEVEL_SAMPLE_PARAMS_DEFAULTS.copy() + unexpected_kwargs = set(random_slices_kwargs).difference(params) + if unexpected_kwargs: + unexpected_key = next(iter(unexpected_kwargs)) + raise TypeError(f"get_random_sample_slices() got an unexpected keyword argument {unexpected_key!r}") + params.update(random_slices_kwargs) + + if params["method"] != "full_random": + raise ValueError(f"get_random_sample_slices : wrong method {params['method']}") + + chunk_size = params["chunk_size"] + if chunk_size is None: + chunk_duration = params["chunk_duration"] + if chunk_duration is None: + raise ValueError("get_random_sample_slices need chunk_size or chunk_duration") + chunk_size = chunk_duration_to_chunk_size(chunk_duration, recording) + + # Match the effective chunk-size adjustment performed by get_random_sample_slices. + margin_frames = params["margin_frames"] + for segment_index in range(recording.get_num_segments()): + chunk_size_limit = recording.get_num_samples(segment_index) - 2 * margin_frames + if chunk_size > chunk_size_limit: + chunk_size = chunk_size_limit - 1 + + return { + "method": params["method"], + "num_chunks_per_segment": params["num_chunks_per_segment"], + "chunk_size": chunk_size, + "margin_frames": margin_frames, + "seed": params["seed"], + } + + +def _noise_level_params_match(stored_params, requested_params): + """Check whether a cached estimate is compatible with a sampling request. + + A requested ``seed=None`` is a wildcard, while all other sampling parameters + must match exactly. + """ + if stored_params is None: + return False + + stored_params = dict(stored_params) + requested_params = dict(requested_params) + requested_seed = requested_params.pop("seed") + stored_seed = stored_params.pop("seed", None) + return stored_params == requested_params and (requested_seed is None or requested_seed == stored_seed) + + def _noise_level_chunk(segment_index, start_frame, end_frame, worker_ctx): recording = worker_ctx["recording"] @@ -531,8 +598,8 @@ def get_noise_levels( If True, returned noise levels are scaled to uV method : "mad" | "std" | "rms", default: "mad" The method to use to estimate noise levels - force_recompute : bool - If True, noise levels are recomputed even if they are already stored in the recording extractor + force_recompute : bool, default: False + If True, noise levels are recomputed even if a compatible estimate is already stored random_slices_kwargs : dict Options transmitted to get_random_sample_slices(), please read documentation from this function for more details. @@ -549,52 +616,67 @@ def get_noise_levels( else: key = f"noise_level_{method}_raw" - if key in recording.get_property_keys() and not force_recompute: - noise_levels = recording.get_property(key=key) - else: - # This is to keep backward compatibility - # lets keep for a while and remove this maybe in 0.103.0 - # chunk_size used to be in the signature and now is ambiguous - random_slices_kwargs_, job_kwargs = split_job_kwargs(kwargs) - if len(random_slices_kwargs_) > 0 or "chunk_size" in job_kwargs: - msg = ( - "get_noise_levels(recording, num_chunks_per_segment=20) is deprecated\n" - "Now, you need to use get_noise_levels(recording, random_slices_kwargs=dict(num_chunks_per_segment=20, chunk_size=1000))\n" - "Please read get_random_sample_slices() documentation for more options." - ) - # if the user use both the old and the new behavior then an error is raised - assert len(random_slices_kwargs) == 0, msg - warnings.warn(msg) - random_slices_kwargs = random_slices_kwargs_ - if "chunk_size" in job_kwargs: - random_slices_kwargs["chunk_size"] = job_kwargs["chunk_size"] - - slices = get_random_sample_slices(recording, **random_slices_kwargs) - - noise_levels_chunks = [] - - def append_noise_chunk(res): - noise_levels_chunks.append(res) - - func = _noise_level_chunk - init_func = _noise_level_chunk_init - init_args = (recording, return_in_uV, method) - executor = TimeSeriesChunkExecutor( - recording, - func, - init_func, - init_args, - job_name="noise_level", - verbose=False, - gather_func=append_noise_chunk, - **job_kwargs, + # This is to keep backward compatibility with the legacy keyword arguments. + # ``chunk_size`` used to be part of the get_noise_levels signature and is + # ambiguous with the job-level chunk_size option. + random_slices_kwargs = {} if random_slices_kwargs is None else dict(random_slices_kwargs) + random_slices_kwargs_legacy, job_kwargs = split_job_kwargs(kwargs) + if len(random_slices_kwargs_legacy) > 0 or "chunk_size" in job_kwargs: + msg = ( + "get_noise_levels(recording, num_chunks_per_segment=20) is deprecated\n" + "Now, you need to use get_noise_levels(recording, random_slices_kwargs=dict(num_chunks_per_segment=20, chunk_size=1000))\n" + "Please read get_random_sample_slices() documentation for more options." ) - executor.run(slices=slices) - noise_levels_chunks = np.stack(noise_levels_chunks) - noise_levels = np.mean(noise_levels_chunks, axis=0) + # If the user uses both the old and the new behavior then an error is raised. + assert len(random_slices_kwargs) == 0, msg + warnings.warn(msg) + random_slices_kwargs = random_slices_kwargs_legacy + if "chunk_size" in job_kwargs: + random_slices_kwargs["chunk_size"] = job_kwargs["chunk_size"] + + requested_params = _normalize_noise_level_params(recording, random_slices_kwargs) + default_params = _normalize_noise_level_params(recording, {}) + noise_levels_params = getattr(recording, "_noise_levels_params", {}) + + if key in recording.get_property_keys() and not force_recompute: + stored_params = noise_levels_params.get(key) + if stored_params is None: + # Preserve manually supplied and loaded properties for the default call. + if requested_params == default_params: + return recording.get_property(key=key) + elif _noise_level_params_match(stored_params, requested_params): + return recording.get_property(key=key) + + slices = get_random_sample_slices(recording, **random_slices_kwargs) - # set property - recording.set_property(key, noise_levels) + noise_levels_chunks = [] + + def append_noise_chunk(res): + noise_levels_chunks.append(res) + + func = _noise_level_chunk + init_func = _noise_level_chunk_init + init_args = (recording, return_in_uV, method) + executor = TimeSeriesChunkExecutor( + recording, + func, + init_func, + init_args, + job_name="noise_level", + verbose=False, + gather_func=append_noise_chunk, + **job_kwargs, + ) + executor.run(slices=slices) + noise_levels_chunks = np.stack(noise_levels_chunks) + noise_levels = np.mean(noise_levels_chunks, axis=0) + + # Store the property and the parameters that produced it. set_property clears + # stale metadata, so attach the new metadata only after setting the property. + recording.set_property(key, noise_levels) + if not hasattr(recording, "_noise_levels_params"): + recording._noise_levels_params = {} + recording._noise_levels_params[key] = requested_params return noise_levels diff --git a/src/spikeinterface/core/tests/test_recording_tools.py b/src/spikeinterface/core/tests/test_recording_tools.py index 0166c3bfdd..7464b795fe 100644 --- a/src/spikeinterface/core/tests/test_recording_tools.py +++ b/src/spikeinterface/core/tests/test_recording_tools.py @@ -1,4 +1,5 @@ import numpy as np +import spikeinterface.core.recording_tools as recording_tools from spikeinterface.core import NumpyRecording, generate_recording, MockRecording @@ -229,6 +230,279 @@ def test_get_noise_levels_output(): assert np.allclose(std_estimated_with_std, [std, std], rtol=1e-2, atol=1e-3) +def test_get_noise_levels_cache_parameters(monkeypatch): + """Sampling parameters must be part of the noise-level cache semantics.""" + recording = generate_recording(num_channels=2, sampling_frequency=1000.0, durations=[2.0]) + computed_chunks = [] + + def fake_noise_level_chunk(segment_index, start_frame, end_frame, worker_ctx): + value = float(len(computed_chunks) + 1) + computed_chunks.append((segment_index, start_frame, end_frame)) + return np.full(recording.get_num_channels(), value) + + monkeypatch.setattr(recording_tools, "_noise_level_chunk", fake_noise_level_chunk) + sampling_kwargs = {"num_chunks_per_segment": 1, "chunk_size": 100, "seed": 42} + + first = get_noise_levels( + recording, return_in_uV=False, random_slices_kwargs=sampling_kwargs, n_jobs=1, progress_bar=False + ) + np.testing.assert_array_equal(first, [1.0, 1.0]) + + # Identical parameters hit the cache, and seed=None is a wildcard for the seed. + cached = get_noise_levels( + recording, return_in_uV=False, random_slices_kwargs=sampling_kwargs, n_jobs=1, progress_bar=False + ) + cached_any_seed = get_noise_levels( + recording, + return_in_uV=False, + random_slices_kwargs={**sampling_kwargs, "seed": None}, + n_jobs=1, + progress_bar=False, + ) + np.testing.assert_array_equal(cached, first) + np.testing.assert_array_equal(cached_any_seed, first) + assert len(computed_chunks) == 1 + + # A different explicit seed and a different chunk layout each trigger recomputation. + different_seed = get_noise_levels( + recording, + return_in_uV=False, + random_slices_kwargs={**sampling_kwargs, "seed": 123}, + n_jobs=1, + progress_bar=False, + ) + np.testing.assert_array_equal(different_seed, [2.0, 2.0]) + different_layout = get_noise_levels( + recording, + return_in_uV=False, + random_slices_kwargs={**sampling_kwargs, "seed": 123, "num_chunks_per_segment": 2}, + n_jobs=1, + progress_bar=False, + ) + np.testing.assert_array_equal(different_layout, [3.5, 3.5]) + assert len(computed_chunks) == 4 + + # The new layout can be reused when the caller does not care about the seed. + cached_any_seed = get_noise_levels( + recording, + return_in_uV=False, + random_slices_kwargs={**sampling_kwargs, "seed": None, "num_chunks_per_segment": 2}, + n_jobs=1, + progress_bar=False, + ) + np.testing.assert_array_equal(cached_any_seed, different_layout) + assert len(computed_chunks) == 4 + + # force_recompute remains an unconditional override. + forced = get_noise_levels( + recording, + return_in_uV=False, + force_recompute=True, + random_slices_kwargs={**sampling_kwargs, "seed": 123, "num_chunks_per_segment": 2}, + n_jobs=1, + progress_bar=False, + ) + np.testing.assert_array_equal(forced, [5.5, 5.5]) + assert len(computed_chunks) == 6 + + # Overwriting a computed property turns it back into user-supplied state. + key = "noise_level_mad_raw" + manual_values = np.full(recording.get_num_channels(), 9.0) + recording.set_property(key, manual_values) + assert key not in recording._noise_levels_params + np.testing.assert_array_equal( + get_noise_levels(recording, return_in_uV=False, n_jobs=1, progress_bar=False), manual_values + ) + + # An explicit sampling request recomputes the manually supplied value. + get_noise_levels(recording, return_in_uV=False, random_slices_kwargs=sampling_kwargs, n_jobs=1, progress_bar=False) + assert key in recording._noise_levels_params + recording.delete_property(key) + assert key not in recording._noise_levels_params + + +def test_noise_levels_cache_invalidated_on_gain_change(): + recording = generate_recording(num_channels=4, durations=[5.0]) + recording.set_channel_gains([1.0] * 4) + recording.set_channel_offsets([0.0] * 4) + + sampling_kwargs = {"seed": 42} + + scaled_before = {} + raw_before = {} + + for method in ("mad", "std", "rms"): + scaled_before[method] = get_noise_levels( + recording, + method=method, + return_in_uV=True, + random_slices_kwargs=sampling_kwargs, + n_jobs=1, + progress_bar=False, + ).copy() + + raw_before[method] = get_noise_levels( + recording, + method=method, + return_in_uV=False, + random_slices_kwargs=sampling_kwargs, + n_jobs=1, + progress_bar=False, + ).copy() + + recording.set_channel_gains([10.0] * 4) + + for method in ("mad", "std", "rms"): + scaled_after = get_noise_levels( + recording, + method=method, + return_in_uV=True, + random_slices_kwargs=sampling_kwargs, + n_jobs=1, + progress_bar=False, + ) + + raw_after = get_noise_levels( + recording, + method=method, + return_in_uV=False, + random_slices_kwargs=sampling_kwargs, + n_jobs=1, + progress_bar=False, + ) + + np.testing.assert_allclose(scaled_after, scaled_before[method] * 10, rtol=1e-4) + np.testing.assert_array_equal(raw_after, raw_before[method]) + + +def test_noise_levels_cache_invalidated_on_offset_change(): + recording = generate_recording(num_channels=4, durations=[5.0]) + recording.set_channel_gains([1.0] * 4) + recording.set_channel_offsets([0.0] * 4) + + sampling_kwargs = {"seed": 42} + + rms_before = get_noise_levels( + recording, + method="rms", + return_in_uV=True, + random_slices_kwargs=sampling_kwargs, + n_jobs=1, + progress_bar=False, + ).copy() + + std_before = get_noise_levels( + recording, + method="std", + return_in_uV=True, + random_slices_kwargs=sampling_kwargs, + n_jobs=1, + progress_bar=False, + ).copy() + + mad_before = get_noise_levels( + recording, + method="mad", + return_in_uV=True, + random_slices_kwargs=sampling_kwargs, + n_jobs=1, + progress_bar=False, + ).copy() + + recording.set_channel_offsets([100.0] * 4) + + rms_after = get_noise_levels( + recording, + method="rms", + return_in_uV=True, + random_slices_kwargs=sampling_kwargs, + n_jobs=1, + progress_bar=False, + ).copy() + + std_after = get_noise_levels( + recording, + method="std", + return_in_uV=True, + random_slices_kwargs=sampling_kwargs, + n_jobs=1, + progress_bar=False, + ).copy() + + mad_after = get_noise_levels( + recording, + method="mad", + return_in_uV=True, + random_slices_kwargs=sampling_kwargs, + n_jobs=1, + progress_bar=False, + ).copy() + + assert not np.allclose(rms_before, rms_after) + np.testing.assert_allclose(std_after, std_before, rtol=1e-5) + np.testing.assert_allclose(mad_after, mad_before, rtol=1e-5) + + +def test_noise_levels_cache_invalidated_on_calibration_property_deletion(): + recording = generate_recording(num_channels=4, durations=[5.0]) + recording.set_channel_gains([1.0] * 4) + recording.set_channel_offsets([0.0] * 4) + + sampling_kwargs = {"seed": 42} + + for method in ("mad", "std", "rms"): + get_noise_levels( + recording, + method=method, + return_in_uV=True, + random_slices_kwargs=sampling_kwargs, + n_jobs=1, + progress_bar=False, + ) + + scaled_keys = [ + "noise_level_mad_scaled", + "noise_level_std_scaled", + "noise_level_rms_scaled", + ] + + for key in scaled_keys: + assert key in recording.get_property_keys() + + recording.delete_property("gain_to_uV") + + for key in scaled_keys: + assert key not in recording.get_property_keys() + + +def test_get_noise_levels_normalizes_chunk_duration(monkeypatch): + """Equivalent duration- and frame-based chunk specifications share the cache.""" + recording = generate_recording(num_channels=1, sampling_frequency=1000.0, durations=[2.0]) + computed_chunks = [] + + def fake_noise_level_chunk(segment_index, start_frame, end_frame, worker_ctx): + computed_chunks.append((segment_index, start_frame, end_frame)) + return np.ones(recording.get_num_channels()) + + monkeypatch.setattr(recording_tools, "_noise_level_chunk", fake_noise_level_chunk) + by_duration = get_noise_levels( + recording, + return_in_uV=False, + random_slices_kwargs={"num_chunks_per_segment": 1, "seed": 0}, + n_jobs=1, + progress_bar=False, + ) + by_frames = get_noise_levels( + recording, + return_in_uV=False, + random_slices_kwargs={"num_chunks_per_segment": 1, "chunk_size": 500, "seed": 0}, + n_jobs=1, + progress_bar=False, + ) + np.testing.assert_array_equal(by_frames, by_duration) + assert len(computed_chunks) == 1 + + def test_get_chunk_with_margin(): rec = generate_recording(num_channels=1, sampling_frequency=1000.0, durations=[10.0]) rec_seg = rec.segments[0] diff --git a/src/spikeinterface/metrics/quality/misc_metrics.py b/src/spikeinterface/metrics/quality/misc_metrics.py index 016a2d24f9..5a6f31dbbb 100644 --- a/src/spikeinterface/metrics/quality/misc_metrics.py +++ b/src/spikeinterface/metrics/quality/misc_metrics.py @@ -1362,8 +1362,8 @@ def compute_sd_ratio( If true, will take into account that the template itself impacts the standard deviation of the noise, and will make a rough estimation of what that impact is (and remove it). seed : int | None, default: None - Seed for the random chunks the noise SD is computed on. With a seed, the noise SD is always drawn - again, and replaces any value `get_noise_levels` has cached on the recording. + Seed for the random chunks the noise SD is computed on. Different explicit seeds trigger a cache miss; + ``None`` accepts any cached estimate with the same chunk layout. **job_kwargs : dict, default: {} Keyword arguments sent to get_noise_levels. @@ -1405,13 +1405,12 @@ def compute_sd_ratio( ) return {unit_id: np.nan for unit_id in unit_ids} job_kwargs["progress_bar"] = False - # with a seed, draw again rather than reuse a noise SD cached on the recording - seed_kwargs = {} if seed is None else dict(random_slices_kwargs=dict(seed=seed), force_recompute=True) + random_slices_kwargs = {} if seed is None else dict(seed=seed) noise_levels = get_noise_levels( sorting_analyzer.recording, return_in_uV=sorting_analyzer.return_in_uV, method="std", - **seed_kwargs, + random_slices_kwargs=random_slices_kwargs, **job_kwargs, ) diff --git a/src/spikeinterface/preprocessing/highpass_spatial_filter.py b/src/spikeinterface/preprocessing/highpass_spatial_filter.py index 838672ab4b..873cfa7176 100644 --- a/src/spikeinterface/preprocessing/highpass_spatial_filter.py +++ b/src/spikeinterface/preprocessing/highpass_spatial_filter.py @@ -127,11 +127,10 @@ def __init__( # Compute or retrieve RMS values if rms_values is None: - if "noise_level_rms_raw" in recording.get_property_keys(): - rms_values = recording.get_property("noise_level_rms_raw") - else: - random_slice_kwargs = {} if random_slice_kwargs is None else random_slice_kwargs - rms_values = get_noise_levels(recording, method="rms", return_in_uV=False, **random_slice_kwargs) + random_slice_kwargs = {} if random_slice_kwargs is None else random_slice_kwargs + rms_values = get_noise_levels( + recording, method="rms", return_in_uV=False, random_slices_kwargs=random_slice_kwargs + ) # Pre-compute spatial filtering parameters butter_kwargs = dict(btype="highpass", N=highpass_butter_order, Wn=highpass_butter_wn) diff --git a/src/spikeinterface/preprocessing/tests/test_highpass_spatial_filter.py b/src/spikeinterface/preprocessing/tests/test_highpass_spatial_filter.py index b4579451a1..9116ccec02 100644 --- a/src/spikeinterface/preprocessing/tests/test_highpass_spatial_filter.py +++ b/src/spikeinterface/preprocessing/tests/test_highpass_spatial_filter.py @@ -123,6 +123,29 @@ def test_highpass_spatial_filter_with_dead_channels(): assert result.shape == traces.shape +def test_highpass_spatial_filter_uses_noise_level_cache(monkeypatch): + """The spatial filter must route RMS cache decisions through get_noise_levels.""" + import importlib + + highpass_spatial_filter_module = importlib.import_module("spikeinterface.preprocessing.highpass_spatial_filter") + recording = generate_recording(num_channels=32, sampling_frequency=30000.0, durations=[0.5]) + # Even an existing property must be checked against the requested sampling parameters. + recording.set_property("noise_level_rms_raw", np.ones(recording.get_num_channels())) + received_kwargs = {} + + def fake_get_noise_levels(recording_arg, **kwargs): + received_kwargs.update(kwargs) + return np.ones(recording_arg.get_num_channels()) + + monkeypatch.setattr(highpass_spatial_filter_module, "get_noise_levels", fake_get_noise_levels) + sampling_kwargs = {"seed": 123, "num_chunks_per_segment": 3, "chunk_size": 100} + spre.highpass_spatial_filter(recording, n_channel_pad=2, random_slice_kwargs=sampling_kwargs) + + assert received_kwargs["method"] == "rms" + assert received_kwargs["return_in_uV"] is False + assert received_kwargs["random_slices_kwargs"] == sampling_kwargs + + @pytest.mark.parametrize("dtype", [np.int16, np.float32, np.float64]) def test_dtype_stability(dtype): """