Skip to content

Commit 33477e5

Browse files
Update README, CHANGELOG, environment
1 parent 633cb6c commit 33477e5

3 files changed

Lines changed: 310 additions & 295 deletions

File tree

CHANGELOG.md

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
# June 5, 2025 (v0.2.0)
2+
### Refactoring and upgrades for PyTorch >= v2.5.1
3+
4+
- Refactored ShEPhERD (aided by Matthew Cox's fork: https://github.com/mcox3406/shepherd/)
5+
- Updated import statements: throughout repo to import directly from `shepherd` assuming local install.
6+
- Fix depreciation warnings:
7+
- `torch.load()` -> `torch.load(weights_only=True)`
8+
- `@torch.cuda.amp.autocast(enabled=False)` -> `@torch.amp.autocast('cuda', enabled=False)`
9+
- Training scripts
10+
- Updated `src/shepherd/datasets.py` for higher versions of PyG. Required changes to the batching functionality for edges (still backwards compatible).
11+
- Slight changes to `training/train.py` for upgraded versions of PyTorch Lightning.
12+
- Model checkpoints have been UPDATED for PyTorch Lightning v2.5.1
13+
- The original checkpoints for PyTorch Lightning v1.2 can be found in previous commits (`c3d5ec0` or before), the original publication Release, or at the Dropbox data link: https://www.dropbox.com/scl/fo/rgn33g9kwthnjt27bsc3m/ADGt-CplyEXSU7u5MKc0aTo?rlkey=fhi74vkktpoj1irl84ehnw95h&e=1&st=wn46d6o2&dl=0
14+
- Created a basic unconditional generation test script
15+
- Updated the environment and relevant files to be compatible with PyTorch >= v2.5.1
16+
- Bug fix for `shepherd.datasets.HeteroDataset.__getitem__` where x3 point extraction should use `get_x2_data`
17+
18+
#### Additional notes
19+
Thank you to Matthew Cox for his contributions in the updated code.
20+
21+
22+
# January 13, 2025
23+
- Added the ability to do partial inpainting for pharmacophores at inference.

README.md

Lines changed: 81 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,14 @@ Note that *ShEPhERD* has a sister repository, [shepherd-score](https://github.co
55

66
The preprint can be found on arXiv: [ShEPhERD: Diffusing shape, electrostatics, and pharmacophores for bioisosteric drug design](https://arxiv.org/abs/2411.04130)
77

8+
### **Important** notice for current repository
9+
This repository has undergone a major refactor to accommodate inference with PyTorch 2.5, primarily for ease-of-use. To maintain reproducibility for training and inference, the original code can be found under commit `c3d5ec0` or the Release titled "Publication code v0.1.0". The model checkpoints used for publication can be found in those binaries or at the following Dropbox [link](https://www.dropbox.com/scl/fo/rgn33g9kwthnjt27bsc3m/ADGt-CplyEXSU7u5MKc0aTo?rlkey=fhi74vkktpoj1irl84ehnw95h&e=1&st=wn46d6o2&dl=0) where training data can also be found. The checkpoints were converted with `python -m pytorch_lightning.utilities.upgrade_checkpoint <chkpt_path>`
10+
Slight changes have also been made to the training code to adhere to Pytorch Lightning >2.0 and new versions of PyTorch Geometric.
11+
12+
We would like to acknowledge Matthew Cox for his contributions in updating this codebase.
13+
814
<p align="center">
9-
<img width="400" src="./shepherd_logo.svg">
15+
<img width="400" src="./docs/images/shepherd_logo.svg">
1016
</p>
1117

1218
<sub><sup>1</sup> **ShEPhERD**: **S**hape, **E**lectrostatics, and **Ph**armacophores **E**xplicit **R**epresentation **D**iffusion</sub>
@@ -23,93 +29,113 @@ The preprint can be found on arXiv: [ShEPhERD: Diffusing shape, electrostatics,
2329

2430
```
2531
.
26-
├── shepherd_score_utils/ # dependencies from shepherd-score Github repository
27-
├── shepherd_chkpts/ # trained model checkpoints (from pytorch lightning)
28-
├── paper_experiments/ # inference scripts for all experiments in preprint
29-
├── samples/ # empty dir to hold outputs from paper_experiments/*.py
30-
├── jobs/ # empty dir to hold ouputs from train.py
31-
├── conformers/ # conditional target structures for experiments, and (sample) training data
32-
├── parameters/ # hyperparameter specifications for all models in preprint
33-
├── model/ # model architecture
34-
│ │ ├── equiformer_operations.py # select E3NN operations from (original) Equiformer
35-
│ │ ├── equiformer_v2_encoder.py # slightly customized Equiformer-V2 module
36-
│ │ └── model.py # module definitions and forward passes
37-
│ ├── utils/ # misc. functions for forward passes
38-
│ ├── egnn/ # customized re-implementation of EGNN
39-
│ └── equiformer_v2/ # clone of equiformer_v2 codebase, with slight modifications for shepherd
40-
├── train.py # main training script
41-
├── lightning_module.py # pytorch-lightning modules and training pipeline
42-
├── datasets.py # torch_geometric dataset class (for training)
43-
├── inference.py # inference functions; see Jupyter notebooks for example uses
44-
├── RUNME_conditional_generation_MOSESaq.ipynb # Jupyter notebook for conditional generation, using MOSES_aq P(x1,x3,x4) model
45-
├── RUNME_unconditional_generation.ipynb # Jupyter notebook for unconditional generation, for all models
32+
├── src/ # source code package
33+
│ └── shepherd/
34+
│ ├── lightning_module.py # pytorch-lightning modules
35+
│ ├── datasets.py # torch_geometric dataset class (for training)
36+
│ ├── inference.py # inference functions
37+
│ ├── extract.py # for extracting field properties
38+
│ ├── shepherd_score_utils/ # dependencies from shepherd-score Github repository
39+
│ └── model/
40+
│ ├── equiformer_operations.py # select E3NN operations from (original) Equiformer
41+
│ ├── equiformer_v2_encoder.py # slightly customized Equiformer-V2 module
42+
│ ├── model.py # module definitions and forward passes
43+
│ ├── utils/ # misc. functions for forward passes
44+
│ ├── egnn/ # customized re-implementation of EGNN
45+
│ └── equiformer_v2/ # clone of equiformer_v2 with slight modifications
46+
├── training/ # training scripts and configs
47+
│ ├── train.py # main training script
48+
│ ├── parameters/ # hyperparameter specifications for all models in preprint
49+
│ └── jobs/ # empty dir to hold outputs from train.py
50+
├── data/
51+
│ ├── shepherd_chkpts/ # trained model checkpoints (from pytorch lightning)
52+
│ └── conformers/ # conditional target structures for experiments, and (sample) training data
53+
├── examples/ # examples and experiments
54+
│ ├── RUNME_conditional_generation_MOSESaq.ipynb # Jupyter notebook for conditional generation
55+
│ ├── RUNME_unconditional_generation.ipynb # Jupyter notebook for unconditional generation
56+
│ ├── basic_inference/ # basic inference example
57+
│ └── paper_experiments/ # inference scripts for all experiments in preprint
58+
├── docs/
59+
│ └── images/
60+
├── docker/ # Docker configuration
61+
│ ├── Dockerfile # Docker image definition
62+
│ └── shepherd_env.yml # conda environment for Docker
63+
├── pyproject.toml # Python project configuration
64+
├── setup.py # package setup script
4665
├── environment.yml # conda environment requirements
66+
├── LICENSE # license file
67+
├── CHANGELOG.md # changelog
4768
└── README.md
4869
```
4970

5071

5172
## Environment
5273

53-
`environment.yml` contains the conda environment that we used for training and running *ShEPhERD*.
54-
55-
**We** followed these steps to create a suitable conda environment, which worked on our Linux system. Please note that this exact installation procedure may depend on your system, particularly your cuda version.
56-
74+
### Requirements
75+
```
76+
python>=3.9
77+
rdkit>=2023.03,<2025.03
78+
torch>=2.5.1
79+
numpy>1.2,<2.0
80+
open3d>=0.18
81+
xtb>=6.6
82+
pandas==2.2.3
5783
```
58-
conda create --name shepherd python=3.8.13
59-
source activate shepherd
60-
conda install merv::envvar-pythonnousersite-true
61-
source deactivate
62-
63-
source activate shepherd
6484

65-
conda config --append channels conda-forge
85+
### Example Environment set-up
6686

67-
pip cache purge
68-
pip3 cache purge
69-
export TMPDIR='/var/tmp'
87+
`environment.yml` contains the updated conda environment for *ShEPhERD* and compatibility with PyTorch >=2.5.
7088

71-
conda install pytorch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 cudatoolkit==11.3.1 -c pytorch
72-
conda install pyg=2.2.0 -c pyg
89+
**We** followed these steps to create a suitable conda environment, which worked on our Linux system. Please note that this exact installation procedure may depend on your system, particularly your cuda version.
7390

74-
pip install e3nn
91+
```
92+
conda create -n shepherd python=3.9
93+
conda activate shepherd
94+
pip install uv
7595
76-
pip install jupyterlab
96+
# download pytorch considering your cuda version
97+
uv pip install torchvision==0.20.1 torchaudio==2.5.1 --index-url https://download.pytorch.org/whl/cu124
98+
uv pip install torch-scatter torch-sparse torch-cluster torch-spline-conv torch-geometric -f https://data.pyg.org/whl/torch-2.5.0+cu124.html
7799
78-
pip install pip==24.0
79-
pip install pytorch-lightning==1.6.3
80-
pip install setuptools==59.5.0
100+
uv pip install pytorch-lightning
101+
uv pip install pandas==2.2.3
81102
82-
pip install rdkit
103+
uv pip install rdkit==2024.09.6 open3d matplotlib jupyterlab
104+
# There may be issues the environment does not set up xTB properly.
105+
# If this is the case, please install from source.
83106
conda install xtb
84-
pip install open3d
85-
conda install h5py
86107
87-
pip install numpy --upgrade
108+
# cd to this repo and do a developer install
109+
# This will install additional requirements found in .toml and not covered above
110+
pip install -e .
88111
```
89112

90113

91114
## Training and inference data
92-
`conformers/` contains the 3D structures of the natural products, PDB ligands, and fragments that we used in our experiments in the preprint. It also includes the 100 test-set structures from GDB-17 that we used in our conditional generation evaluations.
115+
`data/conformers/` contains the 3D structures of the natural products, PDB ligands, and fragments that we used in our experiments in the preprint. It also includes the 100 test-set structures from GDB-17 that we used in our conditional generation evaluations.
93116

94-
`conformers/gdb/example_molblock_charges.pkl` contains *sample* training data from our *ShEPhERD*-GDB-17 training dataset.
95-
`conformers/moses_aq/example_molblock_charges.pkl` contains *sample* training data from our *ShEPhERD*-MOSES_aq training dataset.
117+
`data/conformers/gdb/example_molblock_charges.pkl` contains *sample* training data from our *ShEPhERD*-GDB-17 training dataset.
118+
`data/conformers/moses_aq/example_molblock_charges.pkl` contains *sample* training data from our *ShEPhERD*-MOSES_aq training dataset.
96119

97-
The full training data for both datasets (<10GB each) can be accessed from this Dropbox link: https://www.dropbox.com/scl/fo/rgn33g9kwthnjt27bsc3m/ADGt-CplyEXSU7u5MKc0aTo?rlkey=fhi74vkktpoj1irl84ehnw95h&e=1&st=wn46d6o2&dl=0
120+
The full training data for both datasets (<10GB each) can be accessed from this Dropbox link: [https://www.dropbox.com/scl/fo/rgn33g9kwthnjt27bsc3m/ADGt-CplyEXSU7u5MKc0aTo?rlkey=fhi74vkktpoj1irl84ehnw95h&e=1&st=wn46d6o2&dl=0](https://www.dropbox.com/scl/fo/rgn33g9kwthnjt27bsc3m/ADGt-CplyEXSU7u5MKc0aTo?rlkey=fhi74vkktpoj1irl84ehnw95h&e=1&st=wn46d6o2&dl=0)
98121

99122

100123
## Training
101-
`train.py` is our main training script. It can be run from the command line by specifying a parameter file and a seed. All of our parameter files are held in `parameters/`. As an example, one may re-train the P(x1,x3,x4) model on ShEPhERD-MOSES-aq by calling:
124+
`training/train.py` is our main training script. It can be run from the command line by specifying a parameter file and a seed. All of our parameter files are held in `training/parameters/`. To run training, first `cd` into the `training` directory. As an example, one may re-train the P(x1,x3,x4) model on ShEPhERD-MOSES-aq by calling:
102125

103-
`python train.py params_x1x3x4_diffusion_mosesaq_20240824 0`
126+
```
127+
cd training
128+
python train.py params_x1x3x4_diffusion_mosesaq_20240824 0
129+
```
104130

105-
Note that the trained checkpoints in `shepherd_chkpts/` were obtained after training each model for ~2 weeks on 2 V100 gpus.
131+
The trained checkpoints in `data/shepherd_chkpts/` were obtained after training each model for ~2 weeks on 2 V100 gpus. Note that the checkpoints found in this folder have been converted for PyTorch Lightning v2.5. The original, unmodified checkpoints can be found in the original "Publication" release binaries or at the aforementioned data Dropbox link.
106132

107133

108134
## Inference
109135

110-
The simplest way to run inference is to follow the Jupyter notebooks `RUNME_unconditional_generation_MOSESaq.ipynb` and `RUNME_conditional_generation_MOSESaq.ipynb`.
136+
The simplest way to run inference is to follow the Jupyter notebooks `examples/RUNME_unconditional_generation.ipynb` and `examples/RUNME_conditional_generation_MOSESaq.ipynb`.
111137

112-
`paper_experiments/` also contain scripts that we used to run the experiments in our preprint. Each of these scripts should be copied into the parent directory (same directory as this README) before being called from the command line. Some of the scripts (`paper_experiments/run_inference_*_unconditional_*_.py`) take a few additional command-line arguments, which are detailed in those corresponding scripts by argparse commands.
138+
`examples/paper_experiments/` also contain scripts that we used to run the experiments in our preprint. Some of the scripts (`examples/paper_experiments/run_inference_*_unconditional_*_.py`) take a few additional command-line arguments, which are detailed in those corresponding scripts by argparse commands.
113139

114140
The inference script now supports conditional generation of molecules that contain a superset of the target profile's pharmacophores via partial inpainting. [1/13/2025]
115141

0 commit comments

Comments
 (0)