get_linked_cell Subroutine

subroutine get_linked_cell(mol, n_xyz, head, nxt, lat_inv, cell_w)
   !> Stucture type
   type(structure_type), intent(in) :: mol

   !> Number of linked cells along the x, y, and z directions
   integer, intent(in) :: n_xyz(3)

   !> Flattened cell-to-atom chain heads list
   integer, intent(out) :: head(:)

   !> Next atom in each linked-cell chain, indexed by atom
   integer, intent(out) :: nxt(:)

   !> Inverse of the lattice matrix
   real(wp), intent(in), optional :: lat_inv(3, 3)

   !> Width of each grid cell
   real(wp), intent(in), optional :: cell_w(3)

   integer :: iat, ix, iy, iz, ic
   real(wp) :: fract(3), min_xyz(3)

   head = 0

   if (any(mol%periodic) .and. present(lat_inv)) then
      !$omp parallel do default(none) &
      !$omp private(iat, fract, ix, iy, iz, ic) &
      !$omp shared(head, nxt, n_xyz, lat_inv, mol)
      do iat = 1, mol%nat
         fract(:) = matmul(lat_inv, mol%xyz(:, iat))
         fract(:) = fract(:) - floor(fract(:))

         ix = min(n_xyz(1), max(1, int(fract(1) * n_xyz(1)) + 1))
         iy = min(n_xyz(2), max(1, int(fract(2) * n_xyz(2)) + 1))
         iz = min(n_xyz(3), max(1, int(fract(3) * n_xyz(3)) + 1))

         ic = ix + n_xyz(1)*(iy-1) + n_xyz(1)*n_xyz(2)*(iz-1)

         !$omp atomic capture
         nxt(iat) = head(ic)
         head(ic) = iat
         !$omp end atomic
      end do
      !$omp end parallel do
   else if (present(cell_w)) then
      min_xyz = minval(mol%xyz, dim=2) - buffer

      !$omp parallel do default(none) &
      !$omp private(iat, ix, iy, iz, ic) &
      !$omp shared(head, nxt, n_xyz, cell_w, mol, min_xyz)
      do iat = 1, mol%nat
         ix = min(n_xyz(1), &
            & max(1, int((mol%xyz(1, iat) - min_xyz(1))/cell_w(1)) + 1))
         iy = min(n_xyz(2), &
            & max(1, int((mol%xyz(2, iat) - min_xyz(2))/cell_w(2)) + 1))
         iz = min(n_xyz(3), &
            & max(1, int((mol%xyz(3, iat) - min_xyz(3))/cell_w(3)) + 1))

         ic = ix + n_xyz(1)*(iy-1) + n_xyz(1)*n_xyz(2)*(iz-1)

         !$omp atomic capture
         nxt(iat) = head(ic)
         head(ic) = iat
         !$omp end atomic

      end do
      !$omp end parallel do
   end if

end subroutine get_linked_cell