diff --git a/docs/source/user_guide/benchmarks/biomolecules.rst b/docs/source/user_guide/benchmarks/biomolecules.rst new file mode 100644 index 000000000..75dfb7f89 --- /dev/null +++ b/docs/source/user_guide/benchmarks/biomolecules.rst @@ -0,0 +1,58 @@ +============ +Biomolecules +============ + +Protein folding stability +========================= + +Summary +------- + +Performance in keeping small proteins folded during molecular dynamics. For each +protein, an NVT molecular dynamics simulation is run at 300 K starting from the +native (folded) reference conformation, and the ability of the model to retain the +fold is measured along the trajectory. The benchmark uses a small set of well +characterised proteins (chignolin, tryptophan cage, and an orexin/hypocretin +fragment) with experimental reference structures. + +Metrics +------- + +1. RMSD + +The root mean square deviation of the C-alpha atoms from the native reference +structure is computed for each frame of the trajectory, then averaged over the +trajectory and across all proteins. A lower RMSD indicates the fold is retained. + +2. TM score + +The TM score of each trajectory frame against the native reference structure is +computed, then averaged over the trajectory and across all proteins. A TM score +closer to 1 indicates that the global fold is preserved. + +3. Radius of gyration deviation + +The maximum absolute deviation of the radius of gyration from the initial folded +state along the trajectory, taken across all proteins. A lower value indicates the +protein does not unfold or collapse. + +A line plot shows the RMSD from the reference structure along the trajectory, +averaged across the proteins, for each model. + +Computational cost +------------------ + +High: one MD simulation per protein. Faster inference can be achieved using the +jax-accelerated simulations in MLIP Audit directly. + +Data availability +----------------- + +Input structures: + +* MLIP Audit benchmark suite, InstaDeep. Native reference structures taken from the + Protein Data Bank (chignolin 1UAO, tryptophan cage 2JOF, orexin-B 1CQ0). + +Reference data: + +* Experimental reference structures (X-ray and NMR) from the Protein Data Bank. diff --git a/docs/source/user_guide/benchmarks/index.rst b/docs/source/user_guide/benchmarks/index.rst index af7bb5976..2d0a2d068 100644 --- a/docs/source/user_guide/benchmarks/index.rst +++ b/docs/source/user_guide/benchmarks/index.rst @@ -20,3 +20,4 @@ Benchmarks conformers molecular_dynamics defect + biomolecules diff --git a/ml_peg/analysis/biomolecules/protein_folding_stability/analyse_protein_folding_stability.py b/ml_peg/analysis/biomolecules/protein_folding_stability/analyse_protein_folding_stability.py new file mode 100644 index 000000000..d9c0b2709 --- /dev/null +++ b/ml_peg/analysis/biomolecules/protein_folding_stability/analyse_protein_folding_stability.py @@ -0,0 +1,298 @@ +"""Analyse the protein folding stability benchmark.""" + +from __future__ import annotations + +import json +from pathlib import Path + +from ase.calculators.calculator import Calculator +from ase.io import read +import numpy as np +import pytest + +pytest.importorskip("mlipaudit", reason="Please install `mlipaudit` extra") +from mlipaudit.benchmarks.folding_stability.folding_stability import STRUCTURE_NAMES +from mlipaudit.io import load_model_output_from_disk + +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, +) +from ml_peg.app import APP_ROOT +from ml_peg.calcs import CALCS_ROOT +from ml_peg.calcs.utils.mlipaudit import MlPegFoldingStabilityBenchmark +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 = MlPegFoldingStabilityBenchmark.name + +CALC_PATH = CALCS_ROOT / "biomolecules" / "protein_folding_stability" / "outputs" +OUT_PATH = APP_ROOT / "data" / "biomolecules" / "protein_folding_stability" + +METRICS_CONFIG_PATH = Path(__file__).with_name("metrics.yml") +DEFAULT_THRESHOLDS, DEFAULT_TOOLTIPS, DEFAULT_WEIGHTS = load_metrics_config( + METRICS_CONFIG_PATH +) + + +def structure_xyz(structure_name: str) -> Path: + """ + Get the path to a starting structure saved by the calculation. + + Parameters + ---------- + structure_name + Name of the structure. + + Returns + ------- + Path + Path to the structure's starting geometry. + """ + return CALC_PATH / BENCHMARK / "starting_structures" / f"{structure_name}.xyz" + + +def check_dataset() -> None: + """ + Check the input structures saved by the calculation are available. + + The calculation copies the downloaded input data into its outputs, so the + analysis does not need to download it again. + + Raises + ------ + ValueError + If any starting structure is missing from the calculation outputs. + """ + for structure_name in STRUCTURE_NAMES: + if not structure_xyz(structure_name).exists(): + raise ValueError( + f"{structure_xyz(structure_name)} does not exist. " + "Please run the calculation." + ) + + +@pytest.fixture +def analyze_results() -> dict: + """ + Run the mlipaudit analysis for each model. + + Returns + ------- + dict + Mapping of model name to its ``FoldingStabilityResult``. + """ + check_dataset() + + results = {} + for model_name in MODELS: + output_dir = CALC_PATH / model_name / BENCHMARK + if not (output_dir / "model_output.zip").exists(): + continue + benchmark = MlPegFoldingStabilityBenchmark( + force_field=Calculator(), + data_input_dir=CALC_PATH, + run_mode="standard", + ) + benchmark.model_output = load_model_output_from_disk( + CALC_PATH / model_name, MlPegFoldingStabilityBenchmark + ) + results[model_name] = benchmark.analyze() + return results + + +@pytest.fixture +def struct_info() -> dict: + """ + Write per-structure element info to ``info.json`` for filtering. + + Elements are stored as one list per structure, so individual structures can + be excluded once partial filtering is supported. The order follows + ``STRUCTURE_NAMES``. + + Returns + ------- + dict + Mapping with the per-structure lists of elements. + """ + check_dataset() + + info = { + "systems": list(STRUCTURE_NAMES), + "elements": [ + sorted(set(read(structure_xyz(name)).get_chemical_symbols())) + for name in STRUCTURE_NAMES + ], + } + + OUT_PATH.mkdir(parents=True, exist_ok=True) + with (OUT_PATH / "info.json").open("w", encoding="utf-8") as f: + json.dump(info, f, indent=1) + + return info + + +@pytest.fixture +@plot_scatter( + title="RMSD from reference structure along trajectory", + x_label="Frame", + y_label="RMSD / Å", + show_line=True, + show_markers=False, + filename=str(OUT_PATH / "figure_rmsd_trajectory.json"), +) +def rmsd_trajectories(analyze_results) -> dict[str, tuple[list, list]]: + """ + Get the RMSD trajectory averaged across structures for each model. + + Parameters + ---------- + analyze_results + Mapping of model name to its ``FoldingStabilityResult``. + + Returns + ------- + dict[str, tuple[list, list]] + Per-model ``(frame, mean RMSD)`` profiles across the trajectory. + """ + results = {} + for model_name, result in analyze_results.items(): + if result.failed: + continue + trajectories = [ + molecule.rmsd_trajectory + for molecule in result.molecules + if not molecule.failed and molecule.rmsd_trajectory is not None + ] + if not trajectories: + continue + num_frames = min(len(traj) for traj in trajectories) + stacked = np.array([traj[:num_frames] for traj in trajectories]) + mean_rmsd = stacked.mean(axis=0) + results[model_name] = (list(range(num_frames)), mean_rmsd.tolist()) + return results + + +@pytest.fixture +def get_avg_rmsd(analyze_results) -> dict[str, float]: + """ + Get the average RMSD for each model. + + Parameters + ---------- + analyze_results + Mapping of model name to its ``FoldingStabilityResult``. + + Returns + ------- + dict[str, float] + Average RMSD from the reference structure, averaged across molecules. + """ + return { + model_name: (result.avg_rmsd if result.avg_rmsd is not None else np.nan) + for model_name, result in analyze_results.items() + } + + +@pytest.fixture +def get_avg_tm_score(analyze_results) -> dict[str, float]: + """ + Get the average TM score for each model. + + Parameters + ---------- + analyze_results + Mapping of model name to its ``FoldingStabilityResult``. + + Returns + ------- + dict[str, float] + Average TM score against the reference structure, averaged across molecules. + """ + return { + model_name: (result.avg_tm_score if result.avg_tm_score is not None else np.nan) + for model_name, result in analyze_results.items() + } + + +@pytest.fixture +def get_rgyr_deviation(analyze_results) -> dict[str, float]: + """ + Get the maximum radius of gyration deviation for each model. + + Parameters + ---------- + analyze_results + Mapping of model name to its ``FoldingStabilityResult``. + + Returns + ------- + dict[str, float] + Maximum absolute deviation of the radius of gyration from the initial + state, taken across molecules, in Angstrom. + """ + return { + model_name: ( + result.max_abs_deviation_radius_of_gyration + if result.max_abs_deviation_radius_of_gyration is not None + else np.nan + ) + for model_name, result in analyze_results.items() + } + + +@pytest.fixture +@build_table( + filename=OUT_PATH / "protein_folding_stability_metrics_table.json", + metric_tooltips=DEFAULT_TOOLTIPS, + thresholds=DEFAULT_THRESHOLDS, + weights=DEFAULT_WEIGHTS, + mlip_name_map=DISPERSION_NAME_MAP, +) +def metrics( + rmsd_trajectories, + get_avg_rmsd: dict[str, float], + get_avg_tm_score: dict[str, float], + get_rgyr_deviation: dict[str, float], +) -> dict[str, dict]: + """ + Get all metrics. + + Parameters + ---------- + rmsd_trajectories + Per-model averaged RMSD trajectories (triggers the RMSD line plot). + get_avg_rmsd + Average RMSD values for all models. + get_avg_tm_score + Average TM scores for all models. + get_rgyr_deviation + Maximum radius of gyration deviations for all models. + + Returns + ------- + dict[str, dict] + Metric names and values for all models. + """ + return { + "RMSD": get_avg_rmsd, + "TM Score": get_avg_tm_score, + "Rgyr Deviation": get_rgyr_deviation, + } + + +def test_protein_folding_stability(metrics: dict[str, dict], struct_info: dict) -> None: + """ + Run protein folding stability analysis. + + Parameters + ---------- + metrics : dict[str, dict] + Protein folding stability metric results provided by fixtures. + struct_info : dict + Element info written to ``info.json`` for filtering. + """ diff --git a/ml_peg/analysis/biomolecules/protein_folding_stability/metrics.yml b/ml_peg/analysis/biomolecules/protein_folding_stability/metrics.yml new file mode 100644 index 000000000..f2daca9a3 --- /dev/null +++ b/ml_peg/analysis/biomolecules/protein_folding_stability/metrics.yml @@ -0,0 +1,22 @@ +metrics: + RMSD: + good: 0.0 + bad: 2.0 + unit: Å + weight: 1 + tooltip: Root mean square deviation of the C-alpha atoms from the native reference structure, averaged over the trajectory and across all proteins. Lower is better. + level_of_theory: Experiment + TM Score: + good: 1.0 + bad: 0.5 + unit: null + weight: 1 + tooltip: TM score of each trajectory frame against the native reference structure, averaged over the trajectory and across all proteins. Closer to 1 indicates the fold is retained. + level_of_theory: Experiment + Rgyr Deviation: + good: 0.0 + bad: 2.0 + unit: Å + weight: 1 + tooltip: Maximum absolute deviation of the radius of gyration from the initial folded state along the trajectory, taken across all proteins. Lower is better. + level_of_theory: Experiment diff --git a/ml_peg/app/biomolecules/biomolecules.yml b/ml_peg/app/biomolecules/biomolecules.yml new file mode 100644 index 000000000..8ec5ee4f1 --- /dev/null +++ b/ml_peg/app/biomolecules/biomolecules.yml @@ -0,0 +1,3 @@ +title: Biomolecules +description: Folding stability and dynamics of proteins and other biomolecules +weight: 1 diff --git a/ml_peg/app/biomolecules/protein_folding_stability/app_protein_folding_stability.py b/ml_peg/app/biomolecules/protein_folding_stability/app_protein_folding_stability.py new file mode 100644 index 000000000..d735a4622 --- /dev/null +++ b/ml_peg/app/biomolecules/protein_folding_stability/app_protein_folding_stability.py @@ -0,0 +1,67 @@ +"""Run protein folding stability 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 = "ProteinFoldingStability" +DOCS_URL = "https://ddmms.github.io/ml-peg/user_guide/benchmarks/biomolecules.html#protein-folding-stability" +DATA_PATH = APP_ROOT / "data" / "biomolecules" / "protein_folding_stability" + + +class ProteinFoldingStabilityApp(BaseApp): + """Protein folding stability benchmark app layout and callbacks.""" + + def register_callbacks(self) -> None: + """Register callbacks to app.""" + scatter = read_plot( + DATA_PATH / "figure_rmsd_trajectory.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={"RMSD": scatter}, + ) + + +def get_app() -> ProteinFoldingStabilityApp: + """ + Get protein folding stability benchmark app layout and callback registration. + + Returns + ------- + ProteinFoldingStabilityApp + Benchmark layout and callback registration. + """ + return ProteinFoldingStabilityApp( + name="Protein Folding Stability", + framework_ids="mlip_audit", + description=( + "Performance in keeping small proteins folded during molecular " + "dynamics started from their native conformation. The RMSD, TM " + "score, and radius of gyration relative to experimental reference " + "structures are tracked along the trajectory." + ), + docs_url=DOCS_URL, + table_path=DATA_PATH / "protein_folding_stability_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) diff --git a/ml_peg/calcs/biomolecules/protein_folding_stability/calc_protein_folding_stability.py b/ml_peg/calcs/biomolecules/protein_folding_stability/calc_protein_folding_stability.py new file mode 100644 index 000000000..b57448b38 --- /dev/null +++ b/ml_peg/calcs/biomolecules/protein_folding_stability/calc_protein_folding_stability.py @@ -0,0 +1,84 @@ +""" +Assess protein folding stability during molecular dynamics. + +A molecular dynamics simulation is run for each of a set of small proteins +starting from their native (folded) conformation, and the ability of the model +to keep each protein folded is measured along the trajectory via the RMSD, +TM score, and radius of gyration relative to the reference structure. +""" + +from __future__ import annotations + +from pathlib import Path +import shutil +from typing import Any +from warnings import warn + +import pytest + +pytest.importorskip("mlipaudit", reason="Please install `mlipaudit` extra") +from mlipaudit.benchmarks.folding_stability.folding_stability import ( + STRUCTURE_NAMES, + FoldingStabilityModelOutput, +) +from mlipaudit.io import write_model_output_to_disk + +from ml_peg.calcs.utils.mlipaudit import MlPegFoldingStabilityBenchmark +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_folding_stability(mlip: tuple[str, Any]) -> None: + """ + Benchmark protein folding stability 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_folding_stability/protein_folding_stability.zip", + filename="protein_folding_stability.zip", + ) + + # Save the input data to the calculation outputs so the analysis is self + # contained and does not need to download it again. + name = MlPegFoldingStabilityBenchmark.name + shutil.copytree(data_input_dir / name, OUT_PATH / name, dirs_exist_ok=True) + + benchmark = MlPegFoldingStabilityBenchmark( + 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 folding stability benchmark for " + f"{model_name}: {exc}", + stacklevel=2, + ) + # Structures with a ``None`` simulation state are treated as failed by + # analyze(), so this yields a failed result for every structure. + benchmark.model_output = FoldingStabilityModelOutput( + structure_names=list(STRUCTURE_NAMES), + simulation_states=[None] * len(STRUCTURE_NAMES), + ) + + write_model_output_to_disk( + MlPegFoldingStabilityBenchmark.name, + benchmark.model_output, + OUT_PATH / model_name, + ) diff --git a/ml_peg/calcs/utils/mlipaudit.py b/ml_peg/calcs/utils/mlipaudit.py index e725c3f88..3536a43f4 100644 --- a/ml_peg/calcs/utils/mlipaudit.py +++ b/ml_peg/calcs/utils/mlipaudit.py @@ -5,6 +5,9 @@ from mlipaudit.benchmarks.conformer_selection.conformer_selection import ( ConformerSelectionBenchmark, ) +from mlipaudit.benchmarks.folding_stability.folding_stability import ( + FoldingStabilityBenchmark, +) from mlipaudit.benchmarks.tautomers.tautomers import TautomersBenchmark @@ -19,6 +22,19 @@ class MlPegConformerSelectionBenchmark(ConformerSelectionBenchmark): skip_if_elements_missing = False +class MlPegFoldingStabilityBenchmark(FoldingStabilityBenchmark): + """ + ``FoldingStabilityBenchmark`` 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 + + class MlPegTautomersBenchmark(TautomersBenchmark): """ TautomersBenchmark wired up for ml-peg's ASE calculators. diff --git a/pyproject.toml b/pyproject.toml index b09fcdf99..79b9ac99c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -71,6 +71,9 @@ mlipaudit = [ "jax<0.11; python_version >= '3.11'", "jaxlib<0.11; python_version >= '3.11'", ] +mlipaudit = [ + "mlipaudit; python_version >= '3.11'", +] orb = [ "orb-models == 0.6.2; sys_platform != 'win32' and python_version >= '3.12'", ]