Coverage for gpaw/spinorbit.py: 77%
483 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-20 00:19 +0000
1from __future__ import annotations
2from math import nan
3from operator import attrgetter
4from pathlib import Path
5from typing import (TYPE_CHECKING, Callable, Dict, Iterable, Iterator, List,
6 Optional, Tuple)
8import numpy as np
9from ase.units import Bohr, Ha, alpha
11from gpaw.band_descriptor import BandDescriptor
12from gpaw.grid_descriptor import GridDescriptor
13from gpaw.ibz2bz import IBZ2BZMaps
14from gpaw.kpoint import KPoint
15from gpaw.kpt_descriptor import KPointDescriptor
16from gpaw.mpi import broadcast_array, serial_comm
17from gpaw.occupations import OccupationNumberCalculator, ParallelLayout
18from gpaw.projections import Projections
19from gpaw.setup import Setup
20from gpaw.typing import Array1D, Array2D, Array3D, Array4D, ArrayND
21from gpaw.utilities.partition import AtomPartition
23if TYPE_CHECKING:
24 from gpaw.calculator import GPAW # noqa
25 from gpaw.new.ase_interface import ASECalculator
27_L_vlmm: List[List[np.ndarray]] = [] # see get_L_vlmm() below
30class WaveFunction:
31 def __init__(self,
32 eigenvalues: Array1D,
33 projections: Projections,
34 bz_index: int = None):
35 self.eig_m = eigenvalues
36 self.projections = projections
37 self.spin_projection_mv: Optional[Array2D] = None
38 self.v_mn: Optional[Array2D] = None
39 self.f_m = np.empty_like(self.eig_m)
40 self.f_m[:] = nan
41 self.bz_index = bz_index
43 def transform(self, IBZ2BZMap, K) -> WaveFunction:
44 """Transforms PAW projections from IBZ to BZ k-point."""
45 projections = IBZ2BZMap.map_projections(self.projections)
46 return WaveFunction(self.eig_m.copy(), projections, K)
48 def redistribute_atoms(self,
49 atom_partition: AtomPartition
50 ) -> WaveFunction:
51 projections = self.projections.redist(atom_partition)
52 return WaveFunction(self.eig_m.copy(), projections, self.bz_index)
54 def add_soc(self,
55 dVL_avii: Dict[int, Array3D],
56 s_vss: List[Array2D],
57 C_ss: Array2D) -> None:
58 """Evaluate H in a basis of S_z eigenstates."""
59 if self.projections.bcomm.rank > 0:
60 return
62 M = self.projections.nbands
63 H_mm = np.zeros((M, M), complex)
64 for a, dVL_vii in dVL_avii.items():
65 ni = dVL_vii.shape[1]
66 H_ssii = np.zeros((2, 2, ni, ni), complex)
67 H_ssii[0, 0] = dVL_vii[2]
68 H_ssii[0, 1] = dVL_vii[0] - 1.0j * dVL_vii[1]
69 H_ssii[1, 0] = dVL_vii[0] + 1.0j * dVL_vii[1]
70 H_ssii[1, 1] = -dVL_vii[2]
72 # Tranform to theta, phi basis
73 H_ssii = np.tensordot(C_ss, H_ssii, (0, 1))
74 H_ssii = np.tensordot(C_ss.T.conj(), H_ssii, (1, 1))
75 H_ssii *= Ha
77 P_msi = self.projections[a]
78 for s1 in range(2):
79 for s2 in range(2):
80 H_ii = H_ssii[s1, s2]
81 P1_mi = P_msi[:, s1]
82 P2_mi = P_msi[:, s2]
83 H_mm += np.dot(np.dot(P1_mi.conj(), H_ii), P2_mi.T)
85 domain_comm = self.projections.atom_partition.comm
86 domain_comm.sum(H_mm, 0)
87 if domain_comm.rank == 0:
88 H_mm += np.diag(self.eig_m)
89 self.eig_m, v_nm = np.linalg.eigh(H_mm)
90 else:
91 v_nm = np.empty((M, M), complex)
93 domain_comm.broadcast(v_nm, 0)
95 P_mI = self.projections.matrix.array
96 P_mI[:] = v_nm.T.copy().dot(P_mI)
98 sx_m = []
99 sy_m = []
100 sz_m = []
102 v_msn = v_nm.copy().reshape((M // 2, 2, M)).T.copy()
103 for v_sn in v_msn:
104 sx_m.append(np.trace(v_sn.T.conj().dot(s_vss[0]).dot(v_sn)))
105 sy_m.append(np.trace(v_sn.T.conj().dot(s_vss[1]).dot(v_sn)))
106 sz_m.append(np.trace(v_sn.T.conj().dot(s_vss[2]).dot(v_sn)))
108 self.spin_projection_mv = np.array([sx_m, sy_m, sz_m]).real.T.copy()
109 self.v_mn = v_nm.T
111 def wavefunctions(self, calc, periodic=True):
112 kd = calc.wfs.kd
113 assert kd.nibzkpts == kd.nbzkpts
115 # For spinors the number of bands is doubled and a
116 # spin dimension is added
117 Ns = calc.wfs.nspins
118 Nm, Nn = self.v_mn.shape
120 if calc.wfs.collinear:
121 u_snR = [[calc.wfs.get_wave_function_array(n, self.bz_index, s,
122 periodic=periodic)
123 for n in range(Nn // 2)]
124 for s in range(Ns)]
125 u_msR = np.empty((Nm, 2) + u_snR[0][0].shape, complex)
126 np.einsum('mn, nabc -> mabc', self.v_mn[:, ::2], u_snR[0],
127 out=u_msR[:, 0])
128 np.einsum('mn, nabc -> mabc', self.v_mn[:, 1::2], u_snR[-1],
129 out=u_msR[:, 1])
130 else:
131 u_nsR = np.array(
132 [calc.wfs.get_wave_function_array(n, self.bz_index, 0,
133 periodic=periodic)
134 for n in range(Nn)])
135 u_msR = np.einsum('mn, nsxyz -> msxyz',
136 self.v_mn, u_nsR)
138 return u_msR
140 @property
141 def P_amj(self):
142 M = self.projections.nbands
143 return {
144 a: P_msi.transpose((0, 2, 1)).copy().reshape((M, -1))
145 for a, P_msi in self.projections.items()}
147 def pdos_weights(self,
148 a: int,
149 indices: List[int]
150 ) -> Array3D:
151 """PDOS weights."""
152 dos_ms = np.zeros((self.projections.nbands, 2))
154 P_amsi = self.projections
156 if a in P_amsi:
157 dos_ms[:, :] = (abs(P_amsi[a][:, :, indices])**2).sum(2)
159 return dos_ms
162class BZWaveFunctions:
163 """Container for eigenvalues and PAW projections (all of BZ)."""
164 def __init__(self,
165 kd: KPointDescriptor,
166 wfs: Dict[int, WaveFunction],
167 occ: Optional[OccupationNumberCalculator],
168 nelectrons: float,
169 n_aj: List[List[int]],
170 l_aj: List[List[int]]):
171 self.wfs = wfs
172 self.occ = occ
173 self.nelectrons = nelectrons
174 self.n_aj = n_aj
175 self.l_aj = l_aj
177 self.nbzkpts = kd.nbzkpts
179 # Initialize ranks:
180 self.ranks = np.zeros(kd.nbzkpts, int)
181 for k in wfs:
182 self.ranks[k] = kd.comm.rank
183 kd.comm.sum(self.ranks)
185 wf = next(iter(wfs.values())) # get the first WaveFunction object
187 self.shape = (kd.nbzkpts, wf.projections.nbands)
188 self.domain_comm = wf.projections.atom_partition.comm
189 self.bcomm = wf.projections.bcomm
190 self.kpt_comm = kd.comm
192 self.fermi_level = self._calculate_occ_numbers_and_fermi_level()
194 self.size = kd.N_c
195 self.bz2ibz_map = np.arange(self.nbzkpts)
197 def weights(self):
198 return np.zeros(len(self)) + 1 / self.nbzkpts
200 def _calculate_occ_numbers_and_fermi_level(self) -> float:
201 if self.occ is not None:
202 eig_im = [wf.eig_m for wf in self]
203 weight = 1.0 / self.nbzkpts
204 weight_i = [weight] * len(eig_im)
206 f_im, (fermi_level,), _ = self.occ.calculate(
207 self.nelectrons,
208 eig_im,
209 weight_i)
210 for wf, f_m in zip(self, f_im):
211 wf.f_m[:] = f_m
212 else:
213 fermi_level = 0.0
215 if self.domain_comm.rank == 0:
216 fermi_level = self.bcomm.sum_scalar(fermi_level)
217 fermi_level = self.domain_comm.sum_scalar(fermi_level)
219 return fermi_level
221 def calculate_band_energy(self) -> float:
222 """Calculate sum over occupied eigenvalues."""
223 if self.domain_comm.rank == 0 and self.bcomm.rank == 0:
224 weight = 1.0 / self.nbzkpts
225 e_band = sum(wf.eig_m.dot(wf.f_m) for wf in self) * weight
226 e_band = self.kpt_comm.sum_scalar(e_band)
227 else:
228 e_band = 0.0
230 if self.domain_comm.rank == 0:
231 e_band = self.bcomm.sum_scalar(e_band)
232 e_band = self.domain_comm.sum_scalar(e_band)
233 return e_band
235 def __iter__(self):
236 for bz_index in sorted(self.wfs):
237 yield self[bz_index]
239 def __getitem__(self, bz_index):
240 return self.wfs[bz_index]
242 def __len__(self):
243 return len(self.wfs)
245 def eigenvalues(self,
246 broadcast: bool = True
247 ) -> Array2D:
248 """Eigenvalues in eV for the whole BZ."""
249 return self._collect(attrgetter('eig_m'), broadcast=broadcast)
251 def occupation_numbers(self,
252 broadcast: bool = True
253 ) -> Array2D:
254 """Occupation numbers for the whole BZ."""
255 return self._collect(attrgetter('f_m'), broadcast=broadcast)
257 def eigenvectors(self,
258 broadcast: bool = True
259 ) -> Array4D:
260 """Eigenvectors for the whole BZ."""
261 nbands = self.shape[1]
262 assert nbands % 2 == 0
263 return self._collect(attrgetter('v_mn'), (nbands,), complex,
264 broadcast=broadcast)
266 def spin_projections(self,
267 broadcast: bool = True
268 ) -> Array3D:
269 """Spin projections for the whole BZ."""
270 return self._collect(attrgetter('spin_projection_mv'), (3,),
271 broadcast=broadcast)
273 def get_atomic_density_matrices(self):
274 """Returns atomic density matrices for each atom."""
275 from gpaw.new.wave_functions import add_to_4component_density_matrix
277 assert self.domain_comm.size == 1
278 assert self.bcomm.size == 1
280 D_asii = {} # Could also be an AtomArrays object?
281 for a, l_j in enumerate(self.l_aj):
282 ni = (2 * np.array(l_j) + 1).sum()
283 D_sii = np.zeros([4, ni, ni], dtype=complex)
284 for wfs, weight in zip(self.wfs.values(), self.weights()):
285 add_to_4component_density_matrix(D_sii, wfs.projections[a],
286 wfs.f_m * weight, np)
287 self.kpt_comm.sum(D_sii)
288 D_asii[a] = D_sii
290 return D_asii
292 def get_orbital_magnetic_moments(self):
293 """Returns the orbital magnetic moment vector for each atom a."""
294 D_asii = self.get_atomic_density_matrices()
296 from gpaw.new.orbmag import calculate_orbmag_from_density
297 return calculate_orbmag_from_density(D_asii, self.n_aj, self.l_aj)
299 def pdos_weights(self,
300 a: int,
301 indices: List[int],
302 broadcast: bool = True
303 ) -> Array4D:
304 """Projections for PDOS.
306 Returns (nbzkpts, nbands, 2)-shaped ndarray
307 of the square of absolute value of the projections.
308 """
309 def func(wf):
310 return wf.pdos_weights(a, indices)
312 return self._collect(func,
313 (2,),
314 broadcast=broadcast,
315 sum_over_domain=True)
317 def _collect(self,
318 func: Callable[[WaveFunction], ArrayND],
319 shape: Tuple[int, ...] = None,
320 dtype=float,
321 broadcast: bool = True,
322 sum_over_domain: bool = False) -> ArrayND:
323 """Helper method for collecting (and broadcasting) ndarrays."""
325 total_shape = self.shape + (shape or ())
327 if broadcast:
328 array_kmx = self._collect(func,
329 shape,
330 dtype,
331 False,
332 sum_over_domain)
333 if array_kmx.ndim == 0:
334 array_kmx = np.empty(total_shape, dtype)
335 return broadcast_array(array_kmx,
336 self.kpt_comm, self.bcomm, self.domain_comm)
338 if self.bcomm.rank != 0:
339 return np.empty(shape=())
341 if not sum_over_domain and self.domain_comm.rank != 0:
342 return np.empty(shape=())
344 comm = self.kpt_comm
345 if comm.rank == 0:
346 array_kmx = np.empty(total_shape, dtype)
347 for k, rank in enumerate(self.ranks):
348 if rank == 0:
349 array_kmx[k] = func(self[k])
350 else:
351 comm.receive(array_kmx[k], rank)
352 if sum_over_domain:
353 self.domain_comm.sum(array_kmx)
354 if self.domain_comm.rank == 0:
355 return array_kmx
356 else:
357 return np.empty(shape=())
359 for k, rank in enumerate(self.ranks):
360 if rank == comm.rank:
361 comm.send(func(self[k]), 0)
363 return np.empty(shape=())
366def soc_eigenstates_raw(ibzwfs: Iterable[Tuple[int, WaveFunction]],
367 dVL_avii: Dict[int, Array3D],
368 ibz2bzmaps: IBZ2BZMaps,
369 atom_partition,
370 theta: float = 0.0,
371 phi: float = 0.0) -> Dict[int, WaveFunction]:
373 theta *= np.pi / 180
374 phi *= np.pi / 180
376 # Hamiltonian with SO in KS basis
377 # The even indices in H_mm are spin up along \hat n defined by \theta, phi
378 # Basis change matrix for constructing Pauli matrices in \theta,\phi basis:
379 # \sigma_i^n = C^\dag\sigma_i C
380 C_ss = np.array([[np.cos(theta / 2) * np.exp(-1.0j * phi / 2),
381 -np.sin(theta / 2) * np.exp(-1.0j * phi / 2)],
382 [np.sin(theta / 2) * np.exp(1.0j * phi / 2),
383 np.cos(theta / 2) * np.exp(1.0j * phi / 2)]])
385 sx_ss = np.array([[0, 1], [1, 0]], complex)
386 sy_ss = np.array([[0, -1.0j], [1.0j, 0]], complex)
387 sz_ss = np.array([[1, 0], [0, -1]], complex)
388 s_vss = [C_ss.T.conj() @ sx_ss @ C_ss,
389 C_ss.T.conj() @ sy_ss @ C_ss,
390 C_ss.T.conj() @ sz_ss @ C_ss]
392 bzwfs = {}
393 for ibz_index, ibzwf in ibzwfs:
394 for K in np.nonzero(ibz2bzmaps.kd.bz2ibz_k == ibz_index)[0]:
395 bzwf = ibzwf.transform(ibz2bzmaps[K], K)
397 # Redistribute to match dVL_avii:
398 bzwf = bzwf.redistribute_atoms(atom_partition)
400 bzwf.add_soc(dVL_avii, s_vss, C_ss)
401 bzwfs[int(K)] = bzwf # MYPY: signedinteger[_32Bit | _64Bit] -> int
403 return bzwfs
406def extract_ibz_wave_functions(kpt_qs: List[List[KPoint]],
407 bd: BandDescriptor,
408 gd: GridDescriptor,
409 n1: int,
410 n2: int,
411 eigenvalues: Array3D = None
412 ) -> Iterator[Tuple[int, WaveFunction]]:
413 """Yield tuples of IBZ-index and wave functions.
415 All atoms and bands will be on rank == 0 of gd.comm and bd.comm
416 respectively. This makes slicing the bands (from n1 to n2-1)
417 and symmetry operations on the projections easier.
418 """
420 nproj_a = kpt_qs[0][0].projections.nproj_a
422 collinear = kpt_qs[0][0].s is not None
424 nbands = n2 - n1
425 if collinear:
426 nbands *= 2
428 # All atoms on rank-0:
429 atom_partition = AtomPartition(gd.comm, np.zeros_like(nproj_a))
431 # All bands on rank-0 (nrow * ncol = 1 * 1):
432 bdist = (bd.comm, 1, 1)
434 for kpt_s in kpt_qs:
435 # Collect bands and atoms to bd.comm.rank == 0 and gd.comm.rank == 0:
436 if eigenvalues is None:
437 eig_sn = [bd.collect(kpt.eps_n) for kpt in kpt_s]
438 P_snI = [kpt.projections.collect() for kpt in kpt_s]
440 projections = Projections(
441 nbands=nbands,
442 nproj_a=nproj_a,
443 atom_partition=atom_partition,
444 bdist=bdist,
445 collinear=False)
447 if bd.comm.rank == 0 and gd.comm.rank == 0:
448 P1_nI = P_snI[0]
449 P2_nI = P_snI[-1]
450 assert P1_nI is not None
451 assert P2_nI is not None
452 if collinear:
453 eig_m = np.empty((n2 - n1) * 2)
454 if eigenvalues is None:
455 eig_m[::2] = eig_sn[0][n1:n2] * Ha
456 eig_m[1::2] = eig_sn[-1][n1:n2] * Ha
457 else:
458 for s in range(2):
459 eig_m[s::2] = eigenvalues[-s, kpt_s[-s].k]
460 projections.array[:] = 0.0
461 projections.array[::2, 0] = P1_nI[n1:n2]
462 projections.array[1::2, 1] = P2_nI[n1:n2]
463 else:
464 eig_m = eig_sn[0][n1:n2] * Ha
465 projections.matrix.array[:] = P1_nI[n1:n2]
466 else:
467 eig_m = np.empty(0)
469 ibz_index = kpt_s[0].k
471 yield ibz_index, WaveFunction(eig_m, projections)
474def soc_eigenstates(calc: ASECalculator | GPAW | str | Path,
475 n1: int = None,
476 n2: int = None,
477 scale: float = 1.0,
478 theta: float = 0.0, # degrees
479 phi: float = 0.0, # degrees
480 eigenvalues: Array3D = None, # eV
481 occcalc: OccupationNumberCalculator = None,
482 projected: bool = False
483 ) -> BZWaveFunctions:
484 """Calculate SOC eigenstates.
486 Parameters:
487 calc: Calculator
488 GPAW calculator or path to gpw-file.
489 n1, n2: int
490 Range of bands to include (n1 <= n < n2). Default is all
491 bands available.
492 scale: float
493 Scale the spinorbit coupling by this amount.
494 theta: float
495 Angle in degrees.
496 phi: float
497 Angle in degrees.
498 eigenvalues: ndarray
499 Optionally use these eigenvalues instead for those from *calc*.
500 The shape must be: (nspins, nibzkpts, n2 - n1). Units: eV.
501 occcalc:
502 Occupation-number calculator. By default, the one from *calc*
503 will be used.
505 Returns a BZWaveFunctions object covering the whole BZ.
506 """
508 from gpaw.calculator import GPAW # noqa
510 if isinstance(calc, (str, Path)):
511 calc = GPAW(calc)
513 n1 = n1 or 0
514 n2 = n2 or 0
515 if n2 <= 0:
516 if eigenvalues is None:
517 nbands = calc.get_number_of_bands()
518 else:
519 nbands = eigenvalues.shape[2]
520 n2 += nbands
522 # <phi_i|dV_adr / r * L_v|phi_j>
523 dVL_avii = {a: soc(calc.wfs.setups[a],
524 calc.hamiltonian.xc, D_sp) * scale
525 for a, D_sp in calc.density.D_asp.items()}
527 if projected:
528 dVL_avii = {a: projected_soc(dVL_vii, theta=theta, phi=phi)
529 for a, dVL_vii in dVL_avii.items()}
531 kd = calc.wfs.kd
532 bd = calc.wfs.bd
533 gd = calc.wfs.gd
534 atom_partition = calc.density.atom_partition
536 if eigenvalues is not None:
537 assert eigenvalues.shape == (kd.nspins, kd.nibzkpts, n2 - n1)
539 ibzwfs = extract_ibz_wave_functions(calc.wfs.kpt_qs,
540 bd, gd, n1, n2, eigenvalues)
541 ibz2bzmaps = IBZ2BZMaps.from_calculator(calc)
543 bzwfs = soc_eigenstates_raw(ibzwfs,
544 dVL_avii,
545 ibz2bzmaps,
546 atom_partition,
547 theta, phi)
549 if bd.comm.rank == 0 and gd.comm.rank == 0:
550 parallel_layout = ParallelLayout(BandDescriptor(1),
551 kd.comm,
552 serial_comm)
553 occcalc = occcalc or calc.wfs.occupations
554 occcalc = occcalc.copy(bz2ibzmap=list(range(kd.nbzkpts)),
555 parallel_layout=parallel_layout)
556 else:
557 occcalc = None
559 n_aj = [setup.n_j for setup in calc.wfs.setups]
560 l_aj = [setup.l_j for setup in calc.wfs.setups]
562 return BZWaveFunctions(kd, bzwfs, occcalc, calc.wfs.nvalence, n_aj, l_aj)
565def soc(a: Setup, xc, D_sp: Array2D) -> Array3D:
566 """<phi_i|dU^a/dr / r * L_v|phi_j>"""
567 v_g = get_radial_potential_derivative(a, xc, D_sp)
568 Ng = len(v_g)
569 phi_jg = a.data.phi_jg
571 Lx_lmm, Ly_lmm, Lz_lmm = get_L_vlmm()
573 dVL_vii = np.zeros((3, a.ni, a.ni), complex)
574 N1 = 0
575 for j1, l1 in enumerate(a.l_j):
576 Nm = 2 * l1 + 1
577 N2 = 0
578 for j2, l2 in enumerate(a.l_j):
579 if l1 == l2:
580 f_g = phi_jg[j1][:Ng] * v_g * phi_jg[j2][:Ng]
581 c = a.xc_correction.rgd.integrate(f_g, n=-1) / (4 * np.pi)
582 dVL_vii[0, N1:N1 + Nm, N2:N2 + Nm] = c * Lx_lmm[l1]
583 dVL_vii[1, N1:N1 + Nm, N2:N2 + Nm] = c * Ly_lmm[l1]
584 dVL_vii[2, N1:N1 + Nm, N2:N2 + Nm] = c * Lz_lmm[l1]
585 else:
586 pass
587 N2 += 2 * l2 + 1
588 N1 += Nm
589 return dVL_vii * alpha**2 / 4.0
592def projected_soc(dVL_vii: Array3D,
593 theta: float = 0,
594 phi: float = 0) -> Array3D:
595 """
596 Optional Args:
597 theta (float): The angle from z-axis in degrees
598 phi (float): The angle from x-axis in degrees
599 """
600 theta *= np.pi / 180
601 phi *= np.pi / 180
602 n_v = np.array([np.sin(theta) * np.cos(phi),
603 np.sin(theta) * np.sin(phi),
604 np.cos(theta)])
605 dVL_vii = (np.dot(dVL_vii.T, n_v)[:, :, np.newaxis] * n_v).T
606 return dVL_vii
609def get_radial_potential_derivative(a: Setup, xc, D_sp: Array2D) -> Array1D:
610 """Calculates (dU/dr) for the effective potential.
611 Below, f_g denotes dU/dr which is also the negative of the radial force"""
613 rgd = a.xc_correction.rgd
614 r_g = rgd.r_g.copy()
615 r_g[0] = 1.0e-12
616 dr_g = rgd.dr_g
618 B_pq = a.xc_correction.B_pqL[:, :, 0]
619 n_qg = a.xc_correction.n_qg
620 D_sq = np.dot(D_sp, B_pq)
621 n_sg = np.dot(D_sq, n_qg) / (4 * np.pi)**0.5
622 Ns = len(D_sp)
623 if Ns == 4:
624 Ns = 1
625 n_sg[:Ns] += a.xc_correction.nc_g / Ns
627 # Coulomb force from nucleus
628 fc_g = a.Z / r_g**2
630 # Hartree force
631 rho_g = 4 * np.pi * r_g**2 * dr_g * np.sum(n_sg[:Ns], axis=0)
632 fh_g = -np.array([np.sum(rho_g[:ig]) for ig in range(len(r_g))]) / r_g**2
634 f_g = fc_g + fh_g
636 # xc force
637 if xc.type != 'GLLB':
638 v_sg = np.zeros_like(n_sg)
639 xc.calculate_spherical(rgd, n_sg, v_sg)
640 fxc_g = np.mean([rgd.derivative(v_g) for v_g in v_sg[:Ns]],
641 axis=0)
642 f_g += fxc_g
644 return f_g
647def get_anisotropy(calc, theta=0.0, phi=0.0, nbands=0, width=None):
648 """Calculates the sum of occupied spinorbit eigenvalues.
650 Returns the result relative to the sum of eigenvalues without
651 spinorbit coupling.
652 """
653 raise RuntimeError('Please use BZWaveFunctions.calculate_band_energy() '
654 'instead.')
657def get_magnetic_moments(calc, theta=0.0, phi=0.0, nbands=None, width=None):
658 """Calculates the magnetic moments inside all PAW spheres"""
660 raise RuntimeError(
661 'This function has no tests. It is very likely that it no longer '
662 'works correctly after merging !677.')
664 from gpaw.utilities import unpack_hermitian
666 if nbands is None:
667 nbands = calc.get_number_of_bands()
668 Nk = len(calc.get_ibz_k_points())
670 C_ss = np.array([[np.cos(theta / 2) * np.exp(-1.0j * phi / 2),
671 -np.sin(theta / 2) * np.exp(-1.0j * phi / 2)],
672 [np.sin(theta / 2) * np.exp(1.0j * phi / 2),
673 np.cos(theta / 2) * np.exp(1.0j * phi / 2)]])
674 sx_ss = np.array([[0, 1], [1, 0]], complex)
675 sy_ss = np.array([[0, -1.0j], [1.0j, 0]], complex)
676 sz_ss = np.array([[1, 0], [0, -1]], complex)
677 sx_ss = C_ss.T.conj().dot(sx_ss).dot(C_ss)
678 sy_ss = C_ss.T.conj().dot(sy_ss).dot(C_ss)
679 sz_ss = C_ss.T.conj().dot(sz_ss).dot(C_ss)
681 states = soc_eigenstates(calc,
682 theta=theta,
683 phi=phi,
684 return_wfs=True,
685 bands=range(nbands))
686 e_km = states['eigenvalues']
687 v_knm = states['eigenstates']
689 from gpaw.occupations import occupation_numbers
690 if width is None:
691 assert calc.wfs.occupations.name == 'fermi-dirac'
692 width = calc.wfs.occupations._width
693 if width == 0.0:
694 width = 1.e-6
695 weight_k = calc.get_k_point_weights() / 2
696 ne = calc.wfs.setups.nvalence - calc.density.charge
697 f_km = occupation_numbers({'name': 'fermi-dirac', 'width': width},
698 e_km[np.newaxis],
699 weight_k=weight_k,
700 nelectrons=ne)[0][0]
702 m_v = np.zeros(3, complex)
703 for ik in range(Nk):
704 ut0_nG = np.array([calc.wfs.get_wave_function_array(n, ik, 0)
705 for n in range(nbands)])
706 ut1_nG = np.array([calc.wfs.get_wave_function_array(n, ik, 1)
707 for n in range(nbands)])
708 mocc = np.where(f_km[ik] * Nk - 1.0e-6 < 0.0)[0][0]
709 for m in range(mocc + 1):
710 f = f_km[ik, m]
711 ut0_G = np.dot(v_knm[ik][::2, m], np.swapaxes(ut0_nG, 0, 2))
712 ut1_G = np.dot(v_knm[ik][1::2, m], np.swapaxes(ut1_nG, 0, 2))
713 ut_sG = np.array([ut0_G, ut1_G])
715 mx_G = np.zeros(np.shape(ut0_G), complex)
716 my_G = np.zeros(np.shape(ut0_G), complex)
717 mz_G = np.zeros(np.shape(ut0_G), complex)
718 for s1 in range(2):
719 for s2 in range(2):
720 mx_G += ut_sG[s1].conj() * sx_ss[s1, s2] * ut_sG[s2]
721 my_G += ut_sG[s1].conj() * sy_ss[s1, s2] * ut_sG[s2]
722 mz_G += ut_sG[s1].conj() * sz_ss[s1, s2] * ut_sG[s2]
723 m_v += calc.wfs.gd.integrate(np.array([mx_G, my_G, mz_G])) * f
725 m_av = []
726 for a in range(len(calc.atoms)):
727 N0_p = calc.density.setups[a].N0_p.copy()
728 N0_ij = unpack_hermitian(N0_p)
729 Dx_ij = np.zeros_like(N0_ij, complex)
730 Dy_ij = np.zeros_like(N0_ij, complex)
731 Dz_ij = np.zeros_like(N0_ij, complex)
732 Delta_p = calc.density.setups[a].Delta_pL[:, 0].copy()
733 Delta_ij = unpack_hermitian(Delta_p)
734 for ik in range(Nk):
735 P_ami = ... # get_spinorbit_projections(calc, ik, v_knm[ik])
736 P_smi = np.array([P_ami[a][:, ::2], P_ami[a][:, 1::2]])
737 P_smi = np.dot(C_ss, np.swapaxes(P_smi, 0, 1))
739 P0_mi = P_smi[0]
740 P1_mi = P_smi[1]
741 f_mm = np.diag(f_km[ik])
743 Dx_ij += P0_mi.conj().T.dot(f_mm).dot(P1_mi)
744 Dx_ij += P1_mi.conj().T.dot(f_mm).dot(P0_mi)
745 Dy_ij -= 1.0j * P0_mi.conj().T.dot(f_mm).dot(P1_mi)
746 Dy_ij += 1.0j * P1_mi.conj().T.dot(f_mm).dot(P0_mi)
747 Dz_ij += P0_mi.conj().T.dot(f_mm).dot(P0_mi)
748 Dz_ij -= P1_mi.conj().T.dot(f_mm).dot(P1_mi)
750 mx = np.sum(N0_ij * Dx_ij).real
751 my = np.sum(N0_ij * Dy_ij).real
752 mz = np.sum(N0_ij * Dz_ij).real
754 m_av.append([mx, my, mz])
755 m_v[0] += np.sum(Delta_ij * Dx_ij) * (4 * np.pi)**0.5
756 m_v[1] += np.sum(Delta_ij * Dy_ij) * (4 * np.pi)**0.5
757 m_v[2] += np.sum(Delta_ij * Dz_ij) * (4 * np.pi)**0.5
759 return m_v.real, m_av
762def get_parity_eigenvalues(calc, ik=0, spin_orbit=False, bands=None, Nv=None,
763 inversion_center=[0, 0, 0], deg_tol=1.0e-6,
764 tol=1.0e-6):
765 """Calculates parity eigenvalues at time-reversal invariant k-points.
766 Only works in plane wave mode.
767 """
769 assert len(calc.get_bz_k_points()) == len(calc.get_ibz_k_points())
771 kpt_c = calc.get_ibz_k_points()[ik]
772 if Nv is None:
773 Nv = int(calc.get_number_of_electrons() / 2)
775 if bands is None:
776 bands = range(calc.get_number_of_bands())
778 # Find degenerate subspaces
779 eig_n = calc.get_eigenvalues(kpt=ik)[bands]
780 e_in = []
781 used_n = []
782 for n1, e1 in enumerate(eig_n):
783 if n1 not in used_n:
784 n_n = []
785 for n2, e2 in enumerate(eig_n):
786 if np.abs(e1 - e2) < deg_tol:
787 n_n.append(n2)
788 used_n.append(n2)
789 e_in.append(n_n)
791 print()
792 print(f' Inversion center at: {inversion_center}')
793 print(f' Calculating inversion eigenvalues at k = {kpt_c}')
794 print()
796 center_v = np.array(inversion_center) / Bohr
797 G_Gv = calc.wfs.pd.get_reciprocal_vectors(q=ik, add_q=True)
799 psit_nG = np.array([calc.wfs.kpt_u[ik].psit_nG[n]
800 for n in bands])
801 if spin_orbit:
802 n1 = bands[0]
803 n2 = bands[-1] + 1
804 assert (bands == np.arange(n1, n2)).all()
805 soc = soc_eigenstates(calc, n1=n1, n2=n2)
806 v_kmn = soc.eigenvectors()
807 psit0_mG = np.dot(v_kmn[ik, :, ::2], psit_nG)
808 psit1_mG = np.dot(v_kmn[ik, :, 1::2], psit_nG)
809 for n in range(len(bands)):
810 psit_nG[n] /= (np.sum(np.abs(psit_nG[n])**2))**0.5
811 if spin_orbit:
812 for n in range(2 * len(bands)):
813 A = np.sum(np.abs(psit0_mG[n])**2) + np.sum(np.abs(psit1_mG[n])**2)
814 psit0_mG[n] /= A**0.5
815 psit1_mG[n] /= A**0.5
817 P_GG = np.ones((len(G_Gv), len(G_Gv)), float)
818 for iG, G_v in enumerate(G_Gv):
819 P_GG[iG] -= ((G_Gv[:] + G_v).round(6)).any(axis=1)
820 assert (P_GG == P_GG.T).all()
822 phase_G = np.exp(-2.0j * np.dot(G_Gv, center_v))
824 p_n = []
825 print('n P_n')
826 for n_n in e_in:
827 if spin_orbit:
828 # The dimension of parity matrix is doubled with spinorbit
829 m_m = [2 * n_n[0] + i for i in range(2 * len(n_n))]
830 Ppsit0_mG = np.dot(P_GG, psit0_mG[m_m].T).T
831 Ppsit0_mG[:] *= phase_G
832 Ppsit1_mG = np.dot(P_GG, psit1_mG[m_m].T).T
833 Ppsit1_mG[:] *= phase_G
834 P_nn = np.dot(psit0_mG[m_m].conj(), np.array(Ppsit0_mG).T)
835 P_nn += np.dot(psit1_mG[m_m].conj(), np.array(Ppsit1_mG).T)
836 else:
837 Ppsit_nG = np.dot(P_GG, psit_nG[n_n].T).T
838 Ppsit_nG[:] *= phase_G
839 P_nn = np.dot(psit_nG[n_n].conj(), np.array(Ppsit_nG).T)
840 P_eig = np.linalg.eigh(P_nn)[0]
841 if np.allclose(np.abs(P_eig), 1, tol):
842 P_n = np.sign(P_eig).astype(int).tolist()
843 if spin_orbit:
844 # Only include one of the degenerate pair of eigenvalues
845 Pm = np.sign(P_eig).tolist().count(-1)
846 Pp = np.sign(P_eig).tolist().count(1)
847 P_n = Pm // 2 * [-1] + Pp // 2 * [1]
848 print(f'{str(n_n)[1:-1]}: {str(P_n)[1:-1]}')
849 p_n += P_n
850 else:
851 print(f' {n_n} are not parity eigenstates')
852 print(f' P_n: {P_eig}')
853 print(f' e_n: {eig_n[n_n]}')
854 p_n += [0 for n in n_n]
856 return np.ravel(p_n)
859def get_L_vlmm():
860 if len(_L_vlmm) == 3:
861 return _L_vlmm
863 s = np.array([[0.0]])
864 p = np.zeros((3, 3), complex) # y, z, x
865 p[0, 1] = -1.0j
866 p[1, 0] = 1.0j
867 d = np.zeros((5, 5), complex) # xy, yz, z^2, xz, x^2-y^2
868 d[0, 3] = -1.0j
869 d[3, 0] = 1.0j
870 d[1, 2] = -3**0.5 * 1.0j
871 d[2, 1] = 3**0.5 * 1.0j
872 d[1, 4] = -1.0j
873 d[4, 1] = 1.0j
874 f = np.zeros((7, 7), complex)
875 f[0, 5] = -0.5 * 6**0.5 * 1.0j
876 f[5, 0] = 0.5 * 6**0.5 * 1.0j
877 f[1, 4] = -0.5 * 10**0.5 * 1.0j
878 f[4, 1] = 0.5 * 10**0.5 * 1.0j
879 f[1, 6] = -0.5 * 6**0.5 * 1.0j
880 f[6, 1] = 0.5 * 6**0.5 * 1.0j
881 f[2, 3] = -6**0.5 * 1.0j
882 f[3, 2] = 6**0.5 * 1.0j
883 f[2, 5] = -0.5 * 10**0.5 * 1.0j
884 f[5, 2] = 0.5 * 10**0.5 * 1.0j
885 _L_vlmm.append([s, p, d, f])
887 p = np.zeros((3, 3), complex) # y, z, x
888 p[1, 2] = -1.0j
889 p[2, 1] = 1.0j
890 d = np.zeros((5, 5), complex) # xy, yz, z^2, xz, x^2-y^2
891 d[0, 1] = 1.0j
892 d[1, 0] = -1.0j
893 d[2, 3] = -3**0.5 * 1.0j
894 d[3, 2] = 3**0.5 * 1.0j
895 d[3, 4] = -1.0j
896 d[4, 3] = 1.0j
897 f = np.zeros((7, 7), complex)
898 f[0, 1] = 0.5 * 6**0.5 * 1.0j
899 f[1, 0] = -0.5 * 6**0.5 * 1.0j
900 f[1, 2] = 0.5 * 10**0.5 * 1.0j
901 f[2, 1] = -0.5 * 10**0.5 * 1.0j
902 f[3, 4] = -6**0.5 * 1.0j
903 f[4, 3] = 6**0.5 * 1.0j
904 f[4, 5] = -0.5 * 10**0.5 * 1.0j
905 f[5, 4] = 0.5 * 10**0.5 * 1.0j
906 f[5, 6] = -0.5 * 6**0.5 * 1.0j
907 f[6, 5] = 0.5 * 6**0.5 * 1.0j
908 _L_vlmm.append([s, p, d, f])
910 p = np.zeros((3, 3), complex) # y, z, x
911 p[0, 2] = 1.0j
912 p[2, 0] = -1.0j
913 d = np.zeros((5, 5), complex) # xy, yz, z^2, xz, x^2-y^2
914 d[0, 4] = 2.0j
915 d[4, 0] = -2.0j
916 d[1, 3] = 1.0j
917 d[3, 1] = -1.0j
918 f = np.zeros((7, 7), complex)
919 f[0, 6] = 3.0j
920 f[6, 0] = -3.0j
921 f[1, 5] = 2.0j
922 f[5, 1] = -2.0j
923 f[2, 4] = 1.0j
924 f[4, 2] = -1.0j
925 _L_vlmm.append([s, p, d, f])
926 return _L_vlmm