How to use global parameter and PBC at the same time?

[xiaowei-xie2]

Hi,

I was trying to add a global parameter to my torchForce, but I could not figure out how to make it work along with the PBC. Could you help me have a look?

Here is the structure of my torchForce class:

class ForceModule(torch.nn.Module):
def __init__(self):
    ...
def forward(self, coordinates: torch.Tensor, cell: torch.Tensor, scale: float, RcNL: float = 10.0):
     return energy

Then when I try to load the model with

force = TorchForce('model.pt')
force.setUsesPeriodicBoundaryConditions(True)
force.addGlobalParameter('scale', 0.9)

and try to run a MD with it, I get the following error:

Traceback (most recent call last):
  File "/home/xie1/torchSANI_jit/scripts/openmm_test_3.py", line 128, in <module>
    state = simulation.context.getState(getEnergy=True, groups={10})
  File "/home/xie1/psi4conda/envs/nnpops_env/lib/python3.10/site-packages/openmm/openmm.py", line 12111, in getState
    state = _openmm.Context_getState(self, types, enforcePeriodicBox, groups_mask)
openmm.OpenMMException: forward() Expected a value of type 'float' for argument 'scale' but instead found type 'Tensor'.
Position: 3
Declaration: forward(__torch__.___torch_mangle_0.ForceModule self, Tensor coordinates, Tensor cell, float scale, float RcNL=10.) -> Tensor

Thank you so much!

[peastman]

You declare that the scale parameter should be a float:

def forward(self, coordinates: torch.Tensor, cell: torch.Tensor, scale: float, RcNL: float = 10.0):

But when OpenMM calls your module, it passes all the arguments as Tensors, leading to the exception

openmm.OpenMMException: forward() Expected a value of type 'float' for argument 'scale' but instead found type 'Tensor'.

If you change your declaration to scale: torch.Tensor, it should work.

[xiaowei-xie2]

That worked, thank you! A follow-up question: can I get the derivative w.r.t. that global parameter along the MD trajectory?

[peastman]

No, but that would be a reasonable feature to add. Is that something you would find useful?

[xiaowei-xie2]

Yes, I think that would be super helpful for any FEP type calculations with ML force fields! Do you know any workarounds at the moment? If I can write out the derivative expression myself, what would be a way to save that value at each MD frame?

[peastman]

I created a feature request for it (#141). I'll try to do it soon. I don't think it should be difficult.

[xiaowei-xie2]

Thank you so much!

[peastman]

The implementation is at #143. If you care to try it out, please let us know whether it works for you.

[xiaowei-xie2]

Wow, thank you! Do I need to build from source from your branch to test it out?

[peastman]

Yes, if that's not too difficult. Otherwise you'll need to wait for the next release.

[xiaowei-xie2]

No problem, I will give it a try soon and let you know!