feat:add diatomic test

lambench/metrics/post_process.py

@@ -16,6 +16,7 @@
     exp_average,
     aggregated_nve_md_results,
     aggregated_inference_efficiency_results,
+    aggregated_diatomics_results,
     get_leaderboard_models,
 )
 
@@ -173,6 +174,10 @@ def process_applicability_task_for_one_model(model: BaseLargeAtomModel):
             applicability_results[record.task_name] = (
                 aggregated_inference_efficiency_results(record.metrics)
             )
+        elif record.task_name == "homonuclear_diatomics":
+            applicability_results[record.task_name] = aggregated_diatomics_results(
+                record.metrics
+            )
     return applicability_results
 
 

lambench/metrics/utils.py

@@ -1,3 +1,6 @@
+import json
+from functools import lru_cache
+
 import numpy as np
 import yaml
 from typing import Optional, Literal
@@ -7,6 +10,14 @@
 from lambench.workflow.entrypoint import gather_model_params, gather_model
 from datetime import datetime
 
+_DIATOMICS_JSON = (
+    Path(__file__).resolve().parent.parent
+    / "tasks"
+    / "calculator"
+    / "diatomics"
+    / "diatomics.json"
+)
+
 #############################
 # General utility functions #
 #############################
@@ -151,6 +162,39 @@ def aggregated_inference_efficiency_results(
     }
 
 
+@lru_cache
+def _diatomics_molecule_names() -> tuple[str, ...]:
+    with open(_DIATOMICS_JSON) as fh:
+        return tuple(entry["name"] for entry in json.load(fh))
+
+
+def aggregated_diatomics_results(results: dict[str, dict]) -> dict[str, float]:
+    """
+    Aggregate per-molecule diatomics results.
+
+    avg_roughness: arithmetic mean of curvature RMSE (eV/Ų), used on the leaderboard.
+    Requires a finite roughness for every molecule in diatomics.json; otherwise None.
+    """
+    if not results:
+        return {"avg_roughness": None}
+
+    names = _diatomics_molecule_names()
+    if not names:
+        return {"avg_roughness": None}
+
+    roughness_values = []
+    for name in names:
+        mol_results = results.get(name)
+        if mol_results is None:
+            return {"avg_roughness": None}
+        roughness = mol_results.get("roughness")
+        if roughness is None or not np.isfinite(roughness):
+            return {"avg_roughness": None}
+        roughness_values.append(roughness)
+
+    return {"avg_roughness": float(np.mean(roughness_values))}
+
+
 ####################################
 # Visualization utility functions #
 ####################################

lambench/metrics/vishelper/metrics_calculations.py

@@ -197,6 +197,19 @@ 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]:
+        """
+        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.
+        """
+        raw = self.fetcher.fetch_diatomics_results()
+        return {
+            model: metrics["avg_roughness"]
+            for model, metrics in raw.items()
+            if metrics is not None and metrics.get("avg_roughness") is not None
+        }
+
     def calculate_efficiency_results(self) -> dict[str, float]:
         efficiency_results = self.fetcher.fetch_inference_efficiency_results()
         # filter out models with missing efficiency results
@@ -223,12 +236,14 @@ def summarize_final_rankings(self):
         )
         stability_results = self.calculate_stability_results()
         efficiency_results = self.calculate_efficiency_results()
+        roughness_results = self.calculate_diatomics_roughness_results()
         if not generalizability_ood or not generalizability_downstream:
             logging.warning(
                 "Missing data for generalizability metrics (ood or downstream)"
             )
             return
 
+        # Diatomics is not required for inclusion: missing scores stay None.
         shared_models = (
             set(generalizability_ood.keys())
             .intersection(set(generalizability_downstream.keys()))
@@ -255,6 +270,9 @@ 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
+            ],
         }
 
         # Create DataFrame with models as index
@@ -273,8 +291,9 @@ def summarize_final_rankings(self):
                 "Generalizability-PC Error ↓",
                 "Applicability-Instability ↓",
                 "Applicability-Efficiency ↑",
+                "Applicability-Roughness ↓",
             ],
-            ascending=[True, True, True, False],
+            ascending=[True, True, True, False, True],
         )
         print(
             "Final Rankings:\n",

lambench/metrics/vishelper/results_fetcher.py

@@ -12,6 +12,7 @@
     get_domain_to_direct_task_mapping,
     get_leaderboard_models,
     aggregated_inference_efficiency_results,
+    aggregated_diatomics_results,
 )
 from lambench.models.basemodel import BaseLargeAtomModel
 import pandas as pd
@@ -127,6 +128,24 @@ 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."""
+        results = {}
+        for model in self.leaderboard_models:
+            task_results = CalculatorRecord.query(
+                model_name=model.model_name, task_name="homonuclear_diatomics"
+            )
+            if len(task_results) != 1:
+                logging.warning(
+                    f"Expected one record for {model.model_name} and homonuclear_diatomics, "
+                    f"but got {len(task_results)}"
+                )
+                continue
+            results[model.model_metadata.pretty_name] = aggregated_diatomics_results(
+                task_results[0].metrics
+            )
+        return results
+
     def fetch_downstream_results(self) -> pd.DataFrame:
         """Returns downstream task results as a DataFrame with models as rows and task metrics as columns."""
 

lambench/models/ase_models.py

@@ -84,6 +84,9 @@ def __init__(self, *args, **kwargs):
     @property
     def calc(self) -> Calculator:
         """ASE Calculator with the model loaded."""
+        if self._calc is not None:
+            return self._calc
+
         calculator_dispatch = {
             "MACE": self._init_mace_calculator,
             "ORB": self._init_orb_calculator,
@@ -332,6 +335,10 @@ def evaluate(
                 from lambench.tasks.calculator.surface.surface import run_inference
 
                 assert task.test_data is not None
+                return {"metrics": run_inference(self, task.test_data)}
+            elif task.task_name == "homonuclear_diatomics":
+                from lambench.tasks.calculator.diatomics.diatomics import run_inference
+
                 return {"metrics": run_inference(self, task.test_data)}
             else:
                 raise NotImplementedError(f"Task {task.task_name} is not implemented.")

lambench/tasks/calculator/calculator_tasks.yml

@@ -51,3 +51,6 @@ interface:
 surface:
   test_data: /bohr/lambench-surfaces-43ll/v1/surface
   calculator_params: null
+homonuclear_diatomics:
+  test_data: null
+  calculator_params: null