Repository: mne-tools/mne-python

Issue:
Reading single epochs from a fif file
I noticed that reading single epochs is significantly slower when the raw data is saved as a fif file vs. an edf file. Also, the reading time from fif files linearly increases with how far in the recording the epochs are, which is not the case with edf files.

I'm interested in getting this reading time as low as possible since I have to repeatedly load single epochs. I was wondering whether this difference between fif and edf is expected (maybe part of it is due to the lower precision of edf files).

#### Steps to reproduce
I created a gist that creates .fif, .edf and .hdf5 files out of the same fake raw data, and then times the loading of each individual epoch: https://gist.github.com/hubertjb/1b434688dba296d25eba8fab2866a50f

I included direct reading from an hdf5 file as a baseline. 

#### Expected results
I was expecting fif and edf to yield similar results. Specifically, I was expecting the loading time to be more or less constant across the recording for fif. 

#### Actual results
As shown below, reading from fif is slower than from edf, and reading times increase as the epochs get further from the beginning of the recording.
![image](https://user-images.githubusercontent.com/6500858/76554847-efc2fb80-646c-11ea-908a-ccaa2da9bf22.png)

Mean loading times are
{'edf': 0.0028056514029409373,
 'fif': 0.008060473086787204,
 'hdf5': 0.0009250746053807876}

#### Additional information
Platform:      Linux-4.15.0-70-generic-x86_64-with-debian-buster-sid
Python:        3.7.6 (default, Jan  8 2020, 19:59:22)  [GCC 7.3.0]
Executable:    /storage/store/work/hjacobba/miniconda3/envs/braindecode/bin/python
CPU:           x86_64: 72 cores
Memory:        376.5 GB

mne:           0.20.dev0
numpy:         1.18.1 {blas=openblas, lapack=openblas}
scipy:         1.4.1
matplotlib:    3.1.3 {backend=module://ipykernel.pylab.backend_inline}

sklearn:       0.22.1
numba:         0.48.0
nibabel:       Not found
cupy:          Not found
pandas:        1.0.1
dipy:          Not found
mayavi:        Not found
pyvista:       Not found
vtk:           Not found

@agramfort @larsoner 


Discussion:
> I was expecting fif and edf to yield similar results. Specifically, I was expecting the loading time to be more or less constant across the recording for fif.

In FIF the raw data are stored in "tags" of (usually) 1 sec duration so in principle I don't see why it should be slower. I agree it's worth looking into. @hubertjb do you want to look into it? If not I can give it a shot
... actually from looking at the FIF code I already suspect a culprit. I'll open a WIP PR.

Resolving patch:
diff --git a/mne/io/fiff/raw.py b/mne/io/fiff/raw.py
--- a/mne/io/fiff/raw.py
+++ b/mne/io/fiff/raw.py
@@ -102,10 +102,10 @@ def __init__(self, fname, allow_maxshield=False, preload=False,
 
         # Add annotations for in-data skips
         for extra in self._raw_extras:
-            start = np.array([e['first'] for e in extra if e['ent'] is None])
-            stop = np.array([e['last'] for e in extra if e['ent'] is None])
+            mask = [ent is None for ent in extra['ent']]
+            start = extra['bounds'][:-1][mask]
+            stop = extra['bounds'][1:][mask] - 1
             duration = (stop - start + 1.) / self.info['sfreq']
-
             annot = Annotations(onset=(start / self.info['sfreq']),
                                 duration=duration,
                                 description='BAD_ACQ_SKIP',
@@ -263,6 +263,26 @@ def _read_raw_file(self, fname, allow_maxshield, preload,
 
             next_fname = _get_next_fname(fid, fname_rep, tree)
 
+        # reformat raw_extras to be a dict of list/ndarray rather than
+        # list of dict (faster access)
+        raw_extras = {key: [r[key] for r in raw_extras]
+                      for key in raw_extras[0]}
+        for key in raw_extras:
+            if key != 'ent':  # dict or None
+                raw_extras[key] = np.array(raw_extras[key], int)
+        if not np.array_equal(raw_extras['last'][:-1],
+                              raw_extras['first'][1:] - 1):
+            raise RuntimeError('FIF file appears to be broken')
+        bounds = np.cumsum(np.concatenate(
+            [raw_extras['first'][:1], raw_extras['nsamp']]))
+        raw_extras['bounds'] = bounds
+        assert len(raw_extras['bounds']) == len(raw_extras['ent']) + 1
+        # store the original buffer size
+        buffer_size_sec = np.median(raw_extras['nsamp']) / info['sfreq']
+        del raw_extras['first']
+        del raw_extras['last']
+        del raw_extras['nsamp']
+
         raw.last_samp = first_samp - 1
         raw.orig_format = orig_format
 
@@ -278,10 +298,6 @@ def _read_raw_file(self, fname, allow_maxshield, preload,
                     float(raw.first_samp) / info['sfreq'],
                     float(raw.last_samp) / info['sfreq']))
 
-        # store the original buffer size
-        buffer_size_sec = np.median(
-            [r['nsamp'] for r in raw_extras]) / info['sfreq']
-
         raw.info = info
         raw.verbose = verbose
 
@@ -296,10 +312,10 @@ def _dtype(self):
             return self._dtype_
         dtype = None
         for raw_extra, filename in zip(self._raw_extras, self._filenames):
-            for this in raw_extra:
-                if this['ent'] is not None:
+            for ent in raw_extra['ent']:
+                if ent is not None:
                     with _fiff_get_fid(filename) as fid:
-                        fid.seek(this['ent'].pos, 0)
+                        fid.seek(ent.pos, 0)
                         tag = read_tag_info(fid)
                         if tag is not None:
                             if tag.type in (FIFF.FIFFT_COMPLEX_FLOAT,
@@ -318,52 +334,29 @@ def _dtype(self):
 
     def _read_segment_file(self, data, idx, fi, start, stop, cals, mult):
         """Read a segment of data from a file."""
-        stop -= 1
-        offset = 0
         with _fiff_get_fid(self._filenames[fi]) as fid:
-            for this in self._raw_extras[fi]:
-                #  Do we need this buffer
-                if this['last'] >= start:
-                    #  The picking logic is a bit complicated
-                    if stop > this['last'] and start < this['first']:
-                        #    We need the whole buffer
-                        first_pick = 0
-                        last_pick = this['nsamp']
-                        logger.debug('W')
-
-                    elif start >= this['first']:
-                        first_pick = start - this['first']
-                        if stop <= this['last']:
-                            #   Something from the middle
-                            last_pick = this['nsamp'] + stop - this['last']
-                            logger.debug('M')
-                        else:
-                            #   From the middle to the end
-                            last_pick = this['nsamp']
-                            logger.debug('E')
-                    else:
-                        #    From the beginning to the middle
-                        first_pick = 0
-                        last_pick = stop - this['first'] + 1
-                        logger.debug('B')
-
-                    #   Now we are ready to pick
-                    picksamp = last_pick - first_pick
-                    if picksamp > 0:
-                        # only read data if it exists
-                        if this['ent'] is not None:
-                            one = read_tag(fid, this['ent'].pos,
-                                           shape=(this['nsamp'],
-                                                  self.info['nchan']),
-                                           rlims=(first_pick, last_pick)).data
-                            one.shape = (picksamp, self.info['nchan'])
-                            _mult_cal_one(data[:, offset:(offset + picksamp)],
-                                          one.T, idx, cals, mult)
-                        offset += picksamp
-
-                #   Done?
-                if this['last'] >= stop:
-                    break
+            bounds = self._raw_extras[fi]['bounds']
+            ents = self._raw_extras[fi]['ent']
+            use = (stop > bounds[:-1]) & (start < bounds[1:])
+            offset = 0
+            for ei in np.where(use)[0]:
+                first = bounds[ei]
+                last = bounds[ei + 1]
+                nsamp = last - first
+                ent = ents[ei]
+                first_pick = max(start - first, 0)
+                last_pick = min(nsamp, stop - first)
+                picksamp = last_pick - first_pick
+                # only read data if it exists
+                if ent is not None:
+                    one = read_tag(fid, ent.pos,
+                                   shape=(nsamp, self.info['nchan']),
+                                   rlims=(first_pick, last_pick)).data
+                    one.shape = (picksamp, self.info['nchan'])
+                    _mult_cal_one(data[:, offset:(offset + picksamp)],
+                                  one.T, idx, cals, mult)
+                offset += picksamp
+            assert offset == stop - start
 
     def fix_mag_coil_types(self):
         """Fix Elekta magnetometer coil types.


Test patch:
diff --git a/mne/io/fiff/tests/test_raw_fiff.py b/mne/io/fiff/tests/test_raw_fiff.py
--- a/mne/io/fiff/tests/test_raw_fiff.py
+++ b/mne/io/fiff/tests/test_raw_fiff.py
@@ -431,32 +431,25 @@ def test_split_files(tmpdir):
     raw_crop = raw_1.copy().crop(0, 1)
     raw_crop.save(split_fname, buffer_size_sec=1., overwrite=True)
     raw_read = read_raw_fif(split_fname)
-    assert len(raw_read._raw_extras[0]) == 1
-    assert raw_read._raw_extras[0][0]['nsamp'] == 301
+    assert_array_equal(np.diff(raw_read._raw_extras[0]['bounds']), (301,))
     assert_allclose(raw_crop[:][0], raw_read[:][0])
     # 2 buffers required
     raw_crop.save(split_fname, buffer_size_sec=0.5, overwrite=True)
     raw_read = read_raw_fif(split_fname)
-    assert len(raw_read._raw_extras[0]) == 2
-    assert raw_read._raw_extras[0][0]['nsamp'] == 151
-    assert raw_read._raw_extras[0][1]['nsamp'] == 150
+    assert_array_equal(np.diff(raw_read._raw_extras[0]['bounds']), (151, 150))
     assert_allclose(raw_crop[:][0], raw_read[:][0])
     # 2 buffers required
     raw_crop.save(split_fname,
                   buffer_size_sec=1. - 1.01 / raw_crop.info['sfreq'],
                   overwrite=True)
     raw_read = read_raw_fif(split_fname)
-    assert len(raw_read._raw_extras[0]) == 2
-    assert raw_read._raw_extras[0][0]['nsamp'] == 300
-    assert raw_read._raw_extras[0][1]['nsamp'] == 1
+    assert_array_equal(np.diff(raw_read._raw_extras[0]['bounds']), (300, 1))
     assert_allclose(raw_crop[:][0], raw_read[:][0])
     raw_crop.save(split_fname,
                   buffer_size_sec=1. - 2.01 / raw_crop.info['sfreq'],
                   overwrite=True)
     raw_read = read_raw_fif(split_fname)
-    assert len(raw_read._raw_extras[0]) == 2
-    assert raw_read._raw_extras[0][0]['nsamp'] == 299
-    assert raw_read._raw_extras[0][1]['nsamp'] == 2
+    assert_array_equal(np.diff(raw_read._raw_extras[0]['bounds']), (299, 2))
     assert_allclose(raw_crop[:][0], raw_read[:][0])