from functools import cached_property
import itertools as it
from typing import Union, Any, Callable
from math import e
import numpy as np
import sympy
from sympy import Matrix, lambdify, sympify, Symbol
from ._common import compute_dk
__all__ = ["Kmodel"]
[docs]
class Kmodel:
r"""Base class for calculating k-space proprties of a momentum dependent Hamiltonian.
Attributes:
hamiltonian: The k-space hamiltonian matrix.
params: the parameters used in the definition of ``hamiltonian``
k_1: The momentum values in the first dimension.
k_2: The momentum values in the second dimension.
k_3: The momentum values in the third dimension.
zone: The zone of the momentum space or the Brillouin zone
real_space: False corresponds to reduced Bloch phase representation.
True corresponds to cartesian momenta representation.
For py4mulas integration the reduced momenta representation is recomaended.
If hamiltonian is given as a string or a sympy expression,
the user should be carefull to appropriately handle the underlying representation.
Exemple:
>>> System = Kmodel(model, k_1=[0], k_2=[0], k_3=[0], params={}, zone=shape)
"""
def __init__(
self,
hamiltonian: Any,
k_1: Union[list, np.ndarray] = None,
k_2: Union[list, np.ndarray] = None,
k_3: Union[list, np.ndarray] = None,
zone: callable = None,
params: dict = dict(),
real_space: bool = False,
):
if all(k is None for k in (k_1, k_2, k_3)):
raise ValueError("k-space must be at least 1D.")
self.k_1 = k_1
self.k_2 = k_2
self.k_3 = k_3
self.zone = zone
self._k_vectors = momentum_vectors(self.k_1, self.k_2, self.k_3, zone=self.zone)
self.hamiltonian = hamiltonian
self._params = params
self.real_space = real_space
self.norbs = Symbolic(
self.hamiltonian, params=self._params, real_space=self.real_space
).norbs
self.dim = len(self._k_vectors[0])
@cached_property
def symbolic_ham(self):
# Lambdify hamiltonian and get its size (norbs)
return Symbolic(
self.hamiltonian, params=self.params, real_space=self.real_space
)
@cached_property
def H(self):
return self.symbolic_ham.lambdified()
@property
def symbolic_H(self):
return self.symbolic_ham.H_expression
@property
def velocities(self):
# Make symbolic velocities along the available directions
return self.symbolic_ham.make_velocities()
@property
def params(self):
return self._params
@params.setter
def params(self, new_params):
self._params.update(new_params)
self.__dict__.pop("H", None)
self.__dict__.pop("symbolic_ham", None)
@property
def k_vectors(self):
return self._k_vectors
@k_vectors.setter
def k_vectors(self, vectors):
self._k_vectors = vectors
@property
def bounds(self):
# cartezian bounds for integration
return bounds_from_vectors(self.k_vectors)
@property
def dk(self):
k_1 = self.k_1
k_2 = self.k_2
k_3 = self.k_3
if k_1 is not None:
_dk = k_1[1] - k_1[0]
if k_2 is not None:
_dk *= k_2[1] - k_2[0]
if k_3 is not None:
_dk *= k_3[1] - k_3[0]
return _dk
k_vectors = self.k_vectors
return compute_dk(k_vectors, self.dim)
# -------------------------------- Utilities ---------------------------------
def bounds_from_vectors(k_vectors):
k_vectors = np.asarray(k_vectors)
bounds = []
dim = k_vectors.shape[1]
for i in range(dim):
column_i = k_vectors[:, i]
bounds.append((min(column_i), max(column_i)))
return bounds
def momentum_vectors(
k_1: np.ndarray, k_2: np.ndarray, k_3: np.ndarray, zone: Callable = None
):
"""Generates the momentum vectors in a given `zone`.
If provided, `zone` should be a function
of the underlying momentum variables.
"""
k_s = []
for element in [k_1, k_2, k_3]:
if element is not None:
k_s.append(element)
ksites = it.product(*k_s)
if zone is None:
vectors = list(ksites)
else:
assert callable(zone)
vectors = [k for k in ksites if zone(*k)]
return vectors
def from_kwant(syst: Any, params: dict = None, real_space: bool = False) -> str:
"""Converts an infinite kwant Builder into model if kwant is available
Args:
syst: Should be an infinite kwant Builder
params: Sarameters used in defining the system
real_space: Specifies whether cartesian coordinates
or reduced momenta are used.
Raises:
ImportError: If kwant is not available
ValueError: If `syst` is not a kwant Builder
Returns:
A sympy expression as string
"""
try:
import kwant
from kwant.qsymm import builder_to_model
except ImportError as error:
message = (
"Kwant is not available. Please install it, "
"in order to be able to define Kwant Builders. "
)
raise ImportError(message) from error
if isinstance(syst, kwant.builder.Builder):
model = builder_to_model(syst, params=params, real_space=real_space)
return str(model.tosympy())
raise ValueError(
f"unrecognized model type: {type(model)}, \
model should be an infinite kwant Builder, \
a Sympy Matrix or an str"
)
def extract_momenta(atoms):
momenta = []
for elem in atoms:
if elem[0] == "k" and elem[-1] in "xyz":
momenta.append(elem)
return momenta
# ------------------------------ Symbolic arsenal ---------------------------------
class Symbolic:
"""Symbolic properties of model.
Attributes:
hamiltonian: A sympy Matrix, an str or a kwant Builder
params: Used in defining `hamiltonian`
real_space: Specifies whether cartesian coordinates
or reduced momenta are used.
If an expression is provided we should be sure it is written in
the appropriate coordinates.
"""
def __init__(
self, hamiltonian: Any, params: dict = None, real_space: bool = False
) -> None:
if params is None:
params = {}
params = dict(params)
params.update({"e": e})
if not isinstance(hamiltonian, (sympy.MatrixBase, str)):
hamiltonian = from_kwant(hamiltonian, params=params, real_space=real_space)
self.H_expression = hamiltonian
ham_matrix = Matrix(sympify(hamiltonian, params))
self.hamiltonian = ham_matrix.subs(params)
self.norbs = self.hamiltonian.shape[0]
self.atoms = sorted([s.name for s in self.hamiltonian.atoms(Symbol)])
self.momenta = extract_momenta(self.atoms)
ensure_params(params, self.atoms, self.momenta)
self._arg_symbols = [Symbol(a) for a in self.atoms]
def lambdified(self) -> callable:
"""Return the lambdified Hamiltonian H(*args)."""
return lambdify_matrix(self.hamiltonian, self.momenta)
def make_velocities(self) -> dict:
r"""Returns dict: coord -> lambdified velocity matrix :math:`dH/dk_i`.
Keys (i = last char of momentum symbol).
"""
# compute derivatives once
velocities = {}
for momentum in self.momenta:
coord = momentum[-1]
dH = self.hamiltonian.diff(Symbol(momentum))
velocities[coord] = lambdify_matrix(dH, self.momenta)
return velocities
def ensure_params(params, atoms, momenta):
for atom in atoms:
if atom not in momenta:
if atom not in params:
raise ValueError("some params are not asigned in the input dictionary")
def lambdify_matrix(expr, vars):
n, m = expr.shape
funcs = [
[lambdify(vars, expr[i, j], modules="numpy") for j in range(m)]
for i in range(n)
]
def f(*args):
args = [np.asarray(a) for a in args]
bshape = np.broadcast_arrays(*args)[0].shape
dtype = np.result_type(*args, complex)
if bshape == ():
out = np.empty((n, m), dtype=dtype)
for i in range(n):
for j in range(m):
out[i, j] = funcs[i][j](*args)
return out
out = np.empty(bshape + (n, m), dtype=dtype)
for i in range(n):
for j in range(m):
val = np.asarray(funcs[i][j](*args))
out[..., i, j] = np.broadcast_to(val, bshape)
return out
return f