torchani 0.2.1__tar.gz

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

@@ -0,0 +1,49 @@
+# -*- coding: utf-8 -*-
+"""TorchANI is a PyTorch implementation of `ANI`_, created and maintained by
+the `Roitberg group`_.  TorchANI contains classes like
+:class:`AEVComputer`, :class:`ANIModel`, and :class:`EnergyShifter` that can
+be pipelined to compute molecular energies from the 3D coordinates of
+molecules.  It also include tools to: deal with ANI datasets(e.g. `ANI-1`_,
+`ANI-1x`_, `ANI-1ccx`_, etc.) at :attr:`torchani.data`, import various file
+formats of NeuroChem at :attr:`torchani.neurochem`, help working with ignite
+at :attr:`torchani.ignite`, and more at :attr:`torchani.utils`.
+
+.. _ANI:
+    http://pubs.rsc.org/en/Content/ArticleLanding/2017/SC/C6SC05720A#!divAbstract
+
+.. _Roitberg group:
+    https://roitberg.chem.ufl.edu/
+
+.. _ANI-1:
+    https://www.nature.com/articles/sdata2017193
+
+.. _ANI-1x:
+    https://aip.scitation.org/doi/abs/10.1063/1.5023802
+
+.. _ANI-1ccx:
+    https://doi.org/10.26434/chemrxiv.6744440.v1
+"""
+
+from .utils import EnergyShifter
+from .nn import ANIModel, Ensemble
+from .aev import AEVComputer
+from . import ignite
+from . import utils
+from . import neurochem
+from . import data
+from pkg_resources import get_distribution, DistributionNotFound
+
+try:
+    __version__ = get_distribution(__name__).version
+except DistributionNotFound:
+    # package is not installed
+    pass
+
+__all__ = ['AEVComputer', 'EnergyShifter', 'ANIModel', 'Ensemble',
+           'ignite', 'utils', 'neurochem', 'data']
+
+try:
+    from . import ase  # noqa: F401
+    __all__.append('ase')
+except ImportError:
+    pass

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

@@ -0,0 +1,320 @@
+import torch
+import itertools
+import math
+from . import utils
+
+
+def _cutoff_cosine(distances, cutoff):
+    return torch.where(
+        distances <= cutoff,
+        0.5 * torch.cos(math.pi * distances / cutoff) + 0.5,
+        torch.zeros_like(distances)
+    )
+
+
+def default_neighborlist(species, coordinates, cutoff):
+    """Default neighborlist computer"""
+
+    vec = coordinates.unsqueeze(2) - coordinates.unsqueeze(1)
+    # vec has hape (conformations, atoms, atoms, 3) storing Rij vectors
+
+    distances = vec.norm(2, -1)
+    # distances has shape (conformations, atoms, atoms) storing Rij distances
+
+    padding_mask = (species == -1).unsqueeze(1)
+    distances = distances.masked_fill(padding_mask, math.inf)
+
+    distances, indices = distances.sort(-1)
+
+    min_distances, _ = distances.flatten(end_dim=1).min(0)
+    in_cutoff = (min_distances <= cutoff).nonzero().flatten()[1:]
+    indices = indices.index_select(-1, in_cutoff)
+
+    # TODO: remove this workaround after gather support broadcasting
+    atoms = coordinates.shape[1]
+    species_ = species.unsqueeze(1).expand(-1, atoms, -1)
+    neighbor_species = species_.gather(-1, indices)
+
+    neighbor_distances = distances.index_select(-1, in_cutoff)
+
+    # TODO: remove this workaround when gather support broadcasting
+    # https://github.com/pytorch/pytorch/pull/9532
+    indices_ = indices.unsqueeze(-1).expand(-1, -1, -1, 3)
+    neighbor_coordinates = vec.gather(-2, indices_)
+    return neighbor_species, neighbor_distances, neighbor_coordinates
+
+
+class AEVComputer(torch.nn.Module):
+    r"""The AEV computer that takes coordinates as input and outputs aevs.
+
+    Arguments:
+        Rcr (float): :math:`R_C` in equation (2) when used at equation (3)
+            in the `ANI paper`_.
+        Rca (float): :math:`R_C` in equation (2) when used at equation (4)
+            in the `ANI paper`_.
+        EtaR (:class:`torch.Tensor`): The 1D tensor of :math:`\eta` in
+            equation (3) in the `ANI paper`_.
+        ShfR (:class:`torch.Tensor`): The 1D tensor of :math:`R_s` in
+            equation (3) in the `ANI paper`_.
+        EtaA (:class:`torch.Tensor`): The 1D tensor of :math:`\eta` in
+            equation (4) in the `ANI paper`_.
+        Zeta (:class:`torch.Tensor`): The 1D tensor of :math:`\zeta` in
+            equation (4) in the `ANI paper`_.
+        ShfA (:class:`torch.Tensor`): The 1D tensor of :math:`R_s` in
+            equation (4) in the `ANI paper`_.
+        ShfZ (:class:`torch.Tensor`): The 1D tensor of :math:`\theta_s` in
+            equation (4) in the `ANI paper`_.
+        num_species (int): Number of supported atom types.
+        neighborlist_computer (:class:`collections.abc.Callable`): initial
+            value of :attr:`neighborlist`
+
+    Attributes:
+        neighborlist (:class:`collections.abc.Callable`): The callable
+            (species:Tensor, coordinates:Tensor, cutoff:float)
+            -> Tuple[Tensor, Tensor, Tensor] that returns the species,
+            distances and relative coordinates of neighbor atoms. The input
+            species and coordinates tensor have the same shape convention as
+            the input of :class:`AEVComputer`. The returned neighbor
+            species and coordinates tensor must have shape ``(C, A, N)`` and
+            ``(C, A, N, 3)`` correspoindingly, where ``C`` is the number of
+            conformations in a chunk, ``A`` is the number of atoms, and ``N``
+            is the maximum number of neighbors that an atom could have.
+
+    .. _ANI paper:
+        http://pubs.rsc.org/en/Content/ArticleLanding/2017/SC/C6SC05720A#!divAbstract
+    """
+
+    def __init__(self, Rcr, Rca, EtaR, ShfR, EtaA, Zeta, ShfA, ShfZ,
+                 num_species, neighborlist_computer=default_neighborlist):
+        super(AEVComputer, self).__init__()
+        self.Rcr = Rcr
+        self.Rca = Rca
+        # convert constant tensors to a ready-to-broadcast shape
+        # shape convension (..., EtaR, ShfR)
+        self.register_buffer('EtaR', EtaR.view(-1, 1))
+        self.register_buffer('ShfR', ShfR.view(1, -1))
+        # shape convension (..., EtaA, Zeta, ShfA, ShfZ)
+        self.register_buffer('EtaA', EtaA.view(-1, 1, 1, 1))
+        self.register_buffer('Zeta', Zeta.view(1, -1, 1, 1))
+        self.register_buffer('ShfA', ShfA.view(1, 1, -1, 1))
+        self.register_buffer('ShfZ', ShfZ.view(1, 1, 1, -1))
+
+        self.num_species = num_species
+        self.neighborlist = neighborlist_computer
+
+    def radial_sublength(self):
+        """Returns the length of radial subaev of a single species"""
+        return self.EtaR.numel() * self.ShfR.numel()
+
+    def radial_length(self):
+        """Returns the length of full radial aev"""
+        return self.num_species * self.radial_sublength()
+
+    def angular_sublength(self):
+        """Returns the length of angular subaev of a single species"""
+        return self.EtaA.numel() * self.Zeta.numel() * self.ShfA.numel() * \
+            self.ShfZ.numel()
+
+    def angular_length(self):
+        """Returns the length of full angular aev"""
+        s = self.num_species
+        return (s * (s + 1)) // 2 * self.angular_sublength()
+
+    def aev_length(self):
+        """Returns the length of full aev"""
+        return self.radial_length() + self.angular_length()
+
+    def _radial_subaev_terms(self, distances):
+        """Compute the radial subAEV terms of the center atom given neighbors
+
+        This correspond to equation (3) in the `ANI paper`_. This function just
+        compute the terms. The sum in the equation is not computed.
+        The input tensor have shape (conformations, atoms, N), where ``N``
+        is the number of neighbor atoms within the cutoff radius and output
+        tensor should have shape
+        (conformations, atoms, ``self.radial_sublength()``)
+
+        .. _ANI paper:
+            http://pubs.rsc.org/en/Content/ArticleLanding/2017/SC/C6SC05720A#!divAbstract
+        """
+        distances = distances.unsqueeze(-1).unsqueeze(-1)
+        fc = _cutoff_cosine(distances, self.Rcr)
+        # Note that in the equation in the paper there is no 0.25
+        # coefficient, but in NeuroChem there is such a coefficient.
+        # We choose to be consistent with NeuroChem instead of the paper here.
+        ret = 0.25 * torch.exp(-self.EtaR * (distances - self.ShfR)**2) * fc
+        # At this point, ret now have shape
+        # (conformations, atoms, N, ?, ?) where ? depend on constants.
+        # We then should flat the last 4 dimensions to view the subAEV as one
+        # dimension vector
+        return ret.flatten(start_dim=-2)
+
+    def _angular_subaev_terms(self, vectors1, vectors2):
+        """Compute the angular subAEV terms of the center atom given neighbor pairs.
+
+        This correspond to equation (4) in the `ANI paper`_. This function just
+        compute the terms. The sum in the equation is not computed.
+        The input tensor have shape (conformations, atoms, N), where N
+        is the number of neighbor atom pairs within the cutoff radius and
+        output tensor should have shape
+        (conformations, atoms, ``self.angular_sublength()``)
+
+        .. _ANI paper:
+            http://pubs.rsc.org/en/Content/ArticleLanding/2017/SC/C6SC05720A#!divAbstract
+        """
+        vectors1 = vectors1.unsqueeze(
+            -1).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
+        vectors2 = vectors2.unsqueeze(
+            -1).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
+        distances1 = vectors1.norm(2, dim=-5)
+        distances2 = vectors2.norm(2, dim=-5)
+
+        # 0.95 is multiplied to the cos values to prevent acos from
+        # returning NaN.
+        cos_angles = 0.95 * \
+            torch.nn.functional.cosine_similarity(
+                vectors1, vectors2, dim=-5)
+        angles = torch.acos(cos_angles)
+
+        fcj1 = _cutoff_cosine(distances1, self.Rca)
+        fcj2 = _cutoff_cosine(distances2, self.Rca)
+        factor1 = ((1 + torch.cos(angles - self.ShfZ)) / 2) ** self.Zeta
+        factor2 = torch.exp(-self.EtaA *
+                            ((distances1 + distances2) / 2 - self.ShfA) ** 2)
+        ret = 2 * factor1 * factor2 * fcj1 * fcj2
+        # At this point, ret now have shape
+        # (conformations, atoms, N, ?, ?, ?, ?) where ? depend on constants.
+        # We then should flat the last 4 dimensions to view the subAEV as one
+        # dimension vector
+        return ret.flatten(start_dim=-4)
+
+    def _terms_and_indices(self, species, coordinates):
+        """Returns radial and angular subAEV terms, these terms will be sorted
+        according to their distances to central atoms, and only these within
+        cutoff radius are valid. The returned indices stores the source of data
+        before sorting.
+        """
+        max_cutoff = max([self.Rcr, self.Rca])
+        species_, distances, vec = self.neighborlist(species, coordinates,
+                                                     max_cutoff)
+        radial_terms = self._radial_subaev_terms(distances)
+
+        vec = self._combinations(vec, -2)
+        angular_terms = self._angular_subaev_terms(*vec)
+
+        # Returned tensors has shape:
+        # (conformations, atoms, neighbors, ``self.radial_sublength()``)
+        # (conformations, atoms, pairs, ``self.angular_sublength()``)
+        # (conformations, atoms, neighbors)
+        # (conformations, atoms, pairs)
+        return radial_terms, angular_terms, species_
+
+    def _combinations(self, tensor, dim=0):
+        # TODO: remove this when combinations is merged into PyTorch
+        # https://github.com/pytorch/pytorch/pull/9393
+        n = tensor.shape[dim]
+        if n == 0:
+            return tensor, tensor
+        r = torch.arange(n, dtype=torch.long, device=tensor.device)
+        grid_x, grid_y = torch.meshgrid([r, r])
+        index1 = grid_y.masked_select(
+            torch.triu(torch.ones(n, n, device=tensor.device),
+                       diagonal=1) == 1)
+        index2 = grid_x.masked_select(
+            torch.triu(torch.ones(n, n, device=tensor.device),
+                       diagonal=1) == 1)
+        return tensor.index_select(dim, index1), \
+            tensor.index_select(dim, index2)
+
+    def _compute_mask_r(self, species_r):
+        """Get mask of radial terms for each supported species from indices"""
+        mask_r = (species_r.unsqueeze(-1) ==
+                  torch.arange(self.num_species, device=self.EtaR.device))
+        return mask_r
+
+    def _compute_mask_a(self, species_a, present_species):
+        """Get mask of angular terms for each supported species from indices"""
+        species_a1, species_a2 = self._combinations(species_a, -1)
+        mask_a1 = (species_a1.unsqueeze(-1) == present_species).unsqueeze(-1)
+        mask_a2 = (species_a2.unsqueeze(-1).unsqueeze(-1) == present_species)
+        mask = mask_a1 * mask_a2
+        mask_rev = mask.permute(0, 1, 2, 4, 3)
+        mask_a = (mask + mask_rev) > 0
+        return mask_a
+
+    def _assemble(self, radial_terms, angular_terms, present_species,
+                  mask_r, mask_a):
+        """Returns radial and angular AEV computed from terms according
+        to the given partition information.
+
+        Arguments:
+            radial_terms (:class:`torch.Tensor`): shape (conformations, atoms,
+                neighbors, ``self.radial_sublength()``)
+            angular_terms (:class:`torch.Tensor`): shape (conformations, atoms,
+                pairs, ``self.angular_sublength()``)
+            present_species (:class:`torch.Tensor`):  Long tensor for species
+                of atoms present in the molecules.
+            mask_r (:class:`torch.Tensor`): shape (conformations, atoms,
+                neighbors, supported species)
+            mask_a (:class:`torch.Tensor`): shape (conformations, atoms,
+                pairs, present species, present species)
+        """
+        conformations = radial_terms.shape[0]
+        atoms = radial_terms.shape[1]
+
+        # assemble radial subaev
+        present_radial_aevs = (
+            radial_terms.unsqueeze(-2) *
+            mask_r.unsqueeze(-1).type(radial_terms.dtype)
+        ).sum(-3)
+        # present_radial_aevs has shape
+        # (conformations, atoms, present species, radial_length)
+        radial_aevs = present_radial_aevs.flatten(start_dim=2)
+
+        # assemble angular subaev
+        rev_indices = {present_species[i].item(): i
+                       for i in range(len(present_species))}
+        angular_aevs = []
+        zero_angular_subaev = torch.zeros(
+            conformations, atoms, self.angular_sublength(),
+            dtype=self.EtaR.dtype, device=self.EtaR.device)
+        for s1, s2 in itertools.combinations_with_replacement(
+                                        range(self.num_species), 2):
+            if s1 in rev_indices and s2 in rev_indices:
+                i1 = rev_indices[s1]
+                i2 = rev_indices[s2]
+                mask = mask_a[..., i1, i2].unsqueeze(-1).type(self.EtaR.dtype)
+                subaev = (angular_terms * mask).sum(-2)
+            else:
+                subaev = zero_angular_subaev
+            angular_aevs.append(subaev)
+
+        return radial_aevs, torch.cat(angular_aevs, dim=2)
+
+    def forward(self, species_coordinates):
+        """Compute AEVs
+
+        Arguments:
+            species_coordinates (tuple): Two tensors: species and coordinates.
+                species must have shape ``(C, A)`` and coordinates must have
+                shape ``(C, A, 3)``, where ``C`` is the number of conformations
+                in a chunk, and ``A`` is the number of atoms.
+
+        Returns:
+            tuple: Species and AEVs. species are the species from the input
+            unchanged, and AEVs is a tensor of shape
+            ``(C, A, self.aev_length())``
+        """
+        species, coordinates = species_coordinates
+
+        present_species = utils.present_species(species)
+
+        radial_terms, angular_terms, species_ = \
+            self._terms_and_indices(species, coordinates)
+        mask_r = self._compute_mask_r(species_)
+        mask_a = self._compute_mask_a(species_, present_species)
+
+        radial, angular = self._assemble(radial_terms, angular_terms,
+                                         present_species, mask_r, mask_a)
+        fullaev = torch.cat([radial, angular], dim=2)
+        return species, fullaev

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

@@ -0,0 +1,149 @@
+# -*- coding: utf-8 -*-
+"""Tools for interfacing with `ASE`_.
+
+.. _ASE:
+    https://wiki.fysik.dtu.dk/ase
+"""
+
+import math
+import torch
+import ase.neighborlist
+from . import utils
+import ase.calculators.calculator
+import ase.units
+import copy
+
+
+class NeighborList:
+    """ASE neighborlist computer
+
+    Arguments:
+        cell: same as in :class:`ase.Atoms`
+        pbc: same as in :class:`ase.Atoms`
+    """
+
+    def __init__(self, cell=None, pbc=None):
+        # wrap `cell` and `pbc` with `ase.Atoms`
+        a = ase.Atoms('He', [[0, 0, 0]], cell=cell, pbc=pbc)
+        self.pbc = a.get_pbc()
+        self.cell = a.get_cell(complete=True)
+
+    def __call__(self, species, coordinates, cutoff):
+        conformations = species.shape[0]
+        max_atoms = species.shape[1]
+        neighbor_species = []
+        neighbor_distances = []
+        neighbor_vecs = []
+        for i in range(conformations):
+            s = species[i].unsqueeze(0)
+            c = coordinates[i].unsqueeze(0)
+            s, c = utils.strip_redundant_padding(s, c)
+            s = s.squeeze()
+            c = c.squeeze()
+            atoms = s.shape[0]
+            atoms_object = ase.Atoms(
+                ['He'] * atoms,  # chemical symbols are not important here
+                positions=c.detach().numpy(),
+                pbc=self.pbc,
+                cell=self.cell)
+            idx1, idx2, shift = ase.neighborlist.neighbor_list(
+                'ijS', atoms_object, cutoff)
+            # NB: The absolute distance and distance vectors computed by
+            # `neighbor_list`can not be used since it does not preserve
+            # gradient information
+            idx1 = torch.from_numpy(idx1).to(coordinates.device)
+            idx2 = torch.from_numpy(idx2).to(coordinates.device)
+            D = c.index_select(0, idx2) - c.index_select(0, idx1)
+            shift = torch.from_numpy(shift).to(coordinates.device) \
+                         .to(coordinates.dtype)
+            cell = torch.from_numpy(self.cell).to(coordinates.device) \
+                        .to(coordinates.dtype)
+            D += shift @ cell
+            d = D.norm(2, -1)
+            neighbor_species1 = []
+            neighbor_distances1 = []
+            neighbor_vecs1 = []
+            for i in range(atoms):
+                this_atom_indices = (idx1 == i).nonzero().flatten()
+                neighbor_indices = idx2[this_atom_indices]
+                neighbor_species1.append(s[neighbor_indices])
+                neighbor_distances1.append(d[this_atom_indices])
+                neighbor_vecs1.append(D.index_select(0, this_atom_indices))
+            for i in range(max_atoms - atoms):
+                neighbor_species1.append(torch.full((1,), -1))
+                neighbor_distances1.append(torch.full((1,), math.inf))
+                neighbor_vecs1.append(torch.full((1, 3), 0))
+            neighbor_species1 = torch.nn.utils.rnn.pad_sequence(
+                neighbor_species1, padding_value=-1)
+            neighbor_distances1 = torch.nn.utils.rnn.pad_sequence(
+                neighbor_distances1, padding_value=math.inf)
+            neighbor_vecs1 = torch.nn.utils.rnn.pad_sequence(
+                neighbor_vecs1, padding_value=0)
+            neighbor_species.append(neighbor_species1)
+            neighbor_distances.append(neighbor_distances1)
+            neighbor_vecs.append(neighbor_vecs1)
+        neighbor_species = torch.nn.utils.rnn.pad_sequence(
+            neighbor_species, batch_first=True, padding_value=-1)
+        neighbor_distances = torch.nn.utils.rnn.pad_sequence(
+            neighbor_distances, batch_first=True, padding_value=math.inf)
+        neighbor_vecs = torch.nn.utils.rnn.pad_sequence(
+            neighbor_vecs, batch_first=True, padding_value=0)
+        return neighbor_species.permute(0, 2, 1), \
+            neighbor_distances.permute(0, 2, 1), \
+            neighbor_vecs.permute(0, 2, 1, 3)
+
+
+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.
+        _default_neighborlist (bool): Whether to ignore pbc setting and always
+            use default neighborlist computer. This is for internal use only.
+    """
+
+    implemented_properties = ['energy', 'forces']
+
+    def __init__(self, species, aev_computer, model, energy_shifter,
+                 _default_neighborlist=False):
+        super(Calculator, self).__init__()
+        self._default_neighborlist = _default_neighborlist
+        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.copy(aev_computer)
+        self.model = model
+        self.energy_shifter = energy_shifter
+
+        self.device = self.aev_computer.EtaR.device
+        self.dtype = self.aev_computer.EtaR.dtype
+
+        self.whole = torch.nn.Sequential(
+            self.aev_computer,
+            self.model,
+            self.energy_shifter
+        )
+
+    def calculate(self, atoms=None, properties=['energy'],
+                  system_changes=ase.calculators.calculator.all_changes):
+        super(Calculator, self).calculate(atoms, properties, system_changes)
+        if not self._default_neighborlist:
+            self.aev_computer.neighborlist = NeighborList(
+                cell=self.atoms.get_cell(complete=True),
+                pbc=self.atoms.get_pbc())
+        species = self.species_to_tensor(self.atoms.get_chemical_symbols())
+        species = species.unsqueeze(0)
+        coordinates = torch.tensor(self.atoms.get_positions(wrap=True))
+        coordinates = coordinates.unsqueeze(0).to(self.device).to(self.dtype) \
+                                 .requires_grad_('forces' in properties)
+        _, energy = self.whole((species, coordinates))
+        energy *= ase.units.Hartree
+        self.results['energy'] = energy.item()
+        if 'forces' in properties:
+            forces = -torch.autograd.grad(energy.squeeze(), coordinates)[0]
+            self.results['forces'] = forces.squeeze().numpy()