diff --git a/docs/source/user_guide/benchmarks/index.rst b/docs/source/user_guide/benchmarks/index.rst index a5a4e6cd1..0bae54c5a 100644 --- a/docs/source/user_guide/benchmarks/index.rst +++ b/docs/source/user_guide/benchmarks/index.rst @@ -17,4 +17,5 @@ Benchmarks tm_complexes conformers molecular_dynamics + molecular_reactions defect diff --git a/docs/source/user_guide/benchmarks/molecular_reactions.rst b/docs/source/user_guide/benchmarks/molecular_reactions.rst new file mode 100644 index 000000000..89fe97b41 --- /dev/null +++ b/docs/source/user_guide/benchmarks/molecular_reactions.rst @@ -0,0 +1,53 @@ +=================== +Molecular reactions +=================== + +Grambow barrier heights +======================= + +Summary +------- + +Performance in predicting reaction barrier heights for elementary organic +reactions from the Grambow dataset, comprising almost 12,000 reactant, product +and transition state triplets of small neutral molecules containing H, C, N and +O. The benchmark is targeted at barrier heights and reaction energies. + +Metrics +------- + +1. Barrier height MAE + +For each reaction, a single point energy is computed for the reactant, product +and transition state. The activation energy (barrier height) is the transition +state energy minus the reactant energy, and the reaction energy is +the product energy minus the reactant energy. The reported metrics are the mean +absolute errors of the activation energy and the reaction energy against the +reference. + +2. Grambow score + +The overall score is the mlipaudit soft-threshold score combining the +activation energy and reaction energy errors. + +A density scatter plot shows the predicted against reference activation energies +on clicking the barrier height column. + +Computational cost +------------------ + +Medium: around 36,000 single point inference calls (three states per reaction). Minutes on GPU, tens of minutes to hours on CPU. + +Data availability +----------------- + +Input structures: + +* Grambow, C.A., Pattanaik, L. & Green, W.H. Reactants, products, and + transition states of elementary chemical reactions based on quantum + chemistry. Sci Data 7, 137 (2020). DOI: 10.1038/s41597-020-0460-4 + +Reference data: + +* Same as input data +* :math:`\omega B97X-D3/def2-TZVP` level of theory. diff --git a/ml_peg/analysis/molecular_reactions/grambow_barrier_heights/analyse_grambow_barrier_heights.py b/ml_peg/analysis/molecular_reactions/grambow_barrier_heights/analyse_grambow_barrier_heights.py new file mode 100644 index 000000000..ff6e017ef --- /dev/null +++ b/ml_peg/analysis/molecular_reactions/grambow_barrier_heights/analyse_grambow_barrier_heights.py @@ -0,0 +1,251 @@ +""" +Analyse the Grambow reaction barrier benchmark. + +Grambow, C.A., Pattanaik, L. & Green, W.H. +Reactants, products, and transition states of elementary chemical reactions +based on quantum chemistry. +Sci Data 7, 137 (2020). +DOI: 10.1038/s41597-020-0460-4 +""" + +from __future__ import annotations + +import json +from pathlib import Path + +from ase.calculators.calculator import Calculator +import pytest + +pytest.importorskip("mlipaudit", reason="Please install `mlipaudit` extra") +from mlipaudit.benchmarks.reactivity.reactivity import ReactivityModelOutput + +from ml_peg.analysis.utils.decorators import build_table, plot_density_scatter +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 MlPegGrambowBarrierHeightsBenchmark +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_reactions" / "grambow_barrier_heights" / "outputs" +OUT_PATH = APP_ROOT / "data" / "molecular_reactions" / "grambow_barrier_heights" + +METRICS_CONFIG_PATH = Path(__file__).with_name("metrics.yml") +DEFAULT_THRESHOLDS, DEFAULT_TOOLTIPS, DEFAULT_WEIGHTS = load_metrics_config( + METRICS_CONFIG_PATH +) + + +@pytest.fixture +def analyze_results() -> dict: + """ + Run the mlipaudit analysis for each model. + + Returns + ------- + dict + Mapping of model name to its ``ReactivityResult``. + """ + results = {} + for model_name in MODELS: + path = CALC_PATH / model_name / "model_output.json" + if not path.exists(): + continue + benchmark = MlPegGrambowBarrierHeightsBenchmark( + force_field=Calculator(), + data_input_dir=CALC_PATH, + run_mode="standard", + ) + benchmark.model_output = ReactivityModelOutput.model_validate_json( + path.read_text() + ) + results[model_name] = benchmark.analyze() + return results + + +@pytest.fixture +def struct_info() -> dict: + """ + Write the combined element set to ``info.json`` for filtering. + + Returns + ------- + dict + Mapping with the sorted list of elements present in the dataset. + """ + benchmark = MlPegGrambowBarrierHeightsBenchmark( + force_field=Calculator(), + data_input_dir=CALC_PATH, + run_mode="standard", + ) + elements = sorted( + { + symbol + for reaction in benchmark._grambow_data.values() + for molecule in ( + reaction.reactants, + reaction.products, + reaction.transition_state, + ) + for symbol in molecule.atom_symbols + } + ) + info = {"elements": elements} + OUT_PATH.mkdir(parents=True, exist_ok=True) + (OUT_PATH / "info.json").write_text(json.dumps(info, indent=1)) + return info + + +@pytest.fixture +@plot_density_scatter( + filename=OUT_PATH / "figure_barrier_density.json", + title="Reaction barrier density plot", + x_label="Reference activation energy / kcal/mol", + y_label="Predicted activation energy / kcal/mol", + annotation_metadata={"system_count": "Reactions"}, +) +def barrier_density(analyze_results) -> dict[str, dict]: + """ + Density scatter inputs for the reaction barrier heights. + + Parameters + ---------- + analyze_results + Mapping of model name to its ``ReactivityResult``. + + Returns + ------- + dict[str, dict] + Mapping of model name to density-scatter data. + """ + density_inputs: dict[str, dict] = {} + for model_name, result in analyze_results.items(): + reactions = result.reaction_results.values() + ref = [reaction.activation_energy_ref for reaction in reactions] + pred = [reaction.activation_energy_pred for reaction in reactions] + density_inputs[model_name] = { + "ref": ref, + "pred": pred, + "meta": {"system_count": len(pred)}, + } + return density_inputs + + +@pytest.fixture +def get_barrier_mae(analyze_results) -> dict[str, float]: + """ + Get the activation energy (barrier height) MAE for each model. + + Parameters + ---------- + analyze_results + Mapping of model name to its ``ReactivityResult``. + + Returns + ------- + dict[str, float] + Mean absolute activation energy error in kcal/mol for each model. + """ + return { + model_name: result.mae_activation_energy + for model_name, result in analyze_results.items() + } + + +@pytest.fixture +def get_reaction_energy_mae(analyze_results) -> dict[str, float]: + """ + Get the reaction energy (enthalpy) MAE for each model. + + Parameters + ---------- + analyze_results + Mapping of model name to its ``ReactivityResult``. + + Returns + ------- + dict[str, float] + Mean absolute enthalpy of reaction error in kcal/mol for each model. + """ + return { + model_name: result.mae_enthalpy_of_reaction + for model_name, result in analyze_results.items() + } + + +@pytest.fixture +def get_score(analyze_results) -> dict[str, float]: + """ + Get the mlipaudit benchmark score for each model. + + Parameters + ---------- + analyze_results + Mapping of model name to its ``ReactivityResult``. + + Returns + ------- + dict[str, float] + The mlipaudit soft-threshold score (0 to 1) for each model. + """ + return {model_name: result.score for model_name, result in analyze_results.items()} + + +@pytest.fixture +@build_table( + filename=OUT_PATH / "grambow_barrier_heights_metrics_table.json", + metric_tooltips=DEFAULT_TOOLTIPS, + thresholds=DEFAULT_THRESHOLDS, + weights=DEFAULT_WEIGHTS, + mlip_name_map=DISPERSION_NAME_MAP, +) +def metrics( + barrier_density, + get_barrier_mae: dict[str, float], + get_reaction_energy_mae: dict[str, float], + get_score: dict[str, float], +) -> dict[str, dict]: + """ + Get all metrics. + + Parameters + ---------- + barrier_density + Density scatter inputs for the reaction barriers (triggers the plot). + get_barrier_mae + Activation energy MAEs for all models. + get_reaction_energy_mae + Reaction energy MAEs for all models. + get_score + The mlipaudit benchmark scores for all models. + + Returns + ------- + dict[str, dict] + Metric names and values for all models. + """ + return { + "Barrier Height MAE": get_barrier_mae, + "Reaction Energy MAE": get_reaction_energy_mae, + "Grambow Score": get_score, + } + + +def test_grambow_barrier_heights( + metrics: dict[str, dict], barrier_density: dict[str, dict], struct_info: dict +) -> None: + """ + Run Grambow barrier heights analysis. + + Parameters + ---------- + metrics : dict[str, dict] + Grambow barrier heights metric results provided by fixtures. + barrier_density : dict[str, dict] + Density scatter inputs for the reaction barriers. + struct_info : dict + Element info written to ``info.json`` for filtering. + """ diff --git a/ml_peg/analysis/molecular_reactions/grambow_barrier_heights/metrics.yml b/ml_peg/analysis/molecular_reactions/grambow_barrier_heights/metrics.yml new file mode 100644 index 000000000..514e02abf --- /dev/null +++ b/ml_peg/analysis/molecular_reactions/grambow_barrier_heights/metrics.yml @@ -0,0 +1,22 @@ +metrics: + Barrier Height MAE: + good: 0.0 + bad: 2.0 + unit: kcal/mol + weight: 0 + tooltip: Mean absolute error of the reaction activation energy (barrier height). + level_of_theory: wB97X-D3/def2-TZVP + Reaction Energy MAE: + good: 0.0 + bad: 2.0 + unit: kcal/mol + weight: 0 + tooltip: Mean absolute error of the enthalpy of reaction. + level_of_theory: wB97X-D3/def2-TZVP + Grambow Score: + good: 1.0 + bad: 0.0 + unit: null + weight: 1 + tooltip: mlipaudit soft-threshold score combining activation energy and reaction energy errors (failed reactions score 0). + level_of_theory: wB97X-D3/def2-TZVP diff --git a/ml_peg/app/molecular_reactions/grambow_barrier_heights/app_grambow_barrier_heights.py b/ml_peg/app/molecular_reactions/grambow_barrier_heights/app_grambow_barrier_heights.py new file mode 100644 index 000000000..6c0451ca9 --- /dev/null +++ b/ml_peg/app/molecular_reactions/grambow_barrier_heights/app_grambow_barrier_heights.py @@ -0,0 +1,78 @@ +"""Run Grambow barrier heights 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_cell +from ml_peg.app.utils.load import read_density_plot_for_model +from ml_peg.models import current_models +from ml_peg.models.get_models import get_model_names + +MODELS = get_model_names(current_models) +BENCHMARK_NAME = "GrambowBarrierHeights" +DOCS_URL = "https://ddmms.github.io/ml-peg/user_guide/benchmarks/molecular_reactions.html#grambow-barrier-heights" +DATA_PATH = APP_ROOT / "data" / "molecular_reactions" / "grambow_barrier_heights" + + +class GrambowBarrierHeightsApp(BaseApp): + """Grambow barrier heights benchmark app layout and callbacks.""" + + def register_callbacks(self) -> None: + """Register callbacks to app.""" + density_plots: dict[str, dict] = {} + for model in MODELS: + graph = read_density_plot_for_model( + filename=DATA_PATH / "figure_barrier_density.json", + model=model, + id=f"{BENCHMARK_NAME}-{model}-barrier-figure", + ) + if graph is not None: + density_plots[model] = {"Barrier Height MAE": graph} + + plot_from_table_cell( + table_id=self.table_id, + plot_id=f"{BENCHMARK_NAME}-figure-placeholder", + cell_to_plot=density_plots, + ) + + +def get_app() -> GrambowBarrierHeightsApp: + """ + Get Grambow barrier heights benchmark app layout and callback registration. + + Returns + ------- + GrambowBarrierHeightsApp + Benchmark layout and callback registration. + """ + return GrambowBarrierHeightsApp( + name="Grambow Barrier Heights", + framework_ids="mlip_audit", + description=( + "Performance in predicting reaction barrier heights (and reaction " + "energies) for elementary organic reactions from the Grambow " + "dataset. Reference data from wB97X-D3/def2-TZVP calculations." + ), + docs_url=DOCS_URL, + table_path=DATA_PATH / "grambow_barrier_heights_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, + suppress_callback_exceptions=True, + ) + 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/app/utils/frameworks.yml b/ml_peg/app/utils/frameworks.yml index 27b079925..4bb28f12a 100644 --- a/ml_peg/app/utils/frameworks.yml +++ b/ml_peg/app/utils/frameworks.yml @@ -11,6 +11,13 @@ 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" + logo: "https://raw.githubusercontent.com/instadeepai/mlipaudit/mlpeg-migration/InstaDeep_Logo.png" + mace-multihead: label: Multihead Cross Learning color: "#7c3aed" diff --git a/ml_peg/calcs/molecular_reactions/grambow_barrier_heights/calc_grambow_barrier_heights.py b/ml_peg/calcs/molecular_reactions/grambow_barrier_heights/calc_grambow_barrier_heights.py new file mode 100644 index 000000000..4d84f1841 --- /dev/null +++ b/ml_peg/calcs/molecular_reactions/grambow_barrier_heights/calc_grambow_barrier_heights.py @@ -0,0 +1,84 @@ +""" +Compute the Grambow reaction barrier dataset. + +Grambow, C.A., Pattanaik, L. & Green, W.H. +Reactants, products, and transition states of elementary chemical reactions +based on quantum chemistry. +Sci Data 7, 137 (2020). +DOI: 10.1038/s41597-020-0460-4 +""" + +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.reactivity.reactivity import ( + GRAMBOW_DATASET_FILENAME, + ReactivityModelOutput, +) + +from ml_peg.calcs.utils.mlipaudit import MlPegGrambowBarrierHeightsBenchmark +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_grambow_barrier_heights(mlip: tuple[str, Any]) -> None: + """ + Benchmark the Grambow reaction barrier dataset. + + Parameters + ---------- + mlip + Name of model and model object to get calculator. + """ + model_name, model = mlip + calc = model.get_calculator(precision="high") + calc = model.add_d3_calculator(calc) + + data_input_dir = download_s3_data( + key="inputs/molecular_reactions/grambow_barrier_heights/grambow_barrier_heights.zip", + filename="grambow_barrier_heights.zip", + ) + + out_path = OUT_PATH / model_name + out_path.mkdir(parents=True, exist_ok=True) + + benchmark_name = MlPegGrambowBarrierHeightsBenchmark.name + dataset_dir = OUT_PATH / benchmark_name + dataset_dir.mkdir(parents=True, exist_ok=True) + shutil.copy( + data_input_dir / benchmark_name / GRAMBOW_DATASET_FILENAME, + dataset_dir, + ) + + benchmark = MlPegGrambowBarrierHeightsBenchmark( + force_field=calc, + data_input_dir=data_input_dir, + run_mode="standard", + ) + try: + benchmark.run_model() + except Exception as exc: + warn( + f"Error running Grambow barrier heights benchmark for {model_name}: {exc}", + stacklevel=2, + ) + benchmark.model_output = ReactivityModelOutput( + reaction_ids=[], energy_predictions=[] + ) + + (out_path / "model_output.json").write_text( + benchmark.model_output.model_dump_json() + ) diff --git a/ml_peg/calcs/utils/mlipaudit.py b/ml_peg/calcs/utils/mlipaudit.py new file mode 100644 index 000000000..2c62b3f11 --- /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.reactivity.reactivity import ReactivityBenchmark + + +class MlPegGrambowBarrierHeightsBenchmark(ReactivityBenchmark): + """ + ReactivityBenchmark 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