diff --git a/docs/source/user_guide/benchmarks/molecular_dynamics.rst b/docs/source/user_guide/benchmarks/molecular_dynamics.rst index 24f3fb3b3..75b601e1d 100644 --- a/docs/source/user_guide/benchmarks/molecular_dynamics.rst +++ b/docs/source/user_guide/benchmarks/molecular_dynamics.rst @@ -76,3 +76,46 @@ Reference data: * Same as input data * 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/app/utils/frameworks.yml b/ml_peg/app/utils/frameworks.yml index 27b079925..386d4ceac 100644 --- a/ml_peg/app/utils/frameworks.yml +++ b/ml_peg/app/utils/frameworks.yml @@ -11,6 +11,12 @@ mlip_arena: url: "https://huggingface.co/spaces/atomind/mlip-arena" logo: "https://huggingface.co/front/assets/huggingface_logo-noborder.svg" +mlip_audit: + label: MLIP Audit + color: "#1d4ed8" + text_color: "#ffffff" + url: "https://github.com/instadeepai/mlipaudit" + mace-multihead: label: Multihead Cross Learning color: "#7c3aed" 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 new file mode 100644 index 000000000..6bb133c5f --- /dev/null +++ b/ml_peg/calcs/utils/mlipaudit.py @@ -0,0 +1,20 @@ +"""Adapters for using mlipaudit benchmarks with ml-peg's ASE calculators.""" + +from __future__ import annotations + +from mlipaudit.benchmarks.bond_length_distribution.bond_length_distribution import ( + BondLengthDistributionBenchmark, +) + + +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 diff --git a/pyproject.toml b/pyproject.toml index 23088ff82..231965742 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -64,6 +64,9 @@ mace = [ mattersim = [ "mattersim==1.2.2", ] +mlipaudit = [ + "mlipaudit; python_version >= '3.11'", +] orb = [ "orb-models == 0.6.2; sys_platform != 'win32' and python_version >= '3.12'", ] @@ -206,6 +209,10 @@ conflicts = [ { extra = "mattersim" }, { extra = "grace" }, ], + [ + { extra = "mlipaudit" }, + { extra = "grace" }, + ], ] constraint-dependencies = [ @@ -217,3 +224,4 @@ module-root = "" [tool.uv.sources] asemolec = { git = "https://github.com/imagdau/aseMolec.git" } +mlipaudit = { git = "https://github.com/instadeepai/MLIPAudit.git", branch = "mlpeg-migration" }