feat: add ring planarity benchmark

docs/source/user_guide/benchmarks/molecular_dynamics.rst

@@ -87,7 +87,6 @@ Summary
 Benchmark of the density of water-ethanol mixtures for different concentrations of ethanol, compare to experiment.
 1 ns of NPT MD on about 120 water/ethanol molecules for 6 concentrations.
 
-
 Metrics
 -------
 
@@ -116,3 +115,45 @@ Packmol generated
 Reference data:
 * M. Southard and D. Green, Perry’s Chemical Engineers’ Handbook, 9th Edition. McGraw-Hill Education, 2018.
 * Experimental
+
+
+Ring planarity
+==============
+
+Summary
+-------
+
+Performance in maintaining planar aromatic rings during molecular dynamics of small
+organic molecules. For each molecule, an NVT molecular dynamics simulation is run at 300 K
+starting from a QM-optimised reference geometry (selected from QM9), and the deviation of
+the ring atoms from their best-fit plane is measured along the trajectory.
+
+Metrics
+-------
+
+1. Planarity deviation
+
+At each frame of the trajectory, the ring atoms are fitted to a plane and the root mean
+square deviation of the atoms from that plane is calculated. This is averaged over the
+trajectory and across all molecules. Aromatic rings are planar, so a well behaved potential
+keeps this deviation small; a lower deviation is better.
+
+A histogram shows the distribution of the sampled planarity 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 reference geometries of the aromatic molecules.

ml_peg/analysis/molecular_dynamics/ring_planarity/analyse_ring_planarity.py

@@ -0,0 +1,200 @@
+"""Analyse the aromatic ring planarity 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 MlPegRingPlanarityBenchmark
+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 = MlPegRingPlanarityBenchmark.name
+DATASET_FILENAME = "ring_planarity_data.json"
+
+CALC_PATH = CALCS_ROOT / "molecular_dynamics" / "ring_planarity" / "outputs"
+OUT_PATH = APP_ROOT / "data" / "molecular_dynamics" / "ring_planarity"
+
+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 ring planarity input data.
+    """
+    return download_s3_data(
+        key="inputs/molecular_dynamics/ring_planarity/ring_planarity.zip",
+        filename="ring_planarity.zip",
+    )
+
+
+@pytest.fixture
+def analyze_results() -> dict:
+    """
+    Run the mlipaudit analysis for each model.
+
+    Returns
+    -------
+    dict
+        Mapping of model name to its ``RingPlanarityResult``.
+    """
+    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 = MlPegRingPlanarityBenchmark(
+            force_field=Calculator(),
+            data_input_dir=data_input_dir,
+            run_mode="standard",
+        )
+        benchmark.model_output = load_model_output_from_disk(
+            CALC_PATH / model_name, MlPegRingPlanarityBenchmark
+        )
+        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_ring_planarity_hist.json"),
+    title="Ring planarity deviation distribution",
+    x_label="Planarity deviation / Å",
+    y_label="Probability density",
+    bins=50,
+)
+def deviation_distributions(analyze_results) -> dict[str, np.ndarray]:
+    """
+    Collect the planarity deviations sampled along each model's trajectories.
+
+    Parameters
+    ----------
+    analyze_results
+        Mapping of model name to its ``RingPlanarityResult``.
+
+    Returns
+    -------
+    dict[str, np.ndarray]
+        Per-model flat array of ring planarity 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_mae_deviation(analyze_results) -> dict[str, float]:
+    """
+    Get the mean planarity deviation for each model.
+
+    Parameters
+    ----------
+    analyze_results
+        Mapping of model name to its ``RingPlanarityResult``.
+
+    Returns
+    -------
+    dict[str, float]
+        Mean planarity deviation of the ring atoms over the trajectories, in Angstrom.
+    """
+    return {
+        model_name: result.mae_deviation
+        for model_name, result in analyze_results.items()
+    }
+
+
+@pytest.fixture
+@build_table(
+    filename=OUT_PATH / "ring_planarity_metrics_table.json",
+    metric_tooltips=DEFAULT_TOOLTIPS,
+    thresholds=DEFAULT_THRESHOLDS,
+    weights=DEFAULT_WEIGHTS,
+    mlip_name_map=DISPERSION_NAME_MAP,
+)
+def metrics(
+    deviation_distributions,
+    get_mae_deviation: dict[str, float],
+) -> dict[str, dict]:
+    """
+    Get all metrics.
+
+    Parameters
+    ----------
+    deviation_distributions
+        Per-model deviation arrays (triggers the histogram plot).
+    get_mae_deviation
+        Mean planarity deviations for all models.
+
+    Returns
+    -------
+    dict[str, dict]
+        Metric names and values for all models.
+    """
+    return {
+        "Planarity Deviation": get_mae_deviation,
+    }
+
+
+def test_ring_planarity(metrics: dict[str, dict], struct_info: None) -> None:
+    """
+    Run ring planarity analysis.
+
+    Parameters
+    ----------
+    metrics : dict[str, dict]
+        Ring planarity metric results provided by fixtures.
+    struct_info : None
+        Element info written to ``info.json`` for filtering.
+    """

ml_peg/analysis/molecular_dynamics/ring_planarity/metrics.yml

@@ -0,0 +1,8 @@
+metrics:
+  Planarity Deviation:
+    good: 0.0
+    bad: 0.05
+    unit: Å
+    weight: 1
+    tooltip: Mean RMSD of the ring atoms from their best-fit plane, averaged over the MD trajectory and across all molecules.
+    level_of_theory: DFT

ml_peg/app/molecular_dynamics/ring_planarity/app_ring_planarity.py

@@ -0,0 +1,66 @@
+"""Run ring planarity 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 = "RingPlanarity"
+DOCS_URL = "https://ddmms.github.io/ml-peg/user_guide/benchmarks/molecular_dynamics.html#ring-planarity"
+DATA_PATH = APP_ROOT / "data" / "molecular_dynamics" / "ring_planarity"
+
+
+class RingPlanarityApp(BaseApp):
+    """Ring planarity benchmark app layout and callbacks."""
+
+    def register_callbacks(self) -> None:
+        """Register callbacks to app."""
+        histogram = read_plot(
+            DATA_PATH / "figure_ring_planarity_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={"Planarity Deviation": histogram},
+        )
+
+
+def get_app() -> RingPlanarityApp:
+    """
+    Get ring planarity benchmark app layout and callback registration.
+
+    Returns
+    -------
+    RingPlanarityApp
+        Benchmark layout and callback registration.
+    """
+    return RingPlanarityApp(
+        name="Ring Planarity",
+        framework_ids="mlip_audit",
+        description=(
+            "Performance in maintaining planar aromatic rings during molecular "
+            "dynamics of small organic molecules. Reference geometries are taken "
+            "from QM-optimised structures."
+        ),
+        docs_url=DOCS_URL,
+        table_path=DATA_PATH / "ring_planarity_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"