def testForce(self):
# Create a random cloud of particles.
numParticles = 10
system = mm.System()
positions = np.random.rand(numParticles, 3)
for i in range(numParticles):
system.addParticle(1.0)
force = ot.TorchForce("../../tests/central.pt")
system.addForce(force)
# Compute the forces and energy.
integ = mm.VerletIntegrator(1.0)
context = mm.Context(system, integ,
mm.Platform.getPlatformByName('Reference'))
context.setPositions(positions)
state = context.getState(getEnergy=True, getForces=True)
# See if the energy and forces are correct. The network defines a potential of the form E(r) = |r|^2
expectedEnergy = np.sum(positions * positions)
assert np.allclose(
expectedEnergy,
state.getPotentialEnergy().value_in_unit(unit.kilojoules_per_mole))
assert np.allclose(-2 * positions, state.getForces(asNumpy=True))
def testModuleArguments(deviceString, precision):
if pt.cuda.device_count() < 1 and deviceString == 'cuda:0':
pytest.skip('A CUDA device is not available')
if pt.cuda.device_count() < 2 and deviceString == 'cuda:1':
pytest.skip('Two CUDA devices are not available')
class TestModule(pt.nn.Module):
def __init__(self, device, dtype, positions):
super().__init__()
self.device = device
self.dtype = dtype
self.register_buffer('positions', pt.tensor(positions).to(dtype))
def forward(self, positions):
assert self.positions.device == self.device
assert positions.device == self.device
assert positions.dtype == self.dtype
assert pt.all(positions == self.positions)
return pt.sum(positions)
with NamedTemporaryFile() as fd:
numParticles = 10
system = mm.System()
positions = np.random.rand(numParticles, 3)
for _ in range(numParticles):
system.addParticle(1.0)
device = pt.device(deviceString)
if device.type == 'cpu' or precision == 'double':
dtype = pt.float64
else:
dtype = pt.float32
module = TestModule(device, dtype, positions)
pt.jit.script(module).save(fd.name)
force = ot.TorchForce(fd.name)
system.addForce(force)
integrator = mm.VerletIntegrator(1.0)
platform = mm.Platform.getPlatformByName(device.type.upper())
properties = {}
if device.type == 'cuda':
properties['DeviceIndex'] = str(device.index)
properties['Precision'] = precision
context = mm.Context(system, integrator, platform, properties)
context.setPositions(positions)
context.getState(getEnergy=True, getForces=True)
def addForces(self,
topology: openmm.app.Topology,
system: openmm.System,
atoms: Optional[Iterable[int]],
forceGroup: int,
filename: str = 'animodel.pt',
**args):
# Create the TorchANI model.
import torchani
import torch
import openmmtorch
if self.name == 'ani1ccx':
model = torchani.models.ANI1ccx()
elif self.name == 'ani2x':
model = torchani.models.ANI2x()
else:
raise ValueError('Unsupported ANI model: ' + self.name)
# Create the PyTorch model that will be invoked by OpenMM.
includedAtoms = list(topology.atoms())
if atoms is not None:
includedAtoms = [includedAtoms[i] for i in atoms]
elements = [atom.element.symbol for atom in includedAtoms]
species = model.species_to_tensor(elements).unsqueeze(0)
class ANIForce(torch.nn.Module):
def __init__(self, model, species, atoms, periodic):
super(ANIForce, self).__init__()
self.model = model
self.species = species
self.energyScale = torchani.units.hartree2kjoulemol(1)
if atoms is None:
self.indices = None
else:
self.indices = torch.tensor(sorted(atoms),
dtype=torch.int64)
if periodic:
self.pbc = torch.tensor([True, True, True],
dtype=torch.bool)
else:
self.pbc = None
def forward(self,
positions,
boxvectors: Optional[torch.Tensor] = None):
positions = positions.to(torch.float32)
if self.indices is not None:
positions = positions[self.indices]
if boxvectors is None:
_, energy = self.model(
(self.species, 10.0 * positions.unsqueeze(0)))
else:
boxvectors = boxvectors.to(torch.float32)
_, energy = self.model(
(self.species, 10.0 * positions.unsqueeze(0)),
cell=10.0 * boxvectors,
pbc=self.pbc)
return self.energyScale * energy
aniForce = ANIForce(model, species, atoms,
topology.getPeriodicBoxVectors() is not None)
# Convert it to TorchScript and save it.
module = torch.jit.script(aniForce)
module.save(filename)
# Create the TorchForce and add it to the System.
force = openmmtorch.TorchForce(filename)
force.setForceGroup(forceGroup)
if topology.getPeriodicBoxVectors() is not None:
force.setUsesPeriodicBoundaryConditions(True)
system.addForce(force)
import simtk.openmm as mm
import simtk.unit as unit
import openmmtorch as ommt
# you may need to add libtorch to LD_LIBRARY_PATH before running the script
system = mm.System()
for i in range(3):
system.addParticle(1.0)
f = ommt.TorchForce('../tests/central.pt')
system.addForce(f)
integrator = mm.VerletIntegrator(0.001)
platform = mm.Platform.getPlatformByName('CUDA')
context = mm.Context(system, integrator, platform)
positions = [mm.Vec3(3, 0, 0), mm.Vec3(0, 4, 0), mm.Vec3(3, 4, 0)]
context.setPositions(positions)
print(context.getState(getEnergy=True).getPotentialEnergy())
# should give 50.0 kJ/mol