states_elec_restart.F90

Go to the documentation of this file.

!! Copyright (C) 2002-2016 M. Marques, A. Castro, A. Rubio, G. Bertsch, X. Andrade
!!
!! This program is free software; you can redistribute it and/or modify
!! it under the terms of the GNU General Public License as published by
!! the Free Software Foundation; either version 2, or (at your option)
!! any later version.
!!
!! This program is distributed in the hope that it will be useful,
!! but WITHOUT ANY WARRANTY; without even the implied warranty of
!! MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
!! GNU General Public License for more details.
!!
!! You should have received a copy of the GNU General Public License
!! along with this program; if not, write to the Free Software
!! Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
!! 02110-1301, USA.
!!
 
#include "global.h"
 
 
module states_elec_restart_oct_m
#ifdef HAVE_ADIOS2
  use adios2
#endif
  use batch_oct_m
  use calc_mode_par_oct_m
  use debug_oct_m
  use electron_space_oct_m
  use global_oct_m
  use io_function_oct_m
  use, intrinsic :: iso_fortran_env
  use kpoints_oct_m
  use lalg_basic_oct_m
  use linear_response_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 mpi_oct_m
  use multicomm_oct_m
  use namespace_oct_m
  use output_low_oct_m
  use parser_oct_m
  use par_vec_oct_m
  use profiling_oct_m
  use restart_oct_m
  use smear_oct_m
  use space_oct_m
  use states_abst_oct_m
  use states_elec_calc_oct_m
  use states_elec_oct_m
  use states_elec_dim_oct_m
  use string_oct_m
  use types_oct_m
  use wfs_elec_oct_m
 
  implicit none
 
 
  private
  public ::                              &
    states_elec_look_and_load,           &
    states_elec_dump,                    &
    states_elec_load,                    &
    states_elec_dump_rho,                &
    states_elec_load_rho,                &
    states_elec_dump_spin,               &
    states_elec_load_spin,               &
    states_elec_dump_frozen,             &
    states_elec_load_frozen,             &
    states_elec_read_user_def_orbitals,  &
    states_elec_transform
 
contains
 
  ! ---------------------------------------------------------
  subroutine states_elec_look_and_load(restart, namespace, space, st, mesh, kpoints, fixed_occ, is_complex, packed)
    type(restart_t),             intent(in)    :: restart
    type(namespace_t),           intent(in)    :: namespace
    class(space_t),              intent(in)    :: space
    type(states_elec_t), target, intent(inout) :: st
    class(mesh_t),               intent(in)    :: mesh
    type(kpoints_t),             intent(in)    :: kpoints
    logical,                     intent(in)    :: fixed_occ
    logical,           optional, intent(in)    :: is_complex
    logical,           optional, intent(in)    :: packed
 
    integer :: nkpt, dim, nst, ierr
    real(real64), allocatable :: new_occ(:,:)
 
    push_sub(states_elec_look_and_load)
 
    !check how many wfs we have
    call states_elec_look(restart, nkpt, dim, nst, ierr)
    if (ierr /= 0) then
      message(1) = "Unable to read states information."
      call messages_fatal(1, namespace=namespace)
    end if
 
    if (st%parallel_in_states) then
      message(1) = "Internal error: cannot use states_elec_look_and_load when parallel in states."
      call messages_fatal(1, namespace=namespace)
    end if
 
    ! Resize st%occ, retaining information there
    allocate(new_occ(1:nst, 1:st%nik))
    new_occ(:,:) = m_zero
    new_occ(1:min(nst, st%nst),:) = st%occ(1:min(nst, st%nst),:)
    safe_deallocate_a(st%occ)
    call move_alloc(new_occ, st%occ)
 
    ! FIXME: This wrong, one cannot just change the number of states
    ! without updating the internal structures, in the case of parallel in states.
    ! I think it is right now without state parallelism.
    st%nst      = nst
    st%st_start = 1
    st%st_end   = nst
    st%lnst     = nst
 
    safe_deallocate_a(st%node)
    safe_allocate(st%node(1:st%nst))
    st%node(:)  = 0
 
    safe_deallocate_a(st%eigenval)
    safe_allocate(st%eigenval(1:st%nst, 1:st%nik))
    st%eigenval = huge(st%eigenval)
 
    if (present(is_complex)) then
      if (is_complex) then
        call states_elec_allocate_wfns(st, mesh, type_cmplx, packed=optional_default(packed, .false.))
      else
        call states_elec_allocate_wfns(st, mesh, type_float, packed=optional_default(packed, .false.))
      end if
    else
      ! allow states_elec_allocate_wfns to decide for itself whether complex or real needed
      call states_elec_allocate_wfns(st, mesh, packed=optional_default(packed, .false.))
    end if
 
    if (st%d%ispin == spinors) then
      safe_allocate(st%spin(1:3, 1:st%nst, 1:st%nik))
      st%spin = m_zero
    end if
 
    ! load wavefunctions
    call states_elec_load(restart, namespace, space, st, mesh, kpoints, fixed_occ, ierr)
    if (ierr /= 0) then
      message(1) = "Unable to read wavefunctions."
      call messages_fatal(1, namespace=namespace)
    end if
 
    pop_sub(states_elec_look_and_load)
  end subroutine states_elec_look_and_load
 
 
  ! ---------------------------------------------------------
  subroutine states_elec_dump(restart, space, st, mesh, kpoints, ierr, iter, lr, verbose)
    type(restart_t),             intent(in)  :: restart
    class(space_t),              intent(in)  :: space
    type(states_elec_t), target, intent(in)  :: st
    class(mesh_t),               intent(in)  :: mesh
    type(kpoints_t),             intent(in)  :: kpoints
    integer,                     intent(out) :: ierr
    integer,           optional, intent(in)  :: iter
    type(lr_t),        optional, intent(in)  :: lr
    logical,           optional, intent(in)  :: verbose
 
    integer :: iunit_wfns, iunit_occs, iunit_states
    integer :: err, err2(2), ik, idir, ist, idim, itot
    integer :: root(1:P_STRATEGY_MAX)
    character(len=MAX_PATH_LEN) :: filename
    character(len=500) :: lines(3)
    logical :: lr_wfns_are_associated, should_write, verbose_
    real(real64)   :: kpoint(space%dim)
    real(real64),  allocatable :: dpsi(:)
    complex(real64),  allocatable :: zpsi(:)
    integer :: restart_file_format
 
    push_sub(states_elec_dump)
 
    verbose_ = optional_default(verbose, .true.)
 
    ierr = 0
 
    if (restart%skip()) then
      pop_sub(states_elec_dump)
      return
    end if
 
    if (verbose_) then
      message(1) = "Info: Writing states."
      call print_date(trim(message(1))//' ')
    end if
 
    call profiling_in("RESTART_WRITE")
 
    restart_file_format = restart%file_format_states
    if (present(lr)) then
      lr_wfns_are_associated = (allocated(lr%ddl_psi) .and. states_are_real(st)) .or. &
        (allocated(lr%zdl_psi) .and. states_are_complex(st))
      assert(lr_wfns_are_associated)
      ! always use obf format for linear response
      restart_file_format = option__restartfileformatstates__obf
    end if
 
    call restart_block_signals()
 
    iunit_states = restart%open('states')
    write(lines(1), '(a20,1i10)')  'nst=                ', st%nst
    write(lines(2), '(a20,1i10)')  'dim=                ', st%d%dim
    write(lines(3), '(a20,1i10)')  'nik=                ', st%nik
    call restart%write(iunit_states, lines, 3, err)
    if (err /= 0) ierr = ierr + 1
    call restart%close(iunit_states)
 
 
    iunit_wfns = restart%open('wfns')
    lines(1) = '#     #k-point            #st            #dim    filename'
    if (states_are_real(st)) then
      lines(2) = '%Real_Wavefunctions'
    else
      lines(2) = '%Complex_Wavefunctions'
    end if
    call restart%write(iunit_wfns, lines, 2, err)
    if (err /= 0) ierr = ierr + 2
 
 
    iunit_occs = restart%open('occs')
    lines(1) = '# occupations | eigenvalue[a.u.] | Im(eigenvalue) [a.u.] | k-points | k-weights | filename | ik | ist | idim'
    lines(2) = '%Occupations_Eigenvalues_K-Points'
    call restart%write(iunit_occs, lines, 2, err)
    if (err /= 0) ierr = ierr + 4
 
 
    if (states_are_real(st)) then
      safe_allocate(dpsi(1:mesh%np))
    else
      safe_allocate(zpsi(1:mesh%np))
    end if
 
    itot = 1
    root = -1
    err2 = 0
    do ik = 1, st%nik
      kpoint(1:space%dim) = &
        kpoints%get_point(st%d%get_kpoint_index(ik), absolute_coordinates = .true.)
 
      do ist = 1, st%nst
        do idim = 1, st%d%dim
          root(p_strategy_domains) = mod(itot - 1, mesh%mpi_grp%size)
          write(filename,'(i10.10)') itot
 
          write(lines(1), '(i8,a,i8,a,i8,3a)') ik, ' | ', ist, ' | ', idim, ' | "', trim(filename), '"'
          call restart%write(iunit_wfns, lines, 1, err)
          if (err /= 0) err2(1) = err2(1) + 1
 
          write(lines(1), '(e23.16,a,e23.16)') st%occ(ist,ik), ' | ', st%eigenval(ist, ik)
          write(lines(1), '(a,a,e23.16)') trim(lines(1)), ' | ', m_zero
          do idir = 1, space%dim
            write(lines(1), '(a,a,e23.16)') trim(lines(1)), ' | ', kpoint(idir)
          end do
          write(lines(1), '(a,a,e23.16,a,i10.10,3(a,i8))') trim(lines(1)), &
            ' | ', st%kweights(ik), ' | ', itot, ' | ', ik, ' | ', ist, ' | ', idim
          call restart%write(iunit_occs, lines, 1, err)
          if (err /= 0) err2(1) = err2(1) + 1
 
          should_write = st%st_start <= ist .and. ist <= st%st_end
 
          if (restart_file_format == option__restartfileformatstates__obf) then
            if (should_write) then
              if (st%d%kpt%start <= ik .and. ik <= st%d%kpt%end) then
                if (states_are_real(st)) then
                  if (.not. present(lr)) then
                    call states_elec_get_state(st, mesh, idim, ist, ik, dpsi)
                    call restart%write_mesh_function(filename, mesh, dpsi, err, root = root)
                  else
                    call restart%write_mesh_function(filename, mesh, &
                      lr%ddl_psi(:, idim, ist, ik), err, root = root)
                  end if
                else
                  if (.not. present(lr)) then
                    call states_elec_get_state(st, mesh, idim, ist, ik, zpsi)
                    call restart%write_mesh_function(filename, mesh, zpsi, err, root = root)
                  else
                    call restart%write_mesh_function(filename, mesh, &
                      lr%zdl_psi(:, idim, ist, ik), err, root = root)
                  end if
                end if
              else
                err = 0
              end if
            end if
            if (err /= 0) err2(2) = err2(2) + 1
          end if
 
          itot = itot + 1
        end do ! st%d%dim
      end do ! st%nst
    end do ! st%nik
    if (err2(1) /= 0) ierr = ierr + 8
    if (err2(2) /= 0) ierr = ierr + 16
 
    safe_deallocate_a(dpsi)
    safe_deallocate_a(zpsi)
 
    if (restart_file_format == option__restartfileformatstates__adios2) then
      if (present(lr)) then
        ! this should never occur because we force obf format for linear response
        call messages_not_implemented("ADIOS2 restart file format with linear response.")
      end if
      call states_elec_dump_adios2(restart, st, mesh, ierr)
    end if
 
    lines(1) = '%'
    call restart%write(iunit_occs, lines, 1, err)
    if (err /= 0) ierr = ierr + 32
    call restart%write(iunit_wfns, lines, 1, err)
    if (err /= 0) ierr = ierr + 64
    if (present(iter)) then
      write(lines(1),'(a,i7)') 'Iter = ', iter
      call restart%write(iunit_wfns, lines, 1, err)
      if (err /= 0) ierr = ierr + 128
    end if
 
    call restart%close(iunit_wfns)
    call restart%close(iunit_occs)
 
    if (verbose_) then
      message(1) = "Info: Finished writing states."
      call print_date(trim(message(1))//' ')
    end if
 
    call restart_unblock_signals()
 
    call profiling_out("RESTART_WRITE")
    pop_sub(states_elec_dump)
  end subroutine states_elec_dump
 
  ! ---------------------------------------------------------
  subroutine states_elec_dump_adios2(restart, st, mesh, ierr)
    type(restart_t),             intent(in)  :: restart
    type(states_elec_t), target, intent(in)  :: st
    class(mesh_t),               intent(in)  :: mesh
    integer,                     intent(out) :: ierr
 
#ifdef HAVE_ADIOS2
    type(adios2_adios) :: adios
    type(adios2_io) :: io
    type(adios2_variable) :: var, var_indices
    type(adios2_attribute) :: attribute
    type(adios2_engine) :: engine
    integer :: adios2_type, ik, ib, ip, adios2_mode
    integer(int64), allocatable :: global_indices(:)
    class(wfs_elec_t), pointer :: psib
#endif
    push_sub(states_elec_dump_adios2)
 
#ifdef HAVE_ADIOS2
#ifdef HAVE_MPI
    call adios2_init(adios, restart%mpi_grp%comm%MPI_VAL, ierr)
#else
    call adios2_init(adios, ierr)
#endif
    call check_error(.not. adios%valid .or. ierr /= adios2_error_none, "Problem initializing ADIOS2.")
    call adios2_declare_io(io, adios, "writer", ierr)
    call check_error(.not. io%valid .or. ierr /= adios2_error_none, "Problem initializing ADIOS2.")
    ! do not compute min and max
    call adios2_set_parameter(io, "StatsLevel", "0", ierr)
    call check_error(ierr /= adios2_error_none, "Problem setting ADIOS2 parameter.")
    ! open restart file
    call adios2_open(engine, io, trim(restart%dir())//"/"//"wavefunctions.bp", adios2_mode_write, ierr)
    call check_error(.not. engine%valid .or. ierr /= adios2_error_none, "Problem opening ADIOS2 restart file.")
    if (states_are_real(st)) then
      adios2_type = adios2_type_dp
    else
      adios2_type = adios2_type_complex_dp
    end if
    ! define one variable for all states
    call adios2_define_variable(var, io, "wavefunctions", adios2_type, 3, &
      [int(st%nst * st%d%dim, int64), int(mesh%np_global, int64), int(st%nik, int64)], & ! global size of wavefunctions
      [0_int64, 0_int64, 0_int64], & ! local selection of the variable, overwritten by set_selection call below
      [1_int64, 1_int64, 1_int64], &
      adios2_variable_dims, ierr)
    call check_error(.not. var%valid .or. ierr /= adios2_error_none, "Problem creating ADIOS2 variable (wavefunctions).")
    ! save ParDomains in an attribute
    call adios2_define_attribute(attribute, io, "ParDomains", mesh%mpi_grp%size, ierr)
    call check_error(.not. attribute%valid .or. ierr /= adios2_error_none, "Problem creating ADIOS2 attribute (ParDomains).")
    ! define one variable for all states
    call adios2_define_variable(var_indices, io, "global_indices", adios2_type_integer8, 1, &
      [int(mesh%np_global, int64)], & ! global size of array
      [int(mesh%pv%xlocal-1, int64)], & ! local offset
      [int(mesh%np, int64)], & ! local size
      adios2_variable_dims, ierr)
    call check_error(.not. var_indices%valid .or. ierr /= adios2_error_none, "Problem creating ADIOS2 variable (global_indices).")
    ! save grid indices
    safe_allocate(global_indices(mesh%np))
    do ip = 1, mesh%np
      global_indices(ip) = par_vec_local2global(mesh%pv, ip)
    end do
 
    call adios2_begin_step(engine, ierr)
    call adios2_put(engine, var_indices, global_indices, ierr)
 
    do ik = st%d%kpt%start, st%d%kpt%end
      do ib = st%group%block_start, st%group%block_end
        ! make sure we have a batch in BATCH_PACKED status
        select case (st%group%psib(ib, ik)%status())
        case (batch_not_packed)
          safe_allocate(psib)
          call st%group%psib(ib, ik)%copy_to(psib)
          ! transform the unpacked batch to a packed batch on the CPU
          call psib%do_pack(batch_packed)
          ! we need to use sync mode in this case because we use a temporary batch
          adios2_mode = adios2_mode_sync
        case (batch_packed)
          psib => st%group%psib(ib, ik)
          adios2_mode = adios2_mode_deferred
        case (batch_device_packed)
          safe_allocate(psib)
          call st%group%psib(ib, ik)%copy_to(psib, copy_data=.true.)
          ! copy to CPU
          call psib%do_unpack(force=.true.)
          ! make sure the batch is in packed state on the CPU
          call psib%do_pack(batch_packed)
          ! we need to use sync mode in this case because we use a temporary batch
          adios2_mode = adios2_mode_sync
        end select
        assert(psib%status() == batch_packed)
 
        ! select part of variable to write, correspond to the memory of the current batch
        ! first argument gives the start in the global array (-1 needed for C indexing in ADIOS2)
        ! second argument gives the size of the local batch
        call adios2_set_selection(var, 3, &
          [int((states_elec_block_min(st, ib)-1)*st%d%dim, int64), int(mesh%pv%xlocal-1, int64), int(ik-1, int64)], &
          [int(st%group%psib(ib, ik)%nst_linear, int64), int(mesh%np, int64), 1_int64], ierr)
        ! set memory layout of one batch, take into account padding and ghost/boundary points
        if (states_are_real(st)) then
          call adios2_set_memory_selection(var, 3, [0_int64, 0_int64, 0_int64], &
            [int(size(psib%dff_pack, 1), int64), &
            int(size(psib%dff_pack, 2), int64), 1_int64], ierr)
          call adios2_put(engine, var, psib%dff_pack, adios2_mode, ierr)
        else
          call adios2_set_memory_selection(var, 3, [0_int64, 0_int64, 0_int64], &
            [int(size(psib%zff_pack, 1), int64), &
            int(size(psib%zff_pack, 2), int64), 1_int64], ierr)
          call adios2_put(engine, var, psib%zff_pack, adios2_mode, ierr)
        end if
        select case (st%group%psib(ib, ik)%status())
        case (batch_not_packed, batch_device_packed)
          call psib%end()
          safe_deallocate_p(psib)
        case (batch_packed)
          nullify(psib)
        end select
      end do
    end do
    call adios2_end_step(engine, ierr)
 
    ! deallocate only after end_step because only now everything has been written
    safe_deallocate_a(global_indices)
 
    ! close restart file
    call adios2_close(engine, ierr)
    call check_error(ierr /= adios2_error_none, "Problem closing ADIOS2 engine.")
    call adios2_finalize(adios, ierr)
    call check_error(ierr /= adios2_error_none, "Problem finalizing ADIOS2.")
#else
    ! avoid empty routine
    ierr = 0
#endif
 
    pop_sub(states_elec_dump_adios2)
  end subroutine states_elec_dump_adios2
 
 
  ! ---------------------------------------------------------
  subroutine states_elec_load(restart, namespace, space, st, mesh, kpoints, fixed_occ, &
    ierr, iter, lr, lowest_missing, label, verbose, skip)
    type(restart_t),             intent(in)    :: restart
    type(namespace_t),           intent(in)    :: namespace
    class(space_t),              intent(in)    :: space
    type(states_elec_t), target, intent(inout) :: st
    class(mesh_t),               intent(in)    :: mesh
    type(kpoints_t),             intent(in)    :: kpoints
    logical,                     intent(in)    :: fixed_occ
    integer,                     intent(out)   :: ierr
    integer,           optional, intent(out)   :: iter
    type(lr_t),        optional, intent(inout) :: lr
    integer,           optional, intent(out)   :: lowest_missing(:, :)
    character(len=*),  optional, intent(in)    :: label
    logical,           optional, intent(in)    :: verbose
    logical,           optional, intent(in)    :: skip(:)
 
    integer              :: states_elec_file, wfns_file, occ_file, err, ik, ist, idir, idim
    integer              :: idone, iread, ntodo
    character(len=12)    :: filename
    character(len=1)     :: char
    logical, allocatable :: filled(:, :, :)
    character(len=256)   :: lines(3), label_
    character(len=50)    :: str
 
    real(real64)         :: my_occ, imev, my_kweight
    logical              :: read_occ, lr_allocated, verbose_
    logical              :: integral_occs
    real(real64), allocatable   :: dpsi(:)
    complex(real64), allocatable   :: zpsi(:)
    character(len=256), allocatable :: restart_file(:, :, :)
    logical,            allocatable :: restart_file_present(:, :, :)
    real(real64)         :: kpoint(space%dim), read_kpoint(space%dim)
 
    integer              :: iread_tmp
    integer, allocatable :: lowest_missing_tmp(:, :)
    integer :: restart_file_format
 
    push_sub(states_elec_load)
 
    ierr = 0
 
    ! make sure these intent(out)`s are initialized no matter what
    if (present(lowest_missing)) lowest_missing = 1
    if (present(iter)) iter = 0
 
    if (present(skip)) then
      assert(ubound(skip, dim = 1) == st%nst)
    end if
 
    if (restart%skip()) then
      ierr = -1
      pop_sub(states_elec_load)
      return
    end if
 
    call profiling_in("RESTART_READ")
 
    verbose_ = optional_default(verbose, .true.)
 
    if (present(label)) then
      label_ = trim(label)
    else
      if (present(lr)) then
        label_ = " for linear response"
      else
        label_ = ""
      end if
    end if
 
    message(1) = 'Info: Reading states'
    if (len(trim(label_)) > 0) then
      message(1) = trim(message(1)) // trim(label_)
    end if
    message(1) = trim(message(1)) // "."
    if (verbose_) call print_date(trim(message(1))//' ')
 
    if (.not. present(lr)) then
      st%fromScratch = .false. ! obviously, we are using restart info
    end if
 
    ! If one restarts a GS calculation changing the %Occupations block, one
    ! cannot read the occupations, otherwise these overwrite the ones from
    ! the input file. restart_fixed_occ makes that we do use the ones in the file.
    integral_occs = .true. ! only used if restart_fixed_occ
    if (fixed_occ) then
      read_occ = .true.
      st%fixed_occ = .true.
    else
      read_occ = .not. st%fixed_occ
    end if
 
    if (.not. present(lr)) then
      st%eigenval(:, :) = m_zero
      ! to be filled in from reading afterward
    end if
 
    if (.not. present(lr) .and. read_occ) then
      st%occ(:, :) = m_zero
      ! to be filled in from reading afterward
    end if
 
    restart_file_format = restart%file_format_states
    ! sanity check
    if (present(lr)) then
      lr_allocated = (allocated(lr%ddl_psi) .and. states_are_real(st)) .or. &
        (allocated(lr%zdl_psi) .and. states_are_complex(st))
      assert(lr_allocated)
      ! always use obf format for linear response
      restart_file_format = option__restartfileformatstates__obf
    end if
 
    states_elec_file  = restart%open('states')
    ! sanity check on spin/k-points. Example file 'states':
    ! nst=                         2
    ! dim=                         1
    ! nik=                         2
    call restart%read(states_elec_file, lines, 3, err)
    if (err /= 0) then
      ierr = ierr - 2
    else
      read(lines(2), *) str, idim
      read(lines(3), *) str, ik
      if (idim == 2 .and. st%d%dim == 1) then
        write(message(1),'(a)') 'Incompatible restart information: saved calculation is spinors, this one is not.'
        call messages_warning(1, namespace=namespace)
        ierr = ierr - 2**2
      end if
      if (idim == 1 .and. st%d%dim == 2) then
        write(message(1),'(a)') 'Incompatible restart information: this calculation is spinors, saved one is not.'
        call messages_warning(1, namespace=namespace)
        ierr = ierr - 2**3
      end if
      if (ik < st%nik) then
        write(message(1),'(a)') 'Incompatible restart information: not enough k-points.'
        write(message(2),'(2(a,i6))') 'Expected ', st%nik, ' > Read ', ik
        call messages_warning(2, namespace=namespace)
      end if
      ! We will check that they are the right k-points later, so we do not need to put a specific error here.
    end if
    call restart%close(states_elec_file)
 
 
    ! open files to read
    wfns_file = restart%open('wfns')
    occ_file = restart%open('occs')
    call restart%read(wfns_file, lines, 2, err)
    if (err /= 0) then
      ierr = ierr - 2**5
    else if (states_are_real(st)) then
      read(lines(2), '(a)') str
      if (str(2:8) == 'Complex') then
        message(1) = "Cannot read real states from complex wavefunctions."
        call messages_warning(1, namespace=namespace)
        ierr = ierr - 2**6
      else if (str(2:5) /= 'Real') then
        message(1) = "Restart file 'wfns' does not specify real/complex; cannot check compatibility."
        call messages_warning(1, namespace=namespace)
      end if
    end if
    ! complex can be restarted from real, so there is no problem.
 
    ! Skip two lines.
    call restart%read(occ_file, lines, 2, err)
    if (err /= 0) ierr = ierr - 2**7
 
    ! If any error occured up to this point then it is not worth continuing,
    ! as there something fundamentally wrong with the restart files
    if (ierr /= 0) then
      call restart%close(wfns_file)
      call restart%close(occ_file)
      call profiling_out("RESTART_READ")
      pop_sub(states_elec_load)
      return
    end if
 
    if (states_are_real(st)) then
      safe_allocate(dpsi(1:mesh%np))
    else
      safe_allocate(zpsi(1:mesh%np))
    end if
 
    safe_allocate(restart_file(1:st%d%dim, st%st_start:st%st_end, 1:st%nik))
    safe_allocate(restart_file_present(1:st%d%dim, st%st_start:st%st_end, 1:st%nik))
    restart_file_present = .false.
 
    ! Next we read the list of states from the files.
    ! Errors in reading the information of a specific state from the files are ignored
    ! at this point, because later we will skip reading the wavefunction of that state.
    do
      call restart%read(wfns_file, lines, 1, err)
      if (err == 0) then
        read(lines(1), '(a)') char
        if (char == '%') then
          !We reached the end of the file
          exit
        else
          read(lines(1), *) ik, char, ist, char, idim, char, filename
        end if
      end if
 
      if (err /= 0 .or. index_is_wrong()) then
        call restart%read(occ_file, lines, 1, err)
        cycle
      end if
 
      if (ist >= st%st_start .and. ist <= st%st_end .and. &
        ik >= st%d%kpt%start .and. ik <= st%d%kpt%end) then
 
        restart_file(idim, ist, ik) = trim(filename)
        restart_file_present(idim, ist, ik) = .true.
      end if
 
      call restart%read(occ_file, lines, 1, err)
      if (.not. present(lr)) then ! do not read eigenvalues or occupations when reading linear response
        ! # occupations | eigenvalue[a.u.] | Im(eigenvalue) [a.u.] | k-points | k-weights | filename | ik | ist | idim
 
        if (err == 0) then
          read(lines(1), *) my_occ, char, st%eigenval(ist, ik), char, imev, char, &
            (read_kpoint(idir), char, idir = 1, space%dim), my_kweight
          ! we do not want to read the k-weights, we have already set them appropriately
        else
          ! There is a problem with this states information, so we skip it.
          if (ist >= st%st_start .and. ist <= st%st_end .and. &
            ik >= st%d%kpt%start .and. ik <= st%d%kpt%end) then
            restart_file_present(idim, ist, ik) = .false.
          end if
          cycle
        end if
 
        kpoint(1:space%dim) = &
          kpoints%get_point(st%d%get_kpoint_index(ik), absolute_coordinates = .true.)
        ! FIXME: maybe should ignore ik and just try to match actual vector k-points?
        if (any(abs(kpoint(1:space%dim) - read_kpoint(1:space%dim)) > 1e-12_real64)) then
          ! write only once for each k-point so as not to be too verbose
          if (ist == 1) then
            write(message(1),'(a,i6)') 'Incompatible restart information: k-point mismatch for ik ', ik
            write(message(2),'(a,99f18.12)') '  Expected : ', kpoint(1:space%dim)
            write(message(3),'(a,99f18.12)') '  Read     : ', read_kpoint(1:space%dim)
            call messages_warning(3, namespace=namespace)
          end if
          if (ist >= st%st_start .and. ist <= st%st_end .and. &
            ik >= st%d%kpt%start .and. ik <= st%d%kpt%end) then
            restart_file_present(idim, ist, ik) = .false.
          end if
        end if
 
        if (read_occ) then
          st%occ(ist, ik) = my_occ
          integral_occs = integral_occs .and. &
            abs((st%occ(ist, ik) - st%smear%el_per_state) * st%occ(ist, ik))  <=  m_epsilon
        end if
      end if
    end do
 
    if (present(iter)) then
      call restart%read(wfns_file, lines, 1, err)
      if (err /= 0) then
        ierr = ierr - 2**8
      else
        read(lines(1), *) filename, filename, iter
      end if
    end if
 
    call restart%close(wfns_file)
    call restart%close(occ_file)
 
    ! At this point, the restart file should be available
    ! if not, we fall back to the obf format
    if (restart_file_format == option__restartfileformatstates__adios2) then
      if (.not.loct_dir_exists(trim(restart%dir())//"/"//"wavefunctions.bp")) then
        message(1) = "ADIOS2 restart file not found, falling back to obf format."
        call messages_warning(1, namespace=namespace)
        restart_file_format = option__restartfileformatstates__obf
      end if
    end if
 
    ! Now we read the wavefunctions. At this point we need to have all the information from the
    ! states, occs, and wfns files in order to avoid serialisation of reads, as restart_read
    ! works as a barrier.
    safe_allocate(filled(1:st%d%dim, st%st_start:st%st_end, st%d%kpt%start:st%d%kpt%end))
    filled = .false.
 
    if (present(lowest_missing)) lowest_missing = st%nst + 1
 
    iread = 0
    if (st%system_grp%is_root() .and. verbose_) then
      idone = 0
      ntodo = st%lnst*st%d%kpt%nlocal*st%d%dim
      call loct_progress_bar(-1, ntodo)
    end if
    if (restart_file_format == option__restartfileformatstates__obf) then
      do ik = st%d%kpt%start, st%d%kpt%end
        do ist = st%st_start, st%st_end
          if (present(skip)) then
            if (skip(ist)) cycle
          end if
 
          do idim = 1, st%d%dim
 
            if (.not. restart_file_present(idim, ist, ik)) then
              if (present(lowest_missing)) then
                lowest_missing(idim, ik) = min(lowest_missing(idim, ik), ist)
              end if
              cycle
            end if
 
            if (states_are_real(st)) then
              call restart%read_mesh_function(restart_file(idim, ist, ik), mesh, dpsi, err)
              if (.not. present(lr)) then
                call states_elec_set_state(st, mesh, idim, ist, ik, dpsi)
              else
                call lalg_copy(mesh%np, dpsi, lr%ddl_psi(:, idim, ist, ik))
              end if
            else
              call restart%read_mesh_function(restart_file(idim, ist, ik), mesh, zpsi, err)
              if (.not. present(lr)) then
                call states_elec_set_state(st, mesh, idim, ist, ik, zpsi)
              else
                call lalg_copy(mesh%np, zpsi, lr%zdl_psi(:, idim, ist, ik))
              end if
            end if
 
            if (err == 0) then
              filled(idim, ist, ik) = .true.
              iread = iread + 1
            else if (present(lowest_missing)) then
              lowest_missing(idim, ik) = min(lowest_missing(idim, ik), ist)
            end if
 
            if (st%system_grp%is_root() .and. verbose_) then
              idone = idone + 1
              call loct_progress_bar(idone, ntodo)
            end if
 
          end do
        end do
      end do
    else if (restart_file_format == option__restartfileformatstates__adios2) then
      if (present(lr)) then
        ! this should never occur because we force obf format for linear response
        call messages_not_implemented("ADIOS2 restart file format with linear response.")
      end if
      if (restart%has_map()) then
        message(1) = "ADIOS2 restart file format does not support restarting from a different grid."
        message(2) = "Please set RestartFileFormatStates = obf."
        call messages_fatal(2, namespace=namespace)
      end if
      call states_elec_load_adios2(restart, st, mesh, iread, filled, restart_file_present, ierr, skip, lowest_missing)
    end if
 
    safe_deallocate_a(dpsi)
    safe_deallocate_a(zpsi)
    safe_deallocate_a(restart_file)
    safe_deallocate_a(restart_file_present)
 
    if (st%system_grp%is_root() .and. verbose_) then
      call messages_new_line()
    end if
 
    if (st%parallel_in_states .or. st%d%kpt%parallel) then
      iread_tmp = iread
      call st%st_kpt_mpi_grp%allreduce(iread_tmp, iread, 1, mpi_integer, mpi_sum)
    end if
 
    if (st%d%kpt%parallel) then
      ! make sure all tasks know lowest_missing over all k-points
      if (present(lowest_missing)) then
        safe_allocate(lowest_missing_tmp(1:st%d%dim, 1:st%nik))
        lowest_missing_tmp = lowest_missing
        call st%d%kpt%mpi_grp%allreduce(lowest_missing_tmp(1,1), lowest_missing(1,1), st%d%dim*st%nik, &
          mpi_integer, mpi_min)
        safe_deallocate_a(lowest_missing_tmp)
      end if
    end if
 
    if (fixed_occ .and. iread == st%nst * st%nik * st%d%dim) then
      ! reset to overwrite whatever smearing may have been set earlier
      ! only do this if all states could be loaded from restart files
      call smear_init(st%smear, namespace, st%d%ispin, fixed_occ = .true., integral_occs = integral_occs, kpoints = kpoints)
    end if
 
    if (.not. present(lr) .and. .not. present(skip)) call fill_random()
    ! it is better to initialize lr wfns to zero
 
    safe_deallocate_a(filled)
 
    if (ierr == 0 .and. iread /= st%nst * st%nik * st%d%dim) then
      if (iread > 0) then
        ierr = iread
      else
        ierr = -1
      end if
      ! otherwise ierr = 0 would mean either all was read correctly, or nothing at all was read!
 
      if (.not. present(lr)) then
        write(str, '(a,i5)') 'Reading states.'
      else
        write(str, '(a,i5)') 'Reading states information for linear response.'
      end if
      call messages_print_with_emphasis(msg=trim(str), namespace=namespace)
      if (.not. present(skip)) then
        write(message(1),'(a,i6,a,i6,a)') 'Only ', iread,' files out of ', &
          st%nst * st%nik * st%d%dim, ' could be read.'
      else
        write(message(1),'(a,i6,a,i6,a)') 'Only ', iread,' files out of ', &
          st%nst * st%nik * st%d%dim, ' were loaded.'
        ierr = 0
      end if
      call messages_info(1, namespace=namespace)
      call messages_print_with_emphasis(namespace=namespace)
    end if
 
    message(1) = 'Info: States reading done.'
    if (verbose_) call print_date(trim(message(1))//' ')
 
    call profiling_out("RESTART_READ")
    pop_sub(states_elec_load)
 
  contains
 
    ! ---------------------------------------------------------
    subroutine fill_random() !< put random function in orbitals that could not be read.
      push_sub(states_elec_load.fill_random)
 
      do ik = st%d%kpt%start, st%d%kpt%end
 
        do ist = st%st_start, st%st_end
          do idim = 1, st%d%dim
            if (filled(idim, ist, ik)) cycle
 
            call states_elec_generate_random(st, mesh, kpoints, ist, ist, ik, ik)
          end do
        end do
      end do
 
      pop_sub(states_elec_load.fill_random)
    end subroutine fill_random
 
    ! ---------------------------------------------------------
 
    logical function index_is_wrong() !< .true. if the index (idim, ist, ik) is not present in st structure...
      push_sub(states_elec_load.index_is_wrong)
 
      if (idim > st%d%dim .or. idim < 1 .or.   &
        ist   > st%nst   .or. ist  < 1 .or.   &
        ik    > st%nik .or. ik   < 1) then
        index_is_wrong = .true.
      else
        index_is_wrong = .false.
      end if
 
      pop_sub(states_elec_load.index_is_wrong)
    end function index_is_wrong
 
  end subroutine states_elec_load
 
  ! ---------------------------------------------------------
  subroutine states_elec_load_adios2(restart, st, mesh, iread, filled, restart_file_present, ierr, skip, lowest_missing)
    type(restart_t),             intent(in)    :: restart
    type(states_elec_t), target, intent(inout) :: st
    class(mesh_t),               intent(in)    :: mesh
    integer,                     intent(out)   :: iread
    logical,                     intent(out)   :: filled(1:, st%st_start:, st%d%kpt%start:)
    logical,                     intent(in)    :: restart_file_present(:, :, :)
    integer,                     intent(out)   :: ierr
    logical,           optional, intent(in)    :: skip(:)
    integer,           optional, intent(inout) :: lowest_missing(:, :)
 
#ifdef HAVE_ADIOS2
    type(adios2_adios) :: adios
    type(adios2_io) :: io
    type(adios2_variable) :: var, var_indices
    type(adios2_attribute) :: attribute
    type(adios2_engine) :: engine
    integer :: type_file, type_requested, ib, ik, ndims, pardomains
    integer :: nst_max
    integer(int64), allocatable :: var_shape(:)
    integer(int64), allocatable :: global_indices(:)
#endif
 
    push_sub(states_elec_load_adios2)
#ifdef HAVE_ADIOS2
#ifdef HAVE_MPI
    call adios2_init(adios, restart%mpi_grp%comm%MPI_VAL, ierr)
#else
    call adios2_init(adios, ierr)
#endif
    call check_error(.not. adios%valid .or. ierr /= adios2_error_none, "Problem initializing ADIOS2.")
    call adios2_declare_io(io, adios, "reader", ierr)
    call check_error(.not. io%valid .or. ierr /= adios2_error_none, "Problem initializing ADIOS2.")
    ! open restart file
    call adios2_open(engine, io, trim(restart%dir())//"/"//"wavefunctions.bp", adios2_mode_read, ierr)
    call check_error(.not. engine%valid .or. ierr /= adios2_error_none, "Problem opening ADIOS2 restart file.")
 
    call adios2_begin_step(engine, ierr)
    ! get ParDomains from attribute
    call adios2_inquire_attribute(attribute, io, "ParDomains", ierr)
    call check_error(.not. attribute%valid .or. ierr /= adios2_error_none, "Problem inquiring ADIOS2 attribute.")
    call adios2_attribute_data(pardomains, attribute, ierr)
    if (states_are_real(st)) then
      type_requested = adios2_type_dp
    else
      type_requested = adios2_type_complex_dp
    end if
    call adios2_inquire_variable(var, io, 'wavefunctions', ierr)
    call check_error(.not. var%valid .or. ierr /= adios2_error_none, "Problem loading ADIOS2 variable (wavefunctions).")
    ! get variable type and shape
    call adios2_variable_type(type_file, var, ierr)
    call adios2_variable_ndims(ndims, var, ierr)
    call adios2_variable_shape(var_shape, ndims, var, ierr)
    assert(ndims == 3)
    nst_max = st%nst
    if (var_shape(1) < int(st%nst * st%d%dim, int64)) then
      ! more states requested than available in the restart data, only read up to nst_max
      nst_max = int(var_shape(1), int32) / st%d%dim
    end if
    if (var_shape(2) /= int(mesh%np_global, int64)) then
      ! this should never happen because we check above if there is a map which is only allocated if the grid is different
      message(1) = "Error: trying to restart with a different number of grid points. Not supported with ADIOS2 format."
      call messages_fatal(1)
    end if
    if (present(skip)) then
      ! read state by state if skip parameter present, otherwise batch by batch which is much better
      ! for the performance
      call states_elec_adios2_read_state_by_state(st, mesh, engine, var, type_requested, type_file, &
        filled, restart_file_present, iread, skip, lowest_missing)
    else
      call states_elec_adios2_read_batch_by_batch(st, mesh, engine, var, var_shape, type_requested, type_file, &
        nst_max, filled, iread)
    end if
    ! communicate points for mesh parallelization
    if (mesh%mpi_grp%size > 1 .or. pardomains > 1) then
      call adios2_inquire_variable(var_indices, io, 'global_indices', ierr)
      call check_error(.not. var_indices%valid .or. ierr /= adios2_error_none, &
        "Problem loading ADIOS2 variable (global_indices).")
      safe_allocate(global_indices(mesh%np))
      call adios2_set_selection(var_indices, 1, &
        [int(mesh%pv%xlocal-1, int64)], &
        [int(mesh%np, int64)], ierr)
      call check_error(ierr /= adios2_error_none, "Problem setting ADIOS2 selection (global_indices).")
      call adios2_get(engine, var_indices, global_indices, ierr)
      call check_error(ierr /= adios2_error_none, "Problem in ADIOS2 get (global_indices).")
    end if
    call adios2_end_step(engine, ierr)
    call check_error(ierr /= adios2_error_none, "Problem in ADIOS2 end_step.")
    if (mesh%mpi_grp%size > 1 .or. pardomains > 1) then
      do ik = st%d%kpt%start, st%d%kpt%end
        do ib = st%group%block_start, st%group%block_end
          if (states_are_real(st)) then
            call dmesh_batch_exchange_points(mesh, st%group%psib(ib, ik), forward_map=global_indices)
          else
            call zmesh_batch_exchange_points(mesh, st%group%psib(ib, ik), forward_map=global_indices)
          end if
        end do
      end do
      safe_deallocate_a(global_indices)
    end if
 
    ! close restart file
    call adios2_close(engine, ierr)
    call check_error(ierr /= adios2_error_none, "Problem closing ADIOS2 engine.")
    call adios2_finalize(adios, ierr)
    call check_error(ierr /= adios2_error_none, "Problem finalizing ADIOS2.")
#else
    ! avoid an empty routine
    ierr = 0
#endif
    pop_sub(states_elec_load_adios2)
  end subroutine states_elec_load_adios2
 
 
#ifdef HAVE_ADIOS2
  ! ---------------------------------------------------------
  subroutine states_elec_adios2_read_state_by_state(st, mesh, engine, var, type_requested, type_file, &
    filled, restart_file_present, iread, skip, lowest_missing)
    type(states_elec_t), target, intent(inout) :: st
    class(mesh_t),               intent(in)    :: mesh
    type(adios2_engine),         intent(inout) :: engine
    type(adios2_variable),       intent(inout) :: var
    integer,                     intent(in)    :: type_requested
    integer,                     intent(in)    :: type_file
    logical,                     intent(inout) :: filled(1:, st%st_start:, st%d%kpt%start:)
    logical,                     intent(in)    :: restart_file_present(:, :, :)
    integer,                     intent(inout) :: iread
    logical,           optional, intent(in)    :: skip(:)
    integer,           optional, intent(inout) :: lowest_missing(:, :)
 
    integer :: ik, ist, idim, ierr
    real(real64), allocatable :: dpsi(:)
    complex(real64), allocatable :: zpsi(:)
 
    push_sub(states_elec_adios2_read_state_by_state)
 
    if (type_requested == adios2_type_dp .and. type_file == adios2_type_complex_dp) then
      message(1) = "Error: calculation requested real states for restarting, but restart file has complex states."
      call messages_fatal(1)
    end if
 
    do ik = st%d%kpt%start, st%d%kpt%end
      ! the buffer handed to adios2_get must match the type stored in the file,
      ! which is not necessarily the type of the states we are reading into
      if (type_file == adios2_type_dp) then
        safe_allocate(dpsi(1:mesh%np))
      else
        safe_allocate(zpsi(1:mesh%np))
      end if
      do ist = st%st_start, st%st_end
        if (present(skip)) then
          if (skip(ist)) cycle
        end if
 
        do idim = 1, st%d%dim
          if (.not. restart_file_present(idim, ist, ik)) then
            if (present(lowest_missing)) then
              lowest_missing(idim, ik) = min(lowest_missing(idim, ik), ist)
            end if
            cycle
          end if
 
          ! read only one state
          call adios2_set_selection(var, 3, &
            [int((ist-1)*st%d%dim+idim-1, int64), int(mesh%pv%xlocal-1, int64), int(ik-1, int64)], &
            [1_int64, int(mesh%np, int64), 1_int64], ierr)
          ! use sync mode here to immediately read into dpsi/zpsi
          if (type_file == adios2_type_dp) then
            call adios2_get(engine, var, dpsi, adios2_mode_sync, ierr)
            call check_error(ierr /= adios2_error_none, "Problem in ADIOS2 get dpsi.")
            if (type_requested == adios2_type_complex_dp) then
              ! real restart data read into complex states
              call states_elec_set_state(st, mesh, idim, ist, ik, cmplx(dpsi, m_zero, real64))
            else
              call states_elec_set_state(st, mesh, idim, ist, ik, dpsi)
            end if
          else
            call adios2_get(engine, var, zpsi, adios2_mode_sync, ierr)
            call check_error(ierr /= adios2_error_none, "Problem in ADIOS2 get zpsi.")
            call states_elec_set_state(st, mesh, idim, ist, ik, zpsi)
          end if
 
          filled(idim, ist, ik) = .true.
          iread = iread + 1
        end do
      end do
      safe_deallocate_a(dpsi)
      safe_deallocate_a(zpsi)
    end do
 
    pop_sub(states_elec_adios2_read_state_by_state)
  end subroutine states_elec_adios2_read_state_by_state
 
  ! ---------------------------------------------------------
  subroutine states_elec_adios2_read_batch_by_batch(st, mesh, engine, var, var_shape, type_requested, type_file, &
    nst_max, filled, iread)
    type(states_elec_t), target, intent(inout) :: st
    class(mesh_t),               intent(in)    :: mesh
    type(adios2_engine),         intent(inout) :: engine
    type(adios2_variable),       intent(inout) :: var
    integer(int64),              intent(in)    :: var_shape(:)
    integer,                     intent(in)    :: type_requested
    integer,                     intent(in)    :: type_file
    integer,                     intent(in)    :: nst_max
    logical,                     intent(inout) :: filled(1:, st%st_start:, st%d%kpt%start:)
    integer,                     intent(inout) :: iread
 
    integer :: ib, ik, ist, ip, adios2_mode, nst_linear, ierr
    class(wfs_elec_t), pointer :: psib
    real(real64), allocatable :: dff_pack(:, :)
 
    push_sub(states_elec_adios2_read_batch_by_batch)
 
    do ik = st%d%kpt%start, st%d%kpt%end
      ! if there are more k points requested than available, skip the ones that are not available
      ! this can, e.g., happen for unocc calculations with a k points path
      if (ik > var_shape(3)) cycle
      do ib = st%group%block_start, st%group%block_end
        select case (st%group%psib(ib, ik)%status())
        case (batch_not_packed)
          safe_allocate(psib)
          call st%group%psib(ib, ik)%copy_to(psib, copy_data=.false.)
          call psib%do_pack(batch_packed, copy=.false.)
          adios2_mode = adios2_mode_sync
        case (batch_packed)
          psib => st%group%psib(ib, ik)
          !adios2_mode = adios2_mode_deferred
          ! switch to sync mode for the time being, seems to solve some errors
          ! at the moment, it is unclear why
          adios2_mode = adios2_mode_sync
        case (batch_device_packed)
          safe_allocate(psib)
          call st%group%psib(ib, ik)%copy_to(psib)
          call psib%do_unpack(force=.true., copy=.false.)
          ! make sure we have a packed batch
          call psib%do_pack(batch_packed)
          adios2_mode = adios2_mode_sync
        end select
        assert(psib%status() == batch_packed)
 
        ! read only up to the highest available state in the restart data
        nst_linear = st%group%psib(ib, ik)%nst_linear
        if (states_elec_block_max(st, ib) > nst_max) then
          nst_linear = (nst_max - states_elec_block_min(st, ib) + 1) * st%d%dim
        end if
        if (nst_linear <= 0) then
          select case (st%group%psib(ib, ik)%status())
          case (batch_not_packed)
            call psib%end()
            safe_deallocate_p(psib)
          case (batch_packed)
            nullify(psib)
          case (batch_device_packed)
            call psib%end()
            safe_deallocate_p(psib)
          end select
          cycle
        end if
 
        ! select part of variable to write, correspond to the memory of the current batch
        ! first argument gives the start in the global array
        ! second argument gives the size of the local batch
        call adios2_set_selection(var, 3, &
          [int((states_elec_block_min(st, ib)-1)*st%d%dim, int64), int(mesh%pv%xlocal-1, int64), int(ik-1, int64)], &
          [int(nst_linear, int64), int(mesh%np, int64), 1_int64], ierr)
        call check_error(ierr /= adios2_error_none, "Problem setting ADIOS2 selection.")
 
        if (type_requested == adios2_type_dp .and. type_file == adios2_type_dp) then
          ! set memory layout of one batch, take into account padding and ghost/boundary points
          call adios2_set_memory_selection(var, 3, [0_int64, 0_int64, 0_int64], &
            [int(size(psib%dff_pack, 1), int64), &
            int(size(psib%dff_pack, 2), int64), 1_int64], ierr)
          call check_error(ierr /= adios2_error_none, "Problem setting ADIOS2 memory selection (reading real).")
 
          call adios2_get(engine, var, psib%dff_pack, adios2_mode, ierr)
          call check_error(ierr /= adios2_error_none, "Problem in ADIOS2 get (reading real).")
        else if (type_requested == adios2_type_complex_dp .and. type_file == adios2_type_complex_dp) then
          ! set memory layout of one batch, take into account padding and ghost/boundary points
          call adios2_set_memory_selection(var, 3, [0_int64, 0_int64, 0_int64], &
            [int(size(psib%zff_pack, 1), int64), &
            int(size(psib%zff_pack, 2), int64), 1_int64], ierr)
          call check_error(ierr /= adios2_error_none, "Problem setting ADIOS2 memory selection (reading complex).")
 
          call adios2_get(engine, var, psib%zff_pack, adios2_mode, ierr)
          call check_error(ierr /= adios2_error_none, "Problem in ADIOS2 get (reading complex).")
        else if (type_requested == adios2_type_complex_dp .and. type_file == adios2_type_dp) then
          ! read into a local real variable and copy to the complex batch buffer
          safe_allocate(dff_pack(size(psib%zff_pack, 1), size(psib%zff_pack, 2)))
          ! set memory layout of one batch, take into account padding and ghost/boundary points
          call adios2_set_memory_selection(var, 3, [0_int64, 0_int64, 0_int64], &
            [int(size(psib%zff_pack, 1), int64), &
            int(size(psib%zff_pack, 2), int64), 1_int64], ierr)
          call check_error(ierr /= adios2_error_none, "Problem setting ADIOS2 memory selection (reading complex from real).")
 
          ! use sync mode here to immediately read into dff_pack
          call adios2_get(engine, var, dff_pack, adios2_mode_sync, ierr)
          call check_error(ierr /= adios2_error_none, "Problem in ADIOS2 get (reading complex from real).")
          ! copy to batch memory
          !$omp parallel do private(ist)
          do ip = 1, mesh%np
            do ist = 1, nst_linear
              psib%zff_pack(ist, ip) = cmplx(dff_pack(ist, ip), m_zero, real64)
            end do
          end do
          safe_deallocate_a(dff_pack)
        else
          message(1) = "Error: requested to read complex states for restarting, but calculation has real states."
          call messages_fatal(1)
        end if
        ! make sure iread and filled are set correctly
        iread = iread + nst_linear
        do ist = 1, nst_linear
          filled(st%group%psib(ib, ik)%linear_to_idim(ist), st%group%psib(ib, ik)%linear_to_ist(ist), ik) = .true.
        end do
 
        select case (st%group%psib(ib, ik)%status())
        case (batch_not_packed)
          call psib%do_unpack()
          call psib%copy_data_to(mesh%np, st%group%psib(ib, ik))
          call psib%end()
          safe_deallocate_p(psib)
        case (batch_packed)
          nullify(psib)
        case (batch_device_packed)
          call psib%do_pack()
          call psib%copy_data_to(mesh%np, st%group%psib(ib, ik))
          call psib%end()
          safe_deallocate_p(psib)
        end select
      end do
    end do
 
    pop_sub(states_elec_adios2_read_batch_by_batch)
  end subroutine states_elec_adios2_read_batch_by_batch
 
  ! ---------------------------------------------------------
  ! only used for adios2 error checking
  subroutine check_error(condition, error_message)
    logical, intent(in) :: condition
    character(len=*), intent(in) :: error_message
    if (condition) then
      message(1) = error_message
      call messages_fatal(1)
    end if
  end subroutine check_error
#endif
 
  subroutine states_elec_dump_rho(restart, st, mesh, ierr, iter)
    type(restart_t),      intent(in)    :: restart
    type(states_elec_t),  intent(in)    :: st
    class(mesh_t),        intent(in)    :: mesh
    integer,              intent(out)   :: ierr
    integer,    optional, intent(in)    :: iter
 
    integer :: iunit, isp, err, err2(2)
    character(len=MAX_PATH_LEN) :: filename
    character(len=300) :: lines(2)
 
    push_sub(states_elec_dump_rho)
 
    ierr = 0
 
    if (restart%skip()) then
      pop_sub(states_elec_dump_rho)
      return
    end if
 
    message(1) = "Debug: Writing density restart."
    call messages_info(1, debug_only=.true.)
 
    call restart_block_signals()
 
    !write the densities
    iunit = restart%open('density')
    lines(1) = '#     #spin    #nspin    filename'
    lines(2) = '%densities'
    call restart%write(iunit, lines, 2, err)
    if (err /= 0) ierr = ierr + 1
 
    err2 = 0
    do isp = 1, st%d%nspin
      filename = get_filename_with_spin('density', st%d%nspin, isp)
      write(lines(1), '(i8,a,i8,a)') isp, ' | ', st%d%nspin, ' | "'//trim(adjustl(filename))//'"'
      call restart%write(iunit, lines, 1, err)
      if (err /= 0) err2(1) = err2(1) + 1
 
      call restart%write_mesh_function(filename, mesh, st%rho(:,isp), err)
      if (err /= 0) err2(2) = err2(2) + 1
 
    end do
    if (err2(1) /= 0) ierr = ierr + 2
    if (err2(2) /= 0) ierr = ierr + 4
 
    lines(1) = '%'
    call restart%write(iunit, lines, 1, err)
    if (err /= 0) ierr = ierr + 8
    if (present(iter)) then
      write(lines(1),'(a,i7)') 'Iter = ', iter
      call restart%write(iunit, lines, 1, err)
      if (err /= 0) ierr = ierr + 16
    end if
    call restart%close(iunit)
 
    call restart_unblock_signals()
 
    message(1) = "Debug: Writing density restart done."
    call messages_info(1, debug_only=.true.)
 
    pop_sub(states_elec_dump_rho)
  end subroutine states_elec_dump_rho
 
 
  ! ---------------------------------------------------------
  subroutine states_elec_load_rho(restart, st, mesh, ierr)
    type(restart_t),      intent(in)    :: restart
    type(states_elec_t),  intent(inout) :: st
    class(mesh_t),        intent(in)    :: mesh
    integer,              intent(out)   :: ierr
 
    integer              :: err, err2, isp
    character(len=MAX_PATH_LEN) :: filename
 
    push_sub(states_elec_load_rho)
 
    ierr = 0
 
    if (restart%skip()) then
      ierr = -1
      pop_sub(states_elec_load_rho)
      return
    end if
 
    message(1) = "Debug: Reading density restart."
    call messages_info(1, debug_only=.true.)
 
    ! skip for now, since we know what the files are going to be called
    !read the densities
!    iunit_rho = io_open(trim(dir)//'/density', restart%namespace, action='write')
!    call iopar_read(st%dom_st_kpt_mpi_grp, iunit_rho, line, err)
!    call iopar_read(st%dom_st_kpt_mpi_grp, iunit_rho, line, err)
!   we could read the iteration 'iter' too, not sure if that is useful.
 
    err2 = 0
    do isp = 1, st%d%nspin
      filename = get_filename_with_spin('density', st%d%nspin, isp)
!      if (st%dom_st_kpt_mpi_grp%is_root()) then
!        read(iunit_rho, '(i8,a,i8,a)') isp, ' | ', st%d%nspin, ' | "'//trim(adjustl(filename))//'"'
!      end if
      call restart%read_mesh_function(filename, mesh, st%rho(:,isp), err)
      if (err /= 0) err2 = err2 + 1
 
    end do
    if (err2 /= 0) ierr = ierr + 1
 
    message(1) = "Debug: Reading density restart done."
    call messages_info(1, debug_only=.true.)
 
    pop_sub(states_elec_load_rho)
  end subroutine states_elec_load_rho
 
  subroutine states_elec_dump_frozen(restart, space, st, mesh, ierr)
    type(restart_t),      intent(in)    :: restart
    class(space_t),       intent(in)    :: space
    type(states_elec_t),  intent(in)    :: st
    class(mesh_t),        intent(in)    :: mesh
    integer,              intent(out)   :: ierr
 
    integer :: isp, err, err2(2), idir
    character(len=MAX_PATH_LEN) :: filename
 
    push_sub(states_elec_dump_frozen)
 
    ierr = 0
 
    assert(allocated(st%frozen_rho))
 
    if (restart%skip()) then
      pop_sub(states_elec_dump_frozen)
      return
    end if
 
    message(1) = "Debug: Writing frozen densities restart."
    call messages_info(1, debug_only=.true.)
 
    call restart_block_signals()
 
    err2 = 0
    do isp = 1, st%d%nspin
      filename = get_filename_with_spin('frozen_rho', st%d%nspin, isp)
 
      call restart%write_mesh_function(filename, mesh, st%frozen_rho(:,isp), err)
      if (err /= 0) err2(2) = err2(2) + 1
 
      if (allocated(st%frozen_tau)) then
        filename = get_filename_with_spin('frozen_tau', st%d%nspin, isp)
        call restart%write_mesh_function(filename, mesh, st%frozen_tau(:,isp), err)
        if (err /= 0) err2 = err2 + 1
      end if
 
      if (allocated(st%frozen_gdens)) then
        do idir = 1, space%dim
          if (st%d%nspin == 1) then
            write(filename, fmt='(a,i1)') 'frozen_gdens-dir', idir
          else
            write(filename, fmt='(a,i1,a,i1)') 'frozen_gdens-dir', idir, '-', isp
          end if
          call restart%write_mesh_function(filename, mesh, st%frozen_gdens(:,idir,isp), err)
          if (err /= 0) err2 = err2 + 1
        end do
      end if
 
      if (allocated(st%frozen_ldens)) then
        filename = get_filename_with_spin('frozen_ldens', st%d%nspin, isp)
        call restart%write_mesh_function(filename, mesh, st%frozen_ldens(:,isp), err)
        if (err /= 0) err2 = err2 + 1
      end if
 
    end do
    if (err2(1) /= 0) ierr = ierr + 2
    if (err2(2) /= 0) ierr = ierr + 4
 
    call restart_unblock_signals()
 
    message(1) = "Debug: Writing frozen densities restart done."
    call messages_info(1, debug_only=.true.)
 
    pop_sub(states_elec_dump_frozen)
  end subroutine states_elec_dump_frozen
 
 
  ! ---------------------------------------------------------
  subroutine states_elec_load_frozen(restart, space, st, mesh, ierr)
    type(restart_t),      intent(in)    :: restart
    class(space_t),       intent(in)    :: space
    type(states_elec_t),  intent(inout) :: st
    class(mesh_t),        intent(in)    :: mesh
    integer,              intent(out)   :: ierr
 
    integer              :: err, err2, isp, idir
    character(len=MAX_PATH_LEN) :: filename
 
    push_sub(states_elec_load_frozen)
 
    assert(allocated(st%frozen_rho))
 
    ierr = 0
 
    if (restart%skip()) then
      ierr = -1
      pop_sub(states_elec_load_frozen)
      return
    end if
 
    message(1) = "Debug: Reading densities restart."
    call messages_info(1, debug_only=.true.)
 
    err2 = 0
    do isp = 1, st%d%nspin
      filename = get_filename_with_spin('frozen_rho', st%d%nspin, isp)
      call restart%read_mesh_function(filename, mesh, st%frozen_rho(:,isp), err)
      if (err /= 0) err2 = err2 + 1
 
      if (allocated(st%frozen_tau)) then
        filename = get_filename_with_spin('frozen_tau', st%d%nspin, isp)
        call restart%read_mesh_function(filename, mesh, st%frozen_tau(:,isp), err)
        if (err /= 0) err2 = err2 + 1
      end if
 
      if (allocated(st%frozen_gdens)) then
        do idir = 1, space%dim
          if (st%d%nspin == 1) then
            write(filename, fmt='(a,i1)') 'frozen_gdens-dir', idir
          else
            write(filename, fmt='(a,i1,a,i1)') 'frozen_gdens-dir', idir, '-', isp
          end if
          call restart%read_mesh_function(filename, mesh, st%frozen_gdens(:,idir,isp), err)
          if (err /= 0) err2 = err2 + 1
        end do
      end if
 
      if (allocated(st%frozen_ldens)) then
        filename = get_filename_with_spin('frozen_ldens', st%d%nspin, isp)
        call restart%read_mesh_function(filename, mesh, st%frozen_ldens(:,isp), err)
        if (err /= 0) err2 = err2 + 1
      end if
 
    end do
    if (err2 /= 0) ierr = ierr + 1
 
    message(1) = "Debug: Reading frozen densities restart done."
    call messages_info(1, debug_only=.true.)
 
    pop_sub(states_elec_load_frozen)
  end subroutine states_elec_load_frozen
 
 
  ! ---------------------------------------------------------
  subroutine states_elec_read_user_def_orbitals(mesh, namespace, space, st)
    class(mesh_t),       intent(in)    :: mesh
    type(namespace_t),   intent(in)    :: namespace
    class(space_t),      intent(in)    :: space
    type(states_elec_t), intent(inout) :: st
 
    type(block_t) :: blk
    integer :: ip, id, is, ik, nstates, state_from, ierr, ncols
    integer :: ib, idim, inst, inik, normalize
    real(real64) :: xx(space%dim), rr, psi_re, psi_im
    character(len=150) :: filename
    complex(real64), allocatable :: zpsi(:, :)
 
    integer, parameter ::           &
      state_from_formula  = 1,      &
      state_from_file     = -10010, &
      normalize_yes       = 1,      &
      normalize_no        = 0
 
    push_sub(states_elec_read_user_def_orbitals)
 
    !%Variable UserDefinedStates
    !%Type block
    !%Section States
    !%Description
    !% Instead of using the ground state as initial state for
    !% time-propagations it might be interesting in some cases
    !% to specify alternate states. Like with user-defined
    !% potentials, this block allows you to specify formulas for
    !% the orbitals at <i>t</i>=0.
    !%
    !% Example:
    !%
    !% <tt>%UserDefinedStates
    !% <br>&nbsp;&nbsp; 1 | 1 | 1 | formula | "exp(-r^2)*exp(-i*0.2*x)" | normalize_yes
    !% <br>%</tt>
    !%
    !% The first column specifies the component of the spinor,
    !% the second column the number of the state and the third
    !% contains <i>k</i>-point and spin quantum numbers. Column four
    !% indicates that column five should be interpreted as a formula
    !% for the corresponding orbital.
    !%
    !% Alternatively, if column four states <tt>file</tt> the state will
    !% be read from the file given in column five.
    !%
    !% <tt>%UserDefinedStates
    !% <br>&nbsp;&nbsp; 1 | 1 | 1 | file | "/path/to/file" | normalize_no
    !% <br>%</tt>
    !%
    !% Octopus reads first the ground-state orbitals from
    !% the <tt>restart/gs</tt> directory. Only the states that are
    !% specified in the above block will be overwritten with
    !% the given analytic expression for the orbital.
    !%
    !% The sixth (optional) column indicates whether <tt>Octopus</tt> should renormalize
    !% the orbital. The default (no sixth column given) is to renormalize.
    !%
    !%Option file -10010
    !% Read initial orbital from file.
    !% Accepted file formats, detected by extension: obf, ncdf and csv (real only).
    !% For csv files, the following formatting rules apply:
    !%
    !%  The functions in this file read an array from an ascii matrix (csv) file.
    !%  Format with values "valueXYZ" as follows
    !%
    !%  File values.csv:
    !%  --------
    !%  value111 value211 value311 value411
    !%  value121 value221 value321 value421
    !%  value131 value231 value331 value431
    !%
    !%  value112 value212 value312 value412
    !%  value122 value222 value322 value422
    !%  value132 value232 value332 value432
    !%
    !%  value113 value213 value313 value413
    !%  value123 value223 value323 value423
    !%  value133 value233 value333 value433
    !%  --------
    !%
    !%  That is, every line scans the x-coordinate, every XY-plane as a table of values
    !%  and all XY-planes separated by an empty row.
    !%
    !%  The given matrix is interpolated/stretched to fit the calculation
    !%  box defined in input file.
    !%
    !%  Note that there must not be any empty line at the end of the file.
    !%
    !%  Calculation box shape must be "parallelepiped".
    !%
    !%  The delimiter can be a tab, a comma or a space.
    !%Option formula 1
    !% Calculate initial orbital by given analytic expression.
    !%Option normalize_yes 1
    !% Normalize orbitals (default).
    !%Option normalize_no 0
    !% Do not normalize orbitals.
    !%End
    if (parse_block(namespace, 'UserDefinedStates', blk) == 0) then
 
      call messages_print_with_emphasis(msg=trim('Substitution of orbitals'), namespace=namespace)
 
      ! find out how many lines (i.e. states) the block has
      nstates = parse_block_n(blk)
 
      safe_allocate(zpsi(1:mesh%np, 1:st%d%dim))
 
      ! read all lines
      do ib = 1, nstates
        ! Check that number of columns is five or six.
        ncols = parse_block_cols(blk, ib - 1)
        if (ncols  <  5 .or. ncols > 6) then
          message(1) = 'Each line in the UserDefinedStates block must have'
          message(2) = 'five or six columns.'
          call messages_fatal(2, namespace=namespace)
        end if
 
        call parse_block_integer(blk, ib - 1, 0, idim)
        call parse_block_integer(blk, ib - 1, 1, inst)
        call parse_block_integer(blk, ib - 1, 2, inik)
 
        ! Calculate from expression or read from file?
        call parse_block_integer(blk, ib - 1, 3, state_from)
 
        ! loop over all states
        do id = 1, st%d%dim
          do is = 1, st%nst
            do ik = 1, st%nik
 
              if (.not.(id == idim .and. is == inst .and. ik == inik )) cycle
 
              ! We first parse the value, such that the parsing and stdout looks correct
              ! However, if rank 0 is not responsible for the parsing, the variables in the formula
              ! will still appear as not parsed (pure k-point/state parallelization case).
              ! This could be solved by a dummy call to parse_expression
              select case (state_from)
              case (state_from_formula)
                ! parse formula string
                call parse_block_string(                            &
                  blk, ib - 1, 4, st%user_def_states(id, is, ik))
 
                write(message(1), '(a,3i5)') 'Substituting state of orbital with k, ist, dim = ', ik, is, id
                write(message(2), '(2a)') '  with the expression:'
                write(message(3), '(2a)') '  ',trim(st%user_def_states(id, is, ik))
                call messages_info(3, namespace=namespace)
 
                ! convert to C string
                call conv_to_c_string(st%user_def_states(id, is, ik))
 
              case (state_from_file)
                ! The input format can be coded in column four now. As it is
                ! not used now, we just say "file".
                ! Read the filename.
                call parse_block_string(blk, ib - 1, 4, filename)
 
                write(message(1), '(a,3i5)') 'Substituting state of orbital with k, ist, dim = ', ik, is, id
                write(message(2), '(2a)') '  with data from file:'
                write(message(3), '(2a)') '  ',trim(filename)
                call messages_info(3, namespace=namespace)
              end select
 
 
              ! does the block entry match and is this node responsible?
              if (.not.(st%st_start  <=  is .and. st%st_end >= is          &
                .and. st%d%kpt%start  <=  ik .and. st%d%kpt%end >= ik)) cycle
 
              select case (state_from)
 
              case (state_from_formula)
                ! fill states with user-defined formulas
                do ip = 1, mesh%np
                  xx = mesh%x(:, ip)
                  rr = norm2(xx)
 
                  ! parse user-defined expressions
                  call parse_expression(psi_re, psi_im, space%dim, xx, rr, m_zero, st%user_def_states(id, is, ik))
                  ! fill state
                  zpsi(ip, 1) = cmplx(psi_re, psi_im, real64)
                end do
 
              case (state_from_file)
                ! finally read the state
                call zio_function_input(filename, namespace, space, mesh, zpsi(:, 1), ierr)
                if (ierr > 0) then
                  message(1) = 'Could not read the file!'
                  write(message(2),'(a,i1)') 'Error code: ', ierr
                  call messages_fatal(2, namespace=namespace)
                end if
 
              case default
                message(1) = 'Wrong entry in UserDefinedStates, column 4.'
                message(2) = 'You may state "formula" or "file" here.'
                call messages_fatal(2, namespace=namespace)
              end select
 
              ! normalize orbital
              if (parse_block_cols(blk, ib - 1) == 6) then
                call parse_block_integer(blk, ib - 1, 5, normalize)
              else
                normalize = 1
              end if
              select case (normalize)
              case (normalize_no)
                call states_elec_set_state(st, mesh, id, is, ik, zpsi(:, 1))
              case (normalize_yes)
                assert(st%d%dim == 1)
                call zmf_normalize(mesh, st%d%dim, zpsi)
                call states_elec_set_state(st, mesh, is, ik, zpsi)
              case default
                message(1) = 'The sixth column in UserDefinedStates may either be'
                message(2) = '"normalize_yes" or "normalize_no"'
                call messages_fatal(2, namespace=namespace)
              end select
 
            end do
          end do
        end do
 
      end do
 
      safe_deallocate_a(zpsi)
 
      call parse_block_end(blk)
      call messages_print_with_emphasis(namespace=namespace)
 
    else
      call messages_variable_is_block(namespace, 'UserDefinedStates')
    end if
 
    pop_sub(states_elec_read_user_def_orbitals)
  end subroutine states_elec_read_user_def_orbitals
 
  ! ---------------------------------------------------------
  ! This is needed as for the generalized Bloch theorem we need to label
  ! the states from the expectation value of Sz computed from the GS.
  subroutine states_elec_dump_spin(restart, st, ierr)
    type(restart_t),      intent(in)  :: restart
    type(states_elec_t),  intent(in)  :: st
    integer,              intent(out) :: ierr
 
    integer :: iunit_spin
    integer :: err, err2(2), ik, ist
    character(len=300) :: lines(3)
 
    push_sub(states_elec_dump_spin)
 
    ierr = 0
 
    if (restart%skip()) then
      pop_sub(states_elec_dump_spin)
      return
    end if
 
    call profiling_in("RESTART_WRITE_SPIN")
 
    call restart_block_signals()
 
    iunit_spin = restart%open('spin')
    lines(1) = '#     #k-point            #st       #spin(x) spin(y) spin(z)'
    call restart%write(iunit_spin, lines, 1, err)
    if (err /= 0) ierr = ierr + 1
 
    err2 = 0
    do ik = 1, st%nik
      do ist = 1, st%nst
        write(lines(1), '(i8,a,i8,3(a,f18.12))') ik, ' | ', ist, ' | ', &
          st%spin(1,ist,ik), ' | ', st%spin(2,ist,ik),' | ', st%spin(3,ist,ik)
        call restart%write(iunit_spin, lines, 1, err)
        if (err /= 0) err2(1) = err2(1) + 1
      end do ! st%nst
    end do ! st%nik
    lines(1) = '%'
    call restart%write(iunit_spin, lines, 1, err)
    if (err2(1) /= 0) ierr = ierr + 8
    if (err2(2) /= 0) ierr = ierr + 16
 
    call restart%close(iunit_spin)
 
    call restart_unblock_signals()
 
    call profiling_out("RESTART_WRITE_SPIN")
    pop_sub(states_elec_dump_spin)
  end subroutine states_elec_dump_spin
 
 
 
 
  ! ---------------------------------------------------------
  subroutine states_elec_load_spin(restart, st, ierr)
    type(restart_t),            intent(in)    :: restart
    type(states_elec_t),        intent(inout) :: st
    integer,                    intent(out)   :: ierr
 
    integer              :: spin_file, err, ik, ist
    character(len=256)   :: lines(3)
    real(real64)         :: spin(3)
    character(len=1)     :: char
 
 
    push_sub(states_elec_load_spin)
 
    ierr = 0
 
    if (restart%skip()) then
      ierr = -1
      pop_sub(states_elec_load_spin)
      return
    end if
 
    call profiling_in("RESTART_READ_SPIN")
 
    ! open files to read
    spin_file = restart%open('spin')
    ! Skip two lines.
    call restart%read(spin_file, lines, 1, err)
    if (err /= 0) ierr = ierr - 2**7
 
    ! If any error occured up to this point then it is not worth continuing,
    ! as there something fundamentally wrong with the restart files
    if (ierr /= 0) then
      call restart%close(spin_file)
      call profiling_out("RESTART_READ_SPIN")
      pop_sub(states_elec_load_spin)
      return
    end if
 
    ! Next we read the list of states from the files.
    ! Errors in reading the information of a specific state from the files are ignored
    ! at this point, because later we will skip reading the wavefunction of that state.
    do
      call restart%read(spin_file, lines, 1, err)
      if (err /= 0) cycle
 
      read(lines(1), '(a)') char
      if (char == '%') then
        !We reached the end of the file
        exit
      else
        read(lines(1), *) ik, char, ist, char,  spin(1), char, spin(2), char, spin(3)
      end if
 
      st%spin(1:3, ist, ik) = spin(1:3)
    end do
 
    call restart%close(spin_file)
 
    call profiling_out("RESTART_READ_SPIN")
    pop_sub(states_elec_load_spin)
  end subroutine states_elec_load_spin
 
  ! ---------------------------------------------------------
  subroutine states_elec_transform(st, namespace, space, restart, mesh, kpoints, prefix)
    type(states_elec_t),        intent(inout) :: st
    type(namespace_t),          intent(in)    :: namespace
    class(space_t),             intent(in)    :: space
    type(restart_t),            intent(inout) :: restart
    class(mesh_t),              intent(in)    :: mesh
    type(kpoints_t),            intent(in)    :: kpoints
    character(len=*), optional, intent(in)    :: prefix
 
    type(states_elec_t) :: stin
    type(block_t) :: blk
    complex(real64), allocatable :: rotation_matrix(:,:), psi(:, :)
    integer :: ist, jst, ncols, iqn
    character(len=256) :: block_name
 
    push_sub(states_elec_transform)
 
    block_name = trim(optional_default(prefix, "")) // "TransformStates"
 
    !%Variable TransformStates
    !%Type block
    !%Default no
    !%Section States
    !%Description
    !% Before starting the <tt>td</tt> calculation, the initial states (that are
    !% read from the <tt>restart/gs</tt> directory, which should have been
    !% generated in a previous ground-state calculation) can be "transformed"
    !% among themselves. The block <tt>TransformStates</tt> gives the transformation matrix
    !% to be used. The number of rows and columns of the matrix should equal the number
    !% of the states present in the time-dependent calculation (the independent
    !% spin and <i>k</i>-point subspaces are all transformed equally); the number of
    !% columns should be equal to the number of states present in the
    !% <tt>restart/gs</tt> directory. This number may be different: for example,
    !% one could have run previously in <tt>unocc</tt> mode in order to obtain unoccupied
    !% Kohn-Sham states, and therefore <tt>restart/gs</tt> will contain more states.
    !% These states can be used in the transformation.
    !%
    !% Note that the code will not check the orthonormality of the new states!
    !% Occupations are also fixed to the one of the previous states
    !%
    !% Each line provides the coefficients of the new states, in terms of
    !% the old ones. The coefficients are complex, but the imaginary part will be
    !% ignored for real wavefunctions.
    !% Note: This variable cannot be used when parallel in states.
    !%End
    if (parse_is_defined(namespace, trim(block_name))) then
      if (parse_block(namespace, trim(block_name), blk) == 0) then
        if (st%parallel_in_states) then
          call messages_not_implemented(trim(block_name) // " parallel in states", namespace=namespace)
        end if
        if (parse_block_n(blk) /= st%nst) then
          message(1) = "Number of rows in block " // trim(block_name) // " must equal number of states in this calculation."
          call messages_fatal(1, namespace=namespace)
        end if
        call states_elec_copy(stin, st, exclude_wfns = .true.)
        call states_elec_look_and_load(restart, namespace, space, stin, mesh, kpoints, fixed_occ = .true.)
 
        ! FIXME: rotation matrix should be R_TYPE
        safe_allocate(rotation_matrix(1:stin%nst, 1:stin%nst))
        safe_allocate(psi(1:mesh%np, 1:st%d%dim))
 
        rotation_matrix = diagonal_matrix(stin%nst, m_z1)
 
        do ist = 1, st%nst
          ncols = parse_block_cols(blk, ist-1)
          if (ncols /= stin%nst) then
            write(message(1),'(a,i6,a,i6,3a,i6,a)') "Number of columns (", ncols, ") in row ", ist, " of block ", &
              trim(block_name), " must equal number of states (", stin%nst, ") read from gs restart."
            call messages_fatal(1, namespace=namespace)
          end if
          do jst = 1, stin%nst
            call parse_block_cmplx(blk, ist - 1, jst - 1, rotation_matrix(jst, ist))
          end do
        end do
 
        call parse_block_end(blk)
 
        do iqn = st%d%kpt%start, st%d%kpt%end
          if (states_are_real(st)) then
            call states_elec_rotate(stin, mesh, real(rotation_matrix, real64) , iqn)
          else
            call states_elec_rotate(stin, mesh, rotation_matrix, iqn)
          end if
 
          do ist = st%st_start, st%st_end
            call states_elec_get_state(stin, mesh, ist, iqn, psi)
            call states_elec_set_state(st, mesh, ist, iqn, psi)
          end do
 
        end do
 
        safe_deallocate_a(rotation_matrix)
        safe_deallocate_a(psi)
 
        call states_elec_end(stin)
 
      else
        call messages_variable_is_block(namespace, trim(block_name))
      end if
 
    end if
 
    pop_sub(states_elec_transform)
  end subroutine states_elec_transform
 
end module states_elec_restart_oct_m
 
 
!! Local Variables:
!! mode: f90
!! coding: utf-8
!! End: