import numpy as np
from .calculator import Calculator
from ..formula import covariant as frml
from ..result import KBandResult, TABresult
# The base classes for Tabulating
# particular calculators are below
[docs]
class Tabulator(Calculator):
def __init__(self, Formula, ibands=None, kwargs_formula=None, **kwargs):
self.Formula = Formula
self.ibands = np.array(ibands) if (ibands is not None) else None
self.kwargs_formula = kwargs_formula if kwargs_formula is not None else {}
super().__init__(**kwargs)
def __call__(self, data_K):
formula = self.Formula(data_K, **self.kwargs_formula)
nk = data_K.nk
NB = data_K.num_wann
ibands = self.ibands
if ibands is None:
ibands = np.arange(NB)
band_groups = data_K.get_bands_in_range_groups(
-np.inf, np.inf, degen_thresh=self.degen_thresh, degen_Kramers=self.degen_Kramers, sea=False)
# bands_groups is a digtionary (ib1,ib2):E
# now select only the needed groups
band_groups = [
[n for n in groups.keys() if np.any((ibands >= n[0]) * (ibands < n[1]))] for groups in band_groups
] # select only the needed groups
group = [[] for _ in range(nk)]
for ik in range(nk):
for ib in ibands:
for n in band_groups[ik]:
if n[1] > ib >= n[0]:
group[ik].append(n)
break
rslt = np.zeros((nk, len(ibands)) + (3,) * formula.ndim)
for ik in range(nk):
values = {}
for n in band_groups[ik]:
inn = np.arange(n[0], n[1])
out = np.concatenate((np.arange(0, n[0]), np.arange(n[1], NB)))
values[n] = formula.trace(ik, inn, out) / (n[1] - n[0])
for ib, b in enumerate(ibands):
rslt[ik, ib] = values[group[ik][ib]]
return KBandResult(rslt, transformTR=formula.transformTR, transformInv=formula.transformInv)
[docs]
class TabulatorAll(Calculator):
"""
TabulatorAll - a pack of all k-resolved calculators (Tabulators)
"""
def __init__(self, tabulators, ibands=None, mode="grid", save_mode="bin", print_comment=False):
""" tabulators - dict 'key':tabulator
one of them should be "Energy" """
self.tabulators = tabulators
mode = mode.lower()
assert mode in ("grid", "path")
self.mode = mode
self.save_mode = save_mode
if "Energy" not in self.tabulators.keys():
self.tabulators["Energy"] = Energy()
if ibands is not None:
ibands = np.array(ibands)
for k, v in self.tabulators.items():
if hasattr(v, 'ibands'):
if v.ibands is not None:
try:
assert len(v.ibands) == len(ibands)
assert np.all(v.ibands == ibands)
except AssertionError:
raise ValueError(
f"tabulator {k} has ibands={v.ibands} not equal to ibands={ibands} required in TabulatorAll")
else:
v.ibands = ibands
self.comment = (self.__doc__ + "\n Includes the following tabulators : \n" + "-" * 50 + "\n" + "\n".join(
f""" "{key}" : {val} : {val.comment}\n""" for key, val in self.tabulators.items()) +
"\n" + "-" * 50 + "\n")
self._set_comment(print_comment)
def __call__(self, data_K):
return TABresult(
kpoints=data_K.kpoints_all.copy(),
mode=self.mode,
recip_lattice=data_K.system.recip_lattice,
save_mode=self.save_mode,
results={k: v(data_K)
for k, v in self.tabulators.items()})
@property
def allow_path(self):
return self.mode == "path"
@property
def allow_grid(self):
return self.mode == "grid"
###############################################
###############################################
###############################################
###############################################
#### ######
#### Implemented calculators ######
#### ######
###############################################
###############################################
###############################################
###############################################
[docs]
class Energy(Tabulator):
def __init__(self, **kwargs):
super().__init__(frml.Eavln, **kwargs)
[docs]
class Velocity(Tabulator):
def __init__(self, **kwargs):
super().__init__(frml.Velocity, **kwargs)
[docs]
class InvMass(Tabulator):
def __init__(self, **kwargs):
super().__init__(frml.InvMass, **kwargs)
[docs]
class Der3E(Tabulator):
def __init__(self, **kwargs):
super().__init__(frml.Der3E, **kwargs)
[docs]
class BerryCurvature(Tabulator):
def __init__(self, **kwargs):
super().__init__(frml.Omega, **kwargs)
[docs]
class DerBerryCurvature(Tabulator):
r"Derivative of Berry curvature :math:`X_{ab}\partial_b\Omega_a`"
def __init__(self, **kwargs):
super().__init__(frml.DerOmega, **kwargs)
[docs]
class Der2BerryCurvature(Tabulator):
r"Second Derivative of Berry curvature :math:`X_{ab}\partial_bc\Omega_a`"
def __init__(self, **kwargs):
super().__init__(frml.Der2Omega, **kwargs)
[docs]
class Spin(Tabulator):
r""" Spin expectation :math:` \langle u | \mathbf{\sigma} | u \rangle`"""
def __init__(self, **kwargs):
super().__init__(frml.Spin, **kwargs)
[docs]
class DerSpin(Tabulator):
def __init__(self, **kwargs):
super().__init__(frml.DerSpin, **kwargs)
[docs]
class Der2Spin(Tabulator):
def __init__(self, **kwargs):
super().__init__(frml.Der2Spin, **kwargs)
[docs]
class OrbitalMoment(Tabulator):
def __init__(self, **kwargs):
super().__init__(frml.morb, **kwargs)
[docs]
class DerOrbitalMoment(Tabulator):
def __init__(self, **kwargs):
super().__init__(frml.Dermorb, **kwargs)
[docs]
class DerOrbitalMoment_test(Tabulator):
def __init__(self, **kwargs):
super().__init__(frml.DerMorb_test, **kwargs)
[docs]
class Der2OrbitalMoment(Tabulator):
def __init__(self, **kwargs):
super().__init__(frml.Der2morb, **kwargs)
[docs]
class SpinBerry(Tabulator):
def __init__(self, **kwargs):
super().__init__(frml.SpinOmega, **kwargs)