Skip to content

Commit 65cbd8d

Browse files
Merge pull request #11 from coleygroup/refactor
Refactor and update repo for PyTorch v2.5
2 parents b81042f + 33477e5 commit 65cbd8d

246 files changed

Lines changed: 1261 additions & 728 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

.gitignore

Lines changed: 177 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,177 @@
1+
# macOS
2+
*.DS_Store
3+
4+
# Byte-compiled / optimized / DLL files
5+
__pycache__/
6+
*.py[cod]
7+
*$py.class
8+
9+
# C extensions
10+
*.so
11+
12+
# Distribution / packaging
13+
.Python
14+
build/
15+
develop-eggs/
16+
dist/
17+
downloads/
18+
eggs/
19+
.eggs/
20+
lib/
21+
lib64/
22+
parts/
23+
sdist/
24+
var/
25+
wheels/
26+
share/python-wheels/
27+
*.egg-info/
28+
.installed.cfg
29+
*.egg
30+
MANIFEST
31+
32+
# PyInstaller
33+
# Usually these files are written by a python script from a template
34+
# before PyInstaller builds the exe, so as to inject date/other infos into it.
35+
*.manifest
36+
*.spec
37+
38+
# Installer logs
39+
pip-log.txt
40+
pip-delete-this-directory.txt
41+
42+
# Unit test / coverage reports
43+
htmlcov/
44+
.tox/
45+
.nox/
46+
.coverage
47+
.coverage.*
48+
.cache
49+
nosetests.xml
50+
coverage.xml
51+
*.cover
52+
*.py,cover
53+
.hypothesis/
54+
.pytest_cache/
55+
cover/
56+
57+
# Translations
58+
*.mo
59+
*.pot
60+
61+
# Django stuff:
62+
*.log
63+
local_settings.py
64+
db.sqlite3
65+
db.sqlite3-journal
66+
67+
# Flask stuff:
68+
instance/
69+
.webassets-cache
70+
71+
# Scrapy stuff:
72+
.scrapy
73+
74+
# Sphinx documentation
75+
docs/_build/
76+
77+
# PyBuilder
78+
.pybuilder/
79+
target/
80+
81+
# Jupyter Notebook
82+
.ipynb_checkpoints
83+
84+
# IPython
85+
profile_default/
86+
ipython_config.py
87+
88+
# pyenv
89+
# For a library or package, you might want to ignore these files since the code is
90+
# intended to run in multiple environments; otherwise, check them in:
91+
# .python-version
92+
93+
# pipenv
94+
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
95+
# However, in case of collaboration, if having platform-specific dependencies or dependencies
96+
# having no cross-platform support, pipenv may install dependencies that don't work, or not
97+
# install all needed dependencies.
98+
#Pipfile.lock
99+
100+
# UV
101+
# Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
102+
# This is especially recommended for binary packages to ensure reproducibility, and is more
103+
# commonly ignored for libraries.
104+
#uv.lock
105+
106+
# poetry
107+
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
108+
# This is especially recommended for binary packages to ensure reproducibility, and is more
109+
# commonly ignored for libraries.
110+
# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
111+
#poetry.lock
112+
113+
# pdm
114+
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
115+
#pdm.lock
116+
# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
117+
# in version control.
118+
# https://pdm.fming.dev/latest/usage/project/#working-with-version-control
119+
.pdm.toml
120+
.pdm-python
121+
.pdm-build/
122+
123+
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
124+
__pypackages__/
125+
126+
# Celery stuff
127+
celerybeat-schedule
128+
celerybeat.pid
129+
130+
# SageMath parsed files
131+
*.sage.py
132+
133+
# Environments
134+
.env
135+
.venv
136+
env/
137+
venv/
138+
ENV/
139+
env.bak/
140+
venv.bak/
141+
142+
# Spyder project settings
143+
.spyderproject
144+
.spyproject
145+
146+
# Rope project settings
147+
.ropeproject
148+
149+
# mkdocs documentation
150+
/site
151+
152+
# mypy
153+
.mypy_cache/
154+
.dmypy.json
155+
dmypy.json
156+
157+
# Pyre type checker
158+
.pyre/
159+
160+
# pytype static type analyzer
161+
.pytype/
162+
163+
# Cython debug symbols
164+
cython_debug/
165+
166+
# PyCharm
167+
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
168+
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
169+
# and can be added to the global gitignore or merged into this file. For a more nuclear
170+
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
171+
#.idea/
172+
173+
# Ruff stuff:
174+
.ruff_cache/
175+
176+
# PyPI configuration file
177+
.pypirc

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

File renamed without changes.

conformers/fragment_merging/fragment_merge_condition.pickle renamed to data/conformers/fragment_merging/fragment_merge_condition.pickle

File renamed without changes.
File renamed without changes.
File renamed without changes.
File renamed without changes.
File renamed without changes.
File renamed without changes.

0 commit comments

Comments
 (0)