pes_flux.F90

Go to the documentation of this file.

!! Copyright (C) 2015 P. Wopperer and U. De Giovannini
!!
!! 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 pes_flux_oct_m
  use accel_oct_m
  use absorbing_boundaries_oct_m
  use batch_oct_m
  use batch_ops_oct_m
  use boundaries_oct_m
  use box_oct_m
  use box_parallelepiped_oct_m
  use box_sphere_oct_m
  use calc_mode_par_oct_m
  use comm_oct_m
  use debug_oct_m
  use derivatives_oct_m
  use electron_space_oct_m
  use ext_partner_list_oct_m
  use global_oct_m
  use kpoints_oct_m
  use lalg_basic_oct_m
  use interaction_partner_oct_m
  use io_oct_m
  use, intrinsic :: iso_fortran_env
  use lasers_oct_m
  use loct_oct_m
  use loct_math_oct_m
  use math_oct_m
  use mesh_interpolation_oct_m
  use mesh_oct_m
  use messages_oct_m
  use multicomm_oct_m
  use mpi_oct_m
  use namespace_oct_m
  use parser_oct_m
  use phase_oct_m
  use profiling_oct_m
  use restart_oct_m
  use space_oct_m
  use states_elec_oct_m
  use states_elec_dim_oct_m
  use string_oct_m
  use types_oct_m
  use unit_oct_m
  use unit_system_oct_m
  use utils_oct_m
  use varinfo_oct_m
  use wfs_elec_oct_m
 
  implicit none
 
  private
 
  public ::                    &
    pes_flux_t,                &
    pes_flux_init,             &
    pes_flux_end,              &
    pes_flux_calc,             &
    pes_flux_output,           &
    pes_flux_load,             &
    pes_flux_dump,             &
    pes_flux_reciprocal_mesh_gen, &
    pes_flux_pmesh,            &
    pes_flux_map_from_states,  &
    pes_flux_out_energy,       &
    pes_flux_out_vmap
 
  type pes_flux_t
    private
    
 
    integer            :: nkpnts
    integer            :: nkpnts_start, nkpnts_end
    integer            :: nk
    integer            :: nk_start, nk_end
 
    ! spherical momentum grid
    integer            :: nstepsthetak, nstepsphik
    real(real64)       :: thetak_rng(1:2)
    real(real64)       :: phik_rng(1:2)
    integer            :: nstepsomegak
    integer            :: nstepsomegak_start, nstepsomegak_end
    real(real64)       :: ktransf(1:3,1:3)
 
    real(real64), allocatable :: klinear(:,:)
    !                                                    !< polar:     klinear(nk,1)    defines the radial part of the k-mesh
    !                                                    !< cartesian: klinear(nk,1:3)  defines the the k-mesh along each axes
 
    real(real64)       :: dk
    real(real64), allocatable :: kcoords_cub(:,:,:)
    real(real64), allocatable :: kcoords_sph(:,:,:)
    integer            :: kgrid
 
    ! Surface relates quantities
    integer, public    :: surf_shape
    integer            :: nsrfcpnts
    integer            :: nsrfcpnts_start, nsrfcpnts_end
    real(real64), allocatable :: srfcnrml(:,:)
    real(real64), allocatable :: rcoords(:,:)
    integer, allocatable :: srfcpnt(:)
    integer, allocatable :: rankmin(:)
    integer            :: lmax
    complex(real64), allocatable :: ylm_r(:,:,:)
    complex(real64), allocatable :: ylm_k(:,:,:)
    real(real64), allocatable :: j_l(:,:)
    real(real64)       :: radius
    integer, allocatable :: face_idx_range(:,:)
    !                                                  !! 1 (2) start (end) idx of face nface in rcoords(:,:)
 
 
    complex(real64),   allocatable :: bvk_phase(:,:)
    complex(real64),   allocatable :: expkr(:,:,:,:)
    complex(real64),   allocatable :: expkr_perp(:,:)
    !                                                  !! direction perpendicular to the face
    real(real64),   allocatable :: LLr(:,:)
    integer, allocatable :: NN(:,:)
    integer, allocatable :: Lkpuvz_inv(:,:,:)
    !                                                  !! (u,v parametric on face and z perpendicular)
 
    ! Surface and time integration
    integer          :: tdsteps
    !                                                  !< = mesh%mpi_grp%size (PES_SPHERICAL)
    integer          :: max_iter
    integer          :: save_iter
    integer          :: itstep
 
    complex(real64), allocatable :: wf(:,:,:,:,:)
    complex(real64), allocatable :: gwf(:,:,:,:,:,:)
    real(real64),    allocatable :: veca(:,:)
    complex(real64), allocatable :: conjgphase_prev(:,:)
    complex(real64), allocatable :: spctramp_cub(:,:,:,:)
    complex(real64), allocatable :: spctramp_sph(:,:,:,:,:)
 
    integer, public  :: ll(3)
    !                                                  !< mesh. Used when working with semi-periodic systems
 
    type(mesh_interpolation_t) :: interp
 
    logical          :: parallel_in_momentum
    logical          :: arpes_grid
    logical          :: surf_interp
    logical          :: use_symmetry
    logical          :: runtime_output
    logical          :: anisotrpy_correction
 
    integer          :: par_strategy
    integer          :: dim
    integer          :: pdim
 
    ! Buffers for GPU acceleration
    type(accel_mem_t) :: pt_buff
    type(accel_mem_t) :: interp_buff
    type(accel_mem_t) :: spacing_buff
    type(accel_mem_t) :: coords_buff
 
  end type pes_flux_t
 
  integer, parameter ::   &
    PES_CUBIC      = 1,     &
    pes_spherical  = 2,     &
    pes_plane      = 3
 
  integer, parameter ::   &
    PES_POLAR      = 1,     &
    pes_cartesian  = 2
 
 
contains
 
 
  ! ---------------------------------------------------------
  subroutine pes_flux_init(this, namespace, space, mesh, st, ext_partners, kpoints, abb, save_iter, max_iter)
    type(pes_flux_t),             intent(inout) :: this
    type(namespace_t),            intent(in)    :: namespace
    class(space_t),               intent(in)    :: space
    type(mesh_t),                 intent(in)    :: mesh
    type(states_elec_t),          intent(in)    :: st
    type(partner_list_t),         intent(in)    :: ext_partners
    type(kpoints_t),              intent(in)    :: kpoints
    type(absorbing_boundaries_t), intent(in)    :: abb
    integer,                      intent(in)    :: save_iter
    integer,                      intent(in)    :: max_iter
 
    type(block_t)      :: blk
    real(real64)       :: border(space%dim)       ! distance of surface from border
    real(real64)       :: offset(space%dim)       ! offset for border
    integer            :: stst, stend, kptst, kptend, sdim, mdim, pdim
    integer            :: imdim
    integer            :: isp
    integer            :: il
 
    integer            :: nstepsphir, nstepsthetar
    integer            :: ll, mm
    integer            :: default_shape
    real(real64)       :: fc_ptdens
 
    integer            :: par_strategy
    logical            :: use_symmetry
 
    type(lasers_t), pointer :: lasers
 
    push_sub_with_profile(pes_flux_init)
 
    stst   = st%st_start
    stend  = st%st_end
    kptst  = st%d%kpt%start
    kptend = st%d%kpt%end
    sdim   = st%d%dim
    mdim   = space%dim
    pdim   = space%periodic_dim
 
    this%surf_interp = .false.
 
 
    lasers => list_get_lasers(ext_partners)
    if(associated(lasers)) then
      do il = 1, lasers%no_lasers
        if (laser_kind(lasers%lasers(il)) /= e_field_vector_potential) then
          message(1) = 't-surff only works in velocity gauge.'
          call messages_fatal(1, namespace=namespace)
        end if
      end do
    end if
 
    message(1) = 'Info: Calculating PES using t-surff technique.'
    call messages_info(1, namespace=namespace)
 
    ! -----------------------------------------------------------------
    ! Setting up r-mesh (the surface)
    ! -----------------------------------------------------------------
    !%Variable PES_Flux_Shape
    !%Type integer
    !%Section Time-Dependent::PhotoElectronSpectrum
    !%Description
    !% The shape of the surface.
    !%Option cub 1
    !% Uses a parallelepiped as surface. All surface points are grid points.
    !% Choose the location of the surface with variable <tt>PES_Flux_Lsize</tt>
    !% (default for 1D and 2D).
    !%Option sph 2
    !% Constructs a sphere with radius <tt>PES_Flux_Radius</tt>. Points on the sphere
    !% are interpolated by trilinear interpolation (default for 3D).
    !%Option pln 3
    !% This option is for periodic systems.
    !% Constructs a plane perpendicular to the non-periodic dimension
    !% at <tt>PES_Flux_Lsize</tt>.
    !%End
 
    default_shape = pes_spherical
    if (space%is_periodic()) then
      default_shape = pes_plane
    else if (mdim <= 2) then
      default_shape = pes_cubic
    else
      select type (box => mesh%box)
      type is (box_parallelepiped_t)
        default_shape = pes_cubic
      end select
    end if
 
    call parse_variable(namespace, 'PES_Flux_Shape', default_shape, this%surf_shape)
    if (.not. varinfo_valid_option('PES_Flux_Shape', this%surf_shape, is_flag = .true.)) then
      call messages_input_error(namespace,'PES_Flux_Shape')
    end if
    if (this%surf_shape == pes_spherical .and. mdim /= 3) then
      message(1) = 'Spherical grid works only in 3d.'
      call messages_fatal(1, namespace=namespace)
    end if
    call messages_print_var_option('PES_Flux_Shape', this%surf_shape, namespace=namespace)
 
 
    !%Variable PES_Flux_AnisotropyCorrection
    !%Type logical
    !%Section Time-Dependent::PhotoElectronSpectrum
    !%Description
    !% Apply anisotropy correction.
    !%
    !%End
    if (this%surf_shape == pes_cubic) then
      call parse_variable(namespace, 'PES_Flux_AnisotropyCorrection', .false., this%anisotrpy_correction)
      call messages_print_var_value('PES_Flux_AnisotropyCorrection', this%anisotrpy_correction, namespace=namespace)
    else
      this%anisotrpy_correction = .false.
    end if
 
    !%Variable PES_Flux_Offset
    !%Type block
    !%Section Time-Dependent::PhotoElectronSpectrum
    !%Description
    !% Shifts the surface for <tt>PES_Flux_Shape = cub</tt>. The syntax is:
    !%
    !% <tt>%PES_Flux_Offset
    !% <br>&nbsp;&nbsp;xshift | yshift | zshift
    !% <br>%
    !% </tt>
    !%End
    offset = m_zero
    if (parse_block(namespace, 'PES_Flux_Offset', blk) == 0) then
      do imdim = 1, mdim
        call parse_block_float(blk, 0, imdim - 1, offset(imdim))
      end do
      call parse_block_end(blk)
    end if
 
    !%Variable PES_Flux_Lmax
    !%Type integer
    !%Default 80
    !%Section Time-Dependent::PhotoElectronSpectrum
    !%Description
    !% Maximum order of the spherical harmonic to be integrated on an equidistant spherical
    !% grid (to be changed to Gauss-Legendre quadrature).
    !%End
    if (this%surf_shape == pes_spherical) then
      call parse_variable(namespace, 'PES_Flux_Lmax', 80, this%lmax)
      if (this%lmax < 1) call messages_input_error(namespace,'PES_Flux_Lmax', 'must be > 0')
      call messages_print_var_value('PES_Flux_Lmax', this%lmax, namespace=namespace)
    end if
 
 
    if (this%surf_shape == pes_cubic .or. this%surf_shape == pes_plane) then
 
 
      !%Variable PES_Flux_Lsize
      !%Type block
      !%Section Time-Dependent::PhotoElectronSpectrum
      !%Description
      !% For <tt>PES_Flux_Shape = cub</tt> sets the dimensions along each direction. The syntax is:
      !%
      !% <tt>%PES_Flux_Lsize
      !% <br>&nbsp;&nbsp;xsize | ysize | zsize
      !% <br>%
      !% </tt>
      !%
      !% where <tt>xsize</tt>, <tt>ysize</tt>, and <tt>zsize</tt> are with respect to the origin. The surface can
      !% be shifted with <tt>PES_Flux_Offset</tt>.
      !% If <tt>PES_Flux_Shape = pln</tt>, specifies the position of two planes perpendicular to
      !% the non-periodic dimension symmetrically placed at <tt>PES_Flux_Lsize</tt> distance from
      !% the origin.
      !%End
      if (parse_block(namespace, 'PES_Flux_Lsize', blk) == 0) then
        do imdim = 1, mdim
          call parse_block_float(blk, 0, imdim - 1, border(imdim))
        end do
        ! Snap the face to the closest grid point
        border(1:mdim) = int(border(1:mdim) / mesh%spacing(1:mdim)) * mesh%spacing(1:mdim)
        call parse_block_end(blk)
 
      else if (parse_is_defined(namespace, 'PES_Flux_Lsize')) then
        border(mdim)  = mesh%box%bounding_box_l(mdim) * m_half
        if (space%is_periodic()) then
          ! the cube sides along the periodic directions are out of the simulation box
          border(1:pdim)= mesh%box%bounding_box_l(1:pdim) * m_two
          call parse_variable(namespace, 'PES_Flux_Lsize', border(mdim), border(mdim))
          ! Snap the plane to the closest grid point
          border(mdim) = floor(border(mdim) / mesh%spacing(mdim)) * mesh%spacing(mdim)
        else
          call parse_variable(namespace, 'PES_Flux_Lsize', border(mdim), border(mdim))
          border(1:mdim - 1) = border(mdim)
          border(1:mdim)     = floor(border(1:mdim) / mesh%spacing(1:mdim)) * mesh%spacing(1:mdim)
        end if
      else
        select type (box => mesh%box)
        type is (box_sphere_t)
          border(1:mdim) = box%radius / sqrt(m_two) * m_half
        type is (box_parallelepiped_t)
          border(1:mdim) = box%half_length(1:mdim) * m_half
        class default
          call messages_write('PES_Flux_Lsize not specified. No default values available for this box shape.')
          call messages_new_line()
          call messages_write('Specify the location of the parallelepiped with block PES_Flux_Lsize.')
          call messages_fatal(namespace=namespace)
        end select
        call messages_write('PES_Flux_Lsize not specified. Using default values.')
        call messages_info(namespace=namespace)
      end if
 
      call messages_print_var_value('PES_Flux_Lsize', border(1:mdim), namespace=namespace)
 
      !%Variable PES_Flux_Face_Dens
      !%Type block
      !%Section Time-Dependent::PhotoElectronSpectrum
      !%Description
      !% Define the number of points density per unit of area (in au) on the
      !% face of the 'cub' surface.
      !%End
      if (parse_is_defined(namespace, 'PES_Flux_Face_Dens')) then
        this%surf_interp = .true.
        call parse_variable(namespace, 'PES_Flux_Face_Dens', m_one, fc_ptdens)
        call messages_print_var_value('PES_Flux_Face_Dens', fc_ptdens, namespace=namespace)
      end if
 
    else
 
      this%surf_interp = .true.
 
      !%Variable PES_Flux_Radius
      !%Type float
      !%Section Time-Dependent::PhotoElectronSpectrum
      !%Description
      !% The radius of the sphere, if <tt>PES_Flux_Shape == sph</tt>.
      !%End
      if (parse_is_defined(namespace, 'PES_Flux_Radius')) then
        call parse_variable(namespace, 'PES_Flux_Radius', m_zero, this%radius)
        if (this%radius <= m_zero) call messages_input_error(namespace, 'PES_Flux_Radius')
        call messages_print_var_value('PES_Flux_Radius', this%radius, namespace=namespace)
      else
        select type (box => mesh%box)
        type is (box_sphere_t)
          this%radius = box%radius
        type is (box_parallelepiped_t)
          this%radius = minval(box%half_length(1:mdim))
        class default
          message(1) = 'PES_Flux_Radius not specified. No default values available for this box shape.'
          message(2) = 'Specify the radius of the sphere with variable PES_Flux_Radius.'
          call messages_fatal(2, namespace=namespace)
        end select
        message(1) = 'PES_Flux_Radius not specified. Using default values.'
        call messages_info(1, namespace=namespace)
        call messages_print_var_value('PES_Flux_Radius', this%radius, namespace=namespace)
      end if
 
      !%Variable PES_Flux_StepsThetaR
      !%Type integer
      !%Default: 2 <tt>PES_Flux_Lmax</tt> + 1
      !%Section Time-Dependent::PhotoElectronSpectrum
      !%Description
      !% Number of steps in <math>\theta</math> (<math>0 \le \theta \le \pi</math>) for the spherical surface.
      !%End
      call parse_variable(namespace, 'PES_Flux_StepsThetaR', 2 * this%lmax + 1, nstepsthetar)
      if (nstepsthetar < 0) call messages_input_error(namespace, 'PES_Flux_StepsThetaR')
      call messages_print_var_value("PES_Flux_StepsThetaR", nstepsthetar, namespace=namespace)
 
      !%Variable PES_Flux_StepsPhiR
      !%Type integer
      !%Default: 2 <tt>PES_Flux_Lmax</tt> + 1
      !%Section Time-Dependent::PhotoElectronSpectrum
      !%Description
      !% Number of steps in <math>\phi</math> (<math>0 \le \phi \le 2 \pi</math>) for the spherical surface.
      !%End
      call parse_variable(namespace, 'PES_Flux_StepsPhiR', 2 * this%lmax + 1, nstepsphir)
      if (nstepsphir < 0) call messages_input_error(namespace, 'PES_Flux_StepsPhiR')
      call messages_print_var_value("PES_Flux_StepsPhiR", nstepsphir, namespace=namespace)
 
    end if
 
 
    if (this%surf_interp)  call mesh_interpolation_init(this%interp, mesh)
 
    ! -----------------------------------------------------------------
    ! Get the surface points
    ! -----------------------------------------------------------------
    if (this%surf_shape == pes_cubic .or. this%surf_shape == pes_plane) then
      call pes_flux_getcube(this, mesh, abb, border, offset, fc_ptdens)
    else
      ! equispaced grid in theta & phi (Gauss-Legendre would optimize to nstepsthetar = this%lmax & nstepsphir = 2*this%lmax + 1):
      ! nstepsthetar = M_TWO * this%lmax + 1
      ! nstepsphir   = M_TWO * this%lmax + 1
      call pes_flux_getsphere(this, mesh, nstepsthetar, nstepsphir)
    end if
 
 
    !%Variable PES_Flux_Parallelization
    !%Type flag
    !%Section Time-Dependent::PhotoElectronSpectrum
    !%Description
    !% The parallelization strategy to be used to calculate the PES spectrum
    !% using the resources available in the domain parallelization pool.
    !% This option is not available without domain parallelization.
    !% Parallelization over k-points and states is always enabled when available.
    !%Option pf_none bit(1)
    !% No parallelization.
    !%Option pf_time bit(2)
    !% Parallelize time integration. This requires to store some quantities over a
    !% number of time steps equal to the number of cores available.
    !%Option pf_momentum bit(3)
    !% Parallelize over the final momentum grid. This strategy has a much lower
    !% memory footprint than the one above (time) but seems to provide a smaller
    !% speedup.
    !%Option pf_surface bit(4)
    !% Parallelize over surface points.
    !%
    !%
    !% Option pf_time and pf_surface can be combined: pf_time + pf_surface.
    !%
    !%End
 
    this%par_strategy      = option__pes_flux_parallelization__pf_none
    if (mesh%parallel_in_domains) then
 
      if (this%surf_shape == pes_spherical) then
        this%par_strategy = option__pes_flux_parallelization__pf_time  &
          + option__pes_flux_parallelization__pf_surface
      else
        this%par_strategy = option__pes_flux_parallelization__pf_surface
        if (space%dim == 1) this%par_strategy = option__pes_flux_parallelization__pf_time
      end if
      par_strategy = this%par_strategy
      call parse_variable(namespace, 'PES_Flux_Parallelization', par_strategy, this%par_strategy)
      if (.not. varinfo_valid_option('PES_Flux_Parallelization', this%par_strategy, is_flag = .true.)) then
        call messages_input_error(namespace,'PES_Flux_Parallelization')
      end if
 
    end if
 
    !Sanity check
    if (bitand(this%par_strategy, option__pes_flux_parallelization__pf_surface) /= 0 .and. &
      bitand(this%par_strategy, option__pes_flux_parallelization__pf_momentum) /= 0) then
      call messages_input_error(namespace,'PES_Flux_Parallelization', "Cannot combine pf_surface and pf_momentum")
    end if
 
    write(message(1),'(a)') 'Input: [PES_Flux_Parallelization = '
    if (bitand(this%par_strategy, option__pes_flux_parallelization__pf_none) /= 0) then
      write(message(1),'(a,1x,a)') trim(message(1)), 'pf_none '
    end if
    if (bitand(this%par_strategy, option__pes_flux_parallelization__pf_time) /= 0) then
      write(message(1),'(a,1x,a)') trim(message(1)), 'pf_time '
    end if
    if (bitand(this%par_strategy, option__pes_flux_parallelization__pf_momentum) /= 0) then
      write(message(1),'(a,1x,a)') trim(message(1)), 'pf_momentum'
    end if
    if (bitand(this%par_strategy, option__pes_flux_parallelization__pf_surface) /= 0) then
      write(message(1),'(a,1x,a)') trim(message(1)), 'pf_surface'
    end if
    write(message(1),'(2a)') trim(message(1)), ']'
    call messages_info(1, namespace=namespace)
 
 
 
    ! distribute the surface points on nodes,
    ! since mesh domains may have different numbers of surface points.
    if (bitand(this%par_strategy, option__pes_flux_parallelization__pf_surface) /= 0) then
 
      call pes_flux_distribute(1, this%nsrfcpnts, this%nsrfcpnts_start, this%nsrfcpnts_end, mesh%mpi_grp%comm)
 
      if (this%surf_shape /= pes_spherical) call pes_flux_distribute_facepnts_cub(this, mesh)
 
!   Keep this because is useful for debug but not enough to bother having it polished out
!       if (debug%info) then
!         call mpi_world%barrier()
!         write(*,*) &
!           'Debug: surface points on node ', mpi_world%rank, ' : ', this%nsrfcpnts_start, this%nsrfcpnts_end
!         call mpi_world%barrier()
!       end if
 
    else
      this%nsrfcpnts_start = 1
      this%nsrfcpnts_end   = this%nsrfcpnts
 
    end if
 
!   Keep this because is useful for debug but not enough to bother having it polished out
!     if (debug%info .and. mpi_world%is_root()) then
!       do isp = 1, this%nsrfcpnts
!         write(223,*) isp, this%rcoords(:, isp), this%srfcnrml(:,isp)
!       end do
!       if (this%nsrfcpnts > 0) flush(223)
!     end if
 
    ! Generate the momentum space mesh grid
    call pes_flux_reciprocal_mesh_gen(this, namespace, space, st, kpoints, mesh%mpi_grp%comm)
 
 
    !%Variable PES_Flux_UseSymmetries
    !%Type logical
    !%Section Time-Dependent::PhotoElectronSpectrum
    !%Description
    !% Use surface and momentum grid symmetries to speed up calculation and
    !% lower memory footprint.
    !% By default available only when the surface shape matches the grid symmetry i.e.:
    !% PES_Flux_Shape = m_cub or m_pln and PES_Flux_Momenutum_Grid = m_cartesian
    !% or
    !% PES_Flux_Shape = m_sph and PES_Flux_Momenutum_Grid = m_polar
    !%End
    use_symmetry = .false.
    if ((this%surf_shape == pes_cubic .or. this%surf_shape == pes_plane) &
      .and. this%kgrid == pes_cartesian .and. mdim == 3) use_symmetry = .true.
    if (this%surf_shape == pes_spherical .and. this%kgrid == pes_polar) use_symmetry = .true.
    call parse_variable(namespace, 'PES_Flux_UseSymmetries', use_symmetry, this%use_symmetry)
    call messages_print_var_value('PES_Flux_UseSymmetries', this%use_symmetry, namespace=namespace)
 
 
 
    ! get the values of the spherical harmonics for the surface points for PES_SPHERICAL
    if (this%surf_shape == pes_spherical) then
      safe_allocate(this%ylm_r(this%nsrfcpnts_start:this%nsrfcpnts_end, 0:this%lmax, -this%lmax:this%lmax))
      this%ylm_r = m_z0
 
      if (this%nsrfcpnts_start > 0) then
        do ll = 0, this%lmax
          do mm = -ll, ll
            do isp = this%nsrfcpnts_start, this%nsrfcpnts_end
              call ylmr_cmplx(this%rcoords(1:3, isp), ll, mm, this%ylm_r(isp, ll, mm))
              this%ylm_r(isp, ll, mm) = conjg(this%ylm_r(isp, ll, mm))
            end do
          end do
        end do
      end if
 
    else
 
      call pes_flux_integrate_cub_tabulate(this, space, mesh, st, kpoints)
 
    end if
 
 
    !%Variable PES_Flux_RuntimeOutput
    !%Type logical
    !%Section Time-Dependent::PhotoElectronSpectrum
    !%Description
    !% Write output in ascii format at runtime.
    !%
    !%End
    call parse_variable(namespace, 'PES_Flux_RuntimeOutput', .false., this%runtime_output)
    call messages_print_var_value('PES_Flux_RuntimeOutput', this%runtime_output, namespace=namespace)
 
 
 
    ! -----------------------------------------------------------------
    ! Options for time integration
    ! -----------------------------------------------------------------
    this%max_iter  = max_iter
    this%save_iter = save_iter
    this%itstep    = 0
    this%tdsteps = 1
 
    if (mesh%parallel_in_domains .and. bitand(this%par_strategy, option__pes_flux_parallelization__pf_time) /= 0) then
      this%tdsteps = mesh%mpi_grp%size
    end if
 
    ! -----------------------------------------------------------------
    ! Allocations
    ! -----------------------------------------------------------------
 
    safe_allocate(this%wf(0:this%nsrfcpnts, stst:stend, 1:sdim, kptst:kptend, 1:this%tdsteps))
    this%wf = m_z0
 
    safe_allocate(this%gwf(0:this%nsrfcpnts, stst:stend, 1:sdim, kptst:kptend, 1:this%tdsteps, 1:mdim))
    this%gwf = m_z0
 
    safe_allocate(this%veca(1:mdim, 1:this%tdsteps))
    this%veca = m_zero
 
    if (this%surf_shape == pes_spherical) then
      safe_allocate(this%spctramp_sph(1:this%nstepsomegak, stst:stend, 1:sdim, kptst:kptend, 1:this%nk))
      this%spctramp_sph = m_z0
 
      safe_allocate(this%conjgphase_prev(1:this%nstepsomegak, 1:this%nk))
 
    else
      safe_allocate(this%spctramp_cub(stst:stend, 1:sdim, kptst:kptend, 1:this%nkpnts))
      this%spctramp_cub = m_z0
 
      safe_allocate(this%conjgphase_prev(1:this%nkpnts, kptst:kptend))
 
    end if
 
    this%conjgphase_prev = m_z1
 
    call messages_write('Info: Total number of surface points = ')
    call messages_write(this%nsrfcpnts)
    call messages_info(namespace=namespace)
 
    call messages_write('Info: Total number of momentum points =  ')
    call messages_write(this%nkpnts)
    call messages_info(namespace=namespace)
 
    pop_sub_with_profile(pes_flux_init)
  end subroutine pes_flux_init
 
  ! ---------------------------------------------------------
  subroutine pes_flux_end(this)
    type(pes_flux_t), intent(inout) :: this
 
    push_sub(pes_flux_end)
 
    if (this%surf_shape == pes_spherical) then
      safe_deallocate_a(this%kcoords_sph)
      safe_deallocate_a(this%ylm_k)
      safe_deallocate_a(this%j_l)
      safe_deallocate_a(this%ylm_r)
      safe_deallocate_a(this%conjgphase_prev)
      safe_deallocate_a(this%spctramp_sph)
    else
      safe_deallocate_a(this%kcoords_cub)
      safe_deallocate_a(this%conjgphase_prev)
      safe_deallocate_a(this%spctramp_cub)
 
      if (.not. this%surf_interp) then
        safe_deallocate_a(this%srfcpnt)
      end if
      safe_deallocate_a(this%rankmin)
 
      safe_deallocate_a(this%face_idx_range)
      safe_deallocate_a(this%LLr)
      safe_deallocate_a(this%NN)
 
      safe_deallocate_a(this%expkr)
      safe_deallocate_a(this%expkr_perp)
      safe_deallocate_a(this%bvk_phase)
 
      safe_deallocate_a(this%Lkpuvz_inv)
 
    end if
 
    safe_deallocate_a(this%klinear)
 
    safe_deallocate_a(this%srfcnrml)
    safe_deallocate_a(this%rcoords)
 
    safe_deallocate_a(this%wf)
    safe_deallocate_a(this%gwf)
    safe_deallocate_a(this%veca)
 
    if(accel_buffer_is_allocated(this%pt_buff)) then
      call accel_release_buffer(this%pt_buff)
    end if
    if(accel_buffer_is_allocated(this%coords_buff)) then
      call accel_release_buffer(this%coords_buff)
    end if
    if(accel_buffer_is_allocated(this%interp_buff)) then
      call accel_release_buffer(this%interp_buff)
    end if
    if(accel_buffer_is_allocated(this%spacing_buff)) then
      call accel_release_buffer(this%spacing_buff)
    end if
 
    pop_sub(pes_flux_end)
  end subroutine pes_flux_end
 
  ! ---------------------------------------------------------
  subroutine pes_flux_reciprocal_mesh_gen(this, namespace, space, st, kpoints, comm, post)
    type(pes_flux_t),    intent(inout) :: this
    type(namespace_t),   intent(in)    :: namespace
    class(space_t),      intent(in)    :: space
    type(states_elec_t), intent(in)    :: st
    type(kpoints_t),     intent(in)    :: kpoints
    type(mpi_comm),      intent(in)    :: comm
    logical, optional,   intent(in)    :: post
 
    integer           :: mdim, pdim
    integer           :: kptst, kptend
    integer           :: ikp, ikpt
    integer           :: ll, mm, idim, idir, ispin, nspin
    integer           :: ikk, ith, iph, iomk,ie, ik1, ik2, ik3, kgrid_block_dim
    real(real64)      :: kmax(space%dim), kmin(space%dim), kact, thetak, phik
    type(block_t)     :: blk
 
    real(real64)      :: emin, emax, de , kvec(1:3), xx(3)
    integer           :: nkp_out, nkmin, nkmax
 
    integer             :: kgrid
    real(real64), allocatable  :: gpoints(:,:), gpoints_reduced(:,:)
    real(real64)        :: dk(1:3), kpoint(1:3), dthetak, dphik
    logical             :: use_enegy_grid, arpes_grid
 
    push_sub(pes_flux_reciprocal_mesh_gen)
 
    kptst  = st%d%kpt%start
    kptend = st%d%kpt%end
    mdim   = space%dim
    pdim   = space%periodic_dim
 
    this%dim  = mdim
    this%pdim = pdim
 
    !%Variable PES_Flux_Momenutum_Grid
    !%Type integer
    !%Section Time-Dependent::PhotoElectronSpectrum
    !%Description
    !% Decides how the grid in momentum space is generated.
    !%Option polar 1
    !% The grid is in polar coordinates with the zenith axis is along z.
    !% The grid parameters are defined by PES_Flux_Kmax, PES_Flux_DeltaK,
    !% PES_Flux_StepsThetaK, PES_Flux_StepsPhiK.
    !% This is the default choice for PES_Flux_Shape = sph or cub.
    !%Option cartesian 2
    !% The grid is in cartesian coordinates with parameters defined by
    !% PES_Flux_ARPES_grid, PES_Flux_EnergyGrid.
    !% This is the default choice for PES_Flux_Shape = sph or cub.
    !%End
 
    ! default values
    select case (this%surf_shape)
    case (pes_spherical)
      kgrid = pes_polar
    case (pes_plane)
      kgrid = pes_cartesian
    case (pes_cubic)
      kgrid = pes_cartesian
      if (mdim == 1)  kgrid = pes_polar
 
    case default
      assert(.false.)
 
    end select
 
    call parse_variable(namespace, 'PES_Flux_Momenutum_Grid', kgrid, this%kgrid)
    if (.not. varinfo_valid_option('PES_Flux_Momenutum_Grid', this%kgrid, is_flag = .true.)) then
      call messages_input_error(namespace,'PES_Flux_Momenutum_Grid')
    end if
    call messages_print_var_option('PES_Flux_Momenutum_Grid', this%kgrid, namespace=namespace)
 
 
 
    !Check availability of the calculation requested
    if (this%surf_shape == pes_spherical) then
      if (this%kgrid == pes_cartesian) then
        call messages_not_implemented('Cartesian momentum grid with a spherical surface', namespace=namespace)
      end if
      if (space%is_periodic()) then
        call messages_not_implemented('Spherical surface flux for periodic systems', namespace=namespace)
      end if
      if (mdim == 1) then
        call messages_not_implemented('Spherical surface flux for one-dimensional systems', namespace=namespace)
      end if
    end if
 
    if (this%surf_shape == pes_cubic) then
      if (space%is_periodic()) then
        call messages_not_implemented('Use of cubic surface for periodic systems (use pln)', namespace=namespace)
      end if
    end if
 
 
 
    !%Variable PES_Flux_EnergyGrid
    !%Type block
    !%Section Time-Dependent::PhotoElectronSpectrum
    !%Description
    !% The block <tt>PES_Flux_EnergyGrid</tt> specifies the energy grid
    !% in momentum space.
    !% <tt><br>%PES_Flux_EnergyGrid
    !% <br>&nbsp;&nbsp; Emin | Emax | DeltaE
    !% <br>%</tt>
    !%End
 
    emin = 0
    emax = 10
    de   = 0.1_real64
    use_enegy_grid = .false.
 
    if (parse_block(namespace, 'PES_Flux_EnergyGrid', blk) == 0) then
 
      call parse_block_float(blk, 0, 0, emin)
      call parse_block_float(blk, 0, 1, emax)
      call parse_block_float(blk, 0, 2, de)
 
      call parse_block_end(blk)
      use_enegy_grid = .true.
 
 
      call messages_write("Energy grid (Emin, Emax, DE) [")
      call messages_write(trim(units_abbrev(units_out%energy)))
      call messages_write("]:  (")
      call messages_write(units_from_atomic(units_out%energy, emin),fmt = '(f8.3)')
      call messages_write(", ")
      call messages_write(units_from_atomic(units_out%energy, emax), fmt = '(f8.3)')
      call messages_write(", ")
      call messages_write(units_from_atomic(units_out%energy, de), fmt = '(e9.2)')
      call messages_write(")")
      call messages_info(namespace=namespace)
 
      kmax(1:mdim) = sqrt(m_two*emax)
      kmin(1:mdim) = sqrt(m_two*emin)
      this%dk = sqrt(m_two*de)
 
    end if
 
    kgrid_block_dim = 1
    !ugly hack (since this input variable is properly handled below) but effective
    call parse_variable(namespace, 'PES_Flux_ARPES_grid', .false., this%arpes_grid)
    if (.not. use_enegy_grid .or. this%arpes_grid) then
 
      !%Variable PES_Flux_Kmax
      !%Type float
      !%Default 1.0
      !%Section Time-Dependent::PhotoElectronSpectrum
      !%Description
      !% The maximum value of the photoelectron momentum.
      !% For cartesian momentum grids one can specify a value different
      !% for cartesian direction using a block input.
      !%End
      if (parse_block(namespace, 'PES_Flux_Kmax', blk) == 0) then
        if (this%kgrid == pes_cartesian) then
          do idim = 1, mdim
            call parse_block_float(blk, 0, idim - 1, kmax(idim))
          end do
          kgrid_block_dim = mdim
        else
          message(1) = 'Wrong block format for PES_Flux_Kmax and non-cartesian grid'
          call messages_fatal(1, namespace=namespace)
        end if
      else
        call parse_variable(namespace, 'PES_Flux_Kmax', m_one, kmax(1))
        kmax(:)=kmax(1)
      end if
 
      !%Variable PES_Flux_Kmin
      !%Type float
      !%Default 0.0
      !%Section Time-Dependent::PhotoElectronSpectrum
      !%Description
      !% The minimum value of the photoelectron momentum.
      !% For cartesian momentum grids one can specify a value different
      !% for cartesian direction using a block input.
      !%End
      if (parse_block(namespace, 'PES_Flux_Kmin', blk) == 0) then
        if (this%kgrid == pes_cartesian) then
          do idim = 1, mdim
            call parse_block_float(blk, 0, idim - 1, kmin(idim))
          end do
          kgrid_block_dim = mdim
        else
          message(1) = 'Wrong block format for PES_Flux_Kmin and non-cartesian grid'
          call messages_fatal(1, namespace=namespace)
        end if
      else
        call parse_variable(namespace, 'PES_Flux_Kmin', m_zero, kmin(1))
        kmin(:)=kmin(1)
      end if
 
 
      !%Variable PES_Flux_DeltaK
      !%Type float
      !%Default 0.02
      !%Section Time-Dependent::PhotoElectronSpectrum
      !%Description
      !% Spacing of the the photoelectron momentum grid.
      !%End
      call parse_variable(namespace, 'PES_Flux_DeltaK', 0.02_real64, this%dk)
      if (this%dk <= m_zero) call messages_input_error(namespace,'PES_Flux_DeltaK')
 
    end if
 
    do idim = 1, kgrid_block_dim
      if (kgrid_block_dim == 1) then
        call messages_write("Momentum linear grid (Pmin, Pmax, DP) [")
      else
        call messages_write("Momentum linear grid (Pmin, Pmax, DP) "//index2axis(idim)//"-axis [")
      end if
      call messages_write(trim(units_abbrev(units_out%mass*units_out%velocity)))
      call messages_write("]:  (")
      call messages_write(kmin(idim),fmt = '(f8.3)')
      call messages_write(", ")
      call messages_write(kmax(idim), fmt = '(f8.3)')
      call messages_write(", ")
      call messages_write(this%dk, fmt = '(e9.2)')
      call messages_write(")")
      call messages_info(namespace=namespace)
    end do
 
 
 
!     if (this%surf_shape == PES_SPHERICAL .or. this%surf_shape == PES_CUBIC) then
    if (this%kgrid == pes_polar) then
      ! -----------------------------------------------------------------
      ! Setting up k-mesh
      ! 1D = 2 points, 2D = polar coordinates, 3D = spherical coordinates
      ! -----------------------------------------------------------------
 
 
      !%Variable PES_Flux_ThetaK
      !%Type block
      !%Section Time-Dependent::PhotoElectronSpectrum
      !%Description
      !% Define the grid points on theta (<math>0 \le \theta \le \pi</math>) when
      !% using a spherical grid in momentum.
      !% The block defines the maximum and minimum values of theta and the number of
      !% of points for the discretization.
      !%
      !% <tt>%PES_Flux_ThetaK
      !% <br>&nbsp;&nbsp; theta_min | theta_max  | npoints
      !% <br>%
      !% </tt>
      !%
      !% By default theta_min=0, theta_max = pi, npoints = 45.
      !%End
      this%nstepsthetak = 45
      this%thetak_rng(1:2) = (/m_zero, m_pi/)
      if (parse_block(namespace, 'PES_Flux_ThetaK', blk) == 0) then
        call parse_block_float(blk, 0, 0, this%thetak_rng(1))
        call parse_block_float(blk, 0, 1, this%thetak_rng(2))
        call parse_block_integer(blk, 0, 2, this%nstepsthetak)
        call parse_block_end(blk)
        do idim = 1,2
          if (this%thetak_rng(idim) < m_zero .or. this%thetak_rng(idim) > m_pi) then
            call messages_input_error(namespace,'PES_Flux_ThetaK')
          end if
        end do
        if (this%nstepsthetak < 0) call messages_input_error(namespace,'PES_Flux_ThetaK')
      else
 
        !%Variable PES_Flux_StepsThetaK
        !%Type integer
        !%Default 45
        !%Section Time-Dependent::PhotoElectronSpectrum
        !%Description
        !% Number of steps in <math>\theta</math> (<math>0 \le \theta \le \pi</math>) for the spherical grid in k.
        !%End
        call parse_variable(namespace, 'PES_Flux_StepsThetaK', 45, this%nstepsthetak)
        if (this%nstepsthetak < 0) call messages_input_error(namespace,'PES_Flux_StepsThetaK')
 
        ! should do this at some point
        !       call messages_obsolete_variable(namespace, 'PES_Flux_StepsThetaK')
      end if
 
 
 
      !%Variable PES_Flux_PhiK
      !%Type block
      !%Section Time-Dependent::PhotoElectronSpectrum
      !%Description
      !% Define the grid points on theta (<math>0 \le \theta \le 2\pi</math>) when
      !% using a spherical grid in momentum.
      !% The block defines the maximum and minimum values of theta and the number of
      !% of points for the discretization.
      !%
      !% <tt>%PES_Flux_PhiK
      !% <br>&nbsp;&nbsp; theta_min | theta_max  | npoints
      !% <br>%
      !% </tt>
      !%
      !% By default theta_min=0, theta_max = pi, npoints = 90.
      !%End
      this%nstepsphik = 90
      this%phik_rng(1:2) = (/m_zero, 2 * m_pi/)
      if (mdim == 1) this%phik_rng(1:2) = (/m_pi, m_zero/)
      if (parse_block(namespace, 'PES_Flux_PhiK', blk) == 0) then
        call parse_block_float(blk, 0, 0, this%phik_rng(1))
        call parse_block_float(blk, 0, 1, this%phik_rng(2))
        call parse_block_integer(blk, 0, 2, this%nstepsphik)
        call parse_block_end(blk)
        do idim = 1,2
          if (this%phik_rng(idim) < m_zero .or. this%phik_rng(idim) > m_two * m_pi) then
            call messages_input_error(namespace,'PES_Flux_PhiK')
          end if
        end do
        if (this%nstepsphik < 0) call messages_input_error(namespace,'PES_Flux_PhiK')
 
      else
 
        !%Variable PES_Flux_StepsPhiK
        !%Type integer
        !%Default 90
        !%Section Time-Dependent::PhotoElectronSpectrum
        !%Description
        !% Number of steps in <math>\phi</math> (<math>0 \le \phi \le 2 \pi</math>) for the spherical grid in k.
        !%End
        call parse_variable(namespace, 'PES_Flux_StepsPhiK', 90, this%nstepsphik)
        if (this%nstepsphik < 0) call messages_input_error(namespace,'PES_Flux_StepsPhiK')
        if (this%nstepsphik == 0) this%nstepsphik = 1
      end if
 
 
      dthetak = m_zero
      if (mdim == 3)  dthetak = abs(this%thetak_rng(2) - this%thetak_rng(1)) / (this%nstepsthetak)
      dphik = abs(this%phik_rng(2) - this%phik_rng(1)) / (this%nstepsphik)
 
 
      select case (mdim)
      case (1)
        this%nstepsthetak = 0
        this%nstepsphik   = 2
        this%nstepsomegak = this%nstepsphik
 
      case (2)
        this%nstepsthetak = 0
        this%nstepsomegak = this%nstepsphik
 
      case (3)
        if (this%nstepsthetak <= 1) then
          this%nstepsphik   = 1
          this%nstepsthetak = 1
        end if
        ! count the omegak samples
        iomk = 0
        do ith = 0, this%nstepsthetak
          thetak = ith * dthetak + this%thetak_rng(1)
          do iph = 0, this%nstepsphik - 1
            phik = iph * dphik + this%phik_rng(1)
            iomk = iomk + 1
            if (thetak < m_epsilon .or. abs(thetak-m_pi) < m_epsilon) exit
          end do
        end do
        this%nstepsomegak  = iomk
 
      case default
        assert(.false.)
 
      end select
 
      write(message(1),'(a)') "Polar momentum grid:"
      call messages_info(1, namespace=namespace)
      if (mdim == 3) then
        write(message(1),'(a,f12.6,a,f12.6,a, i6)') &
          "  Theta = (", this%thetak_rng(1), ",",this%thetak_rng(2), &
          "); n = ", this%nstepsthetak
        call messages_info(1, namespace=namespace)
      end if
      write(message(1),'(a,f12.6,a,f12.6,a, i6)') &
        "  Phi   = (", this%phik_rng(1), ",",this%phik_rng(2), &
        "); n = ", this%nstepsphik
      call messages_info(1, namespace=namespace)
 
 
 
      if (use_enegy_grid) then
        this%nk     = nint(abs(emax - emin) / de)
      else
        this%nk     = nint(abs(kmax(1) - kmin(1)) / this%dk)
      end if
      this%nkpnts = this%nstepsomegak * this%nk
 
      this%ll(1)      = this%nk
      this%ll(2)      = this%nstepsphik
      this%ll(3)      = this%nstepsthetak !- 1
      this%ll(mdim+1:3) = 1
 
      safe_allocate(this%klinear(1:this%nk, 1))
 
 
    else
      ! Cartesian
 
      dk(1:mdim) = m_one/kpoints%nik_axis(1:mdim)
 
      this%arpes_grid = .false.
      if (space%is_periodic()) then
        !%Variable PES_Flux_ARPES_grid
        !%Type logical
        !%Section Time-Dependent::PhotoElectronSpectrum
        !%Description
        !% Use a curvilinear momentum space grid that compensates the transformation
        !% used to obtain ARPES. With this choice ARPES data is laid out on a Cartesian
        !% regular grid.
        !% By default true when <tt>PES_Flux_Shape = pln</tt> and a <tt>KPointsPath</tt>
        !% is specified.
        !%End
        arpes_grid = kpoints%have_zero_weight_path()
        call parse_variable(namespace, 'PES_Flux_ARPES_grid', arpes_grid, this%arpes_grid)
        call messages_print_var_value("PES_Flux_ARPES_grid", this%arpes_grid, namespace=namespace)
      end if
 
 
 
 
      this%ll(:) = 1
 
      if (kpoints%have_zero_weight_path()) then
 
        if (this%arpes_grid) then
          nkmax = floor(emax / de)
          nkmin = floor(emin / de)
 
        else
          nkmax = floor(kmax(mdim) / this%dk)
          nkmin = floor(kmin(mdim) / this%dk)
 
        end if
 
        this%ll(mdim) = abs(nkmax - nkmin) + 1
 
        this%nk = this%ll(mdim)
 
 
      else
 
        if (.not. this%arpes_grid) then
          this%ll(1:mdim) = floor(abs(kmax(1:mdim) - kmin(1:mdim))/this%dk) + 1
          this%nk = maxval(this%ll(1:mdim))
 
        else
 
          nkmax = floor(emax / de)
          nkmin = floor(emin / de)
          this%nk = abs(nkmax - nkmin) + 1
 
          this%ll(1:pdim) = floor(abs(kmax(1:pdim) - kmin(1:pdim))/this%dk) + 1
          this%ll(mdim) = this%nk
 
        end if
 
        safe_allocate(this%klinear(1:maxval(this%ll(1:mdim)), 1:mdim))
        this%klinear = m_zero
 
      end if
 
      ! Total number of points
      this%nkpnts = product(this%ll(1:mdim))
 
 
    end if
 
 
 
 
 
    this%parallel_in_momentum = .false.
 
    ! Create the grid
    select case (this%kgrid)
 
    case (pes_polar)
 
      if (use_enegy_grid) then
        do ie = 1, this%nk
          this%klinear(ie, 1) = sqrt(m_two * (ie * de + emin))
        end do
      else
        do ikk = 1, this%nk
          this%klinear(ikk, 1) = ikk * this%dk + kmin(1)
        end do
      end if
 
 
      if (this%surf_shape == pes_spherical) then
 
        if (optional_default(post, .false.)) then
          pop_sub(pes_flux_reciprocal_mesh_gen)
          return
        end if
 
        ! we split the k-mesh in radial & angular part
        call pes_flux_distribute(1, this%nk, this%nk_start, this%nk_end, comm)
        if ((this%nk_end - this%nk_start + 1) < this%nk) this%parallel_in_momentum = .true.
        call pes_flux_distribute(1, this%nstepsomegak, this%nstepsomegak_start, this%nstepsomegak_end, comm)
 
!   Keep this because is useful for debug but not enough to bother having it polished out
!         if (debug%info) then
!           call mpi_world%barrier()
!           write(*,*) &
!             'Debug: momentum points on node ', mpi_world%rank, ' : ', this%nk_start, this%nk_end
!           call mpi_world%barrier()
!           write(*,*) &
!             'Debug: momentum directions on node ', mpi_world%rank, ' : ', this%nstepsomegak_start, this%nstepsomegak_end
!           call mpi_world%barrier()
!         end if
        safe_allocate(this%j_l(this%nk_start:this%nk_end, 0:this%lmax))
        this%j_l = m_zero
 
        safe_allocate(this%kcoords_sph(1:3, this%nk_start:this%nk_end, 1:this%nstepsomegak))
        this%kcoords_sph = m_zero
 
        safe_allocate(this%ylm_k(this%nstepsomegak_start:this%nstepsomegak_end, -this%lmax:this%lmax, 0:this%lmax))
        this%ylm_k = m_z0
 
        ! spherical harmonics & kcoords_sph
        iomk = 0
        do ith = 0, this%nstepsthetak
          thetak = ith * dthetak + this%thetak_rng(1)
          do iph = 0, this%nstepsphik - 1
            phik = iph * dphik + this%phik_rng(1)
            iomk = iomk + 1
            if (iomk >= this%nstepsomegak_start .and. iomk <= this%nstepsomegak_end) then
              xx(1) = cos(phik) * sin(thetak)
              xx(2) = sin(phik) * sin(thetak)
              xx(3) = cos(thetak)
              !$omp parallel do private(ll, mm)
              do ll = 0, this%lmax
                do mm = -ll, ll
                  call ylmr_cmplx(xx, ll, mm, this%ylm_k(iomk, mm, ll))
                end do
              end do
            end if
            if (this%nk_start > 0) then
              this%kcoords_sph(1, this%nk_start:this%nk_end, iomk) = cos(phik) * sin(thetak)
              this%kcoords_sph(2, this%nk_start:this%nk_end, iomk) = sin(phik) * sin(thetak)
              this%kcoords_sph(3, this%nk_start:this%nk_end, iomk) = cos(thetak)
            end if
            if (thetak < m_epsilon .or. abs(thetak-m_pi) < m_epsilon) exit
          end do
        end do
 
        if (this%nk_start > 0) then
          ! Bessel functions & kcoords_sph
          do ikk = this%nk_start, this%nk_end
            kact = this%klinear(ikk,1)
            do ll = 0, this%lmax
              this%j_l(ikk, ll) = loct_sph_bessel(ll, kact * this%radius) * &
                m_two * m_pi / (m_two * m_pi)**m_three / m_two
            end do
            this%kcoords_sph(:, ikk, :) = kact * this%kcoords_sph(:, ikk, :)
          end do
        end if
 
      else
        !planar or cubic surface
 
        ! no distribution
        this%nkpnts_start = 1
        this%nkpnts_end   = this%nkpnts
 
 
        safe_allocate(this%kcoords_cub(1:mdim, this%nkpnts_start:this%nkpnts_end, 1))
 
        this%kcoords_cub = m_zero
 
        if (mdim == 1) then
 
          ikp = 0
          do ikk = -this%nk, this%nk
            if (ikk == 0) cycle
            ikp = ikp + 1
            kact = sign(1,ikk) * this%klinear(abs(ikk), 1)
            this%kcoords_cub(1, ikp, 1) = kact
          end do
 
        else !mdim = 2,3
          thetak = m_pi / m_two
 
          ikp = 0
          do ikk = 1, this%nk
            do ith = 0, this%nstepsthetak
              if (mdim == 3) thetak = ith * dthetak + this%thetak_rng(1)
              do iph = 0, this%nstepsphik - 1
                ikp = ikp + 1
                phik = iph * dphik + this%phik_rng(1)
                kact = this%klinear(ikk,1)
                this%kcoords_cub(1, ikp,1) = kact * cos(phik) * sin(thetak)
                this%kcoords_cub(2, ikp,1) = kact * sin(phik) * sin(thetak)
                if (mdim == 3) this%kcoords_cub(3, ikp,1) = kact * cos(thetak)
                if (mdim == 3 .and. (thetak < m_epsilon .or. abs(thetak-m_pi) < m_epsilon)) exit
              end do
            end do
          end do
 
        end if
      end if
 
    case (pes_cartesian)
 
      this%nkpnts_start = 1
      this%nkpnts_end   = this%nkpnts
 
      if (kpoints%have_zero_weight_path()) then
        ! its a special case because we are generating a different (1D) grid for
        ! each kpoint and then we combine it in postprocessing
 
        ! Note that kptst:kptend includes spin
        safe_allocate(this%kcoords_cub(1:mdim, this%nkpnts_start:this%nkpnts_end, kptst:kptend))
        this%kcoords_cub = m_zero
 
        nkp_out = 0
 
        ! For the below code to work, we need to reinitialize ikp for each spin
        nspin = 1
        if (st%d%ispin == spin_polarized) nspin = 2
 
        do ispin = 1, nspin
          do ikpt = (kptst-1)+ispin, kptend, nspin
            ikp = 0
            do ikk = nkmin, nkmax
              kvec(1:mdim) = - kpoints%get_point(st%d%get_kpoint_index(ikpt))
              call fill_non_periodic_dimension(this)
            end do
          end do
        end do
 
      else
 
        safe_allocate(this%kcoords_cub(1:mdim, this%nkpnts_start:this%nkpnts_end, 1))
 
        safe_allocate(this%Lkpuvz_inv(1:this%ll(1), 1:this%ll(2), 1:this%ll(3)))
 
        do idir = 1, mdim
          do ikk = 1, this%ll(idir)
            this%klinear(ikk, idir) = ikk * this%dk + kmin(idir)
          end do
        end do
 
 
        if (.not. this%arpes_grid) then
          ! Normal velocity map
 
 
          ikpt = 1
          ikp = 0
          kvec(:) = m_zero
          do ik3 = 1, this%ll(3)
            if (mdim>2) kvec(3) = this%klinear(ik3, 3)
            do ik2 = 1, this%ll(2)
              if (mdim>1) kvec(2) = this%klinear(ik2, 2)
              do ik1 = 1, this%ll(1)
                ikp = ikp + 1
                kvec(1) = this%klinear(ik1, 1)
                this%kcoords_cub(1:mdim, ikp, ikpt) =  kvec(1:mdim)
                this%Lkpuvz_inv(ik1, ik2, ik3) = ikp
 
              end do
            end do
          end do
 
        else ! we want an ARPES-friendly grid layout
 
          nkp_out = 0
          ikpt = 1
          ikp = 0
          kvec(:) = m_zero
          do ikk = nkmin, nkmax ! this is going to be turned into energy
 
            ! loop over periodic directions
            select case (pdim)
            case (1)
 
              do ik1 = 1, this%ll(1)
                kvec(1) = this%klinear(ik1,1)
                kvec(1:pdim) = matmul(kpoints%latt%klattice_primitive(1:pdim, 1:pdim), kvec(1:pdim))
                call fill_non_periodic_dimension(this)
              end do
 
            case (2)
 
              do ik2 = 1, this%ll(2)
                do ik1 = 1, this%ll(1)
                  kvec(1:2) = (/this%klinear(ik1,1), this%klinear(ik2,2)/)
                  kvec(1:pdim) = matmul(kpoints%latt%klattice_primitive(1:pdim, 1:pdim), kvec(1:pdim))
                  call fill_non_periodic_dimension(this)
 
                  this%Lkpuvz_inv(ik1, ik2, ikk - nkmin + 1) = ikp
 
 
                end do
              end do
 
            case default
              assert(.false.)
 
            end select
 
          end do
 
 
        end if
 
      end if
 
!   Keep this because is useful for debug but not enough to bother having it polished out
!       if (debug%info .and. mpi_world%is_root()) then
!         ! this does not work for parallel in kpoint
!         ! you need to gather kcoords_pln
!         if (kpoints_have_zero_weight_path(sb%kpoints)) then
!           write(229,*) "#   ikpt (kpoint index),   ikp (momentum index),   this%kcoords_cub(1:mdim, ikp, ikpt)"
!           do ikpt = kptst, kptend
!             do ikp = 1, this%nkpnts
!               write(229,*) ikpt, ikp, this%kcoords_cub(1:mdim, ikp, ikpt)
!             end do
!           end do
!         else
!           write(229,*) "#   ikp (momentum index),   this%kcoords_cub(1:mdim, ikp, ikpt)"
!           do ikp = 1, this%nkpnts
!             write(229,*) ikp, this%kcoords_cub(1:mdim, ikp, 1)
!           end do
!         end if
!         flush(229)
!
!
!         if (mdim == 3) then
!           write(230,*) "#   ik1, ik2, ik3,  this%Lkpuvz_inv(ik1,ik2,ik3) "
!           do ik1 = 1, this%ll(1)
!             do ik2 = 1, this%ll(2)
!               do ik3 = 1, this%ll(3)
!                 write(230,*) ik1, ik2, ik3,  this%Lkpuvz_inv(ik1,ik2,ik3)
!               end do
!             end do
!           end do
!
!           flush(230)
!         end if
!
!       end if
 
      if (this%arpes_grid) then
        call messages_write("Number of points with E<p//^2/2 = ")
        call messages_write(nkp_out)
        call messages_write(" [of ")
        call messages_write(this%nkpnts*kpoints_number(kpoints))
        call messages_write("]")
        call messages_info(namespace=namespace)
      end if
 
      safe_deallocate_a(gpoints)
      safe_deallocate_a(gpoints_reduced)
 
    case default
      assert(.false.)
 
    end select
 
 
    !%Variable PES_Flux_GridTransformMatrix
    !%Type block
    !%Section Time-Dependent::PhotoElectronSpectrum
    !%Description
    !% Define an optional transformation matrix for the momentum grid.
    !%
    !% <tt>%PES_Flux_GridTransformMatrix
    !% <br>&nbsp;&nbsp; M_11 | M_12  | M_13
    !% <br>&nbsp;&nbsp; M_21 | M_22  | M_23
    !% <br>&nbsp;&nbsp; M_31 | M_32  | M_33
    !% <br>%
    !% </tt>
    !%End
    this%ktransf(:,:) = m_zero
    do idim = 1,mdim
      this%ktransf(idim,idim) = m_one
    end do
 
    if (parse_block(namespace, 'PES_Flux_GridTransformMatrix', blk) == 0) then
      do idim = 1,mdim
        do idir = 1, mdim
          call parse_block_float(blk, idir - 1, idim - 1, this%ktransf(idim, idir))
        end do
      end do
      call parse_block_end(blk)
 
      write(message(1),'(a)') 'Momentum grid transformation matrix :'
      do idir = 1, space%dim
        write(message(1 + idir),'(9f12.6)') ( this%ktransf(idim, idir), idim = 1, mdim)
      end do
      call messages_info(1 + mdim, namespace=namespace)
 
 
      !Apply the transformation
      if (this%surf_shape == pes_spherical) then
 
        iomk = 0
        do ith = 0, this%nstepsthetak
          do iph = 0, this%nstepsphik - 1
            iomk = iomk + 1
            do ikk = this%nk_start, this%nk_end
              kvec(1:mdim) = this%kcoords_sph(1:mdim, ikk, iomk)
              this%kcoords_sph(1:mdim, ikk, iomk) = matmul(this%ktransf(1:mdim, 1:mdim), kvec(1:mdim))
            end do
            if (ith == 0 .or. ith == this%nstepsthetak) exit
          end do
        end do
 
      else !planar or cubic surface
 
        do ikpt = kptst, kptend + 1
          if (ikpt == kptend + 1) then
            kpoint(1:space%dim) = m_zero
          else
            kpoint(1:space%dim) = kpoints%get_point(st%d%get_kpoint_index(ikpt))
          end if
 
          do ikp = 1, this%nkpnts
 
            kvec(1:mdim) = this%kcoords_cub(1:mdim, ikp, ikpt) - kpoint(1:mdim)
            this%kcoords_cub(1:mdim, ikp, ikpt) =  matmul(this%ktransf(1:mdim, 1:mdim), kvec(1:mdim)) &
              + kpoint(1:mdim)
          end do
        end do
 
      end if
 
 
    end if
 
 
 
 
    pop_sub(pes_flux_reciprocal_mesh_gen)
 
  contains
 
    ! Fill the non-periodic direction
    subroutine fill_non_periodic_dimension(this)
      type(pes_flux_t),   intent(inout) :: this
 
      integer :: sign
      real(real64)   :: kpar(1:pdim), val
 
      ikp = ikp + 1
 
      sign = 1
      if (ikk /= 0) sign= ikk / abs(ikk)
 
      kpar(1:pdim) = kvec(1:pdim)
      val = abs(ikk) * de * m_two - sum(kpar(1:pdim)**2)
      if (val >= 0) then
        kvec(mdim) =  sign * sqrt(val)
      else  ! if E < p//^2/2
        !FIXME: Should handle the exception doing something smarter than this
        kvec(mdim) = sqrt(val) ! set to NaN
        nkp_out = nkp_out + 1
      end if
 
      this%kcoords_cub(1:mdim, ikp, ikpt) =  kvec(1:mdim)
 
 
    end subroutine fill_non_periodic_dimension
 
 
 
  end subroutine pes_flux_reciprocal_mesh_gen
 
  ! ---------------------------------------------------------
  subroutine pes_flux_calc(this, space, mesh, st, der, ext_partners, kpoints, iter, dt, stopping)
    type(pes_flux_t),         intent(inout)            :: this
    class(space_t),           intent(in)               :: space
    type(mesh_t),             intent(in)               :: mesh
    type(states_elec_t),      intent(inout), target    :: st
    type(derivatives_t),      intent(in)               :: der
    type(partner_list_t),     intent(in)               :: ext_partners
    type(kpoints_t),          intent(in)               :: kpoints
    integer,                  intent(in)               :: iter
    real(real64),             intent(in)               :: dt
    logical,                  intent(in)               :: stopping
 
    integer            :: stst, stend, kptst, kptend, sdim, mdim
    integer            :: ist, ik, isdim, imdim
    integer            :: isp
    integer            :: il, ip_local
    integer            :: ib, minst, maxst, nst_linear, ist_linear
    integer            :: k_offset, n_boundary_points
    logical            :: contains_isp
 
    complex(real64), allocatable :: interp_values(:)
    complex(real64), allocatable :: interp_values_accel(:, :)
    complex(real64), allocatable :: wf_tmp_accel(:, :), gwf_tmp_accel(:, :, :)
    complex(real64), allocatable :: gpsi(:, :), psi(:)
 
    type(wfs_elec_t)          :: gpsib(space%dim)
    type(wfs_elec_t), pointer :: psib
 
    type(phase_t) :: phase
 
    type(lasers_t), pointer :: lasers
 
    if (iter == 0) return
 
    push_sub_with_profile(pes_flux_calc)
 
    stst   = st%st_start
    stend  = st%st_end
    kptst  = st%d%kpt%start
    kptend = st%d%kpt%end
    sdim   = st%d%dim
    mdim   = space%dim
 
    safe_allocate(psi(1:mesh%np_part))
    safe_allocate(gpsi(1:mesh%np_part, 1:mdim))
 
    if (.not. accel_is_enabled() .and. this%surf_interp) then
      safe_allocate(interp_values(1:this%nsrfcpnts))
    end if
 
    this%itstep = this%itstep + 1
 
    ! get and save current laser field
    lasers => list_get_lasers(ext_partners)
    if(associated(lasers)) then
      do il = 1, lasers%no_lasers
        call laser_field(lasers%lasers(il), this%veca(1:mdim, this%itstep), iter*dt)
      end do
    end if
    this%veca(:, this%itstep) = - this%veca(:, this%itstep)
 
!  Ideally one could directly access uniform_vector_potential
!       this%veca(:, this%itstep) = hm%hm_base%uniform_vector_potential(:)
 
    call phase%init(mesh, st%d%kpt, kpoints, st%d, space)
 
    do ik = kptst, kptend
      do ib = st%group%block_start, st%group%block_end
        minst = states_elec_block_min(st, ib)
        maxst = states_elec_block_max(st, ib)
 
        if (this%surf_shape == pes_plane) then
          safe_allocate(psib)
          call st%group%psib(ib, ik)%copy_to(psib)
 
          call phase%apply_to(mesh, mesh%np_part, conjugate=.false., psib=psib, src=st%group%psib(ib, ik), async=.true.)
 
          k_offset = psib%ik - kptst
          n_boundary_points = int(mesh%np_part - mesh%np)
          call boundaries_set(der%boundaries, mesh, psib, phase_correction=phase%phase_corr(:, psib%ik), &
            buff_phase_corr=phase%buff_phase_corr, offset=k_offset*n_boundary_points, async=.true.)
        else
          call boundaries_set(der%boundaries, mesh, st%group%psib(ib, ik))
          psib => st%group%psib(ib, ik)
        end if
 
        call zderivatives_batch_grad(der, psib, gpsib, set_bc=.false.)
        call accel_finish()
 
        if (accel_is_enabled()) then
 
          nst_linear = psib%nst_linear
 
          if(this%surf_interp) then
 
            safe_allocate(interp_values_accel(1:this%nsrfcpnts, 1:nst_linear))
 
            call mesh_interpolation_evaluate(this%interp, this%nsrfcpnts, nst_linear, psib, &
              this%rcoords(1:mdim, 1:this%nsrfcpnts), this%coords_buff, &
              interp_values_accel(1:this%nsrfcpnts, 1:nst_linear), this%interp_buff, this%spacing_buff, this%pt_buff)
 
            do ist = minst, maxst
              do isdim = 1, sdim
                ist_linear = psib%ist_idim_to_linear((/ist - minst + 1, isdim/))
                this%wf(1:this%nsrfcpnts, ist, isdim, ik, this%itstep) = &
                  st%occ(ist, ik) * interp_values_accel(1:this%nsrfcpnts, ist_linear)
              end do
            end do
 
            do imdim = 1, mdim
              call mesh_interpolation_evaluate(this%interp, this%nsrfcpnts, nst_linear, &
                gpsib(imdim), this%rcoords(1:mdim, 1:this%nsrfcpnts), this%coords_buff, &
                interp_values_accel(1:this%nsrfcpnts, 1:nst_linear), this%interp_buff, this%spacing_buff, this%pt_buff)
              do ist = minst, maxst
                do isdim = 1, sdim
                  ist_linear = psib%ist_idim_to_linear((/ist - minst + 1, isdim/))
                  this%gwf(1:this%nsrfcpnts, ist, isdim, ik, this%itstep, imdim) = &
                    st%occ(ist, ik) * interp_values_accel(1:this%nsrfcpnts, ist_linear)
                end do
              end do
            end do
 
            safe_deallocate_a(interp_values_accel)
 
          else
 
            call profiling_in(tostring(pes_flux_calc_cubic_reduction))
 
            safe_allocate(wf_tmp_accel(1:mesh%np_part, 1:nst_linear))
            safe_allocate(gwf_tmp_accel(1:mdim, 1:mesh%np_part, 1:nst_linear))
 
            call accel_read_buffer(psib%ff_device, mesh%np_part, nst_linear, wf_tmp_accel(:, :))
            do imdim = 1, mdim
              call accel_read_buffer(gpsib(imdim)%ff_device, mesh%np_part, nst_linear, gwf_tmp_accel(imdim, :, :))
            end do
 
            do ist = minst, maxst
              contains_isp = .true.
              do isp = 1, this%nsrfcpnts
                if (mesh%parallel_in_domains) then
                  contains_isp = mesh%mpi_grp%rank == this%rankmin(isp)
                end if
 
                if (.not. contains_isp) cycle
 
                ip_local = this%srfcpnt(isp)
                do isdim = 1, sdim
                  ist_linear = psib%ist_idim_to_linear((/ist - minst + 1, isdim/))
                  this%wf(isp, ist, isdim, ik, this%itstep) = &
                    st%occ(ist, ik) * wf_tmp_accel(ip_local, ist_linear)
                  this%gwf(isp, ist, isdim, ik, this%itstep, 1:mdim) = &
                    st%occ(ist, ik) * gwf_tmp_accel(1:mdim, ip_local, ist_linear)
                end do
              end do
 
              if (mesh%parallel_in_domains) then
                do isdim = 1, sdim
                  call mesh%allreduce(this%wf(1:this%nsrfcpnts, ist, isdim, ik, this%itstep))
                  do imdim = 1, mdim
                    call mesh%allreduce(this%gwf(1:this%nsrfcpnts, ist, isdim, ik, this%itstep, imdim))
                  end do
                end do
              end if
            end do
 
            safe_deallocate_a(wf_tmp_accel)
            safe_deallocate_a(gwf_tmp_accel)
 
            call profiling_out(tostring(pes_flux_calc_cubic_reduction))
 
          end if
 
        else
 
          do ist = minst, maxst
            do isdim = 1, sdim
 
              call batch_get_state(psib, (/ist, isdim/), mesh%np_part, psi(:))
 
              do imdim = 1, mdim
                call batch_get_state(gpsib(imdim), (/ist, isdim/), mesh%np_part, gpsi(:, imdim))
              end do
 
              if (this%surf_interp) then
 
                call mesh_interpolation_evaluate(this%interp, this%nsrfcpnts, psi(1:mesh%np_part), &
                  this%rcoords(1:mdim, 1:this%nsrfcpnts), interp_values(1:this%nsrfcpnts))
                this%wf(1:this%nsrfcpnts, ist, isdim, ik, this%itstep) = st%occ(ist, ik) * interp_values(1:this%nsrfcpnts)
                do imdim = 1, mdim
                  call mesh_interpolation_evaluate(this%interp, this%nsrfcpnts, gpsi(1:mesh%np_part, imdim), &
                    this%rcoords(1:mdim, 1:this%nsrfcpnts), interp_values(1:this%nsrfcpnts))
                  this%gwf(1:this%nsrfcpnts, ist, isdim, ik, this%itstep, imdim) = &
                    st%occ(ist, ik) * interp_values(1:this%nsrfcpnts)
                end do
 
              else
 
                contains_isp = .true.
                do isp = 1, this%nsrfcpnts
                  if (mesh%parallel_in_domains) then
                    contains_isp = mesh%mpi_grp%rank == this%rankmin(isp)
                  end if
 
                  if (.not. contains_isp) cycle
 
                  ip_local = this%srfcpnt(isp)
                  this%wf(isp, ist, isdim, ik, this%itstep) = st%occ(ist, ik) * psi(ip_local)
                  this%gwf(isp, ist, isdim, ik, this%itstep, 1:mdim) = &
                    st%occ(ist, ik) * gpsi(ip_local, 1:mdim)
                end do
 
                if (mesh%parallel_in_domains) then
                  call mesh%allreduce(this%wf(1:this%nsrfcpnts, ist, isdim, ik, this%itstep))
                  do imdim = 1, mdim
                    call mesh%allreduce(this%gwf(1:this%nsrfcpnts, ist, isdim, ik, this%itstep, imdim))
                  end do
                end if
 
              end if
            end do
          end do
 
        end if
 
        do imdim = 1, mdim
          call gpsib(imdim)%end(copy = .false.)
        end do
 
        if (this%surf_shape == pes_plane) then
          call psib%end(copy = .false.)
          safe_deallocate_p(psib)
        end if
      end do
    end do
 
    if (this%surf_interp) then
      safe_deallocate_a(interp_values)
    end if
 
    safe_deallocate_a(psi)
    safe_deallocate_a(gpsi)
 
    call phase%end()
 
    if (this%itstep == this%tdsteps .or. mod(iter, this%save_iter) == 0 .or. iter == this%max_iter .or. stopping) then
      if (this%surf_shape == pes_cubic .or. this%surf_shape == pes_plane) then
        call pes_flux_integrate_cub(this, space, mesh, st, kpoints, dt)
      else
        call pes_flux_integrate_sph(this, mesh, st, dt)
      end if
 
      ! clean up fields
      this%wf  = m_z0
      this%gwf  = m_z0
      this%veca = m_zero
      this%itstep = 0
    end if
 
    pop_sub_with_profile(pes_flux_calc)
  end subroutine pes_flux_calc
 
 
 
  ! ---------------------------------------------------------
  subroutine pes_flux_integrate_cub_tabulate(this, space, mesh, st, kpoints)
    type(pes_flux_t),     intent(inout) :: this
    class(space_t),       intent(in)    :: space
    type(mesh_t),         intent(in)    :: mesh
    type(states_elec_t),  intent(in)    :: st
    type(kpoints_t),      intent(in)    :: kpoints
 
    integer            :: kptst, kptend,  mdim
    integer            :: ik, j1, j2, jvec(1:2)
    integer            :: isp, ikp, ikp_start, ikp_end
    integer            :: ik_map
 
    integer            :: idir, pdim, nfaces, ifc, n_dir
    complex(real64)    :: tmp
    real(real64)       :: kpoint(1:3), vec(1:3), lvec(1:3)
 
    push_sub(pes_flux_integrate_cub_tabulate)
 
    if (kpoints%have_zero_weight_path()) then
      kptst     = st%d%kpt%start
      kptend    = st%d%kpt%end
    else
      kptst     = 1
      kptend    = 1
    end if
 
    mdim = space%dim
    pdim = space%periodic_dim
 
    ikp_start = this%nkpnts_start
    ikp_end   = this%nkpnts_end
 
 
    if (this%surf_shape == pes_plane) then
      ! This is not general but should work in the specific case where it is relevant
      !i.e. when the system is semiperiodic in <=2 dimensions
      this%srfcnrml(:, 1:this%nsrfcpnts) = this%srfcnrml(:, 1:this%nsrfcpnts) * mesh%coord_system%surface_element(3)
    end if
 
    if (.not. this%use_symmetry .or. kpoints%have_zero_weight_path()) then
 
      safe_allocate(this%expkr(1:this%nsrfcpnts,ikp_start:ikp_end,kptst:kptend,1))
 
      do ik = kptst, kptend
        do ikp = ikp_start, ikp_end
          do isp = 1, this%nsrfcpnts
            this%expkr(isp,ikp,ik,1) = exp(-m_zi * sum(this%rcoords(1:mdim,isp) &
              * this%kcoords_cub(1:mdim, ikp, ik))) &
              * (m_two * m_pi)**(-mdim / m_two)
 
          end do
        end do
      end do
 
 
 
    else !do something smart to exploit symmetries
 
      nfaces = mdim*2
      if (this%surf_shape == pes_plane) nfaces = 1
 
 
      safe_allocate(this%expkr(1:this%nsrfcpnts,maxval(this%ll(1:mdim)),kptst:kptend,1:mdim))
      this%expkr(:,:,:,:) = m_z1
 
      do ik = kptst, kptend !should only be ik=1
        do idir = 1, mdim
          do ikp = 1, this%ll(idir)
            do isp = 1, this%nsrfcpnts
              this%expkr(isp,ikp,ik,idir) = exp(-m_zi * this%rcoords(idir,isp) &
                * this%klinear(ikp, idir)) &
                * (m_two * m_pi)**(-m_one / m_two)
 
            end do
          end do
        end do
      end do
 
      safe_allocate(this%expkr_perp(1:maxval(this%ll(1:mdim)), 1:nfaces))
      this%expkr_perp(:,:) = m_z1
 
      do ifc = 1, nfaces
        if (this%face_idx_range(ifc, 1) < 0) cycle ! this face have no local surface point
        isp = this%face_idx_range(ifc, 1)
        do idir = 1, mdim
          if (abs(this%srfcnrml(idir, isp)) >= m_epsilon) n_dir = idir
        end do
 
        do ikp = 1, this%ll(n_dir)
          this%expkr_perp(ikp, ifc) = exp(- m_zi * this%rcoords(n_dir,isp) &
            * this%klinear(ikp, n_dir)) &
            * (m_two * m_pi)**(-m_one / m_two)
 
 
        end do
      end do
 
    end if
 
 
    if (space%is_periodic()) then
      !Tabulate the Born-von Karman phase
      safe_allocate(this%bvk_phase(ikp_start:ikp_end,st%d%kpt%start:st%d%kpt%end))
 
      this%bvk_phase(:,:) = m_z0
      vec(:) = m_zero
 
      do ik = st%d%kpt%start, st%d%kpt%end
        kpoint(1:mdim) = kpoints%get_point(st%d%get_kpoint_index(ik))
        if (kpoints%have_zero_weight_path()) then
          ik_map = ik
        else
          ik_map = 1
        end if
        do ikp = ikp_start, ikp_end
          vec(1:pdim) = this%kcoords_cub(1:pdim, ikp, ik_map) + kpoint(1:pdim)
          do j1 = 0, kpoints%nik_axis(1) - 1
            do j2 = 0, kpoints%nik_axis(2) - 1
              jvec(1:2)=(/j1, j2/)
              lvec(1:pdim) = matmul(kpoints%latt%rlattice(1:pdim, 1:2), jvec(1:2))
              tmp = sum(lvec(1:pdim) * vec(1:pdim))
              this%bvk_phase(ikp, ik) =  this%bvk_phase(ikp,ik) &
                +  exp(m_zi * tmp)
 
            end do
          end do
 
        end do
      end do
      this%bvk_phase(:,:) = this%bvk_phase(:,:) * m_one / product(kpoints%nik_axis(1:pdim))
 
 
!   Keep this because is useful for debug but not enough to bother having it polished out
!       if (debug%info .and. mpi_world%is_root()) then
!         write(225,*) "ik, ikp, this%bvk_phase(ikp,ik)"
!         do ik = st%d%kpt%start, st%d%kpt%end
!           do ikp = ikp_start, ikp_end
!             write(225,*) ik, ikp, this%bvk_phase(ikp,ik)
!           end do
!         end do
!         flush(225)
!       end if
 
    end if
 
    pop_sub(pes_flux_integrate_cub_tabulate)
 
  end subroutine pes_flux_integrate_cub_tabulate
 
 
  ! ---------------------------------------------------------
  pure function get_ikp(this, ikpu, ikpv, ikpz, n_dir) result(ikp)
    type(pes_flux_t),        intent(in)  :: this
    integer,                 intent(in)  :: ikpu
    integer,                 intent(in)  :: ikpv
    integer,                 intent(in)  :: ikpz
    integer,                 intent(in)  :: n_dir
    integer                              :: ikp
 
    select case (n_dir)
    case (1)
      ikp = this%Lkpuvz_inv(ikpz, ikpu, ikpv)
    case (2)
      ikp = this%Lkpuvz_inv(ikpu, ikpz, ikpv)
    case (3)
      ikp = this%Lkpuvz_inv(ikpu, ikpv, ikpz)
 
    case default
      ! should die here but cannot use assert in a pure function
 
    end select
 
  end function get_ikp
 
  ! ---------------------------------------------------------
  subroutine pes_flux_distribute_facepnts_cub(this, mesh)
    type(pes_flux_t), intent(inout) :: this
    type(mesh_t),        intent(in) :: mesh
 
 
    integer :: mdim, nfaces, ifc, ifp_start, ifp_end
 
    push_sub(pes_flux_distribute_facepnts_cub)
 
    mdim      = mesh%box%dim
 
    nfaces = mdim*2
    if (this%surf_shape == pes_plane) nfaces = 1
 
    do ifc = 1, nfaces
 
      ifp_start = this%face_idx_range(ifc, 1)
      ifp_end   = this%face_idx_range(ifc, 2)
 
 
      if (this%nsrfcpnts_start <= ifp_end) then ! the local domain could be in this face
 
        if (this%nsrfcpnts_start >= ifp_start) then
          if (this%nsrfcpnts_start <= ifp_end) then
            this%face_idx_range(ifc, 1) = this%nsrfcpnts_start
          else
            this%face_idx_range(ifc, 1:2) = -1  ! the local domain is not in this face
          end if
        end if
 
        if (this%nsrfcpnts_end <= ifp_end) then
          if (this%nsrfcpnts_end >= ifp_start) then
            this%face_idx_range(ifc, 2) = this%nsrfcpnts_end
          else
            this%face_idx_range(ifc, 1:2) = -1  ! the local domain is not in this face
          end if
        end if
 
      else
 
        this%face_idx_range(ifc, 1:2) = -1  ! the local domain is not in this face
 
      end if
 
!   Keep this because is useful for debug but not enough to bother having it polished out
!       if (debug%info) then
!         print *, mpi_world%rank, ifc, ifp_start, ifp_end, this%face_idx_range(ifc, 1:2), this%nsrfcpnts_start, this%nsrfcpnts_end
!       end if
 
    end do
 
 
    pop_sub(pes_flux_distribute_facepnts_cub)
  end subroutine pes_flux_distribute_facepnts_cub
 
 
  ! ---------------------------------------------------------
  subroutine pes_flux_integrate_cub(this, space, mesh, st, kpoints, dt)
    type(pes_flux_t),    intent(inout) :: this
    class(space_t),      intent(in)    :: space
    type(mesh_t),        intent(in)    :: mesh
    type(states_elec_t), intent(inout) :: st
    type(kpoints_t),     intent(in)    :: kpoints
    real(real64),        intent(in)    :: dt
 
    integer            :: stst, stend, kptst, kptend, sdim, mdim
    integer            :: ist, ik, isdim, imdim, ik_map
    integer            :: ikp, itstep
    integer            :: idir, n_dir, nfaces
    complex(real64), allocatable :: Jk_cub(:,:,:,:), spctramp_cub(:,:,:,:)
    integer            :: ikp_start, ikp_end, isp_start, isp_end
    real(real64)       :: vec, kpoint(1:3)
    integer            :: ifc
    complex(real64), allocatable :: wfpw(:), gwfpw(:)
    complex(real64), allocatable :: phase(:,:),vphase(:,:)
 
    integer            :: tdstep_on_node
    integer            :: nfp
 
    !Symmetry helpers
    integer            :: ikpu, ikpv, ikpz, dir_on_face(1:2)
    complex(real64)    :: face_int_gwf, face_int_wf
 
 
    push_sub(pes_flux_integrate_cub)
 
    ! this routine is parallelized over time steps and surface points
 
    call profiling_in('pes_flux_integrate_cub')
 
    call messages_write("Debug: calculating pes_flux cub surface integral (accelerated direct expression)")
    call messages_info(debug_only=.true.)
 
 
    tdstep_on_node = 1
    if (bitand(this%par_strategy, option__pes_flux_parallelization__pf_time) /= 0) then
      if (mesh%parallel_in_domains) tdstep_on_node = mesh%mpi_grp%rank + 1
    end if
 
    stst      = st%st_start
    stend     = st%st_end
    kptst     = st%d%kpt%start
    kptend    = st%d%kpt%end
    sdim      = st%d%dim
    mdim      = mesh%box%dim
 
    ikp_start = this%nkpnts_start
    ikp_end   = this%nkpnts_end
 
 
 
    safe_allocate(jk_cub(stst:stend, 1:sdim, kptst:kptend, ikp_start:ikp_end))
    safe_allocate(spctramp_cub(stst:stend, 1:sdim, kptst:kptend, ikp_start:ikp_end))
    spctramp_cub = m_z0
 
 
    safe_allocate(vphase(ikp_start:ikp_end, kptst:kptend))
    safe_allocate(phase(ikp_start:ikp_end, kptst:kptend))
    vphase(:,:) = m_z0
    phase(:,:) = m_z0
 
 
    nfaces = mdim * 2
    if (this%surf_shape == pes_plane) nfaces = 1
 
    safe_allocate( wfpw(ikp_start:ikp_end))
    safe_allocate(gwfpw(ikp_start:ikp_end))
 
 
    do ifc = 1, nfaces
 
      if (this%face_idx_range(ifc, 1)<0) cycle ! this face have no local surface point
 
 
      isp_start = this%face_idx_range(ifc, 1)
      isp_end   = this%face_idx_range(ifc, 2)
 
 
      nfp = isp_end - isp_start + 1 ! faces can have a different number of points
 
      wfpw = m_z0
      gwfpw = m_z0
 
      ! get the directions normal to the surface and parallel to it
      imdim = 1
      do idir = 1, mdim
        if (abs(this%srfcnrml(idir, isp_start)) >= m_epsilon) then
          n_dir = idir
        else
          dir_on_face(imdim)=idir
          imdim = imdim + 1
        end if
      end do
 
 
      ! get the previous Volkov phase
      vphase(ikp_start:ikp_end,:) = this%conjgphase_prev(ikp_start:ikp_end,:)
 
      jk_cub(:, :, :, :) = m_z0
 
      do itstep = 1, this%itstep
 
        do ik = kptst, kptend
 
          if (kpoints%have_zero_weight_path()) then
            ik_map = ik
          else
            ik_map = 1
          end if
 
          kpoint(1:mdim) = kpoints%get_point(st%d%get_kpoint_index(ik))
          do ikp = ikp_start, ikp_end
            vec = sum((this%kcoords_cub(1:mdim, ikp, ik_map) - this%veca(1:mdim, itstep) / p_c)**2)
            vphase(ikp, ik) = vphase(ikp, ik) * exp(m_zi * vec * dt / m_two)
 
 
            if (space%is_periodic()) then
              phase(ikp, ik)  = vphase(ikp, ik) *  this%bvk_phase(ikp,ik)
            else
              phase(ikp, ik)  = vphase(ikp, ik)
            end if
 
          end do
 
          if (itstep /= tdstep_on_node) cycle
 
          do ist = stst, stend
            do isdim = 1, sdim
 
 
              ! calculate the surface integrals
              if (.not. this%use_symmetry .or. kpoints%have_zero_weight_path()) then
 
                do ikp = ikp_start , ikp_end
 
                  gwfpw(ikp) = &
                    sum(this%gwf(isp_start:isp_end, ist, isdim, ik, itstep, n_dir) &
                    * this%expkr(isp_start:isp_end, ikp, ik_map, 1)              &
                    * this%srfcnrml(n_dir, isp_start:isp_end))
 
 
                  wfpw(ikp) = &
                    sum(this%wf(isp_start:isp_end, ist, isdim, ik, itstep)        &
                    * this%expkr(isp_start:isp_end, ikp,ik_map, 1)              &
                    * this%srfcnrml(n_dir, isp_start:isp_end))
                end do
 
              else
                !$omp parallel do private (ikpv,ikpz,face_int_gwf,face_int_wf) shared(gwfpw, wfpw)
                do ikpu = 1, this%ll(dir_on_face(1))
                  do ikpv = 1, this%ll(dir_on_face(2))
 
 
                    face_int_gwf = sum(this%gwf(isp_start:isp_end, ist, isdim, ik, itstep, n_dir) &
                      * this%expkr(isp_start:isp_end, ikpu, ik_map, dir_on_face(1)) &
                      * this%expkr(isp_start:isp_end, ikpv, ik_map, dir_on_face(2)) &
                      * this%srfcnrml(n_dir, isp_start:isp_end))
 
                    face_int_wf  = sum(this%wf(isp_start:isp_end, ist, isdim, ik, itstep) &
                      * this%expkr(isp_start:isp_end, ikpu, ik_map, dir_on_face(1)) &
                      * this%expkr(isp_start:isp_end, ikpv, ik_map, dir_on_face(2)) &
                      * this%srfcnrml(n_dir, isp_start:isp_end))
 
                    do ikpz = 1, this%ll(n_dir)
 
                      gwfpw(get_ikp(this, ikpu, ikpv, ikpz, n_dir)) = face_int_gwf &
                        * this%expkr_perp(ikpz, ifc)
                      wfpw( get_ikp(this, ikpu, ikpv, ikpz, n_dir)) = face_int_wf &
                        * this%expkr_perp(ikpz, ifc)
 
                    end do
 
                  end do
                end do
 
 
              end if
 
 
              ! Sum it up
              jk_cub(ist, isdim, ik, ikp_start:ikp_end) = jk_cub(ist, isdim, ik, ikp_start:ikp_end) +  &
                phase(ikp_start:ikp_end, ik) * (wfpw(ikp_start:ikp_end) * &
                (m_two * this%veca(n_dir, itstep) / p_c  - this%kcoords_cub(n_dir, ikp_start:ikp_end, ik_map)) + &
                m_zi * gwfpw(ikp_start:ikp_end))
 
 
            end do ! isdim
          end do ! ist
        end do ! is
 
      end do !istep
 
 
      spctramp_cub(:,:,:,:) = spctramp_cub(:,:,:,:) + jk_cub(:,:,:,:) * m_half
 
 
    end do ! face loop
 
 
    this%conjgphase_prev(ikp_start:ikp_end, :) = vphase(ikp_start:ikp_end, :)
 
 
    if (mesh%parallel_in_domains .and.(bitand(this%par_strategy, option__pes_flux_parallelization__pf_time) /= 0 &
      .or. bitand(this%par_strategy, option__pes_flux_parallelization__pf_surface) /= 0)) then
      call mesh%allreduce(spctramp_cub)
    end if
 
 
 
    this%spctramp_cub = this%spctramp_cub + spctramp_cub * dt
 
    safe_deallocate_a(gwfpw)
    safe_deallocate_a(wfpw)
 
    safe_deallocate_a(jk_cub)
    safe_deallocate_a(spctramp_cub)
    safe_deallocate_a(phase)
    safe_deallocate_a(vphase)
 
    call profiling_out('pes_flux_integrate_cub')
 
    pop_sub(pes_flux_integrate_cub)
  end subroutine pes_flux_integrate_cub
 
 
 
  ! ---------------------------------------------------------
  subroutine pes_flux_integrate_sph(this, mesh, st, dt)
    type(pes_flux_t),    intent(inout) :: this
    type(mesh_t),        intent(in)    :: mesh
    type(states_elec_t), intent(inout) :: st
    real(real64),        intent(in)    :: dt
 
    integer            :: stst, stend, kptst, kptend, sdim
    integer            :: ist, ik, isdim, imdim
    integer            :: itstep, lmax
    integer            :: ll, mm, ip
    complex(real64), allocatable :: s1_node(:,:,:,:,:,:)
    complex(real64), allocatable :: s2_node(:,:,:,:,:)
    complex(real64), allocatable :: integ11_t(:,:,:), integ12_t(:,:,:)
    complex(real64), allocatable :: integ21_t(:,:), integ22_t(:,:)
    complex(real64), allocatable :: spctramp_sph(:,:,:,:,:)
    integer            :: ikk, ikk_start, ikk_end
    integer            :: isp_start, isp_end
    integer            :: iomk, iomk_start, iomk_end
    complex(real64), allocatable :: phase_act(:,:)
    real(real64)       :: vec
    complex(real64)    :: weight_x, weight_y, weight_z
    integer            :: tdstep_on_node
 
    push_sub_with_profile(pes_flux_integrate_sph)
 
    ! this routine is parallelized over time steps and surface points
 
    call messages_write("Debug: calculating pes_flux sph surface integral")
    call messages_info(debug_only=.true.)
 
    tdstep_on_node = 1
    if (mesh%parallel_in_domains) tdstep_on_node = mesh%mpi_grp%rank + 1
 
    stst   = st%st_start
    stend  = st%st_end
    kptst  = st%d%kpt%start
    kptend = st%d%kpt%end
    sdim   = st%d%dim
 
    lmax       = this%lmax
    ikk_start  = this%nk_start
    ikk_end    = this%nk_end
    isp_start  = this%nsrfcpnts_start
    isp_end    = this%nsrfcpnts_end
    iomk_start = this%nstepsomegak_start
    iomk_end   = this%nstepsomegak_end
 
    ! spectral amplitude for k-points on node
    safe_allocate(integ11_t(1:this%nstepsomegak, 1:3, 0:lmax))
    safe_allocate(integ21_t(1:this%nstepsomegak, 0:lmax))
    safe_allocate(integ12_t(1:this%nstepsomegak, 1:3, ikk_start:ikk_end))
    safe_allocate(integ22_t(1:this%nstepsomegak, ikk_start:ikk_end))
 
    safe_allocate(spctramp_sph(1:this%nstepsomegak, stst:stend, 1:sdim, kptst:kptend, 1:this%nk))
    spctramp_sph = m_z0
 
    ! get the previous Volkov phase
    call profiling_in("PES_FLU_SPH_PREVOLKOV")
    safe_allocate(phase_act(1:this%nstepsomegak, ikk_start:ikk_end))
    if (ikk_start > 0) then
      phase_act(:, ikk_start:ikk_end) = this%conjgphase_prev(:, ikk_start:ikk_end)
    else
      phase_act(:, ikk_start:ikk_end) = m_z0
    end if
    call profiling_out("PES_FLU_SPH_PREVOLKOV")
 
    !TODO: The storage (-lmax:lmax, 0:lmax) is not efficient for large l
 
    ! surface integral S_lm (for time step on node)
    safe_allocate(s1_node(1:3, -lmax:lmax, 0:lmax, stst:stend, 1:sdim, kptst:kptend))
    safe_allocate(s2_node(-lmax:lmax, 0:lmax, stst:stend, 1:sdim, kptst:kptend))
 
    do itstep = 1, this%itstep
      call profiling_in("PES_FLU_SPH_INT")
 
      !$omp parallel do private(ist, isdim, ll, mm, ip, weight_x, weight_y, weight_z) collapse(4)
      do ik = kptst, kptend
        do ist = stst, stend
          do isdim = 1, sdim
            do ll = 0, lmax
              do mm = -ll, ll
                s1_node(1:3, mm, ll, ist, isdim, ik) = m_zero
                s2_node(mm, ll, ist, isdim, ik) = m_zero
                ! surface integrals
                do ip = isp_start, isp_end
                  weight_x = this%ylm_r(ip, ll, mm) * this%srfcnrml(1, ip)
                  weight_y = this%ylm_r(ip, ll, mm) * this%srfcnrml(2, ip)
                  weight_z = this%ylm_r(ip, ll, mm) * this%srfcnrml(3, ip)
                  s1_node(1, mm, ll, ist, isdim, ik) = s1_node(1, mm, ll, ist, isdim, ik) + &
                    weight_x * this%wf(ip, ist, isdim, ik, itstep)
                  s1_node(2, mm, ll, ist, isdim, ik) = s1_node(2, mm, ll, ist, isdim, ik) + &
                    weight_y * this%wf(ip, ist, isdim, ik, itstep)
                  s1_node(3, mm, ll, ist, isdim, ik) = s1_node(3, mm, ll, ist, isdim, ik) + &
                    weight_z * this%wf(ip, ist, isdim, ik, itstep)
 
                  s2_node(mm, ll, ist, isdim, ik) = s2_node(mm, ll, ist, isdim, ik) &
                    + weight_x * this%gwf(ip, ist, isdim, ik, itstep, 1) &
                    + weight_y * this%gwf(ip, ist, isdim, ik, itstep, 2) &
                    + weight_z * this%gwf(ip, ist, isdim, ik, itstep, 3)
                end do
              end do
            end do
          end do
        end do
      end do
      !$omp end parallel do
      if (mesh%parallel_in_domains .and. bitand(this%par_strategy, option__pes_flux_parallelization__pf_surface) /= 0) then
        call mesh%allreduce(s1_node)
        call mesh%allreduce(s2_node)
      end if
      call profiling_out("PES_FLU_SPH_INT")
 
      call profiling_in("PES_FLU_SPH_VOLKOV")
      ! calculate the current Volkov phase
      !$omp parallel do private(iomk, vec)
      do ikk = ikk_start, ikk_end
        do iomk = 1, this%nstepsomegak
          vec = sum((this%kcoords_sph(1:3, ikk, iomk) - this%veca(1:3, itstep) / p_c)**2)
          phase_act(iomk, ikk) = phase_act(iomk, ikk) * exp(m_zi * vec * dt * m_half)
        end do
      end do
      !$omp end parallel do
      call profiling_out("PES_FLU_SPH_VOLKOV")
 
      do ik = kptst, kptend
        do ist = stst, stend
          do isdim = 1, sdim
            integ12_t = m_z0
            integ22_t = m_z0
            integ11_t = m_z0
            integ21_t = m_z0
 
            call profiling_in("PES_FLU_SPH_INTEGXX")
 
            !$omp parallel do private(mm, imdim, ikk)
            do ll = 0, lmax
              do mm = -ll, ll
                ! multiply with spherical harmonics & sum over all mm
                do imdim = 1, 3
                  !$omp simd
                  do ikk = iomk_start, iomk_end
                    integ11_t(ikk, imdim, ll) = integ11_t(ikk, imdim, ll) + &
                      s1_node(imdim, mm, ll, ist, isdim, ik) * this%ylm_k(ikk, mm, ll)
                  end do
                end do
                !$omp simd
                do ikk = iomk_start, iomk_end
                  integ21_t(ikk, ll) = integ21_t(ikk, ll) + &
                    s2_node(mm, ll, ist, isdim, ik) * this%ylm_k(ikk, mm, ll)
                end do
              end do
            end do
            !$omp end parallel do
 
            call mesh%allreduce(integ11_t)
            call mesh%allreduce(integ21_t)
 
            !$omp parallel private(weight_x, imdim, ikk)
            do ll = 0, lmax
              weight_x = ( - m_zi)**ll
              ! multiply with Bessel function & sum over all ll
              !$omp do
              do ikk = ikk_start, ikk_end
                do imdim = 1, 3
                  integ12_t(:, imdim, ikk) = integ12_t(:, imdim, ikk) + this%j_l(ikk, ll) * weight_x * integ11_t(:, imdim, ll)
                end do
                integ22_t(:, ikk) = integ22_t(:, ikk) + this%j_l(ikk, ll) * weight_x * integ21_t(:, ll)
              end do
              !$omp end do
            end do
            !$omp end parallel
 
            call profiling_out("PES_FLU_SPH_INTEGXX")
 
            call profiling_in("PES_FLU_SPH_INTEGXX2")
 
            !$omp parallel do private(imdim)
            do ikk = ikk_start, ikk_end
              ! sum over dimensions
              do imdim = 1, 3
                spctramp_sph(:, ist, isdim, ik, ikk) = &
                  spctramp_sph(:, ist, isdim, ik, ikk) + &
                  phase_act(:, ikk) * (integ12_t(:, imdim, ikk) * &
                  (m_two * this%veca(imdim, itstep)  / p_c - this%kcoords_sph(imdim, ikk,:)))
              end do
              spctramp_sph(:, ist, isdim, ik, ikk) = &
                spctramp_sph(:, ist, isdim, ik, ikk) + phase_act(:, ikk) * integ22_t(:, ikk) * m_zi
            end do
            !$omp end parallel do
 
            call profiling_out("PES_FLU_SPH_INTEGXX2")
          end do
        end do
      end do
    end do
 
    safe_deallocate_a(s1_node)
    safe_deallocate_a(s2_node)
 
    safe_deallocate_a(integ11_t)
    safe_deallocate_a(integ12_t)
    safe_deallocate_a(integ21_t)
    safe_deallocate_a(integ22_t)
 
    ! save the Volkov phase and the spectral amplitude
    this%conjgphase_prev = m_z0
 
    if (ikk_start > 0) then
      this%conjgphase_prev(:, ikk_start:ikk_end) = phase_act(:, ikk_start:ikk_end)
    end if
    safe_deallocate_a(phase_act)
 
    if (mesh%parallel_in_domains .and. bitand(this%par_strategy, option__pes_flux_parallelization__pf_time) /= 0) then
      call mesh%allreduce(this%conjgphase_prev)
      call mesh%allreduce(spctramp_sph)
    end if
 
    !TODO: Use axpy here
    do ikk = 1, this%nk
      do ik = kptst, kptend
        call lalg_axpy(this%nstepsomegak, stend-stst+1, sdim, dt, spctramp_sph(1:this%nstepsomegak, stst:stend, 1:sdim, ik, ikk), &
          this%spctramp_sph(1:this%nstepsomegak, stst:stend, 1:sdim, ik, ikk))
      end do
    end do
    safe_deallocate_a(spctramp_sph)
 
    pop_sub_with_profile(pes_flux_integrate_sph)
  end subroutine pes_flux_integrate_sph
 
  ! ---------------------------------------------------------
  subroutine pes_flux_getcube(this, mesh, abb, border, offset, fc_ptdens)
    type(mesh_t),                 intent(in)    :: mesh
    type(pes_flux_t),             intent(inout) :: this
    type(absorbing_boundaries_t), intent(in)    :: abb
    real(real64),                 intent(in)    :: border(:)
    real(real64),                 intent(in)    :: offset(:)
    real(real64),                 intent(in)    :: fc_ptdens
 
    integer, allocatable  :: which_surface(:)
    real(real64)          :: xx(mesh%box%dim), dd, area, dS(mesh%box%dim, 2), factor
    integer               :: mdim, imdim, idir, isp, pm,nface, idim, ndir, iu,iv, iuv(1:2)
    integer(int64)           :: ip_global
    integer               :: npface
    integer               :: rankmin, nsurfaces
    logical               :: in_ab
    integer               :: ip_local, NN(mesh%box%dim, 2), idx(mesh%box%dim, 2)
    integer               :: isp_end, isp_start, ifc, n_dir, nfaces, mindim
    real(real64)          :: RSmax(2), RSmin(2), RS(2), dRS(2), weight
 
 
    push_sub(pes_flux_getcube)
 
    ! this routine works on absolute coordinates
 
 
    mdim = mesh%box%dim
 
 
 
    safe_allocate(this%face_idx_range(1:mdim * 2, 1:2))
    safe_allocate(this%LLr(mdim, 1:2))
    safe_allocate(this%NN(mdim, 1:2))
 
    this%face_idx_range(:, :) = 0
    this%LLr(:, :) = m_zero
    this%NN(:, :) = 1
    nn(:, :) = 1
 
    if (this%surf_interp) then
      ! Create a surface with points not on the mesh
 
      idx(:,:) = 0
 
      mindim  = 1
      factor = m_two
      if (this%surf_shape == pes_plane) then
        mindim = mdim  ! We only have two planes along the non periodic dimension
        ! this is due to the fact that for semiperiodic systems we multiply border
        ! by two to prenvet the creation of surfaces at the edges
        factor = m_one
      end if
 
 
      this%nsrfcpnts = 0
 
      do ndir = mdim, mindim, -1
        area = m_one
        do idir=1, mdim
          if (idir == ndir) cycle
          area = area * border(idir) * factor
        end do
 
        npface = int(fc_ptdens * area) ! number of points on the face
 
        idim = 1
        do idir=1, mdim
          if (idir == ndir) cycle
          nn(ndir, idim) = int(sqrt(npface * (border(idir) * factor)**2 / area))
          ds(ndir, idim) = border(idir) * factor / nn(ndir, idim)
          this%LLr(ndir, idim) = nn(ndir, idim) * ds(ndir, idim)
          idx(ndir, idim) = idir
          idim = idim + 1
        end do
        this%nsrfcpnts = this%nsrfcpnts + 2 * product(nn(ndir, 1:mdim - 1))
      end do
 
      this%NN(1:mdim, :) = nn(1:mdim, :)
 
 
      assert(this%nsrfcpnts > 0) !at this point should be satisfied otherwise the point density is way too small
 
 
      safe_allocate(this%srfcnrml(1:mdim, 0:this%nsrfcpnts))
      safe_allocate(this%rcoords(1:mdim, 0:this%nsrfcpnts))
 
      this%srfcnrml(:, :) = m_zero
      this%rcoords(:, :)  = m_zero
 
 
      isp = 1
      ifc = 1
      do ndir = mdim, mindim, -1
 
        !Up face
        this%face_idx_range(ifc,1) = isp
        do iu = 1, nn(ndir,1)
          do iv = 1, nn(ndir,2)
            this%rcoords(ndir, isp)  =  border(ndir)
            this%srfcnrml(ndir, isp) =  product(ds(ndir,1:mdim-1))
            iuv =(/iu, iv/)
            do idim = 1, mdim-1
              this%rcoords(idx(ndir, idim), isp) = (-nn(ndir, idim) * m_half - m_half + iuv(idim)) * ds(ndir, idim)
            end do
            isp = isp + 1
          end do
        end do
        this%face_idx_range(ifc, 2) = isp-1
 
        ifc = ifc + 1
 
        !Down face
        this%face_idx_range(ifc, 1) = isp
        do iu = 1, nn(ndir, 1)
          do iv = 1, nn(ndir, 2)
            this%rcoords(ndir, isp)  = -border(ndir)
            this%srfcnrml(ndir, isp) = -product(ds(ndir, 1:mdim - 1))
            iuv =(/iu, iv/)
            do idim = 1, mdim-1
              this%rcoords(idx(ndir, idim), isp) = (-nn(ndir, idim) * m_half -m_half + iuv(idim)) * ds(ndir, idim)
            end do
            isp = isp + 1
          end do
        end do
        this%face_idx_range(ifc, 2) = isp-1
 
        ifc = ifc + 1
      end do
 
      do isp = 1, this%nsrfcpnts
        this%rcoords(1:mdim, isp) = mesh%coord_system%to_cartesian(this%rcoords(1:mdim, isp))
      end do
 
    else
      ! Surface points are on the mesh
 
      nfaces = mdim * 2
      if (this%surf_shape == pes_plane) nfaces = 2
 
      in_ab = .false.
 
      safe_allocate(which_surface(1:mesh%np_global))
      which_surface = 0
 
      ! get the surface points
      this%nsrfcpnts = 0
      do ip_local = 1, mesh%np
        ip_global = mesh_local2global(mesh, ip_local)
 
        nsurfaces = 0
 
        xx(1:mdim) = mesh%x(ip_local, 1:mdim) - offset(1:mdim)
 
        ! eventually check whether we are in absorbing zone
        select case (abb%abtype)
        case (mask_absorbing)
          in_ab = .not. is_close(abb%mf(ip_local), m_one)
        case (imaginary_absorbing)
          in_ab = abs(abb%mf(ip_local)) > m_epsilon
        case default
          assert(.false.)
        end select
 
        ! check whether the point is inside the cube
        if (all(abs(xx(1:mdim)) <= border(1:mdim)) .and. .not. in_ab) then
          ! check whether the point is close to any border
          do imdim = 1, mdim
            if (this%surf_shape == pes_plane) then
              dd = border(imdim) - xx(imdim)
            else
              dd = border(imdim) - abs(xx(imdim))
            end if
            if (dd < mesh%spacing(imdim) / m_two) then
              nsurfaces = nsurfaces + 1
              which_surface(ip_global) = int(sign(m_one, xx(imdim))) * imdim  ! +-x=+-1, +-y=+-2, +-z=+-3
            end if
          end do
 
          ! check whether the point is close to one border only (not an edge)
          if (nsurfaces == 1) then
            this%nsrfcpnts = this%nsrfcpnts + 1
          else
            which_surface(ip_global) = 0
          end if
        end if
 
      end do
 
      if (mesh%parallel_in_domains) then
        assert(mesh%np_global < huge(0_int32))
        call mesh%allreduce(this%nsrfcpnts)
        call mesh%allreduce(which_surface)
      end if
 
      safe_allocate(this%srfcpnt(1:this%nsrfcpnts))
      safe_allocate(this%srfcnrml(1:mdim, 0:this%nsrfcpnts))
      safe_allocate(this%rcoords(1:mdim, 0:this%nsrfcpnts))
      safe_allocate(this%rankmin(1:this%nsrfcpnts))
 
      this%srfcnrml = m_zero
      this%rcoords  = m_zero
 
      this%face_idx_range(:,:) = 0
 
      isp = 0
      nface = 0
      do idir = mdim, 1, -1 ! Start counting from z to simplify implementation for semiperiodic systems (pln)
        do pm = 1, -1, -2
          nface = nface + 1
          this%face_idx_range(nface, 1) = isp + 1
 
          do ip_global = 1, mesh%np_global
            if (abs(which_surface(ip_global)) == idir .and. sign(1, which_surface(ip_global)) == pm) then
              isp = isp + 1
              ! coordinate of surface point
              xx(1:mdim) = mesh_x_global(mesh, ip_global)
              this%rcoords(1:mdim, isp) = xx(1:mdim)
              ! local ip & node which has the surface point
              this%srfcpnt(isp) = mesh_nearest_point(mesh, this%rcoords(1:mdim, isp), dd, rankmin)
              this%rankmin(isp) = rankmin
              ! surface normal
              this%srfcnrml(idir, isp) = sign(1, which_surface(ip_global))
              ! add the surface element (of directions orthogonal to the normal vector)
              do imdim = 1, mdim
                if (imdim == idir) cycle
                this%srfcnrml(idir, isp) = this%srfcnrml(idir, isp) * mesh%spacing(imdim)
              end do
            end if
          end do
 
          this%face_idx_range(nface,2) = isp
 
        end do
      end do
 
 
      !Get dimensions and spacing on each (pair of) face
      do ifc = 1, nint((nfaces + 0.5) / 2)
        isp_start = this%face_idx_range(ifc, 1)
        isp_end   = this%face_idx_range(ifc, 2)
 
        !retrieve face normal direction
        n_dir = 0
        do idir = 1, mdim
          if (abs(this%srfcnrml(idir, isp_start)) >= m_epsilon) n_dir = idir
        end do
 
        !retrieve the spacing on the face
        drs(:) = m_zero
        idim = 1
        do idir = 1, mdim
          if (idir == n_dir) cycle
          drs(idim)= mesh%spacing(idir)
          idim = idim + 1
        end do
 
 
        !retrieve the dimensions of the face
        rsmin = m_zero
        rsmax = m_zero
        do isp = isp_start, isp_end
          !measure in reduced coordinates
          xx = mesh%coord_system%from_cartesian(this%rcoords(1:mdim, isp))
          idim = 1
          do idir = 1, mdim
            if (idir == n_dir) cycle
            rs(idim)=xx(idir)
            if (rs(idim) < rsmin(idim)) rsmin(idim) = rs(idim)
            if (rs(idim) > rsmax(idim)) rsmax(idim) = rs(idim)
            idim = idim + 1
          end do
        end do
 
 
        do idir = 1, mdim - 1
          this%LLr(n_dir, idir) = rsmax(idir) - rsmin(idir) + drs(idir)
          if (drs(idir) > m_zero) this%NN(n_dir, idir) = int(this%LLr(n_dir, idir) / drs(idir))
        end do
 
      end do
 
 
    end if
 
    if (this%anisotrpy_correction) then
      !Compensate the fact that the angular distribution of points is not uniform
      !and have peaks in correspondence to the edges and corners of the parallelepiped
      do isp = 1, this%nsrfcpnts
        weight = sum(this%rcoords(1:mdim, isp) * this%srfcnrml(1:mdim, isp))
        weight = weight/sum(this%rcoords(1:mdim, isp)**2) / sum(this%srfcnrml(1:mdim, isp)**2)
        this%srfcnrml(1:mdim, isp) = this%srfcnrml(1:mdim, isp) * abs(weight)
      end do
 
    end if
 
 
    safe_deallocate_a(which_surface)
 
    pop_sub(pes_flux_getcube)
  end subroutine pes_flux_getcube
 
  ! ---------------------------------------------------------
  subroutine pes_flux_getsphere(this, mesh, nstepsthetar, nstepsphir)
    type(pes_flux_t), intent(inout) :: this
    type(mesh_t),     intent(in)    :: mesh
    integer,          intent(inout) :: nstepsthetar, nstepsphir
 
    integer :: mdim
    integer :: isp, ith, iph
    real(real64)   :: thetar, phir
    real(real64)   :: weight, dthetar
 
    push_sub(pes_flux_getsphere)
 
    mdim = mesh%box%dim
 
    if (nstepsphir == 0) nstepsphir = 1
 
    if (nstepsthetar <= 1) then
      nstepsphir = 1
      nstepsthetar = 1
    end if
    dthetar = m_pi / nstepsthetar
    this%nsrfcpnts  = nstepsphir * (nstepsthetar - 1) + 2
 
    safe_allocate(this%srfcnrml(1:mdim, 0:this%nsrfcpnts))
    safe_allocate(this%rcoords(1:mdim, 0:this%nsrfcpnts))
 
    this%srfcnrml = m_zero
    this%rcoords  = m_zero
 
    ! initializing spherical grid
    isp = 0
    do ith = 0, nstepsthetar
      thetar  = ith * dthetar
 
      if (ith == 0 .or. ith == nstepsthetar) then
        weight = (m_one - cos(dthetar / m_two)) * m_two * m_pi
      else
        weight = abs(cos(thetar - dthetar / m_two) - cos(thetar + dthetar / m_two)) &
          * m_two * m_pi / nstepsphir
      end if
 
      do iph = 0, nstepsphir - 1     ! 2*pi is the same than zero
        isp = isp + 1
        phir = iph * m_two * m_pi / nstepsphir
        this%srfcnrml(1, isp) = cos(phir) * sin(thetar)
        if (mdim >= 2) this%srfcnrml(2, isp) = sin(phir) * sin(thetar)
        if (mdim == 3) this%srfcnrml(3, isp) = cos(thetar)
        this%rcoords(1:mdim, isp) = this%radius * this%srfcnrml(1:mdim, isp)
        ! here we also include the surface elements
        this%srfcnrml(1:mdim, isp) = this%radius**m_two * weight * this%srfcnrml(1:mdim, isp)
        if (ith == 0 .or. ith == nstepsthetar) exit
      end do
    end do
 
    pop_sub(pes_flux_getsphere)
  end subroutine pes_flux_getsphere
 
 
 
 
  ! ---------------------------------------------------------
  subroutine pes_flux_distribute(istart_global, iend_global, istart, iend, comm)
    integer,          intent(in)    :: istart_global
    integer,          intent(in)    :: iend_global
    integer,          intent(inout) :: istart
    integer,          intent(inout) :: iend
    type(mpi_comm),   intent(in)    :: comm
 
#if defined(HAVE_MPI)
    integer, allocatable :: dimrank(:)
    integer              :: mpisize, mpirank, irest, irank
    integer              :: inumber
#endif
 
    push_sub(pes_flux_distribute)
 
#if defined(HAVE_MPI)
    call mpi_comm_size(comm, mpisize)
    call mpi_comm_rank(comm, mpirank)
 
    safe_allocate(dimrank(0:mpisize-1))
 
    inumber = iend_global - istart_global + 1
    irest = mod(inumber, mpisize)
    do irank = 0, mpisize - 1
      if (irest > 0) then
        dimrank(irank) = int(inumber / mpisize) + 1
        irest = irest - 1
      else
        dimrank(irank) = int(inumber / mpisize)
      end if
    end do
 
    iend   = istart_global + sum(dimrank(:mpirank)) - 1
    istart = iend - dimrank(mpirank) + 1
 
    if (istart > iend) then
      iend = 0
      istart = 0
    end if
 
    safe_deallocate_a(dimrank)
 
#else
    istart = istart_global
    iend   = iend_global
#endif
 
    pop_sub(pes_flux_distribute)
  end subroutine pes_flux_distribute
 
 
#include "pes_flux_out_inc.F90"
 
end module pes_flux_oct_m
 
!! Local Variables:
!! mode: f90
!! coding: utf-8
!! End: