Add energy derivatives to README

README.md

@@ -256,6 +256,72 @@ to return forces.
 torch_force.setOutputsForces(True)
 ```
 
+Computing energy derivatives with respect to global parameters
+--------------------------------------------------------------
+
+Its possible to query `TorchForce` for the derivative of the energy with respect to global parameters. In order to do so the global parameters must be registered as energy derivatives. This is done by calling `addEnergyParameterDerivative()` for each parameter.
+
+The parameter derivatives can be queried by calling `getEnergyParameterDerivatives()` on the `State` object returned by `Context.getState()`. The result is a dictionary with the parameter names as keys and the derivatives as values.
+
+```python
+import torch as pt
+from torch import Tensor
+from openmmtorch import TorchForce
+import openmm as mm
+
+class ForceWithParameters(pt.nn.Module):
+
+ def __init__(self):
+ super(ForceWithParameters, self).__init__()
+
+ def forward(
+ self, positions: Tensor, parameter1: Tensor, parameter2: Tensor
+ ) -> Tensor:
+ x2 = positions.pow(2).sum(dim=1)
+ u_harmonic = ((parameter1 + parameter2**2) * x2).sum()
+ return u_harmonic
+
+
+def example():
+ numParticles = 10
+ system = mm.System()
+ positions = np.random.rand(numParticles, 3)
+ for _ in range(numParticles):
+ system.addParticle(1.0)
+
+ pt_force = ForceWithParameters()
+ model = pt.jit.script(pt_force)
+ tforce = TorchForce(model)
+ parameter1 = 1.0
+ parameter2 = 1.0
+ force.setOutputsForces(False)
+ force.addGlobalParameter("parameter1", parameter1)
+ force.addEnergyParameterDerivative("parameter1")
+ force.addGlobalParameter("parameter2", parameter2)
+ force.addEnergyParameterDerivative("parameter2")
+ system.addForce(force)
+ integ = mm.VerletIntegrator(1.0)
+ platform = mm.Platform.getPlatformByName(platform)
+ context = mm.Context(system, integ, platform)
+ context.setPositions(positions)
+ state = context.getState(
+ getEnergy=True, getForces=True, getParameterDerivatives=True
+ )
+ # The network defines a potential of the form E(r) = (parameter1 + parameter2**2)*|r|^2
+ r2 = np.sum(positions * positions)
+ expectedEnergy = (parameter1 + parameter2**2) * r2
+ assert np.allclose(
+ r2,
+ state.getEnergyParameterDerivatives()["parameter1"],
+ )
+ assert np.allclose(
+ 2 * parameter2 * r2,
+ state.getEnergyParameterDerivatives()["parameter2"],
+ )
+```
+
+
+
 Recording the model into a CUDA graph
 -------------------------------------