from abc import ABC, abstractmethod
from typing import Union, Optional
import numpy as np
import sympy as sp
from .models import Kmodel
from ._common import other_than, invert, _to_eigenbasis
__all__ = [
"Identity",
"Velocity",
"SpinOperator",
"SpinDensity",
"SpinCurrent",
"OrbitalDensity",
"OrbitalCurrent",
"MagneticOctupoleDensity",
"MagneticOctupoleCurrent",
"QuadripoleDensity",
"QuadripoleCurrent",
]
class KspaceOpera(ABC):
def __init__(self, model: Kmodel):
self.model = model
self.dim = model.dim
self.norbs = model.norbs
@abstractmethod
def __call__(self, k_args: list[np.ndarray], en: np.ndarray, psi: np.ndarray):
pass
[docs]
class Identity(KspaceOpera):
r"""Computes the identity operator :math:`\langle \psi_m | \psi_n \rangle`
Args:
model: An instance of :class:`~py4mulas.models.Kmodel`
"""
def __init__(self, model: Kmodel):
super().__init__(model)
[docs]
def __call__(self, k_args: list[np.ndarray], en: np.ndarray, psi: np.ndarray):
return np.eye(self.norbs)[None, ...]
[docs]
class Velocity(KspaceOpera):
r"""Computes velocity operator along a given :attr:`direction`
Args:
model: An instance of :class:`~py4mulas.models.Kmodel`
direction: The direction along which velocity is
computed.
"""
def __init__(self, model: Kmodel, direction: str):
super().__init__(model)
self.direction = direction
[docs]
def __call__(self, k_args: list[np.ndarray], en: np.ndarray, psi: np.ndarray):
r"""Computes the velocity matrix along :attr:`direction`.
If a list or a tuple of floats is given a linear combination
of velocities is considered. Then the elements of :attr:`direction`
give the respective coefficients along :math:`x`, :math:`y`, :math:`z`.
"""
direction = self.direction
velocities = self.model.velocities
if isinstance(direction, str):
return _to_eigenbasis(psi, velocities[direction](*k_args))
elif isinstance(direction, (list, tuple)):
combination = 0
for axis, (coord, velo) in enumerate(velocities.items()):
combination += direction[axis] * velo(*k_args)
return _to_eigenbasis(psi, combination)
else:
raise TypeError(
"direction should be a string from 'xyz', a list or tuple of floats"
)
class OrbitalOperator(KspaceOpera):
def __init__(
self,
model: Kmodel,
direction: str = None,
gamma: Union[str, np.ndarray, sp.MatrixBase] = "z",
):
super().__init__(model)
if isinstance(gamma, sp.MatrixBase):
gamma = np.array(gamma.tolist(), dtype=complex)
self.gamma = gamma
if direction is not None:
self.direction = model.velocities[direction]
if isinstance(gamma, str):
assert gamma in "xyz", "gamma should be x, y or z"
direction1, direction2 = other_than(gamma)
# these are the directions in the plane perpendicular to gamma
self.vel1 = model.velocities[direction1]
self.vel2 = model.velocities[direction2]
def density(self, k_args: list[np.ndarray], en: np.ndarray, psi: np.ndarray):
r"""Computes the eigen-basis projected orbital density.
Args:
k_args: Momentum arguments to pass to the velocity operator
en: Model's eigenenergies
psi: Model's eigenvectors
"""
if isinstance(self.gamma, np.ndarray):
# then it's similar to SpinDensity computation
return SpinOperator(self.model, gamma=self.gamma, direction=None)(
k_args, en, psi
)
v1 = self.vel1(*k_args)
v2 = self.vel2(*k_args)
return orbital_matrix(en, psi, v1, v2, norbs=self.norbs)
def __call__(self, k_args: list[np.ndarray], en: np.ndarray, psi: np.ndarray):
r"""Computes the eigen basis projected orbital current
flowing along a given direction :attr:`direction`.
Args:
k_args: Momentum arguments to pass to velocity operator
en: Model's eigenenergies
psi: Model's eigenvectors
"""
if self.direction is None:
return self.density(k_args, en, psi)
velo = self.direction(*k_args)
v_left = _to_eigenbasis(psi, velo)
l_gamma = self.density(k_args, en, psi)
return 0.5 * (v_left @ l_gamma + l_gamma @ v_left)
[docs]
class SpinOperator(KspaceOpera):
"""Base class for SpinDensity and SpinCurrent operators.
Attributes:
model: An instance of :class:`~py4mulas.models.Kmodel`
direction: The direction of flow for the current.
It should be specified for SpinCurrent.
pauli: Angular momentum polarization. A pauli or an orbital matrix.
"""
def __init__(
self,
model: Kmodel,
direction: str = None,
pauli: Union[np.ndarray, sp.MatrixBase] = None,
):
super().__init__(model)
if isinstance(pauli, sp.MatrixBase):
pauli = np.array(pauli.tolist(), dtype=complex)
self.pauli = pauli
self.direction = direction
if direction is not None:
self.velo = model.velocities[direction]
[docs]
def band_resolved(
self, n: Union[int, list, tuple] = 0, m: Optional[int] = None
) -> np.ndarray:
"""For computing a SpinOperator expectation value for a specific band or a list of bands.
If `m` is given the expectation value between `n` and `m` is computed.
Args:
n: A band index or a list of bands.
m: A band index for the ket Bloch vector.
Returns:
the expectation values over py4mulas
"""
dim = self.dim
pauli = self.pauli
model = self.model
if n is None:
raise ValueError(
"a band index, list or tuple of band indeces must be given"
)
else:
if not isinstance(n, (int, list, tuple)):
raise ValueError("n should be an integer, list or tuple")
# make summation symbols to be used in np.einsum
contractions = self._make_contractions(n)
k_vectors = np.asarray(model.k_vectors)
k_args = [k_vectors[:, d] for d in range(dim)]
if self.direction is None:
opera_matrix = pauli
else:
velo = self.velo(*k_args)
opera_matrix = 0.5 * (velo @ pauli + pauli @ velo)
# a rediagonalization is redendent but should be fine here
# as this method is designed for visualization purposes
hk = model.H(*k_args)
_, psi = np.linalg.eigh(hk)
if isinstance(n, int):
if m is None:
m = n
psi_n = psi[:, :, n]
psi_m = psi[:, :, m]
return np.einsum(contractions, psi_n.conj(), opera_matrix, psi_m)
assert m is None
n_list = list(n)
psi_list = psi[:, :, n_list]
return np.einsum(contractions, psi_list.conj(), opera_matrix, psi_list)
def _make_contractions(self, n: Union[int, list, tuple]) -> str:
"""Make sum symbols for einsum(' ', psi.conj(), opera, psi)"""
if isinstance(n, int):
if self.direction is None:
# compute <psi_n|O|psi_m>
# O has shape (norbs, norbs)
contractions = "bp,pq,bq -> b"
else:
# O is (k, norbs, norbs)
contractions = "bp,bpq,bq -> b"
elif isinstance(n, (list, tuple)):
# a sum over the list index should be performed
# \sum_{i in list} <psi_i|O|psi_i>
if self.direction is None:
contractions = "bpi,pq,bqi -> b"
else:
contractions = "bpi,bpq,bqi -> b"
else:
raise ValueError("band index should be an integer, list or tuple")
return contractions
[docs]
def __call__(self, k_args: list[np.ndarray], en: np.ndarray, psi: np.ndarray):
r"""Calculates the current polarized along :attr:`pauli` and flowing
along :attr:`direction`.
Args:
pauli: Spin matrix of the same shape as the hamiltonian.
direction: Current flow direction.
"""
pauli = self.pauli # (n, n)
if self.direction is None:
return np.einsum("bpi,pq,bqj->bij", psi.conj(), pauli, psi)
velo = self.velo(*k_args)
current_operator = 0.5 * (velo @ pauli + pauli @ velo)
return _to_eigenbasis(psi, current_operator)
class MagneticOctupole(KspaceOpera):
def __init__(
self,
model: Kmodel,
a: Union[np.ndarray, sp.MatrixBase, str],
b: Union[np.ndarray, sp.MatrixBase, str],
c: Union[np.ndarray, sp.MatrixBase],
direction: str = None,
):
super().__init__(model)
self.a = a
self.b = b
self.c = c
if isinstance(a, sp.MatrixBase):
self.a = np.array(a.tolist(), dtype=complex)
if isinstance(b, sp.MatrixBase):
self.b = np.array(b.tolist(), dtype=complex)
if isinstance(c, sp.MatrixBase):
self.c = np.array(c.tolist(), dtype=complex)
self.norbs = model.norbs
self.direction = direction
if direction is not None:
self.velo = model.velocities[direction]
assert isinstance(c, (np.ndarray, sp.MatrixBase)), "c should be an array"
if isinstance(a, str):
assert isinstance(b, str)
self.a_vel1, self.a_vel2 = self.perpendicular_velocities(a)
self.b_vel1, self.b_vel2 = self.perpendicular_velocities(b)
else:
if not isinstance(self.a, (np.ndarray, sp.MatrixBase)):
raise ValueError("a and b can only be str or arrays")
if self.a.shape != (self.norbs, self.norbs):
raise ValueError(
f"expected a shape ({self.norbs}, {self.norbs}) got {self.a.shape}"
)
self.model = model
def perpendicular_velocities(self, a):
a_direction1, a_direction2 = other_than(a)
# these are the directions in the plane perpendicular to a
a_vel1 = self.model.velocities[a_direction1]
a_vel2 = self.model.velocities[a_direction2]
return a_vel1, a_vel2
def density(
self, k_args: list[np.ndarray], en: np.ndarray, psi: np.ndarray
) -> np.ndarray:
r"""Computes the octupole density projected in eigenbasis.
Args:
k_args: Momentum args to pass to velocity operators
en: Eigen eneries matrix
psi: Eigenvectors
"""
if isinstance(self.a, str):
a_vel1 = self.a_vel1(*k_args)
a_vel2 = self.a_vel2(*k_args)
b_vel1 = self.b_vel1(*k_args)
b_vel2 = self.b_vel2(*k_args)
La = orbital_matrix(en, psi, a_vel1, a_vel2, norbs=self.norbs)
Lb = orbital_matrix(en, psi, b_vel1, b_vel2, norbs=self.norbs)
if self.c is not None:
# we shall transform spin operator into eigen basis representation
Spin = np.einsum(
"bpi,pq,bqj->bij", psi.conj(), self.c, psi, optimize=True
)
else:
La = self.a
Lb = self.b
if self.c is not None:
Spin = self.c
O_ab_c = La @ Lb + Lb @ La
if self.c is not None:
O_ab_c @= Spin # units of hbar^{-2}
return O_ab_c
def __call__(self, k_args, en, psi):
if self.direction is None:
return self.density(k_args, en, psi)
velo = self.velo(*k_args)
J_abc = self.density(k_args, en, psi)
if isinstance(self.a, np.ndarray):
v_J_abc = 0.5 * (velo @ J_abc + J_abc @ velo)
return _to_eigenbasis(psi, v_J_abc)
velo = _to_eigenbasis(psi, velo)
v_J_abc = 0.5 * (velo @ J_abc + J_abc @ velo)
return v_J_abc
def orbital_matrix(
en: np.ndarray, psi: np.ndarray, v1: np.ndarray, v2: np.ndarray, norbs: int
) -> np.ndarray:
r"""Computes the orbital matrix :math:`L` whose component
is along the direction erpendicular to `v1` and `v2`.
Args:
en: Energy eigen values
psi: Eigen states
v1: Velocity matrix along direction1
v2: Velocity matrix along direction2
norbs: Number of orbitals
"""
v1_mat = _to_eigenbasis(psi, v1)
v2_mat = _to_eigenbasis(psi, v2)
idx = np.arange(norbs)
v1_mat[:, idx, idx] = 0
v2_mat[:, idx, idx] = 0
energy_factor = invert(en[..., None, :] - en[..., :, None])
v1_tilde = v1_mat * energy_factor
v2_tilde = v2_mat * energy_factor
l_gamma = v1_tilde @ v2_mat
l_gamma -= v2_tilde @ v1_mat
l_gamma += v2_mat @ v1_tilde
l_gamma -= v1_mat @ v2_tilde
# units [e / g_L \mu_B] or [2 m / (g_L \hbar)]
return (-1j / 4.0) * l_gamma
[docs]
class SpinCurrent(SpinOperator):
"""Cumputes spin current operator
Attributes:
model: An instance of :class:`~py4mulas.models.Kmodel`
direction: The current direction
pauli: The polarization direction of the spin current
"""
def __init__(self, model: Kmodel, pauli: np.ndarray = None, direction: str = None):
super().__init__(model, pauli=pauli, direction=direction)
[docs]
class SpinDensity(SpinOperator):
"""Cumputes spin density operator
Attributes:
model: An instance of :class:`~py4mulas.models.Kmodel`
pauli: The spin matrix along the polarization direction
"""
def __init__(self, model: Kmodel, pauli: np.ndarray = None):
super().__init__(model, pauli=pauli, direction=None)
[docs]
class OrbitalCurrent(OrbitalOperator):
r"""Cumputes orbital current operator polarized along :attr:`gamma`
flowing in :attr:`direction`
Attributes:
model: An instance of :class:`~py4mulas.models.Kmodel`
direction: The current flow direction
gamma: The polarization direction of the orbital current
"""
def __init__(self, model: Kmodel, direction: str, gamma: Union[np.ndarray, str]):
super().__init__(model, direction=direction, gamma=gamma)
[docs]
class OrbitalDensity(OrbitalOperator):
"""Cumputes orbital density operator
Attributes:
model: An instance of :class:`~py4mulas.models.Kmodel`
gamma: The polarization direction of the orbital angular momentum
"""
def __init__(self, model: Kmodel, gamma=Union[np.ndarray, str]):
super().__init__(model, direction=None, gamma=gamma)
[docs]
class MagneticOctupoleCurrent(MagneticOctupole):
r"""Cumputes octupole current operator
:math:`\frac{1}{2}\{v, O_{ab}^c\}`
with
:math:`v` being the velocity operator along the flow ``direction`` and
:math:`O_{ab}^c=\{L_a, L_b\} S_c`.
:math:`L_a` is the orbital operator polarized along :math:`a`
and :math:`S_c` is the spin matrix written in the same basis as :math:`L`.
See. Han et al [Phys. Rev. Lett. 135, 076705 (2025)]
Attributes:
model: An instance of :class:`~py4mulas.models.Kmodel`
direction: The current flow direction
a: An orbital matrix or a polarization direction.
In the atomic center approximation, we have the orbital matrix. This works perfectly.
If a string is given then :math:`L_a` will be constructed
from itenerent orbital moment expression of modern theory of orbital magnetism.
The last approach is not fully tested. This should be only used when really intended.
b: An orbital matrix or a polarization direction.
In the atomic center approximation, we have the orbital matrix. This works perfectly.
If a string is given then :math:`L_a` will be constructed
from itenerent orbital moment expression of modern theory of orbital magnetism.
The last approach is not fully tested. This should be only used when really intended.
c: A spin angular momentum matrix.
"""
def __init__(
self,
model: Kmodel,
a: Union[np.ndarray, sp.MatrixBase, str],
b: Union[np.ndarray, sp.MatrixBase, str],
c: Union[np.ndarray, sp.MatrixBase],
direction: str,
):
super().__init__(model, a, b, c, direction=direction)
[docs]
class MagneticOctupoleDensity(MagneticOctupole):
r"""Cumputes otupole density operator
:math:`O_{ab}^c=\{L_a, L_b\} S_c`.
:math:`L_a` being the orbital operator polarized along :math:`a`
Attributes:
model: An instance of :class:`~py4mulas.models.Kmodel`
a: An orbital matrix or a polarization direction.
In the atomic center approximation, we have the orbital matrix. This works perfectly.
If a string is given then :math:`L_a` will be constructed
from itenerent orbital moment expression of modern theory of orbital magnetism.
The last approach is not fully tested. This should be only used when really intended.
b: An orbital matrix or a polarization direction.
In the atomic center approximation, we have the orbital matrix. This works perfectly.
If a string is given then :math:`L_b` will be constructed
from itenerent orbital moment expression of modern theory of orbital magnetism.
The last approach is not fully tested. This should be only used when really intended.
c: A spin angular momentum matrix.
"""
def __init__(
self,
model: Kmodel,
a: Union[np.ndarray, str],
b: Union[np.ndarray, str],
c: Union[np.ndarray, sp.MatrixBase],
):
super().__init__(model, a, b, c, direction=None)
[docs]
class QuadripoleCurrent(MagneticOctupole):
r"""Cumputes the quadripole current operator
:math:`\frac{1}{2}\{v, Q_{ab}\}`
with
:math:`v` the velocity operator along the flow ``direction`` and
:math:`Q_{ab}=\{L_a,L_b\}`, and
:math:`L_a` being the orbital operator polarized along :math:`a`
Attributes:
model: An instance of :class:`~py4mulas.models.Kmodel`
direction: The current flow direction
a: An orbital matrix or a polarization direction.
In the atomic center approximation, we have the orbital matrix. This works perfectly.
If a string is given then :math:`L_a` will be constructed
from itenerent orbital moment expression of modern theory of orbital magnetism.
The last approach is not fully tested. This should be only used when really intended.
b: An orbital matrix or a polarization direction.
In the atomic center approximation, we have the orbital matrix. This works perfectly.
If a string is given then :math:`L_b` will be constructed
from itenerent orbital moment expression of modern theory of orbital magnetism.
The last approach is not fully tested. This should be only used when really intended.
"""
def __init__(
self,
model: Kmodel,
a: Union[np.ndarray, sp.MatrixBase, str],
b: Union[np.ndarray, sp.MatrixBase, str],
direction: str,
):
super().__init__(model, a, b, None, direction=direction)
[docs]
class QuadripoleDensity(MagneticOctupole):
r"""Cumputes quadripole density operator
:math:`Q_{ab}=\{L_a,L_b\}`
with :math:`L_a` being the orbital operator polarized along :math:`a`
Attributes:
model: An instance of :class:`~py4mulas.models.Kmodel`
a: An orbital matrix or a polarization direction.
In the atomic center approximation, we have the orbital matrix. This works perfectly.
If a string is given then :math:`L_a` will be constructed
from itenerent orbital moment expression of modern theory of orbital magnetism.
The last approach is not fully tested. This should be only used when really intended.
b: An orbital matrix or a polarization direction.
In the atomic center approximation, we have the orbital matrix. This works perfectly.
If a string is given then :math:`L_a` will be constructed
from itenerent orbital moment expression of modern theory of orbital magnetism.
The last approach is not fully tested. This should be only used when really intended.
"""
def __init__(
self, model: Kmodel, a: Union[np.ndarray, str], b: Union[np.ndarray, str]
):
super().__init__(model, a, b, None, direction=None)