feat: add protein sampling benchmark

docs/source/user_guide/benchmarks/biomolecules.rst

@@ -0,0 +1,48 @@
+============
+Biomolecules
+============
+
+Protein sampling
+================
+
+Summary
+-------
+
+Performance in exploring the conformational space of small proteins during molecular
+dynamics. A short molecular dynamics simulation is run for each of a set of small
+proteins (chignolin, Trp-cage, and an orexin beta fragment), and the sampled backbone
+dihedral angles are compared to reference distributions. This probes whether a model
+samples physically reasonable protein conformations rather than only reproducing
+static reference geometries.
+
+Metrics
+-------
+
+1. Backbone Dihedral RMSD
+2. Backbone Hellinger Distance
+3. Backbone Outliers Ratio
+
+For each system, the sampled backbone (phi/psi) dihedral angles are collected across the
+trajectory and binned into a distribution per residue type. This distribution is
+compared to a reference distribution using the root mean square deviation (RMSD) and the
+Hellinger distance. The outliers ratio measures the fraction of sampled dihedrals that
+lie far from any point in the reference data. All three metrics are averaged over residue
+types and over the stable systems, and lower values are better.
+
+Computational cost
+------------------
+
+High: each model runs several molecular dynamics simulations of solvated proteins.
+
+Data availability
+-----------------
+
+Input structures:
+
+* Experimental structures from the Protein Data Bank (chignolin 1UAO, Trp-cage 2JOF,
+  orexin beta 1CQ0).
+
+Reference data:
+
+* Reference backbone and side-chain dihedral distributions derived from molecular
+  dynamics reference simulations.

docs/source/user_guide/benchmarks/index.rst

@@ -18,4 +18,5 @@ Benchmarks
     tm_complexes
     conformers
     molecular_dynamics
+    biomolecules
     defect

ml_peg/analysis/biomolecules/protein_sampling/analyse_protein_sampling.py

@@ -0,0 +1,203 @@
+"""Analyse the protein conformational sampling 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 pytest
+
+from ml_peg.analysis.utils.decorators import build_table
+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 MlPegSamplingBenchmark
+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 = MlPegSamplingBenchmark.name
+
+CALC_PATH = CALCS_ROOT / "biomolecules" / "protein_sampling" / "outputs"
+OUT_PATH = APP_ROOT / "data" / "biomolecules" / "protein_sampling"
+
+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 protein sampling input data.
+    """
+    return download_s3_data(
+        key="inputs/biomolecules/protein_sampling/protein_sampling.zip",
+        filename="protein_sampling.zip",
+    )
+
+
+@pytest.fixture
+def analyze_results() -> dict:
+    """
+    Run the mlipaudit analysis for each model.
+
+    Returns
+    -------
+    dict
+        Mapping of model name to its ``SamplingResult``.
+    """
+    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 = MlPegSamplingBenchmark(
+            force_field=Calculator(),
+            data_input_dir=data_input_dir,
+            run_mode="standard",
+        )
+        benchmark.model_output = load_model_output_from_disk(
+            CALC_PATH / model_name, MlPegSamplingBenchmark
+        )
+        results[model_name] = benchmark.analyze()
+    return results
+
+
+@pytest.fixture
+def struct_info() -> None:
+    """Write the combined element set to ``info.json`` for filtering."""
+    elements = sorted(MlPegSamplingBenchmark.required_elements)
+
+    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
+def get_rmsd_backbone(analyze_results) -> dict[str, float | None]:
+    """
+    Get the mean backbone dihedral distribution RMSD for each model.
+
+    Parameters
+    ----------
+    analyze_results
+        Mapping of model name to its ``SamplingResult``.
+
+    Returns
+    -------
+    dict[str, float | None]
+        Backbone dihedral distribution RMSD averaged over residues and systems.
+    """
+    return {
+        model_name: result.rmsd_backbone_total
+        for model_name, result in analyze_results.items()
+    }
+
+
+@pytest.fixture
+def get_hellinger_backbone(analyze_results) -> dict[str, float | None]:
+    """
+    Get the mean backbone dihedral Hellinger distance for each model.
+
+    Parameters
+    ----------
+    analyze_results
+        Mapping of model name to its ``SamplingResult``.
+
+    Returns
+    -------
+    dict[str, float | None]
+        Backbone dihedral Hellinger distance averaged over residues and systems.
+    """
+    return {
+        model_name: result.hellinger_distance_backbone_total
+        for model_name, result in analyze_results.items()
+    }
+
+
+@pytest.fixture
+def get_outliers_ratio_backbone(analyze_results) -> dict[str, float | None]:
+    """
+    Get the mean backbone dihedral outliers ratio for each model.
+
+    Parameters
+    ----------
+    analyze_results
+        Mapping of model name to its ``SamplingResult``.
+
+    Returns
+    -------
+    dict[str, float | None]
+        Fraction of sampled backbone dihedrals lying far from the reference data,
+        averaged over residues and systems.
+    """
+    return {
+        model_name: result.outliers_ratio_backbone_total
+        for model_name, result in analyze_results.items()
+    }
+
+
+@pytest.fixture
+@build_table(
+    filename=OUT_PATH / "protein_sampling_metrics_table.json",
+    metric_tooltips=DEFAULT_TOOLTIPS,
+    thresholds=DEFAULT_THRESHOLDS,
+    weights=DEFAULT_WEIGHTS,
+    mlip_name_map=DISPERSION_NAME_MAP,
+)
+def metrics(
+    get_rmsd_backbone: dict[str, float | None],
+    get_hellinger_backbone: dict[str, float | None],
+    get_outliers_ratio_backbone: dict[str, float | None],
+) -> dict[str, dict]:
+    """
+    Get all metrics.
+
+    Parameters
+    ----------
+    get_rmsd_backbone
+        Backbone dihedral distribution RMSD for all models.
+    get_hellinger_backbone
+        Backbone dihedral Hellinger distance for all models.
+    get_outliers_ratio_backbone
+        Backbone dihedral outliers ratio for all models.
+
+    Returns
+    -------
+    dict[str, dict]
+        Metric names and values for all models.
+    """
+    return {
+        "Backbone Dihedral RMSD": get_rmsd_backbone,
+        "Backbone Hellinger Distance": get_hellinger_backbone,
+        "Backbone Outliers Ratio": get_outliers_ratio_backbone,
+    }
+
+
+def test_protein_sampling(metrics: dict[str, dict], struct_info: None) -> None:
+    """
+    Run protein sampling analysis.
+
+    Parameters
+    ----------
+    metrics : dict[str, dict]
+        Protein sampling metric results provided by fixtures.
+    struct_info : None
+        Element info written to ``info.json`` for filtering.
+    """

ml_peg/analysis/biomolecules/protein_sampling/metrics.yml

@@ -0,0 +1,22 @@
+metrics:
+  Backbone Dihedral RMSD:
+    good: 0.0
+    bad: 0.001
+    unit: null
+    weight: 1
+    tooltip: RMSD between the sampled and reference backbone (phi/psi) dihedral distributions, averaged over residue types and systems. Lower is better.
+    level_of_theory: MD reference
+  Backbone Hellinger Distance:
+    good: 0.0
+    bad: 1.0
+    unit: null
+    weight: 1
+    tooltip: Hellinger distance between the sampled and reference backbone (phi/psi) dihedral distributions, averaged over residue types and systems. Lower is better.
+    level_of_theory: MD reference
+  Backbone Outliers Ratio:
+    good: 0.0
+    bad: 0.1
+    unit: null
+    weight: 1
+    tooltip: Fraction of sampled backbone dihedrals lying far from any reference data point, averaged over residue types and systems. Lower is better.
+    level_of_theory: MD reference

ml_peg/app/biomolecules/biomolecules.yml

@@ -0,0 +1,3 @@
+title: Biomolecules
+description: Conformational sampling and stability of proteins and other biomolecules
+weight: 1

ml_peg/app/biomolecules/protein_sampling/app_protein_sampling.py

@@ -0,0 +1,49 @@
+"""Run protein sampling benchmark app."""
+
+from __future__ import annotations
+
+from dash import Dash
+
+from ml_peg.app import APP_ROOT
+from ml_peg.app.base_app import BaseApp
+
+BENCHMARK_NAME = "ProteinSampling"
+DOCS_URL = "https://ddmms.github.io/ml-peg/user_guide/benchmarks/biomolecules.html#protein-sampling"
+DATA_PATH = APP_ROOT / "data" / "biomolecules" / "protein_sampling"
+
+
+class ProteinSamplingApp(BaseApp):
+    """Protein sampling benchmark app layout and callbacks."""
+
+
+def get_app() -> ProteinSamplingApp:
+    """
+    Get protein sampling benchmark app layout and callback registration.
+
+    Returns
+    -------
+    ProteinSamplingApp
+        Benchmark layout and callback registration.
+    """
+    return ProteinSamplingApp(
+        name="Protein Sampling",
+        framework_ids="mlip_audit",
+        description=(
+            "Performance in exploring protein conformational space during molecular "
+            "dynamics. Sampled backbone dihedral distributions are compared against "
+            "reference distributions via distribution RMSD, Hellinger distance, and "
+            "the ratio of outlying conformations."
+        ),
+        docs_url=DOCS_URL,
+        table_path=DATA_PATH / "protein_sampling_metrics_table.json",
+        info_path=DATA_PATH / "info.json",
+        extra_components=[],
+    )
+
+
+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=8071, 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"