def plot_ax(self, ax, exchange_xy=False, *args, **kwargs):
"""
Helper function to plot self on axis ax.
Args:
ax: `matplotlib` axis.
exchange_xy: True to exchange the axis in the plot
args: Positional arguments passed to ax.plot
kwargs: Keyword arguments passed to `matplotlib`. Accepts also:
============== ===============================================================
kwargs Meaning
============== ===============================================================
cplx_mode string defining the data to print.
Possible choices are (case-insensitive): `re` for the real part
"im" for the imaginary part, "abs" for the absolute value.
"angle" to display the phase of the complex number in radians.
Options can be concatenated with "-"
============== ===============================================================
Returns:
List of lines added.
"""
cplx_mode = kwargs.pop("cplx_mode", "re-im")
lines = []
from abipy.tools.plotting_utils import data_from_cplx_mode
for c in cplx_mode.lower().split("-"):
xx, yy = self.mesh, data_from_cplx_mode(c, self.values)
if exchange_xy:
xx, yy = yy, xx
lines.extend(ax.plot(xx, yy, *args, **kwargs))
return lines
def plot_w(self, gvec1, gvec2=None, waxis="real", cplx_mode="re-im", ax=None, **kwargs):
"""
Plot the frequency dependence of W_{g1, g2}
Args:
gvec1, gvec2:
waxis: "real" to plot along the real axis, "imag" for the imaginary axis.
cplx_mode: string defining the data to print.
Possible choices are (case-insensitive): `re` for the real part
"im" for the imaginary part, "abs" for the absolute value.
"angle" will display the phase of the complex number in radians.
Options can be concatenated with "-" e.g. "re-im"
ax: matplotlib :class:`Axes` or None if a new figure should be created.
Returns:
matplotlib figure.
"""
# Select data to plot
ig1 = self.gindex(gvec1)
ig2 = ig1 if gvec2 is None else self.gindex(gvec2)
ax, fig, plt = get_ax_fig_plt(ax)
if waxis == "real":
if self.nrew == 0: return fig
xx = (self.real_wpts * Ha_to_eV).real
yy = self.wggmat_realw[:, ig1, ig2]
elif waxis == "imag":
if self.nimw == 0: return fig
xx = (self.imag_wpts * Ha_to_eV).imag
yy = self.wggmat_imagw[:, ig1, ig2]
else:
raise ValueError("Wrong value for waxis: %s" % waxis)
color_cmode = dict(re="red", im="blue", abs="black", angle="green")
linewidth = kwargs.pop("linewidth", 2)
linestyle = kwargs.pop("linestyle", "solid")
lines = []
for c in cplx_mode.lower().split("-"):
l, = ax.plot(xx, data_from_cplx_mode(c, yy),
color=color_cmode[c], linewidth=linewidth,
linestyle=linestyle,
label=self.latex_label(c))
lines.append(l)
ax.grid(True)
ax.set_xlabel("$\omega$ [eV]")
ax.set_title("%s, qpoint: %s" % (self.etsf_name, self.qpoint))
#ax.legend(loc="best")
ax.legend(loc="upper right")
return fig
def plot_ggmat(self, cplx_mode="abs", wpos=None, **kwargs):
"""
Use imshow for plotting W_GG'.
Args:
cplx_mode:
wpos: List of frequency indices to plot. If None, the first
frequency is used (usually w=0). if wpos == "all"
all frequencies are shown (use it carefully)
Other possible values: "real" if only real frequencies are wanted.
"imag" for imaginary frequencies only.
Returns:
`matplotlib` figure.
"""
# Get wpos indices.
choice_wpos = {
None: [0],
"all": range(self.nw),
"real": range(self.nrew),
"imag": range(self.nrew, self.nw)
}
if any(wpos == k for k in choice_wpos):
wpos = choice_wpos[wpos]
else:
if isinstance(wpos, int): wpos = [wpos]
wpos = np.array(wpos)
# Build plotter.
plotter = ArrayPlotter()
for iw in wpos:
label = cplx_mode + " $\omega = %s" % self.wpts[iw]
data = data_from_cplx_mode(cplx_mode, self.wggmat[iw, :, :])
plotter.add_array(label, data)
return plotter.plot(**kwargs)
def plot_ggmat(self, cplx_mode="abs", wpos=None, **kwargs):
"""
Use imshow for plotting W_GG'.
Args:
cplx_mode:
wpos: List of frequency indices to plot. If None, the first
frequency is used (usually w=0). if wpos == "all"
all frequencies are shown (use it carefully)
Other possible values: "real" if only real frequencies are wanted.
"imag" for imaginary frequencies only.
Returns:
`matplotlib` figure.
"""
# Get wpos indices.
choice_wpos = {
None: [0],
"all": range(self.nw),
"real": range(self.nrew),
"imag": range(self.nrew, self.nw)}
if any(wpos == k for k in choice_wpos):
wpos = choice_wpos[wpos]
else:
if isinstance(wpos, int): wpos = [wpos]
wpos = np.array(wpos)
# Build plotter.
plotter = ArrayPlotter()
for iw in wpos:
label = cplx_mode + " $\omega = %s" % self.wpts[iw]
data = data_from_cplx_mode(cplx_mode, self.wggmat[iw,:,:])
plotter.add_array(label, data)
return plotter.plot(**kwargs)
def plot_w(self,
gvec1,
gvec2=None,
waxis="real",
cplx_mode="re-im",
ax=None,
**kwargs):
"""
Plot the frequency dependence of W_{g1, g2}
Args:
gvec1, gvec2:
waxis: "real" to plot along the real axis, "imag" for the imaginary axis.
cplx_mode: string defining the data to print.
Possible choices are (case-insensitive): `re` for the real part
"im" for the imaginary part, "abs" for the absolute value.
"angle" will display the phase of the complex number in radians.
Options can be concatenated with "-" e.g. "re-im"
ax: matplotlib :class:`Axes` or None if a new figure should be created.
Returns:
matplotlib figure.
"""
# Select data to plot
ig1 = self.gindex(gvec1)
ig2 = ig1 if gvec2 is None else self.gindex(gvec2)
ax, fig, plt = get_ax_fig_plt(ax)
if waxis == "real":
if self.nrew == 0: return fig
xx = (self.real_wpts * Ha_to_eV).real
yy = self.wggmat_realw[:, ig1, ig2]
elif waxis == "imag":
if self.nimw == 0: return fig
xx = (self.imag_wpts * Ha_to_eV).imag
yy = self.wggmat_imagw[:, ig1, ig2]
else:
raise ValueError("Wrong value for waxis: %s" % waxis)
color_cmode = dict(re="red", im="blue", abs="black", angle="green")
linewidth = kwargs.pop("linewidth", 2)
linestyle = kwargs.pop("linestyle", "solid")
lines = []
for c in cplx_mode.lower().split("-"):
l, = ax.plot(xx,
data_from_cplx_mode(c, yy),
color=color_cmode[c],
linewidth=linewidth,
linestyle=linestyle,
label=self.latex_label(c))
lines.append(l)
ax.grid(True)
ax.set_xlabel("$\omega$ [eV]")
ax.set_title("%s, qpoint: %s" % (self.etsf_name, self.qpoint))
#ax.legend(loc="best")
ax.legend(loc="upper right")
return fig
