Optimize EnergyShifter of TorchANI

CMakeLists.txt

@@ -30,10 +30,12 @@ foreach(TEST_PATH ${TEST_PATHS})
 endforeach()
 
 add_test(TestBatchedNN         pytest -v ${CMAKE_SOURCE_DIR}/src/pytorch/TestBatchedNN.py)
+add_test(TestEnergyShifter     pytest -v ${CMAKE_SOURCE_DIR}/src/pytorch/TestEnergyShifter.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/EnergyShifter.py
               src/pytorch/SymmetryFunctions.py
         DESTINATION ${Python_SITEARCH}/${NAME})

README.md

@@ -77,6 +77,7 @@ import torchani
 
 from NNPOps.SymmetryFunctions import TorchANISymmetryFunctions
 from NNPOps.BatchedNN import TorchANIBatchedNN
+from NNPOps.EnergyShifter import TorchANIEnergyShifter
 
 device = torch.device('cuda')
 
@@ -89,6 +90,7 @@ positions = torch.tensor(molecule.xyz * 10, dtype=torch.float32, requires_grad=T
 nnp = torchani.models.ANI2x(periodic_table_index=True).to(device)
 nnp.aev_computer = TorchANISymmetryFunctions(nnp.aev_computer).to(device)
 nnp.neural_networks = TorchANIBatchedNN(nnp.species_converter, nnp.neural_networks, species).to(device)
+nnp.energy_shifter = TorchANIEnergyShifter(nnp.species_converter, nnp.energy_shifter, species).to(device)
 
 # Compute energy and forces
 energy = nnp((species, positions)).energies

src/pytorch/EnergyShifter.py

@@ -0,0 +1,53 @@
+#
+# 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 NamedTuple, Optional, Tuple
+
+class SpeciesEnergies(NamedTuple):
+    species: Tensor
+    energies: Tensor
+
+class TorchANIEnergyShifter(torch.nn.Module):
+
+    from torchani.nn import SpeciesConverter # https://github.com/openmm/NNPOps/issues/44
+    from torchani.utils import EnergyShifter # https://github.com/openmm/NNPOps/issues/44
+
+    def __init__(self, converter: SpeciesConverter, shifter: EnergyShifter, atomicNumbers: Tensor) -> None:
+
+        super().__init__()
+
+        # Convert atomic numbers to a list of species
+        species = converter((atomicNumbers, torch.empty(0))).species
+
+        # Compute atomic self energies
+        self.register_buffer('self_energies', shifter.sae(species))
+
+    def forward(self, species_energies: Tuple[Tensor, Tensor],
+                cell: Optional[Tensor] = None,
+                pbc: Optional[Tensor] = None) -> SpeciesEnergies:
+
+        species, energies = species_energies
+
+        return SpeciesEnergies(species, energies + self.self_energies)

src/pytorch/TestEnergyShifter.py

@@ -0,0 +1,105 @@
+#
+# 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')
+
+def test_import():
+    import NNPOps
+    import NNPOps.EnergyShifter
+
+@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.EnergyShifter import TorchANIEnergyShifter
+
+    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.energy_shifter = TorchANIEnergyShifter(nnp.species_converter, nnp.energy_shifter, 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
+    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.EnergyShifter import TorchANIEnergyShifter
+
+    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.energy_shifter = TorchANIEnergyShifter(nnp_ref.species_converter, nnp_ref.energy_shifter, 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