flash-eeg 0.1.0

GPU-accelerated EEG preprocessing transforms

flash-eeg

GPU-accelerated EEG preprocessing. Drop-in replacement for MNE/Scipy transforms.

Install

pip install flash-eeg

Usage

import flash_eeg as feeg
import torch

x = torch.randn(32, 8, 7500, device="cuda")  # [batch, channels, samples]

# Functional API
spec = feeg.spectrogram(x)     # STFT spectrogram
morlet = feeg.morlet(x)        # Morlet wavelet
conn = feeg.connectivity(x)    # WPLI connectivity
img = feeg.reshape(x)          # Signal to image

# Class API (for training loops)
spec_fn = feeg.Spectrogram(device="cuda")
for batch in dataloader:
    out = spec_fn(batch)

# Both APIs have identical performance after first call.
# Functional API uses lru_cache internally - same params reuse the compiled module.

Options

# Output size
feeg.spectrogram(x, output_size=128)

# FFT parameters
feeg.spectrogram(x, n_fft=512, hop_length=64, freq_min=1.0, freq_max=50.0)
feeg.morlet(x, n_freqs=30, freq_min=1.0, freq_max=100.0)
feeg.connectivity(x, n_fft=1024, num_tapers=5)

# Force compile on/off (auto-detects A100/H100/H200 by default)
feeg.spectrogram(x, compile=True)

Benchmarks

Tested on NVIDIA H200 vs CPU baselines (MNE/Scipy with full parallelization).

Transform	CPU	GPU (compiled)	Speedup
STFT Spectrogram	7.17s	6.3ms	1,147x
Morlet Wavelet	59.2s	41.1ms	1,440x
WPLI Connectivity	72.1s	1.1ms	66,731x
Reshape	6.96s	0.86ms	8,085x

Batch size 1024, 8 channels, 7500 samples (30s @ 250Hz). H200 GPU vs parallelized CPU (MNE/Scipy).

Requirements

• PyTorch >= 2.0
• CUDA GPU (for acceleration)
• scipy (for DPSS tapers in connectivity)

License

MIT