propagator_rk.F90

Go to the documentation of this file.

!! Copyright (C) 2002-2006 M. Marques, A. Castro, A. Rubio, G. Bertsch
!!
!! 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 propagator_rk_oct_m
  use batch_ops_oct_m
  use comm_oct_m
  use debug_oct_m
  use density_oct_m
  use electron_space_oct_m
  use ext_partner_list_oct_m
  use forces_oct_m
  use gauge_field_oct_m
  use grid_oct_m
  use global_oct_m
  use hamiltonian_elec_oct_m
  use interaction_partner_oct_m
  use ion_dynamics_oct_m
  use ions_oct_m
  use, intrinsic :: iso_fortran_env
  use ks_potential_oct_m
  use lalg_basic_oct_m
  use lda_u_oct_m
  use mesh_oct_m
  use mesh_function_oct_m
  use messages_oct_m
  use mpi_oct_m
  use namespace_oct_m
  use oct_exchange_oct_m
  use opt_control_state_oct_m
  use parser_oct_m
  use perturbation_ionic_oct_m
  use poisson_oct_m
  use potential_interpolation_oct_m
  use profiling_oct_m
  use propagator_base_oct_m
  use space_oct_m
  use sparskit_oct_m
  use states_elec_oct_m
  use v_ks_oct_m
  use propagation_ops_elec_oct_m
  use xc_oct_m
 
  implicit none
 
  private
 
  public ::                    &
    td_explicit_runge_kutta4,  &
    td_runge_kutta2,           &
    td_runge_kutta4,           &
    propagator_rk_end
 
  class(mesh_t),            pointer     :: mesh_p
  type(hamiltonian_elec_t), pointer     :: hm_p
  type(states_elec_t),      pointer     :: st_p
  type(propagator_base_t),  pointer     :: tr_p
  type(namespace_t),        pointer     :: namespace_p
  type(electron_space_t),   pointer     :: space_p
  type(partner_list_t),     pointer     :: ext_partners_p
  integer                               :: dim_op
  real(real64)                          :: t_op, dt_op
  real(real64),             allocatable :: vpsl1_op(:), vpsl2_op(:)
  logical :: move_ions_op
  type(xc_copied_potentials_t) :: vhxc1_op, vhxc2_op
 
contains
 
  subroutine td_explicit_runge_kutta4(ks, namespace, space, hm, gr, st, time, dt, ions_dyn, ions, ext_partners, qcchi)
    type(v_ks_t),                        target, intent(inout) :: ks
    type(namespace_t),                           intent(in)    :: namespace
    type(electron_space_t),                      intent(in)    :: space
    type(hamiltonian_elec_t),            target, intent(inout) :: hm
    type(grid_t),                        target, intent(in)    :: gr
    type(states_elec_t),                 target, intent(inout) :: st
    real(real64),                                intent(in)    :: time
    real(real64),                                intent(in)    :: dt
    type(ion_dynamics_t),                        intent(inout) :: ions_dyn
    type(ions_t),                                intent(inout) :: ions
    type(partner_list_t),                        intent(in)    :: ext_partners
    type(opt_control_state_t), optional, target, intent(inout) :: qcchi
 
    type(states_elec_t), pointer :: chi
    real(real64), pointer :: q(:, :), p(:, :)
 
    integer :: np_part, np, kp1, kp2, st1, st2, nspin, ik, ist, iatom, ib
    complex(real64), allocatable :: zphi(:, :, :, :), zchi(:, :, :, :), dvpsi(:, :, :)
    type(states_elec_t) :: hst, stphi, inh, hchi, stchi
    logical :: propagate_chi
    real(real64), allocatable :: pos0(:, :), vel0(:, :), &
      posk(:, :), velk(:, :), &
      pos(:, :), vel(:, :), &
      posfinal(:, :), velfinal(:, :), &
      pos0t(:, :), vel0t(:, :), &
      poskt(:, :), velkt(:, :), &
      post(:, :), velt(:, :), &
      posfinalt(:, :), velfinalt(:, :), &
      coforce(:, :)
 
    push_sub(td_explicit_runge_kutta4)
 
    propagate_chi = present(qcchi)
    if (propagate_chi) then
      chi => opt_control_point_qs(qcchi)
      q => opt_control_point_q(qcchi)
      p => opt_control_point_p(qcchi)
    end if
 
    st1 = st%st_start
    st2 = st%st_end
    kp1 = st%d%kpt%start
    kp2 = st%d%kpt%end
    np_part = gr%np_part
    np = gr%np
    nspin = hm%d%nspin
 
    safe_allocate(zphi(1:np_part, 1:st%d%dim, st1:st2, kp1:kp2))
    if (propagate_chi) then
      safe_allocate(zchi(1:np_part, 1:st%d%dim, st1:st2, kp1:kp2))
    end if
    if (ion_dynamics_ions_move(ions_dyn)) then
      safe_allocate(pos(1:ions%space%dim, 1:ions%natoms))
      safe_allocate(vel(1:ions%space%dim, 1:ions%natoms))
      safe_allocate(pos0(1:ions%space%dim, 1:ions%natoms))
      safe_allocate(vel0(1:ions%space%dim, 1:ions%natoms))
      safe_allocate(posk(1:ions%space%dim, 1:ions%natoms))
      safe_allocate(velk(1:ions%space%dim, 1:ions%natoms))
      safe_allocate(posfinal(1:ions%space%dim, 1:ions%natoms))
      safe_allocate(velfinal(1:ions%space%dim, 1:ions%natoms))
 
      if (propagate_chi) then
        safe_allocate(post(1:ions%space%dim, 1:ions%natoms))
        safe_allocate(velt(1:ions%space%dim, 1:ions%natoms))
        safe_allocate(pos0t(1:ions%space%dim, 1:ions%natoms))
        safe_allocate(vel0t(1:ions%space%dim, 1:ions%natoms))
        safe_allocate(poskt(1:ions%space%dim, 1:ions%natoms))
        safe_allocate(velkt(1:ions%space%dim, 1:ions%natoms))
        safe_allocate(posfinalt(1:ions%space%dim, 1:ions%natoms))
        safe_allocate(velfinalt(1:ions%space%dim, 1:ions%natoms))
        safe_allocate(coforce(1:ions%natoms, 1:ions%space%dim))
      end if
    end if
 
    call states_elec_copy(hst, st)
    call states_elec_copy(stphi, st)
    call states_elec_get_state(st, gr, zphi)
 
    if (propagate_chi) then
      call states_elec_copy(hchi, chi)
      call states_elec_copy(stchi, chi)
      call states_elec_get_state(chi, gr, zchi)
    end if
 
    if (ion_dynamics_ions_move(ions_dyn)) then
      pos0 = ions%pos
      vel0 = ions%vel
      posfinal = pos0
      velfinal = vel0
 
      if (propagate_chi) then
        do iatom = 1, ions%natoms
          pos0t(1:ions%space%dim, iatom) = q(iatom, 1:ions%space%dim)
          vel0t(1:ions%space%dim, iatom) = p(iatom, 1:ions%space%dim) / ions%mass(iatom)
        end do
        posfinalt = pos0t
        velfinalt = vel0t
      end if
    end if
 
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    ! Stage 1.
    !
 
    !
    call states_elec_set_state(stphi, gr, zphi)
    if (propagate_chi) then
      call states_elec_set_state(stchi, gr, zchi)
    end if
    if (ion_dynamics_ions_move(ions_dyn)) then
      pos = pos0
      vel = vel0
      if (propagate_chi) then
        post = pos0t
        velt = vel0t
      end if
    end if
 
    call f_psi(time - dt)
    if (propagate_chi) call f_chi(time - dt)
    if (ion_dynamics_ions_move(ions_dyn)) call f_ions(time - dt)
 
    call update_state(m_one/6.0_real64)
 
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    ! Stage 2.
 
    call states_elec_set_state(stphi, gr, zphi)
    if (propagate_chi) then
      call states_elec_set_state(stchi, gr, zchi)
    end if
 
    do ik = stphi%d%kpt%start, stphi%d%kpt%end
      do ib = stphi%group%block_start, stphi%group%block_end
        call batch_axpy(gr%np, -m_half*m_zi*dt, hst%group%psib(ib, ik), stphi%group%psib(ib, ik))
        if (propagate_chi) then
          call batch_axpy(gr%np, -m_half*m_zi*dt, hchi%group%psib(ib, ik), stchi%group%psib(ib, ik))
        end if
      end do
    end do
 
    if (ion_dynamics_ions_move(ions_dyn)) then
      pos = pos0 + m_half * posk
      vel = vel0 + m_half * velk
      if (propagate_chi) then
        post = pos0t + m_half * poskt
        velt = vel0t + m_half * velkt
      end if
    end if
 
    call f_psi(time-m_half*dt)
    if (propagate_chi) call f_chi(time-m_half*dt)
    if (ion_dynamics_ions_move(ions_dyn)) call f_ions(time-m_half*dt)
    call update_state(m_third)
 
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    ! Stage 3.
 
    call states_elec_set_state(stphi, gr, zphi)
    if (propagate_chi) then
      call states_elec_set_state(stchi, gr, zchi)
    end if
 
    do ik = stphi%d%kpt%start, stphi%d%kpt%end
      do ib = stphi%group%block_start, stphi%group%block_end
        call batch_axpy(gr%np, -m_half*m_zi*dt, hst%group%psib(ib, ik), stphi%group%psib(ib, ik))
        if (propagate_chi) then
          call batch_axpy(gr%np, -m_half*m_zi*dt, hchi%group%psib(ib, ik), stchi%group%psib(ib, ik))
        end if
      end do
    end do
 
    if (ion_dynamics_ions_move(ions_dyn)) then
      pos = pos0 + m_half * posk
      vel = vel0 + m_half * velk
      if (propagate_chi) then
        post = pos0t + m_half * poskt
        velt = vel0t + m_half * velkt
      end if
    end if
 
    call f_psi(time-m_half*dt)
    if (propagate_chi) call f_chi(time-m_half*dt)
    if (ion_dynamics_ions_move(ions_dyn)) call f_ions(time-m_half*dt)
 
    call update_state(m_third)
 
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    ! Stage 4.
 
    call states_elec_set_state(stphi, gr, zphi)
    if (propagate_chi) then
      call states_elec_set_state(stchi, gr, zchi)
    end if
 
    do ik = stphi%d%kpt%start, stphi%d%kpt%end
      do ib = stphi%group%block_start, stphi%group%block_end
        call batch_axpy(gr%np, -m_zi*dt, hst%group%psib(ib, ik), stphi%group%psib(ib, ik))
        if (propagate_chi) then
          call batch_axpy(gr%np, -m_zi*dt, hchi%group%psib(ib, ik), stchi%group%psib(ib, ik))
        end if
      end do
    end do
 
    if (ion_dynamics_ions_move(ions_dyn)) then
      pos = pos0 + posk
      vel = vel0 + velk
      if (propagate_chi) then
        post = pos0t + poskt
        velt = vel0t + velkt
      end if
    end if
 
    call f_psi(time)
    if (propagate_chi) call f_chi(time)
    if (ion_dynamics_ions_move(ions_dyn)) call f_ions(time)
 
    call update_state(m_one / 6.0_real64)
 
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    ! Collect the results.
 
    call density_calc(st, gr, st%rho)
    if (ion_dynamics_ions_move(ions_dyn)) then
      ions%pos = posfinal
      ions%vel = velfinal
      call hamiltonian_elec_epot_generate(hm, namespace,  space, gr, ions, ext_partners, st, time)
      !call forces_calculate(gr, namespace, ions, hm, stphi, time, dt)
      call ions%update_kinetic_energy()
 
      if (propagate_chi) then
        do iatom = 1, ions%natoms
          q(iatom, 1:ions%space%dim) = posfinalt(1:ions%space%dim, iatom)
          p(iatom, 1:ions%space%dim) = ions%mass(iatom) * velfinalt(1:ions%space%dim, iatom)
        end do
      end if
    end if
 
    call states_elec_end(hst)
    call states_elec_end(stphi)
    safe_deallocate_a(zphi)
    if (propagate_chi) then
      call states_elec_end(hchi)
      call states_elec_end(stchi)
      safe_deallocate_a(zchi)
      nullify(chi)
      nullify(p)
      nullify(q)
    end if
 
    if (ion_dynamics_ions_move(ions_dyn)) then
      safe_deallocate_a(pos)
      safe_deallocate_a(vel)
      safe_deallocate_a(pos0)
      safe_deallocate_a(vel0)
      safe_deallocate_a(posk)
      safe_deallocate_a(velk)
      safe_deallocate_a(posfinal)
      safe_deallocate_a(velfinal)
 
      if (propagate_chi) then
        safe_deallocate_a(post)
        safe_deallocate_a(velt)
        safe_deallocate_a(pos0t)
        safe_deallocate_a(vel0t)
        safe_deallocate_a(poskt)
        safe_deallocate_a(velkt)
        safe_deallocate_a(posfinalt)
        safe_deallocate_a(velfinalt)
        safe_deallocate_a(coforce)
      end if
    end if
    pop_sub(td_explicit_runge_kutta4)
 
  contains
 
    subroutine f_psi(tau)
      real(real64), intent(in) :: tau
 
      if (ion_dynamics_ions_move(ions_dyn)) then
        ions%pos = pos
        ions%vel = vel
        call hamiltonian_elec_epot_generate(hm, namespace,  space, gr, ions, ext_partners, stphi, time = tau)
      end if
      if (.not. oct_exchange_enabled(hm%oct_exchange)) then
        call density_calc(stphi, gr, stphi%rho)
        call v_ks_calc(ks, namespace, space, hm, stphi, ions, ext_partners, &
          calc_current = list_has_gauge_field(ext_partners), &
          time = tau, calc_energy = .false., calc_eigenval = .false.)
      else
        call hm%update(gr, namespace, space, ext_partners, time = tau)
      end if
      call lda_u_update_occ_matrices(hm%lda_u, namespace, gr, st, hm%hm_base, hm%phase, hm%energy)
      call zhamiltonian_elec_apply_all(hm, namespace, gr, stphi, hst)
    end subroutine f_psi
 
    subroutine f_ions(tau)
      real(real64), intent(in) :: tau
 
      call forces_calculate(gr, namespace, ions, hm, ext_partners, stphi, ks, t = tau, dt = dt)
      do iatom = 1, ions%natoms
        posk(:, iatom) = dt * vel(:, iatom)
        velk(:, iatom) = dt * ions%tot_force(:, iatom) / ions%mass(iatom)
      end do
      if (propagate_chi) then
        call forces_costate_calculate(gr, namespace, ions, hm, stphi, stchi, coforce, transpose(post))
        do iatom = 1, ions%natoms
          poskt(:, iatom) = dt * velt(:, iatom)
          velkt(:, iatom) = dt * coforce(iatom, :) / ions%mass(iatom)
        end do
      end if
    end subroutine f_ions
 
    subroutine f_chi(tau)
      real(real64), intent(in) :: tau
 
      if (hm%theory_level /= independent_particles) call oct_exchange_set(hm%oct_exchange, stphi, gr)
      call prepare_inh()
      call hamiltonian_elec_adjoint(hm)
 
      call propagation_ops_elec_update_hamiltonian(namespace, space, st, gr, hm, ext_partners, tau)
 
      call zhamiltonian_elec_apply_all(hm, namespace, gr, stchi, hchi)
      call hamiltonian_elec_not_adjoint(hm)
 
 
      call apply_inh()
      if (hm%theory_level /= independent_particles) call oct_exchange_remove(hm%oct_exchange)
      if (ion_dynamics_ions_move(ions_dyn)) call hamiltonian_elec_remove_inh(hm)
    end subroutine f_chi
 
    subroutine apply_inh()
      integer :: ib
 
      if (hamiltonian_elec_inh_term(hm)) then
        do ik = kp1, kp2
          do ib = 1, st%group%block_start, st%group%block_end
            call batch_axpy(np, m_zi, hm%inh_st%group%psib(ib, ik), hchi%group%psib(ib, ik))
          end do
        end do
      end if
    end subroutine apply_inh
 
    subroutine prepare_inh()
      integer :: idir
      complex(real64), allocatable :: psi(:, :), inhpsi(:, :)
      type(perturbation_ionic_t), pointer :: pert
 
      if (ion_dynamics_ions_move(ions_dyn)) then
        call states_elec_copy(inh, st)
 
        safe_allocate(psi(1:gr%np_part, 1:st%d%dim))
        safe_allocate(inhpsi(1:gr%np_part, 1:st%d%dim))
        safe_allocate(dvpsi(1:gr%np_part, 1:st%d%dim, 1:space%dim))
 
        do ik = 1, st%nik
          do ist = 1, st%nst
 
            inhpsi = m_z0
            call states_elec_get_state(stphi, gr, ist, ik, psi)
 
            do iatom = 1, ions%natoms
              do idir = 1, space%dim
                pert => perturbation_ionic_t(namespace, ions)
                call pert%setup_atom(iatom)
                call pert%setup_dir(idir)
                call pert%zapply(namespace, space, gr, hm, ik, psi(:, :), dvpsi(:, :, idir))
                dvpsi(:, :, idir) = - dvpsi(:, :, idir)
                inhpsi(1:gr%np, 1:stphi%d%dim) = inhpsi(1:gr%np, 1:stphi%d%dim) &
                  + st%occ(ist, ik)*post(idir, iatom)*dvpsi(1:gr%np, 1:stphi%d%dim, idir)
                safe_deallocate_p(pert)
              end do
            end do
 
            call states_elec_set_state(inh, gr, ist, ik, inhpsi)
 
          end do
        end do
 
        call hamiltonian_elec_set_inh(hm, inh)
        call states_elec_end(inh)
 
        safe_deallocate_a(psi)
        safe_deallocate_a(inhpsi)
        safe_deallocate_a(dvpsi)
 
      end if
    end subroutine prepare_inh
 
    subroutine update_state(epsilon)
      real(real64), intent(in) :: epsilon
 
      do ik = stphi%d%kpt%start, stphi%d%kpt%end
        do ib = stphi%group%block_start, stphi%group%block_end
          call batch_axpy(gr%np, -m_zi*dt*epsilon, hst%group%psib(ib, ik), st%group%psib(ib, ik))
          if (propagate_chi) then
            call batch_axpy(gr%np, -m_zi*dt*epsilon, hchi%group%psib(ib, ik), chi%group%psib(ib, ik))
          end if
        end do
      end do
 
      if (ion_dynamics_ions_move(ions_dyn)) then
        posfinal = posfinal + posk * epsilon
        velfinal = velfinal + velk * epsilon
        if (propagate_chi) then
          posfinalt = posfinalt + poskt * epsilon
          velfinalt = velfinalt + velkt * epsilon
        end if
      end if
    end subroutine update_state
 
  end subroutine td_explicit_runge_kutta4
 
 
  subroutine td_runge_kutta2(ks, namespace, space, hm, gr, st, tr, time, dt, ions_dyn, ions, ext_partners)
    type(v_ks_t),             target, intent(inout) :: ks
    type(namespace_t),        target, intent(in)    :: namespace
    type(electron_space_t),   target, intent(in)    :: space
    type(hamiltonian_elec_t), target, intent(inout) :: hm
    type(grid_t),             target, intent(in)    :: gr
    type(states_elec_t),      target, intent(inout) :: st
    type(propagator_base_t),  target, intent(inout) :: tr
    real(real64),                     intent(in)    :: time
    real(real64),                     intent(in)    :: dt
    type(ion_dynamics_t),             intent(inout) :: ions_dyn
    type(ions_t),                     intent(inout) :: ions
    type(partner_list_t),     target, intent(in)    :: ext_partners
 
    integer :: np_part, np, kp1, kp2, st1, st2, nspin, ik, ist, idim, j, ip
    integer :: i
    real(real64) :: dres
    complex(real64), allocatable :: zphi(:, :, :, :)
    complex(real64), allocatable :: zpsi(:), rhs(:)
    complex(real64), allocatable :: k2(:, :, :, :), oldk2(:, :, :, :), rhs1(:, :, :, :)
    complex(real64), allocatable :: inhpsi(:)
    type(ion_state_t) :: ions_state
 
    push_sub(td_runge_kutta2)
 
    st1 = st%st_start
    st2 = st%st_end
    kp1 = st%d%kpt%start
    kp2 = st%d%kpt%end
    np_part = gr%np_part
    np = gr%np
    nspin = hm%d%nspin
    move_ions_op = ion_dynamics_ions_move(ions_dyn)
 
    sp_np = np
    sp_dim = st%d%dim
    sp_st1 = st1
    sp_st2 = st2
    sp_kp1 = kp1
    sp_kp2 = kp2
    sp_parallel = st%parallel_in_states .or. st%d%kpt%parallel
    if (sp_parallel) call mpi_grp_copy(sp_grp, st%st_kpt_mpi_grp)
 
    ! define pointer and variables for usage in td_rk2op, td_rk2opt routines
    mesh_p    => gr
    hm_p      => hm
    tr_p      => tr
    st_p      => st
    namespace_p => namespace
    space_p   => space
    ext_partners_p  => ext_partners
    dt_op = dt
    t_op  = time - dt/m_two
    dim_op = st%d%dim
 
    safe_allocate(k2(1:np_part, 1:st%d%dim, st1:st2, kp1:kp2))
    safe_allocate(oldk2(1:np_part, 1:st%d%dim, st1:st2, kp1:kp2))
    safe_allocate(zphi(1:gr%np_part, st%d%dim, st1:st2, kp1:kp2))
    safe_allocate(rhs1(1:np_part, 1:st%d%dim, st1:st2, kp1:kp2))
    safe_allocate(rhs(1:tr%tdsk_size))
    safe_allocate(zpsi(1:tr%tdsk_size))
    safe_allocate(vpsl1_op(1:np))
 
    ! First, we get the state that we want to propagate. For the moment being, only one state.
    do ik = kp1, kp2
      do ist = st1, st2
        call states_elec_get_state(st, gr, ist, ik, zphi(:, :, ist, ik))
      end do
    end do
 
    if (oct_exchange_enabled(hm%oct_exchange)) then
      call oct_exchange_prepare(hm%oct_exchange, gr, zphi, hm%xc, hm%psolver, namespace)
    end if
 
    call propagation_ops_elec_update_hamiltonian(namespace, space, st, gr, hm, ext_partners, time - dt)
 
    rhs1 = m_z0
    do ik = kp1, kp2
      do ist = st1, st2
        call zhamiltonian_elec_apply_single(hm_p, namespace, mesh_p, zphi(:, :, ist, ik), rhs1(:, :, ist, ik), ist, ik)
      end do
    end do
    do ik = kp1, kp2
      do ist = st1, st2
        if (oct_exchange_enabled(hm%oct_exchange)) then
          call oct_exchange_operator(hm%oct_exchange, namespace, gr, rhs1(:, :, ist, ik), ist, ik)
        end if
      end do
    end do
 
    rhs1 = zphi - m_zi * m_half * dt * rhs1
 
    if (hamiltonian_elec_inh_term(hm)) then
      safe_allocate(inhpsi(1:gr%np))
      do ik = kp1, kp2
        do ist = st1, st2
          do idim = 1, st%d%dim
            call states_elec_get_state(hm%inh_st, gr, idim, ist, ik, inhpsi)
            do ip = 1, gr%np
              rhs1(ip, idim, ist, ik) = rhs1(ip, idim, ist, ik) + dt * inhpsi(ip)
            end do
          end do
        end do
      end do
      safe_deallocate_a(inhpsi)
    end if
 
    k2 = zphi
 
    i = 1
    do
      oldk2 = k2
 
      ! Set the Hamiltonian at the final time of the propagation
      if (.not. oct_exchange_enabled(hm_p%oct_exchange)) then
        do ik = kp1, kp2
          do ist = st1, st2
            call states_elec_set_state(st, gr, ist, ik, k2(:, :, ist, ik))
          end do
        end do
        call density_calc(st, gr, st%rho)
        call v_ks_calc(ks, namespace, space, hm, st, ions, ext_partners, &
          calc_current = list_has_gauge_field(ext_partners),&
          calc_energy = .false., calc_eigenval = .false.)
      end if
      if (ion_dynamics_ions_move(ions_dyn)) then
        call ion_dynamics_save_state(ions_dyn, ions, ions_state)
        call ion_dynamics_propagate(ions_dyn, ions, time, dt, namespace)
        call hamiltonian_elec_epot_generate(hm, namespace, space, gr, ions, ext_partners, st, time = time)
        vpsl1_op = hm%ep%vpsl
      end if
 
      call propagation_ops_elec_update_hamiltonian(namespace, space, st, gr, hm, ext_partners, time)
 
      if (.not. oct_exchange_enabled(hm_p%oct_exchange)) then
        if (i == 1) then
          call hm%ks_pot%get_interpolated_potentials(tr%vks_old, 0, storage=vhxc1_op)
          i = i + 1
        else
          call hm%ks_pot%store_potentials(vhxc1_op)
        end if
        t_op  = time
      else
        call hm%ks_pot%store_potentials(vhxc1_op)
      end if
 
      if (ion_dynamics_ions_move(ions_dyn)) then
        call ion_dynamics_restore_state(ions_dyn, ions, ions_state)
      end if
 
      j = 1
      do ik = kp1, kp2
        do ist = st1, st2
          do idim = 1, st%d%dim
            rhs(j:j+np-1) = rhs1(1:np, idim, ist, ik)
            j = j + np
          end do
        end do
      end do
 
      ! Now we populate an initial guess for the equation.
      j = 1
      do ik = kp1, kp2
        do ist = st1, st2
          do idim = 1, st%d%dim
            zpsi(j:j+np-1) = k2(1:np, idim, ist, ik)
            j = j + np
          end do
        end do
      end do
 
      t_op  = time - dt
      call zsparskit_solver_run(namespace, tr%tdsk, td_rk2op, td_rk2opt, zpsi, rhs)
 
      k2 = m_z0
      j = 1
      do ik = kp1, kp2
        do ist = st1, st2
          do idim = 1, st%d%dim
            k2(1:np, idim, ist, ik) = zpsi(j:j+np-1)
            j = j + np
          end do
        end do
      end do
 
      dres = m_zero
      do ik = kp1, kp2
        do ist = st1, st2
          do idim = 1, st%d%dim
            dres = dres + (zmf_nrm2(gr, k2(:, idim, ist, ik) - oldk2(:, idim, ist, ik), reduce = .false.))**2
          end do
        end do
      end do
 
      call comm_allreduce(st%dom_st_kpt_mpi_grp, dres)
 
      if (sqrt(dres) < tr%scf_threshold) exit
    end do
 
    zphi = k2
    do ik = kp1, kp2
      do ist = st1, st2
        call states_elec_set_state(st, gr, ist, ik, zphi(:, :, ist, ik))
      end do
    end do
 
    call density_calc(st, gr, st%rho)
 
    safe_deallocate_a(k2)
    safe_deallocate_a(oldk2)
    safe_deallocate_a(zphi)
    safe_deallocate_a(rhs1)
    safe_deallocate_a(zpsi)
    safe_deallocate_a(rhs)
    safe_deallocate_a(vpsl1_op)
 
    pop_sub(td_runge_kutta2)
  end subroutine td_runge_kutta2
 
  !----------------------------------------------------------------------------
 
  subroutine td_runge_kutta4(ks, namespace, space, hm, gr, st, tr, time, dt, ions_dyn, ions, ext_partners)
    type(v_ks_t),             target, intent(inout) :: ks
    type(namespace_t),        target, intent(in)    :: namespace
    type(electron_space_t),   target, intent(in)    :: space
    type(hamiltonian_elec_t), target, intent(inout) :: hm
    type(grid_t),             target, intent(in)    :: gr
    type(states_elec_t),      target, intent(inout) :: st
    type(propagator_base_t),  target, intent(inout) :: tr
    real(real64),                     intent(in)    :: time
    real(real64),                     intent(in)    :: dt
    type(ion_dynamics_t),             intent(inout) :: ions_dyn
    type(ions_t),                     intent(inout) :: ions
    type(partner_list_t),     target, intent(in)    :: ext_partners
 
    integer :: idim, ip, ist, ik, j, kp1, kp2, st1, st2, nspin
    real(real64) :: dres
    complex(real64), allocatable :: inhpsi(:)
    complex(real64), allocatable :: zpsi(:), rhs(:)
    complex(real64), allocatable :: zphi(:, :, :, :)
    type(ion_state_t) :: ions_state
 
    real(real64) :: a(2, 2), c(2), b(2)
    complex(real64), allocatable :: k1(:, :, :, :), k2(:, :, :, :), oldk1(:, :, :, :), &
      oldk2(:, :, :, :), yn1(:, :, :, :), yn2(:, :, :, :), &
      rhs1(:, :, :, :), rhs2(:, :, :, :)
 
    push_sub(td_runge_kutta4)
 
    c(1) = m_half - sqrt(m_three)/6.0_real64
    c(2) = m_half + sqrt(m_three)/6.0_real64
    b(1) = m_half
    b(2) = m_half
    a(1, 1) = m_fourth
    a(1, 2) = m_fourth - sqrt(m_three)/6.0_real64
    a(2, 1) = m_fourth + sqrt(m_three)/6.0_real64
    a(2, 2) = m_fourth
 
    st1 = st%st_start
    st2 = st%st_end
    kp1 = st%d%kpt%start
    kp2 = st%d%kpt%end
    nspin = hm%d%nspin
    move_ions_op = ion_dynamics_ions_move(ions_dyn)
 
    sp_np = gr%np
    sp_dim = st%d%dim
    sp_st1 = st1
    sp_st2 = st2
    sp_kp1 = kp1
    sp_kp2 = kp2
    sp_parallel = st%parallel_in_states .or. st%d%kpt%parallel
    if (sp_parallel) call mpi_grp_copy(sp_grp, st%st_kpt_mpi_grp)
 
    ! define pointer and variables for usage in td_rk4op, td_rk4opt routines
    mesh_p    => gr
    hm_p      => hm
    tr_p      => tr
    st_p      => st
    namespace_p => namespace
    space_p   => space
    ext_partners_p  => ext_partners
    dt_op = dt
    t_op  = time - dt/m_two
    dim_op = st%d%dim
 
    safe_allocate(vpsl1_op(1:gr%np))
    safe_allocate(vpsl2_op(1:gr%np))
    safe_allocate(k1(1:gr%np_part, 1:st%d%dim, st1:st2, kp1:kp2))
    safe_allocate(k2(1:gr%np_part, 1:st%d%dim, st1:st2, kp1:kp2))
    safe_allocate(oldk1(1:gr%np_part, 1:st%d%dim, st1:st2, kp1:kp2))
    safe_allocate(oldk2(1:gr%np_part, 1:st%d%dim, st1:st2, kp1:kp2))
    safe_allocate(yn1(1:gr%np_part, 1:st%d%dim, st1:st2, kp1:kp2))
    safe_allocate(yn2(1:gr%np_part, 1:st%d%dim, st1:st2, kp1:kp2))
    safe_allocate(rhs1(1:gr%np_part, 1:st%d%dim, st1:st2, kp1:kp2))
    safe_allocate(rhs2(1:gr%np_part, 1:st%d%dim, st1:st2, kp1:kp2))
    safe_allocate(rhs(1:tr%tdsk_size))
    safe_allocate(zpsi(1:tr%tdsk_size))
    safe_allocate(zphi(1:gr%np_part, st%d%dim, st1:st2, kp1:kp2))
 
    ! First, we get the state that we want to propagate. For the moment being, only one state.
    do ik = kp1, kp2
      do ist = st1, st2
        call states_elec_get_state(st, gr, ist, ik, zphi(:, :, ist, ik))
      end do
    end do
    k1 = m_z0
    k2 = m_z0
 
    do
 
      oldk1 = k1
      oldk2 = k2
 
      yn1 = zphi + a(1, 1) * k1 + a(1, 2) * k2
      yn2 = zphi + a(2, 1) * k1 + a(2, 2) * k2
 
      ! Set the Hamiltonian at time-dt + c(1) * dt
      do ik = kp1, kp2
        do ist = st1, st2
          call states_elec_set_state(st, gr, ist, ik, yn1(:, :, ist, ik))
        end do
      end do
      call density_calc(st, gr, st%rho)
      call v_ks_calc(ks, namespace, space, hm, st, ions, ext_partners, &
        calc_current = list_has_gauge_field(ext_partners), &
        calc_energy = .false., calc_eigenval = .false.)
      if (ion_dynamics_ions_move(ions_dyn)) then
        call ion_dynamics_save_state(ions_dyn, ions, ions_state)
        call ion_dynamics_propagate(ions_dyn, ions, time - dt + c(1)*dt, c(1)*dt, namespace)
        call hamiltonian_elec_epot_generate(hm, namespace, space, gr, ions, ext_partners, st, time = time - dt + c(1)*dt)
        vpsl1_op = hm%ep%vpsl
      end if
 
      call propagation_ops_elec_update_hamiltonian(namespace, space, st, gr, hm, ext_partners, time - dt + c(1)*dt)
 
      call hm%ks_pot%store_potentials(vhxc1_op)
      t_op  = time - dt + c(1) * dt
      rhs1 = m_z0
      do ik = kp1, kp2
        do ist = st1, st2
          call zhamiltonian_elec_apply_single(hm_p, namespace, mesh_p, zphi(:, :, ist, ik), rhs1(:, :, ist, ik), ist, ik)
          if (hamiltonian_elec_inh_term(hm)) then
            safe_allocate(inhpsi(1:gr%np))
            do idim = 1, st%d%dim
              call states_elec_get_state(hm%inh_st, gr, idim, ist, ik, inhpsi)
              do ip = 1, gr%np
                rhs1(ip, idim, ist, ik) = rhs1(ip, idim, ist, ik) + m_zi * inhpsi(ip)
              end do
            end do
            safe_deallocate_a(inhpsi)
          end if
        end do
      end do
      rhs1 = - m_zi * dt * rhs1
      if (ion_dynamics_ions_move(ions_dyn)) then
        call ion_dynamics_restore_state(ions_dyn, ions, ions_state)
      end if
 
      ! Set the Hamiltonian at time-dt + c(2) * dt
      do ik = kp1, kp2
        do ist = st1, st2
          call states_elec_set_state(st, gr, ist, ik, yn2(:, :, ist, ik))
        end do
      end do
      call density_calc(st, gr, st%rho)
      call v_ks_calc(ks, namespace, space, hm, st, ions, ext_partners, &
        calc_current = list_has_gauge_field(ext_partners), &
        calc_energy = .false., calc_eigenval = .false.)
      if (ion_dynamics_ions_move(ions_dyn)) then
        call ion_dynamics_save_state(ions_dyn, ions, ions_state)
        call ion_dynamics_propagate(ions_dyn, ions, time - dt + c(2)*dt, c(2)*dt, namespace)
        call hamiltonian_elec_epot_generate(hm, namespace, space, gr, ions, ext_partners, st, time = time - dt + c(2)*dt)
        vpsl2_op = hm%ep%vpsl
      end if
 
      call propagation_ops_elec_update_hamiltonian(namespace, space, st, gr, hm, ext_partners, time - dt + c(2)*dt)
 
      call hm%ks_pot%store_potentials(vhxc2_op)
      t_op  = time - dt + c(2) * dt
      rhs2 = m_z0
      do ik = kp1, kp2
        do ist = st1, st2
          call zhamiltonian_elec_apply_single(hm_p, namespace, mesh_p, zphi(:, :, ist, ik), rhs2(:, :, ist, ik), ist, ik)
          if (hamiltonian_elec_inh_term(hm)) then
            safe_allocate(inhpsi(1:gr%np))
            do idim = 1, st%d%dim
              call states_elec_get_state(hm%inh_st, gr, idim, ist, ik, inhpsi)
              do ip = 1, gr%np
                rhs2(ip, idim, ist, ik) = rhs2(ip, idim, ist, ik) + m_zi * inhpsi(ip)
              end do
            end do
            safe_deallocate_a(inhpsi)
          end if
        end do
      end do
      rhs2 = -m_zi * dt * rhs2
      if (ion_dynamics_ions_move(ions_dyn)) then
        call ion_dynamics_restore_state(ions_dyn, ions, ions_state)
      end if
 
      j = 1
      do ik = kp1, kp2
        do ist = st1, st2
          do idim = 1, st%d%dim
            call lalg_copy(gr%np, rhs1(1:gr%np, idim, ist, ik), rhs(j:j+gr%np-1))
            j = j + gr%np
          end do
        end do
      end do
      do ik = kp1, kp2
        do ist = st1, st2
          do idim = 1, st%d%dim
            call lalg_copy(gr%np, rhs2(1:gr%np, idim, ist, ik), rhs(j:j+gr%np-1))
            j = j + gr%np
          end do
        end do
      end do
 
      ! Now we populate an initial guess for the equation.
      j = 1
      do ik = kp1, kp2
        do ist = st1, st2
          do idim = 1, st%d%dim
            call lalg_copy(gr%np, k1(1:gr%np, idim, ist, ik), zpsi(j:j+gr%np-1))
            j = j + gr%np
          end do
        end do
      end do
      do ik = kp1, kp2
        do ist = st1, st2
          do idim = 1, st%d%dim
            call lalg_copy(gr%np, k2(1:gr%np, idim, ist, ik), zpsi(j:j+gr%np-1))
            j = j + gr%np
          end do
        end do
      end do
 
      t_op  = time - dt
      call zsparskit_solver_run(namespace, tr%tdsk, td_rk4op, td_rk4opt, zpsi, rhs)
 
      k1 = m_z0
      k2 = m_z0
      j = 1
      do ik = kp1, kp2
        do ist = st1, st2
          do idim = 1, st%d%dim
            call lalg_copy(gr%np, zpsi(j:j+gr%np-1), k1(1:gr%np, idim, ist, ik))
            j = j + gr%np
          end do
        end do
      end do
      do ik = kp1, kp2
        do ist = st1, st2
          do idim = 1, st%d%dim
            call lalg_copy(gr%np, zpsi(j:j+gr%np-1), k2(1:gr%np, idim, ist, ik))
            j = j + gr%np
          end do
        end do
      end do
 
      dres = m_zero
      do ik = kp1, kp2
        do ist = st1, st2
          do idim = 1, st%d%dim
            dres = dres + (zmf_nrm2(gr, k1(:, idim, ist, ik) - oldk1(:, idim, ist, ik)))**2
            dres = dres + (zmf_nrm2(gr, k2(:, idim, ist, ik) - oldk2(:, idim, ist, ik)))**2
          end do
        end do
      end do
      if (sp_parallel) call comm_allreduce(sp_grp, dres)
      !write(*, *) 'Residual = ', dres
 
      if (sqrt(dres) < tr%scf_threshold) exit
    end do
 
 
    zphi = zphi + b(1) * k1 + b(2) * k2
    do ik = kp1, kp2
      do ist = st1, st2
        call states_elec_set_state(st, gr, ist, ik, zphi(:, :, ist, ik))
      end do
    end do
 
    call density_calc(st, gr, st%rho)
 
    safe_deallocate_a(rhs1)
    safe_deallocate_a(rhs2)
    safe_deallocate_a(k1)
    safe_deallocate_a(k2)
    safe_deallocate_a(oldk1)
    safe_deallocate_a(oldk2)
    safe_deallocate_a(yn1)
    safe_deallocate_a(yn2)
    safe_deallocate_a(vpsl1_op)
    safe_deallocate_a(vpsl2_op)
    safe_deallocate_a(zpsi)
    safe_deallocate_a(rhs)
 
    pop_sub(td_runge_kutta4)
  end subroutine td_runge_kutta4
 
  ! ---------------------------------------------------------
  subroutine td_rk4op(xre, xim, yre, yim)
    real(real64), intent(in)  :: xre(:)
    real(real64), intent(in)  :: xim(:)
    real(real64), intent(out) :: yre(:)
    real(real64), intent(out) :: yim(:)
 
    integer :: idim, j, ik, ist, kp1, kp2, st1, st2, dim, k, jj
    complex(real64), allocatable :: zpsi(:, :)
    complex(real64), allocatable :: opzpsi(:, :)
    real(real64) :: a(2, 2), c(2)
    integer :: np_part, np
 
    push_sub(td_rk4op)
 
    np_part = mesh_p%np_part
    np = mesh_p%np
    st1 = st_p%st_start
    st2 = st_p%st_end
    kp1 = st_p%d%kpt%start
    kp2 = st_p%d%kpt%end
    dim = st_p%d%dim
 
    safe_allocate(zpsi(1:np_part, 1:dim))
    safe_allocate(opzpsi(1:np_part, 1:dim))
 
    a(1, 1) = m_fourth
    a(1, 2) = m_fourth - sqrt(m_three)/6.0_real64
    a(2, 1) = m_fourth + sqrt(m_three)/6.0_real64
    a(2, 2) = m_fourth
 
    c(1) = m_half - sqrt(m_three)/6.0_real64
    c(2) = m_half + sqrt(m_three)/6.0_real64
 
    zpsi = m_z0
 
    call hm_p%ks_pot%restore_potentials(vhxc1_op)
    if (move_ions_op) hm_p%ep%vpsl = vpsl1_op
    call propagation_ops_elec_update_hamiltonian(namespace_p, space_p, st_p, mesh_p, hm_p, ext_partners_p, t_op + c(1)*dt_op)
    j = 1
    k = np * (kp2 - kp1 + 1) * (st2 - st1 + 1) * dim + 1
    do ik = kp1, kp2
      do ist = st1, st2
        jj = j
        do idim = 1, dim
          zpsi(1:np, idim) = a(1, 1) * cmplx(xre(j:j+np-1), xim(j:j+np-1), real64)  + &
            a(1, 2) * cmplx(xre(k:k+np-1), xim(k:k+np-1), real64)
          j = j + np
          k = k + np
        end do
 
        call zhamiltonian_elec_apply_single(hm_p, namespace_p, mesh_p, zpsi, opzpsi, ist, ik)
 
        do idim = 1, dim
          yre(jj:jj+np-1) = xre(jj:jj+np-1) - aimag(dt_op * opzpsi(1:np, idim))
          yim(jj:jj+np-1) = xim(jj:jj+np-1) + real(dt_op * opzpsi(1:np, idim), real64)
          jj = jj + np
        end do
      end do
    end do
 
    call hm_p%ks_pot%restore_potentials(vhxc2_op)
 
    if (move_ions_op) hm_p%ep%vpsl = vpsl2_op
    call propagation_ops_elec_update_hamiltonian(namespace_p, space_p, st_p, mesh_p, hm_p, ext_partners_p, t_op + c(2)*dt_op)
    j = 1
    k = np * (kp2 - kp1 + 1) * (st2 - st1 + 1) * dim + 1
    do ik = kp1, kp2
      do ist = st1, st2
        jj = k
        do idim = 1, dim
          zpsi(1:np, idim) = a(2, 1) * cmplx(xre(j:j+np-1), xim(j:j+np-1), real64)  + &
            a(2, 2) * cmplx(xre(k:k+np-1), xim(k:k+np-1), real64)
          j = j + np
          k = k + np
        end do
 
        call zhamiltonian_elec_apply_single(hm_p, namespace_p, mesh_p, zpsi, opzpsi, ist, ik)
 
        do idim = 1, dim
          yre(jj:jj+np-1) = xre(jj:jj+np-1) - aimag(dt_op * opzpsi(1:np, idim))
          yim(jj:jj+np-1) = xim(jj:jj+np-1) + real(dt_op * opzpsi(1:np, idim), real64)
          jj = jj + np
        end do
      end do
    end do
 
    safe_deallocate_a(zpsi)
    safe_deallocate_a(opzpsi)
    pop_sub(td_rk4op)
  end subroutine td_rk4op
  ! ---------------------------------------------------------
 
 
  ! ---------------------------------------------------------
  subroutine td_rk4opt(xre, xim, yre, yim)
    real(real64), intent(in)  :: xre(:)
    real(real64), intent(in)  :: xim(:)
    real(real64), intent(out) :: yre(:)
    real(real64), intent(out) :: yim(:)
 
    integer :: idim, j, ik, ist, kp1, kp2, st1, st2, dim, k, jj
    complex(real64), allocatable :: zpsi(:, :)
    complex(real64), allocatable :: opzpsi(:, :)
    real(real64) :: a(2, 2), c(2)
    integer :: np_part, np
 
    push_sub(td_rk4opt)
 
    np_part = mesh_p%np_part
    np = mesh_p%np
    st1 = st_p%st_start
    st2 = st_p%st_end
    kp1 = st_p%d%kpt%start
    kp2 = st_p%d%kpt%end
    dim = st_p%d%dim
 
    safe_allocate(zpsi(1:np_part, 1:dim))
    safe_allocate(opzpsi(1:np_part, 1:dim))
 
    a(1, 1) = m_fourth
    a(1, 2) = m_fourth - sqrt(m_three)/6.0_real64
    a(2, 1) = m_fourth + sqrt(m_three)/6.0_real64
    a(2, 2) = m_fourth
 
    c(1) = m_half - sqrt(m_three)/6.0_real64
    c(2) = m_half + sqrt(m_three)/6.0_real64
 
    zpsi = m_z0
 
    call hm_p%ks_pot%restore_potentials(vhxc1_op)
    if (move_ions_op) hm_p%ep%vpsl = vpsl1_op
 
    call propagation_ops_elec_update_hamiltonian(namespace_p, space_p, st_p, mesh_p, hm_p, ext_partners_p, t_op + c(1)*dt_op)
 
    j = 1
    k = np * (kp2 - kp1 + 1) * (st2 - st1 + 1) * dim + 1
    do ik = kp1, kp2
      do ist = st1, st2
        jj = j
        do idim = 1, dim
          zpsi(1:np, idim) = a(1, 1) * cmplx(xre(j:j+np-1), -xim(j:j+np-1), real64)  + &
            a(1, 2) * cmplx(xre(k:k+np-1), -xim(k:k+np-1), real64)
          j = j + np
          k = k + np
        end do
 
        call zhamiltonian_elec_apply_single(hm_p, namespace_p, mesh_p, zpsi, opzpsi, ist, ik)
 
        do idim = 1, dim
          yre(jj:jj+np-1) = xre(jj:jj+np-1) - aimag(dt_op * opzpsi(1:np, idim))
          yim(jj:jj+np-1) = xim(jj:jj+np-1) - real(dt_op * opzpsi(1:np, idim), real64)
          jj = jj + np
        end do
      end do
    end do
 
    call hm_p%ks_pot%restore_potentials(vhxc2_op)
    if (move_ions_op) hm_p%ep%vpsl = vpsl2_op
 
    call propagation_ops_elec_update_hamiltonian(namespace_p, space_p, st_p, mesh_p, hm_p, ext_partners_p, t_op + c(2)*dt_op)
 
    j = 1
    k = np * (kp2 - kp1 + 1) * (st2 - st1 + 1) * dim + 1
    do ik = kp1, kp2
      do ist = st1, st2
        jj = k
        do idim = 1, dim
          zpsi(1:np, idim) = a(2, 1) * cmplx(xre(j:j+np-1), -xim(j:j+np-1), real64)  + &
            a(2, 2) * cmplx(xre(k:k+np-1), -xim(k:k+np-1), real64)
          j = j + np
          k = k + np
        end do
 
        call zhamiltonian_elec_apply_single(hm_p, namespace_p, mesh_p, zpsi, opzpsi, ist, ik)
 
        do idim = 1, dim
          yre(jj:jj+np-1) = xre(jj:jj+np-1) - aimag(dt_op * opzpsi(1:np, idim))
          yim(jj:jj+np-1) = xim(jj:jj+np-1) - real(dt_op * opzpsi(1:np, idim), real64)
          jj = jj + np
        end do
      end do
    end do
 
    safe_deallocate_a(zpsi)
    safe_deallocate_a(opzpsi)
    pop_sub(td_rk4opt)
  end subroutine td_rk4opt
  ! ---------------------------------------------------------
 
 
  ! ---------------------------------------------------------
  subroutine td_rk2op(xre, xim, yre, yim)
    real(real64), intent(in)  :: xre(:)
    real(real64), intent(in)  :: xim(:)
    real(real64), intent(out) :: yre(:)
    real(real64), intent(out) :: yim(:)
 
    integer :: np_part, np, st1, st2, kp1, kp2, dim, idim, ik, ist, jj, j
    complex(real64), allocatable :: zpsi(:, :)
    complex(real64), allocatable :: opzpsi(:, :)
    complex(real64), allocatable :: zpsi_(:, :, :, :)
 
    push_sub(td_rk2op)
 
    np_part = mesh_p%np_part
    np = mesh_p%np
    st1 = st_p%st_start
    st2 = st_p%st_end
    kp1 = st_p%d%kpt%start
    kp2 = st_p%d%kpt%end
    dim = st_p%d%dim
 
    safe_allocate(zpsi(1:np_part, 1:dim))
    safe_allocate(opzpsi(1:np_part, 1:dim))
    safe_allocate(zpsi_(1:np_part, 1:dim, st1:st2, kp1:kp2))
 
    zpsi = m_z0
    opzpsi = m_z0
 
    call hm_p%ks_pot%restore_potentials(vhxc1_op)
    if (move_ions_op) hm_p%ep%vpsl = vpsl1_op
    call propagation_ops_elec_update_hamiltonian(namespace_p, space_p, st_p, mesh_p, hm_p, ext_partners_p, t_op + dt_op)
 
    if (oct_exchange_enabled(hm_p%oct_exchange)) then
      zpsi_ = m_z0
      j = 1
      do ik = kp1, kp2
        do ist = st1, st2
          jj = j
          do idim = 1, dim
            zpsi_(1:np, idim, ist, ik) = cmplx(xre(j:j+np-1), xim(j:j+np-1), real64)
            j = j + np
          end do
        end do
      end do
      call oct_exchange_prepare(hm_p%oct_exchange, mesh_p, zpsi_, hm_p%xc, hm_p%psolver, namespace_p)
    end if
 
    j = 1
    do ik = kp1, kp2
      do ist = st1, st2
        jj = j
        do idim = 1, dim
          zpsi(1:np, idim) = cmplx(xre(j:j+np-1), xim(j:j+np-1), real64)
          j = j + np
        end do
        call zhamiltonian_elec_apply_single(hm_p, namespace_p, mesh_p, zpsi, opzpsi, ist, ik)
        do idim = 1, dim
          yre(jj:jj+np-1) = xre(jj:jj+np-1) - aimag(dt_op * m_half * opzpsi(1:np, idim))
          yim(jj:jj+np-1) = xim(jj:jj+np-1) + real(dt_op * m_half * opzpsi(1:np, idim), real64)
          jj = jj + np
        end do
      end do
    end do
 
    if (oct_exchange_enabled(hm_p%oct_exchange)) then
      j = 1
      do ik = kp1, kp2
        do ist = st1, st2
          jj = j
          do idim = 1, dim
            zpsi(1:np, idim) = cmplx(xre(j:j+np-1), xim(j:j+np-1), real64)
            j = j + np
          end do
          opzpsi = m_z0
          call oct_exchange_operator(hm_p%oct_exchange, namespace_p, mesh_p, opzpsi, ist, ik)
 
          do idim = 1, dim
            yre(jj:jj+np-1) = yre(jj:jj+np-1) - aimag(dt_op * m_half * opzpsi(1:np, idim))
            yim(jj:jj+np-1) = yim(jj:jj+np-1) + real(dt_op * m_half * opzpsi(1:np, idim), real64)
            jj = jj + np
          end do
        end do
      end do
    end if
 
    safe_deallocate_a(zpsi)
    safe_deallocate_a(opzpsi)
 
    pop_sub(td_rk2op)
  end subroutine td_rk2op
  ! ---------------------------------------------------------
 
 
  ! ---------------------------------------------------------
  subroutine td_rk2opt(xre, xim, yre, yim)
    real(real64), intent(in)  :: xre(:)
    real(real64), intent(in)  :: xim(:)
    real(real64), intent(out) :: yre(:)
    real(real64), intent(out) :: yim(:)
 
    integer :: np_part, np, st1, st2, kp1, kp2, dim, idim, ik, ist, jj, j
    complex(real64), allocatable :: zpsi(:, :)
    complex(real64), allocatable :: opzpsi(:, :)
    complex(real64), allocatable :: zpsi_(:, :, :, :)
 
    push_sub(td_rk2opt)
 
    np_part = mesh_p%np_part
    np = mesh_p%np
    st1 = st_p%st_start
    st2 = st_p%st_end
    kp1 = st_p%d%kpt%start
    kp2 = st_p%d%kpt%end
    dim = st_p%d%dim
 
    safe_allocate(zpsi(1:np_part, 1:dim))
    safe_allocate(opzpsi(1:np_part, 1:dim))
    safe_allocate(zpsi_(1:np_part, 1:dim, st1:st2, kp1:kp2))
 
    zpsi = m_z0
    opzpsi = m_z0
 
    call hm_p%ks_pot%restore_potentials(vhxc1_op)
    if (move_ions_op) hm_p%ep%vpsl = vpsl1_op
    call propagation_ops_elec_update_hamiltonian(namespace_p, space_p, st_p, mesh_p, hm_p, ext_partners_p, t_op + dt_op)
 
    if (oct_exchange_enabled(hm_p%oct_exchange)) then
      zpsi_ = m_z0
      j = 1
      do ik = kp1, kp2
        do ist = st1, st2
          jj = j
          do idim = 1, dim
            zpsi_(1:np, idim, ist, ik) = cmplx(xre(j:j+np-1), -xim(j:j+np-1), real64)
            j = j + np
          end do
        end do
      end do
      call oct_exchange_prepare(hm_p%oct_exchange, mesh_p, zpsi_, hm_p%xc, hm_p%psolver, namespace_p)
    end if
 
    j = 1
    do ik = kp1, kp2
      do ist = st1, st2
        jj = j
        do idim = 1, dim
          zpsi(1:np, idim) = cmplx(xre(j:j+np-1), -xim(j:j+np-1), real64)
          j = j + np
        end do
        call zhamiltonian_elec_apply_single(hm_p, namespace_p, mesh_p, zpsi, opzpsi, ist, ik)
 
        do idim = 1, dim
          yre(jj:jj+np-1) = xre(jj:jj+np-1) - aimag(dt_op * m_half * opzpsi(1:np, idim))
          yim(jj:jj+np-1) = xim(jj:jj+np-1) - real(dt_op * m_half * opzpsi(1:np, idim), real64)
          jj = jj + np
        end do
      end do
    end do
 
    if (oct_exchange_enabled(hm_p%oct_exchange)) then
      j = 1
      do ik = kp1, kp2
        do ist = st1, st2
          jj = j
          do idim = 1, dim
            zpsi(1:np, idim) = cmplx(xre(j:j+np-1), xim(j:j+np-1), real64)
            j = j + np
          end do
          opzpsi = m_z0
          call oct_exchange_operator(hm_p%oct_exchange, namespace_p, mesh_p, opzpsi, ist, ik)
 
          do idim = 1, dim
            yre(jj:jj+np-1) = yre(jj:jj+np-1) - aimag(dt_op * m_half * opzpsi(1:np, idim))
            yim(jj:jj+np-1) = yim(jj:jj+np-1) - real(dt_op * m_half * opzpsi(1:np, idim), real64)
            jj = jj + np
          end do
        end do
      end do
    end if
 
    safe_deallocate_a(zpsi)
    safe_deallocate_a(opzpsi)
    pop_sub(td_rk2opt)
  end subroutine td_rk2opt
  ! ---------------------------------------------------------
 
  subroutine propagator_rk_end()
    push_sub(propagator_rk_end)
 
    call xc_copied_potentials_end(vhxc1_op)
    call xc_copied_potentials_end(vhxc2_op)
 
    pop_sub(propagator_rk_end)
  end subroutine propagator_rk_end
end module propagator_rk_oct_m
 
!! Local Variables:
!! mode: f90
!! coding: utf-8
!! End: