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 Gauge-covariant selection of fixed trial orbitals for Wannier90.
10 : !> \author Thomas D. Kuehne, tkuehne@cp2k.org
11 : ! **************************************************************************************************
12 : MODULE wannier90_projections
13 : #if defined (__HAS_IEEE_EXCEPTIONS)
14 : USE ieee_exceptions, ONLY: IEEE_ALL, ieee_get_halting_mode, ieee_set_halting_mode
15 : #endif
16 : USE kinds, ONLY: dp
17 : USE message_passing, ONLY: mp_para_env_type
18 : #include "./base/base_uses.f90"
19 :
20 : IMPLICIT NONE
21 : PRIVATE
22 :
23 : CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'wannier90_projections'
24 : PUBLIC :: wannier90_select_projections, wannier90_projection_rank, wannier90_hybrid_candidates, &
25 : wannier90_window_mask
26 :
27 : CONTAINS
28 :
29 : ! **************************************************************************************************
30 : !> \brief Restrict fixed trial selection to the bands admitted by disentanglement windows.
31 : !> \param eigenvalues Exported eigenvalues in the same units as the windows.
32 : !> \param num_wann Number of target Wannier functions.
33 : !> \param windows Outer minimum/maximum and frozen minimum/maximum energies.
34 : !> \param window_set Whether each bound is explicitly set.
35 : !> \param allowed Bands used to select the physical trial orbitals.
36 : ! **************************************************************************************************
37 28 : SUBROUTINE wannier90_window_mask(eigenvalues, num_wann, windows, window_set, allowed)
38 : REAL(KIND=dp), INTENT(IN) :: eigenvalues(:)
39 : INTEGER, INTENT(IN) :: num_wann
40 : REAL(KIND=dp), INTENT(IN) :: windows(4)
41 : LOGICAL, INTENT(IN) :: window_set(4)
42 : LOGICAL, INTENT(OUT) :: allowed(:)
43 :
44 56 : LOGICAL :: frozen(SIZE(eigenvalues))
45 :
46 28 : CPASSERT(SIZE(allowed) == SIZE(eigenvalues))
47 140 : allowed = .TRUE.
48 140 : IF (window_set(1)) allowed = allowed .AND. eigenvalues >= windows(1)
49 140 : IF (window_set(2)) allowed = allowed .AND. eigenvalues <= windows(2)
50 140 : IF (COUNT(allowed) < num_wann) THEN
51 0 : CPABORT('Wannier90 outer window contains too few bands.')
52 : END IF
53 28 : IF (window_set(4)) THEN
54 130 : frozen = eigenvalues <= windows(4)
55 26 : IF (window_set(3)) frozen = frozen .AND. eigenvalues >= windows(3)
56 130 : IF (COUNT(frozen) > num_wann) THEN
57 0 : CPABORT('Wannier90 frozen window contains too many bands.')
58 : END IF
59 130 : IF (ANY(frozen .AND. .NOT. allowed)) THEN
60 0 : CPABORT('Wannier90 frozen states lie outside the outer window.')
61 : END IF
62 260 : IF (COUNT(frozen) == num_wann) allowed = frozen
63 : END IF
64 28 : END SUBROUTINE wannier90_window_mask
65 :
66 : ! **************************************************************************************************
67 : !> \brief Form tetrahedral s/p trials from the full AO bank before selecting projections.
68 : !> \param ao Band/AO/k-point projections on the source rank, empty elsewhere.
69 : !> \param groups AO indices of one s and three p functions for each radial shell pair.
70 : !> \param candidates Hybrid trials and AOs not belonging to any hybrid group.
71 : ! **************************************************************************************************
72 12 : SUBROUTINE wannier90_hybrid_candidates(ao, groups, candidates)
73 : COMPLEX(KIND=dp), INTENT(IN) :: ao(:, :, :)
74 : INTEGER, INTENT(IN) :: groups(:, :)
75 : COMPLEX(KIND=dp), ALLOCATABLE, INTENT(OUT) :: candidates(:, :, :)
76 :
77 : REAL(KIND=dp), PARAMETER :: hybrid(4, 4) = 0.5_dp*RESHAPE([1.0_dp, 1.0_dp, 1.0_dp, 1.0_dp, &
78 : 1.0_dp, 1.0_dp, -1.0_dp, -1.0_dp, 1.0_dp, -1.0_dp, 1.0_dp, -1.0_dp, 1.0_dp, -1.0_dp, &
79 : -1.0_dp, 1.0_dp], [4, 4])
80 :
81 : INTEGER :: iao, ic, ig, ik, nao, nb, nk
82 12 : LOGICAL, ALLOCATABLE :: grouped(:)
83 :
84 12 : nb = SIZE(ao, 1)
85 12 : nao = SIZE(ao, 2)
86 12 : nk = SIZE(ao, 3)
87 12 : IF (nb == 0) THEN
88 6 : ALLOCATE (candidates(0, 0, 0))
89 : RETURN
90 : END IF
91 6 : CPASSERT(SIZE(groups, 1) == 4)
92 132 : CPASSERT(ALL(groups > 0) .AND. ALL(groups <= nao))
93 18 : ALLOCATE (grouped(nao))
94 6 : grouped = .FALSE.
95 18 : DO ig = 1, SIZE(groups, 2)
96 60 : DO ic = 1, 4
97 252 : CPASSERT(COUNT(groups(:, ig) == groups(ic, ig)) == 1)
98 : END DO
99 66 : grouped(groups(:, ig)) = .TRUE.
100 : END DO
101 70 : ALLOCATE (candidates(nb, COUNT(.NOT. grouped) + 4*SIZE(groups, 2), nk))
102 6 : ic = 0
103 46 : DO iao = 1, nao
104 40 : IF (grouped(iao)) CYCLE
105 7 : ic = ic + 1
106 55 : candidates(:, ic, :) = ao(:, iao, :)
107 : END DO
108 : ! Include every radial s/p pair. Selecting only already-pivoted AO columns
109 : ! would exclude physical trials that the original QR happened to discard.
110 18 : DO ig = 1, SIZE(groups, 2)
111 84 : DO ik = 1, nk
112 6996 : candidates(:, ic + 1:ic + 4, ik) = MATMUL(ao(:, groups(:, ig), ik), hybrid)
113 : END DO
114 18 : ic = ic + 4
115 : END DO
116 6 : DEALLOCATE (grouped)
117 : END SUBROUTINE wannier90_hybrid_candidates
118 :
119 : ! **************************************************************************************************
120 : !> \brief Select the same trial orbitals at every k-point using a mesh-wide pivoted QR.
121 : !> \param candidates Band/candidate/k-point projections, stored only on the source rank.
122 : !> \param para_env Global parallel environment.
123 : !> \param amn Selected projection matrices, allocated with the target dimensions on all ranks.
124 : !> \param pivots Selected candidate indices, allocated on all ranks.
125 : !> \param min_svalue Smallest singular value of the selected projections over the mesh.
126 : !> \param anchor_kpoint Optional fixed QR anchor; zero selects the stacked mesh. Without an
127 : !> explicit anchor, try individual k-points if the stacked selection loses rank.
128 : ! **************************************************************************************************
129 32 : SUBROUTINE wannier90_select_projections(candidates, para_env, amn, pivots, min_svalue, anchor_kpoint)
130 : COMPLEX(KIND=dp), INTENT(IN) :: candidates(:, :, :)
131 : TYPE(mp_para_env_type), POINTER :: para_env
132 : COMPLEX(KIND=dp), INTENT(OUT) :: amn(:, :, :)
133 : INTEGER, INTENT(OUT) :: pivots(:)
134 : REAL(KIND=dp), INTENT(OUT) :: min_svalue
135 : INTEGER, INTENT(IN), OPTIONAL :: anchor_kpoint
136 :
137 32 : COMPLEX(KIND=dp), ALLOCATABLE :: qr(:, :)
138 : INTEGER :: anchor, best_defect, defect, &
139 : first_anchor, ic, ik, iw, last_anchor, &
140 : nb, nc, nk, nr, nw, previous_defect, &
141 : saved_pivot
142 32 : INTEGER, ALLOCATABLE :: best_pivots(:)
143 :
144 32 : nb = SIZE(amn, 1)
145 32 : nw = SIZE(amn, 2)
146 32 : nk = SIZE(amn, 3)
147 32 : CPASSERT(nb >= nw .AND. nw > 0 .AND. SIZE(pivots) == nw)
148 32 : first_anchor = 0
149 32 : last_anchor = nk
150 32 : IF (PRESENT(anchor_kpoint)) THEN
151 10 : CPASSERT(anchor_kpoint >= 0 .AND. anchor_kpoint <= nk)
152 : first_anchor = anchor_kpoint
153 : last_anchor = anchor_kpoint
154 : END IF
155 32 : min_svalue = 0.0_dp
156 32 : IF (para_env%is_source()) THEN
157 16 : nc = SIZE(candidates, 2)
158 16 : CPASSERT(SIZE(candidates, 1) == nb .AND. SIZE(candidates, 3) == nk)
159 16 : CPASSERT(nc >= nw)
160 48 : ALLOCATE (best_pivots(nw))
161 32 : best_defect = HUGE(0)
162 32 : DO anchor = first_anchor, last_anchor
163 25 : nr = nb*nk
164 25 : IF (anchor > 0) nr = nb
165 100 : ALLOCATE (qr(nr, nc))
166 25 : IF (anchor == 0) THEN
167 71 : DO ik = 1, nk
168 1507 : qr((ik - 1)*nb + 1:ik*nb, :) = candidates(:, :, ik)
169 : END DO
170 : ELSE
171 138 : qr(:, :) = candidates(:, :, anchor)
172 : END IF
173 : ! A unitary change of the MO gauge acts on rows, not on the trial orbitals.
174 : ! The selected columns therefore represent one physical trial set on the full mesh.
175 25 : CALL select_qr_pivots(qr, pivots)
176 113 : DO ik = 1, nk
177 609 : amn(:, :, ik) = candidates(:, pivots, ik)
178 : END DO
179 25 : DEALLOCATE (qr)
180 :
181 25 : min_svalue = wannier90_projection_rank(amn, defect)
182 25 : IF (defect < best_defect) THEN
183 55 : best_defect = defect
184 55 : best_pivots(:) = pivots(:)
185 : END IF
186 32 : IF (min_svalue >= 1.0e-10_dp) EXIT
187 : END DO
188 16 : IF (best_defect > 0 .AND. .NOT. PRESENT(anchor_kpoint)) THEN
189 : ! QR optimizes one anchor, not rank everywhere. Each accepted exchange
190 : ! strictly reduces the number of missing trial directions across the mesh.
191 6 : pivots = best_pivots
192 : DO
193 2 : previous_defect = best_defect
194 4 : repair: DO iw = 1, nw
195 3 : saved_pivot = pivots(iw)
196 12 : DO ic = 1, nc
197 19 : IF (ANY(pivots == ic)) CYCLE
198 4 : pivots(iw) = ic
199 14 : DO ik = 1, nk
200 74 : amn(:, :, ik) = candidates(:, pivots, ik)
201 : END DO
202 4 : min_svalue = wannier90_projection_rank(amn, defect)
203 4 : IF (defect < best_defect) THEN
204 : best_defect = defect
205 : EXIT repair
206 : END IF
207 10 : pivots(iw) = saved_pivot
208 : END DO
209 : END DO repair
210 2 : IF (best_defect == 0 .OR. best_defect == previous_defect) EXIT
211 : END DO
212 7 : DO ik = 1, nk
213 37 : amn(:, :, ik) = candidates(:, pivots, ik)
214 : END DO
215 2 : min_svalue = wannier90_projection_rank(amn)
216 : END IF
217 16 : DEALLOCATE (best_pivots)
218 : END IF
219 1780 : CALL para_env%bcast(amn)
220 152 : CALL para_env%bcast(pivots)
221 32 : CALL para_env%bcast(min_svalue)
222 :
223 32 : END SUBROUTINE wannier90_select_projections
224 :
225 : ! **************************************************************************************************
226 : !> \brief Pivoted Householder QR with a reproducible choice among numerically tied columns.
227 : !> \param qr Stacked or single-k-point candidate matrix, overwritten by its factorization.
228 : !> \param pivots Original indices of the selected columns.
229 : ! **************************************************************************************************
230 25 : SUBROUTINE select_qr_pivots(qr, pivots)
231 : COMPLEX(KIND=dp), INTENT(INOUT) :: qr(:, :)
232 : INTEGER, INTENT(OUT) :: pivots(:)
233 :
234 : REAL(KIND=dp), PARAMETER :: tie_tolerance = 1.0e-12_dp
235 :
236 : COMPLEX(KIND=dp) :: tau(1)
237 25 : COMPLEX(KIND=dp), ALLOCATABLE :: column(:), work(:)
238 : INTEGER :: ic, info, ip, iw, nc, nr, original
239 25 : INTEGER, ALLOCATABLE :: order(:)
240 : REAL(KIND=dp) :: largest
241 25 : REAL(KIND=dp), ALLOCATABLE :: norms(:)
242 :
243 25 : nr = SIZE(qr, 1)
244 25 : nc = SIZE(qr, 2)
245 25 : CPASSERT(SIZE(pivots) <= MIN(nr, nc))
246 225 : ALLOCATE (column(nr), work(MAX(1, nc)), order(nc), norms(nc))
247 259 : order(:) = [(ic, ic=1, nc)]
248 73 : DO iw = 1, SIZE(pivots)
249 248 : DO ic = iw, nc
250 1937 : norms(ic) = SUM(ABS(qr(iw:nr, ic))**2)
251 : END DO
252 248 : largest = MAXVAL(norms(iw:nc))
253 : original = HUGE(0)
254 : ip = iw
255 248 : DO ic = iw, nc
256 200 : IF (norms(ic) < largest*(1.0_dp - tie_tolerance)) CYCLE
257 113 : IF (order(ic) < original) THEN
258 48 : ip = ic
259 48 : original = order(ic)
260 : END IF
261 : END DO
262 312 : column(:) = qr(:, iw)
263 312 : qr(:, iw) = qr(:, ip)
264 312 : qr(:, ip) = column
265 48 : order(ip) = order(iw)
266 48 : order(iw) = original
267 48 : pivots(iw) = original
268 : ! LAPACK constructs and applies the reflector; only pivot tie-breaking differs.
269 48 : CALL zgeqrf(nr - iw + 1, 1, qr(iw, iw), nr, tau, work, SIZE(work), info)
270 48 : CPASSERT(info == 0)
271 73 : IF (iw < nc) THEN
272 : CALL zunmqr('L', 'C', nr - iw + 1, nc - iw, 1, qr(iw, iw), nr, tau, &
273 48 : qr(iw, iw + 1), nr, work, SIZE(work), info)
274 48 : CPASSERT(info == 0)
275 : END IF
276 : END DO
277 25 : DEALLOCATE (column, work, order, norms)
278 25 : END SUBROUTINE select_qr_pivots
279 :
280 : ! **************************************************************************************************
281 : !> \brief Validate the trial rank independently at every k-point.
282 : !> \param amn Band/trial/k-point projections, possibly rectangular.
283 : !> \param rank_defect Total number of missing trial directions summed over the mesh.
284 : !> \return Smallest singular value over all k-points.
285 : ! **************************************************************************************************
286 33 : FUNCTION wannier90_projection_rank(amn, rank_defect) RESULT(min_svalue)
287 : COMPLEX(KIND=dp), INTENT(IN) :: amn(:, :, :)
288 : INTEGER, INTENT(OUT), OPTIONAL :: rank_defect
289 : REAL(KIND=dp) :: min_svalue
290 :
291 : COMPLEX(KIND=dp) :: query(1), unused_u(1, 1), unused_v(1, 1)
292 33 : COMPLEX(KIND=dp), ALLOCATABLE :: matrix(:, :), work(:)
293 : INTEGER :: ik, info, lwork, nb, nw
294 33 : REAL(KIND=dp), ALLOCATABLE :: rwork(:), svalues(:)
295 : #if defined (__HAS_IEEE_EXCEPTIONS)
296 : LOGICAL :: halt(SIZE(IEEE_ALL))
297 : #endif
298 :
299 33 : nb = SIZE(amn, 1)
300 33 : nw = SIZE(amn, 2)
301 33 : CPASSERT(nb >= nw .AND. nw > 0 .AND. SIZE(amn, 3) > 0)
302 264 : ALLOCATE (matrix(nb, nw), svalues(nw), rwork(5*nw))
303 245 : matrix(:, :) = amn(:, :, 1)
304 33 : CALL zgesvd('N', 'N', nb, nw, matrix, nb, svalues, unused_u, 1, unused_v, 1, query, -1, rwork, info)
305 33 : CPASSERT(info == 0)
306 33 : lwork = MAX(1, INT(REAL(query(1), KIND=dp)))
307 99 : ALLOCATE (work(lwork))
308 33 : min_svalue = HUGE(0.0_dp)
309 33 : IF (PRESENT(rank_defect)) rank_defect = 0
310 140 : DO ik = 1, SIZE(amn, 3)
311 753 : matrix(:, :) = amn(:, :, ik)
312 : ! The values-only SVD uses DLASQ, whose IEEECK probe deliberately raises exceptions.
313 : #if defined (__HAS_IEEE_EXCEPTIONS)
314 : CALL ieee_get_halting_mode(IEEE_ALL, halt)
315 : CALL ieee_set_halting_mode(IEEE_ALL, .FALSE.)
316 : #endif
317 107 : CALL zgesvd('N', 'N', nb, nw, matrix, nb, svalues, unused_u, 1, unused_v, 1, work, lwork, rwork, info)
318 : #if defined (__HAS_IEEE_EXCEPTIONS)
319 : CALL ieee_set_halting_mode(IEEE_ALL, halt)
320 : #endif
321 107 : CPASSERT(info == 0)
322 309 : min_svalue = MIN(min_svalue, MINVAL(svalues))
323 418 : IF (PRESENT(rank_defect)) rank_defect = rank_defect + COUNT(svalues < 1.0e-10_dp)
324 : END DO
325 33 : DEALLOCATE (matrix, work, rwork, svalues)
326 33 : END FUNCTION wannier90_projection_rank
327 :
328 : END MODULE wannier90_projections
|
