Implement NNPOps Optimised ANI

[not-matt]

These changes simplify the use of NNPOps OptimisedTorchANI in OpenMM simulations, allowing ANI2x to be used with ~6x speedup on my machine. It also allows periodic boundary conditions again for CUDA (see aiqm/torchani#607 (comment))

• Faster ANI for small, unsolvated systems
• Allows large solvated systems to be run using mixed ANI/AMBER

Thanks to @jchodera for his example code openmm/openmm-torch#59 (comment)
Addresses #14

Demonstration
https://github.com/meyresearch/ANI-Peptides/blob/main/demos/ANI_minimal.ipynb

Important to note that this PR lacks an option to choose between classic TorchANI and OptimisedTorchANI, instead enforcing NNPops as the only method. I'd be happy to implement a way to choose between the two, but I'm not sure how this would best be implemented to fit this library's structure?

eg.

openmmml.MLPotential('ani2x_optimised')

or perhaps

openmmml.MLPotential('ani2x', optimised=True)

[raimis]

Note that #20 tries to do the same.

[raimis]

@not-matt, regarding the performance, could you run the benchmarks and compare with #20 (comment)

[not-matt]

With the bench slightly modified to fix a couple of bugs (#20 (comment)), I got 0.0031 s/step using NNPops. This is in line with the results from that PR - the implementations are pretty much identical under the hood.

[raimis]

Regarding the optimization keyword, we have 3 cases:

• Pure PyTorch implementation of ANI-2x

openmmml.MLPotential('ani2x', optimise=None)

• ANI-2x accelerated with cuAEV (https://github.com/aiqm/torchani/tree/master/torchani/cuaev)

openmmml.MLPotential('ani2x', optimise='cuaev')

• ANI-2x accelerated with NNPOps (by default)

openmmml.MLPotential('ani2x', optimise='nnpops')

I think, this will be the simplest API for a user.

Ping: @jchodera @peastman @yueyericardo @dominicrufa]

[peastman]

How about calling it implementation, since that's really what you're selecting: which implementation to use. I'd also suggest that if you want the default TorchANI implementation the value would be 'torchani' rather than None. That will be clearer.

[raimis]

If no more comments, the API should look like this:

openmmml.MLPotential('ani2x', implementation='torchani')
openmmml.MLPotential('ani2x', implementation='cuaev')
openmmml.MLPotential('ani2x', implementation='nnpops') # default

@not-matt could you implement?

[not-matt]

Yep, can do.

@raimis cuaev is enabled like so?

self.model.aev_computer.use_cuda_extension = True

[raimis]

@raimis cuaev is enabled like so?

self.model.aev_computer.use_cuda_extension = True

@yueyericardo could you confirm?

[yueyericardo]

confirmed!

[raimis]

@not-matt any progress?

[not-matt]

@raimis Yep just about there. Trying to test it thoroughly with each different optimisation option.

[raimis]

This doesn't work

[raimis]

atomic_numbers is a more meaningful name for this variable.

[raimis]

Also, it would be good to have a test to make sure things don't get broken.

[peastman]

Is this ready to merge?

[raimis]

I think, we still need some test to make sure the things keep working.

[peastman]

What else has to be done before we can merge this? It would be really nice if we could get it in before building the OpenMM 8.0 beta, so we can get beta users to try out the full stack including optimized kernels.

[peastman]

This is superseded by #35.