install PyTorch 1.10.0 (e.g., with CUDA 11.3)

July 17, 2026 · View on GitHub

NucleusDiff: Manifold-Constrained Nucleus-Level Denoising Diffusion Model for Structure-Based Drug Design

The official implementation of the PNAS 2025 paper Manifold-Constrained Nucleus-Level Denoising Diffusion Model for Structure-Based Drug Design.

Shengchao Liu^*, Liang Yan^*, Weitao Du, Weiyang Liu, Zhuoxinran Li, Hongyu Guo,
Christian Borgs, Jennifer Chayes, Anima Anandkumar

Proceedings of the National Academy of Sciences (PNAS), 2025
^*Equal contribution

Scan the QR code or join via the Slack invitation link

NucleusDiff inference process

Update & News

2024-06-18 · 🌐 Project page launched. Explore the model, visualizations, paper, and code on the NucleusDiff project page.
2024-09-16 · 📄 arXiv v1 released. The first preprint of NucleusDiff is available as arXiv:2409.10584v1.
2024-09-30 · 🔄 arXiv v2 released. The revised preprint is available as arXiv:2409.10584v2.
2025-08-11 · 🎉 Accepted by PNAS. NucleusDiff was accepted by the Proceedings of the National Academy of Sciences; read the PNAS paper.
2025-09-10 · 💻 Code open sourced. We released the official NucleusDiff implementation as v1.0.0 on GitHub.
2025-10-20 · 📰 Featured by Caltech News. Caltech highlighted NucleusDiff in “New AI Model for Drug Design Brings More Physics to Bear in Predictions”.
2026-07-14 · 💬 SciGenAI Slack channel launched. We opened a dedicated NucleusDiff channel for real-time questions, discussions, code contributions, and collaboration. Join the SciGenAI Slack community—everyone is welcome!
2026-07-16 · 🧹 Codebase cleanup. We cleaned 17 Python files by removing unused imports and variable bindings, organizing imports, simplifying redundant string formatting, removing obsolete comments, and trimming trailing whitespace. This maintenance update does not change the model architecture, numerical computations, command-line defaults, or any data, training, evaluation, and inference logic. See commit 1cde91c.
2026-07-16 · 🤗 Hugging Face mirror released. The pretrained checkpoint and all project data artifacts are now available from the NucleusDiff Hugging Face repository.

NucleusDiff is a manifold-constrained denoising diffusion model for structure-based drug design. It jointly models atomic nuclei and their surrounding electron-cloud manifolds to reduce atomic collisions while generating high-affinity ligands.

Model & Data Downloads

The pretrained checkpoint and data artifacts are hosted in the NucleusDiff Hugging Face repository. Google Drive remains available as a backup mirror.

Install the Hugging Face CLI:

pip install -U huggingface_hub

Run the following commands from the NucleusDiff repository root. The original model/ and data/ directory structure will be preserved.

Download the pretrained checkpoint:

hf download LiangYan3612/NucleusDiff \
   model/nucleusdiff_pretrained_model.pt \
   --local-dir .

Download all data artifacts:

hf download LiangYan3612/NucleusDiff \
   --include "data/*" \
   --local-dir .

Hugging Face repository layout:

Type	Path	Description
Checkpoint	`model/nucleusdiff_pretrained_model.pt`	Pretrained NucleusDiff model
Training data	`data/crossdocked_v1.1_rmsd1.0_pocket10_processed_w_manifold_data_version.lmdb`	Preprocessed CrossDocked manifold LMDB
Data split	`data/crossdocked_pocket10_pose_w_manifold_data_split.pt`	NucleusDiff train/validation/test split
Filtered data	`data/crossdocked_v1.1_rmsd1.0.tar.gz`	Filtered CrossDocked structures
Reference split	`data/split_by_name.pt`	Reference CrossDocked split
Docking test set	`data/test_set.zip`	Protein test set for docking evaluation
Therapeutic targets	`data/real_world.zip`	Therapeutic-target evaluation data
Metadata	`data/affinity_info.pkl`	Affinity metadata
Metadata	`data/test_vina_crossdock_dict.pkl`	CrossDocked Vina evaluation metadata

Update & News
Overview
Model & Data Downloads
Installation
Data Preparation
CrossDocked2020 Experiments
Therapeutic Target Experiments
Universal Inference
SciGenAI Community
Citation

1. Installation

We recommend using the Conda environments below to reproduce our results. For full evaluation or training logs, contact yanliangfdu@gmail.com.

1.1 Main experiment dependencies

The code has been tested in the following environment:

Package	Version
Python	3.8.13
PyTorch	1.12.1
CUDA	11.0
PyTorch Geometric	2.5.2
RDKit	2021.03.1b1

Install via Conda and Pip:

conda create -n "nucleusdiff" python=3.8.13
conda activate nucleusdiff
conda install pytorch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 cudatoolkit=11.3 -c pytorch
pip install torch_geometric
pip install https://data.pyg.org/whl/torch-1.12.0%2Bcu113/pyg_lib-0.3.1%2Bpt112cu113-cp38-cp38-linux_x86_64.whl
pip install https://data.pyg.org/whl/torch-1.12.0%2Bcu113/torch_cluster-1.6.0%2Bpt112cu113-cp38-cp38-linux_x86_64.whl
pip install https://data.pyg.org/whl/torch-1.12.0%2Bcu113/torch_scatter-2.1.0%2Bpt112cu113-cp38-cp38-linux_x86_64.whl
pip install https://data.pyg.org/whl/torch-1.12.0%2Bcu113/torch_sparse-0.6.16%2Bpt112cu113-cp38-cp38-linux_x86_64.whl
pip install https://data.pyg.org/whl/torch-1.12.0%2Bcu113/torch_spline_conv-1.2.1%2Bpt112cu113-cp38-cp38-linux_x86_64.whl
conda install rdkit/label/nightly::rdkit
conda install openbabel tensorboard pyyaml easydict python-lmdb -c conda-forge
pip install wandb
pip install pytorch-lightning==2.1.3
pip install matplotlib
pip install numpy==1.23
pip install accelerate
pip install transformers


# For Vina Docking
pip install meeko==0.1.dev3 scipy pdb2pqr vina==1.2.2
python -m pip install git+https://github.com/Valdes-Tresanco-MS/AutoDockTools_py3

The code should work with PyTorch >= 1.9.0 and PyG >= 2.0. You can adjust the package versions for your environment.

1.2 Manifold preprocessing dependencies

Use this separate environment only if you want to process the CrossDocked manifold dataset from scratch.

# We recommend using conda for environment management
conda create -n Manifold python=3.7.3
conda activate Manifold

pip install -r ./crossdock_manifold_data_preparation/requirements.txt
# install PyMesh for surface mesh processing
PYMESH_PATH="~/PyMesh" # substitute with your own PyMesh path
git clone https://github.com/PyMesh/PyMesh.git $PYMESH_PATH 
cd $PYMESH_PATH 
git submodule update --init
apt-get update
# make sure you have these libraries installed before building PyMesh
apt-get install cmake libgmp-dev libmpfr-dev libgmpxx4ldbl libboost-dev libboost-thread-dev libopenmpi-dev
cd $PYMESH_PATH/third_party
python build.py all # build third party dependencies
cd $PYMESH_PATH
mkdir build
cd build
cmake ..
make -j # check for missing third-party dependencies if failed to make
cd $PYMESH_PATH
python setup.py install
python -c "import pymesh; pymesh.test()"

# install meshplot
conda install -c conda-forge meshplot

# install libigl
conda install -c conda-forge igl

# download MSMS
MSMS_PATH="~/MSMS" # substitute with your own MSMS path
wget https://ccsb.scripps.edu/msms/download/933/ -O msms_i86_64Linux2_2.6.1.tar.gz
mkdir -p $MSMS_PATH # mark this directory as your $MSMS_bin for later use
tar zxvf msms_i86_64Linux2_2.6.1.tar.gz -C $MSMS_PATH

# install PyTorch 1.10.0 (e.g., with CUDA 11.3)
conda install pytorch==1.10.0 cudatoolkit=11.3 -c pytorch -c conda-forge
pip install torch-scatter -f https://data.pyg.org/whl/torch-1.10.0+cu113.html

# install Manifold
pip install -e .

2. Data Preparation

2.1 CrossDocked2020 data

Download the training and evaluation files from the NucleusDiff Hugging Face data directory. The original Google Drive folder is retained as a backup mirror.
To train the model from scratch, download the preprocessed LMDB and split files:
- crossdocked_v1.1_rmsd1.0_pocket10_processed_w_manifold_data_version.lmdb
- crossdocked_pocket10_pose_w_manifold_data_split.pt
To evaluate the model on the test set, you need to download and unzip the test_set.zip. It includes the original PDB files that will be used in Vina Docking.
If you want to process the dataset from scratch, you need to download CrossDocked2020 v1.1 from here, save it into ./data/CrossDocked2020, and run the scripts in ./crossdock_data_preparation:

Process CrossDocked2020 from scratch

step1_clean_crossdocked.py filters the original dataset and keeps entries with RMSD < 1 Å. It generates index.pkl and a directory containing the filtered data (corresponding to crossdocked_v1.1_rmsd1.0.tar.gz in Google Drive). You can skip this step if you downloaded the preprocessed LMDB file.

python ./crossdock_data_preparation/step1_clean_crossdocked.py \
       --source "./data/CrossDocked2020" \
       --dest "./data/crossdocked_v1.1_rmsd1.0" \
       --rmsd_thr 1.0

step2_extract_pockets.py clips each protein file to a 10 Å region around the binding molecule.

python ./crossdock_data_preparation/step2_extract_pockets.py \
       --source "./data/crossdocked_v1.1_rmsd1.0" \
       --dest "./data/crossdocked_v1.1_rmsd1.0_pocket10"

step3_split_pl_dataset.py creates the training and test splits. We use the same split_by_name.pt as AR and Pocket2Mol; it is also available in the Google Drive data folder.

python ./crossdock_data_preparation/step3_split_pl_dataset.py \
       --path "./data/crossdocked_v1.1_rmsd1.0_pocket10" \
       --dest "./data/crossdocked_pocket10_pose_split.pt" \
       --fixed_split "./data/split_by_name.pt"

2.2 CrossDocked manifold data

Activate the manifold preprocessing environment:

conda activate Manifold

Prepare input for MSMS:

python step1_convert_npz_to_xyzrn.py \
       --crossdock_source [path/to/crossdock_pocket10_auxdata/] \
       --out_root "./data/crossdocked_pocket10_mesh"

Run MSMS to generate molecular surfaces:

python step2_compute_msms.py \
       --data_root "./data/crossdocked_pocket10_mesh" \
       --msms-bin [path/to/MSMS/dir]/msms.x86_64Linux2.2.6.1

Refine the surface meshes:

python step3_refine_mesh.py \
       --data_root "./data/crossdocked_pocket10_mesh"

2.3 Generate the final LMDB and split files

python ./datasets/pl_pair_dataset.py \
       --data_root "./data/crossdocked_v1.1_rmsd1.0_pocket10"

3. CrossDocked2020 Experiments

3.1 Training

python train.py \
       --lr 0.001 \
       --device "cuda:0" \
       --wandb_project_name "nucleusdiff_train" \
       --loss_mesh_constained_weight 1

Note: The pretrained checkpoint is available from Hugging Face, with Google Drive retained as a backup mirror.

3.2 Inference

python sample_for_crossdock.py \
       --ckpt_path "./logs_diffusion/nucleusdiff_train" \
       --ckpt_it 100000 \
       --cuda_device 0 \
       --data_id 0

You can also speed up sampling with multiple GPUs, e.g.:

python sample_for_crossdock.py \
       --ckpt_path "./logs_diffusion/nucleusdiff_train" \
       --ckpt_it 100000 \
       --cuda_device 0 \
       --data_id 0

python sample_for_crossdock.py \
       --ckpt_path "./logs_diffusion/nucleusdiff_train" \
       --ckpt_it 100000 \
       --cuda_device 1 \
       --data_id 1 

python sample_for_crossdock.py \
       --ckpt_path "./logs_diffusion/nucleusdiff_train" \
       --ckpt_it 100000 \
       --cuda_device 2 \
       --data_id 2

python sample_for_crossdock.py \
       --ckpt_path "./logs_diffusion/nucleusdiff_train" \
       --ckpt_it 100000 \
       --cuda_device 3 \
       --data_id 3

3.3 General metrics

python ./evaluation/evaluate_for_crossdock_on_general_metrics.py \
        --sample_path "./result_output" \
        --eval_step -1 \
        --protein_root "./data/test_set" \
        --docking_mode "vina_dock"

3.4 Collision metrics

python ./evaluation/evaluate_for_crossdock_on_collision_metrics.py \
        --sample_path "./result_output" \
        --eval_step -1

4. Therapeutic Target Experiments

4.1 Data preparation

If you want to process the dataset from scratch, download data/real_world.zip, extract it into ./data, and run the scripts in ./covid_19_data_preparation:

python ./covid_19_data_preparation/extract_pockets_for_real_world.py \
        --source "./data/real_world" \
        --dest "./real_world_test_extract_pockets"

4.2 Inference

python sample_for_covid_19.py \
        --checkpoint [path/to/nucleusdiff/checkpoint] \
        --pdb_path "./real_world_test_extract_pockets/CDK2/cdk2_ligand_pocket10.pdb" \
        --result_path "./read_world_cdk2_test" \
        --sample_num_atoms "real_world_testing" \
        --inference_num_atoms 30

4.3 General metrics

python ./evaluation/evaluate_for_covid_19_on_general_metrics.py \
        --sample_path "./read_world_cdk2_test" \
        --protein_root "./real_world/cdk2_processed.pdb" \
        --ligand_filename "CDK2" \
        --docking_mode "vina_dock"

4.4 Collision metrics

python ./evaluation/evaluate_for_covid_19_on_collision_metrics.py \
        --sample_path "./read_world_cdk2_test" \
        --model "nucleusdiff_train" \
        --target "cdk2_test"

5. Universal Inference for a Specified Protein

Use sample_for_specific_protein.py to generate ligands for an arbitrary single protein pocket PDB.

5.1 Input preparation

Prepare a pocket PDB centered at the binding site (e.g., 10 Å around the ligand or binding residues).
You may reuse the script in 4.1: ./covid_19_data_preparation/extract_pockets_for_real_world.py.
Example pocket file: ./specific_protein/3cl_ligand_pocket10.pdb.

5.2 Inference

python sample_for_specific_protein.py \
        --checkpoint ./model/nucleusdiff_pretrained_model.pt \
        --pdb_path ./specific_protein/3cl_ligand_pocket10.pdb \
        --result_path ./results_specific_protein \
        --sample_num_atoms real_world_testing \
        --inference_num_atoms 30 \
        --num_samples 1000 \
        --num_steps 1000 \
        --device cuda:0

Key arguments:

--checkpoint: path to a NucleusDiff checkpoint (.pt).
--pdb_path: pocket PDB for your target protein.
--result_path: output directory.
--sample_num_atoms: set to real_world_testing to use a fixed atom count.
--inference_num_atoms: atoms per generated ligand when using real_world_testing.
--num_samples: number of ligands to generate.
--num_steps: diffusion steps (trade-off between quality and speed).
--device: GPU device, e.g., cuda:0.

5.3 Outputs

${result_path}/sample_{test_time}.pt: raw tensors and sampling trajectories.
${result_path}/sdf/*.sdf: reconstructed molecules in SDF format.

Run python sample_for_specific_protein.py --help for the complete list of options and defaults.

SciGenAI Community

Join the NucleusDiff Slack channel

SciGenAI is a community platform for discussion and collaboration around generative AI for science. It connects researchers, students, engineers, and open-source contributors working across NucleusDiff and other AI4Science projects.

The workspace includes a dedicated NucleusDiff channel where you can:

Ask NucleusDiff questions and get real-time community support.
Discuss the paper, implementation details, experiments, and reproducibility.
Coordinate code contributions, commits, issues, code reviews, and pull requests.
Connect with other generative-AI-for-science projects and find collaborators.

Option 1 — Invitation link: Join the SciGenAI Slack community

Option 2 — QR code: Scan the code below with your phone. The QR image is also clickable.

Citation

If you find this work useful, please cite:

@article{liu2025manifold,
  title={Manifold-constrained nucleus-level denoising diffusion model for structure-based drug design},
  author={Liu, Shengchao and Yan, Liang and Du, Weitao and Liu, Weiyang and Li, Zhuoxinran and Guo, Hongyu and Borgs, Christian and Chayes, Jennifer and Anandkumar, Anima},
  journal={Proceedings of the National Academy of Sciences},
  volume={122},
  number={41},
  pages={e2415666122},
  year={2025},
  publisher={National Academy of Sciences}
}