Schnet in OpenMM

[keano130]

As in openmm/openmm-torch#33, I am trying to implement a SchNet neural network into OpenMM for md simulations. I wrote the ForceModule as in https://github.com/openmm/openmm-torch/blob/master/README.md such that i could use torch.jit.script. However, I start getting errors due to the code in Schnetpack: for example the function atom_distances in https://github.com/atomistic-machine-learning/schnetpack/blob/master/src/schnetpack/nn/neighbors.py sometimes returns 1 variable and in other cases 2, which causes errors for torch.jit.script.
Adapting the schnetpack code such that i can compile it to a torch_script seems futile as this method is not that efficient in the first place (as the neighborlist is rebuild in every iteration).
In order to have an efficient version it seems that NNPOps should do the trick, however, for as far as I understand it there is not yet a python wrapper for the schnet in NNPOps. I have little experience with C and python wrappers, so I could definitely be mistaken.
Therefore, I have the question whether or not it is possible to rebuild the schnet network in python using NNPOps, or if there is a simpler solution to use a SchNet neural network in OpenMM md Simulations?

[raimis]

Yes, currently there is no Python wrapper for the SchNet operations. Following the ANI as an example, it shouldn't be very hard to implement it.

Alternatively, you can try the SchNet implementation in PyTorch Geometric (https://pytorch-geometric.readthedocs.io/en/latest/_modules/torch_geometric/nn/models/schnet.html). It should be faster that SchNetPack, but its conda packages not binary compatible with OpenMM ones.