diff --git a/sisl/geom/neighs/fneighs.f90 b/sisl/geom/neighs/fneighs.f90 index ce5b5b06d..d3ec53fa7 100644 --- a/sisl/geom/neighs/fneighs.f90 +++ b/sisl/geom/neighs/fneighs.f90 @@ -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, & @@ -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 diff --git a/sisl/geom/neighs/neighfinder.py b/sisl/geom/neighs/neighfinder.py index d9c1096e0..4881ae717 100644 --- a/sisl/geom/neighs/neighfinder.py +++ b/sisl/geom/neighs/neighfinder.py @@ -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) @@ -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 ) @@ -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.