Line data Source code
1 : !--------------------------------------------------------------------------------------------------!
2 : ! CP2K: A general program to perform molecular dynamics simulations !
3 : ! Copyright 2000-2026 CP2K developers group <https://cp2k.org> !
4 : ! !
5 : ! SPDX-License-Identifier: GPL-2.0-or-later !
6 : !--------------------------------------------------------------------------------------------------!
7 :
8 : ! **************************************************************************************************
9 : !> \brief Definition and initialisation of the ps_wavelet data type.
10 : !> \history 01.2014 Renamed from ps_wavelet_types to disentangle dependencies (Ole Schuett)
11 : !> \author Florian Schiffmann (09.2007,fschiff)
12 : ! **************************************************************************************************
13 : MODULE ps_wavelet_methods
14 :
15 : USE bibliography, ONLY: Genovese2006,&
16 : Genovese2007,&
17 : cite_reference
18 : USE kinds, ONLY: dp
19 : USE ps_wavelet_kernel, ONLY: createKernel
20 : USE ps_wavelet_types, ONLY: WAVELET0D,&
21 : WAVELET2D,&
22 : ps_wavelet_release,&
23 : ps_wavelet_type
24 : USE ps_wavelet_util, ONLY: F_FFT_dimensions,&
25 : PSolver,&
26 : P_FFT_dimensions,&
27 : S_FFT_dimensions
28 : USE pw_grid_types, ONLY: pw_grid_type
29 : USE pw_poisson_types, ONLY: pw_poisson_parameter_type
30 : USE pw_types, ONLY: pw_r3d_rs_type
31 : USE util, ONLY: get_limit
32 : #include "../base/base_uses.f90"
33 :
34 : IMPLICIT NONE
35 :
36 : PRIVATE
37 :
38 : CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'ps_wavelet_methods'
39 :
40 : ! *** Public data types ***
41 :
42 : PUBLIC :: ps_wavelet_create, &
43 : cp2k_distribution_to_z_slices, &
44 : z_slices_to_cp2k_distribution, &
45 : ps_wavelet_solve
46 :
47 : CONTAINS
48 :
49 : ! **************************************************************************************************
50 : !> \brief creates the ps_wavelet_type which is needed for the link to
51 : !> the Poisson Solver of Luigi Genovese
52 : !> \param poisson_params ...
53 : !> \param wavelet wavelet to create
54 : !> \param pw_grid the grid that is used to create the wavelet kernel
55 : !> \author Flroian Schiffmann
56 : ! **************************************************************************************************
57 856 : SUBROUTINE ps_wavelet_create(poisson_params, wavelet, pw_grid)
58 : TYPE(pw_poisson_parameter_type), INTENT(IN) :: poisson_params
59 : TYPE(ps_wavelet_type), POINTER :: wavelet
60 : TYPE(pw_grid_type), POINTER :: pw_grid
61 :
62 : CHARACTER(len=*), PARAMETER :: routineN = 'ps_wavelet_create'
63 :
64 : INTEGER :: handle
65 : REAL(KIND=dp) :: hx, hy, hz
66 :
67 856 : CALL timeset(routineN, handle)
68 :
69 856 : CALL cite_reference(Genovese2006)
70 856 : CALL cite_reference(Genovese2007)
71 :
72 856 : IF (ASSOCIATED(wavelet)) THEN
73 0 : CALL ps_wavelet_release(wavelet)
74 : NULLIFY (wavelet)
75 : END IF
76 :
77 7704 : ALLOCATE (wavelet)
78 :
79 : NULLIFY (wavelet%karray, wavelet%rho_z_sliced)
80 :
81 856 : wavelet%geocode = poisson_params%wavelet_geocode
82 856 : wavelet%method = poisson_params%wavelet_method
83 856 : wavelet%special_dimension = poisson_params%wavelet_special_dimension
84 856 : wavelet%itype_scf = poisson_params%wavelet_scf_type
85 856 : wavelet%datacode = "D"
86 :
87 856 : CALL set_wavelet_axis(wavelet)
88 856 : hx = pw_grid%dr(wavelet%axis(1))
89 856 : hy = pw_grid%dr(wavelet%axis(2))
90 856 : hz = pw_grid%dr(wavelet%axis(3))
91 :
92 856 : IF (poisson_params%wavelet_method == WAVELET0D) THEN
93 524 : IF (hx /= hy) THEN
94 0 : CPABORT("Poisson solver for non cubic cells not yet implemented")
95 : END IF
96 524 : IF (hz /= hy) THEN
97 0 : CPABORT("Poisson solver for non cubic cells not yet implemented")
98 : END IF
99 : END IF
100 :
101 856 : CALL RS_z_slice_distribution(wavelet, pw_grid)
102 :
103 856 : CALL timestop(handle)
104 856 : END SUBROUTINE ps_wavelet_create
105 :
106 : ! **************************************************************************************************
107 : !> \brief ...
108 : !> \param wavelet ...
109 : !> \param pw_grid ...
110 : ! **************************************************************************************************
111 856 : SUBROUTINE RS_z_slice_distribution(wavelet, pw_grid)
112 :
113 : TYPE(ps_wavelet_type), POINTER :: wavelet
114 : TYPE(pw_grid_type), POINTER :: pw_grid
115 :
116 : CHARACTER(len=*), PARAMETER :: routineN = 'RS_z_slice_distribution'
117 :
118 : CHARACTER(LEN=1) :: geocode
119 : INTEGER :: handle, iproc, m1, m2, m3, md1, md2, &
120 : md3, n1, n2, n3, nd1, nd2, nd3, nproc, &
121 : nx, ny, nz, z_dim
122 : REAL(KIND=dp) :: hx, hy, hz
123 :
124 856 : CALL timeset(routineN, handle)
125 2568 : nproc = PRODUCT(pw_grid%para%group%num_pe_cart)
126 856 : iproc = pw_grid%para%group%mepos
127 856 : geocode = wavelet%geocode
128 856 : CALL get_wavelet_grid(wavelet, pw_grid, nx, ny, nz, hx, hy, hz)
129 :
130 : !calculate Dimensions for the z-distributed density and for the kernel
131 :
132 856 : IF (geocode == 'P') THEN
133 326 : CALL P_FFT_dimensions(nx, ny, nz, m1, m2, m3, n1, n2, n3, md1, md2, md3, nd1, nd2, nd3, nproc)
134 530 : ELSE IF (geocode == 'S') THEN
135 6 : CALL S_FFT_dimensions(nx, ny, nz, m1, m2, m3, n1, n2, n3, md1, md2, md3, nd1, nd2, nd3, nproc)
136 524 : ELSE IF (geocode == 'F') THEN
137 524 : CALL F_FFT_dimensions(nx, ny, nz, m1, m2, m3, n1, n2, n3, md1, md2, md3, nd1, nd2, nd3, nproc)
138 : END IF
139 :
140 856 : wavelet%PS_grid(1) = md1
141 856 : wavelet%PS_grid(2) = md3
142 856 : wavelet%PS_grid(3) = md2
143 856 : z_dim = md2/nproc
144 : !!!!!!!!! indices y and z are interchanged !!!!!!!
145 4280 : ALLOCATE (wavelet%rho_z_sliced(md1, md3, z_dim))
146 :
147 : CALL createKernel(geocode, nx, ny, nz, hx, hy, hz, wavelet%itype_scf, iproc, nproc, wavelet%karray, &
148 856 : pw_grid%para%group)
149 :
150 856 : CALL timestop(handle)
151 1712 : END SUBROUTINE RS_z_slice_distribution
152 :
153 : ! **************************************************************************************************
154 : !> \brief ...
155 : !> \param density ...
156 : !> \param wavelet ...
157 : !> \param pw_grid ...
158 : ! **************************************************************************************************
159 33295 : SUBROUTINE cp2k_distribution_to_z_slices(density, wavelet, pw_grid)
160 :
161 : TYPE(pw_r3d_rs_type), INTENT(IN) :: density
162 : TYPE(ps_wavelet_type), POINTER :: wavelet
163 : TYPE(pw_grid_type), POINTER :: pw_grid
164 :
165 : CHARACTER(len=*), PARAMETER :: routineN = 'cp2k_distribution_to_z_slices'
166 :
167 : INTEGER :: dest, handle, i, ii, iproc, j, k, l, &
168 : local_z_dim, loz, m, m2, md2, nproc, &
169 : should_warn
170 33295 : INTEGER, ALLOCATABLE, DIMENSION(:) :: rcount, rdispl, scount, sdispl, tmp
171 : INTEGER, DIMENSION(2) :: cart_pos, lox, loy
172 : INTEGER, DIMENSION(3) :: lb, ub
173 : REAL(KIND=dp) :: max_val_low, max_val_up
174 33295 : REAL(KIND=dp), DIMENSION(:), POINTER :: rbuf, sbuf
175 :
176 33295 : CALL timeset(routineN, handle)
177 :
178 33295 : CPASSERT(ASSOCIATED(wavelet))
179 :
180 33295 : IF (.NOT. wavelet_axis_is_identity(wavelet)) THEN
181 36 : CALL warn_density_edges(density, wavelet, pw_grid)
182 36 : CALL cp2k_distribution_to_z_slices_permuted(density, wavelet, pw_grid)
183 36 : CALL timestop(handle)
184 36 : RETURN
185 : END IF
186 :
187 99777 : nproc = PRODUCT(pw_grid%para%group%num_pe_cart)
188 33259 : iproc = pw_grid%para%group%mepos
189 33259 : md2 = wavelet%PS_grid(3)
190 33259 : m2 = pw_grid%npts(3)
191 133036 : lb(:) = pw_grid%bounds_local(1, :)
192 133036 : ub(:) = pw_grid%bounds_local(2, :)
193 33259 : local_z_dim = MAX((md2/nproc), 1)
194 :
195 199554 : ALLOCATE (sbuf(PRODUCT(pw_grid%npts_local)))
196 199554 : ALLOCATE (rbuf(PRODUCT(wavelet%PS_grid)/nproc))
197 232813 : ALLOCATE (scount(nproc), sdispl(nproc), rcount(nproc), rdispl(nproc), tmp(nproc))
198 :
199 956352052 : rbuf = 0.0_dp
200 33259 : ii = 1
201 982922 : DO k = lb(3), ub(3)
202 34523189 : DO j = lb(2), ub(2)
203 968147046 : DO i = lb(1), ub(1)
204 933657116 : sbuf(ii) = density%array(i, j, k)
205 967197383 : ii = ii + 1
206 : END DO
207 : END DO
208 : END DO
209 :
210 33259 : should_warn = 0
211 33259 : IF (wavelet%geocode == 'S' .OR. wavelet%geocode == 'F') THEN
212 15936 : max_val_low = 0._dp
213 15936 : max_val_up = 0._dp
214 17108942 : IF (lb(2) == pw_grid%bounds(1, 2)) max_val_low = MAXVAL(ABS(density%array(:, lb(2), :)))
215 17108942 : IF (ub(2) == pw_grid%bounds(2, 2)) max_val_up = MAXVAL(ABS(density%array(:, ub(2), :)))
216 15936 : IF (max_val_low >= 0.0001_dp) should_warn = 1
217 15936 : IF (max_val_up >= 0.0001_dp) should_warn = 1
218 15936 : IF (wavelet%geocode == 'F') THEN
219 15918 : max_val_low = 0._dp
220 15918 : max_val_up = 0._dp
221 16898574 : IF (lb(1) == pw_grid%bounds(1, 1)) max_val_low = MAXVAL(ABS(density%array(lb(1), :, :)))
222 16898574 : IF (ub(1) == pw_grid%bounds(2, 1)) max_val_up = MAXVAL(ABS(density%array(ub(1), :, :)))
223 15918 : IF (max_val_low >= 0.0001_dp) should_warn = 1
224 15918 : IF (max_val_up >= 0.0001_dp) should_warn = 1
225 15918 : max_val_low = 0._dp
226 15918 : max_val_up = 0._dp
227 17081708 : IF (lb(3) == pw_grid%bounds(1, 3)) max_val_low = MAXVAL(ABS(density%array(:, :, lb(3))))
228 17081708 : IF (ub(3) == pw_grid%bounds(2, 3)) max_val_up = MAXVAL(ABS(density%array(:, :, ub(3))))
229 15918 : IF (max_val_low >= 0.0001_dp) should_warn = 1
230 15918 : IF (max_val_up >= 0.0001_dp) should_warn = 1
231 : END IF
232 : END IF
233 :
234 33259 : CALL pw_grid%para%group%max(should_warn)
235 33259 : IF (should_warn > 0 .AND. iproc == 0) THEN
236 4544 : CPWARN("Density non-zero on the edges of the unit cell: wrong results in WAVELET solver")
237 : END IF
238 79694 : DO i = 0, pw_grid%para%group%num_pe_cart(1) - 1
239 126129 : DO j = 0, pw_grid%para%group%num_pe_cart(2) - 1
240 139305 : cart_pos = [i, j]
241 46435 : CALL pw_grid%para%group%rank_cart(cart_pos, dest)
242 46435 : IF ((ub(1) >= lb(1)) .AND. (ub(2) >= lb(2))) THEN
243 46435 : IF (dest*local_z_dim <= m2) THEN
244 46435 : IF ((dest + 1)*local_z_dim <= m2) THEN
245 41822 : scount(dest + 1) = ABS((ub(1) - lb(1) + 1)*(ub(2) - lb(2) + 1)*local_z_dim)
246 : ELSE
247 4613 : scount(dest + 1) = ABS((ub(1) - lb(1) + 1)*(ub(2) - lb(2) + 1)*MOD(m2, local_z_dim))
248 : END IF
249 : ELSE
250 0 : scount(dest + 1) = 0
251 : END IF
252 : ELSE
253 0 : scount(dest + 1) = 0
254 : END IF
255 46435 : lox = get_limit(pw_grid%npts(1), pw_grid%para%group%num_pe_cart(1), i)
256 46435 : loy = get_limit(pw_grid%npts(2), pw_grid%para%group%num_pe_cart(2), j)
257 92870 : IF ((lox(2) >= lox(1)) .AND. (loy(2) >= loy(1))) THEN
258 46435 : IF (iproc*local_z_dim <= m2) THEN
259 46435 : IF ((iproc + 1)*local_z_dim <= m2) THEN
260 41822 : rcount(dest + 1) = ABS((lox(2) - lox(1) + 1)*(loy(2) - loy(1) + 1)*local_z_dim)
261 : ELSE
262 4613 : rcount(dest + 1) = ABS((lox(2) - lox(1) + 1)*(loy(2) - loy(1) + 1)*MOD(m2, local_z_dim))
263 : END IF
264 : ELSE
265 0 : rcount(dest + 1) = 0
266 : END IF
267 : ELSE
268 0 : rcount(dest + 1) = 0
269 : END IF
270 :
271 : END DO
272 : END DO
273 33259 : sdispl(1) = 0
274 33259 : rdispl(1) = 0
275 46435 : DO i = 2, nproc
276 13176 : sdispl(i) = sdispl(i - 1) + scount(i - 1)
277 46435 : rdispl(i) = rdispl(i - 1) + rcount(i - 1)
278 : END DO
279 2846327961 : CALL pw_grid%para%group%alltoall(sbuf, scount, sdispl, rbuf, rcount, rdispl)
280 : !!!! and now, how to put the right cubes to the right position!!!!!!
281 :
282 980841882 : wavelet%rho_z_sliced = 0.0_dp
283 :
284 79694 : DO i = 0, pw_grid%para%group%num_pe_cart(1) - 1
285 126129 : DO j = 0, pw_grid%para%group%num_pe_cart(2) - 1
286 139305 : cart_pos = [i, j]
287 46435 : CALL pw_grid%para%group%rank_cart(cart_pos, dest)
288 :
289 46435 : lox = get_limit(pw_grid%npts(1), pw_grid%para%group%num_pe_cart(1), i)
290 46435 : loy = get_limit(pw_grid%npts(2), pw_grid%para%group%num_pe_cart(2), j)
291 92870 : IF (iproc*local_z_dim <= m2) THEN
292 46435 : IF ((iproc + 1)*local_z_dim <= m2) THEN
293 : loz = local_z_dim
294 : ELSE
295 4613 : loz = MOD(m2, local_z_dim)
296 : END IF
297 46435 : ii = 1
298 996098 : DO k = 1, loz
299 34536365 : DO l = loy(1), loy(2)
300 968147046 : DO m = lox(1), lox(2)
301 933657116 : wavelet%rho_z_sliced(m, l, k) = rbuf(ii + rdispl(dest + 1))
302 967197383 : ii = ii + 1
303 : END DO
304 : END DO
305 : END DO
306 : END IF
307 : END DO
308 : END DO
309 :
310 33259 : DEALLOCATE (sbuf, rbuf, scount, sdispl, rcount, rdispl, tmp)
311 :
312 33259 : CALL timestop(handle)
313 :
314 66590 : END SUBROUTINE cp2k_distribution_to_z_slices
315 :
316 : ! **************************************************************************************************
317 : !> \brief ...
318 : !> \param density ...
319 : !> \param wavelet ...
320 : !> \param pw_grid ...
321 : ! **************************************************************************************************
322 33295 : SUBROUTINE z_slices_to_cp2k_distribution(density, wavelet, pw_grid)
323 :
324 : TYPE(pw_r3d_rs_type), INTENT(IN) :: density
325 : TYPE(ps_wavelet_type), POINTER :: wavelet
326 : TYPE(pw_grid_type), POINTER :: pw_grid
327 :
328 : INTEGER :: dest, i, ii, iproc, j, k, l, &
329 : local_z_dim, loz, m, m2, md2, nproc
330 33295 : INTEGER, ALLOCATABLE, DIMENSION(:) :: rcount, rdispl, scount, sdispl, tmp
331 : INTEGER, DIMENSION(2) :: cart_pos, lox, loy, min_x, min_y
332 : INTEGER, DIMENSION(3) :: lb, ub
333 33295 : REAL(KIND=dp), DIMENSION(:), POINTER :: rbuf, sbuf
334 :
335 0 : CPASSERT(ASSOCIATED(wavelet))
336 :
337 33295 : IF (.NOT. wavelet_axis_is_identity(wavelet)) THEN
338 36 : CALL z_slices_to_cp2k_distribution_permuted(density, wavelet, pw_grid)
339 36 : RETURN
340 : END IF
341 :
342 99777 : nproc = PRODUCT(pw_grid%para%group%num_pe_cart)
343 33259 : iproc = pw_grid%para%group%mepos
344 33259 : md2 = wavelet%PS_grid(3)
345 33259 : m2 = pw_grid%npts(3)
346 :
347 133036 : lb(:) = pw_grid%bounds_local(1, :)
348 133036 : ub(:) = pw_grid%bounds_local(2, :)
349 :
350 33259 : local_z_dim = MAX((md2/nproc), 1)
351 :
352 199554 : ALLOCATE (rbuf(PRODUCT(pw_grid%npts_local)))
353 199554 : ALLOCATE (sbuf(PRODUCT(wavelet%PS_grid)/nproc))
354 232813 : ALLOCATE (scount(nproc), sdispl(nproc), rcount(nproc), rdispl(nproc), tmp(nproc))
355 33259 : scount = 0
356 33259 : rcount = 0
357 933690375 : rbuf = 0.0_dp
358 33259 : ii = 1
359 33259 : IF (iproc*local_z_dim <= m2) THEN
360 33259 : IF ((iproc + 1)*local_z_dim <= m2) THEN
361 : loz = local_z_dim
362 : ELSE
363 2816 : loz = MOD(m2, local_z_dim)
364 : END IF
365 : ELSE
366 : loz = 0
367 : END IF
368 :
369 33259 : min_x = get_limit(pw_grid%npts(1), pw_grid%para%group%num_pe_cart(1), 0)
370 33259 : min_y = get_limit(pw_grid%npts(2), pw_grid%para%group%num_pe_cart(2), 0)
371 79694 : DO i = 0, pw_grid%para%group%num_pe_cart(1) - 1
372 126129 : DO j = 0, pw_grid%para%group%num_pe_cart(2) - 1
373 139305 : cart_pos = [i, j]
374 46435 : CALL pw_grid%para%group%rank_cart(cart_pos, dest)
375 46435 : IF ((ub(1) >= lb(1)) .AND. (ub(2) >= lb(2))) THEN
376 46435 : IF (dest*local_z_dim <= m2) THEN
377 46435 : IF ((dest + 1)*local_z_dim <= m2) THEN
378 41822 : rcount(dest + 1) = ABS((ub(1) - lb(1) + 1)*(ub(2) - lb(2) + 1)*local_z_dim)
379 : ELSE
380 4613 : rcount(dest + 1) = ABS((ub(1) - lb(1) + 1)*(ub(2) - lb(2) + 1)*MOD(m2, local_z_dim))
381 : END IF
382 : ELSE
383 0 : rcount(dest + 1) = 0
384 : END IF
385 : ELSE
386 0 : rcount(dest + 1) = 0
387 : END IF
388 46435 : lox = get_limit(pw_grid%npts(1), pw_grid%para%group%num_pe_cart(1), i)
389 46435 : loy = get_limit(pw_grid%npts(2), pw_grid%para%group%num_pe_cart(2), j)
390 92870 : IF ((lox(2) >= lox(1)) .AND. (loy(2) >= loy(1))) THEN
391 46435 : scount(dest + 1) = ABS((lox(2) - lox(1) + 1)*(loy(2) - loy(1) + 1)*loz)
392 996098 : DO k = lox(1) - min_x(1) + 1, lox(2) - min_x(1) + 1
393 34536365 : DO l = loy(1) - min_y(1) + 1, loy(2) - min_y(1) + 1
394 968147046 : DO m = 1, loz
395 933657116 : sbuf(ii) = wavelet%rho_z_sliced(k, l, m)
396 967197383 : ii = ii + 1
397 : END DO
398 : END DO
399 : END DO
400 : ELSE
401 0 : scount(dest + 1) = 0
402 : END IF
403 : END DO
404 : END DO
405 33259 : sdispl(1) = 0
406 33259 : rdispl(1) = 0
407 46435 : DO i = 2, nproc
408 13176 : sdispl(i) = sdispl(i - 1) + scount(i - 1)
409 46435 : rdispl(i) = rdispl(i - 1) + rcount(i - 1)
410 : END DO
411 2823666284 : CALL pw_grid%para%group%alltoall(sbuf, scount, sdispl, rbuf, rcount, rdispl)
412 :
413 : !!!! and now, how to put the right cubes to the right position!!!!!!
414 :
415 79694 : DO i = 0, pw_grid%para%group%num_pe_cart(1) - 1
416 126129 : DO j = 0, pw_grid%para%group%num_pe_cart(2) - 1
417 139305 : cart_pos = [i, j]
418 46435 : CALL pw_grid%para%group%rank_cart(cart_pos, dest)
419 92870 : IF (dest*local_z_dim <= m2) THEN
420 46435 : IF ((dest + 1)*local_z_dim <= m2) THEN
421 : loz = local_z_dim
422 : ELSE
423 4613 : loz = MOD(m2, local_z_dim)
424 : END IF
425 46435 : ii = 1
426 46435 : IF (lb(3) + (dest*local_z_dim) <= ub(3)) THEN
427 996098 : DO m = lb(1), ub(1)
428 34536365 : DO l = lb(2), ub(2)
429 968147046 : DO k = lb(3) + (dest*local_z_dim), lb(3) + (dest*local_z_dim) + loz - 1
430 933657116 : density%array(m, l, k) = rbuf(ii + rdispl(dest + 1))
431 967197383 : ii = ii + 1
432 : END DO
433 : END DO
434 : END DO
435 : END IF
436 : END IF
437 : END DO
438 : END DO
439 33259 : DEALLOCATE (sbuf, rbuf, scount, sdispl, rcount, rdispl, tmp)
440 :
441 66590 : END SUBROUTINE z_slices_to_cp2k_distribution
442 :
443 : ! **************************************************************************************************
444 : !> \brief Set the internal Wavelet solver axes for the requested boundary conditions.
445 : !> \param wavelet ...
446 : ! **************************************************************************************************
447 856 : SUBROUTINE set_wavelet_axis(wavelet)
448 :
449 : TYPE(ps_wavelet_type), POINTER :: wavelet
450 :
451 3424 : wavelet%axis = [1, 2, 3]
452 :
453 856 : IF (wavelet%method == WAVELET2D) THEN
454 6 : SELECT CASE (wavelet%special_dimension)
455 : CASE (1)
456 8 : wavelet%axis = [2, 1, 3]
457 : CASE (2)
458 8 : wavelet%axis = [1, 2, 3]
459 : CASE (3)
460 8 : wavelet%axis = [1, 3, 2]
461 : CASE DEFAULT
462 6 : CPABORT("Invalid isolated dimension for WAVELET 2D")
463 : END SELECT
464 : END IF
465 :
466 856 : END SUBROUTINE set_wavelet_axis
467 :
468 : ! **************************************************************************************************
469 : !> \brief Return grid sizes and spacings in the internal Wavelet solver axes.
470 : !> \param wavelet ...
471 : !> \param pw_grid ...
472 : !> \param nx ...
473 : !> \param ny ...
474 : !> \param nz ...
475 : !> \param hx ...
476 : !> \param hy ...
477 : !> \param hz ...
478 : ! **************************************************************************************************
479 34151 : SUBROUTINE get_wavelet_grid(wavelet, pw_grid, nx, ny, nz, hx, hy, hz)
480 :
481 : TYPE(ps_wavelet_type), POINTER :: wavelet
482 : TYPE(pw_grid_type), POINTER :: pw_grid
483 : INTEGER, INTENT(OUT) :: nx, ny, nz
484 : REAL(KIND=dp), INTENT(OUT) :: hx, hy, hz
485 :
486 34151 : nx = pw_grid%npts(wavelet%axis(1))
487 34151 : ny = pw_grid%npts(wavelet%axis(2))
488 34151 : nz = pw_grid%npts(wavelet%axis(3))
489 34151 : hx = pw_grid%dr(wavelet%axis(1))
490 34151 : hy = pw_grid%dr(wavelet%axis(2))
491 34151 : hz = pw_grid%dr(wavelet%axis(3))
492 :
493 34151 : END SUBROUTINE get_wavelet_grid
494 :
495 : ! **************************************************************************************************
496 : !> \brief Return whether CP2K and internal Wavelet solver axes are identical.
497 : !> \param wavelet ...
498 : !> \return ...
499 : ! **************************************************************************************************
500 66590 : FUNCTION wavelet_axis_is_identity(wavelet) RESULT(is_identity)
501 :
502 : TYPE(ps_wavelet_type), POINTER :: wavelet
503 : LOGICAL :: is_identity
504 :
505 266180 : is_identity = ALL(wavelet%axis == [1, 2, 3])
506 :
507 66590 : END FUNCTION wavelet_axis_is_identity
508 :
509 : ! **************************************************************************************************
510 : !> \brief Warn if the density is non-zero at isolated Wavelet cell edges.
511 : !> \param density ...
512 : !> \param wavelet ...
513 : !> \param pw_grid ...
514 : ! **************************************************************************************************
515 36 : SUBROUTINE warn_density_edges(density, wavelet, pw_grid)
516 :
517 : TYPE(pw_r3d_rs_type), INTENT(IN) :: density
518 : TYPE(ps_wavelet_type), POINTER :: wavelet
519 : TYPE(pw_grid_type), POINTER :: pw_grid
520 :
521 : INTEGER :: idir, iproc, should_warn
522 :
523 36 : should_warn = 0
524 36 : iproc = pw_grid%para%group%mepos
525 :
526 36 : IF (wavelet%geocode == 'S') THEN
527 36 : CALL update_edge_warning(density, pw_grid, wavelet%special_dimension, should_warn)
528 0 : ELSE IF (wavelet%geocode == 'F') THEN
529 0 : DO idir = 1, 3
530 0 : CALL update_edge_warning(density, pw_grid, idir, should_warn)
531 : END DO
532 : END IF
533 :
534 36 : CALL pw_grid%para%group%max(should_warn)
535 36 : IF (should_warn > 0 .AND. iproc == 0) THEN
536 0 : CPWARN("Density non-zero on the edges of the unit cell: wrong results in WAVELET solver")
537 : END IF
538 :
539 36 : END SUBROUTINE warn_density_edges
540 :
541 : ! **************************************************************************************************
542 : !> \brief Update the density edge warning for one real-space direction.
543 : !> \param density ...
544 : !> \param pw_grid ...
545 : !> \param direction ...
546 : !> \param should_warn ...
547 : ! **************************************************************************************************
548 36 : SUBROUTINE update_edge_warning(density, pw_grid, direction, should_warn)
549 :
550 : TYPE(pw_r3d_rs_type), INTENT(IN) :: density
551 : TYPE(pw_grid_type), POINTER :: pw_grid
552 : INTEGER, INTENT(IN) :: direction
553 : INTEGER, INTENT(INOUT) :: should_warn
554 :
555 : INTEGER, DIMENSION(3) :: lb, ub
556 : REAL(KIND=dp) :: max_val_low, max_val_up
557 :
558 144 : lb(:) = pw_grid%bounds_local(1, :)
559 144 : ub(:) = pw_grid%bounds_local(2, :)
560 144 : IF (.NOT. ALL(ub >= lb)) RETURN
561 :
562 36 : max_val_low = 0._dp
563 36 : max_val_up = 0._dp
564 54 : SELECT CASE (direction)
565 : CASE (1)
566 26748 : IF (lb(1) == pw_grid%bounds(1, 1)) max_val_low = MAXVAL(ABS(density%array(lb(1), :, :)))
567 26748 : IF (ub(1) == pw_grid%bounds(2, 1)) max_val_up = MAXVAL(ABS(density%array(ub(1), :, :)))
568 : CASE (2)
569 0 : IF (lb(2) == pw_grid%bounds(1, 2)) max_val_low = MAXVAL(ABS(density%array(:, lb(2), :)))
570 0 : IF (ub(2) == pw_grid%bounds(2, 2)) max_val_up = MAXVAL(ABS(density%array(:, ub(2), :)))
571 : CASE (3)
572 27234 : IF (lb(3) == pw_grid%bounds(1, 3)) max_val_low = MAXVAL(ABS(density%array(:, :, lb(3))))
573 27234 : IF (ub(3) == pw_grid%bounds(2, 3)) max_val_up = MAXVAL(ABS(density%array(:, :, ub(3))))
574 : CASE DEFAULT
575 36 : CPABORT("Invalid WAVELET isolated dimension")
576 : END SELECT
577 :
578 36 : IF (max_val_low >= 0.0001_dp) should_warn = 1
579 36 : IF (max_val_up >= 0.0001_dp) should_warn = 1
580 :
581 : END SUBROUTINE update_edge_warning
582 :
583 : ! **************************************************************************************************
584 : !> \brief Convert counts into zero-based displacements for all-to-all communication.
585 : !> \param counts ...
586 : !> \param displacements ...
587 : ! **************************************************************************************************
588 144 : SUBROUTINE set_displacements(counts, displacements)
589 :
590 : INTEGER, DIMENSION(:), INTENT(IN) :: counts
591 : INTEGER, DIMENSION(:), INTENT(OUT) :: displacements
592 :
593 : INTEGER :: i
594 :
595 144 : displacements(1) = 0
596 288 : DO i = 2, SIZE(counts)
597 288 : displacements(i) = displacements(i - 1) + counts(i - 1)
598 : END DO
599 :
600 144 : END SUBROUTINE set_displacements
601 :
602 : ! **************************************************************************************************
603 : !> \brief Return the distributed grid owner of a compact one-based grid index.
604 : !> \param index ...
605 : !> \param npts ...
606 : !> \param nparts ...
607 : !> \return ...
608 : ! **************************************************************************************************
609 11337408 : FUNCTION grid_owner(index, npts, nparts) RESULT(owner)
610 :
611 : INTEGER, INTENT(IN) :: index, npts, nparts
612 : INTEGER :: owner
613 :
614 : INTEGER :: ipart
615 : INTEGER, DIMENSION(2) :: limits
616 :
617 14171760 : DO ipart = 0, nparts - 1
618 14171760 : limits = get_limit(npts, nparts, ipart)
619 14171760 : IF (index >= limits(1) .AND. index <= limits(2)) THEN
620 11337408 : owner = ipart
621 : RETURN
622 : END IF
623 : END DO
624 0 : CPABORT("Grid index outside distributed bounds")
625 :
626 0 : END FUNCTION grid_owner
627 :
628 : ! **************************************************************************************************
629 : !> \brief Return the Wavelet z-slice owner of a compact one-based grid index.
630 : !> \param index ...
631 : !> \param local_z_dim ...
632 : !> \param nproc ...
633 : !> \return ...
634 : ! **************************************************************************************************
635 5668704 : FUNCTION z_slice_owner(index, local_z_dim, nproc) RESULT(owner)
636 :
637 : INTEGER, INTENT(IN) :: index, local_z_dim, nproc
638 : INTEGER :: owner
639 :
640 5668704 : owner = MIN((index - 1)/local_z_dim, nproc - 1)
641 :
642 5668704 : END FUNCTION z_slice_owner
643 :
644 : ! **************************************************************************************************
645 : !> \brief Return the CP2K real-space rank owning a real-space x/y grid point.
646 : !> \param ix ...
647 : !> \param iy ...
648 : !> \param pw_grid ...
649 : !> \return ...
650 : ! **************************************************************************************************
651 5668704 : FUNCTION cp2k_rank_owner(ix, iy, pw_grid) RESULT(rank)
652 :
653 : INTEGER, INTENT(IN) :: ix, iy
654 : TYPE(pw_grid_type), POINTER :: pw_grid
655 : INTEGER :: rank
656 :
657 : INTEGER, DIMENSION(2) :: cart_pos
658 :
659 : cart_pos(1) = grid_owner(ix - pw_grid%bounds(1, 1) + 1, pw_grid%npts(1), &
660 5668704 : pw_grid%para%group%num_pe_cart(1))
661 : cart_pos(2) = grid_owner(iy - pw_grid%bounds(1, 2) + 1, pw_grid%npts(2), &
662 5668704 : pw_grid%para%group%num_pe_cart(2))
663 5668704 : CALL pw_grid%para%group%rank_cart(cart_pos, rank)
664 :
665 5668704 : END FUNCTION cp2k_rank_owner
666 :
667 : ! **************************************************************************************************
668 : !> \brief Transfer a CP2K real-space grid into a permuted Wavelet z-sliced layout.
669 : !> \param density ...
670 : !> \param wavelet ...
671 : !> \param pw_grid ...
672 : ! **************************************************************************************************
673 36 : SUBROUTINE cp2k_distribution_to_z_slices_permuted(density, wavelet, pw_grid)
674 :
675 : TYPE(pw_r3d_rs_type), INTENT(IN) :: density
676 : TYPE(ps_wavelet_type), POINTER :: wavelet
677 : TYPE(pw_grid_type), POINTER :: pw_grid
678 :
679 : INTEGER :: dest, i, idir, ii, j, k, local_z_dim, &
680 : md2, nproc, nrecv, nsend
681 : INTEGER, ALLOCATABLE, DIMENSION(:) :: coord_rbuf, coord_rcount, coord_rdispl, coord_sbuf, &
682 36 : coord_scount, coord_sdispl, rcount, rdispl, scount, sdispl, send_pos
683 : INTEGER, DIMENSION(3) :: lb, q, u, ub
684 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:) :: rbuf, sbuf
685 :
686 108 : nproc = PRODUCT(pw_grid%para%group%num_pe_cart)
687 36 : md2 = wavelet%PS_grid(3)
688 36 : local_z_dim = MAX(md2/nproc, 1)
689 144 : lb(:) = pw_grid%bounds_local(1, :)
690 144 : ub(:) = pw_grid%bounds_local(2, :)
691 :
692 252 : ALLOCATE (scount(nproc), sdispl(nproc), rcount(nproc), rdispl(nproc), send_pos(nproc))
693 36 : scount = 0
694 1980 : DO k = lb(3), ub(3)
695 106956 : DO j = lb(2), ub(2)
696 2941272 : DO i = lb(1), ub(1)
697 11337408 : u = [i, j, k]
698 11337408 : DO idir = 1, 3
699 11337408 : q(idir) = u(wavelet%axis(idir)) - pw_grid%bounds(1, wavelet%axis(idir)) + 1
700 : END DO
701 2834352 : dest = z_slice_owner(q(3), local_z_dim, nproc)
702 2939328 : scount(dest + 1) = scount(dest + 1) + 1
703 : END DO
704 : END DO
705 : END DO
706 :
707 36 : CALL pw_grid%para%group%alltoall(scount, rcount, 1)
708 36 : CALL set_displacements(scount, sdispl)
709 36 : CALL set_displacements(rcount, rdispl)
710 108 : nsend = SUM(scount)
711 108 : nrecv = SUM(rcount)
712 :
713 180 : ALLOCATE (sbuf(MAX(nsend, 1)), rbuf(MAX(nrecv, 1)))
714 180 : ALLOCATE (coord_sbuf(MAX(3*nsend, 1)), coord_rbuf(MAX(3*nrecv, 1)))
715 180 : ALLOCATE (coord_scount(nproc), coord_sdispl(nproc), coord_rcount(nproc), coord_rdispl(nproc))
716 108 : coord_scount(:) = 3*scount(:)
717 108 : coord_rcount(:) = 3*rcount(:)
718 108 : coord_sdispl(:) = 3*sdispl(:)
719 108 : coord_rdispl(:) = 3*rdispl(:)
720 108 : send_pos(:) = sdispl(:) + 1
721 :
722 1980 : DO k = lb(3), ub(3)
723 106956 : DO j = lb(2), ub(2)
724 2941272 : DO i = lb(1), ub(1)
725 11337408 : u = [i, j, k]
726 11337408 : DO idir = 1, 3
727 11337408 : q(idir) = u(wavelet%axis(idir)) - pw_grid%bounds(1, wavelet%axis(idir)) + 1
728 : END DO
729 2834352 : dest = z_slice_owner(q(3), local_z_dim, nproc)
730 2834352 : ii = send_pos(dest + 1)
731 2834352 : sbuf(ii) = density%array(i, j, k)
732 2834352 : coord_sbuf(3*ii - 2) = q(1)
733 2834352 : coord_sbuf(3*ii - 1) = q(2)
734 2834352 : coord_sbuf(3*ii) = q(3)
735 2939328 : send_pos(dest + 1) = ii + 1
736 : END DO
737 : END DO
738 : END DO
739 :
740 36 : CALL pw_grid%para%group%alltoall(sbuf, scount, sdispl, rbuf, rcount, rdispl)
741 : CALL pw_grid%para%group%alltoall(coord_sbuf, coord_scount, coord_sdispl, &
742 36 : coord_rbuf, coord_rcount, coord_rdispl)
743 :
744 2887848 : wavelet%rho_z_sliced = 0.0_dp
745 2834388 : DO ii = 1, nrecv
746 2834352 : q(1) = coord_rbuf(3*ii - 2)
747 2834352 : q(2) = coord_rbuf(3*ii - 1)
748 2834352 : q(3) = coord_rbuf(3*ii)
749 2834388 : wavelet%rho_z_sliced(q(1), q(2), q(3) - pw_grid%para%group%mepos*local_z_dim) = rbuf(ii)
750 : END DO
751 :
752 0 : DEALLOCATE (sbuf, rbuf, coord_sbuf, coord_rbuf, coord_scount, coord_sdispl, coord_rcount, &
753 36 : coord_rdispl, scount, sdispl, rcount, rdispl, send_pos)
754 :
755 36 : END SUBROUTINE cp2k_distribution_to_z_slices_permuted
756 :
757 : ! **************************************************************************************************
758 : !> \brief Transfer a permuted Wavelet z-sliced layout back to a CP2K real-space grid.
759 : !> \param density ...
760 : !> \param wavelet ...
761 : !> \param pw_grid ...
762 : ! **************************************************************************************************
763 36 : SUBROUTINE z_slices_to_cp2k_distribution_permuted(density, wavelet, pw_grid)
764 :
765 : TYPE(pw_r3d_rs_type), INTENT(IN) :: density
766 : TYPE(ps_wavelet_type), POINTER :: wavelet
767 : TYPE(pw_grid_type), POINTER :: pw_grid
768 :
769 : INTEGER :: dest, i, idir, ii, j, k, local_z, &
770 : local_z_dim, md2, n1, n2, n3, nproc, &
771 : nrecv, nsend, z_end, z_start
772 36 : INTEGER, ALLOCATABLE, DIMENSION(:) :: coord_rbuf, coord_rcount, coord_rdispl, coord_sbuf, &
773 36 : coord_scount, coord_sdispl, rcount, rdispl, scount, sdispl, send_pos
774 : INTEGER, DIMENSION(3) :: lb, q, u, ub
775 36 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:) :: rbuf, sbuf
776 :
777 108 : nproc = PRODUCT(pw_grid%para%group%num_pe_cart)
778 36 : md2 = wavelet%PS_grid(3)
779 36 : local_z_dim = MAX(md2/nproc, 1)
780 36 : n1 = pw_grid%npts(wavelet%axis(1))
781 36 : n2 = pw_grid%npts(wavelet%axis(2))
782 36 : n3 = pw_grid%npts(wavelet%axis(3))
783 36 : z_start = pw_grid%para%group%mepos*local_z_dim + 1
784 36 : z_end = MIN((pw_grid%para%group%mepos + 1)*local_z_dim, n3)
785 144 : lb(:) = pw_grid%bounds_local(1, :)
786 144 : ub(:) = pw_grid%bounds_local(2, :)
787 :
788 252 : ALLOCATE (scount(nproc), sdispl(nproc), rcount(nproc), rdispl(nproc), send_pos(nproc))
789 36 : scount = 0
790 1008 : DO k = z_start, z_end
791 53496 : DO j = 1, n2
792 2887812 : DO i = 1, n1
793 11337408 : q = [i, j, k]
794 11337408 : DO idir = 1, 3
795 11337408 : u(wavelet%axis(idir)) = q(idir) + pw_grid%bounds(1, wavelet%axis(idir)) - 1
796 : END DO
797 2834352 : dest = cp2k_rank_owner(u(1), u(2), pw_grid)
798 2886840 : scount(dest + 1) = scount(dest + 1) + 1
799 : END DO
800 : END DO
801 : END DO
802 :
803 36 : CALL pw_grid%para%group%alltoall(scount, rcount, 1)
804 36 : CALL set_displacements(scount, sdispl)
805 36 : CALL set_displacements(rcount, rdispl)
806 108 : nsend = SUM(scount)
807 108 : nrecv = SUM(rcount)
808 :
809 180 : ALLOCATE (sbuf(MAX(nsend, 1)), rbuf(MAX(nrecv, 1)))
810 180 : ALLOCATE (coord_sbuf(MAX(3*nsend, 1)), coord_rbuf(MAX(3*nrecv, 1)))
811 180 : ALLOCATE (coord_scount(nproc), coord_sdispl(nproc), coord_rcount(nproc), coord_rdispl(nproc))
812 108 : coord_scount(:) = 3*scount(:)
813 108 : coord_rcount(:) = 3*rcount(:)
814 108 : coord_sdispl(:) = 3*sdispl(:)
815 108 : coord_rdispl(:) = 3*rdispl(:)
816 108 : send_pos(:) = sdispl(:) + 1
817 :
818 1008 : DO k = z_start, z_end
819 53496 : DO j = 1, n2
820 2887812 : DO i = 1, n1
821 11337408 : q = [i, j, k]
822 11337408 : DO idir = 1, 3
823 11337408 : u(wavelet%axis(idir)) = q(idir) + pw_grid%bounds(1, wavelet%axis(idir)) - 1
824 : END DO
825 2834352 : dest = cp2k_rank_owner(u(1), u(2), pw_grid)
826 2834352 : ii = send_pos(dest + 1)
827 2834352 : local_z = k - pw_grid%para%group%mepos*local_z_dim
828 2834352 : sbuf(ii) = wavelet%rho_z_sliced(i, j, local_z)
829 2834352 : coord_sbuf(3*ii - 2) = u(1)
830 2834352 : coord_sbuf(3*ii - 1) = u(2)
831 2834352 : coord_sbuf(3*ii) = u(3)
832 2886840 : send_pos(dest + 1) = ii + 1
833 : END DO
834 : END DO
835 : END DO
836 :
837 36 : CALL pw_grid%para%group%alltoall(sbuf, scount, sdispl, rbuf, rcount, rdispl)
838 : CALL pw_grid%para%group%alltoall(coord_sbuf, coord_scount, coord_sdispl, &
839 36 : coord_rbuf, coord_rcount, coord_rdispl)
840 :
841 2834388 : DO ii = 1, nrecv
842 2834352 : u(1) = coord_rbuf(3*ii - 2)
843 2834352 : u(2) = coord_rbuf(3*ii - 1)
844 2834352 : u(3) = coord_rbuf(3*ii)
845 : IF (u(1) >= lb(1) .AND. u(1) <= ub(1) .AND. u(2) >= lb(2) .AND. u(2) <= ub(2) .AND. &
846 2834388 : u(3) >= lb(3) .AND. u(3) <= ub(3)) THEN
847 2834352 : density%array(u(1), u(2), u(3)) = rbuf(ii)
848 : END IF
849 : END DO
850 :
851 0 : DEALLOCATE (sbuf, rbuf, coord_sbuf, coord_rbuf, coord_scount, coord_sdispl, coord_rcount, &
852 36 : coord_rdispl, scount, sdispl, rcount, rdispl, send_pos)
853 :
854 36 : END SUBROUTINE z_slices_to_cp2k_distribution_permuted
855 :
856 : ! **************************************************************************************************
857 : !> \brief ...
858 : !> \param wavelet ...
859 : !> \param pw_grid ...
860 : ! **************************************************************************************************
861 66590 : SUBROUTINE ps_wavelet_solve(wavelet, pw_grid)
862 :
863 : TYPE(ps_wavelet_type), POINTER :: wavelet
864 : TYPE(pw_grid_type), POINTER :: pw_grid
865 :
866 : CHARACTER(len=*), PARAMETER :: routineN = 'ps_wavelet_solve'
867 :
868 : CHARACTER(LEN=1) :: geocode
869 : INTEGER :: handle, iproc, nproc, nx, ny, nz
870 : REAL(KIND=dp) :: hx, hy, hz
871 :
872 33295 : CALL timeset(routineN, handle)
873 99885 : nproc = PRODUCT(pw_grid%para%group%num_pe_cart)
874 33295 : iproc = pw_grid%para%group%mepos
875 33295 : geocode = wavelet%geocode
876 33295 : CALL get_wavelet_grid(wavelet, pw_grid, nx, ny, nz, hx, hy, hz)
877 :
878 : CALL PSolver(geocode, iproc, nproc, nx, ny, nz, hx, hy, hz, &
879 33295 : wavelet%rho_z_sliced, wavelet%karray, pw_grid)
880 33295 : CALL timestop(handle)
881 33295 : END SUBROUTINE ps_wavelet_solve
882 :
883 : END MODULE ps_wavelet_methods
|
