Implement NNPOps Optimised ANI

openmmml/models/anipotential.py

@@ -63,55 +63,77 @@ def addForces(self,
                   atoms: Optional[Iterable[int]],
                   forceGroup: int,
                   filename: str = 'animodel.pt',
+                  implementation: str = "nnpops",
                   **args):
         # Create the TorchANI model.
-
         import torchani
         import torch
         import openmmtorch
+        from NNPOps import OptimizedTorchANI
+
         if self.name == 'ani1ccx':
-            model = torchani.models.ANI1ccx()
+            model = torchani.models.ANI1ccx(periodic_table_index=True)
         elif self.name == 'ani2x':
-            model = torchani.models.ANI2x()
+            model = torchani.models.ANI2x(periodic_table_index=True)
         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)
+        atomic_numbers = [atom.element.atomic_number for atom in includedAtoms]
+        species = torch.tensor(atomic_numbers).unsqueeze(0)
+
+        if implementation == "nnpops":
+            from NNPOps import OptimizedTorchANI
+            device = torch.device('cuda')
+            model = OptimizedTorchANI(model, species).to(device)
+        elif implementation == "cuaev":
+            self.model.aev_computer.use_cuda_extension = True
+        elif implementation == "torchani":
+            # nothing extra to do here
+            pass
+        else:
+            raise ValueError(f"Unsupported implementation: {implementation}")
 
         class ANIForce(torch.nn.Module):
 
-            def __init__(self, model, species, atoms, periodic):
+            def __init__(self, model, species, atoms):
                 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
+
+                # Accelerate the model
+                self.model = OptimizedTorchANI(model, self.atomic_numbers)
+                self.pbc = torch.tensor([True, True, True], dtype=torch.bool)
 
             def forward(self, positions, boxvectors: Optional[torch.Tensor] = None):
+                # Prepare the positions
                 positions = positions.to(torch.float32)
+
                 if self.indices is not None:
                     positions = positions[self.indices]
+
+                positions = positions.unsqueeze(0) * 10.0 # nm --> Å
+
+                # Run ANI to get the potential energy
                 if boxvectors is None:
-                    _, energy = self.model((self.species, 10.0*positions.unsqueeze(0)))
+                    _, energy = self.model((self.atomic_numbers, positions))
                 else:
+                    self.pbc = self.pbc.to(positions.device)
                     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
+                    _, energy = self.model((self.atomic_numbers, positions), cell=10.0*boxvectors, pbc=self.pbc)
+
+                return energy * self.energyScale # Hartree --> kJ/mol
 
-        aniForce = ANIForce(model, species, atoms, topology.getPeriodicBoxVectors() is not None)
+        aniForce = ANIForce(model, species, atoms)
 
         # Convert it to TorchScript and save it.