diff --git a/docs/source/user_guide/benchmarks/actinides.rst b/docs/source/user_guide/benchmarks/actinides.rst new file mode 100644 index 000000000..4df7b2cb1 --- /dev/null +++ b/docs/source/user_guide/benchmarks/actinides.rst @@ -0,0 +1,36 @@ +================== +Actinides +================== + +Plutonium Dioxide +================== + +Summary +------- + +General performance on Plutonium Dioxide against DFT+U calculations. The DFT+U calculations are evaluted on samples in the temperature range 0-1200K and been have parameterized to correctly predict the lattice constant (within 0.3%) and thermal expansion at low temperature. + +Metrics +------- + +1. Energy MAE (PBE+U) + +Mean absolute error of energy predictions (per atom). + +2. Force MAE (PBE+U) + +Mean absolute error of force (individual components) predictions against DFT+U calculations. + +3. Stress MAE (PBE+U) + +Mean absolute error of stress (individual tensor components) predictions against DFT+U calculations. + +Computational cost +------------------ + +Low + +Data availability +----------------- + +Reference data: availabile in repo, for specific calculation details contact willdavie2002@gmail.com. diff --git a/ml_peg/analysis/actinides/plutonium_dioxide/analyse_plutonium_dioxide.py b/ml_peg/analysis/actinides/plutonium_dioxide/analyse_plutonium_dioxide.py new file mode 100644 index 000000000..30f77a051 --- /dev/null +++ b/ml_peg/analysis/actinides/plutonium_dioxide/analyse_plutonium_dioxide.py @@ -0,0 +1,350 @@ +"""Plutonium Dioxide benchmark against DFT+U.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +from ase import io, units +import numpy as np +import pytest + +from ml_peg.analysis.utils.decorators import build_table, plot_density_scatter +from ml_peg.analysis.utils.utils import ( + build_density_inputs, + load_metrics_config, + mae, + write_density_trajectories, +) +from ml_peg.app import APP_ROOT +from ml_peg.calcs import CALCS_ROOT +from ml_peg.models.get_models import get_model_names +from ml_peg.models.models import current_models + +MODELS = get_model_names(current_models) +CALC_PATH = CALCS_ROOT / "actinides" / "plutonium_dioxide" / "outputs" +OUT_PATH = APP_ROOT / "data" / "actinides" / "plutonium_dioxide" + +METRICS_CONFIG_PATH = Path(__file__).with_name("metrics.yml") +DEFAULT_THRESHOLDS, DEFAULT_TOOLTIPS, DEFAULT_WEIGHTS = load_metrics_config( + METRICS_CONFIG_PATH +) + +EV_TO_KJ_PER_MOL = units.mol / units.kJ + + +@pytest.fixture +def puo2_stats() -> dict[str, dict[str, Any]]: + """ + Load and cache statistics per model. + + Returns + ------- + dict[str, dict[str, Any]] + Processed information per model (energy, force, stress, labels). + """ + OUT_PATH.mkdir(parents=True, exist_ok=True) + stats: dict[str, dict[str, Any]] = {} + + for model_name in MODELS: + model_dir = CALC_PATH / model_name + if not model_dir.exists(): + continue + + struct_dir = OUT_PATH / model_name + struct_dir.mkdir(parents=True, exist_ok=True) + + energies_ref, energies_pred = [], [] + forces_ref, forces_pred = [], [] + stress_ref, stress_pred = [], [] + energy_labels: list[str] = [] + force_labels: list[str] = [] + stress_labels: list[str] = [] + excluded = 0 + frame_idx = 0 + + for xyz_file in sorted(model_dir.glob("*.xyz")): + frames = io.read(xyz_file, ":") + for atoms in frames: + label = str(frame_idx) + natoms = atoms.get_number_of_atoms() + e_ref = atoms.info.get("energy_xtb") + f_ref = atoms.arrays.get("forces_xtb") + s_ref = atoms.info.get("REF_stress") + + io.write(struct_dir / f"{label}.xyz", atoms) + + if e_ref is not None: + energies_ref.append(e_ref / natoms) + energies_pred.append(atoms.get_total_energy() / natoms) + energy_labels.append(label) + + if f_ref is not None: + forces_ref.append(f_ref.ravel()) + forces_pred.append(atoms.get_forces().ravel()) + force_labels.extend([label] * (natoms * 3)) + + if s_ref is not None: + s_ref_flat = np.asarray(s_ref).ravel().tolist() + stress_ref.extend(s_ref_flat) + stress_pred.extend(atoms.get_stress(voigt=False).ravel()) + stress_labels.extend([label] * len(s_ref_flat)) + + frame_idx += 1 + + stats[model_name] = { + "energies": { + "ref": energies_ref, + "pred": energies_pred, + }, + "forces": { + "ref": np.concatenate(forces_ref).tolist() if forces_ref else [], + "pred": np.concatenate(forces_pred).tolist() if forces_pred else [], + }, + "stress": { + "ref": stress_ref, + "pred": stress_pred, + }, + "excluded": excluded, + "energy_labels": energy_labels, + "force_labels": force_labels, + "stress_labels": stress_labels, + } + return stats + + +@pytest.fixture +def energy_mae(puo2_stats: dict[str, dict[str, Any]]) -> dict[str, float | None]: + """ + Mean absolute error for energy predictions. + + Parameters + ---------- + puo2_stats + Aggregrated energy/force/stress per model. + + Returns + ------- + dict[str, float | None] + MAE values for each model (``None`` if no data). + """ + results: dict[str, float | None] = {} + for model_name, props in puo2_stats.items(): + energies = props.get("energies", {}) + results[model_name] = mae(energies.get("ref", []), energies.get("pred", [])) + return results + + +@pytest.fixture +def forces_mae(puo2_stats: dict[str, dict[str, Any]]) -> dict[str, float | None]: + """ + Mean absolute error for force predictions. + + Parameters + ---------- + puo2_stats + Aggregrated energy/force/stress per model. + + Returns + ------- + dict[str, float | None] + MAE values for each model (``None`` if no data). + """ + results: dict[str, float | None] = {} + for model_name, props in puo2_stats.items(): + forces = props.get("forces", {}) + results[model_name] = mae(forces.get("ref", []), forces.get("pred", [])) + return results + + +@pytest.fixture +def stress_mae(puo2_stats: dict[str, dict[str, Any]]) -> dict[str, float | None]: + """ + Mean absolute error for stress predictions. + + Parameters + ---------- + puo2_stats + Aggregrated energy/force/stress per model. + + Returns + ------- + dict[str, float | None] + MAE values for each model (``None`` if no data). + """ + results: dict[str, float | None] = {} + for model_name, props in puo2_stats.items(): + stress = props.get("stress", {}) + results[model_name] = mae(stress.get("ref", []), stress.get("pred", [])) + return results + + +# Density plots for each metric. + + +@pytest.fixture +@plot_density_scatter( + filename=OUT_PATH / "figure_energy_density.json", + title="Relative Energy Plutonium Dioxide", + x_label="PBE+U Reference Energy / eV / Atom", + y_label="Predicted Energy / eV / Atom", + annotation_metadata={"excluded": "Excluded"}, +) +def energy_density(puo2_stats: dict[str, dict[str, Any]]) -> dict[str, dict]: + """ + Density scatter input for energy. + + Parameters + ---------- + puo2_stats + Aggregrated energy/force/stress per model. + + Returns + ------- + dict[str, dict] + Mapping of model name to density-scatter data. + """ + result = build_density_inputs( + list(puo2_stats.keys()), + puo2_stats, + "energies", + metric_fn=mae, + ) + for model_name, model_stats in puo2_stats.items(): + write_density_trajectories( + labels_list=model_stats["energy_labels"], + ref_vals=model_stats["energies"]["ref"], + pred_vals=model_stats["energies"]["pred"], + struct_dir=OUT_PATH / model_name, + traj_dir=OUT_PATH / model_name / "density_traj_energy", + struct_filename_builder=lambda label: f"{label}.xyz", + ) + return result + + +@pytest.fixture +@plot_density_scatter( + filename=OUT_PATH / "figure_force_density.json", + title="Forces Plutonium Dioxide", + x_label="PBE+U Reference Forces / eV / Å", + y_label="Predicted Forces / eV / Å", + annotation_metadata={"excluded": "Excluded"}, +) +def force_density(puo2_stats: dict[str, dict[str, Any]]) -> dict[str, dict]: + """ + Density scatter input for force. + + Parameters + ---------- + puo2_stats + Aggregrated energy/force/stress per model. + + Returns + ------- + dict[str, dict] + Mapping of model name to density-scatter data. + """ + result = build_density_inputs( + list(puo2_stats.keys()), + puo2_stats, + "forces", + metric_fn=mae, + ) + for model_name, model_stats in puo2_stats.items(): + write_density_trajectories( + labels_list=model_stats["force_labels"], + ref_vals=model_stats["forces"]["ref"], + pred_vals=model_stats["forces"]["pred"], + struct_dir=OUT_PATH / model_name, + traj_dir=OUT_PATH / model_name / "density_traj_force", + struct_filename_builder=lambda label: f"{label}.xyz", + ) + return result + + +@pytest.fixture +@plot_density_scatter( + filename=OUT_PATH / "figure_stress_density.json", + title="Stress Plutonium Dioxide", + x_label="PBE+U Reference Stress / eV / ų", + y_label="Predicted Stress / eV / ų", + annotation_metadata={"excluded": "Excluded"}, +) +def stress_density(puo2_stats: dict[str, dict[str, Any]]) -> dict[str, dict]: + """ + Density scatter input for stress. + + Parameters + ---------- + puo2_stats + Aggregrated energy/force/stress per model. + + Returns + ------- + dict[str, dict] + Mapping of model name to density-scatter data. + """ + result = build_density_inputs( + list(puo2_stats.keys()), + puo2_stats, + "stress", + metric_fn=mae, + ) + for model_name, model_stats in puo2_stats.items(): + write_density_trajectories( + labels_list=model_stats["stress_labels"], + ref_vals=model_stats["stress"]["ref"], + pred_vals=model_stats["stress"]["pred"], + struct_dir=OUT_PATH / model_name, + traj_dir=OUT_PATH / model_name / "density_traj_stress", + struct_filename_builder=lambda label: f"{label}.xyz", + ) + return result + + +@pytest.fixture +@build_table( + filename=OUT_PATH / "puo2_metrics_table.json", + metric_tooltips=DEFAULT_TOOLTIPS, + thresholds=DEFAULT_THRESHOLDS, + weights=DEFAULT_WEIGHTS, +) +def metrics( + energy_mae: dict[str, float | None], + forces_mae: dict[str, float | None], + stress_mae: dict[str, float | None], +) -> dict[str, dict]: + """ + Metric table. + + Parameters + ---------- + energy_mae + Energy MAE per model. + forces_mae + Force MAE per model. + stress_mae + Stress MAE per model. + + Returns + ------- + dict[str, dict] + Mapping of metric name to model-value dictionaries. + """ + return { + "Energy MAE": energy_mae, + "Force MAE": forces_mae, + "Stress MAE": stress_mae, + } + + +def test_puo2(metrics: dict[str, dict]) -> None: + """ + Run puo2 analysis. + + Parameters + ---------- + metrics + Benchmark metric values. + """ + return diff --git a/ml_peg/analysis/actinides/plutonium_dioxide/metrics.yml b/ml_peg/analysis/actinides/plutonium_dioxide/metrics.yml new file mode 100644 index 000000000..2fdd7f4e1 --- /dev/null +++ b/ml_peg/analysis/actinides/plutonium_dioxide/metrics.yml @@ -0,0 +1,19 @@ +metrics: + Energy MAE: + good: 0 + bad: 3 + unit: eV / Atom + tooltip: Mean absolute error of energy. + level_of_theory: PBE+U + Force MAE: + good: 0 + bad: 0.5 + unit: eV / Å + tooltip: Mean absolute error of individual force components. + level_of_theory: PBE+U + Stress MAE: + good: 0 + bad: 0.01 + unit: eV / ų + tooltip: Mean absolute error of individual stress components. + level_of_theory: PBE+U diff --git a/ml_peg/app/actinides/plutonium_dioxide/app_plutonium_dioxide.py b/ml_peg/app/actinides/plutonium_dioxide/app_plutonium_dioxide.py new file mode 100644 index 000000000..8413f30dd --- /dev/null +++ b/ml_peg/app/actinides/plutonium_dioxide/app_plutonium_dioxide.py @@ -0,0 +1,111 @@ +"""Run PuO2 force benchmark app.""" + +from __future__ import annotations + +from dash import Dash +from dash.html import Div + +from ml_peg.app import APP_ROOT +from ml_peg.app.base_app import BaseApp +from ml_peg.app.utils.build_callbacks import plot_from_table_cell, struct_from_scatter +from ml_peg.app.utils.load import collect_traj_assets, read_density_plot_for_model +from ml_peg.models.get_models import get_model_names +from ml_peg.models.models import current_models + +MODELS = get_model_names(current_models) +BENCHMARK_NAME = "Plutonium Dioxide" +DOCS_URL = "https://ddmms.github.io/ml-peg/user_guide/benchmarks/actinides.html#plutonium_dioxide" +DATA_PATH = APP_ROOT / "data" / "actinides" / "plutonium_dioxide" + + +class PuO2App(BaseApp): + """PuO2 benchmark app.""" + + def register_callbacks(self) -> None: + """Register callbacks to app.""" + density_plots = {} + + for model in MODELS: + # We explicitly map the exact Column Header strings from metrics.yml + # to their respective JSON figure files. + plots = { + "Energy MAE": read_density_plot_for_model( + filename=DATA_PATH / "figure_energy_density.json", + model=model, + id=f"{BENCHMARK_NAME}-{model}-energy-figure", + ), + "Force MAE": read_density_plot_for_model( + filename=DATA_PATH / "figure_force_density.json", + model=model, + id=f"{BENCHMARK_NAME}-{model}-force-figure", + ), + "Stress MAE": read_density_plot_for_model( + filename=DATA_PATH / "figure_stress_density.json", + model=model, + id=f"{BENCHMARK_NAME}-{model}-stress-figure", + ), + } + + # This comprehension is key: it prevents the app from trying to + # link to a plot if the data is missing for that specific model. + model_plots = {k: v for k, v in plots.items() if v is not None} + + if model_plots: + density_plots[model] = model_plots + + # This links the sanitized ID to the compiled dictionary of figures + plot_from_table_cell( + table_id=self.table_id, + plot_id=f"{BENCHMARK_NAME}-figure-placeholder", + cell_to_plot=density_plots, + ) + + assets_prefix = "assets/actinides/plutonium_dioxide" + + for prop, scatter_suffix in [ + ("density_traj_energy", "energy-figure"), + ("density_traj_force", "force-figure"), + ("density_traj_stress", "stress-figure"), + ]: + struct_trajs = collect_traj_assets( + data_path=DATA_PATH, + assets_prefix=assets_prefix, + models=MODELS, + traj_dirname=prop, + ) + for model in struct_trajs: + struct_from_scatter( + scatter_id=f"{BENCHMARK_NAME}-{model}-{scatter_suffix}", + struct_id=f"{BENCHMARK_NAME}-struct-placeholder", + structs=struct_trajs[model], + mode="traj", + ) + + +def get_app() -> PuO2App: + """ + Get PuO2 benchmark app instance with layout and callbacks. + + Returns + ------- + PuO2App + Benchmark layout and callback registration. + """ + return PuO2App( + name=BENCHMARK_NAME, + description=("Basic performance metrics on plutonium (IV) dioxide"), + docs_url=DOCS_URL, + table_path=DATA_PATH / "puo2_metrics_table.json", + extra_components=[ + Div(id=f"{BENCHMARK_NAME}-figure-placeholder"), + Div(id=f"{BENCHMARK_NAME}-struct-placeholder"), + ], + ) + + +if __name__ == "__main__": + full_app = Dash(__name__) + puo2_app = get_app() + full_app.layout = puo2_app.layout + puo2_app.register_callbacks() + full_app.run(port=8050, debug=True) diff --git a/ml_peg/calcs/actinides/plutonium_dioxide/calc_plutonium_dioxide.py b/ml_peg/calcs/actinides/plutonium_dioxide/calc_plutonium_dioxide.py new file mode 100644 index 000000000..9ae55db70 --- /dev/null +++ b/ml_peg/calcs/actinides/plutonium_dioxide/calc_plutonium_dioxide.py @@ -0,0 +1,62 @@ +"""Calculate plutonium dioxide metrics.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +from ase import units +from ase.io import read, write +import pytest + +from ml_peg.calcs.utils.utils import download_s3_data +from ml_peg.models.get_models import load_models +from ml_peg.models.models import current_models + +MODELS = load_models(current_models) + +DATA_PATH = Path(__file__).parent / "data" +OUT_PATH = Path(__file__).parent / "outputs" + +EV_TO_KJ_PER_MOL = units.mol / units.kJ + + +@pytest.mark.parametrize("mlip", MODELS.items()) +def test_puo2_parity(mlip: tuple[str, Any]) -> None: + """ + Generate data for MAE analysis and density plots. + + Parameters + ---------- + mlip : tuple[str, Any] + A tuple of (model_name, model_object). + """ + model_name, model = mlip + + # Download data. + puo2_data_dir = ( + download_s3_data( + key="inputs/actinides/plutonium_dioxide/plutonium_dioxide.zip", + filename="plutonium_dioxide.zip", + ) + / "plutonium_dioxide" + ) + + ref_file = puo2_data_dir / "dft_ref_data.xyz" + ref_structures = read(ref_file, ":") + + results_to_save = [] + + for atoms in ref_structures: + atoms.calc = ( + model.get_calculator() + ) # must be called each time as number of atoms changes. + atoms.get_potential_energy() + atoms.get_forces() + atoms.get_stress() + + results_to_save.append(atoms) + + write_dir = OUT_PATH / model_name + write_dir.mkdir(parents=True, exist_ok=True) + write(write_dir / "puo2_results.xyz", results_to_save)