dm_propagation.F90

Go to the documentation of this file.

!! Copyright (C) 2024 - 2025 S. Pal, Z. Nie, U. De Giovannini
!! 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 dm_propagation_oct_m
  use batch_ops_oct_m
  use blas_oct_m
  use debug_oct_m
  use density_oct_m
  use distributed_oct_m
  use eigensolver_oct_m
  use electron_space_oct_m
  use global_oct_m
  use grid_oct_m
  use hamiltonian_elec_oct_m
  use interaction_partner_oct_m
  use ions_oct_m
  use, intrinsic :: iso_fortran_env
  use kpoints_oct_m
  use lalg_adv_oct_m
  use lalg_basic_oct_m
  use mesh_batch_oct_m
  use mesh_function_oct_m
  use messages_oct_m
  use mpi_distribute_oct_m
  use mpi_lib_oct_m
  use mpi_oct_m
  use multicomm_oct_m
  use namespace_oct_m
  use parser_oct_m
  use profiling_oct_m
  use restart_oct_m
  use space_oct_m
  use states_elec_calc_oct_m
  use states_elec_oct_m
  use states_elec_restart_oct_m
  use td_write_low_oct_m
  use unit_oct_m
  use unit_system_oct_m
  use v_ks_oct_m
  use varinfo_oct_m
  use write_iter_oct_m
 
  implicit none
 
  private
  public ::                             &
    dmp_t,                              &
    dm_propagation_init_run,            &
    dm_propagation_run,                 &
    dm_end_run
 
  type dmp_t
    integer                             :: calculation_mode
    integer                             :: basis
    logical                             :: unitary_transform
    type(states_elec_t)                 :: adiabatic_st
    integer                             :: strategy
    logical                             :: othn
    type(restart_t)                     :: restart_dump
    integer                             :: istart, iend
    complex(real64), allocatable        :: rho_offd(:, :)
    ! 2-Times Model
    real(real64)                        :: tmodel(2)
    real(real64), allocatable           :: occ_gs(:, :)
    ! Uniform Decay
    real(real64)                        :: uniform(2)
    ! EPW Data Parsing and Shared Memory Configuration
    character(len=256)                  :: epw_file
    integer                             :: iunit
    integer                             :: estart, eend, wnst
    integer                             :: ia
    integer                             :: na(3)
    real(real64)                        :: astep(3)
    integer(int64)                      :: num
    integer, allocatable                :: kmap(:)
    type(mpi_grp_t)                     :: intranode_grp, internode_grp
#ifdef HAVE_MPI
    type(MPI_Win)                       :: window_trans_rate
#endif
  contains
    procedure :: init => dmp_init
    procedure :: update_trans_rate => dm_propagation_update_trans_rate
    procedure :: write_dm_proj_basis => td_write_dm_proj_basis
  end type dmp_t
 
  ! Defined separately because the VOLATILE attribute is not allowed for derived-type components.
  real(real32), pointer, volatile :: ave_trans(:)
 
contains
 
  subroutine dmp_init(this, namespace, st, space, hm)
    class(dmp_t),             intent(inout) :: this
    type(namespace_t),        intent(in)    :: namespace
    type(states_elec_t),      intent(in)    :: st
    class(space_t),           intent(in)    :: space
    type(hamiltonian_elec_t), intent(in)    :: hm
 
    type(block_t)       :: blk
    integer             :: ncols, nempty
    real(real64)        :: nempty_percent
 
    push_sub(dmp_init)
 
    !%Variable TDDMPropagationBasis
    !%Type integer
    !%Default Adiabatic
    !%Section Time-Dependent
    !%Description
    !% Decides the basis set for the density matrix propagation.
    !%Option Adiabatic 01
    !% Instantaneous eigenstates of the Hamiltonian.
    !%Option Groundstate 02
    !% Eigenstates of the Hamiltonian at t=0.
    !%End
    call parse_variable(namespace, 'TDDMPropagationBasis', option__tddmpropagationbasis__adiabatic, this%basis)
    if (.not. varinfo_valid_option('TDDMPropagationBasis', this%basis)) then
      call messages_input_error(namespace, 'TDDMPropagationBasis')
    endif
    call messages_print_var_option('TDDMPropagationBasis', this%basis, namespace=namespace)
 
    !%Variable TDDMOrthogonal
    !%Type logical
    !%Default no
    !%Section Time-Dependent
    !%Description
    !% Use a fully orthonormalized basis when constructing and
    !% damping the density matrix.
    !%End
    call parse_variable(namespace, 'TDDMOrthogonal', .true., this%othn)
    call messages_print_var_value('TDDMOrthogonal', this%othn, namespace=namespace)
 
    !%Variable TDDMUnitaryTransformFix
    !%Type logical
    !%Default yes
    !%Section Time-Dependent
    !%Description
    !% Applies an additional unitary transformation to the damped wavefunctions
    !% to maximize their overlap with the reference wavefunctions after damping.
    !%End
    call parse_variable(namespace, 'TDDMUnitaryTransformFix', .true., this%unitary_transform)
    call messages_print_var_value('TDDMUnitaryTransformFix', this%unitary_transform, namespace=namespace)
 
    !%Variable TDDMBandLowest
    !%Type integer
    !%Default 1
    !%Section Time-Dependent
    !%Description
    !% The starting DFT band index used for density matrix propagation.
    !%End
    call parse_variable(namespace, 'TDDMBandLowest', 1, this%istart)
    if (this%istart < 1) then
      message(1) = "Input: TDDMBandLowest must be positive."
      call messages_fatal(1, namespace=namespace)
    end if
 
    !%Variable TDDMBandHighest
    !%Type integer
    !%Default -1
    !%Section Time-Dependent
    !%Description
    !% The ending DFT band index used for density matrix propagation.
    !% Defaults to the total number of states.
    !%End
    call parse_variable(namespace, 'TDDMBandHighest', st%nst, this%iend)
    if (this%iend < this%istart .or. this%iend > st%nst) then
      message(1) = "Input: TDDMBandHighest must be within [TDDMBandLowest, total states]."
      call messages_fatal(1, namespace=namespace)
    end if
    write(message(1), '(a, i0, a, i0)') 'Input: Propagating bands ', this%istart, ' - ', this%iend
    call messages_info(1, namespace=namespace)
 
    this%strategy = -1
 
    !%Variable TDDMPropagation_uniform_decay
    !%Type block
    !%Section Time-Dependent
    !%Description
    !% The intra-k transition rates between any pair of states are taken to be constant.
    !% The population dynamics are therefore determined solely by the state occupations
    !% and the bath temperature, following Fermi’s golden rule.
    !% The first column of the block specifies the characteristic lifetime, and the
    !% second column gives the bath temperature (in Kelvin).
    !% <tt>%TDDMPropagation_uniform_decay
    !% <br>&nbsp;&nbsp; Time | Temp
    !% <br>%</tt>
    !%End
    if (parse_block(namespace, 'TDDMPropagation_uniform_decay', blk) == 0) then
      if (this%strategy /= -1) then
        message(1) = "Multiple dissipation strategies are not allowed."
        call messages_fatal(1, namespace=namespace)
      end if
      this%strategy = 1
      ncols = parse_block_cols(blk, 0)
      if (ncols == 2) then
        call parse_block_float(blk, 0, 0, this%uniform(1), units_inp%time)
        call parse_block_float(blk, 0, 1, this%uniform(2))
        call parse_block_end(blk)
 
        message(1) = "Info: TDDMPropagation uniform decay approximation:"
        write(message(2),'(a, f0.2, 1x, 2a, F0.2, 1x, 2a)') '      [lifetime, temperature] = [', &
          this%uniform(1), trim(units_abbrev(units_out%time)), ', ', this%uniform(2), trim(units_abbrev(unit_kelvin)), ']'
        call messages_info(2, namespace=namespace)
      else
        message(1) = "Input: TDDMPropagation_uniform_decay block must have 2 columns."
        call messages_fatal(1, namespace=namespace)
      end if
    end if
 
    !%Variable TDDMPropagation_2Times
    !%Type block
    !%Section Time-Dependent
    !%Description
    !% Two times approximation for the jump operator in the master equation.
    !% S. A. Sato <i>et al.</i>, <i>Phys. Rev. B</i> 99, 214302 (2019).
    !%
    !% <tt>%TDDMPropagation_2Times
    !% <br>&nbsp;&nbsp; t1 | t2
    !% <br>%</tt>
    !%End
    if (parse_block(namespace, 'TDDMPropagation_2Times', blk) == 0) then
      if (this%strategy /= -1) then
        message(1) = "Multiple dissipation strategies are not allowed."
        call messages_fatal(1, namespace=namespace)
      end if
      this%strategy = 2
      ncols = parse_block_cols(blk, 0)
      if (ncols == 2) then
        call parse_block_float(blk, 0, 0, this%tmodel(1), units_inp%time)
        call parse_block_float(blk, 0, 1, this%tmodel(2), units_inp%time)
        call parse_block_end(blk)
 
        message(1) = "Info: TDDMPropagation 2-times approximation:"
        write(message(2),'(a, f12.6, a, f12.6, a,a)') &
          '      [t1, t2] = [', this%tmodel(1) , ', ', this%tmodel(2), '] ', &
          trim(units_abbrev(units_out%time))
        call messages_info(2, namespace=namespace)
      else
        message(1) = "Input: TDDMPropagation_2Times block must have 2 columns."
        call messages_fatal(1, namespace=namespace)
      end if
      ! for two times model, only orthonormalized basis is allowed
      if (.not. this%othn) then
        this%othn = .true.
        message(1) = "Overriding input: TDDMOrthogonal set to yes for TDDMPropagation_2Times."
        call messages_warning(1, namespace=namespace)
      end if
    end if
 
    !%Variable TDDMPropagation_from_epw
    !%Type string
    !%Default "-"
    !%Section Time-Dependent
    !%Description
    !% Specifies the transition-rate file generated by EPW. Once loaded, the
    !% simulation includes all intra-k and inter-k scattering processes, with
    !% the electron dynamics governed by Fermi’s Golden Rule. This option
    !% requires EPWBandLowest to be set.
    !%End
    call parse_variable(namespace, 'TDDMPropagation_from_epw', '-', this%epw_file)
    if (trim(this%epw_file) /= '-') then
      if (bitand(hm%kpoints%method, kpoints_monkh_pack) == 0) then
        write(message(1),'(a)') 'Only Monkhorst-Pack k-point meshes are supported in the EPW-DMPropagation strategy.'
        call messages_fatal(1, namespace = namespace)
      end if
      if (hm%kpoints%reduced%nshifts /= 1) then
        write(message(1),'(a)') 'Multiple Monkhorst-Pack shifts are not supported in the EPW-DMPropagation strategy.'
        call messages_fatal(1, namespace = namespace)
      end if
      if (space%dim /= 3) then
        write(message(1),'(a)') 'Only 3D systems are supported in the EPW-DMPropagation strategy.'
        call messages_fatal(1, namespace = namespace)
      end if
 
      if (this%strategy /= -1) then
        message(1) = "Multiple dissipation strategies are not allowed."
        call messages_fatal(1, namespace=namespace)
      end if
      this%strategy = 3
      write(message(1),'(a, a)') "Info: TDDMPropagation transition rates from file: ", trim(this%epw_file)
      call messages_info(1, namespace=namespace)
    end if
 
    !%Variable EPWBandLowest
    !%Type integer
    !%Default -1
    !%Section Time-Dependent
    !%Description
    !% The starting DFT band index used for damping.
    !%
    !% Since EPW transition rates are calculated for a subset of DFT bands,
    !% this variable tells the code which global DFT band corresponds
    !% to damping band #1.
    !%
    !% Example: If you damp DFT bands 10 through 15, set EPWBandLowest = 10.
    !%End
    if (this%strategy == 3) then
      call parse_variable(namespace, 'EPWBandLowest', 0, this%estart)
      if (this%estart < this%istart .or. this%estart > this%iend) then
        write(message(1), '(a)') 'Input: EPWBandLowest must be within [TDDMBandLowest, TDDMBandHighest].'
        call messages_fatal(1, namespace=namespace)
      end if
    end if
 
    ! Sanity checks
    if (this%calculation_mode == option__tddmpropagation__collision_integral .and. this%strategy /= 3) then
      message(1) = "Warning: TDDMPropagation_from_epw is required for collision integral."
      message(2) = "         Add TDDMPropagation_from_epw and rerun."
      call messages_fatal(2, namespace=namespace)
    end if
 
    call parse_variable(namespace, 'ExtraStates', 0, nempty)
    call parse_variable(namespace, 'ExtraStatesInPercent', m_zero, nempty_percent)
    if (nempty == 0 .and. nempty_percent < m_epsilon) then
      message(1) = "Warning: TDDMPropagation requires a number of empty states."
      message(2) = "         Add ExtraStates and rerun."
      call messages_fatal(2, namespace=namespace)
    end if
 
    if (st%parallel_in_states) then
      message(1) = "Warning: TDDMPropagation does not support parallel states."
      message(2) = "         Remove ParallelStates and rerun."
      call messages_fatal(2, namespace=namespace)
    end if
 
    if (st%d%ispin == spin_polarized .and. this%strategy == 3) then
      call messages_not_implemented('Spin-polarized TDDMPropagation with EPW', namespace=namespace)
    end if
 
    pop_sub(dmp_init)
 
  end subroutine dmp_init
 
  subroutine dm_propagation_init_run(dmp, namespace, space, gr, ions, st, hm, mc, from_scratch)
    type(dmp_t),              intent(inout)   :: dmp
    type(namespace_t),        intent(in)      :: namespace
    type(electron_space_t),   intent(in)      :: space
    type(grid_t),             intent(in)      :: gr
    type(ions_t),             intent(in)      :: ions
    type(states_elec_t),      intent(in)      :: st
    type(hamiltonian_elec_t), intent(in)      :: hm
    type(multicomm_t),        intent(in)      :: mc
    logical,                  intent(in)      :: from_scratch
 
    integer               :: ierr, act_nst
    type(restart_t)       :: restart_load
 
    push_sub(dm_propagation_init_run)
 
    ! By default, wavefunctions are copied. In principle, exclude_eigenval should be true.
    ! But we need its allocation
    call states_elec_copy(dmp%adiabatic_st, st, special=.true.)
 
    if (from_scratch .or. dmp%basis == option__tddmpropagationbasis__groundstate) then
      call restart_load%init(namespace, restart_gs, restart_type_load, mc, ierr, mesh=gr, exact=.true.)
    else
      call restart_load%init(namespace, restart_dm, restart_type_load, mc, ierr, mesh=gr, exact=.true.)
    end if
 
    if (ierr == 0) then
      call states_elec_load(restart_load, namespace, space, dmp%adiabatic_st, gr, hm%kpoints, fixed_occ=st%restart_fixed_occ, &
        ierr=ierr, label = ": DM-basis")
    else
      message(1) = 'Unable to read DM-basis wavefunctions.'
      call messages_fatal(1, namespace=namespace)
    end if
 
    call restart_load%end()
 
    ! TDDMPropagation_2Times
    if (dmp%strategy == 2) then
      safe_allocate(dmp%occ_gs(1:st%nst, 1:st%nik))
      dmp%occ_gs = dmp%adiabatic_st%occ
    end if
 
    ! dmp%ave_trans: replicated on each node, shared within the node
    if (dmp%strategy == 3) call iopar_open_trans_rate(namespace, ions, hm, st%system_grp, dmp)
 
    ! record only adiabatic states
    if (dmp%basis == option__tddmpropagationbasis__adiabatic) then
      call dmp%restart_dump%init(namespace, restart_dm, restart_type_dump, mc, ierr, mesh=gr)
    end if
 
    ! record off-diagonals
    act_nst = dmp%iend - dmp%istart + 1
    safe_allocate(dmp%rho_offd(1:act_nst*(act_nst-1)/2, st%d%kpt%start:st%d%kpt%end))
 
    pop_sub(dm_propagation_init_run)
  end subroutine dm_propagation_init_run
 
  subroutine dm_end_run(system_grp, dmp)
    type(mpi_grp_t),  intent(in)     :: system_grp
    type(dmp_t),      intent(inout)  :: dmp
 
    push_sub(dm_end_run)
 
    if (dmp%basis == option__tddmpropagationbasis__adiabatic) then
      call dmp%restart_dump%end()
    end if
 
    call states_elec_end(dmp%adiabatic_st)
 
    if (dmp%strategy == 3) call iopar_close_trans_rate(system_grp, dmp)
 
    safe_deallocate_a(dmp%occ_gs)
    safe_deallocate_a(dmp%rho_offd)
 
    pop_sub(dm_end_run)
  end subroutine dm_end_run
 
  subroutine dm_propagation_run(dmp, namespace, space, gr, ions, st, mc, hm, ks, iter, dt, ext_partners, update_energy)
    type(dmp_t),              intent(inout)   :: dmp
    type(namespace_t),        intent(in)      :: namespace
    type(electron_space_t),   intent(in)      :: space
    type(grid_t),             intent(in)      :: gr
    type(ions_t),             intent(in)      :: ions
    type(states_elec_t),      intent(inout)   :: st
    type(multicomm_t),        intent(in)      :: mc
    type(hamiltonian_elec_t), intent(inout)   :: hm
    type(v_ks_t),             intent(inout)   :: ks
    integer,                  intent(in)      :: iter
    real(real64),             intent(in)      :: dt
    type(partner_list_t),     intent(in)      :: ext_partners
    logical, optional,        intent(in)      :: update_energy
 
    real(real64), parameter       :: zero = 1.0e-15_real64
    type(eigensolver_t)           :: eigens
    real(real64)                  :: nrm2_tdks(st%nst, st%nik), population(3), leak
    integer                       :: ik, act_nst
    logical                       :: update_energy_
    type(states_elec_t)           :: resd_st
    complex(real64), allocatable  :: overlap_ad_ks(:, :, :), overlap_resd_ks(:, :, :), rho_mat_k(:, :, :)
    integer, allocatable          :: nresd_k(:)
 
    push_sub_with_profile(dm_propagation_run)
 
    ! hm is updated after each td propagation.
    if (dmp%basis == option__tddmpropagationbasis__adiabatic) then
      call eigensolver_init(eigens, namespace, gr, dmp%adiabatic_st, hm, mc, space)
      eigens%converged = 0
      call eigens%run(namespace, gr, dmp%adiabatic_st, hm, space, ext_partners, iter)
    end if
 
    act_nst = dmp%iend - dmp%istart + 1
 
    safe_allocate(overlap_ad_ks(1:st%nst, 1:st%nst, st%d%kpt%start:st%d%kpt%end))
    safe_allocate(overlap_resd_ks(1:st%nst, 1:st%nst, st%d%kpt%start:st%d%kpt%end))
    safe_allocate(rho_mat_k(2*act_nst, 2*act_nst, st%d%kpt%start:st%d%kpt%end))
    safe_allocate(nresd_k(st%d%kpt%start:st%d%kpt%end))
 
    population = 0.0_real64
 
    ! Create a copy of the TDKS states to initialise the residuals
    call states_elec_copy(resd_st, st, exclude_eigenval = .true.)
 
    ! Store all density matrices for potential future communication
    do ik = st%d%kpt%start, st%d%kpt%end
      ! pop = \sum_{i,k} f_{i,k} <\psi_{i,k} | \psi_{i,k}>
      call total_population(ik, st, gr, nrm2_tdks(:, ik), population(1))
      ! S^{\phi,\psi}_{ij} = <\phi_i | \psi_j>
      call calc_overlap_state(ik, gr, dmp%adiabatic_st, st, dmp%istart, dmp%iend, dmp%istart, dmp%iend, overlap_ad_ks(:, :, ik))
      ! population before damping, set occupation as well
      call population_in_adiabatic(ik, dmp%adiabatic_st, st, overlap_ad_ks(:, :, ik), dmp%istart, dmp%iend, population(2))
 
      ! residual |d_j> = |\psi_j> - \sum_i S^{\phi,\psi}_{ij} |\phi_i>
      ! S^{d,\psi}_{ij} = <d_i | \psi_j>
      call construct_residuals(gr, namespace, dmp%adiabatic_st, st, ik, dmp%othn, dmp%istart, dmp%iend, &
        overlap_ad_ks(:, :, ik), nrm2_tdks(:, ik), nresd_k(ik), overlap_resd_ks(:, :, ik), resd_st)
 
      ! Active block: istart:iend, nst+istart:nst+istart+nresd-1, allocated shape: (2*nst, 2*nst, :)
      call construct_density_matrix(nresd_k(ik), ik, st, dmp%istart, dmp%iend, overlap_ad_ks(:, :, ik), &
        overlap_resd_ks(:, :, ik), rho_mat_k(:, :, ik), dmp%rho_offd(:, ik))
    end do
 
    ! broadcast the instanteous occupation for Pauli blocking
    call broadcast_occupation(dmp%adiabatic_st%occ, dmp%adiabatic_st%d%kpt, st%parallel_in_states, dmp%istart, dmp%iend)
 
    ! dissipate the density matrix
    call dissipation(hm, st, namespace, nresd_k, dt, dmp, rho_mat_k)
 
    do ik = st%d%kpt%start, st%d%kpt%end
      ! reconstruct wavefunctions. Active block: (nresd+nst, nresd+nst, :), allocated shape: (2*nst, 2*nst, :)
      call update_st(dmp, ik, gr, namespace, nresd_k(ik), overlap_ad_ks(:, :, ik), overlap_resd_ks(:, :, ik), nrm2_tdks(:, ik), &
        resd_st, st, rho_mat_k(:, :, ik), population(3))
    end do
 
    call states_elec_end(resd_st)
 
    safe_deallocate_a(overlap_ad_ks)
    safe_deallocate_a(overlap_resd_ks)
 
    call broadcast_occupation(st%occ, st%d%kpt, st%parallel_in_states, dmp%istart, dmp%iend)
 
    safe_deallocate_a(rho_mat_k)
    safe_deallocate_a(nresd_k)
 
    call st%d%kpt%mpi_grp%allreduce_inplace(population, 3, mpi_double_precision, mpi_sum)
 
    write(message(1), '(a,E21.14,a,E21.14,a,E21.14,a)') 'DMPopulation: ', population(1), &
      ' (', population(2), ' in basis ) before damping; ', population(3), ' after damping'
    call messages_info(1)
 
    leak = population(1) - population(3)
    if (abs(leak) > zero) then
      write(message(1), '(a,E15.8,a,I8)') 'Leakage: ', leak, '(after damping) at time step', iter
      call messages_info(1)
    end if
    dmp%adiabatic_st%qtot = population(1)
 
    ! renew hamiltonian
    call density_calc(st, gr, st%rho)
    update_energy_ = optional_default(update_energy, .true.)
    call v_ks_calc(ks, namespace, space, hm, st, ions, ext_partners, calc_eigenval = update_energy_, &
      time = abs(iter * dt), calc_energy = update_energy_)
 
    if (dmp%basis == option__tddmpropagationbasis__adiabatic) then
      call eigensolver_end(eigens)
    end if
 
    pop_sub_with_profile(dm_propagation_run)
 
  end subroutine dm_propagation_run
 
  subroutine total_population(ik, st, gr, nrm2, pop)
    integer,                  intent(in)    :: ik
    type(states_elec_t),      intent(in)    :: st
    type(grid_t),             intent(in)    :: gr
    real(real64), contiguous, intent(out)   :: nrm2(:)
    real(real64),             intent(inout) :: pop
 
    integer                           :: ib, i_minst, i_maxst
 
    push_sub(total_population)
 
    assert(pop >= 0.0_real64)
    assert(size(nrm2) == st%nst)
 
    do ib = st%group%block_start, st%group%block_end
      i_minst = states_elec_block_min(st, ib)
      i_maxst = states_elec_block_max(st, ib)
      call mesh_batch_nrm2(gr, st%group%psib(ib, ik), nrm2(i_minst:i_maxst), reduce = .false.)
    end do
    nrm2(1:st%nst) = nrm2(1:st%nst)**2
    if (gr%parallel_in_domains) then
      call gr%allreduce(nrm2, dim = st%nst)
    end if
 
    pop = pop + sum(st%occ(1:st%nst, ik) * nrm2(1:st%nst)) * st%kweights(ik)
 
    pop_sub(total_population)
 
  end subroutine total_population
 
  subroutine calc_overlap_single(ik, gr, st, psi, istart, iend, overlap, reduce)
    integer,                         intent(in)      :: ik
    type(grid_t),                    intent(in)      :: gr
    type(states_elec_t),             intent(in)      :: st
    complex(real64), contiguous,     intent(in)      :: psi(:,:)
    integer,                         intent(in)      :: istart
    integer,                         intent(in)      :: iend
    complex(real64),                 intent(out)     :: overlap(:)
    logical, optional,               intent(in)      :: reduce
 
    integer                          :: ib, i_minst, i_maxst, ist
    complex(real64), allocatable     :: phi(:, :)
 
    push_sub_with_profile(calc_overlap_single)
 
    assert(iend >= istart)
    assert(ubound(overlap, dim = 1) == st%nst)
 
    safe_allocate(phi(1:gr%np, st%d%dim))
 
    overlap = m_z0
 
    do ib = st%group%block_start, st%group%block_end
      i_minst = states_elec_block_min(st, ib)
      i_maxst = states_elec_block_max(st, ib)
      if (i_minst >= istart .and. i_maxst <= iend) then
        call zmesh_batch_mf_dotp(gr, st%group%psib(ib, ik), psi, overlap(i_minst:i_maxst), reduce=.false., &
          nst=i_maxst-i_minst+1)
      else
        do ist = max(i_minst, istart), min(i_maxst, iend)
          call states_elec_get_state(st, gr, ist, ik, phi)
          overlap(ist) = zmf_dotp(gr, st%d%dim, phi, psi, reduce = .false.)
        end do
      end if
    end do
 
    safe_deallocate_a(phi)
 
    if (optional_default(reduce, .false.)) then
      call gr%allreduce(overlap, dim = st%nst)
    end if
 
    pop_sub_with_profile(calc_overlap_single)
  end subroutine calc_overlap_single
 
  subroutine calc_overlap_state(ik, gr, ad_st, st, istart_ad, iend_ad, istart_st, iend_st, overlap)
    integer,                         intent(in)      :: ik
    type(grid_t),                    intent(in)      :: gr
    type(states_elec_t),             intent(in)      :: ad_st
    type(states_elec_t),             intent(in)      :: st
    integer,                         intent(in)      :: istart_ad
    integer,                         intent(in)      :: iend_ad
    integer,                         intent(in)      :: istart_st
    integer,                         intent(in)      :: iend_st
    complex(real64),                 intent(out)     :: overlap(:, :)
 
    integer                           :: ist, jst, ib, jb, i_minst, i_maxst, j_minst, j_maxst
    complex(real64), allocatable      :: psi(:, :), dot(:)
 
    push_sub_with_profile(calc_overlap_state)
 
    assert(istart_ad <= iend_ad)
    assert(istart_st <= iend_st)
    assert(ubound(overlap, dim = 1) == ad_st%nst)
    assert(ubound(overlap, dim = 2) == st%nst)
 
    safe_allocate(psi(1:gr%np, st%d%dim))
    safe_allocate(dot(1:st%nst))
 
    overlap = m_z0
 
    do jb = st%group%block_start, st%group%block_end
      j_minst = states_elec_block_min(st, jb)
      j_maxst = states_elec_block_max(st, jb)
 
      ! full st batch
      if (j_minst >= istart_st .and. j_maxst <= iend_st) then
        do ib = ad_st%group%block_start, ad_st%group%block_end
          i_minst = states_elec_block_min(ad_st, ib)
          i_maxst = states_elec_block_max(ad_st, ib)
 
          ! full ad_st batch, full st batch
          if (i_minst >= istart_ad .and. i_maxst <= iend_ad) then
            call zmesh_batch_dotp_matrix(gr, ad_st%group%psib(ib, ik), st%group%psib(jb, ik), overlap, reduce=.false.)
 
            ! partial ad_st batch, full st batch
          else
            do ist = max(i_minst, istart_ad), min(i_maxst, iend_ad)
              call states_elec_get_state(ad_st, gr, ist, ik, psi)
              call zmesh_batch_mf_dotp(gr, st%group%psib(jb, ik), psi, dot(j_minst:j_maxst), reduce=.false., &
                nst=j_maxst-j_minst+1)
              overlap(ist, j_minst:j_maxst) = conjg(dot(j_minst:j_maxst))
            end do
          end if
        end do
 
        ! partial st batch
      else
        do jst = max(j_minst, istart_st), min(j_maxst, iend_st)
          call states_elec_get_state(st, gr, jst, ik, psi)
          call calc_overlap_single(ik, gr, ad_st, psi, istart_ad, iend_ad, overlap(:,jst), reduce=.false.)
        end do
      end if
    end do
 
    if (gr%parallel_in_domains) then
      call gr%allreduce(overlap)
    end if
 
    safe_deallocate_a(psi)
    safe_deallocate_a(dot)
 
    pop_sub_with_profile(calc_overlap_state)
 
  end subroutine calc_overlap_state
 
  subroutine population_in_adiabatic(ik, ad_st, st, overlap, istart, iend, pop)
    integer,                        intent(in)      :: ik
    type(states_elec_t),            intent(inout)   :: ad_st
    type(states_elec_t),            intent(in)      :: st
    complex(real64), contiguous,    intent(in)      :: overlap(:, :)
    integer,                        intent(in)      :: istart
    integer,                        intent(in)      :: iend
    real(real64),                   intent(inout)   :: pop
 
    integer                       :: ist, jst
 
    push_sub(population_in_adiabatic)
 
    assert(ubound(overlap, dim = 1) == ad_st%nst)
    assert(ubound(overlap, dim = 2) == st%nst)
 
    ! Assign ad_st occupation and store in dm_proj_basis (also used for hopping rate calculation)
    ad_st%occ(1:ad_st%nst, ik) = 0.0_real64
    do jst = istart, iend
      do ist = istart, iend
        ad_st%occ(ist, ik) = ad_st%occ(ist, ik) + &
          (real(overlap(ist, jst))**2 + aimag(overlap(ist, jst))**2) * st%occ(jst, ik)
      end do
    end do
 
    pop = pop + sum(ad_st%occ(istart:iend, ik)) * ad_st%kweights(ik)
 
    pop_sub(population_in_adiabatic)
 
  end subroutine population_in_adiabatic
 
  subroutine construct_residuals(gr, namespace, ad_st, st, ik, othn, istart, iend, &
    overlap_ad_ks, nrm2_tdks, nresd, overlap_resd_ks, resd)
    type(grid_t),                    intent(in)     :: gr
    type(namespace_t),               intent(in)     :: namespace
    type(states_elec_t),             intent(in)     :: ad_st
    type(states_elec_t),             intent(in)     :: st
    integer,                         intent(in)     :: ik
    logical,                         intent(in)     :: othn
    integer,                         intent(in)     :: istart
    integer,                         intent(in)     :: iend
    complex(real64), contiguous,     intent(in)     :: overlap_ad_ks(:, :)
    real(real64), contiguous,        intent(in)     :: nrm2_tdks(:)
    integer,                         intent(out)    :: nresd
    complex(real64), contiguous,     intent(out)    :: overlap_resd_ks(:, :)
    type(states_elec_t),             intent(inout)  :: resd
 
    integer                           :: ib, i, j, i_minst, i_maxst
    complex(real64), allocatable      :: d_j(:, :), phi(:, :)
    real(real64), parameter           :: small_rho = 1.0e-14_real64  
    real(real64), parameter           :: small_resd = 1.0e-7_real64  
    real(real64)                      :: nrm_dj, nrm2_dj
    logical, allocatable              :: mask(:)
 
    push_sub_with_profile(construct_residuals)
 
    assert(istart <= iend)
    assert(ubound(overlap_ad_ks, dim = 1) == ad_st%nst)
    assert(ubound(overlap_ad_ks, dim = 2) == st%nst)
    assert(ubound(overlap_resd_ks, dim = 1) == resd%nst)
    assert(ubound(overlap_resd_ks, dim = 2) == st%nst)
 
    safe_allocate(d_j(1:gr%np, st%d%dim))
    safe_allocate(phi(1:gr%np, ad_st%d%dim))
    safe_allocate(mask(1:st%nst))
 
    overlap_resd_ks = m_z0
 
    if (.not. othn) then
 
      ! Iterate over TDKS batches, construct the residuals
      do j = istart, iend
        call states_elec_get_state(st, gr, j, ik, d_j)
        ! Compute |\psi_j > - sum_i S^{\phi,\psi}_{ij} |\phi_i >
        do ib = ad_st%group%block_start, ad_st%group%block_end
          i_minst = states_elec_block_min(ad_st, ib)
          i_maxst = states_elec_block_max(ad_st, ib)
          if (i_minst >= istart .and. i_maxst <= iend) then
            call zbatch_axpy_function(gr%np, -overlap_ad_ks(i_minst:i_maxst, j), &
              ad_st%group%psib(ib, ik), d_j, nst=i_maxst-i_minst+1)
          else
            do i = max(i_minst, istart), min(i_maxst, iend)
              call states_elec_get_state(ad_st, gr, i, ik, phi)
              call lalg_axpy(gr%np, ad_st%d%dim, -overlap_ad_ks(i, j), phi, d_j)
            end do
          end if
        end do
 
        ! reset the residuals
        call states_elec_set_state(resd, gr, j, ik, d_j)
 
        ! Overlap S^{d,\psi} = < d | \psi>
        ! This on-the-fly calculation keeps d_j in cache, avoiding the massive overhead of re-reading the grid later.
        call calc_overlap_single(ik, gr, st, d_j, istart, iend, overlap_resd_ks(j,:), reduce=.false.)
      enddo
 
      nresd = iend - istart + 1
 
    else
 
      ! Construct the orthonormalized residuals
      nresd = 0
 
      ! Iterate over TDKS states, construct the linearly-independent terms as residuals
      do j = istart, iend
        ! nrm2_dj = nrm2_tdks - \sum_i |S^{\phi,\psi}_{ij}|^2
        nrm2_dj = nrm2_tdks(j) - sum(real(overlap_ad_ks(istart:iend, j))**2) - sum(aimag(overlap_ad_ks(istart:iend, j))**2)
        if (nrm2_dj > small_rho) then
          ! scale the TDKS wavefunction to make the residual terms approximately normalized
          call states_elec_get_state(st, gr, j, ik, d_j)
          nrm_dj = sqrt(nrm2_dj)
          call lalg_scal(gr%np, resd%d%dim, m_one / nrm_dj, d_j)
 
          ! substract the overlap with ad states
          ! Compute 1/|d_j| |\psi_j > - 1/|d_j| sum_i S^{\phi,\psi}_{ij} |\phi_i >
          do ib = ad_st%group%block_start, ad_st%group%block_end
            i_minst = states_elec_block_min(ad_st, ib)
            i_maxst = states_elec_block_max(ad_st, ib)
            if (i_minst >= istart .and. i_maxst <= iend) then
              call zbatch_axpy_function(gr%np, -overlap_ad_ks(i_minst:i_maxst, j) / nrm_dj, &
                ad_st%group%psib(ib, ik), d_j, nst=i_maxst-i_minst+1)
            else
              do i = max(i_minst, istart), min(i_maxst, iend)
                call states_elec_get_state(ad_st, gr, i, ik, phi)
                call lalg_axpy(gr%np, ad_st%d%dim, -overlap_ad_ks(i, j) / nrm_dj, phi, d_j)
              end do
            end if
          end do
 
          mask = .true.
          mask(istart:iend) = .false.
          call zstates_elec_orthogonalize_single_batch(ad_st, gr, st%nst, ik, d_j, mask=mask)
 
          ! substract the overlap with previous residuals
          if (nresd > 0) then
            mask = .true.
            mask(istart:istart+nresd-1) = .false.
            call zstates_elec_orthogonalize_single_batch(resd, gr, istart+nresd-1, ik, d_j, mask=mask)
          end if
          ! screen the small terms
          nrm_dj = zmf_nrm2(gr, resd%d%dim, d_j, reduce=.true.)
          if (nrm_dj > small_resd) then
            call lalg_scal(gr%np, resd%d%dim, m_one / nrm_dj, d_j)
            nresd = nresd + 1
            call states_elec_set_state(resd, gr, istart+nresd-1, ik, d_j)
 
            ! Overlap matrix S^{d,\psi}_{ji} = < d_j | \psi_i >
            call calc_overlap_single(ik, gr, st, d_j, istart, iend, overlap_resd_ks(istart+nresd-1,:), reduce=.false.)
          end if
        end if
      enddo
    end if
 
    if (gr%parallel_in_domains) then
      call gr%allreduce(overlap_resd_ks)
    end if
 
    overlap_resd_ks = conjg(overlap_resd_ks)
 
    safe_deallocate_a(d_j)
    safe_deallocate_a(phi)
    safe_deallocate_a(mask)
 
    pop_sub_with_profile(construct_residuals)
 
  end subroutine construct_residuals
 
  subroutine construct_density_matrix(nresd, ik, st, istart, iend, overlap_ad_ks, overlap_resd_ks, rho_mat, rho_offd)
    integer,                            intent(in)     :: nresd
    integer,                            intent(in)     :: ik
    type(states_elec_t),                intent(in)     :: st
    integer,                            intent(in)     :: istart
    integer,                            intent(in)     :: iend
    complex(real64), contiguous,        intent(in)     :: overlap_ad_ks(:, :)
    complex(real64), contiguous,        intent(in)     :: overlap_resd_ks(:, :)
    complex(real64), contiguous,        intent(out)    :: rho_mat(:, :)
    complex(real64), contiguous,        intent(out)    :: rho_offd(:)
 
    integer                           :: ist, jst, lst, act_nst, idx
    real(real64)                      :: sqrt_f
    complex(real64), allocatable      :: s_ad_scaled(:, :), s_resd_scaled(:, :)
 
    push_sub(construct_density_matrix)
 
    act_nst = iend - istart + 1
 
    assert(istart <= iend)
    assert(ubound(overlap_ad_ks, dim = 1) == st%nst)
    assert(ubound(overlap_ad_ks, dim = 2) == st%nst)
    assert(ubound(overlap_resd_ks, dim = 1) == st%nst)
    assert(ubound(overlap_resd_ks, dim = 2) == st%nst)
    assert(ubound(rho_mat, dim = 1) == 2*act_nst)
    assert(ubound(rho_mat, dim = 2) == 2*act_nst)
    assert(ubound(rho_offd, dim = 1) == act_nst*(act_nst-1)/2)
 
    safe_allocate(s_ad_scaled(st%nst, st%nst))
    if (nresd > 0) then
      safe_allocate(s_resd_scaled(st%nst, st%nst))
    end if
 
    ! initialize to zero
    rho_mat = m_z0
    s_ad_scaled = m_z0
    if (nresd > 0) then
      s_resd_scaled = m_z0
    end if
 
    ! scale S matrix
    do lst = istart, iend
      sqrt_f = sqrt(st%occ(lst, ik))
      s_ad_scaled(istart:iend, lst) = overlap_ad_ks(istart:iend, lst) * sqrt_f
      if (nresd > 0) then
        s_resd_scaled(istart:istart+nresd-1, lst) = overlap_resd_ks(istart:istart+nresd-1, lst) * sqrt_f
      end if
    end do
 
    ! ad-ad
    call blas_herk('U', 'N', act_nst, act_nst, 1.0_real64, s_ad_scaled(istart,istart), st%nst, 0.0_real64, &
      rho_mat(1, 1), 2*act_nst)
    if (nresd > 0) then
      ! resd-resd
      call blas_herk('U', 'N', nresd, act_nst, 1.0_real64, s_resd_scaled(istart,istart), st%nst, 0.0_real64, &
        rho_mat(act_nst+1, act_nst+1), 2*act_nst)
      ! ad-resd
      call blas_gemm('N', 'C', act_nst, nresd, act_nst, m_z1, s_ad_scaled(istart, istart), st%nst, &
        s_resd_scaled(istart,istart), st%nst, m_z0, rho_mat(1, act_nst+1), 2*act_nst)
    end if
 
    !$omp parallel do private(jst, ist)
    do jst = 1, act_nst
      ! mirror ad-ad lower triangle
      do ist = 1, jst - 1
        rho_mat(jst, ist) = conjg(rho_mat(ist, jst))
      end do
 
      ! mirror ad-resd to resd-ad (bottom-left block)
      do ist = 1, nresd
        rho_mat(act_nst+ist, jst) = conjg(rho_mat(jst, act_nst+ist))
      end do
    end do
    !$omp end parallel do
 
    !$omp parallel do private(jst, ist)
    do jst = 1, nresd
      ! mirror resd-resd lower triangle
      do ist = 1, jst - 1
        rho_mat(act_nst+jst, act_nst+ist) = conjg(rho_mat(act_nst+ist, act_nst+jst))
      end do
    end do
    !$omp end parallel do
 
    safe_deallocate_a(s_ad_scaled)
    safe_deallocate_a(s_resd_scaled)
 
    ! off-diagonal elements
    idx = 0
    do jst = 1, act_nst
      do ist = 1, jst - 1
        idx = idx + 1
        rho_offd(idx) = rho_mat(ist, jst)
      end do
    end do
 
    pop_sub(construct_density_matrix)
  end subroutine construct_density_matrix
 
  subroutine broadcast_occupation(occ, kpt, parstate, istart, iend)
    real(real64), contiguous, intent(inout)   :: occ(:, :)
    type(distributed_t),      intent(in)      :: kpt
    logical,                  intent(in)      :: parstate
    integer,                  intent(in)      :: istart
    integer,                  intent(in)      :: iend
 
    integer                       :: incount, act_nst
    integer, allocatable          :: rdispls(:), recvcnts(:)
    real(real64), allocatable     :: sendbuffer(:, :), recvbuffer(:, :)
 
    push_sub_with_profile(broadcast_occupation)
 
    assert(istart <= iend)
    assert(ubound(occ, dim = 1) >= iend)
    assert(ubound(occ, dim = 2) == kpt%nglobal())
    assert(.not. parstate)
 
    if (kpt%parallel) then
      act_nst = iend - istart + 1
 
      safe_allocate(recvcnts(1:kpt%mpi_grp%size))
      safe_allocate(rdispls(1:kpt%mpi_grp%size))
      safe_allocate(sendbuffer(1:act_nst, kpt%nlocal()))
      safe_allocate(recvbuffer(1:act_nst, 1:kpt%nglobal()))
 
      incount = act_nst * kpt%nlocal()
      recvcnts(:) = act_nst * kpt%num(:)
      sendbuffer(1:act_nst, 1:kpt%nlocal()) = occ(istart:iend, kpt%start:kpt%end)
 
      call mpi_displacements(recvcnts, rdispls)
      ! send buffer and receive buffer can not be the same address
      ! recvbuffer should be contiguous in memory
      call kpt%mpi_grp%allgatherv(sendbuffer, incount, mpi_double_precision, &
        recvbuffer, recvcnts, rdispls, mpi_double_precision)
 
      occ(istart:iend, 1:kpt%nglobal()) = recvbuffer(1:act_nst, 1:kpt%nglobal())
 
      safe_deallocate_a(sendbuffer)
      safe_deallocate_a(recvbuffer)
      safe_deallocate_a(recvcnts)
      safe_deallocate_a(rdispls)
    end if
 
    pop_sub_with_profile(broadcast_occupation)
 
  end subroutine broadcast_occupation
 
  subroutine dissipation(hm, st, namespace, nresd_k, dt, dmp, rho_mat_k)
    type(hamiltonian_elec_t), intent(in)     :: hm
    type(states_elec_t),      intent(in)     :: st
    type(namespace_t),        intent(in)     :: namespace
    integer,                  intent(in)     :: nresd_k(:)
    real(real64),             intent(in)     :: dt
    type(dmp_t),              intent(inout)  :: dmp
    complex(real64),          intent(inout)  :: rho_mat_k(:, :, :)
 
    integer                   :: act_nst
 
    push_sub_with_profile(dissipation)
 
    act_nst = dmp%iend - dmp%istart + 1
 
    assert(ubound(nresd_k, dim = 1) == st%d%kpt%nlocal())
    assert(ubound(rho_mat_k, dim = 1) == 2*act_nst)
    assert(ubound(rho_mat_k, dim = 2) == 2*act_nst)
    assert(ubound(rho_mat_k, dim = 3) == st%d%kpt%nlocal())
 
    if (dmp%calculation_mode == option__tddmpropagation__master_equation) then
      select case (dmp%strategy)
      case (1)
        call lindblad_uniform(dmp, st%d%kpt, nresd_k, dt, rho_mat_k)
      case (2)
        call lindblad_2times(dmp, st%d%kpt, nresd_k, dt, rho_mat_k)
      case (3)
        call lindblad_from_epw(dmp, hm, st%d%kpt, st%system_grp, namespace, nresd_k, dt, rho_mat_k)
      end select
    else if (dmp%calculation_mode == option__tddmpropagation__collision_integral) then
      call collision_from_epw(dmp, hm, st%d%kpt, st%system_grp, namespace, nresd_k, dt, rho_mat_k)
    end if
 
    pop_sub_with_profile(dissipation)
 
  end subroutine dissipation
 
  subroutine lindblad_uniform(dmp, kpt, nresd_k, dt, rho_mat_k)
    type(dmp_t),         intent(in)     :: dmp
    type(distributed_t), intent(in)     :: kpt
    integer,             intent(in)     :: nresd_k(:)
    real(real64),        intent(in)     :: dt
    complex(real64),     intent(inout)  :: rho_mat_k(:, :, :)
 
    real(real64)                  :: coeff
    real(real64), allocatable     :: rtrans(:, :)
    complex(real64), allocatable  :: rho_in(:, :), rho_out(:, :), rho_res(:, :), rho_tmp(:, :)
    integer                       :: iorder, nst, ik, ik_, nresd, act_nst
    integer, parameter            :: norder = 4
 
    push_sub_with_profile(lindblad_uniform)
 
    nst = dmp%adiabatic_st%nst
    act_nst = dmp%iend - dmp%istart + 1
 
    ! lindblad operate on the density matrix
    safe_allocate(rtrans(1:act_nst, 1:act_nst))
    safe_allocate(rho_in(1:2*act_nst, 1:2*act_nst))
    safe_allocate(rho_out(1:2*act_nst, 1:2*act_nst))
    safe_allocate(rho_res(1:2*act_nst, 1:2*act_nst))
 
    do ik = kpt%start, kpt%end
      ik_ = ik - kpt%start + 1
      nresd = nresd_k(ik_)
 
      call transition_rate_uniform(dmp%uniform, dmp%adiabatic_st, dmp%istart, dmp%iend, ik, rtrans)
 
      rho_res = 0.0_real64
      rho_in(1:2*act_nst, 1:2*act_nst) = rho_mat_k(1:2*act_nst, 1:2*act_nst, ik_)
 
      coeff = m_one
      do iorder = 1, norder-1
 
        call lindblad_operator_uniform(nresd, dmp%istart, dmp%iend, rtrans, rho_in, rho_out)
 
        coeff = coeff * dt / iorder
        rho_res = rho_res + coeff * rho_out
        ! swap the pointer between rho_in and rho_out, so new rho_in is current rho_out,
        ! and new rho_out is the prior input that will get overwritten
        call move_alloc(rho_in, rho_tmp)
        call move_alloc(rho_out, rho_in)
        call move_alloc(rho_tmp, rho_out)
      end do
      ! iorder == norder
      call lindblad_operator_uniform(nresd, dmp%istart, dmp%iend, rtrans, rho_in, rho_out)
      coeff = coeff * dt / norder
      rho_res = rho_res + coeff * rho_out
 
      rho_mat_k(1:2*act_nst, 1:2*act_nst, ik_) = rho_mat_k(1:2*act_nst, 1:2*act_nst, ik_) + rho_res
    end do
 
    safe_deallocate_a(rho_in)
    safe_deallocate_a(rho_out)
    safe_deallocate_a(rtrans)
    safe_deallocate_a(rho_res)
 
    pop_sub_with_profile(lindblad_uniform)
 
  end subroutine lindblad_uniform
 
  subroutine transition_rate_uniform(uniform, ad_st, istart, iend, ik, rtrans)
    real(real64), contiguous,    intent(in)     :: uniform(:)
    type(states_elec_t),         intent(in)     :: ad_st
    integer,                     intent(in)     :: istart
    integer,                     intent(in)     :: iend
    integer,                     intent(in)     :: ik
    real(real64), contiguous,    intent(out)    :: rtrans(:, :)
 
    integer                       :: ist, jst, act_nst
    real(real64)                  :: rate_character, omega, inv_omega, nph, delta_e
    real(real64), parameter       :: small_e = 0.002_real64/(m_two*p_ry), large_e = 0.5_real64/(m_two*p_ry)
    real(real64)                  :: unocc_ist, unocc_jst
 
    push_sub(transition_rate_uniform)
 
    act_nst = iend - istart + 1
 
    assert(ubound(rtrans, dim = 1) == act_nst)
    assert(ubound(rtrans, dim = 2) == act_nst)
 
    rate_character = 1 / uniform(1)
    omega = units_to_atomic(unit_kelvin, uniform(2))
    inv_omega = 1.0_real64 / max(omega, 1.0e-12_real64)
 
    rtrans = m_zero
    ! rtrans(ist, jst) is the transition rate from jst to ist
    do ist = istart, iend
      unocc_ist = 1.0_real64 - ad_st%occ(ist, ik) / ad_st%smear%el_per_state
 
      do jst = ist + 1, iend
        delta_e = ad_st%eigenval(jst, ik) - ad_st%eigenval(ist, ik)
        if (delta_e < small_e) cycle
 
        if (delta_e < large_e) then
          nph = 1.0_real64/(exp(delta_e*inv_omega) - 1.0_real64)
        else
          nph = m_zero
        end if
        unocc_jst = 1.0_real64 - ad_st%occ(jst, ik) / ad_st%smear%el_per_state
 
        rtrans(ist-istart+1, jst-istart+1) = merge(rate_character * unocc_ist * (nph + 1), m_zero, unocc_ist > m_zero)
        rtrans(jst-istart+1, ist-istart+1) = merge(rate_character * unocc_jst * nph, m_zero, unocc_jst > m_zero)
      end do
    end do
 
    pop_sub(transition_rate_uniform)
 
  end subroutine transition_rate_uniform
 
  subroutine lindblad_operator_uniform(nresd, istart, iend, rtrans, den_mat, l_mat)
    integer,                         intent(in)    :: nresd
    integer,                         intent(in)    :: istart
    integer,                         intent(in)    :: iend
    real(real64), contiguous,        intent(in)    :: rtrans(:, :)
    complex(real64), contiguous,     intent(in)    :: den_mat(:, :)
    complex(real64), contiguous,     intent(out)   :: l_mat(:, :)
 
    integer                     :: ist, jst, lst, act_nst
 
    push_sub_with_profile(lindblad_operator_uniform)
 
    act_nst = iend - istart + 1
 
    assert(ubound(rtrans, dim = 1) == act_nst)
    assert(ubound(rtrans, dim = 2) == act_nst)
    assert(ubound(den_mat, dim = 1) == 2*act_nst)
    assert(ubound(den_mat, dim = 2) == 2*act_nst)
    assert(ubound(l_mat, dim = 1) == 2*act_nst)
    assert(ubound(l_mat, dim = 2) == 2*act_nst)
 
    l_mat = 0.0_real64
 
    do ist = 1, act_nst
      do jst = 1, act_nst
        if (ist == jst) cycle
        do lst = 1, act_nst+nresd
          l_mat(ist, lst) = l_mat(ist, lst) - m_half * rtrans(jst, ist) * den_mat(ist, lst)
          l_mat(lst, ist) = l_mat(lst, ist) - m_half * rtrans(jst, ist) * den_mat(lst, ist)
        end do
        l_mat(jst, jst) = l_mat(jst, jst) + rtrans(jst, ist) * den_mat(ist, ist)
      end do
    end do
 
    pop_sub_with_profile(lindblad_operator_uniform)
 
  end subroutine lindblad_operator_uniform
 
  subroutine lindblad_2times(dmp, kpt, nresd_k, dt, rho_mat_k)
    type(dmp_t),         intent(in)     :: dmp
    type(distributed_t), intent(in)     :: kpt
    integer,             intent(in)     :: nresd_k(:)
    real(real64),        intent(in)     :: dt
    complex(real64),     intent(inout)  :: rho_mat_k(:, :, :)
 
    real(real64)                  :: decay_T1, decay_T2
    integer                       :: ist, jst, ik, ik_, nresd, act_nst
 
    push_sub_with_profile(lindblad_2times)
 
    act_nst = dmp%iend - dmp%istart + 1
 
    decay_t1 = exp(-dt / dmp%tmodel(1))
    decay_t2 = exp(-dt / dmp%tmodel(2))
 
    do ik = kpt%start, kpt%end
      ik_ = ik - kpt%start + 1
      nresd = nresd_k(ik_)
 
      do ist = dmp%istart, dmp%iend
        rho_mat_k(ist-dmp%istart+1, ist-dmp%istart+1, ik_) = dmp%occ_gs(ist, ik_) + &
          (rho_mat_k(ist-dmp%istart+1, ist-dmp%istart+1, ik_) - dmp%occ_gs(ist, ik_)) * decay_t1
      end do
      do ist = act_nst+1, act_nst+nresd
        rho_mat_k(ist, ist, ik_) = rho_mat_k(ist, ist, ik_) * decay_t1
      end do
      do ist = 1, act_nst+nresd
        do jst = ist+1, act_nst+nresd
          rho_mat_k(jst, ist, ik_) = rho_mat_k(jst, ist, ik_) * decay_t2
          rho_mat_k(ist, jst, ik_) = rho_mat_k(ist, jst, ik_) * decay_t2
        end do
      end do
    end do
 
    pop_sub_with_profile(lindblad_2times)
  end subroutine lindblad_2times
 
  subroutine lindblad_from_epw(dmp, hm, kpt, system_grp, namespace, nresd_k, dt, rho_mat_k)
    type(dmp_t),              intent(inout)  :: dmp
    type(hamiltonian_elec_t), intent(in)     :: hm
    type(distributed_t),      intent(in)     :: kpt
    type(mpi_grp_t),          intent(in)     :: system_grp
    type(namespace_t),        intent(in)     :: namespace
    integer,                  intent(in)     :: nresd_k(:)
    real(real64),             intent(in)     :: dt
    complex(real64),          intent(inout)  :: rho_mat_k(:, :, :)
 
    real(real64)                  :: coeff
    real(real64), allocatable     :: rho_diag(:, :)
    complex(real64), allocatable  :: rho_in(:, :, :), rho_out(:, :, :), rho_tmp(:, :, :)
    integer, parameter            :: norder = 4
    integer                       :: ist, iorder, ik, ik_, nik, act_nst
 
    push_sub_with_profile(lindblad_from_epw)
 
    nik = dmp%adiabatic_st%nik
    act_nst = dmp%iend - dmp%istart + 1
 
    call dmp%update_trans_rate(hm, system_grp, namespace)
 
    ! the diagonal elements of L\rho is initialized to occupation
    ! SAFE_ALLOCATE_SOURCE_A(rho_diag, dmp%adiabatic_st%occ)
    safe_allocate(rho_diag(1:act_nst, nik))
    rho_diag(1:act_nst, 1:nik) = dmp%adiabatic_st%occ(dmp%istart:dmp%iend, 1:nik)
 
    ! L_D[rho_in] = rho_out
    safe_allocate_source(rho_in, rho_mat_k)
    safe_allocate(rho_out(1:2*act_nst, 1:2*act_nst, 1:kpt%nlocal()))
 
    ! dimensition of rho_out should be modified later
    coeff = 1.0_real64
    do iorder = 1, norder - 1
 
      coeff = coeff * dt / iorder
 
      do ik = kpt%start, kpt%end
        ik_ = ik - kpt%start + 1
 
        call lindblad_operator_epw(dmp, ik, hm%kpoints%nik_skip, nresd_k(ik_), rho_diag, rho_in(:, :, ik_), rho_out(:, :, ik_))
 
        call lalg_axpy(2*act_nst, 2*act_nst, coeff, rho_out(:, :, ik_), rho_mat_k(:, :, ik_))
      end do
 
      ! update rho_diag independently
      do ik = kpt%start, kpt%end
        ik_ = ik - kpt%start + 1
        do ist = dmp%estart, dmp%eend
          rho_diag(ist-dmp%istart+1, ik) = real(rho_out(ist-dmp%istart+1, ist-dmp%istart+1, ik_))
        end do
      end do
      ! broadcast occupation for all processes, modify later
      ! call broadcast_occupation(rho_diag, kpt, nst, dmp%adiabatic_st%parallel_in_states)
      call broadcast_occupation(rho_diag, kpt, dmp%adiabatic_st%parallel_in_states, &
        dmp%estart-dmp%istart+1, dmp%eend-dmp%istart+1)
 
      ! swap the pointer between rho_in and rho_out, so new rho_in is current rho_out,
      ! and new rho_out is the prior input that will get overwritten
      call move_alloc(rho_in, rho_tmp)
      call move_alloc(rho_out, rho_in)
      call move_alloc(rho_tmp, rho_out)
    end do
 
    ! iorder == norder
    coeff = coeff * dt / norder
 
    do ik = kpt%start, kpt%end
      ik_ = ik - kpt%start + 1
 
      call lindblad_operator_epw(dmp, ik, hm%kpoints%nik_skip, nresd_k(ik_), rho_diag, rho_in(:, :, ik_), rho_out(:, :, ik_))
 
      call lalg_axpy(2*act_nst, 2*act_nst, coeff, rho_out(:, :, ik_), rho_mat_k(:, :, ik_))
    end do
 
 
    safe_deallocate_a(rho_in)
    safe_deallocate_a(rho_out)
    safe_deallocate_a(rho_diag)
 
    pop_sub_with_profile(lindblad_from_epw)
  end subroutine lindblad_from_epw
 
  subroutine collision_from_epw(dmp, hm, kpt, system_grp, namespace, nresd_k, dt, rho_mat_k)
    type(dmp_t),              intent(inout)  :: dmp
    type(hamiltonian_elec_t), intent(in)     :: hm
    type(distributed_t),      intent(in)     :: kpt
    type(mpi_grp_t),          intent(in)     :: system_grp
    type(namespace_t),        intent(in)     :: namespace
    integer,                  intent(in)     :: nresd_k(:)
    real(real64),             intent(in)     :: dt
    complex(real64),          intent(inout)  :: rho_mat_k(:, :, :)
 
    real(real64)                  :: gam, gam_in, gam_out
    real(real64), allocatable     :: gam_bnd(:)
    real(real64), parameter       :: gthresh = 1.0e-8_real64
    integer                       :: ist, jst, ik, nresd, ik_, nik_skip, act_nst
 
    push_sub_with_profile(collision_from_epw)
 
    nik_skip = hm%kpoints%nik_skip
    act_nst = dmp%iend - dmp%istart + 1
 
    safe_allocate(gam_bnd(2*dmp%adiabatic_st%nst))
 
    call dmp%update_trans_rate(hm, system_grp, namespace)
 
    ! off-diagonals
    do ik = kpt%start, kpt%end
      ik_ = ik - kpt%start + 1
      nresd = nresd_k(ik_)
 
      ! pre-calculate gamma
      gam_bnd = 0.0_real64
      do ist = dmp%estart, dmp%eend
        call lifetime(dmp, ik, ist, nik_skip, gam_bnd(ist))
      end do
 
      do ist = 1, act_nst + nresd
        do jst = ist + 1, act_nst + nresd
          gam = -(gam_bnd(ist+dmp%istart-1) + gam_bnd(jst+dmp%istart-1)) / 2.0_real64
 
          ! damping
          rho_mat_k(jst, ist, ik_) = rho_mat_k(jst, ist, ik_) * exp(gam * dt)
          rho_mat_k(ist, jst, ik_) = conjg(rho_mat_k(jst, ist, ik_))
        end do
      end do
    end do
 
    ! diagonals, \dot{f} = -gam_out * f + gam_in, gives analytical solution
    do ik = kpt%start, kpt%end
      ik_ = ik - kpt%start + 1
      do ist = dmp%estart, dmp%eend
        call get_gamma(dmp, ik, ist, nik_skip, gam_in, gam_out)
        if (gam_out * dt > gthresh) then
          rho_mat_k(ist-dmp%istart+1, ist-dmp%istart+1, ik_) = rho_mat_k(ist-dmp%istart+1, ist-dmp%istart+1, ik_) *&
            exp(-gam_out * dt) + (gam_in / gam_out) * (1.0_real64 - exp(-gam_out * dt))
        else
          ! when denominator divergent
          rho_mat_k(ist-dmp%istart+1, ist-dmp%istart+1, ik_) = rho_mat_k(ist-dmp%istart+1, ist-dmp%istart+1, ik_) *&
            (1.0_real64 - gam_out * dt) + gam_in * dt
        end if
      end do
    end do
 
    safe_deallocate_a(gam_bnd)
 
    pop_sub_with_profile(collision_from_epw)
  end subroutine collision_from_epw
 
  subroutine dm_propagation_update_trans_rate(this, hm, system_grp, namespace)
    class(dmp_t),              intent(inout)  :: this
    type(hamiltonian_elec_t),  intent(in)     :: hm
    type(mpi_grp_t),           intent(in)     :: system_grp
    type(namespace_t),         intent(in)     :: namespace
 
    integer                   :: ia
 
    push_sub_with_profile(dmp_propagation_update_trans_rate)
 
    ia = get_vector_field_index(this, hm, namespace)
 
    if (ia /= this%ia) then
      call iopar_read_trans_rate(ia, system_grp, namespace, this)
      this%ia = ia
    end if
 
    pop_sub_with_profile(dmp_propagation_update_trans_rate)
 
  end subroutine dm_propagation_update_trans_rate
 
  subroutine iopar_open_trans_rate(namespace, ions, hm, system_grp, dmp)
    type(namespace_t),        intent(in)    :: namespace
    type(ions_t),             intent(in)    :: ions
    type(hamiltonian_elec_t), intent(in)    :: hm
    type(mpi_grp_t),          intent(in)    :: system_grp
    type(dmp_t),              intent(inout) :: dmp
 
    integer               :: ierr, iostat, idim, idir, iqq, totq
    integer               :: epw_nk(3), oct_nk(3)
    real(real64)          :: oct_s(3), at(3,3)
 
    push_sub(iopar_open_trans_rate)
 
    ! metadata
    if (system_grp%is_root()) then
      open(newunit=dmp%iunit, file=trim(dmp%epw_file), form='unformatted', access='stream', status='old', &
        action='read', iostat=iostat)
      if (iostat /= 0) then
        dmp%iunit = -1
        write(message(1), '(a,a)') 'Error opening file: ', trim(dmp%epw_file)
        call messages_fatal(1, namespace=namespace)
      end if
      !
      read(dmp%iunit, iostat=ierr) iqq, totq, dmp%wnst, epw_nk(1:3), at, oct_nk(1:3), oct_s(1:3), dmp%astep(1:3), dmp%na(1:3)
      if (ierr /= 0) then
        write(message(1), '(a,a)') 'Error reading header from: ', trim(dmp%epw_file)
        call messages_fatal(1, namespace=namespace)
      end if
      !
    end if
 
    call system_grp%bcast(dmp%wnst, 1, mpi_integer, 0)
    call system_grp%bcast(at, 9, mpi_double_precision, 0)
    call system_grp%bcast(oct_s, 3, mpi_double_precision, 0)
    call system_grp%bcast(dmp%astep, 3, mpi_double_precision, 0)
    call system_grp%bcast(dmp%na, 3, mpi_integer, 0)
    call system_grp%bcast(epw_nk, 3, mpi_integer, 0)
    call system_grp%bcast(oct_nk, 3, mpi_integer, 0)
 
    if (any(oct_nk /= hm%kpoints%nik_axis)) then
      write(message(1), '(a, a)') 'Inconsistent k-point mesh in KPointsGrid and ', trim(dmp%epw_file)
      call messages_fatal(1, namespace=namespace)
    end if
    if (any(abs(oct_s - hm%kpoints%full%shifts(:, 1)) > m_epsilon)) then
      write(message(1), '(a, a)') 'Inconsistent k-point mesh shifts in KPointsGrid and ', trim(dmp%epw_file)
      call messages_fatal(1, namespace=namespace)
    end if
    !
    write(message(1), '(3(a,i0), a)') 'Info: Averaged transition rates obtained from a ', epw_nk(1), ' x ', &
      epw_nk(2), ' x ', epw_nk(3), ' EPW k-point mesh'
    call messages_info(1, namespace=namespace)
    !
    dmp%eend = dmp%estart + dmp%wnst - 1
    write(message(1), '(a, i0, a, i0)') 'Info: Damping band ', dmp%estart, ' - ', dmp%eend
    call messages_info(1, namespace=namespace)
    if (dmp%eend > dmp%iend) then
      write(message(1), '(a)') 'The highest damping band cannot exceed TDDMBandHighest.'
      call messages_fatal(1, namespace=namespace)
    end if
    !
    write(message(1),'(a)') '  EPW Lattice Vectors [1/alat]'
    do idim = 1, 3
      write(message(1+idim),'(3f12.6)') (at(idir, idim), idir = 1, 3)
    end do
    call messages_info(4, namespace=namespace)
    if (any(abs(at - ions%latt%rlattice_primitive) > m_epsilon)) then
      write(message(1),'(a)') 'Lattice settings are not fully consistent with those in EPW'
      call messages_warning(1, namespace=namespace)
    end if
 
    call build_epw_kmap(namespace, hm%kpoints, dmp)
 
    ! initialize vector field grid index
    dmp%ia = -1
 
    ! shared memory for transition rates
    dmp%num = int(product(oct_nk), kind=int64)**2 * int(dmp%wnst, kind=int64)**2
#ifdef HAVE_MPI
    ! create shared memory window and fill it only on root
    call create_intranode_communicator(system_grp, dmp%intranode_grp, dmp%internode_grp)
    ! We inline the logic of lmpi_create_shared_memory_window because single precision is not supported.
    call slmpi_create_shared_memory_window(dmp%num, dmp%intranode_grp, dmp%window_trans_rate, ave_trans)
#else
    safe_allocate(ave_trans(1:dmp%num))
#endif
 
    pop_sub(iopar_open_trans_rate)
  end subroutine iopar_open_trans_rate
 
  subroutine iopar_read_trans_rate(ia, system_grp, namespace, dmp)
    integer,                  intent(in)      :: ia
    type(mpi_grp_t),          intent(in)      :: system_grp
    type(namespace_t),        intent(in)      :: namespace
    type(dmp_t),              intent(inout)   :: dmp
 
    integer(int64), parameter   :: header_bytes = 168 ! head info in binary file
    integer(int64), parameter   :: epw_bytes = 4      ! single precision
    integer(int64)              :: offset
    integer                     :: ierr
 
    push_sub(iopar_read_trans_rate)
 
    if (system_grp%is_root()) then
      ! single precision
      offset = header_bytes + int(ia - 1, kind=int64) * dmp%num * epw_bytes + 1_int64
      read(dmp%iunit, pos=offset, iostat=ierr) ave_trans
      if (ierr /= 0) then
        write(message(1), '(a,a)') 'Error reading transition rates from: ', trim(dmp%epw_file)
        call messages_fatal(1, namespace=namespace)
      end if
    end if
 
#ifdef HAVE_MPI
    ! now broadcast the global arrays to local rank 0 on each node
    if (dmp%intranode_grp%rank == 0) then
      call dmp%internode_grp%bcast(ave_trans(1), dmp%num, mpi_real, 0)
    end if
    call lmpi_sync_shared_memory_window(dmp%window_trans_rate, dmp%intranode_grp)
#endif
 
    pop_sub(iopar_read_trans_rate)
  end subroutine iopar_read_trans_rate
 
  subroutine iopar_close_trans_rate(system_grp, dmp)
    type(mpi_grp_t), intent(in)    :: system_grp
    type(dmp_t),     intent(inout) :: dmp
 
    push_sub(iopar_close_trans_rate)
 
    safe_deallocate_a(dmp%kmap)
 
    if (system_grp%is_root()) then
      close(dmp%iunit)
    end if
 
#ifdef HAVE_MPI
    call lmpi_destroy_shared_memory_window(dmp%window_trans_rate)
    nullify(ave_trans)
#else
    safe_deallocate_p(ave_trans)
#endif
 
    pop_sub(iopar_close_trans_rate)
 
  end subroutine iopar_close_trans_rate
 
  subroutine build_epw_kmap(namespace, kpoints, dmp)
    type(namespace_t),    intent(in)    :: namespace
    type(kpoints_t),      intent(in)    :: kpoints
    type(dmp_t),          intent(inout) :: dmp
 
    integer                 :: ik, idx, nik, nik_mp
    integer                 :: kidx_(3)
    real(real64)            :: kred(3), kidx(3)
    real(real64), parameter :: tol = 1.0d-10
 
    push_sub(build_epw_kmap)
 
    nik = dmp%adiabatic_st%nik
    nik_mp = nik - kpoints%nik_skip
    safe_allocate(dmp%kmap(nik))
 
    ! mp k-points
    dmp%kmap = nik_mp + 1
 
    do ik = 1, nik
      ! Map reduced coordinate into an integer MP-grid coordinate
      kred = kpoints%get_point(ik, .false.)
      kidx = (kred + 0.5_real64) * real(kpoints%nik_axis, kind=real64) - kpoints%full%shifts(:, 1)
 
      ! for MP k-points
      if (ik <= nik_mp) then
        if (any(abs(kidx - nint(kidx)) > tol)) then
          write(message(1), '(a)') 'K-point mesh is not compatible with EPW input'
          call messages_fatal(1, namespace=namespace)
        end if
      end if
 
      ! Fold into central cell in EPW
      kidx_ = modulo(nint(kidx), kpoints%nik_axis)
      ! Flatten into a single, contiguous index
      idx = kidx_(1) * kpoints%nik_axis(2) * kpoints%nik_axis(3) + kidx_(2) * kpoints%nik_axis(3) + &
        kidx_(3) + 1
 
      dmp%kmap(ik) = idx
    end do
 
    ! sanity check
    assert(all(dmp%kmap <= nik_mp))
 
    pop_sub(build_epw_kmap)
  end subroutine build_epw_kmap
 
  function get_vector_field_index(dmp, hm, namespace) result(aidx)
    type(dmp_t),              intent(in)  :: dmp
    type(hamiltonian_elec_t), intent(in)  :: hm
    type(namespace_t),        intent(in)  :: namespace
 
    real(real64)           :: ared(3), approx(3)
    integer                :: apoint_idx(3)
    integer                :: aidx, idim
 
    push_sub(get_vector_field_index)
 
    ! Find the nearest discrete vector field grid point on the EPW vector field mesh
    if (allocated(hm%hm_base%uniform_vector_potential)) then
      call kpoints_to_reduced(hm%ions%latt, hm%hm_base%uniform_vector_potential, ared)
    else
      ared = 0.0_real64
    end if
    do idim = 1, 3
      if (dmp%astep(idim) > m_zero) then
        apoint_idx(idim) = nint(ared(idim) / dmp%astep(idim))
        approx(idim) = apoint_idx(idim) * dmp%astep(idim)
      else
        apoint_idx(idim) = 0
        approx(idim) = 0.0_real64
      end if
    end do
    if (any(abs(apoint_idx) > dmp%na)) then
      write(message(1), '(a, 3F8.3)') 'Vector potential exceeds mesh range: ', ared
      call messages_warning(1, namespace=namespace)
      where (apoint_idx < -dmp%na)
        apoint_idx = -dmp%na
      end where
      where (apoint_idx > dmp%na)
        apoint_idx = dmp%na
      end where
    end if
    ! Get the flattened 1D index
    ! Step 1: Find the relative offset by subtracting the minimum index (-na).
    !         This shifts the range [-na, na] to an offset [0, 2*na].
    !         e.g., if range is [-2, 2], the offset for 1 is 1 - (-2) = 3.
    ! Step 2: Flatten the 3D offsets into 1D, and finally add 1 for Fortran 1-based indexing.
    aidx = (apoint_idx(1) + dmp%na(1)) * (2 * dmp%na(2) + 1) * (2 * dmp%na(3) + 1) + &
      (apoint_idx(2) + dmp%na(2)) * (2 * dmp%na(3) + 1) + &
      (apoint_idx(3) + dmp%na(3)) + 1
 
    write(message(1), '(a, x, 3F7.3)') 'Approximated vector field damping grid:', approx
    call messages_info(1, namespace=namespace)
 
    pop_sub(get_vector_field_index)
  end function get_vector_field_index
 
  function get_trans_rate(dmp, nik_mp, jbnd, ibnd, k, kq, p_block) result(res)
    type(dmp_t),        intent(in) :: dmp
    integer,            intent(in) :: nik_mp
    integer,            intent(in) :: jbnd
    integer,            intent(in) :: ibnd
    integer,            intent(in) :: k
    integer,            intent(in) :: kq
    logical, optional,  intent(in) :: p_block
 
    real(real64)      :: unocc, res
    integer           :: k_, kq_
    integer(int64)    :: idx
 
    push_sub(get_trans_rate)
 
    k_ = dmp%kmap(k)
    kq_ = dmp%kmap(kq)
    ! Flatten 4D transition rate index to 1D: (kq, k, ibnd, jbnd) -> idx
    ! Matrix dimensions are [nik_mp, nik_mp, wnst, wnst]
    ! The layout follows the hierarchy: kq is the slowest index, jbnd is the fastest.
    idx = int(kq_-1, kind=int64) * int(nik_mp, kind=int64) * int(dmp%wnst, kind=int64)**2 + &
      int((k_-1) * dmp%wnst**2, kind=int64) + &
      int((ibnd-dmp%estart)*dmp%wnst + (jbnd-dmp%estart) + 1, kind=int64)
 
    assert(idx >= 1 .and. idx <= dmp%num)
 
    res = real(ave_trans(idx), kind=real64)
    ! Pauli blocking
    if (optional_default(p_block, .true.)) then
      unocc = 1.0_real64 - dmp%adiabatic_st%occ(jbnd, kq) / dmp%adiabatic_st%smear%el_per_state
      res = res * max(unocc, 0.0_real64)
    end if
 
    pop_sub(get_trans_rate)
  end function get_trans_rate
 
  subroutine lindblad_operator_epw(dmp, ik, nik_skip, nresd, rho_diag, den_mat, l_mat)
    type(dmp_t),                  intent(in)  :: dmp
    integer,                      intent(in)  :: ik
    integer,                      intent(in)  :: nik_skip
    integer,                      intent(in)  :: nresd
    real(real64), contiguous,     intent(in)  :: rho_diag(:, :)
    complex(real64), contiguous,  intent(in)  :: den_mat(:, :)
    complex(real64), contiguous,  intent(out) :: l_mat(:, :)
 
    integer                     :: lst, nst, nik_mp, ikq, ibnd, jbnd, act_nst
    real(real64)                :: tr
    push_sub(lindblad_operator_epw)
 
    nst = dmp%adiabatic_st%nst
    nik_mp = dmp%adiabatic_st%nik - nik_skip
    act_nst = dmp%iend - dmp%istart + 1
 
    assert(ubound(rho_diag, dim = 1) == act_nst)
    assert(ubound(rho_diag, dim = 2) == dmp%adiabatic_st%nik)
    assert(ubound(den_mat, dim = 1) == 2*act_nst)
    assert(ubound(den_mat, dim = 2) == 2*act_nst)
    assert(ubound(l_mat, dim = 1) == 2*act_nst)
    assert(ubound(l_mat, dim = 2) == 2*act_nst)
 
    l_mat = 0.0_real64
 
    ! there are two types of damping: 1. mp-mp k-grid scattering; 2. highsympath-mp k-grid scattering
    ! in both case, ikq is the mp grid.
    do ikq = 1, nik_mp
      do jbnd = dmp%estart, dmp%eend
        do ibnd = dmp%estart, dmp%eend
 
          tr = get_trans_rate(dmp, nik_mp, jbnd, ibnd, ik, ikq, p_block=.true.)
          do lst = 1, act_nst+nresd
            l_mat(ibnd-dmp%istart+1, lst) = l_mat(ibnd-dmp%istart+1, lst) - m_half * tr * den_mat(ibnd-dmp%istart+1, lst)
            l_mat(lst, ibnd-dmp%istart+1) = l_mat(lst, ibnd-dmp%istart+1) - m_half * tr * den_mat(lst, ibnd-dmp%istart+1)
          end do
          l_mat(ibnd-dmp%istart+1, ibnd-dmp%istart+1) = l_mat(ibnd-dmp%istart+1, ibnd-dmp%istart+1) + &
            get_trans_rate(dmp, nik_mp, ibnd, jbnd, ikq, ik, p_block=.true.) * rho_diag(jbnd-dmp%istart+1, ikq)
        end do
      end do
    end do
 
    pop_sub(lindblad_operator_epw)
 
  end subroutine lindblad_operator_epw
 
  subroutine lifetime(dmp, ik, ibnd, nik_skip, gam)
    type(dmp_t),         intent(in)        :: dmp
    integer,             intent(in)        :: ik
    integer,             intent(in)        :: ibnd
    integer,             intent(in)        :: nik_skip
    real(real64),        intent(out)       :: gam
 
    integer                     :: nik_mp, ikq, jbnd
    real(real64)                :: tr
    push_sub(lifetime)
 
    assert(ibnd >= dmp%estart .and. ibnd <= dmp%eend)
    nik_mp = dmp%adiabatic_st%nik - nik_skip
 
    ! there are two types of damping: 1. mp-mp k-grid scattering; 2. highsympath-mp k-grid scattering
    ! in both case, ikq is the mp grid.
    gam = 0.0_real64
    do ikq = 1, nik_mp
      do jbnd = dmp%estart, dmp%eend
        tr = get_trans_rate(dmp, nik_mp, jbnd, ibnd, ik, ikq, p_block=.true.)
        gam = gam + tr
        tr = get_trans_rate(dmp, nik_mp, ibnd, jbnd, ikq, ik, p_block=.false.)
        gam = gam + tr * dmp%adiabatic_st%occ(jbnd, ikq) / dmp%adiabatic_st%smear%el_per_state
      end do
    end do
 
    pop_sub(lifetime)
 
  end subroutine lifetime
 
  subroutine get_gamma(dmp, ik, ibnd, nik_skip, gam_in, gam_out)
    type(dmp_t),         intent(in)        :: dmp
    integer,             intent(in)        :: ik
    integer,             intent(in)        :: ibnd
    integer,             intent(in)        :: nik_skip
    real(real64),        intent(out)       :: gam_in
    real(real64),        intent(out)       :: gam_out
 
    integer                     :: nik_mp, ikq, jbnd
    real(real64)                :: tr
    push_sub(get_gamma)
 
    assert(ibnd >= dmp%estart .and. ibnd <= dmp%eend)
    nik_mp = dmp%adiabatic_st%nik - nik_skip
 
    ! there are two types of damping: 1. mp-mp k-grid scattering; 2. highsympath-mp k-grid scattering
    ! in both case, ikq is the mp grid.
    gam_in = 0.0_real64
    gam_out = 0.0_real64
    do ikq = 1, nik_mp
      do jbnd = dmp%estart, dmp%eend
        tr = get_trans_rate(dmp, nik_mp, jbnd, ibnd, ik, ikq, p_block=.true.)
        gam_out = gam_out + tr
        tr = get_trans_rate(dmp, nik_mp, ibnd, jbnd, ikq, ik, p_block=.true.)
        gam_in = gam_in + tr * dmp%adiabatic_st%occ(jbnd, ikq)
      end do
    end do
 
    pop_sub(get_gamma)
 
  end subroutine get_gamma
 
  subroutine update_st(dmp, ik,  gr, namespace, nresd, overlap_ad_ks, overlap_resd_ks, nrm2_tdks, resd, st, rho_mat, pop)
    type(dmp_t),                  intent(inout) :: dmp
    integer,                      intent(in)    :: ik
    type(grid_t),                 intent(in)    :: gr
    type(namespace_t),            intent(in)    :: namespace
    integer,                      intent(in)    :: nresd
    complex(real64), contiguous,  intent(in)    :: overlap_ad_ks(:, :)
    complex(real64), contiguous,  intent(in)    :: overlap_resd_ks(:, :)
    real(real64), contiguous,     intent(in)    :: nrm2_tdks(:)
    type(states_elec_t),          intent(inout) :: resd
    type(states_elec_t),          intent(inout) :: st
    complex(real64), contiguous,  intent(inout) :: rho_mat(:, :)
    real(real64),                 intent(inout) :: pop
 
    real(real64), allocatable             :: occ_act(:)
    complex(real64), allocatable          :: overlap_ad_resd(:, :), overlap_act(:, :)
    integer                               :: ist, jst, nst, istart, iend, act_nst
 
    push_sub_with_profile(update_st)
 
    nst = dmp%adiabatic_st%nst
    istart = dmp%istart
    iend = dmp%iend
    act_nst = iend - istart + 1
 
    assert(ubound(overlap_ad_ks, dim = 1) == nst)
    assert(ubound(overlap_ad_ks, dim = 2) == nst)
    assert(ubound(overlap_resd_ks, dim = 1) == nst)
    assert(ubound(overlap_resd_ks, dim = 2) == nst)
    assert(ubound(rho_mat, dim = 1) == 2*act_nst)
    assert(ubound(rho_mat, dim = 2) == 2*act_nst)
    assert(is_hermitian(2*act_nst, rho_mat))
 
    safe_allocate(occ_act(1:act_nst+nresd))
 
    ! Pack the active blocks of the density matrix into a contiguous sub-matrix
    if (dmp%othn) then
      call lalg_eigensolve(act_nst+nresd, rho_mat, occ_act)
 
      call check_set_positive(act_nst+nresd, act_nst, ik, namespace, occ_act)
    else
      assert(nresd == act_nst)
 
      ! generalized eigenvalue problem
      safe_allocate(overlap_ad_resd(1:nst, 1:nst))
      safe_allocate(overlap_act(1:act_nst+nresd, 1:act_nst+nresd))
 
      ! ad-ad block
      overlap_act = m_z0
      do ist = 1, act_nst
        overlap_act(ist, ist) = m_z1
      end do
 
      ! ad-resd block
      call calc_overlap_state(ik, gr, dmp%adiabatic_st, resd, istart, iend, istart, iend, overlap_ad_resd)
      do jst = 1, act_nst
        do ist = 1, act_nst
          overlap_act(ist, act_nst+jst) = overlap_ad_resd(ist+istart-1, jst+istart-1)
          overlap_act(act_nst+jst, ist) = conjg(overlap_ad_resd(ist+istart-1, jst+istart-1))
        end do
      end do
 
      ! resd-resd block
      ! _i<d | d>_j = S^{d,\psi}_{ij} - sum_k S^{\phi,\psi}_{kj} S^{d,\phi}_{ik}
      overlap_act(act_nst+1:act_nst+nresd, act_nst+1:act_nst+nresd) = overlap_resd_ks(istart:iend, istart:iend)
      call blas_gemm('C', 'N', nresd, nresd, act_nst, -m_z1, overlap_ad_resd(istart, istart), nst, &
        overlap_ad_ks(istart, istart), nst, m_z1, overlap_act(act_nst+1, act_nst+1), act_nst+nresd)
 
      ! only the upper triangle of overlap is needed
      call lalg_geneigensolve(act_nst+nresd, rho_mat, overlap_act, occ_act, preserve_mat=.false.)
 
      call check_set_positive(act_nst+nresd, act_nst, ik, namespace, occ_act)
 
      safe_deallocate_a(overlap_ad_resd)
      safe_deallocate_a(overlap_act)
 
    end if
 
    if (dmp%unitary_transform) then
      call update_wfc_occ_procrustes(ik, gr, istart, iend, nresd, dmp%adiabatic_st, resd, overlap_ad_ks, overlap_resd_ks, &
        nrm2_tdks, occ_act, rho_mat, st, pop)
    else
      call update_wfc_occ(ik, gr, istart, iend, nresd, dmp%adiabatic_st, resd, occ_act, rho_mat, st, pop)
    end if
 
    pop = pop + (sum(st%occ(1:istart-1,ik) * nrm2_tdks(1:istart-1)) + &
      sum(st%occ(iend+1:nst,ik) * nrm2_tdks(iend+1:nst))) * st%kweights(ik)
 
    safe_deallocate_a(occ_act)
 
    pop_sub_with_profile(update_st)
 
  end subroutine update_st
 
  subroutine check_set_positive(locc, npos, ik, namespace, occ)
    integer,                    intent(in)       :: locc
    integer,                    intent(in)       :: npos
    integer,                    intent(in)       :: ik
    type(namespace_t),          intent(in)       :: namespace
    real(real64),               intent(inout)    :: occ(:)
 
    real(real64), parameter               :: small_eig = 1.0e-13_real64   
    real(real64)                          :: occ_min
    integer                               :: i, istart
 
    assert(ubound(occ, dim = 1) == locc)
 
    istart = locc - npos + 1
    occ_min = occ(istart)
 
    ! Single-pass memory traversal: find the minimum and floor negative values in-place
    do i = istart, locc
      if (occ(i) < occ_min) occ_min = occ(i)
      if (occ(i) < 0.0_real64) occ(i) = 0.0_real64
    end do
 
    ! Abort if unphysical negative eigenvalues exceed the acceptable noise threshold
    if (occ_min < -small_eig) then
      write(message(1), '(a, E9.2, a, I0)') 'Density matrix has negative eigenvalues of ', occ_min, ' at ik = ', ik
      call messages_fatal(1, namespace=namespace)
    end if
  end subroutine check_set_positive
 
  subroutine update_wfc_occ(ik, gr, istart, iend, nresd, ad_st, resd, occ, v_mat, st, pop)
    integer,                      intent(in)       :: ik
    type(grid_t),                 intent(in)       :: gr
    integer,                      intent(in)       :: istart
    integer,                      intent(in)       :: iend
    integer,                      intent(in)       :: nresd
    type(states_elec_t),          intent(in)       :: ad_st
    type(states_elec_t),          intent(in)       :: resd
    real(real64), contiguous,     intent(in)       :: occ(:)
    complex(real64), contiguous,  intent(in)       :: v_mat(:, :)
    type(states_elec_t),          intent(inout)    :: st
    real(real64),                 intent(inout)    :: pop
 
    complex(real64), allocatable  :: psi_j(:, :), phi(:, :)
    integer                       :: i, j, ib, i_minst, i_maxst, act_nst
 
    push_sub_with_profile(update_wfc_occ)
 
    act_nst = iend - istart + 1
 
    assert(ubound(v_mat, dim = 1) == 2*act_nst)
    assert(ubound(v_mat, dim = 2) == 2*act_nst)
 
    safe_allocate(psi_j(1:gr%np, st%d%dim))
    safe_allocate(phi(1:gr%np, st%d%dim))
 
    do j = act_nst+nresd, nresd+1, -1
      psi_j = (0.0_real64, 0.0_real64)
      ! ad block
      do ib = ad_st%group%block_start, ad_st%group%block_end
        i_minst = states_elec_block_min(ad_st, ib)
        i_maxst = states_elec_block_max(ad_st, ib)
        if (i_minst >= istart .and. i_maxst <= iend) then
          call zbatch_axpy_function(gr%np, v_mat(i_minst-istart+1:i_maxst-istart+1, j), &
            ad_st%group%psib(ib, ik), psi_j)
        else
          do i = max(i_minst, istart), min(i_maxst, iend)
            call states_elec_get_state(ad_st, gr, i, ik, phi)
            call lalg_axpy(gr%np, ad_st%d%dim, v_mat(i-istart+1, j), phi, psi_j)
          end do
        end if
      end do
      ! resd block
      if (nresd > 0) then
        do ib = resd%group%block_start, resd%group%block_end
          i_minst = states_elec_block_min(resd, ib)
          i_maxst = states_elec_block_max(resd, ib)
          if (i_minst >= istart .and. i_maxst <= istart+nresd-1) then
            call zbatch_axpy_function(gr%np, v_mat(act_nst+i_minst-istart+1:act_nst+i_maxst-istart+1, j), &
              resd%group%psib(ib, ik), psi_j)
          else
            do i = max(i_minst, istart), min(i_maxst, istart+nresd-1)
              call states_elec_get_state(resd, gr, i, ik, phi)
              call lalg_axpy(gr%np, resd%d%dim, v_mat(act_nst+i-istart+1, j), phi, psi_j)
            end do
          end if
        end do
      end if
 
      ! reset TDKS state
      call states_elec_set_state(st, gr, act_nst+nresd+istart-j, ik, psi_j)
    end do
 
    ! eigenvalues in ascending order, so inverse
    st%occ(istart:iend, ik) = occ(act_nst+nresd:nresd+1:-1)
    pop = pop + sum(st%occ(istart:iend, ik)) * st%kweights(ik)
 
    safe_deallocate_a(psi_j)
    safe_deallocate_a(phi)
    pop_sub_with_profile(update_wfc_occ)
  end subroutine update_wfc_occ
 
 
  subroutine update_wfc_occ_procrustes(ik, gr, istart, iend, nresd, ad_st, resd, overlap_ad_ks, overlap_resd_ks, &
    nrm2_tdks, occ_tilde, v_mat, st, pop)
    integer,                      intent(in)       :: ik
    type(grid_t),                 intent(in)       :: gr
    integer,                      intent(in)       :: istart
    integer,                      intent(in)       :: iend
    integer,                      intent(in)       :: nresd
    type(states_elec_t),          intent(in)       :: ad_st
    type(states_elec_t),          intent(in)       :: resd
    complex(real64), contiguous,  intent(in)       :: overlap_ad_ks(:, :)
    complex(real64), contiguous,  intent(in)       :: overlap_resd_ks(:, :)
    real(real64), contiguous,     intent(in)       :: nrm2_tdks(:)
    real(real64), contiguous,     intent(in)       :: occ_tilde(:)
    complex(real64), contiguous,  intent(inout)    :: v_mat(:, :)
    type(states_elec_t),          intent(inout)    :: st
    real(real64),                 intent(inout)    :: pop
 
    integer                       :: ist, jst, ib, i_minst, i_maxst, nst, act_nst
    complex(real64), allocatable  :: overlap_procrus(:, :), uproj(:, :), utrans(:, :)
    complex(real64), allocatable  :: uu(:, :), vt(:, :)
    complex(real64), allocatable  :: psi(:, :), phi(:, :)
    real(real64), parameter       :: small_occ = 5.0e-15_real64 ! for those zero occupations, assume them a small quantity
    real(real64)                  :: sg_values(1:iend-istart+1), rocc_tilde(1:iend-istart+1+nresd), rocc(1:iend-istart+1)
    real(real64)                  :: qtot_transform, nrm2, occ
 
    push_sub_with_profile(update_wfc_occ_procrustes)
 
    nst = ad_st%nst
    act_nst = iend - istart + 1
 
    assert(istart <= iend)
    assert(ubound(overlap_ad_ks, dim = 1) == nst)
    assert(ubound(overlap_ad_ks, dim = 2) == nst)
    assert(ubound(overlap_resd_ks, dim = 1) == nst)
    assert(ubound(overlap_resd_ks, dim = 2) == nst)
    assert(ubound(nrm2_tdks, dim = 1) == nst)
    assert(ubound(occ_tilde, dim = 1) == act_nst+nresd)
    assert(ubound(v_mat, dim = 1) == 2*act_nst)
    assert(ubound(v_mat, dim = 2) == 2*act_nst)
 
    ! set rocc for both TDKS and new wfcs. For those zero occupations, set them a small quantity
    rocc_tilde(1:act_nst+nresd) = sqrt(max(occ_tilde(1:act_nst+nresd), small_occ))
    rocc(1:act_nst) = sqrt(max(st%occ(istart:iend, ik), small_occ))
 
    safe_allocate(overlap_procrus(1:act_nst+nresd, 1:act_nst))
 
    ! Procrustes Overlap: S \propto (C_{top})^\dagger S_{ad\_ks} + (C_{bot})^\dagger S_{resd\_ks}
    ! Computed via in-place BLAS GEMM accumulation ('C' for conjugate transpose),
    ! avoiding the temporary memory allocation and copy overhead of native matmul.
    call blas_gemm('C', 'N', act_nst+nresd, act_nst, act_nst, m_z1, v_mat(1, 1), 2*act_nst, &
      overlap_ad_ks(istart, istart), nst, m_z0, overlap_procrus(1, 1), act_nst+nresd)
    call blas_gemm('C', 'N', act_nst+nresd, act_nst, nresd, m_z1, v_mat(1+act_nst, 1), 2*act_nst, &
      overlap_resd_ks(istart, istart), nst, m_z1, overlap_procrus(1, 1), act_nst+nresd)
    do jst = 1, act_nst
      do ist = 1, act_nst+nresd
        overlap_procrus(ist, jst) = overlap_procrus(ist, jst) * rocc_tilde(ist) * rocc(jst)
      end do
    end do
 
    safe_allocate(uproj(1:act_nst+nresd, 1:act_nst))
    safe_allocate(uu(1:act_nst+nresd, 1:act_nst+nresd))
    safe_allocate(vt(1:act_nst, 1:act_nst))
 
    call lalg_singular_value_decomp(act_nst+nresd, act_nst, overlap_procrus, uu, vt, sg_values)
    uproj = matmul(uu(:,1:act_nst), vt)
 
    safe_deallocate_a(overlap_procrus)
    safe_deallocate_a(vt)
    safe_deallocate_a(uu)
 
    safe_allocate(utrans(1:act_nst+nresd, 1:act_nst))
 
    ! update transformation matrix
    do ist = 1, act_nst+nresd
      call blas_scal(2*act_nst, cmplx(rocc_tilde(ist), 0.0, real64), v_mat(1, ist), 1)
    end do
    ! utrans = matmul(v_mat(1:act_nst+nresd,1:act_nst+nresd), uproj)
    call blas_gemm('N', 'N', act_nst+nresd, act_nst, act_nst+nresd, m_z1, v_mat(1, 1), 2*act_nst, &
      uproj(1, 1), act_nst+nresd, m_z0, utrans(1, 1), act_nst+nresd)
 
    safe_deallocate_a(uproj)
    safe_allocate(psi(1:gr%np, st%d%dim))
    safe_allocate(phi(1:gr%np, st%d%dim))
 
    ! update state and occupations
    qtot_transform = 0.0_real64
    do jst = istart, iend
      psi = (0.0_real64, 0.0_real64)
      ! ad basis
      do ib = ad_st%group%block_start, ad_st%group%block_end
        i_minst = states_elec_block_min(ad_st, ib)
        i_maxst = states_elec_block_max(ad_st, ib)
        if (i_minst >= istart .and. i_maxst <= iend) then
          call zbatch_axpy_function(gr%np, utrans(i_minst-istart+1:i_maxst-istart+1, jst-istart+1), &
            ad_st%group%psib(ib, ik), psi, nst=i_maxst-i_minst+1)
        else
          do ist = max(i_minst, istart), min(i_maxst, iend)
            call states_elec_get_state(ad_st, gr, ist, ik, phi)
            call lalg_axpy(gr%np, ad_st%d%dim, utrans(ist-istart+1, jst-istart+1), phi, psi)
          end do
        end if
      end do
      ! resd basis
      do ib = resd%group%block_start, resd%group%block_end
        i_minst = states_elec_block_min(resd, ib)
        i_maxst = states_elec_block_max(resd, ib)
        if (i_minst >= istart .and. i_maxst <= istart+nresd-1) then
          call zbatch_axpy_function(gr%np, utrans(act_nst+i_minst-istart+1:act_nst+i_maxst-istart+1, jst-istart+1), &
            resd%group%psib(ib, ik), psi, nst=i_maxst-i_minst+1)
        else
          do ist = max(i_minst, istart), min(i_maxst, istart+nresd-1)
            call states_elec_get_state(resd, gr, ist, ik, phi)
            call lalg_axpy(gr%np, st%d%dim, utrans(act_nst+ist-istart+1, jst-istart+1), phi, psi)
          end do
        end if
      end do
      nrm2 = real(zmf_dotp(gr, st%d%dim, psi, psi, reduce = .true.), real64)
      qtot_transform = qtot_transform + nrm2
      occ = nrm2 / nrm2_tdks(jst)
      st%occ(jst, ik) = occ
      call lalg_scal(gr%np, st%d%dim, 1.0_real64/sqrt(occ), psi)
      call states_elec_set_state(st, gr, jst, ik, psi)
    end do
 
    pop = pop + qtot_transform * st%kweights(ik)
 
    safe_deallocate_a(utrans)
    safe_deallocate_a(psi)
    safe_deallocate_a(phi)
 
    pop_sub_with_profile(update_wfc_occ_procrustes)
 
  end subroutine update_wfc_occ_procrustes
 
  logical function is_hermitian(n, mat)
    integer,                 intent(in) :: n
    complex(real64),         intent(in) :: mat(:, :)
    real(real64), parameter             :: tol = 1.0e-14_real64
 
    assert(ubound(mat, dim=1) == n)
    assert(ubound(mat, dim=2) == n)
 
    is_hermitian = maxval(abs(mat - transpose(conjg(mat)))) <= tol
  end function is_hermitian
 
  subroutine td_write_dm_proj_basis(this, out_dm_proj_basis, iter)
    class(dmp_t),        intent(in)    :: this
    type(c_ptr),         intent(inout) :: out_dm_proj_basis
    integer,             intent(in)    :: iter
 
    character(len=80) :: aux
    integer                       :: incount, ik, n_offd
    integer, allocatable          :: rdispls(:), recvcnts(:)
    complex(real64), allocatable  :: rho_offd(:, :)
 
    push_sub(td_write_dm_proj_basis)
 
    if (iter == 0) then
      if (this%adiabatic_st%system_grp%is_root()) then
        call td_write_print_header_init(out_dm_proj_basis)
 
        write(aux, '(a15,i8)') "# nspin        ", this%adiabatic_st%d%nspin
        call write_iter_string(out_dm_proj_basis, aux)
        call write_iter_nl(out_dm_proj_basis)
 
        write(aux, '(a15,i8)') "# nik          ", this%adiabatic_st%nik
        call write_iter_string(out_dm_proj_basis, aux)
        call write_iter_nl(out_dm_proj_basis)
 
        write(aux, '(a15,2i8)') "# istart, iend ", this%istart, this%iend
        call write_iter_string(out_dm_proj_basis, aux)
        call write_iter_nl(out_dm_proj_basis)
 
        write(aux, '(a)') "# loop sequence: "
        call write_iter_string(out_dm_proj_basis, aux)
        call write_iter_nl(out_dm_proj_basis)
        write(aux, '(a)') "# do ik = 1, nik "
        call write_iter_string(out_dm_proj_basis, aux)
        call write_iter_nl(out_dm_proj_basis)
        write(aux, '(a)') "#   do ist = istart, iend"
        call write_iter_string(out_dm_proj_basis, aux)
        call write_iter_nl(out_dm_proj_basis)
        write(aux, '(a)') "#     write E(ist, ik)"
        call write_iter_string(out_dm_proj_basis, aux)
        call write_iter_nl(out_dm_proj_basis)
 
        write(aux, '(a)') "# do ik = 1, nik "
        call write_iter_string(out_dm_proj_basis, aux)
        call write_iter_nl(out_dm_proj_basis)
        write(aux, '(a)') "#   do ist = istart, iend"
        call write_iter_string(out_dm_proj_basis, aux)
        call write_iter_nl(out_dm_proj_basis)
        write(aux, '(a)') "#     write occ(ist, ik)"
        call write_iter_string(out_dm_proj_basis, aux)
        call write_iter_nl(out_dm_proj_basis)
 
        write(aux, '(a)') "# do ik = 1, nik "
        call write_iter_string(out_dm_proj_basis, aux)
        call write_iter_nl(out_dm_proj_basis)
        write(aux, '(a)') "#   write Re(rho_off_diagonal(:)), Im(rho_off_diagonal(:))"
        call write_iter_string(out_dm_proj_basis, aux)
        call write_iter_nl(out_dm_proj_basis)
        write(aux, '(a)') "#   ! where 1D array rho_off_diagonal is packed as:"
        call write_iter_string(out_dm_proj_basis, aux)
        call write_iter_nl(out_dm_proj_basis)
        write(aux, '(a)') "#   ! do jst = istart, iend      ! column"
        call write_iter_string(out_dm_proj_basis, aux)
        call write_iter_nl(out_dm_proj_basis)
        write(aux, '(a)') "#   !   do ist = istart, jst-1   ! row"
        call write_iter_string(out_dm_proj_basis, aux)
        call write_iter_nl(out_dm_proj_basis)
 
        write(aux, '(a15)') "# w(ik)        "
        call write_iter_string(out_dm_proj_basis, aux)
        call write_iter_double(out_dm_proj_basis, this%adiabatic_st%kweights(1:this%adiabatic_st%nik), this%adiabatic_st%nik)
        call write_iter_nl(out_dm_proj_basis)
 
        call td_write_print_header_end(out_dm_proj_basis)
      end if
    else
 
      ! gather rho_offd
      n_offd = ubound(this%rho_offd, dim = 1)
 
      if (this%adiabatic_st%d%kpt%parallel) then
        if (this%adiabatic_st%d%kpt%mpi_grp%rank == 0) then
          safe_allocate(rho_offd(1:n_offd, this%adiabatic_st%nik))
        end if
        safe_allocate(recvcnts(1:this%adiabatic_st%d%kpt%mpi_grp%size))
        safe_allocate(rdispls(1:this%adiabatic_st%d%kpt%mpi_grp%size))
 
        incount = n_offd * this%adiabatic_st%d%kpt%nlocal()
        recvcnts(:) = n_offd * this%adiabatic_st%d%kpt%num(:)
 
        call mpi_displacements(recvcnts, rdispls)
        call this%adiabatic_st%d%kpt%mpi_grp%gatherv(this%rho_offd(:, this%adiabatic_st%d%kpt%start:this%adiabatic_st%d%kpt%end), &
          incount, mpi_double_complex, rho_offd, recvcnts, rdispls, mpi_double_complex, 0)
 
        safe_deallocate_a(recvcnts)
        safe_deallocate_a(rdispls)
      else
        safe_allocate_source(rho_offd(1:n_offd, this%adiabatic_st%nik), this%rho_offd)
      end if
 
      ! output
      if (this%adiabatic_st%system_grp%is_root()) then
        call write_iter_string(out_dm_proj_basis, "# ------------")
        call write_iter_start(out_dm_proj_basis)
        call write_iter_double(out_dm_proj_basis, this%adiabatic_st%qtot, 1)
        call write_iter_nl(out_dm_proj_basis)
        do ik = 1, this%adiabatic_st%nik
          call write_iter_double(out_dm_proj_basis, this%adiabatic_st%eigenval(this%istart:this%iend, ik), this%iend-this%istart+1)
          call write_iter_nl(out_dm_proj_basis)
        end do
        do ik = 1, this%adiabatic_st%nik
          call write_iter_double(out_dm_proj_basis, this%adiabatic_st%occ(this%istart:this%iend, ik), this%iend-this%istart+1)
          call write_iter_nl(out_dm_proj_basis)
        end do
        do ik = 1, this%adiabatic_st%nik
          call write_iter_double(out_dm_proj_basis, real(rho_offd(1:n_offd, ik)), n_offd)
          call write_iter_double(out_dm_proj_basis, aimag(rho_offd(1:n_offd, ik)), n_offd)
          call write_iter_nl(out_dm_proj_basis)
        end do
      end if
 
      safe_deallocate_a(rho_offd)
    end if
 
    pop_sub(td_write_dm_proj_basis)
  end subroutine td_write_dm_proj_basis
 
#ifdef HAVE_MPI
  subroutine slmpi_create_shared_memory_window(number_of_elements, intranode_grp, window, array)
    use iso_c_binding
    integer(int64),        intent(in)  :: number_of_elements
    type(mpi_grp_t),       intent(in)  :: intranode_grp
    type(mpi_win),         intent(out) :: window
    real(real32), pointer, intent(out) :: array(:)
 
    type(c_ptr) :: ptr
    integer(kind=MPI_ADDRESS_KIND) :: window_size
    integer :: disp_unit
 
    assert(not_in_openmp())
 
    ! allocate only on rank 0 of each node
    if (intranode_grp%rank == 0) then
      window_size = number_of_elements * 4_mpi_address_kind
    else
      window_size = 0_mpi_address_kind
    end if
    assert(number_of_elements * 4 < huge(0_mpi_address_kind))
    assert(window_size >= 0)
    disp_unit = 4
    call mpi_win_allocate_shared(window_size, disp_unit, mpi_info_null, &
      intranode_grp%comm, ptr, window)
    ! get pointer on all ranks
    if (intranode_grp%rank /= 0) then
      call mpi_win_shared_query(window, 0, window_size, disp_unit, ptr)
    end if
    ! get fortran pointer
    call c_f_pointer(ptr, array, [number_of_elements])
 
    ! start access epoch
    call mpi_win_lock_all(mpi_mode_nocheck, window)
 
  end subroutine slmpi_create_shared_memory_window
#endif
 
end module dm_propagation_oct_m