From 1fefd7111ebdc7af1944abc5f86b4a2fb783752a Mon Sep 17 00:00:00 2001 From: lwalew Date: Fri, 3 Jul 2026 17:31:27 +0200 Subject: [PATCH 1/6] feat: add bond length distribution benchmark Migrate the bond length distribution benchmark from the MLIP Audit suite. An NVT molecular dynamics simulation is run for each of a set of small organic molecules, and the deviation of a tracked covalent bond from its QM-optimised reference length is measured over the trajectory. Lower average deviation is better. Adds the calculation, analysis (histogram of deviations plus a metrics table), Dash app, documentation, and the mlipaudit dependency wiring (optional extra, git source and MLIP Audit framework badge). # Conflicts: # docs/source/user_guide/benchmarks/molecular_dynamics.rst # ml_peg/app/utils/frameworks.yml # ml_peg/calcs/utils/mlipaudit.py --- .../benchmarks/molecular_dynamics.rst | 43 ++++ .../analyse_bond_length_distribution.py | 200 ++++++++++++++++++ .../bond_length_distribution/metrics.yml | 8 + .../app_bond_length_distribution.py | 66 ++++++ .../calc_bond_length_distribution.py | 67 ++++++ ml_peg/calcs/utils/mlipaudit.py | 16 ++ pyproject.toml | 3 + 7 files changed, 403 insertions(+) create mode 100644 ml_peg/analysis/molecular_dynamics/bond_length_distribution/analyse_bond_length_distribution.py create mode 100644 ml_peg/analysis/molecular_dynamics/bond_length_distribution/metrics.yml create mode 100644 ml_peg/app/molecular_dynamics/bond_length_distribution/app_bond_length_distribution.py create mode 100644 ml_peg/calcs/molecular_dynamics/bond_length_distribution/calc_bond_length_distribution.py diff --git a/docs/source/user_guide/benchmarks/molecular_dynamics.rst b/docs/source/user_guide/benchmarks/molecular_dynamics.rst index 55f2c0d49..96f9a59b7 100644 --- a/docs/source/user_guide/benchmarks/molecular_dynamics.rst +++ b/docs/source/user_guide/benchmarks/molecular_dynamics.rst @@ -116,3 +116,46 @@ Packmol generated Reference data: * M. Southard and D. Green, Perry’s Chemical Engineers’ Handbook, 9th Edition. McGraw-Hill Education, 2018. * Experimental + + +Bond length distribution +======================== + +Summary +------- + +Performance in maintaining physically reasonable covalent bond lengths during molecular +dynamics of small organic molecules. For each of a set of molecules covering the C-C, C=C, +C#C, C-N, C-O, C=O and C-F bond types, an NVT molecular dynamics simulation is run at 300 K +starting from a QM-optimised reference geometry, and the deviation of a tracked bond from +its reference length is measured along the trajectory. + +Metrics +------- + +1. Bond length deviation + +The length of the tracked bond is measured at each frame of the trajectory, and its absolute +deviation from the reference bond length is averaged over the trajectory and across all +molecules. A well behaved potential keeps bonds close to their reference length, so a lower +deviation is better. + +A histogram shows the distribution of the sampled bond length deviations for each model. + +Computational cost +------------------ + +High: one MD simulation per molecule, each 1,000,000 steps. Faster inference can be achieved +using the jax-accelerated simulations in MLIP Audit directly. + +Data availability +----------------- + +Input structures: + +* MLIP Audit benchmark suite, InstaDeep. Reference geometries selected from the QM9 dataset + (Ramakrishnan et al., Scientific Data 1, 140022, 2014). + +Reference data: + +* QM-optimised equilibrium bond lengths of the reference geometries. diff --git a/ml_peg/analysis/molecular_dynamics/bond_length_distribution/analyse_bond_length_distribution.py b/ml_peg/analysis/molecular_dynamics/bond_length_distribution/analyse_bond_length_distribution.py new file mode 100644 index 000000000..196f8af94 --- /dev/null +++ b/ml_peg/analysis/molecular_dynamics/bond_length_distribution/analyse_bond_length_distribution.py @@ -0,0 +1,200 @@ +"""Analyse the covalent bond length distribution 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 numpy as np +import pytest + +from ml_peg.analysis.utils.decorators import build_table, plot_hist +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 MlPegBondLengthDistributionBenchmark +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 = MlPegBondLengthDistributionBenchmark.name +DATASET_FILENAME = "bond_length_distribution.json" + +CALC_PATH = CALCS_ROOT / "molecular_dynamics" / "bond_length_distribution" / "outputs" +OUT_PATH = APP_ROOT / "data" / "molecular_dynamics" / "bond_length_distribution" + +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 bond length distribution input data. + """ + return download_s3_data( + key="inputs/molecular_dynamics/bond_length_distribution/bond_length_distribution.zip", + filename="bond_length_distribution.zip", + ) + + +@pytest.fixture +def analyze_results() -> dict: + """ + Run the mlipaudit analysis for each model. + + Returns + ------- + dict + Mapping of model name to its ``BondLengthDistributionResult``. + """ + 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 = MlPegBondLengthDistributionBenchmark( + force_field=Calculator(), + data_input_dir=data_input_dir, + run_mode="standard", + ) + benchmark.model_output = load_model_output_from_disk( + CALC_PATH / model_name, MlPegBondLengthDistributionBenchmark + ) + results[model_name] = benchmark.analyze() + return results + + +@pytest.fixture +def struct_info() -> None: + """Write the combined element set to ``info.json`` for filtering.""" + data_path = _data_input_dir() / BENCHMARK / DATASET_FILENAME + with open(data_path, encoding="utf-8") as f: + data = json.load(f) + + elements = sorted( + {symbol for molecule in data.values() for symbol in molecule["atom_symbols"]} + ) + + 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 +@plot_hist( + filename=str(OUT_PATH / "figure_bond_length_hist.json"), + title="Bond length deviation distribution", + x_label="Bond length deviation / Å", + y_label="Probability density", + bins=50, +) +def deviation_distributions(analyze_results) -> dict[str, np.ndarray]: + """ + Collect the bond length deviations sampled along each model's trajectories. + + Parameters + ---------- + analyze_results + Mapping of model name to its ``BondLengthDistributionResult``. + + Returns + ------- + dict[str, np.ndarray] + Per-model flat array of bond length deviations across all molecules. + """ + results = {} + for model_name, result in analyze_results.items(): + if result.failed: + continue + deviations = [ + value + for molecule in result.molecules + if molecule.deviation_trajectory is not None + for value in molecule.deviation_trajectory + ] + if deviations: + results[model_name] = np.array(deviations) + return results + + +@pytest.fixture +def get_avg_deviation(analyze_results) -> dict[str, float]: + """ + Get the average bond length deviation for each model. + + Parameters + ---------- + analyze_results + Mapping of model name to its ``BondLengthDistributionResult``. + + Returns + ------- + dict[str, float] + Mean absolute bond length deviation over the trajectories, in Angstrom. + """ + return { + model_name: result.avg_deviation + for model_name, result in analyze_results.items() + } + + +@pytest.fixture +@build_table( + filename=OUT_PATH / "bond_length_distribution_metrics_table.json", + metric_tooltips=DEFAULT_TOOLTIPS, + thresholds=DEFAULT_THRESHOLDS, + weights=DEFAULT_WEIGHTS, + mlip_name_map=DISPERSION_NAME_MAP, +) +def metrics( + deviation_distributions, + get_avg_deviation: dict[str, float], +) -> dict[str, dict]: + """ + Get all metrics. + + Parameters + ---------- + deviation_distributions + Per-model deviation arrays (triggers the histogram plot). + get_avg_deviation + Average bond length deviations for all models. + + Returns + ------- + dict[str, dict] + Metric names and values for all models. + """ + return { + "Bond Length Deviation": get_avg_deviation, + } + + +def test_bond_length_distribution(metrics: dict[str, dict], struct_info: None) -> None: + """ + Run bond length distribution analysis. + + Parameters + ---------- + metrics : dict[str, dict] + Bond length metric results provided by fixtures. + struct_info : None + Element info written to ``info.json`` for filtering. + """ diff --git a/ml_peg/analysis/molecular_dynamics/bond_length_distribution/metrics.yml b/ml_peg/analysis/molecular_dynamics/bond_length_distribution/metrics.yml new file mode 100644 index 000000000..b450a13fa --- /dev/null +++ b/ml_peg/analysis/molecular_dynamics/bond_length_distribution/metrics.yml @@ -0,0 +1,8 @@ +metrics: + Bond Length Deviation: + good: 0.0 + bad: 0.05 + unit: Å + weight: 1 + tooltip: Mean absolute deviation of a tracked covalent bond from its QM-optimised reference length, averaged over the MD trajectory and across all molecules. + level_of_theory: DFT diff --git a/ml_peg/app/molecular_dynamics/bond_length_distribution/app_bond_length_distribution.py b/ml_peg/app/molecular_dynamics/bond_length_distribution/app_bond_length_distribution.py new file mode 100644 index 000000000..faa4dc226 --- /dev/null +++ b/ml_peg/app/molecular_dynamics/bond_length_distribution/app_bond_length_distribution.py @@ -0,0 +1,66 @@ +"""Run bond length distribution 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 = "BondLength" +DOCS_URL = "https://ddmms.github.io/ml-peg/user_guide/benchmarks/molecular_dynamics.html#bond-length-distribution" +DATA_PATH = APP_ROOT / "data" / "molecular_dynamics" / "bond_length_distribution" + + +class BondLengthApp(BaseApp): + """Bond length distribution benchmark app layout and callbacks.""" + + def register_callbacks(self) -> None: + """Register callbacks to app.""" + histogram = read_plot( + DATA_PATH / "figure_bond_length_hist.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={"Bond Length Deviation": histogram}, + ) + + +def get_app() -> BondLengthApp: + """ + Get bond length distribution benchmark app layout and callback registration. + + Returns + ------- + BondLengthApp + Benchmark layout and callback registration. + """ + return BondLengthApp( + name="Bond Length Distribution", + framework_ids="mlip_audit", + description=( + "Performance in maintaining physically reasonable covalent bond " + "lengths during molecular dynamics of small organic molecules. " + "Reference bond lengths are taken from QM-optimised geometries." + ), + docs_url=DOCS_URL, + table_path=DATA_PATH / "bond_length_distribution_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/molecular_dynamics/bond_length_distribution/calc_bond_length_distribution.py b/ml_peg/calcs/molecular_dynamics/bond_length_distribution/calc_bond_length_distribution.py new file mode 100644 index 000000000..a6b4069f9 --- /dev/null +++ b/ml_peg/calcs/molecular_dynamics/bond_length_distribution/calc_bond_length_distribution.py @@ -0,0 +1,67 @@ +""" +Measure covalent bond length deviations during molecular dynamics. + +A molecular dynamics simulation is run for each of a set of small organic +molecules, and the deviation of a tracked covalent bond from its reference +equilibrium length is measured over the trajectory. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any +from warnings import warn + +from mlipaudit.benchmarks.bond_length_distribution.bond_length_distribution import ( + BondLengthDistributionModelOutput, +) +from mlipaudit.io import write_model_output_to_disk +import pytest + +from ml_peg.calcs.utils.mlipaudit import MlPegBondLengthDistributionBenchmark +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_bond_length_distribution(mlip: tuple[str, Any]) -> None: + """ + Benchmark covalent bond length deviations 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/molecular_dynamics/bond_length_distribution/bond_length_distribution.zip", + filename="bond_length_distribution.zip", + ) + + benchmark = MlPegBondLengthDistributionBenchmark( + force_field=calc, + data_input_dir=data_input_dir, + run_mode="standard", + ) + try: + benchmark.run_model() + except Exception as exc: + warn( + f"Error running bond length distribution benchmark for {model_name}: {exc}", + stacklevel=2, + ) + # An empty set of molecules is treated as a failed benchmark by analyze(). + benchmark.model_output = BondLengthDistributionModelOutput(molecules=[]) + + write_model_output_to_disk( + "bond_length_distribution", 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..683955dd3 100644 --- a/ml_peg/calcs/utils/mlipaudit.py +++ b/ml_peg/calcs/utils/mlipaudit.py @@ -2,12 +2,28 @@ from __future__ import annotations + +from mlipaudit.benchmarks.bond_length_distribution.bond_length_distribution import ( + BondLengthDistributionBenchmark, +) from mlipaudit.benchmarks.conformer_selection.conformer_selection import ( ConformerSelectionBenchmark, ) from mlipaudit.benchmarks.tautomers.tautomers import TautomersBenchmark +class MlPegBondLengthDistributionBenchmark(BondLengthDistributionBenchmark): + """ + ``BondLengthDistributionBenchmark`` 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 MlPegConformerSelectionBenchmark(ConformerSelectionBenchmark): """ ConformerSelectionBenchmark 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 544ad8ea7fcd0d96bcf18770c9b9c1feda541d5f Mon Sep 17 00:00:00 2001 From: lwalew Date: Mon, 3 Aug 2026 10:05:32 +0200 Subject: [PATCH 2/6] fix: mlipaudit spurious dependency --- pyproject.toml | 3 --- 1 file changed, 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 79b9ac99c..b09fcdf99 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -71,9 +71,6 @@ 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 2ed74d87501145f46c9e9d68749fd28b9b0a81db Mon Sep 17 00:00:00 2001 From: lwalew Date: Mon, 3 Aug 2026 10:17:54 +0200 Subject: [PATCH 3/6] feat: better imports --- ml_peg/calcs/utils/mlipaudit.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/ml_peg/calcs/utils/mlipaudit.py b/ml_peg/calcs/utils/mlipaudit.py index 683955dd3..6be426784 100644 --- a/ml_peg/calcs/utils/mlipaudit.py +++ b/ml_peg/calcs/utils/mlipaudit.py @@ -2,14 +2,11 @@ from __future__ import annotations - -from mlipaudit.benchmarks.bond_length_distribution.bond_length_distribution import ( +from mlipaudit.benchmarks import ( BondLengthDistributionBenchmark, -) -from mlipaudit.benchmarks.conformer_selection.conformer_selection import ( ConformerSelectionBenchmark, + TautomersBenchmark, ) -from mlipaudit.benchmarks.tautomers.tautomers import TautomersBenchmark class MlPegBondLengthDistributionBenchmark(BondLengthDistributionBenchmark): @@ -24,6 +21,7 @@ class MlPegBondLengthDistributionBenchmark(BondLengthDistributionBenchmark): skip_if_elements_missing = False + class MlPegConformerSelectionBenchmark(ConformerSelectionBenchmark): """ ConformerSelectionBenchmark wired up for ml-peg's ASE calculators. From 37738a3cd279128d95195c03cdd9d5ddeab12ea0 Mon Sep 17 00:00:00 2001 From: lwalew Date: Mon, 3 Aug 2026 10:22:33 +0200 Subject: [PATCH 4/6] feat: skip if mlipaudit not installed --- .../analyse_bond_length_distribution.py | 4 +++- .../bond_length_distribution/calc_bond_length_distribution.py | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/ml_peg/analysis/molecular_dynamics/bond_length_distribution/analyse_bond_length_distribution.py b/ml_peg/analysis/molecular_dynamics/bond_length_distribution/analyse_bond_length_distribution.py index 196f8af94..23bbc7c01 100644 --- a/ml_peg/analysis/molecular_dynamics/bond_length_distribution/analyse_bond_length_distribution.py +++ b/ml_peg/analysis/molecular_dynamics/bond_length_distribution/analyse_bond_length_distribution.py @@ -6,10 +6,12 @@ from pathlib import Path from ase.calculators.calculator import Calculator -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.io import load_model_output_from_disk + from ml_peg.analysis.utils.decorators import build_table, plot_hist from ml_peg.analysis.utils.utils import ( build_dispersion_name_map, diff --git a/ml_peg/calcs/molecular_dynamics/bond_length_distribution/calc_bond_length_distribution.py b/ml_peg/calcs/molecular_dynamics/bond_length_distribution/calc_bond_length_distribution.py index a6b4069f9..2bfe667c1 100644 --- a/ml_peg/calcs/molecular_dynamics/bond_length_distribution/calc_bond_length_distribution.py +++ b/ml_peg/calcs/molecular_dynamics/bond_length_distribution/calc_bond_length_distribution.py @@ -12,11 +12,13 @@ from typing import Any from warnings import warn +import pytest + +pytest.importorskip("mlipaudit", reason="Please install `mlipaudit` extra") from mlipaudit.benchmarks.bond_length_distribution.bond_length_distribution import ( BondLengthDistributionModelOutput, ) from mlipaudit.io import write_model_output_to_disk -import pytest from ml_peg.calcs.utils.mlipaudit import MlPegBondLengthDistributionBenchmark from ml_peg.calcs.utils.utils import download_s3_data From 5412964d52d1a4f5f1f8df8cfa0c7dd5445f6419 Mon Sep 17 00:00:00 2001 From: lwalew Date: Mon, 3 Aug 2026 11:36:02 +0200 Subject: [PATCH 5/6] feat: set calculator precision to `high` --- .../bond_length_distribution/calc_bond_length_distribution.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ml_peg/calcs/molecular_dynamics/bond_length_distribution/calc_bond_length_distribution.py b/ml_peg/calcs/molecular_dynamics/bond_length_distribution/calc_bond_length_distribution.py index 2bfe667c1..fad59a50b 100644 --- a/ml_peg/calcs/molecular_dynamics/bond_length_distribution/calc_bond_length_distribution.py +++ b/ml_peg/calcs/molecular_dynamics/bond_length_distribution/calc_bond_length_distribution.py @@ -41,7 +41,7 @@ def test_bond_length_distribution(mlip: tuple[str, Any]) -> None: Name of model and model object to get calculator. """ model_name, model = mlip - calc = model.get_calculator() + calc = model.get_calculator(precision="high") calc = model.add_d3_calculator(calc) data_input_dir = download_s3_data( From afc279c93e8379e739b7cf1698d365ad462fd298 Mon Sep 17 00:00:00 2001 From: lwalew Date: Mon, 3 Aug 2026 12:47:16 +0200 Subject: [PATCH 6/6] feat: reuse calc output and filter per molecule Save the downloaded dataset to the calculation outputs and read it from there during analysis, so the analysis no longer re-downloads the input data from S3. Raise a clear error if the calculation has not been run. Store elements as one list per molecule, in dataset order, so individual molecules can be excluded once partial filtering is supported. Report a failed model as NaN rather than None, and add a GPU cost estimate to the docs. Co-Authored-By: Claude Opus 5 (1M context) --- .../benchmarks/molecular_dynamics.rst | 7 +- .../analyse_bond_length_distribution.py | 75 ++++++++++++------- .../calc_bond_length_distribution.py | 11 +++ 3 files changed, 65 insertions(+), 28 deletions(-) diff --git a/docs/source/user_guide/benchmarks/molecular_dynamics.rst b/docs/source/user_guide/benchmarks/molecular_dynamics.rst index 96f9a59b7..d12855ab4 100644 --- a/docs/source/user_guide/benchmarks/molecular_dynamics.rst +++ b/docs/source/user_guide/benchmarks/molecular_dynamics.rst @@ -145,8 +145,11 @@ A histogram shows the distribution of the sampled bond length deviations for eac Computational cost ------------------ -High: one MD simulation per molecule, each 1,000,000 steps. Faster inference can be achieved -using the jax-accelerated simulations in MLIP Audit directly. +High: 8 molecules of 4-13 atoms, one MD simulation each, of 1,000,000 steps, i.e. 1 ns at a +1 fs timestep. The molecules are small, so the cost per step is dominated by per-call +overhead rather than by system size, and tests are likely to take a couple of hours per +model on GPU. Faster inference can be achieved using the jax-accelerated simulations in +MLIP Audit directly. Data availability ----------------- diff --git a/ml_peg/analysis/molecular_dynamics/bond_length_distribution/analyse_bond_length_distribution.py b/ml_peg/analysis/molecular_dynamics/bond_length_distribution/analyse_bond_length_distribution.py index 23bbc7c01..b50900ada 100644 --- a/ml_peg/analysis/molecular_dynamics/bond_length_distribution/analyse_bond_length_distribution.py +++ b/ml_peg/analysis/molecular_dynamics/bond_length_distribution/analyse_bond_length_distribution.py @@ -10,6 +10,9 @@ import pytest pytest.importorskip("mlipaudit", reason="Please install `mlipaudit` extra") +from mlipaudit.benchmarks.bond_length_distribution.bond_length_distribution import ( + BOND_LENGTH_DISTRIBUTION_DATASET_FILENAME, +) from mlipaudit.io import load_model_output_from_disk from ml_peg.analysis.utils.decorators import build_table, plot_hist @@ -20,7 +23,6 @@ from ml_peg.app import APP_ROOT from ml_peg.calcs import CALCS_ROOT from ml_peg.calcs.utils.mlipaudit import MlPegBondLengthDistributionBenchmark -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 @@ -28,7 +30,6 @@ DISPERSION_NAME_MAP = build_dispersion_name_map(MODELS) BENCHMARK = MlPegBondLengthDistributionBenchmark.name -DATASET_FILENAME = "bond_length_distribution.json" CALC_PATH = CALCS_ROOT / "molecular_dynamics" / "bond_length_distribution" / "outputs" OUT_PATH = APP_ROOT / "data" / "molecular_dynamics" / "bond_length_distribution" @@ -39,19 +40,21 @@ ) -def _data_input_dir() -> Path: +def check_dataset() -> None: """ - Download and return the benchmark input data directory. + Check the dataset saved by the calculation is available. - Returns - ------- - Path - Directory containing the extracted bond length distribution input data. + The calculation copies the downloaded dataset into its outputs, so the + analysis does not need to download the input data again. + + Raises + ------ + ValueError + If the dataset is missing from the calculation outputs. """ - return download_s3_data( - key="inputs/molecular_dynamics/bond_length_distribution/bond_length_distribution.zip", - filename="bond_length_distribution.zip", - ) + dataset_path = CALC_PATH / BENCHMARK / BOND_LENGTH_DISTRIBUTION_DATASET_FILENAME + if not dataset_path.exists(): + raise ValueError(f"{dataset_path} does not exist. Please run the calculation.") @pytest.fixture @@ -64,7 +67,7 @@ def analyze_results() -> dict: dict Mapping of model name to its ``BondLengthDistributionResult``. """ - data_input_dir = _data_input_dir() + check_dataset() results = {} for model_name in MODELS: @@ -73,7 +76,7 @@ def analyze_results() -> dict: continue benchmark = MlPegBondLengthDistributionBenchmark( 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,19 +87,37 @@ def analyze_results() -> dict: @pytest.fixture -def struct_info() -> None: - """Write the combined element set to ``info.json`` for filtering.""" - data_path = _data_input_dir() / BENCHMARK / DATASET_FILENAME - with open(data_path, encoding="utf-8") as f: - data = json.load(f) - - elements = sorted( - {symbol for molecule in data.values() for symbol in molecule["atom_symbols"]} +def struct_info() -> dict: + """ + Write per-molecule element info to ``info.json`` for filtering. + + Elements are stored as one list per molecule, so individual molecules can be + excluded once partial filtering is supported. The order follows the dataset, + matching the order of the molecules in ``analyze()``'s results. + + Returns + ------- + dict + Mapping with the per-molecule lists of elements. + """ + check_dataset() + + benchmark = MlPegBondLengthDistributionBenchmark( + force_field=Calculator(), + data_input_dir=CALC_PATH, + run_mode="standard", ) + data = benchmark._bond_length_distribution_data + info = { + "molecules": list(data), + "elements": [sorted(set(molecule.atom_symbols)) for molecule in data.values()], + } 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 @@ -152,7 +173,9 @@ def get_avg_deviation(analyze_results) -> dict[str, float]: Mean absolute bond length deviation over the trajectories, in Angstrom. """ return { - model_name: result.avg_deviation + model_name: ( + result.avg_deviation if result.avg_deviation is not None else np.nan + ) for model_name, result in analyze_results.items() } @@ -189,7 +212,7 @@ def metrics( } -def test_bond_length_distribution(metrics: dict[str, dict], struct_info: None) -> None: +def test_bond_length_distribution(metrics: dict[str, dict], struct_info: dict) -> None: """ Run bond length distribution analysis. @@ -197,6 +220,6 @@ def test_bond_length_distribution(metrics: dict[str, dict], struct_info: None) - ---------- metrics : dict[str, dict] Bond length 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/molecular_dynamics/bond_length_distribution/calc_bond_length_distribution.py b/ml_peg/calcs/molecular_dynamics/bond_length_distribution/calc_bond_length_distribution.py index fad59a50b..d54701e75 100644 --- a/ml_peg/calcs/molecular_dynamics/bond_length_distribution/calc_bond_length_distribution.py +++ b/ml_peg/calcs/molecular_dynamics/bond_length_distribution/calc_bond_length_distribution.py @@ -9,6 +9,7 @@ from __future__ import annotations from pathlib import Path +import shutil from typing import Any from warnings import warn @@ -16,6 +17,7 @@ pytest.importorskip("mlipaudit", reason="Please install `mlipaudit` extra") from mlipaudit.benchmarks.bond_length_distribution.bond_length_distribution import ( + BOND_LENGTH_DISTRIBUTION_DATASET_FILENAME, BondLengthDistributionModelOutput, ) from mlipaudit.io import write_model_output_to_disk @@ -49,6 +51,15 @@ def test_bond_length_distribution(mlip: tuple[str, Any]) -> None: filename="bond_length_distribution.zip", ) + dataset_dir = OUT_PATH / MlPegBondLengthDistributionBenchmark.name + dataset_dir.mkdir(parents=True, exist_ok=True) + shutil.copy( + data_input_dir + / MlPegBondLengthDistributionBenchmark.name + / BOND_LENGTH_DISTRIBUTION_DATASET_FILENAME, + dataset_dir, + ) + benchmark = MlPegBondLengthDistributionBenchmark( force_field=calc, data_input_dir=data_input_dir,