From efdd0d9ea76cc92cb5306dde0449b5d34a22cca9 Mon Sep 17 00:00:00 2001 From: DrDavie1 <101512473+DrDavie1@users.noreply.github.com> Date: Sun, 1 Mar 2026 17:12:32 +0000 Subject: [PATCH 1/6] add puo2 benchmark --- .../user_guide/benchmarks/actinides.rst | 36 ++ .../analyse_plutonium_dioxide.py | 312 ++++++++++++++++++ .../actinides/plutonium_dioxide/metrics.yml | 19 ++ .../app_plutonium_dioxide.py | 87 +++++ .../calc_plutonium_dioxide.py | 66 ++++ 5 files changed, 520 insertions(+) create mode 100644 docs/source/user_guide/benchmarks/actinides.rst create mode 100644 ml_peg/analysis/actinides/plutonium_dioxide/analyse_plutonium_dioxide.py create mode 100644 ml_peg/analysis/actinides/plutonium_dioxide/metrics.yml create mode 100644 ml_peg/app/actinides/plutonium_dioxide/app_plutonium_dioxide.py create mode 100644 ml_peg/calcs/actinides/plutonium_dioxide/calc_plutonium_dioxide.py 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..8d52bbee4 --- /dev/null +++ b/ml_peg/analysis/actinides/plutonium_dioxide/analyse_plutonium_dioxide.py @@ -0,0 +1,312 @@ +"""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 +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). + """ + 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 + + energies_ref, energies_pred = [], [] + forces_ref, forces_pred = [], [] + stress_ref, stress_pred = [], [] + excluded = 0 + + for xyz_file in sorted(model_dir.glob("*.xyz")): + frames = io.read(xyz_file, ":") + for atoms in frames: + 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") + + if e_ref is not None: + energies_ref.append(e_ref / natoms) + energies_pred.append(atoms.get_total_energy() / natoms) + + if f_ref is not None: + forces_ref.append(f_ref.ravel()) + forces_pred.append(atoms.get_forces().ravel()) + + if s_ref is not None: + stress_ref.extend(s_ref.tolist()) + stress_pred.extend(atoms.get_stress(voigt=False).ravel()) + + stats[model_name] = { + "energies": { + "ref": energies_ref, + "pred": energies_pred, + }, + "forces": { + "ref": np.concatenate(forces_ref).tolist(), + "pred": np.concatenate(forces_pred).tolist(), + }, + "stress": { + "ref": stress_ref, + "pred": stress_pred, + }, + "excluded": excluded, + } + 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. + """ + return build_density_inputs( + list(puo2_stats.keys()), + puo2_stats, + "energies", + metric_fn=mae, + ) + + +@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. + """ + return build_density_inputs( + list(puo2_stats.keys()), + puo2_stats, + "forces", + metric_fn=mae, + ) + + +@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. + """ + return build_density_inputs( + list(puo2_stats.keys()), + puo2_stats, + "stress", + metric_fn=mae, + ) + + +@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], + energy_density: dict[str, dict], + force_density: dict[str, dict], + stress_density: dict[str, dict], +) -> None: + """ + Run puo2 analysis. + + Parameters + ---------- + metrics + Benchmark metric values. + energy_density + Density scatter inputs for energy. + force_density + Density scatter inputs for forces. + stress_density + Density scatter inputs for stress. + """ + assert metrics is not None + assert energy_density is not None + assert force_density is not None + assert stress_density is not None + + 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..01ffae16a --- /dev/null +++ b/ml_peg/app/actinides/plutonium_dioxide/app_plutonium_dioxide.py @@ -0,0 +1,87 @@ +"""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 +from ml_peg.app.utils.load import 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, + ) + + +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")], + ) + + +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..8a7f66fed --- /dev/null +++ b/ml_peg/calcs/actinides/plutonium_dioxide/calc_plutonium_dioxide.py @@ -0,0 +1,66 @@ +"""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 + + # puo2_data_dir = ( + # Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")) + # / "puo2_data") + + # Download data. + puo2_data_dir = ( + download_s3_data( + key="inputs/actinides/isomer_complexes/puo2.zip", + filename="puo2_data.zip", + ) + / "puo2_data" + ) + + 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) From 01a7b0bb6fa35f1b36ecc3e9cd3b82e171b435ce Mon Sep 17 00:00:00 2001 From: DrDavie1 <101512473+DrDavie1@users.noreply.github.com> Date: Sun, 1 Mar 2026 17:21:38 +0000 Subject: [PATCH 2/6] edit --- .../calcs/actinides/plutonium_dioxide/calc_plutonium_dioxide.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ml_peg/calcs/actinides/plutonium_dioxide/calc_plutonium_dioxide.py b/ml_peg/calcs/actinides/plutonium_dioxide/calc_plutonium_dioxide.py index 8a7f66fed..28c862596 100644 --- a/ml_peg/calcs/actinides/plutonium_dioxide/calc_plutonium_dioxide.py +++ b/ml_peg/calcs/actinides/plutonium_dioxide/calc_plutonium_dioxide.py @@ -40,7 +40,7 @@ def test_puo2_parity(mlip: tuple[str, Any]) -> None: # Download data. puo2_data_dir = ( download_s3_data( - key="inputs/actinides/isomer_complexes/puo2.zip", + key="inputs/actinides/plutonium_dioxide/puo2.zip", filename="puo2_data.zip", ) / "puo2_data" From a941f0b7090999c5943d2e790c4a0933660df3c9 Mon Sep 17 00:00:00 2001 From: DrDavie1 <101512473+DrDavie1@users.noreply.github.com> Date: Fri, 13 Mar 2026 12:06:53 +0000 Subject: [PATCH 3/6] s3download --- .../plutonium_dioxide/calc_plutonium_dioxide.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/ml_peg/calcs/actinides/plutonium_dioxide/calc_plutonium_dioxide.py b/ml_peg/calcs/actinides/plutonium_dioxide/calc_plutonium_dioxide.py index 28c862596..9ae55db70 100644 --- a/ml_peg/calcs/actinides/plutonium_dioxide/calc_plutonium_dioxide.py +++ b/ml_peg/calcs/actinides/plutonium_dioxide/calc_plutonium_dioxide.py @@ -33,17 +33,13 @@ def test_puo2_parity(mlip: tuple[str, Any]) -> None: """ model_name, model = mlip - # puo2_data_dir = ( - # Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")) - # / "puo2_data") - # Download data. puo2_data_dir = ( download_s3_data( - key="inputs/actinides/plutonium_dioxide/puo2.zip", - filename="puo2_data.zip", + key="inputs/actinides/plutonium_dioxide/plutonium_dioxide.zip", + filename="plutonium_dioxide.zip", ) - / "puo2_data" + / "plutonium_dioxide" ) ref_file = puo2_data_dir / "dft_ref_data.xyz" From a4474e3b55ad284bdbf91beea37a5e44ff4b1492 Mon Sep 17 00:00:00 2001 From: joehart2001 Date: Mon, 16 Mar 2026 17:36:34 +0000 Subject: [PATCH 4/6] add density to structure visualisation for energy, forces and stresses --- .../analyse_plutonium_dioxide.py | 70 ++++++++++++++++--- .../app_plutonium_dioxide.py | 30 +++++++- 2 files changed, 89 insertions(+), 11 deletions(-) diff --git a/ml_peg/analysis/actinides/plutonium_dioxide/analyse_plutonium_dioxide.py b/ml_peg/analysis/actinides/plutonium_dioxide/analyse_plutonium_dioxide.py index 8d52bbee4..a8b572b3f 100644 --- a/ml_peg/analysis/actinides/plutonium_dioxide/analyse_plutonium_dioxide.py +++ b/ml_peg/analysis/actinides/plutonium_dioxide/analyse_plutonium_dioxide.py @@ -10,7 +10,12 @@ 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 +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 @@ -36,7 +41,7 @@ def puo2_stats() -> dict[str, dict[str, Any]]: Returns ------- dict[str, dict[str, Any]] - Processed information per model (energy, force, stress). + Processed information per model (energy, force, stress, labels). """ OUT_PATH.mkdir(parents=True, exist_ok=True) stats: dict[str, dict[str, Any]] = {} @@ -46,30 +51,46 @@ def puo2_stats() -> dict[str, dict[str, Any]]: 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: - stress_ref.extend(s_ref.tolist()) + 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": { @@ -77,14 +98,17 @@ def puo2_stats() -> dict[str, dict[str, Any]]: "pred": energies_pred, }, "forces": { - "ref": np.concatenate(forces_ref).tolist(), - "pred": np.concatenate(forces_pred).tolist(), + "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 @@ -180,12 +204,22 @@ def energy_density(puo2_stats: dict[str, dict[str, Any]]) -> dict[str, dict]: dict[str, dict] Mapping of model name to density-scatter data. """ - return build_density_inputs( + 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 @@ -210,12 +244,22 @@ def force_density(puo2_stats: dict[str, dict[str, Any]]) -> dict[str, dict]: dict[str, dict] Mapping of model name to density-scatter data. """ - return build_density_inputs( + 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 @@ -240,12 +284,22 @@ def stress_density(puo2_stats: dict[str, dict[str, Any]]) -> dict[str, dict]: dict[str, dict] Mapping of model name to density-scatter data. """ - return build_density_inputs( + 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 diff --git a/ml_peg/app/actinides/plutonium_dioxide/app_plutonium_dioxide.py b/ml_peg/app/actinides/plutonium_dioxide/app_plutonium_dioxide.py index 01ffae16a..8413f30dd 100644 --- a/ml_peg/app/actinides/plutonium_dioxide/app_plutonium_dioxide.py +++ b/ml_peg/app/actinides/plutonium_dioxide/app_plutonium_dioxide.py @@ -7,8 +7,8 @@ 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 -from ml_peg.app.utils.load import read_density_plot_for_model +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 @@ -60,6 +60,27 @@ def register_callbacks(self) -> None: 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: """ @@ -75,7 +96,10 @@ def get_app() -> PuO2App: 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")], + extra_components=[ + Div(id=f"{BENCHMARK_NAME}-figure-placeholder"), + Div(id=f"{BENCHMARK_NAME}-struct-placeholder"), + ], ) From 010c4cb58ff2642794ee72d4e739330990be80cb Mon Sep 17 00:00:00 2001 From: Joseph Hart <92541539+joehart2001@users.noreply.github.com> Date: Mon, 16 Mar 2026 21:06:39 +0000 Subject: [PATCH 5/6] Apply suggestion from @joehart2001 --- .../analyse_plutonium_dioxide.py | 18 +----------------- 1 file changed, 1 insertion(+), 17 deletions(-) diff --git a/ml_peg/analysis/actinides/plutonium_dioxide/analyse_plutonium_dioxide.py b/ml_peg/analysis/actinides/plutonium_dioxide/analyse_plutonium_dioxide.py index a8b572b3f..30f77a051 100644 --- a/ml_peg/analysis/actinides/plutonium_dioxide/analyse_plutonium_dioxide.py +++ b/ml_peg/analysis/actinides/plutonium_dioxide/analyse_plutonium_dioxide.py @@ -338,12 +338,7 @@ def metrics( } -def test_puo2( - metrics: dict[str, dict], - energy_density: dict[str, dict], - force_density: dict[str, dict], - stress_density: dict[str, dict], -) -> None: +def test_puo2(metrics: dict[str, dict]) -> None: """ Run puo2 analysis. @@ -351,16 +346,5 @@ def test_puo2( ---------- metrics Benchmark metric values. - energy_density - Density scatter inputs for energy. - force_density - Density scatter inputs for forces. - stress_density - Density scatter inputs for stress. """ - assert metrics is not None - assert energy_density is not None - assert force_density is not None - assert stress_density is not None - return From c69a44248f07a19ee6a471456dad6328e9a36e55 Mon Sep 17 00:00:00 2001 From: DrDavie1 <101512473+DrDavie1@users.noreply.github.com> Date: Mon, 23 Mar 2026 16:45:12 +0000 Subject: [PATCH 6/6] update docs --- docs/source/user_guide/benchmarks/actinides.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/user_guide/benchmarks/actinides.rst b/docs/source/user_guide/benchmarks/actinides.rst index 4df7b2cb1..777ffaf01 100644 --- a/docs/source/user_guide/benchmarks/actinides.rst +++ b/docs/source/user_guide/benchmarks/actinides.rst @@ -33,4 +33,4 @@ Low Data availability ----------------- -Reference data: availabile in repo, for specific calculation details contact willdavie2002@gmail.com. +Reference data: availabile in repo. Data and complete calculation details will be released in an upcoming publication. For now, please contact willdavie2002@gmail.com.