From 54f310bc70c597eab013f652f14d29fb72863cde Mon Sep 17 00:00:00 2001
From: Javi <javi@galatolo.xyz>
Date: Thu, 17 Jul 2025 14:32:14 +0000
Subject: [PATCH 1/6] initial mssv commit

---
 physioex/preprocess/mssv.py | 379 ++++++++++++++++++++++++++++++++++++
 1 file changed, 379 insertions(+)
 create mode 100644 physioex/preprocess/mssv.py

diff --git a/physioex/preprocess/mssv.py b/physioex/preprocess/mssv.py
new file mode 100644
index 00000000..2e6630b3
--- /dev/null
+++ b/physioex/preprocess/mssv.py
@@ -0,0 +1,379 @@
+import os
+from pathlib import Path
+from typing import List, Tuple
+import random
+
+import numpy as np
+import pandas as pd
+from loguru import logger
+from tqdm import tqdm
+import datalad.api as dl
+import openneuro as on
+from scipy.signal import filtfilt, firwin, resample
+
+from physioex.preprocess.preprocessor import Preprocessor
+from physioex.preprocess.utils.signal import xsleepnet_preprocessing_mouse
+
+from physioex.preprocess.utils.mousedata import get_channels, read_channel_signal
+
+
+
+def process_recording(edf_path, tsv_path):
+
+    fs = 100
+    epoch_second = 4
+
+    available_channels = get_channels(edf_path)
+    
+    try:
+        stages = pd.read_csv(tsv_path, sep='\t')['stage'].values
+        stages = stages - 1
+    except Exception as e:
+        print(f"Error reading file: {tsv_path}")
+        print(f"skipping subject")
+        return None, None
+    
+    eeg_candidates = [ch for ch in available_channels if 'EEG' in ch.upper()]
+    eeg_channel = random.choice(eeg_candidates) if eeg_candidates else None
+
+    eeg1, old_fs = read_channel_signal(edf_path, eeg_channel)
+
+    # Parametri
+    Nfir = 100
+
+    # Creazione del filtro FIR bandpass
+    b_band = firwin(Nfir + 1, [0.3, 40], pass_zero=False, fs=old_fs)
+
+    # Applicazione del filtro al segnale EEG
+    eeg1 = filtfilt(b_band, 1, eeg1)
+
+    if fs != old_fs:
+        eeg1 = resample(eeg1, int(len(eeg1) * fs / old_fs))
+     
+    eeg2 = eeg1.copy() # only working with one EEG channel for now, but filling eeg2 for shape coherence with other datasets
+    
+    emg_candidates = [ch for ch in available_channels if 'EMG' in ch.upper()]
+    emg_channel = random.choice(eeg_candidates) if emg_candidates else None
+    if emg_channel is None:
+        print(f"Error: no EMG channel found in {edf_path}")
+        print(f"Available channels: {available_channels}")
+        return None, None
+    else:
+        emg, old_fs = read_channel_signal(edf_path, emg_channel)
+
+    # filtering and resampling
+    b_band = firwin(Nfir + 1, 10, pass_zero=False, fs=old_fs)
+    emg = filtfilt(b_band, 1, emg)
+
+    if fs != old_fs:
+        emg = resample(emg, int(len(emg) * fs / old_fs))
+
+    expected_epochs = len(eeg1) // (epoch_second * fs)
+
+    # checking coherence of the signals with stages
+    if expected_epochs > len(stages):
+        expected_epochs = len(stages)
+    else:
+        stages = stages[:expected_epochs]
+    total_samples = expected_epochs * epoch_second * fs
+    eeg1 = eeg1[:total_samples]
+    eeg2 = eeg2[:total_samples]
+    emg = emg[:total_samples]
+    stages = np.array(stages)
+    # print stages distribution
+    # print(f'Stages distribution: {np.bincount(stages)}')
+
+    # buffer the signals into epochs
+    signal = np.array([eeg1, eeg2, emg])
+    signal = np.transpose(signal).reshape(expected_epochs, epoch_second * fs, 3)
+
+    # find the epochs associated with stages < 0 or > 2
+    invalid_epochs = np.where(np.logical_or(stages < 0, stages > 2))[0]
+
+    # remove the invalid epochs
+    stages = np.delete(stages, invalid_epochs)
+    signal = np.delete(signal, invalid_epochs, axis=0)
+    
+    signal = np.transpose(signal, (0, 2, 1))
+
+    return signal.astype(np.float32), stages.astype(int)
+
+
+class RecordingsIterator:
+    '''
+    So different recordings from same subject are saved separately.
+    
+    '''
+    def __init__(self, data):
+        if isinstance(data, pd.DataFrame):
+            self._iter = data.iterrows()  # (index, row)
+            self.len = len(data)
+        else:
+            raise TypeError("Unsupported data type")
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        _, row = next(self._iter)
+        return row
+    
+    def __len__(self):
+        return self.len
+
+
+class MSSVPreprocessor(Preprocessor):
+
+    def __init__(
+        self,
+        preprocessors_name: List[str] = ["xsleepnet_mouse"],
+        preprocessors=[xsleepnet_preprocessing_mouse],
+        preprocessor_shape=[[3, 17, 129]],
+        data_folder: str = None,
+    ):
+
+        super().__init__(
+            dataset_name="mssv",
+            signal_shape=[3, 400],
+            preprocessors_name=preprocessors_name,
+            preprocessors=preprocessors,
+            preprocessors_shape=preprocessor_shape,
+            data_folder=data_folder,
+        )
+        
+        self.source_dataset = os.path.join(self.dataset_folder, 'mssv_openneuro')
+        
+        self.split_subjects_table = None
+        self._iterator = None
+        
+    @logger.catch
+    def download_dataset(self) -> None:
+        """
+        Downloads the dataset if it is not already present on disk.
+
+        """
+        # pass
+
+        # if not os.listdir(self.source_dataset):
+        #     raise NotImplementedError(
+        #         "❌ Automatic download of MSSV is not supported yet. "
+        #         f"Please download the dataset manually from https://openneuro.org/datasets/ds006366/versions/1.0.0/download# in the directory {self.source_dataset}"
+        #     )
+
+        if not os.path.exists(self.source_dataset):
+
+            os.makedirs(self.source_dataset)
+
+        print(
+            "The openneuro downloader is very unstable and might crash before finishing. "
+            "Rerun the script as many times as needed to resume the download until "
+            "the download finishes and the actual preprocessing starts."
+        )
+        print("")
+
+        on.download(dataset='ds006366', target_dir=self.source_dataset)
+
+
+    @logger.catch
+    def get_subjects_records(self) -> np.ndarray:
+        """
+        Finds all .edf files in the data folder and extracts the subject ID from the file name.
+
+        Returns:
+            np.ndarray: An array of unique subject IDs.
+        """
+        
+        if self.split_subjects_table is None:
+
+            participants_path = os.path.join(self.source_dataset, 'participants.tsv')
+            self.participants = pd.read_csv(participants_path, sep='\t')
+            
+            split_subjects_rows = []    
+                
+            for _, row in self.participants.iterrows():
+                subj_folder = os.path.join(self.source_dataset, row['participant_id'])
+                
+                eeg_files = [f for f in os.listdir(subj_folder + '/eeg') if f.endswith('eeg.edf')]
+                
+                for f in eeg_files:
+                    run = f.split('_')[-2]
+                    
+                    split_subjects_rows.append({
+                        'edf_path': f,
+                        'real_subject': row['participant_id'],
+                        'run': run,
+                        'lab': row['lab']
+                    })
+                    
+            self.split_subjects_table = pd.DataFrame(split_subjects_rows)
+        
+        self._iterator = RecordingsIterator(self.split_subjects_table)
+                  
+        return self._iterator
+
+    
+    @logger.catch
+    def read_subject_record(self, record: pd.Series) -> Tuple[np.array, np.array]:
+        """
+        Reads all recordings belonging to 'record', processes and concatenates them.
+
+        Args:
+            record (pd.Series): The row representing the subject's recording.
+
+        Returns:
+            Tuple[np.array, np.array]: A tuple containing the signal and labels with shapes
+            [n_windows, n_channels, n_timestamps] and [n_windows], respectively. If the record
+            should be skipped, the function should return None, None.
+        """
+                
+        subject_folder = os.path.join(self.source_dataset, record['real_subject'])
+        
+        edf_path = os.path.join(subject_folder, 'eeg', record['edf_path'])
+        tsv_path = os.path.join(subject_folder, 'eeg', record['edf_path'].replace('_eeg.edf', '_events.tsv'))
+        
+        signal, stages = process_recording(edf_path, tsv_path)
+                
+        return signal, stages
+    
+    def customize_table(self, table) -> pd.DataFrame:
+        """
+        Customizes the dataset table before saving it.
+
+
+        Parameters:
+            table (pd.DataFrame): The dataset table to be customized.
+
+        Returns:
+            pd.DataFrame: The customized dataset table.
+        """
+        
+        return table.join(self.split_subjects_table[['real_subject', 'run', 'lab']])
+    
+    
+    def get_sets(self, k=4) -> Tuple[List[np.array], List[np.array], List[np.array]]:
+        """
+        Performs K-Fold splitting using a greedy allocation strategy,
+        stratified by 'lab' to ensure balanced lab distribution in each fold.
+        
+        The greedy strategy also assigns each subject to the set that has the lowest proportion filled 
+        relative to its target allocation ratio. This is done because some of the mice have much
+        more epochs than others. This ensures a correct distribution of sleep epochs according to
+        the predefined ratios, while keeping mice segregated.
+
+        Args:
+            k (int): Number of folds.
+
+        Returns:
+            Tuple[List[np.array], List[np.array], List[np.array]]:
+            Lists of train, validation, and test sets for each fold.
+        """
+
+        # 1. Aggregate per subject
+        subject_groups = self.table.groupby('real_subject')
+        subject_ids = np.array(list(subject_groups.groups.keys()))
+        subject_durations = subject_groups['num_windows'].sum().values
+        subject_labs = subject_groups['lab'].first().values  # assumes consistent lab per subject
+
+        total_duration = subject_durations.sum()
+        train_ratio = 0.7
+        val_ratio = 0.15
+        test_ratio = 1 - train_ratio - val_ratio
+
+        used_test_subjects = set()
+
+        all_train_folds, all_val_folds, all_test_folds = [], [], []
+
+        np.random.seed(42)
+        unique_labs = np.unique(subject_labs)
+
+        for fold in range(k):
+            train_subjects, val_subjects, test_subjects = [], [], []
+
+            for lab in unique_labs:
+                lab_mask = subject_labs == lab
+                lab_subject_ids = subject_ids[lab_mask]
+                lab_durations = subject_durations[lab_mask]
+
+                # Shuffle lab-specific subjects
+                perm = np.random.permutation(len(lab_subject_ids))
+                lab_subject_ids = lab_subject_ids[perm]
+                lab_durations = lab_durations[perm]
+
+                lab_total_duration = lab_durations.sum()
+                lab_test_dur, lab_val_dur = 0, 0
+                lab_test, lab_val, lab_train = [], [], []
+
+                for subj, dur in zip(lab_subject_ids, lab_durations):
+                    if subj in used_test_subjects:
+                        continue
+                    if lab_test_dur + dur <= test_ratio * lab_total_duration:
+                        lab_test.append(subj)
+                        lab_test_dur += dur
+                        used_test_subjects.add(subj)
+
+                remaining = [s for s in lab_subject_ids if s not in lab_test]
+
+                for subj in remaining:
+                    idx = np.where(subject_ids == subj)[0][0]
+                    dur = subject_durations[idx]
+                    if lab_val_dur + dur <= val_ratio * lab_total_duration:
+                        lab_val.append(subj)
+                        lab_val_dur += dur
+
+                lab_train = [s for s in remaining if s not in lab_val]
+
+                train_subjects.extend(lab_train)
+                val_subjects.extend(lab_val)
+                test_subjects.extend(lab_test)
+
+            # Convert to arrays
+            train_subjects = np.array(train_subjects)
+            val_subjects = np.array(val_subjects)
+            test_subjects = np.array(test_subjects)
+
+            # Map back to full table: get all recordings for each subject
+            train = self.table[self.table['real_subject'].isin(train_subjects)]
+            val = self.table[self.table['real_subject'].isin(val_subjects)]
+            test = self.table[self.table['real_subject'].isin(test_subjects)]
+
+            # Compute durations
+            train_dur = train['num_windows'].sum()
+            val_dur = val['num_windows'].sum()
+            test_dur = test['num_windows'].sum()
+
+            train_prop = train_dur / total_duration
+            val_prop = val_dur / total_duration
+            test_prop = test_dur / total_duration
+
+            # Store row indices corresponding to each subject set
+            train_indices = train.index.values
+            val_indices = val.index.values
+            test_indices = test.index.values
+
+            all_train_folds.append(train_indices)
+            all_val_folds.append(val_indices)
+            all_test_folds.append(test_indices)
+
+            print(f"\n===== Fold {fold + 1} =====")
+            print(f"Train Subjects ({len(train_subjects)}): {train_subjects}")
+            print(f"Validation Subjects ({len(val_subjects)}): {val_subjects}")
+            print(f"Test Subjects ({len(test_subjects)}): {test_subjects}")
+            print(f"Epoch Distribution: Train {train_prop:.2%}, Val {val_prop:.2%}, Test {test_prop:.2%}")
+
+            # Optional: per-lab distribution check
+            def lab_stats(name, df):
+                return f"{name} lab distribution:\n{df.groupby('lab')['num_windows'].sum() / df['num_windows'].sum()}\n"
+
+            print(lab_stats("Train", train))
+            print(lab_stats("Validation", val))
+            print(lab_stats("Test", test))
+
+        return all_train_folds, all_val_folds, all_test_folds
+
+
+
+if __name__ == "__main__":
+
+    p = MSSVPreprocessor(data_folder="/home/coder/sleep/sleep-data/")
+
+    p.run()

From 9ab20c3bf5e57a70d02a11e0503c31bf8a45a9f3 Mon Sep 17 00:00:00 2001
From: Javi <javi@galatolo.xyz>
Date: Fri, 22 Aug 2025 14:39:11 +0000
Subject: [PATCH 2/6] fs=100 in sleepyrat

---
 physioex/preprocess/sleepyrat.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/physioex/preprocess/sleepyrat.py b/physioex/preprocess/sleepyrat.py
index 19280864..2f91f3f7 100644
--- a/physioex/preprocess/sleepyrat.py
+++ b/physioex/preprocess/sleepyrat.py
@@ -28,7 +28,7 @@ def __init__(
 
         super().__init__(
             dataset_name="sleepyrat",
-            signal_shape=[3, 512],
+            signal_shape=[3, 400],
             preprocessors_name=preprocessors_name,
             preprocessors=preprocessors,
             preprocessors_shape=preprocessor_shape,
@@ -95,7 +95,7 @@ def process_recording(self, edf_path, stages):
         mousedata:process_sleepdata_file, adapted for SleepyRat dataset
         '''
         
-        fs = 128
+        fs = 100
         epoch_second = 4
 
         # get the file name of the absolute path filename without the extension

From d91f63bcfc2fe8870e463cbaf06389d411e94e05 Mon Sep 17 00:00:00 2001
From: Javi <javi@galatolo.xyz>
Date: Fri, 22 Aug 2025 14:45:06 +0000
Subject: [PATCH 3/6] bug in mssv emg; split mssv per lab;

---
 physioex/preprocess/mssv.py | 29 ++++++++++++++++++++++-------
 1 file changed, 22 insertions(+), 7 deletions(-)

diff --git a/physioex/preprocess/mssv.py b/physioex/preprocess/mssv.py
index 2e6630b3..217f3ebc 100644
--- a/physioex/preprocess/mssv.py
+++ b/physioex/preprocess/mssv.py
@@ -50,10 +50,18 @@ def process_recording(edf_path, tsv_path):
     if fs != old_fs:
         eeg1 = resample(eeg1, int(len(eeg1) * fs / old_fs))
      
-    eeg2 = eeg1.copy() # only working with one EEG channel for now, but filling eeg2 for shape coherence with other datasets
+    if len(eeg_candidates) > 1:
+        eeg_candidates_other = [ch for ch in eeg_candidates if ch != eeg_channel]
+        eeg2_channel = random.choice(eeg_candidates_other)
+        eeg2, old_fs = read_channel_signal(edf_path, eeg2_channel)
+        eeg2 = filtfilt(b_band, 1, eeg2)
+        if fs != old_fs:
+            eeg2 = resample(eeg2, int(len(eeg2) * fs / old_fs))
+    else:
+        eeg2 = eeg1.copy()  # only one EEG channel available
     
     emg_candidates = [ch for ch in available_channels if 'EMG' in ch.upper()]
-    emg_channel = random.choice(eeg_candidates) if emg_candidates else None
+    emg_channel = random.choice(emg_candidates) if emg_candidates else None
     if emg_channel is None:
         print(f"Error: no EMG channel found in {edf_path}")
         print(f"Available channels: {available_channels}")
@@ -95,6 +103,8 @@ def process_recording(edf_path, tsv_path):
     signal = np.delete(signal, invalid_epochs, axis=0)
     
     signal = np.transpose(signal, (0, 2, 1))
+    
+    signal = signal * 1e6
 
     return signal.astype(np.float32), stages.astype(int)
 
@@ -126,6 +136,7 @@ class MSSVPreprocessor(Preprocessor):
 
     def __init__(
         self,
+        lab: str,
         preprocessors_name: List[str] = ["xsleepnet_mouse"],
         preprocessors=[xsleepnet_preprocessing_mouse],
         preprocessor_shape=[[3, 17, 129]],
@@ -133,7 +144,7 @@ def __init__(
     ):
 
         super().__init__(
-            dataset_name="mssv",
+            dataset_name="mssv/" + lab,
             signal_shape=[3, 400],
             preprocessors_name=preprocessors_name,
             preprocessors=preprocessors,
@@ -141,7 +152,8 @@ def __init__(
             data_folder=data_folder,
         )
         
-        self.source_dataset = os.path.join(self.dataset_folder, 'mssv_openneuro')
+        self.source_dataset = os.path.join(os.path.dirname(self.dataset_folder), 'openneuro')
+        self.lab = lab 
         
         self.split_subjects_table = None
         self._iterator = None
@@ -187,6 +199,7 @@ def get_subjects_records(self) -> np.ndarray:
 
             participants_path = os.path.join(self.source_dataset, 'participants.tsv')
             self.participants = pd.read_csv(participants_path, sep='\t')
+            self.participants = self.participants[self.participants['lab'] == self.lab]
             
             split_subjects_rows = []    
                 
@@ -374,6 +387,8 @@ def lab_stats(name, df):
 
 if __name__ == "__main__":
 
-    p = MSSVPreprocessor(data_folder="/home/coder/sleep/sleep-data/")
-
-    p.run()
+    MSSVPreprocessor(data_folder="/home/coder/sleep/sleep-data/", lab='lab_1').run()
+    MSSVPreprocessor(data_folder="/home/coder/sleep/sleep-data/", lab='lab_2').run()
+    MSSVPreprocessor(data_folder="/home/coder/sleep/sleep-data/", lab='lab_3').run()
+    MSSVPreprocessor(data_folder="/home/coder/sleep/sleep-data/", lab='lab_4').run()
+    MSSVPreprocessor(data_folder="/home/coder/sleep/sleep-data/", lab='lab_5').run()
\ No newline at end of file

From 7db31bd73327a6a06b9f19f4b44bd17fa8f919df Mon Sep 17 00:00:00 2001
From: Javier Garcia <jgciudad@outlook.com>
Date: Mon, 21 Jul 2025 12:44:38 +0200
Subject: [PATCH 4/6] more time-efficient voting strategy; support for variable
 number of channels and prototypes in prototype model

---
 physioex/train/networks/base.py      | 19 +++++++++++++-----
 physioex/train/networks/config.yaml  |  5 ++++-
 physioex/train/networks/prototype.py | 29 +++++++++++++---------------
 3 files changed, 31 insertions(+), 22 deletions(-)

diff --git a/physioex/train/networks/base.py b/physioex/train/networks/base.py
index a7e9d207..2f10a5da 100755
--- a/physioex/train/networks/base.py
+++ b/physioex/train/networks/base.py
@@ -8,11 +8,20 @@
 import torchmetrics as tm
 
 
-def voting_strategy(model: torch.nn.Module, inputs: torch.Tensor, L: int):
-    embeddings, outputs = model.encode(inputs)
-
-    outputs = torch.zeros_like(outputs)
-    embeddings = torch.zeros_like(embeddings)
+def voting_strategy( model : torch.nn.Module, inputs : torch.Tensor, L : int  ):
+
+    batch_size, night_length, n_channels, _, _ = inputs.size()
+    
+    embeddings_sample, _ = model.encode(inputs[:, 0:L])   
+    embeddings_dim = embeddings_sample.shape[-1]
+    
+    outputs = torch.zeros(
+        batch_size, night_length, model.n_classes, device=inputs.device, dtype=inputs.dtype 
+    )
+    
+    embeddings = torch.zeros(
+        batch_size, night_length, embeddings_dim, device=inputs.device, dtype=inputs.dtype 
+    )
 
     # input shape is ( bach_size, night_lenght, n_channels, ... )
     # segment the input in self.L segments with a sliding window of stride 1 and size self.L
diff --git a/physioex/train/networks/config.yaml b/physioex/train/networks/config.yaml
index 3a93c635..1e0d5e08 100644
--- a/physioex/train/networks/config.yaml
+++ b/physioex/train/networks/config.yaml
@@ -38,6 +38,7 @@ protosleeptransformer:
   target_transform: null
   model_kwargs:
     weights : [0.75, 0.25]
+    n_prototypes : 50
 
 protoseqsleepnet:
   model: physioex.train.networks.protoseqsleepnet:ProtoSeqSleepNet
@@ -45,7 +46,7 @@ protoseqsleepnet:
   target_transform: null
   model_kwargs:
     weights : [0.75, 0.25] #[0.99, 0.01]
-
+    n_prototypes : 50
 
 protosleeptransformer.1:
   model: physioex.train.networks.protosleeptransformer:ProtoSleepTransformerNet
@@ -53,6 +54,7 @@ protosleeptransformer.1:
   target_transform: null
   model_kwargs:
     weights : [1, 0]
+    n_prototypes : 50
 
 protoseqsleepnet.1:
   model: physioex.train.networks.protoseqsleepnet:ProtoSeqSleepNet
@@ -60,6 +62,7 @@ protoseqsleepnet.1:
   target_transform: null
   model_kwargs:
     weights : [1, 0]
+    n_prototypes : 50
 
 
 default:
diff --git a/physioex/train/networks/prototype.py b/physioex/train/networks/prototype.py
index 74c161ba..3725169d 100644
--- a/physioex/train/networks/prototype.py
+++ b/physioex/train/networks/prototype.py
@@ -48,26 +48,19 @@ def compute_loss(
         loss = loss + proto_loss + commit_loss
 
         proto_acc = self.wacc(proto_y, targets)
+        
+        channel_acc = [self.wacc(mcy[:, i], targets) for i in range(self.nn.in_channels)]
 
-        eeg_acc = self.wacc(mcy[:, 0], targets)
-        eog_acc = self.wacc(mcy[:, 1], targets)
-        emg_acc = self.wacc(mcy[:, 2], targets)
+        self.nn.channels_proba = channel_acc
+        
+        for i, c_acc in enumerate(channel_acc):
+            self.log(f"{log}/c{i}_acc", c_acc, sync_dist=True)
 
-        mc_loss = (
-            self.loss(mcy[:, 0], targets)
-            + self.loss(mcy[:, 1], targets)
-            + self.loss(mcy[:, 2], targets)
-        )
+        mc_loss = sum(self.loss(mcy[:, i], targets) for i in range(self.nn.in_channels))
         loss = loss + mc_loss
 
-        self.nn.channels_proba = [eeg_acc, eog_acc, emg_acc]
-
         self.log(f"{log}/p_acc", proto_acc, sync_dist=True)
 
-        self.log(f"{log}/eeg_acc", eeg_acc, sync_dist=True)
-        self.log(f"{log}/eog_acc", eog_acc, sync_dist=True)
-        self.log(f"{log}/emg_acc", emg_acc, sync_dist=True)
-
         if log == "val":
             self.log(f"{log}_acc", self.wacc(outputs, targets), sync_dist=True)
 
@@ -112,6 +105,8 @@ class ProtoSleepNet(nn.Module):
     def __init__(self, module_config=module_config):
         super(ProtoSleepNet, self).__init__()
 
+        self.in_channels = module_config["in_channels"]
+
         self.time_masking = TimeMasking(
             hidden_size=128,  # hidden size of the epoch encoder
             L=29,  # length of the time masking window
@@ -126,15 +121,17 @@ def __init__(self, module_config=module_config):
 
         self.channel_mixer = _initialize_residual_transformer(self.channel_mixer)
 
+        self.n_prototypes = module_config["n_prototypes"]
+        
         self.prototype = SimVQ(
             dim=128,
-            codebook_size= module_config.get("n_prototypes", 50),
+            codebook_size=self.n_prototypes,
             rotation_trick=True,  # use rotation trick from Fifty et al.
             channel_first=False,
         )
 
         self.channels_dropout = ChannelsDropout(dropout_prob=0.5)
-        self.channels_proba = [0.7, 0.7, 0.7]
+        self.channels_proba = [0.7 for _ in range(self.in_channels)]
 
         self.clf = nn.Linear(128, 5)
 

From 884c50bc326ea3c6e1b69a0071e36c1530d1ca65 Mon Sep 17 00:00:00 2001
From: Javier Garcia <jgciudad@outlook.com>
Date: Wed, 20 Aug 2025 15:25:33 +0200
Subject: [PATCH 5/6] bug in fold selection in eval and test sets

---
 physioex/data/datamodule.py   |  2 +-
 physioex/data/dataset.py      | 44 +++++++++++++++--------------------
 physioex/train/bin/parser.py  |  4 ++--
 physioex/train/utils/test.py  |  2 +-
 physioex/train/utils/train.py |  2 +-
 5 files changed, 24 insertions(+), 30 deletions(-)

diff --git a/physioex/data/datamodule.py b/physioex/data/datamodule.py
index fa0a0412..625af178 100644
--- a/physioex/data/datamodule.py
+++ b/physioex/data/datamodule.py
@@ -59,7 +59,7 @@ def __init__(
         sequence_length: int = 21,
         target_transform: Callable = None,
         task: str = "sleep",
-        folds: Union[int, List[int]] = -1,
+        folds: Union[int, List[int]] = 0,
         data_folder: str = None,
         num_workers: int = 0,
         data_prefetch: bool = True,
diff --git a/physioex/data/dataset.py b/physioex/data/dataset.py
index 00e34cf8..92a71fb1 100644
--- a/physioex/data/dataset.py
+++ b/physioex/data/dataset.py
@@ -88,8 +88,8 @@ def __init__(
             self.mean, self.std = self.readers[0].reader.mean, self.readers[0].reader.std
         
         self.dataset_idx = np.array(self.dataset_idx, dtype=np.int8)
-        # set the table fold to a random fold by default
-        self.split()
+        # set the table fold to fold 0 by default
+        self.split(0)
         self.target_transform = target_transform
 
         self.len = offset
@@ -107,36 +107,30 @@ def set_scaling(self, mean : torch.Tensor, std : torch.Tensor ):
     def get_scaling(self):
         return self.mean, self.std
 
-    def split(self, fold: int = -1, dataset_idx: int = -1):
+    def split(self, fold: int = 0, dataset_idx: int = -1):
+        assert fold >= 0, "ERR: fold must be >= 0. fold=-1 (randomly selected fold) is deprecated."
         assert dataset_idx < len(self.tables), "ERR: dataset_idx out of range"
 
-        # if fold is -1, set the split to a random fold for each dataset
-        if fold == -1 and dataset_idx == -1:
-            for i, table in enumerate(self.tables):
-                num_folds = [col for col in table.columns if "fold_" in col]
-                num_folds = len(num_folds)
-                selected_fold = np.random.randint(0, num_folds)
-
-                self.tables[i]["split"] = self.tables[i][f"fold_{selected_fold}"].map(
-                    {"train": 0, "valid": 1, "test": 2}
-                )
-        elif fold == -1 and dataset_idx != -1:
-            num_folds = [
-                col for col in self.tables[dataset_idx].columns if "fold_" in col
-            ]
-            num_folds = len(num_folds)
-            selcted_fold = np.random.randint(0, num_folds)
-
-            self.tables[dataset_idx]["split"] = table[f"fold_{selcted_fold}"].map(
-                {"train": 0, "valid": 1, "test": 2}
-            )
-        elif fold != -1 and dataset_idx == -1:
+        if dataset_idx == -1:
+            # Apply to all datasets
             for i, table in enumerate(self.tables):
+                fold_columns = [col for col in table.columns if "fold_" in col]
+                num_folds = len(fold_columns)
+                if fold >= num_folds:
+                    raise ValueError(f"ERR: fold {fold} is out of range for dataset {i}. Available folds: 0-{num_folds-1} (total: {num_folds} folds)")
+                
                 self.tables[i]["split"] = table[f"fold_{fold}"].map(
                     {"train": 0, "valid": 1, "test": 2}
                 )
         else:
-            self.tables[dataset_idx]["split"] = self.tables[dataset_idx][f"fold_{fold}"].map(
+            # Apply to specific dataset
+            table = self.tables[dataset_idx]
+            fold_columns = [col for col in table.columns if "fold_" in col]
+            num_folds = len(fold_columns)
+            if fold >= num_folds:
+                raise ValueError(f"ERR: fold {fold} is out of range for dataset {dataset_idx}. Available folds: 0-{num_folds-1} (total: {num_folds} folds)")
+            
+            self.tables[dataset_idx]["split"] = table[f"fold_{fold}"].map(
                 {"train": 0, "valid": 1, "test": 2}
             )
 
diff --git a/physioex/train/bin/parser.py b/physioex/train/bin/parser.py
index aee044e2..dc5a40c1 100644
--- a/physioex/train/bin/parser.py
+++ b/physioex/train/bin/parser.py
@@ -136,9 +136,9 @@ class PhysioExParser:
         "--fold",
         "-fd",
         type=int,
-        default=-1,
+        default=0,
         required=False,
-        help="Fold number to use. Default: -1 (random)",
+        help="Fold number to use. Default: Fold 0",
     )
 
     parser.add_argument(
diff --git a/physioex/train/utils/test.py b/physioex/train/utils/test.py
index 413b8d4e..5e465b82 100644
--- a/physioex/train/utils/test.py
+++ b/physioex/train/utils/test.py
@@ -21,7 +21,7 @@ def test(
     model_class=None,
     model_config: dict = None,
     batch_size: int = 128,
-    fold: int = -1,
+    fold: int = 0,
     checkpoint_path: str = None,
     results_path: str = None,
     num_nodes: int = 1,
diff --git a/physioex/train/utils/train.py b/physioex/train/utils/train.py
index dbe2daea..624021ec 100644
--- a/physioex/train/utils/train.py
+++ b/physioex/train/utils/train.py
@@ -25,7 +25,7 @@ def train(
     model_class: Type[SleepModule] = None,
     model_config: dict = None,
     batch_size: int = 128,
-    fold: int = -1,
+    fold: int = 0,
     num_validations: int = 10,
     checkpoint_path: str = None,
     max_epochs: int = 10,

From b78d42666e7b23b43aa5433183f2bfd62b12a7f0 Mon Sep 17 00:00:00 2001
From: Javier Garcia <jgciudad@outlook.com>
Date: Wed, 11 Feb 2026 17:25:02 +0100
Subject: [PATCH 6/6] revert bugs from #51391e9

---
 physioex/data/datamodule.py   |  2 +-
 physioex/train/bin/train.py   |  2 --
 physioex/train/utils/test.py  |  1 -
 physioex/train/utils/train.py | 12 +++++++++++-
 4 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/physioex/data/datamodule.py b/physioex/data/datamodule.py
index 625af178..764890dd 100644
--- a/physioex/data/datamodule.py
+++ b/physioex/data/datamodule.py
@@ -90,7 +90,7 @@ def __init__(
 
         if isinstance(eval_datasets, list):
             self.eval_dataset = PhysioExDataset(
-                datasets=datasets,
+                datasets=eval_datasets,
                 preprocessing=preprocessing,
                 selected_channels=selected_channels,
                 sequence_length=-1,
diff --git a/physioex/train/bin/train.py b/physioex/train/bin/train.py
index c6968800..844d1397 100755
--- a/physioex/train/bin/train.py
+++ b/physioex/train/bin/train.py
@@ -27,7 +27,6 @@ def train_script():
         "model": None,
         "batch_size": parser["batch_size"],
         "fold": parser["fold"],
-        "hpc": parser["hpc"],
         "num_validations": parser["num_validations"],
         "checkpoint_path": (
             parser["checkpoint_dir"]
@@ -54,7 +53,6 @@ def train_script():
             model_class=parser["model"],
             model_config=parser["model_kwargs"],
             batch_size=parser["batch_size"],
-            hpc=parser["hpc"],
             num_nodes=parser["num_nodes"],
             checkpoint_path=best_checkpoint,
             results_path=parser["results_path"],
diff --git a/physioex/train/utils/test.py b/physioex/train/utils/test.py
index 5e465b82..4aa705d4 100644
--- a/physioex/train/utils/test.py
+++ b/physioex/train/utils/test.py
@@ -35,7 +35,6 @@ def test(
 
     datamodule_kwargs["batch_size"] = batch_size
     datamodule_kwargs["folds"] = fold
-    datamodule_kwargs["num_nodes"] = num_nodes
 
     ##### DataModule Setup #####
     if isinstance(datasets, PhysioExDataModule):
diff --git a/physioex/train/utils/train.py b/physioex/train/utils/train.py
index 624021ec..24981af7 100644
--- a/physioex/train/utils/train.py
+++ b/physioex/train/utils/train.py
@@ -42,7 +42,6 @@ def train(
 
     datamodule_kwargs["batch_size"] = batch_size
     datamodule_kwargs["folds"] = fold
-    datamodule_kwargs["num_nodes"] = num_nodes
 
     if checkpoint_path is None:
         checkpoint_path = "models/" + str(uuid.uuid4())
@@ -102,6 +101,17 @@ def train(
         CSVLogger(save_dir=checkpoint_path),
     ]
 
+    ########### Trainer Setup ############
+    from lightning.pytorch.accelerators import find_usable_cuda_devices
+
+    try :
+        devices = find_usable_cuda_devices(-1)
+        logger.info( f"Available devices: {devices}")
+        effective_batch_size = batch_size * num_nodes *  len(devices)
+
+    except :
+        devices = "auto"
+        effective_batch_size = batch_size * num_nodes
     num_steps = datamodule.__len__() // effective_batch_size
     val_check_interval = max(1, num_steps // num_validations)