From 09cf846466b0d530082b519d342a735f24fd15a1 Mon Sep 17 00:00:00 2001 From: Leon Wehrhan Date: Thu, 2 Jul 2026 15:06:40 +0200 Subject: [PATCH 1/3] feat: add torsion net benchmark --- .../user_guide/benchmarks/conformers.rst | 45 +++ .../TorsionNet500/analyse_TorsionNet500.py | 266 ++++++++++++++++++ .../conformers/TorsionNet500/metrics.yml | 15 + .../TorsionNet500/app_TorsionNet500.py | 68 +++++ ml_peg/app/utils/frameworks.yml | 7 + .../TorsionNet500/calc_TorsionNet500.py | 67 +++++ ml_peg/calcs/utils/mlipaudit.py | 16 ++ 7 files changed, 484 insertions(+) create mode 100644 ml_peg/analysis/conformers/TorsionNet500/analyse_TorsionNet500.py create mode 100644 ml_peg/analysis/conformers/TorsionNet500/metrics.yml create mode 100644 ml_peg/app/conformers/TorsionNet500/app_TorsionNet500.py create mode 100644 ml_peg/calcs/conformers/TorsionNet500/calc_TorsionNet500.py create mode 100644 ml_peg/calcs/utils/mlipaudit.py 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..147a7a1fe --- /dev/null +++ b/ml_peg/analysis/conformers/TorsionNet500/analyse_TorsionNet500.py @@ -0,0 +1,266 @@ +"""Analyse TorsionNet500 dihedral scan benchmark.""" + +from __future__ import annotations + +import json +from pathlib import Path + +from ase.calculators.calculator import Calculator +from mlipaudit.benchmarks.dihedral_scan.dihedral_scan import DihedralScanModelOutput +import pytest + +from ml_peg.analysis.utils.decorators import build_table, plot_parity +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 [fragment.fragment_name for fragment in output.fragments] + return [] + + +def _barrier(profile: list[float]) -> float: + """ + Get the torsion barrier height of an energy profile. + + Parameters + ---------- + profile + Energy profile along the dihedral scan in kcal/mol. + + Returns + ------- + float + Difference between the maximum and minimum energy in the profile. + """ + return max(profile) - min(profile) + + +@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 +@plot_parity( + filename=OUT_PATH / "figure_torsionnet500.json", + title="Torsion barrier heights", + x_label="Predicted barrier height / kcal/mol", + y_label="Reference barrier height / kcal/mol", + hoverdata={ + "Fragment": labels(), + }, +) +def barrier_heights(analyze_results) -> dict[str, list]: + """ + Get predicted and reference torsion barrier heights. + + Parameters + ---------- + analyze_results + Mapping of model name to its ``DihedralScanResult``. + + Returns + ------- + dict[str, list] + Dictionary of reference and predicted barrier heights, aligned to the + ordered fragment names. + """ + ids = labels() + ref_map: dict[str, float] = {} + pred_maps: dict[str, dict[str, float]] = {} + + for model_name, result in analyze_results.items(): + pred_maps[model_name] = {} + for fragment in result.fragments: + if fragment.failed: + continue + pred_maps[model_name][fragment.fragment_name] = _barrier( + fragment.predicted_energy_profile + ) + ref_map.setdefault( + fragment.fragment_name, _barrier(fragment.reference_energy_profile) + ) + + results = {"ref": [ref_map.get(i) for i in ids]} + for model_name in MODELS: + model_preds = pred_maps.get(model_name, {}) + results[model_name] = [model_preds.get(i) for i in ids] + 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()} + + +@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( + barrier_heights, get_mae: dict[str, float], get_score: dict[str, float] +) -> dict[str, dict]: + """ + Get all metrics. + + Parameters + ---------- + barrier_heights + Reference and predicted barrier heights (triggers the parity plot). + 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], barrier_heights, struct_info: dict +) -> None: + """ + Run TorsionNet500 analysis. + + Parameters + ---------- + metrics : dict[str, dict] + TorsionNet500 metric results provided by fixtures. + barrier_heights + Reference and predicted barrier heights (triggers the parity plot). + 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..fb87570e8 --- /dev/null +++ b/ml_peg/app/conformers/TorsionNet500/app_TorsionNet500.py @@ -0,0 +1,68 @@ +"""Run TorsionNet500 dihedral scan 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 = "TorsionNet500" +DOCS_URL = ( + "https://ddmms.github.io/ml-peg/user_guide/benchmarks/conformers.html#torsionnet500" +) +DATA_PATH = APP_ROOT / "data" / "conformers" / "TorsionNet500" + + +class TorsionNet500App(BaseApp): + """TorsionNet500 dihedral scan benchmark app layout and callbacks.""" + + def register_callbacks(self) -> None: + """Register callbacks to app.""" + scatter = read_plot( + DATA_PATH / "figure_torsionnet500.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={"Barrier Height MAE": scatter}, + ) + + +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", + 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=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..fc09f68ca --- /dev/null +++ b/ml_peg/calcs/conformers/TorsionNet500/calc_TorsionNet500.py @@ -0,0 +1,67 @@ +""" +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 + +from mlipaudit.benchmarks.dihedral_scan.dihedral_scan import DihedralScanModelOutput +import pytest + +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 From 8c84c46bb4e9d4f5127ee62aef6c42071ec69728 Mon Sep 17 00:00:00 2001 From: joehart2001 Date: Thu, 23 Jul 2026 23:23:09 +0100 Subject: [PATCH 2/3] add structure saving and visualisation --- .../TorsionNet500/analyse_TorsionNet500.py | 19 ++++++++++++++-- .../TorsionNet500/app_TorsionNet500.py | 22 ++++++++++++++++++- .../TorsionNet500/calc_TorsionNet500.py | 6 ++++- 3 files changed, 43 insertions(+), 4 deletions(-) diff --git a/ml_peg/analysis/conformers/TorsionNet500/analyse_TorsionNet500.py b/ml_peg/analysis/conformers/TorsionNet500/analyse_TorsionNet500.py index 147a7a1fe..586d43e30 100644 --- a/ml_peg/analysis/conformers/TorsionNet500/analyse_TorsionNet500.py +++ b/ml_peg/analysis/conformers/TorsionNet500/analyse_TorsionNet500.py @@ -5,7 +5,9 @@ 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 pytest @@ -43,7 +45,7 @@ def labels() -> list: path = CALC_PATH / model_name / "model_output.json" if path.exists(): output = DihedralScanModelOutput.model_validate_json(path.read_text()) - return [fragment.fragment_name for fragment in output.fragments] + return sorted(fragment.fragment_name for fragment in output.fragments) return [] @@ -99,7 +101,10 @@ def analyze_results() -> dict: @pytest.fixture def struct_info() -> dict: """ - Write the combined element set to ``info.json`` for filtering. + Write ``info.json`` for filtering and one ``.xyz`` scan per fragment. + + Each fragment is written as a multi-frame trajectory (one frame per + dihedral scan conformer) so the app can step through the torsion scan. Returns ------- @@ -125,6 +130,16 @@ def struct_info() -> dict: info = {"elements": elements} OUT_PATH.mkdir(parents=True, exist_ok=True) (OUT_PATH / "info.json").write_text(json.dumps(info, indent=1)) + + structs_dir = OUT_PATH / "mock" + structs_dir.mkdir(parents=True, exist_ok=True) + for fragment_name, fragment in benchmark._torsion_net_500.items(): + images = [ + Atoms(symbols=fragment.atom_symbols, positions=coords) + for coords in fragment.conformer_coordinates + ] + write(structs_dir / f"{fragment_name}.xyz", images) + return info diff --git a/ml_peg/app/conformers/TorsionNet500/app_TorsionNet500.py b/ml_peg/app/conformers/TorsionNet500/app_TorsionNet500.py index fb87570e8..3c0d14dae 100644 --- a/ml_peg/app/conformers/TorsionNet500/app_TorsionNet500.py +++ b/ml_peg/app/conformers/TorsionNet500/app_TorsionNet500.py @@ -7,7 +7,10 @@ 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.build_callbacks import ( + plot_from_table_column, + struct_from_scatter, +) from ml_peg.app.utils.load import read_plot BENCHMARK_NAME = "TorsionNet500" @@ -27,12 +30,28 @@ def register_callbacks(self) -> None: id=f"{BENCHMARK_NAME}-figure", ) + struct_dir = DATA_PATH / "mock" + if struct_dir.exists(): + labels = sorted([f.stem for f in struct_dir.glob("*.xyz")]) + structs = [ + f"/assets/conformers/TorsionNet500/mock/{label}.xyz" for label in labels + ] + else: + structs = [] + plot_from_table_column( table_id=self.table_id, plot_id=f"{BENCHMARK_NAME}-figure-placeholder", column_to_plot={"Barrier Height MAE": scatter}, ) + struct_from_scatter( + scatter_id=f"{BENCHMARK_NAME}-figure", + struct_id=f"{BENCHMARK_NAME}-struct-placeholder", + structs=structs, + mode="struct", + ) + def get_app() -> TorsionNet500App: """ @@ -56,6 +75,7 @@ def get_app() -> TorsionNet500App: info_path=DATA_PATH / "info.json", extra_components=[ Div(id=f"{BENCHMARK_NAME}-figure-placeholder"), + Div(id=f"{BENCHMARK_NAME}-struct-placeholder"), ], ) diff --git a/ml_peg/calcs/conformers/TorsionNet500/calc_TorsionNet500.py b/ml_peg/calcs/conformers/TorsionNet500/calc_TorsionNet500.py index fc09f68ca..abbfdcc8e 100644 --- a/ml_peg/calcs/conformers/TorsionNet500/calc_TorsionNet500.py +++ b/ml_peg/calcs/conformers/TorsionNet500/calc_TorsionNet500.py @@ -13,9 +13,13 @@ from typing import Any from warnings import warn -from mlipaudit.benchmarks.dihedral_scan.dihedral_scan import DihedralScanModelOutput 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 53a14d8b23dfd78ad6b7c336924dbd620e4319bb Mon Sep 17 00:00:00 2001 From: joehart2001 Date: Thu, 23 Jul 2026 23:55:11 +0100 Subject: [PATCH 3/3] add torsion profile and structure vis --- .../TorsionNet500/analyse_TorsionNet500.py | 308 +++++++++++++----- .../TorsionNet500/app_TorsionNet500.py | 203 ++++++++++-- 2 files changed, 411 insertions(+), 100 deletions(-) diff --git a/ml_peg/analysis/conformers/TorsionNet500/analyse_TorsionNet500.py b/ml_peg/analysis/conformers/TorsionNet500/analyse_TorsionNet500.py index 586d43e30..f3a5676ce 100644 --- a/ml_peg/analysis/conformers/TorsionNet500/analyse_TorsionNet500.py +++ b/ml_peg/analysis/conformers/TorsionNet500/analyse_TorsionNet500.py @@ -9,9 +9,11 @@ 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, plot_parity +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 @@ -49,23 +51,6 @@ def labels() -> list: return [] -def _barrier(profile: list[float]) -> float: - """ - Get the torsion barrier height of an energy profile. - - Parameters - ---------- - profile - Energy profile along the dihedral scan in kcal/mol. - - Returns - ------- - float - Difference between the maximum and minimum energy in the profile. - """ - return max(profile) - min(profile) - - @pytest.fixture def analyze_results() -> dict: """ @@ -101,10 +86,7 @@ def analyze_results() -> dict: @pytest.fixture def struct_info() -> dict: """ - Write ``info.json`` for filtering and one ``.xyz`` scan per fragment. - - Each fragment is written as a multi-frame trajectory (one frame per - dihedral scan conformer) so the app can step through the torsion scan. + Write the combined element set to ``info.json`` for filtering. Returns ------- @@ -130,32 +112,13 @@ def struct_info() -> dict: info = {"elements": elements} OUT_PATH.mkdir(parents=True, exist_ok=True) (OUT_PATH / "info.json").write_text(json.dumps(info, indent=1)) - - structs_dir = OUT_PATH / "mock" - structs_dir.mkdir(parents=True, exist_ok=True) - for fragment_name, fragment in benchmark._torsion_net_500.items(): - images = [ - Atoms(symbols=fragment.atom_symbols, positions=coords) - for coords in fragment.conformer_coordinates - ] - write(structs_dir / f"{fragment_name}.xyz", images) - return info @pytest.fixture -@plot_parity( - filename=OUT_PATH / "figure_torsionnet500.json", - title="Torsion barrier heights", - x_label="Predicted barrier height / kcal/mol", - y_label="Reference barrier height / kcal/mol", - hoverdata={ - "Fragment": labels(), - }, -) -def barrier_heights(analyze_results) -> dict[str, list]: +def fragment_profiles(analyze_results) -> dict[str, dict[str, list]]: """ - Get predicted and reference torsion barrier heights. + Get per-fragment barrier heights and torsion profile for each model. Parameters ---------- @@ -164,30 +127,46 @@ def barrier_heights(analyze_results) -> dict[str, list]: Returns ------- - dict[str, list] - Dictionary of reference and predicted barrier heights, aligned to the - ordered fragment names. + 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). """ - ids = labels() - ref_map: dict[str, float] = {} - pred_maps: dict[str, dict[str, float]] = {} - - for model_name, result in analyze_results.items(): - pred_maps[model_name] = {} - for fragment in result.fragments: - if fragment.failed: - continue - pred_maps[model_name][fragment.fragment_name] = _barrier( - fragment.predicted_energy_profile - ) - ref_map.setdefault( - fragment.fragment_name, _barrier(fragment.reference_energy_profile) - ) - - results = {"ref": [ref_map.get(i) for i in ids]} + results = {} for model_name in MODELS: - model_preds = pred_maps.get(model_name, {}) - results[model_name] = [model_preds.get(i) for i in ids] + 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 @@ -230,6 +209,185 @@ def get_score(analyze_results) -> dict[str, float]: 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", @@ -238,16 +396,12 @@ def get_score(analyze_results) -> dict[str, float]: weights=DEFAULT_WEIGHTS, mlip_name_map=DISPERSION_NAME_MAP, ) -def metrics( - barrier_heights, get_mae: dict[str, float], get_score: dict[str, float] -) -> dict[str, dict]: +def metrics(get_mae: dict[str, float], get_score: dict[str, float]) -> dict[str, dict]: """ Get all metrics. Parameters ---------- - barrier_heights - Reference and predicted barrier heights (triggers the parity plot). get_mae Mean absolute barrier height errors for all models. get_score @@ -265,7 +419,11 @@ def metrics( def test_torsionnet500( - metrics: dict[str, dict], barrier_heights, struct_info: dict + metrics: dict[str, dict], + fragment_scatter_figures: None, + torsion_curve_figures: None, + torsion_trajectories: None, + struct_info: dict, ) -> None: """ Run TorsionNet500 analysis. @@ -274,8 +432,12 @@ def test_torsionnet500( ---------- metrics : dict[str, dict] TorsionNet500 metric results provided by fixtures. - barrier_heights - Reference and predicted barrier heights (triggers the parity plot). + 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/app/conformers/TorsionNet500/app_TorsionNet500.py b/ml_peg/app/conformers/TorsionNet500/app_TorsionNet500.py index 3c0d14dae..f094d40f1 100644 --- a/ml_peg/app/conformers/TorsionNet500/app_TorsionNet500.py +++ b/ml_peg/app/conformers/TorsionNet500/app_TorsionNet500.py @@ -2,22 +2,148 @@ from __future__ import annotations -from dash import Dash -from dash.html import Div +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_column, - struct_from_scatter, + 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): @@ -25,32 +151,47 @@ class TorsionNet500App(BaseApp): def register_callbacks(self) -> None: """Register callbacks to app.""" - scatter = read_plot( - DATA_PATH / "figure_torsionnet500.json", - id=f"{BENCHMARK_NAME}-figure", - ) - - struct_dir = DATA_PATH / "mock" - if struct_dir.exists(): - labels = sorted([f.stem for f in struct_dir.glob("*.xyz")]) - structs = [ - f"/assets/conformers/TorsionNet500/mock/{label}.xyz" for label in labels + # 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")) ] - else: - structs = [] - plot_from_table_column( + # 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}-figure-placeholder", - column_to_plot={"Barrier Height MAE": scatter}, + plot_id=f"{BENCHMARK_NAME}-scatter-placeholder", + cell_to_plot=scatter_cell_to_plot, ) - struct_from_scatter( - scatter_id=f"{BENCHMARK_NAME}-figure", - struct_id=f"{BENCHMARK_NAME}-struct-placeholder", - structs=structs, - mode="struct", - ) + # 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: @@ -72,11 +213,19 @@ def get_app() -> TorsionNet500App: ), docs_url=DOCS_URL, table_path=DATA_PATH / "torsionnet500_metrics_table.json", - info_path=DATA_PATH / "info.json", extra_components=[ - Div(id=f"{BENCHMARK_NAME}-figure-placeholder"), + 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, )