torchani.neighbors

Modular neighborlists to improve scaling for large systems

Functions

adaptive_list
all_pairs
atom_image_converters
cell_list
compute_bounding_cell	Compute the rectangular unit cell that minimally bounds a set of coords.
coords_to_fractional
coords_to_grid_idx3
count_atoms_in_buckets
discard_inter_molecule_pairs	Discard neighbors that don't belong to the same molecule
discard_outside_cutoff	Discard neighbors with distances that lie outside of the given cutoff
flatten_idx3
image_pairs_within
lower_image_pairs_between
narrow_down	Takes a set of potential neighbor idxs and narrows it down to true neighbors
neighbors_to_triples	Converts output of a neighborlist calculation into triples of atoms
reconstruct_shifts	Reconstruct shift values used to calculate neighbors
setup_grid

Classes

AdaptiveList	Compute pairs of neighbors using the best algorithm for the system size
AllPairs	Compute pairs of neighbors.
CellList	Compute pairs of neighbors using the 'Cell List' algorithm.
FastCellList	This class is experimental and requires the compiled Cell-List extension
Neighborlist	Base class for modules that compute pairs of neighbors.
Neighbors	Holds pairs of atoms that are neighbors.
Triples	Holds groups of 3 atoms that are neighbors.
VerletCellList	Compute pairs of neighbors.

class torchani.neighbors.Neighbors(indices, distances, diff_vectors)

Holds pairs of atoms that are neighbors. Result of a neighborlist calculation

indices: Tensor

Long tensor with idxs of neighbor pairs. Shape (2, pairs)

distances: Tensor

The associated pair distances. Shape is (pairs,)

diff_vectors: Tensor

The associated difference vectors. Shape is (pairs, 3)

class torchani.neighbors.Triples(central_idxs, side_idxs, diff_signs, distances, diff_vectors)

Holds groups of 3 atoms that are neighbors. Result of neighbors_to_triples

central_idxs: Tensor

Long tensor with central idxs. Shape (triples,)

side_idxs: Tensor

Long tensor with side idxs. Shape (2, triples)

diff_signs: Tensor

Long tensor of diff vector directions. Shape (2, triples)

distances: Tensor

Similar to Neighbors, for triples. (2, triples)

diff_vectors: Tensor

Similar to Neighbors, for triples. (2, triples, 3)

torchani.neighbors.discard_inter_molecule_pairs(neighbors, molecule_idxs)

Discard neighbors that don’t belong to the same molecule

torchani.neighbors.discard_outside_cutoff(neighbors, cutoff)

Discard neighbors with distances that lie outside of the given cutoff

torchani.neighbors.narrow_down(cutoff, elem_idxs, coords, neighbor_idxs, shifts=None)

Takes a set of potential neighbor idxs and narrows it down to true neighbors

torchani.neighbors.compute_bounding_cell(coords, eps=0.001, displace=True, square=False)

Compute the rectangular unit cell that minimally bounds a set of coords.

Optionally displace the coords so that they fully lie inside the cell. Displacing the coordinates causes small floating point differences which have a negligible effect on energies and forces. The eps value is used to add padding to the cell.

class torchani.neighbors.Neighborlist(*args, **kwargs)

Base class for modules that compute pairs of neighbors. Can support PBC.

Subclasses must override Neighborlist.forward

forward(cutoff, species, coords, cell=None, pbc=None)

Calculate all pairs of atoms that are neighbors, given a cutoff.

Parameters:

• cutoff (float) – Cutoff value for the neighborlist. Pairs further away than this are not included.

• species (Tensor) – An int torch.Tensor that stores the element indices of a batch of molecules, (for example after conversion with torchani.nn.SpeciesConverter). Shape is (molecules, 3).

• coords (Tensor) – A float tensor with the coordinates of a batch of molecules. Shape is (molecules, atoms, 3). All ANI models use Angstrom.

• cell (Tensor | None) – A float tensor with unit cell vectors in its rows. Only use this with PBC. A cell with dimensions 10, 15, 20 (in Angstrom for all ANI models) in the x, y, and z directions is given by torch.tensor([[10., 0., 0.],[0., 15., 0.],[0., 0., 20.]]).

• pbc (Tensor | None) – A bool tensor that stores whether periodic boundary conditions (PBC) are enabled for the x, y, z directions. pbc=torch.tensor([True, True, True]) fully enables PBC and pbc=None (or torch.tensor([False, False, False])). fully disables it.

Returns:

typing.NamedTuple with all pairs of atoms that are neighbors.

Return type:

Neighbors

class torchani.neighbors.AllPairs(*args, **kwargs)

Compute pairs of neighbors. Uses a naive algorithm.

This is a naive implementation, with \(O(N^2)\) scaling. It computes all pairs and then discards those that are further away from the cutoff. Supports PBC.

forward(cutoff, species, coords, cell=None, pbc=None)

Calculate all pairs of atoms that are neighbors, given a cutoff.

Parameters:

• cutoff (float) – Cutoff value for the neighborlist. Pairs further away than this are not included.

• species (Tensor) – An int torch.Tensor that stores the element indices of a batch of molecules, (for example after conversion with torchani.nn.SpeciesConverter). Shape is (molecules, 3).

• coords (Tensor) – A float tensor with the coordinates of a batch of molecules. Shape is (molecules, atoms, 3). All ANI models use Angstrom.

• cell (Tensor | None) – A float tensor with unit cell vectors in its rows. Only use this with PBC. A cell with dimensions 10, 15, 20 (in Angstrom for all ANI models) in the x, y, and z directions is given by torch.tensor([[10., 0., 0.],[0., 15., 0.],[0., 0., 20.]]).

• pbc (Tensor | None) – A bool tensor that stores whether periodic boundary conditions (PBC) are enabled for the x, y, z directions. pbc=torch.tensor([True, True, True]) fully enables PBC and pbc=None (or torch.tensor([False, False, False])). fully disables it.

Returns:

typing.NamedTuple with all pairs of atoms that are neighbors.

Return type:

Neighbors

class torchani.neighbors.FastCellList

This class is experimental and requires the compiled Cell-List extension

forward(cutoff, species, coords, cell=None, pbc=None)

Calculate all pairs of atoms that are neighbors, given a cutoff.

Parameters:

• cutoff (float) – Cutoff value for the neighborlist. Pairs further away than this are not included.

• species (Tensor) – An int torch.Tensor that stores the element indices of a batch of molecules, (for example after conversion with torchani.nn.SpeciesConverter). Shape is (molecules, 3).

• coords (Tensor) – A float tensor with the coordinates of a batch of molecules. Shape is (molecules, atoms, 3). All ANI models use Angstrom.

• cell (Tensor | None) – A float tensor with unit cell vectors in its rows. Only use this with PBC. A cell with dimensions 10, 15, 20 (in Angstrom for all ANI models) in the x, y, and z directions is given by torch.tensor([[10., 0., 0.],[0., 15., 0.],[0., 0., 20.]]).

• pbc (Tensor | None) – A bool tensor that stores whether periodic boundary conditions (PBC) are enabled for the x, y, z directions. pbc=torch.tensor([True, True, True]) fully enables PBC and pbc=None (or torch.tensor([False, False, False])). fully disables it.

Returns:

typing.NamedTuple with all pairs of atoms that are neighbors.

Return type:

Neighbors

class torchani.neighbors.CellList(*args, **kwargs)

Compute pairs of neighbors using the ‘Cell List’ algorithm.

This is a linearly scaling implementation that uses the ‘cell list’ algorithm. It subdivides space into cells and then computes all pairs within each cell and between each cell and neighboring cells. and then discards those that are further away from the cutoff.

forward(cutoff, species, coords, cell=None, pbc=None)

Calculate all pairs of atoms that are neighbors, given a cutoff.

Parameters:

• cutoff (float) – Cutoff value for the neighborlist. Pairs further away than this are not included.

• species (Tensor) – An int torch.Tensor that stores the element indices of a batch of molecules, (for example after conversion with torchani.nn.SpeciesConverter). Shape is (molecules, 3).

• coords (Tensor) – A float tensor with the coordinates of a batch of molecules. Shape is (molecules, atoms, 3). All ANI models use Angstrom.

• cell (Tensor | None) – A float tensor with unit cell vectors in its rows. Only use this with PBC. A cell with dimensions 10, 15, 20 (in Angstrom for all ANI models) in the x, y, and z directions is given by torch.tensor([[10., 0., 0.],[0., 15., 0.],[0., 0., 20.]]).

• pbc (Tensor | None) – A bool tensor that stores whether periodic boundary conditions (PBC) are enabled for the x, y, z directions. pbc=torch.tensor([True, True, True]) fully enables PBC and pbc=None (or torch.tensor([False, False, False])). fully disables it.

Returns:

typing.NamedTuple with all pairs of atoms that are neighbors.

Return type:

Neighbors

class torchani.neighbors.AdaptiveList(threshold=190, threshold_nopbc=1770)

Compute pairs of neighbors using the best algorithm for the system size

This is a linearly scaling implementation that uses the ‘cell list’ algorithm for large system sizes and the naive ‘all pairs’ for small sizes.

forward(cutoff, species, coords, cell=None, pbc=None)

Calculate all pairs of atoms that are neighbors, given a cutoff.

Parameters:

• cutoff (float) – Cutoff value for the neighborlist. Pairs further away than this are not included.

• species (Tensor) – An int torch.Tensor that stores the element indices of a batch of molecules, (for example after conversion with torchani.nn.SpeciesConverter). Shape is (molecules, 3).

• coords (Tensor) – A float tensor with the coordinates of a batch of molecules. Shape is (molecules, atoms, 3). All ANI models use Angstrom.

• cell (Tensor | None) – A float tensor with unit cell vectors in its rows. Only use this with PBC. A cell with dimensions 10, 15, 20 (in Angstrom for all ANI models) in the x, y, and z directions is given by torch.tensor([[10., 0., 0.],[0., 15., 0.],[0., 0., 20.]]).

• pbc (Tensor | None) – A bool tensor that stores whether periodic boundary conditions (PBC) are enabled for the x, y, z directions. pbc=torch.tensor([True, True, True]) fully enables PBC and pbc=None (or torch.tensor([False, False, False])). fully disables it.

Returns:

typing.NamedTuple with all pairs of atoms that are neighbors.

Return type:

Neighbors

class torchani.neighbors.VerletCellList(skin=1.0)

Compute pairs of neighbors. Uses a cell-list algorithm with ‘verlet’ skin.

forward(cutoff, species, coords, cell=None, pbc=None)

Calculate all pairs of atoms that are neighbors, given a cutoff.

Parameters:

• cutoff (float) – Cutoff value for the neighborlist. Pairs further away than this are not included.

• species (Tensor) – An int torch.Tensor that stores the element indices of a batch of molecules, (for example after conversion with torchani.nn.SpeciesConverter). Shape is (molecules, 3).

• coords (Tensor) – A float tensor with the coordinates of a batch of molecules. Shape is (molecules, atoms, 3). All ANI models use Angstrom.

• cell (Tensor | None) – A float tensor with unit cell vectors in its rows. Only use this with PBC. A cell with dimensions 10, 15, 20 (in Angstrom for all ANI models) in the x, y, and z directions is given by torch.tensor([[10., 0., 0.],[0., 15., 0.],[0., 0., 20.]]).

• pbc (Tensor | None) – A bool tensor that stores whether periodic boundary conditions (PBC) are enabled for the x, y, z directions. pbc=torch.tensor([True, True, True]) fully enables PBC and pbc=None (or torch.tensor([False, False, False])). fully disables it.

Returns:

typing.NamedTuple with all pairs of atoms that are neighbors.

Return type:

Neighbors

torchani.neighbors.neighbors_to_triples(neighbors)

Converts output of a neighborlist calculation into triples of atoms

torchani.neighbors.reconstruct_shifts(coords, neighbors)

Reconstruct shift values used to calculate neighbors