"""Plotting helpers for UFP example and analysis workflows."""
from __future__ import annotations
from pathlib import Path
from typing import Any, Sequence
import numpy as np
[docs]
def plot_pair_histogram(
histogram: np.ndarray,
bin_edges: np.ndarray,
*,
cutoff: float,
title: str = "Pair-distance distribution",
) -> tuple[Any, Any]:
"""Plot a pair-distance histogram."""
import matplotlib.pyplot as plt
centers = 0.5 * (bin_edges[1:] + bin_edges[:-1])
fig, ax = plt.subplots(figsize=(5.0, 3.0))
ax.plot(centers, histogram, color="black", linewidth=1.5)
ax.axvline(cutoff, color="tab:red", linestyle="--", linewidth=1.0)
ax.set_xlabel("Pair distance (A)")
ax.set_ylabel("Probability density")
ax.set_title(title)
fig.tight_layout()
return fig, ax
[docs]
def plot_pair_curve(
distances: np.ndarray,
energies: np.ndarray,
*,
title: str,
xlabel: str = "Pair distance (A)",
) -> tuple[Any, Any]:
"""Plot a pair potential curve."""
import matplotlib.pyplot as plt
fig, ax = plt.subplots(figsize=(5.0, 3.0))
ax.plot(distances, energies, color="tab:blue", linewidth=1.5)
ax.axhline(0.0, color="black", linewidth=0.8)
ax.set_xlabel(xlabel)
ax.set_ylabel("Interaction energy (eV)")
ax.set_title(title)
fig.tight_layout()
return fig, ax
[docs]
def density_scatter(
x: np.ndarray,
y: np.ndarray,
*,
xlabel: str,
ylabel: str,
title: str,
) -> tuple[Any, Any]:
"""Draw a compact density-style parity plot."""
import matplotlib.pyplot as plt
fig, ax = plt.subplots(figsize=(3.6, 3.6))
ax.hexbin(x, y, gridsize=55, mincnt=1, cmap="viridis")
lower = float(min(np.min(x), np.min(y)))
upper = float(max(np.max(x), np.max(y)))
ax.plot([lower, upper], [lower, upper], color="white", linewidth=1.0)
ax.set_xlabel(xlabel)
ax.set_ylabel(ylabel)
ax.set_title(title)
fig.tight_layout()
return fig, ax
def _prediction_energy_arrays(data) -> tuple[np.ndarray, np.ndarray, str]:
"""Return per-atom energy arrays from prediction ``.npz`` contents."""
if "system_sizes" in data.files:
sizes = np.asarray(data["system_sizes"], dtype=float)
elif "sizes" in data.files:
sizes = np.asarray(data["sizes"], dtype=float)
else:
sizes = None
if sizes is not None:
return (
np.asarray(data["true_energies"], dtype=float) / sizes,
np.asarray(data["predicted_energies"], dtype=float) / sizes,
"energy per atom (eV/atom)",
)
has_per_atom_energy = (
"true_energy_per_atom" in data.files
and "predicted_energy_per_atom" in data.files
)
if has_per_atom_energy:
return (
np.asarray(data["true_energy_per_atom"], dtype=float),
np.asarray(data["predicted_energy_per_atom"], dtype=float),
"energy per atom (eV/atom)",
)
return (
np.asarray(data["true_energies"], dtype=float),
np.asarray(data["predicted_energies"], dtype=float),
"energy (eV)",
)
[docs]
def plot_prediction_density_arrays(
true_energies: np.ndarray,
predicted_energies: np.ndarray,
true_force_components: np.ndarray,
predicted_force_components: np.ndarray,
*,
title_prefix: str,
) -> tuple[Any, Any]:
"""Plot energy and force density-parity panels from prediction arrays."""
import matplotlib.pyplot as plt
fig, axes = plt.subplots(1, 2, figsize=(7.2, 3.4))
axes[0].hexbin(true_energies, predicted_energies, gridsize=55, mincnt=1)
axes[0].set_xlabel("True energy per atom (eV/atom)")
axes[0].set_ylabel("Predicted energy per atom (eV/atom)")
axes[0].set_title(f"{title_prefix} energies")
axes[1].hexbin(
true_force_components,
predicted_force_components,
gridsize=55,
mincnt=1,
)
axes[1].set_xlabel("True force component (eV/A)")
axes[1].set_ylabel("Predicted force component (eV/A)")
axes[1].set_title(f"{title_prefix} forces")
for ax in axes:
x_min, x_max = ax.get_xlim()
y_min, y_max = ax.get_ylim()
lower = min(x_min, y_min)
upper = max(x_max, y_max)
ax.plot([lower, upper], [lower, upper], color="white", linewidth=1.0)
fig.tight_layout()
return fig, axes
[docs]
def plot_prediction_density_file(
npz_path: Path,
*,
title_prefix: str,
) -> tuple[Any, Any]:
"""Plot energy and force density-parity panels from a prediction ``.npz``."""
data = np.load(npz_path)
true_energy, predicted_energy, _ = _prediction_energy_arrays(data)
return plot_prediction_density_arrays(
true_energy,
predicted_energy,
data["true_force_components"],
data["predicted_force_components"],
title_prefix=title_prefix,
)
[docs]
def save_prediction_density_plots(
npz_path: Path,
output_directory: Path,
*,
prefix: str,
split: str,
title_prefix: str,
dpi: int = 200,
) -> tuple[Path, Path]:
"""Write separate energy and force density plots from one prediction file."""
import matplotlib.pyplot as plt
output_directory.mkdir(parents=True, exist_ok=True)
data = np.load(npz_path)
true_energy, predicted_energy, energy_label = _prediction_energy_arrays(data)
fig, _ = density_scatter(
true_energy,
predicted_energy,
xlabel=f"True {energy_label}",
ylabel=f"Predicted {energy_label}",
title=f"{title_prefix} energies",
)
energy_path = output_directory / f"{prefix}_{split}_energy_density.png"
fig.savefig(energy_path, dpi=dpi)
plt.close(fig)
fig, _ = density_scatter(
data["true_force_components"],
data["predicted_force_components"],
xlabel="True force component (eV/A)",
ylabel="Predicted force component (eV/A)",
title=f"{title_prefix} forces",
)
force_path = output_directory / f"{prefix}_{split}_force_density.png"
fig.savefig(force_path, dpi=dpi)
plt.close(fig)
return energy_path, force_path
[docs]
def save_uncertainty_calibration_plots(
metrics: dict[str, object],
output_directory: Path,
*,
prefix: str,
dpi: int = 200,
) -> tuple[Path, Path, Path]:
"""Write standard energy uncertainty calibration diagnostic plots."""
import matplotlib.pyplot as plt
output_directory.mkdir(parents=True, exist_ok=True)
residual = np.asarray(metrics["energy_residual_per_atom"], dtype=float)
predicted_std = np.asarray(metrics["predicted_energy_std_per_atom"], dtype=float)
normalized = np.asarray(metrics["normalized_energy_residual"], dtype=float)
nominal = np.asarray(metrics["nominal_coverage"], dtype=float)
empirical = np.asarray(metrics["empirical_coverage"], dtype=float)
fig, ax = plt.subplots(figsize=(4.2, 3.6))
ax.scatter(1000.0 * predicted_std, 1000.0 * np.abs(residual), s=16, alpha=0.75)
upper = float(
max(
np.max(1000.0 * predicted_std),
np.max(1000.0 * np.abs(residual)),
1.0,
)
)
ax.plot([0.0, upper], [0.0, upper], color="black", linewidth=0.9)
ax.set_xlabel("Predicted energy std (meV/atom)")
ax.set_ylabel("Absolute energy residual (meV/atom)")
ax.set_title("Residual vs uncertainty")
fig.tight_layout()
residual_path = output_directory / f"{prefix}_residual_vs_std.png"
fig.savefig(residual_path, dpi=dpi)
plt.close(fig)
fig, ax = plt.subplots(figsize=(4.2, 3.6))
ax.hist(normalized, bins=30, color="tab:blue", alpha=0.75, density=True)
ax.axvline(0.0, color="black", linewidth=0.9)
ax.set_xlabel("Normalized energy residual")
ax.set_ylabel("Density")
ax.set_title("Normalized residuals")
fig.tight_layout()
histogram_path = output_directory / f"{prefix}_normalized_residuals.png"
fig.savefig(histogram_path, dpi=dpi)
plt.close(fig)
fig, ax = plt.subplots(figsize=(4.2, 3.6))
ax.plot(nominal, empirical, marker="o", linewidth=1.5)
ax.plot([0.0, 1.0], [0.0, 1.0], color="black", linewidth=0.9)
ax.set_xlim(0.0, 1.0)
ax.set_ylim(0.0, 1.0)
ax.set_xlabel("Nominal coverage")
ax.set_ylabel("Empirical coverage")
ax.set_title("Coverage calibration")
fig.tight_layout()
coverage_path = output_directory / f"{prefix}_coverage.png"
fig.savefig(coverage_path, dpi=dpi)
plt.close(fig)
return residual_path, histogram_path, coverage_path
[docs]
def save_force_uncertainty_calibration_plots(
metrics: dict[str, object],
output_directory: Path,
*,
prefix: str,
dpi: int = 200,
) -> tuple[Path, Path, Path]:
"""Write standard force uncertainty calibration diagnostic plots."""
import matplotlib.pyplot as plt
output_directory.mkdir(parents=True, exist_ok=True)
residual = np.asarray(metrics["force_residual"], dtype=float)
predicted_std = np.asarray(metrics["predicted_force_std"], dtype=float)
normalized = np.asarray(metrics["normalized_force_residual"], dtype=float)
nominal = np.asarray(metrics["nominal_coverage"], dtype=float)
empirical = np.asarray(metrics["empirical_coverage"], dtype=float)
fig, ax = plt.subplots(figsize=(4.2, 3.6))
ax.scatter(1000.0 * predicted_std, 1000.0 * np.abs(residual), s=8, alpha=0.5)
upper = float(
max(
np.max(1000.0 * predicted_std),
np.max(1000.0 * np.abs(residual)),
1.0,
)
)
ax.plot([0.0, upper], [0.0, upper], color="black", linewidth=0.9)
ax.set_xlabel("Predicted force std (meV/A)")
ax.set_ylabel("Absolute force residual (meV/A)")
ax.set_title("Force residual vs uncertainty")
fig.tight_layout()
residual_path = output_directory / f"{prefix}_force_residual_vs_std.png"
fig.savefig(residual_path, dpi=dpi)
plt.close(fig)
fig, ax = plt.subplots(figsize=(4.2, 3.6))
ax.hist(normalized, bins=40, color="tab:orange", alpha=0.75, density=True)
ax.axvline(0.0, color="black", linewidth=0.9)
ax.set_xlabel("Normalized force residual")
ax.set_ylabel("Density")
ax.set_title("Force normalized residuals")
fig.tight_layout()
histogram_path = output_directory / f"{prefix}_force_normalized_residuals.png"
fig.savefig(histogram_path, dpi=dpi)
plt.close(fig)
fig, ax = plt.subplots(figsize=(4.2, 3.6))
ax.plot(nominal, empirical, marker="o", linewidth=1.5)
ax.plot([0.0, 1.0], [0.0, 1.0], color="black", linewidth=0.9)
ax.set_xlim(0.0, 1.0)
ax.set_ylim(0.0, 1.0)
ax.set_xlabel("Nominal coverage")
ax.set_ylabel("Empirical coverage")
ax.set_title("Force coverage calibration")
fig.tight_layout()
coverage_path = output_directory / f"{prefix}_force_coverage.png"
fig.savefig(coverage_path, dpi=dpi)
plt.close(fig)
return residual_path, histogram_path, coverage_path
[docs]
def plot_onebody_values(
values_by_label: dict[str, float],
*,
title: str = "One-body reference energies",
) -> tuple[Any, Any]:
"""Plot fitted one-body reference values as a small bar chart."""
import matplotlib.pyplot as plt
labels = tuple(values_by_label)
values = np.asarray([values_by_label[label] for label in labels], dtype=float)
fig, ax = plt.subplots(figsize=(max(3.2, 0.6 * len(labels)), 3.0))
ax.bar(labels, values, color="tab:gray")
ax.axhline(0.0, color="black", linewidth=0.8)
ax.set_ylabel("One-body energy (eV)")
ax.set_title(title)
fig.tight_layout()
return fig, ax
[docs]
def plot_pair_components(
interaction_model: Any,
channels: Sequence[tuple[tuple[str, str], float, float]],
*,
title: str,
n_points: int = 1000,
) -> tuple[Any, Any]:
"""Plot two-body dimer curves for a sequence of element-pair channels."""
import matplotlib.pyplot as plt
from ufp.analysis.outputs import dimer_curve
fig, ax = plt.subplots(figsize=(5.6, 3.6))
for symbols, r_min, r_max in channels:
distances, energies = dimer_curve(
interaction_model,
symbols,
r_min=r_min,
r_max=r_max,
n_points=n_points,
)
ax.plot(distances, energies, linewidth=1.5, label="-".join(symbols))
ax.axhline(0.0, color="black", linewidth=0.8)
ax.set_xlabel("Pair distance (A)")
ax.set_ylabel("Two-body contribution (eV)")
ax.set_title(title)
ax.legend(fontsize=8)
fig.tight_layout()
return fig, ax
[docs]
def plot_threebody_surface(
x: np.ndarray,
y: np.ndarray,
values: np.ndarray,
*,
title: str,
xlabel: str = "r_ij (A)",
ylabel: str = "r_ik (A)",
colorbar_label: str = "Three-body energy (eV)",
) -> tuple[Any, Any]:
"""Plot one precomputed three-body diagnostic surface."""
import matplotlib.pyplot as plt
fig, ax = plt.subplots(figsize=(4.2, 3.6))
image = ax.contourf(x, y, values, levels=21, cmap="coolwarm")
ax.set_xlabel(xlabel)
ax.set_ylabel(ylabel)
ax.set_title(title)
fig.colorbar(image, ax=ax, label=colorbar_label)
fig.tight_layout()
return fig, ax
__all__ = [
"density_scatter",
"plot_onebody_values",
"plot_pair_components",
"plot_pair_curve",
"plot_pair_histogram",
"plot_prediction_density_arrays",
"plot_prediction_density_file",
"plot_threebody_surface",
"save_prediction_density_plots",
"save_force_uncertainty_calibration_plots",
"save_uncertainty_calibration_plots",
]