Optimize TorchANI (#41)

* Optimize EnergyShifter of TorchANI

* Install EnergyShifter and enable testing

* Copy the model to a correct device

* Update READM.md

* Fix typo

* Simplify code

* Simplify more

* Implement TorchANISpeciesConverter

* Add test for TorchANISpeciesConverter

* Update README.md

* Implement OptimizedTorchANI

* Add tests for OptimizeTorchANI

* Import directly OptimizedTorchANI

* Fix a file name

* Add an example for OptimizedTorchANI

* Fix the test tolerances

* Fix a type in the example

* Move the top-level import of TorchANI

* Update the doc string (#44)

CMakeLists.txt

@@ -30,14 +30,16 @@ foreach(TEST_PATH ${TEST_PATHS})
 endforeach()
 
 add_test(TestBatchedNN         pytest -v ${CMAKE_SOURCE_DIR}/src/pytorch/TestBatchedNN.py)
-add_test(TestSpeciesConverter  pytest -v ${CMAKE_SOURCE_DIR}/src/pytorch/TestSpeciesConverter.py)
 add_test(TestEnergyShifter     pytest -v ${CMAKE_SOURCE_DIR}/src/pytorch/TestEnergyShifter.py)
+add_test(TestOptimizedTorchANI pytest -v ${CMAKE_SOURCE_DIR}/src/pytorch/TestOptimizedTorchANI.py)
+add_test(TestSpeciesConverter  pytest -v ${CMAKE_SOURCE_DIR}/src/pytorch/TestSpeciesConverter.py)
 add_test(TestSymmetryFunctions pytest -v ${CMAKE_SOURCE_DIR}/src/pytorch/TestSymmetryFunctions.py)
 
 install(TARGETS ${LIBRARY} DESTINATION ${Python_SITEARCH}/${NAME})
 install(FILES src/pytorch/__init__.py
               src/pytorch/BatchedNN.py
-              src/pytorch/SpeciesConverter.py
               src/pytorch/EnergyShifter.py
+              src/pytorch/OptimizedTorchANI.py
+              src/pytorch/SpeciesConverter.py
               src/pytorch/SymmetryFunctions.py
         DESTINATION ${Python_SITEARCH}/${NAME})

README.md

@@ -80,6 +80,8 @@ from NNPOps.SymmetryFunctions import TorchANISymmetryFunctions
 from NNPOps.BatchedNN import TorchANIBatchedNN
 from NNPOps.EnergyShifter import TorchANIEnergyShifter
 
+from NNPOps import OptimizedTorchANI
+
 device = torch.device('cuda')
 
 # Load a molecule
@@ -100,4 +102,16 @@ energy.backward()
 forces = -positions.grad.clone()
 
 print(energy, forces)
-```
+
+# Alternatively, all the optimizations can be applied with OptimizedTorchANI
+nnp2 = torchani.models.ANI2x(periodic_table_index=True).to(device)
+nnp2 = OptimizedTorchANI(nnp2, species).to(device)
+
+# Compute energy and forces again
+energy = nnp2((species, positions)).energies
+positions.grad.zero_()
+energy.backward()
+forces = -positions.grad.clone()
+
+print(energy, forces)
+```

src/pytorch/OptimizedTorchANI.py

@@ -0,0 +1,56 @@
+#
+# Copyright (c) 2020-2021 Acellera
+# Authors: Raimondas Galvelis
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+
+import torch
+from torch import Tensor
+from typing import Optional, Tuple
+
+from NNPOps.BatchedNN import TorchANIBatchedNN
+from NNPOps.EnergyShifter import TorchANIEnergyShifter, SpeciesEnergies
+from NNPOps.SpeciesConverter import TorchANISpeciesConverter
+from NNPOps.SymmetryFunctions import TorchANISymmetryFunctions
+
+class OptimizedTorchANI(torch.nn.Module):
+
+    from torchani.models import BuiltinModel # https://github.com/openmm/NNPOps/issues/44
+
+    def __init__(self, model: BuiltinModel, atomicNumbers: Tensor) -> None:
+
+        super().__init__()
+
+        # Optimize the components of an ANI model
+        self.species_converter = TorchANISpeciesConverter(model.species_converter, atomicNumbers)
+        self.aev_computer = TorchANISymmetryFunctions(model.aev_computer)
+        self.neural_networks = TorchANIBatchedNN(model.species_converter, model.neural_networks, atomicNumbers)
+        self.energy_shifter = TorchANIEnergyShifter(model.species_converter, model.energy_shifter, atomicNumbers)
+
+    def forward(self, species_coordinates: Tuple[Tensor, Tensor],
+                cell: Optional[Tensor] = None,
+                pbc: Optional[Tensor] = None) -> SpeciesEnergies:
+
+        species_coordinates = self.species_converter(species_coordinates)
+        species_aevs = self.aev_computer(species_coordinates, cell=cell, pbc=pbc)
+        species_energies = self.neural_networks(species_aevs)
+        species_energies = self.energy_shifter(species_energies)
+
+        return species_energies

src/pytorch/TestOptimizedTorchANI.py

@@ -0,0 +1,104 @@
+#
+# Copyright (c) 2020-2021 Acellera
+# Authors: Raimondas Galvelis
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+
+import mdtraj
+import os
+import pytest
+import tempfile
+import torch
+import torchani
+
+molecules = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'molecules')
+
+@pytest.mark.parametrize('deviceString', ['cpu', 'cuda'])
+@pytest.mark.parametrize('molFile', ['1hvj', '1hvk', '2iuz', '3hkw', '3hky', '3lka', '3o99'])
+def test_compare_with_native(deviceString, molFile):
+
+    if deviceString == 'cuda' and not torch.cuda.is_available():
+        pytest.skip('CUDA is not available')
+
+    from NNPOps import OptimizedTorchANI
+
+    device = torch.device(deviceString)
+
+    mol = mdtraj.load(os.path.join(molecules, f'{molFile}_ligand.mol2'))
+    atomicNumbers = torch.tensor([[atom.element.atomic_number for atom in mol.top.atoms]], device=device)
+    atomicPositions = torch.tensor(mol.xyz * 10, dtype=torch.float32, requires_grad=True, device=device)
+
+    nnp = torchani.models.ANI2x(periodic_table_index=True).to(device)
+    energy_ref = nnp((atomicNumbers, atomicPositions)).energies
+    energy_ref.backward()
+    grad_ref = atomicPositions.grad.clone()
+
+    nnp = OptimizedTorchANI(nnp, atomicNumbers).to(device)
+    energy = nnp((atomicNumbers, atomicPositions)).energies
+    atomicPositions.grad.zero_()
+    energy.backward()
+    grad = atomicPositions.grad.clone()
+
+    energy_error = torch.abs((energy - energy_ref)/energy_ref)
+    grad_error = torch.max(torch.abs((grad - grad_ref)/grad_ref))
+
+    assert energy_error < 5e-7
+    if molFile == '3o99':
+        assert grad_error < 7e-3
+    else:
+        assert grad_error < 5e-3
+
+@pytest.mark.parametrize('deviceString', ['cpu', 'cuda'])
+@pytest.mark.parametrize('molFile', ['1hvj', '1hvk', '2iuz', '3hkw', '3hky', '3lka', '3o99'])
+def test_model_serialization(deviceString, molFile):
+
+    if deviceString == 'cuda' and not torch.cuda.is_available():
+        pytest.skip('CUDA is not available')
+
+    from NNPOps import OptimizedTorchANI
+
+    device = torch.device(deviceString)
+
+    mol = mdtraj.load(os.path.join(molecules, f'{molFile}_ligand.mol2'))
+    atomicNumbers = torch.tensor([[atom.element.atomic_number for atom in mol.top.atoms]], device=device)
+    atomicPositions = torch.tensor(mol.xyz * 10, dtype=torch.float32, requires_grad=True, device=device)
+
+    nnp_ref = torchani.models.ANI2x(periodic_table_index=True).to(device)
+    nnp_ref = OptimizedTorchANI(nnp_ref, atomicNumbers).to(device)
+
+    energy_ref = nnp_ref((atomicNumbers, atomicPositions)).energies
+    energy_ref.backward()
+    grad_ref = atomicPositions.grad.clone()
+
+    with tempfile.NamedTemporaryFile() as fd:
+
+        torch.jit.script(nnp_ref).save(fd.name)
+        nnp = torch.jit.load(fd.name)
+
+        energy = nnp((atomicNumbers, atomicPositions)).energies
+        atomicPositions.grad.zero_()
+        energy.backward()
+        grad = atomicPositions.grad.clone()
+
+    energy_error = torch.abs((energy - energy_ref)/energy_ref)
+    grad_error = torch.max(torch.abs((grad - grad_ref)/grad_ref))
+
+    assert energy_error < 5e-7
+    assert grad_error < 5e-3

src/pytorch/__init__.py

@@ -0,0 +1 @@
+from NNPOps.OptimizedTorchANI import OptimizedTorchANI