diff --git a/docs/source/user_guide/benchmarks/biomolecules.rst b/docs/source/user_guide/benchmarks/biomolecules.rst new file mode 100644 index 000000000..7227a22d0 --- /dev/null +++ b/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. diff --git a/docs/source/user_guide/benchmarks/index.rst b/docs/source/user_guide/benchmarks/index.rst index ea0dcafb2..423f7fd6a 100644 --- a/docs/source/user_guide/benchmarks/index.rst +++ b/docs/source/user_guide/benchmarks/index.rst @@ -18,4 +18,5 @@ Benchmarks tm_complexes conformers molecular_dynamics + biomolecules defect diff --git a/ml_peg/analysis/biomolecules/protein_sampling/analyse_protein_sampling.py b/ml_peg/analysis/biomolecules/protein_sampling/analyse_protein_sampling.py new file mode 100644 index 000000000..733362eca --- /dev/null +++ b/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. + """ diff --git a/ml_peg/analysis/biomolecules/protein_sampling/metrics.yml b/ml_peg/analysis/biomolecules/protein_sampling/metrics.yml new file mode 100644 index 000000000..1357b0369 --- /dev/null +++ b/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 diff --git a/ml_peg/app/biomolecules/biomolecules.yml b/ml_peg/app/biomolecules/biomolecules.yml new file mode 100644 index 000000000..ac8d4082b --- /dev/null +++ b/ml_peg/app/biomolecules/biomolecules.yml @@ -0,0 +1,3 @@ +title: Biomolecules +description: Conformational sampling and stability of proteins and other biomolecules +weight: 1 diff --git a/ml_peg/app/biomolecules/protein_sampling/app_protein_sampling.py b/ml_peg/app/biomolecules/protein_sampling/app_protein_sampling.py new file mode 100644 index 000000000..26feead00 --- /dev/null +++ b/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) diff --git a/ml_peg/app/utils/frameworks.yml b/ml_peg/app/utils/frameworks.yml index 27b079925..4bb28f12a 100644 --- a/ml_peg/app/utils/frameworks.yml +++ b/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" diff --git a/ml_peg/calcs/biomolecules/protein_sampling/calc_protein_sampling.py b/ml_peg/calcs/biomolecules/protein_sampling/calc_protein_sampling.py new file mode 100644 index 000000000..aea49739c --- /dev/null +++ b/ml_peg/calcs/biomolecules/protein_sampling/calc_protein_sampling.py @@ -0,0 +1,72 @@ +""" +Sample protein conformations with molecular dynamics. + +A molecular dynamics simulation is run for each of a set of small proteins, and +the sampled backbone and side-chain dihedral angles are compared to reference +distributions to assess how well the model explores conformational space. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any +from warnings import warn + +from mlipaudit.benchmarks.sampling.sampling import ( + STRUCTURE_NAMES, + SamplingModelOutput, +) +from mlipaudit.io import write_model_output_to_disk +import pytest + +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) + +OUT_PATH = Path(__file__).parent / "outputs" + + +@pytest.mark.parametrize("mlip", MODELS.items()) +def test_protein_sampling(mlip: tuple[str, Any]) -> None: + """ + Benchmark protein conformational sampling during MD. + + Parameters + ---------- + mlip + Name of model and model object to get calculator. + """ + model_name, model = mlip + calc = model.get_calculator() + calc = model.add_d3_calculator(calc) + + data_input_dir = download_s3_data( + key="inputs/biomolecules/protein_sampling/protein_sampling.zip", + filename="protein_sampling.zip", + ) + + benchmark = MlPegSamplingBenchmark( + force_field=calc, + data_input_dir=data_input_dir, + run_mode="standard", + ) + try: + benchmark.run_model() + except Exception as exc: + warn( + f"Error running protein sampling benchmark for {model_name}: {exc}", + stacklevel=2, + ) + # Simulation states of None for every system are treated as failed + # simulations by analyze(), which then reports a failed benchmark. + benchmark.model_output = SamplingModelOutput( + structure_names=list(STRUCTURE_NAMES), + simulation_states=[None] * len(STRUCTURE_NAMES), + ) + + write_model_output_to_disk( + MlPegSamplingBenchmark.name, benchmark.model_output, OUT_PATH / model_name + ) diff --git a/ml_peg/calcs/utils/mlipaudit.py b/ml_peg/calcs/utils/mlipaudit.py new file mode 100644 index 000000000..b1f47b7c1 --- /dev/null +++ b/ml_peg/calcs/utils/mlipaudit.py @@ -0,0 +1,18 @@ +"""Adapters for using mlipaudit benchmarks with ml-peg's ASE calculators.""" + +from __future__ import annotations + +from mlipaudit.benchmarks.sampling.sampling import SamplingBenchmark + + +class MlPegSamplingBenchmark(SamplingBenchmark): + """ + ``SamplingBenchmark`` wired up for ml-peg's ASE calculators. + + ``skip_if_elements_missing`` is disabled because ml-peg's ASE ``Calculator`` + objects do not expose the set of elements the underlying model supports, so + the benchmark cannot decide up front whether to skip. Missing element errors + are instead handled at runtime. + """ + + skip_if_elements_missing = False diff --git a/pyproject.toml b/pyproject.toml index ec7590a04..f86aa7c47 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -64,6 +64,9 @@ mace = [ mattersim = [ "mattersim==1.2.2", ] +mlipaudit = [ + "mlipaudit; python_version >= '3.11'", +] orb = [ "orb-models == 0.6.2; sys_platform != 'win32' and python_version >= '3.12'", ] @@ -206,6 +209,10 @@ conflicts = [ { extra = "mattersim" }, { extra = "grace" }, ], + [ + { extra = "mlipaudit" }, + { extra = "grace" }, + ], ] constraint-dependencies = [ @@ -217,3 +224,4 @@ module-root = "" [tool.uv.sources] asemolec = { git = "https://github.com/imagdau/aseMolec.git" } +mlipaudit = { git = "https://github.com/instadeepai/MLIPAudit.git", branch = "mlpeg-migration" }