Skip to content

Commit

Permalink
enh: find xyz atom neighbors
Browse files Browse the repository at this point in the history
  • Loading branch information
@pfebrer
pfebrer committed Feb 22, 2022
1 parent d1b681a commit 1620824
Show file tree
Hide file tree
Showing 2 changed files with 209 additions and 7 deletions.
150 changes: 148 additions & 2 deletions sisl/geom/neighs/fneighs.f90
View file
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,152 @@ subroutine build_table(nbins, bin_indices, Nat, list, heads, counts)
end do

end subroutine build_table

subroutine get_close(search_xyz, indices, iscs, heads, list, max_npairs, &
xyz, cell, pbc, thresholds, N_ind, N, neighs, split_indices)
! Gets the atoms that are close to given positions
!
! INPUTS
! ---------
! search_xyz:
! The coordinates for which we want to look for atoms that
! are close.
! indices:
! For each position index (first dimension), the indices of the
! 8 bins that contain potential neighbours. Bin indices are
! 0-indexed (python indices).
! iscs:
! Supercell indices of each bin index in `indices`. Supercell
! indices are 0-indexed (python indices).
! heads:
! For each bin, the index of `list` where we can find the
! first atom that is contained in it.
! This array is constructed by `build_table`.
! list:
! contains the list of atoms, modified to encode all bin
! locations. Each item in the list contains the index of
! the next atom that we can find in the same bin.
! Indices are 0-indexed (python indices). If an item is
! -1, that there are no more atoms in the same bin.
! This array is constructed by `build_table`.
! max_npairs:
! The number of maximum pairs that can be found.
! It is used to allocate the `neighs` array. This is computed
! in python with the help of the `counts` array constructed
! by `build_table`.
! xyz:
! the cartesian coordinates of all atoms in the geometry.
! cell:
! the lattice vectors of the geometry, where cell(i, :) is the
! ith lattice vector.
! pbc:
! for each direction, whether periodic boundary conditions should
! be considered.
! thresholds:
! the threshold radius for each atom in the geometry.
! N_ind:
! The number of positions for which we will look for neighbours.
! No need to pass it when calling from python, it will be
! inferred from the shape of `search_xyz` and `indices`.
! N:
! total number of atoms. No need to pass it when calling from
! python. It will be inferred from the shapes of the arrays.
!
!
! OUTPUTS
! --------
! neighs:
! contains the position indices and atom indices that are
! "neighbours". All indices are 0-indexed (python indices).
! split_indices:
! array containing the breakpoints in `neighs`. These
! breakpoints delimit the position where one can find pairs
! for each 8 bin search. It can be used later to quickly
! split `neighs` on the multiple individual searches.
implicit none

real(8), intent(in) :: search_xyz(N_ind, 3)
integer, intent(in) :: indices(N_ind, 8), iscs(N_ind, 8, 3)
integer, intent(in) :: heads(:), list(N)
integer, intent(in) :: max_npairs
real(8), intent(in) :: xyz(N, 3), cell(3,3), thresholds(N)
logical, intent(in) :: pbc(3)
integer, intent(in) :: N, N_ind
integer, intent(out) :: neighs(max_npairs, 5)
integer, intent(out) :: split_indices(N_ind)

! Auxiliar variables
integer :: search_index
integer :: bin_index, neigh_at, i_pair, j
real(8) :: dist, threshold, ref_xyz(3)
integer :: isc(3)
logical :: is_neigh_cell

! Initialize the pair index, which will tell us the location
! to store each pair that we find. Each time we store a new pair
! this is incremented by 1.
i_pair = 1

! Loop through all searches that we need to perform
do search_index=1, N_ind
! Loop through the 8 bins we have to explore for this search
do j=1, 8
! Find the bin index. Sum 1 to get the fortran index.
bin_index = indices(search_index, j) + 1

! Get the first atom index in this bin
neigh_at = heads(bin_index)

! If there are no atoms in this bin, do not even bother
! checking supercell indices.
if (neigh_at == -1) cycle

! Find the supercell indices for this bin
isc(:) = iscs(search_index, j, :)
! And check if this bin corresponds to a neighbor cell.
is_neigh_cell = count(isc /= 0) > 0

ref_xyz = search_xyz(search_index, :)
if (is_neigh_cell) then
! If we are looking at a neighbouring cell in a direction
! where there are no periodic boundary conditions, go to
! next bin.
if (any(isc /= 0 .and. .not. pbc)) cycle
! Otherwise, move the position to the neighbor cell. We do this
! instead of moving potential neighbors to the unit cell
! because in this way we reduce the number of operations.
ref_xyz = ref_xyz - matmul(isc, cell)
end if

! Loop through all atoms that are in this bin.
! If neigh_at == -1, this means no more atoms are in this bin.
do while (neigh_at >= 0)
! Calculate the distance between the atom and the potential
! neighbor.
dist = norm2(ref_xyz - xyz(neigh_at + 1, :))

! Get the threshold for this pair of atoms
threshold = thresholds(neigh_at + 1)

if (dist < threshold) then
! Store the pair of neighbours.
neighs(i_pair, 1) = search_index
neighs(i_pair, 2) = neigh_at
neighs(i_pair, 3:) = isc
! Increment the pair index
i_pair = i_pair + 1
end if

! Get the next atom in this bin. Sum 1 to get fortran index.
neigh_at = list(neigh_at + 1)
end do
end do

! We have finished this search, store the breakpoint.
split_indices(search_index) = i_pair - 1
end do

end subroutine get_close

subroutine get_pairs(at_indices, indices, iscs, heads, list, max_npairs, &
self_interaction, xyz, cell, pbc, thresholds, sphere_overlap, N_ind, N, &
Expand Down Expand Up @@ -119,8 +265,8 @@ subroutine get_pairs(at_indices, indices, iscs, heads, list, max_npairs, &
! OUTPUTS
! --------
! neighs:
! contains the atom indices of all pairs that are potential
! neighbours. Atom indices are 0-indexed (python indices).
! contains the atom indices of all pairs of neighbor atoms.
! Atom indices are 0-indexed (python indices).
! split_indices:
! array containing the breakpoints in `neighs`. These
! breakpoints delimit the position where one can find pairs
Expand Down
66 changes: 61 additions & 5 deletions sisl/geom/neighs/neighfinder.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -292,7 +292,7 @@ def find_neighbours(self, atoms=None, as_pairs=False, self_interaction=False, pb
Each pair `ij` means that `j` is a neighbour of `i`.
The three extra columns are the supercell indices of atom `j`.
If `as_pairs=False`:
A list containing a numpy array of shape (n_neighs) for each atom.
A list containing a numpy array of shape (n_neighs, 4) for each atom.
"""
# Sanitize atoms
atoms = self.geometry._sanitize_atoms(atoms)
Expand All @@ -307,13 +307,13 @@ def find_neighbours(self, atoms=None, as_pairs=False, self_interaction=False, pb
# Get atom counts
at_counts = self._get_search_atom_counts(search_indices)

total_pairs = at_counts.sum()
max_pairs = at_counts.sum()
if not self_interaction:
total_pairs -= search_indices.shape[0]
max_pairs -= search_indices.shape[0]

# Find the candidate pairs
# Find the neighbour pairs
neighbour_pairs, split_ind = _fneighs.get_pairs(
atoms, search_indices, isc, self._heads, self._list, total_pairs, self_interaction,
atoms, search_indices, isc, self._heads, self._list, max_pairs, self_interaction,
self.geometry.xyz, self.geometry.cell, pbc, thresholds, self._sphere_overlap
)

Expand Down Expand Up @@ -385,6 +385,62 @@ def find_all_unique_pairs(self, self_interaction=False, pbc=(True, True, True)):

return candidate_pairs[:n_pairs]

def find_close(self, xyz, as_pairs=False, pbc=(True, True, True)):
"""Find all atoms that are close to some coordinates in space.
This routine only executes the action of finding neighbours,
the parameters of the search (`geometry`, `R`, `sphere_overlap`...)
are defined when the finder is initialized or by calling `setup_finder`.
Parameters
----------
xyz: array-like of shape ([npoints], 3)
the coordinates for which neighbours are desired.
as_pairs: bool, optional
If `True` pairs are returned, where the first item is the index
of the point and the second one is the atom index.
Otherwise a list containing the neighbours for each point is returned.
pbc: bool or array-like of shape (3, )
whether periodic conditions should be considered.
If a single bool is passed, all directions use that value.
Returns
----------
np.ndarray or list:
If `as_pairs=True`:
A `np.ndarray` of shape (n_pairs, 5) is returned.
Each pair `ij` means that `j` is a neighbour of `i`.
The three extra columns are the supercell indices of atom `j`.
If `as_pairs=False`:
A list containing a numpy array of shape (n_neighs, 4) for each atom.
"""
# Cast R and pbc into arrays of appropiate shape and type.
thresholds = np.full(self.geometry.na, self._R, dtype=np.float64)
pbc = np.full(3, pbc, dtype=np.bool)

xyz = np.atleast_2d(xyz)
# Get search indices
search_indices, isc = self._get_search_indices(xyz.dot(self.geometry.icell.T) % 1, cartesian=False)

# Get atom counts
at_counts = self._get_search_atom_counts(search_indices)

max_pairs = at_counts.sum()

# Find the neighbour pairs
neighbour_pairs, split_ind = _fneighs.get_close(
xyz, search_indices, isc, self._heads, self._list, max_pairs,
self.geometry.xyz, self.geometry.cell, pbc, thresholds
)

if as_pairs:
# Just return the neighbour pairs
return neighbour_pairs[:split_ind[-1]]
else:
# Split to get the neighbours of each position
return np.split(neighbour_pairs[:, 1:], split_ind, axis=0)[:-1]
return

def find_neighbours_old(self, atoms=None, as_pairs=False, self_interaction=False):
"""Find neighbours as specified in the finder.
Expand Down

0 comments on commit 1620824

Please sign in to comment.