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 2 : PROGRAM kpoint_lattice_fft_unittest
8 2 : USE fft_tools, ONLY: finalize_fft,&
9 : init_fft
10 : USE kinds, ONLY: dp
11 : USE kpoint_lattice_fft, ONLY: cell_to_k_grid_fft,&
12 : k_grid_to_cell_execute,&
13 : k_grid_to_cell_fft,&
14 : k_grid_to_cell_prepare,&
15 : k_grid_to_cell_release,&
16 : k_grid_to_cell_work_type,&
17 : regular_kpoint_grid
18 : USE lattice_preconditioner_operator, ONLY: &
19 : apply_lattice_inverse_dense, apply_lattice_state_inverse_dense, build_lattice_inverse, &
20 : build_lattice_local_correction, lattice_fft_auto_select, lattice_fft_cost_ratio, &
21 : lattice_fft_storage_ratio, lattice_grid, lattice_local_cost_ratio, &
22 : lattice_local_storage_ratio
23 : USE mathconstants, ONLY: gaussi,&
24 : twopi,&
25 : z_zero
26 : USE message_passing, ONLY: mp_comm_type,&
27 : mp_world_finalize,&
28 : mp_world_init
29 :
30 : IMPLICIT NONE
31 :
32 : INTEGER, PARAMETER :: base_block = 2, base_cell = 8, base_n = 16, &
33 : base_ncol = 3, ncell = 31, nkp = 64, &
34 : nround = 16, nselected = 3
35 : INTEGER, PARAMETER :: large_block = 1000, large_ncol = 1000, &
36 : large_num_pe = 16
37 2 : COMPLEX(KIND=dp), ALLOCATABLE, DIMENSION(:, :, :) :: inverse_k
38 2 : COMPLEX(KIND=dp), ALLOCATABLE, DIMENSION(:, :, :, :) :: state_inverse_k
39 : COMPLEX(KIND=dp), DIMENSION(2, 3, nkp) :: direct, transformed
40 : COMPLEX(KIND=dp), DIMENSION(2, 3, nround) :: reconstructed, round_k
41 : COMPLEX(KIND=dp), DIMENSION(2, 3, nselected) :: selected_transformed
42 : INTEGER :: d, i, icell, info, j, jcell, ik, ix, iy, iz
43 2 : INTEGER, ALLOCATABLE, DIMENSION(:) :: correction_indices
44 : INTEGER, DIMENSION(3) :: base_dims, delta, large_dims
45 : INTEGER, DIMENSION(3, base_cell) :: base_index
46 : INTEGER, DIMENSION(3) :: nkp_grid
47 : INTEGER, DIMENSION(3, ncell) :: index_to_cell
48 : INTEGER, DIMENSION(3, nround) :: round_index
49 : INTEGER, DIMENSION(nselected) :: selected_kpoints
50 : LOGICAL :: used_fft
51 : REAL(KIND=dp), DIMENSION(3, 3) :: hmat
52 : REAL(KIND=dp) :: base_residual_norm, corrected_projection_error, &
53 : corrected_residual_norm, cost_ratio, &
54 : expected_cost_ratio, imaginary_residual, &
55 : large_block_real, large_ncell_real, &
56 : large_ncol_real, projection_error
57 : REAL(KIND=dp), DIMENSION(3, base_cell) :: base_xkp
58 : REAL(KIND=dp), DIMENSION(base_n, base_n) :: base_overlap
59 : REAL(KIND=dp), DIMENSION(base_n, base_n) :: defect_operator
60 : REAL(KIND=dp), DIMENSION(base_n, base_ncol) :: base_input, base_input_rotated, &
61 : base_output, base_output_rotated, &
62 : base_state_output, defect_residual
63 2 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :) :: coarse_inverse, defect_base_output, &
64 2 : defect_corrected_output, local_operator, &
65 2 : selected_output, selected_rhs
66 : REAL(KIND=dp), DIMENSION(base_ncol, base_ncol) :: column_rotation
67 : REAL(KIND=dp), DIMENSION(3, nkp) :: xkp, xkp_bad
68 : REAL(KIND=dp), DIMENSION(2, 3, ncell) :: values_rs
69 : REAL(KIND=dp), DIMENSION(2, 3, nround) :: round_rs
70 : TYPE(mp_comm_type) :: mp_comm
71 8 : TYPE(k_grid_to_cell_work_type) :: inverse_work
72 : REAL(KIND=dp), DIMENSION(9) :: subset_weights
73 : COMPLEX(KIND=dp), DIMENSION(2, 3, ncell) :: weighted_reference, weighted_result
74 :
75 2 : CALL mp_world_init(mp_comm)
76 2 : CALL check_image_reconstruction(mp_comm)
77 : CALL init_fft("FFTSG", alltoall=.FALSE., pool_limit=4, &
78 2 : wisdom_file="", plan_style=1)
79 :
80 2 : nkp_grid = [4, 4, 4]
81 : hmat = RESHAPE([4.0_dp, 0.1_dp, 0.2_dp, &
82 : 0.0_dp, 5.0_dp, 0.3_dp, &
83 2 : 0.0_dp, 0.0_dp, 6.0_dp], [3, 3])
84 :
85 2 : ik = 0
86 10 : DO iz = 0, nkp_grid(3) - 1
87 42 : DO iy = 0, nkp_grid(2) - 1
88 168 : DO ix = 0, nkp_grid(1) - 1
89 128 : ik = ik + 1
90 : xkp(:, ik) = [(REAL(ix, dp) + 0.5_dp)/REAL(nkp_grid(1), dp), &
91 : REAL(iy, dp)/REAL(nkp_grid(2), dp), &
92 544 : (REAL(iz, dp) + 0.25_dp)/REAL(nkp_grid(3), dp)]
93 : END DO
94 : END DO
95 : END DO
96 2 : IF (.NOT. regular_kpoint_grid(xkp, nkp_grid)) ERROR STOP "Regular k-point grid not recognized"
97 :
98 64 : DO icell = 1, ncell
99 : index_to_cell(:, icell) = [MODULO(3*icell + 1, 11) - 5, &
100 : MODULO(5*icell + 2, 9) - 4, &
101 248 : MODULO(7*icell + 3, 7) - 3]
102 250 : DO d = 1, 3
103 186 : values_rs(1, d, icell) = SIN(0.13_dp*REAL(icell + 2*d, dp))
104 248 : values_rs(2, d, icell) = COS(0.17_dp*REAL(2*icell - d, dp))
105 : END DO
106 : END DO
107 :
108 2 : CALL cell_to_k_grid_fft(values_rs, index_to_cell, xkp, nkp_grid, transformed, used_fft)
109 2 : IF (.NOT. used_fft) ERROR STOP "Complete regular grid did not use lattice FFT"
110 2 : CALL direct_reference(values_rs, index_to_cell, xkp, direct)
111 1282 : IF (MAXVAL(ABS(transformed - direct)) > 2.0E-12_dp) THEN
112 0 : ERROR STOP "Shifted lattice FFT differs from direct transform"
113 : END IF
114 :
115 2 : selected_kpoints = [2, 31, 64]
116 : CALL cell_to_k_grid_fft(values_rs, index_to_cell, xkp, nkp_grid, selected_transformed, &
117 2 : used_fft, selected_kpoints=selected_kpoints)
118 2 : IF (.NOT. used_fft) ERROR STOP "Selected distributed output did not use lattice FFT"
119 62 : IF (MAXVAL(ABS(selected_transformed - direct(:, :, selected_kpoints))) > 2.0E-12_dp) THEN
120 0 : ERROR STOP "Selected lattice-FFT output differs from the full transform"
121 : END IF
122 :
123 2 : CALL cell_to_k_grid_fft(values_rs, index_to_cell, xkp, nkp_grid, transformed, used_fft, 2, hmat)
124 2 : IF (.NOT. used_fft) ERROR STOP "Derivative did not use lattice FFT"
125 2 : CALL direct_reference(values_rs, index_to_cell, xkp, direct, 2, hmat)
126 1282 : IF (MAXVAL(ABS(transformed - direct)) > 5.0E-12_dp) THEN
127 0 : ERROR STOP "Lattice-FFT derivative differs from direct transform"
128 : END IF
129 :
130 2 : nkp_grid = [4, 4, 1]
131 2 : ik = 0
132 10 : DO iy = 0, nkp_grid(2) - 1
133 42 : DO ix = 0, nkp_grid(1) - 1
134 32 : ik = ik + 1
135 : xkp(:, ik) = [(REAL(ix, dp) + 0.5_dp)/REAL(nkp_grid(1), dp), &
136 136 : REAL(iy, dp)/REAL(nkp_grid(2), dp), 0.0_dp]
137 : END DO
138 : END DO
139 : CALL cell_to_k_grid_fft(values_rs, index_to_cell, xkp(:, 1:ik), nkp_grid, &
140 2 : transformed(:, :, 1:ik), used_fft)
141 2 : IF (.NOT. used_fft) ERROR STOP "Singleton grid direction did not use padded lattice FFT"
142 2 : CALL direct_reference(values_rs, index_to_cell, xkp(:, 1:ik), direct(:, :, 1:ik))
143 322 : IF (MAXVAL(ABS(transformed(:, :, 1:ik) - direct(:, :, 1:ik))) > 2.0E-12_dp) THEN
144 0 : ERROR STOP "Padded lattice FFT differs from direct transform"
145 : END IF
146 :
147 : icell = 0
148 10 : DO iy = 0, nkp_grid(2) - 1
149 42 : DO ix = 0, nkp_grid(1) - 1
150 32 : icell = icell + 1
151 128 : round_index(:, icell) = [ix, iy, 0]
152 328 : round_rs(:, :, icell) = values_rs(:, :, icell)
153 : END DO
154 : END DO
155 2 : CALL cell_to_k_grid_fft(round_rs, round_index, xkp(:, 1:ik), nkp_grid, round_k, used_fft)
156 2 : IF (.NOT. used_fft) ERROR STOP "Round-trip forward transform did not use lattice FFT"
157 2 : CALL k_grid_to_cell_fft(round_k, xkp(:, 1:ik), nkp_grid, round_index, reconstructed, used_fft)
158 2 : IF (.NOT. used_fft) ERROR STOP "Round-trip inverse transform did not use lattice FFT"
159 322 : IF (MAXVAL(ABS(reconstructed - CMPLX(round_rs, KIND=dp))) > 3.0E-12_dp) THEN
160 0 : ERROR STOP "Forward/inverse lattice-FFT round trip is inconsistent"
161 : END IF
162 :
163 : ! An incomplete shifted grid uses the supplied weights exactly, including
164 : ! negative translations. Reuse the preparation for separate matrix batches.
165 2 : subset_weights = [(REAL(ik, dp)/45.0_dp, ik=1, 9)]
166 2 : weighted_reference = z_zero
167 64 : DO icell = 1, ncell
168 622 : DO ik = 1, 9
169 : weighted_reference(:, :, icell) = weighted_reference(:, :, icell) + subset_weights(ik)* &
170 7316 : EXP(-gaussi*twopi*SUM(xkp(:, ik)*REAL(index_to_cell(:, icell), dp)))*round_k(:, :, ik)
171 : END DO
172 : END DO
173 : CALL k_grid_to_cell_prepare(inverse_work, xkp(:, 1:9), nkp_grid, index_to_cell, &
174 2 : weights=subset_weights, allow_incomplete=.TRUE.)
175 6 : DO i = 1, 2
176 4 : CALL k_grid_to_cell_execute(inverse_work, round_k(i:i, :, 1:9), weighted_result(i:i, :, :), used_fft)
177 6 : IF (.NOT. used_fft) ERROR STOP "Prepared incomplete grid did not use FFT"
178 : END DO
179 622 : IF (MAXVAL(ABS(weighted_result - weighted_reference)) > 3.0E-12_dp) THEN
180 0 : ERROR STOP "Prepared weighted inverse differs from explicit phase sum"
181 : END IF
182 : CALL k_grid_to_cell_fft(round_k(:, :, 1:9), xkp(:, 1:9), nkp_grid, index_to_cell, weighted_result, &
183 2 : used_fft, weights=subset_weights, allow_incomplete=.TRUE.)
184 622 : IF (MAXVAL(ABS(weighted_result - weighted_reference)) > 3.0E-12_dp) THEN
185 0 : ERROR STOP "One-shot and prepared inverse transforms disagree"
186 : END IF
187 : CALL k_grid_to_cell_prepare(inverse_work, xkp(:, 1:9), nkp_grid, index_to_cell, &
188 20 : weights=2.0_dp*subset_weights, allow_incomplete=.TRUE.)
189 2 : CALL k_grid_to_cell_execute(inverse_work, round_k(:, :, 1:9), weighted_result, used_fft)
190 622 : IF (MAXVAL(ABS(weighted_result - 2.0_dp*weighted_reference)) > 6.0E-12_dp) THEN
191 0 : ERROR STOP "Repreparing inverse transform retained stale weights"
192 : END IF
193 2 : CALL k_grid_to_cell_release(inverse_work)
194 2 : CALL k_grid_to_cell_release(inverse_work)
195 :
196 2 : nkp_grid = [4, 4, 4]
197 :
198 2 : xkp_bad = xkp
199 2 : xkp_bad(1, nkp) = xkp_bad(1, nkp) + 0.03125_dp
200 2 : CALL cell_to_k_grid_fft(values_rs, index_to_cell, xkp_bad, nkp_grid, transformed, used_fft)
201 2 : IF (used_fft) ERROR STOP "Irregular k-point list incorrectly used lattice FFT"
202 2 : CALL direct_reference(values_rs, index_to_cell, xkp_bad, direct)
203 1282 : IF (MAXVAL(ABS(transformed - direct)) > 1.0E-14_dp) THEN
204 0 : ERROR STOP "Direct lattice-transform fallback is inconsistent"
205 : END IF
206 :
207 : ! Verify the block-circulant inverse against direct matrix application and right covariance.
208 2 : base_dims = [4, 2, 1]
209 2 : CALL lattice_grid(base_dims, base_index, base_xkp)
210 2 : base_overlap = 0.0_dp
211 18 : DO jcell = 1, base_cell
212 146 : DO icell = 1, base_cell
213 512 : delta = MODULO(base_index(:, icell) - base_index(:, jcell), base_dims)
214 384 : DO i = 1, base_block
215 512 : IF (ALL(delta == 0)) THEN
216 : base_overlap((icell - 1)*base_block + i, (jcell - 1)*base_block + i) = &
217 32 : 2.8_dp + 0.2_dp*REAL(i, KIND=dp)
218 224 : ELSE IF ((delta(1) == 1 .OR. delta(1) == base_dims(1) - 1) .AND. &
219 : delta(2) == 0) THEN
220 64 : base_overlap((icell - 1)*base_block + i, (jcell - 1)*base_block + i) = -0.12_dp
221 160 : ELSE IF (delta(1) == 0 .AND. delta(2) == 1) THEN
222 32 : base_overlap((icell - 1)*base_block + i, (jcell - 1)*base_block + i) = -0.08_dp
223 : END IF
224 : END DO
225 208 : IF (ALL(delta == 0)) THEN
226 16 : base_overlap((icell - 1)*base_block + 1, (jcell - 1)*base_block + 2) = 0.15_dp
227 16 : base_overlap((icell - 1)*base_block + 2, (jcell - 1)*base_block + 1) = 0.15_dp
228 : END IF
229 : END DO
230 : END DO
231 8 : DO j = 1, base_ncol
232 104 : DO i = 1, base_n
233 102 : base_input(i, j) = SIN(0.11_dp*REAL(2*i + 3*j, KIND=dp))
234 : END DO
235 : END DO
236 2 : CALL build_lattice_inverse(base_overlap, base_dims, inverse_k, projection_error, used_fft, info)
237 2 : IF (info /= 0) ERROR STOP "Positive lattice overlap failed block Cholesky inversion"
238 2 : IF (.NOT. used_fft) ERROR STOP "Block-circulant overlap construction did not use lattice FFT"
239 2 : IF (projection_error > 2.0E-15_dp) ERROR STOP "Exact block-circulant overlap was changed by projection"
240 2 : IF (lattice_fft_auto_select(base_dims, base_block, base_ncol, 1, 1, &
241 : projection_error, used_fft, info)) THEN
242 0 : ERROR STOP "Cost-aware lattice-FFT AUTO accepted a tiny eight-cell operator"
243 : END IF
244 2 : IF (.NOT. lattice_fft_auto_select(base_dims, 6, 8, 1, 2, &
245 : projection_error, used_fft, info)) THEN
246 0 : ERROR STOP "Cost-aware lattice-FFT AUTO rejected an economical eight-cell operator"
247 : END IF
248 2 : IF (lattice_fft_auto_select([4, 1, 1], 6, 8, 1, 2, &
249 : projection_error, used_fft, info)) THEN
250 0 : ERROR STOP "Conservative lattice-FFT AUTO accepted fewer than eight cells"
251 : END IF
252 2 : IF (lattice_fft_auto_select(base_dims, 6, 8, 1, 2, 2.0E-3_dp, used_fft, info)) THEN
253 0 : ERROR STOP "Conservative lattice-FFT AUTO accepted a large projection error"
254 : END IF
255 2 : IF (lattice_fft_auto_select(base_dims, 6, 8, 17, 2, &
256 : projection_error, used_fft, info)) THEN
257 0 : ERROR STOP "Storage-aware lattice-FFT AUTO accepted too many state operators"
258 : END IF
259 2 : IF (lattice_fft_cost_ratio(base_dims, base_block, base_ncol, 1) <= 0.8_dp) THEN
260 0 : ERROR STOP "Lattice-FFT cost estimate missed the tiny-batch overhead"
261 : END IF
262 2 : large_dims = [100, 100, 1]
263 8 : large_ncell_real = REAL(PRODUCT(large_dims), KIND=dp)
264 2 : large_block_real = REAL(large_block, KIND=dp)
265 2 : large_ncol_real = REAL(large_ncol, KIND=dp)
266 : expected_cost_ratio = ( &
267 : large_ncell_real*large_block_real**2*large_ncol_real + &
268 : 4.0_dp*large_ncell_real*large_block_real*large_ncol_real* &
269 : LOG(large_ncell_real)/LOG(2.0_dp) + &
270 : 8.0_dp*large_ncell_real*large_block_real + &
271 : 2.0_dp*large_ncell_real*large_block_real*large_ncol_real* &
272 : LOG(REAL(large_num_pe + 1, KIND=dp))/LOG(2.0_dp))/ &
273 2 : ((large_ncell_real*large_block_real)**2*large_ncol_real)
274 2 : cost_ratio = lattice_fft_cost_ratio(large_dims, large_block, large_ncol, large_num_pe)
275 2 : IF (ABS(cost_ratio - expected_cost_ratio) > 1.0E-14_dp*expected_cost_ratio) THEN
276 0 : ERROR STOP "Lattice-FFT cost estimate overflowed for a large problem"
277 : END IF
278 2 : IF (ABS(lattice_fft_storage_ratio(base_dims, 1) - 0.25_dp) > 1.0E-14_dp) THEN
279 0 : ERROR STOP "Lattice-FFT storage estimate is inconsistent"
280 : END IF
281 2 : IF (ABS(lattice_local_cost_ratio(base_dims, 1) - 0.28125_dp) > 1.0E-14_dp) THEN
282 0 : ERROR STOP "Lattice local-correction cost estimate is inconsistent"
283 : END IF
284 2 : IF (ABS(lattice_local_storage_ratio(base_dims, 3, 1) - 0.421875_dp) > 1.0E-14_dp) THEN
285 0 : ERROR STOP "Lattice local-correction storage estimate is inconsistent"
286 : END IF
287 : CALL apply_lattice_inverse_dense(inverse_k, base_dims, base_input, base_output, &
288 2 : used_fft, imaginary_residual)
289 2 : IF (.NOT. used_fft) ERROR STOP "Block-circulant inverse application did not use lattice FFT"
290 2 : IF (imaginary_residual > 3.0E-13_dp) ERROR STOP "Lattice inverse did not reconstruct a real matrix"
291 1742 : IF (MAXVAL(ABS(MATMUL(base_overlap, base_output) - base_input)) > 2.0E-11_dp) THEN
292 0 : ERROR STOP "FFT block-circulant inverse differs from the dense operator inverse"
293 : END IF
294 :
295 2 : ALLOCATE (state_inverse_k(base_block, base_block, base_cell, base_ncol))
296 8 : DO j = 1, base_ncol
297 344 : state_inverse_k(:, :, :, j) = (0.5_dp + 0.25_dp*REAL(j, KIND=dp))*inverse_k
298 : END DO
299 : CALL apply_lattice_state_inverse_dense(state_inverse_k, base_dims, base_input, &
300 2 : base_state_output, used_fft, imaginary_residual)
301 2 : IF (.NOT. used_fft) ERROR STOP "State-dependent inverse application did not use lattice FFT"
302 8 : DO j = 1, base_ncol
303 102 : IF (MAXVAL(ABS(base_state_output(:, j) - &
304 2 : (0.5_dp + 0.25_dp*REAL(j, KIND=dp))*base_output(:, j))) > 3.0E-12_dp) THEN
305 0 : ERROR STOP "State-dependent lattice inverse used the wrong column operator"
306 : END IF
307 : END DO
308 2 : DEALLOCATE (state_inverse_k)
309 :
310 2 : column_rotation = 0.0_dp
311 2 : column_rotation(1, 1) = COS(0.31_dp)
312 2 : column_rotation(1, 2) = -SIN(0.31_dp)
313 2 : column_rotation(2, 1) = SIN(0.31_dp)
314 2 : column_rotation(2, 2) = COS(0.31_dp)
315 2 : column_rotation(3, 3) = 1.0_dp
316 314 : base_input_rotated = MATMUL(base_input, column_rotation)
317 : CALL apply_lattice_inverse_dense(inverse_k, base_dims, base_input_rotated, &
318 2 : base_output_rotated, used_fft)
319 416 : IF (MAXVAL(ABS(base_output_rotated - MATMUL(base_output, column_rotation))) > 3.0E-12_dp) THEN
320 0 : ERROR STOP "FFT lattice base is not right-covariant under orbital rotations"
321 : END IF
322 :
323 : ! Add one localized positive defect and verify the balanced cell correction.
324 2 : defect_operator = base_overlap
325 2 : defect_operator(5, 5) = defect_operator(5, 5) + 0.90_dp
326 2 : defect_operator(6, 6) = defect_operator(6, 6) + 0.60_dp
327 2 : CALL build_lattice_inverse(defect_operator, base_dims, inverse_k, projection_error, used_fft, info)
328 2 : IF (info /= 0 .OR. .NOT. used_fft) ERROR STOP "Defective lattice operator failed construction"
329 2 : IF (projection_error <= 1.0E-3_dp) ERROR STOP "Localized defect did not break lattice translations"
330 2 : IF (lattice_fft_auto_select(base_dims, 6, 8, 1, 2, projection_error, used_fft, info)) THEN
331 0 : ERROR STOP "Lattice-FFT AUTO accepted a localized defect"
332 : END IF
333 : CALL build_lattice_local_correction( &
334 : defect_operator, base_dims, 1, correction_indices, local_operator, coarse_inverse, &
335 2 : corrected_projection_error, info)
336 2 : IF (info /= 0) ERROR STOP "Local correction coarse inverse failed"
337 6 : IF (ANY(correction_indices /= [5, 6])) ERROR STOP "Local correction selected the wrong defect cell"
338 2 : IF (corrected_projection_error >= projection_error) THEN
339 0 : ERROR STOP "Local correction did not reduce the non-circulant residual"
340 : END IF
341 :
342 2 : ALLOCATE (defect_base_output(base_n, base_ncol), defect_corrected_output(base_n, base_ncol))
343 2 : CALL apply_lattice_inverse_dense(inverse_k, base_dims, base_input, defect_base_output, used_fft)
344 : CALL apply_balanced_local_inverse(inverse_k, base_dims, correction_indices, local_operator, &
345 2 : coarse_inverse, base_input, defect_corrected_output)
346 3280 : defect_residual = MATMUL(defect_operator, defect_base_output) - base_input
347 104 : base_residual_norm = NORM2(defect_residual)
348 1742 : defect_residual = MATMUL(defect_operator, defect_corrected_output) - base_input
349 104 : corrected_residual_norm = NORM2(defect_residual)
350 2 : IF (corrected_residual_norm >= base_residual_norm) THEN
351 0 : ERROR STOP "Balanced local correction did not improve inverse application"
352 : END IF
353 104 : IF (SUM(base_input*defect_corrected_output) <= 0.0_dp) THEN
354 0 : ERROR STOP "Balanced local correction lost positive definiteness"
355 : END IF
356 :
357 0 : ALLOCATE (selected_rhs(base_n, SIZE(correction_indices)), &
358 8 : selected_output(base_n, SIZE(correction_indices)))
359 70 : selected_rhs(:, :) = local_operator
360 : CALL apply_balanced_local_inverse(inverse_k, base_dims, correction_indices, local_operator, &
361 2 : coarse_inverse, selected_rhs, selected_output)
362 2 : selected_rhs = 0.0_dp
363 6 : DO i = 1, SIZE(correction_indices)
364 6 : selected_rhs(correction_indices(i), i) = 1.0_dp
365 : END DO
366 70 : IF (MAXVAL(ABS(selected_output - selected_rhs)) > 3.0E-11_dp) THEN
367 0 : ERROR STOP "Balanced local correction is not exact in its selected subspace"
368 : END IF
369 0 : DEALLOCATE (coarse_inverse, correction_indices, defect_base_output, defect_corrected_output, &
370 2 : local_operator, selected_output, selected_rhs)
371 2 : DEALLOCATE (inverse_k)
372 :
373 2 : CALL finalize_fft(mp_comm, "")
374 2 : CALL mp_world_finalize()
375 :
376 : CONTAINS
377 :
378 : ! **************************************************************************************************
379 : !> \brief Check triangular image storage against rskp_transform for both neighbor orientations.
380 : !> \param comm MPI communicator
381 : ! **************************************************************************************************
382 2 : SUBROUTINE check_image_reconstruction(comm)
383 : USE cp_dbcsr_api, ONLY: &
384 : dbcsr_create, dbcsr_distribution_new, dbcsr_distribution_release, dbcsr_distribution_type, &
385 : dbcsr_finalize, dbcsr_finalize_lib, dbcsr_get_block_p, dbcsr_init_lib, dbcsr_p_type, &
386 2 : dbcsr_put_block, dbcsr_release, dbcsr_type, dbcsr_type_antisymmetric, dbcsr_type_symmetric
387 : USE cp_log_handling, ONLY: cp_add_default_logger, cp_logger_create, cp_logger_release, &
388 : cp_logger_type, cp_rm_default_logger
389 : USE kpoint_k_r_trafo_simple, ONLY: replicate_rs_matrices, rs_to_kp
390 : USE kpoint_methods, ONLY: rskp_transform
391 : USE kpoint_types, ONLY: kpoint_type
392 : USE machine, ONLY: default_output_unit
393 : USE message_passing, ONLY: mp_para_env_release, mp_para_env_type
394 : USE qs_neighbor_list_types, ONLY: add_neighbor_list, add_neighbor_node, &
395 : allocate_neighbor_list_set, deallocate_neighbor_list_set, neighbor_list_p_type
396 : TYPE(mp_comm_type), INTENT(IN) :: comm
397 :
398 : CHARACTER :: matrix_type, mirror_type
399 : COMPLEX(dp) :: expected(5, 5), phase, RESULT(5, 5)
400 : INTEGER :: atom, ik, image(3), img, io, mode, &
401 : neighbor, orientation, owner
402 2 : INTEGER, POINTER :: col_dist(:), images(:, :), &
403 2 : process_grid(:, :), row_dist(:), &
404 2 : sizes(:)
405 : LOGICAL :: found
406 : REAL(dp) :: expanded(1, 5, 5, 3), point(3), sign, &
407 : values(2, 3)
408 2 : REAL(dp), POINTER :: cosine(:, :), sine(:, :)
409 : TYPE(cp_logger_type), POINTER :: logger
410 : TYPE(dbcsr_distribution_type) :: distribution
411 2 : TYPE(dbcsr_p_type), POINTER :: source(:, :)
412 : TYPE(dbcsr_type) :: cosine_matrix, sine_matrix
413 : TYPE(kpoint_type), POINTER :: points
414 : TYPE(mp_para_env_type), POINTER :: para_env
415 : TYPE(neighbor_list_p_type) :: neighbors
416 :
417 0 : ALLOCATE (para_env)
418 2 : CALL para_env%from_dup(comm)
419 2 : io = -1
420 2 : IF (comm%mepos == 0) io = default_output_unit
421 2 : NULLIFY (logger)
422 : CALL cp_logger_create(logger, para_env=para_env, default_global_unit_nr=io, &
423 2 : close_global_unit_on_dealloc=.FALSE.)
424 2 : CALL cp_add_default_logger(logger)
425 2 : CALL dbcsr_init_lib(comm%get_handle(), io)
426 6 : ALLOCATE (process_grid(0:comm%num_pe - 1, 0:0), row_dist(2), col_dist(2), sizes(2))
427 10 : process_grid(:, 0) = [(img, img=0, comm%num_pe - 1)]
428 6 : col_dist = 0
429 6 : sizes = [2, 3]
430 30 : ALLOCATE (points, source(1, 3), images(3, 3))
431 4 : ALLOCATE (points%cell_to_index(-1:1, 0:0, 0:0), points%sab_nl(1))
432 8 : DO img = 1, 3
433 6 : points%cell_to_index(img - 2, 0, 0) = img
434 24 : images(:, img) = [img - 2, 0, 0]
435 8 : ALLOCATE (source(1, img)%matrix)
436 : END DO
437 2 : values = RESHAPE([0.3_dp, -0.2_dp, 0.7_dp, 0.1_dp, -0.4_dp, 0.6_dp], [2, 3])
438 6 : DO mode = 1, 2
439 4 : matrix_type = MERGE(dbcsr_type_symmetric, dbcsr_type_antisymmetric, mode == 1)
440 4 : mirror_type = MERGE(dbcsr_type_antisymmetric, dbcsr_type_symmetric, mode == 1)
441 4 : sign = MERGE(1.0_dp, -1.0_dp, mode == 1)
442 14 : DO orientation = 1, 2
443 8 : owner = MODULO(orientation - 1, comm%num_pe)
444 24 : row_dist = [owner, MODULO(owner + 1, comm%num_pe)]
445 : CALL dbcsr_distribution_new(distribution, group=comm%get_handle(), pgrid=process_grid, &
446 8 : row_dist=row_dist, col_dist=col_dist)
447 8 : atom = orientation
448 8 : neighbor = 3 - orientation
449 32 : image = [3 - 2*orientation, 0, 0]
450 8 : CALL allocate_neighbor_list_set(points%sab_nl(1)%neighbor_list_set, symmetric=.TRUE.)
451 8 : CALL add_neighbor_list(points%sab_nl(1)%neighbor_list_set, atom, neighbors%neighbor_list)
452 32 : CALL add_neighbor_node(neighbors%neighbor_list, neighbor, image, REAL(image, dp))
453 32 : DO img = 1, 3
454 : CALL dbcsr_create(source(1, img)%matrix, name="Image reconstruction", dist=distribution, &
455 24 : matrix_type=matrix_type, row_blk_size=sizes, col_blk_size=sizes)
456 24 : IF (comm%mepos == owner .AND. img == image(1) + 2) THEN
457 4 : CALL dbcsr_put_block(source(1, img)%matrix, 1, 2, values)
458 : END IF
459 32 : CALL dbcsr_finalize(source(1, img)%matrix)
460 : END DO
461 : CALL dbcsr_create(cosine_matrix, name="Image cosine", dist=distribution, matrix_type=matrix_type, &
462 8 : row_blk_size=sizes, col_blk_size=sizes)
463 : CALL dbcsr_create(sine_matrix, name="Image sine", dist=distribution, matrix_type=mirror_type, &
464 8 : row_blk_size=sizes, col_blk_size=sizes)
465 8 : IF (comm%mepos == owner) THEN
466 40 : CALL dbcsr_put_block(cosine_matrix, 1, 2, 0.0_dp*values)
467 40 : CALL dbcsr_put_block(sine_matrix, 1, 2, 0.0_dp*values)
468 : END IF
469 8 : CALL dbcsr_finalize(cosine_matrix)
470 8 : CALL dbcsr_finalize(sine_matrix)
471 8 : expanded = 0.0_dp
472 8 : CALL replicate_rs_matrices(source, points, expanded, points%cell_to_index)
473 24 : DO ik = 1, 2
474 64 : point = [0.137_dp, 0.217_dp, 0.089_dp]*REAL(ik - 1, dp)
475 16 : CALL rs_to_kp(expanded(1, :, :, :), RESULT, images, point)
476 16 : phase = EXP(gaussi*twopi*point(1))
477 16 : expected = z_zero
478 160 : expected(1:2, 3:5) = phase*values
479 272 : expected(3:5, 1:2) = sign*CONJG(TRANSPOSE(expected(1:2, 3:5)))
480 496 : IF (MAXVAL(ABS(RESULT - expected)) > 1.e-12_dp) THEN
481 0 : ERROR STOP "Image reconstruction depends on neighbor orientation"
482 : END IF
483 : CALL rskp_transform(cosine_matrix, sine_matrix, source, 1, point, &
484 16 : points%cell_to_index, points%sab_nl)
485 16 : CALL dbcsr_get_block_p(cosine_matrix, 1, 2, cosine, found)
486 16 : IF (comm%mepos == owner .AND. .NOT. found) ERROR STOP "Missing image cosine block"
487 24 : IF (found) THEN
488 8 : CALL dbcsr_get_block_p(sine_matrix, 1, 2, sine, found)
489 8 : IF (.NOT. found) ERROR STOP "Missing image sine block"
490 80 : IF (MAXVAL(ABS(RESULT(1:2, 3:5) - CMPLX(cosine, sine, KIND=dp))) > 1.e-12_dp) THEN
491 0 : ERROR STOP "Image reconstruction differs from rskp_transform"
492 : END IF
493 80 : cosine = 0.0_dp
494 80 : sine = 0.0_dp
495 : END IF
496 : END DO
497 8 : CALL dbcsr_release(cosine_matrix)
498 8 : CALL dbcsr_release(sine_matrix)
499 32 : DO img = 1, 3
500 32 : CALL dbcsr_release(source(1, img)%matrix)
501 : END DO
502 8 : CALL dbcsr_distribution_release(distribution)
503 12 : CALL deallocate_neighbor_list_set(points%sab_nl(1)%neighbor_list_set)
504 : END DO
505 : END DO
506 8 : DO img = 1, 3
507 8 : DEALLOCATE (source(1, img)%matrix)
508 : END DO
509 2 : DEALLOCATE (points%cell_to_index, points%sab_nl)
510 2 : DEALLOCATE (points, source, images)
511 2 : DEALLOCATE (process_grid, row_dist, col_dist, sizes)
512 2 : CALL dbcsr_finalize_lib()
513 2 : CALL cp_rm_default_logger()
514 2 : CALL cp_logger_release(logger)
515 2 : CALL mp_para_env_release(para_env)
516 4 : END SUBROUTINE check_image_reconstruction
517 :
518 : ! **************************************************************************************************
519 : !> \brief Apply the balanced two-level inverse used by the production local correction.
520 : !> \param inverse_k block-circulant base inverse
521 : !> \param lattice_dims replicated-cell dimensions
522 : !> \param correction_indices selected coordinate basis
523 : !> \param operator_columns columns A U
524 : !> \param coarse_inverse inverse of U^T A U
525 : !> \param input right-hand sides
526 : !> \param output corrected inverse action
527 : ! **************************************************************************************************
528 20 : SUBROUTINE apply_balanced_local_inverse(inverse_k, lattice_dims, correction_indices, &
529 4 : operator_columns, coarse_inverse, input, output)
530 :
531 : COMPLEX(KIND=dp), DIMENSION(:, :, :), INTENT(IN) :: inverse_k
532 : INTEGER, DIMENSION(3), INTENT(IN) :: lattice_dims
533 : INTEGER, DIMENSION(:), INTENT(IN) :: correction_indices
534 : REAL(KIND=dp), DIMENSION(:, :), INTENT(IN) :: operator_columns, coarse_inverse, input
535 : REAL(KIND=dp), DIMENSION(:, :), INTENT(OUT) :: output
536 :
537 : INTEGER :: i
538 : LOGICAL :: used_fft
539 4 : REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :) :: base_input, coarse_input, &
540 4 : coarse_projected, coarse_residual, &
541 : coarse_rhs
542 :
543 : ALLOCATE (base_input(SIZE(input, 1), SIZE(input, 2)), &
544 : coarse_input(SIZE(correction_indices), SIZE(input, 2)), &
545 : coarse_projected(SIZE(correction_indices), SIZE(input, 2)), &
546 : coarse_residual(SIZE(correction_indices), SIZE(input, 2)), &
547 52 : coarse_rhs(SIZE(correction_indices), SIZE(input, 2)))
548 12 : DO i = 1, SIZE(correction_indices)
549 32 : coarse_input(i, :) = input(correction_indices(i), :)
550 : END DO
551 74 : coarse_rhs(:, :) = MATMUL(coarse_inverse, coarse_input)
552 528 : base_input(:, :) = input - MATMUL(operator_columns, coarse_rhs)
553 4 : CALL apply_lattice_inverse_dense(inverse_k, lattice_dims, base_input, output, used_fft)
554 4 : IF (.NOT. used_fft) ERROR STOP "Balanced correction base did not use lattice FFT"
555 352 : coarse_projected(:, :) = MATMUL(TRANSPOSE(operator_columns), output)
556 74 : coarse_residual(:, :) = MATMUL(coarse_inverse, coarse_projected)
557 34 : output(correction_indices, :) = output(correction_indices, :) + coarse_rhs - coarse_residual
558 4 : DEALLOCATE (base_input, coarse_input, coarse_projected, coarse_residual, coarse_rhs)
559 :
560 2 : END SUBROUTINE apply_balanced_local_inverse
561 :
562 : ! **************************************************************************************************
563 : !> \brief ...
564 : !> \param values_rs ...
565 : !> \param index_to_cell ...
566 : !> \param xkp ...
567 : !> \param values_k ...
568 : !> \param deriv_direction ...
569 : !> \param hmat ...
570 : ! **************************************************************************************************
571 8 : SUBROUTINE direct_reference(values_rs, index_to_cell, xkp, values_k, deriv_direction, hmat)
572 :
573 : REAL(KIND=dp), DIMENSION(:, :, :), INTENT(IN) :: values_rs
574 : INTEGER, DIMENSION(:, :), INTENT(IN) :: index_to_cell
575 : REAL(KIND=dp), DIMENSION(:, :), INTENT(IN) :: xkp
576 : COMPLEX(KIND=dp), DIMENSION(:, :, :), INTENT(OUT) :: values_k
577 : INTEGER, INTENT(IN), OPTIONAL :: deriv_direction
578 : REAL(KIND=dp), DIMENSION(3, 3), INTENT(IN), &
579 : OPTIONAL :: hmat
580 :
581 : COMPLEX(KIND=dp) :: factor
582 : INTEGER :: icell, ik
583 : REAL(KIND=dp), DIMENSION(3) :: cell_index_real, cell_vector
584 :
585 4168 : values_k = z_zero
586 424 : DO ik = 1, SIZE(xkp, 2)
587 13320 : DO icell = 1, SIZE(values_rs, 3)
588 51584 : factor = EXP(gaussi*twopi*SUM(xkp(:, ik)*REAL(index_to_cell(:, icell), KIND=dp)))
589 12896 : IF (PRESENT(deriv_direction)) THEN
590 15872 : cell_index_real = REAL(index_to_cell(:, icell), KIND=dp)
591 51584 : cell_vector = MATMUL(hmat, cell_index_real)
592 3968 : factor = factor*gaussi*cell_vector(deriv_direction)
593 : END IF
594 129376 : values_k(:, :, ik) = values_k(:, :, ik) + factor*values_rs(:, :, icell)
595 : END DO
596 : END DO
597 :
598 8 : END SUBROUTINE direct_reference
599 :
600 : END PROGRAM kpoint_lattice_fft_unittest
|
