requirements.txt

# NeuroMamba: A State-Space Foundation Model for Functional MRI
# Core dependencies for training and inference.
#
# ── IMPORTANT install order ───────────────────────────────────────────────
#   1. Install PyTorch FIRST (platform-specific — see README "Installation").
#         pip install torch torchvision torchaudio
#      (On ARM64 / Blackwell, use the matching CUDA wheel index; on Aurora XPU
#       use the vendor stack. See README.)
#   2. pip install -r requirements.txt
#   3. Install the Mamba CUDA kernels LAST, WITHOUT build isolation
#      (their setup.py does `import torch`, which fails under pip build
#       isolation — this is why they are NOT listed in this file):
#         pip install mamba-ssm causal-conv1d --no-build-isolation
# ──────────────────────────────────────────────────────────────────────────

# Deep Learning Framework
# (install separately first — pinned ranges here are informational)
torch>=2.0.0
torchvision>=0.15.0
torchaudio>=2.0.0

# Mamba State-Space Models
#   NOT pip-installable from this file (build isolation has no torch).
#   After torch is installed:
#       pip install mamba-ssm causal-conv1d --no-build-isolation

# Training Infrastructure
pytorch-lightning>=1.9.0,<2.0.0
torchmetrics>=0.11.0
deepspeed>=0.12.0
transformers>=4.30.0
mpi4py>=3.1.0          # required: datasets_hdf5 inits MPI at import time

# Data Processing
h5py>=3.8.0
numpy>=1.24.0,<2.0.0
scipy>=1.10.0
pandas>=2.0.0

# Neuroimaging
nibabel>=5.0.0
nilearn>=0.10.0
monai>=1.2.0

# Modeling utilities
einops>=0.7.0
mup>=1.0.0             # muP scaling (NeuroMamba scaling-law experiments)
psmpy>=0.3.13          # propensity-score matching in data_utils
timm>=0.9.0

# Utilities
tqdm>=4.65.0
pytz>=2023.3

# Visualization
matplotlib>=3.7.0
seaborn>=0.12.0

# Experiment Tracking
wandb>=0.15.0
tensorboard>=2.13.0

# Machine Learning
scikit-learn>=1.3.0

# Configuration & Utilities
PyYAML>=6.0

# Testing (optional)
pytest>=7.3.0

# ── Aurora (ALCF) ONLY — do NOT install on x86/ARM CUDA systems ────────────
# These are provided by the Intel oneAPI / Aurora software stack. The code
# imports them inside try/except, so they are optional everywhere else.
#   intel_extension_for_pytorch
#   oneccl_bindings_for_pytorch