get_dispersion3_hessian Subroutine

subroutine get_dispersion3_hessian(error, mol, disp, param, cutoff, hessian, partition)
   !DEC$ ATTRIBUTES DLLEXPORT :: get_dispersion3_hessian

   !> Error handling
   type(error_type), allocatable, intent(out) :: error

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

   !> Dispersion model
   class(dispersion_model), intent(in) :: disp

   !> Damping parameters
   class(damping_param), intent(in) :: param

   !> Realspace cutoffs
   type(realspace_cutoff), intent(in) :: cutoff

   !> ATM dispersion Hessian
   real(wp), intent(out) :: hessian(:, :)

   !> Optional externally assigned work partition
   type(work_partition), intent(in), optional :: partition

   integer :: mref, nat, ndim, ii, jj, icn, jcn
   real(wp), allocatable, target :: dcndr(:, :, :)
   real(wp), allocatable :: cn(:), dcndL(:, :, :)
   real(wp), allocatable :: q(:)
   real(wp), allocatable :: gwvec(:, :, :), gwdcn(:, :, :), gwd2cn(:, :, :)
   real(wp), allocatable :: c6(:, :), dc6dcn(:, :)
   real(wp), allocatable :: d2c6dcn2(:, :), d2c6dcnij(:, :)
   real(wp), allocatable :: dEdcn(:), dEdcndr(:, :), dEdcndcn(:, :), cnwork(:, :)
   real(wp), allocatable :: drt(:, :), dEdcndrt(:, :)
   real(wp), pointer :: dr(:, :)
   real(wp) :: tmp
   real(wp), allocatable :: lattr(:, :)

   nat = mol%nat
   ndim = 3*nat
   mref = maxval(disp%ref)

   ! mctc-lib currently evaluates Cartesian CN derivatives only when both
   ! dcndr and dcndL are requested.  Keep dcndL even though the ATM Hessian
   ! itself only consumes dcndr.
   allocate(cn(nat), dcndr(3, nat, nat), dcndL(3, 3, nat))
   call get_lattice_points(mol%periodic, mol%lattice, cutoff%cn, lattr)
   call get_coordination_number(mol, lattr, cutoff%cn, disp%rcov, disp%en, cn, dcndr, dcndL, error)
   if (allocated(error)) return

   ! The D4 ATM contribution is evaluated with charge-independent C6
   ! coefficients, matching the q=0 branch in get_dispersion.
   allocate(q(nat), source=0.0_wp)
   allocate(gwvec(mref, nat, disp%ncoup), gwdcn(mref, nat, disp%ncoup), &
      & gwd2cn(mref, nat, disp%ncoup))
   allocate(c6(nat, nat), dc6dcn(nat, nat), d2c6dcn2(nat, nat), d2c6dcnij(nat, nat))

   select type(disp)
   type is (d4_model)
      call disp%weight_references_hessian(mol, cn, q, gwvec, gwdcn, gwd2cn)
      call disp%get_atomic_c6_hessian(mol, gwvec, gwdcn, gwd2cn, c6, &
         & dc6dcn, d2c6dcn2, d2c6dcnij)
   type is (d4s_model)
      call disp%weight_references_hessian(mol, cn, q, gwvec, gwdcn, gwd2cn)
      call disp%get_atomic_c6_hessian(mol, gwvec, gwdcn, gwd2cn, c6, &
         & dc6dcn, d2c6dcn2, d2c6dcnij)
   class default
      call fatal_error(error, "Analytical ATM Hessian is only implemented for D4 and D4S models")
      return
   end select

   hessian(:, :) = 0.0_wp
   allocate(dEdcn(nat), source=0.0_wp)
   allocate(dEdcndr(ndim, nat), source=0.0_wp)
   allocate(dEdcndcn(nat, nat), source=0.0_wp)

   call get_lattice_points(mol%periodic, mol%lattice, cutoff%disp3, lattr)
   select type(param)
   class is (rational_damping_param)
      call param%get_dispersion3_hessian(mol, lattr, cutoff%disp3, cutoff%width3, &
         & disp%r4r2, c6, dc6dcn, d2c6dcn2, d2c6dcnij, hessian, dEdcn, &
         & dEdcndr, dEdcndcn, partition)
   class default
      call fatal_error(error, "Analytical ATM Hessian is only implemented for rational damping")
      return
   end select

   call get_lattice_points(mol%periodic, mol%lattice, cutoff%cn, lattr)
   call add_coordination_number_hessian(mol, lattr, cutoff%cn, disp%rcov, disp%en, &
      & dEdcn, hessian, error)
   if (allocated(error)) return

   dr(1:ndim, 1:nat) => dcndr
   allocate(cnwork(nat, ndim), drt(nat, ndim), dEdcndrt(nat, ndim))

   ! Assemble the CN chain rule explicitly under OpenMP.  Intrinsic MATMUL is
   ! not guaranteed to use the program's OpenMP team (and is often serial with
   ! a single-threaded BLAS), which otherwise leaves a sizeable serial tail in
   ! analytical ATM Hessian calculations.
   !$omp parallel default(none) &
   !$omp shared(nat, ndim, dEdcndr, dEdcndcn, dr, drt, dEdcndrt, cnwork, hessian) &
   !$omp private(ii, jj, icn, jcn, tmp)

   ! Store the Cartesian-by-CN matrices with CN as the contiguous dimension.
   ! Both contractions below then stream through contiguous columns instead of
   ! repeatedly traversing the second dimension of Fortran arrays.
   !$omp do collapse(2) schedule(static)
   do ii = 1, ndim
      do icn = 1, nat
         drt(icn, ii) = dr(ii, icn)
         dEdcndrt(icn, ii) = dEdcndr(ii, icn)
      end do
   end do
   !$omp end do

   ! cnwork = dEdcndcn * transpose(dr).  Keeping jcn as the accumulation
   ! order preserves the scalar result while the icn loop is vectorizable.
   !$omp do schedule(static)
   do jj = 1, ndim
      cnwork(:, jj) = 0.0_wp
      do jcn = 1, nat
         tmp = drt(jcn, jj)
         !$omp simd
         do icn = 1, nat
            cnwork(icn, jj) = cnwork(icn, jj) + dEdcndcn(icn, jcn)*tmp
         end do
      end do
   end do
   !$omp end do

   ! H += dEdcndr*dr^T + dr*dEdcndr^T + dr*dEdcndcn*dr^T
   !$omp do collapse(2) schedule(static)
   do jj = 1, ndim
      do ii = 1, ndim
         tmp = 0.0_wp
         !$omp simd reduction(+:tmp)
         do icn = 1, nat
            tmp = tmp + dEdcndrt(icn, ii)*drt(icn, jj) &
               & + drt(icn, ii)*dEdcndrt(icn, jj) &
               & + drt(icn, ii)*cnwork(icn, jj)
         end do
         hessian(ii, jj) = hessian(ii, jj) + tmp
      end do
   end do
   !$omp end do
   !$omp end parallel

end subroutine get_dispersion3_hessian