isdf.F90

Go to the documentation of this file.

!! Copyright (C) 2024 - 2025 A. Buccheri
!!
!! This program is free software; you can redistribute it and/or modify
!! it under the terms of the GNU General Public License as published by
!! the Free Software Foundation; either version 2, or (at your option)
!! any later version.
!!
!! This program is distributed in the hope that it will be useful,
!! but WITHOUT ANY WARRANTY; without even the implied warranty of
!! MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
!! GNU General Public License for more details.
!!
!! You should have received a copy of the GNU General Public License
!! along with this program; if not, write to the Free Software
!! Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
!! 02110-1301, USA.
 
#include "global.h"
 
module isdf_oct_m
  use, intrinsic :: iso_fortran_env, only: real64
  use accel_oct_m,             only: accel_is_enabled
  use batch_oct_m,             only: batch_not_packed, batch_packed, batch_device_packed
  use batch_ops_oct_m,         only: batch_set_state, batch_get_state
  use centroids_oct_m
  use comm_oct_m
  use debug_oct_m
  use distributed_oct_m,       only: distributed_t, distributed_init
  use exchange_operator_oct_m, only: exchange_operator_t
  use global_oct_m
  use io_oct_m
  use isdf_options_oct_m,      only: isdf_options_t
  use isdf_utils_oct_m
  use kpoints_oct_m
  use lalg_basic_oct_m
  use lalg_adv_oct_m
  use math_oct_m
  use mesh_oct_m
  use mesh_function_oct_m,     only: dmf_dotp
  use messages_oct_m
  use mpi_oct_m,               only: mpi_world, mpi_double_precision, mpi_sum
  use namespace_oct_m
  use profiling_oct_m
  use poisson_oct_m
  use space_oct_m
  use states_abst_oct_m,       only: states_are_real
  use states_elec_oct_m
  use xc_cam_oct_m,            only: xc_cam_t
 
  implicit none
  private
  public :: &
    isdf_interpolation_vectors, &
    isdf_gram_matrix, &
    isdf_ace_compute_potentials
 
  ! TODO(Alex) Issue #1195 Extend ISDF to spin-polarised systems
  integer, private, parameter   :: ik = 1
 
contains
 
  subroutine isdf_ace_compute_potentials(exxop, namespace, space, mesh, st, Vx_on_st, kpoints)
    type(exchange_operator_t), intent(in ) :: exxop
    !                                                         ISDF interpolation points, and cam parameters.
    type(namespace_t),         intent(in ) :: namespace
    class(space_t),            intent(in ) :: space
    class(mesh_t),             intent(in ) :: mesh
    type(states_elec_t),       intent(in ) :: st
    type(kpoints_t),           intent(in ) :: kpoints
 
    type(states_elec_t),       intent(out) :: Vx_on_st
 
    real(real64),  allocatable :: psi_mu(:, :), P_r_mu(:, :), isdf_vectors(:, :)
    integer                    :: nstates
    type(distributed_t)        :: isdf_dist
 
    push_sub(isdf_ace_compute_potentials)
 
    ! TODO(Alex) Issue #1195 Extend ISDF to spin-polarised and periodic systems
    assert(kpoints%gamma_only())
    assert(.not. space%is_periodic())
    nstates = exxop%isdf%n_ks_states
 
    call isdf_interpolation_vectors(exxop%isdf, namespace, mesh, st, exxop%isdf%centroids, psi_mu, &
      p_r_mu, isdf_vectors)
 
    !Distribute ISDF vectors along states_kpt (spin) communicator
    call distributed_init(isdf_dist, size(isdf_vectors, 2), st%st_kpt_mpi_grp%comm)
 
    call disdf_ace_apply_exchange_op(exxop, namespace, mesh, st, psi_mu, p_r_mu, isdf_vectors, isdf_dist, vx_on_st)
    safe_deallocate_a(psi_mu)
    safe_deallocate_a(p_r_mu)
    safe_deallocate_a(isdf_vectors)
 
    pop_sub(isdf_ace_compute_potentials)
 
  end subroutine isdf_ace_compute_potentials
 
 
  subroutine isdf_interpolation_vectors(isdf, namespace, mesh, st, centroids, psi_mu, P_r_mu, isdf_vectors)
    type(isdf_options_t),       intent(in ) :: isdf
    type(namespace_t),          intent(in ) :: namespace
    class(mesh_t),              intent(in ) :: mesh
    type(states_elec_t),        intent(in ) :: st
    type(centroids_t),          intent(in ) :: centroids
 
    real(real64),  allocatable, intent(out) :: psi_mu(:, :)
    ! defined at interpolation points: \f$ \psi_i(\mathbf{r}_\mu) \f$
    real(real64), allocatable,  intent(out) :: P_r_mu(:, :)
    ! \f$P_{\mathbf{r},\mu}\f$, with size (np, n_int)
    real(real64), allocatable,  intent(out) :: isdf_vectors(:, :)
 
    character(len=6)           :: np_char
    integer                    :: nocc, n_int_g, rank, i, j
    real(real64),  allocatable :: zct(:, :)
    real(real64),  allocatable :: p_mu_nu(:, :)
    !                                                 with both variables defined at interpolation points.
    real(real64),  allocatable :: cct(:, :)
    !                                                 Gets overwritten with its inverse.
 
    push_sub_with_profile(isdf_interpolation_vectors)
 
    ! TODO(Alex) Issue #1195 Extend ISDF to spin-polarised systems
    if (st%d%nspin > 1) then
      call messages_not_implemented("ISDF for SPIN_POLARIZED and SPINOR calculations", namespace)
    endif
 
    ! TODO(Alex) Issue #1196 Template ISDF handle both real and complex states
    if (.not. states_are_real(st)) then
      call messages_not_implemented("ISDF handling of complex states", namespace)
    endif
 
    ! TODO(Alex) Issue #1276 Implement ISDF on GPU
    if (accel_is_enabled()) then
      call messages_not_implemented("ISDF not supported on GPU", namespace)
    end if
 
    ! For debug file naming
    write(np_char, '(I6)') mpi_world%size
 
    ! Total number of interpolation points
    n_int_g = centroids%npoints_global()
 
    ! Max number of states used in ISDF expansion
    nocc = isdf%n_ks_states
 
    if (st%st_start <= nocc .and. nocc <= st%st_end) then
      write(message(1),'(a, 1x, I3, 1x, a, 1x, I3)') "ISDF: Computing ISDF vectors up to state", &
      & nocc, " on process", st%mpi_grp%rank
      call messages_info(1, namespace=namespace, debug_only=.true.)
    endif
 
    ! psi_mu allocated within the routine as shape varies w.r.t. PACKED or UNPACKED
    call dphi_at_interpolation_points(mesh, st, centroids, nocc, psi_mu)
    if (debug%info) call output_psi_mu_for_all_states(namespace, st, nocc, psi_mu)
 
    safe_allocate(p_r_mu(1:mesh%np, 1:n_int_g))
    call dquasi_density_matrix_at_mesh_centroid_points(st, nocc, psi_mu, p_r_mu)
    if (debug%info) call output_matrix(namespace, "p_r_mu_np"//trim(adjustl(np_char))//".txt", p_r_mu)
 
    safe_allocate(zct(1:mesh%np, 1:n_int_g))
    call pair_product_coefficient_matrix(p_r_mu, zct)
    if (debug%info) call output_matrix(namespace, "zct_np"//trim(adjustl(np_char))//".txt", zct)
 
    ! Note(Alex) Might be more efficient to mask p_r_mu -> p_mu_nu than perform GEMM and allreduce
    safe_allocate(p_mu_nu(1:n_int_g, 1:n_int_g))
    ! Contract over the state index, ist: P_mu_nu = [psi_ist_mu]^T @ psi_ist_nu
    if (local_number_of_states(st, nocc) > 0) then
      call lalg_gemm(psi_mu, psi_mu, p_mu_nu, transa='T')
    else
      ! Process st%mpi_grp%rank contains no states that contribute to ISDF expansion
      ! so avoid GEMM, and initialise to zero for allreduce
      do j = 1, n_int_g
        do i = 1, n_int_g
          p_mu_nu(i, j) = 0.0_real64
        enddo
      enddo
    endif
 
    ! States may be distributed so all elements of p_mu_nu only contain a partial contribution from {ist}
    call comm_allreduce(st%mpi_grp, p_mu_nu)
    if (debug%info) call output_matrix(namespace, "p_mu_nu_np"//trim(adjustl(np_char))//".txt", p_mu_nu)
 
    safe_allocate(cct(1:n_int_g, 1:n_int_g))
    call coefficient_product_matrix(p_mu_nu, cct)
    if (debug%info) then
      call output_matrix(namespace, "cct_np"//trim(adjustl(np_char))//".txt", cct)
      assert(is_symmetric(cct))
    endif
    safe_deallocate_a(p_mu_nu)
 
    ! [CC^T]^{-1}, mutating cct in-place
    ! NOTE, CC^T is extremely ill-conditioned once a critical number of interpolation points
    ! is exceeded. Using the pseudo-inverse (SVD) circumvents this problem
    write(message(1),'(a)') "ISDF: Inverting [CC^T]"
    call messages_info(1, namespace=namespace, debug_only=.true.)
    call lalg_svd_inverse(n_int_g, n_int_g, cct)
    call symmetrize_matrix(n_int_g, cct)
 
    if (isdf%check_n_interp) then
      rank = lalg_matrix_rank_svd(cct, preserve_mat=.true.)
      write(message(1),'(a, I7)') "ISDF: Rank of CC^T is ", rank
      if (rank <  n_int_g) then
        write(message(2),'(a)') " - This rank is the optimal ISDFNpoints to run the calculation with"
      else
        write(message(2),'(a)') " - This suggests that ISDFNpoints is either optimal, or could be larger"
      endif
      call messages_info(2, namespace=namespace)
    endif
 
    ! zeta = [ZC^T][CC^T]^-1
    safe_allocate(isdf_vectors(1:mesh%np, 1:n_int_g))
    call lalg_gemm(mesh%np, n_int_g, n_int_g, 1.0_real64, zct, cct, 0.0_real64, isdf_vectors)
 
    ! ISDF vectors are distributed on the mesh, so do not output in that case
    if (debug%info .and. .not. mesh%parallel_in_domains) then
      call output_matrix(namespace, "isdf_np"//trim(adjustl(np_char))//".txt", isdf_vectors)
    endif
 
    safe_deallocate_a(zct)
    safe_deallocate_a(cct)
 
    pop_sub_with_profile(isdf_interpolation_vectors)
 
  end subroutine isdf_interpolation_vectors
 
 
  subroutine pair_product_coefficient_matrix(p_phi, zct, p_psi)
    real(real64), target,           contiguous, intent(in ) :: p_phi(:, :)
    !                                                                       \f$P^{\varphi}(\mathbf{r}, \mathbf{r}_\mu)\f$
    real(real64), target, optional, contiguous, intent(in ) :: p_psi(:, :)
    !                                                                       \f$P^{\psi}(\mathbf{r}, \mathbf{r}_\mu)\f$
 
    real(real64),                   contiguous, intent(out) :: zct(:, :)
 
    integer :: np, n_int_g, ip, i_mu
    real(real64), pointer, contiguous :: p_2(:, :)
 
    push_sub_with_profile(pair_product_coefficient_matrix)
 
    write(message(1),'(a)') "ISDF: Constructing Z C^T"
    call messages_info(1, debug_only=.true.)
 
    if (present(p_psi)) then
      p_2 => p_psi
    else
      p_2 => p_phi
    endif
 
    ! Quasi-density matrices require the same shape for element-wise multiplication
    assert(all(shape(p_phi) == shape(p_2)))
    ! zct should be allocated, and its shape should be consistent with the quasi-density matrices
    assert(all(shape(p_phi) == shape(zct)))
 
    np = size(p_phi, 1)
    n_int_g = size(p_phi, 2)
 
    ! Construct ZC^T
    !$omp parallel
    do i_mu = 1, n_int_g
      !$omp do simd
      do ip = 1, np
        zct(ip, i_mu) = p_phi(ip, i_mu) * p_2(ip, i_mu)
      enddo
      !$omp end do simd nowait
    enddo
    !$omp end parallel
    nullify(p_2)
 
    pop_sub_with_profile(pair_product_coefficient_matrix)
 
  end subroutine pair_product_coefficient_matrix
 
 
  subroutine coefficient_product_matrix(p_phi, cct, p_psi)
    real(real64), target, contiguous,           intent(in ) :: p_phi(:, :)
    !                                                                         \f$P^{\varphi}(\mathbf{r}_\mu, \mathbf{r}_\nu)\f$
    real(real64), target, contiguous, optional, intent(in ) :: p_psi(:, :)
    !                                                                         \f$P^{\psi}(\mathbf{r}_\mu, \mathbf{r}_\nu)\f$
 
    real(real64),         contiguous,           intent(out) :: cct(:, :)
    !                                                                         Array should be allocated by the caller
 
    integer :: n_int_g, i_mu, i_nu
    real(real64), contiguous, pointer :: p_2(:, :)
 
    push_sub_with_profile(coefficient_product_matrix)
 
    write(message(1),'(a)') "ISDF: Construct CC^T"
    call messages_info(1, debug_only=.true.)
 
    if (present(p_psi)) then
      p_2 => p_psi
    else
      p_2 => p_phi
    endif
 
    ! Quasi-density matrices require the same shape for element-wise multiplication
    assert(all(shape(p_phi) == shape(p_2)))
    ! cct should be allocated, and its shape should be consistent with the quasi-density matrices
    assert(all(shape(p_phi) == shape(cct)))
    n_int_g = size(p_phi, 1)
 
    ! Construct CC^T
    !$omp parallel do collapse(2) default(shared)
    do i_nu = 1, n_int_g
      do i_mu = 1, n_int_g
        cct(i_mu, i_nu) = p_phi(i_mu, i_nu) * p_2(i_mu, i_nu)
      enddo
    enddo
    !$omp end parallel do
    nullify(p_2)
 
    pop_sub_with_profile(coefficient_product_matrix)
 
  end subroutine coefficient_product_matrix
 
 
  subroutine isdf_gram_matrix(mesh, isdf_vectors, gram_matrix)
    class(mesh_t), intent(in)  :: mesh
    real(real64),  contiguous, intent(in ) :: isdf_vectors(:, :)
    real(real64),  contiguous, intent(out) :: gram_matrix(:, :)
 
    integer :: n_int, i, j
 
    push_sub(isdf_gram_matrix)
 
    assert(mesh%np ==  size(isdf_vectors, 1))
    n_int = size(isdf_vectors, 2)
    assert(all(shape(gram_matrix) == [n_int, n_int]))
 
    ! It would be more efficient to use DGEMM, but dmf_dotp ensures the correct volume element
 
    ! Diagonal elements
    do i = 1, n_int
      gram_matrix(i, i) = dmf_dotp(mesh, isdf_vectors(:, i), isdf_vectors(:, i), reduce=.false.)
    enddo
 
    ! Upper triangle elements
    do j = 2, n_int
      do i = 1, j - 1
        gram_matrix(i, j) = dmf_dotp(mesh, isdf_vectors(:, i), isdf_vectors(:, j), reduce=.false.)
        ! Lower triangle from symmetry
        gram_matrix(j, i) = gram_matrix(i, j)
      enddo
    enddo
 
    call mesh%allreduce(gram_matrix)
 
    pop_sub(isdf_gram_matrix)
 
  end subroutine isdf_gram_matrix
 
#include "real.F90"
#include "isdf_inc.F90"
#include "undef.F90"
 
end module isdf_oct_m
 
!! Local Variables:
!! mode: f90
!! coding: utf-8
!! End: