density.F90

Go to the documentation of this file.

!! Copyright (C) 2002-2010 M. Marques, A. Castro, A. Rubio, G. Bertsch, X. Andrade
!!
!! 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 density_oct_m
  use accel_oct_m
  use batch_oct_m
  use batch_ops_oct_m
  use iso_c_binding
  use comm_oct_m
  use debug_oct_m
  use derivatives_oct_m
  use electron_space_oct_m
  use global_oct_m
  use grid_oct_m
  use, intrinsic :: iso_fortran_env
  use kpoints_oct_m
  use lalg_basic_oct_m
  use math_oct_m
  use mesh_oct_m
  use mesh_function_oct_m
  use messages_oct_m
  use mpi_oct_m
  use multicomm_oct_m
  use namespace_oct_m
  use parser_oct_m
  use profiling_oct_m
  use smear_oct_m
  use space_oct_m
  use states_abst_oct_m
  use states_elec_oct_m
  use states_elec_dim_oct_m
  use types_oct_m
  use wfs_elec_oct_m
 
  implicit none
 
  private
 
  public ::                           &
    density_calc_t,                   &
    density_calc_init,                &
    density_calc_accumulate,          &
    density_calc_end,                 &
    density_calc,                     &
    states_elec_freeze_orbitals,      &
    states_elec_total_density,        &
    states_elec_sync_buff_density,    &
    ddensity_accumulate_grad,         &
    zdensity_accumulate_grad,         &
    states_elec_freeze_redistribute_states,&
    states_elec_freeze_adjust_qtot
 
  type density_calc_t
    private
    real(real64), contiguous, pointer :: density(:, :)
    !                                                  !! dimensions (1:np, 1:d%nspin)
    type(states_elec_t),  pointer :: st
    type(grid_t),         pointer :: gr
    type(accel_mem_t)             :: buff_density
    integer(int64)                :: pnp
    logical                       :: packed
  end type density_calc_t
 
contains
 
  !
  subroutine density_calc_init(this, st, gr, density)
    type(density_calc_t),               intent(out)   :: this
    type(states_elec_t),      target,   intent(in)    :: st
    type(grid_t),             target,   intent(in)    :: gr
    real(real64), contiguous, target,   intent(out)   :: density(:, :)
 
    integer :: ip, ispin
    logical :: correct_size
 
    push_sub(density_calc_init)
 
    this%st => st
    this%gr => gr
 
    this%density => density
    !$omp parallel private(ip)
    do ispin = 1, st%d%nspin
      !$omp do
      do ip = 1, gr%np
        this%density(ip, ispin) = m_zero
      end do
      !$omp end do nowait
    end do
    !$omp end parallel
 
    this%packed = .false.
 
    correct_size = ubound(this%density, dim = 1) == this%gr%np .or. &
      ubound(this%density, dim = 1) == this%gr%np_part
    assert(correct_size)
 
    pop_sub(density_calc_init)
  end subroutine density_calc_init
 
  ! ---------------------------------------------------
  !
  subroutine density_calc_pack(this, async)
    type(density_calc_t),           intent(inout)   :: this
    logical, optional,              intent(in)      :: async
 
    push_sub(density_calc_pack)
 
    this%packed = .true.
    this%pnp = accel_padded_size(this%gr%np)
    call accel_create_buffer(this%buff_density, accel_mem_read_write, type_float, this%pnp*this%st%d%nspin)
 
    ! set to zero
    call accel_set_buffer_to_zero(this%buff_density, type_float, this%pnp*this%st%d%nspin, async=async)
 
    pop_sub(density_calc_pack)
  end subroutine density_calc_pack
 
  ! ---------------------------------------------------
  !
  subroutine density_calc_state(this, psib, istin)
    type(density_calc_t),         intent(inout) :: this
    type(wfs_elec_t),             intent(in)    :: psib
    integer,                      intent(in)    :: istin
 
    integer :: ist, ip, ispin, istin_, bsize, gsize
    complex(real64) :: term, psi1, psi2
    real(real64), allocatable :: weight(:)
    type(accel_mem_t) :: buff_weight
    type(accel_kernel_t), pointer :: kernel
 
    push_sub(density_calc_state)
    call profiling_in("CALC_STATE_DENSITY")
 
    ispin = this%st%d%get_spin_index(psib%ik)
 
    istin_= 0
 
    safe_allocate(weight(1:psib%nst))
    do ist = 1, psib%nst
      weight(ist) = this%st%kweights(psib%ik)*this%st%occ(psib%ist(ist), psib%ik)
      if (psib%ist(ist) == istin) istin_ = ist
    end do
 
    assert(istin_ > 0)
 
    if (abs(weight(istin_)) <= m_epsilon) then
      call profiling_out("CALC_STATE_DENSITY")
      pop_sub(density_calc_state)
      return
    end if
 
    !fix ist for the remaining code
    ist = istin_
 
 
    select case (psib%status())
    case (batch_not_packed)
      select case (this%st%d%ispin)
      case (unpolarized, spin_polarized)
        if (states_are_real(this%st)) then
          !$omp parallel do simd schedule(static)
          do ip = 1, this%gr%np
            this%density(ip, ispin) = this%density(ip, ispin) + weight(ist)*psib%dff(ip, 1, ist)**2
          end do
        else
          !$omp parallel do schedule(static)
          do ip = 1, this%gr%np
            this%density(ip, ispin) = this%density(ip, ispin) + weight(ist)* &
              real(conjg(psib%zff(ip, 1, ist))*psib%zff(ip, 1, ist), real64)
          end do
        end if
      case (spinors)
        !$omp parallel do schedule(static) private(psi1, psi2, term)
        do ip = 1, this%gr%np
          psi1 = psib%zff(ip, 1, ist)
          psi2 = psib%zff(ip, 2, ist)
          this%density(ip, 1) = this%density(ip, 1) + weight(ist)*real(conjg(psi1)*psi1, real64)
          this%density(ip, 2) = this%density(ip, 2) + weight(ist)*real(conjg(psi2)*psi2, real64)
          term = weight(ist)*psi1*conjg(psi2)
          this%density(ip, 3) = this%density(ip, 3) + real(term, real64)
          this%density(ip, 4) = this%density(ip, 4) + aimag(term)
        end do
      end select
 
    case (batch_packed)
 
      select case (this%st%d%ispin)
      case (unpolarized, spin_polarized)
        if (states_are_real(this%st)) then
          !$omp parallel do schedule(static)
          do ip = 1, this%gr%np
            this%density(ip, ispin) = this%density(ip, ispin) + weight(ist)*psib%dff_pack(ist, ip)**2
          end do
        else
 
          !$omp parallel do schedule(static)
          do ip = 1, this%gr%np
            this%density(ip, ispin) = this%density(ip, ispin) + weight(ist)* &
              real(conjg(psib%zff_pack(ist, ip))*psib%zff_pack(ist, ip), real64)
          end do
        end if
      case (spinors)
        assert(mod(psib%nst_linear, 2) == 0)
        !$omp parallel do schedule(static) private(psi1, psi2, term)
        do ip = 1, this%gr%np
          psi1 = psib%zff_pack(2*ist - 1, ip)
          psi2 = psib%zff_pack(2*ist,     ip)
          term = weight(ist)*psi1*conjg(psi2)
 
          this%density(ip, 1) = this%density(ip, 1) + weight(ist)*real(conjg(psi1)*psi1, real64)
          this%density(ip, 2) = this%density(ip, 2) + weight(ist)*real(conjg(psi2)*psi2, real64)
          this%density(ip, 3) = this%density(ip, 3) + real(term, real64)
          this%density(ip, 4) = this%density(ip, 4) + aimag(term)
        end do
      end select
 
    case (batch_device_packed)
      if (.not. this%packed) call density_calc_pack(this, async=.true.)
 
      call accel_create_buffer(buff_weight, accel_mem_read_only, type_float, psib%nst)
      call accel_write_buffer(buff_weight, psib%nst, weight, async=.true.)
 
      select case (this%st%d%ispin)
      case (unpolarized, spin_polarized)
        if (states_are_real(this%st)) then
          kernel => kernel_density_real
        else
          kernel => kernel_density_complex
        end if
 
        call accel_set_kernel_arg(kernel, 0, psib%nst)
        call accel_set_kernel_arg(kernel, 1, this%gr%np)
        call accel_set_kernel_arg(kernel, 2, int(this%pnp*(ispin - 1), int32))
        call accel_set_kernel_arg(kernel, 3, buff_weight)
        call accel_set_kernel_arg(kernel, 4, psib%ff_device)
        call accel_set_kernel_arg(kernel, 5, int(log2(psib%pack_size(1)), int32))
        call accel_set_kernel_arg(kernel, 6, this%buff_density)
 
      case (spinors)
        kernel => kernel_density_spinors
 
        call accel_set_kernel_arg(kernel, 0, psib%nst)
        call accel_set_kernel_arg(kernel, 1, this%gr%np)
        call accel_set_kernel_arg(kernel, 2, int(this%pnp, int32))
        call accel_set_kernel_arg(kernel, 3, buff_weight)
        call accel_set_kernel_arg(kernel, 4, psib%ff_device)
        call accel_set_kernel_arg(kernel, 5, int(log2(psib%pack_size(1)), int32))
        call accel_set_kernel_arg(kernel, 6, this%buff_density)
      end select
 
      ! Compute the grid
      bsize = accel_kernel_block_size(kernel)
      call accel_grid_size(this%gr%np, bsize, gsize)
 
      call accel_kernel_run(kernel, (/gsize/), (/bsize/))
 
      call accel_finish()
      call accel_free_buffer(buff_weight)
    end select
 
    safe_deallocate_a(weight)
 
    call profiling_out("CALC_STATE_DENSITY")
 
    pop_sub(density_calc_state)
  end subroutine density_calc_state
 
  ! ---------------------------------------------------
  !
  subroutine density_calc_accumulate(this, psib, async)
    type(density_calc_t),         intent(inout) :: this
    type(wfs_elec_t),             intent(in)    :: psib
    logical, optional,            intent(in)    :: async
 
    integer :: ist, ip, ispin, bsize, gsize
    complex(real64)   :: term, psi1, psi2
    real(real64), allocatable :: weight(:)
    type(accel_mem_t) :: buff_weight
    type(accel_kernel_t), pointer :: kernel
 
    push_sub(density_calc_accumulate)
    call profiling_in("CALC_DENSITY")
 
    ispin = this%st%d%get_spin_index(psib%ik)
 
    safe_allocate(weight(1:psib%nst))
    do ist = 1, psib%nst
      weight(ist) = this%st%kweights(psib%ik)*this%st%occ(psib%ist(ist), psib%ik)
    end do
 
    select case (psib%status())
    case (batch_not_packed)
      select case (this%st%d%ispin)
      case (unpolarized, spin_polarized)
        if (states_are_real(this%st)) then
          do ist = 1, psib%nst
            if (abs(weight(ist)) <= m_epsilon) cycle
            !$omp parallel do simd schedule(static)
            do ip = 1, this%gr%np
              this%density(ip, ispin) = this%density(ip, ispin) + weight(ist)*psib%dff(ip, 1, ist)**2
            end do
          end do
        else
          do ist = 1, psib%nst
            if (abs(weight(ist)) <= m_epsilon) cycle
            !$omp parallel do schedule(static)
            do ip = 1, this%gr%np
              this%density(ip, ispin) = this%density(ip, ispin) + weight(ist)* &
                real(conjg(psib%zff(ip, 1, ist))*psib%zff(ip, 1, ist), real64)
            end do
          end do
        end if
      case (spinors)
        do ist = 1, psib%nst
          if (abs(weight(ist)) <= m_epsilon) cycle
          !$omp parallel do schedule(static) private(psi1, psi2, term)
          do ip = 1, this%gr%np
            psi1 = psib%zff(ip, 1, ist)
            psi2 = psib%zff(ip, 2, ist)
            this%density(ip, 1) = this%density(ip, 1) + weight(ist)*real(conjg(psi1)*psi1, real64)
            this%density(ip, 2) = this%density(ip, 2) + weight(ist)*real(conjg(psi2)*psi2, real64)
            term = weight(ist)*psi1*conjg(psi2)
            this%density(ip, 3) = this%density(ip, 3) + real(term, real64)
            this%density(ip, 4) = this%density(ip, 4) + aimag(term)
          end do
        end do
      end select
 
    case (batch_packed)
 
      select case (this%st%d%ispin)
      case (unpolarized, spin_polarized)
        if (states_are_real(this%st)) then
          !$omp parallel do schedule(static) private(ist)
          do ip = 1, this%gr%np
            do ist = 1, psib%nst
              this%density(ip, ispin) = this%density(ip, ispin) + weight(ist)*psib%dff_pack(ist, ip)**2
            end do
          end do
        else
          !$omp parallel do schedule(static) private(ist)
          do ip = 1, this%gr%np
            do ist = 1, psib%nst
              this%density(ip, ispin) = this%density(ip, ispin) + weight(ist)* &
                real(conjg(psib%zff_pack(ist, ip))*psib%zff_pack(ist, ip), real64)
            end do
          end do
        end if
      case (spinors)
        assert(mod(psib%nst_linear, 2) == 0)
        !$omp parallel do schedule(static) private(ist, psi1, psi2, term)
        do ip = 1, this%gr%np
          do ist = 1, psib%nst
            psi1 = psib%zff_pack(2*ist - 1, ip)
            psi2 = psib%zff_pack(2*ist,     ip)
            term = weight(ist)*psi1*conjg(psi2)
 
            this%density(ip, 1) = this%density(ip, 1) + weight(ist)*real(conjg(psi1)*psi1, real64)
            this%density(ip, 2) = this%density(ip, 2) + weight(ist)*real(conjg(psi2)*psi2, real64)
            this%density(ip, 3) = this%density(ip, 3) + real(term, real64)
            this%density(ip, 4) = this%density(ip, 4) + aimag(term)
          end do
        end do
      end select
 
    case (batch_device_packed)
      if (.not. this%packed) call density_calc_pack(this)
 
      call accel_create_buffer(buff_weight, accel_mem_read_only, type_float, psib%nst)
      call accel_write_buffer(buff_weight, psib%nst, weight, async=.true.)
 
      select case (this%st%d%ispin)
      case (unpolarized, spin_polarized)
        if (states_are_real(this%st)) then
          kernel => kernel_density_real
        else
          kernel => kernel_density_complex
        end if
 
        call accel_set_kernel_arg(kernel, 0, psib%nst)
        call accel_set_kernel_arg(kernel, 1, this%gr%np)
        call accel_set_kernel_arg(kernel, 2, int(this%pnp, int32)*(ispin - 1))
        call accel_set_kernel_arg(kernel, 3, buff_weight)
        call accel_set_kernel_arg(kernel, 4, psib%ff_device)
        call accel_set_kernel_arg(kernel, 5, int(log2(psib%pack_size(1)), int32))
        call accel_set_kernel_arg(kernel, 6, this%buff_density)
 
      case (spinors)
        kernel => kernel_density_spinors
 
        call accel_set_kernel_arg(kernel, 0, psib%nst)
        call accel_set_kernel_arg(kernel, 1, this%gr%np)
        call accel_set_kernel_arg(kernel, 2, int(this%pnp, int32))
        call accel_set_kernel_arg(kernel, 3, buff_weight)
        call accel_set_kernel_arg(kernel, 4, psib%ff_device)
        call accel_set_kernel_arg(kernel, 5, int(log2(psib%pack_size(1)), int32))
        call accel_set_kernel_arg(kernel, 6, this%buff_density)
      end select
 
      ! Compute the grid
      bsize = accel_kernel_block_size(kernel)
      call accel_grid_size(this%gr%np, bsize, gsize)
 
      call accel_kernel_run(kernel, (/gsize/), (/bsize/))
 
      call accel_free_buffer(buff_weight)
 
      if(.not. optional_default(async, .false.)) call accel_finish()
 
    end select
 
    safe_deallocate_a(weight)
 
 
    if (this%st%d%ispin /= spinors) then
      if (states_are_real(this%st)) then
        ! 1 add (1), 2 mult (1)
        call profiling_count_operations(psib%nst*this%gr%np*(1 + 2))
      else
        ! 1 add (2), 1 complex mult (6), 1 real-complex mult (2)
        call profiling_count_operations(psib%nst*this%gr%np*(2 + 6 + 2))
      end if
    else
      ! 4 add (2), 3 complex mult (6), 3 real-complex mult (2)
      call profiling_count_operations(psib%nst*this%gr%np*(4*2 + 3*6 + 3*2))
    end if
 
    call profiling_out("CALC_DENSITY")
 
    pop_sub(density_calc_accumulate)
  end subroutine density_calc_accumulate
 
  ! ---------------------------------------------------
  !
  subroutine density_calc_end(this, allreduce, symmetrize, buff_density)
    type(density_calc_t), intent(inout)        :: this
    logical, optional,    intent(in)           :: allreduce
    logical, optional,    intent(in)           :: symmetrize
    type(accel_mem_t), optional, intent(inout) :: buff_density
 
    integer :: ispin, ip
    real(real64), allocatable :: tmpdensity(:)
    logical :: was_packed
 
    logical:: allreduce_local, symmetrize_local
 
    push_sub(density_calc_end)
 
    call profiling_in("CALC_DENSITY_END")
 
    allreduce_local = (this%st%parallel_in_states .or. this%st%d%kpt%parallel) .and. optional_default(allreduce, .true.)
    symmetrize_local = this%st%symmetrize_density .and. optional_default(symmetrize, .true.)
 
    was_packed = this%packed  ! Save before modifying
    if (this%packed) then
      this%packed = .false.
 
      safe_allocate(tmpdensity(1:this%gr%np))
 
      ! the density is in device memory
      do ispin = 1, this%st%d%nspin
        call accel_read_buffer(this%buff_density, int(this%gr%np, int64), tmpdensity, offset = (ispin - 1)*this%pnp)
 
        !$omp parallel do simd
        do ip = 1, this%gr%np
          this%density(ip, ispin) = this%density(ip, ispin) + tmpdensity(ip)
        end do
      end do
 
      safe_deallocate_a(tmpdensity)
 
      ! If buff_density is not needed, free the GPU buffer now
      if (.not. present(buff_density)) then
        call accel_free_buffer(this%buff_density)
      end if
    end if
 
    ! reduce over states and k-points
    if (allreduce_local) then
      call profiling_in("DENSITY_REDUCE")
      call comm_allreduce(this%st%st_kpt_mpi_grp, this%density, dim = (/this%gr%np, this%st%d%nspin/))
      call profiling_out("DENSITY_REDUCE")
    end if
 
    if (symmetrize_local) then
      do ispin = 1, this%st%d%nspin
        call dgrid_symmetrize_scalar_field(this%gr, this%density(:, ispin))
      end do
    end if
 
    if (present(buff_density) .and. was_packed) then
      if (allreduce_local .or. symmetrize_local) then
        do ispin = 1, this%st%d%nspin
          call accel_write_buffer(this%buff_density, int(this%gr%np, int64), &
            this%density(1:this%gr%np, ispin), &
            offset=int((ispin - 1)*this%pnp, int64))
        end do
      end if
 
      ! Free any previously stored buffer before overwriting
      if (buff_density%allocated) call accel_free_buffer(buff_density)
 
      call accel_move_buffer(this%buff_density, buff_density)
    end if
 
    call profiling_out("CALC_DENSITY_END")
 
    pop_sub(density_calc_end)
  end subroutine density_calc_end
 
 
  ! ---------------------------------------------------------
  !
  subroutine density_calc(st, gr, density, istin)
    type(states_elec_t),      intent(inout) :: st
    type(grid_t),             intent(in)    :: gr
    real(real64), contiguous, intent(out)   :: density(:, :)
    integer, optional,        intent(in)    :: istin
 
    integer :: ik, ib
    type(density_calc_t) :: dens_calc
 
    push_sub(density_calc)
 
    assert(ubound(density, dim = 1) == gr%np .or. ubound(density, dim = 1) == gr%np_part)
 
    call density_calc_init(dens_calc, st, gr, density)
 
    if (.not. present(istin)) then
      ! ordinary density accumulate
      do ik = st%d%kpt%start, st%d%kpt%end
        do ib = st%group%block_start, st%group%block_end
          call density_calc_accumulate(dens_calc, st%group%psib(ib, ik), async=.true.)
        end do
      end do
      call accel_finish()
    else
      ! state-resolved density
      do ik = st%d%kpt%start, st%d%kpt%end
        do ib = st%group%block_start, st%group%block_end
          if (st%group%psib(ib, ik)%ist(1) <= istin .and. st%group%psib(ib, ik)%ist(st%group%psib(ib, ik)%nst) >= istin) then
            call density_calc_state(dens_calc, st%group%psib(ib, ik), istin)
          end if
        end do
      end do
    end if
 
    call density_calc_end(dens_calc, allreduce = .not. present(istin), buff_density = st%buff_density)
 
    pop_sub(density_calc)
  end subroutine density_calc
 
  ! ---------------------------------------------------------
  subroutine states_elec_freeze_orbitals(st, namespace, space, gr, mc, kpoints, n, family_is_mgga)
    type(states_elec_t), intent(inout) :: st
    type(namespace_t),   intent(in)    :: namespace
    class(space_t),      intent(in)    :: space
    type(grid_t),        intent(in)    :: gr
    type(multicomm_t),   intent(in)    :: mc
    type(kpoints_t),     intent(in)    :: kpoints
    integer,             intent(in)    :: n
    logical,             intent(in)    :: family_is_mgga
 
    integer :: ist, ik, ib, nblock
    integer :: nodeto, nodefr
    type(states_elec_t) :: staux
    complex(real64), allocatable :: psi(:, :, :), rec_buffer(:,:)
    type(wfs_elec_t)  :: psib
    type(density_calc_t) :: dens_calc
 
    push_sub(states_elec_freeze_orbitals)
 
    if (n >= st%nst) then
      write(message(1),'(a)') 'Attempting to freeze a number of orbitals which is larger or equal to'
      write(message(2),'(a)') 'the total number. The program has to stop.'
      call messages_fatal(2, namespace=namespace)
    end if
 
    assert(states_are_complex(st))
 
    if (.not. allocated(st%frozen_rho)) then
      safe_allocate(st%frozen_rho(1:gr%np, 1:st%d%nspin))
    end if
 
    call density_calc_init(dens_calc, st, gr, st%frozen_rho)
 
    do ik = st%d%kpt%start, st%d%kpt%end
      if (n < st%st_start) cycle
 
      do ib =  st%group%block_start, st%group%block_end
        !We can use the full batch
        if (states_elec_block_max(st, ib) <= n) then
 
          call density_calc_accumulate(dens_calc, st%group%psib(ib, ik))
          if (states_elec_block_max(st, ib) == n) exit
 
        else !Here we only use a part of this batch
 
          nblock = n - states_elec_block_min(st, ib) + 1
 
          safe_allocate(psi(1:gr%np, 1:st%d%dim, 1:nblock))
 
          do ist = 1, nblock
            call states_elec_get_state(st, gr, states_elec_block_min(st, ib) + ist - 1, ik, psi(:, :, ist))
          end do
 
          call wfs_elec_init(psib, st%d%dim, states_elec_block_min(st, ib), n, psi, ik)
          call density_calc_accumulate(dens_calc, psib)
          call psib%end()
          safe_deallocate_a(psi)
 
          exit
 
        end if
 
      end do
    end do
 
    call density_calc_end(dens_calc)
 
    if (family_is_mgga) then
      if (.not. allocated(st%frozen_tau)) then
        safe_allocate(st%frozen_tau(1:gr%np, 1:st%d%nspin))
      end if
      if (.not. allocated(st%frozen_gdens)) then
        safe_allocate(st%frozen_gdens(1:gr%np, 1:space%dim, 1:st%d%nspin))
      end if
      if (.not. allocated(st%frozen_ldens)) then
        safe_allocate(st%frozen_ldens(1:gr%np, 1:st%d%nspin))
      end if
 
      call states_elec_calc_quantities(gr, st, kpoints, .true., kinetic_energy_density = st%frozen_tau, &
        density_gradient = st%frozen_gdens, density_laplacian = st%frozen_ldens, st_end = n)
    end if
 
 
    call states_elec_copy(staux, st)
 
    call states_elec_freeze_redistribute_states(st, namespace, gr, mc, n)
 
    safe_allocate(psi(1:gr%np, 1:st%d%dim, 1))
    safe_allocate(rec_buffer(1:gr%np, 1:st%d%dim))
 
    do ik = st%d%kpt%start, st%d%kpt%end
      ! loop over all frozen states
      do ist = 1, st%nst
        ! new distribution: frozen states
        nodeto = st%node(ist)
        ! old distribution: full states
        nodefr = staux%node(ist+n)
 
        if (st%mpi_grp%rank == nodeto .and. st%mpi_grp%rank == nodefr) then
          ! do a simple copy on the same rank, also for serial version
          call states_elec_get_state(staux, gr, ist+n, ik, psi(:, :, 1))
          call states_elec_set_state(st, gr, ist, ik, psi(:, :, 1))
        else
          ! we can use blocking send/recv calls here because we always have different ranks
          if (st%mpi_grp%rank == nodefr) then
            call states_elec_get_state(staux, gr, ist+n, ik, psi(:, :, 1))
            call st%mpi_grp%send(psi(1, 1, 1), gr%np*st%d%dim, mpi_double_complex, nodeto, ist)
          end if
          if (st%mpi_grp%rank == nodeto) then
            call st%mpi_grp%recv(rec_buffer(1, 1), gr%np*st%d%dim, mpi_double_complex, nodefr, ist)
            call states_elec_set_state(st, gr, ist, ik, rec_buffer(:, :))
          end if
        end if
      end do
    end do
 
    safe_deallocate_a(psi)
    safe_deallocate_a(rec_buffer)
 
    do ik = 1, st%nik
      do ist = 1, st%nst
        st%occ(ist, ik) = staux%occ(n+ist, ik)
        st%eigenval(ist, ik) = staux%eigenval(n+ist, ik)
      end do
    end do
 
    call states_elec_end(staux)
    pop_sub(states_elec_freeze_orbitals)
  end subroutine states_elec_freeze_orbitals
 
  ! ---------------------------------------------------------
  subroutine states_elec_freeze_redistribute_states(st, namespace, mesh, mc, nn)
    type(states_elec_t), intent(inout) :: st
    type(namespace_t),   intent(in)    :: namespace
    class(mesh_t),       intent(in)    :: mesh
    type(multicomm_t),   intent(in)    :: mc
    integer,             intent(in)    :: nn
 
    push_sub(states_elec_freeze_redistribute_states)
 
    st%nst = st%nst - nn
 
    call states_elec_deallocate_wfns(st)
    call states_elec_distribute_nodes(st, namespace, mc)
    call states_elec_allocate_wfns(st, mesh, type_cmplx, packed=.true.)
 
    safe_deallocate_a(st%eigenval)
    safe_allocate(st%eigenval(1:st%nst, 1:st%nik))
    st%eigenval = huge(st%eigenval)
 
    safe_deallocate_a(st%occ)
    safe_allocate(st%occ     (1:st%nst, 1:st%nik))
    st%occ      = m_zero
 
 
    pop_sub(states_elec_freeze_redistribute_states)
  end subroutine states_elec_freeze_redistribute_states
 
  ! ---------------------------------------------------------
  subroutine states_elec_freeze_adjust_qtot(st)
    type(states_elec_t), intent(inout) :: st
 
    integer :: ik, ist
 
    push_sub(states_elec_freeze_adjust_qtot)
 
    ! Change the smearing method by fixing the occupations to
    ! that of the ground-state such that the unfrozen states inherit
    ! those values.
    st%smear%method = smear_fixed_occ
 
    ! Set total charge
    st%qtot = m_zero
    do ik = st%d%kpt%start, st%d%kpt%end
      do ist = st%st_start, st%st_end
        st%qtot = st%qtot + st%occ(ist, ik) * st%kweights(ik)
      end do
    end do
 
    if (st%parallel_in_states .or. st%d%kpt%parallel) then
      call comm_allreduce(st%st_kpt_mpi_grp, st%qtot)
    end if
 
 
    pop_sub(states_elec_freeze_adjust_qtot)
  end subroutine states_elec_freeze_adjust_qtot
 
 
  ! ---------------------------------------------------------
  !
  subroutine states_elec_total_density(st, mesh, total_rho)
    type(states_elec_t),      intent(in)  :: st
    class(mesh_t),            intent(in)  :: mesh
    real(real64), contiguous, intent(out) :: total_rho(:,:)
 
    integer :: is
 
    push_sub(states_elec_total_density)
 
    call lalg_copy(mesh%np, st%d%nspin, st%rho, total_rho)
 
    if (allocated(st%rho_core)) then
      do is = 1, st%d%spin_channels
        call lalg_axpy(mesh%np, m_one/st%d%spin_channels, st%rho_core, total_rho(:, is))
      end do
    end if
 
    ! Add, if it exists, the frozen density from the inner orbitals.
    if (allocated(st%frozen_rho)) then
      call lalg_axpy(mesh%np, st%d%nspin, m_one, st%frozen_rho, total_rho)
    end if
 
    pop_sub(states_elec_total_density)
  end subroutine states_elec_total_density
 
  ! ---------------------------------------------------------
  subroutine states_elec_sync_buff_density(st, mesh)
    type(states_elec_t),      intent(inout) :: st
    class(mesh_t),            intent(in)     :: mesh
 
    integer :: is
    integer(int64) :: pnp
 
    push_sub(states_elec_sync_buff_density)
 
    if (accel_is_enabled() .and. st%buff_density%allocated) then
      pnp = accel_padded_size(mesh%np)
      do is = 1, st%d%nspin
        call accel_write_buffer(st%buff_density, int(mesh%np, int64), st%rho(1:mesh%np, is), &
          offset = (is - 1)*pnp)
      end do
    end if
 
    pop_sub(states_elec_sync_buff_density)
  end subroutine states_elec_sync_buff_density
 
#include "undef.F90"
#include "real.F90"
#include "density_inc.F90"
 
#include "undef.F90"
#include "complex.F90"
#include "density_inc.F90"
 
end module density_oct_m
 
 
!! Local Variables:
!! mode: f90
!! coding: utf-8
!! End: