From 8f47905846a75510900a0fcea9183e534a64b806 Mon Sep 17 00:00:00 2001 From: lmacenul <72354884+lmacenul@users.noreply.github.com> Date: Thu, 30 Nov 2023 18:24:47 +0100 Subject: [PATCH] LRUJ plot (#280) * Here's the version that's giving me the error I sent in Teams. * Mmkay * A working file now. Generates a plot. * Ploot modifications. * Changed this file, actually... --- abipy/electrons/lruj.py | 167 ++++++++++++++++++++++++++++++++-------- 1 file changed, 134 insertions(+), 33 deletions(-) diff --git a/abipy/electrons/lruj.py b/abipy/electrons/lruj.py index 24e8021bd..6c1f10c25 100644 --- a/abipy/electrons/lruj.py +++ b/abipy/electrons/lruj.py @@ -44,23 +44,31 @@ # self.r = r = EtsfReader(filepath) - +#=============================================================================================================== +#=============================================================================================================== @dataclasses.dataclass(kw_only=True) class LrujResults: """ This object stores the results produced by lruj. """ + natom: int npert: int - maxdeg: int + ndata: int pawujat: int macro_uj: int - dmatpuopt: int + diem_token: str diem: float alphas: np.ndarray occ_unscr: np.ndarray occ_scr: np.ndarray chi0_coefficients: dict chi_coefficients: dict + maxdeg: int + dmatpuopt: int + pert_name: str + parname: str + metric: str + fit_df: pd.DataFrame @classmethod def from_file(cls, filepath: PathLike): @@ -81,12 +89,27 @@ def from_file(cls, filepath: PathLike): if in_doc and line.startswith("..."): data = yaml_safe_load("".join(yaml_lines)) - print("data:", data) break if in_doc: yaml_lines.append(line) + natom = data['natom'] + ndata = data['ndata'] + pawujat = data['pawujat'] + macro_uj = data['macro_uj'] + diem_token = data['diem_token'] + diem = data['diem'] + npert = ndata - 1 + if macro_uj==4: + pert_name = r"$\beta$" + metric = r"M $(n^{\uparrow} - n^{\downarrow})$" + parname = 'J' + else: + pert_name = r"$\alpha$" + metric = r"N $(n^{\uparrow} + n^{\downarrow})$" + parname = 'U' + chi0_coefficients = {} chi_coefficients = {} for k, v in data.items(): @@ -99,7 +122,6 @@ def from_file(cls, filepath: PathLike): degree = int(k.replace(magic, "")) chi_coefficients[degree] = v - #print(f"{chi0_coefficients=}") #print(f"{chi_coefficients=}") def find(header, dtype=None): @@ -110,12 +132,8 @@ def find(header, dtype=None): return i, after raise ValueError(f"Cannot find {header=} in {filepath=}") - _, npert = find("Number of perturbations detected:", dtype=int) _, maxdeg = find("Maximum degree of polynomials analyzed:", dtype=int) - _, pawujat = find("Index of perturbed atom:", dtype=int) - _, macro_uj = find("Value of macro_uj:", dtype=int) _, dmatpuopt = find("Value of dmatpuopt:", dtype=int) - _, diem = find("Mixing constant factored out of Chi0:", dtype=float) # Parse the section with perturbations and occupations. """ @@ -126,18 +144,25 @@ def find(header, dtype=None): 0.0000000000 8.6380182458 8.6380182458 -0.1500000676 8.6964981922 8.6520722003 + -OR- + + Perturbations Magnetizations + --------------- ----------------------------- + beta [eV] Unscreened Screened + --------------- ----------------------------- + """ - i, _ = find("alpha [eV] Unscreened Screened") - i += 2 + i, _ = find("Perturbations",dtype=None) + i += 4 vals = [] - for ipert in range(npert): + for ipert in range(ndata): vals.append([float(t) for t in lines[i+ipert].split()]) - vals = np.reshape(vals, (npert, 3)) + vals = np.reshape(vals, (ndata, 3)) alphas, occ_unscr, occ_scr = vals[:,0], vals[:,1], vals[:,2] """ RMS Errors --------------------------------------- - Regression Chi0 [eV^-1] Chi [eV^-1] U [eV] | Chi0 [eV^-1] Chi [eV^-1] U [eV] + Regression Chi0 [eV^-1] Chi [eV^-1] J [eV] | Chi0 [eV^-1] Chi [eV^-1] J [eV] --------------------------------------------------------|--------------------------------------- Linear: -0.8594082 -0.0949434 9.3689952 | 0.0023925 0.0000878 0.1139297 Quadratic: -0.8574665 -0.0955791 9.2963099 | 0.0023777 0.0000129 0.0681722 @@ -146,53 +171,129 @@ def find(header, dtype=None): """ i, _ = find("Regression Chi0 [eV^-1]") i += 2 - keys = ["Chi0", "Chi", "U", "rms_Chi0", "rms_Chi", "rms_U"] + keys = ["Chi0", "Chi", "HP", "rms_Chi0", "rms_Chi", "rms_HP"] dict_list = [] for irow in range(maxdeg): l = lines[i+irow].replace("|", " ") tokens = l.split() - d = dict(zip(keys, [float(t) for t in tokens[1:]])) + d = dict(zip(keys, [float(t) for t in tokens[-6:]])) d["degree"] = irow + 1 dict_list.append(d) fit_df = pd.DataFrame(dict_list) - #print("fit_df:\n", fit_df) # Build instance from locals dict. - data = locals() - return cls(**{k: data[k] for k in [field.name for field in dataclasses.fields(cls)]}) + _data = locals() + return cls(**{k: _data[k] for k in [field.name for field in dataclasses.fields(cls)]}) +#=============================================================================================================== +#=============================================================================================================== @add_fig_kwargs - def plot(self, ax=None, fontsize=12, **kwargs) -> Figure: + def plot(self, ax=None, degrees="all", inset=True, insetdegree=1, insetlocale="lower left", ptcolor0='k', ptcolor='k', gradcolor1='#3575D5',gradcolor2='#FDAE7B', ptitle="default", fontsize=12, **kwargs) -> Figure: """ Plot Args: ax: |matplotlib-Axes| or None if a new figure should be created. + degrees: List of degrees to plot. If None, no polynomials are plotted. + ptcolor0: Color of unscreened response data points (default: black) + ptcolor: Color of screened response data points (default: black) + gradcolor1: Hex code of linear regression color (default: Blue #3575D5) + gradcolor2: Hex code of color of regression of maximum degree in list (default: Salmon #FDAE7B) + ptitle: Plot title (default: "Linear Response for atom ) + inset: Plots inset with LR parameter for polynomial fit of degree (default: True) + insetdegree: Degree of polynomial fit information printed in the inset (default: 1) + insetlocale: Position of inset in the plot. Standard matplotlob locations. (default: "lower left") """ + import seaborn as sns + sns.set() ax, fig, plt = get_ax_fig_plt(ax=ax) - # Plot raw data. - ax.scatter(self.alphas, self.occ_unscr, marker="o", c="b", label="Unscreened") - ax.scatter(self.alphas, self.occ_scr, marker="x", c="r", label="Screened") - ax.axvline(x=0.0, color="k", linestyle="--", lw=0.5) + # Plot data + yshift = self.occ_unscr[np.where(self.alphas == 0.0000)] * (self.diem - 1.0) + data0 = 1.0/self.diem * (self.occ_unscr + yshift) + ax.scatter(self.alphas, data0, s=70, color=ptcolor0, facecolors='none', linewidths=2, label="Unscreened") + ax.scatter(self.alphas, self.occ_scr, s=70, color=ptcolor, label="Screened") + ax.axvline(x=0.0, color="white", linestyle="--", lw=0.5) - # Plot regression fit (only linear part) - xstart, xstop = self.alphas.min(), self.alphas.max() - xstart, xstop = 0.9 * xstart, 1.1 * xstop + # Generate mesh for polynomial functions + xstart, xstop = 1.1 * self.alphas.min(), 1.1 * self.alphas.max() xs = np.arange(xstart, xstop, step=0.01) - #for ideg in range(maxdeg): - # ax.plot(xs, lin_coeff * xs, color=, label="Degree {ideg}") + + #Prepare colors and coefficients for polynomials the use wants to plot + if degrees == "all": + degrees = self.chi0_coefficients.keys() + + def hex_to_RGB(hex_str): + return [int(hex_str[i:i+2], 16) for i in range(1,6,2)] + + def get_color_gradient(c1, c2, n): + assert n > 1 + c1_rgb = np.array(hex_to_RGB(c1))/255 + c2_rgb = np.array(hex_to_RGB(c2))/255 + mix_pcts = [x/(n-1) for x in range(n)] + rgb_colors = [((1-mix)*c1_rgb + (mix*c2_rgb)) for mix in mix_pcts] + return ["#" + "".join([format(int(round(val*255)), "02x") for val in item]) for item in rgb_colors] + + linecolors=get_color_gradient(gradcolor1,gradcolor2,len(degrees)) + + # Plot polynomial functions + def poly0(coeffs): + return lambda x: 1.0/self.diem * (sum((coeff*x**i for i,coeff in enumerate(coeffs))) + yshift) + + def poly(coeffs): + return lambda x: sum((coeff*x**i for i,coeff in enumerate(coeffs))) + + for degree in degrees: + polynomial0=poly0(self.chi0_coefficients[degree]) + polynomial=poly(self.chi_coefficients[degree]) + if degree == 1: + Labelstring='Linear' + elif degree == 2: + Labelstring='Quadratic' + elif degree == 3: + Labelstring='Cubic' + else: + Labelstring=' '.join(['Degree',str(degree)]) + + if insetdegree==degree: + deginfo = ' '.join(['Parameters for',Labelstring,'fit']) + insetcolor = linecolors[degree-1] + + ax.plot(xs,polynomial0(xs),color=linecolors[degree-1],linewidth=2.0,linestyle='dashed') + ax.plot(xs,polynomial(xs),color=linecolors[degree-1],linewidth=2.0,label=Labelstring) ax.legend(loc="best", fontsize=fontsize, shadow=True) + if ptitle=="default": + ptitle=' '.join(["Linear Response",self.parname,"on atom",str(self.pawujat)]) + plt.title(ptitle) ax.grid(True) - ax.set_xlabel(r"$\alpha$ (eV)") - #ylabel = r"$N(n^{\up} + n^{\down})" is U else r"$N(n^{\up} - n^{\down})" - #ax.set_ylabel(ylabel) - + ax.set_xlabel(' '.join([self.pert_name,"(eV)"])) + ax.set_ylabel(self.metric) + + # Generate inset with numerical information on LR + if inset: + from matplotlib.offsetbox import AnchoredText + select = self.fit_df["degree"] == insetdegree + def dfvalue(keywo): + tablevalue = self.fit_df[select][keywo] + return "%.4f" % tablevalue.values[0] + + X0str = ' '.join([r"$\chi_0$","=",dfvalue("Chi0"),r"$\pm$",dfvalue("rms_Chi0"),r"[eV]$^{-1}$"]) + Xstr = ' '.join([r"$\chi$","=",dfvalue("Chi"),r"$\pm$",dfvalue("rms_Chi"),r"[eV]$^{-1}$"]) + HPstr = ' '.join([self.parname,"=",dfvalue("HP"),r"$\pm$",dfvalue("rms_HP"),r"[eV]"]) + insettxt = '\n'.join([deginfo,X0str,Xstr,HPstr]) + parambox = AnchoredText(insettxt,loc=insetlocale) + parambox.patch.set_linewidth(2) + parambox.patch.set_edgecolor(insetcolor) + parambox.patch.set_facecolor('white') + ax.add_artist(parambox) return fig + +#=============================================================================================================== +#=============================================================================================================== class LrujAnalyzer: """ Analyzes multiple sets of LRUJ files.