PBC not applied in TorchMDNetPotential #115
Description
I’ve been using TorchMDNetPotential in OpenMM-ML and noticed that periodic boundary conditions are currently not taken into account inside torchmdnetpotential.py. For systems with a periodic box this becomes an issue, since the forces are computed as if the system were non-periodic.
A small change would solve it. Something like:
def forward(self, positions: torch.Tensor, boxvectors: Optional[torch.Tensor] = None):
positions = positions.to(torch.float32).to(self.numbers.device)
if self.subset:
positions = positions[self.indices]
**if boxvectors is not None:
cell = boxvectors.to(torch.float32)/self.lengthScale
else:
cell = None**
energy = self.model(z=self.numbers, pos=positions/self.lengthScale, batch=self.batch, q=self.charge, **box=cell**)[0]
return energy*self.energyScale
# Create the TorchForce and add it to the System.
module = torch.jit.script(TorchMDNetForce(model, numbers, charge, atoms, batch, self.lengthScale, self.energyScale)).to(torch.device('cpu'))
**isPeriodic = (topology.getPeriodicBoxVectors() is not None) or system.usesPeriodicBoundaryConditions()**
force = openmmtorch.TorchForce(module)
if use_cudagraphs:
force.setProperty("useCUDAGraphs", "true")
force.setForceGroup(forceGroup)
**force.setUsesPeriodicBoundaryConditions(isPeriodic)**
system.addForce(force)