From 59ee2d88ab18d25479e05c538abb2ac8a128d676 Mon Sep 17 00:00:00 2001 From: lwalew Date: Thu, 23 Jul 2026 14:19:42 +0200 Subject: [PATCH 1/3] feat: add protein sampling benchmark # Conflicts: # ml_peg/calcs/utils/mlipaudit.py --- .../user_guide/benchmarks/biomolecules.rst | 48 +++++ docs/source/user_guide/benchmarks/index.rst | 1 + .../analyse_protein_sampling.py | 203 ++++++++++++++++++ .../biomolecules/protein_sampling/metrics.yml | 22 ++ ml_peg/app/biomolecules/biomolecules.yml | 3 + .../protein_sampling/app_protein_sampling.py | 49 +++++ .../protein_sampling/calc_protein_sampling.py | 72 +++++++ ml_peg/calcs/utils/mlipaudit.py | 14 ++ pyproject.toml | 3 + 9 files changed, 415 insertions(+) create mode 100644 docs/source/user_guide/benchmarks/biomolecules.rst create mode 100644 ml_peg/analysis/biomolecules/protein_sampling/analyse_protein_sampling.py create mode 100644 ml_peg/analysis/biomolecules/protein_sampling/metrics.yml create mode 100644 ml_peg/app/biomolecules/biomolecules.yml create mode 100644 ml_peg/app/biomolecules/protein_sampling/app_protein_sampling.py create mode 100644 ml_peg/calcs/biomolecules/protein_sampling/calc_protein_sampling.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..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 af7bb5976..571b1f702 100644 --- a/docs/source/user_guide/benchmarks/index.rst +++ b/docs/source/user_guide/benchmarks/index.rst @@ -19,4 +19,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/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 index e725c3f88..38c8fd131 100644 --- a/ml_peg/calcs/utils/mlipaudit.py +++ b/ml_peg/calcs/utils/mlipaudit.py @@ -5,6 +5,7 @@ from mlipaudit.benchmarks.conformer_selection.conformer_selection import ( ConformerSelectionBenchmark, ) +from mlipaudit.benchmarks.sampling.sampling import SamplingBenchmark from mlipaudit.benchmarks.tautomers.tautomers import TautomersBenchmark @@ -19,6 +20,19 @@ class MlPegConformerSelectionBenchmark(ConformerSelectionBenchmark): skip_if_elements_missing = False +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 + + 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 9929df6c03e7eb0a3c092139150a2e4d73544a61 Mon Sep 17 00:00:00 2001 From: lwalew Date: Mon, 3 Aug 2026 12:17:40 +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) --- .../biomolecules/protein_sampling/analyse_protein_sampling.py | 4 +++- .../biomolecules/protein_sampling/calc_protein_sampling.py | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/ml_peg/analysis/biomolecules/protein_sampling/analyse_protein_sampling.py b/ml_peg/analysis/biomolecules/protein_sampling/analyse_protein_sampling.py index 733362eca..c474b8b09 100644 --- a/ml_peg/analysis/biomolecules/protein_sampling/analyse_protein_sampling.py +++ b/ml_peg/analysis/biomolecules/protein_sampling/analyse_protein_sampling.py @@ -6,9 +6,11 @@ from pathlib import Path from ase.calculators.calculator import Calculator -from mlipaudit.io import load_model_output_from_disk import pytest +pytest.importorskip("mlipaudit", reason="Please install `mlipaudit` extra") +from mlipaudit.io import load_model_output_from_disk + from ml_peg.analysis.utils.decorators import build_table from ml_peg.analysis.utils.utils import ( build_dispersion_name_map, diff --git a/ml_peg/calcs/biomolecules/protein_sampling/calc_protein_sampling.py b/ml_peg/calcs/biomolecules/protein_sampling/calc_protein_sampling.py index aea49739c..9f1d22b8f 100644 --- a/ml_peg/calcs/biomolecules/protein_sampling/calc_protein_sampling.py +++ b/ml_peg/calcs/biomolecules/protein_sampling/calc_protein_sampling.py @@ -12,12 +12,14 @@ from typing import Any from warnings import warn +import pytest + +pytest.importorskip("mlipaudit", reason="Please install `mlipaudit` extra") 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 423e604500342a3a320fd77ae8c2b36a9b7b1d51 Mon Sep 17 00:00:00 2001 From: lwalew Date: Mon, 3 Aug 2026 12:55:21 +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_sampling.py | 113 +++++++++++++----- .../protein_sampling/calc_protein_sampling.py | 6 + 2 files changed, 91 insertions(+), 28 deletions(-) diff --git a/ml_peg/analysis/biomolecules/protein_sampling/analyse_protein_sampling.py b/ml_peg/analysis/biomolecules/protein_sampling/analyse_protein_sampling.py index c474b8b09..4eda1c247 100644 --- a/ml_peg/analysis/biomolecules/protein_sampling/analyse_protein_sampling.py +++ b/ml_peg/analysis/biomolecules/protein_sampling/analyse_protein_sampling.py @@ -6,9 +6,12 @@ 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.sampling.sampling import STRUCTURE_NAMES from mlipaudit.io import load_model_output_from_disk from ml_peg.analysis.utils.decorators import build_table @@ -19,7 +22,6 @@ 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 @@ -37,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 sampling 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_sampling/protein_sampling.zip", - filename="protein_sampling.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 @@ -62,7 +86,7 @@ def analyze_results() -> dict: dict Mapping of model name to its ``SamplingResult``. """ - data_input_dir = _data_input_dir() + check_dataset() results = {} for model_name in MODELS: @@ -71,7 +95,7 @@ def analyze_results() -> dict: continue benchmark = MlPegSamplingBenchmark( 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( @@ -82,17 +106,38 @@ def analyze_results() -> dict: @pytest.fixture -def struct_info() -> None: - """Write the combined element set to ``info.json`` for filtering.""" - elements = sorted(MlPegSamplingBenchmark.required_elements) +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({"elements": elements}, f, indent=1) + json.dump(info, f, indent=1) + + return info @pytest.fixture -def get_rmsd_backbone(analyze_results) -> dict[str, float | None]: +def get_rmsd_backbone(analyze_results) -> dict[str, float]: """ Get the mean backbone dihedral distribution RMSD for each model. @@ -103,17 +148,21 @@ def get_rmsd_backbone(analyze_results) -> dict[str, float | None]: Returns ------- - dict[str, float | None] + dict[str, float] Backbone dihedral distribution RMSD averaged over residues and systems. """ return { - model_name: result.rmsd_backbone_total + model_name: ( + result.rmsd_backbone_total + if result.rmsd_backbone_total is not None + else np.nan + ) for model_name, result in analyze_results.items() } @pytest.fixture -def get_hellinger_backbone(analyze_results) -> dict[str, float | None]: +def get_hellinger_backbone(analyze_results) -> dict[str, float]: """ Get the mean backbone dihedral Hellinger distance for each model. @@ -124,17 +173,21 @@ def get_hellinger_backbone(analyze_results) -> dict[str, float | None]: Returns ------- - dict[str, float | None] + dict[str, float] Backbone dihedral Hellinger distance averaged over residues and systems. """ return { - model_name: result.hellinger_distance_backbone_total + model_name: ( + result.hellinger_distance_backbone_total + if result.hellinger_distance_backbone_total is not None + else np.nan + ) for model_name, result in analyze_results.items() } @pytest.fixture -def get_outliers_ratio_backbone(analyze_results) -> dict[str, float | None]: +def get_outliers_ratio_backbone(analyze_results) -> dict[str, float]: """ Get the mean backbone dihedral outliers ratio for each model. @@ -145,12 +198,16 @@ def get_outliers_ratio_backbone(analyze_results) -> dict[str, float | None]: Returns ------- - dict[str, float | None] + dict[str, float] Fraction of sampled backbone dihedrals lying far from the reference data, averaged over residues and systems. """ return { - model_name: result.outliers_ratio_backbone_total + model_name: ( + result.outliers_ratio_backbone_total + if result.outliers_ratio_backbone_total is not None + else np.nan + ) for model_name, result in analyze_results.items() } @@ -164,9 +221,9 @@ def get_outliers_ratio_backbone(analyze_results) -> dict[str, float | None]: 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], + get_rmsd_backbone: dict[str, float], + get_hellinger_backbone: dict[str, float], + get_outliers_ratio_backbone: dict[str, float], ) -> dict[str, dict]: """ Get all metrics. @@ -192,7 +249,7 @@ def metrics( } -def test_protein_sampling(metrics: dict[str, dict], struct_info: None) -> None: +def test_protein_sampling(metrics: dict[str, dict], struct_info: dict) -> None: """ Run protein sampling analysis. @@ -200,6 +257,6 @@ def test_protein_sampling(metrics: dict[str, dict], struct_info: None) -> None: ---------- metrics : dict[str, dict] Protein sampling 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_sampling/calc_protein_sampling.py b/ml_peg/calcs/biomolecules/protein_sampling/calc_protein_sampling.py index 9f1d22b8f..17c9e2de9 100644 --- a/ml_peg/calcs/biomolecules/protein_sampling/calc_protein_sampling.py +++ b/ml_peg/calcs/biomolecules/protein_sampling/calc_protein_sampling.py @@ -9,6 +9,7 @@ from __future__ import annotations from pathlib import Path +import shutil from typing import Any from warnings import warn @@ -50,6 +51,11 @@ def test_protein_sampling(mlip: tuple[str, Any]) -> None: filename="protein_sampling.zip", ) + # Save the input data to the calculation outputs so the analysis is self + # contained and does not need to download it again. + name = MlPegSamplingBenchmark.name + shutil.copytree(data_input_dir / name, OUT_PATH / name, dirs_exist_ok=True) + benchmark = MlPegSamplingBenchmark( force_field=calc, data_input_dir=data_input_dir,