ENH : add units parameter to read_raw_edf in case units is missing fr…

…om the file (#11099)

* ENH : add units parameter to read_raw_edf in case units is missing from the file

* cleanup

* update what's new + flake8

* adding units to bdf too

* fix docstring

doc/changes/latest.inc

@@ -35,6 +35,7 @@ Enhancements
 - Add HTML representation for `~mne.Evoked` in Jupyter Notebooks (:gh:`11075` by `Valerii Chirkov`_ and `Andrew Quinn`_)
 - Add support for ``temperature`` and ``gsr`` (galvanic skin response, i.e., electrodermal activity) channel types (:gh:`11090` by `Eric Larson`_)
 - Allow :func:`mne.beamformer.make_dics` to take ``pick_ori='vector'`` to compute vector source estimates (:gh:`19080` by `Alex Rockhill`_)
+- Add ``units`` parameter to :func:`mne.io.read_raw_edf` in case units are missing from the file (:gh:`11099` by `Alex Gramfort`_)
 - Add ``on_missing`` functionality to all of our classes that have a ``drop_channels`` method, to control what happens when channel names are not in the object (:gh:`11077` by `Andrew Quinn`_)
 
 Bugs

mne/io/edf/edf.py

@@ -19,7 +19,7 @@
 
 from ...utils import verbose, logger, warn
 from ..utils import _blk_read_lims, _mult_cal_one
-from ..base import BaseRaw
+from ..base import BaseRaw, _get_scaling
 from ..meas_info import _empty_info, _unique_channel_names
 from ..constants import FIFF
 from ...filter import resample
@@ -83,6 +83,7 @@ class RawEDF(BaseRaw):
 
         .. versionadded:: 1.1
     %(preload)s
+    %(units_edf_bdf_io)s
     %(verbose)s
 
     See Also
@@ -132,14 +133,30 @@ class RawEDF(BaseRaw):
     @verbose
     def __init__(self, input_fname, eog=None, misc=None, stim_channel='auto',
                  exclude=(), infer_types=False, preload=False, include=None,
-                 verbose=None):
+                 units=None, *, verbose=None):
         logger.info('Extracting EDF parameters from {}...'.format(input_fname))
         input_fname = os.path.abspath(input_fname)
         info, edf_info, orig_units = _get_info(input_fname, stim_channel, eog,
                                                misc, exclude, infer_types,
                                                preload, include)
         logger.info('Creating raw.info structure...')
 
+        if units is not None and isinstance(units, str):
+            units = {ch_name: units for ch_name in info['ch_names']}
+        elif units is None:
+            units = dict()
+
+        for k, (this_ch, this_unit) in enumerate(orig_units.items()):
+            if this_unit != "" and this_unit in units:
+                raise ValueError(f'Unit for channel {this_ch} is present in '
+                                 'the file. Cannot overwrite it with the '
+                                 'units argument.')
+            if this_unit == "" and this_ch in units:
+                orig_units[this_ch] = units[this_ch]
+                ch_type = edf_info["ch_types"][k]
+                scaling = _get_scaling(ch_type.lower(), orig_units[this_ch])
+                edf_info["units"][k] /= scaling
+
         # Raw attributes
         last_samps = [edf_info['nsamples'] - 1]
         super().__init__(info, preload, filenames=[input_fname],
@@ -1282,7 +1299,7 @@ def _find_tal_idx(ch_names):
 @fill_doc
 def read_raw_edf(input_fname, eog=None, misc=None, stim_channel='auto',
                  exclude=(), infer_types=False, include=None, preload=False,
-                 verbose=None):
+                 units=None, *, verbose=None):
     """Reader function for EDF or EDF+ files.
 
     Parameters
@@ -1322,6 +1339,7 @@ def read_raw_edf(input_fname, eog=None, misc=None, stim_channel='auto',
 
         .. versionadded:: 1.1
     %(preload)s
+    %(units_edf_bdf_io)s
     %(verbose)s
 
     Returns
@@ -1384,13 +1402,13 @@ def read_raw_edf(input_fname, eog=None, misc=None, stim_channel='auto',
     return RawEDF(input_fname=input_fname, eog=eog, misc=misc,
                   stim_channel=stim_channel, exclude=exclude,
                   infer_types=infer_types, preload=preload, include=include,
-                  verbose=verbose)
+                  units=units, verbose=verbose)
 
 
 @fill_doc
 def read_raw_bdf(input_fname, eog=None, misc=None, stim_channel='auto',
                  exclude=(), infer_types=False, include=None, preload=False,
-                 verbose=None):
+                 units=None, *, verbose=None):
     """Reader function for BDF files.
 
     Parameters
@@ -1430,6 +1448,7 @@ def read_raw_bdf(input_fname, eog=None, misc=None, stim_channel='auto',
 
         .. versionadded:: 1.1
     %(preload)s
+    %(units_edf_bdf_io)s
     %(verbose)s
 
     Returns
@@ -1485,7 +1504,7 @@ def read_raw_bdf(input_fname, eog=None, misc=None, stim_channel='auto',
     return RawEDF(input_fname=input_fname, eog=eog, misc=misc,
                   stim_channel=stim_channel, exclude=exclude,
                   infer_types=infer_types, preload=preload, include=include,
-                  verbose=verbose)
+                  units=units, verbose=verbose)
 
 
 @fill_doc

mne/utils/docs.py

@@ -3508,6 +3508,14 @@ def _reflow_param_docstring(docstring, has_first_line=True, width=75):
     channel-type-specific default unit.
 """
 
+docdict['units_edf_bdf_io'] = """
+units : dict | str
+    The units of the channels as stored in the file. This argument
+    is useful only if the units are missing from the original file.
+    If a dict, it must map a channel name to its unit, and if str
+    it is assumed that all channels have the same units.
+"""
+
 docdict['units_topomap'] = """
 units : dict | str | None
     The unit of the channel type used for colorbar label. If