diff --git a/docs/source/user_guide/benchmarks/molecular_dynamics.rst b/docs/source/user_guide/benchmarks/molecular_dynamics.rst index 55f2c0d49..d3b4af660 100644 --- a/docs/source/user_guide/benchmarks/molecular_dynamics.rst +++ b/docs/source/user_guide/benchmarks/molecular_dynamics.rst @@ -87,7 +87,6 @@ Summary Benchmark of the density of water-ethanol mixtures for different concentrations of ethanol, compare to experiment. 1 ns of NPT MD on about 120 water/ethanol molecules for 6 concentrations. - Metrics ------- @@ -116,3 +115,45 @@ Packmol generated Reference data: * M. Southard and D. Green, Perry’s Chemical Engineers’ Handbook, 9th Edition. McGraw-Hill Education, 2018. * Experimental + + +Ring planarity +============== + +Summary +------- + +Performance in maintaining planar aromatic rings during molecular dynamics of small +organic molecules. For each molecule, an NVT molecular dynamics simulation is run at 300 K +starting from a QM-optimised reference geometry (selected from QM9), and the deviation of +the ring atoms from their best-fit plane is measured along the trajectory. + +Metrics +------- + +1. Planarity deviation + +At each frame of the trajectory, the ring atoms are fitted to a plane and the root mean +square deviation of the atoms from that plane is calculated. This is averaged over the +trajectory and across all molecules. Aromatic rings are planar, so a well behaved potential +keeps this deviation small; a lower deviation is better. + +A histogram shows the distribution of the sampled planarity 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 reference geometries of the aromatic molecules. diff --git a/ml_peg/analysis/molecular_dynamics/ring_planarity/analyse_ring_planarity.py b/ml_peg/analysis/molecular_dynamics/ring_planarity/analyse_ring_planarity.py new file mode 100644 index 000000000..af3506e3f --- /dev/null +++ b/ml_peg/analysis/molecular_dynamics/ring_planarity/analyse_ring_planarity.py @@ -0,0 +1,223 @@ +"""Analyse the aromatic ring planarity benchmark.""" + +from __future__ import annotations + +import json +from pathlib import Path + +from ase.calculators.calculator import Calculator +import numpy as np +import pytest + +pytest.importorskip("mlipaudit", reason="Please install `mlipaudit` extra") +from mlipaudit.benchmarks.ring_planarity.ring_planarity import RING_PLANARITY_DATASET +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, + load_metrics_config, +) +from ml_peg.app import APP_ROOT +from ml_peg.calcs import CALCS_ROOT +from ml_peg.calcs.utils.mlipaudit import MlPegRingPlanarityBenchmark +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 = MlPegRingPlanarityBenchmark.name + +CALC_PATH = CALCS_ROOT / "molecular_dynamics" / "ring_planarity" / "outputs" +OUT_PATH = APP_ROOT / "data" / "molecular_dynamics" / "ring_planarity" + +METRICS_CONFIG_PATH = Path(__file__).with_name("metrics.yml") +DEFAULT_THRESHOLDS, DEFAULT_TOOLTIPS, DEFAULT_WEIGHTS = load_metrics_config( + METRICS_CONFIG_PATH +) + + +def check_dataset() -> None: + """ + Check the dataset saved by the calculation is available. + + 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. + """ + dataset_path = CALC_PATH / BENCHMARK / RING_PLANARITY_DATASET + if not dataset_path.exists(): + raise ValueError(f"{dataset_path} does not exist. Please run the calculation.") + + +@pytest.fixture +def analyze_results() -> dict: + """ + Run the mlipaudit analysis for each model. + + Returns + ------- + dict + Mapping of model name to its ``RingPlanarityResult``. + """ + check_dataset() + + results = {} + for model_name in MODELS: + output_dir = CALC_PATH / model_name / BENCHMARK + if not (output_dir / "model_output.zip").exists(): + continue + benchmark = MlPegRingPlanarityBenchmark( + force_field=Calculator(), + data_input_dir=CALC_PATH, + run_mode="standard", + ) + benchmark.model_output = load_model_output_from_disk( + CALC_PATH / model_name, MlPegRingPlanarityBenchmark + ) + results[model_name] = benchmark.analyze() + return results + + +@pytest.fixture +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 = MlPegRingPlanarityBenchmark( + force_field=Calculator(), + data_input_dir=CALC_PATH, + run_mode="standard", + ) + data = benchmark._qm9_structures + 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(info, f, indent=1) + + return info + + +@pytest.fixture +@plot_hist( + filename=str(OUT_PATH / "figure_ring_planarity_hist.json"), + title="Ring planarity deviation distribution", + x_label="Planarity deviation / Å", + y_label="Probability density", + bins=50, +) +def deviation_distributions(analyze_results) -> dict[str, np.ndarray]: + """ + Collect the planarity deviations sampled along each model's trajectories. + + Parameters + ---------- + analyze_results + Mapping of model name to its ``RingPlanarityResult``. + + Returns + ------- + dict[str, np.ndarray] + Per-model flat array of ring planarity 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_mae_deviation(analyze_results) -> dict[str, float]: + """ + Get the mean planarity deviation for each model. + + Parameters + ---------- + analyze_results + Mapping of model name to its ``RingPlanarityResult``. + + Returns + ------- + dict[str, float] + Mean planarity deviation of the ring atoms over the trajectories, in Angstrom. + """ + return { + model_name: ( + result.mae_deviation if result.mae_deviation is not None else np.nan + ) + for model_name, result in analyze_results.items() + } + + +@pytest.fixture +@build_table( + filename=OUT_PATH / "ring_planarity_metrics_table.json", + metric_tooltips=DEFAULT_TOOLTIPS, + thresholds=DEFAULT_THRESHOLDS, + weights=DEFAULT_WEIGHTS, + mlip_name_map=DISPERSION_NAME_MAP, +) +def metrics( + deviation_distributions, + get_mae_deviation: dict[str, float], +) -> dict[str, dict]: + """ + Get all metrics. + + Parameters + ---------- + deviation_distributions + Per-model deviation arrays (triggers the histogram plot). + get_mae_deviation + Mean planarity deviations for all models. + + Returns + ------- + dict[str, dict] + Metric names and values for all models. + """ + return { + "Planarity Deviation": get_mae_deviation, + } + + +def test_ring_planarity(metrics: dict[str, dict], struct_info: dict) -> None: + """ + Run ring planarity analysis. + + Parameters + ---------- + metrics : dict[str, dict] + Ring planarity metric results provided by fixtures. + struct_info : dict + Element info written to ``info.json`` for filtering. + """ diff --git a/ml_peg/analysis/molecular_dynamics/ring_planarity/metrics.yml b/ml_peg/analysis/molecular_dynamics/ring_planarity/metrics.yml new file mode 100644 index 000000000..e8e694e93 --- /dev/null +++ b/ml_peg/analysis/molecular_dynamics/ring_planarity/metrics.yml @@ -0,0 +1,8 @@ +metrics: + Planarity Deviation: + good: 0.0 + bad: 0.05 + unit: Å + weight: 1 + tooltip: Mean RMSD of the ring atoms from their best-fit plane, averaged over the MD trajectory and across all molecules. + level_of_theory: DFT diff --git a/ml_peg/app/molecular_dynamics/ring_planarity/app_ring_planarity.py b/ml_peg/app/molecular_dynamics/ring_planarity/app_ring_planarity.py new file mode 100644 index 000000000..c1e30ae14 --- /dev/null +++ b/ml_peg/app/molecular_dynamics/ring_planarity/app_ring_planarity.py @@ -0,0 +1,66 @@ +"""Run ring planarity 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 = "RingPlanarity" +DOCS_URL = "https://ddmms.github.io/ml-peg/user_guide/benchmarks/molecular_dynamics.html#ring-planarity" +DATA_PATH = APP_ROOT / "data" / "molecular_dynamics" / "ring_planarity" + + +class RingPlanarityApp(BaseApp): + """Ring planarity benchmark app layout and callbacks.""" + + def register_callbacks(self) -> None: + """Register callbacks to app.""" + histogram = read_plot( + DATA_PATH / "figure_ring_planarity_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={"Planarity Deviation": histogram}, + ) + + +def get_app() -> RingPlanarityApp: + """ + Get ring planarity benchmark app layout and callback registration. + + Returns + ------- + RingPlanarityApp + Benchmark layout and callback registration. + """ + return RingPlanarityApp( + name="Ring Planarity", + framework_ids="mlip_audit", + description=( + "Performance in maintaining planar aromatic rings during molecular " + "dynamics of small organic molecules. Reference geometries are taken " + "from QM-optimised structures." + ), + docs_url=DOCS_URL, + table_path=DATA_PATH / "ring_planarity_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/ring_planarity/calc_ring_planarity.py b/ml_peg/calcs/molecular_dynamics/ring_planarity/calc_ring_planarity.py new file mode 100644 index 000000000..2526a419a --- /dev/null +++ b/ml_peg/calcs/molecular_dynamics/ring_planarity/calc_ring_planarity.py @@ -0,0 +1,78 @@ +""" +Measure the planarity of aromatic rings during molecular dynamics. + +A molecular dynamics simulation is run for each of a set of small organic +molecules with aromatic rings, and the deviation of the ring atoms from a +perfect plane is measured over the trajectory. +""" + +from __future__ import annotations + +from pathlib import Path +import shutil +from typing import Any +from warnings import warn + +import pytest + +pytest.importorskip("mlipaudit", reason="Please install `mlipaudit` extra") +from mlipaudit.benchmarks.ring_planarity.ring_planarity import ( + RING_PLANARITY_DATASET, + RingPlanarityModelOutput, +) +from mlipaudit.io import write_model_output_to_disk + +from ml_peg.calcs.utils.mlipaudit import MlPegRingPlanarityBenchmark +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_ring_planarity(mlip: tuple[str, Any]) -> None: + """ + Benchmark aromatic ring planarity 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/ring_planarity/ring_planarity.zip", + filename="ring_planarity.zip", + ) + + dataset_dir = OUT_PATH / MlPegRingPlanarityBenchmark.name + dataset_dir.mkdir(parents=True, exist_ok=True) + shutil.copy( + data_input_dir / MlPegRingPlanarityBenchmark.name / RING_PLANARITY_DATASET, + dataset_dir, + ) + + benchmark = MlPegRingPlanarityBenchmark( + force_field=calc, + data_input_dir=data_input_dir, + run_mode="standard", + ) + try: + benchmark.run_model() + except Exception as exc: + warn( + f"Error running ring planarity benchmark for {model_name}: {exc}", + stacklevel=2, + ) + # An empty set of molecules is treated as a failed benchmark by analyze(). + benchmark.model_output = RingPlanarityModelOutput(molecules=[]) + + write_model_output_to_disk( + "ring_planarity", 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..141cc8160 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.ring_planarity.ring_planarity import RingPlanarityBenchmark from mlipaudit.benchmarks.tautomers.tautomers import TautomersBenchmark @@ -19,6 +20,19 @@ class MlPegConformerSelectionBenchmark(ConformerSelectionBenchmark): skip_if_elements_missing = False +class MlPegRingPlanarityBenchmark(RingPlanarityBenchmark): + """ + ``RingPlanarityBenchmark`` 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'", ]