braindecode 1.3.2.dev173417708__tar.gz → 1.3.2.dev173801889__tar.gz

{braindecode-1.3.2.dev173417708/braindecode.egg-info → braindecode-1.3.2.dev173801889}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.3.2.dev173417708
+Version: 1.3.2.dev173801889
 Summary: Deep learning software to decode EEG, ECG or MEG signals
 Author-email: Robin Tibor Schirrmeister <robintibor@gmail.com>, Bruno Aristimunha Pinto <b.aristimunha@gmail.com>, Alexandre Gramfort <agramfort@meta.com>
 Maintainer-email: Alexandre Gramfort <agramfort@meta.com>, Bruno Aristimunha Pinto <b.aristimunha@gmail.com>, Robin Tibor Schirrmeister <robintibor@gmail.com>

{braindecode-1.3.2.dev173417708 → braindecode-1.3.2.dev173801889}/braindecode/datasets/base.py

@@ -579,6 +579,89 @@ class EEGWindowsDataset(RecordDataset):
             self.metadata,
         )
 
+    def to_epochs_dataset(self) -> WindowsDataset:
+        """Converts this :class:`EEGWindowsDataset` to :class:`WindowsDataset` with ``mne.Epochs``.
+
+        In Braindecode, the data can either be stored as ``mne.io.Raw`` (in :class:`EEGWindowsDataset`)
+        or as ``mne.Epochs`` (in :class:`WindowsDataset`). This function converts from the first type to the second,
+        which can be useful for reducing disk space when you want to save a dataset.
+
+        Returns
+        -------
+        WindowsDataset
+            A new :class:`WindowsDataset` with ``mne.Epochs``.
+
+        Raises
+        ------
+        ValueError
+            If the dataset is not compatible with conversion. This includes:
+            - If targets are not obtained from metadata
+            - If windows have inconsistent sizes or there are no windows to convert
+            - If raw.first_samp!=0 (not supported for simplicity)
+        """
+        # Check the targets:
+        if self.targets_from != "metadata":
+            raise ValueError(
+                "to_epochs_dataset only works if targets are obtained from metadata."
+            )
+
+        # Check and get window sizes:
+        i_start_in_trial = self.crop_inds[:, 1]
+        i_stop_in_trial = self.crop_inds[:, 2]
+        sizes = np.unique(i_stop_in_trial - i_start_in_trial)
+        if len(sizes) > 1:
+            raise ValueError("Windows have inconsistent sizes.")
+        if len(sizes) != 1:
+            raise ValueError("No windows to convert.")
+        input_window_seconds = (sizes[0] - 1) / self.raw.info["sfreq"]
+
+        # Check raw.first_samp:
+        if self.raw.first_samp != 0:
+            raise ValueError(
+                f"to_epochs_dataset only works if raw.first_samp is 0, found {self.raw.first_samp=}"
+            )
+
+        # Create events and epochs:
+        events = np.zeros((len(self), 3), dtype=int)
+        events[:, 0] = i_start_in_trial
+        events[:, 2] = 1
+        epochs = mne.Epochs(
+            raw=self.raw,
+            events=events,
+            event_id={"window": 1},
+            tmin=0,
+            tmax=input_window_seconds,
+            metadata=self.metadata.copy(),
+            baseline=None,
+        )
+        # we skip dropping bads to be consistent with EEGWindowsDataset:
+        epochs._bad_dropped = True
+
+        # Populate new WindowsDataset:
+        windows = WindowsDataset(
+            epochs,
+            description=self.description,
+            transform=self.transform,
+            targets_from=self.targets_from,
+            last_target_only=self.last_target_only,
+        )
+
+        window_kwargs = list(getattr(self, "window_kwargs", []))
+        window_kwargs.append((self.to_epochs_dataset.__name__, {}))
+        setattr(windows, "window_kwargs", window_kwargs)
+        if hasattr(self, "raw_preproc_kwargs"):
+            setattr(
+                windows, "raw_preproc_kwargs", list(getattr(self, "raw_preproc_kwargs"))
+            )
+        if hasattr(self, "window_preproc_kwargs"):
+            setattr(
+                windows,
+                "window_preproc_kwargs",
+                list(getattr(self, "window_preproc_kwargs")),
+            )
+
+        return windows
+
 
 @register_dataset
 class WindowsDataset(RecordDataset):
@@ -617,6 +700,15 @@ class WindowsDataset(RecordDataset):
         last_target_only: bool = True,
     ):
         super().__init__(description, transform)
+        self._fast_disk = self._can_use_fast_get_epoch_from_raw(windows)
+        if not (self._fast_disk or windows.preload):
+            warnings.warn(
+                "The provided mne.Epochs object does not meet the requirements for "
+                "fast epoch access. This may lead to slow data loading from disk. "
+                "Consider preloading the epochs or checking the conditions in "
+                "WindowsDataset._can_use_fast_get_epoch_from_raw.",
+                UserWarning,
+            )
         self.windows = windows
         self.last_target_only = last_target_only
         if targets_from not in ("metadata", "channels"):
@@ -633,6 +725,20 @@ class WindowsDataset(RecordDataset):
         self.raw_preproc_kwargs: list[dict[str, Any]] = []
         self.window_preproc_kwargs: list[dict[str, Any]] = []
 
+    @staticmethod
+    def _can_use_fast_get_epoch_from_raw(epochs: mne.BaseEpochs) -> bool:
+        """Check if we can use the fast _get_epoch_from_raw method,
+        or if we need to use the slow get_data method."""
+        return (
+            not epochs.preload
+            and epochs._bad_dropped
+            and epochs.detrend is None
+            and not epochs._do_baseline  # baseline is None
+            and (epochs._decim == 1)
+            and epochs._offset is None
+            and epochs._projector is None
+        )
+
     def __getitem__(self, index: int):
         """Get a window and its target.
 
@@ -650,7 +756,10 @@ class WindowsDataset(RecordDataset):
         np.ndarray
             Crop indices.
         """
-        X = self.windows.get_data(item=index)[0].astype("float32")
+        if self._fast_disk:
+            X = self.windows._get_epoch_from_raw(index).astype("float32")
+        else:
+            X = self.windows.get_data(item=index)[0].astype("float32")
         if self.transform is not None:
             X = self.transform(X)
         if self.targets_from == "metadata":
@@ -706,7 +815,9 @@ class BaseConcatDataset(ConcatDataset, HubDatasetMixin, Generic[T]):
 
     def __init__(
         self,
-        list_of_ds: list[T | BaseConcatDataset[T]],
+        list_of_ds: (
+            list[T] | list[BaseConcatDataset[T]] | list[T | BaseConcatDataset[T]]
+        ),
         target_transform: Callable | None = None,
         *,
         lazy: bool = False,
@@ -1283,3 +1394,40 @@ class BaseConcatDataset(ConcatDataset, HubDatasetMixin, Generic[T]):
 
     def _repr_html_(self):
         return self._build_repr().to_html()
+
+    def to_epochs_dataset(self) -> BaseConcatDataset[WindowsDataset]:
+        """Converts this :class:`BaseConcatDataset` such that all datasets are :class:`WindowsDataset` with ``mne.Epochs``.
+
+        In Braindecode, the data can either be stored as ``mne.io.Raw`` (in :class:`EEGWindowsDataset`)
+        or as ``mne.Epochs`` (in :class:`WindowsDataset`). This function converts all the underlying datasets to
+        :class:`WindowsDataset` with ``mne.Epochs``.
+        This can be useful for reducing disk space when you want to save a dataset.
+
+        Returns
+        -------
+        BaseConcatDataset[WindowsDataset]
+            A new :class:`BaseConcatDataset` where all datasets are :class:`WindowsDataset` with ``mne.Epochs``.
+
+        Raises
+        ------
+        ValueError
+            If any of the underlying datasets is a :class:`RawDataset` or any other type that is not
+            :class:`EEGWindowsDataset` or :class:`WindowsDataset`, as they cannot be converted to epochs.
+        """
+        datasets = self.datasets
+        if not all(
+            isinstance(ds, (EEGWindowsDataset, WindowsDataset)) for ds in datasets
+        ):
+            raise ValueError(
+                "All datasets must be EEGWindowsDataset or WindowsDataset to convert to WindowsDataset."
+            )
+        new_datasets: list[WindowsDataset] = []
+        for ds in datasets:
+            if isinstance(ds, EEGWindowsDataset):
+                new_ds = ds.to_epochs_dataset()
+                new_datasets.append(new_ds)
+            elif isinstance(ds, WindowsDataset):
+                new_datasets.append(ds)
+        return BaseConcatDataset(
+            new_datasets, target_transform=self.target_transform, lazy=self._lazy
+        )

{braindecode-1.3.2.dev173417708 → braindecode-1.3.2.dev173801889}/braindecode/datasets/bids/hub.py

@@ -27,6 +27,7 @@ The format follows a BIDS-inspired sourcedata structure:
 #
 # License: BSD (3-clause)
 
+import contextlib
 import json
 import logging
 import tempfile
@@ -69,6 +70,8 @@ huggingface_hub = _soft_import(
 
 log = logging.getLogger(__name__)
 
+_LOCK_FILE = "format_info.json"
+
 
 class HubDatasetMixin:
     """
@@ -84,7 +87,6 @@ class HubDatasetMixin:
     >>> dataset = NMT(path=path, preload=True)
     >>> dataset.push_to_hub(
     ...     repo_id="username/nmt-dataset",
-    ...     commit_message="Add NMT dataset"
     ... )
     >>>
     >>> # Load dataset from Hub
@@ -96,13 +98,14 @@ class HubDatasetMixin:
     def push_to_hub(
         self,
         repo_id: str,
-        commit_message: Optional[str] = None,
         private: bool = False,
         token: Optional[str] = None,
-        create_pr: bool = False,
         compression: str = "blosc",
         compression_level: int = 5,
         pipeline_name: str = "braindecode",
+        chunk_size: int = 5_000_000,
+        local_cache_dir: str | Path | None = None,
+        **kwargs,
     ) -> str:
         """
         Upload the dataset to the Hugging Face Hub in BIDS-like Zarr format.
@@ -116,20 +119,41 @@ class HubDatasetMixin:
         ----------
         repo_id : str
             Repository ID on the Hugging Face Hub (e.g., "username/dataset-name").
-        commit_message : str | None
-            Commit message. If None, a default message is generated.
         private : bool, default=False
             Whether to create a private repository.
         token : str | None
             Hugging Face API token. If None, uses cached token.
-        create_pr : bool, default=False
-            Whether to create a Pull Request instead of directly committing.
         compression : str, default="blosc"
             Compression algorithm for Zarr. Options: "blosc", "zstd", "gzip", None.
         compression_level : int, default=5
             Compression level (0-9). Level 5 provides optimal balance.
         pipeline_name : str, default="braindecode"
             Name of the processing pipeline for BIDS sourcedata.
+        chunk_size : int, default=5_000_000
+            Number of samples per chunk in Zarr along the time/window dimension.
+            Larger chunk sizes create fewer but larger chunks/files. This parameter
+            is used for both continuous data (e.g., RawDataset, EEGWindowsDataset)
+            and pre-cut windows (WindowsDataset). For WindowsDataset, multiple
+            windows may be stored in a single chunk depending on their duration
+            and the chosen ``chunk_size``.
+        local_cache_dir : str | Path | None
+            Local directory to use for temporary files during upload. If None, uses
+            the system temp directory and cleans it up after upload. If provided,
+            the directory is used as a persistent cache:
+
+            - If the directory is empty (or does not exist), the cache is built
+              there and a lock file (``format_info.json``) is written once
+              the cache is complete, before the upload starts. The file
+              contains the zarr conversion parameters as JSON.
+            - If the lock file is present and its JSON parameters match the
+              current call, cache creation is skipped and the upload resumes
+              directly (useful for retrying interrupted uploads).
+            - If the lock file is present but its JSON parameters differ from
+              the current call, a ``ValueError`` is raised.
+            - If the directory is non-empty but the lock file is absent, a
+              ``ValueError`` is raised listing the files found.
+        **kwargs
+            Additional arguments passed to huggingface_hub.upload_large_folder().
 
         Returns
         -------
@@ -149,7 +173,6 @@ class HubDatasetMixin:
         >>> # Upload with BIDS-like structure
         >>> url = dataset.push_to_hub(
         ...     repo_id="myusername/nmt-dataset",
-        ...     commit_message="Upload NMT EEG dataset"
         ... )
         """
         if huggingface_hub is False or zarr is False:
@@ -159,13 +182,12 @@ class HubDatasetMixin:
             )
 
         # Create API instance
-        _ = huggingface_hub.HfApi(token=token)
+        hf_api = huggingface_hub.HfApi(token=token)
 
         # Create repository if it doesn't exist
         try:
-            huggingface_hub.create_repo(
+            hf_api.create_repo(
                 repo_id=repo_id,
-                token=token,
                 private=private,
                 repo_type="dataset",
                 exist_ok=True,
@@ -173,74 +195,68 @@ class HubDatasetMixin:
         except Exception as e:
             raise RuntimeError(f"Failed to create repository: {e}")
 
-        # Create a temporary directory for upload
-        with tempfile.TemporaryDirectory() as tmpdir:
-            tmp_path = Path(tmpdir)
-
-            # Create BIDS-like sourcedata structure
-            log.info("Creating BIDS-like sourcedata structure...")
-            bids_layout = hub_format.BIDSSourcedataLayout(
-                tmp_path, pipeline_name=pipeline_name
-            )
-            sourcedata_dir = bids_layout.create_structure()
-
-            # Save dataset_description.json
-            bids_layout.save_dataset_description()
-
-            # Save participants.tsv
-            descriptions = [ds.description for ds in self.datasets]
-            bids_layout.save_participants(descriptions)
-
-            # Save BIDS sidecar files for each recording
-            self._save_bids_sidecar_files(bids_layout)
-
-            # Convert dataset to Zarr format inside sourcedata
-            log.info("Converting dataset to Zarr format...")
-            dataset_path = sourcedata_dir / "dataset.zarr"
-
-            self._convert_to_zarr_inline(
-                dataset_path,
-                compression,
-                compression_level,
-            )
-
-            # Save dataset metadata (README.md)
-            self._save_dataset_card(tmp_path)
-
-            # Save format info
-            format_info_path = tmp_path / "format_info.json"
-            with open(format_info_path, "w", encoding="utf-8") as f:
-                format_info = self._get_format_info_inline()
-                json.dump(
-                    {
-                        "format": "zarr",
-                        "pipeline_name": pipeline_name,
-                        "compression": compression,
-                        "compression_level": compression_level,
-                        "braindecode_version": braindecode.__version__,
-                        **format_info,
-                    },
-                    f,
-                    indent=2,
-                )
+        format_info = self._get_format_info_inline()
+        format_info_lock = {
+            "format": "zarr",
+            "pipeline_name": pipeline_name,
+            "compression": compression,
+            "compression_level": compression_level,
+            "chunk_size": chunk_size,
+            "braindecode_version": braindecode.__version__,
+            **format_info,
+        }
 
-            # Default commit message
-            if commit_message is None:
-                commit_message = (
-                    f"Upload EEG dataset in BIDS-like "
-                    f"Zarr format ({len(self.datasets)} recordings)"
-                )
+        # Determine upload directory and whether to build the cache
+        with contextlib.ExitStack() as stack:
+            if local_cache_dir is None:
+                tmpdir = stack.enter_context(tempfile.TemporaryDirectory())
+                tmp_path = Path(tmpdir)
+                build_cache = True
+            else:
+                tmp_path = Path(local_cache_dir)
+                lock_path = tmp_path / _LOCK_FILE
+                if lock_path.exists():
+                    with open(lock_path, "r", encoding="utf-8") as _f:
+                        _lock_params = json.load(_f)
+                    if _lock_params != format_info_lock:
+                        raise ValueError(
+                            f"Lock file found at '{lock_path}' but its "
+                            f"parameters {_lock_params} differ from the "
+                            f"current call parameters {format_info_lock}. "
+                            "Provide an empty directory or match the "
+                            "original parameters."
+                        )
+                    log.info(
+                        f"Lock file found at '{lock_path}', skipping cache "
+                        "creation and resuming upload."
+                    )
+                    build_cache = False
+                else:
+                    if tmp_path.exists():
+                        existing = list(tmp_path.iterdir())
+                        if existing:
+                            entries = ", ".join(p.name for p in existing)
+                            raise ValueError(
+                                f"local_cache_dir '{tmp_path}' is not empty and "
+                                f"has no lock file. Found: {entries}. Provide an "
+                                "empty directory or one previously prepared by "
+                                "push_to_hub()."
+                            )
+                    else:
+                        tmp_path.mkdir(parents=True)
+                    build_cache = True
+
+            if build_cache:
+                self._build_local_cache(tmp_path, format_info_lock)
 
             # Upload folder to Hub
             log.info(f"Uploading to Hugging Face Hub ({repo_id})...")
             try:
-                url = huggingface_hub.upload_folder(
+                url = hf_api.upload_large_folder(
                     repo_id=repo_id,
                     folder_path=str(tmp_path),
                     repo_type="dataset",
-                    commit_message=commit_message,
-                    token=token,
-                    create_pr=create_pr,
+                    **kwargs,
                 )
                 log.info(f"Dataset uploaded successfully to {repo_id}")
                 log.info(f"URL: https://huggingface.co/datasets/{repo_id}")
@@ -248,6 +264,55 @@ class HubDatasetMixin:
             except Exception as e:
                 raise RuntimeError(f"Failed to upload dataset: {e}")
 
+    def _build_local_cache(
+        self,
+        tmp_path,
+        format_info_lock,
+    ):
+        """Build the local cache directory with the dataset in Zarr format and BIDS-like structure.
+        This folder will be uploaded to the Hub"""
+        compression = format_info_lock["compression"]
+        compression_level = format_info_lock["compression_level"]
+        pipeline_name = format_info_lock["pipeline_name"]
+        chunk_size = format_info_lock["chunk_size"]
+
+        # Create BIDS-like sourcedata structure
+        log.info("Creating BIDS-like sourcedata structure...")
+        bids_layout = hub_format.BIDSSourcedataLayout(
+            tmp_path, pipeline_name=pipeline_name
+        )
+        sourcedata_dir = bids_layout.create_structure()
+
+        # Save dataset_description.json
+        bids_layout.save_dataset_description()
+
+        # Save participants.tsv
+        descriptions = [ds.description for ds in self.datasets]
+        bids_layout.save_participants(descriptions)
+
+        # Save BIDS sidecar files for each recording
+        self._save_bids_sidecar_files(bids_layout)
+
+        # Convert dataset to Zarr format inside sourcedata
+        log.info("Converting dataset to Zarr format...")
+        dataset_path = sourcedata_dir / "dataset.zarr"
+
+        self._convert_to_zarr_inline(
+            dataset_path,
+            compression,
+            compression_level,
+            chunk_size,
+        )
+
+        # Save dataset metadata (README.md)
+        self._save_dataset_card(tmp_path)
+
+        # Save format info
+        # This marks the cache as complete
+        format_info_path = tmp_path / _LOCK_FILE
+        with open(format_info_path, "w", encoding="utf-8") as f:
+            json.dump(format_info_lock, f, indent=2)
+
     def _save_dataset_card(self, path: Path, bids_inspired: bool = True) -> None:
         """Generate and save a dataset card (README.md) with metadata.
 
@@ -465,7 +530,7 @@ class HubDatasetMixin:
             )
 
             # Load format info
-            format_info_path = Path(dataset_dir) / "format_info.json"
+            format_info_path = Path(dataset_dir) / _LOCK_FILE
             if format_info_path.exists():
                 with open(format_info_path, "r") as f:
                     format_info = json.load(f)
@@ -595,6 +660,7 @@ class HubDatasetMixin:
         output_path: Path,
         compression: str,
         compression_level: int,
+        chunk_size: int = 5_000_000,
     ) -> None:
         """Convert dataset to Zarr format (inline implementation)."""
 
@@ -660,7 +726,14 @@ class HubDatasetMixin:
 
                 # Save using inlined function
                 _save_windows_to_zarr(
-                    grp, data, metadata, description, info_dict, compressor, target_name
+                    grp,
+                    data,
+                    metadata,
+                    description,
+                    info_dict,
+                    compressor,
+                    target_name,
+                    chunk_size,
                 )
 
             elif dataset_type == "EEGWindowsDataset":
@@ -682,6 +755,7 @@ class HubDatasetMixin:
                     targets_from,
                     last_target_only,
                     compressor,
+                    chunk_size,
                 )
 
             elif dataset_type == "RawDataset":
@@ -693,7 +767,13 @@ class HubDatasetMixin:
 
                 # Save using inlined function
                 _save_raw_to_zarr(
-                    grp, raw, description, info_dict, target_name, compressor
+                    grp,
+                    raw,
+                    description,
+                    info_dict,
+                    target_name,
+                    compressor,
+                    chunk_size,
                 )
 
     def _get_format_info_inline(self):
@@ -765,11 +845,16 @@ class HubDatasetMixin:
 
                 # Convert to MNE objects and create dataset
                 info = Info.from_json_dict(info_dict)
+                targets = metadata["target"].values
+                if np.issubdtype(targets.dtype, np.integer):
+                    event_ids = targets
+                else:
+                    event_ids = np.ones(len(metadata), dtype=int)
                 events = np.column_stack(
                     [
-                        metadata["i_start_in_trial"].values,
+                        metadata["i_start_in_trial"].values.astype(int),
                         np.zeros(len(metadata), dtype=int),
-                        metadata["target"].values,
+                        event_ids,
                     ]
                 )
                 epochs = mne.EpochsArray(data, info, events=events, metadata=metadata)

{braindecode-1.3.2.dev173417708 → braindecode-1.3.2.dev173801889}/braindecode/datasets/bids/hub_io.py

@@ -45,16 +45,25 @@ def _restore_nan_from_json(obj):
 
 
 def _save_windows_to_zarr(
-    grp, data, metadata, description, info, compressor, target_name
+    grp, data, metadata, description, info, compressor, target_name, chunk_size
 ):
     """Save windowed data to Zarr group (low-level function)."""
     data_array = data.astype(np.float32)
     compressors_list = [compressor] if compressor is not None else None
 
+    max_windows_per_chunk = max(
+        1, chunk_size // (data_array.shape[1] * data_array.shape[2])
+    )
+    n_windows_per_chunk = min(data_array.shape[0], max_windows_per_chunk)
+
     grp.create_array(
         "data",
         data=data_array,
-        chunks=(1, data_array.shape[1], data_array.shape[2]),
+        chunks=(
+            n_windows_per_chunk,
+            data_array.shape[1],
+            data_array.shape[2],
+        ),
         compressors=compressors_list,
     )
 
@@ -70,17 +79,28 @@ def _save_windows_to_zarr(
 
 
 def _save_eegwindows_to_zarr(
-    grp, raw, metadata, description, info, targets_from, last_target_only, compressor
+    grp,
+    raw,
+    metadata,
+    description,
+    info,
+    targets_from,
+    last_target_only,
+    compressor,
+    chunk_size,
 ):
     """Save EEG continuous raw data to Zarr group (low-level function)."""
     continuous_data = raw.get_data()
     continuous_float = continuous_data.astype(np.float32)
     compressors_list = [compressor] if compressor is not None else None
 
+    max_samples_per_chunk = max(1, chunk_size // continuous_float.shape[0])
+    n_samples_per_chunk = min(continuous_float.shape[1], max_samples_per_chunk)
+
     grp.create_array(
         "data",
         data=continuous_float,
-        chunks=(continuous_float.shape[0], min(10000, continuous_float.shape[1])),
+        chunks=(continuous_float.shape[0], n_samples_per_chunk),
         compressors=compressors_list,
     )
 
@@ -143,16 +163,19 @@ def _load_eegwindows_from_zarr(grp, preload):
     return data, metadata, description, info_dict, targets_from, last_target_only
 
 
-def _save_raw_to_zarr(grp, raw, description, info, target_name, compressor):
+def _save_raw_to_zarr(grp, raw, description, info, target_name, compressor, chunk_size):
     """Save RawDataset continuous raw data to Zarr group (low-level function)."""
     continuous_data = raw.get_data()
     continuous_float = continuous_data.astype(np.float32)
     compressors_list = [compressor] if compressor is not None else None
 
+    max_samples_per_chunk = max(1, chunk_size // continuous_float.shape[0])
+    n_samples_per_chunk = min(continuous_float.shape[1], max_samples_per_chunk)
+
     grp.create_array(
         "data",
         data=continuous_float,
-        chunks=(continuous_float.shape[0], min(10000, continuous_float.shape[1])),
+        chunks=(continuous_float.shape[0], n_samples_per_chunk),
         compressors=compressors_list,
     )