Optimize SpeciesConverter of TorchANI

[raimis]

This assume that a molecule is constant, so the atomic number conversion to species can be pre-computed.

• Implement TorchANISpeciesConverter
• Tests
• Documentation

[raimis]

@peastman could you review? One more small optimization like #38.

[peastman]

Like with your EnergyShifter class, I don't think this repository is the right place for this. NNPOps is a collection of optimized kernels for bottleneck operations. It isn't meant to be a grab bag of minor patches to other libraries.

[peastman]

Also, NNPOps is supposed to sit just above PyTorch in the stack, but below libraries that make use of it like TorchANI or TorchMD-Net. That means it can't import TorchANI. Otherwise it leads to circular dependencies. And circular dependencies lead to suffering.

[raimis]

Otherwise it leads to circular dependencies. And circular dependencies lead to suffering.

TorchANI doesn't import NNPOps, so there are no circular dependencies.

Like with your EnergyShifter class, I don't think this repository is the right place for this. NNPOps is a collection of optimized kernels for bottleneck operations. It isn't meant to be a grab bag of minor patches to other libraries.

As you should remember, in March we tried to engage with TorchANI developers, but we haven't reached an agreement.

[peastman]

TorchANI doesn't import NNPOps, so there are no circular dependencies.

But it might in the future. We intend that it should be able to. The point is that when you design an architecture, you need to define the position of every module within the stack. Each one can use things lower in the stack, but not anything higher up.

The NNPOps Python module is meant to sit just above PyTorch, but below any library that might potentially use it (TorchANI, SchNetPack, TorchMD-Net, etc.). When you install it, that shouldn't require installing anything beyond PyTorch. If someone wants to use it, that shouldn't require them (and all their users) to install TorchANI. When you import it, that shouldn't import TorchANI.

As you should remember, in March we tried to engage with TorchANI developers, but we haven't reached an agreement.

I don't think we discussed these two features? I suggest opening an issue on the TorchANI github suggesting them. You can offer to open a PR implementing them.

[raimis]

But it might in the future. We intend that it should be able to. The point is that when you design an architecture, you need to define the position of every module within the stack. Each one can use things lower in the stack, but not anything higher up.

The current goal is to get the ML support into OpenMM (https://github.com/orgs/openmm/projects/1) as soon as possible. What will happen in the future, we will deal in the future.

The NNPOps Python module is meant to sit just above PyTorch, but below any library that might potentially use it (TorchANI, SchNetPack, TorchMD-Net, etc.). When you install it, that shouldn't require installing anything beyond PyTorch. If someone wants to use it, that shouldn't require them (and all their users) to install TorchANI. When you import it, that shouldn't import TorchANI.

This PR doesn't change the situation: SymmetryFunctions and BarchedNN already depend on TorchANI. If we decide to move these classes somewhere else, it won't be a problem to move this one and #38 along.

[peastman]

What will happen in the future, we will deal in the future.

No! That's the exact opposite of how you design robust software! You need to have a clear architecture and rigorously maintain it. Trust me, if you don't, you will come to regret it!

SymmetryFunctions and BarchedNN already depend on TorchANI.

Then we need to fix that ASAP, because it breaks our architecture. For example, if the SchNetPack developers wanted to adopt our optimized kernel for computing SchNet, it would become impossible to import schnetpack if you didn't have TorchANI installed. There's no way they would accept that. They would rightly view it as a dealbreaker.

[peastman]

I think it might help to step back and consider what the purpose of NNPOps is. Fortunately, the goal is clearly stated right at the start of the README:

The goal of this repository is to promote the use of neural network potentials (NNPs) by providing highly optimized, open source implementations of bottleneck operations that appear in popular potentials.

That's it. Anything that doesn't fit within that goal (optimized implementations of bottleneck operations) doesn't belong in NNPOps. It goes somewhere else. Next, the README lists the core design principles. Here is the first one.

Each operation is entirely self contained, consisting of only a few source files that can easily be incorporated into any code that needs to use it.

That's critical. NNPOps is meant to be self contained. You should be able to use it without bringing in lots of other dependencies.

Now consider this PR. It doesn't implement any calculations at all. It's just a trivial wrapper around a TorchANI class. Its only function is to invoke the TorchANI object and stash the result in a field. That's it. I suppose you must have written this because it was useful for some code you were writing, but it has nothing to do with any of the goals of NNPOps. This isn't the right place for it.

[raimis]

Check #41 for the full picture. Specifically, the example: https://github.com/openmm/NNPOps/blob/2a5da7c24590a2bd21786ad0ad4df5092690028d/README.md#example

[peastman]

The first thing I would do is open one or more issues on the TorchANI repo suggesting these changes and offering to provide code. This class doesn't implement any kind of calculation. It's just a workaround for a particular flaw in TorchANI (that it repeats the calculation on every step). The more of these changes we can push upstream to them, the better.

For anything they don't want to accept, the next thing I would do is move the imports inside the classes. You should be able to import NNPOps even if TorchANI isn't installed. In practice that means we should never import torchani at the top level. It should only be imported inside the specific classes or functions that use it.

Finally we should consider whether NNPOps is the best place for all of these, or whether, for example, OpenMM-ML would be a better place.

[raimis]

Let's move the discussion to #40. It is more general than this PR.

[raimis]

@peastman can we proceed with this?

[peastman]

Go ahead.