diff --git a/docs/source/user_guide/benchmarks/conformers.rst b/docs/source/user_guide/benchmarks/conformers.rst index 5e9ed6d30..f40a2af99 100644 --- a/docs/source/user_guide/benchmarks/conformers.rst +++ b/docs/source/user_guide/benchmarks/conformers.rst @@ -40,3 +40,48 @@ Reference data: * Same as input data * :math:`PNO-LCCSD(T)-F12/ AVQZ` level of theory: a local, explicitly correlated coupled cluster method. + + +TorsionNet500 +============= + +Summary +------- + +Performance in predicting torsion energy barriers for 500 drug-like organic +molecules, each scanned across a systematically sampled dihedral angle. + +Metrics +------- + +1. Barrier height error + +For each fragment, a single point energy is computed for every conformer along +the dihedral scan. The predicted profile is aligned to the reference at its +minimum, and the barrier height is taken as the difference between the maximum +and minimum energy of the profile. The reported metric is the mean absolute +error of the barrier height across all fragments. The overall score is the +mlipaudit per-fragment soft-threshold score on the barrier height error. + +A parity plot of the predicted against reference barrier heights is shown on +clicking the metric column. + +Computational cost +------------------ + +Medium: 12,000 single point inference calls on small molecular systems. + +Data availability +----------------- + +Input structures: + +* TorsionNet: A Deep Neural Network to Rapidly Predict Small-Molecule + Torsional Energy Profiles with the Accuracy of Quantum Mechanics. + Rai, B. K. et al. J. Chem. Inf. Model. 2022, 62 (4), 785-800. + DOI: 10.1021/acs.jcim.1c01346 + +Reference data: + +* Recomputed by InstaDeep at the :math:`\omega B97M-D3(BJ)/def2-TZVPPD` + level of theory (MLIP Audit benchmark suite). diff --git a/ml_peg/analysis/conformers/TorsionNet500/analyse_TorsionNet500.py b/ml_peg/analysis/conformers/TorsionNet500/analyse_TorsionNet500.py new file mode 100644 index 000000000..f3a5676ce --- /dev/null +++ b/ml_peg/analysis/conformers/TorsionNet500/analyse_TorsionNet500.py @@ -0,0 +1,443 @@ +"""Analyse TorsionNet500 dihedral scan benchmark.""" + +from __future__ import annotations + +import json +from pathlib import Path + +from ase import Atoms +from ase.calculators.calculator import Calculator +from ase.io import write +from mlipaudit.benchmarks.dihedral_scan.dihedral_scan import DihedralScanModelOutput +import numpy as np +import plotly.graph_objects as go +import pytest + +from ml_peg.analysis.utils.decorators import build_table +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 MlPegDihedralScanBenchmark +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 / "conformers" / "TorsionNet500" / "outputs" +OUT_PATH = APP_ROOT / "data" / "conformers" / "TorsionNet500" + +METRICS_CONFIG_PATH = Path(__file__).with_name("metrics.yml") +DEFAULT_THRESHOLDS, DEFAULT_TOOLTIPS, DEFAULT_WEIGHTS = load_metrics_config( + METRICS_CONFIG_PATH +) + + +def labels() -> list: + """ + Get the ordered list of fragment names. + + Returns + ------- + list + List of all dihedral scan fragment names. + """ + for model_name in MODELS: + path = CALC_PATH / model_name / "model_output.json" + if path.exists(): + output = DihedralScanModelOutput.model_validate_json(path.read_text()) + return sorted(fragment.fragment_name for fragment in output.fragments) + return [] + + +@pytest.fixture +def analyze_results() -> dict: + """ + Run the mlipaudit analysis for each model. + + Returns + ------- + dict + Mapping of model name to its ``DihedralScanResult``. + """ + data_input_dir = download_s3_data( + key="inputs/conformers/TorsionNet500/TorsionNet500.zip", + filename="TorsionNet500.zip", + ) + + results = {} + for model_name in MODELS: + path = CALC_PATH / model_name / "model_output.json" + if not path.exists(): + continue + benchmark = MlPegDihedralScanBenchmark( + force_field=Calculator(), + data_input_dir=data_input_dir, + run_mode="standard", + ) + benchmark.model_output = DihedralScanModelOutput.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. + """ + data_input_dir = download_s3_data( + key="inputs/conformers/TorsionNet500/TorsionNet500.zip", + filename="TorsionNet500.zip", + ) + benchmark = MlPegDihedralScanBenchmark( + force_field=Calculator(), + data_input_dir=data_input_dir, + run_mode="standard", + ) + elements = sorted( + { + symbol + for fragment in benchmark._torsion_net_500.values() + for symbol in fragment.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 +def fragment_profiles(analyze_results) -> dict[str, dict[str, list]]: + """ + Get per-fragment barrier heights and torsion profile for each model. + + Parameters + ---------- + analyze_results + Mapping of model name to its ``DihedralScanResult``. + + Returns + ------- + dict[str, dict[str, list]] + Per model, the fragment labels, the predicted and reference barrier + heights for each torsion scan (for the parity plot), and the + mean-centered angle/energy profile for each scan (for plotting torsion + curves). + """ + results = {} + for model_name in MODELS: + labels_list = [] + pred_barrier = [] + ref_barrier = [] + profiles = [] + result = analyze_results.get(model_name) + if result is not None: + for fragment in result.fragments: + if fragment.failed: + continue + angle = np.asarray(fragment.distance_profile) + ref_energy = np.asarray(fragment.reference_energy_profile) + model_energy = np.asarray(fragment.predicted_energy_profile) + labels_list.append(fragment.fragment_name) + pred_barrier.append(float(model_energy.max() - model_energy.min())) + ref_barrier.append(float(ref_energy.max() - ref_energy.min())) + # Mean-center so only the relative torsional profile is compared. + ref_rel_energy = ref_energy - ref_energy.mean() + model_rel_energy = model_energy - model_energy.mean() + order = np.argsort(angle) + profiles.append( + { + "angle": angle[order].tolist(), + "ref_rel_energy": ref_rel_energy[order].tolist(), + "model_rel_energy": model_rel_energy[order].tolist(), + } + ) + results[model_name] = { + "labels": labels_list, + "pred_barrier": pred_barrier, + "ref_barrier": ref_barrier, + "profiles": profiles, + } + return results + + +@pytest.fixture +def get_mae(analyze_results) -> dict[str, float]: + """ + Get the barrier height mean absolute error for each model. + + Parameters + ---------- + analyze_results + Mapping of model name to its ``DihedralScanResult``. + + Returns + ------- + dict[str, float] + Mean absolute barrier height error in kcal/mol for each model. + """ + return { + model_name: result.mae_barrier_height + 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 ``DihedralScanResult``. + + Returns + ------- + dict[str, float] + The mlipaudit per-fragment soft-threshold score (0 to 1) for each model. + """ + return {model_name: result.score for model_name, result in analyze_results.items()} + + +def plot_fragment_parity_figure( + model_name: str, + labels: list[str], + pred_barrier: list[float], + ref_barrier: list[float], +) -> go.Figure: + """ + Build a barrier height parity plot for one model. + + Parameters + ---------- + model_name + Name of the model the parity plot is for. + labels + Fragment labels, in scatter point order. + pred_barrier + Predicted torsion barrier heights, in the same order as ``labels``. + ref_barrier + Reference torsion barrier heights, in the same order as ``labels``. + + Returns + ------- + go.Figure + Parity plot of predicted against reference barrier height, one point + per fragment, with a ``y = x`` reference line. + """ + fig = go.Figure() + fig.add_trace( + go.Scatter( + x=pred_barrier, + y=ref_barrier, + mode="markers", + customdata=labels, + hovertemplate=( + "Fragment: %{customdata}
" + "Predicted: %{x:.3f} kcal/mol
" + "Reference: %{y:.3f} kcal/mol
" + ), + showlegend=False, + ) + ) + lims = [0.0, max([*pred_barrier, *ref_barrier], default=1.0)] + fig.add_trace( + go.Scatter(x=lims, y=lims, mode="lines", showlegend=False, hoverinfo="skip") + ) + fig.update_layout( + title={"text": f"Barrier height parity - {model_name}"}, + xaxis={"title": {"text": "Predicted barrier height / kcal/mol"}}, + yaxis={"title": {"text": "Reference barrier height / kcal/mol"}}, + ) + return fig + + +@pytest.fixture +def fragment_scatter_figures(fragment_profiles: dict[str, dict[str, list]]) -> None: + """ + Save per-model barrier height parity plots for the app. + + Parameters + ---------- + fragment_profiles + Per-fragment barrier heights, labels, and profiles for each model. + """ + for model_name in MODELS: + labels = fragment_profiles[model_name]["labels"] + if not labels: + continue + out_dir = OUT_PATH / model_name + out_dir.mkdir(parents=True, exist_ok=True) + parity_fig = plot_fragment_parity_figure( + model_name, + labels, + fragment_profiles[model_name]["pred_barrier"], + fragment_profiles[model_name]["ref_barrier"], + ) + parity_fig.write_json(out_dir / "fragment_barrier_parity.json") + + +def plot_torsion_curve_figure( + model_name: str, label: str, profile: dict[str, list] +) -> go.Figure: + """ + Build a torsion energy profile plot for one fragment. + + Parameters + ---------- + model_name + Name of the model the profile was predicted with. + label + Fragment label the profile belongs to. + profile + Mean-centered ``angle``, ``ref_rel_energy``, and ``model_rel_energy`` for + the scan. + + Returns + ------- + go.Figure + Line plot of relative energy against dihedral angle. + """ + fig = go.Figure() + fig.add_trace( + go.Scatter( + x=profile["angle"], + y=profile["ref_rel_energy"], + mode="lines+markers", + name="Reference", + ) + ) + fig.add_trace( + go.Scatter( + x=profile["angle"], + y=profile["model_rel_energy"], + mode="lines+markers", + name=model_name, + ) + ) + fig.update_layout( + title={"text": f"{label} - {model_name}"}, + xaxis={"title": {"text": "Dihedral angle / deg"}}, + yaxis={"title": {"text": "Relative energy / kcal/mol"}}, + ) + return fig + + +@pytest.fixture +def torsion_curve_figures(fragment_profiles: dict[str, dict[str, list]]) -> None: + """ + Save per-fragment torsion energy profile plots for the app. + + Parameters + ---------- + fragment_profiles + Per-fragment barrier height error, labels, and profiles for each model. + """ + for model_name in MODELS: + labels = fragment_profiles[model_name]["labels"] + profiles = fragment_profiles[model_name]["profiles"] + if not labels: + continue + out_dir = OUT_PATH / model_name / "torsion_curves" + out_dir.mkdir(parents=True, exist_ok=True) + for label, profile in zip(labels, profiles, strict=True): + fig = plot_torsion_curve_figure(model_name, label, profile) + fig.write_json(out_dir / f"{label}.json") + + +@pytest.fixture +def torsion_trajectories() -> None: + """ + Save per-fragment torsion scan trajectories for the app's structure viewer. + + Geometries are identical across all models, since only single-point energies + are calculated on the same reference conformers, so this only needs writing + once, from the input dataset, angle-sorted to match the torsion curves. + """ + data_input_dir = download_s3_data( + key="inputs/conformers/TorsionNet500/TorsionNet500.zip", + filename="TorsionNet500.zip", + ) + benchmark = MlPegDihedralScanBenchmark( + force_field=Calculator(), + data_input_dir=data_input_dir, + run_mode="standard", + ) + out_dir = OUT_PATH / "torsion_trajectories" + out_dir.mkdir(parents=True, exist_ok=True) + for fragment_name, fragment in benchmark._torsion_net_500.items(): + angle = np.array([state[0] for state in fragment.dft_energy_profile]) + order = np.argsort(angle) + images = [ + Atoms( + symbols=fragment.atom_symbols, + positions=fragment.conformer_coordinates[i], + ) + for i in order + ] + write(out_dir / f"{fragment_name}.xyz", images) + + +@pytest.fixture +@build_table( + filename=OUT_PATH / "torsionnet500_metrics_table.json", + metric_tooltips=DEFAULT_TOOLTIPS, + thresholds=DEFAULT_THRESHOLDS, + weights=DEFAULT_WEIGHTS, + mlip_name_map=DISPERSION_NAME_MAP, +) +def metrics(get_mae: dict[str, float], get_score: dict[str, float]) -> dict[str, dict]: + """ + Get all metrics. + + Parameters + ---------- + get_mae + Mean absolute barrier height errors 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_mae, + "Torsion Score": get_score, + } + + +def test_torsionnet500( + metrics: dict[str, dict], + fragment_scatter_figures: None, + torsion_curve_figures: None, + torsion_trajectories: None, + struct_info: dict, +) -> None: + """ + Run TorsionNet500 analysis. + + Parameters + ---------- + metrics : dict[str, dict] + TorsionNet500 metric results provided by fixtures. + fragment_scatter_figures + Per-model fragment barrier-error scatter figures (side-effect only). + torsion_curve_figures + Per-fragment torsion curve figures (side-effect only). + torsion_trajectories + Per-fragment torsion scan trajectories (side-effect only). + struct_info : dict + Element info written to ``info.json`` for filtering. + """ diff --git a/ml_peg/analysis/conformers/TorsionNet500/metrics.yml b/ml_peg/analysis/conformers/TorsionNet500/metrics.yml new file mode 100644 index 000000000..963e3bfe9 --- /dev/null +++ b/ml_peg/analysis/conformers/TorsionNet500/metrics.yml @@ -0,0 +1,15 @@ +metrics: + Barrier Height MAE: + good: 0.0 + bad: 2.0 + unit: kcal/mol + weight: 0 + tooltip: Mean absolute error of the torsion barrier height (max minus min of the dihedral energy profile). + level_of_theory: wB97M-D3(BJ)/def2-TZVPPD + Torsion Score: + good: 1.0 + bad: 0.0 + unit: null + weight: 1 + tooltip: mlipaudit per-fragment soft-threshold score on the barrier height error (failed fragments score 0). + level_of_theory: wB97M-D3(BJ)/def2-TZVPPD diff --git a/ml_peg/app/conformers/TorsionNet500/app_TorsionNet500.py b/ml_peg/app/conformers/TorsionNet500/app_TorsionNet500.py new file mode 100644 index 000000000..f094d40f1 --- /dev/null +++ b/ml_peg/app/conformers/TorsionNet500/app_TorsionNet500.py @@ -0,0 +1,237 @@ +"""Run TorsionNet500 dihedral scan benchmark app.""" + +from __future__ import annotations + +from pathlib import Path + +from dash import Dash, Input, Output, State, callback +from dash.dcc import Store +from dash.html import Div, Iframe + +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, + plot_with_download_controls, +) +from ml_peg.app.utils.load import read_plot +from ml_peg.app.utils.weas import generate_weas_html +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 = "TorsionNet500" +DOCS_URL = ( + "https://ddmms.github.io/ml-peg/user_guide/benchmarks/conformers.html#torsionnet500" +) +DATA_PATH = APP_ROOT / "data" / "conformers" / "TorsionNet500" +INFO_PATH = DATA_PATH / "info.json" +ASSETS_DIR = "/assets/conformers/TorsionNet500" + + +def _register_curve_callback( + scatter_id: str, + plot_id: str, + curve_dir: Path, + labels: list[str], + struct_plot_id: str, +) -> None: + """ + Attach callbacks that show a torsion curve and structure on point clicks. + + Unlike `plot_from_scatter`, this reads the clicked fragment's curve JSON from + disk on demand instead of requiring every curve to be pre-loaded into a + `plots_list` up front - with up to 500 fragments per model, pre-loading all of + them for every model would be needlessly expensive when only one is ever shown + at a time. + + A second callback shows the 3D structure at the torsion angle of whichever + point on the curve is clicked, using WEAS trajectory mode. The curve's + fragment label is tracked in a `Store` so the structure callback (triggered by + clicks on the curve, whose id is fixed and reused across fragments) knows which + fragment's trajectory file to load. + + Parameters + ---------- + scatter_id + ID of the per-model fragment-barrier scatter plot. + plot_id + ID of the shared placeholder Div where curves are rendered. + curve_dir + Directory containing this model's per-fragment torsion curve JSON files. + labels + Fragment labels, in the same order as the scatter's points. + struct_plot_id + ID of the shared placeholder Div where structures are rendered. + """ + label_store_id = f"{scatter_id}-curve-label" + + @callback( + Output(plot_id, "children", allow_duplicate=True), + Output(label_store_id, "data", allow_duplicate=True), + Output(struct_plot_id, "children", allow_duplicate=True), + Input(scatter_id, "clickData"), + prevent_initial_call="initial_duplicate", + ) + def show_curve(click_data): + """ + Register callback to show a torsion curve when a scatter point is clicked. + + Also clears any structure shown from a previously displayed curve, so a + stale structure from a different fragment/model doesn't linger on screen + until the new curve is itself clicked. + + Parameters + ---------- + click_data + Clicked data point in the fragment-barrier scatter plot. + + Returns + ------- + tuple[Div, str | None, Div] + Torsion curve plot on scatter click, the clicked fragment's label, and + an empty Div to clear any previously displayed structure. + """ + if not click_data: + return Div(), None, Div() + point = click_data["points"][0] + # The y = x reference line trace carries no customdata; ignore clicks on + # it so they don't map to an unrelated fragment. + if point.get("customdata") is None: + return Div(), None, Div() + idx = point["pointNumber"] + if idx < 0 or idx >= len(labels): + return Div(), None, Div() + label = labels[idx] + curve_path = curve_dir / f"{label}.json" + graph = read_plot(curve_path, id=f"{scatter_id}-curve") + return plot_with_download_controls(graph), label, Div() + + @callback( + Output(struct_plot_id, "children", allow_duplicate=True), + Input(f"{scatter_id}-curve", "clickData"), + State(label_store_id, "data"), + prevent_initial_call="initial_duplicate", + ) + def show_struct(click_data, label): + """ + Register callback to show a structure when a point on the curve is clicked. + + Parameters + ---------- + click_data + Clicked data point in the torsion curve plot. + label + Fragment label of the currently displayed curve. + + Returns + ------- + Div + Structure at the clicked dihedral angle, in WEAS trajectory mode. + """ + if not click_data or not label: + return Div() + idx = click_data["points"][0]["pointNumber"] + traj_path = f"{ASSETS_DIR}/torsion_trajectories/{label}.xyz" + return Div( + Iframe( + srcDoc=generate_weas_html(traj_path, mode="traj", index=idx), + style={ + "height": "550px", + "width": "100%", + "border": "1px solid #ddd", + "borderRadius": "5px", + }, + ) + ) + + +class TorsionNet500App(BaseApp): + """TorsionNet500 dihedral scan benchmark app layout and callbacks.""" + + def register_callbacks(self) -> None: + """Register callbacks to app.""" + # Build a per-fragment barrier-error scatter per model (fragment index vs + # that fragment's barrier height error), and record fragment labels in + # scatter-point order so a later click on a point can be mapped back to + # its curve file. + scatter_cell_to_plot: dict[str, dict[str, Div]] = {} + fragment_labels: dict[str, list[str]] = {} + for model_name in MODELS: + barrier_parity_path = ( + DATA_PATH / model_name / "fragment_barrier_parity.json" + ) + curve_dir = DATA_PATH / model_name / "torsion_curves" + if not barrier_parity_path.exists() or not curve_dir.exists(): + continue + scatter_cell_to_plot[model_name] = { + "Barrier Height MAE": read_plot( + barrier_parity_path, + id=f"{BENCHMARK_NAME}-{model_name}-barrier-figure", + ), + } + fragment_labels[model_name] = [ + curve_file.stem for curve_file in sorted(curve_dir.glob("*.json")) + ] + + # Clicking a model's Barrier Height MAE cell shows that model's + # per-fragment barrier-error scatter. + plot_from_table_cell( + table_id=self.table_id, + plot_id=f"{BENCHMARK_NAME}-scatter-placeholder", + cell_to_plot=scatter_cell_to_plot, + ) + + # Clicking a point on the scatter shows that fragment's torsion curve. + for model_name, labels in fragment_labels.items(): + curve_dir = DATA_PATH / model_name / "torsion_curves" + _register_curve_callback( + scatter_id=f"{BENCHMARK_NAME}-{model_name}-barrier-figure", + plot_id=f"{BENCHMARK_NAME}-curve-placeholder", + curve_dir=curve_dir, + labels=labels, + struct_plot_id=f"{BENCHMARK_NAME}-struct-placeholder", + ) + + +def get_app() -> TorsionNet500App: + """ + Get TorsionNet500 benchmark app layout and callback registration. + + Returns + ------- + TorsionNet500App + Benchmark layout and callback registration. + """ + return TorsionNet500App( + name=BENCHMARK_NAME, + framework_ids="mlip_audit", + description=( + "Performance in predicting torsion energy barriers for drug-like " + "molecules from systematic dihedral scans. Reference data from " + "wB97M-D3(BJ)/def2-TZVPPD calculations." + ), + docs_url=DOCS_URL, + table_path=DATA_PATH / "torsionnet500_metrics_table.json", + extra_components=[ + Div(id=f"{BENCHMARK_NAME}-scatter-placeholder"), + Div(id=f"{BENCHMARK_NAME}-curve-placeholder"), + Div(id=f"{BENCHMARK_NAME}-struct-placeholder"), + # One Store per model, remembering which fragment's curve is currently + # displayed, so a click on the curve (whose id is fixed and reused + # across fragments) can be mapped back to the right trajectory file. + *[ + Store(id=f"{BENCHMARK_NAME}-{model_name}-barrier-figure-curve-label") + for model_name in MODELS + ], + ], + 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=8069, 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/conformers/TorsionNet500/calc_TorsionNet500.py b/ml_peg/calcs/conformers/TorsionNet500/calc_TorsionNet500.py new file mode 100644 index 000000000..abbfdcc8e --- /dev/null +++ b/ml_peg/calcs/conformers/TorsionNet500/calc_TorsionNet500.py @@ -0,0 +1,71 @@ +""" +Compute the TorsionNet500 dihedral scan dataset. + +TorsionNet: A Deep Neural Network to Rapidly Predict Small-Molecule +Torsional Energy Profiles with the Accuracy of Quantum Mechanics. +Rai, B. K. et al. J. Chem. Inf. Model. 2022, 62 (4), 785-800. +DOI: 10.1021/acs.jcim.1c01346 +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any +from warnings import warn + +import pytest + +# Optional extra (ml-peg[mlipaudit]); skip if not installed. +pytest.importorskip("mlipaudit", reason="Please install `mlipaudit` extra") + +from mlipaudit.benchmarks.dihedral_scan.dihedral_scan import DihedralScanModelOutput + +from ml_peg.calcs.utils.mlipaudit import MlPegDihedralScanBenchmark +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_torsionnet500(mlip: tuple[str, Any]) -> None: + """ + Benchmark the TorsionNet500 dihedral scan 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/conformers/TorsionNet500/TorsionNet500.zip", + filename="TorsionNet500.zip", + ) + + out_path = OUT_PATH / model_name + out_path.mkdir(parents=True, exist_ok=True) + + benchmark = MlPegDihedralScanBenchmark( + force_field=calc, + data_input_dir=data_input_dir, + run_mode="standard", + ) + try: + benchmark.run_model() + except Exception as exc: + warn( + f"Error running dihedral scan benchmark for {model_name}: {exc}", + stacklevel=2, + ) + benchmark.model_output = DihedralScanModelOutput(fragments=[]) + + (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..89359cc4d --- /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.dihedral_scan.dihedral_scan import DihedralScanBenchmark + + +class MlPegDihedralScanBenchmark(DihedralScanBenchmark): + """ + DihedralScanBenchmark 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