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 epot_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 nonlocal_pseudopotential_oct_m
  use parser_oct_m
  use phase_oct_m
  use profiling_oct_m
  use projector_oct_m
  use projector_matrix_oct_m
  use ps_oct_m, only: proj_uspp
  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 => null()  
    type(grid_t),         pointer :: gr => null()  
    type(epot_t),         pointer :: ep => null()  
    type(nonlocal_pseudopotential_t), pointer :: vnl => null() 
    type(phase_t),        pointer :: phase => null() 
    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, ep, vnl, phase)
    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(:, :)
    type(epot_t), optional,   target,   intent(in)    :: ep
    type(nonlocal_pseudopotential_t), optional, target, intent(in) :: vnl
    type(phase_t), optional,  target,   intent(in)    :: phase
 
    integer :: ip, ispin
    logical :: correct_size
 
    push_sub(density_calc_init)
 
    this%st => st
    this%gr => gr
    nullify(this%ep)
    nullify(this%vnl)
    nullify(this%phase)
 
    if (st%requires_generalized_overlap) then
      if (.not. present(ep) .or. .not. present(vnl)) then
        call messages_not_implemented("density calculation with ultrasoft/PAW augmentation density")
      end if
      this%ep => ep
      this%vnl => vnl
      if (present(phase)) this%phase => phase
    end if
 
    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
 
    if (this%st%requires_generalized_overlap) then
      call density_calc_accumulate_uspp(this, psib, ispin, weight)
    end if
 
    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
 
    ! Ultrasoft/PAW augmentation density: n(r) += sum_ij Q_ij(r) <psi|beta_i><beta_j|psi>.
    ! This must be done on the batch path used by the SCF (density_calc), not only in
    ! density_calc_state, otherwise the augmentation charge is silently dropped.
    if (this%st%requires_generalized_overlap) then
      call density_calc_accumulate_uspp(this, psib, ispin, weight)
    end if
 
    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_accumulate_uspp(this, psib, ispin, weight)
    type(density_calc_t), intent(inout) :: this
    type(wfs_elec_t),     intent(in)    :: psib
    integer,              intent(in)    :: ispin
    real(real64),         intent(in)    :: weight(:)
 
    integer :: imat, iatom, iproj, jproj, ltot, mtot, is, ip, pair_index, lm, np, lmaxq
    real(real64) :: coeff, val, cval
    real(real64), allocatable :: beta_density(:, :, :)
    real(real64), allocatable :: acc(:)
    type(projector_matrix_t), pointer :: pmat
    type(projector_t),        pointer :: proj
    type(wfs_elec_t)             :: phased_psib
    logical                      :: do_phase
 
    push_sub_with_profile(density_calc_accumulate_uspp)
 
    assert(associated(this%ep))
    assert(associated(this%vnl))
    assert(this%st%d%ispin /= spinors)
    assert(this%st%d%dim == 1)
    assert(.not. this%gr%use_curvilinear)
    assert(this%vnl%apply_projector_matrices)
 
    if (psib%status() == batch_device_packed) then
      call messages_not_implemented("USPP augmentation density for accelerator batches")
    end if
 
    ! The projections <beta|psi> that build the augmentation density matrix must
    ! include the Bloch phase exp(ik.r) at k /= 0, exactly as in the Hamiltonian and
    ! overlap apply.
    do_phase = .false.
    if (associated(this%phase)) do_phase = this%phase%is_allocated()
 
    if (states_are_real(this%st)) then
      call dnonlocal_pseudopotential_density_matrix(this%vnl, this%gr, this%st%d, .false., psib, weight, beta_density)
    else if (do_phase .and. .not. psib%has_phase) then
      call psib%copy_to(phased_psib)
      call this%phase%apply_to(this%gr, this%gr%np, .false., phased_psib, src = psib)
      phased_psib%has_phase = .true.
      call znonlocal_pseudopotential_density_matrix(this%vnl, this%gr, this%st%d, .false., phased_psib, weight, beta_density)
      call phased_psib%end()
    else
      call znonlocal_pseudopotential_density_matrix(this%vnl, this%gr, this%st%d, .false., psib, weight, beta_density)
    end if
 
    do imat = 1, this%vnl%nprojector_matrices
      pmat => this%vnl%projector_matrices(imat)
      iatom = this%vnl%projector_to_atom(imat)
      proj => this%ep%proj(iatom)
 
      assert(projector_is(proj, proj_uspp))
      assert(allocated(proj%uspp_qfunc))
      assert(proj%uspp_qfunc_nbeta > 0)
      assert(proj%uspp_qfunc_lmax >= 0)
 
      np = pmat%npoints
      lmaxq = proj%uspp_qfunc_lmax
 
      ! Real spherical harmonics Y_LM at the sphere points depend only on the (fixed)
      ! sphere geometry. By default they are computed once and cached in the projector
      ! matrix (shared with nonlocal_pseudopotential_update_dij; reset when projectors
      ! rebuild); with USPPSaveMemory the cache is freed at the end of the call.
      ! Index lm = L*(L+1) + M + 1, L = 0..lmaxq, M = -L..L.
      if (.not. allocated(pmat%uspp_ylm)) then
        safe_allocate(pmat%uspp_ylm(1:np, 1:(lmaxq + 1)**2))
        do ltot = 0, lmaxq
          do mtot = -ltot, ltot
            lm = ltot*(ltot + 1) + mtot + 1
            do is = 1, np
              call ylmr_real(pmat%position(1:3, is), ltot, mtot, pmat%uspp_ylm(is, lm))
            end do
          end do
        end do
      end if
 
      safe_allocate(acc(1:np))
      acc(1:np) = m_zero
      do iproj = 1, pmat%nprojs
        assert(pmat%beta_index(iproj) >= 1 .and. pmat%beta_index(iproj) <= proj%uspp_qfunc_nbeta)
        do jproj = 1, pmat%nprojs
          assert(pmat%beta_index(jproj) >= 1 .and. pmat%beta_index(jproj) <= proj%uspp_qfunc_nbeta)
          val = beta_density(iproj, jproj, imat)
          if (abs(val) <= m_epsilon) cycle
 
          pair_index = (pmat%beta_index(iproj) - 1)*proj%uspp_qfunc_nbeta + pmat%beta_index(jproj)
 
          ! Gaunt coefficients depend only on the angular momenta, so hoist them out
          ! of the grid-point loop. The remaining inner loop over sphere points is
          ! contiguous and branch-free (a zero Q_ij^L simply contributes zero).
          do ltot = 0, lmaxq
            do mtot = -ltot, ltot
              coeff = real_gaunt(pmat%l(iproj), pmat%m(iproj), pmat%l(jproj), pmat%m(jproj), ltot, mtot)
              if (abs(coeff) <= 1.0e-12_real64) cycle
              cval = val*coeff
              lm = ltot*(ltot + 1) + mtot + 1
              !$omp simd
              do is = 1, np
                acc(is) = acc(is) + cval*pmat%uspp_ylm(is, lm)*proj%uspp_qfunc(is, pair_index, ltot)
              end do
            end do
          end do
        end do
      end do
 
      ! Scatter the accumulated augmentation density onto the mesh (once per atom).
      do is = 1, np
        ip = pmat%map(is)
        this%density(ip, ispin) = this%density(ip, ispin) + acc(is)
      end do
 
      safe_deallocate_a(acc)
      ! With USPPSaveMemory we do not retain the Y_LM cache between calls.
      if (this%vnl%uspp_save_memory) then
        safe_deallocate_a(pmat%uspp_ylm)
      end if
    end do
 
    safe_deallocate_a(beta_density)
 
    pop_sub_with_profile(density_calc_accumulate_uspp)
  end subroutine density_calc_accumulate_uspp
 
  ! ---------------------------------------------------
  !
  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, ep, vnl, phase)
    type(states_elec_t),      intent(inout) :: st
    type(grid_t),             intent(in)    :: gr
    real(real64), contiguous, intent(out)   :: density(:, :)
    integer, optional,        intent(in)    :: istin
    type(epot_t), optional, target, intent(in) :: ep
    type(nonlocal_pseudopotential_t), optional, target, intent(in) :: vnl
    type(phase_t), optional, target, intent(in) :: phase
 
    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)
 
    if (st%requires_generalized_overlap .and. present(istin)) then
      call messages_not_implemented("state-resolved density with ultrasoft/PAW augmentation density")
    end if
 
    call density_calc_init(dens_calc, st, gr, density, ep, vnl, phase)
 
    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 (st%requires_generalized_overlap) then
      call messages_not_implemented("frozen-orbital density with ultrasoft/PAW augmentation density")
    end if
 
    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: