`TorchForce` inside `CustomCVForce` fails with `nnpops` implementation

[dominicrufa]

I have been trying to create a minimally-modified implementation to get nnpops working on a conda-installable openmm-ml environment, and i am encountering an error in TestMLPotential.py with the implementation.

If you drop in/replace the anipotential.py with this and run the test, (making sure to set the platform to CUDA), you should see that the test fails on line 26,

======================================================================
ERROR: testCreateMixedSystem (__main__.TestMLPotential)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/lila/home/rufad/github/openmm-ml/test/TestMLPotential.py", line 26, in testCreateMixedSystem
    interpEnergy1 = interpContext.getState(getEnergy=True).getPotentialEnergy().value_in_unit(unit.kilojoules_per_mole)
  File "/home/rufad/miniconda3/envs/omm_dev2/lib/python3.9/site-packages/openmm/openmm.py", line 12318, in getState
    state = _openmm.Context_getState(self, types, enforcePeriodicBox, groups_mask)
openmm.OpenMMException: Error invoking kernel: CUDA_ERROR_INVALID_HANDLE (400)

meaning that the createMixedSystem(..., interpolate=False) is compatible with nnpops, but interpolate=True is not. I have found a way around this (by taking the TorchForce out of the customCVForce and adding a scale parameter to the ANIForce manually), but I don't want to open a PR attempting to merge my modifications unless the compatibility issue with nnpops can't be fixed an easier way.

[dominicrufa]

@peastman , @raimis: turns out ^this is actually the crux of this issue

[peastman]

This may also be related to https://github.com/openmm/openmm/issues/3588.

[dominicrufa]

in the interest of getting openmm-ml working with nnpops, i think the easiest thing would be to just extract the TorchForce from the customCV (as in my fork) and run with it until a more stable solution is found. is that agreeable?

[dominicrufa]

@peastman, have you found a stable workaround for this issue, or should we proceed with removing the TorchForce from the CVForce?

[dominicrufa]

tagging @jchodera

[peastman]

Let me make one more effort at trying to figure it out. I suspect this is related to https://github.com/openmm/openmm/issues/3588, and most especially to https://github.com/openmm/openmm/issues/3588#issuecomment-1151731824. If I can't figure out a proper solution, we can use your suggested workaround.

[jchodera]

@dominicrufa : Can you test the PR here? https://github.com/openmm/openmm-torch/pull/78

[dominicrufa]

@peastman , I am unable to source install the PR and make it compatible with the rest of the conda-installed bits needed for the nnpops test. if you can source install it locally, can you run the test i mentioned at the top of this issue?

[peastman]

It runs without problem now. It looks like that fixes it!

[dominicrufa]

@peastman , awesome! can you merge the PR and make it conda-installable? i can add a PR with the nnpops implementation for anipotential.py here and modify the test to test for the torchani and nnpops implementation if you'd like.

[raimis]

@peastman could reopen this. I don't permissions.

[peastman]

I just added you as a maintainer. You should have full permissions now.

[raimis]

@dominicrufa I can reproduce the issue and created a minimal test https://github.com/openmm/openmm-torch/pull/80/commits/63a49c170232531c12edf5db2d2f5255af1dfa6f.

[dominicrufa]

@raimis , great! let me know if/when i can conda-install the fix to write the openmm-ml tests.

[jchodera]

@dominicrufa : Can you confirm the latest release fixes this? If so, can you close this?

[dominicrufa]

yes, this is fixed. closing