Openmm integration

[chrisiacovella]

Pull Request Summary

This PR revamps the neighbor list (allowing us to change the strategy for a JITTED torchscript potential) and provide a wrapper allowing us to run this with openmm -torch.

Key changes

Notable points that this PR has either accomplished or will accomplish.

• Neighborlists include "only_unique_pairs" as a buffer
• Neighborlists for inference combined into a single class, with the ability to set the strategy either at run time or via function call (allowing a single model file)
• wrapper that follows the structure of TorchForce for openmm
• This creates a submodule for openmm integration (via modelforge.openmm)
• Tests for modelforge-openmm wrapping/running openmm

Associated Issue(s)

• Neighborlists in TorchScript models #297

Pull Request Checklist

• Issue(s) raised/addressed and linked
• Includes appropriate unit test(s)
• Appropriate docstring(s) added/updated
• Appropriate .rst doc file(s) added/updated
• PR is ready for review

[chrisiacovella]

Note this is currently failing due to reading in a checkpoint file; I added the unique_pairs variable to the buffers in the neighbor list, as that seemed to be a very import thing to set. However, it is also set up stream (and encoded likely in another location). So I wonder if this is necessary.

[chrisiacovella]

The switch from a named tuple to a class for NNPInput actually presents some problems when we go to use the torch script models (i.e., passing input after calling torch.jit.script). Torch script doesn't really handle getting a class as input well. it sees the the NNPInput class as a different type due to being different Python compilation unit (i.e., compiled using modelforge and then compiled by the wrapper class).

NNPInput is more of a convenience thing (has validation, sets default shapes, etc.), but for inference we just need to pass tensors directly.

[wiederm]

@chrisiacovella , NNPInput is torch scripted right now; we could also set a flag not to jit that class (there is no benefit to it at that point).
An alternative is to pass input as keyword arguments, which should also be fine and need minimal changes.

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 36.95652% with 58 lines in your changes missing coverage. Please review.

Project coverage is 34.22%. Comparing base (426171a) to head (95fae2e).

❗ There is a different number of reports uploaded between BASE (426171a) and HEAD (95fae2e). Click for more details.

HEAD has 7 uploads less than BASE

Flag	BASE (426171a)	HEAD (95fae2e)
unittests	8	1

Additional details and impacted files

[wiederm]

Great work!

I'd love to better understand the control flow from input to property with this module, I left a comment.

[wiederm]

Great work!