Repository: mne-tools/mne-python

Issue:
BUG?: reject and flat parameters not saved with epochs.save
It would be useful to have for historical purposes, especially when using `autoreject` to set per-subject reject params. This seems like an oversight? If so we'll need some new FIFF constants to store the dicts.


Discussion:
+1

👍 

Resolving patch:
diff --git a/mne/epochs.py b/mne/epochs.py
--- a/mne/epochs.py
+++ b/mne/epochs.py
@@ -53,6 +53,15 @@
                     _check_event_id, _gen_events)
 
 
+def _pack_reject_params(epochs):
+    reject_params = dict()
+    for key in ('reject', 'flat', 'reject_tmin', 'reject_tmax'):
+        val = getattr(epochs, key, None)
+        if val is not None:
+            reject_params[key] = val
+    return reject_params
+
+
 def _save_split(epochs, fname, part_idx, n_parts, fmt):
     """Split epochs."""
     # insert index in filename
@@ -147,6 +156,11 @@ def _save_split(epochs, fname, part_idx, n_parts, fmt):
     write_string(fid, FIFF.FIFF_MNE_EPOCHS_DROP_LOG,
                  json.dumps(epochs.drop_log))
 
+    reject_params = _pack_reject_params(epochs)
+    if reject_params:
+        write_string(fid, FIFF.FIFF_MNE_EPOCHS_REJECT_FLAT,
+                     json.dumps(reject_params))
+
     write_int(fid, FIFF.FIFF_MNE_EPOCHS_SELECTION,
               epochs.selection)
 
@@ -2207,6 +2221,7 @@ def _read_one_epoch_file(f, tree, preload):
         baseline = None
         selection = None
         drop_log = None
+        reject_params = {}
         for k in range(my_epochs['nent']):
             kind = my_epochs['directory'][k].kind
             pos = my_epochs['directory'][k].pos
@@ -2236,6 +2251,9 @@ def _read_one_epoch_file(f, tree, preload):
             elif kind == FIFF.FIFF_MNE_EPOCHS_DROP_LOG:
                 tag = read_tag(fid, pos)
                 drop_log = json.loads(tag.data)
+            elif kind == FIFF.FIFF_MNE_EPOCHS_REJECT_FLAT:
+                tag = read_tag(fid, pos)
+                reject_params = json.loads(tag.data)
 
         if bmin is not None or bmax is not None:
             baseline = (bmin, bmax)
@@ -2295,7 +2313,7 @@ def _read_one_epoch_file(f, tree, preload):
             drop_log = [[] for _ in range(len(events))]
 
     return (info, data, data_tag, events, event_id, metadata, tmin, tmax,
-            baseline, selection, drop_log, epoch_shape, cals)
+            baseline, selection, drop_log, epoch_shape, cals, reject_params)
 
 
 @verbose
@@ -2394,7 +2412,8 @@ def __init__(self, fname, proj=True, preload=True,
             fid, tree, _ = fiff_open(fname)
             next_fname = _get_next_fname(fid, fname, tree)
             (info, data, data_tag, events, event_id, metadata, tmin, tmax,
-             baseline, selection, drop_log, epoch_shape, cals) = \
+             baseline, selection, drop_log, epoch_shape, cals,
+             reject_params) = \
                 _read_one_epoch_file(fid, tree, preload)
 
             # here we ignore missing events, since users should already be
@@ -2437,7 +2456,7 @@ def __init__(self, fname, proj=True, preload=True,
             info, data, events, event_id, tmin, tmax, baseline, raw=raw,
             proj=proj, preload_at_end=False, on_missing='ignore',
             selection=selection, drop_log=drop_log, filename=fname,
-            metadata=metadata, verbose=verbose)
+            metadata=metadata, verbose=verbose, **reject_params)
         # use the private property instead of drop_bad so that epochs
         # are not all read from disk for preload=False
         self._bad_dropped = True
diff --git a/mne/io/constants.py b/mne/io/constants.py
--- a/mne/io/constants.py
+++ b/mne/io/constants.py
@@ -884,6 +884,7 @@ def __setattr__(self, attr, val):  # noqa: D105
 # MNE epochs bookkeeping
 FIFF.FIFF_MNE_EPOCHS_SELECTION     = 3800  # the epochs selection
 FIFF.FIFF_MNE_EPOCHS_DROP_LOG      = 3801  # the drop log
+FIFF.FIFF_MNE_EPOCHS_REJECT_FLAT   = 3802  # rejection and flat params
 
 # MNE annotations
 FIFF.FIFFB_MNE_ANNOTATIONS         = 3810  # annotations block


Test patch:
diff --git a/mne/io/tests/test_constants.py b/mne/io/tests/test_constants.py
--- a/mne/io/tests/test_constants.py
+++ b/mne/io/tests/test_constants.py
@@ -14,6 +14,8 @@
 from mne.utils import _fetch_file, requires_good_network
 
 
+commit = '4da59a8ed8b84e9941e739c9994a7348c13d3a48'  # mne-tools/fiff-constants
+
 # These are oddities that we won't address:
 iod_dups = (355, 359)  # these are in both MEGIN and MNE files
 tag_dups = (3501, 3507)  # in both MEGIN and MNE files
@@ -23,6 +25,8 @@
                      'has_key', 'iteritems', 'iterkeys', 'itervalues',  # Py2
                      'viewitems', 'viewkeys', 'viewvalues',  # Py2
                      )
+_tag_ignore_names = (  # for fiff-constants pending updates
+)
 _ignore_incomplete_enums = (  # XXX eventually we could complete these
     'bem_surf_id', 'cardinal_point_cardiac', 'cond_model', 'coord',
     'dacq_system', 'diffusion_param', 'gantry_type', 'map_surf',
@@ -67,8 +71,8 @@ def test_constants(tmpdir):
     """Test compensation."""
     tmpdir = str(tmpdir)  # old pytest...
     dest = op.join(tmpdir, 'fiff.zip')
-    _fetch_file('https://codeload.github.com/mne-tools/fiff-constants/zip/'
-                '066541057377b694a7d65fc18838fd80d8ffc284', dest)
+    _fetch_file('https://codeload.github.com/mne-tools/fiff-constants/zip/' +
+                commit, dest)
     names = list()
     with zipfile.ZipFile(dest, 'r') as ff:
         for name in ff.namelist():
@@ -259,7 +263,7 @@ def test_constants(tmpdir):
             check = 'tags'
         else:
             unknowns.append((name, val))
-        if check is not None:
+        if check is not None and name not in _tag_ignore_names:
             assert val in fif[check], '%s: %s, %s' % (check, val, name)
             if val in con[check]:
                 msg = "%s='%s'  ?" % (name, con[check][val])
diff --git a/mne/tests/test_epochs.py b/mne/tests/test_epochs.py
--- a/mne/tests/test_epochs.py
+++ b/mne/tests/test_epochs.py
@@ -26,7 +26,7 @@
 from mne.epochs import (
     bootstrap, equalize_epoch_counts, combine_event_ids, add_channels_epochs,
     EpochsArray, concatenate_epochs, BaseEpochs, average_movements)
-from mne.utils import (requires_pandas, run_tests_if_main,
+from mne.utils import (requires_pandas, run_tests_if_main, object_diff,
                        requires_version, catch_logging, _FakeNoPandas)
 from mne.chpi import read_head_pos, head_pos_to_trans_rot_t
 
@@ -664,21 +664,12 @@ def test_read_epochs_bad_events():
     assert evoked
 
 
-@pytest.mark.timeout(60)  # can take > 30 s on Travis
-@pytest.mark.slowtest
-def test_read_write_epochs(tmpdir):
+def test_io_epochs_basic(tmpdir):
     """Test epochs from raw files with IO as fif file."""
     raw, events, picks = _get_data(preload=True)
-    tempdir = str(tmpdir)
-    temp_fname = op.join(tempdir, 'test-epo.fif')
-    temp_fname_no_bl = op.join(tempdir, 'test_no_bl-epo.fif')
     baseline = (None, 0)
     epochs = Epochs(raw, events, event_id, tmin, tmax, picks=picks,
                     baseline=baseline, preload=True)
-    epochs_orig = epochs.copy()
-    epochs_no_bl = Epochs(raw, events, event_id, tmin, tmax, picks=picks,
-                          baseline=None, preload=True)
-    assert (epochs_no_bl.baseline is None)
     evoked = epochs.average()
     data = epochs.get_data()
 
@@ -723,6 +714,10 @@ def test_read_write_epochs(tmpdir):
     assert (n_dec_min <= n_dec <= n_dec_min + 1)
     assert (evoked_dec.info['sfreq'] == evoked.info['sfreq'] / 4)
 
+
+@pytest.mark.parametrize('proj', (True, 'delayed', False))
+def test_epochs_io_proj(tmpdir, proj):
+    """Test epochs I/O with projection."""
     # Test event access on non-preloaded data (#2345)
 
     # due to reapplication of the proj matrix, this is our quality limit
@@ -733,22 +728,29 @@ def test_read_write_epochs(tmpdir):
     events[::2, 1] = 1
     events[1::2, 2] = 2
     event_ids = dict(a=1, b=2)
-    for proj in (True, 'delayed', False):
-        epochs = Epochs(raw, events, event_ids, tmin, tmax, picks=picks,
-                        proj=proj, reject=reject)
-        assert_equal(epochs.proj, proj if proj != 'delayed' else False)
-        data1 = epochs.get_data()
-        epochs2 = epochs.copy().apply_proj()
-        assert_equal(epochs2.proj, True)
-        data2 = epochs2.get_data()
-        assert_allclose(data1, data2, **tols)
-        epochs.save(temp_fname)
-        epochs_read = read_epochs(temp_fname, preload=False)
-        assert_allclose(epochs.get_data(), epochs_read.get_data(), **tols)
-        assert_allclose(epochs['a'].get_data(),
-                        epochs_read['a'].get_data(), **tols)
-        assert_allclose(epochs['b'].get_data(),
-                        epochs_read['b'].get_data(), **tols)
+    temp_fname = op.join(str(tmpdir), 'test-epo.fif')
+
+    epochs = Epochs(raw, events, event_ids, tmin, tmax, picks=picks,
+                    proj=proj, reject=reject, flat=dict(),
+                    reject_tmin=tmin + 0.01, reject_tmax=tmax - 0.01)
+    assert_equal(epochs.proj, proj if proj != 'delayed' else False)
+    data1 = epochs.get_data()
+    epochs2 = epochs.copy().apply_proj()
+    assert_equal(epochs2.proj, True)
+    data2 = epochs2.get_data()
+    assert_allclose(data1, data2, **tols)
+    epochs.save(temp_fname)
+    epochs_read = read_epochs(temp_fname, preload=False)
+    assert_allclose(epochs.get_data(), epochs_read.get_data(), **tols)
+    assert_allclose(epochs['a'].get_data(),
+                    epochs_read['a'].get_data(), **tols)
+    assert_allclose(epochs['b'].get_data(),
+                    epochs_read['b'].get_data(), **tols)
+    assert epochs.reject is not None
+    assert object_diff(epochs.reject, reject) == ''
+    assert epochs.flat is None  # empty dict is functionally the same
+    assert epochs.reject_tmin == tmin + 0.01
+    assert epochs.reject_tmax == tmax - 0.01
 
     # ensure we don't leak file descriptors
     epochs_read = read_epochs(temp_fname, preload=False)
@@ -757,109 +759,127 @@ def test_read_write_epochs(tmpdir):
     epochs_copy.get_data()
     del epochs_copy
 
-    # test IO
-    for preload in (False, True):
-        epochs = epochs_orig.copy()
-        epochs.save(temp_fname)
-        epochs_no_bl.save(temp_fname_no_bl)
-        epochs_read = read_epochs(temp_fname, preload=preload)
-        epochs_no_bl.save(temp_fname_no_bl)
-        epochs_read = read_epochs(temp_fname)
-        epochs_no_bl_read = read_epochs(temp_fname_no_bl)
-        pytest.raises(ValueError, epochs.apply_baseline, baseline=[1, 2, 3])
-        epochs_with_bl = epochs_no_bl_read.copy().apply_baseline(baseline)
-        assert (isinstance(epochs_with_bl, BaseEpochs))
-        assert (epochs_with_bl.baseline == baseline)
-        assert (epochs_no_bl_read.baseline != baseline)
-        assert (str(epochs_read).startswith('<Epochs'))
-
-        epochs_no_bl_read.apply_baseline(baseline)
-        assert_array_equal(epochs_no_bl_read.times, epochs.times)
-        assert_array_almost_equal(epochs_read.get_data(), epochs.get_data())
-        assert_array_almost_equal(epochs.get_data(),
-                                  epochs_no_bl_read.get_data())
-        assert_array_equal(epochs_read.times, epochs.times)
-        assert_array_almost_equal(epochs_read.average().data, evoked.data)
-        assert_equal(epochs_read.proj, epochs.proj)
-        bmin, bmax = epochs.baseline
-        if bmin is None:
-            bmin = epochs.times[0]
-        if bmax is None:
-            bmax = epochs.times[-1]
-        baseline = (bmin, bmax)
-        assert_array_almost_equal(epochs_read.baseline, baseline)
-        assert_array_almost_equal(epochs_read.tmin, epochs.tmin, 2)
-        assert_array_almost_equal(epochs_read.tmax, epochs.tmax, 2)
-        assert_equal(epochs_read.event_id, epochs.event_id)
-
-        epochs.event_id.pop('1')
-        epochs.event_id.update({'a:a': 1})  # test allow for ':' in key
-        epochs.save(op.join(tempdir, 'foo-epo.fif'))
-        epochs_read2 = read_epochs(op.join(tempdir, 'foo-epo.fif'),
-                                   preload=preload)
-        assert_equal(epochs_read2.event_id, epochs.event_id)
-        assert_equal(epochs_read2['a:a'].average().comment, 'a:a')
-
-        # add reject here so some of the epochs get dropped
-        epochs = Epochs(raw, events, event_id, tmin, tmax, picks=picks,
-                        reject=reject)
-        epochs.save(temp_fname)
-        # ensure bad events are not saved
-        epochs_read3 = read_epochs(temp_fname, preload=preload)
-        assert_array_equal(epochs_read3.events, epochs.events)
-        data = epochs.get_data()
-        assert (epochs_read3.events.shape[0] == data.shape[0])
-
-        # test copying loaded one (raw property)
-        epochs_read4 = epochs_read3.copy()
-        assert_array_almost_equal(epochs_read4.get_data(), data)
-        # test equalizing loaded one (drop_log property)
-        epochs_read4.equalize_event_counts(epochs.event_id)
-
-        epochs.drop([1, 2], reason='can we recover orig ID?')
-        epochs.save(temp_fname)
-        epochs_read5 = read_epochs(temp_fname, preload=preload)
-        assert_array_equal(epochs_read5.selection, epochs.selection)
-        assert_equal(len(epochs_read5.selection), len(epochs_read5.events))
-        assert_array_equal(epochs_read5.drop_log, epochs.drop_log)
 
-        if preload:
-            # Test that one can drop channels on read file
-            epochs_read5.drop_channels(epochs_read5.ch_names[:1])
-
-        # test warnings on bad filenames
-        epochs_badname = op.join(tempdir, 'test-bad-name.fif.gz')
-        with pytest.warns(RuntimeWarning, match='-epo.fif'):
-            epochs.save(epochs_badname)
-        with pytest.warns(RuntimeWarning, match='-epo.fif'):
-            read_epochs(epochs_badname, preload=preload)
-
-        # test loading epochs with missing events
-        epochs = Epochs(raw, events, dict(foo=1, bar=999), tmin, tmax,
-                        picks=picks, on_missing='ignore')
-        epochs.save(temp_fname)
-        epochs_read = read_epochs(temp_fname, preload=preload)
-        assert_allclose(epochs.get_data(), epochs_read.get_data(), **tols)
-        assert_array_equal(epochs.events, epochs_read.events)
-        assert_equal(set(epochs.event_id.keys()),
-                     {str(x) for x in epochs_read.event_id.keys()})
-
-        # test saving split epoch files
-        epochs.save(temp_fname, split_size='7MB')
-        epochs_read = read_epochs(temp_fname, preload=preload)
-        assert_allclose(epochs.get_data(), epochs_read.get_data(), **tols)
-        assert_array_equal(epochs.events, epochs_read.events)
-        assert_array_equal(epochs.selection, epochs_read.selection)
-        assert_equal(epochs.drop_log, epochs_read.drop_log)
-
-        # Test that having a single time point works
-        epochs.load_data().crop(0, 0)
-        assert_equal(len(epochs.times), 1)
-        assert_equal(epochs.get_data().shape[-1], 1)
-        epochs.save(temp_fname)
-        epochs_read = read_epochs(temp_fname, preload=preload)
-        assert_equal(len(epochs_read.times), 1)
-        assert_equal(epochs.get_data().shape[-1], 1)
+@pytest.mark.parametrize('preload', (False, True))
+def test_epochs_io_preload(tmpdir, preload):
+    """Test epochs I/O with preloading."""
+    # due to reapplication of the proj matrix, this is our quality limit
+    # for some tests
+    tols = dict(atol=1e-3, rtol=1e-20)
+
+    raw, events, picks = _get_data(preload=True)
+    tempdir = str(tmpdir)
+    temp_fname = op.join(tempdir, 'test-epo.fif')
+    temp_fname_no_bl = op.join(tempdir, 'test_no_bl-epo.fif')
+    baseline = (None, 0)
+    epochs = Epochs(raw, events, event_id, tmin, tmax, picks=picks,
+                    baseline=baseline, preload=True)
+    evoked = epochs.average()
+    epochs.save(temp_fname)
+
+    epochs_no_bl = Epochs(raw, events, event_id, tmin, tmax, picks=picks,
+                          baseline=None, preload=True)
+    assert (epochs_no_bl.baseline is None)
+    epochs_no_bl.save(temp_fname_no_bl)
+
+    epochs_read = read_epochs(temp_fname, preload=preload)
+    epochs_no_bl.save(temp_fname_no_bl)
+    epochs_read = read_epochs(temp_fname)
+    epochs_no_bl_read = read_epochs(temp_fname_no_bl)
+    pytest.raises(ValueError, epochs.apply_baseline, baseline=[1, 2, 3])
+    epochs_with_bl = epochs_no_bl_read.copy().apply_baseline(baseline)
+    assert (isinstance(epochs_with_bl, BaseEpochs))
+    assert (epochs_with_bl.baseline == baseline)
+    assert (epochs_no_bl_read.baseline != baseline)
+    assert (str(epochs_read).startswith('<Epochs'))
+
+    epochs_no_bl_read.apply_baseline(baseline)
+    assert_array_equal(epochs_no_bl_read.times, epochs.times)
+    assert_array_almost_equal(epochs_read.get_data(), epochs.get_data())
+    assert_array_almost_equal(epochs.get_data(),
+                              epochs_no_bl_read.get_data())
+    assert_array_equal(epochs_read.times, epochs.times)
+    assert_array_almost_equal(epochs_read.average().data, evoked.data)
+    assert_equal(epochs_read.proj, epochs.proj)
+    bmin, bmax = epochs.baseline
+    if bmin is None:
+        bmin = epochs.times[0]
+    if bmax is None:
+        bmax = epochs.times[-1]
+    baseline = (bmin, bmax)
+    assert_array_almost_equal(epochs_read.baseline, baseline)
+    assert_array_almost_equal(epochs_read.tmin, epochs.tmin, 2)
+    assert_array_almost_equal(epochs_read.tmax, epochs.tmax, 2)
+    assert_equal(epochs_read.event_id, epochs.event_id)
+
+    epochs.event_id.pop('1')
+    epochs.event_id.update({'a:a': 1})  # test allow for ':' in key
+    epochs.save(op.join(tempdir, 'foo-epo.fif'))
+    epochs_read2 = read_epochs(op.join(tempdir, 'foo-epo.fif'),
+                               preload=preload)
+    assert_equal(epochs_read2.event_id, epochs.event_id)
+    assert_equal(epochs_read2['a:a'].average().comment, 'a:a')
+
+    # add reject here so some of the epochs get dropped
+    epochs = Epochs(raw, events, event_id, tmin, tmax, picks=picks,
+                    reject=reject)
+    epochs.save(temp_fname)
+    # ensure bad events are not saved
+    epochs_read3 = read_epochs(temp_fname, preload=preload)
+    assert_array_equal(epochs_read3.events, epochs.events)
+    data = epochs.get_data()
+    assert (epochs_read3.events.shape[0] == data.shape[0])
+
+    # test copying loaded one (raw property)
+    epochs_read4 = epochs_read3.copy()
+    assert_array_almost_equal(epochs_read4.get_data(), data)
+    # test equalizing loaded one (drop_log property)
+    epochs_read4.equalize_event_counts(epochs.event_id)
+
+    epochs.drop([1, 2], reason='can we recover orig ID?')
+    epochs.save(temp_fname)
+    epochs_read5 = read_epochs(temp_fname, preload=preload)
+    assert_array_equal(epochs_read5.selection, epochs.selection)
+    assert_equal(len(epochs_read5.selection), len(epochs_read5.events))
+    assert_array_equal(epochs_read5.drop_log, epochs.drop_log)
+
+    if preload:
+        # Test that one can drop channels on read file
+        epochs_read5.drop_channels(epochs_read5.ch_names[:1])
+
+    # test warnings on bad filenames
+    epochs_badname = op.join(tempdir, 'test-bad-name.fif.gz')
+    with pytest.warns(RuntimeWarning, match='-epo.fif'):
+        epochs.save(epochs_badname)
+    with pytest.warns(RuntimeWarning, match='-epo.fif'):
+        read_epochs(epochs_badname, preload=preload)
+
+    # test loading epochs with missing events
+    epochs = Epochs(raw, events, dict(foo=1, bar=999), tmin, tmax,
+                    picks=picks, on_missing='ignore')
+    epochs.save(temp_fname)
+    epochs_read = read_epochs(temp_fname, preload=preload)
+    assert_allclose(epochs.get_data(), epochs_read.get_data(), **tols)
+    assert_array_equal(epochs.events, epochs_read.events)
+    assert_equal(set(epochs.event_id.keys()),
+                 {str(x) for x in epochs_read.event_id.keys()})
+
+    # test saving split epoch files
+    epochs.save(temp_fname, split_size='7MB')
+    epochs_read = read_epochs(temp_fname, preload=preload)
+    assert_allclose(epochs.get_data(), epochs_read.get_data(), **tols)
+    assert_array_equal(epochs.events, epochs_read.events)
+    assert_array_equal(epochs.selection, epochs_read.selection)
+    assert_equal(epochs.drop_log, epochs_read.drop_log)
+
+    # Test that having a single time point works
+    epochs.load_data().crop(0, 0)
+    assert_equal(len(epochs.times), 1)
+    assert_equal(epochs.get_data().shape[-1], 1)
+    epochs.save(temp_fname)
+    epochs_read = read_epochs(temp_fname, preload=preload)
+    assert_equal(len(epochs_read.times), 1)
+    assert_equal(epochs.get_data().shape[-1], 1)
 
 
 def test_split_saving(tmpdir):