Coverage for gpaw/elph/gmatrix.py: 84%
164 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-14 00:18 +0000
1r"""Module for calculating electron-phonon matrix.
3Electron-phonon interaction::
5 __
6 \ l + +
7 H = ) g c c ( a + a ),
8 el-ph /_ ij i j l l
9 l,ij
11where the electron phonon coupling is given by::
13 ______
14 l / hbar ___
15 g = /------- < i | \ / V * e | j > .
16 ij \/ 2 M w 'u eff l
17 l
19Here, l denotes the vibrational mode, w_l and e_l is the frequency and
20mass-scaled polarization vector, respectively, M is an effective mass, i, j are
21electronic state indices and nabla_u denotes the gradient wrt atomic
22displacements. The implementation supports calculations of the el-ph coupling
23in both finite and periodic systems, i.e. expressed in a basis of molecular
24orbitals or Bloch states.
25"""
26from typing import Optional
28import ase.units as units
29import numpy as np
30from ase import Atoms
31from ase.phonons import Phonons
32from ase.utils.filecache import MultiFileJSONCache
33from ase.utils.timing import Timer, timer
35from gpaw.calculator import GPAW
36from gpaw.mpi import world
37from gpaw.typing import ArrayND
39from .supercell import Supercell
41OPTIMIZE = "optimal"
44class ElectronPhononMatrix:
45 """Class for containing the electron-phonon matrix"""
47 def __init__(self, atoms: Atoms, supercell_cache: str, phonon,
48 load_sc_as_needed: bool = True, indices=None) -> None:
49 """Initialize with base class args and kwargs.
51 Parameters
52 ----------
53 atoms: Atoms
54 Primitive cell object
55 supercell_cache: str
56 Name of JSON cache containing supercell matrix
57 phonon: str, dict, :class:`~ase.phonons.Phonons`
58 Can be either name of phonon cache generated with
59 electron-phonon DisplacementRunner or dictonary
60 of arguments used in Phonons run or Phonons object.
61 load_sc_as_needed: bool
62 Load supercell matrix elements only as needed.
63 Greatly reduces memory requirement for large systems,
64 but introduces huge filesystem overhead
65 indices: list
66 List of atoms (indices) to use. Default: Use all.
67 """
68 if not load_sc_as_needed:
69 assert indices is None, "Use 'load_sc_as_needed' with 'indices'"
71 self.timer = Timer()
73 self.atoms = atoms
74 if indices is None:
75 indices = np.arange(len(atoms))
76 if isinstance(indices, np.ndarray):
77 self.indices = indices.tolist()
79 self._set_supercell_cache(supercell_cache, load_sc_as_needed)
81 self.timer.start("Read phonons")
82 self._set_phonon_cache(phonon, atoms)
84 if set(self.phonon.indices) != set(self.indices):
85 self.phonon.set_atoms(self.indices)
86 self.phonon.D_N = None
88 self._read_phonon_cache()
89 self.timer.stop("Read phonons")
91 def _set_supercell_cache(self, supercell_cache: str,
92 load_sc_as_needed: bool):
93 self.supercell_cache = MultiFileJSONCache(supercell_cache)
94 self.R_cN = self._get_lattice_vectors()
96 if load_sc_as_needed:
97 self._yield_g_NNMM = self._yield_g_NNMM_as_needed
98 self.g_xNNMM = None
99 else:
100 self.g_xsNNMM, _ = Supercell.load_supercell_matrix(supercell_cache)
101 self._yield_g_NNMM = self._yield_g_NNMM_from_var
103 def _set_phonon_cache(self, phonon, atoms):
104 if isinstance(phonon, Phonons):
105 self.phonon = phonon
106 elif isinstance(phonon, str):
107 info = MultiFileJSONCache(phonon)["info"]
108 assert "dr_version" in info, "use valid cache created by elph"
109 # our version of phonons
110 self.phonon = Phonons(atoms, supercell=info["supercell"],
111 name=phonon, delta=info["delta"],
112 center_refcell=True)
113 elif isinstance(phonon, dict):
114 # this would need to be updated if Phonon defaults change
115 supercell = phonon.get("supercell", (1, 1, 1))
116 name = phonon.get("name", "phonon")
117 delta = phonon.get("delta", 0.01)
118 center_refcell = phonon.get("center_refcell", False)
119 self.phonon = Phonons(atoms, supercell=supercell, name=name,
120 delta=delta, center_refcell=center_refcell)
121 else:
122 raise TypeError
124 def _read_phonon_cache(self):
125 if self.phonon.D_N is None:
126 self.phonon.read(symmetrize=10)
128 def _yield_g_NNMM_as_needed(self, x, s):
129 return self.supercell_cache[str(x)][s]
131 def _yield_g_NNMM_from_var(self, x, s):
132 return self.g_xsNNMM[x, s]
134 def _get_lattice_vectors(self):
135 """Recover lattice vectors of elph calculation"""
136 supercell = self.supercell_cache["info"]["supercell"]
137 ph = Phonons(self.atoms, supercell=supercell, center_refcell=True)
138 return ph.compute_lattice_vectors()
140 @classmethod
141 def _gather_all_wfc(cls, wfs, s):
142 """Return complete wave function on rank 0"""
143 c_knM = np.zeros((wfs.kd.nbzkpts, wfs.bd.nbands, wfs.setups.nao),
144 dtype=complex)
145 for k in range(wfs.kd.nbzkpts):
146 for n in range(wfs.bd.nbands):
147 c_knM[k, n] = wfs.get_wave_function_array(n, k, s, False)
148 return c_knM
150 @timer("Bloch matrix q k")
151 def _bloch_matrix(self, var1: ArrayND, C2_nM: ArrayND,
152 k_c: ArrayND, q_c: ArrayND,
153 prefactor: bool, s: Optional[int] = None) -> ArrayND:
154 """Calculates elph matrix entry for a given k and q.
156 The first argument must either be
157 C1_nM, the ket wavefunction at k_c
158 OR
159 or a preprocessed g_xNMn, where the ket side was taken care of.
160 """
161 if var1.ndim == 2:
162 C1_nM = var1
163 precalc = False
164 assert s is not None
165 elif var1.ndim == 4:
166 g_xNMn = var1
167 precalc = True
168 else:
169 raise ValueError("var1 must be C1_nM or g_xNMn")
170 omega_ql, u_ql = self.phonon.band_structure([q_c], modes=True)
171 u_l = u_ql[0]
172 assert len(u_l.shape) == 3
174 # Defining system sizes
175 nmodes = u_l.shape[0]
176 nbands = C2_nM.shape[0]
177 nao = C2_nM.shape[1]
178 ndisp = 3 * len(self.indices)
180 # Allocate array for couplings
181 g_lnn = np.zeros((nmodes, nbands, nbands), dtype=complex)
183 # Mass scaled polarization vectors
184 u_lx = u_l.reshape(nmodes, ndisp)
186 # Multiply phase factors
187 phase_m = np.exp(2.0j * np.pi * np.einsum("i,im->m", k_c + q_c,
188 self.R_cN))
189 if not precalc:
190 phase_n = np.exp(-2.0j * np.pi * np.einsum("i,in->n", k_c,
191 self.R_cN))
193 # Do each cartesian component separately
194 for i, a in enumerate(self.indices):
195 for v in range(3):
196 xinput = 3 * a + v
197 xoutput = 3 * i + v
198 if not precalc:
199 g_NNMM = self._yield_g_NNMM(xinput, s)
200 assert nao == g_NNMM.shape[-1]
201 # some of these things take a long time. make it fast
202 with self.timer("g_MM"):
203 g_MM = np.einsum("mnop,m,n->op", g_NNMM, phase_m,
204 phase_n, optimize=OPTIMIZE)
205 assert g_MM.shape[0] == g_MM.shape[1]
206 with self.timer("g_nn"):
207 g_nn = np.dot(C2_nM.conj(), np.dot(g_MM, C1_nM.T))
208 # g_nn = np.einsum('no,op,mp->nm', C2_nM.conj(),
209 # g_MM, C1_nM)
210 else:
211 with self.timer("g_Mn"):
212 g_Mn = np.einsum("mon,m->on", g_xNMn[xoutput], phase_m)
213 with self.timer("g_nn"):
214 g_nn = np.dot(C2_nM.conj(), g_Mn)
215 with self.timer("g_lnn"):
216 # g_lnn += np.einsum('i,kl->ikl', u_lx[:, x], g_nn,
217 # optimize=OPTIMIZE)
218 g_lnn += np.multiply.outer(u_lx[:, xoutput], g_nn)
220 # Multiply prefactor sqrt(hbar / 2 * M * omega) in units of Bohr
221 if prefactor:
222 # potential BUG: M needs to be unit cell mass according to
223 # some sources
224 amu = units._amu # atomic mass unit
225 me = units._me # electron mass
226 g_lnn /= np.sqrt(2 * amu / me / units.Hartree *
227 omega_ql[0, :, np.newaxis, np.newaxis])
228 # Convert to eV
229 return g_lnn * units.Hartree # eV
230 else:
231 return g_lnn * units.Hartree / units.Bohr # eV / Ang
233 @timer("g ket part")
234 def _precalculate_ket(self, c_knM, kd, s: int):
235 g_xNkMn = []
236 phase_kn = np.exp(-2.0j * np.pi * np.einsum("ki,in->kn", kd.bzk_kc,
237 self.R_cN))
238 for a in self.indices:
239 for v in range(3):
240 x = 3 * a + v
241 g_NNMM = self._yield_g_NNMM_as_needed(x, s)
242 g_NkMM = np.einsum("mnop,kn->mkop", g_NNMM, phase_kn,
243 optimize=OPTIMIZE)
244 g_NkMn = np.einsum("mkop,knp->mkon", g_NkMM, c_knM,
245 optimize=OPTIMIZE)
246 g_xNkMn.append(g_NkMn)
247 return np.array(g_xNkMn)
249 def bloch_matrix(self, calc: GPAW, k_qc: ArrayND = None,
250 savetofile: bool = True, prefactor: bool = True,
251 accoustic: bool = True) -> ArrayND:
252 r"""Calculate el-ph coupling in the Bloch basis for the electrons.
254 This function calculates the electron-phonon coupling between the
255 specified Bloch states, i.e.::
257 ______
258 mnl / hbar ^
259 g = /------- < m k + q | e . grad V | n k >
260 kq \/ 2 M w ql q
261 ql
263 In case the ``prefactor=False`` is given, the bare matrix
264 element (in units of eV / Ang) without the sqrt prefactor is returned.
266 Parameters
267 ----------
268 calc: GPAW
269 Converged calculator object containing the LCAO wavefuntions
270 (don't use point group symmetry)
271 k_qc: np.ndarray
272 q-vectors of the phonons. Must only contain values comenserate
273 with k-point sampling of calculator. Default: all kpoints used.
274 savetofile: bool
275 If true (default), saves matrix to gsqklnn.npy
276 prefactor: bool
277 if false, don't multiply with sqrt prefactor (Default: True)
278 accoustic: bool
279 if True, for 3 accoustic modes set g=0 for q=0 (Default: True)
280 """
281 kd = calc.wfs.kd
282 assert kd.nbzkpts == kd.nibzkpts, "Elph matrix requires FULL BZ"
284 wfs = calc.wfs
285 if k_qc is None:
286 k_qc = kd.get_bz_q_points(first=True)
287 elif not isinstance(k_qc, np.ndarray):
288 k_qc = np.array(k_qc)
289 assert k_qc.ndim == 2
291 g_sqklnn = np.zeros([wfs.nspins, k_qc.shape[0], kd.nbzkpts,
292 3 * len(self.indices), wfs.bd.nbands,
293 wfs.bd.nbands], dtype=complex)
295 for s in range(wfs.nspins):
296 # Collect all wfcs on rank 0
297 with self.timer("Gather wavefunctions to root"):
298 c_knM = self._gather_all_wfc(wfs, s)
300 # precalculate k (ket) of g
301 g_xNkMn = self._precalculate_ket(c_knM, kd, s)
303 for q, q_c in enumerate(k_qc):
304 # Find indices of k+q for the k-points
305 kplusq_k = kd.find_k_plus_q(q_c) # works on FBZ
306 # Note: calculations require use of FULL BZ,
307 # so NO symmetry
308 print("Spin {}/{}; q-point {}/{}".format(
309 s + 1, wfs.nspins, q + 1, len(k_qc)))
311 for k in range(kd.nbzkpts):
312 k_c = kd.bzk_kc[k]
313 kplusq_c = k_c + q_c
314 kplusq_c -= kplusq_c.round()
315 # print(kplusq_c, kd.bzk_kc[kplusq_k[k]])
316 assert np.allclose(kplusq_c, kd.bzk_kc[kplusq_k[k]])
317 ckplusq_nM = c_knM[kplusq_k[k]]
318 g_lnn = self._bloch_matrix(g_xNkMn[:, :, k], ckplusq_nM,
319 k_c, q_c, prefactor)
320 if np.allclose(q_c, [0.0, 0.0, 0.0]) and accoustic:
321 g_lnn[0:3] = 0.0
322 g_sqklnn[s, q, k] += g_lnn
324 if world.rank == 0 and savetofile:
325 np.save("gsqklnn.npy", g_sqklnn)
327 return g_sqklnn
329 def __del__(self):
330 if world.rank == 0:
331 try:
332 self.timer.write()
333 except ValueError:
334 pass
337# def lcao_matrix(self, u_l, omega_l):
338# """Calculate the el-ph coupling in the electronic LCAO basis.
340# For now, only works for Gamma-point phonons.
342# This method is not tested.
344# Parameters
345# ----------
346# u_l: ndarray
347# Mass-scaled polarization vectors (in units of 1 / sqrt(amu)) of
348# the phonons.
349# omega_l: ndarray
350# Vibrational frequencies in eV.
351# """
353# # Supercell matrix (Hartree / Bohr)
354# assert self.g_xsNNMM is not None, "Load supercell matrix."
355# assert self.g_xsNNMM.shape[2:4] == (1, 1)
356# g_xsMM = self.g_xsNNMM[:, :, 0, 0, :, :]
357# # Number of atomic orbitals
358# # nao = g_xMM.shape[-1]
359# # Number of phonon modes
360# nmodes = u_l.shape[0]
362# #
363# u_lx = u_l.reshape(nmodes, 3 * len(self.atoms))
364# # np.dot uses second to last index of second array
365# g_lsMM = np.dot(u_lx, g_xsMM.transpose(2, 0, 1, 3))
367# # Multiply prefactor sqrt(hbar / 2 * M * omega) in units of Bohr
368# amu = units._amu # atomic mass unit
369# me = units._me # electron mass
370# g_lsMM /= np.sqrt(2 * amu / me / units.Hartree *
371# omega_l[:, :, np.newaxis, np.newaxis])
372# # Convert to eV
373# g_lsMM *= units.Hartree
375# return g_lsMM