MRG: Add add_noise to simulation (#5859)

* ENH: Add add_noise to simulation

* FIX: Minor fixes

* FIX: Actually add values

* FIX: Heisenbug

doc/python_reference.rst

@@ -876,6 +876,7 @@ Simulation
 .. autosummary::
    :toctree: generated/
 
+   add_noise
    simulate_evoked
    simulate_raw
    simulate_stc

doc/whats_new.rst

@@ -31,6 +31,8 @@ Changelog
 
 - :func:`mne.io.read_raw_edf` now detects analog stim channels labeled ``'STATUS'`` and sets them as stim channel. :func:`mne.io.read_raw_edf` no longer synthesize TAL annotations into stim channel but stores them in ``raw.annotations`` when reading by `Joan Massich`_
 
+- Add `mne.simulation.add_noise` for ad-hoc noise addition to `io.Raw`, `Epochs`, and `Evoked` instances, by `Eric Larson`_
+
 - Add ``drop_refs=True`` parameter to :func:`set_bipolar_reference` to drop/keep anode and cathode channels after applying the reference by `Cristóbal Moënne-Loccoz`_.
 
 - Add use of :func:`scipy.signal.windows.dpss` for faster multitaper window computations in PSD functions by `Eric Larson`_

mne/simulation/__init__.py

@@ -1,6 +1,6 @@
 """Data simulation code."""
 
-from .evoked import simulate_evoked, simulate_noise_evoked
+from .evoked import simulate_evoked, simulate_noise_evoked, add_noise
 from .raw import simulate_raw
 from .source import select_source_in_label, simulate_stc, simulate_sparse_stc
 from .metrics import source_estimate_quantification

mne/simulation/evoked.py

@@ -8,9 +8,10 @@
 
 import numpy as np
 
-from ..io.pick import pick_channels_cov
+from ..io.pick import pick_channels_cov, pick_info
 from ..forward import apply_forward
-from ..utils import check_random_state, verbose
+from ..utils import (logger, verbose, check_random_state, _check_preload,
+                     deprecated, _validate_type)
 
 
 @verbose
@@ -75,12 +76,14 @@ def simulate_evoked(fwd, stc, info, cov, nave=30, iir_filter=None,
     """
     evoked = apply_forward(fwd, stc, info, use_cps=use_cps)
     if nave < np.inf:
-        noise = simulate_noise_evoked(evoked, cov, iir_filter, random_state)
+        noise = _simulate_noise_evoked(evoked, cov, iir_filter, random_state)
         evoked.data += noise.data / math.sqrt(nave)
         evoked.nave = np.int(nave)
     return evoked
 
 
+@deprecated('simulate_noise_evoked is deprecated in 0.18 and will be removed '
+            'in 0.19, use add_noise instead')
 def simulate_noise_evoked(evoked, cov, iir_filter=None, random_state=None):
     """Create noise as a multivariate Gaussian.
 
@@ -106,10 +109,82 @@ def simulate_noise_evoked(evoked, cov, iir_filter=None, random_state=None):
     -----
     .. versionadded:: 0.10.0
     """
+    return _simulate_noise_evoked(evoked, cov, iir_filter, random_state)
+
+
+def _simulate_noise_evoked(evoked, cov, iir_filter, random_state):
     noise = evoked.copy()
-    noise.data = _generate_noise(evoked.info, cov, iir_filter, random_state,
-                                 evoked.data.shape[1])[0]
-    return noise
+    noise.data[:] = 0
+    return _add_noise(noise, cov, iir_filter, random_state,
+                      allow_subselection=False)
+
+
+@verbose
+def add_noise(inst, cov, iir_filter=None, random_state=None,
+              verbose=None):
+    """Create noise as a multivariate Gaussian.
+
+    The spatial covariance of the noise is given from the cov matrix.
+
+    Parameters
+    ----------
+    inst : instance of Evoked, Epochs, or Raw
+        Instance to which to add noise.
+    cov : instance of Covariance
+        The noise covariance.
+    iir_filter : None | array-like
+        IIR filter coefficients (denominator).
+    random_state : None | int | np.random.RandomState
+        To specify the random generator state.
+    verbose : bool, str, int, or None
+        If not None, override default verbose level (see :func:`mne.verbose`
+        and :ref:`Logging documentation <tut_logging>` for more).
+
+    Returns
+    -------
+    inst : instance of Evoked, Epochs, or Raw
+        The instance, modified to have additional noise.
+
+    Notes
+    -----
+    Only channels in both ``inst.info['ch_names']`` and
+    ``cov['names']`` will have noise added to them.
+
+    This function operates inplace on ``inst``.
+
+    .. versionadded:: 0.18.0
+    """
+    # We always allow subselection here
+    return _add_noise(inst, cov, iir_filter, random_state)
+
+
+def _add_noise(inst, cov, iir_filter, random_state, allow_subselection=True):
+    """Add noise, possibly with channel subselection."""
+    from ..cov import Covariance
+    from ..io import BaseRaw
+    from ..epochs import BaseEpochs
+    from ..evoked import Evoked
+    _validate_type(cov, Covariance, 'cov')
+    _validate_type(inst, (BaseRaw, BaseEpochs, Evoked),
+                   'inst', 'Raw, Epochs, or Evoked')
+    _check_preload(inst, 'Adding noise')
+    data = inst._data
+    assert data.ndim in (2, 3)
+    if data.ndim == 2:
+        data = data[np.newaxis]
+    # Subselect if necessary
+    info = inst.info
+    picks = slice(None)
+    if allow_subselection:
+        use_chs = list(set(info['ch_names']) & set(cov['names']))
+        picks = np.where(np.in1d(info['ch_names'], use_chs))[0]
+        logger.info('Adding noise to %d/%d channels (%d channels in cov)'
+                    % (len(picks), len(info['chs']), len(cov['names'])))
+        info = pick_info(inst.info, picks)
+    for epoch in data:
+        epoch[picks] += _generate_noise(info, cov, iir_filter, random_state,
+                                        epoch.shape[1])[0]
+    return inst
 
 
 def _generate_noise(info, cov, iir_filter, random_state, n_samples, zi=None):

mne/simulation/tests/test_evoked.py

@@ -6,16 +6,17 @@
 
 import numpy as np
 from numpy.testing import (assert_array_almost_equal, assert_array_equal,
-                           assert_equal)
+                           assert_equal, assert_allclose)
 import pytest
 
 from mne import (read_cov, read_forward_solution, convert_forward_solution,
-                 pick_types_forward, read_evokeds)
+                 pick_types_forward, read_evokeds, pick_types, EpochsArray,
+                 compute_covariance, compute_raw_covariance)
 from mne.datasets import testing
-from mne.simulation import simulate_sparse_stc, simulate_evoked
+from mne.simulation import simulate_sparse_stc, simulate_evoked, add_noise
 from mne.io import read_raw_fif
 from mne.cov import regularize
-from mne.utils import run_tests_if_main
+from mne.utils import run_tests_if_main, catch_logging
 
 data_path = testing.data_path(download=False)
 fwd_fname = op.join(data_path, 'MEG', 'sample',
@@ -79,4 +80,52 @@ def test_simulate_evoked():
                   cov, nave=nave, iir_filter=None)
 
 
+@testing.requires_testing_data
+def test_add_noise():
+    """Test noise addition."""
+    rng = np.random.RandomState(0)
+    data_path = testing.data_path()
+    raw = read_raw_fif(data_path + '/MEG/sample/sample_audvis_trunc_raw.fif')
+    raw.del_proj()
+    picks = pick_types(raw.info, eeg=True, exclude=())
+    cov = compute_raw_covariance(raw, picks=picks)
+    with pytest.raises(RuntimeError, match='to be loaded'):
+        add_noise(raw, cov)
+    raw.crop(0, 1).load_data()
+    with pytest.raises(TypeError, match='Raw, Epochs, or Evoked'):
+        add_noise(0., cov)
+    with pytest.raises(TypeError, match='Covariance'):
+        add_noise(raw, 0.)
+    # test a no-op (data preserved)
+    orig_data = raw[:][0]
+    zero_cov = cov.copy()
+    zero_cov['data'].fill(0)
+    add_noise(raw, zero_cov)
+    new_data = raw[:][0]
+    assert_allclose(orig_data, new_data, atol=1e-30)
+    # set to zero to make comparisons easier
+    raw._data[:] = 0.
+    epochs = EpochsArray(np.zeros((1, len(raw.ch_names), 100)),
+                         raw.info.copy())
+    epochs.info['bads'] = []
+    evoked = epochs.average(picks=np.arange(len(raw.ch_names)))
+    for inst in (raw, epochs, evoked):
+        with catch_logging() as log:
+            add_noise(inst, cov, random_state=rng, verbose=True)
+        log = log.getvalue()
+        want = ('to {0}/{1} channels ({0}'
+                .format(len(cov['names']), len(raw.ch_names)))
+        assert want in log
+        if inst is evoked:
+            inst = EpochsArray(inst.data[np.newaxis], inst.info)
+        if inst is raw:
+            cov_new = compute_raw_covariance(inst, picks=picks,
+                                             verbose='error')  # samples
+        else:
+            cov_new = compute_covariance(inst, verbose='error')  # avg ref
+        assert cov['names'] == cov_new['names']
+        r = np.corrcoef(cov['data'].ravel(), cov_new['data'].ravel())[0, 1]
+        assert r > 0.99
+
+
 run_tests_if_main()

mne/tests/test_docstring_parameters.py

@@ -242,7 +242,6 @@ def test_tabs():
 read_tag
 requires_sample_data
 rescale
-simulate_noise_evoked
 source_estimate_quantification
 whiten_evoked
 write_fiducials