torchani 0.4__tar.gz → 0.7__tar.gz

opt/hostedtoolcache/Python/{3.7.0 → 3.7.2}/x64/lib/python3.7/site-packages/torchani/__init__.py

@@ -30,11 +30,9 @@ from .aev import AEVComputer
 from . import utils
 from . import neurochem
 from . import models
+from . import optim
 from pkg_resources import get_distribution, DistributionNotFound
 import sys
-if sys.version_info[0] > 2:
-    from . import ignite
-    from . import data
 
 try:
     __version__ = get_distribution(__name__).version
@@ -43,10 +41,20 @@ except DistributionNotFound:
     pass
 
 __all__ = ['AEVComputer', 'EnergyShifter', 'ANIModel', 'Ensemble',
-           'ignite', 'utils', 'neurochem', 'data', 'models']
+           'utils', 'neurochem', 'models', 'optim']
 
 try:
     from . import ase  # noqa: F401
     __all__.append('ase')
 except ImportError:
     pass
+
+
+if sys.version_info[0] > 2:
+    try:
+        from . import ignite  # noqa: F401
+        __all__.append('ignite')
+        from . import data  # noqa: F401
+        __all__.append('data')
+    except ImportError:
+        pass

opt/hostedtoolcache/Python/{3.7.0 → 3.7.2}/x64/lib/python3.7/site-packages/torchani/aev.py

@@ -169,10 +169,11 @@ def neighbor_pairs(padding_mask, coordinates, cell, shifts, cutoff):
     distances.masked_fill_(padding_mask, math.inf)
     in_cutoff = (distances <= cutoff).nonzero()
     molecule_index, pair_index = in_cutoff.unbind(1)
+    molecule_index *= num_atoms
     atom_index1 = p1_all[pair_index]
     atom_index2 = p2_all[pair_index]
     shifts = shifts_all.index_select(0, pair_index)
-    return molecule_index, atom_index1, atom_index2, shifts
+    return molecule_index + atom_index1, molecule_index + atom_index2, shifts
 
 
 # torch.jit.script
@@ -219,7 +220,7 @@ def cumsum_from_zero(input_):
 
 
 # torch.jit.script
-def triple_by_molecule(molecule_index, atom_index1, atom_index2):
+def triple_by_molecule(atom_index1, atom_index2):
     """Input: indices for pairs of atoms that are close to each other.
     each pair only appear once, i.e. only one of the pairs (1, 2) and
     (2, 1) exists.
@@ -230,45 +231,40 @@ def triple_by_molecule(molecule_index, atom_index1, atom_index2):
     central atom 0, 1, 2, 3, 4 and for cental atom 0, its pairs of neighbors
     are (1, 2), (1, 3), (1, 4), (2, 3), (2, 4), (3, 4)
     """
-    # convert representation from pair to central-other
-    n = molecule_index.shape[0]
-    mi = molecule_index.repeat(2)
+    # convert representation from pair to central-others
+    n = atom_index1.shape[0]
     ai1 = torch.cat([atom_index1, atom_index2])
+    sorted_ai1, rev_indices = ai1.sort()
 
     # sort and compute unique key
-    mi_ai1 = torch.stack([mi, ai1], dim=1)
-    m_ac, rev_indices, counts = torch._unique_dim2_temporary_will_remove_soon(mi_ai1, dim=0, sorted=True, return_inverse=True, return_counts=True)
-    uniqued_molecule_index, uniqued_central_atom_index = m_ac.unbind(1)
+    uniqued_central_atom_index, counts = torch.unique_consecutive(sorted_ai1, return_counts=True)
 
     # do local combinations within unique key, assuming sorted
     pair_sizes = counts * (counts - 1) // 2
     total_size = pair_sizes.sum()
-    molecule_index = torch.repeat_interleave(uniqued_molecule_index, pair_sizes)
-    central_atom_index = torch.repeat_interleave(uniqued_central_atom_index, pair_sizes)
+    pair_indices = torch.repeat_interleave(pair_sizes)
+    central_atom_index = uniqued_central_atom_index.index_select(0, pair_indices)
     cumsum = cumsum_from_zero(pair_sizes)
-    cumsum = torch.repeat_interleave(cumsum, pair_sizes)
-    sorted_local_pair_index = torch.arange(total_size, device=molecule_index.device) - cumsum
+    cumsum = cumsum.index_select(0, pair_indices)
+    sorted_local_pair_index = torch.arange(total_size, device=cumsum.device) - cumsum
     sorted_local_index1, sorted_local_index2 = convert_pair_index(sorted_local_pair_index)
     cumsum = cumsum_from_zero(counts)
-    cumsum = torch.repeat_interleave(cumsum, pair_sizes)
+    cumsum = cumsum.index_select(0, pair_indices)
     sorted_local_index1 += cumsum
     sorted_local_index2 += cumsum
 
     # unsort result from last part
-    argsort = rev_indices.argsort()
-    local_index1 = argsort[sorted_local_index1]
-    local_index2 = argsort[sorted_local_index2]
+    local_index1 = rev_indices[sorted_local_index1]
+    local_index2 = rev_indices[sorted_local_index2]
 
     # compute mapping between representation of central-other to pair
-    sign1 = torch.where(local_index1 < n, torch.ones_like(local_index1), -torch.ones_like(local_index1))
-    sign2 = torch.where(local_index2 < n, torch.ones_like(local_index2), -torch.ones_like(local_index2))
-    pair_index1 = torch.where(local_index1 < n, local_index1, local_index1 - n)
-    pair_index2 = torch.where(local_index2 < n, local_index2, local_index2 - n)
-    return molecule_index, central_atom_index, pair_index1, pair_index2, sign1, sign2
+    sign1 = ((local_index1 < n) * 2).to(torch.long) - 1
+    sign2 = ((local_index2 < n) * 2).to(torch.long) - 1
+    return central_atom_index, local_index1 % n, local_index2 % n, sign1, sign2
 
 
 # torch.jit.script
-def compute_aev(species, coordinates, cell, pbc_switch, triu_index, constants, sizes):
+def compute_aev(species, coordinates, cell, shifts, triu_index, constants, sizes):
     Rcr, EtaR, ShfR, Rca, ShfZ, EtaA, Zeta, ShfA = constants
     num_species, radial_sublength, radial_length, angular_sublength, angular_length, aev_length = sizes
     num_molecules = species.shape[0]
@@ -276,34 +272,35 @@ def compute_aev(species, coordinates, cell, pbc_switch, triu_index, constants, s
     num_species_pairs = angular_length // angular_sublength
     cutoff = max(Rcr, Rca)
 
-    shifts = compute_shifts(cell, pbc_switch, cutoff)
-    molecule_index, atom_index1, atom_index2, shifts = neighbor_pairs(species == -1, coordinates, cell, shifts, cutoff)
-    species1 = species[molecule_index, atom_index1]
-    species2 = species[molecule_index, atom_index2]
+    atom_index1, atom_index2, shifts = neighbor_pairs(species == -1, coordinates, cell, shifts, cutoff)
+    species = species.flatten()
+    coordinates = coordinates.flatten(0, 1)
+    species1 = species[atom_index1]
+    species2 = species[atom_index2]
     shift_values = torch.mm(shifts.to(cell.dtype), cell)
 
-    vec = coordinates[molecule_index, atom_index1, :] - coordinates[molecule_index, atom_index2, :] + shift_values
+    vec = coordinates.index_select(0, atom_index1) - coordinates.index_select(0, atom_index2) + shift_values
     distances = vec.norm(2, -1)
 
     # compute radial aev
     radial_terms_ = radial_terms(Rcr, EtaR, ShfR, distances)
     radial_aev = radial_terms_.new_zeros(num_molecules * num_atoms * num_species, radial_sublength)
-    index1 = (molecule_index * num_atoms + atom_index1) * num_species + species2
-    index2 = (molecule_index * num_atoms + atom_index2) * num_species + species1
-    radial_aev.scatter_add_(0, index1.unsqueeze(1).expand(-1, radial_sublength), radial_terms_)
-    radial_aev.scatter_add_(0, index2.unsqueeze(1).expand(-1, radial_sublength), radial_terms_)
+    index1 = atom_index1 * num_species + species2
+    index2 = atom_index2 * num_species + species1
+    radial_aev.index_add_(0, index1, radial_terms_)
+    radial_aev.index_add_(0, index2, radial_terms_)
     radial_aev = radial_aev.reshape(num_molecules, num_atoms, radial_length)
 
     # compute angular aev
-    molecule_index, central_atom_index, pair_index1, pair_index2, sign1, sign2 = triple_by_molecule(molecule_index, atom_index1, atom_index2)
+    central_atom_index, pair_index1, pair_index2, sign1, sign2 = triple_by_molecule(atom_index1, atom_index2)
     vec1 = vec.index_select(0, pair_index1) * sign1.unsqueeze(1).to(vec.dtype)
     vec2 = vec.index_select(0, pair_index2) * sign2.unsqueeze(1).to(vec.dtype)
     species1_ = torch.where(sign1 == 1, species2[pair_index1], species1[pair_index1])
     species2_ = torch.where(sign2 == 1, species2[pair_index2], species1[pair_index2])
     angular_terms_ = angular_terms(Rca, ShfZ, EtaA, Zeta, ShfA, vec1, vec2)
     angular_aev = angular_terms_.new_zeros(num_molecules * num_atoms * num_species_pairs, angular_sublength)
-    index = (molecule_index * num_atoms + central_atom_index) * num_species_pairs + triu_index[species1_, species2_]
-    angular_aev.scatter_add_(0, index.unsqueeze(1).expand(-1, angular_sublength), angular_terms_)
+    index = central_atom_index * num_species_pairs + triu_index[species1_, species2_]
+    angular_aev.index_add_(0, index, angular_terms_)
     angular_aev = angular_aev.reshape(num_molecules, num_atoms, angular_length)
     return torch.cat([radial_aev, angular_aev], dim=-1)
 
@@ -366,6 +363,15 @@ class AEVComputer(torch.nn.Module):
 
         self.register_buffer('triu_index', triu_index(num_species))
 
+        # Set up default cell and compute default shifts.
+        # These values are used when cell and pbc switch are not given.
+        cutoff = max(self.Rcr, self.Rca)
+        default_cell = torch.eye(3, dtype=self.EtaR.dtype, device=self.EtaR.device)
+        default_pbc = torch.zeros(3, dtype=torch.uint8, device=self.EtaR.device)
+        default_shifts = compute_shifts(default_cell, default_pbc, cutoff)
+        self.register_buffer('default_cell', default_cell)
+        self.register_buffer('default_shifts', default_shifts)
+
     def constants(self):
         return self.Rcr, self.EtaR, self.ShfR, self.Rca, self.ShfZ, self.EtaA, self.Zeta, self.ShfA
 
@@ -403,9 +409,11 @@ class AEVComputer(torch.nn.Module):
         """
         if len(input) == 2:
             species, coordinates = input
-            cell = torch.eye(3, dtype=self.EtaR.dtype, device=self.EtaR.device)
-            pbc = torch.zeros(3, dtype=torch.uint8, device=self.EtaR.device)
+            cell = self.default_cell
+            shifts = self.default_shifts
         else:
             assert len(input) == 4
             species, coordinates, cell, pbc = input
-        return species, compute_aev(species, coordinates, cell, pbc, self.triu_index, self.constants(), self.sizes)
+            cutoff = max(self.Rcr, self.Rca)
+            shifts = compute_shifts(cell, pbc, cutoff)
+        return species, compute_aev(species, coordinates, cell, shifts, self.triu_index, self.constants(), self.sizes)

opt/hostedtoolcache/Python/3.7.2/x64/lib/python3.7/site-packages/torchani/ase.py

@@ -0,0 +1,115 @@
+# -*- coding: utf-8 -*-
+"""Tools for interfacing with `ASE`_.
+
+.. _ASE:
+    https://wiki.fysik.dtu.dk/ase
+"""
+
+from __future__ import absolute_import
+import torch
+import ase.neighborlist
+from . import utils
+import ase.calculators.calculator
+import ase.units
+import copy
+import numpy
+
+
+class Calculator(ase.calculators.calculator.Calculator):
+    """TorchANI calculator for ASE
+
+    Arguments:
+        species (:class:`collections.abc.Sequence` of :class:`str`):
+            sequence of all supported species, in order.
+        aev_computer (:class:`torchani.AEVComputer`): AEV computer.
+        model (:class:`torchani.ANIModel` or :class:`torchani.Ensemble`):
+            neural network potential models.
+        energy_shifter (:class:`torchani.EnergyShifter`): Energy shifter.
+        dtype (:class:`torchani.EnergyShifter`): data type to use,
+            by dafault ``torch.float64``.
+        overwrite (bool): After wrapping atoms into central box, whether
+            to replace the original positions stored in :class:`ase.Atoms`
+            object with the wrapped positions.
+    """
+
+    implemented_properties = ['energy', 'forces', 'stress', 'free_energy']
+
+    def __init__(self, species, aev_computer, model, energy_shifter, dtype=torch.float64, overwrite=False):
+        super(Calculator, self).__init__()
+        self.species_to_tensor = utils.ChemicalSymbolsToInts(species)
+        # aev_computer.neighborlist will be changed later, so we need a copy to
+        # make sure we do not change the original object
+        self.aev_computer = copy.deepcopy(aev_computer)
+        self.model = copy.deepcopy(model)
+        self.energy_shifter = copy.deepcopy(energy_shifter)
+        self.overwrite = overwrite
+
+        self.device = self.aev_computer.EtaR.device
+        self.dtype = dtype
+
+        self.whole = torch.nn.Sequential(
+            self.aev_computer,
+            self.model,
+            self.energy_shifter
+        ).to(dtype)
+
+    @staticmethod
+    def strain(tensor, displacement, surface_normal_axis):
+        rest_axes = {0, 1, 2} - set([surface_normal_axis])
+        displacement_normal = displacement[surface_normal_axis]
+        displacement_of_tensor = torch.zeros_like(tensor)
+        displacement_of_tensor[..., surface_normal_axis] = tensor[..., surface_normal_axis] * displacement_normal
+        for axis in rest_axes:
+            displacement_axis = displacement[axis]
+            displacement_of_tensor[..., axis] = tensor[..., surface_normal_axis] * displacement_axis
+        return displacement_of_tensor
+
+    def calculate(self, atoms=None, properties=['energy'],
+                  system_changes=ase.calculators.calculator.all_changes):
+        super(Calculator, self).calculate(atoms, properties, system_changes)
+        cell = torch.tensor(self.atoms.get_cell(complete=True),
+                            dtype=self.dtype, device=self.device)
+        pbc = torch.tensor(self.atoms.get_pbc().astype(numpy.uint8), dtype=torch.uint8,
+                           device=self.device)
+        pbc_enabled = bool(pbc.any().item())
+        species = self.species_to_tensor(self.atoms.get_chemical_symbols()).to(self.device)
+        species = species.unsqueeze(0)
+        coordinates = torch.tensor(self.atoms.get_positions())
+        coordinates = coordinates.unsqueeze(0).to(self.device).to(self.dtype) \
+                                 .requires_grad_('forces' in properties)
+
+        if pbc_enabled:
+            coordinates = utils.map2central(cell, coordinates, pbc)
+            if self.overwrite and atoms is not None:
+                atoms.set_positions(coordinates.detach().cpu().reshape(-1, 3).numpy())
+
+        if 'stress' in properties:
+            displacements = torch.zeros(3, 3, requires_grad=True,
+                                        dtype=self.dtype, device=self.device)
+            displacement_x, displacement_y, displacement_z = displacements
+            strain_x = self.strain(coordinates, displacement_x, 0)
+            strain_y = self.strain(coordinates, displacement_y, 1)
+            strain_z = self.strain(coordinates, displacement_z, 2)
+            coordinates = coordinates + strain_x + strain_y + strain_z
+
+        if pbc_enabled:
+            if 'stress' in properties:
+                strain_x = self.strain(cell, displacement_x, 0)
+                strain_y = self.strain(cell, displacement_y, 1)
+                strain_z = self.strain(cell, displacement_z, 2)
+                cell = cell + strain_x + strain_y + strain_z
+            _, energy = self.whole((species, coordinates, cell, pbc))
+        else:
+            _, energy = self.whole((species, coordinates))
+        energy *= ase.units.Hartree
+        self.results['energy'] = energy.item()
+        self.results['free_energy'] = energy.item()
+
+        if 'forces' in properties:
+            forces = -torch.autograd.grad(energy.squeeze(), coordinates)[0]
+            self.results['forces'] = forces.squeeze().to('cpu').numpy()
+
+        if 'stress' in properties:
+            volume = self.atoms.get_volume()
+            stress = torch.autograd.grad(energy.squeeze(), displacements)[0] / volume
+            self.results['stress'] = stress.cpu().numpy()

opt/hostedtoolcache/Python/{3.7.0 → 3.7.2}/x64/lib/python3.7/site-packages/torchani/ignite.py

@@ -4,7 +4,7 @@ from __future__ import absolute_import
 import torch
 from . import utils
 from torch.nn.modules.loss import _Loss
-from ignite.metrics import Metric, RootMeanSquaredError
+from ignite.metrics import Metric, RootMeanSquaredError, MeanAbsoluteError
 from ignite.contrib.metrics.regression import MaximumAbsoluteError
 
 
@@ -111,10 +111,15 @@ def RMSEMetric(key):
     return DictMetric(key, RootMeanSquaredError())
 
 
-def MAEMetric(key):
+def MaxAEMetric(key):
     """Create max absolute error metric on key."""
     return DictMetric(key, MaximumAbsoluteError())
 
 
+def MAEMetric(key):
+    """Create max absolute error metric on key."""
+    return DictMetric(key, MeanAbsoluteError())
+
+
 __all__ = ['Container', 'MSELoss', 'TransformedLoss', 'RMSEMetric',
-           'MAEMetric']
+           'MaxAEMetric']