torchani 0.6__tar.gz → 2.1.1__tar.gz

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

@@ -1,12 +1,14 @@
 # -*- coding: utf-8 -*-
-"""TorchANI is a PyTorch implementation of `ANI`_, created and maintained by
+"""`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-1x`_, `ANI-1ccx`_, `ANI-2x`_) at :attr:`torchani.data`, import various file
+formats of NeuroChem at :attr:`torchani.neurochem`, and more at :attr:`torchani.utils`.
+
+.. _TorchANI:
+    https://doi.org/10.26434/chemrxiv.12218294.v1
 
 .. _ANI:
     http://pubs.rsc.org/en/Content/ArticleLanding/2017/SC/C6SC05720A#!divAbstract
@@ -22,19 +24,19 @@ at :attr:`torchani.ignite`, and more at :attr:`torchani.utils`.
 
 .. _ANI-1ccx:
     https://doi.org/10.26434/chemrxiv.6744440.v1
+
+.. _ANI-2x:
+    https://doi.org/10.26434/chemrxiv.11819268.v1
 """
 
 from .utils import EnergyShifter
-from .nn import ANIModel, Ensemble
+from .nn import ANIModel, Ensemble, SpeciesConverter
 from .aev import AEVComputer
 from . import utils
 from . import neurochem
 from . import models
+from . import units
 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
@@ -42,11 +44,17 @@ except DistributionNotFound:
     # package is not installed
     pass
 
-__all__ = ['AEVComputer', 'EnergyShifter', 'ANIModel', 'Ensemble',
-           'ignite', 'utils', 'neurochem', 'data', 'models']
+__all__ = ['AEVComputer', 'EnergyShifter', 'ANIModel', 'Ensemble', 'SpeciesConverter',
+           'utils', 'neurochem', 'models', 'units']
 
 try:
     from . import ase  # noqa: F401
     __all__.append('ase')
 except ImportError:
     pass
+
+try:
+    from . import data  # noqa: F401
+    __all__.append('data')
+except ImportError:
+    pass

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

@@ -1,24 +1,21 @@
-from __future__ import division
 import torch
-from . import _six  # noqa:F401
-import math
 from torch import Tensor
-from typing import Tuple
+import math
+from typing import Tuple, Optional, NamedTuple
+from torch.jit import Final
+
 
+class SpeciesAEV(NamedTuple):
+    species: Tensor
+    aevs: Tensor
 
-# @torch.jit.script
-def cutoff_cosine(distances, cutoff):
-    # type: (Tensor, float) -> Tensor
-    return torch.where(
-        distances <= cutoff,
-        0.5 * torch.cos(math.pi * distances / cutoff) + 0.5,
-        torch.zeros_like(distances)
-    )
 
+def cutoff_cosine(distances: Tensor, cutoff: float) -> Tensor:
+    # assuming all elements in distances are smaller than cutoff
+    return 0.5 * torch.cos(distances * (math.pi / cutoff)) + 0.5
 
-# @torch.jit.script
-def radial_terms(Rcr, EtaR, ShfR, distances):
-    # type: (float, Tensor, Tensor, Tensor) -> Tensor
+
+def radial_terms(Rcr: float, EtaR: Tensor, ShfR: Tensor, distances: Tensor) -> Tensor:
     """Compute the radial subAEV terms of the center atom given neighbors
 
     This correspond to equation (3) in the `ANI paper`_. This function just
@@ -39,14 +36,13 @@ def radial_terms(Rcr, EtaR, ShfR, distances):
     ret = 0.25 * torch.exp(-EtaR * (distances - 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
+    # We then should flat the last 2 dimensions to view the subAEV as one
     # dimension vector
     return ret.flatten(start_dim=-2)
 
 
-# @torch.jit.script
-def angular_terms(Rca, ShfZ, EtaA, Zeta, ShfA, vectors1, vectors2):
-    # type: (float, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor) -> Tensor
+def angular_terms(Rca: float, ShfZ: Tensor, EtaA: Tensor, Zeta: Tensor,
+                  ShfA: Tensor, vectors12: Tensor) -> Tensor:
     """Compute the angular subAEV terms of the center atom given neighbor pairs.
 
     This correspond to equation (4) in the `ANI paper`_. This function just
@@ -59,21 +55,18 @@ def angular_terms(Rca, ShfZ, EtaA, Zeta, ShfA, vectors1, vectors2):
     .. _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)
+    vectors12 = vectors12.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
+    distances12 = vectors12.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)
+    cos_angles = 0.95 * torch.nn.functional.cosine_similarity(vectors12[0], vectors12[1], dim=-5)
     angles = torch.acos(cos_angles)
 
-    fcj1 = cutoff_cosine(distances1, Rca)
-    fcj2 = cutoff_cosine(distances2, Rca)
+    fcj12 = cutoff_cosine(distances12, Rca)
     factor1 = ((1 + torch.cos(angles - ShfZ)) / 2) ** Zeta
-    factor2 = torch.exp(-EtaA * ((distances1 + distances2) / 2 - ShfA) ** 2)
-    ret = 2 * factor1 * factor2 * fcj1 * fcj2
+    factor2 = torch.exp(-EtaA * (distances12.sum(0) / 2 - ShfA) ** 2)
+    ret = 2 * factor1 * factor2 * fcj12.prod(0)
     # 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
@@ -81,8 +74,7 @@ def angular_terms(Rca, ShfZ, EtaA, Zeta, ShfA, vectors1, vectors2):
     return ret.flatten(start_dim=-4)
 
 
-# @torch.jit.script
-def compute_shifts(cell, pbc, cutoff):
+def compute_shifts(cell: Tensor, pbc: Tensor, cutoff: float) -> Tensor:
     """Compute the shifts of unit cell along the given cell vectors to make it
     large enough to contain all pairs of neighbor atoms with PBC under
     consideration
@@ -99,11 +91,10 @@ def compute_shifts(cell, pbc, cutoff):
         :class:`torch.Tensor`: long tensor of shifts. the center cell and
             symmetric cells are not included.
     """
-    # type: (Tensor, Tensor, float) -> Tensor
     reciprocal_cell = cell.inverse().t()
     inv_distances = reciprocal_cell.norm(2, -1)
     num_repeats = torch.ceil(cutoff * inv_distances).to(torch.long)
-    num_repeats = torch.where(pbc, num_repeats, torch.zeros_like(num_repeats))
+    num_repeats = torch.where(pbc, num_repeats, num_repeats.new_zeros(()))
     r1 = torch.arange(1, num_repeats[0] + 1, device=cell.device)
     r2 = torch.arange(1, num_repeats[1] + 1, device=cell.device)
     r3 = torch.arange(1, num_repeats[2] + 1, device=cell.device)
@@ -125,8 +116,8 @@ def compute_shifts(cell, pbc, cutoff):
     ])
 
 
-# @torch.jit.script
-def neighbor_pairs(padding_mask, coordinates, cell, shifts, cutoff):
+def neighbor_pairs(padding_mask: Tensor, coordinates: Tensor, cell: Tensor,
+                   shifts: Tensor, cutoff: float) -> Tuple[Tensor, Tensor]:
     """Compute pairs of atoms that are neighbors
 
     Arguments:
@@ -139,88 +130,91 @@ def neighbor_pairs(padding_mask, coordinates, cell, shifts, cutoff):
         cutoff (float): the cutoff inside which atoms are considered pairs
         shifts (:class:`torch.Tensor`): tensor of shape (?, 3) storing shifts
     """
-    # type: (Tensor, Tensor, Tensor, Tensor, float) -> Tuple[Tensor, Tensor, Tensor, Tensor]
-
     coordinates = coordinates.detach()
     cell = cell.detach()
     num_atoms = padding_mask.shape[1]
+    num_mols = padding_mask.shape[0]
     all_atoms = torch.arange(num_atoms, device=cell.device)
 
     # Step 2: center cell
-    p1_center, p2_center = torch.combinations(all_atoms).unbind(-1)
-    shifts_center = shifts.new_zeros(p1_center.shape[0], 3)
+    # torch.triu_indices is faster than combinations
+    p12_center = torch.triu_indices(num_atoms, num_atoms, 1, device=cell.device)
+    shifts_center = shifts.new_zeros((p12_center.shape[1], 3))
 
     # Step 3: cells with shifts
     # shape convention (shift index, molecule index, atom index, 3)
     num_shifts = shifts.shape[0]
     all_shifts = torch.arange(num_shifts, device=cell.device)
-    shift_index, p1, p2 = torch.cartesian_prod(all_shifts, all_atoms, all_atoms).unbind(-1)
-    shifts_outide = shifts.index_select(0, shift_index)
+    prod = torch.cartesian_prod(all_shifts, all_atoms, all_atoms).t()
+    shift_index = prod[0]
+    p12 = prod[1:]
+    shifts_outside = shifts.index_select(0, shift_index)
 
     # Step 4: combine results for all cells
-    shifts_all = torch.cat([shifts_center, shifts_outide])
-    p1_all = torch.cat([p1_center, p1])
-    p2_all = torch.cat([p2_center, p2])
-    shift_values = torch.mm(shifts_all.to(cell.dtype), cell)
+    shifts_all = torch.cat([shifts_center, shifts_outside])
+    p12_all = torch.cat([p12_center, p12], dim=1)
+    shift_values = shifts_all.to(cell.dtype) @ cell
 
     # step 5, compute distances, and find all pairs within cutoff
-    distances = (coordinates.index_select(1, p1_all) - coordinates.index_select(1, p2_all) + shift_values).norm(2, -1)
-    padding_mask = (padding_mask.index_select(1, p1_all)) | (padding_mask.index_select(1, p2_all))
+    selected_coordinates = coordinates.index_select(1, p12_all.view(-1)).view(num_mols, 2, -1, 3)
+    distances = (selected_coordinates[:, 0, ...] - selected_coordinates[:, 1, ...] + shift_values).norm(2, -1)
+    padding_mask = padding_mask.index_select(1, p12_all.view(-1)).view(2, -1).any(0)
     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]
+    atom_index12 = p12_all[:, pair_index]
     shifts = shifts_all.index_select(0, pair_index)
-    return molecule_index + atom_index1, molecule_index + atom_index2, shifts
+    return molecule_index + atom_index12, shifts
+
+
+def neighbor_pairs_nopbc(padding_mask: Tensor, coordinates: Tensor, cutoff: float) -> Tensor:
+    """Compute pairs of atoms that are neighbors (doesn't use PBC)
 
+    This function bypasses the calculation of shifts and duplication
+    of atoms in order to make calculations faster
+
+    Arguments:
+        padding_mask (:class:`torch.Tensor`): boolean tensor of shape
+            (molecules, atoms) for padding mask. 1 == is padding.
+        coordinates (:class:`torch.Tensor`): tensor of shape
+            (molecules * atoms, 3) for atom coordinates.
+        cutoff (float): the cutoff inside which atoms are considered pairs
+    """
+    coordinates = coordinates.detach()
+    current_device = coordinates.device
+    num_atoms = padding_mask.shape[1]
+    num_mols = padding_mask.shape[0]
+    p12_all = torch.triu_indices(num_atoms, num_atoms, 1, device=current_device)
+    p12_all_flattened = p12_all.view(-1)
 
-# torch.jit.script
-def triu_index(num_species):
-    species = torch.arange(num_species)
-    species1, species2 = torch.combinations(species, r=2, with_replacement=True).unbind(-1)
-    pair_index = torch.arange(species1.shape[0])
+    pair_coordinates = coordinates.index_select(1, p12_all_flattened).view(num_mols, 2, -1, 3)
+    distances = (pair_coordinates[:, 0, ...] - pair_coordinates[:, 1, ...]).norm(2, -1)
+    padding_mask = padding_mask.index_select(1, p12_all_flattened).view(num_mols, 2, -1).any(dim=1)
+    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_index12 = p12_all[:, pair_index] + molecule_index
+    return atom_index12
+
+
+def triu_index(num_species: int) -> Tensor:
+    species1, species2 = torch.triu_indices(num_species, num_species).unbind(0)
+    pair_index = torch.arange(species1.shape[0], dtype=torch.long)
     ret = torch.zeros(num_species, num_species, dtype=torch.long)
     ret[species1, species2] = pair_index
     ret[species2, species1] = pair_index
     return ret
 
 
-# torch.jit.script
-def convert_pair_index(index):
-    """Let's say we have a pair:
-    index: 0 1 2 3 4 5 6 7 8 9 ...
-    elem1: 0 0 1 0 1 2 0 1 2 3 ...
-    elem2: 1 2 2 3 3 3 4 4 4 4 ...
-    This function convert index back to elem1 and elem2
-
-    To implement this, divide it into groups, the first group contains 1
-    elements, the second contains 2 elements, ..., the nth group contains
-    n elements.
-
-    Let's say we want to compute the elem1 and elem2 for index i. We first find
-    the number of complete groups contained in index 0, 1, ..., i - 1
-    (all inclusive, not including i), then i will be in the next group. Let's
-    say there are N complete groups, then these N groups contains
-    N * (N + 1) / 2 elements, solving for the largest N that satisfies
-    N * (N + 1) / 2 <= i, will get the N we want.
-    """
-    n = (torch.sqrt(1.0 + 8.0 * index.to(torch.float)) - 1.0) / 2.0
-    n = torch.floor(n).to(torch.long)
-    num_elems = n * (n + 1) / 2
-    return index - num_elems, n + 1
-
-
-# torch.jit.script
-def cumsum_from_zero(input_):
-    cumsum = torch.cumsum(input_, dim=0)
-    cumsum = torch.cat([input_.new_tensor([0]), cumsum[:-1]])
+def cumsum_from_zero(input_: Tensor) -> Tensor:
+    cumsum = torch.zeros_like(input_)
+    torch.cumsum(input_[:-1], dim=0, out=cumsum[1:])
     return cumsum
 
 
-# torch.jit.script
-def triple_by_molecule(atom_index1, atom_index2):
+def triple_by_molecule(atom_index12: Tensor) -> Tuple[Tensor, Tensor, Tensor]:
     """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.
@@ -232,74 +226,87 @@ def triple_by_molecule(atom_index1, atom_index2):
     are (1, 2), (1, 3), (1, 4), (2, 3), (2, 4), (3, 4)
     """
     # convert representation from pair to central-others
-    n = atom_index1.shape[0]
-    ai1 = torch.cat([atom_index1, atom_index2])
+    ai1 = atom_index12.view(-1)
     sorted_ai1, rev_indices = ai1.sort()
 
     # sort and compute unique key
-    uniqued_central_atom_index, counts = torch.unique_consecutive(sorted_ai1, return_counts=True)
+    unique_results = torch.unique_consecutive(sorted_ai1, return_inverse=True, return_counts=True)
+    uniqued_central_atom_index = unique_results[0]
+    counts = unique_results[-1]
 
-    # do local combinations within unique key, assuming sorted
+    # compute central_atom_index
     pair_sizes = counts * (counts - 1) // 2
-    total_size = pair_sizes.sum()
     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 = 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 = cumsum.index_select(0, pair_indices)
-    sorted_local_index1 += cumsum
-    sorted_local_index2 += cumsum
+
+    # do local combinations within unique key, assuming sorted
+    m = counts.max().item() if counts.numel() > 0 else 0
+    n = pair_sizes.shape[0]
+    intra_pair_indices = torch.tril_indices(m, m, -1, device=ai1.device).unsqueeze(1).expand(-1, n, -1)
+    mask = (torch.arange(intra_pair_indices.shape[2], device=ai1.device) < pair_sizes.unsqueeze(1)).flatten()
+    sorted_local_index12 = intra_pair_indices.flatten(1, 2)[:, mask]
+    sorted_local_index12 += cumsum_from_zero(counts).index_select(0, pair_indices)
 
     # unsort result from last part
-    local_index1 = rev_indices[sorted_local_index1]
-    local_index2 = rev_indices[sorted_local_index2]
+    local_index12 = rev_indices[sorted_local_index12]
 
     # compute mapping between representation of central-other to pair
-    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
+    n = atom_index12.shape[1]
+    sign12 = ((local_index12 < n).to(torch.int8) * 2) - 1
+    return central_atom_index, local_index12 % n, sign12
 
 
-# torch.jit.script
-def compute_aev(species, coordinates, cell, shifts, triu_index, constants, sizes):
+def compute_aev(species: Tensor, coordinates: Tensor, triu_index: Tensor,
+                constants: Tuple[float, Tensor, Tensor, float, Tensor, Tensor, Tensor, Tensor],
+                sizes: Tuple[int, int, int, int, int], cell_shifts: Optional[Tuple[Tensor, Tensor]]) -> Tensor:
     Rcr, EtaR, ShfR, Rca, ShfZ, EtaA, Zeta, ShfA = constants
-    num_species, radial_sublength, radial_length, angular_sublength, angular_length, aev_length = sizes
+    num_species, radial_sublength, radial_length, angular_sublength, angular_length = sizes
     num_molecules = species.shape[0]
     num_atoms = species.shape[1]
     num_species_pairs = angular_length // angular_sublength
-    cutoff = max(Rcr, Rca)
+    coordinates_ = coordinates
+    coordinates = coordinates_.flatten(0, 1)
+
+    # PBC calculation is bypassed if there are no shifts
+    if cell_shifts is None:
+        atom_index12 = neighbor_pairs_nopbc(species == -1, coordinates_, Rcr)
+        selected_coordinates = coordinates.index_select(0, atom_index12.view(-1)).view(2, -1, 3)
+        vec = selected_coordinates[0] - selected_coordinates[1]
+    else:
+        cell, shifts = cell_shifts
+        atom_index12, shifts = neighbor_pairs(species == -1, coordinates_, cell, shifts, Rcr)
+        shift_values = shifts.to(cell.dtype) @ cell
+        selected_coordinates = coordinates.index_select(0, atom_index12.view(-1)).view(2, -1, 3)
+        vec = selected_coordinates[0] - selected_coordinates[1] + shift_values
 
-    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)
+    species12 = species[atom_index12]
 
-    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 = 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_terms_.new_zeros((num_molecules * num_atoms * num_species, radial_sublength))
+    index12 = atom_index12 * num_species + species12.flip(0)
+    radial_aev.index_add_(0, index12[0], radial_terms_)
+    radial_aev.index_add_(0, index12[1], radial_terms_)
     radial_aev = radial_aev.reshape(num_molecules, num_atoms, radial_length)
 
+    # Rca is usually much smaller than Rcr, using neighbor list with cutoff=Rcr is a waste of resources
+    # Now we will get a smaller neighbor list that only cares about atoms with distances <= Rca
+    even_closer_indices = (distances <= Rca).nonzero().flatten()
+    atom_index12 = atom_index12.index_select(1, even_closer_indices)
+    species12 = species12.index_select(1, even_closer_indices)
+    vec = vec.index_select(0, even_closer_indices)
+
     # compute angular aev
-    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 = central_atom_index * num_species_pairs + triu_index[species1_, species2_]
+    central_atom_index, pair_index12, sign12 = triple_by_molecule(atom_index12)
+    species12_small = species12[:, pair_index12]
+    vec12 = vec.index_select(0, pair_index12.view(-1)).view(2, -1, 3) * sign12.unsqueeze(-1)
+    species12_ = torch.where(sign12 == 1, species12_small[1], species12_small[0])
+    angular_terms_ = angular_terms(Rca, ShfZ, EtaA, Zeta, ShfA, vec12)
+    angular_aev = angular_terms_.new_zeros((num_molecules * num_atoms * num_species_pairs, angular_sublength))
+    index = central_atom_index * num_species_pairs + triu_index[species12_[0], species12_[1]]
     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)
@@ -330,14 +337,24 @@ class AEVComputer(torch.nn.Module):
     .. _ANI paper:
         http://pubs.rsc.org/en/Content/ArticleLanding/2017/SC/C6SC05720A#!divAbstract
     """
-    __constants__ = ['Rcr', 'Rca', 'num_species', 'radial_sublength',
-                     'radial_length', 'angular_sublength', 'angular_length',
-                     'aev_length']
+    Rcr: Final[float]
+    Rca: Final[float]
+    num_species: Final[int]
+
+    radial_sublength: Final[int]
+    radial_length: Final[int]
+    angular_sublength: Final[int]
+    angular_length: Final[int]
+    aev_length: Final[int]
+    sizes: Final[Tuple[int, int, int, int, int]]
 
     def __init__(self, Rcr, Rca, EtaR, ShfR, EtaA, Zeta, ShfA, ShfZ, num_species):
         super(AEVComputer, self).__init__()
         self.Rcr = Rcr
         self.Rca = Rca
+        assert Rca <= Rcr, "Current implementation of AEVComputer assumes Rca <= Rcr"
+        self.num_species = num_species
+
         # convert constant tensors to a ready-to-broadcast shape
         # shape convension (..., EtaR, ShfR)
         self.register_buffer('EtaR', EtaR.view(-1, 1))
@@ -348,7 +365,6 @@ class AEVComputer(torch.nn.Module):
         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
         # The length of radial subaev of a single species
         self.radial_sublength = self.EtaR.numel() * self.ShfR.numel()
         # The length of full radial aev
@@ -359,15 +375,15 @@ class AEVComputer(torch.nn.Module):
         self.angular_length = (self.num_species * (self.num_species + 1)) // 2 * self.angular_sublength
         # The length of full aev
         self.aev_length = self.radial_length + self.angular_length
-        self.sizes = self.num_species, self.radial_sublength, self.radial_length, self.angular_sublength, self.angular_length, self.aev_length
+        self.sizes = self.num_species, self.radial_sublength, self.radial_length, self.angular_sublength, self.angular_length
 
-        self.register_buffer('triu_index', triu_index(num_species))
+        self.register_buffer('triu_index', triu_index(num_species).to(device=self.EtaR.device))
 
         # 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_pbc = torch.zeros(3, dtype=torch.bool, 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)
@@ -375,23 +391,33 @@ class AEVComputer(torch.nn.Module):
     def constants(self):
         return self.Rcr, self.EtaR, self.ShfR, self.Rca, self.ShfZ, self.EtaA, self.Zeta, self.ShfA
 
-    # @torch.jit.script_method
-    def forward(self, input):
+    def forward(self, input_: Tuple[Tensor, Tensor],
+                cell: Optional[Tensor] = None,
+                pbc: Optional[Tensor] = None) -> SpeciesAEV:
         """Compute AEVs
 
         Arguments:
-            input (tuple): Can be one of the following two cases:
+            input_ (tuple): Can be one of the following two cases:
 
                 If you don't care about periodic boundary conditions at all,
-                then input can be a tuple of 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 molecules
-                in a chunk, and ``A`` is the number of atoms.
+                then input can be a tuple of two tensors: species, coordinates.
+                species must have shape ``(N, A)``, coordinates must have shape
+                ``(N, A, 3)`` where ``N`` is the number of molecules in a batch,
+                and ``A`` is the number of atoms.
+
+                .. warning::
+
+                    The species must be indexed in 0, 1, 2, 3, ..., not the element
+                    index in periodic table. Check :class:`torchani.SpeciesConverter`
+                    if you want periodic table indexing.
+
+                .. note:: The coordinates, and cell are in Angstrom.
 
                 If you want to apply periodic boundary conditions, then the input
-                would be a tuple of four tensors: species, coordinates, cell, pbc
-                where species and coordinates are the same as described above, cell
-                is a tensor of shape (3, 3) of the three vectors defining unit cell:
+                would be a tuple of two tensors (species, coordinates) and two keyword
+                arguments `cell=...` , and `pbc=...` where species and coordinates are
+                the same as described above, cell is a tensor of shape (3, 3) of the
+                three vectors defining unit cell:
 
                 .. code-block:: python
 
@@ -403,17 +429,17 @@ class AEVComputer(torch.nn.Module):
                 for that direction.
 
         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())``
+            NamedTuple: Species and AEVs. species are the species from the input
+            unchanged, and AEVs is a tensor of shape ``(N, A, self.aev_length())``
         """
-        if len(input) == 2:
-            species, coordinates = input
-            cell = self.default_cell
-            shifts = self.default_shifts
+        species, coordinates = input_
+
+        if cell is None and pbc is None:
+            aev = compute_aev(species, coordinates, self.triu_index, self.constants(), self.sizes, None)
         else:
-            assert len(input) == 4
-            species, coordinates, cell, pbc = input
+            assert (cell is not None and pbc is not None)
             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)
+            aev = compute_aev(species, coordinates, self.triu_index, self.constants(), self.sizes, (cell, shifts))
+
+        return SpeciesAEV(species, aev)