subroutine setup_m_loc(kmesh_info, print_output, m_matrix_local, m_matrix_orig_local, u_matrix, &
num_bands, num_kpts, num_wann, timer, dist_k, error, comm)
!================================================!
!
! map m_matrix_orig_local to m_matrix_local
! at entry, m_matrix_local is not allocated
!
!================================================!
use w90_comms, only: w90_comm_type, mpirank
use w90_constants, only: dp, cmplx_0, cmplx_1
use w90_error
use w90_io, only: io_stopwatch_start, io_stopwatch_stop
use w90_types, only: kmesh_info_type, print_output_type, timer_list_type
! arguments
integer, intent(in) :: num_bands, num_kpts, num_wann
integer, intent(in) :: dist_k(:)
complex(kind=dp), intent(in) :: u_matrix(:, :, :) ! (num_wann, num_wann, num_kpts) -- full array duplicated on all ranks
complex(kind=dp), intent(in) :: m_matrix_orig_local(:, :, :, :) ! (num_bands, num_bands, nntot, num_kpts) -- only local kpts
complex(kind=dp), intent(inout) :: m_matrix_local(:, :, :, :) ! (num_wann, num_wann, nntot, rank_kpts) -- only local kpts
type(kmesh_info_type), intent(in) :: kmesh_info
type(print_output_type), intent(in) :: print_output
type(w90_comm_type), intent(in) :: comm
type(timer_list_type), intent(inout) :: timer
type(w90_error_type), allocatable, intent(out) :: error
! internal variables
complex(kind=dp), allocatable :: cwb(:, :), cww(:, :)
integer :: nkp, nkp2, nn, ierr, nkp_global, nkrank
integer, allocatable :: global_k(:)
integer :: ikg, ikl, my_node_id
if (print_output%timing_level > 1) call io_stopwatch_start('dis: setup_m_loc', timer)
! local-global k index mapping
my_node_id = mpirank(comm)
nkrank = count(dist_k == my_node_id)
allocate (global_k(nkrank), stat=ierr)
if (ierr /= 0) then
call set_error_alloc(error, 'Error in allocating global_k in setup_m_loc', comm)
return
end if
global_k = huge(1); ikl = 1
do ikg = 1, num_kpts
if (dist_k(ikg) == my_node_id) then
global_k(ikl) = ikg ! global [1,num_kpts] index corresponding to local [1,nk_this_node] index
ikl = ikl + 1
end if
end do
allocate (cwb(num_wann, num_bands), stat=ierr)
if (ierr /= 0) then
call set_error_alloc(error, 'Error in allocating cwb in setup_m_loc', comm)
return
end if
allocate (cww(num_wann, num_wann), stat=ierr)
if (ierr /= 0) then
call set_error_alloc(error, 'Error in allocating cww in setup_m_loc', comm)
return
end if
do nkp = 1, nkrank
nkp_global = global_k(nkp)
do nn = 1, kmesh_info%nntot
nkp2 = kmesh_info%nnlist(nkp_global, nn)
call zgemm('C', 'N', num_wann, num_wann, num_wann, cmplx_1, u_matrix(:, :, nkp_global), &
num_wann, m_matrix_orig_local(:, :, nn, nkp), num_bands, cmplx_0, cwb, num_wann)
call zgemm('N', 'N', num_wann, num_wann, num_wann, cmplx_1, cwb, num_wann, &
u_matrix(:, :, nkp2), num_wann, cmplx_0, cww, num_wann)
m_matrix_local(1:num_wann, 1:num_wann, nn, nkp) = cww(:, :)
end do
end do
deallocate (cwb, stat=ierr)
if (ierr /= 0) then
call set_error_dealloc(error, 'Error in deallocating cwb in setup_m_loc', comm)
return
end if
deallocate (cww, stat=ierr)
if (ierr /= 0) then
call set_error_dealloc(error, 'Error in deallocating cww in setup_m_loc', comm)
return
end if
deallocate (global_k, stat=ierr)
if (ierr /= 0) then
call set_error_dealloc(error, 'Error in deallocating global_k in setup_m_loc', comm)
return
end if
if (print_output%timing_level > 1) call io_stopwatch_stop('dis: setup_m_loc', timer)
return
!================================================!
end subroutine setup_m_loc