sponge.partition.NeighbourList

class sponge.partition.NeighbourList(system: Molecule, cutoff: float = None, pace: int = 20, exclude_index: Tensor = None, num_neighbours: int = None, num_cell_cut: int = 1, cutoff_scale: float = 1.2, cell_cap_scale: float = 1.25, grid_num_scale: float = 2, use_grids: bool = False, cast_fp16: bool = False)

Neighbour list

Parameters

• system (Molecule) – Simulation system.

• cutoff (float, optional) – Cut-off distance. If None is given under periodic boundary condition (PBC), the cutoff will be assigned with the default value of 1 nm. Default: None.

• pace (int, optional) – Update frequency for neighbour list. Default: 20

• exclude_index (Tensor, optional) – Tensor of the indices of the neighbouring atoms which could be excluded from the neighbour list. The shape of Tensor is \((B, A, Ex)\), and the data type is int. Default: None.

• num_neighbours (int, optional) – Maximum number of neighbours. If None is given, this value will be calculated by the ratio of the number of neighbouring grids to the total number of grids. Default: None.

• num_cell_cut (int, optional) – Number of subdivision of grid cells according to cutoff. Default: 1

• cutoff_scale (float, optional) – Factor to scale cutoff distance. Default: 1.2

• cell_cap_scale (float, optional) – Scale factor for cell_capacity. Default: 1.25

• grid_num_scale (float, optional) – Scale factor to calculate num_neighbours by ratio of grids. If num_neighbours is not None, it will not be used. Default: 2

• use_grids (bool, optional) – Whether to use grids to calculate the neighbour list. Default: None.

• cast_fp16 (bool, optional) – If this is set to True, the data will be cast to float16 before sort. For use with some devices that only support sorting of float16 data. Default: False.

Note

• B: Batchsize, i.e. number of walkers of the simulation.

• A: Number of the atoms in the simulation system.

• N: Number of the maximum neighbouring atoms.

• D: Dimension of position coordinates.

• Ex: Maximum number of excluded neighbour atoms.

Supported Platforms:

Ascend GPU

Examples

>>> import sponge
>>> from sponge.partition import NeighbourList
>>> from sponge.system import Molecule
>>> system = Molecule(template='water.spce.yaml')
>>> neighbourlist = NeighbourList(system, 0.5)
>>> neighbourlist(system.coordinate, system.pbc_box)
(Tensor(shape=[1, 3, 2], dtype=Int64, value=
 [[[1, 2],
 [0, 2],
 [0, 1]]]),
 Tensor(shape=[1, 3, 2, 3], dtype=Float32, value=
 [[[[ 8.16490427e-02,  5.77358976e-02,  0.00000000e+00],
     [-8.16490427e-02,  5.77358976e-02,  0.00000000e+00]],
 [[-8.16490427e-02, -5.77358976e-02,  0.00000000e+00],
     [-1.63298085e-01,  0.00000000e+00,  0.00000000e+00]],
 [[ 8.16490427e-02, -5.77358976e-02,  0.00000000e+00],
     [ 1.63298085e-01,  0.00000000e+00,  0.00000000e+00]]]]),
 Tensor(shape=[1, 3, 2], dtype=Float32, value=
 [[[ 1.00000001e-01,  1.00000001e-01],
 [ 1.00000001e-01,  1.63298085e-01],
 [ 1.00000001e-01,  1.63298085e-01]]]),
 Tensor(shape=[1, 3, 2], dtype=Bool, value=
 [[[ True,  True],
 [ True,  True],
 [ True,  True]]]))
>>> neighbourlist.calculate(system.coordinate, system.pbc_box)
(Tensor(shape=[1, 3, 2], dtype=Int64, value=
 [[[1, 2],
 [0, 2],
 [0, 1]]]),
 Tensor(shape=[1, 3, 2], dtype=Bool, value=
 [[[ True,  True],
 [ True,  True],
 [ True,  True]]]))
>>> neighbourlist.pace
20

calculate(coordinate: Tensor, pbc_box: Tensor = None)

calculate neighbour list.

Parameters

• coordinate (Tensor) – Tensor of shape \((B, A, D)\). Data type is float. Position coordinate.

• pbc_box (Tensor, optional) – Tensor of shape \((B, D)\). Data type is float. Size of PBC box.

Returns

neigh_idx (Tensor), Tensor of shape \((B, A, N)\). Data type is int. Index of neighbouring atoms of each atoms in system.

neigh_mask (Tensor), Tensor of shape \((B, A, N)\). Data type is bool. Mask for neighbour list neigh_idx.

Note

• B: Batchsize, i.e. number of walkers of the simulation.

• A: Number of the atoms in the simulation system.

• N: Number of the maximum neighbouring atoms.

• D: Dimension of position coordinates.

get_neighbour_list()

get neighbour list.

Returns

neigh_idx (Tensor), Tensor of shape \((B, A, N)\). Data type is int. Index of neighbouring atoms of each atoms in system.

neigh_mask (Tensor):, Tensor of shape \((B, A, N)\). Data type is bool. Mask for neighbour list neigh_idx.

Note

• B: Batchsize, i.e. number of walkers of the simulation.

• A: Number of the atoms in the simulation system.

• N: Number of the maximum neighbouring atoms.

property pace: int

Update frequency for neighbour list

Returns

int, update pace

print_info()

print information of neighbour list

set_exclude_index(exclude_index: Tensor)

set exclude index

Parameters

exclude_index (Tensor) – Tensor of shape \((B, A, Ex)\). Data type is int.

update(coordinate: Tensor, pbc_box: Tensor = None)

update neighbour list.

Parameters

• coordinate (Tensor) – Tensor of shape \((B, A, D)\). Data type is float. Position coordinate.

• pbc_box (Tensor, optional) – Tensor of shape \((B, D)\). Data type is float. Size of PBC box.

Returns

neigh_idx (Tensor), Tensor of shape \((B, A, N)\). Data type is int. Index of neighbouring atoms of each atoms in system.

neigh_mask (Tensor), Tensor of shape \((B, A, N)\). Data type is bool. Mask for neighbour list neigh_idx.

Note

• B: Batchsize, i.e. number of walkers of the simulation.

• A: Number of the atoms in the simulation system.

• N: Number of the maximum neighbouring atoms.

• D: Dimension of position coordinates.