Source code for py4mulas.formulas
from abc import ABC, abstractmethod
from collections.abc import Iterable
import inspect
import os
from typing import Optional, Union, Callable
import numpy as np
from .mu_kernels import MuKernel
from .models import Kmodel
from .utils import KspaceOptions
from ._common import compute_dk
__all__ = ["KuboFormula", "batcher"]
class _Prebuilder(ABC):
def __init__(self) -> None:
# initialize the store
self.have_store = False
self._store_signature = None
def _get_store_signature(self, eta: float) -> float:
"""Ensures that changing eta rebuilds everything
in case kernel depends on it
"""
return float(eta)
@abstractmethod
def _build_store(self, eta):
pass
def _ensure_bulid_store(self, eta: Optional[float] = None) -> None:
if eta is None:
if not (self.have_store):
# discard eta
self._build_store(eta)
else:
sig = self._get_store_signature(eta)
if not (self.have_store) or (self._store_signature != sig):
self._build_store(eta)
class ArrayStore:
def __init__(
self,
shape: tuple,
dtype=complex,
memmap: bool = False,
name: str = "opera_kernel",
):
if not memmap:
self.store = np.empty(shape, dtype=complex)
else:
obj_id = id(self)
temp_dir = os.path.join(os.getcwd(), "_py4temps")
os.makedirs(temp_dir, exist_ok=True)
self.file_name = "./_py4temps/" + f"{name}_{obj_id}.dat"
self.store = np.memmap(self.file_name, dtype=dtype, mode="w+", shape=shape)
def del_file(self):
try:
if os.path.exists(self.file_name):
os.remove(self.file_name)
except Exception:
pass
class OkernelFactory:
def __init__(
self,
kernel_shapes: list[tuple],
memmap: bool = False,
name: str = "opera_kernel",
) -> None:
self.kernels = [
ArrayStore(shape=shape_i, dtype=complex, memmap=memmap, name=name)
for shape_i in kernel_shapes
]
def del_files(self):
for kernel in self.kernels:
kernel.del_file()
class _MainKernel:
"""Computes k_args, eta dependent kernel to be optionally stored.
Args:
opera_kernel: A collection of opera_kernels to store.
Every kernel element is expected to have the signature (k_args, energy, psi, eta).
H: The lambdified hamiltonian of the model.
"""
def __init__(self, opera_kernel: Union[Callable, list[Callable]], H: Callable):
try:
name = opera_kernel.__name__
except AttributeError:
name = None
if name == "_opera_kernel":
self.is__opera_kernel = True
self.opera_kernel = opera_kernel
self.length = 1 # None
else:
self.is__opera_kernel = False
self.opera_kernel = (
opera_kernel if isinstance(opera_kernel, list) else [opera_kernel]
)
self.length = len(self.opera_kernel)
self.H = H
def __call__(
self, k_args: list[np.ndarray], eta: float
) -> tuple[list[np.ndarray], np.ndarray]:
hk = self.H(*k_args)
energy, psi = np.linalg.eigh(hk)
if not self.is__opera_kernel:
return [
kernel(k_args, energy, psi, eta) for kernel in self.opera_kernel
], energy
kernels = self.opera_kernel(k_args, energy, psi, eta)
if not isinstance(kernels, list):
kernels = [kernels]
return kernels, energy
[docs]
class KuboFormula(_Prebuilder):
r"""Computes an arbitrary response formula
.. math::
\sigma = \sum_k\sum_{mn} K_{mn}(k) O_{mn}(k)
with :math:`K_{mn}` being the matrix elements of the energy kernel K
which may depend on transport properties such as :math:`\mu`, :math:`T`, :math:`\eta`
and :math:`E` but should not nvolve :math:`\psi`.
The memeber :math:`O_{mn}` is the operator kernel which may depend on :math:`\psi`.
It is assumed to depend on py4mulas operators.
This can be used to implement an arbitrary formula with tenorial contractions.
The contractions einsum (subscripts) should be provided as an argument.
This separation enables to store the operator kernel for the whole kspace.
So varying transport properties becomes quite cheap.
In cases, where these kernels have further symmetries or simplifications
we should provide the shape of the matrix :math:`O` through ``kernel_norbs``.
This is exploited in :class:`py4mulas.responses.Kubo`.
Attributes:
model: An instance of :class:`~py4mulas.models.Kmodel`
kspace_options: A dictionary specifying chunk_size and precomp.
If chunk_size is not given, a default is used.
Note that if precomp is True, the energy and temperature scans are momentum space free.
If eta is to be changed, then set precomp = False, as the kernels are eta-dependent.
contractions: Explicit np.einsum subscripts to be used for the computation.
mu_kernel: Energy kernel.
If this is not given _mu_kernel method should be implemented.
opera_kernel: Operator kernel to be optionally stored.
If this is not given the _opera_kernel method should be implemented.
If opera_kernel returns a list of kernel arrays, their length should be equal to mu_kernel length.
Possibly, each of these opera_kernels can take a list of mu_kernels. In this case mu_kernel
can be provided as a list of lists. Besides len(mu_kernell) should always equal
the length of returned opera_kernels.
Note:
Enabling the opera_kernel to be a list is mainly for making the computation of these kernels
more performent. For instance one can avoid rediagonalization of the hamiltonian.
Example:
>>> class OperaKernel:
>>> def __init__(self, alpha, beta, **kwargs):
>>> self.alpha = alpha
>>> self.beta = beta
>>> def __call__(self, k_args, energy, psi, eta):
>>> alpha = self.alpha(k_args, energy, psi)
>>> beta = self.beta(k_args, energy, psi)
>>> kernel = beta * np.swapaxes(alpha, 1, 2)
>>> return kernel
>>> opera_kernel = OperaKernel(alpha, beta)
>>> mu_kernel = [KuboKernel('inter_band'), KuboKernel('intra_band')]
>>> kspace_options = dict(chunk_size=1000, precomp=True)
>>> formula = KuboFormula(kmodel, kspace_options=kspace_options, mu_kernel=mu_kernel, opera_kernel=opera_kernel)
>>> response = formula(mu=0, temperature=0, eta=0, k_resolved=False)
"""
def __init__(
self,
kmodel: Kmodel,
kspace_options: Optional[dict] = None,
contractions: Optional[Union[list[str], str]] = None,
opera_kernel: Optional[Callable] = None,
mu_kernel: Optional[
Union[list[Callable], Callable, list[list[Callable]]]
] = None,
) -> None:
super().__init__()
# extract contractions with compatibility check with mu_kernels
self.contractions = _extract_contractions(
contractions, mu_kernel, k_resolved=False
)
self.kmodel = kmodel
if opera_kernel is None:
opera_kernel = self._opera_kernel
self.opera_kernel = _MainKernel(opera_kernel, H=lambda *args: self.H(*args))
self.mu_kernel = mu_kernel
self.dim = kmodel.dim
self.norbs = kmodel.norbs
self._k_vectors = np.asarray(kmodel.k_vectors, dtype=float)
self.bounds = kmodel.bounds
self._chunk_size = _read_from_data(kspace_options, "chunk_size")
self.precomp = _read_from_data(kspace_options, "precomp")
self.memmap = _read_from_data(kspace_options, "memmap")
self.k_prefactor = (2 * np.pi) ** (1 - self.dim)
self.kernel_norbs = _get_kernel_norbs(self.contractions, self.norbs)
self.Ek = None
self.factory = None
@property
def H(self):
return self.kmodel.H
@property
def k_vectors(self):
return self._k_vectors
@k_vectors.setter
def k_vectors(self, vectors):
self._k_vectors = vectors
# reinitialize the store once k_vectors are altered
self.Ek = None
self.factory = None
self.have_store = False
@property
def num_k(self):
return self.k_vectors.shape[0]
@property
def e_shape(self):
return (self.num_k, self.norbs)
@property
def kernel_shape(self):
return [(self.num_k,) + norbs_i for norbs_i in self.kernel_norbs]
@property
def chunk_size(self):
if self._chunk_size is None:
size = self.num_k
else:
size = self._chunk_size
return size
@property
def _dk(self):
return compute_dk(self.k_vectors, self.dim)
@property
def prefactor(self):
try:
return self.k_prefactor * self._dk
except IndexError:
# single k_vectors
return self.k_prefactor
def _build_store(self, eta):
"""Precompute and store operator kernel with energies."""
k_vectors = self.k_vectors
num_k = self.num_k
dim = self.dim
chunk_size = self.chunk_size
opera_kernel = self.opera_kernel
self.Ek = ArrayStore(
shape=self.e_shape, memmap=self.memmap, dtype=float, name="energy"
)
self.factory = OkernelFactory(
self.kernel_shape, memmap=self.memmap, name="opera_kernel"
)
for start in range(0, num_k, chunk_size):
end = start + chunk_size
k_chunk = k_vectors[start:end] # shape: (chunk_size, dim)
k_args = [k_chunk[:, d] for d in range(dim)]
# TODO: check _opera_kernel taking eta=None is always clean
input_kernel, self.Ek.store[start:end] = opera_kernel(k_args, eta)
if not isinstance(input_kernel, Iterable):
input_kernel = [input_kernel]
for factory_kernel, kernel in zip(self.factory.kernels, input_kernel):
factory_kernel.store[start:end] = kernel
self.have_store = True
if eta is not None:
self._store_signature = self._get_store_signature(eta)
def _mu_kernel(self) -> Union[MuKernel, list[MuKernel]]:
r"""The kernel which can be called with :math:`E`, :math:`\mu` and :math:`T`.
Typically a distribution function.
But can also be :math:`\eta` dependent.
Returns:
A list of :class:`~py4mulas.mu_kernels.KuboKernel`
"""
# kernel K
pass
def _opera_kernel(
self, k_args: np.ndarray, eta: float
) -> tuple[Union[Callable, list[Callable]], np.ndarray]:
"""The main kernal of the formula, which involves the operators product.
Args:
k_args: Momentum arguments
eta: Broadening
Note:
This should return [kernel_i for i in range(n)], Ek
Optionally it may return a single kernel with Ek.
"""
# kernel O
pass
[docs]
def integrand(
self,
*k: Union[list, tuple, np.ndarray],
mu: float,
temperature: float,
eta: float,
) -> complex:
"""unvectorized integrand, for single k evaluation"""
prefactor = self.k_prefactor
self.k_vectors = [k]
return prefactor * self(
mu=mu, temperature=temperature, eta=eta, k_resolved=True
)
[docs]
def __call__(
self,
mu: float = 0.0,
temperature: float = 0.0,
eta: float = 0.0,
k_resolved: bool = False,
) -> complex:
r"""Computes the response fomula at :math:`\mu`, :math:`T` and :math:`\eta`.
Args:
mu: Chemical potential.
temperature: Temperature.
eta: Broadening (the infinitesimal parameter).
k_resolved: Specifies whether we want a summed transport response
or momentum resolved.
Returns:
A numpy arry, if k_resolved is `True`
or a complex number if it is `False`.
"""
mu_kernel = self.mu_kernel
opera_kernel = self.opera_kernel
if mu_kernel is None:
mu_kernel = self._mu_kernel()
if not isinstance(mu_kernel, Iterable):
mu_kernel = [mu_kernel]
if opera_kernel.length != len(mu_kernel):
raise ValueError(
"each operator kernel must have an energy kernel "
"or a list of energy kernels"
)
contractions = self.contractions
if self.precomp:
if _is_eta_independent(mu_kernel):
# then in contrast _opera_kernel should be eta dependent.
# their will be no precomputation advantage for varied eta.
# But is good for varying \mu and T.
self._ensure_bulid_store(eta)
else:
# _opera_kernel is eta independent and precomputation
# can be frozen regardless of eta values.
# this is disigned to enable eta independent precomputation
# when an external _mu_kernel (eta dependent) is provided.
# This is a good compromise as one can still:
# (i) use reduced opera_kernel but with eta dependent precomputation
# which is fast for varying \mu and T,
# slow for variying eta. See above.
# (ii) use unoptimized opera_kernel but with eta independent mu_kernels.
# This should be faster when eta is to be varied.
self._ensure_bulid_store(eta=None)
dim = self.dim
num_k = self.num_k
k_vectors = self.k_vectors
chunk_size = self.chunk_size
prefactor = self.prefactor
if k_resolved:
chunk_size = num_k
prefactor = 1
contractions = _extract_contractions(
self.contractions, mu_kernel, k_resolved=True
)
chunk_sum = 0
for start in range(0, num_k, chunk_size):
end = min(start + chunk_size, num_k)
if self.precomp:
Ek = self.Ek.store[start:end]
kernel_list = self.factory.kernels
else:
k_chunk = k_vectors[start:end] # shape: (chunk_size, dim)
k_args = [k_chunk[:, d] for d in range(dim)]
kernel_list, Ek = opera_kernel(k_args, eta)
if not isinstance(kernel_list, Iterable):
kernel_list = [kernel_list]
for i, opera_kernel in enumerate(kernel_list):
if isinstance(opera_kernel, ArrayStore):
o_kernel = opera_kernel.store[start:end]
else:
o_kernel = opera_kernel
mu_kernel_term = mu_kernel[i]
kernel_contractions = contractions[i]
e_kernel = mu_kernel_term(Ek, mu, temperature, eta)
chunk_sum += np.einsum(
kernel_contractions, o_kernel, e_kernel, optimize=True
)
if self.memmap:
self.Ek.del_file()
self.factory.del_files()
return prefactor * chunk_sum
def _extract_contractions(
contractions: Union[list[str], str],
mu_kernel: Union[Callable, list[Callable]],
k_resolved: bool = False,
):
"""Extracts and verify validity of contractions from mu_kernel or input contractions
Args:
contractions: np.einsum subscripts
mu_kernel: The mu dependent kernels
k_resolved: Specifies whether the contractions are used for k_resolved calculation or not
Raises:
AttributeError: If contractions is None and ``mu_kernels`` not implementing
contractions as an attribute.
Returns:
list[str]: verified contractions
"""
collected_contractions = []
if contractions is not None:
if not isinstance(contractions, (list, tuple)):
contractions = [contractions]
for elem in contractions:
assert isinstance(elem, str), (
"contractions should be strings or list of strings"
)
collected_contractions.append(
_set_contractions(elem, k_resolved=k_resolved)
)
else:
mu_kernel = mu_kernel if isinstance(mu_kernel, Iterable) else [mu_kernel]
for kernel in mu_kernel:
try:
_contractions = _set_contractions(kernel.contractions, k_resolved=k_resolved)
collected_contractions.append(_contractions)
except AttributeError:
raise AttributeError(
"if contractions are not provided mu_kernels "
"should have contractions as an attribute"
)
return _valid(collected_contractions)
def _valid(contractions: list[str]) -> list[str]:
"""Policy vor valid contractions
Args:
contractions: The contractions to be verified
Raises:
ValueError: If contractions are explicite np.einsum subscripts
ValueError: If the first subscript index is not labeled k, for both operands
Returns:
list[str]: :var:`contractions`, if the input subscripts are valid
"""
_contractions = contractions
for i, elem in enumerate(_contractions):
if "->" not in elem:
raise ValueError("contractions should be explicite np.einsum subscripts")
left, right = elem.split(",")
if not left[0] == "k":
raise ValueError(
f"first operand of the {i}th subscripts is expected to be named k but was labeled {left[0]}"
)
if not right[0] == "k":
raise ValueError(
f"second operand of the {i}th subscripts is expected to be named k but was labeled {right[0]}"
)
return contractions
def _set_contractions(contractions: str, k_resolved: bool = False) -> str:
if not k_resolved:
return contractions
if contractions[-1] == ">":
# then we should add k for k_resolved
contractions += "k"
# keep first axis (momentum)
# otherwise the provided contractions should
# appropriately handle the momentum indice ("k").
# See Berry Curvature implementation for instance.
# If k_resolved = False, this distingtion does not really matter,
# as a summation over k is anyway performed.
# TODO: add a contractions input checker!
return contractions
def _get_kernel_norbs(
contractions: list[str], norbs: int
) -> list[tuple]:
if not isinstance(contractions, (tuple, list)):
left, right = contractions.split(",")
kernel_norbs = [(len(left) - 1) * (norbs,)]
else:
kernel_norbs = []
for i, cont in enumerate(contractions):
left, right = cont.split(",")
kernel_norbs.append((len(left) - 1) * (norbs,))
return kernel_norbs
[docs]
def batcher(formula: KuboFormula, data: list[tuple], n: int = 10) -> np.ndarray:
r"""Computes formula, partitioning the kspace into `n` batches, for a full set of `data`.
This can be used when the kspace is too large, to reduce memory load and still
use precomp.
Args:
formula: An instance of :class:`~py4mulas.formulas.KuboFormula`
or :class:`~py4mulas.responses.Kubo`
data: List of tuples in the form of (:math:`\mu`, :math:`T`, :math:`\eta`)
to be passed as arguments to formula.
n: Number of wanted kspace batches
Raises:
ValueError: When precomp is set to `False`
Returns:
An array in the same order of the ``data``
"""
if not formula.precomp:
raise ValueError("partitioner only works for precomptuing formulas.")
k_vectors = np.asarray(formula.kmodel.k_vectors, dtype="float64")
# TODO: make k_vectors array like at level of Kmodel
splitted = _split(k_vectors, n)
response = 0
for sub_kvectors in splitted:
formula.k_vectors = sub_kvectors
sub_result = np.empty(len(data), dtype="float64")
for i, elem in enumerate(data):
mu_T_eta = dict(zip(("mu", "temperature", "eta"), elem))
sub_result[i] = formula(**mu_T_eta)
response += sub_result
return response
def _split(k_vectors: np.ndarray, n: int) -> list[np.ndarray]:
"""Splits k_vectors into n batches
Args:
k_vectors: Array of momentum vectors
n: Number of batches
Returns:
A list of `n` arrays
"""
return np.array_split(k_vectors, n)
def _read_from_data(
data: Union[dict, KspaceOptions], param: str
) -> Union[bool, int, tuple, None]:
"""Reads the value of ``param`` from a dictionary or a dataclass.
Args:
param: The param we read
data: The input data from which we read
Returns:
The values are booleans for `precomplute`,
int for `chunk_size`
"""
if data is None:
data = dict(precomp=False, chunk_size=None)
if isinstance(data, dict):
data = KspaceOptions(**data)
elif not isinstance(data, KspaceOptions):
raise TypeError("unrecognized input data")
return getattr(data, param)
def _flatten_list(input_list):
flat = []
for elem in input_list:
if isinstance(elem, list):
flat.extend(elem)
else:
flat.append(elem)
return flat
def _is_eta_independent(
mu_kernels: Union[list[MuKernel], MuKernel, list[Union[list[MuKernel], MuKernel]]],
) -> bool:
mu_kernels = _flatten_list(mu_kernels)
for kernel in mu_kernels:
sig = inspect.signature(kernel)
params = sig.parameters
if "eta" in params:
default = params["eta"].default
if default is inspect.Parameter.empty or default is not None:
return False
return True