add version tag

.github/README.md

@@ -10,7 +10,7 @@
 > [!NOTE]
 > If you're interested in joining the effort, please reach out to Yuan at [cyrusyc@berkeley.edu](mailto:cyrusyc@berkeley.edu).
 
-MLIP Arena is an open-source platform for benchmarking machine learning interatomic potentials (MLIPs). The platform provides a unified interface for users to evaluate the performance of their models on a variety of tasks, including single-point density functional theory calculations and molecular dynamics simulations. The platform is designed to be extensible, allowing users to contribute new models, benchmarks, and training data to the platform.
+MLIP Arena is a platform for evaluating foundation machine learning interatomic potentials (MLIPs) beyond conventional energy and force error metrics. It focuses on revealing the underlying physics and chemistry learned by these models and assessing their performance in molecular dynamics (MD) simulations. The platform's benchmarks are specifically designed to evaluate the readiness and reliability of open-source, open-weight models in accurately reproducing both qualitative and quantitative behaviors of atomic systems.
 
 ## Contribute
 

mlip_arena/models/chgnet.py

@@ -1,54 +0,0 @@
-from typing import Optional, Tuple
-
-import numpy as np
-import torch
-from ase import Atoms
-from ase.calculators.calculator import all_changes
-from huggingface_hub import hf_hub_download
-from torch_geometric.data import Data
-
-from mlip_arena.models import MLIP, MLIPCalculator
-
-# TODO: WIP
-
-
-class CHGNet(MLIPCalculator):
-    def __init__(
-        self,
-        device: torch.device | None = None,
-        restart=None,
-        atoms=None,
-        directory=".",
-        **kwargs,
-    ):
-        self.device = device or torch.device(
-            "cuda" if torch.cuda.is_available() else "cpu"
-        )
-
-        super().__init__(
-            model=model, restart=restart, atoms=atoms, directory=directory, **kwargs
-        )
-
-        self.name: str = self.__class__.__name__
-        self.implemented_properties = ["energy", "forces", "stress"]
-
-    def calculate(
-        self, atoms: Atoms, properties: list[str], system_changes: list = all_changes
-    ):
-        """Calculate energies and forces for the given Atoms object"""
-        super().calculate(atoms, properties, system_changes)
-
-        output = self.forward(atoms)
-
-        self.results = {}
-        if "energy" in properties:
-            self.results["energy"] = output["energy"].item()
-        if "forces" in properties:
-            self.results["forces"] = output["forces"].cpu().detach().numpy()
-        if "stress" in properties:
-            self.results["stress"] = output["stress"].cpu().detach().numpy()
-
-    def forward(self, x: Data | Atoms) -> dict[str, torch.Tensor]:
-        """Implement data conversion, graph creation, and model forward pass"""
-        # TODO
-        raise NotImplementedError

mlip_arena/models/externals.py

@@ -1,283 +0,0 @@
-from __future__ import annotations
-
-import os
-from pathlib import Path
-from typing import Literal
-
-import matgl
-import requests
-import torch
-from alignn.ff.ff import AlignnAtomwiseCalculator, get_figshare_model_ff, default_path
-from ase import Atoms
-from chgnet.model.dynamics import CHGNetCalculator
-from chgnet.model.model import CHGNet as CHGNetModel
-from fairchem.core import OCPCalculator
-from mace.calculators import MACECalculator
-from matgl.ext.ase import PESCalculator
-from orb_models.forcefield import pretrained
-from orb_models.forcefield.calculator import ORBCalculator
-from sevenn.sevennet_calculator import SevenNetCalculator
-
-
-# Avoid circular import
-def get_freer_device() -> torch.device:
-    """Get the GPU with the most free memory, or use MPS if available.
-    s
-        Returns:
-            torch.device: The selected GPU device or MPS.
-
-        Raises:
-            ValueError: If no GPU or MPS is available.
-    """
-    device_count = torch.cuda.device_count()
-    if device_count > 0:
-        # If CUDA GPUs are available, select the one with the most free memory
-        mem_free = [
-            torch.cuda.get_device_properties(i).total_memory
-            - torch.cuda.memory_allocated(i)
-            for i in range(device_count)
-        ]
-        free_gpu_index = mem_free.index(max(mem_free))
-        device = torch.device(f"cuda:{free_gpu_index}")
-        print(
-            f"Selected GPU {device} with {mem_free[free_gpu_index] / 1024**2:.2f} MB free memory from {device_count} GPUs"
-        )
-    elif torch.backends.mps.is_available():
-        # If no CUDA GPUs are available but MPS is, use MPS
-        print("No GPU available. Using MPS.")
-        device = torch.device("mps")
-    else:
-        # Fallback to CPU if neither CUDA GPUs nor MPS are available
-        print("No GPU or MPS available. Using CPU.")
-        device = torch.device("cpu")
-
-    return device
-
-
-class MACE_MP_Medium(MACECalculator):
-    def __init__(
-        self,
-        checkpoint="http://tinyurl.com/5yyxdm76",
-        device: str | None = None,
-        default_dtype="float32",
-        **kwargs,
-    ):
-        cache_dir = Path.home() / ".cache" / "mace"
-        checkpoint_url_name = "".join(
-            c for c in os.path.basename(checkpoint) if c.isalnum() or c in "_"
-        )
-        cached_model_path = f"{cache_dir}/{checkpoint_url_name}"
-        if not os.path.isfile(cached_model_path):
-            import urllib
-
-            os.makedirs(cache_dir, exist_ok=True)
-            _, http_msg = urllib.request.urlretrieve(checkpoint, cached_model_path)
-            if "Content-Type: text/html" in http_msg:
-                raise RuntimeError(
-                    f"Model download failed, please check the URL {checkpoint}"
-                )
-        model = cached_model_path
-
-        device = device or str(get_freer_device())
-
-        super().__init__(
-            model_paths=model, device=device, default_dtype=default_dtype, **kwargs
-        )
-
-
-# TODO: could share the same class with MACE_MP_Medium
-class MACE_OFF_Medium(MACECalculator):
-    def __init__(
-        self,
-        checkpoint="https://github.com/ACEsuit/mace-off/raw/main/mace_off23/MACE-OFF23_medium.model?raw=true",
-        device: str | None = None,
-        default_dtype="float32",
-        **kwargs,
-    ):
-        cache_dir = Path.home() / ".cache" / "mace"
-        checkpoint_url_name = "".join(
-            c for c in os.path.basename(checkpoint) if c.isalnum() or c in "_"
-        )
-        cached_model_path = f"{cache_dir}/{checkpoint_url_name}"
-        if not os.path.isfile(cached_model_path):
-            import urllib
-
-            os.makedirs(cache_dir, exist_ok=True)
-            _, http_msg = urllib.request.urlretrieve(checkpoint, cached_model_path)
-            if "Content-Type: text/html" in http_msg:
-                raise RuntimeError(
-                    f"Model download failed, please check the URL {checkpoint}"
-                )
-        model = cached_model_path
-
-        device = device or str(get_freer_device())
-
-        super().__init__(
-            model_paths=model, device=device, default_dtype=default_dtype, **kwargs
-        )
-
-
-class CHGNet(CHGNetCalculator):
-    def __init__(
-        self,
-        checkpoint: CHGNetModel | None = None,  # TODO: specifiy version
-        device: str | None = None,
-        stress_weight: float | None = 1 / 160.21766208,
-        on_isolated_atoms: Literal["ignore", "warn", "error"] = "warn",
-        **kwargs,
-    ) -> None:
-        use_device = device or str(get_freer_device())
-        super().__init__(
-            model=checkpoint,
-            use_device=use_device,
-            stress_weight=stress_weight,
-            on_isolated_atoms=on_isolated_atoms,
-            **kwargs,
-        )
-
-    def calculate(
-        self,
-        atoms: Atoms | None = None,
-        properties: list | None = None,
-        system_changes: list | None = None,
-    ) -> None:
-        super().calculate(atoms, properties, system_changes)
-
-        # for ase.io.write compatibility
-        self.results.pop("crystal_fea", None)
-
-
-class M3GNet(PESCalculator):
-    def __init__(
-        self,
-        checkpoint="M3GNet-MP-2021.2.8-PES",
-        # TODO: cannot assign device
-        state_attr: torch.Tensor | None = None,
-        stress_weight: float = 1.0,
-        **kwargs,
-    ) -> None:
-        potential = matgl.load_model(checkpoint)
-        super().__init__(potential, state_attr, stress_weight, **kwargs)
-
-
-class EquiformerV2(OCPCalculator):
-    def __init__(
-        self,
-        checkpoint="EquiformerV2-lE4-lF100-S2EFS-OC22",  # TODO: import from registry
-        # TODO: cannot assign device
-        local_cache="/tmp/ocp/",
-        cpu=False,
-        seed=0,
-        **kwargs,
-    ) -> None:
-        super().__init__(
-            model_name=checkpoint,
-            local_cache=local_cache,
-            cpu=cpu,
-            seed=seed,
-            **kwargs,
-        )
-
-    def calculate(self, atoms: Atoms, properties, system_changes) -> None:
-        super().calculate(atoms, properties, system_changes)
-
-        self.results.update(
-            force=atoms.get_forces(),
-        )
-
-
-class EquiformerV2OC20(OCPCalculator):
-    def __init__(
-        self,
-        checkpoint="EquiformerV2-31M-S2EF-OC20-All+MD",  # TODO: import from registry
-        # TODO: cannot assign device
-        local_cache="/tmp/ocp/",
-        cpu=False,
-        seed=0,
-        **kwargs,
-    ) -> None:
-        super().__init__(
-            model_name=checkpoint,
-            local_cache=local_cache,
-            cpu=cpu,
-            seed=seed,
-            **kwargs,
-        )
-
-
-class eSCN(OCPCalculator):
-    def __init__(
-        self,
-        checkpoint="eSCN-L6-M3-Lay20-S2EF-OC20-All+MD",  # TODO: import from registry
-        # TODO: cannot assign device
-        local_cache="/tmp/ocp/",
-        cpu=False,
-        seed=0,
-        **kwargs,
-    ) -> None:
-        super().__init__(
-            model_name=checkpoint,
-            local_cache=local_cache,
-            cpu=cpu,
-            seed=seed,
-            **kwargs,
-        )
-
-    def calculate(self, atoms: Atoms, properties, system_changes) -> None:
-        super().calculate(atoms, properties, system_changes)
-
-        self.results.update(
-            force=atoms.get_forces(),
-        )
-
-
-class ALIGNN(AlignnAtomwiseCalculator):
-    def __init__(self, device=None, **kwargs) -> None:
-        # TODO: cannot control version
-        # _ = get_figshare_model_ff(dir_path=dir_path)
-        model_path = default_path()
-
-        device = device or get_freer_device()
-        super().__init__(path=model_path, device=device, **kwargs)
-
-
-class SevenNet(SevenNetCalculator):
-    def __init__(
-        self,
-        checkpoint="7net-0",  # TODO: import from registry
-        device=None,
-        **kwargs,
-    ):
-        device = device or get_freer_device()
-        super().__init__(checkpoint, device=device, **kwargs)
-
-
-class ORB(ORBCalculator):
-    def __init__(
-        self,
-        checkpoint="orbff-v1-20240827.ckpt",
-        device=None,
-        **kwargs,
-    ):
-        device = device or get_freer_device()
-
-        cache_dir = Path.home() / ".cache" / "orb"
-        cache_dir.mkdir(parents=True, exist_ok=True)
-        ckpt_path = cache_dir / "orbff-v1-20240827.ckpt"
-
-        url = f"https://storage.googleapis.com/orbitalmaterials-public-models/forcefields/{checkpoint}"
-
-        if not ckpt_path.exists():
-            print(f"Downloading ORB model from {url} to {ckpt_path}...")
-            try:
-                response = requests.get(url, stream=True, timeout=120)
-                response.raise_for_status()
-                with open(ckpt_path, "wb") as f:
-                    for chunk in response.iter_content(chunk_size=8192):
-                        f.write(chunk)
-                print("Download completed.")
-            except requests.exceptions.RequestException as e:
-                raise RuntimeError("Failed to download ORB model.") from e
-
-        orbff = pretrained.orb_v1(weights_path=ckpt_path, device=device)
-        super().__init__(orbff, device=device, **kwargs)

mlip_arena/models/registry.yaml

@@ -50,7 +50,7 @@ M3GNet:
   class: M3GNet
   family: matgl
   package: matgl==1.1.2
-  checkpoint:
+  checkpoint: M3GNet-MP-2021.2.8-PES
   username: cyrusyc
   last-update: 2024-07-08T00:00:00
   datetime: 2024-07-08T00:00:00
@@ -135,6 +135,7 @@ eqV2(OMat):
   prediction: EFS
   nvt: true
   npt: true
+  date: 2024-10-18
   github: https://github.com/FAIR-Chem/fairchem
   doi: https://arxiv.org/abs/2410.12771
 
@@ -206,6 +207,8 @@ MACE-OFF(M):
   module: externals
   class: MACE_OFF_Medium
   family: mace-off
+  package: mace-torch==0.3.4
+  checkpoint: MACE-OFF23_medium.model
   username: cyrusyc
   last-update: 2024-03-25T14:30:00
   datetime: 2024-03-25T14:30:00 # TODO: Fake datetime
@@ -228,7 +231,7 @@ ALIGNN:
   class: ALIGNN
   family: alignn
   package: alignn==2024.5.27
-  checkpoint:
+  checkpoint: 2024.5.27
   username: cyrusyc
   last-update: 2024-07-08T00:00:00
   datetime: 2024-07-08T00:00:00

pyproject.toml

@@ -163,6 +163,7 @@ ignore = [
     "FBT003",  # boolean positional variable in function call
     "PERF203",  # `try`-`except` within a loop incurs performance overhead (no overhead in Py 3.11+)
     "F405",  # 'module' may be undefined, or defined from star imports
+    "C0301",  # Line too long
 ]
 fixable = ["ALL"]
 pydocstyle.convention = "google"

serve/leaderboard.py

@@ -28,6 +28,7 @@ table = pd.DataFrame(
         "Training Set",
         "Code",
         "Paper",
+        "Checkpoint",
         "First Release",
     ]
 )
@@ -46,6 +47,7 @@ for model in MODELS:
         "Training Set": metadata.get("datasets", []),
         "Code": metadata.get("github", None) if metadata else None,
         "Paper": metadata.get("doi", None) if metadata else None,
+        "Checkpoint": metadata.get("checkpoint", None),
         "First Release": metadata.get("date", None),
     }
     table = pd.concat([table, pd.DataFrame([new_row])], ignore_index=True)
@@ -73,11 +75,11 @@ st.markdown(
 
 ### :red[Introduction]
 
-Foundation machine learning interatomic potentials (fMLIPs) trained on extensive databases containing millions of density functional theory (DFT) calculations have shown remarkable zero-shot predictive capabilities for complex atomic interactions. These potentials derive quantum mechanical insights with high accuracy, expressivity, and generalizability, significantly outperforming classical empirical force fields while maintaining comparable computational efficiency.
+Foundation machine learning interatomic potentials (fMLIPs), trained on extensive databases containing millions of density functional theory (DFT) calculations, have demonstrated remarkable zero-shot predictive capabilities for complex atomic interactions. These potentials derive quantum mechanical insights with high accuracy, expressivity, and generalizability, significantly outperforming classical empirical force fields while maintaining comparable computational efficiency.
 
-However, MLIPs trained on atomic energy and force labels do not neccessarily capture the correct atomic interactions, even though they usually excel in error-based metrics for bulk systems. To drive further advancements in this field, it is crucial to establish mechanisms that ensure fair and transparent benchmarking practices that go beyond basic regression metrics.
+However, MLIPs trained on atomic energy and force labels do not necessarily capture the correct atomic interactions, even though they often excel in error-based metrics for bulk systems. To drive further advancements in this field, it is crucial to establish mechanisms that ensure fair and transparent benchmarking practices beyond basic regression metrics.
 
-MLIP Arena aims to provide an equitable and transparent platform for benchmarking MLIPs in a collaborative, crowd-sourced environment. Its primary goal is to uncover the learned physics and chemistry of open-source, open-weight MLIPs. The benchmarks are designed to be agnostic to both the underlying architecture and the specific training targets, such as density functionals, ensuring a comprehensive and unbiased evaluation.
+MLIP Arena aims to provide a fair and transparent platform for benchmarking MLIPs in a crowdsourced setting. Its primary goal is to uncover the learned physics and chemistry of open-source, open-weight MLIPs. The benchmarks are designed to be agnostic to both the underlying architecture and specific training targets, such as density functionals, ensuring a cross-comparable and unbiased evaluation.
 
 """,
     unsafe_allow_html=True,

serve/tasks/stability.py

@@ -19,7 +19,7 @@ st.markdown("""
 Stable and accurate molecular dynamics (MD) simulations are important for understanding the properties of matters.
 However, many MLIPs have unphysical potential energy surface (PES) at the short-range interatomic distances or under many-body effect. These are often manifested as softened repulsion and hole in the PES and can lead to incorrect and sampling of the phase space.
 
-Here, we analyze the stability of the MD simulations under high pressure conditions by gradually increasing the pressure from 0 to 1000 GPa at 300K until the system crashes or completes 100 ps trajectory. This benchmark also explores faster the far-from-equilibrium dynamics of the system and the "durability" of the MLIPs under extreme conditions.
+Here, we analyze the stability of the MD simulations under high pressure conditions by gradually increasing the pressure from 0 to 1000 GPa at 300K until the system crashes or completes 100 ps trajectory. This benchmark also explores faster the far-from-equilibrium dynamics of the system and the durability of the MLIPs under extreme conditions.
 """)
 
 st.markdown("### Methods")
@@ -80,7 +80,6 @@ method_color_mapping = {
 ###
 
 # Determine the bin edges for the entire dataset to keep them consistent across groups
-# bins = np.histogram_bin_edges(df['total_steps'], bins=10)
 
 max_steps = df["total_steps"].max()
 max_target_steps = df["target_steps"].max()
@@ -89,14 +88,11 @@ bins = np.append(np.arange(0, max_steps + 1, max_steps // 10), max_target_steps)
 bin_labels = [f"{bins[i]}-{bins[i+1]}" for i in range(len(bins) - 1)]
 
 num_bins = len(bin_labels)
-# colormap = px.colors.sequential.Darkmint_r
 colormap = px.colors.sequential.YlOrRd_r
 indices = np.linspace(0, len(colormap) - 1, num_bins, dtype=int)
 bin_colors = [colormap[i] for i in indices]
-# bin_colors[-1] = px.colors.sequential.Greens[-1]
 
 # Initialize a dictionary to hold the counts for each method and bin range
-# counts_per_method = {method: [0] * len(bin_labels) for method in df["method"].unique()}
 counts_per_method = {method: [0] * len(bin_labels) for method in df["method"].unique()}
 
 
@@ -160,7 +156,7 @@ plot_md_steps(counts_per_method, count_or_percetange)
 
 st.markdown(
     """
-> The histogram shows the distribution of the total number of MD steps before the system crashes or completes the trajectory. :red[The color of the bins indicates the number of steps in the bin]. :blue[The height of the bars indicates the percentage of each bin among all the runs].
+> The histogram shows the distribution of the total number of MD steps before the system crashes or completes the trajectory. :red[The color of the bins indicates the number of steps in the bin]. :blue[The height of the bars indicates the number or percentage of each bin among all the runs].
 """
 )