Repository: mne-tools/mne-python

Issue:
Accessing data after applying gradient compensation & picking doesn't work
#### Steps to reproduce
```python
import os.path as op
from mne.io import read_raw_fif

base_dir = op.join('/Users/hoechenberger/Development/mne-python/mne', 'io',
                   'tests', 'data')
raw_ctf_fname = op.join(base_dir, 'test_ctf_raw.fif')

raw = read_raw_fif(raw_ctf_fname, preload=False)
raw.apply_gradient_compensation(3)
raw.copy().pick_types(meg=True).get_data()
```

#### Expected results
```
Opening raw data file /Users/hoechenberger/Development/mne-python/mne/io/tests/data/test_ctf_raw.fif...
    Read 5 compensation matrices
    Range : 24000 ... 31200 =     10.000 ...    13.000 secs
Ready.
Current compensation grade : 0
Reading 0 ... 7200  =      0.000 ...     3.000 secs...
Compensator constructed to change 0 -> 3
Applying compensator to loaded data
```

#### Actual results
```python
Opening raw data file /Users/hoechenberger/Development/mne-python/mne/io/tests/data/test_ctf_raw.fif...
    Read 5 compensation matrices
    Range : 24000 ... 31200 =     10.000 ...    13.000 secs
Ready.
Current compensation grade : 0
Compensator constructed to change 0 -> 3
Traceback (most recent call last):
  File "/private/tmp/mwe.py", line 10, in <module>
    raw.copy().pick_types(meg=True).get_data()
  File "<decorator-gen-166>", line 21, in get_data
  File "/Users/hoechenberger/Development/mne-python/mne/io/base.py", line 823, in get_data
    data, times = self[picks, start:stop]
  File "/Users/hoechenberger/Development/mne-python/mne/io/base.py", line 772, in __getitem__
    data = self._read_segment(start=start, stop=stop, sel=sel,
  File "<decorator-gen-163>", line 21, in _read_segment
  File "/Users/hoechenberger/Development/mne-python/mne/io/base.py", line 375, in _read_segment
    mult = self._comp[idx] * cals
ValueError: operands could not be broadcast together with shapes (304,309) (1,304) 
```

#### Additional information
If the data is preloaded:
```python
raw = read_raw_fif(raw_ctf_fname, preload=True)
```
the problem does not occur.

Also if I skip the `pick_types()` it does work.

NumPy is 1.19.1.


Discussion:
I can look into this in the coming days

Resolving patch:
diff --git a/mne/channels/channels.py b/mne/channels/channels.py
--- a/mne/channels/channels.py
+++ b/mne/channels/channels.py
@@ -942,8 +942,10 @@ def _pick_drop_channels(self, idx):
 
         pick_info(self.info, idx, copy=False)
 
-        if getattr(self, '_projector', None) is not None:
-            self._projector = self._projector[idx][:, idx]
+        for key in ('_comp', '_projector'):
+            mat = getattr(self, key, None)
+            if mat is not None:
+                setattr(self, key, mat[idx][:, idx])
 
         # All others (Evoked, Epochs, Raw) have chs axis=-2
         axis = -3 if isinstance(self, (AverageTFR, EpochsTFR)) else -2


Test patch:
diff --git a/mne/io/tests/test_compensator.py b/mne/io/tests/test_compensator.py
--- a/mne/io/tests/test_compensator.py
+++ b/mne/io/tests/test_compensator.py
@@ -5,6 +5,7 @@
 import os.path as op
 import numpy as np
 from numpy.testing import assert_allclose
+import pytest
 
 from mne import Epochs, read_evokeds, pick_types
 from mne.io.compensator import make_compensator, get_current_comp
@@ -15,8 +16,8 @@
 ctf_comp_fname = op.join(base_dir, 'test_ctf_comp_raw.fif')
 
 
-def test_compensation(tmpdir):
-    """Test compensation."""
+def test_compensation_identity():
+    """Test compensation identity."""
     raw = read_raw_fif(ctf_comp_fname)
     assert get_current_comp(raw.info) == 3
     comp1 = make_compensator(raw.info, 3, 1, exclude_comp_chs=False)
@@ -37,10 +38,24 @@ def test_compensation(tmpdir):
             assert_allclose(np.dot(comp1, comp2), desired, atol=1e-12)
             assert_allclose(np.dot(comp2, comp1), desired, atol=1e-12)
 
+
+@pytest.mark.parametrize('preload', (True, False))
+@pytest.mark.parametrize('pick', (False, True))
+def test_compensation_apply(tmpdir, preload, pick):
+    """Test applying compensation."""
     # make sure that changing the comp doesn't modify the original data
-    raw2 = read_raw_fif(ctf_comp_fname)
+    raw = read_raw_fif(ctf_comp_fname, preload=preload)
+    assert raw._comp is None
+    raw2 = raw.copy()
     raw2.apply_gradient_compensation(2)
+    if pick:
+        raw2.pick([0] + list(range(2, len(raw.ch_names))))
+        raw.pick([0] + list(range(2, len(raw.ch_names))))
     assert get_current_comp(raw2.info) == 2
+    if preload:
+        assert raw2._comp is None
+    else:
+        assert raw2._comp.shape == (len(raw2.ch_names),) * 2
     fname = op.join(tmpdir, 'ctf-raw.fif')
     raw2.save(fname)
     raw2 = read_raw_fif(fname)