isdf_utils.F90

Go to the documentation of this file.

!! Copyright (C) 2025. Alexander 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, U
 
#include "global.h"
 
module isdf_utils_oct_m
  use, intrinsic :: iso_fortran_env, only: real64, int64
  use batch_oct_m,          only: batch_packed
  use comm_oct_m,           only: comm_allreduce
  use debug_oct_m
  use distributed_oct_m,    only: distributed_t
  use global_oct_m
  use io_oct_m,             only: io_open, io_close
  use io_function_oct_m
  use messages_oct_m
  use mpi_oct_m,            only: mpi_world, mpi_double_precision
  use mpi_debug_oct_m
  use mpi_distribute_oct_m, only: mpi_displacements
  use namespace_oct_m
  use poisson_oct_m,        only: poisson_t, dpoisson_solve
  use profiling_oct_m
  use states_elec_oct_m,    only: states_elec_t, states_elec_block_min, states_elec_block_max
 
  implicit none
 
  private
  public :: &
    output_matrix, &
    isdf_potential, &
    local_number_of_states, &
    gather_psi_mu_over_states, &
    output_psi_mu_for_all_states
 
  ! TODO(Alex) Issue #1195 Extend ISDF to spin-polarised systems
  integer, private, parameter   :: ik = 1
 
contains
 
  subroutine output_matrix(namespace, fname, matrix)
    type(namespace_t),        intent(in) :: namespace
    character(len=*),         intent(in) :: fname
    real(real64), contiguous, intent(in) :: matrix(:, :)
 
    integer :: i, j, unit
 
    push_sub(output_matrix)
 
    write(message(1),'(a)') "ISDF. Outputting:  " // trim(adjustl(fname))
    call messages_info(1, namespace=namespace, debug_only=.true.)
 
    if (mpi_world%is_root()) then
 
      open(newunit=unit, file=trim(adjustl(fname)))
      do j = 1, size(matrix, 2)
        do i = 1, size(matrix, 1)
          write(unit, *) matrix(i, j)
        enddo
      enddo
      close(unit)
 
    endif
 
    pop_sub(output_matrix)
 
  end subroutine output_matrix
 
  subroutine isdf_potential(namespace, poisson_solver, isdf_vectors, isdf_dist, V_r_nu)
    type(namespace_t),   intent(in )             :: namespace
    type(poisson_t),     intent(in )             :: poisson_solver
    real(real64),        intent(in ), contiguous :: isdf_vectors(:, :)
    type(distributed_t), intent(in)              :: isdf_dist
 
    real(real64),        intent(out), contiguous :: V_r_nu(:, :)
 
    integer                   :: np, i_nu
    integer,      allocatable :: recv_counts(:), displs(:)
    real(real64), allocatable :: V_r_nu_local(:, :)
 
    push_sub_with_profile(isdf_potential)
 
    ! Local buffer
    np = size(v_r_nu, 1)
    safe_allocate(v_r_nu_local(1:np, isdf_dist%start:isdf_dist%end))
 
    do i_nu = isdf_dist%start, isdf_dist%end
      call dpoisson_solve(poisson_solver, namespace, v_r_nu_local(:, i_nu), isdf_vectors(:, i_nu), all_nodes=.false.)
    enddo
 
    ! As an initial step, gather V_r_nu rather than keeping it distributed
    safe_allocate(recv_counts(1:isdf_dist%mpi_grp%size))
    safe_allocate(displs(1:isdf_dist%mpi_grp%size))
 
    ! isdf_dist%num corresponds to number of columns per process
    recv_counts(:) = np * isdf_dist%num(:)
    call mpi_displacements(recv_counts, displs)
 
    call isdf_dist%mpi_grp%allgatherv(v_r_nu_local, size(v_r_nu_local), mpi_double_precision, &
      v_r_nu, recv_counts, displs, mpi_double_precision)
 
    safe_deallocate_a(v_r_nu_local)
 
    pop_sub_with_profile(isdf_potential)
 
  end subroutine isdf_potential
 
 
  integer function local_number_of_states(st, max_state)
    type(states_elec_t), intent(in) :: st
    integer,             intent(in) :: max_state
 
    integer :: minst, maxst
 
    push_sub(local_number_of_states)
 
    minst = states_elec_block_min(st, st%group%block_start)
    maxst = min(states_elec_block_max(st, st%group%block_end), max_state)
 
    ! Handles when max_state is part of a block with index < (block_start of current process)
    ! resulting in no states on this process being used in the ISDF expansion
    local_number_of_states = max(maxst - minst + 1, 0)
 
    pop_sub(local_number_of_states)
 
  end function local_number_of_states
 
 
  subroutine gather_psi_mu_over_states(st, psi_mu, psi_global)
    type(states_elec_t),       intent(in ) :: st
    real(real64), contiguous,  intent(in ) :: psi_mu(:, :)
    real(real64), contiguous,  intent(out) :: psi_global(:, :)
 
    integer :: max_state, n_int_g, ic, ist, st_end, ist_local
 
    push_sub(gather_psi_mu_over_states)
 
    if (st%group%psib(st%group%block_start, ik)%status() /= batch_packed) then
      message(1) = "Developer Error: Trying to output psi_mu when not BATCH_PACKED"
      call messages_fatal(1)
    endif
 
    ! Total number of interpolation points
    n_int_g = size(psi_global, 2)
    assert(size(psi_mu, 2) == size(psi_global, 2))
 
    ! Total number of states used in ISDF
    max_state = size(psi_global, 1)
 
    ! All elements must be zeroed, such that allreduce does not
    ! sum contributions from uninitialised elements
    do ic = 1, n_int_g
      do ist = 1, max_state
        psi_global(ist, ic) = 0.0_real64
      enddo
    enddo
 
    ! Ensure states does not iterate beyond the max state used in
    ! ISDF expansion
    st_end = min(st%st_end, max_state)
 
    ! Sanity check on number of local states held by psi_mu
    if (size(psi_mu, 1) > 0) then
      assert(st_end - st%st_start + 1 == size(psi_mu, 1))
    endif
 
    ! Fill a section of psi_global using psi_mu
    do ic = 1, n_int_g
      do ist = st%st_start, st_end
        ist_local = ist - st%st_start + 1
        psi_global(ist, ic) = psi_mu(ist_local, ic)
      enddo
    enddo
 
    call comm_allreduce(st%mpi_grp, psi_global)
 
    pop_sub(gather_psi_mu_over_states)
 
  end subroutine gather_psi_mu_over_states
 
 
  subroutine output_psi_mu_for_all_states(namespace, st, max_state, psi_mu)
    type(namespace_t),         intent(in) :: namespace
    type(states_elec_t),       intent(in) :: st
    integer,                   intent(in) :: max_state
    real(real64), contiguous,  intent(in) :: psi_mu(:, :)
 
    real(real64),       allocatable :: psi_global(:, :)
    integer                         :: n_int_g, ic, ist, unit
    character(len=2)                :: np_char
 
    push_sub(output_psi_mu_for_all_states)
 
    if (st%group%psib(st%group%block_start, ik)%status() /= batch_packed) then
      message(1) = "Trying to output psi_mu when not BATCH_PACKED"
      call messages_fatal(1, namespace=namespace)
    endif
 
    write(message(1),'(a)') "ISDF: Writing psi_mu (all states/DD)"
    call messages_info(1)
 
    n_int_g = size(psi_mu, 2)
    safe_allocate(psi_global(1:max_state, 1:n_int_g))
 
    call gather_psi_mu_over_states(st, psi_mu, psi_global)
 
    write(np_char, '(I2)') st%mpi_grp%size
    unit = io_open("global_psi_mu_np"//trim(adjustl(np_char))//".txt", namespace, action='write')
 
    do ic = 1, n_int_g
      do ist = 1, max_state
        write(unit, *) psi_global(ist, ic)
      enddo
    enddo
 
    call io_close(unit)
 
    safe_deallocate_a(psi_global)
 
    pop_sub(output_psi_mu_for_all_states)
 
  end subroutine output_psi_mu_for_all_states
 
 
end module isdf_utils_oct_m
 
!! Local Variables:
!! mode: f90
!! coding: utf-8
!! End: