Added plot_erange_conv method to the VpqRobot (#315)

* sync with abipy/develop

* some beautification for plot_scf_cycle in vpq

* staged some work on the VPQ robot

* staged some work on the VPQ robot: erange filter convergence

abipy/eph/vpq.py

@@ -165,18 +165,20 @@ def params(self) -> dict:
         ngkpt, shifts = ksampling.mpdivs, ksampling.shifts
         nkbz = np.prod(ngkpt)
 
+        avg_g = r.read_variable("vpq_avg_g")[:]
+        e_frohl = r.read_variable("e_frohl")[:] # in Ha
+
         d = dict(
-            nkbz=nkbz,
-            ngkpt=ngkpt,
-            nksmall=min(ngkpt),
-            cbrt_ngkpt=np.cbrt(np.prod(ngkpt)),
-            frohl_ntheta=r.frohl_ntheta,
-            #("invsc_size", 1.0 / (nkbz * ((abu.Ang_Bohr * self.structure.lattice.volume) ** (1/3)))),
+            #nkbz=nkbz,
+            #nksmall=min(ngkpt),
+            #cbrt_nkbz=np.cbrt(nkbz),
+            #frohl_ntheta=r.frohl_ntheta,
+            avg_g = bool(avg_g),
+            e_frohl = e_frohl * abu.Ha_eV,
+            ngkpt=tuple(ngkpt),
+            invsc_linsize = 1. / np.cbrt(nkbz * self.structure.lattice.volume),
         )
 
-        #keys = ["e_pol", "e_el", "e_ph", "e_elph", "eps"]
-        #energies = np.array(scf[nstep - 1], dtype=float) * HA2EV
-
         return d
 
     def __str__(self) -> str:
@@ -246,6 +248,7 @@ class Polaron:
     nq: int            # Number of q-points in B_qnu (including filtering if any).
     bstart: int        # First band starts at bstart.
     bstop: int         # Last band (python convention)
+    erange: double     # Filtering value (in Ha)
     varpeq: VpqFile
 
     @classmethod
@@ -256,6 +259,7 @@ def from_vpq(cls, varpeq: VpqFile, spin: int) -> Polaron:
         r = varpeq.r
         nstates, nk, nq, nb = r.nstates, r.nk_spin[spin], r.nq_spin[spin], r.nb_spin[spin]
         bstart, bstop = r.brange_spin[spin]
+        erange = r.erange_spin[spin]
 
         data = locals()
         return cls(**{k: data[k] for k in [field.name for field in dataclasses.fields(Polaron)]})
@@ -303,6 +307,7 @@ def scf_df_state(self) -> list[pd.DataFrame]:
         # int cvflag_spin(nsppol, nstates) ;
         #   0 --> calculation is not converged
         #   1 --> calculation is converged
+
         spin = self.spin
         r = self.varpeq.r
         nstep2cv = r.read_variable("nstep2cv_spin")[spin]
@@ -326,6 +331,8 @@ def ufact_k(k):
             df = pd.DataFrame(dct)
             # Add metadata to the attrs dictionary
             df.attrs["converged"] = bool(cvflag[pstate])
+            df.attrs["use_filter"] = bool(abs(self.erange) > 1e-8)
+            df.attrs["filter_value"] = self.erange * abu.Ha_eV
             df_list.append(df)
 
         return df_list
@@ -338,9 +345,14 @@ def get_final_results_df(self, with_params: bool = False) -> pd.DataFrame:
         row_list = []
         for pstate in range(self.nstates):
             df = self.scf_df_state[pstate]
-            row = {"pstate": pstate, "spin": self.spin}
+            row = {"formula": self.structure.reduced_formula,
+                   "spgroup": self.structure.get_space_group_info()[1],
+                   "polaron": self.varpeq.r.vpq_pkind,
+                   "pstate": pstate, "spin": self.spin}
             row.update(df.iloc[-1].to_dict())
             row["converged"] = df.attrs["converged"]
+            row["use_filter"] = df.attrs["use_filter"]
+            row["filter_value"] = df.attrs["filter_value"]
             if with_params:
                 row.update(self.varpeq.params)
             row_list.append(row)
@@ -539,15 +551,23 @@ def plot_scf_cycle(self, ax_mat=None, fontsize=8, **kwargs) -> Figure:
                             ax.plot(xs, ys, label=entry.latex, c='k')
                             ax.set_yscale("log")
                     else:
+                        if entry.name == "E_pol":
+                        # Solid line for the *variational* quantity, also put it on top
+                            ls, zord = '-', 10
+                        else:
+                        # Dashed lines for non-variational, put them below
+                            ls, zord = '--', 0
+
                         if iax == 1:
                             energy_like =True
                             # Plot values linear scale.
-                            ax.plot(xs, ys, label=entry.latex)
+                            ax.plot(xs, ys, label=entry.latex, linestyle=ls, zorder=zord)
                         elif iax == 2:
                             energy_like = True
                             # Plot deltas in logscale.
                             # (remove the last point for pretty-plotting)
-                            ax.plot(xs[:-1], np.abs(ys - ys[-1])[:-1], label=entry.latex)
+                            ax.plot(xs[:-1], np.abs(ys - ys[-1])[:-1], label=entry.latex,
+                                    linestyle=ls, zorder=zord)
                             ax.set_yscale("log")
 
                 ax.set_xlim(1, niter)
@@ -1085,7 +1105,7 @@ def __init__(self, filepath: PathLike):
         self.vpq_pkind = self.read_string("vpq_pkind")
         #self.vpq_aseed = self.read_string("vpq_aseed")
         self.ngqpt = self.read_value("gstore_ngqpt")
-        self.frohl_ntheta = self.read_value("frohl_ntheta")
+        #self.frohl_ntheta = self.read_value("frohl_ntheta")
 
         # Read important variables.
         #self.completed = self.read_value("gstore_completed")
@@ -1101,8 +1121,7 @@ def __init__(self, filepath: PathLike):
         self.brange_spin = self.read_value("brange_spin")
         self.brange_spin[:,0] -= 1
         self.nb_spin = self.brange_spin[:,1] - self.brange_spin[:,0]
-
-        #self.erange_spin = self.read_value("gstore_erange_spin")
+        self.erange_spin = self.read_value("erange_spin")
         # Total number of k/q points for each spin after filtering (if any)
         #self.glob_spin_nq = self.read_value("gstore_glob_nq_spin")
         #self.glob_nk_spin = self.read_value("gstore_glob_nk_spin")
@@ -1128,7 +1147,7 @@ class VpqRobot(Robot, RobotWithEbands):
     .. inheritance-diagram:: VpqRobot
     """
 
-    EXT = "VARPEQ"
+    EXT = "VPQ"
 
     def __str__(self) -> str:
         return self.to_string()
@@ -1167,14 +1186,121 @@ def get_final_results_df(self, spin: int = None, sortby: str = None, with_params
         return df
 
     #@add_fig_kwargs
-    #def plot_erange_conv(self, fontsize=12, **kwargs) -> Figure:
-    #    """
-    #    Plot the convergence of the results wrt to the value of erange.
+    def plot_erange_conv(self, ax_mat=None, spin: int = 0, pstate: int = 0,
+                         **kwargs) -> List(Figure):
+        """
+        Plot the convergence of the results wrt to the value of erange.
+
+        Args:
+            fontsize: fontsize for legends and titles
+        """
+
+        df = self.get_final_results_df(spin)
+
+        # check if dataframe contains entries with efilter
+        df = df[df["use_filter"] & (df["pstate"] == pstate)]
+        if df.empty:
+            raise RuntimeError("No entries with energy filtering.")
+
+
+        # Check if df contains information about multuple systems, polarons, etc
+        systems = set(df["formula"])
+        spgroups = set(df["spgroup"])
+        polarons = set(df["polaron"])
+
+        # For each system, spgroup and polaron type, we will plot convergence at a fixed ngkpt
+        # Entries with single filtering value for a fixed ngkpt are skipped
+
+        # count number of plots
+        entries = {}
+        for sys in systems:
+            for spg in spgroups:
+                for pol in polarons:
+                    filtered_df = df[(df["formula"] == sys) &
+                                     (df["spgroup"] == spg) &
+                                     (df["polaron"] == pol)]
+                    entry_keys = (sys, spg, pol)
+
+                    for scell in set(filtered_df["ngkpt"]):
+
+                        count = (df["ngkpt"] == scell).sum()
+                        # only if we ecnounter multiple entries for single scell (for convergence)
+                        if count > 1:
+                            if entry_keys in entries:
+                                entries[entry_keys].append(scell)
+                            else:
+                                entries[entry_keys] = [scell]
+
+        if not entries:
+            raise RuntimeError("Not enough data for convergence with energy filteing.")
+
+        # For each entry, plot convergence wrt erange for each ngkpt
+        fig_list = []
+        for system_keys, scell_list in entries.items():
+
+            formula, spg, pol = system_keys
+
+            entry_df = df[(df["formula"] == formula) &
+                          (df["spgroup"] == spg) &
+                          (df["polaron"] == pol)]
+
+            nrows, ncols = len(scell_list), 2
+            ax_mat, fig, plt = get_axarray_fig_plt(None, nrows=nrows, ncols=ncols,
+                                                   sharex=True, sharey=False, squeeze=False)
+
+            for iax, scell in enumerate(scell_list):
+                scell_df = entry_df[entry_df["ngkpt"] == scell]
+
+                frohich_correction = [True, False]
+
+                for avg_g in frohich_correction:
+                    _df = scell_df[scell_df["avg_g"] == avg_g].sort_values("filter_value")
+
+                    if _df.empty:
+                        continue
+
+                    epol = _df["E_pol"].to_numpy()
+                    eps = _df["epsilon"].to_numpy()
+                    filter = _df["filter_value"].to_numpy()
+
+                    frohlich_label = " + LR correction" if avg_g else ""
+
+                    # Convergence
+                    ax_mat[iax,0].plot(filter, epol, 's-', **kwargs)
+                    ax_mat[iax,0].plot(filter, eps, 's-', **kwargs)
+
+                    # Relative error
+                    ax_mat[iax,1].plot(filter[:-1], np.abs((epol - epol[-1])/epol[-1])[:-1]*100, 's-',
+                                       label=r'$E_{pol}$' + frohlich_label, **kwargs)
+                    ax_mat[iax,1].plot(filter[:-1], np.abs((eps - eps[-1])/eps[-1])[:-1]*100, 's-',
+                                       label=r"$\varepsilon$" + frohlich_label, **kwargs)
+                    ax_mat[iax,1].set_yscale("log")
+
+                    ax_mat[iax,0].set_ylabel("Energy (eV)")
+                    ax_mat[iax,1].set_ylabel("Relative error (%)")
+
+                    ax_mat[iax,1].legend(title=f"k-mesh = {scell}")
+
+                    for icol in range(ncols):
+                        ax_mat[iax,icol].set_xlim(0)
+                        ax_mat[iax,icol].grid()
+
+
+            ax_mat[0,0].set_title("Energy convergence")
+            ax_mat[0,1].set_title("Relative error")
+            for icol in range(ncols):
+                ax_mat[nrows-1,icol].set_xlabel("Filter value (eV)")
+
+            title = f"{formula}, space group {spg}, {pol} polaron"
+            fig.suptitle(title)
+            fig.tight_layout()
+
+            fig_list.append(fig)
+
+        return fig_list
+
+
 
-    #    Args:
-    #        colormap: Color map. Have a look at the colormaps here and decide which one you like:
-    #        fontsize: fontsize for legends and titles
-    #    """
     #    fig = self.plot_convergence(self, item: Union[str, Callable],
     #                                sortby=None, hue=None, abs_conv=None,
     #                                ax=None, fontsize=8, **kwargs)