Skip to content

ShEPhERD-score v1.3.1

Choose a tag to compare

View all tags
@kentoabeywardane kentoabeywardane released this 16 Mar 00:26
· 46 commits to main since this release
v1.3.1
faa34a0
This commit was created on GitHub.com and signed with GitHub’s verified signature.
GPG key ID: B5690EEEBB952194
Verified
Learn about vigilant mode.

What's Changed

  • Implements general speed improvements alignment speed. Introduces analytical gradients for PyTorch scoring. by @kentoabeywardane in #20

This PR introduces general speed improvements for alignment speeds particularly for PyTorch.

Alignment

PyTorch

  • Precomputes self-overlaps during optimization (Vol, Vol+Avoid, Vol+ESP, Surf, Surf+ESP, Pharm), which leads to ~1.2-1.7x faster alignments
  • Added analytical gradients in PyTorch scoring (Vol, Vol+Avoid, Vol+ESP, Surf, Surf+ESP, Pharm), which leads to an additional ~1.5-2.5x speedup (depending on the scoring function).
    • Now the default in MoleculePair
  • Did not implement torch.compile since the overhead of compilation was quite high, but might be worth it if batching is introduced.

Note: Currently analytical gradients for extended point-style pharmacophore scoring are not implemented

Jax

  • Precomputes self-overlaps during optimization (Vol, Vol+Avoid) which leads to slight speedup
    • Now the default in MoleculePair
  • Did NOT implement analytical gradients
    • Lead to slower alignments for gaussian overlaps (0.8x speedup).
    • Might be worth it to decrease memory overhead

Note: Vol+Avoid for PyTorch (autograd) vs PyTorch (analytical) vs Jax (autograd) can lead to very slightly different scores/alignments due to different handling of ReLU and adam implementations (differences in scores ~<0.001 when averaged across 100 different pairs) PyTorch analytical is the default now which may lead to some changes to results. We don't see many noticeable differences for test cases for the other scoring functions.

Test cases

  • Updated test cases for analytical implementations to compare accuracy with autograd implementations

Module reorganization

  • Created alignment directory which alignment_utils.py was moved to, and split up the alignment.py file into jax, torch, and torch+analytical modules.

Speedup summarized

Example: Align 100 pairs of randomly selected molecules (different sizes). Bolded rows are changes.

Jax is slower because it must compile a new computational graph for each new shape. This could be improved by introducing a padding mechanism. PyTorch can also be parallelized across cpus.

Volume results

algorithm Relative Speedup
Jax (before) 1
Jax 1.1
PyTorch (before) 1.3
PyTorch (autograd) 1.8
PyTorch (analytical) 3.0

Pharmacophore results

algorithm Relative Speedup
Jax (before) 1
Jax (before - vectorized) 5.7
PyTorch (before) 3.2
PyTorch (autograd) 6.3
PyTorch (analytical) 13.5

Example: Align 100 pairs of the same sized molecules (e.g., same sized tensors). Bolded rows are changes.

This favors algorithms like jax that compile code for shapes it has seen.

Volume Results

algorithm Relative Speedup
Jax (before) 6.8
PyTorch (before) 1
PyTorch (autograd) 1.5
PyTorch (analytical) 2.6

Pharmacophore Results

algorithm Relative Speedup
Jax (before) 2.8
Jax (before - vectorized) 22.5
PyTorch (before) 1
PyTorch (autograd) 1.3
PyTorch (analytical) 3.6

Full Changelog: v1.3.0...v1.3.1

Contributors

kentoabeywardane
Assets 2
Loading