From e1ce8c0d1882dc76f354d8606e98b2cff79b2820 Mon Sep 17 00:00:00 2001 From: Fabian Zills <46721498+PythonFZ@users.noreply.github.com> Date: Thu, 10 Jul 2025 19:51:27 +0200 Subject: [PATCH 1/4] Create model_dataset_analysis.py --- ipsuite/analysis/model_dataset_analysis.py | 313 +++++++++++++++++++++ 1 file changed, 313 insertions(+) create mode 100644 ipsuite/analysis/model_dataset_analysis.py diff --git a/ipsuite/analysis/model_dataset_analysis.py b/ipsuite/analysis/model_dataset_analysis.py new file mode 100644 index 00000000..f3bed6da --- /dev/null +++ b/ipsuite/analysis/model_dataset_analysis.py @@ -0,0 +1,313 @@ +from pathlib import Path +import json +import os + +import ase +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import seaborn as sns +import zntrack +from apax.nodes import Apax +from ipsuite.utils.metrics import get_full_metrics + + +# TODO: uncertainty for metrics? +class ModelEvaluationAnalysis(zntrack.Node): + data: list[ase.Atoms] = zntrack.deps() + models: list[Apax] = zntrack.deps() + + figure: Path = zntrack.outs_path(zntrack.nwd / "figure.png") + metrics: list = zntrack.metrics() + + def run(self): + self.metrics = [] + + energies = np.array([x.get_potential_energy() for x in self.data]) + forces = np.array([x.get_forces() for x in self.data]) + + for model in self.models: + calc = model.get_calculator() + if hasattr(calc, "batch_eval"): + frames = calc.batch_eval(self.data, 1) + else: + frames = [] + for atoms in self.data: + atoms.calc = calc + atoms.get_potential_energy() + frames.append(atoms) + self.metrics.append( + { + "energy": get_full_metrics( + energies, np.array([x.get_potential_energy() for x in frames]) + ), + "forces": get_full_metrics( + forces, np.array([x.get_forces() for x in frames]) + ), + } + ) + + fig, ax = plt.subplots(2, 1, figsize=(8, 6), sharex=True) + fig.suptitle("Model evolution evaluation") + ax[0].plot([x["energy"]["mae"] for x in self.metrics], marker="o") + ax[0].set_ylabel("MAE Energy") + ax[1].plot([x["forces"]["mae"] for x in self.metrics], marker="o") + ax[1].set_xlabel("Model iteration") + ax[1].set_ylabel("MAE Forces") + fig.savefig(self.figure, bbox_inches="tight") + + +class ClusterdDataEvaluation(zntrack.Node): + data: list[list[ase.Atoms]] = zntrack.deps() + labels: list[str] = zntrack.params() + model: Apax = zntrack.deps() + + outs_path: Path = zntrack.outs_path(zntrack.nwd / "outs") + + def run(self): + os.makedirs(self.outs_path, exist_ok=True) + + calc = self.model.get_calculator() + + all_metrics = [] + all_energies_true = [] + all_energies_pred = [] + all_forces_true = [] + all_forces_pred = [] + + for i, (data_cluster, label) in enumerate(zip(self.data, self.labels)): + # Get true values + energies_true = np.array([x.get_potential_energy() for x in data_cluster]) + forces_true = np.array([x.get_forces() for x in data_cluster]) + + # Run predictions + if hasattr(calc, "batch_eval"): + frames = calc.batch_eval(data_cluster, 1) + else: + frames = [] + for atoms in data_cluster: + atoms.calc = calc + atoms.get_potential_energy() + frames.append(atoms) + + # Get predicted values + energies_pred = np.array([x.get_potential_energy() for x in frames]) + forces_pred = np.array([x.get_forces() for x in frames]) + + # Calculate metrics + energy_metrics = get_full_metrics(energies_true, energies_pred) + forces_metrics = get_full_metrics(forces_true, forces_pred) + + cluster_metrics = { + "label": label, + "energy": energy_metrics, + "forces": forces_metrics, + "n_structures": len(data_cluster) + } + all_metrics.append(cluster_metrics) + + # Store for plotting + all_energies_true.extend(energies_true) + all_energies_pred.extend(energies_pred) + all_forces_true.extend(forces_true.flatten()) + all_forces_pred.extend(forces_pred.flatten()) + + + # Save all metrics + with open(self.outs_path / "metrics.json", "w") as f: + json.dump(all_metrics, f, indent=2) + + # Create visualization + fig, axes = plt.subplots(2, 2, figsize=(12, 10)) + fig.suptitle("Model Performance by Data Cluster", fontsize=16) + + # Create color map + colors = plt.cm.Set1(np.linspace(0, 1, len(self.labels))) + + # Energy relative error plot + ax1 = axes[0, 0] + start_idx = 0 + for i, (data_cluster, label) in enumerate(zip(self.data, self.labels)): + end_idx = start_idx + len(data_cluster) + true_energies = np.array(all_energies_true[start_idx:end_idx]) + pred_energies = np.array(all_energies_pred[start_idx:end_idx]) + + # Calculate per-atom energies + true_energies_per_atom = true_energies / np.array([len(atoms) for atoms in data_cluster]) + pred_energies_per_atom = pred_energies / np.array([len(atoms) for atoms in data_cluster]) + + # Calculate difference: predicted - true + diff = pred_energies_per_atom - true_energies_per_atom + + ax1.scatter( + true_energies_per_atom, + diff, + color=colors[i], label=label, alpha=0.6 + ) + start_idx = end_idx + + ax1.axhline(y=0, color='k', linestyle='--', alpha=0.5) + ax1.set_xlabel("True Energy per Atom / eV") + ax1.set_ylabel("Difference: Pred - True / eV") + ax1.set_title("Energy Difference Plot") + ax1.legend() + + # Forces relative error plot + ax2 = axes[0, 1] + start_idx = 0 + for i, (data_cluster, label) in enumerate(zip(self.data, self.labels)): + n_forces = len(data_cluster) * 3 * len(data_cluster[0]) + end_idx = start_idx + n_forces + true_forces = np.array(all_forces_true[start_idx:end_idx]) + pred_forces = np.array(all_forces_pred[start_idx:end_idx]) + + # Calculate difference: predicted - true + diff = pred_forces - true_forces + + ax2.scatter( + true_forces, + diff, + color=colors[i], label=label, alpha=0.6, s=1 + ) + start_idx = end_idx + + ax2.axhline(y=0, color='k', linestyle='--', alpha=0.5) + ax2.set_xlabel("True Forces / eV/Å") + ax2.set_ylabel("Difference: Pred - True / eV/Å") + ax2.set_title("Forces Difference Plot") + + # Energy MAE by cluster + ax3 = axes[1, 0] + energy_maes = [m["energy"]["mae"] for m in all_metrics] + ax3.bar(range(len(self.labels)), energy_maes, color=colors) + ax3.set_xlabel("Data Cluster") + ax3.set_ylabel("Energy MAE / eV") + ax3.set_title("Energy MAE by Cluster") + ax3.set_xticks(range(len(self.labels))) + ax3.set_xticklabels(self.labels, rotation=45) + + # Forces MAE by cluster + ax4 = axes[1, 1] + forces_maes = [m["forces"]["mae"] for m in all_metrics] + ax4.bar(range(len(self.labels)), forces_maes, color=colors) + ax4.set_xlabel("Data Cluster") + ax4.set_ylabel("Forces MAE / eV/Å") + ax4.set_title("Forces MAE by Cluster") + ax4.set_xticks(range(len(self.labels))) + ax4.set_xticklabels(self.labels, rotation=45) + + plt.tight_layout() + plt.savefig(self.outs_path / "evaluation_plots.png", dpi=300, bbox_inches="tight") + plt.close() + + # Create jointplots for error distribution by label + + # Combined energy error distribution with all labels + all_true_energies_per_atom = [] + all_energy_errors_per_atom = [] + all_energy_labels = [] + + start_idx = 0 + for i, (data_cluster, label) in enumerate(zip(self.data, self.labels)): + end_idx = start_idx + len(data_cluster) + + true_energies = np.array(all_energies_true[start_idx:end_idx]) + pred_energies = np.array(all_energies_pred[start_idx:end_idx]) + + true_energies_per_atom = true_energies / np.array([len(atoms) for atoms in data_cluster]) + pred_energies_per_atom = pred_energies / np.array([len(atoms) for atoms in data_cluster]) + energy_errors_per_atom = pred_energies_per_atom - true_energies_per_atom + + all_true_energies_per_atom.extend(true_energies_per_atom) + all_energy_errors_per_atom.extend(energy_errors_per_atom) + all_energy_labels.extend([label] * len(true_energies_per_atom)) + + start_idx = end_idx + + # Create DataFrame for combined energy plot + energy_df = pd.DataFrame({ + 'true_energy': all_true_energies_per_atom, + 'energy_error': all_energy_errors_per_atom, + 'label': all_energy_labels + }) + + # Create combined energy jointplot with hue + g = sns.jointplot( + data=energy_df, + x='true_energy', + y='energy_error', + hue='label', + height=7 + ) + + g.ax_joint.axhline(y=0, color='red', linestyle='--', alpha=0.7) + g.set_axis_labels('True Energy per Atom / eV', 'Energy Error: Pred - True / eV') + g.figure.suptitle('Combined Energy Error Distribution - All Labels', y=1.02) + + plt.tight_layout() + plt.savefig(self.outs_path / "energy_error_distribution_combined.png", dpi=300, bbox_inches="tight") + plt.close() + + # Combined forces error distribution with all labels + all_true_forces = [] + all_forces_errors = [] + all_forces_labels = [] + + start_idx = 0 + for i, (data_cluster, label) in enumerate(zip(self.data, self.labels)): + n_forces = len(data_cluster) * 3 * len(data_cluster[0]) + end_idx = start_idx + n_forces + + true_forces = np.array(all_forces_true[start_idx:end_idx]) + pred_forces = np.array(all_forces_pred[start_idx:end_idx]) + forces_errors = pred_forces - true_forces + + # Sample for better visualization + n_sample = min(5000, len(true_forces)) # Smaller sample for combined plot + if len(true_forces) > n_sample: + idx = np.random.choice(len(true_forces), n_sample, replace=False) + true_forces_sample = true_forces[idx] + forces_errors_sample = forces_errors[idx] + else: + true_forces_sample = true_forces + forces_errors_sample = forces_errors + + all_true_forces.extend(true_forces_sample) + all_forces_errors.extend(forces_errors_sample) + all_forces_labels.extend([label] * len(true_forces_sample)) + + start_idx = end_idx + + # Create DataFrame for combined forces plot + forces_df = pd.DataFrame({ + 'true_forces': all_true_forces, + 'forces_error': all_forces_errors, + 'label': all_forces_labels + }) + + # Create combined forces jointplot with hue + g = sns.jointplot( + data=forces_df, + x='true_forces', + y='forces_error', + hue='label', + height=7 + ) + + g.ax_joint.axhline(y=0, color='red', linestyle='--', alpha=0.7) + g.set_axis_labels('True Forces / eV/Å', 'Force Error: Pred - True / eV/Å') + g.figure.suptitle('Combined Forces Error Distribution - All Labels', y=1.02) + + plt.tight_layout() + plt.savefig(self.outs_path / "forces_error_distribution_combined.png", dpi=300, bbox_inches="tight") + plt.close() + + # Save summary metrics + summary_metrics = { + "overall_energy_mae": np.mean(energy_maes), + "overall_forces_mae": np.mean(forces_maes), + "cluster_metrics": all_metrics + } + + with open(self.outs_path / "summary_metrics.json", "w") as f: + json.dump(summary_metrics, f, indent=2) From 26ee727a5aa58f4162d141682422a9dd4da99aef Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 10 Jul 2025 17:53:12 +0000 Subject: [PATCH 2/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- ipsuite/analysis/model_dataset_analysis.py | 190 +++++++++++---------- 1 file changed, 98 insertions(+), 92 deletions(-) diff --git a/ipsuite/analysis/model_dataset_analysis.py b/ipsuite/analysis/model_dataset_analysis.py index f3bed6da..082fadbc 100644 --- a/ipsuite/analysis/model_dataset_analysis.py +++ b/ipsuite/analysis/model_dataset_analysis.py @@ -1,6 +1,6 @@ -from pathlib import Path import json import os +from pathlib import Path import ase import matplotlib.pyplot as plt @@ -9,6 +9,7 @@ import seaborn as sns import zntrack from apax.nodes import Apax + from ipsuite.utils.metrics import get_full_metrics @@ -66,20 +67,20 @@ class ClusterdDataEvaluation(zntrack.Node): def run(self): os.makedirs(self.outs_path, exist_ok=True) - + calc = self.model.get_calculator() - + all_metrics = [] all_energies_true = [] all_energies_pred = [] all_forces_true = [] all_forces_pred = [] - + for i, (data_cluster, label) in enumerate(zip(self.data, self.labels)): # Get true values energies_true = np.array([x.get_potential_energy() for x in data_cluster]) forces_true = np.array([x.get_forces() for x in data_cluster]) - + # Run predictions if hasattr(calc, "batch_eval"): frames = calc.batch_eval(data_cluster, 1) @@ -89,41 +90,40 @@ def run(self): atoms.calc = calc atoms.get_potential_energy() frames.append(atoms) - + # Get predicted values energies_pred = np.array([x.get_potential_energy() for x in frames]) forces_pred = np.array([x.get_forces() for x in frames]) - + # Calculate metrics energy_metrics = get_full_metrics(energies_true, energies_pred) forces_metrics = get_full_metrics(forces_true, forces_pred) - + cluster_metrics = { "label": label, "energy": energy_metrics, "forces": forces_metrics, - "n_structures": len(data_cluster) + "n_structures": len(data_cluster), } all_metrics.append(cluster_metrics) - + # Store for plotting all_energies_true.extend(energies_true) all_energies_pred.extend(energies_pred) all_forces_true.extend(forces_true.flatten()) all_forces_pred.extend(forces_pred.flatten()) - - + # Save all metrics with open(self.outs_path / "metrics.json", "w") as f: json.dump(all_metrics, f, indent=2) - + # Create visualization fig, axes = plt.subplots(2, 2, figsize=(12, 10)) fig.suptitle("Model Performance by Data Cluster", fontsize=16) - + # Create color map colors = plt.cm.Set1(np.linspace(0, 1, len(self.labels))) - + # Energy relative error plot ax1 = axes[0, 0] start_idx = 0 @@ -131,27 +131,29 @@ def run(self): end_idx = start_idx + len(data_cluster) true_energies = np.array(all_energies_true[start_idx:end_idx]) pred_energies = np.array(all_energies_pred[start_idx:end_idx]) - + # Calculate per-atom energies - true_energies_per_atom = true_energies / np.array([len(atoms) for atoms in data_cluster]) - pred_energies_per_atom = pred_energies / np.array([len(atoms) for atoms in data_cluster]) - + true_energies_per_atom = true_energies / np.array( + [len(atoms) for atoms in data_cluster] + ) + pred_energies_per_atom = pred_energies / np.array( + [len(atoms) for atoms in data_cluster] + ) + # Calculate difference: predicted - true diff = pred_energies_per_atom - true_energies_per_atom - + ax1.scatter( - true_energies_per_atom, - diff, - color=colors[i], label=label, alpha=0.6 + true_energies_per_atom, diff, color=colors[i], label=label, alpha=0.6 ) start_idx = end_idx - - ax1.axhline(y=0, color='k', linestyle='--', alpha=0.5) + + ax1.axhline(y=0, color="k", linestyle="--", alpha=0.5) ax1.set_xlabel("True Energy per Atom / eV") ax1.set_ylabel("Difference: Pred - True / eV") ax1.set_title("Energy Difference Plot") ax1.legend() - + # Forces relative error plot ax2 = axes[0, 1] start_idx = 0 @@ -160,22 +162,18 @@ def run(self): end_idx = start_idx + n_forces true_forces = np.array(all_forces_true[start_idx:end_idx]) pred_forces = np.array(all_forces_pred[start_idx:end_idx]) - + # Calculate difference: predicted - true diff = pred_forces - true_forces - - ax2.scatter( - true_forces, - diff, - color=colors[i], label=label, alpha=0.6, s=1 - ) + + ax2.scatter(true_forces, diff, color=colors[i], label=label, alpha=0.6, s=1) start_idx = end_idx - - ax2.axhline(y=0, color='k', linestyle='--', alpha=0.5) + + ax2.axhline(y=0, color="k", linestyle="--", alpha=0.5) ax2.set_xlabel("True Forces / eV/Å") ax2.set_ylabel("Difference: Pred - True / eV/Å") ax2.set_title("Forces Difference Plot") - + # Energy MAE by cluster ax3 = axes[1, 0] energy_maes = [m["energy"]["mae"] for m in all_metrics] @@ -185,7 +183,7 @@ def run(self): ax3.set_title("Energy MAE by Cluster") ax3.set_xticks(range(len(self.labels))) ax3.set_xticklabels(self.labels, rotation=45) - + # Forces MAE by cluster ax4 = axes[1, 1] forces_maes = [m["forces"]["mae"] for m in all_metrics] @@ -195,73 +193,79 @@ def run(self): ax4.set_title("Forces MAE by Cluster") ax4.set_xticks(range(len(self.labels))) ax4.set_xticklabels(self.labels, rotation=45) - + plt.tight_layout() plt.savefig(self.outs_path / "evaluation_plots.png", dpi=300, bbox_inches="tight") plt.close() - + # Create jointplots for error distribution by label - + # Combined energy error distribution with all labels all_true_energies_per_atom = [] all_energy_errors_per_atom = [] all_energy_labels = [] - + start_idx = 0 for i, (data_cluster, label) in enumerate(zip(self.data, self.labels)): end_idx = start_idx + len(data_cluster) - + true_energies = np.array(all_energies_true[start_idx:end_idx]) pred_energies = np.array(all_energies_pred[start_idx:end_idx]) - - true_energies_per_atom = true_energies / np.array([len(atoms) for atoms in data_cluster]) - pred_energies_per_atom = pred_energies / np.array([len(atoms) for atoms in data_cluster]) + + true_energies_per_atom = true_energies / np.array( + [len(atoms) for atoms in data_cluster] + ) + pred_energies_per_atom = pred_energies / np.array( + [len(atoms) for atoms in data_cluster] + ) energy_errors_per_atom = pred_energies_per_atom - true_energies_per_atom - + all_true_energies_per_atom.extend(true_energies_per_atom) all_energy_errors_per_atom.extend(energy_errors_per_atom) all_energy_labels.extend([label] * len(true_energies_per_atom)) - + start_idx = end_idx - + # Create DataFrame for combined energy plot - energy_df = pd.DataFrame({ - 'true_energy': all_true_energies_per_atom, - 'energy_error': all_energy_errors_per_atom, - 'label': all_energy_labels - }) - + energy_df = pd.DataFrame( + { + "true_energy": all_true_energies_per_atom, + "energy_error": all_energy_errors_per_atom, + "label": all_energy_labels, + } + ) + # Create combined energy jointplot with hue g = sns.jointplot( - data=energy_df, - x='true_energy', - y='energy_error', - hue='label', - height=7 + data=energy_df, x="true_energy", y="energy_error", hue="label", height=7 ) - - g.ax_joint.axhline(y=0, color='red', linestyle='--', alpha=0.7) - g.set_axis_labels('True Energy per Atom / eV', 'Energy Error: Pred - True / eV') - g.figure.suptitle('Combined Energy Error Distribution - All Labels', y=1.02) - + + g.ax_joint.axhline(y=0, color="red", linestyle="--", alpha=0.7) + g.set_axis_labels("True Energy per Atom / eV", "Energy Error: Pred - True / eV") + g.figure.suptitle("Combined Energy Error Distribution - All Labels", y=1.02) + plt.tight_layout() - plt.savefig(self.outs_path / "energy_error_distribution_combined.png", dpi=300, bbox_inches="tight") + plt.savefig( + self.outs_path / "energy_error_distribution_combined.png", + dpi=300, + bbox_inches="tight", + ) plt.close() - + # Combined forces error distribution with all labels all_true_forces = [] all_forces_errors = [] all_forces_labels = [] - + start_idx = 0 for i, (data_cluster, label) in enumerate(zip(self.data, self.labels)): n_forces = len(data_cluster) * 3 * len(data_cluster[0]) end_idx = start_idx + n_forces - + true_forces = np.array(all_forces_true[start_idx:end_idx]) pred_forces = np.array(all_forces_pred[start_idx:end_idx]) forces_errors = pred_forces - true_forces - + # Sample for better visualization n_sample = min(5000, len(true_forces)) # Smaller sample for combined plot if len(true_forces) > n_sample: @@ -271,43 +275,45 @@ def run(self): else: true_forces_sample = true_forces forces_errors_sample = forces_errors - + all_true_forces.extend(true_forces_sample) all_forces_errors.extend(forces_errors_sample) all_forces_labels.extend([label] * len(true_forces_sample)) - + start_idx = end_idx - + # Create DataFrame for combined forces plot - forces_df = pd.DataFrame({ - 'true_forces': all_true_forces, - 'forces_error': all_forces_errors, - 'label': all_forces_labels - }) - + forces_df = pd.DataFrame( + { + "true_forces": all_true_forces, + "forces_error": all_forces_errors, + "label": all_forces_labels, + } + ) + # Create combined forces jointplot with hue g = sns.jointplot( - data=forces_df, - x='true_forces', - y='forces_error', - hue='label', - height=7 + data=forces_df, x="true_forces", y="forces_error", hue="label", height=7 ) - - g.ax_joint.axhline(y=0, color='red', linestyle='--', alpha=0.7) - g.set_axis_labels('True Forces / eV/Å', 'Force Error: Pred - True / eV/Å') - g.figure.suptitle('Combined Forces Error Distribution - All Labels', y=1.02) - + + g.ax_joint.axhline(y=0, color="red", linestyle="--", alpha=0.7) + g.set_axis_labels("True Forces / eV/Å", "Force Error: Pred - True / eV/Å") + g.figure.suptitle("Combined Forces Error Distribution - All Labels", y=1.02) + plt.tight_layout() - plt.savefig(self.outs_path / "forces_error_distribution_combined.png", dpi=300, bbox_inches="tight") + plt.savefig( + self.outs_path / "forces_error_distribution_combined.png", + dpi=300, + bbox_inches="tight", + ) plt.close() - + # Save summary metrics summary_metrics = { "overall_energy_mae": np.mean(energy_maes), "overall_forces_mae": np.mean(forces_maes), - "cluster_metrics": all_metrics + "cluster_metrics": all_metrics, } - + with open(self.outs_path / "summary_metrics.json", "w") as f: json.dump(summary_metrics, f, indent=2) From 51e15ded6a022d6881e932491c5179574ee4f08e Mon Sep 17 00:00:00 2001 From: Fabian Zills <46721498+PythonFZ@users.noreply.github.com> Date: Thu, 10 Jul 2025 19:57:48 +0200 Subject: [PATCH 3/4] Update model_dataset_analysis.py --- ipsuite/analysis/model_dataset_analysis.py | 198 +++++++++++---------- 1 file changed, 100 insertions(+), 98 deletions(-) diff --git a/ipsuite/analysis/model_dataset_analysis.py b/ipsuite/analysis/model_dataset_analysis.py index 082fadbc..43c4d2d9 100644 --- a/ipsuite/analysis/model_dataset_analysis.py +++ b/ipsuite/analysis/model_dataset_analysis.py @@ -1,6 +1,6 @@ +from pathlib import Path import json import os -from pathlib import Path import ase import matplotlib.pyplot as plt @@ -9,7 +9,6 @@ import seaborn as sns import zntrack from apax.nodes import Apax - from ipsuite.utils.metrics import get_full_metrics @@ -67,20 +66,20 @@ class ClusterdDataEvaluation(zntrack.Node): def run(self): os.makedirs(self.outs_path, exist_ok=True) - + calc = self.model.get_calculator() - + all_metrics = [] all_energies_true = [] all_energies_pred = [] all_forces_true = [] all_forces_pred = [] - + for i, (data_cluster, label) in enumerate(zip(self.data, self.labels)): # Get true values energies_true = np.array([x.get_potential_energy() for x in data_cluster]) forces_true = np.array([x.get_forces() for x in data_cluster]) - + # Run predictions if hasattr(calc, "batch_eval"): frames = calc.batch_eval(data_cluster, 1) @@ -90,40 +89,41 @@ def run(self): atoms.calc = calc atoms.get_potential_energy() frames.append(atoms) - + # Get predicted values energies_pred = np.array([x.get_potential_energy() for x in frames]) forces_pred = np.array([x.get_forces() for x in frames]) - + # Calculate metrics energy_metrics = get_full_metrics(energies_true, energies_pred) forces_metrics = get_full_metrics(forces_true, forces_pred) - + cluster_metrics = { "label": label, "energy": energy_metrics, "forces": forces_metrics, - "n_structures": len(data_cluster), + "n_structures": len(data_cluster) } all_metrics.append(cluster_metrics) - + # Store for plotting all_energies_true.extend(energies_true) all_energies_pred.extend(energies_pred) all_forces_true.extend(forces_true.flatten()) all_forces_pred.extend(forces_pred.flatten()) - + + # Save all metrics with open(self.outs_path / "metrics.json", "w") as f: json.dump(all_metrics, f, indent=2) - + # Create visualization fig, axes = plt.subplots(2, 2, figsize=(12, 10)) fig.suptitle("Model Performance by Data Cluster", fontsize=16) - + # Create color map colors = plt.cm.Set1(np.linspace(0, 1, len(self.labels))) - + # Energy relative error plot ax1 = axes[0, 0] start_idx = 0 @@ -131,29 +131,27 @@ def run(self): end_idx = start_idx + len(data_cluster) true_energies = np.array(all_energies_true[start_idx:end_idx]) pred_energies = np.array(all_energies_pred[start_idx:end_idx]) - + # Calculate per-atom energies - true_energies_per_atom = true_energies / np.array( - [len(atoms) for atoms in data_cluster] - ) - pred_energies_per_atom = pred_energies / np.array( - [len(atoms) for atoms in data_cluster] - ) - + true_energies_per_atom = true_energies / np.array([len(atoms) for atoms in data_cluster]) + pred_energies_per_atom = pred_energies / np.array([len(atoms) for atoms in data_cluster]) + # Calculate difference: predicted - true diff = pred_energies_per_atom - true_energies_per_atom - + ax1.scatter( - true_energies_per_atom, diff, color=colors[i], label=label, alpha=0.6 + true_energies_per_atom, + diff, + color=colors[i], label=label, alpha=0.6 ) start_idx = end_idx - - ax1.axhline(y=0, color="k", linestyle="--", alpha=0.5) + + ax1.axhline(y=0, color='k', linestyle='--', alpha=0.5) ax1.set_xlabel("True Energy per Atom / eV") ax1.set_ylabel("Difference: Pred - True / eV") ax1.set_title("Energy Difference Plot") ax1.legend() - + # Forces relative error plot ax2 = axes[0, 1] start_idx = 0 @@ -162,110 +160,116 @@ def run(self): end_idx = start_idx + n_forces true_forces = np.array(all_forces_true[start_idx:end_idx]) pred_forces = np.array(all_forces_pred[start_idx:end_idx]) - + # Calculate difference: predicted - true diff = pred_forces - true_forces - - ax2.scatter(true_forces, diff, color=colors[i], label=label, alpha=0.6, s=1) + + ax2.scatter( + true_forces, + diff, + color=colors[i], label=label, alpha=0.6, s=1 + ) start_idx = end_idx - - ax2.axhline(y=0, color="k", linestyle="--", alpha=0.5) + + ax2.axhline(y=0, color='k', linestyle='--', alpha=0.5) ax2.set_xlabel("True Forces / eV/Å") ax2.set_ylabel("Difference: Pred - True / eV/Å") ax2.set_title("Forces Difference Plot") - + # Energy MAE by cluster ax3 = axes[1, 0] energy_maes = [m["energy"]["mae"] for m in all_metrics] + overall_energy_mae = np.mean(energy_maes) ax3.bar(range(len(self.labels)), energy_maes, color=colors) + ax3.axhline(y=overall_energy_mae, color='red', linestyle='--', alpha=0.7, + label=f'Overall MAE: {overall_energy_mae:.4f} eV') ax3.set_xlabel("Data Cluster") ax3.set_ylabel("Energy MAE / eV") ax3.set_title("Energy MAE by Cluster") ax3.set_xticks(range(len(self.labels))) ax3.set_xticklabels(self.labels, rotation=45) - + ax3.legend() + # Forces MAE by cluster ax4 = axes[1, 1] forces_maes = [m["forces"]["mae"] for m in all_metrics] + overall_forces_mae = np.mean(forces_maes) ax4.bar(range(len(self.labels)), forces_maes, color=colors) + ax4.axhline(y=overall_forces_mae, color='red', linestyle='--', alpha=0.7, + label=f'Overall MAE: {overall_forces_mae:.4f} eV/Å') ax4.set_xlabel("Data Cluster") ax4.set_ylabel("Forces MAE / eV/Å") ax4.set_title("Forces MAE by Cluster") ax4.set_xticks(range(len(self.labels))) ax4.set_xticklabels(self.labels, rotation=45) - + ax4.legend() + plt.tight_layout() plt.savefig(self.outs_path / "evaluation_plots.png", dpi=300, bbox_inches="tight") plt.close() - + # Create jointplots for error distribution by label - + # Combined energy error distribution with all labels all_true_energies_per_atom = [] all_energy_errors_per_atom = [] all_energy_labels = [] - + start_idx = 0 for i, (data_cluster, label) in enumerate(zip(self.data, self.labels)): end_idx = start_idx + len(data_cluster) - + true_energies = np.array(all_energies_true[start_idx:end_idx]) pred_energies = np.array(all_energies_pred[start_idx:end_idx]) - - true_energies_per_atom = true_energies / np.array( - [len(atoms) for atoms in data_cluster] - ) - pred_energies_per_atom = pred_energies / np.array( - [len(atoms) for atoms in data_cluster] - ) + + true_energies_per_atom = true_energies / np.array([len(atoms) for atoms in data_cluster]) + pred_energies_per_atom = pred_energies / np.array([len(atoms) for atoms in data_cluster]) energy_errors_per_atom = pred_energies_per_atom - true_energies_per_atom - + all_true_energies_per_atom.extend(true_energies_per_atom) all_energy_errors_per_atom.extend(energy_errors_per_atom) all_energy_labels.extend([label] * len(true_energies_per_atom)) - + start_idx = end_idx - + # Create DataFrame for combined energy plot - energy_df = pd.DataFrame( - { - "true_energy": all_true_energies_per_atom, - "energy_error": all_energy_errors_per_atom, - "label": all_energy_labels, - } - ) - + energy_df = pd.DataFrame({ + 'true_energy': all_true_energies_per_atom, + 'energy_error': all_energy_errors_per_atom, + 'label': all_energy_labels + }) + # Create combined energy jointplot with hue g = sns.jointplot( - data=energy_df, x="true_energy", y="energy_error", hue="label", height=7 + data=energy_df, + x='true_energy', + y='energy_error', + hue='label', + height=7 ) - - g.ax_joint.axhline(y=0, color="red", linestyle="--", alpha=0.7) - g.set_axis_labels("True Energy per Atom / eV", "Energy Error: Pred - True / eV") - g.figure.suptitle("Combined Energy Error Distribution - All Labels", y=1.02) - + + g.ax_joint.axhline(y=0, color='red', linestyle='--', alpha=0.7) + g.set_axis_labels('True Energy per Atom / eV', 'Energy Error: Pred - True / eV') + g.figure.suptitle('Combined Energy Error Distribution - All Labels', y=1.02) + plt.tight_layout() - plt.savefig( - self.outs_path / "energy_error_distribution_combined.png", - dpi=300, - bbox_inches="tight", - ) + plt.savefig(self.outs_path / "energy_error_distribution_combined.png", dpi=300, bbox_inches="tight") plt.close() - + # Combined forces error distribution with all labels all_true_forces = [] all_forces_errors = [] all_forces_labels = [] - + start_idx = 0 for i, (data_cluster, label) in enumerate(zip(self.data, self.labels)): n_forces = len(data_cluster) * 3 * len(data_cluster[0]) end_idx = start_idx + n_forces - + true_forces = np.array(all_forces_true[start_idx:end_idx]) pred_forces = np.array(all_forces_pred[start_idx:end_idx]) forces_errors = pred_forces - true_forces - + # Sample for better visualization n_sample = min(5000, len(true_forces)) # Smaller sample for combined plot if len(true_forces) > n_sample: @@ -275,45 +279,43 @@ def run(self): else: true_forces_sample = true_forces forces_errors_sample = forces_errors - + all_true_forces.extend(true_forces_sample) all_forces_errors.extend(forces_errors_sample) all_forces_labels.extend([label] * len(true_forces_sample)) - + start_idx = end_idx - + # Create DataFrame for combined forces plot - forces_df = pd.DataFrame( - { - "true_forces": all_true_forces, - "forces_error": all_forces_errors, - "label": all_forces_labels, - } - ) - + forces_df = pd.DataFrame({ + 'true_forces': all_true_forces, + 'forces_error': all_forces_errors, + 'label': all_forces_labels + }) + # Create combined forces jointplot with hue g = sns.jointplot( - data=forces_df, x="true_forces", y="forces_error", hue="label", height=7 + data=forces_df, + x='true_forces', + y='forces_error', + hue='label', + height=7 ) - - g.ax_joint.axhline(y=0, color="red", linestyle="--", alpha=0.7) - g.set_axis_labels("True Forces / eV/Å", "Force Error: Pred - True / eV/Å") - g.figure.suptitle("Combined Forces Error Distribution - All Labels", y=1.02) - + + g.ax_joint.axhline(y=0, color='red', linestyle='--', alpha=0.7) + g.set_axis_labels('True Forces / eV/Å', 'Force Error: Pred - True / eV/Å') + g.figure.suptitle('Combined Forces Error Distribution - All Labels', y=1.02) + plt.tight_layout() - plt.savefig( - self.outs_path / "forces_error_distribution_combined.png", - dpi=300, - bbox_inches="tight", - ) + plt.savefig(self.outs_path / "forces_error_distribution_combined.png", dpi=300, bbox_inches="tight") plt.close() - + # Save summary metrics summary_metrics = { "overall_energy_mae": np.mean(energy_maes), "overall_forces_mae": np.mean(forces_maes), - "cluster_metrics": all_metrics, + "cluster_metrics": all_metrics } - + with open(self.outs_path / "summary_metrics.json", "w") as f: json.dump(summary_metrics, f, indent=2) From 7af5b596c28fb39aeecc93bec1ffaff148838bfb Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 10 Jul 2025 17:58:05 +0000 Subject: [PATCH 4/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- ipsuite/analysis/model_dataset_analysis.py | 208 +++++++++++---------- 1 file changed, 112 insertions(+), 96 deletions(-) diff --git a/ipsuite/analysis/model_dataset_analysis.py b/ipsuite/analysis/model_dataset_analysis.py index 43c4d2d9..7b4b6b37 100644 --- a/ipsuite/analysis/model_dataset_analysis.py +++ b/ipsuite/analysis/model_dataset_analysis.py @@ -1,6 +1,6 @@ -from pathlib import Path import json import os +from pathlib import Path import ase import matplotlib.pyplot as plt @@ -9,6 +9,7 @@ import seaborn as sns import zntrack from apax.nodes import Apax + from ipsuite.utils.metrics import get_full_metrics @@ -66,20 +67,20 @@ class ClusterdDataEvaluation(zntrack.Node): def run(self): os.makedirs(self.outs_path, exist_ok=True) - + calc = self.model.get_calculator() - + all_metrics = [] all_energies_true = [] all_energies_pred = [] all_forces_true = [] all_forces_pred = [] - + for i, (data_cluster, label) in enumerate(zip(self.data, self.labels)): # Get true values energies_true = np.array([x.get_potential_energy() for x in data_cluster]) forces_true = np.array([x.get_forces() for x in data_cluster]) - + # Run predictions if hasattr(calc, "batch_eval"): frames = calc.batch_eval(data_cluster, 1) @@ -89,41 +90,40 @@ def run(self): atoms.calc = calc atoms.get_potential_energy() frames.append(atoms) - + # Get predicted values energies_pred = np.array([x.get_potential_energy() for x in frames]) forces_pred = np.array([x.get_forces() for x in frames]) - + # Calculate metrics energy_metrics = get_full_metrics(energies_true, energies_pred) forces_metrics = get_full_metrics(forces_true, forces_pred) - + cluster_metrics = { "label": label, "energy": energy_metrics, "forces": forces_metrics, - "n_structures": len(data_cluster) + "n_structures": len(data_cluster), } all_metrics.append(cluster_metrics) - + # Store for plotting all_energies_true.extend(energies_true) all_energies_pred.extend(energies_pred) all_forces_true.extend(forces_true.flatten()) all_forces_pred.extend(forces_pred.flatten()) - - + # Save all metrics with open(self.outs_path / "metrics.json", "w") as f: json.dump(all_metrics, f, indent=2) - + # Create visualization fig, axes = plt.subplots(2, 2, figsize=(12, 10)) fig.suptitle("Model Performance by Data Cluster", fontsize=16) - + # Create color map colors = plt.cm.Set1(np.linspace(0, 1, len(self.labels))) - + # Energy relative error plot ax1 = axes[0, 0] start_idx = 0 @@ -131,27 +131,29 @@ def run(self): end_idx = start_idx + len(data_cluster) true_energies = np.array(all_energies_true[start_idx:end_idx]) pred_energies = np.array(all_energies_pred[start_idx:end_idx]) - + # Calculate per-atom energies - true_energies_per_atom = true_energies / np.array([len(atoms) for atoms in data_cluster]) - pred_energies_per_atom = pred_energies / np.array([len(atoms) for atoms in data_cluster]) - + true_energies_per_atom = true_energies / np.array( + [len(atoms) for atoms in data_cluster] + ) + pred_energies_per_atom = pred_energies / np.array( + [len(atoms) for atoms in data_cluster] + ) + # Calculate difference: predicted - true diff = pred_energies_per_atom - true_energies_per_atom - + ax1.scatter( - true_energies_per_atom, - diff, - color=colors[i], label=label, alpha=0.6 + true_energies_per_atom, diff, color=colors[i], label=label, alpha=0.6 ) start_idx = end_idx - - ax1.axhline(y=0, color='k', linestyle='--', alpha=0.5) + + ax1.axhline(y=0, color="k", linestyle="--", alpha=0.5) ax1.set_xlabel("True Energy per Atom / eV") ax1.set_ylabel("Difference: Pred - True / eV") ax1.set_title("Energy Difference Plot") ax1.legend() - + # Forces relative error plot ax2 = axes[0, 1] start_idx = 0 @@ -160,116 +162,128 @@ def run(self): end_idx = start_idx + n_forces true_forces = np.array(all_forces_true[start_idx:end_idx]) pred_forces = np.array(all_forces_pred[start_idx:end_idx]) - + # Calculate difference: predicted - true diff = pred_forces - true_forces - - ax2.scatter( - true_forces, - diff, - color=colors[i], label=label, alpha=0.6, s=1 - ) + + ax2.scatter(true_forces, diff, color=colors[i], label=label, alpha=0.6, s=1) start_idx = end_idx - - ax2.axhline(y=0, color='k', linestyle='--', alpha=0.5) + + ax2.axhline(y=0, color="k", linestyle="--", alpha=0.5) ax2.set_xlabel("True Forces / eV/Å") ax2.set_ylabel("Difference: Pred - True / eV/Å") ax2.set_title("Forces Difference Plot") - + # Energy MAE by cluster ax3 = axes[1, 0] energy_maes = [m["energy"]["mae"] for m in all_metrics] overall_energy_mae = np.mean(energy_maes) ax3.bar(range(len(self.labels)), energy_maes, color=colors) - ax3.axhline(y=overall_energy_mae, color='red', linestyle='--', alpha=0.7, - label=f'Overall MAE: {overall_energy_mae:.4f} eV') + ax3.axhline( + y=overall_energy_mae, + color="red", + linestyle="--", + alpha=0.7, + label=f"Overall MAE: {overall_energy_mae:.4f} eV", + ) ax3.set_xlabel("Data Cluster") ax3.set_ylabel("Energy MAE / eV") ax3.set_title("Energy MAE by Cluster") ax3.set_xticks(range(len(self.labels))) ax3.set_xticklabels(self.labels, rotation=45) ax3.legend() - + # Forces MAE by cluster ax4 = axes[1, 1] forces_maes = [m["forces"]["mae"] for m in all_metrics] overall_forces_mae = np.mean(forces_maes) ax4.bar(range(len(self.labels)), forces_maes, color=colors) - ax4.axhline(y=overall_forces_mae, color='red', linestyle='--', alpha=0.7, - label=f'Overall MAE: {overall_forces_mae:.4f} eV/Å') + ax4.axhline( + y=overall_forces_mae, + color="red", + linestyle="--", + alpha=0.7, + label=f"Overall MAE: {overall_forces_mae:.4f} eV/Å", + ) ax4.set_xlabel("Data Cluster") ax4.set_ylabel("Forces MAE / eV/Å") ax4.set_title("Forces MAE by Cluster") ax4.set_xticks(range(len(self.labels))) ax4.set_xticklabels(self.labels, rotation=45) ax4.legend() - + plt.tight_layout() plt.savefig(self.outs_path / "evaluation_plots.png", dpi=300, bbox_inches="tight") plt.close() - + # Create jointplots for error distribution by label - + # Combined energy error distribution with all labels all_true_energies_per_atom = [] all_energy_errors_per_atom = [] all_energy_labels = [] - + start_idx = 0 for i, (data_cluster, label) in enumerate(zip(self.data, self.labels)): end_idx = start_idx + len(data_cluster) - + true_energies = np.array(all_energies_true[start_idx:end_idx]) pred_energies = np.array(all_energies_pred[start_idx:end_idx]) - - true_energies_per_atom = true_energies / np.array([len(atoms) for atoms in data_cluster]) - pred_energies_per_atom = pred_energies / np.array([len(atoms) for atoms in data_cluster]) + + true_energies_per_atom = true_energies / np.array( + [len(atoms) for atoms in data_cluster] + ) + pred_energies_per_atom = pred_energies / np.array( + [len(atoms) for atoms in data_cluster] + ) energy_errors_per_atom = pred_energies_per_atom - true_energies_per_atom - + all_true_energies_per_atom.extend(true_energies_per_atom) all_energy_errors_per_atom.extend(energy_errors_per_atom) all_energy_labels.extend([label] * len(true_energies_per_atom)) - + start_idx = end_idx - + # Create DataFrame for combined energy plot - energy_df = pd.DataFrame({ - 'true_energy': all_true_energies_per_atom, - 'energy_error': all_energy_errors_per_atom, - 'label': all_energy_labels - }) - + energy_df = pd.DataFrame( + { + "true_energy": all_true_energies_per_atom, + "energy_error": all_energy_errors_per_atom, + "label": all_energy_labels, + } + ) + # Create combined energy jointplot with hue g = sns.jointplot( - data=energy_df, - x='true_energy', - y='energy_error', - hue='label', - height=7 + data=energy_df, x="true_energy", y="energy_error", hue="label", height=7 ) - - g.ax_joint.axhline(y=0, color='red', linestyle='--', alpha=0.7) - g.set_axis_labels('True Energy per Atom / eV', 'Energy Error: Pred - True / eV') - g.figure.suptitle('Combined Energy Error Distribution - All Labels', y=1.02) - + + g.ax_joint.axhline(y=0, color="red", linestyle="--", alpha=0.7) + g.set_axis_labels("True Energy per Atom / eV", "Energy Error: Pred - True / eV") + g.figure.suptitle("Combined Energy Error Distribution - All Labels", y=1.02) + plt.tight_layout() - plt.savefig(self.outs_path / "energy_error_distribution_combined.png", dpi=300, bbox_inches="tight") + plt.savefig( + self.outs_path / "energy_error_distribution_combined.png", + dpi=300, + bbox_inches="tight", + ) plt.close() - + # Combined forces error distribution with all labels all_true_forces = [] all_forces_errors = [] all_forces_labels = [] - + start_idx = 0 for i, (data_cluster, label) in enumerate(zip(self.data, self.labels)): n_forces = len(data_cluster) * 3 * len(data_cluster[0]) end_idx = start_idx + n_forces - + true_forces = np.array(all_forces_true[start_idx:end_idx]) pred_forces = np.array(all_forces_pred[start_idx:end_idx]) forces_errors = pred_forces - true_forces - + # Sample for better visualization n_sample = min(5000, len(true_forces)) # Smaller sample for combined plot if len(true_forces) > n_sample: @@ -279,43 +293,45 @@ def run(self): else: true_forces_sample = true_forces forces_errors_sample = forces_errors - + all_true_forces.extend(true_forces_sample) all_forces_errors.extend(forces_errors_sample) all_forces_labels.extend([label] * len(true_forces_sample)) - + start_idx = end_idx - + # Create DataFrame for combined forces plot - forces_df = pd.DataFrame({ - 'true_forces': all_true_forces, - 'forces_error': all_forces_errors, - 'label': all_forces_labels - }) - + forces_df = pd.DataFrame( + { + "true_forces": all_true_forces, + "forces_error": all_forces_errors, + "label": all_forces_labels, + } + ) + # Create combined forces jointplot with hue g = sns.jointplot( - data=forces_df, - x='true_forces', - y='forces_error', - hue='label', - height=7 + data=forces_df, x="true_forces", y="forces_error", hue="label", height=7 ) - - g.ax_joint.axhline(y=0, color='red', linestyle='--', alpha=0.7) - g.set_axis_labels('True Forces / eV/Å', 'Force Error: Pred - True / eV/Å') - g.figure.suptitle('Combined Forces Error Distribution - All Labels', y=1.02) - + + g.ax_joint.axhline(y=0, color="red", linestyle="--", alpha=0.7) + g.set_axis_labels("True Forces / eV/Å", "Force Error: Pred - True / eV/Å") + g.figure.suptitle("Combined Forces Error Distribution - All Labels", y=1.02) + plt.tight_layout() - plt.savefig(self.outs_path / "forces_error_distribution_combined.png", dpi=300, bbox_inches="tight") + plt.savefig( + self.outs_path / "forces_error_distribution_combined.png", + dpi=300, + bbox_inches="tight", + ) plt.close() - + # Save summary metrics summary_metrics = { "overall_energy_mae": np.mean(energy_maes), "overall_forces_mae": np.mean(forces_maes), - "cluster_metrics": all_metrics + "cluster_metrics": all_metrics, } - + with open(self.outs_path / "summary_metrics.json", "w") as f: json.dump(summary_metrics, f, indent=2)