feat: water rdf benchmark

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.

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.
+    """

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

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)

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"