Score diatomics on PBE bond length, wall, and force flips

.pre-commit-config.yaml

@@ -10,5 +10,5 @@ repos:
   rev: v0.16.9
   hooks:
     - id: ruff
-      args: ["--fix"]
+      args: ["--fix", "--select", "E4,E7,E9,F"]
     - id: ruff-format

lambench/metrics/utils.py

@@ -1,7 +1,5 @@
-import json
 from collections import defaultdict
 from datetime import datetime
-from functools import lru_cache
 from pathlib import Path
 from typing import Literal
 
@@ -11,14 +9,6 @@
 import lambench
 from lambench.workflow.entrypoint import gather_model, gather_model_params
 
-_DIATOMICS_JSON = (
-    Path(__file__).resolve().parent.parent
-    / "tasks"
-    / "calculator"
-    / "diatomics"
-    / "diatomics.json"
-)
-
 #############################
 # General utility functions #
 #############################
@@ -140,7 +130,7 @@ def aggregated_nve_md_results(results: dict[str, dict[str, float]]) -> dict[str,
 ## Inference efficiency utility functions
 def aggregated_inference_efficiency_results(
     results: dict[str, dict[str, float]],
-) -> dict[str, float]:
+) -> dict[str, float | None]:
     system_level_avg = []
     system_level_std = []
     system_level_success_rate = []
@@ -155,49 +145,106 @@ def aggregated_inference_efficiency_results(
     if success_count != len(results):
         return {"average_time": None, "std_time": None, "success_rate": 0.0}
     return {
-        "average_time": np.round(np.mean(system_level_avg), 6),
-        "standard_deviation": np.round(
-            np.sqrt(np.mean(np.square(system_level_std))), 6
+        "average_time": float(np.round(np.mean(system_level_avg), 6)),
+        "standard_deviation": float(
+            np.round(np.sqrt(np.mean(np.square(system_level_std))), 6)
         ),
-        "success_rate": np.round(np.mean(system_level_success_rate), 2),
+        "success_rate": float(np.round(np.mean(system_level_success_rate), 2)),
     }
 
 
-@lru_cache
-def _diatomics_molecule_names() -> tuple[str, ...]:
-    with open(_DIATOMICS_JSON) as fh:
-        return tuple(entry["name"] for entry in json.load(fh))
+## Diatomic utility functions
+def _empty_diatomics_agg() -> dict[str, float | None]:
+    return {
+        "bond_length_mae": None,
+        "well_depth_mae": None,
+        "avg_roughness": None,
+        "coverage": 0.0,
+        "score": None,
+    }
+
 
+def _mean_or_none(values: list[float]) -> float | None:
+    if not values:
+        return None
+    return float(np.mean(values))
 
-def _empty_diatomics_agg() -> dict[str, float | None]:
-    return {"avg_roughness": None}
 
+def _dummy_hat(value: float | None, dummy: float | None) -> float | None:
+    if (
+        value is None
+        or dummy is None
+        or not np.isfinite(value)
+        or not np.isfinite(dummy)
+        or dummy <= 0
+    ):
+        return None
+    return float(min(value / dummy, 1.0))
 
-def aggregated_diatomics_results(results: dict[str, dict]) -> dict[str, float]:
-    """
-    Aggregate per-molecule diatomics results.
 
-    avg_roughness: mean slope MAE relative to a constant dummy, capped at 1.
-    Requires a finite roughness for every molecule in diatomics.json; otherwise None.
+def _element_is_excluded(record: dict | None) -> bool:
+    if not isinstance(record, dict) or record.get("model_excluded", True):
+        return True
+    roughness = record.get("roughness")
+    return roughness is None or not np.isfinite(roughness) or not 0 <= roughness <= 1
+
+
+def aggregated_diatomics_results(results: dict[str, dict]) -> dict[str, float | None]:
+    """Aggregate curves using the reference metadata saved during inference.
+
+    bond_length_mae is in Å and well_depth_mae is in eV. avg_roughness is
+    the mean original finite-difference slope error normalized by a constant
+    energy dummy. Bond and well errors use their blind dummy scales. Each
+    term is capped at 1.
+    score is their equal-weight average, reported only for complete coverage.
+    Missing or model_excluded elements make score None. Raw metrics and
+    coverage are retained for diagnostics. A term without samples also makes
+    score None (for example, a reference set with no bound elements).
     """
     if not results:
         return _empty_diatomics_agg()
-
-    names = _diatomics_molecule_names()
+    reference = results["reference"]
+    names = reference["elements"]
+    dummy = reference["dummy_scales"]
+    curves = results["curves"]
     if not names:
         return _empty_diatomics_agg()
 
-    roughness_values = []
+    bond_errors: list[float] = []
+    well_errors: list[float] = []
+    roughness_values: list[float] = []
+    n_excluded = 0
     for name in names:
-        mol_results = results.get(name)
-        if mol_results is None:
-            return _empty_diatomics_agg()
-        roughness = mol_results.get("roughness")
-        if roughness is None or not np.isfinite(roughness):
-            return _empty_diatomics_agg()
-        roughness_values.append(roughness)
-
-    return {"avg_roughness": float(np.mean(roughness_values))}
+        record = curves.get(name)
+        if _element_is_excluded(record):
+            n_excluded += 1
+            continue
+        assert isinstance(record, dict)
+        bond = record.get("bond_length_error")
+        well = record.get("well_depth_error")
+        if bond is not None and np.isfinite(bond):
+            bond_errors.append(float(bond))
+        if well is not None and np.isfinite(well):
+            well_errors.append(float(well))
+        roughness_values.append(float(record["roughness"]))
+
+    coverage = (len(names) - n_excluded) / len(names)
+    bond_mae = _mean_or_none(bond_errors)
+    well_mae = _mean_or_none(well_errors)
+    avg_roughness = _mean_or_none(roughness_values)
+    bond_hat = _dummy_hat(bond_mae, dummy["bond_length_mae"])
+    well_hat = _dummy_hat(well_mae, dummy["well_depth_mae"])
+    if n_excluded or bond_hat is None or well_hat is None or avg_roughness is None:
+        score = None
+    else:
+        score = float((bond_hat + well_hat + avg_roughness) / 3.0)
+    return {
+        "bond_length_mae": bond_mae,
+        "well_depth_mae": well_mae,
+        "avg_roughness": avg_roughness,
+        "coverage": float(coverage),
+        "score": score,
+    }
 
 
 ####################################

lambench/metrics/vishelper/metrics_calculations.py

@@ -199,17 +199,17 @@ def _calculate_instability_error(self, cell: dict, lambda_0: float = 5e-4) -> fl
             else:
                 return np.clip(np.log10(slope / lambda_0), a_min=0, a_max=None)
 
-    def calculate_diatomics_roughness_results(self) -> dict[str, float]:
+    def calculate_diatomics_results(self) -> dict[str, float]:
         """
-        Leaderboard scores for homonuclear diatomics: avg_roughness (lower is better).
-        Models with missing results are omitted from this dict; the final ranking
-        still includes them with a null roughness column.
+        Leaderboard scores for homonuclear diatomics: complete-coverage score
+        (lower is better). Models with missing results are omitted from this
+        dict; the final ranking still includes them with a null column.
         """
         raw = self.fetcher.fetch_diatomics_results()
         return {
-            model: metrics["avg_roughness"]
+            model: metrics["score"]
             for model, metrics in raw.items()
-            if metrics is not None and metrics.get("avg_roughness") is not None
+            if metrics is not None and metrics.get("score") is not None
         }
 
     def calculate_efficiency_results(self) -> dict[str, float]:
@@ -238,7 +238,7 @@ def summarize_final_rankings(self):
         )
         stability_results = self.calculate_stability_results()
         efficiency_results = self.calculate_efficiency_results()
-        roughness_results = self.calculate_diatomics_roughness_results()
+        diatomics_results = self.calculate_diatomics_results()
         if not generalizability_ood or not generalizability_downstream:
             logging.warning(
                 "Missing data for generalizability metrics (ood or downstream)"
@@ -272,8 +272,8 @@ def summarize_final_rankings(self):
             "Applicability-Efficiency ↑": [
                 efficiency_results[model] for model in shared_models
             ],
-            "Applicability-Roughness ↓": [
-                roughness_results.get(model) for model in shared_models
+            "Applicability-Diatomics ↓": [
+                diatomics_results.get(model) for model in shared_models
             ],
         }
 
@@ -293,7 +293,7 @@ def summarize_final_rankings(self):
                 "Generalizability-PC Error ↓",
                 "Applicability-Instability ↓",
                 "Applicability-Efficiency ↑",
-                "Applicability-Roughness ↓",
+                "Applicability-Diatomics ↓",
             ],
             ascending=[True, True, True, False, True],
         )

lambench/metrics/vishelper/results_fetcher.py

@@ -130,7 +130,7 @@ def fetch_inference_efficiency_results(self) -> dict[str, dict[str, float]]:
         return results
 
     def fetch_diatomics_results(self) -> dict[str, dict]:
-        """Returns aggregated diatomics roughness results for all leaderboard models."""
+        """Returns aggregated homonuclear-diatomic scores for leaderboard models."""
         results = {}
         for model in self.leaderboard_models:
             task_results = CalculatorRecord.query(