eegdash 0.3.7.dev183881899__tar.gz → 0.3.8__tar.gz

{eegdash-0.3.7.dev183881899 → eegdash-0.3.8}/MANIFEST.in

@@ -1,7 +1,7 @@
 include README.md
 include LICENSE
 
-include eegdash/dataset_summary.csv
+include eegdash/dataset/dataset_summary.csv
 
 recursive-include docs *.ipynb *.rst conf.py Makefile
 recursive-exclude docs *checkpoint.ipynb

{eegdash-0.3.7.dev183881899/eegdash.egg-info → eegdash-0.3.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: eegdash
-Version: 0.3.7.dev183881899
+Version: 0.3.8
 Summary: EEG data for machine learning
 Author-email: Young Truong <dt.young112@gmail.com>, Arnaud Delorme <adelorme@gmail.com>, Aviv Dotan <avivd220@gmail.com>, Oren Shriki <oren70@gmail.com>, Bruno Aristimunha <b.aristimunha@gmail.com>
 License-Expression: GPL-3.0-only
@@ -40,6 +40,7 @@ Requires-Dist: h5io>=0.2.4
 Requires-Dist: pymatreader
 Requires-Dist: eeglabio
 Requires-Dist: tabulate
+Requires-Dist: docstring_inheritance
 Provides-Extra: tests
 Requires-Dist: pytest; extra == "tests"
 Requires-Dist: pytest-cov; extra == "tests"
@@ -106,7 +107,7 @@ EEGDash queries return a **Pytorch Dataset** formatted to facilitate machine lea
 
 ## Data preprocessing
 
-EEGDash datasets are processed using the popular [BrainDecode](https://braindecode.org/stable/index.html) library. In fact, EEGDash datasets are BrainDecode datasets, which are themselves PyTorch datasets. This means that any preprocessing possible on BrainDecode datasets is also possible on EEGDash datasets. Refer to [BrainDecode](https://braindecode.org/stable/index.html) tutorials for guidance on preprocessing EEG data.
+EEGDash datasets are processed using the popular [braindecode](https://braindecode.org/stable/index.html) library. In fact, EEGDash datasets are braindecode datasets, which are themselves PyTorch datasets. This means that any preprocessing possible on braindecode datasets is also possible on EEGDash datasets. Refer to [braindecode](https://braindecode.org/stable/index.html) tutorials for guidance on preprocessing EEG data.
 
 ## EEG-Dash usage
 
@@ -129,7 +130,7 @@ ds_NDARDB033FW5 = EEGDashDataset(
 )
 ```
 
-This will search and download the metadata for the task **RestingState** for subject **NDARDB033FW5** in BIDS dataset **ds005514**. The actual data will not be downloaded at this stage. Following standard practice, data is only downloaded once it is processed. The **ds_NDARDB033FW5** object is a fully functional BrainDecode dataset, which is itself a PyTorch dataset. This [tutorial](https://github.com/sccn/EEGDash/blob/develop/notebooks/tutorial_eoec.ipynb) shows how to preprocess the EEG data, extracting portions of the data containing eyes-open and eyes-closed segments, then perform eyes-open vs. eyes-closed classification using a (shallow) deep-learning model. 
+This will search and download the metadata for the task **RestingState** for subject **NDARDB033FW5** in BIDS dataset **ds005514**. The actual data will not be downloaded at this stage. Following standard practice, data is only downloaded once it is processed. The **ds_NDARDB033FW5** object is a fully functional braindecode dataset, which is itself a PyTorch dataset. This [tutorial](https://github.com/sccn/EEGDash/blob/develop/notebooks/tutorial_eoec.ipynb) shows how to preprocess the EEG data, extracting portions of the data containing eyes-open and eyes-closed segments, then perform eyes-open vs. eyes-closed classification using a (shallow) deep-learning model. 
 
 To use the data from multiple subjects, enter:
 
@@ -145,7 +146,13 @@ This will search and download the metadata for the task 'RestingState' for all s
 
 ### Automatic caching
 
-EEGDash automatically caches the downloaded data in the .eegdash_cache folder of the current directory from which the script is called. This means that if you run the tutorial [scripts](https://github.com/sccn/EEGDash/tree/develop/notebooks), the data will only be downloaded the first time the script is executed.
+By default, EEGDash caches downloaded data under a single, consistent folder:
+
+- If ``EEGDASH_CACHE_DIR`` is set in your environment, that path is used.
+- Else, if MNE’s ``MNE_DATA`` config is set, that path is used to align with other EEG tooling.
+- Otherwise, ``.eegdash_cache`` in the current working directory is used.
+
+This means that if you run the tutorial [scripts](https://github.com/sccn/EEGDash/tree/develop/notebooks), the data will only be downloaded the first time the script is executed and reused thereafter.
 
 ## Education -- Coming soon...
 
@@ -159,4 +166,3 @@ EEG-DaSh is a collaborative initiative between the United States and Israel, sup
 
 
 
-

{eegdash-0.3.7.dev183881899 → eegdash-0.3.8}/README.md

@@ -36,7 +36,7 @@ EEGDash queries return a **Pytorch Dataset** formatted to facilitate machine lea
 
 ## Data preprocessing
 
-EEGDash datasets are processed using the popular [BrainDecode](https://braindecode.org/stable/index.html) library. In fact, EEGDash datasets are BrainDecode datasets, which are themselves PyTorch datasets. This means that any preprocessing possible on BrainDecode datasets is also possible on EEGDash datasets. Refer to [BrainDecode](https://braindecode.org/stable/index.html) tutorials for guidance on preprocessing EEG data.
+EEGDash datasets are processed using the popular [braindecode](https://braindecode.org/stable/index.html) library. In fact, EEGDash datasets are braindecode datasets, which are themselves PyTorch datasets. This means that any preprocessing possible on braindecode datasets is also possible on EEGDash datasets. Refer to [braindecode](https://braindecode.org/stable/index.html) tutorials for guidance on preprocessing EEG data.
 
 ## EEG-Dash usage
 
@@ -59,7 +59,7 @@ ds_NDARDB033FW5 = EEGDashDataset(
 )
 ```
 
-This will search and download the metadata for the task **RestingState** for subject **NDARDB033FW5** in BIDS dataset **ds005514**. The actual data will not be downloaded at this stage. Following standard practice, data is only downloaded once it is processed. The **ds_NDARDB033FW5** object is a fully functional BrainDecode dataset, which is itself a PyTorch dataset. This [tutorial](https://github.com/sccn/EEGDash/blob/develop/notebooks/tutorial_eoec.ipynb) shows how to preprocess the EEG data, extracting portions of the data containing eyes-open and eyes-closed segments, then perform eyes-open vs. eyes-closed classification using a (shallow) deep-learning model. 
+This will search and download the metadata for the task **RestingState** for subject **NDARDB033FW5** in BIDS dataset **ds005514**. The actual data will not be downloaded at this stage. Following standard practice, data is only downloaded once it is processed. The **ds_NDARDB033FW5** object is a fully functional braindecode dataset, which is itself a PyTorch dataset. This [tutorial](https://github.com/sccn/EEGDash/blob/develop/notebooks/tutorial_eoec.ipynb) shows how to preprocess the EEG data, extracting portions of the data containing eyes-open and eyes-closed segments, then perform eyes-open vs. eyes-closed classification using a (shallow) deep-learning model. 
 
 To use the data from multiple subjects, enter:
 
@@ -75,7 +75,13 @@ This will search and download the metadata for the task 'RestingState' for all s
 
 ### Automatic caching
 
-EEGDash automatically caches the downloaded data in the .eegdash_cache folder of the current directory from which the script is called. This means that if you run the tutorial [scripts](https://github.com/sccn/EEGDash/tree/develop/notebooks), the data will only be downloaded the first time the script is executed.
+By default, EEGDash caches downloaded data under a single, consistent folder:
+
+- If ``EEGDASH_CACHE_DIR`` is set in your environment, that path is used.
+- Else, if MNE’s ``MNE_DATA`` config is set, that path is used to align with other EEG tooling.
+- Otherwise, ``.eegdash_cache`` in the current working directory is used.
+
+This means that if you run the tutorial [scripts](https://github.com/sccn/EEGDash/tree/develop/notebooks), the data will only be downloaded the first time the script is executed and reused thereafter.
 
 ## Education -- Coming soon...
 
@@ -89,4 +95,3 @@ EEG-DaSh is a collaborative initiative between the United States and Israel, sup
 
 
 
-

{eegdash-0.3.7.dev183881899 → eegdash-0.3.8}/docs/Makefile

@@ -11,7 +11,7 @@ help:
 
 .PHONY: apidoc
 apidoc:
-	@sphinx-apidoc -o "$(APIDIR)" "../$(PKG)" -f -e -M
+	@python -m sphinx.ext.apidoc -o "$(APIDIR)" "../$(PKG)" -f -e -M
 
 # Standard build runs examples
 html: apidoc
@@ -28,4 +28,4 @@ Makefile: ;
 
 %: Makefile
 	@python prepare_summary_tables.py ../eegdash/ $(BUILDDIR)
-	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
+	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)

eegdash-0.3.8/docs/build/html/_downloads/06c8d94b7e0b8be2de39fdc122dd12bb/tutorial_challenge_2.ipynb

@@ -0,0 +1,388 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "# For tips on running notebooks in Google Colab:\n# `pip install eegdash`\n%matplotlib inline"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "\n# Challenge 2: Predicting the p-factor from EEG\n   :depth: 2\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "<img src=\"https://colab.research.google.com/assets/colab-badge.svg\" target=\"https://colab.research.google.com/github/eeg2025/startkit/blob/main/challenge_2.ipynb\" alt=\"Open In Colab\">\n\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Overview\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "# The psychopathology factor (P-factor) is a widely recognized construct in mental health research, representing a common underlying dimension of psychopathology across various disorders.\n# Currently, the P-factor is often assessed using self-report questionnaires or clinician ratings, which can be subjective, prone to bias, and time-consuming.\n# **The Challenge 2** consists of developing a model to predict the P-factor from EEG recordings.\n\n# The challenge encourages learning physiologically meaningful signal representations and discovery of reproducible biomarkers.\n# Models of any size should emphasize robust, interpretable features that generalize across subjects,\n# sessions, and acquisition sites.\n\n# Unlike a standard in-distribution classification task, this regression problem stresses out-of-distribution robustness\n# and extrapolation. The goal is not only to minimize error on seen subjects, but also to transfer effectively to unseen data.\n\n# Ensure the dataset is available locally. If not, see the\n# `dataset download guide <https://eeg2025.github.io/data/#downloading-the-data>`__."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Contents of this start kit\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "# 0. Understand the P-factor regression task.\n# 1. Loading the data.\n# 2. Wrap the data into a PyTorch-compatible dataset.\n# 3. Define, train and save a model.\n\n# .. note:: If you need additional explanations on the\n#    `EEGChallengeDataset <https://eeglab.org/EEGDash/api/eegdash.html#eegdash.EEGChallengeDataset>`__\n#    class, dataloading, `braindecode <https://braindecode.org/stable/models/models_table.html>`__'s\n#    deep learning models, or brain decoding in general, please refer to the\n#    start-kit of challenge 1 which delves deeper into these topics.\n\n# More contents will be released during the competition inside the [`eegdash`](https://eeglab.org/EEGDash/overview.html) [examples webpage](https://eeglab.org/EEGDash/generated/auto_examples/index.html).\n\n# We will establish direct communication channels with you over the next two months, resolve platform issues, and discuss other details.\n\n# We are going to do weekly one-hour support time to help teams that may have difficulties with any technical aspects that may arise.\n\n# We will have closer communication starting next week, 08/09/2025.\n\n# Before we begin, I just want to make a deal with you, ok?\n\n# This is a community competition with a strong open-source foundation. When I say open-source, I mean volunteer work.\n\n# So, if you see something that does not work or could be improved, first, **please be kind**, and we will fix it together on GitHub, okay?\n\n# The entire decoding community will only go further when we stop solving the same problems over and over again, and start working together."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "::\n\n    ----\n    | The tutorial assumes some prior knowledge of deep learning concepts and EEG concepts:|\n    | --- |\n    |* Standard neural network architectures, e.g., convolutional neural networks|\n    |* Optimization by batch gradient descent and backpropagation|\n    |* Overfitting, early stopping, regularisation |\n    |* Some knowledge of pytorch |\n    |* That you know what EEG is |\n    |* That you have basic familiarity with EEG preprocessing |\n    |* Like and support open-source :) |\n    ----\n\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "Install dependencies on colab\n -----------------------------\n tags:[\"hide-input\"]\n!pip install braindecode\n!pip install eegdash\n\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Imports\ntags:[\"hide-input\"]\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "from pathlib import Path\nimport math\nimport os\nimport random\nfrom joblib import Parallel, delayed\n\nimport torch\nfrom torch.utils.data import DataLoader\nfrom torch import optim\nfrom torch.nn.functional import l1_loss\nfrom braindecode.preprocessing import create_fixed_length_windows\nfrom braindecode.datasets.base import EEGWindowsDataset, BaseConcatDataset, BaseDataset\nfrom braindecode.models import EEGNeX\nfrom eegdash import EEGChallengeDataset"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "<div class=\"alert alert-danger\"><h4>Warning</h4><p>In case of Colab, before starting, make sure you're on a GPU instance\n   for faster training! If running on Google Colab, please request a GPU runtime\n   by clicking `Runtime/Change runtime type` in the top bar menu, then selecting\n   'T4 GPU' under 'Hardware accelerator'.</p></div>\n\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Identify whether a CUDA-enabled GPU is available\ntags:[\"hide-input\"]\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nif device == \"cuda\":\n    msg = \"CUDA-enabled GPU found. Training should be faster.\"\nelse:\n    msg = (\n        \"No GPU found. Training will be carried out on CPU, which might be \"\n        \"slower.\\n\\nIf running on Google Colab, you can request a GPU runtime by\"\n        \" clicking\\n`Runtime/Change runtime type` in the top bar menu, then \"\n        \"selecting 'T4 GPU'\\nunder 'Hardware accelerator'.\"\n    )\nprint(msg)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 0. Understanding the P-factor regression task.\n\nThe psychopathology factor (P-factor) is a widely recognized construct in mental health research, representing a common underlying dimension of psychopathology across various disorders.\nThe P-factor is thought to reflect the shared variance among different psychiatric conditions, suggesting that individuals with higher P-factor scores may be more vulnerable to a range of mental health issues.\nCurrently, the P-factor is often assessed using self-report questionnaires or clinician ratings, which can be subjective, prone to bias, and time-consuming.\nIn the dataset of this challenge, the P-factor was assessed using the Child Behavior Checklist (CBCL) [McElroy et al., (2017)](https://doi.org/10.1111/jcpp.12849).\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "# The goal of Challenge 2 is to develop a model to predict the P-factor from EEG recordings.\n# **The feasibility of using EEG data for this purpose is still an open question**.\n# The solution may involve finding meaningful representations of the EEG data that correlate with the P-factor scores.\n# The challenge encourages learning physiologically meaningful signal representations and discovery of reproducible biomarkers.\n# If contestants are successful in this task, it could pave the way for more objective and efficient assessments of the P-factor in clinical settings."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 1. Loading the data\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "# A. Define local path and (down)load the data\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\n# In this challenge 2 example, we load the EEG 2025 release using EEG Dash.\n\n# > **Note:** in this example notebook, we load the contrast change detection task from one mini release only as an example. Naturally, you are encouraged to train your models on all complete releases, using data from all the tasks you deem relevant."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "The first step is to define the cache folder!\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "DATA_DIR = Path(\"~/eegdash/eeg2025_competition\")\n\n# Creating the path if it does not exist\nDATA_DIR.mkdir(parents=True, exist_ok=True)\n\n# We define the list of releases to load.\n# Here, only release 5 is loaded.\nrelease_list = [\"R5\"]\n\nall_datasets_list = [\n    EEGChallengeDataset(\n        release=release,\n        task=\"contrastChangeDetection\",\n        mini=True,\n        description_fields=[\n            \"subject\",\n            \"session\",\n            \"run\",\n            \"task\",\n            \"age\",\n            \"gender\",\n            \"sex\",\n            \"p_factor\",\n        ],\n        cache_dir=DATA_DIR,\n    )\n    for release in release_list\n]\nprint(\"Datasets loaded\")\nsub_rm = [\"NDARWV769JM7\"]"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "### B. Combine the datasets into a single one\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "# Here, we combine the datasets from the different releases into a single\n# ``BaseConcatDataset`` object."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "all_datasets = BaseConcatDataset(all_datasets_list)\nprint(all_datasets.description)\n\nraws = Parallel(n_jobs=os.cpu_count())(\n    delayed(lambda d: d.raw)(d) for d in all_datasets.datasets\n)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "### C. Inspect your data\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "# We can check what is inside the dataset consuming the\n# MNE-object inside the Braindecode dataset.\n#\n# The following snippet, if uncommented, will show the first 10 seconds of the raw EEG signal.\n# We can also inspect the data further by looking at the events and annotations.\n# We strongly recommend you to take a look into the details and check how the events are structured."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "raw = all_datasets.datasets[0].raw  # mne.io.Raw object\nprint(raw.info)\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "# raw.plot(duration=10, scalings=\"auto\", show=True)\n\n# print(raw.annotations)\n\nSFREQ = 100"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 2. Wrap the data into a PyTorch-compatible dataset\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "# The class below defines a dataset wrapper that will extract 2-second windows,\n# uniformly sampled over the whole signal. In addition, it will add useful information\n# about the extracted windows, such as the p-factor, the subject or the task."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "class DatasetWrapper(BaseDataset):\n    def __init__(self, dataset: EEGWindowsDataset, crop_size_samples: int, seed=None):\n        self.dataset = dataset\n        self.crop_size_samples = crop_size_samples\n        self.rng = random.Random(seed)\n\n    def __len__(self):\n        return len(self.dataset)\n\n    def __getitem__(self, index):\n        X, _, crop_inds = self.dataset[index]\n\n        # P-factor label:\n        p_factor = self.dataset.description[\"p_factor\"]\n        p_factor = float(p_factor)\n\n        # Additional information:\n        infos = {\n            \"subject\": self.dataset.description[\"subject\"],\n            \"sex\": self.dataset.description[\"sex\"],\n            \"age\": float(self.dataset.description[\"age\"]),\n            \"task\": self.dataset.description[\"task\"],\n            \"session\": self.dataset.description.get(\"session\", None) or \"\",\n            \"run\": self.dataset.description.get(\"run\", None) or \"\",\n        }\n\n        # Randomly crop the signal to the desired length:\n        i_window_in_trial, i_start, i_stop = crop_inds\n        assert i_stop - i_start >= self.crop_size_samples, f\"{i_stop=} {i_start=}\"\n        start_offset = self.rng.randint(0, i_stop - i_start - self.crop_size_samples)\n        i_start = i_start + start_offset\n        i_stop = i_start + self.crop_size_samples\n        X = X[:, start_offset : start_offset + self.crop_size_samples]\n\n        return X, p_factor, (i_window_in_trial, i_start, i_stop), infos"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "We filter out certain recordings, create fixed length windows and finally make use of our `DatasetWrapper`.\n\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "Filter out recordings that are too short\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "all_datasets = BaseConcatDataset(\n    [\n        ds\n        for ds in all_datasets.datasets\n        if ds.description.subject not in sub_rm\n        and ds.raw.n_times >= 4 * SFREQ\n        and len(ds.raw.ch_names) == 129\n        and not math.isnan(ds.description[\"p_factor\"])\n    ]\n)\n\n# Create 4-seconds windows with 2-seconds stride\nwindows_ds = create_fixed_length_windows(\n    all_datasets,\n    window_size_samples=4 * SFREQ,\n    window_stride_samples=2 * SFREQ,\n    drop_last_window=True,\n)\n\n# Wrap each sub-dataset in the windows_ds\nwindows_ds = BaseConcatDataset(\n    [DatasetWrapper(ds, crop_size_samples=2 * SFREQ) for ds in windows_ds.datasets]\n)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 3. Define, train and save a model\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "# Now we have our pytorch dataset necessary for the training!\n#\n# Below, we define a simple EEGNeX model from Braindecode.\n# All the braindecode models expect the input to be of shape (batch_size, n_channels, n_times)\n# and have a test coverage about the behavior of the model.\n# However, you can use any pytorch model you want."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "Initialize model\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "model = EEGNeX(n_chans=129, n_outputs=1, n_times=2 * SFREQ).to(device)\n\n# Specify optimizer\noptimizer = optim.Adamax(params=model.parameters(), lr=0.002)\n\nprint(model)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "Finally, we can train our model. Here we define a simple training loop using pure PyTorch.\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "# > In this example, we only train for a single epoch. Feel free to increase the number of epochs."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "Create PyTorch Dataloader\n\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "collapsed": false
+      },
+      "outputs": [],
+      "source": [
+        "num_workers = (\n    0  # Set num_workers to 0 to avoid multiprocessing issues in notebooks/tutorials.\n)\ndataloader = DataLoader(\n    windows_ds, batch_size=128, shuffle=True, num_workers=num_workers\n)\n\nn_epochs = 1\n\n# Train model for 1 epoch\nfor epoch in range(n_epochs):\n    for idx, batch in enumerate(dataloader):\n        # Reset gradients\n        optimizer.zero_grad()\n\n        # Unpack the batch\n        X, y, crop_inds, infos = batch\n        X = X.to(dtype=torch.float32, device=device)\n        y = y.to(dtype=torch.float32, device=device).unsqueeze(1)\n\n        # Forward pass\n        y_pred = model(X)\n\n        # Compute loss\n        loss = l1_loss(y_pred, y)\n        print(f\"Epoch {0} - step {idx}, loss: {loss.item()}\")\n\n        # Gradient backpropagation\n        loss.backward()\n        optimizer.step()\n\n# Finally, we can save the model for later use\ntorch.save(model.state_dict(), \"weights_challenge_2.pt\")\nprint(\"Model saved as 'weights_challenge_2.pt'\")"
+      ]
+    }
+  ],
+  "metadata": {
+    "kernelspec": {
+      "display_name": "Python 3",
+      "language": "python",
+      "name": "python3"
+    },
+    "language_info": {
+      "codemirror_mode": {
+        "name": "ipython",
+        "version": 3
+      },
+      "file_extension": ".py",
+      "mimetype": "text/x-python",
+      "name": "python",
+      "nbconvert_exporter": "python",
+      "pygments_lexer": "ipython3",
+      "version": "3.11.0"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}