Implement NNPOps Optimised ANI

openmmml/models/anipotential.py

@@ -63,6 +63,7 @@ def addForces(self,
                   atoms: Optional[Iterable[int]],
                   forceGroup: int,
                   filename: str = 'animodel.pt',
+                  implementation: str = "nnpops",
                   **args):
         # Create the TorchANI model.
         import torchani
@@ -82,14 +83,26 @@ def addForces(self,
         if atoms is not None:
             includedAtoms = [includedAtoms[i] for i in atoms]
         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, atomic_numbers, atoms, periodic):
-                super().__init__()
-
-                # Store the atomic numbers
-                self.atomic_numbers = torch.tensor(atomic_numbers).unsqueeze(0)
+            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:
@@ -120,7 +133,7 @@ def forward(self, positions, boxvectors: Optional[torch.Tensor] = None):
 
                 return energy * self.energyScale # Hartree --> kJ/mol
 
-        aniForce = ANIForce(model, atomic_numbers, atoms, topology.getPeriodicBoxVectors() is not None)
+        aniForce = ANIForce(model, species, atoms)
 
         # Convert it to TorchScript and save it.