import numpy as np
import ase # Annotations
from ase.calculators.calculator import PropertyNotImplementedError
from ase.utils import jsonable
def calculate_band_structure(atoms, path=None, scf_kwargs=None,
bs_kwargs=None, kpts_tol=1e-6, cell_tol=1e-6):
"""Calculate band structure.
The purpose of this function is to abstract a band structure calculation
so the workflow does not depend on the calculator.
First trigger SCF calculation if necessary, then set arguments
on the calculator for band structure calculation, then return
calculated band structure.
The difference from get_band_structure() is that the latter
expects the calculation to already have been done."""
if path is None:
path = atoms.cell.bandpath()
from ase.lattice import celldiff # Should this be a method on cell?
if any(path.cell.any(1) != atoms.pbc):
raise ValueError('The band path\'s cell, {}, does not match the '
'periodicity {} of the atoms'
.format(path.cell, atoms.pbc))
cell_err = celldiff(path.cell, atoms.cell.uncomplete(atoms.pbc))
if cell_err > cell_tol:
raise ValueError('Atoms and band path have different unit cells. '
'Please reduce atoms to standard form. '
'Cell lengths and angles are {} vs {}'
.format(atoms.cell.cellpar(), path.cell.cellpar()))
calc = atoms.calc
if calc is None:
raise ValueError('Atoms have no calculator')
if scf_kwargs is not None:
calc.set(**scf_kwargs)
# Proposed standard mechanism for calculators to advertise that they
# use the bandpath keyword to handle band structures rather than
# a double (SCF + BS) run.
use_bandpath_kw = getattr(calc, 'accepts_bandpath_keyword', False)
if use_bandpath_kw:
calc.set(bandpath=path)
atoms.get_potential_energy()
return calc.band_structure()
atoms.get_potential_energy()
if hasattr(calc, 'get_fermi_level'):
# What is the protocol for a calculator to tell whether
# it has fermi_energy?
eref = calc.get_fermi_level()
else:
eref = 0.0
if bs_kwargs is None:
bs_kwargs = {}
calc.set(kpts=path, **bs_kwargs)
calc.results.clear() # XXX get rid of me
# Calculators are too inconsistent here:
# * atoms.get_potential_energy() will fail when total energy is
# not in results after BS calculation (Espresso)
# * calc.calculate(atoms) doesn't ask for any quantity, so some
# calculators may not calculate anything at all
# * 'bandstructure' is not a recognized property we can ask for
try:
atoms.get_potential_energy()
except PropertyNotImplementedError:
pass
ibzkpts = calc.get_ibz_k_points()
kpts_err = np.abs(path.kpts - ibzkpts).max()
if kpts_err > kpts_tol:
raise RuntimeError('Kpoints of calculator differ from those '
'of the band path we just used; '
'err={} > tol={}'.format(kpts_err, kpts_tol))
bs = get_band_structure(atoms, path=path, reference=eref)
return bs
def get_band_structure(atoms=None, calc=None, path=None, reference=None):
"""Create band structure object from Atoms or calculator."""
# path and reference are used internally at the moment, but
# the exact implementation will probably change. WIP.
#
# XXX We throw away info about the bandpath when we create the calculator.
# If we have kept the bandpath, we can provide it as an argument here.
# It would be wise to check that the bandpath kpoints are the same as
# those stored in the calculator.
atoms = atoms if atoms is not None else calc.atoms
calc = calc if calc is not None else atoms.calc
kpts = calc.get_ibz_k_points()
energies = []
for s in range(calc.get_number_of_spins()):
energies.append([calc.get_eigenvalues(kpt=k, spin=s)
for k in range(len(kpts))])
energies = np.array(energies)
if path is None:
from ase.dft.kpoints import (
BandPath,
find_bandpath_kinks,
resolve_custom_points,
)
standard_path = atoms.cell.bandpath(npoints=0)
# Kpoints are already evaluated, we just need to put them into
# the path (whether they fit our idea of what the path is, or not).
#
# Depending on how the path was established, the kpoints might
# be valid high-symmetry points, but since there are multiple
# high-symmetry points of each type, they may not coincide
# with ours if the bandpath was generated by another code.
#
# Here we hack it so the BandPath has proper points even if they
# come from some weird source.
#
# This operation (manually hacking the bandpath) is liable to break.
# TODO: Make it available as a proper (documented) bandpath method.
kinks = find_bandpath_kinks(atoms.cell, kpts, eps=1e-5)
pathspec, special_points = resolve_custom_points(
kpts[kinks], standard_path.special_points, eps=1e-5)
path = BandPath(standard_path.cell,
kpts=kpts,
path=pathspec,
special_points=special_points)
# XXX If we *did* get the path, now would be a good time to check
# that it matches the cell! Although the path can only be passed
# because we internally want to not re-evaluate the Bravais
# lattice type. (We actually need an eps parameter, too.)
if reference is None:
# Fermi level should come from the GS calculation, not the BS one!
reference = calc.get_fermi_level()
if reference is None:
# Fermi level may not be available, e.g., with non-Fermi smearing.
# XXX Actually get_fermi_level() should raise an error when Fermi
# level wasn't available, so we should fix that.
reference = 0.0
return BandStructure(path=path,
energies=energies,
reference=reference)
class BandStructurePlot:
def __init__(self, bs):
self.bs = bs
self.ax = None
self.xcoords = None
def plot(self, ax=None, emin=-10, emax=5, filename=None,
show=False, ylabel=None, colors=None, point_colors=None,
label=None, loc=None,
cmap=None, cmin=-1.0, cmax=1.0, sortcolors=False,
colorbar=True, clabel='$s_z$', cax=None,
**plotkwargs):
"""Plot band-structure.
ax: Axes
MatPlotLib Axes object. Will be created if not supplied.
emin, emax: float
Minimum and maximum energy above reference.
filename: str
If given, write image to a file.
show: bool
Show the image (not needed in notebooks).
ylabel: str
The label along the y-axis. Defaults to 'energies [eV]'
colors: sequence of str
A sequence of one or two color specifications, depending on
whether there is spin.
Default: green if no spin, yellow and blue if spin is present.
point_colors: ndarray
An array of numbers of the shape (nspins, n_kpts, nbands) which
are then mapped onto colors by the colormap (see ``cmap``).
``colors`` and ``point_colors`` are mutually exclusive
label: str or list of str
Label for the curves on the legend. A string if one spin is
present, a list of two strings if two spins are present.
Default: If no spin is given, no legend is made; if spin is
present default labels 'spin up' and 'spin down' are used, but
can be suppressed by setting ``label=False``.
loc: str
Location of the legend.
If ``point_colors`` is given, the following arguments can be specified.
cmap:
Only used if colors is an array of numbers. A matplotlib
colormap object, or a string naming a standard colormap.
Default: The matplotlib default, typically 'viridis'.
cmin, cmax: float
Minimal and maximal values used for colormap translation.
Default: -1.0 and 1.0
colorbar: bool
Whether to make a colorbar.
clabel: str
Label for the colorbar (default 's_z', set to None to suppress.
cax: Axes
Axes object used for plotting colorbar. Default: split off a
new one.
sortcolors (bool or callable):
Sort points so highest color values are in front. If a callable is
given, then it is called on the color values to determine the sort
order.
Any additional keyword arguments are passed directly to matplotlib's
plot() or scatter() methods, depending on whether point_colors is
given.
"""
import matplotlib.pyplot as plt
if colors is not None and point_colors is not None:
raise ValueError("Don't give both 'color' and 'point_color'")
if self.ax is None:
ax = self.prepare_plot(ax, emin, emax, ylabel)
e_skn = self.bs.energies
nspins = len(e_skn)
if point_colors is None:
# Normal band structure plot
if colors is None:
if len(e_skn) == 1:
colors = 'g'
else:
colors = 'yb'
elif (len(colors) != nspins):
raise ValueError(
"colors should be a sequence of {nspin} colors"
)
# Default values for label
if label is None and nspins == 2:
label = ['spin up', 'spin down']
if label:
if nspins == 1 and isinstance(label, str):
label = [label]
elif len(label) != nspins:
raise ValueError(
f'label should be a list of {nspins} strings'
)
for spin, e_kn in enumerate(e_skn):
kwargs = dict(color=colors[spin])
kwargs.update(plotkwargs)
lbl = None # Retain lbl=None if label=False
if label:
lbl = label[spin]
ax.plot(self.xcoords, e_kn[:, 0], label=lbl, **kwargs)
for e_k in e_kn.T[1:]:
ax.plot(self.xcoords, e_k, **kwargs)
show_legend = label is not None or nspins == 2
else:
# A color per datapoint.
kwargs = dict(vmin=cmin, vmax=cmax, cmap=cmap, s=1)
kwargs.update(plotkwargs)
shape = e_skn.shape
xcoords = np.zeros(shape)
xcoords += self.xcoords[np.newaxis, :, np.newaxis]
if sortcolors:
if callable(sortcolors):
perm = sortcolors(point_colors).argsort(axis=None)
else:
perm = point_colors.argsort(axis=None)
e_skn = e_skn.ravel()[perm].reshape(shape)
point_colors = point_colors.ravel()[perm].reshape(shape)
xcoords = xcoords.ravel()[perm].reshape(shape)
things = ax.scatter(xcoords, e_skn, c=point_colors, **kwargs)
if colorbar:
cbar = plt.colorbar(things, cax=cax)
if clabel:
cbar.set_label(clabel)
show_legend = False
self.finish_plot(filename, show, loc, show_legend)
return ax
def prepare_plot(self, ax=None, emin=-10, emax=5, ylabel=None):
import matplotlib.pyplot as plt
if ax is None:
ax = plt.figure().add_subplot(111)
def pretty(kpt):
if kpt == 'G':
kpt = r'$\Gamma$'
elif len(kpt) == 2:
kpt = kpt[0] + '$_' + kpt[1] + '$'
return kpt
self.xcoords, label_xcoords, orig_labels = self.bs.get_labels()
label_xcoords = list(label_xcoords)
labels = [pretty(name) for name in orig_labels]
i = 1
while i < len(labels):
if label_xcoords[i - 1] == label_xcoords[i]:
labels[i - 1] = labels[i - 1] + ',' + labels[i]
labels.pop(i)
label_xcoords.pop(i)
else:
i += 1
for x in label_xcoords[1:-1]:
ax.axvline(x, color='0.5')
ylabel = ylabel if ylabel is not None else 'energies [eV]'
ax.set_xticks(label_xcoords)
ax.set_xticklabels(labels)
ax.set_ylabel(ylabel)
ax.axhline(self.bs.reference, color='k', ls=':')
ax.axis(xmin=0, xmax=self.xcoords[-1], ymin=emin, ymax=emax)
self.ax = ax
return ax
def finish_plot(self, filename, show, loc, show_legend=False):
import matplotlib.pyplot as plt
if show_legend:
leg = plt.legend(loc=loc)
leg.get_frame().set_alpha(1)
if filename:
plt.savefig(filename)
if show:
plt.show()
[docs]
@jsonable('bandstructure')
class BandStructure:
"""A band structure consists of an array of eigenvalues and a bandpath.
BandStructure objects support JSON I/O.
"""
def __init__(self, path, energies, reference=0.0):
self._path = path
self._energies = np.asarray(energies)
assert self.energies.shape[0] in [1, 2] # spins x kpts x bands
assert self.energies.shape[1] == len(path.kpts)
assert np.isscalar(reference)
self._reference = reference
@property
def energies(self) -> np.ndarray:
"""The energies of this band structure.
This is a numpy array of shape (nspins, nkpoints, nbands)."""
return self._energies
@property
def path(self) -> 'ase.dft.kpoints.BandPath':
"""The :class:`~ase.dft.kpoints.BandPath` of this band structure."""
return self._path
@property
def reference(self) -> float:
"""The reference energy.
Semantics may vary; typically a Fermi energy or zero,
depending on how the band structure was created."""
return self._reference
[docs]
def subtract_reference(self) -> 'BandStructure':
"""Return new band structure with reference energy subtracted."""
return BandStructure(self.path, self.energies - self.reference,
reference=0.0)
def todict(self):
return dict(path=self.path,
energies=self.energies,
reference=self.reference)
[docs]
def get_labels(self, eps=1e-5):
""""See :func:`ase.dft.kpoints.labels_from_kpts`."""
return self.path.get_linear_kpoint_axis(eps=eps)
[docs]
def plot(self, *args, **kwargs):
"""Plot this band structure."""
bsp = BandStructurePlot(self)
return bsp.plot(*args, **kwargs)
def __repr__(self):
return ('{}(path={!r}, energies=[{} values], reference={})'
.format(self.__class__.__name__, self.path,
'{}x{}x{}'.format(*self.energies.shape),
self.reference))