feat: add bond length distribution benchmark

docs/source/user_guide/benchmarks/molecular_dynamics.rst

@@ -76,3 +76,46 @@ Reference data:
 
 * Same as input data
 * Experimental
+
+
+Bond length distribution
+========================
+
+Summary
+-------
+
+Performance in maintaining physically reasonable covalent bond lengths during molecular
+dynamics of small organic molecules. For each of a set of molecules covering the C-C, C=C,
+C#C, C-N, C-O, C=O and C-F bond types, an NVT molecular dynamics simulation is run at 300 K
+starting from a QM-optimised reference geometry, and the deviation of a tracked bond from
+its reference length is measured along the trajectory.
+
+Metrics
+-------
+
+1. Bond length deviation
+
+The length of the tracked bond is measured at each frame of the trajectory, and its absolute
+deviation from the reference bond length is averaged over the trajectory and across all
+molecules. A well behaved potential keeps bonds close to their reference length, so a lower
+deviation is better.
+
+A histogram shows the distribution of the sampled bond length deviations for each model.
+
+Computational cost
+------------------
+
+High: one MD simulation per molecule, each 1,000,000 steps. Faster inference can be achieved
+using the jax-accelerated simulations in MLIP Audit directly.
+
+Data availability
+-----------------
+
+Input structures:
+
+* MLIP Audit benchmark suite, InstaDeep. Reference geometries selected from the QM9 dataset
+  (Ramakrishnan et al., Scientific Data 1, 140022, 2014).
+
+Reference data:
+
+* QM-optimised equilibrium bond lengths of the reference geometries.

ml_peg/analysis/molecular_dynamics/bond_length_distribution/analyse_bond_length_distribution.py

@@ -0,0 +1,200 @@
+"""Analyse the covalent bond length distribution benchmark."""
+
+from __future__ import annotations
+
+import json
+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_hist
+from ml_peg.analysis.utils.utils import (
+    build_dispersion_name_map,
+    load_metrics_config,
+)
+from ml_peg.app import APP_ROOT
+from ml_peg.calcs import CALCS_ROOT
+from ml_peg.calcs.utils.mlipaudit import MlPegBondLengthDistributionBenchmark
+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 = MlPegBondLengthDistributionBenchmark.name
+DATASET_FILENAME = "bond_length_distribution.json"
+
+CALC_PATH = CALCS_ROOT / "molecular_dynamics" / "bond_length_distribution" / "outputs"
+OUT_PATH = APP_ROOT / "data" / "molecular_dynamics" / "bond_length_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 bond length distribution input data.
+    """
+    return download_s3_data(
+        key="inputs/molecular_dynamics/bond_length_distribution/bond_length_distribution.zip",
+        filename="bond_length_distribution.zip",
+    )
+
+
+@pytest.fixture
+def analyze_results() -> dict:
+    """
+    Run the mlipaudit analysis for each model.
+
+    Returns
+    -------
+    dict
+        Mapping of model name to its ``BondLengthDistributionResult``.
+    """
+    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 = MlPegBondLengthDistributionBenchmark(
+            force_field=Calculator(),
+            data_input_dir=data_input_dir,
+            run_mode="standard",
+        )
+        benchmark.model_output = load_model_output_from_disk(
+            CALC_PATH / model_name, MlPegBondLengthDistributionBenchmark
+        )
+        results[model_name] = benchmark.analyze()
+    return results
+
+
+@pytest.fixture
+def struct_info() -> None:
+    """Write the combined element set to ``info.json`` for filtering."""
+    data_path = _data_input_dir() / BENCHMARK / DATASET_FILENAME
+    with open(data_path, encoding="utf-8") as f:
+        data = json.load(f)
+
+    elements = sorted(
+        {symbol for molecule in data.values() for symbol in molecule["atom_symbols"]}
+    )
+
+    OUT_PATH.mkdir(parents=True, exist_ok=True)
+    with (OUT_PATH / "info.json").open("w", encoding="utf-8") as f:
+        json.dump({"elements": elements}, f, indent=1)
+
+
+@pytest.fixture
+@plot_hist(
+    filename=str(OUT_PATH / "figure_bond_length_hist.json"),
+    title="Bond length deviation distribution",
+    x_label="Bond length deviation / Å",
+    y_label="Probability density",
+    bins=50,
+)
+def deviation_distributions(analyze_results) -> dict[str, np.ndarray]:
+    """
+    Collect the bond length deviations sampled along each model's trajectories.
+
+    Parameters
+    ----------
+    analyze_results
+        Mapping of model name to its ``BondLengthDistributionResult``.
+
+    Returns
+    -------
+    dict[str, np.ndarray]
+        Per-model flat array of bond length deviations across all molecules.
+    """
+    results = {}
+    for model_name, result in analyze_results.items():
+        if result.failed:
+            continue
+        deviations = [
+            value
+            for molecule in result.molecules
+            if molecule.deviation_trajectory is not None
+            for value in molecule.deviation_trajectory
+        ]
+        if deviations:
+            results[model_name] = np.array(deviations)
+    return results
+
+
+@pytest.fixture
+def get_avg_deviation(analyze_results) -> dict[str, float]:
+    """
+    Get the average bond length deviation for each model.
+
+    Parameters
+    ----------
+    analyze_results
+        Mapping of model name to its ``BondLengthDistributionResult``.
+
+    Returns
+    -------
+    dict[str, float]
+        Mean absolute bond length deviation over the trajectories, in Angstrom.
+    """
+    return {
+        model_name: result.avg_deviation
+        for model_name, result in analyze_results.items()
+    }
+
+
+@pytest.fixture
+@build_table(
+    filename=OUT_PATH / "bond_length_distribution_metrics_table.json",
+    metric_tooltips=DEFAULT_TOOLTIPS,
+    thresholds=DEFAULT_THRESHOLDS,
+    weights=DEFAULT_WEIGHTS,
+    mlip_name_map=DISPERSION_NAME_MAP,
+)
+def metrics(
+    deviation_distributions,
+    get_avg_deviation: dict[str, float],
+) -> dict[str, dict]:
+    """
+    Get all metrics.
+
+    Parameters
+    ----------
+    deviation_distributions
+        Per-model deviation arrays (triggers the histogram plot).
+    get_avg_deviation
+        Average bond length deviations for all models.
+
+    Returns
+    -------
+    dict[str, dict]
+        Metric names and values for all models.
+    """
+    return {
+        "Bond Length Deviation": get_avg_deviation,
+    }
+
+
+def test_bond_length_distribution(metrics: dict[str, dict], struct_info: None) -> None:
+    """
+    Run bond length distribution analysis.
+
+    Parameters
+    ----------
+    metrics : dict[str, dict]
+        Bond length metric results provided by fixtures.
+    struct_info : None
+        Element info written to ``info.json`` for filtering.
+    """

ml_peg/analysis/molecular_dynamics/bond_length_distribution/metrics.yml

@@ -0,0 +1,8 @@
+metrics:
+  Bond Length Deviation:
+    good: 0.0
+    bad: 0.05
+    unit: Å
+    weight: 1
+    tooltip: Mean absolute deviation of a tracked covalent bond from its QM-optimised reference length, averaged over the MD trajectory and across all molecules.
+    level_of_theory: DFT

ml_peg/app/molecular_dynamics/bond_length_distribution/app_bond_length_distribution.py

@@ -0,0 +1,66 @@
+"""Run bond length 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 = "BondLength"
+DOCS_URL = "https://ddmms.github.io/ml-peg/user_guide/benchmarks/molecular_dynamics.html#bond-length-distribution"
+DATA_PATH = APP_ROOT / "data" / "molecular_dynamics" / "bond_length_distribution"
+
+
+class BondLengthApp(BaseApp):
+    """Bond length distribution benchmark app layout and callbacks."""
+
+    def register_callbacks(self) -> None:
+        """Register callbacks to app."""
+        histogram = read_plot(
+            DATA_PATH / "figure_bond_length_hist.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={"Bond Length Deviation": histogram},
+        )
+
+
+def get_app() -> BondLengthApp:
+    """
+    Get bond length distribution benchmark app layout and callback registration.
+
+    Returns
+    -------
+    BondLengthApp
+        Benchmark layout and callback registration.
+    """
+    return BondLengthApp(
+        name="Bond Length Distribution",
+        framework_ids="mlip_audit",
+        description=(
+            "Performance in maintaining physically reasonable covalent bond "
+            "lengths during molecular dynamics of small organic molecules. "
+            "Reference bond lengths are taken from QM-optimised geometries."
+        ),
+        docs_url=DOCS_URL,
+        table_path=DATA_PATH / "bond_length_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,12 @@ 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"
+
 mace-multihead:
   label: Multihead Cross Learning
   color: "#7c3aed"