kpoint_lattice_fft_unittest.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                                                      !
!--------------------------------------------------------------------------------------------------!
PROGRAM kpoint_lattice_fft_unittest
   USE fft_tools,                       ONLY: finalize_fft,&
                                              init_fft
   USE kinds,                           ONLY: dp
   USE kpoint_lattice_fft,              ONLY: cell_to_k_grid_fft,&
                                              k_grid_to_cell_fft,&
                                              regular_kpoint_grid
   USE lattice_preconditioner_operator, ONLY: &
        apply_lattice_inverse_dense, apply_lattice_state_inverse_dense, build_lattice_inverse, &
        build_lattice_local_correction, lattice_fft_auto_select, lattice_fft_cost_ratio, &
        lattice_fft_storage_ratio, lattice_grid, lattice_local_cost_ratio, &
        lattice_local_storage_ratio
   USE mathconstants,                   ONLY: gaussi,&
                                              twopi,&
                                              z_zero
   USE message_passing,                 ONLY: mp_comm_type,&
                                              mp_world_finalize,&
                                              mp_world_init
 
   IMPLICIT NONE
 
   INTEGER, PARAMETER                                  :: base_block = 2, base_cell = 8, base_n = 16, &
                                                          base_ncol = 3, ncell = 31, nkp = 64, &
                                                          nround = 16, nselected = 3
   INTEGER, PARAMETER                                  :: large_block = 1000, large_ncol = 1000, &
                                                          large_num_pe = 16
   COMPLEX(KIND=dp), ALLOCATABLE, DIMENSION(:, :, :)   :: inverse_k
   COMPLEX(KIND=dp), ALLOCATABLE, DIMENSION(:, :, :, :) :: state_inverse_k
   COMPLEX(KIND=dp), DIMENSION(2, 3, nkp)              :: direct, transformed
   COMPLEX(KIND=dp), DIMENSION(2, 3, nround)           :: reconstructed, round_k
   COMPLEX(KIND=dp), DIMENSION(2, 3, nselected)         :: selected_transformed
   INTEGER                                             :: d, i, icell, info, j, jcell, ik, ix, iy, iz
   INTEGER, ALLOCATABLE, DIMENSION(:)                  :: correction_indices
   INTEGER, DIMENSION(3)                               :: base_dims, delta, large_dims
   INTEGER, DIMENSION(3, base_cell)                    :: base_index
   INTEGER, DIMENSION(3)                               :: nkp_grid
   INTEGER, DIMENSION(3, ncell)                        :: index_to_cell
   INTEGER, DIMENSION(3, nround)                       :: round_index
   INTEGER, DIMENSION(nselected)                       :: selected_kpoints
   LOGICAL                                             :: used_fft
   REAL(kind=dp), DIMENSION(3, 3)                      :: hmat
   REAL(kind=dp)                                       :: base_residual_norm, corrected_projection_error, &
                                                          corrected_residual_norm, cost_ratio, &
                                                          expected_cost_ratio, imaginary_residual, &
                                                          large_block_real, large_ncell_real, &
                                                          large_ncol_real, projection_error
   REAL(kind=dp), DIMENSION(3, base_cell)              :: base_xkp
   REAL(kind=dp), DIMENSION(base_n, base_n)            :: base_overlap
   REAL(kind=dp), DIMENSION(base_n, base_n)            :: defect_operator
   REAL(kind=dp), DIMENSION(base_n, base_ncol)          :: base_input, base_input_rotated, &
                                                           base_output, base_output_rotated, &
                                                           base_state_output, defect_residual
   REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :)         :: coarse_inverse, defect_base_output, &
                                                          defect_corrected_output, local_operator, &
                                                          selected_output, selected_rhs
   REAL(kind=dp), DIMENSION(base_ncol, base_ncol)       :: column_rotation
   REAL(kind=dp), DIMENSION(3, nkp)                    :: xkp, xkp_bad
   REAL(kind=dp), DIMENSION(2, 3, ncell)               :: values_rs
   REAL(kind=dp), DIMENSION(2, 3, nround)              :: round_rs
   TYPE(mp_comm_type)                                  :: mp_comm
 
   CALL mp_world_init(mp_comm)
   CALL init_fft("FFTSG", alltoall=.false., pool_limit=4, &
                 wisdom_file="", plan_style=1)
 
   nkp_grid = [4, 4, 4]
   hmat = reshape([4.0_dp, 0.1_dp, 0.2_dp, &
                   0.0_dp, 5.0_dp, 0.3_dp, &
                   0.0_dp, 0.0_dp, 6.0_dp], [3, 3])
 
   ik = 0
   DO iz = 0, nkp_grid(3) - 1
      DO iy = 0, nkp_grid(2) - 1
         DO ix = 0, nkp_grid(1) - 1
            ik = ik + 1
            xkp(:, ik) = [(real(ix, dp) + 0.5_dp)/real(nkp_grid(1), dp), &
                          REAL(iy, dp)/REAL(nkp_grid(2), dp), &
                          (REAL(iz, dp) + 0.25_dp)/REAL(nkp_grid(3), dp)]
         END DO
      END DO
   END DO
   IF (.NOT. regular_kpoint_grid(xkp, nkp_grid)) error stop "Regular k-point grid not recognized"
 
   DO icell = 1, ncell
      index_to_cell(:, icell) = [modulo(3*icell + 1, 11) - 5, &
                                 modulo(5*icell + 2, 9) - 4, &
                                 modulo(7*icell + 3, 7) - 3]
      DO d = 1, 3
         values_rs(1, d, icell) = sin(0.13_dp*real(icell + 2*d, dp))
         values_rs(2, d, icell) = cos(0.17_dp*real(2*icell - d, dp))
      END DO
   END DO
 
   CALL cell_to_k_grid_fft(values_rs, index_to_cell, xkp, nkp_grid, transformed, used_fft)
   IF (.NOT. used_fft) error stop "Complete regular grid did not use lattice FFT"
   CALL direct_reference(values_rs, index_to_cell, xkp, direct)
   IF (maxval(abs(transformed - direct)) > 2.0e-12_dp) THEN
      error stop "Shifted lattice FFT differs from direct transform"
   END IF
 
   selected_kpoints = [2, 31, 64]
   CALL cell_to_k_grid_fft(values_rs, index_to_cell, xkp, nkp_grid, selected_transformed, &
                           used_fft, selected_kpoints=selected_kpoints)
   IF (.NOT. used_fft) error stop "Selected distributed output did not use lattice FFT"
   IF (maxval(abs(selected_transformed - direct(:, :, selected_kpoints))) > 2.0e-12_dp) THEN
      error stop "Selected lattice-FFT output differs from the full transform"
   END IF
 
   CALL cell_to_k_grid_fft(values_rs, index_to_cell, xkp, nkp_grid, transformed, used_fft, 2, hmat)
   IF (.NOT. used_fft) error stop "Derivative did not use lattice FFT"
   CALL direct_reference(values_rs, index_to_cell, xkp, direct, 2, hmat)
   IF (maxval(abs(transformed - direct)) > 5.0e-12_dp) THEN
      error stop "Lattice-FFT derivative differs from direct transform"
   END IF
 
   nkp_grid = [4, 4, 1]
   ik = 0
   DO iy = 0, nkp_grid(2) - 1
      DO ix = 0, nkp_grid(1) - 1
         ik = ik + 1
         xkp(:, ik) = [(real(ix, dp) + 0.5_dp)/real(nkp_grid(1), dp), &
                       REAL(iy, dp)/REAL(nkp_grid(2), dp), 0.0_dp]
      END DO
   END DO
   CALL cell_to_k_grid_fft(values_rs, index_to_cell, xkp(:, 1:ik), nkp_grid, &
                           transformed(:, :, 1:ik), used_fft)
   IF (.NOT. used_fft) error stop "Singleton grid direction did not use padded lattice FFT"
   CALL direct_reference(values_rs, index_to_cell, xkp(:, 1:ik), direct(:, :, 1:ik))
   IF (maxval(abs(transformed(:, :, 1:ik) - direct(:, :, 1:ik))) > 2.0e-12_dp) THEN
      error stop "Padded lattice FFT differs from direct transform"
   END IF
 
   icell = 0
   DO iy = 0, nkp_grid(2) - 1
      DO ix = 0, nkp_grid(1) - 1
         icell = icell + 1
         round_index(:, icell) = [ix, iy, 0]
         round_rs(:, :, icell) = values_rs(:, :, icell)
      END DO
   END DO
   CALL cell_to_k_grid_fft(round_rs, round_index, xkp(:, 1:ik), nkp_grid, round_k, used_fft)
   IF (.NOT. used_fft) error stop "Round-trip forward transform did not use lattice FFT"
   CALL k_grid_to_cell_fft(round_k, xkp(:, 1:ik), nkp_grid, round_index, reconstructed, used_fft)
   IF (.NOT. used_fft) error stop "Round-trip inverse transform did not use lattice FFT"
   IF (maxval(abs(reconstructed - cmplx(round_rs, kind=dp))) > 3.0e-12_dp) THEN
      error stop "Forward/inverse lattice-FFT round trip is inconsistent"
   END IF
 
   nkp_grid = [4, 4, 4]
 
   xkp_bad = xkp
   xkp_bad(1, nkp) = xkp_bad(1, nkp) + 0.03125_dp
   CALL cell_to_k_grid_fft(values_rs, index_to_cell, xkp_bad, nkp_grid, transformed, used_fft)
   IF (used_fft) error stop "Irregular k-point list incorrectly used lattice FFT"
   CALL direct_reference(values_rs, index_to_cell, xkp_bad, direct)
   IF (maxval(abs(transformed - direct)) > 1.0e-14_dp) THEN
      error stop "Direct lattice-transform fallback is inconsistent"
   END IF
 
   ! Verify the block-circulant inverse against direct matrix application and right covariance.
   base_dims = [4, 2, 1]
   CALL lattice_grid(base_dims, base_index, base_xkp)
   base_overlap = 0.0_dp
   DO jcell = 1, base_cell
      DO icell = 1, base_cell
         delta = modulo(base_index(:, icell) - base_index(:, jcell), base_dims)
         DO i = 1, base_block
            IF (all(delta == 0)) THEN
               base_overlap((icell - 1)*base_block + i, (jcell - 1)*base_block + i) = &
                  2.8_dp + 0.2_dp*real(i, kind=dp)
            ELSE IF ((delta(1) == 1 .OR. delta(1) == base_dims(1) - 1) .AND. &
                     delta(2) == 0) THEN
               base_overlap((icell - 1)*base_block + i, (jcell - 1)*base_block + i) = -0.12_dp
            ELSE IF (delta(1) == 0 .AND. delta(2) == 1) THEN
               base_overlap((icell - 1)*base_block + i, (jcell - 1)*base_block + i) = -0.08_dp
            END IF
         END DO
         IF (all(delta == 0)) THEN
            base_overlap((icell - 1)*base_block + 1, (jcell - 1)*base_block + 2) = 0.15_dp
            base_overlap((icell - 1)*base_block + 2, (jcell - 1)*base_block + 1) = 0.15_dp
         END IF
      END DO
   END DO
   DO j = 1, base_ncol
      DO i = 1, base_n
         base_input(i, j) = sin(0.11_dp*real(2*i + 3*j, kind=dp))
      END DO
   END DO
   CALL build_lattice_inverse(base_overlap, base_dims, inverse_k, projection_error, used_fft, info)
   IF (info /= 0) error stop "Positive lattice overlap failed block Cholesky inversion"
   IF (.NOT. used_fft) error stop "Block-circulant overlap construction did not use lattice FFT"
   IF (projection_error > 2.0e-15_dp) error stop "Exact block-circulant overlap was changed by projection"
   IF (lattice_fft_auto_select(base_dims, base_block, base_ncol, 1, 1, &
                               projection_error, used_fft, info)) THEN
      error stop "Cost-aware lattice-FFT AUTO accepted a tiny eight-cell operator"
   END IF
   IF (.NOT. lattice_fft_auto_select(base_dims, 6, 8, 1, 2, &
                                     projection_error, used_fft, info)) THEN
      error stop "Cost-aware lattice-FFT AUTO rejected an economical eight-cell operator"
   END IF
   IF (lattice_fft_auto_select([4, 1, 1], 6, 8, 1, 2, &
                               projection_error, used_fft, info)) THEN
      error stop "Conservative lattice-FFT AUTO accepted fewer than eight cells"
   END IF
   IF (lattice_fft_auto_select(base_dims, 6, 8, 1, 2, 2.0e-3_dp, used_fft, info)) THEN
      error stop "Conservative lattice-FFT AUTO accepted a large projection error"
   END IF
   IF (lattice_fft_auto_select(base_dims, 6, 8, 17, 2, &
                               projection_error, used_fft, info)) THEN
      error stop "Storage-aware lattice-FFT AUTO accepted too many state operators"
   END IF
   IF (lattice_fft_cost_ratio(base_dims, base_block, base_ncol, 1) <= 0.8_dp) THEN
      error stop "Lattice-FFT cost estimate missed the tiny-batch overhead"
   END IF
   large_dims = [100, 100, 1]
   large_ncell_real = real(product(large_dims), kind=dp)
   large_block_real = real(large_block, kind=dp)
   large_ncol_real = real(large_ncol, kind=dp)
   expected_cost_ratio = ( &
                         large_ncell_real*large_block_real**2*large_ncol_real + &
                         4.0_dp*large_ncell_real*large_block_real*large_ncol_real* &
                         log(large_ncell_real)/log(2.0_dp) + &
                         8.0_dp*large_ncell_real*large_block_real + &
                         2.0_dp*large_ncell_real*large_block_real*large_ncol_real* &
                         log(real(large_num_pe + 1, kind=dp))/log(2.0_dp))/ &
                         ((large_ncell_real*large_block_real)**2*large_ncol_real)
   cost_ratio = lattice_fft_cost_ratio(large_dims, large_block, large_ncol, large_num_pe)
   IF (abs(cost_ratio - expected_cost_ratio) > 1.0e-14_dp*expected_cost_ratio) THEN
      error stop "Lattice-FFT cost estimate overflowed for a large problem"
   END IF
   IF (abs(lattice_fft_storage_ratio(base_dims, 1) - 0.25_dp) > 1.0e-14_dp) THEN
      error stop "Lattice-FFT storage estimate is inconsistent"
   END IF
   IF (abs(lattice_local_cost_ratio(base_dims, 1) - 0.28125_dp) > 1.0e-14_dp) THEN
      error stop "Lattice local-correction cost estimate is inconsistent"
   END IF
   IF (abs(lattice_local_storage_ratio(base_dims, 3, 1) - 0.421875_dp) > 1.0e-14_dp) THEN
      error stop "Lattice local-correction storage estimate is inconsistent"
   END IF
   CALL apply_lattice_inverse_dense(inverse_k, base_dims, base_input, base_output, &
                                    used_fft, imaginary_residual)
   IF (.NOT. used_fft) error stop "Block-circulant inverse application did not use lattice FFT"
   IF (imaginary_residual > 3.0e-13_dp) error stop "Lattice inverse did not reconstruct a real matrix"
   IF (maxval(abs(matmul(base_overlap, base_output) - base_input)) > 2.0e-11_dp) THEN
      error stop "FFT block-circulant inverse differs from the dense operator inverse"
   END IF
 
   ALLOCATE (state_inverse_k(base_block, base_block, base_cell, base_ncol))
   DO j = 1, base_ncol
      state_inverse_k(:, :, :, j) = (0.5_dp + 0.25_dp*real(j, kind=dp))*inverse_k
   END DO
   CALL apply_lattice_state_inverse_dense(state_inverse_k, base_dims, base_input, &
                                          base_state_output, used_fft, imaginary_residual)
   IF (.NOT. used_fft) error stop "State-dependent inverse application did not use lattice FFT"
   DO j = 1, base_ncol
      IF (maxval(abs(base_state_output(:, j) - &
                     (0.5_dp + 0.25_dp*real(j, kind=dp))*base_output(:, j))) > 3.0e-12_dp) THEN
         error stop "State-dependent lattice inverse used the wrong column operator"
      END IF
   END DO
   DEALLOCATE (state_inverse_k)
 
   column_rotation = 0.0_dp
   column_rotation(1, 1) = cos(0.31_dp)
   column_rotation(1, 2) = -sin(0.31_dp)
   column_rotation(2, 1) = sin(0.31_dp)
   column_rotation(2, 2) = cos(0.31_dp)
   column_rotation(3, 3) = 1.0_dp
   base_input_rotated = matmul(base_input, column_rotation)
   CALL apply_lattice_inverse_dense(inverse_k, base_dims, base_input_rotated, &
                                    base_output_rotated, used_fft)
   IF (maxval(abs(base_output_rotated - matmul(base_output, column_rotation))) > 3.0e-12_dp) THEN
      error stop "FFT lattice base is not right-covariant under orbital rotations"
   END IF
 
   ! Add one localized positive defect and verify the balanced cell correction.
   defect_operator = base_overlap
   defect_operator(5, 5) = defect_operator(5, 5) + 0.90_dp
   defect_operator(6, 6) = defect_operator(6, 6) + 0.60_dp
   CALL build_lattice_inverse(defect_operator, base_dims, inverse_k, projection_error, used_fft, info)
   IF (info /= 0 .OR. .NOT. used_fft) error stop "Defective lattice operator failed construction"
   IF (projection_error <= 1.0e-3_dp) error stop "Localized defect did not break lattice translations"
   IF (lattice_fft_auto_select(base_dims, 6, 8, 1, 2, projection_error, used_fft, info)) THEN
      error stop "Lattice-FFT AUTO accepted a localized defect"
   END IF
   CALL build_lattice_local_correction( &
      defect_operator, base_dims, 1, correction_indices, local_operator, coarse_inverse, &
      corrected_projection_error, info)
   IF (info /= 0) error stop "Local correction coarse inverse failed"
   IF (any(correction_indices /= [5, 6])) error stop "Local correction selected the wrong defect cell"
   IF (corrected_projection_error >= projection_error) THEN
      error stop "Local correction did not reduce the non-circulant residual"
   END IF
 
   ALLOCATE (defect_base_output(base_n, base_ncol), defect_corrected_output(base_n, base_ncol))
   CALL apply_lattice_inverse_dense(inverse_k, base_dims, base_input, defect_base_output, used_fft)
   CALL apply_balanced_local_inverse(inverse_k, base_dims, correction_indices, local_operator, &
                                     coarse_inverse, base_input, defect_corrected_output)
   defect_residual = matmul(defect_operator, defect_base_output) - base_input
   base_residual_norm = norm2(defect_residual)
   defect_residual = matmul(defect_operator, defect_corrected_output) - base_input
   corrected_residual_norm = norm2(defect_residual)
   IF (corrected_residual_norm >= base_residual_norm) THEN
      error stop "Balanced local correction did not improve inverse application"
   END IF
   IF (sum(base_input*defect_corrected_output) <= 0.0_dp) THEN
      error stop "Balanced local correction lost positive definiteness"
   END IF
 
   ALLOCATE (selected_rhs(base_n, SIZE(correction_indices)), &
             selected_output(base_n, SIZE(correction_indices)))
   selected_rhs(:, :) = local_operator
   CALL apply_balanced_local_inverse(inverse_k, base_dims, correction_indices, local_operator, &
                                     coarse_inverse, selected_rhs, selected_output)
   selected_rhs = 0.0_dp
   DO i = 1, SIZE(correction_indices)
      selected_rhs(correction_indices(i), i) = 1.0_dp
   END DO
   IF (maxval(abs(selected_output - selected_rhs)) > 3.0e-11_dp) THEN
      error stop "Balanced local correction is not exact in its selected subspace"
   END IF
   DEALLOCATE (coarse_inverse, correction_indices, defect_base_output, defect_corrected_output, &
               local_operator, selected_output, selected_rhs)
   DEALLOCATE (inverse_k)
 
   CALL finalize_fft(mp_comm, "")
   CALL mp_world_finalize()
 
CONTAINS
 
! **************************************************************************************************
!> \brief Apply the balanced two-level inverse used by the production local correction.
!> \param inverse_k block-circulant base inverse
!> \param lattice_dims replicated-cell dimensions
!> \param correction_indices selected coordinate basis
!> \param operator_columns columns A U
!> \param coarse_inverse inverse of U^T A U
!> \param input right-hand sides
!> \param output corrected inverse action
! **************************************************************************************************
   SUBROUTINE apply_balanced_local_inverse(inverse_k, lattice_dims, correction_indices, &
                                           operator_columns, coarse_inverse, input, output)
 
      COMPLEX(KIND=dp), DIMENSION(:, :, :), INTENT(IN)   :: inverse_k
      INTEGER, DIMENSION(3), INTENT(IN)                  :: lattice_dims
      INTEGER, DIMENSION(:), INTENT(IN)                  :: correction_indices
      REAL(KIND=dp), DIMENSION(:, :), INTENT(IN)         :: operator_columns, coarse_inverse, input
      REAL(KIND=dp), DIMENSION(:, :), INTENT(OUT)        :: output
 
      INTEGER                                            :: i
      LOGICAL                                            :: used_fft
      REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :)        :: base_input, coarse_input, &
                                                            coarse_projected, coarse_residual, &
                                                            coarse_rhs
 
      ALLOCATE (base_input(SIZE(input, 1), SIZE(input, 2)), &
                coarse_input(SIZE(correction_indices), SIZE(input, 2)), &
                coarse_projected(SIZE(correction_indices), SIZE(input, 2)), &
                coarse_residual(SIZE(correction_indices), SIZE(input, 2)), &
                coarse_rhs(SIZE(correction_indices), SIZE(input, 2)))
      DO i = 1, SIZE(correction_indices)
         coarse_input(i, :) = input(correction_indices(i), :)
      END DO
      coarse_rhs(:, :) = matmul(coarse_inverse, coarse_input)
      base_input(:, :) = input - matmul(operator_columns, coarse_rhs)
      CALL apply_lattice_inverse_dense(inverse_k, lattice_dims, base_input, output, used_fft)
      IF (.NOT. used_fft) error stop "Balanced correction base did not use lattice FFT"
      coarse_projected(:, :) = matmul(transpose(operator_columns), output)
      coarse_residual(:, :) = matmul(coarse_inverse, coarse_projected)
      output(correction_indices, :) = output(correction_indices, :) + coarse_rhs - coarse_residual
      DEALLOCATE (base_input, coarse_input, coarse_projected, coarse_residual, coarse_rhs)
 
   END SUBROUTINE apply_balanced_local_inverse
 
! **************************************************************************************************
!> \brief ...
!> \param values_rs ...
!> \param index_to_cell ...
!> \param xkp ...
!> \param values_k ...
!> \param deriv_direction ...
!> \param hmat ...
! **************************************************************************************************
   SUBROUTINE direct_reference(values_rs, index_to_cell, xkp, values_k, deriv_direction, hmat)
 
      REAL(KIND=dp), DIMENSION(:, :, :), INTENT(IN)      :: values_rs
      INTEGER, DIMENSION(:, :), INTENT(IN)               :: index_to_cell
      REAL(KIND=dp), DIMENSION(:, :), INTENT(IN)         :: xkp
      COMPLEX(KIND=dp), DIMENSION(:, :, :), INTENT(OUT)  :: values_k
      INTEGER, INTENT(IN), OPTIONAL                      :: deriv_direction
      REAL(KIND=dp), DIMENSION(3, 3), INTENT(IN), &
         OPTIONAL                                        :: hmat
 
      COMPLEX(KIND=dp)                                   :: factor
      INTEGER                                            :: icell, ik
      REAL(KIND=dp), DIMENSION(3)                        :: cell_index_real, cell_vector
 
      values_k = z_zero
      DO ik = 1, SIZE(xkp, 2)
         DO icell = 1, SIZE(values_rs, 3)
            factor = exp(gaussi*twopi*sum(xkp(:, ik)*real(index_to_cell(:, icell), kind=dp)))
            IF (PRESENT(deriv_direction)) THEN
               cell_index_real = real(index_to_cell(:, icell), kind=dp)
               cell_vector = matmul(hmat, cell_index_real)
               factor = factor*gaussi*cell_vector(deriv_direction)
            END IF
            values_k(:, :, ik) = values_k(:, :, ik) + factor*values_rs(:, :, icell)
         END DO
      END DO
 
   END SUBROUTINE direct_reference
 
END PROGRAM kpoint_lattice_fft_unittest