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