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 debug_oct_m
  use electron_space_oct_m
  use exponential_oct_m
  use global_oct_m
  use grid_oct_m
  use io_function_oct_m
  use, intrinsic :: iso_fortran_env
  use hamiltonian_elec_oct_m
  use ions_oct_m
  use kpoints_oct_m
  use lalg_basic_oct_m
  use lalg_adv_oct_m
  use math_oct_m
  use mesh_oct_m
  use mesh_function_oct_m
  use messages_oct_m
  use multicomm_oct_m
  use mpi_oct_m
  use namespace_oct_m
  use output_low_oct_m
  use parser_oct_m
  use profiling_oct_m
  use restart_oct_m
  use space_oct_m
  use states_elec_oct_m
  use states_elec_restart_oct_m
  use unit_oct_m
  use unit_system_oct_m
  use varinfo_oct_m
  use lapack_oct_m
  use eigensolver_oct_m
  use propagation_ops_elec_oct_m
  use interaction_partner_oct_m
  use v_ks_oct_m
  use density_oct_m
  use batch_ops_oct_m
  use states_elec_parallel_oct_m
  use energy_calc_oct_m
 
  implicit none
 
  private
 
  public ::                             &
    dmp_t,                              &
    dm_propagation_init_run,            &
    dm_propagation_run
 
  type dmp_t
    
    integer                   :: calculation_mode
    integer                   :: basis
    logical                   :: unitary_transform
    real(real64)              :: t(2)
    ! Initialised from GS KS states, so included in this class so adiabatic_st can be returned
    ! from td_init_run. Note, one could alternatively do in dm_propagation_run:
    !  if(iter==1) call states_elec_copy(dmp%adiabatic_st, st, special=.true.)
    ! as I am not sure if the restart calls are required.
    type(states_elec_t)       :: adiabatic_st
    !
  contains
    procedure :: init => dmp_init
  end type dmp_t
 
 
contains
 
  subroutine dm_propagation_init_run(adiabatic_st, namespace, space, gr, st, hm, mc)
    type(states_elec_t),      intent(inout)   :: adiabatic_st
    type(namespace_t),        intent(in)      :: namespace
    type(electron_space_t),   intent(in)      :: space
    type(grid_t),             intent(in)      :: gr
    type(states_elec_t),      intent(in)      :: st
    type(hamiltonian_elec_t), intent(in)      :: hm
    type(multicomm_t),        intent(in)      :: mc
 
    integer         :: ierr
    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(adiabatic_st, st, special=.true.)
    call restart_init(restart_load, namespace, restart_gs, restart_type_load, mc, ierr, mesh=gr, exact =.true.)
    if (ierr == 0) then
      call states_elec_load(restart_load, namespace, space, adiabatic_st, gr, hm%kpoints, &
        ierr, label = ': Adiabatic states')
    end if
    call restart_end(restart_load)
 
    pop_sub(dm_propagation_init_run)
 
  end subroutine dm_propagation_init_run
 
 
  subroutine dmp_init(this, namespace, st)
    class(dmp_t),        intent(out) :: this
    type(namespace_t),   intent(in)  :: namespace
    type(states_elec_t), intent(in)  :: st
 
    integer             :: cols_t1t2_block
    type(block_t)       :: blk
 
    push_sub(dmp_init)
 
    !%Variable TDDMPropagation_2Times
    !%Type block
    !%Section Time-Dependent
    !%Description
    !% Two times approximation for the jump operator in the master equation.
    !%
    !% <tt>%TDDMPropagation_2Times
    !% <br>&nbsp;&nbsp; t1 | t2
    !% <br>%</tt>
    !%End
 
    cols_t1t2_block = 0
    !TODO: maybe define some more reasonable defaults
    this%t(1) = m_zero
    this%t(2) = m_one
 
    if (parse_block(namespace, 'TDDMPropagation_2Times', blk) == 0) then
      cols_t1t2_block = parse_block_cols(blk, 0)
      if (cols_t1t2_block == 2) then
        call parse_block_float(blk, 0, 0, this%t(1), units_inp%time)
        call parse_block_float(blk, 0, 1, this%t(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%t(1) , ', ', this%t(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
    end if
 
    !%Variable TDDMPropagationBasis
    !%Type integer
    !%Default Adiabatic
    !%Section Time-Dependent
    !%Description
    !% Decides the basis set for the density matrix propagation.
    !%Option Adiabatic 01
    !% Istantaneous 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 TDDMUnitaryTransformFix
    !%Type logical
    !%Default no
    !%Section Time-Dependent
    !%Description
    !% Applies an additional unitary transformation to the damped wavefunctions
    !% to maximize their overlap with the reference wavefunctions before damping.
    !%End
    call parse_variable(namespace, 'TDDMUnitaryTransformFix', .false., this%unitary_transform)
 
    !Sanity checks
    if (all(st%occ >= m_min_occ)) then
      message(1) = "Warning: TDDMPropagation requires a number of empty states."
      message(2) = "         Add ExtraStates and rerun."
      call messages_fatal(2, namespace=namespace)
    else 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 == spinors) then
      call messages_not_implemented('TDDMPropagation with spinors', namespace=namespace)
    end if
 
    pop_sub(dmp_init)
 
  end subroutine dmp_init
 
 
  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-10_real64
    type(eigensolver_t)           :: eigens
    real(real64)                  :: population(3), leak
    integer                       :: ik
    logical                       :: update_energy_
    complex(real64), allocatable  :: resd(:, :, :), rho_mat(:, :), conj_psi_phi(:, :)
 
    push_sub_with_profile(dm_propagation_run)
 
    ! Update the hamiltonian. hm is updated after each td propagation.
    ! But ions may move afterthen
    call hm%update(gr, namespace, space, ext_partners, time=iter * dt)
 
    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
 
    ! Compute everything
    safe_allocate(resd(1:gr%np, 1:st%d%dim, 1:st%nst))
    safe_allocate(rho_mat(1:2*st%nst, 1:2*st%nst))
    safe_allocate(conj_psi_phi(1:st%nst, 1:2*st%nst)) ! (i, j) = <\phi, d_j | \psi_i >
 
    population = 0.0_real64
 
    do ik = st%d%kpt%start, st%d%kpt%end
      population(1) = population(1) + sum(st%occ(:, ik) * st%kweights(ik))
 
      if (debug%info) then
        call orthogonality_check_ks(ik, st, gr)
      end if
 
      call construct_residuals(dmp%adiabatic_st, ik, st, gr, conj_psi_phi, resd)
 
      call construct_density_matrix(ik, st, conj_psi_phi, rho_mat)
 
      ! population before damping
      call population_in_adiabatic(ik, st, rho_mat, population(2))
 
      call dissipation(dmp,dmp%adiabatic_st, st, ik, dt, rho_mat)
 
      call update_st(dmp, ik, gr, resd, st, rho_mat)
 
      ! population after damping
      population(3) = population(3) + sum(st%occ(:, ik))*st%kweights(ik)
    end do
 
    safe_deallocate_a(resd)
    safe_deallocate_a(rho_mat)
    safe_deallocate_a(conj_psi_phi)
 
    call st%d%kpt%mpi_grp%allreduce_inplace(population, 3, mpi_double_precision, mpi_sum)
 
    write(message(1), '(a,E15.8,a,E15.8,a,E15.8,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
 
    ! 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 orthogonality_check_ks(ik, st, gr)
    integer, intent(in)               :: ik
    type(states_elec_t), intent(in)   :: st
    type(grid_t), intent(in)          :: gr
 
    integer                               :: ist, jst, lst
    real(real64)                          :: dqtot
    real(real64)                          :: occ(1:st%nst)
    real(real64), parameter               :: ZERO = 1.0e-12_real64
    complex(real64), allocatable          :: psii(:, :), psij(:, :), overlap(:, :), rho_mat(:, :)
 
    push_sub(orthogonality_check_ks)
 
    safe_allocate(psii(1:gr%np, 1:st%d%dim))
    safe_allocate(psij(1:gr%np, 1:st%d%dim))
    safe_allocate(overlap(1:st%nst, 1:st%nst))
    safe_allocate(rho_mat(1:st%nst, 1:st%nst))
 
    do jst = 1, st%nst
      call states_elec_get_state(st, gr, jst, ik, psij)
      overlap(jst, jst) = zmf_dotp(gr, st%d%dim, psij, psij, reduce = .false.)
      do ist = jst + 1, st%nst
        call states_elec_get_state(st, gr, ist, ik, psii)
        overlap(ist, jst) = zmf_dotp(gr, st%d%dim, psii, psij, reduce = .false.)
      end do
    end do
    call lower_triangular_to_hermitian(st%nst, overlap)
    call gr%allreduce(overlap)
 
 
    safe_deallocate_a(psii)
    safe_deallocate_a(psij)
 
    do jst = 1, st%nst
      do ist = jst, st%nst
        rho_mat(ist, jst) = overlap(ist, 1) * overlap(1, jst) * st%occ(1, ik)
        do lst = 2, st%nst
          rho_mat(ist, jst) = rho_mat(ist, jst) + overlap(ist, lst) * overlap(lst, jst) * st%occ(lst, ik)
        end do
      end do
    end do
    call lower_triangular_to_hermitian(st%nst, rho_mat)
 
 
    call lalg_geneigensolve(st%nst, rho_mat, overlap, occ, preserve_mat=.false.)
    dqtot = sum(st%occ(:, ik)) - sum(occ)
    if (dqtot > zero) then
      write(message(1),'(a,E12.5,a,i4)') 'TDKS wavefunctions is not fullly orthonormalized with ', dqtot,  &
        'electrons difference at ik= ', ik
      call messages_info(1,all_nodes=.true.)
    end if
 
    safe_deallocate_a(overlap)
    safe_deallocate_a(rho_mat)
 
    pop_sub(orthogonality_check_ks)
  end subroutine orthogonality_check_ks
 
  subroutine construct_residuals(adiabatic_st, ik, st, gr, conj_psi_phi, resd)
    type(states_elec_t), intent(in)   :: adiabatic_st
    integer, intent(in)               :: ik
    type(states_elec_t), intent(in)   :: st
    type(grid_t), intent(in)          :: gr
    complex(real64), intent(out)      :: conj_psi_phi(:, :)
    complex(real64), intent(out)      :: resd(:, :, :)
 
    integer                       :: ist, jst
    complex(real64), allocatable  :: psi(:, :), phi(:, :)
 
    push_sub(construct_residuals)
 
    assert(ubound(conj_psi_phi, dim = 1) == st%nst)
    assert(ubound(conj_psi_phi, dim = 2) == 2*st%nst)
    assert(ubound(resd, dim = 1) == gr%np)
    assert(ubound(resd, dim = 2) == st%d%dim)
    assert(ubound(resd, dim = 3) == st%nst)
 
    safe_allocate(psi(1:gr%np, 1:st%d%dim))
    safe_allocate(phi(1:gr%np, 1:st%d%dim))
 
    do jst = 1, st%nst
      call states_elec_get_state(adiabatic_st, gr, jst, ik, phi)
      do ist = 1, st%nst
        call states_elec_get_state(st, gr, ist, ik, psi)
        conj_psi_phi(ist, jst) = zmf_dotp(gr, st%d%dim, phi, psi, reduce = .false.)
      end do
    end do
    call gr%allreduce(conj_psi_phi(:, 1:st%nst))
 
    do jst = 1, st%nst
      call states_elec_get_state(st, gr, jst, ik, resd(:, :, jst))
      do ist = 1, st%nst
        call states_elec_get_state(adiabatic_st, gr, ist, ik, phi)
        call lalg_axpy(gr%np, st%d%dim, -conj_psi_phi(jst, ist), phi, resd(:, :, jst))
      end do
    end do
 
    do ist = 1, st%nst
      call states_elec_get_state(st, gr, ist, ik, psi)
      do jst = 1+st%nst, 2*st%nst
        conj_psi_phi(ist, jst) = zmf_dotp(gr, st%d%dim, resd(:, :, jst-st%nst), psi, reduce = .false.)
      end do
    end do
    call gr%allreduce(conj_psi_phi(:, 1+st%nst:2*st%nst))
 
    safe_deallocate_a(psi)
    safe_deallocate_a(phi)
 
    pop_sub(construct_residuals)
  end subroutine construct_residuals
 
  subroutine construct_density_matrix(ik, st, conj_psi_phi, rho_mat)
    integer, intent(in)               :: ik
    type(states_elec_t), intent(in)   :: st
    complex(real64), intent(in)       :: conj_psi_phi(:, :)
    complex(real64), intent(out)      :: rho_mat(:, :)
 
    integer                       :: ist, jst, lst
 
    push_sub(construct_density_matrix)
 
    assert(ubound(conj_psi_phi, dim = 1) == st%nst)
    assert(ubound(conj_psi_phi, dim = 2) == 2*st%nst)
    assert(ubound(rho_mat, dim = 1) == 2*st%nst)
    assert(ubound(rho_mat, dim = 2) == 2*st%nst)
 
    do jst = 1, 2*st%nst
      do ist = jst, 2*st%nst
        rho_mat(ist, jst) = conj_psi_phi(1, ist) * conjg(conj_psi_phi(1, jst)) * st%occ(1, ik)
        do lst = 2, st%nst
          rho_mat(ist, jst) = rho_mat(ist, jst) + conj_psi_phi(lst, ist) * conjg(conj_psi_phi(lst, jst)) * st%occ(lst, ik)
        end do
      end do
    end do
    call lower_triangular_to_hermitian(2*st%nst, rho_mat)
 
    pop_sub(construct_density_matrix)
  end subroutine construct_density_matrix
 
  subroutine population_in_adiabatic(ik, st, rho_mat, pop)
    integer, intent(in)               :: ik
    type(states_elec_t), intent(in)   :: st
    complex(real64), intent(in)       :: rho_mat(:, :)
    real(real64),intent(inout)        :: pop
 
    integer                       :: ist
    real(real64)                  :: qtot_basis
 
    push_sub(population_in_adiabatic)
 
    assert(ubound(rho_mat, dim = 1) == 2*st%nst)
    assert(ubound(rho_mat, dim = 2) == 2*st%nst)
 
    qtot_basis = m_zero
    do ist = 1, st%nst
      qtot_basis = qtot_basis + real(rho_mat(ist, ist), real64)
    end do
    pop = pop + qtot_basis * st%kweights(ik)
 
    pop_sub(population_in_adiabatic)
  end subroutine population_in_adiabatic
 
  subroutine update_st(dmp, ik,  gr, resd, st, rho_mat)
    type(dmp_t), intent(in)               :: dmp
    integer, intent(in)                   :: ik
    type(grid_t), intent(in)              :: gr
    complex(real64), intent(in)           :: resd(:, :, :)
    type(states_elec_t), intent(inout)    :: st
    complex(real64), intent(inout)        :: rho_mat(:, :)
 
    real(real64), parameter               :: ZERO = 1.0e-12_real64
    real(real64)                          :: occ(1:2*st%nst)
    integer                               :: ist, jst
    complex(real64), allocatable          :: phii(:, :), phij(:, :)
    complex(real64), allocatable          :: overlap(:, :)
 
    push_sub(update_st)
 
    assert(ubound(resd, dim = 1) == gr%np)
    assert(ubound(resd, dim = 2) == st%d%dim)
    assert(ubound(resd, dim = 3) == st%nst)
    assert(ubound(rho_mat, dim = 1) == 2*st%nst)
    assert(ubound(rho_mat, dim = 2) == 2*st%nst)
 
    safe_allocate(overlap(1:2*st%nst, 1:2*st%nst))
 
    if (maxval(abs(rho_mat - transpose(conjg(rho_mat)))) > zero) then
      write(message(1),'(a,E12.5,a,i4)') "rho_mat is not Hermitian! Max deviation = ", &
        maxval(abs(rho_mat - transpose(conjg(rho_mat)))), ' at ik= ', ik
      call messages_fatal(1)
    end if
 
    safe_allocate(phii(1:gr%np, 1:st%d%dim))
    safe_allocate(phij(1:gr%np, 1:st%d%dim))
 
    do jst = 1, 2*st%nst
      if (jst <= st%nst) then
        call states_elec_get_state(dmp%adiabatic_st, gr, jst, ik, phij)
      else
        phij = resd(:, :, jst-st%nst)
      end if
      overlap(jst, jst) = zmf_dotp(gr, st%d%dim, phij, phij, reduce = .false.)
 
      do ist = jst + 1, 2*st%nst
        if (ist <= st%nst) then
          call states_elec_get_state(dmp%adiabatic_st, gr, ist, ik, phii)
        else
          phii = resd(:, :, ist-st%nst)
        end if
 
        overlap(ist, jst) = zmf_dotp(gr, st%d%dim, phii, phij, reduce = .false.)
 
      end do
    end do
    call lower_triangular_to_hermitian(2*st%nst, overlap)
    call gr%allreduce(overlap)
 
    safe_deallocate_a(phii)
    safe_deallocate_a(phij)
 
    call lalg_geneigensolve(2*st%nst, rho_mat, overlap, occ, preserve_mat=.false.)
 
    safe_deallocate_a(overlap)
 
    if (dmp%unitary_transform) then
      call update_wfc_occ_procrustes(dmp%adiabatic_st, ik, st, gr, resd, rho_mat, occ)
    else
      call update_wfc_occ(dmp%adiabatic_st, ik, st, gr, resd, rho_mat, occ)
    end if
 
    pop_sub(update_st)
  end subroutine update_st
 
  subroutine update_wfc_occ(adiabatic_st, ik, st, gr, resd, rho_mat, occ)
    type(states_elec_t), intent(in)       :: adiabatic_st
    integer, intent(in)                   :: ik
    type(states_elec_t), intent(inout)    :: st
    type(grid_t), intent(in)              :: gr
    complex(real64), intent(in)           :: resd(:, :, :)
    complex(real64), intent(in)           :: rho_mat(:, :)
    real(real64), intent(in)              :: occ(1:2*st%nst)
 
    integer                       :: ist, jst
    complex(real64), allocatable  :: psi(:, :), phij(:, :)
 
    push_sub(update_wfc_occ)
 
    assert(ubound(resd, dim = 1) == gr%np)
    assert(ubound(resd, dim = 2) == st%d%dim)
    assert(ubound(resd, dim = 3) == st%nst)
    assert(ubound(rho_mat, dim = 1) == 2*st%nst)
    assert(ubound(rho_mat, dim = 2) == 2*st%nst)
 
    safe_allocate(psi(1:gr%np, 1:st%d%dim))
    safe_allocate(phij(1:gr%np, 1:st%d%dim))
 
    ! update wave functions
    do ist = 1, st%nst
      psi = m_z0
      do jst = 1, st%nst
        call states_elec_get_state(adiabatic_st, gr, jst, ik, phij)
        call lalg_axpy(gr%np, st%d%dim, rho_mat(jst, ist+st%nst), phij, psi)
      end do
      do jst = 1+st%nst, 2*st%nst
        call lalg_axpy(gr%np, st%d%dim, rho_mat(jst, ist+st%nst), resd(:, :, jst-st%nst), psi)
      end do
      call zmf_normalize(gr, st%d%dim, psi)
      call states_elec_set_state(st, gr, ist, ik, psi)
    end do
 
    ! update occupation
    st%occ(1:st%nst, ik) = occ(1+st%nst:2*st%nst)
 
    safe_deallocate_a(psi)
    safe_deallocate_a(phij)
 
    pop_sub(update_wfc_occ)
  end subroutine update_wfc_occ
 
  subroutine update_wfc_occ_procrustes(adiabatic_st, ik, st, gr, resd, rho_mat, occ)
    type(states_elec_t), intent(in)       :: adiabatic_st
    integer, intent(in)                   :: ik
    type(states_elec_t), intent(inout)    :: st
    type(grid_t), intent(in)              :: gr
    complex(real64), intent(in)           :: resd(:, :, :)
    complex(real64), intent(in)           :: rho_mat(:, :)
    real(real64), intent(in)              :: occ(1:2*st%nst)
 
    integer                       :: ist, jst, lst
    complex(real64)               :: uji
    complex(real64), allocatable  :: overlap_procrus(:, :), uproj(:, :)
    complex(real64), allocatable  :: uu(:, :), vt(:, :)
    real(real64)                  :: sg_values(1:st%nst), norm
    real(real64), parameter       :: ZERO = 1.0e-12_real64
    complex(real64), allocatable  :: psi(:, :), psi_tilde(:, :), phi(:, :)
 
    push_sub(update_wfc_occ_procrustes)
 
    assert(ubound(resd, dim = 1) == gr%np)
    assert(ubound(resd, dim = 2) == st%d%dim)
    assert(ubound(resd, dim = 3) == st%nst)
    assert(ubound(rho_mat, dim = 1) == 2*st%nst)
    assert(ubound(rho_mat, dim = 2) == 2*st%nst)
 
    safe_allocate(psi(1:gr%np, 1:st%d%dim))
    safe_allocate(psi_tilde(1:gr%np, 1:st%d%dim))
    safe_allocate(phi(1:gr%np, 1:st%d%dim))
    safe_allocate(overlap_procrus(1:2*st%nst,1:st%nst))
    safe_allocate(uproj(1:2*st%nst, 1:st%nst))
    safe_allocate(uu(1:2*st%nst, 1:2*st%nst))
    safe_allocate(vt(1:st%nst, 1:st%nst))
 
    do ist = 1, 2*st%nst
      psi_tilde = m_z0
      do jst = 1, st%nst
        call states_elec_get_state(adiabatic_st, gr, jst, ik, phi)
        call lalg_axpy(gr%np, st%d%dim, rho_mat(jst, ist), phi, psi_tilde)
      end do
      do jst = 1+st%nst, 2*st%nst
        call lalg_axpy(gr%np, st%d%dim, rho_mat(jst, ist), resd(:, :, jst-st%nst), psi_tilde)
      end do
      call zmf_normalize(gr, st%d%dim, psi_tilde)
 
      do jst = 1, st%nst
        call states_elec_get_state(st, gr, jst, ik, psi)
        overlap_procrus(ist, jst) = zmf_dotp(gr, st%d%dim, psi_tilde, psi, reduce = .false.)
      end do
    end do
    call gr%allreduce(overlap_procrus)
 
    call lalg_singular_value_decomp(2*st%nst, st%nst, overlap_procrus, uu, vt, sg_values)
    uproj = matmul(uu(:,1:st%nst), vt)
 
    safe_deallocate_a(overlap_procrus)
    safe_deallocate_a(uu)
    safe_deallocate_a(vt)
 
    ! update wavefunctions and occupation
    do ist = 1, st%nst
      psi = m_z0
      psi_tilde = m_z0
      do jst = 1, 2*st%nst
        if (occ(jst) > zero) then
          uji = uproj(jst, ist)*sqrt(occ(jst))
          do lst = 1, st%nst
            call states_elec_get_state(adiabatic_st, gr, lst, ik, phi)
            call lalg_axpy(gr%np, st%d%dim, uproj(jst, ist)*rho_mat(lst, jst), phi, psi)
            call lalg_axpy(gr%np, st%d%dim, uji*rho_mat(lst, jst), phi, psi_tilde)
          end do
          do lst = 1+st%nst, 2*st%nst
            call lalg_axpy(gr%np, st%d%dim, uproj(jst, ist)*rho_mat(lst, jst), resd(:, :, lst-st%nst), psi)
            call lalg_axpy(gr%np, st%d%dim, uji*rho_mat(lst, jst), resd(:, :, lst-st%nst), psi_tilde)
          end do
        else
          do lst = 1, st%nst
            call states_elec_get_state(adiabatic_st, gr, lst, ik, phi)
            call lalg_axpy(gr%np, st%d%dim, uproj(jst, ist)*rho_mat(lst, jst), phi, psi)
          end do
          do lst = 1+st%nst, 2*st%nst
            call lalg_axpy(gr%np, st%d%dim, uproj(jst, ist)*rho_mat(lst, jst), resd(:, :, lst-st%nst), psi)
          end do
        end if
      end do
      call zmf_normalize(gr, st%d%dim, psi)
      call states_elec_set_state(st, gr, ist, ik, psi)
      norm = zmf_nrm2(gr, st%d%dim, psi_tilde, reduce = .true.)**2
      st%occ(ist, ik) = norm
    end do
 
    safe_deallocate_a(psi)
    safe_deallocate_a(psi_tilde)
    safe_deallocate_a(phi)
    safe_deallocate_a(uproj)
 
    pop_sub(update_wfc_occ_procrustes)
  end subroutine update_wfc_occ_procrustes
 
  subroutine dissipation(dmp, bst, st, ik, dt, nn)
    type(dmp_t), intent(in) :: dmp
    type(states_elec_t), intent(in) :: bst
    type(states_elec_t), intent(in) :: st
    integer, intent(in) :: ik
    real(real64), intent(in) :: dt
    complex(real64), intent(inout) :: nn(:,:)
 
    complex(real64), allocatable :: DMatrix(:,:,:,:), nn_current(:,:)
    real(real64) :: trace_before_dissipation, trace_after_dissipation, difference
    real(real64) :: tol = 1.0e-10_real64
    complex(real64) :: gamma
    integer :: i, a, b, c, d, m, n
    integer, parameter :: order = 4
    real(real64) :: t2
 
    push_sub_with_profile(dissipation)
    safe_allocate(nn_current(1:bst%nst, 1:bst%nst))
    safe_allocate(dmatrix(1:bst%nst, 1:bst%nst, 1:bst%nst, 1:bst%nst))
    ! t2 is same for all k-points
    t2 = dmp%t(2)*sqrt(real(st%nik, real64)*st%nst/2)
 
    dmatrix = m_z0
    do m = bst%st_start, bst%st_end
      gamma = bst%occ(m, ik)/t2
      do n = bst%st_start, bst%st_end
        if  (m /= n) then
          dmatrix(m, m, n, n) = dmatrix(m,m,n,n) + gamma
          do a = bst%st_start, bst%st_end
            dmatrix(a, n, a, n) = dmatrix(a, n, a, n) - m_half*gamma
            dmatrix(n, a, n, a) = dmatrix(n, a, n, a) - m_half*gamma
          end do ! a
        end if
      end do ! n
    end do ! m
 
    trace_before_dissipation = 0.0_real64
    do a = 1, bst%nst
      trace_before_dissipation  = trace_before_dissipation + real(nn(a, a))
    end do
 
    do i = 1, order
      nn_current = m_z0
      do d= 1, bst%nst
        do c= 1, bst%nst
          do b= 1, bst%nst
            do a = 1, bst%nst
              nn_current(a,b) = nn_current(a,b) + dmatrix(a,b,c,d) * nn(c,d)
            end do
          end do
        end do
      end do
      ! Add to output nn
      do a = 1, bst%nst
        do b = 1, bst%nst
          nn(a,b) = nn(a,b) + (nn_current(a,b)*dt)/ factorial(i)
        end do
      end do
    end do
 
    call symmetrize_matrix(bst%nst, nn)
 
    trace_after_dissipation =  0.0_real64
    do a = 1, bst%nst
      trace_after_dissipation =trace_after_dissipation +  real(nn(a, a))
    end do
 
    difference = trace_after_dissipation - trace_before_dissipation
    if (abs(difference) > tol) then
      write(message(1), '(A, ES12.5)') "Trace difference is non-zero in multiply_dissipator!  = ", abs(difference)
      call messages_fatal(1)
    end if
 
    safe_deallocate_a(nn_current)
    safe_deallocate_a(dmatrix)
 
    pop_sub_with_profile(dissipation)
 
  end subroutine dissipation
 
end module dm_propagation_oct_m