floquet_utils.F

Go to the documentation of this file.

!--------------------------------------------------------------------------------------------------!
!   CP2K: A general program to perform molecular dynamics simulations                              !
!   Copyright 2000-2026 CP2K developers group <https://cp2k.org>                                   !
!                                                                                                  !
!   SPDX-License-Identifier: GPL-2.0-or-later                                                      !
!--------------------------------------------------------------------------------------------------!
 
! **************************************************************************************************
!> \brief Helper routines for the Floquet-Bloch band-structure calculation (floquet_main).
!> \par History
!> \author Shridhar Shanbhag (27.01.2026)
! **************************************************************************************************
MODULE floquet_utils
   USE cp_blacs_env,                    ONLY: cp_blacs_env_create,&
                                              cp_blacs_env_type
   USE cp_cfm_types,                    ONLY: cp_cfm_get_submatrix,&
                                              cp_cfm_set_submatrix,&
                                              cp_cfm_type
   USE cp_files,                        ONLY: close_file,&
                                              open_file
   USE floquet_types,                   ONLY: floquet_active_band_count,&
                                              floquet_env_type
   USE input_constants,                 ONLY: tensor_small_cell_full_kp
   USE kinds,                           ONLY: default_string_length,&
                                              dp,&
                                              int_8
   USE machine,                         ONLY: m_memory_details
   USE mathconstants,                   ONLY: gaussi,&
                                              pi,&
                                              z_one,&
                                              z_zero
   USE message_passing,                 ONLY: mp_comm_split_type_shared,&
                                              mp_comm_type,&
                                              mp_para_env_type
   USE physcon,                         ONLY: a_bohr,&
                                              evolt,&
                                              kelvin
   USE post_scf_bandstructure_types,    ONLY: post_scf_bandstructure_type
   USE qs_environment_types,            ONLY: get_qs_env,&
                                              qs_environment_type
   USE qs_moments,                      ONLY: qs_moment_kpoints_deep
   USE util,                            ONLY: sort
#include "./base/base_uses.f90"
 
   IMPLICIT NONE
 
   PRIVATE
 
   CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'floquet_utils'
 
   ! Number of full Floquet-Hamiltonian copies (each n_f_size^2 complex(dp), 16 bytes per
   ! element) that the subgroup must be able to hold, distributed across its ranks.
   ! We size for 8 to leave headroom.
   INTEGER, PARAMETER, PRIVATE :: n_fm_work_copies = 8
 
   ! Public subroutines
   PUBLIC :: build_floquet_matrix, &
             compute_e_k_p_k, &
             make_floquet_subgroups, &
             distribute_floquet_kp_data, &
             floquet_sector_weights, &
             check_floquet_convergence, &
             calculate_floquet_observables, &
             write_floquet_header, &
             write_floquet_results
 
CONTAINS
 
! **************************************************************************************************
!> \brief Off-diagonal coupling block of H_F(k) for one k-point and spin, from precomputed data.
!> \param bs_env ...
!> \param p_k_kp_spin ...
!> \param off_diag_m ...
! **************************************************************************************************
   SUBROUTINE build_off_diagonal_matrix(bs_env, p_k_kp_spin, off_diag_m)
      TYPE(post_scf_bandstructure_type), POINTER         :: bs_env
      COMPLEX(KIND=dp), DIMENSION(:, :, :), INTENT(IN)   :: p_k_kp_spin
      COMPLEX(KIND=dp), DIMENSION(:, :), INTENT(OUT)     :: off_diag_m
 
      CHARACTER(LEN=*), PARAMETER :: routineN = 'build_off_diagonal_matrix'
 
      COMPLEX(KIND=dp), DIMENSION(3)                     :: efactor
      INTEGER                                            :: handle, i, i_dir, j, n_mo
      REAL(KIND=dp)                                      :: amplitude, omega
      REAL(KIND=dp), DIMENSION(3)                        :: e_vec, phi, polarisation
 
      CALL timeset(routinen, handle)
 
      ! Builds the matrix that occupies the off diagonal blocks in the Floquet Matrix H_F
 
      n_mo = SIZE(p_k_kp_spin, 2)
 
      polarisation(:) = bs_env%floquet_polarisation(:)
      IF (sqrt(sum(polarisation**2)) < epsilon(0.0_dp)) THEN
         cpabort("Invalid (too small) polarisation vector specified for POLARISATION")
      END IF
 
      amplitude = bs_env%floquet_amplitude
      e_vec(:) = amplitude*polarisation
      phi(:) = pi*bs_env%floquet_phi(:)
      omega = bs_env%floquet_omega
 
      ! E_factor(α) = (i E(α) exp(i·φ(α)))/(2ω), where α = x,y,z
      DO i_dir = 1, 3
         efactor(i_dir) = gaussi*e_vec(i_dir)*cmplx(cos(phi(i_dir)), sin(phi(i_dir)), kind=dp)/(2*omega)
      END DO
 
      ! off_diag_uv = Σ_α [p_uv^α(k) · E_factor(α)], where α = x,y,z
      off_diag_m(:, :) = z_zero
      DO i_dir = 1, 3
         DO i = 1, n_mo
            DO j = 1, n_mo
               off_diag_m(i, j) = off_diag_m(i, j) + &
                                  p_k_kp_spin(i_dir, i, j)*efactor(i_dir)
            END DO
         END DO
      END DO
 
      CALL timestop(handle)
 
   END SUBROUTINE build_off_diagonal_matrix
 
! **************************************************************************************************
!> \brief Diagonal block of H_F(k) for Floquet sector f_index.
!> \param bs_env ...
!> \param e_k_kp_spin ...
!> \param f_index Floquet sector index
!> \param diag_e ...
! **************************************************************************************************
   SUBROUTINE build_diagonal_matrix(bs_env, e_k_kp_spin, f_index, diag_e)
      TYPE(post_scf_bandstructure_type), POINTER         :: bs_env
      REAL(KIND=dp), DIMENSION(:), INTENT(IN)            :: e_k_kp_spin
      INTEGER, INTENT(IN)                                :: f_index
      COMPLEX(KIND=dp), DIMENSION(:, :), INTENT(OUT)     :: diag_e
 
      CHARACTER(LEN=*), PARAMETER :: routineN = 'build_diagonal_matrix'
 
      INTEGER                                            :: handle, i, n_mo
      REAL(KIND=dp)                                      :: omega
 
      CALL timeset(routinen, handle)
      ! Builds the diagonal block of the H_F for Floquet sector index f_index
      ! Equilibrium band energies shifted by f_index·ħΩ, i.e., diag_e(n,n) = ε_{nk} + f_index·ħΩ.
      n_mo = SIZE(e_k_kp_spin)
      omega = bs_env%floquet_omega
      diag_e(:, :) = z_zero
      DO i = 1, n_mo
         diag_e(i, i) = e_k_kp_spin(i) + f_index*omega
      END DO
 
      CALL timestop(handle)
 
   END SUBROUTINE build_diagonal_matrix
 
! **************************************************************************************************
!> \brief Builds the Floquet-Bloch Hamiltonian H_F(k) for a single k-point and spin channel from
!>        the band data currently held in floquet_env.
!> \param bs_env ...
!> \param floquet_env ...
!> \param floquet_matrix ...
! **************************************************************************************************
   SUBROUTINE build_floquet_matrix(bs_env, floquet_env, floquet_matrix)
      TYPE(post_scf_bandstructure_type), POINTER         :: bs_env
      TYPE(floquet_env_type), INTENT(IN)                 :: floquet_env
      TYPE(cp_cfm_type), INTENT(IN)                      :: floquet_matrix
 
      CHARACTER(LEN=*), PARAMETER :: routinen = 'build_floquet_matrix'
 
      COMPLEX(KIND=dp), ALLOCATABLE, DIMENSION(:, :)     :: conj_off_diag_m, diag_e, off_diag_m
      INTEGER                                            :: f_index, handle, i, i_f, max_f_index, &
                                                            n_fbands, n_mo
 
      CALL timeset(routinen, handle)
 
      n_mo = floquet_env%n_mo
      max_f_index = floquet_env%max_f_index
      n_fbands = 1 + 2*max_f_index
 
      ! Creates the floquet matrix H_F by placing the diagonal and off diagonal blocks
      ALLOCATE (diag_e(n_mo, n_mo), source=z_zero)
      ALLOCATE (off_diag_m(n_mo, n_mo), source=z_zero)
      ALLOCATE (conj_off_diag_m(n_mo, n_mo), source=z_zero)
 
      CALL build_off_diagonal_matrix(bs_env, floquet_env%p_k_kp_spin, off_diag_m)
      conj_off_diag_m(:, :) = conjg(transpose(off_diag_m(:, :)))
 
      floquet_matrix%local_data(:, :) = z_zero
      DO i = 1, n_fbands
         i_f = 1 + (i - 1)*n_mo
         f_index = i - max_f_index - 1
         CALL build_diagonal_matrix(bs_env, floquet_env%e_k_kp_spin, f_index, diag_e)
         CALL cp_cfm_set_submatrix(floquet_matrix, diag_e, i_f, i_f)
         IF (i > 1) THEN
            CALL cp_cfm_set_submatrix(floquet_matrix, off_diag_m, i_f, i_f - n_mo)
            CALL cp_cfm_set_submatrix(floquet_matrix, conj_off_diag_m, i_f - n_mo, i_f)
         END IF
      END DO
      DEALLOCATE (diag_e, off_diag_m, conj_off_diag_m)
 
      CALL timestop(handle)
   END SUBROUTINE build_floquet_matrix
 
! **************************************************************************************************
!> \brief Make this k-point/spin's band data available to every rank of the owning subgroup.
!> \param para_env the global parallel environment
!> \param para_env_sub the subgroup parallel environment
!> \param ispin the spin channel
!> \param ikp the DOS k-point index
!> \param e_k band energies for all DOS k-points, distributed
!> \param p_k momentum matrix elements, distributed
!> \param floquet_env ...
! **************************************************************************************************
   SUBROUTINE distribute_floquet_kp_data(para_env, para_env_sub, ispin, ikp, e_k, p_k, &
                                         floquet_env)
      TYPE(mp_para_env_type), POINTER                    :: para_env, para_env_sub
      INTEGER, INTENT(IN)                                :: ispin, ikp
      REAL(kind=dp), DIMENSION(:, :, :), INTENT(IN)      :: e_k
      COMPLEX(KIND=dp), DIMENSION(:, :, :, :, :), &
         INTENT(IN)                                      :: p_k
      TYPE(floquet_env_type), INTENT(INOUT)              :: floquet_env
 
      CHARACTER(LEN=*), PARAMETER :: routinen = 'distribute_floquet_kp_data'
 
      INTEGER                                            :: handle, loc_idx, local_src, nhi, nlo, &
                                                            owner
 
      CALL timeset(routinen, handle)
 
      ! Identify the owner's rank within the subgroup (all subgroup ranks agree on local_src).
      owner = mod(ikp - 1, para_env%num_pe)
 
      local_src = -1
      IF (para_env%mepos == owner) local_src = para_env_sub%mepos
      CALL para_env_sub%max(local_src)
 
      ! The owner copies its single-spin slice of this k-point's distributed band data into the
      ! broadcast buffers, then broadcasts them to the whole subgroup. Only the active band
      ! window [n_lo, n_hi] (the MAX_NMO states closest to the VBM) is taken from the full-band
      ! e_k/p_k arrays; the broadcast buffers are sized to this window (floquet_env%n_mo).
      IF (para_env%mepos == owner) THEN
         loc_idx = ceiling(real(ikp)/para_env%num_pe)
         nlo = floquet_env%n_lo(ispin)
         nhi = floquet_env%n_hi(ispin)
         floquet_env%e_k_kp_spin(:) = e_k(ispin, loc_idx, nlo:nhi)
         floquet_env%p_k_kp_spin(:, :, :) = p_k(ispin, loc_idx, :, nlo:nhi, nlo:nhi)
      END IF
      CALL para_env_sub%bcast(floquet_env%e_k_kp_spin, local_src)
      CALL para_env_sub%bcast(floquet_env%p_k_kp_spin, local_src)
 
      CALL timestop(handle)
 
   END SUBROUTINE distribute_floquet_kp_data
 
! **************************************************************************************************
!> \brief Precompute, on the global communicator, the band quantities needed to assemble the
!>        Floquet-Bloch Hamiltonian for all DOS k-points (all spins), distributed across ranks.
!> \param qs_env ...
!> \param bs_env ...
!> \param e_k ...
!> \param p_k ...
! **************************************************************************************************
   SUBROUTINE compute_e_k_p_k(qs_env, bs_env, e_k, p_k)
      TYPE(qs_environment_type), POINTER                 :: qs_env
      TYPE(post_scf_bandstructure_type), POINTER         :: bs_env
      REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :, :), &
         INTENT(OUT)                                     :: e_k
      COMPLEX(KIND=dp), ALLOCATABLE, &
         DIMENSION(:, :, :, :, :), INTENT(OUT)           :: p_k
 
      CHARACTER(LEN=*), PARAMETER                        :: routinen = 'compute_e_k_p_k'
 
      INTEGER                                            :: handle, ikp, nkp_only_bs, nkp_start
      REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :)        :: xkp_all
 
      CALL timeset(routinen, handle)
 
      ! The band-structure k-points are the last nkp_only_bs entries of kpoints_DOS%xkp, stored
      ! after the nkp_only_DOS DOS-only k-points.
      nkp_start = bs_env%nkp_only_DOS
      nkp_only_bs = bs_env%nkp_only_bs
 
      ! Collect the coordinates of all band-structure k-points (global index nkp_start + local).
      ALLOCATE (xkp_all(3, nkp_only_bs))
      DO ikp = 1, nkp_only_bs
         xkp_all(:, ikp) = bs_env%kpoints_DOS%xkp(1:3, nkp_start + ikp)
      END DO
 
      ! Calculate and distribute the results across ranks
      ! One k-point per rank, round-robin by global rank
      ! So the results for ikp are stored in mepos==MOD(ikp-1,num_pe)
      CALL qs_moment_kpoints_deep(qs_env, xkp_all, momentum=p_k, energy=e_k, do_parallel=.true.)
 
      DEALLOCATE (xkp_all)
 
      CALL timestop(handle)
   END SUBROUTINE compute_e_k_p_k
 
! **************************************************************************************************
!> \brief Finds the number of MPI ranks that share a physical node, determined by splitting the
!>        given communicator into node-local communicators via MPI's shared-memory split type.
!> \param para_env ...
!> \return the number of ranks on the calling rank's node
! **************************************************************************************************
   FUNCTION find_ranks_per_node(para_env) RESULT(ranks_per_node)
      TYPE(mp_para_env_type), POINTER                    :: para_env
      INTEGER                                            :: ranks_per_node
 
      TYPE(mp_comm_type)                                 :: node_comm
 
      CALL node_comm%from_split_type(para_env, mp_comm_split_type_shared, key=para_env%mepos)
      ranks_per_node = max(1, node_comm%num_pe)
      CALL node_comm%free()
 
   END FUNCTION find_ranks_per_node
 
! **************************************************************************************************
!> \brief Choose the optimal number of MPI ranks per subgroup for Floquet calculations.
!> \param qs_env ...
!> \param bs_env ...
!> \return the chosen number of ranks per subgroup (1 .. num_pe)
! **************************************************************************************************
   FUNCTION floquet_determine_subgroup_size(qs_env, bs_env) RESULT(group_size)
      TYPE(qs_environment_type), POINTER                 :: qs_env
      TYPE(post_scf_bandstructure_type), POINTER         :: bs_env
      INTEGER                                            :: group_size
 
      CHARACTER(LEN=*), PARAMETER :: routinen = 'floquet_determine_subgroup_size'
 
      INTEGER                                            :: g_load, g_mem, handle, n_f_size, &
                                                            nkp_only_bs, ranks_per_node, unit_nr
      INTEGER(KIND=int_8)                                :: buffers, cached, memfree, memlikelyfree, &
                                                            memtotal, needed_bytes, slab, &
                                                            sreclaimable, usable_per_rank
      REAL(kind=dp)                                      :: mem_fill_fraction
      TYPE(mp_para_env_type), POINTER                    :: para_env
 
      CALL timeset(routinen, handle)
 
      CALL get_qs_env(qs_env, para_env=para_env)
 
      unit_nr = bs_env%unit_nr
      ! Size for the active band window (the MAX_NMO band count, or the full count when MAX_NMO <= 0)
      n_f_size = floquet_active_band_count(bs_env)*(1 + 2*bs_env%max_floquet_index)
      nkp_only_bs = bs_env%nkp_only_bs
      mem_fill_fraction = bs_env%floquet_mem_fill_fraction
 
      ! We choose the subgroup size to be the larger of two floors:
      ! 1. Memory floor G_mem: the subgroup must jointly hold the Floquet working set,
      !    G_mem = ceil(needed_bytes/usable per-rank memory), where the usable memory is
      !    the mem_fill_fraction of the per-rank FREE memory.
      ! 2. K-point floor G_load = floor(num_pe / nkp_only_bs): the LARGEST G that still
      !    yields at least nkp_only_bs subgroups (floor(num_pe/G) >= nkp_only_bs). When ranks
      !    outnumber k-points this hands every k-point its own (larger, hence faster)
      !    subgroup instead of leaving spare ranks idle.
 
      ! usable_per_rank = a fraction of the per-rank FREE memory, not including SCF and other data
      CALL m_memory_details(memtotal, memfree, buffers, cached, slab, sreclaimable, memlikelyfree)
      ranks_per_node = find_ranks_per_node(para_env)
      usable_per_rank = int(mem_fill_fraction*real(memfree, dp), int_8)/int(ranks_per_node, int_8)
 
      ! Total memory the subgroup must hold: n_fm_work_copies copies of the Floquet
      ! Hamiltonian, each n_f_size^2 complex(dp) numbers at 16 bytes each.
      needed_bytes = 16_int_8*int(n_fm_work_copies, int_8)*int(n_f_size, int_8)**2
 
      IF (usable_per_rank <= 0_int_8) THEN
         ! Memory info unavailable (e.g. no /proc found) use one subgroup spanning all ranks
         g_mem = para_env%num_pe
      ELSE
         ! Memory floor: smallest G with needed_bytes/G <= usable_per_rank
         g_mem = int(min((needed_bytes + usable_per_rank - 1_int_8)/usable_per_rank, &
                         int(para_env%num_pe, int_8)))
      END IF
 
      IF (g_mem > para_env%num_pe) cpwarn("Total Memory Likely Insufficient, process may be killed")
 
      ! K-point floor: floor(num_pe / nkp_only_bs) = largest G that gives at least nkp_only_bs subgroups
      g_load = para_env%num_pe/max(nkp_only_bs, 1)
 
      group_size = max(g_mem, g_load)
      group_size = min(group_size, para_env%num_pe)
      group_size = max(group_size, 1)
      CALL para_env%max(group_size)
 
      IF (unit_nr > 0) THEN
         WRITE (unit_nr, '(T2,A)') ""
         WRITE (unit_nr, '(T2,A,T66,I15)') "FLOQUET MEMORY | Detected MPI ranks per node:", ranks_per_node
         WRITE (unit_nr, '(T2,A,T66,I12,A)') "FLOQUET MEMORY | Free memory per rank:", &
            (memlikelyfree/int(ranks_per_node, int_8))/1048576_int_8, " MB"
         WRITE (unit_nr, '(T2,A,T66,I12,A)') "FLOQUET MEMORY | Usable per rank (reserve applied):", &
            usable_per_rank/1048576_int_8, " MB"
         WRITE (unit_nr, '(T2,A,T66,I15)') "FLOQUET MEMORY | Floquet copies needed:", n_fm_work_copies
         WRITE (unit_nr, '(T2,A,T66,I12,A)') "FLOQUET MEMORY | Total Floquet working set:", &
            needed_bytes/1048576_int_8, " MB"
         WRITE (unit_nr, '(T2,A,T66,I12,A)') "FLOQUET MEMORY | Floquet working set per rank:", &
            (needed_bytes/int(max(group_size, 1), int_8))/1048576_int_8, " MB"
         WRITE (unit_nr, '(T2,A,T66,I15)') "FLOQUET MEMORY | Memory floor (ranks/subgroup):", g_mem
         WRITE (unit_nr, '(T2,A,T66,I15)') "FLOQUET MEMORY | K-point floor (ranks/subgroup):", g_load
         WRITE (unit_nr, '(T2,A,T66,I15)') "FLOQUET MEMORY | Chosen MPI ranks per subgroup:", group_size
      END IF
 
      CALL timestop(handle)
   END FUNCTION floquet_determine_subgroup_size
 
! **************************************************************************************************
!> \brief Split the para_env into subgroups so that the Floquet Hamiltonian is distributed across
!>        the ranks of each subgroup. Also creates the subgroup BLACS context.
!> \param qs_env ...
!> \param bs_env ...
!> \param para_env_sub the created subgroup parallel environment
!> \param blacs_env_sub the created subgroup BLACS context
!> \param group_distribution subgroup index of every global rank, 0:num_pe-1
!> \param ngroups the number of subgroups created
! **************************************************************************************************
   SUBROUTINE make_floquet_subgroups(qs_env, bs_env, para_env_sub, blacs_env_sub, &
                                     group_distribution, ngroups)
      TYPE(qs_environment_type), POINTER                 :: qs_env
      TYPE(post_scf_bandstructure_type), POINTER         :: bs_env
      TYPE(mp_para_env_type), POINTER                    :: para_env_sub
      TYPE(cp_blacs_env_type), POINTER                   :: blacs_env_sub
      INTEGER, ALLOCATABLE, DIMENSION(:), INTENT(OUT)    :: group_distribution
      INTEGER, INTENT(OUT)                               :: ngroups
 
      CHARACTER(LEN=*), PARAMETER :: routinen = 'make_floquet_subgroups'
 
      INTEGER                                            :: group_size, handle, n_spin, nkp_only_bs, &
                                                            stride_kp, unit_nr
      TYPE(mp_para_env_type), POINTER                    :: para_env
 
      CALL timeset(routinen, handle)
 
      CALL get_qs_env(qs_env, para_env=para_env)
      unit_nr = bs_env%unit_nr
      nkp_only_bs = bs_env%nkp_only_bs
      n_spin = bs_env%n_spin
 
      ! We split the global communicator into subgroups with the Floquet Hamiltonian distributed
      ! across the processes in each subgroup and the k-points distributed across subgroups.
      group_size = floquet_determine_subgroup_size(qs_env, bs_env)
 
      IF (nkp_only_bs < para_env%num_pe) THEN
         stride_kp = para_env%num_pe/group_size
      ELSE
         stride_kp = 1
      END IF
      ALLOCATE (group_distribution(0:para_env%num_pe - 1))
      ALLOCATE (para_env_sub)
      CALL para_env_sub%from_split(comm=para_env, ngroups=ngroups, &
                                   group_distribution=group_distribution, &
                                   subgroup_min_size=group_size, stride=stride_kp)
      IF (unit_nr > 0) THEN
         WRITE (unit_nr, '(T2,A,T66,I15)') "FLOQUET MEMORY | Subgroup rank stride:", stride_kp
         WRITE (unit_nr, '(T2,A,T66,I15)') "FLOQUET MEMORY | Number of subgroups:", ngroups
         WRITE (unit_nr, '(/,T2,A,I5,A,I1,A)') "FLOQUET CALCULATIONS PROGRESS OUT OF", nkp_only_bs, &
            " K-POINTS AND ", n_spin, " SPINS"
      END IF
 
      NULLIFY (blacs_env_sub)
      CALL cp_blacs_env_create(blacs_env_sub, para_env_sub)
 
      CALL timestop(handle)
 
   END SUBROUTINE make_floquet_subgroups
 
! **************************************************************************************************
!> \brief Central and boundary-sector weights of every Floquet eigenvector, stored into
!>        floquet_env%w0 and floquet_env%wE.
!> \param floquet_env ...
!> \param floquet_eigenvectors Floquet eigenvectors
! **************************************************************************************************
   SUBROUTINE floquet_sector_weights(floquet_env, floquet_eigenvectors)
      TYPE(floquet_env_type), INTENT(INOUT)              :: floquet_env
      TYPE(cp_cfm_type), INTENT(IN)                      :: floquet_eigenvectors
 
      CHARACTER(LEN=*), PARAMETER :: routinen = 'floquet_sector_weights'
 
      COMPLEX(KIND=dp), ALLOCATABLE, DIMENSION(:, :)     :: v_block
      INTEGER                                            :: handle, max_f_index, n_f_size, n_mo
 
      CALL timeset(routinen, handle)
 
      n_mo = floquet_env%n_mo
      max_f_index = floquet_env%max_f_index
      n_f_size = floquet_env%n_f_size
 
      ALLOCATE (v_block(n_mo, n_f_size))
 
      ! Evaluate w0_α = Σ_{n=1..n_mo} |<n,m=0|α>|^2 (central sector)
      CALL cp_cfm_get_submatrix(floquet_eigenvectors, v_block, 1 + n_mo*max_f_index, 1)
      floquet_env%w0(:) = sum(abs(v_block)**2, dim=1)
 
      ! Evaluate wE_α = Σ_{n=1..n_mo} (|<n,m=-M|α>|^2 + |<n,m=+M|α>|^2) (outermost rungs)
      floquet_env%wE(:) = 0.0_dp
      IF (max_f_index >= 1) THEN
         CALL cp_cfm_get_submatrix(floquet_eigenvectors, v_block, 1, 1)
         floquet_env%wE(:) = sum(abs(v_block)**2, dim=1)
 
         CALL cp_cfm_get_submatrix(floquet_eigenvectors, v_block, 1 + n_mo*2*max_f_index, 1)
         floquet_env%wE(:) = floquet_env%wE(:) + sum(abs(v_block)**2, dim=1)
      END IF
 
      DEALLOCATE (v_block)
 
      ! The orthonormal rows of a unitary matrix are such that Σ_α w0_α = n_mo.
      cpassert(abs(sum(floquet_env%w0) - real(n_mo, dp)) < 1.0e-6_dp*real(n_mo, dp))
 
      CALL timestop(handle)
 
   END SUBROUTINE floquet_sector_weights
 
! **************************************************************************************************
!> \brief Checks that MAX_FLOQUET_INDEX is large enough and the Floquet Hamiltonian was
!>        truncated far enough away from the central sector to leave it unperturbed.
!> \param bs_env ...
!> \param floquet_env holds the central- (w0) and outermost-sector (wE) weights
! **************************************************************************************************
   SUBROUTINE check_floquet_convergence(bs_env, floquet_env)
      TYPE(post_scf_bandstructure_type), POINTER         :: bs_env
      TYPE(floquet_env_type), INTENT(IN)                 :: floquet_env
 
      CHARACTER(LEN=*), PARAMETER :: routinen = 'check_floquet_convergence'
 
      CHARACTER(LEN=default_string_length)               :: msg
      INTEGER                                            :: handle
      REAL(kind=dp)                                      :: leak
 
      ! The truncation is converged if m=0 states have small weight on outermost rungs m = ±M
      CALL timeset(routinen, handle)
 
      ! MAX_FLOQUET_INDEX = 0 means no sidebands at all
      IF (bs_env%max_floquet_index < 1) THEN
         CALL timestop(handle)
         RETURN
      END IF
 
      ! No cleanly m=0-dominated state exists: the drive has hybridised every band with
      ! its sidebands, strong field but not a truncation failure.
      IF (.NOT. any(floquet_env%w0 > 0.5_dp)) THEN
         cpwarn("Floquet: no m=0-dominated state; cannot assess MAX_FLOQUET_INDEX convergence.")
         CALL timestop(handle)
         RETURN
      END IF
 
      leak = maxval(floquet_env%wE, mask=(floquet_env%w0 > 0.5_dp))
 
      IF (bs_env%eps_floquet > 0.0_dp .AND. leak > bs_env%eps_floquet) THEN
         WRITE (msg, '(A,ES10.2E2,A,ES10.2E2)') &
            "MAX_FLOQUET_INDEX is too small. Leak: ", leak, "exceeds EPS_FLOQUET: ", bs_env%eps_floquet
         cpabort(trim(msg))
      END IF
 
      CALL timestop(handle)
 
   END SUBROUTINE check_floquet_convergence
 
! **************************************************************************************************
!> \brief Energy origin for the Floquet calculations, chosen to be the valence band maximum.
!> \param bs_env ...
!> \return the valence band maximum
! **************************************************************************************************
   FUNCTION floquet_reference_energy(bs_env) RESULT(mu)
      TYPE(post_scf_bandstructure_type), POINTER         :: bs_env
      REAL(kind=dp)                                      :: mu
 
      IF (bs_env%do_gw .OR. &
          bs_env%gw_implementation == tensor_small_cell_full_kp) THEN
         mu = bs_env%band_edges_scf%VBM
      ELSE IF (bs_env%n_spin == 1) THEN
         mu = bs_env%eigenval_scf_Gamma(bs_env%n_occ(1), 1)
      ELSE
         mu = max(bs_env%eigenval_scf_Gamma(bs_env%n_occ(1), 1), &
                  bs_env%eigenval_scf_Gamma(bs_env%n_occ(2), 2))
      END IF
 
   END FUNCTION floquet_reference_energy
 
! **************************************************************************************************
!> \brief Calculate all Floquet observables for one k-point and spin: the DOS (a_k), the m=0 bands,
!>        the quasi-energies, and store in floquet_env.
!> \param bs_env ...
!> \param para_env_sub ...
!> \param ispin ...
!> \param ikp ...
!> \param floquet_env holds eigenvalues, w0 (input) and a_k, m0_energies, m0_weights,
!>        quasi_energies (output)
! **************************************************************************************************
   SUBROUTINE calculate_floquet_observables(bs_env, para_env_sub, ispin, ikp, floquet_env)
      TYPE(post_scf_bandstructure_type), POINTER         :: bs_env
      TYPE(mp_para_env_type), POINTER                    :: para_env_sub
      INTEGER, INTENT(IN)                                :: ispin, ikp
      TYPE(floquet_env_type), INTENT(INOUT)              :: floquet_env
 
      CHARACTER(LEN=*), PARAMETER :: routinen = 'calculate_floquet_observables'
 
      INTEGER                                            :: handle, i, i_e, j, n_e, n_f_size, n_mo
      INTEGER, ALLOCATABLE, DIMENSION(:)                 :: work
      REAL(kind=dp)                                      :: broad, cum, e_min, energy, energy_step, &
                                                            mu, omega, target_level
 
      CALL timeset(routinen, handle)
 
      n_f_size = floquet_env%n_f_size
      n_mo = floquet_env%n_mo
      n_e = floquet_env%n_E
 
      mu = floquet_reference_energy(bs_env)
      omega = bs_env%floquet_omega
      broad = bs_env%broadening_floquet
      energy_step = bs_env%energy_step_floquet
      e_min = mu - bs_env%energy_window_floquet
 
      ! (1) Spectral function / DOS, A(k,E) = -(1/π) Im Tr_KS G^R_00(k,E),
      ! where G^R_00(k,E) = (E + iη - H_F)^(-1) is the retarded Green's function
      ! With H_F = Σ_α λ_α |α><α|, we simplify
      ! A(k,E) = -(1/π) Σ_α w0_α Im[1/(E + iη - λ_α)], η = broadening/2.
      ! This avoids the per-energy dense inversion of the full H_F.
 
      DO i_e = 1, n_e
         energy = e_min + i_e*energy_step
         floquet_env%a_k(i_e) = -sum(floquet_env%w0(:)* &
                                     aimag(z_one/(energy + gaussi*broad/2.0_dp - &
                                                  floquet_env%eigenvalues(:))))/pi
      END DO
 
      ! (2) m=0 band structure by weighted count
      cum = 0.0_dp
      j = 0
      DO i = 1, n_mo
         target_level = real(i, dp) - 0.5_dp
         DO WHILE (cum < target_level .AND. j < n_f_size)
            j = j + 1
            cum = cum + floquet_env%w0(j)
         END DO
         floquet_env%m0_energies(i) = floquet_env%eigenvalues(j) - mu
         floquet_env%m0_weights(i) = floquet_env%w0(j)
      END DO
 
      ! (3) Fold the m=0 bands (already relative to the VBM) into the first Floquet Brillouin zone.
      floquet_env%quasi_energies(:) = floquet_env%m0_energies &
                                      - omega*real(ceiling(floquet_env%m0_energies/omega - 0.5_dp), dp)
 
      ! Sort the quasi-energies
      ALLOCATE (work(n_mo))
      CALL sort(floquet_env%quasi_energies, n_mo, work)
      DEALLOCATE (work)
 
      ! Store this result on the subgroup source only, so the global sum below picks up
      ! each work item only once
      IF (para_env_sub%is_source()) THEN
         floquet_env%all_quasi_energies(:, ispin, ikp) = floquet_env%quasi_energies(:)
         floquet_env%all_a_k(:, ispin, ikp) = floquet_env%a_k(:)
         floquet_env%all_m0_energies(:, ispin, ikp) = floquet_env%m0_energies(:)
         floquet_env%all_m0_weights(:, ispin, ikp) = floquet_env%m0_weights(:)
      END IF
 
      CALL timestop(handle)
 
   END SUBROUTINE calculate_floquet_observables
 
! **************************************************************************************************
!> \brief Print the Floquet header with the input parameters and a short description of the output.
!> \param bs_env ...
! **************************************************************************************************
   SUBROUTINE write_floquet_header(bs_env)
      TYPE(post_scf_bandstructure_type), POINTER         :: bs_env
 
      CHARACTER(LEN=*), PARAMETER :: routinen = 'write_floquet_header'
 
      INTEGER                                            :: handle, n_active, n_f_size, unit_nr
 
      CALL timeset(routinen, handle)
 
      unit_nr = bs_env%unit_nr
      ! Active band window (the MAX_NMO states closest to the VBM, or all bands when MAX_NMO <= 0)
      n_active = floquet_active_band_count(bs_env)
      n_f_size = n_active*(1 + 2*bs_env%max_floquet_index)
 
      IF (unit_nr > 0) THEN
 
         WRITE (unit_nr, '(T2,A)') ' '
         WRITE (unit_nr, '(T2,A)') repeat('-', 79)
         WRITE (unit_nr, '(T2,A,A78)') '-', '-'
         WRITE (unit_nr, '(T2,A,A51,A27)') '-', 'FLOQUET BANDSTRUCTURE CALCULATION', '-'
         WRITE (unit_nr, '(T2,A,A78)') '-', '-'
         WRITE (unit_nr, '(T2,A)') repeat('-', 79)
         WRITE (unit_nr, '(T2,A)') ' '
 
         WRITE (unit_nr, '(T2,A,T37,A,T67,ES12.2E2)') "FLOQUET PARAMETERS", "Amplitude [V/m]:", &
            bs_env%floquet_amplitude*evolt/a_bohr
         WRITE (unit_nr, '(T37,A,T67,F12.4)') "Frequency [eV]:", bs_env%floquet_omega*evolt
         WRITE (unit_nr, '(T37,A)') repeat("-", 44)
         WRITE (unit_nr, '(T37,A,T55,3F8.4)') "Polarisation:", bs_env%floquet_polarisation(1:3)
         WRITE (unit_nr, '(T37,A,T55,3F8.4)') "Phase offsets:", pi*bs_env%floquet_phi(1:3)
         WRITE (unit_nr, '(T37,A)') repeat("-", 44)
         WRITE (unit_nr, '(T37,A,T67,I12)') "Max Floquet index:", bs_env%max_floquet_index
         WRITE (unit_nr, '(T37,A,T67,I12)') "Max KS states (MAX_NMO, 0=all):", bs_env%max_nmo_floquet
         WRITE (unit_nr, '(T37,A,T67,I12)') "Active bands used (of total):", n_active
         WRITE (unit_nr, '(T37,A,T67,I12)') "Total number of bands:", bs_env%n_ao
         WRITE (unit_nr, '(T37,A,T67,I12)') "Floquet Hamiltonian Size:", n_f_size
         WRITE (unit_nr, '(T37,A)') repeat("-", 44)
         WRITE (unit_nr, '(T37,A,T67,F12.4)') &
            "Energy window [eV]:", bs_env%energy_window_floquet*evolt
         WRITE (unit_nr, '(T37,A,T67,F12.4)') "Energy step [eV]:", bs_env%energy_step_floquet*evolt
         WRITE (unit_nr, '(T37,A,T67,F12.4)') "Broadening [eV]:", bs_env%broadening_floquet*evolt
         WRITE (unit_nr, '(T37,A)') repeat("-", 44)
         WRITE (unit_nr, '(A)') ""
 
         WRITE (unit_nr, '(T2,A)') &
            "We construct the Floquet-Bloch Hamiltonian and diagonalise it. Projecting the"
         WRITE (unit_nr, '(T2,A)') &
            "eigenvectors onto the Floquet sectors gives the weights w = Σ_n |<n,m|α>|²,"
         WRITE (unit_nr, '(T2,A)') &
            "from which all of the following are obtained."
         WRITE (unit_nr, '(A)') ""
         WRITE (unit_nr, '(T2,A)') &
            "The m=0 bands are the eigenvectors with the largest central-sector weight; they"
         WRITE (unit_nr, '(T2,A)') &
            "reduce to the equilibrium bands at zero field and are stored, with their"
         WRITE (unit_nr, '(T2,A)') &
            " weights, in FLOQUET_BANDSTRUCTURE.bs"
         WRITE (unit_nr, '(A)') ""
         WRITE (unit_nr, '(T2,A)') &
            "Folding those bands into the first Floquet Brillouin zone, relative to the VBM,"
         WRITE (unit_nr, '(T2,A)') &
            "gives the quasi-energies stored in QUASI_ENERGIES.bs"
         WRITE (unit_nr, '(A)') ""
         WRITE (unit_nr, '(T2,A)') &
            "The k-resolved density of states is obtained by computing the trace of"
         WRITE (unit_nr, '(T2,A)') &
            "the retarded Green's function and stored in FLOQUET_DOS.out"
         WRITE (unit_nr, '(T2,A)') &
            "DOS(ω,k) = -1/π*Im[Tr_KS(G^R(ω,k))]"
         WRITE (unit_nr, '(A)') ""
      END IF
 
      CALL timestop(handle)
 
   END SUBROUTINE write_floquet_header
 
! **************************************************************************************************
!> \brief Sum the accumulated results across MPI ranks and write the m=0 band structure,
!>        the quasi-energies and the Floquet DOS to their files.
!> \param bs_env ...
!> \param floquet_env ...
! **************************************************************************************************
   SUBROUTINE write_floquet_results(bs_env, floquet_env)
      TYPE(post_scf_bandstructure_type), POINTER         :: bs_env
      TYPE(floquet_env_type), INTENT(INOUT)              :: floquet_env
 
      CHARACTER(LEN=*), PARAMETER :: routinen = 'write_floquet_results'
 
      CHARACTER(LEN=default_string_length)               :: fname
      INTEGER                                            :: bunit, handle, i, i_e, ikp_for_file, &
                                                            ispin, n_e, n_mo, n_spin, nkp_only_bs, &
                                                            nkp_start, qunit, wunit
      REAL(kind=dp)                                      :: e_min, energy, energy_step, f_occ, kt, &
                                                            mu, x
      REAL(kind=dp), DIMENSION(3)                        :: xkp
 
      CALL timeset(routinen, handle)
 
      ! Collect all results on the global communicator and write them
      CALL bs_env%para_env%sum(floquet_env%all_quasi_energies)
      CALL bs_env%para_env%sum(floquet_env%all_a_k)
      CALL bs_env%para_env%sum(floquet_env%all_m0_energies)
      CALL bs_env%para_env%sum(floquet_env%all_m0_weights)
 
      IF (bs_env%para_env%is_source()) THEN
 
         n_mo = floquet_env%n_mo
         n_spin = floquet_env%n_spin
         nkp_only_bs = floquet_env%nkp_only_bs
         n_e = floquet_env%n_E
         nkp_start = bs_env%nkp_only_DOS
 
         ! All three share one energy origin, the VBM.
         mu = floquet_reference_energy(bs_env)
         energy_step = bs_env%energy_step_floquet
         e_min = mu - bs_env%energy_window_floquet
 
         ! k_B T in Hartree (only used when TEMPERATURE > 0); E[K] = E[Hartree]*kelvin.
         kt = bs_env%floquet_temperature/kelvin
 
         ! Each file is opened once (REPLACE) and written for every band-structure k-point and spin.
 
         ! The m=0 band structure
         WRITE (fname, "(2A)") trim(bs_env%floquet_bs_file), ".bs"
         CALL open_file(trim(fname), unit_number=bunit, file_status="REPLACE", file_action="WRITE")
         WRITE (bunit, "(A)") "# Floquet m=0 (central-sector) band structure"
         WRITE (bunit, "(A)") "# (in units of eV, relative to the VBM, not folded)"
         WRITE (bunit, "(A)") "# w = sum_n |<n,m=0|alpha>|^2 in [0,1]: w~1 clean m=0 replica,"
         WRITE (bunit, "(A)") "# w~0.5 hybridised with a sideband (drive near resonance)"
         WRITE (bunit, "(A)") "# The band column is the absolute Kohn-Sham band index; only the"
         WRITE (bunit, "(A)") "# MAX_NMO band window closest to the VBM is reported (all bands if 0)."
 
         ! Quasi-energies
         WRITE (fname, "(2A)") trim(bs_env%floquet_qe_file), ".bs"
         CALL open_file(trim(fname), unit_number=qunit, file_status="REPLACE", file_action="WRITE")
         WRITE (qunit, "(A)") "# Quasi-energies obtained by diagonalising the Floquet Hamiltonian"
         WRITE (qunit, "(A)") "# (in units of eV, the m=0 bands relative to the VBM, folded to"
         WRITE (qunit, "(A)") "# the first Floquet Brillouin zone -hbar*Omega/2 < e <= hbar*Omega/2)"
 
         ! DOS
         WRITE (fname, "(2A)") trim(bs_env%floquet_dos_file), ".out"
         CALL open_file(trim(fname), unit_number=wunit, file_status="REPLACE", file_action="WRITE")
         WRITE (wunit, "(A)") "# Floquet Density of States: D(ω,k) = -1/π*Im[Tr_KS(G^R(ω,k))]"
 
         DO ikp_for_file = 1, nkp_only_bs
            xkp(1:3) = bs_env%kpoints_DOS%xkp(1:3, nkp_start + ikp_for_file)
            DO ispin = 1, n_spin
 
               ! <floquet_bs_file>.bs the m=0 band structure (absolute energies) and its weights
               WRITE (bunit, "(A,I0,T10,A,I0,A,T24,3(1X,F14.8))") &
                  "#  Spin ", ispin, "  Point ", ikp_for_file, ": ", xkp(1:3)
               WRITE (bunit, "(A)") "#  KS band       Energy [eV]         m=0 weight"
               DO i = 1, n_mo
                  WRITE (bunit, "(I8,F21.8,F17.5)") floquet_env%n_lo(ispin) - 1 + i, &
                     floquet_env%all_m0_energies(i, ispin, ikp_for_file)*evolt, &
                     floquet_env%all_m0_weights(i, ispin, ikp_for_file)
               END DO
 
               ! <floquet_qe_file>.bs Quasi-energies (folded to -ħΩ/2 < ε ≤ ħΩ/2)
               WRITE (qunit, "(A,I0,T10,A,I0,A,T24,3(1X,F14.8))") &
                  "#  Spin ", ispin, "  Point ", ikp_for_file, ": ", xkp(1:3)
               WRITE (qunit, "(A)") "#  Floquet band  Quasi-energy [eV]"
               DO i = 1, n_mo
                  WRITE (qunit, "(I8,F21.8)") i, &
                     floquet_env%all_quasi_energies(i, ispin, ikp_for_file)*evolt
               END DO
 
               ! <floquet_dos_file>.out the k-resolved DOS A(k,ω), and (if TEMPERATURE > 0) the
               ! occupied spectral weight f(E)*A(k,ω), f the Fermi-Dirac occupation at the VBM.
               WRITE (wunit, "(A,I0,T10,A,I0,A,T24,3(1X,F14.8))") &
                  "#  Spin ", ispin, "  Point ", ikp_for_file, ": ", xkp(1:3)
               IF (bs_env%floquet_temperature > 0.0_dp) THEN
                  WRITE (wunit, "(A)") "#Energy-VBM (eV)  A(ω,k) = DOS (1/eV)   f*A = occupied DOS (1/eV)"
                  DO i_e = 1, n_e
                     energy = e_min + i_e*energy_step
                     x = (energy - mu)/kt                    ! (E - VBM)/k_B T, dimensionless
                     IF (x > 40.0_dp) THEN                   ! guard EXP overflow at low T
                        f_occ = 0.0_dp
                     ELSE IF (x < -40.0_dp) THEN
                        f_occ = 1.0_dp
                     ELSE
                        f_occ = 1.0_dp/(exp(x) + 1.0_dp)
                     END IF
                     WRITE (wunit, "(2X,3G13.4)") (energy - mu)*evolt, &
                        floquet_env%all_a_k(i_e, ispin, ikp_for_file)/evolt, &
                        f_occ*floquet_env%all_a_k(i_e, ispin, ikp_for_file)/evolt
                  END DO
               ELSE
                  WRITE (wunit, "(A)") "#Energy-VBM (eV)  A(ω,k) = DOS (1/eV)"
                  DO i_e = 1, n_e
                     energy = e_min + i_e*energy_step
                     WRITE (wunit, "(2X,2G13.4)") (energy - mu)*evolt, &
                        floquet_env%all_a_k(i_e, ispin, ikp_for_file)/evolt
                  END DO
               END IF
 
            END DO
         END DO
 
         CALL close_file(bunit)
         CALL close_file(qunit)
         CALL close_file(wunit)
      END IF
 
      CALL timestop(handle)
 
   END SUBROUTINE write_floquet_results
 
END MODULE floquet_utils