Repository: mne-tools/mne-python

Issue:
apply_hilbert for source estimated data 
### Describe the new feature or enhancement

Currently apply_hilbert can be applied only to raw/epochs/epoched data through, for example, epochs. apply_hilbert.
The idea is to change this whole function into the time_frequency module. Then the epochs/raw/evoked/source data can be an argument to the function  (EX: mne.time_frequency.EpochsTFR) 

### Describe your proposed implementation

It can be a new class in the tfr.py file. It would be based on the ancient .apply_hilbert function.

### Describe possible alternatives

The alternative that is being used now is script hilbert function for source data or apply_hilbert for other types.
The proposed implementation will make the function independent of the data type/object

### Additional context

I am happy to propose an implementation. I just want to check with maintainers for directions before starting. As this is a huge change in a lot of classes/files. I think it should be discussed before starting the implementation.

Discussion:
BabaSanfour: working to fix the failed tests

BabaSanfour: All tests passed.
I think the functions are decent. 
For SourceEstimate.resample the function overwrite the MixinFilter.resample so I did not update the latter. However I think that SourceEstimate.resample can be integrated into MixinFilter.resample. Not sure if this would have any implication on speed but I think it would improve the code. 
BabaSanfour: I have addressed all of the suggestions. I also added one function to test to `stc.apply_function` in `test_source_estimate.py` . For the `object.apply_hilbert` function, the `test_evoked.py` has the test function for raw, epochs and evoked objects, I wonder if the tests should be added in the same function. Also, to test `object.apply_function` there is pre-existing data in `../io/tests/data` . Should we use for example raw data to generate sources? or is there a way to get access for stc_test_data ?

Patch:
diff --git a/doc/changes/devel/12323.newfeature.rst b/doc/changes/devel/12323.newfeature.rst
new file mode 100644
index 000000000..f10fdf5cf
--- /dev/null
+++ b/doc/changes/devel/12323.newfeature.rst
@@ -0,0 +1,1 @@
+Add :meth:`~mne.SourceEstimate.savgol_filter`, :meth:`~mne.SourceEstimate.filter`, :meth:`~mne.SourceEstimate.apply_hilbert`, and :meth:`~mne.SourceEstimate.apply_function` methods to :class:`mne.SourceEstimate` and related classes, by `Hamza Abdelhedi`_.
\ No newline at end of file
diff --git a/mne/filter.py b/mne/filter.py
index 99fdf7f3b..477434a7c 100644
--- a/mne/filter.py
+++ b/mne/filter.py
@@ -2476,7 +2476,7 @@ class FilterMixin:
 
         Returns
         -------
-        inst : instance of Epochs or Evoked
+        inst : instance of Epochs, Evoked or SourceEstimate
             The object with the filtering applied.
 
         See Also
@@ -2489,6 +2489,8 @@ class FilterMixin:
 
             https://gist.github.com/larsoner/bbac101d50176611136b
 
+        When working on SourceEstimates the sample rate of the original data is inferred from tstep.
+
         .. versionadded:: 0.9.0
 
         References
@@ -2504,13 +2506,19 @@ class FilterMixin:
         >>> evoked.savgol_filter(10.)  # low-pass at around 10 Hz # doctest:+SKIP
         >>> evoked.plot()  # doctest:+SKIP
         """  # noqa: E501
+        from .source_estimate import _BaseSourceEstimate
+
         _check_preload(self, "inst.savgol_filter")
+        if not isinstance(self, _BaseSourceEstimate):
+            s_freq = self.info["sfreq"]
+        else:
+            s_freq = 1 / self.tstep
         h_freq = float(h_freq)
-        if h_freq >= self.info["sfreq"] / 2.0:
+        if h_freq >= s_freq / 2.0:
             raise ValueError("h_freq must be less than half the sample rate")
 
         # savitzky-golay filtering
-        window_length = (int(np.round(self.info["sfreq"] / h_freq)) // 2) * 2 + 1
+        window_length = (int(np.round(s_freq / h_freq)) // 2) * 2 + 1
         logger.info("Using savgol length %d" % window_length)
         self._data[:] = signal.savgol_filter(
             self._data, axis=-1, polyorder=5, window_length=window_length
@@ -2537,7 +2545,7 @@ class FilterMixin:
         *,
         verbose=None,
     ):
-        """Filter a subset of channels.
+        """Filter a subset of channels/vertices.
 
         Parameters
         ----------
@@ -2561,7 +2569,7 @@ class FilterMixin:
 
         Returns
         -------
-        inst : instance of Epochs, Evoked, or Raw
+        inst : instance of Epochs, Evoked, SourceEstimate, or Raw
             The filtered data.
 
         See Also
@@ -2598,6 +2606,9 @@ class FilterMixin:
                   ``len(picks) * n_times`` additional time points need to
                   be temporarily stored in memory.
 
+        When working on SourceEstimates the sample rate of the original
+        data is inferred from tstep.
+
         For more information, see the tutorials
         :ref:`disc-filtering` and :ref:`tut-filter-resample` and
         :func:`mne.filter.create_filter`.
@@ -2606,11 +2617,16 @@ class FilterMixin:
         """
         from .annotations import _annotations_starts_stops
         from .io import BaseRaw
+        from .source_estimate import _BaseSourceEstimate
 
         _check_preload(self, "inst.filter")
+        if not isinstance(self, _BaseSourceEstimate):
+            update_info, picks = _filt_check_picks(self.info, picks, l_freq, h_freq)
+            s_freq = self.info["sfreq"]
+        else:
+            s_freq = 1.0 / self.tstep
         if pad is None and method != "iir":
             pad = "edge"
-        update_info, picks = _filt_check_picks(self.info, picks, l_freq, h_freq)
         if isinstance(self, BaseRaw):
             # Deal with annotations
             onsets, ends = _annotations_starts_stops(
@@ -2629,7 +2645,7 @@ class FilterMixin:
             use_verbose = verbose if si == max_idx else "error"
             filter_data(
                 self._data[:, start:stop],
-                self.info["sfreq"],
+                s_freq,
                 l_freq,
                 h_freq,
                 picks,
@@ -2646,9 +2662,10 @@ class FilterMixin:
                 pad=pad,
                 verbose=use_verbose,
             )
-        # update info if filter is applied to all data channels,
+        # update info if filter is applied to all data channels/vertices,
         # and it's not a band-stop filter
-        _filt_update_info(self.info, update_info, l_freq, h_freq)
+        if not isinstance(self, _BaseSourceEstimate):
+            _filt_update_info(self.info, update_info, l_freq, h_freq)
         return self
 
     @verbose
@@ -2703,7 +2720,7 @@ class FilterMixin:
         from .evoked import Evoked
 
         # Should be guaranteed by our inheritance, and the fact that
-        # mne.io.BaseRaw overrides this method
+        # mne.io.BaseRaw and _BaseSourceEstimate overrides this method
         assert isinstance(self, (BaseEpochs, Evoked))
 
         sfreq = float(sfreq)
@@ -2740,13 +2757,13 @@ class FilterMixin:
     def apply_hilbert(
         self, picks=None, envelope=False, n_jobs=None, n_fft="auto", *, verbose=None
     ):
-        """Compute analytic signal or envelope for a subset of channels.
+        """Compute analytic signal or envelope for a subset of channels/vertices.
 
         Parameters
         ----------
         %(picks_all_data_noref)s
         envelope : bool
-            Compute the envelope signal of each channel. Default False.
+            Compute the envelope signal of each channel/vertex. Default False.
             See Notes.
         %(n_jobs)s
         n_fft : int | None | str
@@ -2758,19 +2775,19 @@ class FilterMixin:
 
         Returns
         -------
-        self : instance of Raw, Epochs, or Evoked
+        self : instance of Raw, Epochs, Evoked or SourceEstimate
             The raw object with transformed data.
 
         Notes
         -----
         **Parameters**
 
-        If ``envelope=False``, the analytic signal for the channels defined in
+        If ``envelope=False``, the analytic signal for the channels/vertices defined in
         ``picks`` is computed and the data of the Raw object is converted to
         a complex representation (the analytic signal is complex valued).
 
         If ``envelope=True``, the absolute value of the analytic signal for the
-        channels defined in ``picks`` is computed, resulting in the envelope
+        channels/vertices defined in ``picks`` is computed, resulting in the envelope
         signal.
 
         .. warning: Do not use ``envelope=True`` if you intend to compute
@@ -2803,7 +2820,15 @@ class FilterMixin:
         by computing the analytic signal in sensor space, applying the MNE
         inverse, and computing the envelope in source space.
         """
+        from .source_estimate import _BaseSourceEstimate
+
+        if not isinstance(self, _BaseSourceEstimate):
+            use_info = self.info
+        else:
+            use_info = len(self._data)
         _check_preload(self, "inst.apply_hilbert")
+        picks = _picks_to_idx(use_info, picks, exclude=(), with_ref_meg=False)
+
         if n_fft is None:
             n_fft = len(self.times)
         elif isinstance(n_fft, str):
@@ -2819,7 +2844,6 @@ class FilterMixin:
                 f"{len(self.times)})"
             )
         dtype = None if envelope else np.complex128
-        picks = _picks_to_idx(self.info, picks, exclude=(), with_ref_meg=False)
         args, kwargs = (), dict(n_fft=n_fft, envelope=envelope)
 
         data_in = self._data
diff --git a/mne/source_estimate.py b/mne/source_estimate.py
index 2bd1ef48d..e99f260c6 100644
--- a/mne/source_estimate.py
+++ b/mne/source_estimate.py
@@ -17,13 +17,14 @@ from scipy.spatial.distance import cdist, pdist
 
 from ._fiff.constants import FIFF
 from ._fiff.meas_info import Info
-from ._fiff.pick import pick_types
+from ._fiff.pick import _picks_to_idx, pick_types
 from ._freesurfer import _get_atlas_values, _get_mri_info_data, read_freesurfer_lut
 from .baseline import rescale
 from .cov import Covariance
 from .evoked import _get_peak
-from .filter import resample
+from .filter import FilterMixin, _check_fun, resample
 from .fixes import _safe_svd
+from .parallel import parallel_func
 from .source_space._source_space import (
     SourceSpaces,
     _check_volume_labels,
@@ -42,6 +43,7 @@ from .utils import (
     _check_option,
     _check_pandas_index_arguments,
     _check_pandas_installed,
+    _check_preload,
     _check_src_normal,
     _check_stc_units,
     _check_subject,
@@ -494,7 +496,7 @@ def _verify_source_estimate_compat(a, b):
         )
 
 
-class _BaseSourceEstimate(TimeMixin):
+class _BaseSourceEstimate(TimeMixin, FilterMixin):
     _data_ndim = 2
 
     @verbose
@@ -642,6 +644,57 @@ class _BaseSourceEstimate(TimeMixin):
             verbose=verbose,
         )
 
+    @verbose
+    def apply_function(
+        self, fun, picks=None, dtype=None, n_jobs=None, verbose=None, **kwargs
+    ):
+        """Apply a function to a subset of vertices.
+
+        %(applyfun_summary_stc)s
+
+        Parameters
+        ----------
+        %(fun_applyfun_stc)s
+        %(picks_all)s
+        %(dtype_applyfun)s
+        %(n_jobs)s Ignored if ``vertice_wise=False`` as the workload
+            is split across vertices.
+        %(verbose)s
+        %(kwargs_fun)s
+
+        Returns
+        -------
+        self : instance of _BaseSourceEstimate
+            The SourceEstimate object with transformed data.
+        """
+        _check_preload(self, "source_estimate.apply_function")
+        picks = _picks_to_idx(len(self._data), picks, exclude=(), with_ref_meg=False)
+
+        if not callable(fun):
+            raise ValueError("fun needs to be a function")
+
+        data_in = self._data
+        if dtype is not None and dtype != self._data.dtype:
+            self._data = self._data.astype(dtype)
+
+        # check the dimension of the source estimate data
+        _check_option("source_estimate.ndim", self._data.ndim, [2, 3])
+
+        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(fun, data_in[idx, :], **kwargs)
+        else:
+            # use parallel function
+            data_picks_new = parallel(
+                p_fun(fun, data_in[p, :], **kwargs) for p in picks
+            )
+            for pp, p in enumerate(picks):
+                self._data[p, :] = data_picks_new[pp]
+
+        return self
+
     @verbose
     def apply_baseline(self, baseline=(None, 0), *, verbose=None):
         """Baseline correct source estimate data.
diff --git a/mne/utils/check.py b/mne/utils/check.py
index 461d7e526..13eca1e0b 100644
--- a/mne/utils/check.py
+++ b/mne/utils/check.py
@@ -294,10 +294,11 @@ def _check_preload(inst, msg):
     """Ensure data are preloaded."""
     from ..epochs import BaseEpochs
     from ..evoked import Evoked
+    from ..source_estimate import _BaseSourceEstimate
     from ..time_frequency import _BaseTFR
     from ..time_frequency.spectrum import BaseSpectrum
 
-    if isinstance(inst, (_BaseTFR, Evoked, BaseSpectrum)):
+    if isinstance(inst, (_BaseTFR, Evoked, BaseSpectrum, _BaseSourceEstimate)):
         pass
     else:
         name = "epochs" if isinstance(inst, BaseEpochs) else "raw"
diff --git a/mne/utils/docs.py b/mne/utils/docs.py
index cb7c027c0..02bb6825b 100644
--- a/mne/utils/docs.py
+++ b/mne/utils/docs.py
@@ -216,11 +216,11 @@ Operates in place.
 # raw/epochs/evoked apply_function method
 # apply_function method summary
 applyfun_summary = """\
-The function ``fun`` is applied to the channels defined in ``picks``.
+The function ``fun`` is applied to the channels or vertices defined in ``picks``.
 The {} object's data is modified in-place. If the function returns a different
 data type (e.g. :py:obj:`numpy.complex128`) it must be specified
 using the ``dtype`` parameter, which causes the data type of **all** the data
-to change (even if the function is only applied to channels in ``picks``).{}
+to change (even if the function is only applied to channels/vertices in ``picks``).{}
 
 .. note:: If ``n_jobs`` > 1, more memory is required as
           ``len(picks) * n_times`` additional time points need to
@@ -236,6 +236,7 @@ applyfun_preload = (
 docdict["applyfun_summary_epochs"] = applyfun_summary.format("epochs", applyfun_preload)
 docdict["applyfun_summary_evoked"] = applyfun_summary.format("evoked", "")
 docdict["applyfun_summary_raw"] = applyfun_summary.format("raw", applyfun_preload)
+docdict["applyfun_summary_stc"] = applyfun_summary.format("source estimate", "")
 
 docdict["area_alpha_plot_psd"] = """\
 area_alpha : float
@@ -1600,6 +1601,9 @@ docdict["fun_applyfun"] = applyfun_fun_base.format(
 docdict["fun_applyfun_evoked"] = applyfun_fun_base.format(
     " because it will apply channel-wise"
 )
+docdict["fun_applyfun_stc"] = applyfun_fun_base.format(
+    " because it will apply vertex-wise"
+)
 
 docdict["fwd"] = """
 fwd : instance of Forward


Test patch:
diff --git a/mne/conftest.py b/mne/conftest.py
index 80380d1a3..17ff64b32 100644
--- a/mne/conftest.py
+++ b/mne/conftest.py
@@ -797,8 +797,6 @@ def mixed_fwd_cov_evoked(_evoked_cov_sphere, _all_src_types_fwd):
 
 
 @pytest.fixture(scope="session")
-@pytest.mark.slowtest
-@pytest.mark.parametrize(params=[testing._pytest_param()])
 def src_volume_labels():
     """Create a 7mm source space with labels."""
     pytest.importorskip("nibabel")
diff --git a/mne/source_space/tests/test_source_space.py b/mne/source_space/tests/test_source_space.py
index 4a1e20eef..a2648459f 100644
--- a/mne/source_space/tests/test_source_space.py
+++ b/mne/source_space/tests/test_source_space.py
@@ -679,6 +679,7 @@ def test_source_space_from_label(tmp_path, pass_ids):
     _compare_source_spaces(src, src_from_file, mode="approx")
 
 
+@pytest.mark.slowtest
 @testing.requires_testing_data
 def test_source_space_exclusive_complete(src_volume_labels):
     """Test that we produce exclusive and complete labels."""
diff --git a/mne/tests/test_evoked.py b/mne/tests/test_evoked.py
index fbf4c0123..17586b1a4 100644
--- a/mne/tests/test_evoked.py
+++ b/mne/tests/test_evoked.py
@@ -23,6 +23,7 @@ from scipy import fftpack
 from mne import (
     Epochs,
     EpochsArray,
+    SourceEstimate,
     combine_evoked,
     create_info,
     equalize_channels,
@@ -917,7 +918,7 @@ def test_evoked_baseline(tmp_path):
 
 
 def test_hilbert():
-    """Test hilbert on raw, epochs, and evoked."""
+    """Test hilbert on raw, epochs, evoked and SourceEstimate data."""
     raw = read_raw_fif(raw_fname).load_data()
     raw.del_proj()
     raw.pick(raw.ch_names[:2])
@@ -927,10 +928,17 @@ def test_hilbert():
         epochs.apply_hilbert()
     epochs.load_data()
     evoked = epochs.average()
+    # Create SourceEstimate stc data
+    verts = [np.arange(10), np.arange(90)]
+    data = np.random.default_rng(0).normal(size=(100, 10))
+    stc = SourceEstimate(data, verts, 0, 1e-1, "foo")
+
     raw_hilb = raw.apply_hilbert()
     epochs_hilb = epochs.apply_hilbert()
     evoked_hilb = evoked.copy().apply_hilbert()
     evoked_hilb_2_data = epochs_hilb.get_data(copy=False).mean(0)
+    stc_hilb = stc.copy().apply_hilbert()
+    stc_hilb_env = stc.copy().apply_hilbert(envelope=True)
     assert_allclose(evoked_hilb.data, evoked_hilb_2_data)
     # This one is only approximate because of edge artifacts
     evoked_hilb_3 = Epochs(raw_hilb, events).average()
@@ -941,6 +949,8 @@ def test_hilbert():
     # envelope=True mode
     evoked_hilb_env = evoked.apply_hilbert(envelope=True)
     assert_allclose(evoked_hilb_env.data, np.abs(evoked_hilb.data))
+    assert len(stc_hilb.data) == len(stc.data)
+    assert_allclose(stc_hilb_env.data, np.abs(stc_hilb.data))
 
 
 def test_apply_function_evk():
diff --git a/mne/tests/test_source_estimate.py b/mne/tests/test_source_estimate.py
index 1ed5c9062..08e08761c 100644
--- a/mne/tests/test_source_estimate.py
+++ b/mne/tests/test_source_estimate.py
@@ -72,7 +72,7 @@ from mne.minimum_norm import (
     read_inverse_operator,
 )
 from mne.morph_map import _make_morph_map_hemi
-from mne.source_estimate import _get_vol_mask, grade_to_tris
+from mne.source_estimate import _get_vol_mask, _make_stc, grade_to_tris
 from mne.source_space._source_space import _get_src_nn
 from mne.transforms import apply_trans, invert_transform, transform_surface_to
 from mne.utils import (
@@ -2055,3 +2055,31 @@ def test_label_extraction_subject(kind):
         stc.subject = None
         with pytest.raises(ValueError, match=r"label\.sub.*not match.* sour"):
             extract_label_time_course(stc, labels_fs, src)
+
+
+def test_apply_function_stc():
+    """Check the apply_function method for source estimate data."""
+    # Create a sample _BaseSourceEstimate object
+    n_vertices = 100
+    n_times = 200
+    vertices = [np.array(np.arange(50)), np.array(np.arange(50, 100))]
+    tmin = 0.0
+    tstep = 0.001
+    data = np.random.default_rng(0).normal(size=(n_vertices, n_times))
+
+    stc = _make_stc(data, vertices, tmin=tmin, tstep=tstep, src_type="surface")
+
+    # A sample function to apply to the data
+    def fun(data_row, **kwargs):
+        return 2 * data_row
+
+    # Test applying the function to all vertices without parallelization
+    stc_copy = stc.copy()
+    stc.apply_function(fun)
+    for idx in range(n_vertices):
+        assert_allclose(stc.data[idx, :], 2 * stc_copy.data[idx, :])
+
+    # Test applying the function with parallelization
+    stc.apply_function(fun, n_jobs=2)
+    for idx in range(n_vertices):
+        assert_allclose(stc.data[idx, :], 4 * stc_copy.data[idx, :])