import sys
import re
import logging
import numpy as np
import matplotlib.pyplot as plt
[docs]class ProcarPlot:
"""A depeciated class ot plot the band structure
"""
[docs] def __init__(self, bands, spd, kpoints=None):
self.bands = bands.transpose()
self.spd = spd.transpose()
self.kpoints = kpoints
return
[docs] def plotBands(
self,
size=0.02,
marker="o",
ticks=None,
color="blue",
discontinuities=[],
figsize=(13, 9),
ax=None,
):
if not ax:
fig = plt.figure(figsize=figsize)
fig.tight_layout()
ax = fig.add_subplot(111)
else:
fig = plt.gca()
if size is not None:
size = size / 2
if self.kpoints is not None:
xaxis = [0]
#### MODIFIED FOR DISCONTINOUS BANDS ####
if ticks:
ticks, ticksNames = list(zip(*ticks))
# counters for number of discontinuities
icounter = 1
ii = 0
for i in range(1, len(self.kpoints) - len(discontinuities)):
d = self.kpoints[icounter] - self.kpoints[icounter - 1]
d = np.sqrt(np.dot(d, d))
xaxis.append(d + xaxis[-1])
icounter += 1
ii += 1
if ii in discontinuities:
icounter += 1
ii += 1
xaxis.append(xaxis[-1])
xaxis = np.array(xaxis)
# plotting
for i_tick in range(len(ticks) - 1):
x = xaxis[ticks[i_tick] : ticks[i_tick + 1] + 1]
y = self.bands.transpose()[ticks[i_tick] : ticks[i_tick + 1] + 1, :]
ax.plot(x, y, "r-", marker=marker, markersize=size, color=color)
#### END OF MODIFIED DISCONTINUOUS BANDS ####
# if ticks are not given
else:
for i in range(1, len(self.kpoints)):
d = self.kpoints[i - 1] - self.kpoints[i]
d = np.sqrt(np.dot(d, d))
xaxis.append(d + xaxis[-1])
xaxis = np.array(xaxis)
ax.plot(
xaxis,
self.bands.transpose(),
"r-",
marker=marker,
markersize=size,
color=color,
)
# if self.kpoints is None
else:
xaxis = np.arange(len(self.bands))
ax.plot(
xaxis,
self.bands.transpose(),
"r-",
marker=marker,
markersize=size,
color=color,
)
ax.set_xlim(xaxis.min(), xaxis.max())
# Handling ticks
if ticks:
# added for meta-GGA calculations
if ticks[0] > 0:
ax.set_xlim(left=xaxis[ticks[0]])
ticks = [xaxis[x] for x in ticks]
ax.set_xticks(ticks)
ax.set_xticklabels(ticksNames)
for xc in ticks:
ax.axvline(x=xc, color="k")
ax.axhline(color="black", linestyle="--")
return fig, ax
[docs] def parametricPlot(
self,
cmap="jet",
vmin=None,
vmax=None,
mask=None,
ticks=None,
discontinuities=[],
ax=None,
figsize=(13, 9),
plot_bar=True,
linewidth=1,
):
from matplotlib.collections import LineCollection
import matplotlib
# fig = plt.figure() # use plt.gca() since it won't create a new figure for band comparison
if ax is None:
fig = plt.figure(figsize=figsize)
fig.tight_layout()
ax = fig.add_subplot(111)
else:
fig = ax.get_figure()
bsize, ksize = self.bands.shape
# print self.bands
if mask is not None:
mbands = np.ma.masked_array(self.bands, np.abs(self.spd) < mask)
else:
# Faking a mask, all elemtnet are included
mbands = np.ma.masked_array(self.bands, False)
if vmin is None:
vmin = self.spd.min()
if vmax is None:
vmax = self.spd.max()
print("normalizing to : ", (vmin, vmax))
norm = matplotlib.colors.Normalize(vmin, vmax)
if self.kpoints is not None:
xaxis = [0]
#### MODIFIED FOR DISCONTINOUS BANDS ####
if ticks:
ticks, ticksNames = list(zip(*ticks))
# counters for number of discontinuities
icounter = 1
ii = 0
for i in range(1, len(self.kpoints) - len(discontinuities)):
d = self.kpoints[icounter] - self.kpoints[icounter - 1]
d = np.sqrt(np.dot(d, d))
xaxis.append(d + xaxis[-1])
icounter += 1
ii += 1
if ii in discontinuities:
icounter += 1
ii += 1
xaxis.append(xaxis[-1])
xaxis = np.array(xaxis)
# plotting
for y, z in zip(mbands, self.spd):
points = np.array([xaxis, y]).T.reshape(-1, 1, 2)
segments = np.concatenate([points[:-1], points[1:]], axis=1)
lc = LineCollection(segments, cmap=plt.get_cmap(cmap), norm=norm)
lc.set_array(z)
lc.set_linewidth(linewidth)
ax.add_collection(lc)
if plot_bar:
cb = fig.colorbar(lc, ax=ax)
cb.ax.tick_params(labelsize=20)
ax.set_xlim(xaxis.min(), xaxis.max())
ax.set_ylim(mbands.min(), mbands.max())
#### END OF MODIFIED DISCONTINUOUS BANDS ####
# if ticks are not given
else:
for i in range(1, len(self.kpoints)):
d = self.kpoints[i - 1] - self.kpoints[i]
d = np.sqrt(np.dot(d, d))
xaxis.append(d + xaxis[-1])
xaxis = np.array(xaxis)
# plotting
for y, z in zip(mbands, self.spd):
points = np.array([xaxis, y]).T.reshape(-1, 1, 2)
segments = np.concatenate([points[:-1], points[1:]], axis=1)
lc = LineCollection(segments, cmap=plt.get_cmap(cmap), norm=norm)
lc.set_array(z)
lc.set_linewidth(linewidth)
ax.add_collection(lc)
if plot_bar:
cb = fig.colorbar(lc, ax=ax)
cb.ax.tick_params(labelsize=20)
ax.set_xlim(xaxis.min(), xaxis.max())
ax.set_ylim(mbands.min(), mbands.max())
# if self.kpoints is None
else:
xaxis = np.arange(ksize)
for y, z in zip(mbands, self.spd):
points = np.array([xaxis, y]).T.reshape(-1, 1, 2)
segments = np.concatenate([points[:-1], points[1:]], axis=1)
lc = LineCollection(segments, cmap=plt.get_cmap(cmap), norm=norm)
lc.set_array(z)
lc.set_linewidth(linewidth)
ax.add_collection(lc)
if plot_bar:
cb = fig.colorbar(lc, ax=ax)
cb.ax.tick_params(labelsize=20)
ax.set_xlim(xaxis.min(), xaxis.max())
ax.set_ylim(mbands.min(), mbands.max())
# handling ticks
if ticks:
# added for meta-GGA calculations
if ticks[0] > 0:
ax.set_xlim(left=xaxis[ticks[0]])
ticks = [xaxis[x] for x in ticks]
ax.set_xticks(ticks)
ax.set_xticklabels(ticksNames)
for xc in ticks:
ax.axvline(x=xc, color="black")
ax.axhline(color="black", linestyle="--")
return fig, ax
[docs] def scatterPlot(
self,
size=50,
mask=None,
cmap="hot_r",
vmax=None,
vmin=None,
marker="o",
ticks=None,
discontinuities=[],
ax=None,
):
bsize, ksize = self.bands.shape
# plotting
if not ax:
fig = plt.figure(figsize=(13, 9))
fig.tight_layout()
ax = fig.add_subplot(111)
else:
fig = ax.get_figure()
if self.kpoints is not None:
xaxis = [0]
#### MODIFIED FOR DISCONTINOUS BANDS ####
if ticks:
ticks, ticksNames = list(zip(*ticks))
# counters for number of discontinuities
icounter = 1
ii = 0
for i in range(1, len(self.kpoints) - len(discontinuities)):
d = self.kpoints[icounter] - self.kpoints[icounter - 1]
d = np.sqrt(np.dot(d, d))
xaxis.append(d + xaxis[-1])
icounter += 1
ii += 1
if ii in discontinuities:
icounter += 1
ii += 1
xaxis.append(xaxis[-1])
xaxis = np.array(xaxis)
xaxis.shape = (1, ksize)
xaxis = xaxis.repeat(bsize, axis=0)
if mask is not None:
mbands = np.ma.masked_array(self.bands, np.abs(self.spd) < mask)
else:
mbands = self.bands
scatter = ax.scatter(
xaxis,
mbands,
c=self.spd,
s=size,
linewidths=0,
cmap=cmap,
vmax=vmax,
vmin=vmin,
marker=marker,
edgecolors="none",
)
fig.colorbar(scatter)
ax.set_xlim(xaxis.min(), xaxis.max())
#### END OF MODIFIED DISCONTINUOUS BANDS ####
# if ticks are not given
else:
for i in range(1, len(self.kpoints)):
d = self.kpoints[i - 1] - self.kpoints[i]
d = np.sqrt(np.dot(d, d))
xaxis.append(d + xaxis[-1])
xaxis = np.array(xaxis)
xaxis.shape = (1, ksize)
xaxis = xaxis.repeat(bsize, axis=0)
if mask is not None:
mbands = np.ma.masked_array(self.bands, np.abs(self.spd) < mask)
else:
mbands = self.bands
scatter = ax.scatter(
xaxis,
mbands,
c=self.spd,
s=size,
linewidths=0,
cmap=cmap,
vmax=vmax,
vmin=vmin,
marker=marker,
edgecolors="none",
)
plt.colorbar(scatter)
ax.set_xlim(xaxis.min(), xaxis.max())
# if kpoints is None
else:
xaxis = np.arange(ksize)
xaxis.shape = (1, ksize)
xaxis = xaxis.repeat(bsize, axis=0)
if mask is not None:
mbands = np.ma.masked_array(self.bands, np.abs(self.spd) < mask)
else:
mbands = self.bands
ax.scatter(
xaxis,
mbands,
c=self.spd,
s=size,
linewidths=0,
cmap=cmap,
vmax=vmax,
vmin=vmin,
marker=marker,
edgecolors="none",
)
plt.colorbar(ax=ax)
ax.set_xlim(xaxis.min(), xaxis.max())
# handling ticks
if ticks:
# added for meta-GGA calculations
if ticks[0] > 0:
ax.set_xlim(left=xaxis[0, ticks[0]])
ticks = [xaxis[0, x] for x in ticks]
ax.set_xticks(ticks)
ax.set_xticklabels(ticksNames)
ax.axhline(color="black", linestyle="--")
return fig, ax
[docs] def atomicPlot(self, cmap="hot_r", vmin=None, vmax=None, ax=None):
"""
Just a handler to parametricPlot. Useful to plot energy levels.
It adds a fake k-point. Shouldn't be invoked with more than one
k-point
ax not implemented here, not need
"""
print("Atomic plot: bands.shape :", self.bands.shape)
print("Atomic plot: spd.shape :", self.spd.shape)
print("Atomic plot: kpoints.shape:", self.kpoints.shape)
self.bands = np.hstack((self.bands, self.bands))
self.spd = np.hstack((self.spd, self.spd))
self.kpoints = np.vstack((self.kpoints, self.kpoints))
self.kpoints[0][-1] += 1
print("Atomic plot: bands.shape :", self.bands.shape)
print("Atomic plot: spd.shape :", self.spd.shape)
print("Atomic plot: kpoints.shape:", self.kpoints.shape)
print("Foooooooooooooooooo", self.kpoints)
fig, ax1 = self.parametricPlot(cmap, vmin, vmax, ax=ax)
# plt.gca().xaxis.set_major_locator(plt.NullLocator())
ax1.xaxis.set_major_locator(plt.NullLocator())
# labels on each band
for i in range(len(self.bands[:, 0])):
# print i, self.bands[i]
ax1.text(0, self.bands[i, 0], str(i + 1))
bbox = txt.get_window_extent()
print('bbox', bbox)
return fig, ax1