from typing import Union, Optional
import matplotlib.pyplot as plt
from matplotlib.colors import Normalize
from matplotlib.collections import LineCollection
import numpy as np
from .models import Kmodel
from .operators import KspaceOpera
[docs]
def kplot(
model: Kmodel,
quantity: np.ndarray,
cmap: str = "seismic",
label: str = None,
title: str = None,
show: bool = True,
file: str = None,
ax: plt.Axes = None,
s: int = 10,
vmin: float = None,
vmax: float = None,
fontsize: float = 12,
figsize: tuple = None,
alpha: float = 0.8,
k_vectors: np.ndarray = None,
):
"""Plots ``quantity`` over the kspace of ``model``. A scatter plot.
Args:
model: An instance of :class:`~py4mulas.models.Kmodel`
quantity: The quantity to plot
cmap: Colormap
label: Colorbar label
title: Title of the plot
show: Whether to show the fig or not
file: Name of file to whish the figure can be saved
ax: Axes to be used for the plot
s: Marker size
vmin: Minimum value of `quantity`
vmax: Maximum value of `quantity`
fontsize: The fontsize to be used
figsize: The size of the figure
alpha: Alpha blending value
k_vectors: Momentum vectors over which `quantity` is plotted
"""
# set fontsize
plt.rcParams.update({"font.size": fontsize})
dim = model.dim
if k_vectors is None:
k_vectors = model.k_vectors
if not vmin:
vmin = quantity.min()
if not vmax:
vmax = quantity.max()
if not ax:
fig, ax = _py4mulas_fig(figsize=figsize, dim=dim)
ax.set_xlabel(r"$k_x$")
if dim > 1:
ax.set_ylabel(r"$k_y$")
if dim == 3:
ax.set_zlabel(r"$k_z$")
if title:
ax.set_title(title, )
else:
fig = ax.figure
if dim == 3:
x, y, z = _get_kpoints(k_vectors, dim)
sc = ax.scatter(
x, y, z, c=quantity, cmap=cmap, s=s, alpha=alpha, vmin=vmin, vmax=vmax
)
else:
x, y = _get_kpoints(k_vectors, dim)
sc = ax.scatter(
x, y, c=quantity, cmap=cmap, s=s, alpha=alpha, vmin=vmin, vmax=vmax)
cbar = fig.colorbar(sc, ax=ax)
if label:
cbar.set_label(label, )
if file:
plt.savefig(file)
if show:
plt.show()
[docs]
def plot_bands(
model: Kmodel,
path: Optional[list[tuple]] = None,
operator: Optional[KspaceOpera] = None,
path_labels: list[str] = None,
n_per_segment: int = 70,
ax: plt.Axes = None,
cmap: str = "viridis",
file: str = False,
show: bool = True,
figsize: tuple = (8, 6),
fontsize: float = 14,
vmin: float = None,
vmax: float = None,
emin: float = None,
emax: float = None,
cbar: bool = True,
bands: Union[list, int, tuple] = None,
):
"""Plots the energy bands of a model along a given path,
optionally with the expectation value of a py4mulas operator
Args:
model: An instance of :class:`~py4mulas.models.Kmodel`
path: a path in kspace
operator: An instance of :class:`~py4mulas.operators.KspaceOpera`,
an operator whose expectation value is to be plotted
path_labels: Labels corresponding to the momentum points in ``path``.
n_per_segment: Number of points for each segment of ``path``.
ax: Axes to be used for the plot.
cmap: A custum colormap
file: Name of file to which plot can be saved.
show: Whether display the plot or not.
figsize: Size of figure.
fontsize: Font size for ticks and labeling.
vmin: Minimum value of the operator expectation value.
vmax: Maxmum value of the operator expectation value.
emin: Minimum energy to be considered.
emax: Maximum energy to be considered.
cbar: Whether to show the colorbar or not.
bands: A sellection of bands to be shown.
"""
norbs = model.norbs
ranges = range(norbs)
if bands is not None:
if isinstance(bands, (list, tuple)):
for i in bands:
if not isinstance(i, int):
raise ValueError("bands should be an iterator of integers")
ranges = bands
elif isinstance(bands, int):
ranges = [bands]
else:
raise TypeError("bands should be an integer, tuple or a list")
plt.rcParams.update({"font.size": fontsize})
k_path = model.k_vectors
if path is not None:
k_path, _ = _interpolate_k_path(path, n_per_segment=n_per_segment)
bands, projection = projected_spectrum(model, k_path, operator)
if not ax:
_, ax = _py4mulas_fig(figsize=figsize, dim=None)
if path is not None:
tick_positions = [n_per_segment * i for i in range(len(path))]
ax.set_xticks(tick_positions)
if path_labels is not None:
ax.set_xticklabels(path_labels)
ax.set_xlabel(r"$k$")
ax.set_ylabel(r"$Energy$")
x = range(len(k_path))
if projection is not None:
colors = projection.real
if not vmin:
vmin = colors.min()
if not vmax:
vmax = colors.max()
norm = Normalize(vmin=vmin, vmax=vmax)
# Create segments from x, y for line collection
for i in ranges:
y = bands[:, i]
ci = colors[:, i].real
points = np.array([x, y]).T.reshape(-1, 1, 2)
segments = np.concatenate([points[:-1], points[1:]], axis=1)
# Create a LineCollection
lc = LineCollection(segments, cmap=cmap, norm=norm)
lc.set_array(ci[:-1]) # Set the array that maps to the colormap
lc.set_linewidth(1.5)
ax.add_collection(lc)
if i == min(ranges):
if cbar:
cbar = plt.colorbar(lc, ax=ax)
else:
for i in ranges:
y = bands[:, i]
ax.plot(x, y)
if emin is not None:
ymin = emin
else:
ymin = np.array(bands).min()
if emax is not None:
ymax = emax
else:
ymax = np.array(bands).max()
# margins
shift = (ymax - ymin) / 100.0
ax.set_ylim(ymin - shift, ymax + shift)
if file:
plt.savefig(file)
if show:
plt.show()
def _py4mulas_fig(figsize=None, dim=2):
if figsize:
assert isinstance(figsize, (list, tuple)), "figsize can only be a list or tuple"
assert len(figsize) == 2, "figsize must be a tuple or list of two real numbers"
else:
figsize = (8, 8)
if dim == 3:
fig = plt.figure(figsize=figsize, constrained_layout=True)
ax = fig.add_subplot(111, projection="3d")
else:
if dim == 1: # rescale for 1d
figsize = (8, 1)
fig, ax = plt.subplots(figsize=figsize, constrained_layout=True)
return fig, ax
def _get_kpoints(k_vectors, dim):
k_vectors = np.array(k_vectors)
x = k_vectors[:, 0]
if dim == 1:
return x, np.zeros_like(x)
else:
y = k_vectors[:, 1]
if dim == 2:
return x, y
else:
z = k_vectors[:, 2]
return x, y, z
def _interpolate_k_path(points, n_per_segment=100):
"""Interpolates a list of k-points into a smooth path.
Returns:
k_points: Array of kx, ky along the path
k_dist: Aistance along the path
k_nodes: Indices of segment ends (for labeling)
"""
k_path = []
k_dist = [0]
k_nodes = [0]
for i in range(len(points) - 1):
p1, p2 = points[i], points[i + 1]
segment = np.linspace(p1, p2, n_per_segment, endpoint=False)
k_path.append(segment)
dist = np.linalg.norm(np.array(p2) - np.array(p1))
k_dist += list(
np.linspace(k_dist[-1], k_dist[-1] + dist, n_per_segment, endpoint=False)[
1:
]
)
k_nodes.append(len(k_dist))
k_path = np.concatenate(k_path, axis=0)
return k_path, k_nodes
[docs]
def projected_spectrum(
model: Kmodel, path: list[np.ndarray] = None, operator: KspaceOpera = None
) -> tuple[np.ndarray]:
"""Compute eigenvalues and optionally operator weight along a path.
Args:
model: An instnce of Kmodel
path: A list of kpoits
operator: An instance of `~py4mulas.operators.KspaceOpera`
Raises:
TypeError: When operator is not An instance of `~py4mulas.operators.KspaceOpera`
Returns:
Array of eigenvaluses, and corresponding operator projection
"""
dim = len(path[0])
n = model.norbs
k_block = np.asarray(path)
k_args = [k_block[:, d] for d in range(dim)]
hk = model.H(*k_args)
energies, psi = np.linalg.eigh(hk)
if operator is None:
return energies, None
if isinstance(operator, KspaceOpera):
operator_matrix = operator(k_args, energies, psi)
return energies, np.einsum("kii->ki", operator_matrix)
elif isinstance(operator, np.ndarray):
assert operator.shape == (n, n)
projection = np.einsum(
"kni,nm,kmi->ki", psi.conj(), operator, psi, optimize=True
)
return energies, projection
else:
raise TypeError("Unsupported operator type")