get_dispersion_spme Subroutine

subroutine get_dispersion_spme(mol, lowrank, ghost, terms, nterm, mesh, gwvec, &
      & gwdcn, energies, dEdcn, gradient, sigma, partition)

   !> Molecular structure data
   class(structure_type), intent(in) :: mol

   !> Separable representation of the C6 coefficients
   class(d3_lowrank_c6), intent(in) :: lowrank

   !> Atoms excluded from the dispersion calculation
   logical, intent(in) :: ghost(:)

   !> Terms of the damped pair potential for each pair of species
   type(fourier_term), intent(in) :: terms(:, :, :)

   !> Number of terms for each pair of species
   integer, intent(in) :: nterm(:, :)

   !> Number of mesh points along each direction
   integer, intent(in) :: mesh(3)

   !> Weighting function for the atomic reference systems
   real(wp), intent(in) :: gwvec(:, :)

   !> Derivative of the weighting function w.r.t. the coordination number
   real(wp), intent(in), optional :: gwdcn(:, :)

   !> Dispersion energy
   real(wp), intent(inout) :: energies(:)

   !> Derivative of the energy w.r.t. the coordination number
   real(wp), intent(inout), optional :: dEdcn(:)

   !> Dispersion gradient
   real(wp), intent(inout), optional :: gradient(:, :)

   !> Dispersion virial
   real(wp), intent(inout), optional :: sigma(:, :)

   !> Work partition over the terms of the low-rank expansion
   type(work_partition), intent(in), optional :: partition

   logical :: grad
   type(fft_mesh) :: fmesh
   integer :: nat, nid, rank, iat, izp, jzp, il, it, ic, jc
   integer :: i1, i2, i3, j1, j2, j3, g1, g2, g3
   real(wp) :: vol, rec(3, 3), invlat(3, 3), kvec(3), knorm, erecip
   real(wp) :: pval, dpval, phi, dphi, wgt, dwgt(3), sfac, qq, escale
   real(wp) :: dudr(3, 3), tvec(3), dtvec(3)
   integer, allocatable :: base(:, :)
   real(wp), allocatable :: theta(:, :, :), dtheta(:, :, :)
   real(wp), allocatable :: bfac(:, :), c6l(:, :), dc6ldcn(:, :)
   real(wp), allocatable :: phik(:, :, :, :), dphik(:, :, :, :)
   real(wp), allocatable :: pot(:, :), esum(:), dsum(:), gsum(:, :), ssum(:, :)
   complex(wp), allocatable :: qgrid(:, :, :, :), tgrid(:, :, :)

   grad = present(gwdcn) .and. present(dEdcn) .and. present(gradient) &
      & .and. present(sigma)

   nat = mol%nat
   nid = mol%nid
   rank = lowrank%rank

   invlat(:, :) = matinv_3x3(mol%lattice)
   rec(:, :) = 2.0_wp * pi * transpose(invlat)
   vol = abs(matdet_3x3(mol%lattice))

   allocate(c6l(rank, nat))
   allocate(dc6ldcn(rank, nat), source=0.0_wp)
   if (grad) then
      call lowrank%get_weights(mol, ghost, gwvec, c6l, gwdcn, dc6ldcn)
   else
      call lowrank%get_weights(mol, ghost, gwvec, c6l)
   end if

   allocate(esum(nat), dsum(nat), source=0.0_wp)
   allocate(gsum(3, nat), source=0.0_wp)
   allocate(ssum(3, 3), source=0.0_wp)

   ! the unrestricted double sum contains the interaction of every atom with
   ! itself, which the value of the potential at the origin removes again
   do iat = 1, nat
      izp = mol%id(iat)
      pval = 0.0_wp
      do it = 1, nterm(izp, izp)
         pval = pval + get_potential_zero(terms(it, izp, izp))
      end do
      do il = 1, rank
         if (.not.owns_index(partition, il)) cycle
         esum(iat) = esum(iat) &
            & + 0.5_wp * lowrank%lambda(il) * c6l(il, iat)**2 * pval
         dsum(iat) = dsum(iat) &
            & + lowrank%lambda(il) * c6l(il, iat) * dc6ldcn(il, iat) * pval
      end do
   end do

   call new_fft_mesh(fmesh, mesh)
   call get_spline_weights(mol, mesh, base, theta, dtheta)
   call get_euler_factors(mesh, bfac)

   ! the transform of the pair potential depends on the species pair but not on
   ! the term of the expansion, so it is tabulated once for the whole mesh
   call get_mesh_kernel(mesh, rec, terms, nterm, nid, grad, phik, dphik)

   do ic = 1, 3
      do jc = 1, 3
         dudr(jc, ic) = mesh(ic) * rec(jc, ic) / (2.0_wp * pi)
      end do
   end do

   allocate(qgrid(mesh(1), mesh(2), mesh(3), nid))
   allocate(tgrid(mesh(1), mesh(2), mesh(3)))
   allocate(pot(nat, nid))

   erecip = 0.0_wp
   escale = -0.5_wp / vol

   ! the terms of the expansion decouple, so only one set of meshes is live and
   ! a partition can hand out whole terms
   do il = 1, rank
      if (.not.owns_index(partition, il)) cycle

      qgrid(:, :, :, :) = cmplx(0.0_wp, 0.0_wp, wp)
      do iat = 1, nat
         if (ghost(iat)) cycle
         izp = mol%id(iat)
         do j3 = 1, spline_order
            g3 = modulo(base(3, iat) + j3, mesh(3)) + 1
            do j2 = 1, spline_order
               g2 = modulo(base(2, iat) + j2, mesh(2)) + 1
               do j1 = 1, spline_order
                  g1 = modulo(base(1, iat) + j1, mesh(1)) + 1
                  wgt = theta(j1, 1, iat) * theta(j2, 2, iat) * theta(j3, 3, iat)
                  qgrid(g1, g2, g3, izp) = qgrid(g1, g2, g3, izp) &
                     & + cmplx(c6l(il, iat) * wgt, 0.0_wp, wp)
               end do
            end do
         end do
      end do

      do izp = 1, nid
         call fft_3d(fmesh, qgrid(:, :, :, izp), +1)
      end do

      ! contract the structure factors with the kernel to obtain the potential
      ! each species experiences, and accumulate energy and virial on the way
      do izp = 1, nid
         tgrid(:, :, :) = cmplx(0.0_wp, 0.0_wp, wp)
         !$omp parallel do schedule(runtime) default(none) collapse(2) &
         !$omp shared(mesh, nid, izp, phik, bfac, qgrid, tgrid) &
         !$omp private(i1, i2, i3, jzp, sfac)
         do i3 = 1, mesh(3)
            do i2 = 1, mesh(2)
               do i1 = 1, mesh(1)
                  sfac = bfac(i1, 1) * bfac(i2, 2) * bfac(i3, 3)
                  do jzp = 1, nid
                     tgrid(i1, i2, i3) = tgrid(i1, i2, i3) + sfac &
                        & * phik(pair_index(izp, jzp), i1, i2, i3) &
                        & * qgrid(i1, i2, i3, jzp)
                  end do
               end do
            end do
         end do

         call fft_3d(fmesh, tgrid, -1)

         ! gather the potential back onto the atoms of this species
         do iat = 1, nat
            pot(iat, izp) = 0.0_wp
         end do
         do iat = 1, nat
            if (ghost(iat) .or. mol%id(iat) /= izp) cycle
            wgt = 0.0_wp
            dwgt(:) = 0.0_wp
            do j3 = 1, spline_order
               g3 = modulo(base(3, iat) + j3, mesh(3)) + 1
               do j2 = 1, spline_order
                  g2 = modulo(base(2, iat) + j2, mesh(2)) + 1
                  do j1 = 1, spline_order
                     g1 = modulo(base(1, iat) + j1, mesh(1)) + 1
                     sfac = real(tgrid(g1, g2, g3), wp)
                     tvec(1) = theta(j1, 1, iat)
                     tvec(2) = theta(j2, 2, iat)
                     tvec(3) = theta(j3, 3, iat)
                     wgt = wgt + tvec(1) * tvec(2) * tvec(3) * sfac
                     if (grad) then
                        dtvec(1) = dtheta(j1, 1, iat) * tvec(2) * tvec(3)
                        dtvec(2) = tvec(1) * dtheta(j2, 2, iat) * tvec(3)
                        dtvec(3) = tvec(1) * tvec(2) * dtheta(j3, 3, iat)
                        dwgt(:) = dwgt(:) + matmul(dudr, dtvec) * sfac
                     end if
                  end do
               end do
            end do
            pot(iat, izp) = wgt
            esum(iat) = esum(iat) + escale * lowrank%lambda(il) * c6l(il, iat) * wgt
            erecip = erecip + escale * lowrank%lambda(il) * c6l(il, iat) * wgt
            if (grad) then
               dsum(iat) = dsum(iat) &
                  & + 2.0_wp * escale * lowrank%lambda(il) * dc6ldcn(il, iat) * wgt
               gsum(:, iat) = gsum(:, iat) &
                  & + 2.0_wp * escale * lowrank%lambda(il) * c6l(il, iat) * dwgt(:)
            end if
         end do
      end do

      if (grad) then
         !$omp parallel do schedule(runtime) default(none) collapse(2) &
         !$omp shared(mesh, nid, rec, bfac, dphik, qgrid, lowrank, il, vol) &
         !$omp private(i1, i2, i3, j1, j2, j3, izp, jzp, kvec, knorm, qq, sfac, ic) &
         !$omp reduction(+:ssum)
         do i3 = 1, mesh(3)
            do i2 = 1, mesh(2)
               do i1 = 1, mesh(1)
                  j1 = i1 - 1 - mesh(1) * ((2*(i1 - 1)) / mesh(1))
                  j2 = i2 - 1 - mesh(2) * ((2*(i2 - 1)) / mesh(2))
                  j3 = i3 - 1 - mesh(3) * ((2*(i3 - 1)) / mesh(3))
                  kvec(:) = rec(:, 1)*j1 + rec(:, 2)*j2 + rec(:, 3)*j3
                  knorm = norm2(kvec)
                  if (knorm <= 0.0_wp) cycle
                  sfac = bfac(i1, 1) * bfac(i2, 2) * bfac(i3, 3)
                  qq = 0.0_wp
                  do izp = 1, nid
                     do jzp = 1, nid
                        qq = qq + lowrank%lambda(il) &
                           & * dphik(pair_index(izp, jzp), i1, i2, i3) * sfac &
                           & * real(qgrid(i1, i2, i3, izp) &
                           &        * conjg(qgrid(i1, i2, i3, jzp)), wp)
                     end do
                  end do
                  do ic = 1, 3
                     ssum(ic, :) = ssum(ic, :) &
                        & + 0.5_wp * qq * kvec(ic) * kvec(:) / (knorm * vol)
                  end do
               end do
            end do
         end do
      end if
   end do

   energies(:) = energies(:) + esum(:)
   if (grad) then
      do ic = 1, 3
         ssum(ic, ic) = ssum(ic, ic) - erecip
      end do
      dEdcn(:) = dEdcn(:) + dsum(:)
      gradient(:, :) = gradient(:, :) + gsum(:, :)
      sigma(:, :) = sigma(:, :) + ssum(:, :)
   end if

end subroutine get_dispersion_spme