Coverage for gpaw/elph/supercell.py: 91%
174 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-20 00:19 +0000
1"""Module for electron-phonon supercell properties."""
3import numpy as np
4from typing import Tuple
6from ase import Atoms
7from ase.parallel import parprint
8from ase.units import Bohr
9from ase.utils.filecache import MultiFileJSONCache
11from gpaw.calculator import GPAW
12from gpaw.lcao.tightbinding import TightBinding
13from gpaw.typing import ArrayND
14from gpaw.utilities import unpack_hermitian
15from gpaw.utilities.tools import tri2full
17from .filter import fourier_filter
19sc_version = 1
20# v1: saves natom, supercell, g_sNNMM.shape and dtype
23class VersionError(Exception):
24 """Error raised for wrong cache versions."""
25 pass
28class Supercell:
29 """Class for supercell-related stuff."""
31 def __init__(self, atoms: Atoms, supercell_name: str = "supercell",
32 supercell: tuple = (1, 1, 1), indices=None) -> None:
33 """Initialize supercell class.
35 Parameters
36 ----------
37 atoms: Atoms
38 The atoms to work on. Primitive cell.
39 supercell_name: str
40 User specified name of the generated JSON cache.
41 Default is 'supercell'.
42 supercell: tuple
43 Size of supercell given by the number of repetitions (l, m, n) of
44 the small unit cell in each direction.
45 """
46 self.atoms = atoms
47 self.supercell_name = supercell_name
48 self.supercell = supercell
49 if indices is None:
50 self.indices = np.arange(len(atoms))
51 else:
52 self.indices = indices
54 def _calculate_supercell_entry(self, a, v, V1t_sG, dH1_asp, wfs,
55 dH_asp) -> ArrayND:
56 kpt_u = wfs.kpt_u
57 setups = wfs.setups
58 nao = setups.nao
59 bfs = wfs.basis_functions
60 dtype = wfs.dtype
61 nspins = wfs.nspins
63 # Array for different k-point components
64 g_sqMM = np.zeros((nspins, len(kpt_u) // nspins, nao, nao), dtype)
66 # 1) Gradient of effective potential
67 for kpt in kpt_u:
68 # Matrix elements
69 # Note: somehow this part does not work with gd-parallelisation
70 geff_MM = np.zeros((nao, nao), dtype)
71 bfs.calculate_potential_matrix(V1t_sG[kpt.s], geff_MM, q=kpt.q)
72 tri2full(geff_MM, "L")
73 # wfs.gd.comm.sum(geff_MM)
74 # print(world.rank, a, v, kpt.k, geff_MM)
75 g_sqMM[kpt.s, kpt.q] += geff_MM
76 # print(wfs.kd.comm.rank, wfs.gd.comm.rank, wfs.bd.comm.rank,
77 # "\n", geff_MM)
79 # 2) Gradient of non-local part (projectors)
80 P_aqMi = getattr(wfs, 'P_aqMi', None)
81 # 2a) dH^a part has contributions from all atoms
82 for kpt in kpt_u:
83 # Matrix elements
84 gp_MM = np.zeros((nao, nao), dtype)
85 for a_, dH1_sp in dH1_asp.items():
86 if a_ not in bfs.my_atom_indices:
87 continue
88 dH1_ii = unpack_hermitian(dH1_sp[kpt.s])
89 if P_aqMi is None:
90 P_Mi = kpt.P_aMi[a_]
91 else:
92 P_Mi = P_aqMi[a_][kpt.q]
93 gp_MM += P_Mi.conj() @ dH1_ii @ P_Mi.T
94 # wfs.gd.comm.sum(gp_MM)
95 g_sqMM[kpt.s, kpt.q] += gp_MM
97 # 2b) dP^a part has only contributions from the same atoms
98 # For the contribution from the derivative of the projectors
99 manytci = wfs.manytci
100 dPdR_aqvMi = manytci.P_aqMi(bfs.my_atom_indices, derivative=True)
101 dH_ii = unpack_hermitian(dH_asp[a][kpt.s])
102 for kpt in kpt_u:
103 gp_MM = np.zeros((nao, nao), dtype)
104 if a in bfs.my_atom_indices:
105 if P_aqMi is None:
106 P_Mi = kpt.P_aMi[a]
107 else:
108 P_Mi = P_aqMi[a][kpt.q]
109 dP_Mi = dPdR_aqvMi[a][kpt.q][v]
110 P1HP_MM = dP_Mi.conj() @ dH_ii @ P_Mi.T
111 # Matrix elements
112 gp_MM += P1HP_MM + P1HP_MM.T.conjugate()
113 # wfs.gd.comm.sum(gp_MM)
114 # print(world.rank, a,v, kpt.k, bfs.my_atom_indices, gp_MM)
115 g_sqMM[kpt.s, kpt.q] += gp_MM
117 return g_sqMM
119 def calculate_supercell_matrix(
120 self, calc: GPAW, fd_name: str = "elph", filter: str = None
121 ) -> None:
122 """Calculate matrix elements of the el-ph coupling in the LCAO basis.
124 This function calculates the matrix elements between LCAOs and local
125 atomic gradients of the effective potential. The matrix elements are
126 calculated for the supercell used to obtain finite-difference
127 approximations to the derivatives of the effective potential wrt to
128 atomic displacements.
130 The resulting g_xsNNMM is saved into a JSON cache.
132 Parameters
133 ----------
134 calc: GPAW
135 LCAO calculator for the calculation of the supercell matrix.
136 fd_name: str
137 User specified name of the finite difference JSON cache.
138 Default is 'elph'.
139 filter: str
140 Fourier filter atomic gradients of the effective potential. The
141 specified components (``normal`` or ``umklapp``) are removed
142 (default: None).
143 """
145 assert calc.wfs.mode == "lcao", "LCAO mode required."
146 assert not calc.symmetry.point_group, \
147 "Point group symmetry not supported"
149 # JSON cache
150 supercell_cache = MultiFileJSONCache(self.supercell_name)
152 # Supercell atoms
153 atoms_N = self.atoms * self.supercell
155 # Initialize calculator if required and extract useful quantities
156 if (not hasattr(calc.wfs, "S_qMM") or
157 not hasattr(calc.wfs.basis_functions, "M_a")):
158 calc.initialize(atoms_N)
159 calc.initialize_positions(atoms_N)
161 # Extract useful objects from the calculator
162 wfs = calc.wfs
163 gd = calc.wfs.gd
164 kd = calc.wfs.kd
165 bd = calc.wfs.bd
166 nao = wfs.setups.nao
167 nspins = wfs.nspins
168 # FIXME: Domain parallelisation broken
169 assert gd.comm.size == 1
170 # FIXME: Band parallelisation broken - M is band parallel
171 assert bd.comm.size == 1
173 # Calculate finite-difference gradients (in Hartree / Bohr)
174 V1t_xsG, dH1_xasp = self.calculate_gradient(fd_name, self.indices)
176 # Equilibrium atomic Hamiltonian matrix (projector coefficients)
177 fd_cache = MultiFileJSONCache(fd_name)
178 dH_asp = fd_cache["eq"]["dH_all_asp"]
180 # Check that the grid is the same as in the calculator
181 assert np.all(V1t_xsG.shape[-3:] == (gd.N_c + gd.pbc_c - 1)), \
182 "Mismatch in grids."
184 # Save basis information, after we checked the data is kosher
185 with supercell_cache.lock("basis") as handle:
186 if handle is not None:
187 basis_info = self.set_basis_info(calc)
188 handle.save(basis_info)
190 # Fourier filter the atomic gradients of the effective potential
191 if filter is not None:
192 for s in range(nspins):
193 fourier_filter(self.atoms, self.supercell, V1t_xsG[:, s],
194 components=filter)
196 if kd.gamma:
197 print("WARNING: Gamma-point calculation. \
198 Overlap with neighboring cell cannot be removed")
199 else:
200 # Bloch to real-space converter
201 tb = TightBinding(atoms_N, calc)
203 # Calculate < i k | grad H | j k >, i.e. matrix elements in LCAO basis
205 # Do each cartesian component separately
206 for i, a in enumerate(self.indices):
207 for v in range(3):
208 # Corresponding array index
209 xoutput = 3 * a + v
210 xinput = 3 * i + v
212 # If exist already, don't recompute
213 with supercell_cache.lock(str(xoutput)) as handle:
214 if handle is None:
215 continue
217 parprint("%s-gradient of atom %u" %
218 (["x", "y", "z"][v], a))
220 g_sqMM = self._calculate_supercell_entry(
221 a, v, V1t_xsG[xinput], dH1_xasp[xinput], wfs, dH_asp
222 )
224 # Extract R_c=(0, 0, 0) block by Fourier transforming
225 if kd.gamma or kd.N_c is None:
226 g_sMM = g_sqMM[:, 0]
227 else:
228 # Convert to array
229 g_sMM_tmp = []
230 for s in range(nspins):
231 g_MM = tb.bloch_to_real_space(g_sqMM[s],
232 R_c=(0, 0, 0))
233 g_sMM_tmp.append(g_MM[0]) # [0] because of above
234 g_sMM = np.array(g_sMM_tmp)
235 del g_sMM_tmp
237 # Reshape to global unit cell indices
238 N = np.prod(self.supercell)
239 # Number of basis function in the primitive cell
240 assert (nao % N) == 0, "Alarm ...!"
241 nao_cell = nao // N
242 g_sNMNM = g_sMM.reshape((nspins, N, nao_cell, N, nao_cell))
243 g_sNNMM = g_sNMNM.swapaxes(2, 3).copy()
244 handle.save(g_sNNMM)
245 if xinput == 0:
246 with supercell_cache.lock("info") as handle:
247 if handle is not None:
248 info = {
249 "sc_version": sc_version,
250 "natom": len(self.atoms),
251 "supercell": self.supercell,
252 "gshape": g_sNNMM.shape,
253 "gtype": g_sNNMM.dtype.name,
254 }
255 handle.save(info)
257 def set_basis_info(self, *args) -> dict:
258 """Store LCAO basis info for atoms in reference cell in attribute.
260 Parameters
261 ----------
262 args: tuple
263 If the LCAO calculator is not available (e.g. if the supercell is
264 loaded from file), the ``load_supercell_matrix`` member function
265 provides the required info as arguments.
267 """
268 assert len(args) in (1, 2)
269 if len(args) == 1:
270 calc = args[0]
271 setups = calc.wfs.setups
272 bfs = calc.wfs.basis_functions
273 nao_a = [setups[a].nao for a in range(len(self.atoms))]
274 M_a = [bfs.M_a[a] for a in range(len(self.atoms))]
275 else:
276 M_a = args[0]
277 nao_a = args[1]
278 return {"M_a": M_a, "nao_a": nao_a}
280 @classmethod
281 def calculate_gradient(cls, fd_name: str,
282 indices=None) -> Tuple[ArrayND, list]:
283 """Calculate gradient of effective potential and projector coefs.
285 This function loads the generated json files and calculates
286 finite-difference derivatives.
288 Parameters
289 ----------
290 fd_name: str
291 name of finite difference JSON cache
292 """
293 cache = MultiFileJSONCache(fd_name)
294 if "dr_version" not in cache["info"]:
295 print("Cache created with old version. Use electronphonon.py")
296 raise VersionError
297 natom = cache["info"]["natom"]
298 delta = cache["info"]["delta"]
300 # Array and dict for finite difference derivatives
301 V1t_xsG = []
302 dH1_xasp = []
304 if indices is None:
305 indices = np.arange(natom)
307 x = 0
308 for a in indices:
309 for v in "xyz":
310 name = "%d%s" % (a, v)
311 # Potential and atomic density matrix for atomic displacement
312 Vtm_sG = cache[name + "-"]["Vt_sG"]
313 dHm_asp = cache[name + "-"]["dH_all_asp"]
314 Vtp_sG = cache[name + "+"]["Vt_sG"]
315 dHp_asp = cache[name + "+"]["dH_all_asp"]
317 # FD derivatives in Hartree / Bohr
318 V1t_sG = (Vtp_sG - Vtm_sG) / (2 * delta / Bohr)
319 V1t_xsG.append(V1t_sG)
321 dH1_asp = {}
322 for atom in dHm_asp.keys():
323 dH1_asp[atom] = (dHp_asp[atom] - dHm_asp[atom]) / (
324 2 * delta / Bohr
325 )
326 dH1_xasp.append(dH1_asp)
327 x += 1
328 return np.array(V1t_xsG), dH1_xasp
330 @classmethod
331 def load_supercell_matrix(cls, name: str = "supercell"
332 ) -> Tuple[ArrayND, dict]:
333 """Load supercell matrix from cache.
335 Parameters
336 ----------
337 name: str
338 User specified name of the cache.
339 """
340 # TODO: load by indices?
341 supercell_cache = MultiFileJSONCache(name)
342 if "sc_version" not in supercell_cache["info"]:
343 print("Cache created with old version. Use electronphonon.py")
344 raise VersionError
345 shape = supercell_cache["info"]["gshape"]
346 dtype = supercell_cache["info"]["gtype"]
347 natom = supercell_cache["info"]["natom"]
348 nx = natom * 3
349 g_xsNNMM = np.empty([nx, ] + list(shape), dtype=dtype)
350 for x in range(nx):
351 g_xsNNMM[x] = supercell_cache[str(x)]
352 basis_info = supercell_cache["basis"]
353 return g_xsNNMM, basis_info