Python plotting_utils.get_ax_fig_plt函数代码示例

    def plot_hist(self, struct_type, ax=None, **kwargs):
        """
        Histogram plot.
        """
        #if codes is None: codes = ["ae"]
        ax, fig, plt = get_ax_fig_plt(ax)
        import seaborn as sns

        codes = ["this", "gbrv_paw"] #, "gbrv_uspp", "pslib", "vasp"]
        new = self[self["struct_type"] == struct_type].copy()
        for code in codes:
            values = (100 * (new[code] - new["ae"]) / new["ae"]).dropna()
            sns.distplot(values, ax=ax, rug=True, hist=False, label=code)

            if code == "this":
                # Add text with Mean or (MARE/RMSRE)
                text = []; app = text.append
                app("MARE = %.2f" % values.abs().mean())
                app("RMSRE = %.2f" % np.sqrt((values**2).mean()))
                ax.text(0.8, 0.8, "\n".join(text), transform=ax.transAxes)

        ax.grid(True)
        ax.set_xlabel("relative error %")
        ax.set_xlim(-0.8, 0.8)

        return fig

    def plot_errors_for_elements(self, ax=None, **kwargs):
        """
        Plot the relative errors associated to the chemical elements.
        """
        dict_list = []
        for idx, row in self.iterrows():
            rerr = 100 * (row["this"] - row["ae"]) / row["ae"]
            for symbol in set(species_from_formula(row.formula)):
                dict_list.append(dict(
                    element=symbol,
                    rerr=rerr,
                    formula=row.formula,
                    struct_type=row.struct_type,
                    ))

        frame = DataFrame(dict_list)
        order = sort_symbols_by_Z(set(frame["element"]))
        #print_frame(frame)

        import seaborn as sns
        ax, fig, plt = get_ax_fig_plt(ax=ax)

        # Draw violinplot
        #sns.violinplot(x="element", y="rerr", order=order, data=frame, ax=ax, orient="v")

        # Box plot
        ax = sns.boxplot(x="element", y="rerr", data=frame, ax=ax, order=order, whis=np.inf, color="c")
        # Add in points to show each observation
        sns.stripplot(x="element", y="rerr", data=frame, ax=ax, order=order,
                      jitter=True, size=5, color=".3", linewidth=0)

        sns.despine(left=True)
        ax.set_ylabel("Relative error %")
        ax.grid(True)
        return fig

    def plot_qpgaps(self, ax=None, spin=None, kpoint=None, hspan=0.01, **kwargs):
        """
        Plot the QP gaps as function of the convergence parameter.

        Args:
            ax: matplotlib :class:`Axes` or None if a new figure should be created.
            spin:
            kpoint: 
            hspan:
            kwargs:

        Returns:
            `matplotlib` figure
        """
        spin_range = range(self.nsppol) if spin is None else torange(spin)
        kpoints_for_plot = self.computed_gwkpoints  #if kpoint is None else KpointList.as_kpoints(kpoint)

        self.prepare_plot()

        ax, fig, plt = get_ax_fig_plt(ax)

        xx = self.xvalues
        for spin in spin_range:
            for kpoint in kpoints_for_plot:
                label = "spin %d, kpoint %s" % (spin, repr(kpoint))
                gaps = self.extract_qpgaps(spin, kpoint)
                ax.plot(xx, gaps, "o-", label=label, **kwargs)

                if hspan is not None:
                    last = gaps[-1]
                    ax.axhspan(last-hspan, last+hspan, facecolor='0.5', alpha=0.5)

        self.decorate_ax(ax)
        return fig

    def plot_den_formfact(self, ecut=60, ax=None, **kwargs):
        """
        Plot the density form factor as function of ecut (Ha units). Return matplotlib Figure.
        """
        ax, fig, plt = get_ax_fig_plt(ax)

        lines, legends = [], []
        for name, rho in self.densities.items():
            if name == "rhoC": continue
            form = rho.get_intr2j0(ecut=ecut) / (4 * np.pi)
            line, = ax.plot(form.mesh, form.values, linewidth=self.linewidth, markersize=self.markersize)
            lines.append(line); legends.append(name)

            intg = rho.r2f_integral()[-1]
            print("r2 f integral: ", intg)
            print("form_factor(0): ", name, form.values[0])

        # Plot vloc(q)
        #for l, pot in self.potentials.items():
        #    if l != -1: continue
        #    form = pot.get_intr2j0(ecut=ecut)
        #    mask = np.where(np.abs(form.values) > 20); form.values[mask] = 20
        #    line, = ax.plot(form.mesh, form.values, linewidth=self.linewidth, markersize=self.markersize)
        #    lines.append(line); legends.append("Vloc(q)")

        decorate_ax(ax, xlabel="Ecut [Ha]", ylabel="$n(q)$", title="Form factor, l=0 ", lines=lines, legends=legends)
        return fig

    def plot_tcore_rspace(self, ax=None, ders=(0, 1, 2, 3), rmax=3.0,  **kwargs):
        """
        Plot the model core and its derivatives in real space.

        Args:
            ax: matplotlib :class:`Axes` or None if a new figure should be created.
            ders: Tuple used to select the derivatives to be plotted.
            rmax: Max radius for plot in Bohr. None is full grid is wanted.

        Returns:
            matplotlib figure.
        """
        ax, fig, plt = get_ax_fig_plt(ax)

        linewidth = kwargs.pop("linewidth", 2.0)
        rmeshes, coresd = self.reader.read_coresd(rmax=rmax)

        for rmesh, mcores in zip(rmeshes, coresd):
            for der, values in enumerate(mcores):
                if der not in ders: continue
                yvals, fact, = rescale(values)
                ax.plot(rmesh, yvals, color=self.color_der[der], linewidth=linewidth,
                        linestyle=self.linestyles_der[der],
                        label=mklabel("\\tilde{n}_c", der, "r") + " x %.4f" % fact)

        ax.grid(True)
        ax.set_xlabel("r [Bohr]")
        ax.set_title("Model core in r-space")
        if kwargs.get("with_legend", False): ax.legend(loc="best")

        return fig

    def plot_der_densities(self, ax=None, order=1, **kwargs):
        """
        Plot the derivatives of the densitiers on axis ax.
        Used to analyze possible derivative discontinuities
        """
        ax, fig, plt = get_ax_fig_plt(ax)

        from scipy.interpolate import UnivariateSpline

        lines, legends = [], []
        for name, rho in self.densities.items():
            if name != "rhoM": continue
            # Need linear mesh for finite_difference --> Spline input densities on lin_rmesh
            lin_rmesh, h = np.linspace(rho.rmesh[0], rho.rmesh[-1], num=len(rho.rmesh) * 4, retstep=True)
            spline = UnivariateSpline(rho.rmesh, rho.values, s=0)
            lin_values = spline(lin_rmesh)
            vder = finite_diff(lin_values, h, order=order, acc=4)
            line, = ax.plot(lin_rmesh, vder) #, **self._wf_pltopts(l, "ae"))
            lines.append(line)

            legends.append("%s-order derivative of %s" % (order, name))

        decorate_ax(ax, xlabel="r [Bohr]", ylabel="$D^%s \n(r)$" % order, title="Derivative of the charge densities",
                    lines=lines, legends=legends)
        return fig

    def plot_radial_wfs(self, ax=None, **kwargs):
        """
        Plot ae and ps radial wavefunctions on axis ax.

        lselect: List to select l channels.
        """
        ax, fig, plt = get_ax_fig_plt(ax)

        ae_wfs, ps_wfs = self.radial_wfs.ae, self.radial_wfs.ps
        lselect = kwargs.get("lselect", [])

        lines, legends = [], []
        for nlk, ae_wf in ae_wfs.items():
            ps_wf, l, k = ps_wfs[nlk], nlk.l, nlk.k
            if l in lselect: continue
            #print(nlk)

            ae_line, = ax.plot(ae_wf.rmesh, ae_wf.values, **self._wf_pltopts(l, "ae"))
            ps_line, = ax.plot(ps_wf.rmesh, ps_wf.values, **self._wf_pltopts(l, "ps"))

            lines.extend([ae_line, ps_line])
            if k is None:
                legends.extend(["AE l=%s" % l, "PS l=%s" % l])
            else:
                legends.extend(["AE l=%s, k=%s" % (l, k), "PS l=%s, k=%s" % (l, k)])

        decorate_ax(ax, xlabel="r [Bohr]", ylabel="$\phi(r)$", title="Wave Functions",
                    lines=lines, legends=legends)

        return fig

    def plot_key(self, key, ax=None, **kwargs):
        """Plot a singol quantity specified by key."""
        ax, fig, plt = get_ax_fig_plt(ax)

        # key --> self.plot_key()
        getattr(self, "plot_" + key)(ax=ax, **kwargs)
        self._mplt.show()

    def plot_der_potentials(self, ax=None, order=1, **kwargs):
        """
        Plot the derivatives of vl and vloc potentials on axis ax.
        Used to analyze the derivative discontinuity introduced by the RRKJ method at rc.
        """
        ax, fig, plt = get_ax_fig_plt(ax)
        from abipy.tools.derivatives import finite_diff
        from scipy.interpolate import UnivariateSpline
        lines, legends = [], []
        for l, pot in self.potentials.items():
            # Need linear mesh for finite_difference --> Spline input potentials on lin_rmesh
            lin_rmesh, h = np.linspace(pot.rmesh[0], pot.rmesh[-1], num=len(pot.rmesh) * 4, retstep=True)
            spline = UnivariateSpline(pot.rmesh, pot.values, s=0)
            lin_values = spline(lin_rmesh)
            vder = finite_diff(lin_values, h, order=order, acc=4)
            line, = ax.plot(lin_rmesh, vder, **self._wf_pltopts(l, "ae"))
            lines.append(line)

            if l == -1:
                legends.append("%s-order derivative Vloc" % order)
            else:
                legends.append("$s-order derivative PS l=%s" % str(l))

        decorate_ax(ax, xlabel="r [Bohr]", ylabel="$D^%s \phi(r)$" % order,
                    title="Derivative of the ion Pseudopotentials",
                    lines=lines, legends=legends)
        return fig

    def cpuwall_histogram(self, ax=None, **kwargs):
        ax, fig, plt = get_ax_fig_plt(ax=ax)

        nk = len(self.sections)
        ind = np.arange(nk)  # the x locations for the groups
        width = 0.35         # the width of the bars

        cpu_times = self.get_values("cpu_time")
        rects1 = plt.bar(ind, cpu_times, width, color='r')

        wall_times = self.get_values("wall_time")
        rects2 = plt.bar(ind + width, wall_times, width, color='y')

        # Add ylable and title
        ax.set_ylabel('Time (s)')

        #if title:
        #    plt.title(title)
        #else:
        #    plt.title('CPU-time and Wall-time for the different sections of the code')

        ticks = self.get_values("name")
        ax.set_xticks(ind + width, ticks)

        ax.legend((rects1[0], rects2[0]), ('CPU', 'Wall'), loc="best")

        return fig

    def plot(self, cplx_mode="abs", color_map="jet", **kwargs):
        """
        Args:
            cplx_mode: "abs" for absolute value, "re", "im", "angle"
            color_map: matplotlib colormap
        """
        ax, fig, plt = get_ax_fig_plt(None)
        plt.axis("off")

        # Grid parameters
        num_plots, ncols, nrows = len(self), 1, 1
        if num_plots > 1:
            ncols = 2
            nrows = (num_plots//ncols) + (num_plots % ncols)

        # use origin to place the [0,0] index of the array in the lower left corner of the axes. 
        for n, (label, arr) in enumerate(self.items()):
            fig.add_subplot(nrows, ncols, n)
            data = data_from_cplx_mode(cplx_mode, arr)
            img = plt.imshow(data, interpolation='nearest', cmap=color_map, origin='lower')

            # make a color bar
            plt.colorbar(img, cmap=color_map)
            plt.title(label + " (%s)" % cplx_mode)

            # Set grid
            plt.grid(True, color='white')


        return fig

    def plot(self, ax=None, **kwargs):
        """
        Uses Matplotlib to plot the energy curve.

        Args:
            ax: :class:`Axes` object. If ax is None, a new figure is produced.

        ================  ==============================================================
        kwargs            Meaning
        ================  ==============================================================
        style             
        color
        text
        label
        ================  ==============================================================

        Returns:
            Matplotlib figure.
        """
        ax, fig, plt = get_ax_fig_plt(ax)

        vmin, vmax = self.volumes.min(), self.volumes.max()
        emin, emax = self.energies.min(), self.energies.max()

        vmin, vmax = (vmin - 0.01 * abs(vmin), vmax + 0.01 * abs(vmax))
        emin, emax = (emin - 0.01 * abs(emin), emax + 0.01 * abs(emax))

        color = kwargs.pop("color", "r")
        label = kwargs.pop("label", None)

        # Plot input data.
        ax.plot(self.volumes, self.energies, linestyle="None", marker="o", color=color) #, label="Input Data")

        # Plot EOS.
        vfit = np.linspace(vmin, vmax, 100)
        if label is None:
            label = self.name + ' fit'

        if self.eos_name == "deltafactor":
            xx = vfit**(-2./3.)
            ax.plot(vfit, np.polyval(self.eos_params, xx), linestyle="dashed", color=color, label=label)
        else:
            ax.plot(vfit, self.func(vfit, *self.eos_params), linestyle="dashed", color=color, label=label)

        # Set xticks and labels.
        ax.grid(True)
        ax.set_xlabel("Volume $\AA^3$")
        ax.set_ylabel("Energy (eV)")

        ax.legend(loc="best", shadow=True)

        # Add text with fit parameters.
        if kwargs.pop("text", True):
            text = []; app = text.append
            app("Min Volume = %1.2f $\AA^3$" % self.v0)
            app("Bulk modulus = %1.2f eV/$\AA^3$ = %1.2f GPa" % (self.b0, self.b0_GPa))
            app("B1 = %1.2f" % self.b1)
            fig.text(0.4, 0.5, "\n".join(text), transform=ax.transAxes)

        return fig

    def plot(self, ax=None, *args, **kwargs):
        """
        Plot all the components of the tensor

        Args:
            ax: matplotlib :class:`Axes` or None if a new figure should be created.

        ==============  ==============================================================
        kwargs          Meaning
        ==============  ==============================================================
        red_coords      True to plot the reduced coordinate tensor (Default: True)
        ==============  ==============================================================

        Returns:
            matplotlib figure
        """
        red_coords = kwargs.pop("red_coords", True)
        ax, fig, plt = get_ax_fig_plt(ax)

        ax.grid(True)
        ax.set_xlabel('Frequency [eV]')
        ax.set_ylabel('Dielectric tensor')

        #if not kwargs:
        #    kwargs = {"color": "black", "linewidth": 2.0}

        # Plot the 6 independent components
        for icomponent in [0,4,8,1,2,5]:
            self.plot_ax(ax, icomponent, red_coords, *args, **kwargs)

        return fig

    def plot_ffspl(self, ax=None, ecut_ffnl=None, ders=(0,), with_qn=0, with_fact=False, **kwargs):
        """
        Plot the nonlocal part of the pseudopotential in q space.

        Args:
            ax: matplotlib :class:`Axes` or None if a new figure should be created.
            ecut_ffnl: Max cutoff energy for ffnl plot (optional)
            ders: Tuple used to select the derivatives to be plotted.
            with_qn:

        Returns:
            matplotlib figure.
        """
        ax, fig, plt = get_ax_fig_plt(ax)

        color = kwargs.pop("color", "black")
        linewidth = kwargs.pop("linewidth", 2.0)

        color_l = {-1: "black", 0: "red", 1: "blue", 2: "green", 3: "orange"}
        linestyles_n = ["solid", '-', '--', '-.', ":"]
        scale = None
        l_seen = set()

        qmesh, vlspl = self.reader.read_vlspl()

        all_projs = self.reader.read_projectors()
        for itypat, projs_type in enumerate(all_projs): 
            # Loop over the projectors for this atom type.
            for p in projs_type:
                for der, values in enumerate(p.data):
                    if der == 1: der = 2
                    if der not in ders: continue
                    #yvals, fact = rescale(values, scale=scale)
                    label = None
                    if p.l not in l_seen:
                        l_seen.add(p.l)
                        label = mklabel("v_{nl}", der, "q") + ", l=%d" % p.l

                    stop = len(p.ecuts)
                    if ecut_ffnl is not None:
                        stop = find_gt(p.ecuts, ecut_ffnl)

                    #values = p.ekb * p.values - vlspl[itypat, 0, :]
                    values = vlspl[itypat, 0, :] + p.sign_sqrtekb * p.values

                    #print(values.min(), values.max())
                    ax.plot(p.ecuts[:stop], values[:stop], color=color_l[p.l], linewidth=linewidth, 
                            linestyle=linestyles_n[p.n], label=label)

        ax.grid(True)
        ax.set_xlabel("Ecut [Hartree]")
        ax.set_title("ffnl(q)")
        if kwargs.get("with_legend", True):
            ax.legend(loc="best")

        ax.axhline(y=0, linewidth=linewidth, color='k', linestyle="solid")

        return fig

    def plot_efficiency(self, key="wall_time", what="gb", nmax=5, ax=None, **kwargs):
        ax, fig, plt = get_ax_fig_plt(ax=ax)

        timers = self.timers()
        peff = self.pefficiency()

        # Table with the parallel efficiency for all the sections.
        #pprint_table(peff.totable())

        n = len(timers)
        xx = np.arange(n)

        ax.set_color_cycle(['g', 'b', 'c', 'm', 'y', 'k'])

        legend_entries = []

        # Plot sections with good efficiency.
        lines = []
        if "g" in what:
            good = peff.good_sections(key=key, nmax=nmax)
            for g in good:
                #print(g, peff[g])
                yy = peff[g][key]
                line, = ax.plot(xx, yy, "-->", linewidth=3.0, markersize=10)
                lines.append(line)
                legend_entries.append(g)

        # Plot sections with bad efficiency.
        if "b" in what:
            bad = peff.bad_sections(key=key, nmax=nmax)
            for b in bad:
                #print(b, peff[b])
                yy = peff[b][key]
                line, = ax.plot(xx, yy, "-.<", linewidth=3.0, markersize=10)
                lines.append(line)
                legend_entries.append(b)

        if "total" not in legend_entries:
            yy = peff["total"][key]
            total_line, = ax.plot(xx, yy, "r", linewidth=3.0, markersize=10)
            lines.append(total_line)
            legend_entries.append("total")

        ax.legend(lines, legend_entries, loc="best", shadow=True)

        #ax.set_title(title)
        ax.set_xlabel('Total_NCPUs')
        ax.set_ylabel('Efficiency')
        ax.grid(True)

        # Set xticks and labels.
        labels = ["MPI = %d, OMP = %d" % (t.mpi_nprocs, t.omp_nthreads) for t in timers]
        ax.set_xticks(xx)
        ax.set_xticklabels(labels, fontdict=None, minor=False, rotation=15)

        return fig