Add homonuclear diatomic rank (#8)

* add homonuclear diatomic rank

* formatting

serve/leaderboard.py

@@ -1,3 +1,4 @@
+import importlib
 from pathlib import Path
 
 import pandas as pd
@@ -6,26 +7,29 @@
 from mlip_arena.models import REGISTRY as MODELS
 from mlip_arena.tasks import REGISTRY as TASKS
 
-import importlib
-
 DATA_DIR = Path("mlip_arena/tasks/diatomics")
 
-dfs = [pd.read_json(DATA_DIR / MODELS[model].get("family") /  "homonuclear-diatomics.json") for model in MODELS]
+dfs = [
+    pd.read_json(DATA_DIR / MODELS[model].get("family") / "homonuclear-diatomics.json")
+    for model in MODELS
+]
 df = pd.concat(dfs, ignore_index=True)
 
 
-table = pd.DataFrame(columns=[
-    "Model",
-    "Element Coverage",
-    # "No. of reversed forces",
-    # "Energy-consistent forces",
-    "Prediction",
-    "NVT",
-    "NPT",
-    "Code",
-    "Paper",
-    "First Release",
-])
+table = pd.DataFrame(
+    columns=[
+        "Model",
+        "Element Coverage",
+        # "No. of reversed forces",
+        # "Energy-consistent forces",
+        "Prediction",
+        "NVT",
+        "NPT",
+        "Code",
+        "Paper",
+        "First Release",
+    ]
+)
 
 for model in MODELS:
     rows = df[df["method"] == model]
@@ -48,22 +52,27 @@
 
 
 s = table.style.background_gradient(
-    cmap="PuRd",
-    subset=["Element Coverage"],
-    vmin=0, vmax=120
+    cmap="PuRd", subset=["Element Coverage"], vmin=0, vmax=120
 )
 
-st.warning("MLIP Arena is currently in **pre-alpha**. The results are not stable. Please interpret them with care.", icon="⚠️")
-st.info("Contributions are welcome. For more information, visit https://github.com/atomind-ai/mlip-arena.", icon="🤗")
+st.warning(
+    "MLIP Arena is currently in **pre-alpha**. The results are not stable. Please interpret them with care.",
+    icon="⚠️",
+)
+st.info(
+    "Contributions are welcome. For more information, visit https://github.com/atomind-ai/mlip-arena.",
+    icon="🤗",
+)
 
 st.markdown(
-"""
+    """
 <h1 style='text-align: center;'>⚔️ MLIP Arena Leaderboard ⚔️</h1>
 
 MLIP Arena is a platform for benchmarking foundation machine learning interatomic potentials (MLIPs), mainly for disclosing the learned physics and chemistry of the models and their performance on molecular dynamics (MD) simulations.
 The benchmarks are designed to evaluate the readiness and reliability of open-source, open-weight models to reproduce the qualitatively or quantitatively correct physics.
-""", unsafe_allow_html=True)
-
+""",
+    unsafe_allow_html=True,
+)
 
 
 st.dataframe(
@@ -85,19 +94,20 @@
 
 
 for task in TASKS:
+    if TASKS[task]["rank-page"] is None:
+        continue
 
     st.header(task, divider=True)
-
-    if TASKS[task]['rank-page'] is None:
-        st.write("Rank for this task is not available yet")
-        continue
+    st.page_link(
+        f"tasks/{TASKS[task]['task-page']}.py",
+        label="Link to task page",
+        icon=":material/link:",
+    )
 
     task_module = importlib.import_module(f"ranks.{TASKS[task]['rank-page']}")
 
-    # task_module = importlib.import_module(f".ranks", TASKS[task]["task-page"])
-
     #  Call the function from the imported module
-    if hasattr(task_module, 'get_rank_page'):
-        task_module.get_rank_page()
+    if hasattr(task_module, "render"):
+        task_module.render()
     else:
-        st.write("Results for the task are not available yet.")
+        st.write("Results for the task are not available yet.")

serve/ranks/homonuclear-diatomics.py

@@ -1,11 +1,134 @@
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import plotly.colors as pcolors
+import plotly.graph_objects as go
 import streamlit as st
+from ase.data import chemical_symbols
+from plotly.subplots import make_subplots
+from scipy.interpolate import CubicSpline
+
+from mlip_arena.models import REGISTRY as MODELS
+
+valid_models = [
+    model
+    for model, metadata in MODELS.items()
+    if Path(__file__).stem in metadata.get("gpu-tasks", [])
+]
+
+DATA_DIR = Path("mlip_arena/tasks/diatomics")
+
+dfs = [
+    pd.read_json(DATA_DIR / MODELS[model].get("family") / "homonuclear-diatomics.json")
+    for model in valid_models
+]
+df = pd.concat(dfs, ignore_index=True)
+
+table = pd.DataFrame()
+
+for model in valid_models:
+    rows = df[df["method"] == model]
+    metadata = MODELS.get(model, {})
+
+    new_row = {
+        "Model": model,
+        "Conservation deviation [eV/Å]": rows["conservation-deviation"].mean(),
+        "Spearman's coeff. (Energy - repulsion)": rows[
+            "spearman-repulsion-energy"
+        ].mean(),
+        "Spearman's coeff. (Force - descending)": rows[
+            "spearman-descending-force"
+        ].mean(),
+        "Tortuosity": rows["tortuosity"].mean(),
+        "Energy jump [eV]": rows["energy-jump"].mean(),
+        "Force flips": rows["force-flip-times"].mean(),
+        "Spearman's coeff. (Energy - attraction)": rows[
+            "spearman-attraction-energy"
+        ].mean(),
+        "Spearman's coeff. (Force - ascending)": rows[
+            "spearman-ascending-force"
+        ].mean(),
+    }
+
+    table = pd.concat([table, pd.DataFrame([new_row])], ignore_index=True)
+
+table.set_index("Model", inplace=True)
+
+table.sort_values("Conservation deviation [eV/Å]", ascending=True, inplace=True)
+table["Rank"] = np.argsort(table["Conservation deviation [eV/Å]"].to_numpy())
+
+table.sort_values(
+    "Spearman's coeff. (Energy - repulsion)", ascending=True, inplace=True
+)
+table["Rank"] += np.argsort(table["Spearman's coeff. (Energy - repulsion)"].to_numpy())
+
+table.sort_values(
+    "Spearman's coeff. (Force - descending)", ascending=True, inplace=True
+)
+table["Rank"] += np.argsort(table["Spearman's coeff. (Force - descending)"].to_numpy())
+
+table.sort_values("Tortuosity", ascending=True, inplace=True)
+table["Rank"] += np.argsort(table["Tortuosity"].to_numpy())
+
+table.sort_values("Energy jump [eV]", ascending=True, inplace=True)
+table["Rank"] += np.argsort(table["Energy jump [eV]"].to_numpy())
+
+table.sort_values("Force flips", ascending=True, inplace=True)
+table["Rank"] += np.argsort(table["Force flips"].to_numpy())
+
+table.sort_values("Rank", ascending=True, inplace=True)
+
+table["Rank aggr."] = table["Rank"]
+
+table["Rank"] = np.argsort(table["Rank"].to_numpy()) + 1
+
+# table.drop(columns=["rank"], inplace=True)
+# table = table.rename(columns={"Rank": "Rank Aggr."})
+
+table = table.reindex(
+    columns=[
+        "Rank",
+        "Rank aggr.",
+        "Conservation deviation [eV/Å]",
+        "Spearman's coeff. (Energy - repulsion)",
+        "Spearman's coeff. (Force - descending)",
+        "Tortuosity",
+        "Energy jump [eV]",
+        "Force flips",
+        "Spearman's coeff. (Energy - attraction)",
+        "Spearman's coeff. (Force - ascending)",
+    ]
+)
 
-def get_rank_page():
-    # st.markdown("""HIHI""")
-    pass
+s = (
+    table.style.background_gradient(
+        cmap="viridis_r",
+        subset=["Conservation deviation [eV/Å]"],
+        gmap=np.log(table["Conservation deviation [eV/Å]"].to_numpy()),
+    )
+    .background_gradient(
+        cmap="Reds",
+        subset=[
+            "Spearman's coeff. (Energy - repulsion)",
+            "Spearman's coeff. (Force - descending)",
+        ],
+        # vmin=-1, vmax=-0.5
+    )
+    .background_gradient(
+        cmap="RdPu",
+        subset=["Tortuosity", "Energy jump [eV]", "Force flips"],
+    )
+    .background_gradient(
+        cmap="Blues",
+        subset=["Rank", "Rank aggr."],
+    )
+)
 
-# def plot():
-#     pass
 
-# if __name__ == '__main__':
-#     pass
+def render():
+    st.dataframe(
+        s,
+        use_container_width=True,
+    )
+    # return table