Add Matbench diatomic evaluation metrics

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

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,25 +387,25 @@ 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()
 
 
 @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

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.