from abc import ABC, abstractmethod
import functools
import itertools
from typing import Union, Any, Callable, Optional
import warnings
import numpy as np
from ..integrator import Nquad
from ..formulas import KuboFormula
from .executors import KspaceExecutor
from ._common import Vectors, Bounds
__all__ = ["MuTEtaComputer", "HamParamComputer"]
class _Computer(ABC):
r"""A base class for parallelly computing a response formula using an :attr:`executor`.
Attributes:
formula: An instance of :class:`~py4mulas.formulas.KuboFormula`
executor: An instance of :class:`~py4mulas.executors.KspaceExecutor`
opts: The options to be passed to scipy integrator
method: Either `discrete` for discrete sum or `scipy` for adaptive integration.
Note:
The performance of 'scipy' method is not yet fully tested.
"""
def __init__(
self,
formula: KuboFormula,
executor: KspaceExecutor,
method: str = "discrete",
opts: dict = None,
):
self.formula = formula
self.opts = opts
self.method = method
if (executor.k_vectors is None) and (executor.k_bounds is None):
warnings.warn(
"If both k_vectors and k_bounds are not specified in the executor, "
"default model k_vectors or k_bounds are assumed "
"for discrete or scipy methods respectively."
)
if method == "discrete" and opts is not None:
warnings.warn("opts are ignored; these are only used for scipy integration")
if method == "discrete" and executor.k_vectors is None:
executor.k_vectors = self.formula.kmodel.k_vectors
if method == "scipy" and executor.k_bounds is None:
executor.k_bounds = self.formula.kmodel.bounds
self.executor = executor
@abstractmethod
def discret_response(
self, k_vectors: Union[Vectors, tuple[Vectors]], **kwargs: float
) -> complex:
pass
@abstractmethod
def nquad_response(
self, k_bounds: Union[Bounds, list[Bounds]], **kwargs: float
) -> complex:
pass
@abstractmethod
def __call__(
self, mu: float = 0, temperature: float = 0, eta: float = 0
) -> np.ndarray:
pass
[docs]
class MuTEtaComputer(_Computer):
r"""Computes a response formula using an :attr:`executor`
with precomputation of the operator kernel for a set of ``data``.
The energy kernel is computed once for each set of ``k_vectors``.
Therefore, the loop over ``data`` is made cheap for each worker.
Args:
formula: An instance of :class:`~py4mulas.formulas.KuboFormula`
executor: An instance of :class:`~py4mulas.executors.KspaceExecutor`
opts: The options to be passed to scipy integrator
method: Either `discrete` for discrete sum or `scipy` for adaptive integration.
"""
def __init__(
self,
formula: KuboFormula,
executor: KspaceExecutor,
opts: Optional[dict] = None,
method: str = "discrete",
):
super().__init__(formula=formula, executor=executor, opts=opts, method=method)
[docs]
def discret_response(
self, k_vectors: Union[Vectors, tuple[Vectors]], data: list
) -> np.ndarray:
formula = self.formula
shape = len(data)
if isinstance(k_vectors, tuple):
formula = self.formula
response = 0
for sub_kvecs in k_vectors:
formula.k_vectors = np.asarray(sub_kvecs)
sub_result = np.empty(shape, 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
formula = self.formula
formula.k_vectors = np.asarray(k_vectors)
sub_result = np.empty(shape, dtype="float64")
for i, elem in enumerate(data):
mu_T_eta = dict(zip(("mu", "temperature", "eta"), elem))
sub_result[i] = formula(**mu_T_eta)
return sub_result
[docs]
def nquad_response(
self, k_bounds: Union[Bounds, list[Bounds]], data: list
) -> np.ndarray:
formula = self.formula
shape = len(data)
if _is_list_of_tuples(k_bounds):
result = np.empty(shape, dtype="float64")
for i, elem in enumerate(data):
mu_T_eta = dict(zip(("mu", "temperature", "eta"), elem))
result[i], _ = Nquad(formula, bounds=k_bounds, opts=self.opts)(
**mu_T_eta
)
return result
response = 0
for bounds_i in k_bounds:
sub_integral = np.empty(shape, dtype="float64")
for i, elem in enumerate(data):
mu_T_eta = dict(zip(("mu", "temperature", "eta"), elem))
sub_integral[i], _ = Nquad(formula, bounds=bounds_i, opts=self.opts)(
**mu_T_eta
)
response += sub_integral
return response
[docs]
def __call__(
self,
mu: Union[float, list] = 0,
temperature: Union[float, list] = 0,
eta: Union[float, list] = 0,
) -> np.ndarray:
r"""Executes parallelly a transport formula for a combunation of transport parameters (:math:`\mu`, :math:`T`, :math:`\eta`)
Args:
mu: Chemical potential :math:`\mu`
temperature: System's temperature
eta: Broadening :math:`\eta`
Example:
>>> mu = np.linspace(0, 1, 10)
>>> temperature = np.linspace(0, 1, 10)
>>> eta = np.linspace(0, 0.1, 10)
>>> result = py4mulas.mpi.computers.MuTEtaComputer(some_formula, some_executor)(mu, temperature, eta)
Returns:
np.ndarray: The transport response as a 1d array ordered exactly
itertools.product(mu, temperature, eta). If these are all given as floats the result is
an array containing a single element.
"""
data = Mu_T_Eta_Data(mu, temperature, eta).data
if self.method == "scipy":
response = functools.partial(self.nquad_response, data=data)
else:
response = functools.partial(self.discret_response, data=data)
return self.executor(response)
[docs]
class HamParamComputer(_Computer):
r"""Computes a response formula using an :attr:`executor`
for varied Hamiltonian params.
Args:
formula: An instance of :class:`~py4mulas.formulas.KuboFormula`
executor: An instance of :class:`~py4mulas.executors.KspaceExecutor`
params: A dictionary containing the parameters to be changed. It should be in the form:
``{'param1':[...], 'param2':[...]}``.
Example:
>>> params = {'param1':np.linspace(0, 1, 10), 'param2':[0, 1]}
>>> result = py4mulas.mpi.computers.HamParamComputer(some_formula, some_executor, params)(mu=0, temperature=0, eta=0)
"""
def __init__(
self,
formula: KuboFormula,
executor: KspaceExecutor,
params: dict[str, Union[list, np.ndarray]],
method: str = "discrete",
opts: dict = None,
):
super().__init__(formula=formula, executor=executor, opts=opts, method=method)
if not isinstance(params, dict):
raise TypeError("params should be a dictionary")
data = HamParamData(params)
self.shape = data.length
self.names = data.names
self.params = list(data.data)
[docs]
def discret_response(
self, k_vectors: Union[Vectors, tuple[Vectors]], **kwargs: float
) -> np.ndarray:
formula = self.formula
if isinstance(k_vectors, tuple):
response = 0
for sub_kvecs in k_vectors:
formula.k_vectors = np.asarray(sub_kvecs)
sub_result = np.empty(self.shape, dtype="float64")
for i, elem in enumerate(self.params):
formula.kmodel.params = zip(self.names, elem)
sub_result[i] = formula(**kwargs)
response += sub_result
return response
formula.k_vectors = np.asarray(k_vectors)
sub_result = np.empty(self.shape, dtype="float64")
for i, elem in enumerate(self.params):
formula.kmodel.params = zip(self.names, elem)
sub_result[i] = formula(**kwargs)
return sub_result
[docs]
def nquad_response(
self, k_bounds: Union[Bounds, list[Bounds]], **kwargs: float
) -> np.ndarray:
formula = self.formula
if _is_list_of_tuples(k_bounds):
result = np.empty(self.shape, dtype="float64")
for i, elem in enumerate(self.params):
formula.kmodel.params = zip(self.names, elem)
result[i], _ = Nquad(formula, bounds=k_bounds, opts=self.opts)(**kwargs)
return result
response = 0
for bounds_i in k_bounds:
sub_integral = np.empty(self.shape, dtype="float64")
for i, elem in enumerate(self.params):
formula.kmodel.params = zip(self.names, elem)
sub_integral[i], _ = Nquad(formula, bounds=bounds_i, opts=self.opts)(
**kwargs
)
response += sub_integral
return response
def _response(self) -> Callable:
if self.method == "scipy":
return self.nquad_response
return self.discret_response
[docs]
def __call__(
self, mu: float = 0, temperature: float = 0, eta: float = 0
) -> np.ndarray:
r"""Gives the parallelly computed response formula at (:math:`\mu`, :math:`T`, :math:`\eta`)
Args:
mu: Chemical potential :math:`\mu`
temperature: System's temperature
eta: Broadening :math:`\eta`
Returns:
np.ndarray: The transport response as a 1d array
ordered exactly as itertools.product(:math:`l_1`, :math:`l2`, ...)
with :math:`l_i` being the list of values for the :math:`i` th parameter in :attr:`~py4mulas.mpi.computers.HamParamComputer.params`
"""
kwargs = dict(mu=mu, temperature=temperature, eta=eta)
response = functools.partial(self._response(), **kwargs)
return self.executor(response)
def _is_list_of_tuples(data: Any) -> bool:
sample = data[0]
return isinstance(sample, tuple) and isinstance(sample[0], (int, float))
def _process_ham_params(params_dict):
all_data = []
names = []
length = 0
for param_i, data_i in params_dict.items():
names.append(param_i)
all_data.append(data_i)
length += len(data_i)
return length, names, itertools.product(*all_data)
class HamParamData:
def __init__(self, input_dict):
self.length, self.names, self.data = _process_ham_params(input_dict)
class Mu_T_Eta_Data:
def __init__(self, mu, temperature, eta):
mus = mu
temps = temperature
etas = eta
if isinstance(mu, (float, int)):
mus = [mu]
if isinstance(temperature, (float, int)):
temps = [temperature]
if isinstance(eta, (float, int)):
etas = [eta]
self.data = list(itertools.product(mus, temps, etas))