Coverage for gpaw/ibz2bz.py: 98%
116 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-14 00:18 +0000
1from collections.abc import Sequence
2import numpy as np
3from gpaw.utilities.blas import gemmdot
4from gpaw.grid_descriptor import GridDescriptor
5from gpaw.projections import Projections
8class IBZ2BZMaps(Sequence):
9 """Sequence of data maps from k-points in the IBZ to the full BZ.
11 The sequence is indexed by the BZ index K."""
13 def __init__(self, kd, spos_ac, R_asii, N_c):
14 """Construct the IBZ2BZMapper.
16 Parameters
17 ----------
18 kd : KPointDescriptor
19 spos_ac : np.array
20 Scaled atomic positions
21 R_asii : list
22 Rotations of the PAW projections under symmetry transformations
23 N_c : iterable of 3 ints
24 Number of grid points along axes
25 """
26 self.kd = kd
27 self.spos_ac = spos_ac
28 self.R_asii = R_asii
30 # Scaled coordinates of the real-space grid
31 self.r_cR = np.array(np.meshgrid(*[np.linspace(0, 1, N, endpoint=False)
32 for N in N_c], indexing='ij'))
34 @classmethod
35 def from_calculator(cls, calc):
36 R_asii = calc.setups.atomrotations.get_R_asii()
37 return cls(calc.wfs.kd, calc.spos_ac, R_asii, calc.wfs.gd.N_c)
39 def __len__(self):
40 return len(self.kd.bzk_kc)
42 def __getitem__(self, K):
43 kd = self.kd
44 return IBZ2BZMap(kd.ibzk_kc[kd.bz2ibz_k[K]],
45 *self.get_symmetry_transformations(kd.sym_k[K]),
46 kd.time_reversal_k[K],
47 self.spos_ac,
48 self.kd.bzk_kc[K],
49 self.r_cR)
51 def get_symmetry_transformations(self, s):
52 return (self.get_rotation_matrix(s),
53 self.get_atomic_permutations(s),
54 self.get_projections_rotations(s))
56 def get_rotation_matrix(self, s):
57 """Coordinate rotation matrix, mapping IBZ -> K."""
58 U_cc = self.kd.symmetry.op_scc[s]
59 return U_cc
61 def get_atomic_permutations(self, s):
62 """Permutations of atomic indices in the IBZ -> K map."""
63 b_a = self.kd.symmetry.a_sa[s] # Atom a is mapped onto atom b
64 return b_a
66 def get_projections_rotations(self, s):
67 """Rotations of the PAW projections for the IBZ -> K mapping."""
68 R_aii = [R_sii[s] for R_sii in self.R_asii]
69 return R_aii
72class IBZ2BZMap:
73 """Functionality to map orbitals from the IBZ to a specific k-point K."""
75 def __init__(self, ik_c, U_cc, b_a, R_aii,
76 time_reversal, spos_ac, k_c, r_cR):
77 """Construct the IBZ2BZMap."""
78 self.ik_c = ik_c
79 self.k_c = k_c # k_c in 1:st BZ that IBZ-kpoint is mapped to
80 self.r_cR = r_cR
81 self.U_cc = U_cc
82 self.b_a = b_a
83 self.R_aii = R_aii
84 self.time_reversal = time_reversal
86 self.spos_ac = spos_ac
88 def map_kpoint(self):
89 """Get the relative k-point coordinates after the IBZ -> K mapping.
91 NB: The mapped k-point can lie outside the BZ, but will always be
92 related to kd.bzk_kc[K] by a reciprocal lattice vector.
93 """
94 # Apply symmetry operations to the irreducible k-point
95 sign = 1 - 2 * self.time_reversal
96 k_c = sign * self.U_cc @ self.ik_c
98 return k_c
100 def map_pseudo_wave(self, ut_R):
101 """Map the periodic part of the pseudo wave from the IBZ -> K.
103 For the symmetry operator U, which maps K = U ik, where ik is the
104 irreducible BZ k-point and K does not necessarily lie within the 1BZ,
106 psit_K(r) = psit_ik(U^T r)
108 The mapping takes place on the coarse real-space grid.
109 """
110 # Apply symmetry operations to the periodic part of the pseudo wave
111 if not (self.U_cc == np.eye(3)).all():
112 N_c = ut_R.shape
113 i_cr = np.dot(self.U_cc.T, np.indices(N_c).reshape((3, -1)))
114 i = np.ravel_multi_index(i_cr, N_c, 'wrap')
115 utout_R = ut_R.ravel()[i].reshape(N_c)
116 else:
117 utout_R = ut_R.copy()
118 if self.time_reversal:
119 utout_R = utout_R.conj()
121 assert utout_R is not ut_R, \
122 "We don't want the output array to point back at the input array"
124 return utout_R
126 def map_pseudo_wave_to_BZ(self, ut_R):
127 """Map the periodic part of wave function from IBZ -> k in BZ.
129 Parameters
130 ----------
131 ut_R: np.array
132 Periodic part of pseudo wf at IBZ k-point
133 """
134 # Map the pseudo wave and K-point using symmetry operations
135 utout_R = self.map_pseudo_wave(ut_R)
136 kpt_shift_c = self.map_kpoint() - self.k_c
138 # Check if the resulting K-point already resides inside the BZ
139 if np.allclose(kpt_shift_c, 0.0):
140 return utout_R
142 # Check that the mapped k-point differ from the BZ K-point by
143 # a reciprocal lattice vector G
144 assert np.allclose((kpt_shift_c - np.round(kpt_shift_c)), 0.0)
146 # Add a phase e^iG.r to the periodic part of the wf
147 return utout_R * np.exp(2j * np.pi * gemmdot(kpt_shift_c, self.r_cR))
149 def map_projections(self, projections):
150 """Perform IBZ -> K mapping of the PAW projections.
152 The projections of atom "a" are mapped onto an atom related to atom
153 "b" by a lattice vector:
155 r_b = U^T r_a (or equivalently: r_b^T = r_a^T U)
157 This means that when mapping
159 psi_K(r) = psi_ik(U^T r),
161 we need to generate the projections at atom a for k-point K based on
162 the projections at atom b for k-point ik.
163 """
164 mapped_projections = projections.new()
165 if mapped_projections.atom_partition.comm.rank == 0:
166 for a, (b, U_ii) in enumerate(zip(self.b_a, self.U_aii)):
167 # Map projections
168 Pin_ni = projections[b]
169 Pout_ni = Pin_ni @ U_ii
170 if self.time_reversal:
171 Pout_ni = np.conj(Pout_ni)
173 # Store output projections
174 I1, I2 = mapped_projections.map[a]
175 mapped_projections.array[..., I1:I2] = Pout_ni
176 else:
177 assert len(mapped_projections.indices) == 0
179 return mapped_projections
181 @property
182 def U_aii(self):
183 """Phase corrected rotation matrices for the PAW projections."""
184 U_aii = []
185 for a, R_ii in enumerate(self.R_aii):
186 # The symmetry transformation maps atom "a" to a position which is
187 # related to atom "b" by a lattice vector (but which does not
188 # necessarily lie within the unit cell)
189 b = self.b_a[a]
190 cell_shift_c = self.spos_ac[a] @ self.U_cc - self.spos_ac[b]
191 assert np.allclose(cell_shift_c.round(), cell_shift_c, atol=1e-6)
192 # This means, that when we want to extract the projections at K for
193 # atom a according to psi_K(r_a) = psi_ik(U^T r_a), we need the
194 # projections at U^T r_a for k-point ik. Since we only have the
195 # projections within the unit cell we need to multiply them with a
196 # phase factor according to the cell shift.
197 phase_factor = np.exp(2j * np.pi * self.ik_c @ cell_shift_c)
198 U_ii = R_ii.T * phase_factor
199 U_aii.append(U_ii)
201 return U_aii
204def get_overlap(bands,
205 gd: GridDescriptor,
206 ut1_nR, ut2_nR,
207 proj1: Projections,
208 proj2: Projections,
209 dO_aii):
210 """Computes the overlap between all-electron wave functions.
212 Parameters
213 ----------
214 bands: list of integers
215 Band indices to calculate overlap for
216 dO_aii: dict
217 Overlap coefficients for the partial waves in the PAW setups.
218 Can be extracted via get_overlap_coefficients().
219 """
220 return _get_overlap(bands, gd.dv, ut1_nR, ut2_nR, proj1, proj2, dO_aii)
223def _get_overlap(bands, dv, ut1_nR, ut2_nR, proj1, proj2, dO_aii):
224 nbands = len(bands)
225 ut1_nR = np.reshape(ut1_nR[bands], (nbands, -1))
226 ut2_nR = np.reshape(ut2_nR[bands], (nbands, -1))
227 M_nn = (ut1_nR.conj() @ ut2_nR.T) * dv
229 for a in proj1.map:
230 dO_ii = dO_aii[a]
231 if proj1.collinear:
232 P1_ni = proj1[a][bands]
233 P2_ni = proj2[a][bands]
234 M_nn += P1_ni.conj() @ (dO_ii) @ (P2_ni.T)
235 else:
236 # We have an additional spinor index s to sum over
237 P1_nsi = proj1[a][bands]
238 P2_nsi = proj2[a][bands]
239 assert P1_nsi.shape[1] == 2
240 for s in range(2):
241 M_nn += P1_nsi[:, s].conj() @ (dO_ii) @ (P2_nsi[:, s].T)
243 return M_nn
246def get_overlap_coefficients(wfs):
247 dO_aii = {}
248 for a in wfs.kpt_u[0].projections.map:
249 dO_aii[a] = wfs.setups[a].dO_ii
250 return dO_aii
253def get_phase_shifted_overlap_coefficients(dO_aii, spos_ac, K_c):
254 """Apply phase shift to overlap coefficients.
256 Applies a Bloch phase factor e^(iK.R_a) to the overlap coefficients dO_aii
257 at each (scaled) atomic position spos_ac (R_a) for the input (scaled) wave
258 vector K_c.
259 """
260 new_dO_aii = {}
261 for a, dO_ii in dO_aii.items():
262 phase_factor = np.exp(2j * np.pi * K_c @ spos_ac[a])
263 new_dO_aii[a] = phase_factor * dO_ii
264 return new_dO_aii