gcorso / DiffDock

Implementation of DiffDock: Diffusion Steps, Twists, and Turns for Molecular Docking
https://arxiv.org/abs/2210.01776
MIT License
976 stars 238 forks source link

vectorize so3, so that I don't have to wait for 20 minutes #218

Closed tornikeo closed 2 months ago

tornikeo commented 2 months ago

Importing utils/so3.py takes absolutely forever. Yes, it's a one-time call, but I don't want to suffer it when I move this every time I use this code on a fresh remote VM. So I vectorized the heaviest calculations

I tested this locally, but I don't see a tests/ directory to check this PR automatically.

Also, are you guys open to me adding a few automated tests for this repo? The recent BELKA challenge on kaggle will likely generate a lot of PRs and issues, and having test suite will make approving and checking good changes lot easier.

jsilter commented 2 months ago

Have you done a runtime comparison that you could share?

I'm definitely open to reviewing automated tests if you'd like to add some.

jsilter commented 2 months ago

Included in v1.1.2