From 56f569f42d513928dfcb510ae150ff38f4919fae Mon Sep 17 00:00:00 2001 From: lwalew Date: Thu, 23 Jul 2026 14:20:45 +0200 Subject: [PATCH 1/3] feat: add protein folding stability benchmark # Conflicts: # ml_peg/calcs/utils/mlipaudit.py --- .../user_guide/benchmarks/biomolecules.rst | 58 ++++ docs/source/user_guide/benchmarks/index.rst | 1 + .../analyse_protein_folding_stability.py | 251 ++++++++++++++++++ .../protein_folding_stability/metrics.yml | 22 ++ ml_peg/app/biomolecules/biomolecules.yml | 3 + .../app_protein_folding_stability.py | 67 +++++ .../calc_protein_folding_stability.py | 76 ++++++ ml_peg/calcs/utils/mlipaudit.py | 16 ++ pyproject.toml | 3 + 9 files changed, 497 insertions(+) create mode 100644 docs/source/user_guide/benchmarks/biomolecules.rst create mode 100644 ml_peg/analysis/biomolecules/protein_folding_stability/analyse_protein_folding_stability.py create mode 100644 ml_peg/analysis/biomolecules/protein_folding_stability/metrics.yml create mode 100644 ml_peg/app/biomolecules/biomolecules.yml create mode 100644 ml_peg/app/biomolecules/protein_folding_stability/app_protein_folding_stability.py create mode 100644 ml_peg/calcs/biomolecules/protein_folding_stability/calc_protein_folding_stability.py 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..b87c2070e --- /dev/null +++ b/ml_peg/analysis/biomolecules/protein_folding_stability/analyse_protein_folding_stability.py @@ -0,0 +1,251 @@ +"""Analyse the protein folding stability benchmark.""" + +from __future__ import annotations + +from pathlib import Path + +from ase.calculators.calculator import Calculator +from mlipaudit.benchmarks.folding_stability.folding_stability import STRUCTURE_NAMES +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_scatter +from ml_peg.analysis.utils.utils import ( + build_dispersion_name_map, + load_metrics_config, + write_struct_info, +) +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.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 = 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 _data_input_dir() -> Path: + """ + Download and return the benchmark input data directory. + + Returns + ------- + Path + Directory containing the extracted protein folding stability input data. + """ + return download_s3_data( + key="inputs/biomolecules/protein_folding_stability/protein_folding_stability.zip", + filename="protein_folding_stability.zip", + ) + + +@pytest.fixture +def analyze_results() -> dict: + """ + Run the mlipaudit analysis for each model. + + Returns + ------- + dict + Mapping of model name to its ``FoldingStabilityResult``. + """ + 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 = MlPegFoldingStabilityBenchmark( + force_field=Calculator(), + data_input_dir=data_input_dir, + 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() -> None: + """Write the combined element set to ``info.json`` for filtering.""" + data_input_dir = _data_input_dir() + write_struct_info( + data_path=[ + data_input_dir / BENCHMARK / "starting_structures" / f"{name}.xyz" + for name in STRUCTURE_NAMES + ], + out_path=OUT_PATH, + ) + + +@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 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 + 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 + 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: None) -> None: + """ + Run protein folding stability analysis. + + Parameters + ---------- + metrics : dict[str, dict] + Protein folding stability metric results provided by fixtures. + struct_info : None + 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..d821270ad --- /dev/null +++ b/ml_peg/calcs/biomolecules/protein_folding_stability/calc_protein_folding_stability.py @@ -0,0 +1,76 @@ +""" +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 +from typing import Any +from warnings import warn + +from mlipaudit.benchmarks.folding_stability.folding_stability import ( + STRUCTURE_NAMES, + FoldingStabilityModelOutput, +) +from mlipaudit.io import write_model_output_to_disk +import pytest + +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", + ) + + 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'", ] From 265a92eb8b46d406546b96bc576ba64669723472 Mon Sep 17 00:00:00 2001 From: lwalew Date: Mon, 3 Aug 2026 12:17:41 +0200 Subject: [PATCH 2/3] feat: skip if mlipaudit not installed Guard the module-level `mlipaudit` imports with `pytest.importorskip` so collection skips instead of erroring when the optional `mlipaudit` extra is not installed. Placed ahead of the `ml_peg.calcs.utils.mlipaudit` import, which pulls in `mlipaudit` unconditionally. Co-Authored-By: Claude Opus 5 (1M context) --- .../analyse_protein_folding_stability.py | 6 ++++-- .../calc_protein_folding_stability.py | 4 +++- 2 files changed, 7 insertions(+), 3 deletions(-) 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 index b87c2070e..f439f050f 100644 --- 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 @@ -5,11 +5,13 @@ from pathlib import Path from ase.calculators.calculator import Calculator -from mlipaudit.benchmarks.folding_stability.folding_stability import STRUCTURE_NAMES -from mlipaudit.io import load_model_output_from_disk 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, 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 index d821270ad..ac4286b6a 100644 --- 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 @@ -13,12 +13,14 @@ 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 -import pytest from ml_peg.calcs.utils.mlipaudit import MlPegFoldingStabilityBenchmark from ml_peg.calcs.utils.utils import download_s3_data From 03cd1901ad58510d1f1650740192a42f59ca8ee3 Mon Sep 17 00:00:00 2001 From: lwalew Date: Mon, 3 Aug 2026 12:55:18 +0200 Subject: [PATCH 3/3] feat: reuse calc output and filter per structure Save the downloaded input data to the calculation outputs and read it from there during analysis, so the analysis no longer re-downloads from S3. Raise a clear error if the calculation has not been run. Store elements as one list per structure rather than a single union, so individual structures can be excluded once partial filtering is supported. Report a failed model as NaN rather than None. Co-Authored-By: Claude Opus 5 (1M context) --- .../analyse_protein_folding_stability.py | 93 ++++++++++++++----- .../calc_protein_folding_stability.py | 6 ++ 2 files changed, 75 insertions(+), 24 deletions(-) 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 index f439f050f..d9c0b2709 100644 --- 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 @@ -2,9 +2,11 @@ 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 @@ -16,12 +18,10 @@ from ml_peg.analysis.utils.utils import ( build_dispersion_name_map, load_metrics_config, - write_struct_info, ) 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.calcs.utils.utils import download_s3_data from ml_peg.models import current_models from ml_peg.models.get_models import load_models @@ -39,19 +39,41 @@ ) -def _data_input_dir() -> Path: +def structure_xyz(structure_name: str) -> Path: """ - Download and return the benchmark input data directory. + Get the path to a starting structure saved by the calculation. + + Parameters + ---------- + structure_name + Name of the structure. Returns ------- Path - Directory containing the extracted protein folding stability input data. + Path to the structure's starting geometry. + """ + return CALC_PATH / BENCHMARK / "starting_structures" / f"{structure_name}.xyz" + + +def check_dataset() -> None: """ - return download_s3_data( - key="inputs/biomolecules/protein_folding_stability/protein_folding_stability.zip", - filename="protein_folding_stability.zip", - ) + 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 @@ -64,7 +86,7 @@ def analyze_results() -> dict: dict Mapping of model name to its ``FoldingStabilityResult``. """ - data_input_dir = _data_input_dir() + check_dataset() results = {} for model_name in MODELS: @@ -73,7 +95,7 @@ def analyze_results() -> dict: continue benchmark = MlPegFoldingStabilityBenchmark( force_field=Calculator(), - data_input_dir=data_input_dir, + data_input_dir=CALC_PATH, run_mode="standard", ) benchmark.model_output = load_model_output_from_disk( @@ -84,16 +106,34 @@ def analyze_results() -> dict: @pytest.fixture -def struct_info() -> None: - """Write the combined element set to ``info.json`` for filtering.""" - data_input_dir = _data_input_dir() - write_struct_info( - data_path=[ - data_input_dir / BENCHMARK / "starting_structures" / f"{name}.xyz" +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=OUT_PATH, - ) + } + + 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 @@ -153,7 +193,8 @@ def get_avg_rmsd(analyze_results) -> dict[str, float]: Average RMSD from the reference structure, averaged across molecules. """ return { - model_name: result.avg_rmsd for model_name, result in analyze_results.items() + model_name: (result.avg_rmsd if result.avg_rmsd is not None else np.nan) + for model_name, result in analyze_results.items() } @@ -173,7 +214,7 @@ def get_avg_tm_score(analyze_results) -> dict[str, float]: Average TM score against the reference structure, averaged across molecules. """ return { - model_name: result.avg_tm_score + 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() } @@ -195,7 +236,11 @@ def get_rgyr_deviation(analyze_results) -> dict[str, float]: state, taken across molecules, in Angstrom. """ return { - model_name: result.max_abs_deviation_radius_of_gyration + 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() } @@ -240,7 +285,7 @@ def metrics( } -def test_protein_folding_stability(metrics: dict[str, dict], struct_info: None) -> None: +def test_protein_folding_stability(metrics: dict[str, dict], struct_info: dict) -> None: """ Run protein folding stability analysis. @@ -248,6 +293,6 @@ def test_protein_folding_stability(metrics: dict[str, dict], struct_info: None) ---------- metrics : dict[str, dict] Protein folding stability metric results provided by fixtures. - struct_info : None + struct_info : dict Element info written to ``info.json`` for filtering. """ 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 index ac4286b6a..b57448b38 100644 --- 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 @@ -10,6 +10,7 @@ from __future__ import annotations from pathlib import Path +import shutil from typing import Any from warnings import warn @@ -51,6 +52,11 @@ def test_protein_folding_stability(mlip: tuple[str, Any]) -> None: 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,