Skip to content

Commit

Permalink
Implement abimpl.py compare
Browse files Browse the repository at this point in the history
  • Loading branch information
@gmatteo
gmatteo committed Aug 25, 2023
1 parent 40c7c8a commit 357056c
Show file tree
Hide file tree
Showing 2 changed files with 95 additions and 59 deletions.
140 changes: 87 additions & 53 deletions abipy/ml/aseml.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -559,11 +559,11 @@ def plot_forces(self, fontsize=8, **kwargs):
ax.scatter(xs, ys, marker="o")
linear_fit_ax(ax, xs, ys)
ax.grid(True)
f_str = f"$F_{direction}"
f_tex = f"$F_{direction}$"
if icol == 0:
ax.set_ylabel(f"{key2} {f_str}")
ax.set_ylabel(f"{key2} {f_tex}")
if irow == 2:
ax.set_xlabel(f"{key1} {f_str}")
ax.set_xlabel(f"{key1} {f_tex}")
ax.set_title(f"{key1}/{key2} MAE: {stats.MAE:.6f}", fontsize=fontsize)

if "title" not in kwargs: fig.suptitle(f"Cartesian forces in ev/Ang for {self.structure.latex_formula}")
Expand All @@ -587,9 +587,9 @@ def plot_stresses(self, fontsize=6, **kwargs):
ax.scatter(xs, ys, marker="o")
linear_fit_ax(ax, xs, ys)
ax.grid(True)
s_tex = "$" + f"\sigma_{voigt_comp}" + "$"
if icol == 0
ax.set_ylabel(f"{key2} {f_str}")
s_tex = "$\sigma_{%s}$" % voigt_comp
if icol == 0:
ax.set_ylabel(f"{key2} {s_tex}")
if irow == (len(self.ALL_VOIGT_COMPS) - 1):
ax.set_xlabel(f"{key1} {s_tex}")
ax.set_title(f"{key1}/{key2} MAE: {stats.MAE:.6f}", fontsize=fontsize)
Expand All @@ -598,60 +598,71 @@ def plot_stresses(self, fontsize=6, **kwargs):
return fig

@add_fig_kwargs
def plot_traj(self, what_list=("energy", "forces", "stress"),
delta_mode=True, with_stress=True, fontsize=6, **kwargs):
def plot_energies_traj(self, delta_mode=True, fontsize=6, markersize=2, **kwargs):
"""
Plot results along the trajectory.
Plot energies along the trajectory.
Args:
delta_mode: True to plot difference instead of absolute values.
with_stress: True to plot stresses.
"""
what_list = list_strings(what_list)

key_pairs = self.get_key_pairs()
nrows, ncols = 3 if with_stress else 2, len(key_pairs)
nrows, ncols = 1, len(key_pairs)
ax_mat, fig, plt = get_axarray_fig_plt(None, nrows=nrows, ncols=ncols,
sharex=True, sharey="row", squeeze=False)

atom1_cmap = plt.get_cmap("viridis")
atom2_cmap = plt.get_cmap("jet")
markersize = 2

for icol, (key1, key2) in enumerate(key_pairs):

row_count = -1
row_count += 1
ax = ax_mat[row_count, icol]

# Plot delta energy along the trajectory.
e1, e2 = self.xy_energies_for_keys(key1, key2, sort=False)
stats = diff_stats(e1, e2)

ax = ax_mat[0, icol]
if delta_mode:
abs_delta_ene = np.abs(e1 - e2)
ax.plot(abs_delta_ene, marker="o", markersize=markersize)
# Plot delta energy along the trajectory.
ax.plot(np.abs(e1 - e2), marker="o", markersize=markersize)
ax.set_yscale("log")
else:
ax.plot(e1, marker="o", color="red", legend=key1, markersize=markersize)
ax.plot(e2, marker="o", color="blue", legend=key2, markersize=markersize)

set_grid_legend(ax, fontsize, # xlabel='trajectory',
ylabel="$|\Delta_E|$ (eV)" if delta_mode else ylabel="$E$ (eV)",
grid=True, legend=not delta_mode,
legend_loc="left",
title=f"{key1}/{key2} MAE: {stats.MAE:.6f}")
set_grid_legend(ax, fontsize, xlabel='trajectory',
ylabel="$|\Delta_E|$" if delta_mode else "$E$",
grid=True, legend=not delta_mode, legend_loc="left",
title=f"{key1}/{key2} MAE: {stats.MAE:.6f} eV")

head = "$\Delta$-Energy in eV" if delta_mode else "Energy in eV"
if "title" not in kwargs: fig.suptitle(f"{head} for {self.structure.latex_formula}")

return fig

@add_fig_kwargs
def plot_forces_traj(self, delta_mode=True, fontsize=6, markersize=2, **kwargs):
"""
Plot forces along the trajectory.
Args:
delta_mode: True to plot difference instead of absolute values.
"""
# Fx,Fy,Fx along rows, pairs along columns.
key_pairs = self.get_key_pairs()
nrows, ncols = 3, len(key_pairs)
ax_mat, fig, plt = get_axarray_fig_plt(None, nrows=nrows, ncols=ncols,
sharex=True, sharey="row", squeeze=False)

atom1_cmap = plt.get_cmap("viridis")
atom2_cmap = plt.get_cmap("jet")
marker_idir = {0: ">", 1: "<", 2: "^"}

# Plot (delta of) forces along the trajectory.
for icol, (key1, key2) in enumerate(key_pairs):
# Arrays of shape: [nsteps, natom, 3]
f1_tad, f2_tad = self.traj_forces_for_keys(key1, key2)

row_count += 1
ax = ax_mat[row_count, icol]
marker_idir = {0: ">", 1: "<", 2: "^"}
for iatom in range(self.natom):
for idir, direction in enumerate(("x", "y", "z")):
for idir, direction in enumerate(("x", "y", "z")):
last_row = idir == 2
xs, ys = self.xy_forces_for_keys(key1, key2, direction)
stats = diff_stats(xs, ys)
ax = ax_mat[idir, icol]
ax.set_title(f"{key1}/{key2} MAE: {stats.MAE:.6f}", fontsize=fontsize)
for iatom in range(self.natom):
if delta_mode:
# Plot delta of forces along the trajectory.
style = dict(marker=marker_idir[idir], markersize=markersize,
color=atom1_cmap(float(iatom) / self.natom))
abs_delta_force = np.abs(f1_tad[:,iatom,idir] - f2_tad[:,iatom,idir])
Expand All @@ -666,34 +677,57 @@ def plot_traj(self, what_list=("energy", "forces", "stress"),
label=f"${key1}\, F_{direction}$" if (iatom, idir) == (0,0) else None)
ax.plot(f2_tad[:,iatom,idir], **f2_style,
label=f"${key2}\, F_{direction}$" if (iatom, idir) == (0,0) else None)
ax.set_ylim(-1, +1)
#ax.set_ylim(-1, +1)

set_grid_legend(ax, fontsize, xlabel='trajectory' if last_row else None,
grid=True, legend=not delta_mode, legend_loc="upper left",
ylabel=f"$|\Delta F_{direction}|$" if delta_mode else f"$F_{direction}$")

set_grid_legend(ax, fontsize, # xlabel='trajectory',
grid=True, legend=True, legend_loc="upper left",
ylabel="$|\Delta_F|$ (eV/Ang)" if delta_mode else "$F$ (eV/Ang)")
head = "$\Delta$-forces in eV/Ang" if delta_mode else "Forces in eV/Ang"
if "title" not in kwargs: fig.suptitle(f"{head} for {self.structure.latex_formula}")

return fig

@add_fig_kwargs
def plot_stress_traj(self, delta_mode=True, markersize=2, fontsize=6, **kwargs):
"""
Plot results along the trajectory.
Args:
delta_mode: True to plot difference instead of absolute values.
"""
# Sxx,Syy,Szz,... along rows, pairs along columns.
key_pairs = self.get_key_pairs()
nrows, ncols = 6, len(key_pairs)
ax_mat, fig, plt = get_axarray_fig_plt(None, nrows=nrows, ncols=ncols,
sharex=True, sharey="row", squeeze=False)

marker_voigt = {"xx": ">", "yy": "<", "zz": "^", "yz": 1, "xz": 2, "xy":3}
for icol, (key1, key2) in enumerate(key_pairs):
# Plot (delta of) stresses along the trajectory.
row_count += 1
ax = ax_mat[row_count, icol]
marker_voigt = {"xx": ">", "yy": "<", "zz": "^", "yz": 1, "xz": 2, "xy":3}
for iv, voigt_comp in enumerate(self.ALL_VOIGT_COMPS):
voigt_comp_tex = "{" + f"{voigt_comp}" + "}"
xs, ys = self.xy_stress_for_keys(key1, key2, voigt_comp)
stats = diff_stats(xs, ys)
last_row = iv == len(self.ALL_VOIGT_COMPS) - 1
ax = ax_mat[iv, icol]
ax.set_title(f"{key1}/{key2} MAE: {stats.MAE:.6f}", fontsize=fontsize)
voigt_comp_tex = "{" + voigt_comp + "}"
s1, s2 = self.xy_stress_for_keys(key1, key2, voigt_comp, sort=False)
s_style = dict(marker=marker_voigt[voigt_comp], markersize=markersize)
if delta_mode:
abs_delta_stress = np.abs(s1 - s2)
ax.plot(abs_delta_stress, **s_style, label=f"$\Delta \sigma_{voigt_comp_tex}$")
ax.plot(abs_delta_stress, **s_style, label=f"$|\Delta \sigma_{voigt_comp_tex}|$")
ax.set_yscale("log")
else:
ax.plot(s1, **s_style, label=f"${key1}\, \sigma_{voigt_comp_tex}$" if iv == 0 else None)
ax.plot(s2, **s_style, label=f"${key2}\, \sigma_{voigt_comp_tex}$" if iv == 0 else None)
ax.set_ylim(-1, +1)
ax.plot(s1, **s_style, label=f"${key1}\,\sigma_{voigt_comp_tex}$" if iv == 0 else None)
ax.plot(s2, **s_style, label=f"${key2}\,\sigma_{voigt_comp_tex}$" if iv == 0 else None)
#ax.set_ylim(-1, +1)

set_grid_legend(ax, fontsize, xlabel='trajectory',
grid=True, legend=True, legend_loc="upper left",
ylabel="$|\Delta_\sigma|$ (eV/Ang$^2$)" if delta_mode else "$\sigma$ (eV/Ang$^2$)")
set_grid_legend(ax, fontsize, xlabel='trajectory' if last_row else None,
grid=True, legend=not delta_mode, legend_loc="upper left",
ylabel=f"$|\Delta \sigma_{voigt_comp_tex}|$ " if delta_mode else "$\sigma$ ")

head = "$\Delta$-terms" if delta_mode else "GS terms"
head = r"$\Delta \sigma$ (eV/Ang$^2$)" if delta_mode else "Stress tensor (eV/Ang$^2$)"
if "title" not in kwargs: fig.suptitle(f"{head} for {self.structure.latex_formula}")

return fig
Expand Down
14 changes: 8 additions & 6 deletions abipy/scripts/abiml.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -123,8 +123,8 @@ def main(ctx, nn_name, seaborn):
"""Script to perform calculations with ML potentials."""
ctx.ensure_object(dict)

#if seaborn:
if True:
if seaborn:
#if True:
# Activate seaborn settings for plots
import seaborn as sns
sns.set(context=seaborn, style='darkgrid', palette='deep',
Expand Down Expand Up @@ -496,8 +496,9 @@ def scan_relax(ctx, filepath,
@click.option("-nns", '--nn-names', type=str, multiple=True, show_default=True, help='ML potentials to be used',
#default=["m3gnet", "chgnet"],
default=["chgnet"])
@click.option("-e", '--exposer', show_default=True, help='Plotting backend: mpl for matplotlib, panel for web-based',
type=click.Choice(["mpl", "panel"]))
@click.option("-e", '--exposer', default="mpl", show_default=True, type=click.Choice(["mpl", "panel"]),
help='Plotting backend: mpl for matplotlib, panel for web-based')

@add_nprocs_opt
@add_workdir_verbose_opts
def compare(ctx, filepath,
Expand All @@ -523,8 +524,9 @@ def compare(ctx, filepath,
from abipy.tools.plotting import Exposer
with Exposer.as_exposer(exposer, title=os.path.basename(filepath)) as e:
with_stress = True
e(c.plot_traj(delta_mode=False, with_stress=with_stress, show=False))
e(c.plot_traj(delta_mode=True, with_stress=with_stress, show=False))
e(c.plot_energies_traj(delta_mode=True, show=False))
e(c.plot_forces_traj(delta_mode=True, show=False))
e(c.plot_stress_traj(delta_mode=True, show=False))
e(c.plot_energies(show=False))
e(c.plot_forces(show=False))
if with_stress:
Expand Down

0 comments on commit 357056c

Please sign in to comment.