sponge.partition.DistanceNeighbours

class sponge.partition.DistanceNeighbours(cutoff: float, num_neighbours: int = None, atom_mask: Tensor = None, exclude_index: Tensor = None, use_pbc: bool = None, cutoff_scale: float = 1.2, large_dis: float = 10000.0, cast_fp16: bool = False)

Neighbour list calculated by distance

Parameters

• cutoff (float) – Cutoff distance.

• num_neighbours (int, optional) – 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.

• atom_mask (Tensor, optional) – Tensor of shape \((B, A)\). Data type is bool_. Mask of atoms in the system. Default: None.

• exclude_index (Tensor, optional) – Tensor of shape \((B, A, Ex)\). Data type is int32. Index of neighbour atoms which could be excluded from the neighbour list. Default: None.

• use_pbc (bool, optional) – Whether to use periodic boundary condition. Default: None.

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

• large_dis (float, optional) – A large number to fill in the distances to the masked neighbouring atoms. Default: 1e4

• 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: Number of simulation walker.

• A: Number of atoms in system.

• N: Number of neighbour atoms.

• Ex: Maximum number of excluded neighbour atoms.

Supported Platforms:

Ascend GPU

Examples

>>> import sponge
>>> from sponge.partition import DistanceNeighbours
>>> import mindspore
>>> from mindspore import Tensor
>>> distance_neighbours = DistanceNeighbours(0.5, use_pbc=False)
>>> coordinate = Tensor([[[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]])
>>> distance_neighbours(coordinate)
(Tensor(shape=[1, 2, 1], dtype=Float32, value=
 [[[ 1.00000000e+00],
 [ 1.00000000e+00]]]),
 Tensor(shape=[1, 2, 1], dtype=Int64, value=
 [[[0],
 [1]]]),
 Tensor(shape=[1, 2, 1], dtype=Bool, value=
 [[[False],
 [False]]]))
>>> distance = Tensor([[[1.0, 0.0, 0.0], [0.0, 0.0, 0.0]]])
>>> distance_neighbours.calc_max_neighbours(distance, 0.5)
Tensor(shape=[], dtype=Int32, value= 2)
>>> distance_neighbours.check_neighbour_list()
DistanceNeighbours<
  (get_distance): GetDistance<>
  >

static calc_max_neighbours(distances: Tensor, cutoff: float)

Calculate the maximum number of neighbouring atoms.

Parameters

• distances (Tensor) – Tensor of shape \((B, A, N)\). Data type is float.

• cutoff (float) – Cutoff distance.

check_neighbour_list()

Check the number of neighbouring atoms in neighbour list.

print_info()

Print the information of neighbour list.

set_exclude_index(exclude_index: Tensor)

Set the indices of atoms to be excluded from the neighbor list.

Parameters

exclude_index (Tensor) – Tensor of shape \((B, A, Ex)\). Data type is int. Index of the atoms that should be excluded from the neighbour list.

set_num_neighbours(coordinate: Tensor, pbc_box: Tensor = None, scale_factor: float = 1.25)

Set maximum number of neighbouring atoms.

Parameters

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

• pbc_box (Tensor, optional) – Tensor of shape \((B, D)\). Data type is bool. Periodic boundary condition box. Default: None.

• scale_factor (float, optional) – Factor to scale the cutoff distance. Default: 1.25.