from abc import abstractmethod
import matplotlib.pyplot as plt
import numpy as np
from matplotlib import colormaps
from directupsampling.effqh.phonons import BasePhonons, ExactPhonons, MeshPhonons
from directupsampling.eos import EOS
from directupsampling.fecontrib import FreeEnergyContrib
from directupsampling.fesurface import FreeEnergySurface
from directupsampling.thermo import Thermo
from directupsampling.ti import ThermoInt
from directupsampling.tifit import ThermoIntFit
[docs]
class EOSPlot:
def __init__(self, eos: EOS):
self.eos = eos
[docs]
def plot(self):
if self.eos.raw_data is not None:
es_raw = self.eos.raw_data.energy
vs_raw = self.eos.raw_data.volume
plt.plot(vs_raw, es_raw, "x")
vmin, vmax = np.min(vs_raw), np.max(vs_raw)
else:
veq = self.eos.eq_volume
vmin, vmax = veq * 0.92, veq * 1.12
vs = np.linspace(vmin, vmax, 101)
plt.plot(vs, self.eos(vs))
plt.xlabel("Volume (Ang$^3$)")
plt.ylabel("Energy (eV/atom)")
[docs]
def plot_pressure(self, show=False):
import matplotlib.pyplot as plt
if self.eos.raw_data is not None:
vs_raw = self.eos.raw_data.volume
vmin, vmax = np.min(vs_raw), np.max(vs_raw)
else:
veq = self.eos.eq_volume
vmin, vmax = veq * 0.92, veq * 1.12
vs = np.linspace(vmin, vmax, 101)
plt.plot(vs, self.eos.get_pressure(vs, unit="GPa"))
plt.xlabel("Volume (Ang$^3$)")
plt.ylabel("Pressure (GPa)")
[docs]
class PhononsPlot:
def __init__(self, phonons: BasePhonons):
self.phonons = phonons
[docs]
@abstractmethod
def plot_phonon_dispersion(self): ...
[docs]
@abstractmethod
def plot_frequencies(self): ...
[docs]
class ExactPhononsPlot(PhononsPlot):
[docs]
def plot_phonon_dispersion(self):
raise NotImplementedError
[docs]
def plot_frequencies(self):
vs = np.linspace(*self.phonons.volume_bounds, 101)
freqs = self.phonons.get_frequencies(vs)
plt.plot(vs, freqs.T)
[docs]
class MeshPhononsPlot(PhononsPlot):
[docs]
def plot_phonon_dispersion(self):
ph = self.phonons._phonopy_handler
ph.plot_phonon_dispersion()
[docs]
def plot_frequencies(self):
vs = np.linspace(*self.phonons.volume_bounds, 101)
freqs = self.phonons.get_frequencies(vs)
plt.plot(vs, freqs.T)
[docs]
class FEContribPlot:
def __init__(self, free_energy_contrib: FreeEnergyContrib, max_lines: int = 20):
self.fe_contrib = free_energy_contrib
self.max_lines = max_lines
[docs]
def plot(self):
fig, axs = plt.subplots(
nrows=1,
ncols=2,
figsize=(10, 5),
sharey=True,
constrained_layout=True,
)
plt.sca(axs[0])
self.plot_vs_temperature()
plt.sca(axs[1])
self.plot_vs_volume()
plt.ylabel(None)
[docs]
def plot_vs_volume(self):
"""Plot as function of volume for fixed temperatures."""
fec = self.fe_contrib
if fec.parametrization.raw_data is not None:
self.plot_raw_data_vs_volume() # Sets self.fix_temps
plt.gca().set_prop_cycle(None)
else:
self.fix_temps = np.linspace(*fec.bounds[1], 5)
vmin, vmax = fec.bounds[0]
vmesh = np.linspace(vmin, vmax, 101)
for temp in self.reduce_fix_values(self.fix_temps):
ymesh = fec.get_free_energy(vmesh, temp)
label = f"{temp} {fec.units['temperature']}"
plt.plot(vmesh, ymesh, lw=1, label=label)
plt.xlabel(f"{'Volume'} ({fec.units['volume']})")
self.set_common()
[docs]
def plot_vs_temperature(self):
"""Plot as function of temperature for fixed volumes."""
fec = self.fe_contrib
if fec.parametrization.raw_data is not None:
self.plot_raw_data_vs_temperature() # Sets self.fix_vols
plt.gca().set_prop_cycle(None)
else:
self.fix_vols = np.linspace(*fec.bounds[0], 5)
tmin, tmax = fec.bounds[1]
tmesh = np.linspace(tmin, tmax, 101)
for volume in self.reduce_fix_values(self.fix_vols):
ymesh = fec.get_free_energy(volume, tmesh)
label = f"{volume} {fec.units['volume']}"
plt.plot(tmesh, ymesh, lw=1, label=label)
plt.xlabel(f"{'Temperature'} ({fec.units['temperature']})")
self.set_common()
[docs]
def plot_raw_data_vs_volume(self):
data = self.fe_contrib.parametrization.raw_data
fix_temps = np.unique(data.temperature)
self.fix_temps = fix_temps
for temp in fix_temps:
mask = data.temperature == temp
x = data.volume[mask]
y = data.free_energy[mask]
e = data.error[mask] if len(data.error) > 0 else None
l = f"{temp} {self.fe_contrib.units['temperature']}"
style = dict(capsize=5, marker=".", ms=5, mew=1.5, ls="")
plt.errorbar(x, y, yerr=e, **style, clip_on=False, label=l)
[docs]
def plot_raw_data_vs_temperature(self):
data = self.fe_contrib.parametrization.raw_data
fix_vols = np.unique(data.volume)
self.fix_vols = fix_vols
for vol in fix_vols:
mask = data.volume == vol
x = data.temperature[mask]
y = data.free_energy[mask]
e = data.error[mask] if len(data.error) > 0 else None
l = f"{vol} {self.fe_contrib.units['volume']}"
style = dict(capsize=5, marker=".", ms=5, mew=1.5, ls="")
plt.errorbar(x, y, yerr=e, **style, clip_on=False, label=l)
[docs]
def reduce_fix_values(self, values: list):
if len(values) > self.max_lines:
values = values[:: int(np.round(len(values) / self.max_lines))]
return values
[docs]
def set_common(self):
fec = self.fe_contrib
plt.ylabel(f"F$_{{\\rm {fec.name}}}$ ({fec.units['energy']})")
hl = _combine_duplicate_legend_labels(*plt.gca().get_legend_handles_labels())
plt.legend(*hl, bbox_to_anchor=(0.5, 1.0), loc="lower center", ncol=3)
def _combine_duplicate_legend_labels(handles: list, labels: list):
new_handles, new_labels = [], []
for handle, label in zip(handles, labels):
if label not in new_labels:
new_labels.append(label)
new_handles.append(handle)
else:
ind = np.where(label == np.array(new_labels))[0][0]
new_handles[ind] = (new_handles[ind],) + (handle,)
return new_handles, new_labels
[docs]
class FreeEnergySurfacePlot:
def __init__(self, free_energy_surface: FreeEnergySurface):
self.fes = free_energy_surface
[docs]
def plot_vs_volume(self, temperature, names=None):
if names is None:
names = [_.name for _ in self.fes.contributions]
vmin, vmax = self.fes.bounds[0]
v = np.linspace(vmin, vmax, 101)
y = np.zeros(v.shape)
for c in self.fes.contributions:
if isinstance(c, EOS):
y += c.get_energy(v) - c.eq_energy
for c in self.fes.contributions:
if not isinstance(c, EOS) and c.name in names:
y += c.get_free_energy(v, temperature)
l = f"F$_{{\\rm {c.name}}}$"
plt.plot(v, y, lw=1, label=l)
plt.xlabel("Volume (Ang$^3$)")
plt.ylabel("F (eV)")
plt.legend(bbox_to_anchor=(0.5, 1.0), loc="lower center", ncol=3)
[docs]
class ThermoIntFitPlot:
def __init__(self, tifit: ThermoIntFit):
self.tifit = tifit
[docs]
def plot(self, **kwargs):
self.plot_raw_data(**kwargs)
xs = np.linspace(0, 1, 1001)
ys = self.tifit.best_fit["fit"](xs)
plt.plot(xs, ys, **kwargs)
[docs]
def plot_all(self):
self.plot_raw_data(color="black")
xs = np.linspace(0, 1, 1001)
fit_results = self.tifit.fit_results
for fit_type in fit_results:
ys = fit_results[fit_type]["fit"](xs)
err = fit_results[fit_type]["fit error"]
plt.plot(xs, ys, label=f"{fit_type}, \u00b1{err:.2}")
plt.legend()
[docs]
def plot_raw_data(self, **kwargs):
raw_data = self.tifit.raw_data
err = raw_data.error if len(raw_data.error) > 0 else None
kwargs.pop("label", None)
plt.errorbar(
raw_data.lambda_,
raw_data.energy,
yerr=err,
capsize=3,
ls="none",
marker="o",
markersize=4,
markerfacecolor="none",
**kwargs,
)
def _make_colorbar(cmap, max_t):
sm = plt.cm.ScalarMappable(cmap=cmap, norm=plt.Normalize(vmin=0, vmax=max_t))
cbar = plt.colorbar(
sm,
ax=plt.gca(),
location="right",
shrink=0.7,
)
cbar.set_label("Temperature (K)")
[docs]
class ThermoIntPlot:
def __init__(self, ti: ThermoInt) -> None:
self.ti = ti
[docs]
def plot(self, **kwargs: dict) -> None:
vt_grid = self.ti.vt_grid
ti_fits = self.ti.fits
cmap = colormaps["inferno"]
max_t = np.max([_.temperature for _ in vt_grid]) * 1.1
lambdas_fit = np.linspace(0.0, 1.0, 101)
axarr = plt.gcf().subplots(
1,
len(ThermoIntFit.fit_types),
sharex=True,
sharey=True,
)
for ax, fit_type in zip(axarr, ThermoIntFit.fit_types, strict=True):
for vtp in vt_grid:
try:
ax.plot(
ti_fits[vtp].raw_data.lambda_,
ti_fits[vtp].raw_data.energy,
color=cmap(float(vtp.temperature) / max_t),
ls="none",
marker=".",
mfc="none",
)
ax.plot(
lambdas_fit,
ti_fits[vtp].fit_results[fit_type]["fit"](lambdas_fit),
color=cmap(float(vtp.temperature) / max_t),
label=vtp,
)
except Exception:
continue
ax.set_xlabel(r"$\lambda$")
ax.set_ylabel(r"$\Delta E$")
ax.label_outer()
ax.set_title(fit_type)
ax = axarr.ravel()[-1]
if len(list(vt_grid)) <= 10:
ax.legend(loc="center right", fontsize=8)
else:
_make_colorbar(cmap, max_t)
[docs]
class ThermoPlot:
def __init__(self, thermo: Thermo) -> None:
self.thermo = thermo
[docs]
def plot(self, properties: list[str] | None = None) -> None:
thermo = self.thermo
axs = plt.gcf().subplots(1, 3)
if properties is None:
properties = [
"heat_capacity_isobaric",
"volume_expansion_coefficient",
"bulk_modulus_adiabatic",
]
for ax, prop in zip(axs.flat, properties):
if prop not in thermo.properties.columns:
raise ValueError(f"Property '{prop}' not found.")
name = prop.replace("_", " ").capitalize()
ax.plot(thermo.properties.index, thermo.properties[prop], lw=1, label=name)
ax.set_title(name)
ax.set_xlabel("Temperature (K)")
[docs]
def plot(obj, **kwargs):
"""Plot a directupsampling object.
Returns
-------
matplotlib.pyplot
"""
if isinstance(obj, EOS):
plt.figure()
plotter = EOSPlot(obj)
plotter.plot(**kwargs)
elif isinstance(obj, ExactPhonons):
plotter = ExactPhononsPlot(obj)
plotter.plot_frequencies(**kwargs)
elif isinstance(obj, MeshPhonons):
plotter = MeshPhononsPlot(obj)
plotter.plot_frequencies(**kwargs)
elif isinstance(obj, FreeEnergyContrib):
plotter = FEContribPlot(obj)
plotter.plot(**kwargs)
elif isinstance(obj, FreeEnergySurface):
plt.figure()
plotter = FreeEnergySurfacePlot(obj)
plotter.plot_vs_volume(**kwargs)
elif isinstance(obj, ThermoIntFit):
plt.figure()
plotter = ThermoIntFitPlot(obj)
plotter.plot(**kwargs)
elif isinstance(obj, ThermoInt):
plt.figure()
plotter = ThermoIntPlot(obj)
plotter.plot(**kwargs)
elif isinstance(obj, Thermo):
plt.figure()
plotter = ThermoPlot(obj)
plotter.plot(**kwargs)
else:
raise NotImplementedError
return plt