from __future__ import annotations
from math import nan
from operator import attrgetter
from pathlib import Path
from typing import (TYPE_CHECKING, Callable, Dict, Iterable, Iterator, List,
Optional, Tuple)
import numpy as np
from ase.units import Bohr, Ha, alpha
from gpaw.band_descriptor import BandDescriptor
from gpaw.grid_descriptor import GridDescriptor
from gpaw.kpoint import KPoint
from gpaw.kpt_descriptor import KPointDescriptor
from gpaw.mpi import broadcast_array, serial_comm
from gpaw.occupations import OccupationNumberCalculator, ParallelLayout
from gpaw.projections import Projections
from gpaw.setup import Setup
from gpaw.typing import Array1D, Array2D, Array3D, Array4D, ArrayND
from gpaw.ibz2bz import IBZ2BZMaps
from gpaw.utilities.partition import AtomPartition
if TYPE_CHECKING:
from gpaw.calculator import GPAW # noqa
from gpaw.new.ase_interface import ASECalculator
_L_vlmm: List[List[np.ndarray]] = [] # see get_L_vlmm() below
[docs]class WaveFunction:
def __init__(self,
eigenvalues: Array1D,
projections: Projections,
bz_index: int = None):
self.eig_m = eigenvalues
self.projections = projections
self.spin_projection_mv: Optional[Array2D] = None
self.v_mn: Optional[Array2D] = None
self.f_m = np.empty_like(self.eig_m)
self.f_m[:] = nan
self.bz_index = bz_index
def redistribute_atoms(self,
atom_partition: AtomPartition
) -> WaveFunction:
projections = self.projections.redist(atom_partition)
return WaveFunction(self.eig_m.copy(), projections, self.bz_index)
[docs] def add_soc(self,
dVL_avii: Dict[int, Array3D],
s_vss: List[Array2D],
C_ss: Array2D) -> None:
"""Evaluate H in a basis of S_z eigenstates."""
if self.projections.bcomm.rank > 0:
return
M = self.projections.nbands
H_mm = np.zeros((M, M), complex)
for a, dVL_vii in dVL_avii.items():
ni = dVL_vii.shape[1]
H_ssii = np.zeros((2, 2, ni, ni), complex)
H_ssii[0, 0] = dVL_vii[2]
H_ssii[0, 1] = dVL_vii[0] - 1.0j * dVL_vii[1]
H_ssii[1, 0] = dVL_vii[0] + 1.0j * dVL_vii[1]
H_ssii[1, 1] = -dVL_vii[2]
# Tranform to theta, phi basis
H_ssii = np.tensordot(C_ss, H_ssii, (0, 1))
H_ssii = np.tensordot(C_ss.T.conj(), H_ssii, (1, 1))
H_ssii *= Ha
P_msi = self.projections[a]
for s1 in range(2):
for s2 in range(2):
H_ii = H_ssii[s1, s2]
P1_mi = P_msi[:, s1]
P2_mi = P_msi[:, s2]
H_mm += np.dot(np.dot(P1_mi.conj(), H_ii), P2_mi.T)
domain_comm = self.projections.atom_partition.comm
domain_comm.sum(H_mm, 0)
if domain_comm.rank == 0:
H_mm += np.diag(self.eig_m)
self.eig_m, v_nm = np.linalg.eigh(H_mm)
else:
v_nm = np.empty((M, M), complex)
domain_comm.broadcast(v_nm, 0)
P_mI = self.projections.matrix.array
P_mI[:] = v_nm.T.copy().dot(P_mI)
sx_m = []
sy_m = []
sz_m = []
v_msn = v_nm.copy().reshape((M // 2, 2, M)).T.copy()
for v_sn in v_msn:
sx_m.append(np.trace(v_sn.T.conj().dot(s_vss[0]).dot(v_sn)))
sy_m.append(np.trace(v_sn.T.conj().dot(s_vss[1]).dot(v_sn)))
sz_m.append(np.trace(v_sn.T.conj().dot(s_vss[2]).dot(v_sn)))
self.spin_projection_mv = np.array([sx_m, sy_m, sz_m]).real.T.copy()
self.v_mn = v_nm.T
def wavefunctions(self, calc, periodic=True):
kd = calc.wfs.kd
assert kd.nibzkpts == kd.nbzkpts
# For spinors the number of bands is doubled and a
# spin dimension is added
Ns = calc.wfs.nspins
Nm, Nn = self.v_mn.shape
if calc.wfs.collinear:
u_snR = [[calc.wfs.get_wave_function_array(n, self.bz_index, s,
periodic=periodic)
for n in range(Nn // 2)]
for s in range(Ns)]
u_msR = np.empty((Nm, 2) + u_snR[0][0].shape, complex)
np.einsum('mn, nabc -> mabc', self.v_mn[:, ::2], u_snR[0],
out=u_msR[:, 0])
np.einsum('mn, nabc -> mabc', self.v_mn[:, 1::2], u_snR[-1],
out=u_msR[:, 1])
else:
u_nsR = np.array(
[calc.wfs.get_wave_function_array(n, self.bz_index, 0,
periodic=periodic)
for n in range(Nn)])
u_msR = np.einsum('mn, nsxyz -> msxyz',
self.v_mn, u_nsR)
return u_msR
@property
def P_amj(self):
M = self.projections.nbands
return {
a: P_msi.transpose((0, 2, 1)).copy().reshape((M, -1))
for a, P_msi in self.projections.items()}
[docs] def pdos_weights(self,
a: int,
indices: List[int]
) -> Array3D:
"""PDOS weights."""
dos_ms = np.zeros((self.projections.nbands, 2))
P_amsi = self.projections
if a in P_amsi:
dos_ms[:, :] = (abs(P_amsi[a][:, :, indices])**2).sum(2)
return dos_ms
[docs]class BZWaveFunctions:
"""Container for eigenvalues and PAW projections (all of BZ)."""
def __init__(self,
kd: KPointDescriptor,
wfs: Dict[int, WaveFunction],
occ: Optional[OccupationNumberCalculator],
nelectrons: float,
n_aj: List[List[int]],
l_aj: List[List[int]]):
self.wfs = wfs
self.occ = occ
self.nelectrons = nelectrons
self.n_aj = n_aj
self.l_aj = l_aj
self.nbzkpts = kd.nbzkpts
# Initialize ranks:
self.ranks = np.zeros(kd.nbzkpts, int)
for k in wfs:
self.ranks[k] = kd.comm.rank
kd.comm.sum(self.ranks)
wf = next(iter(wfs.values())) # get the first WaveFunction object
self.shape = (kd.nbzkpts, wf.projections.nbands)
self.domain_comm = wf.projections.atom_partition.comm
self.bcomm = wf.projections.bcomm
self.kpt_comm = kd.comm
self.fermi_level = self._calculate_occ_numbers_and_fermi_level()
self.size = kd.N_c
self.bz2ibz_map = np.arange(self.nbzkpts)
def weights(self):
return np.zeros(len(self)) + 1 / self.nbzkpts
def _calculate_occ_numbers_and_fermi_level(self) -> float:
if self.occ is not None:
eig_im = [wf.eig_m for wf in self]
weight = 1.0 / self.nbzkpts
weight_i = [weight] * len(eig_im)
f_im, (fermi_level,), _ = self.occ.calculate(
self.nelectrons,
eig_im,
weight_i)
for wf, f_m in zip(self, f_im):
wf.f_m[:] = f_m
else:
fermi_level = 0.0
if self.domain_comm.rank == 0 and self.bcomm.rank == 0:
fermi_level = self.bcomm.sum_scalar(fermi_level)
fermi_level = self.domain_comm.sum_scalar(fermi_level)
return fermi_level
[docs] def calculate_band_energy(self) -> float:
"""Calculate sum over occupied eigenvalues."""
if self.domain_comm.rank == 0 and self.bcomm.rank == 0:
weight = 1.0 / self.nbzkpts
e_band = sum(wf.eig_m.dot(wf.f_m) for wf in self) * weight
e_band = self.kpt_comm.sum_scalar(e_band)
else:
e_band = 0.0
if self.domain_comm.rank == 0:
e_band = self.bcomm.sum_scalar(e_band)
e_band = self.domain_comm.sum_scalar(e_band)
return e_band
def __iter__(self):
for bz_index in sorted(self.wfs):
yield self[bz_index]
def __getitem__(self, bz_index):
return self.wfs[bz_index]
def __len__(self):
return len(self.wfs)
[docs] def eigenvalues(self,
broadcast: bool = True
) -> Array2D:
"""Eigenvalues in eV for the whole BZ."""
return self._collect(attrgetter('eig_m'), broadcast=broadcast)
[docs] def occupation_numbers(self,
broadcast: bool = True
) -> Array2D:
"""Occupation numbers for the whole BZ."""
return self._collect(attrgetter('f_m'), broadcast=broadcast)
[docs] def eigenvectors(self,
broadcast: bool = True
) -> Array4D:
"""Eigenvectors for the whole BZ."""
nbands = self.shape[1]
assert nbands % 2 == 0
return self._collect(attrgetter('v_mn'), (nbands,), complex,
broadcast=broadcast)
[docs] def spin_projections(self,
broadcast: bool = True
) -> Array3D:
"""Spin projections for the whole BZ."""
return self._collect(attrgetter('spin_projection_mv'), (3,),
broadcast=broadcast)
[docs] def get_atomic_density_matrices(self):
"""Returns atomic density matrices for each atom."""
from gpaw.new.wave_functions import add_to_4component_density_matrix
assert self.domain_comm.size == 1
assert self.bcomm.size == 1
D_asii = {} # Could also be an AtomArrays object?
for a, l_j in enumerate(self.l_aj):
ni = (2 * np.array(l_j) + 1).sum()
D_sii = np.zeros([4, ni, ni], dtype=complex)
for wfs, weight in zip(self.wfs.values(), self.weights()):
add_to_4component_density_matrix(D_sii, wfs.projections[a],
wfs.f_m * weight, np)
self.kpt_comm.sum(D_sii)
D_asii[a] = D_sii
return D_asii
[docs] def get_orbital_magnetic_moments(self):
"""Returns the orbital magnetic moment vector for each atom a."""
D_asii = self.get_atomic_density_matrices()
from gpaw.new.orbmag import calculate_orbmag_from_density
return calculate_orbmag_from_density(D_asii, self.n_aj, self.l_aj)
[docs] def pdos_weights(self,
a: int,
indices: List[int],
broadcast: bool = True
) -> Array4D:
"""Projections for PDOS.
Returns (nbzkpts, nbands, 2)-shaped ndarray
of the square of absolute value of the projections.
"""
def func(wf):
return wf.pdos_weights(a, indices)
return self._collect(func,
(2,),
broadcast=broadcast,
sum_over_domain=True)
def _collect(self,
func: Callable[[WaveFunction], ArrayND],
shape: Tuple[int, ...] = None,
dtype=float,
broadcast: bool = True,
sum_over_domain: bool = False) -> ArrayND:
"""Helper method for collecting (and broadcasting) ndarrays."""
total_shape = self.shape + (shape or ())
if broadcast:
array_kmx = self._collect(func,
shape,
dtype,
False,
sum_over_domain)
if array_kmx.ndim == 0:
array_kmx = np.empty(total_shape, dtype)
return broadcast_array(array_kmx,
self.kpt_comm, self.bcomm, self.domain_comm)
if self.bcomm.rank != 0:
return np.empty(shape=())
if not sum_over_domain and self.domain_comm.rank != 0:
return np.empty(shape=())
comm = self.kpt_comm
if comm.rank == 0:
array_kmx = np.empty(total_shape, dtype)
for k, rank in enumerate(self.ranks):
if rank == 0:
array_kmx[k] = func(self[k])
else:
comm.receive(array_kmx[k], rank)
if sum_over_domain:
self.domain_comm.sum(array_kmx)
if self.domain_comm.rank == 0:
return array_kmx
else:
return np.empty(shape=())
for k, rank in enumerate(self.ranks):
if rank == comm.rank:
comm.send(func(self[k]), 0)
return np.empty(shape=())
def soc_eigenstates_raw(ibzwfs: Iterable[Tuple[int, WaveFunction]],
dVL_avii: Dict[int, Array3D],
ibz2bzmaps: IBZ2BZMaps,
atom_partition,
theta: float = 0.0,
phi: float = 0.0) -> Dict[int, WaveFunction]:
theta *= np.pi / 180
phi *= np.pi / 180
# Hamiltonian with SO in KS basis
# The even indices in H_mm are spin up along \hat n defined by \theta, phi
# Basis change matrix for constructing Pauli matrices in \theta,\phi basis:
# \sigma_i^n = C^\dag\sigma_i C
C_ss = np.array([[np.cos(theta / 2) * np.exp(-1.0j * phi / 2),
-np.sin(theta / 2) * np.exp(-1.0j * phi / 2)],
[np.sin(theta / 2) * np.exp(1.0j * phi / 2),
np.cos(theta / 2) * np.exp(1.0j * phi / 2)]])
sx_ss = np.array([[0, 1], [1, 0]], complex)
sy_ss = np.array([[0, -1.0j], [1.0j, 0]], complex)
sz_ss = np.array([[1, 0], [0, -1]], complex)
s_vss = [C_ss.T.conj() @ sx_ss @ C_ss,
C_ss.T.conj() @ sy_ss @ C_ss,
C_ss.T.conj() @ sz_ss @ C_ss]
bzwfs = {}
for ibz_index, ibzwf in ibzwfs:
for K in np.nonzero(ibz2bzmaps.kd.bz2ibz_k == ibz_index)[0]:
bzwf = ibzwf.transform(ibz2bzmaps[K], K)
# Redistribute to match dVL_avii:
bzwf = bzwf.redistribute_atoms(atom_partition)
bzwf.add_soc(dVL_avii, s_vss, C_ss)
bzwfs[K] = bzwf
return bzwfs
def extract_ibz_wave_functions(kpt_qs: List[List[KPoint]],
bd: BandDescriptor,
gd: GridDescriptor,
n1: int,
n2: int,
eigenvalues: Array3D = None
) -> Iterator[Tuple[int, WaveFunction]]:
"""Yield tuples of IBZ-index and wave functions.
All atoms and bands will be on rank == 0 of gd.comm and bd.comm
respectively. This makes slicing the bands (from n1 to n2-1)
and symmetry operations on the projections easier.
"""
nproj_a = kpt_qs[0][0].projections.nproj_a
collinear = kpt_qs[0][0].s is not None
nbands = n2 - n1
if collinear:
nbands *= 2
# All atoms on rank-0:
atom_partition = AtomPartition(gd.comm, np.zeros_like(nproj_a))
# All bands on rank-0 (nrow * ncol = 1 * 1):
bdist = (bd.comm, 1, 1)
for kpt_s in kpt_qs:
# Collect bands and atoms to bd.comm.rank == 0 and gd.comm.rank == 0:
if eigenvalues is None:
eig_sn = [bd.collect(kpt.eps_n) for kpt in kpt_s]
P_snI = [kpt.projections.collect() for kpt in kpt_s]
projections = Projections(
nbands=nbands,
nproj_a=nproj_a,
atom_partition=atom_partition,
bdist=bdist,
collinear=False)
if bd.comm.rank == 0 and gd.comm.rank == 0:
P1_nI = P_snI[0]
P2_nI = P_snI[-1]
assert P1_nI is not None
assert P2_nI is not None
if collinear:
eig_m = np.empty((n2 - n1) * 2)
if eigenvalues is None:
eig_m[::2] = eig_sn[0][n1:n2] * Ha
eig_m[1::2] = eig_sn[-1][n1:n2] * Ha
else:
for s in range(2):
eig_m[s::2] = eigenvalues[-s, kpt_s[-s].k]
projections.array[:] = 0.0
projections.array[::2, 0] = P1_nI[n1:n2]
projections.array[1::2, 1] = P2_nI[n1:n2]
else:
eig_m = eig_sn[0][n1:n2] * Ha
projections.matrix.array[:] = P1_nI[n1:n2]
else:
eig_m = np.empty(0)
ibz_index = kpt_s[0].k
yield ibz_index, WaveFunction(eig_m, projections)
[docs]def soc_eigenstates(calc: ASECalculator | GPAW | str | Path,
n1: int = None,
n2: int = None,
scale: float = 1.0,
theta: float = 0.0, # degrees
phi: float = 0.0, # degrees
eigenvalues: Array3D = None, # eV
occcalc: OccupationNumberCalculator = None,
projected: bool = False
) -> BZWaveFunctions:
"""Calculate SOC eigenstates.
Parameters:
calc: Calculator
GPAW calculator or path to gpw-file.
n1, n2: int
Range of bands to include (n1 <= n < n2). Default is all
bands available.
scale: float
Scale the spinorbit coupling by this amount.
theta: float
Angle in degrees.
phi: float
Angle in degrees.
eigenvalues: ndarray
Optionally use these eigenvalues instead for those from *calc*.
The shape must be: (nspins, nibzkpts, n2 - n1). Units: eV.
occcalc:
Occupation-number calculator. By default, the one from *calc*
will be used.
Returns a BZWaveFunctions object covering the whole BZ.
"""
from gpaw.calculator import GPAW # noqa
if isinstance(calc, (str, Path)):
calc = GPAW(calc)
n1 = n1 or 0
n2 = n2 or 0
if n2 <= 0:
if eigenvalues is None:
nbands = calc.get_number_of_bands()
else:
nbands = eigenvalues.shape[2]
n2 += nbands
# <phi_i|dV_adr / r * L_v|phi_j>
dVL_avii = {a: soc(calc.wfs.setups[a],
calc.hamiltonian.xc, D_sp) * scale
for a, D_sp in calc.density.D_asp.items()}
if projected:
dVL_avii = {a: projected_soc(dVL_vii, theta=theta, phi=phi)
for a, dVL_vii in dVL_avii.items()}
kd = calc.wfs.kd
bd = calc.wfs.bd
gd = calc.wfs.gd
atom_partition = calc.density.atom_partition
if eigenvalues is not None:
assert eigenvalues.shape == (kd.nspins, kd.nibzkpts, n2 - n1)
ibzwfs = extract_ibz_wave_functions(calc.wfs.kpt_qs,
bd, gd, n1, n2, eigenvalues)
ibz2bzmaps = IBZ2BZMaps.from_calculator(calc)
bzwfs = soc_eigenstates_raw(ibzwfs,
dVL_avii,
ibz2bzmaps,
atom_partition,
theta, phi)
if bd.comm.rank == 0 and gd.comm.rank == 0:
parallel_layout = ParallelLayout(BandDescriptor(1),
kd.comm,
serial_comm)
occcalc = occcalc or calc.wfs.occupations
occcalc = occcalc.copy(bz2ibzmap=list(range(kd.nbzkpts)),
parallel_layout=parallel_layout)
else:
occcalc = None
n_aj = [setup.n_j for setup in calc.wfs.setups]
l_aj = [setup.l_j for setup in calc.wfs.setups]
return BZWaveFunctions(kd, bzwfs, occcalc, calc.wfs.nvalence, n_aj, l_aj)
def soc(a: Setup, xc, D_sp: Array2D) -> Array3D:
"""<phi_i|dV_adr / r * L_v|phi_j>"""
v_g = get_radial_potential(a, xc, D_sp)
Ng = len(v_g)
phi_jg = a.data.phi_jg
Lx_lmm, Ly_lmm, Lz_lmm = get_L_vlmm()
dVL_vii = np.zeros((3, a.ni, a.ni), complex)
N1 = 0
for j1, l1 in enumerate(a.l_j):
Nm = 2 * l1 + 1
N2 = 0
for j2, l2 in enumerate(a.l_j):
if l1 == l2:
f_g = phi_jg[j1][:Ng] * v_g * phi_jg[j2][:Ng]
c = a.xc_correction.rgd.integrate(f_g) / (4 * np.pi)
dVL_vii[0, N1:N1 + Nm, N2:N2 + Nm] = c * Lx_lmm[l1]
dVL_vii[1, N1:N1 + Nm, N2:N2 + Nm] = c * Ly_lmm[l1]
dVL_vii[2, N1:N1 + Nm, N2:N2 + Nm] = c * Lz_lmm[l1]
else:
pass
N2 += 2 * l2 + 1
N1 += Nm
return dVL_vii * alpha**2 / 4.0
def projected_soc(dVL_vii: Array3D,
theta: float = 0,
phi: float = 0) -> Array3D:
"""
Optional Args:
theta (float): The angle from z-axis in degrees
phi (float): The angle from x-axis in degrees
"""
theta *= np.pi / 180
phi *= np.pi / 180
n_v = np.array([np.sin(theta) * np.cos(phi),
np.sin(theta) * np.sin(phi),
np.cos(theta)])
dVL_vii = (np.dot(dVL_vii.T, n_v)[:, :, np.newaxis] * n_v).T
return dVL_vii
def get_radial_potential(a: Setup, xc, D_sp: Array2D) -> Array1D:
"""Calculates (dV/dr)/r for the effective potential.
Below, f_g denotes dV/dr = minus the radial force"""
rgd = a.xc_correction.rgd
r_g = rgd.r_g.copy()
r_g[0] = 1.0e-12
dr_g = rgd.dr_g
B_pq = a.xc_correction.B_pqL[:, :, 0]
n_qg = a.xc_correction.n_qg
D_sq = np.dot(D_sp, B_pq)
n_sg = np.dot(D_sq, n_qg) / (4 * np.pi)**0.5
Ns = len(D_sp)
if Ns == 4:
Ns = 1
n_sg[:Ns] += a.xc_correction.nc_g / Ns
# Coulomb force from nucleus
fc_g = a.Z / r_g**2
# Hartree force
rho_g = 4 * np.pi * r_g**2 * dr_g * np.sum(n_sg[:Ns], axis=0)
fh_g = -np.array([np.sum(rho_g[:ig]) for ig in range(len(r_g))]) / r_g**2
f_g = fc_g + fh_g
# xc force
if xc.type != 'GLLB':
v_sg = np.zeros_like(n_sg)
xc.calculate_spherical(rgd, n_sg, v_sg)
fxc_g = np.mean([rgd.derivative(v_g) for v_g in v_sg[:Ns]],
axis=0)
f_g += fxc_g
return f_g / r_g
def get_anisotropy(calc, theta=0.0, phi=0.0, nbands=0, width=None):
"""Calculates the sum of occupied spinorbit eigenvalues.
Returns the result relative to the sum of eigenvalues without
spinorbit coupling.
"""
raise RuntimeError('Please use BZWaveFunctions.calculate_band_energy() '
'instead.')
def get_magnetic_moments(calc, theta=0.0, phi=0.0, nbands=None, width=None):
"""Calculates the magnetic moments inside all PAW spheres"""
raise RuntimeError(
'This function has no tests. It is very likely that it no longer '
'works correctly after merging !677.')
from gpaw.utilities import unpack
if nbands is None:
nbands = calc.get_number_of_bands()
Nk = len(calc.get_ibz_k_points())
C_ss = np.array([[np.cos(theta / 2) * np.exp(-1.0j * phi / 2),
-np.sin(theta / 2) * np.exp(-1.0j * phi / 2)],
[np.sin(theta / 2) * np.exp(1.0j * phi / 2),
np.cos(theta / 2) * np.exp(1.0j * phi / 2)]])
sx_ss = np.array([[0, 1], [1, 0]], complex)
sy_ss = np.array([[0, -1.0j], [1.0j, 0]], complex)
sz_ss = np.array([[1, 0], [0, -1]], complex)
sx_ss = C_ss.T.conj().dot(sx_ss).dot(C_ss)
sy_ss = C_ss.T.conj().dot(sy_ss).dot(C_ss)
sz_ss = C_ss.T.conj().dot(sz_ss).dot(C_ss)
states = soc_eigenstates(calc,
theta=theta,
phi=phi,
return_wfs=True,
bands=range(nbands))
e_km = states['eigenvalues']
v_knm = states['eigenstates']
from gpaw.occupations import occupation_numbers
if width is None:
assert calc.wfs.occupations.name == 'fermi-dirac'
width = calc.wfs.occupations._width
if width == 0.0:
width = 1.e-6
weight_k = calc.get_k_point_weights() / 2
ne = calc.wfs.setups.nvalence - calc.density.charge
f_km = occupation_numbers({'name': 'fermi-dirac', 'width': width},
e_km[np.newaxis],
weight_k=weight_k,
nelectrons=ne)[0][0]
m_v = np.zeros(3, complex)
for ik in range(Nk):
ut0_nG = np.array([calc.wfs.get_wave_function_array(n, ik, 0)
for n in range(nbands)])
ut1_nG = np.array([calc.wfs.get_wave_function_array(n, ik, 1)
for n in range(nbands)])
mocc = np.where(f_km[ik] * Nk - 1.0e-6 < 0.0)[0][0]
for m in range(mocc + 1):
f = f_km[ik, m]
ut0_G = np.dot(v_knm[ik][::2, m], np.swapaxes(ut0_nG, 0, 2))
ut1_G = np.dot(v_knm[ik][1::2, m], np.swapaxes(ut1_nG, 0, 2))
ut_sG = np.array([ut0_G, ut1_G])
mx_G = np.zeros(np.shape(ut0_G), complex)
my_G = np.zeros(np.shape(ut0_G), complex)
mz_G = np.zeros(np.shape(ut0_G), complex)
for s1 in range(2):
for s2 in range(2):
mx_G += ut_sG[s1].conj() * sx_ss[s1, s2] * ut_sG[s2]
my_G += ut_sG[s1].conj() * sy_ss[s1, s2] * ut_sG[s2]
mz_G += ut_sG[s1].conj() * sz_ss[s1, s2] * ut_sG[s2]
m_v += calc.wfs.gd.integrate(np.array([mx_G, my_G, mz_G])) * f
m_av = []
for a in range(len(calc.atoms)):
N0_p = calc.density.setups[a].N0_p.copy()
N0_ij = unpack(N0_p)
Dx_ij = np.zeros_like(N0_ij, complex)
Dy_ij = np.zeros_like(N0_ij, complex)
Dz_ij = np.zeros_like(N0_ij, complex)
Delta_p = calc.density.setups[a].Delta_pL[:, 0].copy()
Delta_ij = unpack(Delta_p)
for ik in range(Nk):
P_ami = ... # get_spinorbit_projections(calc, ik, v_knm[ik])
P_smi = np.array([P_ami[a][:, ::2], P_ami[a][:, 1::2]])
P_smi = np.dot(C_ss, np.swapaxes(P_smi, 0, 1))
P0_mi = P_smi[0]
P1_mi = P_smi[1]
f_mm = np.diag(f_km[ik])
Dx_ij += P0_mi.conj().T.dot(f_mm).dot(P1_mi)
Dx_ij += P1_mi.conj().T.dot(f_mm).dot(P0_mi)
Dy_ij -= 1.0j * P0_mi.conj().T.dot(f_mm).dot(P1_mi)
Dy_ij += 1.0j * P1_mi.conj().T.dot(f_mm).dot(P0_mi)
Dz_ij += P0_mi.conj().T.dot(f_mm).dot(P0_mi)
Dz_ij -= P1_mi.conj().T.dot(f_mm).dot(P1_mi)
mx = np.sum(N0_ij * Dx_ij).real
my = np.sum(N0_ij * Dy_ij).real
mz = np.sum(N0_ij * Dz_ij).real
m_av.append([mx, my, mz])
m_v[0] += np.sum(Delta_ij * Dx_ij) * (4 * np.pi)**0.5
m_v[1] += np.sum(Delta_ij * Dy_ij) * (4 * np.pi)**0.5
m_v[2] += np.sum(Delta_ij * Dz_ij) * (4 * np.pi)**0.5
return m_v.real, m_av
def get_parity_eigenvalues(calc, ik=0, spin_orbit=False, bands=None, Nv=None,
inversion_center=[0, 0, 0], deg_tol=1.0e-6,
tol=1.0e-6):
"""Calculates parity eigenvalues at time-reversal invariant k-points.
Only works in plane wave mode.
"""
assert len(calc.get_bz_k_points()) == len(calc.get_ibz_k_points())
kpt_c = calc.get_ibz_k_points()[ik]
if Nv is None:
Nv = int(calc.get_number_of_electrons() / 2)
if bands is None:
bands = range(calc.get_number_of_bands())
# Find degenerate subspaces
eig_n = calc.get_eigenvalues(kpt=ik)[bands]
e_in = []
used_n = []
for n1, e1 in enumerate(eig_n):
if n1 not in used_n:
n_n = []
for n2, e2 in enumerate(eig_n):
if np.abs(e1 - e2) < deg_tol:
n_n.append(n2)
used_n.append(n2)
e_in.append(n_n)
print()
print(f' Inversion center at: {inversion_center}')
print(f' Calculating inversion eigenvalues at k = {kpt_c}')
print()
center_v = np.array(inversion_center) / Bohr
G_Gv = calc.wfs.pd.get_reciprocal_vectors(q=ik, add_q=True)
psit_nG = np.array([calc.wfs.kpt_u[ik].psit_nG[n]
for n in bands])
if spin_orbit:
n1 = bands[0]
n2 = bands[-1] + 1
assert (bands == np.arange(n1, n2)).all()
soc = soc_eigenstates(calc, n1=n1, n2=n2)
v_kmn = soc.eigenvectors()
psit0_mG = np.dot(v_kmn[ik, :, ::2], psit_nG)
psit1_mG = np.dot(v_kmn[ik, :, 1::2], psit_nG)
for n in range(len(bands)):
psit_nG[n] /= (np.sum(np.abs(psit_nG[n])**2))**0.5
if spin_orbit:
for n in range(2 * len(bands)):
A = np.sum(np.abs(psit0_mG[n])**2) + np.sum(np.abs(psit1_mG[n])**2)
psit0_mG[n] /= A**0.5
psit1_mG[n] /= A**0.5
P_GG = np.ones((len(G_Gv), len(G_Gv)), float)
for iG, G_v in enumerate(G_Gv):
P_GG[iG] -= ((G_Gv[:] + G_v).round(6)).any(axis=1)
assert (P_GG == P_GG.T).all()
phase_G = np.exp(-2.0j * np.dot(G_Gv, center_v))
p_n = []
print('n P_n')
for n_n in e_in:
if spin_orbit:
# The dimension of parity matrix is doubled with spinorbit
m_m = [2 * n_n[0] + i for i in range(2 * len(n_n))]
Ppsit0_mG = np.dot(P_GG, psit0_mG[m_m].T).T
Ppsit0_mG[:] *= phase_G
Ppsit1_mG = np.dot(P_GG, psit1_mG[m_m].T).T
Ppsit1_mG[:] *= phase_G
P_nn = np.dot(psit0_mG[m_m].conj(), np.array(Ppsit0_mG).T)
P_nn += np.dot(psit1_mG[m_m].conj(), np.array(Ppsit1_mG).T)
else:
Ppsit_nG = np.dot(P_GG, psit_nG[n_n].T).T
Ppsit_nG[:] *= phase_G
P_nn = np.dot(psit_nG[n_n].conj(), np.array(Ppsit_nG).T)
P_eig = np.linalg.eigh(P_nn)[0]
if np.allclose(np.abs(P_eig), 1, tol):
P_n = np.sign(P_eig).astype(int).tolist()
if spin_orbit:
# Only include one of the degenerate pair of eigenvalues
Pm = np.sign(P_eig).tolist().count(-1)
Pp = np.sign(P_eig).tolist().count(1)
P_n = Pm // 2 * [-1] + Pp // 2 * [1]
print(f'{str(n_n)[1:-1]}: {str(P_n)[1:-1]}')
p_n += P_n
else:
print(f' {n_n} are not parity eigenstates')
print(f' P_n: {P_eig}')
print(f' e_n: {eig_n[n_n]}')
p_n += [0 for n in n_n]
return np.ravel(p_n)
def get_L_vlmm():
if len(_L_vlmm) == 3:
return _L_vlmm
s = np.array([[0.0]])
p = np.zeros((3, 3), complex) # y, z, x
p[0, 1] = -1.0j
p[1, 0] = 1.0j
d = np.zeros((5, 5), complex) # xy, yz, z^2, xz, x^2-y^2
d[0, 3] = -1.0j
d[3, 0] = 1.0j
d[1, 2] = -3**0.5 * 1.0j
d[2, 1] = 3**0.5 * 1.0j
d[1, 4] = -1.0j
d[4, 1] = 1.0j
_L_vlmm.append([s, p, d])
p = np.zeros((3, 3), complex) # y, z, x
p[1, 2] = -1.0j
p[2, 1] = 1.0j
d = np.zeros((5, 5), complex) # xy, yz, z^2, xz, x^2-y^2
d[0, 1] = 1.0j
d[1, 0] = -1.0j
d[2, 3] = -3**0.5 * 1.0j
d[3, 2] = 3**0.5 * 1.0j
d[3, 4] = -1.0j
d[4, 3] = 1.0j
_L_vlmm.append([s, p, d])
p = np.zeros((3, 3), complex) # y, z, x
p[0, 2] = 1.0j
p[2, 0] = -1.0j
d = np.zeros((5, 5), complex) # xy, yz, z^2, xz, x^2-y^2
d[0, 4] = 2.0j
d[4, 0] = -2.0j
d[1, 3] = 1.0j
d[3, 1] = -1.0j
_L_vlmm.append([s, p, d])
return _L_vlmm