diff --git a/docs/source/user_guide/benchmarks/molecular_dynamics.rst b/docs/source/user_guide/benchmarks/molecular_dynamics.rst index 24f3fb3b3..258e8e42b 100644 --- a/docs/source/user_guide/benchmarks/molecular_dynamics.rst +++ b/docs/source/user_guide/benchmarks/molecular_dynamics.rst @@ -76,3 +76,45 @@ Reference data: * Same as input data * Experimental + + +Stability +========= + +Summary +------- + +Performance in running stable molecular dynamics across a diverse set of +systems: small molecules (containing H/C/N/O, sulfur, and halogens), peptides +in vacuum (neurotensin, PDB: 2LNF; cyclic oxytocin, PDB: 7OFG), a protein in +vacuum (PDB: 1A7M), and solvated peptides with and without counter-ions. + +Metrics +------- + +1. Success rate + +Each system is run with a short NVT molecular dynamics simulation at 300 K. The +fraction of simulations that complete without error is reported. A simulation +counts as failed if it raises an error during integration (for example, a +numerical instability). Higher is better, with 1.0 meaning every system ran to +completion. + +A scatter plot below the table shows, for each system and model, the fraction of +the trajectory completed before it became unstable: 1.0 for a fully stable run, a +smaller value when the run exploded or hydrogens drifted, and 0.0 for systems that +failed to run at all. Individual models can be toggled via the legend. + +Computational cost +------------------ + +High: tests are likely to take many hours on GPU. Faster simulation times can be +achieved using the jax accelerated simulations in MLIP Audit directly. + +Data availability +----------------- + +Input structures: + +* MLIP Audit benchmark suite, InstaDeep. + Structures derived from PDB entries 2LNF, 7OFG, and 1A7M. diff --git a/ml_peg/analysis/molecular_dynamics/stability/analyse_stability.py b/ml_peg/analysis/molecular_dynamics/stability/analyse_stability.py new file mode 100644 index 000000000..e60a40823 --- /dev/null +++ b/ml_peg/analysis/molecular_dynamics/stability/analyse_stability.py @@ -0,0 +1,205 @@ +"""Analyse molecular dynamics stability benchmark.""" + +from __future__ import annotations + +from pathlib import Path + +from ase.calculators.calculator import Calculator +from mlipaudit.benchmarks.stability.stability import STRUCTURE_NAMES, STRUCTURES +from mlipaudit.io import load_model_output_from_disk +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 MlPegStabilityBenchmark +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) + +CALC_PATH = CALCS_ROOT / "molecular_dynamics" / "stability" / "outputs" +OUT_PATH = APP_ROOT / "data" / "molecular_dynamics" / "stability" + +METRICS_CONFIG_PATH = Path(__file__).with_name("metrics.yml") +DEFAULT_THRESHOLDS, DEFAULT_TOOLTIPS, DEFAULT_WEIGHTS = load_metrics_config( + METRICS_CONFIG_PATH +) + + +def _fraction_completed(result) -> float: + """ + Get how far a simulation got before it became unstable. + + Parameters + ---------- + result + A ``StabilityStructureResult`` for a single structure. + + Returns + ------- + float + Fraction of the trajectory completed: 1.0 if stable, 0.0 if the + simulation did not run, otherwise the frame at which it exploded or + drifted divided by the total number of frames. + """ + if result.failed: + return 0.0 + if result.exploded_frame != -1: + return result.exploded_frame / result.num_frames + if result.drift_frame != -1: + return result.drift_frame / result.num_frames + return 1.0 + + +@pytest.fixture +def structure_results() -> dict[str, list]: + """ + Analyse each stored trajectory into per-structure stability results. + + Returns + ------- + dict[str, list] + List of ``StabilityStructureResult`` objects for each model. + """ + data_input_dir = download_s3_data( + key="inputs/molecular_dynamics/stability/stability.zip", + filename="stability.zip", + ) + + results = {} + for model_name in MODELS: + output_dir = CALC_PATH / model_name / MlPegStabilityBenchmark.name + if not (output_dir / "model_output.zip").exists(): + continue + benchmark = MlPegStabilityBenchmark( + force_field=Calculator(), + data_input_dir=data_input_dir, + run_mode="standard", + ) + benchmark.model_output = load_model_output_from_disk( + CALC_PATH / model_name, MlPegStabilityBenchmark + ) + results[model_name] = benchmark.analyze().structure_results + return results + + +@pytest.fixture +def success_rate(structure_results: dict[str, list]) -> dict[str, float]: + """ + Get the fraction of MD simulations that completed without error. + + Parameters + ---------- + structure_results + Per-structure stability results for all models. + + Returns + ------- + dict[str, float] + Fraction of successful simulations for each model. + """ + return { + model_name: sum(not r.failed for r in results) / len(results) + for model_name, results in structure_results.items() + } + + +@pytest.fixture +@plot_scatter( + title="Trajectory stability", + x_label="System", + y_label="Fraction of trajectory completed", + hovertemplate="%{x}
Completed: %{y:.2f}%{fullData.name}", + filename=str(OUT_PATH / "stability_progress.json"), +) +def progress(structure_results: dict[str, list]) -> dict[str, tuple[list, list]]: + """ + Get per-structure trajectory progress for each model. + + For every system, reports the fraction of the trajectory completed before it + exploded, drifted or failed to run. + + Parameters + ---------- + structure_results + Per-structure stability results for all models. + + Returns + ------- + dict[str, tuple[list, list]] + Structure names and completed fractions for each model. + """ + return { + model_name: ( + [r.structure_name for r in results], + [_fraction_completed(r) for r in results], + ) + for model_name, results in structure_results.items() + } + + +@pytest.fixture +@build_table( + filename=OUT_PATH / "stability_metrics_table.json", + metric_tooltips=DEFAULT_TOOLTIPS, + thresholds=DEFAULT_THRESHOLDS, + mlip_name_map=DISPERSION_NAME_MAP, +) +def metrics(success_rate: dict[str, float]) -> dict[str, dict]: + """ + Get all metrics. + + Parameters + ---------- + success_rate + Fraction of successful simulations for all models. + + Returns + ------- + dict[str, dict] + Metric names and values for all models. + """ + return { + "Success Rate": success_rate, + } + + +@pytest.fixture +def element_info() -> None: + """Write element info for all benchmark systems, used by the app element filter.""" + data_input_dir = download_s3_data( + key="inputs/molecular_dynamics/stability/stability.zip", + filename="stability.zip", + ) + xyz_paths = [ + data_input_dir / MlPegStabilityBenchmark.name / STRUCTURES[name]["xyz"] + for name in STRUCTURE_NAMES + ] + write_struct_info(data_path=xyz_paths, out_path=OUT_PATH) + + +def test_stability( + metrics: dict[str, dict], + progress: dict[str, tuple[list, list]], + element_info: None, +) -> None: + """ + Run stability analysis. + + Parameters + ---------- + metrics : dict[str, dict] + Stability metric results provided by fixtures. + progress : dict[str, tuple[list, list]] + Per-structure trajectory progress provided by fixtures. + element_info : None + Element info written for the app element filter. + """ diff --git a/ml_peg/analysis/molecular_dynamics/stability/metrics.yml b/ml_peg/analysis/molecular_dynamics/stability/metrics.yml new file mode 100644 index 000000000..7c9f0542c --- /dev/null +++ b/ml_peg/analysis/molecular_dynamics/stability/metrics.yml @@ -0,0 +1,6 @@ +metrics: + Success Rate: + good: 1.0 + bad: 0.0 + unit: null + tooltip: Fraction of MD simulations that completed without error. diff --git a/ml_peg/analysis/utils/decorators.py b/ml_peg/analysis/utils/decorators.py index a0c3ebe6f..a082ab9ee 100644 --- a/ml_peg/analysis/utils/decorators.py +++ b/ml_peg/analysis/utils/decorators.py @@ -448,6 +448,7 @@ def plot_scatter( show_line: bool = False, show_markers: bool = True, hoverdata: dict | None = None, + hovertemplate: str | None = None, filename: str = "scatter.json", highlight_range: dict = None, ) -> Callable: @@ -468,6 +469,9 @@ def plot_scatter( Whether to show markers on the plot. Default is True. hoverdata Hover data dictionary. Default is `{}`. + hovertemplate + Base hovertemplate string. Overrides the default reference/prediction + template. Any ``hoverdata`` keys are appended to it. Default is `None`. filename Filename to save plot as JSON. Default is "scatter.json". highlight_range @@ -513,11 +517,15 @@ def plot_scatter_wrapper(*args, **kwargs) -> dict[str, Any]: """ results = func(*args, **kwargs) - hovertemplate = "Pred: %{x}
" + "Ref: %{y}
" + template = ( + hovertemplate + if hovertemplate is not None + else "Pred: %{x}
" + "Ref: %{y}
" + ) customdata = [] if hoverdata: for i, key in enumerate(hoverdata): - hovertemplate += f"{key}: %{{customdata[{i}]}}
" + template += f"{key}: %{{customdata[{i}]}}
" customdata = list(zip(*hoverdata.values(), strict=True)) modes = [] @@ -538,7 +546,7 @@ def plot_scatter_wrapper(*args, **kwargs) -> dict[str, Any]: name=name, mode=mode, customdata=customdata, - hovertemplate=hovertemplate, + hovertemplate=template, ) ) diff --git a/ml_peg/app/molecular_dynamics/molecular_dynamics.yml b/ml_peg/app/molecular_dynamics/molecular_dynamics.yml index abcdd622f..5b893e61e 100644 --- a/ml_peg/app/molecular_dynamics/molecular_dynamics.yml +++ b/ml_peg/app/molecular_dynamics/molecular_dynamics.yml @@ -1,3 +1,3 @@ title: Molecular Dynamics -description: Liquid Densities, Water Densities +description: Liquid Densities, Water Densities, Stability weight: 0 diff --git a/ml_peg/app/molecular_dynamics/stability/app_stability.py b/ml_peg/app/molecular_dynamics/stability/app_stability.py new file mode 100644 index 000000000..fcdf6b169 --- /dev/null +++ b/ml_peg/app/molecular_dynamics/stability/app_stability.py @@ -0,0 +1,58 @@ +"""Run molecular dynamics stability app.""" + +from __future__ import annotations + +from dash import Dash + +from ml_peg.app import APP_ROOT +from ml_peg.app.base_app import BaseApp +from ml_peg.app.utils.load import read_plot + +BENCHMARK_NAME = "Stability" +DOCS_URL = "https://ddmms.github.io/ml-peg/user_guide/benchmarks/molecular_dynamics.html#stability" +DATA_PATH = APP_ROOT / "data" / "molecular_dynamics" / "stability" +INFO_PATH = DATA_PATH / "info.json" + + +class StabilityApp(BaseApp): + """Stability benchmark app layout and callbacks.""" + + def register_callbacks(self) -> None: + """Register callbacks to app.""" + + +def get_app() -> StabilityApp: + """ + Get stability benchmark app layout and callback registration. + + Returns + ------- + StabilityApp + Benchmark layout and callback registration. + """ + return StabilityApp( + name=BENCHMARK_NAME, + framework_ids="mlip_audit", + description=( + "Fraction of short molecular dynamics simulations that complete " + "without error, across small molecules, peptides and proteins in " + "vacuum and solvent." + ), + docs_url=DOCS_URL, + table_path=DATA_PATH / "stability_metrics_table.json", + extra_components=[ + read_plot( + filename=DATA_PATH / "stability_progress.json", + id=f"{BENCHMARK_NAME}-progress-figure", + ) + ], + info_path=INFO_PATH, + ) + + +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=8067, debug=True) diff --git a/ml_peg/app/utils/frameworks.yml b/ml_peg/app/utils/frameworks.yml index 27b079925..2d0f5537c 100644 --- a/ml_peg/app/utils/frameworks.yml +++ b/ml_peg/app/utils/frameworks.yml @@ -26,3 +26,10 @@ mace-polar-1: url: "https://arxiv.org/abs/2602.19411" icon: "📄" tooltip: "Benchmark from the MACE-POLAR-1 paper (arXiv:2602.19411) - click to read" + +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" diff --git a/ml_peg/calcs/molecular_dynamics/stability/calc_stability.py b/ml_peg/calcs/molecular_dynamics/stability/calc_stability.py new file mode 100644 index 000000000..be9404b28 --- /dev/null +++ b/ml_peg/calcs/molecular_dynamics/stability/calc_stability.py @@ -0,0 +1,54 @@ +""" +Run molecular dynamics stability simulations. + +Short MD runs for small molecules, peptides and proteins in vacuum and +solvent, checking how many simulations complete without error. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +from mlipaudit.io import write_model_output_to_disk +import pytest + +from ml_peg.calcs.utils.mlipaudit import MlPegStabilityBenchmark +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_stability(mlip: tuple[str, Any]) -> None: + """ + Benchmark molecular dynamics stability. + + Parameters + ---------- + mlip + Name of model and model object to get calculator. + """ + model_name, model = mlip + calc = model.get_calculator(precision="low") + calc = model.get_calculator(precision="low") + + data_input_dir = download_s3_data( + key="inputs/molecular_dynamics/stability/stability.zip", + filename="stability.zip", + ) + + benchmark = MlPegStabilityBenchmark( + force_field=calc, + data_input_dir=data_input_dir, + run_mode="standard", + ) + benchmark.run_model() + + write_model_output_to_disk( + "stability", 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..2772c7014 --- /dev/null +++ b/ml_peg/calcs/utils/mlipaudit.py @@ -0,0 +1,16 @@ +"""Adapters for using mlipaudit benchmarks with ml-peg's ASE calculators.""" + +from __future__ import annotations + +from mlipaudit.benchmarks.stability.stability import StabilityBenchmark + + +class MlPegStabilityBenchmark(StabilityBenchmark): + """ + StabilityBenchmark wired up for ml-peg's ASE calculators. + + ``skip_if_elements_missing`` is disabled because ASE ``Calculator`` objects + do not expose ``allowed_atomic_numbers``. + """ + + skip_if_elements_missing = False