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)