eegdash 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

eegdash/data_config.py

@@ -1,5 +1,8 @@
 config = {
     "required_fields": ["data_name"],
+    # Default set of user-facing primary record attributes expected in the database. Records
+    # where any of these are missing will be loaded with the respective attribute set to None.
+    # Additional fields may be returned if they are present in the database, notably bidsdependencies.
     "attributes": {
         "data_name": "str",
         "dataset": "str",
@@ -11,9 +14,12 @@ config = {
         "sampling_frequency": "float",
         "modality": "str",
         "nchans": "int",
-        "ntimes": "int",
+        "ntimes": "int",  # note: this is really the number of seconds in the data, rounded down
     },
+    # queryable descriptive fields for a given recording
     "description_fields": ["subject", "session", "run", "task", "age", "gender", "sex"],
+    # list of filenames that may be present in the BIDS dataset directory that are used
+    # to load and interpret a given BIDS recording.
     "bids_dependencies_files": [
         "dataset_description.json",
         "participants.tsv",

eegdash/data_utils.py

@@ -1,9 +1,9 @@
 import json
+import logging
 import os
 import re
-import sys
-import tempfile
 from pathlib import Path
+from typing import Any
 
 import mne
 import mne_bids
@@ -12,7 +12,7 @@ import pandas as pd
 import s3fs
 from bids import BIDSLayout
 from joblib import Parallel, delayed
-from mne._fiff.utils import _find_channels, _read_segments_file
+from mne._fiff.utils import _read_segments_file
 from mne.io import BaseRaw
 from mne_bids import (
     BIDSPath,
@@ -20,51 +20,62 @@ from mne_bids import (
 
 from braindecode.datasets import BaseDataset
 
+logger = logging.getLogger("eegdash")
 
-class EEGDashBaseDataset(BaseDataset):
-    """Returns samples from an mne.io.Raw object along with a target.
 
-    Dataset which serves samples from an mne.io.Raw object along with a target.
-    The target is unique for the dataset, and is obtained through the
-    `description` attribute.
+class EEGDashBaseDataset(BaseDataset):
+    """A single EEG recording hosted on AWS S3 and cached locally upon first access.
 
-    Parameters
-    ----------
-    raw : mne.io.Raw
-        Continuous data.
-    description : dict | pandas.Series | None
-        Holds additional description about the continuous signal / subject.
-    target_name : str | tuple | None
-        Name(s) of the index in `description` that should be used to provide the
-        target (e.g., to be used in a prediction task later on).
-    transform : callable | None
-        On-the-fly transform applied to the example before it is returned.
+    This is a subclass of braindecode's BaseDataset, which can consequently be used in
+    conjunction with the preprocessing and training pipelines of braindecode.
     """
 
     AWS_BUCKET = "s3://openneuro.org"
 
-    def __init__(self, record, cache_dir, **kwargs):
+    def __init__(
+        self,
+        record: dict[str, Any],
+        cache_dir: str,
+        s3_bucket: str | None = None,
+        **kwargs,
+    ):
+        """Create a new EEGDashBaseDataset instance. Users do not usually need to call this
+        directly -- instead use the EEGDashDataset class to load a collection of these
+        recordings from a local BIDS folder or using a database query.
+
+        Parameters
+        ----------
+        record : dict
+            A fully resolved metadata record for the data to load.
+        cache_dir : str
+            A local directory where the data will be cached.
+        kwargs : dict
+            Additional keyword arguments to pass to the BaseDataset constructor.
+
+        """
         super().__init__(None, **kwargs)
         self.record = record
         self.cache_dir = Path(cache_dir)
         bids_kwargs = self.get_raw_bids_args()
+
         self.bidspath = BIDSPath(
             root=self.cache_dir / record["dataset"],
             datatype="eeg",
             suffix="eeg",
             **bids_kwargs,
         )
+        self.s3_bucket = s3_bucket if s3_bucket else self.AWS_BUCKET
         self.s3file = self.get_s3path(record["bidspath"])
         self.filecache = self.cache_dir / record["bidspath"]
         self.bids_dependencies = record["bidsdependencies"]
         self._raw = None
-        # if os.path.exists(self.filecache):
-        #     self.raw = mne_bids.read_raw_bids(self.bidspath, verbose=False)
 
-    def get_s3path(self, filepath):
-        return f"{self.AWS_BUCKET}/{filepath}"
+    def get_s3path(self, filepath: str) -> str:
+        """Helper to form an AWS S3 URI for the given relative filepath."""
+        return f"{self.s3_bucket}/{filepath}"
 
-    def _download_s3(self):
+    def _download_s3(self) -> None:
+        """Fetch the given data from its S3 location and cache it locally."""
         self.filecache.parent.mkdir(parents=True, exist_ok=True)
         filesystem = s3fs.S3FileSystem(
             anon=True, client_kwargs={"region_name": "us-east-2"}
@@ -72,7 +83,10 @@ class EEGDashBaseDataset(BaseDataset):
         filesystem.download(self.s3file, self.filecache)
         self.filenames = [self.filecache]
 
-    def _download_dependencies(self):
+    def _download_dependencies(self) -> None:
+        """Download all BIDS dependency files (metadata files, recording sidecar files)
+        from S3 and cache them locally.
+        """
         filesystem = s3fs.S3FileSystem(
             anon=True, client_kwargs={"region_name": "us-east-2"}
         )
@@ -83,11 +97,15 @@ class EEGDashBaseDataset(BaseDataset):
                 filepath.parent.mkdir(parents=True, exist_ok=True)
                 filesystem.download(s3path, filepath)
 
-    def get_raw_bids_args(self):
+    def get_raw_bids_args(self) -> dict[str, Any]:
+        """Helper to restrict the metadata record to the fields needed to locate a BIDS
+        recording.
+        """
         desired_fields = ["subject", "session", "task", "run"]
         return {k: self.record[k] for k in desired_fields if self.record[k]}
 
-    def check_and_get_raw(self):
+    def check_and_get_raw(self) -> None:
+        """Download the S3 file and BIDS dependencies if not already cached."""
         if not os.path.exists(self.filecache):  # not preload
             if self.bids_dependencies:
                 self._download_dependencies()
@@ -95,9 +113,10 @@ class EEGDashBaseDataset(BaseDataset):
         if self._raw is None:
             self._raw = mne_bids.read_raw_bids(self.bidspath, verbose=False)
 
-    def __getitem__(self, index):
-        # self.check_and_get_raw()
+    # === BaseDataset and PyTorch Dataset interface ===
 
+    def __getitem__(self, index):
+        """Main function to access a sample from the dataset."""
         X = self.raw[:, index][0]
         y = None
         if self.target_name is not None:
@@ -108,14 +127,20 @@ class EEGDashBaseDataset(BaseDataset):
             X = self.transform(X)
         return X, y
 
-    def __len__(self):
+    def __len__(self) -> int:
+        """Return the number of samples in the dataset."""
         if self._raw is None:
+            # FIXME: this is a bit strange and should definitely not change as a side effect
+            #  of accessing the data (which it will, since ntimes is the actual length but rounded down)
             return int(self.record["ntimes"] * self.record["sampling_frequency"])
         else:
             return len(self._raw)
 
     @property
     def raw(self):
+        """Return the MNE Raw object for this recording. This will perform the actual
+        retrieval if not yet done so.
+        """
         if self._raw is None:
             self.check_and_get_raw()
         return self._raw
@@ -126,50 +151,44 @@ class EEGDashBaseDataset(BaseDataset):
 
 
 class EEGDashBaseRaw(BaseRaw):
-    r"""MNE Raw object from EEG-Dash connection with Openneuro S3 file.
+    """Wrapper around the MNE BaseRaw class that automatically fetches the data from S3
+    (when _read_segment is called) and caches it locally. Currently for internal use.
 
     Parameters
     ----------
     input_fname : path-like
         Path to the S3 file
-    eog : list | tuple | 'auto'
-        Names or indices of channels that should be designated EOG channels.
-        If 'auto', the channel names containing ``EOG`` or ``EYE`` are used.
-        Defaults to empty tuple.
-    %(preload)s
-        Note that preload=False will be effective only if the data is stored
-        in a separate binary file.
-    %(uint16_codec)s
-    %(montage_units)s
-    %(verbose)s
+    metadata : dict
+        The metadata record for the recording (e.g., from the database).
+    preload : bool
+        Whether to pre-loaded the data before the first access.
+    cache_dir : str
+        Local path under which the data will be cached.
+    bids_dependencies : list
+        List of additional BIDS metadata files that should be downloaded and cached
+        alongside the main recording file.
+    verbose : str | int | None
+        Optionally the verbosity level for MNE logging (see MNE documentation for possible values).
 
     See Also
     --------
     mne.io.Raw : Documentation of attributes and methods.
 
-    Notes
-    -----
-    .. versionadded:: 0.11.0
     """
 
     AWS_BUCKET = "s3://openneuro.org"
 
     def __init__(
         self,
-        input_fname,
-        metadata,
-        eog=(),
-        preload=False,
+        input_fname: str,
+        metadata: dict[str, Any],
+        preload: bool = False,
         *,
-        cache_dir="./.eegdash_cache",
-        bids_dependencies: list = [],
-        uint16_codec=None,
-        montage_units="auto",
-        verbose=None,
+        cache_dir: str = "./.eegdash_cache",
+        bids_dependencies: list[str] = [],
+        verbose: Any = None,
     ):
-        """
-        Get to work with S3 endpoint first, no caching
-        """
+        """Get to work with S3 endpoint first, no caching"""
         # Create a simple RawArray
         sfreq = metadata["sfreq"]  # Sampling frequency
         n_times = metadata["n_times"]
@@ -237,6 +256,20 @@ class EEGDashBaseRaw(BaseRaw):
 
 
 class EEGBIDSDataset:
+    """A one-stop shop interface to a local BIDS dataset containing EEG recordings.
+
+    This is mainly tailored to the needs of EEGDash application and is used to centralize
+    interactions with the BIDS dataset, such as parsing the metadata.
+
+    Parameters
+    ----------
+    data_dir : str | Path
+        The path to the local BIDS dataset directory.
+    dataset : str
+        A name for the dataset.
+
+    """
+
     ALLOWED_FILE_FORMAT = ["eeglab", "brainvision", "biosemi", "european"]
     RAW_EXTENSIONS = {
         ".set": [".set", ".fdt"],  # eeglab
@@ -270,19 +303,13 @@ class EEGBIDSDataset:
             "Unable to construct EEG dataset. No EEG recordings found."
         )
         assert self.check_eeg_dataset(), ValueError("Dataset is not an EEG dataset.")
-        # temp_dir = (Path().resolve() / 'data')
-        # if not os.path.exists(temp_dir):
-        #     os.mkdir(temp_dir)
-        # if not os.path.exists(temp_dir / f'{dataset}_files.npy'):
-        #     self.files = self.get_files_with_extension_parallel(self.bidsdir, extension=self.RAW_EXTENSION[self.raw_format])
-        #     np.save(temp_dir / f'{dataset}_files.npy', self.files)
-        # else:
-        #     self.files = np.load(temp_dir / f'{dataset}_files.npy', allow_pickle=True)
-
-    def check_eeg_dataset(self):
+
+    def check_eeg_dataset(self) -> bool:
+        """Check if the dataset is EEG."""
         return self.get_bids_file_attribute("modality", self.files[0]).lower() == "eeg"
 
-    def get_recordings(self, layout: BIDSLayout):
+    def get_recordings(self, layout: BIDSLayout) -> list[str]:
+        """Get a list of all EEG recording files in the BIDS layout."""
         files = []
         for ext, exts in self.RAW_EXTENSIONS.items():
             files = layout.get(extension=ext, return_type="filename")
@@ -290,11 +317,15 @@ class EEGBIDSDataset:
                 break
         return files
 
-    def get_relative_bidspath(self, filename):
+    def get_relative_bidspath(self, filename: str) -> str:
+        """Make the given file path relative to the BIDS directory."""
         bids_parent_dir = self.bidsdir.parent.absolute()
         return str(Path(filename).relative_to(bids_parent_dir))
 
-    def get_property_from_filename(self, property, filename):
+    def get_property_from_filename(self, property: str, filename: str) -> str:
+        """Parse a property out of a BIDS-compliant filename. Returns an empty string
+        if not found.
+        """
         import platform
 
         if platform.system() == "Windows":
@@ -303,25 +334,38 @@ class EEGBIDSDataset:
             lookup = re.search(rf"{property}-(.*?)[_\/]", filename)
         return lookup.group(1) if lookup else ""
 
-    def merge_json_inheritance(self, json_files):
-        """
-        Merge list of json files found by get_bids_file_inheritance,
-        expecting the order (from left to right) is from lowest level to highest level,
-        and return a merged dictionary
+    def merge_json_inheritance(self, json_files: list[str | Path]) -> dict:
+        """Internal helper to merge list of json files found by get_bids_file_inheritance,
+        expecting the order (from left to right) is from lowest
+        level to highest level, and return a merged dictionary
         """
         json_files.reverse()
         json_dict = {}
         for f in json_files:
-            json_dict.update(json.load(open(f)))
+            json_dict.update(json.load(open(f)))  # FIXME: should close file
         return json_dict
 
-    def get_bids_file_inheritance(self, path, basename, extension):
-        """
-        Get all files with given extension that applies to the basename file
-        following the BIDS inheritance principle in the order of lowest level first
-        @param
-            basename: bids file basename without _eeg.set extension for example
-            extension: e.g. channels.tsv
+    def get_bids_file_inheritance(
+        self, path: str | Path, basename: str, extension: str
+    ) -> list[Path]:
+        """Get all file paths that apply to the basename file in the specified directory
+        and that end with the specified suffix, recursively searching parent directories
+        (following the BIDS inheritance principle in the order of lowest level first).
+
+        Parameters
+        ----------
+        path : str | Path
+            The directory path to search for files.
+        basename : str
+            BIDS file basename without _eeg.set extension for example
+        extension : str
+            Only consider files that end with the specified suffix; e.g. channels.tsv
+
+        Returns
+        -------
+        list[Path]
+            A list of file paths that match the given basename and extension.
+
         """
         top_level_files = ["README", "dataset_description.json", "participants.tsv"]
         bids_files = []
@@ -352,17 +396,25 @@ class EEGBIDSDataset:
             )
             return bids_files
 
-    def get_bids_metadata_files(self, filepath, metadata_file_extension):
-        """
-        (Wrapper for self.get_bids_file_inheritance)
-        Get all BIDS metadata files that are associated with the given filepath, following the BIDS inheritance principle.
-
-        Args:
-            filepath (str or Path): The filepath to get the associated metadata files for.
-            metadata_files_extensions (list): A list of file extensions to search for metadata files.
+    def get_bids_metadata_files(
+        self, filepath: str | Path, metadata_file_extension: list[str]
+    ) -> list[Path]:
+        """Retrieve all metadata file paths that apply to a given data file path and that
+        end with a specific suffix (following the BIDS inheritance principle).
+
+        Parameters
+        ----------
+        filepath: str | Path
+            The filepath to get the associated metadata files for.
+        metadata_file_extension : str
+            Consider only metadata files that end with the specified suffix,
+            e.g., channels.tsv or eeg.json
+
+        Returns
+        -------
+        list[Path]:
+            A list of filepaths for all matching metadata files
 
-        Returns:
-            list: A list of filepaths for all the associated metadata files
         """
         if isinstance(filepath, str):
             filepath = Path(filepath)
@@ -376,7 +428,11 @@ class EEGBIDSDataset:
         )
         return meta_files
 
-    def scan_directory(self, directory, extension):
+    def scan_directory(self, directory: str, extension: str) -> list[Path]:
+        """Return a list of file paths that end with the given extension in the specified
+        directory. Ignores certain special directories like .git, .datalad, derivatives,
+        and code.
+        """
         result_files = []
         directory_to_ignore = [".git", ".datalad", "derivatives", "code"]
         with os.scandir(directory) as entries:
@@ -391,14 +447,35 @@ class EEGBIDSDataset:
         return result_files
 
     def get_files_with_extension_parallel(
-        self, directory, extension=".set", max_workers=-1
-    ):
+        self, directory: str, extension: str = ".set", max_workers: int = -1
+    ) -> list[Path]:
+        """Efficiently scan a directory and its subdirectories for files that end with
+        the given extension.
+
+        Parameters
+        ----------
+        directory : str
+            The root directory to scan for files.
+        extension : str
+            Only consider files that end with this suffix, e.g. '.set'.
+        max_workers : int
+            Optionally specify the maximum number of worker threads to use for parallel scanning.
+            Defaults to all available CPU cores if set to -1.
+
+        Returns
+        -------
+        list[Path]:
+            A list of filepaths for all matching metadata files
+
+        """
         result_files = []
         dirs_to_scan = [directory]
 
         # Use joblib.Parallel and delayed to parallelize directory scanning
         while dirs_to_scan:
-            print(f"Scanning {len(dirs_to_scan)} directories...", dirs_to_scan)
+            logger.info(
+                f"Directories to scan: {len(dirs_to_scan)}, files: {dirs_to_scan}"
+            )
             # Run the scan_directory function in parallel across directories
             results = Parallel(n_jobs=max_workers, prefer="threads", verbose=1)(
                 delayed(self.scan_directory)(d, extension) for d in dirs_to_scan
@@ -412,12 +489,18 @@ class EEGBIDSDataset:
                         dirs_to_scan.append(path)  # Queue up subdirectories to scan
                     else:
                         result_files.append(path)  # Add files to the final result
-            print(f"Current number of files: {len(result_files)}")
+            logger.info(f"Found {len(result_files)} files.")
 
         return result_files
 
-    def load_and_preprocess_raw(self, raw_file, preprocess=False):
-        print(f"Loading {raw_file}")
+    def load_and_preprocess_raw(
+        self, raw_file: str, preprocess: bool = False
+    ) -> np.ndarray:
+        """Utility function to load a raw data file with MNE and apply some simple
+        (hardcoded) preprocessing and return as a numpy array. Not meant for purposes
+        other than testing or debugging.
+        """
+        logger.info(f"Loading raw data from {raw_file}")
         EEG = mne.io.read_raw_eeglab(raw_file, preload=True, verbose="error")
 
         if preprocess:
@@ -429,9 +512,6 @@ class EEGBIDSDataset:
             sfreq = 128
             if EEG.info["sfreq"] != sfreq:
                 EEG = EEG.resample(sfreq)
-            # # normalize data to zero mean and unit variance
-            # scalar = preprocessing.StandardScaler()
-            # mat_data = scalar.fit_transform(mat_data.T).T # scalar normalize for each feature and expects shape data x features
 
         mat_data = EEG.get_data()
 
@@ -439,17 +519,22 @@ class EEGBIDSDataset:
             raise ValueError("Expect raw data to be CxT dimension")
         return mat_data
 
-    def get_files(self):
+    def get_files(self) -> list[Path]:
+        """Get all EEG recording file paths (with valid extensions) in the BIDS folder."""
         return self.files
 
-    def resolve_bids_json(self, json_files: list):
-        """
-        Resolve the BIDS JSON files and return a dictionary of the resolved values.
-        Args:
-            json_files (list): A list of JSON files to resolve in order of leaf level first
+    def resolve_bids_json(self, json_files: list[str]) -> dict:
+        """Resolve the BIDS JSON files and return a dictionary of the resolved values.
 
-        Returns:
+        Parameters
+        ----------
+        json_files : list
+            A list of JSON file paths to resolve in order of leaf level first.
+
+        Returns
+        -------
             dict: A dictionary of the resolved values.
+
         """
         if len(json_files) == 0:
             raise ValueError("No JSON files provided")
@@ -461,7 +546,10 @@ class EEGBIDSDataset:
                 json_dict.update(json.load(f))
         return json_dict
 
-    def get_bids_file_attribute(self, attribute, data_filepath):
+    def get_bids_file_attribute(self, attribute: str, data_filepath: str) -> Any:
+        """Retrieve a specific attribute from the BIDS file metadata applicable
+        to the provided recording file path.
+        """
         entities = self.layout.parse_file_entities(data_filepath)
         bidsfile = self.layout.get(**entities)[0]
         attributes = bidsfile.get_entities(metadata="all")
@@ -478,27 +566,32 @@ class EEGBIDSDataset:
         attribute_value = attributes.get(attribute_mapping.get(attribute), None)
         return attribute_value
 
-    def channel_labels(self, data_filepath):
+    def channel_labels(self, data_filepath: str) -> list[str]:
+        """Get a list of channel labels for the given data file path."""
         channels_tsv = pd.read_csv(
             self.get_bids_metadata_files(data_filepath, "channels.tsv")[0], sep="\t"
         )
         return channels_tsv["name"].tolist()
 
-    def channel_types(self, data_filepath):
+    def channel_types(self, data_filepath: str) -> list[str]:
+        """Get a list of channel types for the given data file path."""
         channels_tsv = pd.read_csv(
             self.get_bids_metadata_files(data_filepath, "channels.tsv")[0], sep="\t"
         )
         return channels_tsv["type"].tolist()
 
-    def num_times(self, data_filepath):
+    def num_times(self, data_filepath: str) -> int:
+        """Get the approximate number of time points in the EEG recording based on the BIDS metadata."""
         eeg_jsons = self.get_bids_metadata_files(data_filepath, "eeg.json")
         eeg_json_dict = self.merge_json_inheritance(eeg_jsons)
         return int(
             eeg_json_dict["SamplingFrequency"] * eeg_json_dict["RecordingDuration"]
         )
 
-    def subject_participant_tsv(self, data_filepath):
-        """Get participants_tsv info of a subject based on filepath"""
+    def subject_participant_tsv(self, data_filepath: str) -> dict[str, Any]:
+        """Get BIDS participants.tsv record for the subject to which the given file
+        path corresponds, as a dictionary.
+        """
         participants_tsv = pd.read_csv(
             self.get_bids_metadata_files(data_filepath, "participants.tsv")[0], sep="\t"
         )
@@ -510,12 +603,16 @@ class EEGBIDSDataset:
         subject = f"sub-{self.get_bids_file_attribute('subject', data_filepath)}"
         return participants_tsv.loc[subject].to_dict()
 
-    def eeg_json(self, data_filepath):
+    def eeg_json(self, data_filepath: str) -> dict[str, Any]:
+        """Get BIDS eeg.json metadata for the given data file path."""
         eeg_jsons = self.get_bids_metadata_files(data_filepath, "eeg.json")
         eeg_json_dict = self.merge_json_inheritance(eeg_jsons)
         return eeg_json_dict
 
-    def channel_tsv(self, data_filepath):
+    def channel_tsv(self, data_filepath: str) -> dict[str, Any]:
+        """Get BIDS channels.tsv metadata for the given data file path, as a dictionary
+        of lists and/or single values.
+        """
         channels_tsv = pd.read_csv(
             self.get_bids_metadata_files(data_filepath, "channels.tsv")[0], sep="\t"
         )

eegdash/dataset.py

@@ -0,0 +1,60 @@
+from .api import EEGDashDataset
+
+
+class EEGChallengeDataset(EEGDashDataset):
+    def __init__(
+        self,
+        release: str = "R5",
+        cache_dir: str = ".eegdash_cache",
+        s3_bucket: str | None = "s3://nmdatasets/NeurIPS25/R5_L100",
+        **kwargs,
+    ):
+        """Create a new EEGDashDataset from a given query or local BIDS dataset directory
+        and dataset name. An EEGDashDataset is pooled collection of EEGDashBaseDataset
+        instances (individual recordings) and is a subclass of braindecode's BaseConcatDataset.
+
+        Parameters
+        ----------
+        query : dict | None
+            Optionally a dictionary that specifies the query to be executed; see
+            EEGDash.find() for details on the query format.
+        data_dir : str | list[str] | None
+            Optionally a string or a list of strings specifying one or more local
+            BIDS dataset directories from which to load the EEG data files. Exactly one
+            of query or data_dir must be provided.
+        dataset : str | list[str] | None
+            If data_dir is given, a name or list of names for for the dataset(s) to be loaded.
+        description_fields : list[str]
+            A list of fields to be extracted from the dataset records
+            and included in the returned data description(s). Examples are typical
+            subject metadata fields such as "subject", "session", "run", "task", etc.;
+            see also data_config.description_fields for the default set of fields.
+        cache_dir : str
+            A directory where the dataset will be cached locally.
+        s3_bucket : str | None
+            An optional S3 bucket URI to use instead of the
+            default OpenNeuro bucket for loading data files.
+        kwargs : dict
+            Additional keyword arguments to be passed to the EEGDashBaseDataset
+            constructor.
+
+        """
+        dsnumber_release_map = {
+            "R11": "ds005516",
+            "R10": "ds005515",
+            "R9": "ds005514",
+            "R8": "ds005512",
+            "R7": "ds005511",
+            "R6": "ds005510",
+            "R4": "ds005508",
+            "R5": "ds005509",
+            "R3": "ds005507",
+            "R2": "ds005506",
+            "R1": "ds005505",
+        }
+        super().__init__(
+            query={"dataset": dsnumber_release_map[release]},
+            cache_dir=cache_dir,
+            s3_bucket=s3_bucket,
+            **kwargs,
+        )