(git:5e7fe52)
Loading...
Searching...
No Matches
lattice_preconditioner_operator.F
Go to the documentation of this file.
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 Block-circulant inverse operators for replicated-cell OT preconditioners.
10! **************************************************************************************************
12 USE fft_tools, ONLY: bwfft,&
14 fwfft,&
16 fft_alloc,&
17 fft_dealloc,&
19 USE ieee_arithmetic, ONLY: ieee_is_finite
20 USE kinds, ONLY: dp
22#include "./base/base_uses.f90"
23
24 IMPLICIT NONE
25 PRIVATE
26
27 CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'lattice_preconditioner_operator'
28
29 INTEGER, PARAMETER, PUBLIC :: lattice_fft_auto_min_cells = 8
30 REAL(kind=dp), PARAMETER, PUBLIC :: lattice_fft_auto_max_error = 1.0e-3_dp
31 REAL(kind=dp), PARAMETER, PUBLIC :: lattice_fft_auto_max_cost_ratio = 0.8_dp
32 REAL(kind=dp), PARAMETER, PUBLIC :: lattice_fft_auto_max_storage_ratio = 4.0_dp
33
45
46CONTAINS
47
48! **************************************************************************************************
49!> \brief Enumerate a regular direct/reciprocal lattice with the first direction varying fastest.
50!> \param lattice_dims number of cells in every lattice direction
51!> \param index_to_cell integer direct-lattice coordinates
52!> \param xkp fractional reciprocal-lattice coordinates
53! **************************************************************************************************
54 SUBROUTINE lattice_grid(lattice_dims, index_to_cell, xkp)
55
56 INTEGER, DIMENSION(3), INTENT(IN) :: lattice_dims
57 INTEGER, DIMENSION(:, :), INTENT(OUT) :: index_to_cell
58 REAL(kind=dp), DIMENSION(:, :), INTENT(OUT) :: xkp
59
60 INTEGER :: icell, ix, iy, iz
61
62 IF (any(lattice_dims <= 0)) cpabort("Invalid lattice dimensions")
63 IF (SIZE(index_to_cell, 1) < 3 .OR. SIZE(xkp, 1) < 3 .OR. &
64 SIZE(index_to_cell, 2) /= product(lattice_dims) .OR. &
65 SIZE(xkp, 2) /= product(lattice_dims)) THEN
66 cpabort("Invalid lattice-grid output shape")
67 END IF
68
69 icell = 0
70 DO iz = 0, lattice_dims(3) - 1
71 DO iy = 0, lattice_dims(2) - 1
72 DO ix = 0, lattice_dims(1) - 1
73 icell = icell + 1
74 index_to_cell(:, icell) = [ix, iy, iz]
75 xkp(:, icell) = [real(ix, kind=dp)/real(lattice_dims(1), kind=dp), &
76 REAL(iy, kind=dp)/real(lattice_dims(2), kind=dp), &
77 REAL(iz, kind=dp)/real(lattice_dims(3), kind=dp)]
78 END DO
79 END DO
80 END DO
81
82 END SUBROUTINE lattice_grid
83
84! **************************************************************************************************
85!> \brief Project a cell-major SPD operator onto lattice translations and invert its Fourier blocks.
86!> \param operator_matrix dense SPD matrix with identical cell blocks stored consecutively
87!> \param lattice_dims number of replicated cells in every lattice direction
88!> \param inverse_k inverse Hermitian operator block at every reciprocal-grid point
89!> \param projection_error relative Frobenius error of the block-circulant projection
90!> \param used_fft reports whether construction used the FFT path
91!> \param info zero on success; otherwise the failing Cholesky factorization index
92! **************************************************************************************************
93 SUBROUTINE build_lattice_inverse(operator_matrix, lattice_dims, inverse_k, projection_error, &
94 used_fft, info)
95
96 REAL(kind=dp), DIMENSION(:, :), INTENT(IN) :: operator_matrix
97 INTEGER, DIMENSION(3), INTENT(IN) :: lattice_dims
98 COMPLEX(KIND=dp), ALLOCATABLE, &
99 DIMENSION(:, :, :), INTENT(OUT) :: inverse_k
100 REAL(kind=dp), INTENT(OUT) :: projection_error
101 LOGICAL, INTENT(OUT) :: used_fft
102 INTEGER, INTENT(OUT) :: info
103
104 INTEGER :: n, ncell
105 REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :, :) :: blocks
106
107 IF (SIZE(operator_matrix, 1) /= SIZE(operator_matrix, 2)) THEN
108 cpabort("Lattice preconditioner operator must be square")
109 END IF
110 IF (any(lattice_dims <= 0)) cpabort("Lattice dimensions must be positive")
111 n = SIZE(operator_matrix, 1)
112 ncell = product(lattice_dims)
113 IF (modulo(n, ncell) /= 0) THEN
114 cpabort("AO dimension does not match replicated cells")
115 END IF
116 CALL project_lattice_operator(operator_matrix, lattice_dims, blocks, projection_error)
117 CALL build_lattice_inverse_from_blocks(blocks, lattice_dims, inverse_k, used_fft, info)
118 DEALLOCATE (blocks)
119
120 END SUBROUTINE build_lattice_inverse
121
122! **************************************************************************************************
123!> \brief Transform and invert already projected real-space lattice blocks.
124!> \param blocks translation-averaged real-space blocks
125!> \param lattice_dims number of replicated cells in every lattice direction
126!> \param inverse_k inverse Hermitian operator block at every reciprocal-grid point
127!> \param used_fft reports whether construction used the FFT path
128!> \param info zero on success; otherwise the failing Cholesky factorization index
129! **************************************************************************************************
130 SUBROUTINE build_lattice_inverse_from_blocks(blocks, lattice_dims, inverse_k, used_fft, info)
131
132 REAL(kind=dp), DIMENSION(:, :, :), INTENT(IN) :: blocks
133 INTEGER, DIMENSION(3), INTENT(IN) :: lattice_dims
134 COMPLEX(KIND=dp), ALLOCATABLE, &
135 DIMENSION(:, :, :), INTENT(OUT) :: inverse_k
136 LOGICAL, INTENT(OUT) :: used_fft
137 INTEGER, INTENT(OUT) :: info
138
139 INTEGER :: block_size, ncell
140 INTEGER, ALLOCATABLE, DIMENSION(:, :) :: index_to_cell
141 REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :) :: xkp
142
143 IF (any(lattice_dims <= 0)) cpabort("Lattice dimensions must be positive")
144 ncell = product(lattice_dims)
145 block_size = SIZE(blocks, 1)
146 IF (block_size <= 0 .OR. SIZE(blocks, 2) /= block_size .OR. SIZE(blocks, 3) /= ncell) THEN
147 cpabort("Projected lattice blocks have inconsistent dimensions")
148 END IF
149
150 ALLOCATE (index_to_cell(3, ncell), xkp(3, ncell))
151 CALL lattice_grid(lattice_dims, index_to_cell, xkp)
152 ALLOCATE (inverse_k(block_size, block_size, ncell))
153 CALL cell_to_k_grid_fft(blocks, index_to_cell, xkp, lattice_dims, inverse_k, used_fft)
154 CALL invert_hermitian_blocks(inverse_k, info)
155 DEALLOCATE (index_to_cell, xkp)
156
158
159! **************************************************************************************************
160!> \brief Build a balanced local correction for selected non-circulant cell blocks.
161!> \param operator_matrix dense SPD operator
162!> \param lattice_dims number of replicated cells in every lattice direction
163!> \param local_cell_count number of cell blocks retained in the correction
164!> \param correction_indices selected global AO indices
165!> \param operator_columns columns A U for the selected coordinate basis U
166!> \param coarse_inverse inverse of U^T A U
167!> \param corrected_projection_error residual projection error after removing selected rows/columns
168!> \param info zero on success; otherwise the failing Cholesky factorization index
169! **************************************************************************************************
170 SUBROUTINE build_lattice_local_correction(operator_matrix, lattice_dims, local_cell_count, &
171 correction_indices, operator_columns, coarse_inverse, &
172 corrected_projection_error, info)
173
174 REAL(kind=dp), DIMENSION(:, :), INTENT(IN) :: operator_matrix
175 INTEGER, DIMENSION(3), INTENT(IN) :: lattice_dims
176 INTEGER, INTENT(IN) :: local_cell_count
177 INTEGER, ALLOCATABLE, DIMENSION(:), INTENT(OUT) :: correction_indices
178 REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :), &
179 INTENT(OUT) :: operator_columns, coarse_inverse
180 REAL(kind=dp), INTENT(OUT) :: corrected_projection_error
181 INTEGER, INTENT(OUT) :: info
182
183 INTEGER :: block_size, cell, chosen, i, index, &
184 local_rank, n, ncell, row_first, &
185 row_last
186 INTEGER, ALLOCATABLE, DIMENSION(:) :: selected_cells
187 LOGICAL, ALLOCATABLE, DIMENSION(:) :: selected
188 REAL(kind=dp) :: denominator
189 REAL(kind=dp), ALLOCATABLE, DIMENSION(:) :: scores
190 REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :) :: projected_operator, residual
191 REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :, :) :: blocks
192
193 IF (SIZE(operator_matrix, 1) /= SIZE(operator_matrix, 2)) THEN
194 cpabort("Lattice local-correction operator must be square")
195 END IF
196 IF (any(lattice_dims <= 0)) cpabort("Lattice dimensions must be positive")
197 n = SIZE(operator_matrix, 1)
198 ncell = product(lattice_dims)
199 IF (modulo(n, ncell) /= 0) cpabort("AO dimension does not match replicated cells")
200 IF (local_cell_count <= 0 .OR. local_cell_count >= ncell) THEN
201 cpabort("Lattice local correction requires between one and ncell-1 cells")
202 END IF
203 block_size = n/ncell
204 local_rank = local_cell_count*block_size
205
206 CALL project_lattice_operator(operator_matrix, lattice_dims, blocks, &
207 corrected_projection_error, projected_operator)
208 DEALLOCATE (blocks)
209 ALLOCATE (residual(n, n))
210 residual(:, :) = operator_matrix - projected_operator
211 DEALLOCATE (projected_operator)
212 ALLOCATE (scores(ncell), selected(ncell), selected_cells(local_cell_count))
213 selected = .false.
214 DO chosen = 1, local_cell_count
215 scores = -1.0_dp
216 DO cell = 1, ncell
217 IF (selected(cell)) cycle
218 row_first = (cell - 1)*block_size + 1
219 row_last = cell*block_size
220 scores(cell) = sum(residual(row_first:row_last, :)**2) + &
221 sum(residual(:, row_first:row_last)**2) - &
222 sum(residual(row_first:row_last, row_first:row_last)**2)
223 END DO
224 index = maxloc(scores, dim=1, mask=.NOT. selected)
225 selected(index) = .true.
226 selected_cells(chosen) = index
227 row_first = (index - 1)*block_size + 1
228 row_last = index*block_size
229 residual(row_first:row_last, :) = 0.0_dp
230 residual(:, row_first:row_last) = 0.0_dp
231 END DO
232
233 ALLOCATE (correction_indices(local_rank))
234 index = 0
235 DO chosen = 1, local_cell_count
236 cell = selected_cells(chosen)
237 DO i = 1, block_size
238 index = index + 1
239 correction_indices(index) = (cell - 1)*block_size + i
240 END DO
241 END DO
242
243 denominator = sum(operator_matrix**2)
244 corrected_projection_error = sqrt(sum(residual**2)/max(denominator, tiny(denominator)))
245
246 ALLOCATE (operator_columns(n, local_rank), coarse_inverse(local_rank, local_rank))
247 operator_columns(:, :) = operator_matrix(:, correction_indices)
248 coarse_inverse(:, :) = operator_matrix(correction_indices, correction_indices)
249 CALL dpotrf('U', local_rank, coarse_inverse, local_rank, info)
250 IF (info == 0) CALL dpotri('U', local_rank, coarse_inverse, local_rank, info)
251 IF (info == 0) THEN
252 DO i = 1, local_rank - 1
253 coarse_inverse(i + 1:local_rank, i) = coarse_inverse(i, i + 1:local_rank)
254 END DO
255 END IF
256
257 DEALLOCATE (residual, scores, selected, selected_cells)
258
259 END SUBROUTINE build_lattice_local_correction
260
261! **************************************************************************************************
262!> \brief Project a dense cell-major operator onto its block-circulant translation average.
263!> \param operator_matrix dense square operator
264!> \param lattice_dims number of replicated cells in every lattice direction
265!> \param blocks translation-averaged real-space blocks
266!> \param projection_error relative Frobenius projection error
267!> \param projected_operator optional dense block-circulant projection
268! **************************************************************************************************
269 SUBROUTINE project_lattice_operator(operator_matrix, lattice_dims, blocks, projection_error, &
270 projected_operator)
271
272 REAL(kind=dp), DIMENSION(:, :), INTENT(IN) :: operator_matrix
273 INTEGER, DIMENSION(3), INTENT(IN) :: lattice_dims
274 REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :, :), &
275 INTENT(OUT) :: blocks
276 REAL(kind=dp), INTENT(OUT) :: projection_error
277 REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :), &
278 INTENT(OUT), OPTIONAL :: projected_operator
279
280 INTEGER :: alpha, beta, block_size, delta_index, &
281 icell, jcell, n, ncell
282 INTEGER, ALLOCATABLE, DIMENSION(:, :) :: index_to_cell
283 INTEGER, DIMENSION(3) :: delta
284 REAL(kind=dp) :: denominator, difference, numerator
285 REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :) :: xkp
286
287 IF (SIZE(operator_matrix, 1) /= SIZE(operator_matrix, 2)) THEN
288 cpabort("Lattice preconditioner operator must be square")
289 END IF
290 IF (any(lattice_dims <= 0)) cpabort("Lattice dimensions must be positive")
291 n = SIZE(operator_matrix, 1)
292 ncell = product(lattice_dims)
293 IF (modulo(n, ncell) /= 0) cpabort("AO dimension does not match replicated cells")
294 block_size = n/ncell
295
296 ALLOCATE (index_to_cell(3, ncell), xkp(3, ncell))
297 CALL lattice_grid(lattice_dims, index_to_cell, xkp)
298 ALLOCATE (blocks(block_size, block_size, ncell), source=0.0_dp)
299 DO jcell = 1, ncell
300 DO icell = 1, ncell
301 ! With X(k)=sum_R exp(+ikR) X(R), the convolution block is indexed by R-T.
302 delta = modulo(index_to_cell(:, icell) - index_to_cell(:, jcell), lattice_dims)
303 delta_index = cell_linear_index(delta, lattice_dims)
304 DO beta = 1, block_size
305 DO alpha = 1, block_size
306 blocks(alpha, beta, delta_index) = blocks(alpha, beta, delta_index) + &
307 operator_matrix((icell - 1)*block_size + alpha, &
308 (jcell - 1)*block_size + beta)/ &
309 REAL(ncell, kind=dp)
310 END DO
311 END DO
312 END DO
313 END DO
314
315 IF (PRESENT(projected_operator)) ALLOCATE (projected_operator(n, n))
316 numerator = 0.0_dp
317 denominator = 0.0_dp
318 DO jcell = 1, ncell
319 DO icell = 1, ncell
320 delta = modulo(index_to_cell(:, icell) - index_to_cell(:, jcell), lattice_dims)
321 delta_index = cell_linear_index(delta, lattice_dims)
322 DO beta = 1, block_size
323 DO alpha = 1, block_size
324 IF (PRESENT(projected_operator)) THEN
325 projected_operator((icell - 1)*block_size + alpha, &
326 (jcell - 1)*block_size + beta) = &
327 blocks(alpha, beta, delta_index)
328 END IF
329 difference = operator_matrix((icell - 1)*block_size + alpha, &
330 (jcell - 1)*block_size + beta) - &
331 blocks(alpha, beta, delta_index)
332 numerator = numerator + difference*difference
333 denominator = denominator + &
334 operator_matrix((icell - 1)*block_size + alpha, &
335 (jcell - 1)*block_size + beta)**2
336 END DO
337 END DO
338 END DO
339 END DO
340 projection_error = sqrt(numerator/max(denominator, tiny(denominator)))
341
342 DEALLOCATE (index_to_cell, xkp)
343
344 END SUBROUTINE project_lattice_operator
345
346! **************************************************************************************************
347!> \brief Estimate lattice-FFT application work relative to a dense inverse application.
348!> \param lattice_dims number of explicitly replicated cells
349!> \param block_size number of orbitals in the reference cell
350!> \param rhs_count number of orbital columns transformed together
351!> \param num_pe number of message-passing ranks sharing the operator
352!> \return estimated lattice-to-dense work ratio
353! **************************************************************************************************
354 PURE REAL(KIND=dp) FUNCTION lattice_fft_cost_ratio(lattice_dims, block_size, rhs_count, num_pe) &
355 result(ratio)
356
357 INTEGER, DIMENSION(3), INTENT(IN) :: lattice_dims
358 INTEGER, INTENT(IN) :: block_size, rhs_count, num_pe
359
360 INTEGER :: ncell, nrhs
361 REAL(kind=dp) :: block_size_real, communication_work, &
362 dense_work, fft_work, matrix_work, &
363 ncell_real, nrhs_real, setup_work
364
365 ncell = max(1, product(lattice_dims))
366 nrhs = max(1, rhs_count)
367 IF (block_size <= 0) THEN
368 ratio = huge(ratio)
369 RETURN
370 END IF
371 ncell_real = real(ncell, kind=dp)
372 nrhs_real = real(nrhs, kind=dp)
373 block_size_real = real(block_size, kind=dp)
374
375 ! This deliberately simple model counts real dense multiply-adds as the reference.
376 ! Four scalar operations per FFT butterfly account for the two complex transforms;
377 ! the fixed packing term prevents AUTO from selecting tiny batches. The logarithmic
378 ! communication term represents the input/output all-to-all routing.
379 dense_work = (ncell_real*block_size_real)**2*nrhs_real
380 matrix_work = ncell_real*block_size_real**2*nrhs_real
381 fft_work = 4.0_dp*ncell_real*block_size_real*nrhs_real* &
382 log(real(max(2, ncell), kind=dp))/log(2.0_dp)
383 setup_work = 8.0_dp*ncell_real*block_size_real
384 communication_work = 2.0_dp*ncell_real*block_size_real*nrhs_real* &
385 log(real(max(2, num_pe + 1), kind=dp))/log(2.0_dp)
386 ratio = (matrix_work + fft_work + setup_work + communication_work)/max(dense_work, tiny(dense_work))
387
388 END FUNCTION lattice_fft_cost_ratio
389
390! **************************************************************************************************
391!> \brief Estimate reciprocal-block storage relative to one dense real operator.
392!> \param lattice_dims number of explicitly replicated cells
393!> \param operator_count number of distinct state-dependent reciprocal operators
394!> \return estimated lattice-to-dense storage ratio
395! **************************************************************************************************
396 PURE REAL(KIND=dp) FUNCTION lattice_fft_storage_ratio(lattice_dims, operator_count) RESULT(ratio)
397
398 INTEGER, DIMENSION(3), INTENT(IN) :: lattice_dims
399 INTEGER, INTENT(IN) :: operator_count
400
401 INTEGER :: ncell
402
403 ncell = max(1, product(lattice_dims))
404 ! Reciprocal blocks are complex; the block-size dependence cancels against
405 ! the reference dense matrix. A state-independent operator has count one.
406 ratio = 2.0_dp*real(max(1, operator_count), kind=dp)/real(ncell, kind=dp)
407
408 END FUNCTION lattice_fft_storage_ratio
409
410! **************************************************************************************************
411!> \brief Estimate balanced-local-correction work relative to a dense inverse application.
412!> \param lattice_dims number of replicated cells
413!> \param local_cell_count number of cell blocks retained in the correction
414!> \return estimated correction-to-dense work ratio
415! **************************************************************************************************
416 PURE REAL(KIND=dp) FUNCTION lattice_local_cost_ratio(lattice_dims, local_cell_count) &
417 result(ratio)
418
419 INTEGER, DIMENSION(3), INTENT(IN) :: lattice_dims
420 INTEGER, INTENT(IN) :: local_cell_count
421
422 INTEGER :: ncell
423 REAL(kind=dp) :: local_fraction
424
425 ncell = max(1, product(lattice_dims))
426 IF (local_cell_count <= 0) THEN
427 ratio = 0.0_dp
428 RETURN
429 END IF
430 local_fraction = real(local_cell_count, kind=dp)/real(ncell, kind=dp)
431 ! Both corrected and dense applications scale linearly with the right-hand-side count.
432 ratio = 2.0_dp*(local_fraction + local_fraction**2)
433
434 END FUNCTION lattice_local_cost_ratio
435
436! **************************************************************************************************
437!> \brief Estimate balanced-local-correction storage relative to one dense real operator.
438!> \param lattice_dims number of replicated cells
439!> \param operator_count number of distinct state-dependent local corrections
440!> \param local_cell_count number of cell blocks retained in each correction
441!> \return estimated correction-to-dense storage ratio
442! **************************************************************************************************
443 PURE REAL(KIND=dp) FUNCTION lattice_local_storage_ratio(lattice_dims, operator_count, &
444 local_cell_count) RESULT(ratio)
445
446 INTEGER, DIMENSION(3), INTENT(IN) :: lattice_dims
447 INTEGER, INTENT(IN) :: operator_count, local_cell_count
448
449 INTEGER :: ncell
450 REAL(kind=dp) :: local_fraction
451
452 ncell = max(1, product(lattice_dims))
453 IF (local_cell_count <= 0) THEN
454 ratio = 0.0_dp
455 RETURN
456 END IF
457 local_fraction = real(local_cell_count, kind=dp)/real(ncell, kind=dp)
458 ratio = real(max(1, operator_count), kind=dp)*(local_fraction + local_fraction**2)
459
460 END FUNCTION lattice_local_storage_ratio
461
462! **************************************************************************************************
463!> \brief Conservative error-, work-, and storage-aware AUTO selector.
464!> \param lattice_dims number of explicitly replicated cells
465!> \param block_size number of orbitals in the reference cell
466!> \param rhs_count number of orbital columns transformed together
467!> \param operator_count number of distinct reciprocal operators stored
468!> \param num_pe number of message-passing ranks sharing the operator
469!> \param projection_error relative Frobenius projection error
470!> \param used_fft whether construction used the FFT path
471!> \param info zero when all projected Fourier blocks were positive definite
472!> \return true when AUTO may use the lattice inverse
473! **************************************************************************************************
474 PURE LOGICAL FUNCTION lattice_fft_auto_select(lattice_dims, block_size, rhs_count, operator_count, &
475 num_pe, projection_error, used_fft, info) &
476 result(selected)
477
478 INTEGER, DIMENSION(3), INTENT(IN) :: lattice_dims
479 INTEGER, INTENT(IN) :: block_size, rhs_count, operator_count, &
480 num_pe
481 REAL(kind=dp), INTENT(IN) :: projection_error
482 LOGICAL, INTENT(IN) :: used_fft
483 INTEGER, INTENT(IN) :: info
484
485 selected = product(lattice_dims) >= lattice_fft_auto_min_cells .AND. &
486 projection_error <= lattice_fft_auto_max_error .AND. used_fft .AND. info == 0
487 selected = selected .AND. &
488 lattice_fft_cost_ratio(lattice_dims, block_size, rhs_count, num_pe) <= &
490 lattice_fft_storage_ratio(lattice_dims, operator_count) <= &
492
493 END FUNCTION lattice_fft_auto_select
494
495! **************************************************************************************************
496!> \brief Apply inverse Fourier blocks to a dense cell-major real matrix.
497!> \param inverse_k inverse overlap blocks on a complete reciprocal grid
498!> \param lattice_dims number of replicated cells in every lattice direction
499!> \param matrix_in matrix to precondition
500!> \param matrix_out preconditioned matrix
501!> \param used_fft reports whether both transforms used the FFT path
502!> \param imaginary_residual largest discarded imaginary component
503! **************************************************************************************************
504 SUBROUTINE apply_lattice_inverse_dense(inverse_k, lattice_dims, matrix_in, matrix_out, &
505 used_fft, imaginary_residual)
506
507 COMPLEX(KIND=dp), DIMENSION(:, :, :), INTENT(IN) :: inverse_k
508 INTEGER, DIMENSION(3), INTENT(IN) :: lattice_dims
509 REAL(kind=dp), DIMENSION(:, :), INTENT(IN) :: matrix_in
510 REAL(kind=dp), DIMENSION(:, :), INTENT(OUT) :: matrix_out
511 LOGICAL, INTENT(OUT) :: used_fft
512 REAL(kind=dp), INTENT(OUT), OPTIONAL :: imaginary_residual
513
514 COMPLEX(KIND=dp), ALLOCATABLE, DIMENSION(:, :, :) :: input_k, input_rs, output_k, output_rs
515 INTEGER :: block_size, icell, ncell
516 LOGICAL :: forward_fft, inverse_fft
517
518 IF (any(lattice_dims <= 0)) cpabort("Lattice dimensions must be positive")
519 ncell = product(lattice_dims)
520 block_size = SIZE(inverse_k, 1)
521 IF (SIZE(inverse_k, 2) /= block_size .OR. SIZE(inverse_k, 3) /= ncell) THEN
522 cpabort("Inconsistent inverse blocks in lattice preconditioner")
523 END IF
524 IF (SIZE(matrix_in, 1) /= block_size*ncell .OR. &
525 any(shape(matrix_out) /= shape(matrix_in))) THEN
526 cpabort("Inconsistent dense matrix shape in lattice preconditioner")
527 END IF
528
529 ALLOCATE (input_rs(block_size, SIZE(matrix_in, 2), ncell))
530 DO icell = 1, ncell
531 input_rs(:, :, icell) = cmplx( &
532 matrix_in((icell - 1)*block_size + 1:icell*block_size, :), 0.0_dp, kind=dp)
533 END DO
534
535 ALLOCATE (input_k(block_size, SIZE(matrix_in, 2), ncell))
536 ALLOCATE (output_k(block_size, SIZE(matrix_in, 2), ncell))
537 ALLOCATE (output_rs(block_size, SIZE(matrix_in, 2), ncell))
538 CALL batched_lattice_fft(input_rs, lattice_dims, bwfft, input_k, forward_fft)
539 DO icell = 1, ncell
540 output_k(:, :, icell) = matmul(inverse_k(:, :, icell), input_k(:, :, icell))
541 END DO
542 CALL batched_lattice_fft(output_k, lattice_dims, fwfft, output_rs, inverse_fft)
543 used_fft = forward_fft .AND. inverse_fft
544 IF (PRESENT(imaginary_residual)) imaginary_residual = maxval(abs(aimag(output_rs)))
545
546 DO icell = 1, ncell
547 matrix_out((icell - 1)*block_size + 1:icell*block_size, :) = real(output_rs(:, :, icell), kind=dp)
548 END DO
549
550 DEALLOCATE (input_k, input_rs, output_k, output_rs)
551
552 END SUBROUTINE apply_lattice_inverse_dense
553
554! **************************************************************************************************
555!> \brief Apply one inverse Fourier operator to every column of a dense real matrix.
556!> \param inverse_k state-dependent inverse blocks on a complete reciprocal grid
557!> \param lattice_dims number of replicated cells in every lattice direction
558!> \param matrix_in matrix to precondition, with one state per column
559!> \param matrix_out preconditioned matrix
560!> \param used_fft reports whether both transforms used the FFT path
561!> \param imaginary_residual largest discarded imaginary component
562!> \param state_indices optional global operator index for every input column
563! **************************************************************************************************
564 SUBROUTINE apply_lattice_state_inverse_dense(inverse_k, lattice_dims, matrix_in, matrix_out, &
565 used_fft, imaginary_residual, state_indices)
566
567 COMPLEX(KIND=dp), DIMENSION(:, :, :, :), &
568 INTENT(IN) :: inverse_k
569 INTEGER, DIMENSION(3), INTENT(IN) :: lattice_dims
570 REAL(kind=dp), DIMENSION(:, :), INTENT(IN) :: matrix_in
571 REAL(kind=dp), DIMENSION(:, :), INTENT(OUT) :: matrix_out
572 LOGICAL, INTENT(OUT) :: used_fft
573 REAL(kind=dp), INTENT(OUT), OPTIONAL :: imaginary_residual
574 INTEGER, DIMENSION(:), INTENT(IN), OPTIONAL :: state_indices
575
576 COMPLEX(KIND=dp), ALLOCATABLE, DIMENSION(:, :, :) :: input_k, input_rs, output_k, output_rs
577 INTEGER :: block_size, icell, j, ncell
578 LOGICAL :: forward_fft, inverse_fft
579
580 IF (any(lattice_dims <= 0)) cpabort("Lattice dimensions must be positive")
581 ncell = product(lattice_dims)
582 block_size = SIZE(inverse_k, 1)
583 IF (SIZE(inverse_k, 2) /= block_size .OR. SIZE(inverse_k, 3) /= ncell) THEN
584 cpabort("Inconsistent state-dependent inverse blocks in lattice preconditioner")
585 END IF
586 IF (PRESENT(state_indices)) THEN
587 IF (SIZE(state_indices) /= SIZE(matrix_in, 2) .OR. &
588 any(state_indices < 1) .OR. any(state_indices > SIZE(inverse_k, 4))) THEN
589 cpabort("Invalid state indices in lattice preconditioner")
590 END IF
591 ELSE IF (SIZE(inverse_k, 4) /= SIZE(matrix_in, 2)) THEN
592 cpabort("Inconsistent state-dependent inverse blocks in lattice preconditioner")
593 END IF
594 IF (SIZE(matrix_in, 1) /= block_size*ncell .OR. &
595 any(shape(matrix_out) /= shape(matrix_in))) THEN
596 cpabort("Inconsistent dense state matrix shape in lattice preconditioner")
597 END IF
598
599 ALLOCATE (input_rs(block_size, SIZE(matrix_in, 2), ncell))
600 DO icell = 1, ncell
601 input_rs(:, :, icell) = cmplx( &
602 matrix_in((icell - 1)*block_size + 1:icell*block_size, :), 0.0_dp, kind=dp)
603 END DO
604
605 ALLOCATE (input_k(block_size, SIZE(matrix_in, 2), ncell))
606 ALLOCATE (output_k(block_size, SIZE(matrix_in, 2), ncell))
607 ALLOCATE (output_rs(block_size, SIZE(matrix_in, 2), ncell))
608 CALL batched_lattice_fft(input_rs, lattice_dims, bwfft, input_k, forward_fft)
609 DO icell = 1, ncell
610 DO j = 1, SIZE(matrix_in, 2)
611 IF (PRESENT(state_indices)) THEN
612 output_k(:, j, icell) = &
613 matmul(inverse_k(:, :, icell, state_indices(j)), input_k(:, j, icell))
614 ELSE
615 output_k(:, j, icell) = matmul(inverse_k(:, :, icell, j), input_k(:, j, icell))
616 END IF
617 END DO
618 END DO
619 CALL batched_lattice_fft(output_k, lattice_dims, fwfft, output_rs, inverse_fft)
620 used_fft = forward_fft .AND. inverse_fft
621 IF (PRESENT(imaginary_residual)) imaginary_residual = maxval(abs(aimag(output_rs)))
622
623 DO icell = 1, ncell
624 matrix_out((icell - 1)*block_size + 1:icell*block_size, :) = real(output_rs(:, :, icell), kind=dp)
625 END DO
626
627 DEALLOCATE (input_k, input_rs, output_k, output_rs)
628
630
631! **************************************************************************************************
632!> \brief Transform a complete replicated-cell batch with three batched one-dimensional FFTs.
633!> \param values_in matrices on the complete direct or reciprocal lattice
634!> \param lattice_dims logical lattice dimensions
635!> \param fsign transform direction, BWFFT for cell-to-k and FWFFT for k-to-cell
636!> \param values_out transformed matrices on the complete lattice
637!> \param used_fft reports whether the batched FFT path succeeded
638! **************************************************************************************************
639 SUBROUTINE batched_lattice_fft(values_in, lattice_dims, fsign, values_out, used_fft)
640
641 COMPLEX(KIND=dp), DIMENSION(:, :, :), INTENT(IN) :: values_in
642 INTEGER, DIMENSION(3), INTENT(IN) :: lattice_dims
643 INTEGER, INTENT(IN) :: fsign
644 COMPLEX(KIND=dp), DIMENSION(:, :, :), INTENT(OUT) :: values_out
645 LOGICAL, INTENT(OUT) :: used_fft
646
647 COMPLEX(KIND=dp), ALLOCATABLE, &
648 DIMENSION(:, :, :, :) :: grid
649 INTEGER :: attempt, batch_index, d, i, icell, ix, &
650 iy, iz, j, radix_length, stat
651 INTEGER, DIMENSION(3) :: input_stride, nfft, output_stride, stride
652 LOGICAL :: compatible
653 REAL(kind=dp) :: normalization
654
655 IF (any(lattice_dims <= 0) .OR. SIZE(values_in, 3) /= product(lattice_dims) .OR. &
656 any(shape(values_out) /= shape(values_in))) THEN
657 cpabort("Invalid shape in batched lattice FFT")
658 END IF
659 IF (fsign /= fwfft .AND. fsign /= bwfft) THEN
660 cpabort("Invalid direction in batched lattice FFT")
661 END IF
662 nfft = lattice_dims
663 compatible = .true.
664 DO d = 1, 3
665 compatible = .false.
666 DO attempt = 0, 15
667 IF (nfft(d) >= 3) THEN
668 CALL fft_radix_operations(nfft(d), radix_length, fft_radix_next)
669 IF (radix_length == nfft(d)) THEN
670 compatible = .true.
671 EXIT
672 END IF
673 END IF
674 nfft(d) = nfft(d) + lattice_dims(d)
675 END DO
676 IF (.NOT. compatible) EXIT
677 END DO
678 IF (.NOT. compatible) THEN
679 values_out = cmplx(0.0_dp, 0.0_dp, kind=dp)
680 used_fft = .false.
681 RETURN
682 END IF
683
684 stride = nfft/lattice_dims
685 input_stride = 1
686 output_stride = 1
687 IF (fsign == bwfft) input_stride = stride
688 IF (fsign == fwfft) output_stride = stride
689 ALLOCATE (grid(SIZE(values_in, 1)*SIZE(values_in, 2), nfft(1), nfft(2), nfft(3)))
690 grid = cmplx(0.0_dp, 0.0_dp, kind=dp)
691 icell = 0
692 DO iz = 0, lattice_dims(3) - 1
693 DO iy = 0, lattice_dims(2) - 1
694 DO ix = 0, lattice_dims(1) - 1
695 icell = icell + 1
696 DO j = 1, SIZE(values_in, 2)
697 DO i = 1, SIZE(values_in, 1)
698 batch_index = i + SIZE(values_in, 1)*(j - 1)
699 grid(batch_index, ix*input_stride(1) + 1, iy*input_stride(2) + 1, &
700 iz*input_stride(3) + 1) = values_in(i, j, icell)
701 END DO
702 END DO
703 END DO
704 END DO
705 END DO
706
707 stat = 0
708 DO d = 1, 3
709 CALL transform_grid_dimension(grid, d, fsign, stat)
710 IF (stat /= 0) EXIT
711 END DO
712 IF (stat /= 0) THEN
713 values_out = cmplx(0.0_dp, 0.0_dp, kind=dp)
714 used_fft = .false.
715 DEALLOCATE (grid)
716 RETURN
717 END IF
718
719 normalization = 1.0_dp
720 IF (fsign == fwfft) THEN
721 normalization = real(product(nfft), kind=dp)/real(product(lattice_dims), kind=dp)
722 END IF
723 icell = 0
724 DO iz = 0, lattice_dims(3) - 1
725 DO iy = 0, lattice_dims(2) - 1
726 DO ix = 0, lattice_dims(1) - 1
727 icell = icell + 1
728 DO j = 1, SIZE(values_out, 2)
729 DO i = 1, SIZE(values_out, 1)
730 batch_index = i + SIZE(values_out, 1)*(j - 1)
731 values_out(i, j, icell) = normalization* &
732 grid(batch_index, ix*output_stride(1) + 1, iy*output_stride(2) + 1, &
733 iz*output_stride(3) + 1)
734 END DO
735 END DO
736 END DO
737 END DO
738 END DO
739 used_fft = .true.
740 DEALLOCATE (grid)
741
742 END SUBROUTINE batched_lattice_fft
743
744! **************************************************************************************************
745!> \brief Transform every line of one dimension of a four-dimensional FFT batch.
746!> \param grid batched three-dimensional grid
747!> \param direction grid dimension to transform
748!> \param fsign transform direction
749!> \param stat zero on success
750! **************************************************************************************************
751 SUBROUTINE transform_grid_dimension(grid, direction, fsign, stat)
752
753 COMPLEX(KIND=dp), DIMENSION(:, :, :, :), &
754 INTENT(INOUT) :: grid
755 INTEGER, INTENT(IN) :: direction, fsign
756 INTEGER, INTENT(OUT) :: stat
757
758 COMPLEX(KIND=dp), CONTIGUOUS, DIMENSION(:, :), &
759 POINTER :: line_in, line_out
760 INTEGER :: ibatch, iline, ix, iy, iz, m, n
761 REAL(kind=dp) :: scale
762
763 n = SIZE(grid, direction + 1)
764 m = SIZE(grid)/n
765 CALL fft_alloc(line_in, [n, m])
766 CALL fft_alloc(line_out, [n, m])
767 iline = 0
768 SELECT CASE (direction)
769 CASE (1)
770 DO iz = 1, SIZE(grid, 4)
771 DO iy = 1, SIZE(grid, 3)
772 DO ibatch = 1, SIZE(grid, 1)
773 iline = iline + 1
774 line_in(:, iline) = grid(ibatch, :, iy, iz)
775 END DO
776 END DO
777 END DO
778 CASE (2)
779 DO iz = 1, SIZE(grid, 4)
780 DO ix = 1, SIZE(grid, 2)
781 DO ibatch = 1, SIZE(grid, 1)
782 iline = iline + 1
783 line_in(:, iline) = grid(ibatch, ix, :, iz)
784 END DO
785 END DO
786 END DO
787 CASE (3)
788 DO iy = 1, SIZE(grid, 3)
789 DO ix = 1, SIZE(grid, 2)
790 DO ibatch = 1, SIZE(grid, 1)
791 iline = iline + 1
792 line_in(:, iline) = grid(ibatch, ix, iy, :)
793 END DO
794 END DO
795 END DO
796 CASE DEFAULT
797 cpabort("Invalid batched FFT grid direction")
798 END SELECT
799
800 scale = 1.0_dp
801 IF (fsign == fwfft) scale = 1.0_dp/real(n, kind=dp)
802 CALL fft_1d_many(fsign, n, m, .false., .false., n, n, &
803 line_in, line_out, scale, stat)
804 IF (stat == 0) THEN
805 iline = 0
806 SELECT CASE (direction)
807 CASE (1)
808 DO iz = 1, SIZE(grid, 4)
809 DO iy = 1, SIZE(grid, 3)
810 DO ibatch = 1, SIZE(grid, 1)
811 iline = iline + 1
812 grid(ibatch, :, iy, iz) = line_out(:, iline)
813 END DO
814 END DO
815 END DO
816 CASE (2)
817 DO iz = 1, SIZE(grid, 4)
818 DO ix = 1, SIZE(grid, 2)
819 DO ibatch = 1, SIZE(grid, 1)
820 iline = iline + 1
821 grid(ibatch, ix, :, iz) = line_out(:, iline)
822 END DO
823 END DO
824 END DO
825 CASE (3)
826 DO iy = 1, SIZE(grid, 3)
827 DO ix = 1, SIZE(grid, 2)
828 DO ibatch = 1, SIZE(grid, 1)
829 iline = iline + 1
830 grid(ibatch, ix, iy, :) = line_out(:, iline)
831 END DO
832 END DO
833 END DO
834 END SELECT
835 END IF
836 CALL fft_dealloc(line_in)
837 CALL fft_dealloc(line_out)
838
839 END SUBROUTINE transform_grid_dimension
840
841! **************************************************************************************************
842!> \brief Convert zero-based lattice coordinates to the cell-major linear index.
843!> \param coordinate zero-based lattice coordinate
844!> \param lattice_dims number of cells in every direction
845!> \return one-based linear index
846! **************************************************************************************************
847 PURE INTEGER FUNCTION cell_linear_index(coordinate, lattice_dims) RESULT(index)
848
849 INTEGER, DIMENSION(3), INTENT(IN) :: coordinate, lattice_dims
850
851 index = coordinate(1) + lattice_dims(1)*(coordinate(2) + &
852 lattice_dims(2)*coordinate(3)) + 1
853
854 END FUNCTION cell_linear_index
855
856! **************************************************************************************************
857!> \brief Hermitize, Cholesky factorize, and invert each small complex block.
858!> \param blocks matrix blocks, replaced by their inverse
859!> \param info zero on success; otherwise the one-based failing block index
860! **************************************************************************************************
861 SUBROUTINE invert_hermitian_blocks(blocks, info)
862
863 COMPLEX(KIND=dp), DIMENSION(:, :, :), &
864 INTENT(INOUT) :: blocks
865 INTEGER, INTENT(OUT) :: info
866
867 COMPLEX(KIND=dp) :: value
868 COMPLEX(KIND=dp), ALLOCATABLE, DIMENSION(:) :: solution, work
869 COMPLEX(KIND=dp), ALLOCATABLE, DIMENSION(:, :) :: factor, inverse
870 INTEGER :: i, iblock, j, k, m, n
871 REAL(kind=dp) :: diagonal
872
873 n = SIZE(blocks, 1)
874 cpassert(SIZE(blocks, 2) == n)
875 ALLOCATE (factor(n, n), inverse(n, n), solution(n), work(n))
876 info = 0
877 DO iblock = 1, SIZE(blocks, 3)
878 factor = cmplx(0.0_dp, 0.0_dp, kind=dp)
879 inverse(:, :) = 0.5_dp*(blocks(:, :, iblock) + &
880 conjg(transpose(blocks(:, :, iblock))))
881
882 ! Small-block upper Cholesky factorization, inverse=U^H U.
883 DO j = 1, n
884 diagonal = real(inverse(j, j), kind=dp)
885 DO m = 1, j - 1
886 diagonal = diagonal - abs(factor(m, j))**2
887 END DO
888 IF (.NOT. ieee_is_finite(diagonal) .OR. diagonal <= 0.0_dp) THEN
889 info = iblock
890 DEALLOCATE (factor, inverse, solution, work)
891 RETURN
892 END IF
893 factor(j, j) = cmplx(sqrt(diagonal), 0.0_dp, kind=dp)
894 DO k = j + 1, n
895 value = inverse(j, k)
896 DO m = 1, j - 1
897 value = value - conjg(factor(m, j))*factor(m, k)
898 END DO
899 factor(j, k) = value/factor(j, j)
900 END DO
901 END DO
902
903 ! Solve U^H U x_j=e_j for every inverse column. A small explicit
904 ! triangular solve keeps the complete small-block kernel self-contained.
905 DO j = 1, n
906 work = cmplx(0.0_dp, 0.0_dp, kind=dp)
907 solution = cmplx(0.0_dp, 0.0_dp, kind=dp)
908 DO i = 1, n
909 value = cmplx(0.0_dp, 0.0_dp, kind=dp)
910 IF (i == j) value = cmplx(1.0_dp, 0.0_dp, kind=dp)
911 DO m = 1, i - 1
912 value = value - conjg(factor(m, i))*work(m)
913 END DO
914 work(i) = value/conjg(factor(i, i))
915 END DO
916 DO i = n, 1, -1
917 value = work(i)
918 DO m = i + 1, n
919 value = value - factor(i, m)*solution(m)
920 END DO
921 solution(i) = value/factor(i, i)
922 END DO
923 inverse(:, j) = solution
924 END DO
925 blocks(:, :, iblock) = 0.5_dp*(inverse + conjg(transpose(inverse)))
926 END DO
927 DEALLOCATE (factor, inverse, solution, work)
928
929 END SUBROUTINE invert_hermitian_blocks
930
static GRID_HOST_DEVICE int modulo(int a, int m)
Equivalent of Fortran's MODULO, which always return a positive number. https://gcc....
subroutine, public fft_radix_operations(radix_in, radix_out, operation)
Determine the allowed lengths of FFT's '''.
Definition fft_tools.F:231
integer, parameter, public bwfft
Definition fft_tools.F:146
subroutine, public fft_1d_many(sign, n, m, trans_in, trans_out, ldx_in, ldx_out, zin, zout, scale, stat)
Performs m 1-D forward FFT-s of size n.
Definition fft_tools.F:350
integer, parameter, public fwfft
Definition fft_tools.F:146
integer, parameter, public fft_radix_next
Definition fft_tools.F:147
Defines the basic variable types.
Definition kinds.F:23
integer, parameter, public dp
Definition kinds.F:34
Batched lattice Fourier transforms for complete regular k-point grids.
subroutine, public cell_to_k_grid_fft(values_rs, index_to_cell, xkp, nkp_grid, values_k, used_fft, deriv_direction, hmat, selected_kpoints)
Transform a batch of real matrices from image cells to every supplied k point.
Block-circulant inverse operators for replicated-cell OT preconditioners.
real(kind=dp), parameter, public lattice_fft_auto_max_cost_ratio
subroutine, public build_lattice_local_correction(operator_matrix, lattice_dims, local_cell_count, correction_indices, operator_columns, coarse_inverse, corrected_projection_error, info)
Build a balanced local correction for selected non-circulant cell blocks.
subroutine, public apply_lattice_inverse_dense(inverse_k, lattice_dims, matrix_in, matrix_out, used_fft, imaginary_residual)
Apply inverse Fourier blocks to a dense cell-major real matrix.
subroutine, public apply_lattice_state_inverse_dense(inverse_k, lattice_dims, matrix_in, matrix_out, used_fft, imaginary_residual, state_indices)
Apply one inverse Fourier operator to every column of a dense real matrix.
subroutine, public build_lattice_inverse(operator_matrix, lattice_dims, inverse_k, projection_error, used_fft, info)
Project a cell-major SPD operator onto lattice translations and invert its Fourier blocks.
real(kind=dp), parameter, public lattice_fft_auto_max_error
pure real(kind=dp) function, public lattice_local_storage_ratio(lattice_dims, operator_count, local_cell_count)
Estimate balanced-local-correction storage relative to one dense real operator.
pure real(kind=dp) function, public lattice_fft_storage_ratio(lattice_dims, operator_count)
Estimate reciprocal-block storage relative to one dense real operator.
pure logical function, public lattice_fft_auto_select(lattice_dims, block_size, rhs_count, operator_count, num_pe, projection_error, used_fft, info)
Conservative error-, work-, and storage-aware AUTO selector.
subroutine, public lattice_grid(lattice_dims, index_to_cell, xkp)
Enumerate a regular direct/reciprocal lattice with the first direction varying fastest.
integer, parameter, public lattice_fft_auto_min_cells
subroutine, public build_lattice_inverse_from_blocks(blocks, lattice_dims, inverse_k, used_fft, info)
Transform and invert already projected real-space lattice blocks.
pure real(kind=dp) function, public lattice_local_cost_ratio(lattice_dims, local_cell_count)
Estimate balanced-local-correction work relative to a dense inverse application.
pure real(kind=dp) function, public lattice_fft_cost_ratio(lattice_dims, block_size, rhs_count, num_pe)
Estimate lattice-FFT application work relative to a dense inverse application.
real(kind=dp), parameter, public lattice_fft_auto_max_storage_ratio