PBC conditions increase memory consumption by more than 20 times

[JSLJ23]

Hi developers of torchANI,

I was trying to run some MD simulations of a really small protein, a 10 peptide Chignolin protein with a 4Å buffer water-box which is about a 1000 atoms or so. Turning on the PBC in ASE with ANI-2x as the calculator seems to increase the CUDA memory allocations by more than 20 times leading to a CUDA OOM issue (on a RTX4090 24GB).

It seems that the neural network is being used to compute the potentials of these periodic unit cells but from my understanding, that shouldn't be required. I was thinking along the lines that the neighbouring 26 PBC cells would just copy the positions and of the central unit cell but just add a x, y, z (or combinations for the diagonals) offset and hence the potentials for them are not required to be computed. However, this does not seem to be the case with the ANI-2x calculator in ASE.

May I know if there is a torchANI specific setting that could be turned on to avoid calculating the potentials in the PBC unit cells or is there something I am missing here?

Are there any recommendations to running MD simulations with torchANI in ASE for larger systems with 30-40k atoms with something like an A100(80GB) and not run into OOM issues with these extra PBC computations?

Thank you for your time to address my queries and I hope to hear from you guys soon.

Best regards,
Joshua Soon

[yueyericardo]

Hi Joshua, the current PBC condition within ANI is implemented in a way that is very expensive when running for a very large system. We have implemented several other options (e.g. cell list) in our internal code, we are planning to release it soon.

We also have the lammps interface implemented, and we are able to run up to a 1.5M atom system on A100 (80GB) GPU. It will also be released soon, but there might be still a few months. Please send jinzexue@ufl.edu an email if you are interested, we could share a docker with you for testing.