Repository: mne-tools/mne-python

Issue:
`add_channels` should update `_orig_units` attribute
[`add_channels`](https://mne.tools/stable/generated/mne.io.Raw.html?highlight=add_channels#mne.io.Raw.add_channels) should add the unit the channel is in to `raw._orig_units`.

This came up in:

- https://github.com/mne-tools/mne-bids/pull/1059

 cc @hoechenberger 



```python
# %%
# Imports
import mne
import numpy as np

# Load data for testing
data_path = mne.datasets.testing.data_path()
vhdr = data_path / "montage" / "bv_dig_test.vhdr"
raw = mne.io.read_raw_brainvision(vhdr, preload=True)

# This has 67 units (as many as there are channels)
assert len(raw._orig_units) == len(raw.ch_names)

# Create a new channel and add it
new_info = mne.create_info(
    ch_names=["NewChannel"], sfreq=raw.info["sfreq"], ch_types=["gsr"]
)
new_data = np.atleast_2d(np.array([2.1e-6] * len(raw.times)))
new_channel = mne.io.RawArray(data=new_data, info=new_info, first_samp=raw.first_samp)
raw.add_channels([new_channel], force_update_info=True)

# The units of the new channel are not added to _orig_units
msg = f"failed: {len(raw._orig_units)} != {len(raw.ch_names)}"
assert len(raw._orig_units) == len(raw.ch_names), msg
```



Discussion:
I can try to look tomorrow
why not but it means we need to be consistent that the `_orig_units` attribute is a part the API of Raw. So I feel is not an isolated change.
@agramfort you mean making it a public attribute? I think I'd keep it private and just include this small fix

Resolving patch:
diff --git a/mne/channels/channels.py b/mne/channels/channels.py
--- a/mne/channels/channels.py
+++ b/mne/channels/channels.py
@@ -415,11 +415,9 @@ def rename_channels(self, mapping, allow_duplicates=False, verbose=None):
         if isinstance(self, BaseRaw):
             # whatever mapping was provided, now we can just use a dict
             mapping = dict(zip(ch_names_orig, self.info['ch_names']))
-            if self._orig_units is not None:
-                for old_name, new_name in mapping.items():
-                    if old_name != new_name:
-                        self._orig_units[new_name] = self._orig_units[old_name]
-                        del self._orig_units[old_name]
+            for old_name, new_name in mapping.items():
+                if old_name in self._orig_units:
+                    self._orig_units[new_name] = self._orig_units.pop(old_name)
             ch_names = self.annotations.ch_names
             for ci, ch in enumerate(ch_names):
                 ch_names[ci] = tuple(mapping.get(name, name) for name in ch)
@@ -830,6 +828,9 @@ def _pick_drop_channels(self, idx, *, verbose=None):
 
         if isinstance(self, BaseRaw):
             self.annotations._prune_ch_names(self.info, on_missing='ignore')
+            self._orig_units = {
+                k: v for k, v in self._orig_units.items()
+                if k in self.ch_names}
 
         self._pick_projs()
         return self
@@ -944,6 +945,8 @@ def add_channels(self, add_list, force_update_info=False):
             self._read_picks = [
                 np.concatenate([r, extra_idx]) for r in self._read_picks]
             assert all(len(r) == self.info['nchan'] for r in self._read_picks)
+            for other in add_list:
+                self._orig_units.update(other._orig_units)
         elif isinstance(self, BaseEpochs):
             self.picks = np.arange(self._data.shape[1])
             if hasattr(self, '_projector'):
diff --git a/mne/io/base.py b/mne/io/base.py
--- a/mne/io/base.py
+++ b/mne/io/base.py
@@ -216,7 +216,7 @@ def __init__(self, info, preload=False,
             # Final check of orig_units, editing a unit if it is not a valid
             # unit
             orig_units = _check_orig_units(orig_units)
-        self._orig_units = orig_units
+        self._orig_units = orig_units or dict()  # always a dict
         self._projectors = list()
         self._projector = None
         self._dtype_ = dtype


Test patch:
diff --git a/mne/io/edf/tests/test_edf.py b/mne/io/edf/tests/test_edf.py
--- a/mne/io/edf/tests/test_edf.py
+++ b/mne/io/edf/tests/test_edf.py
@@ -75,6 +75,20 @@ def test_orig_units():
     orig_units = raw._orig_units
     assert len(orig_units) == len(raw.ch_names)
     assert orig_units['A1'] == 'µV'  # formerly 'uV' edit by _check_orig_units
+    del orig_units
+
+    raw.rename_channels(dict(A1='AA'))
+    assert raw._orig_units['AA'] == 'µV'
+    raw.rename_channels(dict(AA='A1'))
+
+    raw_back = raw.copy().pick(raw.ch_names[:1])  # _pick_drop_channels
+    assert raw_back.ch_names == ['A1']
+    assert set(raw_back._orig_units) == {'A1'}
+    raw_back.add_channels([raw.copy().pick(raw.ch_names[1:])])
+    assert raw_back.ch_names == raw.ch_names
+    assert set(raw_back._orig_units) == set(raw.ch_names)
+    raw_back.reorder_channels(raw.ch_names[::-1])
+    assert set(raw_back._orig_units) == set(raw.ch_names)
 
 
 def test_units_params():
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
@@ -1336,11 +1336,12 @@ def test_add_channels():
     """Test raw splitting / re-appending channel types."""
     rng = np.random.RandomState(0)
     raw = read_raw_fif(test_fif_fname).crop(0, 1).load_data()
+    assert raw._orig_units == {}
     raw_nopre = read_raw_fif(test_fif_fname, preload=False)
     raw_eeg_meg = raw.copy().pick_types(meg=True, eeg=True)
-    raw_eeg = raw.copy().pick_types(meg=False, eeg=True)
-    raw_meg = raw.copy().pick_types(meg=True, eeg=False)
-    raw_stim = raw.copy().pick_types(meg=False, eeg=False, stim=True)
+    raw_eeg = raw.copy().pick_types(eeg=True)
+    raw_meg = raw.copy().pick_types(meg=True)
+    raw_stim = raw.copy().pick_types(stim=True)
     raw_new = raw_meg.copy().add_channels([raw_eeg, raw_stim])
     assert (
         all(ch in raw_new.ch_names