Skip to content

Commit

Permalink
enh: moved threshold checking to fortran
Browse files Browse the repository at this point in the history
  • Loading branch information
@pfebrer
pfebrer committed Nov 8, 2021
1 parent d9c9e63 commit 723dc69
Show file tree
Hide file tree
Showing 2 changed files with 365 additions and 23 deletions.
267 changes: 248 additions & 19 deletions sisl/geom/neighs/fneighs.f90
View file
Original file line number Diff line number Diff line change
Expand Up @@ -58,8 +58,9 @@ subroutine build_table(nbins, bin_indices, Nat, list, heads, counts)

end subroutine build_table

subroutine get_pairs(at_indices, indices, heads, list, npairs, &
self_interaction, N_ind, neighs, split_indices)
subroutine get_pairs(at_indices, indices, heads, list, max_npairs, &
self_interaction, xyz, thresholds, sphere_overlap, N_ind, N, &
neighs, split_indices)
! Gets (possibly duplicated) pairs of atoms that can potentially
! be neighbours.
!
Expand Down Expand Up @@ -107,15 +108,18 @@ subroutine get_pairs(at_indices, indices, heads, list, npairs, &
! split `neighs` on the multiple individual searches.
integer, intent(in) :: at_indices(N_ind)
integer, intent(in) :: heads(:), list(:), indices(N_ind, 8)
integer, intent(in) :: npairs
integer, intent(in) :: max_npairs
logical, intent(in) :: self_interaction
integer, intent(in) :: N_ind
integer, intent(out) :: neighs(npairs, 2)
real(8), intent(in) :: xyz(N, 3), thresholds(N)
logical, intent(in) :: sphere_overlap
integer, intent(in) :: N, N_ind
integer, intent(out) :: neighs(max_npairs, 2)
integer, intent(out) :: split_indices(N_ind)

! Auxiliar variables
integer :: search_index
integer :: bin_index, at, neigh_at, i_pair, j
real(8) :: dist, threshold

! Initialize the pair index, which will tell us the location
! to store each pair that we find. Each time we store a new pair
Expand All @@ -135,12 +139,21 @@ subroutine get_pairs(at_indices, indices, heads, list, npairs, &
! 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)
if (.not. (.not. self_interaction .and. neigh_at == at)) then
! Store the pair of potential neighbours.
neighs(i_pair, 1) = at
neighs(i_pair, 2) = neigh_at
! Increment the pair index
i_pair = i_pair + 1
dist = norm2(xyz(at + 1, :) - xyz(neigh_at + 1, :))

threshold = thresholds(at + 1)
if (sphere_overlap) then
threshold = threshold + thresholds(neigh_at + 1)
endif

if (dist < threshold) then
if (.not. (.not. self_interaction .and. neigh_at == at)) then
! Store the pair of potential neighbours.
neighs(i_pair, 1) = at
neighs(i_pair, 2) = neigh_at
! Increment the pair index
i_pair = i_pair + 1
end if
end if

! Get the next atom in this bin. Sum 1 to get fortran index.
Expand All @@ -154,7 +167,126 @@ subroutine get_pairs(at_indices, indices, heads, list, npairs, &

end subroutine get_pairs

subroutine get_all_unique_pairs(indices, heads, list, max_npairs, &
subroutine get_all_unique_pairs(indices, heads, list, &
max_npairs, self_interaction, xyz, thresholds, sphere_overlap, &
N, neighs, npairs)
! Gets all pairs of atoms that are neighbours according to a given
! threshold.
!
! This routine removes duplicates.
!
! INPUTS
! ---------
! indices:
! For each atom index (first dimension), the indices of the
! 8 bins that contain potential neighbours. Bin indices are
! 0-indexed (python indices). The first dimension should be
! as large as the total number of atoms.
! 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`.
! self_interaction: bool, optional
! whether to consider an atom a neighbour of itself.
! N:
! total number of atoms. No need to pass it when calling from
! python. It will be inferred from the shapes of `indices` and
! `list`.
!
! OUTPUTS
! --------
! neighs:
! contains the atom indices of all unique pairs that are
! potential neighbours. First column of atoms is sorted
! in increasing order. Each pair is also sorted in
! increasing order.
! Atom indices are 0-indexed (python indices).
! npairs:
! The number of unique pairs that have been found, which
! will be less than `max_npairs`. This should presumably
! be used to slice the `neighs` array once in python.

implicit none
integer, parameter :: dp = selected_real_kind(p=15)

integer, intent(in) :: indices(N, 8), heads(:), list(N)
integer, intent(in) :: max_npairs
logical, intent(in) :: self_interaction
real(8), intent(in) :: xyz(N, 3), thresholds(N)
logical, intent(in) :: sphere_overlap
integer, intent(in) :: N
integer, intent(out) :: neighs(max_npairs, 2), npairs

! Auxiliar variables.
integer :: at, neigh_at
integer :: bin_index, i_pair, j
real(8) :: dist, threshold

! 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 atoms, for each atom we will perform a search.
do at = 0, N - 1
! 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(at + 1, j) + 1
! Get the first atom index in this bin.
neigh_at = heads(bin_index)

! Loop through all atoms that are in this bin.
! If neigh_at is smaller than at, we already stored
! this pair when performing the search for neigh_at.
! The following atoms will have even lower indices
! So we can just move to the next bin.
do while (neigh_at > at)
dist = norm2(xyz(at + 1, :) - xyz(neigh_at + 1, :))

threshold = thresholds(at + 1)
if (sphere_overlap) then
threshold = threshold + thresholds(neigh_at + 1)
endif

if (dist < threshold) then
! Store the pair of potential neighbours.
neighs(i_pair, 1) = at
neighs(i_pair, 2) = neigh_at
! 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

if (j==1 .and. self_interaction) then
! Add the self interaction
neighs(i_pair, :) = at
! Increment the pair index
i_pair = i_pair + 1
end if
end do
end do

npairs = i_pair - 1

end subroutine get_all_unique_pairs

subroutine get_all_unique_pairs_old(indices, heads, list, max_npairs, &
self_interaction, N, neighs, npairs)
! Gets all pairs of atoms that can potentially be neighbours.
!
Expand Down Expand Up @@ -231,7 +363,7 @@ subroutine get_all_unique_pairs(indices, heads, list, max_npairs, &
! this pair when performing the search for neigh_at.
! The following atoms will have even lower indices
! So we can just move to the next bin.
do while (neigh_at >= at)
do while (neigh_at > at)
! Store the pair of potential neighbours.
neighs(i_pair, 1) = at
neighs(i_pair, 2) = neigh_at
Expand All @@ -242,15 +374,112 @@ subroutine get_all_unique_pairs(indices, heads, list, max_npairs, &
neigh_at = list(neigh_at + 1)
end do

! If we are in the first bin of the search, the last atom that
! we found is itself (smallest neigh_at > at is at). Just bring
! the i_pair index down and the pair will get overwritten.
if (j == 1 .and. .not. self_interaction) then
i_pair = i_pair - 1
! If we are in the first bin of everythingthe search, the next atom that should
! have been found is itself. Add it if self interaction was requested.
if (j == 1 .and. self_interaction) then
neighs(i_pair, :) = at
! Increment the pair index
i_pair = i_pair + 1
end if
end do
end do

npairs = i_pair - 1

end subroutine get_all_unique_pairs
end subroutine get_all_unique_pairs_old

subroutine get_pairs_old(at_indices, indices, heads, list, npairs, &
self_interaction, N_ind, neighs, split_indices)
! Gets (possibly duplicated) pairs of atoms that can potentially
! be neighbours.
!
! INPUTS
! ---------
! at_indices:
! The indices of the atoms that we want to get potential
! neighbours for. Atom indices are 0-indexed (python indices)
! indices:
! For each atom index (first dimension), the indices of the
! 8 bins that contain potential neighbours. Bin 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`.
! npairs:
! The number of pairs that are to be found in total. This
! is computed in python with the help of the `counts` array
! constructed by `build_table`.
! self_interaction: bool, optional
! whether to consider an atom a neighbour of itself.
! N_ind:
! The number of atoms for which we will look for neighbours.
! No need to pass it when calling from python, it will be
! inferred from the shape of `at_indices` and `indices`.
!
!
! OUTPUTS
! --------
! neighs:
! contains the atom indices of all pairs that are potential
! neighbours. 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
! for each 8 bin search. It can be used later to quickly
! split `neighs` on the multiple individual searches.
integer, intent(in) :: at_indices(N_ind)
integer, intent(in) :: heads(:), list(:), indices(N_ind, 8)
integer, intent(in) :: npairs
logical, intent(in) :: self_interaction
integer, intent(in) :: N_ind
integer, intent(out) :: neighs(npairs, 2)
integer, intent(out) :: split_indices(N_ind)

! Auxiliar variables
integer :: search_index
integer :: bin_index, at, neigh_at, i_pair, j

! 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
at = at_indices(search_index)
! 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)

! 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)
if (.not. (.not. self_interaction .and. neigh_at == at)) then
! Store the pair of potential neighbours.
neighs(i_pair, 1) = at
neighs(i_pair, 2) = neigh_at
! 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_pairs_old
Loading

0 comments on commit 723dc69

Please sign in to comment.