nonlocal_overlap_matrix.F90

Go to the documentation of this file.

!! Copyright (C) 2026 N. Tancogne-Dejean
!!
!! 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 nonlocal_overlap_matrix_oct_m
  use accel_oct_m
  use debug_oct_m
  use epot_oct_m
  use global_oct_m
  use lalg_adv_oct_m
  use mesh_oct_m
  use messages_oct_m
  use profiling_oct_m
  use projector_oct_m
  use projector_matrix_oct_m
  use space_oct_m
  use submesh_oct_m
 
  implicit none
 
  private
  public ::             &
    overlap_cluster_t,  &
    overlap_matrix_t
 
  type overlap_cluster_t
    integer :: nmat = 0
    integer, allocatable :: imats(:)
    integer, allocatable :: coff(:)
    integer :: nc = 0
    logical :: is_cmplx = .false.             
    real(real64),    allocatable :: dt(:, :, :)
    complex(real64), allocatable :: zt(:, :, :)
  end type overlap_cluster_t
 
  type overlap_matrix_t
    logical, public :: use_global_coupling = .false.
    logical, public :: coupling_built = .false. 
    integer :: ncluster = 0
    integer, allocatable :: proj_offset(:)
    type(overlap_cluster_t), allocatable :: clusters(:)
  contains
    procedure :: detect => overlap_matrix_detect
    procedure :: build  => overlap_matrix_build
    procedure :: end    => overlap_matrix_end
    procedure :: dapply => doverlap_matrix_apply
    procedure :: zapply => zoverlap_matrix_apply
  end type overlap_matrix_t
 
  logical :: uspp_gpu_approx_warned = .false.
 
contains
 
  !--------------------------------------------------------
  subroutine overlap_matrix_detect(this, nmat, projector_matrices, projector_to_atom, epot, space, self_overlap)
    class(overlap_matrix_t),  intent(inout) :: this
    integer,                  intent(in)    :: nmat
    type(projector_matrix_t), intent(in)    :: projector_matrices(:)
    integer,                  intent(in)    :: projector_to_atom(:)
    type(epot_t),             intent(in)    :: epot
    class(space_t),           intent(in)    :: space
    logical,                  intent(in)    :: self_overlap
 
    integer :: im, jm, ia, ja, ic, root, nroots, kk
    integer, allocatable :: parent(:), root_of(:), csize(:), cluster_of(:), fill(:)
    logical :: any_overlap
 
    push_sub(overlap_matrix_detect)
 
    call overlap_matrix_end(this)
 
    ! projection-vector offset of each matrix` slice (matches ind(imat) in the apply path)
    safe_allocate(this%proj_offset(1:max(nmat, 1)))
    kk = 0
    do im = 1, nmat
      this%proj_offset(im) = kk
      kk = kk + projector_matrices(im)%nprojs
    end do
 
    if (nmat == 0) then
      pop_sub(overlap_matrix_detect)
      return
    end if
 
    ! union-find over projector matrices whose spheres share grid points (distinct atoms only:
    ! a single matrix per atom, so same-atom pairs never arise here)
    safe_allocate(parent(1:nmat))
    do im = 1, nmat
      parent(im) = im
    end do
    any_overlap = .false.
    do im = 1, nmat
      ia = projector_to_atom(im)
      do jm = im + 1, nmat
        ja = projector_to_atom(jm)
        if (ia == ja) cycle
        if (submesh_overlap(epot%proj(ia)%sphere, epot%proj(ja)%sphere, space)) then
          call uf_union(parent, im, jm)
          any_overlap = .true.
        end if
      end do
    end do
 
    ! the per-atom Gram ignores the image cross terms of a self-overlapping sphere
    this%use_global_coupling = any_overlap .or. self_overlap
    if (.not. this%use_global_coupling) then
      safe_deallocate_a(parent)
      pop_sub(overlap_matrix_detect)
      return
    end if
 
    ! materialize the connected components as clusters
    safe_allocate(cluster_of(1:nmat))
    safe_allocate(root_of(1:nmat))
    safe_allocate(csize(1:nmat))
    nroots = 0
    csize = 0
    do im = 1, nmat
      root = uf_find(parent, im)
      ic = 0
      do kk = 1, nroots
        if (root_of(kk) == root) then
          ic = kk
          exit
        end if
      end do
      if (ic == 0) then
        nroots = nroots + 1
        root_of(nroots) = root
        ic = nroots
      end if
      cluster_of(im) = ic
      csize(ic) = csize(ic) + 1
    end do
 
    this%ncluster = nroots
    safe_allocate(this%clusters(1:nroots))
    safe_allocate(fill(1:nroots))
    fill = 0
    do ic = 1, nroots
      this%clusters(ic)%nmat = csize(ic)
      safe_allocate(this%clusters(ic)%imats(1:csize(ic)))
      safe_allocate(this%clusters(ic)%coff(1:csize(ic)))
    end do
    ! matrices are visited in ascending order, so imats(:) comes out sorted
    do im = 1, nmat
      ic = cluster_of(im)
      fill(ic) = fill(ic) + 1
      this%clusters(ic)%imats(fill(ic)) = im
    end do
    do ic = 1, nroots
      kk = 0
      do im = 1, this%clusters(ic)%nmat
        this%clusters(ic)%coff(im) = kk
        kk = kk + projector_matrices(this%clusters(ic)%imats(im))%nprojs
      end do
      this%clusters(ic)%nc = kk
    end do
 
    safe_deallocate_a(parent)
    safe_deallocate_a(cluster_of)
    safe_deallocate_a(root_of)
    safe_deallocate_a(csize)
    safe_deallocate_a(fill)
 
    pop_sub(overlap_matrix_detect)
  end subroutine overlap_matrix_detect
 
  !--------------------------------------------------------
  recursive integer function uf_find(parent, ii) result(rr)
    integer, intent(inout) :: parent(:)
    integer, intent(in)    :: ii
 
    if (parent(ii) /= ii) parent(ii) = uf_find(parent, parent(ii))
    rr = parent(ii)
  end function uf_find
 
  !--------------------------------------------------------
  subroutine uf_union(parent, ii, jj)
    integer, intent(inout) :: parent(:)
    integer, intent(in)    :: ii, jj
 
    integer :: ri, rj
 
    ri = uf_find(parent, ii)
    rj = uf_find(parent, jj)
    if (ri /= rj) parent(rj) = ri
  end subroutine uf_union
 
  ! ----------------------------------------------------------------------------------
  subroutine overlap_matrix_build(this, projector_matrices, projector_phases, mesh, kstart, kend)
    class(overlap_matrix_t),  target,      intent(inout) :: this
    type(projector_matrix_t),              intent(in)    :: projector_matrices(:)
    complex(real64),          allocatable, intent(in)    :: projector_phases(:, :, :, :)
    class(mesh_t),                         intent(in)    :: mesh
    integer,                               intent(in)    :: kstart
    integer,                               intent(in)    :: kend
 
    integer :: c, ik, nc, a, b, oa, ob, ii, jj, nI
    logical :: phased
    type(overlap_cluster_t),  pointer :: cl
    complex(real64), allocatable :: zgI(:, :), zgJ(:, :), mmat(:, :), qmat(:, :), amat(:, :), tmat(:, :)
 
    push_sub(overlap_matrix_build)
 
    if (.not. this%use_global_coupling) then
      pop_sub(overlap_matrix_build)
      return
    end if
 
    ! The exact cross-block inverse is applied on the CPU only. On the GPU the coupling still uses
    ! the per-atom (block-diagonal) buff_mix, i.e. the previous approximation; warn once. (The
    ! device dense coupling is a follow-up that needs GPU hardware to implement and validate.)
    if (accel_is_enabled() .and. .not. uspp_gpu_approx_warned) then
      uspp_gpu_approx_warned = .true.
      message(1) = "Ultrasoft projector spheres of different atoms overlap on the grid."
      message(2) = "The exact inverse overlap S^{-1} is applied on the CPU only; the GPU path"
      message(3) = "still uses the per-atom approximation. Convergence on GPU may be degraded."
      call messages_warning(3)
    end if
 
    phased = allocated(projector_phases)
 
    do c = 1, this%ncluster
      cl => this%clusters(c)
      nc = cl%nc
      ! T is complex if the projectors are phased (k /= 0 / vector potential); otherwise a real
      ! T is stored.
      cl%is_cmplx = phased .or. projector_matrices(cl%imats(1))%is_cmplx
 
      if (cl%is_cmplx) then
        if (.not. allocated(cl%zt)) then
          safe_allocate(cl%zt(1:nc, 1:nc, kstart:kend))
        end if
      else
        if (.not. allocated(cl%dt)) then
          safe_allocate(cl%dt(1:nc, 1:nc, kstart:kend))
        end if
      end if
 
      safe_allocate(mmat(1:nc, 1:nc))
      safe_allocate(qmat(1:nc, 1:nc))
      safe_allocate(amat(1:nc, 1:nc))
      safe_allocate(tmat(1:nc, 1:nc))
 
      do ik = kstart, kend
 
        ! Cluster Gram M(k): local contribution over shared grid points, upper blocks only
        mmat = m_z0
        do a = 1, cl%nmat
          call fill_phased_ket(cl%imats(a), ik, zgi)
          oa = cl%coff(a)
          do b = a, cl%nmat
            call fill_phased_ket(cl%imats(b), ik, zgj)
            ob = cl%coff(b)
            call accumulate_gram(projector_matrices(cl%imats(a)), &
              projector_matrices(cl%imats(b)), zgi, zgj, mmat, oa, ob)
            safe_deallocate_a(zgj)
          end do
          safe_deallocate_a(zgi)
        end do
        ! fill the strict-lower entries by Hermiticity (diagonal blocks are already full)
        do jj = 1, nc
          do ii = jj + 1, nc
            mmat(ii, jj) = conjg(mmat(jj, ii))
          end do
        end do
        ! the sphere may be split across domains
        if (mesh%parallel_in_domains) call mesh%allreduce(mmat, dim = (/nc, nc/))
        mmat = mmat * mesh%volume_element
 
        ! Block-diagonal augmentation charges Q (per atom; no cross-atom q terms)
        qmat = m_z0
        do a = 1, cl%nmat
          ni = projector_matrices(cl%imats(a))%nprojs
          oa = cl%coff(a)
          qmat(oa+1:oa+ni, oa+1:oa+ni) = &
            cmplx(projector_matrices(cl%imats(a))%qmix(1:ni, 1:ni), m_zero, real64)
        end do
 
        ! A = 1 + M Q ; T = Q A^{-1}
        amat = matmul(mmat, qmat)
        do ii = 1, nc
          amat(ii, ii) = amat(ii, ii) + m_z1
        end do
        call lalg_inverse(nc, amat, 'svd')
        tmat = matmul(qmat, amat)
 
        if (cl%is_cmplx) then
          cl%zt(1:nc, 1:nc, ik) = -tmat(1:nc, 1:nc)
        else
          cl%dt(1:nc, 1:nc, ik) = -real(tmat(1:nc, 1:nc), real64)
        end if
 
      end do
 
      safe_deallocate_a(mmat)
      safe_deallocate_a(qmat)
      safe_deallocate_a(amat)
      safe_deallocate_a(tmat)
    end do
 
    this%coupling_built = .true.
 
    pop_sub(overlap_matrix_build)
 
  contains
 
    ! Build the phased ket weight g(p, i) = beta_i(p) * conjg(phase(p)) for projector matrix im.
    subroutine fill_phased_ket(im, ik, zg)
      integer,                      intent(in)  :: im
      integer,                      intent(in)  :: ik
      complex(real64), allocatable, intent(out) :: zg(:, :)
 
      integer :: np, npr, p, i
 
      np = projector_matrices(im)%npoints
      npr = projector_matrices(im)%nprojs
      safe_allocate(zg(1:max(np, 1), 1:npr))
      if (np == 0) return
 
      if (allocated(projector_matrices(im)%zprojectors)) then
        do i = 1, npr
          do p = 1, np
            zg(p, i) = projector_matrices(im)%zprojectors(p, i)
          end do
        end do
      else
        do i = 1, npr
          do p = 1, np
            zg(p, i) = cmplx(projector_matrices(im)%dprojectors(p, i), m_zero, real64)
          end do
        end do
      end if
 
      if (phased) then
        do i = 1, npr
          do p = 1, np
            zg(p, i) = zg(p, i) * conjg(projector_phases(p, 1, im, ik))
          end do
        end do
      end if
    end subroutine fill_phased_ket
 
    ! Accumulate M[oa+i, ob+j] += sum_{shared points} conjg(gI) gJ.
    ! A self-overlapping sphere stores its map as injective regions,
    ! so each region pair is merged; this covers every pair of
    ! sphere points sharing a grid point, periodic images included.
    subroutine accumulate_gram(pmI, pmJ, zgI, zgJ, mmat, oa, ob)
      type(projector_matrix_t), intent(in)    :: pmI
      type(projector_matrix_t), intent(in)    :: pmJ
      complex(real64),          intent(in)    :: zgI(:, :)
      complex(real64),          intent(in)    :: zgJ(:, :)
      complex(real64),          intent(inout) :: mmat(:, :)
      integer,                  intent(in)    :: oa
      integer,                  intent(in)    :: ob
 
      integer :: ri, rj, ip, jp, ipe, jpe, ni, nj, i, j, dd
      complex(real64) :: bra
 
      ni = size(zgi, 2)
      nj = size(zgj, 2)
      do ri = 1, pmi%nregions
        ipe = pmi%regions(ri + 1) - 1
        do rj = 1, pmj%nregions
          jpe = pmj%regions(rj + 1) - 1
          ip = pmi%regions(ri)
          jp = pmj%regions(rj)
          do while (ip <= ipe .and. jp <= jpe)
            dd = pmi%map(ip) - pmj%map(jp)
            if (dd < 0) then
              ip = ip + 1
            else if (dd > 0) then
              jp = jp + 1
            else
              do i = 1, ni
                bra = conjg(zgi(ip, i))
                do j = 1, nj
                  mmat(oa+i, ob+j) = mmat(oa+i, ob+j) + bra*zgj(jp, j)
                end do
              end do
              ip = ip + 1
              jp = jp + 1
            end if
          end do
        end do
      end do
    end subroutine accumulate_gram
 
  end subroutine overlap_matrix_build
 
  !-----------------------------------------------------------------
  subroutine overlap_matrix_end(this)
    class(overlap_matrix_t), intent(inout) :: this
 
    integer :: ic
 
    push_sub(overlap_matrix_end)
 
    if (allocated(this%clusters)) then
      do ic = 1, size(this%clusters)
        safe_deallocate_a(this%clusters(ic)%imats)
        safe_deallocate_a(this%clusters(ic)%coff)
        safe_deallocate_a(this%clusters(ic)%dt)
        safe_deallocate_a(this%clusters(ic)%zt)
      end do
      safe_deallocate_a(this%clusters)
    end if
    safe_deallocate_a(this%proj_offset)
    this%ncluster = 0
    this%use_global_coupling = .false.
    this%coupling_built = .false.
 
    pop_sub(overlap_matrix_end)
  end subroutine overlap_matrix_end
 
#include "undef.F90"
#include "real.F90"
#include "nonlocal_overlap_matrix_inc.F90"
 
#include "undef.F90"
#include "complex.F90"
#include "nonlocal_overlap_matrix_inc.F90"
 
end module nonlocal_overlap_matrix_oct_m
 
!! Local Variables:
!! mode: f90
!! coding: utf-8
!! End: