OSCR

Bidirectional Mamba-2 boosts EEG super-resolution via regression and diffusion.

Code ↔ Paper

18 matches between paragraphs of the paper and lines of its authors' code, computed by the harvester (lexical-v1). Click a colored paragraph or line to see its counterpart.

The 18 matches
  1. [1] § 2 Materials and methods › 2.3 Datasets and preprocessing ↔ utils.py, lines 169–198 · score 0.75 · notch filter, bad channels, MNE, FIR, head, 40 Hz
  2. [2] § 3 Results ↔ models.py, lines 236–371 · score 0.74 · n_layers, d_model, d_state, Bidirectional Mamba blocks, diffusion modeling, electrode
  3. [3] § 2 Materials and methods › 2.1 Architecture of the proposed system ↔ diffusion_conditioning_ablations.py, lines 28–57 · score 0.74 · squaredcos_cap_v2, DDPM scheduler, beta schedule, clipping, timesteps, linearly
  4. [4] § 2 Materials and methods › 2.1 Architecture of the proposed system ↔ ablations.py, lines 28–53 · score 0.72 · squaredcos_cap_v2, DDPM scheduler, beta schedule, timesteps, linearly, prediction
  5. [5] § 3 Results › 3.2 Results for spatio-temporal super-resolution ↔ compare_maser_topomap.py, lines 213–307 · score 0.67 · PSD topomap comparison, MASER x4, DiBiMa, x8, signal
  6. [6] § 3 Results › 3.2 Results for spatio-temporal super-resolution ↔ utils.py, lines 956–1075 · score 0.66 · spectral fidelity, frequency domains, temporal super resolution, waveform, EEG
  7. [7] § 2 Materials and methods › 2.3 Datasets and preprocessing ↔ visualize_input.ipynb, lines 268–407 · score 0.64 · notch filter, MNE, FIR, scoring, 40 Hz, 50 Hz
  8. [8] § 2 Materials and methods › 2.1 Architecture of the proposed system ↔ compare_to_maser.py, lines 235–315 · score 0.63 · MASER models, FLOPs, DiBiMa, inference, devices, position
  9. [9] § 3 Results › 3.3 Super-resolution explainability and downstream classification task ↔ eeg_fid.py, lines 175–205 · score 0.61 · EEG FID scores, DiBiMa, MASER, x4, models
  10. [10] § 2 Materials and methods › 2.4 Metrics ↔ metrics.py, lines 14–27 · score 0.60 · Peak Signal, Noise Ratio, PSNR, Metrics
  11. [11] § 2 Materials and methods › 2.2 Mamba blocks and the bidirectional mamba layer ↔ models.py, lines 667–809 · score 0.58 · bidirectional Mamba, Mamba blocks, fused, BiMamba, kernel, backward
  12. [12] § 3 Results › 3.2 Results for spatio-temporal super-resolution ↔ app.py, lines 157–184 · score 0.58 · PSD topomap, Qualitative comparison, BiMa, super resolution, x4, x8
  13. [13] § 3 Results › 3.3 Super-resolution explainability and downstream classification task ↔ downstream_task/metrics_class.py, lines 17–61 · score 0.58 · ROC AUC, F1 score, binary, accuracy, metrics, downstream
  14. [14] § 3 Results › 3.3 Super-resolution explainability and downstream classification task ↔ eeg_fid.py, lines 25–69 · score 0.57 · EEG FID scores, ResNet50, pretrained
  15. [15] § 3 Results › 3.3 Super-resolution explainability and downstream classification task ↔ downstream_task/metrics.py, lines 94–144 · score 0.57 · F1 score, ROC AUC, curves, accuracy, classification, metrics
  16. [16] § 2 Materials and methods › 2.4 Metrics ↔ metrics.py, lines 99–118 · score 0.53 · Pearson Correlation Coefficient, PCC, Metrics, channel
  17. [17] § 2 Materials and methods › 2.4 Metrics ↔ utils.py, lines 956–1075 · score 0.52 · Pearson Correlation Coefficient, PCC
  18. [18] § 2 Materials and methods › 2.1 Architecture of the proposed system ↔ diffusion_conditioning_ablations.py, lines 214–318 · score 0.52 · internal residuals, Electrode position, mamba blocks, embeddings, layer, resolution

Paper

Loaded from Europe PMC by your browser, not stored by OSCR: doi.org · Europe PMC

The paper is loaded when this pane is shown.

The authors' code

Python · 1,120 lines · 44 KB · MIT · 3 matches

  2. import math
  3. import os
  4. import numpy as np
  5. from sklearn.preprocessing import StandardScaler
  6. import torch
  7. from torch.utils.data import Dataset
  8. import mne
  9. from mne.datasets import eegbci
  10. from scipy import signal
  11. import shutil
  12. import matplotlib.pyplot as plt
  13. from umap import UMAP
  14. import random
  15. import gc
  16. import torch
  17. import scipy
  18. from sklearn.model_selection import train_test_split
  19. from mne.preprocessing import find_bad_channels_lof, find_bad_channels_maxwell
  21. # Set random seeds for reproducibility
  22. seed = 2
  23. np.random.seed(seed)
  24. torch.manual_seed(seed)
  26. #case 1 of https://ieeexplore.ieee.org/document/9796118
  28. seed_channels = ["Fp1", "Fpz", "Fp2", "AF3", "AF4", "F7", "F5", "F3", "F1", "Fz", "F2", "F4", "F6", "F8", "FT7", "FC5", "FC3", "FC1", "FCz", "FC2", "FC4", "FC6", "FT8", "T7", "C5", "C3", "C1", "Cz", "C2", "C4", "C6", "T8", "TP7", "CP5", "CP3", "CP1", "CPz", "CP2", "CP4", "CP6", "TP8", "P7", "P5", "P3", "P1", "Pz", "P2", "P4", "P6", "P8", "PO7", "PO5", "PO3", "POz", "PO4", "PO6", "PO8", "CB1", "O1", "Oz", "O2", "CB2"]
  29. mmi_channels = ['FC5', 'FC3', 'FC1', 'FCz', 'FC2', 'FC4', 'FC6', 'C5', 'C3', 'C1', 'Cz', 'C2', 'C4', 'C6', 'CP5', 'CP3', 'CP1', 'CPz', 'CP2', 'CP4', 'CP6', 'Fp1', 'Fpz', 'Fp2', 'AF7', 'AF3', 'AFz', 'AF4', 'AF8', 'F7', 'F5', 'F3', 'F1', 'Fz', "F2", "F4", "F6", "F8", "FT7", "FT8", "T7", "T8", "T9", "T10", "TP7", "TP8", "P7", "P5", "P3", "P1", "Pz", "P2", "P4", "P6", "P8", "PO7", "PO3", "POz", "PO4", "PO8", "O1", "Oz", "O2", "Iz"]
  32. map_seed_channels = {seed_channels[i]: i for i in range(len(seed_channels))}
  33. map_mmi_channels = {mmi_channels[i]: i for i in range(len(mmi_channels))}
  35. case1_seed = {
  36. 'x2': ['AF3', 'AF4', 'FT7', 'FC5', 'FC3', 'FC1', 'FCz', 'FC2', 'FC4', 'FC6', 'FT8', 'TP7', 'CP5', 'CP3', 'CP1', 'CPz', 'CP2', 'CP4', 'CP6', 'TP8', 'CB1', 'CB2', 'PO3', 'PO4', 'PO5', 'PO7', 'PO6', 'PO8', 'POz', 'O1', 'Oz', 'O2'],
  37. 'x4': ['Fp1', 'Fp2', 'F5', 'Fz', 'F6', 'C3', 'Cz', 'C4', 'T7', 'T8', 'P5', 'Pz', 'P6', 'O1', 'O2'],
  38. 'x8': ['AF3', 'AF4', 'FC5', 'FC6', 'CP5', 'CP6', 'PO5', 'PO6']
  39. }
  40. case1_mmi = {
  41. "x2": ["Fpz", "AF7", "AF3", "AFz", "AF4", "AF8", "F7", "F3", "Fz", "F4", "F8", "FT7", "FC3", "FCz", "FC4", "FT8", "T7", "C3", "Cz", "C4", "T8", "TP7", "CP3", "CPz", "CP4", "TP8", "P7", "P3", "Pz", "P4", "P8", "PO7", "PO3", "POz", "PO4", "PO8", "Oz", "Iz"],
  42. "x4": ["Fp1", "Fp2", "F5", "Fz", "F6", "T7", "C3", "Cz", "C4", "T8", "P5", "Pz", "P6", "O1", "Oz", "O2"],
  43. "x8": ["AF3", "AF4", "FC5", "FC6", "CP5", "CP6", "O1", "O2"]
  44. }
  46. case2_mmi = {
  47. "x2": ["Fp1", "Fp2", "F5", "F1", "F2", "F6", "FC5", "FC1", "FC2", "FC6", "T9", "C5", "C1", "C2", "C6", "T10", "CP5", "CP1", "CP2", "CP6", "P5", "P1", "P2", "P6", "O1", "O2"],
  48. "x4": ["Fpz", "AF3", "AF4", "FC5", "FC1", "FC2", "FC6", "T9", "CP5", "CP1", "CP2", "CP6", "T10", "PO3", "PO4", "Oz"],
  49. "x8": ["Fp1", "Fp2", "Fz", "T7", "T8", "Pz", "O1", "O2"]
  50. }
  51. case2_seed = {
  52. "x2": ["Fpz", "AF3", "AF4", "F7", "F3", "Fz", "F4", "F8", "FT7", "FC3", "FCz", "FC4", "FT8", "T7", "C3", "Cz", "C4", "T8", "TP7", "CP3", "CPz", "CP4", "TP8", "P7", "P3", "Pz", "P4", "P8", "PO7", "PO3", "POz", "PO4", "PO8", "Oz"],
  53. "x4": ["Fpz", "AF3", "AF4", "FC5", "FC1", "FC2", "FC6", "CP5", "CP1", "CP2", "CP6", "CB1", "PO3", "PO4", "CB2", "Oz"],
  54. "x8": ["Fp1", "Fp2", "Fz", "T7", "T8", "Pz", "O1", "O2"]
  55. }
  57. unmask_channels = {
  58. "mmi":{
  59. "x2": [map_mmi_channels[i] for i in case1_mmi['x2']],
  61. "x4": [map_mmi_channels[i] for i in case1_mmi['x4']], # Your 16ch (frontal/central/parietal) #[0, 2, 4, 6, 14, 16, 18, 20, 22, 25, 27, 42, 43, 56, 58, 61]
  63. "x8": [map_mmi_channels[i] for i in case1_mmi['x8']] # Even 10-10 coverage
  64. },
  65. "seed":{
  66. "x2": [map_seed_channels[i] for i in case1_seed['x2']],
  68. "x4": [map_seed_channels[i] for i in case1_seed['x4']], # Frontal/central/parietal
  70. "x8": [map_seed_channels[i] for i in case1_seed['x8']] # Even 10-10 coverage
  71. }
  72. }
  74. #reverse ch_name:i to i:ch_name
  75. unmask_channels = {
  76. dataset: {
  77. key: sorted(value) for key, value in channels.items()
  78. } for dataset, channels in unmask_channels.items()
  79. }
  81. map_runs_dataset = {
  82. "mmi": range(1, 15),
  83. "seed": None # All files in folder
  84. }
  86. map_tasks = {
  87. 'task1': 'open-close right or left fist',
  88. 'task3': 'open-close both fists or both feet',
  89. 'task2': 'imagine right or left fist',
  90. 'task4': 'imagine both fists or both feet'
  91. }
  93. map_labels_seed = {
  94. -1: 0, # negative
  95. 0: 1, # neutral
  96. 1: 2 # positive
  97. }
  99. labels_subject_mapping = {
  100. 1: 1,
  101. 2: 0,
  102. 3: -1,
  103. 4: -1,
  104. 5: 0,
  105. 6: 1,
  106. 7: -1,
  107. 8: 0,
  108. 9: 1,
  109. 10: 1,
  110. 11: 0,
  111. 12: -1,
  112. 13: 0,
  113. 14: 1,
  114. 15: -1
  115. }
  117. map_runs_mmi = {
  118. 1: 'eyes_open',
  119. 2: 'eyes_closed',
  120. 3: map_tasks['task1'],
  121. 4: map_tasks['task2'],
  122. 5: map_tasks['task3'],
  123. 6: map_tasks['task4'],
  124. 7: map_tasks['task1'],
  125. 8: map_tasks['task2'],
  126. 9: map_tasks['task3'],
  127. 10: map_tasks['task4'],
  128. 11: map_tasks['task1'],
  129. 12: map_tasks['task2'],
  130. 13: map_tasks['task3'],
  131. 14: map_tasks['task4']
  132. }
  134. map_labels_mmi = {
  135. 'eyes_open': 0,
  136. 'eyes_closed': 1,
  137. 'open-close right or left fist': 2,
  138. 'open-close both fists or both feet': 3,
  139. 'imagine right or left fist': 4,
  140. 'imagine both fists or both feet': 5
  141. }
  142. map_labels_mmi_rev = {v: k for k, v in map_labels_mmi.items()}
  145. map_annotations = {
  146. "T0": "rest",
  147. "T1": "right or left fist",
  148. "T2": "both fists or both feet"
  149. }
  151. def get_lr_data_temporal(eeg_hr, factor=2):
  152. """Downsample high-res EEG temporally by factor."""
  153. if eeg_hr.ndim == 2:
  154. _, num_samples = eeg_hr.shape
  155. elif eeg_hr.ndim == 3:
  156. _, _, num_samples = eeg_hr.shape
  157. downsampled_length = num_samples // factor
  158. eeg_lr = signal.resample(eeg_hr, downsampled_length, axis=-1)
  159. return eeg_lr
  161. def get_lr_data_spatial(eeg_64, dataset_name, sr_ratio):
  163. """Extract low-res from 64ch EEG using unmask_channels"""
  164. indices = unmask_channels[dataset_name][f"x{sr_ratio}"]
  165. if eeg_64.ndim == 2:
  166. return eeg_64[indices, :]
  167. elif eeg_64.ndim == 3:
  168. return eeg_64[:, indices, :]
  169. # -------------------------------
  170. # Preprocessing MNE
  171. # -------------------------------
  172. def _preprocess_raw(raw, dataset_name, verbose=False):
  173. fs_cut = 40 if dataset_name == "mmi" else 50
  174. raw.pick_types(eeg=True, eog=False, stim=False)
  176. if dataset_name == "seed":
  177. # SKIP dev_head_t - not needed for EEG preprocessing [web:42]
  179. # Bad channel detection & cleaning (your function)
  180. raw = detect_and_clean_seed_trial(raw)
  181. if raw is None:
  182. return None
  184. # SEED filters
  185. raw.notch_filter(50.0, fir_design='firwin', verbose=verbose)
  186. raw.filter(1.0, fs_cut, fir_design='firwin', verbose=verbose)
  188. # Scale (confirm units with raw.plot() first)
  189. data = raw.get_data()
  190. #data = data * 1e-4
  191. data = (data - data.mean(1, keepdims=True)) / (data.std(1, keepdims=True) + 1e-8)
  192. raw._data = data
  194. else:
  195. raw.notch_filter(50.0, fir_design='firwin', verbose=verbose)
  196. raw.filter(0.5, fs_cut, fir_design='firwin', verbose=verbose)
  198. return raw
  201. # Extract valid EEG positions (C, 3), nan-mask if needed
  202. def get_electrode_positions(raw, channel_order=None):
  203. """
  204. Extract electrode positions from raw EEG data.
  206. Parameters
  207. ----------
  208. raw : mne.io.Raw
  209. Raw EEG data with montage set
  210. channel_order : list, optional
  211. Desired channel order for reordering
  213. Returns
  214. -------
  215. locs : torch.Tensor
  216. Electrode positions (C_eeg, 3)
  217. """
  218. picks = mne.pick_types(raw.info, eeg=True)
  219. locs = np.array([raw.info['chs'][p]['loc'][:3] for p in picks]) # (C_eeg, 3)
  221. # Check for NaNs and report which channels are problematic
  222. if np.isnan(locs).any():
  223. nan_channels = []
  224. for i, p in enumerate(picks):
  225. if np.isnan(locs[i]).any():
  226. ch_name = raw.info['chs'][p]['ch_name']
  227. nan_channels.append(ch_name)
  229. # Get available montage channels for debugging
  230. montage = raw.get_montage()
  231. if montage is not None:
  232. available_channels = list(montage.get_positions()['ch_pos'].keys())
  233. print(f"⚠️ Available montage channels: {available_channels[:10]}...") # Show first 10
  235. raise ValueError(
  236. f"NaNs persist after montage; check channel names. "
  237. f"Problematic channels: {nan_channels}. "
  238. f"Raw channels: {[raw.info['chs'][p]['ch_name'] for p in picks]}"
  239. )
  241. if channel_order: # Reorder to model input (e.g., 64 chs)
  242. reorder_idx = [channel_order.index(ch) for ch in raw.ch_names if ch in channel_order]
  243. locs = locs[reorder_idx]
  245. locs = torch.tensor(locs, dtype=torch.float32) # (C_eeg, 3)
  246. return locs
  248. def load_seed_channel_positions(filepath):
  250. channels = []
  251. positions = []
  252. with open(filepath, 'r') as f:
  253. for line in f:
  254. parts = line.strip().split()
  255. if len(parts) >= 4:
  256. ch_name = parts[-1] # Or parts[-1] if label last
  257. theta_deg = float(parts[1]) # Azimuth, negative=left
  258. phi_frac = float(parts[2]) # Elevation fraction (~0-0.6)
  259. phi_rad = np.radians(phi_frac * 90) # Adjust scale to ~90° max
  260. theta_rad = np.radians(theta_deg)
  261. x = np.sin(phi_rad) * np.cos(theta_rad)
  262. y = np.sin(phi_rad) * np.sin(theta_rad)
  263. z = np.cos(phi_rad)
  264. channels.append(ch_name)
  265. positions.append([x, y, z])
  266. return channels, np.array(positions)
  268. def set_montage(signal, dataset_name, pos=None, channel_names=None, fs=160):
  269. """
  270. Set montage for raw EEG data.
  272. For MMI: pos=None, uses standard_1020
  273. For SEED: pos=array of positions from .pos file
  274. """
  276. #print(channel_names)
  278. # Clean channel names
  279. cleaned_names = []
  280. for ch in channel_names:
  281. clean = ch.strip().rstrip('.')
  282. cleaned_names.append(clean)
  284. # Create MNE RawArray
  285. info = mne.create_info(ch_names=cleaned_names, sfreq=fs, ch_types='eeg')
  286. raw = mne.io.RawArray(signal, info)
  288. if dataset_name == "mmi":
  289. # MMI dataset: Use standard montage
  290. try:
  291. montage = mne.channels.make_standard_montage('standard_1020')
  292. #montage_channels = list(montage.get_positions()['ch_pos'].keys())
  293. #print(f"Montage channels available: {montage_channels[:10]}...") # Show first 10
  294. #print(f"Raw channels: {cleaned_names}")
  295. raw.set_montage(montage, on_missing='warn')
  296. except Exception as e:
  297. print(f"⚠️ Error setting standard montage: {e}")
  298. else:
  299. # SEED dataset: Use custom positions from .pos file
  300. try:
  301. ch_pos = {channel_names[i]: pos[i] for i in range(len(channel_names))}
  302. montage = mne.channels.make_dig_montage(
  303. ch_pos=ch_pos,
  304. coord_frame='head'
  305. )
  306. raw.set_montage(montage, on_missing='warn')
  307. except Exception as e:
  308. print(f"⚠️ Error setting custom montage: {e}")
  309. raise
  311. return raw
  314. def download_mmi_data(subject_ids, runs, project_path, demo = False, verbose=False, is_classification=False):
  316. datas = []
  317. labels = []
  318. positions = []
  319. channel_names = None
  321. for i, subject in enumerate(subject_ids):
  322. #print(f"Processing subject {subject}")
  323. print(f"⬇️ Downloading/Reading data for subject: {i+1}/{len(subject_ids)}", end='\r')
  324. if demo:
  325. if i == 1:
  326. break # For demo, process only first subject
  327. for run in runs:
  328. if is_classification and run not in [1, 2]:
  329. continue # Skip non-classification runs
  331. local_path = os.path.join(project_path, f'S{subject:03d}R{run:02d}.edf')
  332. #print(f"Checking local path: {local_path}")
  333. if not os.path.exists(local_path):
  334. print(f"Downloading S{subject:03d}R{run:02d}.edf...")
  335. try:
  336. eegbci.load_data(subject, [run], path=project_path, update_path=True, force_update=False)
  337. downloaded = os.path.join(
  338. os.path.dirname(project_path), "mmi",
  339. 'MNE-eegbci-data', 'files', 'eegmmidb', '1.0.0',
  340. f'S{subject:03d}', f'S{subject:03d}R{run:02d}.edf'
  341. )
  342. os.makedirs(os.path.dirname(local_path), exist_ok=True)
  343. shutil.move(downloaded, local_path)
  344. print(f"Saved to: {local_path}")
  345. except Exception as e:
  346. print(f"Download error: {e}")
  347. continue
  348. try:
  349. raw = mne.io.read_raw_edf(local_path, preload=True, verbose=verbose)
  351. mne.datasets.eegbci.standardize(raw)
  352. signal = raw.get_data()
  353. #print(raw.ch_names)
  354. if channel_names is None:
  355. channel_names = raw.ch_names
  356. else:
  357. if channel_names != raw.ch_names:
  358. raise ValueError("Inconsistent channel names across recordings.")
  360. #we don't have the positions for mmi, we wait to set the montage and then extract them
  361. raw = set_montage(signal, "mmi", pos=None, channel_names=channel_names)
  362. if raw is None:
  363. print(f"⚠️ set_montage returned None for {local_path}")
  364. continue
  366. raw = _preprocess_raw(raw, dataset_name="mmi", verbose=verbose)
  367. if raw is None:
  368. print(f"⚠️ _preprocess_raw returned None for {local_path}")
  369. continue
  371. position = get_electrode_positions(raw, channel_order=None)
  372. if position is None:
  373. print(f"⚠️ No valid electrode positions for {local_path}")
  374. continue
  376. positions.append(position)
  377. data = raw.get_data()
  378. data = data*1e3 # Convert to µV
  379. data = data.astype(np.float32)
  380. datas.append(data)
  381. label = map_runs_mmi[run]
  382. labels.append(map_labels_mmi[label])
  384. except Exception as e:
  385. print(f"⚠️ Exception during processing {local_path}: {e}")
  386. import traceback
  387. traceback.print_exc() # This will show the exact line causing the error
  388. continue
  390. # Clean up downloaded files
  391. path_to_remove = os.path.join(
  392. os.path.dirname(project_path), "mmi",
  393. 'MNE-eegbci-data', 'files', 'eegmmidb'
  394. )
  395. if os.path.exists(path_to_remove):
  396. shutil.rmtree(path_to_remove)
  398. labels = torch.tensor(np.array(labels), dtype=torch.int32)
  399. positions = np.array(positions)
  400. positions = torch.tensor(positions, dtype=torch.float32)
  401. return datas, labels, positions, channel_names
  403. def load_mat(filepath):
  404. mat = scipy.io.loadmat(filepath)
  405. return mat
  407. def load_seed_data(subject_ids, project_path, demo=False, verbose=False):
  409. seed_datapath = os.path.join(project_path, "Preprocessed_EEG")
  410. files = os.listdir(seed_datapath)
  411. if len(files) == 0:
  412. raise ValueError("No SEED data found in the specified path.")
  414. datas = []
  415. labels = []
  416. positions = []
  418. channels, position = load_seed_channel_positions(os.path.join(project_path, "channel_62_pos.locs"))
  419. #print(position)
  420. i = 0
  421. for file in files:
  422. for subject in subject_ids:
  423. print(f"Processing subject {i+1}/{len(subject_ids)}", end='\r')
  424. if demo:
  425. if i == 1:
  426. break # For demo, process only first subject
  427. if file.startswith(f'{int(subject)}_'):
  428. filepath = os.path.join(seed_datapath, file)
  429. try:
  430. data_hr = load_mat(filepath)
  431. sfreq = 200 # Original SEED sampling rate
  432. eeg_keys = [key for key in data_hr.keys() if "eeg" in key.lower()]
  433. if len(eeg_keys) == 0:
  434. print(f"No EEG data found in {filepath}.")
  435. continue
  436. else:
  437. for key in eeg_keys:
  439. data = data_hr[key] # Shape: (channels, samples)
  440. #raw = mne.io.RawArray(data, mne.create_info(ch_names=channels, sfreq=sfreq, ch_types='eeg'))
  442. raw = set_montage(data, "seed", pos=position, channel_names=channels, fs=sfreq)
  443. raw = _preprocess_raw(raw, dataset_name="seed", verbose=verbose)
  444. if raw is None:
  445. continue
  446. data = raw.get_data()
  447. data = data.astype(np.float32)
  448. datas.append(data)
  449. label = map_labels_seed[labels_subject_mapping[subject]]
  450. labels.append(label)
  451. positions.append(torch.tensor(position, dtype=torch.float32))
  453. i += 1
  454. except Exception as e:
  455. print(f"⚠️ Exception during processing {filepath}: {e}")
  457. positions = np.array(positions)
  458. positions = torch.tensor(positions, dtype=torch.float32)
  459. labels = np.array(labels, dtype=np.int32)
  460. labels = torch.tensor(labels)
  461. return datas, labels, positions, channels
  463. def download_eegbci_data(subject_ids, runs, project_path, demo=False, dataset_name="mmi", is_classification=False, verbose=False):
  464. """
  465. Scarica i dati EEG BCI per i soggetti e le sessioni specificate.
  466. Salva i file EDF localmente in project_path.
  467. """
  468. if dataset_name.lower() not in ["mmi", "seed"]:
  469. raise ValueError("Dataset not supported. Use 'mmi' or 'seed'.")
  471. if dataset_name.lower() == "mmi":
  472. return download_mmi_data(subject_ids, runs, project_path, demo=demo, verbose=verbose, is_classification=is_classification)
  473. else:
  474. return load_seed_data(subject_ids, project_path, demo=demo, verbose=verbose)
  476. def clear_memory():
  477. gc.collect()
  478. torch.cuda.empty_cache()
  480. def set_seed(seed):
  481. torch.manual_seed(seed)
  482. np.random.seed(seed)
  483. if torch.cuda.is_available():
  484. torch.cuda.manual_seed_all(seed)
  486. class EEGWindowsDataset(Dataset):
  487. """
  488. Minimal dataset for pre-processed, pre-windowed EEG data.
  489. Assumes windows are already normalized and ready to use.
  490. Only generates LR on-the-fly during __getitem__.
  491. """
  492. def __init__(self, windows, labels, positions, sr_type="temporal", dataset_name="mmi",
  493. target_channels=64, multiplier=2, channel_names=None, fs_hr=160):
  494. """
  495. Args:
  496. windows: Tensor/array (N, C, T) - preprocessed HR windows
  497. labels: Tensor/array (N,) - class labels
  498. positions: Tensor/array (N, C, 3) or (C, 3) - electrode positions
  499. sr_type: 'temporal' or 'spatial'
  500. num_channels: int - for spatial SR (number of LR channels)
  501. multiplier: int - temporal downsampling factor
  502. channel_names: list - optional channel names
  503. """
  505. self.datas_hr = torch.tensor(windows, dtype=torch.float32) if not isinstance(windows, torch.Tensor) else windows.float()
  506. self.labels = torch.tensor(labels, dtype=torch.long) if not isinstance(labels, torch.Tensor) else labels.long()
  508. # Handle positions
  509. if isinstance(positions, torch.Tensor):
  510. self.positions = positions.float()
  511. else:
  512. self.positions = torch.tensor(positions, dtype=torch.float32)
  514. # Broadcast if single position (C, 3) -> (N, C, 3)
  515. if self.positions.dim() == 2:
  516. self.positions = self.positions.unsqueeze(0).expand(len(self.datas_hr), -1, -1)
  518. self.sr_type = sr_type
  519. self.target_channels = target_channels
  520. self.multiplier = multiplier
  521. self.channel_names = channel_names
  522. self.dataset_name = dataset_name
  523. self.num_classes = len(torch.unique(self.labels))
  524. self.ref_position = self.positions[0]
  525. self.fs_hr = fs_hr
  526. self.fs_lr = int(fs_hr//self.multiplier) if sr_type == "temporal" else fs_hr
  527. self.num_channels = math.ceil(self.target_channels / self.multiplier) if sr_type == "spatial" else self.datas_hr.shape[1]
  528. print(f"✅ EEGWindowsDataset: {len(self.datas_hr)} windows, shape {self.datas_hr.shape}")
  530. def _downsample(self, data, factor):
  531. """Temporal downsampling via signal.resample"""
  532. downsampled_length = data.shape[1] // factor
  533. return signal.resample(data, downsampled_length, axis=1)
  535. def __len__(self):
  536. return len(self.datas_hr)
  538. def __getitem__(self, idx):
  539. hr_data = self.datas_hr[idx] # (C, T)
  540. pos = self.positions[idx] # (C, 3)
  541. label = self.labels[idx] # scalar
  543. # Generate LR on-the-fly
  544. if self.sr_type == "temporal":
  545. lr_data = self._downsample(hr_data.numpy(), self.multiplier)
  546. lr_data = torch.tensor(lr_data, dtype=torch.float32)
  547. else: # spatial
  548. lr_data = get_lr_data_spatial(hr_data.clone(), dataset_name=self.dataset_name, sr_ratio=self.multiplier)
  550. return lr_data, hr_data, pos, label
  552. # -------------------------------
  553. # EEGDataset for Super-Resolution EEG
  554. # -------------------------------
  555. class EEGDataset(Dataset):
  556. """
  557. Dataset EEG preprocessato per Super-Resolution EEG (EEG BCI dataset).
  558. Genera segmenti LR (160/sr_factor Hz) e HR (160 Hz) sincronizzati.
  559. """
  560. def __init__(self, subject_ids, data_folder, dataset_name = "mmi", sr_type="temporal", seconds=10, verbose=False, demo=False, num_channels=64, multiplier=2, normalize=False, is_classification=False):
  562. self.project_path = data_folder
  563. self.verbose = verbose
  564. self.normalize = normalize
  565. self.demo = demo
  566. self.seconds = seconds
  567. self.num_channels = num_channels
  568. self.dataset_name = dataset_name
  569. self.runs = map_runs_dataset[self.dataset_name]
  570. self.sr_type = sr_type # 'temporal' or 'spatial'
  571. self.multiplier = multiplier # Downsampling factor for temporal or spatial SR
  572. self.fs_hr = 160 if self.dataset_name == "mmi" else 200
  573. self.hr_window_length = self.fs_hr * self.seconds
  574. self.is_classification = is_classification
  576. self.scaler = StandardScaler() #MinMaxScaler(feature_range=(0, 1))
  578. self.datas, self.labels, self.positions, self.channel_names = download_eegbci_data(
  579. subject_ids, runs=self.runs, project_path=self.project_path,
  580. demo=self.demo, dataset_name=self.dataset_name, verbose=self.verbose, is_classification=self.is_classification
  581. )
  583. self.positions = self.positions.cpu()
  584. self.labels = self.labels.cpu()
  586. self.num_classes = 6 if self.dataset_name == "mmi" else 3 #len(torch.unique(self.labels))
  588. self.ref_position = self.positions[0] # Assuming all raws have same channel positions
  590. self.datas_hr, self.positions = self._split_windows(self.hr_window_length)
  591. print(f"\n✅ Number of segments created: {len(self.datas_hr)}")
  593. #self.datas_hr = self._zscore_normalization(torch.tensor(self.datas_hr)).numpy()
  594. #print(f"\nData z-score normalization complete: {self.datas_hr.shape}")
  596. if self.normalize:
  597. self.datas_hr = self._normalize_data()
  598. print(f"\nData normalization complete: {self.datas_hr.shape}")
  600. def _zscore_normalization(self, data):
  601. # Z-score per channel (common for EEG DL)
  602. mean = data.mean(dim=-1, keepdim=True) # [B,C,1]
  603. std = data.std(dim=-1, keepdim=True)
  604. std = torch.clamp(std, min=1e-5) # Avoid div0
  605. normalized = (data - mean) / std # [-3,3] typical
  606. return normalized
  608. def _normalize_data(self):
  609. # self.datas_hr shape: (N_windows, C, T)
  610. N, C, T = self.datas_hr.shape
  611. # Flatten to (N*C, T) - treat each channel independently
  612. data_reshaped = self.datas_hr.reshape(N * C, T) # (N*C, T)
  613. # Fit scaler per channel
  614. data_norm = self.scaler.fit_transform(data_reshaped)
  615. # Reshape back to (N, C, T)
  616. data_normalized = torch.tensor(data_norm, dtype=torch.float32).reshape(N, C, T)
  617. return data_normalized
  619. def _downsample(self, data, factor):
  620. _, num_samples = data.shape
  621. downsampled_length = num_samples // factor
  622. downsampled_data = signal.resample(data, downsampled_length, axis=1)
  623. return downsampled_data
  625. def _split_windows(self, window_length, stride=None): # stride=window_length for non-overlap
  626. if stride is None:
  627. stride = window_length # Non-overlapping by default
  628. datas_hr = []
  629. positions = []
  630. labels = []
  631. for i, data in enumerate(self.datas): # self.datas = list of (C, T_i) arrays
  632. print(f"Splitting data {i+1}/{len(self.datas)} (len={data.shape[1]})", end='\r')
  633. position = self.positions[i] # (C, 3)
  634. T = data.shape[1]
  635. num_windows = (T - window_length) // stride + 1 # Floor div for valid windows
  636. for j in range(num_windows):
  637. start = j * stride
  638. end = start + window_length
  639. window = data[:, start:end].astype(np.float32) # (C, window_length)
  640. datas_hr.append(window)
  641. positions.append(position) # Same pos for all windows from this trial
  642. labels.append(self.labels[i])
  643. datas_hr = torch.tensor(np.array(datas_hr), dtype=torch.float32) # (N_windows_total, C, W)
  644. positions = torch.stack(positions) # (N_windows_total, C, 3)
  645. self.labels = torch.tensor(np.array(labels), dtype=torch.int32).cpu()
  646. return datas_hr, positions
  648. def __len__(self):
  649. return len(self.datas_hr)
  651. def __getitem__(self, idx):
  652. device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
  653. hr_data = self.datas_hr[idx]
  654. pos = self.positions[idx]
  655. label = self.labels[idx]
  656. if self.sr_type == "temporal":
  657. lr_data = self._downsample(hr_data, self.multiplier)
  658. else: # spatial
  659. if isinstance(hr_data, torch.Tensor):
  660. lr_data = hr_data.clone()
  661. else:
  662. lr_data = hr_data.copy() # numpy array
  663. lr_data = get_lr_data_spatial(lr_data, dataset_name=self.dataset_name, sr_ratio=self.multiplier)
  664. lr_data = torch.tensor(lr_data, dtype=torch.float32)
  665. return lr_data.to(device), hr_data.to(device), pos.to(device), label.to(device)
  667. def train_test_val_split_patients(patients, test_size=0.2, val_size=0.1, random_state=42):
  669. train_patients, test_patients = train_test_split(patients, test_size=test_size, random_state=random_state)
  670. relative_val_size = val_size / (1 - test_size)
  671. train_patients, val_patients = train_test_split(train_patients, test_size=relative_val_size, random_state=random_state)
  673. return train_patients, val_patients, test_patients
  675. def train_test_val_split(windows, labels, positions, test_size=0.2, val_size=0.1, random_state=42):
  677. train_data, test_data, train_labels, test_labels, train_positions, test_positions = train_test_split(
  678. windows, labels, positions, test_size=test_size, random_state=random_state, stratify=labels
  679. )
  680. relative_val_size = val_size / (1 - test_size)
  681. train_data, val_data, train_labels, val_labels, train_positions, val_positions = train_test_split(
  682. train_data, train_labels, train_positions, test_size=relative_val_size, random_state=random_state, stratify=train_labels
  683. )
  685. return (train_data, train_labels, train_positions), (val_data, val_labels, val_positions), (test_data, test_labels, test_positions)
  687. def plot_mean_timeseries(timeseries, save_path=None):
  689. fig = plt.figure(figsize=(12, 4))
  690. for key, value in timeseries.items():
  691. if value.ndim == 3:
  692. value = value[0] # Take the first sample in the batch
  693. print(f"Plotting {key} with shape {value.shape}, min {value.min()}, max {value.max()}")
  694. mean_signal = np.mean(value, axis=0) # Mean across channels
  695. plt.plot(mean_signal, label=key)
  696. plt.title(f'Mean Timeseries Across Channels')
  697. plt.xlabel('Time (samples)')
  698. plt.ylabel('Amplitude')
  699. plt.grid()
  700. plt.legend()
  701. if save_path:
  702. fig.savefig(save_path, dpi=300)
  703. print(f"Saved figure to {save_path}")
  704. plt.close(fig)
  706. def clear_directory(path, ignore=[]):
  707. """Remove all files in the specified directory."""
  708. if os.path.exists(path):
  709. for filename in os.listdir(path):
  710. if filename == '.gitignore' or filename in ignore:
  711. continue # Skip .gitignore files and ignored files/folders
  712. file_path = os.path.join(path, filename)
  713. try:
  714. if os.path.isfile(file_path) or os.path.islink(file_path):
  715. os.unlink(file_path)
  716. elif os.path.isdir(file_path):
  717. shutil.rmtree(file_path)
  718. except Exception as e:
  719. print(f'Failed to delete {file_path}. Reason: {e}')
  721. def add_zero_channels(input_tensor, target_channels=64, dataset_name="mmi", multiplier=2):
  722. """Add zero channels to input_tensor to match target_channels."""
  724. if input_tensor.ndim == 2:
  725. batch_size = None
  726. nchs = input_tensor.size(0)
  727. length = input_tensor.size(1)
  728. else:
  729. batch_size = input_tensor.size(0)
  730. nchs = input_tensor.size(1)
  731. length = input_tensor.size(2)
  733. if batch_size is None:
  734. input_target = torch.zeros((target_channels, length), device=input_tensor.device)
  735. else:
  736. input_target = torch.zeros((batch_size, target_channels, length), device=input_tensor.device)
  738. channels_to_use = unmask_channels[dataset_name][f"x{multiplier}"]
  739. for i, ch in enumerate(channels_to_use):
  740. if batch_size is None:
  741. input_target[ch, :] = input_tensor[i, :]
  742. else:
  743. input_target[:, ch, :] = input_tensor[:, i, :]
  744. return input_target
  746. def compute_latents(model, dataloader, device, split = "train", map_labels=None):
  748. """
  749. Computes latent vectors from model for all data in dataloader.
  751. Args:
  752. model: PyTorch model with return_latent=True support
  753. dataloader: yields (eeg_lr, eeg_hr, pos, label)
  754. device: torch device
  755. """
  757. latent_vectors = []
  758. labels = []
  760. print("Computing latents for training set", end='\r')
  761. for i, (eeg_lr, eeg_hr, pos, label) in enumerate(dataloader):
  762. print(f"Processing batch {i+1}/{len(dataloader)}", end='\r')
  764. eeg_lr = eeg_lr.to(device)
  765. eeg_hr = eeg_hr.to(device)
  766. pos = pos.to(device)
  767. label = label.to(device) # In case needed
  769. with torch.no_grad():
  770. if model.__class__.__name__ == "DiBiMa_Diff":
  771. batch_size = eeg_lr.size(0)
  772. t = torch.full((batch_size,), model.train_scheduler.num_train_timesteps - 1,
  773. device=device, dtype=torch.long)
  774. x_t_hr = torch.randn_like(eeg_hr).to(device)
  775. latent = model(x_t_hr, t, lr=eeg_lr, pos=pos, label=label, return_latent=True)[-1]
  776. else:
  777. latent = model(eeg_lr, return_latent=True)[-1]
  779. # Flatten to (B, D)
  780. latent = latent.reshape(latent.size(0), -1).cpu().numpy()
  781. latent_vectors.append(latent)
  783. # Collect labels per sample
  784. for l in label.flatten():
  785. labels.append(map_labels[l.item()] if map_labels else l.item())
  787. # Memory management
  788. del eeg_lr, eeg_hr, pos, label, latent
  789. torch.cuda.empty_cache()
  790. gc.collect()
  792. return np.vstack(latent_vectors), np.array(labels)
  794. def plot_umap_latent_space(model, dataloader_train, dataloader_test, save_path=None, map_labels=None, seed=42):
  795. """
  796. Plots UMAP projection of model latent space colored by labels.
  798. Args:
  799. model: PyTorch model with return_latent=True support
  800. dataloader_train: yields (eeg_lr, eeg_hr, pos, label)
  801. dataloader_test: yields (eeg_lr, eeg_hr, pos, label)
  802. save_path: Optional save path for figure
  803. map_labels: Optional dict {int: str} for label mapping
  804. seed: Random seed for reproducibility
  805. """
  806. # Collect all latent vectors and labels
  807. latent_vectors_train = []
  808. latent_vectors_test = []
  809. labels_train = []
  810. labels_test = []
  812. model.eval()
  813. device = next(model.parameters()).device
  815. print("Computing latents for training set...")
  816. latent_vectors_train, labels_train = compute_latents(model, dataloader_train, device, split="train", map_labels=map_labels)
  817. print("\nComputing latents for test set...")
  818. latent_vectors_test, labels_test = compute_latents(model, dataloader_test, device, split="test", map_labels=map_labels)
  820. print(f"UMAP on {latent_vectors_train.shape[0]} samples, {latent_vectors_train.shape[1]} dims")
  822. # Fit UMAP ONCE on full dataset
  823. reducer = UMAP(n_neighbors=15, min_dist=0.1, metric='euclidean',
  824. random_state=seed, n_jobs=-1)
  825. embedding_train = reducer.fit_transform(latent_vectors_train)
  826. embedding_test = reducer.transform(latent_vectors_test)
  828. # Plot by unique labels
  829. plt.figure(figsize=(12, 8))
  830. u_labels = np.unique(labels_test)
  832. print("Generating scatter plot...")
  833. for ul in u_labels:
  834. mask = labels_test == ul
  835. if np.sum(mask) > 0:
  836. plt.scatter(embedding_test[mask, 0], embedding_test[mask, 1],
  837. label=str(ul), alpha=0.6, s=20)
  838. print(f" {ul}: {np.sum(mask)} points")
  840. plt.title('UMAP Latent Space Projection', fontsize=14)
  841. plt.xlabel('UMAP 1', fontsize=12)
  842. plt.ylabel('UMAP 2', fontsize=12)
  843. plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
  844. plt.grid(True, alpha=0.3)
  845. plt.tight_layout()
  847. if save_path:
  848. plt.savefig(save_path, dpi=300, bbox_inches='tight')
  849. print(f"Saved: {save_path}")
  851. plt.show()
  852. plt.close()
  853. print("UMAP visualization complete!")
  856. def tensor2raw(eeg_tensor, info):
  857. """
  858. Convert a PyTorch tensor back to an MNE Raw object.
  859. """
  860. eeg_data = eeg_tensor.cpu().numpy()
  861. info = mne.create_info(ch_names=info['ch_names'], sfreq=info['sfreq'], ch_types='eeg')
  862. raw = mne.io.RawArray(eeg_data, info)
  863. return raw
  865. import torch
  866. import torch.nn as nn
  867. import torch.nn.functional as F
  868. class ReconstructionLoss(nn.Module):
  869. def __init__(
  870. self,
  871. lambda_spectral: float = 0.1,
  872. lambda_l2: float = 1e-6, # Reduced
  873. freq_bands: dict = None
  874. ):
  875. super().__init__()
  876. self.lambda_spectral = lambda_spectral
  877. self.lambda_l2 = lambda_l2
  879. # EEG frequency bands (Hz)
  880. self.freq_bands = freq_bands or {
  881. 'delta': (0.5, 4),
  882. 'theta': (4, 8),
  883. 'alpha': (8, 13),
  884. 'beta': (13, 30),
  885. 'gamma': (30, 45)
  886. }
  888. def spectral_loss(self, pred, target, sample_rate=250):
  889. """Preserve spectral power in EEG bands"""
  890. pred_fft = torch.fft.rfft(pred, dim=-1)
  891. target_fft = torch.fft.rfft(target, dim=-1)
  893. # Power spectral density
  894. pred_psd = torch.abs(pred_fft) ** 2
  895. target_psd = torch.abs(target_fft) ** 2
  897. return F.mse_loss(pred_psd, target_psd)
  899. def forward(self, pred, target, model: nn.Module):
  900. # Base reconstruction
  901. mse_loss = F.mse_loss(pred, target)
  903. # Spectral preservation
  904. spectral_loss = self.spectral_loss(pred, target)
  906. # L2 only on convolutional/linear weights, exclude biases/norms
  907. l2_reg = sum(
  908. torch.sum(p ** 2)
  909. for name, p in model.named_parameters()
  910. if 'weight' in name and len(p.shape) >= 2
  911. )
  913. total_loss = (
  914. mse_loss
  915. + self.lambda_spectral * spectral_loss
  916. + self.lambda_l2 * l2_reg
  917. )
  919. return total_loss
  921. def random_matplotlib_color():
  922. """Generate single random RGB color (0-1 range)."""
  923. return tuple(random.random() for _ in range(3))
  925. def generate_colors(n_colors: int = 1, method: str = 'hsv_uniform'):
  926. """
  927. Generate visually distinct colors for 64-channel signals (ECG/PPG).
  929. Args:
  930. n_colors: Number of colors (default 1)
  931. method: 'hsv_uniform' (recommended), 'random', or 'tab20'
  933. Returns:
  934. List of (R,G,B) tuples, matplotlib-ready
  935. """
  936. if method == 'tab20':
  937. # Matplotlib's built-in qualitative colormap (20 colors, repeats)
  938. cmap = plt.cm.tab20(np.linspace(0, 1, n_colors))
  939. return [tuple(color[:3]) for color in cmap]
  941. elif method == 'random':
  942. # Pure random (may have clashes)
  943. return [random_matplotlib_color() for _ in range(n_colors)]
  945. else: # 'hsv_uniform' - best perceptual separation
  946. colors = []
  947. for i in range(n_colors):
  948. hue = i / n_colors
  949. sat = np.clip(0.7 + 0.2 * np.sin(i * np.pi / 8), 0.6, 1.0)
  950. val = np.clip(0.9 + 0.05 * np.sin(i * np.pi / 4), 0.85, 1.0)
  951. color = plt.cm.hsv(hue)[:3]
  952. color = tuple(np.array(color) * np.array([1, sat, val]))
  953. colors.append(color)
  954. return colors
  956. class EEGSuperResolutionLoss(nn.Module):
  957. def __init__(self,
  958. lambda_grad=1.0,
  959. lambda_corr=0.5,
  960. lambda_freq=0.3,
  961. use_freq_loss=True):
  962. """
  963. Composite loss for EEG temporal super-resolution
  965. Args:
  966. lambda_grad: Weight for gradient loss (controls sharp transitions)
  967. lambda_corr: Weight for Pearson correlation loss (temporal coherence)
  968. lambda_freq: Weight for frequency domain loss (spectral fidelity)
  969. use_freq_loss: Whether to include frequency domain loss
  970. """
  971. super().__init__()
  972. self.lambda_grad = lambda_grad
  973. self.lambda_corr = lambda_corr
  974. self.lambda_freq = lambda_freq
  975. self.use_freq_loss = use_freq_loss
  976. self.mse = nn.MSELoss()
  977. self.__name__ = "EEGSuperResolutionLoss"
  979. def gradient_loss(self, pred, target):
  980. """
  981. Gradient difference loss - captures rate of change
  982. Helps preserve sharp transitions
  983. """
  984. # Temporal gradient (first derivative along time axis)
  985. pred_grad = pred[:, :, 1:] - pred[:, :, :-1]
  986. target_grad = target[:, :, 1:] - target[:, :, :-1]
  988. return F.mse_loss(pred_grad, target_grad)
  990. def pearson_correlation_loss(self, pred, target):
  991. """
  992. Pearson correlation coefficient loss
  993. Preserves temporal waveform patterns
  994. """
  995. # Flatten spatial dimensions, keep batch and time
  996. pred_flat = pred.reshape(pred.size(0), -1)
  997. target_flat = target.reshape(target.size(0), -1)
  999. # Center the data
  1000. pred_centered = pred_flat - pred_flat.mean(dim=1, keepdim=True)
  1001. target_centered = target_flat - target_flat.mean(dim=1, keepdim=True)
  1003. # Covariance
  1004. cov = (pred_centered * target_centered).sum(dim=1)
  1006. # Standard deviations
  1007. pred_std = torch.sqrt((pred_centered ** 2).sum(dim=1) + 1e-8)
  1008. target_std = torch.sqrt((target_centered ** 2).sum(dim=1) + 1e-8)
  1010. # Pearson correlation coefficient
  1011. pcc = cov / (pred_std * target_std + 1e-8)
  1013. # Loss is 1 - PCC (want to maximize correlation)
  1014. return (1 - pcc).mean()
  1016. def frequency_domain_loss(self, pred, target):
  1017. """
  1018. FFT-based frequency domain loss
  1019. Preserves spectral characteristics
  1020. """
  1021. # Apply FFT along time dimension
  1022. pred_fft = torch.fft.rfft(pred, dim=-1)
  1023. target_fft = torch.fft.rfft(target, dim=-1)
  1025. # Compute magnitude spectrum loss
  1026. pred_mag = torch.abs(pred_fft)
  1027. target_mag = torch.abs(target_fft)
  1028. mag_loss = F.mse_loss(pred_mag, target_mag)
  1030. # Compute phase loss (optional, helps with temporal alignment)
  1031. pred_phase = torch.angle(pred_fft)
  1032. target_phase = torch.angle(target_fft)
  1033. phase_loss = F.mse_loss(pred_phase, target_phase)
  1035. return mag_loss + 0.1 * phase_loss
  1037. def forward(self, pred, target):
  1038. """
  1039. Combined loss function
  1041. Args:
  1042. pred: Predicted HR EEG (B, C, T)
  1043. target: Ground truth HR EEG (B, C, T)
  1044. """
  1045. # Base MSE loss
  1046. loss_mse = self.mse(pred, target)
  1048. # Gradient loss for sharp transitions
  1049. loss_grad = self.gradient_loss(pred, target)
  1051. # Pearson correlation for temporal coherence
  1052. loss_corr = self.pearson_correlation_loss(pred, target)
  1054. # Frequency domain loss
  1055. if self.use_freq_loss:
  1056. loss_freq = self.frequency_domain_loss(pred, target)
  1057. else:
  1058. loss_freq = 0.0
  1060. # Combined loss
  1061. total_loss = (loss_mse +
  1062. self.lambda_grad * loss_grad +
  1063. self.lambda_corr * loss_corr)
  1065. if self.use_freq_loss:
  1066. total_loss += self.lambda_freq * loss_freq
  1068. # Return total and individual losses for monitoring
  1069. return total_loss, {
  1070. 'mse': loss_mse.item(),
  1071. 'gradient': loss_grad.item(),
  1072. 'correlation': loss_corr.item(),
  1073. 'frequency': loss_freq.item() if self.use_freq_loss else 0.0,
  1074. 'total': total_loss.item()
  1075. }
  1078. def detect_and_clean_seed_trial(raw, reject_threshold=0.16):
  1079. """Version-safe bad channel detection for SEED"""
  1081. # Basic stats without montage dependency
  1082. data = raw.get_data()
  1084. # Flat channels: variance near zero
  1085. variances = data.var(axis=1)
  1086. flat_idx = np.where(variances < 1e-12)[0]
  1088. # Noisy outliers: z-score variance >5
  1089. var_z = (variances - variances.mean()) / (variances.std() + 1e-8)
  1090. noisy_idx = np.where(var_z > 5)[0]
  1092. # Saturated: low first differences (repeats)
  1093. diffs = np.abs(np.diff(data, axis=1)).mean(axis=1)
  1094. sat_idx = np.where(diffs < 1e-10)[0]
  1096. all_bads = list(set(list(flat_idx) + list(noisy_idx) + list(sat_idx)))
  1097. raw.info['bads'] = [raw.ch_names[i] for i in all_bads]
  1099. n_bad_frac = len(all_bads) / raw.info['nchan']
  1100. #print(f"Trial bad fraction: {n_bad_frac:.1%} ({len(all_bads)}/{raw.info['nchan']})")
  1102. if n_bad_frac > reject_threshold:
  1103. return None
  1105. # Interpolate (works without montage if no spatial ops needed)
  1106. try:
  1107. raw.interpolate_bads(reset_bads=False)
  1108. except Exception as e:
  1109. print(f"Interpolation failed (ok for DL): {e}")
  1110. # Still return - bads marked for your model to ignore
  1112. return raw
  1114. if __name__ == "__main__":
  1116. for dataset_name in ["mmi", "seed"]:
  1117. sr_type = "spatial"
  1118. for sr_ratio in ["x2", "x4", "x8"]:
  1119. channels = unmask_channels[dataset_name][sr_ratio]
  1120. print(f"{dataset_name} {sr_type} {sr_ratio}: {len(channels)} channels -> {channels}")

utils.py at commit 8a666bd, under MIT · at the source

Overview

Authors: Ugo Lomoio1,2, Pietro Lió3, Pietro Hiram Guzzi1, Pierangelo Veltri2
ORCID iDs: Pietro Lió, Pietro Hiram Guzzi
  1. Department of Surgical and Medical Sciences, Magna Graecia University, Catanzaro 88100, Italy
  2. DIMES, University of Calabria, Rende 87036, Italy
  3. Computer Science and Technology, University of Cambridge, Cambridge CB2 1TN, United Kingdom
Institutions: Magna Graecia University (Italy); University of Calabria (Italy); University of Cambridge (United Kingdom)
Journal: Bioinformatics (Oxford, England), volume 42, issue 5, article btag169
Dates: received 15 February 2026; accepted 25 March 2026; published online 15 April 2026; in print May 2026
Type: Research article · Language: English
License: CC BY
Identifiers: DOI 10.1093/bioinformatics/btag169 · PMID 41984820 · PMCID PMC13143424 · OpenAlex W7154494588
Open access: gold, a free copy (OpenAlex)
Status: code verified
Categories: EEG (modality), human (organism), methods / tools (subfield)
Methods: Spectral & time-frequency, Preprocessing, Connectivity, Machine learning
MeSH: Electroencephalography*, Signal Processing, Computer-Assisted*, Algorithms, Brain, Humans (* major topic)
Topic: Functional Brain Connectivity Studies (Cognitive Neuroscience, Neuroscience), according to OpenAlex
Citations: cited by 1 paper (Europe PMC); 22 references in the paper

Abstract

Motivations: Electroencephalography (EEG) is a non-invasive method that records brain electrical activity from scalp electrodes, offering millisecond temporal resolution but limited spatial detail due to sparse sensor layouts.

Results: We present DiBiMa-EEGSR, a bidirectional Mamba-2 diffusion framework for spatio-temporal EEG super-resolution that reconstructs high-resolution signals from standard low-density recordings without additional hardware. The method formulates super-resolution as conditional generative inference and integrates a diffusion process with a bidirectional state-space backbone to model long-range temporal dependencies with linear complexity. Conditioning on low-resolution inputs, electrode positions and task labels enables anatomically coherent and context-aware reconstruction. A one-step sampling strategy substantially reduces inference time while preserving fidelity. Across two public benchmarks, the approach improves reconstruction accuracy, spatial coherence and spectral preservation over convolutional, transformer-based and prior diffusion models in both spatial and temporal upsampling tasks, providing a scalable pathway toward high-resolution electrophysiological imaging.

Availability and implementation: Code to reproduce ablation experiments, training and evaluation of the proposed BiMa and DiBiMa EEGSR models are available at https://github.com/UgoLomoio/DiBiMa-EEGSR.git. Model weights are available at https://huggingface.co/Ugo96/DiBiMa-EEGSR while an interactive demo for EEG spatial super-resolution using our models can be found at https://huggingface.co/spaces/Ugo96/DiBiMa-EEGSR-Demo.

Reproduced under the paper's license (CC BY), from the paper cited above.

Repositories

Its files are read in the Code ↔ Paper reader above, with 18 matches between paragraphs and lines of code.

UgoLomoio/DiBiMa-EEGSR

License: MIT
State: the link answers, verified on 28 September 2026
Evidence: files inventoried
Commit: 8a666bdf5d8af48d8cf7b657f2e7afa2d03f9093, 18 February 2026
Languages: Python (26), Jupyter (1)
Size: 140 files, 27 scripts
Software Heritage: not archived
Found in: “Code availability”
Holds: README, license file, environment (requirements.txt), 1 notebook
Not found: CITATION.cff, tests, continuous integration, documentation
Tools: PyTorch (25 files), NumPy (14 files), scikit-learn (13 files), Matplotlib (12 files), MNE-Python (12 files), pandas (10 files), PyTorch Lightning (9 files), SciPy (8 files), Plotly (2 files), UMAP (2 files)
Availability: 1 check, the latest on 28 September 2026: the link answers
  • 28 September 2026: the link answers
29 files
At the source: github.com/UgoLomoio/DiBiMa-EEGSR.git

huggingface.co/ugo96/dibima-eegsr

License: MIT
State: the link answers, verified on 28 September 2026
Evidence: files inventoried
Commit: 59fb893083e774ce0ab7c13a943710520313526e, 10 February 2026
Size: 26 files, 0 scripts
Software Heritage: not archived
Found in: “Code availability”
Holds: README
Not found: license file, CITATION.cff, environment file, tests, continuous integration, documentation
Availability: 1 check, the latest on 28 September 2026: the link answers
  • 28 September 2026: the link answers
1 file
At the source: huggingface.co/Ugo96/DiBiMa-EEGSR

huggingface.co/spaces/ugo96/dibima-eegsr-demo

License: MIT
State: the link answers, verified on 28 September 2026
Evidence: files inventoried
Commit: 1be1ea39e18f068e49b362bb283857d6b994684f, 11 February 2026
Languages: Python (92), CUDA (9), C/C++ (3), C++ (1)
Size: 163 files, 105 scripts
Software Heritage: not archived
Found in: “Code availability”
Holds: README, environment (Dockerfile, requirements.txt, mamba_install/setup.py, maser/requirements.txt), tests
Not found: license file, CITATION.cff, continuous integration, documentation
Availability: 1 check, the latest on 28 September 2026: the link answers
  • 28 September 2026: the link answers
At the source: huggingface.co/spaces/Ugo96/DiBiMa-EEGSR-Demo

Availability and implementation

Code to reproduce ablation experiments, training and evaluation of the proposed BiMa and DiBiMa EEGSR models are available at https://github.com/UgoLomoio/DiBiMa-EEGSR.git. Model weights are available at https://huggingface.co/Ugo96/DiBiMa-EEGSR while an interactive demo for EEG spatial super-resolution using our models can be found at https://huggingface.co/spaces/Ugo96/DiBiMa-EEGSR-Demo.

Reproduced under the paper's license (CC BY), from the paper cited above.

Tracing map

Proposed by the machine: these links were found in the paper and verified at the source, without human review. The map will receive a Zenodo DOI once one of the paper's authors has validated it with their ORCID.

What the map holds:

  • 3 repositories of the authors' code, each at its verified commit, with its license and how the link was found in the paper;
  • 27 scripts, each with its path and the digest of its content;
  • 18 matches between paragraphs of the paper and lines of the code (method lexical-v1);
  • neither the text of the paper nor the code itself.

Its JSON (tracing-map.json) is deposited on Zenodo with its DOI once the map is validated.

Data

Datasets cited

Data availability

Physionet M/MI is an open-source dataset freely available at: https://physionet.org/content/eegmmidb/1.0.0/. SEED dataset is available upon request at: https://bcmi.sjtu.edu.cn/home/seed/downloads.html

Reproduced under the paper's license (CC BY), from the paper cited above.

Versions

The history of this record: each version stored by the harvester or made by a correction of its authors or of the maintainers of its code, and what changed in its facts. The texts of the paper (its abstract, its availability statements) are not part of it; versions that changed only those are not listed.

Version 1, 28 September 2026: the first record

Recorded: type, language, journal, volume, issue, pages, dates, 4 authors, 5 MeSH terms, 17 references.

Cite

This paper

Lomoio, U., Lió, P., Guzzi, P. H., & Veltri, P. (2026). Bidirectional Mamba-2 boosts EEG super-resolution via regression and diffusion. Bioinformatics (Oxford, England), 42(5), btag169. https://doi.org/10.1093/bioinformatics/btag169

BibTeX

@article{lomoio2026bidirectional,
author = {Lomoio, Ugo and Lió, Pietro and Guzzi, Pietro Hiram and Veltri, Pierangelo},
title = {{Bidirectional Mamba-2 boosts EEG super-resolution via regression and diffusion}},
journal = {Bioinformatics (Oxford, England)},
year = {2026},
month = may,
volume = {42},
number = {5},
pages = {btag169},
publisher = {Oxford University Press},
issn = {1367-4803},
doi = {10.1093/bioinformatics/btag169},
url = {https://doi.org/10.1093/bioinformatics/btag169},
pmid = {41984820},
pmcid = {PMC13143424}
}

RIS

TY - JOUR
AU - Lomoio, Ugo
AU - Lió, Pietro
AU - Guzzi, Pietro Hiram
AU - Veltri, Pierangelo
TI - Bidirectional Mamba-2 boosts EEG super-resolution via regression and diffusion
T2 - Bioinformatics (Oxford, England)
J2 - Bioinformatics
PY - 2026
DA - 2026/05/01
VL - 42
IS - 5
SP - btag169
SN - 1367-4803
PB - Oxford University Press
DO - 10.1093/bioinformatics/btag169
UR - https://doi.org/10.1093/bioinformatics/btag169
LA - en
ER -

CSL-JSON

{
"id": "10.1093/bioinformatics/btag169",
"type": "article-journal",
"title": "Bidirectional Mamba-2 boosts EEG super-resolution via regression and diffusion",
"container-title": "Bioinformatics (Oxford, England)",
"author": [
{
"family": "Lomoio",
"given": "Ugo"
},
{
"family": "Lió",
"given": "Pietro"
},
{
"family": "Guzzi",
"given": "Pietro Hiram"
},
{
"family": "Veltri",
"given": "Pierangelo"
}
],
"container-title-short": "Bioinformatics",
"volume": "42",
"issue": "5",
"page": "btag169",
"DOI": "10.1093/bioinformatics/btag169",
"PMID": "41984820",
"PMCID": "PMC13143424",
"ISSN": "1367-4803",
"publisher": "Oxford University Press",
"URL": "https://doi.org/10.1093/bioinformatics/btag169",
"language": "en",
"issued": {
"date-parts": [
[
2026,
5,
1
]
]
}
}

The tracing map gets a citation of its own once an author has validated it and it has a DOI.

Similar papers

The papers with a page that share the most with this one: the tools found in their code, their categories, datasets, cited references and authors, the rarest counting most.

[1] doi:10.1007/s10916-026-02374-5 [code]
Attention-Enhanced U-Net for Sensor-Efficient High-Density EEG Reconstruction in Wearable Brain Monitoring Systems.
Journal: Journal of medical systems
In common: MNE-Python, PyTorch, scikit-learn, 3 other tools, methods / tools, EEG, 5 references
[2] doi:10.1038/s41598-026-56070-y [code]
SSDLabeler: realistic semi-synthetic data generation for multi-label artifact classification in EEG.
Journal: Scientific reports
In common: PyTorch Lightning, PyTorch, scikit-learn, 3 other tools, physionet.org/content/eegmmidb, EEG
[3] doi:10.1038/s42003-026-10957-8 [code]
Brain defence by the extracellular matrix protein Cochlin.
Journal: Communications biology
In common: PyTorch Lightning, UMAP, Plotly, 6 other tools
[4] doi:10.1038/s41467-026-73996-z [code]
Genetic architecture of white matter microstructure captured by unsupervised deep representation learning of fractional anisotropy maps.
Journal: Nature communications
In common: PyTorch Lightning, UMAP, Plotly, 6 other tools
[5] doi:10.1523/eneuro.0023-26.2026 [code]
Real-Time Segmentation and Classification of Birdsong Syllables for Learning Experiments.
Journal: eNeuro
In common: PyTorch Lightning, UMAP, Plotly, 6 other tools
[6] doi:10.1007/s12021-026-09803-3 [code]
NeuroFusion: A Unified Framework for Generalized Visual Stimulus Decoding from fMRI Across Datasets and Subjects.
Journal: Neuroinformatics
In common: PyTorch Lightning, MNE-Python, PyTorch, 5 other tools, methods / tools
[7] doi:10.1162/imag.a.1207 [code]
Investigating the temporal dynamics and modeling of mid-level feature representations in humans.
Journal: Imaging neuroscience (Cambridge, Mass.)
In common: PyTorch Lightning, MNE-Python, PyTorch, 5 other tools, EEG
[8] doi:10.3390/bioengineering13080924 [code]
Deep Learning-Based Temporal Gait Analysis Using a Smartphone IMU in Older Adults with and Without Non-Specific Low Back Pain.
Journal: Bioengineering (Basel, Switzerland)
In common: PyTorch Lightning, Plotly, PyTorch, 5 other tools, methods / tools
[9] doi:10.1038/s41597-025-05174-7 [code]
A large-scale MEG and EEG dataset for object recognition in naturalistic scenes
Journal: n/a
In common: MNE-Python, Plotly, PyTorch, 5 other tools, methods / tools, EEG
[10] doi:10.1038/s41467-026-76939-w [code]
HIPPIE: a generative model for electrophysiological analysis across species, technologies, and modalities.
Journal: Nature communications
In common: PyTorch Lightning, UMAP, PyTorch, 5 other tools, methods / tools

Contribute

The authors of this paper can claim it, correct its record and validate its tracing map, and the maintainers of its code (its owner, or a public member of its organization) correct what it says of their repository; anyone signed in can ask for its removal. Every request goes to OSCR's own machine, which answers it; your account page follows them.

Sign in with ORCID to claim this paper as one of its authors, correct its record or validate its tracing map: when the paper's metadata lists your ORCID iD, you are recognized at once. Maintainers of its code: sign in with GitHub, then claim the repository on your account page.

Request its removal

To ask OSCR to remove this record, the copies of its authors' scripts or its tracing map, use the removal request page: signed in, you say who you are, what to remove and why, then review and confirm the request. Published rules decide every request (how).

Discussion, reproductions, activity

Discussion: questions and error reports about this paper and its code, from signed-in readers and its authors. It opens with sign-in.

Reproductions: reports from readers who ran the authors' code: what they reproduced, with which environment, commit and data. It opens with sign-in.

Activity: what happens around this paper: new versions of its record, its map's validation, discussions and reproductions. It opens with sign-in.