subroutine write_chkpt(common_data, label, istdout, istderr, ierr)
use w90_comms, only: comms_reduce, mpirank
use w90_error_base, only: w90_error_type
use w90_error, only: set_error_alloc, set_error_dealloc, set_error_fatal
use w90_wannier90_readwrite, only: w90_wannier90_readwrite_write_chkpt
implicit none
! arguments
character(len=*), intent(in) :: label ! e.g. 'postdis' or 'postwann' after disentanglement, wannierisation
integer, intent(in) :: istdout, istderr
integer, intent(inout) :: ierr
type(lib_common_type), target, intent(in) :: common_data
! local variables
complex(kind=dp), allocatable :: m(:, :, :, :)
integer, allocatable :: global_k(:)
integer, pointer :: nw, nb, nk, nn
integer :: rank, nkrank, ikg, ikl, istat
type(w90_error_type), allocatable :: error
ierr = 0
rank = mpirank(common_data%comm)
nkrank = count(common_data%dist_kpoints == rank)
nb => common_data%num_bands
nk => common_data%num_kpts
nn => common_data%kmesh_info%nntot
nw => common_data%num_wann
if (.not. associated(common_data%u_matrix_opt)) then
call set_error_fatal(error, &
'Error: u_matrix_opt not associated for write_chkpt call', common_data%comm)
else if (.not. associated(common_data%u_matrix)) then
call set_error_fatal(error, &
'Error: u_matrix not associated for write_chkpt call', common_data%comm)
else if (.not. associated(common_data%m_matrix_local)) then
call set_error_fatal(error, &
'Error: m_matrix_local not set for write_chkpt call', common_data%comm)
end if
if (allocated(error)) then
call prterr(error, ierr, istdout, istderr, common_data%comm)
return
end if
allocate (global_k(nkrank), stat=istat)
if (istat /= 0) then
call set_error_alloc(error, 'Error allocating global_k in write_chkpt', common_data%comm)
call prterr(error, ierr, istdout, istderr, common_data%comm)
return
end if
global_k = huge(1)
ikl = 1
do ikg = 1, nk
if (rank == common_data%dist_kpoints(ikg)) then
global_k(ikl) = ikg
ikl = ikl + 1
end if
end do
! reassemble full m matrix by MPI reduction
!
! allocating and partially assigning the full matrix on all ranks and reducing is a terrible idea
! alternatively, allocate on root and use point-to-point
! or, if required only for checkpoint file writing, then use mpi-io (but needs to be ordered io, alas)
! or, even better, use parallel hdf5. JJ Nov 22
allocate (m(nw, nw, nn, nk), stat=istat) ! all kpts
if (istat /= 0) call set_error_alloc(error, 'Error allocating m in write_chkpt', common_data%comm)
if (allocated(error)) then
call prterr(error, ierr, istdout, istderr, common_data%comm)
return
end if
m(:, :, :, :) = 0.d0
do ikl = 1, nkrank
ikg = global_k(ikl)
m(:, :, :, ikg) = common_data%m_matrix_local(1:nw, 1:nw, :, ikl)
end do
call comms_reduce(m(1, 1, 1, 1), nw*nw*nn*nk, 'SUM', error, common_data%comm)
if (allocated(error)) then
call prterr(error, ierr, istdout, istderr, common_data%comm)
return
end if
if (rank == 0) then
call w90_wannier90_readwrite_write_chkpt(label, common_data%exclude_bands, &
common_data%wannier_data, common_data%kmesh_info, &
common_data%kpt_latt, nk, common_data%dis_manifold, &
nb, nw, common_data%u_matrix, &
common_data%u_matrix_opt, m, common_data%mp_grid, &
common_data%real_lattice, &
common_data%omega%invariant, &
common_data%have_disentangled, &
common_data%print_output%iprint, istdout, &
common_data%seedname)
end if
deallocate (m, stat=istat)
if (istat /= 0) then
call set_error_dealloc(error, 'Error deallocating m in write_chkpt', common_data%comm)
call prterr(error, ierr, istdout, istderr, common_data%comm)
return
end if
end subroutine write_chkpt