diff --git a/doc/changes/dev/14164.bugfix.rst b/doc/changes/dev/14164.bugfix.rst new file mode 100644 index 00000000000..f927d37dcba --- /dev/null +++ b/doc/changes/dev/14164.bugfix.rst @@ -0,0 +1 @@ +Add CUDA support to :meth:`mne.io.Raw.apply_hilbert` and related data classes, allowing :func:`mne.preprocessing.annotate_muscle_zscore` to use ``n_jobs='cuda'``, by `Daria Agafonova`_. diff --git a/mne/cuda.py b/mne/cuda.py index 2be835c016a..aaa211d4dae 100644 --- a/mne/cuda.py +++ b/mne/cuda.py @@ -122,6 +122,67 @@ def _set_cuda_device(device_id, verbose=None): logger.info(f"Now using CUDA device {device_id}") +def _setup_cuda_hilbert(n_jobs, n_fft): + """Set up CUDA for a Hilbert transform.""" + multiplier = None + if isinstance(n_jobs, str): + _check_option("n_jobs", n_jobs, ("cuda",)) + n_jobs = 1 + init_cuda() + if _cuda_capable: + import cupy + + try: + multiplier = cupy.asarray(_hilbert_multiplier(n_fft, np, dtype=np.int8)) + logger.info("Using CUDA for Hilbert transform") + except Exception as exp: + logger.info( + "CUDA not used, could not allocate the Hilbert multiplier " + f'("{exp}"), falling back to n_jobs=None' + ) + else: + logger.info( + "CUDA not used, CUDA could not be initialized, " + "falling back to n_jobs=None" + ) + return n_jobs, multiplier + + +def _cuda_hilbert(x, n_fft, envelope, multiplier): + """Compute an analytic signal on the GPU.""" + import cupy + + if np.iscomplexobj(x): + raise ValueError("x must be real.") + out = _fft_hilbert(cupy.asarray(x), n_fft, envelope, cupy, multiplier) + return cupy.asnumpy(out) + + +def _hilbert_multiplier(n_fft, xp, dtype): + """Create the frequency-domain multiplier for an analytic signal.""" + multiplier = xp.zeros(n_fft, dtype=dtype) + multiplier[0] = 1 + if n_fft % 2 == 0: + multiplier[1 : n_fft // 2] = 2 + multiplier[n_fft // 2] = 1 + else: + multiplier[1 : (n_fft + 1) // 2] = 2 + return multiplier + + +def _fft_hilbert(x, n_fft, envelope, xp, multiplier=None): + """Compute an analytic signal with a NumPy-compatible array module.""" + n_x = x.shape[-1] + x_fft = xp.fft.fft(x, n=n_fft, axis=-1) + if multiplier is None: + multiplier = _hilbert_multiplier(n_fft, xp, x_fft.real.dtype) + x_fft *= multiplier + out = xp.fft.ifft(x_fft, axis=-1)[..., :n_x] + if envelope: + out = xp.abs(out) + return out + + ############################################################################### # Repeated FFT multiplication diff --git a/mne/filter.py b/mne/filter.py index 8f0cedbf043..75578a9fb34 100644 --- a/mne/filter.py +++ b/mne/filter.py @@ -16,10 +16,12 @@ from ._fiff.pick import _picks_to_idx from ._ola import _COLA from .cuda import ( + _cuda_hilbert, _fft_multiply_repeated, _fft_resample, _setup_cuda_fft_multiply_repeated, _setup_cuda_fft_resample, + _setup_cuda_hilbert, _smart_pad, ) from .fixes import _reshape_view @@ -2695,7 +2697,7 @@ def apply_hilbert( envelope : bool Compute the envelope signal of each channel/vertex. Default False. See Notes. - %(n_jobs)s + %(n_jobs_cuda)s n_fft : int | None | str Points to use in the FFT for Hilbert transformation. The signal will be padded with zeros before computing Hilbert, then cut back @@ -2780,17 +2782,22 @@ def apply_hilbert( if dtype is not None and dtype != self._data.dtype: self._data = self._data.astype(dtype) + n_jobs, cuda_multiplier = _setup_cuda_hilbert(n_jobs, n_fft) + if cuda_multiplier is None: + hilbert_fun = _my_hilbert + else: + hilbert_fun = partial(_cuda_hilbert, multiplier=cuda_multiplier) parallel, p_fun, n_jobs = parallel_func(_check_fun, n_jobs) if n_jobs == 1: # modify data inplace to save memory for idx in picks: self._data[..., idx, :] = _check_fun( - _my_hilbert, data_in[..., idx, :], *args, **kwargs + hilbert_fun, data_in[..., idx, :], *args, **kwargs ) else: # use parallel function data_picks_new = parallel( - p_fun(_my_hilbert, data_in[..., p, :], *args, **kwargs) for p in picks + p_fun(hilbert_fun, data_in[..., p, :], *args, **kwargs) for p in picks ) for pp, p in enumerate(picks): self._data[..., p, :] = data_picks_new[pp] diff --git a/mne/preprocessing/artifact_detection.py b/mne/preprocessing/artifact_detection.py index 7d286958ea1..b1774e18558 100644 --- a/mne/preprocessing/artifact_detection.py +++ b/mne/preprocessing/artifact_detection.py @@ -75,7 +75,7 @@ def annotate_muscle_zscore( filter_freq : array-like, shape (2,) The lower and upper frequencies of the band-pass filter. Default is ``(110, 140)``. - %(n_jobs)s + %(n_jobs_cuda)s %(verbose)s Returns diff --git a/mne/preprocessing/tests/test_artifact_detection.py b/mne/preprocessing/tests/test_artifact_detection.py index ccd8893ba11..d3f6e8f75fa 100644 --- a/mne/preprocessing/tests/test_artifact_detection.py +++ b/mne/preprocessing/tests/test_artifact_detection.py @@ -6,10 +6,10 @@ import pytest from numpy.testing import assert_allclose, assert_array_equal -from mne import Annotations, events_from_annotations +from mne import Annotations, create_info, events_from_annotations from mne.chpi import read_head_pos from mne.datasets import testing -from mne.io import read_raw_fif +from mne.io import RawArray, read_raw_fif from mne.preprocessing import ( annotate_break, annotate_movement, @@ -190,6 +190,22 @@ def test_muscle_annotation_without_meeg_data(meas_date): annotate_muscle_zscore(raw, threshold=10) +def test_muscle_annotation_cuda(): + """Test muscle annotation with CUDA filtering and Hilbert transform.""" + rng = np.random.default_rng(0) + raw = RawArray(rng.standard_normal((3, 1000)), create_info(3, 200.0, "eeg")) + annotations, scores = annotate_muscle_zscore( + raw, + threshold=100, + ch_type="eeg", + filter_freq=(20, 40), + n_jobs="cuda", + ) + assert len(annotations) == 0 + assert scores.shape == (raw.n_times,) + assert np.isfinite(scores).all() + + @pytest.mark.parametrize("meas_date", (None, "orig")) @testing.requires_testing_data def test_annotate_breaks(meas_date): diff --git a/mne/tests/test_filter.py b/mne/tests/test_filter.py index 64804f79c4b..2b0852d42ee 100644 --- a/mne/tests/test_filter.py +++ b/mne/tests/test_filter.py @@ -12,11 +12,12 @@ assert_array_equal, assert_array_less, ) -from scipy.signal import butter, freqz, sosfreqz +from scipy.signal import butter, freqz, hilbert, sosfreqz from scipy.signal import resample as sp_resample from mne import Epochs, create_info from mne._fiff.pick import _DATA_CH_TYPES_SPLIT +from mne.cuda import _fft_hilbert from mne.filter import ( _length_factors, _overlap_add_filter, @@ -846,6 +847,39 @@ def test_cuda_fir(): pytest.skip("CUDA not enabled") +@pytest.mark.parametrize("n_times, n_fft", ((99, 99), (100, 128))) +@pytest.mark.parametrize("envelope", (False, True)) +def test_cuda_hilbert(n_times, n_fft, envelope): + """Test CUDA-based Hilbert transforms and CPU fallback.""" + rng = np.random.default_rng(0) + data = rng.standard_normal((3, n_times)) + raw = RawArray(data, create_info(3, 100.0, "eeg")) + expected = raw.copy().apply_hilbert(envelope=envelope, n_jobs=1, n_fft=n_fft) + with catch_logging() as log_file: + got = raw.copy().apply_hilbert( + envelope=envelope, n_jobs="cuda", n_fft=n_fft, verbose="info" + ) + assert_allclose(got.get_data(), expected.get_data(), rtol=1e-7, atol=1e-12) + + from mne.cuda import _cuda_capable + + used_cuda = "Using CUDA for Hilbert transform" in log_file.getvalue() + assert used_cuda == _cuda_capable + + +@pytest.mark.parametrize("n_times, n_fft", ((99, 99), (100, 128))) +@pytest.mark.parametrize("envelope", (False, True)) +def test_fft_hilbert(n_times, n_fft, envelope): + """Test the FFT implementation used by the CUDA path.""" + rng = np.random.default_rng(0) + data = rng.standard_normal((2, 3, n_times)) + got = _fft_hilbert(data, n_fft, envelope, np) + expected = hilbert(data, N=n_fft, axis=-1)[..., :n_times] + if envelope: + expected = np.abs(expected) + assert_allclose(got, expected, rtol=1e-12, atol=1e-12) + + def test_cuda_resampling(): """Test CUDA resampling.""" rng = np.random.default_rng(0)