diff --git a/docs/source/user_guide/benchmarks/physicality.rst b/docs/source/user_guide/benchmarks/physicality.rst index 8bbca2320..c2e42c20f 100644 --- a/docs/source/user_guide/benchmarks/physicality.rst +++ b/docs/source/user_guide/benchmarks/physicality.rst @@ -126,6 +126,30 @@ Metrics positive correlation, so a value of +1, indicating that as atoms get closer together, the energy increases. +Matbench Discovery metrics +-------------------------- + +A separate, opt-in result reports 12 homonuclear metrics without changing the +five-metric homo- and heteronuclear score. The reference-free metrics are tortuosity, +force flips, energy-difference flips, energy jump, force total variation, and force +jump. PBE-relative metrics cover energy and force MAE, repulsive-wall distance MAE, +bond-length error, well-depth error, and vibrational-frequency error. + +Element-specific windows keep the repulsive wall from dominating general metrics. +Scoring covers H-U except Po, At, Rn, Fr, and Ra; known-discontinuous PBE curves keep +reference-free metrics but not PBE-relative ones, and non-finite wall-window curves +are skipped. Projected forces map to ``-force_parallel`` on atom 0 and +``+force_parallel`` on atom 1. JSON results include schema and source-framework +versions. To write strict JSON: + +.. code-block:: python + + from ml_peg.analysis.physicality.diatomics.analyse_diatomics import ( + write_mbd_diatomic_metrics, + ) + + write_mbd_diatomic_metrics("mbd_diatomics_metrics.json") + Computational cost ------------------ @@ -134,7 +158,8 @@ High: Expected to take hours to run on GPU, or around one day for slower MLIPs. Data availability ----------------- -None required; diatomics are generated in ASE. +Predicted diatomics are generated in ASE. The optional PBE-relative metric family +uses the bundled Matbench Discovery DFT reference curves. Oxidation States diff --git a/ml_peg/analysis/physicality/diatomics/analyse_diatomics.py b/ml_peg/analysis/physicality/diatomics/analyse_diatomics.py index 649de320e..85fc0318b 100644 --- a/ml_peg/analysis/physicality/diatomics/analyse_diatomics.py +++ b/ml_peg/analysis/physicality/diatomics/analyse_diatomics.py @@ -10,6 +10,14 @@ import pytest from scipy.signal import find_peaks +from ml_peg.analysis.physicality.diatomics.metrics import ( + DEFAULT_DFT_REFERENCE_PATH, + DIATOMIC_METRIC_NAMES, + aggregate_finite_means, + calc_diatomic_metrics, + load_dft_reference_curves, + load_ml_peg_curves, +) from ml_peg.analysis.utils.decorators import build_table, periodic_curve_gallery from ml_peg.analysis.utils.utils import load_metrics_config from ml_peg.app import APP_ROOT @@ -22,6 +30,10 @@ OUT_PATH = APP_ROOT / "data" / "physicality" / "diatomics" CURVE_PATH = OUT_PATH / "curves" +RESULT_SCHEMA_VERSION = 1 +SOURCE_FRAMEWORK_ID = "matbench-discovery" +SOURCE_FRAMEWORK_VERSION = "1.3.1" + METRICS_CONFIG_PATH = Path(__file__).with_name("metrics.yml") DEFAULT_THRESHOLDS, DEFAULT_TOOLTIPS, _ = load_metrics_config(METRICS_CONFIG_PATH) @@ -235,6 +247,131 @@ def _load_pair_data() -> dict[str, pd.DataFrame]: return pair_data +def _json_safe_mbd_metrics( + metrics_by_element: dict[str, dict[str, float]], +) -> dict[str, dict[str, float | None]]: + """ + Convert non-finite MBD metrics to strict-JSON null values. + + Parameters + ---------- + metrics_by_element + Metric values grouped by element. + + Returns + ------- + dict[str, dict[str, float | None]] + JSON-safe metric values grouped by element. + """ + return { + element_symbol: { + metric_name: (float(metric_value) if np.isfinite(metric_value) else None) + for metric_name, metric_value in element_metrics.items() + } + for element_symbol, element_metrics in metrics_by_element.items() + } + + +def evaluate_mbd_diatomic_metrics( + pair_data: dict[str, pd.DataFrame] | None = None, + *, + reference_path: str | Path | None = None, + interpolate: bool | int = 200, +) -> dict[str, object]: + """ + Evaluate 12 homonuclear MBD metrics outside the legacy weighted score. + + Parameters + ---------- + pair_data + Optional model-to-dataframe mapping; calculator outputs are loaded by default. + reference_path + Optional DFT reference-curve path. + interpolate + Whether or how many points to use when interpolating curves. + + Returns + ------- + dict[str, object] + Versioned benchmark results and per-model metrics. + """ + resolved_reference_path = Path(reference_path or DEFAULT_DFT_REFERENCE_PATH) + reference_curves = load_dft_reference_curves( + functional="PBE", + ref_path=resolved_reference_path, + ) + model_data = pair_data if pair_data is not None else _load_pair_data() + model_results: dict[str, dict[str, object]] = {} + for model_name, model_dataframe in model_data.items(): + predicted_curves = load_ml_peg_curves( + model_dataframe, include_heteronuclear=False + ) + metrics_by_element = calc_diatomic_metrics( + reference_curves, + predicted_curves, + interpolate=interpolate, + ) + model_results[model_name] = { + "means": aggregate_finite_means(metrics_by_element), + "elements": _json_safe_mbd_metrics(metrics_by_element), + } + + return { + "schema_version": RESULT_SCHEMA_VERSION, + "source": { + "framework": SOURCE_FRAMEWORK_ID, + "version": SOURCE_FRAMEWORK_VERSION, + }, + "curve_scope": "homonuclear", + "weighted_in_legacy_score": False, + "reference": { + "functional": "PBE", + "file": resolved_reference_path.name, + }, + "interpolate": interpolate, + "metric_names": list(DIATOMIC_METRIC_NAMES), + "models": model_results, + } + + +def write_mbd_diatomic_metrics( + output_path: str | Path, + pair_data: dict[str, pd.DataFrame] | None = None, + *, + reference_path: str | Path | None = None, + interpolate: bool | int = 200, +) -> dict[str, object]: + """ + Evaluate MBD metrics and write JSON. + + Parameters + ---------- + output_path + Destination JSON path. + pair_data + Optional model-to-dataframe mapping; calculator outputs are loaded by default. + reference_path + Optional DFT reference-curve path. + interpolate + Whether or how many points to use when interpolating curves. + + Returns + ------- + dict[str, object] + Written benchmark results. + """ + result = evaluate_mbd_diatomic_metrics( + pair_data, + reference_path=reference_path, + interpolate=interpolate, + ) + resolved_output_path = Path(output_path) + resolved_output_path.parent.mkdir(parents=True, exist_ok=True) + with open(resolved_output_path, "w", encoding="utf-8") as file: + json.dump(result, file, indent=2, allow_nan=False) + return result + + @periodic_curve_gallery( curve_dir=CURVE_PATH, periodic_dir=None, @@ -250,12 +387,12 @@ def _load_pair_data() -> dict[str, pd.DataFrame]: ) def persist_diatomics_pair_data() -> dict[str, pd.DataFrame]: """ - Persist curve payloads and return the per-model dataframes. + Persist curve payloads and return per-model dataframes. Returns ------- dict[str, pd.DataFrame] - Mapping of model name to per-pair curve data. + Curve dataframes keyed by model name. """ return _load_pair_data() @@ -263,12 +400,12 @@ def persist_diatomics_pair_data() -> dict[str, pd.DataFrame]: @pytest.fixture def diatomics_pair_data_fixture() -> dict[str, pd.DataFrame]: """ - Load curve data and persist gallery assets for pytest use. + Load curve data and persist gallery assets for pytest. Returns ------- dict[str, pd.DataFrame] - Mapping of model name to per-pair curve data. + Curve dataframes keyed by model name. """ return persist_diatomics_pair_data() @@ -277,20 +414,17 @@ def collect_metrics( pair_data: dict[str, pd.DataFrame] | None = None, ) -> pd.DataFrame: """ - Gather metrics for all models. - - Metrics are averaged across all diatomic pairs (both homonuclear and heteronuclear). + Aggregate metrics across all homo- and heteronuclear pairs by model. Parameters ---------- pair_data - Optional mapping of model names to curve dataframes. When ``None``, - the data is loaded via ``persist_diatomics_pair_data``. + Optional curve dataframes keyed by model name. Returns ------- pd.DataFrame - Aggregated metrics table (all pairs). + One row of aggregated metrics per model. """ metrics_rows: list[dict[str, float | str]] = [] @@ -313,41 +447,21 @@ def diatomics_collection( diatomics_pair_data_fixture: dict[str, pd.DataFrame], ) -> pd.DataFrame: """ - Collect diatomics metrics across all models. + Collect per-model diatomic metrics. Parameters ---------- diatomics_pair_data_fixture - Mapping of model names to curve dataframes generated by the fixture. + Curve dataframes keyed by model name. Returns ------- pd.DataFrame - Aggregated metrics dataframe. + One row of aggregated metrics per model. """ return collect_metrics(diatomics_pair_data_fixture) -@pytest.fixture -def diatomics_metrics_dataframe( - diatomics_collection: pd.DataFrame, -) -> pd.DataFrame: - """ - Provide the aggregated diatomics metrics dataframe. - - Parameters - ---------- - diatomics_collection - Metrics dataframe produced by ``collect_metrics``. - - Returns - ------- - pd.DataFrame - Aggregated diatomics metrics indexed by model. - """ - return diatomics_collection - - @pytest.fixture @build_table( filename=OUT_PATH / "diatomics_metrics_table.json", @@ -356,42 +470,46 @@ def diatomics_metrics_dataframe( weights=None, ) def metrics( - diatomics_metrics_dataframe: pd.DataFrame, + diatomics_collection: pd.DataFrame, ) -> dict[str, dict]: """ - Compute diatomics metrics for all models. + Return metric-name mappings by model. Parameters ---------- - diatomics_metrics_dataframe - Aggregated per-model metrics produced by ``collect_metrics``. + diatomics_collection + Aggregated metrics with one row per model. Returns ------- dict[str, dict] - Mapping of metric names to per-model results. + Model values keyed by metric name. """ - metrics_df = diatomics_metrics_dataframe - metrics_dict: dict[str, dict[str, float | None]] = {} - for column in metrics_df.columns: - if column == "Model": - continue - values = [ - value if pd.notna(value) else None for value in metrics_df[column].tolist() - ] - metrics_dict[column] = dict(zip(metrics_df["Model"], values, strict=False)) - return metrics_dict + return { + column: dict( + zip( + diatomics_collection["Model"], + [ + value if pd.notna(value) else None + for value in diatomics_collection[column] + ], + strict=False, + ) + ) + for column in diatomics_collection + if column != "Model" + } @pytest.mark.framework("mace-multihead") def test_diatomics(metrics: dict[str, dict]) -> None: """ - Run diatomics analysis. + Write diatomic benchmark metadata after fixture evaluation. Parameters ---------- metrics - Benchmark metrics generated by fixtures. + Evaluated metric mappings supplied by pytest. """ mock_data = load_model_data("mock") # Write out info.json diff --git a/ml_peg/analysis/physicality/diatomics/data/README.md b/ml_peg/analysis/physicality/diatomics/data/README.md new file mode 100644 index 000000000..6a7ae1438 --- /dev/null +++ b/ml_peg/analysis/physicality/diatomics/data/README.md @@ -0,0 +1,10 @@ +# Diatomics DFT reference + +`diatomics-dft.json.gz` was copied without modification from +[`matbench_discovery/site/src/lib/diatomics-dft.json.gz`](https://github.com/janosh/matbench-discovery/blob/2c7f9fc42d018711dc2f5df573d225ea6d2d17b2/site/src/lib/diatomics-dft.json.gz) +at Matbench Discovery commit `2c7f9fc42d018711dc2f5df573d225ea6d2d17b2`. + +SHA-256: `1fe6334a82e98208ea74169a3beaf98cd5188bdc7ac40e518697fd36c7196e3d` + +The file contains homonuclear DFT energy and force curves. The imported Matbench +Discovery metrics use its PBE data. Do not modify the compressed file. diff --git a/ml_peg/analysis/physicality/diatomics/data/diatomics-dft.json.gz b/ml_peg/analysis/physicality/diatomics/data/diatomics-dft.json.gz new file mode 100644 index 000000000..58dd61ccb Binary files /dev/null and b/ml_peg/analysis/physicality/diatomics/data/diatomics-dft.json.gz differ diff --git a/ml_peg/analysis/physicality/diatomics/metrics/__init__.py b/ml_peg/analysis/physicality/diatomics/metrics/__init__.py new file mode 100644 index 000000000..a437ab1a2 --- /dev/null +++ b/ml_peg/analysis/physicality/diatomics/metrics/__init__.py @@ -0,0 +1,449 @@ +""" +Diatomic potential-energy metrics adapted from Matbench Discovery. + +The smoothness approach follows Stenczel et al., https://arxiv.org/abs/2401.00096. +""" + +from __future__ import annotations + +from collections.abc import Callable +import logging +from typing import Any + +from ase.data import atomic_numbers, covalent_radii, vdw_alvarez +import numpy as np + +from ml_peg.analysis.physicality.diatomics.metrics.energy import ( + calc_energy_diff_flips, + calc_energy_jump, + calc_pbe_bond_length_error, + calc_pbe_energy_mae, + calc_pbe_vib_freq_error, + calc_pbe_wall_dist_mae, + calc_pbe_well_depth_error, + calc_tortuosity, +) +from ml_peg.analysis.physicality.diatomics.metrics.force import ( + calc_force_flips, + calc_force_jump, + calc_force_mae, + calc_force_total_variation, +) +from ml_peg.analysis.physicality.diatomics.metrics.schema import ( + DEFAULT_DFT_REFERENCE_PATH, + DiatomicCurve, + DiatomicCurves, + curves_from_ml_peg_dataframe, + homo_key, + load_dft_reference_curves, + load_mbd_json, + load_ml_peg_curves, +) + +logger = logging.getLogger(__name__) + +TORTUOSITY = "tortuosity" +FORCE_FLIPS = "force_flips" +ENERGY_JUMP = "energy_jump" +ENERGY_DIFF_FLIPS = "energy_diff_flips" +FORCE_TOTAL_VARIATION = "force_total_variation" +FORCE_JUMP = "force_jump" +PBE_WALL_DIST_MAE = "pbe_wall_dist_mae" +PBE_ENERGY_MAE = "pbe_energy_mae" +PBE_BOND_LENGTH_ERROR = "pbe_bond_length_error" +PBE_WELL_DEPTH_ERROR = "pbe_well_depth_error" +PBE_FORCE_MAE = "pbe_force_mae" +PBE_VIB_FREQ_ERROR = "pbe_vib_freq_error" + +DIATOMIC_METRIC_NAMES: tuple[str, ...] = ( + TORTUOSITY, + FORCE_FLIPS, + ENERGY_JUMP, + ENERGY_DIFF_FLIPS, + FORCE_TOTAL_VARIATION, + FORCE_JUMP, + PBE_WALL_DIST_MAE, + PBE_ENERGY_MAE, + PBE_BOND_LENGTH_ERROR, + PBE_WELL_DEPTH_ERROR, + PBE_FORCE_MAE, + PBE_VIB_FREQ_ERROR, +) +DIATOMIC_METRIC_KEYS = frozenset(DIATOMIC_METRIC_NAMES) + +# Skip H-U elements absent from Materials Project training data. +NON_MP_ELEMENTS = frozenset({"Po", "At", "Rn", "Fr", "Ra"}) +DIATOMIC_WALL_R_MIN_FACTOR = 0.8 +MAX_SCORED_ATOMIC_NUMBER = 92 + + +def find_low_quality_dft_refs( + ref_curves: DiatomicCurves, + *, + min_energy_jump: float = 1.5, + min_energy_flips: int = 3, +) -> set[str]: + """ + Find non-finite or discontinuous DFT references unsuitable for scoring. + + Parameters + ---------- + ref_curves + DFT reference curves. + min_energy_jump + Minimum discontinuity magnitude used by the quality gate. + min_energy_flips + Minimum number of energy-difference flips used by the quality gate. + + Returns + ------- + set[str] + Element symbols with unsuitable reference curves. + """ + low_quality: set[str] = set() + for element_symbol, curve in ref_curves.homo_nuclear.items(): + separations = curve.distances + energies = curve.energies + if separations.size == 0: + continue + radius_min, radius_max = eval_window(element_symbol, float(np.max(separations))) + window_mask = (separations >= radius_min) & (separations <= radius_max) + if window_mask.sum() < 5: + continue # too few in-window points to assess smoothness + if not np.isfinite(energies[window_mask]).all(): + # Non-finite references cannot be scored. + low_quality.add(element_symbol) + continue + if ( + calc_energy_jump(separations[window_mask], energies[window_mask]) + >= min_energy_jump + and calc_energy_diff_flips(separations[window_mask], energies[window_mask]) + >= min_energy_flips + ): + low_quality.add(element_symbol) + return low_quality + + +def eval_window( + elem_symbol: str, + seps_max: float, + *, + r_min_factor: float = 0.9, +) -> tuple[float, float]: + """ + Return the covalent-to-van-der-Waals evaluation window in Å. + + Parameters + ---------- + elem_symbol + Element or pair label. + seps_max + Largest available separation. + r_min_factor + Covalent-radius multiplier for the lower bound. + + Returns + ------- + tuple[float, float] + Lower and upper evaluation bounds. + """ + atomic_number = atomic_numbers[elem_symbol.split("-", maxsplit=1)[0]] + covalent_radius = ( + covalent_radii[atomic_number] if atomic_number < len(covalent_radii) else np.nan + ) + radius_min = r_min_factor * covalent_radius if np.isfinite(covalent_radius) else 0.0 + vdw_radii = vdw_alvarez.vdw_radii + vdw_radius = vdw_radii[atomic_number] if atomic_number < len(vdw_radii) else np.nan + radius_max = ( + min(3.1 * vdw_radius, seps_max) if np.isfinite(vdw_radius) else seps_max + ) + return radius_min, radius_max + + +def calc_diatomic_metrics( + ref_curves: DiatomicCurves | None, + pred_curves: DiatomicCurves, + metrics: dict[str, dict[str, Any]] | None = None, + *, + interpolate: bool | int = False, +) -> dict[str, dict[str, float]]: + """ + Calculate requested metrics for supported homonuclear curves by element. + + Low-quality references receive self-consistency metrics but no ``pbe_*`` metrics. + + Parameters + ---------- + ref_curves + Optional DFT reference curves. + pred_curves + Predicted diatomic curves. + metrics + Optional mapping of requested metric names to keyword arguments. + interpolate + Whether or how many points to use when interpolating paired curves. + + Returns + ------- + dict[str, dict[str, float]] + Metric values grouped by element. + """ + requested_metric_keys = ( + set(metrics) if metrics is not None else set(DIATOMIC_METRIC_KEYS) + ) + unknown_metrics = requested_metric_keys - DIATOMIC_METRIC_KEYS + if unknown_metrics: + raise ValueError( + f"unknown_metrics={unknown_metrics}. " + f"Valid metrics={sorted(DIATOMIC_METRIC_KEYS)}" + ) + metric_kwargs = {key: kwargs.copy() for key, kwargs in (metrics or {}).items()} + for metric_key in (PBE_ENERGY_MAE, PBE_FORCE_MAE): + if metric_key in requested_metric_keys: + metric_kwargs.setdefault(metric_key, {}).setdefault( + "interpolate", interpolate + ) + + low_quality_refs = find_low_quality_dft_refs(ref_curves) if ref_curves else set() + results: dict[str, dict[str, float]] = {} + seen_elements: set[str] = set() + for element_symbol, pred_data in pred_curves.homo_nuclear.items(): + normalized_element = homo_key(element_symbol) + if normalized_element in seen_elements: + raise ValueError( + f"Duplicate homonuclear curve for element {normalized_element!r}" + ) + seen_elements.add(normalized_element) + if ( + normalized_element in NON_MP_ELEMENTS + or atomic_numbers[normalized_element] > MAX_SCORED_ATOMIC_NUMBER + ): + continue # score all models on the same MP-supported element set + # General metrics use the MLIP Arena window; wall metrics extend to + # 0.8 times the covalent radius. + predicted_distances = pred_data.distances + separations_max = float(predicted_distances.max()) + radius_min, radius_max = eval_window(element_symbol, separations_max) + predicted_mask = (predicted_distances >= radius_min) & ( + predicted_distances <= radius_max + ) + if predicted_mask.sum() < 5: # too few points in window for stable metrics + logger.info( + "Skipping %s diatomic metrics: <5 points in eval window", + element_symbol, + ) + continue + predicted_separations = predicted_distances[predicted_mask] + predicted_energies_raw = pred_data.energies + predicted_energies = predicted_energies_raw[predicted_mask] + predicted_forces_raw = pred_data.forces + if not predicted_forces_raw.size: + raise ValueError(f"{element_symbol} diatomic curve is missing forces") + if len(predicted_forces_raw) != len(predicted_distances): + raise ValueError( + f"{element_symbol} diatomic force and distance counts differ: " + f"{len(predicted_forces_raw)} != {len(predicted_distances)}" + ) + predicted_forces = predicted_forces_raw[predicted_mask] + + wall_radius_min = eval_window( + element_symbol, + separations_max, + r_min_factor=DIATOMIC_WALL_R_MIN_FACTOR, + )[0] + # Include a generated DFT endpoint that differs from 0.8*r_cov by one ulp. + wall_radius_min -= 1e-12 + predicted_wall_mask = (predicted_distances >= wall_radius_min) & ( + predicted_distances <= radius_max + ) + if not ( + np.isfinite(predicted_energies_raw[predicted_wall_mask]).all() + and np.isfinite(predicted_forces_raw[predicted_wall_mask]).all() + ): + logger.info( + "Skipping %s diatomic metrics: non-finite wall values", element_symbol + ) + continue + + energy_args = (predicted_separations, predicted_energies) + force_args = (predicted_separations, predicted_forces) + # Calls for metrics that need only the predicted curve. + metric_calls: list[tuple[str, Callable[..., float], tuple[Any, ...]]] = [ + (TORTUOSITY, calc_tortuosity, energy_args), + (ENERGY_DIFF_FLIPS, calc_energy_diff_flips, energy_args), + (ENERGY_JUMP, calc_energy_jump, energy_args), + (FORCE_FLIPS, calc_force_flips, force_args), + (FORCE_TOTAL_VARIATION, calc_force_total_variation, force_args), + (FORCE_JUMP, calc_force_jump, force_args), + ] + + # Add relative metrics only for references that pass the quality gate. + reference_data = ( + ref_curves.homo_nuclear.get(normalized_element) + if ref_curves and normalized_element not in low_quality_refs + else None + ) + if reference_data is not None: + reference_distances = reference_data.distances + reference_mask = (reference_distances >= radius_min) & ( + reference_distances <= radius_max + ) + reference_separations = reference_distances[reference_mask] + reference_energies_raw = reference_data.energies + reference_energies = reference_energies_raw[reference_mask] + if len(reference_separations) >= 2: + pair_args = ( + reference_separations, + reference_energies, + predicted_separations, + predicted_energies, + ) + metric_calls[:0] = [ + (PBE_ENERGY_MAE, calc_pbe_energy_mae, pair_args), + (PBE_BOND_LENGTH_ERROR, calc_pbe_bond_length_error, pair_args), + (PBE_WELL_DEPTH_ERROR, calc_pbe_well_depth_error, pair_args), + ( + PBE_VIB_FREQ_ERROR, + calc_pbe_vib_freq_error, + (element_symbol, *pair_args), + ), + ] + reference_wall_mask = (reference_distances >= wall_radius_min) & ( + reference_distances <= radius_max + ) + if predicted_wall_mask.sum() >= 2 and reference_wall_mask.sum() >= 2: + wall_args = ( + reference_distances[reference_wall_mask], + reference_energies_raw[reference_wall_mask], + predicted_distances[predicted_wall_mask], + predicted_energies_raw[predicted_wall_mask], + ) + metric_calls.insert( + 0, + (PBE_WALL_DIST_MAE, calc_pbe_wall_dist_mae, wall_args), + ) + reference_forces = reference_data.forces + if ( + reference_forces.size + and len(reference_forces) == len(reference_distances) + and len(reference_separations) >= 2 + ): + reference_forces = reference_forces[reference_mask] + force_interpolate = metric_kwargs.get(PBE_FORCE_MAE, {}).get( + "interpolate", False + ) + same_grid = np.array_equal(reference_separations, predicted_separations) + has_overlap = max( + reference_separations.min(), + predicted_separations.min(), + ) < min( + reference_separations.max(), + predicted_separations.max(), + ) + if same_grid or (force_interpolate and has_overlap): + force_pair_args = ( + reference_separations, + reference_forces, + predicted_separations, + predicted_forces, + ) + metric_calls.append( + (PBE_FORCE_MAE, calc_force_mae, force_pair_args) + ) + + results[normalized_element] = { + metric_key: metric_function( + *metric_args, + **metric_kwargs.get(metric_key, {}), + ) + for metric_key, metric_function, metric_args in metric_calls + if metric_key in requested_metric_keys + } + + return results + + +def aggregate_finite_means( + metrics_by_element: dict[str, dict[str, float]], +) -> dict[str, float]: + """ + Average finite values for every metric present, to four significant digits. + + Parameters + ---------- + metrics_by_element + Metric values grouped by element. + + Returns + ------- + dict[str, float] + Finite metric means keyed by metric name. + """ + metric_means: dict[str, float] = {} + metric_names = dict.fromkeys( + metric_name + for element_metrics in metrics_by_element.values() + for metric_name in element_metrics + ) + for metric_name in metric_names: + finite_values = [ + metric_value + for element_metrics in metrics_by_element.values() + if (metric_value := element_metrics.get(metric_name)) is not None + and np.isfinite(metric_value) + ] + if finite_values: + value_scale = max(abs(metric_value) for metric_value in finite_values) + if value_scale == 0: + metric_mean = 0.0 + else: + metric_mean = value_scale * ( + sum(metric_value / value_scale for metric_value in finite_values) + / len(finite_values) + ) + if np.isfinite(metric_mean): + metric_means[metric_name] = float(f"{metric_mean:.4}") + return metric_means + + +__all__ = [ + "DEFAULT_DFT_REFERENCE_PATH", + "DIATOMIC_METRIC_KEYS", + "DIATOMIC_METRIC_NAMES", + "DIATOMIC_WALL_R_MIN_FACTOR", + "ENERGY_DIFF_FLIPS", + "ENERGY_JUMP", + "FORCE_FLIPS", + "FORCE_JUMP", + "FORCE_TOTAL_VARIATION", + "NON_MP_ELEMENTS", + "PBE_BOND_LENGTH_ERROR", + "PBE_ENERGY_MAE", + "PBE_FORCE_MAE", + "PBE_VIB_FREQ_ERROR", + "PBE_WALL_DIST_MAE", + "PBE_WELL_DEPTH_ERROR", + "TORTUOSITY", + "DiatomicCurve", + "DiatomicCurves", + "aggregate_finite_means", + "calc_diatomic_metrics", + "calc_energy_diff_flips", + "calc_energy_jump", + "calc_force_flips", + "calc_force_jump", + "calc_force_mae", + "calc_force_total_variation", + "calc_pbe_bond_length_error", + "calc_pbe_energy_mae", + "calc_pbe_vib_freq_error", + "calc_pbe_wall_dist_mae", + "calc_pbe_well_depth_error", + "calc_tortuosity", + "curves_from_ml_peg_dataframe", + "eval_window", + "find_low_quality_dft_refs", + "load_dft_reference_curves", + "load_mbd_json", + "load_ml_peg_curves", +] diff --git a/ml_peg/analysis/physicality/diatomics/metrics/energy.py b/ml_peg/analysis/physicality/diatomics/metrics/energy.py new file mode 100644 index 000000000..fb14bead3 --- /dev/null +++ b/ml_peg/analysis/physicality/diatomics/metrics/energy.py @@ -0,0 +1,703 @@ +"""Energy-based metrics for diatomic curves.""" + +from __future__ import annotations + +from typing import Literal + +from ase.data import atomic_masses, atomic_numbers +import numpy as np +from numpy.typing import ArrayLike + +PBE_WALL_ENERGY_THRESHOLDS_EV: tuple[float, ...] = (1, 5, 10, 20, 50, 100) + + +def _validate_diatomic_curve( + separations: ArrayLike, + values: ArrayLike, + *, + normalize_energy: bool = False, + value_kind: Literal["energy", "force"] = "energy", +) -> tuple[np.ndarray, np.ndarray]: + """ + Validate, sort, and optionally far-field-normalize a sampled curve. + + Parameters + ---------- + separations + Sample separations. + values + Sampled energy or force values. + normalize_energy + Whether to shift the last energy sample to zero. + value_kind + Kind of sampled values, used for shape validation. + + Returns + ------- + tuple[np.ndarray, np.ndarray] + Sorted separations and values. + """ + separation_array = np.asarray(separations) + value_array = np.asarray(values) + + if separation_array.ndim != 1: + raise ValueError( + f"separations must have shape (n,), got {separation_array.shape}" + ) + if value_kind == "energy" and value_array.ndim != 1: + raise ValueError(f"energy values must have shape (n,), got {value_array.shape}") + if value_kind == "force" and ( + value_array.ndim != 3 or value_array.shape[1:] != (2, 3) + ): + raise ValueError( + f"force values must have shape (n, 2, 3), got {value_array.shape}" + ) + + if len(separation_array) != len(value_array): + raise ValueError( + f"len(separation_array)={len(separation_array)} != " + f"len(value_array)={len(value_array)}" + ) + if len(separation_array) < 2: + raise ValueError( + "Input must have at least 2 points, " + f"got len(separation_array)={len(separation_array)}" + ) + n_separation_nan = int(np.isnan(separation_array).sum()) + n_value_nan = int(np.isnan(value_array).sum()) + if n_separation_nan or n_value_nan: + raise ValueError( + "Input contains NaN values: " + f"n_separation_nan={n_separation_nan}, n_value_nan={n_value_nan}" + ) + n_separation_inf = int(np.isinf(separation_array).sum()) + n_value_inf = int(np.isinf(value_array).sum()) + if n_separation_inf or n_value_inf: + raise ValueError( + "Input contains infinite values: " + f"n_separation_inf={n_separation_inf}, n_value_inf={n_value_inf}" + ) + n_unique = len(np.unique(separation_array)) + if n_unique != len(separation_array): + raise ValueError( + f"separations contains {len(separation_array) - n_unique} duplicates" + ) + + sort_indices = np.argsort(separation_array) + separation_array = separation_array[sort_indices] + value_array = value_array[sort_indices] + + # Normalize energy curves to zero at the largest separation. + if normalize_energy and value_array.ndim == 1: + # The ascending sort places that sample last. + value_array = value_array - value_array[-1] + + return separation_array, value_array + + +def _interpolation_point_count(interpolate: bool | int) -> int: + """ + Return the requested interpolation size, validating a two-point minimum. + + Parameters + ---------- + interpolate + Whether or how many interpolation points to use. + + Returns + ------- + int + Number of interpolation points. + """ + n_points = 100 if interpolate is True else int(interpolate) + if n_points < 2: + raise ValueError("interpolate must request at least 2 points") + return n_points + + +def _common_grid_curve_pair( + separations_ref: ArrayLike, + values_ref: ArrayLike, + separations_pred: ArrayLike, + values_pred: ArrayLike, + *, + interpolate: bool | int, + value_kind: Literal["energy", "force"] = "energy", +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """ + Validate two curves and optionally interpolate their common interval. + + Parameters + ---------- + separations_ref + Reference-curve separations. + values_ref + Reference-curve values. + separations_pred + Predicted-curve separations. + values_pred + Predicted-curve values. + interpolate + Whether or how many common-grid points to use. + value_kind + Kind of sampled values, used for shape validation. + + Returns + ------- + tuple[np.ndarray, np.ndarray, np.ndarray] + Common separations, reference values, and predicted values. + """ + separations_ref, values_ref = _validate_diatomic_curve( + separations_ref, values_ref, value_kind=value_kind + ) + separations_pred, values_pred = _validate_diatomic_curve( + separations_pred, values_pred, value_kind=value_kind + ) + if not interpolate: + if not np.array_equal(separations_ref, separations_pred): + raise ValueError( + "Reference and predicted distances must be same when " + f"interpolate={interpolate}\n" + f"separations_ref={separations_ref}, " + f"separations_pred={separations_pred}" + ) + return separations_ref, values_ref, values_pred + + data_min = max(separations_ref.min(), separations_pred.min()) + data_max = min(separations_ref.max(), separations_pred.max()) + if data_min >= data_max: + curve_label = "force curves" if value_kind == "force" else "curves" + raise ValueError( + f"Cannot interpolate {curve_label} with no overlap: " + f"data_min={data_min}, data_max={data_max}" + ) + common_grid = np.linspace( + data_min, data_max, _interpolation_point_count(interpolate) + ) + + def interpolate_values(separations: np.ndarray, values: np.ndarray) -> np.ndarray: + """ + Interpolate all flattened value components onto ``common_grid``. + + Parameters + ---------- + separations + Source separations. + values + Source values. + + Returns + ------- + np.ndarray + Values interpolated onto the common grid. + """ + flattened_values = values.reshape(len(values), -1) + interpolated = np.column_stack( + [ + np.interp( + common_grid, + separations, + flattened_values[:, component_index], + ) + for component_index in range(flattened_values.shape[1]) + ] + ) + return interpolated.reshape(len(common_grid), *values.shape[1:]) + + return ( + common_grid, + interpolate_values(separations_ref, values_ref), + interpolate_values(separations_pred, values_pred), + ) + + +def _binding_energy(energies: np.ndarray) -> float: + """ + Return well depth relative to the largest sampled separation. + + Parameters + ---------- + energies + Sampled energies ordered by separation. + + Returns + ------- + float + Binding energy. + """ + return float(energies[-1] - np.min(energies)) + + +def _validated_energy_pair( + seps_ref: ArrayLike, + energy_ref: ArrayLike, + seps_pred: ArrayLike, + energy_pred: ArrayLike, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """ + Validate two independently sampled energy curves. + + Parameters + ---------- + seps_ref + Reference separations. + energy_ref + Reference energies. + seps_pred + Predicted separations. + energy_pred + Predicted energies. + + Returns + ------- + tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray] + Sorted reference separations and energies followed by predicted values. + """ + separations_ref, energies_ref = _validate_diatomic_curve(seps_ref, energy_ref) + separations_pred, energies_pred = _validate_diatomic_curve(seps_pred, energy_pred) + return separations_ref, energies_ref, separations_pred, energies_pred + + +def _quadratic_well_fit( + separations: ArrayLike, + energies: ArrayLike, + n_fit_points: int = 5, +) -> tuple[float, float]: + """ + Estimate equilibrium separation and curvature by local quadratic fit. + + Parameters + ---------- + separations + Sample separations. + energies + Sample energies. + n_fit_points + Maximum number of local points to fit. + + Returns + ------- + tuple[float, float] + Equilibrium separation and fitted curvature. + """ + separations, energies = _validate_diatomic_curve(separations, energies) + minimum_index = int(np.argmin(energies)) + if len(separations) < 3: + return float(separations[minimum_index]), np.nan + + start_index = min( + max(0, minimum_index - n_fit_points // 2), + max(0, len(separations) - n_fit_points), + ) + fit_separations = separations[start_index : start_index + n_fit_points] + fit_energies = energies[start_index : start_index + n_fit_points] + if len(fit_separations) < 3: + return float(separations[minimum_index]), np.nan + + quadratic_coefficient, linear_coefficient, _constant_coefficient = np.polyfit( + fit_separations, fit_energies, 2 + ) + curvature = 2 * quadratic_coefficient + if quadratic_coefficient <= 0: + return float(separations[minimum_index]), np.nan + + equilibrium_distance = -linear_coefficient / (2 * quadratic_coefficient) + if fit_separations.min() <= equilibrium_distance <= fit_separations.max(): + return float(equilibrium_distance), float(curvature) + return float(separations[minimum_index]), float(curvature) + + +def _repulsive_radius_at_threshold( + separations: ArrayLike, + energies: ArrayLike, + threshold_ev: float, +) -> float: + """ + Return the repulsive radius at an energy threshold, or NaN if unreached. + + Parameters + ---------- + separations + Sample separations. + energies + Sample energies. + threshold_ev + Energy above the curve minimum. + + Returns + ------- + float + Interpolated repulsive radius or NaN. + """ + separations, energies = _validate_diatomic_curve(separations, energies) + minimum_index = int(np.argmin(energies)) + if minimum_index == 0: + return np.nan + + radii_inward = separations[minimum_index::-1] + energy_above_minimum = energies[minimum_index::-1] - energies[minimum_index] + monotonic_energy = np.maximum.accumulate(energy_above_minimum) + unique_energy, unique_indices = np.unique(monotonic_energy, return_index=True) + if len(unique_energy) < 2 or threshold_ev > unique_energy[-1]: + return np.nan + return float(np.interp(threshold_ev, unique_energy, radii_inward[unique_indices])) + + +def calc_pbe_wall_dist_mae( + seps_ref: ArrayLike, + energy_ref: ArrayLike, + seps_pred: ArrayLike, + energy_pred: ArrayLike, + *, + thresholds_ev: tuple[float, ...] = PBE_WALL_ENERGY_THRESHOLDS_EV, +) -> float: + """ + Calculate mean PBE wall-radius error over reachable energy thresholds. + + A missing predicted crossing receives the full reference-radius error. + + Parameters + ---------- + seps_ref + Reference separations. + energy_ref + Reference energies. + seps_pred + Predicted separations. + energy_pred + Predicted energies. + thresholds_ev + Energy thresholds above the well minimum. + + Returns + ------- + float + Mean absolute wall-radius error. + """ + errors: list[float] = [] + for threshold_ev in thresholds_ev: + radius_ref = _repulsive_radius_at_threshold(seps_ref, energy_ref, threshold_ev) + if not np.isfinite(radius_ref): + continue + radius_pred = _repulsive_radius_at_threshold( + seps_pred, energy_pred, threshold_ev + ) + errors.append( + abs(radius_pred - radius_ref) if np.isfinite(radius_pred) else radius_ref + ) + return float(np.mean(errors)) if errors else np.nan + + +def calc_pbe_energy_mae( + seps_ref: ArrayLike, + energy_ref: ArrayLike, + seps_pred: ArrayLike, + energy_pred: ArrayLike, + *, + interpolate: bool | int = 200, +) -> float: + """ + Calculate PBE energy MAE after optional interpolation and far-field alignment. + + Parameters + ---------- + seps_ref + Reference separations. + energy_ref + Reference energies. + seps_pred + Predicted separations. + energy_pred + Predicted energies. + interpolate + Whether or how many common-grid points to use. + + Returns + ------- + float + Mean absolute energy error. + """ + _, energy_ref, energy_pred = _common_grid_curve_pair( + seps_ref, + energy_ref, + seps_pred, + energy_pred, + interpolate=interpolate, + ) + energy_ref = energy_ref - energy_ref[-1] + energy_pred = energy_pred - energy_pred[-1] + return float(np.mean(np.abs(energy_pred - energy_ref))) + + +def calc_pbe_bond_length_error( + seps_ref: ArrayLike, + energy_ref: ArrayLike, + seps_pred: ArrayLike, + energy_pred: ArrayLike, + *, + min_ref_binding_ev: float = 0.05, +) -> float: + """ + Calculate absolute PBE equilibrium-distance error, or NaN if unbound. + + Parameters + ---------- + seps_ref + Reference separations. + energy_ref + Reference energies. + seps_pred + Predicted separations. + energy_pred + Predicted energies. + min_ref_binding_ev + Minimum reference binding energy required for scoring. + + Returns + ------- + float + Absolute equilibrium-distance error or NaN. + """ + separations_ref, energy_ref, separations_pred, energy_pred = _validated_energy_pair( + seps_ref, energy_ref, seps_pred, energy_pred + ) + if _binding_energy(energy_ref) < min_ref_binding_ev: + return np.nan + reference_distance = _quadratic_well_fit(separations_ref, energy_ref)[0] + predicted_distance = _quadratic_well_fit(separations_pred, energy_pred)[0] + return float(abs(predicted_distance - reference_distance)) + + +def calc_pbe_well_depth_error( + seps_ref: ArrayLike, + energy_ref: ArrayLike, + seps_pred: ArrayLike, + energy_pred: ArrayLike, + *, + min_ref_binding_ev: float = 0.05, +) -> float: + """ + Calculate absolute PBE well-depth error, or NaN if unbound. + + Parameters + ---------- + seps_ref + Reference separations. + energy_ref + Reference energies. + seps_pred + Predicted separations. + energy_pred + Predicted energies. + min_ref_binding_ev + Minimum reference binding energy required for scoring. + + Returns + ------- + float + Absolute well-depth error or NaN. + """ + _, energy_ref, _, energy_pred = _validated_energy_pair( + seps_ref, energy_ref, seps_pred, energy_pred + ) + reference_depth = _binding_energy(energy_ref) + if reference_depth < min_ref_binding_ev: + return np.nan + return float(abs(_binding_energy(energy_pred) - reference_depth)) + + +def _vibrational_wavenumber_cm( + element_symbol: str, + curvature_ev_per_a2: float, +) -> float: + """ + Convert a homonuclear force constant to harmonic wavenumber in cm⁻¹. + + Parameters + ---------- + element_symbol + Element or homonuclear pair label. + curvature_ev_per_a2 + Energy-well curvature in eV/Ų. + + Returns + ------- + float + Harmonic wavenumber in cm⁻¹ or NaN. + """ + if not np.isfinite(curvature_ev_per_a2) or curvature_ev_per_a2 <= 0: + return np.nan + atomic_symbol = element_symbol.split("-", maxsplit=1)[0] + reduced_mass_kg = ( + atomic_masses[atomic_numbers[atomic_symbol]] * 1.66053906660e-27 / 2 + ) + force_constant_n_per_m = curvature_ev_per_a2 * 16.02176634 + angular_frequency_per_second = np.sqrt(force_constant_n_per_m / reduced_mass_kg) + return float(angular_frequency_per_second / (2 * np.pi * 2.99792458e10)) + + +def calc_pbe_vib_freq_error( + elem_symbol: str, + seps_ref: ArrayLike, + energy_ref: ArrayLike, + seps_pred: ArrayLike, + energy_pred: ArrayLike, + *, + min_ref_binding_ev: float = 0.05, +) -> float: + """ + Calculate absolute PBE vibrational-wavenumber error, or NaN if unbound. + + Parameters + ---------- + elem_symbol + Element or homonuclear pair label. + seps_ref + Reference separations. + energy_ref + Reference energies. + seps_pred + Predicted separations. + energy_pred + Predicted energies. + min_ref_binding_ev + Minimum reference binding energy required for scoring. + + Returns + ------- + float + Absolute vibrational-wavenumber error or NaN. + """ + separations_ref, energy_ref, separations_pred, energy_pred = _validated_energy_pair( + seps_ref, energy_ref, seps_pred, energy_pred + ) + if _binding_energy(energy_ref) < min_ref_binding_ev: + return np.nan + reference_curvature = _quadratic_well_fit(separations_ref, energy_ref)[1] + predicted_curvature = _quadratic_well_fit(separations_pred, energy_pred)[1] + reference_wavenumber = _vibrational_wavenumber_cm(elem_symbol, reference_curvature) + predicted_wavenumber = _vibrational_wavenumber_cm(elem_symbol, predicted_curvature) + return float(abs(predicted_wavenumber - reference_wavenumber)) + + +def calc_tortuosity(seps: ArrayLike, energies: ArrayLike) -> float: + """ + Calculate projected arc-chord energy tortuosity, or NaN if constant. + + Parameters + ---------- + seps + Sample separations. + energies + Sample energies. + + Returns + ------- + float + Curve tortuosity or NaN. + """ + _, energies = _validate_diatomic_curve(seps, energies) + + total_energy_variation = np.sum(np.abs(np.diff(energies))) + minimum_energy = np.min(energies) + direct_energy_difference = abs(energies[0] - minimum_energy) + abs( + energies[-1] - minimum_energy + ) + + if direct_energy_difference == 0: + return np.nan + return float(total_energy_variation / direct_energy_difference) + + +def _threshold_diff_signs( + values: np.ndarray, + threshold: float = 1e-3, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """ + Return nonzero thresholded differences, their signs, and flip mask. + + Parameters + ---------- + values + Sample values. + threshold + Magnitudes below this value are treated as zero. + + Returns + ------- + tuple[np.ndarray, np.ndarray, np.ndarray] + Nonzero differences, their signs, and adjacent sign-flip mask. + """ + differences = np.diff(values) + differences[np.abs(differences) < threshold] = 0 + signs = np.sign(differences) + nonzero_mask = signs != 0 + differences, signs = differences[nonzero_mask], signs[nonzero_mask] + flips = np.diff(signs) != 0 + return differences, signs, flips + + +def _jump_magnitude(values: np.ndarray, threshold: float = 1e-3) -> float: + """ + Sum adjacent step magnitudes at sign-flip points. + + Parameters + ---------- + values + Sample values. + threshold + Difference magnitudes below this value are ignored. + + Returns + ------- + float + Total adjacent jump magnitude. + """ + differences, _, flips = _threshold_diff_signs(values, threshold) + return float( + np.abs(differences[:-1][flips]).sum() + np.abs(differences[1:][flips]).sum() + ) + + +def calc_energy_diff_flips( + seps: ArrayLike, + energies: ArrayLike, +) -> float: + """ + Calculate the number of thresholded energy-difference sign flips. + + Parameters + ---------- + seps + Sample separations. + energies + Sample energies. + + Returns + ------- + float + Number of energy-difference sign flips. + """ + _, energies = _validate_diatomic_curve(seps, energies) + _, _, flips = _threshold_diff_signs(energies) + return float(np.sum(flips)) + + +def calc_energy_jump(seps: ArrayLike, energies: ArrayLike) -> float: + """ + Calculate total energy-step magnitude around sign-flip points. + + Parameters + ---------- + seps + Sample separations. + energies + Sample energies. + + Returns + ------- + float + Total energy-step magnitude. + """ + _, energies = _validate_diatomic_curve(seps, energies) + return _jump_magnitude(energies) diff --git a/ml_peg/analysis/physicality/diatomics/metrics/force.py b/ml_peg/analysis/physicality/diatomics/metrics/force.py new file mode 100644 index 000000000..bbc2df10d --- /dev/null +++ b/ml_peg/analysis/physicality/diatomics/metrics/force.py @@ -0,0 +1,144 @@ +"""Force-based metrics for diatomic curves.""" + +from __future__ import annotations + +import numpy as np +from numpy.typing import ArrayLike + +from ml_peg.analysis.physicality.diatomics.metrics.energy import ( + _common_grid_curve_pair, + _jump_magnitude, + _validate_diatomic_curve, +) + + +def calc_force_mae( + seps_ref: ArrayLike, + f_ref: ArrayLike, + seps_pred: ArrayLike, + f_pred: ArrayLike, + *, + interpolate: bool | int = False, +) -> float: + """ + Calculate force MAE, optionally interpolating over the shared range. + + Parameters + ---------- + seps_ref + Reference separations. + f_ref + Reference Cartesian forces. + seps_pred + Predicted separations. + f_pred + Predicted Cartesian forces. + interpolate + Whether or how many common-grid points to use. + + Returns + ------- + float + Mean absolute force error. + """ + _, f_ref, f_pred = _common_grid_curve_pair( + seps_ref, + f_ref, + seps_pred, + f_pred, + interpolate=interpolate, + value_kind="force", + ) + return float(np.mean(np.abs(f_ref - f_pred))) + + +def _radial_forces(seps: ArrayLike, forces: np.ndarray) -> np.ndarray: + """ + Validate a force curve and return first-atom radial forces. + + Parameters + ---------- + seps + Sample separations. + forces + Cartesian forces for both atoms. + + Returns + ------- + np.ndarray + First-atom radial forces. + """ + _, forces = _validate_diatomic_curve(seps, forces, value_kind="force") + return forces[:, 0, 0] # x-component of force on first atom + + +def calc_force_flips( + seps: ArrayLike, + forces: np.ndarray, + threshold: float = 1e-2, # 10meV/A threshold as in reference code +) -> float: + """ + Count thresholded direction changes in the first atom's radial force. + + Parameters + ---------- + seps + Sample separations. + forces + Cartesian forces for both atoms. + threshold + Force magnitudes below this value are treated as zero. + + Returns + ------- + float + Number of radial-force direction changes. + """ + radial_forces = _radial_forces(seps, forces).copy() + radial_forces[np.abs(radial_forces) < threshold] = 0 + force_signs = np.sign(radial_forces[radial_forces != 0]) + return float(np.sum(np.diff(force_signs) != 0)) + + +def calc_force_total_variation( + seps: ArrayLike, + forces: np.ndarray, +) -> float: + """ + Calculate total variation in the first atom's radial force. + + Parameters + ---------- + seps + Sample separations. + forces + Cartesian forces for both atoms. + + Returns + ------- + float + Total radial-force variation. + """ + return float(np.sum(np.abs(np.diff(_radial_forces(seps, forces))))) + + +def calc_force_jump( + seps: ArrayLike, + forces: np.ndarray, +) -> float: + """ + Calculate total radial-force step magnitude around sign-flip points. + + Parameters + ---------- + seps + Sample separations. + forces + Cartesian forces for both atoms. + + Returns + ------- + float + Total radial-force step magnitude. + """ + return _jump_magnitude(_radial_forces(seps, forces), threshold=0) diff --git a/ml_peg/analysis/physicality/diatomics/metrics/schema.py b/ml_peg/analysis/physicality/diatomics/metrics/schema.py new file mode 100644 index 000000000..c4b74e7a3 --- /dev/null +++ b/ml_peg/analysis/physicality/diatomics/metrics/schema.py @@ -0,0 +1,397 @@ +"""Typed diatomic curve schema and local data adapters.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +import gzip +import json +import os +from typing import Any + +import numpy as np +from numpy.typing import ArrayLike +import pandas as pd + +StrPath = str | os.PathLike[str] +DEFAULT_DFT_REFERENCE_PATH = ( + f"{os.path.dirname(os.path.dirname(__file__))}/data/diatomics-dft.json.gz" +) + + +def homo_key(formula: str) -> str: + """ + Collapse a homonuclear pair label such as ``H-H`` to its element key. + + Parameters + ---------- + formula + Element or pair label. + + Returns + ------- + str + Element key for homonuclear pairs, otherwise the original label. + """ + element_1, separator, element_2 = formula.partition("-") + return element_1 if separator and element_1 == element_2 else formula + + +class DiatomicCurve: + """ + Store one validated diatomic energy and Cartesian-force curve. + + Parameters + ---------- + distances + Sample separations. + energies + Sample energies. + forces + Cartesian forces for both atoms. + """ + + distances: np.ndarray + energies: np.ndarray + forces: np.ndarray + + def __init__( + self, + distances: ArrayLike, + energies: ArrayLike, + forces: ArrayLike, + ) -> None: + """ + Convert curve data to arrays and validate shapes and sample counts. + + Parameters + ---------- + distances + Sample separations. + energies + Sample energies. + forces + Cartesian forces for both atoms. + """ + self.distances = np.asarray(distances) + self.energies = np.asarray(energies) + self.forces = np.asarray(forces) + + for name, values in ( + ("distances", self.distances), + ("energies", self.energies), + ): + if values.ndim != 1: + raise ValueError(f"{name} must have shape (n,), got {values.shape}") + + n_distances = len(self.distances) + if (n_energies := len(self.energies)) != n_distances: + raise ValueError( + f"distance and energy counts differ: {n_distances} != {n_energies}" + ) + + if (n_forces := len(self.forces)) != n_distances: + raise ValueError( + f"distance and force counts differ: {n_distances} != {n_forces}" + ) + expected_force_shape = (n_distances, 2, 3) + if self.forces.shape != expected_force_shape: + raise ValueError( + "forces must have shape " + f"{expected_force_shape}, got {self.forces.shape}" + ) + + +@dataclass +class DiatomicCurves: + """Store homo- and heteronuclear curves with their shared or union grid.""" + + distances: np.ndarray + homo_nuclear: dict[str, DiatomicCurve] + hetero_nuclear: dict[str, DiatomicCurve] = field(default_factory=dict) + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> DiatomicCurves: + """ + Parse MBD JSON curves, requiring per-curve grids to be ordered subsets. + + Parameters + ---------- + data + Decoded MBD curve payload. + + Returns + ------- + DiatomicCurves + Validated homo- and heteronuclear curves. + """ + distances = np.asarray(data["distances"]) + grid_position_by_distance = { + float(distance): index for index, distance in enumerate(distances) + } + + def make_curves(section: str) -> dict[str, DiatomicCurve]: + """ + Convert one MBD JSON section to typed curves. + + Parameters + ---------- + section + MBD JSON section name. + + Returns + ------- + dict[str, DiatomicCurve] + Typed curves keyed by normalized formula. + """ + raw_curves = data.get(section, {}) + key_function = homo_key if section.startswith("homo") else str + + def curve_distances( + formula: str, + curve: dict[str, Any], + ) -> np.ndarray: + """ + Return an ordered per-curve subset of the top-level grid. + + Parameters + ---------- + formula + Curve formula used in validation errors. + curve + Raw curve payload. + + Returns + ------- + np.ndarray + Ordered curve-specific distance grid. + """ + curve_distance_array = np.asarray(curve.get("distances", distances)) + # off-grid points map to -1; valid subsets have strictly + # increasing grid positions + grid_positions = np.array( + [ + grid_position_by_distance.get(float(distance), -1) + for distance in curve_distance_array + ] + ) + if (grid_positions < 0).any() or (np.diff(grid_positions) <= 0).any(): + raise ValueError( + f"{formula} curve distances must be an ordered subset " + "of top-level distances" + ) + return curve_distance_array + + return { + key_function(formula): DiatomicCurve( + distances=curve_distances(formula, curve), + energies=curve["energies"], + forces=curve.get("forces", []), + ) + for formula, curve in raw_curves.items() + if len(curve["energies"]) > 0 + } + + return cls( + distances=distances, + homo_nuclear=make_curves("homo-nuclear"), + hetero_nuclear=make_curves("hetero-nuclear"), + ) + + +def _load_json(path: StrPath) -> dict[str, Any]: + """ + Load a JSON or gzipped JSON object. + + Parameters + ---------- + path + JSON or gzipped JSON path. + + Returns + ------- + dict[str, Any] + Decoded JSON object. + """ + string_path = os.fspath(path) + open_function = gzip.open if string_path.endswith(".gz") else open + with open_function(string_path, mode="rt", encoding="utf-8") as file: + return json.load(file) + + +def load_mbd_json(path: StrPath) -> DiatomicCurves: + """ + Load MBD-format predicted curves from JSON or gzipped JSON. + + Parameters + ---------- + path + JSON or gzipped JSON path. + + Returns + ------- + DiatomicCurves + Validated predicted curves. + """ + return DiatomicCurves.from_dict(_load_json(path)) + + +def load_dft_reference_curves( + functional: str = "PBE", + ref_path: StrPath | None = None, +) -> DiatomicCurves: + """ + Load bundled or custom DFT reference curves for one functional. + + Parameters + ---------- + functional + Density functional key in the reference payload. + ref_path + Optional custom reference path. + + Returns + ------- + DiatomicCurves + DFT reference curves. + """ + reference_path = ref_path or DEFAULT_DFT_REFERENCE_PATH + references = _load_json(reference_path)[functional] + return DiatomicCurves( + distances=np.array([]), + homo_nuclear={ + homo_key(formula): DiatomicCurve( + distances=curve["distances"], + energies=curve["energies"], + forces=curve.get("forces", []), + ) + for formula, curve in references.items() + }, + ) + + +def _parse_pair_label(pair_label: str) -> tuple[str, str]: + """ + Parse an ``Element-Element`` pair label. + + Parameters + ---------- + pair_label + Pair label to parse. + + Returns + ------- + tuple[str, str] + First and second element symbols. + """ + elements = pair_label.split("-") + if len(elements) != 2 or not all(elements): + raise ValueError( + f"pair labels must have form 'Element-Element', got {pair_label!r}" + ) + return elements[0], elements[1] + + +def curves_from_ml_peg_dataframe( + dataframe: pd.DataFrame, + *, + include_heteronuclear: bool = True, +) -> DiatomicCurves: + """ + Convert an ml-peg dataframe to x-aligned two-atom force curves. + + Parameters + ---------- + dataframe + Diatomic samples with pair, distance, energy, and projected force columns. + include_heteronuclear + Whether to include heteronuclear pairs. + + Returns + ------- + DiatomicCurves + Converted homo- and heteronuclear curves. + """ + required_columns = {"pair", "distance", "energy", "force_parallel"} + missing_columns = required_columns - set(dataframe.columns) + if missing_columns: + raise ValueError(f"Missing ml-peg diatomics columns: {sorted(missing_columns)}") + + homo_nuclear: dict[str, DiatomicCurve] = {} + hetero_nuclear: dict[str, DiatomicCurve] = {} + for pair_label, pair_dataframe in dataframe.groupby( + "pair", sort=False, dropna=False + ): + string_pair_label = str(pair_label) + element_1, element_2 = _parse_pair_label(string_pair_label) + if element_1 != element_2 and not include_heteronuclear: + continue + sorted_dataframe = pair_dataframe.sort_values("distance") + duplicate_rows = sorted_dataframe[ + sorted_dataframe.duplicated("distance", keep=False) + ] + for distance, duplicate_samples in duplicate_rows.groupby("distance"): + unique_values = duplicate_samples[ + ["energy", "force_parallel"] + ].drop_duplicates() + if len(unique_values) > 1: + raise ValueError( + f"{string_pair_label} has conflicting samples at " + f"distance={distance}" + ) + if not duplicate_rows.empty: + duplicate_distances = duplicate_rows["distance"].unique().tolist() + raise ValueError( + f"{string_pair_label} has duplicate distance values: " + f"{duplicate_distances!r}" + ) + distances = sorted_dataframe["distance"].to_numpy(dtype=float) + energies = sorted_dataframe["energy"].to_numpy(dtype=float) + projected_forces = sorted_dataframe["force_parallel"].to_numpy(dtype=float) + forces = np.zeros((len(distances), 2, 3), dtype=float) + forces[:, 0, 0] = -projected_forces + forces[:, 1, 0] = projected_forces + curve = DiatomicCurve(distances=distances, energies=energies, forces=forces) + if element_1 == element_2: + homo_nuclear[element_1] = curve + else: + hetero_nuclear[string_pair_label] = curve + + included_curves = [*homo_nuclear.values(), *hetero_nuclear.values()] + all_distances = ( + np.concatenate([curve.distances for curve in included_curves]) + if included_curves + else np.array([], dtype=float) + ) + return DiatomicCurves( + distances=np.sort(np.unique(all_distances)), + homo_nuclear=homo_nuclear, + hetero_nuclear=hetero_nuclear, + ) + + +def load_ml_peg_curves( + source: pd.DataFrame | StrPath, + *, + include_heteronuclear: bool = True, +) -> DiatomicCurves: + """ + Load current ml-peg diatomic curves from a dataframe or CSV. + + Parameters + ---------- + source + Diatomic dataframe or CSV path. + include_heteronuclear + Whether to include heteronuclear pairs. + + Returns + ------- + DiatomicCurves + Converted homo- and heteronuclear curves. + """ + dataframe = ( + source if isinstance(source, pd.DataFrame) else pd.read_csv(os.fspath(source)) + ) + return curves_from_ml_peg_dataframe( + dataframe, include_heteronuclear=include_heteronuclear + ) diff --git a/ml_peg/app/utils/utils.py b/ml_peg/app/utils/utils.py index c526db085..e53b6ef3d 100644 --- a/ml_peg/app/utils/utils.py +++ b/ml_peg/app/utils/utils.py @@ -8,11 +8,12 @@ import json from numbers import Number from pathlib import Path -from typing import Any, NotRequired, TypedDict +from typing import Any, TypedDict import dash.dash_table.Format as TableFormat from matplotlib import colormaps import numpy as np +from typing_extensions import NotRequired import yaml from ml_peg.models import MODELS_ROOT diff --git a/pyproject.toml b/pyproject.toml index ec7590a04..4a82eeb36 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,6 +39,7 @@ dependencies = [ "scikit-learn<2,>=1.7.2", "tqdm<5,>=4.67.3", "typer<1.0.0,>=0.19.1", + "typing-extensions>=4.0", ] [project.optional-dependencies] diff --git a/tests/metrics/diatomics/test_analysis.py b/tests/metrics/diatomics/test_analysis.py new file mode 100644 index 000000000..5467ba804 --- /dev/null +++ b/tests/metrics/diatomics/test_analysis.py @@ -0,0 +1,136 @@ +"""Integration tests for ml-peg diatomic analysis entry points.""" + +from __future__ import annotations + +import gzip +import json +from pathlib import Path + +import numpy as np +import pandas as pd +import pytest + +from ml_peg.analysis.physicality.diatomics.analyse_diatomics import ( + DEFAULT_THRESHOLDS, + aggregate_model_metrics, + collect_metrics, + evaluate_mbd_diatomic_metrics, + write_mbd_diatomic_metrics, +) +from ml_peg.analysis.physicality.diatomics.metrics import DIATOMIC_METRIC_KEYS + +pytestmark = pytest.mark.framework("matbench-discovery") + + +def _integration_dataframe() -> pd.DataFrame: + """Build smooth homo- and heteronuclear current-format curves.""" + distances = np.linspace(0.3, 3.0, 20) + energies = 100 * (distances - 1.0) ** 2 - 2 + force_parallel = -200 * (distances - 1.0) + dataframe = pd.DataFrame( + { + "element_1": "H", + "distance": distances, + "energy": energies, + "force_parallel": force_parallel, + } + ) + return pd.concat( + [ + dataframe.assign(pair=pair, element_2=element_2) + for pair, element_2 in (("H-H", "H"), ("H-He", "He")) + ], + ignore_index=True, + ) + + +def _write_reference(path: Path) -> None: + """Write a matching gzipped PBE curve for integration tests.""" + dataframe = _integration_dataframe().query("pair == 'H-H'") + forces = np.zeros((len(dataframe), 2, 3)) + forces[:, 0, 0] = -dataframe["force_parallel"] + forces[:, 1, 0] = dataframe["force_parallel"] + curve = { + "distances": dataframe["distance"].tolist(), + "energies": dataframe["energy"].tolist(), + "forces": forces.tolist(), + } + with gzip.open(path, mode="wt", encoding="utf-8") as file: + json.dump({"PBE": {"H-H": curve}}, file) + + +def test_separate_mbd_evaluation_is_json_safe_and_homonuclear( + tmp_path: Path, +) -> None: + """Separate result reports 12 homonuclear metrics without legacy weighting.""" + reference_path = tmp_path / "reference.json.gz" + _write_reference(reference_path) + pair_data = {"test-model": _integration_dataframe()} + result = evaluate_mbd_diatomic_metrics( + pair_data, + reference_path=reference_path, + interpolate=200, + ) + + assert result["schema_version"] == 1 + assert result["source"] == { + "framework": "matbench-discovery", + "version": "1.3.1", + } + assert result["weighted_in_legacy_score"] is False + assert set(result["metric_names"]) == DIATOMIC_METRIC_KEYS + model_results = result["models"] + assert isinstance(model_results, dict) + assert set(model_results["test-model"]["elements"]) == {"H"} + assert set(model_results["test-model"]["means"]) == DIATOMIC_METRIC_KEYS + json.dumps(result, allow_nan=False) + + output_path = tmp_path / "mbd-metrics.json" + written = write_mbd_diatomic_metrics( + output_path, + pair_data, + reference_path=reference_path, + ) + assert json.loads(output_path.read_text(encoding="utf-8")) == written + + +def test_mbd_evaluation_uses_pair_labels_for_homonuclear_filtering( + tmp_path: Path, +) -> None: + """Ignore inconsistent element columns and classify pairs from their labels.""" + reference_path = tmp_path / "reference.json.gz" + _write_reference(reference_path) + dataframe = _integration_dataframe() + dataframe.loc[dataframe["pair"] == "H-H", "element_2"] = "He" + + result = evaluate_mbd_diatomic_metrics( + {"test-model": dataframe}, reference_path=reference_path + ) + + assert set(result["models"]["test-model"]["elements"]) == {"H"} + + +def test_legacy_metric_regression_remains_unchanged() -> None: + """Legacy five-metric homo-plus-hetero aggregation retains pinned values.""" + pair_dataframe = pd.DataFrame( + { + "pair": ["H-H"] * 5 + ["H-He"] * 5, + "distance": [1, 2, 3, 4, 5] * 2, + "energy": [4, 1, 0, 1, 4] * 2, + "force_parallel": [-2, -1, 0, 1, 2] * 2, + } + ) + expected = { + "Force flips": 1.0, + "Energy minima": 1.0, + "Energy inflections": 0.0, + "ρ(E, repulsion)": -1.0, + "ρ(E, attraction)": 1.0, + } + + assert list(DEFAULT_THRESHOLDS) == list(expected) + assert aggregate_model_metrics(pair_dataframe) == pytest.approx(expected) + collected = collect_metrics({"test-model": pair_dataframe}) + collected_record = collected.to_dict(orient="records")[0] + assert collected_record.pop("Model") == "test-model" + assert collected_record == pytest.approx(expected) diff --git a/tests/metrics/diatomics/test_energy_force.py b/tests/metrics/diatomics/test_energy_force.py new file mode 100644 index 000000000..2fa277e02 --- /dev/null +++ b/tests/metrics/diatomics/test_energy_force.py @@ -0,0 +1,240 @@ +"""Tests for imported diatomic energy and force formulas.""" + +from __future__ import annotations + +from collections.abc import Callable +import re + +import numpy as np +import pytest + +from ml_peg.analysis.physicality.diatomics.metrics.energy import ( + calc_energy_diff_flips, + calc_energy_jump, + calc_pbe_bond_length_error, + calc_pbe_energy_mae, + calc_pbe_vib_freq_error, + calc_pbe_wall_dist_mae, + calc_pbe_well_depth_error, + calc_tortuosity, +) +from ml_peg.analysis.physicality.diatomics.metrics.force import ( + calc_force_flips, + calc_force_jump, + calc_force_mae, + calc_force_total_variation, +) + +pytestmark = pytest.mark.framework("matbench-discovery") + +_LENGTH_ERROR = re.escape("len(separation_array)=2 != len(value_array)=3") +_ENERGY_SHAPE_ERROR = re.escape("energy values must have shape (n,)") +_FORCE_SHAPE_ERROR = re.escape("force values must have shape (n, 2, 3)") + + +def _radial_forces(radial_values: np.ndarray) -> np.ndarray: + """Return equal-and-opposite two-atom forces from radial values.""" + forces = np.zeros((len(radial_values), 2, 3)) + forces[:, 0, 0] = radial_values + forces[:, 1, 0] = -radial_values + return forces + + +@pytest.mark.parametrize( + ("energies", "expected_flips", "expected_jump"), + [ + (np.array([1.0, 2.0, 3.0, 4.0, 5.0]), 0, 0.0), + (np.array([1.0, 3.0, 2.0, 4.0, 5.0]), 2, 6.0), + (np.array([0.0, 2.0, 1.0, 3.0, 0.5]), 3, 10.5), + ], +) +def test_energy_flip_and_jump_formulas( + energies: np.ndarray, + expected_flips: int, + expected_jump: float, +) -> None: + """Energy flips and jumps match hand-computed values.""" + separations = np.arange(1, len(energies) + 1, dtype=float) + assert calc_energy_diff_flips(separations, energies) == expected_flips + assert calc_energy_jump(separations, energies) == pytest.approx(expected_jump) + + +@pytest.mark.parametrize( + ("energies", "expected"), + [ + (np.arange(1, 6, dtype=float), 1.0), + (np.arange(1, 6, dtype=float) ** 2, 1.0), + (np.ones(5), np.nan), + ], +) +def test_tortuosity_formulas(energies: np.ndarray, expected: float) -> None: + """Tortuosity is one for monotone curves and NaN for flat curves.""" + result = calc_tortuosity(np.arange(1, 6), energies) + assert result == pytest.approx(expected, nan_ok=True) + + +def test_source_keyword_argument_names_are_supported() -> None: + """Public metric functions retain Matbench Discovery keyword names.""" + separations = np.arange(1, 6, dtype=float) + energies = np.arange(1, 6, dtype=float) + forces = np.zeros((5, 2, 3)) + assert calc_tortuosity(seps=separations, energies=energies) == pytest.approx(1) + assert calc_force_total_variation(seps=separations, forces=forces) == 0 + energy_kwargs = { + "seps_ref": separations, + "energy_ref": energies, + "seps_pred": separations, + "energy_pred": energies, + } + force_kwargs = { + "seps_ref": separations, + "f_ref": forces, + "seps_pred": separations, + "f_pred": forces, + } + assert calc_pbe_energy_mae(**energy_kwargs) == 0 + assert calc_force_mae(**force_kwargs) == 0 + + +def test_pbe_reference_energy_formulas() -> None: + """PBE-relative metrics match analytic parabolic-well expectations.""" + reference_equilibrium = 1.5 + predicted_equilibrium = 1.6 + predicted_curvature_factor = 1.2 + reference_max = 3.0 + predicted_max = 3.05 + separations_ref = np.linspace(0.5, reference_max, 51) + separations_pred = np.linspace(0.55, predicted_max, 51) + energy_ref = (separations_ref - reference_equilibrium) ** 2 - 2 + energy_pred = ( + predicted_curvature_factor * (separations_pred - predicted_equilibrium) ** 2 + - 1.7 + ) + curve_args = (separations_ref, energy_ref, separations_pred, energy_pred) + expected_depth_error = abs( + predicted_curvature_factor * (predicted_max - predicted_equilibrium) ** 2 + - (reference_max - reference_equilibrium) ** 2 + ) + results = [ + (calc_pbe_wall_dist_mae(*curve_args, thresholds_ev=(1,)), 0.187, 0.01), + (calc_pbe_energy_mae(*curve_args, interpolate=200), 0.118, 0.01), + ( + calc_pbe_bond_length_error(*curve_args), + predicted_equilibrium - reference_equilibrium, + None, + ), + (calc_pbe_well_depth_error(*curve_args), expected_depth_error, 0.01), + (calc_pbe_vib_freq_error("H", *curve_args), 99.1, 1), + ] + for actual, expected, absolute_tolerance in results: + assert actual == pytest.approx(expected, abs=absolute_tolerance) + + +def test_force_metric_formulas() -> None: + """Force metrics match hand-computed radial-force values.""" + separations = np.arange(1, 6, dtype=float) + forces = _radial_forces(np.array([1.0, 2.0, -1.0, 3.0, -2.0])) + assert calc_force_flips(separations, forces) == 3 + assert calc_force_total_variation(separations, forces) == pytest.approx(13) + assert calc_force_jump(separations, forces) == pytest.approx(20) + assert calc_force_mae(separations, forces, separations, forces) == 0 + + +def test_energy_and_force_interpolation_use_shared_range() -> None: + """Interpolated MAEs compare linear curves only on their overlap.""" + separations_ref = np.array([2.0, 3.0, 4.0]) + separations_pred = np.array([2.1, 3.1, 4.1]) + curve_args = (separations_ref, separations_ref, separations_pred, separations_pred) + assert calc_pbe_energy_mae( + *curve_args, + interpolate=20, + ) == pytest.approx(0) + assert calc_force_mae( + separations_ref, + _radial_forces(separations_ref), + separations_pred, + _radial_forces(separations_pred), + interpolate=20, + ) == pytest.approx(0) + + +@pytest.mark.parametrize( + ("metric_function", "separations", "values", "error_match"), + [ + ( + calc_energy_jump, + np.array([1.0, 1.0, 2.0]), + np.arange(3.0), + "contains 1 duplicates", + ), + ( + calc_energy_jump, + np.array([1.0, np.nan, 3.0]), + np.arange(3.0), + "Input contains NaN", + ), + (calc_energy_jump, np.arange(2.0), np.arange(3.0), _LENGTH_ERROR), + (calc_energy_jump, np.arange(3.0), np.zeros((3, 1)), _ENERGY_SHAPE_ERROR), + ( + calc_force_total_variation, + np.arange(3.0), + np.zeros((3, 1, 3)), + _FORCE_SHAPE_ERROR, + ), + ], +) +def test_curve_validation_errors( + metric_function: Callable[[np.ndarray, np.ndarray], float], + separations: np.ndarray, + values: np.ndarray, + error_match: str, +) -> None: + """Curve formulas reject duplicate, non-finite, and malformed inputs.""" + with pytest.raises(ValueError, match=error_match): + metric_function(separations, values) + + +@pytest.mark.parametrize( + "separations_pred", + [np.array([4.0, 5.0, 6.0]), np.array([3.0, 4.0, 5.0])], + ids=["disjoint", "single-shared-point"], +) +def test_force_interpolation_rejects_unusable_overlap( + separations_pred: np.ndarray, +) -> None: + """Force interpolation rejects disjoint and point-only overlap.""" + separations_ref = np.array([1.0, 2.0, 3.0]) + forces = np.zeros((3, 2, 3)) + with pytest.raises(ValueError, match="no overlap"): + calc_force_mae( + separations_ref, forces, separations_pred, forces, interpolate=True + ) + + +@pytest.mark.parametrize( + ("metric_function", "reference_values", "predicted_values"), + [ + ( + calc_pbe_energy_mae, + np.array([0.0, 1.0, 2.0]), + np.array([2.0, 1.0, 0.0]), + ), + (calc_force_mae, np.zeros((3, 2, 3)), np.ones((3, 2, 3))), + ], +) +def test_interpolation_requires_at_least_two_points( + metric_function: Callable[..., float], + reference_values: np.ndarray, + predicted_values: np.ndarray, +) -> None: + """Reject one-point interpolation, which erases far-field energy errors.""" + separations_ref = np.array([1.0, 2.0, 3.0]) + separations_pred = np.array([1.1, 2.1, 3.1]) + with pytest.raises(ValueError, match="at least 2 points"): + metric_function( + separations_ref, + reference_values, + separations_pred, + predicted_values, + interpolate=1, + ) diff --git a/tests/metrics/diatomics/test_metrics.py b/tests/metrics/diatomics/test_metrics.py new file mode 100644 index 000000000..940817ce9 --- /dev/null +++ b/tests/metrics/diatomics/test_metrics.py @@ -0,0 +1,265 @@ +"""Tests for diatomic metric orchestration and aggregation.""" + +from __future__ import annotations + +import hashlib + +import numpy as np +import pytest + +from ml_peg.analysis.physicality.diatomics import metrics +from ml_peg.analysis.physicality.diatomics.metrics import ( + DEFAULT_DFT_REFERENCE_PATH, + DIATOMIC_METRIC_KEYS, + ENERGY_JUMP, + PBE_ENERGY_MAE, + PBE_FORCE_MAE, + TORTUOSITY, + DiatomicCurve, + DiatomicCurves, + aggregate_finite_means, + calc_diatomic_metrics, + eval_window, + find_low_quality_dft_refs, + load_dft_reference_curves, +) + +pytestmark = pytest.mark.framework("matbench-discovery") + + +def _forces_from_energy( + distances: np.ndarray, + energies: np.ndarray, +) -> np.ndarray: + """Construct equal-and-opposite radial forces.""" + forces = np.zeros((len(distances), 2, 3)) + forces[:, 0, 0] = -np.gradient(energies, distances) + forces[:, 1, 0] = -forces[:, 0, 0] + return forces + + +def _make_curves( + curves_by_element: dict[str, np.ndarray], + distances: np.ndarray, +) -> DiatomicCurves: + """Wrap element energy arrays as homonuclear curves.""" + return DiatomicCurves( + distances=distances, + homo_nuclear={ + element_symbol: DiatomicCurve( + distances, energies, _forces_from_energy(distances, energies) + ) + for element_symbol, energies in curves_by_element.items() + }, + ) + + +@pytest.mark.parametrize("prediction_key", ["H", "H-H"]) +def test_all_metrics_and_reference_key_normalization(prediction_key: str) -> None: + """Matching element and pair keys both produce all 12 finite metrics.""" + distances = np.linspace(0.2, 3.0, 50) + energies = 100 * (distances - 1.0) ** 2 - 2 + references = _make_curves({"H": energies}, distances) + predictions = _make_curves({"H": energies}, distances) + if prediction_key == "H-H": + predictions.homo_nuclear[prediction_key] = predictions.homo_nuclear.pop("H") + result = calc_diatomic_metrics(references, predictions, interpolate=200) + assert set(result["H"]) == DIATOMIC_METRIC_KEYS + assert result["H"][PBE_ENERGY_MAE] == pytest.approx(0) + for metric_value in result["H"].values(): + assert np.isfinite(metric_value) + + +def test_duplicate_normalized_element_keys_are_rejected() -> None: + """Reject separate element and pair keys for the same homonuclear curve.""" + distances = np.linspace(0.2, 3.0, 50) + energies = (distances - 1.0) ** 2 + curves = _make_curves({"H": energies}, distances) + curves.homo_nuclear["H-H"] = curves.homo_nuclear["H"] + + with pytest.raises(ValueError, match="Duplicate homonuclear curve"): + calc_diatomic_metrics(None, curves) + + +def test_homonuclear_only_and_non_mp_filtering() -> None: + """Imported metrics ignore heteronuclear curves and unsupported MP elements.""" + distances = np.linspace(0.3, 3.0, 20) + energies = (distances - 1.0) ** 2 + homo_curves = _make_curves( + {"H": energies, "Po": energies, "Og": energies}, distances + ) + homo_curves.hetero_nuclear["H-He"] = DiatomicCurve( + distances, energies, _forces_from_energy(distances, energies) + ) + + result = calc_diatomic_metrics(None, homo_curves) + + assert set(result) == {"H"} + assert TORTUOSITY in result["H"] + + +def test_nonfinite_repulsive_wall_sample_skips_element() -> None: + """Non-finite values in the wider wall window exclude the whole curve.""" + distances = np.linspace(0.5, 3.0, 101) + energies = (distances - 1.5) ** 2 + curves = _make_curves({"C": energies}, distances) + wall_only_index = int(np.flatnonzero((distances >= 0.608) & (distances < 0.684))[0]) + curves.homo_nuclear["C"].energies[wall_only_index] = np.nan + + assert calc_diatomic_metrics(None, curves) == {} + + +def test_low_quality_reference_gate_and_nonfinite_predictions() -> None: + """Jumpy refs lose PBE metrics while non-finite predictions are skipped.""" + distances = np.linspace(0.3, 6.0, 40) + smooth = (distances - 1.5) ** 2 + jumpy = smooth + 5 * (-1) ** np.arange(len(distances)) + nonfinite = smooth.copy() + nonfinite[10] = np.nan + reference_curves = _make_curves( + {"H": smooth, "Ho": jumpy, "Er": np.full(len(distances), np.nan)}, + distances, + ) + predicted_curves = _make_curves( + {"H": smooth, "Ho": smooth, "Er": smooth, "He": nonfinite}, + distances, + ) + + assert find_low_quality_dft_refs(reference_curves) == {"Ho", "Er"} + result = calc_diatomic_metrics(reference_curves, predicted_curves) + assert set(result) == {"H", "Ho", "Er"} + assert PBE_ENERGY_MAE in result["H"] + for gated_element in ("Ho", "Er"): + assert PBE_ENERGY_MAE not in result[gated_element] + assert TORTUOSITY in result[gated_element] + + +def test_missing_forces_raise_clear_error() -> None: + """A scored prediction missing force samples is rejected.""" + distances = np.linspace(0.3, 3.0, 10) + energies = (distances - 1.0) ** 2 + predicted_curves = _make_curves({"H": energies}, distances) + predicted_curves.homo_nuclear["H"].forces = np.array([]) + + with pytest.raises(ValueError, match="H diatomic curve is missing forces"): + calc_diatomic_metrics(None, predicted_curves) + + +def test_full_pipeline_interpolation() -> None: + """Full relative metrics require matching grids unless interpolation is enabled.""" + reference_distances = np.linspace(0.3, 3.0, 20) + predicted_distances = reference_distances * 1.001 + reference_energies = (reference_distances - 1.0) ** 2 + predicted_energies = (predicted_distances - 1.0) ** 2 + reference_curves = _make_curves({"H": reference_energies}, reference_distances) + predicted_curves = _make_curves({"H": predicted_energies}, predicted_distances) + curve_pair = (reference_curves, predicted_curves) + + with pytest.raises(ValueError, match="distances must be same"): + calc_diatomic_metrics(*curve_pair, interpolate=False) + result = calc_diatomic_metrics(*curve_pair, interpolate=200) + assert set(result["H"]) == DIATOMIC_METRIC_KEYS + + +def test_force_interpolation_omits_point_only_overlap() -> None: + """Omit force MAE when reference and prediction ranges only touch.""" + reference_distances = np.linspace(0.4, 1.0, 5) + predicted_distances = np.linspace(1.0, 3.0, 5) + reference_curves = _make_curves({"H": reference_distances**2}, reference_distances) + predicted_curves = _make_curves({"H": predicted_distances**2}, predicted_distances) + result = calc_diatomic_metrics( + reference_curves, + predicted_curves, + metrics={PBE_FORCE_MAE: {"interpolate": True}}, + ) + assert result["H"] == {} + + +def test_eval_window_and_repulsive_exclusion() -> None: + """Element windows use physical radii and exclude deep-overlap spikes.""" + radius_min, radius_max = eval_window("H-H", 2.5) + atomic_number = metrics.atomic_numbers["H"] + assert radius_min == pytest.approx(0.9 * metrics.covalent_radii[atomic_number]) + assert radius_max == pytest.approx(2.5) + + distances = np.linspace(0.1, 6.0, 60) + energies = np.exp(-distances) + energies[distances < 0.2] = 1e6 + result = calc_diatomic_metrics( + None, + _make_curves({"H": energies}, distances), + ) + assert result["H"][ENERGY_JUMP] == pytest.approx(0) + + source_keyword_window = eval_window( + elem_symbol="H-H", seps_max=2.5, r_min_factor=0.8 + ) + assert source_keyword_window[0] == pytest.approx( + 0.8 * metrics.covalent_radii[atomic_number] + ) + + +def test_bundled_pbe_quality_gate_regression() -> None: + """Bundled PBE data retains the eight known jumpy lanthanide references.""" + reference_curves = load_dft_reference_curves() + assert find_low_quality_dft_refs(reference_curves) == { + "Pr", + "Pm", + "Sm", + "Tb", + "Dy", + "Ho", + "Er", + "Tm", + } + self_metrics = calc_diatomic_metrics( + reference_curves, reference_curves, interpolate=200 + ) + assert len(self_metrics) == 87 + mean_metrics = aggregate_finite_means(self_metrics) + assert mean_metrics == { + "tortuosity": 1.043, + "force_flips": 1.632, + "energy_jump": 6.28, + "energy_diff_flips": 2.575, + "force_total_variation": 193.5, + "force_jump": 7.164, + "pbe_wall_dist_mae": 0.0, + "pbe_energy_mae": 0.0, + "pbe_bond_length_error": 0.0, + "pbe_well_depth_error": 0.0, + "pbe_force_mae": 0.0, + "pbe_vib_freq_error": 0.0, + } + + +def test_bundled_pbe_reference_hash() -> None: + """Pin the bundled DFT reference file.""" + with open(DEFAULT_DFT_REFERENCE_PATH, "rb") as file: + digest = hashlib.sha256(file.read()).hexdigest() + + assert digest == "1fe6334a82e98208ea74169a3beaf98cd5188bdc7ac40e518697fd36c7196e3d" + + +def test_finite_mean_aggregation() -> None: + """Aggregation unions metric keys and ignores missing or non-finite values.""" + metrics_by_element = { + "H": {TORTUOSITY: 1.0, ENERGY_JUMP: 2.0}, + "He": {TORTUOSITY: np.nan, ENERGY_JUMP: 4.0}, + "Li": {TORTUOSITY: np.inf}, + } + assert aggregate_finite_means(metrics_by_element) == { + TORTUOSITY: 1.0, + ENERGY_JUMP: 3.0, + } + assert aggregate_finite_means( + {"H": {TORTUOSITY: 1e308}, "He": {TORTUOSITY: 1e308}} + ) == {TORTUOSITY: 1e308} + + +def test_unknown_metric_rejected() -> None: + """The orchestrator rejects names outside the stable 12-key set.""" + distances = np.linspace(0.3, 3.0, 10) + curves = _make_curves({"H": distances**2}, distances) + with pytest.raises(ValueError, match="unknown_metrics"): + calc_diatomic_metrics(None, curves, metrics={"not_a_metric": {}}) diff --git a/tests/metrics/diatomics/test_schema.py b/tests/metrics/diatomics/test_schema.py new file mode 100644 index 000000000..fa5032538 --- /dev/null +++ b/tests/metrics/diatomics/test_schema.py @@ -0,0 +1,201 @@ +"""Tests for diatomic schemas and local data adapters.""" + +from __future__ import annotations + +import gzip +import json +from pathlib import Path + +import numpy as np +import pandas as pd +import pytest + +from ml_peg.analysis.physicality.diatomics.metrics import ( + DiatomicCurve, + DiatomicCurves, + load_dft_reference_curves, + load_mbd_json, + load_ml_peg_curves, +) + +pytestmark = pytest.mark.framework("matbench-discovery") + + +def _curve_payload() -> dict[str, object]: + """Return a minimal two-point MBD curve payload.""" + return { + "energies": [0.2, 0.0], + "forces": [ + [[0.1, 0, 0], [-0.1, 0, 0]], + [[0.0, 0, 0], [0.0, 0, 0]], + ], + } + + +def _ml_peg_dataframe( + *rows: tuple[object, float, float, float], +) -> pd.DataFrame: + """Build an ml-peg dataframe from pair, distance, energy, and force rows.""" + return pd.DataFrame(rows, columns=("pair", "distance", "energy", "force_parallel")) + + +def test_diatomic_classes_parse_arrays() -> None: + """Typed classes convert valid curve arrays.""" + distances = [1.0, 2.0] + energies = [0.1, 0.2] + forces = np.arange(12).reshape(2, 2, 3) + + curve = DiatomicCurve(distances, energies, forces) + assert curve.forces.shape == (2, 2, 3) + assert { + type(curve.distances), + type(curve.energies), + type(curve.forces), + } == {np.ndarray} + + payload = { + "distances": distances, + "homo-nuclear": {"H-H": _curve_payload()}, + "hetero-nuclear": {"H-He": _curve_payload()}, + } + curves = DiatomicCurves.from_dict(payload) + assert list(curves.homo_nuclear) == ["H"] + assert list(curves.hetero_nuclear) == ["H-He"] + np.testing.assert_array_equal(curves.homo_nuclear["H"].distances, distances) + + +@pytest.mark.parametrize( + ("override", "error_match"), + [ + ({"energies": [0.0]}, "distance and energy counts differ"), + ({"forces": np.zeros((1, 2, 3))}, "distance and force counts differ"), + ({"forces": np.zeros((2, 1, 3))}, "forces must have shape"), + ({"distances": [[1.0], [2.0]]}, "distances must have shape"), + ({"energies": [[0.0], [1.0]]}, "energies must have shape"), + ], +) +def test_diatomic_curve_rejects_invalid_shapes_and_counts( + override: dict[str, object], + error_match: str, +) -> None: + """DiatomicCurve rejects malformed arrays and sample-count mismatches.""" + arguments: dict[str, object] = { + "distances": [1.0, 2.0], + "energies": [0.0, 1.0], + "forces": np.zeros((2, 2, 3)), + } + with pytest.raises(ValueError, match=error_match): + DiatomicCurve(**(arguments | override)) + + +@pytest.mark.parametrize( + "bad_distances", + [[0.5, 1.5], [1.0, 1.0], [2.0, 1.0]], + ids=["off-grid", "duplicate", "reordered"], +) +def test_mbd_schema_rejects_invalid_curve_grids( + bad_distances: list[float], +) -> None: + """MBD curves must use ordered subsets of the top-level grid.""" + curve_payload = _curve_payload() | {"distances": bad_distances} + payload = { + "distances": [1.0, 2.0], + "homo-nuclear": {"H-H": curve_payload}, + } + with pytest.raises(ValueError, match="must be an ordered subset"): + DiatomicCurves.from_dict(payload) + + +def test_json_and_gzip_loaders(tmp_path: Path) -> None: + """MBD JSON and gzipped DFT references load into the typed schema.""" + prediction_path = tmp_path / "predictions.json" + prediction_path.write_text( + json.dumps( + { + "distances": [0.7, 1.0], + "homo-nuclear": {"H-H": _curve_payload()}, + } + ), + encoding="utf-8", + ) + prediction_curves = load_mbd_json(prediction_path) + assert list(prediction_curves.homo_nuclear) == ["H"] + + reference_path = tmp_path / "reference.json.gz" + with gzip.open(reference_path, mode="wt", encoding="utf-8") as file: + json.dump({"PBE": {"H-H": _curve_payload() | {"distances": [0.7, 1.0]}}}, file) + reference_curves = load_dft_reference_curves(ref_path=reference_path) + np.testing.assert_array_equal( + reference_curves.homo_nuclear["H"].energies, + [0.2, 0.0], + ) + + +def test_ml_peg_dataframe_and_csv_force_adapter(tmp_path: Path) -> None: + """CSV adapter reconstructs x-axis forces with the expected atom signs.""" + dataframe = _ml_peg_dataframe( + ("H-H", 2, 0, 2.5), + ("H-H", 1, 1, -1.5), + ("H-He", 1, 2, 4), + ("H-He", 2, 1, -3), + ) + csv_path = tmp_path / "diatomics.csv" + dataframe.to_csv(csv_path, index=False) + + for source in (dataframe, csv_path): + curves = load_ml_peg_curves(source) + assert list(curves.homo_nuclear) == ["H"] + assert list(curves.hetero_nuclear) == ["H-He"] + h_curve = curves.homo_nuclear["H"] + np.testing.assert_array_equal(h_curve.distances, [1.0, 2.0]) + np.testing.assert_array_equal(h_curve.forces[:, 0, 0], [1.5, -2.5]) + np.testing.assert_array_equal(h_curve.forces[:, 1, 0], [-1.5, 2.5]) + np.testing.assert_array_equal(h_curve.forces[:, :, 1:], 0) + + homonuclear_only = load_ml_peg_curves(dataframe, include_heteronuclear=False) + assert list(homonuclear_only.homo_nuclear) == ["H"] + assert homonuclear_only.hetero_nuclear == {} + + +@pytest.mark.parametrize( + ("dataframe", "error_match"), + [ + (pd.DataFrame({"pair": ["H-H"]}), "Missing ml-peg diatomics columns"), + (_ml_peg_dataframe(("H2", 1, 0, 0)), "pair labels must have form"), + ( + _ml_peg_dataframe((None, 1, 0, 0), ("H-H", 2, 0, 0)), + "pair labels must have form", + ), + ( + _ml_peg_dataframe((None, 1, 0, 0), (None, 2, 0, 0)), + "pair labels must have form", + ), + ], +) +def test_ml_peg_adapter_rejects_schema_errors( + dataframe: pd.DataFrame, + error_match: str, +) -> None: + """ml-peg adapter reports missing columns and malformed pair labels.""" + with pytest.raises(ValueError, match=error_match): + load_ml_peg_curves(dataframe) + + +@pytest.mark.parametrize( + ("energies", "error_match"), + [ + ([0.0, 0.0], "duplicate distance values"), + ([0.0, 1.0], "conflicting samples"), + ], + ids=["identical", "conflicting"], +) +def test_ml_peg_adapter_rejects_duplicate_distances( + energies: list[float], error_match: str +) -> None: + """Reject duplicate distances instead of retaining an arbitrary sample.""" + dataframe = _ml_peg_dataframe( + ("H-H", 1, energies[0], 0), + ("H-H", 1, energies[1], 0), + ) + with pytest.raises(ValueError, match=error_match): + load_ml_peg_curves(dataframe) diff --git a/uv.lock b/uv.lock index 6f679f234..7cf5f756b 100644 --- a/uv.lock +++ b/uv.lock @@ -5667,6 +5667,7 @@ dependencies = [ { name = "scipy", version = "1.16.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' or (extra == 'extra-6-ml-peg-grace' and extra == 'extra-6-ml-peg-mattersim') or (extra == 'extra-6-ml-peg-grace' and extra == 'extra-6-ml-peg-uma') or (extra == 'extra-6-ml-peg-mace' and extra == 'extra-6-ml-peg-mattersim') or (extra == 'extra-6-ml-peg-mace' and extra == 'extra-6-ml-peg-uma')" }, { name = "tqdm" }, { name = "typer" }, + { name = "typing-extensions" }, ] [package.optional-dependencies] @@ -5758,6 +5759,7 @@ requires-dist = [ { name = "torch-dftd", marker = "extra == 'd3'", specifier = "==0.5.1" }, { name = "tqdm", specifier = ">=4.67.3,<5" }, { name = "typer", specifier = ">=0.19.1,<1.0.0" }, + { name = "typing-extensions", specifier = ">=4.0" }, ] provides-extras = ["asemolec", "chgnet", "d3", "dpa3", "grace", "mace", "mattersim", "orb", "pet-mad", "uma"]