search

openmm / NNPOps

High-performance operations for neural network potentials
Other
81 stars 17 forks source link

Particle Mesh Ewald #103

Closed peastman closed 1 year ago

peastman commented 1 year ago

I'm working on an implementation of PME. This will be useful for models that include explicit Coulomb terms, especially when the charges are predicted by the model and require extra chain rule terms.

The direct space interaction is implemented in PyTorch and is very simple. The reciprocal space part is implemented as a C++ extension. I've written the CPU version. I still have to do the CUDA version.

proteneer commented 1 year ago

Would it be possible to implement the reciprocal space energy calculations in python/torch as well? It would be easier to do tests of things like forces by taking advantage of automatic differentiation (as opposed to just comparing against hard coded floats)

peastman commented 1 year ago

You could try, by I couldn't think of any way to implement it that wouldn't be too horribly inefficient. For example, you either end up with a loop in Python to spread charges from atoms one at a time, or you write a vectorized version that writes to every grid point for every atom. Can you think of a better way?

Letting PyTorch compute derivatives through generic backpropagation is also likely to be extremely expensive. The FFT means that every point in the reciprocal space grid depends on every point in the real space grid.

What sorts of tests would you want to do?

peastman commented 1 year ago

This is ready for review.

I originally hoped to implement the direct space calculation in Python, but I couldn't figure out any way to do it that was compatible with TorchScript. The implementation still isn't especially fast, since it just uses the NNPOps neighbor list structure, which isn't well suited to GPU computation, but it should be plenty fast enough for our purposes. The reciprocal space calculation is based on the implementation in OpenMM. I simplified it a bit, removing optimizations that add a lot of complexity for only modest speedup, but the performance should still be pretty good.

If someone can double check my values for Coulomb's constant, that would be helpful!

RaulPPelaez commented 1 year ago

Really clean implementation @peastman ! Just some comments/questions from me, but looks great!