diff --git a/ml_peg/analysis/molecular_crystal/DMC_ICE13/analyse_DMC_ICE13.py b/ml_peg/analysis/molecular_crystal/DMC_ICE13/analyse_DMC_ICE13.py index 7346ca9bb..b747f5621 100644 --- a/ml_peg/analysis/molecular_crystal/DMC_ICE13/analyse_DMC_ICE13.py +++ b/ml_peg/analysis/molecular_crystal/DMC_ICE13/analyse_DMC_ICE13.py @@ -10,6 +10,7 @@ from ml_peg.analysis.utils.decorators import build_table, plot_parity from ml_peg.analysis.utils.utils import ( build_dispersion_name_map, + get_struct_info, load_metrics_config, mae, ) @@ -28,27 +29,12 @@ METRICS_CONFIG_PATH ) - -def get_polymorph_names() -> list[str]: - """ - Get list of polymorph names. - - Returns - ------- - list[str] - List of polymorph names from structure files. - """ - polymorph_names = [] - for model_name in MODELS: - model_dir = CALC_PATH / model_name - if model_dir.exists(): - xyz_files = sorted(model_dir.glob("*_polymorph.xyz")) - if xyz_files: - for xyz_file in xyz_files: - atoms = read(xyz_file) - polymorph_names.append(atoms.info["polymorph"]) - break - return polymorph_names +INFO = get_struct_info( + calc_path=CALC_PATH, + glob_pattern="*_polymorph.xyz", + info_keys=["polymorph"], + out_path=OUT_PATH, +) @pytest.fixture @@ -58,12 +44,12 @@ def get_polymorph_names() -> list[str]: x_label="Predicted lattice energy / meV", y_label="Reference lattice energy / meV", hoverdata={ - "Polymorph": get_polymorph_names(), + "Polymorph": INFO["polymorph"], }, ) def lattice_energies() -> dict[str, list]: """ - Get lattice energies for all DMC-ICE13 polymorphs. + Get lattice energies for all DMC-ICE13 polymorphs and plot as scatter. Returns ------- @@ -109,8 +95,7 @@ def lattice_energies() -> dict[str, list]: return results -@pytest.fixture -def dmc_ice13_errors(lattice_energies) -> dict[str, float]: +def get_errors(lattice_energies: dict[str, list]) -> dict[str, float]: """ Get mean absolute error for lattice energies. @@ -126,7 +111,7 @@ def dmc_ice13_errors(lattice_energies) -> dict[str, float]: """ results = {} for model_name in MODELS: - if lattice_energies[model_name]: + if lattice_energies.get(model_name): results[model_name] = mae( lattice_energies["ref"], lattice_energies[model_name] ) @@ -135,6 +120,25 @@ def dmc_ice13_errors(lattice_energies) -> dict[str, float]: return results +def get_metrics(lattice_energies: dict[str, list]) -> dict[str, dict]: + """ + Get all DMC-ICE13 metrics. + + Parameters + ---------- + lattice_energies + Dictionary of reference and predicted lattice energies. + + Returns + ------- + dict[str, dict] + Metric names and values for all models. + """ + return { + "MAE": get_errors(lattice_energies), + } + + @pytest.fixture @build_table( filename=OUT_PATH / "dmc_ice13_metrics_table.json", @@ -142,23 +146,21 @@ def dmc_ice13_errors(lattice_energies) -> dict[str, float]: thresholds=DEFAULT_THRESHOLDS, mlip_name_map=DISPERSION_NAME_MAP, ) -def metrics(dmc_ice13_errors: dict[str, float]) -> dict[str, dict]: +def metrics(lattice_energies: dict[str, list]) -> dict[str, dict]: """ Get all DMC-ICE13 metrics. Parameters ---------- - dmc_ice13_errors - Mean absolute errors for all polymorphs. + lattice_energies + Dictionary of reference and predicted lattice energies. Returns ------- dict[str, dict] Metric names and values for all models. """ - return { - "MAE": dmc_ice13_errors, - } + return get_metrics(lattice_energies) def test_dmc_ice13(metrics: dict[str, dict]) -> None: diff --git a/ml_peg/analysis/molecular_crystal/X23/analyse_X23.py b/ml_peg/analysis/molecular_crystal/X23/analyse_X23.py index aded032ee..e57212333 100644 --- a/ml_peg/analysis/molecular_crystal/X23/analyse_X23.py +++ b/ml_peg/analysis/molecular_crystal/X23/analyse_X23.py @@ -11,6 +11,7 @@ from ml_peg.analysis.utils.decorators import build_table, plot_parity from ml_peg.analysis.utils.utils import ( build_dispersion_name_map, + get_struct_info, load_metrics_config, mae, ) @@ -32,27 +33,13 @@ # Unit conversion EV_TO_KJ_PER_MOL = units.mol / units.kJ - -def get_system_names() -> list[str]: - """ - Get list of X23 system names. - - Returns - ------- - list[str] - List of system names from structure files. - """ - system_names = [] - for model_name in MODELS: - model_dir = CALC_PATH / model_name - if model_dir.exists(): - xyz_files = sorted(model_dir.glob("*.xyz")) - if xyz_files: - for xyz_file in xyz_files: - atoms = read(xyz_file) - system_names.append(atoms.info["system"]) - break - return system_names +INFO = get_struct_info( + calc_path=CALC_PATH, + glob_pattern="*.xyz", + index="0", + info_keys=["system"], + out_path=OUT_PATH, +) @pytest.fixture @@ -62,7 +49,7 @@ def get_system_names() -> list[str]: x_label="Predicted lattice energy / kJ/mol", y_label="Reference lattice energy / kJ/mol", hoverdata={ - "System": get_system_names(), + "System": INFO["system"], }, ) def lattice_energies() -> dict[str, list]: @@ -112,8 +99,7 @@ def lattice_energies() -> dict[str, list]: return results -@pytest.fixture -def x23_errors(lattice_energies) -> dict[str, float]: +def get_errors(lattice_energies: dict[str, list]) -> dict[str, float]: """ Get mean absolute error for lattice energies. @@ -129,7 +115,7 @@ def x23_errors(lattice_energies) -> dict[str, float]: """ results = {} for model_name in MODELS: - if lattice_energies[model_name]: + if lattice_energies.get(model_name): results[model_name] = mae( lattice_energies["ref"], lattice_energies[model_name] ) @@ -138,6 +124,25 @@ def x23_errors(lattice_energies) -> dict[str, float]: return results +def get_metrics(lattice_energies: dict[str, list]) -> dict[str, dict]: + """ + Get all X23 metrics. + + Parameters + ---------- + lattice_energies + Dictionary of reference and predicted lattice energies. + + Returns + ------- + dict[str, dict] + Metric names and values for all models. + """ + return { + "MAE": get_errors(lattice_energies), + } + + @pytest.fixture @build_table( filename=OUT_PATH / "x23_metrics_table.json", @@ -145,23 +150,21 @@ def x23_errors(lattice_energies) -> dict[str, float]: thresholds=DEFAULT_THRESHOLDS, mlip_name_map=DISPERSION_NAME_MAP, ) -def metrics(x23_errors: dict[str, float]) -> dict[str, dict]: +def metrics(lattice_energies: dict[str, list]) -> dict[str, dict]: """ Get all X23 metrics. Parameters ---------- - x23_errors - Mean absolute errors for all systems. + lattice_energies + Dictionary of reference and predicted lattice energies. Returns ------- dict[str, dict] Metric names and values for all models. """ - return { - "MAE": x23_errors, - } + return get_metrics(lattice_energies) def test_x23(metrics: dict[str, dict]) -> None: diff --git a/ml_peg/analysis/nebs/li_diffusion/analyse_li_diffusion.py b/ml_peg/analysis/nebs/li_diffusion/analyse_li_diffusion.py index a7434f829..bd442e248 100644 --- a/ml_peg/analysis/nebs/li_diffusion/analyse_li_diffusion.py +++ b/ml_peg/analysis/nebs/li_diffusion/analyse_li_diffusion.py @@ -9,7 +9,7 @@ import pytest from ml_peg.analysis.utils.decorators import build_table, plot_scatter -from ml_peg.analysis.utils.utils import load_metrics_config +from ml_peg.analysis.utils.utils import load_metrics_config, write_struct_info from ml_peg.app import APP_ROOT from ml_peg.calcs import CALCS_ROOT from ml_peg.models import current_models @@ -27,20 +27,20 @@ REF_VALUES = {"path_b": 0.27, "path_c": 2.5} -def plot_nebs(model: str, path: Literal["b", "c"]) -> None: +def plot_nebs(model_name: str, path: Literal["b", "c"]) -> None: """ Plot NEB paths and save all structure files. Parameters ---------- - model + model_name Name of MLIP. path Path "b" or "c" for NEB. """ @plot_scatter( - filename=OUT_PATH / f"figure_{model}_neb_{path.lower()}.json", + filename=OUT_PATH / model_name / f"figure_neb_{path.lower()}.json", title=f"NEB path {path.upper()}", x_label="Image", y_label="Energy / eV", @@ -57,16 +57,16 @@ def plot_neb() -> dict[str, tuple[list[float], list[float]]]: """ results = {} structs = read( - CALC_PATH / f"li_diffusion_{path.lower()}-{model}-neb-band.extxyz", + CALC_PATH / model_name / f"li_diffusion_{path.lower()}-neb-band.extxyz", index=":", ) - results[model] = [ + results[model_name] = [ list(range(len(structs))), [struct.get_potential_energy() for struct in structs], ] - structs_dir = OUT_PATH / model + structs_dir = OUT_PATH / model_name structs_dir.mkdir(parents=True, exist_ok=True) - write(structs_dir / f"{model}-{path.lower()}-neb-band.extxyz", structs) + write(structs_dir / f"{path.lower()}-neb-band.extxyz", structs) return results @@ -88,7 +88,7 @@ def path_b_error() -> dict[str, float]: for model_name in MODELS: plot_nebs(model_name, "b") with open( - CALC_PATH / f"li_diffusion_b-{model_name}-neb-results.dat", encoding="utf8" + CALC_PATH / model_name / "li_diffusion_b-neb-results.dat", encoding="utf8" ) as f: data = f.readlines() pred_barrier, _, _ = tuple(float(x) for x in data[1].split()) @@ -111,7 +111,7 @@ def path_c_error() -> dict[str, float]: for model_name in MODELS: plot_nebs(model_name, "c") with open( - CALC_PATH / f"li_diffusion_c-{model_name}-neb-results.dat", encoding="utf8" + CALC_PATH / model_name / "li_diffusion_c-neb-results.dat", encoding="utf8" ) as f: data = f.readlines() pred_barrier, _, _ = tuple(float(x) for x in data[1].split()) @@ -158,4 +158,8 @@ def test_li_diffusion(metrics: dict[str, dict]) -> None: metrics All Li diffusion metrics. """ - return + write_struct_info( + data_path=CALC_PATH / "mock" / "li_diffusion_b-neb-band.extxyz", + out_path=OUT_PATH, + index=0, + ) diff --git a/ml_peg/analysis/utils/decorators.py b/ml_peg/analysis/utils/decorators.py index bc95d7b42..c5fda8229 100644 --- a/ml_peg/analysis/utils/decorators.py +++ b/ml_peg/analysis/utils/decorators.py @@ -1593,7 +1593,6 @@ def build_table( thresholds: Thresholds, filename: str = "table.json", metric_tooltips: dict[str, str] | None = None, - normalize: bool = True, normalizer: Callable[[float, float, float], float] | None = None, weights: dict[str, float] | None = None, mlip_name_map: dict[str, str] | None = None, @@ -1601,7 +1600,7 @@ def build_table( """ Build DataTable, including optional metric normalisation. - If `normalize` is `True`, by default each metric is normalised to 0-1 scale where: + By default each metric is normalised to 0-1 scale where: - Values <= Y get score 0 - Values >= X get score 1 - Values between Y and X scale linearly, by default. @@ -1616,8 +1615,6 @@ def build_table( Filename to save table. Default is "table.json". metric_tooltips Tooltips for table metric headers. Defaults are set for "MLIP" and "Score". - normalize - Whether to apply normalisation when calculating the score. Default is True. normalizer Optional function to map (value, X, Y) -> normalised score. Default is ml_peg.analysis.utils.utils.normalize_metric. @@ -1726,15 +1723,9 @@ def build_table_wrapper(*args, **kwargs) -> dict[str, Any]: summary_tooltips = { "MLIP": "Model identifier, hover for configuration details.", } - if normalize: - summary_tooltips["Score"] = ( - "Weighted score across metrics, " - "Higher is better (normalised 0 to 1)." - ) - else: - summary_tooltips["Score"] = ( - "Weighted score across metrics, higher is better." - ) + summary_tooltips["Score"] = ( + "Weighted score across metrics, Higher is better (normalised 0 to 1)." + ) if metric_tooltips: tooltip_header = metric_tooltips | summary_tooltips @@ -1746,15 +1737,12 @@ def build_table_wrapper(*args, **kwargs) -> dict[str, Any]: metric_weights.setdefault(column, 1.0) # Calculate scores, including any normalisation - if normalize: - metrics_data = calc_table_scores( - metrics_data=metrics_data, - thresholds=thresholds, - normalizer=normalizer, - weights=metric_weights, - ) - else: - metrics_data = calc_table_scores(metrics_data, weights=metric_weights) + metrics_data = calc_table_scores( + metrics_data=metrics_data, + thresholds=thresholds, + normalizer=normalizer, + weights=metric_weights, + ) table = dash_table.DataTable( metrics_data, diff --git a/ml_peg/analysis/utils/utils.py b/ml_peg/analysis/utils/utils.py index bddffba04..16bfc8133 100644 --- a/ml_peg/analysis/utils/utils.py +++ b/ml_peg/analysis/utils/utils.py @@ -4,9 +4,11 @@ from collections import defaultdict from collections.abc import Callable, Iterable +import json from pathlib import Path from typing import Any +from ase import Atoms from ase.io import read, write from matplotlib import colormaps from matplotlib.colors import Colormap @@ -465,11 +467,14 @@ def calc_table_scores( # Strict mode: require all metrics to be present metrics_row["Score"] = None elif scores_list: - # Calculate weighted average of available metrics - try: - metrics_row["Score"] = np.average(scores_list, weights=weights_list) - except ZeroDivisionError: - metrics_row["Score"] = np.mean(scores_list) + if np.nan in scores_list: + metrics_row["Score"] = np.nan + else: + # Calculate weighted average of available metrics + try: + metrics_row["Score"] = np.average(scores_list, weights=weights_list) + except ZeroDivisionError: + metrics_row["Score"] = np.mean(scores_list) else: metrics_row["Score"] = None @@ -726,7 +731,7 @@ def normalize_metric( try: # Handle NaNs robustly if np.isnan([value, good_threshold, bad_threshold]).any(): - return None + return np.nan except TypeError: return None @@ -738,3 +743,101 @@ def normalize_metric( # Clip to [0, 1] return max(min(1.0, float(t)), 0.0) + + +def get_struct_info( + calc_path: Path, + glob_pattern: str = "*.xyz", + index: str = ":", + info_keys: list[str] | None = None, + write_info: bool = True, + write_structs: bool = True, + out_path: Path | None = None, +) -> dict[str, Any]: + """ + Get info from structure files. + + Parameters + ---------- + calc_path + Path to calculation outputs. + glob_pattern + Glob pattern to match structure files. + index + Index to read from structure files. Default is ":". + info_keys + List of info keys to extract from structure files. Default is None. + write_info + Whether to write out info for each system. Default is True. Requires `out_path`. + write_structs + Whether to write out structure files for each system. Default is `True` if + `out_path` is specified. + out_path + Path to write out info for each system. Required if `write_info` is `True`. + + Returns + ------- + dict[str, Any] + Dictionary of system info for all systems. + """ + info_keys = info_keys or [] + info = {"elements": []} | dict.fromkeys(info_keys, []) + + model_dir = calc_path / "mock" + if not model_dir.exists(): + raise ValueError(f"{model_dir} does not exist. Please run mock calculation.") + files = sorted(model_dir.glob(glob_pattern)) + if not files: + raise ValueError( + f"No file matches in {model_dir}. Please run mock calculation." + ) + + for file in files: + structs = read(file, index=index) + if isinstance(structs, Atoms): + structs = [structs] + for struct in structs: + for key in info_keys: + info[key].append(struct.info[key]) + info["elements"].append(sorted(set(struct.get_chemical_symbols()))) + if write_structs: + (out_path / "mock").mkdir(parents=True, exist_ok=True) + write(out_path / "mock" / file.name, structs) + + if write_info: + if out_path is None: + raise ValueError("`out_path` must be specified if `write_info` is `True`.") + out_path.mkdir(parents=True, exist_ok=True) + out_file = out_path / "info.json" + with out_file.open("w", encoding="utf8") as f: + json.dump(info, f, indent=1) + + return info + + +def write_struct_info(data_path: Path, out_path: Path, index: str = ":") -> None: + """ + Write out element info on structures used in benchmark. + + Parameters + ---------- + data_path + Path to calculation structure file. + out_path + Path to write out info. + index + Index to read from structure files. Default is ":". + """ + elements = [] + if not data_path.exists(): + raise ValueError(f"{data_path} does not exist. Please run mock calculation.") + + structs = read(data_path, index=index) + if isinstance(structs, Atoms): + structs = [structs] + for struct in structs: + elements.append(sorted(set(struct.get_chemical_symbols()))) + + out_path.mkdir(parents=True, exist_ok=True) + with (out_path / "info.json").open("w", encoding="utf8") as f: + json.dump({"elements": elements}, f, indent=1) diff --git a/ml_peg/app/base_app.py b/ml_peg/app/base_app.py index 1fc1ecfce..bd7581eb9 100644 --- a/ml_peg/app/base_app.py +++ b/ml_peg/app/base_app.py @@ -3,8 +3,12 @@ from __future__ import annotations from abc import ABC, abstractmethod +from copy import deepcopy +import json from pathlib import Path +import warnings +from dash.dcc import Store from dash.development.base_component import Component from dash.html import Div @@ -31,7 +35,9 @@ class BaseApp(ABC): URL for online documentation. Default is None. framework_id Framework identifier used for benchmark attribution tags. Default is - ``"ml_peg"``. + `"ml_peg"`. + info_path + Path to json file containing additional info for filtering. Default is None. """ def __init__( @@ -42,6 +48,7 @@ def __init__( extra_components: list[Component], docs_url: str | None = None, framework_id: str = "ml_peg", + info_path: Path | None = None, ): """ Initiaise class. @@ -60,6 +67,9 @@ def __init__( URL to online documentation. Default is None. framework_id Framework identifier used for benchmark attribution tags. + Default is `"ml_peg"`. + info_path + Path to json file containing additional info for filtering. Default is None. """ self.name = name self.description = description @@ -71,7 +81,31 @@ def __init__( self.table = rebuild_table( self.table_path, id=self.table_id, description=description ) + self.original_table = deepcopy(self.table) self.layout = self.build_layout() + if info_path: + self.load_info(info_path) + else: + self.info = None + warnings.warn("No info_path provided.", stacklevel=2) + if hasattr(self, "set_elements"): + self.set_elements() + else: + self.elements = None + + def load_info(self, info_path: Path) -> None: + """ + Load additional info for app. + + Parameters + ---------- + info_path + Path to json file containing additional info for filtering. + """ + if not info_path.exists(): + warnings.warn(f"{info_path} does not exist, skipping.", stacklevel=2) + with open(info_path) as f: + self.info = json.load(f) def build_layout(self) -> Div: """ @@ -91,7 +125,7 @@ def build_layout(self) -> Div: framework_id=self.framework_id, table=self.table, column_widths=getattr(self.table, "column_widths", None), - thresholds=getattr(self.table, "thresholds", None), + thresholds=self.table.thresholds, extra_components=self.extra_components, ) @@ -99,3 +133,53 @@ def build_layout(self) -> Div: def register_callbacks(self): """Register callbacks with app.""" pass + + def filter_table(self, filter_elements: list[str] | None) -> None: + """ + Filter data by elements. + + Parameters + ---------- + filter_elements + List of elements to filter out of data. + + Returns + ------- + dict[str, dict] + Updated benchmark table. + """ + print(f"No filter_data method defined for {self.name}, skipping.") + return self.table.data + + @property + def stores(self) -> list[Store]: + """ + List Stores to be registered with full app. + + Returns + ------- + list[Store] + List of Stores to be registered with full app. + """ + return [ + Store( + id=f"{self.table_id}-computed-store", + storage_type="session", + data=self.table.data, + ), + Store( + id=f"{self.table_id}-raw-data-store", + storage_type="session", + data=self.table.data, + ), + Store( + id=f"{self.table_id}-weight-store", + storage_type="session", + data=self.table.weights, + ), + Store( + id=f"{self.table_id}-thresholds-store", + storage_type="session", + data=self.table.thresholds, + ), + ] diff --git a/ml_peg/app/build_app.py b/ml_peg/app/build_app.py index dc109c6a1..d2ed9a747 100644 --- a/ml_peg/app/build_app.py +++ b/ml_peg/app/build_app.py @@ -14,6 +14,7 @@ from ml_peg.analysis.utils.utils import calc_table_scores, get_table_style from ml_peg.app import APP_ROOT +from ml_peg.app.filter import get_element_filter, get_model_filter from ml_peg.app.utils.build_components import ( build_download_controls, build_faqs, @@ -25,7 +26,10 @@ build_tutorial_button, register_onboarding_callbacks, ) -from ml_peg.app.utils.register_callbacks import register_benchmark_to_category_callback +from ml_peg.app.utils.register_callbacks import ( + register_benchmark_to_category_callback, + register_filter_tables_callback, +) from ml_peg.app.utils.utils import ( build_level_of_theory_warnings, get_framework_config, @@ -342,6 +346,7 @@ def build_sidebar( def get_all_tests( category: str = "*", ) -> tuple[ + dict[str, dict[str, Dash]], dict[str, dict[str, list[Div]]], dict[str, dict[str, DataTable]], dict[str, dict[str, str]], @@ -357,12 +362,13 @@ def get_all_tests( Returns ------- tuple - Layouts, tables, and framework IDs for all categories. + Apps by test name, and layouts, tables, and framework IDs for all categories. """ # Find Python files e.g. app_OC157.py in mlip_tesing.app module. # We will get the category from the parent's parent directory # E.g. ml_peg/app/surfaces/OC157/app_OC157.py -> surfaces tests = APP_ROOT.glob(f"{category}/*/app*.py") + apps = {} layouts = {} tables = {} frameworks = {} @@ -377,15 +383,18 @@ def get_all_tests( f"ml_peg.app.{category_name}.{test_name}.app_{test_name}" ) test_app = test_module.get_app() + apps[test_name] = test_app # Get layouts and tables for each category/test if category_name not in layouts: layouts[category_name] = {} tables[category_name] = {} frameworks[category_name] = {} + layouts[category_name][test_app.name] = test_app.layout tables[category_name][test_app.name] = test_app.table frameworks[category_name][test_app.name] = test_app.framework_id + except FileNotFoundError as err: warnings.warn( f"Unable to load layout for {test_name} in {category_name} category. " @@ -405,7 +414,7 @@ def get_all_tests( ) continue - return layouts, tables, frameworks + return apps, layouts, tables, frameworks def build_category( @@ -439,6 +448,7 @@ def build_category( category_views = {} category_tables = {} category_weights = {} + category_to_title = {} framework_ids: set[str] = set() # `category` corresponds to the category's directory name @@ -458,6 +468,8 @@ def build_category( category_weight = 1 benchmark_weights = {} + category_to_title[category] = category_title + # Build category summary table summary_table = build_summary_table( dict(sorted(all_tables[category].items())), @@ -475,7 +487,6 @@ def build_category( weight_components = build_weight_components( header="Weights", table=summary_table, - include_store=False, include_download_controls=False, column_widths=getattr(summary_table, "column_widths", None), ) @@ -500,15 +511,9 @@ def build_category( "tests": test_entries, } - # Register benchmark table -> category table callbacks - # Category summary table columns add "Score" to name for clarity - for test_name, benchmark_table in all_tables[category].items(): - register_benchmark_to_category_callback( - benchmark_table_id=benchmark_table.id, - category_table_id=f"{category_title}-summary-table", - benchmark_column=test_name + " Score", - model_name_map=getattr(benchmark_table, "model_name_map", None), - ) + # Register callback for all benchmark tables -> category table + # Category summary table columns add "Score" to name for clarity + register_benchmark_to_category_callback(all_tables, category_to_title) return category_views, category_tables, category_weights, framework_ids @@ -814,9 +819,9 @@ def build_summary_table( tooltip_data=tooltip_rows, tooltip_delay=100, tooltip_duration=None, - persistence=True, - persistence_type="session", - persisted_props=["data"], + # persistence=True, + # persistence_type="session", + # persisted_props=["data"], tooltip_header=tooltip_header, editable=False, fill_width=False, @@ -836,6 +841,7 @@ def build_nav( framework_views: dict[str, dict[str, object]], summary_table: DataTable, weight_components: Div, + all_apps: dict[str, Dash], ) -> None: """ Build page layouts and sidebar navigation. @@ -852,6 +858,8 @@ def build_nav( Summary table with score from each category. weight_components Weight sliders, text boxes and reset button. + all_apps + Dictionary of all test apps. """ category_paths = { category_name: _category_to_path(category_name) @@ -869,43 +877,6 @@ def build_nav( framework_id: framework_views[framework_id]["label"] for framework_id in framework_order } - model_options = [{"label": m, "value": m} for m in MODELS] - - model_filter = Details( - [ - Summary( - "Visible models", - style={ - "cursor": "pointer", - "fontWeight": "600", - "fontSize": "11px", - "textTransform": "uppercase", - "letterSpacing": "0.07em", - "color": "#6c757d", - "padding": "5px", - }, - ), - Div( - [ - Dropdown( - id="model-filter-checklist", - options=model_options, - value=MODELS, - multi=True, - maxHeight=600, - optionHeight=10, - placeholder="Select visible models", - closeOnSelect=False, - style={"fontSize": "12px"}, - ), - ], - style={"padding": "8px 12px"}, - ), - ], - id="model-filter-details", - open=True, - style={"marginBottom": "8px", "fontSize": "13px"}, - ) _summary_label_style = { "cursor": "pointer", @@ -977,6 +948,11 @@ def build_nav( ), ] ) + + test_state_stores = [] + for app in all_apps.values(): + test_state_stores.extend(app.stores) + global_state_stores = [ Store( id="summary-table-weight-store", @@ -985,6 +961,7 @@ def build_nav( ), Store(id="cmap-store", storage_type="local", data="viridis_r"), *category_state_stores, + *test_state_stores, ] full_layout = [ @@ -1034,8 +1011,9 @@ def build_nav( sidebar, Div( [ - model_filter, + get_model_filter(MODELS), cmap_selector, + get_element_filter(), Store( id="selected-models-store", storage_type="session", @@ -1261,11 +1239,13 @@ def build_full_app(full_app: Dash, category: str = "*") -> None: Category to build app for. Default is `*`, corresponding to all categories. """ # Get layouts and tables for each test, grouped by categories - all_layouts, all_tables, all_frameworks = get_all_tests(category=category) + all_apps, all_layouts, all_tables, all_frameworks = get_all_tests(category=category) if not all_layouts: raise ValueError("No tests were built successfully") + register_filter_tables_callback(all_apps) + # Combine tests into categories and create category summary cat_views, cat_tables, cat_weights, framework_ids = build_category( all_layouts, all_tables, all_frameworks @@ -1278,7 +1258,6 @@ def build_full_app(full_app: Dash, category: str = "*") -> None: weight_components = build_weight_components( header="Weights", table=summary_table, - include_store=False, include_download_controls=False, column_widths=summary_table.column_widths, ) @@ -1289,5 +1268,6 @@ def build_full_app(full_app: Dash, category: str = "*") -> None: framework_views, summary_table, weight_components, + all_apps, ) register_onboarding_callbacks() diff --git a/ml_peg/app/filter.py b/ml_peg/app/filter.py new file mode 100644 index 000000000..256977cb3 --- /dev/null +++ b/ml_peg/app/filter.py @@ -0,0 +1,108 @@ +"""Build data components.""" + +from __future__ import annotations + +from ase.data import chemical_symbols +from dash.dcc import Dropdown +from dash.html import Details, Div, Summary + + +def get_model_filter(models) -> Details: + """ + Get model filter component. + + Parameters + ---------- + models + List of model names to include in filter options. + + Returns + ------- + Details + Model filter component. + """ + model_options = [{"label": m, "value": m} for m in models] + + return Details( + [ + Summary( + "Visible models", + style={ + "cursor": "pointer", + "fontWeight": "600", + "fontSize": "11px", + "textTransform": "uppercase", + "letterSpacing": "0.07em", + "color": "#6c757d", + "padding": "5px", + }, + ), + Div( + [ + Dropdown( + id="model-filter-checklist", + options=model_options, + value=models, + multi=True, + maxHeight=600, + optionHeight=10, + placeholder="Select visible models", + closeOnSelect=False, + style={"fontSize": "12px"}, + ), + ], + style={"padding": "8px 12px"}, + ), + ], + id="model-filter-details", + open=True, + style={"marginBottom": "8px", "fontSize": "13px"}, + ) + + +def get_element_filter() -> Details: + """ + Get element filter component. + + Returns + ------- + Details + Element filter component. + """ + # Exclude placeholder symbol for index 0 + elements = chemical_symbols[1:] + + return Details( + [ + Summary( + "Filter by elements", + style={ + "cursor": "pointer", + "fontWeight": "600", + "fontSize": "11px", + "textTransform": "uppercase", + "letterSpacing": "0.07em", + "color": "#6c757d", + "padding": "5px", + }, + ), + Div( + [ + Dropdown( + id="element-filter", + options=elements, + value=None, + multi=True, + placeholder="Filter elements", + closeOnSelect=False, + style={"fontSize": "13px"}, + debounce=True, + ), + ], + style={"padding": "8px 12px"}, + ), + ], + id="element-filter-details", + open=True, + style={"marginBottom": "8px", "fontSize": "13px"}, + ) diff --git a/ml_peg/app/molecular_crystal/DMC_ICE13/app_DMC_ICE13.py b/ml_peg/app/molecular_crystal/DMC_ICE13/app_DMC_ICE13.py index f76e7625c..f09f76a40 100644 --- a/ml_peg/app/molecular_crystal/DMC_ICE13/app_DMC_ICE13.py +++ b/ml_peg/app/molecular_crystal/DMC_ICE13/app_DMC_ICE13.py @@ -2,47 +2,58 @@ from __future__ import annotations +import warnings + from dash import Dash +from dash.dcc import Graph from dash.html import Div +import numpy as np +from ml_peg.analysis.molecular_crystal.DMC_ICE13.analyse_DMC_ICE13 import get_metrics from ml_peg.app import APP_ROOT from ml_peg.app.base_app import BaseApp from ml_peg.app.utils.build_callbacks import ( + filter_table, plot_from_table_column, struct_from_scatter, ) from ml_peg.app.utils.load import read_plot -from ml_peg.models import current_models -from ml_peg.models.get_models import get_model_names # Get all models -MODELS = get_model_names(current_models) BENCHMARK_NAME = "DMC-ICE13 Lattice Energies" DOCS_URL = "https://ddmms.github.io/ml-peg/user_guide/benchmarks/molecular_crystal.html#dmc-ice13" DATA_PATH = APP_ROOT / "data" / "molecular_crystal" / "DMC_ICE13" +INFO_PATH = DATA_PATH / "info.json" class DMCICE13App(BaseApp): """DMC-ICE13 benchmark app layout and callbacks.""" - def register_callbacks(self) -> None: - """Register callbacks to app.""" - scatter = read_plot( + def load_data(self) -> None: + """Load data required for filtering.""" + self.data = read_plot( DATA_PATH / "figure_lattice_energies.json", id=f"{BENCHMARK_NAME}-figure", ) + def register_callbacks(self) -> None: + """Register callbacks to app.""" + if not hasattr(self, "data"): + self.load_data() + # Assets dir will be parent directory - individual files for each polymorph - structs_dir = DATA_PATH / MODELS[0] + structs_dir = DATA_PATH / "mock" + if not structs_dir.exists(): + warnings.warn(f"Structures directory {structs_dir} not found", stacklevel=2) structs = [ - f"/assets/molecular_crystal/DMC_ICE13/{MODELS[0]}/{struct_file.stem}.xyz" + f"/assets/molecular_crystal/DMC_ICE13/mock/{struct_file.stem}.xyz" for struct_file in sorted(structs_dir.glob("*.xyz")) ] plot_from_table_column( table_id=self.table_id, plot_id=f"{BENCHMARK_NAME}-figure-placeholder", - column_to_plot={"MAE": scatter}, + column_to_plot={"MAE": self.data}, ) struct_from_scatter( @@ -52,6 +63,65 @@ def register_callbacks(self) -> None: mode="struct", ) + # Ensure data and elements are loaded + if not hasattr(self, "data"): + self.load_data() + if not hasattr(self, "elements"): + self.get_elements() + + filter_table( + table_id=self.table_id, + filter_func=self.filter_data, + filter_kwargs={"data": self.data, "test_elements": self.elements}, + ) + + def get_elements(self) -> None: + """Get element sets for filtering.""" + try: + self.elements = [set(entry) for entry in self.info["elements"]] + except (AttributeError, KeyError, TypeError): + self.elements = [] + warnings.warn("Unable to read elements lists.", stacklevel=2) + + @staticmethod + def filter_data( + filter_elements: set[str], data: Graph, test_elements: list[set[str]] + ) -> dict[str, dict]: + """ + Apply elements filter to data. + + Parameters + ---------- + filter_elements + Set of elements to filter out of data. + data + Scatter plot to filter. + test_elements + List of element for each system. + + Returns + ------- + dict[str, dict] + Metric names and values for all models. + """ + # Get overlap of deselected elements with each system's elements + filtered_indices = [ + not bool(elements & filter_elements) for elements in test_elements + ] + + results = {} + ref_filtered = False + + for plot in data.figure.data: + # Ignore unamed (parity) line + if plot.name: + results[plot.name] = np.array(plot.x)[filtered_indices].tolist() + if not ref_filtered: + results["ref"] = np.array(plot.y)[filtered_indices].tolist() + ref_filtered = True + + return get_metrics(results) + def get_app() -> DMCICE13App: """ @@ -71,6 +141,7 @@ def get_app() -> DMCICE13App: Div(id=f"{BENCHMARK_NAME}-figure-placeholder"), Div(id=f"{BENCHMARK_NAME}-struct-placeholder"), ], + info_path=INFO_PATH, ) diff --git a/ml_peg/app/molecular_crystal/X23/app_X23.py b/ml_peg/app/molecular_crystal/X23/app_X23.py index 437c64f24..3e5bffb5c 100644 --- a/ml_peg/app/molecular_crystal/X23/app_X23.py +++ b/ml_peg/app/molecular_crystal/X23/app_X23.py @@ -2,49 +2,60 @@ from __future__ import annotations +import warnings + from dash import Dash +from dash.dcc import Graph from dash.html import Div +import numpy as np +from ml_peg.analysis.molecular_crystal.X23.analyse_X23 import get_metrics from ml_peg.app import APP_ROOT from ml_peg.app.base_app import BaseApp from ml_peg.app.utils.build_callbacks import ( + filter_table, plot_from_table_column, struct_from_scatter, ) from ml_peg.app.utils.load import read_plot -from ml_peg.models import current_models -from ml_peg.models.get_models import get_model_names # Get all models -MODELS = get_model_names(current_models) BENCHMARK_NAME = "X23 Lattice Energies" DOCS_URL = ( "https://ddmms.github.io/ml-peg/user_guide/benchmarks/molecular_crystal.html#x23" ) DATA_PATH = APP_ROOT / "data" / "molecular_crystal" / "X23" +INFO_PATH = DATA_PATH / "info.json" class X23App(BaseApp): """X23 benchmark app layout and callbacks.""" - def register_callbacks(self) -> None: - """Register callbacks to app.""" - scatter = read_plot( + def load_data(self) -> None: + """Load data required for filtering.""" + self.data = read_plot( DATA_PATH / "figure_lattice_energies.json", id=f"{BENCHMARK_NAME}-figure", ) + def register_callbacks(self) -> None: + """Register callbacks to app.""" + if not hasattr(self, "data"): + self.load_data() + # Assets dir will be parent directory - individual files for each system - structs_dir = DATA_PATH / MODELS[0] + structs_dir = DATA_PATH / "mock" + if not structs_dir.exists(): + warnings.warn(f"Structures directory {structs_dir} not found", stacklevel=2) structs = [ - f"/assets/molecular_crystal/X23/{MODELS[0]}/{struct_file.stem}.xyz" + f"/assets/molecular_crystal/X23/mock/{struct_file.stem}.xyz" for struct_file in sorted(structs_dir.glob("*.xyz")) ] plot_from_table_column( table_id=self.table_id, plot_id=f"{BENCHMARK_NAME}-figure-placeholder", - column_to_plot={"MAE": scatter}, + column_to_plot={"MAE": self.data}, ) struct_from_scatter( @@ -54,6 +65,65 @@ def register_callbacks(self) -> None: mode="struct", ) + # Ensure data and elements are loaded + if not hasattr(self, "data"): + self.load_data() + if not hasattr(self, "elements"): + self.get_elements() + + filter_table( + table_id=self.table_id, + filter_func=self.filter_data, + filter_kwargs={"data": self.data, "test_elements": self.elements}, + ) + + def get_elements(self) -> None: + """Get element sets for filtering from loaded info.""" + try: + self.elements = [set(entry) for entry in self.info["elements"]] + except (AttributeError, KeyError, TypeError): + self.elements = [] + warnings.warn("Unable to read elements lists.", stacklevel=2) + + @staticmethod + def filter_data( + filter_elements: set[str], data: Graph, test_elements: list[set[str]] + ) -> dict[str, dict]: + """ + Apply elements filter to data. + + Parameters + ---------- + filter_elements + Set of elements to filter out of data. + data + Scatter plot to filter. + test_elements + List of element for each system. + + Returns + ------- + dict[str, dict] + Metric names and values for all models. + """ + # Get overlap of deselected elements with each system's elements + filtered_indices = [ + not bool(elements & filter_elements) for elements in test_elements + ] + + results = {} + ref_filtered = False + + for plot in data.figure.data: + # Ignore unamed (parity) line + if plot.name: + results[plot.name] = np.array(plot.x)[filtered_indices].tolist() + if not ref_filtered: + results["ref"] = np.array(plot.y)[filtered_indices].tolist() + ref_filtered = True + + return get_metrics(results) + def get_app() -> X23App: """ @@ -73,6 +143,7 @@ def get_app() -> X23App: Div(id=f"{BENCHMARK_NAME}-figure-placeholder"), Div(id=f"{BENCHMARK_NAME}-struct-placeholder"), ], + info_path=INFO_PATH, ) diff --git a/ml_peg/app/nebs/li_diffusion/app_li_diffusion.py b/ml_peg/app/nebs/li_diffusion/app_li_diffusion.py index 99576f4af..08f6ea08d 100644 --- a/ml_peg/app/nebs/li_diffusion/app_li_diffusion.py +++ b/ml_peg/app/nebs/li_diffusion/app_li_diffusion.py @@ -2,15 +2,14 @@ from __future__ import annotations +import warnings + 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, - struct_from_scatter, -) +from ml_peg.app.utils.build_callbacks import plot_from_table_cell, struct_from_scatter from ml_peg.app.utils.load import read_plot from ml_peg.models import current_models from ml_peg.models.get_models import get_model_names @@ -30,11 +29,11 @@ def register_callbacks(self) -> None: scatter_plots = { model: { "Path B error": read_plot( - DATA_PATH / f"figure_{model}_neb_b.json", + DATA_PATH / model / "figure_neb_b.json", id=f"{BENCHMARK_NAME}-{model}-figure-b", ), "Path C error": read_plot( - DATA_PATH / f"figure_{model}_neb_c.json", + DATA_PATH / model / "figure_neb_c.json", id=f"{BENCHMARK_NAME}-{model}-figure-c", ), } @@ -45,8 +44,8 @@ def register_callbacks(self) -> None: assets_dir = "/assets/nebs/li_diffusion" structs = { model: { - "Path B error": f"{assets_dir}/{model}/{model}-b-neb-band.extxyz", - "Path C error": f"{assets_dir}/{model}/{model}-c-neb-band.extxyz", + "Path B error": f"{assets_dir}/{model}/b-neb-band.extxyz", + "Path C error": f"{assets_dir}/{model}/c-neb-band.extxyz", } for model in MODELS } @@ -66,6 +65,44 @@ def register_callbacks(self) -> None: mode="traj", ) + def set_elements(self) -> None: + """Get element sets for filtering.""" + try: + self.elements = set(self.info["elements"][0]) + except (AttributeError, KeyError, TypeError): + self.elements = set() + warnings.warn("Unable to read elements lists.", stacklevel=2) + + def filter_table(self, filter_elements: list[str] | None) -> dict[str, dict]: + """ + Apply elements filter to data. + + Parameters + ---------- + filter_elements + List of elements to filter out of data. + + Returns + ------- + dict[str, dict] + Metric names and values for all models. + """ + filter_elements = set(filter_elements) if filter_elements else set() + + # Get overlap of deselected elements with each system's elements + if bool(self.elements & filter_elements): + for row in self.table.data: + row["Path B error"] = None + row["Path C error"] = None + else: + for current_row, original_row in zip( + self.table.data, self.original_table.data, strict=True + ): + current_row["Path B error"] = original_row["Path B error"] + current_row["Path C error"] = original_row["Path C error"] + + return self.table.data + def get_app() -> LiDiffusionApp: """ @@ -85,6 +122,7 @@ def get_app() -> LiDiffusionApp: Div(id=f"{BENCHMARK_NAME}-figure-placeholder"), Div(id=f"{BENCHMARK_NAME}-struct-placeholder"), ], + info_path=DATA_PATH / "info.json", ) diff --git a/ml_peg/app/utils/build_components.py b/ml_peg/app/utils/build_components.py index 66de62710..196d612ad 100644 --- a/ml_peg/app/utils/build_components.py +++ b/ml_peg/app/utils/build_components.py @@ -139,7 +139,6 @@ def build_weight_components( *, use_thresholds: bool = False, include_download_controls: bool = True, - include_store: bool = True, column_widths: dict[str, int] | None = None, thresholds: Thresholds | None = None, ) -> Div: @@ -158,10 +157,6 @@ def build_weight_components( recompute Scores consistently. include_download_controls Whether to render download controls in the Score column slot. - include_store - Whether to include this table's weight ``dcc.Store`` in the returned - component. Set to ``False`` when that store is already created elsewhere, - for example in the main app layout. column_widths Optional mapping of table column IDs to pixel widths used to align the inputs with the rendered table. @@ -292,14 +287,6 @@ def build_weight_components( ) layout = [container] - if include_store: - layout.append( - Store( - id=f"{table.id}-weight-store", - storage_type="session", - data=weights, - ) - ) model_levels = getattr(table, "model_levels_of_theory", None) metric_levels = getattr(table, "metric_levels_of_theory", None) @@ -733,10 +720,10 @@ def build_test_layout( description: str, framework_id: str, table: DataTable, + thresholds: Thresholds, extra_components: list[Component] | None = None, docs_url: str | None = None, column_widths: dict[str, int] | None = None, - thresholds: Thresholds | None = None, ) -> Div: """ Build app layout for a test. @@ -752,6 +739,9 @@ def build_test_layout( table Dash Table with metric results. Can include a `weights` attribute to be used by `build_weight_components`. + thresholds + Normalization metadata (metric -> (good, bad, unit)) supplied via the + analysis pipeline. Inline threshold controls are rendered automatically. extra_components List of Dash Components to include after the metrics table. docs_url @@ -759,10 +749,6 @@ def build_test_layout( column_widths Optional column-width mapping inferred from analysis output. Used to align threshold controls beneath the table columns when available. - thresholds - Optional normalization metadata (metric -> (good, bad, unit)) supplied via the - analysis pipeline. When provided, inline threshold controls are rendered - automatically. Returns ------- @@ -821,33 +807,32 @@ def build_test_layout( ) ) - # Inline normalization thresholds when metadata is supplied - threshold_controls = None - if thresholds is not None: - reserved = {"MLIP", "Score", "id"} - metric_columns = [ - col["id"] for col in table.columns if col.get("id") not in reserved - ] - layout_contents.append( - Store( - id=f"{table.id}-raw-data-store", - storage_type="session", - data=table.data, - ) - ) - layout_contents.append( - Store( - id=f"{table.id}-raw-tooltip-store", - storage_type="session", - data=table.tooltip_header, - ) + reserved = {"MLIP", "Score", "id"} + metric_columns = [ + col["id"] for col in table.columns if col.get("id") not in reserved + ] + + layout_contents.append( + Store( + id=f"{table.id}-raw-data-store", + storage_type="session", + data=table.data, ) - threshold_controls = build_threshold_inputs( - table_columns=metric_columns, - thresholds=thresholds, - table_id=table.id, - column_widths=column_widths, + ) + layout_contents.append( + Store( + id=f"{table.id}-raw-tooltip-store", + storage_type="session", + data=table.tooltip_header, ) + ) + + threshold_controls = build_threshold_inputs( + table_columns=metric_columns, + thresholds=thresholds, + table_id=table.id, + column_widths=column_widths, + ) # Add metric-weight controls for every benchmark table metric_weights = build_weight_components( @@ -862,24 +847,21 @@ def build_test_layout( # Build the controls element before the table wrapper so both can go into the # same fit-content div. The controls use width:100% of that wrapper, which # equals the table width, keeping the columns aligned. - if thresholds is not None: - controls_visual = Div( - [ - Div(threshold_controls, style={"marginBottom": "0px"}), - Div(metric_weights, style={"marginTop": "0"}), - ], - style={ - "backgroundColor": "#f8f9fa", - "border": "1px solid #dee2e6", - "borderRadius": "6px", - "padding": "0px 0px 0px 0px", # top right bottom left - "marginTop": "-5px", - "boxSizing": "border-box", - "width": "100%", - }, - ) - else: - controls_visual = metric_weights + controls_visual = Div( + [ + Div(threshold_controls, style={"marginBottom": "0px"}), + Div(metric_weights, style={"marginTop": "0"}), + ], + style={ + "backgroundColor": "#f8f9fa", + "border": "1px solid #dee2e6", + "borderRadius": "6px", + "padding": "0px 0px 0px 0px", # top right bottom left + "marginTop": "-5px", + "boxSizing": "border-box", + "width": "100%", + }, + ) table_section = [ build_download_controls(table.id, row=True), @@ -1145,12 +1127,6 @@ def build_threshold_inputs( ) ) - store = Store( - id=f"{table_id}-thresholds-store", - storage_type="session", - data=default_thresholds, - ) - # Register callbacks for these metrics, pass default_thresholds for reset register_normalization_callbacks( table_id, @@ -1159,9 +1135,4 @@ def build_threshold_inputs( register_toggle=False, ) - return Div( - [ - Div(cells, id=f"{table_id}-threshold-grid", style=container_style), - store, - ] - ) + return Div([Div(cells, id=f"{table_id}-threshold-grid", style=container_style)]) diff --git a/ml_peg/app/utils/filter.py b/ml_peg/app/utils/filter.py new file mode 100644 index 000000000..9839efb00 --- /dev/null +++ b/ml_peg/app/utils/filter.py @@ -0,0 +1,61 @@ +"""Utility functions for filtering data.""" + +from __future__ import annotations + +from collections.abc import Callable + +from dash.dcc import Graph +import numpy as np + + +def filter_parity( + filter_elements: set[str], + data: Graph, + test_elements: list[set[str]], + metric_getter: Callable, + mask_to_getter: bool = False, +) -> dict[str, dict]: + """ + Apply elements filter to data. + + Parameters + ---------- + filter_elements + Set of elements to filter out of data. + data + Scatter plot to filter. + test_elements + List of element for each system. + metric_getter + Function to calculate metric from filtered data. + mask_to_getter + Whether to pass the mask to the metric getter. + + Returns + ------- + dict[str, dict] + Metric names and values for all models. + """ + # Get overlap of deselected elements with each system's elements + filtered_indices = [ + not bool(elements & filter_elements) for elements in test_elements + ] + + results = {} + ref_filtered = False + + for plot in data.figure.data: + # Ignore unamed (parity) line + if plot.name: + print() + print("PLOT NAME", plot.name) + results[plot.name] = np.array(plot.x)[filtered_indices].tolist() + print("RESULTS", results[plot.name]) + print() + if not ref_filtered: + results["ref"] = np.array(plot.y)[filtered_indices].tolist() + ref_filtered = True + + if mask_to_getter: + return metric_getter(results, mask=filtered_indices) + return metric_getter(results) diff --git a/ml_peg/app/utils/load.py b/ml_peg/app/utils/load.py index 225da0758..9fc66490f 100644 --- a/ml_peg/app/utils/load.py +++ b/ml_peg/app/utils/load.py @@ -202,9 +202,9 @@ def rebuild_table( } ], sort_action="native", - persistence=True, - persistence_type="session", - persisted_props=["data"], + # persistence=True, + # persistence_type="session", + # persisted_props=["data"], fill_width=False, ) diff --git a/ml_peg/app/utils/register_callbacks.py b/ml_peg/app/utils/register_callbacks.py index 5bbbef9ec..3a97fe9aa 100644 --- a/ml_peg/app/utils/register_callbacks.py +++ b/ml_peg/app/utils/register_callbacks.py @@ -9,6 +9,7 @@ from dash import ( MATCH, ClientsideFunction, + Dash, Input, Output, State, @@ -18,6 +19,7 @@ dcc, no_update, ) +from dash.dash_table import DataTable from dash.exceptions import PreventUpdate import pandas as pd @@ -276,6 +278,86 @@ def register_category_table_callbacks( model_configs Optional configuration metadata for each model. """ + + @callback( + Output(table_id, "data", allow_duplicate=True), + Output(table_id, "style_data_conditional", allow_duplicate=True), + Input(f"{table_id}-raw-data-store", "data"), + State(f"{table_id}-computed-store", "data"), + State(f"{table_id}-weight-store", "data"), + State(f"{table_id}-thresholds-store", "data"), + State(f"{table_id}-normalized-toggle", "value"), + State("selected-models-store", "data"), + State("cmap-store", "data"), + State(f"{table_id}-raw-tooltip-store", "data"), + State(table_id, "columns"), + prevent_initial_call=True, + optional=True, + ) + def update_table_from_store( + stored_raw_data: list[dict] | None, + stored_computed_data: list[dict] | None, + weights: dict[str, float] | None, + thresholds: dict | None, + toggle_value: list[str] | None, + selected_models: list[str] | None, + cmap_name: str | None, + raw_tooltips: dict[str, str] | None, + current_columns: list[dict] | None, + ) -> list[dict]: + """ + Update visible table from cached data when the raw data store changes. + + Parameters + ---------- + stored_raw_data + Stored raw table data. + stored_computed_data + Stored computed table data. + weights + Stored weights for the table. + thresholds + Stored thresholds for the table. + toggle_value + Value of toggle to show normalised values. + selected_models + List of model names currently selected in the model filter. + cmap_name + Colourmap name from the cmap store. + raw_tooltips + Stored raw tooltip text for the table. + current_columns + Current table columns. + + Returns + ------- + list[dict] + Updated rows for the visible table. + """ + display_rows = get_scores( + stored_raw_data, stored_computed_data, thresholds, toggle_value + ) + scored_rows = calc_metric_scores(stored_raw_data, thresholds=thresholds) + filtered_rows = filter_rows_by_models(display_rows, selected_models) + filtered_scores = filter_rows_by_models(scored_rows, selected_models) + style = ( + get_table_style( + filtered_rows, + scored_data=filtered_scores, + cmap_name=cmap_name or "viridis_r", + ) + if filtered_rows + else [] + ) + style, tooltip_data = apply_level_of_theory_warnings( + filtered_rows, + style, + model_levels=model_levels, + metric_levels=metric_levels, + model_configs=model_configs, + ) + return filtered_rows, style + # Benchmark tables if use_thresholds: @@ -286,7 +368,7 @@ def register_category_table_callbacks( Output(table_id, "columns", allow_duplicate=True), Output(table_id, "tooltip_header", allow_duplicate=True), Output(f"{table_id}-computed-store", "data", allow_duplicate=True), - Output(f"{table_id}-raw-data-store", "data"), + Output(f"{table_id}-raw-data-store", "data", allow_duplicate=True), Input(f"{table_id}-weight-store", "data"), Input(f"{table_id}-thresholds-store", "data"), Input("app-location", "pathname"), @@ -298,6 +380,7 @@ def register_category_table_callbacks( State(f"{table_id}-raw-tooltip-store", "data"), State(table_id, "columns"), prevent_initial_call="initial_duplicate", + optional=True, ) def update_benchmark_table_scores( stored_weights: dict[str, float] | None, @@ -592,83 +675,99 @@ def update_scores_store( def register_benchmark_to_category_callback( - benchmark_table_id: str, - category_table_id: str, - benchmark_column: str, - use_threshold_store: bool = False, - model_name_map: dict[str, str] | None = None, + all_tables: dict[str, dict[str, DataTable]], category_to_title: dict[str, str] ) -> None: """ Propagate a benchmark table's Score into its category summary table column. Parameters ---------- - benchmark_table_id - ID of the benchmark test table (e.g., "OC157-table"). - category_table_id - ID of the category summary table (e.g., "Surfaces-summary-table"). - benchmark_column - Column name in the category summary table corresponding to the benchmark. - use_threshold_store - Whether the benchmark table exposes a normalization store for metrics. - model_name_map - Optional mapping of displayed benchmark MLIP names -> original model names. + all_tables + Tables for all tests, grouped by category. + category_to_title + Dictionary mapping category directory names to their display titles/table IDs. """ - _ = use_threshold_store # cached rows handle normalization - # flag kept for compatibility with existing call sites - name_map = dict(model_name_map or {}) + all_info = {} + for category, tables in all_tables.items(): + all_info[category] = {} + for test_name, benchmark_table in tables.items(): + all_info[category][test_name] = { + "benchmark_table_id": benchmark_table.id, + "benchmark_column": test_name + " Score", + "model_name_map": getattr(benchmark_table, "model_name_map", {}), + } + + outputs = [] + inputs = [] + for category, category_info in sorted(all_info.items()): + category_table_id = f"{category_to_title[category]}-summary-table" + outputs.append( + Output(f"{category_table_id}-computed-store", "data", allow_duplicate=True) + ) - @callback( - Output(f"{category_table_id}-computed-store", "data", allow_duplicate=True), - Input(f"{benchmark_table_id}-computed-store", "data"), - State(f"{category_table_id}-weight-store", "data"), - State(f"{category_table_id}-computed-store", "data"), - prevent_initial_call=True, - ) - def update_category_from_benchmark( - benchmark_computed_store: list[dict] | None, - category_weights: dict[str, float] | None, - category_computed_store: list[dict] | None, - ) -> list[dict]: + inputs.extend( + [ + State(f"{category_table_id}-weight-store", "data"), + State(f"{category_table_id}-computed-store", "data"), + ] + ) + inputs.extend( + [ + Input(f"{table_info['benchmark_table_id']}-computed-store", "data") + for _, table_info in sorted(category_info.items()) + ] + ) + + @callback(outputs, inputs, prevent_initial_call=True) + def update_category_from_benchmark(*args) -> list[list[dict]]: """ - Update cached category summary rows from a benchmark's cached scores. + Update cached category summary rows from all benchmarks' cached scores. Parameters ---------- - benchmark_computed_store - Latest scored benchmark rows emitted by the benchmark table. - category_weights - Stored weights for the category summary metrics. - category_computed_store - Cached scored rows for the category summary. + *args + States and Inputs for all category summary tables and benchmark tables. + Ordered by category. For each category, the weights, computed store, and + benchmark computed stores are listed sequentially. Returns ------- - list[dict] - Refreshed cached rows for the category summary table. + list[list[dict]] + Refreshed cached rows for each category summary table. """ - if not category_computed_store: - raise PreventUpdate - if not benchmark_computed_store: - raise PreventUpdate - category_rows = deepcopy(category_computed_store) + # Rebuild inputs for each category + iterator = iter(args) + + all_category_rows = [] - benchmark_scores: dict[str, float] = {} - for row in benchmark_computed_store: - display_name = row.get("MLIP") - original_name = name_map.get(display_name, display_name) - score = row.get("Score") - if display_name is None or original_name is None or score is None: - continue - benchmark_scores[original_name] = score + for category, category_info in sorted(all_info.items()): + category_weights = next(iterator) + category_rows = deepcopy(next(iterator)) - for row in category_rows: - mlip = row.get("MLIP") - if mlip in benchmark_scores: - row[benchmark_column] = benchmark_scores[mlip] + for test_name, table_info in sorted(category_info.items()): + benchmark_rows = deepcopy(next(iterator)) + name_map = table_info["model_name_map"] - category_rows, _ = update_score_style(category_rows, category_weights) - return category_rows + benchmark_scores = {} + for row in benchmark_rows: + display_name = row.get("MLIP") + original_name = name_map.get(display_name, display_name) + score = row.get("Score") + if display_name is None or original_name is None: + continue + benchmark_scores[original_name] = score + + for row in category_rows: + mlip = row.get("MLIP") + if mlip in benchmark_scores: + row[all_info[category][test_name]["benchmark_column"]] = ( + benchmark_scores[mlip] + ) + + category_rows, _ = update_score_style(category_rows, category_weights) + all_category_rows.append(category_rows) + + return all_category_rows def register_weight_callbacks( @@ -919,6 +1018,7 @@ def sync_threshold_input_styles( State(f"{table_id}", "columns"), State("cmap-store", "data"), prevent_initial_call=True, + optional=True, ) def toggle_normalized_display( show_normalized: list[str] | None, @@ -1090,3 +1190,87 @@ def download_table( Input(f"{table_id}-download-request", "data"), prevent_initial_call=True, ) + + +def register_filter_tables_callback(apps: dict[str, Dash]) -> None: + """ + Update all tables when filter dropdown value changes. + + Parameters + ---------- + apps + Dictionary of test apps to register callbacks for. + """ + app_entries = [] + for app in apps.values(): + app_entries.append( + { + "app": app, + "weight_state": State(f"{app.table_id}-weight-store", "data"), + "threshold_state": State(f"{app.table_id}-thresholds-store", "data"), + } + ) + + outputs = [] + for entry in sorted(app_entries): + app = entry["app"] + outputs.extend( + [ + Output(f"{app.table_id}-computed-store", "data"), + Output(f"{app.table_id}-raw-data-store", "data", allow_duplicate=True), + ] + ) + + states = [] + for entry in app_entries: + states.extend([entry["weight_state"], entry["threshold_state"]]) + + @callback( + outputs, Input("element-filter", "value"), states, prevent_initial_call=True + ) + def recompute_tables(elements, *args): + """ + Recompute all benchmark tables when element filter is applied. + + Parameters + ---------- + elements + List of selected elements to filter by. + *args + Weight and threshold states for each app. + + Returns + ------- + list[list[dict]] + Updated rows for each app's computed store and raw data stores. + """ + # Rebuild inputs for each app + per_app_state = {} + iterator = iter(args) + + for entry in sorted(app_entries): + app = entry["app"] + per_app_state[app.table_id] = { + "weights": next(iterator), + "thresholds": next(iterator), + } + + results = [] + + for entry in sorted(app_entries): + app = entry["app"] + state = per_app_state[app.table_id] + weights = state["weights"] + thresholds = state["thresholds"] + + updated_data = app.filter_table(elements) + + # Update overall table score for new weights and thresholds + metrics_data = calc_table_scores(updated_data, weights, thresholds) + + # Update stored scores per metric + scored_rows = calc_metric_scores(updated_data, thresholds) + + results.extend([scored_rows, metrics_data]) + + return results diff --git a/ml_peg/calcs/bulk_crystal/elasticity/calc_elasticity.py b/ml_peg/calcs/bulk_crystal/elasticity/calc_elasticity.py index 90628dad1..2886f237d 100644 --- a/ml_peg/calcs/bulk_crystal/elasticity/calc_elasticity.py +++ b/ml_peg/calcs/bulk_crystal/elasticity/calc_elasticity.py @@ -141,7 +141,6 @@ def process_result(self, result: dict | None, model_name: str) -> dict: # noqa: dict Dictionary of processed elastic properties. """ - key = get_crystal_system(result["final_structure"]) return { f"bulk_modulus_vrh_{model_name}": ( result["bulk_modulus_vrh"] * eVA3ToGPa @@ -158,7 +157,11 @@ def process_result(self, result: dict | None, model_name: str) -> dict: # noqa: if result is not None else float("nan") ), - f"crystal_system_{model_name}": key, + f"crystal_system_{model_name}": get_crystal_system( + result["final_structure"] + ) + if result is not None + else None, } diff --git a/ml_peg/calcs/bulk_crystal/high_pressure_relaxation/calc_high_pressure_relaxation.py b/ml_peg/calcs/bulk_crystal/high_pressure_relaxation/calc_high_pressure_relaxation.py index 7c274bf89..22e46e4ee 100644 --- a/ml_peg/calcs/bulk_crystal/high_pressure_relaxation/calc_high_pressure_relaxation.py +++ b/ml_peg/calcs/bulk_crystal/high_pressure_relaxation/calc_high_pressure_relaxation.py @@ -9,12 +9,14 @@ import random from typing import Any import urllib.request +from warnings import warn from ase import Atoms from ase.constraints import FixSymmetry from ase.io import write as ase_write from ase.units import GPa from janus_core.calculations.geom_opt import GeomOpt +import numpy as np import pandas as pd from pymatgen.entries.computed_entries import ComputedStructureEntry from pymatgen.io.ase import AseAtomsAdaptor @@ -200,13 +202,17 @@ def relax_with_pressure( converged = True counter += 1 atoms = relaxed - except Exception as e: - print(f"Relaxation failed: {e}") - return None, False, None + except Exception as exc: + warn(f"Relaxation failed: {exc}", stacklevel=2) + return None, False, np.nan if not converged: - return None, False, None + return None, False, np.nan # Calculate enthalpy: H = E + PV - energy = relaxed.get_potential_energy() + try: + energy = relaxed.get_potential_energy() + except Exception as exc: + warn(f"Error calculating energy: {exc}", stacklevel=2) + energy = np.nan volume = relaxed.get_volume() enthalpy = energy + pressure_gpa * GPa * volume n_atoms = len(relaxed) diff --git a/ml_peg/calcs/bulk_crystal/lattice_constants/calc_lattice_constants.py b/ml_peg/calcs/bulk_crystal/lattice_constants/calc_lattice_constants.py index deb3b8c88..2c0a87b93 100644 --- a/ml_peg/calcs/bulk_crystal/lattice_constants/calc_lattice_constants.py +++ b/ml_peg/calcs/bulk_crystal/lattice_constants/calc_lattice_constants.py @@ -6,6 +6,7 @@ import json from pathlib import Path from typing import Any +from warnings import warn from ase import Atoms from ase.build import bulk @@ -136,9 +137,12 @@ def test_lattice_consts(mlip: tuple[str, Any]) -> None: for name, crystal in crystals.items(): crystal.calc = copy(calc) - GeomOpt( - struct=crystal, - fmax=0.03, - write_traj=True, - file_prefix=OUT_PATH / model_name / name, - ).run() + try: + GeomOpt( + struct=crystal, + fmax=0.03, + write_traj=True, + file_prefix=OUT_PATH / model_name / name, + ).run() + except Exception as exc: + warn(f"Error during optimisation for {name}: {exc}", stacklevel=2) diff --git a/ml_peg/calcs/bulk_crystal/low_dimensional_relaxation/calc_low_dimensional_relaxation.py b/ml_peg/calcs/bulk_crystal/low_dimensional_relaxation/calc_low_dimensional_relaxation.py index 6f10dc914..9fb30d558 100644 --- a/ml_peg/calcs/bulk_crystal/low_dimensional_relaxation/calc_low_dimensional_relaxation.py +++ b/ml_peg/calcs/bulk_crystal/low_dimensional_relaxation/calc_low_dimensional_relaxation.py @@ -9,6 +9,7 @@ import random from typing import Any import urllib.request +from warnings import warn from ase import Atoms from ase.io import write as ase_write @@ -335,14 +336,21 @@ def relax_low_dimensional( converged = True counter += 1 atoms = relaxed - except Exception as e: - print(f"Relaxation failed: {e}") - return None, False, None, None + except Exception as exc: + warn(f"Relaxation failed: {exc}", stacklevel=2) + return None, False, np.nan, None if relaxed is None: - return None, False, None, max_force + return None, False, np.nan, max_force - energy_per_atom = relaxed.get_potential_energy() / len(relaxed) + try: + energy_per_atom = relaxed.get_potential_energy() / len(relaxed) + except Exception as exc: + warn( + f"Error calculating energy: {exc}", + stacklevel=2, + ) + energy_per_atom = np.nan return relaxed, converged, energy_per_atom, max_force diff --git a/ml_peg/calcs/conformers/37Conf8/calc_37Conf8.py b/ml_peg/calcs/conformers/37Conf8/calc_37Conf8.py index d9a4dce11..ac7d2e7e6 100644 --- a/ml_peg/calcs/conformers/37Conf8/calc_37Conf8.py +++ b/ml_peg/calcs/conformers/37Conf8/calc_37Conf8.py @@ -8,9 +8,11 @@ from pathlib import Path from typing import Any +from warnings import warn from ase import units from ase.io import read, write +import numpy as np import pandas as pd import pytest from tqdm import tqdm @@ -65,14 +67,26 @@ def test_37conf8_conformer_energies(mlip: tuple[str, Any]) -> None: zero_conf.info["charge"] = 0 zero_conf.info["spin"] = 1 zero_conf.calc = calc - e_model_zero_conf = zero_conf.get_potential_energy() + try: + e_model_zero_conf = zero_conf.get_potential_energy() + except Exception as exc: + warn(f"Error calculating energy for {label}: {exc}", stacklevel=2) + e_model_zero_conf = np.nan else: atoms = read(data_path / "PBEPBE-D3" / f"{label}_PBEPBE-D3.xyz") atoms.info["charge"] = 0 atoms.info["spin"] = 1 atoms.calc = calc - atoms.info["model_rel_energy"] = ( - atoms.get_potential_energy() - e_model_zero_conf - ) + try: + atoms.info["model_rel_energy"] = ( + atoms.get_potential_energy() - e_model_zero_conf + ) + except Exception as exc: + warn( + f"Error calculating relative energy for {label}: {exc}", + stacklevel=2, + ) + atoms.info["model_rel_energy"] = np.nan + atoms.info["ref_energy"] = float(df.iloc[i][2]) * KCAL_TO_EV write(write_dir / f"{label}.xyz", atoms) diff --git a/ml_peg/calcs/conformers/ACONFL/calc_ACONFL.py b/ml_peg/calcs/conformers/ACONFL/calc_ACONFL.py index 9324a0485..7233f8e3c 100644 --- a/ml_peg/calcs/conformers/ACONFL/calc_ACONFL.py +++ b/ml_peg/calcs/conformers/ACONFL/calc_ACONFL.py @@ -11,9 +11,11 @@ from pathlib import Path from typing import Any +from warnings import warn from ase import units from ase.io import read, write +import numpy as np import pytest from tqdm import tqdm @@ -65,9 +67,17 @@ def test_aconfl_conformer_energies(mlip: tuple[str, Any]) -> None: zero_atoms = read(data_path / zero_atoms_label / "struc.xyz") zero_atoms.calc = calc zero_atoms.info.update({"charge": 0, "spin": 1}) - atoms.info["model_rel_energy"] = ( - atoms.get_potential_energy() - zero_atoms.get_potential_energy() - ) + try: + atoms.info["model_rel_energy"] = ( + atoms.get_potential_energy() - zero_atoms.get_potential_energy() + ) + except Exception as exc: + warn( + f"Error calculating energy for {atoms_label} or " + f"{zero_atoms_label}: {exc}", + stacklevel=2, + ) + atoms.info["model_rel_energy"] = np.nan atoms.info["ref_rel_energy"] = ref_rel_energy write_dir = OUT_PATH / model_name diff --git a/ml_peg/calcs/conformers/DipCONFS/calc_DipCONFS.py b/ml_peg/calcs/conformers/DipCONFS/calc_DipCONFS.py index 313733ac6..a90186d84 100644 --- a/ml_peg/calcs/conformers/DipCONFS/calc_DipCONFS.py +++ b/ml_peg/calcs/conformers/DipCONFS/calc_DipCONFS.py @@ -13,9 +13,11 @@ from pathlib import Path from typing import Any +from warnings import warn from ase import units from ase.io import read, write +import numpy as np import pandas as pd import pytest from tqdm import tqdm @@ -97,14 +99,22 @@ def test_dipconfs(mlip: tuple[str, Any]) -> None: # Get zero ref conformer model energy zero_conf = get_atoms(data_path / zero_conf_label / "struc.xyz") zero_conf.calc = calc - e_model_zero_conf = zero_conf.get_potential_energy() + try: + e_model_zero_conf = zero_conf.get_potential_energy() + except Exception as exc: + warn(f"Error calculating energy for {zero_conf_label}: {exc}", stacklevel=2) + e_model_zero_conf = np.nan # Get current conformer model energy atoms = get_atoms(data_path / label / "struc.xyz") atoms.calc = calc - atoms.info["model_rel_energy"] = ( - atoms.get_potential_energy() - e_model_zero_conf - ) + try: + atoms.info["model_rel_energy"] = ( + atoms.get_potential_energy() - e_model_zero_conf + ) + except Exception as exc: + warn(f"Error calculating relative energy for {label}: {exc}", stacklevel=2) + atoms.info["model_rel_energy"] = np.nan atoms.info["ref_energy"] = e_rel_ref write_dir = OUT_PATH / model_name diff --git a/ml_peg/calcs/conformers/Glucose205/calc_Glucose205.py b/ml_peg/calcs/conformers/Glucose205/calc_Glucose205.py index 2c08c03d9..7e0648b47 100644 --- a/ml_peg/calcs/conformers/Glucose205/calc_Glucose205.py +++ b/ml_peg/calcs/conformers/Glucose205/calc_Glucose205.py @@ -10,9 +10,11 @@ from pathlib import Path from typing import Any +from warnings import warn from ase import Atoms, units from ase.io import read, write +import numpy as np import pandas as pd import pytest from tqdm import tqdm @@ -122,7 +124,11 @@ def test_glucose205(mlip: tuple[str, Any]) -> None: data_path / "Glucose_structures" / f"{lowest_conf_label}.xyz" ) conf_lowest.calc = calc - e_conf_lowest_model = conf_lowest.get_potential_energy() + try: + e_conf_lowest_model = conf_lowest.get_potential_energy() + except Exception as exc: + warn(f"Error calculating energy for {lowest_conf_label}: {exc}", stacklevel=2) + e_conf_lowest_model = np.nan for label, e_ref in tqdm(ref_energies.items()): # Skip the reference conformer for which the error is automatically zero @@ -131,9 +137,13 @@ def test_glucose205(mlip: tuple[str, Any]) -> None: atoms = get_atoms(data_path / "Glucose_structures" / f"{label}.xyz") atoms.calc = calc - atoms.info["model_rel_energy"] = ( - atoms.get_potential_energy() - e_conf_lowest_model - ) + try: + atoms.info["model_rel_energy"] = ( + atoms.get_potential_energy() - e_conf_lowest_model + ) + except Exception as exc: + warn(f"Error calculating relative energy for {label}: {exc}", stacklevel=2) + atoms.info["model_rel_energy"] = np.nan atoms.info["ref_energy"] = e_ref write_dir = OUT_PATH / model_name diff --git a/ml_peg/calcs/conformers/MPCONF196/calc_MPCONF196.py b/ml_peg/calcs/conformers/MPCONF196/calc_MPCONF196.py index 145bd924a..a033fabf2 100644 --- a/ml_peg/calcs/conformers/MPCONF196/calc_MPCONF196.py +++ b/ml_peg/calcs/conformers/MPCONF196/calc_MPCONF196.py @@ -9,6 +9,7 @@ from pathlib import Path from typing import Any +from warnings import warn from ase import Atoms, units from ase.io import read, write @@ -135,7 +136,11 @@ def test_mpconf196(mlip: tuple[str, Any]) -> None: atoms = get_atoms(data_path / xyz_fname) atoms.translate(-atoms.get_center_of_mass()) atoms.calc = calc - model_abs_energies.append(atoms.get_potential_energy()) + try: + model_abs_energies.append(atoms.get_potential_energy()) + except Exception as exc: + warn(f"Error calculating energy for {xyz_fname}: {exc}", stacklevel=2) + model_abs_energies.append(np.nan) ref_abs_energies.append(e_ref) current_molecule_labels.append(label) diff --git a/ml_peg/calcs/conformers/Maltose222/calc_Maltose222.py b/ml_peg/calcs/conformers/Maltose222/calc_Maltose222.py index 0b983b612..f3c503f64 100644 --- a/ml_peg/calcs/conformers/Maltose222/calc_Maltose222.py +++ b/ml_peg/calcs/conformers/Maltose222/calc_Maltose222.py @@ -10,9 +10,11 @@ from pathlib import Path from typing import Any +from warnings import warn from ase import units from ase.io import read, write +import numpy as np import pandas as pd import pytest from tqdm import tqdm @@ -122,7 +124,11 @@ def test_maltose222(mlip: tuple[str, Any]) -> None: data_path / "Maltose_structures" / f"{lowest_conf_label}.xyz" ) conf_lowest.calc = calc - e_conf_lowest_model = conf_lowest.get_potential_energy() + try: + e_conf_lowest_model = conf_lowest.get_potential_energy() + except Exception as exc: + warn(f"Error calculating energy for {lowest_conf_label}: {exc}", stacklevel=2) + e_conf_lowest_model = np.nan for label, e_ref in tqdm(ref_energies.items()): # Skip the reference conformer for which the error is automatically zero @@ -131,9 +137,13 @@ def test_maltose222(mlip: tuple[str, Any]) -> None: atoms = get_atoms(data_path / "Maltose_structures" / f"{label}.xyz") atoms.calc = calc - atoms.info["model_rel_energy"] = ( - atoms.get_potential_energy() - e_conf_lowest_model - ) + try: + atoms.info["model_rel_energy"] = ( + atoms.get_potential_energy() - e_conf_lowest_model + ) + except Exception as exc: + warn(f"Error calculating relative energy for {label}: {exc}", stacklevel=2) + atoms.info["model_rel_energy"] = np.nan atoms.info["ref_energy"] = e_ref write_dir = OUT_PATH / model_name diff --git a/ml_peg/calcs/conformers/OpenFF_Tors/calc_OpenFF_Tors.py b/ml_peg/calcs/conformers/OpenFF_Tors/calc_OpenFF_Tors.py index f8f164623..fb4a42180 100644 --- a/ml_peg/calcs/conformers/OpenFF_Tors/calc_OpenFF_Tors.py +++ b/ml_peg/calcs/conformers/OpenFF_Tors/calc_OpenFF_Tors.py @@ -10,6 +10,7 @@ import json from pathlib import Path from typing import Any +from warnings import warn from ase import Atoms, units from ase.io import write @@ -81,14 +82,25 @@ def test_openff_tors(mlip: tuple[str, Any]) -> None: if i == 0: e_ref_zero_conf = ref_energy * units.Hartree - e_model_zero_conf = atoms.get_potential_energy() + try: + e_model_zero_conf = atoms.get_potential_energy() + except Exception as exc: + warn(f"Error calculating energy for {label}: {exc}", stacklevel=2) + e_model_zero_conf = np.nan else: atoms.info["ref_rel_energy"] = ( ref_energy * units.Hartree - e_ref_zero_conf ) - atoms.info["model_rel_energy"] = ( - atoms.get_potential_energy() - e_model_zero_conf - ) + try: + atoms.info["model_rel_energy"] = ( + atoms.get_potential_energy() - e_model_zero_conf + ) + except Exception as exc: + warn( + f"Error calculating relative energy for {label}: {exc}", + stacklevel=2, + ) + atoms.info["model_rel_energy"] = np.nan write_dir = OUT_PATH / model_name write_dir.mkdir(parents=True, exist_ok=True) write(write_dir / f"{label}.xyz", atoms) diff --git a/ml_peg/calcs/conformers/UpU46/calc_UpU46.py b/ml_peg/calcs/conformers/UpU46/calc_UpU46.py index f7d8556cf..0096ec6fe 100644 --- a/ml_peg/calcs/conformers/UpU46/calc_UpU46.py +++ b/ml_peg/calcs/conformers/UpU46/calc_UpU46.py @@ -10,9 +10,11 @@ from pathlib import Path from typing import Any +from warnings import warn from ase import Atoms, units from ase.io import read, write +import numpy as np import pytest from tqdm import tqdm @@ -104,7 +106,11 @@ def test_upu46(mlip: tuple[str, Any]) -> None: conf_lowest = get_atoms(data_path / f"{zero_conf_label}.xyz") conf_lowest.calc = calc - e_conf_lowest_model = conf_lowest.get_potential_energy() + try: + e_conf_lowest_model = conf_lowest.get_potential_energy() + except Exception as exc: + warn(f"Error calculating energy for {zero_conf_label}: {exc}", stacklevel=2) + e_conf_lowest_model = np.nan for label, e_ref in tqdm(ref_energies.items()): # Skip the reference conformer for @@ -114,9 +120,13 @@ def test_upu46(mlip: tuple[str, Any]) -> None: atoms = get_atoms(data_path / f"{label}.xyz") atoms.calc = calc - atoms.info["model_rel_energy"] = ( - atoms.get_potential_energy() - e_conf_lowest_model - ) + try: + atoms.info["model_rel_energy"] = ( + atoms.get_potential_energy() - e_conf_lowest_model + ) + except Exception as exc: + warn(f"Error calculating relative energy for {label}: {exc}", stacklevel=2) + atoms.info["model_rel_energy"] = np.nan atoms.info["ref_energy"] = e_ref atoms.calc = None diff --git a/ml_peg/calcs/conformers/solvMPCONF196/calc_solvMPCONF196.py b/ml_peg/calcs/conformers/solvMPCONF196/calc_solvMPCONF196.py index c6cc078dd..be51974af 100644 --- a/ml_peg/calcs/conformers/solvMPCONF196/calc_solvMPCONF196.py +++ b/ml_peg/calcs/conformers/solvMPCONF196/calc_solvMPCONF196.py @@ -9,6 +9,7 @@ from pathlib import Path from typing import Any +from warnings import warn from ase import Atoms, units from ase.io import read, write @@ -139,7 +140,11 @@ def test_solvmpconf196(mlip: tuple[str, Any]) -> None: ) atoms.translate(-atoms.get_center_of_mass()) atoms.calc = calc - model_abs_energies.append(atoms.get_potential_energy()) + try: + model_abs_energies.append(atoms.get_potential_energy()) + except Exception as exc: + warn(f"Error calculating energy for {xyz_label}: {exc}", stacklevel=2) + model_abs_energies.append(np.nan) ref_abs_energies.append(e_ref) current_molecule_labels.append(label) diff --git a/ml_peg/calcs/conftest.py b/ml_peg/calcs/conftest.py new file mode 100644 index 000000000..7b2855061 --- /dev/null +++ b/ml_peg/calcs/conftest.py @@ -0,0 +1,44 @@ +"""Configure pytest for calculations.""" + +from __future__ import annotations + +from pytest import Config, Parser + +from ml_peg import models + + +def pytest_addoption(parser: Parser) -> None: + """ + Add custom CLI inputs to pytest. + + Parameters + ---------- + parser + Pytest parser object. + """ + parser.addoption( + "--run-mock", + action="store_true", + default=False, + help="Include mock model in tests", + ) + parser.addoption( + "--mock-only", + action="store_true", + default=False, + help="Only run mock model, ignoring other models", + ) + + +def pytest_configure(config: Config) -> None: + """ + Configure pytest to custom CLI inputs. + + Parameters + ---------- + config + Pytest configuration object. + """ + # Set current models from CLI input + models.run_mock = config.getoption("--run-mock") + models.mock_only = config.getoption("--mock-only") diff --git a/ml_peg/calcs/defect/Defectstab/calc_Defectstab.py b/ml_peg/calcs/defect/Defectstab/calc_Defectstab.py index 7f58a9cee..21477f736 100644 --- a/ml_peg/calcs/defect/Defectstab/calc_Defectstab.py +++ b/ml_peg/calcs/defect/Defectstab/calc_Defectstab.py @@ -4,8 +4,10 @@ from pathlib import Path from typing import Any +from warnings import warn from ase.io import read, write +import numpy as np import pytest from ml_peg.calcs.utils.utils import download_s3_data @@ -105,7 +107,14 @@ def test_defectstab(mlip: tuple[str, Any]) -> None: # Assuming existing logic for fe_sia where ref.poscar header has E_bulk # And we need to calculate E_bulk_pred atoms_ref.calc = calc - e_ref_bulk_pred = atoms_ref.get_potential_energy() + try: + e_ref_bulk_pred = atoms_ref.get_potential_energy() + except Exception as exc: + warn( + f"Error calculating bulk reference energy for {ref_file}: {exc}", + stacklevel=2, + ) + e_ref_bulk_pred = np.nan n_bulk = len(atoms_ref) # Save ref calculation @@ -126,7 +135,11 @@ def test_defectstab(mlip: tuple[str, Any]) -> None: e_dft = get_ref_energy(poscar_path) atoms.calc = calc - e_pred = atoms.get_potential_energy() + try: + e_pred = atoms.get_potential_energy() + except Exception as exc: + warn(f"Error calculating energy for {poscar_path}: {exc}", stacklevel=2) + e_pred = np.nan atoms.info["energy_dft"] = e_dft atoms.info["energy_pred"] = e_pred @@ -175,7 +188,11 @@ def test_defectstab(mlip: tuple[str, Any]) -> None: atoms.info.setdefault("spin", 1) e_dft = get_ref_energy(p_path) atoms.calc = calc - e_pred = atoms.get_potential_energy() # Just to verify/store if needed + try: + e_pred = atoms.get_potential_energy() # Just to verify/store if needed + except Exception as exc: + warn(f"Error calculating energy for {p_path}: {exc}", stacklevel=2) + e_pred = np.nan n_atoms = len(atoms) ref_energies[f"E_{el}_dft"] = e_dft / n_atoms @@ -200,7 +217,11 @@ def test_defectstab(mlip: tuple[str, Any]) -> None: atoms_nd.info.setdefault("charge", 0) atoms_nd.info.setdefault("spin", 1) atoms_nd.calc = calc - e_no_defect_pred = atoms_nd.get_potential_energy() + try: + e_no_defect_pred = atoms_nd.get_potential_energy() + except Exception as exc: + warn(f"Error calculating energy for {f}: {exc}", stacklevel=2) + e_no_defect_pred = np.nan break # Process Boron Carbide Subsets @@ -217,7 +238,11 @@ def test_defectstab(mlip: tuple[str, Any]) -> None: e_dft = get_ref_energy(poscar_path) atoms.calc = calc - e_pred = atoms.get_potential_energy() + try: + e_pred = atoms.get_potential_energy() + except Exception as exc: + warn(f"Error calculating energy for {poscar_path}: {exc}", stacklevel=2) + e_pred = np.nan atoms.info["energy_dft"] = e_dft atoms.info["energy_pred"] = e_pred @@ -276,7 +301,11 @@ def test_defectstab(mlip: tuple[str, Any]) -> None: atoms_mai.info.setdefault("charge", 0) atoms_mai.info.setdefault("spin", 1) atoms_mai.calc = calc - e_mai_pred = atoms_mai.get_potential_energy() + try: + e_mai_pred = atoms_mai.get_potential_energy() + except Exception as exc: + warn(f"Error calculating energy for {p}: {exc}", stacklevel=2) + e_mai_pred = np.nan elif "pristine" in p.name: e_pris_dft = get_ref_energy(p) atoms_pris = read(p) @@ -284,7 +313,11 @@ def test_defectstab(mlip: tuple[str, Any]) -> None: atoms_pris.info.setdefault("charge", 0) atoms_pris.info.setdefault("spin", 1) atoms_pris.calc = calc - e_pris_pred = atoms_pris.get_potential_energy() + try: + e_pris_pred = atoms_pris.get_potential_energy() + except Exception as exc: + warn(f"Error calculating energy for {p}: {exc}", stacklevel=2) + e_pris_pred = np.nan for poscar_path in mapi_files: atoms = read(poscar_path) @@ -294,7 +327,11 @@ def test_defectstab(mlip: tuple[str, Any]) -> None: e_dft = get_ref_energy(poscar_path) atoms.calc = calc - e_pred = atoms.get_potential_energy() + try: + e_pred = atoms.get_potential_energy() + except Exception as exc: + warn(f"Error calculating energy for {poscar_path}: {exc}", stacklevel=2) + e_pred = np.nan atoms.info["energy_dft"] = e_dft atoms.info["energy_pred"] = e_pred diff --git a/ml_peg/calcs/defect/Relastab/calc_Relastab.py b/ml_peg/calcs/defect/Relastab/calc_Relastab.py index 1d6df0a48..ba7fbcfaf 100644 --- a/ml_peg/calcs/defect/Relastab/calc_Relastab.py +++ b/ml_peg/calcs/defect/Relastab/calc_Relastab.py @@ -4,8 +4,10 @@ from pathlib import Path from typing import Any +from warnings import warn from ase.io import read, write +import numpy as np import pytest from ml_peg.calcs.utils.utils import download_s3_data @@ -71,8 +73,12 @@ def test_relastab(mlip: tuple[str, Any]) -> None: # Calculate atoms.calc = calc - - atoms.get_potential_energy() + try: + e_pred = atoms.get_potential_energy() + except Exception as exc: + warn(f"Error calculating energy for {poscar}: {exc}", stacklevel=2) + e_pred = np.nan + atoms.info["energy_pred"] = e_pred atoms.info["system"] = poscar.stem atoms.info["subset"] = subset_name diff --git a/ml_peg/calcs/lanthanides/isomer_complexes/calc_isomer_complexes.py b/ml_peg/calcs/lanthanides/isomer_complexes/calc_isomer_complexes.py index f662009e0..afc08b958 100644 --- a/ml_peg/calcs/lanthanides/isomer_complexes/calc_isomer_complexes.py +++ b/ml_peg/calcs/lanthanides/isomer_complexes/calc_isomer_complexes.py @@ -5,9 +5,11 @@ from copy import copy from pathlib import Path from typing import Any +from warnings import warn from ase import units from ase.io import read, write +import numpy as np import pytest from tqdm import tqdm @@ -125,7 +127,11 @@ def test_isomer_complexes(mlip: tuple[str, Any]) -> None: atoms.info["spin_multiplicity"] = entry["multiplicity"] atoms.info["spin"] = entry["multiplicity"] atoms.calc = copy(calc) - atoms.info["model_energy"] = atoms.get_potential_energy() + try: + atoms.info["model_energy"] = atoms.get_potential_energy() + except Exception as exc: + warn(f"Error calculating energy for {entry['xyz']}: {exc}", stacklevel=2) + atoms.info["model_energy"] = np.nan atoms.info["model"] = model_name atoms.info["system"] = entry["system"] diff --git a/ml_peg/calcs/molecular/BMIMCl_RDF/calc_BMIMCl_RDF.py b/ml_peg/calcs/molecular/BMIMCl_RDF/calc_BMIMCl_RDF.py index 850f401d9..b05a95f25 100644 --- a/ml_peg/calcs/molecular/BMIMCl_RDF/calc_BMIMCl_RDF.py +++ b/ml_peg/calcs/molecular/BMIMCl_RDF/calc_BMIMCl_RDF.py @@ -10,6 +10,7 @@ from pathlib import Path from typing import Any +from warnings import warn from ase import units from ase.io import write @@ -17,6 +18,7 @@ from ase.md.velocitydistribution import MaxwellBoltzmannDistribution from ase.optimize import LBFGS import molify +import numpy as np import pytest from tqdm import tqdm @@ -59,7 +61,10 @@ def test_bmimcl_md(mlip: tuple[str, Any]) -> None: box.calc = calc opt = LBFGS(box) - opt.run(fmax=0.1, steps=1000) + try: + opt.run(fmax=0.1, steps=1000) + except Exception as exc: + warn(f"Error during optimisation: {exc}", stacklevel=2) MaxwellBoltzmannDistribution(box, temperature_K=TEMPERATURE) @@ -78,5 +83,10 @@ def test_bmimcl_md(mlip: tuple[str, Any]) -> None: traj_file.unlink() for _ in tqdm(range(STEPS), desc=f"{model_name} MD"): - dyn.run(1) + try: + dyn.run(1) + except Exception as exc: + warn(f"Error during MD step for {model_name}: {exc}", stacklevel=2) + box.info["energy"] = np.nan + break write(traj_file, box, append=True) diff --git a/ml_peg/calcs/molecular/GMTKN55/calc_GMTKN55.py b/ml_peg/calcs/molecular/GMTKN55/calc_GMTKN55.py index 42de08e75..e75a093ad 100644 --- a/ml_peg/calcs/molecular/GMTKN55/calc_GMTKN55.py +++ b/ml_peg/calcs/molecular/GMTKN55/calc_GMTKN55.py @@ -5,6 +5,7 @@ from copy import copy from pathlib import Path from typing import Any +from warnings import warn from ase import Atoms from ase.io import write @@ -97,7 +98,17 @@ def test_gmtkn55(mlip: tuple[str, Any]) -> None: atoms.pbc = False atoms.calc = copy(calc) - atoms.get_potential_energy() + try: + energy = atoms.get_potential_energy() + except Exception as exc: + warn( + f"Error calculating energy for {species_name}: {exc}", + stacklevel=2, + ) + energy = np.nan + + atoms.info["energy"] = energy + atoms.calc = None system_structs.append(atoms) diff --git a/ml_peg/calcs/molecular/Wiggle150/calc_Wiggle150.py b/ml_peg/calcs/molecular/Wiggle150/calc_Wiggle150.py index 764005ab1..b9297a6ac 100644 --- a/ml_peg/calcs/molecular/Wiggle150/calc_Wiggle150.py +++ b/ml_peg/calcs/molecular/Wiggle150/calc_Wiggle150.py @@ -5,10 +5,12 @@ from collections.abc import Iterable from pathlib import Path from typing import Any +from warnings import warn from ase import Atoms, units from ase.calculators.calculator import Calculator from ase.io import read, write +import numpy as np import pytest from tqdm import tqdm @@ -111,7 +113,11 @@ def get_energy(atoms: Atoms, calc: Calculator) -> float: """ atoms_copy = atoms.copy() atoms_copy.calc = calc - energy = atoms_copy.get_potential_energy() + try: + energy = atoms_copy.get_potential_energy() + except Exception as exc: + warn(f"Error calculating energy: {exc}", stacklevel=2) + return np.nan return float(energy) diff --git a/ml_peg/calcs/molecular_crystal/CPOSS209/calc_CPOSS209.py b/ml_peg/calcs/molecular_crystal/CPOSS209/calc_CPOSS209.py index e1a3a5d06..df0df28aa 100644 --- a/ml_peg/calcs/molecular_crystal/CPOSS209/calc_CPOSS209.py +++ b/ml_peg/calcs/molecular_crystal/CPOSS209/calc_CPOSS209.py @@ -5,6 +5,7 @@ from copy import copy from pathlib import Path from typing import Any +from warnings import warn from ase.io import read, write import numpy as np @@ -74,7 +75,13 @@ def test_lattice_energy(mlip: tuple[str, Any]) -> None: solid.info.setdefault("charge", 0) solid.info.setdefault("spin", 1) solid.calc = copy(calc) - solid.get_potential_energy() + try: + solid.get_potential_energy() + except Exception as exc: + warn( + f"Error calculating energy for {crystal_file}: {exc}", stacklevel=2 + ) + solid.info["energy"] = np.nan solid.info["ref"] = ref_crystal solid.info["num_molecules"] = num_mol @@ -132,7 +139,13 @@ def test_lattice_energy(mlip: tuple[str, Any]) -> None: molecule.info.setdefault("charge", 0) molecule.info.setdefault("spin", 1) molecule.calc = copy(calc) - molecule.get_potential_energy() + try: + molecule.get_potential_energy() + except Exception as exc: + warn( + f"Error calculating energy for {molecule_file}: {exc}", stacklevel=2 + ) + molecule.info["energy"] = np.nan # One molecule in each gas phase file molecule.info["num_molecules"] = 1 diff --git a/ml_peg/calcs/molecular_crystal/DMC_ICE13/calc_DMC_ICE13.py b/ml_peg/calcs/molecular_crystal/DMC_ICE13/calc_DMC_ICE13.py index 4d6ae4552..0c83e9b91 100644 --- a/ml_peg/calcs/molecular_crystal/DMC_ICE13/calc_DMC_ICE13.py +++ b/ml_peg/calcs/molecular_crystal/DMC_ICE13/calc_DMC_ICE13.py @@ -6,8 +6,10 @@ import json from pathlib import Path from typing import Any +from warnings import warn from ase.io import read, write +import numpy as np import pytest from ml_peg.calcs.utils.utils import download_s3_data @@ -58,7 +60,11 @@ def test_lattice_energy(mlip: tuple[str, Any]) -> None: # Set default charge and spin water.info.setdefault("charge", 0) water.info.setdefault("spin", 1) - water.get_potential_energy() + try: + water.get_potential_energy() + except Exception as exc: + warn(f"Error calculating energy: {exc}", stacklevel=2) + water.info["energy"] = np.nan for polymorph in polymorphs: polymorph_path = data_dir / polymorph / "POSCAR" @@ -69,7 +75,11 @@ def test_lattice_energy(mlip: tuple[str, Any]) -> None: # Set default charge and spin struct.info.setdefault("charge", 0) struct.info.setdefault("spin", 1) - struct.get_potential_energy() + try: + struct.get_potential_energy() + except Exception as exc: + warn(f"Error calculating energy for {polymorph}: {exc}", stacklevel=2) + struct.info["energy"] = np.nan struct.info["ref"] = ref struct.info["polymorph"] = polymorph diff --git a/ml_peg/calcs/molecular_crystal/X23/calc_X23.py b/ml_peg/calcs/molecular_crystal/X23/calc_X23.py index df5aa92b5..d9680e881 100644 --- a/ml_peg/calcs/molecular_crystal/X23/calc_X23.py +++ b/ml_peg/calcs/molecular_crystal/X23/calc_X23.py @@ -5,6 +5,7 @@ from copy import copy from pathlib import Path from typing import Any +from warnings import warn from ase import units from ase.io import read, write @@ -63,14 +64,25 @@ def test_lattice_energy(mlip: tuple[str, Any]) -> None: # Set default charge and spin molecule.info.setdefault("charge", 0) molecule.info.setdefault("spin", 1) - molecule.get_potential_energy() + try: + molecule.get_potential_energy() + except Exception as exc: + warn( + f"Error calculating energy for {system}: {exc}", + stacklevel=2, + ) + molecule.info["energy"] = np.nan solid = read(solid_path, index=0, format="vasp") solid.calc = copy(calc) # Set default charge and spin solid.info.setdefault("charge", 0) solid.info.setdefault("spin", 1) - solid.get_potential_energy() + try: + solid.get_potential_energy() + except Exception as exc: + warn(f"Error calculating energy for {system}: {exc}", stacklevel=2) + solid.info["energy"] = np.nan ref = np.loadtxt(ref_path)[0] num_molecules = np.loadtxt(num_molecules_path) diff --git a/ml_peg/calcs/molecular_dynamics/liquid_densities/calc_liquid_densities.py b/ml_peg/calcs/molecular_dynamics/liquid_densities/calc_liquid_densities.py index e548e46aa..7d6215ce0 100644 --- a/ml_peg/calcs/molecular_dynamics/liquid_densities/calc_liquid_densities.py +++ b/ml_peg/calcs/molecular_dynamics/liquid_densities/calc_liquid_densities.py @@ -11,10 +11,12 @@ from pathlib import Path import time from typing import Any +from warnings import warn from ase import Atoms, units from ase.io import Trajectory, read from ase.md.nose_hoover_chain import IsotropicMTKNPT +import numpy as np import pytest from ml_peg.calcs.utils.utils import download_s3_data @@ -121,7 +123,11 @@ def run_npt(atoms, calc, output_fname): ) dyn.nsteps = nsteps dyn.attach(log_md, interval=LOG_INTERVAL, dyn=dyn, start_time=time.time()) - dyn.run(steps=NUM_MD_STEPS - nsteps) + try: + dyn.run(steps=NUM_MD_STEPS - nsteps) + except Exception as exc: + warn(f"Error running MD: {exc}", stacklevel=2) + dyn.atoms.info["energy"] = np.nan @pytest.mark.very_slow diff --git a/ml_peg/calcs/molecular_dynamics/water_density/calc_water_density.py b/ml_peg/calcs/molecular_dynamics/water_density/calc_water_density.py index 6f92b5031..daed81224 100644 --- a/ml_peg/calcs/molecular_dynamics/water_density/calc_water_density.py +++ b/ml_peg/calcs/molecular_dynamics/water_density/calc_water_density.py @@ -7,10 +7,12 @@ from pathlib import Path import time from typing import Any +from warnings import warn from ase import Atoms, units from ase.io import Trajectory, read from ase.md.nose_hoover_chain import IsotropicMTKNPT +import numpy as np import pytest from ml_peg.calcs.utils.utils import download_s3_data @@ -117,7 +119,11 @@ def run_npt(atoms, calc, output_fname, temperature): ) dyn.nsteps = nsteps dyn.attach(log_md, interval=LOG_INTERVAL, dyn=dyn, start_time=time.time()) - dyn.run(steps=NUM_MD_STEPS - nsteps) + try: + dyn.run(steps=NUM_MD_STEPS - nsteps) + except Exception as exc: + warn(f"Error running MD: {exc}", stacklevel=2) + dyn.atoms.info["energy"] = np.nan @pytest.mark.very_slow diff --git a/ml_peg/calcs/molecular_reactions/BH2O_36/calc_BH2O_36.py b/ml_peg/calcs/molecular_reactions/BH2O_36/calc_BH2O_36.py index 7f02fbd64..acf107a69 100644 --- a/ml_peg/calcs/molecular_reactions/BH2O_36/calc_BH2O_36.py +++ b/ml_peg/calcs/molecular_reactions/BH2O_36/calc_BH2O_36.py @@ -10,8 +10,10 @@ import json from pathlib import Path from typing import Any +from warnings import warn from ase.io import read, write +import numpy as np import pytest from tqdm import tqdm @@ -106,21 +108,42 @@ def test_bh2o_36(mlip: tuple[str, Any]) -> None: atoms_rct.info["spin"] = 1 atoms_rct.info["ref_energy"] = system["rct"]["energy"] atoms_rct.calc = calc - atoms_rct.info["pred_energy"] = atoms_rct.get_potential_energy() + try: + atoms_rct.info["pred_energy"] = atoms_rct.get_potential_energy() + except Exception as exc: + warn( + f"Error calculating energy for {identifier}: {exc}", + stacklevel=2, + ) + atoms_rct.info["pred_energy"] = np.nan atoms_pro = read(system["pro"]["xyz_path"]) atoms_pro.info["charge"] = int(system["pro"]["charge"]) atoms_pro.info["spin"] = 1 atoms_pro.info["ref_energy"] = system["pro"]["energy"] atoms_pro.calc = calc - atoms_pro.info["pred_energy"] = atoms_pro.get_potential_energy() + try: + atoms_pro.info["pred_energy"] = atoms_pro.get_potential_energy() + except Exception as exc: + warn( + f"Error calculating energy for {identifier}: {exc}", + stacklevel=2, + ) + atoms_pro.info["pred_energy"] = np.nan atoms_ts = read(system["ts"]["xyz_path"]) atoms_ts.info["charge"] = int(system["ts"]["charge"]) atoms_ts.info["spin"] = 1 atoms_ts.info["ref_energy"] = system["ts"]["energy"] atoms_ts.calc = calc - atoms_ts.info["pred_energy"] = atoms_ts.get_potential_energy() + try: + atoms_ts.info["pred_energy"] = atoms_ts.get_potential_energy() + except Exception as exc: + warn( + f"Error calculating energy for {identifier}: {exc}", + stacklevel=2, + ) + atoms_ts.info["pred_energy"] = np.nan write_dir = OUT_PATH / model_name write_dir.mkdir(parents=True, exist_ok=True) diff --git a/ml_peg/calcs/molecular_reactions/BH9/calc_BH9.py b/ml_peg/calcs/molecular_reactions/BH9/calc_BH9.py index d49a7e300..5a68e23b6 100644 --- a/ml_peg/calcs/molecular_reactions/BH9/calc_BH9.py +++ b/ml_peg/calcs/molecular_reactions/BH9/calc_BH9.py @@ -10,9 +10,11 @@ from copy import copy from pathlib import Path from typing import Any +from warnings import warn from ase import units from ase.io import read, write +import numpy as np import pytest from tqdm import tqdm @@ -163,7 +165,11 @@ def test_bh9(mlip: tuple[str, Any]) -> None: # Create list for TS and reactant atoms structs = [process_atoms(xyz_path / f"{label}TS.xyz")] structs[0].calc = calc - structs[0].info["model_energy"] = structs[0].get_potential_energy() + try: + structs[0].info["model_energy"] = structs[0].get_potential_energy() + except Exception as exc: + warn(f"Error calculating energy for {label}: {exc}", stacklevel=2) + structs[0].info["model_energy"] = np.nan structs[0].info["label"] = label # Write both forward and reverse barriers, only forward used in analysis here @@ -173,7 +179,16 @@ def test_bh9(mlip: tuple[str, Any]) -> None: for file in xyz_path.glob(f"{label}R*.xyz"): reactant_atoms = process_atoms(file) reactant_atoms.calc = copy(calc) - reactant_atoms.info["model_energy"] = reactant_atoms.get_potential_energy() + try: + reactant_atoms.info["model_energy"] = ( + reactant_atoms.get_potential_energy() + ) + except Exception as exc: + warn( + f"Error calculating energy for {file}: {exc}", + stacklevel=2, + ) + reactant_atoms.info["model_energy"] = np.nan reactant_atoms.info["label"] = label structs.append(reactant_atoms) diff --git a/ml_peg/calcs/molecular_reactions/CYCLO70/calc_CYCLO70.py b/ml_peg/calcs/molecular_reactions/CYCLO70/calc_CYCLO70.py index 2423bdac2..fde267d8a 100644 --- a/ml_peg/calcs/molecular_reactions/CYCLO70/calc_CYCLO70.py +++ b/ml_peg/calcs/molecular_reactions/CYCLO70/calc_CYCLO70.py @@ -12,9 +12,11 @@ from pathlib import Path from typing import Any +from warnings import warn from ase import units from ase.io import read, write +import numpy as np import pytest from tqdm import tqdm @@ -94,7 +96,15 @@ def test_cyclo70(mlip: tuple[str, Any]) -> None: atoms.info["charge"] = int(atoms.info["charge"]) else: atoms.info["charge"] = 0 - bh_forward_model -= atoms.get_potential_energy() + try: + energy = atoms.get_potential_energy() + except Exception as exc: + warn( + f"Error calculating energy for {atoms_label} in {rxn}: {exc}", + stacklevel=2, + ) + energy = np.nan + bh_forward_model -= energy atoms.info["label"] = atoms_label atoms.calc = None structs_forward.append(atoms) @@ -110,7 +120,15 @@ def test_cyclo70(mlip: tuple[str, Any]) -> None: atoms.info["charge"] = int(atoms.info["charge"]) else: atoms.info["charge"] = 0 - bh_reverse_model -= atoms.get_potential_energy() + try: + energy = atoms.get_potential_energy() + except Exception as exc: + warn( + f"Error calculating energy for {atoms_label} in {rxn}: {exc}", + stacklevel=2, + ) + energy = np.nan + bh_reverse_model -= energy atoms.info["label"] = atoms_label atoms.calc = None structs_reverse.append(atoms) @@ -126,8 +144,16 @@ def test_cyclo70(mlip: tuple[str, Any]) -> None: atoms.info["charge"] = int(atoms.info["charge"]) else: atoms.info["charge"] = 0 - bh_forward_model += atoms.get_potential_energy() - bh_reverse_model += atoms.get_potential_energy() + try: + energy = atoms.get_potential_energy() + except Exception as exc: + warn( + f"Error calculating energy for {atoms_label} in {rxn}: {exc}", + stacklevel=2, + ) + energy = np.nan + bh_forward_model += energy + bh_reverse_model += energy atoms.info["ref_forward_bh"] = bh_forward_ref atoms.info["ref_reverse_bh"] = bh_reverse_ref diff --git a/ml_peg/calcs/molecular_reactions/Criegee22/calc_Criegee22.py b/ml_peg/calcs/molecular_reactions/Criegee22/calc_Criegee22.py index a7fbc1ca8..93f3b59d4 100644 --- a/ml_peg/calcs/molecular_reactions/Criegee22/calc_Criegee22.py +++ b/ml_peg/calcs/molecular_reactions/Criegee22/calc_Criegee22.py @@ -9,9 +9,11 @@ from pathlib import Path from typing import Any +from warnings import warn from ase import units from ase.io import read, write +import numpy as np import pytest from tqdm import tqdm @@ -61,14 +63,28 @@ def test_criegee22(mlip: tuple[str, Any]) -> None: atoms_reac.calc = calc atoms_reac.info["charge"] = 0 atoms_reac.info["spin"] = 1 - atoms_reac.info["model_energy"] = atoms_reac.get_potential_energy() + try: + atoms_reac.info["model_energy"] = atoms_reac.get_potential_energy() + except Exception as exc: + warn( + f"Error calculating model energy for {label}: {exc}", + stacklevel=2, + ) + atoms_reac.info["model_energy"] = np.nan atoms_reac.info["ref_energy"] = 0 atoms_ts = read(data_path / "structures" / f"{label}-TS.xyz") atoms_ts.calc = calc atoms_ts.info["charge"] = 0 atoms_ts.info["spin"] = 1 - atoms_ts.info["model_energy"] = atoms_ts.get_potential_energy() + try: + atoms_ts.info["model_energy"] = atoms_ts.get_potential_energy() + except Exception as exc: + warn( + f"Error calculating model energy for {label}: {exc}", + stacklevel=2, + ) + atoms_ts.info["model_energy"] = np.nan atoms_ts.info["ref_energy"] = bh_ref write_dir = OUT_PATH / model_name diff --git a/ml_peg/calcs/molecular_reactions/RDB7/calc_RDB7.py b/ml_peg/calcs/molecular_reactions/RDB7/calc_RDB7.py index 06c2c6371..9c2e1f476 100644 --- a/ml_peg/calcs/molecular_reactions/RDB7/calc_RDB7.py +++ b/ml_peg/calcs/molecular_reactions/RDB7/calc_RDB7.py @@ -12,6 +12,7 @@ from pathlib import Path from typing import Any +from warnings import warn from ase import Atom, Atoms, units from ase.io import write @@ -118,12 +119,25 @@ def test_rdb87(mlip: tuple[str, Any]) -> None: bh_forward_ref -= get_cc_energy(qm_path) atoms = get_atoms_from_molpro(qm_path) atoms.calc = calc - bh_forward_model -= atoms.get_potential_energy() + try: + energy = atoms.get_potential_energy() + except Exception as exc: + warn( + f"Error calculating energy for {qm_path}: {exc}", + stacklevel=2, + ) + energy = np.nan + bh_forward_model -= energy for qm_path in (data_path / "qm_logs" / f"rxn{label}").glob("ts*"): bh_forward_ref += get_cc_energy(qm_path) atoms = get_atoms_from_molpro(qm_path) atoms.calc = calc - bh_forward_model += atoms.get_potential_energy() + try: + energy = atoms.get_potential_energy() + except Exception as exc: + warn(f"Error calculating energy for {qm_path}: {exc}", stacklevel=2) + energy = np.nan + bh_forward_model += energy atoms.info["model_forward_barrier"] = bh_forward_model atoms.info["ref_forward_barrier"] = bh_forward_ref diff --git a/ml_peg/calcs/nebs/li_diffusion/calc_li_diffusion.py b/ml_peg/calcs/nebs/li_diffusion/calc_li_diffusion.py index 5fad0ebc3..1d5b396e8 100644 --- a/ml_peg/calcs/nebs/li_diffusion/calc_li_diffusion.py +++ b/ml_peg/calcs/nebs/li_diffusion/calc_li_diffusion.py @@ -3,11 +3,13 @@ from __future__ import annotations from pathlib import Path +from warnings import warn from ase import Atoms -from ase.io import read +from ase.io import read, write from janus_core.calculations.geom_opt import GeomOpt from janus_core.calculations.neb import NEB +import numpy as np import pytest from ml_peg.models import current_models @@ -43,11 +45,19 @@ def relaxed_structs() -> dict[str, Atoms]: geomopt = GeomOpt( struct=struct, write_results=True, - file_prefix=OUT_PATH / f"{struct_name}-{model_name}", + file_prefix=OUT_PATH / model_name / struct_name, filter_class=None, + steps=25, ) - geomopt.run() - relaxed_structs[f"{struct_name}-{model_name}"] = geomopt.struct + try: + geomopt.run() + except Exception as exc: + warn( + f"Error for {struct_name} with {model_name}: {exc}", + stacklevel=2, + ) + struct.info["energy"] = np.nan + relaxed_structs[f"{model_name}-{struct_name}"] = geomopt.struct return relaxed_structs @@ -64,23 +74,37 @@ def test_li_diffusion_b(relaxed_structs: dict[str, Atoms], model_name: str) -> N model_name Name of model to use. """ + init_struct = relaxed_structs[f"{model_name}-LiFePO4_start_bc.cif"] + final_struct = relaxed_structs[f"{model_name}-LiFePO4_end_b.cif"] neb = NEB( - init_struct=relaxed_structs[f"LiFePO4_start_bc.cif-{model_name}"], - final_struct=relaxed_structs[f"LiFePO4_end_b.cif-{model_name}"], + init_struct=init_struct, + final_struct=final_struct, n_images=11, interpolator="pymatgen", minimize=True, plot_band=True, write_band=True, - file_prefix=OUT_PATH / f"li_diffusion_b-{model_name}", + file_prefix=OUT_PATH / model_name / "li_diffusion_b", + steps=100, ) # Set default charge and spin for all images - neb.interpolate() - neb.interpolator = None - for image in neb.images: - image.info.setdefault("charge", 0) - image.info.setdefault("spin", 1) - neb.run() + try: + neb.interpolate() + neb.interpolator = None + for image in neb.images: + image.info.setdefault("charge", 0) + image.info.setdefault("spin", 1) + neb.run() + except Exception as e: + print(f"Error running NEB for {model_name} path B: {e}") + # Write out equivalent data + out_dir = OUT_PATH / model_name + write(out_dir / "li_diffusion_b-neb-band.extxyz", [init_struct, final_struct]) + with open( + out_dir / "li_diffusion_b-neb-results.dat", "w", encoding="utf8" + ) as out: + print("#Barrier [eV] | delta E [eV] | Max force [eV/Å] ", file=out) + print(np.nan, np.nan, np.nan, file=out) @pytest.mark.slow @@ -96,20 +120,34 @@ def test_li_diffusion_c(relaxed_structs: dict[str, Atoms], model_name: str) -> N model_name Name of model to use. """ + init_struct = relaxed_structs[f"{model_name}-LiFePO4_start_bc.cif"] + final_struct = relaxed_structs[f"{model_name}-LiFePO4_end_c.cif"] neb = NEB( - init_struct=relaxed_structs[f"LiFePO4_start_bc.cif-{model_name}"], - final_struct=relaxed_structs[f"LiFePO4_end_c.cif-{model_name}"], + init_struct=init_struct, + final_struct=final_struct, n_images=11, interpolator="pymatgen", minimize=True, plot_band=True, write_band=True, - file_prefix=OUT_PATH / f"li_diffusion_c-{model_name}", + file_prefix=OUT_PATH / model_name / "li_diffusion_c", + steps=500, ) # Set default charge and spin for all images - neb.interpolate() - neb.interpolator = None - for image in neb.images: - image.info.setdefault("charge", 0) - image.info.setdefault("spin", 1) - neb.run() + try: + neb.interpolate() + neb.interpolator = None + for image in neb.images: + image.info.setdefault("charge", 0) + image.info.setdefault("spin", 1) + neb.run() + except Exception as e: + print(f"Error running NEB for {model_name} path C: {e}") + # Write out equivalent data + out_dir = OUT_PATH / model_name + write(out_dir / "li_diffusion_c-neb-band.extxyz", [init_struct, final_struct]) + with open( + out_dir / "li_diffusion_c-neb-results.dat", "w", encoding="utf8" + ) as out: + print("#Barrier [eV] | delta E [eV] | Max force [eV/Å] ", file=out) + print(np.nan, np.nan, np.nan, file=out) diff --git a/ml_peg/calcs/nebs/si_defects/calc_si_defects.py b/ml_peg/calcs/nebs/si_defects/calc_si_defects.py index e77ceff34..4d8ea2f0e 100644 --- a/ml_peg/calcs/nebs/si_defects/calc_si_defects.py +++ b/ml_peg/calcs/nebs/si_defects/calc_si_defects.py @@ -6,6 +6,7 @@ from pathlib import Path import sys from typing import Any +from warnings import warn from ase.atoms import Atoms from ase.io import read, write @@ -151,10 +152,18 @@ def test_si_defects(mlip: tuple[str, Any]) -> None: out_atoms = atoms.copy() out_atoms.info["ref_energy_ev"] = ref_energy_ev out_atoms.arrays["ref_forces"] = ref_forces - out_atoms.info["pred_energy_ev"] = float(atoms_pred.get_potential_energy()) - out_atoms.arrays["pred_forces"] = np.asarray( - atoms_pred.get_forces(), dtype=float - ) + try: + pred_energy = float(atoms_pred.get_potential_energy()) + pred_forces = np.asarray(atoms_pred.get_forces(), dtype=float) + except Exception as exc: + warn( + f"Error calculating energy for {case}: {exc}", + stacklevel=2, + ) + pred_energy = np.nan + pred_forces = np.full((len(atoms_pred), 3), np.nan) + out_atoms.info["pred_energy_ev"] = pred_energy + out_atoms.arrays["pred_forces"] = pred_forces results.append(out_atoms) write(out_dir / "si_defects.extxyz", results) diff --git a/ml_peg/calcs/non_covalent_interactions/IONPI19/calc_IONPI19.py b/ml_peg/calcs/non_covalent_interactions/IONPI19/calc_IONPI19.py index f23c27240..34a25af7e 100644 --- a/ml_peg/calcs/non_covalent_interactions/IONPI19/calc_IONPI19.py +++ b/ml_peg/calcs/non_covalent_interactions/IONPI19/calc_IONPI19.py @@ -9,9 +9,11 @@ import os from pathlib import Path from typing import Any +from warnings import warn from ase import Atoms, units from ase.io import read, write +import numpy as np import pytest from tqdm import tqdm @@ -135,7 +137,11 @@ def test_ionpi19(mlip: tuple[str, Any]) -> None: atoms = get_atoms(data_path / label) atoms.info["ref_int_energy"] = ref_energies[system_id] atoms.calc = calc - atoms.info["model_energy"] = atoms.get_potential_energy() + try: + atoms.info["model_energy"] = atoms.get_potential_energy() + except Exception as exc: + warn(f"Error calculating energy for {label}: {exc}", stacklevel=2) + atoms.info["model_energy"] = np.nan write_dir = OUT_PATH / model_name write_dir.mkdir(parents=True, exist_ok=True) diff --git a/ml_peg/calcs/non_covalent_interactions/NCIA_D1200/calc_NCIA_D1200.py b/ml_peg/calcs/non_covalent_interactions/NCIA_D1200/calc_NCIA_D1200.py index bdf56367f..e3f3b1ae0 100644 --- a/ml_peg/calcs/non_covalent_interactions/NCIA_D1200/calc_NCIA_D1200.py +++ b/ml_peg/calcs/non_covalent_interactions/NCIA_D1200/calc_NCIA_D1200.py @@ -9,6 +9,7 @@ from copy import copy from pathlib import Path from typing import Any +from warnings import warn from ase import Atoms, units from ase.io import read, write @@ -135,11 +136,15 @@ def test_ncia_d1200(mlip: tuple[str, Any]) -> None: atoms_a.calc = copy(calc) atoms_b.calc = copy(calc) - atoms.info["model_int_energy"] = ( - atoms.get_potential_energy() - - atoms_a.get_potential_energy() - - atoms_b.get_potential_energy() - ) + try: + atoms.info["model_int_energy"] = ( + atoms.get_potential_energy() + - atoms_a.get_potential_energy() + - atoms_b.get_potential_energy() + ) + except Exception as exc: + warn(f"Error calculating energy for {label}: {exc}", stacklevel=2) + atoms.info["model_int_energy"] = np.nan atoms.info["ref_int_energy"] = ref_energy atoms.calc = None diff --git a/ml_peg/calcs/non_covalent_interactions/NCIA_D442x10/calc_NCIA_D442x10.py b/ml_peg/calcs/non_covalent_interactions/NCIA_D442x10/calc_NCIA_D442x10.py index 8f1529b48..b46178fae 100644 --- a/ml_peg/calcs/non_covalent_interactions/NCIA_D442x10/calc_NCIA_D442x10.py +++ b/ml_peg/calcs/non_covalent_interactions/NCIA_D442x10/calc_NCIA_D442x10.py @@ -9,6 +9,7 @@ from copy import copy from pathlib import Path from typing import Any +from warnings import warn from ase import Atoms, units from ase.io import read, write @@ -133,11 +134,15 @@ def test_ncia_d442x10(mlip: tuple[str, Any]) -> None: atoms_a.calc = copy(calc) atoms_b.calc = copy(calc) - atoms.info["model_int_energy"] = ( - atoms.get_potential_energy() - - atoms_a.get_potential_energy() - - atoms_b.get_potential_energy() - ) + try: + atoms.info["model_int_energy"] = ( + atoms.get_potential_energy() + - atoms_a.get_potential_energy() + - atoms_b.get_potential_energy() + ) + except Exception as exc: + warn(f"Error calculating energy for {label}: {exc}", stacklevel=2) + atoms.info["model_int_energy"] = np.nan atoms.info["ref_int_energy"] = ref_energy atoms.calc = None diff --git a/ml_peg/calcs/non_covalent_interactions/NCIA_HB300SPXx10/calc_NCIA_HB300SPXx10.py b/ml_peg/calcs/non_covalent_interactions/NCIA_HB300SPXx10/calc_NCIA_HB300SPXx10.py index 7ced2ae54..c99bb4f3b 100644 --- a/ml_peg/calcs/non_covalent_interactions/NCIA_HB300SPXx10/calc_NCIA_HB300SPXx10.py +++ b/ml_peg/calcs/non_covalent_interactions/NCIA_HB300SPXx10/calc_NCIA_HB300SPXx10.py @@ -8,9 +8,11 @@ from pathlib import Path from typing import Any +from warnings import warn from ase import Atoms, units from ase.io import read, write +import numpy as np import pytest from tqdm import tqdm @@ -126,11 +128,15 @@ def test_ncia_hb300spxx10(mlip: tuple[str, Any]) -> None: atoms_a.calc = calc atoms_b.calc = calc - atoms.info["model_int_energy"] = ( - atoms.get_potential_energy() - - atoms_a.get_potential_energy() - - atoms_b.get_potential_energy() - ) + try: + atoms.info["model_int_energy"] = ( + atoms.get_potential_energy() + - atoms_a.get_potential_energy() + - atoms_b.get_potential_energy() + ) + except Exception as exc: + warn(f"Error calculating energy for {label}: {exc}", stacklevel=2) + atoms.info["model_int_energy"] = np.nan atoms.info["ref_int_energy"] = ref_energy atoms.calc = None diff --git a/ml_peg/calcs/non_covalent_interactions/NCIA_HB375x10/calc_NCIA_HB375x10.py b/ml_peg/calcs/non_covalent_interactions/NCIA_HB375x10/calc_NCIA_HB375x10.py index 459b1c967..505f3ce2c 100644 --- a/ml_peg/calcs/non_covalent_interactions/NCIA_HB375x10/calc_NCIA_HB375x10.py +++ b/ml_peg/calcs/non_covalent_interactions/NCIA_HB375x10/calc_NCIA_HB375x10.py @@ -9,9 +9,11 @@ from copy import copy from pathlib import Path from typing import Any +from warnings import warn from ase import Atoms, units from ase.io import read, write +import numpy as np import pytest from tqdm import tqdm @@ -132,11 +134,15 @@ def test_ncia_hb375x10(mlip: tuple[str, Any]) -> None: atoms_a.calc = copy(calc) atoms_b.calc = copy(calc) - atoms.info["model_int_energy"] = ( - atoms.get_potential_energy() - - atoms_a.get_potential_energy() - - atoms_b.get_potential_energy() - ) + try: + atoms.info["model_int_energy"] = ( + atoms.get_potential_energy() + - atoms_a.get_potential_energy() + - atoms_b.get_potential_energy() + ) + except Exception as exc: + warn(f"Error calculating energy for {label}: {exc}", stacklevel=2) + atoms.info["model_int_energy"] = np.nan atoms.info["ref_int_energy"] = ref_energy atoms.calc = None diff --git a/ml_peg/calcs/non_covalent_interactions/NCIA_IHB100x10/calc_NCIA_IHB100x10.py b/ml_peg/calcs/non_covalent_interactions/NCIA_IHB100x10/calc_NCIA_IHB100x10.py index c067fbe1e..eb44b1653 100644 --- a/ml_peg/calcs/non_covalent_interactions/NCIA_IHB100x10/calc_NCIA_IHB100x10.py +++ b/ml_peg/calcs/non_covalent_interactions/NCIA_IHB100x10/calc_NCIA_IHB100x10.py @@ -9,9 +9,11 @@ from copy import copy from pathlib import Path from typing import Any +from warnings import warn from ase import Atoms, units from ase.io import read, write +import numpy as np import pytest from tqdm import tqdm @@ -129,11 +131,15 @@ def test_ncia_ihb100x10(mlip: tuple[str, Any]) -> None: atoms_a.calc = copy(calc) atoms_b.calc = copy(calc) - atoms.info["model_int_energy"] = ( - atoms.get_potential_energy() - - atoms_a.get_potential_energy() - - atoms_b.get_potential_energy() - ) + try: + atoms.info["model_int_energy"] = ( + atoms.get_potential_energy() + - atoms_a.get_potential_energy() + - atoms_b.get_potential_energy() + ) + except Exception as exc: + warn(f"Error calculating energy for {label}: {exc}", stacklevel=2) + atoms.info["model_int_energy"] = np.nan atoms.info["ref_int_energy"] = ref_energy atoms.calc = None diff --git a/ml_peg/calcs/non_covalent_interactions/NCIA_R739x5/calc_NCIA_R739x5.py b/ml_peg/calcs/non_covalent_interactions/NCIA_R739x5/calc_NCIA_R739x5.py index 9e23dd278..372511433 100644 --- a/ml_peg/calcs/non_covalent_interactions/NCIA_R739x5/calc_NCIA_R739x5.py +++ b/ml_peg/calcs/non_covalent_interactions/NCIA_R739x5/calc_NCIA_R739x5.py @@ -8,6 +8,7 @@ from pathlib import Path from typing import Any +from warnings import warn from ase import Atoms, units from ase.io import read, write @@ -133,11 +134,15 @@ def test_ncia_r739x5(mlip: tuple[str, Any]) -> None: atoms_a.calc = calc atoms_b.calc = calc - atoms.info["model_int_energy"] = ( - atoms.get_potential_energy() - - atoms_a.get_potential_energy() - - atoms_b.get_potential_energy() - ) + try: + atoms.info["model_int_energy"] = ( + atoms.get_potential_energy() + - atoms_a.get_potential_energy() + - atoms_b.get_potential_energy() + ) + except Exception as exc: + warn(f"Error calculating energy for {label}: {exc}", stacklevel=2) + atoms.info["model_int_energy"] = np.nan atoms.info["ref_int_energy"] = ref_energy atoms.calc = None diff --git a/ml_peg/calcs/non_covalent_interactions/NCIA_SH250x10/calc_NCIA_SH250x10.py b/ml_peg/calcs/non_covalent_interactions/NCIA_SH250x10/calc_NCIA_SH250x10.py index 590f7ceee..7e997a91e 100644 --- a/ml_peg/calcs/non_covalent_interactions/NCIA_SH250x10/calc_NCIA_SH250x10.py +++ b/ml_peg/calcs/non_covalent_interactions/NCIA_SH250x10/calc_NCIA_SH250x10.py @@ -9,9 +9,11 @@ from copy import copy from pathlib import Path from typing import Any +from warnings import warn from ase import Atoms, units from ase.io import read, write +import numpy as np import pytest from tqdm import tqdm @@ -126,11 +128,15 @@ def test_lattice_energy(mlip: tuple[str, Any]) -> None: atoms_a.calc = copy(calc) atoms_b.calc = copy(calc) - atoms.info["model_int_energy"] = ( - atoms.get_potential_energy() - - atoms_a.get_potential_energy() - - atoms_b.get_potential_energy() - ) + try: + atoms.info["model_int_energy"] = ( + atoms.get_potential_energy() + - atoms_a.get_potential_energy() + - atoms_b.get_potential_energy() + ) + except Exception as exc: + warn(f"Error calculating energy for {label}: {exc}", stacklevel=2) + atoms.info["model_int_energy"] = np.nan atoms.info["ref_int_energy"] = ref_energy atoms.calc = None diff --git a/ml_peg/calcs/non_covalent_interactions/QUID/calc_QUID.py b/ml_peg/calcs/non_covalent_interactions/QUID/calc_QUID.py index 9080cc61e..af98e38ea 100644 --- a/ml_peg/calcs/non_covalent_interactions/QUID/calc_QUID.py +++ b/ml_peg/calcs/non_covalent_interactions/QUID/calc_QUID.py @@ -11,6 +11,7 @@ from pathlib import Path from typing import Any +from warnings import warn from ase import Atoms from ase.io import write @@ -104,6 +105,7 @@ def compute_interaction_energy(dataset, label, calc): # List to store dimer and monomers. atoms_list = [] model_int_energy = 0 + for atoms_name, stoich in zip( ["dimer", "small_monomer", "big_monomer"], [1, -1, -1], strict=False ): @@ -112,9 +114,16 @@ def compute_interaction_energy(dataset, label, calc): atoms = Atoms(numbers=atomic_numbers, positions=positions) atoms.info.update({"charge": 0, "spin": 1}) atoms.calc = calc - model_int_energy += atoms.get_potential_energy() * stoich - atoms.calc = None atoms_list.append(atoms) + try: + model_int_energy += atoms.get_potential_energy() * stoich + except Exception as exc: + warn(f"Error calculating energy for {atoms_name}: {exc}", stacklevel=2) + model_int_energy = np.nan + break + + for atoms in atoms_list: + atoms.calc = None atoms_list[0].info.update( { "ref_int_energy": ref_int_energy, diff --git a/ml_peg/calcs/physicality/diatomics/calc_diatomics.py b/ml_peg/calcs/physicality/diatomics/calc_diatomics.py index 6033b7976..be8f6c9e0 100644 --- a/ml_peg/calcs/physicality/diatomics/calc_diatomics.py +++ b/ml_peg/calcs/physicality/diatomics/calc_diatomics.py @@ -6,6 +6,7 @@ import itertools import json from pathlib import Path +from warnings import warn from ase import Atoms from ase.data import chemical_symbols @@ -176,8 +177,16 @@ def run_diatomics(model_name: str, model) -> None: atoms.info.setdefault("spin", 1) atoms.calc = calc - energy = float(atoms.get_potential_energy()) - forces = atoms.get_forces() + try: + energy = float(atoms.get_potential_energy()) + forces = atoms.get_forces() + except Exception as exc: + warn( + f"Error calculating energy for {pair_label}: {exc}", + stacklevel=2, + ) + energy = np.nan + forces = np.full((len(atoms), 3), np.nan) bond_vector = atoms.positions[1] - atoms.positions[0] force_parallel = _project_force(forces, bond_vector) diff --git a/ml_peg/calcs/physicality/extensivity/calc_extensivity.py b/ml_peg/calcs/physicality/extensivity/calc_extensivity.py index 44792a7e9..f94114a8e 100644 --- a/ml_peg/calcs/physicality/extensivity/calc_extensivity.py +++ b/ml_peg/calcs/physicality/extensivity/calc_extensivity.py @@ -5,10 +5,12 @@ from copy import copy from pathlib import Path from typing import Any +from warnings import warn from ase import Atoms from ase.build import bulk, surface from ase.io import write +import numpy as np import pytest from ml_peg.models import current_models @@ -87,14 +89,15 @@ def test_extensivity(mlip: tuple[str, Any]) -> None: slab2_big = slab2.copy() slab2_big.set_cell(tall_cell, scale_atoms=False) - slab1_big.calc = calc - slab1_big.get_potential_energy() - - slab2_big.calc = copy(calc) - slab2_big.get_potential_energy() - - combined.calc = copy(calc) - combined.get_potential_energy() + for struct in (slab1_big, slab2_big, combined): + struct.calc = copy(calc) + try: + energy = struct.get_potential_energy() + except Exception as exc: + warn(f"Error calculating energy: {exc}", stacklevel=2) + energy = np.nan + struct.info["energy"] = energy + struct.calc = None # Write output structures write_dir = OUT_PATH / model_name diff --git a/ml_peg/calcs/physicality/locality/calc_locality.py b/ml_peg/calcs/physicality/locality/calc_locality.py index 419cf3169..a89159bda 100644 --- a/ml_peg/calcs/physicality/locality/calc_locality.py +++ b/ml_peg/calcs/physicality/locality/calc_locality.py @@ -4,6 +4,7 @@ from copy import copy from pathlib import Path +from warnings import warn from ase import Atoms from ase.io import write @@ -152,12 +153,12 @@ def prepared_solute() -> dict[str, Atoms]: solutes = {} for model_name, calc in MODELS.items(): solute = solute.copy() + solute.calc = calc.get_calculator(precision="high") + solutes[model_name] = solute try: - solute.calc = calc.get_calculator(precision="high") solute.get_forces() - solutes[model_name] = solute - # If a model fails, don't block other model tests - except (ModuleNotFoundError, RuntimeError, TypeError): + except Exception as exc: + warn(f"Error preparing solute: {exc} for {model_name}", stacklevel=2) continue return solutes @@ -182,7 +183,11 @@ def test_ghost_atoms(prepared_solute: dict[str, Atoms], model_name: str) -> None system_ghost = prepare_ghost_system(solute, ghost_num, ghost_dist) system_ghost.calc = copy(solute.calc) - system_ghost.get_forces() + try: + system_ghost.get_forces() + except Exception as exc: + warn(f"Error calculating forces: {exc}", stacklevel=2) + system_ghost.arrays["forces"] = np.full((len(system_ghost), 3), np.nan) # Write output structures write_dir = OUT_PATH / model_name @@ -214,7 +219,11 @@ def test_rand_h(prepared_solute: dict[str, Atoms], model_name: str) -> None: for _ in range(rand_trials): system_rand_h = add_random_h(solute, rand_min_dist, rand_max_dist, rng) system_rand_h.calc = copy(solute.calc) - system_rand_h.get_forces() + try: + system_rand_h.get_forces() + except Exception as exc: + warn(f"Error calculating forces: {exc}", stacklevel=2) + system_rand_h.arrays["forces"] = np.full((len(system_rand_h), 3), np.nan) rand_h_structures.append(system_rand_h) # Write output structures diff --git a/ml_peg/calcs/physicality/oxidation_states/calc_oxidation_states.py b/ml_peg/calcs/physicality/oxidation_states/calc_oxidation_states.py index 9bf9101c9..4a825ab18 100644 --- a/ml_peg/calcs/physicality/oxidation_states/calc_oxidation_states.py +++ b/ml_peg/calcs/physicality/oxidation_states/calc_oxidation_states.py @@ -4,6 +4,7 @@ from pathlib import Path from typing import Any +from warnings import warn from ase.geometry.rdf import get_rdf from ase.io import read @@ -67,7 +68,10 @@ def test_iron_oxidation_state_md(mlip: tuple[str, Any]) -> None: file_prefix=out_dir / f"{salt}_{model_name}", restart_every=1000, ) - npt.run() + try: + npt.run() + except Exception as exc: + warn(f"Error during MD for {salt}: {exc}", stacklevel=2) @pytest.mark.parametrize("model_name", MODELS) diff --git a/ml_peg/calcs/physicality/water_slab_dipoles/calc_water_slab_dipoles.py b/ml_peg/calcs/physicality/water_slab_dipoles/calc_water_slab_dipoles.py index f81e78c59..3f6722aef 100644 --- a/ml_peg/calcs/physicality/water_slab_dipoles/calc_water_slab_dipoles.py +++ b/ml_peg/calcs/physicality/water_slab_dipoles/calc_water_slab_dipoles.py @@ -4,6 +4,7 @@ from pathlib import Path from typing import Any +from warnings import warn from ase import units from ase.io import read @@ -84,6 +85,9 @@ def test_water_dipole(mlip: tuple[str, Any]) -> None: print("Starting MD") - md.run() + try: + md.run() + except Exception as exc: + warn(f"Error during MD: {exc}", stacklevel=2) print("Finished MD") diff --git a/ml_peg/calcs/supramolecular/LNCI16/calc_LNCI16.py b/ml_peg/calcs/supramolecular/LNCI16/calc_LNCI16.py index 2a69deafa..1f39c7671 100644 --- a/ml_peg/calcs/supramolecular/LNCI16/calc_LNCI16.py +++ b/ml_peg/calcs/supramolecular/LNCI16/calc_LNCI16.py @@ -3,12 +3,14 @@ from __future__ import annotations from pathlib import Path +from warnings import warn from ase import Atoms, units from ase.calculators.calculator import Calculator from ase.io import read, write import mlipx from mlipx.abc import NodeWithCalculator +import numpy as np from tqdm import tqdm import zntrack @@ -177,12 +179,16 @@ def interaction_energy(frags: dict[str, Atoms], calc: Calculator) -> float: guest_copy = frags["guest"].copy() complex_copy.calc = calc - e_complex = complex_copy.get_potential_energy() host_copy.calc = calc - e_host = host_copy.get_potential_energy() guest_copy.calc = calc - e_guest = guest_copy.get_potential_energy() - return e_complex - e_host - e_guest + try: + e_complex = complex_copy.get_potential_energy() + e_host = host_copy.get_potential_energy() + e_guest = guest_copy.get_potential_energy() + return e_complex - e_host - e_guest + except Exception as exc: + warn(f"Error calculating energies: {exc}", stacklevel=2) + return np.nan @staticmethod def benchmark_lnci16( diff --git a/ml_peg/calcs/supramolecular/S30L/calc_S30L.py b/ml_peg/calcs/supramolecular/S30L/calc_S30L.py index ba899644d..e51ac88ec 100644 --- a/ml_peg/calcs/supramolecular/S30L/calc_S30L.py +++ b/ml_peg/calcs/supramolecular/S30L/calc_S30L.py @@ -3,13 +3,14 @@ from __future__ import annotations from pathlib import Path -import warnings +from warnings import warn from ase import Atoms, units from ase.calculators.calculator import Calculator from ase.io import read, write import mlipx from mlipx.abc import NodeWithCalculator +import numpy as np from tqdm import tqdm import zntrack @@ -59,9 +60,7 @@ def read_charge(folder: Path) -> int: try: return int(f.read_text().strip()) except ValueError: - warnings.warn( - f"Invalid charge in {f} - assuming neutral.", stacklevel=2 - ) + warn(f"Invalid charge in {f} - assuming neutral.", stacklevel=2) return 0 def read_atoms(self, folder: Path, ident: str) -> Atoms: @@ -134,12 +133,16 @@ def interaction_energy(frags: dict[str, Atoms], calc: Calculator) -> float: Interaction energy in eV. """ frags["complex"].calc = calc - e_complex = frags["complex"].get_potential_energy() frags["host"].calc = calc - e_host = frags["host"].get_potential_energy() frags["guest"].calc = calc - e_guest = frags["guest"].get_potential_energy() - return e_complex - e_host - e_guest + try: + e_complex = frags["complex"].get_potential_energy() + e_host = frags["host"].get_potential_energy() + e_guest = frags["guest"].get_potential_energy() + return e_complex - e_host - e_guest + except Exception as exc: + warn(f"Error calculating energies: {exc}", stacklevel=2) + return np.nan @staticmethod def parse_references(path: Path) -> dict[int, float]: @@ -218,8 +221,8 @@ def benchmark_s30l( complex_atoms_list.append(complex_atoms) - except Exception as e: - print(f"Error processing system {idx}: {e}") + except Exception as exc: + warn(f"Error processing system {idx}: {exc}", stacklevel=2) continue return complex_atoms_list diff --git a/ml_peg/calcs/supramolecular/utils/plf547_pla15_utils.py b/ml_peg/calcs/supramolecular/utils/plf547_pla15_utils.py index afecdc7d8..14a6f0d2a 100644 --- a/ml_peg/calcs/supramolecular/utils/plf547_pla15_utils.py +++ b/ml_peg/calcs/supramolecular/utils/plf547_pla15_utils.py @@ -4,6 +4,7 @@ import logging from pathlib import Path +from warnings import warn from ase import Atoms, units from ase.calculators.calculator import Calculator @@ -122,18 +123,22 @@ def get_interaction_energy(fragments: dict[str, Atoms], calc: Calculator) -> flo Interaction energy in eV. """ fragments["complex"].calc = calc - e_complex = fragments["complex"].get_potential_energy() - fragments["complex"].calc = None - fragments["protein"].calc = calc - e_protein = fragments["protein"].get_potential_energy() - fragments["protein"].calc = None - fragments["ligand"].calc = calc - e_ligand = fragments["ligand"].get_potential_energy() - fragments["ligand"].calc = None - - return e_complex - e_protein - e_ligand + try: + e_complex = fragments["complex"].get_potential_energy() + e_protein = fragments["protein"].get_potential_energy() + e_ligand = fragments["ligand"].get_potential_energy() + result = e_complex - e_protein - e_ligand + except Exception as exc: + warn(f"Error calculating energies: {exc}", stacklevel=2) + result = np.nan + finally: + fragments["complex"].calc = None + fragments["protein"].calc = None + fragments["ligand"].calc = None + + return result def mda_atoms_to_ase(atom_list, charge: float, identifier: str) -> Atoms: diff --git a/ml_peg/calcs/surfaces/OC157/calc_OC157.py b/ml_peg/calcs/surfaces/OC157/calc_OC157.py index 9110c5e42..fb74d6e84 100644 --- a/ml_peg/calcs/surfaces/OC157/calc_OC157.py +++ b/ml_peg/calcs/surfaces/OC157/calc_OC157.py @@ -4,12 +4,14 @@ from copy import copy from pathlib import Path +from warnings import warn from ase import Atoms from ase.calculators.calculator import Calculator from ase.io import read, write import mlipx from mlipx.abc import NodeWithCalculator +import numpy as np from tqdm import tqdm import zntrack @@ -123,7 +125,11 @@ def evaluate_energies(triplet: list[Atoms], calc: Calculator) -> None: atoms.info.setdefault("charge", 0) atoms.info.setdefault("spin", 1) - atoms.get_potential_energy() + try: + atoms.get_potential_energy() + except Exception as exc: + warn(f"Error calculating energy: {exc}", stacklevel=2) + atoms.info["energy"] = np.nan def run(self): """Run OC157 energy calculations.""" diff --git a/ml_peg/calcs/surfaces/S24/calc_S24.py b/ml_peg/calcs/surfaces/S24/calc_S24.py index 5709ad060..0f948d5a1 100644 --- a/ml_peg/calcs/surfaces/S24/calc_S24.py +++ b/ml_peg/calcs/surfaces/S24/calc_S24.py @@ -4,12 +4,14 @@ from copy import copy from pathlib import Path +from warnings import warn from ase import Atoms from ase.calculators.calculator import Calculator from ase.io import read, write import mlipx from mlipx.abc import NodeWithCalculator +import numpy as np from tqdm import tqdm import zntrack @@ -114,7 +116,11 @@ def evaluate_energies(atoms_list: list[Atoms], calc: Calculator) -> None: atoms.info.setdefault("charge", 0) atoms.info.setdefault("spin", 1) - atoms.get_potential_energy() + try: + atoms.get_potential_energy() + except Exception as exc: + warn(f"Error calculating energy: {exc}", stacklevel=2) + atoms.info["energy"] = np.nan def run(self): """Run S24 energy calculations.""" @@ -140,9 +146,21 @@ def run(self): sys_id = f"{(i // 3) + 1:03d}" # Store reference energies - surface.info["ref_energy"] = surface.get_potential_energy() - mol_surface.info["ref_energy"] = mol_surface.get_potential_energy() - molecule.info["ref_energy"] = molecule.get_potential_energy() + try: + surface.info["ref_energy"] = surface.get_potential_energy() + except Exception as exc: + warn(f"Error calculating energy for {sys_id}: {exc}", stacklevel=2) + surface.info["ref_energy"] = np.nan + try: + mol_surface.info["ref_energy"] = mol_surface.get_potential_energy() + except Exception as exc: + warn(f"Error calculating energy for {sys_id}: {exc}", stacklevel=2) + mol_surface.info["ref_energy"] = np.nan + try: + molecule.info["ref_energy"] = molecule.get_potential_energy() + except Exception as exc: + warn(f"Error calculating energy for {sys_id}: {exc}", stacklevel=2) + molecule.info["ref_energy"] = np.nan # Store system information surface_formula = surface.get_chemical_formula() @@ -166,12 +184,16 @@ def run(self): self.evaluate_energies(triplet, calc) # Calculate and store adsorption energies - surface_energy = surface.get_potential_energy() - mol_surf_energy = mol_surface.get_potential_energy() - molecule_energy = molecule.get_potential_energy() - pred_ads_energy = self.compute_adsorption_energy( - surface_energy, mol_surf_energy, molecule_energy, adsorbate_count - ) + try: + surface_energy = surface.get_potential_energy() + mol_surf_energy = mol_surface.get_potential_energy() + molecule_energy = molecule.get_potential_energy() + pred_ads_energy = self.compute_adsorption_energy( + surface_energy, mol_surf_energy, molecule_energy, adsorbate_count + ) + except Exception as exc: + warn(f"Error calculating energy for {sys_id}: {exc}", stacklevel=2) + pred_ads_energy = np.nan ref_surface_energy = surface.info["ref_energy"] ref_mol_surf_energy = mol_surface.info["ref_energy"] diff --git a/ml_peg/calcs/surfaces/SBH17/calc_SBH17.py b/ml_peg/calcs/surfaces/SBH17/calc_SBH17.py index eb9b44612..00d0315fc 100644 --- a/ml_peg/calcs/surfaces/SBH17/calc_SBH17.py +++ b/ml_peg/calcs/surfaces/SBH17/calc_SBH17.py @@ -5,6 +5,7 @@ from copy import copy from pathlib import Path from typing import Any +from warnings import warn from ase.io import read, write import numpy as np @@ -61,14 +62,22 @@ def test_surface_barrier(mlip: tuple[str, Any]) -> None: gp.info.setdefault("spin", 1) gp.info["spin"] = int(round(gp.info["spin"])) gp.calc = calc - gp.get_potential_energy() + try: + gp.get_potential_energy() + except Exception as exc: + warn(f"Error calculating energy for {system}: {exc}", stacklevel=2) + gp.info["energy"] = np.nan ts = read(ts_path, index=0) ts.info.setdefault("charge", 0) ts.info.setdefault("spin", 1) ts.info["spin"] = int(round(ts.info["spin"])) ts.calc = copy(calc) - ts.get_potential_energy() + try: + ts.get_potential_energy() + except Exception as exc: + warn(f"Error calculating energy for {system}: {exc}", stacklevel=2) + ts.info["energy"] = np.nan ref = np.loadtxt(ref_path) diff --git a/ml_peg/calcs/surfaces/elemental_slab_oxygen_adsorption/calc_elemental_slab_oxygen_adsorption.py b/ml_peg/calcs/surfaces/elemental_slab_oxygen_adsorption/calc_elemental_slab_oxygen_adsorption.py index c6a379220..a4e610ea2 100644 --- a/ml_peg/calcs/surfaces/elemental_slab_oxygen_adsorption/calc_elemental_slab_oxygen_adsorption.py +++ b/ml_peg/calcs/surfaces/elemental_slab_oxygen_adsorption/calc_elemental_slab_oxygen_adsorption.py @@ -4,12 +4,14 @@ from copy import copy from pathlib import Path +from warnings import warn from ase import Atoms from ase.calculators.calculator import Calculator from ase.io import read, write import mlipx from mlipx.abc import NodeWithCalculator +import numpy as np from tqdm import tqdm import zntrack @@ -105,7 +107,11 @@ def evaluate_energies(atoms_list: list[Atoms], calc: Calculator) -> None: atoms.info.setdefault("charge", 0) atoms.info.setdefault("spin", 1) - atoms.get_potential_energy() + try: + atoms.get_potential_energy() + except Exception as exc: + warn(f"Error calculating energy: {exc}", stacklevel=2) + atoms.info["energy"] = np.nan def run(self): """Run oxygen adsorption energy calculations.""" diff --git a/ml_peg/calcs/surfaces/graphene_wetting_under_strain/calc_graphene_wetting_under_strain.py b/ml_peg/calcs/surfaces/graphene_wetting_under_strain/calc_graphene_wetting_under_strain.py index 88eb8439d..84bcb0c55 100644 --- a/ml_peg/calcs/surfaces/graphene_wetting_under_strain/calc_graphene_wetting_under_strain.py +++ b/ml_peg/calcs/surfaces/graphene_wetting_under_strain/calc_graphene_wetting_under_strain.py @@ -4,8 +4,10 @@ from pathlib import Path from typing import Any +from warnings import warn import ase.io +import numpy as np import pytest from tqdm import tqdm import yaml @@ -66,7 +68,11 @@ def test_graphene_wetting_energy(mlip: tuple[str, Any]) -> None: atoms.calc = calc atoms.info.setdefault("charge", 0) atoms.info.setdefault("spin", 1) - water_energy = atoms.get_potential_energy() + try: + water_energy = atoms.get_potential_energy() + except Exception as exc: + warn(f"Error calculating energy: {exc}", stacklevel=2) + water_energy = np.nan # Iterate through strain conditions for strain in strains: @@ -74,7 +80,14 @@ def test_graphene_wetting_energy(mlip: tuple[str, Any]) -> None: atoms.calc = calc atoms.info.setdefault("charge", 0) atoms.info.setdefault("spin", 1) - graphene_energy = atoms.get_potential_energy() + try: + graphene_energy = atoms.get_potential_energy() + except Exception as exc: + warn( + f"Error calculating energy for strain {strain}: {exc}", + stacklevel=2, + ) + graphene_energy = np.nan # Iterate through orientations for orientation in orientations: @@ -89,7 +102,11 @@ def test_graphene_wetting_energy(mlip: tuple[str, Any]) -> None: atoms.calc = calc atoms.info.setdefault("charge", 0) atoms.info.setdefault("spin", 1) - mlip_potential_energy = atoms.get_potential_energy() + try: + mlip_potential_energy = atoms.get_potential_energy() + except Exception as exc: + warn(f"Error calculating energy: {exc}", stacklevel=2) + mlip_potential_energy = np.nan mlip_adsorption_energy = ( mlip_potential_energy - graphene_energy - water_energy ) diff --git a/ml_peg/calcs/tm_complexes/3dTMV/calc_3dTMV.py b/ml_peg/calcs/tm_complexes/3dTMV/calc_3dTMV.py index 3182fe8af..37fb178a2 100644 --- a/ml_peg/calcs/tm_complexes/3dTMV/calc_3dTMV.py +++ b/ml_peg/calcs/tm_complexes/3dTMV/calc_3dTMV.py @@ -9,9 +9,11 @@ from copy import copy from pathlib import Path from typing import Any +from warnings import warn from ase import units from ase.io import read, write +import numpy as np import pytest from tqdm import tqdm @@ -140,13 +142,21 @@ def test_3dtmv(mlip: tuple[str, Any]) -> None: atoms_ox.info["charge"] = MOLECULAR_DATA[complex_id]["charge_ox"] atoms_ox.info["spin"] = MOLECULAR_DATA[complex_id]["mult_ox"] atoms_ox.calc = copy(calc) - oxidized_energy = atoms_ox.get_potential_energy() + try: + oxidized_energy = atoms_ox.get_potential_energy() + except Exception as exc: + warn(f"Error calculating energy for {complex_id}: {exc}", stacklevel=2) + oxidized_energy = np.nan atoms_in = atoms.copy() atoms_in.info["charge"] = MOLECULAR_DATA[complex_id]["charge_in"] atoms_in.info["spin"] = MOLECULAR_DATA[complex_id]["mult_in"] atoms_in.calc = copy(calc) - initial_energy = atoms_in.get_potential_energy() + try: + initial_energy = atoms_in.get_potential_energy() + except Exception as exc: + warn(f"Error calculating energy for {complex_id}: {exc}", stacklevel=2) + initial_energy = np.nan model_ion_energy = oxidized_energy - initial_energy diff --git a/ml_peg/calcs/utils/gscdb138.py b/ml_peg/calcs/utils/gscdb138.py index 4d0a019ca..0fc26c1e0 100644 --- a/ml_peg/calcs/utils/gscdb138.py +++ b/ml_peg/calcs/utils/gscdb138.py @@ -4,6 +4,7 @@ from pathlib import Path from typing import Any +from warnings import warn from ase import Atoms, units from ase.io import read, write @@ -122,7 +123,11 @@ def run_gscdb138( print(f"Skipping {specy}") continue atoms.calc = calc - energy = atoms.get_potential_energy() + try: + energy = atoms.get_potential_energy() + except Exception as exc: + warn(f"Error calculating energy for {specy}: {exc}", stacklevel=2) + energy = np.nan e_rel_model += stoi * energy atoms.info["model_energy"] = energy atoms.calc = None diff --git a/ml_peg/cli/cli.py b/ml_peg/cli/cli.py index 352b68b10..e727a37a8 100644 --- a/ml_peg/cli/cli.py +++ b/ml_peg/cli/cli.py @@ -176,6 +176,12 @@ def run_calcs( test: Annotated[ str, Option(help="Test to run calculations for. Default is all tests.") ] = "*", + run_mock: Annotated[ + bool, Option(help="Whether to run with mock calculator in addition to models.") + ] = True, + mock_only: Annotated[ + bool, Option(help="Whether to only run mock calculator with no models.") + ] = False, run_slow: Annotated[ bool, Option(help="Whether to run calculations labelled slow.") ] = True, @@ -204,6 +210,10 @@ def run_calcs( test Test to run calculation for. Default is `*`, corresponding to all tests in the category. + run_mock + Whether to run mock calculations. Default is `True`. + mock_only + Whether to only run mock calculations, with no models. Default is `False`. run_slow Whether to run slow calculations. Default is `True`. run_very_slow @@ -230,6 +240,12 @@ def run_calcs( if run_very_slow: options.extend(["--run-very-slow"]) + if run_mock: + options.extend(["--run-mock"]) + + if mock_only: + options.extend(["--mock-only"]) + if models: options.extend(["--models", models]) diff --git a/ml_peg/models/__init__.py b/ml_peg/models/__init__.py index 9e29c36ff..d630ac690 100644 --- a/ml_peg/models/__init__.py +++ b/ml_peg/models/__init__.py @@ -7,3 +7,5 @@ MODELS_ROOT = Path(__file__).parent current_models = None models_file = MODELS_ROOT / "models.yml" +run_mock = False +mock_only = False diff --git a/ml_peg/models/get_models.py b/ml_peg/models/get_models.py index 5700b902b..223fb57ff 100644 --- a/ml_peg/models/get_models.py +++ b/ml_peg/models/get_models.py @@ -111,6 +111,8 @@ def get_subset( def load_models( models: None | str | Iterable = None, filepath: Path | str | None = None, + run_mock: bool | None = None, + mock_only: bool | None = None, ) -> dict[str, Any]: """ Load models for use in calculations. @@ -123,15 +125,35 @@ def load_models( this will be treated as a comma-separated list. filepath Path to YAML file with models. Default is `models_file`. + run_mock + Whether to include mock model in the loaded models. Default is False. + mock_only + Whether to load only mock model, ignoring `models`. Default is False. Returns ------- dict[str, Any] - Loaded models from models.yml. + Loaded models from models.yml and/or loaded mock model. """ - from ml_peg.models.models import FairChemCalc, GenericASECalc, OrbCalc, PetMadCalc - - loaded_models = {} + from ml_peg.models.models import ( + FairChemCalc, + GenericASECalc, + MockCalc, + OrbCalc, + PetMadCalc, + ) + + if run_mock is None: + from ml_peg.models import run_mock + if mock_only is None: + from ml_peg.models import mock_only + + if mock_only and not run_mock: + raise ValueError("Cannot set `mock_only` without `run_mock`") + + loaded_models = {"mock": MockCalc()} if run_mock else {} + if mock_only: + return loaded_models filepath = filepath if filepath else models_file all_models = _load_models_yaml(filepath) diff --git a/ml_peg/models/mock.py b/ml_peg/models/mock.py new file mode 100644 index 000000000..ddc94c37b --- /dev/null +++ b/ml_peg/models/mock.py @@ -0,0 +1,66 @@ +"""Mock calculator class.""" + +from __future__ import annotations + +from ase import Atoms +from ase.calculators.calculator import Calculator, all_changes +import numpy as np + + +class MockCalculator(Calculator): + """A mock calculator that returns zero values for all properties.""" + + implemented_properties = ["energy", "forces", "stress"] + + def calculate( + self, + atoms: Atoms | None = None, + properties: list[str] | None = None, + system_changes: list[str] = all_changes, + **kwargs, + ) -> None: + """ + Define calculation method for mock calculator. + + Parameters + ---------- + atoms + Atoms object to calculate properties for. + properties + List of properties to calculate. + system_changes + List of system changes to consider for calculation. + **kwargs + Any additional keyword arguments. + """ + super().calculate(atoms, properties, system_changes) + + if "energy" in properties: + self.results["energy"] = 0.0 + + if "forces" in properties: + self.results["forces"] = np.zeros((len(self.atoms), 3)) + + if "stress" in properties: + self.results["stress"] = np.zeros(6) + + +class MockErrorCalculator(Calculator): + """A mock calculator that raises an error for all properties.""" + + implemented_properties = ["energy", "forces", "stress"] + results = {} + parameters = {} + + def calculate(self, *args, **kwargs) -> None: + """ + Define calculation method for mock calculator. + + Parameters + ---------- + *args + Any additional positional arguments. + **kwargs + Any additional keyword arguments. + """ + raise ValueError("This is a mock error calculator. All calculations fail.") diff --git a/ml_peg/models/models.py b/ml_peg/models/models.py index e8980b13d..ecc592deb 100644 --- a/ml_peg/models/models.py +++ b/ml_peg/models/models.py @@ -231,7 +231,7 @@ def get_calculator(self, **kwargs) -> Calculator: Returns ------- Calculator - Loaded ASE Orb Calculator. + Loaded ASE fairchem Calculator. """ from fairchem.core import FAIRChemCalculator, pretrained_mlip # torch.serialization.add_safe_globals([slice]) @@ -257,3 +257,29 @@ def available(self) -> bool: return self.model_name in pretrained_mlip._MODEL_CKPTS.checkpoints except Exception: return False + + +@dataclasses.dataclass(kw_only=True) +class MockCalc(SumCalc): + """Dataclass for mock calculator.""" + + model_name: str = "mock" + trained_on_dispersion: bool = True + + def get_calculator(self, **kwargs) -> Calculator: + """ + Prepare and load the calculator. + + Parameters + ---------- + **kwargs + Any additional keyword arguments passed to `get_calculator`. + + Returns + ------- + Calculator + Loaded mock ASE Calculator. + """ + from ml_peg.models.mock import MockCalculator + + return MockCalculator()