ENH Add CUDA support for Hilbert transforms

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`_.

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
 

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]

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

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):

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)