Installation & Data Setup

February 18, 2026 · View on GitHub

← Back to Documentation Index

1. Environment Setup

To set up the environment on Linux (with CUDA 11.7), use Anaconda to create a new environment pxm from environment.yml:

conda env create -f environment.yml
conda activate pxm

Note: If you have a different CUDA version, modify the pytorch-related package versions in environment.yml before creating the environment.

Option A2: Conda for CUDA 12.8

For systems with CUDA 12.8, use the provided base environment file environment_cu128_base.yml:

conda env create -f environment_cu128_base.yml
conda activate pxm_cu128

Then install the CUDA 12.8 PyTorch + PyG stack (copy/paste as a whole):

pip install torch==2.7.0 --index-url https://download.pytorch.org/whl/cu128
pip install lightning torch_geometric
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

Option B: Manual Installation (Pip)

If you need custom versions (e.g., specific CUDA/PyTorch/PyG pins), install packages step by step.

For example, for CUDA 12.6 (Python 3.10):

# PyTorch for CUDA 12.6
pip install torch==2.6.0 --index-url https://download.pytorch.org/whl/cu126
pip install pytorch-lightning==2.6.0
# pip install pytorch-lightning==2.3.0
pip install torch_geometric==2.6.1
# pip install torch_geometric==2.7.0

# PyG extensions (must match torch version + CUDA tag)
pip install torch_scatter torch_sparse torch_cluster -f https://data.pyg.org/whl/torch-2.6.0+cu126.html

# Bio/Chem informatics
pip install biopython==1.83 rdkit==2023.9.3 peptidebuilder==1.1.0
pip install openbabel-wheel==3.1.1.11  # or: conda install -c conda-forge openbabel -y

# Utilities
pip install lmdb==1.7.5 easydict==1.9 numpy==1.24 pandas==1.5.2 scipy==1.10.1
pip install tensorboard==2.20.0  # for training only

2. Data & Model Weights

All training/test data and model weights are available on Zenodo.

For Inference/Sampling (Required)

The model_weights.tar.gz archive contains trained checkpoints. Download and extract it:

wget -c https://zenodo.org/records/17801271/files/model_weights.tar.gz
tar -zxvf model_weights.tar.gz

Creates: data/trained_models/ containing the weights.

(Note: Simple example data is already included in data/examples inside this repo.)

For Benchmarking (Optional)

To run benchmarks on standard test sets (PoseBusters, CrossDocked, etc.), download data_test.tar.gz:

wget -c https://zenodo.org/records/17801271/files/data_test.tar.gz
tar -zxvf data_test.tar.gz

Creates:

  • data/test: Benchmark metadata
  • data/csd, data/geom, data/moad, etc.: Processed test sets.

For Training (Optional)

To train on the reduced demonstration dataset, download data_train_processed.tar.gz:

wget -c https://zenodo.org/records/17801271/files/data_train_processed.tar.gz
tar -zxvf data_train_processed.tar.gz

Creates: data_train with reduced training sets.

Full Dataset Training: The complete processed training data (>500 GB) is not provided as a single archive. To train with the full dataset, follow the instructions in Raw Data Processing to process the raw data yourself.