feat: add md stability benchmark

docs/source/user_guide/benchmarks/molecular_dynamics.rst

@@ -76,3 +76,45 @@ Reference data:
 
 * Same as input data
 * Experimental
+
+
+Stability
+=========
+
+Summary
+-------
+
+Performance in running stable molecular dynamics across a diverse set of
+systems: small molecules (containing H/C/N/O, sulfur, and halogens), peptides
+in vacuum (neurotensin, PDB: 2LNF; cyclic oxytocin, PDB: 7OFG), a protein in
+vacuum (PDB: 1A7M), and solvated peptides with and without counter-ions.
+
+Metrics
+-------
+
+1. Success rate
+
+Each system is run with a short NVT molecular dynamics simulation at 300 K. The
+fraction of simulations that complete without error is reported. A simulation
+counts as failed if it raises an error during integration (for example, a
+numerical instability). Higher is better, with 1.0 meaning every system ran to
+completion.
+
+A scatter plot below the table shows, for each system and model, the fraction of
+the trajectory completed before it became unstable: 1.0 for a fully stable run, a
+smaller value when the run exploded or hydrogens drifted, and 0.0 for systems that
+failed to run at all. Individual models can be toggled via the legend.
+
+Computational cost
+------------------
+
+High: tests are likely to take many 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.
+  Structures derived from PDB entries 2LNF, 7OFG, and 1A7M.

ml_peg/analysis/molecular_dynamics/stability/analyse_stability.py

@@ -0,0 +1,205 @@
+"""Analyse molecular dynamics stability benchmark."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from ase.calculators.calculator import Calculator
+from mlipaudit.benchmarks.stability.stability import STRUCTURE_NAMES, STRUCTURES
+from mlipaudit.io import load_model_output_from_disk
+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 MlPegStabilityBenchmark
+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)
+
+CALC_PATH = CALCS_ROOT / "molecular_dynamics" / "stability" / "outputs"
+OUT_PATH = APP_ROOT / "data" / "molecular_dynamics" / "stability"
+
+METRICS_CONFIG_PATH = Path(__file__).with_name("metrics.yml")
+DEFAULT_THRESHOLDS, DEFAULT_TOOLTIPS, DEFAULT_WEIGHTS = load_metrics_config(
+    METRICS_CONFIG_PATH
+)
+
+
+def _fraction_completed(result) -> float:
+    """
+    Get how far a simulation got before it became unstable.
+
+    Parameters
+    ----------
+    result
+        A ``StabilityStructureResult`` for a single structure.
+
+    Returns
+    -------
+    float
+        Fraction of the trajectory completed: 1.0 if stable, 0.0 if the
+        simulation did not run, otherwise the frame at which it exploded or
+        drifted divided by the total number of frames.
+    """
+    if result.failed:
+        return 0.0
+    if result.exploded_frame != -1:
+        return result.exploded_frame / result.num_frames
+    if result.drift_frame != -1:
+        return result.drift_frame / result.num_frames
+    return 1.0
+
+
+@pytest.fixture
+def structure_results() -> dict[str, list]:
+    """
+    Analyse each stored trajectory into per-structure stability results.
+
+    Returns
+    -------
+    dict[str, list]
+        List of ``StabilityStructureResult`` objects for each model.
+    """
+    data_input_dir = download_s3_data(
+        key="inputs/molecular_dynamics/stability/stability.zip",
+        filename="stability.zip",
+    )
+
+    results = {}
+    for model_name in MODELS:
+        output_dir = CALC_PATH / model_name / MlPegStabilityBenchmark.name
+        if not (output_dir / "model_output.zip").exists():
+            continue
+        benchmark = MlPegStabilityBenchmark(
+            force_field=Calculator(),
+            data_input_dir=data_input_dir,
+            run_mode="standard",
+        )
+        benchmark.model_output = load_model_output_from_disk(
+            CALC_PATH / model_name, MlPegStabilityBenchmark
+        )
+        results[model_name] = benchmark.analyze().structure_results
+    return results
+
+
+@pytest.fixture
+def success_rate(structure_results: dict[str, list]) -> dict[str, float]:
+    """
+    Get the fraction of MD simulations that completed without error.
+
+    Parameters
+    ----------
+    structure_results
+        Per-structure stability results for all models.
+
+    Returns
+    -------
+    dict[str, float]
+        Fraction of successful simulations for each model.
+    """
+    return {
+        model_name: sum(not r.failed for r in results) / len(results)
+        for model_name, results in structure_results.items()
+    }
+
+
+@pytest.fixture
+@plot_scatter(
+    title="Trajectory stability",
+    x_label="System",
+    y_label="Fraction of trajectory completed",
+    hovertemplate="<b>%{x}</b><br>Completed: %{y:.2f}<extra>%{fullData.name}</extra>",
+    filename=str(OUT_PATH / "stability_progress.json"),
+)
+def progress(structure_results: dict[str, list]) -> dict[str, tuple[list, list]]:
+    """
+    Get per-structure trajectory progress for each model.
+
+    For every system, reports the fraction of the trajectory completed before it
+    exploded, drifted or failed to run.
+
+    Parameters
+    ----------
+    structure_results
+        Per-structure stability results for all models.
+
+    Returns
+    -------
+    dict[str, tuple[list, list]]
+        Structure names and completed fractions for each model.
+    """
+    return {
+        model_name: (
+            [r.structure_name for r in results],
+            [_fraction_completed(r) for r in results],
+        )
+        for model_name, results in structure_results.items()
+    }
+
+
+@pytest.fixture
+@build_table(
+    filename=OUT_PATH / "stability_metrics_table.json",
+    metric_tooltips=DEFAULT_TOOLTIPS,
+    thresholds=DEFAULT_THRESHOLDS,
+    mlip_name_map=DISPERSION_NAME_MAP,
+)
+def metrics(success_rate: dict[str, float]) -> dict[str, dict]:
+    """
+    Get all metrics.
+
+    Parameters
+    ----------
+    success_rate
+        Fraction of successful simulations for all models.
+
+    Returns
+    -------
+    dict[str, dict]
+        Metric names and values for all models.
+    """
+    return {
+        "Success Rate": success_rate,
+    }
+
+
+@pytest.fixture
+def element_info() -> None:
+    """Write element info for all benchmark systems, used by the app element filter."""
+    data_input_dir = download_s3_data(
+        key="inputs/molecular_dynamics/stability/stability.zip",
+        filename="stability.zip",
+    )
+    xyz_paths = [
+        data_input_dir / MlPegStabilityBenchmark.name / STRUCTURES[name]["xyz"]
+        for name in STRUCTURE_NAMES
+    ]
+    write_struct_info(data_path=xyz_paths, out_path=OUT_PATH)
+
+
+def test_stability(
+    metrics: dict[str, dict],
+    progress: dict[str, tuple[list, list]],
+    element_info: None,
+) -> None:
+    """
+    Run stability analysis.
+
+    Parameters
+    ----------
+    metrics : dict[str, dict]
+        Stability metric results provided by fixtures.
+    progress : dict[str, tuple[list, list]]
+        Per-structure trajectory progress provided by fixtures.
+    element_info : None
+        Element info written for the app element filter.
+    """

ml_peg/analysis/molecular_dynamics/stability/metrics.yml

@@ -0,0 +1,6 @@
+metrics:
+  Success Rate:
+    good: 1.0
+    bad: 0.0
+    unit: null
+    tooltip: Fraction of MD simulations that completed without error.

ml_peg/analysis/utils/decorators.py

@@ -448,6 +448,7 @@ def plot_scatter(
     show_line: bool = False,
     show_markers: bool = True,
     hoverdata: dict | None = None,
+    hovertemplate: str | None = None,
     filename: str = "scatter.json",
     highlight_range: dict = None,
 ) -> Callable:
@@ -468,6 +469,9 @@ def plot_scatter(
         Whether to show markers on the plot. Default is True.
     hoverdata
         Hover data dictionary. Default is `{}`.
+    hovertemplate
+        Base hovertemplate string. Overrides the default reference/prediction
+        template. Any ``hoverdata`` keys are appended to it. Default is `None`.
     filename
         Filename to save plot as JSON. Default is "scatter.json".
     highlight_range
@@ -513,11 +517,15 @@ def plot_scatter_wrapper(*args, **kwargs) -> dict[str, Any]:
             """
             results = func(*args, **kwargs)
 
-            hovertemplate = "<b>Pred: </b>%{x}<br>" + "<b>Ref: </b>%{y}<br>"
+            template = (
+                hovertemplate
+                if hovertemplate is not None
+                else "<b>Pred: </b>%{x}<br>" + "<b>Ref: </b>%{y}<br>"
+            )
             customdata = []
             if hoverdata:
                 for i, key in enumerate(hoverdata):
-                    hovertemplate += f"<b>{key}: </b>%{{customdata[{i}]}}<br>"
+                    template += f"<b>{key}: </b>%{{customdata[{i}]}}<br>"
                 customdata = list(zip(*hoverdata.values(), strict=True))
 
             modes = []
@@ -538,7 +546,7 @@ def plot_scatter_wrapper(*args, **kwargs) -> dict[str, Any]:
                         name=name,
                         mode=mode,
                         customdata=customdata,
-                        hovertemplate=hovertemplate,
+                        hovertemplate=template,
                     )
                 )
 

ml_peg/app/molecular_dynamics/molecular_dynamics.yml

@@ -1,3 +1,3 @@
 title: Molecular Dynamics
-description: Liquid Densities, Water Densities
+description: Liquid Densities, Water Densities, Stability
 weight: 0

ml_peg/app/molecular_dynamics/stability/app_stability.py

@@ -0,0 +1,58 @@
+"""Run molecular dynamics stability app."""
+
+from __future__ import annotations
+
+from dash import Dash
+
+from ml_peg.app import APP_ROOT
+from ml_peg.app.base_app import BaseApp
+from ml_peg.app.utils.load import read_plot
+
+BENCHMARK_NAME = "Stability"
+DOCS_URL = "https://ddmms.github.io/ml-peg/user_guide/benchmarks/molecular_dynamics.html#stability"
+DATA_PATH = APP_ROOT / "data" / "molecular_dynamics" / "stability"
+INFO_PATH = DATA_PATH / "info.json"
+
+
+class StabilityApp(BaseApp):
+    """Stability benchmark app layout and callbacks."""
+
+    def register_callbacks(self) -> None:
+        """Register callbacks to app."""
+
+
+def get_app() -> StabilityApp:
+    """
+    Get stability benchmark app layout and callback registration.
+
+    Returns
+    -------
+    StabilityApp
+        Benchmark layout and callback registration.
+    """
+    return StabilityApp(
+        name=BENCHMARK_NAME,
+        framework_ids="mlip_audit",
+        description=(
+            "Fraction of short molecular dynamics simulations that complete "
+            "without error, across small molecules, peptides and proteins in "
+            "vacuum and solvent."
+        ),
+        docs_url=DOCS_URL,
+        table_path=DATA_PATH / "stability_metrics_table.json",
+        extra_components=[
+            read_plot(
+                filename=DATA_PATH / "stability_progress.json",
+                id=f"{BENCHMARK_NAME}-progress-figure",
+            )
+        ],
+        info_path=INFO_PATH,
+    )
+
+
+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=8067, debug=True)

ml_peg/app/utils/frameworks.yml

@@ -26,3 +26,10 @@ mace-polar-1:
   url: "https://arxiv.org/abs/2602.19411"
   icon: "📄"
   tooltip: "Benchmark from the MACE-POLAR-1 paper (arXiv:2602.19411) - click to read"
+
+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"