torchani 0.4__tar.gz → 0.6__tar.gz

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/neurochem/__init__.py

@@ -556,7 +556,6 @@ if sys.version_info[0] > 2:
             assert_param('runtype', 'ANNP_CREATE_HDNN_AND_TRAIN')
             assert_param('adptlrn', 'OFF')
             assert_param('tmax', 0)
-            assert_param('nmax', 0)
             assert_param('ntwshr', 0)
 
             # load parameters
@@ -585,6 +584,8 @@ if sys.version_info[0] > 2:
             del params['tbtchsz']
             self.validation_batch_size = params['vbtchsz']
             del params['vbtchsz']
+            self.nmax = params['nmax']
+            del params['nmax']
 
             # construct networks
             input_size, network_setup = network_setup
@@ -700,6 +701,12 @@ if sys.version_info[0] > 2:
                     self.global_epoch = trainer.state.epoch
                     self.global_iteration = trainer.state.iteration
 
+                if self.nmax > 0:
+                    @trainer.on(ignite.engine.Events.EPOCH_COMPLETED)
+                    def terminate_when_nmax_reaches(trainer):
+                        if trainer.state.epoch >= self.nmax:
+                            trainer.terminate()
+
                 if self.tqdm is not None:
                     @trainer.on(ignite.engine.Events.EPOCH_STARTED)
                     def init_tqdm(trainer):

opt/hostedtoolcache/Python/3.7.2/x64/lib/python3.7/site-packages/torchani/resources/ani-1ccx_8x.info

@@ -0,0 +1,5 @@
+ani-1ccx_8x/rHCNO-5.2R_16-3.5A_a4-8.params
+ani-1ccx_8x/sae_linfit.dat
+ani-1ccx_8x/train
+8
+

opt/hostedtoolcache/Python/3.7.2/x64/lib/python3.7/site-packages/torchani/resources/ani-1x_8x.info

@@ -0,0 +1,4 @@
+ani-1x_8x/rHCNO-5.2R_16-3.5A_a4-8.params
+ani-1x_8x/sae_linfit.dat
+ani-1x_8x/train
+8

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

@@ -1,4 +1,5 @@
 import torch
+import math
 
 
 def pad(species):
@@ -205,5 +206,61 @@ class ChemicalSymbolsToInts:
         return torch.tensor(rev, dtype=torch.long)
 
 
-__all__ = ['pad', 'pad_coordinates', 'present_species',
-           'strip_redundant_padding', 'ChemicalSymbolsToInts']
+def hessian(coordinates, energies=None, forces=None):
+    """Compute analytical hessian from the energy graph or force graph.
+
+    Arguments:
+        coordinates (:class:`torch.Tensor`): Tensor of shape `(molecules, atoms, 3)`
+        energies (:class:`torch.Tensor`): Tensor of shape `(molecules,)`, if specified,
+            then `forces` must be `None`. This energies must be computed from
+            `coordinates` in a graph.
+        forces (:class:`torch.Tensor`): Tensor of shape `(molecules, atoms, 3)`, if specified,
+            then `energies` must be `None`. This forces must be computed from
+            `coordinates` in a graph.
+
+    Returns:
+        :class:`torch.Tensor`: Tensor of shape `(molecules, 3A, 3A)` where A is the number of
+        atoms in each molecule
+    """
+    if energies is None and forces is None:
+        raise ValueError('Energies or forces must be specified')
+    if energies is not None and forces is not None:
+        raise ValueError('Energies or forces can not be specified at the same time')
+    if forces is None:
+        forces = -torch.autograd.grad(energies.sum(), coordinates, create_graph=True)[0]
+    flattened_force = forces.flatten(start_dim=1)
+    force_components = flattened_force.unbind(dim=1)
+    return -torch.stack([
+        torch.autograd.grad(f.sum(), coordinates, retain_graph=True)[0].flatten(start_dim=1)
+        for f in force_components
+    ], dim=1)
+
+
+def vibrational_analysis(masses, hessian, unit='cm^-1'):
+    """Computing the vibrational wavenumbers from hessian."""
+    if unit != 'cm^-1':
+        raise ValueError('Only cm^-1 are supported right now')
+    assert hessian.shape[0] == 1, 'Currently only supporting computing one molecule a time'
+    # Solving the eigenvalue problem: Hq = w^2 * T q
+    # where H is the Hessian matrix, q is the normal coordinates,
+    # T = diag(m1, m1, m1, m2, m2, m2, ....) is the mass
+    # We solve this eigenvalue problem through Lowdin diagnolization:
+    # Hq = w^2 * Tq ==> Hq = w^2 * T^(1/2) T^(1/2) q
+    # Letting q' = T^(1/2) q, we then have
+    # T^(-1/2) H T^(1/2) q' = w^2 * q'
+    inv_sqrt_mass = (1 / masses.sqrt()).repeat_interleave(3, dim=1)  # shape (molecule, 3 * atoms)
+    mass_scaled_hessian = hessian * inv_sqrt_mass.unsqueeze(1) * inv_sqrt_mass.unsqueeze(2)
+    if mass_scaled_hessian.shape[0] != 1:
+        raise ValueError('The input should contain only one molecule')
+    mass_scaled_hessian = mass_scaled_hessian.squeeze(0)
+    eigenvalues = torch.symeig(mass_scaled_hessian).eigenvalues
+    angular_frequencies = eigenvalues.sqrt()
+    frequencies = angular_frequencies / (2 * math.pi)
+    # converting from sqrt(hartree / (amu * angstrom^2)) to cm^-1
+    wavenumbers = frequencies * 17092
+    return wavenumbers
+
+
+__all__ = ['pad', 'pad_coordinates', 'present_species', 'hessian',
+           'vibrational_analysis', 'strip_redundant_padding',
+           'ChemicalSymbolsToInts']

opt/hostedtoolcache/Python/{3.7.0/x64/lib/python3.7/site-packages/torchani-0.4-py3.7.egg-info → 3.7.2/x64/lib/python3.7/site-packages/torchani-0.6-py3.7.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 1.0
 Name: torchani
-Version: 0.4
+Version: 0.6
 Summary: PyTorch implementation of ANI
 Home-page: https://github.com/zasdfgbnm/torchani
 Author: Xiang Gao