From d36de082ff941ecfabaa9bf10549d86c7ecf2679 Mon Sep 17 00:00:00 2001 From: Leon Wehrhan Date: Thu, 2 Jul 2026 18:30:47 +0200 Subject: [PATCH] feat: water rdf benchmark --- .../benchmarks/molecular_dynamics.rst | 46 ++++ .../analyse_water_radial_distribution.py | 212 ++++++++++++++++++ .../water_radial_distribution/metrics.yml | 15 ++ .../app_water_radial_distribution.py | 66 ++++++ ml_peg/app/utils/frameworks.yml | 7 + .../calc_water_radial_distribution.py | 65 ++++++ ml_peg/calcs/utils/mlipaudit.py | 18 ++ pyproject.toml | 8 + 8 files changed, 437 insertions(+) create mode 100644 ml_peg/analysis/molecular_dynamics/water_radial_distribution/analyse_water_radial_distribution.py create mode 100644 ml_peg/analysis/molecular_dynamics/water_radial_distribution/metrics.yml create mode 100644 ml_peg/app/molecular_dynamics/water_radial_distribution/app_water_radial_distribution.py create mode 100644 ml_peg/calcs/molecular_dynamics/water_radial_distribution/calc_water_radial_distribution.py create mode 100644 ml_peg/calcs/utils/mlipaudit.py diff --git a/docs/source/user_guide/benchmarks/molecular_dynamics.rst b/docs/source/user_guide/benchmarks/molecular_dynamics.rst index 24f3fb3b3..6c2958029 100644 --- a/docs/source/user_guide/benchmarks/molecular_dynamics.rst +++ b/docs/source/user_guide/benchmarks/molecular_dynamics.rst @@ -76,3 +76,49 @@ Reference data: * Same as input data * Experimental + + +Water radial distribution +========================= + +Summary +------- + +Performance in reproducing the oxygen-oxygen radial distribution function of +liquid water. A short NVT molecular dynamics simulation of a box of 500 water +molecules is run from an equilibrated structure, and the resulting O-O RDF is +compared to the experimental reference. + +Metrics +------- + +1. Peak deviation + +The position of the first solvent peak (the radius at which the RDF is maximal) +is compared to the experimental peak of 2.8. + +2. RDF RMSE + +The root mean square error of the radial distribution function against the +experimental reference, evaluated over the range 2.5-10.0 Å. + +A plot shows the predicted RDF profile of each model against the experimental +reference profile. + +Computational cost +------------------ + +High: tests are likely to take several hours on GPU. Faster simulation times can be +achieved using the jax accelerated simulations in MLIP Audit directly. + +Data availability +----------------- + +Input structures: + +* MLIP Audit benchmark suite, InstaDeep. Equilibrated box of 500 water + molecules. + +Reference data: + +* Experimental oxygen-oxygen radial distribution function. diff --git a/ml_peg/analysis/molecular_dynamics/water_radial_distribution/analyse_water_radial_distribution.py b/ml_peg/analysis/molecular_dynamics/water_radial_distribution/analyse_water_radial_distribution.py new file mode 100644 index 000000000..0d6a92f53 --- /dev/null +++ b/ml_peg/analysis/molecular_dynamics/water_radial_distribution/analyse_water_radial_distribution.py @@ -0,0 +1,212 @@ +"""Analyse the water oxygen-oxygen radial distribution benchmark.""" + +from __future__ import annotations + +from pathlib import Path + +from ase.calculators.calculator import Calculator +from mlipaudit.io import load_model_output_from_disk +import numpy as np +import pytest + +from ml_peg.analysis.utils.decorators import build_table, plot_scatter +from ml_peg.analysis.utils.utils import ( + build_dispersion_name_map, + load_metrics_config, + write_struct_info, +) +from ml_peg.app import APP_ROOT +from ml_peg.calcs import CALCS_ROOT +from ml_peg.calcs.utils.mlipaudit import MlPegWaterRadialDistributionBenchmark +from ml_peg.calcs.utils.utils import download_s3_data +from ml_peg.models import current_models +from ml_peg.models.get_models import load_models + +MODELS = load_models(current_models) +DISPERSION_NAME_MAP = build_dispersion_name_map(MODELS) + +BENCHMARK = MlPegWaterRadialDistributionBenchmark.name +WATERBOX_N500 = "water_box_n500_eq.pdb" +REFERENCE_DATA = "experimental_reference.npz" + +CALC_PATH = CALCS_ROOT / "molecular_dynamics" / "water_radial_distribution" / "outputs" +OUT_PATH = APP_ROOT / "data" / "molecular_dynamics" / "water_radial_distribution" + +METRICS_CONFIG_PATH = Path(__file__).with_name("metrics.yml") +DEFAULT_THRESHOLDS, DEFAULT_TOOLTIPS, DEFAULT_WEIGHTS = load_metrics_config( + METRICS_CONFIG_PATH +) + + +def _data_input_dir() -> Path: + """ + Download and return the benchmark input data directory. + + Returns + ------- + Path + Directory containing the extracted water RDF input data. + """ + return download_s3_data( + key="inputs/molecular_dynamics/water_radial_distribution/water_radial_distribution.zip", + filename="water_radial_distribution.zip", + ) + + +@pytest.fixture +def analyze_results() -> dict: + """ + Run the mlipaudit analysis for each model. + + Returns + ------- + dict + Mapping of model name to its ``WaterRadialDistributionResult``. + """ + data_input_dir = _data_input_dir() + + results = {} + for model_name in MODELS: + output_dir = CALC_PATH / model_name / BENCHMARK + if not (output_dir / "model_output.zip").exists(): + continue + benchmark = MlPegWaterRadialDistributionBenchmark( + force_field=Calculator(), + data_input_dir=data_input_dir, + run_mode="standard", + ) + benchmark.model_output = load_model_output_from_disk( + CALC_PATH / model_name, MlPegWaterRadialDistributionBenchmark + ) + results[model_name] = benchmark.analyze() + return results + + +@pytest.fixture +def struct_info() -> None: + """Write the combined element set to ``info.json`` for filtering.""" + data_input_dir = _data_input_dir() + write_struct_info( + data_path=data_input_dir / BENCHMARK / WATERBOX_N500, + out_path=OUT_PATH, + ) + + +@pytest.fixture +@plot_scatter( + title="Water O-O radial distribution function", + x_label="r / Å", + y_label="g(r)", + show_line=True, + show_markers=False, + filename=str(OUT_PATH / "figure_rdf.json"), +) +def rdf_profiles(analyze_results) -> dict[str, tuple[list, list]]: + """ + Get predicted and reference O-O radial distribution profiles. + + Parameters + ---------- + analyze_results + Mapping of model name to its ``WaterRadialDistributionResult``. + + Returns + ------- + dict[str, tuple[list, list]] + Reference and per-model ``(radii, g(r))`` profiles. + """ + reference = np.load(_data_input_dir() / BENCHMARK / REFERENCE_DATA) + results = {"ref": (reference["r_OO"].tolist(), reference["g_OO"].tolist())} + for model_name, result in analyze_results.items(): + if result.failed: + continue + results[model_name] = (result.radii, result.rdf) + return results + + +@pytest.fixture +def get_peak_deviation(analyze_results) -> dict[str, float]: + """ + Get the first solvent peak deviation for each model. + + Parameters + ---------- + analyze_results + Mapping of model name to its ``WaterRadialDistributionResult``. + + Returns + ------- + dict[str, float] + Deviation of the first solvent peak from the reference range in Angstrom. + """ + return { + model_name: result.peak_deviation + for model_name, result in analyze_results.items() + } + + +@pytest.fixture +def get_rmse(analyze_results) -> dict[str, float]: + """ + Get the RDF profile RMSE for each model. + + Parameters + ---------- + analyze_results + Mapping of model name to its ``WaterRadialDistributionResult``. + + Returns + ------- + dict[str, float] + RMSE of the radial distribution function against the reference. + """ + return {model_name: result.rmse for model_name, result in analyze_results.items()} + + +@pytest.fixture +@build_table( + filename=OUT_PATH / "water_radial_distribution_metrics_table.json", + metric_tooltips=DEFAULT_TOOLTIPS, + thresholds=DEFAULT_THRESHOLDS, + weights=DEFAULT_WEIGHTS, + mlip_name_map=DISPERSION_NAME_MAP, +) +def metrics( + rdf_profiles, + get_peak_deviation: dict[str, float], + get_rmse: dict[str, float], +) -> dict[str, dict]: + """ + Get all metrics. + + Parameters + ---------- + rdf_profiles + Reference and predicted RDF profiles (triggers the RDF plot). + get_peak_deviation + First solvent peak deviations for all models. + get_rmse + RDF profile RMSEs for all models. + + Returns + ------- + dict[str, dict] + Metric names and values for all models. + """ + return { + "Peak Deviation": get_peak_deviation, + "RDF RMSE": get_rmse, + } + + +def test_water_radial_distribution(metrics: dict[str, dict], struct_info: None) -> None: + """ + Run water radial distribution analysis. + + Parameters + ---------- + metrics : dict[str, dict] + Water RDF metric results provided by fixtures. + struct_info : None + Element info written to ``info.json`` for filtering. + """ diff --git a/ml_peg/analysis/molecular_dynamics/water_radial_distribution/metrics.yml b/ml_peg/analysis/molecular_dynamics/water_radial_distribution/metrics.yml new file mode 100644 index 000000000..eee2f72ae --- /dev/null +++ b/ml_peg/analysis/molecular_dynamics/water_radial_distribution/metrics.yml @@ -0,0 +1,15 @@ +metrics: + Peak Deviation: + good: 0.0 + bad: 0.25 + unit: Å + weight: 1 + tooltip: Deviation of the first solvent peak position from the experimental peak at 2.8Å. + level_of_theory: Experiment + RDF RMSE: + good: 0.0 + bad: 0.5 + unit: null + weight: 1 + tooltip: Root mean square error of the O-O radial distribution function against the experimental reference over 2.5-10.0 Å. + level_of_theory: Experiment diff --git a/ml_peg/app/molecular_dynamics/water_radial_distribution/app_water_radial_distribution.py b/ml_peg/app/molecular_dynamics/water_radial_distribution/app_water_radial_distribution.py new file mode 100644 index 000000000..e7fe3f8fa --- /dev/null +++ b/ml_peg/app/molecular_dynamics/water_radial_distribution/app_water_radial_distribution.py @@ -0,0 +1,66 @@ +"""Run water radial distribution 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_column +from ml_peg.app.utils.load import read_plot + +BENCHMARK_NAME = "WaterRDF" +DOCS_URL = "https://ddmms.github.io/ml-peg/user_guide/benchmarks/molecular_dynamics.html#water-radial-distribution" +DATA_PATH = APP_ROOT / "data" / "molecular_dynamics" / "water_radial_distribution" + + +class WaterRDFApp(BaseApp): + """Water radial distribution benchmark app layout and callbacks.""" + + def register_callbacks(self) -> None: + """Register callbacks to app.""" + scatter = read_plot( + DATA_PATH / "figure_rdf.json", + id=f"{BENCHMARK_NAME}-figure", + ) + + plot_from_table_column( + table_id=self.table_id, + plot_id=f"{BENCHMARK_NAME}-figure-placeholder", + column_to_plot={"RDF RMSE": scatter}, + ) + + +def get_app() -> WaterRDFApp: + """ + Get water radial distribution benchmark app layout and callback registration. + + Returns + ------- + WaterRDFApp + Benchmark layout and callback registration. + """ + return WaterRDFApp( + name="Water RDF", + framework_ids="mlip_audit", + description=( + "Performance in reproducing the oxygen-oxygen radial distribution " + "function of liquid water from a short NVT molecular dynamics " + "simulation. Reference data from experiment." + ), + docs_url=DOCS_URL, + table_path=DATA_PATH / "water_radial_distribution_metrics_table.json", + info_path=DATA_PATH / "info.json", + extra_components=[ + Div(id=f"{BENCHMARK_NAME}-figure-placeholder"), + ], + ) + + +if __name__ == "__main__": + full_app = Dash(__name__, assets_folder=DATA_PATH.parent.parent) + benchmark_app = get_app() + full_app.layout = benchmark_app.layout + benchmark_app.register_callbacks() + full_app.run(port=8070, debug=True) diff --git a/ml_peg/app/utils/frameworks.yml b/ml_peg/app/utils/frameworks.yml index 27b079925..4bb28f12a 100644 --- a/ml_peg/app/utils/frameworks.yml +++ b/ml_peg/app/utils/frameworks.yml @@ -11,6 +11,13 @@ mlip_arena: url: "https://huggingface.co/spaces/atomind/mlip-arena" logo: "https://huggingface.co/front/assets/huggingface_logo-noborder.svg" +mlip_audit: + label: MLIP Audit + color: "#1d4ed8" + text_color: "#ffffff" + url: "https://github.com/instadeepai/mlipaudit" + logo: "https://raw.githubusercontent.com/instadeepai/mlipaudit/mlpeg-migration/InstaDeep_Logo.png" + mace-multihead: label: Multihead Cross Learning color: "#7c3aed" diff --git a/ml_peg/calcs/molecular_dynamics/water_radial_distribution/calc_water_radial_distribution.py b/ml_peg/calcs/molecular_dynamics/water_radial_distribution/calc_water_radial_distribution.py new file mode 100644 index 000000000..301fea6dd --- /dev/null +++ b/ml_peg/calcs/molecular_dynamics/water_radial_distribution/calc_water_radial_distribution.py @@ -0,0 +1,65 @@ +""" +Compute the water oxygen-oxygen radial distribution function. + +A short NVT molecular dynamics simulation of a box of 500 water molecules is +run, and the O-O radial distribution function is compared to experiment. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any +from warnings import warn + +from mlipaudit.benchmarks.water_radial_distribution.water_radial_distribution import ( + WaterRadialDistributionModelOutput, +) +from mlipaudit.io import write_model_output_to_disk +import pytest + +from ml_peg.calcs.utils.mlipaudit import MlPegWaterRadialDistributionBenchmark +from ml_peg.calcs.utils.utils import download_s3_data +from ml_peg.models import current_models +from ml_peg.models.get_models import load_models + +MODELS = load_models(current_models) + +OUT_PATH = Path(__file__).parent / "outputs" + + +@pytest.mark.parametrize("mlip", MODELS.items()) +def test_water_radial_distribution(mlip: tuple[str, Any]) -> None: + """ + Benchmark the water radial distribution function. + + Parameters + ---------- + mlip + Name of model and model object to get calculator. + """ + model_name, model = mlip + calc = model.get_calculator() + calc = model.add_d3_calculator(calc) + + data_input_dir = download_s3_data( + key="inputs/molecular_dynamics/water_radial_distribution/water_radial_distribution.zip", + filename="water_radial_distribution.zip", + ) + + benchmark = MlPegWaterRadialDistributionBenchmark( + force_field=calc, + data_input_dir=data_input_dir, + run_mode="standard", + ) + try: + benchmark.run_model() + except Exception as exc: + warn( + f"Error running water RDF benchmark for {model_name}: {exc}", + stacklevel=2, + ) + benchmark.model_output = WaterRadialDistributionModelOutput(failed=True) + + write_model_output_to_disk( + "water_radial_distribution", benchmark.model_output, OUT_PATH / model_name + ) diff --git a/ml_peg/calcs/utils/mlipaudit.py b/ml_peg/calcs/utils/mlipaudit.py new file mode 100644 index 000000000..b99fdddbb --- /dev/null +++ b/ml_peg/calcs/utils/mlipaudit.py @@ -0,0 +1,18 @@ +"""Adapters for using mlipaudit benchmarks with ml-peg's ASE calculators.""" + +from __future__ import annotations + +from mlipaudit.benchmarks.water_radial_distribution.water_radial_distribution import ( + WaterRadialDistributionBenchmark, +) + + +class MlPegWaterRadialDistributionBenchmark(WaterRadialDistributionBenchmark): + """ + WaterRadialDistributionBenchmark wired up for ml-peg's ASE calculators. + + ``skip_if_elements_missing`` is disabled because ASE ``Calculator`` objects + do not expose ``allowed_atomic_numbers``. + """ + + skip_if_elements_missing = False diff --git a/pyproject.toml b/pyproject.toml index 23088ff82..231965742 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -64,6 +64,9 @@ mace = [ mattersim = [ "mattersim==1.2.2", ] +mlipaudit = [ + "mlipaudit; python_version >= '3.11'", +] orb = [ "orb-models == 0.6.2; sys_platform != 'win32' and python_version >= '3.12'", ] @@ -206,6 +209,10 @@ conflicts = [ { extra = "mattersim" }, { extra = "grace" }, ], + [ + { extra = "mlipaudit" }, + { extra = "grace" }, + ], ] constraint-dependencies = [ @@ -217,3 +224,4 @@ module-root = "" [tool.uv.sources] asemolec = { git = "https://github.com/imagdau/aseMolec.git" } +mlipaudit = { git = "https://github.com/instadeepai/MLIPAudit.git", branch = "mlpeg-migration" }