Intuitive aev computer classmethod (#18) (#543)

* Add classmethod for conveniently building AEVComputer

* Fix bug

* Add test

* flake8

* Small change in assert

* rename to zeta for consistency with other code

* trigger tests

* What the hell Dropbox?! -.-
Co-authored-by: Ignacio Pickering <ign.pickering@gmail.com>

tests/test_aev.py

@@ -16,6 +16,31 @@ const_file = os.path.join(path, '../torchani/resources/ani-1x_8x/rHCNO-5.2R_16-3
 N = 97
 
 
+class TestAEVConstructor(unittest.TestCase):
+    # Test that checks that the friendly constructor
+    # reproduces the values from ANI1x with the correct parameters
+    def testCoverLinearly(self):
+        consts = torchani.neurochem.Constants(const_file)
+        aev_computer = torchani.AEVComputer(**consts)
+        ani1x_values = {'radial_cutoff': 5.2,
+                        'angular_cutoff': 3.5,
+                        'radial_eta': 16.0,
+                        'angular_eta': 8.0,
+                        'radial_dist_divisions': 16,
+                        'angular_dist_divisions': 4,
+                        'zeta': 32.0,
+                        'angle_sections': 8,
+                        'num_species': 4}
+        aev_computer_alt = torchani.AEVComputer.cover_linearly(**ani1x_values)
+        constants = aev_computer.constants()
+        constants_alt = aev_computer_alt.constants()
+        for c, ca in zip(constants, constants_alt):
+            if isinstance(c, torch.Tensor):
+                self.assertTrue(torch.isclose(c, ca).all())
+            else:
+                self.assertEqual(c, ca)
+
+
 class TestIsolated(unittest.TestCase):
     # Tests that there is no error when atoms are separated
     # a distance greater than the cutoff radius from all other atoms

torchani/aev.py

 import torch
+
 from torch import Tensor
 import math
 from typing import Tuple, Optional, NamedTuple
@@ -390,6 +391,40 @@ class AEVComputer(torch.nn.Module):
         self.register_buffer('default_cell', default_cell)
         self.register_buffer('default_shifts', default_shifts)
 
+    @classmethod
+    def cover_linearly(cls, radial_cutoff: float, angular_cutoff: float,
+                       radial_eta: float, angular_eta: float,
+                       radial_dist_divisions: int, angular_dist_divisions: int,
+                       zeta: float, angle_sections: int, num_species: int,
+                       angular_start: float = 0.9, radial_start: float = 0.9):
+        r""" Provides a convenient way to linearly fill cutoffs
+
+        This is a user friendly constructor that builds an
+        :class:`torchani.AEVComputer` where the subdivisions along the the
+        distance dimension for the angular and radial sub-AEVs, and the angle
+        sections for the angular sub-AEV, are linearly covered with shifts. By
+        default the distance shifts start at 0.9 Angstroms.
+
+        To reproduce the ANI-1x AEV's the signature ``(5.2, 3.5, 16.0, 8.0, 16, 4, 32.0, 8, 4)``
+        can be used.
+        """
+        # This is intended to be self documenting code that explains the way
+        # the AEV parameters for ANI1x were chosen. This is not necessarily the
+        # best or most optimal way but it is a relatively sensible default.
+        Rcr = radial_cutoff
+        Rca = angular_cutoff
+        EtaR = torch.tensor([float(radial_eta)])
+        EtaA = torch.tensor([float(angular_eta)])
+        Zeta = torch.tensor([float(zeta)])
+
+        ShfR = torch.linspace(radial_start, radial_cutoff, radial_dist_divisions + 1)[:-1]
+        ShfA = torch.linspace(angular_start, angular_cutoff, angular_dist_divisions + 1)[:-1]
+        angle_start = math.pi / (2 * angle_sections)
+
+        ShfZ = (torch.linspace(0, math.pi, angle_sections + 1) + angle_start)[:-1]
+
+        return cls(Rcr, Rca, EtaR, ShfR, EtaA, Zeta, ShfA, ShfZ, num_species)
+
     def constants(self):
         return self.Rcr, self.EtaR, self.ShfR, self.Rca, self.ShfZ, self.EtaA, self.Zeta, self.ShfA