Add new pw types, adjust different interfaces

src/common/cp_array_utils.fypp

@@ -6,22 +6,26 @@
 #!-------------------------------------------------------------------------------------------------!
 #:mute
 
-   #:set nametype1 =['r', 'i', 'logical']
-   #:set type1 = ['REAL(kind=dp)','INTEGER(kind=int_4)','logical']
+   #:set nametype1 =['r', 'c', 'i', 'logical']
+   #:set type1 = ['REAL(kind=dp)','COMPLEX(KIND=dp)','INTEGER(kind=int_4)','logical']
 
-   #:set usekinds = ['dp', 'int_4', '']
+   #:set usekinds = ['dp', 'dp', 'int_4', '']
 
-   #:set defaultFormatType1 = ['"(es14.6)"', '"(i6)"', '"(l1)"']
+   #:set defaultFormatType1 = ['"(es14.6)"', '"(es14.6)"','"(i6)"', '"(l1)"']
 
    #:def lessQnum(el1, el2)
       ${el1}$ < ${el2}$
    #:enddef
 
+   #:def lessQcomplex(el1, el2)
+REAL(${el1}$) < REAL(${el2}$) .OR. (ABS(REAL(${el1}$)-REAL(${el2}$)) < epsilon(MAX(ABS(REAL(${el1}$)), ABS(REAL(${el2}$)))) .and. AIMAG(${el1}$) <  AIMAG(${el2}$))
+   #:enddef
+
    #:def lessQlog(el1, el2)
       (.not. ${el1}$) .and. ${el2}$
    #:enddef
 
-   #:set lessQ = [lessQnum, lessQnum, lessQlog]
+   #:set lessQ = [lessQnum, lessQcomplex, lessQnum, lessQlog]
 
    #:set inst_params = list(zip(nametype1, type1, defaultFormatType1, lessQ))
 

src/pw/cp_linked_list_pw.F

@@ -10,25 +10,29 @@
 MODULE cp_linked_list_pw
    USE kinds, ONLY: dp
    USE pw_types, ONLY: pw_type, pw_p_type
-   USE cp_array_utils, ONLY: cp_3d_r_cp_type
+   USE cp_array_utils, ONLY: cp_3d_r_cp_type, cp_3d_c_cp_type, cp_1d_r_cp_type, cp_1d_c_cp_type
    USE realspace_grid_types, ONLY: realspace_grid_type, realspace_grid_p_type
 #include "../base/base_uses.f90"
 
    #:mute
-      #:set nametype1 = ['3d_r', 'pw', 'rs']
-      #:set type1 = ['REAL(kind=dp),dimension(:,:,:),pointer,contiguous', 'type(pw_type),pointer', 'type(realspace_grid_type),pointer']
+      #:set nametype1 = ['3d_r', '3d_c', '1d_r', '1d_c', 'pw', 'rs']
+      #:set type1 = ['REAL(kind=dp),dimension(:,:,:),pointer,contiguous',&
+         'COMPLEX(kind=dp),dimension(:,:,:),pointer,contiguous',&
+         'REAL(kind=dp),dimension(:),pointer,contiguous',&
+         'COMPLEX(kind=dp),dimension(:),pointer,contiguous',&
+         'type(pw_type),pointer', 'type(realspace_grid_type),pointer']
       #:set type1in = type1
       #:set type1out = type1
 
-      #:set eq = ['=>','=>','=>','=>']
+      #:set eq = ['=>','=>','=>','=>','=>','=>']
       #:set arrayeq = eq
-      #:set type1arrayEl = ['type(cp_3d_r_cp_type)','type(pw_p_type)','type(realspace_grid_p_type)']
+      #:set type1arrayEl = ['type(cp_3d_r_cp_type)','type(cp_3d_c_cp_type)','type(cp_1d_r_cp_type)','type(cp_1d_c_cp_type)','type(pw_p_type)','type(realspace_grid_p_type)']
 
-      #:set arrayEl = ['%array','%pw','%rs_grid']
+      #:set arrayEl = ['%array','%array','%array','%array','%pw','%rs_grid']
 
       #:set private_routines = ''
 
-      #:set default_init = [' => NULL()', ' => NULL()', ' => NULL()']
+      #:set default_init = [' => NULL()', ' => NULL()', ' => NULL()', ' => NULL()', ' => NULL()', ' => NULL()']
 
       #:set common_dir = '../common'
    #:endmute

src/pw/fft_tools.F

@@ -42,7 +42,6 @@ MODULE fft_tools
    USE message_passing,                 ONLY: mp_cart_type,&
                                               mp_comm_null,&
                                               mp_comm_type,&
-                                              mp_para_env_type,&
                                               mp_request_type,&
                                               mp_waitall
    USE offload_api,                     ONLY: offload_free_pinned_mem,&
@@ -214,7 +213,7 @@ END SUBROUTINE init_fft
 !>      10.2007 created [Joost VandeVondele]
 ! **************************************************************************************************
    SUBROUTINE finalize_fft(para_env, wisdom_file)
-      TYPE(mp_para_env_type), POINTER                    :: para_env
+      CLASS(mp_comm_type)                    :: para_env
       CHARACTER(LEN=*), INTENT(IN)                       :: wisdom_file
 
 ! release the FFT scratch pool

src/pw/ps_wavelet_methods.F

@@ -26,7 +26,8 @@ MODULE ps_wavelet_methods
                                               S_FFT_dimensions
    USE pw_grid_types,                   ONLY: pw_grid_type
    USE pw_poisson_types,                ONLY: pw_poisson_parameter_type
-   USE pw_types,                        ONLY: pw_type
+   USE pw_types,                        ONLY: pw_r3d_type,&
+                                              pw_type
    USE util,                            ONLY: get_limit
 #include "../base/base_uses.f90"
 
@@ -43,6 +44,14 @@ MODULE ps_wavelet_methods
              z_slices_to_cp2k_distribution, &
              ps_wavelet_solve
 
+   INTERFACE cp2k_distribution_to_z_slices
+      MODULE PROCEDURE cp2k_distribution_to_z_slices_r3d, cp2k_distribution_to_z_slices_pw
+   END INTERFACE
+
+   INTERFACE z_slices_to_cp2k_distribution
+      MODULE PROCEDURE z_slices_to_cp2k_distribution_pw, z_slices_to_cp2k_distribution_r3d
+   END INTERFACE
+
 CONTAINS
 
 ! **************************************************************************************************
@@ -165,13 +174,135 @@ END SUBROUTINE RS_z_slice_distribution
 !> \param wavelet ...
 !> \param pw_grid ...
 ! **************************************************************************************************
-   SUBROUTINE cp2k_distribution_to_z_slices(density, wavelet, pw_grid)
+   SUBROUTINE cp2k_distribution_to_z_slices_r3d(density, wavelet, pw_grid)
+
+      TYPE(pw_r3d_type), INTENT(IN)                      :: density
+      TYPE(ps_wavelet_type), POINTER                     :: wavelet
+      TYPE(pw_grid_type), POINTER                        :: pw_grid
+
+      CHARACTER(len=*), PARAMETER :: routineN = 'cp2k_distribution_to_z_slices_r3d'
+
+      INTEGER                                            :: dest, handle, i, ii, iproc, j, k, l, &
+                                                            local_z_dim, loz, m, m2, md2, nproc, &
+                                                            should_warn
+      INTEGER, ALLOCATABLE, DIMENSION(:)                 :: rcount, rdispl, scount, sdispl, tmp
+      INTEGER, DIMENSION(2)                              :: cart_pos, lox, loy
+      INTEGER, DIMENSION(3)                              :: lb, ub
+      REAL(KIND=dp)                                      :: max_val_low, max_val_up
+      REAL(KIND=dp), DIMENSION(:), POINTER               :: rbuf, sbuf
+
+      CALL timeset(routineN, handle)
+
+      CPASSERT(ASSOCIATED(wavelet))
+
+      nproc = PRODUCT(pw_grid%para%rs_dims)
+      iproc = pw_grid%para%rs_mpo
+      md2 = wavelet%PS_grid(3)
+      m2 = pw_grid%npts(3)
+      lb(:) = pw_grid%bounds_local(1, :)
+      ub(:) = pw_grid%bounds_local(2, :)
+      local_z_dim = MAX((md2/nproc), 1)
+
+      ALLOCATE (sbuf(PRODUCT(pw_grid%npts_local)))
+      ALLOCATE (rbuf(PRODUCT(wavelet%PS_grid)/nproc))
+      ALLOCATE (scount(nproc), sdispl(nproc), rcount(nproc), rdispl(nproc), tmp(nproc))
+
+      rbuf = 0.0_dp
+      ii = 1
+      DO k = lb(3), ub(3)
+         DO j = lb(2), ub(2)
+            DO i = lb(1), ub(1)
+               sbuf(ii) = density%array(i, j, k)
+               ii = ii + 1
+            END DO
+         END DO
+      END DO
+
+      should_warn = 0
+      IF (wavelet%geocode == 'S' .OR. wavelet%geocode == 'F') THEN
+         max_val_low = 0._dp
+         max_val_up = 0._dp
+         IF (lb(2) == pw_grid%bounds(1, 2)) max_val_low = MAXVAL(ABS(density%array(:, lb(2), :)))
+         IF (ub(2) == pw_grid%bounds(2, 2)) max_val_up = MAXVAL(ABS(density%array(:, ub(2), :)))
+         IF (max_val_low .GE. 0.0001_dp) should_warn = 1
+         IF (max_val_up .GE. 0.0001_dp) should_warn = 1
+         IF (wavelet%geocode == 'F') THEN
+            max_val_low = 0._dp
+            max_val_up = 0._dp
+            IF (lb(1) == pw_grid%bounds(1, 1)) max_val_low = MAXVAL(ABS(density%array(lb(1), :, :)))
+            IF (ub(1) == pw_grid%bounds(2, 1)) max_val_up = MAXVAL(ABS(density%array(ub(1), :, :)))
+            IF (max_val_low .GE. 0.0001_dp) should_warn = 1
+            IF (max_val_up .GE. 0.0001_dp) should_warn = 1
+            max_val_low = 0._dp
+            max_val_up = 0._dp
+            IF (lb(3) == pw_grid%bounds(1, 3)) max_val_low = MAXVAL(ABS(density%array(:, :, lb(3))))
+            IF (ub(3) == pw_grid%bounds(2, 3)) max_val_up = MAXVAL(ABS(density%array(:, :, ub(3))))
+            IF (max_val_low .GE. 0.0001_dp) should_warn = 1
+            IF (max_val_up .GE. 0.0001_dp) should_warn = 1
+         END IF
+      END IF
+
+      CALL pw_grid%para%group%max(should_warn)
+      IF (should_warn > 0 .AND. iproc == 0) &
+         CPWARN("Density non-zero on the edges of the unit cell: wrong results in WAVELET solver")
+
+      DO i = 0, pw_grid%para%rs_dims(1) - 1
+         DO j = 0, pw_grid%para%rs_dims(2) - 1
+            cart_pos = (/i, j/)
+            CALL pw_grid%para%rs_group%rank_cart(cart_pos, dest)
+            IF ((ub(1) .GE. lb(1)) .AND. (ub(2) .GE. lb(2))) THEN
+               IF (dest*local_z_dim .LE. m2) THEN
+                  IF ((dest + 1)*local_z_dim .LE. m2) THEN
+                     scount(dest + 1) = ABS((ub(1) - lb(1) + 1)*(ub(2) - lb(2) + 1)*local_z_dim)
+                  ELSE
+                     scount(dest + 1) = ABS((ub(1) - lb(1) + 1)*(ub(2) - lb(2) + 1)*MOD(m2, local_z_dim))
+                  END IF
+               ELSE
+                  scount(dest + 1) = 0
+               END IF
+            ELSE
+               scount(dest + 1) = 0
+            END IF
+            lox = get_limit(pw_grid%npts(1), pw_grid%para%rs_dims(1), i)
+            loy = get_limit(pw_grid%npts(2), pw_grid%para%rs_dims(2), j)
+            IF (iproc*local_z_dim .LE. m2) THEN
+               IF ((iproc + 1)*local_z_dim .LE. m2) THEN
+                  loz = local_z_dim
+               ELSE
+                  loz = MOD(m2, local_z_dim)
+               END IF
+               ii = 1
+               DO k = 1, loz
+                  DO l = loy(1), loy(2)
+                     DO m = lox(1), lox(2)
+                        wavelet%rho_z_sliced(m, l, k) = rbuf(ii + rdispl(dest + 1))
+                        ii = ii + 1
+                     END DO
+                  END DO
+               END DO
+            END IF
+         END DO
+      END DO
+
+      DEALLOCATE (sbuf, rbuf, scount, sdispl, rcount, rdispl, tmp)
+
+      CALL timestop(handle)
+
+   END SUBROUTINE cp2k_distribution_to_z_slices_r3d
+
+! **************************************************************************************************
+!> \brief ...
+!> \param density ...
+!> \param wavelet ...
+!> \param pw_grid ...
+! **************************************************************************************************
+   SUBROUTINE cp2k_distribution_to_z_slices_pw(density, wavelet, pw_grid)
 
       TYPE(pw_type), INTENT(IN)                          :: density
       TYPE(ps_wavelet_type), POINTER                     :: wavelet
       TYPE(pw_grid_type), POINTER                        :: pw_grid
 
-      CHARACTER(len=*), PARAMETER :: routineN = 'cp2k_distribution_to_z_slices'
+      CHARACTER(len=*), PARAMETER :: routineN = 'cp2k_distribution_to_z_slices_pw'
 
       INTEGER                                            :: dest, handle, i, ii, iproc, j, k, l, &
                                                             local_z_dim, loz, m, m2, md2, nproc, &
@@ -313,15 +444,137 @@ SUBROUTINE cp2k_distribution_to_z_slices(density, wavelet, pw_grid)
 
       CALL timestop(handle)
 
-   END SUBROUTINE cp2k_distribution_to_z_slices
+   END SUBROUTINE cp2k_distribution_to_z_slices_pw
+
+! **************************************************************************************************
+!> \brief ...
+!> \param density ...
+!> \param wavelet ...
+!> \param pw_grid ...
+! **************************************************************************************************
+   SUBROUTINE z_slices_to_cp2k_distribution_r3d(density, wavelet, pw_grid)
+
+      TYPE(pw_r3d_type), INTENT(IN)                      :: density
+      TYPE(ps_wavelet_type), POINTER                     :: wavelet
+      TYPE(pw_grid_type), POINTER                        :: pw_grid
+
+      INTEGER                                            :: dest, i, ii, iproc, j, k, l, &
+                                                            local_z_dim, loz, m, m2, md2, nproc
+      INTEGER, ALLOCATABLE, DIMENSION(:)                 :: rcount, rdispl, scount, sdispl, tmp
+      INTEGER, DIMENSION(2)                              :: cart_pos, lox, loy, min_x, min_y
+      INTEGER, DIMENSION(3)                              :: lb, ub
+      REAL(KIND=dp), DIMENSION(:), POINTER               :: rbuf, sbuf
+
+      CPASSERT(ASSOCIATED(wavelet))
+
+      nproc = PRODUCT(pw_grid%para%rs_dims)
+      iproc = pw_grid%para%rs_mpo
+      md2 = wavelet%PS_grid(3)
+      m2 = pw_grid%npts(3)
+
+      lb(:) = pw_grid%bounds_local(1, :)
+      ub(:) = pw_grid%bounds_local(2, :)
+
+      local_z_dim = MAX((md2/nproc), 1)
+
+      ALLOCATE (rbuf(PRODUCT(pw_grid%npts_local)))
+      ALLOCATE (sbuf(PRODUCT(wavelet%PS_grid)/nproc))
+      ALLOCATE (scount(nproc), sdispl(nproc), rcount(nproc), rdispl(nproc), tmp(nproc))
+      scount = 0
+      rcount = 0
+      rbuf = 0.0_dp
+      ii = 1
+      IF (iproc*local_z_dim .LE. m2) THEN
+         IF ((iproc + 1)*local_z_dim .LE. m2) THEN
+            loz = local_z_dim
+         ELSE
+            loz = MOD(m2, local_z_dim)
+         END IF
+      ELSE
+         loz = 0
+      END IF
+
+      min_x = get_limit(pw_grid%npts(1), pw_grid%para%rs_dims(1), 0)
+      min_y = get_limit(pw_grid%npts(2), pw_grid%para%rs_dims(2), 0)
+      DO i = 0, pw_grid%para%rs_dims(1) - 1
+         DO j = 0, pw_grid%para%rs_dims(2) - 1
+            cart_pos = (/i, j/)
+            CALL pw_grid%para%rs_group%rank_cart(cart_pos, dest)
+            IF ((ub(1) .GE. lb(1)) .AND. (ub(2) .GE. lb(2))) THEN
+               IF (dest*local_z_dim .LE. m2) THEN
+                  IF ((dest + 1)*local_z_dim .LE. m2) THEN
+                     rcount(dest + 1) = ABS((ub(1) - lb(1) + 1)*(ub(2) - lb(2) + 1)*local_z_dim)
+                  ELSE
+                     rcount(dest + 1) = ABS((ub(1) - lb(1) + 1)*(ub(2) - lb(2) + 1)*MOD(m2, local_z_dim))
+                  END IF
+               ELSE
+                  rcount(dest + 1) = 0
+               END IF
+            ELSE
+               rcount(dest + 1) = 0
+            END IF
+            lox = get_limit(pw_grid%npts(1), pw_grid%para%rs_dims(1), i)
+            loy = get_limit(pw_grid%npts(2), pw_grid%para%rs_dims(2), j)
+            IF ((lox(2) .GE. lox(1)) .AND. (loy(2) .GE. loy(1))) THEN
+               scount(dest + 1) = ABS((lox(2) - lox(1) + 1)*(loy(2) - loy(1) + 1)*loz)
+               DO k = lox(1) - min_x(1) + 1, lox(2) - min_x(1) + 1
+                  DO l = loy(1) - min_y(1) + 1, loy(2) - min_y(1) + 1
+                     DO m = 1, loz
+                        sbuf(ii) = wavelet%rho_z_sliced(k, l, m)
+                        ii = ii + 1
+                     END DO
+                  END DO
+               END DO
+            ELSE
+               scount(dest + 1) = 0
+            END IF
+         END DO
+      END DO
+      sdispl(1) = 0
+      rdispl(1) = 0
+      DO i = 2, nproc
+         sdispl(i) = sdispl(i - 1) + scount(i - 1)
+         rdispl(i) = rdispl(i - 1) + rcount(i - 1)
+      END DO
+      CALL pw_grid%para%rs_group%alltoall(sbuf, scount, sdispl, rbuf, rcount, rdispl)
+
+      !!!! and now, how to put the right cubes to the right position!!!!!!
+
+      DO i = 0, pw_grid%para%rs_dims(1) - 1
+         DO j = 0, pw_grid%para%rs_dims(2) - 1
+            cart_pos = (/i, j/)
+            CALL pw_grid%para%rs_group%rank_cart(cart_pos, dest)
+            IF (dest*local_z_dim .LE. m2) THEN
+               IF ((dest + 1)*local_z_dim .LE. m2) THEN
+                  loz = local_z_dim
+               ELSE
+                  loz = MOD(m2, local_z_dim)
+               END IF
+               ii = 1
+               IF (lb(3) + (dest*local_z_dim) .LE. ub(3)) THEN
+                  DO m = lb(1), ub(1)
+                     DO l = lb(2), ub(2)
+                        DO k = lb(3) + (dest*local_z_dim), lb(3) + (dest*local_z_dim) + loz - 1
+                           density%array(m, l, k) = rbuf(ii + rdispl(dest + 1))
+                           ii = ii + 1
+                        END DO
+                     END DO
+                  END DO
+               END IF
+            END IF
+         END DO
+      END DO
+      DEALLOCATE (sbuf, rbuf, scount, sdispl, rcount, rdispl, tmp)
+
+   END SUBROUTINE z_slices_to_cp2k_distribution_r3d
 
 ! **************************************************************************************************
 !> \brief ...
 !> \param density ...
 !> \param wavelet ...
 !> \param pw_grid ...
 ! **************************************************************************************************
-   SUBROUTINE z_slices_to_cp2k_distribution(density, wavelet, pw_grid)
+   SUBROUTINE z_slices_to_cp2k_distribution_pw(density, wavelet, pw_grid)
 
       TYPE(pw_type), INTENT(IN)                          :: density
       TYPE(ps_wavelet_type), POINTER                     :: wavelet
@@ -435,7 +688,7 @@ SUBROUTINE z_slices_to_cp2k_distribution(density, wavelet, pw_grid)
       END DO
       DEALLOCATE (sbuf, rbuf, scount, sdispl, rcount, rdispl, tmp)
 
-   END SUBROUTINE z_slices_to_cp2k_distribution
+   END SUBROUTINE z_slices_to_cp2k_distribution_pw
 
 ! **************************************************************************************************
 !> \brief ...