Merge pull request #97 from deepmodeling/feat/add-md-stability-to-radar

Feat: add MD stability to radar

Author: anyangml

lambench/metrics/visualization.py

@@ -2,9 +2,15 @@
 from lambench.metrics.post_process import DIRECT_TASK_WEIGHTS
 from lambench.models.basemodel import BaseLargeAtomModel
 from lambench.databases.direct_predict_table import DirectPredictRecord
-from lambench.metrics.utils import get_domain_to_direct_task_mapping
-from lambench.metrics.utils import filter_direct_task_results, exp_average
+from lambench.databases.calculator_table import CalculatorRecord
+from lambench.metrics.utils import (
+    get_domain_to_direct_task_mapping,
+    aggregated_nve_md_results,
+    filter_direct_task_results,
+    exp_average,
+)
 from lambench.workflow.entrypoint import gather_models
+import numpy as np
 
 
 def aggregate_domain_results_for_one_model(model: BaseLargeAtomModel):
@@ -52,6 +58,30 @@ def aggregate_domain_results_for_one_model(model: BaseLargeAtomModel):
     return domain_results
 
 
+def fetch_stability_results(model: BaseLargeAtomModel) -> float:
+    """
+    Fetch stability metrics for a model based on NVE MD simulations.
+
+    The stability is measured as the energy drift slope minus the logarithm of the success rate divided by 1000.
+    A lower value indicates better stability.
+    """
+    task_results = CalculatorRecord.query(
+        model_name=model.model_name, task_name="nve_md"
+    )
+
+    if len(task_results) != 1:
+        logging.warning(
+            f"Expected one record for {model.model_name} and nve_md, but got {len(task_results)}"
+        )
+        return None
+
+    metrics = aggregated_nve_md_results(task_results[0].metrics)
+    slope = metrics["slope"]
+    success_rate = metrics["success_rate"]
+
+    return slope - np.log(success_rate) / 1000  # to penalize failed simulations
+
+
 def aggregate_domain_results():
     """
     This function aggregates the results across models and domains.
@@ -67,10 +97,9 @@ def aggregate_domain_results():
     ]
 
     for model in leaderboard_models:
-        results[model.model_name] = aggregate_domain_results_for_one_model(model)
+        domain_results = aggregate_domain_results_for_one_model(model)
+        stability = fetch_stability_results(model)
+        domain_results["Stability"] = stability
+        results[model.model_name] = domain_results
 
     return results
-
-
-if __name__ == "__main__":
-    print(aggregate_domain_results())