states_elec.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 states_elec_oct_m
  use accel_oct_m
  use batch_oct_m
  use batch_ops_oct_m
  use blacs_proc_grid_oct_m
  use boundaries_oct_m
  use batch_oct_m
  use batch_ops_oct_m
  use calc_mode_par_oct_m
  use comm_oct_m
  use debug_oct_m
  use derivatives_oct_m
  use distributed_oct_m
  use electron_space_oct_m
  use global_oct_m
  use grid_oct_m
  use io_oct_m
  use, intrinsic :: iso_fortran_env
  use kpoints_oct_m
  use lalg_basic_oct_m
  use loct_oct_m
  use math_oct_m
  use mesh_oct_m
  use mesh_batch_oct_m
  use mesh_function_oct_m
  use messages_oct_m
  use modelmb_particles_oct_m
  use mpi_oct_m
  use multicomm_oct_m
  use namespace_oct_m
#ifdef HAVE_OPENMP
  use omp_lib
#endif
  use parser_oct_m
  use profiling_oct_m
  use quickrnd_oct_m
  use restart_oct_m
  use smear_oct_m
  use space_oct_m
  use states_abst_oct_m
  use states_elec_group_oct_m
  use states_elec_dim_oct_m
  use types_oct_m
  use unit_oct_m
  use unit_system_oct_m
  use varinfo_oct_m
  use wfs_elec_oct_m
 
  implicit none
 
  private
 
  public ::                           &
    states_elec_t,                         &
    states_elec_init,                      &
    states_elec_look,                      &
    states_elec_densities_init,            &
    states_elec_exec_init,                 &
    states_elec_allocate_wfns,             &
    states_elec_allocate_current,          &
    states_elec_deallocate_wfns,           &
    states_elec_null,                      &
    states_elec_end,                       &
    states_elec_copy,                      &
    states_elec_generate_random,           &
    states_elec_fermi,                     &
    states_elec_eigenvalues_sum,           &
    states_elec_calc_quantities,           &
    state_is_local,                        &
    state_kpt_is_local,                    &
    states_elec_choose_kpoints,            &
    states_elec_distribute_nodes,          &
    states_elec_wfns_memory,               &
    states_elec_get_state,                 &
    states_elec_set_state,                 &
    states_elec_get_points,                &
    states_elec_block_min,                 &
    states_elec_block_max,                 &
    states_elec_block_size,                &
    states_elec_all_states_gpu_packed,     &
    occupied_states,                       &
    states_elec_set_zero,                  &
    states_elec_generate_random_vector,    &
    stress_t,                              &
    kpoints_distribute
 
 
  ! this type must be moved to stress module but due to circular dependency it is not possible now
  type stress_t
    real(real64)  :: total(3,3)           = m_zero
    real(real64)  :: kinetic(3,3)         = m_zero
    real(real64)  :: Hartree(3,3)         = m_zero
    real(real64)  :: xc(3,3)              = m_zero
    real(real64)  :: xc_nlcc(3,3)         = m_zero
    real(real64)  :: ps_local(3,3)        = m_zero
    real(real64)  :: ps_nl(3,3)           = m_zero
    real(real64)  :: ion_ion(3,3)         = m_zero
    real(real64)  :: vdw(3,3)             = m_zero
    real(real64)  :: hubbard(3,3)         = m_zero
 
    real(real64)  :: kinetic_sumrule = m_zero 
    real(real64)  :: hartree_sumrule = m_zero 
  end type stress_t
 
  ! TODO(Alex) Issue #672 Decouple k-point info from `states_elec_dim_t`
 
  !
  type, extends(states_abst_t) :: states_elec_t
    ! Components are public by default
    type(states_elec_dim_t)  :: d
    integer                  :: nst_conv
 
    logical                  :: only_userdef_istates
 
    type(states_elec_group_t) :: group
    integer :: block_size
    logical :: pack_states
 
 
    character(len=1024), allocatable :: user_def_states(:,:,:)
 
    ! TODO(Alex) Issue #821.  Collate current quantities into an object.
    ! the densities and currents (after all we are doing DFT :)
    real(real64), allocatable :: rho(:,:)
    real(real64), allocatable :: rho_core(:)
 
    ! Buffer density
    type(accel_mem_t) :: buff_density
 
    real(real64), allocatable :: current(:, :, :)
    real(real64), allocatable :: current_para(:, :, :)
    real(real64), allocatable :: current_dia(:, :, :)
    real(real64), allocatable :: current_mag(:, :, :)
    real(real64), allocatable :: current_kpt(:,:,:)
 
    ! TODO(Alex) Issue #673. Create frozen density class and replace in states_elec_t
    ! It may be required to "freeze" the deepest orbitals during the evolution; the density
    ! of these orbitals is kept in frozen_rho. It is different from rho_core.
    real(real64), allocatable :: frozen_rho(:, :)
    real(real64), allocatable :: frozen_tau(:, :)
    real(real64), allocatable :: frozen_gdens(:,:,:)
    real(real64), allocatable :: frozen_ldens(:,:)
 
    logical            :: uniform_occ
 
    real(real64), allocatable :: eigenval(:,:)
    logical            :: fixed_occ
    logical            :: restart_fixed_occ
    logical            :: restart_reorder_occs
    real(real64), allocatable :: occ(:,:)
    real(real64), allocatable :: kweights(:)
    integer            :: nik
 
    logical            :: fixed_spins
    real(real64), allocatable :: spin(:, :, :)
 
    real(real64)   :: qtot
    real(real64)   :: val_charge
 
    type(stress_t) :: stress_tensors
 
    logical        :: fromScratch
    type(smear_t)  :: smear
 
    ! TODO(Alex) Issue #823 Move modelmbparticles out of states_elec_t
    type(modelmb_particle_t) :: modelmbparticles
 
    ! TODO(Alex) Issue #824. Package the communicators in a single instance prior to removing
    ! or consider creating a distributed_t instance for each (as distributed_t contains the an instance of mpi_grp_t)
    type(mpi_grp_t)             :: mpi_grp
    type(mpi_grp_t)             :: dom_st_mpi_grp
    type(mpi_grp_t)             :: st_kpt_mpi_grp
    type(mpi_grp_t)             :: dom_st_kpt_mpi_grp
    type(mpi_grp_t)             :: system_grp
    type(blacs_proc_grid_t)     :: dom_st_proc_grid
 
    type(distributed_t)         :: dist
    logical                     :: scalapack_compatible
    logical                     :: parallel_in_states = .false. 
 
    ! TODO(Alex) Issue #820. Remove lnst, st_start, st_end and node, as they are all contained within dist
    integer                     :: lnst
    integer                     :: st_start, st_end
    integer, allocatable        :: node(:)
 
    ! TODO(Alex) Issue #824. Either change from data to a method, or package with `st_kpt_mpi_grp`
    integer, allocatable        :: st_kpt_task(:,:)
 
    type(multicomm_all_pairs_t), private :: ap
    logical                     :: symmetrize_density
    integer                     :: randomization
    integer                     :: orth_method = 0
 
    real(real64)   :: gpu_states_mem
 
  contains
    procedure :: nullify => states_elec_null            
    procedure :: write_info => states_elec_write_info   
    procedure :: pack => states_elec_pack               
    procedure :: unpack => states_elec_unpack           
    procedure :: set_zero => states_elec_set_zero       
    procedure :: dipole => states_elec_calculate_dipole 
  end type states_elec_t
 
  integer, public, parameter :: &
    par_independent = 1,              &
    par_dependent   = 2
 
 
  interface states_elec_get_state
    module procedure dstates_elec_get_state1, zstates_elec_get_state1, dstates_elec_get_state2, zstates_elec_get_state2
    module procedure dstates_elec_get_state3, zstates_elec_get_state3, dstates_elec_get_state4, zstates_elec_get_state4
  end interface states_elec_get_state
 
  interface states_elec_set_state
    module procedure dstates_elec_set_state1, zstates_elec_set_state1, dstates_elec_set_state2, zstates_elec_set_state2
    module procedure dstates_elec_set_state3, zstates_elec_set_state3, dstates_elec_set_state4, zstates_elec_set_state4
  end interface states_elec_set_state
 
  interface states_elec_get_points
    module procedure dstates_elec_get_points1, zstates_elec_get_points1, dstates_elec_get_points2, zstates_elec_get_points2
  end interface states_elec_get_points
 
  interface states_elec_generate_random_vector
    module procedure dstates_elec_generate_random_vector, zstates_elec_generate_random_vector
  end interface
 
contains
 
  ! TODO(Alex): Issue #826. Rename to something like "states_elec_default_wfs_type", or remove
  subroutine states_elec_null(st)
    class(states_elec_t), intent(inout) :: st
 
    push_sub(states_elec_null)
 
    st%wfs_type = type_float ! By default, calculations use real wavefunctions
 
    st%packed = .false.
 
    pop_sub(states_elec_null)
  end subroutine states_elec_null
 
 
  subroutine states_elec_init(st, namespace, space, valence_charge, kpoints, calc_mode_id)
    type(states_elec_t), target, intent(inout) :: st
    type(namespace_t),           intent(in)    :: namespace
    type(electron_space_t),      intent(in)    :: space
    real(real64),                intent(in)    :: valence_charge
    type(kpoints_t),             intent(in)    :: kpoints
    integer, optional,           intent(in)    :: calc_mode_id
 
    real(real64) :: excess_charge, nempty_percent
    integer :: nempty, ntot, default
    integer :: nempty_conv, nempty_conv_default
    logical :: force, adapt_for_chebyshev
    real(real64), parameter :: tol = 1e-13_real64
    integer :: es
    integer, parameter :: rs_chebyshev = 12
    integer(int64), parameter :: chebyshev_compatible_modes(3) = &
      [option__calculationmode__gs, option__calculationmode__go, option__calculationmode__unocc]
 
    push_sub_with_profile(states_elec_init)
 
    st%fromScratch = .true. ! this will be reset if restart_read is called
    call states_elec_null(st)
 
    ! We get the spin dimension from the electronic space
    ! TODO: Remove spin space information from states_elec_dim
    st%d%ispin = space%ispin
 
    ! Use of spinors requires complex wavefunctions.
    if (st%d%ispin == spinors) call states_set_complex(st)
 
    if (st%d%ispin /= unpolarized .and. kpoints%use_time_reversal) then
      message(1) = "Time reversal symmetry is only implemented for unpolarized spins."
      message(2) = "Use KPointsUseTimeReversal = no."
      call messages_fatal(2, namespace=namespace)
    end if
 
 
    !%Variable ExcessCharge
    !%Type float
    !%Default 0.0
    !%Section States
    !%Description
    !% The net charge of the system. A negative value means that we are adding
    !% electrons, while a positive value means we are taking electrons
    !% from the system.
    !%End
    call parse_variable(namespace, 'ExcessCharge', m_zero, excess_charge)
 
    !%Variable TotalStates
    !%Type integer
    !%Default 0
    !%Section States
    !%Description
    !% This variable sets the total number of states that Octopus will
    !% use. This is normally not necessary since by default Octopus
    !% sets the number of states to the minimum necessary to hold the
    !% electrons present in the system. (This default behavior is
    !% obtained by setting <tt>TotalStates</tt> to 0).
    !%
    !% If you want to add some unoccupied states, probably it is more convenient to use the variable
    !% <tt>ExtraStates</tt>.
    !%End
    call parse_variable(namespace, 'TotalStates', 0, ntot)
    if (ntot < 0) then
      write(message(1), '(a,i5,a)') "Input: '", ntot, "' is not a valid value for TotalStates."
      call messages_fatal(1, namespace=namespace)
    end if
 
    !%Variable ExtraStates
    !%Type integer
    !%Default 0
    !%Section States
    !%Description
    !% The number of states is in principle calculated considering the minimum
    !% numbers of states necessary to hold the electrons present in the system.
    !% The number of electrons is
    !% in turn calculated considering the nature of the species supplied in the
    !% <tt>Species</tt> block, and the value of the <tt>ExcessCharge</tt> variable.
    !% However, one may command <tt>Octopus</tt> to use more states, which is necessary if one wants to
    !% use fractional occupational numbers, either fixed from the beginning through
    !% the <tt>Occupations</tt> block or by prescribing
    !% an electronic temperature with <tt>Smearing</tt>, or in order to calculate
    !% excited states (including with <tt>CalculationMode = unocc</tt>).
    !%End
    call parse_variable(namespace, 'ExtraStates', 0, nempty)
    if (nempty < 0) then
      write(message(1), '(a,i5,a)') "Input: '", nempty, "' is not a valid value for ExtraStates."
      message(2) = '(0 <= ExtraStates)'
      call messages_fatal(2, namespace=namespace)
    end if
 
    if (ntot > 0 .and. nempty > 0) then
      message(1) = 'You cannot set TotalStates and ExtraStates at the same time.'
      call messages_fatal(1, namespace=namespace)
    end if
 
    !%Variable ExtraStatesInPercent
    !%Type float
    !%Default 0
    !%Section States
    !%Description
    !% This variable allows to set the number of extra/empty states as percentage of the
    !% used occupied states. For example, a value 35 for ExtraStatesInPercent would amount
    !% to ceiling(35/100 * nstates) extra states, where nstates denotes the amount of occupied
    !% states Octopus is using for the system at hand.
    !%End
    call parse_variable(namespace, 'ExtraStatesInPercent', m_zero, nempty_percent)
    if (nempty_percent < 0) then
      write(message(1), '(a,f8.6,a)') "Input: '", nempty_percent, &
        "' should be a percentage value x (where x is parts in hundred) larger or equal 0"
      call messages_fatal(1, namespace=namespace)
    end if
 
    if (nempty > 0 .and. nempty_percent > 0) then
      message(1) = 'You cannot set ExtraStates and ExtraStatesInPercent at the same time.'
      call messages_fatal(1, namespace=namespace)
    end if
 
    ! Only use extra states for Chebyshev filtering in ground state or geometry optimization runs
    adapt_for_chebyshev = .false.
    if (present(calc_mode_id)) then
      if (any(calc_mode_id == chebyshev_compatible_modes)) then
        ! this variable is documented in electrons/eigensolver.F90
        call parse_variable(namespace, 'Eigensolver', rs_chebyshev, es)
        if (es == rs_chebyshev) then
          ! reuse this condition below
          adapt_for_chebyshev = .true.
        end if
      end if
    end if
    if (adapt_for_chebyshev) then
      if (nempty == 0 .and. is_close(nempty_percent, m_zero)) then
        nempty_percent = 15
        message(1) = 'Chebyshev filtering eigensolver detected. Setting ExtraStatesInPercent = 15'
        call messages_info(1, namespace=namespace)
      end if
    end if
 
    ! For non-periodic systems this should just return the Gamma point
    call states_elec_choose_kpoints(st, kpoints, namespace)
 
    st%val_charge = valence_charge
 
    st%qtot = -(st%val_charge + excess_charge)
 
    if (st%qtot < -m_epsilon) then
      write(message(1),'(a,f12.6,a)') 'Total charge = ', st%qtot, ' < 0'
      message(2) = 'Check Species and ExcessCharge.'
      call messages_fatal(2, only_root_writes = .true., namespace=namespace)
    end if
 
    select case (st%d%ispin)
    case (unpolarized)
      st%d%dim = 1
      st%nst = nint(st%qtot/2)
      if (st%nst*2 - st%qtot < -tol) st%nst = st%nst + 1
      st%d%nspin = 1
      st%d%spin_channels = 1
    case (spin_polarized)
      st%d%dim = 1
      st%nst = nint(st%qtot/2)
      if (st%nst*2 - st%qtot < -tol) st%nst = st%nst + 1
      st%d%nspin = 2
      st%d%spin_channels = 2
    case (spinors)
      st%d%dim = 2
      st%nst = nint(st%qtot)
      if (st%nst - st%qtot < -tol) st%nst = st%nst + 1
      st%d%nspin = 4
      st%d%spin_channels = 2
    end select
 
    if (nempty_percent > 0) then
      nempty = ceiling(nempty_percent * st%nst / 100)
    end if
 
    !%Variable ExtraStatesToConverge
    !%Type integer
    !%Default <tt>ExtraStates</tt> (Default 0)
    !%Section States
    !%Description
    !% For <tt>gs</tt> and <tt>unocc</tt> calculations.
    !% (For the <tt>gs</tt> calculation one needs to set <tt>ConvEigenError=yes</tt>)
    !% Specifies the number of extra states that will be considered for reaching the convergence.
    !% The calculation will consider the number off occupied states plus
    !% <tt>ExtraStatesToConverge</tt> for the convergence criteria.
    !% By default, all extra states need to be converged (For <tt>gs</tt> calculations only with <tt>ConvEigenError=yes</tt>).
    !% Thus, together with <tt>ExtraStates</tt>, one can have some more states which will not be
    !% considered for the convergence criteria, thus making the convergence of the
    !% unocc calculation faster.
    !%
    !% If chebyshev filtering is used, the default is <tt>min(0.8*ExtraStates, ExtraStates - 1)</tt>.
    !%End
    nempty_conv_default = nempty
    if (adapt_for_chebyshev) then
      ! use 80% of the extra states, but at least do not converge one extra state
      nempty_conv_default = min(int(0.8*nempty), nempty - 1)
    end if
    call parse_variable(namespace, 'ExtraStatesToConverge', nempty_conv_default, nempty_conv)
    if (nempty_conv < 0) then
      write(message(1), '(a,i5,a)') "Input: '", nempty_conv, "' is not a valid value for ExtraStatesToConverge."
      message(2) = '(0 <= ExtraStatesToConverge)'
      call messages_fatal(2, namespace=namespace)
    end if
 
    if (nempty_conv > nempty) then
      nempty_conv = nempty
      message(1) = 'You cannot set ExtraStatesToConverge to a higher value than ExtraStates.'
      message(2) = 'Capping ExtraStatesToConverge to ExtraStates.'
      call messages_warning(2, namespace=namespace)
    endif
 
    if (ntot > 0) then
      if (ntot < st%nst) then
        message(1) = 'TotalStates is smaller than the number of states required by the system.'
        call messages_fatal(1, namespace=namespace)
      end if
 
      st%nst = ntot
    end if
 
    st%nst_conv = st%nst + nempty_conv
    st%nst = st%nst + nempty
    if (st%nst == 0) then
      message(1) = "Cannot run with number of states = zero."
      call messages_fatal(1, namespace=namespace)
    end if
 
    !%Variable StatesBlockSize
    !%Type integer
    !%Section Execution::Optimization
    !%Description
    !% Some routines work over blocks of eigenfunctions, which
    !% generally improves performance at the expense of increased
    !% memory consumption. This variable selects the size of the
    !% blocks to be used. If GPUs are used, the default is the
    !% warp size (32 for NVIDIA, 32 or 64 for AMD);
    !% otherwise it is 4.
    !%End
 
    if (accel_is_enabled()) then
      ! use the warp size (usually 32, for some AMD GPUs it is 64)
      default = accel%warp_size
    else
      default = 4
    end if
 
    if (default > pad_pow2(st%nst)) default = pad_pow2(st%nst)
 
    assert(default > 0)
 
    call parse_variable(namespace, 'StatesBlockSize', default, st%block_size)
    if (st%block_size < 1) then
      call messages_write("The variable 'StatesBlockSize' must be greater than 0.")
      call messages_fatal(namespace=namespace)
    end if
 
    st%block_size = min(st%block_size, st%nst)
    conf%target_states_block_size = st%block_size
 
    safe_allocate(st%eigenval(1:st%nst, 1:st%nik))
    st%eigenval = huge(st%eigenval)
 
    ! Periodic systems require complex wavefunctions
    ! but not if it is Gamma-point only
    if (.not. kpoints%gamma_only()) then
      call states_set_complex(st)
    end if
 
    !%Variable OnlyUserDefinedInitialStates
    !%Type logical
    !%Default no
    !%Section States
    !%Description
    !% If true, then only user-defined states from the block <tt>UserDefinedStates</tt>
    !% will be used as initial states for a time-propagation. No attempt is made
    !% to load ground-state orbitals from a previous ground-state run.
    !%End
    call parse_variable(namespace, 'OnlyUserDefinedInitialStates', .false., st%only_userdef_istates)
 
    ! we now allocate some arrays
    safe_allocate(st%occ     (1:st%nst, 1:st%nik))
    st%occ      = m_zero
    ! allocate space for formula strings that define user-defined states
    if (parse_is_defined(namespace, 'UserDefinedStates') .or. parse_is_defined(namespace, 'OCTInitialUserdefined') &
      .or. parse_is_defined(namespace, 'OCTTargetUserdefined')) then
      safe_allocate(st%user_def_states(1:st%d%dim, 1:st%nst, 1:st%nik))
      ! initially we mark all 'formulas' as undefined
      st%user_def_states(1:st%d%dim, 1:st%nst, 1:st%nik) = 'undefined'
    end if
 
    if (st%d%ispin == spinors) then
      safe_allocate(st%spin(1:3, 1:st%nst, 1:st%nik))
    end if
 
    !%Variable StatesRandomization
    !%Type integer
    !%Default par_independent
    !%Section States
    !%Description
    !% The randomization of states can be done in two ways:
    !% i) a parallelisation independent way (default), where the random states are identical,
    !% irrespectively of the number of tasks and
    !% ii) a parallelisation dependent way, which can prevent linear dependency
    !%  to occur for large systems.
    !%Option par_independent 1
    !% Parallelisation-independent randomization of states.
    !%Option par_dependent 2
    !% The randomization depends on the number of taks used in the calculation.
    !%End
    call parse_variable(namespace, 'StatesRandomization', par_independent, st%randomization)
 
 
    call states_elec_read_initial_occs(st, namespace, excess_charge, kpoints)
    call states_elec_read_initial_spins(st, namespace)
 
    ! This test can only be done here, as smear_init is called inside states_elec_read_initial_occs, and
    ! only there smear%photodop is set.
 
    if (st%smear%photodop) then
      if (nempty == 0) then
        write(message(1), '(a,i5,a)') "PhotoDoping requires to specify ExtraStates."
        message(2) = '(0 == ExtraStates)'
        call messages_fatal(2, namespace=namespace)
      end if
    end if
 
    st%st_start = 1
    st%st_end = st%nst
    st%lnst = st%nst
    safe_allocate(st%node(1:st%nst))
    st%node(1:st%nst) = 0
 
    call mpi_grp_init(st%mpi_grp, mpi_comm_undefined)
    st%parallel_in_states = .false.
 
    call distributed_nullify(st%d%kpt, st%nik)
 
    call modelmb_particles_init(st%modelmbparticles, namespace, space, st%nst)
 
    !%Variable SymmetrizeDensity
    !%Type logical
    !%Default no
    !%Section States
    !%Description
    !% When enabled the density is symmetrized. Currently, this can
    !% only be done for periodic systems. (Experimental.)
    !%End
    call parse_variable(namespace, 'SymmetrizeDensity', kpoints%use_symmetries, st%symmetrize_density)
    call messages_print_var_value('SymmetrizeDensity', st%symmetrize_density, namespace=namespace)
 
    !%Variable ForceComplex
    !%Type logical
    !%Default no
    !%Section Execution::Debug
    !%Description
    !% Normally <tt>Octopus</tt> determines automatically the type necessary
    !% for the wavefunctions. When set to yes this variable will
    !% force the use of complex wavefunctions.
    !%
    !% Warning: This variable is designed for testing and
    !% benchmarking and normal users need not use it.
    !%End
    call parse_variable(namespace, 'ForceComplex', .false., force)
 
    if (force) call states_set_complex(st)
 
    st%packed = .false.
 
    pop_sub_with_profile(states_elec_init)
  end subroutine states_elec_init
 
  ! ---------------------------------------------------------
  !
  subroutine states_elec_look(restart, nik, dim, nst, ierr)
    class(restart_t), intent(in)  :: restart
    integer,          intent(out) :: nik
    integer,          intent(out) :: dim
    integer,          intent(out) :: nst
    integer,          intent(out) :: ierr
 
    character(len=256) :: lines(3)
    character(len=20)   :: char
    integer :: iunit
 
    push_sub(states_elec_look)
 
    ierr = 0
 
    iunit = restart%open('states')
    call restart%read(iunit, lines, 3, ierr)
    if (ierr == 0) then
      read(lines(1), *) char, nst
      read(lines(2), *) char, dim
      read(lines(3), *) char, nik
    end if
    call restart%close(iunit)
 
    pop_sub(states_elec_look)
  end subroutine states_elec_look
 
  ! ---------------------------------------------------------
  !
  subroutine states_elec_read_initial_occs(st, namespace, excess_charge, kpoints)
    type(states_elec_t),  intent(inout) :: st
    type(namespace_t),    intent(in)    :: namespace
    real(real64),         intent(in)    :: excess_charge
    type(kpoints_t),      intent(in)    :: kpoints
 
    integer :: ik, ist, ispin, nspin, ncols, nrows, el_per_state, icol, start_pos, spin_n
    type(block_t) :: blk
    real(real64) :: rr, charge
    logical :: integral_occs, unoccupied_states
    real(real64), allocatable :: read_occs(:, :)
    real(real64) :: charge_in_block
 
    push_sub(states_elec_read_initial_occs)
 
    !%Variable RestartFixedOccupations
    !%Type logical
    !%Section States
    !%Description
    !% Setting this variable will make the restart proceed as
    !% if the occupations from the previous calculation had been set via the <tt>Occupations</tt> block,
    !% <i>i.e.</i> fixed. Otherwise, occupations will be determined by smearing.
    !%
    !% Linear response calculations and self-consistent calculations are unaffected by this variable.
    !% This is however used for Sternheimer calculations.
    !%End
    call parse_variable(namespace, 'RestartFixedOccupations', .true., st%restart_fixed_occ)
    ! we will turn on st%fixed_occ if restart_read is ever called
 
    !%Variable Occupations
    !%Type block
    !%Section States
    !%Description
    !% The occupation numbers of the orbitals can be fixed through the use of this
    !% variable. For example:
    !%
    !% <tt>%Occupations
    !% <br>&nbsp;&nbsp;2 | 2 | 2 | 2 | 2
    !% <br>%</tt>
    !%
    !% would fix the occupations of the five states to 2. There can be
    !% at most as many columns as states in the calculation. If there are fewer columns
    !% than states, then the code will assume that the user is indicating the occupations
    !% of the uppermost states where all lower states have full occupation (i.e. 2 for spin-unpolarized
    !% calculations, 1 otherwise) and all higher states have zero occupation. The first column
    !% will be taken to refer to the lowest state such that the occupations would be consistent
    !% with the correct total charge. For example, if there are 8 electrons and 10 states (from
    !% <tt>ExtraStates = 6</tt>), then an abbreviated specification
    !%
    !% <tt>%Occupations
    !% <br>&nbsp;&nbsp;1 | 0 | 1
    !% <br>%</tt>
    !%
    !% would be equivalent to a full specification
    !%
    !% <tt>%Occupations
    !% <br>&nbsp;&nbsp;2 | 2 | 2 | 1 | 0 | 1 | 0 | 0 | 0 | 0
    !% <br>%</tt>
    !%
    !% This is an example of use for constrained density-functional theory,
    !% crudely emulating a HOMO->LUMO+1 optical excitation.
    !% The number of rows should be equal
    !% to the number of k-points times the number of spins. For example, for a finite system
    !% with <tt>SpinComponents == spin_polarized</tt>,
    !% this block should contain two lines, one for each spin channel.
    !% All rows must have the same number of columns.
    !%
    !% The <tt>Occupations</tt> block is useful for the ground state of highly symmetric
    !% small systems (like an open-shell atom), to fix the occupation numbers
    !% of degenerate states in order to help <tt>octopus</tt> to converge. This is to
    !% be used in conjuction with <tt>ExtraStates</tt>. For example, to calculate the
    !% carbon atom, one would do:
    !%
    !% <tt>ExtraStates = 2
    !% <br>%Occupations
    !% <br>&nbsp;&nbsp;2 | 2/3 | 2/3 | 2/3
    !% <br>%</tt>
    !%
    !% If you want the calculation to be spin-polarized (which makes more sense), you could do:
    !%
    !% <tt>ExtraStates = 2
    !% <br>%Occupations
    !% <br>&nbsp;&nbsp; 2/3 | 2/3 | 2/3
    !% <br>&nbsp;&nbsp; 0   |   0 |   0
    !% <br>%</tt>
    !%
    !% Note that in this case the first state is absent, the code will calculate four states
    !% (two because there are four electrons, plus two because <tt>ExtraStates</tt> = 2), and since
    !% it finds only three columns, it will occupy the first state with one electron for each
    !% of the spin options.
    !%
    !% If the sum of occupations is not equal to the total charge set by <tt>ExcessCharge</tt>,
    !% an error message is printed.
    !% If <tt>FromScratch = no</tt> and <tt>RestartFixedOccupations = yes</tt>,
    !% this block will be ignored.
    !%End
 
    integral_occs = .true.
 
    occ_fix: if (parse_block(namespace, 'Occupations', blk) == 0) then
      ! read in occupations
      st%fixed_occ = .true.
 
      ncols = parse_block_cols(blk, 0)
      if (ncols > st%nst) then
        call messages_input_error(namespace, "Occupations", "Too many columns in block Occupations.")
      end if
 
      nrows = parse_block_n(blk)
      if (nrows /= st%nik) then
        call messages_input_error(namespace, "Occupations", "Wrong number of rows in block Occupations.")
      end if
 
      do ik = 1, st%nik - 1
        if (parse_block_cols(blk, ik) /= ncols) then
          call messages_input_error(namespace, "Occupations", &
            "All rows in block Occupations must have the same number of columns.")
        end if
      end do
 
      ! Now we fill all the "missing" states with the maximum occupation.
      if (st%d%ispin == unpolarized) then
        el_per_state = 2
      else
        el_per_state = 1
      end if
 
      safe_allocate(read_occs(1:ncols, 1:st%nik))
 
      charge_in_block = m_zero
      do ik = 1, st%nik
        do icol = 1, ncols
          call parse_block_float(blk, ik - 1, icol - 1, read_occs(icol, ik))
          charge_in_block = charge_in_block + read_occs(icol, ik) * st%kweights(ik)
        end do
      end do
 
      spin_n = 2
      select case (st%d%ispin)
      case (unpolarized)
        spin_n = 2
      case (spin_polarized)
        spin_n = 2
      case (spinors)
        spin_n = 1
      end select
 
      start_pos = nint((st%qtot - charge_in_block)/spin_n)
 
      if (start_pos + ncols > st%nst) then
        message(1) = "To balance charge, the first column in block Occupations is taken to refer to state"
        write(message(2),'(a,i6,a)') "number ", start_pos, " but there are too many columns for the number of states."
        write(message(3),'(a,i6,a)') "Solution: set ExtraStates = ", start_pos + ncols - st%nst
        call messages_fatal(3, namespace=namespace)
      end if
 
      do ik = 1, st%nik
        do ist = 1, start_pos
          st%occ(ist, ik) = el_per_state
        end do
      end do
 
      do ik = 1, st%nik
        do ist = start_pos + 1, start_pos + ncols
          st%occ(ist, ik) = read_occs(ist - start_pos, ik)
          integral_occs = integral_occs .and. &
            abs((st%occ(ist, ik) - el_per_state) * st%occ(ist, ik))  <=  m_epsilon
        end do
      end do
 
      do ik = 1, st%nik
        do ist = start_pos + ncols + 1, st%nst
          st%occ(ist, ik) = m_zero
        end do
      end do
 
      call parse_block_end(blk)
 
      safe_deallocate_a(read_occs)
 
    else
      st%fixed_occ = .false.
      integral_occs = .false.
 
      ! first guess for occupation...paramagnetic configuration
      rr = m_one
      if (st%d%ispin == unpolarized) rr = m_two
 
      st%occ  = m_zero
      st%qtot = -(st%val_charge + excess_charge)
 
      nspin = 1
      if (st%d%nspin == 2) nspin = 2
 
      do ik = 1, st%nik, nspin
        charge = m_zero
        do ist = 1, st%nst
          do ispin = ik, ik + nspin - 1
            st%occ(ist, ispin) = min(rr, -(st%val_charge + excess_charge) - charge)
            charge = charge + st%occ(ist, ispin)
          end do
        end do
      end do
 
    end if occ_fix
 
    !%Variable RestartReorderOccs
    !%Type logical
    !%Default no
    !%Section States
    !%Description
    !% Consider doing a ground-state calculation, and then restarting with new occupations set
    !% with the <tt>Occupations</tt> block, in an attempt to populate the orbitals of the original
    !% calculation. However, the eigenvalues may reorder as the density changes, in which case the
    !% occupations will now be referring to different orbitals. Setting this variable to yes will
    !% try to solve this issue when the restart data is being read, by reordering the occupations
    !% according to the order of the expectation values of the restart wavefunctions.
    !%End
    if (st%fixed_occ) then
      call parse_variable(namespace, 'RestartReorderOccs', .false., st%restart_reorder_occs)
    else
      st%restart_reorder_occs = .false.
    end if
 
    call smear_init(st%smear, namespace, st%d%ispin, st%fixed_occ, integral_occs, kpoints)
 
    unoccupied_states = (st%d%ispin /= spinors .and. st%nst*2 > st%qtot) .or. (st%d%ispin == spinors .and. st%nst > st%qtot)
 
    if (.not. smear_is_semiconducting(st%smear) .and. .not. st%smear%method == smear_fixed_occ) then
      if (.not. unoccupied_states) then
        call messages_write('Smearing needs unoccupied states (via ExtraStates or TotalStates) to be useful.')
        call messages_warning(namespace=namespace)
      end if
    end if
 
    ! sanity check
    charge = m_zero
    do ist = 1, st%nst
      charge = charge + sum(st%occ(ist, 1:st%nik) * st%kweights(1:st%nik))
    end do
    if (abs(charge - st%qtot) > 1e-6_real64) then
      message(1) = "Initial occupations do not integrate to total charge."
      write(message(2), '(6x,f12.6,a,f12.6)') charge, ' != ', st%qtot
      call messages_fatal(2, only_root_writes = .true., namespace=namespace)
    end if
 
    st%uniform_occ = smear_is_semiconducting(st%smear) .and. .not. unoccupied_states
 
    pop_sub(states_elec_read_initial_occs)
  end subroutine states_elec_read_initial_occs
 
 
  ! ---------------------------------------------------------
  !
  subroutine states_elec_read_initial_spins(st, namespace)
    type(states_elec_t), intent(inout) :: st
    type(namespace_t),   intent(in)    :: namespace
 
    integer :: i, j, nrows
    type(block_t) :: blk
 
    push_sub(states_elec_read_initial_spins)
 
    st%fixed_spins = .false.
    if (st%d%ispin /= spinors) then
      pop_sub(states_elec_read_initial_spins)
      return
    end if
 
    !%Variable InitialSpins
    !%Type block
    !%Section States
    !%Description
    !% The spin character of the initial random guesses for the spinors can
    !% be fixed by making use of this block. Note that this will not "fix" the
    !% the spins during the calculation (this cannot be done in spinors mode, in
    !% being able to change the spins is why the spinors mode exists in the first
    !% place).
    !%
    !% This block is meaningless and ignored if the run is not in spinors mode
    !% (<tt>SpinComponents = spinors</tt>).
    !%
    !% The structure of the block is very simple: each column contains the desired
    !% <math>\left< S_x \right>, \left< S_y \right>, \left< S_z \right> </math> for each spinor.
    !% If the calculation is for a periodic system
    !% and there is more than one <i>k</i>-point, the spins of all the <i>k</i>-points are
    !% the same.
    !%
    !% For example, if we have two spinors, and we want one in the <math>S_x</math> "down" state,
    !% and another one in the <math>S_x</math> "up" state:
    !%
    !% <tt>%InitialSpins
    !% <br>&nbsp;&nbsp;&nbsp; 0.5 | 0.0 | 0.0
    !% <br>&nbsp;&nbsp; -0.5 | 0.0 | 0.0
    !% <br>%</tt>
    !%
    !% WARNING: if the calculation is for a system described by pseudopotentials (as
    !% opposed to user-defined potentials or model systems), this option is
    !% meaningless since the random spinors are overwritten by the atomic orbitals.
    !%
    !% This constraint must be fulfilled:
    !% <br><math> \left< S_x \right>^2 + \left< S_y \right>^2 + \left< S_z \right>^2 = \frac{1}{4} </math>
    !%End
    spin_fix: if (parse_block(namespace, 'InitialSpins', blk) == 0) then
      nrows = parse_block_n(blk)
      if (nrows < st%nst) then
        message(1) = "Please specify one row for each state in InitialSpins, also for extra states."
        call messages_fatal(1, namespace=namespace)
      end if
      do i = 1, st%nst
        do j = 1, 3
          call parse_block_float(blk, i-1, j-1, st%spin(j, i, 1))
        end do
        if (abs(sum(st%spin(1:3, i, 1)**2) - m_fourth) > 1.0e-6_real64) call messages_input_error(namespace, 'InitialSpins')
      end do
      call parse_block_end(blk)
      st%fixed_spins = .true.
      do i = 2, st%nik
        st%spin(:, :, i) = st%spin(:, :, 1)
      end do
    end if spin_fix
 
    pop_sub(states_elec_read_initial_spins)
  end subroutine states_elec_read_initial_spins
 
 
  ! ---------------------------------------------------------
  !
  subroutine states_elec_allocate_wfns(st, mesh, wfs_type, skip, packed)
    type(states_elec_t),    intent(inout)   :: st
    class(mesh_t),          intent(in)      :: mesh
    type(type_t), optional, intent(in)      :: wfs_type
    logical,      optional, intent(in)      :: skip(:)
    logical,      optional, intent(in)      :: packed
 
    push_sub(states_elec_allocate_wfns)
 
    if (present(wfs_type)) then
      assert(wfs_type == type_float .or. wfs_type == type_cmplx)
      st%wfs_type = wfs_type
    end if
 
    call states_elec_init_block(st, mesh, skip = skip, packed=packed)
    call states_elec_set_zero(st)
 
    pop_sub(states_elec_allocate_wfns)
  end subroutine states_elec_allocate_wfns
 
  !---------------------------------------------------------------------
  subroutine states_elec_init_block(st, mesh, verbose, skip, packed)
    type(states_elec_t),           intent(inout) :: st
    type(mesh_t),                  intent(in)    :: mesh
    logical, optional,             intent(in)    :: verbose
    logical, optional,             intent(in)    :: skip(:)
    logical, optional,             intent(in)    :: packed
 
    integer :: ib, iqn, ist, istmin, istmax
    logical :: same_node, verbose_, packed_
    integer, allocatable :: bstart(:), bend(:)
 
    push_sub(states_elec_init_block)
 
    safe_allocate(bstart(1:st%nst))
    safe_allocate(bend(1:st%nst))
    safe_allocate(st%group%iblock(1:st%nst))
 
    st%group%iblock = 0
 
    verbose_ = optional_default(verbose, .true.)
    packed_ = optional_default(packed, .false.)
 
    !In case we have a list of state to skip, we do not allocate them
    istmin = 1
    if (present(skip)) then
      do ist = 1, st%nst
        if (.not. skip(ist)) then
          istmin = ist
          exit
        end if
      end do
    end if
 
    istmax = st%nst
    if (present(skip)) then
      do ist = st%nst, istmin, -1
        if (.not. skip(ist)) then
          istmax = ist
          exit
        end if
      end do
    end if
 
    if (present(skip) .and. verbose_) then
      call messages_write('Info: Allocating states from ')
      call messages_write(istmin, fmt = 'i8')
      call messages_write(' to ')
      call messages_write(istmax, fmt = 'i8')
      call messages_info()
    end if
 
    ! count and assign blocks
    ib = 0
    st%group%nblocks = 0
    bstart(1) = istmin
    do ist = istmin, istmax
      ib = ib + 1
 
      st%group%iblock(ist) = st%group%nblocks + 1
 
      same_node = .true.
      if (st%parallel_in_states .and. ist /= istmax) then
        ! We have to avoid that states that are in different nodes end
        ! up in the same block
        same_node = (st%node(ist + 1) == st%node(ist))
      end if
 
      if (ib == st%block_size .or. ist == istmax .or. .not. same_node) then
        ib = 0
        st%group%nblocks = st%group%nblocks + 1
        bend(st%group%nblocks) = ist
        if (ist /= istmax) bstart(st%group%nblocks + 1) = ist + 1
      end if
    end do
 
    safe_allocate(st%group%psib(1:st%group%nblocks, st%d%kpt%start:st%d%kpt%end))
 
    safe_allocate(st%group%block_is_local(1:st%group%nblocks, 1:st%nik))
    st%group%block_is_local = .false.
    st%group%block_start  = -1
    st%group%block_end    = -2  ! this will make that loops block_start:block_end do not run if not initialized
 
    do ib = 1, st%group%nblocks
      if (bstart(ib) >= st%st_start .and. bend(ib) <= st%st_end) then
        if (st%group%block_start == -1) st%group%block_start = ib
        st%group%block_end = ib
        do iqn = st%d%kpt%start, st%d%kpt%end
          st%group%block_is_local(ib, iqn) = .true.
 
          if (states_are_real(st)) then
            call dwfs_elec_init(st%group%psib(ib, iqn), st%d%dim, bstart(ib), bend(ib), mesh%np_part, iqn, &
              special=.true., packed=packed_)
          else
            call zwfs_elec_init(st%group%psib(ib, iqn), st%d%dim, bstart(ib), bend(ib), mesh%np_part, iqn, &
              special=.true., packed=packed_)
          end if
 
        end do
      end if
    end do
 
    safe_allocate(st%group%block_range(1:st%group%nblocks, 1:2))
    safe_allocate(st%group%block_size(1:st%group%nblocks))
 
    st%group%block_range(1:st%group%nblocks, 1) = bstart(1:st%group%nblocks)
    st%group%block_range(1:st%group%nblocks, 2) = bend(1:st%group%nblocks)
    st%group%block_size(1:st%group%nblocks) = bend(1:st%group%nblocks) - bstart(1:st%group%nblocks) + 1
 
    st%group%block_initialized = .true.
 
    safe_allocate(st%group%block_node(1:st%group%nblocks, 1:st%nik))
    st%group%block_node = 0
 
    assert(allocated(st%node))
    assert(all(st%node >= 0) .and. all(st%node < st%mpi_grp%size))
 
    do iqn = st%d%kpt%start, st%d%kpt%end
      do ib = st%group%block_start, st%group%block_end
        st%group%block_node(ib, iqn) = st%st_kpt_mpi_grp%rank
      end do
    end do
 
    call comm_allreduce(st%st_kpt_mpi_grp, st%group%block_node)
 
    if (verbose_) then
      call messages_write('Info: Blocks of states')
      call messages_info()
      do ib = 1, st%group%nblocks
        call messages_write('      Block ')
        call messages_write(ib, fmt = 'i8')
        call messages_write(' contains ')
        call messages_write(st%group%block_size(ib), fmt = 'i8')
        call messages_write(' states')
        if (st%group%block_size(ib) > 0) then
          call messages_write(':')
          call messages_write(st%group%block_range(ib, 1), fmt = 'i8')
          call messages_write(' - ')
          call messages_write(st%group%block_range(ib, 2), fmt = 'i8')
        end if
        call messages_info()
      end do
    end if
 
!!$!!!!DEBUG
!!$    ! some debug output that I will keep here for the moment
!!$    if (st%system_grp%is_root()) then
!!$      print*, "NST       ", st%nst
!!$      print*, "BLOCKSIZE ", st%block_size
!!$      print*, "NBLOCKS   ", st%group%nblocks
!!$
!!$      print*, "==============="
!!$      do ist = 1, st%nst
!!$        print*, st%node(ist), ist, st%group%iblock(ist)
!!$      end do
!!$      print*, "==============="
!!$
!!$      do ib = 1, st%group%nblocks
!!$        print*, ib, bstart(ib), bend(ib)
!!$      end do
!!$
!!$    end if
!!$!!!!ENDOFDEBUG
 
    safe_deallocate_a(bstart)
    safe_deallocate_a(bend)
    pop_sub(states_elec_init_block)
  end subroutine states_elec_init_block
 
 
  ! ---------------------------------------------------------
  subroutine states_elec_deallocate_wfns(st)
    type(states_elec_t), intent(inout) :: st
 
    push_sub(states_elec_deallocate_wfns)
 
    call states_elec_group_end(st%group, st%d)
    ! The wave functions (and their device buffers) are gone, so the batches
    ! are no longer packed. Reset the flag so a subsequent st%pack() actually
    ! repacks instead of returning early on a stale value.
    st%packed = .false.
 
    pop_sub(states_elec_deallocate_wfns)
  end subroutine states_elec_deallocate_wfns
 
 
  ! ---------------------------------------------------------
  subroutine states_elec_densities_init(st, gr)
    type(states_elec_t), target, intent(inout) :: st
    type(grid_t),                intent(in)    :: gr
 
    real(real64) :: fsize
 
    push_sub(states_elec_densities_init)
 
    safe_allocate(st%rho(1:gr%np_part, 1:st%d%nspin))
    st%rho = m_zero
 
    fsize = gr%np_part*8.0_real64*st%block_size
 
    call messages_write('Info: states-block size = ')
    call messages_write(fsize, fmt = '(f10.1)', align_left = .true., units = unit_megabytes, print_units = .true.)
    call messages_info()
 
    pop_sub(states_elec_densities_init)
  end subroutine states_elec_densities_init
 
  !---------------------------------------------------------------------
  subroutine states_elec_allocate_current(st, space, mesh)
    type(states_elec_t), intent(inout) :: st
    class(space_t),      intent(in)    :: space
    class(mesh_t),       intent(in)    :: mesh
 
    push_sub(states_elec_allocate_current)
 
    if (.not. allocated(st%current)) then
      safe_allocate(st%current(1:mesh%np_part, 1:space%dim, 1:st%d%nspin))
      st%current = m_zero
    end if
 
    if (.not. allocated(st%current_para)) then
      safe_allocate(st%current_para(1:mesh%np_part, 1:space%dim, 1:st%d%nspin))
      st%current_para = m_zero
    end if
 
    if (.not. allocated(st%current_dia)) then
      safe_allocate(st%current_dia(1:mesh%np_part, 1:space%dim, 1:st%d%nspin))
      st%current_dia= m_zero
    end if
 
    if (.not. allocated(st%current_mag)) then
      safe_allocate(st%current_mag(1:mesh%np_part, 1:space%dim, 1:st%d%nspin))
      st%current_mag= m_zero
    end if
 
    if (.not. allocated(st%current_kpt)) then
      safe_allocate(st%current_kpt(1:mesh%np,1:space%dim,st%d%kpt%start:st%d%kpt%end))
      st%current_kpt = m_zero
    end if
 
    pop_sub(states_elec_allocate_current)
  end subroutine states_elec_allocate_current
 
  !---------------------------------------------------------------------
  subroutine states_elec_exec_init(st, namespace, mc)
    type(states_elec_t),  intent(inout) :: st
    type(namespace_t),    intent(in)    :: namespace
    type(multicomm_t),    intent(in)    :: mc
 
    integer :: default
 
    push_sub(states_elec_exec_init)
 
    !%Variable StatesPack
    !%Type logical
    !%Section Execution::Optimization
    !%Description
    !% When set to yes, states are stored in packed mode, which improves
    !% performance considerably. Not all parts of the code will profit from
    !% this, but should nevertheless work regardless of how the states are
    !% stored.
    !%
    !% If GPUs are used and this variable is set to yes, Octopus
    !% will store the wave-functions in device (GPU) memory. If
    !% there is not enough memory to store all the wave-functions,
    !% execution will stop with an error.
    !%
    !% See also the related <tt>HamiltonianApplyPacked</tt> variable.
    !%
    !% The default is yes.
    !%End
 
    call parse_variable(namespace, 'StatesPack', .true., st%pack_states)
 
    call messages_print_var_value('StatesPack', st%pack_states, namespace=namespace)
 
    call messages_obsolete_variable(namespace, 'StatesMirror')
 
    !%Variable StatesOrthogonalization
    !%Type integer
    !%Section SCF::Eigensolver
    !%Description
    !% The full orthogonalization method used by some
    !% eigensolvers. The default is <tt>cholesky_serial</tt>, except with state
    !% parallelization, the default is <tt>cholesky_parallel</tt>.
    !%Option cholesky_serial 1
    !% Cholesky decomposition implemented using
    !% BLAS/LAPACK. Can be used with domain parallelization but not
    !% state parallelization.
    !%Option cholesky_parallel 2
    !% Cholesky decomposition implemented using
    !% ScaLAPACK. Compatible with states parallelization.
    !%Option cgs 3
    !% Classical Gram-Schmidt (CGS) orthogonalization.
    !% Can be used with domain parallelization but not state parallelization.
    !% The algorithm is defined in Giraud et al., Computers and Mathematics with Applications 50, 1069 (2005).
    !%Option mgs 4
    !% Modified Gram-Schmidt (MGS) orthogonalization.
    !% Can be used with domain parallelization but not state parallelization.
    !% The algorithm is defined in Giraud et al., Computers and Mathematics with Applications 50, 1069 (2005).
    !%Option drcgs 5
    !% Classical Gram-Schmidt orthogonalization with double-step reorthogonalization.
    !% Can be used with domain parallelization but not state parallelization.
    !% The algorithm is taken from Giraud et al., Computers and Mathematics with Applications 50, 1069 (2005).
    !% According to this reference, this is much more precise than CGS or MGS algorithms.
    !% The MGS version seems not to improve much the stability and would require more communications over the domains.
    !%End
 
    default = option__statesorthogonalization__cholesky_serial
#ifdef HAVE_SCALAPACK
    if (multicomm_strategy_is_parallel(mc, p_strategy_states)) then
      default = option__statesorthogonalization__cholesky_parallel
    end if
#endif
 
    call parse_variable(namespace, 'StatesOrthogonalization', default, st%orth_method)
 
    if (.not. varinfo_valid_option('StatesOrthogonalization', st%orth_method)) then
      call messages_input_error(namespace, 'StatesOrthogonalization')
    end if
    call messages_print_var_option('StatesOrthogonalization', st%orth_method, namespace=namespace)
 
    !%Variable StatesDeviceMemory
    !%Type float
    !%Section Execution::Optimization
    !%Default -512
    !%Description
    !% This variable selects the amount of device memory that
    !% will be used by Octopus to store the states.
    !%
    !% A positive number smaller than 1 indicates a fraction of the total
    !% device memory. A number larger than one indicates an absolute
    !% amount of memory in megabytes. A negative number indicates an
    !% amount of memory in megabytes that would be subtracted from
    !% the total device memory.
    !%End
    call parse_variable(namespace, 'StatesDeviceMemory', -512.0_real64, st%gpu_states_mem)
    call messages_obsolete_variable(namespace, 'StatesCLDeviceMemory', 'StatesDeviceMemory')
 
    pop_sub(states_elec_exec_init)
  end subroutine states_elec_exec_init
 
 
  ! ---------------------------------------------------------
  !
  subroutine states_elec_copy(stout, stin, exclude_wfns, exclude_eigenval, special)
    type(states_elec_t), target, intent(inout) :: stout
    type(states_elec_t),         intent(in)    :: stin
    logical, optional,           intent(in)    :: exclude_wfns
    logical, optional,           intent(in)    :: exclude_eigenval
    logical, optional,           intent(in)    :: special
 
    logical :: exclude_wfns_
 
    push_sub(states_elec_copy)
 
    exclude_wfns_ = optional_default(exclude_wfns, .false.)
 
    call states_elec_null(stout)
 
    call states_elec_dim_copy(stout%d, stin%d)
    safe_allocate_source_a(stout%kweights, stin%kweights)
    stout%nik = stin%nik
 
    call modelmb_particles_copy(stout%modelmbparticles, stin%modelmbparticles)
 
    stout%wfs_type    = stin%wfs_type
    stout%nst         = stin%nst
    stout%block_size  = stin%block_size
    stout%orth_method = stin%orth_method
 
    stout%gpu_states_mem  = stin%gpu_states_mem
    stout%pack_states    = stin%pack_states
 
    stout%only_userdef_istates = stin%only_userdef_istates
 
    if (.not. exclude_wfns_) then
      safe_allocate_source_a(stout%rho, stin%rho)
    end if
 
    stout%uniform_occ = stin%uniform_occ
 
    if (.not. optional_default(exclude_eigenval, .false.)) then
      safe_allocate_source_a(stout%eigenval, stin%eigenval)
      safe_allocate_source_a(stout%occ, stin%occ)
      safe_allocate_source_a(stout%spin, stin%spin)
    end if
 
    ! the call to init_block is done at the end of this subroutine
    ! it allocates iblock, psib, block_is_local
    stout%group%nblocks = stin%group%nblocks
 
    safe_allocate_source_a(stout%user_def_states, stin%user_def_states)
 
    safe_allocate_source_a(stout%current, stin%current)
    safe_allocate_source_a(stout%current_kpt, stin%current_kpt)
    safe_allocate_source_a(stout%rho_core, stin%rho_core)
    safe_allocate_source_a(stout%frozen_rho, stin%frozen_rho)
    safe_allocate_source_a(stout%frozen_tau, stin%frozen_tau)
    safe_allocate_source_a(stout%frozen_gdens, stin%frozen_gdens)
    safe_allocate_source_a(stout%frozen_ldens, stin%frozen_ldens)
 
    stout%fixed_occ = stin%fixed_occ
    stout%restart_fixed_occ = stin%restart_fixed_occ
 
    stout%fixed_spins = stin%fixed_spins
 
    stout%qtot       = stin%qtot
    stout%val_charge = stin%val_charge
 
    call smear_copy(stout%smear, stin%smear)
 
    stout%parallel_in_states = stin%parallel_in_states
    call mpi_grp_copy(stout%mpi_grp, stin%mpi_grp)
    call mpi_grp_copy(stout%system_grp, stin%system_grp)
    call mpi_grp_copy(stout%dom_st_kpt_mpi_grp, stin%dom_st_kpt_mpi_grp)
    call mpi_grp_copy(stout%st_kpt_mpi_grp, stin%st_kpt_mpi_grp)
    call mpi_grp_copy(stout%dom_st_mpi_grp, stin%dom_st_mpi_grp)
    safe_allocate_source_a(stout%node, stin%node)
    safe_allocate_source_a(stout%st_kpt_task, stin%st_kpt_task)
 
#ifdef HAVE_SCALAPACK
    call blacs_proc_grid_copy(stin%dom_st_proc_grid, stout%dom_st_proc_grid)
#endif
 
    call distributed_copy(stin%dist, stout%dist)
 
    stout%scalapack_compatible = stin%scalapack_compatible
 
    stout%lnst       = stin%lnst
    stout%st_start   = stin%st_start
    stout%st_end     = stin%st_end
 
    if (stin%parallel_in_states) call multicomm_all_pairs_copy(stout%ap, stin%ap)
 
    stout%symmetrize_density = stin%symmetrize_density
 
    if (.not. exclude_wfns_) call states_elec_group_copy(stin%d, stin%group, stout%group, special=special)
 
    stout%packed = stin%packed
 
    stout%randomization = stin%randomization
 
    pop_sub(states_elec_copy)
  end subroutine states_elec_copy
 
 
  ! ---------------------------------------------------------
  !
  subroutine states_elec_end(st)
    type(states_elec_t), intent(inout) :: st
 
    push_sub(states_elec_end)
 
    call states_elec_dim_end(st%d)
 
    call modelmb_particles_end(st%modelmbparticles)
 
    ! this deallocates dpsi, zpsi, psib, iblock, iblock
    call states_elec_deallocate_wfns(st)
 
    safe_deallocate_a(st%user_def_states)
 
    safe_deallocate_a(st%rho)
    safe_deallocate_a(st%eigenval)
 
    safe_deallocate_a(st%current)
    safe_deallocate_a(st%current_para)
    safe_deallocate_a(st%current_dia)
    safe_deallocate_a(st%current_mag)
    safe_deallocate_a(st%current_kpt)
    safe_deallocate_a(st%rho_core)
    safe_deallocate_a(st%frozen_rho)
    safe_deallocate_a(st%frozen_tau)
    safe_deallocate_a(st%frozen_gdens)
    safe_deallocate_a(st%frozen_ldens)
    safe_deallocate_a(st%occ)
    safe_deallocate_a(st%spin)
    safe_deallocate_a(st%kweights)
 
 
#ifdef HAVE_SCALAPACK
    call blacs_proc_grid_end(st%dom_st_proc_grid)
#endif
 
    call distributed_end(st%dist)
 
    safe_deallocate_a(st%node)
    safe_deallocate_a(st%st_kpt_task)
 
    if (st%parallel_in_states) then
      safe_deallocate_a(st%ap%schedule)
    end if
 
    if (st%buff_density%allocated) call accel_free_buffer(st%buff_density)
 
    pop_sub(states_elec_end)
  end subroutine states_elec_end
 
 
  ! TODO(Alex) Issue #684. Abstract duplicate code in states_elec_generate_random, to get it to
  ! a point where it can be refactored.
  subroutine states_elec_generate_random(st, mesh, kpoints, ist_start_, ist_end_, ikpt_start_, ikpt_end_, normalized)
    type(states_elec_t),    intent(inout) :: st
    class(mesh_t),          intent(in)    :: mesh
    type(kpoints_t),        intent(in)    :: kpoints
    integer, optional,      intent(in)    :: ist_start_
    integer, optional,      intent(in)    :: ist_end_
    integer, optional,      intent(in)    :: ikpt_start_
    integer, optional,      intent(in)    :: ikpt_end_
    logical, optional,      intent(in)    :: normalized
 
    integer :: ist, ik, id, ist_start, ist_end, jst, ikpt_start, ikpt_end
    complex(real64)   :: alpha, beta, zdotp_tmp
    real(real64), allocatable :: dpsi(:,  :)
    complex(real64), allocatable :: zpsi(:,  :), zpsi2(:)
    integer :: ikpoint, ip
    type(batch_t) :: ffb
 
    logical :: normalized_
 
    normalized_ = optional_default(normalized, .true.)
 
    push_sub(states_elec_generate_random)
 
    ist_start = optional_default(ist_start_, 1)
    ist_end = optional_default(ist_end_, st%nst)
    ikpt_start = optional_default(ikpt_start_, 1)
    ikpt_end = optional_default(ikpt_end_, st%nik)
 
    safe_allocate(dpsi(1:mesh%np, 1:st%d%dim))
    if (states_are_complex(st)) then
      safe_allocate(zpsi(1:mesh%np, 1:st%d%dim))
    end if
 
    select case (st%d%ispin)
    case (unpolarized, spin_polarized)
 
      do ik = ikpt_start, ikpt_end
        ikpoint = st%d%get_kpoint_index(ik)
        do ist = ist_start, ist_end
          if (states_are_real(st).or.kpoints_point_is_gamma(kpoints, ikpoint)) then
            if (st%randomization == par_independent) then
              call dmf_random(mesh, dpsi(:, 1), &
                pre_shift = mesh%pv%xlocal-1, &
                post_shift = mesh%pv%np_global - mesh%pv%xlocal - mesh%np + 1, &
                normalized = normalized)
              ! Ensures that the grid points are properly distributed in the domain parallel case
              if(mesh%parallel_in_domains) then
                call batch_init(ffb, dpsi(:,1))
                call dmesh_batch_exchange_points(mesh, ffb, backward_map = .true.)
                call ffb%end()
              end if
            else
              call dmf_random(mesh, dpsi(:, 1), normalized = normalized)
            end if
            if (.not. state_kpt_is_local(st, ist, ik)) cycle
            if (states_are_complex(st)) then !Gamma point
              do ip = 1, mesh%np
                zpsi(ip,1) = cmplx(dpsi(ip,1), m_zero, real64)
              end do
              call states_elec_set_state(st, mesh, ist,  ik, zpsi)
            else
              call states_elec_set_state(st, mesh, ist,  ik, dpsi)
            end if
          else
            if (st%randomization == par_independent) then
              call zmf_random(mesh, zpsi(:, 1), &
                pre_shift = mesh%pv%xlocal-1, &
                post_shift = mesh%pv%np_global - mesh%pv%xlocal - mesh%np + 1, &
                normalized = normalized)
              ! Ensures that the grid points are properly distributed in the domain parallel case
              if(mesh%parallel_in_domains) then
                call batch_init(ffb, zpsi(:,1))
                call zmesh_batch_exchange_points(mesh, ffb, backward_map = .true.)
                call ffb%end()
              end if
            else
              call zmf_random(mesh, zpsi(:, 1), normalized = normalized)
            end if
            if (.not. state_kpt_is_local(st, ist, ik)) cycle
            call states_elec_set_state(st, mesh, ist,  ik, zpsi)
          end if
        end do
      end do
 
    case (spinors)
 
      assert(states_are_complex(st))
 
      if (st%fixed_spins) then
 
        do ik = ikpt_start, ikpt_end
          ikpoint = st%d%get_kpoint_index(ik)
          do ist = ist_start, ist_end
            if (kpoints_point_is_gamma(kpoints, ikpoint)) then
              if (st%randomization == par_independent) then
                call dmf_random(mesh, dpsi(:, 1), &
                  pre_shift = mesh%pv%xlocal-1, &
                  post_shift = mesh%pv%np_global - mesh%pv%xlocal - mesh%np + 1, &
                  normalized = normalized)
                ! Ensures that the grid points are properly distributed in the domain parallel case
                if(mesh%parallel_in_domains) then
                  call batch_init(ffb, dpsi(:,1))
                  call dmesh_batch_exchange_points(mesh, ffb, backward_map = .true.)
                  call ffb%end()
                end if
              else
                call dmf_random(mesh, dpsi(:, 1), normalized = normalized)
                if (.not. state_kpt_is_local(st, ist, ik)) cycle
              end if
              do ip = 1, mesh%np
                zpsi(ip,1) = cmplx(dpsi(ip,1), m_zero, real64)
              end do
              call states_elec_set_state(st, mesh, ist,  ik, zpsi)
            else
              if (st%randomization == par_independent) then
                call zmf_random(mesh, zpsi(:, 1), &
                  pre_shift = mesh%pv%xlocal-1, &
                  post_shift = mesh%pv%np_global - mesh%pv%xlocal - mesh%np + 1, &
                  normalized = normalized)
                ! Ensures that the grid points are properly distributed in the domain parallel case
                if(mesh%parallel_in_domains) then
                  call batch_init(ffb, zpsi(:,1))
                  call zmesh_batch_exchange_points(mesh, ffb, backward_map = .true.)
                  call ffb%end()
                end if
              else
                call zmf_random(mesh, zpsi(:, 1), normalized = normalized)
                if (.not. state_kpt_is_local(st, ist, ik)) cycle
              end if
            end if
            if (.not. state_kpt_is_local(st, ist, ik)) cycle
            ! In this case, the spinors are made of a spatial part times a vector [alpha beta]^T in
            ! spin space (i.e., same spatial part for each spin component). So (alpha, beta)
            ! determines the spin values. The values of (alpha, beta) can be be obtained
            ! with simple formulae from <Sx>, <Sy>, <Sz>.
            !
            ! Note that here we orthonormalize the orbital part. This ensures that the spinors
            ! are untouched later in the general orthonormalization, and therefore the spin values
            ! of each spinor remain the same.
            safe_allocate(zpsi2(1:mesh%np))
            do jst = ist_start, ist - 1
              call states_elec_get_state(st, mesh, 1, jst, ik, zpsi2)
              zdotp_tmp = zmf_dotp(mesh, zpsi(:, 1), zpsi2)
              call lalg_axpy(mesh%np, -zdotp_tmp, zpsi2, zpsi(:, 1))
            end do
            safe_deallocate_a(zpsi2)
 
            call zmf_normalize(mesh, 1, zpsi)
            zpsi(1:mesh%np, 2) = zpsi(1:mesh%np, 1)
 
            alpha = cmplx(sqrt(m_half + st%spin(3, ist, ik)), m_zero, real64)
            beta  = cmplx(sqrt(m_one - abs(alpha)**2), m_zero, real64)
            if (abs(alpha) > m_epsilon) then
              beta = cmplx(st%spin(1, ist, ik) / abs(alpha), st%spin(2, ist, ik) / abs(alpha), real64)
            end if
            call lalg_scal(mesh%np, alpha, zpsi(:, 1))
            call lalg_scal(mesh%np, beta, zpsi(:, 2))
            call states_elec_set_state(st, mesh, ist,  ik, zpsi)
          end do
        end do
      else
        do ik = ikpt_start, ikpt_end
          do ist = ist_start, ist_end
            do id = 1, st%d%dim
              if (st%randomization == par_independent) then
                call zmf_random(mesh, zpsi(:, id), &
                  pre_shift = mesh%pv%xlocal-1, &
                  post_shift = mesh%pv%np_global - mesh%pv%xlocal - mesh%np + 1, &
                  normalized = .false.)
                ! Ensures that the grid points are properly distributed in the domain parallel case
                if(mesh%parallel_in_domains) then
                  call batch_init(ffb, zpsi(:, id))
                  call zmesh_batch_exchange_points(mesh, ffb, backward_map = .true.)
                  call ffb%end()
                end if
              else
                call zmf_random(mesh, zpsi(:, id), normalized = .false.)
              end if
            end do
            ! We need to generate the wave functions even if not using them in order to be consistent
            ! with the random seed in parallel runs.
            if (.not. state_kpt_is_local(st, ist, ik)) cycle
            ! Note that mf_random normalizes each spin channel independently to 1.
            ! Therefore we need to renormalize the spinor:
            if (normalized_) call zmf_normalize(mesh, st%d%dim, zpsi)
            call states_elec_set_state(st, mesh, ist,  ik, zpsi)
          end do
        end do
      end if
 
    end select
 
    safe_deallocate_a(dpsi)
    safe_deallocate_a(zpsi)
 
    pop_sub(states_elec_generate_random)
  end subroutine states_elec_generate_random
 
  ! ---------------------------------------------------------
  !
  subroutine states_elec_fermi(st, namespace, mesh, compute_spin)
    type(states_elec_t), intent(inout) :: st
    type(namespace_t),   intent(in)    :: namespace
    class(mesh_t),       intent(in)    :: mesh
    logical, optional,   intent(in)    :: compute_spin
 
    integer            :: ist, ik
    real(real64)       :: charge
    complex(real64), allocatable :: zpsi(:, :)
 
    push_sub(states_elec_fermi)
 
    call smear_find_fermi_energy(st%smear, namespace, st%eigenval, st%occ, st%qtot, &
      st%nik, st%nst, st%kweights)
 
    call smear_fill_occupations(st%smear, st%eigenval, st%occ, st%kweights, st%nik, st%nst)
 
    ! check if everything is OK
    charge = m_zero
    do ist = 1, st%nst
      charge = charge + sum(st%occ(ist, 1:st%nik) * st%kweights(1:st%nik))
    end do
    if (abs(charge-st%qtot) > 1e-6_real64) then
      message(1) = 'Occupations do not integrate to total charge.'
      write(message(2), '(6x,f12.8,a,f12.8)') charge, ' != ', st%qtot
      call messages_warning(2, namespace=namespace)
      if (charge < m_epsilon) then
        message(1) = "There don't seem to be any electrons at all!"
        call messages_fatal(1, namespace=namespace)
      end if
    end if
 
    if (st%d%ispin == spinors .and. optional_default(compute_spin,.true.)) then
      assert(states_are_complex(st))
 
      st%spin(:,:,:) = m_zero
 
      safe_allocate(zpsi(1:mesh%np, st%d%dim))
      do ik = st%d%kpt%start, st%d%kpt%end
        do ist = st%st_start, st%st_end
          call states_elec_get_state(st, mesh, ist, ik, zpsi)
          st%spin(1:3, ist, ik) = state_spin(mesh, zpsi)
        end do
      end do
      safe_deallocate_a(zpsi)
 
      if (st%parallel_in_states .or. st%d%kpt%parallel) then
        call comm_allreduce(st%st_kpt_mpi_grp, st%spin)
      end if
 
    end if
 
    pop_sub(states_elec_fermi)
  end subroutine states_elec_fermi
 
 
  ! ---------------------------------------------------------
  !
  function states_elec_eigenvalues_sum(st, alt_eig) result(tot)
    type(states_elec_t),  intent(in) :: st
    real(real64),      optional, intent(in) :: alt_eig(st%st_start:, st%d%kpt%start:)
    real(real64)                     :: tot
 
    integer :: ik
 
    push_sub(states_elec_eigenvalues_sum)
 
    tot = m_zero
    do ik = st%d%kpt%start, st%d%kpt%end
      if (present(alt_eig)) then
        tot = tot + st%kweights(ik) * sum(st%occ(st%st_start:st%st_end, ik) * &
          alt_eig(st%st_start:st%st_end, ik))
      else
        tot = tot + st%kweights(ik) * sum(st%occ(st%st_start:st%st_end, ik) * &
          st%eigenval(st%st_start:st%st_end, ik))
      end if
    end do
 
    if (st%parallel_in_states .or. st%d%kpt%parallel) call comm_allreduce(st%st_kpt_mpi_grp, tot)
 
    pop_sub(states_elec_eigenvalues_sum)
  end function states_elec_eigenvalues_sum
 
 
  subroutine states_elec_distribute_nodes(st, namespace, mc)
    type(states_elec_t),    intent(inout) :: st
    type(namespace_t),      intent(in)    :: namespace
    type(multicomm_t),      intent(in)    :: mc
 
    logical :: default_scalapack_compatible
 
    push_sub(states_elec_distribute_nodes)
 
    ! TODO(Alex) Issue #820. This is superflous. These defaults are set in initialisation of
    ! states, and in the state distribution instance
    ! Defaults.
    st%node(:)            = 0
    st%st_start           = 1
    st%st_end             = st%nst
    st%lnst               = st%nst
    st%parallel_in_states = .false.
 
    call mpi_grp_init(st%mpi_grp, mc%group_comm(p_strategy_states))
    call mpi_grp_init(st%system_grp, mc%master_comm)
    call mpi_grp_init(st%dom_st_kpt_mpi_grp, mc%dom_st_kpt_comm)
    call mpi_grp_init(st%dom_st_mpi_grp, mc%dom_st_comm)
    call mpi_grp_init(st%st_kpt_mpi_grp, mc%st_kpt_comm)
 
    default_scalapack_compatible = calc_mode_par%scalapack_compat() .and. .not. st%d%kpt%parallel
 
    !%Variable ScaLAPACKCompatible
    !%Type logical
    !%Section Execution::Parallelization
    !%Description
    !% Whether to use a layout for states parallelization which is compatible with ScaLAPACK.
    !% The default is yes for <tt>CalculationMode = gs, unocc, go</tt> without k-point parallelization,
    !% and no otherwise. (Setting to other than default is experimental.)
    !% The value must be yes if any ScaLAPACK routines are called in the course of the run;
    !% it must be set by hand for <tt>td</tt> with <tt>TDDynamics = bo</tt>.
    !% This variable has no effect unless you are using states parallelization and have linked ScaLAPACK.
    !% Note: currently, use of ScaLAPACK is not compatible with task parallelization (<i>i.e.</i> slaves).
    !%End
    call parse_variable(namespace, 'ScaLAPACKCompatible', default_scalapack_compatible, st%scalapack_compatible)
 
#ifdef HAVE_SCALAPACK
    if (default_scalapack_compatible .neqv. st%scalapack_compatible) then
      call messages_experimental('Setting ScaLAPACKCompatible to other than default', namespace=namespace)
    end if
 
    if (st%scalapack_compatible) then
      if (multicomm_have_slaves(mc)) then
        call messages_not_implemented("ScaLAPACK usage with task parallelization (slaves)", namespace=namespace)
      end if
      call blacs_proc_grid_init(st%dom_st_proc_grid, st%dom_st_mpi_grp)
    end if
#else
    st%scalapack_compatible = .false.
#endif
 
    if (multicomm_strategy_is_parallel(mc, p_strategy_states)) then
 
#ifdef HAVE_MPI
      call multicomm_create_all_pairs(st%mpi_grp, st%ap)
#endif
 
      if (st%nst < st%mpi_grp%size) then
        message(1) = "Have more processors than necessary"
        write(message(2),'(i4,a,i4,a)') st%mpi_grp%size, " processors and ", st%nst, " states."
        call messages_fatal(2, namespace=namespace)
      end if
 
      call distributed_init(st%dist, st%nst, st%mpi_grp%comm, "states", scalapack_compat = st%scalapack_compatible)
 
      st%parallel_in_states = st%dist%parallel
 
      ! TODO(Alex) Issue #820. Remove lnst, st_start, st_end and node, as they are all contained within dist
      st%st_start = st%dist%start
      st%st_end   = st%dist%end
      st%lnst     = st%dist%nlocal
      st%node(1:st%nst) = st%dist%node(1:st%nst)
 
    end if
 
    call states_elec_kpoints_distribution(st)
 
    pop_sub(states_elec_distribute_nodes)
  end subroutine states_elec_distribute_nodes
 
 
  !
  subroutine states_elec_calc_quantities(gr, st, kpoints, nlcc, &
    kinetic_energy_density, paramagnetic_current, density_gradient, density_laplacian, &
    gi_kinetic_energy_density, st_end)
    type(grid_t),            intent(in)    :: gr
    type(states_elec_t),     intent(in)    :: st
    type(kpoints_t),         intent(in)    :: kpoints
    logical,                 intent(in)    :: nlcc
    real(real64), contiguous, optional, target, intent(out) :: kinetic_energy_density(:,:)
    real(real64), contiguous, optional, target, intent(out) :: paramagnetic_current(:,:,:)
    real(real64), contiguous, optional, intent(out) :: density_gradient(:,:,:)
    real(real64), contiguous, optional, intent(out) :: density_laplacian(:,:)
    real(real64), contiguous, optional, intent(out) :: gi_kinetic_energy_density(:,:)
    integer, optional,       intent(in)    :: st_end
 
    real(real64), pointer, contiguous :: jp(:, :, :)
    real(real64), pointer, contiguous :: tau(:, :)
    complex(real64), allocatable :: wf_psi(:,:), gwf_psi(:,:,:), wf_psi_conj(:,:), lwf_psi(:,:)
    real(real64), allocatable :: abs_wf_psi(:,:), abs_gwf_psi(:,:)
    complex(real64), allocatable :: psi_gpsi(:,:)
    complex(real64)   :: c_tmp
    integer :: is, ik, ist, i_dim, st_dim, ii, st_end_
    real(real64)   :: ww, kpoint(gr%der%dim)
    logical :: something_to_do
 
    call profiling_in("STATES_CALC_QUANTITIES")
 
    push_sub(states_elec_calc_quantities)
 
    st_end_ = min(st%st_end, optional_default(st_end, st%st_end))
 
    something_to_do = present(kinetic_energy_density) .or. present(gi_kinetic_energy_density) .or. &
      present(paramagnetic_current) .or. present(density_gradient) .or. present(density_laplacian)
    assert(something_to_do)
 
    safe_allocate( wf_psi(1:gr%np_part, 1:st%d%dim))
    safe_allocate( wf_psi_conj(1:gr%np_part, 1:st%d%dim))
    safe_allocate(gwf_psi(1:gr%np, 1:gr%der%dim, 1:st%d%dim))
    safe_allocate(abs_wf_psi(1:gr%np, 1:st%d%dim))
    safe_allocate(abs_gwf_psi(1:gr%np, 1:st%d%dim))
    safe_allocate(psi_gpsi(1:gr%np, 1:st%d%dim))
    if (present(density_laplacian)) then
      safe_allocate(lwf_psi(1:gr%np, 1:st%d%dim))
    end if
 
    nullify(tau)
    if (present(kinetic_energy_density)) tau => kinetic_energy_density
 
    nullify(jp)
    if (present(paramagnetic_current)) jp => paramagnetic_current
 
    ! for the gauge-invariant kinetic energy density we need the
    ! current and the kinetic energy density
    if (present(gi_kinetic_energy_density)) then
      if (.not. present(paramagnetic_current) .and. states_are_complex(st)) then
        safe_allocate(jp(1:gr%np, 1:gr%der%dim, 1:st%d%nspin))
      end if
      if (.not. present(kinetic_energy_density)) then
        safe_allocate(tau(1:gr%np, 1:st%d%nspin))
      end if
    end if
 
    if (associated(tau)) tau = m_zero
    if (associated(jp)) jp = m_zero
    if (present(density_gradient)) density_gradient(:,:,:) = m_zero
    if (present(density_laplacian)) density_laplacian(:,:) = m_zero
    if (present(gi_kinetic_energy_density)) gi_kinetic_energy_density = m_zero
 
    do ik = st%d%kpt%start, st%d%kpt%end
 
      kpoint(1:gr%der%dim) = kpoints%get_point(st%d%get_kpoint_index(ik))
      is = st%d%get_spin_index(ik)
 
      do ist = st%st_start, st_end_
        ww = st%kweights(ik)*st%occ(ist, ik)
        if (abs(ww) <= m_epsilon) cycle
 
        ! all calculations will be done with complex wavefunctions
        call states_elec_get_state(st, gr, ist, ik, wf_psi)
 
        do st_dim = 1, st%d%dim
          call boundaries_set(gr%der%boundaries, gr, wf_psi(:, st_dim))
        end do
 
        ! calculate gradient of the wavefunction
        do st_dim = 1, st%d%dim
          call zderivatives_grad(gr%der, wf_psi(:,st_dim), gwf_psi(:,:,st_dim), set_bc = .false.)
        end do
 
        ! calculate the Laplacian of the wavefunction
        if (present(density_laplacian)) then
          do st_dim = 1, st%d%dim
            call zderivatives_lapl(gr%der, wf_psi(:,st_dim), lwf_psi(:,st_dim), ghost_update = .false., set_bc = .false.)
          end do
        end if
 
        ! We precompute some quantites, to avoid to compute it many times
        wf_psi_conj(1:gr%np, 1:st%d%dim) = conjg(wf_psi(1:gr%np,1:st%d%dim))
        !$omp parallel do collapse(2)
        do st_dim = 1, st%d%dim
          do ii = 1, gr%np
            abs_wf_psi(ii, st_dim) = real(wf_psi_conj(ii, st_dim) * wf_psi(ii, st_dim), real64)
          end do
        end do
        !$omp end parallel do
 
        if (present(density_laplacian)) then
          !$omp parallel do
          do ii = 1, gr%np
            density_laplacian(ii, is) = density_laplacian(ii, is) + &
              ww * m_two*real(wf_psi_conj(ii, 1) * lwf_psi(ii, 1), real64)
          end do
          !$omp end parallel do
          if (st%d%ispin == spinors) then
            !$omp parallel do
            do ii = 1, gr%np
              density_laplacian(ii, 2) = density_laplacian(ii, 2) + &
                ww * m_two*real(wf_psi_conj(ii, 2) * lwf_psi(ii, 2), real64)
            end do
            !$omp end parallel do
            !$omp parallel do private(c_tmp)
            do ii = 1, gr%np
              c_tmp = ww*(lwf_psi(ii, 1) * wf_psi_conj(ii, 2) + wf_psi(ii, 1) * conjg(lwf_psi(ii, 2)))
              density_laplacian(ii, 3) = density_laplacian(ii, 3) + real(c_tmp, real64)
              density_laplacian(ii, 4) = density_laplacian(ii, 4) + aimag(c_tmp)
            end do
          end if
        end if
 
        if (associated(tau)) then
          call lalg_axpy(gr%np, ww * sum(kpoint(1:gr%der%dim)**2), abs_wf_psi(:, 1), tau(:, is))
          if (st%d%ispin == spinors) then
            call lalg_axpy(gr%np, ww * sum(kpoint(1:gr%der%dim)**2), abs_wf_psi(:, 2), tau(:, 2))
 
            !$omp parallel do private(c_tmp)
            do ii = 1, gr%np
              c_tmp = ww * sum(kpoint(1:gr%der%dim)**2) * wf_psi(ii, 1) * wf_psi_conj(ii, 2)
              tau(ii, 3) = tau(ii, 3) + real(c_tmp, real64)
              tau(ii, 4) = tau(ii, 4) + aimag(c_tmp)
            end do
          end if
        end if
 
        do i_dim = 1, gr%der%dim
 
          ! We precompute some quantites, to avoid to compute them many times
          !$omp parallel private(st_dim, ii)
          do st_dim = 1, st%d%dim
            !$omp do
            do ii = 1, gr%np
              psi_gpsi(ii, st_dim) = wf_psi_conj(ii, st_dim) * gwf_psi(ii, i_dim, st_dim)
            end do
          end do
          do st_dim = 1, st%d%dim
            !$omp do
            do ii = 1, gr%np
              abs_gwf_psi(ii, st_dim) = real(conjg(gwf_psi(ii, i_dim, st_dim)) &
                * gwf_psi(ii, i_dim, st_dim), real64)
            end do
          end do
          !$omp end parallel
 
          if (present(density_gradient)) then
            !$omp parallel do
            do ii = 1, gr%np
              density_gradient(ii, i_dim, is) = density_gradient(ii, i_dim, is) &
                + ww * m_two * real(psi_gpsi(ii, 1), real64)
            end do
            !$omp end parallel do
            if (st%d%ispin == spinors) then
              !$omp parallel do
              do ii = 1, gr%np
                density_gradient(ii, i_dim, 2) = density_gradient(ii, i_dim, 2)  &
                  + ww * m_two*real(psi_gpsi(ii, 2), real64)
              end do
              !$omp end parallel do
              !$omp parallel do private(c_tmp)
              do ii = 1, gr%np
                c_tmp = ww * (gwf_psi(ii, i_dim, 1) * wf_psi_conj(ii, 2) + wf_psi(ii, 1) * conjg(gwf_psi(ii, i_dim, 2)))
                density_gradient(ii, i_dim, 3) = density_gradient(ii, i_dim, 3) + real(c_tmp, real64)
                density_gradient(ii, i_dim, 4) = density_gradient(ii, i_dim, 4) + aimag(c_tmp)
              end do
            end if
          end if
 
          if (present(density_laplacian)) then
            call lalg_axpy(gr%np, ww*m_two, abs_gwf_psi(:,1), density_laplacian(:,is))
            if (st%d%ispin == spinors) then
              call lalg_axpy(gr%np, ww*m_two, abs_gwf_psi(:,2), density_laplacian(:,2))
              !$omp parallel do private(c_tmp)
              do ii = 1, gr%np
                c_tmp = m_two * ww * gwf_psi(ii, i_dim, 1) * conjg(gwf_psi(ii, i_dim, 2))
                density_laplacian(ii, 3) = density_laplacian(ii, 3) + real(c_tmp, real64)
                density_laplacian(ii, 4) = density_laplacian(ii, 4) + aimag(c_tmp)
              end do
            end if
          end if
 
          ! the expression for the paramagnetic current with spinors is
          !     j = ( jp(1)             jp(3) + i jp(4) )
          !         (-jp(3) + i jp(4)   jp(2)           )
          if (associated(jp)) then
            if (.not.(states_are_real(st))) then
              !$omp parallel do
              do ii = 1, gr%np
                jp(ii, i_dim, is) = jp(ii, i_dim, is) + &
                  ww*(aimag(psi_gpsi(ii, 1)) - abs_wf_psi(ii, 1)*kpoint(i_dim))
              end do
              !$omp end parallel do
              if (st%d%ispin == spinors) then
                !$omp parallel do
                do ii = 1, gr%np
                  jp(ii, i_dim, 2) = jp(ii, i_dim, 2) + &
                    ww*( aimag(psi_gpsi(ii, 2)) - abs_wf_psi(ii, 2)*kpoint(i_dim))
                end do
                !$omp end parallel do
                !$omp parallel do private(c_tmp)
                do ii = 1, gr%np
                  c_tmp = -ww*m_half*m_zi*(gwf_psi(ii, i_dim, 1)*wf_psi_conj(ii, 2) - wf_psi(ii, 1)*conjg(gwf_psi(ii, i_dim, 2)) &
                    - m_two * m_zi*wf_psi(ii, 1)*wf_psi_conj(ii, 2)*kpoint(i_dim))
                  jp(ii, i_dim, 3) = jp(ii, i_dim, 3) + real(c_tmp, real64)
                  jp(ii, i_dim, 4) = jp(ii, i_dim, 4) + aimag(c_tmp)
                end do
              end if
            end if
          end if
 
          ! the expression for the paramagnetic current with spinors is
          !     t = ( tau(1)              tau(3) + i tau(4) )
          !         ( tau(3) - i tau(4)   tau(2)            )
          if (associated(tau)) then
            !$omp parallel do
            do ii = 1, gr%np
              tau(ii, is) = tau(ii, is) + ww*(abs_gwf_psi(ii,1) &
                - m_two*aimag(psi_gpsi(ii, 1))*kpoint(i_dim))
            end do
            !$omp end parallel do
            if (st%d%ispin == spinors) then
              !$omp parallel do
              do ii = 1, gr%np
                tau(ii, 2) = tau(ii, 2) + ww*(abs_gwf_psi(ii, 2) &
                  - m_two*aimag(psi_gpsi(ii, 2))*kpoint(i_dim))
              end do
              !$omp end parallel do
              !$omp parallel do private(c_tmp)
              do ii = 1, gr%np
                c_tmp = ww * ( gwf_psi(ii, i_dim, 1)*conjg(gwf_psi(ii, i_dim, 2))  &
                  + m_zi * (gwf_psi(ii,i_dim,1)*wf_psi_conj(ii, 2)    &
                  - wf_psi(ii, 1)*conjg(gwf_psi(ii,i_dim,2)))*kpoint(i_dim))
                tau(ii, 3) = tau(ii, 3) + real(c_tmp, real64)
                tau(ii, 4) = tau(ii, 4) + aimag(c_tmp)
              end do
            end if
          end if
 
        end do
 
      end do
    end do
 
    safe_deallocate_a(wf_psi_conj)
    safe_deallocate_a(abs_wf_psi)
    safe_deallocate_a(abs_gwf_psi)
    safe_deallocate_a(psi_gpsi)
 
    if (.not. present(gi_kinetic_energy_density)) then
      if (.not. present(paramagnetic_current)) then
        safe_deallocate_p(jp)
      end if
      if (.not. present(kinetic_energy_density)) then
        safe_deallocate_p(tau)
      end if
    end if
 
    if (st%parallel_in_states .or. st%d%kpt%parallel) call reduce_all(st%st_kpt_mpi_grp)
 
    ! We have to symmetrize everything as they are calculated from the
    ! wavefunctions.
    ! This must be done before compute the gauge-invariant kinetic energy density
    if (st%symmetrize_density) then
      do is = 1, st%d%nspin
        if (associated(tau)) then
          call dgrid_symmetrize_scalar_field(gr, tau(:, is), suppress_warning = .true.)
        end if
 
        if (present(density_laplacian)) then
          call dgrid_symmetrize_scalar_field(gr, density_laplacian(:, is), suppress_warning = .true.)
        end if
 
        if (associated(jp)) then
          call dgrid_symmetrize_vector_field(gr, jp(:, :, is), suppress_warning = .true.)
        end if
 
        if (present(density_gradient)) then
          call dgrid_symmetrize_vector_field(gr, density_gradient(:, :, is), suppress_warning = .true.)
        end if
      end do
    end if
 
 
    if (allocated(st%rho_core) .and. nlcc .and. (present(density_laplacian) .or. present(density_gradient))) then
      do ii = 1, gr%np
        wf_psi(ii, 1) = st%rho_core(ii)/st%d%spin_channels
      end do
 
      call boundaries_set(gr%der%boundaries, gr, wf_psi(:, 1))
 
      if (present(density_gradient)) then
        ! calculate gradient of the NLCC
        call zderivatives_grad(gr%der, wf_psi(:,1), gwf_psi(:,:,1), set_bc = .false.)
        do is = 1, st%d%spin_channels
          !$omp parallel do collapse(2)
          do i_dim = 1, gr%der%dim
            do ii = 1, gr%np
              density_gradient(ii, i_dim, is) = density_gradient(ii, i_dim, is) + &
                real(gwf_psi(ii, i_dim, 1), real64)
            end do
          end do
          !$omp end parallel do
        end do
      end if
 
      ! calculate the Laplacian of the wavefunction
      if (present(density_laplacian)) then
        call zderivatives_lapl(gr%der, wf_psi(:,1), lwf_psi(:,1), set_bc = .false.)
 
        !$omp parallel private(is, ii)
        do is = 1, st%d%spin_channels
          !$omp do
          do ii = 1, gr%np
            density_laplacian(ii, is) = density_laplacian(ii, is) + real(lwf_psi(ii, 1))
          end do
        end do
        !$omp end parallel
      end if
    end if
 
    !If we freeze some of the orbitals, we need to had the contributions here
    !Only in the case we are not computing it
    if (allocated(st%frozen_tau) .and. .not. present(st_end) .and. associated(tau)) then
      call lalg_axpy(gr%np, st%d%nspin, m_one, st%frozen_tau, tau)
    end if
    if (allocated(st%frozen_gdens) .and. .not. present(st_end) .and. present(density_gradient)) then
      do is = 1, st%d%nspin
        call lalg_axpy(gr%np, gr%der%dim, m_one, st%frozen_gdens(:,:,is), density_gradient(:,:,is))
      end do
    end if
    if (allocated(st%frozen_ldens) .and. .not. present(st_end) .and. present(density_laplacian)) then
      call lalg_axpy(gr%np, st%d%nspin, m_one, st%frozen_ldens, density_laplacian)
    end if
 
    safe_deallocate_a(wf_psi)
    safe_deallocate_a(gwf_psi)
    safe_deallocate_a(lwf_psi)
 
 
    !We compute the gauge-invariant kinetic energy density
    if (present(gi_kinetic_energy_density)) then
      do is = 1, st%d%nspin
        assert(associated(tau))
        gi_kinetic_energy_density(1:gr%np, is) = tau(1:gr%np, is)
      end do
      if (states_are_complex(st)) then
        assert(associated(jp))
        if (st%d%ispin /= spinors) then
          do is = 1, st%d%nspin
            !$omp parallel do
            do ii = 1, gr%np
              if (st%rho(ii, is) < 1.0e-7_real64) cycle
              gi_kinetic_energy_density(ii, is) = &
                gi_kinetic_energy_density(ii, is) - sum(jp(ii,1:gr%der%dim, is)**2)/st%rho(ii, is)
            end do
          end do
        else ! Note that this is only the U(1) part of the gauge-invariant KED
          !$omp parallel do
          do ii = 1, gr%np
            gi_kinetic_energy_density(ii, 1) = gi_kinetic_energy_density(ii, 1) &
              - sum(jp(ii,1:gr%der%dim, 1)**2)/(safe_tol(st%rho(ii, 1),m_epsilon))
            gi_kinetic_energy_density(ii, 2) = gi_kinetic_energy_density(ii, 2) &
              - sum(jp(ii,1:gr%der%dim, 2)**2)/(safe_tol(st%rho(ii, 2),m_epsilon))
            gi_kinetic_energy_density(ii, 3) = &
              gi_kinetic_energy_density(ii, 3) - sum(jp(ii,1:gr%der%dim, 3)**2 + jp(ii,1:gr%der%dim, 4)**2) &
              /(safe_tol((st%rho(ii, 3)**2 + st%rho(ii, 4)**2), m_epsilon))*st%rho(ii, 3)
            gi_kinetic_energy_density(ii, 4) = &
              gi_kinetic_energy_density(ii, 4) + sum(jp(ii,1:gr%der%dim, 3)**2 + jp(ii,1:gr%der%dim, 4)**2) &
              /(safe_tol((st%rho(ii, 3)**2 + st%rho(ii, 4)**2), m_epsilon))*st%rho(ii, 4)
          end do
        end if
      end if
    end if
 
    if (.not. present(kinetic_energy_density)) then
      safe_deallocate_p(tau)
    end if
    if (.not. present(paramagnetic_current)) then
      safe_deallocate_p(jp)
    end if
 
 
    pop_sub(states_elec_calc_quantities)
 
    call profiling_out("STATES_CALC_QUANTITIES")
 
  contains
 
    subroutine reduce_all(grp)
      type(mpi_grp_t), intent(in)  :: grp
 
      push_sub(states_elec_calc_quantities.reduce_all)
 
      if (associated(tau)) call comm_allreduce(grp, tau, dim = (/gr%np, st%d%nspin/))
 
      if (present(density_laplacian)) call comm_allreduce(grp, density_laplacian, dim = (/gr%np, st%d%nspin/))
 
      do is = 1, st%d%nspin
        if (associated(jp)) call comm_allreduce(grp, jp(:, :, is), dim = (/gr%np, gr%der%dim/))
 
        if (present(density_gradient)) then
          call comm_allreduce(grp, density_gradient(:, :, is), dim = (/gr%np, gr%der%dim/))
        end if
      end do
 
      pop_sub(states_elec_calc_quantities.reduce_all)
    end subroutine reduce_all
 
  end subroutine states_elec_calc_quantities
 
 
  ! ---------------------------------------------------------
  !
  function state_spin(mesh, f1) result(spin)
    type(mesh_t), intent(in) :: mesh
    complex(real64),        intent(in) :: f1(:, :)
    real(real64)             :: spin(1:3)
 
    complex(real64) :: z
 
    push_sub(state_spin)
 
    z = zmf_dotp(mesh, f1(:, 1) , f1(:, 2))
 
    spin(1) = m_two*real(z, real64)
    spin(2) = m_two*aimag(z)
    spin(3) = zmf_nrm2(mesh, f1(:, 1))**2 - zmf_nrm2(mesh, f1(:, 2))**2
    spin = m_half*spin ! spin is half the sigma matrix.
 
    pop_sub(state_spin)
  end function state_spin
 
  ! ---------------------------------------------------------
  !
  logical function state_is_local(st, ist)
    type(states_elec_t), intent(in) :: st
    integer,             intent(in) :: ist
 
    push_sub(state_is_local)
 
    state_is_local = ist >= st%st_start.and.ist <= st%st_end
 
    pop_sub(state_is_local)
  end function state_is_local
 
  ! ---------------------------------------------------------
  !
  logical function state_kpt_is_local(st, ist, ik)
    type(states_elec_t), intent(in) :: st
    integer,             intent(in) :: ist
    integer,             intent(in) :: ik
 
    push_sub(state_kpt_is_local)
 
    state_kpt_is_local = ist >= st%st_start .and. ist <= st%st_end .and. &
      ik >= st%d%kpt%start .and. ik <= st%d%kpt%end
 
    pop_sub(state_kpt_is_local)
  end function state_kpt_is_local
 
 
  ! ---------------------------------------------------------
  real(real64) function states_elec_wfns_memory(st, mesh) result(memory)
    type(states_elec_t), intent(in) :: st
    class(mesh_t),       intent(in) :: mesh
 
    push_sub(states_elec_wfns_memory)
    memory = m_zero
 
    ! orbitals
    memory = memory + real(storage_size(0.0_real64)/8, real64) * &
      real(mesh%np_part_global, real64) *st%d%dim * real(st%nst, real64) *st%d%kpt%nglobal
 
    pop_sub(states_elec_wfns_memory)
  end function states_elec_wfns_memory
 
  ! ---------------------------------------------------------
  !
  subroutine states_elec_pack(st, copy)
    class(states_elec_t),    intent(inout) :: st
    logical,      optional, intent(in)    :: copy
 
    integer :: iqn, ib
    integer(int64) :: max_mem, mem
 
    push_sub(states_elec_pack)
 
    ! nothing to do, already packed
    if (st%packed) then
      pop_sub(states_elec_pack)
      return
    end if
 
    st%packed = .true.
 
    if (accel_is_enabled()) then
      max_mem = accel_global_memory_size()
 
      if (st%gpu_states_mem > m_one) then
        max_mem = int(st%gpu_states_mem, int64)*(1024_8)**2
      else if (st%gpu_states_mem < 0.0_real64) then
        max_mem = max_mem + int(st%gpu_states_mem, int64)*(1024_8)**2
      else
        max_mem = int(st%gpu_states_mem*real(max_mem, real64) , int64)
      end if
    else
      max_mem = huge(max_mem)
    end if
 
    mem = 0
    qnloop: do iqn = st%d%kpt%start, st%d%kpt%end
      do ib = st%group%block_start, st%group%block_end
 
        mem = mem + st%group%psib(ib, iqn)%pack_total_size()
 
        if (mem > max_mem) then
          call messages_write('Not enough CL device memory to store all states simultaneously.', new_line = .true.)
          call messages_write('Only ')
          call messages_write(ib - st%group%block_start)
          call messages_write(' of ')
          call messages_write(st%group%block_end - st%group%block_start + 1)
          call messages_write(' blocks will be stored in device memory.', new_line = .true.)
          call messages_warning()
          exit qnloop
        end if
 
        call st%group%psib(ib, iqn)%do_pack(copy)
      end do
    end do qnloop
 
    pop_sub(states_elec_pack)
  end subroutine states_elec_pack
 
  ! ------------------------------------------------------------
  !
  subroutine states_elec_unpack(st, copy)
    class(states_elec_t),    intent(inout) :: st
    logical,      optional, intent(in)    :: copy
 
    integer :: iqn, ib
 
    push_sub(states_elec_unpack)
 
    if (st%packed) then
      st%packed = .false.
 
      do iqn = st%d%kpt%start, st%d%kpt%end
        do ib = st%group%block_start, st%group%block_end
          if (st%group%psib(ib, iqn)%is_packed()) call st%group%psib(ib, iqn)%do_unpack(copy)
        end do
      end do
    end if
 
    pop_sub(states_elec_unpack)
  end subroutine states_elec_unpack
 
  ! -----------------------------------------------------------
  !
  subroutine states_elec_write_info(st, namespace)
    class(states_elec_t),    intent(in) :: st
    type(namespace_t),       intent(in) :: namespace
 
    push_sub(states_elec_write_info)
 
    call messages_print_with_emphasis(msg="States", namespace=namespace)
 
    write(message(1), '(a,f12.3)') 'Total electronic charge  = ', st%qtot
    write(message(2), '(a,i8)')    'Number of states         = ', st%nst
    write(message(3), '(a,i8)')    'States block-size        = ', st%block_size
    call messages_info(3, namespace=namespace)
 
    call messages_print_with_emphasis(namespace=namespace)
 
    pop_sub(states_elec_write_info)
  end subroutine states_elec_write_info
 
  subroutine states_elec_set_zero(st)
    class(states_elec_t),    intent(inout) :: st
 
    integer :: iqn, ib
 
    push_sub(states_elec_set_zero)
 
    do iqn = st%d%kpt%start, st%d%kpt%end
      do ib = st%group%block_start, st%group%block_end
        call batch_set_zero(st%group%psib(ib, iqn))
      end do
    end do
 
    pop_sub(states_elec_set_zero)
  end subroutine states_elec_set_zero
 
  ! ------------------------------------------------------------
  !
  integer pure function states_elec_block_min(st, ib) result(range)
    type(states_elec_t),    intent(in) :: st
    integer,                intent(in) :: ib
 
    range = st%group%block_range(ib, 1)
  end function states_elec_block_min
 
  ! ------------------------------------------------------------
  !
  integer pure function states_elec_block_max(st, ib) result(range)
    type(states_elec_t),    intent(in) :: st
    integer,                intent(in) :: ib
 
    range = st%group%block_range(ib, 2)
  end function states_elec_block_max
 
  ! ------------------------------------------------------------
  !
  integer pure function states_elec_block_size(st, ib) result(size)
    type(states_elec_t),    intent(in) :: st
    integer,           intent(in) :: ib
 
    size = st%group%block_size(ib)
  end function states_elec_block_size
 
  ! ------------------------------------------------------------
  !
  logical pure function states_elec_all_states_gpu_packed(st) result(all_packed)
    type(states_elec_t),    intent(in) :: st
 
    integer :: ib, ik
 
    all_packed = .true.
    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)%status() /= batch_device_packed) then
          all_packed = .false.
          return
        end if
      end do
    end do
  end function states_elec_all_states_gpu_packed
 
  !
  subroutine occupied_states(st, namespace, ik, n_filled, n_partially_filled, n_half_filled, &
    filled, partially_filled, half_filled)
    type(states_elec_t),    intent(in)  :: st
    type(namespace_t),      intent(in)  :: namespace
    integer,                intent(in)  :: ik
    integer,                intent(out) :: n_filled, n_partially_filled, n_half_filled
    integer,      optional, intent(out) :: filled(:), partially_filled(:), half_filled(:)
 
    integer :: ist
 
    push_sub(occupied_states)
 
    if (present(filled))           filled(:) = 0
    if (present(partially_filled)) partially_filled(:) = 0
    if (present(half_filled))      half_filled(:) = 0
    n_filled = 0
    n_partially_filled = 0
    n_half_filled = 0
 
    select case (st%d%ispin)
    case (unpolarized)
      do ist = 1, st%nst
        if (abs(st%occ(ist, ik) - m_two) < m_min_occ) then
          n_filled = n_filled + 1
          if (present(filled)) filled(n_filled) = ist
        else if (abs(st%occ(ist, ik) - m_one) < m_min_occ) then
          n_half_filled = n_half_filled + 1
          if (present(half_filled)) half_filled(n_half_filled) = ist
        else if (st%occ(ist, ik) > m_min_occ) then
          n_partially_filled = n_partially_filled + 1
          if (present(partially_filled)) partially_filled(n_partially_filled) = ist
        else if (abs(st%occ(ist, ik)) > m_min_occ) then
          write(message(1),*) 'Internal error in occupied_states: Illegal occupation value ', st%occ(ist, ik)
          call messages_fatal(1, namespace=namespace)
        end if
      end do
    case (spin_polarized, spinors)
      do ist = 1, st%nst
        if (abs(st%occ(ist, ik)-m_one) < m_min_occ) then
          n_filled = n_filled + 1
          if (present(filled)) filled(n_filled) = ist
        else if (st%occ(ist, ik) > m_min_occ) then
          n_partially_filled = n_partially_filled + 1
          if (present(partially_filled)) partially_filled(n_partially_filled) = ist
        else if (abs(st%occ(ist, ik)) > m_min_occ) then
          write(message(1),*) 'Internal error in occupied_states: Illegal occupation value ', st%occ(ist, ik)
          call messages_fatal(1, namespace=namespace)
        end if
      end do
    end select
 
    pop_sub(occupied_states)
  end subroutine occupied_states
 
 
  subroutine kpoints_distribute(this, mc)
    type(states_elec_t), intent(inout) :: this
    type(multicomm_t),       intent(in)    :: mc
 
    push_sub(kpoints_distribute)
    call distributed_init(this%d%kpt, this%nik, mc%group_comm(p_strategy_kpoints), "k-points")
 
    pop_sub(kpoints_distribute)
  end subroutine kpoints_distribute
 
 
  ! TODO(Alex) Issue #824. Consider converting this to a function to returns `st_kpt_task`
  ! as this is only called in a couple of places, or package with the `st_kpt_mpi_grp` when split
  ! from st instance
  subroutine states_elec_kpoints_distribution(st)
    type(states_elec_t),    intent(inout) :: st
 
    push_sub(states_elec_kpoints_distribution)
 
    if (.not. allocated(st%st_kpt_task)) then
      safe_allocate(st%st_kpt_task(0:st%st_kpt_mpi_grp%size-1, 1:4))
    end if
 
    st%st_kpt_task(0:st%st_kpt_mpi_grp%size-1, :) = 0
    st%st_kpt_task(st%st_kpt_mpi_grp%rank, :) = [st%st_start, st%st_end, st%d%kpt%start, st%d%kpt%end]
 
    if (st%parallel_in_states .or. st%d%kpt%parallel) then
      call comm_allreduce(st%st_kpt_mpi_grp, st%st_kpt_task)
    end if
 
    pop_sub(states_elec_kpoints_distribution)
  end subroutine states_elec_kpoints_distribution
 
  ! ---------------------------------------------------------
  !
  subroutine states_elec_choose_kpoints(st, kpoints, namespace)
    type(states_elec_t), target, intent(inout) :: st
    type(kpoints_t),             intent(in)    :: kpoints
    type(namespace_t),           intent(in)    :: namespace
 
    integer :: ik, iq
    type(states_elec_dim_t), pointer :: dd
 
 
    push_sub(states_elec_choose_kpoints)
 
    dd => st%d
 
    st%nik = kpoints_number(kpoints)
 
    if (dd%ispin == spin_polarized) st%nik = 2*st%nik
 
    safe_allocate(st%kweights(1:st%nik))
 
    do iq = 1, st%nik
      ik = dd%get_kpoint_index(iq)
      st%kweights(iq) = kpoints%get_weight(ik)
    end do
 
    if (debug%info) call print_kpoints_debug
    pop_sub(states_elec_choose_kpoints)
 
  contains
    subroutine print_kpoints_debug
      integer :: iunit
 
      push_sub(states_elec_choose_kpoints.print_kpoints_debug)
 
      call io_mkdir('debug/', namespace)
      iunit = io_open('debug/kpoints', namespace, action = 'write')
      call kpoints%write_info(iunit=iunit, absolute_coordinates = .true.)
      call io_close(iunit)
 
      pop_sub(states_elec_choose_kpoints.print_kpoints_debug)
    end subroutine print_kpoints_debug
 
  end subroutine states_elec_choose_kpoints
 
  !
  function states_elec_calculate_dipole(this, gr) result(dipole)
    class(states_elec_t), intent(in) :: this
    class(mesh_t),        intent(in) :: gr
    real(real64)                     :: dipole(1:gr%box%dim)
 
    integer :: ispin
    real(real64) :: e_dip(1:gr%box%dim, this%d%spin_channels)
 
    push_sub(states_elec_calculate_dipole)
 
    do ispin = 1, this%d%spin_channels
      call dmf_dipole(gr, this%rho(:, ispin), e_dip(:, ispin))
    end do
 
    dipole(:) = sum(e_dip(:, 1:this%d%spin_channels), 2)  ! dipole moment <mu_el> = \sum_i -e <x_i>
 
    pop_sub(states_elec_calculate_dipole)
  end function states_elec_calculate_dipole
 
 
#include "undef.F90"
#include "real.F90"
#include "states_elec_inc.F90"
 
#include "undef.F90"
#include "complex.F90"
#include "states_elec_inc.F90"
#include "undef.F90"
 
end module states_elec_oct_m
 
 
!! Local Variables:
!! mode: f90
!! coding: utf-8
!! End: