# ============================================================================
# RTX 5090 + CUDA 12.8 Requirements
# ============================================================================
# IMPORTANT: Install PyTorch and PyG packages using the commands below FIRST,
# then install the rest of this file with: pip install -r requirements.txt
# ============================================================================

# STEP 1: Install PyTorch 2.7 stable with CUDA 12.8 support (required for RTX 5090)
# Run this command BEFORE installing this requirements.txt:
# pip3 install torch==2.7.0 torchvision==0.22.0 torchaudio==2.7.0 --index-url https://download.pytorch.org/whl/cu128

# STEP 2: Install PyTorch Geometric and related packages with CUDA 12.8 wheels
# Run this command AFTER installing PyTorch:
# pip install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-2.7.0+cu128.html

# ============================================================================
# Standard packages (install with: pip install -r requirements.txt)
# ============================================================================

# CRITICAL: Pin PyTorch versions to prevent downgrades
torch==2.7.0
torchvision==0.22.0
torchaudio==2.7.0

# PyTorch Geometric (core library, install after the wheels above)
torch_geometric

# E3NN (E(3) equivariant neural networks) - used by NequIP and EquiformerV2
e3nn

# Data science packages
pandas
scikit-learn

# Chemistry/materials packages
mendeleev==0.14.0

# Deep learning frameworks
pytorch-lightning
omegaconf

# OC20 / Transition1x data loading
lmdb
ase
h5py

# Molecular dynamics and materials science
# schnetpack requires --no-deps to avoid PyTorch downgrade on CUDA 12.8:
#   pip install schnetpack>=2.1.0 --no-deps
#   pip install dirsync fasteners torch_ema --no-deps
schnetpack>=2.1.0

# TorchMD-Net, NequIP, GotenNet, PaiNN (schnetpack)
torchmd-net>=2.0
nequip>=0.7.0
gotennet

# Experiment tracking (optional, enable via logging.wandb.enabled=true)
wandb
monty
rdkit