braindecode 0.8.1__py3-none-any.whl → 1.0.0__py3-none-any.whl

braindecode/datasets/tuh.py

@@ -7,20 +7,22 @@ TUH Abnormal EEG Corpus.
 #
 # License: BSD (3-clause)
 
-import re
-import os
+from __future__ import annotations
+
 import glob
+import os
+import re
 import warnings
-from unittest import mock
 from datetime import datetime, timezone
 from typing import Iterable
+from unittest import mock
 
-import pandas as pd
-import numpy as np
 import mne
+import numpy as np
+import pandas as pd
 from joblib import Parallel, delayed
 
-from .base import BaseDataset, BaseConcatDataset
+from .base import BaseConcatDataset, BaseDataset
 
 
 class TUH(BaseConcatDataset):
@@ -44,14 +46,32 @@ class TUH(BaseConcatDataset):
     add_physician_reports: bool
         If True, the physician reports will be read from disk and added to the
         description.
+    rename_channels: bool
+        If True, rename the EEG channels to the standard 10-05 system.
+    set_montage: bool
+        If True, set the montage to the standard 10-05 system.
     n_jobs: int
         Number of jobs to be used to read files in parallel.
     """
-    def __init__(self, path, recording_ids=None, target_name=None,
-                 preload=False, add_physician_reports=False, n_jobs=1):
+
+    def __init__(
+        self,
+        path: str,
+        recording_ids: list[int] | None = None,
+        target_name: str | tuple[str, ...] | None = None,
+        preload: bool = False,
+        add_physician_reports: bool = False,
+        rename_channels: bool = False,
+        set_montage: bool = False,
+        n_jobs: int = 1,
+    ):
+        if set_montage:
+            assert rename_channels, (
+                "If set_montage is True, rename_channels must be True."
+            )
         # create an index of all files and gather easily accessible info
         # without actually touching the files
-        file_paths = glob.glob(os.path.join(path, '**/*.edf'), recursive=True)
+        file_paths = glob.glob(os.path.join(path, "**/*.edf"), recursive=True)
         descriptions = _create_description(file_paths)
         # sort the descriptions chronologicaly
         descriptions = _sort_chronologically(descriptions)
@@ -62,59 +82,139 @@ class TUH(BaseConcatDataset):
                 # of recordings to load
                 recording_ids = range(recording_ids)
             descriptions = descriptions[recording_ids]
+
+        # workaround to ensure warnings are suppressed when running in parallel
+        def create_dataset(*args, **kwargs):
+            with warnings.catch_warnings():
+                warnings.filterwarnings(
+                    "ignore", message=".*not in description. '__getitem__'"
+                )
+                return self._create_dataset(*args, **kwargs)
+
         # this is the second loop (slow)
         # create datasets gathering more info about the files touching them
         # reading the raws and potentially preloading the data
         # disable joblib for tests. mocking seems to fail otherwise
         if n_jobs == 1:
-            base_datasets = [self._create_dataset(
-                descriptions[i], target_name, preload, add_physician_reports)
-                for i in descriptions.columns]
+            base_datasets = [
+                create_dataset(
+                    descriptions[i],
+                    target_name,
+                    preload,
+                    add_physician_reports,
+                    rename_channels,
+                    set_montage,
+                )
+                for i in descriptions.columns
+            ]
         else:
-            base_datasets = Parallel(n_jobs)(delayed(
-                self._create_dataset)(
-                descriptions[i], target_name, preload, add_physician_reports
-            ) for i in descriptions.columns)
+            base_datasets = Parallel(n_jobs)(
+                delayed(create_dataset)(
+                    descriptions[i],
+                    target_name,
+                    preload,
+                    add_physician_reports,
+                    rename_channels,
+                    set_montage,
+                )
+                for i in descriptions.columns
+            )
         super().__init__(base_datasets)
 
     @staticmethod
-    def _create_dataset(description, target_name, preload,
-                        add_physician_reports):
-        file_path = description.loc['path']
+    def _rename_channels(raw):
+        """
+        Renames the EEG channels using mne conventions and sets their type to 'eeg'.
+
+        See https://isip.piconepress.com/publications/reports/2020/tuh_eeg/electrodes/
+        """
+        # remove ref suffix and prefix:
+        # TODO: replace with removesuffix and removeprefix when 3.8 is dropped
+        mapping_strip = {
+            c: c.replace("-REF", "").replace("-LE", "").replace("EEG ", "")
+            for c in raw.ch_names
+        }
+        raw.rename_channels(mapping_strip)
+
+        montage1005 = mne.channels.make_standard_montage("standard_1005")
+        mapping_eeg_names = {
+            c.upper(): c for c in montage1005.ch_names if c.upper() in raw.ch_names
+        }
+
+        # Set channels whose type could not be inferred (defaulted to "eeg") to "misc":
+        non_eeg_names = [c for c in raw.ch_names if c not in mapping_eeg_names]
+        if non_eeg_names:
+            non_eeg_types = raw.get_channel_types(picks=non_eeg_names)
+            mapping_non_eeg_types = {
+                c: "misc" for c, t in zip(non_eeg_names, non_eeg_types) if t == "eeg"
+            }
+            if mapping_non_eeg_types:
+                raw.set_channel_types(mapping_non_eeg_types)
+
+        if mapping_eeg_names:
+            # Set 1005 channels type to "eeg":
+            raw.set_channel_types(
+                {c: "eeg" for c in mapping_eeg_names}, on_unit_change="ignore"
+            )
+            # Fix capitalized EEG channel names:
+            raw.rename_channels(mapping_eeg_names)
+
+    @staticmethod
+    def _set_montage(raw):
+        montage = mne.channels.make_standard_montage("standard_1005")
+        raw.set_montage(montage, on_missing="ignore")
+
+    @staticmethod
+    def _create_dataset(
+        description,
+        target_name,
+        preload,
+        add_physician_reports,
+        rename_channels,
+        set_montage,
+    ):
+        file_path = description.loc["path"]
 
         # parse age and gender information from EDF header
         age, gender = _parse_age_and_gender_from_edf_header(file_path)
-        raw = mne.io.read_raw_edf(file_path, preload=preload)
-
-        meas_date = datetime(1, 1, 1, tzinfo=timezone.utc) \
-            if raw.info['meas_date'] is None else raw.info['meas_date']
+        raw = mne.io.read_raw_edf(
+            file_path, preload=preload, infer_types=True, verbose="error"
+        )
+        if rename_channels:
+            TUH._rename_channels(raw)
+        if set_montage:
+            TUH._set_montage(raw)
+
+        meas_date = (
+            datetime(1, 1, 1, tzinfo=timezone.utc)
+            if raw.info["meas_date"] is None
+            else raw.info["meas_date"]
+        )
         # if this is old version of the data and the year could be parsed from
         # file paths, use this instead as before
-        if 'year' in description:
-            meas_date = meas_date.replace(
-                *description[['year', 'month', 'day']])
+        if "year" in description:
+            meas_date = meas_date.replace(*description[["year", "month", "day"]])
         raw.set_meas_date(meas_date)
 
         d = {
-            'age': int(age),
-            'gender': gender,
+            "age": int(age),
+            "gender": gender,
         }
         # if year exists in description = old version
         # if not, get it from meas_date in raw.info and add to description
         # if meas_date is None, create fake one
-        if 'year' not in description:
-            d['year'] = raw.info['meas_date'].year
-            d['month'] = raw.info['meas_date'].month
-            d['day'] = raw.info['meas_date'].day
+        if "year" not in description:
+            d["year"] = raw.info["meas_date"].year
+            d["month"] = raw.info["meas_date"].month
+            d["day"] = raw.info["meas_date"].day
 
         # read info relevant for preprocessing from raw without loading it
         if add_physician_reports:
             physician_report = _read_physician_report(file_path)
-            d['report'] = physician_report
+            d["report"] = physician_report
         additional_description = pd.Series(d)
         description = pd.concat([description, additional_description])
-        base_dataset = BaseDataset(raw, description,
-                                   target_name=target_name)
+        base_dataset = BaseDataset(raw, description, target_name=target_name)
         return base_dataset
 
 
@@ -126,30 +226,32 @@ def _create_description(file_paths):
 
 def _sort_chronologically(descriptions):
     descriptions.sort_values(
-        ["year", "month", "day", "subject", "session", "segment"],
-        axis=1, inplace=True)
+        ["year", "month", "day", "subject", "session", "segment"], axis=1, inplace=True
+    )
     return descriptions
 
 
 def _read_date(file_path):
-    date_path = file_path.replace('.edf', '_date.txt')
+    date_path = file_path.replace(".edf", "_date.txt")
     # if date file exists, read it
     if os.path.exists(date_path):
-        description = pd.read_json(date_path, typ='series').to_dict()
+        description = pd.read_json(date_path, typ="series").to_dict()
     # otherwise read edf file, extract date and store to file
     else:
-        raw = mne.io.read_raw_edf(file_path, preload=False, verbose='error')
+        raw = mne.io.read_raw_edf(file_path, preload=False, verbose="error")
         description = {
-            'year': raw.info['meas_date'].year,
-            'month': raw.info['meas_date'].month,
-            'day': raw.info['meas_date'].day,
+            "year": raw.info["meas_date"].year,
+            "month": raw.info["meas_date"].month,
+            "day": raw.info["meas_date"].day,
         }
         # if the txt file storing the recording date does not exist, create it
         try:
             pd.Series(description).to_json(date_path)
         except OSError:
-            warnings.warn(f'Cannot save date file to {date_path}. '
-                          f'This might slow down creation of the dataset.')
+            warnings.warn(
+                f"Cannot save date file to {date_path}. "
+                f"This might slow down creation of the dataset."
+            )
     return description
 
 
@@ -158,12 +260,12 @@ def _parse_description_from_file_path(file_path):
     file_path = os.path.normpath(file_path)
     tokens = file_path.split(os.sep)
     # Extract version number and tuh_eeg_abnormal/tuh_eeg from file path
-    if ('train' in tokens) or ('eval' in tokens):  # tuh_eeg_abnormal
+    if ("train" in tokens) or ("eval" in tokens):  # tuh_eeg_abnormal
         abnormal = True
         # Tokens[-2] is channel configuration (always 01_tcp_ar in abnormal)
         # on new versions, or
         #               session (e.g. s004_2013_08_15) on old versions
-        if tokens[-2].split('_')[0][0] == 's':  # s denoting session number
+        if tokens[-2].split("_")[0][0] == "s":  # s denoting session number
             version = tokens[-9]  # Before dec 2022 updata
         else:
             version = tokens[-6]  # After the dec 2022 update
@@ -181,22 +283,24 @@ def _parse_description_from_file_path(file_path):
         # or for abnormal:
         # tuh_eeg_abnormal/v3.0.0/edf/train/normal/
         #     01_tcp_ar/aaaaaaav_s004_t000.edf
-        subject_id = tokens[-1].split('_')[0]
-        session = tokens[-1].split('_')[1]
-        segment = tokens[-1].split('_')[2].split('.')[0]
+        subject_id = tokens[-1].split("_")[0]
+        session = tokens[-1].split("_")[1]
+        segment = tokens[-1].split("_")[2].split(".")[0]
         description = _read_date(file_path)
-        description.update({
-            'path': file_path,
-            'version': version,
-            'subject': subject_id,
-            'session': int(session[1:]),
-            'segment': int(segment[1:]),
-        })
+        description.update(
+            {
+                "path": file_path,
+                "version": version,
+                "subject": subject_id,
+                "session": int(session[1:]),
+                "segment": int(segment[1:]),
+            }
+        )
         if not abnormal:
-            year, month, day = tokens[-3].split('_')[1:]
-            description['year'] = int(year)
-            description['month'] = int(month)
-            description['day'] = int(day)
+            year, month, day = tokens[-3].split("_")[1:]
+            description["year"] = int(year)
+            description["month"] = int(month)
+            description["day"] = int(day)
         return description
     else:  # Old file path structure
         # expect file paths as tuh_eeg/version/file_type/reference/data_split/
@@ -208,43 +312,45 @@ def _parse_description_from_file_path(file_path):
         #     reference/subset/subject/recording session/file
         # v2.0.0/edf/train/normal/01_tcp_ar/000/00000021/
         #     s004_2013_08_15/00000021_s004_t000.edf
-        subject_id = tokens[-1].split('_')[0]
-        session = tokens[-2].split('_')[0]  # string on format 's000'
+        subject_id = tokens[-1].split("_")[0]
+        session = tokens[-2].split("_")[0]  # string on format 's000'
         # According to the example path in the comment 8 lines above,
         # segment is not included in the file name
-        segment = tokens[-1].split('_')[-1].split('.')[0]  # TODO: test with tuh_eeg
-        year, month, day = tokens[-2].split('_')[1:]
+        segment = tokens[-1].split("_")[-1].split(".")[0]  # TODO: test with tuh_eeg
+        year, month, day = tokens[-2].split("_")[1:]
         return {
-            'path': file_path,
-            'version': version,
-            'year': int(year),
-            'month': int(month),
-            'day': int(day),
-            'subject': int(subject_id),
-            'session': int(session[1:]),
-            'segment': int(segment[1:]),
+            "path": file_path,
+            "version": version,
+            "year": int(year),
+            "month": int(month),
+            "day": int(day),
+            "subject": int(subject_id),
+            "session": int(session[1:]),
+            "segment": int(segment[1:]),
         }
 
 
 def _read_physician_report(file_path):
     directory = os.path.dirname(file_path)
-    txt_file = glob.glob(os.path.join(directory, '**/*.txt'), recursive=True)
+    txt_file = glob.glob(os.path.join(directory, "**/*.txt"), recursive=True)
     # check that there is at most one txt file in the same directory
     assert len(txt_file) in [0, 1]
-    report = ''
+    report = ""
     if txt_file:
         txt_file = txt_file[0]
         # somewhere in the corpus, encoding apparently changed
         # first try to read as utf-8, if it does not work use latin-1
         try:
-            with open(txt_file, 'r', encoding='utf-8') as f:
+            with open(txt_file, "r", encoding="utf-8") as f:
                 report = f.read()
         except UnicodeDecodeError:
-            with open(txt_file, 'r', encoding='latin-1') as f:
+            with open(txt_file, "r", encoding="latin-1") as f:
                 report = f.read()
     if not report:
-        raise RuntimeError(f'Could not read physician report ({txt_file}). '
-                           f'Disable option or choose appropriate directory.')
+        raise RuntimeError(
+            f"Could not read physician report ({txt_file}). "
+            f"Disable option or choose appropriate directory."
+        )
     return report
 
 
@@ -292,20 +398,40 @@ class TUHAbnormal(TUH):
     add_physician_reports: bool
         If True, the physician reports will be read from disk and added to the
         description.
+    rename_channels: bool
+        If True, rename the EEG channels to the standard 10-05 system.
+    set_montage: bool
+        If True, set the montage to the standard 10-05 system.
+    n_jobs: int
+        Number of jobs to be used to read files in parallel.
     """
-    def __init__(self, path, recording_ids=None, target_name='pathological',
-                 preload=False, add_physician_reports=False, n_jobs=1):
-        with warnings.catch_warnings():
-            warnings.filterwarnings(
-                "ignore", message=".*not in description. '__getitem__'")
-            super().__init__(path=path, recording_ids=recording_ids,
-                             preload=preload, target_name=target_name,
-                             add_physician_reports=add_physician_reports,
-                             n_jobs=n_jobs)
+
+    def __init__(
+        self,
+        path: str,
+        recording_ids: list[int] | None = None,
+        target_name: str | tuple[str, ...] | None = "pathological",
+        preload: bool = False,
+        add_physician_reports: bool = False,
+        rename_channels: bool = False,
+        set_montage: bool = False,
+        n_jobs: int = 1,
+    ):
+        super().__init__(
+            path=path,
+            recording_ids=recording_ids,
+            preload=preload,
+            target_name=target_name,
+            add_physician_reports=add_physician_reports,
+            rename_channels=rename_channels,
+            set_montage=set_montage,
+            n_jobs=n_jobs,
+        )
         additional_descriptions = []
         for file_path in self.description.path:
-            additional_description = (
-                self._parse_additional_description_from_file_path(file_path))
+            additional_description = self._parse_additional_description_from_file_path(
+                file_path
+            )
             additional_descriptions.append(additional_description)
         additional_descriptions = pd.DataFrame(additional_descriptions)
         self.set_description(additional_descriptions, overwrite=True)
@@ -318,28 +444,45 @@ class TUHAbnormal(TUH):
         #                     reference/subset/subject/recording session/file
         # e.g.            v2.0.0/edf/train/normal/01_tcp_ar/000/00000021/
         #                     s004_2013_08_15/00000021_s004_t000.edf
-        assert ('abnormal' in tokens or 'normal' in tokens), (
-            'No pathology labels found.')
-        assert ('train' in tokens or 'eval' in tokens), (
-            'No train or eval set information found.')
+        assert "abnormal" in tokens or "normal" in tokens, "No pathology labels found."
+        assert "train" in tokens or "eval" in tokens, (
+            "No train or eval set information found."
+        )
         return {
-            'version': tokens[-9],
-            'train': 'train' in tokens,
-            'pathological': 'abnormal' in tokens,
+            "version": tokens[-9],
+            "train": "train" in tokens,
+            "pathological": "abnormal" in tokens,
         }
 
 
 def _fake_raw(*args, **kwargs):
     sfreq = 10
     ch_names = [
-        'EEG A1-REF', 'EEG A2-REF',
-        'EEG FP1-REF', 'EEG FP2-REF', 'EEG F3-REF', 'EEG F4-REF', 'EEG C3-REF',
-        'EEG C4-REF', 'EEG P3-REF', 'EEG P4-REF', 'EEG O1-REF', 'EEG O2-REF',
-        'EEG F7-REF', 'EEG F8-REF', 'EEG T3-REF', 'EEG T4-REF', 'EEG T5-REF',
-        'EEG T6-REF', 'EEG FZ-REF', 'EEG CZ-REF', 'EEG PZ-REF']
+        "EEG A1-REF",
+        "EEG A2-REF",
+        "EEG FP1-REF",
+        "EEG FP2-REF",
+        "EEG F3-REF",
+        "EEG F4-REF",
+        "EEG C3-REF",
+        "EEG C4-REF",
+        "EEG P3-REF",
+        "EEG P4-REF",
+        "EEG O1-REF",
+        "EEG O2-REF",
+        "EEG F7-REF",
+        "EEG F8-REF",
+        "EEG T3-REF",
+        "EEG T4-REF",
+        "EEG T5-REF",
+        "EEG T6-REF",
+        "EEG FZ-REF",
+        "EEG CZ-REF",
+        "EEG PZ-REF",
+    ]
     duration_min = 6
     data = np.random.randn(len(ch_names), duration_min * sfreq * 60)
-    info = mne.create_info(ch_names=ch_names, sfreq=sfreq, ch_types='eeg')
+    info = mne.create_info(ch_names=ch_names, sfreq=sfreq, ch_types="eeg")
     raw = mne.io.RawArray(data=data, info=info)
     return raw
 
@@ -351,54 +494,95 @@ def _get_header(*args, **kwargs):
 
 _TUH_EEG_PATHS = {
     # These are actual file paths and edf headers from the TUH EEG Corpus (v1.1.0 and v1.2.0)
-    'tuh_eeg/v1.1.0/edf/01_tcp_ar/000/00000000/s001_2015_12_30/00000000_s001_t000.edf': b'0       00000000 M 01-JAN-1978 00000000 Age:37                                          ',  # noqa E501
-    'tuh_eeg/v1.1.0/edf/01_tcp_ar/099/00009932/s004_2014_09_30/00009932_s004_t013.edf': b'0       00009932 F 01-JAN-1961 00009932 Age:53                                          ',  # noqa E501
-    'tuh_eeg/v1.1.0/edf/02_tcp_le/000/00000058/s001_2003_02_05/00000058_s001_t000.edf': b'0       00000058 M 01-JAN-2003 00000058 Age:0.0109                                      ',  # noqa E501
-    'tuh_eeg/v1.1.0/edf/03_tcp_ar_a/123/00012331/s003_2014_12_14/00012331_s003_t002.edf': b'0       00012331 M 01-JAN-1975 00012331 Age:39                                          ',  # noqa E501
-    'tuh_eeg/v1.2.0/edf/03_tcp_ar_a/149/00014928/s004_2016_01_15/00014928_s004_t007.edf': b'0       00014928 F 01-JAN-1933 00014928 Age:83                                          ',  # noqa E501
+    "tuh_eeg/v1.1.0/edf/01_tcp_ar/000/00000000/s001_2015_12_30/00000000_s001_t000.edf": b"0       00000000 M 01-JAN-1978 00000000 Age:37                                          ",
+    # noqa E501
+    "tuh_eeg/v1.1.0/edf/01_tcp_ar/099/00009932/s004_2014_09_30/00009932_s004_t013.edf": b"0       00009932 F 01-JAN-1961 00009932 Age:53                                          ",
+    # noqa E501
+    "tuh_eeg/v1.1.0/edf/02_tcp_le/000/00000058/s001_2003_02_05/00000058_s001_t000.edf": b"0       00000058 M 01-JAN-2003 00000058 Age:0.0109                                      ",
+    # noqa E501
+    "tuh_eeg/v1.1.0/edf/03_tcp_ar_a/123/00012331/s003_2014_12_14/00012331_s003_t002.edf": b"0       00012331 M 01-JAN-1975 00012331 Age:39                                          ",
+    # noqa E501
+    "tuh_eeg/v1.2.0/edf/03_tcp_ar_a/149/00014928/s004_2016_01_15/00014928_s004_t007.edf": b"0       00014928 F 01-JAN-1933 00014928 Age:83                                          ",
+    # noqa E501
 }
 _TUH_EEG_ABNORMAL_PATHS = {
     # these are actual file paths and edf headers from TUH Abnormal EEG Corpus (v2.0.0)
-    'tuh_abnormal_eeg/v2.0.0/edf/train/normal/01_tcp_ar/078/00007871/s001_2011_07_05/00007871_s001_t001.edf': b'0       00007871 F 01-JAN-1988 00007871 Age:23                                          ',  # noqa E501
-    'tuh_abnormal_eeg/v2.0.0/edf/train/normal/01_tcp_ar/097/00009777/s001_2012_09_17/00009777_s001_t000.edf': b'0       00009777 M 01-JAN-1986 00009777 Age:26                                          ',  # noqa E501
-    'tuh_abnormal_eeg/v2.0.0/edf/train/abnormal/01_tcp_ar/083/00008393/s002_2012_02_21/00008393_s002_t000.edf': b'0       00008393 M 01-JAN-1960 00008393 Age:52                                          ',  # noqa E501
-    'tuh_abnormal_eeg/v2.0.0/edf/train/abnormal/01_tcp_ar/012/00001200/s003_2010_12_06/00001200_s003_t000.edf': b'0       00001200 M 01-JAN-1963 00001200 Age:47                                          ',  # noqa E501
-    'tuh_abnormal_eeg/v2.0.0/edf/eval/abnormal/01_tcp_ar/059/00005932/s004_2013_03_14/00005932_s004_t000.edf': b'0       00005932 M 01-JAN-1963 00005932 Age:50                                          ',  # noqa E501
+    "tuh_abnormal_eeg/v2.0.0/edf/train/normal/01_tcp_ar/078/00007871/s001_2011_07_05/00007871_s001_t001.edf": b"0       00007871 F 01-JAN-1988 00007871 Age:23                                          ",
+    # noqa E501
+    "tuh_abnormal_eeg/v2.0.0/edf/train/normal/01_tcp_ar/097/00009777/s001_2012_09_17/00009777_s001_t000.edf": b"0       00009777 M 01-JAN-1986 00009777 Age:26                                          ",
+    # noqa E501
+    "tuh_abnormal_eeg/v2.0.0/edf/train/abnormal/01_tcp_ar/083/00008393/s002_2012_02_21/00008393_s002_t000.edf": b"0       00008393 M 01-JAN-1960 00008393 Age:52                                          ",
+    # noqa E501
+    "tuh_abnormal_eeg/v2.0.0/edf/train/abnormal/01_tcp_ar/012/00001200/s003_2010_12_06/00001200_s003_t000.edf": b"0       00001200 M 01-JAN-1963 00001200 Age:47                                          ",
+    # noqa E501
+    "tuh_abnormal_eeg/v2.0.0/edf/eval/abnormal/01_tcp_ar/059/00005932/s004_2013_03_14/00005932_s004_t000.edf": b"0       00005932 M 01-JAN-1963 00005932 Age:50                                          ",
+    # noqa E501
 }
 
 
 class _TUHMock(TUH):
     """Mocked class for testing and examples."""
-    @mock.patch('glob.glob', return_value=_TUH_EEG_PATHS.keys())
-    @mock.patch('mne.io.read_raw_edf', new=_fake_raw)
-    @mock.patch('braindecode.datasets.tuh._read_edf_header',
-                new=_get_header)
-    def __init__(self, mock_glob, path, recording_ids=None, target_name=None,
-                 preload=False, add_physician_reports=False, n_jobs=1):
+
+    @mock.patch("glob.glob", return_value=_TUH_EEG_PATHS.keys())
+    @mock.patch("mne.io.read_raw_edf", new=_fake_raw)
+    @mock.patch("braindecode.datasets.tuh._read_edf_header", new=_get_header)
+    def __init__(
+        self,
+        mock_glob,
+        path: str,
+        recording_ids: list[int] | None = None,
+        target_name: str | tuple[str, ...] | None = None,
+        preload: bool = False,
+        add_physician_reports: bool = False,
+        rename_channels: bool = False,
+        set_montage: bool = False,
+        n_jobs: int = 1,
+    ):
         with warnings.catch_warnings():
-            warnings.filterwarnings(
-                "ignore", message="Cannot save date file")
-            super().__init__(path=path, recording_ids=recording_ids,
-                             target_name=target_name, preload=preload,
-                             add_physician_reports=add_physician_reports,
-                             n_jobs=n_jobs)
+            warnings.filterwarnings("ignore", message="Cannot save date file")
+            super().__init__(
+                path=path,
+                recording_ids=recording_ids,
+                target_name=target_name,
+                preload=preload,
+                add_physician_reports=add_physician_reports,
+                rename_channels=rename_channels,
+                set_montage=set_montage,
+                n_jobs=n_jobs,
+            )
 
 
 class _TUHAbnormalMock(TUHAbnormal):
     """Mocked class for testing and examples."""
-    @mock.patch('glob.glob', return_value=_TUH_EEG_ABNORMAL_PATHS.keys())
-    @mock.patch('mne.io.read_raw_edf', new=_fake_raw)
-    @mock.patch('braindecode.datasets.tuh._read_edf_header',
-                new=_get_header)
-    @mock.patch('braindecode.datasets.tuh._read_physician_report',
-                return_value='simple_test')
-    def __init__(self, mock_glob, mock_report, path, recording_ids=None,
-                 target_name='pathological', preload=False,
-                 add_physician_reports=False, n_jobs=1):
+
+    @mock.patch("glob.glob", return_value=_TUH_EEG_ABNORMAL_PATHS.keys())
+    @mock.patch("mne.io.read_raw_edf", new=_fake_raw)
+    @mock.patch("braindecode.datasets.tuh._read_edf_header", new=_get_header)
+    @mock.patch(
+        "braindecode.datasets.tuh._read_physician_report", return_value="simple_test"
+    )
+    def __init__(
+        self,
+        mock_glob,
+        mock_report,
+        path: str,
+        recording_ids: list[int] | None = None,
+        target_name: str | tuple[str, ...] | None = "pathological",
+        preload: bool = False,
+        add_physician_reports: bool = False,
+        rename_channels: bool = False,
+        set_montage: bool = False,
+        n_jobs: int = 1,
+    ):
         with warnings.catch_warnings():
-            warnings.filterwarnings(
-                "ignore", message="Cannot save date file")
-            super().__init__(path=path, recording_ids=recording_ids,
-                             target_name=target_name, preload=preload,
-                             add_physician_reports=add_physician_reports,
-                             n_jobs=n_jobs)
+            warnings.filterwarnings("ignore", message="Cannot save date file")
+            super().__init__(
+                path=path,
+                recording_ids=recording_ids,
+                target_name=target_name,
+                preload=preload,
+                add_physician_reports=add_physician_reports,
+                rename_channels=rename_channels,
+                set_montage=set_montage,
+                n_jobs=n_jobs,
+            )